146 lines
6.3 KiB
Python
146 lines
6.3 KiB
Python
import os
|
|
import hashlib
|
|
import folder_paths
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import torch
|
|
|
|
def short_hash(path):
|
|
if not os.path.exists(path):
|
|
return "no_file"
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(4096), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()[:10]
|
|
|
|
def load_single_lora(model, clip, lora_name, weight):
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
if not lora_path or not os.path.exists(lora_path):
|
|
return model, clip
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
model, clip = comfy.sd.load_lora_for_models(model, clip, lora, weight, weight)
|
|
return model, clip
|
|
|
|
class CompositeMangoLoader:
|
|
@staticmethod
|
|
def vae_list():
|
|
vaes = folder_paths.get_filename_list("vae")
|
|
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
|
taesd_combinations = {
|
|
"taesd": ("taesd_encoder", "taesd_decoder"),
|
|
"taesdxl": ("taesdxl_encoder", "taesdxl_decoder"),
|
|
"taesd3": ("taesd3_encoder", "taesd3_decoder"),
|
|
"taef1": ("taef1_encoder", "taef1_decoder")
|
|
}
|
|
for name, (enc_prefix, dec_prefix) in taesd_combinations.items():
|
|
if any(v.startswith(enc_prefix) for v in approx_vaes) and any(v.startswith(dec_prefix) for v in approx_vaes):
|
|
vaes.append(name)
|
|
return vaes
|
|
|
|
@staticmethod
|
|
def load_taesd(name):
|
|
sd = {}
|
|
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
|
enc_prefix, dec_prefix = {
|
|
"taesd": ("taesd_encoder", "taesd_decoder"),
|
|
"taesdxl": ("taesdxl_encoder", "taesdxl_decoder"),
|
|
"taesd3": ("taesd3_encoder", "taesd3_decoder"),
|
|
"taef1": ("taef1_encoder", "taef1_decoder")
|
|
}[name]
|
|
|
|
encoder = next(v for v in approx_vaes if v.startswith(enc_prefix))
|
|
decoder = next(v for v in approx_vaes if v.startswith(dec_prefix))
|
|
|
|
enc = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", encoder))
|
|
dec = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", decoder))
|
|
|
|
for k in enc: sd[f"{enc_prefix}.{k}"] = enc[k]
|
|
for k in dec: sd[f"{dec_prefix}.{k}"] = dec[k]
|
|
|
|
sd.update({
|
|
"taesd": (0.18215, 0.0),
|
|
"taesdxl": (0.13025, 0.0),
|
|
"taesd3": (1.5305, 0.0609),
|
|
"taef1": (0.3611, 0.1159)
|
|
}[name])
|
|
sd["vae_scale"], sd["vae_shift"] = torch.tensor(sd["vae_scale"]), torch.tensor(sd["vae_shift"])
|
|
return sd
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"unet_name": (folder_paths.get_filename_list("diffusion_models"),),
|
|
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],),
|
|
"clip_name1": (folder_paths.get_filename_list("text_encoders"),),
|
|
"clip_name2": (folder_paths.get_filename_list("text_encoders"),),
|
|
"type": (["sdxl", "sd3", "flux", "hunyuan_video"],),
|
|
"vae_name": (cls.vae_list(),),
|
|
"LoraName1": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
|
"LoraWeight1": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
|
"LoraName2": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
|
"LoraWeight2": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
|
"LoraName3": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
|
"LoraWeight3": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
|
"LoraName4": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
|
"LoraWeight4": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
|
"LoraName5": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
|
"LoraWeight5": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
|
},
|
|
"optional": {
|
|
"device": (["default", "cpu"], {"advanced": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "LORA_STACK", "STRING", "STRING")
|
|
RETURN_NAMES = ("model", "clip", "vae", "lora_stack", "unet_name", "unet_hash")
|
|
FUNCTION = "load_all"
|
|
CATEGORY = "Mango Node Pack/Loaders"
|
|
|
|
def load_all(self, unet_name, weight_dtype, clip_name1, clip_name2, type, vae_name, device="default", **kwargs):
|
|
# Load UNET
|
|
model_options = {}
|
|
if weight_dtype.startswith("fp8"):
|
|
model_options["dtype"] = getattr(torch, f"float{weight_dtype[2:]}")
|
|
if "fast" in weight_dtype:
|
|
model_options["fp8_optimizations"] = True
|
|
|
|
unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)
|
|
unet_hash = short_hash(unet_path)
|
|
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
|
|
|
# Load Dual CLIP
|
|
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)
|
|
clip_type = getattr(comfy.sd.CLIPType, type.upper())
|
|
|
|
model_options_clip = {}
|
|
if device == "cpu":
|
|
model_options_clip.update({"load_device": "cpu", "offload_device": "cpu"})
|
|
|
|
clip = comfy.sd.load_clip(
|
|
ckpt_paths=[clip_path1, clip_path2],
|
|
clip_type=clip_type,
|
|
model_options=model_options_clip
|
|
)
|
|
|
|
# Load VAE
|
|
if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
|
|
vae_sd = self.load_taesd(vae_name)
|
|
else:
|
|
vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
|
|
vae_sd = comfy.utils.load_torch_file(vae_path)
|
|
vae = comfy.sd.VAE(sd=vae_sd)
|
|
|
|
# Apply LORAs
|
|
lora_stack = []
|
|
for i in range(1, 6):
|
|
lora_name = kwargs.get(f"LoraName{i}")
|
|
lora_weight = kwargs.get(f"LoraWeight{i}", 1.0)
|
|
if lora_name and lora_name != "None":
|
|
model, clip = load_single_lora(model, clip, lora_name, lora_weight)
|
|
lora_stack.append((lora_name, lora_weight, lora_weight))
|
|
|
|
return (model, clip, vae, lora_stack, unet_name, unet_hash)
|