diff --git a/node/node.py b/node/node.py index cae5b4c..ccc25ab 100644 --- a/node/node.py +++ b/node/node.py @@ -24,6 +24,7 @@ class InteriorDesignNode: "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}), }, "optional": { "negative": ("STRING",{"forceInput": True}), # Prompt de descrição do design @@ -36,8 +37,9 @@ class InteriorDesignNode: FUNCTION = "process" # Função que será executada CATEGORY = "interior_design" # Categoria para o ComfyUI - def process(self, image, positive, guidance_scale, num_steps, strength, seed, img_size, img_dept_threshold, img_seg_threshold, negative=None): + def process(self, image, positive, guidance_scale, num_steps, strength, seed, img_size, img_dept_threshold, img_seg_threshold, ip_adapter_scale, negative=None): # Pré-processamento da imagem + self.pipeline.ip_adapter_scale = ip_adapter_scale self.pipeline.seed = seed if negative is not None: self.pipeline.neg_prompt = negative diff --git a/node/pipeline.py b/node/pipeline.py index 016aadf..d48bf6c 100644 --- a/node/pipeline.py +++ b/node/pipeline.py @@ -189,11 +189,49 @@ class ControlNetDepthDesignModelMulti: def __init__(self): """ Initialize your model(s) here """ #os.environ['HF_HUB_OFFLINE'] = "True" - + self.ip_adapter_scale = 0.4 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( + os.path.join(rootfolder, "controlnet_depth"), + torch_dtype=dtype, + use_safetensors=True + ) + + self.controlnet_seg = ControlNetModel.from_pretrained( + os.path.join(rootfolder, "own_controlnet"), + torch_dtype=dtype, + use_safetensors=True + ) + + self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( + "SG161222/Realistic_Vision_V5.1_noVAE", + controlnet=[self.controlnet_depth, self.controlnet_seg], + safety_checker=None, + torch_dtype=dtype + ) + + self.pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", + weight_name="ip-adapter_sd15.bin") + 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: """ @@ -206,7 +244,11 @@ class ControlNetDepthDesignModelMulti: design_image - PIL Image of the same size as the empty room image If the size is not the same the submission will fail. """ - print(prompt) + 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}") + flush() self.generator = torch.Generator(device=device).manual_seed(self.seed) @@ -216,8 +258,8 @@ class ControlNetDepthDesignModelMulti: 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, - seg_image_processor, - image_segmentor)) + self.seg_image_processor, + self.image_segmentor)) 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] @@ -236,13 +278,13 @@ class ControlNetDepthDesignModelMulti: 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, depth_feature_extractor, depth_estimator) + 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 = guide_pipe(pos_prompt, + ip_image = self.guide_pipe(pos_prompt, num_inference_steps=num_steps, negative_prompt=self.neg_prompt, height=new_height_ip, @@ -250,7 +292,7 @@ class ControlNetDepthDesignModelMulti: generator=[self.generator]).images[0] flush() - generated_image = pipe( + generated_image = self.pipe( prompt=pos_prompt, negative_prompt=self.neg_prompt, num_inference_steps=num_steps, @@ -273,33 +315,5 @@ class ControlNetDepthDesignModelMulti: -controlnet_depth = ControlNetModel.from_pretrained( - os.path.join(rootfolder,"controlnet_depth"), torch_dtype=dtype, use_safetensors=True -) - -controlnet_seg = ControlNetModel.from_pretrained( - os.path.join(rootfolder,"own_controlnet"), torch_dtype=dtype, use_safetensors=True -) - -pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( - "SG161222/Realistic_Vision_V5.1_noVAE", - #"models/runwayml--stable-diffusion-inpainting", - controlnet=[controlnet_depth, controlnet_seg], - safety_checker=None, - torch_dtype=dtype -) - -pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", - weight_name="ip-adapter_sd15.bin") -pipe.set_ip_adapter_scale(0.4) -pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) -pipe = pipe.to(device) -guide_pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", - torch_dtype=dtype, use_safetensors=True, variant="fp16") -guide_pipe = guide_pipe.to(device) - -seg_image_processor, image_segmentor = get_segmentation_pipeline() -depth_feature_extractor, depth_estimator = get_depth_pipeline() -depth_estimator = depth_estimator.to(device) diff --git a/pyproject.toml b/pyproject.toml index 94523a5..cc1548d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "interior-design-for-comfyui" description = "This node is based on MykolaL/StableDesign" -version = "1.0.0" +version = "1.1.0" license = {file = "LICENSE"} [project.urls]