Update app.py

app.py

@@ -9,13 +9,15 @@ import torch
 from PIL import Image
 from diffusers import FluxInpaintPipeline
 
+torch.cuda.empty_cache()
+
 MARKDOWN = """
-# FLUX.1 Inpainting using 
-[Black Forest Labs](https://huggingface.co/black-forest-labs)' "FLUX.1-dev"
+# FLUX Inpainting 
+Model used FLUX.1-schnell.
 """
 
 MAX_SEED = np.iinfo(np.int32).max
-IMAGE_SIZE = 1024
+IMAGE_SIZE = 512  
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
@@ -34,32 +36,32 @@ def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
     return image
 
 
-EXAMPLES = [
-    [
-        {
-            "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
-            "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-2.png", stream=True).raw))],
-            "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-2.png", stream=True).raw),
-        },
-        "little lion",
-        42,
-        False,
-        0.85,
-        30
-    ],
-    [
-        {
-            "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
-            "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-3.png", stream=True).raw))],
-            "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-3.png", stream=True).raw),
-        },
-        "tribal tattoos",
-        42,
-        False,
-        0.85,
-        30
-    ]
-]
+# EXAMPLES = [
+#     [
+#         {
+#             "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
+#             "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-2.png", stream=True).raw))],
+#             "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-2.png", stream=True).raw),
+#         },
+#         "little lion",
+#         42,
+#         False,
+#         0.85,
+#         30
+#     ],
+#     [
+#         {
+#             "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
+#             "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-3.png", stream=True).raw))],
+#             "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-3.png", stream=True).raw),
+#         },
+#         "tribal tattoos",
+#         42,
+#         False,
+#         0.85,
+#         30
+#     ]
+# ]
 
 pipe = FluxInpaintPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
@@ -122,17 +124,18 @@ def process(
     if randomize_seed_checkbox:
         seed_slicer = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed_slicer)
-    result = pipe(
-        prompt=input_text,
-        image=resized_image,
-        mask_image=resized_mask,
-        width=width,
-        height=height,
-        strength=strength_slider,
-        generator=generator,
-        num_inference_steps=num_inference_steps_slider
-    ).images[0]
-    print('INFERENCE DONE')
+    with torch.no_grad(), torch.autocast("cuda"):
+        result = pipe(
+            prompt=input_text,
+            image=resized_image,
+            mask_image=resized_mask,
+            width=width,
+            height=height,
+            strength=strength_slider,
+            generator=generator,
+            num_inference_steps=num_inference_steps_slider
+        ).images[0]
+    torch.cuda.empty_cache()
     return result, resized_mask
 
 
@@ -174,9 +177,9 @@ with gr.Blocks() as demo:
                 with gr.Row():
                     strength_slider_component = gr.Slider(
                         label="Strength",
-                        info="Indicates extent to transform the reference `image`. "
-                             "Must be between 0 and 1. `image` is used as a starting "
-                             "point and more noise is added the higher the `strength`.",
+                        # info="Indicates extent to transform the reference `image`. "
+                        #      "Must be between 0 and 1. `image` is used as a starting "
+                        #      "point and more noise is added the higher the `strength`.",
                         minimum=0,
                         maximum=1,
                         step=0.01,
@@ -185,10 +188,10 @@ with gr.Blocks() as demo:
 
                     num_inference_steps_slider_component = gr.Slider(
                         label="Number of inference steps",
-                        info="The number of denoising steps. More denoising steps "
-                             "usually lead to a higher quality image at the",
+                        # info="The number of denoising steps. More denoising steps "
+                        #      "usually lead to a higher quality image at the",
                         minimum=1,
-                        maximum=50,
+                        maximum=20,
                         step=1,
                         value=20,
                     )
@@ -198,25 +201,25 @@ with gr.Blocks() as demo:
             with gr.Accordion("Debug", open=False):
                 output_mask_component = gr.Image(
                     type='pil', image_mode='RGB', label='Input mask', format="png")
-    with gr.Row():
-        gr.Examples(
-            fn=process,
-            examples=EXAMPLES,
-            inputs=[
-                input_image_editor_component,
-                input_text_component,
-                seed_slicer_component,
-                randomize_seed_checkbox_component,
-                strength_slider_component,
-                num_inference_steps_slider_component
-            ],
-            outputs=[
-                output_image_component,
-                output_mask_component
-            ],
-            run_on_click=True,
-            cache_examples=True
-        )
+    # with gr.Row():
+    #     gr.Examples(
+    #         fn=process,
+    #         examples=EXAMPLES,
+    #         inputs=[
+    #             input_image_editor_component,
+    #             input_text_component,
+    #             seed_slicer_component,
+    #             randomize_seed_checkbox_component,
+    #             strength_slider_component,
+    #             num_inference_steps_slider_component
+    #         ],
+    #         outputs=[
+    #             output_image_component,
+    #             output_mask_component
+    #         ],
+    #         run_on_click=True,
+    #         cache_examples=True
+    #     )
 
     submit_button_component.click(
         fn=process,
@@ -234,5 +237,4 @@ with gr.Blocks() as demo:
         ]
     )
 
-demo.launch(debug=False, show_error=True)
-
+demo.launch(debug=False, show_error=True,share= True)