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/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/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 827de11eb..37c542375 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -837,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 39f31fd3e..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 = { @@ -68,7 +69,8 @@ _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 = {