diff --git a/examples/train/configs/overfit_ltx2_t2v.yaml b/examples/train/configs/overfit_ltx2_t2v.yaml new file mode 100644 index 000000000..490f435fd --- /dev/null +++ b/examples/train/configs/overfit_ltx2_t2v.yaml @@ -0,0 +1,90 @@ +# LTX-2 T2V overfitting test config. +# +# Overfits on a single short video (480x832, 81 frames @ 24fps) to +# verify the LTX-2 training plugin works end-to-end. Uses the +# distilled checkpoint (validation is 8 sampling steps, single pass). +# +# Preprocess data first (writes data/ltx2_overfit_preprocessed): +# CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py +# +# Run: +# NUM_GPUS=4 bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v.yaml + +models: + student: + _target_: fastvideo.train.models.ltx2.LTX2Model + init_from: FastVideo/LTX2-Distilled-Diffusers + trainable: true + enable_gradient_checkpointing_type: full + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: data/ltx2_overfit_preprocessed + dataloader_num_workers: 0 + train_batch_size: 1 + # LTX2Model requires 0.0: CFG dropout would zero post-connector + # embeddings, which is not the model's unconditional input. + training_cfg_rate: 0.0 + seed: 42 + num_latent_t: 11 # (81 - 1) / 8 + 1 + num_height: 480 + num_width: 832 + num_frames: 81 + + optimizer: + learning_rate: 5.0e-5 + betas: [0.9, 0.999] + weight_decay: 0.0 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 300 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/ltx2_overfit + # A full training-state checkpoint is ~150GB for the 13B trainable + # video branch; disable saves for the overfit smoke run. + training_state_checkpointing_steps: 0 + checkpoints_total_limit: 1 + + tracker: + trackers: [wandb] + project_name: fastvideo_ltx2 + run_name: ltx2_overfit + + model: + # LTX2Model.predict_noise converts the DiT's x0 output to velocity, + # so the default noise-minus-clean target reproduces the official + # unweighted masked-MSE (mask is all-ones for plain T2V). + precondition_outputs: false + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + validation: + _target_: fastvideo.train.callbacks.validation.ValidationCallback + pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline + dataset_file: data/ltx2_overfit_preprocessed/validation_prompts.json + every_steps: 50 + sampling_steps: [8] + guidance_scale: 1.0 + num_frames: 81 + +# Required so the LTX2T2VConfig pipeline config is resolved from +# init_from (without a `pipeline:` key the loader falls back to a +# generic PipelineConfig and the LTX-2 DiT cannot be constructed). +pipeline: {} diff --git a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py new file mode 100644 index 000000000..08f721da6 --- /dev/null +++ b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Preprocess LTX-2 overfit data into parquet format. + +Encodes videos with the LTX-2 causal video VAE and captions with the +Gemma text encoder (feature extractor + embedding connector) into the +t2v parquet schema expected by the training framework. + +The stored text embeddings are POST-connector: the connector replaces +pad positions with learnable registers and returns an all-valid mask, +so the parquet collate's ones/zeros mask stays semantically correct +and training needs no text encoder at all. + +Videos are resampled to TRAIN_FPS and the preprocessed clip is also +saved as an mp4 next to the parquet so overfit tests can use it as +the SSIM reference. + +Usage: + CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py +""" + +import json +import os +from typing import Any + +import av +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import torch + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.utils import maybe_download_model, verify_model_config_and_directory + +# --- Config --- +NUM_FRAMES = 81 # 8k+1 for temporal compression ratio 8 +MAX_HEIGHT = 480 # divisible by 32 (spatial compression) +MAX_WIDTH = 832 +TRAIN_FPS = 24.0 # matches the LTX-2 preset fps used at validation + +DATA_DIR = os.environ.get("LTX2_OVERFIT_DATA_DIR", "data/cats") +CAPTION_JSON = os.environ.get("LTX2_OVERFIT_CAPTION_JSON", "videos2caption_1_sample.json") +VIDEO_SUBDIR = os.environ.get("LTX2_OVERFIT_VIDEO_SUBDIR", "video") +OUTPUT_DIR = os.environ.get("LTX2_OVERFIT_OUTPUT_DIR", "data/ltx2_overfit_preprocessed") +MODEL_REPO = os.environ.get("LTX2_OVERFIT_MODEL", "FastVideo/LTX2-Distilled-Diffusers") +# The train dataloader samples with drop_last=True across data-parallel +# groups, so the dataset must hold at least num_sp_groups * batch_size +# rows or every rank gets zero batches. Replicate the overfit sample so +# a 4-GPU FSDP run still sees one batch per rank. +NUM_COPIES = int(os.environ.get("LTX2_OVERFIT_NUM_COPIES", "4")) + + +def _init_single_process_distributed() -> None: + """FastVideo component loaders expect an initialized distributed env.""" + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29511") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + from fastvideo.distributed import ( + maybe_init_distributed_environment_and_model_parallel, ) + maybe_init_distributed_environment_and_model_parallel(1, 1) + + +def load_video(path: str, num_frames: int, target_fps: float, height: int, + width: int) -> tuple[torch.Tensor, np.ndarray]: + """Load a video as [1, C, T, H, W] in [-1, 1], resampled to target_fps. + + Also returns the uint8 RGB frames [T, H, W, C] for reference-video export. + """ + with av.open(path) as container: + stream = container.streams.video[0] + native_fps = float(stream.average_rate or target_fps) + decoded = [torch.from_numpy(frame.to_ndarray(format="rgb24")) for frame in container.decode(video=0)] + if not decoded: + raise RuntimeError(f"Could not read any frames from {path}") + raw = torch.stack(decoded).permute(0, 3, 1, 2) # [T, C, H, W] uint8 + step = native_fps / target_fps + wanted = [min(int(round(i * step)), raw.shape[0] - 1) for i in range(num_frames)] + + frames = raw[wanted].float() # [T, C, H, W] in [0, 255] + src_h, src_w = frames.shape[2], frames.shape[3] + scale = max(height / src_h, width / src_w) + new_h, new_w = int(round(src_h * scale)), int(round(src_w * scale)) + frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True) + top = (new_h - height) // 2 + left = (new_w - width) // 2 + frames = frames[:, :, top:top + height, left:left + width] + + frames_np = (frames.permute(0, 2, 3, 1).clamp(0, 255).round().to(torch.uint8).numpy()) + video = frames / 127.5 - 1.0 # [0,255] -> [-1,1] + video = video.permute(1, 0, 2, 3).unsqueeze(0) # [1,C,T,H,W] + return video, frames_np + + +def main() -> None: + _init_single_process_distributed() + + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.loader.component_loader import ( + PipelineComponentLoader, ) + from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig + + device = torch.device("cuda:0") + model_path = maybe_download_model(MODEL_REPO) + model_index = verify_model_config_and_directory(model_path) + + os.makedirs(OUTPUT_DIR, exist_ok=True) + + with open(os.path.join(DATA_DIR, CAPTION_JSON)) as f: + caption_data = json.load(f) + + pipeline_config = LTX2T2VConfig() + fastvideo_args = FastVideoArgs( + model_path=model_path, + pipeline_config=pipeline_config, + num_gpus=1, + tp_size=1, + sp_size=1, + hsdp_shard_dim=1, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + ) + + def load_component(name: str) -> Any: + transformers_or_diffusers, _ = model_index[name] + return PipelineComponentLoader.load_module( + module_name=name, + component_model_path=os.path.join(model_path, name), + transformers_or_diffusers=transformers_or_diffusers, + fastvideo_args=fastvideo_args, + ) + + print("Loading LTX-2 VAE...") + vae = load_component("vae") + vae_dtype = next(vae.parameters()).dtype + print(f"VAE loaded ({sum(p.numel() for p in vae.parameters())/1e6:.0f}M, {vae_dtype})") + + print("Loading Gemma text encoder + tokenizer...") + text_encoder = load_component("text_encoder") + tokenizer = load_component("tokenizer") + tokenizer.padding_side = "left" + if tokenizer.pad_token is None and tokenizer.eos_token is not None: + tokenizer.pad_token = tokenizer.eos_token + + encoder_config = pipeline_config.text_encoder_configs[0] + tokenizer_kwargs = dict(encoder_config.tokenizer_kwargs) + if "max_length" not in tokenizer_kwargs: + tokenizer_kwargs["max_length"] = encoder_config.arch_config.text_len + preprocess_text = pipeline_config.preprocess_text_funcs[0] + + # --- Process each video --- + records = [] + for idx, item in enumerate(caption_data): + video_name = item["path"] + record_id = f"{idx:04d}_{video_name}" + caption = item["cap"][0] if isinstance(item["cap"], list) else item["cap"] + video_path = os.path.join(DATA_DIR, VIDEO_SUBDIR, video_name) + + print(f"\nProcessing: {video_name}") + print(f" Caption: {caption[:80]}...") + + video, frames_np = load_video(video_path, NUM_FRAMES, TRAIN_FPS, MAX_HEIGHT, MAX_WIDTH) + video = video.to(device=device, dtype=vae_dtype) + print(f" Video shape: {video.shape}") + + with torch.no_grad(): + # LTX-2 encode() returns a deterministic distribution whose + # mean is already per-channel normalized; store as-is. + latent = vae.encode(video).mean.squeeze(0).float().cpu() + print(f" Latent shape: {latent.shape}") + + with torch.no_grad(): + pre_embeds, attention_mask = text_encoder.preprocess_text_embeddings( + prompts=[preprocess_text(caption)], + tokenizer=tokenizer, + tokenizer_kwargs=tokenizer_kwargs, + padding_side=encoder_config.arch_config.padding_side, + ) + video_embeds, _audio_embeds, _post_mask = text_encoder.run_connectors( + pre_embeds, + attention_mask, + ) + text_embedding = video_embeds.squeeze(0).float().cpu() + print(f" Text embedding shape: {text_embedding.shape}") + + record = { + "id": record_id, + "vae_latent_bytes": latent.numpy().tobytes(), + "vae_latent_shape": list(latent.shape), + "vae_latent_dtype": str(latent.dtype).replace("torch.", ""), + "text_embedding_bytes": (text_embedding.numpy().tobytes()), + "text_embedding_shape": list(text_embedding.shape), + "text_embedding_dtype": str(text_embedding.dtype).replace("torch.", ""), + "file_name": video_name, + "caption": caption, + "media_type": "video", + "width": MAX_WIDTH, + "height": MAX_HEIGHT, + "num_frames": NUM_FRAMES, + "duration_sec": NUM_FRAMES / TRAIN_FPS, + "fps": TRAIN_FPS, + } + records.append(record) + + # Save the preprocessed clip so overfit tests can compare + # validation output against the memorization target. + import imageio + ref_path = os.path.join(OUTPUT_DIR, f"training_sample_{idx}.mp4") + with imageio.get_writer(ref_path, fps=TRAIN_FPS) as writer: + for frame in frames_np: + writer.append_data(frame) + print(f" Wrote reference clip to {ref_path}") + + # Clean up + del text_encoder, tokenizer, vae + torch.cuda.empty_cache() + + # Write parquet (replicated NUM_COPIES times; see comment at top) + replicated = [] + for copy_idx in range(max(1, NUM_COPIES)): + for r in records: + row = dict(r) + row["id"] = f"{r['id']}_copy{copy_idx}" + replicated.append(row) + table = pa.table( + {k: [r[k] for r in replicated] + for k in replicated[0]}, + schema=pyarrow_schema_t2v, + ) + output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet") + pq.write_table(table, output_path) + print(f"\nWrote {len(replicated)} records " + f"({len(records)} unique x {max(1, NUM_COPIES)} copies) to {output_path}") + + # Write validation prompts for the validation callback + val_prompts = { + "data": [{ + "caption": (item["cap"][0] if isinstance(item["cap"], list) else item["cap"]), + } for item in caption_data], + } + val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json") + with open(val_path, "w") as f: + json.dump(val_prompts, f, indent=2) + print(f"Wrote validation prompts to {val_path}") + + print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py new file mode 100644 index 000000000..1f7d228f0 --- /dev/null +++ b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py @@ -0,0 +1,139 @@ +"""E2E overfit test for LTX-2 fine-tuning on the modular trainer (fastvideo/train). + +Downloads the single-sample cats dataset, preprocesses it with the +LTX-2 VAE + Gemma text encoder into the t2v parquet format, trains +LTX2Model with FineTuneMethod for 300 steps on 4 GPUs, and checks that +the final validation video reproduces the training clip (MS-SSIM +against the preprocessed ground-truth clip) better than the step-0 +baseline. + +GPU assumptions: 4 x large-memory GPUs (developed on 4x GB200 192GB; +the 18.9B-param DiT trains FSDP-sharded with fp32 master weights). +Requires the FastVideo/LTX2-Distilled-Diffusers checkpoint (set +HF_HOME to a cache that contains it, or allow ~60GB of downloads). + +Guarded by FASTVIDEO_NIGHTLY=1 so the default test suite stays fast. +""" + +import glob +import json +import os +import re +import subprocess +import sys +from pathlib import Path + +import pytest + +from huggingface_hub import snapshot_download + +sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) + +from fastvideo.tests.utils import compute_video_ssim_torchvision # noqa: E402 + +NUM_GPUS_TRAINING = "4" + +DATA_DIR = "data" +LOCAL_RAW_DATA_DIR = Path(DATA_DIR) / "cats" +LOCAL_PREPROCESSED_DATA_DIR = Path(DATA_DIR) / "ltx2_overfit_preprocessed" +LOCAL_OUTPUT_DIR = Path(DATA_DIR) / "outputs_ltx2_overfit" + +PREPROCESSING_SCRIPT = "fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py" +TRAINING_CONFIG = "examples/train/configs/overfit_ltx2_t2v.yaml" + + +def download_data(): + os.makedirs(DATA_DIR, exist_ok=True) + print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...") + snapshot_download( + repo_id="wlsaidhi/cats-overfit-merged", + local_dir=str(LOCAL_RAW_DATA_DIR), + repo_type="dataset", + token=os.environ.get("HF_TOKEN"), + ) + assert LOCAL_RAW_DATA_DIR.exists(), ( + f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") + + +def run_preprocessing(): + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = env.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0] + env["LTX2_OVERFIT_DATA_DIR"] = str(LOCAL_RAW_DATA_DIR) + env["LTX2_OVERFIT_OUTPUT_DIR"] = str(LOCAL_PREPROCESSED_DATA_DIR) + cmd = [sys.executable, PREPROCESSING_SCRIPT] + print(f"Running preprocessing: {cmd}") + subprocess.run(cmd, check=True, env=env) + + +def run_training(): + env = os.environ.copy() + env.setdefault("WANDB_MODE", "offline") + cmd = [ + "torchrun", + "--nnodes", "1", + "--nproc_per_node", NUM_GPUS_TRAINING, + "-m", "fastvideo.train.entrypoint.train", + "--config", TRAINING_CONFIG, + "--training.checkpoint.output_dir", str(LOCAL_OUTPUT_DIR), + "--training.data.data_path", str(LOCAL_PREPROCESSED_DATA_DIR), + "--callbacks.validation.dataset_file", + str(LOCAL_PREPROCESSED_DATA_DIR / "validation_prompts.json"), + ] + print(f"Running training: {cmd}") + subprocess.run(cmd, check=True, env=env) + + +def _validation_videos_by_step() -> dict[int, str]: + pattern = os.path.join(str(LOCAL_OUTPUT_DIR), "validation_step_*_video_0.mp4") + videos: dict[int, str] = {} + for path in glob.glob(pattern): + match = re.search(r"validation_step_(\d+)_", os.path.basename(path)) + if match: + videos[int(match.group(1))] = path + return videos + + +@pytest.mark.skipif( + os.environ.get("FASTVIDEO_NIGHTLY") != "1", + reason="nightly e2e overfit test; set FASTVIDEO_NIGHTLY=1 to run", +) +def test_e2e_ltx2_overfit_new_stack(): + download_data() + run_preprocessing() + run_training() + + reference_video = LOCAL_PREPROCESSED_DATA_DIR / "training_sample_0.mp4" + assert reference_video.exists(), ( + f"Reference (preprocessed training clip) not found at {reference_video}") + + videos = _validation_videos_by_step() + assert videos, f"No validation videos found under {LOCAL_OUTPUT_DIR}" + final_step = max(videos) + final_video = videos[final_step] + print(f"Final validation video (step {final_step}): {final_video}") + + final_mean, final_min, final_max = compute_video_ssim_torchvision( + str(reference_video), final_video, use_ms_ssim=True) + print(f"\n===== MS-SSIM vs training clip at step {final_step} =====") + print(f"Mean: {final_mean:.4f} Min: {final_min:.4f} Max: {final_max:.4f}") + + results = {"final_step": final_step, "final_mean_ssim": final_mean} + baseline_step = min(videos) + if baseline_step != final_step: + base_mean, _, _ = compute_video_ssim_torchvision( + str(reference_video), videos[baseline_step], use_ms_ssim=True) + print(f"Baseline (step {baseline_step}) mean MS-SSIM: {base_mean:.4f}") + results["baseline_mean_ssim"] = base_mean + assert final_mean > base_mean, ( + f"Overfitting did not improve similarity to the training clip: " + f"step {final_step} mean SSIM {final_mean:.4f} <= " + f"step {baseline_step} mean SSIM {base_mean:.4f}") + print(json.dumps(results)) + + assert final_mean > 0.5, ( + f"Mean MS-SSIM vs the training clip is below 0.5: {final_mean:.4f}") + + +if __name__ == "__main__": + os.environ.setdefault("FASTVIDEO_NIGHTLY", "1") + test_e2e_ltx2_overfit_new_stack() diff --git a/fastvideo/tests/train/fixtures/ltx2_t2v_finetune_min.yaml b/fastvideo/tests/train/fixtures/ltx2_t2v_finetune_min.yaml new file mode 100644 index 000000000..5754a1f68 --- /dev/null +++ b/fastvideo/tests/train/fixtures/ltx2_t2v_finetune_min.yaml @@ -0,0 +1,53 @@ +# Minimum config to run a single training step of FineTuneMethod on +# LTX2Model for the per-method smoke test. Uses the real +# FastVideo/LTX2-Distilled-Diffusers checkpoint (18.9B params, ~38GB +# in bf16 — needs a large-memory GPU; the test skips below 60GB) with +# tiny synthetic latents. LTX2 requires training_cfg_rate=0 (CFG +# dropout would zero post-connector text embeddings, which is not the +# model's unconditional input) and precondition_outputs=false (the +# plugin converts the DiT's x0 output to velocity). + +models: + student: + _target_: fastvideo.train.models.ltx2.LTX2Model + init_from: FastVideo/LTX2-Distilled-Diffusers + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + dit_precision: bf16 + + distributed: + num_gpus: 1 + sp_size: 1 + tp_size: 1 + + data: + seed: 42 + train_batch_size: 1 + training_cfg_rate: 0.0 + num_latent_t: 3 + num_height: 256 + num_width: 256 + num_frames: 17 + + optimizer: + learning_rate: 1.0e-6 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1 + gradient_accumulation_steps: 1 + + model: + precondition_outputs: false + +# Required so the LTX2T2VConfig pipeline config is resolved from +# init_from (without a `pipeline:` key the loader falls back to a +# generic PipelineConfig and the LTX-2 DiT cannot be constructed). +pipeline: {} diff --git a/fastvideo/tests/train/methods/grad_norm_refs.json b/fastvideo/tests/train/methods/grad_norm_refs.json index a026a5f84..765338d5b 100644 --- a/fastvideo/tests/train/methods/grad_norm_refs.json +++ b/fastvideo/tests/train/methods/grad_norm_refs.json @@ -3,6 +3,9 @@ "H200": 5.2763, "L40S": 4.1553 }, + "test_ltx2_finetune": { + "GB200": 0.3127 + }, "test_matrixgame2_finetune": { "H200": 4.8247 }, diff --git a/fastvideo/tests/train/methods/test_ltx2_finetune.py b/fastvideo/tests/train/methods/test_ltx2_finetune.py new file mode 100644 index 000000000..f36ff970c --- /dev/null +++ b/fastvideo/tests/train/methods/test_ltx2_finetune.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-method GPU smoke test: ``LTX2Model`` + ``FineTuneMethod``. + +Mirrors ``test_wan_finetune.py`` for the LTX-2 plugin. LTX2-specific +differences: the synthetic ``raw_batch`` carries a single +post-connector Gemma embedding ([1024, 3840]) and 128-channel VAE +latents; the DiT is 18.9B params, so the test skips on GPUs with less +than 60GB memory (e.g. the L40S CI runner). +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29522") + +from pathlib import Path + +import pytest +import torch + +from fastvideo.train.methods.fine_tuning.finetune import ( + FineTuneMethod, ) +from fastvideo.train.models.ltx2 import LTX2Model +from fastvideo.train.utils.config import load_run_config + +from .grad_norm_regression import ( + check_grad_norm_regression, + resolve_blocks, +) + +_FIXTURE = str( + Path(__file__).resolve().parent.parent / "fixtures" + / "ltx2_t2v_finetune_min.yaml") + +# Post-connector Gemma embedding width / length. +_LTX2_TEXT_DIM = 3840 +_LTX2_TEXT_LEN = 1024 + +_MIN_GPU_MEMORY_GB = 60 + + +def _gpu_too_small() -> bool: + if not torch.cuda.is_available(): + return True + total = torch.cuda.get_device_properties(0).total_memory + return total < _MIN_GPU_MEMORY_GB * 1024**3 + + +def _build_synthetic_batch( + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + """Tiny synthetic ``raw_batch`` matching ``LTX2Model.prepare_batch``. + + LTX-2 latents are 128-channel (temporal 8x / spatial 32x + compression); text embeddings are post-connector Gemma features. + """ + batch_size = 1 + return { + "text_embedding": + torch.randn(batch_size, + _LTX2_TEXT_LEN, + _LTX2_TEXT_DIM, + device=device, + dtype=dtype), + "text_attention_mask": + torch.ones(batch_size, _LTX2_TEXT_LEN, device=device), + "vae_latent": + torch.randn(batch_size, 128, 3, 8, 8, device=device, dtype=dtype), + } + + +@pytest.mark.usefixtures("distributed_setup") +def test_ltx2_finetune_single_train_step( + monkeypatch: pytest.MonkeyPatch) -> None: + if _gpu_too_small(): + pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB " + "memory (LTX-2 DiT is 18.9B params)") + + cfg = load_run_config(_FIXTURE) + + device = torch.device("cuda:0") + dtype = torch.bfloat16 + + # Feed a synthetic ``raw_batch`` straight into ``single_train_step``, + # so the parquet train dataloader built by ``init_preprocessors`` is + # never iterated. Stub it out so construction does not require a real + # ``training.data.data_path``. + monkeypatch.setattr( + "fastvideo.train.utils.dataloader." + "build_parquet_t2v_train_dataloader", + lambda *args, **kwargs: None, + ) + + model = LTX2Model( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=True, + ) + model.transformer = model.transformer.to(device=device, dtype=dtype) + + method = FineTuneMethod( + cfg=cfg, + role_models={"student": model}, + ) + method.on_train_start() + + batch = _build_synthetic_batch(device, dtype) + loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0) + + loss = loss_map["total_loss"] + assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" + assert torch.isfinite(loss).item(), ( + f"total_loss is not finite: {loss.item()}") + + method.backward(loss_map, outputs, grad_accum_rounds=1) + + # LTX-2 nests its transformer_blocks under the ``model`` submodule. + blocks = resolve_blocks(model.transformer.model) + assert blocks is not None and len(blocks) > 0, ( + "transformer is expected to expose a non-empty block list") + layer0 = blocks[0] + + trainable = [p for p in layer0.parameters() if p.requires_grad] + assert len(trainable) > 0, "layer 0 has no trainable parameters" + + for i, p in enumerate(trainable): + assert p.grad is not None, f"layer 0 param[{i}] has None grad" + assert torch.isfinite(p.grad).all().item(), ( + f"layer 0 param[{i}] grad contains NaN/Inf") + + any_nonzero = any( + p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ( + "all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") + + # Audio / cross-modal parameters must be frozen by default. + audio_trainable = [ + name for name, param in model.transformer.named_parameters() + if param.requires_grad and any( + pattern in name for pattern in ("audio", "a2v", "v2a", "av_ca")) + ] + assert not audio_trainable, ( + f"audio/cross-modal params unexpectedly trainable: " + f"{audio_trainable[:5]}") + + # Device-keyed grad-norm regression on top of the same harness. + # Skips when the current GPU has no seeded reference. + check_grad_norm_regression("test_ltx2_finetune", model.transformer.model) diff --git a/fastvideo/tests/train/models/test_load_ltx2.py b/fastvideo/tests/train/models/test_load_ltx2.py new file mode 100644 index 000000000..ff6f3eba1 --- /dev/null +++ b/fastvideo/tests/train/models/test_load_ltx2.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GPU loading + forward smoke test for ``LTX2Model``. + +Loads the real FastVideo/LTX2-Distilled-Diffusers checkpoint (18.9B at +bf16, ~38GB — skips on GPUs with less than 60GB memory) via +``LTX2Model.__init__`` and runs one transformer forward pass on +synthetic inputs. Catches loader or forward-signature regressions in +``fastvideo.train.models.ltx2.LTX2Model`` and the underlying +``LTX2Transformer3DModel``. + +LTX-2's transformer takes per-token sigma timesteps in [0, 1] shaped +[B, tokens] and a post-connector Gemma text embedding ([1024, 3840]); +it returns the denoised x0 prediction. This mirrors the kwargs in +``LTX2Model._build_distill_input_kwargs``. +""" + +from __future__ import annotations + +import os + +# Required by the ``distributed_setup`` fixture pulled from +# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29519") + +from pathlib import Path + +import pytest +import torch + +from fastvideo.forward_context import set_forward_context +from fastvideo.pipelines import ForwardBatch +from fastvideo.train.models.ltx2 import LTX2Model +from fastvideo.train.utils.config import load_run_config + +_FIXTURE = str( + Path(__file__).resolve().parent.parent / "fixtures" + / "ltx2_t2v_finetune_min.yaml") + +# Post-connector Gemma embedding width / length. +_LTX2_TEXT_DIM = 3840 +_LTX2_TEXT_LEN = 1024 + +_MIN_GPU_MEMORY_GB = 60 + + +def _gpu_too_small() -> bool: + if not torch.cuda.is_available(): + return True + total = torch.cuda.get_device_properties(0).total_memory + return total < _MIN_GPU_MEMORY_GB * 1024**3 + + +@pytest.mark.usefixtures("distributed_setup") +def test_ltx2_model_loads_and_forwards(): + if _gpu_too_small(): + pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB " + "memory (LTX-2 DiT is 18.9B params)") + + cfg = load_run_config(_FIXTURE) + model = LTX2Model( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=False, + ) + + transformer = model.transformer + assert isinstance(transformer, torch.nn.Module) + assert sum(p.numel() for p in transformer.parameters()) > 0 + + device = torch.device("cuda:0") + dtype = torch.bfloat16 + transformer = transformer.to(device=device, dtype=dtype).eval() + + # LTX-2 transformer takes [B, 128, T, H, W] latents, a post-connector + # Gemma embedding, and PER-TOKEN sigmas in [0, 1] ([B, T*H*W] with + # patch size 1x1x1). Small spatial + few frames so this fits next to + # the 18.9B model. + b, c, t, h, w = 1, 128, 3, 8, 8 + tokens = t * h * w + hidden_states = torch.randn(b, c, t, h, w, device=device, dtype=dtype) + encoder_hidden_states = torch.randn(b, + _LTX2_TEXT_LEN, + _LTX2_TEXT_DIM, + device=device, + dtype=dtype) + timestep = torch.full((b, tokens), 0.5, device=device, dtype=torch.float32) + + with torch.no_grad(), torch.autocast(device.type, dtype=dtype), \ + set_forward_context( + current_timestep=timestep * 1000.0, + attn_metadata=None, + forward_batch=ForwardBatch(data_type="video", fps=24.0), + ): + out = transformer( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + encoder_attention_mask=None, + return_dict=False, + ) + + assert torch.is_tensor(out), f"expected a tensor output, got {type(out)}" + assert out.shape == hidden_states.shape, ( + f"denoised output shape {tuple(out.shape)} != input latent shape " + f"{tuple(hidden_states.shape)}") + assert torch.isfinite(out).all().item(), "forward output contains NaN/Inf" diff --git a/fastvideo/train/models/ltx2/__init__.py b/fastvideo/train/models/ltx2/__init__.py new file mode 100644 index 000000000..6b7d14321 --- /dev/null +++ b/fastvideo/train/models/ltx2/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LTX-2 model plugin package.""" + +from fastvideo.train.models.ltx2.ltx2 import ( + LTX2Model as LTX2Model, ) diff --git a/fastvideo/train/models/ltx2/ltx2.py b/fastvideo/train/models/ltx2/ltx2.py new file mode 100644 index 000000000..7327e2f63 --- /dev/null +++ b/fastvideo/train/models/ltx2/ltx2.py @@ -0,0 +1,534 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LTX-2 model plugin (per-role instance). + +Subclasses WanModel but replaces the pieces where LTX-2 differs: + - transformer class name: LTX2Transformer3DModel (blocks nested + under ``transformer.model``, so activation checkpointing must + target the inner module) + - latents come pre-normalized from the LTX-2 VAE encoder + (per-channel statistics applied inside ``encode``), so no + ``normalize_dit_input`` call + - sigma sampling follows the official LTX-2 trainer: stretched + shifted-logit-normal with a token-count-dependent shift + (0.95 @ 1024 tokens -> 2.05 @ 4096 tokens) and a 10% uniform + mixture, instead of scheduler-index sampling + - the DiT consumes PER-TOKEN sigmas in [0, 1] shaped [B, tokens] + (not integer 0-1000 timesteps) and returns the DENOISED x0 + prediction; ``predict_noise`` converts it back to the + framework's velocity convention v = (x_t - x0) / sigma so the + default FineTune target ``noise - clean`` applies unchanged + - temporal RoPE coordinates are divided by fps read from + ``get_forward_context().forward_batch.fps``; training must set + it to the data fps or validation (which uses the preset fps) + would see different temporal frequencies than training + - the ~5.8B audio / cross-modal (a2v, v2a, av_ca) parameters are + frozen by default: video-only forwards never give them + gradients, and freezing keeps them out of DCP checkpoints and + optimizer state +""" + +from __future__ import annotations + +from typing import Any, Literal, TYPE_CHECKING + +import torch + +import fastvideo.envs as envs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.dits.ltx2 import VideoLatentShape +from fastvideo.pipelines import ForwardBatch, TrainingBatch +from fastvideo.training.activation_checkpoint import ( + apply_activation_checkpointing, ) + +from fastvideo.train.models.wan.wan import WanModel +from fastvideo.train.utils.module_state import ( + apply_trainable, ) +from fastvideo.train.utils.moduleloader import ( + load_module_from_path, ) + +if TYPE_CHECKING: + from fastvideo.train.utils.training_config import ( + TrainingConfig, ) + from fastvideo.train.utils.lora import LoraConfig + +logger = init_logger(__name__) + +# Parameters whose names match any of these substrings belong to the +# audio branch or the audio<->video cross-modal machinery. They are +# unused (no gradients) in video-only forwards. +_AUDIO_PARAM_PATTERNS = ("audio", "a2v", "v2a", "av_ca") + +# Official LTX-2 trainer timestep sampler constants +# (ltx_trainer/timestep_samplers.py: ShiftedLogitNormalTimestepSampler). +_SIGMA_MIN_TOKENS = 1024.0 +_SIGMA_MAX_TOKENS = 4096.0 +_SIGMA_MIN_SHIFT = 0.95 +_SIGMA_MAX_SHIFT = 2.05 +_SIGMA_STD = 1.0 +_SIGMA_EPS = 1e-3 +# 0.5% / 99.9% normal percentiles used to stretch the logit-normal +# samples so the sigma distribution covers the full [0, 1] range. +_SIGMA_Z_LO = -2.5758 +_SIGMA_Z_HI = 3.0902 + +# Preset fps for LTX-2 checkpoints; used when a batch carries no fps +# metadata so training RoPE still matches validation inference. +_DEFAULT_ROPE_FPS = 24.0 + + +class LTX2Model(WanModel): + """LTX-2 per-role model for the modular trainer.""" + + _transformer_cls_name: str = "LTX2Transformer3DModel" + + def __init__( + self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + disable_custom_init_weights: bool = False, + enable_gradient_checkpointing_type: str + | None = None, + transformer_override_safetensor: str + | None = None, + lora: LoraConfig | dict[str, Any] | None = None, + train_audio: bool = False, + timestep_uniform_prob: float = 0.1, + ) -> None: + cfg_rate = float(getattr(training_config.data, "training_cfg_rate", 0.0) or 0.0) + if cfg_rate > 0.0: + raise NotImplementedError("LTX2Model only supports training_cfg_rate=0. CFG dropout " + "zeroes the post-connector text embeddings, which is not " + "what LTX-2 inference uses as the unconditional input " + "(an empty prompt encoded through Gemma + connector). Set " + "training.data.training_cfg_rate: 0.0.") + + self._timestep_uniform_prob = float(timestep_uniform_prob) + self._train_audio = bool(train_audio) + self._rope_fps: float = _DEFAULT_ROPE_FPS + self._rope_forward_batch: ForwardBatch | None = None + + super().__init__( + init_from=init_from, + training_config=training_config, + trainable=trainable, + disable_custom_init_weights=(disable_custom_init_weights), + # The LTX-2 sigma sampler below does not use the scheduler; + # keep an unshifted scheduler for the inherited distillation + # helpers (add_noise / num_train_timesteps). + flow_shift=1.0, + enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type), + transformer_override_safetensor=(transformer_override_safetensor), + lora=lora, + ) + + if trainable and not self._train_audio: + self._freeze_audio_parameters() + + # No negative-prompt cache: cfg_rate is forced to 0 above and + # loading Gemma (~23GB) on every rank just for an unused + # negative embedding is wasteful. + self.set_requires_negative_conditioning(False) + + # ------------------------------------------------------------------ + # Loading + # ------------------------------------------------------------------ + + def _load_transformer( + self, + *, + init_from: str, + trainable: bool, + disable_custom_init_weights: bool, + enable_gradient_checkpointing_type: str | None, + training_config: TrainingConfig, + transformer_override_safetensor: str | None = None, + ) -> torch.nn.Module: + transformer = load_module_from_path( + model_path=init_from, + module_type="transformer", + training_config=training_config, + disable_custom_init_weights=(disable_custom_init_weights), + override_transformer_cls_name=(self._transformer_cls_name), + transformer_override_safetensor=(transformer_override_safetensor), + ) + ckpt_type = (enable_gradient_checkpointing_type or getattr( + getattr(training_config, "model", None), + "enable_gradient_checkpointing_type", + None, + )) + if trainable and ckpt_type: + # LTX-2 nests transformer_blocks under ``.model``; applying + # checkpointing at the wrapper level raises because no block + # list is found there. + transformer.model = apply_activation_checkpointing( + transformer.model, + checkpointing_type=ckpt_type, + ) + if self._enable_lora_if_configured(transformer): + return transformer + transformer = apply_trainable(transformer, trainable=trainable) + return transformer + + def _freeze_audio_parameters(self) -> None: + frozen_params = 0 + for name, param in self.transformer.named_parameters(): + if any(pattern in name for pattern in _AUDIO_PARAM_PATTERNS): + param.requires_grad_(False) + frozen_params += param.numel() + logger.info( + "LTX2Model: froze %.2fB audio/cross-modal parameters " + "(video-only training)", + frozen_params / 1e9, + ) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def ensure_negative_conditioning(self) -> None: + raise NotImplementedError("LTX2Model does not implement negative conditioning; " + "training_cfg_rate must stay 0.") + + @torch.no_grad() + def decode_latents( + self, + latents_b_t_c_h_w: torch.Tensor, + ) -> torch.Tensor: + if self.vae is None: + raise RuntimeError("LTX-2 VAE is not initialized") + latents = latents_b_t_c_h_w.permute(0, 2, 1, 3, 4) + # The LTX-2 decoder un-normalizes internally via its + # per-channel statistics buffers; no external denorm. + media = self.vae.decode(latents.to(next(self.vae.parameters()).dtype)) + return (media.float() / 2 + 0.5).clamp(0, 1) + + def _init_timestep_mechanics(self) -> None: + assert self.training_config is not None + tc = self.training_config + # LTX2T2VConfig carries no flow_shift (sigmas are computed + # inline by the pipeline); the inherited shift/clamp helpers + # are unused for fine-tuning, so default to identity shift. + flow_shift = tc.pipeline_config.flow_shift # type: ignore[union-attr] + self.timestep_shift = (float(flow_shift) if flow_shift is not None else 1.0) + self.num_train_timestep = int(self.noise_scheduler.num_train_timesteps) + self.min_timestep = 0 + self.max_timestep = self.num_train_timestep + + # ------------------------------------------------------------------ + # Runtime primitives + # ------------------------------------------------------------------ + + def prepare_batch( + self, + raw_batch: dict[str, Any], + *, + generator: torch.Generator, + latents_source: Literal["data", "zeros"] = "data", + ) -> TrainingBatch: + assert self.training_config is not None + tc = self.training_config + + dtype = self._get_training_dtype() + device = self.device + + training_batch = TrainingBatch() + encoder_hidden_states = raw_batch["text_embedding"] + encoder_attention_mask = raw_batch["text_attention_mask"] + infos = raw_batch.get("info_list") + + if latents_source == "zeros": + batch_size = encoder_hidden_states.shape[0] + vae_config = ( + tc.pipeline_config.vae_config.arch_config # type: ignore[union-attr] + ) + num_channels = getattr( + vae_config, + "z_dim", + getattr(vae_config, "latent_channels", 128), + ) + spatial_compression_ratio = (vae_config.spatial_compression_ratio) + latent_height = (tc.data.num_height // spatial_compression_ratio) + latent_width = (tc.data.num_width // spatial_compression_ratio) + latents = torch.zeros( + batch_size, + num_channels, + tc.data.num_latent_t, + latent_height, + latent_width, + device=device, + dtype=dtype, + ) + elif latents_source == "data": + if "vae_latent" not in raw_batch: + raise ValueError("vae_latent not found in batch " + "and latents_source='data'") + latents = raw_batch["vae_latent"] + latents = latents[:, :, :tc.data.num_latent_t] + latents = latents.to(device, dtype=dtype) + else: + raise ValueError(f"Unknown latents_source: " + f"{latents_source!r}") + + # LTX-2 VAE encode() already applies per-channel normalization + # (scaling_factor is 1.0); latents are used as stored. + training_batch.latents = latents + training_batch.encoder_hidden_states = (encoder_hidden_states.to(device, dtype=dtype)) + training_batch.encoder_attention_mask = (encoder_attention_mask.to(device, dtype=dtype)) + training_batch.infos = infos + + self._update_rope_fps(infos) + + training_batch = self._prepare_dit_inputs(training_batch, generator) + training_batch = self._build_attention_metadata(training_batch) + training_batch.attn_metadata_vsa = None + + return training_batch + + def predict_noise( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + batch: TrainingBatch, + *, + conditional: bool, + cfg_uncond: dict[str, Any] | None = None, + attn_kind: Literal["dense", "vsa"] = "dense", + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, + ) -> torch.Tensor: + if clean_x is not None or aug_t is not None: + raise NotImplementedError("LTX2Model does not support teacher forcing inputs") + device_type = self.device.type + dtype = self._get_training_dtype() + if conditional: + text_dict = batch.conditional_dict + if text_dict is None: + raise RuntimeError("Missing conditional_dict in " + "TrainingBatch") + else: + text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond) + + if attn_kind not in ("dense", "vsa"): + raise ValueError(f"Unknown attn_kind: {attn_kind!r}") + attn_metadata = batch.attn_metadata + + assert batch.sigmas is not None + # finetune.py hands noisy_latents in (B, T, C, H, W); the LTX-2 + # DiT expects (B, C, T, H, W). + noisy_bcthw = noisy_latents.permute(0, 2, 1, 3, 4).to(dtype) + + with torch.autocast(device_type, dtype=dtype), set_forward_context( + current_timestep=batch.timesteps, + attn_metadata=attn_metadata, + forward_batch=self._make_rope_forward_batch(), + ): + input_kwargs = self._build_distill_input_kwargs( + noisy_bcthw, + timestep, + text_dict, + ) + transformer = self._get_transformer(timestep) + denoised = transformer(**input_kwargs) + + if isinstance(denoised, tuple): + denoised = denoised[0] + + # The wrapper returns the denoised x0 prediction + # (sample - velocity * sigma). Convert back to velocity so the + # framework's default target ``noise - clean`` applies. Compute + # in fp32: sigma can be ~1e-3 and bf16 division would amplify + # quantization error. + sigmas = batch.sigmas.float() + velocity = (noisy_bcthw.float() - denoised.float()) / sigmas + return velocity.permute(0, 2, 1, 3, 4) + + def backward( + self, + loss: torch.Tensor, + ctx: Any, + *, + grad_accum_rounds: int, + ) -> None: + timesteps, attn_metadata = ctx + # Re-enter the forward context with the same fps-carrying batch + # so activation-checkpoint recompute sees identical RoPE inputs. + with set_forward_context( + current_timestep=timesteps, + attn_metadata=attn_metadata, + forward_batch=self._make_rope_forward_batch(), + ): + (loss / max(1, int(grad_accum_rounds))).backward() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _update_rope_fps(self, infos: list[dict[str, Any]] | None) -> None: + fps: float | None = None + if infos: + raw_fps = infos[0].get("fps") + if raw_fps: + fps = float(raw_fps) + if fps is None or fps <= 0: + fps = _DEFAULT_ROPE_FPS + if self._rope_forward_batch is None or fps != self._rope_fps: + self._rope_fps = fps + self._rope_forward_batch = ForwardBatch(data_type="video", fps=fps) + + def _make_rope_forward_batch(self) -> ForwardBatch: + if self._rope_forward_batch is None: + self._rope_forward_batch = ForwardBatch( + data_type="video", + fps=self._rope_fps, + ) + return self._rope_forward_batch + + def _sample_ltx2_sigmas( + self, + batch_size: int, + token_count: int, + device: torch.device, + generator: torch.Generator, + ) -> torch.Tensor: + """Official LTX-2 stretched shifted-logit-normal sigma sampler. + + mu is linearly interpolated from the patchified video token + count (unclamped, matching the official trainer), the + logit-normal samples are stretched to cover [0, 1] using the + 0.5%/99.9% percentiles, and 10% of samples are replaced by + U(eps, 1). + """ + slope = ((_SIGMA_MAX_SHIFT - _SIGMA_MIN_SHIFT) / (_SIGMA_MAX_TOKENS - _SIGMA_MIN_TOKENS)) + mu = slope * float(token_count) + (_SIGMA_MIN_SHIFT - slope * _SIGMA_MIN_TOKENS) + + normal = torch.randn( + (batch_size, ), + generator=generator, + device=device, + dtype=torch.float32, + ) * _SIGMA_STD + mu + sigmas = torch.sigmoid(normal) + + lo = torch.sigmoid(torch.tensor(mu + _SIGMA_Z_LO * _SIGMA_STD, device=device)) + hi = torch.sigmoid(torch.tensor(mu + _SIGMA_Z_HI * _SIGMA_STD, device=device)) + raw = (sigmas - lo) / (hi - lo) + stretched = torch.where(raw >= _SIGMA_EPS, raw, 2 * _SIGMA_EPS - raw) + stretched = stretched.clamp(0.0, 1.0) + + if self._timestep_uniform_prob > 0.0: + prob = torch.rand( + (batch_size, ), + generator=generator, + device=device, + ) + uniform = torch.rand( + (batch_size, ), + generator=generator, + device=device, + ) * (1.0 - _SIGMA_EPS) + _SIGMA_EPS + stretched = torch.where( + prob > self._timestep_uniform_prob, + stretched, + uniform, + ) + return stretched + + def _prepare_dit_inputs( + self, + training_batch: TrainingBatch, + generator: torch.Generator, + ) -> TrainingBatch: + assert self.training_config is not None + tc = self.training_config + latents = training_batch.latents + assert isinstance(latents, torch.Tensor) + batch_size = latents.shape[0] + + video_shape = VideoLatentShape.from_torch_shape(latents.shape) + patchifier = getattr(self.transformer, "patchifier", None) + if patchifier is not None: + token_count = int(patchifier.get_token_count(video_shape)) + else: + token_count = int(latents.shape[2] * latents.shape[3] * latents.shape[4]) + self._token_count = token_count + + sigmas = self._sample_ltx2_sigmas( + batch_size, + token_count, + latents.device, + generator, + ) + if int(tc.distributed.sp_size or 1) > 1: + self.sp_group.broadcast(sigmas, src=0) + + noise = torch.randn( + latents.shape, + generator=generator, + device=latents.device, + dtype=latents.dtype, + ) + sigmas_expanded = sigmas.view(-1, 1, 1, 1, 1) + noisy_model_input = ((1.0 - sigmas_expanded) * latents.float() + sigmas_expanded * noise.float()).to( + latents.dtype) + + training_batch.noisy_model_input = noisy_model_input + # Keep the framework-wide "timesteps ~ sigma * 1000" convention + # (used for logging and the forward context); sigmas stay fp32 + # for the velocity conversion. + training_batch.timesteps = sigmas * 1000.0 + training_batch.sigmas = sigmas_expanded + training_batch.noise = noise + training_batch.raw_latent_shape = latents.shape + + training_batch.conditional_dict = { + "encoder_hidden_states": (training_batch.encoder_hidden_states), + "encoder_attention_mask": (training_batch.encoder_attention_mask), + } + + training_batch.latents = (training_batch.latents.permute(0, 2, 1, 3, 4)) + return training_batch + + def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch: + if envs.FASTVIDEO_ATTENTION_BACKEND in ("VIDEO_SPARSE_ATTN", "VMOBA_ATTN"): + raise NotImplementedError("LTX2Model does not support VSA/VMOBA attention backends") + training_batch.attn_metadata = None + return training_batch + + def _build_distill_input_kwargs( + self, + noise_input: torch.Tensor, + timestep: torch.Tensor, + text_dict: dict[str, torch.Tensor] | None, + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, + ) -> dict[str, Any]: + if text_dict is None: + raise ValueError("text_dict cannot be None for LTX-2 forward") + if clean_x is not None or aug_t is not None: + raise NotImplementedError("LTX2Model does not support teacher forcing inputs") + + # noise_input arrives already in (B, C, T, H, W). + batch_size = noise_input.shape[0] + token_count = getattr(self, "_token_count", None) + if token_count is None: + token_count = int(noise_input.shape[2] * noise_input.shape[3] * noise_input.shape[4]) + + # The DiT wants per-token sigmas in [0, 1] shaped [B, tokens] + # (i2v-style conditioning would zero conditioned tokens; plain + # T2V uses the same sigma everywhere). ``timestep`` follows the + # framework's sigma*1000 convention. + sigma = (timestep.to(torch.float32) / 1000.0).view(batch_size, 1) + per_token_timestep = sigma.expand(batch_size, token_count).contiguous() + + return { + "hidden_states": noise_input, + "encoder_hidden_states": (text_dict["encoder_hidden_states"]), + # Post-connector embeddings are all-valid; the connector + # replaced pad positions with learnable registers. + "encoder_attention_mask": None, + "timestep": per_token_timestep, + "return_dict": False, + }