diff --git a/__init__.py b/__init__.py index 3403d19..7e0fb79 100644 --- a/__init__.py +++ b/__init__.py @@ -5,6 +5,8 @@ from .nodes.context.even_larger_context import TT_EvenLargerContext from .nodes.context.large_context_flux import TT_LargeContextFlux from .nodes.context.large_context_sdxl import TT_LargeContextSdxl from .nodes.latent.latent_factory import TT_LatentFactory +from .nodes.loaders.sdxl_models_loader import TT_SdxlModelsLoader +from .nodes.loaders.sdxl_models_loader_advanced import TT_SdxlModelsLoaderAdvanced from .nodes.workflow.flux_workfow_settings import TT_FluxWorkflowSettings from .nodes.workflow.flux_workfow_settings_advanced import TT_FluxWorkflowSettingsAdvanced from .nodes.workflow.sdxl_workfow_settings import TT_SdxlWorkflowSettings @@ -21,6 +23,9 @@ NODE_CLASS_MAPPINGS = { "TT_ContextPassthrough": TT_ContextPassthrough, # Latent "TT_LatentFactory": TT_LatentFactory, + # Loaders + "TT_SdxlModelsLoader": TT_SdxlModelsLoader, + "TT_SdxlModelsLoaderAdvanced": TT_SdxlModelsLoaderAdvanced, # Workflow "TT_FluxWorkflowSettings": TT_FluxWorkflowSettings, "TT_FluxWorkflowSettingsAdvanced": TT_FluxWorkflowSettingsAdvanced, @@ -39,6 +44,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "TT_ContextPassthrough": "TT Context Passthrough", # Latent "TT_LatentFactory": "TT Latent Factory", + # Loaders + "TT_SdxlModelsLoader": "TT SDXL Models Loader", + "TT_SdxlModelsLoaderAdvanced": "TT SDXL Models Loader (Advanced)", # Workflow "TT_FluxWorkflowSettings": "TT FLUX Workflow Settings", "TT_FluxWorkflowSettingsAdvanced": "TT FLUX Workflow Settings (Advanced)", diff --git a/nodes/loaders/loader_helpers.py b/nodes/loaders/loader_helpers.py new file mode 100644 index 0000000..35cc747 --- /dev/null +++ b/nodes/loaders/loader_helpers.py @@ -0,0 +1,77 @@ +import os + +import comfy +import folder_paths +import torch + +from nodes import VAELoader + +VIDEO_TAES = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5"] +IMAGE_TAES = ["taesd", "taesdxl", "taesd3", "taef1"] + + +def load_checkpoint(ckpt_name): + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) + return comfy.sd.load_checkpoint_guess_config( + ckpt_path, output_vae=False, output_clip=False, + embedding_directory=folder_paths.get_folder_paths("embeddings") + ) + + +def load_clip(clip_l, clip_g, device): + clip_type = getattr(comfy.sd.CLIPType, "SDXL", comfy.sd.CLIPType.STABLE_DIFFUSION) + clip_l_path = folder_paths.get_full_path_or_raise("text_encoders", clip_l) + clip_g_path = folder_paths.get_full_path_or_raise("text_encoders", clip_g) + + model_options = {} + if device == "cpu": + model_options["load_device"] = model_options["offload_device"] = torch.device("cpu") + + return comfy.sd.load_clip( + ckpt_paths=[clip_l_path, clip_g_path], + embedding_directory=folder_paths.get_folder_paths("embeddings"), + clip_type=clip_type, + model_options=model_options + ) + + +def apply_lora(loaded_lora, model, clip, lora_name, strength): + if loaded_lora == None: + lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + else: + lora = loaded_lora + + patched_model, patched_clip = comfy.sd.load_lora_for_models(model, clip, lora, strength, strength) + + return (patched_model, patched_clip, lora) + + +def load_vae(vae_name, vae_device, vae_dtype): + dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[vae_dtype] + + device = None + if vae_device == "default": + device = comfy.model_management.get_torch_device() + elif vae_device == "cpu": + device = torch.device("cpu") + + metadata = None + if vae_name == "pixel_space": + state_dict = {} + state_dict["pixel_space_vae"] = torch.tensor(1.0) + + elif vae_name in IMAGE_TAES: + state_dict = VAELoader.load_taesd(vae_name) + + else: + if os.path.splitext(vae_name)[0] in 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) + state_dict, metadata = comfy.utils.load_torch_file(vae_path, return_metadata=True) + + vae = comfy.sd.VAE(sd=state_dict, device=device, dtype=dtype, metadata=metadata) + vae.throw_exception_if_invalid() + + return vae diff --git a/nodes/loaders/sdxl_models_loader.py b/nodes/loaders/sdxl_models_loader.py new file mode 100644 index 0000000..20f02bb --- /dev/null +++ b/nodes/loaders/sdxl_models_loader.py @@ -0,0 +1,41 @@ +import folder_paths + +from nodes import VAELoader +from .loader_helpers import load_checkpoint, load_clip, load_vae + + +class TT_SdxlModelsLoader: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "primary_ckpt": (folder_paths.get_filename_list("checkpoints"),), + "secondary_ckpt": (["None"] + folder_paths.get_filename_list("checkpoints"),), + "primary_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "clip_l": (folder_paths.get_filename_list("text_encoders"),), + "clip_g": (folder_paths.get_filename_list("text_encoders"),), + "clip_device": (["default", "cpu"], {"advanced": True}), + "vae_name": (VAELoader.vae_list(VAELoader),) + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "VAE",) + RETURN_NAMES = ("MODEL", "CLIP", "VAE",) + FUNCTION = "load_models" + CATEGORY = "TenserTensor/Loaders/SDXL" + + def load_models(self, primary_ckpt, secondary_ckpt, primary_weight, clip_l, clip_g, clip_device, vae_name): + model = load_checkpoint(primary_ckpt)[0] + + if secondary_ckpt != "None": + secondary_model = load_checkpoint(secondary_ckpt)[0] + model = model.clone() + key_patches = secondary_model.get_key_patches("diffusion_model.") + for key in key_patches: + model.add_patches({key: key_patches[key]}, 1.0 - primary_weight, primary_weight) + + clip = load_clip(clip_l, clip_g, clip_device) + + vae = load_vae(vae_name, "cpu", "bfloat16") + + return (model, clip, vae,) diff --git a/nodes/loaders/sdxl_models_loader_advanced.py b/nodes/loaders/sdxl_models_loader_advanced.py new file mode 100644 index 0000000..9d41fdb --- /dev/null +++ b/nodes/loaders/sdxl_models_loader_advanced.py @@ -0,0 +1,76 @@ +import folder_paths + +from nodes import VAELoader +from .loader_helpers import load_checkpoint, load_clip, apply_lora, load_vae + + +class TT_SdxlModelsLoaderAdvanced: + def __init__(self): + self.loaded_lora_1 = None + self.loaded_lora_2 = None + self.loaded_lora_3 = None + self.loaded_lora_4 = None + + @classmethod + def INPUT_TYPES(cls): + DEVICES = ["default", "cpu"] + DTYPES = ["bfloat16", "float16", "float32"] + + return { + "required": { + "primary_ckpt": (folder_paths.get_filename_list("checkpoints"),), + "secondary_ckpt": (["None"] + folder_paths.get_filename_list("checkpoints"),), + "primary_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "clip_l": (folder_paths.get_filename_list("text_encoders"),), + "clip_g": (folder_paths.get_filename_list("text_encoders"),), + "clip_device": (DEVICES, {"advanced": True}), + "lora_name_1": (["None"] + folder_paths.get_filename_list("loras"),), + "strength_1": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.1}), + "lora_name_2": (["None"] + folder_paths.get_filename_list("loras"),), + "strength_2": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.1}), + "lora_name_3": (["None"] + folder_paths.get_filename_list("loras"),), + "strength_3": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.1}), + "lora_name_4": (["None"] + folder_paths.get_filename_list("loras"),), + "strength_4": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.1}), + "vae_name": (VAELoader.vae_list(VAELoader),), + "vae_device": (DEVICES, {"advanced": True}), + "vae_dtype": (DTYPES, {"advanced": True}), + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "VAE",) + RETURN_NAMES = ("MODEL", "CLIP", "VAE",) + FUNCTION = "load_models" + CATEGORY = "TenserTensor/Loaders/SDXL" + + def load_models(self, primary_ckpt, secondary_ckpt, primary_weight, clip_l, clip_g, clip_device, + lora_name_1, strength_1, lora_name_2, strength_2, lora_name_3, strength_3, lora_name_4, strength_4, + vae_name, vae_device, vae_dtype): + + model = load_checkpoint(primary_ckpt)[0] + + if secondary_ckpt != "None": + secondary_model = load_checkpoint(secondary_ckpt)[0] + model = model.clone() + key_patches = secondary_model.get_key_patches("diffusion_model.") + for key in key_patches: + model.add_patches({key: key_patches[key]}, 1.0 - primary_weight, primary_weight) + + clip = load_clip(clip_l, clip_g, clip_device) + + loras = [ + (lora_name_1, strength_1, 'loaded_lora_1'), + (lora_name_2, strength_2, 'loaded_lora_2'), + (lora_name_3, strength_3, 'loaded_lora_3'), + (lora_name_4, strength_4, 'loaded_lora_4'), + ] + + for name, strength, attr in loras: + if name != "None" and strength != 0: + cached = getattr(self, attr) + model, clip, lora = apply_lora(cached, model, clip, name, strength) + setattr(self, attr, lora) + + vae = load_vae(vae_name, vae_device, vae_dtype) + + return (model, clip, vae,) diff --git a/web/docs/TT_SdxlModelsLoader.md b/web/docs/TT_SdxlModelsLoader.md new file mode 100644 index 0000000..bfd1051 --- /dev/null +++ b/web/docs/TT_SdxlModelsLoader.md @@ -0,0 +1,32 @@ +## TT_SdxlModelsLoader + +All-in-one model loader for SDXL workflows with checkpoint merging support. + +### Inputs + +| Parameter | Type | Required | Description | +|------------------|--------------|----------|----------------------------------------------------------------| +| `primary_ckpt` | CHECKPOINT | Yes | Primary SDXL checkpoint to load | +| `secondary_ckpt` | CHECKPOINT | Yes | Secondary checkpoint to merge (select "None" to skip) | +| `primary_weight` | FLOAT | Yes | Weight for primary checkpoint in merge (0.0-1.0, default: 1.0) | +| `clip_l` | TEXT_ENCODER | Yes | CLIP-L text encoder model | +| `clip_g` | TEXT_ENCODER | Yes | CLIP-G text encoder model | +| `clip_device` | STRING | Yes | Device for CLIP ("default" or "cpu") - advanced | +| `vae_name` | VAE | Yes | VAE model to load | + +### Outputs + +| Parameter | Type | Description | +|-----------|-------|-------------------------------------------| +| `MODEL` | MODEL | Loaded (and optionally merged) SDXL model | +| `CLIP` | CLIP | Combined CLIP-L and CLIP-G text encoders | +| `VAE` | VAE | Loaded VAE | + +### Usage + +Use this node to load all SDXL model components in one place. Optionally merge two checkpoints by selecting a secondary +checkpoint and adjusting the primary_weight (1.0 = 100% primary, 0.5 = 50/50 mix, 0.0 = 100% secondary). + +--- + +[View on GitHub](https://github.com/tenser-tensor/ComfyUI-TenserTensor) diff --git a/web/docs/TT_SdxlModelsLoaderAdvanced.md b/web/docs/TT_SdxlModelsLoaderAdvanced.md new file mode 100644 index 0000000..bfe7799 --- /dev/null +++ b/web/docs/TT_SdxlModelsLoaderAdvanced.md @@ -0,0 +1,43 @@ +## TT_SdxlModelsLoaderAdvanced + +Extended all-in-one model loader for SDXL with checkpoint merging and LoRA support. + +### Inputs + +| Parameter | Type | Required | Description | +|------------------|--------------|----------|-----------------------------------------------------------------| +| `primary_ckpt` | CHECKPOINT | Yes | Primary SDXL checkpoint to load | +| `secondary_ckpt` | CHECKPOINT | Yes | Secondary checkpoint to merge (select "None" to skip) | +| `primary_weight` | FLOAT | Yes | Weight for primary checkpoint in merge (0.0-1.0, default: 1.0) | +| `clip_l` | TEXT_ENCODER | Yes | CLIP-L text encoder model | +| `clip_g` | TEXT_ENCODER | Yes | CLIP-G text encoder model | +| `clip_device` | STRING | Yes | Device for CLIP ("default" or "cpu") - advanced | +| `lora_name_1` | LORA | Yes | First LoRA to apply (select "None" to skip) | +| `strength_1` | FLOAT | Yes | Strength for first LoRA (-10.0 to 10.0, default: 1.0) | +| `lora_name_2` | LORA | Yes | Second LoRA to apply (select "None" to skip) | +| `strength_2` | FLOAT | Yes | Strength for second LoRA (-10.0 to 10.0, default: 1.0) | +| `lora_name_3` | LORA | Yes | Third LoRA to apply (select "None" to skip) | +| `strength_3` | FLOAT | Yes | Strength for third LoRA (-10.0 to 10.0, default: 1.0) | +| `lora_name_4` | LORA | Yes | Fourth LoRA to apply (select "None" to skip) | +| `strength_4` | FLOAT | Yes | Strength for fourth LoRA (-10.0 to 10.0, default: 1.0) | +| `vae_name` | VAE | Yes | VAE model to load | +| `vae_device` | STRING | Yes | Device for VAE ("default" or "cpu") - advanced | +| `vae_dtype` | STRING | Yes | Data type for VAE ("bfloat16", "float16", "float32") - advanced | + +### Outputs + +| Parameter | Type | Description | +|-----------|-------|-------------------------------------------------------------| +| `MODEL` | MODEL | Loaded model with applied LoRAs and optional merge | +| `CLIP` | CLIP | Combined CLIP-L and CLIP-G text encoders with applied LoRAs | +| `VAE` | VAE | Loaded VAE | + +### Usage + +Extended version of the SDXL loader with support for up to 4 LoRAs and additional VAE configuration options. Load all +model components, optionally merge checkpoints, apply LoRAs sequentially, and configure device/dtype settings for CLIP +and VAE. LoRAs are cached for efficiency. + +--- + +[View on GitHub](https://github.com/tenser-tensor/ComfyUI-TenserTensor)