diff --git a/examples/inference/basic/basic_ltx2.py b/examples/inference/basic/basic_ltx2.py new file mode 100644 index 000000000..d8acd56bc --- /dev/null +++ b/examples/inference/basic/basic_ltx2.py @@ -0,0 +1,34 @@ +from fastvideo import VideoGenerator + +PROMPT = ( + "A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\" The camera slowly " + "pans right, revealing the grandfather in the garden wearing enormous " + "butterfly wings, waving his arms in the air like he's trying to take " + "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " + "The woman covers her face, on the verge of tears. The tone is deadpan, " + "absurd, and quietly tragic." +) + + +def main() -> None: + generator = VideoGenerator.from_pretrained( + "FastVideo/LTX2-Distilled-Diffusers", + num_gpus=1, + ) + + output_path = "outputs_video/ltx2_basic/output_ltx2_distilled_t2v.mp4" + generator.generate_video( + prompt=PROMPT, + output_path=output_path, + save_video=True, + ) + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/models/__init__.py b/fastvideo/configs/models/__init__.py index 286c6b9dd..aa05ba62e 100644 --- a/fastvideo/configs/models/__init__.py +++ b/fastvideo/configs/models/__init__.py @@ -2,5 +2,16 @@ from fastvideo.configs.models.base import ModelConfig from fastvideo.configs.models.dits.base import DiTConfig from fastvideo.configs.models.encoders.base import EncoderConfig from fastvideo.configs.models.vaes.base import VAEConfig +from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig, + LTX2AudioEncoderConfig, + LTX2VocoderConfig) -__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"] +__all__ = [ + "ModelConfig", + "VAEConfig", + "DiTConfig", + "EncoderConfig", + "LTX2AudioEncoderConfig", + "LTX2AudioDecoderConfig", + "LTX2VocoderConfig", +] diff --git a/fastvideo/configs/models/audio/__init__.py b/fastvideo/configs/models/audio/__init__.py new file mode 100644 index 000000000..ad1035536 --- /dev/null +++ b/fastvideo/configs/models/audio/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.configs.models.audio.ltx2_audio_vae import ( + LTX2AudioDecoderConfig, + LTX2AudioEncoderConfig, + LTX2VocoderConfig, +) + +__all__ = [ + "LTX2AudioEncoderConfig", + "LTX2AudioDecoderConfig", + "LTX2VocoderConfig", +] diff --git a/fastvideo/configs/models/audio/ltx2_audio_vae.py b/fastvideo/configs/models/audio/ltx2_audio_vae.py new file mode 100644 index 000000000..829f65a07 --- /dev/null +++ b/fastvideo/configs/models/audio/ltx2_audio_vae.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 audio VAE and vocoder configuration. +""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.base import ArchConfig, ModelConfig + + +@dataclass +class LTX2AudioArchConfig(ArchConfig): + architectures: list[str] = field(default_factory=list) + + +@dataclass +class LTX2AudioEncoderConfig(ModelConfig): + arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig( + architectures=["LTX2AudioEncoder"])) + + +@dataclass +class LTX2AudioDecoderConfig(ModelConfig): + arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig( + architectures=["LTX2AudioDecoder"])) + + +@dataclass +class LTX2VocoderConfig(ModelConfig): + arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig( + architectures=["LTX2Vocoder"])) diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 958e9bc71..c4b21a714 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -3,11 +3,12 @@ from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config from fastvideo.configs.models.dits.longcat import LongCatVideoConfig +from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.stepvideo import StepVideoConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig", "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", - "LongCatVideoConfig" + "LongCatVideoConfig", "LTX2VideoConfig" ] diff --git a/fastvideo/configs/models/dits/ltx2.py b/fastvideo/configs/models/dits/ltx2.py new file mode 100644 index 000000000..2c280a38a --- /dev/null +++ b/fastvideo/configs/models/dits/ltx2.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 Transformer configuration for native FastVideo integration. +""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_ltx2_blocks(name: str, _module) -> bool: + """FSDP shard condition for LTX-2 transformer blocks.""" + return "transformer_blocks" in name + + +@dataclass +class LTX2VideoArchConfig(DiTArchConfig): + """Architecture configuration for LTX-2 video transformer.""" + + _fsdp_shard_conditions: list = field( + default_factory=lambda: [is_ltx2_blocks]) + _compile_conditions: list = field(default_factory=lambda: [is_ltx2_blocks]) + + # Parameter name mapping for weight conversion (hf/comfy -> FastVideo) + param_names_mapping: dict = field( + default_factory=lambda: { + r"^model\.diffusion_model\.(.*)$": r"model.\1", + r"^diffusion_model\.(.*)$": r"model.\1", + r"^model\.(.*)$": r"model.\1", + r"^(.*)$": r"model.\1", + }) + + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + lora_param_names_mapping: dict = field(default_factory=lambda: {}) + + # Core transformer settings (defaults from LTX-2 metadata) + num_attention_heads: int = 32 + attention_head_dim: int = 128 + num_layers: int = 48 + cross_attention_dim: int = 4096 + caption_channels: int = 3840 + norm_eps: float = 1e-6 + attention_type: str = "default" + rope_type: str = "split" + double_precision_rope: bool = True + + positional_embedding_theta: float = 10000.0 + positional_embedding_max_pos: list[int] = field( + default_factory=lambda: [20, 2048, 2048]) + timestep_scale_multiplier: int = 1000 + use_middle_indices_grid: bool = True + + # Patchification (video-only path) + patch_size: tuple[int, int, int] = (1, 1, 1) + num_channels_latents: int = 128 + in_channels: int | None = None + out_channels: int | None = None + + # Audio defaults (reserved for joint AV ports) + audio_num_attention_heads: int = 32 + audio_attention_head_dim: int = 64 + audio_in_channels: int = 128 + audio_out_channels: int = 128 + audio_cross_attention_dim: int = 2048 + audio_positional_embedding_max_pos: list[int] = field( + default_factory=lambda: [20]) + av_ca_timestep_scale_multiplier: int = 1 + + def __post_init__(self): + super().__post_init__() + patch_volume = self.patch_size[0] * self.patch_size[ + 1] * self.patch_size[2] + if self.in_channels is None: + self.in_channels = self.num_channels_latents * patch_volume + if self.out_channels is None: + self.out_channels = self.in_channels + + +@dataclass +class LTX2VideoConfig(DiTConfig): + """Main configuration for LTX-2 transformer.""" + + arch_config: DiTArchConfig = field(default_factory=LTX2VideoArchConfig) + prefix: str = "ltx2" diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index f943114e8..778b15237 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -8,10 +8,11 @@ from fastvideo.configs.models.encoders.llama import LlamaConfig from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config +from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig __all__ = [ "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", - "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config" + "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig" ] diff --git a/fastvideo/configs/models/encoders/gemma.py b/fastvideo/configs/models/encoders/gemma.py new file mode 100644 index 000000000..73c51dd03 --- /dev/null +++ b/fastvideo/configs/models/encoders/gemma.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.encoders.base import ( + TextEncoderArchConfig, + TextEncoderConfig, +) + + +@dataclass +class LTX2GemmaArchConfig(TextEncoderArchConfig): + architectures: list[str] = field( + default_factory=lambda: ["LTX2GemmaTextEncoderModel"]) + hidden_size: int = 3840 + num_hidden_layers: int = 48 + num_attention_heads: int = 30 + text_len: int = 1024 + pad_token_id: int = 0 + eos_token_id: int = 2 + + gemma_model_path: str = "" + gemma_dtype: str = "bfloat16" + padding_side: str = "left" + + feature_extractor_in_features: int = 3840 * 49 + feature_extractor_out_features: int = 3840 + + connector_num_attention_heads: int = 30 + connector_attention_head_dim: int = 128 + connector_num_layers: int = 2 + connector_positional_embedding_theta: float = 10000.0 + connector_positional_embedding_max_pos: list[int] = field( + default_factory=lambda: [4096]) + connector_rope_type: str = "split" + connector_double_precision_rope: bool = False + connector_num_learnable_registers: int | None = 128 + + def __post_init__(self) -> None: + super().__post_init__() + self.tokenizer_kwargs["padding"] = "max_length" + + +@dataclass +class LTX2GemmaConfig(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field( + default_factory=LTX2GemmaArchConfig) + + prefix: str = "ltx2_gemma" diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index ad1b2f206..2e6cc023f 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -2,6 +2,7 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig +from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig from fastvideo.configs.models.vaes.wanvae import WanVAEConfig @@ -12,4 +13,5 @@ __all__ = [ "CosmosVAEConfig", "Cosmos25VAEConfig", "Hunyuan15VAEConfig", + "LTX2VAEConfig", ] diff --git a/fastvideo/configs/models/vaes/ltx2vae.py b/fastvideo/configs/models/vaes/ltx2vae.py new file mode 100644 index 000000000..d659b9d2a --- /dev/null +++ b/fastvideo/configs/models/vaes/ltx2vae.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 VAE configuration. +""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class LTX2VAEArchConfig(VAEArchConfig): + # Mirrors LTX-2 safetensors metadata config under "vae" + _class_name: str = "CausalVideoAutoencoder" + dims: int = 3 + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 128 + encoder_blocks: list = field(default_factory=list) + decoder_blocks: list = field(default_factory=list) + patch_size: int = 4 + norm_layer: str = "pixel_norm" + latent_log_var: str = "uniform" + encoder_spatial_padding_mode: str = "zeros" + decoder_spatial_padding_mode: str = "reflect" + causal_decoder: bool = False + timestep_conditioning: bool = True + use_quant_conv: bool = False + scaling_factor: float = 1.0 + normalize_latent_channels: bool = False + + # Match FastVideo naming for compression ratios (LTX-2 default) + temporal_compression_ratio: int = 8 + spatial_compression_ratio: int = 32 + + +@dataclass +class LTX2VAEConfig(VAEConfig): + arch_config: VAEArchConfig = field(default_factory=LTX2VAEArchConfig) + + # LTX-2 tiling defaults (match ltx_core.video_vae.TilingConfig.default()). + ltx2_spatial_tile_size_in_pixels: int = 512 + ltx2_spatial_tile_overlap_in_pixels: int = 64 + ltx2_temporal_tile_size_in_frames: int = 64 + ltx2_temporal_tile_overlap_in_frames: int = 24 diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 8120ed53e..296ffa8a6 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -4,6 +4,7 @@ from fastvideo.configs.pipelines.cosmos import CosmosConfig from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig +from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig from fastvideo.configs.pipelines.registry import ( get_pipeline_config_cls_from_name) from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig @@ -16,5 +17,6 @@ __all__ = [ "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig", - "CosmosConfig", "Cosmos25Config", "get_pipeline_config_cls_from_name" + "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", + "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/ltx2.py b/fastvideo/configs/pipelines/ltx2.py new file mode 100644 index 000000000..e08952882 --- /dev/null +++ b/fastvideo/configs/pipelines/ltx2.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig, + LTX2AudioDecoderConfig, LTX2VocoderConfig, + VAEConfig) +from fastvideo.configs.models.dits import LTX2VideoConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput, LTX2GemmaConfig +from fastvideo.configs.models.vaes import LTX2VAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text + + +def ltx2_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: + return outputs.last_hidden_state + + +@dataclass +class LTX2T2VConfig(PipelineConfig): + """Configuration for LTX-2 T2V pipeline.""" + + dit_config: DiTConfig = field(default_factory=LTX2VideoConfig) + vae_config: VAEConfig = field(default_factory=LTX2VAEConfig) + vae_tiling: bool = True + vae_sp: bool = False + + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (LTX2GemmaConfig(), )) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (preprocess_text, )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], + ...] = field(default_factory=lambda: + (ltx2_postprocess_text, )) + + dit_precision: str = "bf16" + vae_precision: str = "bf16" + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("bf16", )) + + audio_decoder_config: ModelConfig = field( + default_factory=LTX2AudioDecoderConfig) + vocoder_config: ModelConfig = field(default_factory=LTX2VocoderConfig) + audio_decoder_precision: str = "bf16" + vocoder_precision: str = "bf16" + + def __post_init__(self) -> None: + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index 0b4c148df..baade6da1 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -9,6 +9,7 @@ from fastvideo.configs.pipelines.cosmos import CosmosConfig from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig +from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig from fastvideo.configs.pipelines.turbodiffusion import ( @@ -64,6 +65,9 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "FastVideo/LongCat-Video-T2V-Diffusers": LongCatT2V480PConfig, "FastVideo/LongCat-Video-I2V-Diffusers": LongCatT2V480PConfig, "FastVideo/LongCat-Video-VC-Diffusers": LongCatT2V480PConfig, + # LTX-2 models + "Lightricks/LTX-2": LTX2T2VConfig, + "converted/ltx2_diffusers": LTX2T2VConfig, # TurboDiffusion models "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers": TurboDiffusionT2V_1_3B_Config, "loayrashid/TurboWan2.1-T2V-14B-Diffusers": TurboDiffusionT2V_14B_Config, @@ -102,6 +106,8 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { lambda id: "cosmos25" in id.lower(), "turbodiffusion": lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(), + "ltx2": + lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(), # Add other pipeline architecture detectors } @@ -123,6 +129,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { "wancausaldmdpipeline": SelfForcingWanT2V480PConfig, "stepvideo": StepVideoT2VConfig, "turbodiffusion": TurboDiffusionT2V_1_3B_Config, + "ltx2": LTX2T2VConfig, # Other fallbacks by architecture } diff --git a/fastvideo/configs/sample/ltx2.py b/fastvideo/configs/sample/ltx2.py new file mode 100644 index 000000000..4be8d40bd --- /dev/null +++ b/fastvideo/configs/sample/ltx2.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass + +from fastvideo.configs.sample.base import SamplingParam + + +@dataclass +class LTX2SamplingParam(SamplingParam): + """Default sampling parameters for LTX-2 distilled T2V. + """ + + seed: int = 10 + num_frames: int = 121 + height: int = 1024 + width: int = 1536 + fps: int = 24 + num_inference_steps: int = 8 + guidance_scale: float = 1.0 + # No default negative_prompt for distilled models + negative_prompt: str = "" diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 89673bed8..6b26133d2 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -10,6 +10,7 @@ from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam from fastvideo.configs.sample.cosmos2_5 import Cosmos_Predict2_5_2B_Diffusers_SamplingParam +from fastvideo.configs.sample.ltx2 import LTX2SamplingParam # isort: off from fastvideo.configs.sample.wan import ( @@ -40,48 +41,36 @@ from fastvideo.utils import (maybe_download_model_index, logger = init_logger(__name__) # Registry maps specific model weights to their config classes SAMPLING_PARAM_REGISTRY: dict[str, Any] = { - "FastVideo/FastHunyuan-diffusers": - FastHunyuanSamplingParam, - "hunyuanvideo-community/HunyuanVideo": - HunyuanSamplingParam, + "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam, + "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": Hunyuan15_480P_SamplingParam, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v": Hunyuan15_720P_SamplingParam, - "FastVideo/stepvideo-t2v-diffusers": - StepVideoT2VSamplingParam, + "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, # Wan2.1 - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": - WanT2V_1_3B_SamplingParam, - "Wan-AI/Wan2.1-T2V-14B-Diffusers": - WanT2V_14B_SamplingParam, - "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": - WanI2V_14B_480P_SamplingParam, - "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": - WanI2V_14B_720P_SamplingParam, + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam, + "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam, "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": Wan2_1_Fun_1_3B_InP_SamplingParam, "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": Wan2_1_Fun_1_3B_Control_SamplingParam, # Wan2.2 - "Wan-AI/Wan2.2-TI2V-5B-Diffusers": - Wan2_2_TI2V_5B_SamplingParam, + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers": Wan2_2_TI2V_5B_SamplingParam, - "Wan-AI/Wan2.2-T2V-A14B-Diffusers": - Wan2_2_T2V_A14B_SamplingParam, - "Wan-AI/Wan2.2-I2V-A14B-Diffusers": - Wan2_2_I2V_A14B_SamplingParam, + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, + "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, # FastWan2.1 - "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": - FastWanT2V480P_SamplingParam, + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480P_SamplingParam, # FastWan2.2 - "FastVideo/FastWan2.2-TI2V-5B-Diffusers": - Wan2_2_TI2V_5B_SamplingParam, + "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, # Causal Self-Forcing Wan2.1 "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": @@ -102,12 +91,9 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { Cosmos_Predict2_5_2B_Diffusers_SamplingParam, # MatrixGame2.0 models - "FastVideo/Matrix-Game-2.0-Base-Diffusers": - MatrixGame2_SamplingParam, - "FastVideo/Matrix-Game-2.0-GTA-Diffusers": - MatrixGame2_SamplingParam, - "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": - MatrixGame2_SamplingParam, + "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam, + "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam, + "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGame2_SamplingParam, # TurboDiffusion models "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers": @@ -117,6 +103,10 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "loayrashid/TurboWan2.2-I2V-A14B-Diffusers": TurboDiffusionI2V_A14B_SamplingParam, + # LTX-2 models + "Lightricks/LTX-2": LTX2SamplingParam, + "FastVideo/LTX2-Distilled-Diffusers": LTX2SamplingParam, + # Add other specific weight variants } @@ -144,6 +134,8 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = { lambda id: "cosmos2_5" in id.lower(), "cosmos": lambda id: "cosmos" in id.lower() and "2_5" not in id.lower(), + "ltx2": + lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(), # Add other pipeline architecture detectors } @@ -164,6 +156,7 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = { TurboDiffusionT2V_1_3B_SamplingParam, # Default to T2V for fallback "cosmos25": Cosmos_Predict2_5_2B_Diffusers_SamplingParam, "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam, + "ltx2": LTX2SamplingParam, # Other fallbacks by architecture } diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index d14b7a73d..d586d9215 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -18,6 +18,8 @@ import numpy as np import torch import torchvision from einops import rearrange +import shutil +import tempfile from fastvideo.configs.sample import SamplingParam from fastvideo.fastvideo_args import FastVideoArgs @@ -389,6 +391,11 @@ class VideoGenerator: if batch.save_video: imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") logger.info("Saved video to %s", output_path) + audio = output_batch.extra.get("audio") + audio_sample_rate = output_batch.extra.get("audio_sample_rate") + if (audio is not None and audio_sample_rate is not None and + not self._mux_audio(output_path, audio, audio_sample_rate)): + logger.warning("Audio mux failed; saved video without audio.") if batch.return_frames: return frames @@ -396,6 +403,7 @@ class VideoGenerator: return { "samples": samples, "frames": frames, + "audio": output_batch.extra.get("audio"), "prompts": prompt, "size": (target_height, target_width, batch.num_frames), "generation_time": gen_time, @@ -405,6 +413,98 @@ class VideoGenerator: "trajectory_decoded": output_batch.trajectory_decoded, } + @staticmethod + def _mux_audio( + video_path: str, + audio: torch.Tensor | np.ndarray, + sample_rate: int, + ) -> bool: + """Mux audio into video using PyAV.""" + try: + import av + except ImportError: + logger.warning("PyAV not installed; cannot mux audio. " + "Install with: pip install av") + return False + + if torch.is_tensor(audio): + audio_np = audio.detach().cpu().float().numpy() + else: + audio_np = np.asarray(audio, dtype=np.float32) + + if audio_np.ndim == 1: + audio_np = audio_np[:, None] + elif audio_np.ndim == 2: + if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]: + audio_np = audio_np.T + else: + logger.warning("Unexpected audio shape %s; skipping mux.", + audio_np.shape) + return False + + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_int16 = (audio_np * 32767.0).astype(np.int16) + num_channels = audio_int16.shape[1] + layout = "stereo" if num_channels == 2 else "mono" + + try: + import wave + with tempfile.TemporaryDirectory() as tmpdir: + out_path = os.path.join(tmpdir, "muxed.mp4") + wav_path = os.path.join(tmpdir, "audio.wav") + + # Write audio to WAV file + with wave.open(wav_path, "wb") as wav_file: + wav_file.setnchannels(num_channels) + wav_file.setsampwidth(2) + wav_file.setframerate(sample_rate) + wav_file.writeframes(audio_int16.tobytes()) + + # Open input video and audio + input_video = av.open(video_path) + input_audio = av.open(wav_path) + + # Create output with both streams + output = av.open(out_path, mode="w") + + # Add video stream (copy codec from input) + in_video_stream = input_video.streams.video[0] + out_video_stream = output.add_stream( + codec_name=in_video_stream.codec_context.name, + rate=in_video_stream.average_rate, + ) + out_video_stream.width = in_video_stream.width + out_video_stream.height = in_video_stream.height + out_video_stream.pix_fmt = in_video_stream.pix_fmt + + # Add audio stream (AAC) + out_audio_stream = output.add_stream("aac", rate=sample_rate) + out_audio_stream.layout = layout + + # Remux video (decode and re-encode to be safe) + for frame in input_video.decode(video=0): + for packet in out_video_stream.encode(frame): + output.mux(packet) + for packet in out_video_stream.encode(): + output.mux(packet) + + # Encode audio + for frame in input_audio.decode(audio=0): + frame.pts = None # Let encoder assign PTS + for packet in out_audio_stream.encode(frame): + output.mux(packet) + for packet in out_audio_stream.encode(): + output.mux(packet) + + input_video.close() + input_audio.close() + output.close() + shutil.move(out_path, video_path) + return True + except Exception as e: + logger.warning("Audio mux failed: %s", e) + return False + def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None) -> None: diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index c2afc9f3a..33c06efcc 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -166,6 +166,14 @@ class FastVideoArgs: # Prompt text file for batch processing prompt_txt: str | None = None + # LTX-2 VAE tiling overrides + ltx2_vae_tiling: bool | None = None + ltx2_vae_spatial_tile_size_in_pixels: int | None = None + ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None + ltx2_vae_temporal_tile_size_in_frames: int | None = None + ltx2_vae_temporal_tile_overlap_in_frames: int | None = None + ltx2_initial_latent_path: str | None = None + # model paths for correct deallocation model_paths: dict[str, str] = field(default_factory=dict) model_loaded: dict[str, bool] = field(default_factory=lambda: { @@ -203,8 +211,44 @@ class FastVideoArgs: logger.error("Failed to load V-MoBA config from %s: %s", self.moba_config_path, e) raise + self._apply_ltx2_vae_overrides() self.check_fastvideo_args() + def _apply_ltx2_vae_overrides(self) -> None: + if self.pipeline_config is None: + return + vae_config = self.pipeline_config.vae_config + has_any = any(value is not None for value in ( + self.ltx2_vae_spatial_tile_size_in_pixels, + self.ltx2_vae_spatial_tile_overlap_in_pixels, + self.ltx2_vae_temporal_tile_size_in_frames, + self.ltx2_vae_temporal_tile_overlap_in_frames, + )) + if self.ltx2_vae_tiling is not None and hasattr(self.pipeline_config, + "vae_tiling"): + self.pipeline_config.vae_tiling = self.ltx2_vae_tiling + elif has_any and hasattr(self.pipeline_config, "vae_tiling"): + self.pipeline_config.vae_tiling = True + + if hasattr(vae_config, "ltx2_spatial_tile_size_in_pixels" + ) and self.ltx2_vae_spatial_tile_size_in_pixels is not None: + vae_config.ltx2_spatial_tile_size_in_pixels = ( + self.ltx2_vae_spatial_tile_size_in_pixels) + if hasattr( + vae_config, "ltx2_spatial_tile_overlap_in_pixels" + ) and self.ltx2_vae_spatial_tile_overlap_in_pixels is not None: + vae_config.ltx2_spatial_tile_overlap_in_pixels = ( + self.ltx2_vae_spatial_tile_overlap_in_pixels) + if hasattr(vae_config, "ltx2_temporal_tile_size_in_frames" + ) and self.ltx2_vae_temporal_tile_size_in_frames is not None: + vae_config.ltx2_temporal_tile_size_in_frames = ( + self.ltx2_vae_temporal_tile_size_in_frames) + if hasattr( + vae_config, "ltx2_temporal_tile_overlap_in_frames" + ) and self.ltx2_vae_temporal_tile_overlap_in_frames is not None: + vae_config.ltx2_temporal_tile_overlap_in_frames = ( + self.ltx2_vae_temporal_tile_overlap_in_frames) + @staticmethod def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: # Model and path configuration @@ -325,6 +369,44 @@ class FastVideoArgs: "Path to a text file containing prompts (one per line) for batch processing", ) + # LTX-2 VAE tiling overrides + parser.add_argument( + "--ltx2-vae-tiling", + action=StoreBoolean, + default=FastVideoArgs.ltx2_vae_tiling, + help="Enable LTX-2 VAE tiling overrides.", + ) + parser.add_argument( + "--ltx2-vae-spatial-tile-size-in-pixels", + type=int, + default=FastVideoArgs.ltx2_vae_spatial_tile_size_in_pixels, + help="LTX-2 VAE spatial tile size in pixels.", + ) + parser.add_argument( + "--ltx2-vae-spatial-tile-overlap-in-pixels", + type=int, + default=FastVideoArgs.ltx2_vae_spatial_tile_overlap_in_pixels, + help="LTX-2 VAE spatial tile overlap in pixels.", + ) + parser.add_argument( + "--ltx2-vae-temporal-tile-size-in-frames", + type=int, + default=FastVideoArgs.ltx2_vae_temporal_tile_size_in_frames, + help="LTX-2 VAE temporal tile size in frames.", + ) + parser.add_argument( + "--ltx2-vae-temporal-tile-overlap-in-frames", + type=int, + default=FastVideoArgs.ltx2_vae_temporal_tile_overlap_in_frames, + help="LTX-2 VAE temporal tile overlap in frames.", + ) + parser.add_argument( + "--ltx2-initial-latent-path", + type=str, + default=FastVideoArgs.ltx2_initial_latent_path, + help="Path to load/save a precomputed LTX-2 initial latent.", + ) + # LoRA parameters (inference-time adapter loading) parser.add_argument( "--lora-path", diff --git a/fastvideo/models/audio/__init__.py b/fastvideo/models/audio/__init__.py new file mode 100644 index 000000000..9ba265f8e --- /dev/null +++ b/fastvideo/models/audio/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.models.audio.ltx2_audio_vae import ( + LTX2AudioDecoder, + LTX2AudioEncoder, + LTX2Vocoder, +) + +__all__ = ["LTX2AudioEncoder", "LTX2AudioDecoder", "LTX2Vocoder"] diff --git a/fastvideo/models/audio/ltx2_audio_vae.py b/fastvideo/models/audio/ltx2_audio_vae.py new file mode 100644 index 000000000..fcc047f6c --- /dev/null +++ b/fastvideo/models/audio/ltx2_audio_vae.py @@ -0,0 +1,1327 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Native LTX-2 Audio VAE and Vocoder implementation for FastVideo. +""" + +from __future__ import annotations + +import math +from enum import Enum +from typing import Any, List, NamedTuple, Set, Tuple + +import einops +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# ============================================================================= +# Enums +# ============================================================================= + + +class NormType(Enum): + """Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm).""" + GROUP = "group" + PIXEL = "pixel" + + +class AttentionType(Enum): + """Enum for specifying the attention mechanism type.""" + VANILLA = "vanilla" + LINEAR = "linear" + NONE = "none" + + +class CausalityAxis(Enum): + """Enum for specifying the causality axis in causal convolutions.""" + NONE = None + WIDTH = "width" + HEIGHT = "height" + WIDTH_COMPATIBILITY = "width-compatibility" + + +# ============================================================================= +# Audio Latent Shape +# ============================================================================= + + +class AudioLatentShape(NamedTuple): + """Shape of audio in VAE latent space: (batch, channels, frames, mel_bins).""" + batch: int + channels: int + frames: int + mel_bins: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.mel_bins]) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> AudioLatentShape: + return AudioLatentShape( + batch=shape[0], + channels=shape[1], + frames=shape[2], + mel_bins=shape[3], + ) + + +# ============================================================================= +# Constants +# ============================================================================= + + +LATENT_DOWNSAMPLE_FACTOR = 4 +LRELU_SLOPE = 0.1 + + +# ============================================================================= +# Normalization Layers +# ============================================================================= + + +class PixelNorm(nn.Module): + """ + Per-pixel (per-location) RMS normalization layer. + """ + + def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + mean_sq = torch.mean(x ** 2, dim=self.dim, keepdim=True) + rms = torch.sqrt(mean_sq + self.eps) + return x / rms + + +def build_normalization_layer( + in_channels: int, + *, + num_groups: int = 32, + normtype: NormType = NormType.GROUP, +) -> nn.Module: + """Create a normalization layer based on the normalization type.""" + if normtype == NormType.GROUP: + return nn.GroupNorm( + num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True + ) + if normtype == NormType.PIXEL: + return PixelNorm(dim=1, eps=1e-6) + raise ValueError(f"Invalid normalization type: {normtype}") + + +# ============================================================================= +# Per-Channel Statistics +# ============================================================================= + + +class PerChannelStatistics(nn.Module): + """ + Per-channel statistics for normalizing and denormalizing the latent representation. + """ + + def __init__(self, latent_channels: int = 128) -> None: + super().__init__() + self.register_buffer("std-of-means", torch.empty(latent_channels)) + self.register_buffer("mean-of-means", torch.empty(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer( + "mean-of-means" + ).to(x) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer( + "std-of-means" + ).to(x) + + +# ============================================================================= +# Audio Patchifier +# ============================================================================= + + +class AudioPatchifier: + """Simple patchifier for audio latents.""" + + def __init__( + self, + patch_size: int = 1, + sample_rate: int = 16000, + hop_length: int = 160, + audio_latent_downsample_factor: int = 4, + is_causal: bool = True, + ): + self.patch_size = patch_size + self.sample_rate = sample_rate + self.hop_length = hop_length + self.audio_latent_downsample_factor = audio_latent_downsample_factor + self.is_causal = is_causal + + def patchify(self, audio_latents: torch.Tensor) -> torch.Tensor: + """Flatten audio latent tensor along time: (B, C, T, F) -> (B, T, C*F).""" + return einops.rearrange(audio_latents, "b c t f -> b t (c f)") + + def unpatchify( + self, audio_latents: torch.Tensor, output_shape: AudioLatentShape + ) -> torch.Tensor: + """Restore (B, C, T, F) from flattened patches: (B, T, C*F) -> (B, C, T, F).""" + return einops.rearrange( + audio_latents, + "b t (c f) -> b c t f", + c=output_shape.channels, + f=output_shape.mel_bins, + ) + + +# ============================================================================= +# Causal 2D Convolution +# ============================================================================= + + +class CausalConv2d(nn.Module): + """ + A causal 2D convolution. + Ensures output at time t only depends on inputs at time t and earlier. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | Tuple[int, int], + stride: int = 1, + dilation: int | Tuple[int, int] = 1, + groups: int = 1, + bias: bool = True, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + ) -> None: + super().__init__() + + self.causality_axis = causality_axis + + # Ensure kernel_size and dilation are tuples + kernel_size = nn.modules.utils._pair(kernel_size) + dilation = nn.modules.utils._pair(dilation) + + # Calculate padding dimensions + pad_h = (kernel_size[0] - 1) * dilation[0] + pad_w = (kernel_size[1] - 1) * dilation[1] + + # The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom) + if self.causality_axis == CausalityAxis.NONE: + self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2) + elif self.causality_axis in (CausalityAxis.WIDTH, CausalityAxis.WIDTH_COMPATIBILITY): + self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2) + elif self.causality_axis == CausalityAxis.HEIGHT: + self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0) + else: + raise ValueError(f"Invalid causality_axis: {causality_axis}") + + # The internal convolution layer uses no padding, as we handle it manually + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, self.padding) + return self.conv(x) + + +def make_conv2d( + in_channels: int, + out_channels: int, + kernel_size: int | Tuple[int, int], + stride: int = 1, + padding: Tuple[int, int, int, int] | None = None, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causality_axis: CausalityAxis | None = None, +) -> nn.Module: + """Create a 2D convolution layer that can be either causal or non-causal.""" + if causality_axis is not None: + return CausalConv2d( + in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis + ) + else: + if padding is None: + padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size) + return nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + groups, + bias, + ) + + +# ============================================================================= +# Attention Blocks +# ============================================================================= + + +class AttnBlock(nn.Module): + """Vanilla self-attention block for 2D features.""" + + def __init__( + self, + in_channels: int, + norm_type: NormType = NormType.GROUP, + ) -> None: + super().__init__() + self.in_channels = in_channels + + self.norm = build_normalization_layer(in_channels, normtype=norm_type) + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # Compute attention + b, c, h, w = q.shape + q = q.reshape(b, c, h * w).contiguous() + q = q.permute(0, 2, 1).contiguous() # b, hw, c + k = k.reshape(b, c, h * w).contiguous() # b, c, hw + w_ = torch.bmm(q, k).contiguous() # b, hw, hw + w_ = w_ * (int(c) ** (-0.5)) + w_ = F.softmax(w_, dim=2) + + # Attend to values + v = v.reshape(b, c, h * w).contiguous() + w_ = w_.permute(0, 2, 1).contiguous() # b, hw, hw + h_ = torch.bmm(v, w_).contiguous() # b, c, hw + h_ = h_.reshape(b, c, h, w).contiguous() + + h_ = self.proj_out(h_) + + return x + h_ + + +def make_attn( + in_channels: int, + attn_type: AttentionType = AttentionType.VANILLA, + norm_type: NormType = NormType.GROUP, +) -> nn.Module: + """Factory function for attention blocks.""" + if attn_type == AttentionType.VANILLA: + return AttnBlock(in_channels, norm_type=norm_type) + elif attn_type == AttentionType.NONE: + return nn.Identity() + elif attn_type == AttentionType.LINEAR: + raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.") + else: + raise ValueError(f"Unknown attention type: {attn_type}") + + +# ============================================================================= +# ResNet Blocks (2D for Audio VAE) +# ============================================================================= + + +class ResnetBlock(nn.Module): + """2D ResNet block for audio VAE.""" + + def __init__( + self, + *, + in_channels: int, + out_channels: int | None = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + norm_type: NormType = NormType.GROUP, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + ) -> None: + super().__init__() + self.causality_axis = causality_axis + + if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP: + raise ValueError("Causal ResnetBlock with GroupNorm is not supported.") + + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = build_normalization_layer(in_channels, normtype=norm_type) + self.non_linearity = nn.SiLU() + self.conv1 = make_conv2d( + in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + if temb_channels > 0: + self.temb_proj = nn.Linear(temb_channels, out_channels) + self.norm2 = build_normalization_layer(out_channels, normtype=norm_type) + self.dropout = nn.Dropout(dropout) + self.conv2 = make_conv2d( + out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = make_conv2d( + in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + else: + self.nin_shortcut = make_conv2d( + in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis + ) + + def forward( + self, + x: torch.Tensor, + temb: torch.Tensor | None = None, + ) -> torch.Tensor: + h = x + h = self.norm1(h) + h = self.non_linearity(h) + h = self.conv1(h) + + if temb is not None: + h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None] + + h = self.norm2(h) + h = self.non_linearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x) + + return x + h + + +# ============================================================================= +# Vocoder ResBlocks (1D) +# ============================================================================= + + +class ResBlock1(nn.Module): + """1D ResBlock for vocoder with dilated convolutions.""" + + def __init__( + self, + channels: int, + kernel_size: int = 3, + dilation: Tuple[int, int, int] = (1, 3, 5), + ): + super().__init__() + self.convs1 = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding="same"), + ] + ) + self.convs2 = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv1, conv2 in zip(self.convs1, self.convs2, strict=True): + xt = F.leaky_relu(x, LRELU_SLOPE) + xt = conv1(xt) + xt = F.leaky_relu(xt, LRELU_SLOPE) + xt = conv2(xt) + x = xt + x + return x + + +class ResBlock2(nn.Module): + """1D ResBlock for vocoder (simpler version).""" + + def __init__( + self, + channels: int, + kernel_size: int = 3, + dilation: Tuple[int, int] = (1, 3), + ): + super().__init__() + self.convs = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv in self.convs: + xt = F.leaky_relu(x, LRELU_SLOPE) + xt = conv(xt) + x = xt + x + return x + + +# ============================================================================= +# Downsampling +# ============================================================================= + + +class Downsample(nn.Module): + """Downsampling layer with strided convolution or average pooling.""" + + def __init__( + self, + in_channels: int, + with_conv: bool, + causality_axis: CausalityAxis = CausalityAxis.WIDTH, + ) -> None: + super().__init__() + self.with_conv = with_conv + self.causality_axis = causality_axis + + if self.causality_axis != CausalityAxis.NONE and not self.with_conv: + raise ValueError("causality is only supported when `with_conv=True`.") + + if self.with_conv: + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.with_conv: + # Padding tuple is in the order: (left, right, top, bottom) + if self.causality_axis == CausalityAxis.NONE: + pad = (0, 1, 0, 1) + elif self.causality_axis == CausalityAxis.WIDTH: + pad = (2, 0, 0, 1) + elif self.causality_axis == CausalityAxis.HEIGHT: + pad = (0, 1, 2, 0) + elif self.causality_axis == CausalityAxis.WIDTH_COMPATIBILITY: + pad = (1, 0, 0, 1) + else: + raise ValueError(f"Invalid causality_axis: {self.causality_axis}") + + x = F.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + x = F.avg_pool2d(x, kernel_size=2, stride=2) + + return x + + +def build_downsampling_path( + *, + ch: int, + ch_mult: Tuple[int, ...], + num_resolutions: int, + num_res_blocks: int, + resolution: int, + temb_channels: int, + dropout: float, + norm_type: NormType, + causality_axis: CausalityAxis, + attn_type: AttentionType, + attn_resolutions: Set[int], + resamp_with_conv: bool, +) -> Tuple[nn.ModuleList, int]: + """Build the downsampling path with residual blocks, attention, and downsampling layers.""" + down_modules = nn.ModuleList() + curr_res = resolution + in_ch_mult = (1, *tuple(ch_mult)) + block_in = ch + + for i_level in range(num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + + for _ in range(num_res_blocks): + block.append( + ResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=temb_channels, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) + + down = nn.Module() + down.block = block + down.attn = attn + if i_level != num_resolutions - 1: + down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis) + curr_res = curr_res // 2 + down_modules.append(down) + + return down_modules, block_in + + +# ============================================================================= +# Upsampling +# ============================================================================= + + +class Upsample(nn.Module): + """Upsampling layer with nearest-neighbor interpolation and optional convolution.""" + + def __init__( + self, + in_channels: int, + with_conv: bool, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + ) -> None: + super().__init__() + self.with_conv = with_conv + self.causality_axis = causality_axis + if self.with_conv: + self.conv = make_conv2d( + in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + # Drop FIRST element in the causal axis to undo encoder's padding + if self.causality_axis == CausalityAxis.NONE: + pass + elif self.causality_axis == CausalityAxis.HEIGHT: + x = x[:, :, 1:, :] + elif self.causality_axis == CausalityAxis.WIDTH: + x = x[:, :, :, 1:] + elif self.causality_axis == CausalityAxis.WIDTH_COMPATIBILITY: + pass + else: + raise ValueError(f"Invalid causality_axis: {self.causality_axis}") + + return x + + +def build_upsampling_path( + *, + ch: int, + ch_mult: Tuple[int, ...], + num_resolutions: int, + num_res_blocks: int, + resolution: int, + temb_channels: int, + dropout: float, + norm_type: NormType, + causality_axis: CausalityAxis, + attn_type: AttentionType, + attn_resolutions: Set[int], + resamp_with_conv: bool, + initial_block_channels: int, +) -> Tuple[nn.ModuleList, int]: + """Build the upsampling path with residual blocks, attention, and upsampling layers.""" + up_modules = nn.ModuleList() + block_in = initial_block_channels + curr_res = resolution // (2 ** (num_resolutions - 1)) + + for level in reversed(range(num_resolutions)): + stage = nn.Module() + stage.block = nn.ModuleList() + stage.attn = nn.ModuleList() + block_out = ch * ch_mult[level] + + for _ in range(num_res_blocks + 1): + stage.block.append( + ResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=temb_channels, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) + + if level != 0: + stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis) + curr_res *= 2 + + up_modules.insert(0, stage) + + return up_modules, block_in + + +# ============================================================================= +# Mid Block +# ============================================================================= + + +def build_mid_block( + channels: int, + temb_channels: int, + dropout: float, + norm_type: NormType, + causality_axis: CausalityAxis, + attn_type: AttentionType, + add_attention: bool, +) -> nn.Module: + """Build the middle block with two ResNet blocks and optional attention.""" + mid = nn.Module() + mid.block_1 = ResnetBlock( + in_channels=channels, + out_channels=channels, + temb_channels=temb_channels, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + mid.attn_1 = ( + make_attn(channels, attn_type=attn_type, norm_type=norm_type) + if add_attention + else nn.Identity() + ) + mid.block_2 = ResnetBlock( + in_channels=channels, + out_channels=channels, + temb_channels=temb_channels, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + return mid + + +def run_mid_block(mid: nn.Module, features: torch.Tensor) -> torch.Tensor: + """Run features through the middle block.""" + features = mid.block_1(features, temb=None) + features = mid.attn_1(features) + return mid.block_2(features, temb=None) + + +# ============================================================================= +# Audio Encoder +# ============================================================================= + + +class AudioEncoder(nn.Module): + """ + Encoder that compresses audio spectrograms into latent representations. + """ + + def __init__( + self, + *, + ch: int, + ch_mult: Tuple[int, ...] = (1, 2, 4, 8), + num_res_blocks: int, + attn_resolutions: Set[int], + dropout: float = 0.0, + resamp_with_conv: bool = True, + in_channels: int, + resolution: int, + z_channels: int, + double_z: bool = True, + attn_type: AttentionType = AttentionType.VANILLA, + mid_block_add_attention: bool = True, + norm_type: NormType = NormType.GROUP, + causality_axis: CausalityAxis = CausalityAxis.WIDTH, + sample_rate: int = 16000, + mel_hop_length: int = 160, + n_fft: int = 1024, + is_causal: bool = True, + mel_bins: int = 64, + **_ignore_kwargs, + ) -> None: + super().__init__() + + self.per_channel_statistics = PerChannelStatistics(latent_channels=ch) + self.sample_rate = sample_rate + self.mel_hop_length = mel_hop_length + self.n_fft = n_fft + self.is_causal = is_causal + self.mel_bins = mel_bins + + self.patchifier = AudioPatchifier( + patch_size=1, + audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR, + sample_rate=sample_rate, + hop_length=mel_hop_length, + is_causal=is_causal, + ) + + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.z_channels = z_channels + self.double_z = double_z + self.norm_type = norm_type + self.causality_axis = causality_axis + self.attn_type = attn_type + + # Input convolution + self.conv_in = make_conv2d( + in_channels, + self.ch, + kernel_size=3, + stride=1, + causality_axis=self.causality_axis, + ) + + self.non_linearity = nn.SiLU() + + # Downsampling path + self.down, block_in = build_downsampling_path( + ch=ch, + ch_mult=ch_mult, + num_resolutions=self.num_resolutions, + num_res_blocks=num_res_blocks, + resolution=resolution, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + attn_type=self.attn_type, + attn_resolutions=attn_resolutions, + resamp_with_conv=resamp_with_conv, + ) + + # Mid block + self.mid = build_mid_block( + channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + attn_type=self.attn_type, + add_attention=mid_block_add_attention, + ) + + # Output layers + self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type) + self.conv_out = make_conv2d( + block_in, + 2 * z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + causality_axis=self.causality_axis, + ) + + def forward(self, spectrogram: torch.Tensor) -> torch.Tensor: + """ + Encode audio spectrogram into latent representations. + Args: + spectrogram: Input spectrogram of shape (batch, channels, time, frequency) + Returns: + Encoded latent representation of shape (batch, channels, frames, mel_bins) + """ + h = self.conv_in(spectrogram) + h = self._run_downsampling_path(h) + h = run_mid_block(self.mid, h) + h = self._finalize_output(h) + + return self._normalize_latents(h) + + def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor: + for level in range(self.num_resolutions): + stage = self.down[level] + for block_idx in range(self.num_res_blocks): + h = stage.block[block_idx](h, temb=None) + if stage.attn: + h = stage.attn[block_idx](h) + + if level != self.num_resolutions - 1: + h = stage.downsample(h) + + return h + + def _finalize_output(self, h: torch.Tensor) -> torch.Tensor: + h = self.norm_out(h) + h = self.non_linearity(h) + return self.conv_out(h) + + def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor: + """Normalize encoder latents using per-channel statistics.""" + means = torch.chunk(latent_output, 2, dim=1)[0] + latent_shape = AudioLatentShape( + batch=means.shape[0], + channels=means.shape[1], + frames=means.shape[2], + mel_bins=means.shape[3], + ) + latent_patched = self.patchifier.patchify(means) + latent_normalized = self.per_channel_statistics.normalize(latent_patched) + return self.patchifier.unpatchify(latent_normalized, latent_shape) + + +# ============================================================================= +# Audio Decoder +# ============================================================================= + + +class AudioDecoder(nn.Module): + """ + Symmetric decoder that reconstructs audio spectrograms from latent features. + """ + + def __init__( + self, + *, + ch: int, + out_ch: int, + ch_mult: Tuple[int, ...] = (1, 2, 4, 8), + num_res_blocks: int, + attn_resolutions: Set[int], + resolution: int, + z_channels: int, + norm_type: NormType = NormType.GROUP, + causality_axis: CausalityAxis = CausalityAxis.WIDTH, + dropout: float = 0.0, + mid_block_add_attention: bool = True, + sample_rate: int = 16000, + mel_hop_length: int = 160, + is_causal: bool = True, + mel_bins: int | None = None, + ) -> None: + super().__init__() + + # Internal behavioural defaults + resamp_with_conv = True + attn_type = AttentionType.VANILLA + + # Per-channel statistics for denormalizing latents + self.per_channel_statistics = PerChannelStatistics(latent_channels=ch) + self.sample_rate = sample_rate + self.mel_hop_length = mel_hop_length + self.is_causal = is_causal + self.mel_bins = mel_bins + self.patchifier = AudioPatchifier( + patch_size=1, + audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR, + sample_rate=sample_rate, + hop_length=mel_hop_length, + is_causal=is_causal, + ) + + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.out_ch = out_ch + self.give_pre_end = False + self.tanh_out = False + self.norm_type = norm_type + self.z_channels = z_channels + self.channel_multipliers = ch_mult + self.attn_resolutions = attn_resolutions + self.causality_axis = causality_axis + self.attn_type = attn_type + + base_block_channels = ch * self.channel_multipliers[-1] + base_resolution = resolution // (2 ** (self.num_resolutions - 1)) + self.z_shape = (1, z_channels, base_resolution, base_resolution) + + self.conv_in = make_conv2d( + z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis + ) + self.non_linearity = nn.SiLU() + + self.mid = build_mid_block( + channels=base_block_channels, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + attn_type=self.attn_type, + add_attention=mid_block_add_attention, + ) + + self.up, final_block_channels = build_upsampling_path( + ch=ch, + ch_mult=ch_mult, + num_resolutions=self.num_resolutions, + num_res_blocks=num_res_blocks, + resolution=resolution, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + attn_type=self.attn_type, + attn_resolutions=attn_resolutions, + resamp_with_conv=resamp_with_conv, + initial_block_channels=base_block_channels, + ) + + self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type) + self.conv_out = make_conv2d( + final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + """ + Decode latent features back to audio spectrograms. + Args: + sample: Encoded latent representation of shape (batch, channels, frames, mel_bins) + Returns: + Reconstructed audio spectrogram of shape (batch, channels, time, frequency) + """ + sample, target_shape = self._denormalize_latents(sample) + + h = self.conv_in(sample) + h = run_mid_block(self.mid, h) + h = self._run_upsampling_path(h) + h = self._finalize_output(h) + + return self._adjust_output_shape(h, target_shape) + + def _denormalize_latents( + self, sample: torch.Tensor + ) -> Tuple[torch.Tensor, AudioLatentShape]: + latent_shape = AudioLatentShape( + batch=sample.shape[0], + channels=sample.shape[1], + frames=sample.shape[2], + mel_bins=sample.shape[3], + ) + + sample_patched = self.patchifier.patchify(sample) + sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched) + sample = self.patchifier.unpatchify(sample_denormalized, latent_shape) + + target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR + if self.causality_axis != CausalityAxis.NONE: + target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1) + + target_shape = AudioLatentShape( + batch=latent_shape.batch, + channels=self.out_ch, + frames=target_frames, + mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins, + ) + + return sample, target_shape + + def _adjust_output_shape( + self, + decoded_output: torch.Tensor, + target_shape: AudioLatentShape, + ) -> torch.Tensor: + """Adjust output shape to match target dimensions for variable-length audio.""" + _, _, current_time, current_freq = decoded_output.shape + target_channels = target_shape.channels + target_time = target_shape.frames + target_freq = target_shape.mel_bins + + # Crop first + decoded_output = decoded_output[ + :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq) + ] + + # Calculate padding needed + time_padding_needed = target_time - decoded_output.shape[2] + freq_padding_needed = target_freq - decoded_output.shape[3] + + # Apply padding if needed + if time_padding_needed > 0 or freq_padding_needed > 0: + padding = ( + 0, + max(freq_padding_needed, 0), + 0, + max(time_padding_needed, 0), + ) + decoded_output = F.pad(decoded_output, padding) + + # Final safety crop + decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq] + + return decoded_output + + def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor: + for level in reversed(range(self.num_resolutions)): + stage = self.up[level] + for block_idx, block in enumerate(stage.block): + h = block(h, temb=None) + if stage.attn: + h = stage.attn[block_idx](h) + + if level != 0 and hasattr(stage, "upsample"): + h = stage.upsample(h) + + return h + + def _finalize_output(self, h: torch.Tensor) -> torch.Tensor: + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = self.non_linearity(h) + h = self.conv_out(h) + return torch.tanh(h) if self.tanh_out else h + + +# ============================================================================= +# Vocoder +# ============================================================================= + + +class Vocoder(nn.Module): + """ + Vocoder model for synthesizing audio from Mel spectrograms. + """ + + def __init__( + self, + resblock_kernel_sizes: List[int] | None = None, + upsample_rates: List[int] | None = None, + upsample_kernel_sizes: List[int] | None = None, + resblock_dilation_sizes: List[List[int]] | None = None, + upsample_initial_channel: int = 1024, + stereo: bool = True, + resblock: str = "1", + output_sample_rate: int = 24000, + ): + super().__init__() + + # Initialize default values + if resblock_kernel_sizes is None: + resblock_kernel_sizes = [3, 7, 11] + if upsample_rates is None: + upsample_rates = [6, 5, 2, 2, 2] + if upsample_kernel_sizes is None: + upsample_kernel_sizes = [16, 15, 8, 4, 4] + if resblock_dilation_sizes is None: + resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]] + + self.output_sample_rate = output_sample_rate + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + in_channels = 128 if stereo else 64 + self.conv_pre = nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3) + resblock_class = ResBlock1 if resblock == "1" else ResBlock2 + + self.ups = nn.ModuleList() + for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)): + self.ups.append( + nn.ConvTranspose1d( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + kernel_size, + stride, + padding=(kernel_size - stride) // 2, + ) + ) + + self.resblocks = nn.ModuleList() + for i, _ in enumerate(self.ups): + ch = upsample_initial_channel // (2 ** (i + 1)) + for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): + self.resblocks.append(resblock_class(ch, kernel_size, dilations)) + + out_channels = 2 if stereo else 1 + final_channels = upsample_initial_channel // (2**self.num_upsamples) + self.conv_post = nn.Conv1d(final_channels, out_channels, 7, 1, padding=3) + + self.upsample_factor = math.prod(layer.stride[0] for layer in self.ups) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the vocoder. + Args: + x: Input Mel spectrogram tensor of shape (batch, channels, time, mel_bins) + Returns: + Audio waveform tensor of shape (batch, out_channels, audio_length) + """ + x = x.transpose(2, 3) # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time) + + if x.dim() == 4: # stereo + assert x.shape[1] == 2, "Input must have 2 channels for stereo" + x = einops.rearrange(x, "b s c t -> b (s c) t") + + x = self.conv_pre(x) + + for i in range(self.num_upsamples): + x = F.leaky_relu(x, LRELU_SLOPE) + x = self.ups[i](x) + start = i * self.num_kernels + end = start + self.num_kernels + + # Evaluate all resblocks and average + block_outputs = torch.stack( + [self.resblocks[idx](x) for idx in range(start, end)], + dim=0, + ) + + x = block_outputs.mean(dim=0) + + x = self.conv_post(F.leaky_relu(x)) + return torch.tanh(x) + + +# ============================================================================= +# Configurators +# ============================================================================= + + +class AudioEncoderConfigurator: + """Factory for AudioEncoder from checkpoint config.""" + + @classmethod + def from_config(cls, config: dict) -> AudioEncoder: + audio_vae_cfg = config.get("audio_vae", {}) + model_cfg = audio_vae_cfg.get("model", {}) + model_params = model_cfg.get("params", {}) + ddconfig = model_params.get("ddconfig", {}) + preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) + stft_cfg = preprocessing_cfg.get("stft", {}) + mel_cfg = preprocessing_cfg.get("mel", {}) + variables_cfg = audio_vae_cfg.get("variables", {}) + + sample_rate = model_params.get("sampling_rate", 16000) + mel_hop_length = stft_cfg.get("hop_length", 160) + n_fft = stft_cfg.get("filter_length", 1024) + is_causal = stft_cfg.get("causal", True) + mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins") + + return AudioEncoder( + ch=ddconfig.get("ch", 128), + ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), + num_res_blocks=ddconfig.get("num_res_blocks", 2), + attn_resolutions=set(ddconfig.get("attn_resolutions", {8, 16, 32})), + resolution=ddconfig.get("resolution", 256), + z_channels=ddconfig.get("z_channels", 8), + double_z=ddconfig.get("double_z", True), + dropout=ddconfig.get("dropout", 0.0), + resamp_with_conv=ddconfig.get("resamp_with_conv", True), + in_channels=ddconfig.get("in_channels", 2), + attn_type=AttentionType(ddconfig.get("attn_type", "vanilla")), + mid_block_add_attention=ddconfig.get("mid_block_add_attention", True), + norm_type=NormType(ddconfig.get("norm_type", "pixel")), + causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), + sample_rate=sample_rate, + mel_hop_length=mel_hop_length, + n_fft=n_fft, + is_causal=is_causal, + mel_bins=mel_bins, + ) + + +class AudioDecoderConfigurator: + """Factory for AudioDecoder from checkpoint config.""" + + @classmethod + def from_config(cls, config: dict) -> AudioDecoder: + audio_vae_cfg = config.get("audio_vae", {}) + model_cfg = audio_vae_cfg.get("model", {}) + model_params = model_cfg.get("params", {}) + ddconfig = model_params.get("ddconfig", {}) + preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) + stft_cfg = preprocessing_cfg.get("stft", {}) + mel_cfg = preprocessing_cfg.get("mel", {}) + variables_cfg = audio_vae_cfg.get("variables", {}) + + sample_rate = model_params.get("sampling_rate", 16000) + mel_hop_length = stft_cfg.get("hop_length", 160) + is_causal = stft_cfg.get("causal", True) + mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins") + + return AudioDecoder( + ch=ddconfig.get("ch", 128), + out_ch=ddconfig.get("out_ch", 2), + ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), + num_res_blocks=ddconfig.get("num_res_blocks", 2), + attn_resolutions=set(ddconfig.get("attn_resolutions", {8, 16, 32})), + resolution=ddconfig.get("resolution", 256), + z_channels=ddconfig.get("z_channels", 8), + norm_type=NormType(ddconfig.get("norm_type", "pixel")), + causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), + dropout=ddconfig.get("dropout", 0.0), + mid_block_add_attention=ddconfig.get("mid_block_add_attention", True), + sample_rate=sample_rate, + mel_hop_length=mel_hop_length, + is_causal=is_causal, + mel_bins=mel_bins, + ) + + +class VocoderConfigurator: + """Factory for Vocoder from checkpoint config.""" + + @classmethod + def from_config(cls, config: dict) -> Vocoder: + vocoder_cfg = config.get("vocoder", {}) + return Vocoder( + resblock_kernel_sizes=vocoder_cfg.get("resblock_kernel_sizes", [3, 7, 11]), + upsample_rates=vocoder_cfg.get("upsample_rates", [6, 5, 2, 2, 2]), + upsample_kernel_sizes=vocoder_cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), + resblock_dilation_sizes=vocoder_cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), + upsample_initial_channel=vocoder_cfg.get("upsample_initial_channel", 1024), + stereo=vocoder_cfg.get("stereo", True), + resblock=vocoder_cfg.get("resblock", "1"), + output_sample_rate=vocoder_cfg.get("output_sample_rate", 24000), + ) + + +# ============================================================================= +# Public API Wrappers +# ============================================================================= + + +class LTX2AudioEncoder(nn.Module): + """Public wrapper for Audio Encoder with native FastVideo implementation.""" + + def __init__(self, config: dict[str, Any]): + super().__init__() + self.model: AudioEncoder = AudioEncoderConfigurator.from_config(config) + + def forward(self, spectrogram: torch.Tensor) -> torch.Tensor: + return self.model(spectrogram) + + +class LTX2AudioDecoder(nn.Module): + """Public wrapper for Audio Decoder with native FastVideo implementation.""" + + def __init__(self, config: dict[str, Any]): + super().__init__() + self.model: AudioDecoder = AudioDecoderConfigurator.from_config(config) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + return self.model(sample) + + +class LTX2Vocoder(nn.Module): + """Public wrapper for Vocoder with native FastVideo implementation.""" + + def __init__(self, config: dict[str, Any]): + super().__init__() + self.model: Vocoder = VocoderConfigurator.from_config(config) + + def forward(self, spectrogram: torch.Tensor) -> torch.Tensor: + return self.model(spectrogram) + + +def decode_audio( + latent: torch.Tensor, audio_decoder: AudioDecoder, vocoder: Vocoder +) -> torch.Tensor: + """ + Decode an audio latent representation using the provided audio decoder and vocoder. + """ + decoded_audio = audio_decoder(latent) + decoded_audio = vocoder(decoded_audio).squeeze(0).float() + return decoded_audio diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py new file mode 100644 index 000000000..b8e0aecc7 --- /dev/null +++ b/fastvideo/models/dits/ltx2.py @@ -0,0 +1,1800 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 transformer implementation +""" + +from dataclasses import dataclass, replace +from enum import Enum +import functools +import math +import os +from pathlib import Path +from typing import Any, Optional, Tuple, Callable + +import numpy as np + +import torch +import torch.nn as nn +from einops import rearrange, repeat + +from fastvideo.attention.backends.sdpa import SDPAMetadata +from fastvideo.attention.layer import LocalAttention +from fastvideo.configs.models.dits import LTX2VideoConfig +from fastvideo.forward_context import get_forward_context, set_forward_context +from fastvideo.models.dits.base import CachableDiT +from fastvideo.platforms import AttentionBackendEnum + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + """Sinusoidal timestep embedding used by LTX-2 AdaLN.""" + if len(timesteps.shape) != 1: + raise ValueError("Timesteps should be a 1d-array") + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + emb = scale * emb + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class TimestepEmbedding(torch.nn.Module): + """Two-layer MLP to project timestep embeddings.""" + def __init__( + self, + in_channels: int, + time_embed_dim: int, + out_dim: int | None = None, + post_act_fn: str | None = None, + cond_proj_dim: int | None = None, + sample_proj_bias: bool = True, + ): + super().__init__() + self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias) + self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) if cond_proj_dim is not None else None + self.act = torch.nn.SiLU() + time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim + self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) + self.post_act = None if post_act_fn is None else None + + def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor: + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + if self.act is not None: + sample = self.act(sample) + sample = self.linear_2(sample) + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class Timesteps(torch.nn.Module): + """Sinusoidal timestep embedding wrapper with scaling knobs.""" + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + return get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): + """PixArt-Alpha timestep embedding used by LTX-2 AdaLN.""" + def __init__(self, embedding_dim: int, size_emb_dim: int): + super().__init__() + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + return timesteps_emb + + +class AdaLayerNormSingle(torch.nn.Module): + """AdaLN-single modulation that emits scale/shift/gates for LTX-2.""" + def __init__(self, embedding_dim: int, embedding_coefficient: int = 6): + super().__init__() + self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim, + size_emb_dim=embedding_dim // 3, + ) + self.silu = torch.nn.SiLU() + self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True) + + def forward( + self, + timestep: torch.Tensor, + hidden_dtype: Optional[torch.dtype] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) + return self.linear(self.silu(embedded_timestep)), embedded_timestep + + +class PixArtAlphaTextProjection(torch.nn.Module): + """Caption projection MLP used by LTX-2.""" + def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"): + super().__init__() + if out_features is None: + out_features = hidden_size + self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) + if act_fn == "gelu_tanh": + self.act_1 = torch.nn.GELU(approximate="tanh") + elif act_fn == "silu": + self.act_1 = torch.nn.SiLU() + else: + raise ValueError(f"Unknown activation function: {act_fn}") + self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) + + def forward(self, caption: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +class GELUApprox(nn.Module): + """Linear + tanh-approximate GELU used by LTX-2 FFN.""" + def __init__(self, in_features: int, out_features: int): + super().__init__() + self.proj = nn.Linear(in_features, out_features) + self.act = nn.GELU(approximate="tanh") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(self.proj(x)) + + +class FeedForward(nn.Module): + """LTX-2 FFN: GELUApprox -> Identity -> Linear.""" + def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: + super().__init__() + inner_dim = int(dim * mult) + project_in = GELUApprox(dim, inner_dim) + self.net = nn.Sequential(project_in, nn.Identity(), nn.Linear(inner_dim, dim_out)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class VideoLatentShape(tuple): + """Helper for (B, C, T, H, W) latent shapes.""" + @property + def batch(self) -> int: + return self[0] + + @property + def channels(self) -> int: + return self[1] + + @property + def frames(self) -> int: + return self[2] + + @property + def height(self) -> int: + return self[3] + + @property + def width(self) -> int: + return self[4] + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "VideoLatentShape": + return VideoLatentShape(shape) + + def to_torch_shape(self) -> torch.Size: + return torch.Size(self) + + +class AudioLatentShape(tuple): + """Helper for (B, C, T, F) audio latent shapes.""" + + @property + def batch(self) -> int: + return self[0] + + @property + def channels(self) -> int: + return self[1] + + @property + def frames(self) -> int: + return self[2] + + @property + def mel_bins(self) -> int: + return self[3] + + def to_torch_shape(self) -> torch.Size: + return torch.Size(self) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "AudioLatentShape": + return AudioLatentShape(shape) + + @staticmethod + def from_duration( + batch: int, + duration: float, + channels: int, + mel_bins: int, + sample_rate: int, + hop_length: int, + audio_latent_downsample_factor: int, + ) -> "AudioLatentShape": + latents_per_second = float(sample_rate) / float( + hop_length) / float(audio_latent_downsample_factor) + return AudioLatentShape( + (batch, channels, round(duration * latents_per_second), mel_bins)) + + +class VideoLatentPatchifier: + """Patchify/unpatchify latent tokens for LTX-2.""" + def __init__(self, patch_size: int): + self._patch_size = (1, patch_size, patch_size) + + @property + def patch_size(self) -> Tuple[int, int, int]: + return self._patch_size + + def get_token_count(self, tgt_shape: VideoLatentShape) -> int: + return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size) + + def patchify(self, latents: torch.Tensor) -> torch.Tensor: + return rearrange( + latents, + "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", + p1=self._patch_size[0], + p2=self._patch_size[1], + p3=self._patch_size[2], + ) + + def unpatchify(self, latents: torch.Tensor, output_shape: VideoLatentShape) -> torch.Tensor: + patch_grid_frames = output_shape.frames // self._patch_size[0] + patch_grid_height = output_shape.height // self._patch_size[1] + patch_grid_width = output_shape.width // self._patch_size[2] + return rearrange( + latents, + "b (f h w) (c p q) -> b c f (h p) (w q)", + f=patch_grid_frames, + h=patch_grid_height, + w=patch_grid_width, + p=self._patch_size[1], + q=self._patch_size[2], + ) + + def get_patch_grid_bounds( + self, + output_shape: VideoLatentShape, + device: Optional[torch.device] = None, + ) -> torch.Tensor: + frames = output_shape.frames + height = output_shape.height + width = output_shape.width + batch_size = output_shape.batch + + grid_coords = torch.meshgrid( + torch.arange(start=0, end=frames, step=self._patch_size[0], device=device), + torch.arange(start=0, end=height, step=self._patch_size[1], device=device), + torch.arange(start=0, end=width, step=self._patch_size[2], device=device), + indexing="ij", + ) + + patch_starts = torch.stack(grid_coords, dim=0) + patch_size_delta = torch.tensor( + self._patch_size, + device=patch_starts.device, + dtype=patch_starts.dtype, + ).view(3, 1, 1, 1) + patch_ends = patch_starts + patch_size_delta + latent_coords = torch.stack((patch_starts, patch_ends), dim=-1) + latent_coords = repeat( + latent_coords, + "c f h w bounds -> b c (f h w) bounds", + b=batch_size, + bounds=2, + ) + return latent_coords + + +class AudioLatentPatchifier: + """Patchify/unpatchify audio latents and compute timing bounds.""" + + def __init__( + self, + patch_size: int, + sample_rate: int, + hop_length: int, + audio_latent_downsample_factor: int, + is_causal: bool = True, + shift: int = 0, + ) -> None: + self.hop_length = hop_length + self.sample_rate = sample_rate + self.audio_latent_downsample_factor = audio_latent_downsample_factor + self.is_causal = is_causal + self.shift = shift + self._patch_size = (1, patch_size, patch_size) + + def get_token_count(self, tgt_shape: AudioLatentShape) -> int: + return tgt_shape.frames + + def patchify(self, latents: torch.Tensor) -> torch.Tensor: + return rearrange(latents, "b c t f -> b t (c f)") + + def unpatchify( + self, + latents: torch.Tensor, + output_shape: AudioLatentShape, + ) -> torch.Tensor: + return rearrange( + latents, + "b t (c f) -> b c t f", + c=output_shape.channels, + f=output_shape.mel_bins, + ) + + def get_patch_grid_bounds( + self, + output_shape: AudioLatentShape, + device: Optional[torch.device] = None, + ) -> torch.Tensor: + start_timings = self._get_audio_latent_time_in_sec( + self.shift, + output_shape.frames + self.shift, + torch.float32, + device, + ) + start_timings = start_timings.unsqueeze(0).expand(output_shape.batch, + -1).unsqueeze(1) + + end_timings = self._get_audio_latent_time_in_sec( + self.shift + 1, + output_shape.frames + self.shift + 1, + torch.float32, + device, + ) + end_timings = end_timings.unsqueeze(0).expand(output_shape.batch, + -1).unsqueeze(1) + + return torch.stack([start_timings, end_timings], dim=-1) + + def _get_audio_latent_time_in_sec( + self, + start_latent: int, + end_latent: int, + dtype: torch.dtype, + device: Optional[torch.device], + ) -> torch.Tensor: + resolved_device = device or torch.device("cpu") + audio_latent_frame = torch.arange( + start_latent, end_latent, dtype=dtype, device=resolved_device) + audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor + if self.is_causal: + causal_offset = 1 + audio_mel_frame = ( + audio_mel_frame + causal_offset - + self.audio_latent_downsample_factor).clamp(min=0) + return audio_mel_frame * self.hop_length / self.sample_rate + + +def _get_pixel_coords( + latent_coords: torch.Tensor, + scale_factors: tuple[int, int, int], + fps: float | None, + causal_fix: bool = True, +) -> torch.Tensor: + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 + scale_tensor = torch.tensor( + scale_factors, + device=latent_coords.device, + dtype=torch.float32, + ).view(*broadcast_shape) + pixel_coords = latent_coords.to(torch.float32) * scale_tensor + if causal_fix: + pixel_coords[:, 0, ...] = ( + pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0) + if fps: + pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps + return pixel_coords + + +def _to_denoised( + sample: torch.Tensor, + velocity: torch.Tensor, + sigma: torch.Tensor, + calc_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(calc_dtype) + while sigma.ndim < sample.ndim: + sigma = sigma.unsqueeze(-1) + return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) + + +def _debug_block_log_line(message: str) -> None: + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1": + return + log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "") + if not log_path: + return + log_dir = os.path.dirname(log_path) + if log_dir: + os.makedirs(log_dir, exist_ok=True) + with open(log_path, "a", encoding="utf-8") as f: + f.write(message + "\n") + + +def _debug_transformer_args(prefix: str, args: "TransformerArgs | None") -> None: + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1" or args is None: + return + pe_cos, pe_sin = args.positional_embeddings + cross_cos = None + cross_sin = None + if args.cross_positional_embeddings is not None: + cross_cos, cross_sin = args.cross_positional_embeddings + mask = args.context_mask + if mask is None: + mask_summary = "mask=None" + else: + finite = torch.isfinite(mask) + finite_sum = mask[finite].sum().item() if finite.any() else 0.0 + mask_summary = ( + f"mask_min={mask.min().item():.6f} " + f"mask_max={mask.max().item():.6f} " + f"mask_finite_sum={finite_sum:.6f} " + f"mask_finite_count={finite.sum().item()}" + ) + _debug_block_log_line( + f"{prefix}:x_sum={args.x.float().sum().item():.6f} " + f"context_sum={args.context.float().sum().item():.6f} " + f"t_sum={args.timesteps.float().sum().item():.6f} " + f"emb_sum={args.embedded_timestep.float().sum().item():.6f} " + f"pe_cos_sum={pe_cos.float().sum().item():.6f} " + f"pe_sin_sum={pe_sin.float().sum().item():.6f} " + f"cross_pe_cos_sum={(cross_cos.float().sum().item() if cross_cos is not None else 0.0):.6f} " + f"cross_pe_sin_sum={(cross_sin.float().sum().item() if cross_sin is not None else 0.0):.6f} " + f"{mask_summary}" + ) + + +class LTXRopeType(Enum): + """LTX-2 rotary variants (interleaved vs split).""" + INTERLEAVED = "interleaved" + SPLIT = "split" + + +DEFAULT_LTX2_SCALE_FACTORS = (8, 32, 32) +DEFAULT_LTX2_AUDIO_CHANNELS = 8 +DEFAULT_LTX2_AUDIO_MEL_BINS = 16 +DEFAULT_LTX2_AUDIO_SAMPLE_RATE = 16000 +DEFAULT_LTX2_AUDIO_HOP_LENGTH = 160 +DEFAULT_LTX2_AUDIO_DOWNSAMPLE = 4 +# The vocoder upsamples from mel spectrograms (16kHz sample rate) to 24kHz audio +DEFAULT_LTX2_VOCODER_OUTPUT_SAMPLE_RATE = 24000 + + +def apply_ltx_rotary_emb( + input_tensor: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, +) -> torch.Tensor: + """Apply LTX-2 rotary embeddings to a tensor.""" + if rope_type == LTXRopeType.INTERLEAVED: + return _apply_ltx_interleaved_rotary_emb(input_tensor, *freqs_cis) + if rope_type == LTXRopeType.SPLIT: + return _apply_ltx_split_rotary_emb(input_tensor, *freqs_cis) + raise ValueError(f"Invalid rope type: {rope_type}") + + +def _apply_ltx_interleaved_rotary_emb( + input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor +) -> torch.Tensor: + t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) + t1, t2 = t_dup.unbind(dim=-1) + t_dup = torch.stack((-t2, t1), dim=-1) + input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") + return input_tensor * cos_freqs + input_tensor_rot * sin_freqs + + +def _apply_ltx_split_rotary_emb( + input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor +) -> torch.Tensor: + needs_reshape = False + if input_tensor.ndim != 4 and cos_freqs.ndim == 4: + b, h, t, _ = cos_freqs.shape + input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2) + needs_reshape = True + + split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) + first_half_input = split_input[..., :1, :] + second_half_input = split_input[..., 1:, :] + + output = split_input * cos_freqs.unsqueeze(-2) + first_half_output = output[..., :1, :] + second_half_output = output[..., 1:, :] + + first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input) + second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input) + + output = rearrange(output, "... d r -> ... (d r)") + if needs_reshape: + output = output.swapaxes(1, 2).reshape(b, t, -1) + return output + + +@functools.lru_cache(maxsize=5) +def generate_ltx_freq_grid_np( + positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int +) -> torch.Tensor: + """Generate LTX-2 rotary frequencies with high-precision numpy.""" + theta = positional_embedding_theta + start = 1 + end = theta + n_elem = 2 * positional_embedding_max_pos_count + pow_indices = np.power( + theta, + np.linspace( + np.log(start) / np.log(theta), + np.log(end) / np.log(theta), + inner_dim // n_elem, + dtype=np.float64, + ), + ) + return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32) + + +@functools.lru_cache(maxsize=5) +def generate_ltx_freq_grid_pytorch( + positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int +) -> torch.Tensor: + """Generate LTX-2 rotary frequencies in torch for speed.""" + theta = positional_embedding_theta + start = 1 + end = theta + n_elem = 2 * positional_embedding_max_pos_count + indices = theta ** ( + torch.linspace( + math.log(start, theta), + math.log(end, theta), + inner_dim // n_elem, + dtype=torch.float32, + ) + ) + indices = indices.to(dtype=torch.float32) + return indices * math.pi / 2 + + +def _ltx_get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor: + n_pos_dims = indices_grid.shape[1] + if n_pos_dims != len(max_pos): + raise ValueError( + f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})" + ) + return torch.stack( + [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)], + dim=-1, + ) + + +def _ltx_generate_freqs( + indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool +) -> torch.Tensor: + if use_middle_indices_grid: + indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1] + indices_grid = (indices_grid_start + indices_grid_end) / 2.0 + elif len(indices_grid.shape) == 4: + indices_grid = indices_grid[..., 0] + + fractional_positions = _ltx_get_fractional_positions(indices_grid, max_pos) + indices = indices.to(device=fractional_positions.device) + freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2) + return freqs + + +def _ltx_split_freqs_cis( + freqs: torch.Tensor, pad_size: int, num_attention_heads: int +) -> tuple[torch.Tensor, torch.Tensor]: + cos_freq = freqs.cos() + sin_freq = freqs.sin() + + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) + cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) + sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) + + b = cos_freq.shape[0] + t = cos_freq.shape[1] + cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1) + sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1) + cos_freq = torch.swapaxes(cos_freq, 1, 2) + sin_freq = torch.swapaxes(sin_freq, 1, 2) + return cos_freq, sin_freq + + +def _ltx_interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]: + cos_freq = freqs.cos().repeat_interleave(2, dim=-1) + sin_freq = freqs.sin().repeat_interleave(2, dim=-1) + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size]) + cos_freq = torch.cat([cos_padding, cos_freq], dim=-1) + sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) + return cos_freq, sin_freq + + +def precompute_ltx_freqs_cis( + indices_grid: torch.Tensor, + dim: int, + out_dtype: torch.dtype, + theta: float = 10000.0, + max_pos: list[int] | None = None, + use_middle_indices_grid: bool = False, + num_attention_heads: int = 32, + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, + freq_grid_generator: Callable[[float, int, int], torch.Tensor] = generate_ltx_freq_grid_pytorch, +) -> tuple[torch.Tensor, torch.Tensor]: + """Precompute LTX-2 rotary cos/sin grids for (t, x, y) positions.""" + if max_pos is None: + max_pos = [20, 2048, 2048] + + indices = freq_grid_generator(theta, indices_grid.shape[1], dim) + freqs = _ltx_generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid) + + if rope_type == LTXRopeType.SPLIT: + expected_freqs = dim // 2 + current_freqs = freqs.shape[-1] + pad_size = expected_freqs - current_freqs + cos_freq, sin_freq = _ltx_split_freqs_cis(freqs, pad_size, num_attention_heads) + else: + n_elem = 2 * indices_grid.shape[1] + cos_freq, sin_freq = _ltx_interleaved_freqs_cis(freqs, dim % n_elem) + return cos_freq.to(out_dtype), sin_freq.to(out_dtype) + + +@dataclass(frozen=True) +class TransformerArgs: + """Pack transformer inputs for LTX-2 blocks.""" + x: torch.Tensor + context: torch.Tensor + context_mask: torch.Tensor | None + timesteps: torch.Tensor + embedded_timestep: torch.Tensor + positional_embeddings: torch.Tensor + cross_positional_embeddings: torch.Tensor | None + cross_scale_shift_timestep: torch.Tensor | None + cross_gate_timestep: torch.Tensor | None + enabled: bool + + +@dataclass(frozen=True) +class Modality: + """Lightweight modality container for LTX-2 inputs.""" + enabled: bool + latent: torch.Tensor + timesteps: torch.Tensor + positions: torch.Tensor + context: torch.Tensor + context_mask: torch.Tensor | None = None + + +class TransformerArgsPreprocessor: + """Prepare LTX-2 transformer inputs (patchify, AdaLN, rope).""" + def __init__( # noqa: PLR0913 + self, + patchify_proj: torch.nn.Linear, + adaln: AdaLayerNormSingle, + caption_projection: PixArtAlphaTextProjection, + inner_dim: int, + max_pos: list[int], + num_attention_heads: int, + use_middle_indices_grid: bool, + timestep_scale_multiplier: int, + double_precision_rope: bool, + positional_embedding_theta: float, + rope_type: LTXRopeType, + ) -> None: + self.patchify_proj = patchify_proj + self.adaln = adaln + self.caption_projection = caption_projection + self.inner_dim = inner_dim + self.max_pos = max_pos + self.num_attention_heads = num_attention_heads + self.use_middle_indices_grid = use_middle_indices_grid + self.timestep_scale_multiplier = timestep_scale_multiplier + self.double_precision_rope = double_precision_rope + self.positional_embedding_theta = positional_embedding_theta + self.rope_type = rope_type + + def _prepare_timestep( + self, timestep: torch.Tensor, batch_size: int, hidden_dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + timestep = timestep * self.timestep_scale_multiplier + timestep, embedded_timestep = self.adaln(timestep.flatten(), hidden_dtype=hidden_dtype) + timestep = timestep.view(batch_size, -1, timestep.shape[-1]) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1]) + return timestep, embedded_timestep + + def _prepare_context( + self, + context: torch.Tensor, + x: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + batch_size = x.shape[0] + if context.device != x.device: + context = context.to(x.device) + if attention_mask is not None and attention_mask.device != x.device: + attention_mask = attention_mask.to(x.device) + context = self.caption_projection(context) + context = context.view(batch_size, -1, x.shape[-1]) + return context, attention_mask + + def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: + if attention_mask is None or torch.is_floating_point(attention_mask): + return attention_mask + return (attention_mask - 1).to(x_dtype).reshape( + (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) + ) * torch.finfo(x_dtype).max + + def _prepare_positional_embeddings( + self, + positions: torch.Tensor, + inner_dim: int, + max_pos: list[int], + use_middle_indices_grid: bool, + num_attention_heads: int, + x_dtype: torch.dtype, + ) -> torch.Tensor: + if self.double_precision_rope: + freq_grid_generator = generate_ltx_freq_grid_np + else: + freq_grid_generator = generate_ltx_freq_grid_pytorch + return precompute_ltx_freqs_cis( + positions, + dim=inner_dim, + out_dtype=x_dtype, + theta=self.positional_embedding_theta, + max_pos=max_pos, + use_middle_indices_grid=use_middle_indices_grid, + num_attention_heads=num_attention_heads, + rope_type=self.rope_type, + freq_grid_generator=freq_grid_generator, + ) + + def prepare(self, modality: Modality) -> TransformerArgs: + x = self.patchify_proj(modality.latent) + timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype) + context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask) + attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype) + pe = self._prepare_positional_embeddings( + positions=modality.positions, + inner_dim=self.inner_dim, + max_pos=self.max_pos, + use_middle_indices_grid=self.use_middle_indices_grid, + num_attention_heads=self.num_attention_heads, + x_dtype=modality.latent.dtype, + ) + return TransformerArgs( + x=x, + context=context, + context_mask=attention_mask, + timesteps=timestep, + embedded_timestep=embedded_timestep, + positional_embeddings=pe, + cross_positional_embeddings=None, + cross_scale_shift_timestep=None, + cross_gate_timestep=None, + enabled=modality.enabled, + ) + + +class MultiModalTransformerArgsPreprocessor: + """Prepare transformer args for audio/video cross-attention paths.""" + + def __init__( # noqa: PLR0913 + self, + patchify_proj: torch.nn.Linear, + adaln: AdaLayerNormSingle, + caption_projection: PixArtAlphaTextProjection, + cross_scale_shift_adaln: AdaLayerNormSingle, + cross_gate_adaln: AdaLayerNormSingle, + inner_dim: int, + max_pos: list[int], + num_attention_heads: int, + cross_pe_max_pos: int, + use_middle_indices_grid: bool, + audio_cross_attention_dim: int, + timestep_scale_multiplier: int, + double_precision_rope: bool, + positional_embedding_theta: float, + rope_type: LTXRopeType, + av_ca_timestep_scale_multiplier: int, + ) -> None: + self.simple_preprocessor = TransformerArgsPreprocessor( + patchify_proj=patchify_proj, + adaln=adaln, + caption_projection=caption_projection, + inner_dim=inner_dim, + max_pos=max_pos, + num_attention_heads=num_attention_heads, + use_middle_indices_grid=use_middle_indices_grid, + timestep_scale_multiplier=timestep_scale_multiplier, + double_precision_rope=double_precision_rope, + positional_embedding_theta=positional_embedding_theta, + rope_type=rope_type, + ) + self.cross_scale_shift_adaln = cross_scale_shift_adaln + self.cross_gate_adaln = cross_gate_adaln + self.cross_pe_max_pos = cross_pe_max_pos + self.audio_cross_attention_dim = audio_cross_attention_dim + self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier + + def prepare(self, modality: Modality) -> TransformerArgs: + transformer_args = self.simple_preprocessor.prepare(modality) + cross_pe = self.simple_preprocessor._prepare_positional_embeddings( + positions=modality.positions[:, 0:1, :], + inner_dim=self.audio_cross_attention_dim, + max_pos=[self.cross_pe_max_pos], + use_middle_indices_grid=True, + num_attention_heads=self.simple_preprocessor.num_attention_heads, + x_dtype=modality.latent.dtype, + ) + + cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( + timestep=modality.timesteps, + timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, + batch_size=transformer_args.x.shape[0], + hidden_dtype=modality.latent.dtype, + ) + + return replace( + transformer_args, + cross_positional_embeddings=cross_pe, + cross_scale_shift_timestep=cross_scale_shift_timestep, + cross_gate_timestep=cross_gate_timestep, + ) + + def _prepare_cross_attention_timestep( + self, + timestep: torch.Tensor, + timestep_scale_multiplier: int, + batch_size: int, + hidden_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + timestep = timestep * timestep_scale_multiplier + + av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier + + scale_shift_timestep, _ = self.cross_scale_shift_adaln( + timestep.flatten(), + hidden_dtype=hidden_dtype, + ) + scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) + gate_noise_timestep, _ = self.cross_gate_adaln( + timestep.flatten() * av_ca_factor, + hidden_dtype=hidden_dtype, + ) + gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) + + return scale_shift_timestep, gate_noise_timestep + + +@dataclass +class TransformerConfig: + """Attention/FFN dims for LTX-2 transformer blocks.""" + dim: int + heads: int + d_head: int + context_dim: int + + +class LTXSelfAttention(nn.Module): + """LTX-2 attention block with RMSNorm + FastVideo LocalAttention.""" + def __init__( + self, + query_dim: int, + context_dim: int | None, + heads: int, + dim_head: int, + norm_eps: float, + rope_type: LTXRopeType, + supported_attention_backends: tuple[AttentionBackendEnum, ...], + ) -> None: + super().__init__() + inner_dim = dim_head * heads + context_dim = query_dim if context_dim is None else context_dim + + self.heads = heads + self.dim_head = dim_head + self.rope_type = rope_type + + self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + self.to_q = nn.Linear(query_dim, inner_dim, bias=True) + self.to_k = nn.Linear(context_dim, inner_dim, bias=True) + self.to_v = nn.Linear(context_dim, inner_dim, bias=True) + self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()) + + self.attn = LocalAttention( + num_heads=heads, + head_size=dim_head, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + self.attn_masked = LocalAttention( + num_heads=heads, + head_size=dim_head, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA,), + ) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + pe: tuple[torch.Tensor, torch.Tensor] | None = None, + k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + q = self.to_q(x) + context = x if context is None else context + k = self.to_k(context) + v = self.to_v(context) + + q = self.q_norm(q) + k = self.k_norm(k) + + if pe is not None: + q = apply_ltx_rotary_emb(q, pe, self.rope_type) + k = apply_ltx_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) + + b, q_len, _ = q.shape + k_len = k.shape[1] + q = q.view(b, q_len, self.heads, self.dim_head) + k = k.view(b, k_len, self.heads, self.dim_head) + v = v.view(b, k_len, self.heads, self.dim_head) + if mask is not None: + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(1) + try: + forward_ctx = get_forward_context() + current_timestep = forward_ctx.current_timestep + forward_batch = forward_ctx.forward_batch + except AssertionError: + current_timestep = 0 + forward_batch = None + attn_metadata = SDPAMetadata( + current_timestep=current_timestep, + attn_mask=mask, + ) + with set_forward_context( + current_timestep=current_timestep, + attn_metadata=attn_metadata, + forward_batch=forward_batch, + ): + out = self.attn_masked(q, k, v) + else: + out = self.attn(q, k, v) + out = out.reshape(b, q_len, -1) + return self.to_out(out) + + +class BasicAVTransformerBlock(torch.nn.Module): + """LTX-2 transformer block (audio/video + cross attention + FFN).""" + + def __init__( + self, + idx: int, + video: TransformerConfig | None = None, + audio: TransformerConfig | None = None, + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, + norm_eps: float = 1e-6, + ): + super().__init__() + self.idx = idx + + if video is not None: + self.attn1 = LTXSelfAttention( + query_dim=video.dim, + context_dim=None, + heads=video.heads, + dim_head=video.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + ) + self.attn2 = LTXSelfAttention( + query_dim=video.dim, + context_dim=video.context_dim, + heads=video.heads, + dim_head=video.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + ) + self.ff = FeedForward(video.dim, dim_out=video.dim) + self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim)) + + if audio is not None: + self.audio_attn1 = LTXSelfAttention( + query_dim=audio.dim, + context_dim=None, + heads=audio.heads, + dim_head=audio.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + ) + self.audio_attn2 = LTXSelfAttention( + query_dim=audio.dim, + context_dim=audio.context_dim, + heads=audio.heads, + dim_head=audio.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + ) + self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim) + self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim)) + + if audio is not None and video is not None: + self.audio_to_video_attn = LTXSelfAttention( + query_dim=video.dim, + context_dim=audio.dim, + heads=audio.heads, + dim_head=audio.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + ) + self.video_to_audio_attn = LTXSelfAttention( + query_dim=audio.dim, + context_dim=video.dim, + heads=audio.heads, + dim_head=audio.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + ) + self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim)) + self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim)) + + self.norm_eps = norm_eps + + def get_ada_values( + self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice + ) -> tuple[torch.Tensor, ...]: + num_ada_params = scale_shift_table.shape[0] + ada_values = ( + scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) + + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] + ).unbind(dim=2) + return ada_values + + def get_av_ca_ada_values( + self, + scale_shift_table: torch.Tensor, + batch_size: int, + scale_shift_timestep: torch.Tensor, + gate_timestep: torch.Tensor, + num_scale_shift_values: int = 4, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + scale_shift_ada_values = self.get_ada_values( + scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None) + ) + gate_ada_values = self.get_ada_values( + scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None) + ) + + scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values] + gate_ada_values = [t.squeeze(2) for t in gate_ada_values] + + return (*scale_shift_chunks, *gate_ada_values) + + def forward( + self, + video: TransformerArgs | None, + audio: TransformerArgs | None, + ) -> tuple[TransformerArgs | None, TransformerArgs | None]: + vx = video.x if video is not None else None + ax = audio.x if audio is not None else None + + run_vx = video is not None and video.enabled and vx.numel() > 0 + run_ax = audio is not None and audio.enabled and ax.numel() > 0 + run_a2v = run_vx and run_ax + run_v2a = run_ax and run_vx + + if run_vx: + vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( + self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3) + ) + norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa + vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa + vx = vx + self.attn2( + torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps), + context=video.context, + mask=video.context_mask, + ) + + if run_ax: + ashift_msa, ascale_msa, agate_msa = self.get_ada_values( + self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3) + ) + norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa + ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa + ax = ax + self.audio_attn2( + torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps), + context=audio.context, + mask=audio.context_mask, + ) + + if run_a2v or run_v2a: + vx_norm3 = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) + ax_norm3 = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) + + ( + scale_ca_audio_hidden_states_a2v, + shift_ca_audio_hidden_states_a2v, + scale_ca_audio_hidden_states_v2a, + shift_ca_audio_hidden_states_v2a, + gate_out_v2a, + ) = self.get_av_ca_ada_values( + self.scale_shift_table_a2v_ca_audio, + ax.shape[0], + audio.cross_scale_shift_timestep, + audio.cross_gate_timestep, + ) + + ( + scale_ca_video_hidden_states_a2v, + shift_ca_video_hidden_states_a2v, + scale_ca_video_hidden_states_v2a, + shift_ca_video_hidden_states_v2a, + gate_out_a2v, + ) = self.get_av_ca_ada_values( + self.scale_shift_table_a2v_ca_video, + vx.shape[0], + video.cross_scale_shift_timestep, + video.cross_gate_timestep, + ) + + if run_a2v: + vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v + ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v + vx = vx + ( + self.audio_to_video_attn( + vx_scaled, + context=ax_scaled, + pe=video.cross_positional_embeddings, + k_pe=audio.cross_positional_embeddings, + ) + * gate_out_a2v + ) + + if run_v2a: + ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a + vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a + ax = ax + ( + self.video_to_audio_attn( + ax_scaled, + context=vx_scaled, + pe=audio.cross_positional_embeddings, + k_pe=video.cross_positional_embeddings, + ) + * gate_out_v2a + ) + + if run_vx: + vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( + self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, None) + ) + vx_scaled = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp + vx = vx + self.ff(vx_scaled) * vgate_mlp + + if run_ax: + ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values( + self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, None) + ) + ax_scaled = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp + ax = ax + self.audio_ff(ax_scaled) * agate_mlp + + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": + video_sum = vx.float().sum().item() if vx is not None else 0.0 + audio_sum = ax.float().sum().item() if ax is not None else 0.0 + _debug_block_log_line( + f"fastvideo:block={self.idx}:video_sum={video_sum:.6f} " + f"audio_sum={audio_sum:.6f}" + ) + + return ( + replace(video, x=vx) if video is not None else None, + replace(audio, x=ax) if audio is not None else None, + ) + + +class LTXModelType(Enum): + """Model type flags for LTX-2.""" + AudioVideo = "ltx av model" + VideoOnly = "ltx video only model" + AudioOnly = "ltx audio only model" + + def is_video_enabled(self) -> bool: + return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly) + + def is_audio_enabled(self) -> bool: + return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly) + + +class LTXModel(torch.nn.Module): + """Core LTX-2 transformer stack for audio/video latents.""" + + def __init__( # noqa: PLR0913 + self, + *, + model_type: LTXModelType = LTXModelType.AudioVideo, + num_attention_heads: int = 32, + attention_head_dim: int = 128, + in_channels: int = 128, + out_channels: int = 128, + num_layers: int = 48, + cross_attention_dim: int = 4096, + norm_eps: float = 1e-06, + caption_channels: int = 3840, + positional_embedding_theta: float = 10000.0, + positional_embedding_max_pos: list[int] | None = None, + timestep_scale_multiplier: int = 1000, + use_middle_indices_grid: bool = True, + audio_num_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + audio_in_channels: int = 128, + audio_out_channels: int = 128, + audio_cross_attention_dim: int = 2048, + audio_positional_embedding_max_pos: list[int] | None = None, + av_ca_timestep_scale_multiplier: int = 1, + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, + double_precision_rope: bool = False, + ): + super().__init__() + self._enable_gradient_checkpointing = False + self.use_middle_indices_grid = use_middle_indices_grid + self.rope_type = rope_type + self.double_precision_rope = double_precision_rope + self.timestep_scale_multiplier = timestep_scale_multiplier + self.positional_embedding_theta = positional_embedding_theta + self.model_type = model_type + cross_pe_max_pos = None + + if model_type.is_video_enabled(): + if positional_embedding_max_pos is None: + positional_embedding_max_pos = [20, 2048, 2048] + self.positional_embedding_max_pos = positional_embedding_max_pos + self.num_attention_heads = num_attention_heads + self.inner_dim = num_attention_heads * attention_head_dim + self._init_video( + in_channels=in_channels, + out_channels=out_channels, + caption_channels=caption_channels, + norm_eps=norm_eps, + ) + + if model_type.is_audio_enabled(): + if audio_positional_embedding_max_pos is None: + audio_positional_embedding_max_pos = [20] + self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos + self.audio_num_attention_heads = audio_num_attention_heads + self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim + self._init_audio( + in_channels=audio_in_channels, + out_channels=audio_out_channels, + caption_channels=caption_channels, + norm_eps=norm_eps, + ) + + if model_type.is_video_enabled() and model_type.is_audio_enabled(): + cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) + self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier + self.audio_cross_attention_dim = audio_cross_attention_dim + self._init_audio_video(num_scale_shift_values=4) + + self._init_preprocessors(cross_pe_max_pos) + self._init_transformer_blocks( + num_layers=num_layers, + attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0, + cross_attention_dim=cross_attention_dim, + audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0, + audio_cross_attention_dim=audio_cross_attention_dim, + norm_eps=norm_eps, + ) + + def _init_video( + self, + in_channels: int, + out_channels: int, + caption_channels: int, + norm_eps: float, + ) -> None: + self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) + self.adaln_single = AdaLayerNormSingle(self.inner_dim) + self.caption_projection = PixArtAlphaTextProjection( + in_features=caption_channels, + hidden_size=self.inner_dim, + ) + 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) + + def _init_audio( + self, + in_channels: int, + out_channels: int, + caption_channels: int, + norm_eps: float, + ) -> None: + self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) + self.audio_adaln_single = AdaLayerNormSingle(self.audio_inner_dim) + self.audio_caption_projection = PixArtAlphaTextProjection( + in_features=caption_channels, + hidden_size=self.audio_inner_dim, + ) + 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) + + def _init_audio_video(self, num_scale_shift_values: int) -> None: + self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle( + self.inner_dim, + embedding_coefficient=num_scale_shift_values, + ) + self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle( + self.audio_inner_dim, + embedding_coefficient=num_scale_shift_values, + ) + self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle( + self.inner_dim, + embedding_coefficient=1, + ) + self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle( + self.audio_inner_dim, + embedding_coefficient=1, + ) + + def _init_preprocessors(self, cross_pe_max_pos: int | None = None) -> None: + if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): + self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor( + patchify_proj=self.patchify_proj, + adaln=self.adaln_single, + caption_projection=self.caption_projection, + cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single, + cross_gate_adaln=self.av_ca_a2v_gate_adaln_single, + inner_dim=self.inner_dim, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + cross_pe_max_pos=cross_pe_max_pos, + use_middle_indices_grid=self.use_middle_indices_grid, + audio_cross_attention_dim=self.audio_cross_attention_dim, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, + ) + self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor( + patchify_proj=self.audio_patchify_proj, + adaln=self.audio_adaln_single, + caption_projection=self.audio_caption_projection, + cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single, + cross_gate_adaln=self.av_ca_v2a_gate_adaln_single, + inner_dim=self.audio_inner_dim, + max_pos=self.audio_positional_embedding_max_pos, + num_attention_heads=self.audio_num_attention_heads, + cross_pe_max_pos=cross_pe_max_pos, + use_middle_indices_grid=self.use_middle_indices_grid, + audio_cross_attention_dim=self.audio_cross_attention_dim, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, + ) + elif self.model_type.is_video_enabled(): + self.video_args_preprocessor = TransformerArgsPreprocessor( + patchify_proj=self.patchify_proj, + adaln=self.adaln_single, + caption_projection=self.caption_projection, + inner_dim=self.inner_dim, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + use_middle_indices_grid=self.use_middle_indices_grid, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + ) + elif self.model_type.is_audio_enabled(): + self.audio_args_preprocessor = TransformerArgsPreprocessor( + patchify_proj=self.audio_patchify_proj, + adaln=self.audio_adaln_single, + caption_projection=self.audio_caption_projection, + inner_dim=self.audio_inner_dim, + max_pos=self.audio_positional_embedding_max_pos, + num_attention_heads=self.audio_num_attention_heads, + use_middle_indices_grid=self.use_middle_indices_grid, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + ) + + def _init_transformer_blocks( + self, + num_layers: int, + attention_head_dim: int, + cross_attention_dim: int, + audio_attention_head_dim: int, + audio_cross_attention_dim: int, + norm_eps: float, + ) -> None: + video_config = ( + TransformerConfig( + dim=self.inner_dim, + heads=self.num_attention_heads, + d_head=attention_head_dim, + context_dim=cross_attention_dim, + ) + if self.model_type.is_video_enabled() + else None + ) + audio_config = ( + TransformerConfig( + dim=self.audio_inner_dim, + heads=self.audio_num_attention_heads, + d_head=audio_attention_head_dim, + context_dim=audio_cross_attention_dim, + ) + if self.model_type.is_audio_enabled() + else None + ) + self.transformer_blocks = torch.nn.ModuleList( + [ + BasicAVTransformerBlock( + idx=idx, + video=video_config, + audio=audio_config, + rope_type=self.rope_type, + norm_eps=norm_eps, + ) + for idx in range(num_layers) + ] + ) + + def _process_transformer_blocks( + self, + video: TransformerArgs | None, + audio: TransformerArgs | None, + ) -> tuple[TransformerArgs | None, TransformerArgs | None]: + for block in self.transformer_blocks: + video, audio = block(video=video, audio=audio) + return video, audio + + def _process_output( + self, + scale_shift_table: torch.Tensor, + norm_out: torch.nn.LayerNorm, + proj_out: torch.nn.Linear, + x: torch.Tensor, + embedded_timestep: torch.Tensor, + ) -> torch.Tensor: + scale_shift_values = ( + scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None] + ) + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + x = norm_out(x) + x = x * (1 + scale) + shift + x = proj_out(x) + return x + + def forward( + self, + video: Modality | None, + audio: Modality | None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": + _debug_block_log_line( + "fastvideo:patchify_proj" + f":video_w_sum={self.patchify_proj.weight.float().sum().item():.6f} " + f"video_b_sum={self.patchify_proj.bias.float().sum().item():.6f} " + f"audio_w_sum={self.audio_patchify_proj.weight.float().sum().item():.6f} " + f"audio_b_sum={self.audio_patchify_proj.bias.float().sum().item():.6f}" + ) + if not self.model_type.is_video_enabled() and video is not None: + raise ValueError("Video is not enabled for this model") + if not self.model_type.is_audio_enabled() and audio is not None: + raise ValueError("Audio is not enabled for this model") + + video_args = self.video_args_preprocessor.prepare(video) if video is not None else None + audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None + _debug_transformer_args("fastvideo:prep_video", video_args) + _debug_transformer_args("fastvideo:prep_audio", audio_args) + video_out, audio_out = self._process_transformer_blocks(video_args, audio_args) + + vx = ( + self._process_output( + self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep + ) + if video_out is not None + else None + ) + ax = ( + self._process_output( + self.audio_scale_shift_table, + self.audio_norm_out, + self.audio_proj_out, + audio_out.x, + audio_out.embedded_timestep, + ) + if audio_out is not None + else None + ) + return vx, ax + + +class LTX2Transformer3DModel(CachableDiT): + """ + LTX-2 transformer using native FastVideo LTX-2 modules. + """ + + param_names_mapping = LTX2VideoConfig().param_names_mapping + reverse_param_names_mapping = LTX2VideoConfig().reverse_param_names_mapping + lora_param_names_mapping = LTX2VideoConfig().lora_param_names_mapping + + def __init__(self, config: LTX2VideoConfig, hf_config: dict[str, Any]): + super().__init__(config=config, hf_config=hf_config) + + arch = config.arch_config + + model_type = LTXModelType.AudioVideo + self.model = LTXModel( + model_type=model_type, + num_attention_heads=arch.num_attention_heads, + attention_head_dim=arch.attention_head_dim, + in_channels=arch.in_channels, + out_channels=arch.out_channels, + num_layers=arch.num_layers, + cross_attention_dim=arch.cross_attention_dim, + norm_eps=arch.norm_eps, + caption_channels=arch.caption_channels, + positional_embedding_theta=arch.positional_embedding_theta, + positional_embedding_max_pos=arch.positional_embedding_max_pos, + timestep_scale_multiplier=arch.timestep_scale_multiplier, + use_middle_indices_grid=arch.use_middle_indices_grid, + rope_type=LTXRopeType(arch.rope_type), + double_precision_rope=arch.double_precision_rope, + audio_num_attention_heads=arch.audio_num_attention_heads, + audio_attention_head_dim=arch.audio_attention_head_dim, + audio_in_channels=arch.audio_in_channels, + audio_out_channels=arch.audio_out_channels, + audio_cross_attention_dim=arch.audio_cross_attention_dim, + audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos, + av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier, + ) + + self.patchifier = VideoLatentPatchifier(patch_size=arch.patch_size[1]) + self.audio_patchifier = AudioLatentPatchifier( + patch_size=DEFAULT_LTX2_AUDIO_MEL_BINS, + sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE, + hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH, + audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE, + is_causal=True, + shift=0, + ) + + self.hidden_size = arch.num_attention_heads * arch.attention_head_dim + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.num_channels_latents + + if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1": + detail_path = os.getenv("LTX2_PIPELINE_DEBUG_DETAIL_PATH", "") + if detail_path: + self._attach_debug_detail_hooks(detail_path) + + def _attach_debug_detail_hooks(self, log_path: str) -> None: + path = Path(log_path) + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + path.unlink() + + def _format_sum(tensor: torch.Tensor | None) -> str: + if tensor is None: + return "None" + return f"{tensor.float().sum().item():.6f}" + + def _hook_factory(block_idx: int, name: str): + def _hook(_module, _inputs, outputs): # noqa: ANN001 + if isinstance(outputs, tuple): + out = outputs[0] + else: + out = outputs + out_sum = _format_sum(out if torch.is_tensor(out) else None) + with path.open("a", encoding="utf-8") as f: + f.write(f"fastvideo:{block_idx}:{name}:out_sum={out_sum}\n") + return _hook + + for block in self.model.transformer_blocks: + idx = block.idx + for name in ( + "attn1", + "attn2", + "ff", + "audio_attn1", + "audio_attn2", + "audio_ff", + "audio_to_video_attn", + "video_to_audio_attn", + ): + if hasattr(block, name): + getattr(block, name).register_forward_hook( + _hook_factory(idx, name)) + + def _output_hook(label: str): + def _hook(_module, _inputs, outputs): # noqa: ANN001 + out = outputs[0] if isinstance(outputs, tuple) else outputs + out_sum = _format_sum(out if torch.is_tensor(out) else None) + with path.open("a", encoding="utf-8") as f: + f.write(f"fastvideo:output:{label}:out_sum={out_sum}\n") + return _hook + + if hasattr(self.model, "proj_out"): + self.model.proj_out.register_forward_hook(_output_hook("proj_out")) + if hasattr(self.model, "audio_proj_out"): + self.model.audio_proj_out.register_forward_hook( + _output_hook("audio_proj_out")) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + guidance=None, + audio_hidden_states: torch.Tensor | None = None, + audio_encoder_hidden_states: torch.Tensor | None = None, + audio_timestep: torch.Tensor | None = None, + audio_encoder_attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + if isinstance(encoder_hidden_states, list): + encoder_hidden_states = encoder_hidden_states[0] + + video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) + positions = self.patchifier.get_patch_grid_bounds( + video_shape, device=hidden_states.device + ) + fps = None + try: + forward_ctx = get_forward_context() + except AssertionError: + forward_ctx = None + if forward_ctx is not None and forward_ctx.forward_batch is not None: + fps_value = forward_ctx.forward_batch.fps + if isinstance(fps_value, list): + fps_value = fps_value[0] if fps_value else None + if fps_value is not None: + fps = float(fps_value) + positions = _get_pixel_coords( + positions, + DEFAULT_LTX2_SCALE_FACTORS, + fps=fps, + causal_fix=True, + ).to(hidden_states.dtype) + latents = self.patchifier.patchify(hidden_states) + video_modality = Modality( + enabled=True, + latent=latents, + timesteps=timestep, + positions=positions, + context=encoder_hidden_states, + context_mask=encoder_attention_mask, + ) + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": + video_head = latents.flatten()[:8].float().tolist() + video_flat = latents.float().flatten() + video_checksum = (video_flat * torch.arange(video_flat.numel(), device=video_flat.device)).sum().item() + _debug_block_log_line( + "fastvideo:modality_video" + f":latent_sum={latents.float().sum().item():.6f} " + f"latent_shape={tuple(latents.shape)} " + f"positions_sum={positions.float().sum().item():.6f} " + f"positions_shape={tuple(positions.shape)} " + f"latent_head={video_head} " + f"latent_checksum={video_checksum:.6f}" + ) + audio_modality = None + audio_shape = None + if audio_hidden_states is not None and audio_encoder_hidden_states is not None and audio_timestep is not None: + audio_shape = AudioLatentShape.from_torch_shape( + audio_hidden_states.shape) + audio_positions = self.audio_patchifier.get_patch_grid_bounds( + audio_shape, device=audio_hidden_states.device) + audio_latents = self.audio_patchifier.patchify( + audio_hidden_states) + audio_modality = Modality( + enabled=True, + latent=audio_latents, + timesteps=audio_timestep, + positions=audio_positions, + context=audio_encoder_hidden_states, + context_mask=audio_encoder_attention_mask, + ) + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": + audio_head = audio_latents.flatten()[:8].float().tolist() + audio_flat = audio_latents.float().flatten() + audio_checksum = (audio_flat * torch.arange(audio_flat.numel(), device=audio_flat.device)).sum().item() + _debug_block_log_line( + "fastvideo:modality_audio" + f":latent_sum={audio_latents.float().sum().item():.6f} " + f"latent_shape={tuple(audio_latents.shape)} " + f"positions_sum={audio_positions.float().sum().item():.6f} " + f"positions_shape={tuple(audio_positions.shape)} " + f"latent_head={audio_head} " + f"latent_checksum={audio_checksum:.6f}" + ) + video_out, audio_out = self.model( + video=video_modality, + audio=audio_modality, + ) + if video_out is not None and video_modality is not None: + video_out = _to_denoised( + video_modality.latent, + video_out, + video_modality.timesteps, + ) + if audio_out is not None and audio_modality is not None: + audio_out = _to_denoised( + audio_modality.latent, + audio_out, + audio_modality.timesteps, + ) + video_out = self.patchifier.unpatchify( + video_out, output_shape=video_shape) + if audio_out is None or audio_shape is None: + return video_out + audio_out = self.audio_patchifier.unpatchify( + audio_out, output_shape=audio_shape) + return video_out, audio_out diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py new file mode 100644 index 000000000..a279e1669 --- /dev/null +++ b/fastvideo/models/encoders/gemma.py @@ -0,0 +1,563 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from dataclasses import dataclass +import os +from typing import Iterable + +import torch +from torch import nn +from transformers import Gemma3ForConditionalGeneration + +from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.dits.ltx2 import ( + FeedForward, + LTXRopeType, + apply_ltx_rotary_emb, + generate_ltx_freq_grid_np, + generate_ltx_freq_grid_pytorch, + precompute_ltx_freqs_cis, +) +from fastvideo.models.loader.weight_utils import default_weight_loader +from fastvideo.platforms import AttentionBackendEnum + + +def _debug_log_line(message: str) -> None: + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1": + return + log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "") + if not log_path: + return + log_dir = os.path.dirname(log_path) + if log_dir: + os.makedirs(log_dir, exist_ok=True) + with open(log_path, "a", encoding="utf-8") as f: + f.write(message + "\n") + + +def _debug_gemma_log_line(message: str) -> None: + log_path = os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", "") + if not log_path: + return + log_dir = os.path.dirname(log_path) + if log_dir: + os.makedirs(log_dir, exist_ok=True) + with open(log_path, "a", encoding="utf-8") as f: + f.write(message + "\n") + + +@dataclass(frozen=True) +class GemmaConnectorConfig: + num_attention_heads: int + attention_head_dim: int + num_layers: int + positional_embedding_theta: float + positional_embedding_max_pos: list[int] + rope_type: LTXRopeType + double_precision_rope: bool + num_learnable_registers: int | None + + +class GemmaFeaturesExtractorProjLinear(nn.Module): + """Linear projection that aggregates stacked Gemma hidden states.""" + + def __init__(self, in_features: int, out_features: int) -> None: + super().__init__() + self.aggregate_embed = nn.Linear(in_features, out_features, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.aggregate_embed(x) + + +class _BasicTransformerBlock1D(nn.Module): + """1D transformer block for connector processing.""" + + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + rope_type: LTXRopeType, + norm_eps: float = 1e-6, + ) -> None: + super().__init__() + self.attn1 = _GemmaAttention( + query_dim=dim, + context_dim=None, + heads=heads, + dim_head=dim_head, + norm_eps=norm_eps, + rope_type=rope_type, + ) + self.ff = FeedForward(dim, dim_out=dim) + self.norm_eps = norm_eps + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + pe: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + norm_hidden_states = torch.nn.functional.rms_norm( + hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps + ) + if norm_hidden_states.ndim == 4: + norm_hidden_states = norm_hidden_states.squeeze(1) + + attn_output = self.attn1( + norm_hidden_states, + mask=attention_mask, + pe=pe, + ) + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + norm_hidden_states = torch.nn.functional.rms_norm( + hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps + ) + ff_output = self.ff(norm_hidden_states) + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + return hidden_states + + +class _GemmaAttention(nn.Module): + """Attention implementation aligned with LTX-2 text encoder.""" + + def __init__( + self, + query_dim: int, + context_dim: int | None, + heads: int, + dim_head: int, + norm_eps: float, + rope_type: LTXRopeType, + ) -> None: + super().__init__() + inner_dim = dim_head * heads + context_dim = query_dim if context_dim is None else context_dim + + self.heads = heads + self.dim_head = dim_head + self.rope_type = rope_type + + self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + self.to_q = nn.Linear(query_dim, inner_dim, bias=True) + self.to_k = nn.Linear(context_dim, inner_dim, bias=True) + self.to_v = nn.Linear(context_dim, inner_dim, bias=True) + self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + pe: tuple[torch.Tensor, torch.Tensor] | None = None, + k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + q = self.to_q(x) + context = x if context is None else context + k = self.to_k(context) + v = self.to_v(context) + + q = self.q_norm(q) + k = self.k_norm(k) + + if pe is not None: + q = apply_ltx_rotary_emb(q, pe, self.rope_type) + k = apply_ltx_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) + + b, q_len, _ = q.shape + k_len = k.shape[1] + q = q.view(b, q_len, self.heads, self.dim_head).transpose(1, 2) + k = k.view(b, k_len, self.heads, self.dim_head).transpose(1, 2) + v = v.view(b, k_len, self.heads, self.dim_head).transpose(1, 2) + + if mask is not None: + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(1) + + out = torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + ) + out = out.transpose(1, 2).reshape(b, q_len, -1) + return self.to_out(out) + + +class Embeddings1DConnector(nn.Module): + """Transformer connector that refines Gemma embeddings for LTX-2.""" + + _supports_gradient_checkpointing = True + + def __init__(self, config: GemmaConnectorConfig) -> None: + super().__init__() + self.num_attention_heads = config.num_attention_heads + self.inner_dim = config.num_attention_heads * config.attention_head_dim + self.positional_embedding_theta = config.positional_embedding_theta + self.positional_embedding_max_pos = config.positional_embedding_max_pos + self.rope_type = config.rope_type + self.double_precision_rope = config.double_precision_rope + self.transformer_1d_blocks = nn.ModuleList( + [ + _BasicTransformerBlock1D( + dim=self.inner_dim, + heads=config.num_attention_heads, + dim_head=config.attention_head_dim, + rope_type=config.rope_type, + ) + for _ in range(config.num_layers) + ] + ) + self.num_learnable_registers = config.num_learnable_registers + if self.num_learnable_registers: + self.learnable_registers = nn.Parameter( + torch.rand( + self.num_learnable_registers, + self.inner_dim, + dtype=torch.bfloat16, + ) + * 2.0 + - 1.0 + ) + + def _replace_padded_with_learnable_registers( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert hidden_states.shape[1] % self.num_learnable_registers == 0, ( + f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by " + f"num_learnable_registers {self.num_learnable_registers}." + ) + + num_registers_duplications = ( + hidden_states.shape[1] // self.num_learnable_registers + ) + learnable_registers = torch.tile( + self.learnable_registers, (num_registers_duplications, 1) + ) + attention_mask_binary = ( + attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0 + ).int() + + non_zero_hidden_states = hidden_states[ + :, attention_mask_binary.squeeze().bool(), : + ] + non_zero_nums = non_zero_hidden_states.shape[1] + pad_length = hidden_states.shape[1] - non_zero_nums + adjusted_hidden_states = torch.nn.functional.pad( + non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0 + ) + flipped_mask = torch.flip(attention_mask_binary, dims=[1]) + hidden_states = flipped_mask * adjusted_hidden_states + ( + 1 - flipped_mask + ) * learnable_registers + + attention_mask = torch.full_like( + attention_mask, + 0.0, + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + return hidden_states, attention_mask + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.num_learnable_registers: + hidden_states, attention_mask = ( + self._replace_padded_with_learnable_registers( + hidden_states, attention_mask + ) + ) + + indices_grid = torch.arange( + hidden_states.shape[1], + dtype=torch.float32, + device=hidden_states.device, + ) + indices_grid = indices_grid[None, None, :] + freq_grid_generator = ( + generate_ltx_freq_grid_np + if self.double_precision_rope + else generate_ltx_freq_grid_pytorch + ) + freqs_cis = precompute_ltx_freqs_cis( + indices_grid=indices_grid, + dim=self.inner_dim, + out_dtype=hidden_states.dtype, + theta=self.positional_embedding_theta, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + rope_type=self.rope_type, + freq_grid_generator=freq_grid_generator, + ) + + for block in self.transformer_1d_blocks: + hidden_states = block( + hidden_states, attention_mask=attention_mask, pe=freqs_cis + ) + + hidden_states = torch.nn.functional.rms_norm( + hidden_states, (hidden_states.shape[-1],), eps=1e-6 + ) + + return hidden_states, attention_mask + + +class LTX2GemmaTextEncoderModel(TextEncoder): + _supported_attention_backends = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + + def __init__(self, config: TextEncoderConfig) -> None: + super().__init__(config) + arch = config.arch_config + + self.feature_extractor_linear = GemmaFeaturesExtractorProjLinear( + in_features=arch.feature_extractor_in_features, + out_features=arch.feature_extractor_out_features, + ) + + connector_config = GemmaConnectorConfig( + num_attention_heads=arch.connector_num_attention_heads, + attention_head_dim=arch.connector_attention_head_dim, + num_layers=arch.connector_num_layers, + positional_embedding_theta=arch.connector_positional_embedding_theta, + positional_embedding_max_pos=arch.connector_positional_embedding_max_pos, + rope_type=LTXRopeType(arch.connector_rope_type), + double_precision_rope=arch.connector_double_precision_rope, + num_learnable_registers=arch.connector_num_learnable_registers, + ) + self.embeddings_connector = Embeddings1DConnector(connector_config) + self.audio_embeddings_connector = Embeddings1DConnector(connector_config) + + 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 + + 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."): + continue + yield name, param + + @property + def gemma_model(self) -> Gemma3ForConditionalGeneration: + 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( + gemma_path, + local_files_only=True, + torch_dtype=dtype, + ) + # Configure model-level attention implementation when using TORCH_SDPA. + # Note: torch.backends.cuda.enable_*_sdp() settings should be configured + # at application/pipeline initialization level, not here, to avoid + # unexpected side effects across the application. + if os.getenv("FASTVIDEO_ATTENTION_BACKEND") == "TORCH_SDPA": + if hasattr(self._gemma_model.config, "attn_implementation"): + self._gemma_model.config.attn_implementation = "sdpa" + if hasattr(self._gemma_model.config, "_attn_implementation"): + self._gemma_model.config._attn_implementation = "sdpa" + device = next(self.feature_extractor_linear.parameters()).device + self._gemma_model.to(device=device) + self._gemma_model.eval() + return self._gemma_model + + def _run_feature_extractor( + self, + hidden_states: tuple[torch.Tensor, ...], + attention_mask: torch.Tensor, + padding_side: str, + ) -> torch.Tensor: + encoded_text_features = torch.stack(hidden_states, dim=-1) + if os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", ""): + for idx, layer in enumerate(hidden_states): + _debug_gemma_log_line( + f"fastvideo:gemma_hidden_state_{idx}" + f":sum={layer.float().sum().item():.6f}" + ) + _debug_gemma_log_line( + "fastvideo:gemma_hidden_states_stack" + f":sum={encoded_text_features.float().sum().item():.6f}" + ) + encoded_text_features_dtype = encoded_text_features.dtype + sequence_lengths = attention_mask.sum(dim=-1) + normed_text_features = _norm_and_concat_padded_batch( + encoded_text_features, sequence_lengths, padding_side=padding_side + ) + return self.feature_extractor_linear( + normed_text_features.to(encoded_text_features_dtype) + ) + + def _convert_to_additive_mask( + self, attention_mask: torch.Tensor, dtype: torch.dtype + ) -> torch.Tensor: + return (attention_mask - 1).to(dtype).reshape( + (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) + ) * torch.finfo(dtype).max + + def _run_connectors( + self, + encoded_input: torch.Tensor, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + connector_attention_mask = self._convert_to_additive_mask( + attention_mask, encoded_input.dtype + ) + encoded, encoded_connector_attention_mask = self.embeddings_connector( + encoded_input, connector_attention_mask + ) + + attention_mask = (encoded_connector_attention_mask < 0.000001).to( + torch.int64 + ) + attention_mask = attention_mask.reshape( + [encoded.shape[0], encoded.shape[1], 1] + ) + encoded = encoded * attention_mask + + encoded_for_audio, _ = self.audio_embeddings_connector( + encoded_input, connector_attention_mask + ) + + return encoded, encoded_for_audio, attention_mask.squeeze(-1) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + if input_ids is None: + raise ValueError("input_ids is required for Gemma text encoding.") + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) + + model = self.gemma_model + input_ids = input_ids.to(device=model.device) + attention_mask = attention_mask.to(device=model.device) + outputs = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + ) + + encoded_inputs = self._run_feature_extractor( + outputs.hidden_states, + attention_mask, + padding_side=self.padding_side, + ) + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": + _debug_log_line( + "fastvideo:gemma_feature" + f":sum={encoded_inputs.float().sum().item():.6f} " + f"shape={tuple(encoded_inputs.shape)}" + ) + video_encoding, audio_encoding, attention_mask = self._run_connectors( + encoded_inputs, attention_mask + ) + if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": + _debug_log_line( + "fastvideo:gemma_video_encoding" + f":sum={video_encoding.float().sum().item():.6f} " + f"shape={tuple(video_encoding.shape)}" + ) + _debug_log_line( + "fastvideo:gemma_audio_encoding" + f":sum={audio_encoding.float().sum().item():.6f} " + f"shape={tuple(audio_encoding.shape)}" + ) + + hidden_states = (audio_encoding, ) if output_hidden_states else None + return BaseEncoderOutput( + last_hidden_state=video_encoding, + hidden_states=hidden_states, + attention_mask=attention_mask, + ) + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> set[str]: + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if name == "aggregate_embed.weight": + name = "feature_extractor_linear.aggregate_embed.weight" + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +def _norm_and_concat_padded_batch( + encoded_text: torch.Tensor, + sequence_lengths: torch.Tensor, + padding_side: str = "right", +) -> torch.Tensor: + b, t, d, l = encoded_text.shape + device = encoded_text.device + + token_indices = torch.arange(t, device=device)[None, :] + if padding_side == "right": + mask = token_indices < sequence_lengths[:, None] + elif padding_side == "left": + start_indices = t - sequence_lengths[:, None] + mask = token_indices >= start_indices + else: + raise ValueError( + f"padding_side must be 'left' or 'right', got {padding_side}" + ) + + mask = mask.reshape(b, t, 1, 1) + eps = 1e-6 + + masked = encoded_text.masked_fill(~mask, 0.0) + denom = (sequence_lengths * d).view(b, 1, 1, 1) + mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps) + + x_min = encoded_text.masked_fill(~mask, float("inf")).amin( + dim=(1, 2), keepdim=True + ) + x_max = encoded_text.masked_fill(~mask, float("-inf")).amax( + dim=(1, 2), keepdim=True + ) + range_ = x_max - x_min + + normed = 8 * (encoded_text - mean) / (range_ + eps) + normed = normed.reshape(b, t, -1) + + mask_flattened = mask.reshape(b, t, 1).expand(-1, -1, d * l) + normed = normed.masked_fill(~mask_flattened, 0.0) + return normed diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index d61c574fa..37c542375 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -80,6 +80,9 @@ class ComponentLoader(ABC): "transformer": (TransformerLoader, "diffusers"), "transformer_2": (TransformerLoader, "diffusers"), "vae": (VAELoader, "diffusers"), + "audio_vae": (AudioDecoderLoader, "diffusers"), + "audio_decoder": (AudioDecoderLoader, "diffusers"), + "vocoder": (VocoderLoader, "diffusers"), "text_encoder": (TextEncoderLoader, "transformers"), "text_encoder_2": (TextEncoderLoader, "transformers"), "tokenizer": (TokenizerLoader, "transformers"), @@ -242,6 +245,47 @@ class TextEncoderLoader(ComponentLoader): model_config.pop("model_type", None) model_config.pop("tokenizer_class", None) model_config.pop("torch_dtype", None) + repo_root = os.path.dirname(model_path) + index_path = os.path.join(repo_root, "model_index.json") + gemma_path = "" + gemma_path_from_candidate = False + if os.path.isfile(index_path): + try: + with open(index_path, encoding="utf-8") as f: + model_index = json.load(f) + gemma_path = model_index.get("gemma_model_path", "") + except json.JSONDecodeError: + gemma_path = "" + if not gemma_path: + candidate = os.path.normpath(os.path.join(model_path, "gemma")) + if os.path.isdir(candidate): + gemma_path = candidate + gemma_path_from_candidate = True + model_config["gemma_model_path"] = gemma_path + 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) + ) + transformer_config_path = os.path.join( + repo_root, "transformer", "config.json" + ) + if os.path.isfile(transformer_config_path): + try: + with open(transformer_config_path, encoding="utf-8") as f: + transformer_config = json.load(f) + if ( + "connector_double_precision_rope" not in model_config + or not model_config["connector_double_precision_rope"] + ): + if transformer_config.get("double_precision_rope") is True: + model_config["connector_double_precision_rope"] = True + if "connector_rope_type" not in model_config: + rope_type = transformer_config.get("rope_type") + if rope_type is not None: + model_config["connector_rope_type"] = rope_type + except json.JSONDecodeError: + pass logger.info("HF Model config: %s", model_config) # @TODO(Wei): Better way to handle this? @@ -489,8 +533,20 @@ class TokenizerLoader(ComponentLoader): # in v0, this was same string as encoder_name "ClipTextModel" # TODO(will): pass these tokenizer kwargs from inference args? Maybe # other method of config? - padding_size="right", ) + padding_side = None + if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"): + try: + arch_config = fastvideo_args.pipeline_config.text_encoder_configs[ + 0 + ].arch_config + padding_side = getattr(arch_config, "padding_side", None) + except Exception: + padding_side = None + if padding_side: + tokenizer.padding_side = padding_side + if tokenizer.pad_token is None and tokenizer.eos_token is not None: + tokenizer.pad_token = tokenizer.eos_token logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__) return tokenizer @@ -501,15 +557,12 @@ class VAELoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the VAE based on the model path, and inference args.""" config = get_diffusers_config(model=model_path) - class_name = config.pop("_class_name") + class_name = config.get("_class_name") assert class_name is not None, ( "Model config does not contain a _class_name attribute. Only diffusers format is supported." ) fastvideo_args.model_paths["vae"] = model_path - vae_config = fastvideo_args.pipeline_config.vae_config - vae_config.update_model_arch(config) - from fastvideo.platforms import current_platform if fastvideo_args.vae_cpu_offload: @@ -543,8 +596,29 @@ class VAELoader(ComponentLoader): vae.load_state_dict(sd, strict=False) return vae.eval() - vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) - vae = vae_cls(vae_config).to(target_device) + # LTX-2 uses CausalVideoAutoencoder with nested "vae" config + if class_name == "CausalVideoAutoencoder" and "vae" in config: + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vae = vae_cls(config).to(target_device) + if hasattr(vae, "set_tiling_config"): + vae_config = fastvideo_args.pipeline_config.vae_config + vae.set_tiling_config( + spatial_tile_size_in_pixels=getattr( + vae_config, "ltx2_spatial_tile_size_in_pixels", 512), + spatial_tile_overlap_in_pixels=getattr( + vae_config, "ltx2_spatial_tile_overlap_in_pixels", 64), + temporal_tile_size_in_frames=getattr( + vae_config, "ltx2_temporal_tile_size_in_frames", 64), + temporal_tile_overlap_in_frames=getattr( + vae_config, + "ltx2_temporal_tile_overlap_in_frames", 24), + ) + else: + config.pop("_class_name", None) + vae_config = fastvideo_args.pipeline_config.vae_config + vae_config.update_model_arch(config) + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vae = vae_cls(vae_config).to(target_device) # Find all safetensors files safetensors_list = glob.glob( @@ -553,17 +627,101 @@ class VAELoader(ComponentLoader): raise ValueError(f"No safetensors files found in {model_path}") # Common case: a single `.safetensors` checkpoint file. # Some models may be sharded into multiple files; in that case we merge. - if len(safetensors_list) == 1: - loaded = safetensors_load_file(safetensors_list[0]) - else: - loaded = {} - for sf_file in safetensors_list: - loaded.update(safetensors_load_file(sf_file)) + loaded = {} + for sf_file in safetensors_list: + loaded.update(safetensors_load_file(sf_file)) + + # LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping + if class_name == "CausalVideoAutoencoder" and "vae" in config: + per_channel_prefixes = ( + "per_channel_statistics.", + "vae.per_channel_statistics.", + ) + remapped = {} + for key, tensor in loaded.items(): + remapped[key] = tensor + for prefix in per_channel_prefixes: + if key.startswith(prefix): + suffix = key[len(prefix):] + remapped.setdefault( + f"encoder.per_channel_statistics.{suffix}", + tensor, + ) + remapped.setdefault( + f"decoder.per_channel_statistics.{suffix}", + tensor, + ) + break + loaded = remapped + vae.load_state_dict(loaded, strict=False) return vae.eval() +class AudioDecoderLoader(ComponentLoader): + """Loader for LTX-2 audio decoder (audio_vae component).""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + config = get_diffusers_config(model=model_path) + class_name = config.pop("_class_name", None) or "LTX2AudioDecoder" + + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + target_device = get_local_torch_device() + + precision = getattr( + fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16" + ) + with set_default_torch_dtype(PRECISION_TO_TYPE[precision]): + audio_decoder = model_cls(config).to(target_device) + + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors") + ) + loaded: dict[str, torch.Tensor] = {} + for sf_file in safetensors_list: + loaded.update(safetensors_load_file(sf_file)) + + decoder_state = {} + for name, tensor in loaded.items(): + if name.startswith("decoder."): + decoder_state[name.replace("decoder.", "")] = tensor + elif name.startswith("per_channel_statistics."): + decoder_state[name] = tensor + + target_module = getattr(audio_decoder, "model", audio_decoder) + target_module.load_state_dict(decoder_state, strict=False) + return audio_decoder.eval() + + +class VocoderLoader(ComponentLoader): + """Loader for LTX-2 vocoder.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + config = get_diffusers_config(model=model_path) + class_name = config.pop("_class_name", None) or "LTX2Vocoder" + + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + target_device = get_local_torch_device() + + precision = getattr( + fastvideo_args.pipeline_config, "vocoder_precision", "bf16" + ) + with set_default_torch_dtype(PRECISION_TO_TYPE[precision]): + vocoder = model_cls(config).to(target_device) + + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors") + ) + loaded: dict[str, torch.Tensor] = {} + for sf_file in safetensors_list: + loaded.update(safetensors_load_file(sf_file)) + + target_module = getattr(vocoder, "model", vocoder) + target_module.load_state_dict(loaded, strict=False) + return vocoder.eval() + + class TransformerLoader(ComponentLoader): """Loader for transformer.""" @@ -679,7 +837,18 @@ class TransformerLoader(ComponentLoader): model = model.eval() if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload: - enable_layerwise_offload(model) + # Check if model has nn.ModuleList for layerwise offload compatibility + has_module_list = any( + isinstance(m, nn.ModuleList) for m in model.children() + ) + if has_module_list: + enable_layerwise_offload(model) + else: + logger.warning( + "Layerwise offload requested but model %s does not have " + "nn.ModuleList structure. Skipping layerwise offload.", + cls_name + ) return model diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 4fe9b1f39..2bf86ab3a 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -33,6 +33,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"), "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1) "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2) + "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"), } _IMAGE_TO_VIDEO_DIT_MODELS = { @@ -54,6 +55,7 @@ _TEXT_ENCODER_MODELS = { "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"), "Qwen2_5_VLForConditionalGeneration": ("encoders", "reason1", "Reason1TextEncoder"), + "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"), } _IMAGE_ENCODER_MODELS: dict[str, tuple] = { @@ -67,7 +69,14 @@ _VAE_MODELS = { ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"), "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"), - "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo") + "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"), + "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"), +} + +_AUDIO_MODELS = { + "LTX2AudioEncoder": ("audio", "ltx2_audio_vae", "LTX2AudioEncoder"), + "LTX2AudioDecoder": ("audio", "ltx2_audio_vae", "LTX2AudioDecoder"), + "LTX2Vocoder": ("audio", "ltx2_audio_vae", "LTX2Vocoder"), } _SCHEDULERS = { @@ -91,6 +100,7 @@ _FAST_VIDEO_MODELS = { **_TEXT_ENCODER_MODELS, **_IMAGE_ENCODER_MODELS, **_VAE_MODELS, + **_AUDIO_MODELS, **_SCHEDULERS, } diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py new file mode 100644 index 000000000..e0a01db83 --- /dev/null +++ b/fastvideo/models/vaes/ltx2vae.py @@ -0,0 +1,1804 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 Video VAE implementation +""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, replace +from enum import Enum +from typing import Any, Callable, Iterator, List, NamedTuple, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fastvideo.models.vaes.common import DiagonalGaussianDistribution + + +# ============================================================================= +# Enums +# ============================================================================= + + +class NormLayerType(Enum): + GROUP_NORM = "group_norm" + PIXEL_NORM = "pixel_norm" + + +class LogVarianceType(Enum): + PER_CHANNEL = "per_channel" + UNIFORM = "uniform" + CONSTANT = "constant" + NONE = "none" + + +class PaddingModeType(Enum): + ZEROS = "zeros" + REFLECT = "reflect" + REPLICATE = "replicate" + CIRCULAR = "circular" + + +# ============================================================================= +# Tiling Data Structures +# ============================================================================= + + +@dataclass(frozen=True) +class SpatialTilingConfig: + """Configuration for dividing each frame into spatial tiles with optional overlap.""" + tile_size_in_pixels: int + tile_overlap_in_pixels: int = 0 + + def __post_init__(self) -> None: + if self.tile_size_in_pixels < 64: + raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}") + if self.tile_size_in_pixels % 32 != 0: + raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}") + if self.tile_overlap_in_pixels % 32 != 0: + raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}") + if self.tile_overlap_in_pixels >= self.tile_size_in_pixels: + raise ValueError( + f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}" + ) + + +@dataclass(frozen=True) +class TemporalTilingConfig: + """Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap.""" + tile_size_in_frames: int + tile_overlap_in_frames: int = 0 + + def __post_init__(self) -> None: + if self.tile_size_in_frames < 16: + raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}") + if self.tile_size_in_frames % 8 != 0: + raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}") + if self.tile_overlap_in_frames % 8 != 0: + raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}") + if self.tile_overlap_in_frames >= self.tile_size_in_frames: + raise ValueError( + f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}" + ) + + +@dataclass(frozen=True) +class TilingConfig: + """Configuration for splitting video into tiles with optional overlap.""" + spatial_config: SpatialTilingConfig | None = None + temporal_config: TemporalTilingConfig | None = None + + @classmethod + def default(cls) -> "TilingConfig": + return cls( + spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64), + temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24), + ) + + +@dataclass(frozen=True) +class DimensionIntervals: + """Intervals which a single dimension of the latent space is split into.""" + starts: List[int] + ends: List[int] + left_ramps: List[int] + right_ramps: List[int] + + +@dataclass(frozen=True) +class LatentIntervals: + """Intervals which the latent tensor of given shape is split into.""" + original_shape: torch.Size + dimension_intervals: Tuple[DimensionIntervals, ...] + + +class VideoLatentShape(NamedTuple): + """Shape of the tensor representing video in VAE latent space.""" + batch: int + channels: int + frames: int + height: int + width: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.height, self.width]) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "VideoLatentShape": + return VideoLatentShape( + batch=shape[0], + channels=shape[1], + frames=shape[2], + height=shape[3], + width=shape[4], + ) + + def upscale(self, time_scale: int, spatial_scale: int) -> "VideoLatentShape": + return self._replace( + channels=3, + frames=(self.frames - 1) * time_scale + 1, + height=self.height * spatial_scale, + width=self.width * spatial_scale, + ) + + +# Operation to split a single dimension of the tensor into intervals based on the length along the dimension. +SplitOperation = Callable[[int], DimensionIntervals] +# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension. +MappingOperation = Callable[[DimensionIntervals], Tuple[List[slice], List[torch.Tensor | None]]] + + +class Tile(NamedTuple): + """Represents a single tile.""" + in_coords: Tuple[slice, ...] + out_coords: Tuple[slice, ...] + masks_1d: Tuple[torch.Tensor | None, ...] + + @property + def blend_mask(self) -> torch.Tensor: + num_dims = len(self.out_coords) + per_dimension_masks: List[torch.Tensor] = [] + + for dim_idx in range(num_dims): + mask_1d = self.masks_1d[dim_idx] + view_shape = [1] * num_dims + if mask_1d is None: + one = torch.ones(1) + view_shape[dim_idx] = 1 + per_dimension_masks.append(one.view(*view_shape)) + continue + + view_shape[dim_idx] = mask_1d.shape[0] + per_dimension_masks.append(mask_1d.view(*view_shape)) + + combined_mask = per_dimension_masks[0] + for mask in per_dimension_masks[1:]: + combined_mask = combined_mask * mask + + return combined_mask + + +# ============================================================================= +# Tiling Helper Functions +# ============================================================================= + + +def compute_trapezoidal_mask_1d( + length: int, + ramp_left: int, + ramp_right: int, + left_starts_from_0: bool = False, +) -> torch.Tensor: + """Generate a 1D trapezoidal blending mask with linear ramps.""" + if length <= 0: + raise ValueError("Mask length must be positive.") + + ramp_left = max(0, min(ramp_left, length)) + ramp_right = max(0, min(ramp_right, length)) + + mask = torch.ones(length) + + if ramp_left > 0: + interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2 + fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1] + if not left_starts_from_0: + fade_in = fade_in[1:] + mask[:ramp_left] *= fade_in + + if ramp_right > 0: + fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1] + mask[-ramp_right:] *= fade_out + + return mask.clamp_(0, 1) + + +def default_split_operation(length: int) -> DimensionIntervals: + return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0]) + + +DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation + + +def default_mapping_operation( + _intervals: DimensionIntervals, +) -> Tuple[List[slice], List[torch.Tensor | None]]: + return [slice(0, None)], [None] + + +DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation + + +def split_in_spatial(size: int, overlap: int) -> SplitOperation: + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap) + starts = [i * (size - overlap) for i in range(amount)] + ends = [start + size for start in starts] + ends[-1] = dimension_size + left_ramps = [0] + [overlap] * (amount - 1) + right_ramps = [overlap] * (amount - 1) + [0] + return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps) + + return split + + +def split_in_temporal(size: int, overlap: int) -> SplitOperation: + non_causal_split = split_in_spatial(size, overlap) + + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + intervals = non_causal_split(dimension_size) + starts = list(intervals.starts) + starts[1:] = [s - 1 for s in starts[1:]] + left_ramps = list(intervals.left_ramps) + left_ramps[1:] = [r + 1 for r in left_ramps[1:]] + return replace(intervals, starts=starts, left_ramps=left_ramps) + + return split + + +def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]: + start = begin * scale + stop = 1 + (end - 1) * scale + left_ramp_scaled = 1 + (left_ramp - 1) * scale + right_ramp_scaled = right_ramp * scale + + return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp_scaled, right_ramp_scaled, True) + + +def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]: + start = begin * scale + stop = end * scale + left_ramp_scaled = left_ramp * scale + right_ramp_scaled = right_ramp * scale + + return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp_scaled, right_ramp_scaled, False) + + +def to_mapping_operation( + map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]], + scale: int, +) -> MappingOperation: + def map_op(intervals: DimensionIntervals) -> Tuple[List[slice], List[torch.Tensor | None]]: + output_slices: List[slice] = [] + masks_1d: List[torch.Tensor | None] = [] + number_of_slices = len(intervals.starts) + for i in range(number_of_slices): + start = intervals.starts[i] + end = intervals.ends[i] + left_ramp = intervals.left_ramps[i] + right_ramp = intervals.right_ramps[i] + output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale) + output_slices.append(output_slice) + masks_1d.append(mask_1d) + return output_slices, masks_1d + + return map_op + + +def create_tiles_from_intervals_and_mappers( + intervals: LatentIntervals, + mappers: List[MappingOperation], +) -> List[Tile]: + full_dim_input_slices = [] + full_dim_output_slices = [] + full_dim_masks_1d = [] + for axis_index in range(len(intervals.original_shape)): + dimension_intervals = intervals.dimension_intervals[axis_index] + starts = dimension_intervals.starts + ends = dimension_intervals.ends + input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)] + output_slices, masks_1d = mappers[axis_index](dimension_intervals) + full_dim_input_slices.append(input_slices) + full_dim_output_slices.append(output_slices) + full_dim_masks_1d.append(masks_1d) + + tiles = [] + tile_in_coords = list(itertools.product(*full_dim_input_slices)) + tile_out_coords = list(itertools.product(*full_dim_output_slices)) + tile_mask_1ds = list(itertools.product(*full_dim_masks_1d)) + for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True): + tiles.append( + Tile( + in_coords=in_coord, + out_coords=out_coord, + masks_1d=mask_1d, + ) + ) + return tiles + + +def create_tiles( + latent_shape: torch.Size, + splitters: List[SplitOperation], + mappers: List[MappingOperation], +) -> List[Tile]: + if len(splitters) != len(latent_shape): + raise ValueError( + f"Number of splitters must be equal to number of dimensions in latent shape, " + f"got {len(splitters)} and {len(latent_shape)}" + ) + if len(mappers) != len(latent_shape): + raise ValueError( + f"Number of mappers must be equal to number of dimensions in latent shape, " + f"got {len(mappers)} and {len(latent_shape)}" + ) + intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)] + latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals)) + return create_tiles_from_intervals_and_mappers(latent_intervals, mappers) + + +# ============================================================================= +# Utility Functions +# ============================================================================= + + +def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange spatial dimensions into channels (space-to-depth). + + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. + patch_size_t: Temporal patch size. Default=1 (no temporal patching). + + For 5D: (B, C, F, H, W) -> (B, C*patch_size_hw^2*patch_size_t, F/patch_size_t, H/patch_size_hw, W/patch_size_hw) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b c (f p) (h q) (w r) -> b (c p r q) f h w", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange channels back into spatial dimensions (depth-to-space). + + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. + patch_size_t: Temporal patch size. Default=1 (no temporal expansion). + + For 5D: (B, C*patch_size_hw^2*patch_size_t, F, H, W) -> (B, C, F*patch_size_t, H*patch_size_hw, W*patch_size_hw) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b (c p r q) f h w -> b c (f p) (h q) (w r)", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + return x + + +# ============================================================================= +# Normalization Layers +# ============================================================================= + + +class PixelNorm(nn.Module): + """ + Per-pixel (per-location) RMS normalization layer. + Normalizes along the channel dimension using root-mean-square. + """ + + def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) + return x / torch.sqrt(mean_sq + self.eps) + + +# ============================================================================= +# Per-Channel Statistics for Latent Normalization +# ============================================================================= + + +class PerChannelStatistics(nn.Module): + """ + Per-channel statistics for normalizing and denormalizing the latent representation. + Statistics are computed over the dataset and stored in the model checkpoint. + """ + + def __init__(self, latent_channels: int = 128): + super().__init__() + self.register_buffer("std-of-means", torch.ones(latent_channels)) + self.register_buffer("mean-of-means", torch.zeros(latent_channels)) + self.register_buffer("mean-of-stds", torch.ones(latent_channels)) + self.register_buffer("mean-of-stds_over_std-of-means", torch.ones(latent_channels)) + self.register_buffer("channel", torch.arange(latent_channels, dtype=torch.float32)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x) + mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x) + return x * std + mean + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x) + mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x) + return (x - mean) / std + + +# ============================================================================= +# Causal 3D Convolution +# ============================================================================= + + +class CausalConv3d(nn.Module): + """ + Causal 3D convolution that pads temporally by repeating the first frame. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 3, + stride: int | Tuple[int, int, int] = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ) -> None: + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + + kernel_size = (kernel_size, kernel_size, kernel_size) + self.time_kernel_size = kernel_size[0] + + dilation = (dilation, 1, 1) + + height_pad = kernel_size[1] // 2 + width_pad = kernel_size[2] // 2 + padding = (0, height_pad, width_pad) + + self.conv = nn.Conv3d( + in_channels, + out_channels, + kernel_size, + stride=stride, + dilation=dilation, + padding=padding, + padding_mode=spatial_padding_mode.value, + groups=groups, + bias=bias, + ) + + def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor: + if causal: + first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) + x = torch.cat((first_frame_pad, x), dim=2) + else: + first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) + last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) + x = torch.cat((first_frame_pad, x, last_frame_pad), dim=2) + x = self.conv(x) + return x + + @property + def weight(self) -> torch.Tensor: + return self.conv.weight + + +def make_conv_nd( + dims: int, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int | Tuple[int, int, int] = 1, + padding: int = 0, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, +) -> nn.Module: + """Create a convolution layer (2D or 3D, causal or not).""" + if dims == 2: + return nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + elif dims == 3: + if causal: + return CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + groups=groups, + bias=bias, + spatial_padding_mode=spatial_padding_mode, + ) + return nn.Conv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + else: + raise ValueError(f"unsupported dimensions: {dims}") + + +def make_linear_nd( + dims: int, + in_channels: int, + out_channels: int, + bias: bool = True, +) -> nn.Module: + """Create a 1x1 convolution (pointwise linear).""" + if dims == 2: + return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) + elif dims == 3: + return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) + else: + raise ValueError(f"unsupported dimensions: {dims}") + + +# ============================================================================= +# ResNet Blocks +# ============================================================================= + + +class ResnetBlock3D(nn.Module): + """A 3D ResNet block with optional timestep conditioning and noise injection.""" + + def __init__( + self, + dims: int, + in_channels: int, + out_channels: int | None = None, + dropout: float = 0.0, + groups: int = 32, + eps: float = 1e-6, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.inject_noise = inject_noise + + if norm_layer == NormLayerType.GROUP_NORM: + self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.norm1 = PixelNorm() + + self.non_linearity = nn.SiLU() + + self.conv1 = make_conv_nd( + dims, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + + if inject_noise: + self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1))) + + if norm_layer == NormLayerType.GROUP_NORM: + self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.norm2 = PixelNorm() + + self.dropout = nn.Dropout(dropout) + + self.conv2 = make_conv_nd( + dims, + out_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + + if inject_noise: + self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1))) + + self.conv_shortcut = ( + make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels) + if in_channels != out_channels + else nn.Identity() + ) + + self.norm3 = ( + nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True) + if in_channels != out_channels + else nn.Identity() + ) + + self.timestep_conditioning = timestep_conditioning + + if timestep_conditioning: + self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5) + + def _feed_spatial_noise( + self, + hidden_states: torch.Tensor, + per_channel_scale: torch.Tensor, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + spatial_shape = hidden_states.shape[-2:] + device = hidden_states.device + dtype = hidden_states.dtype + + spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None] + scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...] + hidden_states = hidden_states + scaled_noise + return hidden_states + + def forward( + self, + input_tensor: torch.Tensor, + causal: bool = True, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + hidden_states = input_tensor + batch_size = hidden_states.shape[0] + + hidden_states = self.norm1(hidden_states) + if self.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True") + ada_values = self.scale_shift_table[None, ..., None, None, None].to( + device=hidden_states.device, dtype=hidden_states.dtype + ) + timestep.reshape( + batch_size, + 4, + -1, + timestep.shape[-3], + timestep.shape[-2], + timestep.shape[-1], + ) + shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1) + hidden_states = hidden_states * (1 + scale1) + shift1 + + hidden_states = self.non_linearity(hidden_states) + hidden_states = self.conv1(hidden_states, causal=causal) + + if self.inject_noise: + hidden_states = self._feed_spatial_noise( + hidden_states, + self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype), + generator=generator, + ) + + hidden_states = self.norm2(hidden_states) + + if self.timestep_conditioning: + hidden_states = hidden_states * (1 + scale2) + shift2 + + hidden_states = self.non_linearity(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states, causal=causal) + + if self.inject_noise: + hidden_states = self._feed_spatial_noise( + hidden_states, + self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype), + generator=generator, + ) + + input_tensor = self.norm3(input_tensor) + input_tensor = self.conv_shortcut(input_tensor) + output_tensor = input_tensor + hidden_states + return output_tensor + + +class UNetMidBlock3D(nn.Module): + """A 3D UNet mid-block with multiple residual blocks.""" + + def __init__( + self, + dims: int, + in_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_groups: int = 32, + norm_layer: NormLayerType = NormLayerType.GROUP_NORM, + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + attention_head_dim: int | None = None, # unused, for compatibility + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + self.timestep_conditioning = timestep_conditioning + + if timestep_conditioning: + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim=in_channels * 4, size_emb_dim=0 + ) + + self.res_blocks = nn.ModuleList( + [ + ResnetBlock3D( + dims=dims, + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + norm_layer=norm_layer, + inject_noise=inject_noise, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + causal: bool = True, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + timestep_embed = None + if self.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True") + batch_size = hidden_states.shape[0] + timestep_embed = self.time_embedder( + timestep=timestep.flatten(), + hidden_dtype=hidden_states.dtype, + ) + timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1) + + for resnet in self.res_blocks: + hidden_states = resnet( + hidden_states, + causal=causal, + timestep=timestep_embed, + generator=generator, + ) + + return hidden_states + + +# ============================================================================= +# Timestep Embedding (simplified version for decoder conditioning) +# ============================================================================= + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): + """Timestep embeddings for decoder conditioning.""" + + def __init__(self, embedding_dim: int, size_emb_dim: int = 0): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + return timesteps_emb + + +class Timesteps(nn.Module): + """Sinusoidal timestep embeddings.""" + + def __init__(self, num_channels: int, flip_sin_to_cos: bool = True, downscale_freq_shift: float = 0): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + half_dim = self.num_channels // 2 + exponent = -math.log(10000) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - self.downscale_freq_shift) + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + if self.flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + return emb + + +class TimestepEmbedding(nn.Module): + """MLP for timestep embeddings.""" + + def __init__(self, in_channels: int, time_embed_dim: int): + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample = self.linear_1(sample) + sample = self.act(sample) + sample = self.linear_2(sample) + return sample + + +# ============================================================================= +# Sampling Blocks (Downsampling / Upsampling) +# ============================================================================= + + +class SpaceToDepthDownsample(nn.Module): + """Downsampling via space-to-depth with residual connection.""" + + def __init__( + self, + dims: int, + in_channels: int, + out_channels: int, + stride: Tuple[int, int, int], + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + self.stride = stride + self.group_size = in_channels * math.prod(stride) // out_channels + self.conv = make_conv_nd( + dims=dims, + in_channels=in_channels, + out_channels=out_channels // math.prod(stride), + kernel_size=3, + stride=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + + def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor: + if self.stride[0] == 2: + x = torch.cat([x[:, :, :1, :, :], x], dim=2) + + x_in = rearrange( + x, + "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size) + x_in = x_in.mean(dim=2) + + x = self.conv(x, causal=causal) + x = rearrange( + x, + "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + + x = x + x_in + return x + + +class DepthToSpaceUpsample(nn.Module): + """Upsampling via depth-to-space (pixel shuffle).""" + + def __init__( + self, + dims: int, + in_channels: int, + stride: Tuple[int, int, int], + residual: bool = False, + out_channels_reduction_factor: int = 1, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + self.stride = stride + self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor + self.conv = make_conv_nd( + dims=dims, + in_channels=in_channels, + out_channels=self.out_channels, + kernel_size=3, + stride=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + self.residual = residual + self.out_channels_reduction_factor = out_channels_reduction_factor + + def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor: + if self.residual: + x_in = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor + x_in = x_in.repeat(1, num_repeat, 1, 1, 1) + if self.stride[0] == 2: + x_in = x_in[:, :, 1:, :, :] + + x = self.conv(x, causal=causal) + x = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + if self.stride[0] == 2: + x = x[:, :, 1:, :, :] + if self.residual: + x = x + x_in + return x + + +# ============================================================================= +# Block Factory Functions +# ============================================================================= + + +def _make_encoder_block( + block_name: str, + block_config: dict[str, Any], + in_channels: int, + convolution_dimensions: int, + norm_layer: NormLayerType, + norm_num_groups: int, + spatial_padding_mode: PaddingModeType, +) -> Tuple[nn.Module, int]: + """Create an encoder block based on the block name and config.""" + out_channels = in_channels + + if block_name == "res_x": + block = UNetMidBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + num_layers=block_config["num_layers"], + resnet_eps=1e-6, + resnet_groups=norm_num_groups, + norm_layer=norm_layer, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "res_x_y": + out_channels = in_channels * block_config.get("multiplier", 2) + block = ResnetBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + eps=1e-6, + groups=norm_num_groups, + norm_layer=norm_layer, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_time": + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 1, 1), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_space": + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(1, 2, 2), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all": + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 2, 2), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all_x_y": + out_channels = in_channels * block_config.get("multiplier", 2) + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 2, 2), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all_res": + out_channels = in_channels * block_config.get("multiplier", 2) + block = SpaceToDepthDownsample( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + stride=(2, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_space_res": + out_channels = in_channels * block_config.get("multiplier", 2) + block = SpaceToDepthDownsample( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + stride=(1, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_time_res": + out_channels = in_channels * block_config.get("multiplier", 2) + block = SpaceToDepthDownsample( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + stride=(2, 1, 1), + spatial_padding_mode=spatial_padding_mode, + ) + else: + raise ValueError(f"unknown block: {block_name}") + + return block, out_channels + + +def _make_decoder_block( + block_name: str, + block_config: dict[str, Any], + in_channels: int, + convolution_dimensions: int, + norm_layer: NormLayerType, + timestep_conditioning: bool, + norm_num_groups: int, + spatial_padding_mode: PaddingModeType, +) -> Tuple[nn.Module, int]: + """Create a decoder block based on the block name and config.""" + out_channels = in_channels + + if block_name == "res_x": + block = UNetMidBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + num_layers=block_config["num_layers"], + resnet_eps=1e-6, + resnet_groups=norm_num_groups, + norm_layer=norm_layer, + inject_noise=block_config.get("inject_noise", False), + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "attn_res_x": + block = UNetMidBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + num_layers=block_config["num_layers"], + resnet_groups=norm_num_groups, + norm_layer=norm_layer, + inject_noise=block_config.get("inject_noise", False), + timestep_conditioning=timestep_conditioning, + attention_head_dim=block_config["attention_head_dim"], + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "res_x_y": + out_channels = in_channels // block_config.get("multiplier", 2) + block = ResnetBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + eps=1e-6, + groups=norm_num_groups, + norm_layer=norm_layer, + inject_noise=block_config.get("inject_noise", False), + timestep_conditioning=False, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_time": + block = DepthToSpaceUpsample( + dims=convolution_dimensions, + in_channels=in_channels, + stride=(2, 1, 1), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_space": + block = DepthToSpaceUpsample( + dims=convolution_dimensions, + in_channels=in_channels, + stride=(1, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all": + out_channels = in_channels // block_config.get("multiplier", 1) + block = DepthToSpaceUpsample( + dims=convolution_dimensions, + in_channels=in_channels, + stride=(2, 2, 2), + residual=block_config.get("residual", False), + out_channels_reduction_factor=block_config.get("multiplier", 1), + spatial_padding_mode=spatial_padding_mode, + ) + else: + raise ValueError(f"unknown layer: {block_name}") + + return block, out_channels + + +# ============================================================================= +# Video Encoder +# ============================================================================= + + +class VideoEncoder(nn.Module): + """ + LTX-2 Video Encoder. Encodes video frames into a latent representation. + """ + + _DEFAULT_NORM_NUM_GROUPS = 32 + + def __init__( + self, + convolution_dimensions: int = 3, + in_channels: int = 3, + out_channels: int = 128, + encoder_blocks: list[tuple[str, int]] | list[tuple[str, dict[str, Any]]] = [], + patch_size: int = 4, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + latent_log_var: LogVarianceType = LogVarianceType.UNIFORM, + encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + + self.patch_size = patch_size + self.norm_layer = norm_layer + self.latent_channels = out_channels + self.latent_log_var = latent_log_var + self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS + + self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels) + + in_channels = in_channels * patch_size**2 + feature_channels = out_channels + + self.conv_in = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=feature_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=encoder_spatial_padding_mode, + ) + + self.down_blocks = nn.ModuleList([]) + + for block_name, block_params in encoder_blocks: + block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params + + block, feature_channels = _make_encoder_block( + block_name=block_name, + block_config=block_config, + in_channels=feature_channels, + convolution_dimensions=convolution_dimensions, + norm_layer=norm_layer, + norm_num_groups=self._norm_num_groups, + spatial_padding_mode=encoder_spatial_padding_mode, + ) + + self.down_blocks.append(block) + + if norm_layer == NormLayerType.GROUP_NORM: + self.conv_norm_out = nn.GroupNorm( + num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6 + ) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.conv_norm_out = PixelNorm() + + self.conv_act = nn.SiLU() + + conv_out_channels = out_channels + if latent_log_var == LogVarianceType.PER_CHANNEL: + conv_out_channels *= 2 + elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}: + conv_out_channels += 1 + elif latent_log_var != LogVarianceType.NONE: + raise ValueError(f"Invalid latent_log_var: {latent_log_var}") + + self.conv_out = make_conv_nd( + dims=convolution_dimensions, + in_channels=feature_channels, + out_channels=conv_out_channels, + kernel_size=3, + padding=1, + causal=True, + spatial_padding_mode=encoder_spatial_padding_mode, + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + frames_count = sample.shape[2] + if ((frames_count - 1) % 8) != 0: + raise ValueError( + "Invalid number of frames: Encode input must have 1 + 8 * x frames " + "(e.g., 1, 9, 17, ...). Please check your input." + ) + + sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) + sample = self.conv_in(sample) + + for down_block in self.down_blocks: + sample = down_block(sample) + + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if self.latent_log_var == LogVarianceType.UNIFORM: + if sample.shape[1] < 2: + raise ValueError( + f"Invalid channel count for UNIFORM mode: expected at least 2, got {sample.shape[1]}" + ) + means = sample[:, :-1, ...] + logvar = sample[:, -1:, ...] + num_channels = means.shape[1] + repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2) + repeated_logvar = logvar.repeat(*repeat_shape) + sample = torch.cat([means, repeated_logvar], dim=1) + elif self.latent_log_var == LogVarianceType.CONSTANT: + sample = sample[:, :-1, ...] + approx_ln_0 = -30 + sample = torch.cat( + [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0], + dim=1, + ) + + means, _ = torch.chunk(sample, 2, dim=1) + return self.per_channel_statistics.normalize(means) + + +# ============================================================================= +# Video Decoder +# ============================================================================= + + +class VideoDecoder(nn.Module): + """ + LTX-2 Video Decoder. Decodes latent representation into video frames. + """ + + _DEFAULT_NORM_NUM_GROUPS = 32 + + def __init__( + self, + convolution_dimensions: int = 3, + in_channels: int = 128, + out_channels: int = 3, + decoder_blocks: list[tuple[str, int | dict]] = [], + patch_size: int = 4, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + causal: bool = False, + timestep_conditioning: bool = False, + decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT, + ): + super().__init__() + + self.patch_size = patch_size + out_channels = out_channels * patch_size**2 + self.causal = causal + self.timestep_conditioning = timestep_conditioning + self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS + + self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) + + self.decode_noise_scale = 0.025 + self.decode_timestep = 0.05 + + feature_channels = in_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": + feature_channels = feature_channels * block_config.get("multiplier", 2) + if block_name == "compress_all": + feature_channels = feature_channels * block_config.get("multiplier", 1) + + self.conv_in = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=feature_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=decoder_spatial_padding_mode, + ) + + self.up_blocks = nn.ModuleList([]) + + for block_name, block_params in list(reversed(decoder_blocks)): + block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params + + block, feature_channels = _make_decoder_block( + block_name=block_name, + block_config=block_config, + in_channels=feature_channels, + convolution_dimensions=convolution_dimensions, + norm_layer=norm_layer, + timestep_conditioning=timestep_conditioning, + norm_num_groups=self._norm_num_groups, + spatial_padding_mode=decoder_spatial_padding_mode, + ) + + self.up_blocks.append(block) + + if norm_layer == NormLayerType.GROUP_NORM: + self.conv_norm_out = nn.GroupNorm( + num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6 + ) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.conv_norm_out = PixelNorm() + + self.conv_act = nn.SiLU() + self.conv_out = make_conv_nd( + dims=convolution_dimensions, + in_channels=feature_channels, + out_channels=out_channels, + kernel_size=3, + padding=1, + causal=True, + spatial_padding_mode=decoder_spatial_padding_mode, + ) + + if timestep_conditioning: + self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0)) + self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim=feature_channels * 2, size_emb_dim=0 + ) + self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels)) + + def forward( + self, + sample: torch.Tensor, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + batch_size = sample.shape[0] + + if self.timestep_conditioning: + noise = ( + torch.randn( + sample.size(), + generator=generator, + dtype=sample.dtype, + device=sample.device, + ) + * self.decode_noise_scale + ) + sample = noise + (1.0 - self.decode_noise_scale) * sample + + sample = self.per_channel_statistics.un_normalize(sample) + + if timestep is None and self.timestep_conditioning: + timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype) + + sample = self.conv_in(sample, causal=self.causal) + + scaled_timestep = None + if self.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True") + scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample) + + for up_block in self.up_blocks: + if isinstance(up_block, UNetMidBlock3D): + block_kwargs = { + "causal": self.causal, + "timestep": scaled_timestep if self.timestep_conditioning else None, + "generator": generator, + } + sample = up_block(sample, **block_kwargs) + elif isinstance(up_block, ResnetBlock3D): + sample = up_block(sample, causal=self.causal, generator=generator) + else: + sample = up_block(sample, causal=self.causal) + + sample = self.conv_norm_out(sample) + + if self.timestep_conditioning: + embedded_timestep = self.last_time_embedder( + timestep=scaled_timestep.flatten(), + hidden_dtype=sample.dtype, + ) + embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1) + ada_values = self.last_scale_shift_table[None, ..., None, None, None].to( + device=sample.device, dtype=sample.dtype + ) + embedded_timestep.reshape( + batch_size, + 2, + -1, + embedded_timestep.shape[-3], + embedded_timestep.shape[-2], + embedded_timestep.shape[-1], + ) + shift, scale = ada_values.unbind(dim=1) + sample = sample * (1 + scale) + shift + + sample = self.conv_act(sample) + sample = self.conv_out(sample, causal=self.causal) + sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) + + return sample + + +# ============================================================================= +# Configurators (for loading from config dict) +# ============================================================================= + + +class VideoEncoderConfigurator: + """Configurator for creating a video VAE Encoder from a configuration dictionary.""" + + @classmethod + def from_config(cls, config: dict) -> VideoEncoder: + config = config.get("vae", config) + convolution_dimensions = config.get("dims", 3) + in_channels = config.get("in_channels", 3) + latent_channels = config.get("latent_channels", 128) + encoder_spatial_padding_mode = PaddingModeType(config.get("encoder_spatial_padding_mode", "zeros")) + encoder_blocks = config.get("encoder_blocks", []) + patch_size = config.get("patch_size", 4) + norm_layer_str = config.get("norm_layer", "pixel_norm") + latent_log_var_str = config.get("latent_log_var", "uniform") + + return VideoEncoder( + convolution_dimensions=convolution_dimensions, + in_channels=in_channels, + out_channels=latent_channels, + encoder_blocks=encoder_blocks, + patch_size=patch_size, + norm_layer=NormLayerType(norm_layer_str), + latent_log_var=LogVarianceType(latent_log_var_str), + encoder_spatial_padding_mode=encoder_spatial_padding_mode, + ) + + +class VideoDecoderConfigurator: + """Configurator for creating a video VAE Decoder from a configuration dictionary.""" + + @classmethod + def from_config(cls, config: dict) -> VideoDecoder: + config = config.get("vae", config) + convolution_dimensions = config.get("dims", 3) + latent_channels = config.get("latent_channels", 128) + decoder_spatial_padding_mode = PaddingModeType(config.get("decoder_spatial_padding_mode", "reflect")) + out_channels = config.get("out_channels", 3) + decoder_blocks = config.get("decoder_blocks", []) + patch_size = config.get("patch_size", 4) + norm_layer_str = config.get("norm_layer", "pixel_norm") + causal = config.get("causal_decoder", False) + timestep_conditioning = config.get("timestep_conditioning", True) + + return VideoDecoder( + convolution_dimensions=convolution_dimensions, + in_channels=latent_channels, + out_channels=out_channels, + decoder_blocks=decoder_blocks, + patch_size=patch_size, + norm_layer=NormLayerType(norm_layer_str), + causal=causal, + timestep_conditioning=timestep_conditioning, + decoder_spatial_padding_mode=decoder_spatial_padding_mode, + ) + + +# ============================================================================= +# Public API (Wrapper Classes) +# ============================================================================= + + +class LTX2VideoEncoder(nn.Module): + """LTX-2 Video Encoder wrapper for FastVideo compatibility.""" + + def __init__(self, config: dict[str, Any]): + super().__init__() + self.model: VideoEncoder = VideoEncoderConfigurator.from_config(config) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + return self.model(sample) + + +class LTX2VideoDecoder(nn.Module): + """LTX-2 Video Decoder wrapper for FastVideo compatibility.""" + + def __init__(self, config: dict[str, Any]): + super().__init__() + self.model: VideoDecoder = VideoDecoderConfigurator.from_config(config) + + def forward( + self, + sample: torch.Tensor, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + return self.model(sample, timestep=timestep, generator=generator) + + +class LTX2CausalVideoAutoencoder(nn.Module): + """ + LTX-2 VAE that exposes FastVideo's VAE encode/decode interface. + Supports tiled decoding to reduce memory usage for high-resolution videos. + """ + + # LTX-2 VAE scale factors + TIME_SCALE: int = 8 + SPATIAL_SCALE: int = 32 + + def __init__(self, config: dict[str, Any]): + super().__init__() + self.config = config + self.encoder = VideoEncoderConfigurator.from_config(config) + self.decoder = VideoDecoderConfigurator.from_config(config) + self._use_tiling: bool = False + + def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: + means = self.encoder(x) + zeros = torch.zeros_like(means) + return DiagonalGaussianDistribution(torch.cat([means, zeros], dim=1), deterministic=True) + + def decode( + self, + z: torch.Tensor, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Decode latents to video, using tiling if enabled.""" + if self._use_tiling: + # Collect all chunks from tiled decode and concatenate + chunks = list(self.tiled_decode(z, TilingConfig.default(), timestep, generator)) + return torch.cat(chunks, dim=2) # Concatenate along temporal dimension + return self.decoder(z, timestep=timestep, generator=generator) + + def enable_tiling(self) -> None: + """Enable tiled decoding with default configuration.""" + self._use_tiling = True + + def disable_tiling(self) -> None: + """Disable tiled decoding.""" + self._use_tiling = False + + def _prepare_tiles( + self, + latent: torch.Tensor, + tiling_config: TilingConfig | None, + ) -> List[Tile]: + """Prepare tiles for tiled decoding based on tiling configuration.""" + splitters: List[SplitOperation] = [DEFAULT_SPLIT_OPERATION] * 5 + mappers: List[MappingOperation] = [DEFAULT_MAPPING_OPERATION] * 5 + + if tiling_config is not None and tiling_config.spatial_config is not None: + cfg = tiling_config.spatial_config + tile_size = cfg.tile_size_in_pixels // self.SPATIAL_SCALE + overlap = cfg.tile_overlap_in_pixels // self.SPATIAL_SCALE + splitters[3] = split_in_spatial(tile_size, overlap) + splitters[4] = split_in_spatial(tile_size, overlap) + mappers[3] = to_mapping_operation(map_spatial_slice, self.SPATIAL_SCALE) + mappers[4] = to_mapping_operation(map_spatial_slice, self.SPATIAL_SCALE) + + if tiling_config is not None and tiling_config.temporal_config is not None: + cfg = tiling_config.temporal_config + tile_size = cfg.tile_size_in_frames // self.TIME_SCALE + overlap = cfg.tile_overlap_in_frames // self.TIME_SCALE + splitters[2] = split_in_temporal(tile_size, overlap) + mappers[2] = to_mapping_operation(map_temporal_slice, self.TIME_SCALE) + + return create_tiles(latent.shape, splitters, mappers) + + def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]: + """Group tiles by their temporal output slice.""" + if not tiles: + return [] + + groups = [] + current_slice = tiles[0].out_coords[2] + current_group = [] + + for tile in tiles: + tile_slice = tile.out_coords[2] + if tile_slice == current_slice: + current_group.append(tile) + else: + groups.append(current_group) + current_slice = tile_slice + current_group = [tile] + + if current_group: + groups.append(current_group) + + return groups + + def _accumulate_temporal_group_into_buffer( + self, + group_tiles: List[Tile], + buffer: torch.Tensor, + latent: torch.Tensor, + timestep: torch.Tensor | None, + generator: torch.Generator | None, + ) -> torch.Tensor: + """Decode and accumulate all tiles of a temporal group into a local buffer.""" + temporal_slice = group_tiles[0].out_coords[2] + weights = torch.zeros_like(buffer) + + for tile in group_tiles: + decoded_tile = self.decoder(latent[tile.in_coords], timestep, generator) + mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype) + temporal_offset = tile.out_coords[2].start - temporal_slice.start + expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start + decoded_temporal_len = decoded_tile.shape[2] + + actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset) + + chunk_coords = ( + slice(None), # batch + slice(None), # channels + slice(temporal_offset, temporal_offset + actual_temporal_len), + tile.out_coords[3], # height + tile.out_coords[4], # width + ) + + decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :] + mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask + + buffer[chunk_coords] += decoded_slice * mask_slice + weights[chunk_coords] += mask_slice + + return weights + + def tiled_decode( + self, + latent: torch.Tensor, + tiling_config: TilingConfig | None = None, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + """ + Decode a latent tensor into video frames using tiled processing. + Splits the latent tensor into tiles, decodes each tile individually, + and yields video chunks as they become available. + + Args: + latent: Input latent tensor (B, C, F', H', W'). + tiling_config: Tiling configuration for the latent tensor. + timestep: Optional timestep for decoder conditioning. + generator: Optional random generator for deterministic decoding. + + Yields: + Video chunks (B, C, T, H, W) by temporal slices. + """ + full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale( + self.TIME_SCALE, self.SPATIAL_SCALE + ) + tiles = self._prepare_tiles(latent, tiling_config) + temporal_groups = self._group_tiles_by_temporal_slice(tiles) + + previous_chunk = None + previous_weights = None + previous_temporal_slice = None + + for temporal_group_tiles in temporal_groups: + curr_temporal_slice = temporal_group_tiles[0].out_coords[2] + + temporal_tile_buffer_shape = full_video_shape._replace( + frames=curr_temporal_slice.stop - curr_temporal_slice.start, + ) + + buffer = torch.zeros( + temporal_tile_buffer_shape.to_torch_shape(), + device=latent.device, + dtype=latent.dtype, + ) + + curr_weights = self._accumulate_temporal_group_into_buffer( + group_tiles=temporal_group_tiles, + buffer=buffer, + latent=latent, + timestep=timestep, + generator=generator, + ) + + # Blend with previous temporal chunk if it exists + if previous_chunk is not None: + if previous_temporal_slice.stop > curr_temporal_slice.start: + overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start + temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None) + + previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :] + previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[ + :, :, slice(0, overlap_len), :, : + ] + + buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :] + curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[ + :, :, temporal_overlap_slice, :, : + ] + + # Yield the non-overlapping part of the previous chunk + previous_weights = previous_weights.clamp(min=1e-8) + yield_len = curr_temporal_slice.start - previous_temporal_slice.start + yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :] + + # Update state for next iteration + previous_chunk = buffer + previous_weights = curr_weights + previous_temporal_slice = curr_temporal_slice + + # Yield any remaining chunk + if previous_chunk is not None: + previous_weights = previous_weights.clamp(min=1e-8) + yield previous_chunk / previous_weights diff --git a/fastvideo/pipelines/basic/ltx2/__init__.py b/fastvideo/pipelines/basic/ltx2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py new file mode 100644 index 000000000..72f51e735 --- /dev/null +++ b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 text-to-video pipeline. +""" + +import os +from typing import Any + +from transformers import AutoTokenizer + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import PipelineComponentLoader +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import (DecodingStage, InputValidationStage, + LTX2AudioDecodingStage, + LTX2DenoisingStage, + LTX2LatentPreparationStage, + TextEncodingStage) + +logger = init_logger(__name__) + + +class LTX2Pipeline(ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "transformer", + "vae", + "audio_vae", + "vocoder", + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + self.add_stage( + stage_name="input_validation_stage", + stage=InputValidationStage(), + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LTX2LatentPreparationStage( + transformer=self.get_module("transformer"), ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=LTX2DenoisingStage( + transformer=self.get_module("transformer"), ), + ) + + self.add_stage( + stage_name="audio_decoding_stage", + stage=LTX2AudioDecodingStage( + audio_decoder=self.get_module("audio_vae"), + vocoder=self.get_module("vocoder"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae")), + ) + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + tokenizer = self.get_module("tokenizer") + if tokenizer is not None: + tokenizer.padding_side = "left" + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + def load_modules( + self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, Any] | None = None, + ) -> dict[str, Any]: + model_index = self._load_config(self.model_path) + logger.info("Loading pipeline modules from config: %s", model_index) + + model_index.pop("_class_name") + model_index.pop("_diffusers_version") + model_index.pop("workload_type", None) + + if len(model_index) <= 1: + raise ValueError( + "model_index.json must contain at least one pipeline module") + + required_modules = self.required_config_modules + modules: dict[str, Any] = {} + + for module_name, module_spec in model_index.items(): + if not isinstance(module_spec, list) or len(module_spec) < 1: + continue + transformers_or_diffusers = module_spec[0] + if transformers_or_diffusers is None: + if module_name in self.required_config_modules: + self.required_config_modules.remove(module_name) + continue + if module_name not in required_modules: + continue + if loaded_modules is not None and module_name in loaded_modules: + modules[module_name] = loaded_modules[module_name] + continue + + component_model_path = os.path.join(self.model_path, module_name) + if module_name == "tokenizer" and not os.path.isdir( + component_model_path): + gemma_path = os.path.join(self.model_path, "text_encoder", + "gemma") + if os.path.isdir(gemma_path): + component_model_path = gemma_path + else: + raise ValueError( + "Tokenizer directory missing and Gemma weights were not found." + ) + + module = PipelineComponentLoader.load_module( + module_name=module_name, + component_model_path=component_model_path, + transformers_or_diffusers=transformers_or_diffusers, + fastvideo_args=fastvideo_args, + ) + logger.info("Loaded module %s from %s", module_name, + component_model_path) + modules[module_name] = module + + if "tokenizer" in required_modules and "tokenizer" not in modules: + gemma_path = os.path.join(self.model_path, "text_encoder", "gemma") + if os.path.isdir(gemma_path): + modules["tokenizer"] = AutoTokenizer.from_pretrained( + gemma_path, local_files_only=True) + + for module_name in required_modules: + if module_name not in modules or modules[module_name] is None: + raise ValueError( + f"Required module {module_name} was not loaded properly") + + return modules + + +EntryClass = LTX2Pipeline diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index a2f325dd2..a2e027854 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -35,6 +35,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "LongCatPipeline": "longcat", "LongCatImageToVideoPipeline": "longcat", "LongCatVideoContinuationPipeline": "longcat", + "LTX2Pipeline": "ltx2", } _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = { diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index a8136bafe..4a4164f43 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -22,6 +22,10 @@ from fastvideo.pipelines.stages.input_validation import InputValidationStage from fastvideo.pipelines.stages.latent_preparation import ( Cosmos25LatentPreparationStage, CosmosLatentPreparationStage, LatentPreparationStage) +from fastvideo.pipelines.stages.ltx2_audio_decoding import LTX2AudioDecodingStage +from fastvideo.pipelines.stages.ltx2_denoising import LTX2DenoisingStage +from fastvideo.pipelines.stages.ltx2_latent_preparation import ( + LTX2LatentPreparationStage) from fastvideo.pipelines.stages.matrixgame_denoising import ( MatrixGameCausalDenoisingStage) from fastvideo.pipelines.stages.stepvideo_encoding import ( @@ -44,6 +48,8 @@ __all__ = [ "LatentPreparationStage", "CosmosLatentPreparationStage", "Cosmos25LatentPreparationStage", + "LTX2LatentPreparationStage", + "LTX2AudioDecodingStage", "ConditioningStage", "DenoisingStage", "DmdDenoisingStage", @@ -51,6 +57,7 @@ __all__ = [ "MatrixGameCausalDenoisingStage", "CosmosDenoisingStage", "Cosmos25DenoisingStage", + "LTX2DenoisingStage", "EncodingStage", "DecodingStage", "ImageEncodingStage", diff --git a/fastvideo/pipelines/stages/ltx2_audio_decoding.py b/fastvideo/pipelines/stages/ltx2_audio_decoding.py new file mode 100644 index 000000000..6248ba062 --- /dev/null +++ b/fastvideo/pipelines/stages/ltx2_audio_decoding.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Audio decoding stage for LTX-2 pipelines. +""" + +from __future__ import annotations + +import os + +import torch + +from fastvideo.models.dits.ltx2 import DEFAULT_LTX2_VOCODER_OUTPUT_SAMPLE_RATE + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +class LTX2AudioDecodingStage(PipelineStage): + """Decode LTX-2 audio latents into a waveform.""" + + def __init__(self, audio_decoder, vocoder) -> None: + super().__init__() + self.audio_decoder = audio_decoder + self.vocoder = vocoder + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + audio_latents = batch.extra.get("ltx2_audio_latents") + if audio_latents is None: + return batch + + device = get_local_torch_device() + self.audio_decoder = self.audio_decoder.to(device) + self.vocoder = self.vocoder.to(device) + audio_latents = audio_latents.to(device) + + disable_autocast = os.getenv("LTX2_DISABLE_AUDIO_AUTOCAST", "1") == "1" + with torch.no_grad(), torch.autocast( + device_type="cuda", + dtype=audio_latents.dtype, + enabled=not disable_autocast, + ): + decoded_spec = self.audio_decoder(audio_latents) + audio_wave = self.vocoder(decoded_spec).squeeze(0).float() + + # Move to CPU for pickling across process boundary + batch.extra["audio"] = audio_wave.cpu() + batch.extra[ + "audio_sample_rate"] = DEFAULT_LTX2_VOCODER_OUTPUT_SAMPLE_RATE + return batch + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("audio_latents", batch.extra.get("ltx2_audio_latents"), + V.none_or_tensor) + return result diff --git a/fastvideo/pipelines/stages/ltx2_denoising.py b/fastvideo/pipelines/stages/ltx2_denoising.py new file mode 100644 index 000000000..98d04f6e1 --- /dev/null +++ b/fastvideo/pipelines/stages/ltx2_denoising.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 denoising stage using the native sigma schedule. +""" + +from __future__ import annotations + +import math +import os + +import torch +from tqdm.auto import tqdm + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.logger import init_logger +from fastvideo.models.dits.ltx2 import ( + AudioLatentShape, DEFAULT_LTX2_AUDIO_CHANNELS, + DEFAULT_LTX2_AUDIO_DOWNSAMPLE, DEFAULT_LTX2_AUDIO_HOP_LENGTH, + DEFAULT_LTX2_AUDIO_MEL_BINS, DEFAULT_LTX2_AUDIO_SAMPLE_RATE, + VideoLatentShape) +from fastvideo.utils import PRECISION_TO_TYPE + +BASE_SHIFT_ANCHOR = 1024 +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 +] + +logger = init_logger(__name__) + + +def _ltx2_sigmas( + steps: int, + latent: torch.Tensor | None, + device: torch.device, + max_shift: float = 2.05, + base_shift: float = 0.95, + stretch: bool = True, + terminal: float = 0.1, +) -> torch.Tensor: + tokens = math.prod( + latent.shape[2:]) if latent is not None else MAX_SHIFT_ANCHOR + sigmas = torch.linspace(1.0, + 0.0, + steps + 1, + device=device, + dtype=torch.float32) + + mm = (max_shift - base_shift) / (MAX_SHIFT_ANCHOR - BASE_SHIFT_ANCHOR) + b = base_shift - mm * BASE_SHIFT_ANCHOR + sigma_shift = tokens * mm + b + + numerator = math.exp(sigma_shift) + sigmas = torch.where( + sigmas != 0, + numerator / (numerator + (1 / sigmas - 1)), + torch.zeros_like(sigmas), + ) + + if stretch: + non_zero_mask = sigmas != 0 + non_zero_sigmas = sigmas[non_zero_mask] + one_minus_z = 1.0 - non_zero_sigmas + scale_factor = one_minus_z[-1] / (1.0 - terminal) + stretched = 1.0 - (one_minus_z / scale_factor) + sigmas = sigmas.clone() + sigmas[non_zero_mask] = stretched + + return sigmas + + +class LTX2DenoisingStage(PipelineStage): + """Run the LTX-2 denoising loop over the sigma schedule.""" + + def __init__(self, transformer) -> None: + super().__init__() + self.transformer = transformer + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + if batch.latents is None: + raise ValueError("Latents must be provided before denoising.") + + latents = batch.latents + prompt_embeds = batch.prompt_embeds[0] + prompt_mask = None + + neg_prompt_embeds = None + neg_prompt_mask = None + # Only load negative prompts if CFG is actually enabled + if batch.do_classifier_free_guidance: + assert batch.negative_prompt_embeds is not None, ( + "CFG is enabled but negative_prompt_embeds is None") + neg_prompt_embeds = batch.negative_prompt_embeds[0] + + # Ensure text conditioning is on the same device as latents. + if prompt_embeds.device != latents.device: + prompt_embeds = prompt_embeds.to(latents.device) + if prompt_mask is not None and prompt_mask.device != latents.device: + prompt_mask = prompt_mask.to(latents.device) + if neg_prompt_embeds is not None and neg_prompt_embeds.device != latents.device: + neg_prompt_embeds = neg_prompt_embeds.to(latents.device) + if neg_prompt_mask is not None and neg_prompt_mask.device != latents.device: + neg_prompt_mask = neg_prompt_mask.to(latents.device) + + target_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.dit_precision] + disable_autocast = os.getenv("LTX2_DISABLE_AUTOCAST", "1") == "1" + autocast_enabled = (target_dtype != torch.float32 + ) and not fastvideo_args.disable_autocast and ( + not disable_autocast) + + # Use official distilled sigma schedule for 8 steps (distilled models) + use_distilled_sigmas = os.getenv("LTX2_USE_DISTILLED_SIGMAS", + "1") == "1" + if use_distilled_sigmas and batch.num_inference_steps == 8: + sigmas = torch.tensor( + DISTILLED_SIGMA_VALUES, + device=latents.device, + dtype=torch.float32, + ) + logger.info("[LTX2] Using official distilled sigma schedule") + else: + sigmas = _ltx2_sigmas( + steps=batch.num_inference_steps, + latent=None, + device=latents.device, + ) + 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 + timestep_template = torch.ones( + (latents.shape[0], token_count), + device=latents.device, + dtype=torch.float32, + ) + audio_prompt_embeds = batch.extra.get("ltx2_audio_prompt_embeds") + audio_neg_embeds = batch.extra.get("ltx2_audio_negative_embeds") + audio_context_p = audio_prompt_embeds[0] if audio_prompt_embeds else None + audio_context_n = audio_neg_embeds[0] if audio_neg_embeds else None + audio_latents = None + audio_timestep_template = None + if audio_context_p is not None: + fps_value = batch.fps + if isinstance(fps_value, list): + fps_value = fps_value[0] if fps_value else None + if fps_value is None: + fps_value = 1.0 + duration = float(batch.num_frames) / float(fps_value) + audio_shape = AudioLatentShape.from_duration( + batch=latents.shape[0], + duration=duration, + channels=DEFAULT_LTX2_AUDIO_CHANNELS, + mel_bins=DEFAULT_LTX2_AUDIO_MEL_BINS, + sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE, + hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH, + audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE, + ) + audio_generator = None + if fastvideo_args.ltx2_initial_latent_path and batch.seed is not None: + audio_generator = torch.Generator( + device=latents.device).manual_seed(batch.seed) + elif batch.generator is not None: + if isinstance(batch.generator, list): + audio_generator = batch.generator[0] + else: + audio_generator = batch.generator + if audio_generator is not None and audio_generator.device.type != latents.device.type: + if batch.seed is None: + audio_generator = torch.Generator(device=latents.device) + else: + audio_generator = torch.Generator( + device=latents.device).manual_seed(batch.seed) + audio_patch_shape = ( + audio_shape.batch, + audio_shape.frames, + audio_shape.channels * audio_shape.mel_bins, + ) + audio_latents_patch = torch.randn( + audio_patch_shape, + generator=audio_generator, + device=latents.device, + dtype=latents.dtype, + ) + if hasattr(self.transformer, "audio_patchifier"): + audio_latents = self.transformer.audio_patchifier.unpatchify( + audio_latents_patch, audio_shape) + else: + audio_latents = audio_latents_patch.view( + audio_shape.batch, + audio_shape.frames, + audio_shape.channels, + audio_shape.mel_bins, + ).permute(0, 2, 1, 3).contiguous() + audio_timestep_template = torch.ones( + (latents.shape[0], audio_shape.frames), + device=latents.device, + dtype=torch.float32, + ) + logger.info( + "[LTX2] Denoising start: steps=%d dtype=%s guidance=%s " + "sigmas_shape=%s latents_shape=%s", + batch.num_inference_steps, + target_dtype, + batch.guidance_scale, + tuple(sigmas.shape), + tuple(latents.shape), + ) + + for step_index in tqdm(range(len(sigmas) - 1)): + sigma = sigmas[step_index] + sigma_next = sigmas[step_index + 1] + timestep = timestep_template * sigma + audio_timestep = (audio_timestep_template * sigma + if audio_timestep_template is not None else None) + + with torch.autocast( + device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled, + ), set_forward_context( + current_timestep=sigma.item(), + attn_metadata=None, + forward_batch=batch, + ): + pos_outputs = self.transformer( + hidden_states=latents.to(target_dtype), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_mask, + timestep=timestep, + audio_hidden_states=audio_latents, + audio_encoder_hidden_states=audio_context_p, + audio_timestep=audio_timestep, + ) + if isinstance(pos_outputs, tuple): + pos_denoised, pos_audio = pos_outputs + else: + pos_denoised = pos_outputs + pos_audio = None + + # Only run negative pass if CFG is enabled + if batch.do_classifier_free_guidance: + neg_outputs = self.transformer( + hidden_states=latents.to(target_dtype), + encoder_hidden_states=neg_prompt_embeds, + encoder_attention_mask=neg_prompt_mask, + timestep=timestep, + audio_hidden_states=audio_latents, + audio_encoder_hidden_states=audio_context_n, + audio_timestep=audio_timestep, + ) + if isinstance(neg_outputs, tuple): + neg_denoised, neg_audio = neg_outputs + else: + neg_denoised = neg_outputs + neg_audio = None + pos_denoised = pos_denoised + (batch.guidance_scale - 1) * ( + pos_denoised - neg_denoised) + if pos_audio is not None and neg_audio is not None: + pos_audio = pos_audio + (batch.guidance_scale - + 1) * (pos_audio - neg_audio) + + sigma_value = sigma.to(torch.float32) if isinstance( + sigma, torch.Tensor) else torch.tensor( + float(sigma), + device=latents.device, + dtype=torch.float32, + ) + dt = sigma_next - sigma + 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 + batch.extra["ltx2_audio_latents"] = audio_latents + logger.info("[LTX2] Denoising done.") + return batch + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, + [V.is_tensor, V.with_dims(5)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("num_inference_steps", batch.num_inference_steps, + V.positive_int) + return result diff --git a/fastvideo/pipelines/stages/ltx2_latent_preparation.py b/fastvideo/pipelines/stages/ltx2_latent_preparation.py new file mode 100644 index 000000000..3459f5f28 --- /dev/null +++ b/fastvideo/pipelines/stages/ltx2_latent_preparation.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Latent preparation stage for LTX-2 pipelines. +""" + +from pathlib import Path + +import torch +from diffusers.utils.torch_utils import randn_tensor + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +class LTX2LatentPreparationStage(PipelineStage): + """Prepare initial LTX-2 latents without relying on a diffusers scheduler.""" + + def __init__(self, transformer) -> None: + super().__init__() + self.transformer = transformer + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + latent_num_frames = self._adjust_video_length(batch, fastvideo_args) + + if not batch.prompt_embeds: + batch_size = 1 + elif isinstance(batch.prompt, list): + batch_size = len(batch.prompt) + elif batch.prompt is not None: + batch_size = 1 + else: + batch_size = batch.prompt_embeds[0].shape[0] + + batch_size *= batch.num_videos_per_prompt + + if not batch.prompt_embeds: + transformer_dtype = next(self.transformer.parameters()).dtype + device = get_local_torch_device() + dummy_prompt = torch.zeros( + batch_size, + 0, + self.transformer.hidden_size, + device=device, + dtype=transformer_dtype, + ) + batch.prompt_embeds = [dummy_prompt] + batch.negative_prompt_embeds = [] + batch.do_classifier_free_guidance = False + + dtype = batch.prompt_embeds[0].dtype + device = get_local_torch_device() + generator = batch.generator + latents = batch.latents + num_frames = latent_num_frames if latent_num_frames is not None else batch.num_frames + height = batch.height + width = batch.width + latent_path = fastvideo_args.ltx2_initial_latent_path + + if height is None or width is None: + raise ValueError("Height and width must be provided") + + spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + if height % spatial_ratio != 0 or width % spatial_ratio != 0: + raise ValueError( + f"Height and width must be divisible by {spatial_ratio} " + f"but are {height} and {width}.") + shape = ( + batch_size, + self.transformer.num_channels_latents, + num_frames, + height // spatial_ratio, + width // spatial_ratio, + ) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, " + f"but requested an effective batch size of {batch_size}.") + + if latents is None: + if latent_path: + loaded_latents = self._load_initial_latent( + latent_path, device, dtype) + if loaded_latents is not None: + latents = loaded_latents + else: + latents = randn_tensor( + shape, + generator=generator, + device=device, + dtype=dtype, + ) + self._save_initial_latent(latent_path, latents) + else: + latents = randn_tensor( + shape, + generator=generator, + device=device, + dtype=dtype, + ) + else: + latents = latents.to(device) + + batch.latents = latents + batch.raw_latent_shape = shape + return batch + + def _adjust_video_length(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> int | None: + if not fastvideo_args.pipeline_config.vae_config.use_temporal_scaling_frames: + return None + temporal_scale_factor = (fastvideo_args.pipeline_config.vae_config. + arch_config.temporal_compression_ratio) + video_length = batch.num_frames + return int((video_length - 1) // temporal_scale_factor + 1) + + def _load_initial_latent( + self, + latent_path: str, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + path = Path(latent_path) + if not path.exists(): + return None + payload = torch.load(path, map_location=device) + if isinstance(payload, dict): + if "video_latent" in payload: + latent = payload["video_latent"] + elif "latent" in payload: + latent = payload["latent"] + else: + latent = None + else: + latent = payload + if not torch.is_tensor(latent): + raise TypeError(f"Expected tensor for initial latent in {path}") + logger.info("[LTX2] Loaded initial latent from %s", path) + return latent.to(device=device, dtype=dtype) + + def _save_initial_latent(self, latent_path: str, + latents: torch.Tensor) -> None: + path = Path(latent_path) + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + return + torch.save({"video_latent": latents.detach().cpu()}, path) + logger.info("[LTX2] Saved initial latent to %s", path) + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check( + "prompt_or_embeds", + None, + lambda _: V.string_or_list_strings(batch.prompt) or not batch. + prompt_embeds or V.list_not_empty(batch.prompt_embeds), + ) + if batch.prompt_embeds: + result.add_check("prompt_embeds", batch.prompt_embeds, + V.list_of_tensors) + result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt, + V.positive_int) + result.add_check("generator", batch.generator, + V.generator_or_list_generators) + result.add_check("num_frames", batch.num_frames, V.positive_int) + result.add_check("height", batch.height, V.positive_int) + result.add_check("width", batch.width, V.positive_int) + result.add_check("latents", batch.latents, V.none_or_tensor) + return result + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, + [V.is_tensor, V.with_dims(5)]) + result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple) + return result diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 74b115ff0..87252fde2 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -11,14 +11,11 @@ from typing import Any from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context -from fastvideo.logger import init_logger from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult -logger = init_logger(__name__) - class TextEncodingStage(PipelineStage): """ @@ -39,6 +36,7 @@ class TextEncodingStage(PipelineStage): super().__init__() self.tokenizers = tokenizers self.text_encoders = text_encoders + self._last_audio_embeds: list[torch.Tensor] | None = None @torch.no_grad() def forward( @@ -70,6 +68,8 @@ class TextEncodingStage(PipelineStage): encoder_index=all_indices, return_attention_mask=True, ) + if self._last_audio_embeds is not None: + batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds for pe in prompt_embeds_list: batch.prompt_embeds.append(pe) @@ -86,6 +86,9 @@ class TextEncodingStage(PipelineStage): encoder_index=all_indices, return_attention_mask=True, ) + if self._last_audio_embeds is not None: + batch.extra[ + "ltx2_audio_negative_embeds"] = self._last_audio_embeds assert batch.negative_prompt_embeds is not None for ne in neg_embeds_list: @@ -184,10 +187,13 @@ class TextEncodingStage(PipelineStage): embeds_list: list[torch.Tensor] = [] attn_masks_list: list[torch.Tensor] = [] + audio_embeds_list: list[torch.Tensor] = [] preprocess_funcs = fastvideo_args.pipeline_config.preprocess_text_funcs postprocess_funcs = fastvideo_args.pipeline_config.postprocess_text_funcs encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs + is_ltx2 = getattr(fastvideo_args.pipeline_config.dit_config, "prefix", + "") == "ltx2" if return_type not in ("list", "dict", "stack"): raise ValueError( @@ -259,6 +265,11 @@ class TextEncodingStage(PipelineStage): except Exception: prompt_embeds, attention_mask = postprocess_func( outputs, attention_mask) + if is_ltx2 and getattr(outputs, "hidden_states", None): + audio_embed = outputs.hidden_states[0] + if dtype is not None: + audio_embed = audio_embed.to(dtype=dtype) + audio_embeds_list.append(audio_embed) if dtype is not None: prompt_embeds = prompt_embeds.to(dtype=dtype) @@ -266,6 +277,7 @@ class TextEncodingStage(PipelineStage): if return_attention_mask: attn_masks_list.append(attention_mask) + self._last_audio_embeds = audio_embeds_list if is_ltx2 else None return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index dbeb49821..714630968 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -67,17 +67,23 @@ def run_test(pytest_command: str): sys.exit(result.returncode) -@app.function(gpu="H100:1", image=image, timeout=1200, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})]) +@app.function(gpu="H100:1", + image=image, + timeout=1200, + secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})], + volumes={"/root/data": model_vol}) def run_encoder_tests(): - run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/encoders -vs") + run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/encoders -vs") -@app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})]) +@app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})], + volumes={"/root/data": model_vol}) def run_vae_tests(): - run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/vaes -vs") + run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/vaes -vs") -@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})]) +@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})], + volumes={"/root/data": model_vol}) def run_transformer_tests(): - run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs") + run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs") @app.function( gpu="L40S:4", @@ -89,13 +95,21 @@ def run_transformer_tests(): def run_ssim_tests(): run_test("export HF_HOME='/root/data/.cache' && export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs") -@app.function(gpu="L40S:4", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})]) +@app.function(gpu="L40S:4", + image=image, + timeout=900, + secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})], + volumes={"/root/data": model_vol}) def run_training_tests(): - run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP") + run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP") -@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})]) +@app.function(gpu="L40S:2", + image=image, + timeout=900, + secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})], + volumes={"/root/data": model_vol}) def run_training_lora_tests(): - run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP") + run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP") @app.function(gpu="H100:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})]) def run_training_tests_VSA(): diff --git a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 new file mode 100644 index 000000000..7acf878af Binary files /dev/null and b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 new file mode 100644 index 000000000..545fb2d69 Binary files /dev/null and b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 differ diff --git a/fastvideo/tests/ssim/test_inference_similarity.py b/fastvideo/tests/ssim/test_inference_similarity.py index ba6b7afb0..6f8b34538 100644 --- a/fastvideo/tests/ssim/test_inference_similarity.py +++ b/fastvideo/tests/ssim/test_inference_similarity.py @@ -80,9 +80,48 @@ WAN_I2V_PARAMS = { "text-encoder-precision": ("fp32",) } +# LTX-2 distilled one-stage params (no refine/upscale) +# Official defaults: height=512, width=768, num_frames=121, fps=24, seed=10 +# Using num_frames=41 for faster CI (still valid: 41 = 8×5 + 1) +LTX2_T2V_PARAMS = { + "num_gpus": 2, + "model_path": "FastVideo/LTX2-Distilled-Diffusers", + "height": 512, + "width": 768, + "num_frames": 41, # Shorter for CI; official default is 121 + "num_inference_steps": 8, # Distilled uses 8 steps + "guidance_scale": 1.0, # No CFG for distilled + "embedded_cfg_scale": 6, + "seed": 1024, + "sp_size": 2, + "tp_size": 1, + "fps": 24, + "neg_prompt": ( + "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, " + "excessive noise, grainy texture, poor lighting, flickering, motion blur, distorted " + "proportions, unnatural skin tones, deformed facial features, asymmetrical face, " + "missing facial features, extra limbs, disfigured hands, wrong hand count, artifacts " + "around text, inconsistent perspective, camera shake, incorrect depth of field, " + "background too sharp, background clutter, distracting reflections, harsh shadows, " + "inconsistent lighting direction, color banding, cartoonish rendering, 3D CGI look, " + "unrealistic materials, uncanny valley effect, incorrect ethnicity, wrong gender, " + "exaggerated expressions, wrong gaze direction, mismatched lip sync, silent or muted " + "audio, distorted voice, robotic voice, echo, background noise, off-sync audio, " + "incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward " + "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, " + "flat lighting, inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." + ), + "ltx2_vae_tiling": True, + "ltx2_vae_spatial_tile_size_in_pixels": 512, + "ltx2_vae_spatial_tile_overlap_in_pixels": 64, + "ltx2_vae_temporal_tile_size_in_frames": 64, + "ltx2_vae_temporal_tile_overlap_in_frames": 24, +} + MODEL_TO_PARAMS = { "FastHunyuan-diffusers": HUNYUAN_PARAMS, "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS, + # "ltx2_diffusers": LTX2_T2V_PARAMS, } I2V_MODEL_TO_PARAMS = { @@ -229,18 +268,26 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id): init_kwargs = { "num_gpus": BASE_PARAMS["num_gpus"], - "flow_shift": BASE_PARAMS["flow_shift"], "sp_size": BASE_PARAMS["sp_size"], "tp_size": BASE_PARAMS["tp_size"], "use_fsdp_inference": True, "dit_cpu_offload": False, "dit_layerwise_offload": False, } + if "flow_shift" in BASE_PARAMS: + init_kwargs["flow_shift"] = BASE_PARAMS["flow_shift"] if BASE_PARAMS.get("vae_sp"): init_kwargs["vae_sp"] = True init_kwargs["vae_tiling"] = True if "text-encoder-precision" in BASE_PARAMS: init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"] + # LTX2-specific VAE tiling parameters + if BASE_PARAMS.get("ltx2_vae_tiling"): + init_kwargs["ltx2_vae_tiling"] = True + init_kwargs["ltx2_vae_spatial_tile_size_in_pixels"] = BASE_PARAMS.get("ltx2_vae_spatial_tile_size_in_pixels", 512) + init_kwargs["ltx2_vae_spatial_tile_overlap_in_pixels"] = BASE_PARAMS.get("ltx2_vae_spatial_tile_overlap_in_pixels", 64) + init_kwargs["ltx2_vae_temporal_tile_size_in_frames"] = BASE_PARAMS.get("ltx2_vae_temporal_tile_size_in_frames", 64) + init_kwargs["ltx2_vae_temporal_tile_overlap_in_frames"] = BASE_PARAMS.get("ltx2_vae_temporal_tile_overlap_in_frames", 24) generation_kwargs = { "num_inference_steps": num_inference_steps, diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index 2e3268334..cc09c5898 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -132,9 +132,13 @@ class MultiprocExecutor(Executor): else: logging_info = None + # Get extra dict (contains audio, etc.) + extra = responses[0].get("extra", {}) + result_batch = ForwardBatch(data_type=forward_batch.data_type, output=output, - logging_info=logging_info) + logging_info=logging_info, + extra=extra) return result_batch @@ -648,7 +652,8 @@ class WorkerMultiprocProc: logging_info = output_batch.logging_info self.pipe.send({ "output_batch": output_batch.output.cpu(), - "logging_info": logging_info + "logging_info": logging_info, + "extra": output_batch.extra, }) else: result = self.worker.execute_method( diff --git a/pyproject.toml b/pyproject.toml index d46492542..a704bd8dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "diffusers>=0.33.1", "torch>=2.9.1", "torchvision", + "torchaudio", # Acceleration & Optimization "accelerate==1.0.1", diff --git a/scripts/checkpoint_conversion/convert_ltx2_weights.py b/scripts/checkpoint_conversion/convert_ltx2_weights.py new file mode 100644 index 000000000..bfcbf58e2 --- /dev/null +++ b/scripts/checkpoint_conversion/convert_ltx2_weights.py @@ -0,0 +1,444 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Convert LTX-2 weights to FastVideo naming conventions and split by component. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import re +import shutil +from collections import OrderedDict +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import load_file, save_file + +try: + from huggingface_hub import snapshot_download +except ImportError: # pragma: no cover - optional dependency + snapshot_download = None + + +PARAM_NAME_MAP: dict[str, str] = { + r"^model\.diffusion_model\.(.*)$": r"\1", +} + +COMPONENT_PREFIXES: dict[str, tuple[str, ...]] = { + "transformer": ("model.diffusion_model.",), + "vae": ("vae.",), + "audio_vae": ("audio_vae.",), + "vocoder": ("vocoder.",), + "text_embedding_projection": ("text_embedding_projection.", "model.text_embedding_projection."), +} + + +def _find_shards(model_path: Path) -> list[Path]: + if model_path.is_file(): + return [model_path] + + index_files = list(model_path.glob("*.safetensors.index.json")) + if index_files: + with index_files[0].open("r", encoding="utf-8") as f: + index = json.load(f) + return sorted({model_path / shard for shard in index["weight_map"].values()}) + return sorted(Path(p) for p in glob.glob(str(model_path / "*.safetensors"))) + + +def _apply_mapping(key: str) -> str: + for pattern, replacement in PARAM_NAME_MAP.items(): + if re.match(pattern, key): + return re.sub(pattern, replacement, key) + return key + + +def _load_weights(shards: list[Path]) -> dict[str, torch.Tensor]: + weights: dict[str, torch.Tensor] = {} + for shard in shards: + weights.update(load_file(str(shard))) + return weights + + +def _read_metadata_config(path: Path) -> dict: + with safe_open(str(path), framework="pt") as f: + metadata = f.metadata() + if not metadata or "config" not in metadata: + return {} + return json.loads(metadata["config"]) + + +def _filter_transformer_config(config: dict) -> dict: + transformer = config.get("transformer", {}) + allowed = { + "num_attention_heads", + "attention_head_dim", + "num_layers", + "cross_attention_dim", + "caption_channels", + "norm_eps", + "attention_type", + "positional_embedding_theta", + "positional_embedding_max_pos", + "timestep_scale_multiplier", + "use_middle_indices_grid", + "rope_type", + "frequencies_precision", + "in_channels", + "out_channels", + "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", + } + filtered = {k: v for k, v in transformer.items() if k in allowed} + if "frequencies_precision" in filtered: + filtered["double_precision_rope"] = filtered["frequencies_precision"] == "float64" + del filtered["frequencies_precision"] + return filtered + + +def _build_text_embedding_projection_config( + gemma_model_path: str = "", +) -> dict: + return { + "architectures": ["LTX2GemmaTextEncoderModel"], + "hidden_size": 3840, + "num_hidden_layers": 48, + "num_attention_heads": 30, + "text_len": 1024, + "pad_token_id": 0, + "eos_token_id": 2, + "gemma_model_path": gemma_model_path, + "gemma_dtype": "bfloat16", + "padding_side": "left", + "feature_extractor_in_features": 3840 * 49, + "feature_extractor_out_features": 3840, + "connector_num_attention_heads": 30, + "connector_attention_head_dim": 128, + "connector_num_layers": 2, + "connector_positional_embedding_theta": 10000.0, + "connector_positional_embedding_max_pos": [4096], + "connector_rope_type": "split", + "connector_double_precision_rope": True, + "connector_num_learnable_registers": 128, + } + + +def _wrap_component_config( + component_name: str, + component_config: dict | None, + class_name: str | None = None, +) -> dict | None: + if component_config is None: + return None + wrapped = {component_name: component_config} + if class_name is not None: + wrapped["_class_name"] = class_name + return wrapped + + +def _split_component_weights(weights: dict[str, torch.Tensor]) -> dict[str, OrderedDict]: + components: dict[str, OrderedDict] = {name: OrderedDict() for name in COMPONENT_PREFIXES} + for key, value in weights.items(): + if key.startswith("model.diffusion_model.audio_embeddings_connector."): + new_key = key.replace("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.") + components["text_embedding_projection"][new_key] = value + continue + if key.startswith("model.diffusion_model.video_embeddings_connector."): + new_key = key.replace("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.") + components["text_embedding_projection"][new_key] = value + continue + + matched = False + for component, prefixes in COMPONENT_PREFIXES.items(): + for prefix in prefixes: + if key.startswith(prefix): + new_key = key[len(prefix):] + components[component][new_key] = value + matched = True + break + if matched: + break + return {name: weights for name, weights in components.items() if weights} + + +def _write_component( + output_dir: Path, + name: str, + weights: OrderedDict, + config: dict | None, + dir_name: str | None = None, +) -> None: + component_dir = output_dir / (dir_name or name) + component_dir.mkdir(parents=True, exist_ok=True) + output_file = component_dir / "model.safetensors" + save_file(weights, str(output_file)) + print(f"Saved {name} weights to {output_file}") + + if config is not None: + config_path = component_dir / "config.json" + with config_path.open("w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + f.write("\n") + print(f"Saved {name} config to {config_path}") + + +def _build_model_index( + transformer_class_name: str, + vae_class_name: str, + pipeline_class_name: str, + diffusers_version: str, +) -> dict: + return { + "_class_name": pipeline_class_name, + "_diffusers_version": diffusers_version, + "transformer": ["diffusers", transformer_class_name], + "vae": ["diffusers", vae_class_name], + "text_encoder": ["transformers", "LTX2GemmaTextEncoderModel"], + "tokenizer": ["transformers", "AutoTokenizer"], + "audio_vae": ["diffusers", "LTX2AudioDecoder"], + "vocoder": ["diffusers", "LTX2Vocoder"], + } + + +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" + with model_index_path.open("w", encoding="utf-8") as f: + json.dump(model_index, f, indent=2) + f.write("\n") + print(f"Saved model_index.json to {model_index_path}") + + +def convert_components( + source_path: Path, + output_dir: Path, + metadata_config: dict, + transformer_class_name: str, + components_to_write: set[str] | None = None, + emit_diffusers_repo: bool = True, + pipeline_class_name: str = "LTX2Pipeline", + diffusers_version: str = "0.33.0.dev0", + gemma_model_path: str = "", +) -> None: + shards = _find_shards(source_path) + if not shards: + raise FileNotFoundError(f"No safetensors found in {source_path}") + + weights = _load_weights(shards) + split_weights = _split_component_weights(weights) + if components_to_write is not None: + split_weights = {name: weights for name, weights in split_weights.items() if name in components_to_write} + + transformer_weights = split_weights.get("transformer", OrderedDict()) + converted_transformer = OrderedDict() + for key, value in transformer_weights.items(): + new_key = _apply_mapping(f"model.diffusion_model.{key}") + converted_transformer[new_key] = value + split_weights["transformer"] = converted_transformer + + transformer_config = _filter_transformer_config(metadata_config) + if transformer_config: + transformer_config["_class_name"] = transformer_class_name + + component_configs: dict[str, dict | None] = { + "transformer": transformer_config or None, + "vae": _wrap_component_config( + "vae", + metadata_config.get("vae"), + class_name="CausalVideoAutoencoder", + ), + "audio_vae": _wrap_component_config( + "audio_vae", + metadata_config.get("audio_vae"), + class_name="LTX2AudioDecoder", + ), + "vocoder": _wrap_component_config( + "vocoder", + metadata_config.get("vocoder"), + class_name="LTX2Vocoder", + ), + "text_embedding_projection": _build_text_embedding_projection_config( + gemma_model_path=gemma_model_path + ), + } + + output_dir.mkdir(parents=True, exist_ok=True) + for name, component_weights in split_weights.items(): + _write_component(output_dir, name, component_weights, component_configs.get(name)) + if emit_diffusers_repo and name == "text_embedding_projection": + _write_component( + output_dir, + name, + component_weights, + component_configs.get(name), + dir_name="text_encoder", + ) + if emit_diffusers_repo: + required_for_index = { + "transformer", + "vae", + "audio_vae", + "vocoder", + "text_embedding_projection", + } + if components_to_write is not None and not required_for_index.issubset(components_to_write): + print("Skipping model_index.json; not all diffusers components were written.") + return + if not required_for_index.issubset(split_weights.keys()): + print("Skipping model_index.json; missing diffusers components in weights.") + return + vae_class_name = (component_configs.get("vae") or {}).get( + "_class_name", "CausalVideoAutoencoder" + ) + model_index = _build_model_index( + transformer_class_name=transformer_class_name, + vae_class_name=vae_class_name, + pipeline_class_name=pipeline_class_name, + diffusers_version=diffusers_version, + ) + _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}") + return + + with config_path.open("r", encoding="utf-8") as f: + config = json.load(f) + + config["_class_name"] = class_name + with config_path.open("w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + f.write("\n") + print(f"Updated _class_name in {config_path} -> {class_name}") + + +def maybe_download(repo_id: str, target_dir: Path, token: str | None, allow_patterns: str | None) -> Path: + if snapshot_download is None: + raise RuntimeError("huggingface_hub is required for --download") + target_dir.mkdir(parents=True, exist_ok=True) + snapshot_download( + repo_id=repo_id, + local_dir=str(target_dir), + local_dir_use_symlinks=False, + token=token, + allow_patterns=allow_patterns, + ) + return target_dir + + +def main() -> None: + parser = argparse.ArgumentParser(description="Convert LTX-2 weights to FastVideo format") + parser.add_argument("--source", type=str, help="Path to transformer weights directory") + parser.add_argument("--output", type=str, required=True, help="Output directory for converted weights") + parser.add_argument("--download", type=str, help="HF repo id to download before conversion") + parser.add_argument("--allow-patterns", type=str, help="Limit HF download to matching files") + parser.add_argument("--token", type=str, default=os.getenv("HF_TOKEN"), help="HF token (or set HF_TOKEN)") + parser.add_argument("--update-config", action="store_true", help="Update source config.json _class_name") + parser.add_argument("--class-name", type=str, default="LTX2Transformer3DModel") + parser.add_argument( + "--diffusers-repo", + action=argparse.BooleanOptionalAction, + default=True, + help="Emit a diffusers-style repo layout with model_index.json.", + ) + parser.add_argument( + "--pipeline-class-name", + type=str, + default="LTX2Pipeline", + help="Pipeline class name for model_index.json.", + ) + parser.add_argument( + "--diffusers-version", + type=str, + default="0.33.0.dev0", + help="Diffusers version for model_index.json.", + ) + parser.add_argument( + "--transformer-only", + action="store_true", + help="Only convert transformer weights (no component split).", + ) + parser.add_argument( + "--components", + type=str, + default="", + help=( + "Comma-separated component list to write " + "(transformer,vae,audio_vae,vocoder,text_embedding_projection)." + ), + ) + parser.add_argument( + "--gemma-path", + type=str, + default="", + help="Optional local Gemma model path to copy into the output repo.", + ) + + args = parser.parse_args() + + if args.download: + if args.source: + raise ValueError("Use either --download or --source, not both.") + source_dir = maybe_download(args.download, Path(args.output) / "download", args.token, args.allow_patterns) + else: + if not args.source: + raise ValueError("--source is required when not using --download") + source_dir = Path(args.source) + + output_dir = Path(args.output) + shards = _find_shards(source_dir) + if not shards: + raise FileNotFoundError(f"No safetensors found in {source_dir}") + metadata_path = shards[0] + metadata_config = _read_metadata_config(metadata_path) + components_to_write: set[str] | None = None + if args.transformer_only: + components_to_write = {"transformer"} + elif args.components: + components_to_write = { + component.strip() + for component in args.components.split(",") + if component.strip() + } + + gemma_model_path = "" + if args.gemma_path: + gemma_src = Path(args.gemma_path) + if not gemma_src.is_dir(): + raise ValueError(f"--gemma-path must be a directory: {gemma_src}") + gemma_dest = output_dir / "text_encoder" / "gemma" + if gemma_dest.exists(): + shutil.rmtree(gemma_dest) + gemma_dest.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(gemma_src, gemma_dest) + gemma_model_path = "gemma" + + convert_components( + source_dir, + output_dir, + metadata_config, + args.class_name, + components_to_write=components_to_write, + emit_diffusers_repo=args.diffusers_repo, + pipeline_class_name=args.pipeline_class_name, + diffusers_version=args.diffusers_version, + gemma_model_path=gemma_model_path, + ) + + if args.update_config: + if source_dir.is_dir(): + update_transformer_config(source_dir / "config.json", args.class_name) + + +if __name__ == "__main__": + main() diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 000000000..6fc6b9d8a --- /dev/null +++ b/tests/README.md @@ -0,0 +1,7 @@ +# Tests + +- `tests/local_tests/` are local-only tests that require a checked-out `LTX-2/` + directory under the repo root (`FastVideo/LTX-2`). Without that repo, they + will skip or fail. +- The CI-backed test suite still lives in `fastvideo/tests/`. +- Eventually, all tests will move under `tests/`. diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/local_tests/README.md b/tests/local_tests/README.md new file mode 100644 index 000000000..c6b976d54 --- /dev/null +++ b/tests/local_tests/README.md @@ -0,0 +1,7 @@ +# Local LTX-2 Tests + +These tests depend on a checked-out `LTX-2/` directory under the repo root +(`FastVideo/LTX-2`). Without that local repo, the tests will skip or fail. + +For the CI-backed test suite, see `fastvideo/tests/`. Those are the tests +currently exercised in CI. (Eventually all tests will move to `tests/`.) diff --git a/tests/local_tests/__init__.py b/tests/local_tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/local_tests/encoders/__init__.py b/tests/local_tests/encoders/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/local_tests/encoders/test_ltx2_gemma_encoder.py b/tests/local_tests/encoders/test_ltx2_gemma_encoder.py new file mode 100644 index 000000000..8fbd65c08 --- /dev/null +++ b/tests/local_tests/encoders/test_ltx2_gemma_encoder.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +import sys +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file +from torch.testing import assert_close + +from fastvideo.configs.models.encoders import LTX2GemmaConfig +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.models.loader.component_loader import TextEncoderLoader + +repo_root = Path(__file__).resolve().parents[3] +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + + +def _load_connector_weights(path: str) -> dict[str, torch.Tensor]: + weights = load_file(path) + mapped: dict[str, torch.Tensor] = {} + for name, tensor in weights.items(): + if name == "aggregate_embed.weight": + mapped["feature_extractor_linear.aggregate_embed.weight"] = tensor + elif name.startswith("embeddings_connector."): + mapped[name] = tensor + elif name.startswith("audio_embeddings_connector."): + mapped[name] = tensor + return mapped + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="LTX-2 Gemma encoder parity test requires CUDA.", +) +def test_ltx2_gemma_text_encoder_parity(): + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + text_encoder_path = os.getenv( + "LTX2_TEXT_ENCODER_PATH", + str(diffusers_root / "text_encoder"), + ) + gemma_model_path = str(Path(text_encoder_path) / "gemma") + if not os.path.isdir(text_encoder_path): + pytest.skip(f"LTX-2 text encoder weights not found at {text_encoder_path}") + if not gemma_model_path or not os.path.isdir(gemma_model_path): + pytest.skip("Gemma weights not found in text_encoder/gemma.") + + try: + from ltx_core.text_encoders.gemma.embeddings_connector import ( + Embeddings1DConnector, + ) + from ltx_core.text_encoders.gemma.encoders.av_encoder import ( + AVGemmaTextEncoderModel, + ) + from ltx_core.text_encoders.gemma.feature_extractor import ( + GemmaFeaturesExtractorProjLinear, + ) + from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer + from transformers import Gemma3ForConditionalGeneration + except Exception as exc: + pytest.skip(f"LTX-2 Gemma import failed: {exc}") + + device = torch.device("cuda:0") + precision = torch.bfloat16 + + tokenizer = LTXVGemmaTokenizer(gemma_model_path, max_length=1024) + gemma_model = Gemma3ForConditionalGeneration.from_pretrained( + gemma_model_path, + local_files_only=True, + torch_dtype=precision, + ).to(device) + gemma_model.eval() + + ref_model = AVGemmaTextEncoderModel( + feature_extractor_linear=GemmaFeaturesExtractorProjLinear(), + embeddings_connector=Embeddings1DConnector(), + audio_embeddings_connector=Embeddings1DConnector(), + tokenizer=tokenizer, + model=gemma_model, + dtype=precision, + ).to(device) + ref_model.eval() + + connector_weights = _load_connector_weights( + os.path.join(text_encoder_path, "model.safetensors") + ) + ref_model.load_state_dict(connector_weights, strict=False) + + prompt = "A fast moving train in a snowy landscape." + token_pairs = tokenizer.tokenize_with_weights(prompt)["gemma"] + input_ids = torch.tensor( + [[t[0] for t in token_pairs]], device=device, dtype=torch.long + ) + attention_mask = torch.tensor( + [[t[1] for t in token_pairs]], device=device, dtype=torch.long + ) + + args = FastVideoArgs( + model_path=text_encoder_path, + pipeline_config=PipelineConfig( + text_encoder_configs=(LTX2GemmaConfig(),), + text_encoder_precisions=("bf16",), + ), + ) + loader = TextEncoderLoader() + fastvideo_model = loader.load(text_encoder_path, args).to(device) + fastvideo_model.eval() + + with torch.no_grad(): + ref_video, ref_audio, ref_mask = ref_model(prompt, padding_side="left") + with set_forward_context(current_timestep=0, attn_metadata=None): + fastvideo_out = fastvideo_model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + + assert_close(ref_video, fastvideo_out.last_hidden_state, atol=1e-2, rtol=1e-2) + assert_close(ref_audio, fastvideo_out.hidden_states[0], atol=1e-2, rtol=1e-2) + assert torch.equal(ref_mask, fastvideo_out.attention_mask) diff --git a/tests/local_tests/encoders/test_ltx2_gemma_parity.py b/tests/local_tests/encoders/test_ltx2_gemma_parity.py new file mode 100644 index 000000000..5f8de5c94 --- /dev/null +++ b/tests/local_tests/encoders/test_ltx2_gemma_parity.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from pathlib import Path +import sys + +import pytest +import torch +from torch.testing import assert_close + +from safetensors.torch import load_file + +from fastvideo.configs.models.encoders import LTX2GemmaConfig +from fastvideo.models.encoders.gemma import LTX2GemmaTextEncoderModel +from fastvideo.models.loader.component_loader import get_diffusers_config + + +repo_root = Path(__file__).resolve().parents[3] +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + + +def _init_log_paths() -> tuple[Path, Path]: + base_dir = Path(os.getenv("LTX2_DEBUG_DIR", "ltx2_debug")) + fastvideo_log = Path(os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", base_dir / "fastvideo_gemma.log")) + reference_log = Path(os.getenv("LTX2_REFERENCE_GEMMA_LOG", base_dir / "reference_gemma.log")) + fastvideo_log.parent.mkdir(parents=True, exist_ok=True) + reference_log.parent.mkdir(parents=True, exist_ok=True) + fastvideo_log.write_text("") + reference_log.write_text("") + return fastvideo_log, reference_log + + +def _log_line(path: Path, message: str) -> None: + with path.open("a", encoding="utf-8") as f: + f.write(message + "\n") + + +def _attach_encoder_logging( + encoder: torch.nn.Module, + log_path: Path, + label: str, +) -> None: + def _format_sum(tensor: torch.Tensor | None) -> str: + if tensor is None: + return "None" + return f"{tensor.float().sum().item():.6f}" + + def _hook_factory(name: str): + def _hook(_module, _inputs, outputs): # noqa: ANN001 + out = outputs[0] if isinstance(outputs, tuple) else outputs + out_sum = _format_sum(out if torch.is_tensor(out) else None) + _log_line(log_path, f"{label}:{name}:sum={out_sum}") + return _hook + + def _pre_hook_factory(name: str): + def _hook(_module, inputs): # noqa: ANN001 + tensor = inputs[0] if inputs else None + in_sum = _format_sum(tensor if torch.is_tensor(tensor) else None) + _log_line(log_path, f"{label}:{name}:in_sum={in_sum}") + return _hook + + def _attach_block_detail(block: torch.nn.Module, prefix: str) -> None: + block.register_forward_pre_hook(_pre_hook_factory(f"{prefix}:input")) + block.register_forward_hook(_hook_factory(f"{prefix}:output")) + if hasattr(block, "attn1"): + block.attn1.register_forward_hook( + _hook_factory(f"{prefix}:attn1")) + if hasattr(block, "ff"): + block.ff.register_forward_hook(_hook_factory(f"{prefix}:ff")) + + if hasattr(encoder, "feature_extractor_linear"): + encoder.feature_extractor_linear.register_forward_pre_hook( + _pre_hook_factory("feature_extractor_linear")) + encoder.feature_extractor_linear.register_forward_hook(_hook_factory("feature_extractor_linear")) + if hasattr(encoder, "embeddings_connector"): + encoder.embeddings_connector.register_forward_hook(_hook_factory("embeddings_connector")) + for idx, block in enumerate(encoder.embeddings_connector.transformer_1d_blocks): + _attach_block_detail(block, f"embeddings_block_{idx}") + if hasattr(encoder, "audio_embeddings_connector"): + encoder.audio_embeddings_connector.register_forward_hook(_hook_factory("audio_embeddings_connector")) + for idx, block in enumerate(encoder.audio_embeddings_connector.transformer_1d_blocks): + _attach_block_detail(block, f"audio_embeddings_block_{idx}") + + +def _log_register_sums(encoder: torch.nn.Module, log_path: Path, label: str) -> None: + def _sum_param(module: torch.nn.Module, name: str) -> float | None: + if not hasattr(module, "learnable_registers"): + return None + param = getattr(module, "learnable_registers") + if not torch.is_tensor(param): + return None + return param.float().sum().item() + + if hasattr(encoder, "embeddings_connector"): + reg_sum = _sum_param(encoder.embeddings_connector, "learnable_registers") + if reg_sum is not None: + _log_line(log_path, f"{label}:embeddings_registers:sum={reg_sum:.6f}") + if hasattr(encoder, "audio_embeddings_connector"): + reg_sum = _sum_param(encoder.audio_embeddings_connector, "learnable_registers") + if reg_sum is not None: + _log_line(log_path, f"{label}:audio_registers:sum={reg_sum:.6f}") + + +def _log_param_sums(encoder: torch.nn.Module, log_path: Path, label: str) -> None: + def _log_param(name: str, tensor: torch.Tensor | None) -> None: + if tensor is None: + _log_line(log_path, f"{label}:param:{name}:sum=None") + return + _log_line(log_path, f"{label}:param:{name}:sum={tensor.float().sum().item():.6f}") + + if hasattr(encoder, "feature_extractor_linear"): + _log_param( + "feature_extractor_linear.aggregate_embed.weight", + encoder.feature_extractor_linear.aggregate_embed.weight, + ) + if hasattr(encoder, "embeddings_connector"): + block0 = encoder.embeddings_connector.transformer_1d_blocks[0] + _log_param("embeddings_block0.attn1.to_q.weight", block0.attn1.to_q.weight) + _log_param("embeddings_block0.attn1.to_k.weight", block0.attn1.to_k.weight) + _log_param("embeddings_block0.attn1.to_v.weight", block0.attn1.to_v.weight) + _log_param("embeddings_block0.ff.net.0.proj.weight", block0.ff.net[0].proj.weight) + if hasattr(encoder, "audio_embeddings_connector"): + block0 = encoder.audio_embeddings_connector.transformer_1d_blocks[0] + _log_param("audio_block0.attn1.to_q.weight", block0.attn1.to_q.weight) + _log_param("audio_block0.attn1.to_k.weight", block0.attn1.to_k.weight) + _log_param("audio_block0.attn1.to_v.weight", block0.attn1.to_v.weight) + _log_param("audio_block0.ff.net.0.proj.weight", block0.ff.net[0].proj.weight) + + +def _log_gemma_param_sums( + gemma_model: torch.nn.Module | None, + log_path: Path, + label: str, +) -> None: + if gemma_model is None: + _log_line(log_path, f"{label}:gemma_param:embed_tokens.weight:sum=None") + return + tokens = None + if hasattr(gemma_model, "get_input_embeddings"): + try: + tokens = gemma_model.get_input_embeddings() + except Exception: + tokens = None + if tokens is None: + embed = getattr(gemma_model, "model", None) + if embed is not None: + tokens = getattr(embed, "embed_tokens", None) + if tokens is None or not hasattr(tokens, "weight"): + _log_line(log_path, f"{label}:gemma_param:embed_tokens.weight:sum=None") + return + _log_line( + log_path, + f"{label}:gemma_param:embed_tokens.weight:sum={tokens.weight.float().sum().item():.6f}", + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="LTX-2 Gemma parity test requires CUDA.", +) +def test_ltx2_gemma_parity(): + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + fastvideo_log, reference_log = _init_log_paths() + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + official_path = Path( + os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + ) + text_encoder_path = diffusers_root / "text_encoder" + gemma_path = text_encoder_path / "gemma" + + if not official_path.exists(): + pytest.skip(f"LTX-2 weights not found at {official_path}") + if not text_encoder_path.exists(): + pytest.skip(f"LTX-2 text encoder not found at {text_encoder_path}") + if not gemma_path.exists(): + pytest.skip(f"LTX-2 Gemma weights not found at {gemma_path}") + + try: + from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder + from ltx_core.model.transformer.attention import Attention, AttentionFunction + from ltx_core.text_encoders.gemma import ( + AV_GEMMA_TEXT_ENCODER_KEY_OPS, + AVGemmaTextEncoderModelConfigurator, + module_ops_from_gemma_root, + ) + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + device = torch.device("cuda:0") + precision = torch.bfloat16 + + reference_builder = SingleGPUModelBuilder( + model_path=str(official_path), + model_class_configurator=AVGemmaTextEncoderModelConfigurator, + model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS, + module_ops=module_ops_from_gemma_root(str(gemma_path)), + ) + reference_encoder = reference_builder.build( + device=device, dtype=precision + ).to(device=device, dtype=precision) + if hasattr(reference_encoder.model, "config"): + if hasattr(reference_encoder.model.config, "attn_implementation"): + reference_encoder.model.config.attn_implementation = "sdpa" + if hasattr(reference_encoder.model.config, "_attn_implementation"): + reference_encoder.model.config._attn_implementation = "sdpa" + for module in reference_encoder.modules(): + if isinstance(module, Attention): + module.attention_function = AttentionFunction.PYTORCH + reference_encoder.eval() + _attach_encoder_logging(reference_encoder, reference_log, "reference") + _log_register_sums(reference_encoder, reference_log, "reference") + _log_param_sums(reference_encoder, reference_log, "reference") + _log_gemma_param_sums(reference_encoder.model, reference_log, "reference") + _log_line( + reference_log, + "reference:gemma_config:attn_impl=" + f"{getattr(reference_encoder.model.config, 'attn_implementation', None)} " + f"dtype={reference_encoder.model.dtype}", + ) + + diffusers_config = get_diffusers_config(model=str(text_encoder_path)) + encoder_config = LTX2GemmaConfig() + encoder_config.update_model_arch(diffusers_config) + encoder_config.arch_config.gemma_model_path = str(gemma_path) + fastvideo_encoder = LTX2GemmaTextEncoderModel(encoder_config).to( + device=device, dtype=precision + ) + if hasattr(fastvideo_encoder.gemma_model, "config"): + if hasattr(fastvideo_encoder.gemma_model.config, "attn_implementation"): + fastvideo_encoder.gemma_model.config.attn_implementation = "sdpa" + if hasattr(fastvideo_encoder.gemma_model.config, "_attn_implementation"): + fastvideo_encoder.gemma_model.config._attn_implementation = "sdpa" + + official_weights = load_file(str(official_path)) + fastvideo_weights: dict[str, torch.Tensor] = {} + for name, tensor in official_weights.items(): + mapped_name = AV_GEMMA_TEXT_ENCODER_KEY_OPS.apply_to_key(name) + if mapped_name is None: + continue + fastvideo_weights[mapped_name] = tensor + fastvideo_encoder.load_weights(fastvideo_weights.items()) + fastvideo_encoder.eval() + _attach_encoder_logging(fastvideo_encoder, fastvideo_log, "fastvideo") + _log_register_sums(fastvideo_encoder, fastvideo_log, "fastvideo") + _log_param_sums(fastvideo_encoder, fastvideo_log, "fastvideo") + _log_gemma_param_sums(fastvideo_encoder.gemma_model, fastvideo_log, "fastvideo") + _log_line( + fastvideo_log, + "fastvideo:gemma_config:attn_impl=" + f"{getattr(fastvideo_encoder.gemma_model.config, 'attn_implementation', None)} " + f"dtype={fastvideo_encoder.gemma_model.dtype}", + ) + + prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers." + ref_tokenizer = reference_encoder.tokenizer + if ref_tokenizer is None: + pytest.skip("Reference tokenizer is not initialized.") + token_pairs = ref_tokenizer.tokenize_with_weights(prompt)["gemma"] + input_ids = torch.tensor([[t[0] for t in token_pairs]], device=device) + attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=device) + + with torch.no_grad(), torch.backends.cuda.sdp_kernel( + enable_flash=False, + enable_mem_efficient=False, + enable_math=True, + ): + ref_outputs = reference_encoder.model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + use_cache=False, + ) + _log_line( + reference_log, + "reference:gemma_hidden_last:sum=" + f"{ref_outputs.hidden_states[-1].float().sum().item():.6f}", + ) + ref_projected = reference_encoder._run_feature_extractor( + ref_outputs.hidden_states, + attention_mask=attention_mask, + padding_side="left", + ) + # Compare rotary embeddings between implementations. + from ltx_core.model.transformer.rope import precompute_freqs_cis + from fastvideo.models.dits.ltx2 import precompute_ltx_freqs_cis + + seq_len = ref_projected.shape[1] + indices_grid = torch.arange(seq_len, device=device, dtype=torch.float32)[None, None, :] + ref_cos, ref_sin = precompute_freqs_cis( + indices_grid=indices_grid, + dim=reference_encoder.embeddings_connector.inner_dim, + out_dtype=ref_projected.dtype, + theta=reference_encoder.embeddings_connector.positional_embedding_theta, + max_pos=reference_encoder.embeddings_connector.positional_embedding_max_pos, + num_attention_heads=reference_encoder.embeddings_connector.num_attention_heads, + rope_type=reference_encoder.embeddings_connector.rope_type, + ) + fast_cos, fast_sin = precompute_ltx_freqs_cis( + indices_grid=indices_grid, + dim=fastvideo_encoder.embeddings_connector.inner_dim, + out_dtype=ref_projected.dtype, + theta=fastvideo_encoder.embeddings_connector.positional_embedding_theta, + max_pos=fastvideo_encoder.embeddings_connector.positional_embedding_max_pos, + num_attention_heads=fastvideo_encoder.embeddings_connector.num_attention_heads, + rope_type=fastvideo_encoder.embeddings_connector.rope_type, + ) + _log_line( + reference_log, + "reference:rope:cos_sum=" + f"{ref_cos.float().sum().item():.6f} sin_sum={ref_sin.float().sum().item():.6f} " + f"theta={reference_encoder.embeddings_connector.positional_embedding_theta} " + f"max_pos={reference_encoder.embeddings_connector.positional_embedding_max_pos} " + f"rope_type={reference_encoder.embeddings_connector.rope_type}" + ) + _log_line( + fastvideo_log, + "fastvideo:rope:cos_sum=" + f"{fast_cos.float().sum().item():.6f} sin_sum={fast_sin.float().sum().item():.6f} " + f"theta={fastvideo_encoder.embeddings_connector.positional_embedding_theta} " + f"max_pos={fastvideo_encoder.embeddings_connector.positional_embedding_max_pos} " + f"rope_type={fastvideo_encoder.embeddings_connector.rope_type}" + ) + ref_video, ref_audio, _ = reference_encoder._run_connectors( + ref_projected, attention_mask + ) + fast_video_from_ref, fast_audio_from_ref, _ = fastvideo_encoder._run_connectors( + ref_projected, attention_mask + ) + _log_line( + fastvideo_log, + "fastvideo:connector_on_ref:video_sum=" + f"{fast_video_from_ref.float().sum().item():.6f}", + ) + _log_line( + fastvideo_log, + "fastvideo:connector_on_ref:audio_sum=" + f"{fast_audio_from_ref.float().sum().item():.6f}", + ) + + fast_outputs = fastvideo_encoder.gemma_model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + use_cache=False, + ) + _log_line( + fastvideo_log, + "fastvideo:gemma_hidden_last:sum=" + f"{fast_outputs.hidden_states[-1].float().sum().item():.6f}", + ) + fast_projected = fastvideo_encoder._run_feature_extractor( + fast_outputs.hidden_states, + attention_mask=attention_mask, + padding_side=fastvideo_encoder.padding_side, + ) + fast_video, fast_audio, _ = fastvideo_encoder._run_connectors( + fast_projected, attention_mask + ) + + assert ref_video.shape == fast_video.shape + assert ref_audio.shape == fast_audio.shape + assert torch.isfinite(ref_video).all(), "Reference Gemma produced non-finite video embeddings." + assert torch.isfinite(ref_audio).all(), "Reference Gemma produced non-finite audio embeddings." + assert torch.isfinite(fast_video).all(), "FastVideo Gemma produced non-finite video embeddings." + assert torch.isfinite(fast_audio).all(), "FastVideo Gemma produced non-finite audio embeddings." + assert_close(ref_video, fast_video, atol=3e-1, rtol=5e-2) + assert_close(ref_audio, fast_audio, atol=3e-1, rtol=5e-2) diff --git a/tests/local_tests/pipelines/__init__.py b/tests/local_tests/pipelines/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/local_tests/pipelines/test_ltx2_pipeline_smoke.py b/tests/local_tests/pipelines/test_ltx2_pipeline_smoke.py new file mode 100644 index 000000000..1369934f6 --- /dev/null +++ b/tests/local_tests/pipelines/test_ltx2_pipeline_smoke.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from pathlib import Path +import sys +import tempfile + +import pytest +import torch + +from torch.testing import assert_close + +from fastvideo import VideoGenerator +from fastvideo.models.dits.ltx2 import ( + AudioLatentShape, + DEFAULT_LTX2_AUDIO_CHANNELS, + DEFAULT_LTX2_AUDIO_DOWNSAMPLE, + DEFAULT_LTX2_AUDIO_HOP_LENGTH, + DEFAULT_LTX2_AUDIO_MEL_BINS, + DEFAULT_LTX2_AUDIO_SAMPLE_RATE, +) +from fastvideo.models.loader.component_loader import PipelineComponentLoader + + +def _log_tensor_stats(label: str, tensor: torch.Tensor) -> None: + tensor_f32 = tensor.float() + print( + f"[LTX2 SMOKE] {label}: shape={tuple(tensor.shape)} " + f"dtype={tensor.dtype} device={tensor.device} " + f"min={tensor_f32.min().item():.6f} max={tensor_f32.max().item():.6f} " + f"mean={tensor_f32.mean().item():.6f} sum={tensor_f32.sum().item():.6f}" + ) + +def _truncate_debug_logs() -> None: + for env_var in ( + "LTX2_PIPELINE_DEBUG_PATH", + "LTX2_REFERENCE_DEBUG_PATH", + "LTX2_PIPELINE_DEBUG_DETAIL_PATH", + "LTX2_REFERENCE_DEBUG_DETAIL_PATH", + ): + log_path = os.getenv(env_var, "") + if not log_path: + continue + log_dir = os.path.dirname(log_path) + if log_dir: + os.makedirs(log_dir, exist_ok=True) + with open(log_path, "w", encoding="utf-8") as f: + f.write("") + + +def _run_audio_decode_smoke( + diffusers_path: str, + fastvideo_args, + device: torch.device, + num_frames: int, + fps: float, +) -> None: + audio_decoder_path = os.path.join(diffusers_path, "audio_vae") + vocoder_path = os.path.join(diffusers_path, "vocoder") + if not os.path.isdir(audio_decoder_path): + pytest.skip(f"Missing LTX-2 audio decoder at {audio_decoder_path}") + if not os.path.isdir(vocoder_path): + pytest.skip(f"Missing LTX-2 vocoder at {vocoder_path}") + + audio_decoder = PipelineComponentLoader.load_module( + module_name="audio_decoder", + component_model_path=audio_decoder_path, + transformers_or_diffusers="diffusers", + fastvideo_args=fastvideo_args, + ) + vocoder = PipelineComponentLoader.load_module( + module_name="vocoder", + component_model_path=vocoder_path, + transformers_or_diffusers="diffusers", + fastvideo_args=fastvideo_args, + ) + + duration = float(num_frames) / float(fps) + audio_shape = AudioLatentShape.from_duration( + batch=1, + duration=duration, + channels=DEFAULT_LTX2_AUDIO_CHANNELS, + mel_bins=DEFAULT_LTX2_AUDIO_MEL_BINS, + sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE, + hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH, + audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE, + ) + audio_dtype = next(audio_decoder.parameters()).dtype + audio_latents = torch.randn( + audio_shape.to_torch_shape(), + device=device, + dtype=audio_dtype, + ) + + with torch.no_grad(): + decoded_spec = audio_decoder(audio_latents) + audio_wave = vocoder(decoded_spec) + + assert audio_wave.ndim == 3 + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="LTX-2 pipeline smoke test requires CUDA.", +) +def test_ltx2_pipeline_smoke(): + repo_root = Path(__file__).resolve().parents[3] + debug_dir = repo_root / "ltx2_debug" + os.environ.setdefault( + "LTX2_PIPELINE_DEBUG_PATH", + str(debug_dir / "fastvideo_pipeline.log"), + ) + os.environ.setdefault( + "LTX2_REFERENCE_DEBUG_PATH", + str(debug_dir / "reference_pipeline.log"), + ) + os.environ.setdefault( + "LTX2_PIPELINE_DEBUG_DETAIL_PATH", + str(debug_dir / "fastvideo_pipeline_detail.log"), + ) + os.environ.setdefault( + "LTX2_REFERENCE_DEBUG_DETAIL_PATH", + str(debug_dir / "reference_pipeline_detail.log"), + ) + os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + os.environ.setdefault("LTX2_REFERENCE_ATTN", "pytorch") + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + _truncate_debug_logs() + + ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" + if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + ltx_pipelines_path = repo_root / "LTX-2" / "packages" / "ltx-pipelines" / "src" + if ltx_pipelines_path.exists() and str(ltx_pipelines_path) not in sys.path: + sys.path.insert(0, str(ltx_pipelines_path)) + os.environ["PYTHONPATH"] = str(repo_root) + + diffusers_path = os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + gemma_model_path = os.path.join(diffusers_path, "text_encoder", "gemma") + official_path = os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + + if not os.path.isdir(diffusers_path): + pytest.skip(f"Missing LTX-2 diffusers repo at {diffusers_path}") + if not os.path.isfile(os.path.join(diffusers_path, "model_index.json")): + pytest.skip(f"Missing model_index.json in {diffusers_path}") + + if not gemma_model_path or not os.path.isdir(gemma_model_path): + pytest.skip("Gemma weights not found in text_encoder/gemma.") + if not os.path.isfile(official_path): + pytest.skip(f"Missing LTX-2 official weights at {official_path}") + + try: + from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + from ltx_core.model.transformer import attention as ltx_attention + except ImportError as exc: + pytest.skip(f"LTX-2 pipeline import failed: {exc}") + ltx_attention.memory_efficient_attention = None + ltx_attention.flash_attn_interface = None + + device = torch.device("cuda:0") + prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers." + negative_prompt = "low quality, blurry, distorted, artifacts, jpeg compression" + seed = 42 + height = 64 + width = 96 + num_frames = 9 + fps = 12.0 + steps = 4 + guidance_scale = 4.0 + + with tempfile.TemporaryDirectory() as tmpdir: + latent_path = str(Path(tmpdir) / "ltx2_initial_latent.pt") + + generator = VideoGenerator.from_pretrained( + diffusers_path, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + pin_cpu_memory=False, + ltx2_vae_tiling=False, + ltx2_initial_latent_path=latent_path, + ) + result = generator.generate_video( + prompt=prompt, + negative_prompt=negative_prompt, + output_path="outputs_video/ltx2_smoke", + save_video=False, + height=height, + width=width, + num_frames=num_frames, + fps=fps, + num_inference_steps=steps, + guidance_scale=guidance_scale, + seed=seed, + ) + generator.shutdown() + _run_audio_decode_smoke( + diffusers_path=diffusers_path, + fastvideo_args=generator.fastvideo_args, + device=device, + num_frames=num_frames, + fps=fps, + ) + + fastvideo_out = result["samples"] + fastvideo_out = fastvideo_out.to(device=device, dtype=torch.float32) + _log_tensor_stats("fastvideo_video", fastvideo_out) + + ref_pipeline = TI2VidOneStagePipeline( + checkpoint_path=official_path, + gemma_root=gemma_model_path, + loras=[], + device=device, + fp8transformer=False, + ) + original_text_encoder = ref_pipeline.model_ledger.text_encoder + original_transformer = ref_pipeline.model_ledger.transformer + original_video_decoder = ref_pipeline.model_ledger.video_decoder + + def _patched_text_encoder(): + encoder = original_text_encoder() + try: + from ltx_core.model.transformer.attention import ( # type: ignore + Attention, + AttentionFunction, + ) + except ImportError: + return encoder + if hasattr(encoder, "model") and hasattr(encoder.model, "config"): + if hasattr(encoder.model.config, "attn_implementation"): + encoder.model.config.attn_implementation = "sdpa" + if hasattr(encoder.model.config, "_attn_implementation"): + encoder.model.config._attn_implementation = "sdpa" + for module in encoder.modules(): + if isinstance(module, Attention): + module.attention_function = AttentionFunction.PYTORCH + return encoder + + ref_pipeline.model_ledger.text_encoder = _patched_text_encoder + if os.getenv("LTX2_DISABLE_VAE_NOISE", "1") == "1": + def _patched_video_decoder(): + decoder = original_video_decoder() + if hasattr(decoder, "decode_noise_scale"): + decoder.decode_noise_scale = 0.0 + return decoder + + ref_pipeline.model_ledger.video_decoder = _patched_video_decoder + if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1": + from ..transformers.test_ltx2 import ( + _attach_block_detail_logging, + ) + + def _patched_transformer(): + model = original_transformer() + core = getattr(model, "velocity_model", model) + _attach_block_detail_logging( + core, + Path(os.environ["LTX2_REFERENCE_DEBUG_DETAIL_PATH"]), + "reference", + True, + ) + return model + + ref_pipeline.model_ledger.transformer = _patched_transformer + with torch.no_grad(): + ref_video_iter, _ = ref_pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=fps, + num_inference_steps=steps, + cfg_guidance_scale=guidance_scale, + images=[], + enhance_prompt=False, + initial_video_latent_path=latent_path, + ) + ref_chunks = list(ref_video_iter) + ref_video = torch.cat( + [chunk if torch.is_tensor(chunk) else torch.from_numpy(chunk) for chunk in ref_chunks], + dim=0, + ) + ref_video = ref_video.to(torch.float32) / 255.0 + ref_video = ref_video.permute(3, 0, 1, 2).unsqueeze(0) + _log_tensor_stats("reference_video", ref_video) + + assert ref_video.shape == fastvideo_out.shape + assert_close(ref_video, fastvideo_out, atol=2 / 255, rtol=1e-3) diff --git a/tests/local_tests/transformers/__init__.py b/tests/local_tests/transformers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/local_tests/transformers/test_ltx2.py b/tests/local_tests/transformers/test_ltx2.py new file mode 100644 index 000000000..14a3d8366 --- /dev/null +++ b/tests/local_tests/transformers/test_ltx2.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from pathlib import Path +import sys + +import pytest +import torch +from torch.testing import assert_close + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29513") + +repo_root = Path(__file__).resolve().parents[3] +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.configs.models.dits import LTX2VideoConfig +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.models.loader.component_loader import TransformerLoader + + +def _read_transformer_config(config: dict) -> dict: + transformer_config = config.get("transformer", {}) + if not transformer_config: + raise ValueError("Missing transformer config in LTX-2 metadata.") + return transformer_config + + +def _infer_patch_params(in_channels: int) -> tuple[int, int]: + patch_size = 1 + num_channels_latents = 128 + for candidate in (8, 16, 32, 64, 128): + if in_channels % candidate != 0: + continue + patch_volume = in_channels // candidate + root = int(round(patch_volume**0.5)) + if root * root == patch_volume: + patch_size = root + num_channels_latents = candidate + break + print( + f"[LTX2 TEST] Inferred patch_size={patch_size}, " + f"num_channels_latents={num_channels_latents}" + ) + return patch_size, num_channels_latents + + +def _attach_block_sum_logging( + model: torch.nn.Module, + log_path: Path, + label: str, + enabled: bool, +) -> None: + if not enabled: + return + + log_path.parent.mkdir(parents=True, exist_ok=True) + if log_path.exists(): + log_path.unlink() + + def _format_sum(tensor: torch.Tensor | None) -> str: + if tensor is None: + return "None" + return f"{tensor.float().sum().item():.6f}" + + def _hook(module, inputs, outputs): # noqa: ANN001 + if isinstance(outputs, tuple): + video_args, audio_args = outputs + video_sum = _format_sum(video_args.x if video_args is not None else None) + audio_sum = _format_sum(audio_args.x if audio_args is not None else None) + else: + video_sum = _format_sum(outputs) + audio_sum = "None" + with log_path.open("a", encoding="utf-8") as f: + f.write(f"{label}:{module.idx}:video_sum={video_sum},audio_sum={audio_sum}\n") + + for block in model.transformer_blocks: + block.register_forward_hook(_hook) + + +def _attach_block_detail_logging( + model: torch.nn.Module, + log_path: Path, + label: str, + enabled: bool, +) -> None: + if not enabled: + return + + log_path.parent.mkdir(parents=True, exist_ok=True) + if log_path.exists(): + log_path.unlink() + + def _format_sum(tensor: torch.Tensor | None) -> str: + if tensor is None: + return "None" + return f"{tensor.float().sum().item():.6f}" + + def _hook_factory(block_idx: int, name: str): + def _hook(_module, _inputs, outputs): # noqa: ANN001 + if isinstance(outputs, tuple): + out = outputs[0] + else: + out = outputs + out_sum = _format_sum(out if torch.is_tensor(out) else None) + with log_path.open("a", encoding="utf-8") as f: + f.write(f"{label}:{block_idx}:{name}:out_sum={out_sum}\n") + return _hook + + for block in model.transformer_blocks: + idx = block.idx + for name in ( + "attn1", + "attn2", + "ff", + "audio_attn1", + "audio_attn2", + "audio_ff", + "audio_to_video_attn", + "video_to_audio_attn", + ): + if hasattr(block, name): + getattr(block, name).register_forward_hook(_hook_factory(idx, name)) + + def _output_hook(name: str): + def _hook(_module, _inputs, outputs): # noqa: ANN001 + out = outputs[0] if isinstance(outputs, tuple) else outputs + out_sum = _format_sum(out if torch.is_tensor(out) else None) + with log_path.open("a", encoding="utf-8") as f: + f.write(f"{label}:output:{name}:out_sum={out_sum}\n") + return _hook + + for name in ("proj_out", "audio_proj_out"): + if hasattr(model, name): + getattr(model, name).register_forward_hook(_output_hook(name)) + + +def test_ltx2_transformer_parity(): + torch.manual_seed(42) + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + official_path = Path( + os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + ) + fastvideo_path = Path( + os.getenv( + "LTX2_FASTVIDEO_PATH", + str(diffusers_root / "transformer"), + ) + ) + if not official_path.exists(): + pytest.skip(f"LTX-2 official weights not found at {official_path}") + if not fastvideo_path.exists(): + pytest.skip(f"FastVideo converted weights not found at {fastvideo_path}") + + try: + from ltx_core.components.patchifiers import VideoLatentPatchifier + from ltx_core.guidance.perturbations import BatchedPerturbationConfig + from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader + from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder + from ltx_core.model.transformer import (LTXModelConfigurator, + LTXV_MODEL_COMFY_RENAMING_MAP) + from ltx_core.model.transformer.modality import Modality + from ltx_core.types import VideoLatentShape + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + config_loader = SafetensorsModelStateDictLoader() + metadata = config_loader.metadata(str(official_path)) + transformer_config = _read_transformer_config(metadata) + + config = LTX2VideoConfig() + cfg = config.arch_config + cfg.num_attention_heads = transformer_config.get("num_attention_heads", + cfg.num_attention_heads) + cfg.attention_head_dim = transformer_config.get("attention_head_dim", + cfg.attention_head_dim) + cfg.num_layers = transformer_config.get("num_layers", cfg.num_layers) + cfg.cross_attention_dim = transformer_config.get( + "cross_attention_dim", cfg.cross_attention_dim) + cfg.caption_channels = transformer_config.get("caption_channels", + cfg.caption_channels) + cfg.norm_eps = transformer_config.get("norm_eps", cfg.norm_eps) + cfg.attention_type = transformer_config.get("attention_type", + cfg.attention_type) + cfg.positional_embedding_theta = transformer_config.get( + "positional_embedding_theta", cfg.positional_embedding_theta) + cfg.positional_embedding_max_pos = transformer_config.get( + "positional_embedding_max_pos", cfg.positional_embedding_max_pos) + cfg.timestep_scale_multiplier = transformer_config.get( + "timestep_scale_multiplier", cfg.timestep_scale_multiplier) + cfg.use_middle_indices_grid = transformer_config.get( + "use_middle_indices_grid", cfg.use_middle_indices_grid) + cfg.rope_type = transformer_config.get("rope_type", cfg.rope_type) + cfg.double_precision_rope = transformer_config.get( + "double_precision_rope", + transformer_config.get("frequencies_precision", "") + == "float64", + ) + cfg.audio_num_attention_heads = transformer_config.get( + "audio_num_attention_heads", cfg.audio_num_attention_heads) + cfg.audio_attention_head_dim = transformer_config.get( + "audio_attention_head_dim", cfg.audio_attention_head_dim) + cfg.audio_in_channels = transformer_config.get("audio_in_channels", + cfg.audio_in_channels) + cfg.audio_out_channels = transformer_config.get("audio_out_channels", + cfg.audio_out_channels) + cfg.audio_cross_attention_dim = transformer_config.get( + "audio_cross_attention_dim", cfg.audio_cross_attention_dim) + cfg.audio_positional_embedding_max_pos = transformer_config.get( + "audio_positional_embedding_max_pos", + cfg.audio_positional_embedding_max_pos, + ) + cfg.av_ca_timestep_scale_multiplier = transformer_config.get( + "av_ca_timestep_scale_multiplier", cfg.av_ca_timestep_scale_multiplier) + cfg.in_channels = transformer_config.get("in_channels", cfg.in_channels) + cfg.out_channels = transformer_config.get("out_channels", cfg.out_channels) + + patch_size, num_channels_latents = _infer_patch_params(cfg.in_channels) + cfg.patch_size = (1, patch_size, patch_size) + cfg.num_channels_latents = num_channels_latents + + if not torch.cuda.is_available(): + pytest.skip("LTX-2 transformer parity test requires CUDA for attention backends.") + + device = torch.device("cuda:0") + precision = torch.bfloat16 + precision_str = "bf16" + + args = FastVideoArgs( + model_path=str(fastvideo_path), + dit_cpu_offload=True, + use_fsdp_inference=False, + pipeline_config=PipelineConfig(dit_config=config, dit_precision=precision_str), + ) + args.device = device + + loader = TransformerLoader() + fastvideo_model = loader.load(str(fastvideo_path), args).to(device=device, dtype=precision) + + reference_builder = SingleGPUModelBuilder( + model_class_configurator=LTXModelConfigurator, + model_path=str(official_path), + model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, + ) + reference_model = reference_builder.build( + device=device, dtype=precision).to(device=device, dtype=precision) + reference_model.set_gradient_checkpointing(False) + + fastvideo_model.eval() + reference_model.eval() + + debug_logs = os.getenv("LTX2_DEBUG_LOGS", "0") == "1" + _attach_block_sum_logging( + fastvideo_model.model, + repo_root / "ltx2_debug" / "fastvideo.log", + "fastvideo", + debug_logs, + ) + _attach_block_sum_logging( + reference_model, + repo_root / "ltx2_debug" / "reference.log", + "reference", + debug_logs, + ) + _attach_block_detail_logging( + fastvideo_model.model, + repo_root / "ltx2_debug" / "fastvideo_detail.log", + "fastvideo", + os.getenv("LTX2_DEBUG_DETAIL", "0") == "1", + ) + _attach_block_detail_logging( + reference_model, + repo_root / "ltx2_debug" / "reference_detail.log", + "reference", + os.getenv("LTX2_DEBUG_DETAIL", "0") == "1", + ) + + patchifier = VideoLatentPatchifier(patch_size=cfg.patch_size[1]) + batch_size = 1 + frames = 4 + height = cfg.patch_size[1] * 4 + width = cfg.patch_size[2] * 4 + hidden_states = torch.randn( + batch_size, + cfg.num_channels_latents, + frames, + height, + width, + device=device, + dtype=precision, + ) + encoder_hidden_states = torch.randn( + batch_size, + 16, + cfg.caption_channels, + device=device, + dtype=precision, + ) + timestep = torch.tensor([500], device=device, dtype=precision) + + video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) + positions = patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) + latents = patchifier.patchify(hidden_states) + + video = Modality( + enabled=True, + latent=latents, + timesteps=timestep, + positions=positions, + context=encoder_hidden_states, + context_mask=None, + ) + + with torch.no_grad(): + ref_out, _ = reference_model( + video=video, + audio=None, + perturbations=BatchedPerturbationConfig.empty(batch_size), + ) + ref_out = patchifier.unpatchify(ref_out, output_shape=video_shape) + print(f"[LTX2 TEST] Reference model output shape: {ref_out.shape}") + with set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=None, + ): + fastvideo_out = fastvideo_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + ) + print(f"[LTX2 TEST] FastVideo model output shape: {fastvideo_out.shape}") + assert ref_out.shape == fastvideo_out.shape + assert ref_out.dtype == fastvideo_out.dtype + assert_close(ref_out, fastvideo_out, atol=1e-4, rtol=1e-4) diff --git a/tests/local_tests/transformers/test_ltx2_audio.py b/tests/local_tests/transformers/test_ltx2_audio.py new file mode 100644 index 000000000..c67c8cb74 --- /dev/null +++ b/tests/local_tests/transformers/test_ltx2_audio.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from pathlib import Path +import sys + +import pytest +import torch +from torch.testing import assert_close + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29513") +# Force TORCH_SDPA backend for parity testing - both FastVideo and LTX-2 reference +# will use PyTorch's scaled_dot_product_attention for consistent results +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + +repo_root = Path(__file__).resolve().parents[3] +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.configs.models.dits import LTX2VideoConfig +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.models.dits.ltx2 import Modality as FastVideoModality +from fastvideo.models.loader.component_loader import TransformerLoader +from .test_ltx2 import ( + _attach_block_detail_logging, + _attach_block_sum_logging, + _infer_patch_params, + _read_transformer_config, +) + + +def test_ltx2_transformer_audio_parity(): + torch.manual_seed(42) + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + official_path = Path( + os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + ) + fastvideo_path = Path( + os.getenv( + "LTX2_FASTVIDEO_PATH", + str(diffusers_root / "transformer"), + ) + ) + if not official_path.exists(): + pytest.skip(f"LTX-2 official weights not found at {official_path}") + if not fastvideo_path.exists(): + pytest.skip(f"FastVideo converted weights not found at {fastvideo_path}") + + try: + from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier + from ltx_core.guidance.perturbations import BatchedPerturbationConfig + from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader + from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder + from ltx_core.model.transformer import (LTXModelConfigurator, + LTXV_MODEL_COMFY_RENAMING_MAP) + from ltx_core.model.transformer.modality import Modality + from ltx_core.types import AudioLatentShape, VideoLatentShape + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + # Load config from metadata using same approach as test_ltx2.py + config_loader = SafetensorsModelStateDictLoader() + metadata = config_loader.metadata(str(official_path)) + transformer_config = _read_transformer_config(metadata) + + config = LTX2VideoConfig() + cfg = config.arch_config + cfg.num_attention_heads = transformer_config.get("num_attention_heads", + cfg.num_attention_heads) + cfg.attention_head_dim = transformer_config.get("attention_head_dim", + cfg.attention_head_dim) + cfg.num_layers = transformer_config.get("num_layers", cfg.num_layers) + cfg.cross_attention_dim = transformer_config.get( + "cross_attention_dim", cfg.cross_attention_dim) + cfg.caption_channels = transformer_config.get("caption_channels", + cfg.caption_channels) + cfg.norm_eps = transformer_config.get("norm_eps", cfg.norm_eps) + cfg.attention_type = transformer_config.get("attention_type", + cfg.attention_type) + cfg.positional_embedding_theta = transformer_config.get( + "positional_embedding_theta", cfg.positional_embedding_theta) + cfg.positional_embedding_max_pos = transformer_config.get( + "positional_embedding_max_pos", cfg.positional_embedding_max_pos) + cfg.timestep_scale_multiplier = transformer_config.get( + "timestep_scale_multiplier", cfg.timestep_scale_multiplier) + cfg.use_middle_indices_grid = transformer_config.get( + "use_middle_indices_grid", cfg.use_middle_indices_grid) + cfg.rope_type = transformer_config.get("rope_type", cfg.rope_type) + cfg.double_precision_rope = transformer_config.get( + "double_precision_rope", + transformer_config.get("frequencies_precision", "") + == "float64", + ) + cfg.audio_num_attention_heads = transformer_config.get( + "audio_num_attention_heads", cfg.audio_num_attention_heads) + cfg.audio_attention_head_dim = transformer_config.get( + "audio_attention_head_dim", cfg.audio_attention_head_dim) + cfg.audio_in_channels = transformer_config.get("audio_in_channels", + cfg.audio_in_channels) + cfg.audio_out_channels = transformer_config.get("audio_out_channels", + cfg.audio_out_channels) + cfg.audio_cross_attention_dim = transformer_config.get( + "audio_cross_attention_dim", cfg.audio_cross_attention_dim) + cfg.audio_positional_embedding_max_pos = transformer_config.get( + "audio_positional_embedding_max_pos", + cfg.audio_positional_embedding_max_pos, + ) + cfg.av_ca_timestep_scale_multiplier = transformer_config.get( + "av_ca_timestep_scale_multiplier", cfg.av_ca_timestep_scale_multiplier) + cfg.in_channels = transformer_config.get("in_channels", cfg.in_channels) + cfg.out_channels = transformer_config.get("out_channels", cfg.out_channels) + + patch_size, num_channels_latents = _infer_patch_params(cfg.in_channels) + cfg.patch_size = (1, patch_size, patch_size) + cfg.num_channels_latents = num_channels_latents + + if not torch.cuda.is_available(): + pytest.skip("LTX-2 transformer parity test requires CUDA for attention backends.") + + device = torch.device("cuda:0") + precision = torch.bfloat16 + precision_str = "bf16" + + args = FastVideoArgs( + model_path=str(fastvideo_path), + dit_cpu_offload=True, + use_fsdp_inference=False, + pipeline_config=PipelineConfig(dit_config=config, dit_precision=precision_str), + ) + args.device = device + + loader = TransformerLoader() + fastvideo_model = loader.load(str(fastvideo_path), args).to(device=device, dtype=precision) + + # Use SingleGPUModelBuilder to load the reference model (same as test_ltx2.py) + reference_builder = SingleGPUModelBuilder( + model_class_configurator=LTXModelConfigurator, + model_path=str(official_path), + model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, + ) + reference_model = reference_builder.build( + device=device, dtype=precision).to(device=device, dtype=precision) + reference_model.set_gradient_checkpointing(False) + + fastvideo_model.eval() + reference_model.eval() + + debug_logs = os.getenv("LTX2_DEBUG_LOGS", "0") == "1" + _attach_block_sum_logging( + fastvideo_model.model, + repo_root / "ltx2_debug" / "fastvideo_audio.log", + "fastvideo", + debug_logs, + ) + _attach_block_sum_logging( + reference_model, + repo_root / "ltx2_debug" / "reference_audio.log", + "reference", + debug_logs, + ) + _attach_block_detail_logging( + fastvideo_model.model, + repo_root / "ltx2_debug" / "fastvideo_audio_detail.log", + "fastvideo", + os.getenv("LTX2_DEBUG_DETAIL", "0") == "1", + ) + _attach_block_detail_logging( + reference_model, + repo_root / "ltx2_debug" / "reference_audio_detail.log", + "reference", + os.getenv("LTX2_DEBUG_DETAIL", "0") == "1", + ) + + patchifier = VideoLatentPatchifier(patch_size=cfg.patch_size[1]) + audio_patchifier = AudioPatchifier(patch_size=1) + batch_size = 1 + frames = 4 + height = cfg.patch_size[1] * 4 + width = cfg.patch_size[2] * 4 + hidden_states = torch.randn( + batch_size, + cfg.num_channels_latents, + frames, + height, + width, + device=device, + dtype=precision, + ) + encoder_hidden_states = torch.randn( + batch_size, + 16, + cfg.caption_channels, + device=device, + dtype=precision, + ) + timestep = torch.tensor([500], device=device, dtype=precision) + + video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) + positions = patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) + latents = patchifier.patchify(hidden_states) + + audio_frames = 16 + audio_channels = 8 + audio_mel_bins = 16 + audio_latents = torch.randn( + batch_size, + audio_channels, + audio_frames, + audio_mel_bins, + device=device, + dtype=precision, + ) + audio_shape = AudioLatentShape.from_torch_shape(audio_latents.shape) + audio_positions = audio_patchifier.get_patch_grid_bounds(audio_shape, device=audio_latents.device) + audio_tokens = audio_patchifier.patchify(audio_latents) + + video = Modality( + enabled=True, + latent=latents, + timesteps=timestep, + positions=positions, + context=encoder_hidden_states, + context_mask=None, + ) + audio = Modality( + enabled=True, + latent=audio_tokens, + timesteps=timestep, + positions=audio_positions, + context=encoder_hidden_states, + context_mask=None, + ) + + fastvideo_video = FastVideoModality( + enabled=True, + latent=latents, + timesteps=timestep, + positions=positions, + context=encoder_hidden_states, + context_mask=None, + ) + fastvideo_audio = FastVideoModality( + enabled=True, + latent=audio_tokens, + timesteps=timestep, + positions=audio_positions, + context=encoder_hidden_states, + context_mask=None, + ) + + with torch.no_grad(): + _, ref_audio_out = reference_model( + video=video, + audio=audio, + perturbations=BatchedPerturbationConfig.empty(batch_size), + ) + ref_audio_out = audio_patchifier.unpatchify(ref_audio_out, output_shape=audio_shape) + with set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=None, + ): + _, fastvideo_audio_out = fastvideo_model.model( + video=fastvideo_video, + audio=fastvideo_audio, + ) + fastvideo_audio_out = audio_patchifier.unpatchify(fastvideo_audio_out, output_shape=audio_shape) + + assert ref_audio_out.shape == fastvideo_audio_out.shape + assert ref_audio_out.dtype == fastvideo_audio_out.dtype + # With TORCH_SDPA backend for both, use same tolerance as video parity test + assert_close(ref_audio_out, fastvideo_audio_out, atol=1e-4, rtol=1e-4) diff --git a/tests/local_tests/vaes/__init__.py b/tests/local_tests/vaes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/local_tests/vaes/test_ltx2_audio_vae.py b/tests/local_tests/vaes/test_ltx2_audio_vae.py new file mode 100644 index 000000000..e29f8b884 --- /dev/null +++ b/tests/local_tests/vaes/test_ltx2_audio_vae.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +from pathlib import Path +import sys + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import load_file +from torch.testing import assert_close + + +repo_root = Path(__file__).resolve().parents[3] +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.models.audio.ltx2_audio_vae import ( + LTX2AudioDecoder, + LTX2AudioEncoder, + LTX2Vocoder, +) + + +def _load_metadata(path: Path) -> dict: + with safe_open(str(path), framework="pt") as f: + meta = f.metadata() + if not meta or "config" not in meta: + raise KeyError("Missing config metadata in safetensors file.") + return json.loads(meta["config"]) + + +def _load_weights(path: Path) -> dict[str, torch.Tensor]: + print(f"[LTX2 AUDIO VAE TEST] Loading weights from {path}") + return load_file(str(path)) + + +def _select_audio_vae_weights( + weights: dict[str, torch.Tensor], prefix: str +) -> dict[str, torch.Tensor]: + filtered: dict[str, torch.Tensor] = {} + alt_prefix = prefix.replace("audio_vae.", "") + for name, tensor in weights.items(): + if name.startswith(prefix): + filtered[name.replace(prefix, "")] = tensor + elif alt_prefix and name.startswith(alt_prefix): + filtered[name.replace(alt_prefix, "")] = tensor + elif name.startswith("audio_vae.per_channel_statistics."): + filtered[name.replace("audio_vae.", "")] = tensor + elif name.startswith("per_channel_statistics."): + filtered[name] = tensor + print(f"[LTX2 AUDIO VAE TEST] Selected {len(filtered)} tensors for {prefix}") + return filtered + + +def _select_vocoder_weights( + weights: dict[str, torch.Tensor] +) -> dict[str, torch.Tensor]: + if any(name.startswith("vocoder.") for name in weights): + filtered = { + name.replace("vocoder.", ""): tensor + for name, tensor in weights.items() + if name.startswith("vocoder.") + } + else: + filtered = dict(weights) + print(f"[LTX2 AUDIO VAE TEST] Selected {len(filtered)} tensors for vocoder.") + return filtered + + +def _load_into_model( + model: torch.nn.Module, weights: dict[str, torch.Tensor] +) -> tuple[int, list[str]]: + model_state = model.state_dict() + filtered = { + k: v + for k, v in weights.items() + if k in model_state and model_state[k].shape == v.shape + } + missing = [k for k in model_state.keys() if k not in filtered] + print( + f"[LTX2 AUDIO VAE TEST] Loading {len(filtered)} / {len(model_state)} tensors " + f"from {len(weights)} available" + ) + if not filtered: + return 0, missing + model.load_state_dict(filtered, strict=False) + return len(filtered), missing + + +def test_ltx2_audio_vae_vocoder_parity(): + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + official_path = Path( + os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + ) + audio_vae_path = Path( + os.getenv("LTX2_AUDIO_VAE_PATH", str(diffusers_root / "audio_vae")) + ) + vocoder_path = Path( + os.getenv("LTX2_VOCODER_PATH", str(diffusers_root / "vocoder")) + ) + if not official_path.exists(): + pytest.skip(f"LTX-2 weights not found at {official_path}") + if not audio_vae_path.exists(): + pytest.skip(f"LTX-2 audio VAE weights not found at {audio_vae_path}") + if not vocoder_path.exists(): + pytest.skip(f"LTX-2 vocoder weights not found at {vocoder_path}") + + config = _load_metadata(official_path) + if "audio_vae" not in config or "vocoder" not in config: + pytest.skip("Audio VAE or vocoder config not found in safetensors metadata.") + + try: + from ltx_core.model.audio_vae import ( + AudioDecoderConfigurator, + AudioEncoderConfigurator, + VocoderConfigurator, + ) + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + ref_weights = _load_weights(official_path) + encoder_weights = _select_audio_vae_weights(ref_weights, "audio_vae.encoder.") + decoder_weights = _select_audio_vae_weights(ref_weights, "audio_vae.decoder.") + vocoder_weights = _select_vocoder_weights(ref_weights) + if not encoder_weights or not decoder_weights or not vocoder_weights: + pytest.skip("Audio VAE or vocoder weights not found in safetensors file.") + + fastvideo_audio_weights_path = audio_vae_path / "model.safetensors" + fastvideo_vocoder_weights_path = vocoder_path / "model.safetensors" + if not fastvideo_audio_weights_path.exists(): + pytest.skip( + f"FastVideo audio VAE weights not found at {fastvideo_audio_weights_path}" + ) + if not fastvideo_vocoder_weights_path.exists(): + pytest.skip( + f"FastVideo vocoder weights not found at {fastvideo_vocoder_weights_path}" + ) + fastvideo_audio_weights = _load_weights(fastvideo_audio_weights_path) + fastvideo_encoder_weights = _select_audio_vae_weights( + fastvideo_audio_weights, "encoder." + ) + fastvideo_decoder_weights = _select_audio_vae_weights( + fastvideo_audio_weights, "decoder." + ) + fastvideo_vocoder_weights = _select_vocoder_weights( + _load_weights(fastvideo_vocoder_weights_path) + ) + if (not fastvideo_encoder_weights or not fastvideo_decoder_weights + or not fastvideo_vocoder_weights): + pytest.skip("FastVideo audio VAE/vocoder weights not found in diffusers files.") + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 if torch.cuda.is_available() else torch.float32 + + fastvideo_encoder = LTX2AudioEncoder(config).to(device=device, dtype=precision) + fastvideo_decoder = LTX2AudioDecoder(config).to(device=device, dtype=precision) + fastvideo_vocoder = LTX2Vocoder(config).to(device=device, dtype=precision) + + ref_encoder = AudioEncoderConfigurator.from_config(config).to( + device=device, dtype=precision + ) + ref_decoder = AudioDecoderConfigurator.from_config(config).to( + device=device, dtype=precision + ) + ref_vocoder = VocoderConfigurator.from_config(config).to( + device=device, dtype=precision + ) + + loaded_fastvideo_encoder, missing_fastvideo_encoder = _load_into_model( + fastvideo_encoder.model, fastvideo_encoder_weights + ) + loaded_ref_encoder, missing_ref_encoder = _load_into_model( + ref_encoder, encoder_weights + ) + loaded_fastvideo_decoder, missing_fastvideo_decoder = _load_into_model( + fastvideo_decoder.model, fastvideo_decoder_weights + ) + loaded_ref_decoder, missing_ref_decoder = _load_into_model( + ref_decoder, decoder_weights + ) + loaded_fastvideo_vocoder, missing_fastvideo_vocoder = _load_into_model( + fastvideo_vocoder.model, fastvideo_vocoder_weights + ) + loaded_ref_vocoder, missing_ref_vocoder = _load_into_model( + ref_vocoder, vocoder_weights + ) + + if min( + loaded_fastvideo_encoder, + loaded_ref_encoder, + loaded_fastvideo_decoder, + loaded_ref_decoder, + loaded_fastvideo_vocoder, + loaded_ref_vocoder, + ) == 0: + pytest.skip("Failed to load audio VAE or vocoder weights.") + if ( + missing_fastvideo_encoder + or missing_ref_encoder + or missing_fastvideo_decoder + or missing_ref_decoder + or missing_fastvideo_vocoder + or missing_ref_vocoder + ): + print( + f"[LTX2 AUDIO VAE TEST] Missing encoder keys: {len(missing_fastvideo_encoder)}" + ) + print( + f"[LTX2 AUDIO VAE TEST] Missing decoder keys: {len(missing_fastvideo_decoder)}" + ) + print( + f"[LTX2 AUDIO VAE TEST] Missing vocoder keys: {len(missing_fastvideo_vocoder)}" + ) + pytest.skip("Missing audio VAE/vocoder keys; cannot ensure parity.") + + fastvideo_encoder.model.eval() + fastvideo_decoder.model.eval() + fastvideo_vocoder.model.eval() + ref_encoder.eval() + ref_decoder.eval() + ref_vocoder.eval() + + ddconfig = config["audio_vae"]["model"]["params"]["ddconfig"] + in_channels = ddconfig.get("in_channels", 2) + resolution = ddconfig.get("resolution", 256) + mel_bins = ddconfig.get("mel_bins", 64) + batch_size = 1 + spectrogram = torch.randn( + batch_size, + in_channels, + resolution, + mel_bins, + device=device, + dtype=precision, + ) + + with torch.no_grad(): + ref_latents = ref_encoder(spectrogram) + fast_latents = fastvideo_encoder(spectrogram) + + assert ref_latents.shape == fast_latents.shape + assert ref_latents.dtype == fast_latents.dtype + assert torch.isfinite(ref_latents).all(), "Reference encoder produced non-finite latents." + assert torch.isfinite(fast_latents).all(), "FastVideo encoder produced non-finite latents." + assert_close(ref_latents, fast_latents, atol=1e-2, rtol=1e-2) + + with torch.no_grad(): + ref_decoded = ref_decoder(ref_latents) + fast_decoded = fastvideo_decoder(ref_latents) + + assert ref_decoded.shape == fast_decoded.shape + assert ref_decoded.dtype == fast_decoded.dtype + assert torch.isfinite(ref_decoded).all(), "Reference decoder produced non-finite output." + assert torch.isfinite(fast_decoded).all(), "FastVideo decoder produced non-finite output." + assert_close(ref_decoded, fast_decoded, atol=1e-2, rtol=1e-2) + + with torch.no_grad(): + ref_audio = ref_vocoder(ref_decoded) + fast_audio = fastvideo_vocoder(ref_decoded) + + assert ref_audio.shape == fast_audio.shape + assert ref_audio.dtype == fast_audio.dtype + assert torch.isfinite(ref_audio).all(), "Reference vocoder produced non-finite audio." + assert torch.isfinite(fast_audio).all(), "FastVideo vocoder produced non-finite audio." + assert_close(ref_audio, fast_audio, atol=1e-2, rtol=1e-2) diff --git a/tests/local_tests/vaes/test_ltx2_vae.py b/tests/local_tests/vaes/test_ltx2_vae.py new file mode 100644 index 000000000..cff218203 --- /dev/null +++ b/tests/local_tests/vaes/test_ltx2_vae.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +from pathlib import Path +import sys + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import load_file +from torch.testing import assert_close + + +repo_root = Path(__file__).resolve().parents[3] +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.models.vaes.ltx2vae import LTX2VideoDecoder, LTX2VideoEncoder + + +def _load_metadata(path: Path) -> dict: + with safe_open(str(path), framework="pt") as f: + meta = f.metadata() + if not meta or "config" not in meta: + raise KeyError("Missing config metadata in safetensors file.") + return json.loads(meta["config"]) + + +def _load_weights(path: Path) -> dict[str, torch.Tensor]: + print(f"[LTX2 VAE TEST] Loading weights from {path}") + return load_file(str(path)) + + +def _select_vae_weights(weights: dict[str, torch.Tensor], prefix: str) -> dict[str, torch.Tensor]: + filtered: dict[str, torch.Tensor] = {} + alt_prefix = prefix.replace("vae.", "") + for name, tensor in weights.items(): + if name.startswith(prefix): + filtered[name.replace(prefix, "")] = tensor + elif alt_prefix and name.startswith(alt_prefix): + filtered[name.replace(alt_prefix, "")] = tensor + elif name.startswith("vae.per_channel_statistics."): + filtered[name.replace("vae.", "")] = tensor + elif name.startswith("per_channel_statistics."): + filtered[name] = tensor + print(f"[LTX2 VAE TEST] Selected {len(filtered)} tensors for {prefix}") + return filtered + + +def _load_into_model(model: torch.nn.Module, weights: dict[str, torch.Tensor]) -> tuple[int, list[str]]: + model_state = model.state_dict() + filtered = { + k: v + for k, v in weights.items() + if k in model_state and model_state[k].shape == v.shape + } + missing = [k for k in model_state.keys() if k not in filtered] + print( + f"[LTX2 VAE TEST] Loading {len(filtered)} / {len(model_state)} tensors " + f"from {len(weights)} available" + ) + if not filtered: + return 0, missing + model.load_state_dict(filtered, strict=False) + return len(filtered), missing + + +def test_ltx2_vae_parity(): + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + official_path = Path( + os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + ) + fastvideo_path = Path( + os.getenv("LTX2_VAE_PATH", str(diffusers_root / "vae")) + ) + if not official_path.exists(): + pytest.skip(f"LTX-2 weights not found at {official_path}") + if not fastvideo_path.exists(): + pytest.skip(f"LTX-2 diffusers VAE not found at {fastvideo_path}") + + config = _load_metadata(official_path) + if "vae" not in config: + pytest.skip("VAE config not found in safetensors metadata.") + + try: + from ltx_core.model.video_vae import VideoDecoderConfigurator, VideoEncoderConfigurator + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + ref_weights = _load_weights(official_path) + encoder_weights = _select_vae_weights(ref_weights, "vae.encoder.") + decoder_weights = _select_vae_weights(ref_weights, "vae.decoder.") + if not encoder_weights or not decoder_weights: + pytest.skip("VAE weights not found in safetensors file.") + + fastvideo_weights_path = fastvideo_path / "model.safetensors" + if not fastvideo_weights_path.exists(): + pytest.skip(f"FastVideo VAE weights not found at {fastvideo_weights_path}") + fastvideo_weights = _load_weights(fastvideo_weights_path) + fastvideo_encoder_weights = _select_vae_weights( + fastvideo_weights, "encoder." + ) + fastvideo_decoder_weights = _select_vae_weights( + fastvideo_weights, "decoder." + ) + if not fastvideo_encoder_weights or not fastvideo_decoder_weights: + pytest.skip("FastVideo VAE weights not found in diffusers file.") + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 if torch.cuda.is_available() else torch.float32 + + fastvideo_encoder = LTX2VideoEncoder(config).to(device=device, dtype=precision) + fastvideo_decoder = LTX2VideoDecoder(config).to(device=device, dtype=precision) + ref_encoder = VideoEncoderConfigurator.from_config(config).to(device=device, dtype=precision) + ref_decoder = VideoDecoderConfigurator.from_config(config).to(device=device, dtype=precision) + + loaded_fastvideo_encoder, missing_fastvideo_encoder = _load_into_model( + fastvideo_encoder.model, fastvideo_encoder_weights + ) + loaded_ref_encoder, missing_ref_encoder = _load_into_model(ref_encoder, encoder_weights) + loaded_fastvideo_decoder, missing_fastvideo_decoder = _load_into_model( + fastvideo_decoder.model, fastvideo_decoder_weights + ) + loaded_ref_decoder, missing_ref_decoder = _load_into_model(ref_decoder, decoder_weights) + + if min( + loaded_fastvideo_encoder, + loaded_ref_encoder, + loaded_fastvideo_decoder, + loaded_ref_decoder, + ) == 0: + pytest.skip("Failed to load VAE weights into one or more models.") + if ( + missing_fastvideo_encoder + or missing_ref_encoder + or missing_fastvideo_decoder + or missing_ref_decoder + ): + print(f"[LTX2 VAE TEST] Missing encoder keys: {len(missing_fastvideo_encoder)}") + print(f"[LTX2 VAE TEST] Missing decoder keys: {len(missing_fastvideo_decoder)}") + pytest.skip("Missing VAE keys; cannot ensure parity.") + + fastvideo_encoder.model.eval() + fastvideo_decoder.model.eval() + ref_encoder.eval() + ref_decoder.eval() + + fastvideo_decoder.model.decode_noise_scale = 0.0 + ref_decoder.decode_noise_scale = 0.0 + + batch_size = 1 + frames = 9 + height = 64 + width = 64 + video = torch.randn( + batch_size, + 3, + frames, + height, + width, + device=device, + dtype=precision, + ) + + with torch.no_grad(): + ref_latents = ref_encoder(video) + fast_latents = fastvideo_encoder(video) + + assert ref_latents.shape == fast_latents.shape + assert ref_latents.dtype == fast_latents.dtype + assert torch.isfinite(ref_latents).all(), "Reference encoder produced non-finite latents." + assert torch.isfinite(fast_latents).all(), "FastVideo encoder produced non-finite latents." + assert_close(ref_latents, fast_latents, atol=1e-2, rtol=1e-2) + + timestep = torch.tensor([0.05], device=device, dtype=precision) + with torch.no_grad(): + ref_decoded = ref_decoder(ref_latents, timestep=timestep) + fast_decoded = fastvideo_decoder(fast_latents, timestep=timestep) + + assert ref_decoded.shape == fast_decoded.shape + assert ref_decoded.dtype == fast_decoded.dtype + assert torch.isfinite(ref_decoded).all(), "Reference decoder produced non-finite output." + assert torch.isfinite(fast_decoded).all(), "FastVideo decoder produced non-finite output." + assert_close(ref_decoded, fast_decoded, atol=1e-2, rtol=1e-2) diff --git a/tests/local_tests/vaes/test_ltx2_vae_official.py b/tests/local_tests/vaes/test_ltx2_vae_official.py new file mode 100644 index 000000000..c02da5875 --- /dev/null +++ b/tests/local_tests/vaes/test_ltx2_vae_official.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +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] +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.configs.models.vaes import LTX2VAEConfig +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.loader.component_loader import VAELoader + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="LTX-2 VAE parity test requires CUDA.", +) +def test_ltx2_vae_parity_official(): + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + official_path = Path( + os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + ) + fastvideo_path = Path( + os.getenv("LTX2_VAE_PATH", str(diffusers_root / "vae")) + ) + if not official_path.exists(): + pytest.skip(f"LTX-2 weights not found at {official_path}") + if not fastvideo_path.exists(): + pytest.skip(f"LTX-2 diffusers VAE not found at {fastvideo_path}") + + try: + from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder + from ltx_core.model.video_vae import ( + VAE_DECODER_COMFY_KEYS_FILTER, + VAE_ENCODER_COMFY_KEYS_FILTER, + VideoDecoderConfigurator, + VideoEncoderConfigurator, + ) + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + device = torch.device("cuda:0") + precision = torch.bfloat16 + precision_str = "bf16" + + args = FastVideoArgs( + model_path=str(fastvideo_path), + vae_cpu_offload=False, + pipeline_config=PipelineConfig( + vae_config=LTX2VAEConfig(), + vae_precision=precision_str, + ), + ) + + loader = VAELoader() + fastvideo_vae = loader.load(str(fastvideo_path), args).to( + device=device, dtype=precision + ) + + encoder_builder = SingleGPUModelBuilder( + model_class_configurator=VideoEncoderConfigurator, + model_path=str(official_path), + model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER, + ) + decoder_builder = SingleGPUModelBuilder( + model_class_configurator=VideoDecoderConfigurator, + model_path=str(official_path), + model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER, + ) + ref_encoder = encoder_builder.build( + device=device, dtype=precision + ).to(device=device, dtype=precision) + ref_decoder = decoder_builder.build( + device=device, dtype=precision + ).to(device=device, dtype=precision) + + fastvideo_vae.encoder.eval() + fastvideo_vae.decoder.eval() + ref_encoder.eval() + ref_decoder.eval() + + if hasattr(fastvideo_vae.decoder, "decode_noise_scale"): + fastvideo_vae.decoder.decode_noise_scale = 0.0 + if hasattr(ref_decoder, "decode_noise_scale"): + ref_decoder.decode_noise_scale = 0.0 + + batch_size = 1 + frames = 9 + height = 64 + width = 64 + video = torch.randn( + batch_size, + 3, + frames, + height, + width, + device=device, + dtype=precision, + ) + + with torch.no_grad(): + ref_latents = ref_encoder(video) + fast_latents = fastvideo_vae.encoder(video) + + assert ref_latents.shape == fast_latents.shape + assert ref_latents.dtype == fast_latents.dtype + assert torch.isfinite(ref_latents).all(), "Reference encoder produced non-finite latents." + assert torch.isfinite(fast_latents).all(), "FastVideo encoder produced non-finite latents." + assert_close(ref_latents, fast_latents, atol=1e-2, rtol=1e-2) + + timestep = torch.tensor([0.05], device=device, dtype=precision) + with torch.no_grad(): + ref_decoded = ref_decoder(ref_latents, timestep=timestep) + fast_decoded = fastvideo_vae.decoder(fast_latents, timestep=timestep) + + assert ref_decoded.shape == fast_decoded.shape + assert ref_decoded.dtype == fast_decoded.dtype + assert torch.isfinite(ref_decoded).all(), "Reference decoder produced non-finite output." + assert torch.isfinite(fast_decoded).all(), "FastVideo decoder produced non-finite output." + assert_close(ref_decoded, fast_decoded, atol=1e-2, rtol=1e-2)