from .utils import PipelineFactory, Text2ImgPipelineCreator, Img2ImgPipelineCreator, InpaintPipelineCreator from .utils import Text2ImageInference, InpaintInference, Img2ImgInference from .utils import convert_images_to_tensors, filter_lora, scale_lora, invert_mask from diffusers.utils import load_image import folder_paths import torch import numpy as np import os import cv2 from PIL import Image '''To do - Investigate the model caching - Deduce the model architecture, SDXL or SD (see B-Lora Comfy code) - change pipeline instanciation ''' # Registering the needed pipelines PipelineFactory.register_creator('text2img', Text2ImgPipelineCreator) PipelineFactory.register_creator('img2img', Img2ImgPipelineCreator) PipelineFactory.register_creator('inpainting', InpaintPipelineCreator) class Text2ImgStableDiffusionPipeline: ''' This node generates the Stable Diffusion Pipeline for Text2Img. Required inputs : - is_sdxl : whether the loaded model has an SDXL base - low_vram : whether to activate low VRAM options - model : the model that has to be instanciated Optional outputs: - vae : which VAE to use - copy the name of the VAE from Hugging Face - controlnet_model : which ControlNet model to use - copy the name of the ControlNet from Hugging Face Ouputs: Pipeline ''' def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "is_sdxl": ("BOOLEAN", {"default": True}), "low_vram": ("BOOLEAN", {"default": False}), "model": (folder_paths.get_filename_list("checkpoints"),), }, "optional": { "vae": ("STRING", {"multiline": False}), "controlnet_model" : ("STRING", {"multiline": False}), } } RETURN_TYPES = ("PIPELINE",) FUNCTION = "create_text2img_pipeline" CATEGORY = "Diffusers-in-Comfy/Pipelines" def create_text2img_pipeline(self, is_sdxl, low_vram, model, vae, controlnet_model): factory = PipelineFactory() pipeline = factory.get_pipeline(low_vram, 'text2img', model, vae, controlnet_model, is_sdxl, torch_dtype=torch.float16) return (pipeline,) class Img2ImgStableDiffusionPipeline: ''' This node generates the Stable Diffusion Pipeline for Img2Img. Required inputs : - is_sdxl : whether the loaded model has an SDXL base - low_vram : whether to activate low VRAM options - model : the model that has to be instanciated Optional outputs: - vae : which VAE to use - copy the name of the VAE from Hugging Face - controlnet_model : which ControlNet model to use - copy the name of the ControlNet from Hugging Face Ouputs: Pipeline ''' def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "is_sdxl": ("BOOLEAN", {"default": True}), "low_vram": ("BOOLEAN", {"default": False}), "model": (folder_paths.get_filename_list("checkpoints"),), }, "optional": { "vae": ("STRING", {"multiline": False}), "controlnet_model" : ("STRING", {"multiline": False}), } } RETURN_TYPES = ("PIPELINE",) FUNCTION = "create_img2img_pipeline" CATEGORY = "Diffusers-in-Comfy/Pipelines" def create_img2img_pipeline(self, is_sdxl, low_vram, model, vae, controlnet_model): factory = PipelineFactory() pipeline = factory.get_pipeline(low_vram, 'img2img', model, vae, controlnet_model, is_sdxl, torch_dtype=torch.float16) return (pipeline,) class InpaintingStableDiffusionPipeline: ''' This node generates the Stable Diffusion Pipeline for inpainting. Required inputs : - is_sdxl : whether the loaded model has an SDXL base - low_vram : whether to activate low VRAM options - model : the model that has to be instanciated Optional outputs: - vae : which VAE to use - copy the name of the VAE from Hugging Face - controlnet_model : which ControlNet model to use - copy the name of the ControlNet from Hugging Face Ouputs: Pipeline ''' def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "is_sdxl": ("BOOLEAN", {"default": True}), "low_vram": ("BOOLEAN", {"default": False}), "model": (folder_paths.get_filename_list("checkpoints"),), }, "optional": { "vae": ("STRING", {"multiline": False}), "controlnet_model" : ("STRING", {"multiline": False}), } } RETURN_TYPES = ("PIPELINE",) FUNCTION = "create_inpaint_pipeline" CATEGORY = "Diffusers-in-Comfy/Pipelines" def create_inpaint_pipeline(self, is_sdxl, low_vram, model, vae, controlnet_model): factory = PipelineFactory() pipeline = factory.get_pipeline(low_vram, 'inpainting', model, vae, controlnet_model, is_sdxl, torch_dtype=torch.float16) return (pipeline,) class MakeCanny: ''' This node allows you to transform your image into a canny mask. Required inputs: - image_path : the string path to your image on your computer, or an URL - high_threshold : high threshold for the edges to be calculated - low_threshold : low_threshold for the edges to be calculated Output: - an image "Canny" that can then be processed further down the Comfy Workflow - a tensor 'Canny Preview" that is an image converted to tensor to be displayed by Comfy's preview image node ''' def __init__(self) -> None: pass @classmethod def INPUT_TYPES(cls): return {"required": { "image_path" : ("STRING", {"multiline": True}), "high_threshold": ("INT", {"default": 200, "min":0, "max":255}), "low_threshold": ("INT", {"default": 100, "min":0, "max":255}), }} RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("Canny", "Canny Preview") FUNCTION = "create_canny" CATEGORY = "Diffusers-in-Comfy/Utils" def create_canny(self, image_path, low_threshold, high_threshold): original_image = load_image(image_path) image = np.array(original_image) image = cv2.Canny(image, low_threshold, high_threshold) image = image[:, :, None] image = np.concatenate([image, image, image], axis=2) canny = Image.fromarray(image) return (canny, convert_images_to_tensors([canny]),) class GenerateInpaintImage: """ This class proceeds to the inference of the image based on the pipeline and its components. Required inputs: - pipeline : the Stable Diffusion Pipeline - seed : seed for the generation of the image - positive : the positive prompt - negative : the negative prompt - steps : the number of inference steps - width : width of the generated image - height : height of the generated image - cfg : guidance scale Optional inputs: - controlnet_image : the image coming out of the Canny Node - controlnet_scale : weight of the canny to control the inferred image - input_image : the image that has to be inpainted - mask_image : the mask of the area where the image has to be inpainted Output: - an image """ 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}), "cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), "input_image" : ("STRING", {"multiline": True}), "mask_image" : ("STRING", {"multiline": True}), "mask_invert" : ("BOOLEAN",{"default":True}), }, "optional": { "controlnet_image" : ("IMAGE",), "controlnet_scale": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step":0.1, "round": 0.01}), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Result",) FUNCTION = "generate_image" CATEGORY = "Diffusers-in-Comfy/Inference" def generate_image(self, pipeline, seed, steps, cfg, positive, negative, width, height, input_image, mask_image, mask_invert, controlnet_image=None, controlnet_scale=None, ): inpaint_inferer = InpaintInference() print(f' at this stage input image is {input_image}') images = inpaint_inferer.infer_image(pipeline, seed, steps, cfg, positive, negative, width, height, input_image, mask_image, mask_invert, controlnet_image=controlnet_image, controlnet_scale=controlnet_scale, ) return (images,) class GenerateImg2Image: """ This class proceeds to the inference of the image based on the pipeline and its components. Required inputs: - pipeline : the Stable Diffusion Pipeline - seed : seed for the generation of the image - positive : the positive prompt - negative : the negative prompt - steps : the number of inference steps - width : width of the generated image - height : height of the generated image - cfg : guidance scale - input_image : the image that is going to be used as prompt Optional inputs: - controlnet_image : the image coming out of the Canny Node - controlnet_scale : weight of the canny to control the inferred image Output: - an image """ 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}), "cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), "input_image" : ("STRING", {"multiline": True}), }, "optional": { "controlnet_image" : ("IMAGE",), "controlnet_scale": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step":0.1, "round": 0.01}), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Result",) FUNCTION = "generate_image" CATEGORY = "Diffusers-in-Comfy/Inference" def generate_image(self, pipeline, seed, steps, cfg, positive, negative, width, height, input_image, controlnet_image=None, controlnet_scale=None, ): img2img_inferer = Img2ImgInference() images = img2img_inferer.infer_image(pipeline, seed, steps, cfg, positive, negative, width, height, input_image, controlnet_image=controlnet_image, controlnet_scale=controlnet_scale, ) return (images,) class GenerateText2Image: """ This class proceeds to the inference of the image based on the pipeline and its components. Required inputs: - pipeline : the Stable Diffusion Pipeline - seed : seed for the generation of the image - positive : the positive prompt - negative : the negative prompt - steps : the number of inference steps - width : width of the generated image - height : height of the generated image - cfg : guidance scale Optional inputs: - controlnet_image : the image coming out of the Canny Node - controlnet_scale : weight of the canny to control the inferred image Output: - an image """ 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}), "cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), }, "optional": { "controlnet_image" : ("IMAGE",), "controlnet_scale": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step":0.1, "round": 0.01}), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("Result",) FUNCTION = "generate_image" CATEGORY = "Diffusers-in-Comfy/Inference" def generate_image(self, pipeline, seed, steps, cfg, positive, negative, width, height, controlnet_image=None, controlnet_scale=None): txt2img_inferer = Text2ImageInference() images = txt2img_inferer.infer_image(pipeline, seed, steps, cfg, positive, negative, width, height, controlnet_image=controlnet_image, controlnet_scale=controlnet_scale) return (images,) class LoRALoader: """ This node allows to load a LoRA to the pipeline. Required inputs: - pipeline : the Stable diffusion pipeline to connect the LoRA to - lora_name : name of your lora in the Lora folder. Only write the namen with the extension, not the total path - lora_scale : weight of the lora on the generated output Output: - pipeline with the LoRA loaded to the Unet """ 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/Components" 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] # to remove the extension from the name pipeline.load_lora_weights(lora_path, adapter_name=lora_name) pipeline.set_adapters([lora_name], adapter_weights=[lora_scale]) return (pipeline,) class BLoRALoader: """ This node allows to load a B-Lora to the Unet. Required: - pipeline : the Stable Diffusion pipeline to which the B-Lora must be applied - style_lora_name : name of the lora you want to apply as style lora (only write the name + extension inside your lora folder, not total path) - style_lora_scale : weight of the style lora on the generated output - content_lora_name : name of the lora you want to apply as content lora (only write the name + extension inside your lora folder, not total path) - content_lora_scale : weight of the content lora on the generated output Output: - pipeline with the B-LoRA loaded into its Unet. """ 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/Components" 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 = filter_lora(style_lora_state_dict, self.BLOCKS['style']) style_lora = 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 = filter_lora(content_lora_state_dict, self.BLOCKS['content']) content_lora = 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 = { "Text2ImgStableDiffusionPipeline": Text2ImgStableDiffusionPipeline, "Img2ImgStableDiffusionPipeline" : Img2ImgStableDiffusionPipeline, "InpaintingStableDiffusionPipeline": InpaintingStableDiffusionPipeline, "GenerateTxt2Image" : GenerateText2Image, "GenerateInpaintImage": GenerateInpaintImage, "GenerateImg2Image" : GenerateImg2Image, "LoRALoader" : LoRALoader, "BLoRALoader" : BLoRALoader, "MakeCanny": MakeCanny, } NODE_DISPLAY_NAME_MAPPINGS = { "Text2ImgStableDiffusionPipeline" : "Text2ImgStableDiffusionPipeline", "Img2ImgStableDiffusionPipeline" : "Img2ImgStableDiffusionPipeline", "InpaintingStableDiffusionPipeline": "InpaintingStableDiffusionPipeline", "GenerateTxt2Image" : "GenerateTxt2Image", "GenerateInpaintImage": "GenerateInpaintImage", "GenerateImg2Image" : "GenerateImg2Image", "LoRALoader" : "LoRALoader", "BLoRALoader" : "BLoRALoader", "MakeCanny": "MakeCanny", }