1451 lines
58 KiB
Python
1451 lines
58 KiB
Python
from functools import partial
|
|
import json
|
|
import os
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import torchaudio
|
|
from transformers import Gemma3Config
|
|
import yaml
|
|
from toolkit.config_modules import GenerateImageConfig, ModelConfig
|
|
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
|
|
from toolkit.dto import DTO
|
|
from toolkit.models.base_model import BaseModel
|
|
from toolkit.basic import flush
|
|
from toolkit.prompt_utils import PromptEmbeds
|
|
from toolkit.samplers.custom_flowmatch_sampler import (
|
|
CustomFlowMatchEulerDiscreteScheduler,
|
|
)
|
|
from accelerate import init_empty_weights
|
|
from toolkit.accelerator import unwrap_model
|
|
from toolkit.paths import MODELS_PATH
|
|
from toolkit.util.mixed_precision import attach_per_op_casting, pin_stored_fp32
|
|
from safetensors.torch import load_file
|
|
from PIL import Image
|
|
import huggingface_hub
|
|
|
|
try:
|
|
from diffusers import LTX2Pipeline, LTX2ImageToVideoPipeline
|
|
from diffusers.pipelines.ltx2.export_utils import encode_video
|
|
from transformers import (
|
|
Gemma3ForConditionalGeneration,
|
|
GemmaTokenizerFast,
|
|
)
|
|
from toolkit.models.v2.diffusion_models.ltx2 import (
|
|
LTX2TextConnectors,
|
|
LTX2VideoTransformer3DModel,
|
|
LTX2Vocoder,
|
|
LTX2VocoderWithBWE,
|
|
)
|
|
from toolkit.models.v2.vae.ltx2 import (
|
|
LTX2AudioVAE as AutoencoderKLLTX2Audio,
|
|
LTX2VideoVAE as AutoencoderKLLTX2Video,
|
|
)
|
|
from toolkit.models.v2.text_encoders.gemma3 import Gemma3TextEncoder
|
|
from .convert_ltx2_to_diffusers import (
|
|
get_model_state_dict_from_combined_ckpt,
|
|
convert_ltx2_transformer,
|
|
convert_ltx2_video_vae,
|
|
convert_ltx2_audio_vae,
|
|
convert_ltx2_vocoder,
|
|
convert_ltx2_connectors,
|
|
split_transformer_and_connector_state_dict,
|
|
dequantize_state_dict,
|
|
convert_comfy_gemma3_to_transformers,
|
|
convert_lora_original_to_diffusers,
|
|
convert_lora_diffusers_to_original,
|
|
)
|
|
except ImportError as e:
|
|
print("Diffusers import error:", e)
|
|
raise ImportError(
|
|
"Diffusers is out of date. Update diffusers to the latest version by doing pip uninstall diffusers and then pip install -r requirements.txt"
|
|
)
|
|
|
|
|
|
scheduler_config = {
|
|
"base_image_seq_len": 1024,
|
|
"base_shift": 0.95,
|
|
"invert_sigmas": False,
|
|
"max_image_seq_len": 4096,
|
|
"max_shift": 2.05,
|
|
"num_train_timesteps": 1000,
|
|
"shift": 1.0,
|
|
"shift_terminal": 0.1,
|
|
"stochastic_sampling": False,
|
|
"time_shift_type": "exponential",
|
|
"use_beta_sigmas": False,
|
|
"use_dynamic_shifting": True,
|
|
"use_exponential_sigmas": False,
|
|
"use_karras_sigmas": False,
|
|
}
|
|
|
|
dit_prefix = "model.diffusion_model."
|
|
vae_prefix = "vae."
|
|
audio_vae_prefix = "audio_vae."
|
|
vocoder_prefix = "vocoder."
|
|
base_te_path = "Lightricks/gemma-3-12b-it-qat-q4_0-unquantized"
|
|
|
|
HF_TOKEN = os.getenv("HF_TOKEN", None)
|
|
|
|
|
|
def new_save_image_function(
|
|
self: GenerateImageConfig,
|
|
image, # will contain a dict that can be dumped ditectly into encode_video, just add output_path to it.
|
|
count: int = 0,
|
|
max_count: int = 0,
|
|
**kwargs,
|
|
):
|
|
# this replaces gen image config save image function so we can save the video with sound from ltx2
|
|
image["output_path"] = self.get_image_path(count, max_count)
|
|
# make sample directory if it does not exist
|
|
os.makedirs(os.path.dirname(image["output_path"]), exist_ok=True)
|
|
encode_video(**image)
|
|
flush()
|
|
|
|
|
|
def blank_log_image_function(self, *args, **kwargs):
|
|
# todo handle wandb logging of videos with audio
|
|
return
|
|
|
|
|
|
class ComboVae(torch.nn.Module):
|
|
"""Combines video and audio VAEs for joint encoding and decoding."""
|
|
|
|
def __init__(
|
|
self,
|
|
vae: AutoencoderKLLTX2Video,
|
|
audio_vae: AutoencoderKLLTX2Audio,
|
|
) -> None:
|
|
super().__init__()
|
|
self.vae = vae
|
|
self.audio_vae = audio_vae
|
|
|
|
@property
|
|
def device(self):
|
|
return self.vae.device
|
|
|
|
@property
|
|
def dtype(self):
|
|
return self.vae.dtype
|
|
|
|
@property
|
|
def config(self):
|
|
return self.vae.config
|
|
|
|
def encode(
|
|
self,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
return self.vae.encode(*args, **kwargs)
|
|
|
|
def decode(
|
|
self,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
return self.vae.decode(*args, **kwargs)
|
|
|
|
|
|
class AudioProcessor(torch.nn.Module):
|
|
"""Converts audio waveforms to log-mel spectrograms with optional resampling."""
|
|
|
|
def __init__(
|
|
self,
|
|
sample_rate: int,
|
|
mel_bins: int,
|
|
mel_hop_length: int,
|
|
n_fft: int,
|
|
) -> None:
|
|
super().__init__()
|
|
self.sample_rate = sample_rate
|
|
self.mel_transform = torchaudio.transforms.MelSpectrogram(
|
|
sample_rate=sample_rate,
|
|
n_fft=n_fft,
|
|
win_length=n_fft,
|
|
hop_length=mel_hop_length,
|
|
f_min=0.0,
|
|
f_max=sample_rate / 2.0,
|
|
n_mels=mel_bins,
|
|
window_fn=torch.hann_window,
|
|
center=True,
|
|
pad_mode="reflect",
|
|
power=1.0,
|
|
mel_scale="slaney",
|
|
norm="slaney",
|
|
)
|
|
|
|
def resample_waveform(
|
|
self,
|
|
waveform: torch.Tensor,
|
|
source_rate: int,
|
|
target_rate: int,
|
|
) -> torch.Tensor:
|
|
"""Resample waveform to target sample rate if needed."""
|
|
if source_rate == target_rate:
|
|
return waveform
|
|
resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
|
|
return resampled.to(device=waveform.device, dtype=waveform.dtype)
|
|
|
|
def waveform_to_mel(
|
|
self,
|
|
waveform: torch.Tensor,
|
|
waveform_sample_rate: int,
|
|
) -> torch.Tensor:
|
|
"""Convert waveform to log-mel spectrogram [batch, channels, time, n_mels]."""
|
|
waveform = self.resample_waveform(
|
|
waveform, waveform_sample_rate, self.sample_rate
|
|
)
|
|
|
|
mel = self.mel_transform(waveform)
|
|
mel = torch.log(torch.clamp(mel, min=1e-5))
|
|
|
|
mel = mel.to(device=waveform.device, dtype=waveform.dtype)
|
|
return mel.permute(0, 1, 3, 2).contiguous()
|
|
|
|
|
|
class LTX2Model(BaseModel):
|
|
arch = "ltx2"
|
|
ltx_version = "2.0"
|
|
ltx_te_path = None
|
|
|
|
def __init__(
|
|
self,
|
|
device,
|
|
model_config: ModelConfig,
|
|
dtype="bf16",
|
|
custom_pipeline=None,
|
|
noise_scheduler=None,
|
|
**kwargs,
|
|
):
|
|
super().__init__(
|
|
device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
|
|
)
|
|
self.is_flow_matching = True
|
|
self.is_transformer = True
|
|
self.target_lora_modules = ["LTX2VideoTransformer3DModel"]
|
|
# defines if the model supports model paths. Only some will
|
|
self.supports_model_paths = True
|
|
# use the new format on this new model by default
|
|
self.use_old_lokr_format = False
|
|
self.audio_processor = None
|
|
|
|
# gemma needs left side padding
|
|
self.te_padding_side = "left"
|
|
|
|
# loss mask for i2v conditioning (1 = train, 0 = conditioned token), set per step in get_noise_prediction
|
|
self._i2v_loss_mask = None
|
|
|
|
# invalidate older caches
|
|
self.latent_space_version = f"{self.arch}_v2"
|
|
|
|
# static method to get the noise scheduler
|
|
@staticmethod
|
|
def get_train_scheduler():
|
|
return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
|
|
|
|
def get_bucket_divisibility(self):
|
|
return 32
|
|
|
|
def load_model(self):
|
|
dtype = self.torch_dtype
|
|
self.print_and_status_update("Loading LTX2 model")
|
|
model_path = self.model_config.name_or_path
|
|
base_model_path = self.model_config.extras_name_or_path
|
|
|
|
combined_state_dict = None
|
|
|
|
self.print_and_status_update("Loading transformer")
|
|
|
|
if not os.path.exists(model_path) and model_path.endswith(".safetensors"):
|
|
# download the model from the Hugging Face Hub if it is not a local path
|
|
splits = model_path.split("/")
|
|
if len(splits) < 3:
|
|
raise ValueError(
|
|
f"Invalid model path: {model_path}. Must be in the format 'repo_id/repo/filename.safetensors' or 'repo_id/repo/subfolder/filename.safetensors' to download from the Hugging Face Hub."
|
|
)
|
|
rel_path = "/".join(splits[2:])
|
|
# use the file from the models folder if it is already there
|
|
local_candidates = [
|
|
os.path.join(MODELS_PATH, rel_path),
|
|
os.path.join(MODELS_PATH, splits[-1]),
|
|
]
|
|
for candidate in local_candidates:
|
|
if os.path.exists(candidate):
|
|
model_path = candidate
|
|
break
|
|
else:
|
|
# download the model from the hub into the models folder
|
|
model_path = huggingface_hub.hf_hub_download(
|
|
repo_id="/".join(splits[:2]),
|
|
filename=rel_path,
|
|
token=HF_TOKEN,
|
|
local_dir=MODELS_PATH,
|
|
)
|
|
|
|
# if we have a safetensors file it is a mono checkpoint
|
|
if os.path.exists(model_path) and model_path.endswith(".safetensors"):
|
|
combined_state_dict = load_file(model_path)
|
|
combined_state_dict = dequantize_state_dict(combined_state_dict)
|
|
|
|
if combined_state_dict is not None:
|
|
original_dit_ckpt = get_model_state_dict_from_combined_ckpt(
|
|
combined_state_dict, dit_prefix
|
|
)
|
|
transformer = convert_ltx2_transformer(
|
|
original_dit_ckpt, version=self.ltx_version
|
|
)
|
|
transformer = transformer.to(dtype)
|
|
# the transformer holds these tensors (assign=True); drop the dict refs
|
|
# so each block's bf16 original frees as it quantizes instead of the
|
|
# whole dit staying in RAM. Connector keys stay — converted later.
|
|
transformer_sd, _ = split_transformer_and_connector_state_dict(
|
|
original_dit_ckpt
|
|
)
|
|
for key in transformer_sd:
|
|
combined_state_dict.pop(dit_prefix + key, None)
|
|
del transformer_sd, original_dit_ckpt
|
|
flush()
|
|
else:
|
|
if os.path.exists(model_path):
|
|
# check if the path is a full checkpoint.
|
|
te_folder_path = os.path.join(model_path, "text_encoder")
|
|
# if we have the te, this folder is a full checkpoint, use it as the base
|
|
if os.path.exists(te_folder_path):
|
|
base_model_path = model_path
|
|
|
|
transformer = LTX2VideoTransformer3DModel.load_model(model_path, dtype=dtype)
|
|
|
|
# quantize + offload + placement, all driven by model_config
|
|
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
|
|
|
|
flush()
|
|
|
|
self.print_and_status_update("Loading text encoder")
|
|
if (
|
|
self.model_config.te_name_or_path is not None
|
|
and self.model_config.te_name_or_path.endswith(".safetensors")
|
|
):
|
|
# load from comfyui gemma3 checkpoint
|
|
tokenizer = GemmaTokenizerFast.from_pretrained(base_te_path)
|
|
|
|
with init_empty_weights():
|
|
text_encoder = Gemma3TextEncoder(
|
|
Gemma3Config(
|
|
**{
|
|
"boi_token_index": 255999,
|
|
"bos_token_id": 2,
|
|
"eoi_token_index": 256000,
|
|
"eos_token_id": 106,
|
|
"image_token_index": 262144,
|
|
"initializer_range": 0.02,
|
|
"mm_tokens_per_image": 256,
|
|
"model_type": "gemma3",
|
|
"pad_token_id": 0,
|
|
"text_config": {
|
|
"attention_bias": False,
|
|
"attention_dropout": 0.0,
|
|
"attn_logit_softcapping": None,
|
|
"cache_implementation": "hybrid",
|
|
"final_logit_softcapping": None,
|
|
"head_dim": 256,
|
|
"hidden_activation": "gelu_pytorch_tanh",
|
|
"hidden_size": 3840,
|
|
"initializer_range": 0.02,
|
|
"intermediate_size": 15360,
|
|
"max_position_embeddings": 131072,
|
|
"model_type": "gemma3_text",
|
|
"num_attention_heads": 16,
|
|
"num_hidden_layers": 48,
|
|
"num_key_value_heads": 8,
|
|
"query_pre_attn_scalar": 256,
|
|
"rms_norm_eps": 1e-06,
|
|
"rope_local_base_freq": 10000,
|
|
"rope_scaling": {"factor": 8.0, "rope_type": "linear"},
|
|
"rope_theta": 1000000,
|
|
"sliding_window": 1024,
|
|
"sliding_window_pattern": 6,
|
|
"torch_dtype": "bfloat16",
|
|
"use_cache": True,
|
|
"vocab_size": 262208,
|
|
},
|
|
"torch_dtype": "bfloat16",
|
|
"transformers_version": "4.51.3",
|
|
"unsloth_fixed": True,
|
|
"vision_config": {
|
|
"attention_dropout": 0.0,
|
|
"hidden_act": "gelu_pytorch_tanh",
|
|
"hidden_size": 1152,
|
|
"image_size": 896,
|
|
"intermediate_size": 4304,
|
|
"layer_norm_eps": 1e-06,
|
|
"model_type": "siglip_vision_model",
|
|
"num_attention_heads": 16,
|
|
"num_channels": 3,
|
|
"num_hidden_layers": 27,
|
|
"patch_size": 14,
|
|
"torch_dtype": "bfloat16",
|
|
"vision_use_head": False,
|
|
},
|
|
}
|
|
)
|
|
)
|
|
te_state_dict = load_file(self.model_config.te_name_or_path)
|
|
te_state_dict = convert_comfy_gemma3_to_transformers(te_state_dict)
|
|
for key in te_state_dict:
|
|
te_state_dict[key] = te_state_dict[key].to(dtype)
|
|
|
|
text_encoder.load_state_dict(te_state_dict, assign=True, strict=True)
|
|
del te_state_dict
|
|
flush()
|
|
elif self.model_config.te_name_or_path is not None:
|
|
# a repo or folder
|
|
tokenizer = GemmaTokenizerFast.from_pretrained(
|
|
self.model_config.te_name_or_path
|
|
)
|
|
text_encoder = Gemma3TextEncoder.load_model(
|
|
self.model_config.te_name_or_path, dtype=dtype, subfolder=""
|
|
)
|
|
elif self.ltx_te_path is not None:
|
|
# pull from model specific te
|
|
tokenizer = GemmaTokenizerFast.from_pretrained(self.ltx_te_path)
|
|
text_encoder = Gemma3TextEncoder.load_model(
|
|
self.ltx_te_path, dtype=dtype, subfolder=""
|
|
)
|
|
else:
|
|
# using combo hf repo
|
|
tokenizer = GemmaTokenizerFast.from_pretrained(
|
|
self.model_config.name_or_path, subfolder="tokenizer"
|
|
)
|
|
text_encoder = Gemma3TextEncoder.load_model(
|
|
self.model_config.name_or_path, dtype=dtype
|
|
)
|
|
|
|
# remove the vision tower
|
|
text_encoder.model.vision_tower = None
|
|
flush()
|
|
|
|
# quantize + offload + placement, all driven by model_config
|
|
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
|
|
text_encoder.to(self.device_torch, dtype=dtype)
|
|
flush()
|
|
|
|
self.print_and_status_update("Loading VAEs and other components")
|
|
if combined_state_dict is not None:
|
|
original_vae_ckpt = get_model_state_dict_from_combined_ckpt(
|
|
combined_state_dict, vae_prefix
|
|
)
|
|
vae = convert_ltx2_video_vae(
|
|
original_vae_ckpt, version=self.ltx_version
|
|
).to(dtype)
|
|
del original_vae_ckpt
|
|
original_audio_vae_ckpt = get_model_state_dict_from_combined_ckpt(
|
|
combined_state_dict, audio_vae_prefix
|
|
)
|
|
audio_vae = convert_ltx2_audio_vae(
|
|
original_audio_vae_ckpt, version=self.ltx_version
|
|
).to(dtype)
|
|
del original_audio_vae_ckpt
|
|
original_connectors_ckpt = get_model_state_dict_from_combined_ckpt(
|
|
combined_state_dict, dit_prefix
|
|
)
|
|
connectors = convert_ltx2_connectors(
|
|
original_connectors_ckpt, version=self.ltx_version
|
|
).to(dtype)
|
|
del original_connectors_ckpt
|
|
original_vocoder_ckpt = get_model_state_dict_from_combined_ckpt(
|
|
combined_state_dict, vocoder_prefix
|
|
)
|
|
vocoder = convert_ltx2_vocoder(
|
|
original_vocoder_ckpt, version=self.ltx_version
|
|
).to(dtype)
|
|
del original_vocoder_ckpt
|
|
del combined_state_dict
|
|
flush()
|
|
else:
|
|
vae = AutoencoderKLLTX2Video.load_model(base_model_path, dtype=dtype)
|
|
audio_vae = AutoencoderKLLTX2Audio.load_model(base_model_path, dtype=dtype)
|
|
|
|
connectors = LTX2TextConnectors.load_model(base_model_path, dtype=dtype)
|
|
|
|
vocoder_cls = LTX2Vocoder
|
|
if self.ltx_version in ("2.3", "2.5"):
|
|
vocoder_cls = LTX2VocoderWithBWE
|
|
|
|
vocoder = vocoder_cls.load_model(base_model_path, dtype=dtype)
|
|
|
|
self.noise_scheduler = LTX2Model.get_train_scheduler()
|
|
|
|
self.print_and_status_update("Making pipe")
|
|
|
|
pipe: LTX2Pipeline = LTX2Pipeline(
|
|
scheduler=self.noise_scheduler,
|
|
vae=vae,
|
|
audio_vae=audio_vae,
|
|
text_encoder=None,
|
|
tokenizer=tokenizer,
|
|
connectors=connectors,
|
|
transformer=None,
|
|
vocoder=vocoder,
|
|
)
|
|
# for quantization, it works best to do these after making the pipe
|
|
pipe.text_encoder = text_encoder
|
|
pipe.transformer = transformer
|
|
|
|
self.print_and_status_update("Preparing Model")
|
|
|
|
text_encoder = [pipe.text_encoder]
|
|
tokenizer = [pipe.tokenizer]
|
|
|
|
# leave it on cpu for now
|
|
if not self.low_vram:
|
|
pipe.transformer = pipe.transformer.to(self.device_torch)
|
|
|
|
flush()
|
|
# low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
|
|
# to the gpu on demand
|
|
if not self.low_vram:
|
|
text_encoder[0].to(self.device_torch)
|
|
text_encoder[0].requires_grad_(False)
|
|
text_encoder[0].eval()
|
|
flush()
|
|
|
|
# save it to the model class
|
|
self.vae = ComboVae(pipe.vae, pipe.audio_vae)
|
|
self.text_encoder = text_encoder # list of text encoders
|
|
self.tokenizer = tokenizer # list of tokenizers
|
|
self.model = pipe.transformer
|
|
self.pipeline = pipe
|
|
|
|
self.audio_processor = AudioProcessor(
|
|
sample_rate=pipe.audio_sampling_rate,
|
|
mel_bins=audio_vae.config.mel_bins,
|
|
mel_hop_length=pipe.audio_hop_length,
|
|
n_fft=1024, # todo get this from vae if we can, I couldnt find it.
|
|
).to(self.device_torch, dtype=torch.float32)
|
|
|
|
self.print_and_status_update("Model Loaded")
|
|
|
|
@torch.no_grad()
|
|
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
|
|
if device is None:
|
|
device = self.vae_device_torch
|
|
if dtype is None:
|
|
dtype = self.vae_torch_dtype
|
|
|
|
if self.pipeline.vae.device == torch.device("cpu"):
|
|
self.pipeline.vae.to(device)
|
|
self.pipeline.vae.eval()
|
|
self.pipeline.vae.requires_grad_(False)
|
|
|
|
if self.model_config.low_vram:
|
|
self.pipeline.vae.tile_sample_min_num_frames = 64
|
|
self.pipeline.vae.tile_sample_stride_num_frames = 16
|
|
# they check the wrong flat on encode currently so set both to future proof
|
|
self.pipeline.vae.use_framewise_decoding = True
|
|
self.pipeline.vae.use_framewise_encoding = True
|
|
|
|
image_list = [image.to(device, dtype=dtype) for image in image_list]
|
|
|
|
# Normalize shapes
|
|
norm_images = []
|
|
for image in image_list:
|
|
if image.ndim == 3:
|
|
# (C, H, W) -> (C, 1, H, W)
|
|
norm_images.append(image.unsqueeze(1))
|
|
elif image.ndim == 4:
|
|
# (T, C, H, W) -> (C, T, H, W)
|
|
norm_images.append(image.permute(1, 0, 2, 3))
|
|
else:
|
|
raise ValueError(f"Invalid image shape: {image.shape}")
|
|
|
|
# Stack to (B, C, T, H, W)
|
|
images = torch.stack(norm_images)
|
|
|
|
latents = self.pipeline.vae.encode(images).latent_dist.sample()
|
|
|
|
# Normalize latents across the channel dimension [B, C, F, H, W]
|
|
scaling_factor = 1.0
|
|
latents_mean = self.pipeline.vae.latents_mean.view(1, -1, 1, 1, 1).to(
|
|
latents.device, latents.dtype
|
|
)
|
|
latents_std = self.pipeline.vae.latents_std.view(1, -1, 1, 1, 1).to(
|
|
latents.device, latents.dtype
|
|
)
|
|
latents = (latents - latents_mean) * scaling_factor / latents_std
|
|
|
|
if self.model_config.low_vram:
|
|
self.pipeline.vae.use_framewise_decoding = False
|
|
self.pipeline.vae.use_framewise_encoding = False
|
|
|
|
return latents.to(device, dtype=dtype)
|
|
|
|
def get_generation_pipeline(self):
|
|
scheduler = LTX2Model.get_train_scheduler()
|
|
|
|
pipeline: LTX2Pipeline = LTX2Pipeline(
|
|
scheduler=scheduler,
|
|
vae=unwrap_model(self.pipeline.vae),
|
|
audio_vae=unwrap_model(self.pipeline.audio_vae),
|
|
text_encoder=None,
|
|
tokenizer=unwrap_model(self.pipeline.tokenizer),
|
|
connectors=unwrap_model(self.pipeline.connectors),
|
|
transformer=None,
|
|
vocoder=unwrap_model(self.pipeline.vocoder),
|
|
)
|
|
pipeline.transformer = unwrap_model(self.model)
|
|
pipeline.text_encoder = unwrap_model(self.text_encoder[0])
|
|
|
|
pipeline = pipeline.to(self.device_torch)
|
|
|
|
return pipeline
|
|
|
|
def generate_single_image(
|
|
self,
|
|
pipeline: LTX2Pipeline,
|
|
gen_config: GenerateImageConfig,
|
|
conditional_embeds: PromptEmbeds,
|
|
unconditional_embeds: PromptEmbeds,
|
|
generator: torch.Generator,
|
|
extra: dict,
|
|
):
|
|
if self.model.device == torch.device("cpu"):
|
|
self.model.to(self.device_torch)
|
|
|
|
# handle control image
|
|
if gen_config.ctrl_img is not None:
|
|
# switch to image to video pipeline
|
|
pipeline = LTX2ImageToVideoPipeline(
|
|
scheduler=pipeline.scheduler,
|
|
vae=pipeline.vae,
|
|
audio_vae=pipeline.audio_vae,
|
|
text_encoder=pipeline.text_encoder,
|
|
tokenizer=pipeline.tokenizer,
|
|
connectors=pipeline.connectors,
|
|
transformer=pipeline.transformer,
|
|
vocoder=pipeline.vocoder,
|
|
)
|
|
|
|
is_video = gen_config.num_frames > 1
|
|
# override the generate single image to handle video + audio generation
|
|
if is_video:
|
|
gen_config._orig_save_image_function = gen_config.save_image
|
|
gen_config.save_image = partial(new_save_image_function, gen_config)
|
|
gen_config.log_image = partial(blank_log_image_function, gen_config)
|
|
# set output extension to mp4
|
|
gen_config.output_ext = "mp4"
|
|
|
|
# reactivate progress bar since this is slooooow
|
|
pipeline.set_progress_bar_config(disable=False)
|
|
pipeline = pipeline.to(self.device_torch)
|
|
|
|
# make sure dimensions are valid
|
|
bd = self.get_bucket_divisibility()
|
|
gen_config.height = (gen_config.height // bd) * bd
|
|
gen_config.width = (gen_config.width // bd) * bd
|
|
|
|
# handle control image
|
|
if gen_config.ctrl_img is not None:
|
|
control_img = Image.open(gen_config.ctrl_img).convert("RGB")
|
|
# resize the control image
|
|
control_img = control_img.resize(
|
|
(gen_config.width, gen_config.height), Image.LANCZOS
|
|
)
|
|
# add the control image to the extra dict
|
|
extra["image"] = control_img
|
|
|
|
# frames must be divisible by 8 then + 1. so 1, 9, 17, 25, etc.
|
|
if gen_config.num_frames != 1:
|
|
if (gen_config.num_frames - 1) % 8 != 0:
|
|
gen_config.num_frames = ((gen_config.num_frames - 1) // 8) * 8 + 1
|
|
|
|
if self.low_vram:
|
|
# set vae to tile decode
|
|
# pipeline.vae.enable_tiling(
|
|
# tile_sample_min_height=256,
|
|
# tile_sample_min_width=256,
|
|
# tile_sample_min_num_frames=8,
|
|
# tile_sample_stride_height=224,
|
|
# tile_sample_stride_width=224,
|
|
# tile_sample_stride_num_frames=4,
|
|
# )
|
|
self.pipeline.vae.tile_sample_min_num_frames = 16
|
|
self.pipeline.vae.tile_sample_stride_num_frames = 8
|
|
self.pipeline.vae.use_framewise_decoding = True
|
|
|
|
# We only encode and store the minimum prompt tokens, but need them padded to 1024 for LTX2
|
|
conditional_embeds = self.pad_embeds(conditional_embeds)
|
|
unconditional_embeds = self.pad_embeds(unconditional_embeds)
|
|
|
|
if self.ltx_version in ("2.3", "2.5"):
|
|
extra["stg_scale"] = 1.0
|
|
extra["modality_scale"] = 3.0
|
|
extra["guidance_rescale"] = 0.7
|
|
extra["audio_guidance_scale"] = 7.0
|
|
extra["audio_stg_scale"] = 1.0
|
|
extra["audio_modality_scale"] = 3.0
|
|
extra["audio_guidance_rescale"] = 0.7
|
|
extra["spatio_temporal_guidance_blocks"] = [28]
|
|
extra["use_cross_timestep"] = (
|
|
True # they dont set this in some examples in diffusers, but I believe it should always be true for 2.3
|
|
)
|
|
|
|
video, audio = pipeline(
|
|
prompt_embeds=conditional_embeds.text_embeds.to(
|
|
self.device_torch, dtype=self.torch_dtype
|
|
),
|
|
prompt_attention_mask=conditional_embeds.attention_mask.to(
|
|
self.device_torch
|
|
),
|
|
negative_prompt_embeds=unconditional_embeds.text_embeds.to(
|
|
self.device_torch, dtype=self.torch_dtype
|
|
),
|
|
negative_prompt_attention_mask=unconditional_embeds.attention_mask.to(
|
|
self.device_torch
|
|
),
|
|
height=gen_config.height,
|
|
width=gen_config.width,
|
|
num_inference_steps=gen_config.num_inference_steps,
|
|
guidance_scale=gen_config.guidance_scale,
|
|
latents=gen_config.latents,
|
|
num_frames=gen_config.num_frames,
|
|
generator=generator,
|
|
return_dict=False,
|
|
output_type="np" if is_video else "pil",
|
|
**extra,
|
|
)
|
|
if self.low_vram:
|
|
# Restore no tiling
|
|
# pipeline.vae.use_tiling = False
|
|
self.pipeline.vae.use_framewise_decoding = False
|
|
|
|
if is_video:
|
|
# redurn as a dict, we will handle it with an override function
|
|
video = (video * 255).round().astype("uint8")
|
|
video = torch.from_numpy(video)
|
|
return {
|
|
"video": video[0],
|
|
"fps": gen_config.fps,
|
|
"audio": audio[0].float().cpu(),
|
|
"audio_sample_rate": pipeline.vocoder.config.output_sampling_rate, # should be 24000
|
|
"output_path": None,
|
|
}
|
|
else:
|
|
# shape = [1, frames, channels, height, width]
|
|
# make sure this is right
|
|
video = video[0] # list of pil images
|
|
audio = audio[0] # tensor
|
|
if gen_config.num_frames > 1:
|
|
return video # return the frames.
|
|
else:
|
|
# get just the first image
|
|
img = video[0]
|
|
return img
|
|
|
|
def encode_audio(self, audio_data_list):
|
|
# audio_date_list is a list of {"waveform": waveform[C, L], "sample_rate": int(sample_rate)}
|
|
if self.pipeline.audio_vae.device == torch.device("cpu"):
|
|
self.pipeline.audio_vae.to(self.device_torch)
|
|
|
|
output_tensor = None
|
|
audio_num_frames = None
|
|
|
|
# do them seperatly for now
|
|
for audio_data in audio_data_list:
|
|
waveform = audio_data["waveform"].to(
|
|
device=self.device_torch, dtype=torch.float32
|
|
)
|
|
sample_rate = audio_data["sample_rate"]
|
|
|
|
# Add batch dimension if needed: [channels, samples] -> [batch, channels, samples]
|
|
if waveform.dim() == 2:
|
|
waveform = waveform.unsqueeze(0)
|
|
|
|
if waveform.shape[1] == 1:
|
|
# make sure it is stereo
|
|
waveform = waveform.repeat(1, 2, 1)
|
|
|
|
# Convert waveform to mel spectrogram using AudioProcessor
|
|
mel_spectrogram = self.audio_processor.waveform_to_mel(
|
|
waveform, waveform_sample_rate=sample_rate
|
|
)
|
|
mel_spectrogram = mel_spectrogram.to(dtype=self.torch_dtype)
|
|
|
|
# Encode mel spectrogram to latents
|
|
latents = self.pipeline.audio_vae.encode(
|
|
mel_spectrogram.to(self.device_torch, dtype=self.torch_dtype)
|
|
).latent_dist.sample()
|
|
|
|
if audio_num_frames is None:
|
|
audio_num_frames = latents.shape[2] # (latents is [B, C, T, F])
|
|
|
|
packed_latents = self.pipeline._pack_audio_latents(
|
|
latents,
|
|
# patch_size=self.pipeline.transformer.config.audio_patch_size,
|
|
# patch_size_t=self.pipeline.transformer.config.audio_patch_size_t,
|
|
) # [B, L, C * M]
|
|
if output_tensor is None:
|
|
output_tensor = packed_latents
|
|
else:
|
|
output_tensor = torch.cat([output_tensor, packed_latents], dim=0)
|
|
|
|
# normalize latents, opposite of (latents * latents_std) + latents_mean
|
|
latents_mean = self.pipeline.audio_vae.latents_mean
|
|
latents_std = self.pipeline.audio_vae.latents_std
|
|
output_tensor = (output_tensor - latents_mean) / latents_std
|
|
return output_tensor
|
|
|
|
def pad_embeds(self, embeds: PromptEmbeds):
|
|
# ltx-2 connector requires 1024 tokens for good results. Any smaller and it degrades.
|
|
target_length = 1024
|
|
current_length = embeds.text_embeds.shape[1]
|
|
if current_length < target_length:
|
|
pad_length = target_length - current_length
|
|
pad_tensor = torch.zeros(
|
|
(embeds.text_embeds.shape[0], pad_length, embeds.text_embeds.shape[2]),
|
|
device=embeds.text_embeds.device,
|
|
dtype=embeds.text_embeds.dtype,
|
|
)
|
|
embeds.text_embeds = torch.cat([pad_tensor, embeds.text_embeds], dim=1)
|
|
if embeds.attention_mask is not None:
|
|
pad_mask = torch.zeros(
|
|
(embeds.attention_mask.shape[0], pad_length),
|
|
device=embeds.attention_mask.device,
|
|
dtype=embeds.attention_mask.dtype,
|
|
)
|
|
embeds.attention_mask = torch.cat(
|
|
[pad_mask, embeds.attention_mask], dim=1
|
|
)
|
|
return embeds
|
|
|
|
def get_noise_prediction(
|
|
self,
|
|
latent_model_input: torch.Tensor,
|
|
timestep: torch.Tensor, # 0 to 1000 scale
|
|
text_embeddings: PromptEmbeds,
|
|
batch: "DataLoaderBatchDTO" = None,
|
|
**kwargs,
|
|
):
|
|
audio_target = None
|
|
with torch.no_grad():
|
|
if self.model.device == torch.device("cpu"):
|
|
self.model.to(self.device_torch)
|
|
|
|
# We only encode and store the minimum prompt tokens, but need them padded to 1024 for LTX2
|
|
text_embeddings = self.pad_embeds(text_embeddings)
|
|
|
|
batch_size, C, latent_num_frames, latent_height, latent_width = (
|
|
latent_model_input.shape
|
|
)
|
|
|
|
video_timestep = timestep.clone()
|
|
self._i2v_loss_mask = None
|
|
|
|
# i2v from first frame
|
|
if batch.dataset_config.do_i2v and batch.num_frames > 1:
|
|
# check to see if we had it cached
|
|
if batch.first_frame_latents is not None:
|
|
init_latents = batch.first_frame_latents.to(
|
|
self.device_torch, dtype=self.torch_dtype
|
|
)
|
|
else:
|
|
# extract the first frame and encode it
|
|
# videos come in (bs, num_frames, channels, height, width)
|
|
# images come in (bs, channels, height, width)
|
|
frames = batch.tensor
|
|
if len(frames.shape) == 4:
|
|
first_frames = frames
|
|
elif len(frames.shape) == 5:
|
|
first_frames = frames[:, 0]
|
|
else:
|
|
raise ValueError(f"Unknown frame shape {frames.shape}")
|
|
# first frame doesnt have time dim, add it back
|
|
init_latents = self.encode_images(
|
|
first_frames, device=self.device_torch, dtype=self.torch_dtype
|
|
)
|
|
|
|
# expand the latents to match video frames
|
|
init_latents = init_latents.repeat(1, 1, latent_num_frames, 1, 1)
|
|
mask_shape = (
|
|
batch_size,
|
|
1,
|
|
latent_num_frames,
|
|
latent_height,
|
|
latent_width,
|
|
)
|
|
# First condition is image latents and those should be kept clean.
|
|
conditioning_mask = torch.zeros(
|
|
mask_shape, device=self.device_torch, dtype=self.torch_dtype
|
|
)
|
|
conditioning_mask[:, :, 0] = 1.0
|
|
|
|
# use conditioning mask to replace latents
|
|
latent_model_input = (
|
|
init_latents * conditioning_mask
|
|
+ latent_model_input * (1 - conditioning_mask)
|
|
)
|
|
|
|
# conditioned tokens are clean with timestep 0 and their prediction is
|
|
# discarded at inference, so they must not contribute to the loss
|
|
self._i2v_loss_mask = 1.0 - conditioning_mask
|
|
|
|
packed_conditioning_mask = self.pipeline._pack_latents(
|
|
conditioning_mask,
|
|
patch_size=self.pipeline.transformer_spatial_patch_size,
|
|
patch_size_t=self.pipeline.transformer_temporal_patch_size,
|
|
).squeeze(-1)
|
|
|
|
# set video timestep
|
|
video_timestep = timestep.unsqueeze(-1) * (1 - packed_conditioning_mask)
|
|
|
|
frame_rate = batch.dataset_config.fps
|
|
# check frame dimension
|
|
# Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
|
|
packed_latents = self.pipeline._pack_latents(
|
|
latent_model_input,
|
|
patch_size=self.pipeline.transformer_spatial_patch_size,
|
|
patch_size_t=self.pipeline.transformer_temporal_patch_size,
|
|
)
|
|
|
|
# audio only trains for video batches from datasets that asked for
|
|
# it. Cached latents can carry audio after do_audio was turned off,
|
|
# and image (single frame) batches must never pick up a soundtrack.
|
|
do_audio = (
|
|
batch.dataset_config is not None
|
|
and batch.dataset_config.do_audio
|
|
and getattr(batch, "num_frames", 1) > 1
|
|
)
|
|
if do_audio and (
|
|
batch.audio_latents is not None or batch.audio_tensor is not None
|
|
):
|
|
if batch.audio_latents is not None:
|
|
# we have audio latents cached
|
|
raw_audio_latents = batch.audio_latents.to(
|
|
self.device_torch, dtype=self.torch_dtype
|
|
)
|
|
else:
|
|
# we have audio waveforms to encode
|
|
# use audio from the batch if available
|
|
raw_audio_latents = self.encode_audio(batch.audio_data)
|
|
|
|
audio_num_frames = raw_audio_latents.shape[1]
|
|
# the audio noise is drawn once per step and shared by every
|
|
# pass (prior, primary, cfg/guidance, preservation) so they all
|
|
# see the same soundtrack and every pass's target matches. It
|
|
# rides on the latents DTO.
|
|
audio_noise = (
|
|
batch.latents.get("audio_noise")
|
|
if isinstance(batch.latents, DTO)
|
|
else None
|
|
)
|
|
if (
|
|
audio_noise is not None
|
|
and audio_noise.shape == raw_audio_latents.shape
|
|
):
|
|
audio_noise = audio_noise.to(
|
|
raw_audio_latents.device, dtype=raw_audio_latents.dtype
|
|
)
|
|
else:
|
|
audio_noise = torch.randn_like(raw_audio_latents)
|
|
if batch.latents is not None:
|
|
batch.latents = DTO(batch.latents, audio_noise=audio_noise)
|
|
audio_target = (audio_noise - raw_audio_latents).detach()
|
|
audio_latents = self.add_noise(
|
|
raw_audio_latents,
|
|
audio_noise,
|
|
timestep,
|
|
).to(self.device_torch, dtype=self.torch_dtype)
|
|
else:
|
|
# no audio
|
|
num_mel_bins = self.pipeline.audio_vae.config.mel_bins
|
|
# latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
|
|
num_channels_latents_audio = (
|
|
self.pipeline.audio_vae.config.latent_channels
|
|
)
|
|
duration_s = batch.num_frames / frame_rate
|
|
audio_latents_per_second = (
|
|
self.pipeline.audio_sampling_rate
|
|
/ self.pipeline.audio_hop_length
|
|
/ float(self.pipeline.audio_vae_temporal_compression_ratio)
|
|
)
|
|
audio_num_frames = round(duration_s * audio_latents_per_second)
|
|
audio_latents = self.pipeline.prepare_audio_latents(
|
|
batch_size,
|
|
num_channels_latents=num_channels_latents_audio,
|
|
audio_latent_length=audio_num_frames,
|
|
num_mel_bins=num_mel_bins,
|
|
noise_scale=0.0,
|
|
dtype=torch.float32,
|
|
device=self.transformer.device,
|
|
generator=None,
|
|
latents=None,
|
|
)
|
|
|
|
if self.pipeline.connectors.device != self.transformer.device:
|
|
self.pipeline.connectors.to(self.transformer.device)
|
|
|
|
# Padding side for default Gemma3-12B text encoder
|
|
tokenizer_padding_side = "left"
|
|
if getattr(self, "tokenizer", None) is not None:
|
|
tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
|
|
(
|
|
connector_prompt_embeds,
|
|
connector_audio_prompt_embeds,
|
|
connector_attention_mask,
|
|
) = self.pipeline.connectors(
|
|
text_embeddings.text_embeds,
|
|
text_embeddings.attention_mask.to(self.transformer.dtype),
|
|
padding_side=tokenizer_padding_side,
|
|
)
|
|
|
|
# compute video and audio positional ids
|
|
video_coords = self.transformer.rope.prepare_video_coords(
|
|
packed_latents.shape[0],
|
|
latent_num_frames,
|
|
latent_height,
|
|
latent_width,
|
|
packed_latents.device,
|
|
fps=frame_rate,
|
|
)
|
|
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
|
|
audio_latents.shape[0], audio_num_frames, audio_latents.device
|
|
)
|
|
|
|
# use_cross_timestep - Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
|
|
# calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
|
|
# `False` is the legacy LTX-2.0 behavior.
|
|
use_cross_timestep = self.ltx_version in ("2.3", "2.5")
|
|
|
|
noise_pred_video, noise_pred_audio = self.transformer(
|
|
hidden_states=packed_latents,
|
|
audio_hidden_states=audio_latents.to(self.transformer.dtype),
|
|
encoder_hidden_states=connector_prompt_embeds,
|
|
audio_encoder_hidden_states=connector_audio_prompt_embeds,
|
|
timestep=video_timestep,
|
|
sigma=timestep, # Used by LTX-2.3
|
|
audio_timestep=timestep,
|
|
encoder_attention_mask=connector_attention_mask,
|
|
audio_encoder_attention_mask=connector_attention_mask,
|
|
num_frames=latent_num_frames,
|
|
height=latent_height,
|
|
width=latent_width,
|
|
fps=frame_rate,
|
|
audio_num_frames=audio_num_frames,
|
|
video_coords=video_coords,
|
|
audio_coords=audio_coords,
|
|
isolate_modalities=False,
|
|
spatio_temporal_guidance_blocks=None,
|
|
perturbation_mask=None,
|
|
use_cross_timestep=use_cross_timestep,
|
|
attention_kwargs=None,
|
|
return_dict=False,
|
|
)
|
|
|
|
unpacked_output = self.pipeline._unpack_latents(
|
|
latents=noise_pred_video,
|
|
num_frames=latent_num_frames,
|
|
height=latent_height,
|
|
width=latent_width,
|
|
patch_size=self.pipeline.transformer_spatial_patch_size,
|
|
patch_size_t=self.pipeline.transformer_temporal_patch_size,
|
|
)
|
|
|
|
if audio_target is not None:
|
|
# every pass's DTO carries its own audio stream and target
|
|
return DTO(
|
|
unpacked_output,
|
|
audio=noise_pred_audio,
|
|
audio_target=audio_target,
|
|
)
|
|
return unpacked_output
|
|
|
|
def get_prompt_embeds(self, prompt: str) -> PromptEmbeds:
|
|
if self.pipeline.text_encoder.device != self.device_torch:
|
|
self.pipeline.text_encoder.to(self.device_torch)
|
|
|
|
device = self.device_torch
|
|
scale_factor = 8
|
|
batch_size = len(prompt)
|
|
# Gemma expects left padding for chat-style prompts
|
|
self.tokenizer[0].padding_side = "left"
|
|
if self.tokenizer[0].pad_token is None:
|
|
self.tokenizer[0].pad_token = self.tokenizer[0].eos_token
|
|
|
|
prompt = [p.strip() for p in prompt]
|
|
text_inputs = self.tokenizer[0](
|
|
prompt,
|
|
# padding="max_length",
|
|
padding="longest",
|
|
max_length=1024,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt",
|
|
)
|
|
text_input_ids = text_inputs.input_ids
|
|
prompt_attention_mask = text_inputs.attention_mask
|
|
|
|
text_input_ids = text_input_ids.to(device)
|
|
prompt_attention_mask = prompt_attention_mask.to(device)
|
|
|
|
text_encoder_outputs = self.text_encoder[0](
|
|
input_ids=text_input_ids,
|
|
attention_mask=prompt_attention_mask,
|
|
output_hidden_states=True,
|
|
)
|
|
text_encoder_hidden_states = text_encoder_outputs.hidden_states
|
|
text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1)
|
|
prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(
|
|
dtype=self.torch_dtype
|
|
) # Pack to 3D
|
|
|
|
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
|
_, seq_len, _ = prompt_embeds.shape
|
|
prompt_embeds = prompt_embeds.repeat(1, 1, 1)
|
|
prompt_embeds = prompt_embeds.view(batch_size * 1, seq_len, -1)
|
|
|
|
prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
|
|
prompt_attention_mask = prompt_attention_mask.repeat(1, 1)
|
|
|
|
pe = PromptEmbeds([prompt_embeds, None])
|
|
pe.attention_mask = prompt_attention_mask
|
|
return pe
|
|
|
|
def get_model_has_grad(self):
|
|
return False
|
|
|
|
def get_te_has_grad(self):
|
|
return False
|
|
|
|
def save_model(self, output_path, meta, save_dtype):
|
|
transformer: LTX2VideoTransformer3DModel = unwrap_model(self.model)
|
|
transformer.save_pretrained(
|
|
save_directory=os.path.join(output_path, "transformer"),
|
|
safe_serialization=True,
|
|
)
|
|
|
|
meta_path = os.path.join(output_path, "aitk_meta.yaml")
|
|
with open(meta_path, "w") as f:
|
|
yaml.dump(meta, f)
|
|
|
|
def get_loss_target(self, *args, **kwargs):
|
|
noise = kwargs.get("noise")
|
|
batch = kwargs.get("batch")
|
|
return (noise - batch.latents).detach()
|
|
|
|
def scale_loss(self, loss):
|
|
# zero out the loss on i2v conditioned tokens, renormalized so the loss
|
|
# magnitude matches unconditioned batches (masked mean)
|
|
if self._i2v_loss_mask is not None:
|
|
loss_mask = self._i2v_loss_mask.to(loss.device, dtype=loss.dtype)
|
|
loss = loss * loss_mask / loss_mask.mean().clamp(min=1e-8)
|
|
self._i2v_loss_mask = None
|
|
return loss
|
|
|
|
def get_base_model_version(self):
|
|
return "ltx2"
|
|
|
|
def get_transformer_block_names(self) -> Optional[List[str]]:
|
|
return ["transformer_blocks"]
|
|
|
|
lora_keys_use_comfy_prefix = True
|
|
|
|
def convert_lora_weights_before_save(self, state_dict):
|
|
state_dict = super().convert_lora_weights_before_save(state_dict)
|
|
return convert_lora_diffusers_to_original(state_dict, version=self.ltx_version)
|
|
|
|
def convert_lora_weights_before_load(self, state_dict):
|
|
state_dict = convert_lora_original_to_diffusers(
|
|
state_dict, version=self.ltx_version
|
|
)
|
|
return super().convert_lora_weights_before_load(state_dict)
|
|
|
|
|
|
class LTX23Model(LTX2Model):
|
|
arch = "ltx2.3"
|
|
ltx_version = "2.3"
|
|
ltx_te_path = base_te_path
|
|
|
|
|
|
# LTX-2.5 ships as ComfyUI-style split files (no diffusers folders, no mono
|
|
# checkpoint). Files are used in place when present under MODELS_PATH and
|
|
# downloaded to exactly these locations only when missing, so the models
|
|
# folder stays shareable with a ComfyUI install. The int8 ConvRot files are
|
|
# the defaults; bf16 variants stay selectable via model_kwargs overrides.
|
|
COMFY_LTX25_REPO = "Lightricks/LTX-2.5"
|
|
COMFY_LTX25_FILES = {
|
|
"dit": "diffusion_models/ltx-2.5-22b-dev-transformer-comfy-int8-convrot.safetensors",
|
|
"text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors",
|
|
# the "-conv-" file is the classic conv VAE; the default 2.5 vae file is a
|
|
# new diffusion-decoder VAE that diffusers has no class for
|
|
"video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors",
|
|
# bundles the BWE vocoder alongside the audio VAE
|
|
"audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
|
|
}
|
|
|
|
|
|
class LTX25Model(LTX2Model):
|
|
arch = "ltx2.5"
|
|
ltx_version = "2.5"
|
|
ltx_te_path = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# ComfyUI-style file resolution (toolkit/models/v2/resolver.py)
|
|
# ------------------------------------------------------------------
|
|
def _resolve_comfy_file(self, component: str) -> str:
|
|
from toolkit.models.v2.resolver import (
|
|
repo_id_from_name_or_path,
|
|
resolve_comfy_file,
|
|
)
|
|
|
|
return resolve_comfy_file(
|
|
COMFY_LTX25_FILES[component],
|
|
repo_id=repo_id_from_name_or_path(
|
|
self.model_config.name_or_path, COMFY_LTX25_REPO
|
|
),
|
|
override_path=self.model_config.model_kwargs.get(
|
|
f"{component}_path", None
|
|
),
|
|
hf_token=HF_TOKEN,
|
|
status_fn=self.print_and_status_update,
|
|
)
|
|
|
|
def _resolve_named_file(self, path: str, component: str) -> str:
|
|
from toolkit.models.v2.resolver import resolve_named_file
|
|
|
|
return resolve_named_file(path, component=component, hf_token=HF_TOKEN)
|
|
|
|
def _resolve_dit_path(self) -> str:
|
|
name_or_path = self.model_config.name_or_path
|
|
if name_or_path and name_or_path.endswith(".safetensors"):
|
|
return self._resolve_named_file(name_or_path, "transformer")
|
|
return self._resolve_comfy_file("dit")
|
|
|
|
def _resolve_te_path(self) -> str:
|
|
te_name_or_path = self.model_config.te_name_or_path
|
|
if te_name_or_path is not None:
|
|
return self._resolve_named_file(te_name_or_path, "text encoder")
|
|
return self._resolve_comfy_file("text_encoder")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Loading
|
|
# ------------------------------------------------------------------
|
|
def _load_quantized_module(self, module, state_dict, name: str) -> int:
|
|
"""Attach pre-quantized (int8 ConvRot) linears onto the toolkit's
|
|
quantization backends and load the rest of the (meta-built) module
|
|
from the state dict. Works unchanged for bf16 checkpoints, where no
|
|
quant markers exist and everything strict-loads."""
|
|
from toolkit.util.comfy_quant_import import import_comfy_quantized_layers
|
|
from toolkit.models.v2._mixin import OstrisModelMixin
|
|
|
|
state_dict, num_quantized = import_comfy_quantized_layers(
|
|
module, state_dict, orig_dtype=self.torch_dtype
|
|
)
|
|
if num_quantized:
|
|
self.print_and_status_update(
|
|
f" - attached {num_quantized} pre-quantized ConvRot layers to {name}"
|
|
)
|
|
# whitelist for quantized weights + leftover-meta check
|
|
OstrisModelMixin._load_state_dict_with_quantized(module, state_dict)
|
|
return num_quantized
|
|
|
|
def _load_gemma4_text_encoder(self, te_path: str, te_state_dict: dict):
|
|
"""Build the Gemma-4 12B text stack from the single comfy file. Only
|
|
the text decoder is loaded — the unified checkpoint's vision/audio
|
|
tower pieces and the connector projections are used elsewhere or
|
|
dropped, matching how the Gemma-3 vision tower was discarded."""
|
|
from safetensors import safe_open
|
|
from transformers import Gemma4TextConfig
|
|
from toolkit.models.v2.text_encoders.gemma3 import Gemma4TextEncoder
|
|
|
|
with safe_open(te_path, framework="pt") as f:
|
|
metadata = f.metadata() or {}
|
|
gemma_config = json.loads(metadata["gemma_config"])
|
|
text_config = {
|
|
k: v for k, v in gemma_config["text_config"].items() if k != "dtype"
|
|
}
|
|
|
|
with init_empty_weights():
|
|
text_encoder = Gemma4TextEncoder(Gemma4TextConfig(**text_config))
|
|
|
|
def strip_model_prefix(key: str) -> str:
|
|
return key[len("model.") :] if key.startswith("model.") else key
|
|
|
|
te_sd = {
|
|
strip_model_prefix(k): v
|
|
for k, v in te_state_dict.items()
|
|
if k.startswith("model.")
|
|
}
|
|
num_quantized = self._load_quantized_module(text_encoder, te_sd, "text encoder")
|
|
return text_encoder, num_quantized
|
|
|
|
def _load_gemma4_tokenizer(self, te_path: str, te_state_dict: dict):
|
|
"""The comfy file embeds the tokenizer and its configs as uint8
|
|
tensors; extract them next to the file once and load from there."""
|
|
from transformers import AutoTokenizer
|
|
|
|
assets_dir = os.path.splitext(te_path)[0] + "_hf_assets"
|
|
assets = {
|
|
"tokenizer.json": "tokenizer_json",
|
|
"tokenizer_config.json": "hf_asset__tokenizer_config.json",
|
|
"chat_template.jinja": "hf_asset__chat_template.jinja",
|
|
}
|
|
os.makedirs(assets_dir, exist_ok=True)
|
|
for filename, tensor_key in assets.items():
|
|
out_path = os.path.join(assets_dir, filename)
|
|
blob = te_state_dict.get(tensor_key, None)
|
|
if blob is None or os.path.exists(out_path):
|
|
continue
|
|
with open(out_path, "wb") as f:
|
|
f.write(bytes(blob.cpu().numpy().tobytes()))
|
|
# the embedded tokenizer.json has an empty post-processor (ComfyUI
|
|
# prepends BOS in its own wrapper); restore the standard Gemma
|
|
# behavior so blank prompts still yield a token
|
|
return AutoTokenizer.from_pretrained(assets_dir, add_bos_token=True)
|
|
|
|
def load_model(self):
|
|
dtype = self.torch_dtype
|
|
self.print_and_status_update("Loading LTX-2.5 model")
|
|
|
|
# ---- transformer + embedding connectors (one comfy file) ----
|
|
dit_path = self._resolve_dit_path()
|
|
te_path = self._resolve_te_path()
|
|
self.print_and_status_update(
|
|
f"Loading transformer from {os.path.basename(dit_path)}"
|
|
)
|
|
combined = load_file(dit_path)
|
|
dit_sd = get_model_state_dict_from_combined_ckpt(combined, dit_prefix)
|
|
del combined
|
|
|
|
# the per-modality text projections ride in the text encoder file but
|
|
# belong to the connectors module
|
|
te_state_dict = load_file(te_path)
|
|
for key in list(te_state_dict.keys()):
|
|
if key.startswith("text_embedding_projection."):
|
|
dit_sd[key] = te_state_dict.pop(key)
|
|
|
|
transformer, transformer_sd = convert_ltx2_transformer(
|
|
dit_sd, version=self.ltx_version, load=False
|
|
)
|
|
num_quantized_dit = self._load_quantized_module(
|
|
transformer, transformer_sd, "transformer"
|
|
)
|
|
del transformer_sd
|
|
# dit_sd still references every transformer tensor (assign=True sharing);
|
|
# drop them so each block frees as it quantizes. Connector keys stay for
|
|
# the convert_ltx2_connectors call below.
|
|
trans_sd, _ = split_transformer_and_connector_state_dict(dit_sd)
|
|
for key in trans_sd:
|
|
dit_sd.pop(key, None)
|
|
del trans_sd
|
|
if num_quantized_dit == 0:
|
|
transformer = transformer.to(dtype)
|
|
else:
|
|
# mixed-precision comfy file (fp32 scale_shift tables next to bf16
|
|
# weights): comfy-style per-op input casting, with the stored-fp32
|
|
# pieces pinned so parent .to(dtype) casts cannot downcast them
|
|
attach_per_op_casting(transformer)
|
|
pin_stored_fp32(transformer)
|
|
flush()
|
|
|
|
if num_quantized_dit:
|
|
# pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
|
|
transformer.aitk_is_quantized = True
|
|
# quantize + offload + placement, all driven by model_config
|
|
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
|
|
flush()
|
|
|
|
self.print_and_status_update("Loading connectors")
|
|
connectors, connectors_sd = convert_ltx2_connectors(
|
|
dit_sd, version=self.ltx_version, load=False
|
|
)
|
|
num_quantized_connectors = self._load_quantized_module(
|
|
connectors, connectors_sd, "connectors"
|
|
)
|
|
del connectors_sd, dit_sd
|
|
if num_quantized_connectors == 0:
|
|
connectors = connectors.to(dtype)
|
|
else:
|
|
attach_per_op_casting(connectors)
|
|
pin_stored_fp32(connectors)
|
|
flush()
|
|
|
|
# ---- text encoder (Gemma-4 12B, single comfy file) ----
|
|
self.print_and_status_update("Loading text encoder")
|
|
tokenizer = self._load_gemma4_tokenizer(te_path, te_state_dict)
|
|
text_encoder, num_quantized_te = self._load_gemma4_text_encoder(
|
|
te_path, te_state_dict
|
|
)
|
|
del te_state_dict
|
|
flush()
|
|
|
|
if num_quantized_te:
|
|
# pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
|
|
text_encoder.aitk_is_quantized = True
|
|
# quantize + offload + placement, all driven by model_config
|
|
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
|
|
|
|
text_encoder.to(self.device_torch)
|
|
text_encoder.requires_grad_(False)
|
|
text_encoder.eval()
|
|
flush()
|
|
|
|
# ---- VAEs + vocoder ----
|
|
self.print_and_status_update("Loading VAEs and other components")
|
|
video_vae_path = self._resolve_comfy_file("video_vae")
|
|
vae = convert_ltx2_video_vae(
|
|
load_file(video_vae_path), version=self.ltx_version
|
|
).to(dtype)
|
|
flush()
|
|
|
|
audio_vae_path = self._resolve_comfy_file("audio_vae")
|
|
audio_combined = load_file(audio_vae_path)
|
|
audio_sd = get_model_state_dict_from_combined_ckpt(
|
|
audio_combined, audio_vae_prefix
|
|
)
|
|
audio_vae = convert_ltx2_audio_vae(audio_sd, version=self.ltx_version).to(dtype)
|
|
vocoder_sd = get_model_state_dict_from_combined_ckpt(
|
|
audio_combined, vocoder_prefix
|
|
)
|
|
vocoder = convert_ltx2_vocoder(vocoder_sd, version=self.ltx_version).to(dtype)
|
|
del audio_combined, audio_sd, vocoder_sd
|
|
flush()
|
|
|
|
self.noise_scheduler = LTX2Model.get_train_scheduler()
|
|
|
|
self.print_and_status_update("Making pipe")
|
|
pipe: LTX2Pipeline = LTX2Pipeline(
|
|
scheduler=self.noise_scheduler,
|
|
vae=vae,
|
|
audio_vae=audio_vae,
|
|
text_encoder=None,
|
|
tokenizer=tokenizer,
|
|
connectors=connectors,
|
|
transformer=None,
|
|
vocoder=vocoder,
|
|
)
|
|
pipe.text_encoder = text_encoder
|
|
pipe.transformer = transformer
|
|
|
|
self.print_and_status_update("Preparing Model")
|
|
text_encoder = [pipe.text_encoder]
|
|
tokenizer = [pipe.tokenizer]
|
|
|
|
if not self.low_vram:
|
|
pipe.transformer = pipe.transformer.to(self.device_torch)
|
|
flush()
|
|
|
|
self.vae = ComboVae(pipe.vae, pipe.audio_vae)
|
|
self.text_encoder = text_encoder # list of text encoders
|
|
self.tokenizer = tokenizer # list of tokenizers
|
|
self.model = pipe.transformer
|
|
self.pipeline = pipe
|
|
|
|
self.audio_processor = AudioProcessor(
|
|
sample_rate=pipe.audio_sampling_rate,
|
|
mel_bins=audio_vae.config.mel_bins,
|
|
mel_hop_length=pipe.audio_hop_length,
|
|
n_fft=1024, # todo get this from vae if we can, I couldnt find it.
|
|
).to(self.device_torch, dtype=torch.float32)
|
|
|
|
self.print_and_status_update("Model Loaded")
|