nodes/RodrigoSKohl-InteriorDesign-for-ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: update to version 2.0.0 with new features

Browse Source
This commit is contained in:
rodrigoskohl
2025-07-19 22:27:48 -03:00
parent 9443c413e1
commit 484cb1dbbb
10 changed files with 3374 additions and 597 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node/node.py
+117 -64
View File
@@ -1,84 +1,137 @@
from typing import Literal
import torch
import torch.nn.functional as F
from comfy.model_management import get_torch_device
from .pipeline import ControlNetDepthDesignModelMulti # Seu pipeline original
from torchvision import transforms
from torchvision.transforms import ToPILImage
from .pipeline import SimplifiedSegmenter
import torch.nn.functional as F
class InteriorDesignNode:
def __init__(self):
self.device = get_torch_device()
self.dtype = torch.float16 if self.device == "cuda" else torch.float32
self.pipeline = ControlNetDepthDesignModelMulti() # Seu pipeline original
self.pipeline = SimplifiedSegmenter(
device=self.device
)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"positive": ("STRING",{"forceInput": True}), # Prompt de descrição do design
"guidance_scale": ("FLOAT", {"default": 10, "min": 1.0, "max": 20.0, "step": 0.1}),
"num_steps": ("INT", {"default": 50, "min": 1, "max": 100}),
"strength": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 1.0}),
"seed": ("INT", {"default": 40, "min": 0, "max": 0xffffffffffffffff}),
"img_size": ("INT", {"default": 768, "min": 256, "max": 768, "step": 8}),
"img_dept_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 0.99, "step": 0.01}),
"img_seg_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 0.99, "step": 0.01}),
"ip_adapter_scale": ("FLOAT", {"default": 0.40, "min": 0.01, "max": 1.0, "step": 0.01}),
"ip_adapter_model":(["ip-adapter_sd15.bin","ip-adapter-plus_sd15.bin","ip-adapter_sd15_light.bin"], {})
"control_items": ("CONTROL_ITEMS",),
},
"optional": {
"negative": ("STRING",{"forceInput": True}), # Prompt de descrição do design
}
}
RETURN_TYPES = ("IMAGE",) # Retorna a imagem gerada
RETURN_NAMES = ("IMAGE",) # Nome do retorno
RETURN_TYPES = ("IMAGE", "MASK")
RETURN_NAMES = ("IMAGE", "MASK")
FUNCTION = "process"
CATEGORY = "segmentation"
FUNCTION = "process" # Função que será executada
CATEGORY = "interior_design" # Categoria para o ComfyUI
def process(self, image, control_items):
image = image.permute(0, 3, 1, 2).to(self.device)
def process(self, image, positive, guidance_scale, num_steps, strength, seed, img_size, img_dept_threshold, img_seg_threshold, ip_adapter_scale, ip_adapter_model, negative=None):
# Pré-processamento da imagem
self.pipeline.ip_adapter_scale = ip_adapter_scale
self.pipeline.ip_adapter_model = ip_adapter_model
self.pipeline.seed = seed
if negative is not None:
self.pipeline.neg_prompt = negative
image = tensor_to_pil(image)
# Chama o pipeline para gerar o design de interiores
design_image = self.pipeline.generate_design(
empty_room_image=image,
prompt=positive,
guidance_scale=guidance_scale,
num_steps=num_steps,
strength=strength,
img_size=img_size,
img_dept_threshold=img_dept_threshold,
img_seg_threshold=img_seg_threshold,
seg_cond, mask = self.pipeline(images=image,control_items=control_items)
seg_cond = seg_cond.permute(0, 2, 3, 1)
mask = mask.permute(0, 2, 3, 1)
mask = mask.squeeze(-1)
#seg_cond = seg_cond.float() / 255.0 (its already normalized in the pipeline)
return (seg_cond, mask,)
class ImageNormalize:
def __init__(self):
self.device = get_torch_device()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"target_size": ("INT", {
"default": 512,
"min": 2,
"max": 4096,
"step": 1,
"display": "number",
}),
"multiple": ("INT", {
"default": 8,
"min": 2,
"max": 128,
"step": 1,
"display": "number",
}),
"mode": (["bilinear", "nearest", "bicubic", "trilinear","area","nearest-exact"],),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("IMAGE",)
FUNCTION = "resize_tensor_to_target"
CATEGORY = "image"
def resize_tensor_to_target(
self,
images: torch.Tensor,
target_size: int,
multiple: int,
mode: Literal[
"nearest", "linear", "bilinear", "bicubic", "trilinear", "area", "nearest-exact"
] = "bilinear"
) -> torch.Tensor:
if images.dim() == 4 and images.shape[-1] == 3:
images = images.permute(0, 3, 1, 2)
b, c, h, w = images.shape
max_dim = max(h, w)
scale = target_size / max_dim if max_dim > target_size else 1.0
new_h = int(h * scale)
new_w = int(w * scale)
new_h = ((new_h + multiple - 1) // multiple) * multiple
new_w = ((new_w + multiple - 1) // multiple) * multiple
resized = F.interpolate(
images, size=(new_h, new_w), mode=mode,
align_corners=False if mode in ['linear', 'bilinear', 'bicubic', 'trilinear'] else None
)
design_image = pil_to_tensor(design_image, self.device, self.dtype)
return (design_image,)
def tensor_to_pil(image_tensor):
print(f"Tensor: {image_tensor.shape},{image_tensor.dtype}")
if image_tensor.ndimension() == 4:
image_tensor = image_tensor.squeeze(0) # Remove a dimensão do batch
image_tensor = image_tensor.permute(2, 0, 1) # Troca de [altura, largura, canais] -> [canais, altura, largura]
if image_tensor.shape[0] == 4:
image_tensor = image_tensor[:3]
to_pil = ToPILImage()
pil_image = to_pil(image_tensor)
print(f"Imagem PIL: {pil_image.size},{pil_image.mode}")
return pil_image
resized = resized.permute(0, 2, 3, 1)
def pil_to_tensor(image_pil, device, dtype):
print(f"Imagem PIL: {image_pil.size},{image_pil.mode}")
to_tensor = transforms.ToTensor()
tensor = to_tensor(image_pil)
tensor = tensor.permute(1, 2, 0) # Troca de [canais, altura, largura] -> [altura, largura, canais]
tensor = tensor.unsqueeze(0) # Adiciona a dimensão do batch
tensor = tensor.to(device, dtype=dtype)
print(f"Tensor: {tensor.shape},{tensor.dtype}")
return tensor
return (resized,)
class ControlItems:
def __init__(self):
self.device = get_torch_device()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"window": ("BOOLEAN", {"default": True}),
"door": ("BOOLEAN", {"default": True}),
"staircase": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = ("CONTROL_ITEMS",)
RETURN_NAMES = ("CONTROL_ITEMS",)
FUNCTION = "process"
CATEGORY = "control"
def process(self, window: bool, door: bool, staircase: bool):
control_items = []
if window:
control_items.append("windowpane;window")
if door:
control_items.append("door;double;door")
if staircase:
control_items.append("stairway;staircase")
return (control_items,)
node/pipeline.py
+99 -292
View File
@@ -1,319 +1,126 @@
from typing import Tuple, Union, List
from typing import Tuple,List, Union
import numpy as np
from PIL import Image
import torch
from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline
from diffusers import ControlNetModel, UniPCMultistepScheduler, AutoPipelineForText2Image
from transformers import AutoImageProcessor, UperNetForSemanticSegmentation, AutoModelForDepthEstimation
from .colors import ade_palette
from .utils import map_colors_rgb
from diffusers import StableDiffusionXLPipeline
import gc
import os
from huggingface_hub import hf_hub_download
import folder_paths
import torch
from torch import Tensor
from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
# -------------------
# Classe principal
# -------------------
repo_id = "MykolaL/StableDesign"
file_name = "diffusion_pytorch_model.safetensors"
models = ["controlnet_depth", "own_controlnet"]
rootfolder = folder_paths.folder_names_and_paths["controlnet"][0][0]
class SimplifiedSegmenter:
def __init__(self, device: str = "cuda"):
self.device = device
def check_and_download(model_name,rootfolder):
hf_hub_download(
repo_id=repo_id,
filename=file_name,
subfolder= model_name,
local_dir=rootfolder,
repo_type="space"
)
hf_hub_download(
repo_id=repo_id,
filename="config.json",
subfolder= model_name,
local_dir=rootfolder,
repo_type="space"
)
for model_name in models:
model_path = os.path.join(rootfolder, model_name)
check_and_download(model_name,rootfolder)
device = "cuda"
dtype = torch.float16
def filter_items(
colors_list: Union[List, np.ndarray],
items_list: Union[List, np.ndarray],
items_to_remove: Union[List, np.ndarray]
) -> Tuple[Union[List, np.ndarray], Union[List, np.ndarray]]:
"""
Filters items and their corresponding colors from given lists, excluding
specified items.
Args:
colors_list: A list or numpy array of colors corresponding to items.
items_list: A list or numpy array of items.
items_to_remove: A list or numpy array of items to be removed.
Returns:
A tuple of two lists or numpy arrays: filtered colors and filtered
items.
"""
filtered_colors = []
filtered_items = []
for color, item in zip(colors_list, items_list):
if item not in items_to_remove:
filtered_colors.append(color)
filtered_items.append(item)
return filtered_colors, filtered_items
def get_segmentation_pipeline(
) -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
"""Method to load the segmentation pipeline
Returns:
Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]: segmentation pipeline
"""
image_processor = AutoImageProcessor.from_pretrained(
"openmmlab/upernet-convnext-small"
)
image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
"openmmlab/upernet-convnext-small"
)
return image_processor, image_segmentor
def segment_image(
image: Image,
image_processor: AutoImageProcessor,
image_segmentor: UperNetForSemanticSegmentation
) -> Image:
"""
Segments an image using a semantic segmentation model.
Args:
image (Image): The input image to be segmented.
image_processor (AutoImageProcessor): The processor to prepare the
image for segmentation.
image_segmentor (UperNetForSemanticSegmentation): The semantic
segmentation model used to identify different segments in the image.
Returns:
Image: The segmented image with each segment colored differently based
on its identified class.
"""
# image_processor, image_segmentor = get_segmentation_pipeline()
pixel_values = image_processor(image, return_tensors="pt").pixel_values
with torch.no_grad():
outputs = image_segmentor(pixel_values)
seg = image_processor.post_process_semantic_segmentation(
outputs, target_sizes=[image.size[::-1]])[0]
color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
palette = np.array(ade_palette())
for label, color in enumerate(palette):
color_seg[seg == label, :] = color
color_seg = color_seg.astype(np.uint8)
seg_image = Image.fromarray(color_seg).convert('RGB')
return seg_image
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
def get_depth_image(
image: Image,
feature_extractor: AutoImageProcessor,
depth_estimator: AutoModelForDepthEstimation
) -> Image:
image_to_depth = feature_extractor(images=image, return_tensors="pt").to(device)
with torch.no_grad():
depth_map = depth_estimator(**image_to_depth).predicted_depth
width, height = image.size
depth_map = torch.nn.functional.interpolate(
depth_map.unsqueeze(1).float(),
size=(height, width),
mode="bicubic",
align_corners=False,
)
depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
depth_map = (depth_map - depth_min) / (depth_max - depth_min)
image = torch.cat([depth_map] * 3, dim=1)
image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
return image
def resize_dimensions(dimensions, target_size):
"""
Resize PIL to target size while maintaining aspect ratio
If smaller than target size leave it as is
"""
width, height = dimensions
# Check if both dimensions are smaller than the target size
if width < target_size and height < target_size:
return dimensions
# Determine the larger side
if width > height:
# Calculate the aspect ratio
aspect_ratio = height / width
# Resize dimensions
return (target_size, int(target_size * aspect_ratio))
else:
# Calculate the aspect ratio
aspect_ratio = width / height
# Resize dimensions
return (int(target_size * aspect_ratio), target_size)
def flush():
gc.collect()
torch.cuda.empty_cache()
class ControlNetDepthDesignModelMulti:
""" Produces random noise images """
def __init__(self):
""" Initialize your model(s) here """
#os.environ['HF_HUB_OFFLINE'] = "True"
self.ip_adapter_scale = 0.4
self.ip_adapter_model = "ip-adapter_sd15.bin"
self.seed = 323*111
self.neg_prompt = "window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner"
self.control_items = ["windowpane;window", "door;double;door"]
self.additional_quality_suffix = "interior design, 4K, high resolution, photorealistic"
self.controlnet_depth = ControlNetModel.from_pretrained(
rootfolder,subfolder="controlnet_depth",
torch_dtype=dtype,
use_safetensors=True
print("🔄 Loading segmentator...")
self.image_processor = AutoImageProcessor.from_pretrained(
"openmmlab/upernet-convnext-small"
)
self.controlnet_seg = ControlNetModel.from_pretrained(
rootfolder,subfolder="own_controlnet",
torch_dtype=dtype,
use_safetensors=True
self.image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
"openmmlab/upernet-convnext-small"
)
self.image_segmentor.eval()
self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
"SG161222/Realistic_Vision_V5.1_noVAE",
controlnet=[self.controlnet_depth, self.controlnet_seg],
safety_checker=None,
torch_dtype=dtype,
)
def tensor_to_pil(self, tensor: Tensor) -> Image.Image:
if tensor.max() > 1.0:
arr = tensor.cpu().permute(1, 2, 0).numpy().astype(np.uint8)
else:
arr = (tensor.cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8)
return Image.fromarray(arr)
self.pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models",
weight_name=self.ip_adapter_model,)
self.pipe.set_ip_adapter_scale(self.ip_adapter_scale)
self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
self.pipe = self.pipe.to(device)
self.guide_pipe = StableDiffusionXLPipeline.from_pretrained(
"segmind/SSD-1B",
torch_dtype=dtype,
use_safetensors=True,
variant="fp16"
)
self.guide_pipe = self.guide_pipe.to(device)
self.seg_image_processor, self.image_segmentor = get_segmentation_pipeline()
self.depth_feature_extractor, self.depth_estimator = get_depth_pipeline()
self.depth_estimator = self.depth_estimator.to(device)
def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9, img_size: int = 640, img_dept_threshold=0.5, img_seg_threshold=0.5 ) -> Image:
def segment_image(self,
image: Image,
) -> Image:
"""
Given an image of an empty room and a prompt
generate the designed room according to the prompt
Inputs -
empty_room_image - An RGB PIL Image of the empty room
prompt - Text describing the target design elements of the room
Returns -
design_image - PIL Image of the same size as the empty room image
If the size is not the same the submission will fail.
Segments an image using a semantic segmentation model.
Args:
image (Image): The input image to be segmented.
image_processor (AutoImageProcessor): The processor to prepare the
image for segmentation.
image_segmentor (UperNetForSemanticSegmentation): The semantic
segmentation model used to identify different segments in the image.
Returns:
Image: The segmented image with each segment colored differently based
on its identified class.
"""
print(f"Positive Prompt: {prompt}")
print(f"Negative Prompt: {self.neg_prompt}")
print(f"Seed: {self.seed}")
print(f"IP Adapter Scale: {self.ip_adapter_scale}")
# image_processor, image_segmentor = get_segmentation_pipeline()
pixel_values = self.image_processor(image, return_tensors="pt").pixel_values
image_size = image.size[::-1]
print(f"Image size: {image_size}, pixel_values shape: {pixel_values.shape}")
with torch.no_grad():
outputs = self.image_segmentor(pixel_values)
seg = self.image_processor.post_process_semantic_segmentation(
outputs, target_sizes=[image_size])[0]
color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
palette = np.array(ade_palette())
for label, color in enumerate(palette):
color_seg[seg == label, :] = color
color_seg = color_seg.astype(np.uint8)
seg_image = Image.fromarray(color_seg).convert('RGB')
return seg_image
@staticmethod
def filter_items(
colors_list: Union[List, np.ndarray],
items_list: Union[List, np.ndarray],
items_to_remove: Union[List, np.ndarray]
) -> Tuple[Union[List, np.ndarray], Union[List, np.ndarray]]:
"""
Filters items and their corresponding colors from given lists, excluding
specified items.
Args:
colors_list: A list or numpy array of colors corresponding to items.
items_list: A list or numpy array of items.
items_to_remove: A list or numpy array of items to be removed.
Returns:
A tuple of two lists or numpy arrays: filtered colors and filtered
items.
"""
filtered_colors = []
filtered_items = []
for color, item in zip(colors_list, items_list):
if item not in items_to_remove:
filtered_colors.append(color)
filtered_items.append(item)
return filtered_colors, filtered_items
flush()
self.generator = torch.Generator(device=device).manual_seed(self.seed)
pos_prompt = prompt + f', {self.additional_quality_suffix}'
orig_w, orig_h = empty_room_image.size
new_width, new_height = resize_dimensions(empty_room_image.size, img_size)
input_image = empty_room_image.resize((new_width, new_height))
real_seg = np.array(segment_image(input_image,
self.seg_image_processor,
self.image_segmentor))
def __call__(self, images: Tensor, control_items: List) -> Tuple[Tensor, Tensor]:
"""
images: Tensor (B,C,H,W) ou (C,H,W)
Returns:
seg_cond_tensor: (B,3,H,W) uint8
mask_tensor: (B,1,H,W) float 0/1
"""
print(f'Control items: {control_items}')
if images.dim() == 4:
input_image = self.tensor_to_pil(images[0])
else:
input_image = self.tensor_to_pil(images)
real_seg = np.array(self.segment_image(input_image))
unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
unique_colors = [tuple(color) for color in unique_colors]
segment_items = [map_colors_rgb(i) for i in unique_colors]
chosen_colors, segment_items = filter_items(
chosen_colors, segment_items = self.filter_items(
colors_list=unique_colors,
items_list=segment_items,
items_to_remove=self.control_items
items_to_remove=control_items
)
mask = np.zeros_like(real_seg)
mask = np.zeros((real_seg.shape[0], real_seg.shape[1]), dtype=np.uint8)
for color in chosen_colors:
color_matches = (real_seg == color).all(axis=2)
mask[color_matches] = 1
image_np = np.array(input_image)
image = Image.fromarray(image_np).convert("RGB")
mask_image = Image.fromarray((mask * 255).astype(np.uint8)).convert("RGB")
segmentation_cond_image = Image.fromarray(real_seg).convert("RGB")
image_depth = get_depth_image(image, self.depth_feature_extractor, self.depth_estimator)
# generate image that would be used as IP-adapter
flush()
new_width_ip = int(new_width / 8) * 8
new_height_ip = int(new_height / 8) * 8
ip_image = self.guide_pipe(pos_prompt,
num_inference_steps=num_steps,
negative_prompt=self.neg_prompt,
height=new_height_ip,
width=new_width_ip,
generator=[self.generator]).images[0]
flush()
generated_image = self.pipe(
prompt=pos_prompt,
negative_prompt=self.neg_prompt,
num_inference_steps=num_steps,
strength=strength,
guidance_scale=guidance_scale,
generator=[self.generator],
image=image,
mask_image=mask_image,
ip_adapter_image=ip_image,
control_image=[image_depth, segmentation_cond_image],
controlnet_conditioning_scale=[img_dept_threshold, img_seg_threshold]
).images[0]
flush()
design_image = generated_image.resize(
(orig_w, orig_h), Image.Resampling.LANCZOS
)
return design_image
seg_cond_image = Image.fromarray(real_seg).convert("RGB")
seg_cond_tensor = torch.from_numpy(np.array(seg_cond_image)).permute(2,0,1).unsqueeze(0).to(self.device).float() / 255.0 # shape (1,3,H,W)
mask_tensor = torch.from_numpy(mask).unsqueeze(0).unsqueeze(0).float() # shape (1,1,H,W)
mask_tensor = mask_tensor.to('cpu') # move to cpu (avoid bug when use the mask on IP Adapter?)
print(f'Shape of seg_cond_tensor: {seg_cond_tensor.shape}, mask_tensor: {mask_tensor.shape}')
return seg_cond_tensor, mask_tensor
node/utils.py
-1
View File
@@ -1,5 +1,4 @@
import gc
import numpy as np
from PIL import Image
import torch