459 lines
19 KiB
Python
459 lines
19 KiB
Python
import folder_paths
|
|
import torch
|
|
import comfy.samplers
|
|
import comfy.sample
|
|
import comfy.sampler_helpers
|
|
import comfy.model_sampling
|
|
import comfy.latent_formats
|
|
import comfy.sd
|
|
import comfy.clip_vision
|
|
import comfy.supported_models
|
|
from comfy.utils import load_torch_file
|
|
|
|
# Documentation: Self-documenting code
|
|
# Instructions for use: Obvious
|
|
# Expected results: Fork desync
|
|
# adapted from https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py
|
|
|
|
clip_types = ["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "hunyuan_dit", "flux", "mochi", "ltxv", "hunyuan_video", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace"]
|
|
|
|
class BaseModelLoader:
|
|
@staticmethod
|
|
def load_taesd(name):
|
|
sd = {}
|
|
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
|
|
|
encoder = next(filter(lambda a: a.startswith(f"{name}_encoder."), approx_vaes))
|
|
decoder = next(filter(lambda a: a.startswith(f"{name}_decoder."), approx_vaes))
|
|
|
|
enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder))
|
|
for k in enc:
|
|
sd[f"taesd_encoder.{k}"] = enc[k]
|
|
|
|
dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder))
|
|
for k in dec:
|
|
sd[f"taesd_decoder.{k}"] = dec[k]
|
|
|
|
# VAE scale and shift mapping
|
|
vae_params = {
|
|
"taesd": (0.18215, 0.0),
|
|
"taesdxl": (0.13025, 0.0),
|
|
"taesd3": (1.5305, 0.0609),
|
|
"taef1": (0.3611, 0.1159)
|
|
}
|
|
|
|
if name in vae_params:
|
|
scale, shift = vae_params[name]
|
|
sd["vae_scale"] = torch.tensor(scale)
|
|
sd["vae_shift"] = torch.tensor(shift)
|
|
|
|
return sd
|
|
|
|
@staticmethod
|
|
def guess_clip_type(model):
|
|
import comfy.model_base as mb
|
|
|
|
type_map = [
|
|
(mb.SDXLRefiner, "sdxl"),
|
|
(mb.SDXL, "sdxl"),
|
|
(mb.SD15_instructpix2pix, "stable_diffusion"),
|
|
(mb.SDXL_instructpix2pix, "sdxl"),
|
|
(mb.StableCascade_C, "stable_cascade"),
|
|
(mb.StableCascade_B, "stable_cascade"),
|
|
(mb.Flux, "flux"),
|
|
(mb.LTXV, "ltxv"),
|
|
(mb.HunyuanDiT, "hunyuan_dit"),
|
|
(mb.HunyuanVideo, "hunyuan_video"),
|
|
(mb.HunyuanVideoI2V, "hunyuan_video"),
|
|
(mb.HunyuanVideoSkyreelsI2V, "hunyuan_video"),
|
|
(mb.PixArt, "pixart"),
|
|
(mb.CosmosVideo, "cosmos"),
|
|
(mb.Lumina2, "lumina2"),
|
|
(mb.WAN21, "wan"),
|
|
(mb.WAN21_Vace, "wan"),
|
|
(mb.WAN21_Camera, "wan"),
|
|
(mb.HiDream, "hidream"),
|
|
(mb.Chroma, "chroma"),
|
|
(mb.ACEStep, "ace"),
|
|
(mb.SD3, "sd3"),
|
|
(mb.GenmoMochi, "mochi"),
|
|
]
|
|
|
|
for cls, clip_type in type_map:
|
|
if isinstance(model, cls):
|
|
return clip_type.upper()
|
|
|
|
# fallback
|
|
known_types = {
|
|
"stable_diffusion", "stable_cascade", "sd3", "stable_audio", "hunyuan_dit", "flux", "mochi", "ltxv",
|
|
"hunyuan_video", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace"
|
|
}
|
|
|
|
class_name = model.__class__.__name__.lower()
|
|
for t in known_types:
|
|
if t in class_name:
|
|
return t.upper()
|
|
|
|
default_clip_type = "stable_diffusion"
|
|
return default_clip_type.upper()
|
|
|
|
@staticmethod
|
|
def get_model_files():
|
|
return [f for f in folder_paths.get_filename_list("checkpoints") +
|
|
folder_paths.get_filename_list("diffusion_models")
|
|
if f.endswith((".ckpt", ".safetensors", ".sft", ".pt"))]
|
|
|
|
@staticmethod
|
|
def get_weight_options():
|
|
return ["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"]
|
|
|
|
@staticmethod
|
|
def get_clip_options():
|
|
return [".use_ckpt_clip"] + folder_paths.get_filename_list("text_encoders")
|
|
|
|
@staticmethod
|
|
def vae_list():
|
|
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_enc = True
|
|
elif v.startswith("taef1_decoder."):
|
|
f1_taesd_dec = True
|
|
|
|
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")
|
|
return vaes
|
|
|
|
def process_weight_dtype(self, weight_dtype):
|
|
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
|
|
return model_options
|
|
|
|
def load_checkpoint(self, model_name, output_vae, output_clip, model_options):
|
|
try:
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", model_name)
|
|
|
|
out = None
|
|
try:
|
|
out = comfy.sd.load_checkpoint_guess_config(
|
|
ckpt_path,
|
|
output_vae=output_vae,
|
|
output_clip=output_clip,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
model_options=model_options
|
|
)
|
|
except RuntimeError as e:
|
|
if "ERROR: Could not detect model type of:" in str(e):
|
|
error_msg = ""
|
|
if output_vae is True:
|
|
error_msg += "Model/Checkpoint file does not contain a VAE\n"
|
|
if output_clip is True:
|
|
error_msg += "Model/Checkpoint file does not contain a CLIP\n"
|
|
if error_msg != "":
|
|
raise ValueError(error_msg)
|
|
else:
|
|
out = (comfy.sd.load_diffusion_model(ckpt_path, model_options),)
|
|
else:
|
|
raise e
|
|
|
|
return out
|
|
|
|
except FileNotFoundError:
|
|
ckpt_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
|
|
|
|
model = comfy.sd.load_diffusion_model(ckpt_path, model_options=model_options)
|
|
return (model, )
|
|
|
|
|
|
|
|
def load_vae(self, vae_name, ckpt_out):
|
|
if vae_name == ".use_ckpt_vae":
|
|
if ckpt_out[2] is None:
|
|
raise ValueError("Model does not have a VAE")
|
|
return ckpt_out[2]
|
|
elif vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
|
|
sd = self.load_taesd(vae_name)
|
|
return comfy.sd.VAE(sd=sd)
|
|
elif vae_name == ".none":
|
|
return None
|
|
else:
|
|
vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
|
|
sd = comfy.utils.load_torch_file(vae_path)
|
|
return comfy.sd.VAE(sd=sd)
|
|
|
|
|
|
def load_clipvision(ckpt_path):
|
|
sd = load_torch_file(ckpt_path)
|
|
clip_vision = comfy.clip_vision.load(ckpt_path)
|
|
return clip_vision
|
|
|
|
class FluxLoader(BaseModelLoader):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model_name": (s.get_model_files(),),
|
|
"weight_dtype": (s.get_weight_options(),),
|
|
"clip_name1": (s.get_clip_options(),),
|
|
"clip_name2_opt": ([".none"] + folder_paths.get_filename_list("text_encoders"),),
|
|
"vae_name": ([".use_ckpt_vae"] + s.vae_list(),),
|
|
"clip_vision_name": ([".none"] + folder_paths.get_filename_list("clip_vision"),),
|
|
"style_model_name": ([".none"] + folder_paths.get_filename_list("style_models"),),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "CLIP_VISION", "STYLE_MODEL")
|
|
RETURN_NAMES = ("model", "clip", "vae", "clip_vision", "style_model")
|
|
FUNCTION = "main"
|
|
CATEGORY = "RES4LYF/loaders"
|
|
|
|
def main(self, model_name, weight_dtype, clip_name1, clip_name2_opt, vae_name, clip_vision_name, style_model_name):
|
|
model_options = self.process_weight_dtype(weight_dtype)
|
|
|
|
torch.manual_seed(42)
|
|
torch.cuda.manual_seed_all(42)
|
|
|
|
if clip_name1 == ".use_ckpt_clip" and clip_name2_opt != ".none":
|
|
raise ValueError("Cannot specify both \".use_ckpt_clip\" and another clip")
|
|
|
|
output_vae = vae_name == ".use_ckpt_vae"
|
|
output_clip = clip_name1 == ".use_ckpt_clip"
|
|
ckpt_out = self.load_checkpoint(model_name, output_vae, output_clip, model_options)
|
|
|
|
if clip_name1 == ".use_ckpt_clip":
|
|
if ckpt_out[1] is None:
|
|
raise ValueError("Model does not have a clip")
|
|
clip = ckpt_out[1]
|
|
else:
|
|
clip_paths = [folder_paths.get_full_path_or_raise("text_encoders", clip_name1)]
|
|
if clip_name2_opt != ".none":
|
|
clip_paths.append(folder_paths.get_full_path_or_raise("text_encoders", clip_name2_opt))
|
|
clip = comfy.sd.load_clip(clip_paths,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
clip_type=comfy.sd.CLIPType.FLUX)
|
|
|
|
clip_vision = None if clip_vision_name == ".none" else \
|
|
load_clipvision(folder_paths.get_full_path_or_raise("clip_vision", clip_vision_name))
|
|
|
|
style_model = None if style_model_name == ".none" else \
|
|
comfy.sd.load_style_model(folder_paths.get_full_path_or_raise("style_models", style_model_name))
|
|
|
|
vae = self.load_vae(vae_name, ckpt_out)
|
|
|
|
return (ckpt_out[0], clip, vae, clip_vision, style_model)
|
|
|
|
class SD35Loader(BaseModelLoader):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model_name": (s.get_model_files(),),
|
|
"weight_dtype": (s.get_weight_options(),),
|
|
"clip_name1": (s.get_clip_options(),),
|
|
"clip_name2_opt": ([".none"] + folder_paths.get_filename_list("text_encoders"),),
|
|
"clip_name3_opt": ([".none"] + folder_paths.get_filename_list("text_encoders"),),
|
|
"vae_name": ([".use_ckpt_vae"] + folder_paths.get_filename_list("vae") + ["taesd", "taesdxl", "taesd3", "taef1"],),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
RETURN_NAMES = ("model", "clip", "vae")
|
|
FUNCTION = "main"
|
|
CATEGORY = "RES4LYF/loaders"
|
|
|
|
def main(self, model_name, weight_dtype, clip_name1, clip_name2_opt, clip_name3_opt, vae_name):
|
|
model_options = self.process_weight_dtype(weight_dtype)
|
|
|
|
torch.manual_seed(42)
|
|
torch.cuda.manual_seed_all(42)
|
|
|
|
if clip_name1 == ".use_ckpt_clip" and (clip_name2_opt != ".none" or clip_name3_opt != ".none"):
|
|
raise ValueError("Cannot specify both \".use_ckpt_clip\" and another clip")
|
|
|
|
output_vae = vae_name == ".use_ckpt_vae"
|
|
output_clip = clip_name1 == ".use_ckpt_clip"
|
|
ckpt_out = self.load_checkpoint(model_name, output_vae, output_clip, model_options)
|
|
|
|
if clip_name1 == ".use_ckpt_clip":
|
|
if ckpt_out[1] is None:
|
|
raise ValueError("Model does not have a clip")
|
|
clip = ckpt_out[1]
|
|
else:
|
|
clip_paths = [folder_paths.get_full_path_or_raise("text_encoders", clip_name1)]
|
|
for clip_name in [clip_name2_opt, clip_name3_opt]:
|
|
if clip_name != ".none":
|
|
clip_paths.append(folder_paths.get_full_path_or_raise("text_encoders", clip_name))
|
|
clip = comfy.sd.load_clip(clip_paths,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
clip_type=comfy.sd.CLIPType.SD3)
|
|
|
|
vae = self.load_vae(vae_name, ckpt_out)
|
|
|
|
return (ckpt_out[0], clip, vae)
|
|
|
|
class RES4LYFModelLoader(BaseModelLoader):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model_name": (s.get_model_files(),),
|
|
"weight_dtype": (s.get_weight_options(),),
|
|
"clip_name1_opt": ([".none"] + s.get_clip_options(),),
|
|
"clip_name2_opt": ([".none"] + folder_paths.get_filename_list("text_encoders"),),
|
|
"clip_name3_opt": ([".none"] + folder_paths.get_filename_list("text_encoders"),),
|
|
"clip_name4_opt": ([".none"] + folder_paths.get_filename_list("text_encoders"),),
|
|
"clip_type": ([".auto"] + clip_types,),
|
|
"vae_name": ([".none", ".use_ckpt_vae"] + folder_paths.get_filename_list("vae") + ["taesd", "taesdxl", "taesd3", "taef1"],),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
RETURN_NAMES = ("model", "clip", "vae")
|
|
FUNCTION = "main"
|
|
CATEGORY = "RES4LYF/loaders"
|
|
|
|
def main(self, model_name, weight_dtype, clip_name1_opt, clip_name2_opt, clip_name3_opt, clip_name4_opt, clip_type, vae_name):
|
|
model_options = self.process_weight_dtype(weight_dtype)
|
|
|
|
torch.manual_seed(42)
|
|
torch.cuda.manual_seed_all(42)
|
|
|
|
if clip_name1_opt == ".use_ckpt_clip" and (clip_name2_opt != ".none" or clip_name3_opt != ".none" or clip_name4_opt != ".none"):
|
|
raise ValueError("Cannot specify both \".use_ckpt_clip\" and another clip")
|
|
|
|
output_vae = vae_name == ".use_ckpt_vae"
|
|
output_clip = clip_name1_opt == ".use_ckpt_clip"
|
|
ckpt_out = self.load_checkpoint(model_name, output_vae, output_clip, model_options)
|
|
|
|
if clip_name1_opt == ".use_ckpt_clip":
|
|
if ckpt_out[1] is None:
|
|
raise ValueError("Model does not have a clip")
|
|
clip = ckpt_out[1]
|
|
elif clip_name1_opt == ".none":
|
|
clip = None
|
|
else:
|
|
clip_paths = [folder_paths.get_full_path_or_raise("text_encoders", clip_name1_opt)]
|
|
for clip_name in [clip_name2_opt, clip_name3_opt, clip_name4_opt]:
|
|
if clip_name != ".none":
|
|
clip_paths.append(folder_paths.get_full_path_or_raise("text_encoders", clip_name))
|
|
if "auto" in clip_type and ckpt_out[0].model is not None:
|
|
sdCLIPType = getattr(comfy.sd.CLIPType, self.guess_clip_type(ckpt_out[0].model), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
|
else:
|
|
sdCLIPType = getattr(comfy.sd.CLIPType, clip_type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
|
clip = comfy.sd.load_clip(clip_paths,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
clip_type=sdCLIPType)
|
|
|
|
vae = self.load_vae(vae_name, ckpt_out)
|
|
|
|
return (ckpt_out[0], clip, vae)
|
|
|
|
from .style_transfer import Retrojector
|
|
import torch.nn as nn
|
|
|
|
class LayerPatcher:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"embedder": (s.get_model_patches(),),
|
|
"gates": (s.get_model_patches(),),
|
|
"last_layer": (s.get_model_patches(),),
|
|
"dtype": (["bfloat16", "float16", "float32", "float64"], {"default": "float64"}),
|
|
#"retrojector": (s.get_model_patches(),),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "main"
|
|
CATEGORY = "RES4LYF/patchers"
|
|
|
|
@staticmethod
|
|
def get_model_patches():
|
|
return [f for f in folder_paths.get_filename_list("diffusion_models") if f.endswith((".safetensors", ".sft"))]
|
|
|
|
def main(self, model, embedder, gates, last_layer, retrojector=None, dtype="float64"):
|
|
|
|
dtype = getattr(torch, dtype)
|
|
|
|
embedder = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("diffusion_models", embedder))
|
|
last_layer = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("diffusion_models", last_layer))
|
|
#retrojector = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("diffusion_models", retrojector))
|
|
|
|
gates = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("diffusion_models", gates))
|
|
m = model.model.diffusion_model
|
|
|
|
if embedder:
|
|
m.x_embedder.proj = nn.Linear(
|
|
m.x_embedder.proj.in_features,
|
|
m.x_embedder.proj.out_features,
|
|
bias=True,
|
|
device=m.x_embedder.proj.weight.data.device,
|
|
dtype=dtype
|
|
)
|
|
m.x_embedder.proj.weight.data = embedder['x_embedder.proj.weight'].to(dtype).cuda()
|
|
m.x_embedder.proj.bias.data = embedder['x_embedder.proj.bias'].to(dtype).cuda()
|
|
|
|
if gates:
|
|
for key, tensor in gates.items():
|
|
#print(f"Patching {key} with shape {tensor.shape}")
|
|
set_nested_attr(model=m, key=key, value=tensor, dtype=dtype)
|
|
|
|
if last_layer:
|
|
m.final_layer.linear.weight.data = last_layer['final_layer.linear.weight'].to(dtype).cuda()
|
|
m.final_layer.linear.bias.data = last_layer['final_layer.linear.bias'].to(dtype).cuda()
|
|
m.final_layer.adaLN_modulation[1].weight.data = last_layer['final_layer.adaLN_modulation.1.weight'].to(dtype).cuda()
|
|
m.final_layer.adaLN_modulation[1].bias.data = last_layer['final_layer.adaLN_modulation.1.bias'].to(dtype).cuda()
|
|
|
|
#if retrojector:
|
|
# m.Retrojector = Retrojector(model.model.diffusion_model.img_in, pinv_dtype=style_dtype, dtype=style_dtype)
|
|
# m.final_layer.linear.weight.data = last_layer['final_layer.linear.weight']
|
|
# m.final_layer.linear.bias.data = last_layer['final_layer.linear.bias']
|
|
# m.final_layer.adaLN_modulation[1].weight.data = last_layer['final_layer.adaLN_modulation.1.weight']
|
|
# m.final_layer.adaLN_modulation[1].bias.data = last_layer['final_layer.adaLN_modulation.1.bias']
|
|
|
|
return (model,)
|
|
|
|
|
|
|
|
def set_nested_attr(model, key, value, dtype):
|
|
parts = key.split(".")
|
|
attr = model
|
|
for p in parts[:-1]:
|
|
if p.isdigit():
|
|
attr = attr[int(p)]
|
|
else:
|
|
attr = getattr(attr, p)
|
|
getattr(attr, parts[-1]).data.copy_(value.to(getattr(attr, parts[-1]).device, dtype=dtype))
|
|
|
|
|
|
|