359 lines
13 KiB
Python
359 lines
13 KiB
Python
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",
|
|
|
|
} |