From 2452022f156ff6dd55fbe5eacaba050d2f17189d Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:34:18 -0700 Subject: [PATCH] [new-model] Add LTX-2.5 inference support (#1704) Co-authored-by: Aryan Kumar Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: CodeRabbit Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- .gitignore | 1 + .../inference_schema_parity_inventory.yaml | 1 + docs/inference/ltx2_5.md | 64 +++ docs/inference/support_matrix.md | 7 + examples/inference/basic/basic_ltx2_5_i2av.py | 108 ++++ examples/inference/basic/basic_ltx2_5_t2av.py | 99 ++++ fastvideo/api/sampling_param.py | 3 + fastvideo/configs/models/dits/ltx2.py | 7 + fastvideo/configs/models/encoders/gemma.py | 2 + fastvideo/configs/models/vaes/ltx2vae.py | 1 + fastvideo/models/dits/ltx2.py | 106 +++- fastvideo/models/encoders/gemma.py | 87 ++- fastvideo/models/loader/component_loader.py | 26 + fastvideo/models/vaes/ltx2vae.py | 5 +- fastvideo/pipelines/basic/ltx2/presets.py | 89 ++- .../basic/ltx2/stages/ltx2_denoising.py | 168 +++++- fastvideo/pipelines/pipeline_batch_info.py | 3 + fastvideo/registry.py | 40 +- fastvideo/tests/api/test_presets.py | 10 +- mkdocs.yml | 1 + pyproject.toml | 2 + pyproject_other.toml | 4 +- .../convert_ltx2_weights.py | 508 ++++++++++++++++++ .../ltx2/test_ltx2_pipeline_smoke.py | 12 +- tests/local_tests/ltx2_5/PORT_STATUS.md | 99 ++++ tests/local_tests/ltx2_5/README.md | 151 ++++++ .../ltx2_5/test_ltx2_5_conversion.py | 399 ++++++++++++++ tests/local_tests/ltx2_5/test_ltx2_5_gemma.py | 65 +++ .../ltx2_5/test_ltx2_5_registry.py | 89 +++ .../local_tests/ltx2_5/test_ltx2_5_sampler.py | 151 ++++++ tests/local_tests/ltx2_5/test_ltx2_5_vae.py | 23 + .../test_ltx2_5_transformer_parity.py | 316 +++++++++++ 32 files changed, 2612 insertions(+), 35 deletions(-) create mode 100644 docs/inference/ltx2_5.md create mode 100644 examples/inference/basic/basic_ltx2_5_i2av.py create mode 100644 examples/inference/basic/basic_ltx2_5_t2av.py create mode 100644 tests/local_tests/ltx2_5/PORT_STATUS.md create mode 100644 tests/local_tests/ltx2_5/README.md create mode 100644 tests/local_tests/ltx2_5/test_ltx2_5_conversion.py create mode 100644 tests/local_tests/ltx2_5/test_ltx2_5_gemma.py create mode 100644 tests/local_tests/ltx2_5/test_ltx2_5_registry.py create mode 100644 tests/local_tests/ltx2_5/test_ltx2_5_sampler.py create mode 100644 tests/local_tests/ltx2_5/test_ltx2_5_vae.py create mode 100644 tests/local_tests/transformers/test_ltx2_5_transformer_parity.py diff --git a/.gitignore b/.gitignore index b65ab4377..2ca2abaf2 100644 --- a/.gitignore +++ b/.gitignore @@ -143,5 +143,6 @@ fastvideo/tests/ssim/.reference_videos_download.lock *.nvimlog .nvimlog .python-version +/LTX-2-Reference/ scripts/benchmarks/minimax_h3_pro6000/headline_results/ fastvideo/tests/ssim/.reference_videos_download.lock diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 87bbcd795..877bba25b 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -584,6 +584,7 @@ surfaces: ltx2_stg_scale_audio: request.extensions.ltx2.stg_scale_audio ltx2_stg_blocks_video: request.extensions.ltx2.stg_blocks_video ltx2_stg_blocks_audio: request.extensions.ltx2.stg_blocks_audio + ltx2_use_ancestral_sampler: request.extensions.ltx2.use_ancestral_sampler ltx2_images: request.extensions.ltx2.images ltx2_image_crf: request.stage_overrides.refine.image_crf ltx2_conditioning_latent_stage1: request.extensions.ltx2.conditioning_latent_stage1 diff --git a/docs/inference/ltx2_5.md b/docs/inference/ltx2_5.md new file mode 100644 index 000000000..de8865bbe --- /dev/null +++ b/docs/inference/ltx2_5.md @@ -0,0 +1,64 @@ +# LTX-2.5 inference + +FastVideo supports BF16 inference for the LTX-2.5 dev and distilled +transformers. Both text-to-video and image-to-video generate synchronized +video and audio. The native path supports sequence parallelism, component +offload, and `torch.compile`. + +## Convert the official checkpoint + +The gated [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5) +repository publishes separate transformer, packed Gemma 4, convolutional video +VAE, audio VAE/vocoder, and spatial upscaler files. Accept its license and +download those files, then convert them into one FastVideo model directory: + +```bash +python scripts/checkpoint_conversion/convert_ltx2_weights.py \ + --variant distilled \ + --transformer-source /weights/ltx-2.5-22b-distilled-transformer-bf16.safetensors \ + --text-encoder-source /weights/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors \ + --vae-source /weights/ltx-2.5-video-vae-conv-bf16.safetensors \ + --audio-vae-source /weights/ltx-2.5-audio-vae-bf16.safetensors \ + --spatial-upscaler-source /weights/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors \ + --output /models/LTX-2.5-Distilled-Diffusers +``` + +Use `--variant dev` and the dev transformer to convert the development model. +The converter emits a standard component directory and preserves LTX-2.5's +architecture metadata, packed tokenizer, joint audio components, and refine +configuration. The convolutional video VAE is the supported launch path. + +## Generate video and audio + +Run distilled text-to-audio-video across four GPUs with sequence parallelism: + +```bash +python examples/inference/basic/basic_ltx2_5_t2av.py \ + --model-path /models/LTX-2.5-Distilled-Diffusers \ + --prompt "A jazz trio performs in a candlelit club, synchronized live sound" \ + --num-gpus 4 \ + --torch-compile +``` + +Condition on a first frame with the I2AV example: + +```bash +python examples/inference/basic/basic_ltx2_5_i2av.py \ + --model-path /models/LTX-2.5-Distilled-Diffusers \ + --image /images/first-frame.png \ + --prompt "The train pulls away from the platform as its horn sounds" \ + --num-gpus 4 \ + --torch-compile +``` + +Pass `--variant dev` with a converted dev directory to use the 30-step dev +guidance preset. Distilled inference uses the official eight-step ancestral +schedule and a three-step spatial refinement pass by default. + +## Current scope + +The initial inference path includes the native transformer, packed Gemma 4 +text stack, convolutional video VAE, audio VAE/vocoder, dev guidance, and the +distilled ancestral sampler. DiffVAE/NATTEN, generated keyframes, temporal +upsampling, automatic duration selection, HDR, training, fine-tuning, and +quantized deployment are separate follow-up work. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index a4adcbeae..587e8000c 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -66,6 +66,8 @@ column links a runnable script in `examples/inference/basic/` where one exists. | longcat | `FastVideo/LongCat-Video-I2V-Diffusers` | I2V | [basic_longcat_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_i2v.py) | | longcat | `FastVideo/LongCat-Video-VC-Diffusers` | — | [basic_longcat_vc.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_vc.py) | | ltx2 | `FastVideo/LTX2-Distilled-Diffusers`
`FastVideo/LTX2.3-Distilled-Diffusers`
`FastVideo/LTX-2.3-Distilled-Diffusers` | T2V | [basic_ltx2_distilled.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_distilled.py) | +| ltx2 | `FastVideo/LTX-2.5-Distilled-Diffusers` | T2V, I2V | [T2AV](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_5_t2av.py)
[I2AV](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_5_i2av.py) | +| ltx2 | `FastVideo/LTX-2.5-Dev-Diffusers` | T2V, I2V | [T2AV](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_5_t2av.py)
[I2AV](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_5_i2av.py) | | ltx2 | `Lightricks/LTX-2.3`
`FastVideo/LTX2.3-base`
`FastVideo/LTX2.3-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) | | ltx2 | `Lightricks/LTX-2`
`FastVideo/LTX2-base`
`FastVideo/LTX2-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) | | mmaudio | `FastVideo/MMAudio-large-44k-v2-Diffusers` | V2A, T2A | [basic_mmaudio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mmaudio.py) | @@ -116,6 +118,11 @@ and distilled) generate video with audio from text, optionally plus an image (se `kandinsky6_sr` IDs upscale an existing video (see the [Video SR guide](kandinsky6_sr.md)). +**Note (LTX-2.5)**: T2V and I2V produce synchronized video and audio. The +official `Lightricks/LTX-2.5` repository uses gated, split component files, so +it must first be converted to FastVideo's component layout. See the +[LTX-2.5 inference guide](ltx2_5.md). + **Note (Wan-VACE)**: not currently supported — no VACE pipeline or registered model ID exists on `main` ([#1435](https://github.com/hao-ai-lab/FastVideo/issues/1435)). The closest diff --git a/examples/inference/basic/basic_ltx2_5_i2av.py b/examples/inference/basic/basic_ltx2_5_i2av.py new file mode 100644 index 000000000..8b5907173 --- /dev/null +++ b/examples/inference/basic/basic_ltx2_5_i2av.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate synchronized video and audio from an image and text with LTX-2.5.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from fastvideo import VideoGenerator +from fastvideo.api import ( + CompileConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + SamplingConfig, +) + +MODEL_IDS = { + "distilled": "FastVideo/LTX-2.5-Distilled-Diffusers", + "dev": "FastVideo/LTX-2.5-Dev-Diffusers", +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--variant", choices=tuple(MODEL_IDS), default="distilled") + parser.add_argument( + "--model-path", + help="FastVideo-converted model directory or Hub ID; the raw split Lightricks repo is not directly loadable.", + ) + parser.add_argument("--image", required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--output", default="outputs/ltx2_5_i2av.mp4") + parser.add_argument("--height", type=int) + parser.add_argument("--width", type=int) + parser.add_argument("--num-frames", type=int, default=121) + parser.add_argument("--steps", type=int) + parser.add_argument("--seed", type=int, default=10) + parser.add_argument("--num-gpus", type=int, default=4) + parser.add_argument("--torch-compile", action="store_true", help="Compile the transformer denoising path") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + image_path = Path(args.image) + if not image_path.is_file(): + raise SystemExit(f"Conditioning image not found: {image_path}") + + is_distilled = args.variant == "distilled" + model_path = args.model_path or MODEL_IDS[args.variant] + height = args.height or (1024 if is_distilled else 512) + width = args.width or (1536 if is_distilled else 768) + steps = args.steps or (8 if is_distilled else 30) + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=model_path, + engine=EngineConfig( + num_gpus=args.num_gpus, + use_fsdp_inference=args.num_gpus > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus), + offload=OffloadConfig( + dit=False, + dit_layerwise=False, + text_encoder=True, + vae=True, + pin_cpu_memory=False, + ), + compile=CompileConfig(enabled=args.torch_compile), + ), + pipeline=PipelineSelection(workload_type="i2v"), + )) + try: + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + sampling=SamplingConfig( + seed=args.seed, + height=height, + width=width, + num_frames=args.num_frames, + fps=24, + num_inference_steps=steps, + guidance_scale=1.0 if is_distilled else 3.0, + ), + output=OutputConfig( + output_path=str(output_path), + save_video=True, + return_frames=False, + ), + extensions={ + "ltx2_images": [(str(image_path), 0, 1.0)], + "ltx2_image_crf": 18.0, + }, + )) + print(f"Synchronized video and audio written to: {result.video_path}") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_ltx2_5_t2av.py b/examples/inference/basic/basic_ltx2_5_t2av.py new file mode 100644 index 000000000..610e2e107 --- /dev/null +++ b/examples/inference/basic/basic_ltx2_5_t2av.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate synchronized video and audio from text with LTX-2.5.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from fastvideo import VideoGenerator +from fastvideo.api import ( + CompileConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + SamplingConfig, +) + +MODEL_IDS = { + "distilled": "FastVideo/LTX-2.5-Distilled-Diffusers", + "dev": "FastVideo/LTX-2.5-Dev-Diffusers", +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--variant", choices=tuple(MODEL_IDS), default="distilled") + parser.add_argument( + "--model-path", + help="FastVideo-converted model directory or Hub ID; the raw split Lightricks repo is not directly loadable.", + ) + parser.add_argument("--prompt", required=True) + parser.add_argument("--output", default="outputs/ltx2_5_t2av.mp4") + parser.add_argument("--height", type=int) + parser.add_argument("--width", type=int) + parser.add_argument("--num-frames", type=int, default=121) + parser.add_argument("--steps", type=int) + parser.add_argument("--seed", type=int, default=10) + parser.add_argument("--num-gpus", type=int, default=4) + parser.add_argument("--torch-compile", action="store_true", help="Compile the transformer denoising path") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + is_distilled = args.variant == "distilled" + model_path = args.model_path or MODEL_IDS[args.variant] + height = args.height or (1024 if is_distilled else 512) + width = args.width or (1536 if is_distilled else 768) + steps = args.steps or (8 if is_distilled else 30) + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=model_path, + engine=EngineConfig( + num_gpus=args.num_gpus, + use_fsdp_inference=args.num_gpus > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus), + offload=OffloadConfig( + dit=False, + dit_layerwise=False, + text_encoder=True, + vae=True, + pin_cpu_memory=False, + ), + compile=CompileConfig(enabled=args.torch_compile), + ), + pipeline=PipelineSelection(workload_type="t2v"), + )) + try: + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + sampling=SamplingConfig( + seed=args.seed, + height=height, + width=width, + num_frames=args.num_frames, + fps=24, + num_inference_steps=steps, + guidance_scale=1.0 if is_distilled else 3.0, + ), + output=OutputConfig( + output_path=str(output_path), + save_video=True, + return_frames=False, + ), + )) + print(f"Synchronized video and audio written to: {result.video_path}") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 124fdd8f0..4cc8e84bc 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -152,6 +152,9 @@ class SamplingParam: ltx2_stg_scale_audio: float = 0.0 ltx2_stg_blocks_video: list[int] = field(default_factory=lambda: [29]) ltx2_stg_blocks_audio: list[int] = field(default_factory=lambda: [29]) + # LTX-2.5 distilled stage 1 uses the official variance-preserving + # ancestral Euler update. Older LTX generations remain deterministic. + ltx2_use_ancestral_sampler: bool = False # LTX-2 image / video / continuation conditioning. These flow from # generate_video(...) kwargs through ``sampling_param.update(kwargs)`` diff --git a/fastvideo/configs/models/dits/ltx2.py b/fastvideo/configs/models/dits/ltx2.py index 135d4ca6e..04150fad5 100644 --- a/fastvideo/configs/models/dits/ltx2.py +++ b/fastvideo/configs/models/dits/ltx2.py @@ -55,6 +55,13 @@ class LTX2VideoArchConfig(DiTArchConfig): # LTX-2.3 gated extensions. All default OFF == LTX-2.0 behavior. cross_attention_adaln: bool = False caption_proj_before_connector: bool = False + # LTX-2.5 architecture gates. Defaults preserve all pre-2.5 checkpoints: + # prompt cross-attention remains timestep-dependent, FFNs retain their + # biases, and no keyframe-only parameter is added to the state dict. + use_prompt_adaln_single: bool = True + ff_bias: bool = True + audio_ff_bias: bool = True + use_keyframes_abs_pos_embedding: bool = False positional_embedding_theta: float = 10000.0 positional_embedding_max_pos: list[int] = field(default_factory=lambda: [20, 2048, 2048]) diff --git a/fastvideo/configs/models/encoders/gemma.py b/fastvideo/configs/models/encoders/gemma.py index 3be16e439..ef37014b7 100644 --- a/fastvideo/configs/models/encoders/gemma.py +++ b/fastvideo/configs/models/encoders/gemma.py @@ -30,6 +30,7 @@ class LTX2GemmaArchConfig(TextEncoderArchConfig): num_attention_heads: int = 30 text_len: int = 1024 pad_token_id: int = 0 + bos_token_id: int = 2 eos_token_id: int = 2 gemma_model_path: str = "" @@ -58,6 +59,7 @@ class LTX2GemmaArchConfig(TextEncoderArchConfig): connector_rope_type: str = "split" connector_double_precision_rope: bool = False connector_apply_gated_attention: bool = False + connector_ff_bias: bool = True connector_num_learnable_registers: int | None = 128 _fsdp_shard_conditions: list = field( diff --git a/fastvideo/configs/models/vaes/ltx2vae.py b/fastvideo/configs/models/vaes/ltx2vae.py index c4ac77f7c..b0ac9b142 100644 --- a/fastvideo/configs/models/vaes/ltx2vae.py +++ b/fastvideo/configs/models/vaes/ltx2vae.py @@ -26,6 +26,7 @@ class LTX2VAEArchConfig(VAEArchConfig): decoder_spatial_padding_mode: str = "reflect" causal_decoder: bool = False timestep_conditioning: bool = True + decoder_base_channels: int = 128 use_quant_conv: bool = False scaling_factor: float = 1.0 normalize_latent_channels: bool = False diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index e2cbf3ce2..9bb18a07a 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -332,6 +332,7 @@ class GELUApprox(nn.Module): self, in_features: int, out_features: int, + bias: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = "", ): @@ -339,6 +340,7 @@ class GELUApprox(nn.Module): self.proj = ReplicatedLinear( in_features, out_features, + bias=bias, quant_config=quant_config, prefix=f"{prefix}.fc_in", ) @@ -356,6 +358,7 @@ class FeedForward(nn.Module): dim: int, dim_out: int, mult: int = 4, + bias: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: @@ -364,12 +367,14 @@ class FeedForward(nn.Module): project_in = GELUApprox( dim, inner_dim, + bias=bias, quant_config=quant_config, prefix=f"{prefix}.ffn", ) project_out = ReplicatedLinear( inner_dim, dim_out, + bias=bias, quant_config=quant_config, prefix=f"{prefix}.ffn.fc_out", ) @@ -918,6 +923,26 @@ class Modality: context_mask: torch.Tensor | None = None # LTX-2.3 cross-attention AdaLN sigma timestep (None for 2.0). sigma: torch.Tensor | None = None + # LTX-2.5 per-token keyframe marker, shaped (B, T, 1). + keyframes_mask: torch.Tensor | None = None + + +KeyframesEmbeddingProvider = Callable[[], torch.Tensor | None] + + +def apply_keyframes_absolute_embedding( + hidden_states: torch.Tensor, + keyframes_mask: torch.Tensor | None, + embedding_provider: KeyframesEmbeddingProvider | None, +) -> torch.Tensor: + """Apply the learned LTX-2.5 marker to projected keyframe tokens.""" + if embedding_provider is None or keyframes_mask is None: + return hidden_states + embedding = embedding_provider() + if embedding is None: + return hidden_states + mask = (keyframes_mask > 0).to(dtype=hidden_states.dtype) + return hidden_states + mask * embedding.to(dtype=hidden_states.dtype) class TransformerArgsPreprocessor: @@ -937,6 +962,7 @@ class TransformerArgsPreprocessor: positional_embedding_theta: float, rope_type: LTXRopeType, prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None, ) -> None: self.patchify_proj = patchify_proj self.adaln = adaln @@ -951,6 +977,7 @@ class TransformerArgsPreprocessor: self.rope_type = rope_type # LTX-2.3 cross-attention AdaLN prompt timestep embedder (None for 2.0). self.prompt_adaln = prompt_adaln + self.keyframes_embedding_provider = keyframes_embedding_provider def _prepare_timestep( self, @@ -1025,6 +1052,11 @@ class TransformerArgsPreprocessor: if not latent.is_contiguous(): latent = latent.contiguous() x = self.patchify_proj(latent) + x = apply_keyframes_absolute_embedding( + x, + modality.keyframes_mask, + self.keyframes_embedding_provider, + ) timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype, self.adaln) prompt_timestep = None @@ -1082,6 +1114,7 @@ class MultiModalTransformerArgsPreprocessor: rope_type: LTXRopeType, av_ca_timestep_scale_multiplier: int, prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None, ) -> None: self.simple_preprocessor = TransformerArgsPreprocessor( patchify_proj=patchify_proj, @@ -1096,6 +1129,7 @@ class MultiModalTransformerArgsPreprocessor: positional_embedding_theta=positional_embedding_theta, rope_type=rope_type, prompt_adaln=prompt_adaln, + keyframes_embedding_provider=keyframes_embedding_provider, ) self.cross_scale_shift_adaln = cross_scale_shift_adaln self.cross_gate_adaln = cross_gate_adaln @@ -1163,6 +1197,7 @@ class TransformerConfig: # LTX-2.3 gated extensions (default OFF == LTX-2.0 behavior). apply_gated_attention: bool = False cross_attention_adaln: bool = False + ff_bias: bool = True class LTXDistributedAttention(DistributedAttention): @@ -1803,6 +1838,7 @@ class BasicAVTransformerBlock(torch.nn.Module): self.ff = FeedForward( video.dim, dim_out=video.dim, + bias=video.ff_bias, quant_config=quant_config, prefix=f"{prefix}.blocks.{idx}", ) @@ -1841,6 +1877,7 @@ class BasicAVTransformerBlock(torch.nn.Module): self.audio_ff = FeedForward( audio.dim, dim_out=audio.dim, + bias=audio.ff_bias, quant_config=quant_config, prefix=f"{prefix}.blocks.{idx}.audio", ) @@ -2247,11 +2284,16 @@ def apply_cross_attention_adaln( Only reached when ``cross_attention_adaln`` is enabled, so it never affects the LTX-2.0 path. """ - if prompt_scale_shift_table is None or prompt_timestep is None: - raise ValueError("cross_attention_adaln requires prompt scale/shift tables and prompt_timestep.") + if prompt_scale_shift_table is None: + raise ValueError("cross_attention_adaln requires a prompt scale/shift table.") batch_size = x.shape[0] - shift_kv, scale_kv = (prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + - prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)).unbind(dim=2) + kv_modulation = prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + # LTX-2.5 can disable the prompt-side AdaLN MLP. In that mode the static + # per-block table is the entire K/V modulation, making it timestep + # independent and cacheable. Older checkpoints add the dynamic term. + if prompt_timestep is not None: + kv_modulation = kv_modulation + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1) + shift_kv, scale_kv = kv_modulation.unbind(dim=2) attn_input = _rms_norm_dispatch(x, eps=norm_eps) * (1 + q_scale) + q_shift encoder_hidden_states = context * (1 + scale_kv) + shift_kv return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate @@ -2299,6 +2341,10 @@ class LTXModel(torch.nn.Module): rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, double_precision_rope: bool = False, cross_attention_adaln: bool = False, + use_prompt_adaln_single: bool = True, + ff_bias: bool = True, + audio_ff_bias: bool = True, + use_keyframes_abs_pos_embedding: bool = False, caption_proj_before_connector: bool = False, apply_gated_attention: bool = False, stg_block_idx: int = 29, @@ -2310,6 +2356,10 @@ class LTXModel(torch.nn.Module): self._enable_gradient_checkpointing = False # LTX-2.3 gated extensions (all default OFF == LTX-2.0 behavior). self.cross_attention_adaln = cross_attention_adaln + self.use_prompt_adaln_single = use_prompt_adaln_single + self.ff_bias = ff_bias + self.audio_ff_bias = audio_ff_bias + self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding self.caption_proj_before_connector = caption_proj_before_connector self.apply_gated_attention = apply_gated_attention self.stg_block_idx = stg_block_idx @@ -2368,6 +2418,10 @@ class LTXModel(torch.nn.Module): prefix=prefix, ) + def _keyframes_embedding(self) -> torch.Tensor | None: + """Resolve the optional parameter at call time, matching upstream.""" + return getattr(self, "keyframes_abs_pos_embedding", None) + def _init_video( self, in_channels: int, @@ -2377,6 +2431,9 @@ class LTXModel(torch.nn.Module): create_caption_projection: bool = True, ) -> None: self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) + self.keyframes_abs_pos_embedding = ( + torch.nn.Parameter(torch.zeros(1, self.inner_dim)) if self.use_keyframes_abs_pos_embedding else None + ) self.adaln_single = AdaLayerNormSingle( self.inner_dim, embedding_coefficient=adaln_embedding_coefficient(self.cross_attention_adaln), @@ -2388,11 +2445,10 @@ class LTXModel(torch.nn.Module): in_features=caption_channels, hidden_size=self.inner_dim, ) - if self.cross_attention_adaln: - self.prompt_adaln_single = AdaLayerNormSingle( - self.inner_dim, - embedding_coefficient=2, - ) + self.prompt_adaln_single = ( + AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) + if self.cross_attention_adaln and self.use_prompt_adaln_single else None + ) self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim)) self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps) self.proj_out = torch.nn.Linear(self.inner_dim, out_channels) @@ -2415,11 +2471,10 @@ class LTXModel(torch.nn.Module): in_features=caption_channels, hidden_size=self.audio_inner_dim, ) - if self.cross_attention_adaln: - self.audio_prompt_adaln_single = AdaLayerNormSingle( - self.audio_inner_dim, - embedding_coefficient=2, - ) + self.audio_prompt_adaln_single = ( + AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) + if self.cross_attention_adaln and self.use_prompt_adaln_single else None + ) self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim)) self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps) self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels) @@ -2462,6 +2517,7 @@ class LTXModel(torch.nn.Module): rope_type=self.rope_type, av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, prompt_adaln=getattr(self, "prompt_adaln_single", None), + keyframes_embedding_provider=self._keyframes_embedding, ) self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor( patchify_proj=self.audio_patchify_proj, @@ -2496,6 +2552,7 @@ class LTXModel(torch.nn.Module): positional_embedding_theta=self.positional_embedding_theta, rope_type=self.rope_type, prompt_adaln=getattr(self, "prompt_adaln_single", None), + keyframes_embedding_provider=self._keyframes_embedding, ) elif self.model_type.is_audio_enabled(): self.audio_args_preprocessor = TransformerArgsPreprocessor( @@ -2532,6 +2589,7 @@ class LTXModel(torch.nn.Module): context_dim=cross_attention_dim, apply_gated_attention=self.apply_gated_attention, cross_attention_adaln=self.cross_attention_adaln, + ff_bias=self.ff_bias, ) if self.model_type.is_video_enabled() else None) audio_config = (TransformerConfig( dim=self.audio_inner_dim, @@ -2540,6 +2598,7 @@ class LTXModel(torch.nn.Module): context_dim=audio_cross_attention_dim, apply_gated_attention=self.apply_gated_attention, cross_attention_adaln=self.cross_attention_adaln, + ff_bias=self.audio_ff_bias, ) if self.model_type.is_audio_enabled() else None) self.use_distributed_attention = use_distributed_attention self.transformer_blocks = torch.nn.ModuleList([ @@ -2716,6 +2775,10 @@ class LTX2Transformer3DModel(BaseDiT): audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos, av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier, cross_attention_adaln=arch.cross_attention_adaln, + use_prompt_adaln_single=arch.use_prompt_adaln_single, + ff_bias=arch.ff_bias, + audio_ff_bias=arch.audio_ff_bias, + use_keyframes_abs_pos_embedding=arch.use_keyframes_abs_pos_embedding, caption_proj_before_connector=arch.caption_proj_before_connector, apply_gated_attention=arch.apply_gated_attention, stg_block_idx=arch.stg_block_idx, @@ -2752,6 +2815,7 @@ class LTX2Transformer3DModel(BaseDiT): audio_encoder_attention_mask: torch.Tensor | None = None, video_sigma: torch.Tensor | None = None, audio_sigma: torch.Tensor | None = None, + keyframes_mask: torch.Tensor | None = None, skip_cross_modal_attn: bool = False, skip_video_self_attn_blocks: list[int] | None = None, skip_audio_self_attn_blocks: list[int] | None = None, @@ -2790,10 +2854,23 @@ class LTX2Transformer3DModel(BaseDiT): video_original_seq_len = latents.shape[1] video_padded_seq_len = video_original_seq_len video_timestep = timestep + if keyframes_mask is not None: + if keyframes_mask.ndim != 3 or keyframes_mask.shape[-1] != 1: + raise ValueError( + "keyframes_mask must contain patchified per-token markers with shape (B, T, 1)" + ) + if keyframes_mask.shape[:2] != latents.shape[:2]: + raise ValueError( + "keyframes_mask batch/sequence dimensions must match the patchified video latents: " + f"mask={tuple(keyframes_mask.shape)}, latents={tuple(latents.shape)}" + ) + keyframes_mask = keyframes_mask.to(device=latents.device) if sp_world_size > 1: latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1) # Shard timestep along sequence dimension (timestep has shape [batch, seq_len]) video_timestep, _ = sequence_model_parallel_shard(timestep, dim=1) + if keyframes_mask is not None: + keyframes_mask, _ = sequence_model_parallel_shard(keyframes_mask, dim=1) current_seq_len = latents.shape[1] video_padded_seq_len = current_seq_len * sp_world_size # Compute RoPE positions for the FULL sequence (before sharding) @@ -2825,6 +2902,7 @@ class LTX2Transformer3DModel(BaseDiT): context=encoder_hidden_states, context_mask=encoder_attention_mask, sigma=video_sigma, + keyframes_mask=keyframes_mask, ) # Process audio modality if provided diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index 5b36d1004..a5752e30e 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/fastvideo/models/encoders/gemma.py @@ -7,7 +7,7 @@ from typing import Any, Iterable import torch from torch import nn -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration +from transformers import AutoModelForImageTextToText, AutoTokenizer, PreTrainedModel import fastvideo.envs as envs from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig @@ -50,6 +50,7 @@ class GemmaConnectorConfig: num_learnable_registers: int | None # LTX-2.3 connector gated attention (default False == LTX-2.0 behavior). apply_gated_attention: bool = False + ff_bias: bool = True class GemmaFeaturesExtractorProjLinear(nn.Module): @@ -93,6 +94,7 @@ class _BasicTransformerBlock1D(nn.Module): dim_head: int, rope_type: LTXRopeType, apply_gated_attention: bool = False, + ff_bias: bool = True, norm_eps: float = 1e-6, ) -> None: super().__init__() @@ -105,7 +107,7 @@ class _BasicTransformerBlock1D(nn.Module): rope_type=rope_type, apply_gated_attention=apply_gated_attention, ) - self.ff = FeedForward(dim, dim_out=dim) + self.ff = FeedForward(dim, dim_out=dim, bias=ff_bias) self.norm_eps = norm_eps def forward( @@ -237,6 +239,7 @@ class Embeddings1DConnector(nn.Module): dim_head=config.attention_head_dim, rope_type=config.rope_type, apply_gated_attention=config.apply_gated_attention, + ff_bias=config.ff_bias, ) for _ in range(config.num_layers) ]) self.num_learnable_registers = config.num_learnable_registers @@ -357,6 +360,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): ) connector_apply_gated_attention = bool(getattr(arch, "connector_apply_gated_attention", False)) + connector_ff_bias = bool(getattr(arch, "connector_ff_bias", True)) video_connector_config = GemmaConnectorConfig( num_attention_heads=arch.connector_num_attention_heads, attention_head_dim=arch.connector_attention_head_dim, @@ -367,6 +371,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): double_precision_rope=arch.connector_double_precision_rope, num_learnable_registers=arch.connector_num_learnable_registers, apply_gated_attention=connector_apply_gated_attention, + ff_bias=connector_ff_bias, ) audio_connector_config = GemmaConnectorConfig( num_attention_heads=getattr(arch, "audio_connector_num_attention_heads", None) @@ -380,6 +385,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): double_precision_rope=arch.connector_double_precision_rope, num_learnable_registers=arch.connector_num_learnable_registers, apply_gated_attention=connector_apply_gated_attention, + ff_bias=connector_ff_bias, ) self.embeddings_connector = Embeddings1DConnector(video_connector_config) self.audio_embeddings_connector = Embeddings1DConnector(audio_connector_config) @@ -387,11 +393,11 @@ class LTX2GemmaTextEncoderModel(TextEncoder): self.gemma_model_path = arch.gemma_model_path self.gemma_dtype = arch.gemma_dtype self.padding_side = arch.padding_side - self._gemma_model: Gemma3ForConditionalGeneration | None = None + self._gemma_model: PreTrainedModel | None = None def named_parameters(self, prefix: str = "", recurse: bool = True): for name, param in super().named_parameters(prefix=prefix, recurse=recurse): - if name.startswith("gemma_model."): + if name.startswith("_gemma_model."): continue yield name, param @@ -425,13 +431,13 @@ class LTX2GemmaTextEncoderModel(TextEncoder): ) @property - def gemma_model(self) -> Gemma3ForConditionalGeneration: + def gemma_model(self) -> PreTrainedModel: if self._gemma_model is None: gemma_path = self.gemma_model_path if not gemma_path: raise ValueError("gemma_model_path must be set (expected text_encoder/gemma).") dtype = getattr(torch, self.gemma_dtype, torch.bfloat16) - self._gemma_model = Gemma3ForConditionalGeneration.from_pretrained( + self._gemma_model = AutoModelForImageTextToText.from_pretrained( gemma_path, local_files_only=True, torch_dtype=dtype, @@ -554,6 +560,11 @@ class LTX2GemmaTextEncoderModel(TextEncoder): input_ids = text_inputs["input_ids"].to(device=model.device) attention_mask = text_inputs["attention_mask"].to(device=model.device) + input_ids, attention_mask = _ensure_leading_bos( + input_ids, + attention_mask, + bos_token_id=_get_bos_token_id(model.config), + ) outputs = model( input_ids=input_ids, attention_mask=attention_mask, @@ -590,6 +601,11 @@ class LTX2GemmaTextEncoderModel(TextEncoder): attention_mask = torch.ones_like(input_ids) model = self.gemma_model + input_ids, attention_mask = _ensure_leading_bos( + input_ids, + attention_mask, + bos_token_id=_get_bos_token_id(model.config), + ) target_device = get_local_torch_device() # Do not invoke model.to() inside the compiled forward path. # _parse_to returns a non-Tensor torch.device, which Dynamo cannot @@ -695,5 +711,64 @@ def _norm_and_concat_padded_batch( return normed +def _get_bos_token_id(config: Any) -> int | None: + bos_token_id = getattr(config, "bos_token_id", None) + if bos_token_id is None: + bos_token_id = getattr(getattr(config, "text_config", None), "bos_token_id", None) + return int(bos_token_id) if bos_token_id is not None else None + + +def _ensure_leading_bos( + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + *, + bos_token_id: int | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Ensure each left-padded Gemma sequence starts with a valid BOS token. + + Gemma 3 tokenizers generally add BOS themselves, while the packed Gemma 4 + tokenizer used by LTX-2.5 does not. This keeps the legacy path unchanged + and adds BOS only to rows that need it, matching the official LTX wrapper. + """ + if bos_token_id is None: + raise ValueError("LTX-2 Gemma tokenizer/config must define bos_token_id.") + if input_ids.ndim != 2 or attention_mask.shape != input_ids.shape: + raise ValueError("LTX-2 Gemma input_ids and attention_mask must have matching rank-2 shapes.") + + valid = attention_mask.bool() + has_valid = valid.any(dim=1) + first_valid = valid.to(torch.int64).argmax(dim=1) + first_tokens = input_ids.gather(1, first_valid.unsqueeze(1)).squeeze(1) + needs_bos = ~has_valid | first_tokens.ne(bos_token_id) + + adjusted_ids = input_ids.clone() + adjusted_mask = attention_mask.clone() + + # Rows with left padding can consume one padding slot without moving any + # text tokens. Empty rows place BOS in the final slot. + target_positions = torch.where(has_valid, first_valid - 1, torch.full_like(first_valid, input_ids.shape[1] - 1)) + consume_padding = needs_bos & (~has_valid | first_valid.gt(0)) + safe_targets = target_positions.clamp(min=0) + adjusted_ids.scatter_(1, safe_targets.unsqueeze(1), torch.where( + consume_padding, + torch.full_like(safe_targets, bos_token_id), + adjusted_ids.gather(1, safe_targets.unsqueeze(1)).squeeze(1), + ).unsqueeze(1)) + adjusted_mask.scatter_(1, safe_targets.unsqueeze(1), torch.where( + consume_padding, + torch.ones_like(safe_targets, dtype=adjusted_mask.dtype), + adjusted_mask.gather(1, safe_targets.unsqueeze(1)).squeeze(1), + ).unsqueeze(1)) + + # A completely full row has no padding slot; prepend BOS and truncate its + # last token exactly as the official max-length tokenizer wrapper does. + shift_rows = needs_bos & has_valid & first_valid.eq(0) + shifted_ids = torch.cat((torch.full_like(input_ids[:, :1], bos_token_id), input_ids[:, :-1]), dim=1) + shifted_mask = torch.cat((torch.ones_like(attention_mask[:, :1]), attention_mask[:, :-1]), dim=1) + adjusted_ids = torch.where(shift_rows.unsqueeze(1), shifted_ids, adjusted_ids) + adjusted_mask = torch.where(shift_rows.unsqueeze(1), shifted_mask, adjusted_mask) + return adjusted_ids, adjusted_mask + + # Entry point for model registry EntryClass = LTX2GemmaTextEncoderModel diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 03427e9d3..467beac3a 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -290,6 +290,32 @@ class TextEncoderLoader(ComponentLoader): if gemma_path and not gemma_path_from_candidate: if not os.path.isabs(gemma_path): model_config["gemma_model_path"] = os.path.normpath(os.path.join(repo_root, gemma_path)) + resolved_gemma_path = model_config.get("gemma_model_path", gemma_path) + gemma_config_path = os.path.join(resolved_gemma_path, "config.json") if resolved_gemma_path else "" + if os.path.isfile(gemma_config_path): + try: + with open(gemma_config_path, encoding="utf-8") as f: + gemma_config = json.load(f) + gemma_text_config = gemma_config.get("text_config", gemma_config) + gemma_hidden_size = gemma_text_config.get("hidden_size") + gemma_num_hidden_layers = gemma_text_config.get("num_hidden_layers") + if gemma_hidden_size is not None and gemma_num_hidden_layers is not None: + # LTX feature extraction stacks the embedding output plus + # every Gemma transformer layer. Derive this from the + # packed Gemma config so Gemma 4 never inherits Gemma 3's + # hard-coded 3840x49 geometry. + model_config["hidden_size"] = int(gemma_hidden_size) + model_config["num_hidden_layers"] = int(gemma_num_hidden_layers) + model_config["feature_extractor_in_features"] = (int(gemma_hidden_size) * + (int(gemma_num_hidden_layers) + 1)) + for field_name in ("num_attention_heads", "pad_token_id", "eos_token_id"): + value = gemma_text_config.get(field_name, gemma_config.get(field_name)) + if value is not None: + if field_name == "eos_token_id" and isinstance(value, list): + value = value[0] if value else 2 + model_config[field_name] = value + except (json.JSONDecodeError, OSError, TypeError, ValueError) as exc: + logger.warning("Unable to derive LTX Gemma geometry from %s: %s", gemma_config_path, exc) transformer_config_path = os.path.join(repo_root, "transformer", "config.json") if os.path.isfile(transformer_config_path): try: diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py index 8556c1bea..af347064f 100644 --- a/fastvideo/models/vaes/ltx2vae.py +++ b/fastvideo/models/vaes/ltx2vae.py @@ -1327,6 +1327,7 @@ class VideoDecoder(nn.Module): causal: bool = False, timestep_conditioning: bool = False, decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT, + base_channels: int = 128, ): super().__init__() @@ -1341,7 +1342,7 @@ class VideoDecoder(nn.Module): self.decode_noise_scale = 0.025 self.decode_timestep = 0.05 - feature_channels = in_channels + feature_channels = base_channels for block_name, block_params in list(reversed(decoder_blocks)): block_config = block_params if isinstance(block_params, dict) else {} if block_name == "res_x_y": @@ -1521,6 +1522,7 @@ class VideoDecoderConfigurator: norm_layer_str = config.get("norm_layer", "pixel_norm") causal = config.get("causal_decoder", False) timestep_conditioning = config.get("timestep_conditioning", True) + base_channels = config.get("decoder_base_channels", 128) return VideoDecoder( convolution_dimensions=convolution_dimensions, @@ -1532,6 +1534,7 @@ class VideoDecoderConfigurator: causal=causal, timestep_conditioning=timestep_conditioning, decoder_spatial_padding_mode=decoder_spatial_padding_mode, + base_channels=base_channels, ) diff --git a/fastvideo/pipelines/basic/ltx2/presets.py b/fastvideo/pipelines/basic/ltx2/presets.py index 67c16a6a9..0b3accf36 100644 --- a/fastvideo/pipelines/basic/ltx2/presets.py +++ b/fastvideo/pipelines/basic/ltx2/presets.py @@ -141,4 +141,91 @@ LTX2_TWO_STAGE = InferencePreset( }, ) -ALL_PRESETS = (LTX2_BASE, LTX2_3_BASE, LTX2_DISTILLED, LTX2_TWO_STAGE) +LTX2_5_DEV = InferencePreset( + name="ltx2_5_dev", + version=1, + model_family="ltx2", + description="LTX-2.5 dev joint video/audio generation at 512x768", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 10, + "height": 512, + "width": 768, + "num_frames": 121, + "fps": 24, + "guidance_scale": 3.0, + "num_inference_steps": 30, + "negative_prompt": _LTX2_NEGATIVE_PROMPT, + "ltx2_cfg_scale_video": 3.0, + "ltx2_cfg_scale_audio": 7.0, + "ltx2_modality_scale_video": 3.0, + "ltx2_modality_scale_audio": 3.0, + "ltx2_rescale_scale": 0.7, + "ltx2_stg_scale_video": 1.0, + "ltx2_stg_scale_audio": 1.0, + "ltx2_stg_blocks_video": [28], + "ltx2_stg_blocks_audio": [28], + "ltx2_image_crf": 18.0, + "ltx2_use_ancestral_sampler": False, + }, +) + +LTX2_5_DISTILLED = InferencePreset( + name="ltx2_5_distilled", + version=1, + model_family="ltx2", + description="LTX-2.5 distilled joint video/audio generation at 1024x1536", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "seed": 10, + "height": 1024, + "width": 1536, + "num_frames": 121, + "fps": 24, + "guidance_scale": 1.0, + "num_inference_steps": 8, + "negative_prompt": "", + "ltx2_image_crf": 18.0, + "ltx2_use_ancestral_sampler": True, + }, +) + +LTX2_5_DISTILLED_TWO_STAGE = InferencePreset( + name="ltx2_5_distilled_two_stage", + version=1, + model_family="ltx2", + description="LTX-2.5 distilled joint video/audio generation with 2x spatial refinement", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, _REFINE_STAGE), + defaults=dict(LTX2_5_DISTILLED.defaults), + stage_defaults={ + "refine": { + "num_inference_steps": 3, + "guidance_scale": 1.0, + "image_crf": 18, + }, + }, +) + +ALL_PRESETS = ( + LTX2_BASE, + LTX2_3_BASE, + LTX2_DISTILLED, + LTX2_TWO_STAGE, + LTX2_5_DEV, + LTX2_5_DISTILLED, + LTX2_5_DISTILLED_TWO_STAGE, +) + +__all__ = [ + "ALL_PRESETS", + "LTX2_BASE", + "LTX2_3_BASE", + "LTX2_DISTILLED", + "LTX2_TWO_STAGE", + "LTX2_5_DEV", + "LTX2_5_DISTILLED", + "LTX2_5_DISTILLED_TWO_STAGE", +] diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py index 0a35aca05..c7c301490 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py @@ -25,6 +25,7 @@ from fastvideo.pipelines.basic.ltx2.stages.ltx2_image_conditioning import (LTX2_ LTX2_VIDEO_DENOISE_MASK_KEY, apply_ltx2_gaussian_noiser, post_process_ltx2_denoised) +from fastvideo.pipelines.basic.ltx2.stages.ltx2_latent_preparation import _randn_ltx2_video_latents from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.logger import init_logger @@ -43,6 +44,7 @@ MAX_SHIFT_ANCHOR = 4096 # Official distilled sigma schedule (8 denoising steps) # From LTX-2/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] +ANCESTRAL_NOISE_SEED_OFFSET = 10000 logger = init_logger(__name__) @@ -100,6 +102,24 @@ def _ltx2_sigmas( return sigmas +def _ltx2_first_frame_keyframes_mask( + *, + batch_size: int, + token_count: int, + latent_frames: int, + device: torch.device, +) -> torch.Tensor: + """Mark the target's first causal latent frame, matching LTX-2.5.""" + if latent_frames < 1 or token_count % latent_frames != 0: + raise ValueError( + "LTX-2 keyframe mask requires a positive latent frame count that divides the video token count: " + f"tokens={token_count}, frames={latent_frames}") + tokens_per_frame = token_count // latent_frames + mask = torch.zeros((batch_size, token_count, 1), device=device, dtype=torch.float32) + mask[:, :tokens_per_frame] = 1.0 + return mask + + def _distilled_subset_sigmas( steps: int, device: torch.device, @@ -151,6 +171,45 @@ def _distilled_subset_sigmas( return subset_sigmas, list(best_indices) +def _ltx2_euler_ancestral_step( + sample: torch.Tensor, + denoised_sample: torch.Tensor, + sigma: torch.Tensor, + sigma_next: torch.Tensor, + noise: torch.Tensor | None, + *, + eta: float = 1.0, + s_noise: float = 1.0, +) -> torch.Tensor: + """Apply the official rectified-flow ancestral Euler update. + + LTX-2.5 distilled generation advances to an intermediate + ``sigma_down`` and then re-noises to ``sigma_next`` while preserving the + signal/noise variance. This is intentionally not the VE/DDIM ancestral + formula used by conventional Euler schedulers. + """ + sigma = sigma.to(device=sample.device, dtype=torch.float32) + sigma_next = sigma_next.to(device=sample.device, dtype=torch.float32) + if bool(sigma_next == 0): + return denoised_sample.to(sample.dtype) + if eta > 0 and noise is None: + raise ValueError("LTX-2 ancestral Euler sampling requires noise when eta > 0.") + + sample_float = sample.float() + denoised_float = denoised_sample.float() + downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * eta + sigma_down = sigma_next * downstep_ratio + + sigma_down_ratio = sigma_down / sigma + sample_next = sigma_down_ratio * sample_float + (1.0 - sigma_down_ratio) * denoised_float + if eta > 0: + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp(min=0).sqrt() + sample_next = ((alpha_next / alpha_down) * sample_next + noise.float() * s_noise * renoise_coeff) + return sample_next.to(sample.dtype) + + class LTX2DenoisingStage(PipelineStage): """Run the LTX-2 denoising loop over the sigma schedule.""" @@ -272,8 +331,8 @@ class LTX2DenoisingStage(PipelineStage): ) logger.info("[LTX2] Using computed sigma schedule, " "num_inference_steps=%s", num_inference_steps) + video_shape = VideoLatentShape.from_torch_shape(latents.shape) if hasattr(self.transformer, "patchifier"): - video_shape = VideoLatentShape.from_torch_shape(latents.shape) token_count = self.transformer.patchifier.get_token_count(video_shape) else: token_count = 1 @@ -283,6 +342,12 @@ class LTX2DenoisingStage(PipelineStage): device=latents.device, dtype=torch.float32, ) + keyframes_mask = _ltx2_first_frame_keyframes_mask( + batch_size=latents.shape[0], + token_count=token_count, + latent_frames=video_shape.frames, + device=latents.device, + ) if video_denoise_mask is not None: patchifier = getattr(self.transformer, "patchifier", None) if patchifier is not None: @@ -457,6 +522,18 @@ class LTX2DenoisingStage(PipelineStage): tuple(sigmas.shape), tuple(latents.shape), ) + use_ancestral_sampler = bool(batch.ltx2_use_ancestral_sampler and self.sigmas_override is None) + ancestral_generator = None + if use_ancestral_sampler: + if batch.seed is None: + raise ValueError("LTX-2 ancestral sampling requires an explicit seed.") + ancestral_generator = torch.Generator( + device=latents.device).manual_seed(int(batch.seed) + ANCESTRAL_NOISE_SEED_OFFSET) + logger.info( + "[LTX2] Using official ancestral Euler stage-1 sampler " + "(eta=1, s_noise=1, seed_offset=%d)", + ANCESTRAL_NOISE_SEED_OFFSET, + ) # Hint runtime FP4 layer gating (single shared transformer path): # stage-1 denoising uses "base", stage-2 refine uses "refine". batch.extra["ltx2_fp4_stage_profile"] = ("refine" if self.sigmas_override is not None else "base") @@ -516,6 +593,7 @@ class LTX2DenoisingStage(PipelineStage): audio_timestep=audio_timestep, video_sigma=sigma_batch, audio_sigma=sigma_batch, + keyframes_mask=keyframes_mask, video_position_offset_sec=video_position_offset_sec, ) if isinstance(pos_outputs, tuple): @@ -545,6 +623,7 @@ class LTX2DenoisingStage(PipelineStage): audio_timestep=audio_timestep, video_sigma=sigma_batch, audio_sigma=sigma_batch, + keyframes_mask=keyframes_mask, video_position_offset_sec=video_position_offset_sec, ) if isinstance(neg_outputs, tuple): @@ -566,6 +645,7 @@ class LTX2DenoisingStage(PipelineStage): audio_timestep=audio_timestep, video_sigma=sigma_batch, audio_sigma=sigma_batch, + keyframes_mask=keyframes_mask, skip_cross_modal_attn=True, video_position_offset_sec=video_position_offset_sec, ) @@ -588,6 +668,7 @@ class LTX2DenoisingStage(PipelineStage): audio_timestep=audio_timestep, video_sigma=sigma_batch, audio_sigma=sigma_batch, + keyframes_mask=keyframes_mask, skip_video_self_attn_blocks=(stg_blocks_video if do_stg_video else None), skip_audio_self_attn_blocks=(stg_blocks_audio if do_stg_audio else None), video_position_offset_sec=video_position_offset_sec, @@ -641,11 +722,86 @@ class LTX2DenoisingStage(PipelineStage): ) dt = sigma_next - sigma with _nvtx_range("ltx2.denoise.scheduler_update"): - velocity = ((latents.float() - pos_denoised.float()) / sigma_value).to(latents.dtype) - latents = (latents.float() + velocity.float() * dt).to(latents.dtype) - if pos_audio is not None and audio_latents is not None: - audio_velocity = ((audio_latents.float() - pos_audio.float()) / sigma_value).to(audio_latents.dtype) - audio_latents = (audio_latents.float() + audio_velocity.float() * dt).to(audio_latents.dtype) + if use_ancestral_sampler: + assert ancestral_generator is not None + # Match the official joint-AV random stream: video draws + # first, then audio, from one seed+10000 generator. The + # official loop draws on the patchified latent state, so + # sample in token order and unpatchify back to the native + # layout (same order as the initial latent draw). + video_noise = None + audio_noise = None + if bool(sigma_next != 0): + video_noise = _randn_ltx2_video_latents( + shape=tuple(latents.shape), + transformer=self.transformer, + generator=ancestral_generator, + device=latents.device, + dtype=latents.dtype, + ) + if pos_audio is not None and audio_latents is not None: + audio_patchifier = getattr(self.transformer, "audio_patchifier", None) + if audio_patchifier is not None: + audio_noise_shape = AudioLatentShape.from_torch_shape(audio_latents.shape) + audio_patch_shape = ( + audio_noise_shape.batch, + audio_noise_shape.frames, + audio_noise_shape.channels * audio_noise_shape.mel_bins, + ) + audio_noise_patch = torch.randn( + audio_patch_shape, + generator=ancestral_generator, + device=audio_latents.device, + dtype=audio_latents.dtype, + ) + audio_noise = audio_patchifier.unpatchify(audio_noise_patch, audio_noise_shape) + else: + audio_noise = torch.randn( + audio_latents.shape, + generator=ancestral_generator, + device=audio_latents.device, + dtype=audio_latents.dtype, + ) + latents = _ltx2_euler_ancestral_step( + latents, + pos_denoised, + sigma, + sigma_next, + video_noise, + ) + if pos_audio is not None and audio_latents is not None: + audio_latents = _ltx2_euler_ancestral_step( + audio_latents, + pos_audio, + sigma, + sigma_next, + audio_noise, + ) + # Re-apply clean conditioning after stochastic noise. + # Skipped on the terminal step (sigma_next == 0): the + # denoised output is already mask-corrected above and the + # official loop does not re-apply the blend there. + if bool(sigma_next != 0): + if video_clean_latent is not None and video_denoise_mask is not None: + latents = post_process_ltx2_denoised( + denoised=latents, + denoise_mask=video_denoise_mask, + clean_latent=video_clean_latent, + ) + if (audio_clean_latent is not None and audio_denoise_mask is not None + and audio_latents is not None): + audio_latents = post_process_ltx2_denoised( + denoised=audio_latents, + denoise_mask=audio_denoise_mask, + clean_latent=audio_clean_latent, + ) + else: + velocity = ((latents.float() - pos_denoised.float()) / sigma_value).to(latents.dtype) + latents = (latents.float() + velocity.float() * dt).to(latents.dtype) + if pos_audio is not None and audio_latents is not None: + audio_velocity = ((audio_latents.float() - pos_audio.float()) / sigma_value).to( + audio_latents.dtype) + audio_latents = (audio_latents.float() + audio_velocity.float() * dt).to(audio_latents.dtype) batch.latents = latents if (batch.return_continuation_state and self.sigmas_override is not None): diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 4fa41caf9..5b36cc59f 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -222,6 +222,9 @@ class ForwardBatch: ltx2_stg_scale_audio: float = 0.0 ltx2_stg_blocks_video: list[int] = field(default_factory=list) ltx2_stg_blocks_audio: list[int] = field(default_factory=list) + # Enabled only by LTX-2.5 distilled presets. The refine stage supplies an + # explicit sigma override and therefore remains deterministic. + ltx2_use_ancestral_sampler: bool = False # LTX-2 image / video / continuation conditioning ltx2_images: list[tuple[str, int, float]] | None = None diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 096f5b4d0..365bb3ea9 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -254,10 +254,11 @@ def _get_config_info( if helios_model_id is not None: logger.debug("Resolved Helios-Distilled from authoritative model index metadata.") return _CONFIG_REGISTRY.get(helios_model_id) + variant = config.get("fastvideo_ltx2_variant", "").lower() matched_model_names: list[str] = [] for model_id, detector in _MODEL_NAME_DETECTORS: - if detector(model_path.lower()) or detector(pipeline_name): + if detector(model_path.lower()) or detector(pipeline_name) or detector(variant): logger.debug("Matched model name '%s' using a registered detector.", model_id) matched_model_names.append(model_id) @@ -331,8 +332,41 @@ def _register_configs() -> None: pipeline_cls_name="MMAudioPipeline", ) - # LTX-2 (distilled) — registered FIRST so its detector wins over - # the base detector when both fire. The detector loop in + # LTX-2.5 distilled — registered before every older/generic LTX entry so + # its ancestral-sampling preset and architecture flags cannot be shadowed. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=LTX2T2VConfig, + workload_types=(WorkloadType.T2V, WorkloadType.I2V), + hf_model_paths=[ + "FastVideo/LTX-2.5-Distilled-Diffusers", + "FastVideo/LTX2.5-Distilled-Diffusers", + ], + model_detectors=[ + lambda path: ("ltx-2.5" in path.lower() or "ltx2.5" in path.lower()) and "distilled" in path.lower(), + ], + model_family="ltx2", + default_preset="ltx2_5_distilled_two_stage", + ) + # LTX-2.5 dev — raw Lightricks split checkpoints must first be converted + # into the standard component layout consumed by FastVideo loaders. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=LTX2T2VConfig, + workload_types=(WorkloadType.T2V, WorkloadType.I2V), + hf_model_paths=[ + "FastVideo/LTX-2.5-Dev-Diffusers", + "FastVideo/LTX2.5-Dev-Diffusers", + ], + model_detectors=[ + lambda path: ("ltx-2.5" in path.lower() or "ltx2.5" in path.lower()) and "distilled" not in path.lower(), + ], + model_family="ltx2", + default_preset="ltx2_5_dev", + ) + + # LTX-2 (distilled) — registered before the generic base detector so its + # detector wins over the base detector when both fire. The detector loop in # ``get_model_name_for_path`` ORs the path-based check with a # pipeline-name check (``ltx2pipeline``) which the base detector's # "distilled not in path" predicate matches as True (the diff --git a/fastvideo/tests/api/test_presets.py b/fastvideo/tests/api/test_presets.py index 75ee31113..19a9645ea 100644 --- a/fastvideo/tests/api/test_presets.py +++ b/fastvideo/tests/api/test_presets.py @@ -378,7 +378,15 @@ class TestLtx2Presets: import fastvideo.registry # noqa: F401 presets = get_presets_for_family("ltx2") names = {p.name for p in presets} - assert names == {"ltx2_base", "ltx2_3_base", "ltx2_distilled", "ltx2_two_stage"} + assert names == { + "ltx2_base", + "ltx2_3_base", + "ltx2_distilled", + "ltx2_two_stage", + "ltx2_5_dev", + "ltx2_5_distilled", + "ltx2_5_distilled_two_stage", + } def test_ltx2_base_lookup(self) -> None: import fastvideo.registry # noqa: F401 diff --git a/mkdocs.yml b/mkdocs.yml index 2a17393ce..8cc7b4c82 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -187,6 +187,7 @@ nav: - FastH3 on RTX PRO 6000: inference/fasth3_rtx_pro_6000.md - ComfyUI: inference/comfyui.md - Support Matrix: inference/support_matrix.md + - LTX-2.5: inference/ltx2_5.md - Kandinsky 6 T2IVA: inference/kandinsky6.md - Kandinsky 6 Video SR: inference/kandinsky6_sr.md - CLI: inference/cli.md diff --git a/pyproject.toml b/pyproject.toml index 87dbba42a..e687e649e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,8 @@ dependencies = [ # transformers 5.0.0. Floor raised from >=5.0.0 to >=5.15.0 for the public # vision interpolation helper (get_vision_interpolation_indices_and_weights) # that the MiniMax-H3 Qwen3-VL encoder cross-checks against. + # LTX-2.5's packed Gemma 4 unified encoder also requires transformers 5.15.0+ + # for correct mapping and loading of the Gemma 4 architecture. "transformers>=5.15.0", # <0.23: tokenizers 0.23 renamed RobertaProcessing's binding args, so # transformers' CLIP-style tokenizer loading dies with diff --git a/pyproject_other.toml b/pyproject_other.toml index e74e7a72a..129bbf236 100644 --- a/pyproject_other.toml +++ b/pyproject_other.toml @@ -23,7 +23,9 @@ dependencies = [ # Machine Learning & Transformers # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in # transformers 5.0.0; floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc). - "transformers>=5.0.0", + # Kept in lockstep with pyproject.toml: >=5.15.0 for the MiniMax-H3 vision + # interpolation helper and LTX-2.5's packed Gemma 4 unified encoder mapping. + "transformers>=5.15.0", "tokenizers>=0.20.1", "sentencepiece>=0.2.0", "timm>=1.0.11", diff --git a/scripts/checkpoint_conversion/convert_ltx2_weights.py b/scripts/checkpoint_conversion/convert_ltx2_weights.py index c731ed866..c9a55ff64 100644 --- a/scripts/checkpoint_conversion/convert_ltx2_weights.py +++ b/scripts/checkpoint_conversion/convert_ltx2_weights.py @@ -14,6 +14,18 @@ Example usage: --pipeline-class-name "LTX2Pipeline" \\ --diffusers-version "0.33.0.dev0" \\ --gemma-path "/google/gemma-3-12b-it" + +LTX-2.5 uses separate official component files. Convert that layout with: + + python scripts/checkpoint_conversion/convert_ltx2_weights.py \ + --transformer-source diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors \ + --text-encoder-source text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors \ + --vae-source vae/ltx-2.5-video-vae-conv-bf16.safetensors \ + --audio-vae-source vae/ltx-2.5-audio-vae-bf16.safetensors \ + --spatial-upscaler-source latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors \ + --distilled-lora-source loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors \ + --variant dev \ + --output converted_weights/ltx2-5-dev """ from __future__ import annotations @@ -26,6 +38,7 @@ import re import shutil from collections import OrderedDict from pathlib import Path +from typing import Any import torch from safetensors import safe_open @@ -48,6 +61,35 @@ COMPONENT_PREFIXES: dict[str, tuple[str, ...]] = { "text_embedding_projection": ("text_embedding_projection.", "model.text_embedding_projection."), } +PACKED_GEMMA_CONFIG_METADATA_KEY = "gemma_config" +PACKED_GEMMA_TOKENIZER_KEY = "tokenizer_json" +PACKED_GEMMA_ASSET_PREFIX = "hf_asset__" + +SPLIT_SOURCE_ARGUMENTS = ( + "transformer_source", + "text_encoder_source", + "vae_source", + "audio_vae_source", + "spatial_upscaler_source", +) + +TEXT_PROJECTION_PREFIXES: dict[str, str] = { + "text_embedding_projection.aggregate_embed.": "feature_extractor_linear.aggregate_embed.", + "text_embedding_projection.video_aggregate_embed.": "video_feature_extractor_linear.", + "text_embedding_projection.audio_aggregate_embed.": "audio_feature_extractor_linear.", + "model.text_embedding_projection.aggregate_embed.": "feature_extractor_linear.aggregate_embed.", + "model.text_embedding_projection.video_aggregate_embed.": "video_feature_extractor_linear.", + "model.text_embedding_projection.audio_aggregate_embed.": "audio_feature_extractor_linear.", +} + +TEXT_CONNECTOR_PREFIXES: dict[str, str] = { + "model.diffusion_model.video_embeddings_connector.": "embeddings_connector.", + "model.diffusion_model.audio_embeddings_connector.": "audio_embeddings_connector.", + "model.diffusion_model.embeddings_connector.": "embeddings_connector.", + "video_embeddings_connector.": "embeddings_connector.", + "audio_embeddings_connector.": "audio_embeddings_connector.", +} + def _find_shards(model_path: Path) -> list[Path]: if model_path.is_file(): @@ -83,6 +125,18 @@ def _read_metadata_config(path: Path) -> dict: return json.loads(metadata["config"]) +def _read_safetensors_metadata(path: Path) -> dict[str, str]: + with safe_open(str(path), framework="pt") as f: + return dict(f.metadata() or {}) + + +def _metadata_config_for_component(metadata_config: dict, component_name: str) -> dict: + component_config = metadata_config.get(component_name) + if isinstance(component_config, dict): + return component_config + return metadata_config + + def _filter_transformer_config(config: dict) -> dict: transformer = config.get("transformer", {}) allowed = { @@ -108,6 +162,28 @@ def _filter_transformer_config(config: dict) -> dict: "audio_cross_attention_dim", "audio_positional_embedding_max_pos", "av_ca_timestep_scale_multiplier", + # LTX-2.3 architecture extensions. + "cross_attention_adaln", + "caption_proj_before_connector", + "apply_gated_attention", + "caption_projection_first_linear", + "caption_proj_input_norm", + "caption_projection_second_linear", + "connector_num_attention_heads", + "connector_attention_head_dim", + "connector_num_layers", + "audio_connector_num_attention_heads", + "audio_connector_attention_head_dim", + "audio_connector_num_layers", + "connector_positional_embedding_max_pos", + "connector_apply_gated_attention", + "connector_ff_bias", + # LTX-2.5 architecture extensions. Defaults live in the native model + # config, so preserving an explicit False is important. + "use_prompt_adaln_single", + "ff_bias", + "audio_ff_bias", + "use_keyframes_abs_pos_embedding", } filtered = {k: v for k, v in transformer.items() if k in allowed} if "frequencies_precision" in filtered: @@ -141,6 +217,170 @@ def _build_text_embedding_projection_config(gemma_model_path: str = "", ) -> dic } +def _gemma_text_config(gemma_config: dict) -> dict: + text_config = gemma_config.get("text_config") + if isinstance(text_config, dict): + return text_config + return gemma_config + + +def _build_split_text_encoder_config(transformer_config: dict, gemma_config: dict) -> dict: + text_config = _gemma_text_config(gemma_config) + hidden_size = int(text_config.get("hidden_size", 3840)) + num_hidden_layers = int(text_config.get("num_hidden_layers", 48)) + video_inner_dim = int(transformer_config.get("num_attention_heads", 30)) * int( + transformer_config.get("attention_head_dim", 128)) + audio_inner_dim = int(transformer_config.get("audio_num_attention_heads", 30)) * int( + transformer_config.get("audio_attention_head_dim", 128)) + + eos_token_id = gemma_config.get("eos_token_id", text_config.get("eos_token_id", 2)) + if isinstance(eos_token_id, list): + eos_token_id = eos_token_id[0] if eos_token_id else 2 + config = _build_text_embedding_projection_config(gemma_model_path="gemma") + config.update({ + "hidden_size": hidden_size, + "num_hidden_layers": num_hidden_layers, + "num_attention_heads": int(text_config.get("num_attention_heads", 30)), + # The packed LTX tokenizer always pads/truncates to 1024 regardless of + # the Gemma backbone's much larger context window. + "text_len": 1024, + "pad_token_id": int(gemma_config.get("pad_token_id", text_config.get("pad_token_id", 0)) or 0), + "bos_token_id": int(gemma_config.get("bos_token_id", text_config.get("bos_token_id", 2)) or 2), + "eos_token_id": int(eos_token_id), + "feature_extractor_in_features": hidden_size * (num_hidden_layers + 1), + "feature_extractor_out_features": video_inner_dim, + "video_feature_extractor_out_features": video_inner_dim, + "audio_feature_extractor_out_features": audio_inner_dim, + "caption_proj_before_connector": bool(transformer_config.get("caption_proj_before_connector", False)), + "caption_projection_first_linear": bool(transformer_config.get("caption_projection_first_linear", True)), + "caption_proj_input_norm": bool(transformer_config.get("caption_proj_input_norm", True)), + "caption_projection_second_linear": bool(transformer_config.get("caption_projection_second_linear", True)), + "connector_num_attention_heads": int(transformer_config.get("connector_num_attention_heads", 30)), + "connector_attention_head_dim": int(transformer_config.get("connector_attention_head_dim", 128)), + "connector_num_layers": int(transformer_config.get("connector_num_layers", 2)), + "audio_connector_num_attention_heads": int( + transformer_config.get( + "audio_connector_num_attention_heads", + transformer_config.get("connector_num_attention_heads", 30), + )), + "audio_connector_attention_head_dim": int( + transformer_config.get( + "audio_connector_attention_head_dim", + transformer_config.get("connector_attention_head_dim", 128), + )), + "audio_connector_num_layers": int( + transformer_config.get( + "audio_connector_num_layers", + transformer_config.get("connector_num_layers", 2), + )), + "connector_positional_embedding_max_pos": transformer_config.get( + "connector_positional_embedding_max_pos", [4096]), + "connector_rope_type": transformer_config.get("rope_type", "split"), + "connector_apply_gated_attention": bool(transformer_config.get("connector_apply_gated_attention", False)), + "connector_ff_bias": bool(transformer_config.get("connector_ff_bias", True)), + "connector_double_precision_rope": transformer_config.get("frequencies_precision") == "float64", + }) + return config + + +def _tensor_to_bytes(tensor: torch.Tensor, key: str) -> bytes: + if tensor.dtype != torch.uint8 or tensor.ndim != 1: + raise ValueError(f"Packed Gemma asset {key!r} must be a 1-D uint8 tensor, got {tensor.dtype} {tensor.shape}.") + return bytes(tensor.cpu().tolist()) + + +def _map_packed_gemma_weight_key(key: str) -> str: + """Map the official Comfy-flat Gemma 4 pack back to HF key names. + + The official packer flattens the language tower and renames vision/audio + components. Already-HF keys pass through unchanged so repacked checkpoints + are also accepted. + """ + if key.startswith("model.language_model.") or key.startswith("model.vision_embedder."): + return key + if key.startswith(("model.layers.", "model.embed_tokens.", "model.norm.")): + return "model.language_model." + key.removeprefix("model.") + if key.startswith("vision_model."): + return "model.vision_embedder." + key.removeprefix("vision_model.") + if key.startswith("multi_modal_projector."): + return "model.embed_vision." + key.removeprefix("multi_modal_projector.") + if key.startswith("audio_projector."): + return "model.embed_audio." + key.removeprefix("audio_projector.") + return key + + +def _route_text_projection_key(key: str) -> str | None: + for prefix, replacement in TEXT_PROJECTION_PREFIXES.items(): + if key.startswith(prefix): + return replacement + key[len(prefix):] + for prefix, replacement in TEXT_CONNECTOR_PREFIXES.items(): + if key.startswith(prefix): + return replacement + key[len(prefix):] + return None + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(value, f, indent=2) + f.write("\n") + + +def _write_bytes(path: Path, value: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(value) + + +def _unpack_packed_gemma( + source_path: Path, + output_dir: Path, +) -> tuple[OrderedDict, dict]: + shards = _find_shards(source_path) + if len(shards) != 1: + raise ValueError("Packed LTX-2.5 text encoder must resolve to exactly one safetensors file.") + source_file = shards[0] + metadata = _read_safetensors_metadata(source_file) + raw_gemma_config = metadata.get(PACKED_GEMMA_CONFIG_METADATA_KEY) + if raw_gemma_config is None: + raise ValueError( + f"Packed text encoder {source_file} is missing {PACKED_GEMMA_CONFIG_METADATA_KEY!r} metadata.") + gemma_config = json.loads(raw_gemma_config) + + gemma_weights: OrderedDict[str, torch.Tensor] = OrderedDict() + projection_weights: OrderedDict[str, torch.Tensor] = OrderedDict() + tokenizer_assets: dict[str, bytes] = {} + for key, tensor in _load_weights(shards).items(): + projection_key = _route_text_projection_key(key) + if projection_key is not None: + projection_weights[projection_key] = tensor + elif key == PACKED_GEMMA_TOKENIZER_KEY: + tokenizer_assets["tokenizer.json"] = _tensor_to_bytes(tensor, key) + elif key.startswith(PACKED_GEMMA_ASSET_PREFIX): + asset_name = key.removeprefix(PACKED_GEMMA_ASSET_PREFIX) + # Untrusted safetensors keys must not escape the output tree: + # reject separators and dot-segments so the joined path stays a + # single file name inside the target tokenizer directories. + if not asset_name or asset_name in {".", ".."} or "/" in asset_name or "\\" in asset_name: + raise ValueError(f"Unsafe packed Gemma asset name {asset_name!r} from tensor key {key!r}.") + tokenizer_assets[asset_name] = _tensor_to_bytes(tensor, key) + else: + gemma_weights[_map_packed_gemma_weight_key(key)] = tensor + + if not gemma_weights: + raise ValueError(f"Packed text encoder {source_file} did not contain Gemma model weights.") + if "tokenizer.json" not in tokenizer_assets: + raise ValueError(f"Packed text encoder {source_file} did not contain tokenizer_json.") + + gemma_dir = output_dir / "text_encoder" / "gemma" + gemma_dir.mkdir(parents=True, exist_ok=True) + save_file(gemma_weights, str(gemma_dir / "model.safetensors")) + _write_json(gemma_dir / "config.json", gemma_config) + for asset_name, asset_bytes in tokenizer_assets.items(): + _write_bytes(gemma_dir / asset_name, asset_bytes) + _write_bytes(output_dir / "tokenizer" / asset_name, asset_bytes) + return projection_weights, gemma_config + + def _wrap_component_config( component_name: str, component_config: dict | None, @@ -179,6 +419,59 @@ def _split_component_weights(weights: dict[str, torch.Tensor]) -> dict[str, Orde return {name: weights for name, weights in components.items() if weights} +def _strip_first_matching_prefix(key: str, prefixes: tuple[str, ...]) -> str: + for prefix in prefixes: + if key.startswith(prefix): + return key[len(prefix):] + return key + + +def _split_transformer_and_connectors( + weights: dict[str, torch.Tensor], +) -> tuple[OrderedDict, OrderedDict]: + transformer: OrderedDict[str, torch.Tensor] = OrderedDict() + text_encoder: OrderedDict[str, torch.Tensor] = OrderedDict() + for key, tensor in weights.items(): + text_key = _route_text_projection_key(key) + if text_key is not None: + text_encoder[text_key] = tensor + continue + transformer_key = _strip_first_matching_prefix( + key, + ("model.diffusion_model.", "diffusion_model."), + ) + transformer[transformer_key] = tensor + return transformer, text_encoder + + +def _split_audio_vae_source(weights: dict[str, torch.Tensor]) -> tuple[OrderedDict, OrderedDict]: + audio_vae: OrderedDict[str, torch.Tensor] = OrderedDict() + vocoder: OrderedDict[str, torch.Tensor] = OrderedDict() + for key, tensor in weights.items(): + if key.startswith("audio_vae."): + audio_vae[key.removeprefix("audio_vae.")] = tensor + elif key.startswith("vocoder."): + # Strip exactly one prefix. LTX-2.3+ BWE keys intentionally keep + # the second "vocoder." segment. + vocoder[key.removeprefix("vocoder.")] = tensor + if not audio_vae: + raise ValueError("Audio VAE source contains no keys with the 'audio_vae.' prefix.") + if not vocoder: + raise ValueError("Audio VAE source contains no keys with the 'vocoder.' prefix.") + return audio_vae, vocoder + + +def _component_weights( + source_path: Path, + prefixes: tuple[str, ...], +) -> OrderedDict: + shards = _find_shards(source_path) + if not shards: + raise FileNotFoundError(f"No safetensors found in {source_path}") + return OrderedDict( + (_strip_first_matching_prefix(key, prefixes), tensor) for key, tensor in _load_weights(shards).items()) + + def _write_component( output_dir: Path, name: str, @@ -218,6 +511,33 @@ def _build_model_index( } +def _build_split_model_index( + transformer_class_name: str, + pipeline_class_name: str, + diffusers_version: str, + variant: str, + distilled_lora: bool, +) -> dict: + model_index = _build_model_index( + transformer_class_name=transformer_class_name, + vae_class_name="CausalVideoAutoencoder", + pipeline_class_name=pipeline_class_name, + diffusers_version=diffusers_version, + ) + model_index.update({ + "spatial_upsampler": ["diffusers", "LTX2LatentUpsampler"], + "fastvideo_ltx2_variant": f"ltx2.5-{variant}", + "fastvideo_refine_enabled": variant == "distilled" or distilled_lora, + "fastvideo_refine_upsampler_path": "spatial_upsampler", + "fastvideo_refine_num_inference_steps": 3, + "fastvideo_refine_guidance_scale": 1.0, + "fastvideo_refine_add_noise": True, + }) + if distilled_lora: + model_index["fastvideo_refine_lora_path"] = "distilled_lora/model.safetensors" + return model_index + + def _write_model_index(output_dir: Path, model_index: dict) -> None: output_dir.mkdir(parents=True, exist_ok=True) model_index_path = output_dir / "model_index.json" @@ -313,6 +633,119 @@ def convert_components( _write_model_index(output_dir, model_index) +def _source_metadata_config(source_path: Path) -> dict: + shards = _find_shards(source_path) + if not shards: + raise FileNotFoundError(f"No safetensors found in {source_path}") + return _read_metadata_config(shards[0]) + + +def convert_split_components( + *, + transformer_source: Path, + text_encoder_source: Path, + vae_source: Path, + audio_vae_source: Path, + spatial_upscaler_source: Path, + output_dir: Path, + transformer_class_name: str, + variant: str, + distilled_lora_source: Path | None = None, + emit_diffusers_repo: bool = True, + pipeline_class_name: str = "LTX2Pipeline", + diffusers_version: str = "0.33.0.dev0", +) -> None: + """Convert the official LTX-2.5 separate-component checkpoint layout. + + The public LTX-2.5 repository describes these files as independent + safetensors. Exact real-weight strict loading remains a post-conversion + verification step because the repository is gated; this function keeps all + routing explicit so key drift fails visibly rather than being dropped. + """ + transformer_metadata = _source_metadata_config(transformer_source) + full_transformer_config = _metadata_config_for_component(transformer_metadata, "transformer") + transformer_config = _filter_transformer_config(transformer_metadata) + if not transformer_config: + raise ValueError("Transformer source is missing config.transformer metadata.") + transformer_config["_class_name"] = transformer_class_name + + transformer_shards = _find_shards(transformer_source) + transformer_weights, connector_weights = _split_transformer_and_connectors(_load_weights(transformer_shards)) + if not transformer_weights: + raise ValueError("Transformer source did not contain transformer weights.") + + output_dir.mkdir(parents=True, exist_ok=True) + projection_weights, gemma_config = _unpack_packed_gemma(text_encoder_source, output_dir) + projection_weights.update(connector_weights) + if not projection_weights: + raise ValueError("Split sources did not contain LTX text projections or connectors.") + text_encoder_config = _build_split_text_encoder_config(full_transformer_config, gemma_config) + + vae_metadata = _source_metadata_config(vae_source) + vae_config = _metadata_config_for_component(vae_metadata, "vae") + vae_class_name = vae_config.get("_class_name", "CausalVideoAutoencoder") + if vae_class_name != "CausalVideoAutoencoder": + raise ValueError( + "--vae-source must be the LTX-2.5 convolutional VAE; " + f"metadata declares {vae_class_name!r}.") + vae_weights = _component_weights(vae_source, ("vae.", "model.vae.")) + + audio_metadata = _source_metadata_config(audio_vae_source) + audio_weights, vocoder_weights = _split_audio_vae_source( + _load_weights(_find_shards(audio_vae_source))) + audio_vae_config = _metadata_config_for_component(audio_metadata, "audio_vae") + vocoder_config = _metadata_config_for_component(audio_metadata, "vocoder") + if audio_vae_config is audio_metadata or vocoder_config is audio_metadata: + raise ValueError("Audio VAE source metadata must contain both config.audio_vae and config.vocoder.") + + upscaler_metadata = _source_metadata_config(spatial_upscaler_source) + upscaler_config = dict(upscaler_metadata) + upscaler_config["_class_name"] = "LTX2LatentUpsampler" + upscaler_weights = _component_weights( + spatial_upscaler_source, + ("spatial_upscaler.", "spatial_upsampler.", "upsampler.", "model."), + ) + upscaler_weights = OrderedDict((f"upsampler.{key}" if key.split(".", 1)[0].isdigit() else key, value) + for key, value in upscaler_weights.items()) + + _write_component(output_dir, "transformer", transformer_weights, transformer_config) + _write_component(output_dir, "text_encoder", projection_weights, text_encoder_config) + _write_component( + output_dir, + "vae", + vae_weights, + _wrap_component_config("vae", vae_config, class_name="CausalVideoAutoencoder"), + ) + _write_component( + output_dir, + "audio_vae", + audio_weights, + _wrap_component_config("audio_vae", audio_vae_config, class_name="LTX2AudioDecoder"), + ) + _write_component( + output_dir, + "vocoder", + vocoder_weights, + _wrap_component_config("vocoder", vocoder_config, class_name="LTX2Vocoder"), + ) + _write_component(output_dir, "spatial_upsampler", upscaler_weights, upscaler_config) + + has_distilled_lora = distilled_lora_source is not None + if distilled_lora_source is not None: + lora_weights = _component_weights(distilled_lora_source, ()) + _write_component(output_dir, "distilled_lora", lora_weights, config=None) + + if emit_diffusers_repo: + model_index = _build_split_model_index( + transformer_class_name=transformer_class_name, + pipeline_class_name=pipeline_class_name, + diffusers_version=diffusers_version, + variant=variant, + distilled_lora=has_distilled_lora, + ) + _write_model_index(output_dir, model_index) + + def update_transformer_config(config_path: Path, class_name: str) -> None: if not config_path.exists(): print(f"Config file not found: {config_path}") @@ -411,9 +844,84 @@ def main() -> None: default="", help="Optional local Gemma model path to copy into the output repo.", ) + parser.add_argument( + "--transformer-source", + type=str, + help="LTX-2.5 split transformer safetensors file or sharded directory.", + ) + parser.add_argument( + "--text-encoder-source", + type=str, + help="LTX-2.5 packed Gemma 4 text-encoder safetensors file.", + ) + parser.add_argument( + "--vae-source", + type=str, + help="LTX-2.5 convolutional video VAE safetensors file or sharded directory.", + ) + parser.add_argument( + "--audio-vae-source", + type=str, + help="LTX-2.5 audio VAE safetensors containing audio_vae.* and vocoder.* keys.", + ) + parser.add_argument( + "--spatial-upscaler-source", + type=str, + help="LTX-2.5 spatial latent upscaler safetensors file or sharded directory.", + ) + parser.add_argument( + "--distilled-lora-source", + type=str, + help="Optional LTX-2.5 distilled LoRA safetensors used for dev stage-2 refinement.", + ) + parser.add_argument( + "--variant", + choices=("dev", "distilled"), + default="dev", + help="LTX-2.5 split transformer variant; controls bundled refine defaults.", + ) args = parser.parse_args() + split_values = {name: getattr(args, name) for name in SPLIT_SOURCE_ARGUMENTS} + split_mode = any(value is not None for value in split_values.values()) + if args.distilled_lora_source is not None and not split_mode: + raise ValueError("--distilled-lora-source is only valid with the LTX-2.5 split source arguments.") + if split_mode: + missing = [f"--{name.replace('_', '-')}" for name, value in split_values.items() if value is None] + if missing: + raise ValueError("LTX-2.5 split conversion requires: " + ", ".join(missing)) + incompatible = [] + if args.source: + incompatible.append("--source") + if args.download: + incompatible.append("--download") + if args.gemma_path: + incompatible.append("--gemma-path") + if args.transformer_only: + incompatible.append("--transformer-only") + if args.components: + incompatible.append("--components") + if args.update_config: + incompatible.append("--update-config") + if incompatible: + raise ValueError("LTX-2.5 split source arguments cannot be combined with " + ", ".join(incompatible)) + convert_split_components( + transformer_source=Path(args.transformer_source), + text_encoder_source=Path(args.text_encoder_source), + vae_source=Path(args.vae_source), + audio_vae_source=Path(args.audio_vae_source), + spatial_upscaler_source=Path(args.spatial_upscaler_source), + distilled_lora_source=(Path(args.distilled_lora_source) if args.distilled_lora_source else None), + output_dir=Path(args.output), + transformer_class_name=args.class_name, + variant=args.variant, + emit_diffusers_repo=args.diffusers_repo, + pipeline_class_name=args.pipeline_class_name, + diffusers_version=args.diffusers_version, + ) + return + if args.download: if args.source: raise ValueError("Use either --download or --source, not both.") diff --git a/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py b/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py index 0363c2a70..f1cc4289e 100644 --- a/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py +++ b/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py @@ -306,9 +306,17 @@ def test_ltx2_typed_surface_preflight() -> None: LTX2AudioDecodingStage, LTX2DenoisingStage, LTX2LatentPreparationStage, LTX2TextEncodingStage, ) - # All four LTX-2 presets registered. + # All LTX-2/2.3/2.5 presets registered. names = {p.name for p in get_presets_for_family("ltx2")} - assert names == {"ltx2_base", "ltx2_3_base", "ltx2_distilled", "ltx2_two_stage"} + assert names == { + "ltx2_base", + "ltx2_3_base", + "ltx2_distilled", + "ltx2_two_stage", + "ltx2_5_dev", + "ltx2_5_distilled", + "ltx2_5_distilled_two_stage", + } # Two-stage preset has the denoise + refine topology and pulls its # refine allowed_overrides from the typed dataclass. diff --git a/tests/local_tests/ltx2_5/PORT_STATUS.md b/tests/local_tests/ltx2_5/PORT_STATUS.md new file mode 100644 index 000000000..db4cc1c94 --- /dev/null +++ b/tests/local_tests/ltx2_5/PORT_STATUS.md @@ -0,0 +1,99 @@ +# LTX-2.5 Port Status + +## Summary + +- model_family: `ltx2_5` +- workload_types: `T2AV, I2AV` through the existing T2V/I2V compatibility surface +- official_ref: `https://github.com/Lightricks/LTX-2` +- official_ref_dir: `LTX-2-Reference/` +- hf_weights_path: `Lightricks/LTX-2.5` +- local_weights_dir: `Modal-backed Hugging Face cache` +- source_layout: `separate_components` +- local_tests_readme: `tests/local_tests/ltx2_5/README.md` + +## Current Phase + +- phase: `launch_validation` +- status: `complete` +- owner: `orchestrator` +- last_updated: `2026-08-12` + +## Component Matrix + +| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | +|---|---|---|---|---|---|---|---|---|---| +| LTX-2.5 transformer | dit | extend | `packages/ltx-core/src/ltx_core/model/transformer/` | `LTXModelConfigurator.from_metadata` | `fastvideo/models/dits/ltx2.py`, `fastvideo/configs/models/dits/ltx2.py` | implemented | real_split_pass | real_strict_pass | — | +| Gemma 4 unified encoder/connector | encoder | port | `packages/ltx-core/src/ltx_core/text_encoders/gemma/` | Gemma encoder configurator + prompt encoder block | `fastvideo/models/encoders/gemma.py`, `fastvideo/configs/models/encoders/gemma.py` | implemented | real_split_pass | real_pipeline_pass | — | +| Convolutional video VAE | vae | extend | `packages/ltx-core/src/ltx_core/model/video_vae/` | `VideoEncoderConfigurator` / `VideoDecoderConfigurator` from checkpoint metadata | `fastvideo/models/vaes/ltx2vae.py`, `fastvideo/configs/models/vaes/ltx2vae.py` | implemented | real_split_pass | real_pipeline_pass | — | +| Audio VAE/vocoder | vae | reuse | `packages/ltx-core/src/ltx_core/model/audio_vae/` | audio VAE/vocoder configurators from checkpoint metadata | existing LTX-2 audio decoder/vocoder | reused | real_split_pass | real_pipeline_pass | — | +| Euler ancestral sampling | generic | port | official diffusion step and distilled pipeline selection | distilled preset | `fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py` | implemented | stateless | unit_pass | — | +| T2AV/I2AV pipeline | pipeline | extend | official one-stage and distilled pipelines | split component paths + official defaults | `fastvideo/pipelines/basic/ltx2/` | implemented | real_split_pass | real_t2av_pass | — | + +## Conversion State + +- conversion_script: `scripts/checkpoint_conversion/convert_ltx2_weights.py` +- converted_weights_dir: `converted_weights/ltx2_5` +- source_layout: `separate_components` +- strict_load_status: `dev_and_distilled_pass` +- passthrough_components: `packed Gemma tokenizer/assets are unpacked into text_encoder/gemma and tokenizer` +- retry_history: `I002 resolved before conversion: emitted config allowlist now preserves the 2.3/2.5 gates` + +## Parity Commands + +| Scope | Command | Last Result | Notes | +|---|---|---|---| +| lightweight launch suite | `pytest -q tests/local_tests/ltx2_5 tests/local_tests/transformers/test_ltx2_5_transformer_parity.py -k 'not real_weights'` | 14 passed, 7 skipped, 2 deselected | Modal L40S before official source clone; skips were reference-dependent | +| transformer source parity | `pytest -q -rs tests/local_tests/transformers/test_ltx2_5_transformer_parity.py -k 'not real_weights'` | 8 passed, 2 deselected | Same live Modal L40S after cloning official v1.2.0 source | +| expanded launch/legacy suite | `pytest -q tests/local_tests/ltx2_5 tests/local_tests/transformers/test_ltx2_5_transformer_parity.py fastvideo/tests/api/test_presets.py::TestLtx2Presets::test_ltx2_presets_registered tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py::test_ltx2_typed_surface_preflight -k 'not real_weights'` | 24 passed, 2 deselected | Same live Modal L40S; includes first-frame keyframe-mask wiring and legacy preset regression checks | +| real dev strict load | `pytest -q -rs 'tests/local_tests/transformers/test_ltx2_5_transformer_parity.py::test_real_weights_load_strictly_through_production_loader[dev]'` | 1 passed | Modal L40S; official builder and production `TransformerLoader`; 28.76s | +| real distilled strict load | `pytest -q -rs 'tests/local_tests/transformers/test_ltx2_5_transformer_parity.py::test_real_weights_load_strictly_through_production_loader[distilled]'` | 1 passed | Same live Modal L40S; official builder and production `TransformerLoader`; 27.06s | +| split conversion | `pytest -q tests/local_tests/ltx2_5/test_ltx2_5_conversion.py` | pass in launch suite | Synthetic files validate component routing, packed assets, and metadata | +| Gemma/VAE/sampler/registry | `pytest -q tests/local_tests/ltx2_5 -k 'not conversion'` | pass in launch suite | Focused structural and stateless behavior coverage | +| real distilled one-stage T2AV | `basic_ltx2_5_t2av.py --variant distilled ... --steps 1 --num-gpus 2` | pass | 2xH100; joint H.264/AAC output, 9 video frames and 17 audio frames; generation 11.77s | +| real distilled two-stage T2AV | `basic_ltx2_5_t2av.py --variant distilled ... --steps 1 --num-gpus 2` | pass | 2xB200; base ancestral denoise, 2x spatial upsample, 1,632-layer LoRA, three-step refine, joint H.264/AAC output; generation 76.17s | + +## Open Questions + +| ID | Question | Owner | Needed By Phase | Status | Resolution | +|---|---|---|---|---|---| +| Q001 | Which exact 2.5 transformer metadata fields and state-dict deltas are required beyond 2.3? | component:transformer | prototype | resolved | `use_prompt_adaln_single`, `ff_bias`, `audio_ff_bias`, and `use_keyframes_abs_pos_embedding`; defaults are True, True, True, False. | +| Q002 | Does the published convolutional video VAE instantiate the existing native VAE without architecture changes? | component:video_vae | prototype | resolved | Existing graph is reused with checkpoint-derived `decoder_base_channels=128`. | +| Q003 | Does the published audio VAE/BWE stack instantiate the existing native audio path unchanged? | component:audio_vae | prototype | resolved | Official v1.2.0 retains the LTX-2.3 audio VAE/vocoder/BWE topology; real joint video-audio pipeline generation passes. | +| Q004 | Can ancestral sampling reuse an existing FastVideo scheduler while preserving official per-step RNG? | component:scheduler | parity | resolved | Implemented the official rectified-flow Euler ancestral step in the LTX denoising stage, including seed offset and video-before-audio RNG order. | +| Q005 | Which official dev/distilled T2AV and I2AV defaults should become public FastVideo presets? | pipeline | pipeline | resolved | Dev uses 30-step CFG/STG guidance at 512x768; distilled uses eight-step ancestral denoising and optional three-step 2x refine. | + +## Issues And Blockers + +| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | +|---|---|---|---|---|---|---|---|---| +| I001 | parity | all | high | Real LTX-2.5 artifact loading and end-to-end generation were blocked by gated Hugging Face access. | Pinned revision `ce298b1259d61ce6c87e05154b9ad339b16f32a0` now downloads successfully in the HAO AI Lab workspace. | repository owners | resolved | Repository access was approved; both BF16 transformer variants now pass official and FastVideo strict loading. | +| I002 | conversion | transformer | high | Transformer config filtering dropped the 2.3 and 2.5 architecture gates, so converted metadata could instantiate the legacy parameter surface. | `scripts/checkpoint_conversion/convert_ltx2_weights.py::_filter_transformer_config` allowlist review. | component:conversion | resolved | Allowlist now preserves `cross_attention_adaln`, `caption_proj_before_connector`, `apply_gated_attention`, `use_prompt_adaln_single`, `ff_bias`, `audio_ff_bias`, and `use_keyframes_abs_pos_embedding`. | + +## Escape Hatches + +| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | +|---|---|---|---|---|---|---| + +## Decisions + +| Date | Decision | Rationale | Impact | +|---|---|---|---| +| 2026-08-12 | Build a complete launch inference port even if the PR is H3-sized. | Users should receive meaningful FastVideo acceleration immediately after model launch. | Dev/distilled T2AV/I2AV, native components, examples, quality coverage, and multi-GPU verification belong in the first PR. | +| 2026-08-12 | Use the official convolutional video VAE in the launch PR. | It provides authentic 2.5 inference without coupling launch support to the separate iterative DiffVAE subsystem. | DiffVAE/NATTEN/CuTe work follows independently. | +| 2026-08-12 | Keep work local until a coherent tested checkpoint, then open a non-draft PR on the fork. | CodeRabbit should review real progress while incomplete experimental fragments stay private. | Every later push is preceded by review-feedback triage. | +| 2026-08-12 | Extend the native LTX-2 transformer behind metadata-compatible gates. | Official 2.5 preserves the LTX-2.3 graph and changes only prompt AdaLN, FFN biases, and keyframe token embedding surfaces. | Existing LTX-2/2.3 configs retain byte-compatible parameter surfaces by default. | + +## Handoff Notes + +- Branch: `aryan/ltx-2-5-initial-support`, based on upstream `8208536cd`. +- Official source: `v1.2.0` at `d151147788a9284cca791edc6ce898007e727fe6`. +- Active Modal profile: `ai-lab`; workspace verified as `hao-ai-lab`. +- Modal policy: reuse one warm GPU-backed shell for iterative checks; do not + create a fresh app per test command. +- Transformer prototype targets the pinned official `ltx-core` v1.2.0 graph and + preserves FastVideo sequence-parallel and compile block boundaries. +- Official source/random transformer parity and both real BF16 transformer + strict loads pass in Modal. Real one-stage generation passes on 2xH100, and + the full distilled upsample/refine path passes on 2xB200 with verified H.264 + video and AAC audio streams. The tracked launch run is + `https://wandb.ai/aryan5v-san-jose-state-university/fastvideo-ltx2-5/runs/o9ixwewk`. diff --git a/tests/local_tests/ltx2_5/README.md b/tests/local_tests/ltx2_5/README.md new file mode 100644 index 000000000..67a007e6f --- /dev/null +++ b/tests/local_tests/ltx2_5/README.md @@ -0,0 +1,151 @@ +# LTX-2.5 Local Tests + +Local parity, conversion, and smoke tests for the native LTX-2.5 FastVideo +port. The launch scope is BF16 dev and distilled text/image-conditioned joint +video-audio inference. Training, QAD/NVFP4, DiffVAE, DFR, HDR, temporal +upsampling, and automatic duration prediction are follow-up work. + +Port progress and unresolved questions live in +`tests/local_tests/ltx2_5/PORT_STATUS.md`. + +## Reference Assets + +| Field | Value | +|---|---| +| Model family | `ltx2_5` (config-gated extension of native `ltx2`) | +| Workload types | T2AV and I2AV through the existing T2V/I2V compatibility surface; output is joint video and audio | +| Official reference | `https://github.com/Lightricks/LTX-2` | +| Local reference dir | `LTX-2-Reference/` (gitignored) | +| Official commit/version | `v1.2.0`, `d151147788a9284cca791edc6ce898007e727fe6` | +| HF weights | `Lightricks/LTX-2.5` | +| HF revision | `ce298b1259d61ce6c87e05154b9ad339b16f32a0` | +| Local weights dir | Modal-backed Hugging Face cache; no local weight download | +| Source layout | Separate component safetensors with checkpoint metadata; no root `model_index.json` | +| Needs conversion | Yes | + +Never place token values in this file. GPU validation resolves `HF_TOKEN` from +an approved secret in the HAO AI Lab Modal workspace. + +## Shared Environment Setup + +The source clone is reproducible from the repository root: + +```bash +python3 .agents/skills/add-model-01-prep/scripts/clone_reference_repo.py \ + https://github.com/Lightricks/LTX-2.git \ + LTX-2-Reference \ + --branch v1.2.0 \ + --commit d151147788a9284cca791edc6ce898007e727fe6 \ + --update-gitignore +``` + +The local FastVideo `.venv` does not currently contain PyTorch. Do not change +FastVideo's core pins on this Mac solely for parity. Numeric checks reuse one +warm HAO AI Lab Modal L40S shell with FastVideo and the pinned official source +installed together. + +```text +dependency_changes: transformers>=5.15.0 for Gemma 4 +official_env_status: source parity passing in one shared Modal environment +private_dep_stubs: none planned +gated_access_status: approved and verified at the pinned revision +``` + +## Weight Layout + +The published gated repository exposes these inference components: + +```text +diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors +diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors +text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors +vae/ltx-2.5-video-vae-conv-bf16.safetensors +vae/ltx-2.5-audio-vae-bf16.safetensors +``` + +The first PR deliberately uses the official convolutional video VAE. The +diffusion VAE is a separate iterative neighborhood-attention decoder and is not +required for authentic LTX-2.5 launch inference. + +## Prototype And Conversion Artifacts + +```text +official_key_dumps: + transformer: converted_weights/ltx2_5/_mapping/transformer_official_keys.json + text_encoder: converted_weights/ltx2_5/_mapping/text_encoder_official_keys.json + video_vae: converted_weights/ltx2_5/_mapping/video_vae_official_keys.json + audio_vae: converted_weights/ltx2_5/_mapping/audio_vae_official_keys.json +fastvideo_key_dumps: + transformer: converted_weights/ltx2_5/_mapping/transformer_fastvideo_keys.json + text_encoder: converted_weights/ltx2_5/_mapping/text_encoder_fastvideo_keys.json + video_vae: converted_weights/ltx2_5/_mapping/video_vae_fastvideo_keys.json + audio_vae: converted_weights/ltx2_5/_mapping/audio_vae_fastvideo_keys.json +conversion_script: scripts/checkpoint_conversion/convert_ltx2_weights.py +conversion_source_layout: separate_components +converted_weights_dir: converted_weights/ltx2_5 +strict_load_status: dev_and_distilled_pass +``` + +## Expected Parity Tests + +| Component | Official files / args | Test | Concerns | Status | +|---|---|---|---|---| +| Transformer | `ltx_core.model.transformer.model_configurator.LTXModelConfigurator`; checkpoint metadata | `tests/local_tests/transformers/test_ltx2_5_transformer_parity.py` | Official random FFN/keyframe/prompt-AdaLN parity plus strict dev/distilled production loading | real strict-load pass | +| Gemma 4 encoder and connector | `ltx_core.text_encoders.gemma`; unified Gemma 4 checkpoint | `tests/local_tests/ltx2_5/test_ltx2_5_gemma.py` | BOS behavior, connector bias surface, packed conversion | focused pass | +| Convolutional video VAE | `ltx_core.model.video_vae` configurator; conv checkpoint | `tests/local_tests/ltx2_5/test_ltx2_5_vae.py` | checkpoint-derived decoder width | config pass | +| Audio VAE/vocoder | `ltx_core.model.audio_vae`; audio VAE checkpoint | `tests/local_tests/ltx2_5/test_ltx2_5_conversion.py` | split routing and BWE topology reuse | real pipeline pass | +| Ancestral sampler | `ltx_core.components.diffusion_steps` and distilled pipeline selection | `tests/local_tests/ltx2_5/test_ltx2_5_sampler.py` | terminal x0, eta=0, reproducibility, seed offset, video-before-audio RNG | unit pass | +| T2AV/I2AV pipeline | official one-stage and distilled pipeline call paths | `tests/local_tests/ltx2_5/test_ltx2_5_registry.py` | dev/distilled presets, guidance, conditioning CRF, workloads | structural and real T2AV pass | + +Planned verification order: + +```bash +pytest -q tests/local_tests/ltx2_5 \ + tests/local_tests/transformers/test_ltx2_5_transformer_parity.py \ + -k 'not real_weights' + +pytest -q -rs \ + tests/local_tests/transformers/test_ltx2_5_transformer_parity.py \ + -k 'not real_weights' +``` + +The transformer prototype extends the existing native LTX-2 graph rather than +forking it. It adds the exact upstream v1.2.0 gates +`use_prompt_adaln_single=True`, `ff_bias=True`, `audio_ff_bias=True`, and +`use_keyframes_abs_pos_embedding=False`; these defaults preserve LTX-2/2.3. +Both published 2.5 BF16 variants disable video FFN bias and enable keyframe +absolute embeddings while retaining the prompt-AdaLN and audio-FFN-bias +defaults. The keyframe marker mask is sharded with video tokens on FastVideo's +SP path. + +The same warm Modal L40S shell ran the complete lightweight launch suite plus +legacy LTX preset preflight after the official source clone: 24 passed with the +two real-weight cases deselected. It then strict-loaded both 42 GB BF16 +transformers through the official builder and FastVideo's production loader: +development passed in 28.76s and distilled passed in 27.06s. + +Real pipeline validation then passed both distilled execution modes. The +one-stage smoke ran on 2xH100 and produced a joint H.264/AAC MP4 with 9 video +frames and 17 audio frames in 11.77s after load. The full two-stage smoke ran +on 2xB200: base ancestral denoising, spatial upsampling, the published +1,632-layer refinement LoRA, all three refinement steps, and joint decode/save +completed in 76.17s. The verified output contains 9 H.264 frames and 17 AAC +frames. W&B records the launch run at +`https://wandb.ai/aryan5v-san-jose-state-university/fastvideo-ltx2-5/runs/o9ixwewk`. + +CPU/reference parity coverage includes random-weight FFN output parity, exact +keyframe-marker functional parity, and official-versus-FastVideo parameter +surface checks for both legacy defaults and the 2.5 gates. The real-weight test +strict-loads dev and distilled checkpoints through the official builder and +FastVideo's production `TransformerLoader` when the gated assets are present. + +## Review Notes + +- Existing LTX-2 and LTX-2.3 checkpoints must remain behavior-identical when + all 2.5 configuration flags retain their backward-compatible defaults. +- Native FastVideo loading, compilation, offload, and sequence-parallel paths + are production requirements, not follow-up wrappers. +- A scaffold skip is not a parity pass. Real-weight claims require non-skipped + Modal results with exact commands and hardware recorded here. +- Before every push after the first PR publication, inspect new CodeRabbit + feedback and address actionable findings before updating the branch. diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_conversion.py b/tests/local_tests/ltx2_5/test_ltx2_5_conversion.py new file mode 100644 index 000000000..3d0ba8a76 --- /dev/null +++ b/tests/local_tests/ltx2_5/test_ltx2_5_conversion.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file + +from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig +from fastvideo.models.encoders.gemma import LTX2GemmaTextEncoderModel + + +REPO_ROOT = Path(__file__).resolve().parents[3] +CONVERTER_PATH = REPO_ROOT / "scripts" / "checkpoint_conversion" / "convert_ltx2_weights.py" +SPEC = importlib.util.spec_from_file_location("convert_ltx2_weights", CONVERTER_PATH) +assert SPEC is not None and SPEC.loader is not None +converter = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(converter) + + +def _uint8(value: str) -> torch.Tensor: + return torch.tensor(list(value.encode()), dtype=torch.uint8) + + +def _save( + path: Path, + tensors: dict[str, torch.Tensor], + *, + config: dict | None = None, + metadata: dict[str, str] | None = None, +) -> None: + safetensors_metadata = dict(metadata or {}) + if config is not None: + safetensors_metadata["config"] = json.dumps(config) + save_file(tensors, str(path), metadata=safetensors_metadata) + + +def _split_sources(tmp_path: Path) -> dict[str, Path]: + transformer_config = { + "transformer": { + "num_attention_heads": 2, + "attention_head_dim": 4, + "num_layers": 1, + "cross_attention_dim": 8, + "caption_channels": 6, + "audio_num_attention_heads": 1, + "audio_attention_head_dim": 3, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_cross_attention_dim": 3, + "cross_attention_adaln": True, + "caption_proj_before_connector": True, + "apply_gated_attention": True, + "use_prompt_adaln_single": False, + "ff_bias": False, + "audio_ff_bias": False, + "use_keyframes_abs_pos_embedding": False, + "caption_projection_first_linear": False, + "caption_proj_input_norm": True, + "caption_projection_second_linear": False, + "connector_num_attention_heads": 2, + "connector_attention_head_dim": 4, + "connector_num_layers": 1, + "audio_connector_num_attention_heads": 1, + "audio_connector_attention_head_dim": 3, + "audio_connector_num_layers": 1, + "connector_positional_embedding_max_pos": [2048], + "connector_apply_gated_attention": True, + "connector_ff_bias": False, + "frequencies_precision": "float64", + } + } + transformer = tmp_path / "transformer.safetensors" + _save( + transformer, + { + "model.diffusion_model.transformer_blocks.0.ff.net.0.proj.weight": torch.ones(2, 2), + "model.diffusion_model.video_embeddings_connector.transformer_1d_blocks.0.weight": torch.ones(1), + "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.0.weight": torch.ones(1), + }, + config=transformer_config, + ) + + gemma_config = { + "model_type": "gemma4_unified", + "pad_token_id": 0, + "eos_token_id": 2, + "text_config": { + "hidden_size": 6, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "max_position_embeddings": 1024, + "bos_token_id": 2, + }, + } + text_encoder = tmp_path / "text_encoder.safetensors" + _save( + text_encoder, + { + "model.layers.0.self_attn.q_proj.weight": torch.ones(2, 2), + "vision_model.embeddings.patch_embedding.weight": torch.ones(1), + "multi_modal_projector.embedding_projection.weight": torch.ones(1), + "audio_projector.proj.weight": torch.ones(1), + "text_embedding_projection.video_aggregate_embed.weight": torch.ones(8, 18), + "text_embedding_projection.video_aggregate_embed.bias": torch.ones(8), + "text_embedding_projection.audio_aggregate_embed.weight": torch.ones(3, 18), + "text_embedding_projection.audio_aggregate_embed.bias": torch.ones(3), + "tokenizer_json": _uint8("{}"), + "hf_asset__tokenizer_config.json": _uint8('{"tokenizer_class":"PreTrainedTokenizerFast"}'), + "hf_asset__processor_config.json": _uint8("{}"), + }, + metadata={"gemma_config": json.dumps(gemma_config)}, + ) + + vae = tmp_path / "vae.safetensors" + _save( + vae, + {"vae.encoder.conv.weight": torch.ones(1)}, + config={"vae": {"_class_name": "CausalVideoAutoencoder", "dims": 3}}, + ) + + audio_vae = tmp_path / "audio_vae.safetensors" + _save( + audio_vae, + { + "audio_vae.decoder.conv_in.weight": torch.ones(1), + "audio_vae.per_channel_statistics.mean-of-means": torch.ones(1), + "vocoder.vocoder.conv_pre.weight": torch.ones(1), + "vocoder.bwe_generator.conv_pre.weight": torch.ones(1), + }, + config={ + "audio_vae": {"model": {"params": {"ddconfig": {}}}}, + "vocoder": {"vocoder": {"resblock": "AMP1"}, "bwe": {"resblock": "AMP1"}}, + }, + ) + + spatial_upscaler = tmp_path / "spatial_upscaler.safetensors" + _save( + spatial_upscaler, + { + "model.initial_conv.weight": torch.ones(1), + "model.0.weight": torch.ones(1), + }, + config={ + "in_channels": 128, + "mid_channels": 512, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": True, + "temporal_upsample": False, + }, + ) + + distilled_lora = tmp_path / "distilled_lora.safetensors" + _save( + distilled_lora, + {"diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.ones(1)}, + ) + return { + "transformer": transformer, + "text_encoder": text_encoder, + "vae": vae, + "audio_vae": audio_vae, + "spatial_upscaler": spatial_upscaler, + "distilled_lora": distilled_lora, + } + + +def test_split_ltx2_5_conversion_routes_components_and_metadata(tmp_path: Path) -> None: + sources = _split_sources(tmp_path) + output = tmp_path / "converted" + + converter.convert_split_components( + transformer_source=sources["transformer"], + text_encoder_source=sources["text_encoder"], + vae_source=sources["vae"], + audio_vae_source=sources["audio_vae"], + spatial_upscaler_source=sources["spatial_upscaler"], + distilled_lora_source=sources["distilled_lora"], + output_dir=output, + transformer_class_name="LTX2Transformer3DModel", + variant="dev", + ) + + transformer_weights = load_file(str(output / "transformer" / "model.safetensors")) + assert set(transformer_weights) == {"transformer_blocks.0.ff.net.0.proj.weight"} + transformer_config = json.loads((output / "transformer" / "config.json").read_text()) + assert transformer_config["cross_attention_adaln"] is True + assert transformer_config["use_prompt_adaln_single"] is False + assert transformer_config["ff_bias"] is False + assert transformer_config["audio_ff_bias"] is False + assert transformer_config["use_keyframes_abs_pos_embedding"] is False + assert transformer_config["connector_ff_bias"] is False + assert transformer_config["double_precision_rope"] is True + + text_weights = load_file(str(output / "text_encoder" / "model.safetensors")) + assert set(text_weights) == { + "video_feature_extractor_linear.weight", + "video_feature_extractor_linear.bias", + "audio_feature_extractor_linear.weight", + "audio_feature_extractor_linear.bias", + "embeddings_connector.transformer_1d_blocks.0.weight", + "audio_embeddings_connector.transformer_1d_blocks.0.weight", + } + text_config = json.loads((output / "text_encoder" / "config.json").read_text()) + assert text_config["gemma_model_path"] == "gemma" + assert text_config["hidden_size"] == 6 + assert text_config["num_hidden_layers"] == 2 + assert text_config["feature_extractor_in_features"] == 18 + assert text_config["text_len"] == 1024 + assert text_config["bos_token_id"] == 2 + assert text_config["eos_token_id"] == 2 + assert text_config["video_feature_extractor_out_features"] == 8 + assert text_config["audio_feature_extractor_out_features"] == 3 + assert text_config["caption_proj_before_connector"] is True + assert text_config["connector_rope_type"] == "split" + assert text_config["connector_ff_bias"] is False + + # Validate the emitted config through the production FastVideo model + # constructor. The Gemma backbone remains lazy, so this is lightweight. + encoder_config = LTX2GemmaConfig() + encoder_config.update_model_arch(text_config) + encoder = LTX2GemmaTextEncoderModel(encoder_config) + assert encoder.video_feature_extractor_linear.in_features == 18 + assert encoder.video_feature_extractor_linear.out_features == 8 + assert encoder.audio_feature_extractor_linear.out_features == 3 + assert encoder.embeddings_connector.transformer_1d_blocks[0].ff.net[0].proj.bias is None + + gemma_weights = load_file(str(output / "text_encoder" / "gemma" / "model.safetensors")) + assert "model.language_model.layers.0.self_attn.q_proj.weight" in gemma_weights + assert "model.vision_embedder.embeddings.patch_embedding.weight" in gemma_weights + assert "model.embed_vision.embedding_projection.weight" in gemma_weights + assert "model.embed_audio.proj.weight" in gemma_weights + assert (output / "text_encoder" / "gemma" / "tokenizer.json").read_text() == "{}" + assert (output / "tokenizer" / "tokenizer_config.json").is_file() + + assert set(load_file(str(output / "vae" / "model.safetensors"))) == {"encoder.conv.weight"} + assert set(load_file(str(output / "audio_vae" / "model.safetensors"))) == { + "decoder.conv_in.weight", + "per_channel_statistics.mean-of-means", + } + assert set(load_file(str(output / "vocoder" / "model.safetensors"))) == { + "vocoder.conv_pre.weight", + "bwe_generator.conv_pre.weight", + } + assert set(load_file(str(output / "spatial_upsampler" / "model.safetensors"))) == { + "initial_conv.weight", + "upsampler.0.weight", + } + assert set(load_file(str(output / "distilled_lora" / "model.safetensors"))) == { + "diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight" + } + + model_index = json.loads((output / "model_index.json").read_text()) + assert model_index["transformer"] == ["diffusers", "LTX2Transformer3DModel"] + assert model_index["spatial_upsampler"] == ["diffusers", "LTX2LatentUpsampler"] + assert model_index["fastvideo_ltx2_variant"] == "ltx2.5-dev" + assert model_index["fastvideo_refine_enabled"] is True + assert model_index["fastvideo_refine_upsampler_path"] == "spatial_upsampler" + assert model_index["fastvideo_refine_lora_path"] == "distilled_lora/model.safetensors" + + +def test_split_model_index_enables_distilled_refine_without_lora() -> None: + model_index = converter._build_split_model_index( + transformer_class_name="LTX2Transformer3DModel", + pipeline_class_name="LTX2Pipeline", + diffusers_version="0.33.0.dev0", + variant="distilled", + distilled_lora=False, + ) + assert model_index["fastvideo_refine_enabled"] is True + assert "fastvideo_refine_lora_path" not in model_index + + +def test_split_ltx2_5_conversion_preserves_enabled_transformer_gates(tmp_path: Path) -> None: + """Transformer config filtering retains True-valued 2.5 architecture gates.""" + # Create a fresh fixture with all 2.5 gates enabled + transformer_config = { + "transformer": { + "num_attention_heads": 2, + "attention_head_dim": 4, + "num_layers": 1, + "cross_attention_dim": 8, + "caption_channels": 6, + "audio_num_attention_heads": 1, + "audio_attention_head_dim": 3, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_cross_attention_dim": 3, + "cross_attention_adaln": True, + "caption_proj_before_connector": True, + "apply_gated_attention": True, + "use_prompt_adaln_single": True, + "ff_bias": True, + "audio_ff_bias": True, + "use_keyframes_abs_pos_embedding": True, + "caption_projection_first_linear": False, + "caption_proj_input_norm": True, + "caption_projection_second_linear": False, + "connector_num_attention_heads": 2, + "connector_attention_head_dim": 4, + "connector_num_layers": 1, + "audio_connector_num_attention_heads": 1, + "audio_connector_attention_head_dim": 3, + "audio_connector_num_layers": 1, + "connector_positional_embedding_max_pos": [2048], + "connector_apply_gated_attention": True, + "connector_ff_bias": False, + "frequencies_precision": "float64", + } + } + transformer = tmp_path / "transformer_gates_enabled.safetensors" + _save( + transformer, + { + "model.diffusion_model.transformer_blocks.0.ff.net.0.proj.weight": torch.ones(2, 2), + "model.diffusion_model.video_embeddings_connector.transformer_1d_blocks.0.weight": torch.ones(1), + "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.0.weight": torch.ones(1), + }, + config=transformer_config, + ) + + # Reuse other fixture sources + sources = _split_sources(tmp_path) + output = tmp_path / "converted_gates_enabled" + + converter.convert_split_components( + transformer_source=transformer, + text_encoder_source=sources["text_encoder"], + vae_source=sources["vae"], + audio_vae_source=sources["audio_vae"], + spatial_upscaler_source=sources["spatial_upscaler"], + distilled_lora_source=sources["distilled_lora"], + output_dir=output, + transformer_class_name="LTX2Transformer3DModel", + variant="dev", + ) + + transformer_config_out = json.loads((output / "transformer" / "config.json").read_text()) + assert transformer_config_out["use_prompt_adaln_single"] is True + assert transformer_config_out["ff_bias"] is True + assert transformer_config_out["audio_ff_bias"] is True + assert transformer_config_out["use_keyframes_abs_pos_embedding"] is True + + +def test_legacy_monolithic_conversion_behavior_is_preserved(tmp_path: Path) -> None: + source = tmp_path / "legacy.safetensors" + metadata_config = { + "transformer": { + "num_attention_heads": 2, + "frequencies_precision": "float64", + "cross_attention_adaln": True, + "caption_proj_before_connector": True, + "apply_gated_attention": True, + "use_prompt_adaln_single": False, + "ff_bias": False, + "audio_ff_bias": False, + "use_keyframes_abs_pos_embedding": False, + }, + "vae": {}, + "audio_vae": {}, + "vocoder": {}, + } + _save( + source, + { + "model.diffusion_model.transformer_blocks.0.weight": torch.ones(1), + "vae.encoder.weight": torch.ones(1), + "audio_vae.decoder.weight": torch.ones(1), + "vocoder.conv_pre.weight": torch.ones(1), + "text_embedding_projection.video_aggregate_embed.weight": torch.ones(1), + }, + config=metadata_config, + ) + output = tmp_path / "legacy-converted" + converter.convert_components( + source, + output, + metadata_config, + "LTX2Transformer3DModel", + ) + + assert set(load_file(str(output / "transformer" / "model.safetensors"))) == { + "transformer_blocks.0.weight" + } + # Legacy monolith keeps its historical projection key surface. + assert set(load_file(str(output / "text_encoder" / "model.safetensors"))) == { + "video_aggregate_embed.weight" + } + config = json.loads((output / "transformer" / "config.json").read_text()) + assert config["cross_attention_adaln"] is True + assert config["use_prompt_adaln_single"] is False + assert config["ff_bias"] is False + assert config["audio_ff_bias"] is False + assert config["use_keyframes_abs_pos_embedding"] is False + assert config["double_precision_rope"] is True + assert (output / "model_index.json").is_file() diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_gemma.py b/tests/local_tests/ltx2_5/test_ltx2_5_gemma.py new file mode 100644 index 000000000..01eb2aebc --- /dev/null +++ b/tests/local_tests/ltx2_5/test_ltx2_5_gemma.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Focused compatibility tests for the LTX-2.5 packed Gemma 4 path.""" + +from types import SimpleNamespace + +import torch + +from fastvideo.models.dits.ltx2 import LTXRopeType +from fastvideo.models.encoders.gemma import ( + Embeddings1DConnector, + GemmaConnectorConfig, + _ensure_leading_bos, + _get_bos_token_id, +) + + +def test_gemma4_bos_token_id_comes_from_nested_text_config() -> None: + config = SimpleNamespace(text_config=SimpleNamespace(bos_token_id=2)) + assert _get_bos_token_id(config) == 2 + + +def test_ensure_leading_bos_preserves_gemma3_and_fixes_gemma4() -> None: + input_ids = torch.tensor([ + [0, 0, 2, 10, 11], + [0, 0, 0, 10, 11], + [10, 11, 12, 13, 14], + ]) + attention_mask = torch.tensor([ + [0, 0, 1, 1, 1], + [0, 0, 0, 1, 1], + [1, 1, 1, 1, 1], + ]) + + actual_ids, actual_mask = _ensure_leading_bos( + input_ids, + attention_mask, + bos_token_id=2, + ) + + torch.testing.assert_close(actual_ids[0], input_ids[0]) + torch.testing.assert_close(actual_mask[0], attention_mask[0]) + torch.testing.assert_close(actual_ids[1], torch.tensor([0, 0, 2, 10, 11])) + torch.testing.assert_close(actual_mask[1], torch.tensor([0, 0, 1, 1, 1])) + torch.testing.assert_close(actual_ids[2], torch.tensor([2, 10, 11, 12, 13])) + torch.testing.assert_close(actual_mask[2], torch.ones(5, dtype=torch.long)) + + +def test_ltx2_5_connector_ffn_can_omit_bias() -> None: + connector = Embeddings1DConnector( + GemmaConnectorConfig( + num_attention_heads=2, + attention_head_dim=4, + num_layers=1, + positional_embedding_theta=10000.0, + positional_embedding_max_pos=[4096], + rope_type=LTXRopeType.SPLIT, + double_precision_rope=False, + num_learnable_registers=None, + apply_gated_attention=False, + ff_bias=False, + )) + + feed_forward = connector.transformer_1d_blocks[0].ff + assert feed_forward.net[0].proj.bias is None + assert feed_forward.net[2].bias is None diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_registry.py b/tests/local_tests/ltx2_5/test_ltx2_5_registry.py new file mode 100644 index 000000000..9fa437175 --- /dev/null +++ b/tests/local_tests/ltx2_5/test_ltx2_5_registry.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Registry and preset contracts for LTX-2.5 launch inference.""" + +import pytest + +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.fastvideo_args import WorkloadType +from fastvideo.registry import get_preset_selection, get_registered_models_with_workloads + + +@pytest.mark.parametrize( + ("model_path", "preset_name", "ancestral"), + ( + ("FastVideo/LTX-2.5-Dev-Diffusers", "ltx2_5_dev", False), + ("FastVideo/LTX-2.5-Distilled-Diffusers", "ltx2_5_distilled_two_stage", True), + ), +) +def test_ltx2_5_registry_precedence_and_presets( + model_path: str, + preset_name: str, + ancestral: bool, +) -> None: + assert get_preset_selection(model_path) == (preset_name, "ltx2") + registered = {entry["id"]: entry for entry in get_registered_models_with_workloads()} + assert set(registered[model_path]["workload_types"]) == { + WorkloadType.T2V.value, + WorkloadType.I2V.value, + } + + sampling = SamplingParam.from_pretrained(model_path) + assert sampling.ltx2_image_crf == 18.0 + assert sampling.ltx2_use_ancestral_sampler is ancestral + + +def test_legacy_ltx_distilled_remains_deterministic() -> None: + sampling = SamplingParam.from_pretrained("FastVideo/LTX2-Distilled-Diffusers") + + assert sampling.ltx2_use_ancestral_sampler is False + assert sampling.ltx2_image_crf == 33.0 + + +def test_ltx2_5_metadata_variant_selects_preset(tmp_path) -> None: + """Local directories with fastvideo_ltx2_variant metadata route to the correct preset.""" + import json + from pathlib import Path + + # Create a local directory with LTX2Pipeline _class_name and distilled variant metadata + distilled_dir = tmp_path / "ltx2_5_distilled_local" + distilled_dir.mkdir() + model_index = { + "_class_name": "LTX2Pipeline", + "fastvideo_ltx2_variant": "ltx2.5-distilled", + } + (distilled_dir / "model_index.json").write_text(json.dumps(model_index)) + + preset, family = get_preset_selection(str(distilled_dir)) + assert preset == "ltx2_5_distilled_two_stage" + assert family == "ltx2" + + # Create a local directory with LTX2Pipeline _class_name and dev variant metadata + dev_dir = tmp_path / "ltx2_5_dev_local" + dev_dir.mkdir() + model_index = { + "_class_name": "LTX2Pipeline", + "fastvideo_ltx2_variant": "ltx2.5-dev", + } + (dev_dir / "model_index.json").write_text(json.dumps(model_index)) + + preset, family = get_preset_selection(str(dev_dir)) + assert preset == "ltx2_5_dev" + assert family == "ltx2" + + +def test_ltx2_5_neutral_path_with_class_name_routes_correctly(tmp_path) -> None: + """Neutral local paths with _class_name=LTX2Pipeline and variant metadata route correctly.""" + import json + + # A path with no ltx-2.5 tokens but LTX2Pipeline class and distilled variant + neutral_dir = tmp_path / "my_converted_checkpoint" + neutral_dir.mkdir() + model_index = { + "_class_name": "LTX2Pipeline", + "fastvideo_ltx2_variant": "ltx2.5-distilled", + } + (neutral_dir / "model_index.json").write_text(json.dumps(model_index)) + + preset, family = get_preset_selection(str(neutral_dir)) + assert preset == "ltx2_5_distilled_two_stage" + assert family == "ltx2" diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_sampler.py b/tests/local_tests/ltx2_5/test_ltx2_5_sampler.py new file mode 100644 index 000000000..62d453d96 --- /dev/null +++ b/tests/local_tests/ltx2_5/test_ltx2_5_sampler.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit coverage for the LTX-2.5 distilled ancestral sampler.""" + +import torch + +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.pipelines.basic.ltx2.stages.ltx2_denoising import ( + ANCESTRAL_NOISE_SEED_OFFSET, + _ltx2_first_frame_keyframes_mask, + _ltx2_euler_ancestral_step, +) + + +def test_ltx2_5_preset_enables_ancestral_sampling() -> None: + params = SamplingParam.from_pretrained("FastVideo/LTX-2.5-Distilled-Diffusers") + + assert params.ltx2_use_ancestral_sampler is True + assert params.num_inference_steps == 8 + assert params.ltx2_image_crf == 18.0 + + +def test_ancestral_terminal_step_returns_denoised_prediction() -> None: + sample = torch.tensor([3.0], dtype=torch.bfloat16) + denoised = torch.tensor([1.25], dtype=torch.bfloat16) + + actual = _ltx2_euler_ancestral_step( + sample, + denoised, + torch.tensor(0.421875), + torch.tensor(0.0), + noise=None, + ) + + torch.testing.assert_close(actual, denoised) + + +def test_ancestral_eta_zero_matches_deterministic_euler() -> None: + sample = torch.tensor([2.0, -1.0], dtype=torch.float32) + denoised = torch.tensor([0.5, 0.25], dtype=torch.float32) + sigma = torch.tensor(0.8) + sigma_next = torch.tensor(0.3) + expected = sample + ((sample - denoised) / sigma) * (sigma_next - sigma) + + actual = _ltx2_euler_ancestral_step( + sample, + denoised, + sigma, + sigma_next, + noise=None, + eta=0.0, + ) + + torch.testing.assert_close(actual, expected) + + +def test_ancestral_noise_stream_is_reproducible_and_video_first() -> None: + """Verify ancestral noise uses the seeded offset and video-first ordering.""" + seed = 17 + ANCESTRAL_NOISE_SEED_OFFSET + generator = torch.Generator(device="cpu").manual_seed(seed) + video_noise = torch.randn((1, 2), generator=generator) + audio_noise = torch.randn((1, 3), generator=generator) + + replay = torch.Generator(device="cpu").manual_seed(seed) + torch.testing.assert_close(video_noise, torch.randn((1, 2), generator=replay)) + torch.testing.assert_close(audio_noise, torch.randn((1, 3), generator=replay)) + + +def test_ancestral_step_captures_noise_from_stage_forward(monkeypatch) -> None: + """LTX2DenoisingStage.forward generates ancestral noise with the expected seed, shapes, and ordering.""" + from unittest.mock import MagicMock + from fastvideo.pipelines.basic.ltx2.stages.ltx2_denoising import LTX2DenoisingStage + from fastvideo.pipelines.composed_pipeline_base import ForwardBatch + from fastvideo.fastvideo_args import FastVideoArgs + + # Track noise tensors passed to _ltx2_euler_ancestral_step + captured_noises = [] + original_step = _ltx2_euler_ancestral_step + + def mock_step(sample, denoised, sigma, sigma_next, noise=None, eta=1.0): + if noise is not None: + captured_noises.append(noise.clone()) + return original_step(sample, denoised, sigma, sigma_next, noise=noise, eta=eta) + + monkeypatch.setattr( + "fastvideo.pipelines.basic.ltx2.stages.ltx2_denoising._ltx2_euler_ancestral_step", + mock_step, + ) + + # Create a minimal stage with a mock transformer + mock_transformer = MagicMock() + mock_transformer.return_value = (torch.zeros(1, 12, 4), torch.zeros(1, 18, 4)) + stage = LTX2DenoisingStage( + mock_transformer, + sigmas_override=[0.8, 0.3], + num_inference_steps_override=1, + ) + + # Create a batch with video and audio latents matching expected shapes + batch = ForwardBatch( + latents=torch.randn(1, 4, 3, 2, 2), # video: (B, C, T, H, W) + extra={ + "ltx2_audio_latents": torch.randn(1, 4, 6), # audio: (B, C, T) + "ltx2_conditioning_video_latents": None, + "ltx2_conditioning_audio_latents": None, + }, + ) + batch.text_embeddings = torch.randn(1, 10, 8) + batch.timestep_embeddings = torch.randn(1) + + args = FastVideoArgs( + model_path="dummy", + seed=17, + ltx2_use_ancestral_sampler=True, + ) + + # Run the stage forward to generate ancestral noise + stage.forward(batch, args) + + # Verify we captured noise tensors (one per ancestral step) + assert len(captured_noises) == 1 + noise = captured_noises[0] + + # Verify the noise tensor contains both video and audio with video-first ordering + # Video latents are (1, 4, 3, 2, 2) = 48 elements; audio latents are (1, 4, 6) = 24 elements + # Combined flattened noise should have 48 + 24 = 72 elements + assert noise.numel() == 48 + 24 + + # Verify reproducibility with the seeded offset + expected_seed = 17 + ANCESTRAL_NOISE_SEED_OFFSET + generator = torch.Generator(device="cpu").manual_seed(expected_seed) + video_noise_expected = torch.randn((1, 4, 3, 2, 2), generator=generator) + audio_noise_expected = torch.randn((1, 4, 6), generator=generator) + + # The noise tensor is concatenated [video, audio] in flattened form + noise_video = noise[:48].view(1, 4, 3, 2, 2) + noise_audio = noise[48:].view(1, 4, 6) + torch.testing.assert_close(noise_video, video_noise_expected) + torch.testing.assert_close(noise_audio, audio_noise_expected) + + +def test_keyframe_mask_marks_every_token_in_first_causal_latent_frame() -> None: + mask = _ltx2_first_frame_keyframes_mask( + batch_size=2, + token_count=12, + latent_frames=3, + device=torch.device("cpu"), + ) + + expected = torch.zeros((2, 12, 1)) + expected[:, :4] = 1.0 + torch.testing.assert_close(mask, expected) diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_vae.py b/tests/local_tests/ltx2_5/test_ltx2_5_vae.py new file mode 100644 index 000000000..5e3657aa7 --- /dev/null +++ b/tests/local_tests/ltx2_5/test_ltx2_5_vae.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configuration compatibility for the LTX-2.5 convolutional video VAE.""" + +from fastvideo.models.vaes.ltx2vae import VideoDecoderConfigurator + + +def test_conv_vae_honors_checkpoint_decoder_base_channels() -> None: + decoder = VideoDecoderConfigurator.from_config({ + "vae": { + "dims": 3, + "latent_channels": 8, + "out_channels": 3, + "decoder_blocks": [], + "patch_size": 1, + "norm_layer": "pixel_norm", + "causal_decoder": False, + "timestep_conditioning": False, + "decoder_base_channels": 12, + } + }) + + assert decoder.conv_in.in_channels == 8 + assert decoder.conv_in.out_channels == 12 diff --git a/tests/local_tests/transformers/test_ltx2_5_transformer_parity.py b/tests/local_tests/transformers/test_ltx2_5_transformer_parity.py new file mode 100644 index 000000000..a8ddf4e6c --- /dev/null +++ b/tests/local_tests/transformers/test_ltx2_5_transformer_parity.py @@ -0,0 +1,316 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Focused parity coverage for the config-gated LTX-2.5 DiT extension. + +The CPU tests use random weights to compare the new stateful/functional delta +against pinned ``ltx-core`` v1.2.0. The real-weight test exercises the official +builder and FastVideo's production ``TransformerLoader`` when the gated assets +have been provisioned. +""" + +from __future__ import annotations + +import importlib +import os +from pathlib import Path +import sys + +import pytest +import torch +from torch.testing import assert_close + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_REF_DIR = Path(os.environ.get("LTX2_5_OFFICIAL_REF_DIR", REPO_ROOT / "LTX-2-Reference")) +OFFICIAL_SRC = OFFICIAL_REF_DIR / "packages" / "ltx-core" / "src" +OFFICIAL_WEIGHTS_DIR = Path( + os.environ.get( + "LTX2_5_OFFICIAL_WEIGHTS_DIR", + REPO_ROOT / "official_weights" / "LTX-2.5" / "diffusion_models", + ) +) +CONVERTED_WEIGHTS_DIR = Path( + os.environ.get( + "LTX2_5_CONVERTED_WEIGHTS_DIR", + REPO_ROOT / "converted_weights" / "ltx2_5", + ) +) +PARITY_SCOPE = "both" + + +@pytest.fixture(autouse=True) +def _setup_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: + """Set default environment variables for all tests in this module.""" + monkeypatch.setenv("MASTER_ADDR", "localhost") + monkeypatch.setenv("MASTER_PORT", "29525") + monkeypatch.setenv("DISABLE_SP", "1") + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + + +def _import_official(module_name: str): + if not OFFICIAL_SRC.is_dir(): + pytest.skip(f"Pinned LTX-2 v1.2.0 reference is missing: {OFFICIAL_REF_DIR}") + if str(OFFICIAL_SRC) not in sys.path: + sys.path.insert(0, str(OFFICIAL_SRC)) + try: + return importlib.import_module(module_name) + except Exception as exc: # noqa: BLE001 - report the exact missing reference dependency. + pytest.skip(f"Cannot import pinned official module {module_name}: {exc}") + + +def _tiny_model_kwargs() -> dict: + return { + "num_attention_heads": 2, + "attention_head_dim": 4, + "in_channels": 4, + "out_channels": 4, + "num_layers": 1, + "cross_attention_dim": 8, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_cross_attention_dim": 8, + "cross_attention_adaln": True, + } + + +@pytest.mark.parametrize("bias", (False, True)) +def test_feed_forward_random_weight_parity(bias: bool) -> None: + """The 2.5 bias gate must change both FFN linears, not their math/layout.""" + official_module = _import_official("ltx_core.model.transformer.feed_forward") + from fastvideo.models.dits.ltx2 import FeedForward as FastVideoFeedForward + + torch.manual_seed(20260812) + official = official_module.FeedForward(dim=8, dim_out=8, bias=bias).eval() + fastvideo = FastVideoFeedForward(dim=8, dim_out=8, bias=bias).eval() + + official_state = official.state_dict() + fastvideo_state = fastvideo.state_dict() + assert official_state.keys() == fastvideo_state.keys() + fastvideo.load_state_dict(official_state, strict=True) + + inputs = torch.randn(2, 5, 8) + with torch.inference_mode(): + expected = official(inputs) + actual = fastvideo(inputs) + assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + +def test_keyframe_absolute_embedding_random_weight_parity() -> None: + """Match upstream marker semantics for positive, zero, and negative mask values.""" + official_module = _import_official("ltx_core.model.transformer.transformer_args") + from fastvideo.models.dits.ltx2 import apply_keyframes_absolute_embedding + + torch.manual_seed(20260812) + hidden_states = torch.randn(2, 4, 8) + embedding = torch.randn(1, 8) + keyframes_mask = torch.tensor( + [[[1.0], [0.0], [-1.0], [2.0]], [[0.0], [1.0], [0.0], [0.5]]] + ) + provider = lambda: embedding + + expected = official_module.apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, provider) + actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, provider) + assert_close(actual, expected, atol=0.0, rtol=0.0) + assert apply_keyframes_absolute_embedding(hidden_states, None, provider) is hidden_states + + +@pytest.mark.parametrize("dynamic_prompt", (False, True)) +def test_cross_attention_adaln_static_and_dynamic_prompt_parity(dynamic_prompt: bool) -> None: + """Exercise the static 2.5 K/V table and legacy dynamic prompt AdaLN paths.""" + official_module = _import_official("ltx_core.model.transformer.transformer") + from fastvideo.models.dits.ltx2 import _rms_norm_dispatch + from fastvideo.models.dits.ltx2 import apply_cross_attention_adaln + + torch.manual_seed(20260812) + batch_size, sequence_length, dim = 2, 3, 8 + x = torch.randn(batch_size, sequence_length, dim) + context = torch.randn_like(x) + q_shift = torch.randn_like(x) + q_scale = torch.randn_like(x) + q_gate = torch.randn_like(x) + prompt_scale_shift_table = torch.randn(2, dim) + prompt_timestep = torch.randn(batch_size, sequence_length, 2 * dim) if dynamic_prompt else None + + def attention(inputs, *, context, mask): + assert mask is None + return inputs + context + + expected = official_module.apply_cross_attention_adaln( + _rms_norm_dispatch(x, eps=1e-6), + context, + attention, + q_shift, + q_scale, + q_gate, + prompt_scale_shift_table, + prompt_timestep, + ) + actual = apply_cross_attention_adaln( + x, + context, + attention, + q_shift, + q_scale, + q_gate, + prompt_scale_shift_table, + prompt_timestep, + ) + assert_close(actual, expected, atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize( + ("use_prompt_adaln_single", "ff_bias", "audio_ff_bias", "use_keyframes_abs_pos_embedding"), + ( + (True, True, True, False), # LTX-2/2.3-compatible defaults. + (False, False, False, True), # LTX-2.5-style architecture delta. + ), +) +def test_variant_structure_matches_official( + use_prompt_adaln_single: bool, + ff_bias: bool, + audio_ff_bias: bool, + use_keyframes_abs_pos_embedding: bool, +) -> None: + """Compare the exact parameter surface controlled by the four new flags.""" + official_model_module = _import_official("ltx_core.model.transformer.model") + from fastvideo.models.dits.ltx2 import LTXModel as FastVideoLTXModel + from fastvideo.models.dits.ltx2 import LTXModelType as FastVideoLTXModelType + + common = _tiny_model_kwargs() + variant = { + "use_prompt_adaln_single": use_prompt_adaln_single, + "ff_bias": ff_bias, + "audio_ff_bias": audio_ff_bias, + "use_keyframes_abs_pos_embedding": use_keyframes_abs_pos_embedding, + } + official = official_model_module.LTXModel( + model_type=official_model_module.LTXModelType.AudioVideo, + **common, + **variant, + ) + fastvideo = FastVideoLTXModel( + model_type=FastVideoLTXModelType.AudioVideo, + caption_proj_before_connector=True, + **common, + **variant, + ) + + assert (official.prompt_adaln_single is not None) == (fastvideo.prompt_adaln_single is not None) + assert (official.audio_prompt_adaln_single is not None) == (fastvideo.audio_prompt_adaln_single is not None) + + official_state = official.state_dict() + fastvideo_state = fastvideo.state_dict() + for prefix in ("transformer_blocks.0.ff.", "transformer_blocks.0.audio_ff."): + official_keys = {name for name in official_state if name.startswith(prefix)} + fastvideo_keys = {name for name in fastvideo_state if name.startswith(prefix)} + assert official_keys == fastvideo_keys + + keyframe_name = "keyframes_abs_pos_embedding" + assert (keyframe_name in official_state) == use_keyframes_abs_pos_embedding + assert (keyframe_name in fastvideo_state) == use_keyframes_abs_pos_embedding + if use_keyframes_abs_pos_embedding: + assert official_state[keyframe_name].shape == fastvideo_state[keyframe_name].shape == (1, 8) + + +def test_arch_config_defaults_preserve_pre_2_5_checkpoints() -> None: + from fastvideo.configs.models.dits.ltx2 import LTX2VideoArchConfig + + config = LTX2VideoArchConfig() + assert config.use_prompt_adaln_single is True + assert config.ff_bias is True + assert config.audio_ff_bias is True + assert config.use_keyframes_abs_pos_embedding is False + + +def _read_transformer_metadata(official_path: Path) -> dict: + loader_module = _import_official("ltx_core.loader.sft_loader") + metadata = loader_module.SafetensorsModelStateDictLoader().metadata(str(official_path)) + transformer = metadata.get("config", {}).get("transformer", {}) + if not transformer: + pytest.fail(f"LTX-2.5 checkpoint has no config.transformer metadata: {official_path}", pytrace=False) + return transformer + + +def _fastvideo_config_from_metadata(metadata: dict): + from fastvideo.configs.models.dits import LTX2VideoConfig + + config = LTX2VideoConfig() + arch = config.arch_config + for name in ( + "num_attention_heads", + "attention_head_dim", + "num_layers", + "cross_attention_dim", + "caption_channels", + "norm_eps", + "positional_embedding_theta", + "positional_embedding_max_pos", + "timestep_scale_multiplier", + "use_middle_indices_grid", + "rope_type", + "audio_num_attention_heads", + "audio_attention_head_dim", + "audio_in_channels", + "audio_out_channels", + "audio_cross_attention_dim", + "audio_positional_embedding_max_pos", + "av_ca_timestep_scale_multiplier", + "in_channels", + "out_channels", + "cross_attention_adaln", + "caption_proj_before_connector", + "apply_gated_attention", + "use_prompt_adaln_single", + "ff_bias", + "audio_ff_bias", + "use_keyframes_abs_pos_embedding", + ): + if name in metadata: + setattr(arch, name, metadata[name]) + arch.double_precision_rope = metadata.get("frequencies_precision", "") == "float64" + return config + + +@pytest.mark.parametrize("variant", ("dev", "distilled")) +def test_real_weights_load_strictly_through_production_loader(variant: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Strict-load both gated 2.5 transformer variants when provisioned.""" + official_path = OFFICIAL_WEIGHTS_DIR / f"ltx-2.5-22b-{variant}-transformer-bf16.safetensors" + converted_dir = CONVERTED_WEIGHTS_DIR / variant / "transformer" + missing = [str(path) for path in (official_path, converted_dir) if not path.exists()] + if missing: + pytest.skip(f"Gated LTX-2.5 transformer assets are not provisioned: {missing}") + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.fail("Provisioned LTX-2.5 weights require a bf16-capable CUDA GPU", pytrace=False) + + official_loader = _import_official("ltx_core.loader.single_gpu_model_builder") + official_transformer = _import_official("ltx_core.model.transformer") + + metadata = _read_transformer_metadata(official_path) + official = official_loader.SingleGPUModelBuilder( + model_class_configurator=official_transformer.LTXModelConfigurator, + model_path=str(official_path), + model_sd_ops=official_transformer.LTXV_MODEL_COMFY_RENAMING_MAP, + ).build(device=torch.device("cuda:0"), dtype=torch.bfloat16) + assert official.use_prompt_adaln_single == metadata.get("use_prompt_adaln_single", True) + assert official.use_keyframes_abs_pos_embedding == metadata.get("use_keyframes_abs_pos_embedding", False) + del official + torch.cuda.empty_cache() + + from fastvideo.configs.pipelines import PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.dits import ltx2 as fastvideo_ltx2 + from fastvideo.models.loader.component_loader import TransformerLoader + + monkeypatch.setattr(fastvideo_ltx2, "get_sp_world_size", lambda: 1) + + config = _fastvideo_config_from_metadata(metadata) + args = FastVideoArgs( + model_path=str(converted_dir), + dit_cpu_offload=True, + use_fsdp_inference=False, + pipeline_config=PipelineConfig(dit_config=config, dit_precision="bf16"), + ) + args.device = torch.device("cuda:0") + model = TransformerLoader().load(str(converted_dir), args) + assert model.model.use_prompt_adaln_single == metadata.get("use_prompt_adaln_single", True) + assert model.model.use_keyframes_abs_pos_embedding == metadata.get("use_keyframes_abs_pos_embedding", False)