107 lines
4.5 KiB
Python
107 lines
4.5 KiB
Python
import folder_paths
|
|
import torch
|
|
import comfy.sd
|
|
import nodes
|
|
|
|
"""Prepack_Model_DualCLIP: load a diffusion UNet and a dual-CLIP pair in one step."""
|
|
|
|
|
|
class PrepackModelDualCLIP:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"unet_name": (folder_paths.get_filename_list("diffusion_models"), {
|
|
"tooltip": "The name of the diffusion model (UNET) to load."
|
|
}),
|
|
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"], {
|
|
"tooltip": "The weight dtype for the diffusion model."
|
|
}),
|
|
"vae_name": (folder_paths.get_filename_list("vae"), {
|
|
"tooltip": "The VAE to load."
|
|
}),
|
|
"clip_name1": (folder_paths.get_filename_list("text_encoders"), {
|
|
"tooltip": "The name of the first CLIP model to load."
|
|
}),
|
|
"clip_name2": (folder_paths.get_filename_list("text_encoders"), {
|
|
"tooltip": "The name of the second CLIP model to load."
|
|
}),
|
|
"type": (["sdxl", "sd3", "flux", "hunyuan_video"], {
|
|
"tooltip": "The type of CLIP configuration to use."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"device": (["default", "cpu"], {"advanced": True, "tooltip": "Device override for CLIP loading/offloading. Use 'cpu' to force CPU; 'default' lets Comfy manage devices."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
RETURN_NAMES = ("model", "clip", "vae")
|
|
OUTPUT_TOOLTIPS = (
|
|
"The loaded diffusion model (UNET) used for denoising latents.",
|
|
"The loaded dual CLIP model used for encoding text prompts.",
|
|
"The loaded VAE for encoding/decoding latents (can be None if loading fails)."
|
|
)
|
|
FUNCTION = "load_prepack"
|
|
|
|
CATEGORY = "💀Prepack"
|
|
DESCRIPTION = "Load a diffusion UNet and a dual-CLIP pair in one step (SDXL, SD3, FLUX, optional Hunyuan Video)."
|
|
|
|
def load_prepack(self, unet_name, weight_dtype, vae_name, clip_name1, clip_name2, type, device="default"):
|
|
model_options = {}
|
|
try:
|
|
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
|
|
except AttributeError:
|
|
# Fallback silently if torch version doesn't support float8 types
|
|
pass
|
|
|
|
try:
|
|
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)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load diffusion model '{unet_name}': {str(e)}")
|
|
|
|
try:
|
|
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)
|
|
|
|
if type == "sdxl":
|
|
clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION
|
|
elif type == "sd3":
|
|
clip_type = comfy.sd.CLIPType.SD3
|
|
elif type == "flux":
|
|
clip_type = comfy.sd.CLIPType.FLUX
|
|
elif type == "hunyuan_video":
|
|
clip_type = getattr(comfy.sd.CLIPType, "HUNYUAN_VIDEO", comfy.sd.CLIPType.STABLE_DIFFUSION)
|
|
else:
|
|
clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION
|
|
|
|
clip_model_options = {}
|
|
if device == "cpu":
|
|
clip_model_options["load_device"] = clip_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=clip_model_options
|
|
)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load CLIP models '{clip_name1}' and '{clip_name2}': {str(e)}")
|
|
|
|
try:
|
|
vae = nodes.VAELoader().load_vae(vae_name)[0]
|
|
except Exception as e:
|
|
print(f"Warning: Failed to load VAE '{vae_name}': {str(e)}. Returning None.")
|
|
vae = None
|
|
|
|
return (model, clip, vae)
|
|
|
|
|