476 lines
18 KiB
Python
476 lines
18 KiB
Python
from __future__ import annotations
|
|
import torch
|
|
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import hashlib
|
|
import inspect
|
|
import traceback
|
|
import math
|
|
import time
|
|
import random
|
|
import logging
|
|
|
|
from PIL import Image, ImageOps, ImageSequence
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
import numpy as np
|
|
import safetensors.torch
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
|
|
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import comfy.controlnet
|
|
|
|
import comfy.clip_vision
|
|
|
|
import comfy.model_management
|
|
|
|
import folder_paths
|
|
from comfy.comfy_types.node_typing import IO
|
|
|
|
def before_node_execution():
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
|
|
def interrupt_processing(value=True):
|
|
comfy.model_management.interrupt_current_processing(value)
|
|
|
|
MAX_RESOLUTION=16384
|
|
|
|
|
|
class CheckpointLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "config_name": (folder_paths.get_filename_list("configs"), ),
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), )},
|
|
"optional": {"any": (IO.ANY, {})}}
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
FUNCTION = "load_checkpoint"
|
|
|
|
CATEGORY = "advanced/loaders"
|
|
DEPRECATED = True
|
|
|
|
def load_checkpoint(self, config_name, ckpt_name, any=None):
|
|
config_path = folder_paths.get_full_path("configs", config_name)
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
|
return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
|
|
|
|
class CheckpointLoaderSimple_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}),
|
|
},
|
|
"optional": {"any": (IO.ANY, {})}
|
|
}
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
OUTPUT_TOOLTIPS = ("The model used for denoising latents.",
|
|
"The CLIP model used for encoding text prompts.",
|
|
"The VAE model used for encoding and decoding images to and from latent space.")
|
|
FUNCTION = "load_checkpoint"
|
|
|
|
CATEGORY = "loaders"
|
|
DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents."
|
|
|
|
def load_checkpoint(self, ckpt_name, any=None):
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
|
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
return out[:3]
|
|
|
|
|
|
class unCLIPCheckpointLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
|
|
},
|
|
"optional": {"any": (IO.ANY, {})}}
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "CLIP_VISION")
|
|
FUNCTION = "load_checkpoint"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_checkpoint(self, ckpt_name, any=None, output_vae=True, output_clip=True):
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
|
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
return out
|
|
|
|
|
|
class LoraLoader_Any:
|
|
def __init__(self):
|
|
self.loaded_lora = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}),
|
|
"clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}),
|
|
"lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The name of the LoRA."}),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
|
|
},
|
|
"optional": {"any": (IO.ANY, {})}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
OUTPUT_TOOLTIPS = ("The modified diffusion model.", "The modified CLIP model.")
|
|
FUNCTION = "load_lora"
|
|
|
|
CATEGORY = "loaders"
|
|
DESCRIPTION = "LoRAs are used to modify diffusion and CLIP models, altering the way in which latents are denoised such as applying styles. Multiple LoRA nodes can be linked together."
|
|
|
|
def load_lora(self, model, clip, lora_name, strength_model, strength_clip, any=None):
|
|
if strength_model == 0 and strength_clip == 0:
|
|
return (model, clip)
|
|
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
lora = None
|
|
if self.loaded_lora is not None:
|
|
if self.loaded_lora[0] == lora_path:
|
|
lora = self.loaded_lora[1]
|
|
else:
|
|
self.loaded_lora = None
|
|
|
|
if lora is None:
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
self.loaded_lora = (lora_path, lora)
|
|
|
|
new_model, new_clip = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
|
|
return (new_model, new_clip)
|
|
|
|
|
|
class LoraLoaderModelOnly_Any(LoraLoader_Any):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
|
},
|
|
"optional": {"any": (IO.ANY, {})}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load_lora_model_only"
|
|
|
|
def load_lora_model_only(self, model, lora_name, strength_model, any=None):
|
|
return (self.load_lora(model, None, lora_name, strength_model, 0, any)[0],)
|
|
|
|
|
|
class VAELoader_Any:
|
|
video_taes = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5"]
|
|
image_taes = ["taesd", "taesdxl", "taesd3", "taef1"]
|
|
@staticmethod
|
|
def vae_list(s):
|
|
vaes = folder_paths.get_filename_list("vae")
|
|
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
|
sdxl_taesd_enc = False
|
|
sdxl_taesd_dec = False
|
|
sd1_taesd_enc = False
|
|
sd1_taesd_dec = False
|
|
sd3_taesd_enc = False
|
|
sd3_taesd_dec = False
|
|
f1_taesd_enc = False
|
|
f1_taesd_dec = False
|
|
|
|
for v in approx_vaes:
|
|
if v.startswith("taesd_decoder."):
|
|
sd1_taesd_dec = True
|
|
elif v.startswith("taesd_encoder."):
|
|
sd1_taesd_enc = True
|
|
elif v.startswith("taesdxl_decoder."):
|
|
sdxl_taesd_dec = True
|
|
elif v.startswith("taesdxl_encoder."):
|
|
sdxl_taesd_enc = True
|
|
elif v.startswith("taesd3_decoder."):
|
|
sd3_taesd_dec = True
|
|
elif v.startswith("taesd3_encoder."):
|
|
sd3_taesd_enc = True
|
|
elif v.startswith("taef1_encoder."):
|
|
f1_taesd_dec = True
|
|
elif v.startswith("taef1_decoder."):
|
|
f1_taesd_enc = True
|
|
else:
|
|
for tae in s.video_taes:
|
|
if v.startswith(tae):
|
|
vaes.append(v)
|
|
|
|
if sd1_taesd_dec and sd1_taesd_enc:
|
|
vaes.append("taesd")
|
|
if sdxl_taesd_dec and sdxl_taesd_enc:
|
|
vaes.append("taesdxl")
|
|
if sd3_taesd_dec and sd3_taesd_enc:
|
|
vaes.append("taesd3")
|
|
if f1_taesd_dec and f1_taesd_enc:
|
|
vaes.append("taef1")
|
|
vaes.append("pixel_space")
|
|
return vaes
|
|
|
|
@staticmethod
|
|
def load_taesd(name):
|
|
sd = {}
|
|
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
|
|
|
encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes))
|
|
decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes))
|
|
|
|
enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder))
|
|
for k in enc:
|
|
sd["taesd_encoder.{}".format(k)] = enc[k]
|
|
|
|
dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder))
|
|
for k in dec:
|
|
sd["taesd_decoder.{}".format(k)] = dec[k]
|
|
|
|
if name == "taesd":
|
|
sd["vae_scale"] = torch.tensor(0.18215)
|
|
sd["vae_shift"] = torch.tensor(0.0)
|
|
elif name == "taesdxl":
|
|
sd["vae_scale"] = torch.tensor(0.13025)
|
|
sd["vae_shift"] = torch.tensor(0.0)
|
|
elif name == "taesd3":
|
|
sd["vae_scale"] = torch.tensor(1.5305)
|
|
sd["vae_shift"] = torch.tensor(0.0609)
|
|
elif name == "taef1":
|
|
sd["vae_scale"] = torch.tensor(0.3611)
|
|
sd["vae_shift"] = torch.tensor(0.1159)
|
|
return sd
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "vae_name": (s.vae_list(s), )}, "optional": {"any": (IO.ANY, {})}}
|
|
RETURN_TYPES = ("VAE",)
|
|
FUNCTION = "load_vae"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
#TODO: scale factor?
|
|
def load_vae(self, vae_name, any=None):
|
|
if vae_name == "pixel_space":
|
|
sd = {}
|
|
sd["pixel_space_vae"] = torch.tensor(1.0)
|
|
elif vae_name in self.image_taes:
|
|
sd = self.load_taesd(vae_name)
|
|
else:
|
|
if os.path.splitext(vae_name)[0] in self.video_taes:
|
|
vae_path = folder_paths.get_full_path_or_raise("vae_approx", vae_name)
|
|
else:
|
|
vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
|
|
sd = comfy.utils.load_torch_file(vae_path)
|
|
vae = comfy.sd.VAE(sd=sd)
|
|
vae.throw_exception_if_invalid()
|
|
return (vae,)
|
|
|
|
|
|
class ControlNetLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"), )}, "optional": {"any": (IO.ANY, {})}}
|
|
|
|
RETURN_TYPES = ("CONTROL_NET",)
|
|
FUNCTION = "load_controlnet"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_controlnet(self, control_net_name, any=None):
|
|
controlnet_path = folder_paths.get_full_path_or_raise("controlnet", control_net_name)
|
|
controlnet = comfy.controlnet.load_controlnet(controlnet_path)
|
|
if controlnet is None:
|
|
raise RuntimeError("ERROR: controlnet file is invalid and does not contain a valid controlnet model.")
|
|
return (controlnet,)
|
|
|
|
|
|
class DiffControlNetLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "model": ("MODEL",),
|
|
"control_net_name": (folder_paths.get_filename_list("controlnet"), )},
|
|
"optional": {"any": (IO.ANY, {})}}
|
|
|
|
RETURN_TYPES = ("CONTROL_NET",)
|
|
FUNCTION = "load_controlnet"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_controlnet(self, model, control_net_name, any=None):
|
|
controlnet_path = folder_paths.get_full_path_or_raise("controlnet", control_net_name)
|
|
controlnet = comfy.controlnet.load_controlnet(controlnet_path, model)
|
|
return (controlnet,)
|
|
|
|
class UNETLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "unet_name": (folder_paths.get_filename_list("unet"), ),
|
|
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],)
|
|
},
|
|
"optional": {"any": (IO.ANY, {})}
|
|
}
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load_unet"
|
|
|
|
CATEGORY = "advanced/loaders"
|
|
|
|
def load_unet(self, unet_name, weight_dtype, any=None):
|
|
model_options = {}
|
|
if weight_dtype == "fp8_e4m3fn":
|
|
model_options["dtype"] = torch.float8_e4m3fn
|
|
elif weight_dtype == "fp8_e4m3fn_fast":
|
|
model_options["dtype"] = torch.float8_e4m3fn
|
|
model_options["fp8_optimizations"] = True
|
|
elif weight_dtype == "fp8_e5m2":
|
|
model_options["dtype"] = torch.float8_e5m2
|
|
|
|
unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)
|
|
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
|
return (model,)
|
|
|
|
class CLIPLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
|
|
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "cosmos", "lumina2", "wan", "hidream", "omnigen2"], ),
|
|
},
|
|
"optional": {
|
|
"device": (["default", "cpu"], {"advanced": True}),
|
|
"any": (IO.ANY, {})
|
|
}}
|
|
RETURN_TYPES = ("CLIP",)
|
|
FUNCTION = "load_clip"
|
|
|
|
CATEGORY = "advanced/loaders"
|
|
|
|
DESCRIPTION = "[Recipes]\n\nstable_diffusion: clip-l\nstable_cascade: clip-g\nsd3: t5 xxl/ clip-g / clip-l\nstable_audio: t5 base\nmochi: t5 xxl\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\n hidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B"
|
|
|
|
def load_clip(self, clip_name, type, any=None, device="default"):
|
|
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
|
|
|
model_options = {}
|
|
if device == "cpu":
|
|
model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
|
|
|
|
clip_path = folder_paths.get_full_path_or_raise("text_encoders", clip_name)
|
|
clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type, model_options=model_options)
|
|
return (clip,)
|
|
|
|
|
|
class DualCLIPLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "clip_name1": (folder_paths.get_filename_list("text_encoders"), ),
|
|
"clip_name2": (folder_paths.get_filename_list("text_encoders"), ),
|
|
"type": (["sdxl", "sd3", "flux", "hidream", "hunyuan_image"], ),
|
|
},
|
|
"optional": {
|
|
"device": (["default", "cpu"], {"advanced": True}),
|
|
"any": (IO.ANY, {})
|
|
}}
|
|
RETURN_TYPES = ("CLIP",)
|
|
FUNCTION = "load_clip"
|
|
|
|
CATEGORY = "advanced/loaders"
|
|
|
|
DESCRIPTION = "[Recipes]\n\nsdxl: clip-l, clip-g\nsd3: clip-l, clip-g / clip-l, t5 / clip-g, t5\nflux: clip-l, t5\nhidream: at least one of t5 or llama, recommended t5 and llama\nhunyuan_image: qwen2.5vl 7b and byt5 small"
|
|
|
|
def load_clip(self, clip_name1, clip_name2, type, any=None, device="default"):
|
|
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
|
|
|
clip_path1 = folder_paths.get_full_path_or_raise("text_encoders", clip_name1)
|
|
clip_path2 = folder_paths.get_full_path_or_raise("text_encoders", clip_name2)
|
|
|
|
model_options = {}
|
|
if device == "cpu":
|
|
model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
|
|
|
|
clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type, model_options=model_options)
|
|
return (clip,)
|
|
|
|
|
|
class CLIPVisionLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ),
|
|
},
|
|
"optional": {"any": (IO.ANY, {})}}
|
|
RETURN_TYPES = ("CLIP_VISION",)
|
|
FUNCTION = "load_clip"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_clip(self, clip_name, any=None):
|
|
clip_path = folder_paths.get_full_path_or_raise("clip_vision", clip_name)
|
|
clip_vision = comfy.clip_vision.load(clip_path)
|
|
return (clip_vision,)
|
|
|
|
|
|
class StyleModelLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"), )}, "optional": {"any": (IO.ANY, {})}}
|
|
|
|
RETURN_TYPES = ("STYLE_MODEL",)
|
|
FUNCTION = "load_style_model"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_style_model(self, style_model_name, any=None):
|
|
style_model_path = folder_paths.get_full_path_or_raise("style_models", style_model_name)
|
|
style_model = comfy.sd.load_style_model(style_model_path)
|
|
return (style_model,)
|
|
|
|
|
|
class GLIGENLoader_Any:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "gligen_name": (folder_paths.get_filename_list("gligen"), ), }, "optional": {"any": (IO.ANY, {})}}
|
|
|
|
RETURN_TYPES = ("GLIGEN",)
|
|
FUNCTION = "load_gligen"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_gligen(self, gligen_name, any=None):
|
|
gligen_path = folder_paths.get_full_path_or_raise("gligen", gligen_name)
|
|
gligen = comfy.sd.load_gligen(gligen_path)
|
|
return (gligen,)
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CheckpointLoader_Any": CheckpointLoader_Any,
|
|
"CheckpointLoaderSimple_Any": CheckpointLoaderSimple_Any,
|
|
"unCLIPCheckpointLoader_Any": unCLIPCheckpointLoader_Any,
|
|
"LoraLoader_Any": LoraLoader_Any,
|
|
"LoraLoaderModelOnly_Any": LoraLoaderModelOnly_Any,
|
|
"VAELoader_Any": VAELoader_Any,
|
|
"ControlNetLoader_Any": ControlNetLoader_Any,
|
|
"DiffControlNetLoader_Any": DiffControlNetLoader_Any,
|
|
"UNETLoader_Any": UNETLoader_Any,
|
|
"CLIPLoader_Any": CLIPLoader_Any,
|
|
"DualCLIPLoader_Any": DualCLIPLoader_Any,
|
|
"CLIPVisionLoader_Any": CLIPVisionLoader_Any,
|
|
"StyleModelLoader_Any": StyleModelLoader_Any,
|
|
"GLIGENLoader_Any": GLIGENLoader_Any,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CheckpointLoader_Any": "Load Checkpoint With Config (Any)",
|
|
"CheckpointLoaderSimple_Any": "Load Checkpoint (Any)",
|
|
"unCLIPCheckpointLoader_Any": "Load unCLIP Checkpoint (Any)",
|
|
"LoraLoader_Any": "Load LoRA (Any)",
|
|
"LoraLoaderModelOnly_Any": "Load LoRA for Model Only (Any)",
|
|
"VAELoader_Any": "Load VAE (Any)",
|
|
"ControlNetLoader_Any": "Load ControlNet Model (Any)",
|
|
"DiffControlNetLoader_Any": "Load ControlNet Model (diff) (Any)",
|
|
"UNETLoader_Any": "Load Diffusion Model (Any)",
|
|
"CLIPLoader_Any": "Load CLIP (Any)",
|
|
"DualCLIPLoader_Any": "Load Dual CLIP (Any)",
|
|
"CLIPVisionLoader_Any": "Load CLIP Vision (Any)",
|
|
"StyleModelLoader_Any": "Load Style Model (Any)",
|
|
"GLIGENLoader_Any": "Load GLIGEN (Any)",
|
|
} |