diff --git a/examples/train/configs/overfit_ltx2_3_t2v.yaml b/examples/train/configs/overfit_ltx2_3_t2v.yaml new file mode 100644 index 000000000..b86e29533 --- /dev/null +++ b/examples/train/configs/overfit_ltx2_3_t2v.yaml @@ -0,0 +1,95 @@ +# LTX-2.3 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 on an LTX-2.3 +# checkpoint (gated attention, cross-attention AdaLN, 4096-d +# post-connector text embeddings, no in-DiT caption projection). +# Uses the distilled checkpoint (validation is 8 sampling steps). +# +# Preprocess data first (writes data/ltx2_3_overfit_preprocessed): +# CUDA_VISIBLE_DEVICES=0 \ +# LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers \ +# LTX2_OVERFIT_OUTPUT_DIR=data/ltx2_3_overfit_preprocessed \ +# python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py +# +# Run: +# NUM_GPUS=4 bash examples/train/run.sh examples/train/configs/overfit_ltx2_3_t2v.yaml + +models: + student: + _target_: fastvideo.train.models.ltx2.LTX2Model + init_from: FastVideo/LTX-2.3-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_3_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_3_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_3_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_3_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 index 08f721da6..0ef6d4bcd 100644 --- a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py @@ -8,7 +8,10 @@ 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. +and training needs no text encoder at all. Captions are encoded via +the encoder's forward() (the exact inference path), which handles both +LTX-2.0 (shared 3840-d features) and LTX-2.3 (separate 4096-d video / +2048-d audio feature extractors). 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 @@ -20,6 +23,7 @@ Usage: import json import os +import shutil from typing import Any import av @@ -105,6 +109,9 @@ def main() -> None: model_index = verify_model_config_and_directory(model_path) os.makedirs(OUTPUT_DIR, exist_ok=True) + # The map-style dataset caches parquet file metadata; a stale cache + # next to a regenerated parquet can crash or serve old rows. + shutil.rmtree(os.path.join(OUTPUT_DIR, "map_style_cache"), ignore_errors=True) with open(os.path.join(DATA_DIR, CAPTION_JSON)) as f: caption_data = json.load(f) @@ -170,17 +177,16 @@ def main() -> None: 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, + # Encode through forward() — the inference text path. The + # two-step preprocess_text_embeddings + run_connectors route + # breaks on LTX-2.3, whose separate audio feature extractor + # is narrower than the video one. + text_inputs = tokenizer([preprocess_text(caption)], **tokenizer_kwargs) + encoder_out = text_encoder( + input_ids=text_inputs["input_ids"].to(device), + attention_mask=text_inputs["attention_mask"].to(device), ) - video_embeds, _audio_embeds, _post_mask = text_encoder.run_connectors( - pre_embeds, - attention_mask, - ) - text_embedding = video_embeds.squeeze(0).float().cpu() + text_embedding = encoder_out.last_hidden_state.squeeze(0).float().cpu() print(f" Text embedding shape: {text_embedding.shape}") record = { diff --git a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py index 1f7d228f0..180d707d2 100644 --- a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py +++ b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py @@ -1,16 +1,17 @@ -"""E2E overfit test for LTX-2 fine-tuning on the modular trainer (fastvideo/train). +"""E2E overfit tests 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. +Parametrized over LTX-2.0 and LTX-2.3 distilled checkpoints. Each case +downloads the single-sample cats dataset, preprocesses it with the +matching 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). +Requires the FastVideo LTX-2 distilled checkpoints (set HF_HOME to a +cache that contains them, or allow ~60GB of downloads per version). Guarded by FASTVIDEO_NIGHTLY=1 so the default test suite stays fast. """ @@ -35,11 +36,23 @@ 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" + +_CASES = { + "ltx2": { + "model": "FastVideo/LTX2-Distilled-Diffusers", + "config": "examples/train/configs/overfit_ltx2_t2v.yaml", + "prep_dir": Path(DATA_DIR) / "ltx2_overfit_preprocessed", + "out_dir": Path(DATA_DIR) / "outputs_ltx2_overfit", + }, + "ltx2_3": { + "model": "FastVideo/LTX-2.3-Distilled-Diffusers", + "config": "examples/train/configs/overfit_ltx2_3_t2v.yaml", + "prep_dir": Path(DATA_DIR) / "ltx2_3_overfit_preprocessed", + "out_dir": Path(DATA_DIR) / "outputs_ltx2_3_overfit", + }, +} def download_data(): @@ -55,17 +68,18 @@ def download_data(): f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") -def run_preprocessing(): +def run_preprocessing(case: dict): 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) + env["LTX2_OVERFIT_OUTPUT_DIR"] = str(case["prep_dir"]) + env["LTX2_OVERFIT_MODEL"] = case["model"] cmd = [sys.executable, PREPROCESSING_SCRIPT] - print(f"Running preprocessing: {cmd}") + print(f"Running preprocessing: {cmd} (model={case['model']})") subprocess.run(cmd, check=True, env=env) -def run_training(): +def run_training(case: dict): env = os.environ.copy() env.setdefault("WANDB_MODE", "offline") cmd = [ @@ -73,18 +87,18 @@ def run_training(): "--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), + "--config", case["config"], + "--training.checkpoint.output_dir", str(case["out_dir"]), + "--training.data.data_path", str(case["prep_dir"]), "--callbacks.validation.dataset_file", - str(LOCAL_PREPROCESSED_DATA_DIR / "validation_prompts.json"), + str(case["prep_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") +def _validation_videos_by_step(out_dir: Path) -> dict[int, str]: + pattern = os.path.join(str(out_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)) @@ -97,27 +111,29 @@ def _validation_videos_by_step() -> dict[int, str]: os.environ.get("FASTVIDEO_NIGHTLY") != "1", reason="nightly e2e overfit test; set FASTVIDEO_NIGHTLY=1 to run", ) -def test_e2e_ltx2_overfit_new_stack(): +@pytest.mark.parametrize("case_id", _CASES.keys()) +def test_e2e_ltx2_overfit_new_stack(case_id: str): + case = _CASES[case_id] download_data() - run_preprocessing() - run_training() + run_preprocessing(case) + run_training(case) - reference_video = LOCAL_PREPROCESSED_DATA_DIR / "training_sample_0.mp4" + reference_video = case["prep_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}" + videos = _validation_videos_by_step(case["out_dir"]) + assert videos, f"No validation videos found under {case['out_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"\n===== MS-SSIM vs training clip at step {final_step} ({case_id}) =====") print(f"Mean: {final_mean:.4f} Min: {final_min:.4f} Max: {final_max:.4f}") - results = {"final_step": final_step, "final_mean_ssim": final_mean} + results = {"case": case_id, "final_step": final_step, "final_mean_ssim": final_mean} baseline_step = min(videos) if baseline_step != final_step: base_mean, _, _ = compute_video_ssim_torchvision( @@ -136,4 +152,5 @@ def test_e2e_ltx2_overfit_new_stack(): if __name__ == "__main__": os.environ.setdefault("FASTVIDEO_NIGHTLY", "1") - test_e2e_ltx2_overfit_new_stack() + for _case_id in _CASES: + test_e2e_ltx2_overfit_new_stack(_case_id) diff --git a/fastvideo/tests/train/fixtures/ltx2_3_t2v_finetune_min.yaml b/fastvideo/tests/train/fixtures/ltx2_3_t2v_finetune_min.yaml new file mode 100644 index 000000000..88dc7fc5c --- /dev/null +++ b/fastvideo/tests/train/fixtures/ltx2_3_t2v_finetune_min.yaml @@ -0,0 +1,51 @@ +# Minimum config to run a single training step of FineTuneMethod on +# LTX2Model with an LTX-2.3 checkpoint (gated attention, cross-attn +# AdaLN, 4096-d post-connector text embeddings, no in-DiT caption +# projection) for the per-method smoke test. 18.9B params, ~38GB in +# bf16 — the test skips below 60GB GPU memory. Same constraints as +# LTX-2.0: training_cfg_rate=0 and precondition_outputs=false. + +models: + student: + _target_: fastvideo.train.models.ltx2.LTX2Model + init_from: FastVideo/LTX-2.3-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 765338d5b..ca07d583b 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_3_finetune": { + "GB200": 2.2158 + }, "test_ltx2_finetune": { "GB200": 0.3127 }, diff --git a/fastvideo/tests/train/methods/test_ltx2_finetune.py b/fastvideo/tests/train/methods/test_ltx2_finetune.py index f36ff970c..29b041e65 100644 --- a/fastvideo/tests/train/methods/test_ltx2_finetune.py +++ b/fastvideo/tests/train/methods/test_ltx2_finetune.py @@ -1,11 +1,13 @@ # 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). +Mirrors ``test_wan_finetune.py`` for the LTX-2 plugin, parametrized +over LTX-2.0 and LTX-2.3 checkpoints. LTX2-specific differences: the +synthetic ``raw_batch`` carries a single post-connector Gemma +embedding (3840-d for 2.0, 4096-d for 2.3 which has no in-DiT caption +projection) 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 @@ -30,12 +32,14 @@ from .grad_norm_regression import ( resolve_blocks, ) -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "ltx2_t2v_finetune_min.yaml") +_FIXTURE_DIR = Path(__file__).resolve().parent.parent / "fixtures" + +# id -> (fixture, post-connector text embedding width, grad-norm ref name) +_CASES = { + "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840, "test_ltx2_finetune"), + "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096, "test_ltx2_3_finetune"), +} -# Post-connector Gemma embedding width / length. -_LTX2_TEXT_DIM = 3840 _LTX2_TEXT_LEN = 1024 _MIN_GPU_MEMORY_GB = 60 @@ -51,6 +55,7 @@ def _gpu_too_small() -> bool: def _build_synthetic_batch( device: torch.device, dtype: torch.dtype, + text_dim: int, ) -> dict[str, torch.Tensor]: """Tiny synthetic ``raw_batch`` matching ``LTX2Model.prepare_batch``. @@ -62,7 +67,7 @@ def _build_synthetic_batch( "text_embedding": torch.randn(batch_size, _LTX2_TEXT_LEN, - _LTX2_TEXT_DIM, + text_dim, device=device, dtype=dtype), "text_attention_mask": @@ -73,13 +78,15 @@ def _build_synthetic_batch( @pytest.mark.usefixtures("distributed_setup") +@pytest.mark.parametrize("case", _CASES.keys()) def test_ltx2_finetune_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: + case: str, 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) + fixture_name, text_dim, ref_name = _CASES[case] + cfg = load_run_config(str(_FIXTURE_DIR / fixture_name)) device = torch.device("cuda:0") dtype = torch.bfloat16 @@ -107,7 +114,7 @@ def test_ltx2_finetune_single_train_step( ) method.on_train_start() - batch = _build_synthetic_batch(device, dtype) + batch = _build_synthetic_batch(device, dtype, text_dim) loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0) loss = loss_map["total_loss"] @@ -149,4 +156,4 @@ def test_ltx2_finetune_single_train_step( # 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) + check_grad_norm_regression(ref_name, model.transformer.model) diff --git a/fastvideo/tests/train/models/test_load_ltx2.py b/fastvideo/tests/train/models/test_load_ltx2.py index ff6f3eba1..f7066a45e 100644 --- a/fastvideo/tests/train/models/test_load_ltx2.py +++ b/fastvideo/tests/train/models/test_load_ltx2.py @@ -1,16 +1,17 @@ # 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 +Loads the real LTX-2.0 / LTX-2.3 distilled checkpoints (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 +[B, tokens] and a post-connector Gemma text embedding (3840-d for 2.0; +4096-d for 2.3, which has no in-DiT caption projection); it returns +the denoised x0 prediction. This mirrors the kwargs in ``LTX2Model._build_distill_input_kwargs``. """ @@ -33,12 +34,14 @@ 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") +_FIXTURE_DIR = Path(__file__).resolve().parent.parent / "fixtures" + +# id -> (fixture, post-connector text embedding width) +_CASES = { + "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840), + "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096), +} -# Post-connector Gemma embedding width / length. -_LTX2_TEXT_DIM = 3840 _LTX2_TEXT_LEN = 1024 _MIN_GPU_MEMORY_GB = 60 @@ -52,12 +55,14 @@ def _gpu_too_small() -> bool: @pytest.mark.usefixtures("distributed_setup") -def test_ltx2_model_loads_and_forwards(): +@pytest.mark.parametrize("case", _CASES.keys()) +def test_ltx2_model_loads_and_forwards(case: str): 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) + fixture_name, text_dim = _CASES[case] + cfg = load_run_config(str(_FIXTURE_DIR / fixture_name)) model = LTX2Model( init_from=cfg.models["student"]["init_from"], training_config=cfg.training, @@ -81,7 +86,7 @@ def test_ltx2_model_loads_and_forwards(): 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, + text_dim, device=device, dtype=dtype) timestep = torch.full((b, tokens), 0.5, device=device, dtype=torch.float32) diff --git a/fastvideo/train/models/ltx2/ltx2.py b/fastvideo/train/models/ltx2/ltx2.py index 7327e2f63..0e18e5e4e 100644 --- a/fastvideo/train/models/ltx2/ltx2.py +++ b/fastvideo/train/models/ltx2/ltx2.py @@ -272,6 +272,8 @@ class LTX2Model(WanModel): raise ValueError(f"Unknown latents_source: " f"{latents_source!r}") + self._check_text_embedding_dim(encoder_hidden_states) + # LTX-2 VAE encode() already applies per-channel normalization # (scaling_factor is 1.0); latents are used as stored. training_batch.latents = latents @@ -366,6 +368,28 @@ class LTX2Model(WanModel): # Internal helpers # ------------------------------------------------------------------ + def _check_text_embedding_dim(self, encoder_hidden_states: torch.Tensor) -> None: + """Fail fast when the parquet was preprocessed with a mismatched + text stack (LTX-2.0 stores 3840-d post-connector embeddings; 2.3 + has no in-DiT caption projection and expects 4096-d).""" + # Read the LOADED transformer's config: the transformer loader + # deepcopies pipeline_config.dit_config before applying the + # checkpoint's arch flags, so training_config.pipeline_config + # never sees version flags like caption_proj_before_connector. + arch = self.transformer.config.arch_config + if bool(getattr(arch, "caption_proj_before_connector", False)): + expected_dim = int(arch.cross_attention_dim) + else: + expected_dim = int(arch.caption_channels) + actual_dim = int(encoder_hidden_states.shape[-1]) + if actual_dim != expected_dim: + raise ValueError(f"text_embedding width {actual_dim} does not match the " + f"checkpoint's expected text context dim {expected_dim}. " + "The parquet was likely preprocessed with a different " + "LTX-2 version (2.0 stores 3840-d, 2.3 stores 4096-d); " + "re-run preprocess_ltx2_overfit.py with LTX2_OVERFIT_MODEL " + "set to the same checkpoint as models.student.init_from.") + def _update_rope_fps(self, infos: list[dict[str, Any]] | None) -> None: fps: float | None = None if infos: