608 lines
22 KiB
Python
608 lines
22 KiB
Python
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",
|
|
|
|
} |