Update

README.md

@@ -1,13 +1,30 @@
 ---
-title: FurnitureBlendingDemo
-emoji: 🐢
-colorFrom: purple
-colorTo: pink
+title: FurnitureInpaintingDemo
+emoji: 🌖
+colorFrom: pink
+colorTo: red
 sdk: gradio
-sdk_version: 5.15.0
+python_version: 3.12
+sdk_version: 5.12.0
+suggested_hardware: a100-large
 app_file: app.py
-pinned: false
-short_description: Furniture Blending Demo
+# fullWidth: true
+# header: mini
+# models: blanchon/VirtualUnstagingModel
+# datasets: blanchon/VirtualUnstagingDataset
+tags:
+  - image-generation
+  - image-to-image
+  - furniture
+  - virtual-staging
+  - home-decor
+  - home-design
+pinned: true
+# preload_from_hub:
+#  - blanchon/VirtualUnstagingModel
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# VirtualUnstaging
+
+...

app.py

@@ -0,0 +1,285 @@
+import os
+import numpy as np
+from typing import cast
+import torch
+from PIL import Image, ImageDraw
+from diffusers import DiffusionPipeline
+import gradio as gr
+from gradio.components.image_editor import EditorValue
+import spaces
+
+DEVICE = "cuda"
+
+MAIN_MODEL_REPO_ID = os.getenv("MAIN_MODEL_REPO_ID", None)
+SUB_MODEL_REPO_ID = os.getenv("SUB_MODEL_REPO_ID", None)
+SUB_MODEL_SUBFOLDER = os.getenv("SUB_MODEL_SUBFOLDER", None)
+
+if MAIN_MODEL_REPO_ID is None:
+    raise ValueError("MAIN_MODEL_REPO_ID is not set")
+if SUB_MODEL_REPO_ID is None:
+    raise ValueError("SUB_MODEL_REPO_ID is not set")
+if SUB_MODEL_SUBFOLDER is None:
+    raise ValueError("SUB_MODEL_SUBFOLDER is not set")
+
+pipeline = DiffusionPipeline.from_pretrained(
+    MAIN_MODEL_REPO_ID,
+    torch_dtype=torch.bfloat16,
+    custom_pipeline=SUB_MODEL_REPO_ID,
+).to(DEVICE)
+pipeline.post_init()
+
+
+def crop_divisible_by_16(image: Image.Image) -> Image.Image:
+    w, h = image.size
+    w = w - w % 16
+    h = h - h % 16
+    return image.crop((0, 0, w, h))
+
+
+@spaces.GPU(duration=150)
+def predict(
+    image_and_mask: EditorValue,
+    furniture_reference: Image.Image | None,
+    seed: int = 0,
+    num_inference_steps: int = 28,
+    condition_size: int = 512,
+    target_size: int = 512,
+    condition_scale: float = 1.0,
+    progress: gr.Progress = gr.Progress(track_tqdm=True),  # noqa: ARG001, B008
+) -> Image.Image | None:
+    # ) -> tuple[Image.Image, 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
+
+    pipeline.load(
+        SUB_MODEL_REPO_ID,
+        subfolder=SUB_MODEL_SUBFOLDER,
+    )
+
+    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
+
+    pipeline.load(
+        SUB_MODEL_REPO_ID,
+        subfolder=SUB_MODEL_SUBFOLDER,
+    )
+
+    mask_image = Image.fromarray(mask_np).convert("L")
+    mask_image.thumbnail((target_size, target_size))
+    mask_image_bbox = mask_image.getbbox()
+
+    target_image = Image.fromarray(image_np).convert("RGB")
+    target_image.thumbnail((target_size, target_size))
+    target_image_size = target_image.size
+
+    condition_image = Image.new("RGB", (target_size, target_size), (255, 255, 255))
+    condition_image.paste(target_image, (0, 0))
+
+    # Fill all the bbox area with 255
+    draw = ImageDraw.Draw(condition_image)
+    draw.rectangle(mask_image_bbox, fill="white", outline="white", width=10)
+
+    # # Resize the furniture_reference so it fit in the bbox and overlay it on the condition_image
+    bbox_width = mask_image_bbox[2] - mask_image_bbox[0]
+    bbox_height = mask_image_bbox[3] - mask_image_bbox[1]
+    furniture_reference.thumbnail((bbox_width, bbox_height))
+    # Paste at the bbox center
+    condition_image.paste(
+        furniture_reference,
+        (
+            mask_image_bbox[0] + (bbox_width - furniture_reference.width) // 2,
+            mask_image_bbox[1] + (bbox_height - furniture_reference.height) // 2,
+        ),
+    )
+
+    # Save condition image
+    condition_image.save("condition_image.png")
+
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+
+    final_image = pipeline(
+        condition_image=condition_image,
+        prompt="",
+        num_inference_steps=num_inference_steps,
+        height=target_size,
+        width=target_size,
+        union_cond_attn=True,
+        add_cond_attn=False,
+        latent_lora=False,
+        default_lora=False,
+        condition_scale=condition_scale,
+        generator=generator,
+        max_sequence_length=512,
+    ).images[0]
+
+    final_image_crop = final_image.crop((
+        0,
+        0,
+        target_image_size[0],
+        target_image_size[1],
+    ))
+
+    return final_image_crop
+
+
+intro_markdown = r"""
+# Furniture Inpainting 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(
+                """
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
+                <div>
+                Step 1. Upload a room image ⬇️
+                </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(
+                """
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
+                <div>
+                Step 2. Upload a 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(
+                """
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
+                <div>
+                Step 3. Press Run to launch the generation
+                </div>
+            </div>
+            """,
+                max_height=50,
+            )
+            # image_slider = ImageSlider(
+            #     label="Result",
+            #     interactive=False,
+            # )
+            result = gr.Image(label="Result")
+            run_button = gr.Button("Run")
+
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=100_000,
+                    step=1,
+                    value=0,
+                )
+                condition_scale = gr.Slider(
+                    label="Condition Scale",
+                    minimum=-10.0,
+                    maximum=10.0,
+                    step=0.10,
+                    value=1.0,
+                )
+                with gr.Column():
+                    condition_size = gr.Slider(
+                        label="Condition Size",
+                        minimum=256,
+                        maximum=1024,
+                        step=128,
+                        value=512,
+                    )
+
+                    target_size = gr.Slider(
+                        label="Target Size",
+                        minimum=256,
+                        maximum=1024,
+                        step=128,
+                        value=512,
+                    )
+
+                    num_inference_steps = gr.Slider(
+                        label="Number of inference steps",
+                        minimum=1,
+                        maximum=50,
+                        step=1,
+                        value=28,
+                    )
+
+    run_button.click(
+        fn=predict,
+        inputs=[
+            image_and_mask,
+            condition_image,
+            seed,
+            num_inference_steps,
+            condition_size,
+            target_size,
+            condition_scale,
+        ],
+        # outputs=[image_slider],
+        outputs=[result],
+    )
+
+
+demo.launch()

pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "VirtualStaging"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "accelerate>=1.2.1",
+    "diffusers==0.31.0",
+    "gradio>=5.12.0",
+    "gradio-imageslider>=0.0.20",
+    "peft>=0.14.0",
+    "pillow>=11.1.0",
+    "safetensors>=0.5.2",
+    "sentencepiece>=0.2.0",
+    "spaces>=0.32.0",
+    "transformers>=4.48.0",
+]

requirements.txt

@@ -0,0 +1,10 @@
+diffusers
+transformers
+accelerate
+safetensors
+sentencepiece
+peft
+gradio
+spaces
+pillow
+gradio_imageslider