Update app.py

app.py

@@ -9,15 +9,13 @@ import torch
 from PIL import Image
 from diffusers import FluxInpaintPipeline
 
-torch.cuda.empty_cache()
-
 MARKDOWN = """
-# FLUX Inpainting 
-Model used FLUX.1-schnell.
+# FLUX.1 Inpainting using 
+[Black Forest Labs](https://huggingface.co/black-forest-labs)'s FLUX.1-schnell
 """
 
 MAX_SEED = np.iinfo(np.int32).max
-IMAGE_SIZE = 512  
+IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
@@ -36,35 +34,35 @@ 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)
+    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(DEVICE)
 
 
 def resize_image_dimensions(
@@ -124,18 +122,17 @@ def process(
     if randomize_seed_checkbox:
         seed_slicer = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed_slicer)
-    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()
+    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')
     return result, resized_mask
 
 
@@ -177,9 +174,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,
@@ -188,10 +185,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=20,
+                        maximum=50,
                         step=1,
                         value=20,
                     )
@@ -201,25 +198,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,
@@ -237,4 +234,4 @@ with gr.Blocks() as demo:
         ]
     )
 
-demo.launch(debug=False, show_error=True)
+demo.launch(debug=False, show_error=True, share=True)