from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline, StableDiffusionXLControlNetPipeline, StableDiffusionControlNetPipeline, ControlNetModel import folder_paths import torch import numpy as np import os from torchvision.transforms import ToTensor from diffusers.utils import load_image, make_image_grid import cv2 from PIL import Image '''To do - Investigate the model caching - Le controlnet c'est pas possible de le passer après que le pipe ait été instancié. Il faudrait faire un node spécial CreateControlNetPipe ''' class CreateGenericPipeline: def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "is_sdxl": ("BOOLEAN", {"default": True}), "low_vram": ("BOOLEAN", {"default": True}), "model": (folder_paths.get_filename_list("checkpoints"),), }} RETURN_TYPES = ("PIPELINE",) FUNCTION = "create_pipeline" CATEGORY = "Diffusers-in-Comfy" def create_pipeline(self, is_sdxl, low_vram, model): model_path = folder_paths.get_full_path("checkpoints", model) if is_sdxl: pipeline = StableDiffusionXLPipeline.from_single_file(model_path, torch_dtype=torch.float16).to("cuda") else: pipeline = StableDiffusionPipeline.from_single_file(model_path, torch_dtype=torch.float16).to("cuda") if low_vram: pipeline.enable_xformers_memory_efficient_attention() pipeline.enable_model_cpu_offload() return (pipeline,) class CreateControlNetPipeline: def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "is_sdxl": ("BOOLEAN", {"default": True}), "low_vram": ("BOOLEAN", {"default": True}), "model": (folder_paths.get_filename_list("checkpoints"),), "controlnet_model" : ("STRING", {"multiline":False}), }} RETURN_TYPES = ("PIPELINE",) FUNCTION = "create_controlnet_pipeline" CATEGORY = "Diffusers-in-Comfy/ControlNet" def create_controlnet_pipeline(self, is_sdxl, low_vram, model, controlnet_model): model_path = folder_paths.get_full_path("checkpoints", model) controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16,use_safetensors=True) if is_sdxl: pipeline = StableDiffusionXLControlNetPipeline.from_single_file(model_path, controlnet=controlnet, torch_dtype=torch.float16).to("cuda") else: pipeline = StableDiffusionControlNetPipeline.from_single_file(model_path, controlnet=controlnet, torch_dtype=torch.float16).to("cuda") if low_vram: pipeline.enable_xformers_memory_efficient_attention() pipeline.enable_model_cpu_offload() return (pipeline, ) class GenerateImage: def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "pipeline": ("PIPELINE",), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "positive": ("STRING", {"multiline": True}), "negative": ("STRING", {"multiline": True}), "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), "height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), }} RETURN_TYPES = ("IMAGE",) FUNCTION = "generate_image" CATEGORY = "Diffusers-in-Comfy" def convert_images_to_tensors(self, images): return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images]) def generate_image(self, pipeline, seed, steps, cfg, positive, negative, width, height): generator = torch.Generator(device='cuda').manual_seed(seed) if negative == '': images = pipeline( prompt=positive, generator=generator, num_inference_steps = steps, guidance_scale = cfg, width=width, height=height, ).images else: images = pipeline( prompt=positive, negative_prompt=negative, generator=generator, num_inference_steps = steps, guidance_scale = cfg, width=width, height=height, ).images return (self.convert_images_to_tensors(images),) class GenerateControlledImage: def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "pipeline": ("PIPELINE",), "controlnet_image_path" : ("STRING", {"multiline": True}), "control_scale": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step":0.1, "round": 0.01}), "controlnet_high_threshold": ("INT", {"default": 200, "min":0, "max":255}), "controlnet_low_threshold": ("INT", {"default": 100, "min":0, "max":255}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "positive": ("STRING", {"multiline": True}), "negative": ("STRING", {"multiline": True}), "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), "height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), }} RETURN_TYPES = ("IMAGE",) FUNCTION = "generate_controlled_image" CATEGORY = "Diffusers-in-Comfy/ControlNet" def create_controlnet_image(self, image_path, high, low): original_image = load_image(image_path) image = np.array(original_image) image = cv2.Canny(image, low, high) image = image[:, :, None] image = np.concatenate([image, image, image], axis=2) return Image.fromarray(image) def convert_images_to_tensors(self, images): return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images]) def generate_controlled_image(self, pipeline, controlnet_image_path, controlnet_high_threshold, controlnet_low_threshold, control_scale, seed, steps, cfg, positive, negative, width, height): generator = torch.Generator(device='cuda').manual_seed(seed) control_image = self.create_controlnet_image(controlnet_image_path, controlnet_high_threshold, controlnet_low_threshold) if negative == '': images = pipeline( prompt=positive, image=control_image, controlnet_conditioning_scale=control_scale, generator=generator, num_inference_steps = steps, guidance_scale = cfg, width=width, height=height, ).images else: images = pipeline( prompt=positive, negative_prompt=negative, image=control_image, controlnet_conditioning_scale=control_scale, generator=generator, num_inference_steps = steps, guidance_scale = cfg, width=width, height=height, ).images return (self.convert_images_to_tensors(images),) class LoRALoader: def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required":{ "pipeline" : ("PIPELINE",), "lora_name" : (folder_paths.get_filename_list("loras"),), "lora_scale" : ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.1, "round": 0.01}), }} RETURN_TYPES = ("PIPELINE",) FUNCTION = "load_lora" CATEGORY = "Diffusers-in-Comfy" def load_lora(self, pipeline, lora_name, lora_scale): lora_path = folder_paths.get_full_path("loras", lora_name) lora_name = os.path.splitext(lora_name)[0] pipeline.load_lora_weights(lora_path, adapter_name=lora_name) pipeline.set_adapters([lora_name], adapter_weights=[lora_scale]) return (pipeline,) class BLoRALoader: BLOCKS = { 'content': ['unet.up_blocks.0.attentions.0'], 'style': ['unet.up_blocks.0.attentions.1'], } def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required":{ "pipeline" : ("PIPELINE",), "style_lora_name" : ("STRING", {"multiline": False}), "style_lora_scale" : ("FLOAT", {"default": 1.1, "min": 0.0, "max": 1.1, "step":0.1, "round": 0.01}), "content_lora_name" : ("STRING", {"multiline": False}), "content_lora_scale" : ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.1, "round": 0.01}), }} RETURN_TYPES = ("PIPELINE",) FUNCTION = "load_b_lora_to_unet" CATEGORY = "Diffusers-in-Comfy" def is_belong_to_blocks(self, key, blocks): try: for g in blocks: if g in key: return True return False except Exception as e: raise type(e)(f'failed to is_belong_to_block, due to: {e}') def filter_lora(self, state_dict, blocks_): try: return {k: v for k, v in state_dict.items() if self.is_belong_to_blocks(k, blocks_)} except Exception as e: raise type(e)(f'failed to filter_lora, due to: {e}') def scale_lora(self, state_dict, alpha): try: return {k: v * alpha for k, v in state_dict.items()} except Exception as e: raise type(e)(f'failed to scale_lora, due to: {e}') def load_b_lora_to_unet(self, pipeline, style_lora_name, content_lora_name, style_lora_scale=1., content_lora_scale=1.): if style_lora_name != '' : style_lora_path = folder_paths.get_full_path("loras", style_lora_name) style_lora_state_dict, _ = pipeline.lora_state_dict(style_lora_path) style_lora = self.filter_lora(style_lora_state_dict, self.BLOCKS['style']) style_lora = self.scale_lora(style_lora, style_lora_scale) else: style_lora = {} if content_lora_name != '': content_lora_path = folder_paths.get_full_path("loras", content_lora_name) content_lora_state_dict, _ = pipeline.lora_state_dict(content_lora_path) content_lora = self.filter_lora(content_lora_state_dict, self.BLOCKS['content']) content_lora = self.scale_lora(content_lora, content_lora_scale) else: content_lora = {} merged_lora = {**style_lora, **content_lora} pipeline.load_lora_into_unet(merged_lora, None, pipeline.unet) return (pipeline,) NODE_CLASS_MAPPINGS = { "CreatePipeline": CreateGenericPipeline, "CreateControlNetPipeline": CreateControlNetPipeline, "GenerateImage": GenerateImage, "GenerateControlledImage": GenerateControlledImage, "LoRALoader" : LoRALoader, "BLoRALoader" : BLoRALoader, } NODE_DISPLAY_NAME_MAPPINGS = { "CreatePipeline" : "CreateGenericPipeline", "CreateControlNetPipeline": "CreateControlNetPipeline", "GenerateImage" : "GenerateImage", "GenerateControlledImage": "GenerateControlledImage", "LoRALoader" : "LoRALoader", "BLoRALoader" : "BLoRALoader", }