Update app.py

app.py

@@ -1,154 +1,247 @@
-from typing import Tuple, Dict
+from typing import Tuple
 import requests
 import random
 import numpy as np
 import gradio as gr
+import spaces
 import torch
-from PIL import Image
-from diffusers import StableDiffusionInpaintPipeline
-
-INFO = """
-# FLUX-Based Inpainting 🎨
-
-This interface utilizes a FLUX model variant for precise inpainting. Special thanks to the [Black Forest Labs](https://huggingface.co/black-forest-labs) team 
-and [Gothos](https://github.com/Gothos) for contributing to this advanced solution.
+from PIL import Image, UnidentifiedImageError
+from diffusers import FluxInpaintPipeline
+
+MARKDOWN = """
+# FLUX.1 Inpainting 🔥
+Shoutout to [Black Forest Labs](https://huggingface.co/black-forest-labs) team for 
+creating this amazing model, and a big thanks to [Gothos](https://github.com/Gothos) 
+for taking it to the next level by enabling inpainting with the FLUX.
 """
 
-# Constants
-MAX_SEED_VALUE = np.iinfo(np.int32).max
-TARGET_DIM = 1024
+MAX_SEED = np.iinfo(np.int32).max
+IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Function to clear background
-def clear_background(image: Image.Image, threshold: int = 50) -> Image.Image:
+
+def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
     image = image.convert("RGBA")
-    pixels = image.getdata()
-    processed_data = [
-        (0, 0, 0, 0) if sum(pixel[:3]) / 3 < threshold else pixel for pixel in pixels
-    ]
-    image.putdata(processed_data)
+    data = image.getdata()
+    new_data = []
+    for item in data:
+        avg = sum(item[:3]) / 3
+        if avg < threshold:
+            new_data.append((0, 0, 0, 0))
+        else:
+            new_data.append(item)
+
+    image.putdata(new_data)
     return image
 
-# Sample data examples
+
+def load_image(url: str) -> Image.Image:
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # Raise an HTTPError for bad responses
+        image = Image.open(BytesIO(response.content))
+        return image
+    except requests.HTTPError as http_err:
+        print(f"HTTP error occurred: {http_err}")
+        return None
+    except UnidentifiedImageError:
+        print("Cannot identify image file")
+        return None
+    except Exception as err:
+        print(f"Other error occurred: {err}")
+        return None
+
+
 EXAMPLES = [
     [
         {
-            "background": Image.open(requests.get("https://example.com/doge-1.png", stream=True).raw),
-            "layers": [clear_background(Image.open(requests.get("https://example.com/mask-1.png", stream=True).raw))],
-            "composite": Image.open(requests.get("https://example.com/composite-1.png", stream=True).raw),
+            "background": load_image("https://media.roboflow.com/spaces/doge-2-image.png"),
+            "layers": [remove_background(load_image("https://media.roboflow.com/spaces/doge-2-mask-2.png"))],
+            "composite": load_image("https://media.roboflow.com/spaces/doge-2-composite-2.png"),
         },
-        "desert mirage",
+        "little lion",
         42,
         False,
-        0.75,
-        25
+        0.85,
+        30
     ],
     [
         {
-            "background": Image.open(requests.get("https://example.com/doge-2.png", stream=True).raw),
-            "layers": [clear_background(Image.open(requests.get("https://example.com/mask-2.png", stream=True).raw))],
-            "composite": Image.open(requests.get("https://example.com/composite-2.png", stream=True).raw),
+            "background": load_image("https://media.roboflow.com/spaces/doge-2-image.png"),
+            "layers": [remove_background(load_image("https://media.roboflow.com/spaces/doge-2-mask-3.png"))],
+            "composite": load_image("https://media.roboflow.com/spaces/doge-2-composite-3.png"),
         },
-        "neon city",
-        100,
-        True,
-        0.9,
-        35
+        "tribal tattoos",
+        42,
+        False,
+        0.85,
+        30
     ]
 ]
 
-# Load model
-inpainting_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+pipe = FluxInpaintPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
 
-# Utility to adjust image size
-def get_scaled_dimensions(
-    original_size: Tuple[int, int], max_dim: int = TARGET_DIM
+
+def resize_image_dimensions(
+    original_resolution_wh: Tuple[int, int],
+    maximum_dimension: int = IMAGE_SIZE
 ) -> Tuple[int, int]:
-    width, height = original_size
-    scaling_factor = max_dim / max(width, height)
-    return (int(width * scaling_factor) // 32 * 32, int(height * scaling_factor) // 32 * 32)
+    width, height = original_resolution_wh
+
+    if width > height:
+        scaling_factor = maximum_dimension / width
+    else:
+        scaling_factor = maximum_dimension / height
+
+    new_width = int(width * scaling_factor)
+    new_height = int(height * scaling_factor)
+
+    new_width = new_width - (new_width % 32)
+    new_height = new_height - (new_height % 32)
+
+    return new_width, new_height
+
 
 @spaces.GPU(duration=100)
-def generate_inpainting(
-    input_data: Dict,
-    prompt_text: str,
-    chosen_seed: int,
-    use_random_seed: bool,
-    inpainting_strength: float,
-    steps: int,
+def process(
+    input_image_editor: dict,
+    input_text: str,
+    seed_slicer: int,
+    randomize_seed_checkbox: bool,
+    strength_slider: float,
+    num_inference_steps_slider: int,
     progress=gr.Progress(track_tqdm=True)
 ):
-    if not prompt_text:
-        return gr.Info("Provide a prompt to proceed."), None
+    if not input_text:
+        return None, None, "Please enter a text prompt."
 
-    background = input_data.get("background")
-    mask_layer = input_data.get("layers")[0]
+    image = input_image_editor.get('background')
+    mask = input_image_editor.get('layers', [None])[0]
 
-    if not background:
-        return gr.Info("Background image is missing."), None
+    if not image:
+        return None, None, "Please upload an image."
 
-    if not mask_layer:
-        return gr.Info("Mask layer is missing."), None
+    if not mask:
+        return None, None, "Please draw a mask on the image."
 
-    new_width, new_height = get_scaled_dimensions(background.size)
-    resized_background = background.resize((new_width, new_height), Image.LANCZOS)
-    resized_mask = mask_layer.resize((new_width, new_height), Image.LANCZOS)
+    width, height = resize_image_dimensions(original_resolution_wh=image.size)
+    resized_image = image.resize((width, height), Image.LANCZOS)
+    resized_mask = mask.resize((width, height), Image.LANCZOS)
 
-    if use_random_seed:
-        chosen_seed = random.randint(0, MAX_SEED_VALUE)
-
-    torch.manual_seed(chosen_seed)
-    generated_image = inpainting_pipeline(
-        prompt=prompt_text,
-        image=resized_background,
+    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,
-        strength=inpainting_strength,
-        num_inference_steps=steps,
+        width=width,
+        height=height,
+        strength=strength_slider,
+        generator=generator,
+        num_inference_steps=num_inference_steps_slider
     ).images[0]
+    
+    return result, resized_mask, None
 
-    return generated_image, resized_mask
-
-# Build the Gradio interface
-with gr.Blocks() as flux_app:
-    gr.Markdown(INFO)
 
+with gr.Blocks() as demo:
+    gr.Markdown(MARKDOWN)
     with gr.Row():
         with gr.Column():
-            image_editor = gr.ImageEditor(
-                label="Edit Image",
-                type="pil",
+            input_image_editor_component = gr.ImageEditor(
+                label='Image',
+                type='pil',
                 sources=["upload", "webcam"],
-                brush=gr.Brush(colors=["#FFF"], color_mode="fixed")
-            )
-
-            prompt_box = gr.Text(
-                label="Inpainting Prompt", placeholder="Describe the change you'd like."
-            )
-            run_button = gr.Button(value="Run Inpainting")
-
-            with gr.Accordion("Settings"):
-                seed_slider = gr.Slider(0, MAX_SEED_VALUE, step=1, value=42, label="Seed")
-                random_seed_toggle = gr.Checkbox(label="Randomize Seed", value=True)
-                inpainting_strength_slider = gr.Slider(0.0, 1.0, step=0.01, value=0.85, label="Inpainting Strength")
-                steps_slider = gr.Slider(1, 50, step=1, value=25, label="Inference Steps")
-
+                image_mode='RGB',
+                layers=False,
+                brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"))
+
+            with gr.Row():
+                input_text_component = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your prompt",
+                    container=False,
+                )
+                submit_button_component = gr.Button(
+                    value='Submit', variant='primary', scale=0)
+
+            with gr.Accordion("Advanced Settings", open=False):
+                seed_slicer_component = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=42,
+                )
+
+                randomize_seed_checkbox_component = gr.Checkbox(
+                    label="Randomize seed", value=True)
+
+                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`.",
+                        minimum=0,
+                        maximum=1,
+                        step=0.01,
+                        value=0.85,
+                    )
+
+                    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",
+                        minimum=1,
+                        maximum=50,
+                        step=1,
+                        value=20,
+                    )
         with gr.Column():
-            output_image = gr.Image(label="Output Image")
-            output_mask = gr.Image(label="Processed Mask")
-            
-    run_button.click(
-        generate_inpainting,
-        inputs=[image_editor, prompt_box, seed_slider, random_seed_toggle, inpainting_strength_slider, steps_slider],
-        outputs=[output_image, output_mask]
-    )
-
-    gr.Examples(
-        examples=EXAMPLES,
-        fn=generate_inpainting,
-        inputs=[image_editor, prompt_box, seed_slider, random_seed_toggle, inpainting_strength_slider, steps_slider],
-        outputs=[output_image, output_mask],
-        run_on_click=True,
+            output_image_component = gr.Image(
+                type='pil', image_mode='RGB', label='Generated image', format="png")
+            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
+        )
+
+    submit_button_component.click(
+        fn=process,
+        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
+        ]
     )
 
-flux_app.launch(debug=False, show_error=True)
+demo.launch(debug=False, show_error=True)