Edit all

app.py

@@ -1,29 +1,37 @@
 import spaces
+
 from typing import Tuple, Union, List
 import os
 import gc
 
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageFilter
 
 import torch
-from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline
+
+from diffusers.pipelines.controlnet import (
+    StableDiffusionControlNetInpaintPipeline
+)
+
 from diffusers import (
     ControlNetModel,
     UniPCMultistepScheduler,
 )
+
 from transformers import (
     AutoImageProcessor,
     UperNetForSemanticSegmentation,
     AutoModelForDepthEstimation,
 )
+
 from colors import ade_palette
 from utils import map_colors_rgb
+
 import gradio as gr
 
 
 # ============================================================
-# DEVICE
+# CONFIG
 # ============================================================
 
 device = "cuda"
@@ -31,7 +39,7 @@ dtype = torch.float16
 
 
 # ============================================================
-# CSS
+# UI CSS
 # ============================================================
 
 css = """
@@ -50,47 +58,182 @@ css = """
 
 
 # ============================================================
-# SEGMENTATION
+# UTILS
 # ============================================================
 
-def filter_items(
-    colors_list: Union[List, np.ndarray],
-    items_list: Union[List, np.ndarray],
-    items_to_remove: Union[List, np.ndarray]
-) -> Tuple[Union[List, np.ndarray], Union[List, np.ndarray]]:
+def flush():
+    """
+    Release unused GPU memory.
+    """
+
+    gc.collect()
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def resize_dimensions(
+    dimensions,
+    target_size
+):
+    """
+    Resize while preserving the original aspect ratio.
+    """
+
+    width, height = dimensions
+
+    if width <= target_size and height <= target_size:
+        return dimensions
+
+    if width > height:
+
+        aspect_ratio = height / width
+
+        return (
+            target_size,
+            max(64, int(target_size * aspect_ratio))
+        )
+
+    aspect_ratio = width / height
+
+    return (
+        max(64, int(target_size * aspect_ratio)),
+        target_size
+    )
+
+
+def resize_keep_aspect(
+    image: Image.Image,
+    target_width: int,
+    target_height: int,
+    background=(255, 255, 255)
+):
+    """
+    Resize an image preserving its aspect ratio and place it
+    in a canvas of the requested size.
+
+    This is intentionally NOT a stretch resize.
+    """
+
+    image = image.convert("RGB")
+
+    src_w, src_h = image.size
+
+    scale = min(
+        target_width / src_w,
+        target_height / src_h
+    )
+
+    new_w = max(1, int(src_w * scale))
+    new_h = max(1, int(src_h * scale))
+
+    resized = image.resize(
+        (new_w, new_h),
+        Image.Resampling.LANCZOS
+    )
+
+    canvas = Image.new(
+        "RGB",
+        (target_width, target_height),
+        background
+    )
+
+    x = (target_width - new_w) // 2
+    y = (target_height - new_h) // 2
+
+    canvas.paste(
+        resized,
+        (x, y)
+    )
+
+    return canvas
+
+
+def create_binary_mask_from_segmentation(
+    segmentation: Image.Image,
+    wanted_items: List[str]
+):
+    """
+    Create a binary mask from an ADE20K semantic segmentation.
+
+    Pixels belonging to the requested classes become white.
+    """
+
+    seg_rgb = np.array(segmentation)
+
+    unique_colors = np.unique(
+        seg_rgb.reshape(
+            -1,
+            seg_rgb.shape[2]
+        ),
+        axis=0
+    )
+
+    mask = np.zeros(
+        seg_rgb.shape[:2],
+        dtype=np.uint8
+    )
 
-    filtered_colors = []
-    filtered_items = []
+    for color in unique_colors:
 
-    for color, item in zip(colors_list, items_list):
-        if item not in items_to_remove:
-            filtered_colors.append(color)
-            filtered_items.append(item)
+        color_tuple = tuple(
+            int(v)
+            for v in color
+        )
 
-    return filtered_colors, filtered_items
+        item_name = map_colors_rgb(
+            color_tuple
+        )
 
+        if item_name in wanted_items:
 
-def get_segmentation_pipeline(
-) -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
+            matches = (
+                seg_rgb == color
+            ).all(axis=2)
+
+            mask[matches] = 255
+
+    return Image.fromarray(
+        mask,
+        mode="L"
+    )
+
+
+# ============================================================
+# SEGMENTATION
+# ============================================================
+
+def get_segmentation_pipeline():
+    """
+    Load ADE20K semantic segmentation.
+    """
 
     image_processor = AutoImageProcessor.from_pretrained(
         "openmmlab/upernet-convnext-small"
     )
 
-    image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
-        "openmmlab/upernet-convnext-small"
+    image_segmentor = (
+        UperNetForSemanticSegmentation.from_pretrained(
+            "openmmlab/upernet-convnext-small"
+        )
     )
 
-    return image_processor, image_segmentor
+    return (
+        image_processor,
+        image_segmentor
+    )
 
 
 @torch.inference_mode()
 @spaces.GPU
 def segment_image(
-    image: Image,
+    image: Image.Image,
     image_processor: AutoImageProcessor,
     image_segmentor: UperNetForSemanticSegmentation
-) -> Image:
+) -> Image.Image:
+    """
+    Generate an RGB ADE20K segmentation image.
+    """
 
     pixel_values = image_processor(
         image,
@@ -98,25 +241,50 @@ def segment_image(
     ).pixel_values.to(device)
 
     with torch.no_grad():
-        outputs = image_segmentor(pixel_values)
 
-    seg = image_processor.post_process_semantic_segmentation(
-        outputs, target_sizes=[image.size[::-1]]
-    )[0]
-    
-    # Move segmentation tensor from GPU to CPU before using it with NumPy
-    seg = seg.detach().cpu().numpy()
-    
-    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
+        outputs = image_segmentor(
+            pixel_values
+        )
+
+    seg = (
+        image_processor
+        .post_process_semantic_segmentation(
+            outputs,
+            target_sizes=[image.size[::-1]]
+        )[0]
+    )
+
+    # IMPORTANT:
+    # Move tensor from CUDA to CPU before NumPy.
+    seg = (
+        seg.detach()
+        .cpu()
+        .numpy()
+    )
+
+    color_seg = np.zeros(
+        (
+            seg.shape[0],
+            seg.shape[1],
+            3
+        ),
+        dtype=np.uint8
+    )
 
-    palette = np.array(ade_palette())
+    palette = np.array(
+        ade_palette()
+    )
 
     for label, color in enumerate(palette):
-        color_seg[seg == label, :] = color
 
-    color_seg = color_seg.astype(np.uint8)
+        color_seg[
+            seg == label,
+            :
+        ] = color
 
-    return Image.fromarray(color_seg).convert("RGB")
+    return Image.fromarray(
+        color_seg
+    ).convert("RGB")
 
 
 # ============================================================
@@ -135,16 +303,22 @@ def get_depth_pipeline():
         torch_dtype=dtype
     )
 
-    return feature_extractor, depth_estimator
+    return (
+        feature_extractor,
+        depth_estimator
+    )
 
 
 @torch.inference_mode()
 @spaces.GPU
 def get_depth_image(
-    image: Image,
+    image: Image.Image,
     feature_extractor: AutoImageProcessor,
     depth_estimator: AutoModelForDepthEstimation
-) -> Image:
+) -> Image.Image:
+    """
+    Generate normalized depth image.
+    """
 
     image_to_depth = feature_extractor(
         images=image,
@@ -152,9 +326,13 @@ def get_depth_image(
     ).to(device)
 
     with torch.no_grad():
-        depth_map = depth_estimator(
-            **image_to_depth
-        ).predicted_depth
+
+        depth_map = (
+            depth_estimator(
+                **image_to_depth
+            )
+            .predicted_depth
+        )
 
     width, height = image.size
 
@@ -188,57 +366,190 @@ def get_depth_image(
         dim=1
     )
 
-    image = image.permute(
-        0, 2, 3, 1
-    ).cpu().numpy()[0]
+    image = (
+        image
+        .permute(0, 2, 3, 1)
+        .cpu()
+        .numpy()[0]
+    )
 
-    image = Image.fromarray(
+    return Image.fromarray(
         (image * 255.0)
         .clip(0, 255)
         .astype(np.uint8)
     )
 
-    return image
-
 
 # ============================================================
-# RESIZE
+# REFERENCE ROOM -> FURNITURE EXTRACTION
 # ============================================================
 
-def resize_dimensions(dimensions, target_size):
+REFERENCE_FURNITURE_CLASSES = {
+    # Main furniture
+    "chair",
+    "sofa",
+    "table",
+    "bed",
+    "cabinet",
+    "shelf",
+    "desk",
+    "dresser",
+    "bench",
+    "armchair",
+
+    # Small furniture / objects
+    "ottoman",
+    "stool",
+    "pillow",
+    "lamp",
+    "television",
+    "computer",
+    "monitor",
+
+    # Decoration
+    "plant",
+    "vase",
+    "painting",
+    "mirror",
+    "clock",
+
+    # Floor / textile elements
+    "rug",
+
+    # Additional ADE-like categories
+    "book",
+    "books",
+    "chest of drawers",
+    "counter",
+}
 
-    width, height = dimensions
 
-    if width < target_size and height < target_size:
-        return dimensions
+def extract_furniture_reference(
+    reference_image: Image.Image,
+    segmentation_image: Image.Image
+) -> Image.Image:
+    """
+    Remove the architectural/background part of the reference room.
 
-    if width > height:
-        aspect_ratio = height / width
+    The goal is NOT to extract a single object.
 
-        return (
-            target_size,
-            int(target_size * aspect_ratio)
-        )
+    Instead we preserve the complete set of recognizable
+    furniture/decor elements from the reference room.
+    """
+
+    mask = create_binary_mask_from_segmentation(
+        segmentation_image,
+        list(REFERENCE_FURNITURE_CLASSES)
+    )
 
-    else:
-        aspect_ratio = width / height
+    # Slightly enlarge furniture regions so that thin objects
+    # and object edges are not destroyed by semantic segmentation.
+    mask = mask.filter(
+        ImageFilter.MaxFilter(9)
+    )
 
-        return (
-            int(target_size * aspect_ratio),
-            target_size
-        )
+    reference_rgb = reference_image.convert(
+        "RGB"
+    )
+
+    # Neutral background.
+    #
+    # A neutral light gray is preferable to black because black
+    # can be interpreted as an actual dark object by the
+    # image encoder.
+    background = Image.new(
+        "RGB",
+        reference_rgb.size,
+        (235, 235, 235)
+    )
+
+    furniture_reference = Image.composite(
+        reference_rgb,
+        background,
+        mask
+    )
+
+    return furniture_reference
 
 
 # ============================================================
-# MEMORY CLEANUP
+# TARGET ROOM MASK
 # ============================================================
 
-def flush():
+ARCHITECTURE_CLASSES = {
+    "wall",
+    "floor",
+    "ceiling",
+    "window",
+    "windowpane",
+    "door",
+    "door;double;door",
+    "curtain",
+}
 
-    gc.collect()
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+def create_room_edit_mask(
+    segmentation_image: Image.Image
+) -> Image.Image:
+    """
+    Create the region where furniture can be generated.
+
+    Architecture is kept strongly constrained by ControlNet.
+
+    The mask deliberately does not depend on a manually selected
+    rectangle or furniture placement area.
+    """
+
+    seg_rgb = np.array(
+        segmentation_image
+    )
+
+    unique_colors = np.unique(
+        seg_rgb.reshape(
+            -1,
+            seg_rgb.shape[2]
+        ),
+        axis=0
+    )
+
+    # Start with everything editable.
+    editable = np.ones(
+        seg_rgb.shape[:2],
+        dtype=np.uint8
+    )
+
+    for color in unique_colors:
+
+        color_tuple = tuple(
+            int(v)
+            for v in color
+        )
+
+        item_name = map_colors_rgb(
+            color_tuple
+        )
+
+        if item_name in ARCHITECTURE_CLASSES:
+
+            matches = (
+                seg_rgb == color
+            ).all(axis=2)
+
+            # Windows, doors and architectural areas are
+            # protected from direct inpainting.
+            editable[matches] = 0
+
+    mask = Image.fromarray(
+        editable * 255,
+        mode="L"
+    )
+
+    # Keep furniture transitions soft.
+    mask = mask.filter(
+        ImageFilter.GaussianBlur(2)
+    )
+
+    return mask
 
 
 # ============================================================
@@ -252,38 +563,43 @@ class ControlNetDepthDesignModelMulti:
         self.seed = 323 * 111
 
         self.neg_prompt = (
-            "window, door, low resolution, banner, logo, watermark, "
-            "text, deformed, blurry, out of focus, surreal, ugly, "
-            "beginner, duplicated furniture, extra furniture"
+            "deformed furniture, distorted furniture, "
+            "warped furniture, melted furniture, "
+            "wrong furniture geometry, duplicated furniture, "
+            "extra furniture, floating furniture, "
+            "impossible perspective, impossible proportions, "
+            "bad architecture, distorted walls, distorted floor, "
+            "changed windows, changed doors, "
+            "window, door, text, logo, watermark, "
+            "low resolution, blurry, out of focus, "
+            "surreal, ugly, cartoon, illustration, "
+            "CGI, plastic, oversaturated"
         )
 
-        self.control_items = [
-            "windowpane;window",
-            "door;double;door"
-        ]
-
         self.additional_quality_suffix = (
-            "interior design, 4K, high resolution, photorealistic"
+            "professional interior photography, "
+            "photorealistic interior, realistic materials, "
+            "natural lighting, realistic shadows, "
+            "correct perspective, correct scale, "
+            "high detail, 4K"
         )
 
-    # ========================================================
-    # GENERATE
-    # ========================================================
-
     @spaces.GPU
     def generate_design(
         self,
-        empty_room_image: Image,
-        furniture_image: Image,
+        empty_room_image: Image.Image,
+        furniture_image: Image.Image,
         prompt: str,
-        guidance_scale: int = 10,
-        num_steps: int = 30,
-        strength: float = 0.85,
+        guidance_scale: float = 7.5,
+        num_steps: int = 35,
+        strength: float = 0.62,
         img_size: int = 640,
-        ip_scale: float = 0.75
-    ) -> Image:
+        ip_scale: float = 0.90
+    ) -> Image.Image:
 
-        print("Starting furniture insertion...")
+        print(
+            "Starting complete room furnishing transfer..."
+        )
 
         flush()
 
@@ -293,19 +609,23 @@ class ControlNetDepthDesignModelMulti:
             int(self.seed)
         )
 
-        # ----------------------------------------------------
-        # PROMPT
-        # ----------------------------------------------------
-
         if prompt is None:
             prompt = ""
 
         prompt = prompt.strip()
 
         if prompt == "":
+
             prompt = (
-                "Place the furniture from the reference image "
-                "naturally inside the room."
+                "Transfer the complete furniture and decor "
+                "ensemble from the reference room into the target room. "
+                "Preserve the target room architecture, walls, floor, "
+                "windows, doors and camera viewpoint. "
+                "Recreate the same sofa, chairs, tables, cabinets, "
+                "lamps, rugs, plants, artwork and decorative objects "
+                "from the reference room. "
+                "Adapt their size, position, perspective, lighting "
+                "and shadows naturally to the target room."
             )
 
         pos_prompt = (
@@ -314,178 +634,169 @@ class ControlNetDepthDesignModelMulti:
             + self.additional_quality_suffix
         )
 
-        # ----------------------------------------------------
-        # ORIGINAL ROOM SIZE
-        # ----------------------------------------------------
+        # --------------------------------------------------------
+        # TARGET ROOM
+        # --------------------------------------------------------
 
-        orig_w, orig_h = empty_room_image.size
-
-        # ----------------------------------------------------
-        # ROOM
-        # ----------------------------------------------------
+        orig_w, orig_h = (
+            empty_room_image.size
+        )
 
         new_width, new_height = resize_dimensions(
             empty_room_image.size,
             img_size
         )
 
-        input_image = empty_room_image.convert(
-            "RGB"
-        ).resize(
-            (new_width, new_height)
+        input_image = (
+            empty_room_image
+            .convert("RGB")
+            .resize(
+                (new_width, new_height),
+                Image.Resampling.LANCZOS
+            )
         )
 
-        # ----------------------------------------------------
-        # FURNITURE REFERENCE
-        # ----------------------------------------------------
-
-        furniture_image = furniture_image.convert("RGB")
+        # --------------------------------------------------------
+        # REFERENCE ROOM
+        # --------------------------------------------------------
 
-        ip_width = max(
-            64,
-            int(new_width / 8) * 8
+        reference_image = (
+            furniture_image
+            .convert("RGB")
         )
 
-        ip_height = max(
-            64,
-            int(new_height / 8) * 8
+        # Resize reference WITHOUT stretching it.
+        reference_for_segmentation = resize_keep_aspect(
+            reference_image,
+            new_width,
+            new_height,
+            background=(235, 235, 235)
         )
 
-        furniture_ip = furniture_image.resize(
-            (ip_width, ip_height),
-            Image.Resampling.LANCZOS
+        # --------------------------------------------------------
+        # TARGET SEGMENTATION
+        # --------------------------------------------------------
+
+        target_segmentation = segment_image(
+            input_image,
+            seg_image_processor,
+            image_segmentor
         )
 
-        # ----------------------------------------------------
-        # SEGMENTATION
-        # ----------------------------------------------------
+        # --------------------------------------------------------
+        # REFERENCE SEGMENTATION
+        # --------------------------------------------------------
 
-        real_seg = np.array(
-            segment_image(
-                input_image,
-                seg_image_processor,
-                image_segmentor
-            )
+        print(
+            "Segmenting reference room furniture..."
         )
 
-        unique_colors = np.unique(
-            real_seg.reshape(
-                -1,
-                real_seg.shape[2]
-            ),
-            axis=0
+        reference_segmentation = segment_image(
+            reference_for_segmentation,
+            seg_image_processor,
+            image_segmentor
         )
 
-        unique_colors = [
-            tuple(color)
-            for color in unique_colors
-        ]
-
-        segment_items = [
-            map_colors_rgb(i)
-            for i in unique_colors
-        ]
+        # --------------------------------------------------------
+        # EXTRACT COMPLETE FURNITURE ENSEMBLE
+        # --------------------------------------------------------
 
-        chosen_colors, segment_items = filter_items(
-            colors_list=unique_colors,
-            items_list=segment_items,
-            items_to_remove=self.control_items
+        furniture_only = extract_furniture_reference(
+            reference_for_segmentation,
+            reference_segmentation
         )
 
-        mask = np.zeros_like(real_seg)
-
-        for color in chosen_colors:
-
-            color_matches = (
-                real_seg == color
-            ).all(axis=2)
-
-            mask[color_matches] = 1
-
-        image_np = np.array(input_image)
+        # --------------------------------------------------------
+        # SAVE/PRINT INFORMATION FOR DEBUGGING
+        # --------------------------------------------------------
 
-        image = Image.fromarray(
-            image_np
-        ).convert("RGB")
+        print(
+            "Furniture reference extracted successfully."
+        )
 
-        mask_image = Image.fromarray(
-            (mask * 255).astype(np.uint8)
-        ).convert("RGB")
+        # --------------------------------------------------------
+        # TARGET EDIT MASK
+        # --------------------------------------------------------
 
-        segmentation_cond_image = Image.fromarray(
-            real_seg
-        ).convert("RGB")
+        mask_image = create_room_edit_mask(
+            target_segmentation
+        )
 
-        # ----------------------------------------------------
+        # --------------------------------------------------------
         # DEPTH
-        # ----------------------------------------------------
+        # --------------------------------------------------------
 
         image_depth = get_depth_image(
-            image,
+            input_image,
             depth_feature_extractor,
             depth_estimator
         )
 
         flush()
 
-        # ----------------------------------------------------
+        # --------------------------------------------------------
         # IP-ADAPTER
-        #
-        # IMPORTANT:
-        # Unlike the original StableDesign code,
-        # we use the ACTUAL furniture photo here.
-        # ----------------------------------------------------
+        # --------------------------------------------------------
 
         pipe.set_ip_adapter_scale(
             float(ip_scale)
         )
 
-        # ----------------------------------------------------
-        # GENERATION
-        # ----------------------------------------------------
+        # --------------------------------------------------------
+        # CONTROLNET
+        # --------------------------------------------------------
 
         generated_image = pipe(
             prompt=pos_prompt,
+
             negative_prompt=self.neg_prompt,
 
-            num_inference_steps=num_steps,
+            num_inference_steps=int(
+                num_steps
+            ),
 
-            strength=strength,
+            strength=float(
+                strength
+            ),
 
-            guidance_scale=guidance_scale,
+            guidance_scale=float(
+                guidance_scale
+            ),
 
             generator=self.generator,
 
-            image=image,
+            image=input_image,
 
             mask_image=mask_image,
 
-            ip_adapter_image=furniture_ip,
+            # IMPORTANT:
+            # This is now NOT the entire reference room.
+            #
+            # It is a furniture/decor-focused version
+            # of the reference room.
+            ip_adapter_image=furniture_only,
 
             control_image=[
                 image_depth,
-                segmentation_cond_image
+                target_segmentation
             ],
 
             controlnet_conditioning_scale=[
-                0.5,
-                0.5
+                0.70,
+                0.65
             ]
-
         ).images[0]
 
         flush()
 
-        # ----------------------------------------------------
-        # RESTORE ORIGINAL SIZE
-        # ----------------------------------------------------
-
         design_image = generated_image.resize(
             (orig_w, orig_h),
             Image.Resampling.LANCZOS
         )
 
-        print("Generation completed.")
+        print(
+            "Furniture transfer completed."
+        )
 
         return design_image
 
@@ -497,30 +808,35 @@ class ControlNetDepthDesignModelMulti:
 def create_demo(model):
 
     gr.Markdown(
-        "### StableDesign Furniture"
+        "### StableDesign — Reference Room Furniture Transfer"
     )
 
     gr.Markdown(
-        "Upload a room photo and a furniture/product photo. "
-        "The AI will place the furniture into the room."
+        """
+        Upload two room photographs:
+
+        **1. Target room:** the room whose architecture,
+        camera perspective, walls, floor, windows and doors
+        must be preserved.
+
+        **2. Reference room:** a completely furnished room.
+        The AI extracts the furniture and decorative ensemble
+        from this image and recreates it inside the target room.
+        """
     )
 
     with gr.Row():
 
-        # ----------------------------------------------------
-        # LEFT
-        # ----------------------------------------------------
-
         with gr.Column():
 
             room_image = gr.Image(
-                label="1. Room photo",
+                label="1. Target room",
                 type="pil",
                 elem_id="img-display-input"
             )
 
             furniture_image = gr.Image(
-                label="2. Furniture / product photo",
+                label="2. Furnished reference room",
                 type="pil",
                 elem_id="img-display-input"
             )
@@ -528,12 +844,16 @@ def create_demo(model):
             prompt = gr.Textbox(
                 label="Instructions",
                 value=(
-                    "Place the furniture from the second image "
-                    "naturally in the room. Preserve the furniture's "
-                    "shape, design, color and materials. Match the "
-                    "perspective, scale, lighting and shadows to the room."
+                    "Transfer the complete furniture and decor "
+                    "ensemble from the second image into the first room. "
+                    "Preserve the architecture, walls, floor, windows, "
+                    "doors and camera perspective of the first image. "
+                    "Recreate the same furniture, decorations, rugs, "
+                    "lamps, plants and artwork from the reference room, "
+                    "adapting scale, perspective, lighting and shadows "
+                    "to the target room."
                 ),
-                lines=4
+                lines=6
             )
 
             with gr.Accordion(
@@ -543,9 +863,9 @@ def create_demo(model):
 
                 num_steps = gr.Slider(
                     label="Steps",
-                    minimum=10,
+                    minimum=20,
                     maximum=50,
-                    value=30,
+                    value=35,
                     step=1
                 )
 
@@ -559,25 +879,25 @@ def create_demo(model):
 
                 guidance_scale = gr.Slider(
                     label="Guidance Scale",
-                    minimum=1.0,
-                    maximum=20.0,
+                    minimum=4.0,
+                    maximum=12.0,
                     value=7.5,
                     step=0.5
                 )
 
                 strength = gr.Slider(
                     label="Room editing strength",
-                    minimum=0.3,
-                    maximum=1.0,
-                    value=0.85,
+                    minimum=0.35,
+                    maximum=0.85,
+                    value=0.62,
                     step=0.05
                 )
 
                 ip_scale = gr.Slider(
-                    label="Furniture reference strength",
-                    minimum=0.2,
-                    maximum=1.2,
-                    value=0.75,
+                    label="Reference furniture strength",
+                    minimum=0.40,
+                    maximum=1.20,
+                    value=0.90,
                     step=0.05
                 )
 
@@ -593,30 +913,37 @@ def create_demo(model):
                 a_prompt = gr.Textbox(
                     label="Added Prompt",
                     value=(
-                        "interior design, 4K, high resolution, "
-                        "photorealistic"
+                        "professional interior photography, "
+                        "photorealistic interior, realistic materials, "
+                        "natural lighting, realistic shadows, "
+                        "correct perspective, correct scale, "
+                        "high detail, 4K"
                     )
                 )
 
                 n_prompt = gr.Textbox(
                     label="Negative Prompt",
                     value=(
-                        "window, door, low resolution, banner, logo, "
-                        "watermark, text, deformed, blurry, out of focus, "
-                        "surreal, ugly, beginner, duplicated furniture, "
-                        "extra furniture, distorted furniture"
-                    )
+                        "deformed furniture, distorted furniture, "
+                        "warped furniture, melted furniture, "
+                        "wrong furniture geometry, duplicated furniture, "
+                        "extra furniture, floating furniture, "
+                        "impossible perspective, impossible proportions, "
+                        "bad architecture, distorted walls, "
+                        "distorted floor, changed windows, changed doors, "
+                        "window, door, text, logo, watermark, "
+                        "low resolution, blurry, out of focus, "
+                        "surreal, ugly, cartoon, illustration, "
+                        "CGI, plastic, oversaturated"
+                    ),
+                    lines=6
                 )
 
             submit = gr.Button(
-                "Generate furnished room",
+                "Transfer furniture from reference room",
                 variant="primary"
             )
 
-        # ----------------------------------------------------
-        # RIGHT
-        # ----------------------------------------------------
-
         with gr.Column():
 
             design_image = gr.Image(
@@ -624,10 +951,6 @@ def create_demo(model):
                 elem_id="img-display-output"
             )
 
-    # ========================================================
-    # SUBMIT
-    # ========================================================
-
     def on_submit(
         room,
         furniture,
@@ -643,13 +966,15 @@ def create_demo(model):
     ):
 
         if room is None:
+
             raise gr.Error(
-                "Please upload a photo of the room."
+                "Please upload the target room image."
             )
 
         if furniture is None:
+
             raise gr.Error(
-                "Please upload a photo of the furniture."
+                "Please upload a fully furnished reference room image."
             )
 
         model.seed = int(seed)
@@ -664,11 +989,26 @@ def create_demo(model):
                 room,
                 furniture,
                 text,
-                guidance_scale=guidance_scale,
-                num_steps=int(num_steps),
-                strength=strength,
-                img_size=int(img_size),
-                ip_scale=ip_scale
+
+                guidance_scale=float(
+                    guidance_scale
+                ),
+
+                num_steps=int(
+                    num_steps
+                ),
+
+                strength=float(
+                    strength
+                ),
+
+                img_size=int(
+                    img_size
+                ),
+
+                ip_scale=float(
+                    ip_scale
+                )
             )
 
         return out_img
@@ -698,12 +1038,21 @@ def create_demo(model):
 # LOAD MODELS
 # ============================================================
 
+print(
+    "Loading ControlNet depth..."
+)
+
 controlnet_depth = ControlNetModel.from_pretrained(
     "controlnet_depth",
     torch_dtype=dtype,
     use_safetensors=True
 )
 
+
+print(
+    "Loading ControlNet segmentation..."
+)
+
 controlnet_seg = ControlNetModel.from_pretrained(
     "own_controlnet",
     torch_dtype=dtype,
@@ -711,6 +1060,10 @@ controlnet_seg = ControlNetModel.from_pretrained(
 )
 
 
+print(
+    "Loading Stable Diffusion..."
+)
+
 pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     "SG161222/Realistic_Vision_V5.1_noVAE",
 
@@ -729,36 +1082,71 @@ pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
 # IP-ADAPTER
 # ============================================================
 
+print(
+    "Loading IP-Adapter..."
+)
+
 pipe.load_ip_adapter(
     "h94/IP-Adapter",
     subfolder="models",
     weight_name="ip-adapter_sd15.bin"
 )
 
-pipe.set_ip_adapter_scale(0.75)
+pipe.set_ip_adapter_scale(
+    0.90
+)
+
 
-pipe.scheduler = UniPCMultistepScheduler.from_config(
-    pipe.scheduler.config
+# ============================================================
+# SCHEDULER
+# ============================================================
+
+pipe.scheduler = (
+    UniPCMultistepScheduler.from_config(
+        pipe.scheduler.config
+    )
 )
 
-pipe = pipe.to(device)
+
+pipe = pipe.to(
+    device
+)
 
 
 # ============================================================
-# VISION MODELS
+# SEGMENTATION MODEL
 # ============================================================
 
-seg_image_processor, image_segmentor = (
-    get_segmentation_pipeline()
+print(
+    "Loading segmentation model..."
+)
+
+(
+    seg_image_processor,
+    image_segmentor
+) = get_segmentation_pipeline()
+
+image_segmentor = (
+    image_segmentor.to(device)
 )
 
-image_segmentor = image_segmentor.to(device)
 
-depth_feature_extractor, depth_estimator = (
-    get_depth_pipeline()
+# ============================================================
+# DEPTH MODEL
+# ============================================================
+
+print(
+    "Loading depth model..."
 )
 
-depth_estimator = depth_estimator.to(device)
+(
+    depth_feature_extractor,
+    depth_estimator
+) = get_depth_pipeline()
+
+depth_estimator = (
+    depth_estimator.to(device)
+)
 
 
 # ============================================================
@@ -767,19 +1155,34 @@ depth_estimator = depth_estimator.to(device)
 
 def main():
 
-    model = ControlNetDepthDesignModelMulti()
+    model = (
+        ControlNetDepthDesignModelMulti()
+    )
 
     print(
-        "StableDesign Furniture models loaded successfully."
+        "StableDesign Furniture Transfer "
+        "models loaded successfully."
     )
 
-    title = "# StableDesign Furniture"
+    title = (
+        "# StableDesign — Furniture Transfer"
+    )
 
     description = """
     <p>
-    Upload a photo of an empty room and a photo of a furniture item.
-    The furniture image is used as a visual reference and the room
-    geometry is preserved using depth and semantic segmentation.
+    <b>Target room:</b> upload the room that must keep
+    its architecture and camera perspective.
+    </p>
+
+    <p>
+    <b>Reference room:</b> upload a completely furnished
+    room. The system automatically extracts the visible
+    furniture and decorative elements and uses them as
+    visual references for the target room.
+    </p>
+
+    <p>
+    No manual furniture placement area is required.
     </p>
     """
 
@@ -787,11 +1190,17 @@ def main():
         css=css
     ) as demo:
 
-        gr.Markdown(title)
+        gr.Markdown(
+            title
+        )
 
-        gr.Markdown(description)
+        gr.Markdown(
+            description
+        )
 
-        create_demo(model)
+        create_demo(
+            model
+        )
 
     demo.queue().launch(
         share=False
@@ -799,4 +1208,5 @@ def main():
 
 
 if __name__ == "__main__":
+
     main()