feat: update to version 2.0.0 with new features
This commit is contained in:
rodrigoskohl
@@ -1,14 +1,82 @@
|
||||
<div align="center">
|
||||
|
||||
# Interior Design for Comfyui
|
||||
# 🛋️ Interior Design for ComfyUI v2.0
|
||||
|
||||
</div>
|
||||
|
||||
This node adapts [StableDesign: 2nd place solution for the Generative Interior Design 2024 competition](https://huggingface.co/spaces/MykolaL/StableDesign) to use directly on ComfyUI.
|
||||
This project adapts the original pipeline from
|
||||
[StableDesign: 2nd place solution for the Generative Interior Design 2024 competition](https://huggingface.co/spaces/MykolaL/StableDesign)
|
||||
to work **natively and fully inside ComfyUI**, leveraging modern custom nodes and a modular architecture.
|
||||
|
||||
More about this pipeline works, see this post on [Medium](https://medium.com/@melgor89/generative-interior-design-challenge-2024-2nd-place-solution-6338f19f6fe3)
|
||||
#
|
||||
The pipeline automatically downloads trained models of `controlnet_depth` and `own_controlnet` to the local script of ComfyUI runs in folder `./custom_models`. Other models are downloaded to the `.cache` folder of Hugging Face.
|
||||
Learn more about the original method in this [Medium article](https://medium.com/@melgor89/generative-interior-design-challenge-2024-2nd-place-solution-6338f19f6fe3).
|
||||
|
||||
This node is available on [ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager)
|
||||
|
||||
---
|
||||
|
||||
## ✅ Key Features
|
||||
|
||||
- **Fully restructured for ComfyUI compatibility**
|
||||
- Includes 3 custom nodes developed specifically for this repository:
|
||||
- `"Interior Design Segmentator"`
|
||||
- `"Image Normalize"`
|
||||
- `"Control Items"`
|
||||
- Uses additional advanced custom nodes:
|
||||
- [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2)
|
||||
- [ComfyUI_IPAdapter_plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus)
|
||||
- Reimplements only the **necessary components** from the original pipeline, such as:
|
||||
- **Semantic segmenter** using [`openmmlab/upernet-convnext-small`](https://huggingface.co/openmmlab/upernet-convnext-small)
|
||||
- Supports standard ComfyUI workflows with **ControlNet and IP Adapter integration**
|
||||
- Compatible with other workflows and custom models
|
||||
- Inference is significantly faster than the original pipeline (test the workflow and compare).
|
||||
|
||||
---
|
||||
|
||||
## 🧱 Pipeline Overview
|
||||
|
||||
- The idea was to use as many nodes from the ComfyUI core as possible. However, for certain tasks, I had to rely on custom nodes or adapt my own nodes to stay as faithful as possible to the original pipeline, or alternatively, use dependencies from other custom node repositories.
|
||||
- Uses the **same base models** as the original StableDesign pipeline
|
||||
- Includes **custom ControlNet models**:
|
||||
- `controlnet_depth`
|
||||
- `own_controlnet` (segmentation)
|
||||
- These ControlNets, along with the depth map and segmentation map, are passed into the **UNet** of the diffusion model
|
||||
- Final step includes **refinement with latent noise injection and up/down-scaling**(because of the size limitation of the original pipeline)
|
||||
|
||||
---
|
||||
|
||||
## 📁 Example Workflow
|
||||
|
||||
An example workflow is included in this repository:
|
||||
[`stable-design-for-comfyui.json`](./workflow/stable-design-for-comfyui.json)
|
||||
|
||||
You can load it directly into ComfyUI to test or extend the pipeline.
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 📦 Installation & Dependencies
|
||||
|
||||
This node is available on [ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager)
|
||||
You must also install the following custom nodes manually if not yet available:
|
||||
|
||||
- [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2)
|
||||
- [ComfyUI_IPAdapter_plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus) *(note: unmaintained)*
|
||||
- [ComfyMath](https://github.com/evanspearman/ComfyMath)
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Upcoming Features
|
||||
|
||||
- **Expansion of Control Items**: Add support for more semantic regions and control inputs to increase flexibility and detail in generation.
|
||||
- **Evaluation of CLIP Conditioning**: Reassess the combination of positive and negative CLIP embeddings in ControlNet to verify if results match the original pipeline's behavior.
|
||||
- **Custom IP Adapter Node**: Implement a native IP Adapter node to replace the current dependency on `ComfyUI_IPAdapter_plus`, which is currently unmaintained.
|
||||
|
||||
---
|
||||
|
||||
## 🧪 License & Acknowledgements
|
||||
|
||||
Original pipeline by MykolaL ([StableDesign on Hugging Face](https://huggingface.co/spaces/MykolaL/StableDesign))
|
||||
This adaptation reuses only essential logic, converted to full ComfyUI-compatible workflow format.
|
||||
|
||||
---
|
||||
|
||||
feel free to open issues and PRs
|
||||
|
||||
+9
-3
@@ -1,13 +1,19 @@
|
||||
from .node.node import InteriorDesignNode
|
||||
from .node.node import ImageNormalize
|
||||
from .node.node import ControlItems
|
||||
|
||||
|
||||
# Registra o nó no ComfyUI
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"interior-design-for-comfyui": InteriorDesignNode,
|
||||
"Interior Design Segmentator": InteriorDesignNode,
|
||||
"Image Normalize": ImageNormalize,
|
||||
"Control Items": ControlItems,
|
||||
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"interior-design-for-comfyui": "Interior Design for ComfyUI",
|
||||
"Interior Design Segmentator": "Interior Design Segmentator",
|
||||
"Image Normalize": "Image Resize Normalizer",
|
||||
"Control Items": "Control Items Selector",
|
||||
}
|
||||
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 104 KiB |
+117
-64
@@ -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,)
|
||||
|
||||
+99
-292
@@ -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
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import gc
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "interior-design-for-comfyui"
|
||||
description = "This node is based on MykolaL/StableDesign"
|
||||
version = "1.1.5"
|
||||
version = "2.0.0"
|
||||
license = {file = "LICENSE"}
|
||||
|
||||
[project.urls]
|
||||
|
||||
+5
-12
@@ -1,12 +1,5 @@
|
||||
transformers==4.48.0
|
||||
huggingface_hub==0.30.2
|
||||
accelerate==1.6.0
|
||||
altair==5.5.0
|
||||
diffusers==0.33.1
|
||||
numpy==1.26.4
|
||||
opencv-python==4.11.0.86
|
||||
pillow==11.2.1
|
||||
scipy==1.15.2
|
||||
torchvision
|
||||
transformers==4.48.0
|
||||
xformers==0.0.29.post3
|
||||
torch
|
||||
transformers
|
||||
Pillow
|
||||
numpy
|
||||
scipy
|
||||
@@ -1,217 +0,0 @@
|
||||
{
|
||||
"last_node_id": 30,
|
||||
"last_link_id": 44,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 14,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
936.3425903320312,
|
||||
-25.60385513305664
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
246
|
||||
],
|
||||
"flags": {
|
||||
"collapsed": false
|
||||
},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 43
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
14.436566352844238,
|
||||
-146.12979125976562
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
42
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"bedroom_1.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 26,
|
||||
"type": "interior-design-for-comfyui",
|
||||
"pos": [
|
||||
487.18255615234375,
|
||||
-8.370706558227539
|
||||
],
|
||||
"size": [
|
||||
340.20001220703125,
|
||||
274
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 42
|
||||
},
|
||||
{
|
||||
"name": "positive",
|
||||
"type": "STRING",
|
||||
"link": 44,
|
||||
"widget": {
|
||||
"name": "positive"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "negative",
|
||||
"type": "STRING",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "negative"
|
||||
},
|
||||
"shape": 7
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
43
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "interior-design-for-comfyui"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
10,
|
||||
50,
|
||||
0.9,
|
||||
814014690245587,
|
||||
"randomize",
|
||||
768,
|
||||
0.5,
|
||||
0.5,
|
||||
""
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 29,
|
||||
"type": "StringFunction|pysssss",
|
||||
"pos": [
|
||||
-133.7979736328125,
|
||||
223.61289978027344
|
||||
],
|
||||
"size": [
|
||||
577.1492309570312,
|
||||
354.80126953125
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "STRING",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
44
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "StringFunction|pysssss"
|
||||
},
|
||||
"widgets_values": [
|
||||
"append",
|
||||
"yes",
|
||||
"An elegantly appointed bedroom in the Art Deco style, featuring a grand king-size bed with geometric bedding, a luxurious velvet armchair, and a mirrored nightstand that reflects the room's opulence. Art Deco-inspired artwork adds a touch of glamour",
|
||||
"",
|
||||
"",
|
||||
"An elegantly appointed bedroom in the Art Deco style, featuring a grand king-size bed with geometric bedding, a luxurious velvet armchair, and a mirrored nightstand that reflects the room's opulence. Art Deco-inspired artwork adds a touch of glamour"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
42,
|
||||
5,
|
||||
0,
|
||||
26,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
43,
|
||||
26,
|
||||
0,
|
||||
14,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
44,
|
||||
29,
|
||||
0,
|
||||
26,
|
||||
1,
|
||||
"STRING"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.5706456420278478,
|
||||
"offset": [
|
||||
625.0543518197444,
|
||||
269.4681234254746
|
||||
]
|
||||
},
|
||||
"ue_links": [],
|
||||
"node_versions": {
|
||||
"comfy-core": "v0.3.7-10-g23827ca",
|
||||
"StableDesign-for-ComfyUI": "d469e19f1196cce71737197957efd604c7d45c4f",
|
||||
"comfyui-custom-scripts": "a53ef9b617ed1331640d7a2cd97644995908dc00"
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user