From 3c8c28681e6a25b7b13527c3e8ec05bb98b66d88 Mon Sep 17 00:00:00 2001 From: Limitex <76650151+Limitex@users.noreply.github.com> Date: Wed, 27 Dec 2023 06:30:26 +0900 Subject: [PATCH] Adding a base node --- __init__.py | 2 +- example_node.py | 102 ------------------- nodes.py | 252 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 6 ++ utils.py | 204 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 463 insertions(+), 103 deletions(-) delete mode 100644 example_node.py create mode 100644 nodes.py create mode 100644 requirements.txt create mode 100644 utils.py diff --git a/__init__.py b/__init__.py index c9d83bc..c594d08 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,3 @@ -from .example_node import * +from .nodes import * __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/example_node.py b/example_node.py deleted file mode 100644 index 733014f..0000000 --- a/example_node.py +++ /dev/null @@ -1,102 +0,0 @@ -class Example: - """ - A example node - - Class methods - ------------- - INPUT_TYPES (dict): - Tell the main program input parameters of nodes. - - Attributes - ---------- - RETURN_TYPES (`tuple`): - The type of each element in the output tulple. - RETURN_NAMES (`tuple`): - Optional: The name of each output in the output tulple. - FUNCTION (`str`): - The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() - OUTPUT_NODE ([`bool`]): - If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. - The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. - Assumed to be False if not present. - CATEGORY (`str`): - The category the node should appear in the UI. - execute(s) -> tuple || None: - The entry point method. The name of this method must be the same as the value of property `FUNCTION`. - For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. - """ - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - """ - Return a dictionary which contains config for all input fields. - Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". - Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. - The type can be a list for selection. - - Returns: `dict`: - - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` - - Value input_fields (`dict`): Contains input fields config: - * Key field_name (`string`): Name of a entry-point method's argument - * Value field_config (`tuple`): - + First value is a string indicate the type of field or a list for selection. - + Secound value is a config for type "INT", "STRING" or "FLOAT". - """ - return { - "required": { - "image": ("IMAGE",), - "int_field": ("INT", { - "default": 0, - "min": 0, #Minimum value - "max": 4096, #Maximum value - "step": 64, #Slider's step - "display": "number" # Cosmetic only: display as "number" or "slider" - }), - "float_field": ("FLOAT", { - "default": 1.0, - "min": 0.0, - "max": 10.0, - "step": 0.01, - "round": 0.001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. - "display": "number"}), - "print_to_screen": (["enable", "disable"],), - "string_field": ("STRING", { - "multiline": False, #True if you want the field to look like the one on the ClipTextEncode node - "default": "Hello World!" - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - #RETURN_NAMES = ("image_output_name",) - - FUNCTION = "test" - - #OUTPUT_NODE = False - - CATEGORY = "Example" - - def test(self, image, string_field, int_field, float_field, print_to_screen): - if print_to_screen == "enable": - print(f"""Your input contains: - string_field aka input text: {string_field} - int_field: {int_field} - float_field: {float_field} - """) - #do some processing on the image, in this example I just invert it - image = 1.0 - image - return (image,) - - -# A dictionary that contains all nodes you want to export with their names -# NOTE: names should be globally unique -NODE_CLASS_MAPPINGS = { - "Example": Example -} - -# A dictionary that contains the friendly/humanly readable titles for the nodes -NODE_DISPLAY_NAME_MAPPINGS = { - "Example": "Example Node" -} diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..47e5771 --- /dev/null +++ b/nodes.py @@ -0,0 +1,252 @@ +import json +import os +from .utils import SCHEDULERS, token_auto_concat_embeds, vae_pt_to_vae_diffuser +import numpy as np +import torch +from comfy.model_management import get_torch_device, get_torch_device_name +import folder_paths +from diffusers import StableDiffusionPipeline, AutoencoderKL +from comfy.cli_args import args +from PIL import Image, ImageOps, ImageSequence +from PIL.PngImagePlugin import PngInfo + + +class DiffusersPipelineLoader: + def __init__(self): + self.tmp_dir = folder_paths.get_temp_directory() + self.dtype = torch.float32 + + @classmethod + def INPUT_TYPES(s): + return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), }} + + RETURN_TYPES = ("PIPELINE",) + + FUNCTION = "create_pipeline" + + CATEGORY = "Diffusers" + + def create_pipeline(self, ckpt_name): + ckpt_cache_path = os.path.join(self.tmp_dir, ckpt_name) + + StableDiffusionPipeline.from_single_file( + pretrained_model_link_or_path=folder_paths.get_full_path("checkpoints", ckpt_name), + torch_dtype=self.dtype, + cache_dir=self.tmp_dir, + ).save_pretrained(ckpt_cache_path, safe_serialization=True) + + pipe = StableDiffusionPipeline.from_pretrained( + pretrained_model_name_or_path=ckpt_cache_path, + torch_dtype=self.dtype, + cache_dir=self.tmp_dir, + ) + return ((pipe, ckpt_cache_path),) + +class DiffusersVaeLoader: + def __init__(self): + self.tmp_dir = folder_paths.get_temp_directory() + self.dtype = torch.float32 + + @classmethod + def INPUT_TYPES(s): + return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), ), }} + + RETURN_TYPES = ("AUTOENCODER",) + + FUNCTION = "create_pipeline" + + CATEGORY = "Diffusers" + + def create_pipeline(self, vae_name): + ckpt_cache_path = os.path.join(self.tmp_dir, vae_name) + vae_pt_to_vae_diffuser(folder_paths.get_full_path("vae", vae_name), ckpt_cache_path) + + vae = AutoencoderKL.from_pretrained( + pretrained_model_name_or_path=ckpt_cache_path, + torch_dtype=self.dtype, + cache_dir=self.tmp_dir, + ) + + return (vae,) + +class DiffusersSchedulerLoader: + def __init__(self): + self.tmp_dir = folder_paths.get_temp_directory() + self.dtype = torch.float32 + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pipeline": ("PIPELINE", ), + "scheduler_name": (list(SCHEDULERS.keys()), ), + } + } + + RETURN_TYPES = ("SCHEDULER",) + + FUNCTION = "load_scheduler" + + CATEGORY = "Diffusers" + + def load_scheduler(self, pipeline, scheduler_name): + scheduler = SCHEDULERS[scheduler_name].from_pretrained( + pretrained_model_name_or_path=pipeline[1], + torch_dtype=self.dtype, + cache_dir=self.tmp_dir, + subfolder='scheduler' + ) + return (scheduler,) + +class DiffusersModelMakeup: + def __init__(self): + self.torch_device = get_torch_device() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pipeline": ("PIPELINE", ), + "autoencoder": ("AUTOENCODER", ), + "scheduler": ("SCHEDULER", ), + }, + } + + RETURN_TYPES = ("MAKED_PIPELINE",) + + FUNCTION = "makeup_pipeline" + + CATEGORY = "Diffusers" + + def makeup_pipeline(self, pipeline, autoencoder, scheduler): + pipeline = pipeline[0] + pipeline.vae = autoencoder + pipeline.scheduler = scheduler + pipeline.safety_checker = None if pipeline.safety_checker is None else lambda images, **kwargs: (images, [False]) + pipeline.enable_attention_slicing() + pipeline = pipeline.to(self.torch_device) + return (pipeline,) + +class DiffusersClipTextEncode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "maked_pipeline": ("MAKED_PIPELINE", ), + "positive": ("STRING", {"multiline": True}), + "negative": ("STRING", {"multiline": True}), + }} + + RETURN_TYPES = ("EMBEDS", "EMBEDS", ) + RETURN_NAMES = ("positive", "negative", ) + + FUNCTION = "concat_embeds" + + CATEGORY = "Diffusers" + + def concat_embeds(self, maked_pipeline, positive, negative): + positive_embeds, negative_embeds = token_auto_concat_embeds(maked_pipeline, positive,negative) + + return (positive_embeds, negative_embeds, ) + +class DiffusersSampler: + def __init__(self): + self.torch_device = get_torch_device() + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "maked_pipeline": ("MAKED_PIPELINE", ), + "positive": ("EMBEDS", ), + "negative": ("EMBEDS", ), + "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}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }} + + RETURN_TYPES = ("PIL_IMAGE",) + + FUNCTION = "sample" + + CATEGORY = "Diffusers" + + def sample(self, maked_pipeline, positive, negative, height, width, steps, cfg, seed): + images = maked_pipeline( + prompt_embeds=positive, + height=height, + width=width, + num_inference_steps=steps, + guidance_scale=cfg, + negative_prompt_embeds=negative, + generator=torch.Generator(self.torch_device).manual_seed(seed) + ).images + return (images,) + +class DiffusersSaveImage: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(s): + return {"required": + {"images": ("PIL_IMAGE", ), + "filename_prefix": ("STRING", {"default": "ComfyUI"})}, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + + OUTPUT_NODE = True + + CATEGORY = "Diffusers" + + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + filename_prefix += self.prefix_append + width, height = images[0].size + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, width, height) + results = list() + for image in images: + metadata = None + if not args.disable_metadata: + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + file = f"{filename}_{counter:05}_.png" + image.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + return { "ui": { "images": results } } + +NODE_CLASS_MAPPINGS = { + "DiffusersPipelineLoader": DiffusersPipelineLoader, + "DiffusersVaeLoader": DiffusersVaeLoader, + "DiffusersSchedulerLoader": DiffusersSchedulerLoader, + "DiffusersModelMakeup": DiffusersModelMakeup, + "DiffusersClipTextEncode": DiffusersClipTextEncode, + "DiffusersSampler": DiffusersSampler, + "DiffusersSaveImage": DiffusersSaveImage +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "DiffusersPipelineLoader": "Diffusers Pipeline Loader", + "DiffusersVaeLoader": "Diffusers Vae Loader", + "DiffusersSchedulerLoader": "Diffusers Scheduler Loader", + "DiffusersModelMakeup": "Diffusers Model Makeup", + "DiffusersClipTextEncode": "Diffusers Clip Text Encode", + "DiffusersSampler": "Diffusers Sampler", + "DiffusersSaveImage": "Diffusers Save Image" +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d1b11fa --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +diffusers[torch] +accelerate +transformers +safetensors +omegaconf +pytorch_lightning \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..c387152 --- /dev/null +++ b/utils.py @@ -0,0 +1,204 @@ +import torch +import argparse +import io + +import requests +import torch +from omegaconf import OmegaConf +from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( + assign_to_checkpoint, + conv_attn_to_linear, + create_vae_diffusers_config, + renew_vae_attention_paths, + renew_vae_resnet_paths, +) + +from diffusers import ( + AutoencoderKL, + DDIMScheduler, + DDPMScheduler, + DEISMultistepScheduler, + DPMSolverMultistepScheduler, + DPMSolverSinglestepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + HeunDiscreteScheduler, + KDPM2AncestralDiscreteScheduler, + KDPM2DiscreteScheduler, + UniPCMultistepScheduler, +) + +SCHEDULERS = { + 'DDIM' : DDIMScheduler, + 'DDPM' : DDPMScheduler, + 'DEISMultistep' : DEISMultistepScheduler, + 'DPMSolverMultistep' : DPMSolverMultistepScheduler, + 'DPMSolverSinglestep' : DPMSolverSinglestepScheduler, + 'EulerAncestralDiscrete' : EulerAncestralDiscreteScheduler, + 'EulerDiscrete' : EulerDiscreteScheduler, + 'HeunDiscrete' : HeunDiscreteScheduler, + 'KDPM2AncestralDiscrete' : KDPM2AncestralDiscreteScheduler, + 'KDPM2Discrete' : KDPM2DiscreteScheduler, + 'UniPCMultistep' : UniPCMultistepScheduler +} + +def token_auto_concat_embeds(pipe, positive, negative): + max_length = pipe.tokenizer.model_max_length + positive_length = pipe.tokenizer(positive, return_tensors="pt").input_ids.shape[-1] + negative_length = pipe.tokenizer(negative, return_tensors="pt").input_ids.shape[-1] + + print(f'Token length is model maximum: {max_length}, positive length: {positive_length}, negative length: {negative_length}.') + if max_length < positive_length or max_length < negative_length: + print('Concatenated embedding.') + if positive_length > negative_length: + positive_ids = pipe.tokenizer(positive, return_tensors="pt").input_ids.to("cuda") + negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=positive_ids.shape[-1], return_tensors="pt").input_ids.to("cuda") + else: + negative_ids = pipe.tokenizer(negative, return_tensors="pt").input_ids.to("cuda") + positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=negative_ids.shape[-1], return_tensors="pt").input_ids.to("cuda") + else: + positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda") + negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda") + + positive_concat_embeds = [] + negative_concat_embeds = [] + for i in range(0, positive_ids.shape[-1], max_length): + positive_concat_embeds.append(pipe.text_encoder(positive_ids[:, i: i + max_length])[0]) + negative_concat_embeds.append(pipe.text_encoder(negative_ids[:, i: i + max_length])[0]) + + positive_prompt_embeds = torch.cat(positive_concat_embeds, dim=1) + negative_prompt_embeds = torch.cat(negative_concat_embeds, dim=1) + return positive_prompt_embeds, negative_prompt_embeds + + +def custom_convert_ldm_vae_checkpoint(checkpoint, config): + vae_state_dict = checkpoint + + new_checkpoint = {} + + new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] + new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] + new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] + new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] + new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] + new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] + + new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] + new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] + new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] + new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] + new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] + new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] + + new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] + new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] + new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] + new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] + + # Retrieves the keys for the encoder down blocks only + num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) + down_blocks = { + layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) + } + + # Retrieves the keys for the decoder up blocks only + num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) + up_blocks = { + layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) + } + + for i in range(num_down_blocks): + resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] + + if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: + new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( + f"encoder.down.{i}.downsample.conv.weight" + ) + new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( + f"encoder.down.{i}.downsample.conv.bias" + ) + + paths = renew_vae_resnet_paths(resnets) + meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} + assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) + + mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] + num_mid_res_blocks = 2 + for i in range(1, num_mid_res_blocks + 1): + resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] + + paths = renew_vae_resnet_paths(resnets) + meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} + assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) + + mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] + paths = renew_vae_attention_paths(mid_attentions) + meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} + assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) + conv_attn_to_linear(new_checkpoint) + + for i in range(num_up_blocks): + block_id = num_up_blocks - 1 - i + resnets = [ + key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key + ] + + if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: + new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ + f"decoder.up.{block_id}.upsample.conv.weight" + ] + new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ + f"decoder.up.{block_id}.upsample.conv.bias" + ] + + paths = renew_vae_resnet_paths(resnets) + meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} + assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) + + mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] + num_mid_res_blocks = 2 + for i in range(1, num_mid_res_blocks + 1): + resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] + + paths = renew_vae_resnet_paths(resnets) + meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} + assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) + + mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] + paths = renew_vae_attention_paths(mid_attentions) + meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} + assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) + conv_attn_to_linear(new_checkpoint) + return new_checkpoint + + +def vae_pt_to_vae_diffuser( + checkpoint_path: str, + output_path: str, +): + # Only support V1 + r = requests.get( + " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml" + ) + io_obj = io.BytesIO(r.content) + + original_config = OmegaConf.load(io_obj) + image_size = 512 + device = "cuda" if torch.cuda.is_available() else "cpu" + if checkpoint_path.endswith("safetensors"): + from safetensors import safe_open + + checkpoint = {} + with safe_open(checkpoint_path, framework="pt", device="cpu") as f: + for key in f.keys(): + checkpoint[key] = f.get_tensor(key) + else: + checkpoint = torch.load(checkpoint_path, map_location=device)["state_dict"] + + # Convert the VAE model. + vae_config = create_vae_diffusers_config(original_config, image_size=image_size) + converted_vae_checkpoint = custom_convert_ldm_vae_checkpoint(checkpoint, vae_config) + + vae = AutoencoderKL(**vae_config) + vae.load_state_dict(converted_vae_checkpoint) + vae.save_pretrained(output_path) \ No newline at end of file