nodes/maepopi-Diffusers-in-ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
56677782ccff8a5e20862d6fe0c9fd1f69b4eaa2
maepopi-Diffusers-in-ComfyUI/diffusers_in_comfy.py
T

299 lines
9.9 KiB
Python
Raw Blame History

from diffusers import AutoencoderKL, 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 GenerateStableDiffusionPipeline:
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"),),
},
"optional":
{
"vae": ("STRING", {"multiline": False}),
"controlnet_model" : ("STRING", {"multiline": False}),
}
}
RETURN_TYPES = ("PIPELINE",)
FUNCTION = "create_pipeline"
CATEGORY = "Diffusers-in-Comfy"
def create_pipeline(self, is_sdxl, low_vram, model, vae, controlnet_model):
args = {
"pretrained_model_link_or_path" : folder_paths.get_full_path("checkpoints", model),
"torch_dtype" : torch.float16,
}
if vae != '':
args['vae'] = AutoencoderKL.from_pretrained(vae, torch_dtype=torch.float16, use_safetensors=True)
if controlnet_model != '':
args['controlnet'] = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16, use_safetensors=True)
if is_sdxl:
if 'controlnet' in args:
pipeline = StableDiffusionXLControlNetPipeline.from_single_file(**args)
else:
pipeline = StableDiffusionXLPipeline.from_single_file(**args)
else:
if 'controlnet' in args:
pipeline = StableDiffusionControlNetPipeline.from_single_file(**args)
else:
pipeline = StableDiffusionPipeline.from_single_file(**args)
if low_vram:
pipeline.enable_xformers_memory_efficient_attention()
pipeline.enable_model_cpu_offload()
return (pipeline,)
class MakeCanny:
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"
def convert_images_to_tensors(self, images):
return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images])
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, self.convert_images_to_tensors([canny]),)
class ImageInference:
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}),
},
"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 = "infer_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 infer_image(self, pipeline, seed, steps, cfg, positive, negative, width, height, controlnet_image=None, controlnet_scale=None):
generator = torch.Generator(device='cuda').manual_seed(seed)
print(f'at this stage, controlnet image is {controlnet_image} of type {type(controlnet_image)}')
args = {
"prompt": positive,
"generator": generator,
"num_inference_steps" : steps,
"guidance_scale" : cfg,
"width" : width,
"height" : height,
}
if negative != '':
args['negative'] = negative
if controlnet_image:
args['image'] = controlnet_image
args['controlnet_scale'] = controlnet_scale
images = pipeline(**args).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": GenerateStableDiffusionPipeline,
"GenerateImage": ImageInference,
"LoRALoader" : LoRALoader,
"BLoRALoader" : BLoRALoader,
"MakeCanny": MakeCanny,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"CreatePipeline" : "GenerateStableDiffusionPipeline",
"GenerateImage" : "ImageInference",
"LoRALoader" : "LoRALoader",
"BLoRALoader" : "BLoRALoader",
"MakeCanny": "MakeCanny",
}
Reference in New Issue View Git Blame Copy Permalink