Update

app.py

@@ -0,0 +1,294 @@
+import base64
+import io
+import os
+import zipfile
+from io import BytesIO
+from typing import Literal, cast
+
+import gradio as gr
+import numpy as np
+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")
+
+ENDPOINT = os.environ.get("ENDPOINT", None)
+if not ENDPOINT:
+    raise ValueError("ENDPOINT is not set")
+
+
+def encode_image_as_base64(image: Image.Image) -> str:
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
+def predict(
+    model_type: Literal["schnell", "dev"],
+    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,
+    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")
+        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)
+
+    # if mask_np is empty, return None
+    if np.sum(mask_np) == 0:
+        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,
+        },
+    )
+    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)
+
+    final_image_list: list[Image.Image] = []
+
+    # 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)
+
+    return final_image_list
+
+
+intro_markdown = r"""
+# Furniture Blending Demo
+"""
+
+css = r"""
+#col-left {
+    margin: 0 auto;
+    max-width: 430px;
+}
+#col-mid {
+    margin: 0 auto;
+    max-width: 430px;
+}
+#col-right {
+    margin: 0 auto;
+    max-width: 430px;
+}
+#col-showcase {
+    margin: 0 auto;
+    max-width: 1100px;
+}
+"""
+
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(intro_markdown)
+
+    with gr.Row() as content:
+        with gr.Column(elem_id="col-left"):
+            gr.HTML(
+                r"""
+                <div style="display: flex; justify-content: start; align-items: center; text-align: center; font-size: 20px; height: 50px;">
+                    <div>
+                    🪟 Room image with inpainting mask ⬇️
+                    </div>
+                </div>
+                """,
+                max_height=50,
+            )
+            image_and_mask = gr.ImageMask(
+                label="Image and Mask",
+                layers=False,
+                height="full",
+                width="full",
+                show_fullscreen_button=False,
+                sources=["upload"],
+                show_download_button=False,
+                interactive=True,
+                brush=gr.Brush(default_size=75, colors=["#000000"], color_mode="fixed"),
+                transforms=[],
+            )
+        with gr.Column(elem_id="col-mid"):
+            gr.HTML(
+                r"""
+                <div style="display: flex; justify-content: start; align-items: center; text-align: center; font-size: 20px; height: 50px;">
+                    <div>
+                    🪑 Furniture reference image ⬇️
+                    </div>
+                </div>
+                """,
+                max_height=50,
+            )
+            condition_image = gr.Image(
+                label="Furniture Reference",
+                type="pil",
+                sources=["upload"],
+                image_mode="RGB",
+            )
+        with gr.Column(elem_id="col-right"):
+            gr.HTML(
+                r"""
+                <div style="display: flex; justify-content: start; align-items: center; text-align: center; font-size: 20px; height: 50px;">
+                    <div>
+                    🔥 Press Run ⬇️
+                    </div>
+                </div>
+                """,
+                max_height=50,
+            )
+            results = gr.Gallery(
+                label="Result",
+                format="png",
+                file_types="image",
+                show_label=False,
+                columns=2,
+                allow_preview=True,
+                preview=True,
+            )
+            model_type = gr.Radio(
+                choices=["schnell", "dev"],
+                value="schnell",
+                label="Model Type",
+            )
+            run_button = gr.Button("Run")
+
+            with gr.Accordion("Advanced Settings", open=False):
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    value="",
+                )
+                subfolder = gr.Textbox(
+                    label="Subfolder",
+                    value="",
+                )
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=np.iinfo(np.int32).max,
+                    step=1,
+                    value=0,
+                )
+                num_images_per_prompt = gr.Slider(
+                    label="Number of images per prompt",
+                    minimum=1,
+                    maximum=10,
+                    step=1,
+                    value=4,
+                )
+                crop = gr.Checkbox(
+                    label="Crop",
+                    value=False,
+                )
+                margin = gr.Slider(
+                    label="Margin",
+                    minimum=0,
+                    maximum=256,
+                    step=16,
+                    value=128,
+                )
+                with gr.Column():
+                    max_dimension = gr.Slider(
+                        label="Max Dimension",
+                        minimum=256,
+                        maximum=1024,
+                        step=128,
+                        value=512,
+                    )
+
+                    num_inference_steps = gr.Slider(
+                        label="Number of inference steps",
+                        minimum=4,
+                        maximum=30,
+                        step=2,
+                        value=4,
+                    )
+
+    # Change the number of inference steps based on the model type
+    model_type.change(
+        fn=lambda x: gr.update(value=4 if x == "schnell" else 28),
+        inputs=model_type,
+        outputs=num_inference_steps,
+    )
+
+    run_button.click(
+        fn=predict,
+        inputs=[
+            model_type,
+            image_and_mask,
+            condition_image,
+            prompt,
+            subfolder,
+            seed,
+            num_inference_steps,
+            max_dimension,
+            margin,
+            crop,
+            num_images_per_prompt,
+        ],
+        outputs=[results],
+    )
+
+
+demo.launch()

pyproject.toml

@@ -0,0 +1,13 @@
+[project]
+name = "FurnitureBlendingDemoAPI"
+version = "0.1.0"
+description = "Furniture blending demo API"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "gradio",
+    "gradio-imageslider",
+    "pillow",
+    "requests",
+    "numpy",
+]

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+pillow
+gradio_imageslider
+requests
+numpy