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(
"""
Image 1 — Target room:
the room that will be furnished. Its architecture,
perspective, walls, floor, windows and doors are
preserved using structural controls.
Image 2 — Reference room:
a completely furnished room. The system automatically
analyzes the furniture and decorative elements and
uses them as the visual reference.
No furniture area or placement needs to be selected.
"""
)
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()