Fix all 2

app.py

@@ -1,11 +1,10 @@
 import spaces
 
 from typing import Tuple, Union, List
-import os
 import gc
 
 import numpy as np
-from PIL import Image, ImageFilter
+from PIL import Image, ImageFilter, ImageOps
 
 import torch
 
@@ -22,6 +21,7 @@ from transformers import (
     AutoImageProcessor,
     UperNetForSemanticSegmentation,
     AutoModelForDepthEstimation,
+    CLIPVisionModelWithProjection,
 )
 
 from colors import ade_palette
@@ -31,15 +31,19 @@ import gradio as gr
 
 
 # ============================================================
-# CONFIG
+# GLOBAL CONFIG
 # ============================================================
 
 device = "cuda"
 dtype = torch.float16
 
+# Maximum number of individual furniture regions that will be
+# extracted from the reference image.
+MAX_REFERENCE_OBJECTS = 8
+
 
 # ============================================================
-# UI CSS
+# CSS
 # ============================================================
 
 css = """
@@ -58,26 +62,26 @@ css = """
 
 
 # ============================================================
-# UTILS
+# MEMORY
 # ============================================================
 
 def flush():
-    """
-    Release unused GPU memory.
-    """
-
     gc.collect()
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
 
+# ============================================================
+# IMAGE RESIZING
+# ============================================================
+
 def resize_dimensions(
     dimensions,
     target_size
 ):
     """
-    Resize while preserving the original aspect ratio.
+    Preserve aspect ratio.
     """
 
     width, height = dimensions
@@ -85,34 +89,31 @@ def resize_dimensions(
     if width <= target_size and height <= target_size:
         return dimensions
 
-    if width > height:
+    if width >= height:
 
-        aspect_ratio = height / width
+        ratio = height / width
 
         return (
             target_size,
-            max(64, int(target_size * aspect_ratio))
+            max(64, int(target_size * ratio))
         )
 
-    aspect_ratio = width / height
+    ratio = width / height
 
     return (
-        max(64, int(target_size * aspect_ratio)),
+        max(64, int(target_size * ratio)),
         target_size
     )
 
 
-def resize_keep_aspect(
+def fit_image(
     image: Image.Image,
-    target_width: int,
-    target_height: int,
-    background=(255, 255, 255)
+    width: int,
+    height: int,
+    background=(235, 235, 235)
 ):
     """
-    Resize an image preserving its aspect ratio and place it
-    in a canvas of the requested size.
-
-    This is intentionally NOT a stretch resize.
+    Fit an image into a canvas without stretching it.
     """
 
     image = image.convert("RGB")
@@ -120,12 +121,19 @@ def resize_keep_aspect(
     src_w, src_h = image.size
 
     scale = min(
-        target_width / src_w,
-        target_height / src_h
+        width / src_w,
+        height / src_h
     )
 
-    new_w = max(1, int(src_w * scale))
-    new_h = max(1, int(src_h * scale))
+    new_w = max(
+        1,
+        int(src_w * scale)
+    )
+
+    new_h = max(
+        1,
+        int(src_h * scale)
+    )
 
     resized = image.resize(
         (new_w, new_h),
@@ -134,12 +142,12 @@ def resize_keep_aspect(
 
     canvas = Image.new(
         "RGB",
-        (target_width, target_height),
+        (width, height),
         background
     )
 
-    x = (target_width - new_w) // 2
-    y = (target_height - new_h) // 2
+    x = (width - new_w) // 2
+    y = (height - new_h) // 2
 
     canvas.paste(
         resized,
@@ -149,93 +157,35 @@ def resize_keep_aspect(
     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
-    )
-
-    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 wanted_items:
-
-            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(
+    processor = AutoImageProcessor.from_pretrained(
         "openmmlab/upernet-convnext-small"
     )
 
-    image_segmentor = (
-        UperNetForSemanticSegmentation.from_pretrained(
-            "openmmlab/upernet-convnext-small"
-        )
+    model = UperNetForSemanticSegmentation.from_pretrained(
+        "openmmlab/upernet-convnext-small"
     )
 
-    return (
-        image_processor,
-        image_segmentor
-    )
+    return processor, model
 
 
 @torch.inference_mode()
 @spaces.GPU
 def segment_image(
     image: Image.Image,
-    image_processor: AutoImageProcessor,
-    image_segmentor: UperNetForSemanticSegmentation
-) -> Image.Image:
+    image_processor,
+    image_segmentor
+):
     """
-    Generate an RGB ADE20K segmentation image.
+    ADE20K semantic segmentation.
     """
 
-    pixel_values = image_processor(
+    inputs = image_processor(
         image,
         return_tensors="pt"
     ).pixel_values.to(device)
@@ -243,10 +193,10 @@ def segment_image(
     with torch.no_grad():
 
         outputs = image_segmentor(
-            pixel_values
+            inputs
         )
 
-    seg = (
+    segmentation = (
         image_processor
         .post_process_semantic_segmentation(
             outputs,
@@ -255,17 +205,18 @@ def segment_image(
     )
 
     # IMPORTANT:
-    # Move tensor from CUDA to CPU before NumPy.
-    seg = (
-        seg.detach()
+    # CUDA -> CPU before NumPy.
+    segmentation = (
+        segmentation
+        .detach()
         .cpu()
         .numpy()
     )
 
     color_seg = np.zeros(
         (
-            seg.shape[0],
-            seg.shape[1],
+            segmentation.shape[0],
+            segmentation.shape[1],
             3
         ),
         dtype=np.uint8
@@ -278,7 +229,7 @@ def segment_image(
     for label, color in enumerate(palette):
 
         color_seg[
-            seg == label,
+            segmentation == label,
             :
         ] = color
 
@@ -313,79 +264,83 @@ def get_depth_pipeline():
 @spaces.GPU
 def get_depth_image(
     image: Image.Image,
-    feature_extractor: AutoImageProcessor,
-    depth_estimator: AutoModelForDepthEstimation
-) -> Image.Image:
-    """
-    Generate normalized depth image.
-    """
+    feature_extractor,
+    depth_estimator
+):
 
-    image_to_depth = feature_extractor(
+    inputs = feature_extractor(
         images=image,
         return_tensors="pt"
     ).to(device)
 
     with torch.no_grad():
 
-        depth_map = (
+        depth = (
             depth_estimator(
-                **image_to_depth
+                **inputs
             )
             .predicted_depth
         )
 
     width, height = image.size
 
-    depth_map = torch.nn.functional.interpolate(
-        depth_map.unsqueeze(1).float(),
+    depth = torch.nn.functional.interpolate(
+        depth.unsqueeze(1).float(),
         size=(height, width),
         mode="bicubic",
         align_corners=False,
     )
 
     depth_min = torch.amin(
-        depth_map,
+        depth,
         dim=[1, 2, 3],
         keepdim=True
     )
 
     depth_max = torch.amax(
-        depth_map,
+        depth,
         dim=[1, 2, 3],
         keepdim=True
     )
 
-    depth_map = (
-        depth_map - depth_min
+    depth = (
+        depth - depth_min
     ) / (
         depth_max - depth_min + 1e-8
     )
 
-    image = torch.cat(
-        [depth_map] * 3,
+    depth = torch.cat(
+        [depth] * 3,
         dim=1
     )
 
-    image = (
-        image
+    depth = (
+        depth
         .permute(0, 2, 3, 1)
         .cpu()
         .numpy()[0]
     )
 
     return Image.fromarray(
-        (image * 255.0)
+        (
+            depth * 255
+        )
         .clip(0, 255)
         .astype(np.uint8)
     )
 
 
 # ============================================================
-# REFERENCE ROOM -> FURNITURE EXTRACTION
+# ADE20K OBJECT CLASSES
 # ============================================================
 
-REFERENCE_FURNITURE_CLASSES = {
-    # Main furniture
+# These are the classes we WANT to transfer from the
+# reference room.
+#
+# We intentionally do not include walls, floors, ceiling,
+# windows, doors, etc.
+
+FURNITURE_CLASSES = {
     "chair",
     "sofa",
     "table",
@@ -396,27 +351,19 @@ REFERENCE_FURNITURE_CLASSES = {
     "dresser",
     "bench",
     "armchair",
-
-    # Small furniture / objects
     "ottoman",
     "stool",
-    "pillow",
     "lamp",
-    "television",
-    "computer",
-    "monitor",
-
-    # Decoration
+    "pillow",
     "plant",
     "vase",
     "painting",
     "mirror",
+    "television",
+    "computer",
+    "monitor",
     "clock",
-
-    # Floor / textile elements
     "rug",
-
-    # Additional ADE-like categories
     "book",
     "books",
     "chest of drawers",
@@ -424,136 +371,503 @@ REFERENCE_FURNITURE_CLASSES = {
 }
 
 
-def extract_furniture_reference(
-    reference_image: Image.Image,
-    segmentation_image: Image.Image
-) -> Image.Image:
-    """
-    Remove the architectural/background part of the reference room.
-
-    The goal is NOT to extract a single object.
+# ============================================================
+# SEGMENTATION -> BINARY FURNITURE MASK
+# ============================================================
 
-    Instead we preserve the complete set of recognizable
-    furniture/decor elements from the reference room.
+def get_class_mask(
+    segmentation: Image.Image,
+    wanted_classes
+):
+    """
+    Convert ADE20K RGB segmentation into a binary mask.
     """
 
-    mask = create_binary_mask_from_segmentation(
-        segmentation_image,
-        list(REFERENCE_FURNITURE_CLASSES)
+    seg = np.array(
+        segmentation
     )
 
-    # Slightly enlarge furniture regions so that thin objects
-    # and object edges are not destroyed by semantic segmentation.
-    mask = mask.filter(
-        ImageFilter.MaxFilter(9)
+    unique_colors = np.unique(
+        seg.reshape(
+            -1,
+            3
+        ),
+        axis=0
     )
 
-    reference_rgb = reference_image.convert(
-        "RGB"
+    mask = np.zeros(
+        seg.shape[:2],
+        dtype=np.uint8
     )
 
-    # 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)
+    for color in unique_colors:
+
+        color_tuple = tuple(
+            int(x)
+            for x in color
+        )
+
+        try:
+            item = map_colors_rgb(
+                color_tuple
+            )
+        except Exception:
+            continue
+
+        if item in wanted_classes:
+
+            pixels = (
+                seg == color
+            ).all(axis=2)
+
+            mask[pixels] = 255
+
+    return Image.fromarray(
+        mask,
+        mode="L"
     )
 
-    furniture_reference = Image.composite(
-        reference_rgb,
-        background,
+
+# ============================================================
+# CONNECTED COMPONENTS
+# ============================================================
+
+def connected_components(
+    mask: Image.Image,
+    min_area=150
+):
+    """
+    Find connected regions in a binary mask.
+
+    Pure NumPy implementation so no additional dependency
+    such as OpenCV is required.
+    """
+
+    arr = np.array(
         mask
     )
 
-    return furniture_reference
+    binary = arr > 0
+
+    height, width = binary.shape
+
+    visited = np.zeros(
+        binary.shape,
+        dtype=bool
+    )
+
+    components = []
+
+    # 8-connected neighborhood
+    neighbors = [
+        (-1, -1),
+        (-1, 0),
+        (-1, 1),
+        (0, -1),
+        (0, 1),
+        (1, -1),
+        (1, 0),
+        (1, 1),
+    ]
+
+    for y in range(height):
+
+        for x in range(width):
+
+            if not binary[y, x]:
+                continue
+
+            if visited[y, x]:
+                continue
+
+            stack = [(x, y)]
+
+            visited[y, x] = True
+
+            points = []
+
+            while stack:
+
+                px, py = stack.pop()
+
+                points.append(
+                    (px, py)
+                )
+
+                for dx, dy in neighbors:
+
+                    nx = px + dx
+                    ny = py + dy
+
+                    if nx < 0 or nx >= width:
+                        continue
+
+                    if ny < 0 or ny >= height:
+                        continue
+
+                    if not binary[ny, nx]:
+                        continue
+
+                    if visited[ny, nx]:
+                        continue
+
+                    visited[ny, nx] = True
+
+                    stack.append(
+                        (nx, ny)
+                    )
+
+            if len(points) < min_area:
+                continue
+
+            xs = [
+                p[0]
+                for p in points
+            ]
+
+            ys = [
+                p[1]
+                for p in points
+            ]
+
+            x0 = max(
+                0,
+                min(xs)
+            )
+
+            y0 = max(
+                0,
+                min(ys)
+            )
+
+            x1 = min(
+                width - 1,
+                max(xs)
+            )
+
+            y1 = min(
+                height - 1,
+                max(ys)
+            )
+
+            components.append(
+                {
+                    "area": len(points),
+                    "bbox": (
+                        x0,
+                        y0,
+                        x1 + 1,
+                        y1 + 1
+                    )
+                }
+            )
+
+    components.sort(
+        key=lambda x: x["area"],
+        reverse=True
+    )
+
+    return components
 
 
 # ============================================================
-# TARGET ROOM MASK
+# REFERENCE OBJECT EXTRACTION
 # ============================================================
 
-ARCHITECTURE_CLASSES = {
-    "wall",
-    "floor",
-    "ceiling",
-    "window",
-    "windowpane",
-    "door",
-    "door;double;door",
-    "curtain",
-}
+def extract_reference_objects(
+    reference_image: Image.Image,
+    reference_segmentation: Image.Image
+):
+    """
+    Extract the largest furniture/decor regions from the
+    reference room.
+
+    Returns a list of cropped RGB object images.
+    """
+
+    furniture_mask = get_class_mask(
+        reference_segmentation,
+        FURNITURE_CLASSES
+    )
+
+    # Small dilation to avoid cutting chair legs, lamp edges,
+    # sofa edges, etc.
+    furniture_mask = furniture_mask.filter(
+        ImageFilter.MaxFilter(5)
+    )
+
+    components = connected_components(
+        furniture_mask,
+        min_area=120
+    )
+
+    objects = []
+
+    for component in components:
+
+        if len(objects) >= MAX_REFERENCE_OBJECTS:
+            break
 
+        x0, y0, x1, y1 = component["bbox"]
 
-def create_room_edit_mask(
-    segmentation_image: Image.Image
-) -> Image.Image:
+        box_w = x1 - x0
+        box_h = y1 - y0
+
+        # Ignore extremely tiny regions.
+        if box_w < 20 or box_h < 20:
+            continue
+
+        # Add contextual margin around object.
+        margin_x = max(
+            8,
+            int(box_w * 0.10)
+        )
+
+        margin_y = max(
+            8,
+            int(box_h * 0.10)
+        )
+
+        x0 = max(
+            0,
+            x0 - margin_x
+        )
+
+        y0 = max(
+            0,
+            y0 - margin_y
+        )
+
+        x1 = min(
+            reference_image.width,
+            x1 + margin_x
+        )
+
+        y1 = min(
+            reference_image.height,
+            y1 + margin_y
+        )
+
+        crop = reference_image.crop(
+            (
+                x0,
+                y0,
+                x1,
+                y1
+            )
+        )
+
+        objects.append(
+            crop
+        )
+
+    return objects
+
+
+# ============================================================
+# BUILD REFERENCE BOARD
+# ============================================================
+
+def build_reference_board(
+    reference_image: Image.Image,
+    reference_segmentation: Image.Image,
+    board_size=640
+):
     """
-    Create the region where furniture can be generated.
+    Create a visual board containing:
 
-    Architecture is kept strongly constrained by ControlNet.
+    1. The complete reference room, slightly reduced.
+    2. Individual furniture/decor crops.
 
-    The mask deliberately does not depend on a manually selected
-    rectangle or furniture placement area.
+    This gives IP-Adapter both scene context and object detail.
     """
 
-    seg_rgb = np.array(
-        segmentation_image
+    reference_image = reference_image.convert(
+        "RGB"
     )
 
-    unique_colors = np.unique(
-        seg_rgb.reshape(
-            -1,
-            seg_rgb.shape[2]
+    # --------------------------------------------------------
+    # Extract individual objects
+    # --------------------------------------------------------
+
+    objects = extract_reference_objects(
+        reference_image,
+        reference_segmentation
+    )
+
+    # --------------------------------------------------------
+    # Board
+    # --------------------------------------------------------
+
+    board = Image.new(
+        "RGB",
+        (
+            board_size,
+            board_size
         ),
-        axis=0
+        (245, 245, 245)
     )
 
-    # Start with everything editable.
-    editable = np.ones(
-        seg_rgb.shape[:2],
-        dtype=np.uint8
+    # --------------------------------------------------------
+    # Main reference image
+    # --------------------------------------------------------
+
+    main_size = int(
+        board_size * 0.52
     )
 
-    for color in unique_colors:
+    main_image = fit_image(
+        reference_image,
+        main_size,
+        main_size,
+        background=(235, 235, 235)
+    )
 
-        color_tuple = tuple(
-            int(v)
-            for v in color
+    board.paste(
+        main_image,
+        (10, 10)
+    )
+
+    # --------------------------------------------------------
+    # Object thumbnails
+    # --------------------------------------------------------
+
+    if len(objects) == 0:
+
+        # Fallback:
+        # use the complete reference image.
+        return fit_image(
+            reference_image,
+            board_size,
+            board_size,
+            background=(235, 235, 235)
         )
 
-        item_name = map_colors_rgb(
-            color_tuple
+    thumb_area_x = (
+        board_size
+        - main_size
+        - 30
+    )
+
+    thumb_area_y = (
+        board_size
+        - 20
+    )
+
+    columns = 2
+
+    rows = max(
+        1,
+        int(
+            np.ceil(
+                len(objects) / columns
+            )
         )
+    )
 
-        if item_name in ARCHITECTURE_CLASSES:
+    cell_w = max(
+        64,
+        thumb_area_x // columns
+    )
 
-            matches = (
-                seg_rgb == color
-            ).all(axis=2)
+    cell_h = max(
+        64,
+        thumb_area_y // min(rows, 4)
+    )
 
-            # Windows, doors and architectural areas are
-            # protected from direct inpainting.
-            editable[matches] = 0
+    start_x = (
+        main_size
+        + 20
+    )
 
-    mask = Image.fromarray(
-        editable * 255,
-        mode="L"
+    start_y = 10
+
+    for index, obj in enumerate(objects):
+
+        row = index // columns
+        col = index % columns
+
+        if row >= 4:
+            break
+
+        x = (
+            start_x
+            + col * cell_w
+        )
+
+        y = (
+            start_y
+            + row * cell_h
+        )
+
+        thumb = fit_image(
+            obj,
+            max(32, cell_w - 8),
+            max(32, cell_h - 8),
+            background=(235, 235, 235)
+        )
+
+        board.paste(
+            thumb,
+            (x, y)
+        )
+
+    return board
+
+
+# ============================================================
+# TARGET INPAINT MASK
+# ============================================================
+
+def build_target_inpaint_mask(
+    target_segmentation: Image.Image
+):
+    """
+    The target room is intentionally allowed to change.
+
+    The architecture is preserved by ControlNet rather than
+    by freezing the entire target image.
+
+    We use a mostly-full mask with slightly protected borders.
+    """
+
+    width, height = target_segmentation.size
+
+    mask = Image.new(
+        "L",
+        (
+            width,
+            height
+        ),
+        255
     )
 
-    # Keep furniture transitions soft.
-    mask = mask.filter(
-        ImageFilter.GaussianBlur(2)
+    # Protect a narrow outer border.
+    # This reduces accidental changes at image boundaries.
+    arr = np.array(
+        mask
     )
 
-    return mask
+    border = max(
+        2,
+        int(
+            min(width, height) * 0.015
+        )
+    )
+
+    arr[:border, :] = 0
+    arr[-border:, :] = 0
+    arr[:, :border] = 0
+    arr[:, -border:] = 0
+
+    return Image.fromarray(
+        arr,
+        mode="L"
+    )
 
 
 # ============================================================
-# MAIN MODEL
+# MODEL
 # ============================================================
 
 class ControlNetDepthDesignModelMulti:
@@ -565,50 +879,67 @@ class ControlNetDepthDesignModelMulti:
         self.neg_prompt = (
             "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, "
+            "duplicated furniture, extra furniture, "
+            "floating furniture, furniture without contact "
+            "with the floor, impossible perspective, "
+            "impossible proportions, bad architecture, "
+            "distorted walls, distorted floor, "
             "changed windows, changed doors, "
-            "window, door, text, logo, watermark, "
+            "new windows, new doors, "
+            "text, logo, watermark, "
             "low resolution, blurry, out of focus, "
-            "surreal, ugly, cartoon, illustration, "
-            "CGI, plastic, oversaturated"
+            "surreal, cartoon, illustration, CGI, "
+            "plastic materials, oversaturated"
         )
 
         self.additional_quality_suffix = (
-            "professional interior photography, "
-            "photorealistic interior, realistic materials, "
-            "natural lighting, realistic shadows, "
-            "correct perspective, correct scale, "
-            "high detail, 4K"
+            "photorealistic interior photography, "
+            "realistic furniture materials, "
+            "realistic global illumination, "
+            "natural shadows, "
+            "correct perspective, "
+            "correct object scale, "
+            "high detail, "
+            "high resolution"
         )
 
     @spaces.GPU
     def generate_design(
         self,
-        empty_room_image: Image.Image,
-        furniture_image: Image.Image,
+        target_room: Image.Image,
+        reference_room: Image.Image,
         prompt: str,
-        guidance_scale: float = 7.5,
+        guidance_scale: float = 7.0,
         num_steps: int = 35,
-        strength: float = 0.62,
+        strength: float = 0.72,
         img_size: int = 640,
-        ip_scale: float = 0.90
-    ) -> Image.Image:
+        ip_scale: float = 0.95
+    ):
+
+        print(
+            "=========================================="
+        )
+
+        print(
+            "Starting reference-room transfer..."
+        )
 
         print(
-            "Starting complete room furnishing transfer..."
+            "=========================================="
         )
 
         flush()
 
-        self.generator = torch.Generator(
+        generator = torch.Generator(
             device=device
         ).manual_seed(
             int(self.seed)
         )
 
+        # ----------------------------------------------------
+        # PROMPT
+        # ----------------------------------------------------
+
         if prompt is None:
             prompt = ""
 
@@ -617,15 +948,20 @@ class ControlNetDepthDesignModelMulti:
         if prompt == "":
 
             prompt = (
-                "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 "
+                "A professionally photographed interior "
+                "of the target room furnished with the "
+                "complete furniture and decoration ensemble "
                 "from the reference room. "
-                "Adapt their size, position, perspective, lighting "
-                "and shadows naturally to the target room."
+                "Transfer the sofa, chairs, tables, cabinets, "
+                "lamps, rugs, plants, artwork and decorative "
+                "objects from the reference room. "
+                "Keep the target room architecture, "
+                "walls, ceiling, floor, windows, doors and "
+                "camera viewpoint unchanged. "
+                "Adapt the transferred furniture naturally "
+                "to the target room geometry. "
+                "Preserve realistic proportions, "
+                "perspective, materials, lighting and shadows."
             )
 
         pos_prompt = (
@@ -634,119 +970,151 @@ class ControlNetDepthDesignModelMulti:
             + self.additional_quality_suffix
         )
 
-        # --------------------------------------------------------
-        # TARGET ROOM
-        # --------------------------------------------------------
+        # ----------------------------------------------------
+        # TARGET SIZE
+        # ----------------------------------------------------
 
-        orig_w, orig_h = (
-            empty_room_image.size
+        original_width, original_height = (
+            target_room.size
         )
 
-        new_width, new_height = resize_dimensions(
-            empty_room_image.size,
-            img_size
+        target_width, target_height = (
+            resize_dimensions(
+                target_room.size,
+                img_size
+            )
         )
 
-        input_image = (
-            empty_room_image
+        target = (
+            target_room
             .convert("RGB")
             .resize(
-                (new_width, new_height),
+                (
+                    target_width,
+                    target_height
+                ),
                 Image.Resampling.LANCZOS
             )
         )
 
-        # --------------------------------------------------------
-        # REFERENCE ROOM
-        # --------------------------------------------------------
+        # ----------------------------------------------------
+        # REFERENCE
+        # ----------------------------------------------------
 
-        reference_image = (
-            furniture_image
+        reference = (
+            reference_room
             .convert("RGB")
         )
 
-        # Resize reference WITHOUT stretching it.
-        reference_for_segmentation = resize_keep_aspect(
-            reference_image,
-            new_width,
-            new_height,
+        reference = fit_image(
+            reference,
+            target_width,
+            target_height,
             background=(235, 235, 235)
         )
 
-        # --------------------------------------------------------
+        # ----------------------------------------------------
         # TARGET SEGMENTATION
-        # --------------------------------------------------------
+        # ----------------------------------------------------
+
+        print(
+            "1/5 - Target room segmentation..."
+        )
 
         target_segmentation = segment_image(
-            input_image,
+            target,
             seg_image_processor,
             image_segmentor
         )
 
-        # --------------------------------------------------------
+        # ----------------------------------------------------
         # REFERENCE SEGMENTATION
-        # --------------------------------------------------------
+        # ----------------------------------------------------
 
         print(
-            "Segmenting reference room furniture..."
+            "2/5 - Reference room segmentation..."
         )
 
         reference_segmentation = segment_image(
-            reference_for_segmentation,
+            reference,
             seg_image_processor,
             image_segmentor
         )
 
-        # --------------------------------------------------------
-        # EXTRACT COMPLETE FURNITURE ENSEMBLE
-        # --------------------------------------------------------
-
-        furniture_only = extract_furniture_reference(
-            reference_for_segmentation,
-            reference_segmentation
-        )
-
-        # --------------------------------------------------------
-        # SAVE/PRINT INFORMATION FOR DEBUGGING
-        # --------------------------------------------------------
+        # ----------------------------------------------------
+        # BUILD OBJECT-AWARE REFERENCE BOARD
+        # ----------------------------------------------------
 
         print(
-            "Furniture reference extracted successfully."
+            "3/5 - Extracting furniture and decorations..."
         )
 
-        # --------------------------------------------------------
-        # TARGET EDIT MASK
-        # --------------------------------------------------------
+        reference_board = build_reference_board(
+            reference,
+            reference_segmentation,
+            board_size=max(
+                512,
+                target_width
+            )
+        )
 
-        mask_image = create_room_edit_mask(
-            target_segmentation
+        reference_board = reference_board.resize(
+            (
+                target_width,
+                target_height
+            ),
+            Image.Resampling.LANCZOS
         )
 
-        # --------------------------------------------------------
+        # ----------------------------------------------------
         # DEPTH
-        # --------------------------------------------------------
+        # ----------------------------------------------------
 
-        image_depth = get_depth_image(
-            input_image,
+        print(
+            "4/5 - Calculating target room depth..."
+        )
+
+        target_depth = get_depth_image(
+            target,
             depth_feature_extractor,
             depth_estimator
         )
 
+        # ----------------------------------------------------
+        # INPAINT MASK
+        # ----------------------------------------------------
+
+        # IMPORTANT:
+        #
+        # We DO NOT freeze the target room.
+        #
+        # ControlNet provides the architectural structure.
+        # The diffusion model must have enough freedom to
+        # actually create the transferred furniture.
+
+        mask_image = build_target_inpaint_mask(
+            target_segmentation
+        )
+
         flush()
 
-        # --------------------------------------------------------
+        # ----------------------------------------------------
         # IP-ADAPTER
-        # --------------------------------------------------------
+        # ----------------------------------------------------
 
         pipe.set_ip_adapter_scale(
             float(ip_scale)
         )
 
-        # --------------------------------------------------------
-        # CONTROLNET
-        # --------------------------------------------------------
+        # ----------------------------------------------------
+        # GENERATION
+        # ----------------------------------------------------
+
+        print(
+            "5/5 - Generating furnished target room..."
+        )
 
-        generated_image = pipe(
+        result = pipe(
             prompt=pos_prompt,
 
             negative_prompt=self.neg_prompt,
@@ -763,80 +1131,96 @@ class ControlNetDepthDesignModelMulti:
                 guidance_scale
             ),
 
-            generator=self.generator,
+            generator=generator,
 
-            image=input_image,
+            image=target,
 
             mask_image=mask_image,
 
             # 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,
+            # This is the reference board, not a manually
+            # selected furniture image.
+            ip_adapter_image=reference_board,
 
             control_image=[
-                image_depth,
+                target_depth,
                 target_segmentation
             ],
 
+            # Depth is slightly stronger because it carries
+            # the target room's spatial structure.
             controlnet_conditioning_scale=[
-                0.70,
+                0.80,
                 0.65
             ]
         ).images[0]
 
         flush()
 
-        design_image = generated_image.resize(
-            (orig_w, orig_h),
+        result = result.resize(
+            (
+                original_width,
+                original_height
+            ),
             Image.Resampling.LANCZOS
         )
 
         print(
-            "Furniture transfer completed."
+            "Generation completed."
         )
 
-        return design_image
+        return result
 
 
 # ============================================================
-# GRADIO UI
+# GRADIO
 # ============================================================
 
 def create_demo(model):
 
     gr.Markdown(
-        "### StableDesign — Reference Room Furniture Transfer"
+        "### StableDesign — Reference Room Transfer"
     )
 
     gr.Markdown(
         """
-        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.
+        <p>
+        <b>Image 1 — Target room:</b>
+        the room that will be furnished. Its architecture,
+        perspective, walls, floor, windows and doors are
+        preserved using structural controls.
+        </p>
+
+        <p>
+        <b>Image 2 — Reference room:</b>
+        a completely furnished room. The system automatically
+        analyzes the furniture and decorative elements and
+        uses them as the visual reference.
+        </p>
+
+        <p>
+        No furniture area or placement needs to be selected.
+        </p>
         """
     )
 
     with gr.Row():
 
+        # ====================================================
+        # INPUTS
+        # ====================================================
+
         with gr.Column():
 
-            room_image = gr.Image(
-                label="1. Target room",
+            target_image = gr.Image(
+                label="1. Target / empty room",
                 type="pil",
                 elem_id="img-display-input"
             )
 
-            furniture_image = gr.Image(
-                label="2. Furnished reference room",
+            reference_image = gr.Image(
+                label="2. Fully furnished reference room",
                 type="pil",
                 elem_id="img-display-input"
             )
@@ -844,14 +1228,15 @@ def create_demo(model):
             prompt = gr.Textbox(
                 label="Instructions",
                 value=(
-                    "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."
+                    "Transfer the complete furniture and "
+                    "decoration ensemble from image 2 "
+                    "into image 1. Preserve the architecture, "
+                    "walls, floor, windows, doors and camera "
+                    "perspective of image 1. Recreate the "
+                    "same furniture and decorative elements "
+                    "from image 2, adapting their scale, "
+                    "perspective, lighting and shadows "
+                    "naturally to image 1."
                 ),
                 lines=6
             )
@@ -881,23 +1266,23 @@ def create_demo(model):
                     label="Guidance Scale",
                     minimum=4.0,
                     maximum=12.0,
-                    value=7.5,
+                    value=7.0,
                     step=0.5
                 )
 
                 strength = gr.Slider(
-                    label="Room editing strength",
-                    minimum=0.35,
-                    maximum=0.85,
-                    value=0.62,
+                    label="Furniture generation strength",
+                    minimum=0.45,
+                    maximum=0.90,
+                    value=0.72,
                     step=0.05
                 )
 
                 ip_scale = gr.Slider(
-                    label="Reference furniture strength",
-                    minimum=0.40,
+                    label="Reference room strength",
+                    minimum=0.50,
                     maximum=1.20,
-                    value=0.90,
+                    value=0.95,
                     step=0.05
                 )
 
@@ -913,12 +1298,16 @@ def create_demo(model):
                 a_prompt = gr.Textbox(
                     label="Added Prompt",
                     value=(
-                        "professional interior photography, "
-                        "photorealistic interior, realistic materials, "
-                        "natural lighting, realistic shadows, "
-                        "correct perspective, correct scale, "
-                        "high detail, 4K"
-                    )
+                        "photorealistic interior photography, "
+                        "realistic furniture materials, "
+                        "realistic global illumination, "
+                        "natural shadows, "
+                        "correct perspective, "
+                        "correct object scale, "
+                        "high detail, "
+                        "high resolution"
+                    ),
+                    lines=4
                 )
 
                 n_prompt = gr.Textbox(
@@ -926,99 +1315,118 @@ def create_demo(model):
                     value=(
                         "deformed furniture, distorted furniture, "
                         "warped furniture, melted furniture, "
-                        "wrong furniture geometry, duplicated furniture, "
-                        "extra furniture, floating furniture, "
-                        "impossible perspective, impossible proportions, "
+                        "duplicated furniture, extra furniture, "
+                        "floating furniture, furniture without "
+                        "contact with the floor, 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"
+                        "distorted floor, changed windows, "
+                        "changed doors, new windows, new doors, "
+                        "text, logo, watermark, low resolution, "
+                        "blurry, out of focus, surreal, cartoon, "
+                        "illustration, CGI, plastic materials, "
+                        "oversaturated"
                     ),
                     lines=6
                 )
 
             submit = gr.Button(
-                "Transfer furniture from reference room",
+                "Transfer complete furniture",
                 variant="primary"
             )
 
+        # ====================================================
+        # OUTPUT
+        # ====================================================
+
         with gr.Column():
 
-            design_image = gr.Image(
+            result_image = gr.Image(
                 label="Result",
                 elem_id="img-display-output"
             )
 
+    # ========================================================
+    # SUBMIT
+    # ========================================================
+
     def on_submit(
-        room,
-        furniture,
+        target,
+        reference,
         text,
-        num_steps,
-        guidance_scale,
-        seed,
-        strength,
-        ip_scale,
-        a_prompt,
-        n_prompt,
-        img_size
+        num_steps_value,
+        guidance_value,
+        seed_value,
+        strength_value,
+        ip_value,
+        additional_prompt,
+        negative_prompt,
+        image_size
     ):
 
-        if room is None:
+        if target is None:
 
             raise gr.Error(
-                "Please upload the target room image."
+                "Please upload image 1: the target room."
             )
 
-        if furniture is None:
+        if reference is None:
 
             raise gr.Error(
-                "Please upload a fully furnished reference room image."
+                "Please upload image 2: the fully furnished "
+                "reference room."
             )
 
-        model.seed = int(seed)
+        model.seed = int(
+            seed_value
+        )
 
-        model.neg_prompt = n_prompt
+        model.neg_prompt = (
+            negative_prompt
+        )
 
-        model.additional_quality_suffix = a_prompt
+        model.additional_quality_suffix = (
+            additional_prompt
+        )
 
         with torch.no_grad():
 
-            out_img = model.generate_design(
-                room,
-                furniture,
-                text,
+            output = model.generate_design(
+                target_room=target,
+
+                reference_room=reference,
+
+                prompt=text,
 
                 guidance_scale=float(
-                    guidance_scale
+                    guidance_value
                 ),
 
                 num_steps=int(
-                    num_steps
+                    num_steps_value
                 ),
 
                 strength=float(
-                    strength
+                    strength_value
                 ),
 
                 img_size=int(
-                    img_size
+                    image_size
                 ),
 
                 ip_scale=float(
-                    ip_scale
+                    ip_value
                 )
             )
 
-        return out_img
+        return output
 
     submit.click(
         on_submit,
 
         inputs=[
-            room_image,
-            furniture_image,
+            target_image,
+            reference_image,
             prompt,
             num_steps,
             guidance_scale,
@@ -1030,12 +1438,12 @@ def create_demo(model):
             img_size
         ],
 
-        outputs=design_image
+        outputs=result_image
     )
 
 
 # ============================================================
-# LOAD MODELS
+# LOAD CONTROLNET
 # ============================================================
 
 print(
@@ -1060,40 +1468,67 @@ controlnet_seg = ControlNetModel.from_pretrained(
 )
 
 
+# ============================================================
+# LOAD CLIP VISION FOR IP-ADAPTER PLUS
+# ============================================================
+
+print(
+    "Loading IP-Adapter Plus image encoder..."
+)
+
+image_encoder = (
+    CLIPVisionModelWithProjection
+    .from_pretrained(
+        "h94/IP-Adapter",
+        subfolder="models/image_encoder",
+        torch_dtype=dtype
+    )
+)
+
+
+# ============================================================
+# LOAD STABLE DIFFUSION
+# ============================================================
+
 print(
     "Loading Stable Diffusion..."
 )
 
-pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-    "SG161222/Realistic_Vision_V5.1_noVAE",
+pipe = (
+    StableDiffusionControlNetInpaintPipeline
+    .from_pretrained(
+        "SG161222/Realistic_Vision_V5.1_noVAE",
 
-    controlnet=[
-        controlnet_depth,
-        controlnet_seg
-    ],
+        controlnet=[
+            controlnet_depth,
+            controlnet_seg
+        ],
 
-    safety_checker=None,
+        image_encoder=image_encoder,
 
-    torch_dtype=dtype
+        safety_checker=None,
+
+        torch_dtype=dtype
+    )
 )
 
 
 # ============================================================
-# IP-ADAPTER
+# IP-ADAPTER PLUS
 # ============================================================
 
 print(
-    "Loading IP-Adapter..."
+    "Loading IP-Adapter Plus SD1.5..."
 )
 
 pipe.load_ip_adapter(
     "h94/IP-Adapter",
     subfolder="models",
-    weight_name="ip-adapter_sd15.bin"
+    weight_name="ip-adapter-plus_sd15.safetensors"
 )
 
 pipe.set_ip_adapter_scale(
-    0.90
+    0.95
 )
 
 
@@ -1114,11 +1549,11 @@ pipe = pipe.to(
 
 
 # ============================================================
-# SEGMENTATION MODEL
+# SEGMENTATION
 # ============================================================
 
 print(
-    "Loading segmentation model..."
+    "Loading ADE20K segmentation..."
 )
 
 (
@@ -1132,11 +1567,11 @@ image_segmentor = (
 
 
 # ============================================================
-# DEPTH MODEL
+# DEPTH
 # ============================================================
 
 print(
-    "Loading depth model..."
+    "Loading Depth Anything..."
 )
 
 (
@@ -1160,42 +1595,36 @@ def main():
     )
 
     print(
-        "StableDesign Furniture Transfer "
-        "models loaded successfully."
+        "=========================================="
     )
 
-    title = (
-        "# StableDesign — Furniture Transfer"
+    print(
+        "StableDesign V2 loaded successfully."
     )
 
-    description = """
-    <p>
-    <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>
-    """
+    print(
+        "Reference-room furniture transfer enabled."
+    )
+
+    print(
+        "=========================================="
+    )
 
     with gr.Blocks(
         css=css
     ) as demo:
 
         gr.Markdown(
-            title
+            "# StableDesign V2"
         )
 
         gr.Markdown(
-            description
+            """
+            Upload an empty/target room and a fully furnished
+            reference room. The system automatically analyzes
+            the reference furniture and recreates the furnishing
+            ensemble inside the target room.
+            """
         )
 
         create_demo(