app.py

import spaces

from typing import Tuple, Union, List
import gc

import numpy as np
from PIL import Image, ImageFilter, ImageOps

import torch

from diffusers.pipelines.controlnet import (
    StableDiffusionControlNetInpaintPipeline
)

from diffusers import (
    ControlNetModel,
    UniPCMultistepScheduler,
)

from transformers import (
    AutoImageProcessor,
    UperNetForSemanticSegmentation,
    AutoModelForDepthEstimation,
    CLIPVisionModelWithProjection,
)

from colors import ade_palette
from utils import map_colors_rgb

import gradio as gr


# ============================================================
# 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


# ============================================================
# CSS
# ============================================================

css = """
#img-display-container {
    max-height: 50vh;
}

#img-display-input {
    max-height: 40vh;
}

#img-display-output {
    max-height: 40vh;
}
"""


# ============================================================
# MEMORY
# ============================================================

def flush():
    gc.collect()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# ============================================================
# IMAGE RESIZING
# ============================================================

def resize_dimensions(
    dimensions,
    target_size
):
    """
    Preserve aspect ratio.
    """

    width, height = dimensions

    if width <= target_size and height <= target_size:
        return dimensions

    if width >= height:

        ratio = height / width

        return (
            target_size,
            max(64, int(target_size * ratio))
        )

    ratio = width / height

    return (
        max(64, int(target_size * ratio)),
        target_size
    )


def fit_image(
    image: Image.Image,
    width: int,
    height: int,
    background=(235, 235, 235)
):
    """
    Fit an image into a canvas without stretching it.
    """

    image = image.convert("RGB")

    src_w, src_h = image.size

    scale = min(
        width / src_w,
        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",
        (width, height),
        background
    )

    x = (width - new_w) // 2
    y = (height - new_h) // 2

    canvas.paste(
        resized,
        (x, y)
    )

    return canvas


# ============================================================
# SEGMENTATION
# ============================================================

def get_segmentation_pipeline():

    processor = AutoImageProcessor.from_pretrained(
        "openmmlab/upernet-convnext-small"
    )

    model = UperNetForSemanticSegmentation.from_pretrained(
        "openmmlab/upernet-convnext-small"
    )

    return processor, model


@torch.inference_mode()
@spaces.GPU
def segment_image(
    image: Image.Image,
    image_processor,
    image_segmentor
):
    """
    ADE20K semantic segmentation.
    """

    inputs = image_processor(
        image,
        return_tensors="pt"
    ).pixel_values.to(device)

    with torch.no_grad():

        outputs = image_segmentor(
            inputs
        )

    segmentation = (
        image_processor
        .post_process_semantic_segmentation(
            outputs,
            target_sizes=[image.size[::-1]]
        )[0]
    )

    # IMPORTANT:
    # CUDA -> CPU before NumPy.
    segmentation = (
        segmentation
        .detach()
        .cpu()
        .numpy()
    )

    color_seg = np.zeros(
        (
            segmentation.shape[0],
            segmentation.shape[1],
            3
        ),
        dtype=np.uint8
    )

    palette = np.array(
        ade_palette()
    )

    for label, color in enumerate(palette):

        color_seg[
            segmentation == label,
            :
        ] = color

    return Image.fromarray(
        color_seg
    ).convert("RGB")


# ============================================================
# DEPTH
# ============================================================

def get_depth_pipeline():

    feature_extractor = AutoImageProcessor.from_pretrained(
        "LiheYoung/depth-anything-large-hf",
        torch_dtype=dtype
    )

    depth_estimator = AutoModelForDepthEstimation.from_pretrained(
        "LiheYoung/depth-anything-large-hf",
        torch_dtype=dtype
    )

    return (
        feature_extractor,
        depth_estimator
    )


@torch.inference_mode()
@spaces.GPU
def get_depth_image(
    image: Image.Image,
    feature_extractor,
    depth_estimator
):

    inputs = feature_extractor(
        images=image,
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():

        depth = (
            depth_estimator(
                **inputs
            )
            .predicted_depth
        )

    width, height = image.size

    depth = torch.nn.functional.interpolate(
        depth.unsqueeze(1).float(),
        size=(height, width),
        mode="bicubic",
        align_corners=False,
    )

    depth_min = torch.amin(
        depth,
        dim=[1, 2, 3],
        keepdim=True
    )

    depth_max = torch.amax(
        depth,
        dim=[1, 2, 3],
        keepdim=True
    )

    depth = (
        depth - depth_min
    ) / (
        depth_max - depth_min + 1e-8
    )

    depth = torch.cat(
        [depth] * 3,
        dim=1
    )

    depth = (
        depth
        .permute(0, 2, 3, 1)
        .cpu()
        .numpy()[0]
    )

    return Image.fromarray(
        (
            depth * 255
        )
        .clip(0, 255)
        .astype(np.uint8)
    )


# ============================================================
# ADE20K OBJECT CLASSES
# ============================================================

# 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",
    "bed",
    "cabinet",
    "shelf",
    "desk",
    "dresser",
    "bench",
    "armchair",
    "ottoman",
    "stool",
    "lamp",
    "pillow",
    "plant",
    "vase",
    "painting",
    "mirror",
    "television",
    "computer",
    "monitor",
    "clock",
    "rug",
    "book",
    "books",
    "chest of drawers",
    "counter",
}


# ============================================================
# SEGMENTATION -> BINARY FURNITURE MASK
# ============================================================

def get_class_mask(
    segmentation: Image.Image,
    wanted_classes
):
    """
    Convert ADE20K RGB segmentation into a binary mask.
    """

    seg = np.array(
        segmentation
    )

    unique_colors = np.unique(
        seg.reshape(
            -1,
            3
        ),
        axis=0
    )

    mask = np.zeros(
        seg.shape[:2],
        dtype=np.uint8
    )

    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"
    )


# ============================================================
# 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
    )

    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


# ============================================================
# REFERENCE OBJECT EXTRACTION
# ============================================================

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"]

        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 a visual board containing:

    1. The complete reference room, slightly reduced.
    2. Individual furniture/decor crops.

    This gives IP-Adapter both scene context and object detail.
    """

    reference_image = reference_image.convert(
        "RGB"
    )

    # --------------------------------------------------------
    # Extract individual objects
    # --------------------------------------------------------

    objects = extract_reference_objects(
        reference_image,
        reference_segmentation
    )

    # --------------------------------------------------------
    # Board
    # --------------------------------------------------------

    board = Image.new(
        "RGB",
        (
            board_size,
            board_size
        ),
        (245, 245, 245)
    )

    # --------------------------------------------------------
    # Main reference image
    # --------------------------------------------------------

    main_size = int(
        board_size * 0.52
    )

    main_image = fit_image(
        reference_image,
        main_size,
        main_size,
        background=(235, 235, 235)
    )

    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)
        )

    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
            )
        )
    )

    cell_w = max(
        64,
        thumb_area_x // columns
    )

    cell_h = max(
        64,
        thumb_area_y // min(rows, 4)
    )

    start_x = (
        main_size
        + 20
    )

    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
    )

    # Protect a narrow outer border.
    # This reduces accidental changes at image boundaries.
    arr = np.array(
        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"
    )


# ============================================================
# MODEL
# ============================================================

class ControlNetDepthDesignModelMulti:

    def __init__(self):

        self.seed = 323 * 111

        self.neg_prompt = (
            "deformed furniture, distorted furniture, "
            "warped furniture, melted furniture, "
            "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, "
            "new windows, new doors, "
            "text, logo, watermark, "
            "low resolution, blurry, out of focus, "
            "surreal, cartoon, illustration, CGI, "
            "plastic materials, oversaturated"
        )

        self.additional_quality_suffix = (
            "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,
        target_room: Image.Image,
        reference_room: Image.Image,
        prompt: str,
        guidance_scale: float = 7.0,
        num_steps: int = 35,
        strength: float = 0.72,
        img_size: int = 640,
        ip_scale: float = 0.95
    ):

        print(
            "=========================================="
        )

        print(
            "Starting reference-room transfer..."
        )

        print(
            "=========================================="
        )

        flush()

        generator = torch.Generator(
            device=device
        ).manual_seed(
            int(self.seed)
        )

        # ----------------------------------------------------
        # PROMPT
        # ----------------------------------------------------

        if prompt is None:
            prompt = ""

        prompt = prompt.strip()

        if prompt == "":

            prompt = (
                "A professionally photographed interior "
                "of the target room furnished with the "
                "complete furniture and decoration ensemble "
                "from the reference 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 = (
            prompt
            + ", "
            + self.additional_quality_suffix
        )

        # ----------------------------------------------------
        # TARGET SIZE
        # ----------------------------------------------------

        original_width, original_height = (
            target_room.size
        )

        target_width, target_height = (
            resize_dimensions(
                target_room.size,
                img_size
            )
        )

        target = (
            target_room
            .convert("RGB")
            .resize(
                (
                    target_width,
                    target_height
                ),
                Image.Resampling.LANCZOS
            )
        )

        # ----------------------------------------------------
        # REFERENCE
        # ----------------------------------------------------

        reference = (
            reference_room
            .convert("RGB")
        )

        reference = fit_image(
            reference,
            target_width,
            target_height,
            background=(235, 235, 235)
        )

        # ----------------------------------------------------
        # TARGET SEGMENTATION
        # ----------------------------------------------------

        print(
            "1/5 - Target room segmentation..."
        )

        target_segmentation = segment_image(
            target,
            seg_image_processor,
            image_segmentor
        )

        # ----------------------------------------------------
        # REFERENCE SEGMENTATION
        # ----------------------------------------------------

        print(
            "2/5 - Reference room segmentation..."
        )

        reference_segmentation = segment_image(
            reference,
            seg_image_processor,
            image_segmentor
        )

        # ----------------------------------------------------
        # BUILD OBJECT-AWARE REFERENCE BOARD
        # ----------------------------------------------------

        print(
            "3/5 - Extracting furniture and decorations..."
        )

        reference_board = build_reference_board(
            reference,
            reference_segmentation,
            board_size=max(
                512,
                target_width
            )
        )

        reference_board = reference_board.resize(
            (
                target_width,
                target_height
            ),
            Image.Resampling.LANCZOS
        )

        # ----------------------------------------------------
        # DEPTH
        # ----------------------------------------------------

        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)
        )

        # ----------------------------------------------------
        # GENERATION
        # ----------------------------------------------------

        print(
            "5/5 - Generating furnished target room..."
        )

        result = pipe(
            prompt=pos_prompt,

            negative_prompt=self.neg_prompt,

            num_inference_steps=int(
                num_steps
            ),

            strength=float(
                strength
            ),

            guidance_scale=float(
                guidance_scale
            ),

            generator=generator,

            image=target,

            mask_image=mask_image,

            # IMPORTANT:
            #
            # This is the reference board, not a manually
            # selected furniture image.
            ip_adapter_image=reference_board,

            control_image=[
                target_depth,
                target_segmentation
            ],

            # Depth is slightly stronger because it carries
            # the target room's spatial structure.
            controlnet_conditioning_scale=[
                0.80,
                0.65
            ]
        ).images[0]

        flush()

        result = result.resize(
            (
                original_width,
                original_height
            ),
            Image.Resampling.LANCZOS
        )

        print(
            "Generation completed."
        )

        return result


# ============================================================
# GRADIO
# ============================================================

def create_demo(model):

    gr.Markdown(
        "### StableDesign — Reference Room Transfer"
    )

    gr.Markdown(
        """
        <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():

            target_image = gr.Image(
                label="1. Target / empty room",
                type="pil",
                elem_id="img-display-input"
            )

            reference_image = gr.Image(
                label="2. Fully furnished reference room",
                type="pil",
                elem_id="img-display-input"
            )

            prompt = gr.Textbox(
                label="Instructions",
                value=(
                    "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
            )

            with gr.Accordion(
                "Advanced options",
                open=False
            ):

                num_steps = gr.Slider(
                    label="Steps",
                    minimum=20,
                    maximum=50,
                    value=35,
                    step=1
                )

                img_size = gr.Slider(
                    label="Image size",
                    minimum=384,
                    maximum=768,
                    value=640,
                    step=64
                )

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=4.0,
                    maximum=12.0,
                    value=7.0,
                    step=0.5
                )

                strength = gr.Slider(
                    label="Furniture generation strength",
                    minimum=0.45,
                    maximum=0.90,
                    value=0.80,
                    step=0.05
                )

                ip_scale = gr.Slider(
                    label="Reference room strength",
                    minimum=0.50,
                    maximum=1.20,
                    value=1.10,
                    step=0.05
                )

                seed = gr.Slider(
                    label="Seed",
                    minimum=-1,
                    maximum=2147483647,
                    value=323 * 111,
                    step=1,
                    randomize=True
                )

                a_prompt = gr.Textbox(
                    label="Added Prompt",
                    value=(
                        "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(
                    label="Negative Prompt",
                    value=(
                        "deformed furniture, distorted furniture, "
                        "warped furniture, melted furniture, "
                        "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, 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 complete furniture",
                variant="primary"
            )

        # ====================================================
        # OUTPUT
        # ====================================================

        with gr.Column():

            result_image = gr.Image(
                label="Result",
                elem_id="img-display-output"
            )

    # ========================================================
    # SUBMIT
    # ========================================================

    def on_submit(
        target,
        reference,
        text,
        num_steps_value,
        guidance_value,
        seed_value,
        strength_value,
        ip_value,
        additional_prompt,
        negative_prompt,
        image_size
    ):

        if target is None:

            raise gr.Error(
                "Please upload image 1: the target room."
            )

        if reference is None:

            raise gr.Error(
                "Please upload image 2: the fully furnished "
                "reference room."
            )

        model.seed = int(
            seed_value
        )

        model.neg_prompt = (
            negative_prompt
        )

        model.additional_quality_suffix = (
            additional_prompt
        )

        with torch.no_grad():

            output = model.generate_design(
                target_room=target,

                reference_room=reference,

                prompt=text,

                guidance_scale=float(
                    guidance_value
                ),

                num_steps=int(
                    num_steps_value
                ),

                strength=float(
                    strength_value
                ),

                img_size=int(
                    image_size
                ),

                ip_scale=float(
                    ip_value
                )
            )

        return output

    submit.click(
        on_submit,

        inputs=[
            target_image,
            reference_image,
            prompt,
            num_steps,
            guidance_scale,
            seed,
            strength,
            ip_scale,
            a_prompt,
            n_prompt,
            img_size
        ],

        outputs=result_image
    )


# ============================================================
# LOAD CONTROLNET
# ============================================================

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,
    use_safetensors=True
)


# ============================================================
# 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",

        controlnet=[
            controlnet_depth,
            controlnet_seg
        ],

        image_encoder=image_encoder,

        safety_checker=None,

        torch_dtype=dtype
    )
)


# ============================================================
# IP-ADAPTER PLUS
# ============================================================

print(
    "Loading IP-Adapter Plus SD1.5..."
)

pipe.load_ip_adapter(
    "h94/IP-Adapter",
    subfolder="models",
    weight_name="ip-adapter-plus_sd15.safetensors"
)

pipe.set_ip_adapter_scale(
    0.95
)


# ============================================================
# SCHEDULER
# ============================================================

pipe.scheduler = (
    UniPCMultistepScheduler.from_config(
        pipe.scheduler.config
    )
)


pipe = pipe.to(
    device
)


# ============================================================
# SEGMENTATION
# ============================================================

print(
    "Loading ADE20K segmentation..."
)

(
    seg_image_processor,
    image_segmentor
) = get_segmentation_pipeline()

image_segmentor = (
    image_segmentor.to(device)
)


# ============================================================
# DEPTH
# ============================================================

print(
    "Loading Depth Anything..."
)

(
    depth_feature_extractor,
    depth_estimator
) = get_depth_pipeline()

depth_estimator = (
    depth_estimator.to(device)
)


# ============================================================
# MAIN
# ============================================================

def main():

    model = (
        ControlNetDepthDesignModelMulti()
    )

    print(
        "=========================================="
    )

    print(
        "StableDesign V2 loaded successfully."
    )

    print(
        "Reference-room furniture transfer enabled."
    )

    print(
        "=========================================="
    )

    with gr.Blocks(
        css=css
    ) as demo:

        gr.Markdown(
            "# StableDesign V2"
        )

        gr.Markdown(
            """
            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(
            model
        )

    demo.queue().launch(
        share=False
    )


if __name__ == "__main__":

    main()