Init

README.md

@@ -1,8 +1,8 @@
 ---
-title: FurnitureAdapter
-emoji: 🌖
-colorFrom: purple
-colorTo: indigo
+title: FurnitureDemo
+emoji: 🪑
+colorFrom: blue
+colorTo: white
 sdk: gradio
 python_version: 3.12
 sdk_version: 5.18.0
@@ -18,6 +18,6 @@ pinned: true
 license: mit
 ---
 
-# FurnitureBlendingDemoAPI
+# FurnitureDemo
 
 ...

app.py

@@ -4,7 +4,7 @@ import os
 import zipfile
 from io import BytesIO
 from pathlib import Path
-from typing import Literal, cast
+from typing import Literal, TypedDict, cast
 
 import gradio as gr
 import numpy as np
@@ -12,13 +12,23 @@ import requests
 from gradio.components.image_editor import EditorValue
 from PIL import Image
 
-PASSWORD = os.environ.get("PASSWORD", None)
-if not PASSWORD:
-    raise ValueError("PASSWORD is not set")
+_PASSWORD = os.environ.get("PASSWORD", None)
+if not _PASSWORD:
+    msg = "PASSWORD is not set"
+    raise ValueError(msg)
+PASSWORD = cast("str", _PASSWORD)
 
-ENDPOINT = os.environ.get("ENDPOINT", None)
-if not ENDPOINT:
-    raise ValueError("ENDPOINT is not set")
+_ENDPOINT = os.environ.get("ENDPOINT", None)
+if not _ENDPOINT:
+    msg = "ENDPOINT is not set"
+    raise ValueError(msg)
+ENDPOINT = cast("str", _ENDPOINT)
+
+# Add constants at the top
+THUMBNAIL_MAX_SIZE = 2048
+REFERENCE_MAX_SIZE = 1024
+REQUEST_TIMEOUT = 300  # 5 minutes
+DEFAULT_BRUSH_SIZE = 75
 
 
 def encode_image_as_base64(image: Image.Image) -> str:
@@ -60,99 +70,156 @@ def make_example(image_path: Path, mask_path: Path | None) -> EditorValue:
     }
 
 
+class InputFurnitureBlendingTypedDict(TypedDict):
+    return_type: Literal["zipfile", "s3"]
+    model_type: Literal["schnell", "dev"]
+    room_image_input: str
+    bbox: tuple[int, int, int, int]
+    furniture_reference_image: str
+    prompt: str
+    seed: int
+    num_inference_steps: int
+    max_dimension: int
+    margin: int
+    crop: bool
+    num_images_per_prompt: int
+    bucket: str
+
+
+# Add type hints for the response
+class GenerationResponse(TypedDict):
+    images: list[Image.Image]
+    error: str | None
+
+
+def validate_inputs(
+    image_and_mask: EditorValue | None,
+    furniture_reference: Image.Image | None,
+) -> tuple[Literal[True], None] | tuple[Literal[False], str]:
+    if not image_and_mask:
+        return False, "Please upload an image and draw a mask"
+
+    image_np = cast("np.ndarray", image_and_mask["background"])
+    if np.sum(image_np) == 0:
+        return False, "Please upload an image"
+
+    alpha_channel = cast("np.ndarray", image_and_mask["layers"][0])
+    mask_np = np.where(alpha_channel[:, :, 3] == 0, 0, 255).astype(np.uint8)
+    if np.sum(mask_np) == 0:
+        return False, "Please mark the areas you want to remove"
+
+    if not furniture_reference:
+        return False, "Please upload a furniture reference image"
+
+    return True, None
+
+
+def process_images(
+    image_and_mask: EditorValue,
+    furniture_reference: Image.Image,
+) -> tuple[Image.Image, Image.Image, Image.Image]:
+    image_np = cast("np.ndarray", image_and_mask["background"])
+    alpha_channel = cast("np.ndarray", image_and_mask["layers"][0])
+    mask_np = np.where(alpha_channel[:, :, 3] == 0, 0, 255).astype(np.uint8)
+
+    mask_image = Image.fromarray(mask_np).convert("L")
+    target_image = Image.fromarray(image_np).convert("RGB")
+
+    # Resize images
+    mask_image.thumbnail(
+        (THUMBNAIL_MAX_SIZE, THUMBNAIL_MAX_SIZE), Image.Resampling.LANCZOS
+    )
+    target_image.thumbnail(
+        (THUMBNAIL_MAX_SIZE, THUMBNAIL_MAX_SIZE), Image.Resampling.LANCZOS
+    )
+    furniture_reference.thumbnail(
+        (REFERENCE_MAX_SIZE, REFERENCE_MAX_SIZE), Image.Resampling.LANCZOS
+    )
+
+    return target_image, mask_image, furniture_reference
+
+
 def predict(
     model_type: Literal["schnell", "dev", "pixart"],
     image_and_mask: EditorValue,
     furniture_reference: Image.Image | None,
     prompt: str = "",
-    subfolder: str = "",
     seed: int = 0,
     num_inference_steps: int = 28,
     max_dimension: int = 512,
-    margin: int = 64,
+    margin: int = 128,
     crop: bool = True,
     num_images_per_prompt: int = 1,
 ) -> list[Image.Image] | None:
-    if not image_and_mask:
-        gr.Info("Please upload an image and draw a mask")
-        return None
-    if not furniture_reference:
-        gr.Info("Please upload a furniture reference image")
+    # Validate inputs
+    is_valid, error_message = validate_inputs(image_and_mask, furniture_reference)
+    if not is_valid and error_message:
+        gr.Info(error_message)
         return None
 
     if model_type == "pixart":
         gr.Info("PixArt is not supported yet")
         return None
 
-    image_np = image_and_mask["background"]
-    image_np = cast(np.ndarray, image_np)
-
-    # If the image is empty, return None
-    if np.sum(image_np) == 0:
-        gr.Info("Please upload an image")
-        return None
-
-    alpha_channel = image_and_mask["layers"][0]
-    alpha_channel = cast(np.ndarray, alpha_channel)
-    mask_np = np.where(alpha_channel[:, :, 3] == 0, 0, 255).astype(np.uint8)
+    # Process images
+    target_image, mask_image, furniture_reference = process_images(
+        image_and_mask, cast("Image.Image", furniture_reference)
+    )
 
-    # if mask_np is empty, return None
-    if np.sum(mask_np) == 0:
+    bbox = mask_image.getbbox()
+    if not bbox:
         gr.Info("Please mark the areas you want to remove")
         return None
 
-    mask_image = Image.fromarray(mask_np).convert("L")
-    target_image = Image.fromarray(image_np).convert("RGB")
-
-    # Avoid too big image to be sent to the API
-    mask_image.thumbnail((2048, 2048), Image.Resampling.LANCZOS)
-    target_image.thumbnail((2048, 2048), Image.Resampling.LANCZOS)
-    furniture_reference.thumbnail((1024, 1024), Image.Resampling.LANCZOS)
-
-    room_image_input_base64 = encode_image_as_base64(target_image)
-    room_image_mask_base64 = encode_image_as_base64(mask_image)
-    furniture_reference_base64 = encode_image_as_base64(furniture_reference)
-
-    room_image_input_base64 = "data:image/png;base64," + room_image_input_base64
-    room_image_mask_base64 = "data:image/png;base64," + room_image_mask_base64
-    furniture_reference_base64 = "data:image/png;base64," + furniture_reference_base64
-
-    response = requests.post(
-        ENDPOINT,
-        headers={"accept": "application/json", "Content-Type": "application/json"},
-        json={
-            "model_type": model_type,
-            "room_image_input": room_image_input_base64,
-            "room_image_mask": room_image_mask_base64,
-            "furniture_reference_image": furniture_reference_base64,
-            "prompt": prompt,
-            "subfolder": subfolder,
-            "seed": seed,
-            "num_inference_steps": num_inference_steps,
-            "max_dimension": max_dimension,
-            "condition_scale": 1.0,
-            "margin": margin,
-            "crop": crop,
-            "num_images_per_prompt": num_images_per_prompt,
-            "password": PASSWORD,
-        },
+    # Prepare API request
+    room_image_input_base64 = "data:image/png;base64," + encode_image_as_base64(
+        target_image
+    )
+    furniture_reference_base64 = "data:image/png;base64," + encode_image_as_base64(
+        furniture_reference
     )
-    if response.status_code != 200:
-        gr.Info("An error occurred during the generation")
-        return None
 
-    # Read the returned ZIP file from the response.
-    zip_bytes = io.BytesIO(response.content)
+    body = InputFurnitureBlendingTypedDict(
+        return_type="zipfile",
+        model_type=model_type,
+        room_image_input=room_image_input_base64,
+        bbox=bbox,
+        furniture_reference_image=furniture_reference_base64,
+        prompt=prompt,
+        seed=seed,
+        num_inference_steps=num_inference_steps,
+        max_dimension=max_dimension,
+        margin=margin,
+        crop=crop,
+        num_images_per_prompt=num_images_per_prompt,
+        bucket="furniture-blending",
+    )
 
-    final_image_list: list[Image.Image] = []
+    try:
+        response = requests.post(
+            ENDPOINT,
+            headers={"accept": "application/json", "Content-Type": "application/json"},
+            json=body,
+            timeout=REQUEST_TIMEOUT,
+        )
+        response.raise_for_status()
+    except requests.RequestException as e:
+        gr.Info(f"API request failed: {e!s}")
+        return None
 
-    # Open the ZIP archive.
-    with zipfile.ZipFile(zip_bytes, "r") as zip_file:
-        image_filenames = zip_file.namelist()
-        for filename in image_filenames:
-            with zip_file.open(filename) as file:
-                image = Image.open(file).convert("RGB")
-                final_image_list.append(image)
+    # Process response
+    try:
+        zip_bytes = io.BytesIO(response.content)
+        final_image_list: list[Image.Image] = []
+
+        with zipfile.ZipFile(zip_bytes, "r") as zip_file:
+            for filename in zip_file.namelist():
+                with zip_file.open(filename) as file:
+                    image = Image.open(file).convert("RGB")
+                    final_image_list.append(image)
+    except (OSError, zipfile.BadZipFile) as e:
+        gr.Info(f"Failed to process response: {e!s}")
+        return None
 
     return final_image_list
 
@@ -198,7 +265,7 @@ with gr.Blocks(css=css) as demo:
         </div>
     """)
 
-    with gr.Row() as content:
+    with gr.Row():
         with gr.Column(elem_id="col-left"):
             gr.HTML(
                 r"""
@@ -219,10 +286,14 @@ with gr.Blocks(css=css) as demo:
                 sources=["upload"],
                 show_download_button=False,
                 interactive=True,
-                brush=gr.Brush(default_size=75, colors=["#000000"], color_mode="fixed"),
+                brush=gr.Brush(
+                    default_size=DEFAULT_BRUSH_SIZE,
+                    colors=["#000000"],
+                    color_mode="fixed",
+                ),
                 transforms=[],
             )
-            image_and_mask_examples = gr.Examples(
+            gr.Examples(
                 examples=[
                     make_example(path, None)
                     for path in Path("./examples/scenes").glob("*.png")
@@ -248,7 +319,7 @@ with gr.Blocks(css=css) as demo:
                 sources=["upload"],
                 image_mode="RGB",
             )
-            furniture_examples = gr.Examples(
+            gr.Examples(
                 examples=list(Path("./examples/objects").glob("*.png")),
                 label="Furniture examples",
                 examples_per_page=6,
@@ -268,7 +339,7 @@ with gr.Blocks(css=css) as demo:
             results = gr.Gallery(
                 label="Result",
                 format="png",
-                file_types="image",
+                file_types=["image"],
                 show_label=False,
                 columns=2,
                 allow_preview=True,
@@ -286,10 +357,6 @@ with gr.Blocks(css=css) as demo:
                     label="Prompt",
                     value="",
                 )
-                subfolder = gr.Textbox(
-                    label="Subfolder",
-                    value="",
-                )
                 seed = gr.Slider(
                     label="Seed",
                     minimum=0,
@@ -339,14 +406,23 @@ with gr.Blocks(css=css) as demo:
         outputs=num_inference_steps,
     )
 
+    # Add loading indicator
+    with gr.Row():
+        loading_indicator = gr.HTML(
+            '<div id="loading" style="display:none;">Processing... Please wait.</div>'
+        )
+
+    # Update click handler to show loading state
     run_button.click(
+        fn=lambda: gr.update(visible=True),
+        outputs=[loading_indicator],
+    ).then(
         fn=predict,
         inputs=[
             model_type,
             image_and_mask,
             condition_image,
             prompt,
-            subfolder,
             seed,
             num_inference_steps,
             max_dimension,
@@ -355,7 +431,10 @@ with gr.Blocks(css=css) as demo:
             num_images_per_prompt,
         ],
         outputs=[results],
+    ).then(
+        fn=lambda: gr.update(visible=False),
+        outputs=[loading_indicator],
     )
 
-
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()