feat: update to version 2.0.0 with new features
This commit is contained in:
rodrigoskohl
+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,)
|
||||
|
||||
Reference in New Issue
Block a user