|
import json |
|
from pathlib import Path |
|
|
|
import comfy |
|
import comfy.model_management |
|
import comfy.model_patcher |
|
import folder_paths |
|
import safetensors.torch |
|
import torch |
|
from ltx_video.models.autoencoders.causal_video_autoencoder import ( |
|
CausalVideoAutoencoder, |
|
) |
|
from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier |
|
from ltx_video.models.transformers.transformer3d import Transformer3DModel |
|
from safetensors import safe_open |
|
|
|
from .model import LTXVModel, LTXVModelConfig, LTXVTransformer3D |
|
from .nodes_registry import comfy_node |
|
from .vae import LTXVVAE |
|
|
|
|
|
@comfy_node(name="LTXVLoader") |
|
class LTXVLoader: |
|
@classmethod |
|
def INPUT_TYPES(s): |
|
return { |
|
"required": { |
|
"ckpt_name": ( |
|
folder_paths.get_filename_list("checkpoints"), |
|
{"tooltip": "The name of the checkpoint (model) to load."}, |
|
), |
|
"dtype": (["bfloat16", "float32"], {"default": "bfloat16"}), |
|
} |
|
} |
|
|
|
RETURN_TYPES = ("MODEL", "VAE") |
|
RETURN_NAMES = ("model", "vae") |
|
FUNCTION = "load" |
|
CATEGORY = "lightricks/LTXV" |
|
TITLE = "LTXV Loader" |
|
OUTPUT_NODE = False |
|
|
|
def load(self, ckpt_name, dtype): |
|
dtype_map = {"bfloat16": torch.bfloat16, "float32": torch.float32} |
|
load_device = comfy.model_management.get_torch_device() |
|
offload_device = comfy.model_management.unet_offload_device() |
|
|
|
ckpt_path = Path(folder_paths.get_full_path("checkpoints", ckpt_name)) |
|
|
|
vae_config = None |
|
unet_config = None |
|
with safe_open(ckpt_path, framework="pt", device="cpu") as f: |
|
metadata = f.metadata() |
|
if metadata is not None: |
|
config_metadata = metadata.get("config", None) |
|
if config_metadata is not None: |
|
config_metadata = json.loads(config_metadata) |
|
vae_config = config_metadata.get("vae", None) |
|
unet_config = config_metadata.get("transformer", None) |
|
|
|
weights = safetensors.torch.load_file(ckpt_path, device="cpu") |
|
|
|
vae = self._load_vae(weights, vae_config) |
|
num_latent_channels = vae.first_stage_model.config.latent_channels |
|
|
|
model = self._load_unet( |
|
load_device, |
|
offload_device, |
|
weights, |
|
num_latent_channels, |
|
dtype=dtype_map[dtype], |
|
config=unet_config, |
|
) |
|
return (model, vae) |
|
|
|
def _load_vae(self, weights, config=None): |
|
if config is None: |
|
config = { |
|
"_class_name": "CausalVideoAutoencoder", |
|
"dims": 3, |
|
"in_channels": 3, |
|
"out_channels": 3, |
|
"latent_channels": 128, |
|
"blocks": [ |
|
["res_x", 4], |
|
["compress_all", 1], |
|
["res_x_y", 1], |
|
["res_x", 3], |
|
["compress_all", 1], |
|
["res_x_y", 1], |
|
["res_x", 3], |
|
["compress_all", 1], |
|
["res_x", 3], |
|
["res_x", 4], |
|
], |
|
"scaling_factor": 1.0, |
|
"norm_layer": "pixel_norm", |
|
"patch_size": 4, |
|
"latent_log_var": "uniform", |
|
"use_quant_conv": False, |
|
"causal_decoder": False, |
|
} |
|
vae_prefix = "vae." |
|
vae = LTXVVAE.from_config_and_state_dict( |
|
vae_class=CausalVideoAutoencoder, |
|
config=config, |
|
state_dict={ |
|
key.removeprefix(vae_prefix): value |
|
for key, value in weights.items() |
|
if key.startswith(vae_prefix) |
|
}, |
|
) |
|
return vae |
|
|
|
def _load_unet( |
|
self, |
|
load_device, |
|
offload_device, |
|
weights, |
|
num_latent_channels, |
|
dtype, |
|
config=None, |
|
): |
|
if config is None: |
|
config = { |
|
"_class_name": "Transformer3DModel", |
|
"_diffusers_version": "0.25.1", |
|
"_name_or_path": "PixArt-alpha/PixArt-XL-2-256x256", |
|
"activation_fn": "gelu-approximate", |
|
"attention_bias": True, |
|
"attention_head_dim": 64, |
|
"attention_type": "default", |
|
"caption_channels": 4096, |
|
"cross_attention_dim": 2048, |
|
"double_self_attention": False, |
|
"dropout": 0.0, |
|
"in_channels": 128, |
|
"norm_elementwise_affine": False, |
|
"norm_eps": 1e-06, |
|
"norm_num_groups": 32, |
|
"num_attention_heads": 32, |
|
"num_embeds_ada_norm": 1000, |
|
"num_layers": 28, |
|
"num_vector_embeds": None, |
|
"only_cross_attention": False, |
|
"out_channels": 128, |
|
"project_to_2d_pos": True, |
|
"upcast_attention": False, |
|
"use_linear_projection": False, |
|
"qk_norm": "rms_norm", |
|
"standardization_norm": "rms_norm", |
|
"positional_embedding_type": "rope", |
|
"positional_embedding_theta": 10000.0, |
|
"positional_embedding_max_pos": [20, 2048, 2048], |
|
"timestep_scale_multiplier": 1000, |
|
} |
|
|
|
transformer = Transformer3DModel.from_config(config) |
|
unet_prefix = "model.diffusion_model." |
|
transformer.load_state_dict( |
|
{ |
|
key.removeprefix(unet_prefix): value |
|
for key, value in weights.items() |
|
if key.startswith(unet_prefix) |
|
} |
|
) |
|
transformer.to(dtype).to(load_device).eval() |
|
patchifier = SymmetricPatchifier(1) |
|
diffusion_model = LTXVTransformer3D(transformer, patchifier, None, None, None) |
|
model = LTXVModel( |
|
LTXVModelConfig(num_latent_channels, dtype=dtype), |
|
model_type=comfy.model_base.ModelType.FLOW, |
|
device=comfy.model_management.get_torch_device(), |
|
) |
|
model.diffusion_model = diffusion_model |
|
|
|
patcher = comfy.model_patcher.ModelPatcher(model, load_device, offload_device) |
|
|
|
return patcher |
|
|
