diff --git a/.gitignore b/.gitignore index 6fbe89528..42278d4ef 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,8 @@ cache_dir/ wandb/ venv/ .venv/ +/venv-port-magi-2/ +/archived/ runs/ samples/ Miniconda3-latest-Linux-x86_64.sh diff --git a/.gitmodules b/.gitmodules index d95585605..33e94f2d2 100644 --- a/.gitmodules +++ b/.gitmodules @@ -7,3 +7,12 @@ [submodule "fastvideo/third_party/eval/vbench"] path = fastvideo/third_party/eval/vbench url = https://github.com/Vchitect/VBench.git +[submodule "fastvideo/third_party/magi_attention"] + path = fastvideo/third_party/magi_attention + url = https://github.com/SandAI-org/MagiAttention.git +[submodule "fastvideo/third_party/magi_compiler"] + path = fastvideo/third_party/magi_compiler + url = https://github.com/SandAI-org/MagiCompiler.git +[submodule "fastvideo/third_party/flash_attention"] + path = fastvideo/third_party/flash_attention + url = https://github.com/Dao-AILab/flash-attention.git diff --git a/fastvideo/api/compat.py b/fastvideo/api/compat.py index ea3ff10ac..578347b2f 100644 --- a/fastvideo/api/compat.py +++ b/fastvideo/api/compat.py @@ -157,7 +157,7 @@ def legacy_from_pretrained_to_config( preset_refine["num_inference_steps"] = value elif key == "ltx2_refine_guidance_scale": preset_refine["guidance_scale"] = value - elif key in {"enable_stage_verification", "use_fsdp_inference", "disable_autocast"}: + elif key in {"enable_stage_verification", "use_fsdp_inference", "disable_autocast", "deterministic"}: engine[key] = value elif key == "override_text_encoder_quant": quantization["text_encoder_quant"] = value @@ -244,6 +244,7 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An "enable_stage_verification": engine.enable_stage_verification, "use_fsdp_inference": engine.use_fsdp_inference, "disable_autocast": engine.disable_autocast, + "deterministic": engine.deterministic, } if normalized.pipeline.workload_type is not None: kwargs["workload_type"] = normalized.pipeline.workload_type diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index d3dbb0516..98bd85311 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -81,6 +81,7 @@ class EngineConfig: enable_stage_verification: bool = True use_fsdp_inference: bool = False disable_autocast: bool = False + deterministic: bool = False quantization: QuantizationConfig | None = None diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 64a5c118d..4937d1438 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -10,6 +10,7 @@ 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.magi_human import MagiHumanVideoConfig +from fastvideo.configs.models.dits.magi2 import Magi2PreviewVideoConfig, Magi2RefinerVideoConfig from fastvideo.configs.models.dits.stable_audio import StableAudioConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig @@ -19,5 +20,6 @@ __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", + "Magi2PreviewVideoConfig", "Magi2RefinerVideoConfig", "StableAudioConfig", "GlmImageDiTConfig" ] diff --git a/fastvideo/configs/models/dits/magi2.py b/fastvideo/configs/models/dits/magi2.py new file mode 100644 index 000000000..4710be129 --- /dev/null +++ b/fastvideo/configs/models/dits/magi2.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Architecture configurations for the MAGI-2 preview and refiner DiTs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def _is_transformer_layer(name: str, module) -> bool: + """Select a complete MAGI-2 transformer layer as an FSDP shard unit.""" + del module + parts = name.split(".") + return len(parts) >= 3 and parts[0] == "block" and parts[1] == "layers" and parts[2].isdigit() + + +@dataclass +class Magi2AttentionGatingConfig: + """Control the learned per-head attention output gate.""" + + enable: bool = True + + +@dataclass +class Magi2AttentionSinksConfig: + """Configure the learned Flash Attention sink logits.""" + + enable: bool = True + sink_token_num: int = 1 + + +@dataclass +class Magi2MHCConfig: + """Configure multi-stream hyper-connections (MHC).""" + + enable: bool = True + num_stream: int = 4 + alpha_init: float = 0.01 + + +@dataclass +class Magi2MoEConfig: + """Configure the preview model's multi-head mixture of experts.""" + + num_experts: int = 256 + top_k: int = 6 + score_func: str = "sigmoid" + route_norm: bool = True + route_scale: float = 4.9 + moe_layers: list[int] = field(default_factory=lambda: list(range(2, 38))) + expert_intermediate_size: int = 1280 + shared_expert_intermediate_size: int = 1280 + modality_specific_expert_intermediate_size: int = 1280 + num_heads: int = 12 + split_merge_intermediate_size: int | None = None + + +@dataclass +class Magi2PreviewArchConfig(DiTArchConfig): + """Define the published 114B MAGI-2 preview transformer architecture.""" + + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_layer]) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.FLASH_ATTN, ) + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + num_layers: int = 40 + hidden_size: int = 3072 + head_dim: int = 128 + num_query_groups: int = 24 + video_in_channels: int = 48 + audio_in_channels: int = 64 + text_in_channels: int = 5120 + params_dtype: str = "bfloat16" + intermediate_factor: int = 4 + mm_layers: list[int] = field(default_factory=lambda: [0, 1, 38, 39]) + layer_activation_types: dict[int, str] = field(default_factory=dict) + activation_type: str = "swiglu7" + attn_softcap: float = -1.0 + attn_gating: Magi2AttentionGatingConfig = field(default_factory=Magi2AttentionGatingConfig) + attn_sinks: Magi2AttentionSinksConfig = field(default_factory=Magi2AttentionSinksConfig) + mhc_config: Magi2MHCConfig = field(default_factory=Magi2MHCConfig) + moe_config: Magi2MoEConfig = field(default_factory=Magi2MoEConfig) + + num_heads_q: int = 0 + num_heads_kv: int = 0 + num_attention_heads: int = 0 + num_channels_latents: int = 48 + in_channels: int = 48 + out_channels: int = 48 + + def __post_init__(self) -> None: + """Derive attention dimensions and normalize nested checkpoint dictionaries.""" + super().__post_init__() + if isinstance(self.attn_gating, dict): + self.attn_gating = Magi2AttentionGatingConfig(**self.attn_gating) + if isinstance(self.attn_sinks, dict): + self.attn_sinks = Magi2AttentionSinksConfig(**self.attn_sinks) + if isinstance(self.mhc_config, dict): + self.mhc_config = Magi2MHCConfig(**self.mhc_config) + if isinstance(self.moe_config, dict): + self.moe_config = Magi2MoEConfig(**self.moe_config) + self.num_heads_q = self.hidden_size // self.head_dim + self.num_heads_kv = self.num_query_groups + self.num_attention_heads = self.num_heads_q + + +@dataclass +class Magi2RefinerArchConfig(DiTArchConfig): + """Define the published MAGI-2 1080p refiner transformer architecture.""" + + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_layer]) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.FLASH_ATTN, ) + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + num_layers: int = 30 + hidden_size: int = 4096 + head_dim: int = 128 + num_query_groups: int = 8 + video_in_channels: int = 48 + audio_in_channels: int = 64 + text_in_channels: int = 5120 + checkpoint_qk_layernorm_rope: bool = False + params_dtype: str = "bfloat16" + mm_layers: list[int] = field(default_factory=lambda: [0, 1, 28, 29]) + local_attn_layers: list[int] = field(default_factory=lambda: list(range(30))) + enable_attn_gating: bool = True + activation_type: str = "swiglu7" + layer_activation_types: dict[int, str] = field(default_factory=dict) + post_norm_layers: list[int] = field(default_factory=list) + + num_heads_q: int = 0 + num_heads_kv: int = 0 + num_attention_heads: int = 0 + num_channels_latents: int = 48 + in_channels: int = 48 + out_channels: int = 48 + + def __post_init__(self) -> None: + """Derive the published grouped-query attention dimensions.""" + super().__post_init__() + self.num_heads_q = self.hidden_size // self.head_dim + self.num_heads_kv = self.num_query_groups + self.num_attention_heads = self.num_heads_q + + +@dataclass +class Magi2PreviewVideoConfig(DiTConfig): + """Load the MAGI-2 preview transformer through FastVideo's DiT loader.""" + + arch_config: DiTArchConfig = field(default_factory=Magi2PreviewArchConfig) + prefix: str = "magi2_preview" + + +@dataclass +class Magi2RefinerVideoConfig(DiTConfig): + """Load the MAGI-2 refiner transformer through FastVideo's DiT loader.""" + + arch_config: DiTArchConfig = field(default_factory=Magi2RefinerArchConfig) + prefix: str = "magi2_refiner" diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 7636fce23..187cc91f1 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -9,6 +9,7 @@ from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig +from fastvideo.configs.models.encoders.qwen3_5 import Magi2Qwen35ArchConfig, Magi2Qwen35Config from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig, StableAudioConditionerConfig) from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig @@ -17,5 +18,6 @@ __all__ = [ "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig", - "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig" + "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig", + "Magi2Qwen35ArchConfig", "Magi2Qwen35Config" ] diff --git a/fastvideo/configs/models/encoders/qwen3_5.py b/fastvideo/configs/models/encoders/qwen3_5.py new file mode 100644 index 000000000..a65579981 --- /dev/null +++ b/fastvideo/configs/models/encoders/qwen3_5.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configuration for the MAGI-2 Qwen3.5-27B prompt encoder.""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig + + +@dataclass +class Magi2Qwen35ArchConfig(TextEncoderArchConfig): + """Select Qwen3.5 hidden state -3 and the official token limit.""" + + architectures: list[str] = field(default_factory=lambda: ["Magi2Qwen35TextEncoder"]) + vocab_size: int = 248320 + hidden_size: int = 5120 + intermediate_size: int = 17408 + num_hidden_layers: int = 64 + num_attention_heads: int = 24 + num_key_value_heads: int = 4 + head_dim: int = 256 + hidden_act: str = "silu" + max_position_embeddings: int = 262144 + rms_norm_eps: float = 1e-6 + attention_bias: bool = False + attention_dropout: float = 0.0 + attn_output_gate: bool = True + linear_conv_kernel_dim: int = 4 + linear_key_head_dim: int = 128 + linear_value_head_dim: int = 128 + linear_num_key_heads: int = 16 + linear_num_value_heads: int = 48 + full_attention_interval: int = 4 + rope_theta: float = 10000000.0 + partial_rotary_factor: float = 0.25 + mrope_section: tuple[int, int, int] = (11, 11, 10) + pad_token_id: int | None = None + eos_token_id: int = 248044 + text_len: int = 7000 + hidden_state_skip_layer: int = 2 + output_hidden_states: bool = True + tokenizer_kwargs: dict = field(default_factory=lambda: {"padding_side": "right"}) + _fsdp_shard_conditions: list = field(default_factory=list) + + +@dataclass +class Magi2Qwen35Config(TextEncoderConfig): + """Load the published Qwen3.5 checkpoint through the FastVideo loader.""" + + arch_config: TextEncoderArchConfig = field(default_factory=Magi2Qwen35ArchConfig) + prefix: str = "magi2_qwen35" diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index 6f3ca9768..f30b157fa 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -6,6 +6,8 @@ from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig 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.magi2_turbo_vae import Magi2TurboVAEArchConfig, Magi2TurboVAEConfig +from fastvideo.configs.models.vaes.magi2_wanvae import Magi2WanVAEArchConfig, Magi2WanVAEConfig from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig from fastvideo.configs.models.vaes.wanvae import WanVAEConfig @@ -19,6 +21,10 @@ __all__ = [ "Gen3CVAEConfig", "Hunyuan15VAEConfig", "LTX2VAEConfig", + "Magi2TurboVAEArchConfig", + "Magi2TurboVAEConfig", + "Magi2WanVAEArchConfig", + "Magi2WanVAEConfig", "OobleckVAEArchConfig", "OobleckVAEConfig", "Flux2VAEConfig", diff --git a/fastvideo/configs/models/vaes/magi2_turbo_vae.py b/fastvideo/configs/models/vaes/magi2_turbo_vae.py new file mode 100644 index 000000000..6b172316c --- /dev/null +++ b/fastvideo/configs/models/vaes/magi2_turbo_vae.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configuration for the distilled MAGI-2 Turbo VAE decoder.""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class Magi2TurboVAEArchConfig(VAEArchConfig): + """Describe the latent and output geometry used by Turbo VAE.""" + + architectures: list[str] = field(default_factory=lambda: ["Magi2TurboVAEModel"]) + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 48 + decoder_block_out_channels: tuple[int, ...] = (64, 128, 256, 512) + decoder_causal: bool = False + decoder_layers_per_block: tuple[int, ...] = (2, 2, 2, 3, 3) + patch_size: int = 2 + patch_size_t: int = 1 + resnet_norm_eps: float = 1e-6 + scaling_factor: float = 1.0 + decoder_spatio_temporal_scaling: tuple[bool, ...] = (False, True, True, True) + decoder_spatio_only: tuple[bool, ...] = (False, True, False, False) + decoder_is_dw_conv: tuple[bool, ...] = (False, False, False, False, False) + decoder_dw_kernel_size: int = 5 + temporal_compression_ratio: int = 4 + spatial_compression_ratio: int = 16 + first_chunk_size: int = 7 + step_size: int = 7 + use_unpatchify: bool = True + + +@dataclass +class Magi2TurboVAEConfig(VAEConfig): + """Provide the published Turbo VAE JSON and checkpoint paths.""" + + arch_config: VAEArchConfig = field(default_factory=Magi2TurboVAEArchConfig) + config_path: str = "" + checkpoint_path: str = "" + pretrained_dtype: str = "bfloat16" + load_encoder: bool = False + load_decoder: bool = True + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False diff --git a/fastvideo/configs/models/vaes/magi2_wanvae.py b/fastvideo/configs/models/vaes/magi2_wanvae.py new file mode 100644 index 000000000..c9b707fc9 --- /dev/null +++ b/fastvideo/configs/models/vaes/magi2_wanvae.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan 2.2 VAE configuration for MAGI-2 I2V reference images.""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig, WanVAEConfig + + +MAGI2_WAN_LATENTS_MEAN = ( + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667, +) + +MAGI2_WAN_LATENTS_STD = ( + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744, +) + + +@dataclass +class Magi2WanVAEArchConfig(WanVAEArchConfig): + """Match the published Wan2.2 VAE encoder architecture and normalization.""" + + in_channels: int = 12 + out_channels: int = 12 + base_dim: int = 160 + decoder_base_dim: int | None = 256 + z_dim: int = 48 + temperal_downsample: tuple[bool, ...] = (False, True, True) + is_residual: bool = True + clip_output: bool = False + latents_mean: tuple[float, ...] = MAGI2_WAN_LATENTS_MEAN + latents_std: tuple[float, ...] = MAGI2_WAN_LATENTS_STD + patch_size: int | None = 2 + scale_factor_temporal: int = 4 + scale_factor_spatial: int = 16 + + +@dataclass +class Magi2WanVAEConfig(WanVAEConfig): + """Load only the Wan encoder used by MAGI-2 image conditioning.""" + + arch_config: WanVAEArchConfig = field(default_factory=Magi2WanVAEArchConfig) + load_encoder: bool = True + load_decoder: bool = False diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 3983479de..6fbf50b54 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -9,6 +9,7 @@ from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig +from fastvideo.configs.pipelines.magi2 import Magi2PreviewPipelineConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig from fastvideo.registry import get_pipeline_config_cls_from_name from fastvideo.configs.pipelines.wan import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig, @@ -19,5 +20,6 @@ __all__ = [ "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig", - "Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name" + "Kandinsky5I2VConfig", "Magi2PreviewPipelineConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", + "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/magi2.py b/fastvideo/configs/pipelines/magi2.py new file mode 100644 index 000000000..f5406860a --- /dev/null +++ b/fastvideo/configs/pipelines/magi2.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline configuration for the MAGI-2 Preview 1080p release profile.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models import DiTConfig, VAEConfig +from fastvideo.configs.models.dits import ( + Magi2PreviewVideoConfig, + Magi2RefinerVideoConfig, +) +from fastvideo.configs.models.vaes import Magi2TurboVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +@dataclass +class Magi2PreviewPipelineConfig(PipelineConfig): + """Configure the published 10-second MAGI-2 Preview pipeline.""" + + dit_config: DiTConfig = field(default_factory=Magi2PreviewVideoConfig) + refiner_dit_config: DiTConfig = field(default_factory=Magi2RefinerVideoConfig) + vae_config: VAEConfig = field(default_factory=Magi2TurboVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + text_encoder_configs: tuple = field(default_factory=tuple) + preprocess_text_funcs: tuple = field(default_factory=tuple) + postprocess_text_funcs: tuple = field(default_factory=tuple) + + dit_precision: str = "bf16" + vae_precision: str = "bf16" + text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple) + + preview_width: int = 896 + preview_height: int = 512 + output_width: int = 1920 + output_height: int = 1088 + output_frames: int = 249 + output_fps: int = 25 + + preview_video_channels: int = 48 + preview_video_length: int = 32 + preview_latent_height: int = 32 + preview_latent_width: int = 56 + audio_channels: int = 64 + audio_latent_length: int = 250 + refiner_latent_height: int = 68 + refiner_latent_width: int = 120 + + preview_flow_shift: float = 7.0 + preview_video_guidance_scale: float = 5.0 + preview_audio_guidance_scale: float = 7.0 + refiner_flow_shift: float = 5.0 + refiner_video_guidance_scale: float = 2.0 + refiner_audio_guidance_scale: float = 5.0 + refiner_noise_index: int = 220 + + +__all__ = ["Magi2PreviewPipelineConfig"] diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 1588d65bb..b5db1bd37 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -46,6 +46,9 @@ if TYPE_CHECKING: FASTVIDEO_CFG_GATE_STEP: float = 1.0 FASTVIDEO_HOST_IP: str = "" FASTVIDEO_LOOPBACK_IP: str = "" + MAGI2_DETERMINISTIC: bool = False + MAGI2_SAVE_LATENT_PATH: str | None = None + MAGI_ATTENTION_DETERMINISTIC_MODE: bool = False def get_default_cache_root() -> str: @@ -320,6 +323,14 @@ environment_variables: dict[str, Callable[[], Any]] = { # - Wan2.2 high/low-noise expert switch invalidates the cache. "FASTVIDEO_CFG_GATE_STEP": lambda: float(os.getenv("FASTVIDEO_CFG_GATE_STEP", "1.0")), + + # MAGI-2 reproducibility and post-refiner latent capture. + "MAGI2_DETERMINISTIC": + lambda: os.getenv("MAGI2_DETERMINISTIC", "0") == "1", + "MAGI2_SAVE_LATENT_PATH": + lambda: os.getenv("MAGI2_SAVE_LATENT_PATH"), + "MAGI_ATTENTION_DETERMINISTIC_MODE": + lambda: os.getenv("MAGI_ATTENTION_DETERMINISTIC_MODE", "0") == "1", } # end-env-vars-definition diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index b27706a56..92d1f8645 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -152,6 +152,7 @@ class FastVideoArgs: torch_compile_kwargs_audio_vae: dict[str, Any] = field(default_factory=dict) disable_autocast: bool = False + deterministic: bool = False # VSA parameters VSA_sparsity: float = 0.0 # inference/validation sparsity @@ -596,6 +597,12 @@ class FastVideoArgs: action=StoreBoolean, help="Disable autocast for denoising loop and vae decoding in pipeline sampling", ) + parser.add_argument( + "--deterministic", + action=StoreBoolean, + default=FastVideoArgs.deterministic, + help="Enable deterministic model kernels when the pipeline supports them", + ) # VSA parameters parser.add_argument( diff --git a/fastvideo/models/dits/magi2.py b/fastvideo/models/dits/magi2.py new file mode 100644 index 000000000..b1235b1e1 --- /dev/null +++ b/fastvideo/models/dits/magi2.py @@ -0,0 +1,3523 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MAGI-2 preview model (magi2_preview) - all layers and model definition.""" + +from __future__ import annotations + +import logging +import math +import os +from dataclasses import dataclass +from enum import Enum, IntEnum +from functools import partial +from typing import Callable, List, Literal, Optional, Tuple, TypeAlias, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange, repeat +from flash_attn_interface import flash_attn_3_cuda +from magi_attn_extensions.fa3_interface_with_sink import ( + FA3VarlenFuncWithSink, + fa3_func_with_sink, + fa3_varlen_func_with_sink, +) +from magi_compiler import magi_compile, magi_register_custom_op +from fastvideo.configs.models.dits.magi2 import ( + Magi2MoEConfig, + Magi2PreviewArchConfig, + Magi2PreviewVideoConfig, +) +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.magi2_runtime import psm +from fastvideo.models.dits.magi2_runtime.all_to_all_primitive import ( + batch_scatter_head_gather_seqlen, + scatter_seqlen_gather_head, +) +from fastvideo.models.dits.magi2_runtime.context_parallel import ulysses_scheduler +from fastvideo.models.dits.magi2_runtime.flash_mh_moe import ( + compute_topk_probs_and_indices as _mh_moe_compute_topk, + flash_mh_moe_fwd as _mh_moe_fwd, + flash_mh_moe_global_sort as _mh_moe_global_sort, +) +from torch.library import triton_op, wrap_triton + + +logger = logging.getLogger(__name__) + +ModelConfig = Magi2PreviewArchConfig +MoEConfig = Magi2MoEConfig + + +def env_is_true(name: str) -> bool: + """Interpret one MAGI-2 environment variable as a Boolean flag.""" + return os.environ.get(name, "0").lower() in {"1", "true", "yes", "on"} + + +def _resolve_params_dtype(value: torch.dtype | str) -> torch.dtype: + """Convert a serialized MAGI-2 parameter dtype into a PyTorch dtype.""" + if isinstance(value, torch.dtype): + return value + dtype_by_name = { + "bfloat16": torch.bfloat16, + "torch.bfloat16": torch.bfloat16, + "float16": torch.float16, + "torch.float16": torch.float16, + "float32": torch.float32, + "torch.float32": torch.float32, + } + try: + return dtype_by_name[value] + except KeyError as error: + raise ValueError(f"Unsupported MAGI-2 parameter dtype: {value!r}") from error + +# ============================================================ +# enum +# ============================================================ + + +class MLPActivationType(Enum): + """Enumeration of supported activation functions for MLP""" + + SWIGLU7 = "swiglu7" + + +class Modality(IntEnum): + VIDEO = 0 + AUDIO = 1 + TEXT = 2 + # Special marker modality used ONLY at input to identify time tokens. + # The model will remap TIME -> TEXT internally before passing to ModalityDispatcher, + # so we do NOT need to increase num_modality or change parameter shapes. + TIME = 3 + + +# ============================================================ +# activation +# ============================================================ + + +@torch.compile +def swiglu7( + x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None +): + """Apply the bounded SwiGLU7 activation in the requested output dtype.""" + out_dtype = x.dtype if out_dtype is None else out_dtype + x = x.to(torch.float32) + x_glu, x_linear = x[..., ::2], x[..., 1::2] + # Clamp the input values + x_glu = x_glu.clamp(min=None, max=limit) + x_linear = x_linear.clamp(min=-limit, max=limit) + out_glu = x_glu * torch.sigmoid(alpha * x_glu) + # Note we add an extra bias of 1 to the linear layer (from GPT-OSS) + return (out_glu * (x_linear + 1)).to(out_dtype) + + +def create_activation_func(activation_type: MLPActivationType) -> Callable: + if activation_type is MLPActivationType.SWIGLU7: + return swiglu7 + raise ValueError(f"Unknown activation type: {activation_type}") + + +# ============================================================ +# _common +# ============================================================ + + + + +def _maybe_gather( + input: torch.Tensor, gather_ids: Optional[torch.Tensor] +) -> torch.Tensor: + if gather_ids is None: + return input + gather_ids = gather_ids.to(input.device) + if gather_ids.dim() == 1: + return input.index_select(0, gather_ids) + return torch.gather(input, 0, gather_ids) + + +def _expand_grouped_bias(bias: torch.Tensor, m_splits: list[int]) -> torch.Tensor: + chunks = [] + for expert_idx, split in enumerate(m_splits): + if split > 0: + chunks.append(bias[expert_idx].expand(split, -1)) + return torch.cat(chunks, dim=0) if chunks else bias.new_empty((0, bias.size(-1))) + + +def _m_splits_from( + cu_seqlens: Optional[torch.Tensor], m_splits: Optional[list[int] | torch.Tensor] +) -> list[int]: + if isinstance(m_splits, torch.Tensor): + return [int(v) for v in m_splits.detach().cpu().tolist()] + if m_splits is not None: + return [int(v) for v in m_splits] + if cu_seqlens is None: + raise ValueError("m_splits or cu_seqlens is required for grouped linear") + return [int(v) for v in torch.diff(cu_seqlens).detach().cpu().tolist()] + + +def _torch_grouped_linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + m_splits: list[int], + gather_ids: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Apply expert-specific linear projections to contiguous token groups.""" + ordered_input = _maybe_gather(input, gather_ids) + outputs = [] + start = 0 + for expert_idx, split in enumerate(m_splits): + end = start + split + if end > start: + outputs.append( + F.linear( + ordered_input[start:end], + weight[expert_idx], + None if bias is None else bias[expert_idx], + ) + ) + start = end + return ( + torch.cat(outputs, dim=0) + if outputs + else ordered_input.new_empty((0, weight.size(1))) + ) + + + +class GroupedLinearBase(nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + num_experts: int = 1, + num_patterns: int = 1, + bias: bool = False, + device=None, + dtype: torch.dtype | None = None, + ) -> None: + """Allocate expert-stacked weights for grouped linear projections.""" + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.num_experts = num_experts + self.num_patterns = num_patterns + self.weight = nn.Parameter( + torch.empty( + num_experts * out_features, in_features, device=device, dtype=dtype + ) + ) + if bias: + self.bias = nn.Parameter( + torch.empty(num_experts * out_features, device=device, dtype=dtype) + ) + else: + self.register_parameter("bias", None) + + def forward( + self, + input: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + m_splits: Optional[list[int]] = None, + gather_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Project tokens with one expert or expert-specific grouped weights.""" + weight = self.weight.view(self.num_experts, self.out_features, self.in_features) + bias = ( + self.bias.view(self.num_experts, self.out_features) + if self.bias is not None + else None + ) + if self.num_experts == 1: + return F.linear(input, weight[0], None if bias is None else bias[0]) + m_splits = _m_splits_from(cu_seqlens, m_splits) + return _torch_grouped_linear(input, weight, bias, m_splits, gather_ids) + + +@triton.jit +def _multi_head_rms_norm_fwd_kernel( + x_ptr, + out_ptr, + weight_ptr, + weight_bias, + eps, + x_stride_t, + x_stride_h, + x_stride_d, + out_stride_t, + out_stride_h, + out_stride_d, + weight_stride_h, + weight_stride_d, + T, + D: tl.constexpr, + acc_dtype: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, +): + """Normalize each token head with root mean square statistics in Triton.""" + block_idx_t = tl.program_id(0) + h = tl.program_id(1) + t_start = block_idx_t * BLOCK_SIZE_T + d_block_offs = tl.arange(0, BLOCK_SIZE_D) + + for t_off in tl.static_range(0, BLOCK_SIZE_T): + t = t_start + t_off + + if t < T: + x_start_ptr = x_ptr + t * x_stride_t + h * x_stride_h + out_start_ptr = out_ptr + t * out_stride_t + h * out_stride_h + weight_start_ptr = weight_ptr + h * weight_stride_h + + rms_sum = tl.zeros([], dtype=acc_dtype) + for d_block_start in tl.static_range(0, D, BLOCK_SIZE_D): + d_ptrs = d_block_start + d_block_offs + d_mask = d_ptrs < D + x = tl.load( + x_start_ptr + d_ptrs * x_stride_d, mask=d_mask, other=0.0 + ).to(acc_dtype) + rms_sum += tl.sum(x * x, axis=0) + + rms_norm_inv = tl.rsqrt(rms_sum / D + eps) + + for d_block_start in tl.static_range(0, D, BLOCK_SIZE_D): + d_ptrs = d_block_start + d_block_offs + d_mask = d_ptrs < D + x = tl.load( + x_start_ptr + d_ptrs * x_stride_d, mask=d_mask, other=0.0 + ).to(acc_dtype) + weight = tl.load( + weight_start_ptr + d_ptrs * weight_stride_d, mask=d_mask, other=0.0 + ).to(acc_dtype) + out = x * rms_norm_inv * (weight + weight_bias) + tl.store(out_start_ptr + d_ptrs * out_stride_d, out, mask=d_mask) + + +@triton_op("magi2::multi_head_rms_norm_fwd", mutates_args={}) +def _multi_head_rms_norm_fwd_op( + t: torch.Tensor, + w: torch.Tensor, + weight_bias: float, + eps: float, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Launch multi-head root mean square normalization with Triton.""" + out = torch.empty(t.shape, dtype=out_dtype, device=t.device) + T, H, D = t.shape + block_size_d = min(8192, triton.next_power_of_2(D)) + block_size_t = 1 + grid = (triton.cdiv(T, block_size_t), H) + + wrap_triton(_multi_head_rms_norm_fwd_kernel)[grid]( + t, + out, + w, + weight_bias, + eps, + t.stride(0), + t.stride(1), + t.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + w.stride(0), + w.stride(1), + T, + D, + acc_dtype=tl.float32, + BLOCK_SIZE_T=block_size_t, + BLOCK_SIZE_D=block_size_d, + ) + return out + + +@triton.jit +def _grouped_multi_head_rms_norm_fwd_kernel( + x_ptr, + out_ptr, + weight_ptr, + cu_group_sizes_ptr, + weight_bias, + eps, + x_stride_t, + x_stride_b, + x_stride_h, + x_stride_d, + out_stride_t, + out_stride_b, + out_stride_h, + out_stride_d, + weight_stride_g, + weight_stride_h, + weight_stride_d, + T, + B: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + G: tl.constexpr, + acc_dtype: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, +): + """Normalize modality-grouped token heads with Triton.""" + block_idx_t = tl.program_id(0) + h = tl.program_id(1) + + g = -1 + t_start = -1 + + accum_blocks = 0 + for curr_g in range(G): + g_start = tl.load(cu_group_sizes_ptr + curr_g) + g_end = tl.load(cu_group_sizes_ptr + curr_g + 1) + g_size = g_end - g_start + g_blocks = tl.cdiv(g_size, BLOCK_SIZE_T) + + is_in_group = (block_idx_t >= accum_blocks) & ( + block_idx_t < accum_blocks + g_blocks + ) + if is_in_group: + g = curr_g + rel_block_idx = block_idx_t - accum_blocks + t_start = g_start + rel_block_idx * BLOCK_SIZE_T + + accum_blocks += g_blocks + + if g == -1: + return + + t_end = tl.load(cu_group_sizes_ptr + g + 1) + d_block_offs = tl.arange(0, BLOCK_SIZE_D) + + for t_off in tl.static_range(0, BLOCK_SIZE_T): + t = t_start + t_off + + if (t < t_end) & (t < T): + for b in tl.static_range(0, B): + x_start_ptr = x_ptr + t * x_stride_t + b * x_stride_b + h * x_stride_h + out_start_ptr = ( + out_ptr + t * out_stride_t + b * out_stride_b + h * out_stride_h + ) + weight_start_ptr = ( + weight_ptr + g * weight_stride_g + h * weight_stride_h + ) + + rms_sum = tl.zeros([], dtype=acc_dtype) + for d_block_start in tl.static_range(0, D, BLOCK_SIZE_D): + d_ptrs = d_block_start + d_block_offs + d_mask = d_ptrs < D + x = tl.load( + x_start_ptr + d_ptrs * x_stride_d, mask=d_mask, other=0.0 + ).to(acc_dtype) + rms_sum += tl.sum(x * x, axis=0) + + rms_norm_inv = tl.rsqrt(rms_sum / D + eps) + + for d_block_start in tl.static_range(0, D, BLOCK_SIZE_D): + d_ptrs = d_block_start + d_block_offs + d_mask = d_ptrs < D + x = tl.load( + x_start_ptr + d_ptrs * x_stride_d, mask=d_mask, other=0.0 + ).to(acc_dtype) + weight = tl.load( + weight_start_ptr + d_ptrs * weight_stride_d, + mask=d_mask, + other=0.0, + ).to(acc_dtype) + out = x * rms_norm_inv * (weight + weight_bias) + tl.store(out_start_ptr + d_ptrs * out_stride_d, out, mask=d_mask) + + +@triton_op("magi2::grouped_multi_head_rms_norm_fwd", mutates_args={}) +def _grouped_multi_head_rms_norm_fwd_op( + t: torch.Tensor, + w: torch.Tensor, + cu_group_sizes: torch.Tensor, + weight_bias: float, + eps: float, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Launch grouped multi-head root mean square normalization with Triton.""" + out = torch.empty(t.shape, dtype=out_dtype, device=t.device) + T, B, H, D = t.shape + G = w.shape[0] + block_size_d = min(8192, triton.next_power_of_2(D)) + block_size_t = 1 + grid = (triton.cdiv(T, block_size_t), H, 1) + + wrap_triton(_grouped_multi_head_rms_norm_fwd_kernel)[grid]( + t, + out, + w, + cu_group_sizes, + weight_bias, + eps, + t.stride(0), + t.stride(1), + t.stride(2), + t.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + w.stride(0), + w.stride(1), + w.stride(2), + T, + B, + H, + D, + G, + acc_dtype=tl.float32, + BLOCK_SIZE_T=block_size_t, + BLOCK_SIZE_D=block_size_d, + ) + return out + + +def _shape_as_grouped_heads( + x: torch.Tensor, num_heads: int, head_dim: int +) -> tuple[int, int]: + """Validate a head tensor and derive its sequence and batch dimensions.""" + seqlen = x.shape[0] + if x.shape[-1] == num_heads * head_dim: + hidden = num_heads * head_dim + return seqlen, x.numel() // (seqlen * hidden) + if x.dim() >= 3 and x.shape[-2] == num_heads and x.shape[-1] == head_dim: + hidden = num_heads * head_dim + return seqlen, x.numel() // (seqlen * hidden) + raise AssertionError( + "x must have shape [t, *batch_dims, num_heads * head_dim] " + "or [t, *batch_dims, num_heads, head_dim]." + ) + + +def _multi_head_rms_norm_triton( + x: torch.Tensor, + weight: torch.Tensor, + num_heads: int, + head_dim: int, + eps: float, + weight_bias: float, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Apply Triton root mean square normalization to multi-head tokens.""" + t = x.view(-1, num_heads, head_dim) + w = weight.view(num_heads, head_dim) + return _multi_head_rms_norm_fwd_op(t, w, weight_bias, eps, out_dtype).view_as(x) + + +def _grouped_multi_head_rms_norm_triton( + x: torch.Tensor, + weight: torch.Tensor, + num_groups: int, + num_heads: int, + head_dim: int, + cu_group_sizes: torch.Tensor, + eps: float, + weight_bias: float, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Apply Triton root mean square normalization to grouped token heads.""" + seqlen, batch = _shape_as_grouped_heads(x, num_heads, head_dim) + t = x.view(seqlen, batch, num_heads, head_dim) + w = weight.view(num_groups, num_heads, head_dim) + return _grouped_multi_head_rms_norm_fwd_op( + t, w, cu_group_sizes.contiguous(), weight_bias, eps, out_dtype + ).view_as(x) + + + + +class MultiModalityRMSNorm(nn.Module): + __constants__ = ["dim", "eps", "num_modality"] + dim: int + eps: float + num_modality: int + + def __init__( + self, + dim: int, + eps: float = 1e-6, + device: torch.device | None = None, + num_modality: int = 1, + num_patterns: int = 1, + out_dtype: torch.dtype | None = None, + backend: Literal["torch", "triton"] = "triton", + ): + """Configure modality-specific normalization weights and execution backend.""" + super().__init__() + if backend not in {"torch", "triton"}: + raise ValueError(f"Unsupported RMSNorm backend: {backend}") + self.dim = dim + self.eps = eps + self.num_modality = num_modality + self.num_patterns = num_patterns + self.out_dtype = out_dtype + self.backend = backend + + self.compute_dtype = torch.float32 + self.weight_bias = 1.0 + + self.weight = torch.nn.Parameter( + torch.zeros( + num_patterns * dim * num_modality, device=device, dtype=torch.float32 + ) + ) + + def forward( + self, x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None + ) -> torch.Tensor: + """Normalize tokens with single- or multi-modality weights.""" + target_dtype = self.out_dtype + + if self.num_modality > 1: + assert modality_dispatcher is not None, ( + "modality_dispatcher is required for multi-modality RMSNorm" + ) + return self.forward_multi_experts(x, modality_dispatcher, target_dtype) + else: + return self.forward_single_expert(x, target_dtype) + + def forward_multi_experts( + self, + x: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + target_dtype: torch.dtype | None = None, + ) -> torch.Tensor: + """Normalize modality-grouped tokens with Triton or PyTorch.""" + out_dtype = target_dtype if target_dtype is not None else x.dtype + if self.backend == "triton" and x.is_cuda and not torch.is_grad_enabled(): + return _grouped_multi_head_rms_norm_triton( + x=x, + weight=self.weight, + num_groups=self.num_modality, + num_heads=self.num_patterns, + head_dim=self.dim, + cu_group_sizes=modality_dispatcher.cu_group_sizes.to( + device=x.device, dtype=torch.int32 + ), + eps=self.eps, + weight_bias=self.weight_bias, + out_dtype=out_dtype, + ) + return self._forward_multi_experts_torch_impl( + x, modality_dispatcher, target_dtype + ) + + @torch.compile(dynamic=True, fullgraph=False) + def _forward_multi_experts_torch_impl( + self, + x: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + target_dtype: torch.dtype | None, + ) -> torch.Tensor: + """Apply modality-specific root mean square normalization in PyTorch.""" + t, original_dtype = x.float(), x.dtype + t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) + + weight_chunked = self.weight.chunk(self.num_modality, dim=-1) + t_list = modality_dispatcher.dispatch(t) + + for i in range(self.num_modality): + t_list[i] = t_list[i] * ( + weight_chunked[i].view(self.num_patterns, self.dim) + 1 + ) + + t = modality_dispatcher.undispatch(*t_list) + + out_dtype = target_dtype if target_dtype is not None else original_dtype + return t.to(out_dtype) + + def forward_single_expert( + self, x: torch.Tensor, target_dtype: torch.dtype | None = None + ) -> torch.Tensor: + """Normalize tokens with one root mean square normalization weight set.""" + out_dtype = target_dtype if target_dtype is not None else x.dtype + if ( + self.backend == "triton" + and self.num_patterns > 1 + and x.is_cuda + and not torch.is_grad_enabled() + ): + return _multi_head_rms_norm_triton( + x=x, + weight=self.weight, + num_heads=self.num_patterns, + head_dim=self.dim, + eps=self.eps, + weight_bias=self.weight_bias, + out_dtype=out_dtype, + ) + return self._forward_single_expert_torch_impl(x, target_dtype) + + @torch.compile(dynamic=True, fullgraph=False) + def _forward_single_expert_torch_impl( + self, x: torch.Tensor, target_dtype: torch.dtype | None = None + ) -> torch.Tensor: + t, original_dtype = x.float(), x.dtype + t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) + + out_dtype = target_dtype if target_dtype is not None else original_dtype + + return (t * (self.weight.view(self.num_patterns, self.dim) + 1)).to(out_dtype) + + +class CustomGroupedLinear(GroupedLinearBase): + def __init__( + self, + in_features: int, + out_features: int, + num_experts: int, + num_patterns: int, + num_layers_for_initialization: int, + bias: bool = False, + device=None, + dtype: torch.dtype | None = None, + ) -> None: + """Initialize grouped projection weights for expert token partitions.""" + super().__init__( + in_features=in_features, + out_features=out_features, + num_experts=num_experts, + num_patterns=num_patterns, + bias=bias, + device=device, + dtype=dtype, + ) + + def forward( + self, + input: torch.Tensor, + modality_dispatcher: Optional[Union[ModalityDispatcher, list[int]]] = None, + gather_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Derive token-group boundaries and apply the grouped projection.""" + if isinstance(modality_dispatcher, ModalityDispatcher): + m_splits = modality_dispatcher.group_size_cpu + cu_seqlens = modality_dispatcher.cu_group_sizes + elif isinstance(modality_dispatcher, torch.Tensor): + m_splits = modality_dispatcher.tolist() + cu_seqlens = None + elif isinstance(modality_dispatcher, list): + m_splits = modality_dispatcher + cu_seqlens = None + else: # fall back to base linear which does not need split infos + m_splits = None + cu_seqlens = None + + return super().forward( + input=input, cu_seqlens=cu_seqlens, m_splits=m_splits, gather_ids=gather_ids + ) + + +def create_linear( + in_features: int, + out_features: int, + num_layers: int = 1, + num_experts: int = 1, + num_patterns: int = 1, + bias: bool = True, + device=None, + dtype: torch.dtype | None = None, +) -> CustomGroupedLinear: + """Create the grouped linear layer used by MAGI-2 transformer blocks.""" + return CustomGroupedLinear( + in_features=in_features, + out_features=out_features, + num_experts=num_experts, + num_patterns=num_patterns, + num_layers_for_initialization=num_layers, + bias=bias, + device=device, + dtype=dtype, + ) + + +# ============================================================ +# element_wise_rope +# ============================================================ + +Coords: TypeAlias = tuple[int, int, int] + + +class ElementWiseFourierEmbed(nn.Module): + def __init__( + self, + dim: int, + max_res: int = 224, + temperature: float = 10000.0, + in_pixels: bool = True, + linear_bands: bool = False, + learnable: bool = False, + device: torch.device = torch.device("cpu"), + dtype: torch.dtype = torch.float32, + ): + """Configure frequency bands for coordinate-wise Fourier embeddings. + + Args: + dim: output feature dimension, total number of channels, must be divisible by 6 + max_res: maximum pixel-frequency resolution, used to build the pixel-domain frequency bands + temperature: temperature parameter for the inverse-frequency mode + in_pixels: True -> use pixel frequency bands, False -> use inverse frequency bands + linear_bands: whether the pixel frequency bands are linearly spaced + learnable: whether to make the frequency bands a trainable parameter + """ + super().__init__() + self.dim = dim + self.in_pixels = in_pixels + self.learnable = learnable + self.temperature = temperature + self.max_res = max_res + self.linear_bands = linear_bands + self.device = device + self.dtype = dtype + # make the frequency bands a trainable parameter or register them as a buffer + bands = self.get_default_bands() + if self.learnable: + self.bands = nn.Parameter(bands) + else: + self.register_buffer("bands", bands) + + def forward(self, coords: torch.Tensor) -> torch.Tensor: + """Encode token coordinates as element-wise Fourier features. + + Args: + coords: [L,9], column order is (time, row, col, T, H, W, ref_T, ref_H, ref_W) + Returns: + emb: [L, dim] element-wise Fourier embedding + """ + # use slicing instead of unbind + stack to reduce intermediate tensors + coords_xyz = coords[:, :3] # [L,3] -> (t, h, w) + sizes = coords[:, 3:6] # [L,3] -> (T, H, W) + refs = coords[:, 6:9] # [L,3] -> (ref_T, ref_H, ref_W) + + # compute the scaling factor + scales = (refs - 1) / (sizes - 1) # [L,3] + + # NOTE: for points where both ref and size are 1 the scale is fixed to 1; anything else errors out + scales[(refs == 1) & (sizes == 1)] = 1 + assert not scales.isnan().any(), "scales has nan" + assert not scales.isinf().any(), "scales has inf" + + # center alignment, only for the h, w dims, not for the t dim + centers = (sizes - 1) / 2 # [L,3] + centers[:, 0] = 0 # do not center the time dimension + coords_xyz = coords_xyz - centers # [L,3] + + # project onto the frequency bands in one shot: [L,3,B] + proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands + + # compute sin & cos and concatenate + sin_proj = proj.sin() # [L,3,B] + cos_proj = proj.cos() + + return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) + + def reset_parameters(self): + bands = self.get_default_bands() + self.bands.copy_(bands) + + def get_default_bands(self): + if self.in_pixels: + raise NotImplementedError("in_pixels are not implemented yet") + else: + bands = freq_bands( + self.dim // 8, temperature=self.temperature, step=1, device=self.device + ).to(self.dtype) + return bands + + +def get_coords( + shape: Coords, + ref_feat_shape: Coords, + offset_thw: Coords = (0, 0, 0), + device: torch.device = torch.device("cpu"), + dtype: torch.dtype = torch.float32, + time_positions: Optional[torch.Tensor] = None, +): + """ + Build the feature-grid coordinates together with their original and reference size metadata. + Args: + feat_shape: [T, H, W] original feature-map size + ref_feat_shape: [T_ref, H_ref, W_ref] reference feature-map size + device: device the coordinate tensors live on + time_positions: optional 1D tensor (length == ori_t) used in place of the default arange(ori_t) + as the time-dimension coordinates. Together with ref_t == ori_t this bypasses the linear rescale. + Returns: + coords: tensor of shape (T*H*W, 9) containing (t, h, w, T, H, W, ref_T, ref_H, ref_W) + """ + ori_t, ori_h, ori_w = shape + ref_t, ref_h, ref_w = ref_feat_shape + + offset_t, offset_h, offset_w = offset_thw + # build the index ranges + if time_positions is None: + time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t + else: + time_rng = time_positions.to(device=device, dtype=dtype) + offset_t + height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h + width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w + + # meshgrid to build the 3D grid (T, H, W) + time_grid, height_grid, width_grid = torch.meshgrid( + time_rng, height_rng, width_rng, indexing="ij" + ) + + # stack and flatten + coords_grid = torch.stack([time_grid, height_grid, width_grid], dim=-1) + coords_flat = coords_grid.reshape(-1, 3) + + # build and broadcast the size metadata + meta = torch.tensor( + [ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], device=device, dtype=dtype + ) + meta_expanded = meta.expand(coords_flat.size(0), -1) + + # concatenate and return + return torch.cat([coords_flat, meta_expanded], dim=-1) + + +# ============================================================ +# nd_rotary_pos_embedding +# ============================================================ + + +def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]: + """generate N-D grid in dimension order. + + The ndgrid function is like meshgrid except that the order of the first two input arguments are switched. + + That is, the statement + [X1,X2,X3] = ndgrid(x1,x2,x3) + + produces the same result as + + [X2,X1,X3] = meshgrid(x2,x1,x3) + + This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make + torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy'). + + """ + try: + return torch.meshgrid(*tensors, indexing="ij") + except TypeError: + # old PyTorch < 1.10 will follow this path as it does not have indexing arg, + # the old behaviour of meshgrid was 'ij' + return torch.meshgrid(*tensors) + + +def meshgrid(*tensors) -> Tuple[torch.Tensor, ...]: + """generate N-D grid in spatial dim order. + + The meshgrid function is similar to ndgrid except that the order of the + first two input and output arguments is switched. + + That is, the statement + + [X,Y,Z] = meshgrid(x,y,z) + produces the same result as + + [Y,X,Z] = ndgrid(y,x,z) + Because of this, meshgrid is better suited to problems in two- or three-dimensional Cartesian space, + while ndgrid is better suited to multidimensional problems that aren't spatially based. + """ + + # NOTE: this will throw in PyTorch < 1.10 as meshgrid did not support indexing arg or have + # capability of generating grid in xy order before then. + return torch.meshgrid(*tensors, indexing="xy") + + +def pixel_freq_bands( + num_bands: int, + max_freq: float = 224.0, + linear_bands: bool = True, + device: Optional[torch.device] = None, +): + """Create pixel-domain frequency bands scaled by pi.""" + if linear_bands: + bands = torch.linspace( + 1.0, max_freq / 2, num_bands, dtype=torch.float32, device=device + ) + else: + bands = 2 ** torch.linspace( + 0, math.log(max_freq, 2) - 1, num_bands, dtype=torch.float32, device=device + ) + return bands * torch.pi + + +def freq_bands( + num_bands: int, + temperature: float = 10000.0, + step: int = 2, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Create inverse-frequency bands for sinusoidal embeddings.""" + exp = ( + torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to( + torch.float32 + ) + / num_bands + ) + bands = 1.0 / (temperature**exp) + return bands + + +def build_sincos2d_pos_embed( + feat_shape: List[int], + dim: int = 64, + temperature: float = 10000.0, + reverse_coord: bool = False, + interleave_sin_cos: bool = False, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Build a two-dimensional sinusoidal embedding for a feature grid. + + Args: + feat_shape: + dim: + temperature: + reverse_coord: stack grid order W, H instead of H, W + interleave_sin_cos: sin, cos, sin, cos stack instead of sin, sin, cos, cos + dtype: + device: + + Returns: + + """ + assert dim % 4 == 0, ( + "Embed dimension must be divisible by 4 for sin-cos 2D position embedding" + ) + pos_dim = dim // 4 + bands = freq_bands(pos_dim, temperature=temperature, step=1, device=device) + + if reverse_coord: + feat_shape = feat_shape[::-1] # stack W, H instead of H, W + grid = ( + torch.stack( + ndgrid( + [ + torch.arange(s, device=device, dtype=torch.int64).to(torch.float32) + for s in feat_shape + ] + ) + ) + .flatten(1) + .transpose(0, 1) + ) + pos2 = grid.unsqueeze(-1) * bands.unsqueeze(0) + # FIXME add support for unflattened spatial dim? + + stack_dim = ( + 2 if interleave_sin_cos else 1 + ) # stack sin, cos, sin, cos instead of sin sin cos cos + pos_emb = torch.stack([torch.sin(pos2), torch.cos(pos2)], dim=stack_dim).flatten(1) + return pos_emb.to(dtype=dtype) + + +def build_fourier_pos_embed( + feat_shape: List[int], + bands: Optional[torch.Tensor] = None, + num_bands: int = 64, + max_res: int = 224, + temperature: float = 10000.0, + linear_bands: bool = False, + include_grid: bool = False, + in_pixels: bool = True, + ref_feat_shape: Optional[List[float]] = None, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, +) -> List[torch.Tensor]: + """Build sine and cosine Fourier features for a feature grid. + + Args: + feat_shape: Feature shape for embedding. + bands: Pre-calculated frequency bands. + num_bands: Number of frequency bands (determines output dim). + max_res: Maximum resolution for pixel based freq. + temperature: Temperature for non-pixel freq. + linear_bands: Linear band spacing for pixel based freq. + include_grid: Include the spatial grid in output. + in_pixels: Output in pixel freq. + ref_feat_shape: Reference feature shape for resize / fine-tune. + dtype: Output dtype. + device: Output device. + + Returns: + + """ + if bands is None: + if in_pixels: + bands = pixel_freq_bands( + num_bands, float(max_res), linear_bands=linear_bands, device=device + ) + else: + bands = freq_bands( + num_bands, temperature=temperature, step=1, device=device + ) + else: + if device is None: + device = bands.device + if dtype is None: + dtype = bands.dtype + + if in_pixels: + t = [ + torch.linspace(-1.0, 1.0, steps=s, device=device, dtype=torch.float32) + for s in feat_shape + ] + else: + t = [ + torch.arange(s, device=device, dtype=torch.int64).to(torch.float32) + for s in feat_shape + ] + # align spatial center (H/2,W/2) to (0,0) + t[1] = t[1] - (feat_shape[1] - 1) / 2 + t[2] = t[2] - (feat_shape[2] - 1) / 2 + if ref_feat_shape is not None: + # eva's scheme for resizing rope embeddings (ref shape = pretrain) + # aligning to the endpoint e.g [0,1,2] -> [0, 0.4, 0.8, 1.2, 1.6, 2] + t_rescaled = [] + for x, f, r in zip(t, feat_shape, ref_feat_shape): + if f == 1: + assert r == 1, "ref_feat_shape must be 1 when feat_shape is 1" + t_rescaled.append(x) + else: + t_rescaled.append(x / (f - 1) * (r - 1)) + else: + t_rescaled = t + + grid = torch.stack(ndgrid(t_rescaled), dim=-1) + grid = grid.unsqueeze(-1) + pos = grid * bands + + pos_sin, pos_cos = pos.sin().to(dtype=dtype), pos.cos().to(dtype) + out = [grid, pos_sin, pos_cos] if include_grid else [pos_sin, pos_cos] + return out + + +class FourierEmbed(nn.Module): + def __init__( + self, + max_res: int = 224, + num_bands: int = 64, + concat_grid=True, + keep_spatial=False, + ): + """Configure Fourier bands and the positional-feature output layout.""" + super().__init__() + self.max_res = max_res + self.num_bands = num_bands + self.concat_grid = concat_grid + self.keep_spatial = keep_spatial + self.register_buffer( + "bands", pixel_freq_bands(max_res, num_bands), persistent=False + ) + + def forward(self, x): + """Append Fourier positional features to an image feature tensor.""" + B, C = x.shape[:2] + feat_shape = x.shape[2:] + emb = build_fourier_pos_embed( + feat_shape, + self.bands, + include_grid=self.concat_grid, + dtype=x.dtype, + device=x.device, + ) + emb = torch.cat(emb, dim=-1) + emb = emb.transpose(-1, -2).flatten(len(feat_shape)) + batch_expand = (B,) + (-1,) * (x.ndim - 1) + + # FIXME support nD + if self.keep_spatial: + x = torch.cat( + [x, emb.unsqueeze(0).expand(batch_expand).permute(0, 3, 1, 2)], dim=1 + ) + else: + x = torch.cat( + [x.permute(0, 2, 3, 1), emb.unsqueeze(0).expand(batch_expand)], dim=-1 + ) + x = x.reshape(B, feat_shape.numel(), -1) + + return x + + +def rot(x): + return torch.stack([-x[..., 1::2], x[..., ::2]], -1).reshape(x.shape) + + +def apply_rot_embed(x: torch.Tensor, sin_emb, cos_emb): + if sin_emb.ndim == 3: + return x * cos_emb.unsqueeze(1).expand_as(x) + rot(x) * sin_emb.unsqueeze( + 1 + ).expand_as(x) + return x * cos_emb + rot(x) * sin_emb + + +def apply_rot_embed_list(x: List[torch.Tensor], sin_emb, cos_emb): + if isinstance(x, torch.Tensor): + x = [x] + return [t * cos_emb + rot(t) * sin_emb for t in x] + + +def apply_rot_embed_cat(x: torch.Tensor, emb): + sin_emb, cos_emb = emb.tensor_split(2, -1) + if sin_emb.ndim == 3: + return x * cos_emb.unsqueeze(1).expand_as(x) + rot(x) * sin_emb.unsqueeze( + 1 + ).expand_as(x) + return x * cos_emb.unsqueeze(1).unsqueeze(2) + rot(x) * sin_emb.unsqueeze( + 1 + ).unsqueeze(2) + + +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange( + torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2 + ) + + +def apply_rotary_emb_torch(x, cos, sin, interleaved=False): + """Apply rotary embeddings to the rotary channels of attention heads. + + x: (batch_size, seqlen, nheads, headdim) + cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) + """ + ro_dim = cos.shape[-1] * 2 + assert ro_dim <= x.shape[-1] + cos = repeat( + cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)" + ) + sin = repeat( + sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)" + ) + return torch.cat( + [ + x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, + x[..., ro_dim:], + ], + dim=-1, + ) + + +def apply_keep_indices_nlc(x, pos_embed, keep_indices): + pos_embed = pos_embed.unsqueeze(0).expand(x.shape[0], -1, -1) + pos_embed = pos_embed.gather( + 1, keep_indices.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1]) + ) + return pos_embed + + +def build_rotary_pos_embed( + feat_shape: List[int], + bands: Optional[torch.Tensor] = None, + dim: int = 64, + max_res: int = 224, + temperature: float = 10000.0, + linear_bands: bool = False, + in_pixels: bool = True, + ref_feat_shape: Optional[List[float]] = None, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, +): + """Build split sine and cosine rotary embeddings for a feature grid. + + Args: + feat_shape: Spatial shape of the target tensor for embedding. + bands: Optional pre-generated frequency bands + dim: Output dimension of embedding tensor. + max_res: Maximum resolution for pixel mode. + temperature: Temperature (inv freq) for non-pixel mode + linear_bands: Linearly (instead of log) spaced bands for pixel mode + in_pixels: Pixel vs language (inv freq) mode. + dtype: Output dtype. + device: Output device. + + Returns: + + """ + sin_emb, cos_emb = build_fourier_pos_embed( + feat_shape, + bands=bands, + num_bands=dim // 8, + max_res=max_res, + temperature=temperature, + linear_bands=linear_bands, + in_pixels=in_pixels, + ref_feat_shape=ref_feat_shape, + device=device, + dtype=dtype, + ) + num_spatial_dim = 1 + # this would be much nicer as a .numel() call to torch.Size(), but torchscript sucks + for x in feat_shape: + num_spatial_dim *= x + + # the original timm implementation performs a repeat_interleave here + # but in flash attn the repeat_interleave is done on the fly + # this aligns with the flash attn implementation and does not run repeat_interleave here + # sin_emb = sin_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1) + # cos_emb = cos_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1) + sin_emb = sin_emb.reshape(num_spatial_dim, -1) + cos_emb = cos_emb.reshape(num_spatial_dim, -1) + return sin_emb, cos_emb + + +class RotaryEmbedding(nn.Module): + """Rotary position embedding + + NOTE: This is my initial attempt at impl rotary embedding for spatial use, it has not + been well tested, and will likely change. It will be moved to its own file. + + The following impl/resources were referenced for this impl: + * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py + * https://blog.eleuther.ai/rotary-embeddings/ + """ + + def __init__( + self, + dim, + max_res=224, + temperature=10000, + in_pixels=True, + linear_bands: bool = False, + feat_shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[float]] = None, + ): + """Configure cached or shape-dependent sine and cosine rotary embeddings.""" + super().__init__() + self.dim = dim + self.max_res = max_res + self.temperature = temperature + self.in_pixels = in_pixels + self.feat_shape = feat_shape + self.ref_feat_shape = ref_feat_shape + + if feat_shape is None: + # only cache bands + if in_pixels: + bands = pixel_freq_bands( + dim // 8, float(max_res), linear_bands=linear_bands + ) + else: + bands = freq_bands(dim // 8, temperature=temperature, step=1) + logger.debug("bands: %s", bands) + self.register_buffer("bands", bands, persistent=False) + self.pos_embed_sin = None + self.pos_embed_cos = None + else: + # cache full sin/cos embeddings if shape provided up front + emb_sin, emb_cos = build_rotary_pos_embed( + feat_shape=feat_shape, + dim=dim, + max_res=max_res, + linear_bands=linear_bands, + in_pixels=in_pixels, + ref_feat_shape=self.ref_feat_shape, + temperature=self.temperature, + ) + self.bands = None + self.register_buffer("pos_embed_sin", emb_sin, persistent=False) + self.register_buffer("pos_embed_cos", emb_cos, persistent=False) + + def get_embed(self, shape: Optional[List[int]] = None): + """Return cached rotary embeddings or build them for a target shape.""" + if self.bands is not None: + # rebuild embeddings every call, use if target shape changes + assert shape is not None # type: ignore[unreachable] + return build_rotary_pos_embed( + shape, + self.bands, + in_pixels=self.in_pixels, + temperature=self.temperature, + ) + else: + return self.pos_embed_sin, self.pos_embed_cos + + def forward(self, x): + # assuming channel-first tensor where spatial dim are >= 2 + sin_emb, cos_emb = self.get_embed(x.shape[2:]) + return apply_rot_embed(x, sin_emb, cos_emb) + + +class RotaryEmbeddingCat(nn.Module): + """Rotary position embedding w/ concatenatd sin & cos + + The following impl/resources were referenced for this impl: + * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py + * https://blog.eleuther.ai/rotary-embeddings/ + """ + + def __init__( + self, + dim, + max_res=224, + temperature=10000, + in_pixels=True, + linear_bands: bool = False, + feat_shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[float]] = None, + ): + """Configure concatenated sine and cosine rotary embedding state.""" + super().__init__() + self.dim = dim + self.max_res = max_res + self.temperature = temperature + self.in_pixels = in_pixels + self.linear_bands = linear_bands + self.feat_shape = feat_shape + self.ref_feat_shape = ref_feat_shape + + if feat_shape is None: + self.bands = None + self.pos_embed = None + else: + # cache full sin/cos embeddings if shape provided up front + embeds = build_rotary_pos_embed( + feat_shape=feat_shape, + dim=dim, + max_res=max_res, + linear_bands=linear_bands, + in_pixels=in_pixels, + ref_feat_shape=self.ref_feat_shape, + temperature=self.temperature, + ) + self.bands = None + self.register_buffer("pos_embed", torch.cat(embeds, -1), persistent=False) + + def get_embed( + self, + shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[float]] = None, + ): + """Return concatenated rotary embeddings for a target feature shape.""" + if shape is not None: + # rebuild bands and embeddings every call, use if target shape changes + embeds = build_rotary_pos_embed( + feat_shape=shape, + dim=self.dim, + max_res=self.max_res, + linear_bands=self.linear_bands, + in_pixels=self.in_pixels, + ref_feat_shape=ref_feat_shape + if ref_feat_shape + else self.ref_feat_shape, + temperature=self.temperature, + device=torch.cuda.current_device(), + ) + return torch.cat(embeds, -1) + elif self.pos_embed is not None: + return self.pos_embed # type: ignore[unreachable] + else: + assert False, ( + "get_embed() requires pre-computed pos_embed or valid shape w/ pre-computed bands" + ) + + def forward(self, x): + # assuming channel-first tensor where spatial dim are >= 2 + pos_embed = self.get_embed(x.shape[2:]) + return apply_rot_embed_cat(x, pos_embed) + + +class LearnableRotaryEmbeddingCat(nn.Module): + """Rotary position embedding w/ concatenatd sin & cos + + The following impl/resources were referenced for this impl: + * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py + * https://blog.eleuther.ai/rotary-embeddings/ + """ + + def __init__( + self, + dim, + max_res=224, + temperature=10000, + in_pixels=True, + linear_bands: bool = False, + feat_shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[int]] = None, + ): + """Initialize trainable frequency bands for rotary embeddings.""" + super().__init__() + self.dim = dim + self.max_res = max_res + self.temperature = temperature + self.in_pixels = in_pixels + self.linear_bands = linear_bands + self.feat_shape = feat_shape + self.ref_feat_shape = ref_feat_shape + self.bands = nn.Parameter(self.get_default_bands()) + setattr(self.bands, "sequence_parallel", True) + + def get_default_bands(self): + """Create the initial trainable rotary frequency bands.""" + if self.in_pixels: + bands = pixel_freq_bands( + self.dim // 8, + float(self.max_res), + linear_bands=self.linear_bands, + device=torch.cuda.current_device(), + ) + else: + bands = freq_bands( + self.dim // 8, + temperature=self.temperature, + step=1, + device=torch.cuda.current_device(), + ) + return bands + + def get_embed( + self, shape: Optional[List[int]], ref_feat_shape: Optional[List[float]] = None + ): + """Build concatenated rotary embeddings from learned frequency bands.""" + # rebuild bands and embeddings every call, use if target shape changes + embeds = build_rotary_pos_embed( + feat_shape=shape, # type: ignore[arg-type] + bands=self.bands, # use learned bands + dim=self.dim, + max_res=self.max_res, + linear_bands=self.linear_bands, + in_pixels=self.in_pixels, + ref_feat_shape=ref_feat_shape if ref_feat_shape else self.ref_feat_shape, # type: ignore[arg-type] + temperature=self.temperature, + device=torch.cuda.current_device(), + ) + return torch.cat(embeds, -1) + + +# ============================================================ +# modality_dispatcher +# ============================================================ + + +def seqlens2cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.pad(torch.cumsum(seqlens, dim=0), (1, 0)) + + +class ModalityDispatcher: + """A torch.compile-friendly ModalityDispatcher. + + Design choices: + - dispatch() has no @torch._dynamo.disable(), so it can be traced by Dynamo + - _permute()/_inv_permute() have no side effects + """ + + num_modalities: int + + permuted_modality_mapping: torch.Tensor + modality_mapping: torch.Tensor + permute_mapping: torch.Tensor + inv_permute_mapping: torch.Tensor + + group_size: torch.Tensor + group_size_cpu: list[int] + group_size_cpu_tensor: torch.Tensor + cu_group_sizes: torch.Tensor + expert_ranges: torch.Tensor + + @torch._dynamo.disable + def __init__(self, modality_mapping: torch.Tensor, num_modalities: int): + """Precompute token permutations and group boundaries by modality.""" + self.modality_mapping = modality_mapping + self.num_modalities = num_modalities + + self.permute_mapping = torch.argsort(modality_mapping) + self.inv_permute_mapping = torch.argsort(self.permute_mapping) + self.permuted_modality_mapping = modality_mapping.index_select( + 0, self.permute_mapping + ) + + self.group_size = torch.bincount( + self.permuted_modality_mapping, minlength=num_modalities + ).to(torch.int32) + + # FIXME: these two causes a lot of trouble with torch compile + # need to avoid using them as soon as possible, and eventually remove them + self.group_size_cpu_tensor = self.group_size.to("cpu") + self.group_size_cpu = self.group_size_cpu_tensor.tolist() + + self.cu_group_sizes = seqlens2cu_seqlens(self.group_size) + + end_indices = torch.cumsum(self.group_size, dim=0) + start_indices = torch.cat( + [ + torch.zeros(1, dtype=end_indices.dtype, device=end_indices.device), + end_indices[:-1], + ], + dim=0, + ) + self.expert_ranges = ( + torch.stack([start_indices, end_indices], dim=1).to(torch.int32).cpu() + ) + + def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]: + return list(torch.split(x, self.group_size_cpu, dim=0)) + + def undispatch(self, *processed_groups: torch.Tensor) -> torch.Tensor: + return torch.cat(processed_groups, dim=0) + + def _permute(self, x: torch.Tensor) -> torch.Tensor: + return x.index_select(0, self.permute_mapping) + + def _inv_permute(self, x: torch.Tensor) -> torch.Tensor: + return x.index_select(0, self.inv_permute_mapping) + + +# ============================================================ +# mhc +# ============================================================ + +MHCTensorTuple = tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + +def _use_triton_backend(x: torch.Tensor, backend: str) -> bool: + return backend == "triton" and x.is_cuda and not torch.is_grad_enabled() + + +@triton.jit +def _sigmoid_affine_fwd_kernel( + x_ptr, + alpha_ptr, + bias_ptr, + out_ptr, + n_elements, + matmul_scale, + sigmoid_scale, + n_cols: tl.constexpr, + block_size: tl.constexpr, +): + """Apply an elementwise sigmoid affine transform with Triton.""" + block_start = tl.program_id(0) * block_size + offsets = block_start + tl.arange(0, block_size) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + alpha = tl.load(alpha_ptr) + bias = tl.load(bias_ptr + offsets % n_cols, mask=mask) + out = sigmoid_scale * tl.sigmoid(alpha * matmul_scale * x + bias) + tl.store(out_ptr + offsets, out.to(out_ptr.dtype.element_ty), mask=mask) + + +@triton_op("magi2::mhc_sigmoid_affine_fwd", mutates_args={}) +def _triton_sigmoid_affine_forward( + x: torch.Tensor, + alpha: torch.Tensor, + bias: torch.Tensor, + matmul_scale: float, + sigmoid_scale: float, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Launch the Triton sigmoid affine transform for connection gates.""" + x = x.contiguous() + alpha = alpha.contiguous() + bias = bias.contiguous() + out = torch.empty_like(x, dtype=out_dtype) + block_size = 256 + grid = (triton.cdiv(x.numel(), block_size),) + wrap_triton(_sigmoid_affine_fwd_kernel)[grid]( + x, + alpha, + bias, + out, + x.numel(), + matmul_scale, + sigmoid_scale, + x.shape[-1], + block_size, + ) + return out + + +def _sigmoid_affine( + x: torch.Tensor, + alpha: torch.Tensor, + bias: torch.Tensor, + matmul_scale: float, + sigmoid_scale: float, + out_dtype: torch.dtype, + backend: str, +) -> torch.Tensor: + """Apply the sigmoid affine gate with the selected execution backend.""" + if _use_triton_backend(x, backend): + return _triton_sigmoid_affine_forward( + x, alpha, bias, matmul_scale, sigmoid_scale, out_dtype + ) + return ( + sigmoid_scale * torch.sigmoid(alpha * matmul_scale * x + bias.unsqueeze(0)) + ).to(out_dtype) + + +@triton.jit +def _sinkhorn_knopp_affine_fwd_kernel( + x_ptr, + alpha_ptr, + bias_ptr, + out_ptr, + scale, + sk_eps, + stride_xb, + stride_xr, + stride_xc, + stride_ob, + stride_or, + stride_oc, + n_batch, + n_rows: tl.constexpr, + n_cols: tl.constexpr, + num_iters: tl.constexpr, + block_batch: tl.constexpr, + block_r: tl.constexpr, + block_c: tl.constexpr, +): + """Apply affine logits and Sinkhorn normalization with Triton.""" + batch_pid = tl.program_id(0) + batch_offsets = batch_pid * block_batch + tl.arange(0, block_batch) + batch_mask = batch_offsets < n_batch + row_offsets = tl.arange(0, block_r) + col_offsets = tl.arange(0, block_c) + mask = ( + batch_mask[:, None, None] + & (row_offsets[None, :, None] < n_rows) + & (col_offsets[None, None, :] < n_cols) + ) + + x = tl.load( + x_ptr + + batch_offsets[:, None, None] * stride_xb + + row_offsets[None, :, None] * stride_xr + + col_offsets[None, None, :] * stride_xc, + mask=mask, + other=-float("inf"), + ).to(tl.float32) + alpha = tl.load(alpha_ptr).to(tl.float32) + bias = tl.load( + bias_ptr + row_offsets[:, None] * stride_xr + col_offsets[None, :] * stride_xc, + mask=(row_offsets[:, None] < n_rows) & (col_offsets[None, :] < n_cols), + other=0.0, + ).to(tl.float32) + + h = alpha * scale * x + bias + h_max = tl.max(tl.max(tl.where(mask, h, -float("inf")), axis=2), axis=1) + h_max = tl.where(h_max == -float("inf"), 0.0, h_max) + m = tl.exp(h - h_max[:, None, None]) + for _ in range(num_iters): + col_sum = tl.sum(m, axis=1) + m = m / (col_sum[:, None, :] + sk_eps) + row_sum = tl.sum(m, axis=2) + m = m / (row_sum[:, :, None] + sk_eps) + + tl.store( + out_ptr + + batch_offsets[:, None, None] * stride_ob + + row_offsets[None, :, None] * stride_or + + col_offsets[None, None, :] * stride_oc, + m.to(out_ptr.dtype.element_ty), + mask=mask, + ) + + +@triton_op("magi2::mhc_sinkhorn_knopp_affine_fwd", mutates_args={}) +def _triton_sinkhorn_knopp_affine_forward( + x: torch.Tensor, + alpha: torch.Tensor, + bias: torch.Tensor, + matmul_scale: float, + num_sk_iters: int, + sk_eps: float, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Launch Triton Sinkhorn normalization for residual mixing matrices.""" + x = x.contiguous() + alpha = alpha.contiguous() + bias = bias.contiguous() + t, n_rows, n_cols = x.shape + out = torch.empty_like(x, dtype=out_dtype) + block_r = triton.next_power_of_2(n_rows) + block_c = triton.next_power_of_2(n_cols) + block_batch = max(1, 256 // (block_r * block_c)) + grid = (triton.cdiv(t, block_batch),) + wrap_triton(_sinkhorn_knopp_affine_fwd_kernel)[grid]( + x, + alpha, + bias, + out, + matmul_scale, + sk_eps, + x.stride(0), + x.stride(1), + x.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + t, + n_rows, + n_cols, + num_sk_iters, + block_batch, + block_r, + block_c, + ) + return out + + +def _sinkhorn_knopp(h: torch.Tensor, num_iters: int, eps: float) -> torch.Tensor: + m = torch.exp(h - h.amax(dim=(-2, -1), keepdim=True)) + for _ in range(num_iters): + m = m / (m.sum(dim=-2, keepdim=True) + eps) + m = m / (m.sum(dim=-1, keepdim=True) + eps) + return m + + +def _sinkhorn_knopp_affine( + x: torch.Tensor, + alpha: torch.Tensor, + bias: torch.Tensor, + matmul_scale: float, + num_sk_iters: int, + sk_eps: float, + out_dtype: torch.dtype, + backend: str, +) -> torch.Tensor: + """Build residual mixing matrices with the selected execution backend.""" + if _use_triton_backend(x, backend): + return _triton_sinkhorn_knopp_affine_forward( + x, alpha, bias, matmul_scale, num_sk_iters, sk_eps, out_dtype + ) + h = alpha * matmul_scale * x.to(torch.float32) + bias.unsqueeze(0).to(torch.float32) + return _sinkhorn_knopp(h, num_sk_iters, sk_eps).to(out_dtype) + + +@triton.jit +def _apply_hpre_fwd_kernel( + h_pre_ptr, + x_ptr, + out_ptr, + hpre_stride_t, + hpre_stride_n, + x_stride_t, + x_stride_n, + x_stride_c, + out_stride_t, + out_stride_c, + n_stream: tl.constexpr, + hidden_size: tl.constexpr, + block_size_c: tl.constexpr, +): + """Reduce multi-stream features with pre-connection weights in Triton.""" + t = tl.program_id(0) + n_offsets = tl.arange(0, n_stream) + c_offsets_base = tl.arange(0, block_size_c) + hpre = tl.load(h_pre_ptr + t * hpre_stride_t + n_offsets * hpre_stride_n).to( + tl.float32 + )[:, None] + + for c_start in tl.range(0, hidden_size, block_size_c): + c_offsets = c_start + c_offsets_base + c_mask = c_offsets < hidden_size + x = tl.load( + x_ptr + + t * x_stride_t + + n_offsets[:, None] * x_stride_n + + c_offsets[None, :] * x_stride_c, + mask=c_mask[None, :], + other=0.0, + ).to(tl.float32) + out = tl.sum(hpre * x, axis=0) + tl.store( + out_ptr + t * out_stride_t + c_offsets * out_stride_c, out, mask=c_mask + ) + + +@triton_op("magi2::mhc_apply_hpre_fwd", mutates_args={}) +def _triton_apply_hpre_forward(h_pre: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """Launch Triton reduction of multi-stream features into one stream.""" + _, n_stream, hidden_size = x.shape + out = torch.empty((x.size(0), hidden_size), device=x.device, dtype=x.dtype) + block_size_c = min(1024, triton.next_power_of_2(max(1, hidden_size // 2))) + grid = (x.size(0),) + wrap_triton(_apply_hpre_fwd_kernel)[grid]( + h_pre, + x, + out, + h_pre.stride(0), + h_pre.stride(1), + x.stride(0), + x.stride(1), + x.stride(2), + out.stride(0), + out.stride(1), + n_stream, + hidden_size, + block_size_c, + num_warps=4, + ) + return out + + +def _apply_hpre(h_pre: torch.Tensor, x: torch.Tensor, backend: str) -> torch.Tensor: + if _use_triton_backend(x, backend): + return _triton_apply_hpre_forward(h_pre, x) + return torch.einsum("tn,tnc->tc", h_pre, x) + + +@triton.jit +def _hyper_connect_fwd_kernel( + res_ptr, + out_ptr, + h_post_ptr, + h_res_ptr, + output_ptr, + stride_res_t, + stride_res_n, + stride_res_c, + stride_out_t, + stride_out_c, + stride_h_post_t, + stride_h_post_n, + stride_h_res_t, + stride_h_res_n, + stride_h_res_n2, + stride_output_t, + stride_output_n, + stride_output_c, + n_stream: tl.constexpr, + hidden_size: tl.constexpr, + block_size_c: tl.constexpr, +): + """Combine residual streams with post-connection weights in Triton.""" + t = tl.program_id(0) + i_offsets = tl.arange(0, n_stream) + j_offsets = tl.arange(0, n_stream) + c_offsets_base = tl.arange(0, block_size_c) + + h_post = tl.load(h_post_ptr + t * stride_h_post_t + i_offsets * stride_h_post_n).to( + tl.float32 + ) + h_res = tl.load( + h_res_ptr + + t * stride_h_res_t + + i_offsets[:, None] * stride_h_res_n + + j_offsets[None, :] * stride_h_res_n2 + ).to(tl.float32) + + for c_start in tl.range(0, hidden_size, block_size_c): + c_offsets = c_start + c_offsets_base + c_mask = c_offsets < hidden_size + out_row = tl.load( + out_ptr + t * stride_out_t + c_offsets * stride_out_c, + mask=c_mask, + other=0.0, + ).to(tl.float32) + + for i in tl.static_range(0, n_stream): + acc = tl.zeros([block_size_c], dtype=tl.float32) + for j in tl.static_range(0, n_stream): + h_res_ij = tl.sum( + tl.where( + (tl.arange(0, n_stream)[:, None] == i) + & (tl.arange(0, n_stream)[None, :] == j), + h_res, + 0.0, + ) + ) + res_row = tl.load( + res_ptr + + t * stride_res_t + + j * stride_res_n + + c_offsets * stride_res_c, + mask=c_mask, + other=0.0, + ).to(tl.float32) + acc += h_res_ij * res_row + + h_post_i = tl.sum(tl.where(tl.arange(0, n_stream) == i, h_post, 0.0)) + acc += h_post_i * out_row + tl.store( + output_ptr + + t * stride_output_t + + i * stride_output_n + + c_offsets * stride_output_c, + acc.to(output_ptr.dtype.element_ty), + mask=c_mask, + ) + + +@triton_op("magi2::mhc_hyper_connect_fwd", mutates_args={}) +def _triton_hyper_connect_forward( + res: torch.Tensor, out: torch.Tensor, h_post: torch.Tensor, h_res: torch.Tensor +) -> torch.Tensor: + """Launch Triton residual-stream mixing and output injection.""" + _, n_stream, hidden_size = res.shape + output = torch.empty_like(res) + block_size_c = 256 + grid = (res.size(0),) + wrap_triton(_hyper_connect_fwd_kernel)[grid]( + res, + out, + h_post, + h_res, + output, + res.stride(0), + res.stride(1), + res.stride(2), + out.stride(0), + out.stride(1), + h_post.stride(0), + h_post.stride(1), + h_res.stride(0), + h_res.stride(1), + h_res.stride(2), + output.stride(0), + output.stride(1), + output.stride(2), + n_stream, + hidden_size, + block_size_c, + num_warps=4, + ) + return output + + +def _hyper_connect( + res: torch.Tensor, + out: torch.Tensor, + h_post: torch.Tensor, + h_res: torch.Tensor, + backend: str, +) -> torch.Tensor: + """Mix residual streams and inject one transformer output stream.""" + if _use_triton_backend(res, backend): + return _triton_hyper_connect_forward(res, out, h_post, h_res) + out_mstream = torch.einsum("tn,tc->tnc", h_post, out) + mixed_res = torch.einsum("tij,tjc->tic", h_res, res) + return mixed_res + out_mstream + + +class MHCHandler: + """High-performance mHC helper with local Triton inference kernels.""" + + def __init__( + self, + n_stream: int, + hidden_size: int, + num_sk_iters: int = 20, + sk_eps: float = 1e-12, + dtype: torch.dtype = torch.float32, + backend: str | None = None, + ) -> None: + """Configure multi-stream hyper-connection transforms and backend.""" + self.n = n_stream + self.hidden_size = hidden_size + self.num_sk_iters = num_sk_iters + self.sk_eps = sk_eps + self.dtype = dtype + self.backend = backend or os.environ.get("MAGI2_MHC_BACKEND", "triton") + self.matmul_scale = 1.0 / math.sqrt(float(n_stream * hidden_size)) + + def flatten_mstream(self, x: torch.Tensor) -> torch.Tensor: + self._check_mstream_shape(x) + return x.view(x.size(0), -1) + + def apply_norm_and_compute_h_( + self, + x: torch.Tensor, + norm_fn: Callable[[torch.Tensor], torch.Tensor], + phi_fused: torch.Tensor, + ) -> MHCTensorTuple: + """Normalize streams and project pre-, post-, and residual parameters.""" + self._check_flatten_mstream_shape(x) + h_fused = torch.matmul(norm_fn(x).to(self.dtype), phi_fused) + h_pre, h_post, h_res = torch.split( + h_fused, [self.n, self.n, self.n * self.n], dim=-1 + ) + return h_pre, h_post, h_res.view(-1, self.n, self.n) + + def compute_and_apply_hpre( + self, + x: torch.Tensor, + abh_pre: MHCTensorTuple, + out_dtype: torch.dtype | None = None, + ) -> torch.Tensor: + """Compute pre-connection gates and reduce residual streams.""" + self._check_mstream_shape(x) + alpha, bias, h_pre_ = abh_pre + dtype = out_dtype or x.dtype + h_pre = _sigmoid_affine( + h_pre_, + alpha, + bias, + self.matmul_scale, + sigmoid_scale=1.0, + out_dtype=dtype, + backend=self.backend, + ) + return _apply_hpre(h_pre, x, self.backend) + + def compute_hpost_and_hres( + self, + abh_post: MHCTensorTuple, + abh_res: MHCTensorTuple, + out_dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute output gates and residual-stream mixing matrices.""" + alpha_post, bias_post, h_post_ = abh_post + alpha_res, bias_res, h_res_ = abh_res + dtype = out_dtype or h_post_.dtype + h_post = _sigmoid_affine( + h_post_, + alpha_post, + bias_post, + self.matmul_scale, + sigmoid_scale=2.0, + out_dtype=dtype, + backend=self.backend, + ) + h_res = _sinkhorn_knopp_affine( + h_res_, + alpha_res, + bias_res, + self.matmul_scale, + self.num_sk_iters, + self.sk_eps, + dtype, + self.backend, + ) + return h_post, h_res + + def apply_hpost_and_hres( + self, + res: torch.Tensor, + out: torch.Tensor, + h_post: torch.Tensor, + h_res: torch.Tensor, + ) -> torch.Tensor: + self._check_mstream_shape(res) + self._check_sstream_shape(out) + return _hyper_connect(res, out, h_post, h_res, self.backend) + + def _check_sstream_shape(self, x: torch.Tensor) -> None: + if x.dim() != 2 or x.size(-1) != self.hidden_size: + raise ValueError( + f"Expected single-stream shape (T, {self.hidden_size}), got {tuple(x.shape)}" + ) + + def _check_mstream_shape(self, x: torch.Tensor) -> None: + if x.dim() != 3 or x.size(1) != self.n or x.size(-1) != self.hidden_size: + raise ValueError( + f"Expected multi-stream shape (T, {self.n}, {self.hidden_size}), got {tuple(x.shape)}" + ) + + def _check_flatten_mstream_shape(self, x: torch.Tensor) -> None: + expected = self.n * self.hidden_size + if x.dim() != 2 or x.size(-1) != expected: + raise ValueError( + f"Expected flattened multi-stream shape (T, {expected}), got {tuple(x.shape)}" + ) + + +# ============================================================ +# modeling +# ============================================================ + +os.environ.setdefault("MAGI_ATTENTION_WORKSPACE_BASE", "/tmp") + + +@dataclass +class VarlenHandler: + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + + +def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor: + """Encode scalar diffusion positions with sinusoidal frequencies.""" + position = position.to(torch.float32) * 1000.0 + half = dim // 2 + freqs = torch.exp( + -math.log(10000) + * torch.arange(start=0, end=half, dtype=torch.float32, device=position.device) + / half + ) + args = position[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +def _attention_sinks_for_cp( + attn_sinks: Optional[torch.Tensor], +) -> Optional[torch.Tensor]: + if attn_sinks is None or psm.get_world_size("cp") <= 1: + return attn_sinks + return torch.chunk(attn_sinks, chunks=psm.get_world_size("cp"), dim=-1)[ + dist.get_rank(psm.get_parallel_group("cp")) + ].contiguous() + + +def _fa3_varlen_func_with_sink_inference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + softcap: float, + sink: Optional[torch.Tensor], + sink_layout: str, + deterministic: bool = False, +) -> torch.Tensor: + """Run variable-length FlashAttention 3 and correct attention-sink outputs.""" + if torch.is_grad_enabled(): + return fa3_varlen_func_with_sink( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softcap=softcap, + sink=sink, + sink_layout=sink_layout, + deterministic=deterministic + or os.environ.get("MAGI2_DETERMINISTIC", "0") == "1", + ) + + softmax_scale = q.shape[-1] ** -0.5 + out, softmax_lse, *_ = flash_attn_3_cuda.fwd( + q, + k, + v, + None, # k_new + None, # v_new + None, # qv + None, # out + cu_seqlens_q.contiguous(), + cu_seqlens_k.contiguous(), + None, # cu_seqlens_k_new + None, # seqused_q + None, # seqused_k + max_seqlen_q, + max_seqlen_k, + None, # page_table + None, # kv_batch_idx + None, # leftpad_k + None, # rotary_cos + None, # rotary_sin + None, # seqlens_rotary + None, # q_descale + None, # k_descale + None, # v_descale + softmax_scale, + False, # causal + -1, # window_size_left + -1, # window_size_right + 0, # attention_chunk + softcap, + True, # rotary_interleaved + None, # scheduler_metadata + 1, # num_splits + None, # pack_gqa + 0, # sm_margin + ) + out, _ = FA3VarlenFuncWithSink.correct_out_lse_with_sink( + out=out, lse=softmax_lse, sink=sink, sink_layout=sink_layout + ) + return out + + +def _flash_attn_with_sink_impl( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + cp_split_sizes: list[int], + attn_softcap: float, + attn_sinks: Optional[torch.Tensor], +) -> torch.Tensor: + """Execute sink attention with context-parallel head and sequence exchange.""" + q, k, v = q.squeeze(0).bfloat16(), k.squeeze(0).bfloat16(), v.squeeze(0).bfloat16() + if attn_sinks is not None and attn_sinks.numel() == 0: + attn_sinks = None + attn_sinks = _attention_sinks_for_cp(attn_sinks) + + if psm.get_world_size("cp") > 1: + q, k, v = batch_scatter_head_gather_seqlen( + [q, k, v], cp_split_sizes, psm.get_parallel_group("cp") + ) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + + if cu_seqlens_q is not None: + out = _fa3_varlen_func_with_sink_inference( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softcap=attn_softcap, + sink=attn_sinks, + sink_layout="sh", + deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1", + ) + else: + out = fa3_func_with_sink( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + softcap=attn_softcap, + sink=attn_sinks.unsqueeze(0) if attn_sinks is not None else None, + sink_layout="sh", + deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1", + ) + out = out.squeeze(0) + + if psm.get_world_size("cp") > 1: + out = scatter_seqlen_gather_head( + out, cp_split_sizes, psm.get_parallel_group("cp"), async_op=False + ) + out = rearrange( + out, "(cp sq) hn hd -> sq (cp hn) hd", cp=psm.get_world_size("cp") + ) + return out + + +def _fa3_with_sink_cp_meta( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_sinks: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_k: Optional[torch.Tensor], + cp_split_sizes: torch.Tensor, + attn_softcap: float, +) -> torch.Tensor: + """Describe the compiled context-parallel attention output tensor.""" + return torch.empty_like(q, dtype=torch.bfloat16).squeeze(0) + + +@magi_register_custom_op( + "infra::magi2_fa3_with_sink_cp", + infer_output_meta_fn=_fa3_with_sink_cp_meta, + is_subgraph_boundary=True, +) +def magi2_fa3_with_sink_cp( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_sinks: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_k: Optional[torch.Tensor], + cp_split_sizes: torch.Tensor, + attn_softcap: float, +) -> torch.Tensor: + """FA3 attention with sink tokens and context parallelism for MAGI-2. + + max_seqlen_q/k are derived from cu_seqlens inside this boundary op instead + of being passed as int args, because MagiCompile specializes Python ints as + constants — see test_boundary_int_args_are_specialized. + + cp_split_sizes is passed as a Tensor (created outside compiled code) so its + values flow dynamically through the FX graph. + """ + if cu_seqlens_q is not None: + max_seqlen_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item() + max_seqlen_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).max().item() + else: + max_seqlen_q = q.shape[0] + max_seqlen_k = k.shape[0] + return _flash_attn_with_sink_impl( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + cp_split_sizes=cp_split_sizes.tolist(), + attn_softcap=attn_softcap, + attn_sinks=attn_sinks, + ) + + +def flash_attn_with_sink( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + varlen_handler: VarlenHandler, + cp_split_sizes: torch.Tensor, + attn_softcap: float, + attn_sinks: Optional[torch.Tensor], +) -> torch.Tensor: + """Invoke sink attention using packed variable-length sequence metadata.""" + sinks_tensor = ( + attn_sinks if attn_sinks is not None else torch.empty(0, device=q.device) + ) + return magi2_fa3_with_sink_cp( + q, + k, + v, + sinks_tensor, + cu_seqlens_q=varlen_handler.cu_seqlens_q, + cu_seqlens_k=varlen_handler.cu_seqlens_k, + cp_split_sizes=cp_split_sizes, + attn_softcap=attn_softcap, + ) + + +@dataclass +class AttentionConfig: + hidden_size: int + num_heads_q: int + num_heads_kv: int + head_dim: int + params_dtype: torch.dtype + num_modality: int + num_layers: int + attn_softcap: float = -1.0 + sink_token_num: int = 1 + + +class Attention(nn.Module): + def __init__(self, config: AttentionConfig): + """Build modality-aware projections, normalization, and attention sinks.""" + super().__init__() + self.config = config + self.pre_norm = MultiModalityRMSNorm( + config.hidden_size, eps=1e-6, num_modality=config.num_modality + ) + + self.head_dim = config.head_dim + self.num_heads_q = config.num_heads_q + self.num_heads_kv = config.num_heads_kv + self.q_size = self.num_heads_q * self.head_dim + self.kv_size = self.num_heads_kv * self.head_dim + + self.gating_total_size = self.num_heads_q + + out_dim = self.q_size + self.kv_size * 2 + self.linear_g = create_linear( + config.hidden_size, + self.gating_total_size, + num_experts=config.num_modality, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + ) + + self.linear_qkv = create_linear( + config.hidden_size, + out_dim, + num_experts=config.num_modality, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + ) + self.linear_proj = create_linear( + self.q_size, + config.hidden_size, + bias=False, + num_experts=config.num_modality, + dtype=config.params_dtype, + num_layers=config.num_layers, + ) + + self.sinks = nn.Parameter( + torch.empty( + max(1, config.sink_token_num), config.num_heads_q, dtype=torch.float32 + ) + ) + + self.q_norm = MultiModalityRMSNorm( + config.head_dim, num_modality=config.num_modality, out_dtype=torch.float32 + ) + self.k_norm = MultiModalityRMSNorm( + config.head_dim, num_modality=config.num_modality, out_dtype=torch.float32 + ) + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + varlen_handler: VarlenHandler, + modality_dispatcher: ModalityDispatcher, + cp_split_sizes: torch.Tensor, + ) -> torch.Tensor: + """Apply modality-aware sink attention to packed token features.""" + x = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher) + g = self.linear_g(x, modality_dispatcher=modality_dispatcher) + qkv = self.linear_qkv(x, modality_dispatcher=modality_dispatcher) + q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=1) + + q = q.view(hidden_states.size(0), self.num_heads_q, self.head_dim).clone() + k = k.view(hidden_states.size(0), self.num_heads_kv, self.head_dim).clone() + v = v.view(hidden_states.size(0), self.num_heads_kv, self.head_dim) + g = g.view(hidden_states.size(0), self.num_heads_q, 1).clone() + + q = self.q_norm(q, modality_dispatcher=modality_dispatcher) + k = self.k_norm(k, modality_dispatcher=modality_dispatcher) + + q = modality_dispatcher._inv_permute(q).unsqueeze(0) + k = modality_dispatcher._inv_permute(k).unsqueeze(0) + v = modality_dispatcher._inv_permute(v).unsqueeze(0) + + sin_emb, cos_emb = rope.tensor_split(2, -1) + q = apply_rotary_emb_torch(q, cos_emb, sin_emb) + k = apply_rotary_emb_torch(k, cos_emb, sin_emb) + out = flash_attn_with_sink( + q, + k, + v, + varlen_handler, + cp_split_sizes, + self.config.attn_softcap, + self.sinks, + ) + + out = modality_dispatcher._permute(out) + out = out * torch.sigmoid(g) + out = out.reshape(-1, self.q_size).to(torch.bfloat16) + return self.linear_proj(out, modality_dispatcher=modality_dispatcher) + + +@dataclass +class MLPConfig: + hidden_size: int + intermediate_size: int + activation_type: MLPActivationType + params_dtype: torch.dtype + num_modality: int = 1 + num_layers: int = 1 + gated_act: bool = False + + +class MLP(nn.Module): + def __init__(self, config: MLPConfig) -> None: + """Build modality-aware gated feed-forward projections.""" + super().__init__() + self.config = config + self.pre_norm = MultiModalityRMSNorm( + config.hidden_size, num_modality=config.num_modality + ) + intermediate_size_up = ( + config.intermediate_size * 2 + if config.gated_act + else config.intermediate_size + ) + self.up_gate_proj = create_linear( + config.hidden_size, + intermediate_size_up, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + num_experts=config.num_modality, + ) + self.down_proj = create_linear( + config.intermediate_size, + config.hidden_size, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + num_experts=config.num_modality, + ) + self.activation_func = create_activation_func(config.activation_type) + + def forward( + self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher + ) -> torch.Tensor: + x = self.pre_norm(x, modality_dispatcher=modality_dispatcher) + x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher) + x = self.activation_func(x) + x = self.down_proj(x, modality_dispatcher=modality_dispatcher) + return x + + +@dataclass +class CoreMultiHeadMoEConfig: + hidden_size: int + num_heads: int + num_experts: int + top_k: int + expert_intermediate_size: int + num_layers: int + params_dtype: torch.dtype + score_func: str = "sigmoid" + route_norm: bool = True + route_scale: float = 1.0 + recompute_once_in_bwd: bool = True + + +_magi2_moe_registry: dict[int, "CoreMultiHeadMoE"] = {} + + +class CoreMultiHeadMoE(nn.Module): + def __init__(self, config: CoreMultiHeadMoEConfig) -> None: + """Build expert-parallel multi-head mixture-of-experts routing state.""" + super().__init__() + if config.hidden_size % config.num_heads != 0: + raise ValueError("hidden_size must be divisible by num_heads") + self._moe_key = len(_magi2_moe_registry) + _magi2_moe_registry[self._moe_key] = self + self.config = config + self.num_heads = config.num_heads + self.num_experts = config.num_experts + self.topk = config.top_k + self.d_head = config.hidden_size // config.num_heads + self.d_expert = config.expert_intermediate_size + self.flatten_num_experts = self.num_heads * self.num_experts + self.ep_size = psm.get_world_size("ep") + + # When num_heads is not divisible by ep_size, pad to next multiple. + # Ranks beyond the real head range hold dummy (zero) weights and + # participate in all-to-all but produce discarded output. + if self.ep_size > 1 and self.num_heads % self.ep_size != 0: + import math + + self.padded_num_heads = ( + math.ceil(self.num_heads / self.ep_size) * self.ep_size + ) + self.ep_pad_heads = self.padded_num_heads - self.num_heads + ep_rank = psm.get_local_rank("ep") + real_heads_end = self.num_heads + my_head_start = ep_rank * (self.padded_num_heads // self.ep_size) + self.has_real_moe_heads = my_head_start < real_heads_end + else: + self.padded_num_heads = self.num_heads + self.ep_pad_heads = 0 + self.has_real_moe_heads = True + + if self.ep_pad_heads > 0: + self.register_buffer( + "_ep_pad_zeros", + torch.zeros( + 1, self.ep_pad_heads, self.d_head, dtype=config.params_dtype + ), + persistent=False, + ) + + self.local_num_heads = self.padded_num_heads // self.ep_size + self.local_flatten_num_experts = self.local_num_heads * self.num_experts + + self.gate = nn.Parameter( + torch.empty( + self.local_flatten_num_experts, self.d_head, dtype=torch.float32 + ) + ) + self.W_gate = nn.Parameter( + torch.empty( + self.local_flatten_num_experts, + self.d_head, + self.d_expert, + dtype=config.params_dtype, + ) + ) + self.W_up = nn.Parameter( + torch.empty( + self.local_flatten_num_experts, + self.d_head, + self.d_expert, + dtype=config.params_dtype, + ) + ) + self.W_down = nn.Parameter( + torch.empty( + self.local_flatten_num_experts, + self.d_expert, + self.d_head, + dtype=config.params_dtype, + ) + ) + self.router = nn.Module() + self.router.register_buffer( + "expert_bias", + torch.zeros(self.local_flatten_num_experts, dtype=torch.float32), + ) + self.router.register_buffer( + "expert_bias_ema", + torch.zeros(self.local_flatten_num_experts, dtype=torch.float32), + ) + + if not self.has_real_moe_heads: + with torch.no_grad(): + self.gate.zero_() + self.W_gate.zero_() + self.W_up.zero_() + self.W_down.zero_() + + def _route(self, x_heads: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Select top-k experts and probabilities for each token head.""" + num_heads = x_heads.size(1) + gate = self.gate.view(num_heads, self.num_experts, self.d_head).float() + router_logits = torch.einsum("shd,hed->hse", x_heads.float(), gate) + expert_bias = self.router.expert_bias.view(num_heads, self.num_experts) + topk_probs, topk_indices = _mh_moe_compute_topk( + router_logits=router_logits, + top_k=self.topk, + score_func=self.config.score_func, + expert_bias=expert_bias, + route_norm=self.config.route_norm, + ) + topk_probs = topk_probs * self.config.route_scale + return topk_probs, topk_indices + + def _flash_forward( + self, + x_heads: torch.Tensor, + topk_probs: torch.Tensor, + topk_indices: torch.Tensor, + ) -> torch.Tensor: + """Execute sorted fused expert computation for routed token heads.""" + gather_ids, probs, expert_offsets = _mh_moe_global_sort( + topk_probs=topk_probs, + topk_indices=topk_indices, + num_experts=self.num_experts, + ) + return _mh_moe_fwd( + x=x_heads, + gather_ids=gather_ids, + probs=probs, + expert_offsets=expert_offsets, + W_gate=self.W_gate, + W_up=self.W_up, + W_down=self.W_down, + ) + + def _forward_impl(self, x: torch.Tensor) -> torch.Tensor: + """Run expert-parallel routing and fused experts across token heads.""" + x_heads = x.view(-1, self.num_heads, self.d_head) + + if self.ep_pad_heads > 0: + x_heads = torch.cat( + [x_heads, self._ep_pad_zeros.expand(x_heads.shape[0], -1, -1)], dim=1 + ) + + if self.ep_size > 1: + ep_group = psm.get_parallel_group("ep") + if ep_group is None: + raise RuntimeError( + f"MAGI-2 EP group is not initialized but ep_size={self.ep_size}" + ) + x_heads = ep_dispatch(x_heads, ep_group=ep_group) + + if self.has_real_moe_heads: + topk_probs, topk_indices = self._route(x_heads) + out = self._flash_forward(x_heads, topk_probs, topk_indices) + else: + out = torch.zeros_like(x_heads) + + if self.ep_size > 1: + out = ep_undispatch(out, ep_group=ep_group) + + if self.ep_pad_heads > 0: + out = out[:, : self.num_heads, :] + + return out.reshape(-1, self.num_heads * self.d_head) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return magi2_mh_moe(x, self._moe_key) + + +def _magi2_mh_moe_meta(x: torch.Tensor, moe_key: int) -> torch.Tensor: + moe_module: CoreMultiHeadMoE = _magi2_moe_registry[moe_key] + if moe_module.ep_pad_heads > 0: + padded_stride = moe_module.padded_num_heads * moe_module.d_head + return torch.empty_strided( + x.shape, (padded_stride, 1), dtype=x.dtype, device=x.device + ) + return torch.empty_like(x) + + +@magi_register_custom_op( + "infra::magi2_mh_moe", + infer_output_meta_fn=_magi2_mh_moe_meta, + is_subgraph_boundary=True, +) +def magi2_mh_moe(x: torch.Tensor, moe_key: int) -> torch.Tensor: + """MoE forward with EP dispatch/undispatch for MAGI-2. + + Module is looked up from _magi2_moe_registry by a stable int key assigned + during __init__ (monotonic counter, not id(self)), so Dynamo's int + specialization produces deterministic guards that survive cache reuse + across process restarts. + """ + moe_module: CoreMultiHeadMoE = _magi2_moe_registry[moe_key] + return moe_module._forward_impl(x) + + +@dataclass +class MultiHeadMoELayerConfig: + hidden_size: int + num_modality: int + num_layers: int + activation_type: MLPActivationType + params_dtype: torch.dtype + shared_expert_intermediate_sizes: int + modality_specific_expert_intermediate_sizes: int + shared_expert_params_dtype: torch.dtype + split_merge_intermediate_size: Optional[int] + moe_config: MoEConfig + + +class MultiHeadMoELayer(nn.Module): + def __init__(self, config: MultiHeadMoELayerConfig) -> None: + """Build routed and shared expert branches for one transformer layer.""" + super().__init__() + self.config = config + moe_hidden_size = config.split_merge_intermediate_size or config.hidden_size + self.moe_mlp = CoreMultiHeadMoE( + CoreMultiHeadMoEConfig( + hidden_size=moe_hidden_size, + num_heads=config.moe_config.num_heads, + num_experts=config.moe_config.num_experts, + top_k=config.moe_config.top_k, + expert_intermediate_size=config.moe_config.expert_intermediate_size, + num_layers=config.num_layers, + params_dtype=config.params_dtype, + score_func=config.moe_config.score_func, + route_norm=config.moe_config.route_norm, + route_scale=config.moe_config.route_scale, + ) + ) + self.is_swiglu = config.activation_type in [MLPActivationType.SWIGLU7] + self.activation_func = create_activation_func(config.activation_type) + self.shared_expert_fc1, self.shared_expert_fc2 = ( + self._create_shared_expert_pair( + config, config.shared_expert_intermediate_sizes, num_modality=1 + ) + ) + ( + self.modality_specific_shared_expert_fc1, + self.modality_specific_shared_expert_fc2, + ) = self._create_shared_expert_pair( + config, + config.modality_specific_expert_intermediate_sizes, + num_modality=config.num_modality, + ) + self.pre_norm = MultiModalityRMSNorm( + config.hidden_size, num_modality=config.num_modality + ) + mid = config.split_merge_intermediate_size or config.hidden_size + self.split_linear = create_linear( + config.hidden_size, + mid, + num_layers=config.num_layers, + dtype=config.shared_expert_params_dtype, + bias=False, + ) + self.merge_linear = create_linear( + mid, + config.hidden_size, + num_layers=config.num_layers, + dtype=config.shared_expert_params_dtype, + bias=False, + ) + + @staticmethod + def _create_shared_expert_pair( + config: MultiHeadMoELayerConfig, intermediate_size: int, num_modality: int = 1 + ) -> tuple[nn.Module, nn.Module]: + """Build input and output projections for one shared expert branch.""" + factor = 2 if config.activation_type in [MLPActivationType.SWIGLU7] else 1 + fc1 = create_linear( + config.hidden_size, + intermediate_size * factor, + num_layers=config.num_layers, + dtype=config.shared_expert_params_dtype, + num_experts=num_modality, + bias=False, + ) + fc2 = create_linear( + intermediate_size, + config.hidden_size, + num_layers=config.num_layers, + dtype=config.shared_expert_params_dtype, + num_experts=num_modality, + bias=False, + ) + return fc1, fc2 + + def _shared_expert_forward( + self, norm_output: torch.Tensor, modality_dispatcher: ModalityDispatcher + ) -> torch.Tensor: + """Run global and modality-specific shared expert branches.""" + x1 = self.shared_expert_fc1(norm_output) + x2 = self.modality_specific_shared_expert_fc1( + norm_output, modality_dispatcher=modality_dispatcher + ) + x = torch.concat([x1, x2], dim=-1) + x = self.activation_func(x) + x1, x2 = x.split( + [ + self.config.shared_expert_intermediate_sizes, + self.config.modality_specific_expert_intermediate_sizes, + ], + dim=-1, + ) + x1 = self.shared_expert_fc2(x1) + x2 = self.modality_specific_shared_expert_fc2( + x2.contiguous(), modality_dispatcher=modality_dispatcher + ) + return x1 + x2 + + def forward( + self, hidden_states: torch.Tensor, modality_dispatcher: ModalityDispatcher + ) -> torch.Tensor: + norm_output = self.pre_norm( + hidden_states, modality_dispatcher=modality_dispatcher + ) + moe_input = self.split_linear(norm_output) + moe_out = self.moe_mlp(moe_input) + moe_out = self.merge_linear(moe_out) + return moe_out + self._shared_expert_forward(norm_output, modality_dispatcher) + + +@dataclass +class AdapterConfig: + hidden_size: int + num_attention_heads: int + text_in_channels: int + video_in_channels: int + audio_in_channels: int + virtual_width_factor: int = 1 + + +class PreAdapter(nn.Module): + def __init__(self, config: AdapterConfig) -> None: + """Build modality-specific input projections and rotary embeddings.""" + super().__init__() + self.config = config + self.adapter_dim = config.hidden_size * config.virtual_width_factor + self.video_embedder = nn.Linear( + config.video_in_channels, self.adapter_dim, bias=True, dtype=torch.float32 + ) + self.text_embedder = nn.Linear( + config.text_in_channels, self.adapter_dim, bias=True, dtype=torch.float32 + ) + self.audio_embedder = nn.Linear( + config.audio_in_channels, self.adapter_dim, bias=True, dtype=torch.float32 + ) + self.rope = ElementWiseFourierEmbed( + config.hidden_size // config.num_attention_heads, + in_pixels=False, + learnable=False, + ) + + def forward( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + video_idx: torch.Tensor, + audio_idx: torch.Tensor, + text_idx: torch.Tensor, + time_idx: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Project packed modality tokens and build their rotary embeddings.""" + rope = self.rope(coords_mapping) + output_x = torch.zeros( + x.shape[0], self.adapter_dim, device=x.device, dtype=torch.float32 + ) + if text_idx.numel() > 0: + output_x.index_copy_( + 0, + text_idx, + self.text_embedder( + x.index_select(0, text_idx)[:, : self.config.text_in_channels] + ), + ) + if audio_idx.numel() > 0: + output_x.index_copy_( + 0, + audio_idx, + self.audio_embedder( + x.index_select(0, audio_idx)[:, : self.config.audio_in_channels] + ), + ) + if video_idx.numel() > 0: + output_x.index_copy_( + 0, + video_idx, + self.video_embedder( + x.index_select(0, video_idx)[:, : self.config.video_in_channels] + ), + ) + return output_x, rope + + +class PostAdapter(nn.Module): + def __init__(self, config: AdapterConfig) -> None: + """Build modality-specific output normalization and projections.""" + super().__init__() + self.config = config + self.adapter_dim = config.hidden_size * config.virtual_width_factor + self.final_norm_video = MultiModalityRMSNorm(self.adapter_dim) + self.final_norm_audio = MultiModalityRMSNorm(self.adapter_dim) + self.final_linear_video = nn.Linear( + self.adapter_dim, config.video_in_channels, bias=False, dtype=torch.float32 + ) + self.final_linear_audio = nn.Linear( + self.adapter_dim, config.audio_in_channels, bias=False, dtype=torch.float32 + ) + self.final_out_dim = max(config.video_in_channels, config.audio_in_channels) + + def forward( + self, x: torch.Tensor, video_idx: torch.Tensor, audio_idx: torch.Tensor + ) -> torch.Tensor: + """Project packed video and audio features into channel outputs.""" + x_out = torch.zeros( + x.shape[0], self.final_out_dim, device=x.device, dtype=torch.float32 + ) + if video_idx.numel() > 0: + x_video = self.final_norm_video( + x.index_select(0, video_idx).to(self.final_norm_video.weight.dtype) + ) + x_video = self.final_linear_video(x_video).to(torch.float32) + x_out[:, : self.config.video_in_channels].index_copy_(0, video_idx, x_video) + if audio_idx.numel() > 0: + x_audio = self.final_norm_audio( + x.index_select(0, audio_idx).to(self.final_norm_audio.weight.dtype) + ) + x_audio = self.final_linear_audio(x_audio).to(torch.float32) + x_out[:, : self.config.audio_in_channels].index_copy_(0, audio_idx, x_audio) + return x_out + + +class TransformerLayer(nn.Module): + def __init__(self, model_config: ModelConfig, layer_idx: int) -> None: + """Build attention, expert, and multi-stream connections for one layer.""" + super().__init__() + num_modality = 3 if layer_idx in model_config.mm_layers else 1 + self.config_mhc = model_config.mhc_config + self.attention = Attention( + AttentionConfig( + hidden_size=model_config.hidden_size, + num_heads_q=model_config.num_heads_q, + num_heads_kv=model_config.num_heads_kv, + head_dim=model_config.head_dim, + params_dtype=model_config.params_dtype, + num_modality=num_modality, + num_layers=model_config.num_layers, + attn_softcap=model_config.attn_softcap, + sink_token_num=model_config.attn_sinks.sink_token_num, + ) + ) + activation_type = MLPActivationType( + model_config.layer_activation_types.get( + layer_idx, model_config.activation_type + ) + ) + gated_act = activation_type in [MLPActivationType.SWIGLU7] + if layer_idx in model_config.moe_config.moe_layers: + moe_cfg = model_config.moe_config + self.mlp = MultiHeadMoELayer( + MultiHeadMoELayerConfig( + hidden_size=model_config.hidden_size, + num_modality=3, + num_layers=model_config.num_layers, + activation_type=activation_type, + params_dtype=model_config.params_dtype, + shared_expert_intermediate_sizes=moe_cfg.shared_expert_intermediate_size, + modality_specific_expert_intermediate_sizes=moe_cfg.modality_specific_expert_intermediate_size, + shared_expert_params_dtype=torch.bfloat16, + split_merge_intermediate_size=moe_cfg.split_merge_intermediate_size, + moe_config=moe_cfg, + ) + ) + else: + intermediate_size = ( + int(model_config.hidden_size * model_config.intermediate_factor * 2 / 3) + // 128 + * 128 + if gated_act + else int(model_config.hidden_size * model_config.intermediate_factor) + ) + self.mlp = MLP( + MLPConfig( + hidden_size=model_config.hidden_size, + intermediate_size=intermediate_size, + activation_type=activation_type, + params_dtype=model_config.params_dtype, + num_modality=num_modality, + num_layers=model_config.num_layers, + gated_act=gated_act, + ) + ) + self._init_mhc(model_config, num_modality) + + def _init_mhc(self, model_config: ModelConfig, num_modality: int) -> None: + """Allocate hyper-connection parameters for attention and feed-forward stages.""" + n = self.config_mhc.num_stream + c = model_config.hidden_size + dtype = torch.float32 + self.mhc_alpha_pre_attn = nn.Parameter( + torch.full((1,), float(self.config_mhc.alpha_init), dtype=dtype) + ) + self.mhc_alpha_post_attn = nn.Parameter( + torch.full((1,), float(self.config_mhc.alpha_init), dtype=dtype) + ) + self.mhc_alpha_res_attn = nn.Parameter( + torch.full((1,), float(self.config_mhc.alpha_init), dtype=dtype) + ) + self.mhc_alpha_pre_mlp = nn.Parameter( + torch.full((1,), float(self.config_mhc.alpha_init), dtype=dtype) + ) + self.mhc_alpha_post_mlp = nn.Parameter( + torch.full((1,), float(self.config_mhc.alpha_init), dtype=dtype) + ) + self.mhc_alpha_res_mlp = nn.Parameter( + torch.full((1,), float(self.config_mhc.alpha_init), dtype=dtype) + ) + self.mhc_bias_pre_attn = nn.Parameter(torch.empty(n, dtype=dtype)) + self.mhc_bias_post_attn = nn.Parameter(torch.empty(n, dtype=dtype)) + self.mhc_bias_pre_mlp = nn.Parameter(torch.empty(n, dtype=dtype)) + self.mhc_bias_post_mlp = nn.Parameter(torch.empty(n, dtype=dtype)) + self.mhc_bias_res_attn = nn.Parameter(torch.empty(n, n, dtype=dtype)) + self.mhc_bias_res_mlp = nn.Parameter(torch.empty(n, n, dtype=dtype)) + self.mhc_phi_fused_attn = nn.Parameter( + torch.empty(n * c, n + n + (n * n), dtype=dtype) + ) + self.mhc_phi_fused_mlp = nn.Parameter( + torch.empty(n * c, n + n + (n * n), dtype=dtype) + ) + self.mhc_norm = MultiModalityRMSNorm( + c * n, num_modality=num_modality, out_dtype=dtype + ) + self.mhc_handler = MHCHandler( + n_stream=n, hidden_size=c, num_sk_iters=20, sk_eps=1e-12, dtype=dtype + ) + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + varlen_handler: VarlenHandler, + modality_dispatcher: ModalityDispatcher, + cp_split_sizes: torch.Tensor, + ) -> torch.Tensor: + """Run attention and feed-forward stages through multi-stream connections.""" + n = self.config_mhc.num_stream + s = hidden_states.reshape(hidden_states.shape[0], n, -1) + assert s.is_contiguous() + + h_post_attn_, h_res_attn_, attn_in = self._forward_mhc_pre_attn( + s, modality_dispatcher + ) + attn_out = self.attention( + attn_in, rope, varlen_handler, modality_dispatcher, cp_split_sizes + ) + + s_, h_post_mlp_, h_res_mlp_, mlp_in = self._forward_mhc_attn_to_mlp( + s, attn_out, h_post_attn_, h_res_attn_, modality_dispatcher + ) + mlp_out = self.mlp(mlp_in, modality_dispatcher) + + s = self._forward_mhc_post_mlp(s_, mlp_out, h_post_mlp_, h_res_mlp_) + return s.reshape(s.shape[0], -1) + + def _forward_mhc_pre_attn( + self, s: torch.Tensor, modality_dispatcher: ModalityDispatcher + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute pre-attention stream gates and the attention input.""" + s_flat = self.mhc_handler.flatten_mstream(s) + h_pre_attn_, h_post_attn_, h_res_attn_ = ( + self.mhc_handler.apply_norm_and_compute_h_( + x=s_flat, + norm_fn=partial(self.mhc_norm, modality_dispatcher=modality_dispatcher), + phi_fused=self.mhc_phi_fused_attn, + ) + ) + attn_in = self.mhc_handler.compute_and_apply_hpre( + x=s, + abh_pre=(self.mhc_alpha_pre_attn, self.mhc_bias_pre_attn, h_pre_attn_), + out_dtype=s.dtype, + ) + return h_post_attn_, h_res_attn_, attn_in + + def _forward_mhc_attn_to_mlp( + self, + s: torch.Tensor, + attn_out: torch.Tensor, + h_post_attn_: torch.Tensor, + h_res_attn_: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Mix attention output into residual streams and prepare the MLP input.""" + h_post_attn, h_res_attn = self.mhc_handler.compute_hpost_and_hres( + abh_post=(self.mhc_alpha_post_attn, self.mhc_bias_post_attn, h_post_attn_), + abh_res=(self.mhc_alpha_res_attn, self.mhc_bias_res_attn, h_res_attn_), + out_dtype=s.dtype, + ) + + s_ = self.mhc_handler.apply_hpost_and_hres( + res=s, out=attn_out, h_post=h_post_attn, h_res=h_res_attn + ) + s_flat = self.mhc_handler.flatten_mstream(s_) + h_pre_mlp_, h_post_mlp_, h_res_mlp_ = ( + self.mhc_handler.apply_norm_and_compute_h_( + x=s_flat, + norm_fn=partial(self.mhc_norm, modality_dispatcher=modality_dispatcher), + phi_fused=self.mhc_phi_fused_mlp, + ) + ) + mlp_in = self.mhc_handler.compute_and_apply_hpre( + x=s_, + abh_pre=(self.mhc_alpha_pre_mlp, self.mhc_bias_pre_mlp, h_pre_mlp_), + out_dtype=s.dtype, + ) + return s_, h_post_mlp_, h_res_mlp_, mlp_in + + def _forward_mhc_post_mlp( + self, + s_: torch.Tensor, + mlp_out: torch.Tensor, + h_post_mlp_: torch.Tensor, + h_res_mlp_: torch.Tensor, + ) -> torch.Tensor: + """Mix the feed-forward output back into residual streams.""" + h_post_mlp, h_res_mlp = self.mhc_handler.compute_hpost_and_hres( + abh_post=(self.mhc_alpha_post_mlp, self.mhc_bias_post_mlp, h_post_mlp_), + abh_res=(self.mhc_alpha_res_mlp, self.mhc_bias_res_mlp, h_res_mlp_), + out_dtype=s_.dtype, + ) + return self.mhc_handler.apply_hpost_and_hres( + res=s_, out=mlp_out, h_post=h_post_mlp, h_res=h_res_mlp + ) + + +@magi_compile( + dynamic_arg_dims={ + "x": 0, + "rope": 0, + "varlen_handler.cu_seqlens_q": 0, + "varlen_handler.cu_seqlens_k": 0, + "modality_dispatcher.modality_mapping": 0, + "modality_dispatcher.permuted_modality_mapping": 0, + "modality_dispatcher.permute_mapping": 0, + "modality_dispatcher.inv_permute_mapping": 0, + } +) +class TransformerBlock(nn.Module): + def __init__(self, model_config: ModelConfig) -> None: + super().__init__() + self.layers = nn.ModuleList( + [TransformerLayer(model_config, i) for i in range(model_config.num_layers)] + ) + + def forward( + self, + x: torch.Tensor, + rope: torch.Tensor, + varlen_handler: VarlenHandler, + modality_dispatcher: ModalityDispatcher, + cp_split_sizes: torch.Tensor, + ) -> torch.Tensor: + """Apply every transformer layer to context-parallel packed tokens.""" + x = x.to(torch.bfloat16) + x = modality_dispatcher._permute(x) + for layer in self.layers: + x = layer(x, rope, varlen_handler, modality_dispatcher, cp_split_sizes) + x = modality_dispatcher._inv_permute(x) + return x + + +def _initialize_magi2_skip_load_weights(module: nn.Module) -> None: + """Initialize placeholder parameters when checkpoint loading is skipped.""" + with torch.no_grad(): + for name, submodule in module.named_modules(): + if isinstance(submodule, CustomGroupedLinear): + if name.endswith(".split_linear"): + submodule.weight.fill_(1e-3) + else: + submodule.weight.zero_() + if submodule.bias is not None: + submodule.bias.zero_() + elif isinstance(submodule, MultiModalityRMSNorm): + submodule.weight.zero_() + elif isinstance(submodule, Attention): + submodule.sinks.zero_() + elif isinstance(submodule, CoreMultiHeadMoE): + nn.init.normal_(submodule.gate, mean=0.0, std=0.02) + submodule.W_gate.zero_() + submodule.W_up.zero_() + submodule.W_down.zero_() + submodule.router.expert_bias.zero_() + submodule.router.expert_bias_ema.zero_() + elif isinstance(submodule, TransformerLayer): + for param_name, parameter in submodule.named_parameters(recurse=False): + if param_name.startswith(("mhc_bias_", "mhc_phi_fused_")): + parameter.zero_() + + +def _validate_inference_model_config(model_config: ModelConfig) -> None: + if not model_config.mhc_config.enable: + raise ValueError("MAGI-2 inference expects mhc_config.enable=true.") + if not model_config.attn_gating.enable: + raise ValueError( + "MAGI-2 inference expects scalar attention gating to be enabled." + ) + if not model_config.attn_sinks.enable: + raise ValueError("MAGI-2 inference expects sh-layout attention sinks.") + + +class Transformer(BaseDiT): + _config = Magi2PreviewVideoConfig() + _fsdp_shard_conditions = _config._fsdp_shard_conditions + _compile_conditions = _config._compile_conditions + _supported_attention_backends = _config._supported_attention_backends + param_names_mapping = _config.param_names_mapping + reverse_param_names_mapping = _config.reverse_param_names_mapping + lora_param_names_mapping = _config.lora_param_names_mapping + + def __init__( + self, + config: Magi2PreviewVideoConfig, + hf_config: dict | None = None, + ep_size: int = 1, + **kwargs, + ) -> None: + """Construct the production MAGI-2 preview transformer with official module names.""" + del kwargs + super().__init__(config=config, hf_config=hf_config or {}) + model_config = config.arch_config + model_config.params_dtype = _resolve_params_dtype(model_config.params_dtype) + _validate_inference_model_config(model_config) + self.config = model_config + self.ep_size = ep_size + self.hidden_size = model_config.hidden_size + self.num_attention_heads = model_config.num_heads_q + self.num_channels_latents = model_config.num_channels_latents + adapter_config = AdapterConfig( + hidden_size=model_config.hidden_size, + num_attention_heads=model_config.num_heads_q, + text_in_channels=model_config.text_in_channels, + video_in_channels=model_config.video_in_channels, + audio_in_channels=model_config.audio_in_channels, + virtual_width_factor=model_config.mhc_config.num_stream, + ) + self.pre_adapter = PreAdapter(adapter_config) + self.post_adapter = PostAdapter(adapter_config) + self.block = TransformerBlock(model_config) + if env_is_true("SKIP_LOAD_MODEL"): + self.initialize_for_skip_load() + self.__post_init__() + + def initialize_for_skip_load(self) -> None: + _initialize_magi2_skip_load_weights(self) + + def forward( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + modality_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + time_token_sequence: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Transform packed video, audio, text, and time-token features.""" + x = ulysses_scheduler().dispatch(x) + coords_mapping = ulysses_scheduler().dispatch(coords_mapping) + modality_mapping = ulysses_scheduler().dispatch(modality_mapping) + if time_token_sequence is not None: + time_token_sequence = ulysses_scheduler().dispatch(time_token_sequence) + cp_split_sizes = ulysses_scheduler().cp_split_sizes + cp_split_sizes_t = torch.tensor( + cp_split_sizes, device=x.device, dtype=torch.long + ) + + time_mask = modality_mapping == Modality.TIME + modality_mapping = modality_mapping.clone() + modality_mapping[time_mask] = Modality.TEXT + modality_dispatcher = ModalityDispatcher(modality_mapping, 3) + video_idx = (modality_mapping == Modality.VIDEO).nonzero().flatten() + audio_idx = (modality_mapping == Modality.AUDIO).nonzero().flatten() + text_idx = (modality_mapping == Modality.TEXT).nonzero().flatten() + time_idx = time_mask.nonzero().flatten() + + x, rope = self.pre_adapter( + x, coords_mapping, video_idx, audio_idx, text_idx, time_idx + ) + if time_token_sequence is not None and time_token_sequence.shape[-1] > 0: + x[:, : time_token_sequence.shape[-1]] = time_token_sequence.to(x.dtype) + + x = self.block( + x, + rope, + varlen_handler=varlen_handler, + modality_dispatcher=modality_dispatcher, + cp_split_sizes=cp_split_sizes_t, + ) + x_out = self.post_adapter(x, video_idx, audio_idx) + x_out = ulysses_scheduler().undispatch(x_out) + return x_out + + +# ----------------------------------------------------------------------------- +# MoE routing +# ----------------------------------------------------------------------------- + +"""Top-k routing and global sort for multi-head MoE. + +Multi-head MoE routing with pure-torch implementations. +Sufficient for inference; training may benefit from fused CUDA/Triton variants. +""" + +from typing import Literal + +import torch + + +# ----------------------------------------------------------------------------- +# MoE expert-parallel dispatch +# ----------------------------------------------------------------------------- + +import torch + +# ========================================== +# EP Dispatch Autograd Functions +# ========================================== + + +def ep_dispatch(x: torch.Tensor, ep_group: dist.ProcessGroup) -> torch.Tensor: + """ + Dispatch tokens to the Rank with the specified Head. + + Input tensor has shape ``(S, H, D)``. + Output tensor has shape ``(S * ep_size, H / ep_size, D)``. + """ + ep_size = dist.get_world_size(ep_group) if ep_group is not None else 1 + if ep_size == 1: + return x + + S, H, D = x.shape + assert H % ep_size == 0, ( + f"Number of heads H ({H}) must be divisible by ep_size ({ep_size})" + ) + + x = x.contiguous().view(S, ep_size, H // ep_size, D) + x = x.permute(1, 0, 2, 3).contiguous() + + out = torch.empty_like(x) + dist.all_to_all_single(out, x, group=ep_group) + + return out.view(ep_size * S, H // ep_size, D) + + + +def ep_undispatch(x: torch.Tensor, ep_group: dist.ProcessGroup) -> torch.Tensor: + """ + Undispatch tokens back to the original Rank and concatenate back to the complete Head. + + Input tensor has shape ``(S * ep_size, H / ep_size, D)``. + Output tensor has shape ``(S, H, D)``. + """ + ep_size = dist.get_world_size(ep_group) if ep_group is not None else 1 + if ep_size == 1: + return x + + total_S, H_per_ep, D = x.shape + assert total_S % ep_size == 0, ( + f"Total sequence length ({total_S}) must be divisible by ep_size ({ep_size})" + ) + + S = total_S // ep_size + + x = x.contiguous().view(ep_size, S, H_per_ep, D) + + out = torch.empty_like(x) + dist.all_to_all_single(out, x, group=ep_group) + + out = out.permute(1, 0, 2, 3).contiguous() + return out.view(S, H_per_ep * ep_size, D) + + +# ----------------------------------------------------------------------------- +# MoE Triton forward +# ----------------------------------------------------------------------------- + +"""Plain multi-head MoE forward implementation (inference only).""" + + +import torch +import triton + +# ═══════════════════════════════════════════════════════════════════════════ +# SwiGLU7 device function +# ═══════════════════════════════════════════════════════════════════════════ + +SWIGLU7_ALPHA = tl.constexpr(1.702) +SWIGLU7_LIMIT = tl.constexpr(7.0) +SWIGLU7_BIAS = tl.constexpr(1.0) + + +@triton.jit +def swiglu7_fwd_func(gate, up, out_dtype: tl.constexpr): + """Clamped, weighted, biased SwiGLU (GPT-OSS variant). Returns hidden only. + + hidden = swish(clamp(gate)) * (clamp(up) + 1), with a fast sigmoid swish. + """ + gate_clamped = tl.minimum(gate, SWIGLU7_LIMIT) + up_clamped = tl.maximum(tl.minimum(up, SWIGLU7_LIMIT), -SWIGLU7_LIMIT) + + sig = tl.sigmoid(SWIGLU7_ALPHA * gate_clamped) + swish = gate_clamped * sig + swiglu = swish * (up_clamped + SWIGLU7_BIAS) + + return swiglu.to(out_dtype) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Expert-id binary search +# ═══════════════════════════════════════════════════════════════════════════ + + +@triton.jit +def binary_search_expert_id( + cu_expert_tile_counts_ptr, + tile_id, + NUM_EXPERTS: tl.constexpr, + LOG2_NUM_EXPERTS: tl.constexpr, +): + """Return the expert id owning ``tile_id`` given the cumulative tile counts.""" + lo = 0 + hi = NUM_EXPERTS + for _ in tl.static_range(0, LOG2_NUM_EXPERTS + 1): + mid = (lo + hi + 1) // 2 + below = tl.load(cu_expert_tile_counts_ptr + mid) <= tile_id + lo = tl.where(below, mid, lo) + hi = tl.where(below, hi, mid - 1) + return lo + + +Magi2PreviewDiT = Transformer +EntryClass = Magi2PreviewDiT diff --git a/fastvideo/models/dits/magi2_checkpointing.py b/fastvideo/models/dits/magi2_checkpointing.py new file mode 100644 index 000000000..f4fc7ba92 --- /dev/null +++ b/fastvideo/models/dits/magi2_checkpointing.py @@ -0,0 +1,178 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint loading for MAGI-2. + +Weights ship as a single HuggingFace-style safetensors directory holding the +full, unsharded expert stack. Every expert-parallel rank reads that same +directory but pulls only the slice of each MoE tensor it owns, so one copy on +disk serves any ``ep_size``. +""" + +import json +import os +from collections import defaultdict +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import torch +from safetensors import safe_open +from tqdm.auto import tqdm + +from fastvideo.models.dits.magi2_runtime import psm + +_INDEX_NAME = "model.safetensors.index.json" +_SINGLE_NAME = "model.safetensors" + +# MoE tensors are concatenated over (head, expert) along dim 0 and split across +# expert-parallel ranks; everything else is replicated. +_EP_SHARDED_SUFFIXES = ( + "moe_mlp.gate", + "moe_mlp.W_gate", + "moe_mlp.W_up", + "moe_mlp.W_down", + "moe_mlp.router.expert_bias", + "moe_mlp.router.expert_bias_ema", +) + + +def print_rank_0(message: str) -> None: + """Print one checkpoint message from global rank zero.""" + if psm.get_global_rank() == 0: + print(message) + + +def read_safetensors_weight_map(checkpoint_dir: str) -> dict[str, str]: + """Map every tensor name in a HuggingFace-style directory to its file.""" + index_path = os.path.join(checkpoint_dir, _INDEX_NAME) + if os.path.exists(index_path): + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + return {key: os.path.join(checkpoint_dir, shard) for key, shard in weight_map.items()} + + single_path = os.path.join(checkpoint_dir, _SINGLE_NAME) + if not os.path.exists(single_path): + raise FileNotFoundError(f"No {_INDEX_NAME} or {_SINGLE_NAME} under {checkpoint_dir}") + with safe_open(single_path, framework="pt") as f: + return {key: single_path for key in f} + + +def load_safetensors_dir(checkpoint_dir: str, desc: str = "Loading shards") -> dict[str, torch.Tensor]: + """Load every tensor from a HuggingFace-style safetensors directory.""" + return _load_by_shard(read_safetensors_weight_map(checkpoint_dir), reader=_read_whole, desc=desc) + + +def _is_ep_sharded_key(key: str) -> bool: + return ".moe_mlp." in key and key.endswith(_EP_SHARDED_SUFFIXES) + + +def _read_whole(handle: Any, key: str, _target: torch.Tensor | None) -> torch.Tensor: + return handle.get_tensor(key) + + +def _read_ep_slice(handle: Any, key: str, target: torch.Tensor | None) -> torch.Tensor: + """Read only the (head, expert) rows this expert-parallel rank owns. + + ``num_heads`` is padded up to a multiple of ``ep_size``, so the trailing + ranks may fall wholly or partly past the real experts and get zeros. + """ + if target is None or not _is_ep_sharded_key(key): + return _read_whole(handle, key, target) + + ep_size = psm.get_world_size("ep") + view = handle.get_slice(key) + real_rows = view.get_shape()[0] + shard_rows = target.shape[0] + if ep_size <= 1 or shard_rows == real_rows: + return handle.get_tensor(key) + + start = psm.get_local_rank("ep") * shard_rows + if start >= real_rows: + return torch.zeros_like(target, device="cpu") + + chunk = view[start : min(start + shard_rows, real_rows)] + if chunk.shape[0] == shard_rows: + return chunk + padding = torch.zeros((shard_rows - chunk.shape[0], *chunk.shape[1:]), dtype=chunk.dtype) + return torch.cat([chunk, padding], dim=0) + + +def _load_by_shard( + weight_map: Mapping[str, str], + reader, + desc: str, + targets: Mapping[str, torch.Tensor] | None = None, +) -> dict[str, torch.Tensor]: + """Read tensors grouped by file so each shard is opened exactly once.""" + by_shard: dict[str, list[str]] = defaultdict(list) + for key, path in weight_map.items(): + by_shard[path].append(key) + + def read_shard(item: tuple[str, list[str]]) -> dict[str, torch.Tensor]: + path, keys = item + with safe_open(path, framework="pt") as handle: + return {key: reader(handle, key, None if targets is None else targets.get(key)) for key in keys} + + state_dict: dict[str, torch.Tensor] = {} + with ThreadPoolExecutor() as executor: + futures = list(executor.map(read_shard, by_shard.items())) + for shard_state in tqdm(futures, desc=desc, total=len(futures), disable=psm.get_global_rank() != 0): + state_dict.update(shard_state) + return state_dict + + +def _apply_router_bias_ema(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Pair the router bias with exponential moving average (EMA) model weights. + + The checkpoint stores both ``expert_bias`` (unaveraged training buffer) and + ``expert_bias_ema`` (EMA-smoothed buffer). Inference uses the EMA model + weights, so the router must use the paired EMA bias; otherwise routing + biases mismatch the weights and quality degrades. Set the env var + ``MAGI2_ROUTER_BIAS_SOURCE=main`` to keep the raw ``expert_bias`` instead. + """ + if (os.environ.get("MAGI2_ROUTER_BIAS_SOURCE") or "ema").strip().lower() == "main": + print_rank_0("[magi2] Router bias source: expert_bias (raw, MAGI2_ROUTER_BIAS_SOURCE=main)") + return state_dict + + copied = 0 + for key, value in list(state_dict.items()): + if not key.endswith("moe_mlp.router.expert_bias_ema"): + continue + target_key = key[: -len("_ema")] + target = state_dict.get(target_key) + if target is None or tuple(target.shape) != tuple(value.shape): + continue + state_dict[target_key] = value.clone() + copied += 1 + + print_rank_0(f"[magi2] Router bias: expert_bias_ema -> expert_bias ({copied} tensor(s))") + return state_dict + + +def load_magi2_model_state_dict(model: torch.nn.Module, checkpoint_dir: str) -> Mapping[str, Any]: + """Load the rank-local preview state and select EMA router biases.""" + if not checkpoint_dir: + raise ValueError("checkpoint_dir must point at a safetensors checkpoint directory") + + print_rank_0(f"[magi2] Loading preview weights from {checkpoint_dir}") + targets = model.state_dict() + weight_map = {key: path for key, path in read_safetensors_weight_map(checkpoint_dir).items() if key in targets} + + state_dict = _load_by_shard(weight_map, reader=_read_ep_slice, desc="Loading shards", targets=targets) + for key, value in state_dict.items(): + target = targets[key] + if value.dtype != target.dtype and value.is_floating_point() and target.is_floating_point(): + state_dict[key] = value.to(dtype=target.dtype) + return _apply_router_bias_ema(state_dict) diff --git a/fastvideo/models/dits/magi2_loader.py b/fastvideo/models/dits/magi2_loader.py new file mode 100644 index 000000000..c27950053 --- /dev/null +++ b/fastvideo/models/dits/magi2_loader.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Production checkpoint loaders for the MAGI-2 preview and refiner DiTs.""" + +from __future__ import annotations + +import gc + +import torch + +from fastvideo.configs.models.dits.magi2 import Magi2PreviewVideoConfig, Magi2RefinerVideoConfig +from fastvideo.models.dits.magi2 import Magi2PreviewDiT +from fastvideo.models.dits.magi2_checkpointing import load_magi2_model_state_dict, load_safetensors_dir +from fastvideo.models.dits.magi2_refiner import Magi2RefinerDiT + + +def load_magi2_preview_model( + checkpoint_dir: str, + config: Magi2PreviewVideoConfig, + device: torch.device | str, +) -> Magi2PreviewDiT: + """Build the EP-local preview model and strictly load its official state.""" + model = Magi2PreviewDiT(config=config) + state_dict = load_magi2_model_state_dict(model, checkpoint_dir) + incompatible = model.load_state_dict(state_dict, strict=True) + if incompatible.missing_keys or incompatible.unexpected_keys: + raise RuntimeError( + "Strict MAGI-2 preview loading reported incompatible keys: " + f"missing={incompatible.missing_keys}, unexpected={incompatible.unexpected_keys}" + ) + del state_dict + gc.collect() + model = model.to(device=device).eval() + torch.cuda.empty_cache() + return model + + +def load_magi2_refiner_model( + checkpoint_dir: str, + config: Magi2RefinerVideoConfig, + device: torch.device | str, +) -> Magi2RefinerDiT: + """Build a replicated refiner and strictly load every official tensor.""" + model = Magi2RefinerDiT(config=config) + state_dict = load_safetensors_dir(checkpoint_dir, desc="Loading MAGI-2 refiner shards") + incompatible = model.load_state_dict(state_dict, strict=True) + if incompatible.missing_keys or incompatible.unexpected_keys: + raise RuntimeError( + "Strict MAGI-2 refiner loading reported incompatible keys: " + f"missing={incompatible.missing_keys}, unexpected={incompatible.unexpected_keys}" + ) + del state_dict + gc.collect() + model = model.to(device=device).eval() + torch.cuda.empty_cache() + return model diff --git a/fastvideo/models/dits/magi2_refiner.py b/fastvideo/models/dits/magi2_refiner.py new file mode 100644 index 000000000..348cb98a2 --- /dev/null +++ b/fastvideo/models/dits/magi2_refiner.py @@ -0,0 +1,2352 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MAGI-2 super-resolution model.""" + +import importlib +import logging +import math +import os +from dataclasses import dataclass, field +from enum import Enum, IntEnum +from typing import ( + Callable, + Dict, + List, + Literal, + Optional, + Tuple, +) + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from flash_attn.layers.rotary import apply_rotary_emb +from magi_compiler import magi_compile +from magi_compiler.api import magi_register_custom_op +from magi_compiler.config import CompileConfig +from torch import Tensor +from torch.nn import Parameter + +from fastvideo.configs.models.dits.magi2 import Magi2RefinerArchConfig, Magi2RefinerVideoConfig +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.magi2_runtime import psm +from fastvideo.models.dits.magi2_runtime.all_to_all_primitive import ( + batch_scatter_head_gather_seqlen, + scatter_seqlen_gather_head, +) +from fastvideo.models.dits.magi2_runtime.context_parallel import ulysses_scheduler + + +logger = logging.getLogger(__name__) + +ModelConfig = Magi2RefinerArchConfig + + +def _resolve_params_dtype(value: torch.dtype | str) -> torch.dtype: + """Convert a serialized MAGI-2 refiner parameter dtype into a PyTorch dtype.""" + if isinstance(value, torch.dtype): + return value + dtype_by_name = { + "bfloat16": torch.bfloat16, + "torch.bfloat16": torch.bfloat16, + "float16": torch.float16, + "torch.float16": torch.float16, + "float32": torch.float32, + "torch.float32": torch.float32, + } + try: + return dtype_by_name[value] + except KeyError as error: + raise ValueError(f"Unsupported MAGI-2 refiner parameter dtype: {value!r}") from error + +# ============================================================ +# _merge_enum +# ============================================================ + + +# Define the MLP activation type +class MLPActivationType(Enum): + """Enumeration of supported activation functions for MLP""" + + SWIGLU7 = "swiglu7" + + +class Modality(IntEnum): + VIDEO = 0 + AUDIO = 1 + TEXT = 2 + + +# ============================================================ +# _merge_activation +# ============================================================ + + +def swiglu7( + x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None +): + """Apply the clamped SwiGLU7 activation and cast to the requested dtype.""" + out_dtype = x.dtype if out_dtype is None else out_dtype + x = x.to(torch.float32) + x_glu, x_linear = x[..., ::2], x[..., 1::2] + # Clamp the input values + x_glu = x_glu.clamp(min=None, max=limit) + x_linear = x_linear.clamp(min=-limit, max=limit) + out_glu = x_glu * torch.sigmoid(alpha * x_glu) + # Note we add an extra bias of 1 to the linear layer (from GPT-OSS) + return (out_glu * (x_linear + 1)).to(out_dtype) + + +def create_activation_func(activation_type: MLPActivationType) -> Callable: + if activation_type is MLPActivationType.SWIGLU7: + return swiglu7 + raise ValueError(f"Unknown activation type: {activation_type}") + + +# ============================================================ +# _merge_schemas +# ============================================================ + + +@dataclass +class FFAHandler: + q_ranges: torch.Tensor + k_ranges: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + attn_type_map: torch.Tensor + softmax_scale: Optional[float] = None + auto_range_merge: bool = False + sparse_load: bool = False + + +@dataclass +class MoEConfig: + num_experts: int = 1 + num_shared_experts: int = 0 + type: Literal["token_choice", "expert_choice", "video_only_expert_choice"] = ( + "token_choice" + ) + score_func: Literal["softmax", "sigmoid"] = "softmax" + route_norm: bool = False + route_scale: float = 1.0 + top_k: Optional[int] = None # for token_choice + capacity_factor: Optional[int] = ( + None # for expert_choice and video_only_expert_choice + ) + moe_layers: list[int] = field(default_factory=list) + with_zero_experts: bool = True + + def __post_init__(self): + if self.top_k is not None and self.capacity_factor is not None: + raise ValueError("top_k and capacity_factor cannot be set at the same time") + + +# ============================================================ +# _merge_modality_dispatcher +# ============================================================ + + +class ModalityDispatcher: + permuted_modality_mapping: torch.Tensor + group_size: torch.Tensor + num_modalities: int + + @torch._dynamo.disable + def __init__(self, modality_mapping: torch.Tensor, num_modalities: int): + """ + Initialize the dispatcher. + Runs once at creation and precomputes all required mappings. + """ + self.modality_mapping = modality_mapping + self.num_modalities = num_modalities + + self.permuted_modality_mapping = self._precompute_permute_mapping( + modality_mapping + ) + + self.group_size = torch.bincount( + self.permuted_modality_mapping, minlength=num_modalities + ).to(torch.int32) + group_size_cpu = [int(x) for x in self.group_size.to("cpu").tolist()] + + # Carrier tensor: shape encodes per-modality token counts. Only the + # shape is ever read, so it lives on meta and owns no storage -- + # a real tensor would allocate num_video * num_text elements. + # tolist() above triggers a graph break when called inside + # @torch.compile, so this code always runs in eager. + # mark_unbacked prevents Dynamo from unifying symbols when some + # modalities have 0 tokens (e.g. total == video when text=0). + self._size_carrier = torch.empty(group_size_cpu, device="meta") + if not torch.compiler.is_compiling(): + for i in range(num_modalities): + torch._dynamo.decorators.mark_unbacked(self._size_carrier, i) + + @property + def group_size_cpu(self) -> list[int]: + return [self._size_carrier.shape[i] for i in range(self.num_modalities)] + + def _precompute_permute_mapping(self, modality_mapping): + """Cache permutations that group tokens by modality.""" + # 1. Compute forward and inverse permutation mappings + # argsort is efficient O(N log N) + self.permute_mapping = torch.argsort(modality_mapping) + self.inv_permute_mapping = torch.argsort(self.permute_mapping) + + # 2. Compute group sizes + # bincount is very efficient counting + permuted_modality_mapping = modality_mapping[self.permute_mapping] + + return permuted_modality_mapping + + def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]: + grouped_tensors = torch.split(x, self.group_size_cpu, dim=0) + return list(grouped_tensors) + + def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor: + return torch.cat(processed_groups, dim=0) + + @staticmethod + def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor: + """Apply forward permutation to tensor.""" + return x[permute_mapping] + + @staticmethod + def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor: + """Apply inverse permutation to tensor.""" + return x[inv_permute_mapping] + + +# ============================================================ +# _merge__common +# ============================================================ + +DYNAMIC_LORA_RANK = 0 + + +def register_lora_module(cls_or_module): + return cls_or_module + + +class MultiModalityRMSNorm(nn.Module): + __constants__ = ["dim", "eps", "num_modality"] + dim: int + eps: float + num_modality: int + + def __init__( + self, + dim: int, + eps: float = 1e-6, + device: torch.device | None = None, + num_modality: int = 1, + ): + """Initialize modality-specific root-mean-square scaling weights.""" + super().__init__() + self.dim = dim + self.eps = eps + self.num_modality = num_modality + + self.weight = torch.nn.Parameter( + torch.zeros(dim * num_modality, device=device, dtype=torch.float32) + ) + if num_modality > 1: + self.forward = self.forward_multi_experts + else: + self.forward = self.forward_single_expert + + self.reset_parameters() + + def reset_parameters(self): + nn.init.zeros_(self.weight) + + def rms(self, x: torch.Tensor) -> torch.Tensor: + t = x.float() + t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) + return t + + def forward_multi_experts( + self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher + ) -> torch.Tensor: + """Normalize tokens and apply the scale for each modality group.""" + original_dtype = x.dtype + t = self.rms(x) + + weight_chunked = self.weight.chunk(self.num_modality, dim=0) + t_list = modality_dispatcher.dispatch(t) + for i in range(self.num_modality): + t_list[i] = t_list[i] * (weight_chunked[i] + 1) + t = modality_dispatcher.undispatch(*t_list) + + return t.to(original_dtype) + + def forward_single_expert( + self, x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None + ) -> torch.Tensor: + t, original_dtype = x.float(), x.dtype + t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) + return (t * (self.weight + 1)).to(original_dtype) + + +def _bf16_compute_linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + output_dtype: Optional[torch.dtype], + compute_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Linear forward in bf16 compute precision, cast output to output_dtype.""" + input_cast = input.to(compute_dtype) + weight_cast = weight.to(compute_dtype) + output = torch.matmul(input_cast, weight_cast.t()) + if bias is not None: + output = output + bias.to(compute_dtype) + return output.to(output_dtype) + + +@register_lora_module +class BaseLinear(nn.Module): + __constants__ = ["in_features", "out_features", "num_layers", "num_experts"] + in_features: int + out_features: int + num_layers_for_initialization: int + num_experts: int + weight: Tensor + + def __init__( + self, + in_features, + out_features, + num_layers_for_initialization, + num_experts, + bias=True, + device=None, + dtype=None, + ): + """Initialize stacked projection weights for one or more experts.""" + super().__init__() + factory_kwargs = {"device": device, "dtype": torch.bfloat16} + self.in_features = in_features + self.out_features = out_features + self.num_layers_for_initialization = num_layers_for_initialization + self.num_experts = num_experts + self.use_bias = bias + self.weight = Parameter( + torch.empty((out_features * num_experts, in_features), **factory_kwargs) + ) + if bias: + self.bias = Parameter( + torch.empty(out_features * num_experts, **factory_kwargs) + ) + else: + self.register_parameter("bias", None) + + def forward( + self, + input: torch.Tensor, + output_dtype: Optional[torch.dtype] = None, + modality_dispatcher: Optional[ModalityDispatcher] = None, + ) -> torch.Tensor: + output_dtype = input.dtype if output_dtype is None else output_dtype + return _bf16_compute_linear( + input, self.weight, self.bias, output_dtype, torch.bfloat16 + ) + + # ── LoRA Bypass ── + + def lora_layout(self) -> Dict[str, Tuple]: + return { + "lora_A": (DYNAMIC_LORA_RANK, self.in_features), + "lora_B": (self.out_features, DYNAMIC_LORA_RANK), + } + + def lora_forward(self, input, *args, **kwargs): + out = self.orig_forward(input, *args, **kwargs) + x = input.to(self.lora_A.dtype) + delta = torch.matmul( + torch.matmul(x, self.lora_A.T) * self.lora_scale, self.lora_B.T + ) + return out + delta.to(out.dtype) + + +@register_lora_module +class NativeMoELinear(BaseLinear): + def forward( + self, + input: torch.Tensor, + output_dtype: Optional[torch.dtype] = None, + modality_dispatcher: Optional[ModalityDispatcher] = None, + ) -> torch.Tensor: + """Apply the matching linear expert to each dispatched modality group.""" + output_dtype = input.dtype if output_dtype is None else output_dtype + + input_list = modality_dispatcher.dispatch(input) # type: ignore + weight_chunked = self.weight.chunk(self.num_experts, dim=0) + + if self.bias is not None: + bias_chunked = self.bias.chunk(self.num_experts, dim=0) + + for i in range(self.num_experts): + input_list[i] = _bf16_compute_linear( + input_list[i], + weight_chunked[i], + bias_chunked[i] if self.bias is not None else None, + output_dtype, + torch.bfloat16, + ) + return modality_dispatcher.undispatch(*input_list) # type: ignore + + # ── LoRA Bypass (overrides BaseLinear: 3-D per-expert) ── + + def lora_layout(self) -> Dict[str, Tuple]: + return { + "lora_A": (self.num_experts, DYNAMIC_LORA_RANK, self.in_features), + "lora_B": (self.num_experts, self.out_features, DYNAMIC_LORA_RANK), + } + + def lora_forward(self, input, output_dtype=None, modality_dispatcher=None): + """Add per-expert low-rank updates to the mixture-of-experts output.""" + out = self.orig_forward( + input, output_dtype=output_dtype, modality_dispatcher=modality_dispatcher + ) + num_experts = self.num_experts + input_list = modality_dispatcher.dispatch(input) + deltas = [] + for i in range(num_experts): + x_i = input_list[i].to(self.lora_A.dtype) + delta_i = torch.matmul( + torch.matmul(x_i, self.lora_A[i].T) * self.lora_scale, self.lora_B[i].T + ) + deltas.append(delta_i.to(out.dtype)) + return out + modality_dispatcher.undispatch(*deltas) + + + + + +def create_linear( + in_features, + out_features, + num_layers=1, + num_experts=1, + bias=True, + device=None, + dtype=None, +) -> BaseLinear | NativeMoELinear: + """Construct a dense or modality-expert linear projection.""" + if num_experts == 1: + return BaseLinear( + in_features, out_features, num_layers, num_experts, bias, device, dtype + ) + return NativeMoELinear( + in_features, out_features, num_layers, num_experts, bias, device, dtype + ) + + +# ============================================================ +# _merge_element_wise_rope +# ============================================================ + + +class ElementWiseFourierEmbed(nn.Module): + def __init__( + self, + dim: int, + max_res: int = 224, + temperature: float = 10000.0, + in_pixels: bool = True, + linear_bands: bool = False, + learnable: bool = False, + device: torch.device = torch.device("cpu"), + dtype: torch.dtype = torch.float32, + ): + """Configure frequency bands for coordinate-wise Fourier embeddings. + + Args: + dim: Output feature dimension, total channels, must be divisible by 6 + max_res: Max pixel frequency resolution for generating pixel-domain frequency bands + temperature: Temperature parameter for inverse frequency mode + in_pixels: True -> pixel frequency bands, False -> inverse frequency bands + linear_bands: Whether pixel frequency bands are linearly spaced + learnable: Whether to make frequency bands learnable parameters + """ + super().__init__() + self.dim = dim + self.in_pixels = in_pixels + self.learnable = learnable + self.temperature = temperature + self.max_res = max_res + self.linear_bands = linear_bands + self.device = device + self.dtype = dtype + # Make frequency bands learnable parameters or register as buffer + bands = self.get_default_bands() + if self.learnable: + self.bands = nn.Parameter(bands) + else: + self.register_buffer("bands", bands) + + def forward(self, coords: torch.Tensor) -> torch.Tensor: + """Encode token coordinates as element-wise Fourier features. + + Args: + coords: [L,9], column order is (time, row, col, T, H, W, ref_T, ref_H, ref_W) + Returns: + emb: [L, dim] element-wise Fourier embedding + """ + # Use slicing instead of unbind + stack to reduce intermediate tensors + coords_xyz = coords[:, :3] # [L,3] -> (t, h, w) + sizes = coords[:, 3:6] # [L,3] -> (T, H, W) + refs = coords[:, 6:9] # [L,3] -> (ref_T, ref_H, ref_W) + + # Compute scaling factors + scales = (refs - 1) / (sizes - 1) # [L,3] + + # NOTE: For points where both ref and size are 1, scale factor is 1; otherwise error + scales[(refs == 1) & (sizes == 1)] = 1 + assert not scales.isnan().any(), "scales has nan" + assert not scales.isinf().any(), "scales has inf" + + # Center-align, only for h,w dimensions, not for t dimension + centers = (sizes - 1) / 2 # [L,3] + centers[:, 0] = 0 # No centering for temporal dimension + coords_xyz = coords_xyz - centers # [L,3] + + # Project to frequency bands in one pass: [L,3,B] + proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands + + # Compute sin & cos and concatenate + sin_proj = proj.sin() # [L,3,B] + cos_proj = proj.cos() + + return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) + + def reset_parameters(self): + bands = self.get_default_bands() + self.bands.copy_(bands) + + def get_default_bands(self): + if self.in_pixels: + raise NotImplementedError("in_pixels are not implemented yet") + else: + bands = freq_bands( + self.dim // 8, temperature=self.temperature, step=1, device=self.device + ).to(self.dtype) + return bands + + +def get_coords( + shape: list[int], + ref_feat_shape: list[int], + device: torch.device = torch.device("cpu"), + dtype: torch.dtype = torch.float32, +): + """ + Generate feature grid coordinates with original and reference size info + Args: + feat_shape: [T, H, W] original feature map shape + ref_feat_shape: [T_ref, H_ref, W_ref] reference feature map shape + device: Device for coordinate tensors + Returns: + coords: Tensor of shape (T*H*W, 9), containing (t, h, w, T, H, W, ref_T, ref_H, ref_W) + """ + ori_t, ori_h, ori_w = shape + ref_t, ref_h, ref_w = ref_feat_shape + + # Generate index ranges + time_rng = torch.arange(ori_t, device=device, dtype=dtype) + height_rng = torch.arange(ori_h, device=device, dtype=dtype) + width_rng = torch.arange(ori_w, device=device, dtype=dtype) + + # Create 3D meshgrid (T, H, W) + time_grid, height_grid, width_grid = torch.meshgrid( + time_rng, height_rng, width_rng, indexing="ij" + ) + + # Stack and flatten + coords_grid = torch.stack([time_grid, height_grid, width_grid], dim=-1) + coords_flat = coords_grid.reshape(-1, 3) + + # Construct and expand size metadata + meta = torch.tensor( + [ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], device=device, dtype=dtype + ) + meta_expanded = meta.expand(coords_flat.size(0), -1) + + # Concatenate and return + return torch.cat([coords_flat, meta_expanded], dim=-1) + + +# ============================================================ +# _merge_nd_rotary_pos_embedding +# ============================================================ + + +def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]: + """generate N-D grid in dimension order. + + The ndgrid function is like meshgrid except that the order of the first two input arguments are switched. + + That is, the statement + [X1,X2,X3] = ndgrid(x1,x2,x3) + + produces the same result as + + [X2,X1,X3] = meshgrid(x2,x1,x3) + + This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make + torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy'). + + """ + try: + return torch.meshgrid(*tensors, indexing="ij") + except TypeError: + # old PyTorch < 1.10 will follow this path as it does not have indexing arg, + # the old behaviour of meshgrid was 'ij' + return torch.meshgrid(*tensors) + + +def meshgrid(*tensors) -> Tuple[torch.Tensor, ...]: + """generate N-D grid in spatial dim order. + + The meshgrid function is similar to ndgrid except that the order of the + first two input and output arguments is switched. + + That is, the statement + + [X,Y,Z] = meshgrid(x,y,z) + produces the same result as + + [Y,X,Z] = ndgrid(y,x,z) + Because of this, meshgrid is better suited to problems in two- or three-dimensional Cartesian space, + while ndgrid is better suited to multidimensional problems that aren't spatially based. + """ + + # NOTE: this will throw in PyTorch < 1.10 as meshgrid did not support indexing arg or have + # capability of generating grid in xy order before then. + return torch.meshgrid(*tensors, indexing="xy") + + +def pixel_freq_bands( + num_bands: int, + max_freq: float = 224.0, + linear_bands: bool = True, + device: Optional[torch.device] = None, +): + """Generate pixel-domain Fourier frequency bands scaled by pi.""" + if linear_bands: + bands = torch.linspace( + 1.0, max_freq / 2, num_bands, dtype=torch.float32, device=device + ) + else: + bands = 2 ** torch.linspace( + 0, math.log(max_freq, 2) - 1, num_bands, dtype=torch.float32, device=device + ) + return bands * torch.pi + + +def freq_bands( + num_bands: int, + temperature: float = 10000.0, + step: int = 2, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Generate inverse-frequency bands with exponential spacing.""" + exp = ( + torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to( + torch.float32 + ) + / num_bands + ) + bands = 1.0 / (temperature**exp) + return bands + + +def build_sincos2d_pos_embed( + feat_shape: List[int], + dim: int = 64, + temperature: float = 10000.0, + reverse_coord: bool = False, + interleave_sin_cos: bool = False, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Build a two-dimensional sinusoidal embedding for a feature grid. + + Args: + feat_shape: + dim: + temperature: + reverse_coord: stack grid order W, H instead of H, W + interleave_sin_cos: sin, cos, sin, cos stack instead of sin, sin, cos, cos + dtype: + device: + + Returns: + + """ + assert dim % 4 == 0, ( + "Embed dimension must be divisible by 4 for sin-cos 2D position embedding" + ) + pos_dim = dim // 4 + bands = freq_bands(pos_dim, temperature=temperature, step=1, device=device) + + if reverse_coord: + feat_shape = feat_shape[::-1] # stack W, H instead of H, W + grid = ( + torch.stack( + ndgrid( + [ + torch.arange(s, device=device, dtype=torch.int64).to(torch.float32) + for s in feat_shape + ] + ) + ) + .flatten(1) + .transpose(0, 1) + ) + pos2 = grid.unsqueeze(-1) * bands.unsqueeze(0) + # FIXME add support for unflattened spatial dim? + + stack_dim = ( + 2 if interleave_sin_cos else 1 + ) # stack sin, cos, sin, cos instead of sin sin cos cos + pos_emb = torch.stack([torch.sin(pos2), torch.cos(pos2)], dim=stack_dim).flatten(1) + return pos_emb.to(dtype=dtype) + + +def build_fourier_pos_embed( + feat_shape: List[int], + bands: Optional[torch.Tensor] = None, + num_bands: int = 64, + max_res: int = 224, + temperature: float = 10000.0, + linear_bands: bool = False, + include_grid: bool = False, + in_pixels: bool = True, + ref_feat_shape: Optional[List[float]] = None, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, +) -> List[torch.Tensor]: + """Build sine and cosine Fourier features for a feature grid. + + Args: + feat_shape: Feature shape for embedding. + bands: Pre-calculated frequency bands. + num_bands: Number of frequency bands (determines output dim). + max_res: Maximum resolution for pixel based freq. + temperature: Temperature for non-pixel freq. + linear_bands: Linear band spacing for pixel based freq. + include_grid: Include the spatial grid in output. + in_pixels: Output in pixel freq. + ref_feat_shape: Reference feature shape for resize / fine-tune. + dtype: Output dtype. + device: Output device. + + Returns: + + """ + if bands is None: + if in_pixels: + bands = pixel_freq_bands( + num_bands, float(max_res), linear_bands=linear_bands, device=device + ) + else: + bands = freq_bands( + num_bands, temperature=temperature, step=1, device=device + ) + else: + if device is None: + device = bands.device + if dtype is None: + dtype = bands.dtype + + if in_pixels: + t = [ + torch.linspace(-1.0, 1.0, steps=s, device=device, dtype=torch.float32) + for s in feat_shape + ] + else: + t = [ + torch.arange(s, device=device, dtype=torch.int64).to(torch.float32) + for s in feat_shape + ] + # align spatial center (H/2,W/2) to (0,0) + t[1] = t[1] - (feat_shape[1] - 1) / 2 + t[2] = t[2] - (feat_shape[2] - 1) / 2 + if ref_feat_shape is not None: + # eva's scheme for resizing rope embeddings (ref shape = pretrain) + # aligning to the endpoint e.g [0,1,2] -> [0, 0.4, 0.8, 1.2, 1.6, 2] + t_rescaled = [] + for x, f, r in zip(t, feat_shape, ref_feat_shape): + if f == 1: + assert r == 1, "ref_feat_shape must be 1 when feat_shape is 1" + t_rescaled.append(x) + else: + t_rescaled.append(x / (f - 1) * (r - 1)) + else: + t_rescaled = t + + grid = torch.stack(ndgrid(t_rescaled), dim=-1) + grid = grid.unsqueeze(-1) + pos = grid * bands + + pos_sin, pos_cos = pos.sin().to(dtype=dtype), pos.cos().to(dtype) + out = [grid, pos_sin, pos_cos] if include_grid else [pos_sin, pos_cos] + return out + + +class FourierEmbed(nn.Module): + def __init__( + self, + max_res: int = 224, + num_bands: int = 64, + concat_grid=True, + keep_spatial=False, + ): + """Initialize fixed Fourier bands for input feature grids.""" + super().__init__() + self.max_res = max_res + self.num_bands = num_bands + self.concat_grid = concat_grid + self.keep_spatial = keep_spatial + self.register_buffer( + "bands", pixel_freq_bands(max_res, num_bands), persistent=False + ) + + def forward(self, x): + """Append Fourier positional features to an image feature tensor.""" + B, C = x.shape[:2] + feat_shape = x.shape[2:] + emb = build_fourier_pos_embed( + feat_shape, + self.bands, + include_grid=self.concat_grid, + dtype=x.dtype, + device=x.device, + ) + emb = torch.cat(emb, dim=-1) + emb = emb.transpose(-1, -2).flatten(len(feat_shape)) + batch_expand = (B,) + (-1,) * (x.ndim - 1) + + # FIXME support nD + if self.keep_spatial: + x = torch.cat( + [x, emb.unsqueeze(0).expand(batch_expand).permute(0, 3, 1, 2)], dim=1 + ) + else: + x = torch.cat( + [x.permute(0, 2, 3, 1), emb.unsqueeze(0).expand(batch_expand)], dim=-1 + ) + x = x.reshape(B, feat_shape.numel(), -1) + + return x + + +def rot(x): + return torch.stack([-x[..., 1::2], x[..., ::2]], -1).reshape(x.shape) + + +def apply_rot_embed(x: torch.Tensor, sin_emb, cos_emb): + if sin_emb.ndim == 3: + return x * cos_emb.unsqueeze(1).expand_as(x) + rot(x) * sin_emb.unsqueeze( + 1 + ).expand_as(x) + return x * cos_emb + rot(x) * sin_emb + + +def apply_rot_embed_list(x: List[torch.Tensor], sin_emb, cos_emb): + if isinstance(x, torch.Tensor): + x = [x] + return [t * cos_emb + rot(t) * sin_emb for t in x] + + +def apply_rot_embed_cat(x: torch.Tensor, emb): + sin_emb, cos_emb = emb.tensor_split(2, -1) + if sin_emb.ndim == 3: + return x * cos_emb.unsqueeze(1).expand_as(x) + rot(x) * sin_emb.unsqueeze( + 1 + ).expand_as(x) + return x * cos_emb.unsqueeze(1).unsqueeze(2) + rot(x) * sin_emb.unsqueeze( + 1 + ).unsqueeze(2) + + +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange( + torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2 + ) + + +def apply_rotary_emb_torch(x, cos, sin, interleaved=False): + """Apply rotary embeddings to the rotary channels of attention heads. + + x: (batch_size, seqlen, nheads, headdim) + cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) + """ + ro_dim = cos.shape[-1] * 2 + assert ro_dim <= x.shape[-1] + cos = repeat( + cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)" + ) + sin = repeat( + sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)" + ) + return torch.cat( + [ + x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, + x[..., ro_dim:], + ], + dim=-1, + ) + + +def apply_keep_indices_nlc(x, pos_embed, keep_indices): + pos_embed = pos_embed.unsqueeze(0).expand(x.shape[0], -1, -1) + pos_embed = pos_embed.gather( + 1, keep_indices.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1]) + ) + return pos_embed + + +def build_rotary_pos_embed( + feat_shape: List[int], + bands: Optional[torch.Tensor] = None, + dim: int = 64, + max_res: int = 224, + temperature: float = 10000.0, + linear_bands: bool = False, + in_pixels: bool = True, + ref_feat_shape: Optional[List[float]] = None, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, +): + """Build split sine and cosine rotary embeddings for a feature grid. + + Args: + feat_shape: Spatial shape of the target tensor for embedding. + bands: Optional pre-generated frequency bands + dim: Output dimension of embedding tensor. + max_res: Maximum resolution for pixel mode. + temperature: Temperature (inv freq) for non-pixel mode + linear_bands: Linearly (instead of log) spaced bands for pixel mode + in_pixels: Pixel vs language (inv freq) mode. + dtype: Output dtype. + device: Output device. + + Returns: + + """ + sin_emb, cos_emb = build_fourier_pos_embed( + feat_shape, + bands=bands, + num_bands=dim // 8, + max_res=max_res, + temperature=temperature, + linear_bands=linear_bands, + in_pixels=in_pixels, + ref_feat_shape=ref_feat_shape, + device=device, + dtype=dtype, + ) + num_spatial_dim = 1 + # this would be much nicer as a .numel() call to torch.Size(), but torchscript sucks + for x in feat_shape: + num_spatial_dim *= x + + # In the original timm implementation, a repeat_interleave is performed + # But in flash attn, the repeat_interleave is done on the fly + # Here we align with flash attn implementation and skip repeat_interleave + # sin_emb = sin_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1) + # cos_emb = cos_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1) + sin_emb = sin_emb.reshape(num_spatial_dim, -1) + cos_emb = cos_emb.reshape(num_spatial_dim, -1) + return sin_emb, cos_emb + + +class RotaryEmbedding(nn.Module): + """Rotary position embedding + + NOTE: This is my initial attempt at impl rotary embedding for spatial use, it has not + been well tested, and will likely change. It will be moved to its own file. + + The following impl/resources were referenced for this impl: + * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py + * https://blog.eleuther.ai/rotary-embeddings/ + """ + + def __init__( + self, + dim, + max_res=224, + temperature=10000, + in_pixels=True, + linear_bands: bool = False, + feat_shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[float]] = None, + ): + """Initialize cached bands or shape-specific rotary embeddings.""" + super().__init__() + self.dim = dim + self.max_res = max_res + self.temperature = temperature + self.in_pixels = in_pixels + self.feat_shape = feat_shape + self.ref_feat_shape = ref_feat_shape + + if feat_shape is None: + # only cache bands + if in_pixels: + bands = pixel_freq_bands( + dim // 8, float(max_res), linear_bands=linear_bands + ) + else: + bands = freq_bands(dim // 8, temperature=temperature, step=1) + logger.debug("bands: %s", bands) + self.register_buffer("bands", bands, persistent=False) + self.pos_embed_sin = None + self.pos_embed_cos = None + else: + # cache full sin/cos embeddings if shape provided up front + emb_sin, emb_cos = build_rotary_pos_embed( + feat_shape=feat_shape, + dim=dim, + max_res=max_res, + linear_bands=linear_bands, + in_pixels=in_pixels, + ref_feat_shape=self.ref_feat_shape, + temperature=self.temperature, + ) + self.bands = None + self.register_buffer("pos_embed_sin", emb_sin, persistent=False) + self.register_buffer("pos_embed_cos", emb_cos, persistent=False) + + def get_embed(self, shape: Optional[List[int]] = None): + """Return cached rotary embeddings or build them for a feature shape.""" + if self.bands is not None: + # rebuild embeddings every call, use if target shape changes + assert shape is not None # type: ignore[unreachable] + return build_rotary_pos_embed( + shape, + self.bands, + in_pixels=self.in_pixels, + temperature=self.temperature, + ) + else: + return self.pos_embed_sin, self.pos_embed_cos + + def forward(self, x): + # assuming channel-first tensor where spatial dim are >= 2 + sin_emb, cos_emb = self.get_embed(x.shape[2:]) + return apply_rot_embed(x, sin_emb, cos_emb) + + +class RotaryEmbeddingCat(nn.Module): + """Rotary position embedding w/ concatenatd sin & cos + + The following impl/resources were referenced for this impl: + * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py + * https://blog.eleuther.ai/rotary-embeddings/ + """ + + def __init__( + self, + dim, + max_res=224, + temperature=10000, + in_pixels=True, + linear_bands: bool = False, + feat_shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[float]] = None, + ): + """Initialize optional shape-specific concatenated rotary embeddings.""" + super().__init__() + self.dim = dim + self.max_res = max_res + self.temperature = temperature + self.in_pixels = in_pixels + self.linear_bands = linear_bands + self.feat_shape = feat_shape + self.ref_feat_shape = ref_feat_shape + + if feat_shape is None: + self.bands = None + self.pos_embed = None + else: + # cache full sin/cos embeddings if shape provided up front + embeds = build_rotary_pos_embed( + feat_shape=feat_shape, + dim=dim, + max_res=max_res, + linear_bands=linear_bands, + in_pixels=in_pixels, + ref_feat_shape=self.ref_feat_shape, + temperature=self.temperature, + ) + self.bands = None + self.register_buffer("pos_embed", torch.cat(embeds, -1), persistent=False) + + def get_embed( + self, + shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[float]] = None, + ): + """Build concatenated rotary embeddings or return cached values.""" + if shape is not None: + # rebuild bands and embeddings every call, use if target shape changes + embeds = build_rotary_pos_embed( + feat_shape=shape, + dim=self.dim, + max_res=self.max_res, + linear_bands=self.linear_bands, + in_pixels=self.in_pixels, + ref_feat_shape=ref_feat_shape + if ref_feat_shape + else self.ref_feat_shape, + temperature=self.temperature, + device=torch.cuda.current_device(), + ) + return torch.cat(embeds, -1) + elif self.pos_embed is not None: + return self.pos_embed # type: ignore[unreachable] + else: + assert False, ( + "get_embed() requires pre-computed pos_embed or valid shape w/ pre-computed bands" + ) + + def forward(self, x): + # assuming channel-first tensor where spatial dim are >= 2 + pos_embed = self.get_embed(x.shape[2:]) + return apply_rot_embed_cat(x, pos_embed) + + +class LearnableRotaryEmbeddingCat(nn.Module): + """Rotary position embedding w/ concatenatd sin & cos + + The following impl/resources were referenced for this impl: + * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py + * https://blog.eleuther.ai/rotary-embeddings/ + """ + + def __init__( + self, + dim, + max_res=224, + temperature=10000, + in_pixels=True, + linear_bands: bool = False, + feat_shape: Optional[List[int]] = None, + ref_feat_shape: Optional[List[int]] = None, + ): + """Initialize learnable rotary frequency bands for spatial features.""" + super().__init__() + self.dim = dim + self.max_res = max_res + self.temperature = temperature + self.in_pixels = in_pixels + self.linear_bands = linear_bands + self.feat_shape = feat_shape + self.ref_feat_shape = ref_feat_shape + self.bands = nn.Parameter(self.get_default_bands()) + setattr(self.bands, "sequence_parallel", True) + + def get_default_bands(self): + """Create initial bands for the configured coordinate domain.""" + if self.in_pixels: + bands = pixel_freq_bands( + self.dim // 8, + float(self.max_res), + linear_bands=self.linear_bands, + device=torch.cuda.current_device(), + ) + else: + bands = freq_bands( + self.dim // 8, + temperature=self.temperature, + step=1, + device=torch.cuda.current_device(), + ) + return bands + + def get_embed( + self, shape: Optional[List[int]], ref_feat_shape: Optional[List[float]] = None + ): + """Build concatenated rotary embeddings from learned frequency bands.""" + # rebuild bands and embeddings every call, use if target shape changes + embeds = build_rotary_pos_embed( + feat_shape=shape, # type: ignore[arg-type] + bands=self.bands, # use learned bands + dim=self.dim, + max_res=self.max_res, + linear_bands=self.linear_bands, + in_pixels=self.in_pixels, + ref_feat_shape=ref_feat_shape if ref_feat_shape else self.ref_feat_shape, # type: ignore[arg-type] + temperature=self.temperature, + device=torch.cuda.current_device(), + ) + return torch.cat(embeds, -1) + + +# ============================================================ +# _merge_custom_ops +# ============================================================ + +HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None +HAS_FA3 = importlib.util.find_spec("flash_attn_interface") is not None + + +def is_hopper_arch() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major >= 9 + + +@magi_register_custom_op( + "magi2::flash_attn_func", infer_output_meta_fn=["query"], is_subgraph_boundary=True +) +def flash_attn_func( + query: torch.Tensor, key: torch.Tensor, value: torch.Tensor +) -> torch.Tensor: + """Select and run the FlashAttention kernel that matches GPU support.""" + if HAS_FA3 and is_hopper_arch(): + from flash_attn_interface import flash_attn_func as fa3_flash_attn_func + + return fa3_flash_attn_func(query, key, value) + else: + from flash_attn.flash_attn_interface import ( + flash_attn_func as fa2_flash_attn_func, + ) + + return fa2_flash_attn_func(query, key, value) + + +def _split_q_range_with_no_overlap( + q_ranges: torch.Tensor, k_ranges: torch.Tensor +) -> Tuple[List[List[int]], List[List[List[int]]]]: + """Split overlapping query spans and retain their associated key ranges.""" + range_boundary = torch.unique(q_ranges, sorted=True).tolist() + candidates = [ + [start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:]) + ] + q_ranges = q_ranges.tolist() + k_ranges = k_ranges.tolist() + for q_range, k_range in zip(q_ranges, k_ranges): + q_start, q_end = q_range + for q_range_cand in candidates: + if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end: + q_range_cand[2].append(k_range) + q_ranges_out = [] + k_ranges_out = [] + for q_range_cand in candidates: + if len(q_range_cand[2]) > 0: + q_ranges_out.append(q_range_cand[0:2]) + k_ranges_out.append(q_range_cand[2]) + return q_ranges_out, k_ranges_out + + +def _flash_attn_with_correction( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q_ranges: List[List[int]], + k_range_list: List[List[List[int]]], +): + """Combine disjoint key-range attention with log-sum-exp correction.""" + output = torch.zeros_like(query) + output_lse = torch.zeros( + (query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device + ) + + from flash_attn.flash_attn_interface import flash_attn_func + + for q_range, k_ranges in zip(q_ranges, k_range_list): + q_start, q_end = q_range + qo_out, qo_lse = None, None + for k_range in k_ranges: + k_start, k_end = k_range + cur_qo_out, cur_qo_lse, _ = flash_attn_func( + query[q_start:q_end].unsqueeze(0), + key[k_start:k_end].unsqueeze(0), + value[k_start:k_end].unsqueeze(0), + return_attn_probs=True, + ) + cur_qo_out, cur_qo_lse = cur_qo_out.squeeze(0), cur_qo_lse.squeeze(0) + + if qo_out is None: + qo_out = cur_qo_out + qo_lse = cur_qo_lse + else: + qo_lse[qo_lse == torch.inf] = -torch.inf + cur_qo_lse[cur_qo_lse == torch.inf] = -torch.inf + max_lse = torch.max(qo_lse, cur_qo_lse) + qo_se, cur_qo_se = ( + torch.exp(qo_lse - max_lse), + torch.exp(cur_qo_lse - max_lse), + ) + sum_se = qo_se + cur_qo_se + qo_scale, cur_qo_scale = qo_se / sum_se, cur_qo_se / sum_se + + qo_out = qo_out * qo_scale.permute(1, 0).unsqueeze( + -1 + ) + cur_qo_out * cur_qo_scale.permute(1, 0).unsqueeze(-1) + qo_lse = torch.log(sum_se) + max_lse + + output[q_start:q_end] = qo_out + output_lse[q_start:q_end, :] = qo_lse.permute(1, 0) + return output, output_lse + + +def _flex_flash_attn_func_meta( + query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *args, **kwargs +) -> Tuple[torch.Tensor, torch.Tensor]: + output = torch.empty_like(query) + output_lse = torch.empty( + (query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device + ) + return output, output_lse + + +_PARAM_RENAMES = {"auto_range_merge": "range_merge", "sparse_load": "block_sparse"} + + +def _adapt_magi_attn_kwargs(fn, kwargs: dict) -> dict: + """Translate legacy kwarg names to current magi_attention API and vice versa.""" + import inspect + + params = set(inspect.signature(fn).parameters) + out = {} + for key, value in kwargs.items(): + if key in params: + out[key] = value + elif key in _PARAM_RENAMES and _PARAM_RENAMES[key] in params: + out[_PARAM_RENAMES[key]] = value + return out + +@magi_register_custom_op( + "magi2::flex_flash_attn_func", + infer_output_meta_fn=_flex_flash_attn_func_meta, + is_subgraph_boundary=True, +) +def flex_flash_attn_func( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q_ranges: torch.Tensor, + k_ranges: torch.Tensor, + attn_type_map: torch.Tensor | None = None, + max_seqlen_q: int | None = None, + auto_range_merge: bool = False, + sparse_load: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Run range-restricted attention through the available kernel.""" + if HAS_MAGI_ATTENTION and is_hopper_arch(): + from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func + + kwargs = _adapt_magi_attn_kwargs( + magi_flex_flash_attn_func, + {"max_seqlen_q": max_seqlen_q, "auto_range_merge": auto_range_merge, "sparse_load": sparse_load}, + ) + output, meta = magi_flex_flash_attn_func( + query, key, value, q_ranges, k_ranges, attn_type_map, **kwargs + ) + return output, meta.lse + return _custom_flex_flash_attn_func(query, key, value, q_ranges, k_ranges) + + +def _custom_flex_flash_attn_func( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q_ranges: torch.Tensor, + k_ranges: torch.Tensor, + **kwargs, +): + q_ranges, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges) + return _flash_attn_with_correction(query, key, value, q_ranges, k_range_list) + + +# ============================================================ +# _merge_local_attn +# ============================================================ + + +# num_video_tokens = 26000 +# num_audio_and_txt_tokens = 101+500 +# num_frames = 26 +# frame_base_recept_field = 2 +def calc_local_qk_range( + num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field +): + """Build local-video and global cross-modality query and key ranges.""" + token_per_frame = num_video_tokens // num_frames + total_tokens = num_video_tokens + num_audio_and_txt_tokens + + q_range_list = [] + k_range_list = [] + + for i in range(num_frames): + local_q_range = torch.tensor([i * token_per_frame, (i + 1) * token_per_frame]) + local_k_range = torch.tensor( + [ + (i - frame_receptive_field) * token_per_frame, + (i + frame_receptive_field + 1) * token_per_frame, + ] + ) + + q_range_list.append(local_q_range) + k_range_list.append(local_k_range) + local_q_range = torch.stack(q_range_list, dim=0) + local_k_range = torch.stack(k_range_list, dim=0) + + local_k_range[local_k_range < 0] = 0 + local_k_range[local_k_range > num_video_tokens] = num_video_tokens + + video_q_range = torch.tensor([[0, num_video_tokens]]) + video_k_range = torch.tensor( + [[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]] + ) + + at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]]) + at_k_ranges = torch.tensor([[0, total_tokens]]) + + q_ranges = ( + torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0) + .to(torch.int32) + .to("cuda", non_blocking=True) + ) + k_ranges = ( + torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0) + .to(torch.int32) + .to("cuda", non_blocking=True) + ) + + return (q_ranges, k_ranges) + + +def calc_local_attn_ffa_handler( + num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field +): + """Create the attention descriptor for local video and cross-modality ranges.""" + q_ranges, k_ranges = calc_local_qk_range( + num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field + ) + max_seqlen_q = num_video_tokens + num_audio_and_txt_tokens + max_seqlen_k = num_video_tokens + num_audio_and_txt_tokens + attn_type_map = torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32) + softmax_scale = None + + ffa_handler = FFAHandler( + q_ranges=q_ranges, + k_ranges=k_ranges, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + attn_type_map=attn_type_map, + softmax_scale=softmax_scale, + ) + return ffa_handler + + +def _env_flag(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.lower() in {"1", "true", "yes", "on"} + + +def _magi_compile_enabled(_=None) -> bool: + return not _env_flag("MAGI2_DISABLE_MAGI_COMPILE", False) and not _env_flag( + "TORCH_COMPILE_DISABLE", False + ) + + +# ============================================================ +# _merge_modeling +# ============================================================ + + +def build_dataclass_from_config(dataclass_cls, config, overrides=None): + """Build a dataclass instance from a pydantic config, applying optional overrides.""" + fields = ( + {f.name for f in dataclass_cls.__dataclass_fields__.values()} + if hasattr(dataclass_cls, "__dataclass_fields__") + else set() + ) + kwargs = {} + for field_name in fields: + if overrides and field_name in overrides: + kwargs[field_name] = overrides[field_name] + elif hasattr(config, field_name): + kwargs[field_name] = getattr(config, field_name) + return dataclass_cls(**kwargs) + + +def register_model(name, aliases=None): + """No-op decorator for inference (model registry not needed).""" + + def decorator(cls): + return cls + + return decorator + + +@dataclass +class VarlenHandler(object): + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + + +def _flash_attn_with_cp_meta(q: torch.Tensor, *args, **kwargs) -> torch.Tensor: + return torch.empty_like(q, dtype=torch.bfloat16).squeeze(0) + + +@magi_register_custom_op( + "magi2::flash_attn_with_cp", + infer_output_meta_fn=_flash_attn_with_cp_meta, + is_subgraph_boundary=True, +) +def flash_attn_with_cp( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cp_split_sizes: List[int] +) -> torch.Tensor: + """Run dense attention with context-parallel head and sequence exchange. + + Args: + q: bsz, seq_len, n_heads, head_dim + k: bsz, seq_len, n_heads, head_dim + v: bsz, seq_len, n_heads, head_dim + rope: seq_len, rotary_dim + + Returns: + self_attn_out: bsz, seq_len, n_heads, head_dim + """ + q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16) + + if psm.get_world_size("cp") > 1: + q, k, v = batch_scatter_head_gather_seqlen( + [q.squeeze(0), k.squeeze(0), v.squeeze(0)], + cp_split_sizes, + psm.get_parallel_group("cp"), + ) + q = q.unsqueeze(0) + k = k.unsqueeze(0) + v = v.unsqueeze(0) + + self_attn_out = torch.ops.magi2.flash_attn_func(q, k, v).squeeze(0) + + if psm.get_world_size("cp") > 1: + self_attn_out = scatter_seqlen_gather_head( + self_attn_out, cp_split_sizes, psm.get_parallel_group("cp"), async_op=False + ) + self_attn_out = rearrange( + self_attn_out, "(cp sq) hn hd -> sq (cp hn) hd", cp=psm.get_world_size("cp") + ) + + return self_attn_out + + +@magi_register_custom_op( + "magi2::flex_flash_attn_with_cp", + infer_output_meta_fn=_flash_attn_with_cp_meta, + is_subgraph_boundary=True, +) +def flex_flash_attn_with_cp( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_ranges: torch.Tensor, + k_ranges: torch.Tensor, + attn_type_map: torch.Tensor, + max_seqlen_q: int, + auto_range_merge: bool, + sparse_load: bool, + cp_split_sizes: List[int], +) -> torch.Tensor: + """Run range-restricted attention with context-parallel token exchange. + + Args: + q: bsz, seq_len, n_heads, head_dim + k: bsz, seq_len, n_heads, head_dim + v: bsz, seq_len, n_heads, head_dim + q_ranges: query ranges for flex attention + k_ranges: key ranges for flex attention + attn_type_map: attention type for each q/k range + max_seqlen_q: maximum query range length for tile scheduling + auto_range_merge: merge repeated query ranges before launching magi attention + sparse_load: enable sparse-load optimization in magi attention + + Returns: + self_attn_out: bsz, seq_len, n_heads, head_dim + """ + q, k, v = ( + q.to(torch.bfloat16).squeeze(0), + k.to(torch.bfloat16).squeeze(0), + v.to(torch.bfloat16).squeeze(0), + ) + + if psm.get_world_size("cp") > 1: + q, k, v = batch_scatter_head_gather_seqlen( + [q, k, v], cp_split_sizes, psm.get_parallel_group("cp") + ) + + out, _ = torch.ops.magi2.flex_flash_attn_func( + q, + k, + v, + q_ranges, + k_ranges, + attn_type_map, + max_seqlen_q, + auto_range_merge, + ) + + if psm.get_world_size("cp") > 1: + out = scatter_seqlen_gather_head( + out, cp_split_sizes, psm.get_parallel_group("cp"), async_op=False + ) + out = rearrange( + out, "(cp sq) hn hd -> sq (cp hn) hd", cp=psm.get_world_size("cp") + ) + + return out + + +# Define the Attention module +@dataclass +class AttentionConfig: + """ + Args: + hidden_size (int): Hidden dimension size of the attention module. + num_heads_q (int): Number of query heads. + num_heads_kv (int): Number of key-value heads. + head_dim (int): Dimension size of each head. + params_dtype (torch.dtype): Dtype of the parameters. + + """ + + hidden_size: int + num_heads_q: int + num_heads_kv: int + head_dim: int + params_dtype: torch.dtype + checkpoint_qk_layernorm_rope: bool + num_modality: int + num_layers: int + use_local_attn: bool = False + enable_attn_gating: bool = False + + +class Attention(torch.nn.Module): + """ + Multi-head attention module with support for both self-attention and cross-attention. + """ + + config: AttentionConfig + + def __init__(self, config: AttentionConfig): + """Initialize multimodal attention normalization and projections.""" + super().__init__() + self.config = config + + self.pre_norm = MultiModalityRMSNorm( + config.hidden_size, eps=1e-6, num_modality=config.num_modality + ) + + # Linear projection for query, key, value for self-attention + self.linear_qkv = create_linear( + config.hidden_size, + config.num_heads_q * config.head_dim + + config.num_heads_kv * config.head_dim * 2, + num_experts=config.num_modality, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + ) + if config.enable_attn_gating: + self.linear_g = create_linear( + config.hidden_size, + config.num_heads_q, + num_experts=config.num_modality, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + ) + + # Output projection for self-attention outputs + self.linear_proj = create_linear( + config.num_heads_q * config.head_dim, + config.hidden_size, + bias=False, + num_experts=config.num_modality, + dtype=config.params_dtype, + num_layers=config.num_layers, + ) + + # Layer normalization for query and key in self-attention + self.q_norm = MultiModalityRMSNorm( + config.head_dim, num_modality=config.num_modality + ) + self.k_norm = MultiModalityRMSNorm( + config.head_dim, num_modality=config.num_modality + ) + + self.q_size = config.num_heads_q * config.head_dim + self.kv_size = config.num_heads_kv * config.head_dim + + def reset_parameters(self): + if hasattr(self.linear_proj, "reset_parameters_output_layer"): + self.linear_proj.reset_parameters_output_layer() + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler, # Key for distributed attention + modality_dispatcher: ModalityDispatcher, + cp_split_sizes: List[int], + ) -> torch.Tensor: + """ + Forward pass for the attention module. + + Args: + hidden_states(torch.Tensor): Input tensor for self-attention + rope(torch.Tensor): Rotary position embeddings + + Returns: + out(torch.Tensor): Combined output from self-attention and cross-attention + + + Shape: + - hidden_states: (num_tokens, hidden_size) + - rope: (num_tokens, embed_dim) + - out: (num_tokens, hidden_size) + + """ + # Project input to query, query_xattn, key, value for self-attention + + hidden_states = self.pre_norm( + hidden_states, modality_dispatcher=modality_dispatcher + ).to(torch.bfloat16) + qkv: torch.Tensor = self.linear_qkv( + hidden_states, modality_dispatcher=modality_dispatcher + ).to(torch.float32) + + q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=1) + q = q.view(-1, self.config.num_heads_q, self.config.head_dim) + k = k.view(-1, self.config.num_heads_kv, self.config.head_dim) + v = v.view(-1, self.config.num_heads_kv, self.config.head_dim) + if self.config.enable_attn_gating: + g = self.linear_g( + hidden_states, modality_dispatcher=modality_dispatcher + ).to(torch.float32) + # IMPORTANT: Do NOT use `g.view(k.shape[0], num_heads, -1)` here. + # ModalityDispatcher marks per-modality sizes as unbacked symbols (u0,u1,u2), + # so k.shape[0] = u0+u1+u2. Using -1 in view forces Inductor to guard on + # `u0+u1+u2 > 0`, which is forbidden for unbacked symbols and raises + # GuardOnDataDependentSymNode at compile time. + g = g.unsqueeze(-1) + + q = self.q_norm(q, modality_dispatcher=modality_dispatcher) + k = self.k_norm(k, modality_dispatcher=modality_dispatcher) + + q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0) + k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0) + v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0) + + sin_emb, cos_emb = rope.tensor_split(2, -1) + q = apply_rotary_emb(q, cos_emb, sin_emb) + k = apply_rotary_emb(k, cos_emb, sin_emb) + + if self.config.use_local_attn: + auto_range_merge = bool( + getattr(local_attn_handler, "auto_range_merge", False) + ) + sparse_load = bool(getattr(local_attn_handler, "sparse_load", False)) + self_attn_out = flex_flash_attn_with_cp( + q, + k, + v, + local_attn_handler.q_ranges, + local_attn_handler.k_ranges, + local_attn_handler.attn_type_map, + local_attn_handler.max_seqlen_q, + auto_range_merge, + sparse_load, + cp_split_sizes, + ) + else: + self_attn_out = flash_attn_with_cp(q, k, v, cp_split_sizes) + self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping) + + if self.config.enable_attn_gating: + self_attn_out = self_attn_out * torch.sigmoid(g) + + self_attn_out = self_attn_out.view( + -1, self.config.num_heads_q * self.config.head_dim + ).to(torch.bfloat16) + out = self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher) + + return out + + +@dataclass +class MLPConfig: + """Configuration for the MLP module + + Args: + hidden_size (int): Hidden dimension size. + intermediate_size (int): Intermediate Hidden dimension of the mlp layer. + activation_type (MLPActivationType): Activation function type. + params_dtype (torch.dtype): Dtype of the parameters. + """ + + hidden_size: int + intermediate_size: int + activation_type: MLPActivationType + params_dtype: torch.dtype + num_modality: int = 1 + num_layers: int = 1 + gated_act: bool = False + + +class MLP(torch.nn.Module): + """MLP module with traditional architecture (up-projection, activation, and down-projection)""" + + config: MLPConfig + + def __init__(self, config: MLPConfig): + """Initialize modality-aware gated feed-forward projections.""" + super().__init__() + + num_experts = config.num_modality + + self.pre_norm = MultiModalityRMSNorm( + config.hidden_size, num_modality=config.num_modality + ) + if config.gated_act: + intermediate_size_up = config.intermediate_size * 2 + else: + intermediate_size_up = config.intermediate_size + + # Combined projection for up-projection and gate + self.up_gate_proj = create_linear( + config.hidden_size, + intermediate_size_up, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + num_experts=num_experts, + ) + + # Down-projection back to hidden size + self.down_proj = create_linear( + config.intermediate_size, + config.hidden_size, + bias=False, + dtype=config.params_dtype, + num_layers=config.num_layers, + num_experts=num_experts, + ) + + # Initialize activation function based on configuration + self.activation_func = create_activation_func(config.activation_type) + + def forward( + self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher + ) -> torch.Tensor: + """Forward pass of the MLP module. + + Args: + x (torch.Tensor): Input tensor + + Returns: + output (torch.Tensor): Output tensor + + Shape: + - x: (num_tokens, hidden_size) + - output: (num_tokens, hidden_size) + """ + + x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) + + # Project input to up-projection and gate + x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to( + torch.float32 + ) + + x = self.activation_func(x).to(torch.bfloat16) + + # Element-wise multiplication with gate and down-projection + x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32) + return x + + def extra_repr(self) -> str: + return f"{self.up_gate_proj.weight.shape=}, {self.down_proj.weight.shape=}" + + +# Define the Adapter module +@dataclass +class AdapterConfig: + """ + Configuration for the Adapter module that handles various input embeddings and conditioning + """ + + hidden_size: int + num_attention_heads: int + + # caption + text_in_channels: int + video_in_channels: int + audio_in_channels: int + + params_dtype: torch.dtype + + +class Adapter(torch.nn.Module): + """ + Adapter module that handles input embeddings, timestep embeddings, and caption conditioning + """ + + config: AdapterConfig + + def __init__(self, config: AdapterConfig): + """Initialize modality projections and rotary coordinate embedding.""" + super().__init__() + self.config = config + + # Linear projection for input features + self.video_embedder = nn.Linear( + config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32 + ) + + # Projection for text embeddings for cross-attention + self.text_embedder = nn.Linear( + config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32 + ) + + self.audio_embedder = nn.Linear( + config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32 + ) + + # Initialize rotary position embeddings + self.rope = ElementWiseFourierEmbed( + config.hidden_size // config.num_attention_heads, + in_pixels=False, + learnable=False, + ) + + def forward( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + video_mask: torch.Tensor, + audio_mask: torch.Tensor, + text_mask: torch.Tensor, + ): + """ + Forward pass of the Adapter module + + Args: + x: (num_tokens, max(image_token_channels, txt_token_channels)) + coords_mapping: (num_tokens, ) + video_mask: (num_tokens, ) + audio_mask: (num_tokens, ) + text_mask: (num_tokens, ) + + Returns: + output_x: (num_tokens, hidden_size) + rope: (num_tokens, hidden_size // num_attention_heads) + + """ + + # Calculate rotary position embeddings + rope = self.rope(coords_mapping) + + output_x = torch.zeros( + x.shape[0], self.config.hidden_size, device=x.device, dtype=x.dtype + ) + + output_x[text_mask] = self.text_embedder( + x[text_mask, : self.config.text_in_channels], + # output_dtype=torch.float32 + ) + + # Project input features to hidden dimension + output_x[audio_mask] = self.audio_embedder( + x[audio_mask, : self.config.audio_in_channels], + # output_dtype=torch.float32 + ) + + output_x[video_mask] = self.video_embedder( + x[video_mask, : self.config.video_in_channels], + # output_dtype=torch.float32 + ) + + return output_x, rope + + +class TransformerLayer(torch.nn.Module): + """ + A single transformer layer with attention, MLP, and adaptive layer normalization + """ + + def __init__(self, config: ModelConfig, layer_idx: int): + """Initialize one refiner layer from its indexed architecture settings.""" + super().__init__() + num_modality = 3 if layer_idx in config.mm_layers else 1 + use_local_attn = layer_idx in config.local_attn_layers + self.post_norm = layer_idx in config.post_norm_layers + # Build attention and MLP modules + attention_config = build_dataclass_from_config( + AttentionConfig, + config, + {"num_modality": num_modality, "use_local_attn": use_local_attn}, + ) + self.attention: Attention = Attention(attention_config) + + if layer_idx in config.layer_activation_types: + activation_type = MLPActivationType( + config.layer_activation_types[layer_idx] + ) + else: + activation_type = MLPActivationType(config.activation_type) + if activation_type in [MLPActivationType.SWIGLU7]: + gated_act = True + intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4 + else: + gated_act = False + intermediate_size = config.hidden_size * 4 + mlp_config = build_dataclass_from_config( + MLPConfig, + config, + { + "num_modality": num_modality, + "activation_type": activation_type, + "intermediate_size": intermediate_size, + "gated_act": gated_act, + }, + ) + self.mlp: MLP = MLP(mlp_config) + if self.post_norm: + self.attn_post_norm = MultiModalityRMSNorm( + config.hidden_size, num_modality=num_modality + ) + self.mlp_post_norm = MultiModalityRMSNorm( + config.hidden_size, num_modality=num_modality + ) + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler, # Key for distributed attention + modality_dispatcher: ModalityDispatcher, + cp_split_sizes: List[int], + ) -> torch.Tensor: + """ + Forward pass of the transformer layer + + Args: + hidden_states: Input tensor + rope: Rotary position embeddings + magi_attn_key: Key for distributed self-attention computation + + Returns: + hidden_states(torch.Tensor): Processed hidden states + + Shape: + - hidden_states: (num_tokens, hidden_size) + - rope: (num_tokens, embed_dim) + """ + # Apply Attention + attn_out = self.attention( + hidden_states, + rope, + permute_mapping, + inv_permute_mapping, + varlen_handler, + local_attn_handler, + modality_dispatcher, + cp_split_sizes, + ) + if self.post_norm: + attn_out = self.attn_post_norm( + attn_out, modality_dispatcher=modality_dispatcher + ) + hidden_states = hidden_states + attn_out + + # Apply MLP + mlp_out = self.mlp(hidden_states, modality_dispatcher) + if self.post_norm: + mlp_out = self.mlp_post_norm( + mlp_out, modality_dispatcher=modality_dispatcher + ) + hidden_states = hidden_states + mlp_out + + return hidden_states + + +keep_preview_model_once = True + + +def config_patch(compile_config: CompileConfig) -> CompileConfig: + global keep_preview_model_once + if keep_preview_model_once: + keep_preview_model_once = False + else: + compile_config.offload_config.gpu_resident_weight_ratio = 0.0 + return compile_config + + +@magi_compile(config_patch=config_patch, enable_if=_magi_compile_enabled) +class TransformerBlock(torch.nn.Module): + """ + A stack of transformer layers + """ + + def __init__(self, model_config: ModelConfig): + super().__init__() + # Build all transformer layers + self.layers: list[TransformerLayer] = nn.ModuleList() + for layer_idx in range(model_config.num_layers): + self.layers.append(TransformerLayer(model_config, layer_idx)) + + def forward( + self, + x: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler, # Key for distributed attention + modality_dispatcher: ModalityDispatcher, + cp_split_sizes: List[int], + ) -> torch.Tensor: + """ + Forward pass of the transformer block + + Args: + x(torch.Tensor): Input tensor + rope(torch.Tensor): Rotary position embeddings + local_attn_handler(FFAHandler): Key for distributed attention + video_mask(torch.Tensor): Mask for video modality + audio_mask(torch.Tensor): Mask for audio modality + text_mask(torch.Tensor): Mask for text modality + modality_dispatcher(ModalityDispatcher): Modality dispatcher + Returns: + x(torch.Tensor): Processed input features + + Shape: + - x: (num_tokens, hidden_size) + - rope: (num_tokens, embed_dim) + - modality_mapping: (num_tokens, ) + - video_mask: (num_tokens, ) + - audio_mask: (num_tokens, ) + - text_mask: (num_tokens, ) + """ + + for layer in self.layers: + x = layer( + x, + rope, + permute_mapping, + inv_permute_mapping, + varlen_handler, + local_attn_handler, + modality_dispatcher, + cp_split_sizes, + ) + + return x + + +# Define the Transformer module +@dataclass +class TransformerConfig: + hidden_size: int # Size of hidden representations + video_in_channels: int # Number of input channels + audio_in_channels: int # Number of input channels + text_in_channels: int # Number of input channels + params_dtype: torch.dtype # Data type for parameters + post_process_dtype: torch.dtype # Data type for post-processing + + +@register_model("magi2_refiner", aliases=["magi2_refiner"]) +class PostAdapter(torch.nn.Module): + """Final per-modality normalization and projection back to latent channels.""" + + def __init__(self, config: TransformerConfig) -> None: + """Initialize video and audio normalization and output projections.""" + super().__init__() + self.config = config + self.final_norm_video = MultiModalityRMSNorm(config.hidden_size) + self.final_norm_audio = MultiModalityRMSNorm(config.hidden_size) + self.final_linear_video = nn.Linear( + config.hidden_size, + config.video_in_channels, + bias=False, + dtype=torch.float32, + ) + self.final_linear_audio = nn.Linear( + config.hidden_size, + config.audio_in_channels, + bias=False, + dtype=torch.float32, + ) + + def forward( + self, x: torch.Tensor, video_mask: torch.Tensor, audio_mask: torch.Tensor + ) -> torch.Tensor: + """Project normalized video and audio tokens into a latent tensor.""" + x_video = x[video_mask].to(self.final_norm_video.weight.dtype) + x_video = self.final_norm_video(x_video) + x_video = self.final_linear_video(x_video) + + x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype) + x_audio = self.final_norm_audio(x_audio) + x_audio = self.final_linear_audio(x_audio) + + x_out = torch.zeros( + x.shape[0], + max(self.config.video_in_channels, self.config.audio_in_channels), + device=x.device, + dtype=x.dtype, + ) + x_out[video_mask, : self.config.video_in_channels] = x_video + x_out[audio_mask, : self.config.audio_in_channels] = x_audio + return x_out + + +class Transformer(BaseDiT): + _config = Magi2RefinerVideoConfig() + _fsdp_shard_conditions = _config._fsdp_shard_conditions + _compile_conditions = _config._compile_conditions + _supported_attention_backends = _config._supported_attention_backends + param_names_mapping = _config.param_names_mapping + reverse_param_names_mapping = _config.reverse_param_names_mapping + lora_param_names_mapping = _config.lora_param_names_mapping + + config: TransformerConfig + + def __init__( + self, + config: Magi2RefinerVideoConfig, + hf_config: dict | None = None, + **kwargs, + ) -> None: + """Construct the production MAGI-2 refiner with official module names.""" + del kwargs + super().__init__(config=config, hf_config=hf_config or {}) + self.fastvideo_config = config + model_config = config.arch_config + model_config.params_dtype = _resolve_params_dtype(model_config.params_dtype) + self.hidden_size = model_config.hidden_size + self.num_attention_heads = model_config.num_heads_q + self.num_channels_latents = model_config.num_channels_latents + + self.config = build_dataclass_from_config( + TransformerConfig, model_config, {"post_process_dtype": torch.float32} + ) + adapter_config = build_dataclass_from_config( + AdapterConfig, + model_config, + { + "num_attention_heads": model_config.num_heads_q, + "params_dtype": torch.float32, + }, + ) + self.pre_adapter: Adapter = Adapter(adapter_config) + + self.block: TransformerBlock = TransformerBlock(model_config=model_config) + + self.post_adapter: PostAdapter = PostAdapter(self.config) + self.__post_init__() + + @torch.compile(dynamic=True, fullgraph=False) + def forward( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + modality_mapping: torch.Tensor, + varlen_handler: VarlenHandler, + local_attn_handler: FFAHandler, # Key for distributed attention + ): + """Refine packed video, audio, and text token features. + + Args: + x(torch.Tensor): Input features + coords_mapping(torch.Tensor): Mapping from tokens to coords for rope + modality_mapping(torch.Tensor): Mapping from tokens to modality + + Returns: + x(torch.Tensor): Processed input features + + """ + x = ulysses_scheduler().dispatch(x) + coords_mapping = ulysses_scheduler().dispatch(coords_mapping) + modality_mapping = ulysses_scheduler().dispatch(modality_mapping) + cp_split_sizes = ulysses_scheduler().cp_split_sizes + + modality_dispatcher = ModalityDispatcher(modality_mapping, 3) + permute_mapping, inv_permute_mapping = ( + modality_dispatcher.permute_mapping, + modality_dispatcher.inv_permute_mapping, + ) + # assert modality_dispatcher.is_sorted + video_mask = modality_mapping == Modality.VIDEO + audio_mask = modality_mapping == Modality.AUDIO + text_mask = modality_mapping == Modality.TEXT + + # Process inputs through adapter + x, rope = self.pre_adapter( + x, coords_mapping, video_mask, audio_mask, text_mask + ) + + # NOTE(LLZ):only convert x to params_dtype, rope stays in float32 + x = x.to(self.config.params_dtype) + x = ModalityDispatcher.permute(x, permute_mapping) + x = self.block( + x, + rope, + permute_mapping=permute_mapping, + inv_permute_mapping=inv_permute_mapping, + varlen_handler=varlen_handler, + local_attn_handler=local_attn_handler, + modality_dispatcher=modality_dispatcher, + cp_split_sizes=cp_split_sizes, + ) + x = ModalityDispatcher.inv_permute(x, inv_permute_mapping) + + x_out = self.post_adapter(x, video_mask, audio_mask) + + x_out = ulysses_scheduler().undispatch(x_out) + + return x_out + + +Magi2RefinerDiT = Transformer +EntryClass = Magi2RefinerDiT diff --git a/fastvideo/models/dits/magi2_runtime/__init__.py b/fastvideo/models/dits/magi2_runtime/__init__.py new file mode 100644 index 000000000..6692538c6 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact distributed primitives used by the MAGI-2 inference kernels.""" + +from fastvideo.models.dits.magi2_runtime.psm import initialize_expert_parallel, initialize_model_parallel, psm + +__all__ = ["initialize_expert_parallel", "initialize_model_parallel", "psm"] diff --git a/fastvideo/models/dits/magi2_runtime/all_to_all_primitive.py b/fastvideo/models/dits/magi2_runtime/all_to_all_primitive.py new file mode 100644 index 000000000..de73a7773 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/all_to_all_primitive.py @@ -0,0 +1,142 @@ +# Copyright (c) 2025 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Tuple, Union + +import torch +import torch.distributed as dist +from einops import rearrange + + + +def _divide(numerator: int, denominator: int) -> int: + assert numerator % denominator == 0, ( + f"{numerator} is not divisible by {denominator}" + ) + return numerator // denominator + + +class FakeHandle: + def __init__(self): + pass + + def wait(self): + pass + + + +def scatter_seqlen_gather_head( + tensor: torch.Tensor, + split_sizes: List[int] = None, + group: dist.ProcessGroup = None, + async_op: bool = True, +) -> Union[torch.Tensor, Tuple[torch.Tensor, Union[dist.Work, FakeHandle]]]: + """ + Scatter seq_len and gather head_number, for example: + input: (seq_len * cp, hn, hd) + output: (seq_len, cp * hn, hd) + NOTE: seq_len of output maybe not equal, which depends on split_sizes[rank] + NOTE: rearrange the tensor after communication: (cp, seq, hn, hd) -> (seq, cp * hn, hd) + """ + if group is None or dist.get_world_size(group) == 1: + return tensor, FakeHandle() if async_op else tensor + group_world_size = dist.get_world_size(group) + if split_sizes is None: + assert tensor.shape[0] % group_world_size == 0, ( + f"tensor.shape[0] {tensor.shape[0]} % group_world_size {group_world_size} != 0" + ) + split_sizes = [tensor.shape[0] // group_world_size] * group_world_size + assert tensor.is_contiguous() + assert tensor.dim() == 3, f"tensor must be 3D, but got {tensor.dim()}D" + output = torch.empty( + [group_world_size * split_sizes[dist.get_rank(group)], *tensor.shape[1:]], + device=tensor.device, + dtype=tensor.dtype, + ) + output_split_sizes = [split_sizes[dist.get_rank(group)]] * group_world_size + if async_op: + handle = dist.all_to_all_single( + output, + tensor, + output_split_sizes=output_split_sizes, + input_split_sizes=split_sizes, + group=group, + async_op=True, + ) + return output, handle + else: + dist.all_to_all_single( + output, + tensor, + output_split_sizes=output_split_sizes, + input_split_sizes=split_sizes, + group=group, + async_op=False, + ) + return output + + +def batch_scatter_head_gather_seqlen( + inputs: List[torch.Tensor], + split_sizes: List[int] = None, + group: dist.ProcessGroup = None, +) -> List[torch.Tensor]: + """ + Batch scatter head_number and gather seq_len, for example: + inputs[i] input: (seq_len_i, cp * hn_i, hd) + outputs[i] output: (seq_len_i * cp, hn_i, hd) + NOTE: seq_len of inputs maybe not equal across ranks, which depends on split_sizes[rank] + NOTE: fuse along head dim before communication, and split back after + """ + if group is None or dist.get_world_size(group) == 1: + return inputs + rank = dist.get_rank(group) + group_world_size = dist.get_world_size(group) + if split_sizes is None: + split_sizes = [inputs[0].shape[0]] * group_world_size + assert all(input.shape[0] == split_sizes[rank] for input in inputs), ( + f"inputs[0].shape[0] {inputs[0].shape[0]} != split_sizes[rank] {split_sizes[rank]}" + ) + assert all(input.dim() == 3 for input in inputs), ( + f"inputs[0].dim() {inputs[0].dim()} != 3" + ) + for idx in range(len(inputs)): + _, hn, _ = inputs[idx].shape + if group_world_size % hn == 0 and group_world_size != hn: + inputs[idx] = torch.repeat_interleave( + inputs[idx], repeats=_divide(group_world_size, hn), dim=1 + ) + inputs[idx] = rearrange( + inputs[idx], "seq (cp hn) hd -> (cp seq) hn hd", cp=group_world_size + ).contiguous() + + head_split_number = [input.shape[1] for input in inputs] + fused_input = torch.cat(inputs, dim=1).contiguous() + input_split_sizes = [fused_input.shape[0] // group_world_size] * group_world_size + + fused_output = torch.empty( + [sum(split_sizes), *fused_input.shape[1:]], + device=fused_input.device, + dtype=fused_input.dtype, + ) + dist.all_to_all_single( + fused_output, + fused_input, + output_split_sizes=split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=False, + ) + outputs = torch.split(fused_output, head_split_number, dim=1) + return outputs diff --git a/fastvideo/models/dits/magi2_runtime/context_parallel/__init__.py b/fastvideo/models/dits/magi2_runtime/context_parallel/__init__.py new file mode 100644 index 000000000..00ce516d0 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/context_parallel/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Context-parallel scheduling for MAGI-2.""" + +from fastvideo.models.dits.magi2_runtime.context_parallel.ulysses_scheduler import ulysses_scheduler + +__all__ = ["ulysses_scheduler"] diff --git a/fastvideo/models/dits/magi2_runtime/context_parallel/ulysses_scheduler.py b/fastvideo/models/dits/magi2_runtime/context_parallel/ulysses_scheduler.py new file mode 100644 index 000000000..202a50517 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/context_parallel/ulysses_scheduler.py @@ -0,0 +1,146 @@ +# Copyright (c) 2025 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Generic, List, Optional, TypeVar + +import torch +from torch.utils._pytree import tree_map + +from fastvideo.models.dits.magi2_runtime import psm + +from fastvideo.models.dits.magi2_runtime.gather_scatter_primitive import ( + gather_from_context_parallel_region, + scatter_to_context_parallel_region, +) + +T = TypeVar("T") + + +class UlyssesScheduler(Generic[T]): + """ + A naive implementation of Ulysses scheduler for context parallel processing. + + This scheduler handles tensor dispatching and undispatching operations when tensors + enter and exit the context parallel region. It supports arbitrary nested data structures + containing tensors and applies the same balanced sequence partition to every tensor. + + The scheduler splits input tensors along the sequence dimension across multiple GPUs + in the context parallel group, enabling parallel processing of long sequences. + """ + + def __init__(self): + """Initialize the naive Ulysses scheduler.""" + self._cp_split_sizes: Optional[List[int]] = None + + @property + def cp_split_sizes(self): + """Get the current context parallel split sizes.""" + return self._cp_split_sizes + + def _dispatch(self, x: torch.Tensor) -> torch.Tensor: + """ + Dispatch a tensor to the context parallel region. + + This method splits the tensor along the sequence dimension across the context + parallel group. The balanced split sizes differ by at most one token. + + Args: + x: Input tensor with shape [seq_len, ...] where seq_len is the sequence length. + + Returns: + Dispatched tensor that has been split and distributed across the context parallel group. + + Raises: + AssertionError: If the split sizes change between calls, indicating inconsistent + sequence lengths or context parallel group size. + """ + seq_len = x.shape[0] + cp_world_size = psm.get_world_size("cp") + if seq_len % cp_world_size == 0: + cp_split_sizes = [seq_len // cp_world_size] * cp_world_size + else: + num_ranks_with_one_extra = seq_len % cp_world_size + min_tokens_per_rank = (seq_len - num_ranks_with_one_extra) // cp_world_size + cp_split_sizes = [min_tokens_per_rank + 1] * num_ranks_with_one_extra + [min_tokens_per_rank] * ( + cp_world_size - num_ranks_with_one_extra + ) + if self._cp_split_sizes is not None: + assert ( + self._cp_split_sizes == cp_split_sizes + ), f"cp_split_sizes changed from {self._cp_split_sizes} to {cp_split_sizes}" + self._cp_split_sizes = cp_split_sizes + x = scatter_to_context_parallel_region(x, cp_split_sizes, group=psm.get_parallel_group("cp")) + return x + + def _undispatch(self, x: torch.Tensor) -> torch.Tensor: + """ + Undispatch a tensor from the context parallel region. + + This method gathers the balanced tensor slices from all ranks in the context + parallel group and concatenates them into the original sequence. + + Args: + x: Dispatched tensor from the context parallel region. + + Returns: + Reconstructed tensor with the original sequence length. + """ + x = gather_from_context_parallel_region(x, self._cp_split_sizes, group=psm.get_parallel_group("cp")) + return x + + def dispatch(self, tensors: T) -> T: + """ + Apply dispatch operation to all tensor leaf nodes in a nested data structure. + + This method recursively applies the _dispatch operation to all tensors in the + input data structure, preparing them for context parallel computation. The + structure of the input is preserved in the output. + + Args: + tensors: Arbitrary nested data structure containing tensors (single tensor, + tuple, list, dict, etc.). All tensors should have the same sequence + length in their first dimension. + + Returns: + A new data structure with the same structure as input, where all tensors + have been dispatched to the context parallel region. + """ + return tree_map(self._dispatch, tensors) + + def undispatch(self, tensors: T) -> T: + """ + Apply undispatch operation to all tensor leaf nodes in a nested data structure. + + This method recursively applies the _undispatch operation to all tensors in the + input data structure, reconstructing them from the context parallel region. The + structure of the input is preserved in the output. + + Args: + tensors: Arbitrary nested data structure containing dispatched tensors. + + Returns: + A new data structure with the same structure as input, where all tensors + have been reconstructed from the context parallel region. + """ + output = tree_map(self._undispatch, tensors) + self._cp_split_sizes = None + return output + + +_ULYSSES_SCHEDULER = UlyssesScheduler() + + +def ulysses_scheduler() -> UlyssesScheduler: + assert _ULYSSES_SCHEDULER is not None, "ulysses scheduler is not initialized" + return _ULYSSES_SCHEDULER diff --git a/fastvideo/models/dits/magi2_runtime/fastvideo_parallel.py b/fastvideo/models/dits/magi2_runtime/fastvideo_parallel.py new file mode 100644 index 000000000..5aa450490 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/fastvideo_parallel.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Connect MAGI-2 collectives to FastVideo-owned process groups.""" + +from fastvideo.distributed import get_dp_group, get_sp_group +from fastvideo.models.dits.magi2_runtime import psm as psm_manager +from fastvideo.models.dits.magi2_runtime.psm import bind_process_groups + + +MAGI2_PARALLEL_SIZE = 8 + + +def bind_fastvideo_parallel_state() -> None: + """Use FastVideo sequence parallelism for MAGI-2 context and experts.""" + sequence_group = get_sp_group() + data_group = get_dp_group() + if sequence_group.world_size != MAGI2_PARALLEL_SIZE: + raise ValueError( + "MAGI-2 Preview strict inference requires sp_size=8; " + f"received sp_size={sequence_group.world_size}" + ) + bind_process_groups( + cp_group=sequence_group.device_group, + cp_ranks=sequence_group.ranks, + dp_group=data_group.device_group, + dp_ranks=data_group.ranks, + ep_group=sequence_group.device_group, + ep_ranks=sequence_group.ranks, + ) + if psm_manager.get_world_size("cp") != MAGI2_PARALLEL_SIZE: + raise RuntimeError("MAGI-2 context-parallel binding did not preserve eight ranks") + + +__all__ = ["MAGI2_PARALLEL_SIZE", "bind_fastvideo_parallel_state"] diff --git a/fastvideo/models/dits/magi2_runtime/flash_mh_moe/__init__.py b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/__init__.py new file mode 100644 index 000000000..1a38413bf --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/__init__.py @@ -0,0 +1,11 @@ +# Vendored flash_mh_moe inference subset. +# Routing/sorting: pure PyTorch. +# Forward GEMM: fused Triton kernel with deterministic scatter support. +from .fwd import flash_mh_moe_fwd +from .route import compute_topk_probs_and_indices, flash_mh_moe_global_sort + +__all__ = [ + "compute_topk_probs_and_indices", + "flash_mh_moe_global_sort", + "flash_mh_moe_fwd", +] diff --git a/fastvideo/models/dits/magi2_runtime/flash_mh_moe/fwd.py b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/fwd.py new file mode 100644 index 000000000..16b3ea931 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/fwd.py @@ -0,0 +1,40 @@ +# Copyright (c) 2025-2026 SandAI. All Rights Reserved. +# Apache-2.0. +"""Flash-MH-MoE forward — fused Triton kernel with deterministic scatter.""" + +from __future__ import annotations + +import os + +import torch + + +def _is_deterministic() -> bool: + return os.environ.get("MAGI2_DETERMINISTIC", "0") == "1" + + +def flash_mh_moe_fwd( + x: torch.Tensor, + gather_ids: torch.Tensor, + probs: torch.Tensor, + expert_offsets: torch.Tensor, + W_gate: torch.Tensor, + W_up: torch.Tensor, + W_down: torch.Tensor, +) -> torch.Tensor: + """Flash-MH-MoE forward: gather → gate/up GEMM → SwiGLU7 → down GEMM → scatter. + + When ``MAGI2_DETERMINISTIC=1``, uses sequential scatter for bit-exact reproducibility. + """ + from .triton import mh_moe_fwd_func + + return mh_moe_fwd_func( + x=x, + gather_ids=gather_ids, + probs=probs, + expert_offsets=expert_offsets, + W_gate=W_gate, + W_up=W_up, + W_down=W_down, + deterministic=_is_deterministic(), + ) diff --git a/fastvideo/models/dits/magi2_runtime/flash_mh_moe/route.py b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/route.py new file mode 100644 index 000000000..30d0e08f9 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/route.py @@ -0,0 +1,249 @@ +# Copyright (c) 2025-2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Literal, TypeAlias + +import torch +import torch.nn.functional as F + +TopKScoreFunc: TypeAlias = Literal["softmax", "sigmoid"] + + +# ═══════════════════════════════════════════════════════════════════════════ +# Fused top-k probabilities and indices computation +# ═══════════════════════════════════════════════════════════════════════════ + + +def _compute_topk_probs_and_indices_torch_impl( + router_logits: torch.Tensor, + top_k: int, + score_func: TopKScoreFunc = "softmax", + expert_bias: torch.Tensor | None = None, + num_virtual_group: int = 1, + topk_virtual_group: int = 0, + route_norm: bool = True, + norm_eps: float = 1e-12, + use_score_matrix_for_topk_probs: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select and normalize each token's expert-routing probabilities.""" + H, S, E = router_logits.shape + K = top_k + + # ── Gating scores ── + + match score_func: + case "softmax": + router_scores = F.softmax(router_logits, dim=-1) + case "sigmoid": + router_scores = F.sigmoid(router_logits) + case _: + raise ValueError(f"Unknown score_func: {score_func}") + + # ── Aux-free bias for topk selection (not for actual probs) ── + + topk_scores = router_scores + if expert_bias is not None: + bias = expert_bias.view(H, 1, E) # [H, E] -> [H, 1, E] + topk_scores = topk_scores + bias + + # ── TopK per head per token (with optional virtual group filtering) ── + + if num_virtual_group > 1: + assert ( + topk_virtual_group > 0 and topk_virtual_group <= num_virtual_group + ), f"topk_virtual_group ({topk_virtual_group}) must be in (0, num_virtual_group ({num_virtual_group})]" + assert ( + E % num_virtual_group == 0 + ), f"num_experts ({E}) must be divisible by num_virtual_group ({num_virtual_group})" + experts_per_group = E // num_virtual_group + + # Score each group by sum of its top-2 expert scores + # [H, S, E] -> [H, S, num_virtual_group, experts_per_group] + # -> topk(2, dim=-1) -> [H, S, num_virtual_group, 2] + # -> sum(dim=-1) -> [H, S, num_virtual_group] + group_scores = ( + topk_scores.view(H, S, num_virtual_group, experts_per_group) + .topk(min(2, experts_per_group), dim=-1)[0] + .sum(dim=-1) + ) + + # Select top-k virtual groups + # [H, S, num_virtual_group] -> topk(k, dim=-1) -> [H, S, topk_virtual_group] + _, group_idx = torch.topk( + group_scores, + k=topk_virtual_group, + dim=-1, + sorted=False, + ) + + # Build a mask to select experts from the top-k virtual groups + # group_mask: [H, S, num_virtual_group] + # score_mask: [H, S, num_virtual_group, 1] -> expand(dim=-1) -> [H, S, E] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(-1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(H, S, num_virtual_group, experts_per_group) + .reshape(H, S, E) + ) + + # Mask out scores not in the top-k virtual groups before selecting top-k experts + # topk_scores: [H, S, E] -> masked_fill(~score_mask, 0.0) -> [H, S, E] + # -> topk(k, dim=-1) -> topk_probs: [H, S, K] + scores_for_choice = topk_scores.masked_fill(~score_mask.bool(), 0.0) + _, topk_indices = torch.topk(scores_for_choice, k=K, dim=-1, sorted=False) + else: + _, topk_indices = torch.topk(topk_scores, K, dim=-1) # [H, S, K] + + if use_score_matrix_for_topk_probs: + # The score-matrix path masks unselected experts, normalizes across all + # experts, and then gathers the selected probabilities. Gathering before + # normalization is algebraically equivalent but can differ numerically + # because the two paths normalize different tensor representations. + global_routing_matrix = torch.zeros( + H, + S, + E, + device=topk_indices.device, + dtype=torch.int32, + ).scatter_(-1, topk_indices, 1) + + global_score_matrix = global_routing_matrix * router_scores + + # ── Normalize ── + if route_norm: + global_score_matrix = F.normalize( + global_score_matrix, p=1, dim=-1, eps=norm_eps + ) + + topk_probs = global_score_matrix.gather(-1, topk_indices) + else: + # Expert selection uses biased or virtual-group-masked topk_scores, while + # probability magnitudes come from the unbiased router_scores. + topk_probs = router_scores.gather(-1, topk_indices) # [H, S, K] + + # ── Normalize ── + + if route_norm: + topk_probs = F.normalize(topk_probs, p=1, dim=-1, eps=norm_eps) + + return ( + topk_probs, + topk_indices, + ) + + +score_func_name2id_map = { + "softmax": 0, + "sigmoid": 1, +} + + +def compute_topk_probs_and_indices( + router_logits: torch.Tensor, + top_k: int, + score_func: TopKScoreFunc = "softmax", + expert_bias: torch.Tensor | None = None, + num_virtual_group: int = 1, + topk_virtual_group: int = 0, + route_norm: bool = True, + norm_eps: float = 1e-12, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute the top-k probabilities and indices from router logits. + + Args: + router_logits: Tensor of shape [H, S, E], the router logits for each head, token, and expert. + top_k: The number of top probabilities and indices to compute. + score_func: The function to convert logits to scores, either "softmax" or "sigmoid". + expert_bias: Optional tensor of shape [H, E] to be added to the router scores before top-k selection, + for biasing certain experts. + num_virtual_group: If >1, the experts are divided into this many virtual groups + for a two-stage top-k selection (first select top-k virtual groups, + then select top-k experts within those groups). + topk_virtual_group: If num_virtual_group > 1, + the number of top virtual groups to select before selecting top-k experts. + route_norm: Whether to normalize the top-k probabilities to sum to 1. Defaults to ``True``. + norm_eps: the epsilon used for route_norm. Defaults to ``1e-12``. + + + Returns: + A tuple of (topk_probs, topk_indices): + - topk_probs: Tensor of shape [H, S, top_k], the top-k probabilities for each head and token. + - topk_indices: Tensor of shape [H, S, top_k], + the indices of the top-k experts for each head and token. + """ + return _compute_topk_probs_and_indices_torch_impl( + router_logits=router_logits, + top_k=top_k, + score_func=score_func, + expert_bias=expert_bias, + num_virtual_group=num_virtual_group, + topk_virtual_group=topk_virtual_group, + route_norm=route_norm, + norm_eps=norm_eps, + ) + + + +# ═══════════════════════════════════════════════════════════════════════════ +# Mapping topk probs/indices to sorted gather ids and probs and offsets +# prepared as meta inputs for flash_mh_moe kernels +# ═══════════════════════════════════════════════════════════════════════════ + + +def flash_mh_moe_global_sort( + topk_probs: torch.Tensor, + topk_indices: torch.Tensor, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Map per-head expert ids to global [0, H*E) and argsort into CSR format. + + Args: + topk_probs (torch.Tensor): Routing probabilities with shape ``(H, S, K)``. + topk_indices (torch.Tensor): Per-head expert indices in ``[0, E)`` with shape ``(H, S, K)``. + num_experts (int): E, number of experts per head. + + Returns: + gather_ids (torch.Tensor): Token indices to gather from ``[S]`` with shape ``(T,)`` (int32). + probs_sorted (torch.Tensor): Routing probabilities sorted by expert with shape ``(T,)``. + expert_offsets (torch.Tensor): CSR offsets with shape ``(H*E+1,)`` (long). + """ + H, S, K = topk_indices.shape + E = num_experts + device = topk_indices.device + + head_offset = torch.arange(H, device=device).view(H, 1, 1) * E + global_indices = topk_indices + head_offset + + flat_indices = global_indices.reshape(-1) + flat_probs = topk_probs.reshape(-1) + flat_token_ids = ( + torch.arange(S, device=device).view(1, S, 1).expand(H, S, K).reshape(-1) + ) + + sorted_order = flat_indices.argsort(stable=True) + gather_ids = flat_token_ids[sorted_order].to(torch.int32) + probs_sorted = flat_probs[sorted_order].float() + + expert_counts = torch.zeros(H * E, device=device, dtype=torch.long) + expert_counts.scatter_add_( + 0, + flat_indices[sorted_order], + torch.ones_like(sorted_order, dtype=torch.long), + ) + expert_offsets = torch.zeros(H * E + 1, device=device, dtype=torch.long) + expert_offsets[1:] = expert_counts.cumsum(0) + + return gather_ids, probs_sorted, expert_offsets diff --git a/fastvideo/models/dits/magi2_runtime/flash_mh_moe/triton/__init__.py b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/triton/__init__.py new file mode 100644 index 000000000..4e298e9f1 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/triton/__init__.py @@ -0,0 +1,3 @@ +from .mh_moe_fwd import mh_moe_fwd_func + +__all__ = ["mh_moe_fwd_func"] diff --git a/fastvideo/models/dits/magi2_runtime/flash_mh_moe/triton/mh_moe_fwd.py b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/triton/mh_moe_fwd.py new file mode 100644 index 000000000..64fec5d85 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/flash_mh_moe/triton/mh_moe_fwd.py @@ -0,0 +1,335 @@ +# Copyright (c) 2025-2026 SandAI. All Rights Reserved. +# Apache-2.0. +"""Triton kernel for multi-head MoE forward (inference). + +Multi-head MoE forward: gather → GEMM(gate,up) → SwiGLU7 → GEMM(down) → scatter. +Supports deterministic mode (sequential scatter_back) for bit-exact reproducibility. +""" + +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +__all__ = ["mh_moe_fwd_func"] + +SWIGLU7_ALPHA = tl.constexpr(1.702) +SWIGLU7_LIMIT = tl.constexpr(7.0) +SWIGLU7_BIAS = tl.constexpr(1.0) + + +@triton.jit +def _swiglu7_fwd(gate, up, out_dtype: tl.constexpr): + gate_clamped = tl.minimum(gate, SWIGLU7_LIMIT) + up_clamped = tl.maximum(tl.minimum(up, SWIGLU7_LIMIT), -SWIGLU7_LIMIT) + sig = tl.sigmoid(SWIGLU7_ALPHA * gate_clamped) + swish = gate_clamped * sig + return (swish * (up_clamped + SWIGLU7_BIAS)).to(out_dtype) + + +@triton.jit +def _binary_search_expert_id( + cu_expert_tile_counts_ptr, + tile_id, + NUM_EXPERTS: tl.constexpr, + LOG2_NUM_EXPERTS: tl.constexpr, +): + """Map a flattened output tile to the expert that owns the tile.""" + lo = 0 + hi = NUM_EXPERTS + for _ in tl.static_range(0, LOG2_NUM_EXPERTS + 1): + mid = (lo + hi + 1) // 2 + below = tl.load(cu_expert_tile_counts_ptr + mid) <= tile_id + lo = tl.where(below, mid, lo) + hi = tl.where(below, hi, mid - 1) + return lo + + +@triton.jit +def _mh_moe_fwd_kernel( + x_ptr, + wg_ptr, + wu_ptr, + wd_ptr, + y_ptr, + gather_ids_ptr, + probs_ptr, + expert_offsets_ptr, + cu_expert_tile_counts_ptr, + stride_x_s, + stride_x_h, + stride_x_dh, + stride_wg_ne, + stride_wg_dh, + stride_wg_de, + stride_wu_ne, + stride_wu_dh, + stride_wu_de, + stride_wd_ne, + stride_wd_de, + stride_wd_dh, + stride_y_s, + stride_y_h, + stride_y_dh, + D_HEAD: tl.constexpr, + D_EXPERT: tl.constexpr, + NUM_HEADS: tl.constexpr, + NUM_EXPERTS: tl.constexpr, + LOG2_NUM_EXPERTS: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_DH: tl.constexpr, + BLOCK_DE: tl.constexpr, + ACC_DTYPE: tl.constexpr = tl.float32, + DETERMINISTIC: tl.constexpr = False, +): + """Run the fused gather, expert SwiGLU7 projections, and output scatter.""" + tile_id = tl.program_id(0) + + total_tiles = tl.load(cu_expert_tile_counts_ptr + NUM_EXPERTS) + if tile_id >= total_tiles: + return + + expert_id = _binary_search_expert_id( + cu_expert_tile_counts_ptr, tile_id, NUM_EXPERTS, LOG2_NUM_EXPERTS + ) + expert_id_i64 = expert_id.to(tl.int64) + head_id = expert_id // (NUM_EXPERTS // NUM_HEADS) + + tile_in_expert = tile_id - tl.load(cu_expert_tile_counts_ptr + expert_id) + token_start = tl.load(expert_offsets_ptr + expert_id) + tile_in_expert * BLOCK_T + expert_end = tl.load(expert_offsets_ptr + expert_id + 1) + n = tl.minimum(token_start + BLOCK_T, expert_end) - token_start + + block_dh_offs = tl.arange(0, BLOCK_DH) + block_de_offs = tl.arange(0, BLOCK_DE) + block_t_offs = tl.arange(0, BLOCK_T) + dh_offs = tl.arange(0, D_HEAD) + + block_token_offs = token_start + block_t_offs + block_token_mask = block_t_offs < n + + gather_ids = tl.load( + gather_ids_ptr + block_token_offs, mask=block_token_mask, other=0 + ) + probs = tl.load(probs_ptr + block_token_offs, mask=block_token_mask, other=0.0) + x_offs = gather_ids * stride_x_s + head_id * stride_x_h + + y_acc = tl.zeros([BLOCK_T, D_HEAD], dtype=ACC_DTYPE) + for dexpert_start in tl.range(0, D_EXPERT, BLOCK_DE): + dexpert_block_offs = dexpert_start + block_de_offs + + gate_acc = tl.zeros([BLOCK_T, BLOCK_DE], dtype=ACC_DTYPE) + up_acc = tl.zeros([BLOCK_T, BLOCK_DE], dtype=ACC_DTYPE) + for dhead_start in tl.static_range(0, D_HEAD, BLOCK_DH): + dhead_block_offs = dhead_start + block_dh_offs + + x_block = tl.load( + x_ptr + x_offs[:, None] + dhead_block_offs[None, :] * stride_x_dh, + mask=block_token_mask[:, None], + other=0.0, + ) + w_gate = tl.load( + wg_ptr + + expert_id_i64 * stride_wg_ne + + dhead_block_offs[:, None] * stride_wg_dh + + dexpert_block_offs[None, :] * stride_wg_de + ) + w_up = tl.load( + wu_ptr + + expert_id_i64 * stride_wu_ne + + dhead_block_offs[:, None] * stride_wu_dh + + dexpert_block_offs[None, :] * stride_wu_de + ) + + gate_acc += tl.dot(x_block, w_gate) + up_acc += tl.dot(x_block, w_up) + + hidden = _swiglu7_fwd(gate_acc, up_acc, wd_ptr.dtype.element_ty) + + w_down = tl.load( + wd_ptr + + expert_id_i64 * stride_wd_ne + + dexpert_block_offs[:, None] * stride_wd_de + + dh_offs[None, :] * stride_wd_dh + ) + + y_acc += tl.dot(hidden, w_down) + + y_acc = y_acc * probs[:, None] + + if DETERMINISTIC: + # Store to flat (T, D_HEAD) buffer; host scatter_back accumulates deterministically. + y_ptrs = ( + y_ptr + + block_token_offs[:, None] * stride_y_s + + dh_offs[None, :] * stride_y_dh + ) + tl.store( + y_ptrs, + y_acc.to(y_ptr.dtype.element_ty), + mask=block_token_mask[:, None], + ) + else: + # Scatter-reduce by atomic add into (S, H, D_HEAD). + y_offs = gather_ids * stride_y_s + head_id * stride_y_h + y_ptrs = y_ptr + y_offs[:, None] + dh_offs[None, :] * stride_y_dh + tl.atomic_add( + y_ptrs, + y_acc.to(y_ptr.dtype.element_ty), + mask=block_token_mask[:, None], + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Host-side scatter_back (deterministic path) +# ═══════════════════════════════════════════════════════════════════════════ + + +def _scatter_back( + y_sorted: torch.Tensor, + ref_y: torch.Tensor, + gather_ids: torch.Tensor, + expert_offsets: torch.Tensor, + num_flatten_experts: int, + num_heads: int, +) -> torch.Tensor: + """Scatter prob-weighted expert outputs back to (S, H, D) deterministically.""" + D = ref_y.size(2) + num_experts_per_head = num_flatten_experts // num_heads + + y = torch.zeros_like(ref_y).view(-1, D) # (S*H, D) + + expert_seqlens = expert_offsets.diff() + head_id_values = torch.arange(num_flatten_experts, device=gather_ids.device) // num_experts_per_head + head_ids = torch.repeat_interleave(head_id_values, expert_seqlens) + + if num_heads == 1: + scatter_idx = gather_ids.long() + else: + scatter_idx = gather_ids.long() * num_heads + head_ids.long() + + scatter_idx = scatter_idx.unsqueeze(-1).expand_as(y_sorted) + y.scatter_add_(0, scatter_idx, y_sorted.to(y.dtype)) + + return y.view_as(ref_y) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Public forward function +# ═══════════════════════════════════════════════════════════════════════════ + + +def _select_block_config() -> tuple[int, int, int, int, int]: + """(BLOCK_T, BLOCK_DH, BLOCK_DE, NUM_STAGES, NUM_WARPS) for d_head=256.""" + major, _ = torch.cuda.get_device_capability() + if major >= 10: # Blackwell + return (128, 64, 32, 2, 8) + return (128, 64, 32, 2, 8) # Hopper + + +def mh_moe_fwd_func( + x: torch.Tensor, + gather_ids: torch.Tensor, + probs: torch.Tensor, + expert_offsets: torch.Tensor, + W_gate: torch.Tensor, + W_up: torch.Tensor, + W_down: torch.Tensor, + deterministic: bool = False, +) -> torch.Tensor: + """Triton multi-head MoE forward: gather → gate/up GEMM → SwiGLU7 → down GEMM → scatter. + + When ``deterministic=True``, the kernel writes per-expert outputs to a flat buffer + and the host performs a sequential scatter_back for bit-exact reproducibility. + + Args: + x: (S, H, d_head) + gather_ids: (T,) int32 + probs: (T,) float32 + expert_offsets: (H*E+1,) int32/int64 + W_gate: (H*E, d_head, d_expert) + W_up: (H*E, d_head, d_expert) + W_down: (H*E, d_expert, d_head) + deterministic: use sequential scatter for reproducibility + """ + T = gather_ids.size(0) + if T == 0: + return torch.zeros_like(x) + + num_heads = x.size(1) + d_head, d_expert = x.size(2), W_down.size(1) + + BLOCK_T, BLOCK_DH, BLOCK_DE, NUM_STAGES, NUM_WARPS = _select_block_config() + assert d_head % BLOCK_DH == 0, f"d_head={d_head} must be divisible by {BLOCK_DH}" + assert d_expert % BLOCK_DE == 0, f"d_expert={d_expert} must be divisible by {BLOCK_DE}" + + if deterministic: + y = torch.empty(T, 1, d_head, device=x.device, dtype=x.dtype) + else: + y = torch.zeros_like(x) + + num_flatten_experts = expert_offsets.size(0) - 1 + num_flatten_experts_log2 = max(1, math.ceil(math.log2(max(num_flatten_experts, 1) + 1))) + + expert_token_counts = expert_offsets.diff() + expert_tile_counts = (expert_token_counts + BLOCK_T - 1) // BLOCK_T + cu_expert_tile_counts = torch.cat([ + torch.zeros(1, dtype=torch.int32, device=expert_offsets.device), + torch.cumsum(expert_tile_counts, dim=0, dtype=torch.int32), + ]) + + grid_size = (T + BLOCK_T - 1) // BLOCK_T + num_flatten_experts + + _mh_moe_fwd_kernel[(grid_size,)]( + x_ptr=x, + wg_ptr=W_gate, + wu_ptr=W_up, + wd_ptr=W_down, + y_ptr=y, + gather_ids_ptr=gather_ids, + probs_ptr=probs, + expert_offsets_ptr=expert_offsets, + cu_expert_tile_counts_ptr=cu_expert_tile_counts, + stride_x_s=x.stride(0), + stride_x_h=x.stride(1), + stride_x_dh=x.stride(2), + stride_wg_ne=W_gate.stride(0), + stride_wg_dh=W_gate.stride(1), + stride_wg_de=W_gate.stride(2), + stride_wu_ne=W_up.stride(0), + stride_wu_dh=W_up.stride(1), + stride_wu_de=W_up.stride(2), + stride_wd_ne=W_down.stride(0), + stride_wd_de=W_down.stride(1), + stride_wd_dh=W_down.stride(2), + stride_y_s=y.stride(0), + stride_y_h=y.stride(1), + stride_y_dh=y.stride(2), + D_HEAD=d_head, + D_EXPERT=d_expert, + NUM_HEADS=num_heads, + NUM_EXPERTS=num_flatten_experts, + LOG2_NUM_EXPERTS=num_flatten_experts_log2, + BLOCK_T=BLOCK_T, + BLOCK_DH=BLOCK_DH, + BLOCK_DE=BLOCK_DE, + ACC_DTYPE=tl.float32, + DETERMINISTIC=deterministic, + num_stages=NUM_STAGES, + num_warps=NUM_WARPS, + ) + + if deterministic: + y = _scatter_back( + y_sorted=y.view(T, d_head), + ref_y=x, + gather_ids=gather_ids, + expert_offsets=expert_offsets, + num_flatten_experts=num_flatten_experts, + num_heads=num_heads, + ) + + return y diff --git a/fastvideo/models/dits/magi2_runtime/gather_scatter_primitive.py b/fastvideo/models/dits/magi2_runtime/gather_scatter_primitive.py new file mode 100644 index 000000000..cfd162a8b --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/gather_scatter_primitive.py @@ -0,0 +1,221 @@ +# Copyright (c) 2025 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from functools import partial +from typing import List, Union + +import torch +import torch.distributed as dist +from torch.utils._pytree import tree_map + + + +class Metadata: + def __init__(self, dtype: torch.dtype, numel: int, ndim: int, shape: List[int]): + self.dtype = dtype + self.numel = numel + self.ndim = ndim + self.shape = shape + + def __repr__(self): + return f"Metadata(dtype={self.dtype}, numel={self.numel}, ndim={self.ndim}, shape={self.shape})" + + +def _gather_metadata(tensor_list: List[torch.Tensor], group: dist.ProcessGroup) -> List[List[Metadata]]: + """Exchange tensor shapes and dtypes for an arbitrary distributed gather.""" + dist.get_rank(group) + world_size = dist.get_world_size(group) + + local_rank = torch.distributed.get_rank() % torch.cuda.device_count() + assert ( + local_rank == torch.cuda.current_device() + ), f"local_rank {local_rank} != current_device {torch.cuda.current_device()}" + device = tensor_list[0].device if len(tensor_list) > 0 else torch.device("cuda") + + # ========== Step 1: flatten local tensor list ========== + + # Metadata: [dtype_code, numel, ndim, *shape] + local_metadata = [] + + dtype_map = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2, torch.int32: 3, torch.int64: 4, torch.uint8: 5} + reverse_dtype_map = {v: k for k, v in dtype_map.items()} + + for t in tensor_list: + dtype_code = dtype_map[t.dtype] + shape = list(t.shape) + numel = t.numel() + local_metadata.append(torch.tensor([dtype_code, numel, len(shape)] + shape, dtype=torch.int32, device=device)) + + if local_metadata: + local_metadata_tensor = torch.cat(local_metadata) + else: + local_metadata_tensor = torch.empty(0, dtype=torch.int32, device=device) + local_metadata_tensor = local_metadata_tensor.contiguous() + local_metadata_len = torch.tensor([local_metadata_tensor.numel()], dtype=torch.int32, device=device) + + # ========== Step 2: all_gather metadata lengths ========== + metadata_lens = [torch.empty_like(local_metadata_len) for _ in range(world_size)] + dist.all_gather(metadata_lens, local_metadata_len, group) + + # ========== Step 3: all_gather metadata payloads (with cpu tensor) ========== + metadata_lists = [torch.empty(m.item(), dtype=torch.int32, device=device) for m in metadata_lens] + dist.all_gather(metadata_lists, local_metadata_tensor, group) + + # ========== Step 4: decode metadata and reconstruct tensor list ========== + result = [] + for metadata_list in metadata_lists: + offset = 0 + local_metadata = [] + while offset < metadata_list.numel(): + dtype_code = metadata_list[offset].item() + numel = metadata_list[offset + 1].item() + ndim = metadata_list[offset + 2].item() + shape = metadata_list[offset + 3 : offset + 3 + ndim].tolist() + offset += 3 + ndim + + local_metadata.append(Metadata(reverse_dtype_map[dtype_code], numel, ndim, shape)) + result.append(local_metadata) + + return result + + +def _get_dtype_and_assert_consistency(metadata_lists: List[List[Metadata]]): + dtype_set = set() + for metadata_list in metadata_lists: + for metadata in metadata_list: + dtype_set.add(metadata.dtype) + assert len(dtype_set) == 1, f"Metadata lists are not consistent: {dtype_set}" + return dtype_set.pop() + + +def _get_numel_for_each_rank(metadata_lists: List[List[Metadata]]) -> List[int]: + return [sum(meta.numel for meta in metadata_list) for metadata_list in metadata_lists] + + +def gather_arbitrary_tensor_list(tensor_list: List[torch.Tensor], group: dist.ProcessGroup) -> List[torch.Tensor]: + """ + Magic gather primitive. Provide the following features: + 1. Support tensor list with different length for each rank. + 2. Support arbitrary Tensor, which means the Tensor can have different shapes but same dtype. + 3. Support empty tensor_list in some ranks without padding. + + Args: + tensor_list: A list of tensors to gather. + group: The process group to use. + + Returns: + A list of tensors gathered from all ranks. + """ + + dist.get_rank(group) + world_size = dist.get_world_size(group) + + local_rank = torch.distributed.get_rank() % torch.cuda.device_count() + assert ( + local_rank == torch.cuda.current_device() + ), f"local_rank {local_rank} != current_device {torch.cuda.current_device()}" + device = tensor_list[0].device if len(tensor_list) > 0 else torch.device("cuda") + + # Step 1: Gather metadata + metadata_lists = _gather_metadata(tensor_list, group) + tensor_dtype = _get_dtype_and_assert_consistency(metadata_lists) + + # Step 2: Flatten local tensors into a single 1D buffer + if tensor_list: + flat_tensor = torch.cat([t.flatten() for t in tensor_list], dim=0).contiguous() + else: + flat_tensor = torch.empty(0, dtype=tensor_dtype, device=device) # dummy, will be ignored + + # Step 3: Gather lengths from metadata + all_numels_int = _get_numel_for_each_rank(metadata_lists) + + # Step 4: Allocate buffers and gather flat tensor data + output_flat_tensors = [] + for numel in all_numels_int: + output_flat_tensors.append(torch.empty(numel, dtype=tensor_dtype, device=device)) + dist.all_gather(output_flat_tensors, flat_tensor, group) + + # Step 5: Reconstruct individual tensors using metadata + gathered_tensor_lists = [] + for i in range(world_size): + flat = output_flat_tensors[i] + if flat.numel() == 0: + continue + metadata_list = metadata_lists[i] + offset = 0 + for meta in metadata_list: + numel = meta.numel + t = flat[offset : offset + numel].view(meta.shape).to(meta.dtype) + offset += numel + gathered_tensor_lists.append(t) + + return gathered_tensor_lists + + +def _scatter_to_context_parallel_region(input: torch.Tensor, split_sizes: List[int], group: dist.ProcessGroup = None): + """Split the tensor along its first dimension and keep the + corresponding slice.""" + # Split along first dimension with padding. + rank = dist.get_rank(group) + dim_offset = sum(split_sizes[:rank]) + output = input[dim_offset : dim_offset + split_sizes[rank]].contiguous() + return output + + +def scatter_to_context_parallel_region( + inputs: Union[torch.Tensor, List[torch.Tensor]], split_sizes: List[int] = None, group: dist.ProcessGroup = None +): + """Split the tensor along its first dimension and keep the + corresponding slice.""" + if group is None or torch.distributed.get_world_size(group) == 1: + return inputs + + if split_sizes is None: + assert ( + inputs.shape[0] % dist.get_world_size(group) == 0 + ), f"inputs.shape[0] {inputs.shape[0]} % dist.get_world_size(group) {dist.get_world_size(group)} != 0" + split_sizes = [inputs.shape[0] // dist.get_world_size(group)] * dist.get_world_size(group) + + partial_func = partial(_scatter_to_context_parallel_region, split_sizes=split_sizes, group=group) + return tree_map(partial_func, inputs) + + +def _gather_from_context_parallel_region( + input: Union[torch.Tensor, List[torch.Tensor]], split_sizes: List[int], group: dist.ProcessGroup = None +): + """Gather context-parallel slices and concatenate them in rank order.""" + input = input.contiguous() + dim_size = list(input.size()) + dim_size[0] = sum(split_sizes) + + output = torch.empty(dim_size, dtype=input.dtype, device=input.device) + outputs = list(torch.split(output, split_sizes, dim=0)) + torch.distributed.all_gather(outputs, input, group=group) + output = torch.concat(outputs, dim=0) + + return output + + +def gather_from_context_parallel_region( + inputs: Union[torch.Tensor, List[torch.Tensor]], split_sizes: List[int] = None, group: dist.ProcessGroup = None +): + """Gather tensors and concatinate along the first dimension.""" + + if group is None or torch.distributed.get_world_size(group) == 1: + return inputs + + if split_sizes is None: + split_sizes = [inputs.shape[0] * dist.get_world_size(group)] + partial_func = partial(_gather_from_context_parallel_region, split_sizes=split_sizes, group=group) + return tree_map(partial_func, inputs) diff --git a/fastvideo/models/dits/magi2_runtime/psm.py b/fastvideo/models/dits/magi2_runtime/psm.py new file mode 100644 index 000000000..502750034 --- /dev/null +++ b/fastvideo/models/dits/magi2_runtime/psm.py @@ -0,0 +1,174 @@ +# Copyright (c) 2025-2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Process-group management for MAGI-2 inference. + +Inference uses three parallelism axes: + +* **Context parallel (cp)** splits one sample's token sequence across ranks. +* **Data parallel (dp)** gives different samples to different ranks. +* **Expert parallel (ep)** shards MoE experts across ranks. + +Ranks are numbered cp-major: a cp group is a contiguous block while a dp +group strides across blocks. With ``world_size=8, cp_size=4`` the cp +groups are ``[0,1,2,3]`` and ``[4,5,6,7]``; dp groups are ``[0,4]``, +``[1,5]``, ``[2,6]``, ``[3,7]``. + +All inference code accesses groups through the ``psm`` singleton below. +""" + +from datetime import timedelta +from typing import List, Optional + +import torch.distributed as dist + +_CP_GROUP = None +_CP_RANKS: Optional[List[int]] = None +_DP_GROUP = None +_DP_RANKS: Optional[List[int]] = None +_EP_GROUP = None +_EP_RANKS: Optional[List[int]] = None + + +def bind_process_groups( + cp_group: "dist.ProcessGroup", + cp_ranks: list[int], + dp_group: "dist.ProcessGroup", + dp_ranks: list[int], + ep_group: "dist.ProcessGroup", + ep_ranks: list[int], +) -> None: + """Bind MAGI-2 parallel axes to process groups owned by FastVideo.""" + global _CP_GROUP, _CP_RANKS, _DP_GROUP, _DP_RANKS, _EP_GROUP, _EP_RANKS + + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before binding MAGI-2 groups") + rank = dist.get_rank() + for group_name, group, ranks in ( + ("cp", cp_group, cp_ranks), + ("dp", dp_group, dp_ranks), + ("ep", ep_group, ep_ranks), + ): + if group is None or rank not in ranks: + raise ValueError(f"Rank {rank} does not belong to the MAGI-2 {group_name} group {ranks}") + + _CP_GROUP, _CP_RANKS = cp_group, list(cp_ranks) + _DP_GROUP, _DP_RANKS = dp_group, list(dp_ranks) + _EP_GROUP, _EP_RANKS = ep_group, list(ep_ranks) + + +def initialize_model_parallel(cp_size: int = 1, distributed_timeout_minutes: int = 30) -> None: + """Create cp and dp groups. Must be called after ``dist.init_process_group``.""" + global _CP_GROUP, _CP_RANKS, _DP_GROUP, _DP_RANKS + + assert dist.is_initialized(), "torch.distributed must be initialized first" + assert _CP_GROUP is None, "parallel state is already initialized" + + world_size = dist.get_world_size() + assert world_size % cp_size == 0, f"world_size {world_size} not divisible by cp_size {cp_size}" + dp_size = world_size // cp_size + rank = dist.get_rank() + timeout = timedelta(minutes=distributed_timeout_minutes) + + for i in range(dp_size): + ranks = list(range(i * cp_size, (i + 1) * cp_size)) + group = dist.new_group(ranks, timeout=timeout) + if rank in ranks: + _CP_GROUP, _CP_RANKS = group, ranks + + for i in range(cp_size): + ranks = list(range(i, world_size, cp_size)) + group = dist.new_group(ranks, timeout=timeout) + if rank in ranks: + _DP_GROUP, _DP_RANKS = group, ranks + + +def initialize_expert_parallel(ep_size: int) -> None: + """Create ep groups. Call after ``initialize_model_parallel``.""" + global _EP_GROUP, _EP_RANKS + if ep_size <= 1: + return + world_size = dist.get_world_size() + rank = dist.get_rank() + assert world_size % ep_size == 0 + for i in range(world_size // ep_size): + ranks = list(range(i * ep_size, (i + 1) * ep_size)) + group = dist.new_group(ranks) + if rank in ranks: + _EP_GROUP, _EP_RANKS = group, ranks + + +class ParallelStateManager: + """Minimal psm interface used by MAGI-2 inference.""" + + @staticmethod + def is_initialized() -> bool: + return _CP_GROUP is not None + + @staticmethod + def get_global_rank() -> int: + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + return 0 + + def get_world_size(self, dim: str = "") -> int: + if not dist.is_available() or not dist.is_initialized(): + return 1 + if dim == "cp": + return dist.get_world_size(group=_CP_GROUP) + elif dim == "dp": + return dist.get_world_size(group=_DP_GROUP) + elif dim == "ep": + return dist.get_world_size(group=_EP_GROUP) if _EP_GROUP else dist.get_world_size() + return dist.get_world_size() + + def get_local_rank(self, dim: str) -> int: + if not dist.is_available() or not dist.is_initialized(): + return 0 + if dim == "cp": + return dist.get_rank(group=_CP_GROUP) + elif dim == "dp": + return dist.get_rank(group=_DP_GROUP) + elif dim == "ep": + return dist.get_rank(group=_EP_GROUP) if _EP_GROUP else dist.get_rank() + return 0 + + def is_group_first_rank(self, dim: str) -> bool: + return self.get_local_rank(dim) == 0 + + def is_group_last_rank(self, dim: str) -> bool: + return self.get_local_rank(dim) == (self.get_world_size(dim) - 1) + + def get_parallel_group(self, dim: str) -> Optional["dist.ProcessGroup"]: + if not dist.is_available() or not dist.is_initialized(): + return None + if dim == "cp": + return _CP_GROUP + elif dim == "dp": + return _DP_GROUP + elif dim == "ep": + return _EP_GROUP + return None + + def get_global_ranks(self, dim: str) -> List[int]: + if dim == "cp": + return _CP_RANKS or [self.get_global_rank()] + elif dim == "dp": + return _DP_RANKS or [self.get_global_rank()] + elif dim == "ep": + return _EP_RANKS or list(range(dist.get_world_size())) + return [self.get_global_rank()] + + +psm = ParallelStateManager() diff --git a/fastvideo/models/encoders/qwen3_5.py b/fastvideo/models/encoders/qwen3_5.py new file mode 100644 index 000000000..e24616399 --- /dev/null +++ b/fastvideo/models/encoders/qwen3_5.py @@ -0,0 +1,1001 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +from pathlib import Path +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn.functional as F +from safetensors.torch import safe_open +from tokenizers import Regex, pre_tokenizers +from torch import nn +from transformers import AutoTokenizer + +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.encoders.qwen3_5 import Magi2Qwen35Config +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.loader.utils import set_default_torch_dtype +from fastvideo.models.loader.weight_utils import default_weight_loader + + +_LANGUAGE_MODEL_PREFIX = "model.language_model." + + +def _iter_indexed_language_model_weights( + model_path: Path, + weight_map: dict[str, str], + device: torch.device, +) -> Iterable[tuple[str, torch.Tensor]]: + """Read only Qwen3.5 language-model tensors from indexed checkpoint shards.""" + weights_by_shard: dict[str, list[str]] = {} + for checkpoint_name, shard_name in weight_map.items(): + if checkpoint_name.startswith(_LANGUAGE_MODEL_PREFIX): + weights_by_shard.setdefault(shard_name, []).append(checkpoint_name) + + shard_paths = { + shard_name: model_path / shard_name for shard_name in weights_by_shard + } + missing_shards = [str(path) for path in shard_paths.values() if not path.is_file()] + if missing_shards: + raise FileNotFoundError( + "Missing Qwen3.5 language-model weight shards: " + + ", ".join(sorted(missing_shards)) + ) + + for shard_name in sorted(weights_by_shard): + with safe_open( + str(shard_paths[shard_name]), + framework="pt", + device=str(device), + ) as shard: + for checkpoint_name in sorted(weights_by_shard[shard_name]): + yield checkpoint_name, shard.get_tensor(checkpoint_name) + + +def strip_empty(obj): + if isinstance(obj, dict): + return { + key: value + for key, value in ((key, strip_empty(value)) for key, value in obj.items()) + if value is not None and value != [] and value != {} + } + if isinstance(obj, list): + return [strip_empty(value) for value in obj if value is not None] + return obj + + +def _one_line(value: Any) -> str: + if value is None: + return "" + if not isinstance(value, str): + value = str(value) + return " ".join(value.split()) + + +def _as_dict(value: Any) -> dict[str, Any]: + return value if isinstance(value, dict) else {} + + +def _as_list(value: Any) -> list: + if value is None: + return [] + return value if isinstance(value, list) else [value] + + +def _pick(values: dict, keys: Iterable[str]) -> list[tuple[str, Any]]: + return [ + (key, values[key]) + for key in keys + if key in values and values[key] not in (None, [], {}) + ] + + +def json_to_compact_markdown(raw_json: str | dict[str, Any]) -> str: + """Flatten a structured MAGI-2 prompt while preserving semantic sections.""" + obj = json.loads(raw_json.strip()) if isinstance(raw_json, str) else raw_json + + obj = strip_empty(obj) + if not isinstance(obj, dict) or "global_layer" not in obj: + return ( + raw_json + if isinstance(raw_json, str) + else json.dumps(raw_json, ensure_ascii=False) + ) + + global_layer = _as_dict(obj.get("global_layer")) + dynamic_layer = _as_dict(obj.get("dynamic_layer")) + reference_layer = obj.get("reference_layer", []) + + lines: list[str] = [] + context = _one_line(global_layer.get("context")) + description = _one_line(global_layer.get("description")) + if context or description: + lines.append(f"context: {context}" if context else "context") + if description: + lines.append(description) + + aesthetics = _as_dict(global_layer.get("aesthetics")) + aesthetic_parts = [] + for label, key in ( + ("style", "style"), + ("mood", "mood_atmosphere"), + ("color", "color_scheme"), + ): + value = _one_line(aesthetics.get(key)) + if value: + aesthetic_parts.append(f"{label}={value}") + if aesthetic_parts: + lines.append("aesthetics: " + "; ".join(aesthetic_parts)) + + audio = _as_dict(global_layer.get("audio_baseline")) + dialogue = _as_dict(audio.get("dialogue")) + audio_parts = [] + language = _one_line(dialogue.get("language")) + speakers = _as_list(dialogue.get("speaker_tags")) + ambience = _one_line(audio.get("ambience")) + if language: + audio_parts.append(f"language={language}") + if speakers: + audio_parts.append("speakers=" + ",".join(map(_one_line, speakers))) + if ambience: + audio_parts.append(f"ambience={ambience}") + if audio_parts: + lines.append("audio: " + "; ".join(audio_parts)) + + subjects = global_layer.get("alive_subjects_static") + if isinstance(subjects, list) and subjects: + lines.append("subjects:") + for subject in subjects: + if not isinstance(subject, dict): + continue + sid = _one_line(subject.get("subject_id")) + desc = _one_line(subject.get("description")) + position = _one_line(subject.get("position")) + orientation = _one_line(subject.get("orientation")) + attrs = [] + for key, value in _pick( + _as_dict(subject.get("visual_attributes")), + ["gender", "age_group", "ethnicity", "clothing", "appearance_details"], + ): + value = _one_line(value) + if value: + attrs.append(f"{key}={value}") + row = " - " + " | ".join([part for part in [sid, desc] if part]) + extra = "; ".join([part for part in [position, orientation] if part]) + if extra: + row += f" ({extra})" + if attrs: + row += " :: " + "; ".join(attrs) + lines.append(row) + + objects = global_layer.get("objects_static") + if isinstance(objects, list) and objects: + lines.append("objects:") + for obj_item in objects: + if not isinstance(obj_item, dict): + continue + oid = _one_line(obj_item.get("object_id")) + desc = _one_line(obj_item.get("description")) + shape = _one_line(obj_item.get("shape_and_color")) + position = _one_line(obj_item.get("position")) + row = " - " + " | ".join([part for part in [oid, desc] if part]) + details = "; ".join([part for part in [shape, position] if part]) + if details: + row += " :: " + details + lines.append(row) + + segments = dynamic_layer.get("timeline_segments") + if isinstance(segments, list) and segments: + lines.append("timeline:") + for segment in segments: + if not isinstance(segment, dict): + continue + basic = _as_dict(segment.get("segment_basic_info")) + timestamp = _one_line(basic.get("timestamp_range")) + desc = _one_line(basic.get("segment_description")) + head = f" - {timestamp}" if timestamp else " -" + if desc: + head += f" {desc}" + lines.append(head.rstrip()) + + segment_audio = _as_dict(segment.get("audio")) + for dialogue_line in segment_audio.get("dialogue_lines", []) or []: + if not isinstance(dialogue_line, dict): + continue + speaker = _one_line(dialogue_line.get("speaker")) + text = _one_line(dialogue_line.get("text")) + timestamp = _one_line(dialogue_line.get("timestamp")) + if text: + prefix = ( + f" - dialogue {speaker}: " if speaker else " - dialogue: " + ) + line = prefix + text + if timestamp: + line += f" ({timestamp})" + lines.append(line) + + for alive in segment.get("alive_subjects", []) or []: + if not isinstance(alive, dict): + continue + sid = _one_line(alive.get("subject_id")) + action = _as_dict(alive.get("action")) + parts = [ + _one_line(action.get("primary_action")), + _one_line(action.get("interaction")), + _one_line(action.get("facial_expression")), + ] + parts = [part for part in parts if part] + if parts: + lines.append(f" - action {sid}: " + "; ".join(parts)) + + for obj_item in segment.get("objects", []) or []: + if not isinstance(obj_item, dict): + continue + oid = _one_line(obj_item.get("object_id")) + state = _as_dict(obj_item.get("dynamic_state")) + parts = [ + _one_line(state.get("state_change")), + _one_line(state.get("motion_detail")), + ] + parts = [part for part in parts if part] + if parts: + lines.append(f" - objects {oid}: " + "; ".join(parts)) + + if isinstance(reference_layer, list): + lines.extend(str(line) for line in reference_layer if str(line).strip()) + + return "\n".join(line for line in lines if line.strip()) + + +def _l2norm( + tensor: torch.Tensor, + dim: int = -1, + eps: float = 1e-6, +) -> torch.Tensor: + """Normalize query or key vectors with the Qwen3.5 epsilon convention.""" + inverse_norm = torch.rsqrt((tensor * tensor).sum(dim=dim, keepdim=True) + eps) + return tensor * inverse_norm + + +def _torch_chunk_gated_delta_rule( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + decay: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, +) -> torch.Tensor: + """Evaluate the Qwen3.5 chunked gated-delta recurrence in FP32.""" + output_dtype = query.dtype + query = _l2norm(query) + key = _l2norm(key) + query, key, value, beta, decay = [ + tensor.transpose(1, 2).contiguous().to(torch.float32) + for tensor in (query, key, value, beta, decay) + ] + + batch_size, num_heads, sequence_length, key_head_dim = key.shape + value_head_dim = value.shape[-1] + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + decay = F.pad(decay, (0, pad_size)) + padded_length = sequence_length + pad_size + query = query * (query.shape[-1] ** -0.5) + + value_beta = value * beta.unsqueeze(-1) + key_beta = key * beta.unsqueeze(-1) + query, key, value, key_beta, value_beta = [ + tensor.reshape( + tensor.shape[0], + tensor.shape[1], + -1, + chunk_size, + tensor.shape[-1], + ) + for tensor in (query, key, value, key_beta, value_beta) + ] + decay = decay.reshape(decay.shape[0], decay.shape[1], -1, chunk_size) + diagonal_mask = torch.triu( + torch.ones( + chunk_size, + chunk_size, + dtype=torch.bool, + device=query.device, + ), + diagonal=0, + ) + + decay = decay.cumsum(dim=-1) + decay_mask = ( + (decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float() + ).tril() + attention = -( + (key_beta @ key.transpose(-1, -2)) * decay_mask + ).masked_fill(diagonal_mask, 0) + for row_index in range(1, chunk_size): + row = attention[..., row_index, :row_index].clone() + preceding_rows = attention[..., :row_index, :row_index].clone() + attention[..., row_index, :row_index] = row + ( + row.unsqueeze(-1) * preceding_rows + ).sum(-2) + attention = attention + torch.eye( + chunk_size, + dtype=attention.dtype, + device=attention.device, + ) + value = attention @ value_beta + cumulative_key = attention @ (key_beta * decay.exp().unsqueeze(-1)) + recurrent_state = torch.zeros( + batch_size, + num_heads, + key_head_dim, + value_head_dim, + dtype=value.dtype, + device=value.device, + ) + recurrence_output = torch.zeros_like(value) + causal_mask = torch.triu( + torch.ones( + chunk_size, + chunk_size, + dtype=torch.bool, + device=query.device, + ), + diagonal=1, + ) + + for chunk_index in range(padded_length // chunk_size): + chunk_query = query[:, :, chunk_index] + chunk_key = key[:, :, chunk_index] + chunk_value = value[:, :, chunk_index] + chunk_attention = ( + chunk_query + @ chunk_key.transpose(-1, -2) + * decay_mask[:, :, chunk_index] + ).masked_fill_(causal_mask, 0) + recurrent_value = cumulative_key[:, :, chunk_index] @ recurrent_state + value_delta = chunk_value - recurrent_value + inter_chunk_attention = ( + chunk_query * decay[:, :, chunk_index, :, None].exp() + ) @ recurrent_state + recurrence_output[:, :, chunk_index] = ( + inter_chunk_attention + chunk_attention @ value_delta + ) + recurrent_state = ( + recurrent_state + * decay[:, :, chunk_index, -1, None, None].exp() + + ( + chunk_key + * ( + decay[:, :, chunk_index, -1, None] + - decay[:, :, chunk_index] + ).exp()[..., None] + ).transpose(-1, -2) + @ value_delta + ) + + recurrence_output = recurrence_output.reshape( + recurrence_output.shape[0], + recurrence_output.shape[1], + -1, + recurrence_output.shape[-1], + ) + recurrence_output = recurrence_output[:, :, :sequence_length] + return recurrence_output.transpose(1, 2).contiguous().to(output_dtype) + + +def _rotate_half(tensor: torch.Tensor) -> torch.Tensor: + """Rotate the two halves of each rotary-embedding vector.""" + first_half = tensor[..., : tensor.shape[-1] // 2] + second_half = tensor[..., tensor.shape[-1] // 2 :] + return torch.cat((-second_half, first_half), dim=-1) + + +def _apply_rotary_pos_emb( + query: torch.Tensor, + key: torch.Tensor, + cosine: torch.Tensor, + sine: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply partial rotary embeddings to Qwen3.5 query and key heads.""" + cosine = cosine.unsqueeze(1) + sine = sine.unsqueeze(1) + rotary_dim = cosine.shape[-1] + query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] + key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] + query_rotary = query_rotary * cosine + _rotate_half(query_rotary) * sine + key_rotary = key_rotary * cosine + _rotate_half(key_rotary) * sine + return ( + torch.cat((query_rotary, query_pass), dim=-1), + torch.cat((key_rotary, key_pass), dim=-1), + ) + + +class Qwen35RMSNormGated(nn.Module): + """Apply FP32 root-mean-square normalization before a SiLU gate.""" + + def __init__(self, hidden_size: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward( + self, + hidden_states: torch.Tensor, + gate: torch.Tensor, + ) -> torch.Tensor: + """Normalize value heads and apply their learned input gates.""" + input_dtype = hidden_states.dtype + normalized_states = hidden_states.to(torch.float32) + variance = normalized_states.pow(2).mean(-1, keepdim=True) + normalized_states = normalized_states * torch.rsqrt( + variance + self.variance_epsilon + ) + normalized_states = self.weight * normalized_states.to(input_dtype) + normalized_states = normalized_states * F.silu(gate.to(torch.float32)) + return normalized_states.to(input_dtype) + + +class Qwen35GatedDeltaNet(nn.Module): + """Mix tokens with Qwen3.5 causal convolution and gated delta recurrence.""" + + def __init__(self, config: Magi2Qwen35Config, layer_index: int) -> None: + """Create the projections and recurrence parameters for one linear layer.""" + super().__init__() + self.hidden_size = config.hidden_size + self.num_value_heads = config.linear_num_value_heads + self.num_key_heads = config.linear_num_key_heads + self.key_head_dim = config.linear_key_head_dim + self.value_head_dim = config.linear_value_head_dim + self.key_dim = self.key_head_dim * self.num_key_heads + self.value_dim = self.value_head_dim * self.num_value_heads + self.conv_kernel_size = config.linear_conv_kernel_dim + self.layer_index = layer_index + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + bias=False, + kernel_size=self.conv_kernel_size, + groups=self.conv_dim, + padding=self.conv_kernel_size - 1, + ) + self.dt_bias = nn.Parameter(torch.ones(self.num_value_heads)) + transition = torch.empty(self.num_value_heads).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(transition)) + self.norm = Qwen35RMSNormGated( + self.value_head_dim, + eps=config.rms_norm_eps, + ) + self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + self.in_proj_qkv = nn.Linear( + self.hidden_size, + self.key_dim * 2 + self.value_dim, + bias=False, + ) + self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False) + self.in_proj_b = nn.Linear( + self.hidden_size, + self.num_value_heads, + bias=False, + ) + self.in_proj_a = nn.Linear( + self.hidden_size, + self.num_value_heads, + bias=False, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Transform one complete prompt without recurrent cache state.""" + batch_size, sequence_length, _ = hidden_states.shape + mixed_qkv = self.in_proj_qkv(hidden_states).transpose(1, 2) + gate = self.in_proj_z(hidden_states).reshape( + batch_size, + sequence_length, + -1, + self.value_head_dim, + ) + beta = self.in_proj_b(hidden_states).sigmoid() + decay_input = self.in_proj_a(hidden_states) + mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :sequence_length]) + mixed_qkv = mixed_qkv.transpose(1, 2) + query, key, value = torch.split( + mixed_qkv, + (self.key_dim, self.key_dim, self.value_dim), + dim=-1, + ) + query = query.reshape( + batch_size, + sequence_length, + -1, + self.key_head_dim, + ) + key = key.reshape( + batch_size, + sequence_length, + -1, + self.key_head_dim, + ) + value = value.reshape( + batch_size, + sequence_length, + -1, + self.value_head_dim, + ) + decay = -self.A_log.float().exp() * F.softplus( + decay_input.float() + self.dt_bias + ) + head_repeat = self.num_value_heads // self.num_key_heads + query = query.repeat_interleave(head_repeat, dim=2) + key = key.repeat_interleave(head_repeat, dim=2) + mixed_output = _torch_chunk_gated_delta_rule( + query, + key, + value, + decay, + beta, + ) + mixed_output = mixed_output.reshape(-1, self.value_head_dim) + gate = gate.reshape(-1, self.value_head_dim) + mixed_output = self.norm(mixed_output, gate) + mixed_output = mixed_output.reshape(batch_size, sequence_length, -1) + return self.out_proj(mixed_output) + + +class Qwen35RotaryEmbedding(nn.Module): + """Construct Qwen3.5 partial rotary embeddings in FP32.""" + + def __init__(self, config: Magi2Qwen35Config) -> None: + """Create inverse frequencies for the checkpoint's partial rotary span.""" + super().__init__() + rotary_dim = int(config.head_dim * config.partial_rotary_factor) + inverse_frequency = 1.0 / ( + config.rope_theta + ** ( + torch.arange(0, rotary_dim, 2, dtype=torch.int64).to( + dtype=torch.float32 + ) + / rotary_dim + ) + ) + self.register_buffer("inv_freq", inverse_frequency, persistent=False) + self.mrope_section = config.mrope_section + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return cosine and sine tensors for three identical text coordinates.""" + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand( + 3, + position_ids.shape[0], + -1, + ) + inverse_frequency = self.inv_freq[None, None, :, None].float().expand( + 3, + position_ids.shape[1], + -1, + 1, + ) + expanded_positions = position_ids[:, :, None, :].float() + device_type = ( + hidden_states.device.type + if hidden_states.device.type != "mps" + else "cpu" + ) + with torch.autocast(device_type=device_type, enabled=False): + frequencies = ( + inverse_frequency.float() @ expanded_positions.float() + ).transpose(2, 3) + text_frequencies = frequencies[0] + for dimension, offset in enumerate((1, 2), start=1): + length = self.mrope_section[dimension] * 3 + indices = slice(offset, length, 3) + text_frequencies[..., indices] = frequencies[ + dimension, + ..., + indices, + ] + embedding = torch.cat((text_frequencies, text_frequencies), dim=-1) + cosine = embedding.cos() + sine = embedding.sin() + return ( + cosine.to(dtype=hidden_states.dtype), + sine.to(dtype=hidden_states.dtype), + ) + + +class Qwen35RMSNorm(nn.Module): + """Apply Qwen3.5 one-centered RMS normalization.""" + + def __init__(self, hidden_size: int, eps: float) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(hidden_size)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Normalize in FP32 and cast the scaled output to the input dtype.""" + normalized_states = hidden_states.float() + normalized_states = normalized_states * torch.rsqrt( + normalized_states.pow(2).mean(-1, keepdim=True) + self.eps + ) + normalized_states = normalized_states * (1.0 + self.weight.float()) + return normalized_states.type_as(hidden_states) + + +class Qwen35Attention(nn.Module): + """Apply grouped-query causal self-attention with an output gate.""" + + def __init__(self, config: Magi2Qwen35Config, layer_index: int) -> None: + """Create the query, key, value, gate, and output projections.""" + super().__init__() + self.layer_index = layer_index + self.head_dim = config.head_dim + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.q_proj = nn.Linear( + config.hidden_size, + self.num_attention_heads * self.head_dim * 2, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + self.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = Qwen35RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = Qwen35RMSNorm(self.head_dim, config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """Run the complete-prompt SDPA path selected by Transformers 5.5.0.""" + input_shape = hidden_states.shape[:-1] + query_shape = (*input_shape, self.num_attention_heads, self.head_dim) + key_value_shape = ( + *input_shape, + self.num_key_value_heads, + self.head_dim, + ) + query_states, gate = torch.chunk( + self.q_proj(hidden_states).view( + *input_shape, + self.num_attention_heads, + self.head_dim * 2, + ), + 2, + dim=-1, + ) + gate = gate.reshape(*input_shape, -1) + query_states = self.q_norm(query_states.view(query_shape)).transpose(1, 2) + key_states = self.k_norm( + self.k_proj(hidden_states).view(key_value_shape) + ).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(key_value_shape).transpose(1, 2) + query_states, key_states = _apply_rotary_pos_emb( + query_states, + key_states, + *position_embeddings, + ) + attention_output = F.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=None, + dropout_p=0.0, + is_causal=query_states.shape[2] > 1, + scale=self.scaling, + enable_gqa=True, + ) + attention_output = attention_output.transpose(1, 2).contiguous() + attention_output = attention_output.reshape(*input_shape, -1).contiguous() + attention_output = attention_output * torch.sigmoid(gate) + return self.o_proj(attention_output) + + +class Qwen35MLP(nn.Module): + """Apply the Qwen3.5 SwiGLU feed-forward network.""" + + def __init__(self, config: Magi2Qwen35Config) -> None: + """Create separate gate, value, and output projections for SwiGLU.""" + super().__init__() + self.gate_proj = nn.Linear( + config.hidden_size, + config.intermediate_size, + bias=False, + ) + self.up_proj = nn.Linear( + config.hidden_size, + config.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear( + config.intermediate_size, + config.hidden_size, + bias=False, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project hidden states through the gated SiLU activation.""" + return self.down_proj( + F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states) + ) + + +class Qwen35DecoderLayer(nn.Module): + """Combine one Qwen3.5 token mixer and one feed-forward residual block.""" + + def __init__(self, config: Magi2Qwen35Config, layer_index: int) -> None: + """Select the token mixer and create the layer's residual modules.""" + super().__init__() + self.layer_type = ( + "full_attention" + if (layer_index + 1) % config.full_attention_interval == 0 + else "linear_attention" + ) + if self.layer_type == "linear_attention": + self.linear_attn = Qwen35GatedDeltaNet(config, layer_index) + else: + self.self_attn = Qwen35Attention(config, layer_index) + self.mlp = Qwen35MLP(config) + self.input_layernorm = Qwen35RMSNorm( + config.hidden_size, + config.rms_norm_eps, + ) + self.post_attention_layernorm = Qwen35RMSNorm( + config.hidden_size, + config.rms_norm_eps, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """Apply pre-normalized token mixing and feed-forward residuals.""" + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + if self.layer_type == "linear_attention": + hidden_states = self.linear_attn(hidden_states) + else: + hidden_states = self.self_attn(hidden_states, position_embeddings) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class Magi2Qwen35TextEncoder(TextEncoder): + """Encode MAGI-2 prompts with the text-only Qwen3.5-27B architecture.""" + + supports_hf_from_pretrained = True + + def __init__(self, config: Magi2Qwen35Config) -> None: + """Create the prompt-processing state for the release Qwen3.5 model.""" + super().__init__(config) + self.text_model: nn.Module | None = None + self.max_length = config.text_len + self.skip_layer = config.hidden_state_skip_layer + self.tokenizer = None + + @classmethod + def from_pretrained_local( + cls, + model_path: str, + model_config: Magi2Qwen35Config, + dtype: torch.dtype, + device: torch.device, + ) -> Magi2Qwen35TextEncoder: + """Load the same Transformers Qwen3.5 text path as the official release.""" + from transformers import Qwen3_5TextModel + + with set_default_torch_dtype(dtype): + encoder = cls(model_config) + encoder.text_model = Qwen3_5TextModel.from_pretrained( + model_path, + torch_dtype=dtype, + local_files_only=True, + ) + encoder.to(device) + encoder.tokenizer = AutoTokenizer.from_pretrained( + model_path, + padding_side="right", + local_files_only=True, + ) + cjk_split = pre_tokenizers.Split( + pattern=Regex( + r"([\u1100-\u11ff\u2e80-\ua4cf\ua840-\uD7AF\uF900-\uFAFF\uFE30-\uFE4F\uFF65-\uFFDC\U00020000-\U0002FFFF])" + ), + behavior="isolated", + ) + original_pre_tokenizer = encoder.tokenizer.backend_tokenizer.pre_tokenizer + encoder.tokenizer.backend_tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [cjk_split, original_pre_tokenizer] + ) + encoder.requires_grad_(False) + return encoder.eval() + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + """Load only the text-model tensors from the composite Qwen3.5 checkpoint.""" + parameter_by_name = dict(self.named_parameters()) + loaded_parameters: set[str] = set() + for checkpoint_name, checkpoint_tensor in weights: + if not checkpoint_name.startswith(_LANGUAGE_MODEL_PREFIX): + continue + parameter_name = checkpoint_name[len(_LANGUAGE_MODEL_PREFIX) :] + if parameter_name not in parameter_by_name: + raise KeyError( + f"Unexpected Qwen3.5 text parameter: {checkpoint_name}" + ) + default_weight_loader( + parameter_by_name[parameter_name], + checkpoint_tensor, + ) + loaded_parameters.add(parameter_name) + return loaded_parameters + + def _normalize_prompt(self, prompt: str) -> str: + """Convert structured JSON prompts into the compact MAGI-2 text form.""" + try: + parsed_json = json.loads(prompt) + if isinstance(parsed_json, (dict, list)): + return json_to_compact_markdown(prompt) + except (json.JSONDecodeError, TypeError): + pass + return prompt + + def _require_tokenizer(self): + """Return the tokenizer initialized by the local checkpoint loader.""" + if self.tokenizer is None: + raise RuntimeError( + "Load Magi2Qwen35TextEncoder through from_pretrained_local()." + ) + return self.tokenizer + + def get_target_token_indices( + self, + prompt: str, + target_str: str | None, + ) -> list[int] | None: + """Map one figure marker's character span to tokenizer indices.""" + if not target_str: + return None + prompt = self._normalize_prompt(prompt) + tokenizer = self._require_tokenizer() + inputs = tokenizer( + [prompt], + return_tensors="pt", + padding="longest", + return_offsets_mapping=True, + max_length=self.max_length, + truncation=True, + ) + offsets = inputs["offset_mapping"][0] + start_char_index = prompt.find(target_str) + if start_char_index == -1: + return None + end_char_index = start_char_index + len(target_str) + token_indices: list[int] = [] + for token_index, (start, end) in enumerate(offsets): + if start == 0 and end == 0 and token_index != 0: + continue + if max(start, start_char_index) < min(end, end_char_index): + token_indices.append(token_index) + return token_indices + + def get_special_token( + self, + prompt: str, + target_strs: list[str], + text_feature: torch.Tensor, + ) -> torch.Tensor: + """Pool prompt embeddings over requested figure-marker token spans.""" + embeddings = [] + for target_str in target_strs: + token_indices = self.get_target_token_indices(prompt, target_str) + if token_indices: + embeddings.append( + text_feature[0, token_indices, :].mean(dim=0).clone() + ) + else: + embeddings.append( + torch.zeros( + text_feature.shape[-1], + device=text_feature.device, + dtype=text_feature.dtype, + ) + ) + return torch.stack(embeddings, dim=0) + + @torch.inference_mode() + def encode(self, prompt: str) -> torch.Tensor: + """Tokenize one prompt and return the hidden state before the final two layers.""" + normalized_prompt = self._normalize_prompt(prompt) + tokenizer = self._require_tokenizer() + if self.text_model is None: + raise RuntimeError("The Qwen3.5 text model has not been loaded") + device = next(self.text_model.parameters()).device + inputs = tokenizer( + [normalized_prompt], + return_tensors="pt", + padding="longest", + max_length=self.max_length, + truncation=True, + ).to(device) + outputs = self.text_model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + output_hidden_states=True, + return_dict=True, + ) + if self.skip_layer == 0: + return outputs.last_hidden_state + assert outputs.hidden_states is not None + return outputs.hidden_states[-(self.skip_layer + 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: + """Delegate one prompt to the Transformers model used by MAGI-2.""" + if self.text_model is None: + raise RuntimeError("The Qwen3.5 text model has not been loaded") + outputs = self.text_model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + return_dict=True, + **kwargs, + ) + return BaseEncoderOutput( + last_hidden_state=outputs.last_hidden_state, + hidden_states=outputs.hidden_states, + attention_mask=attention_mask, + ) + + +EntryClass = Magi2Qwen35TextEncoder diff --git a/fastvideo/models/vaes/magi2_audio_vae.py b/fastvideo/models/vaes/magi2_audio_vae.py new file mode 100644 index 000000000..3f5e736b3 --- /dev/null +++ b/fastvideo/models/vaes/magi2_audio_vae.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MAGI-2 Stable Audio Open decoder loading and source-key conversion.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path + +import torch +from safetensors import safe_open +from torch import nn + +from fastvideo.models.vaes.oobleck import OobleckDecoder + +_SOURCE_PREFIX = "pretransform.model." + + +class Magi2AudioVAE(nn.Module): + """Decode MAGI-2 audio latents with the Stable Audio Oobleck decoder.""" + + def __init__( + self, + decoder: OobleckDecoder, + sampling_rate: int, + downsampling_ratio: int, + ) -> None: + """Store the decoder and the published audio sampling geometry.""" + super().__init__() + self.decoder = decoder + self.sampling_rate = sampling_rate + self.hop_length = downsampling_ratio + + def decode(self, latent: torch.Tensor) -> torch.Tensor: + """Decode channel-major audio latents into channel-major waveforms.""" + if latent.ndim == 2: + latent = latent.unsqueeze(0) + return self.decoder(latent) + + +def _source_decoder_to_fastvideo_key( + source_name: str, + decoder_block_count: int, +) -> str: + """Map one Stable Audio sequential decoder key to an Oobleck decoder key.""" + parts = source_name.split(".") + if parts[:2] != ["decoder", "layers"] or len(parts) < 4: + raise ValueError(f"Unsupported Stable Audio decoder key: {source_name}") + + top_level_index = int(parts[2]) + parameter_parts = parts[3:] + if top_level_index == 0: + return ".".join(["decoder", "conv1", *parameter_parts]) + if top_level_index == decoder_block_count + 1: + return ".".join(["decoder", "snake1", *parameter_parts]) + if top_level_index == decoder_block_count + 2: + return ".".join(["decoder", "conv2", *parameter_parts]) + if not 1 <= top_level_index <= decoder_block_count: + raise ValueError(f"Unsupported Stable Audio decoder key: {source_name}") + + if parameter_parts[:1] != ["layers"] or len(parameter_parts) < 3: + raise ValueError(f"Unsupported Stable Audio decoder block key: {source_name}") + block_index = top_level_index - 1 + block_layer_index = int(parameter_parts[1]) + block_parameter_parts = parameter_parts[2:] + if block_layer_index == 0: + target_parts = ["decoder", "block", str(block_index), "snake1"] + elif block_layer_index == 1: + target_parts = ["decoder", "block", str(block_index), "conv_t1"] + elif 2 <= block_layer_index <= 4: + if block_parameter_parts[:1] != ["layers"] or len(block_parameter_parts) < 3: + raise ValueError(f"Unsupported Stable Audio residual key: {source_name}") + residual_layer_index = int(block_parameter_parts[1]) + residual_names = ("snake1", "conv1", "snake2", "conv2") + if residual_layer_index >= len(residual_names): + raise ValueError(f"Unsupported Stable Audio residual key: {source_name}") + target_parts = [ + "decoder", + "block", + str(block_index), + f"res_unit{block_layer_index - 1}", + residual_names[residual_layer_index], + ] + block_parameter_parts = block_parameter_parts[2:] + else: + raise ValueError(f"Unsupported Stable Audio decoder block key: {source_name}") + return ".".join([*target_parts, *block_parameter_parts]) + + +def _load_decoder_state( + decoder_model: Magi2AudioVAE, + checkpoint_path: Path, + decoder_block_count: int, +) -> None: + """Read only decoder tensors and load the mapped state strictly.""" + target_state = decoder_model.state_dict() + mapped_state: dict[str, torch.Tensor] = {} + with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint: + for checkpoint_name in list(checkpoint.keys()): + if not checkpoint_name.startswith(f"{_SOURCE_PREFIX}decoder."): + continue + source_name = checkpoint_name.removeprefix(_SOURCE_PREFIX) + target_name = _source_decoder_to_fastvideo_key( + source_name, + decoder_block_count, + ) + if target_name not in target_state: + raise KeyError( + f"Stable Audio key {checkpoint_name} maps to unknown key {target_name}" + ) + source_tensor = checkpoint.get_tensor(checkpoint_name) + target_tensor = target_state[target_name] + if source_tensor.numel() != target_tensor.numel(): + raise ValueError( + f"Stable Audio tensor size differs for {checkpoint_name}: " + f"source={tuple(source_tensor.shape)}, target={tuple(target_tensor.shape)}" + ) + mapped_state[target_name] = source_tensor.reshape(target_tensor.shape) + incompatible = decoder_model.load_state_dict(mapped_state, strict=True) + if incompatible.missing_keys or incompatible.unexpected_keys: + raise RuntimeError( + "Stable Audio strict load failed: " + f"missing={incompatible.missing_keys}, " + f"unexpected={incompatible.unexpected_keys}" + ) + + +def load_magi2_audio_vae( + component_dir: str | Path, + device: torch.device | str, +) -> Magi2AudioVAE: + """Build and strictly load MAGI-2's published FP32 audio decoder.""" + component_path = Path(component_dir) + config_path = component_path / "model_config.json" + checkpoint_path = component_path / "model.safetensors" + with config_path.open(encoding="utf-8") as config_file: + full_config = json.load(config_file) + sampling_rate = int(full_config["sample_rate"]) + vae_config = full_config["model"]["pretransform"]["config"] + decoder_spec = vae_config["decoder"] + decoder_config = decoder_spec["config"] + if decoder_spec["type"] != "oobleck" or not decoder_config["use_snake"]: + raise ValueError("MAGI-2 audio decoding requires the published Oobleck Snake decoder") + if decoder_config.get("use_nearest_upsample", False): + raise ValueError("MAGI-2 audio decoding requires transposed-convolution upsampling") + if decoder_config.get("final_tanh", True): + raise ValueError("MAGI-2 audio decoding requires final_tanh=false") + + strides = [int(stride) for stride in decoder_config["strides"]] + downsampling_ratio = int(vae_config["downsampling_ratio"]) + if math.prod(strides) != downsampling_ratio: + raise ValueError("Stable Audio strides do not match the declared downsampling ratio") + decoder = OobleckDecoder( + channels=int(decoder_config["channels"]), + input_channels=int(decoder_config["latent_dim"]), + audio_channels=int(decoder_config["out_channels"]), + upsampling_ratios=list(reversed(strides)), + channel_multiples=[int(value) for value in decoder_config["c_mults"]], + ) + audio_vae = Magi2AudioVAE( + decoder=decoder, + sampling_rate=sampling_rate, + downsampling_ratio=downsampling_ratio, + ) + _load_decoder_state(audio_vae, checkpoint_path, len(strides)) + audio_vae.to(device=device, dtype=torch.float32) + audio_vae.requires_grad_(False) + return audio_vae.eval() + + +__all__ = ["Magi2AudioVAE", "load_magi2_audio_vae"] diff --git a/fastvideo/models/vaes/magi2_turbo_vae.py b/fastvideo/models/vaes/magi2_turbo_vae.py new file mode 100644 index 000000000..b2204a05f --- /dev/null +++ b/fastvideo/models/vaes/magi2_turbo_vae.py @@ -0,0 +1,1190 @@ +# Copyright (c) 2025 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import torch +import torch.nn as nn +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from einops import rearrange +from magi_compiler.api import magi_compile + +from fastvideo.configs.models.vaes.magi2_turbo_vae import Magi2TurboVAEConfig + +__all__ = ["Magi2TurboVAEModel", "TurboVAED", "get_turbo_vaed"] + +ACT2CLS = { + "swish": nn.SiLU, + "silu": nn.SiLU, + "mish": nn.Mish, + "gelu": nn.GELU, + "relu": nn.ReLU, +} + + +def get_activation(act_fn: str) -> nn.Module: + act_fn = act_fn.lower() + if act_fn in ACT2CLS: + return ACT2CLS[act_fn]() + else: + raise ValueError( + f"activation function {act_fn} not found in ACT2CLS mapping {list(ACT2CLS.keys())}" + ) + + +def _turbo_unpatchify(x, patch_size): + """ + Unpatchify operation: convert patched representation back to original spatial resolution. + Similar to Wan VAE's unpatchify. + + Args: + x: Input tensor with shape [batch_size, (channels * patch_size * patch_size), frame, height, width] + patch_size: The patch size used during patchification + + Returns: + Tensor with shape [batch_size, channels, frame, height * patch_size, width * patch_size] + """ + if patch_size == 1: + return x + + if x.dim() != 5: + raise ValueError(f"Invalid input shape: {x.shape}") + + # x shape: [batch_size, (channels * patch_size * patch_size), frame, height, width] + batch_size, c_patches, frames, height, width = x.shape + channels = c_patches // (patch_size * patch_size) + + # Reshape to [b, c, patch_size, patch_size, f, h, w] + x = x.view(batch_size, channels, patch_size, patch_size, frames, height, width) + + # Rearrange to [b, c, f, h * patch_size, w * patch_size] + x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous() + x = x.view(batch_size, channels, frames, height * patch_size, width * patch_size) + + return x + + +class RMSNorm(nn.Module): + r""" + RMS Norm as introduced in https://huggingface.co/papers/1910.07467 by Zhang et al. + + Args: + dim (`int`): Number of dimensions to use for `weights`. Only effective when `elementwise_affine` is True. + eps (`float`): Small value to use when calculating the reciprocal of the square-root. + elementwise_affine (`bool`, defaults to `True`): + Boolean flag to denote if affine transformation should be applied. + bias (`bool`, defaults to False): If also training the `bias` param. + """ + + def __init__( + self, dim, eps: float, elementwise_affine: bool = True, bias: bool = False + ): + """Initialize channel-wise root mean square normalization.""" + super().__init__() + + self.eps = eps + self.elementwise_affine = elementwise_affine + + if isinstance(dim, int): + dim = (dim,) + + self.dim = torch.Size(dim) + + self.weight = None + + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim)) + + + def forward(self, hidden_states): + """Normalize each activation vector with its root mean square.""" + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + + if self.weight is not None: + # convert into half-precision if necessary + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + hidden_states = hidden_states * self.weight + else: + hidden_states = hidden_states.to(input_dtype) + + return hidden_states + + +class TurboVAEDConv2dSplitUpsampler(nn.Module): + def __init__( + self, + in_channels: int, + kernel_size: int | tuple[int, int] = 3, + stride: int | tuple[int, int] = 1, + upscale_factor: int = 1, + padding_mode: str = "zeros", + ): + """Build the spatial convolution and pixel-shuffle upsampler.""" + super().__init__() + + self.in_channels = in_channels + self.stride = stride if isinstance(stride, tuple) else (stride, stride) + self.upscale_factor = upscale_factor + + out_channels = in_channels + + self.kernel_size = ( + kernel_size + if isinstance(kernel_size, tuple) + else (kernel_size, kernel_size) + ) + + height_pad = self.kernel_size[0] // 2 + width_pad = self.kernel_size[1] // 2 + padding = (height_pad, width_pad) + + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=self.kernel_size, + stride=1, + padding=padding, + padding_mode=padding_mode, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv(hidden_states) + hidden_states = torch.nn.functional.pixel_shuffle(hidden_states, self.stride[0]) + + return hidden_states + + +class TurboVAEDCausalConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + dilation: int | tuple[int, int, int] = 1, + groups: int = 1, + padding_mode: str = "zeros", + is_causal: bool = False, + ): + """Build a temporally edge-padded three-dimensional convolution.""" + super().__init__() + + assert not is_causal + self.in_channels = in_channels + self.out_channels = out_channels + self.is_causal = is_causal + self.kernel_size = ( + kernel_size + if isinstance(kernel_size, tuple) + else (kernel_size, kernel_size, kernel_size) + ) + + dilation = dilation if isinstance(dilation, tuple) else (dilation, 1, 1) + stride = stride if isinstance(stride, tuple) else (stride, stride, stride) + height_pad = self.kernel_size[1] // 2 + width_pad = self.kernel_size[2] // 2 + padding = (0, height_pad, width_pad) + + self.conv = nn.Conv3d( + in_channels, + out_channels, + self.kernel_size, + stride=stride, + dilation=dilation, + groups=groups, + padding=padding, + padding_mode=padding_mode, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Apply temporal edge padding followed by three-dimensional convolution.""" + time_kernel_size = self.kernel_size[0] + + if time_kernel_size > 1: + pad_left = hidden_states[:, :, :1, :, :].repeat( + (1, 1, (time_kernel_size - 1) // 2, 1, 1) + ) + pad_right = hidden_states[:, :, -1:, :, :].repeat( + (1, 1, (time_kernel_size - 1) // 2, 1, 1) + ) + hidden_states = torch.cat([pad_left, hidden_states, pad_right], dim=2) + + hidden_states = self.conv(hidden_states) + return hidden_states + + +class TurboVAEDCausalDepthwiseSeparableConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + dilation: int | tuple[int, int, int] = 1, + padding_mode: str = "zeros", + is_causal: bool = True, + ): + """Build depthwise and pointwise convolutions for video features.""" + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.is_causal = is_causal + self.kernel_size = ( + kernel_size + if isinstance(kernel_size, tuple) + else (kernel_size, kernel_size, kernel_size) + ) + self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride) + self.dilation = dilation if isinstance(dilation, tuple) else (dilation, 1, 1) + + # Calculate padding for height and width dimensions + height_pad = self.kernel_size[1] // 2 + width_pad = self.kernel_size[2] // 2 + self.padding = (0, height_pad, width_pad) + + # Depthwise Convolution + self.depthwise_conv = nn.Conv3d( + in_channels, + in_channels, + self.kernel_size, + stride=self.stride, + dilation=self.dilation, + groups=in_channels, # Each input channel is convolved separately + padding=self.padding, + padding_mode=padding_mode, + ) + + # Pointwise Convolution + self.pointwise_conv = nn.Conv3d( + in_channels, out_channels, kernel_size=1 + ) # 1x1x1 convolution to mix channels + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Apply temporal padding, depthwise convolution, and channel mixing.""" + time_kernel_size = self.kernel_size[0] + if time_kernel_size > 1: + pad_count = (time_kernel_size - 1) // 2 + pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, pad_count, 1, 1)) + pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, pad_count, 1, 1)) + hidden_states = torch.cat([pad_left, hidden_states, pad_right], dim=2) + + # Apply depthwise convolution + hidden_states = self.depthwise_conv(hidden_states) + # Apply pointwise convolution + hidden_states = self.pointwise_conv(hidden_states) + + return hidden_states + + +class TurboVAEDResNetBlock3d(nn.Module): + r""" + A 3D ResNet block used in the TurboVAED model. + + Args: + in_channels (`int`): + Number of input channels. + out_channels (`int`, *optional*): + Number of output channels. If None, defaults to `in_channels`. + eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + elementwise_affine (`bool`, defaults to `False`): + Whether to enable elementwise affine parameters in the normalization layers. + non_linearity (`str`, defaults to `"swish"`): + Activation function to use. + conv_shortcut (bool, defaults to `False`): + Whether or not to use a convolution shortcut. + """ + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + eps: float = 1e-6, + elementwise_affine: bool = False, + non_linearity: str = "swish", + is_causal: bool = True, + is_upsampler_modified: bool = False, + is_dw_conv: bool = False, + dw_kernel_size: int = 3, + ) -> None: + """Build the residual normalization and convolution branches.""" + super().__init__() + + out_channels = out_channels or in_channels + + self.nonlinearity = get_activation(non_linearity) + + self.conv_operation = ( + TurboVAEDCausalConv3d + if not is_dw_conv + else TurboVAEDCausalDepthwiseSeparableConv3d + ) + self.kernel_size = 3 if not is_dw_conv else dw_kernel_size + + self.is_upsampler_modified = is_upsampler_modified + self.replace_nonlinearity = get_activation("relu") + + self.norm1 = RMSNorm( + in_channels, eps=1e-8, elementwise_affine=elementwise_affine + ) + self.conv1 = self.conv_operation( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=self.kernel_size, + is_causal=is_causal, + ) + + self.norm2 = RMSNorm( + out_channels, eps=1e-8, elementwise_affine=elementwise_affine + ) + self.conv2 = self.conv_operation( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=self.kernel_size, + is_causal=is_causal, + ) + + self.norm3 = None + self.conv_shortcut = None + if in_channels != out_channels: + self.norm3 = RMSNorm( + in_channels, eps=eps, elementwise_affine=elementwise_affine + ) + self.conv_shortcut = self.conv_operation( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + is_causal=is_causal, + ) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """Transform video features and add the residual branch.""" + hidden_states = inputs + + hidden_states = self.norm1(hidden_states) + + if self.is_upsampler_modified: + hidden_states = self.replace_nonlinearity(hidden_states) + else: + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv2(hidden_states) + + if self.norm3 is not None: + inputs = self.norm3(inputs) + + if self.conv_shortcut is not None: + inputs = self.conv_shortcut(inputs) + + hidden_states = hidden_states + inputs + return hidden_states + + +class WanUpsample(nn.Upsample): + r""" + Perform upsampling while ensuring the output tensor has the same data type as the input. + + Args: + x (torch.Tensor): Input tensor to be upsampled. + + Returns: + torch.Tensor: Upsampled tensor with the same data type as the input. + """ + + def forward(self, x): + return super().forward(x.float()).type_as(x) + + +class WanResample(nn.Module): + r""" + A custom resampling module for 2D and 3D data. + + Args: + dim (int): The number of input/output channels. + mode (str): The resampling mode. Must be one of: + - 'none': No resampling (identity operation). + - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. + - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. + """ + + def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: + """Build the requested spatial or spatiotemporal upsampling operation.""" + super().__init__() + self.dim = dim + self.mode = mode + + # default to dim //2 + if upsample_out_dim is None: + upsample_out_dim = dim // 2 + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + self.time_conv = TurboVAEDCausalConv3d(dim, dim * 2, (3, 1, 1)) + else: + self.resample = nn.Identity() + + def forward(self, x, is_first_chunk: bool = True): + """Upsample video features and trim the first temporal overlap frame.""" + b, c, t, h, w = x.shape + if self.mode == "upsample3d": + x = self.time_conv(x) + x = rearrange(x, "b (n_split c) t h w -> b c (t n_split) h w", n_split=2) + assert x.shape == (b, c, t * 2, h, w), "x.shape: {}, expected: {}".format( + x.shape, (b, c, t * 2, h, w) + ) + if is_first_chunk: + x = x[:, :, 1:] + + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=b) + + return x + + +class TurboVAEDMidBlock3d(nn.Module): + r""" + A middle block used in the TurboVAED model. + + Args: + in_channels (`int`): + Number of input channels. + num_layers (`int`, defaults to `1`): + Number of resnet layers. + resnet_eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + resnet_act_fn (`str`, defaults to `"swish"`): + Activation function to use. + is_causal (`bool`, defaults to `True`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + """ + + + def __init__( + self, + in_channels: int, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + is_causal: bool = True, + is_dw_conv: bool = False, + dw_kernel_size: int = 3, + ) -> None: + """Build the decoder middle block from residual layers.""" + super().__init__() + + resnets = [] + for _ in range(num_layers): + resnets.append( + TurboVAEDResNetBlock3d( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + non_linearity=resnet_act_fn, + is_causal=is_causal, + is_dw_conv=is_dw_conv, + dw_kernel_size=dw_kernel_size, + ) + ) + self.resnets = nn.ModuleList(resnets) + + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + r"""Forward method of the `LTXMidBlock3D` class.""" + + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + return hidden_states + + +class TurboVAEDUpBlock3d(nn.Module): + r""" + Up block used in the TurboVAED model. + + Args: + in_channels (`int`): + Number of input channels. + out_channels (`int`, *optional*): + Number of output channels. If None, defaults to `in_channels`. + num_layers (`int`, defaults to `1`): + Number of resnet layers. + resnet_eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + resnet_act_fn (`str`, defaults to `"swish"`): + Activation function to use. + spatio_temporal_scale (`bool`, defaults to `True`): + Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension. + Whether or not to downsample across temporal dimension. + is_causal (`bool`, defaults to `True`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + """ + + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + spatio_temporal_scale: bool = True, + is_causal: bool = True, + is_dw_conv: bool = False, + dw_kernel_size: int = 3, + spatio_only: bool = False, + ): + """Build one decoder stage with optional channel and scale changes.""" + super().__init__() + + out_channels = out_channels or in_channels + + self.conv_in = None + if in_channels != out_channels: + self.conv_in = TurboVAEDResNetBlock3d( + in_channels=in_channels, + out_channels=out_channels, + eps=resnet_eps, + non_linearity=resnet_act_fn, + is_causal=is_causal, + is_dw_conv=is_dw_conv, + dw_kernel_size=dw_kernel_size, + ) + + self.upsamplers = None + if spatio_temporal_scale: + self.upsamplers = nn.ModuleList( + [ + WanResample( + dim=out_channels, + mode="upsample2d" if spatio_only else "upsample3d", + upsample_out_dim=out_channels, + ) + ] + ) + + resnets = [] + for _ in range(num_layers): + resnets.append( + TurboVAEDResNetBlock3d( + in_channels=out_channels, + out_channels=out_channels, + eps=resnet_eps, + non_linearity=resnet_act_fn, + is_causal=is_causal, + is_dw_conv=is_dw_conv, + dw_kernel_size=dw_kernel_size, + is_upsampler_modified=(spatio_temporal_scale), + ) + ) + self.resnets = nn.ModuleList(resnets) + + + def forward( + self, hidden_states: torch.Tensor, is_first_chunk: bool + ) -> torch.Tensor: + """Apply channel projection, upsampling, and residual layers.""" + if self.conv_in is not None: + hidden_states = self.conv_in(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, is_first_chunk=is_first_chunk) + + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + return hidden_states + + +class TurboVAEDDecoder3d(nn.Module): + r""" + The `TurboVAEDDecoder3d` layer of a variational autoencoder that decodes its latent representation into an output + sample. + + Args: + in_channels (`int`, defaults to 128): + Number of latent channels. + out_channels (`int`, defaults to 3): + Number of output channels. + block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + The number of output channels for each block. + spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`: + Whether a block should contain spatio-temporal upscaling layers or not. + layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): + The number of layers per block. + patch_size (`int`, defaults to `4`): + The size of spatial patches. + patch_size_t (`int`, defaults to `1`): + The size of temporal patches. + resnet_norm_eps (`float`, defaults to `1e-6`): + Epsilon value for ResNet normalization layers. + is_causal (`bool`, defaults to `False`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + """ + + def __init__( + self, + in_channels: int = 128, + out_channels: int = 3, + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False), + layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4), + patch_size: int = 4, + patch_size_t: int = 1, + resnet_norm_eps: float = 1e-6, + is_causal: bool = False, + decoder_is_dw_conv: tuple[bool, ...] = (False, False, False, False, False), + decoder_dw_kernel_size: int = 3, + spatio_only: tuple[bool, ...] = (False, False, False, False), + upsampling: bool = False, + use_unpatchify: bool = False, + ) -> None: + """Build the distilled decoder from its published architecture fields.""" + super().__init__() + + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.out_channels = out_channels + + self.upsampling = upsampling + self.use_unpatchify = use_unpatchify + + block_out_channels = tuple(reversed(block_out_channels)) + spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling)) + layers_per_block = tuple(reversed(layers_per_block)) + decoder_is_dw_conv = tuple(reversed(decoder_is_dw_conv)) + spatio_only = tuple(reversed(spatio_only)) + output_channel = block_out_channels[0] + + self.conv_in = TurboVAEDCausalConv3d( + in_channels=in_channels, + out_channels=output_channel, + kernel_size=3, + stride=1, + is_causal=is_causal, + ) + + self.mid_block = TurboVAEDMidBlock3d( + in_channels=output_channel, + num_layers=layers_per_block[0], + resnet_eps=resnet_norm_eps, + is_causal=is_causal, + is_dw_conv=decoder_is_dw_conv[0], + dw_kernel_size=decoder_dw_kernel_size, + ) + + # up blocks + num_block_out_channels = len(block_out_channels) + self.up_blocks = nn.ModuleList([]) + for i in range(num_block_out_channels): + input_channel = output_channel + output_channel = block_out_channels[i] + + up_block = TurboVAEDUpBlock3d( + in_channels=input_channel, + out_channels=output_channel, + num_layers=layers_per_block[i + 1], + resnet_eps=resnet_norm_eps, + spatio_temporal_scale=spatio_temporal_scaling[i], + is_causal=is_causal, + is_dw_conv=decoder_is_dw_conv[i + 1], + dw_kernel_size=decoder_dw_kernel_size, + spatio_only=spatio_only[i], + ) + + self.up_blocks.append(up_block) + + # out + assert self.patch_size == 2 + if not self.use_unpatchify: + self.norm_up_1 = RMSNorm(output_channel, eps=1e-8, elementwise_affine=False) + self.upsampler2d_1 = TurboVAEDConv2dSplitUpsampler( + in_channels=output_channel, kernel_size=3, stride=(2, 2) + ) + output_channel = output_channel // (2 * 2) + + self.conv_act = nn.SiLU() + + # When use_unpatchify=True, conv_out outputs more channels (out_channels * patch_size^2) + # and unpatchify will recover the spatial resolution + conv_out_channels = self.out_channels + if self.use_unpatchify and self.patch_size >= 2: + conv_out_channels = self.out_channels * self.patch_size * self.patch_size + + self.conv_out = TurboVAEDCausalConv3d( + in_channels=output_channel, + out_channels=conv_out_channels, + kernel_size=3, + stride=1, + is_causal=is_causal, + ) + + + def forward( + self, hidden_states: torch.Tensor, is_first_chunk: bool + ) -> torch.Tensor: + """Decode one latent window into its corresponding video frames.""" + hidden_states = self.conv_in(hidden_states) + + hidden_states = self.mid_block(hidden_states) + + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states, is_first_chunk=is_first_chunk) + + if not self.use_unpatchify: + hidden_states = self.norm_up_1(hidden_states) + hidden_states = self.conv_act(hidden_states) + + hidden_states_array = [] + for t in range(hidden_states.shape[2]): + h = self.upsampler2d_1(hidden_states[:, :, t, :, :]) + hidden_states_array.append(h) + hidden_states = torch.stack(hidden_states_array, dim=2) + + # RMSNorm + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + 1e-8) + hidden_states = hidden_states.to(input_dtype) + + hidden_states = self.conv_act(hidden_states) + + hidden_states = self.conv_out(hidden_states) + + if self.use_unpatchify: + hidden_states = _turbo_unpatchify(hidden_states, self.patch_size) + + return hidden_states + + +class TurboVAED(ModelMixin, ConfigMixin): + + @register_to_config + def __init__( + self, + in_channels: int = 3, # useless arg for compatibility, we only use latent channels + out_channels: int = 3, + latent_channels: int = 128, + decoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + decoder_layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4), + decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False), + patch_size: int = 4, + patch_size_t: int = 1, + resnet_norm_eps: float = 1e-6, + scaling_factor: float = 1.0, + decoder_causal: bool = False, + decoder_is_dw_conv: tuple[bool, ...] = (False, False, False, False, False), + decoder_dw_kernel_size: int = 3, + decoder_spatio_only: tuple[bool, ...] = (False, False, False, False), + first_chunk_size: int = 3, + step_size: int = 5, + spatial_compression_ratio: int = 16, + temporal_compression_ratio: int = 4, + use_unpatchify: bool = False, + **kwargs, # absorb unused config keys (e.g. training-only params) + ): + """Build the Turbo VAE decoder and its latent normalization buffers.""" + super().__init__() + + self.decoder = TurboVAEDDecoder3d( + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=decoder_block_out_channels, + spatio_temporal_scaling=decoder_spatio_temporal_scaling, + layers_per_block=decoder_layers_per_block, + patch_size=patch_size, + patch_size_t=patch_size_t, + resnet_norm_eps=resnet_norm_eps, + is_causal=decoder_causal, + decoder_is_dw_conv=decoder_is_dw_conv, + decoder_dw_kernel_size=decoder_dw_kernel_size, + spatio_only=decoder_spatio_only, + use_unpatchify=use_unpatchify, + ) + + self.first_chunk_size = first_chunk_size + self.step_size = step_size + print( + f"[Turbo VAE Decoder] init with first_chunk_size: {first_chunk_size}, step_size: {step_size}" + ) + + self.spatial_compression_ratio = spatial_compression_ratio + self.temporal_compression_ratio = temporal_compression_ratio + + self.z_dim = latent_channels + self.register_buffer( + "mean", + torch.tensor( + [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ], + dtype=torch.float32, + ), + persistent=False, + ) + self.register_buffer( + "std", + torch.tensor( + [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ], + dtype=torch.float32, + ), + persistent=False, + ) + self.register_buffer( + "inverse_std", + 1.0 / self.std, + persistent=False, + ) + + def _apply(self, fn, recurse: bool = True): + """Move normalization buffers while preserving their FP32 values.""" + mean = self.mean.float().clone() + std = self.std.float().clone() + module = super()._apply(fn, recurse=recurse) + target_device = self.mean.device + self.mean = mean.to(device=target_device) + self.std = std.to(device=target_device) + self.inverse_std = (1.0 / std).to(device=target_device) + return module + + def _sliding_window_decode(self, z, output_offload=False): + """Decode latent windows with one-frame context and join their outputs.""" + z_dtype = z.dtype + z = z / self.inverse_std.view(1, self.z_dim, 1, 1, 1) + self.mean.view( + 1, self.z_dim, 1, 1, 1 + ) + z = z.to(z_dtype) + + first_chunk_size = self.first_chunk_size + step = self.step_size + + # Context mapping: 1 latent frame context -> temporal_compression_ratio pixel frames overlap + num_overlap_pixel_frames = 1 * self.temporal_compression_ratio + + _, _, num_frames, _, _ = z.shape + + # 1. Pad frames to satisfy chunking requirements + # The total number of frames must follow the formula: + # num_frames = first_chunk_size + n * step_size + num_padding_frames = 0 + + if num_frames < first_chunk_size: + # if input is shorter than first_chunk_size + num_padding_frames = first_chunk_size - num_frames + elif (num_frames - first_chunk_size) % step != 0: + num_padding_frames = step - (num_frames - first_chunk_size) % step + + if num_padding_frames > 0: + z = torch.cat( + [z, z[:, :, -1:].repeat(1, 1, num_padding_frames, 1, 1)], dim=2 + ) + num_frames = num_frames + num_padding_frames + + # 2. Decode with overlapping windows + # Collect chunks on CPU to avoid GPU OOM for high resolution (e.g., 1080P) when output_offload=True + out_chunks = [] + + if num_frames == first_chunk_size: + # if only one chunk, decode directly + out = self.decoder(z, is_first_chunk=True) + out_chunks.append(out.cpu() if output_offload else out) + del out + else: + # first chunk: attach the right frame + out = self.decoder( + z[:, :, 0 : first_chunk_size + 1, :, :], is_first_chunk=True + ) + out = out[:, :, :-num_overlap_pixel_frames] + out_chunks.append(out.cpu() if output_offload else out) + del out + + # middle chunk: attach the left and right frames + # last chunk: attach the left frame + for i in range(first_chunk_size, num_frames, step): + is_last_chunk = i + step == num_frames + left = i - 1 + right = i + step + 1 if not is_last_chunk else i + step + + assert left >= 0 and right <= num_frames, ( + f"left: {left}, right: {right}, num_frames: {num_frames}" + ) + + out_ = self.decoder(z[:, :, left:right, :, :], is_first_chunk=False) + + if is_last_chunk: + out_ = out_[:, :, num_overlap_pixel_frames:] + else: + out_ = out_[ + :, :, num_overlap_pixel_frames:-num_overlap_pixel_frames + ] + + out_chunks.append(out_.cpu() if output_offload else out_) + del out_ + + # Concatenate chunks (on CPU if output_offload, otherwise on GPU) + out = torch.cat(out_chunks, dim=2) + del out_chunks + + # 3. Remove padded frames + if num_padding_frames > 0: + out = out[:, :, : -num_padding_frames * self.temporal_compression_ratio] + + return out + + @magi_compile(dynamic_arg_dims=({"z": [3, 4]})) + def decode(self, z: torch.Tensor, output_offload: bool = False): + return self._sliding_window_decode(z, output_offload=output_offload) + + +def get_turbo_vaed( + config_path, ckpt_path, device="cuda", weight_dtype=torch.float32 +) -> TurboVAED: + """Load the published Turbo VAE configuration and decoder checkpoint.""" + with open(config_path, encoding="utf-8") as f: + config = json.load(f) + student = TurboVAED.from_config(config) + + ckpt = torch.load(ckpt_path, map_location="cpu") + + if "ema_state_dict" in ckpt: + state_dict = ckpt["ema_state_dict"] + elif "state_dict" in ckpt: + state_dict = ckpt["state_dict"] + else: + state_dict = ckpt + normalized_state_dict = {} + for key, value in state_dict.items(): + if key.startswith("module."): + normalized_state_dict[key[7:]] = value + else: + normalized_state_dict[key] = value + state_dict = normalized_state_dict + + missing, unexpected = student.load_state_dict(state_dict, strict=False) + if len(missing) > 0: + sample_key = next(iter(state_dict.keys())) + if not sample_key.startswith("decoder.") and not sample_key.startswith( + "encoder." + ): + print("Attempting to load state_dict directly into student.decoder...") + student.decoder.load_state_dict(state_dict, strict=False) + + student = student.to(device, dtype=weight_dtype) + student.eval() + student.requires_grad_(False) + return student + + +class Magi2TurboVAEModel(TurboVAED): + """Expose the distilled Turbo VAE through FastVideo's standard VAE loader.""" + + def __init__(self, config: Magi2TurboVAEConfig) -> None: + """Build the published decoder and optionally load its source checkpoint.""" + architecture = config.arch_config + model_kwargs = { + "in_channels": architecture.in_channels, + "out_channels": architecture.out_channels, + "latent_channels": architecture.latent_channels, + "decoder_block_out_channels": tuple( + architecture.decoder_block_out_channels + ), + "decoder_layers_per_block": tuple( + architecture.decoder_layers_per_block + ), + "decoder_spatio_temporal_scaling": tuple( + architecture.decoder_spatio_temporal_scaling + ), + "patch_size": architecture.patch_size, + "patch_size_t": architecture.patch_size_t, + "resnet_norm_eps": architecture.resnet_norm_eps, + "scaling_factor": architecture.scaling_factor, + "decoder_causal": architecture.decoder_causal, + "decoder_is_dw_conv": tuple(architecture.decoder_is_dw_conv), + "decoder_dw_kernel_size": architecture.decoder_dw_kernel_size, + "decoder_spatio_only": tuple(architecture.decoder_spatio_only), + "first_chunk_size": architecture.first_chunk_size, + "step_size": architecture.step_size, + "spatial_compression_ratio": architecture.spatial_compression_ratio, + "temporal_compression_ratio": architecture.temporal_compression_ratio, + "use_unpatchify": architecture.use_unpatchify, + } + if config.config_path: + with open(config.config_path, encoding="utf-8") as config_file: + published_config = json.load(config_file) + model_kwargs.update( + { + key: published_config[key] + for key in model_kwargs + if key in published_config + } + ) + super().__init__(**model_kwargs) + self.fastvideo_config = config + if config.checkpoint_path: + self._load_source_checkpoint(config) + + def _load_source_checkpoint(self, config: Magi2TurboVAEConfig) -> None: + """Load EMA decoder tensors for direct official-checkpoint parity tests.""" + checkpoint = torch.load( + config.checkpoint_path, + map_location="cpu", + weights_only=True, + ) + if "ema_state_dict" in checkpoint: + state_dict = checkpoint["ema_state_dict"] + elif "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + else: + state_dict = checkpoint + state_dict = { + name.removeprefix("module."): tensor + for name, tensor in state_dict.items() + if name.removeprefix("module.").startswith("decoder.") + } + incompatible = self.load_state_dict(state_dict, strict=True) + if incompatible.missing_keys or incompatible.unexpected_keys: + raise RuntimeError( + "Turbo VAE strict load failed: " + f"missing={incompatible.missing_keys}, " + f"unexpected={incompatible.unexpected_keys}" + ) + try: + weight_dtype = getattr(torch, config.pretrained_dtype) + except AttributeError as error: + raise ValueError( + f"Unsupported Turbo VAE dtype: {config.pretrained_dtype!r}" + ) from error + target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.to(device=target_device, dtype=weight_dtype) + self.requires_grad_(False) + + +EntryClass = Magi2TurboVAEModel diff --git a/fastvideo/models/vaes/magi2_wan_loader.py b/fastvideo/models/vaes/magi2_wan_loader.py new file mode 100644 index 000000000..6ea0753c8 --- /dev/null +++ b/fastvideo/models/vaes/magi2_wan_loader.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 SandAI. All Rights Reserved. +"""Strict source-checkpoint loading for the MAGI-2 Wan image encoder.""" + +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fastvideo.configs.models.vaes.magi2_wanvae import Magi2WanVAEConfig + + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """Apply left-padded temporal convolution with an optional feature cache.""" + + def __init__(self, *args, **kwargs): + """Store causal padding separately from the underlying convolution.""" + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + @torch.compile + def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + """Convolve one temporal chunk after prepending the prior feature cache.""" + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return super().forward(x) + + +class RMSNorm(nn.Module): + """Normalize channel vectors with the Wan encoder's root-mean-square rule.""" + + def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False): + """Create the broadcastable scale and optional bias parameters.""" + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + @torch.compile + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Normalize the channel dimension and apply the learned scale.""" + dim = 1 if self.channel_first else -1 + return F.normalize(x, dim=dim) * self.scale * self.gamma + self.bias + + +class Resample(nn.Module): + """Downsample spatial dimensions and selected temporal transitions.""" + + def __init__(self, dim: int, mode: str): + """Build the spatial convolution and optional temporal convolution.""" + if mode not in ("downsample2d", "downsample3d"): + raise ValueError(f"Unsupported MAGI-2 Wan encoder resampling mode: {mode}") + super().__init__() + self.dim = dim + self.mode = mode + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2)), + ) + if mode == "downsample3d": + self.time_conv = CausalConv3d( + dim, + dim, + (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0), + ) + + @torch.compile + def forward( + self, + x: torch.Tensor, + feat_cache: list[torch.Tensor | None] | None = None, + feat_idx: list[int] = [0], # noqa: B006 - compiled nested calls share this cache cursor. + ) -> torch.Tensor: + """Downsample one chunk and update its temporal feature cache.""" + batch_size, channels, time, height, width = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size, t=time) + + if self.mode == "downsample3d" and feat_cache is not None: + index = feat_idx[0] + if feat_cache[index] is None: + feat_cache[index] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], 2)) + feat_cache[index] = cache_x + feat_idx[0] += 1 + + return x + + +class ResidualBlock(nn.Module): + """Apply two cached causal convolutions and an additive shortcut.""" + + def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0): + """Build the normalization, convolution, and shortcut modules.""" + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.residual = nn.Sequential( + RMSNorm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMSNorm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + @torch.compile + def forward( + self, + x: torch.Tensor, + feat_cache: list[torch.Tensor | None] | None = None, + feat_idx: list[int] = [0], # noqa: B006 - compiled nested calls share this cache cursor. + ) -> torch.Tensor: + """Run the residual path while preserving the causal convolution caches.""" + shortcut = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + index = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[index] is not None: + cache_x = torch.cat( + [feat_cache[index][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = layer(x, feat_cache[index]) + feat_cache[index] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + shortcut + + +class AttentionBlock(nn.Module): + """Apply single-head spatial self-attention independently per frame.""" + + def __init__(self, dim: int): + """Build normalization, fused query-key-value projection, and output projection.""" + super().__init__() + self.dim = dim + self.norm = RMSNorm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + @torch.compile + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Attend over the spatial positions of each frame and add the residual.""" + identity = x + batch_size, channels, time, height, width = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + query, key, value = ( + self.to_qkv(x) + .reshape(batch_size * time, 1, channels * 3, -1) + .permute(0, 1, 3, 2) + .contiguous() + .chunk(3, dim=-1) + ) + x = F.scaled_dot_product_attention(query, key, value) + x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) + x = self.proj(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=time) + return x + identity + + +def patchify(x: torch.Tensor, patch_size: int) -> torch.Tensor: + """Move each spatial patch into the channel dimension.""" + if patch_size == 1: + return x + if x.dim() == 4: + return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + if x.dim() == 5: + return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size) + raise ValueError(f"Invalid MAGI-2 Wan image shape: {tuple(x.shape)}") + + +class AvgDown3D(nn.Module): + """Average channel groups after spatial and temporal rearrangement.""" + + def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1): + """Record the factors and channel-group size for the shortcut path.""" + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + if in_channels * self.factor % out_channels != 0: + raise ValueError("MAGI-2 Wan shortcut channels must divide evenly") + self.group_size = in_channels * self.factor // out_channels + + @torch.compile + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Rearrange a video into channel groups and average each group.""" + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + x = F.pad(x, (0, 0, 0, 0, pad_t, 0)) + batch_size, channels, time, height, width = x.shape + x = x.view( + batch_size, + channels, + time // self.factor_t, + self.factor_t, + height // self.factor_s, + self.factor_s, + width // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + batch_size, + channels * self.factor, + time // self.factor_t, + height // self.factor_s, + width // self.factor_s, + ) + x = x.view( + batch_size, + self.out_channels, + self.group_size, + time // self.factor_t, + height // self.factor_s, + width // self.factor_s, + ) + return x.mean(dim=2) + + +class DownResidualBlock(nn.Module): + """Combine residual encoder blocks with an averaged downsample shortcut.""" + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float, + mult: int, + temporal_downsample: bool = False, + down_flag: bool = False, + ): + """Build the main and shortcut paths for one encoder resolution.""" + super().__init__() + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temporal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + downsamples: list[nn.Module] = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + if down_flag: + mode = "downsample3d" if temporal_downsample else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + self.downsamples = nn.Sequential(*downsamples) + + @torch.compile + def forward( + self, + x: torch.Tensor, + feat_cache: list[torch.Tensor | None] | None = None, + feat_idx: list[int] = [0], # noqa: B006 - compiled nested calls share this cache cursor. + ) -> torch.Tensor: + """Run one resolution block and combine its two paths.""" + shortcut_input = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + return x + self.avg_shortcut(shortcut_input) + + +class Encoder3d(nn.Module): + """Encode patchified conditioning images into Gaussian moments.""" + + def __init__( + self, + dim: int, + z_dim: int, + dim_mult: tuple[int, ...], + num_res_blocks: int, + attn_scales: tuple[float, ...], + temporal_downsample: tuple[bool, ...], + dropout: float, + ): + """Build the published MAGI-2 Wan image-encoder hierarchy.""" + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temporal_downsample = temporal_downsample + + dims = [dim * multiplier for multiplier in [1, *dim_mult]] + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + downsamples: list[nn.Module] = [] + for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): + temporal_downsample_flag = temporal_downsample[index] if index < len(temporal_downsample) else False + downsamples.append( + DownResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temporal_downsample=temporal_downsample_flag, + down_flag=index != len(dim_mult) - 1, + ) + ) + self.downsamples = nn.Sequential(*downsamples) + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + self.head = nn.Sequential( + RMSNorm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + @torch.compile + def forward( + self, + x: torch.Tensor, + feat_cache: list[torch.Tensor | None] | None = None, + feat_idx: list[int] = [0], # noqa: B006 - compiled nested calls share this cache cursor. + ) -> torch.Tensor: + """Encode one temporal chunk while updating every causal cache slot.""" + if feat_cache is not None: + index = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[index] is not None: + cache_x = torch.cat( + [feat_cache[index][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = self.conv1(x, feat_cache[index]) + feat_cache[index] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + x = layer(x, feat_cache, feat_idx) if feat_cache is not None else layer(x) + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + index = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[index] is not None: + cache_x = torch.cat( + [feat_cache[index][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = layer(x, feat_cache[index]) + feat_cache[index] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_causal_convolutions(model: nn.Module) -> int: + """Count cache slots required by the encoder's causal convolutions.""" + return sum(isinstance(module, CausalConv3d) for module in model.modules()) + + +class Magi2WanImageEncoder(nn.Module): + """Encode MAGI-2 I2V reference images into normalized 48-channel latents.""" + + def __init__(self, config: Magi2WanVAEConfig): + """Build the encoder and quantization convolution with source checkpoint names.""" + super().__init__() + self.z_dim = config.z_dim + self.patch_size = config.patch_size + self.encoder = Encoder3d( + dim=config.base_dim, + z_dim=config.z_dim * 2, + dim_mult=tuple(config.dim_mult), + num_res_blocks=config.num_res_blocks, + attn_scales=tuple(config.attn_scales), + temporal_downsample=tuple(config.temperal_downsample), + dropout=config.dropout, + ) + self.conv1 = CausalConv3d(config.z_dim * 2, config.z_dim * 2, 1) + self.register_buffer("mean", torch.tensor(config.latents_mean, dtype=torch.float32), persistent=False) + self.register_buffer("std", torch.tensor(config.latents_std, dtype=torch.float32), persistent=False) + self.register_buffer("inverse_std", 1.0 / self.std, persistent=False) + self.clear_cache() + + def clear_cache(self) -> None: + """Reset every causal feature cache before and after an image batch.""" + self._enc_conv_num = count_causal_convolutions(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map: list[torch.Tensor | None] = [None] * self._enc_conv_num + + def set_normalization(self, mean: tuple[float, ...], std: tuple[float, ...]) -> None: + """Create normalization tensors on the encoder device with FP32 arithmetic.""" + device = self.conv1.weight.device + self.mean = torch.tensor(mean, dtype=torch.float32, device=device) + self.std = torch.tensor(std, dtype=torch.float32, device=device) + self.inverse_std = 1.0 / self.std + + def encode(self, video: torch.Tensor) -> torch.Tensor: + """Encode FP32 video produced by the pipeline's BF16 image round trip.""" + if video.dtype != torch.float32: + raise TypeError(f"MAGI-2 Wan encoder input must be torch.float32, received {video.dtype}") + self.clear_cache() + video = patchify(video, patch_size=self.patch_size) + time = video.shape[2] + chunk_count = 1 + (time - 1) // 4 + for chunk_index in range(chunk_count): + self._enc_conv_idx = [0] + if chunk_index == 0: + encoded = self.encoder( + video[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + else: + encoded_chunk = self.encoder( + video[:, :, 1 + 4 * (chunk_index - 1) : 1 + 4 * chunk_index, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + encoded = torch.cat([encoded, encoded_chunk], 2) + latent_mean, _ = self.conv1(encoded).chunk(2, dim=1) + latent = (latent_mean - self.mean.view(1, self.z_dim, 1, 1, 1)) * self.inverse_std.view( + 1, + self.z_dim, + 1, + 1, + 1, + ) + self.clear_cache() + return latent.float() + + def forward(self, video: torch.Tensor) -> torch.Tensor: + """Encode one pipeline-ready reference-image tensor.""" + return self.encode(video) + + +def _is_decoder_tensor(source_name: str) -> bool: + """Identify checkpoint tensors that belong only to video decoding.""" + return source_name.startswith("decoder.") or source_name.startswith("conv2.") + + +def load_magi2_wan_image_encoder( + checkpoint_path: str | Path, + device: torch.device | str, +) -> Magi2WanImageEncoder: + """Build and strictly load every MAGI-2 Wan encoder and quant tensor.""" + config = Magi2WanVAEConfig() + with torch.device("meta"): + model = Magi2WanImageEncoder(config) + target_state = model.state_dict() + source_state = torch.load( + Path(checkpoint_path), + map_location="cpu", + weights_only=True, + mmap=True, + ) + encoder_state: dict[str, torch.Tensor] = {} + for source_name, source_tensor in source_state.items(): + if _is_decoder_tensor(source_name): + continue + if source_name not in target_state: + raise KeyError(f"Unexpected MAGI-2 Wan encoder checkpoint tensor: {source_name}") + target_tensor = target_state[source_name] + if source_tensor.shape != target_tensor.shape or source_tensor.dtype != target_tensor.dtype: + raise ValueError( + f"MAGI-2 Wan encoder tensor metadata differs for {source_name}: " + f"source={tuple(source_tensor.shape)}/{source_tensor.dtype}, " + f"target={tuple(target_tensor.shape)}/{target_tensor.dtype}" + ) + encoder_state[source_name] = source_tensor + missing_names = sorted(set(target_state) - set(encoder_state)) + if missing_names: + raise RuntimeError(f"MAGI-2 Wan encoder checkpoint is missing tensors: {missing_names}") + incompatible = model.load_state_dict(encoder_state, strict=True, assign=True) + if incompatible.missing_keys or incompatible.unexpected_keys: + raise RuntimeError( + "MAGI-2 Wan encoder strict load failed: " + f"missing={incompatible.missing_keys}, unexpected={incompatible.unexpected_keys}" + ) + model.mean = torch.empty(len(config.latents_mean), dtype=torch.float32) + model.std = torch.empty(len(config.latents_std), dtype=torch.float32) + model.inverse_std = torch.empty(len(config.latents_std), dtype=torch.float32) + model.to(device=device, dtype=torch.float32) + model.set_normalization(config.latents_mean, config.latents_std) + model.requires_grad_(False) + return model.eval() + + +__all__ = ["Magi2WanImageEncoder", "load_magi2_wan_image_encoder"] diff --git a/fastvideo/pipelines/basic/magi2/__init__.py b/fastvideo/pipelines/basic/magi2/__init__.py new file mode 100644 index 000000000..51e8bba63 --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MAGI-2 Preview pipeline package.""" diff --git a/fastvideo/pipelines/basic/magi2/magi2_pipeline.py b/fastvideo/pipelines/basic/magi2/magi2_pipeline.py new file mode 100644 index 000000000..b8310fca5 --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/magi2_pipeline.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo pipeline for MAGI-2 Preview text-to-video and image-to-video.""" + +from __future__ import annotations + +import os +import random +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.models.dits.magi2_runtime import psm +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, +) +from fastvideo.pipelines.basic.magi2.stages import ( + Magi2AudioDecodingStage, + Magi2DataProxyConfig, + Magi2InputValidationStage, + Magi2LatentPreparationStage, + Magi2LatentSavingStage, + Magi2PreviewDenoisingStage, + Magi2ReferenceImageStage, + Magi2RefinerDataProxyConfig, + Magi2RefinerStage, + Magi2TextEncodingStage, + Magi2VideoDecodingStage, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase + + +def _configure_deterministic_kernels(seed: int) -> None: + """Enable the official deterministic attention, MoE, and PyTorch paths.""" + os.environ["MAGI2_DETERMINISTIC"] = "1" + os.environ["MAGI_ATTENTION_DETERMINISTIC_MODE"] = "1" + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + + +class Magi2Pipeline(ComposedPipelineBase): + """Generate a 10-second video and stereo audio with MAGI-2 Preview.""" + + is_video_pipeline = True + _required_config_modules = [ + "transformer", + "transformer_2", + "text_encoder", + "image_encoder", + "vae", + "audio_vae", + "scheduler", + ] + + def __init__( + self, + model_path: str, + fastvideo_args: FastVideoArgs | TrainingArgs, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None, + ) -> None: + """Configure deterministic kernels before distributed model imports.""" + deterministic_environment = os.environ.get( + "MAGI2_DETERMINISTIC", + "0", + ) == "1" + if fastvideo_args.deterministic or deterministic_environment: + _configure_deterministic_kernels(seed=42) + fastvideo_args.dit_layerwise_offload = False + fastvideo_args.dit_cpu_offload = True + super().__init__( + model_path=model_path, + fastvideo_args=fastvideo_args, + required_config_modules=required_config_modules, + loaded_modules=loaded_modules, + ) + + def load_modules( + self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Strictly load official components while keeping decoder copies on one rank.""" + self._load_config(self.model_path) + from fastvideo.configs.models.encoders import Magi2Qwen35Config + from fastvideo.configs.models.vaes import Magi2TurboVAEConfig + from fastvideo.models.dits.magi2_loader import ( + load_magi2_preview_model, + load_magi2_refiner_model, + ) + from fastvideo.models.dits.magi2_runtime.fastvideo_parallel import ( + bind_fastvideo_parallel_state, + ) + from fastvideo.models.encoders.qwen3_5 import Magi2Qwen35TextEncoder + from fastvideo.models.vaes.magi2_audio_vae import load_magi2_audio_vae + from fastvideo.models.vaes.magi2_turbo_vae import Magi2TurboVAEModel + from fastvideo.models.vaes.magi2_wan_loader import ( + load_magi2_wan_image_encoder, + ) + + bind_fastvideo_parallel_state() + checkpoint_root = Path(self.model_path) + pipeline_config = fastvideo_args.pipeline_config + provided_modules = loaded_modules or {} + modules: dict[str, Any] = {} + + modules["transformer"] = provided_modules.get("transformer") + if modules["transformer"] is None: + modules["transformer"] = load_magi2_preview_model( + str(checkpoint_root / "transformer"), + pipeline_config.dit_config, + "cpu", + ) + modules["transformer_2"] = provided_modules.get("transformer_2") + if modules["transformer_2"] is None: + modules["transformer_2"] = load_magi2_refiner_model( + str(checkpoint_root / "transformer_2"), + pipeline_config.refiner_dit_config, + "cpu", + ) + + is_decode_rank = psm.is_group_first_rank("cp") + modules["text_encoder"] = provided_modules.get("text_encoder") + modules["image_encoder"] = provided_modules.get("image_encoder") + modules["vae"] = provided_modules.get("vae") + modules["audio_vae"] = provided_modules.get("audio_vae") + if is_decode_rank and modules["text_encoder"] is None: + modules["text_encoder"] = Magi2Qwen35TextEncoder.from_pretrained_local( + str(checkpoint_root / "text_encoder"), + Magi2Qwen35Config(), + torch.bfloat16, + torch.device("cpu"), + ) + if is_decode_rank and modules["image_encoder"] is None: + modules["image_encoder"] = load_magi2_wan_image_encoder( + checkpoint_root / "image_encoder" / "Wan2.2_VAE.pth", + "cpu", + ) + if is_decode_rank and modules["vae"] is None: + turbo_config = Magi2TurboVAEConfig( + config_path=str( + checkpoint_root + / "vae" + / "TurboV3-Wan22-TinyShallow_7_7.json" + ), + checkpoint_path=str(checkpoint_root / "vae" / "checkpoint.ckpt"), + pretrained_dtype="bfloat16", + ) + modules["vae"] = Magi2TurboVAEModel(turbo_config).to("cpu") + torch.cuda.empty_cache() + if is_decode_rank and modules["audio_vae"] is None: + modules["audio_vae"] = load_magi2_audio_vae( + checkpoint_root / "audio_vae", + "cpu", + ) + modules["scheduler"] = provided_modules.get( + "scheduler", + FlowUniPCMultistepScheduler(), + ) + return modules + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Compose conditioning, denoising, refinement, and decoding stages.""" + pipeline_config = fastvideo_args.pipeline_config + self.add_stage( + "input_validation_stage", + Magi2InputValidationStage( + output_frames=pipeline_config.output_frames, + output_height=pipeline_config.output_height, + output_width=pipeline_config.output_width, + output_fps=pipeline_config.output_fps, + ), + ) + self.add_stage( + "reference_image_stage", + Magi2ReferenceImageStage( + image_encoder=self.get_module("image_encoder"), + preview_height=pipeline_config.preview_height, + preview_width=pipeline_config.preview_width, + ), + ) + self.add_stage( + "text_encoding_stage", + Magi2TextEncodingStage(self.get_module("text_encoder")), + ) + self.add_stage( + "latent_preparation_stage", + Magi2LatentPreparationStage( + video_channels=pipeline_config.preview_video_channels, + video_length=pipeline_config.preview_video_length, + video_height=pipeline_config.preview_latent_height, + video_width=pipeline_config.preview_latent_width, + audio_length=pipeline_config.audio_latent_length, + audio_channels=pipeline_config.audio_channels, + ), + ) + self.add_stage( + "preview_denoising_stage", + Magi2PreviewDenoisingStage( + transformer=self.get_module("transformer"), + data_proxy_config=Magi2DataProxyConfig(), + flow_shift=pipeline_config.preview_flow_shift, + video_guidance_scale=pipeline_config.preview_video_guidance_scale, + audio_guidance_scale=pipeline_config.preview_audio_guidance_scale, + ), + ) + self.add_stage( + "refiner_stage", + Magi2RefinerStage( + transformer=self.get_module("transformer_2"), + data_proxy_config=Magi2RefinerDataProxyConfig( + t_patch_size=1, + patch_size=1, + frame_receptive_field=11, + spatial_rope_interpolation="extra", + text_offset=0, + coords_style="v1", + attn_config={ + "mode": "window", + "block_t_size": 8, + "block_size": 4, + "window": { + "level": "block", + "block_mode": "grid", + "block_t_radius": 2, + "block_h_radius": 2, + "block_w_radius": 2, + "win_size": 384, + "frame_receptive_field": -1, + "auto_range_merge": True, + "sparse_load": False, + "full_attn_layers": [], + }, + }, + magi2_refiner_condition_input="none", + ), + latent_height=pipeline_config.refiner_latent_height, + latent_width=pipeline_config.refiner_latent_width, + noise_index=pipeline_config.refiner_noise_index, + flow_shift=pipeline_config.refiner_flow_shift, + video_guidance_scale=pipeline_config.refiner_video_guidance_scale, + audio_guidance_scale=pipeline_config.refiner_audio_guidance_scale, + audio_channels=pipeline_config.audio_channels, + ), + ) + self.add_stage("latent_saving_stage", Magi2LatentSavingStage()) + self.add_stage( + "video_decoding_stage", + Magi2VideoDecodingStage(self.get_module("vae")), + ) + self.add_stage( + "audio_decoding_stage", + Magi2AudioDecodingStage(self.get_module("audio_vae")), + ) + + +EntryClass = Magi2Pipeline + +__all__ = ["Magi2Pipeline"] diff --git a/fastvideo/pipelines/basic/magi2/pipeline_configs.py b/fastvideo/pipelines/basic/magi2/pipeline_configs.py new file mode 100644 index 000000000..a9720312c --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/pipeline_configs.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline-local export for the MAGI-2 configuration class.""" + +from fastvideo.configs.pipelines.magi2 import Magi2PreviewPipelineConfig + +__all__ = ["Magi2PreviewPipelineConfig"] diff --git a/fastvideo/pipelines/basic/magi2/presets.py b/fastvideo/pipelines/basic/magi2/presets.py new file mode 100644 index 000000000..97c73837c --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/presets.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference preset for the published MAGI-2 Preview 1080p profile.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +MAGI2_NEGATIVE_PROMPT = ( + "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, " + "overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly " + "drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy " + "background, three legs, many people in the background, walking backwards, low quality, worst quality, poor " + "quality, noise, background noise, hiss, hum, buzz, crackle, static, compression artifacts, MP3 artifacts, " + "digital clipping, distortion, muffled, muddy, unclear, echo, reverb, room echo, over-reverberated, hollow " + "sound, distant, washed out, harsh, shrill, piercing, grating, tinny, thin sound, boomy, bass-heavy, flat EQ, " + "over-compressed, abrupt cut, jarring transition, sudden silence, looping artifact, music, instrumental, " + "sirens, alarms, crowd noise, unrelated sound effects, chaotic, disorganized, messy, cheap sound, emotionless, " + "flat delivery, deadpan, lifeless, apathetic, robotic, mechanical, monotone, flat intonation, undynamic, boring, " + "reading from a script, AI voice, synthetic, text-to-speech, TTS, insincere, fake emotion, exaggerated, overly " + "dramatic, melodramatic, cheesy, cringey, hesitant, unconfident, tired, weak voice, stuttering, stammering, " + "mumbling, slurred speech, mispronounced, bad articulation, lisp, vocal fry, creaky voice, mouth clicks, lip " + "smacks, wet mouth sounds, heavy breathing, audible inhales, plosives, p-pops, coughing, clearing throat, " + "sneezing, speaking too fast, rushed, speaking too slow, dragged out, unnatural pauses, awkward silence, choppy, " + "disconnected, multiple speakers, two voices, background talking, out of tune, off-key, autotune artifacts") + +_PREVIEW_STAGE = PresetStageSpec( + name="preview", + kind="denoising", + description="Joint video and audio preview denoising pass.", + allowed_overrides=frozenset({"num_inference_steps"}), +) + +_REFINER_STAGE = PresetStageSpec( + name="refiner", + kind="refinement", + description="Spatial-temporal 1080p latent refinement pass.", + allowed_overrides=frozenset({"num_inference_steps_sr"}), +) + +MAGI2_PREVIEW_1080P = InferencePreset( + name="magi2_preview_1080p", + version=1, + model_family="magi2", + description="MAGI-2 Preview text- or image-conditioned 10-second 1080p video with stereo audio.", + workload_type=None, + stage_schemas=(_PREVIEW_STAGE, _REFINER_STAGE), + defaults={ + "seed": 42, + "height": 1088, + "width": 1920, + "num_frames": 249, + "fps": 25, + "num_inference_steps": 100, + "num_inference_steps_sr": 5, + "guidance_scale": 5.0, + "negative_prompt": MAGI2_NEGATIVE_PROMPT, + }, +) + +ALL_PRESETS = (MAGI2_PREVIEW_1080P, ) + +__all__ = ["ALL_PRESETS", "MAGI2_NEGATIVE_PROMPT", "MAGI2_PREVIEW_1080P"] diff --git a/fastvideo/pipelines/basic/magi2/stages/__init__.py b/fastvideo/pipelines/basic/magi2/stages/__init__.py new file mode 100644 index 000000000..84df2ceba --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/__init__.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MAGI-2 pipeline components.""" + +from fastvideo.pipelines.basic.magi2.stages.audio_decoding import ( + MAGI2_AUDIO_TIME_STRETCH, + Magi2AudioDecodingStage, + decode_magi2_audio, + resample_magi2_audio, +) +from fastvideo.pipelines.basic.magi2.stages.conditioning import ( + Magi2ReferenceImageStage, + Magi2TextEncodingStage, +) +from fastvideo.pipelines.basic.magi2.stages.output import ( + Magi2LatentSavingStage, + Magi2VideoDecodingStage, +) +from fastvideo.pipelines.basic.magi2.stages.preview_data_proxy import ( + Magi2DataProxy, + Magi2DataProxyConfig, + ModelInput, + Modality, + VarlenHandler, +) +from fastvideo.pipelines.basic.magi2.stages.refiner_data_proxy import ( + Magi2RefinerDataProxy, + Magi2RefinerDataProxyConfig, + RefinerModelInput, +) +from fastvideo.pipelines.basic.magi2.stages.runtime import ( + Magi2InputValidationStage, + Magi2LatentPreparationStage, +) +from fastvideo.pipelines.basic.magi2.stages.sampling import ( + Magi2PreviewDenoisingStage, + Magi2RefinerStage, + ZeroSNRDDPMDiscretization, +) + +__all__ = [ + "MAGI2_AUDIO_TIME_STRETCH", + "Magi2AudioDecodingStage", + "Magi2DataProxy", + "Magi2DataProxyConfig", + "Magi2InputValidationStage", + "Magi2LatentPreparationStage", + "Magi2LatentSavingStage", + "Magi2PreviewDenoisingStage", + "Magi2ReferenceImageStage", + "Magi2RefinerStage", + "Magi2RefinerDataProxy", + "Magi2RefinerDataProxyConfig", + "Magi2TextEncodingStage", + "Magi2VideoDecodingStage", + "ModelInput", + "Modality", + "RefinerModelInput", + "VarlenHandler", + "ZeroSNRDDPMDiscretization", + "decode_magi2_audio", + "resample_magi2_audio", +] diff --git a/fastvideo/pipelines/basic/magi2/stages/audio_decoding.py b/fastvideo/pipelines/basic/magi2/stages/audio_decoding.py new file mode 100644 index 000000000..6c75861fa --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/audio_decoding.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio decoding for the MAGI-2 joint video-audio pipeline.""" + +from __future__ import annotations + +import numpy as np +import torch +from scipy.signal import resample + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.dits.magi2_runtime import psm +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + +MAGI2_AUDIO_TIME_STRETCH = 441.0 / 512.0 + + +def resample_magi2_audio( + sample_major_audio: np.ndarray, + time_stretching: float = MAGI2_AUDIO_TIME_STRETCH, +) -> np.ndarray: + """Resample sample-major stereo audio with MAGI-2's FFT interpolation.""" + output_length = int(sample_major_audio.shape[0] * time_stretching) + return resample(sample_major_audio, output_length) + + +@torch.inference_mode() +def decode_magi2_audio(audio_vae, sample_major_latent: torch.Tensor) -> np.ndarray: + """Decode a ``[latent samples, channels]`` tensor into stereo samples.""" + waveform = audio_vae.decode(sample_major_latent.T) + sample_major_audio = waveform.squeeze(0).T.cpu().numpy() + return resample_magi2_audio(sample_major_audio) + + +class Magi2AudioDecodingStage(PipelineStage): + """Decode preview audio latents on the context-parallel leader rank.""" + + def __init__(self, audio_vae) -> None: + """Store the published Stable Audio decoder.""" + super().__init__() + self.audio_vae = audio_vae + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Accept the audio latent that the joint denoising stage produced.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default stage verification record.""" + del batch, fastvideo_args + return VerificationResult() + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Decode and store sample-major audio on the output-producing rank.""" + if batch.audio_latents is None: + raise ValueError("MAGI-2 audio decoding requires audio_latents") + if not psm.is_group_first_rank("cp"): + return batch + device = torch.device("cuda", torch.cuda.current_device()) + self.audio_vae.to(device=device, dtype=torch.float32) + sample_major_latent = batch.audio_latents.squeeze(0).to(device) + batch.extra["audio"] = decode_magi2_audio( + self.audio_vae, + sample_major_latent, + ) + batch.extra["audio_sample_rate"] = int(self.audio_vae.sampling_rate) + if fastvideo_args.vae_cpu_offload: + self.audio_vae.to("cpu") + torch.cuda.empty_cache() + return batch + + +__all__ = [ + "MAGI2_AUDIO_TIME_STRETCH", + "Magi2AudioDecodingStage", + "decode_magi2_audio", + "resample_magi2_audio", +] diff --git a/fastvideo/pipelines/basic/magi2/stages/conditioning.py b/fastvideo/pipelines/basic/magi2/stages/conditioning.py new file mode 100644 index 000000000..751d7f9f8 --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/conditioning.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reference-image and Qwen3.5 conditioning stages for MAGI-2.""" + +from __future__ import annotations + +import json +from typing import Any + +import torch +from diffusers.utils import load_image +from diffusers.video_processor import VideoProcessor +from PIL import Image + +from fastvideo.distributed import get_sp_group +from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType +from fastvideo.models.dits.magi2_runtime import psm +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +def _resizepad(image: Image.Image, target_height: int, target_width: int) -> Image.Image: + """Fit an RGB image inside a white canvas with the official letterboxing rule.""" + width, height = image.size + if width <= 0 or height <= 0: + raise ValueError( + f"MAGI-2 received an invalid image size: width={width}, height={height}" + ) + scale = min(target_width / width, target_height / height) + resized_width = max(1, int(round(width * scale))) + resized_height = max(1, int(round(height * scale))) + resized = image.convert("RGB").resize( + (resized_width, resized_height), + resample=Image.Resampling.LANCZOS, + ) + canvas = Image.new("RGB", (target_width, target_height), (255, 255, 255)) + canvas.paste( + resized, + ( + (target_width - resized_width) // 2, + (target_height - resized_height) // 2, + ), + ) + return canvas + + +def _ensure_figure_reference(prompt: str) -> str: + """Address the single I2V image as ``
`` in a plain or JSON prompt.""" + try: + prompt_object = json.loads(prompt) + if not isinstance(prompt_object, dict): + raise ValueError("prompt JSON is not an object") + prompt_object["reference_layer"] = [ + "The first frame refers to
" + ] + return json.dumps(prompt_object, ensure_ascii=False) + except (json.JSONDecodeError, TypeError, ValueError): + return prompt + "reference_layer:The first frame refers to
" + + +class Magi2ReferenceImageStage(PipelineStage): + """Encode one I2V reference image on the context-parallel leader.""" + + def __init__( + self, + image_encoder: Any | None, + preview_height: int, + preview_width: int, + ) -> None: + """Store the Wan encoder and the published preview resolution.""" + super().__init__() + self.image_encoder = image_encoder + self.preview_height = preview_height + self.preview_width = preview_width + self.video_processor = VideoProcessor(vae_scale_factor=32) + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default validation record; ``forward`` gives precise errors.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default validation record for T2V and I2V outputs.""" + del batch, fastvideo_args + return VerificationResult() + + @torch.inference_mode() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Encode and broadcast the official FP32 Wan posterior mean for I2V.""" + if fastvideo_args.workload_type != WorkloadType.I2V: + batch.magi2_ref_image_feat = None + batch.magi2_ref_image_feat_len = None + batch.magi2_ref_image_special_tokens = None + return batch + + is_leader = psm.is_group_first_rank("cp") + payload: dict[str, torch.Tensor] | None = None + if is_leader: + if self.image_encoder is None: + raise RuntimeError("MAGI-2 I2V requires the Wan image encoder") + image = batch.pil_image + if image is None and batch.image_path is not None: + image = load_image(batch.image_path) + if not isinstance(image, Image.Image): + raise TypeError("MAGI-2 I2V requires one PIL image or image path") + + maximum_length = max(self.preview_height, self.preview_width) + if image.width > image.height: + target_width = maximum_length + target_height = int(image.height * maximum_length / image.width) + else: + target_height = maximum_length + target_width = int(image.width * maximum_length / image.height) + image = _resizepad(image, target_height, target_width) + device = torch.device("cuda", torch.cuda.current_device()) + image_tensor = self.video_processor.preprocess( + image, + height=target_height, + width=target_width, + ) + image_tensor = image_tensor.to(device=device, dtype=torch.bfloat16) + image_tensor = image_tensor.unsqueeze(2)[:, :3] + self.image_encoder.to(device=device, dtype=torch.float32) + reference_latent = self.image_encoder.encode(image_tensor.float()) + latent_height = int(reference_latent.shape[-2]) + latent_width = int(reference_latent.shape[-1]) + payload = { + "reference_latent": reference_latent.unsqueeze(1), + "reference_length": torch.tensor( + [[[latent_height, latent_width]]], + device=device, + dtype=torch.long, + ), + } + if fastvideo_args.image_encoder_cpu_offload: + self.image_encoder.to("cpu") + + broadcast_payload = get_sp_group().broadcast_tensor_dict(payload, src=0) + if broadcast_payload is None: + raise RuntimeError("MAGI-2 reference-image broadcast returned no payload") + batch.magi2_ref_image_feat = broadcast_payload["reference_latent"] + batch.magi2_ref_image_feat_len = broadcast_payload["reference_length"] + if not isinstance(batch.prompt, str): + raise TypeError("MAGI-2 accepts one prompt string per request") + batch.prompt = _ensure_figure_reference(batch.prompt) + if is_leader and fastvideo_args.image_encoder_cpu_offload: + torch.cuda.empty_cache() + return batch + + +class Magi2TextEncodingStage(PipelineStage): + """Encode positive and negative prompts on the context-parallel leader.""" + + def __init__(self, text_encoder: Any | None) -> None: + """Store the native Qwen3.5 encoder loaded from the release checkpoint.""" + super().__init__() + self.text_encoder = text_encoder + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record; prompt validation occurs in ``forward``.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record for the broadcast conditioning tensors.""" + del batch, fastvideo_args + return VerificationResult() + + @torch.inference_mode() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Produce Qwen3.5 hidden states and the optional figure-token embedding.""" + if not isinstance(batch.prompt, str): + raise TypeError("MAGI-2 accepts one prompt string per request") + if not isinstance(batch.negative_prompt, str): + raise TypeError("MAGI-2 requires one negative prompt string") + + is_leader = psm.is_group_first_rank("cp") + payload: dict[str, torch.Tensor | None] | None = None + if is_leader: + if self.text_encoder is None: + raise RuntimeError("MAGI-2 requires the Qwen3.5 text encoder") + device = torch.device("cuda", torch.cuda.current_device()) + self.text_encoder.to(device=device) + text_context = self.text_encoder.encode(batch.prompt).to(torch.bfloat16) + negative_context = self.text_encoder.encode(batch.negative_prompt).to( + torch.bfloat16 + ) + special_tokens = None + if batch.magi2_ref_image_feat is not None: + special_tokens = self.text_encoder.get_special_token( + batch.prompt, + ["
"], + text_context, + ).unsqueeze(0) + payload = { + "text_context": text_context, + "negative_context": negative_context, + "special_tokens": special_tokens, + } + if fastvideo_args.text_encoder_cpu_offload: + self.text_encoder.to("cpu") + + broadcast_payload = get_sp_group().broadcast_tensor_dict(payload, src=0) + if broadcast_payload is None: + raise RuntimeError("MAGI-2 text-conditioning broadcast returned no payload") + batch.magi2_text_context = broadcast_payload["text_context"] + batch.magi2_negative_context = broadcast_payload["negative_context"] + batch.magi2_ref_image_special_tokens = broadcast_payload["special_tokens"] + if is_leader and fastvideo_args.text_encoder_cpu_offload: + torch.cuda.empty_cache() + return batch + + +__all__ = ["Magi2ReferenceImageStage", "Magi2TextEncodingStage"] diff --git a/fastvideo/pipelines/basic/magi2/stages/output.py b/fastvideo/pipelines/basic/magi2/stages/output.py new file mode 100644 index 000000000..81ab1327d --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/output.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Post-refiner latent capture and Turbo VAE decoding for MAGI-2.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch + +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.dits.magi2_runtime import psm +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +class Magi2LatentSavingStage(PipelineStage): + """Write each post-refiner video latent on the output-producing rank.""" + + def __init__(self) -> None: + """Initialize the per-process sample counter used in latent filenames.""" + super().__init__() + self.sample_index = 0 + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record for the post-refiner latent.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record after the optional filesystem write.""" + del batch, fastvideo_args + return VerificationResult() + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Save ``latent_N.pt`` when ``MAGI2_SAVE_LATENT_PATH`` is configured.""" + del fastvideo_args + latent_directory = envs.MAGI2_SAVE_LATENT_PATH + if not latent_directory or not psm.is_group_first_rank("cp"): + return batch + if batch.latents is None: + raise ValueError("MAGI-2 latent saving requires a post-refiner latent") + output_directory = Path(latent_directory) + output_directory.mkdir(parents=True, exist_ok=True) + output_path = output_directory / f"latent_{self.sample_index}.pt" + torch.save(batch.latents.detach().cpu(), output_path) + self.sample_index += 1 + return batch + + +class Magi2VideoDecodingStage(PipelineStage): + """Decode the 1080p latent with Turbo VAE on one rank per video.""" + + performance_component_metric = "vae_decode_time_s" + + def __init__(self, turbo_vae: Any | None) -> None: + """Store the distilled Turbo VAE decoder loaded from ``ckpt/turbo_vae``.""" + super().__init__() + self.turbo_vae = turbo_vae + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record; ``forward`` validates the video latent.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record because non-leader ranks intentionally skip decode.""" + del batch, fastvideo_args + return VerificationResult() + + @torch.inference_mode() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Decode and store a CPU float video tensor in ``[B,C,T,H,W]`` layout.""" + if not psm.is_group_first_rank("cp"): + batch.output = None + return batch + if batch.latents is None: + raise ValueError("MAGI-2 Turbo VAE decoding requires a video latent") + if self.turbo_vae is None: + raise RuntimeError("MAGI-2 requires the distilled Turbo VAE decoder") + device = torch.device("cuda", torch.cuda.current_device()) + self.turbo_vae.to(device=device, dtype=torch.bfloat16) + decoder_input = batch.latents.squeeze(0).to(torch.bfloat16) + if decoder_input.dim() == 4: + decoder_input = decoder_input.unsqueeze(0) + decoded_video = self.turbo_vae.decode(decoder_input).float() + batch.output = decoded_video.mul(0.5).add(0.5).clamp(0, 1).cpu() + if fastvideo_args.vae_cpu_offload: + self.turbo_vae.to("cpu") + torch.cuda.empty_cache() + return batch + + +__all__ = ["Magi2LatentSavingStage", "Magi2VideoDecodingStage"] diff --git a/fastvideo/pipelines/basic/magi2/stages/preview_data_proxy.py b/fastvideo/pipelines/basic/magi2/stages/preview_data_proxy.py new file mode 100644 index 000000000..96e400346 --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/preview_data_proxy.py @@ -0,0 +1,892 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pack MAGI-2 preview video, audio, text, and image-conditioning tokens.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import IntEnum +from itertools import chain +from typing import Any, Literal + +import torch +import torch.distributed as dist +from einops import rearrange +from torch.nn import functional as F + +from fastvideo.models.dits.magi2_runtime.psm import psm + + +class Modality(IntEnum): + """Modality identifiers consumed by the MAGI-2 preview transformer.""" + + VIDEO = 0 + AUDIO = 1 + TEXT = 2 + TIME = 3 + + +@dataclass +class VarlenHandler: + """Variable-length attention boundaries for a packed token sequence.""" + + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + + +@dataclass(frozen=True) +class Magi2DataProxyConfig: + """Configuration that controls MAGI-2 preview token packing.""" + + t_patch_size: int = 1 + patch_size: int = 1 + spatial_rope_interpolation: Literal["inter", "extra"] = "extra" + add_time_token: bool = False + time_channel_dim: int = 64 + time_aligned_rope: bool = False + audio_latent_fps: float = 25.0 + time_pos_fps: float = 3.125 + vae_first_latent_is_image: bool = True + video_fps: float = 25.0 + + +@dataclass +class ModelInput: + """Multimodal tensors and lengths consumed by ``Magi2DataProxy``.""" + + x_t: torch.Tensor + audio_x_t: torch.Tensor + audio_feat_len: torch.Tensor | list[int] + txt_feat: torch.Tensor + txt_feat_len: torch.Tensor | list[int] + t: torch.Tensor + ref_audio_feat: torch.Tensor | None = None + ref_audio_feat_len: torch.Tensor | list[int] | None = None + ref_video_feat: torch.Tensor | None = None + ref_video_feat_len: torch.Tensor | list[int] | None = None + per_token_video_t: torch.Tensor | None = None + per_token_audio_t: torch.Tensor | None = None + ref_image_feat: torch.Tensor | None = None + ref_image_feat_len: torch.Tensor | None = None + ref_image_special_token_embedding: torch.Tensor | None = None + + +def _to_int(value: int | torch.Tensor) -> int: + """Convert a scalar tensor or Python integer into an integer.""" + if isinstance(value, torch.Tensor): + return int(value.detach().reshape(-1)[0].item()) + return int(value) + + +def _pad_cat( + tensors: list[torch.Tensor], + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Right-pad feature channels and concatenate token segments.""" + if not tensors: + return torch.empty(0, 0, device=device, dtype=dtype) + max_channel = max(tensor.shape[-1] for tensor in tensors) + return torch.cat( + [ + F.pad( + tensor.to(device=device, dtype=dtype), + (0, max_channel - tensor.shape[-1]), + ) + for tensor in tensors + ], + dim=0, + ) + + +def _segment_paint( + values: list[int], + seqlens: list[int], + dtype: torch.dtype, + device: torch.device, + output_size: int, +) -> torch.Tensor: + """Expand one modality identifier over each packed token segment.""" + mapping = torch.empty(output_size, dtype=dtype, device=device) + offset = 0 + for value, seqlen in zip(values, seqlens, strict=True): + mapping[offset:offset + seqlen] = int(value) + offset += seqlen + return mapping + + +def _seqlens2cu_seqlens( + seqlens: list[int], + device: torch.device | None = None, +) -> torch.Tensor: + """Convert sequence lengths into int32 cumulative boundaries.""" + seqlens_tensor = torch.tensor(seqlens, dtype=torch.int32, device=device) + return F.pad(torch.cumsum(seqlens_tensor, dim=0), (1, 0)) + + +def _ceil_div(dividend: int, divisor: int) -> int: + """Return integer ceiling division.""" + return (dividend + divisor - 1) // divisor + + +def _len_to_list(value: torch.Tensor) -> list[int]: + """Convert one packed reference-image grid shape into Python integers.""" + return [int(entry) for entry in value.detach().to(torch.long).reshape(-1).tolist()] + + +def _get_coords( + shape: tuple[int, int, int], + ref_feat_shape: tuple[int, int, int], + offset_thw: tuple[int, int, int] = (0, 0, 0), + device: torch.device | None = None, + dtype: torch.dtype = torch.float32, + time_positions: torch.Tensor | None = None, +) -> torch.Tensor: + """Build token coordinates with source-grid and reference-grid metadata.""" + if device is None: + device = torch.device("cpu") + ori_t, ori_h, ori_w = shape + ref_t, ref_h, ref_w = ref_feat_shape + offset_t, offset_h, offset_w = offset_thw + if time_positions is None: + time_range = torch.arange(ori_t, device=device, dtype=dtype) + offset_t + else: + time_range = time_positions.to(device=device, dtype=dtype) + offset_t + height_range = torch.arange(ori_h, device=device, dtype=dtype) + offset_h + width_range = torch.arange(ori_w, device=device, dtype=dtype) + offset_w + time_grid, height_grid, width_grid = torch.meshgrid( + time_range, + height_range, + width_range, + indexing="ij", + ) + coords_grid = torch.stack([time_grid, height_grid, width_grid], dim=-1) + coords_flat = coords_grid.reshape(-1, 3) + metadata = torch.tensor( + [ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], + device=device, + dtype=dtype, + ) + return torch.cat([coords_flat, metadata.expand(coords_flat.size(0), -1)], dim=-1) + + +def _sinusoidal_embedding_1d( + dim: int, + position: torch.Tensor, +) -> torch.Tensor: + """Encode normalized diffusion times into cosine-sine feature channels.""" + position = position.to(torch.float32) * 1000.0 + half = dim // 2 + frequencies = torch.exp( + -math.log(10000) + * torch.arange( + start=0, + end=half, + dtype=torch.float32, + device=position.device, + ) + / half + ) + arguments = position[:, None].float() * frequencies[None] + embedding = torch.cat([torch.cos(arguments), torch.sin(arguments)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], + dim=-1, + ) + return embedding + + +@dataclass +class SingleData: + """One sample's token segments and packing metadata.""" + + video_x_t: torch.Tensor + audio_x_t: torch.Tensor + audio_feat_len: int + txt_feat: torch.Tensor + txt_feat_len: int + t: int + h: int + w: int + patch_size: int + t_patch_size: int + spatial_rope_interpolation: Literal["inter", "extra"] + diffusion_t: torch.Tensor | None = None + per_token_video_t: torch.Tensor | None = None + per_token_audio_t: torch.Tensor | None = None + time_channel_dim: int = 0 + vae_first_latent_is_image: bool = True + video_fps: float = 25.0 + time_pos_fps: float = 3.125 + ref_image_feats: list[torch.Tensor] | None = None + ref_image_feat_lens: list[list[int]] | None = None + ref_image_special_tokens: list[torch.Tensor] | None = None + + def __post_init__(self) -> None: + """Trim variable-length segments and derive immutable packing sizes.""" + self.video_token_num = self.video_x_t.shape[0] + self.origin_audio_feat_len = self.audio_x_t.shape[0] + self.audio_x_t = self.audio_x_t[:self.audio_feat_len] + self.txt_feat = self.txt_feat[:self.txt_feat_len] + if self.per_token_audio_t is not None: + self.per_token_audio_t = self.per_token_audio_t[:self.audio_feat_len] + + self.ref_image_feats = self.ref_image_feats or [] + self.ref_image_feat_lens = self.ref_image_feat_lens or [] + self.ref_image_special_tokens = self.ref_image_special_tokens or [] + self.ref_image_token_nums = [ + int(math.prod(feat_len)) for feat_len in self.ref_image_feat_lens + ] + self.ref_image_feats = [ + feat[:token_num] + for feat, token_num in zip( + self.ref_image_feats, + self.ref_image_token_nums, + strict=True, + ) + ] + self.num_ref_images = len(self.ref_image_feats) + self.total_ref_image_feat_len = sum(self.ref_image_token_nums) + self.video_channel = self.video_x_t.shape[-1] + self.audio_channel = self.audio_x_t.shape[-1] + + @property + def device(self) -> torch.device: + """Return the device that owns this sample's packed tokens.""" + return self.video_x_t.device + + @property + def default_dtype(self) -> torch.dtype: + """Return the video-token dtype used for the packed sequence.""" + return self.video_x_t.dtype + + @property + def add_time_token(self) -> bool: + """Return whether this sample appends a standalone time token.""" + return self.diffusion_t is not None + + @property + def total_token_num(self) -> int: + """Return the complete token count for one sample.""" + total = self.video_token_num + self.audio_feat_len + self.txt_feat_len + total += self.total_ref_image_feat_len + self.num_ref_images + return total + (1 if self.add_time_token else 0) + + @property + def feat_to_cat(self) -> list[torch.Tensor]: + """Return token segments in transformer consumption order.""" + tensors = [self.video_x_t, self.audio_x_t, self.txt_feat] + for image_index in range(self.num_ref_images): + tensors.append(self.ref_image_special_tokens[image_index]) + tensors.append(self.ref_image_feats[image_index]) + if self.add_time_token: + assert self.diffusion_t is not None + tensors.append( + self.diffusion_t.to( + device=self.device, + dtype=self.default_dtype, + ).reshape(1, 1) + ) + return tensors + + @property + def token_sequence(self) -> torch.Tensor: + """Return channel-padded tokens for one sample.""" + return _pad_cat(self.feat_to_cat, self.device, self.default_dtype) + + @property + def modality_map_seqlens(self) -> tuple[list[int], list[int]]: + """Return segment lengths and modality identifiers in token order.""" + seqlens = [self.video_token_num, self.audio_feat_len, self.txt_feat_len] + modalities: list[int] = [ + int(Modality.VIDEO), + int(Modality.AUDIO), + int(Modality.TEXT), + ] + for image_index in range(self.num_ref_images): + seqlens.append(1) + modalities.append(int(Modality.TEXT)) + seqlens.append(self.ref_image_token_nums[image_index]) + modalities.append(int(Modality.VIDEO)) + if self.add_time_token: + seqlens.append(1) + modalities.append(int(Modality.TIME)) + return seqlens, modalities + + @property + def modality_mapping(self) -> torch.Tensor: + """Expand segment modality identifiers to one value per token.""" + seqlens, modalities = self.modality_map_seqlens + return _segment_paint( + modalities, + seqlens, + dtype=torch.int32, + device=self.device, + output_size=self.total_token_num, + ) + + def _default_coords( + self, + shape: tuple[int, int, int], + ref_feat_shape: tuple[int, int, int], + offset_thw: tuple[int, int, int] = (0, 0, 0), + time_positions: torch.Tensor | None = None, + ) -> torch.Tensor: + """Build coordinate rows on this sample's device and dtype.""" + return _get_coords( + shape, + ref_feat_shape, + offset_thw=offset_thw, + device=self.device, + dtype=self.default_dtype, + time_positions=time_positions, + ) + + @property + def coords_to_cat(self) -> list[torch.Tensor]: + """Return coordinate segments in the same order as feature segments.""" + t_steps = self.t // self.t_patch_size + h_steps = self.h // self.patch_size + w_steps = self.w // self.patch_size + if self.spatial_rope_interpolation == "inter": + video_h_ref, video_w_ref = 32, 32 + else: + video_h_ref, video_w_ref = h_steps, w_steps + + video_coords = self._default_coords( + (t_steps, h_steps, w_steps), + (t_steps, video_h_ref, video_w_ref), + ) + magic_audio_ref_t = (self.audio_feat_len - 1) // 8 + 1 + audio_coords = self._default_coords( + (self.audio_feat_len, 1, 1), + (magic_audio_ref_t // self.t_patch_size, 1, 1), + ) + coords = [ + video_coords, + audio_coords, + self._default_coords( + (self.txt_feat_len, 1, 1), + (1, 1, 1), + offset_thw=(-self.txt_feat_len, 0, 0), + ), + ] + + for image_index in range(self.num_ref_images): + token_len = self.ref_image_token_nums[image_index] + if len(self.ref_image_feat_lens[image_index]) >= 2: + image_h = int(self.ref_image_feat_lens[image_index][0]) + image_w = int(self.ref_image_feat_lens[image_index][1]) + else: + image_h = image_w = int(math.ceil(math.sqrt(token_len))) + time_offset = t_steps + 2 + image_index + coords.append( + torch.tensor( + [ + [ + time_offset, + -1, + -1, + 1, + image_h, + image_w, + 1, + image_h, + image_w, + ] + ], + device=self.device, + dtype=self.default_dtype, + ) + ) + coords.append( + self._default_coords( + (1, image_h, image_w), + (1, image_h, image_w), + offset_thw=(time_offset, 0, 0), + )[:token_len] + ) + + if self.add_time_token: + coords.append(self._default_coords((1, 1, 1), (1, 1, 1))[:1]) + return coords + + @property + def coords_mapping(self) -> torch.Tensor: + """Return one nine-value coordinate row per token.""" + return torch.cat(self.coords_to_cat, dim=0) + + @property + def time_token_sequence(self) -> torch.Tensor: + """Return per-token diffusion-time features in feature-segment order.""" + if self.time_channel_dim == 0: + return torch.empty(self.total_token_num, 0, device=self.device) + assert self.per_token_video_t is not None + assert self.per_token_audio_t is not None + time_parts = [ + self.per_token_video_t.squeeze(-1), + self.per_token_audio_t.squeeze(-1), + torch.zeros(self.txt_feat_len, device=self.device), + ] + for image_index in range(self.num_ref_images): + time_parts.append(torch.zeros(1, device=self.device)) + time_parts.append( + torch.zeros( + self.ref_image_token_nums[image_index], + device=self.device, + ) + ) + if self.add_time_token: + assert self.diffusion_t is not None + time_parts.append(self.diffusion_t.reshape(1).to(self.device)) + raw_time = torch.cat(time_parts, dim=0) + if self.time_channel_dim == 1: + return raw_time.unsqueeze(-1) + return _sinusoidal_embedding_1d(self.time_channel_dim, raw_time) + + +@dataclass +class SimplePackedData: + """A batch represented as adjacent variable-length sample sequences.""" + + items: list[SingleData] + + def __post_init__(self) -> None: + """Cache token-length summaries shared by attention and depacking.""" + if len(self.items) == 0: + raise ValueError("SimplePackedData must contain at least one item.") + self._total_token_num_list = [item.total_token_num for item in self.items] + self._total_token_num_sum = sum(self._total_token_num_list) + self._total_token_num_max = max(self._total_token_num_list) + self._total_token_cu_seqlens = _seqlens2cu_seqlens( + self._total_token_num_list + ) + + @property + def device(self) -> torch.device: + """Return the device used by every packed sample.""" + return self.items[0].device + + @property + def default_dtype(self) -> torch.dtype: + """Return the token dtype used by every packed sample.""" + return self.items[0].default_dtype + + @property + def token_sequence(self) -> torch.Tensor: + """Return every sample's feature segments as one packed tensor.""" + feature_segments = list( + chain.from_iterable(item.feat_to_cat for item in self.items) + ) + return _pad_cat(feature_segments, self.device, self.default_dtype) + + @property + def modality_mapping(self) -> torch.Tensor: + """Return one modality identifier per token across all samples.""" + seqlens: list[int] = [] + modalities: list[int] = [] + for item in self.items: + item_seqlens, item_modalities = item.modality_map_seqlens + seqlens.extend(item_seqlens) + modalities.extend(item_modalities) + return _segment_paint( + modalities, + seqlens, + dtype=torch.int32, + device=self.device, + output_size=self.total_token_num, + ) + + @property + def coords_mapping(self) -> torch.Tensor: + """Return coordinate rows across all packed samples.""" + coordinate_segments = list( + chain.from_iterable(item.coords_to_cat for item in self.items) + ) + return torch.cat(coordinate_segments, dim=0) + + @property + def time_token_sequence(self) -> torch.Tensor: + """Return diffusion-time features across all packed samples.""" + return torch.cat( + [item.time_token_sequence for item in self.items], + dim=0, + ) + + @property + def total_token_num(self) -> int: + """Return the total number of unpadded tokens.""" + return self._total_token_num_sum + + @property + def cu_seqlen(self) -> torch.Tensor: + """Return cumulative token boundaries for each packed sample.""" + return self._total_token_cu_seqlens.clone() + + @property + def max_seqlen(self) -> int: + """Return the longest unpadded sample sequence.""" + return self._total_token_num_max + + def __getitem__(self, index: int) -> SingleData: + """Return one packed sample descriptor.""" + return self.items[index] + + def depack_token_sequence( + self, + token_sequence: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Restore video and audio tensors from model-output token segments.""" + videos: list[torch.Tensor] = [] + audios: list[torch.Tensor] = [] + token_counts = [item.total_token_num for item in self.items] + for item, token_slice in zip( + self.items, + torch.split(token_sequence, token_counts, dim=0), + strict=True, + ): + video_flat = token_slice[:item.video_token_num, :item.video_channel] + output_channel = item.video_channel // ( + item.t_patch_size * item.patch_size * item.patch_size + ) + video = rearrange( + video_flat, + "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", + T=item.t // item.t_patch_size, + H=item.h // item.patch_size, + W=item.w // item.patch_size, + pT=item.t_patch_size, + pH=item.patch_size, + pW=item.patch_size, + C=output_channel, + ).contiguous() + audio = torch.zeros( + item.origin_audio_feat_len, + item.audio_channel, + device=token_sequence.device, + dtype=token_sequence.dtype, + ) + audio[:item.audio_feat_len] = token_slice[ + item.video_token_num:item.video_token_num + item.audio_feat_len, + :item.audio_channel, + ] + videos.append(video) + audios.append(audio) + return torch.stack(videos, dim=0), torch.stack(audios, dim=0) + + +class Magi2DataProxy: + """Pack MAGI-2 preview inputs and restore preview-model outputs.""" + + def __init__(self, config: Magi2DataProxyConfig) -> None: + self.config = config + self.patch_size = config.patch_size + self.t_patch_size = config.t_patch_size + self._saved_data: dict[str, Any] = {} + + def saved_for_output(self, **kwargs: Any) -> None: + """Save packing metadata that is required to depack model outputs.""" + self._saved_data.update(kwargs) + + def get_saved_data(self, key: str) -> Any: + """Return saved packing metadata by name.""" + return self._saved_data[key] + + def _reduce_max_token_num_for_ep_cp(self, total_token_num: int) -> int: + """Find the maximum token count across active expert and context groups.""" + if not (dist.is_available() and dist.is_initialized()): + return total_token_num + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + local_token_count = torch.tensor( + [total_token_num], + device=device, + dtype=torch.int64, + ) + maximum_counts = [local_token_count] + ep_size = psm.get_world_size("ep") + if ep_size > 1: + ep_maximum = local_token_count.clone() + dist.all_reduce( + ep_maximum, + op=dist.ReduceOp.MAX, + group=psm.get_parallel_group("ep"), + ) + maximum_counts.append(ep_maximum) + cp_size = psm.get_world_size("cp") + if cp_size > 1: + cp_maximum = local_token_count.clone() + dist.all_reduce( + cp_maximum, + op=dist.ReduceOp.MAX, + group=psm.get_parallel_group("cp"), + ) + maximum_counts.append(cp_maximum) + return max(int(maximum.item()) for maximum in maximum_counts) + + def _pad_for_ep_cp( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + modality_mapping: torch.Tensor, + time_token_sequence: torch.Tensor, + varlen_handler: VarlenHandler, + align_to: int = 48, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + VarlenHandler, + int, + ]: + """Pad distributed packed inputs to a shared multiple-of-48 length.""" + cp_size = max(1, psm.get_world_size("cp")) + ep_size = max(1, psm.get_world_size("ep")) + if cp_size <= 1 and ep_size <= 1: + return ( + x, + coords_mapping, + modality_mapping, + time_token_sequence, + varlen_handler, + 0, + ) + + total_token_num = x.shape[0] + target_size = self._reduce_max_token_num_for_ep_cp(total_token_num) + padded_size = _ceil_div(target_size, align_to) * align_to + pad_size = padded_size - total_token_num + if pad_size <= 0: + return ( + x, + coords_mapping, + modality_mapping, + time_token_sequence, + varlen_handler, + 0, + ) + + x = F.pad(x, (0, 0, 0, pad_size), value=0) + pad_coords = _get_coords( + shape=(pad_size, 1, 1), + ref_feat_shape=(pad_size, 1, 1), + offset_thw=(0, 0, 0), + device=x.device, + dtype=torch.float32, + ) + coords_mapping = torch.cat([coords_mapping, pad_coords], dim=0) + modality_mapping = F.pad( + modality_mapping, + (0, pad_size), + value=int(Modality.TEXT), + ) + if time_token_sequence.shape[-1] == 0: + time_token_sequence = time_token_sequence.new_zeros(padded_size, 0) + else: + time_token_sequence = F.pad( + time_token_sequence, + (0, 0, 0, pad_size), + value=0, + ) + padded_boundary = torch.tensor( + [padded_size], + device=x.device, + dtype=torch.int32, + ) + varlen_handler = VarlenHandler( + cu_seqlens_q=torch.cat( + [varlen_handler.cu_seqlens_q, padded_boundary] + ), + cu_seqlens_k=torch.cat( + [varlen_handler.cu_seqlens_k, padded_boundary] + ), + max_seqlen_q=max(varlen_handler.max_seqlen_q, pad_size), + max_seqlen_k=max(varlen_handler.max_seqlen_k, pad_size), + ) + return ( + x, + coords_mapping, + modality_mapping, + time_token_sequence, + varlen_handler, + pad_size, + ) + + def img2tokens(self, x_t: torch.Tensor) -> torch.Tensor: + """Extract non-overlapping 3D patches with channel-major features.""" + kernel_size = (self.t_patch_size, self.patch_size, self.patch_size) + if not all( + size >= kernel + for size, kernel in zip(x_t.shape[2:], kernel_size, strict=True) + ): + return torch.empty( + x_t.shape[0], + 0, + x_t.shape[1] * math.prod(kernel_size), + device=x_t.device, + dtype=x_t.dtype, + ) + patches = ( + x_t.unfold(2, self.t_patch_size, self.t_patch_size) + .unfold(3, self.patch_size, self.patch_size) + .unfold(4, self.patch_size, self.patch_size) + ) + batch_size, channels, t_steps, h_steps, w_steps, _, _, _ = patches.shape + return ( + patches.permute(0, 2, 3, 4, 1, 5, 6, 7) + .reshape(batch_size, t_steps * h_steps * w_steps, -1) + .contiguous() + ) + + def process_input( + self, + data: ModelInput, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + VarlenHandler, + torch.Tensor, + ]: + """Pack a batch without allowing conditional samples to attend each other.""" + batch_size, input_video_channel, video_t, video_h, video_w = data.x_t.shape + video_tokens = self.img2tokens(data.x_t) + audio_tokens = data.audio_x_t.contiguous() + text_tokens = data.txt_feat.contiguous() + + per_token_video_tokens = None + per_token_audio = None + if self.config.time_channel_dim > 0 and data.per_token_video_t is not None: + per_token_video_tokens = self.img2tokens(data.per_token_video_t)[:, :, :1] + per_token_audio = data.per_token_audio_t + + ref_image_data = data.ref_image_feat + ref_image_feat_len = data.ref_image_feat_len + ref_image_special_tokens = data.ref_image_special_token_embedding + has_ref_image = ( + ref_image_data is not None + and ref_image_data.ndim >= 5 + and ref_image_feat_len is not None + and ref_image_feat_len.ndim >= 2 + and ref_image_special_tokens is not None + and ref_image_special_tokens.ndim >= 3 + ) + num_ref_images = ref_image_data.shape[1] if has_ref_image else 0 + ref_device = text_tokens.device + ref_dtype = text_tokens.dtype + + items: list[SingleData] = [] + for batch_index in range(batch_size): + image_features: list[torch.Tensor] = [] + image_feature_lengths: list[list[int]] = [] + image_special_tokens: list[torch.Tensor] = [] + for image_index in range(num_ref_images): + assert ref_image_data is not None + assert ref_image_feat_len is not None + assert ref_image_special_tokens is not None + image_features.append( + self.img2tokens( + ref_image_data[batch_index, image_index].unsqueeze(0) + ).squeeze(0) + ) + image_feature_lengths.append( + _len_to_list(ref_image_feat_len[batch_index, image_index]) + ) + image_special_tokens.append( + ref_image_special_tokens[batch_index, image_index] + .to(device=ref_device, dtype=ref_dtype) + .unsqueeze(0) + ) + + items.append( + SingleData( + video_x_t=video_tokens[batch_index], + audio_x_t=audio_tokens[batch_index], + audio_feat_len=_to_int(data.audio_feat_len[batch_index]), + txt_feat=text_tokens[batch_index], + txt_feat_len=_to_int(data.txt_feat_len[batch_index]), + ref_image_feats=image_features, + ref_image_feat_lens=image_feature_lengths, + ref_image_special_tokens=image_special_tokens, + t=video_t, + h=video_h, + w=video_w, + patch_size=self.patch_size, + t_patch_size=self.t_patch_size, + spatial_rope_interpolation=self.config.spatial_rope_interpolation, + diffusion_t=( + data.t[batch_index] if self.config.add_time_token else None + ), + per_token_video_t=( + per_token_video_tokens[batch_index] + if per_token_video_tokens is not None + else None + ), + per_token_audio_t=( + per_token_audio[batch_index] + if per_token_audio is not None + else None + ), + time_channel_dim=self.config.time_channel_dim, + time_pos_fps=self.config.time_pos_fps, + vae_first_latent_is_image=self.config.vae_first_latent_is_image, + video_fps=self.config.video_fps, + ) + ) + + packed = SimplePackedData(items) + packed_cu_seqlens = packed.cu_seqlen.to( + device=data.x_t.device, + dtype=torch.int32, + ) + varlen_handler = VarlenHandler( + cu_seqlens_q=packed_cu_seqlens, + cu_seqlens_k=packed_cu_seqlens.clone(), + max_seqlen_q=packed.max_seqlen, + max_seqlen_k=packed.max_seqlen, + ) + ( + token_sequence, + coords_mapping, + modality_mapping, + time_token_sequence, + varlen_handler, + pad_size, + ) = self._pad_for_ep_cp( + packed.token_sequence, + packed.coords_mapping, + packed.modality_mapping, + packed.time_token_sequence, + varlen_handler, + ) + self.saved_for_output( + simple_packed_data=packed, + input_video_channel=input_video_channel, + pad_size=pad_size, + ) + return ( + token_sequence, + coords_mapping, + modality_mapping, + varlen_handler, + time_token_sequence, + ) + + def process_output( + self, + x: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Remove distributed padding and restore video and audio tensors.""" + packed: SimplePackedData = self.get_saved_data("simple_packed_data") + pad_size = self.get_saved_data("pad_size") + if pad_size > 0: + x = x[:-pad_size] + return packed.depack_token_sequence(x) diff --git a/fastvideo/pipelines/basic/magi2/stages/refiner_data_proxy.py b/fastvideo/pipelines/basic/magi2/stages/refiner_data_proxy.py new file mode 100644 index 000000000..e513c40ed --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/refiner_data_proxy.py @@ -0,0 +1,1096 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pack MAGI-2 refiner video, audio, text, and reference tokens.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from enum import IntEnum +from typing import Any, Literal + +import torch +from einops import rearrange +from torch.nn import functional as F + + +class Modality(IntEnum): + """Modality identifiers consumed by the MAGI-2 refiner transformer.""" + + VIDEO = 0 + AUDIO = 1 + TEXT = 2 + + +@dataclass +class VarlenHandler: + """Variable-length attention boundaries for a packed refiner sequence.""" + + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seqlen_q: int | torch.Tensor + max_seqlen_k: int | torch.Tensor + + +@dataclass +class WindowLocalAttnHandler: + """Query and key ranges consumed by MAGI-2 window attention.""" + + q_ranges: torch.Tensor + k_ranges: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + attn_type_map: torch.Tensor + softmax_scale: float | None = None + bwd_q_ranges: torch.Tensor | None = None + bwd_k_ranges: torch.Tensor | None = None + bwd_attn_type_map: torch.Tensor | None = None + auto_range_merge: bool = False + sparse_load: bool = False + + +@dataclass(frozen=True) +class BlockLocalAttnRanges: + """Forward and backward attention ranges for video blocks.""" + + fwd_q_ranges: torch.Tensor + fwd_k_ranges: torch.Tensor + bwd_q_ranges: torch.Tensor + bwd_k_ranges: torch.Tensor + max_q_len: int + + +@dataclass(frozen=True) +class Magi2RefinerDataProxyConfig: + """Configuration that controls MAGI-2 refiner token packing.""" + + t_patch_size: int = 1 + patch_size: int = 2 + frame_receptive_field: int = 11 + spatial_rope_interpolation: Literal["inter", "extra"] = "extra" + text_offset: int = 0 + coords_style: Literal["v1", "v2"] = "v1" + attn_config: dict[str, Any] = field(default_factory=dict) + magi2_refiner_condition_input: Literal["none", "concat_x1"] = "none" + + +@dataclass +class RefinerModelInput: + """Multimodal tensors and lengths consumed by the refiner data proxy.""" + + x_t: torch.Tensor + audio_x_t: torch.Tensor + audio_feat_len: torch.Tensor + txt_feat: torch.Tensor + txt_feat_len: torch.Tensor + ref_audio_feat: torch.Tensor + ref_audio_feat_len: torch.Tensor + ref_video_feat: torch.Tensor + ref_video_feat_len: torch.Tensor + + +def _to_int(value: int | torch.Tensor) -> int: + """Convert a scalar tensor or Python integer into an integer.""" + if isinstance(value, torch.Tensor): + return int(value.detach().reshape(-1)[0].item()) + return int(value) + + +def _max_range_len(ranges: torch.Tensor) -> int: + """Return the longest query range in a two-column range tensor.""" + if ranges.numel() == 0: + return 0 + lengths = ranges[:, 1].to(torch.int64) - ranges[:, 0].to(torch.int64) + return int(lengths.max().item()) + + +def _cat_ranges( + *ranges: torch.Tensor | None, + device: torch.device, +) -> torch.Tensor: + """Concatenate non-empty attention ranges as contiguous int32 rows.""" + valid_ranges = [value for value in ranges if value is not None and value.numel() > 0] + if valid_ranges: + return torch.cat(valid_ranges, dim=0).to(dtype=torch.int32).contiguous() + return torch.empty((0, 2), device=device, dtype=torch.int32) + + +class BlockLocalAttn: + """Build scan-window attention ranges over an ordered list of blocks.""" + + def __init__(self, block_ranges_tensor: torch.Tensor, win_size: int) -> None: + self.block_ranges_tensor = block_ranges_tensor.to(dtype=torch.int32).contiguous() + self.win_size = int(win_size) + + @property + def num_blocks(self) -> int: + """Return the number of blocks in the scan order.""" + return int(self.block_ranges_tensor.shape[0]) + + @property + def device(self) -> torch.device: + """Return the device that owns the block ranges.""" + return self.block_ranges_tensor.device + + def build_ranges(self) -> BlockLocalAttnRanges: + """Build centered scan windows and their backward-attention ranges.""" + if self.num_blocks == 0: + empty = torch.empty((0, 2), device=self.device, dtype=torch.int32) + return BlockLocalAttnRanges(empty, empty, empty, empty, 0) + window_width = min(self.num_blocks, max(1, self.win_size)) + query_block = torch.arange(self.num_blocks, device=self.device, dtype=torch.int32) + left_width = window_width // 2 + window_start = (query_block - left_width).clamp( + min=0, + max=self.num_blocks - window_width, + ) + window_end = window_start + window_width + query_ranges = self.block_ranges_tensor[query_block.long()].contiguous() + key_ranges = torch.stack( + [ + self.block_ranges_tensor[window_start.long(), 0], + self.block_ranges_tensor[(window_end - 1).long(), 1], + ], + dim=1, + ).contiguous() + + key_block = torch.arange(self.num_blocks, device=self.device, dtype=torch.int32) + first_query_block = torch.searchsorted(window_end, key_block, right=True, out_int32=True) + last_query_block = torch.searchsorted(window_start, key_block, right=True, out_int32=True) - 1 + backward_query_ranges = torch.stack( + [ + self.block_ranges_tensor[first_query_block.long(), 0], + self.block_ranges_tensor[last_query_block.long(), 1], + ], + dim=1, + ).contiguous() + backward_key_ranges = self.block_ranges_tensor[key_block.long()].contiguous() + return BlockLocalAttnRanges( + query_ranges, + key_ranges, + backward_query_ranges, + backward_key_ranges, + _max_range_len(query_ranges), + ) + + +class BlockGridLocalAttn: + """Build radius-based attention ranges over a three-dimensional block grid.""" + + def __init__( + self, + block_ranges_tensor: torch.Tensor, + block_grid_shape: tuple[int, int, int], + radius: tuple[int, int, int], + ) -> None: + """Validate and store the block grid used to construct local ranges.""" + self.block_ranges_tensor = block_ranges_tensor.to(dtype=torch.int32).contiguous() + self.block_grid_shape = tuple(int(dimension) for dimension in block_grid_shape) + self.radius = tuple(int(value) for value in radius) + expected_num_blocks = math.prod(self.block_grid_shape) + if expected_num_blocks != int(self.block_ranges_tensor.shape[0]): + raise ValueError( + "block_grid_shape product must match block count, got " + f"expected_num_blocks={expected_num_blocks} " + f"num_blocks={self.block_ranges_tensor.shape[0]}" + ) + + @property + def num_blocks(self) -> int: + """Return the number of blocks in the spatial-temporal grid.""" + return int(self.block_ranges_tensor.shape[0]) + + @property + def device(self) -> torch.device: + """Return the device that owns the block ranges.""" + return self.block_ranges_tensor.device + + def build_ranges(self) -> BlockLocalAttnRanges: + """Pair each query block with every in-bounds neighbor in its radius.""" + if self.num_blocks == 0: + empty = torch.empty((0, 2), device=self.device, dtype=torch.int32) + return BlockLocalAttnRanges(empty, empty, empty, empty, 0) + time_blocks, height_blocks, width_blocks = self.block_grid_shape + time_radius, height_radius, width_radius = self.radius + query_block = torch.arange(self.num_blocks, device=self.device, dtype=torch.int64) + query_time = query_block // (height_blocks * width_blocks) + query_height = (query_block // width_blocks) % height_blocks + query_width = query_block % width_blocks + query_ranges_by_block = self.block_ranges_tensor[query_block] + query_range_parts: list[torch.Tensor] = [] + key_range_parts: list[torch.Tensor] = [] + for time_delta in range(-time_radius, time_radius + 1): + key_time = query_time + time_delta + valid_time = (key_time >= 0) & (key_time < time_blocks) + if not bool(valid_time.any()): + continue + for height_delta in range(-height_radius, height_radius + 1): + key_height = query_height + height_delta + valid_time_height = valid_time & (key_height >= 0) & (key_height < height_blocks) + if not bool(valid_time_height.any()): + continue + for width_delta in range(-width_radius, width_radius + 1): + key_width = query_width + width_delta + valid = valid_time_height & (key_width >= 0) & (key_width < width_blocks) + if not bool(valid.any()): + continue + valid_indices = torch.nonzero(valid, as_tuple=False).squeeze(1) + key_block = key_time * (height_blocks * width_blocks) + key_height * width_blocks + key_width + query_range_parts.append(query_ranges_by_block[valid_indices]) + key_range_parts.append(self.block_ranges_tensor[key_block[valid_indices]]) + if not query_range_parts: + empty = torch.empty((0, 2), device=self.device, dtype=torch.int32) + return BlockLocalAttnRanges(empty, empty, empty, empty, 0) + query_ranges = torch.cat(query_range_parts, dim=0).to(dtype=torch.int32).contiguous() + key_ranges = torch.cat(key_range_parts, dim=0).to(dtype=torch.int32).contiguous() + return BlockLocalAttnRanges( + query_ranges, + key_ranges, + query_ranges, + key_ranges, + _max_range_len(query_ranges), + ) + + +def _get_coords( + shape: tuple[int, int, int], + ref_feat_shape: tuple[int, int, int], + offset_thw: tuple[int, int, int] = (0, 0, 0), + device: torch.device | None = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Build token coordinates with source-grid and reference-grid metadata.""" + if device is None: + device = torch.device("cpu") + original_time, original_height, original_width = shape + reference_time, reference_height, reference_width = ref_feat_shape + time_offset, height_offset, width_offset = offset_thw + time_range = torch.arange(original_time, device=device, dtype=dtype) + time_offset + height_range = torch.arange(original_height, device=device, dtype=dtype) + height_offset + width_range = torch.arange(original_width, device=device, dtype=dtype) + width_offset + time_grid, height_grid, width_grid = torch.meshgrid( + time_range, + height_range, + width_range, + indexing="ij", + ) + coordinates = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3) + metadata = torch.tensor( + [ + original_time, + original_height, + original_width, + reference_time, + reference_height, + reference_width, + ], + device=device, + dtype=dtype, + ) + return torch.cat([coordinates, metadata.expand(coordinates.size(0), -1)], dim=-1) + + +@dataclass +class SingleData: + """One refiner sample's token segments and packing metadata.""" + + video_x_t: torch.Tensor + audio_x_t: torch.Tensor + audio_feat_len: int + ref_audio_feat: torch.Tensor + ref_audio_feat_len: int + ref_video_feat: torch.Tensor + ref_video_feat_len: int + txt_feat: torch.Tensor + txt_feat_len: int + t: int + h: int + w: int + patch_size: int + t_patch_size: int + spatial_rope_interpolation: Literal["inter", "extra"] + text_offset: int + coords_style: Literal["v1", "v2"] = "v1" + + def __post_init__(self) -> None: + """Trim variable-length segments and cache channel and token counts.""" + self.audio_feat_len = _to_int(self.audio_feat_len) + self.ref_audio_feat_len = _to_int(self.ref_audio_feat_len) + self.ref_video_feat_len = _to_int(self.ref_video_feat_len) + self.txt_feat_len = _to_int(self.txt_feat_len) + self.video_token_num = self.video_x_t.shape[0] + self.ref_video_feat = self.ref_video_feat[:self.ref_video_feat_len] + self.audio_x_t = self.audio_x_t[:self.audio_feat_len] + self.ref_audio_feat = self.ref_audio_feat[:self.ref_audio_feat_len] + self.txt_feat = self.txt_feat[:self.txt_feat_len] + self.video_channel = self.video_x_t.shape[-1] + self.audio_channel = self.audio_x_t.shape[-1] + self.txt_channel = self.txt_feat.shape[-1] + + @property + def device(self) -> torch.device: + """Return the device that owns this sample's packed tokens.""" + return self.video_x_t.device + + @property + def default_dtype(self) -> torch.dtype: + """Return the video-token dtype used for coordinates.""" + return self.video_x_t.dtype + + @property + def total_token_num(self) -> int: + """Return the complete token count for one refiner sample.""" + return ( + self.video_token_num + + self.audio_feat_len + + self.txt_feat_len + + self.ref_audio_feat_len + + self.ref_video_feat_len + ) + + @property + def token_sequence(self) -> torch.Tensor: + """Return channel-padded token segments in transformer consumption order.""" + token_segments = [ + self.video_x_t, + self.audio_x_t, + self.txt_feat, + self.ref_audio_feat, + self.ref_video_feat, + ] + max_channel = max(segment.shape[-1] for segment in token_segments) + padded_segments = [ + F.pad(segment, (0, max_channel - segment.shape[-1])) + for segment in token_segments + ] + return torch.cat(padded_segments, dim=0) + + @property + def modality_mapping(self) -> torch.Tensor: + """Return one modality identifier per token in transformer order.""" + video_mapping = torch.full( + (self.video_token_num,), + int(Modality.VIDEO), + dtype=torch.int64, + device=self.device, + ) + audio_mapping = torch.full( + (self.audio_feat_len,), + int(Modality.AUDIO), + dtype=torch.int64, + device=self.device, + ) + text_mapping = torch.full( + (self.txt_feat_len,), + int(Modality.TEXT), + dtype=torch.int64, + device=self.device, + ) + reference_audio_mapping = torch.full( + (self.ref_audio_feat_len,), + int(Modality.AUDIO), + dtype=torch.int64, + device=self.device, + ) + reference_video_mapping = torch.full( + (self.ref_video_feat_len,), + int(Modality.VIDEO), + dtype=torch.int64, + device=self.device, + ) + return torch.cat( + [ + video_mapping, + audio_mapping, + text_mapping, + reference_audio_mapping, + reference_video_mapping, + ], + dim=0, + ) + + def _default_coords( + self, + shape: tuple[int, int, int], + ref_feat_shape: tuple[int, int, int], + offset_thw: tuple[int, int, int] = (0, 0, 0), + ) -> torch.Tensor: + """Build coordinate rows on this sample's device and dtype.""" + return _get_coords( + shape=shape, + ref_feat_shape=ref_feat_shape, + offset_thw=offset_thw, + device=self.device, + dtype=self.default_dtype, + ) + + @property + def coords_mapping(self) -> torch.Tensor: + """Return source and reference coordinates for every packed token.""" + time_steps = self.t // self.t_patch_size + height_steps = self.h // self.patch_size + width_steps = self.w // self.patch_size + if self.spatial_rope_interpolation == "inter": + video_reference_shape = (time_steps, 32, 32) + else: + video_reference_shape = (time_steps, height_steps, width_steps) + video_coords = self._default_coords( + shape=(time_steps, height_steps, width_steps), + ref_feat_shape=video_reference_shape, + ) + + reference_video_spatial_size = ( + int(math.ceil(math.sqrt(self.ref_video_feat_len))) + if self.ref_video_feat_len > 0 + else 10 + ) + audio_reference_time = (self.audio_feat_len - 1) // 8 + 1 + audio_coords = self._default_coords( + shape=(self.audio_feat_len, 1, 1), + ref_feat_shape=(audio_reference_time // self.t_patch_size, 1, 1), + ) + text_coords = self._default_coords( + shape=(self.txt_feat_len, 1, 1), + ref_feat_shape=(1, 1, 1), + offset_thw=(-self.txt_feat_len, 0, 0), + ) + reference_audio_time = math.ceil( + ((self.ref_audio_feat_len - 1) // 8 + 1) / self.t_patch_size + ) + if self.ref_audio_feat_len > 1: + reference_audio_time = max(reference_audio_time, 2) + reference_audio_coords = self._default_coords( + shape=(self.ref_audio_feat_len, 1, 1), + ref_feat_shape=(reference_audio_time, 1, 1), + offset_thw=(2 * self.audio_feat_len, 0, 0), + ) + reference_video_coords = self._default_coords( + shape=(1, reference_video_spatial_size, reference_video_spatial_size), + ref_feat_shape=(1, reference_video_spatial_size, reference_video_spatial_size), + offset_thw=(1000, 0, 0), + )[:self.ref_video_feat_len] + return torch.cat( + [ + video_coords, + audio_coords, + text_coords, + reference_audio_coords, + reference_video_coords, + ], + dim=0, + ) + + def depack_token_sequence( + self, + token_sequence: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Restore one sample's video latent and active audio-token output.""" + video_tokens = token_sequence[:self.video_token_num, :self.video_channel] + video = rearrange( + video_tokens, + "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", + H=self.h // self.patch_size, + W=self.w // self.patch_size, + pT=self.t_patch_size, + pH=self.patch_size, + pW=self.patch_size, + ).contiguous() + audio = token_sequence[ + self.video_token_num:self.video_token_num + self.audio_feat_len, + :self.audio_channel, + ] + return video, audio + + +@dataclass +class PackedRefinerBatch: + """A refiner batch represented as adjacent variable-length sequences.""" + + items: list[SingleData] + + @property + def token_sequence(self) -> torch.Tensor: + """Return every sample's tokens as one packed tensor.""" + return torch.cat([item.token_sequence for item in self.items], dim=0) + + @property + def modality_mapping(self) -> torch.Tensor: + """Return modality identifiers across every packed sample.""" + return torch.cat([item.modality_mapping for item in self.items], dim=0) + + @property + def coords_mapping(self) -> torch.Tensor: + """Return coordinate rows across every packed sample.""" + return torch.cat([item.coords_mapping for item in self.items], dim=0) + + @property + def total_token_num(self) -> int: + """Return the total token count across the packed batch.""" + return sum(item.total_token_num for item in self.items) + + def __getitem__(self, index: int) -> SingleData: + """Return one packed sample descriptor.""" + return self.items[index] + + @property + def cu_seqlen(self) -> torch.Tensor: + """Return cumulative token boundaries for each packed sample.""" + lengths = torch.tensor([item.total_token_num for item in self.items]) + return F.pad(torch.cumsum(lengths, dim=0), (1, 0)) + + @property + def max_seqlen(self) -> int: + """Return the longest unpadded sample sequence.""" + return max(item.total_token_num for item in self.items) + + def depack_token_sequence( + self, + token_sequence: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Restore batched video latents and active audio-token outputs.""" + video_outputs: list[torch.Tensor] = [] + audio_outputs: list[torch.Tensor] = [] + token_counts = [item.total_token_num for item in self.items] + for item, token_slice in zip( + self.items, + torch.split(token_sequence, token_counts, dim=0), + strict=True, + ): + video, audio = item.depack_token_sequence(token_slice) + video_outputs.append(video) + audio_outputs.append(audio) + return torch.stack(video_outputs, dim=0), torch.stack(audio_outputs, dim=0) + + +def _build_frame_local_attn_handler( + num_video_tokens: int, + num_audio_and_text_tokens: int, + num_frames: int, + frame_receptive_field: int, + device: torch.device, +) -> WindowLocalAttnHandler: + """Build the refiner's legacy frame-window attention ranges.""" + tokens_per_frame = num_video_tokens // num_frames + total_tokens = num_video_tokens + num_audio_and_text_tokens + query_ranges: list[torch.Tensor] = [] + key_ranges: list[torch.Tensor] = [] + for frame_index in range(num_frames): + query_ranges.append( + torch.tensor( + [ + frame_index * tokens_per_frame, + (frame_index + 1) * tokens_per_frame, + ] + ) + ) + key_ranges.append( + torch.tensor( + [ + (frame_index - frame_receptive_field) * tokens_per_frame, + (frame_index + frame_receptive_field + 1) * tokens_per_frame, + ] + ) + ) + local_query_ranges = torch.stack(query_ranges, dim=0) + local_key_ranges = torch.stack(key_ranges, dim=0) + local_key_ranges[local_key_ranges < 0] = 0 + local_key_ranges[local_key_ranges > num_video_tokens] = num_video_tokens + video_query_range = torch.tensor([[0, num_video_tokens]]) + video_key_range = torch.tensor( + [[num_video_tokens, num_video_tokens + num_audio_and_text_tokens]] + ) + context_query_range = torch.tensor([[num_video_tokens, total_tokens]]) + context_key_range = torch.tensor([[0, total_tokens]]) + packed_query_ranges = torch.cat( + [local_query_ranges, video_query_range, context_query_range], + dim=0, + ).to(device=device, dtype=torch.int32, non_blocking=True) + packed_key_ranges = torch.cat( + [local_key_ranges, video_key_range, context_key_range], + dim=0, + ).to(device=device, dtype=torch.int32, non_blocking=True) + return WindowLocalAttnHandler( + q_ranges=packed_query_ranges, + k_ranges=packed_key_ranges, + max_seqlen_q=total_tokens, + max_seqlen_k=total_tokens, + attn_type_map=torch.zeros( + packed_query_ranges.shape[0], + device=device, + dtype=torch.int32, + ), + ) + + +class Magi2RefinerDataProxy: + """Pack MAGI-2 refiner inputs and restore refiner-model outputs.""" + + def __init__(self, config: Magi2RefinerDataProxyConfig) -> None: + """Initialize patch extraction and local-attention configuration.""" + self.patch_size = config.patch_size + self.t_patch_size = config.t_patch_size + self.frame_receptive_field = config.frame_receptive_field + self.spatial_rope_interpolation = config.spatial_rope_interpolation + self.text_offset = config.text_offset + self.coords_style = config.coords_style + self.attn_config = dict(config.attn_config or {}) + attention_mode = str(self.attn_config.get("mode", "dense")).lower() + if attention_mode == "local": + attention_mode = "window" + if attention_mode not in {"dense", "window"}: + raise ValueError(f"Unsupported refiner attention mode: {attention_mode!r}") + self.use_window_attn = attention_mode == "window" + self.window_attn_config = dict(self.attn_config.get("window", {})) + self._saved_data: dict[str, Any] = {} + + def saved_for_output(self, **kwargs: Any) -> None: + """Save packing metadata that is required to depack model outputs.""" + self._saved_data.update(kwargs) + + def get_saved_data(self, key: str) -> Any: + """Return saved packing metadata by name.""" + return self._saved_data[key] + + def img2tokens(self, video: torch.Tensor) -> torch.Tensor: + """Extract non-overlapping 3D patches with channel-major features.""" + patches = ( + video.unfold(2, self.t_patch_size, self.t_patch_size) + .unfold(3, self.patch_size, self.patch_size) + .unfold(4, self.patch_size, self.patch_size) + ) + batch_size, _channels, time_steps, height_steps, width_steps, _, _, _ = patches.shape + return ( + patches.permute(0, 2, 3, 4, 1, 5, 6, 7) + .reshape(batch_size, time_steps * height_steps * width_steps, -1) + .contiguous() + ) + + @staticmethod + def _dimension_block_sizes( + dimension_size: int, + block_size: int, + device: torch.device, + ) -> torch.Tensor: + """Return valid-token counts for full and edge blocks in one dimension.""" + sizes = torch.full( + (math.ceil(dimension_size / block_size),), + block_size, + device=device, + dtype=torch.int32, + ) + sizes[-1] = dimension_size - block_size * (sizes.numel() - 1) + return sizes + + @classmethod + def _build_video_block_ranges( + cls, + patch_counts: tuple[int, int, int], + block_sizes: tuple[int, int, int], + device: torch.device, + ) -> torch.Tensor: + """Build contiguous token ranges for every spatial-temporal video block.""" + time_count, height_count, width_count = patch_counts + time_block, height_block, width_block = block_sizes + time_sizes = cls._dimension_block_sizes(time_count, time_block, device) + height_sizes = cls._dimension_block_sizes(height_count, height_block, device) + width_sizes = cls._dimension_block_sizes(width_count, width_block, device) + block_token_counts = ( + time_sizes[:, None, None] + * height_sizes[None, :, None] + * width_sizes[None, None, :] + ).flatten() + ends = torch.cumsum(block_token_counts, dim=0) + starts = ends - block_token_counts + return torch.stack([starts, ends], dim=1).contiguous() + + @classmethod + def _build_block_order( + cls, + *, + patch_counts: tuple[int, int, int], + block_sizes: tuple[int, int, int], + valid_token_num: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[int, int, int]]: + """Reorder video tokens into contiguous blocks and retain the inverse order.""" + video_token_num = math.prod(patch_counts) + video_token_indices = torch.arange( + video_token_num, + device=device, + dtype=torch.int32, + ).view(*patch_counts) + padded_counts = tuple( + math.ceil(patch_count / block_size) * block_size + for patch_count, block_size in zip( + patch_counts, + block_sizes, + strict=True, + ) + ) + padded_video_indices = torch.full( + padded_counts, + -1, + dtype=torch.int32, + device=device, + ) + padded_video_indices[ + :patch_counts[0], + :patch_counts[1], + :patch_counts[2], + ] = video_token_indices + time_block, height_block, width_block = block_sizes + padded_video_indices = ( + padded_video_indices.view( + padded_counts[0] // time_block, + time_block, + padded_counts[1] // height_block, + height_block, + padded_counts[2] // width_block, + width_block, + ) + .permute(0, 2, 4, 1, 3, 5) + .contiguous() + .view(-1) + ) + video_order = padded_video_indices[padded_video_indices >= 0] + tail_indices = torch.arange( + video_token_num, + valid_token_num, + device=device, + dtype=torch.int32, + ) + token_order = torch.cat([video_order, tail_indices], dim=0) + token_restore_order = torch.empty_like(token_order) + token_restore_order[token_order] = torch.arange( + token_order.numel(), + device=device, + dtype=torch.int32, + ) + block_grid_shape = ( + math.ceil(patch_counts[0] / time_block), + math.ceil(patch_counts[1] / height_block), + math.ceil(patch_counts[2] / width_block), + ) + video_block_ranges = cls._build_video_block_ranges( + patch_counts, + block_sizes, + device, + ) + return token_order, token_restore_order, video_block_ranges, block_grid_shape + + @staticmethod + def _dense_context_query_ranges( + query_range: tuple[int, int], + block_size: int, + device: torch.device, + ) -> torch.Tensor: + """Split a dense context query range into bounded attention rows.""" + query_start, query_end = query_range + if query_end <= query_start: + return torch.empty((0, 2), device=device, dtype=torch.int32) + query_starts = torch.arange( + query_start, + query_end, + max(1, block_size), + device=device, + dtype=torch.int32, + ) + query_ends = torch.clamp(query_starts + max(1, block_size), max=query_end) + return torch.stack((query_starts, query_ends), dim=1).contiguous() + + def _build_window_local_attn_handler( + self, + item: SingleData, + total_token_num: int, + device: torch.device, + ) -> WindowLocalAttnHandler: + """Build block- or frame-window ranges plus dense cross-modal ranges.""" + window_config = {**self.attn_config, **self.window_attn_config} + level = str(window_config.get("level", "block")).lower() + if level not in {"block", "frame"}: + raise ValueError(f"Unsupported window attention level: {level!r}") + patch_counts = ( + item.t // item.t_patch_size, + item.h // item.patch_size, + item.w // item.patch_size, + ) + tokens_per_frame = patch_counts[1] * patch_counts[2] + video_token_num = math.prod(patch_counts) + tail_range = (video_token_num, total_token_num) + video_range = (0, video_token_num) + + if level == "block": + block_sizes = ( + int(window_config.get("block_t_size", self.attn_config.get("block_t_size", 1))), + int(window_config.get("block_size", self.attn_config.get("block_size", 1))), + int(window_config.get("block_size", self.attn_config.get("block_size", 1))), + ) + _, _, video_block_ranges, block_grid_shape = self._build_block_order( + patch_counts=patch_counts, + block_sizes=block_sizes, + valid_token_num=total_token_num, + device=device, + ) + block_mode = str( + window_config.get( + "block_mode", + window_config.get("window_block_mode", "scan"), + ) + ).lower() + if block_mode == "grid": + window_ranges = BlockGridLocalAttn( + block_ranges_tensor=video_block_ranges, + block_grid_shape=block_grid_shape, + radius=( + int(window_config.get("block_t_radius", 1)), + int(window_config.get("block_h_radius", 1)), + int(window_config.get("block_w_radius", 1)), + ), + ).build_ranges() + elif block_mode == "scan": + window_ranges = BlockLocalAttn( + block_ranges_tensor=video_block_ranges, + win_size=int(window_config.get("win_size", self.attn_config.get("win_size", 1))), + ).build_ranges() + else: + raise ValueError(f"Unsupported block window mode: {block_mode!r}") + video_dense_query_ranges = video_block_ranges + dense_context_query_block_size = max(math.prod(block_sizes), 1) + else: + frame_indices = torch.arange(patch_counts[0], device=device, dtype=torch.int32) + frame_ranges = torch.stack( + ( + frame_indices * tokens_per_frame, + (frame_indices + 1) * tokens_per_frame, + ), + dim=1, + ).contiguous() + radius = int(window_config.get("frame_receptive_field", self.frame_receptive_field)) + if radius < 0: + raise ValueError("frame-level window attention requires frame_receptive_field >= 0") + key_start_frames = (frame_indices - radius).clamp(min=0) + key_end_frames = (frame_indices + radius + 1).clamp(max=patch_counts[0]) + forward_key_ranges = torch.stack( + ( + key_start_frames * tokens_per_frame, + key_end_frames * tokens_per_frame, + ), + dim=1, + ) + window_ranges = BlockLocalAttnRanges( + frame_ranges, + forward_key_ranges, + forward_key_ranges, + frame_ranges, + tokens_per_frame, + ) + video_dense_query_ranges = frame_ranges + dense_context_query_block_size = max(tokens_per_frame, 1) + + dense_query_ranges = torch.empty((0, 2), device=device, dtype=torch.int32) + dense_key_ranges = torch.empty((0, 2), device=device, dtype=torch.int32) + if tail_range[0] < tail_range[1]: + query_parts = [ + video_dense_query_ranges, + self._dense_context_query_ranges( + tail_range, + dense_context_query_block_size, + device, + ), + self._dense_context_query_ranges( + tail_range, + dense_context_query_block_size, + device, + ), + ] + key_parts = [ + torch.tensor(tail_range, device=device, dtype=torch.int32) + .view(1, 2) + .expand(video_dense_query_ranges.shape[0], 2), + torch.tensor(video_range, device=device, dtype=torch.int32) + .view(1, 2) + .expand(query_parts[1].shape[0], 2), + torch.tensor(tail_range, device=device, dtype=torch.int32) + .view(1, 2) + .expand(query_parts[2].shape[0], 2), + ] + dense_query_ranges = torch.cat(query_parts, dim=0).contiguous() + dense_key_ranges = torch.cat(key_parts, dim=0).contiguous() + + query_ranges = _cat_ranges( + dense_query_ranges, + window_ranges.fwd_q_ranges, + device=device, + ) + key_ranges = _cat_ranges( + dense_key_ranges, + window_ranges.fwd_k_ranges, + device=device, + ) + return WindowLocalAttnHandler( + q_ranges=query_ranges, + k_ranges=key_ranges, + max_seqlen_q=max(_max_range_len(query_ranges), window_ranges.max_q_len), + max_seqlen_k=total_token_num, + attn_type_map=torch.zeros( + query_ranges.shape[0], + device=device, + dtype=torch.int32, + ), + bwd_q_ranges=_cat_ranges( + dense_query_ranges, + window_ranges.bwd_q_ranges, + device=device, + ), + bwd_k_ranges=_cat_ranges( + dense_key_ranges, + window_ranges.bwd_k_ranges, + device=device, + ), + bwd_attn_type_map=torch.zeros( + query_ranges.shape[0], + device=device, + dtype=torch.int32, + ), + auto_range_merge=bool(window_config.get("auto_range_merge", False)), + sparse_load=bool(window_config.get("sparse_load", False)), + ) + + def process_input( + self, + data: RefinerModelInput, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + VarlenHandler, + WindowLocalAttnHandler | None, + ]: + """Pack one refiner batch and construct its attention metadata. + + The denoising loop owns classifier-free guidance and diffusion times. + ``UlyssesScheduler`` owns context-parallel sequence splitting. + """ + batch_size, input_video_channel, video_time, video_height, video_width = data.x_t.shape + video_tokens = self.img2tokens(data.x_t) + reference_video_tokens = self.img2tokens(data.ref_video_feat) + audio_tokens = data.audio_x_t.contiguous() + reference_audio_tokens = data.ref_audio_feat.contiguous() + text_tokens = data.txt_feat.contiguous() + + items: list[SingleData] = [] + for batch_index in range(batch_size): + items.append( + SingleData( + video_x_t=video_tokens[batch_index], + audio_x_t=audio_tokens[batch_index], + audio_feat_len=_to_int(data.audio_feat_len[batch_index]), + ref_audio_feat=reference_audio_tokens[batch_index], + ref_audio_feat_len=_to_int(data.ref_audio_feat_len[batch_index]), + txt_feat=text_tokens[batch_index], + txt_feat_len=_to_int(data.txt_feat_len[batch_index]), + t=video_time, + h=video_height, + w=video_width, + patch_size=self.patch_size, + t_patch_size=self.t_patch_size, + spatial_rope_interpolation=self.spatial_rope_interpolation, + text_offset=self.text_offset, + ref_video_feat=reference_video_tokens[batch_index], + ref_video_feat_len=_to_int(data.ref_video_feat_len[batch_index]), + coords_style=self.coords_style, + ) + ) + packed_refiner_batch = PackedRefinerBatch(items=items) + cumulative_lengths = packed_refiner_batch.cu_seqlen.to( + device=data.x_t.device, + dtype=torch.int32, + ) + maximum_length = torch.tensor( + packed_refiner_batch.max_seqlen, + device=data.x_t.device, + dtype=torch.int32, + ) + varlen_handler = VarlenHandler( + cu_seqlens_q=cumulative_lengths, + cu_seqlens_k=cumulative_lengths.clone(), + max_seqlen_q=maximum_length, + max_seqlen_k=maximum_length.clone(), + ) + + token_sequence = packed_refiner_batch.token_sequence + coords_mapping = packed_refiner_batch.coords_mapping + modality_mapping = packed_refiner_batch.modality_mapping + token_restore_order = None + if self.use_window_attn: + assert batch_size == 1, "window attention only supports batch size 1" + item = packed_refiner_batch[0] + local_attn_handler = self._build_window_local_attn_handler( + item, + int(packed_refiner_batch.total_token_num), + token_sequence.device, + ) + level = str(self.window_attn_config.get("level", "block")).lower() + if level == "block": + window_config = {**self.attn_config, **self.window_attn_config} + block_sizes = ( + int(window_config.get("block_t_size", self.attn_config.get("block_t_size", 1))), + int(window_config.get("block_size", self.attn_config.get("block_size", 1))), + int(window_config.get("block_size", self.attn_config.get("block_size", 1))), + ) + patch_counts = ( + item.t // item.t_patch_size, + item.h // item.patch_size, + item.w // item.patch_size, + ) + token_order, token_restore_order, _, _ = self._build_block_order( + patch_counts=patch_counts, + block_sizes=block_sizes, + valid_token_num=int(packed_refiner_batch.total_token_num), + device=token_sequence.device, + ) + token_sequence = token_sequence[token_order] + coords_mapping = coords_mapping[token_order] + modality_mapping = modality_mapping[token_order] + elif self.frame_receptive_field != -1: + assert batch_size == 1, "local attention only supports batch size 1" + item = packed_refiner_batch[0] + local_attn_handler = _build_frame_local_attn_handler( + num_video_tokens=item.video_token_num, + num_audio_and_text_tokens=( + item.audio_feat_len + + item.txt_feat_len + + item.ref_audio_feat_len + ), + num_frames=video_time, + frame_receptive_field=self.frame_receptive_field, + device=token_sequence.device, + ) + else: + local_attn_handler = None + + self.saved_for_output( + packed_refiner_batch=packed_refiner_batch, + input_video_channel=input_video_channel, + token_restore_order=token_restore_order, + ) + return ( + token_sequence, + coords_mapping, + modality_mapping, + varlen_handler, + local_attn_handler, + ) + + def process_output( + self, + output_tokens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Restore block order and depack refiner video and audio outputs.""" + token_restore_order = self.get_saved_data("token_restore_order") + if token_restore_order is not None: + output_tokens = output_tokens[token_restore_order] + packed_refiner_batch: PackedRefinerBatch = self.get_saved_data( + "packed_refiner_batch" + ) + return packed_refiner_batch.depack_token_sequence(output_tokens) diff --git a/fastvideo/pipelines/basic/magi2/stages/runtime.py b/fastvideo/pipelines/basic/magi2/stages/runtime.py new file mode 100644 index 000000000..79255291b --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/runtime.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Request validation and latent initialization for MAGI-2 inference.""" + +from __future__ import annotations + +import random + +import numpy as np +import torch + +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +class Magi2InputValidationStage(PipelineStage): + """Validate the published ten-second 1080p T2V and I2V request profile.""" + + def __init__( + self, + output_frames: int, + output_height: int, + output_width: int, + output_fps: int, + ) -> None: + """Store the fixed output geometry supported by the release checkpoint.""" + super().__init__() + self.output_frames = output_frames + self.output_height = output_height + self.output_width = output_width + self.output_fps = output_fps + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record; ``forward`` emits model-specific errors.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record after request normalization and seeding.""" + del batch, fastvideo_args + return VerificationResult() + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Require release geometry and reset every request's random-number stream.""" + if fastvideo_args.workload_type not in (WorkloadType.T2V, WorkloadType.I2V): + raise ValueError("MAGI-2 supports only T2V and I2V workloads") + if not isinstance(batch.prompt, str) or not batch.prompt.strip(): + raise ValueError("MAGI-2 requires one non-empty prompt string") + expected_geometry = ( + self.output_frames, + self.output_height, + self.output_width, + self.output_fps, + ) + received_geometry = ( + batch.num_frames, + batch.height, + batch.width, + batch.fps, + ) + if received_geometry != expected_geometry: + raise ValueError( + "MAGI-2 Preview supports 249 frames at 1088x1920 and 25 fps; " + f"received frames/height/width/fps={received_geometry}" + ) + if batch.num_inference_steps <= 0 or batch.num_inference_steps_sr <= 0: + raise ValueError("MAGI-2 denoising step counts must be positive") + seed = 42 if batch.seed is None else int(batch.seed) + batch.seed = seed + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + if fastvideo_args.deterministic or envs.MAGI2_DETERMINISTIC: + random.seed(seed) + np.random.seed(seed) + return batch + + +class Magi2LatentPreparationStage(PipelineStage): + """Draw video noise before audio noise with the official release shapes.""" + + def __init__( + self, + video_channels: int, + video_length: int, + video_height: int, + video_width: int, + audio_length: int, + audio_channels: int, + ) -> None: + """Store the preview video and audio latent geometry.""" + super().__init__() + self.video_shape = ( + 1, + video_channels, + video_length, + video_height, + video_width, + ) + self.audio_shape = (1, audio_length, audio_channels) + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record for a validated and seeded request.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record for joint video and audio noise tensors.""" + del batch, fastvideo_args + return VerificationResult() + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Draw FP32 video and audio noise in the release model's RNG order.""" + del fastvideo_args + device = torch.device("cuda", torch.cuda.current_device()) + batch.latents = torch.randn( + self.video_shape, + device=device, + dtype=torch.float32, + ) + batch.audio_latents = torch.randn( + self.audio_shape, + device=device, + dtype=torch.float32, + ) + return batch + + +__all__ = ["Magi2InputValidationStage", "Magi2LatentPreparationStage"] diff --git a/fastvideo/pipelines/basic/magi2/stages/sampling.py b/fastvideo/pipelines/basic/magi2/stages/sampling.py new file mode 100644 index 000000000..49f5921ce --- /dev/null +++ b/fastvideo/pipelines/basic/magi2/stages/sampling.py @@ -0,0 +1,500 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Preview and refiner denoising stages for MAGI-2.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from torch.nn import functional as F + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, +) +from fastvideo.pipelines.basic.magi2.stages.preview_data_proxy import ( + Magi2DataProxy, + Magi2DataProxyConfig, + ModelInput, +) +from fastvideo.pipelines.basic.magi2.stages.refiner_data_proxy import ( + Magi2RefinerDataProxy, + Magi2RefinerDataProxyConfig, + RefinerModelInput, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +def _pad_or_trim( + tensor: torch.Tensor, + target_size: int, + dimension: int, + pad_value: float = 0.0, +) -> torch.Tensor: + """Match one tensor dimension to the shared classifier-free-guidance length.""" + current_size = tensor.size(dimension) + if current_size < target_size: + padding_amount = target_size - current_size + padding = [0] * (2 * tensor.dim()) + padding_dimension = tensor.dim() - 1 - dimension + padding[2 * padding_dimension + 1] = padding_amount + return F.pad(tensor, tuple(padding), "constant", pad_value) + slices = [slice(None)] * tensor.dim() + slices[dimension] = slice(0, target_size) + return tensor[tuple(slices)] + + +class Magi2PreviewDenoisingStage(PipelineStage): + """Run the joint video-audio preview transformer with two-way CFG.""" + + def __init__( + self, + transformer: Any, + data_proxy_config: Magi2DataProxyConfig, + flow_shift: float, + video_guidance_scale: float, + audio_guidance_scale: float, + ) -> None: + """Store the preview transformer and published sampling constants.""" + super().__init__() + self.transformer = transformer + self.data_proxy = Magi2DataProxy(data_proxy_config) + self.flow_shift = flow_shift + self.video_guidance_scale = video_guidance_scale + self.audio_guidance_scale = audio_guidance_scale + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record; ``forward`` validates required tensors.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record for denoised video and audio latents.""" + del batch, fastvideo_args + return VerificationResult() + + @staticmethod + def _prepare_reference_cfg( + reference_latent: torch.Tensor | None, + reference_length: torch.Tensor | None, + special_tokens: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Duplicate I2V conditioning across both classifier-free-guidance halves.""" + if reference_latent is None: + return None, None, None + if reference_length is None or special_tokens is None: + raise ValueError("MAGI-2 I2V conditioning is incomplete") + return ( + torch.cat([reference_latent, reference_latent], dim=0), + torch.cat([reference_length, reference_length], dim=0), + torch.cat([special_tokens, special_tokens], dim=0), + ) + + def _prepare_model_input( + self, + batch: ForwardBatch, + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + timestep: torch.Tensor, + ) -> ModelInput: + """Build the exact two-sample conditional/unconditional preview input.""" + text_context = batch.magi2_text_context + negative_context = batch.magi2_negative_context + if text_context is None or negative_context is None: + raise ValueError("MAGI-2 preview denoising requires text conditioning") + text_length = text_context.shape[1] + negative_length = negative_context.shape[1] + shared_length = max(text_length, negative_length) + text_context = _pad_or_trim(text_context, shared_length, 1) + negative_context = _pad_or_trim(negative_context, shared_length, 1) + + reference_video = torch.empty_like(video_latent) + reference_video = torch.cat( + [reference_video, torch.zeros_like(reference_video)], + dim=0, + ) + reference_video_length = torch.tensor( + [0, 0], + device=video_latent.device, + ) + reference_image, reference_image_length, special_tokens = ( + self._prepare_reference_cfg( + batch.magi2_ref_image_feat, + batch.magi2_ref_image_feat_len, + batch.magi2_ref_image_special_tokens, + ) + ) + + timestep_value = timestep.reshape(-1)[0].to( + device=video_latent.device, + dtype=video_latent.dtype, + ) + normalized_timestep = torch.stack([timestep_value, timestep_value]) / 1000.0 + _, _, video_time, video_height, video_width = video_latent.shape + audio_time = audio_latent.shape[1] + per_token_video_t = normalized_timestep.view(-1, 1, 1, 1, 1).expand( + 2, + 1, + video_time, + video_height, + video_width, + ).clone() + per_token_audio_t = normalized_timestep.view(-1, 1, 1).expand( + 2, + audio_time, + 1, + ).clone() + reference_audio = torch.zeros( + 1, + 0, + audio_latent.shape[-1], + device=video_latent.device, + dtype=torch.float32, + ) + return ModelInput( + x_t=torch.cat([video_latent, video_latent], dim=0), + audio_x_t=torch.cat([audio_latent, audio_latent], dim=0), + audio_feat_len=torch.tensor( + [audio_time, audio_time], + device=video_latent.device, + ), + txt_feat=torch.cat([text_context, negative_context], dim=0), + txt_feat_len=torch.tensor( + [text_length, negative_length], + device=video_latent.device, + ), + t=normalized_timestep, + per_token_video_t=per_token_video_t, + per_token_audio_t=per_token_audio_t, + ref_audio_feat=torch.cat( + [reference_audio, torch.zeros_like(reference_audio)], + dim=0, + ), + ref_audio_feat_len=torch.tensor( + [0, 0], + device=video_latent.device, + ), + ref_video_feat=reference_video, + ref_video_feat_len=reference_video_length, + ref_image_feat=reference_image, + ref_image_feat_len=reference_image_length, + ref_image_special_token_embedding=special_tokens, + ) + + def _predict_velocity( + self, + model_input: ModelInput, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pack tokens, execute the preview transformer, and restore both modalities.""" + transformer_input = self.data_proxy.process_input(model_input) + transformer_output = self.transformer(*transformer_input) + return self.data_proxy.process_output(transformer_output) + + @torch.inference_mode() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Advance the video and audio latents through the published preview schedule.""" + if batch.latents is None or batch.audio_latents is None: + raise ValueError("MAGI-2 preview denoising requires joint noise latents") + device = torch.device("cuda", torch.cuda.current_device()) + self.transformer.to(device) + video_latent = batch.latents.clone() + audio_latent = batch.audio_latents.clone() + video_scheduler = FlowUniPCMultistepScheduler() + audio_scheduler = FlowUniPCMultistepScheduler() + video_scheduler.set_timesteps( + batch.num_inference_steps, + device=device, + shift=self.flow_shift, + ) + audio_scheduler.set_timesteps( + batch.num_inference_steps, + device=device, + shift=self.flow_shift, + ) + video_guidance = torch.tensor( + [self.video_guidance_scale], + device=device, + ) + for timestep in video_scheduler.timesteps: + model_input = self._prepare_model_input( + batch, + video_latent, + audio_latent, + timestep, + ) + video_prediction, audio_prediction = self._predict_velocity(model_input) + conditional_video = video_prediction[0:1] + unconditional_video = video_prediction[1:2] + video_velocity = unconditional_video + video_guidance * ( + conditional_video - unconditional_video + ) + conditional_audio = audio_prediction[0:1] + unconditional_audio = audio_prediction[1:2] + audio_velocity = unconditional_audio + self.audio_guidance_scale * ( + conditional_audio - unconditional_audio + ) + video_latent = video_scheduler.step( + video_velocity, + timestep, + video_latent, + return_dict=False, + )[0] + audio_latent = audio_scheduler.step( + audio_velocity, + timestep, + audio_latent, + return_dict=False, + )[0] + batch.latents = video_latent + batch.audio_latents = audio_latent + if fastvideo_args.dit_cpu_offload: + self.transformer.to("cpu") + torch.cuda.empty_cache() + return batch + + +class ZeroSNRDDPMDiscretization: + """Construct the refiner's published zero-terminal-SNR noise schedule.""" + + def __init__( + self, + linear_start: float = 0.00085, + linear_end: float = 0.0120, + num_timesteps: int = 1000, + ) -> None: + """Precompute the cumulative alpha schedule in FP64 NumPy arithmetic.""" + self.num_timesteps = num_timesteps + betas = torch.linspace( + linear_start**0.5, + linear_end**0.5, + num_timesteps, + dtype=torch.float64, + ) ** 2 + self.alphas_cumprod = np.cumprod(1.0 - betas.numpy(), axis=0) + + def __call__( + self, + count: int, + do_append_zero: bool = True, + device: str = "cpu", + flip: bool = False, + ) -> torch.Tensor: + """Return sampled sigma values with optional terminal zero and reversal.""" + sigmas = self.get_sigmas(count, device=device) + if do_append_zero: + sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) + return torch.flip(sigmas, (0,)) if flip else sigmas + + def get_sigmas(self, count: int, device: str = "cpu") -> torch.Tensor: + """Convert cumulative alphas into the refiner's sigma parameterization.""" + if count < self.num_timesteps: + timesteps = np.linspace( + self.num_timesteps - 1, + 0, + count, + endpoint=False, + ).astype(int)[::-1] + alphas_cumprod = self.alphas_cumprod[timesteps] + elif count == self.num_timesteps: + alphas_cumprod = self.alphas_cumprod + else: + raise ValueError( + f"sigma count must be <= {self.num_timesteps}, received {count}" + ) + alphas = torch.tensor( + alphas_cumprod, + dtype=torch.float32, + device=device, + ).sqrt() + first_alpha = alphas[0].clone() + terminal_alpha = alphas[-1].clone() + alphas -= terminal_alpha + alphas *= first_alpha / (first_alpha - terminal_alpha) + return torch.flip(alphas, (0,)) + + +class Magi2RefinerStage(PipelineStage): + """Upsample preview latents and run the five-step 1080p refiner.""" + + def __init__( + self, + transformer: Any, + data_proxy_config: Magi2RefinerDataProxyConfig, + latent_height: int, + latent_width: int, + noise_index: int, + flow_shift: float, + video_guidance_scale: float, + audio_guidance_scale: float, + audio_channels: int, + ) -> None: + """Store the refiner transformer, token proxy, and release constants.""" + super().__init__() + self.transformer = transformer + self.data_proxy = Magi2RefinerDataProxy(data_proxy_config) + self.latent_height = latent_height + self.latent_width = latent_width + self.noise_index = noise_index + self.flow_shift = flow_shift + self.video_guidance_scale = video_guidance_scale + self.audio_guidance_scale = audio_guidance_scale + self.audio_channels = audio_channels + self.sigmas = ZeroSNRDDPMDiscretization()( + 1000, + do_append_zero=False, + flip=True, + ) + + def verify_input( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record; ``forward`` validates latent conditioning.""" + del batch, fastvideo_args + return VerificationResult() + + def verify_output( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> VerificationResult: + """Return the default record for the post-refiner video latent.""" + del batch, fastvideo_args + return VerificationResult() + + def _predict_velocity( + self, + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + text_context: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pack one text branch, execute the refiner, and restore both modalities.""" + reference_audio = torch.zeros( + 1, + 0, + self.audio_channels, + device=video_latent.device, + dtype=torch.float32, + ) + model_input = RefinerModelInput( + x_t=video_latent, + audio_x_t=audio_latent, + audio_feat_len=torch.tensor([audio_latent.shape[1]]), + txt_feat=text_context.to(torch.float32), + txt_feat_len=torch.tensor([text_context.shape[1]]), + ref_audio_feat=reference_audio, + ref_audio_feat_len=torch.tensor([0]), + ref_video_feat=torch.empty_like(video_latent), + ref_video_feat_len=torch.tensor([0]), + ) + transformer_input = self.data_proxy.process_input(model_input) + transformer_output = self.transformer(*transformer_input) + return self.data_proxy.process_output(transformer_output) + + @torch.inference_mode() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Inject refiner noise and advance the upsampled latent for five steps.""" + if batch.latents is None: + raise ValueError("MAGI-2 refiner requires the preview video latent") + if batch.magi2_text_context is None or batch.magi2_negative_context is None: + raise ValueError("MAGI-2 refiner requires positive and negative text context") + device = torch.device("cuda", torch.cuda.current_device()) + video_latent = F.interpolate( + batch.latents, + size=( + batch.latents.shape[2] * 2 - 1, + self.latent_height, + self.latent_width, + ), + mode="trilinear", + align_corners=True, + ) + noise = torch.randn_like(video_latent, device=video_latent.device) + sigma = self.sigmas.to(video_latent.device)[self.noise_index] + video_latent = video_latent * sigma + noise * (1 - sigma**2) ** 0.5 + refiner_audio = torch.zeros( + 1, + 0, + self.audio_channels, + device=device, + dtype=torch.float32, + ) + video_scheduler = FlowUniPCMultistepScheduler() + audio_scheduler = FlowUniPCMultistepScheduler() + video_scheduler.set_timesteps( + batch.num_inference_steps_sr, + device=device, + shift=self.flow_shift, + ) + audio_scheduler.set_timesteps( + batch.num_inference_steps_sr, + device=device, + shift=self.flow_shift, + ) + video_guidance = torch.tensor( + self.video_guidance_scale, + device=device, + ).expand(1, 1, video_latent.shape[2], 1, 1).clone() + self.transformer.to(device) + for timestep in video_scheduler.timesteps: + conditional_video, conditional_audio = self._predict_velocity( + video_latent, + refiner_audio, + batch.magi2_text_context, + ) + unconditional_video, unconditional_audio = self._predict_velocity( + video_latent, + refiner_audio, + batch.magi2_negative_context, + ) + video_velocity = unconditional_video + video_guidance * ( + conditional_video - unconditional_video + ) + audio_velocity = unconditional_audio + self.audio_guidance_scale * ( + conditional_audio - unconditional_audio + ) + video_latent = video_scheduler.step( + video_velocity, + timestep, + video_latent, + return_dict=False, + )[0] + if audio_velocity.numel() > 0: + refiner_audio = audio_scheduler.step( + audio_velocity, + timestep, + refiner_audio, + return_dict=False, + )[0] + batch.latents = video_latent + if fastvideo_args.dit_cpu_offload: + self.transformer.to("cpu") + torch.cuda.empty_cache() + return batch + + +__all__ = [ + "Magi2PreviewDenoisingStage", + "Magi2RefinerStage", + "ZeroSNRDDPMDiscretization", +] diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 243bd2ae0..ad2197785 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -104,6 +104,11 @@ class ForwardBatch: negative_attention_mask: list[torch.Tensor] | None = None clip_embedding_pos: list[torch.Tensor] | None = None clip_embedding_neg: list[torch.Tensor] | None = None + magi2_text_context: torch.Tensor | None = None + magi2_negative_context: torch.Tensor | None = None + magi2_ref_image_feat: torch.Tensor | None = None + magi2_ref_image_feat_len: torch.Tensor | None = None + magi2_ref_image_special_tokens: torch.Tensor | None = None # Additional text-related parameters max_sequence_length: int | None = None @@ -125,6 +130,7 @@ class ForwardBatch: # Latent tensors latents: torch.Tensor | None = None + audio_latents: torch.Tensor | None = None lq_latents: torch.Tensor | None = None raw_latent_shape: tuple[int, ...] | None = None noise_pred: torch.Tensor | None = None diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 30d885937..82eea856a 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -32,6 +32,7 @@ from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsk from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig +from fastvideo.configs.pipelines.magi2 import Magi2PreviewPipelineConfig from fastvideo.configs.pipelines.flux_2 import ( Flux2KleinPipelineConfig, Flux2PipelineConfig, @@ -299,6 +300,17 @@ def _register_configs() -> None: default_preset="ltx2_base", ) + register_configs( + sampling_param_cls=None, + pipeline_config_cls=Magi2PreviewPipelineConfig, + workload_types=(WorkloadType.T2V, WorkloadType.I2V), + model_detectors=[ + lambda path: "magi2pipeline" in path.lower(), + ], + model_family="magi2", + default_preset="magi2_preview_1080p", + ) + # Stable Audio Open (text-to-audio). Both variants must be loaded # from the FastVideo-curated converted Diffusers-format repos — # the upstream `stabilityai/stable-audio-open-{1.0,small}` repos @@ -1214,6 +1226,8 @@ def _register_presets() -> None: ALL_PRESETS as LONGCAT_PRESETS, ) from fastvideo.pipelines.basic.ltx2.presets import ( ALL_PRESETS as LTX2_PRESETS, ) + from fastvideo.pipelines.basic.magi2.presets import ( + ALL_PRESETS as MAGI2_PRESETS, ) from fastvideo.pipelines.basic.matrixgame2.presets import ( ALL_PRESETS as MATRIXGAME2_PRESETS, ) from fastvideo.pipelines.basic.matrixgame3.presets import ( @@ -1242,6 +1256,7 @@ def _register_presets() -> None: LINGBOTWORLD_PRESETS, LONGCAT_PRESETS, LTX2_PRESETS, + MAGI2_PRESETS, MATRIXGAME2_PRESETS, MATRIXGAME3_PRESETS, SD35_PRESETS, diff --git a/fastvideo/tests/api/test_parser.py b/fastvideo/tests/api/test_parser.py index 64da28dc8..b21030c19 100644 --- a/fastvideo/tests/api/test_parser.py +++ b/fastvideo/tests/api/test_parser.py @@ -129,6 +129,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "enable_stage_verification": True, "use_fsdp_inference": False, "disable_autocast": False, + "deterministic": False, "quantization": None, }, "pipeline": { diff --git a/fastvideo/third_party/flash_attention b/fastvideo/third_party/flash_attention new file mode 160000 index 000000000..b613d9e2c --- /dev/null +++ b/fastvideo/third_party/flash_attention @@ -0,0 +1 @@ +Subproject commit b613d9e2c8475945baff3fd68f2030af1b890acf diff --git a/fastvideo/third_party/magi_attention b/fastvideo/third_party/magi_attention new file mode 160000 index 000000000..2c6413571 --- /dev/null +++ b/fastvideo/third_party/magi_attention @@ -0,0 +1 @@ +Subproject commit 2c6413571c2cac6a80d1f85a434c6713fe0f5286 diff --git a/fastvideo/third_party/magi_compiler b/fastvideo/third_party/magi_compiler new file mode 160000 index 000000000..5950612dd --- /dev/null +++ b/fastvideo/third_party/magi_compiler @@ -0,0 +1 @@ +Subproject commit 5950612ddf1205f9ba9c3238a8f02a078023e15c diff --git a/scripts/checkpoint_conversion/convert_magi2_to_fastvideo.py b/scripts/checkpoint_conversion/convert_magi2_to_fastvideo.py new file mode 100644 index 000000000..2da7547fb --- /dev/null +++ b/scripts/checkpoint_conversion/convert_magi2_to_fastvideo.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Create a FastVideo MAGI-2 repository from the official checkpoint layout. + +The published ``sand-ai/MAGI-2-preview`` snapshot already stores tensors in the +names and dtypes that the FastVideo-native loaders consume. This converter keeps +the tensor files byte-for-byte and changes only the component directory layout. +Hard links avoid a second 286 GiB allocation while producing regular files that +Hugging Face upload tools can publish as a self-contained repository. + +Example: + python scripts/checkpoint_conversion/convert_magi2_to_fastvideo.py \ + --source official_weights/magi2 \ + --output converted_weights/magi2 +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + + +SOURCE_REPOSITORY = "sand-ai/MAGI-2-preview" +SOURCE_REVISION = "2dea51b64db47ee5b4402d36fd90829a0c58913b" + +COMPONENT_DIRECTORY_MAPPING: dict[str, str] = { + "preview": "transformer", + "refiner": "transformer_2", + "text_encoder": "text_encoder", + "vae": "image_encoder", + "turbo_vae": "vae", + "stable-audio-open-1.0": "audio_vae", +} + +REQUIRED_COMPONENT_FILES: dict[str, tuple[str, ...]] = { + "preview": ("model.safetensors.index.json",), + "refiner": ("model.safetensors.index.json",), + "text_encoder": ("config.json", "model.safetensors.index.json", "tokenizer.json"), + "vae": ("Wan2.2_VAE.pth",), + "turbo_vae": ("TurboV3-Wan22-TinyShallow_7_7.json", "checkpoint.ckpt"), + "stable-audio-open-1.0": ("model_config.json", "model.safetensors"), +} + +# The distilled decoder checkpoint contains these feature-matching heads for +# training. MAGI-2 inference strictly loads the 88 ``decoder.*`` tensors. +TURBO_VAE_SKIPPED_KEYS: tuple[str, ...] = ( + "aligned_feature_projection_heads.0.0.conv.weight", + "aligned_feature_projection_heads.0.0.conv.bias", + "aligned_feature_projection_heads.0.1.conv.weight", + "aligned_feature_projection_heads.0.1.conv.bias", + "aligned_feature_projection_heads.1.0.conv.weight", + "aligned_feature_projection_heads.1.0.conv.bias", + "aligned_feature_projection_heads.1.1.conv.weight", + "aligned_feature_projection_heads.1.1.conv.bias", +) + +MODEL_INDEX: dict[str, object] = { + "_class_name": "Magi2Pipeline", + "_diffusers_version": "0.37.0", + "transformer": ["fastvideo.models.dits.magi2", "Magi2PreviewDiT"], + "transformer_2": ["fastvideo.models.dits.magi2_refiner", "Magi2RefinerDiT"], + "text_encoder": ["fastvideo.models.encoders.qwen3_5", "Magi2Qwen35TextEncoder"], + "image_encoder": ["fastvideo.models.vaes.magi2_wan_loader", "Magi2WanImageEncoder"], + "vae": ["fastvideo.models.vaes.magi2_turbo_vae", "Magi2TurboVAEModel"], + "audio_vae": ["fastvideo.models.vaes.magi2_audio_vae", "Magi2AudioVAE"], + "scheduler": [ + "fastvideo.models.schedulers.scheduling_flow_unipc_multistep", + "FlowUniPCMultistepScheduler", + ], +} + +SCHEDULER_CONFIG: dict[str, object] = { + "_class_name": "FlowUniPCMultistepScheduler", + "_diffusers_version": "0.37.0", + "disable_corrector": [], + "lower_order_final": True, + "num_train_timesteps": 1000, + "predict_x0": True, + "prediction_type": "flow_prediction", + "shift": 1.0, + "solver_order": 2, + "solver_type": "bh2", +} + +TRANSFORMER_CONFIGS: dict[str, dict[str, object]] = { + "transformer": { + "_class_name": "Magi2PreviewDiT", + "_diffusers_version": "0.37.0", + "source_subfolder": "preview", + }, + "transformer_2": { + "_class_name": "Magi2RefinerDiT", + "_diffusers_version": "0.37.0", + "source_subfolder": "refiner", + }, +} + + +def _validate_source(source: Path) -> None: + """Require component metadata and every shard referenced by an index.""" + missing_files = [ + str(source / component / relative_path) + for component, relative_paths in REQUIRED_COMPONENT_FILES.items() + for relative_path in relative_paths + if not (source / component / relative_path).is_file() + ] + if missing_files: + raise FileNotFoundError( + "The MAGI-2 source snapshot is incomplete; missing: " + + ", ".join(missing_files) + ) + + missing_shards: list[str] = [] + for component, relative_paths in REQUIRED_COMPONENT_FILES.items(): + for relative_path in relative_paths: + if not relative_path.endswith(".index.json"): + continue + index_path = source / component / relative_path + index_payload = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = index_payload.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"Checkpoint index has no weight_map object: {index_path}") + if not weight_map or not all(isinstance(shard_name, str) for shard_name in weight_map.values()): + raise ValueError(f"Checkpoint index has an invalid weight_map object: {index_path}") + shard_names = set(weight_map.values()) + missing_shards.extend( + str(index_path.parent / shard_name) + for shard_name in sorted(shard_names) + if not (index_path.parent / shard_name).is_file() + ) + if missing_shards: + raise FileNotFoundError( + "The MAGI-2 source snapshot is missing indexed checkpoint shards: " + + ", ".join(missing_shards) + ) + + +def _hardlink_component(source: Path, destination: Path) -> None: + """Replicate one component tree with regular-file hard links.""" + destination.mkdir(parents=True) + for source_path in sorted(source.rglob("*")): + relative_path = source_path.relative_to(source) + destination_path = destination / relative_path + if source_path.is_dir(): + destination_path.mkdir() + continue + if not source_path.is_file(): + raise ValueError(f"Unsupported checkpoint entry: {source_path}") + os.link(source_path, destination_path) + + +def _write_json(path: Path, payload: object) -> None: + """Write stable, reviewable JSON with a trailing newline.""" + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def _conversion_manifest() -> dict[str, object]: + """Describe directory and tensor-name transformations for every component.""" + return { + "format_version": 1, + "source": { + "repository": SOURCE_REPOSITORY, + "revision": SOURCE_REVISION, + }, + "component_directory_mapping": COMPONENT_DIRECTORY_MAPPING, + "tensor_mapping": { + "transformer": "identity; expert_bias_ema replaces expert_bias during inference loading", + "transformer_2": "identity", + "text_encoder": "identity; Transformers Qwen3.5 loads the language-model tensors", + "image_encoder": "identity; ignore decoder.* and conv2.* tensors", + "vae": "strip module.; load decoder.* tensors", + "audio_vae": "strip pretransform.model.; map the published sequential decoder to OobleckDecoder", + }, + "skipped_checkpoint_keys": { + "vae/checkpoint.ckpt": list(TURBO_VAE_SKIPPED_KEYS), + }, + } + + +def convert_checkpoint_layout(source: Path, output: Path) -> None: + """Hard-link official components and add FastVideo repository metadata.""" + source = source.resolve() + output = output.resolve() + _validate_source(source) + if output.exists(): + raise FileExistsError(f"Output path already exists: {output}") + output.mkdir(parents=True) + + for source_name, destination_name in COMPONENT_DIRECTORY_MAPPING.items(): + _hardlink_component(source / source_name, output / destination_name) + + scheduler_dir = output / "scheduler" + scheduler_dir.mkdir() + _write_json(scheduler_dir / "scheduler_config.json", SCHEDULER_CONFIG) + for component_name, config in TRANSFORMER_CONFIGS.items(): + _write_json(output / component_name / "config.json", config) + _write_json(output / "model_index.json", MODEL_INDEX) + _write_json(output / "magi2_conversion_manifest.json", _conversion_manifest()) + + +def main() -> None: + """Parse source and destination paths and create the converted repository.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--source", + type=Path, + required=True, + help="Local snapshot of sand-ai/MAGI-2-preview.", + ) + parser.add_argument( + "--output", + type=Path, + required=True, + help="Destination FastVideo model repository; the path must not exist.", + ) + arguments = parser.parse_args() + convert_checkpoint_layout(arguments.source, arguments.output) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/magi2/PORT_STATUS.md b/tests/local_tests/magi2/PORT_STATUS.md new file mode 100644 index 000000000..db5310582 --- /dev/null +++ b/tests/local_tests/magi2/PORT_STATUS.md @@ -0,0 +1,168 @@ +# MAGI-2 Preview Port Status + +## Completion + +The FastVideo MAGI-2 Preview port supports text-to-video (T2V), image-to-video +(I2V), and joint audio generation. Strict release-profile parity passes against +the official SandAI implementation for both modalities. + +| Field | Value | +| ---------------------- | ------------------------------------------------------------- | +| Branch | `model/port-magi-2` | +| Base commit | `1c04ace57351f7340e8d1a2e8b8f62180856ed16` | +| Official code revision | `073c84f2102ec3c9287623113a103c14402770ad` | +| Checkpoint revision | `2dea51b64db47ee5b4402d36fd90829a0c58913b` | +| Workloads | T2V and I2V with generated audio | +| Release profile | 100 preview steps, 5 refiner steps, seed 42 | +| Distributed topology | 8 Hopper GPUs; sequence, context, and expert parallel width 8 | +| Numerical requirement | Exact shape, dtype, stride, and tensor bytes | +| Full pipeline result | Exact for every captured stage and both decoded outputs | +| Video output | `uint8 [249, 1088, 1920, 3]` at 25 frames per second | +| Audio output | `float32 [441000, 2]` at 44.1 kHz | + +## Component Matrix + +| Component | FastVideo implementation | Strict parity result | +| ------------------------ | ---------------------------------------------------------------- | -------------------------------------------- | +| Qwen3.5 prompt encoder | `fastvideo/models/encoders/qwen3_5.py` | Exact | +| Wan 2.2 image encoder | `fastvideo/models/vaes/magi2_wan_loader.py` | Exact | +| Preview input packing | `fastvideo/pipelines/basic/magi2/stages/preview_data_proxy.py` | Exact | +| Preview transformer | `fastvideo/models/dits/magi2.py` | Exact across 40 layers | +| Refiner input packing | `fastvideo/pipelines/basic/magi2/stages/refiner_data_proxy.py` | Exact | +| Refiner transformer | `fastvideo/models/dits/magi2_refiner.py` | Exact across 30 layers | +| Flow UniPC scheduler | `fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py` | Exact | +| Turbo VAE decoder | `fastvideo/models/vaes/magi2_turbo_vae.py` | Exact across first, middle, and last windows | +| Stable Audio VAE decoder | `fastvideo/models/vaes/magi2_audio_vae.py` | Exact before and after resampling | +| T2V pipeline | `fastvideo/pipelines/basic/magi2/magi2_pipeline.py` | Exact, 100 + 5 steps | +| I2V pipeline | `fastvideo/pipelines/basic/magi2/magi2_pipeline.py` | Exact, 100 + 5 steps | + +## Architecture Fidelity + +The preview `Transformer` preserves the 40-layer hierarchical head-parallel +design described in sections 3.1 and 3.2 of the SandAI MAGI-2 report. The +implementation exchanges fixed head activations between ranks, shards each +layer's experts across the local high-bandwidth interconnect, and uses the +12-head, 256-expert, top-6 MagiMoE routing layout. Routing runs in FP32, expert +computation runs in BF16, and each expert uses the SwiGLU7 feed-forward layout. +The strict preview parity test compares every layer boundary across eight ranks. + +The refiner preserves the 30-layer local-attention design, including the exact +temporal and spatial attention ranges for the 1080p latent grid. The strict +refiner parity test compares every layer boundary across eight ranks. + +## Full-Profile Output Digests + +The official and FastVideo manifests differ only in the `implementation` +label. Every stage-boundary digest and output digest is identical. + +| Workload | Tensor | Shape | SHA-256 | +| -------- | ------------- | ---------------------- | ------------------------------------------------------------------ | +| T2V | Decoded video | `[249, 1088, 1920, 3]` | `0326570a07353cc78d488117f265b881ef6c681a3d660dc83acce2222b62e9a3` | +| T2V | Decoded audio | `[441000, 2]` | `fa62b45da7cfd055829ba79d1fe24d09736ed38e1075098aca71c2d0baad1937` | +| I2V | Decoded video | `[249, 1088, 1920, 3]` | `fd211e0f647e2b3304f636202940b19ec138b520abdfd4bfabba74678097edc9` | +| I2V | Decoded audio | `[441000, 2]` | `32d578e901d9443d21b3659b068992e112e6430c97e1a56231a99499f6fa3db0` | + +The full-profile manifests are stored at: + +- `archived/magi2_parity/validation/pipeline/official/capture.json` +- `archived/magi2_parity/validation/pipeline/fastvideo/capture.json` + +## Compiled-Path Verification + +The official and FastVideo pipelines also ran from separate empty node-local +MagiCompiler caches with one preview step and one refiner step. Their T2V and +I2V manifests match exactly at every captured stage and decoded output. This +reduced schedule verifies compilation and production wiring; the 100 + 5 step +comparison provides the full release-profile numerical result. + +| Workload | Tensor | SHA-256 | +| -------- | ------------- | ------------------------------------------------------------------ | +| T2V | Decoded video | `ca379741f7af306cdc2d2fbbd38dfabfa35f215b613b8c7fe8b9193dcef154f5` | +| T2V | Decoded audio | `2e44652e3c4047560255fac37b606ecc26e44d140ce500ee0a7f6328cf30c22e` | +| I2V | Decoded video | `354a56ace8504a61487f299867c4a31beee2fd60fcce7f5bc1eb6acefdf9f532` | +| I2V | Decoded audio | `622420e1a0777b16804ec1a8e1b53ff5a2ac2dad7edae7823b8d39a8377f67aa` | + +The compiled manifests are stored under +`archived/magi2_parity/validation/pipeline_compiled/`. + +## Checkpoint Conversion + +`scripts/checkpoint_conversion/convert_magi2_to_fastvideo.py` converts the +official mixed checkpoint layout into the component layout consumed by the +FastVideo registry. The converter validates every indexed shard and uses hard +links for tensor files so that both layouts share the same file data. + +| Official component | FastVideo component | Role | +| ------------------------ | ------------------- | ---------------------------------------- | +| `preview/` | `transformer/` | Joint video-audio preview transformer | +| `refiner/` | `transformer_2/` | 1080p video refiner | +| `text_encoder/` | `text_encoder/` | Qwen3.5 prompt encoder and tokenizer | +| `vae/Wan2.2_VAE.pth` | `image_encoder/` | Wan reference-image encoder | +| `turbo_vae/` | `vae/` | Default distilled sliding-window decoder | +| `stable-audio-open-1.0/` | `audio_vae/` | Stable Audio VAE decoder | + +The preview, refiner, Turbo VAE, Wan VAE, and audio VAE production loaders +reject missing tensors, unexpected tensors, and shape mismatches. The Qwen3.5 +loader selects the text backbone from the multimodal checkpoint and leaves the +visual tower and language-model head unused, matching the release pipeline. + +## Runtime Contract + +- Turbo VAE decoding is the default video-decoding path. The decoder uses a + seven-latent first window and seven-latent subsequent windows on one context- + parallel leader rank per video. +- `--deterministic` and `MAGI2_DETERMINISTIC=1` set fixed Python, NumPy, + PyTorch, and CUDA seeds and activate deterministic MAGI-2 and MagiAttention + kernels. +- `MAGI2_SAVE_LATENT_PATH` writes each sample's post-refiner latent from the + leader rank. An empty environment value disables latent saving. +- I2V conditioning preserves the official BF16 image round trip before the Wan + encoder performs FP32 encoding. +- Random tensors are drawn in the official order: preview video noise, preview + audio noise, and refiner video noise. +- The pipeline returns decoded video and generated audio tensors before media + container encoding. + +## Pinned Dependencies + +| Dependency | Repository | Revision | +| ---------------------- | ---------------------------------------------- | ------------------------------------------ | +| MagiAttention | `https://github.com/SandAI-org/MagiAttention` | `2c6413571c2cac6a80d1f85a434c6713fe0f5286` | +| MagiCompiler | `https://github.com/SandAI-org/MagiCompiler` | `5950612ddf1205f9ba9c3238a8f02a078023e15c` | +| Flash Attention Hopper | `https://github.com/Dao-AILab/flash-attention` | `b613d9e2c8475945baff3fd68f2030af1b890acf` | + +The revisions match the build arguments in the published MAGI-2 Dockerfile. +The dependencies live under `fastvideo/third_party/` and are declared in the +repository's `.gitmodules` file. + +## Validation Coverage + +| Validation file | Coverage | +| ------------------------------------------ | ---------------------------------------------------------- | +| `test_magi2_text_encoder_parity.py` | Prompt formatting, CJK splitting, tokenization, embeddings | +| `test_magi2_image_encoder_parity.py` | I2V resize, BF16 round trip, Wan posterior mean | +| `test_magi2_preview_data_proxy_parity.py` | Preview packing, coordinates, and context padding | +| `test_magi2_preview_transformer_parity.py` | T2V and I2V boundaries across all 40 layers | +| `test_magi2_refiner_data_proxy_parity.py` | Local-attention ranges and refiner packing | +| `test_magi2_refiner_transformer_parity.py` | Boundaries across all 30 refiner layers | +| `test_magi2_scheduler_parity.py` | Preview and refiner scheduler arrays and updates | +| `test_magi2_turbo_vae_parity.py` | First, middle, and last temporal decode windows | +| `test_magi2_audio_vae_parity.py` | Stable Audio VAE decoding and 44.1 kHz resampling | +| `test_magi2_runtime_controls.py` | Determinism and leader-only latent saving | +| `test_magi2_checkpoint_conversion.py` | Component mapping and indexed-shard validation | +| `test_magi2_registry_and_metadata.py` | T2V and I2V registry routing | +| `test_magi2_pipeline_parity.py` | Full T2V and I2V stage and decoded-output parity | + +The validation suite and environment setup records are described in +`tests/local_tests/magi2/README.md`. + +## Decisions + +| Decision | Reason | +| ---------------------------------------- | ----------------------------------------------------------- | +| Branch from local `main` | The requested base is local `main` commit `1c04ace...`. | +| Preserve the official dependency pins | Kernel and compiler revisions affect strict parity. | +| Use Turbo VAE by default | The official release and port contract select this decoder. | +| Compare tensors before MP4/AAC encoding | Media codecs operate outside model inference. | +| Use eager mode for full byte-parity runs | Eager execution isolates model math from compiler caching. | +| Keep MagiCompiler enabled in production | The pinned compiler provides the official optimized path. | diff --git a/tests/local_tests/magi2/README.md b/tests/local_tests/magi2/README.md new file mode 100644 index 000000000..f058e9589 --- /dev/null +++ b/tests/local_tests/magi2/README.md @@ -0,0 +1,167 @@ +# MAGI-2 Preview Local Port Validation + +This directory contains the reproducible local validation suite for the MAGI-2 +Preview text-to-video (T2V) and image-to-video (I2V) port. The validation target +is exact tensor equality with the official SandAI implementation for model +components, pipeline stages, terminal video tensors, and terminal audio tensors. + +## Sources + +| Source | Location | Revision | +| ------------------------- | -------------------------------------------------------------------------------------------- | ------------------------------------------ | +| Official inference code | `https://github.com/SandAI-org/MAGI-2-preview` | `073c84f2102ec3c9287623113a103c14402770ad` | +| Official local checkout | `/mnt/weka/shrd/wm/junda/fv-hub/MAGI-2-preview` | `073c84f2102ec3c9287623113a103c14402770ad` | +| Published checkpoints | `https://huggingface.co/sand-ai/MAGI-2-preview` | `2dea51b64db47ee5b4402d36fd90829a0c58913b` | +| Local checkpoint snapshot | `official_weights/magi2` | `2dea51b64db47ee5b4402d36fd90829a0c58913b` | +| MagiAttention | `fastvideo/third_party/magi_attention` | `2c6413571c2cac6a80d1f85a434c6713fe0f5286` | +| MagiCompiler | `fastvideo/third_party/magi_compiler` | `5950612ddf1205f9ba9c3238a8f02a078023e15c` | +| Flash Attention Hopper | `https://github.com/Dao-AILab/flash-attention/tree/b613d9e2c8475945baff3fd68f2030af1b890acf` | `b613d9e2c8475945baff3fd68f2030af1b890acf` | + +The Hugging Face repository is public and does not require an access token. No +Hugging Face token environment variable is used. + +## Checkpoint Layout + +The checkpoint repository has a custom mixed layout without a root +`model_index.json` file. The converter maps the official component directories +into a FastVideo model repository and writes the component metadata. The tensor +files remain byte-identical hard links. + +| Checkpoint path | Pipeline role | +| ------------------------ | ------------------------------------------------------------------ | +| `preview/` | Joint video-audio preview transformer. | +| `refiner/` | 1080p video refiner transformer. | +| `text_encoder/` | Qwen3.5-27B prompt encoder and tokenizer. | +| `vae/Wan2.2_VAE.pth` | Wan 2.2 reference-image encoder for I2V. | +| `turbo_vae/` | Default distilled decoder; first window 7 latents, step 7 latents. | +| `stable-audio-open-1.0/` | Stable Audio variational autoencoder (VAE) audio decoder. | + +The checkpoint snapshot contains 306,721,986,476 bytes. The converted repository +uses hard links, so the source and converted layouts share the same file data on +the local filesystem. + +Run the converter with: + +```bash +venv-port-magi-2/bin/python \ + scripts/checkpoint_conversion/convert_magi2_to_fastvideo.py \ + --source official_weights/magi2 \ + --output converted_weights/magi2 +``` + +## Environment + +The Python environment is `venv-port-magi-2/` at the FastVideo worktree root. +It uses Python 3.12, PyTorch 2.11 with CUDA 12.8, and one numeric software stack +for the official implementation and the FastVideo implementation. + +Setup commands: + +```bash +uv venv --python /mnt/weka/home/junda.su/.local/bin/python3.12 venv-port-magi-2 +UV_TORCH_BACKEND=cu128 uv pip install \ + --python venv-port-magi-2/bin/python \ + pip setuptools wheel ninja torch==2.11.0 torchvision torchaudio +``` + +The official code imports Flash Attention 3 through `flash_attn_interface`. +There is no wheel for the official pinned revision and this Python, CUDA, and +PyTorch combination. Build revision +`b613d9e2c8475945baff3fd68f2030af1b890acf` with `MAX_JOBS=16`. Do not increase +the build job count. + +Initialize the pinned MagiAttention dependency before installation: + +```bash +git -C fastvideo/third_party/magi_attention submodule update --init --recursive +``` + +The environment setup record and command logs are stored under +`archived/magi2_parity/env_setup/`. + +## Official Execution Contract + +Use eight Hopper GPUs. Set both deterministic environment variables for strict +official parity because the official `--deterministic` implementation does not +set the environment variable that controls deterministic mixture-of-experts +(MoE) scatter: + +```bash +MAGI2_DETERMINISTIC=1 \ +MAGI_ATTENTION_DETERMINISTIC_MODE=1 \ +MAGI2_SAVE_LATENT_PATH=archived/magi2_parity/validation/reference_latents \ +torchrun --nproc_per_node=8 inference/pipeline/entry.py \ + --resolution 1080p \ + --seconds 10 \ + --seed 42 \ + --prompt-file assets/sample_000.txt \ + --output archived/magi2_parity/validation/reference_t2v +``` + +Add `--image assets/sample_000.jpeg` for I2V. The FastVideo deterministic option +sets fixed Python, NumPy, PyTorch, and CUDA seeds and activates both MAGI-2 +kernel determinism controls. + +## Required Parity Coverage + +Each component test loads the official component and the production FastVideo +component. A reused FastVideo component requires a non-skipped parity pass. +Strict parity compares shape, dtype, stride, and every tensor value. + +| Scope | Test | +| --------------------------- | ------------------------------------------ | +| Prompt and Qwen3.5 encoding | `test_magi2_text_encoder_parity.py` | +| I2V Wan VAE encoding | `test_magi2_image_encoder_parity.py` | +| Preview token packing | `test_magi2_preview_data_proxy_parity.py` | +| Preview transformer | `test_magi2_preview_transformer_parity.py` | +| Refiner token packing | `test_magi2_refiner_data_proxy_parity.py` | +| Refiner transformer | `test_magi2_refiner_transformer_parity.py` | +| Flow UniPC scheduler | `test_magi2_scheduler_parity.py` | +| Turbo VAE decoding | `test_magi2_turbo_vae_parity.py` | +| Stable Audio VAE decoding | `test_magi2_audio_vae_parity.py` | +| Runtime controls | `test_magi2_runtime_controls.py` | +| Checkpoint conversion | `test_magi2_checkpoint_conversion.py` | +| Registry and metadata | `test_magi2_registry_and_metadata.py` | +| T2V and I2V pipeline stages | `test_magi2_pipeline_parity.py` | +| Decoded video and audio | `test_magi2_pipeline_parity.py` | + +Run the local suite with: + +```bash +venv-port-magi-2/bin/python -m pytest tests/local_tests/magi2 -sv +``` + +The component tests compare prompt normalization, tokenization, image encoding, +scheduler arrays, packed model inputs, every preview and refiner transformer +layer boundary, all three Turbo VAE temporal-window roles, audio decoding before +resampling, and resampled audio. The pipeline test compares stage-boundary +digests for conditioned prompts, reference-image latents, text embeddings, +initial video and audio noise, preview outputs, the noise-injected refiner input, +the post-refiner latent, and decoded outputs. End-to-end parity compares the +decoded `uint8 [T, H, W, 3]` video array and the sample-major stereo audio array +before media encoding. + +## Release-Profile Result + +The official and FastVideo workers ran both T2V and I2V with seed 42, 100 +preview steps, 5 refiner steps, and deterministic MAGI-2 and MagiAttention +kernels. The manifest comparison found one expected metadata difference: the +`implementation` label. All captured tensor metadata and SHA-256 digests match. + +| Artifact | Path | +| ------------------- | ----------------------------------------------------------------------- | +| Official manifest | `archived/magi2_parity/validation/pipeline/official/capture.json` | +| FastVideo manifest | `archived/magi2_parity/validation/pipeline/fastvideo/capture.json` | +| Completion report | `tests/local_tests/magi2/PORT_STATUS.md` | +| Environment record | `archived/magi2_parity/env_setup/SETUP_SUMMARY.md` | + +## Review Requirements + +The port is ready for review when all required component tests and the combined +T2V/I2V pipeline test report non-skipped strict parity passes. Review also +verifies strict checkpoint loading, Turbo VAE default selection, one-rank-per- +video decoding, deterministic controls, post-refiner latent saving through +`MAGI2_SAVE_LATENT_PATH`, T2V behavior, I2V behavior, and audio output. + +Generated tensors, checkpoints, converted weights, compiler caches, and run +logs remain outside version control. diff --git a/tests/local_tests/magi2/_parity_utils.py b/tests/local_tests/magi2/_parity_utils.py new file mode 100644 index 000000000..5e7a6fe81 --- /dev/null +++ b/tests/local_tests/magi2/_parity_utils.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared utilities for MAGI-2 component parity scaffolds.""" +from __future__ import annotations + +import importlib +import json +import os +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_REF_DIR = Path( + os.getenv( + "MAGI2_OFFICIAL_REF_DIR", + "/mnt/weka/shrd/wm/junda/fv-hub/MAGI-2-preview", + ) +) +LOCAL_WEIGHTS_DIR = Path( + os.getenv("MAGI2_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "magi2") +) + + +def require_path(path: Path, purpose: str) -> Path: + """Return a required local path or skip with its concrete purpose.""" + if not path.exists(): + pytest.skip(f"Missing {purpose}: {path}") + return path + + +def require_complete_safetensor_index(component_dir: Path) -> Path: + """Require every shard named by a component's safetensors index.""" + index_path = require_path( + component_dir / "model.safetensors.index.json", + "safetensors index", + ) + with index_path.open(encoding="utf-8") as index_file: + weight_index = json.load(index_file) + shard_names = sorted(set(weight_index["weight_map"].values())) + missing_shards = [name for name in shard_names if not (component_dir / name).is_file()] + if missing_shards: + pytest.skip( + f"Incomplete weights in {component_dir}; missing shards: " + f"{', '.join(missing_shards)}" + ) + return component_dir + + +def import_official_module(module_name: str): + """Import one module from the user-specified official MAGI-2 checkout.""" + require_path(OFFICIAL_REF_DIR, "official MAGI-2 checkout") + reference_path = str(OFFICIAL_REF_DIR) + if reference_path not in sys.path: + sys.path.insert(0, reference_path) + try: + official_module = importlib.import_module(module_name) + except ModuleNotFoundError as exc: + if exc.name == module_name or module_name.startswith(f"{exc.name}."): + raise + pytest.skip( + f"Official module {module_name} requires missing dependency {exc.name!r}" + ) + module_path = Path(official_module.__file__).resolve() + if not module_path.is_relative_to(OFFICIAL_REF_DIR.resolve()): + raise RuntimeError( + f"Imported {module_name} from {module_path}, expected {OFFICIAL_REF_DIR}" + ) + return official_module + + +def assert_tensor_exact( + fastvideo_tensor: torch.Tensor, + official_tensor: torch.Tensor, + tensor_name: str, +) -> None: + """Require identical tensor metadata and identical values.""" + assert fastvideo_tensor.shape == official_tensor.shape, tensor_name + assert fastvideo_tensor.dtype == official_tensor.dtype, tensor_name + assert fastvideo_tensor.stride() == official_tensor.stride(), tensor_name + assert torch.equal(fastvideo_tensor, official_tensor), tensor_name + + +def assert_array_exact( + fastvideo_array: np.ndarray, + official_array: np.ndarray, + array_name: str, +) -> None: + """Require identical NumPy array metadata and identical values.""" + assert fastvideo_array.shape == official_array.shape, array_name + assert fastvideo_array.dtype == official_array.dtype, array_name + assert fastvideo_array.strides == official_array.strides, array_name + assert np.array_equal(fastvideo_array, official_array), array_name diff --git a/tests/local_tests/magi2/_pipeline_parity_worker.py b/tests/local_tests/magi2/_pipeline_parity_worker.py new file mode 100644 index 000000000..1e99bb17a --- /dev/null +++ b/tests/local_tests/magi2/_pipeline_parity_worker.py @@ -0,0 +1,502 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Capture strict MAGI-2 pipeline boundaries from one implementation.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +import random +import sys +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_ROOT = Path( + os.environ.get("MAGI2_OFFICIAL_REF_DIR", REPO_ROOT.parent / "MAGI-2-preview") +) +WEIGHTS_ROOT = Path( + os.environ.get("MAGI2_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "magi2") +) +CONVERTED_ROOT = Path( + os.environ.get("MAGI2_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / "magi2") +) +WORLD_SIZE = 8 +SEED = 42 + + +def _parse_args() -> argparse.Namespace: + """Parse the implementation, denoising counts, and requested modalities.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--implementation", + choices=("official", "fastvideo"), + required=True, + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--preview-steps", type=int, default=100) + parser.add_argument("--refiner-steps", type=int, default=5) + parser.add_argument( + "--cases", + nargs="+", + choices=("t2v", "i2v"), + default=("t2v", "i2v"), + ) + return parser.parse_args() + + +def _seed_all() -> None: + """Reset every random-number source used by the release pipeline.""" + random.seed(SEED) + np.random.seed(SEED) + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + + +def _enable_determinism() -> None: + """Apply the deterministic controls exposed by the official entry point.""" + os.environ["MAGI2_DETERMINISTIC"] = "1" + os.environ["MAGI_ATTENTION_DETERMINISTIC_MODE"] = "1" + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + _seed_all() + torch.use_deterministic_algorithms(True) + + +def _hash_array(array: np.ndarray) -> str: + """Hash the exact bytes of one contiguous NumPy array without a byte copy.""" + contiguous_array = np.ascontiguousarray(array) + return hashlib.sha256(memoryview(contiguous_array).cast("B")).hexdigest() + + +def _tensor_digest(tensor: torch.Tensor | None) -> dict[str, Any] | None: + """Describe and hash a tensor while preserving its source metadata.""" + if tensor is None: + return None + source_tensor = tensor.detach() + byte_array = source_tensor.to(device="cpu").contiguous().view(torch.uint8).numpy() + return { + "shape": list(source_tensor.shape), + "dtype": str(source_tensor.dtype), + "stride": list(source_tensor.stride()), + "sha256": _hash_array(byte_array), + } + + +def _array_digest(array: np.ndarray | None) -> dict[str, Any] | None: + """Describe and hash a NumPy array while preserving its source metadata.""" + if array is None: + return None + source_array = np.asarray(array) + return { + "shape": list(source_array.shape), + "dtype": str(source_array.dtype), + "stride": list(source_array.strides), + "sha256": _hash_array(source_array), + } + + +def _case_inputs(case_name: str) -> tuple[str, str | None]: + """Return the official release prompt and optional reference image path.""" + if case_name == "t2v": + prompt_path = OFFICIAL_ROOT / "assets" / "sample_enhanced_t2v.json" + return prompt_path.read_text(encoding="utf-8").strip(), None + if case_name == "i2v": + prompt_path = OFFICIAL_ROOT / "assets" / "sample_000.txt" + image_path = OFFICIAL_ROOT / "assets" / "sample_000.jpeg" + return prompt_path.read_text(encoding="utf-8").strip(), str(image_path) + raise ValueError(f"Unknown MAGI-2 parity case: {case_name}") + + +def _capture_official_case( + engine: Any, + case_name: str, + preview_steps: int, + refiner_steps: int, + is_capture_rank: bool, +) -> dict[str, Any]: + """Run the official engine and capture every externally visible stage boundary.""" + prompt, image_path = _case_inputs(case_name) + capture: dict[str, Any] = {} + text_call_count = 0 + + original_encode_images = engine._encode_images + original_get_text_embedding = engine.get_text_embedding + original_get_special_token = engine.get_special_token + original_preview_sample = engine.sampler.sample + original_refiner_forward = engine._forward_magi2_refiner + original_refiner_sample = engine.evaluate_magi2_refiner_with_latent + + def capture_images(image: Any, height: int, width: int): + """Record the Wan image-encoder output and its figure identity.""" + encoded_images = original_encode_images(image, height, width) + if is_capture_rank: + reference_latent, reference_length, reference_ids = encoded_images + capture["reference_latent"] = _tensor_digest(reference_latent) + capture["reference_length"] = _tensor_digest(reference_length) + capture["reference_ids"] = _tensor_digest(reference_ids) + return encoded_images + + def capture_text(prompt_text: str) -> torch.Tensor: + """Record positive and negative Qwen3.5 conditioning in call order.""" + nonlocal text_call_count + text_embedding = original_get_text_embedding(prompt_text) + if is_capture_rank: + if text_call_count == 0: + capture["conditioned_prompt"] = prompt_text + capture["positive_text"] = _tensor_digest(text_embedding) + elif text_call_count == 1: + capture["negative_prompt"] = prompt_text + capture["negative_text"] = _tensor_digest(text_embedding) + else: + raise RuntimeError("MAGI-2 encoded more than two prompts in one request") + text_call_count += 1 + return text_embedding + + def capture_special_token( + prompt_text: str, + figure_ids: torch.Tensor, + text_feature: torch.Tensor, + ) -> torch.Tensor: + """Record the pooled text embedding that prefixes image patches.""" + special_token = original_get_special_token( + prompt_text, + figure_ids, + text_feature, + ) + if is_capture_rank: + capture["special_tokens"] = _tensor_digest(special_token.unsqueeze(0)) + return special_token + + def capture_preview(sampler_input: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Record the initial noise and the completed preview latents.""" + if is_capture_rank: + capture["initial_video_noise"] = _tensor_digest(sampler_input.latent) + capture["initial_audio_noise"] = _tensor_digest( + sampler_input.audio_latent + ) + preview_video, preview_audio = original_preview_sample(sampler_input) + if is_capture_rank: + capture["preview_video"] = _tensor_digest(preview_video) + capture["preview_audio"] = _tensor_digest(preview_audio) + return preview_video, preview_audio + + def capture_refiner_forward( + latent_video: torch.Tensor, + latent_audio: torch.Tensor, + txt_feat: torch.Tensor, + ref_audio_feat: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Record the noise-injected 1080p latent entering the refiner.""" + if is_capture_rank and "refiner_input" not in capture: + capture["refiner_input"] = _tensor_digest(latent_video) + return original_refiner_forward( + latent_video, + latent_audio, + txt_feat, + ref_audio_feat, + ) + + def capture_refiner_sample(**kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Record the video latent that the refiner passes to the decoder.""" + refined_video, refined_audio = original_refiner_sample(**kwargs) + if is_capture_rank: + capture["refined_video"] = _tensor_digest(refined_video) + return refined_video, refined_audio + + engine._encode_images = capture_images + engine.get_text_embedding = capture_text + engine.get_special_token = capture_special_token + engine.sampler.sample = capture_preview + engine._forward_magi2_refiner = capture_refiner_forward + engine.evaluate_magi2_refiner_with_latent = capture_refiner_sample + _seed_all() + try: + video, audio = engine.evaluate( + prompt=prompt, + image=image_path, + eval_task_type="text2video" if image_path is None else "image2video", + seconds=10.0, + preview_width=896, + preview_height=512, + br_num_inference_steps=preview_steps, + refiner_width=1920, + refiner_height=1088, + magi2_refiner_num_inference_steps=refiner_steps, + ) + finally: + engine._encode_images = original_encode_images + engine.get_text_embedding = original_get_text_embedding + engine.get_special_token = original_get_special_token + engine.sampler.sample = original_preview_sample + engine._forward_magi2_refiner = original_refiner_forward + engine.evaluate_magi2_refiner_with_latent = original_refiner_sample + + if is_capture_rank: + capture.setdefault("special_tokens", None) + capture["decoded_video"] = _array_digest(video) + capture["decoded_audio"] = _array_digest(audio) + return capture + + +def _load_official_engine() -> Any: + """Initialize the official eight-rank runtime and load every release component.""" + sys.path.insert(0, str(OFFICIAL_ROOT)) + from inference.common.magi2_config import load_config + from inference.infra.checkpoint.load_checkpoint import ( + load_magi2_model, + load_magi2_refiner, + ) + from inference.infra.distributed import ( + initialize_expert_parallel, + initialize_model_parallel, + ) + + config = load_config(str(OFFICIAL_ROOT / "configs" / "magi2_refiner.json")) + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + if dist.get_world_size() != WORLD_SIZE: + raise RuntimeError( + f"MAGI-2 pipeline parity requires {WORLD_SIZE} ranks, " + f"received {dist.get_world_size()}" + ) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + initialize_model_parallel(cp_size=WORLD_SIZE) + initialize_expert_parallel(ep_size=WORLD_SIZE) + + from inference.pipeline.inference_engine import Magi2InferenceEngine + + preview_model = load_magi2_model(config) + refiner_model = load_magi2_refiner(config) + return Magi2InferenceEngine( + model=preview_model, + config=config.evaluation_config, + device=f"cuda:{local_rank}", + weight_dtype=torch.bfloat16, + magi2_refiner=refiner_model, + ) + + +def _canonical_fastvideo_video(video: torch.Tensor | None) -> np.ndarray | None: + """Convert FastVideo's BCHWT float output into the release's THWC bytes.""" + if video is None: + return None + if video.shape[0] != 1: + raise ValueError(f"MAGI-2 parity expects one decoded video, received {video.shape}") + return video[0].permute(1, 2, 3, 0).mul(255).numpy().astype(np.uint8) + + +def _capture_fastvideo_case( + pipeline: Any, + fastvideo_args: Any, + case_name: str, + preview_steps: int, + refiner_steps: int, + is_capture_rank: bool, +) -> dict[str, Any]: + """Run FastVideo stages and capture the boundaries used by the official worker.""" + from fastvideo.fastvideo_args import WorkloadType + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + prompt, image_path = _case_inputs(case_name) + fastvideo_args.workload_type = ( + WorkloadType.T2V if case_name == "t2v" else WorkloadType.I2V + ) + batch = ForwardBatch( + data_type="video", + prompt=prompt, + negative_prompt=pipeline.negative_prompt, + image_path=image_path, + seed=SEED, + num_frames=249, + height=1088, + width=1920, + fps=25, + num_inference_steps=preview_steps, + num_inference_steps_sr=refiner_steps, + ) + capture: dict[str, Any] = {} + refiner_stage = pipeline.refiner_stage + original_refiner_predict = refiner_stage._predict_velocity + + def capture_refiner_predict( + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + text_context: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Record the noise-injected latent before the first refiner forward.""" + if is_capture_rank and "refiner_input" not in capture: + capture["refiner_input"] = _tensor_digest(video_latent) + return original_refiner_predict(video_latent, audio_latent, text_context) + + refiner_stage._predict_velocity = capture_refiner_predict + _seed_all() + try: + for stage_name, stage in pipeline._stage_name_mapping.items(): + batch = stage(batch, fastvideo_args) + if not is_capture_rank: + continue + if stage_name == "reference_image_stage": + capture["conditioned_prompt"] = batch.prompt + capture["reference_latent"] = _tensor_digest( + batch.magi2_ref_image_feat + ) + capture["reference_length"] = _tensor_digest( + batch.magi2_ref_image_feat_len + ) + reference_ids = ( + None + if batch.magi2_ref_image_feat is None + else torch.tensor([[1]], dtype=torch.long) + ) + capture["reference_ids"] = _tensor_digest(reference_ids) + elif stage_name == "text_encoding_stage": + capture["negative_prompt"] = batch.negative_prompt + capture["positive_text"] = _tensor_digest(batch.magi2_text_context) + capture["negative_text"] = _tensor_digest( + batch.magi2_negative_context + ) + capture["special_tokens"] = _tensor_digest( + batch.magi2_ref_image_special_tokens + ) + elif stage_name == "latent_preparation_stage": + capture["initial_video_noise"] = _tensor_digest(batch.latents) + capture["initial_audio_noise"] = _tensor_digest( + batch.audio_latents + ) + elif stage_name == "preview_denoising_stage": + capture["preview_video"] = _tensor_digest(batch.latents) + capture["preview_audio"] = _tensor_digest(batch.audio_latents) + elif stage_name == "refiner_stage": + capture["refined_video"] = _tensor_digest(batch.latents) + elif stage_name == "video_decoding_stage": + canonical_video = _canonical_fastvideo_video(batch.output) + capture["decoded_video"] = _array_digest(canonical_video) + elif stage_name == "audio_decoding_stage": + capture["decoded_audio"] = _array_digest(batch.extra.get("audio")) + finally: + refiner_stage._predict_velocity = original_refiner_predict + return capture + + +def _load_fastvideo_pipeline() -> tuple[Any, Any]: + """Initialize FastVideo's eight-rank runtime and load the converted checkpoint.""" + sys.path.insert(0, str(REPO_ROOT)) + from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType + from fastvideo.pipelines.basic.magi2.magi2_pipeline import Magi2Pipeline + from fastvideo.pipelines.basic.magi2.pipeline_configs import ( + Magi2PreviewPipelineConfig, + ) + from fastvideo.pipelines.basic.magi2.presets import MAGI2_NEGATIVE_PROMPT + + fastvideo_args = FastVideoArgs( + model_path=str(CONVERTED_ROOT), + workload_type=WorkloadType.T2V, + num_gpus=WORLD_SIZE, + tp_size=1, + sp_size=WORLD_SIZE, + pipeline_config=Magi2PreviewPipelineConfig(), + deterministic=True, + dit_cpu_offload=True, + dit_layerwise_offload=False, + text_encoder_cpu_offload=True, + image_encoder_cpu_offload=True, + vae_cpu_offload=True, + enable_stage_verification=True, + ) + pipeline = Magi2Pipeline(str(CONVERTED_ROOT), fastvideo_args) + pipeline.negative_prompt = MAGI2_NEGATIVE_PROMPT + pipeline.post_init() + if dist.get_world_size() != WORLD_SIZE: + raise RuntimeError( + f"MAGI-2 pipeline parity requires {WORLD_SIZE} ranks, " + f"received {dist.get_world_size()}" + ) + return pipeline, fastvideo_args + + +def _write_capture( + output_dir: Path, + implementation: str, + preview_steps: int, + refiner_steps: int, + captures: dict[str, Any], +) -> None: + """Atomically write the rank-zero digest manifest for one implementation.""" + output_dir.mkdir(parents=True, exist_ok=True) + artifact = { + "schema_version": 1, + "implementation": implementation, + "world_size": WORLD_SIZE, + "preview_steps": preview_steps, + "refiner_steps": refiner_steps, + "cases": captures, + } + artifact_path = output_dir / "capture.json" + temporary_path = output_dir / "capture.json.tmp" + temporary_path.write_text( + json.dumps(artifact, indent=2, sort_keys=True), + encoding="utf-8", + ) + os.replace(temporary_path, artifact_path) + + +def main() -> None: + """Load one implementation, capture all cases, and persist rank-zero digests.""" + args = _parse_args() + if args.preview_steps <= 0 or args.refiner_steps <= 0: + raise ValueError("MAGI-2 parity step counts must be positive") + os.environ["MAGI2_CKPT_ROOT"] = str(WEIGHTS_ROOT) + os.environ.pop("MAGI2_SAVE_LATENT_PATH", None) + os.environ.pop("NEGATIVE_PROMPT", None) + os.environ.pop("SKIP_LOAD_MODEL", None) + _enable_determinism() + + if args.implementation == "official": + engine = _load_official_engine() + is_capture_rank = dist.get_rank() == 0 + captures = { + case_name: _capture_official_case( + engine, + case_name, + args.preview_steps, + args.refiner_steps, + is_capture_rank, + ) + for case_name in args.cases + } + else: + pipeline, fastvideo_args = _load_fastvideo_pipeline() + is_capture_rank = dist.get_rank() == 0 + captures = { + case_name: _capture_fastvideo_case( + pipeline, + fastvideo_args, + case_name, + args.preview_steps, + args.refiner_steps, + is_capture_rank, + ) + for case_name in args.cases + } + + if is_capture_rank: + _write_capture( + args.output_dir, + args.implementation, + args.preview_steps, + args.refiner_steps, + captures, + ) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/magi2/_preview_transformer_parity_worker.py b/tests/local_tests/magi2/_preview_transformer_parity_worker.py new file mode 100644 index 000000000..c9ae9dc4e --- /dev/null +++ b/tests/local_tests/magi2/_preview_transformer_parity_worker.py @@ -0,0 +1,364 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run one isolated side of MAGI-2 preview transformer numerical parity.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +from pathlib import Path +import random +import sys +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_ROOT = Path( + os.environ.get("MAGI2_OFFICIAL_REF_DIR", REPO_ROOT.parent / "MAGI-2-preview") +) +WEIGHTS_ROOT = Path( + os.environ.get("MAGI2_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "magi2") +) +WORLD_SIZE = 8 +SEED = 42 + + +def _parse_args() -> argparse.Namespace: + """Parse the implementation and artifact directory for one torchrun job.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--implementation", + choices=("official", "fastvideo"), + required=True, + ) + parser.add_argument("--output-dir", type=Path, required=True) + return parser.parse_args() + + +def _enable_determinism() -> None: + """Apply the deterministic controls used by the official entry point.""" + os.environ["MAGI2_DETERMINISTIC"] = "1" + os.environ["MAGI_ATTENTION_DETERMINISTIC_MODE"] = "1" + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + random.seed(SEED) + np.random.seed(SEED) + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + torch.use_deterministic_algorithms(True) + + +def _initialize_distributed(implementation: str) -> tuple[int, torch.device]: + """Initialize the official eight-rank CP and EP process-group topology.""" + if dist.is_initialized(): + raise RuntimeError("The parity worker requires a fresh distributed process") + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + if dist.get_world_size() != WORLD_SIZE: + raise RuntimeError( + f"MAGI-2 preview parity requires {WORLD_SIZE} ranks, " + f"received {dist.get_world_size()}" + ) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + + if implementation == "official": + sys.path.insert(0, str(OFFICIAL_ROOT)) + from inference.infra.distributed import ( + initialize_expert_parallel, + initialize_model_parallel, + ) + else: + sys.path.insert(0, str(REPO_ROOT)) + from fastvideo.models.dits.magi2_runtime import ( + initialize_expert_parallel, + initialize_model_parallel, + ) + + initialize_model_parallel(cp_size=WORLD_SIZE) + initialize_expert_parallel(ep_size=WORLD_SIZE) + return rank, torch.device("cuda", local_rank) + + +def _load_runtime(implementation: str, device: torch.device) -> tuple[Any, Any, type]: + """Load one implementation's production transformer, proxy, and input type.""" + if implementation == "official": + from inference.common.magi2_config import load_config + from inference.infra.checkpoint.load_checkpoint import load_magi2_model + from inference.pipeline.preview_data_proxy import Magi2DataProxy, ModelInput + + config = load_config(str(OFFICIAL_ROOT / "configs" / "magi2_preview.json")) + config.engine_config.load = str(WEIGHTS_ROOT / "preview") + config.engine_config.cp_size = WORLD_SIZE + config.engine_config.ep_size = WORLD_SIZE + model = load_magi2_model(config) + proxy = Magi2DataProxy(config.evaluation_config.data_proxy_config) + return model, proxy, ModelInput + + from fastvideo.configs.models.dits.magi2 import Magi2PreviewVideoConfig + from fastvideo.models.dits.magi2_loader import load_magi2_preview_model + from fastvideo.pipelines.basic.magi2.stages.preview_data_proxy import ( + Magi2DataProxy, + Magi2DataProxyConfig, + ModelInput, + ) + + with (OFFICIAL_ROOT / "configs" / "magi2_preview.json").open( + encoding="utf-8" + ) as config_file: + proxy_values = json.load(config_file)["evaluation_config"]["data_proxy_config"] + model = load_magi2_preview_model( + checkpoint_dir=str(WEIGHTS_ROOT / "preview"), + config=Magi2PreviewVideoConfig(), + device=device, + ) + proxy = Magi2DataProxy(Magi2DataProxyConfig(**proxy_values)) + return model, proxy, ModelInput + + +def _pattern( + shape: tuple[int, ...], + *, + offset: int, + scale: float, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Create a stable, nonuniform tensor without consuming random numbers.""" + element_count = math.prod(shape) + values = torch.arange(element_count, dtype=torch.float32) + values = ((values + offset).remainder(127) - 63) * scale + return values.reshape(shape).to(dtype=dtype) + + +def _build_model_input( + model_input_type: type, + mode: str, + device: torch.device, +) -> Any: + """Build a compact CFG input that exercises every preview modality.""" + video = _pattern((1, 48, 2, 2, 2), offset=3, scale=1 / 64).repeat(2, 1, 1, 1, 1) + audio = _pattern((1, 4, 64), offset=11, scale=1 / 32).repeat(2, 1, 1) + text = _pattern( + (2, 3, 5120), + offset=29, + scale=1 / 128, + dtype=torch.bfloat16, + ) + normalized_time = torch.tensor([0.75, 0.75], dtype=torch.float32) + input_values: dict[str, Any] = { + "x_t": video.to(device), + "audio_x_t": audio.to(device), + "audio_feat_len": torch.tensor([4, 4], dtype=torch.int32, device=device), + "txt_feat": text.to(device), + "txt_feat_len": torch.tensor([3, 2], dtype=torch.int32, device=device), + "t": normalized_time.to(device), + "per_token_video_t": normalized_time.view(2, 1, 1, 1, 1) + .expand(2, 1, 2, 2, 2) + .clone() + .to(device), + "per_token_audio_t": normalized_time.view(2, 1, 1) + .expand(2, 4, 1) + .clone() + .to(device), + } + if mode == "i2v": + conditioned_image = _pattern( + (1, 1, 48, 1, 2, 2), + offset=47, + scale=1 / 64, + ) + conditioned_token = _pattern( + (1, 1, 5120), + offset=71, + scale=1 / 128, + dtype=torch.bfloat16, + ) + input_values.update( + { + "ref_image_feat": torch.cat( + [conditioned_image, torch.zeros_like(conditioned_image)], + dim=0, + ).to(device), + "ref_image_feat_len": torch.tensor( + [[[2, 2]], [[2, 2]]], + dtype=torch.int32, + device=device, + ), + "ref_image_special_token_embedding": torch.cat( + [conditioned_token, torch.zeros_like(conditioned_token)], + dim=0, + ).to(device), + } + ) + elif mode != "t2v": + raise ValueError(f"Unknown preview parity mode: {mode}") + return model_input_type(**input_values) + + +def _tensor_record(tensor: torch.Tensor) -> dict[str, Any]: + """Copy tensor values to CPU and retain the source tensor metadata.""" + return { + "shape": tuple(tensor.shape), + "dtype": str(tensor.dtype), + "stride": tuple(tensor.stride()), + "layout": str(tensor.layout), + "device_type": tensor.device.type, + "device_index": tensor.device.index, + "requires_grad": tensor.requires_grad, + "is_contiguous": tensor.is_contiguous(), + "value": tensor.detach().to(device="cpu").contiguous(), + } + + +def _packed_input_record(packed_input: tuple[Any, ...]) -> dict[str, Any]: + """Record every tensor and sequence boundary passed into the transformer.""" + tokens, coords, modalities, varlen_handler, time_tokens = packed_input + return { + "tokens": _tensor_record(tokens), + "coords": _tensor_record(coords), + "modalities": _tensor_record(modalities), + "time_tokens": _tensor_record(time_tokens), + "cu_seqlens_q": _tensor_record(varlen_handler.cu_seqlens_q), + "cu_seqlens_k": _tensor_record(varlen_handler.cu_seqlens_k), + "max_seqlen_q": varlen_handler.max_seqlen_q, + "max_seqlen_k": varlen_handler.max_seqlen_k, + } + + +def _capture_model_boundaries( + model: torch.nn.Module, +) -> tuple[dict[str, Any], list[torch.utils.hooks.RemovableHandle]]: + """Attach eager hooks at the adapters and every transformer layer boundary.""" + capture: dict[str, Any] = { + "pre_adapter": {}, + "layer_boundaries": [ + {"layer_index": layer_index} for layer_index in range(len(model.block.layers)) + ], + } + handles: list[torch.utils.hooks.RemovableHandle] = [] + + def capture_pre_adapter( + _module: torch.nn.Module, + _inputs: tuple[Any, ...], + output: tuple[torch.Tensor, torch.Tensor], + ) -> None: + """Record projected multimodal tokens and rotary embeddings.""" + hidden_states, rope = output + capture["pre_adapter"] = { + "hidden_states": _tensor_record(hidden_states), + "rope": _tensor_record(rope), + } + + def make_layer_input_hook(layer_index: int): + """Create a hook that records one layer's input hidden states.""" + + def capture_layer_input( + _module: torch.nn.Module, + inputs: tuple[Any, ...], + ) -> None: + """Record the hidden states entering one transformer layer.""" + capture["layer_boundaries"][layer_index]["input"] = _tensor_record( + inputs[0] + ) + + return capture_layer_input + + def make_layer_output_hook(layer_index: int): + """Create a hook that records one layer's output hidden states.""" + + def capture_layer_output( + _module: torch.nn.Module, + _inputs: tuple[Any, ...], + output: torch.Tensor, + ) -> None: + """Record the hidden states leaving one transformer layer.""" + capture["layer_boundaries"][layer_index]["output"] = _tensor_record(output) + + return capture_layer_output + + def capture_post_adapter( + _module: torch.nn.Module, + _inputs: tuple[Any, ...], + output: torch.Tensor, + ) -> None: + """Record the local video-and-audio projection before CP gathering.""" + capture["post_adapter"] = _tensor_record(output) + + handles.append(model.pre_adapter.register_forward_hook(capture_pre_adapter)) + for layer_index, layer in enumerate(model.block.layers): + handles.append( + layer.register_forward_pre_hook(make_layer_input_hook(layer_index)) + ) + handles.append( + layer.register_forward_hook(make_layer_output_hook(layer_index)) + ) + handles.append(model.post_adapter.register_forward_hook(capture_post_adapter)) + return capture, handles + + +def _run_case( + model: torch.nn.Module, + proxy: Any, + model_input_type: type, + mode: str, + device: torch.device, +) -> dict[str, Any]: + """Run one T2V or I2V forward pass and capture all parity boundaries.""" + model_input = _build_model_input(model_input_type, mode, device) + packed_input = proxy.process_input(model_input) + capture, handles = _capture_model_boundaries(model) + capture["packed_input"] = _packed_input_record(packed_input) + try: + with torch.inference_mode(): + model_output = model(*packed_input) + finally: + for handle in handles: + handle.remove() + video_output, audio_output = proxy.process_output(model_output) + capture["model_output"] = _tensor_record(model_output) + capture["depacked_output"] = { + "video": _tensor_record(video_output), + "audio": _tensor_record(audio_output), + } + if any( + set(layer_capture) != {"layer_index", "input", "output"} + for layer_capture in capture["layer_boundaries"] + ): + raise RuntimeError("A preview transformer layer hook did not execute") + return capture + + +def main() -> None: + """Load one implementation, run both modalities, and save rank-local captures.""" + args = _parse_args() + os.environ.setdefault("MAGI_COMPILE_COMPILE_MODE", "NONE") + os.environ.pop("SKIP_LOAD_MODEL", None) + os.environ["MAGI2_CKPT_ROOT"] = str(WEIGHTS_ROOT) + _enable_determinism() + rank, device = _initialize_distributed(args.implementation) + model, proxy, model_input_type = _load_runtime(args.implementation, device) + artifact = { + "schema_version": 1, + "implementation": args.implementation, + "rank": rank, + "world_size": dist.get_world_size(), + "cases": { + mode: _run_case(model, proxy, model_input_type, mode, device) + for mode in ("t2v", "i2v") + }, + } + args.output_dir.mkdir(parents=True, exist_ok=True) + artifact_path = args.output_dir / f"rank_{rank}.pt" + temporary_path = args.output_dir / f"rank_{rank}.pt.tmp" + torch.save(artifact, temporary_path) + os.replace(temporary_path, artifact_path) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/magi2/_refiner_transformer_parity_worker.py b/tests/local_tests/magi2/_refiner_transformer_parity_worker.py new file mode 100644 index 000000000..23e870493 --- /dev/null +++ b/tests/local_tests/magi2/_refiner_transformer_parity_worker.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run one isolated side of MAGI-2 refiner transformer numerical parity.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +from pathlib import Path +import random +import sys +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_ROOT = Path( + os.environ.get("MAGI2_OFFICIAL_REF_DIR", REPO_ROOT.parent / "MAGI-2-preview") +) +WEIGHTS_ROOT = Path( + os.environ.get("MAGI2_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "magi2") +) +WORLD_SIZE = 8 +SEED = 42 + + +def _parse_args() -> argparse.Namespace: + """Parse the implementation and artifact directory for one torchrun job.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--implementation", + choices=("official", "fastvideo"), + required=True, + ) + parser.add_argument("--output-dir", type=Path, required=True) + return parser.parse_args() + + +def _enable_determinism() -> None: + """Apply the deterministic controls used by the official entry point.""" + os.environ["MAGI2_DETERMINISTIC"] = "1" + os.environ["MAGI_ATTENTION_DETERMINISTIC_MODE"] = "1" + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + random.seed(SEED) + np.random.seed(SEED) + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + torch.use_deterministic_algorithms(True) + + +def _initialize_distributed(implementation: str) -> tuple[int, torch.device]: + """Initialize the official eight-rank CP and EP process-group topology.""" + if dist.is_initialized(): + raise RuntimeError("The parity worker requires a fresh distributed process") + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + if dist.get_world_size() != WORLD_SIZE: + raise RuntimeError( + f"MAGI-2 refiner parity requires {WORLD_SIZE} ranks, " + f"received {dist.get_world_size()}" + ) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + + if implementation == "official": + sys.path.insert(0, str(OFFICIAL_ROOT)) + from inference.infra.distributed import ( + initialize_expert_parallel, + initialize_model_parallel, + ) + else: + sys.path.insert(0, str(REPO_ROOT)) + from fastvideo.models.dits.magi2_runtime import ( + initialize_expert_parallel, + initialize_model_parallel, + ) + + initialize_model_parallel(cp_size=WORLD_SIZE) + initialize_expert_parallel(ep_size=WORLD_SIZE) + return rank, torch.device("cuda", local_rank) + + +def _load_runtime(implementation: str, device: torch.device) -> tuple[Any, Any, type]: + """Load one implementation's production transformer, proxy, and input type.""" + if implementation == "official": + from inference.common.magi2_config import load_config + from inference.infra.checkpoint.load_checkpoint import load_magi2_refiner + from inference.pipeline.inference_engine import EvalInput + from inference.pipeline.refiner_data_proxy import Magi2RefinerDataProxy + + config = load_config(str(OFFICIAL_ROOT / "configs" / "magi2_refiner.json")) + config.engine_config.cp_size = WORLD_SIZE + config.engine_config.ep_size = WORLD_SIZE + config.evaluation_config.magi2_refiner_model_path = str(WEIGHTS_ROOT / "refiner") + model = load_magi2_refiner(config).to(device=device).eval() + proxy = Magi2RefinerDataProxy( + config.evaluation_config.magi2_refiner_data_proxy_config + ) + return model, proxy, EvalInput + + from fastvideo.configs.models.dits.magi2 import Magi2RefinerVideoConfig + from fastvideo.models.dits.magi2_loader import load_magi2_refiner_model + from fastvideo.pipelines.basic.magi2.stages.refiner_data_proxy import ( + Magi2RefinerDataProxy, + Magi2RefinerDataProxyConfig, + RefinerModelInput, + ) + + with (OFFICIAL_ROOT / "configs" / "magi2_refiner.json").open( + encoding="utf-8" + ) as config_file: + proxy_values = json.load(config_file)["evaluation_config"][ + "magi2_refiner_data_proxy_config" + ] + model = load_magi2_refiner_model( + checkpoint_dir=str(WEIGHTS_ROOT / "refiner"), + config=Magi2RefinerVideoConfig(), + device=device, + ) + proxy = Magi2RefinerDataProxy(Magi2RefinerDataProxyConfig(**proxy_values)) + return model, proxy, RefinerModelInput + + +def _pattern( + shape: tuple[int, ...], + *, + offset: int, + scale: float, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Create a stable, nonuniform tensor without consuming random numbers.""" + element_count = math.prod(shape) + values = torch.arange(element_count, dtype=torch.float32) + values = ((values + offset).remainder(127) - 63) * scale + return values.reshape(shape).to(dtype=dtype) + + +def _build_model_input(model_input_type: type, device: torch.device) -> Any: + """Build one packed input that exercises every refiner modality and CP split.""" + input_values = { + "x_t": _pattern((1, 48, 9, 4, 24), offset=3, scale=1 / 64).to(device), + "audio_x_t": _pattern((1, 3, 64), offset=11, scale=1 / 32).to(device), + "audio_feat_len": torch.tensor([3], dtype=torch.int64), + "txt_feat": _pattern((1, 5, 5120), offset=29, scale=1 / 128).to(device), + "txt_feat_len": torch.tensor([5], dtype=torch.int64), + "ref_audio_feat": _pattern((1, 2, 64), offset=47, scale=1 / 32).to(device), + "ref_audio_feat_len": torch.tensor([2], dtype=torch.int64), + "ref_video_feat": _pattern((1, 48, 9, 4, 24), offset=71, scale=1 / 64).to(device), + "ref_video_feat_len": torch.tensor([4], dtype=torch.int64), + } + return model_input_type(**input_values) + + +def _tensor_record(tensor: torch.Tensor) -> dict[str, Any]: + """Copy tensor values to CPU and retain the source tensor metadata.""" + return { + "shape": tuple(tensor.shape), + "dtype": str(tensor.dtype), + "stride": tuple(tensor.stride()), + "layout": str(tensor.layout), + "device_type": tensor.device.type, + "device_index": tensor.device.index, + "requires_grad": tensor.requires_grad, + "is_contiguous": tensor.is_contiguous(), + "value": tensor.detach().to(device="cpu").contiguous(), + } + + +def _scalar_or_tensor_record(value: int | float | torch.Tensor | None) -> Any: + """Record scalar attention metadata while preserving tensor metadata.""" + if isinstance(value, torch.Tensor): + return _tensor_record(value) + return value + + +def _local_attention_record(handler: Any) -> dict[str, Any] | None: + """Record every field that controls the refiner's local attention kernel.""" + if handler is None: + return None + field_names = ( + "q_ranges", + "k_ranges", + "max_seqlen_q", + "max_seqlen_k", + "attn_type_map", + "softmax_scale", + "bwd_q_ranges", + "bwd_k_ranges", + "bwd_attn_type_map", + "auto_range_merge", + "sparse_load", + ) + return { + field_name: _scalar_or_tensor_record(getattr(handler, field_name, None)) + for field_name in field_names + } + + +def _packed_input_record(packed_input: tuple[Any, ...]) -> dict[str, Any]: + """Record every tensor and sequence boundary passed into the transformer.""" + tokens, coords, modalities, varlen_handler, local_attn_handler = packed_input + return { + "tokens": _tensor_record(tokens), + "coords": _tensor_record(coords), + "modalities": _tensor_record(modalities), + "cu_seqlens_q": _tensor_record(varlen_handler.cu_seqlens_q), + "cu_seqlens_k": _tensor_record(varlen_handler.cu_seqlens_k), + "max_seqlen_q": _scalar_or_tensor_record(varlen_handler.max_seqlen_q), + "max_seqlen_k": _scalar_or_tensor_record(varlen_handler.max_seqlen_k), + "local_attention": _local_attention_record(local_attn_handler), + } + + +def _capture_model_boundaries( + model: torch.nn.Module, +) -> tuple[dict[str, Any], list[torch.utils.hooks.RemovableHandle]]: + """Attach eager hooks at the adapters and all 30 transformer layers.""" + capture: dict[str, Any] = { + "pre_adapter": {}, + "layer_boundaries": [ + {"layer_index": layer_index} for layer_index in range(len(model.block.layers)) + ], + "post_adapter": {}, + } + if len(capture["layer_boundaries"]) != 30: + raise RuntimeError( + "MAGI-2 refiner parity requires 30 layer boundaries, " + f"received {len(capture['layer_boundaries'])}" + ) + handles: list[torch.utils.hooks.RemovableHandle] = [] + + def capture_pre_adapter( + _module: torch.nn.Module, + inputs: tuple[Any, ...], + output: tuple[torch.Tensor, torch.Tensor], + ) -> None: + """Record CP-local adapter inputs, hidden states, and rotary embeddings.""" + tokens, coords, video_mask, audio_mask, text_mask = inputs + hidden_states, rope = output + capture["pre_adapter"] = { + "tokens": _tensor_record(tokens), + "coords": _tensor_record(coords), + "video_mask": _tensor_record(video_mask), + "audio_mask": _tensor_record(audio_mask), + "text_mask": _tensor_record(text_mask), + "hidden_states": _tensor_record(hidden_states), + "rope": _tensor_record(rope), + } + + def make_layer_input_hook(layer_index: int): + """Create a hook that records one layer's input hidden states.""" + + def capture_layer_input( + _module: torch.nn.Module, + inputs: tuple[Any, ...], + ) -> None: + """Record the hidden states entering one transformer layer.""" + capture["layer_boundaries"][layer_index]["input"] = _tensor_record( + inputs[0] + ) + + return capture_layer_input + + def make_layer_output_hook(layer_index: int): + """Create a hook that records one layer's output hidden states.""" + + def capture_layer_output( + _module: torch.nn.Module, + _inputs: tuple[Any, ...], + output: torch.Tensor, + ) -> None: + """Record the hidden states leaving one transformer layer.""" + capture["layer_boundaries"][layer_index]["output"] = _tensor_record(output) + + return capture_layer_output + + def capture_post_adapter( + _module: torch.nn.Module, + inputs: tuple[Any, ...], + output: torch.Tensor, + ) -> None: + """Record the CP-local final projection inputs, masks, and output.""" + hidden_states, video_mask, audio_mask = inputs + capture["post_adapter"] = { + "hidden_states": _tensor_record(hidden_states), + "video_mask": _tensor_record(video_mask), + "audio_mask": _tensor_record(audio_mask), + "output": _tensor_record(output), + } + + handles.append(model.pre_adapter.register_forward_hook(capture_pre_adapter)) + for layer_index, layer in enumerate(model.block.layers): + handles.append( + layer.register_forward_pre_hook(make_layer_input_hook(layer_index)) + ) + handles.append(layer.register_forward_hook(make_layer_output_hook(layer_index))) + handles.append(model.post_adapter.register_forward_hook(capture_post_adapter)) + return capture, handles + + +def _run_case( + model: torch.nn.Module, + proxy: Any, + model_input_type: type, + device: torch.device, +) -> dict[str, Any]: + """Run one refiner forward pass and capture all parity boundaries.""" + model_input = _build_model_input(model_input_type, device) + packed_input = proxy.process_input(model_input) + capture, handles = _capture_model_boundaries(model) + capture["packed_input"] = _packed_input_record(packed_input) + try: + with torch.inference_mode(): + model_output = model(*packed_input) + finally: + for handle in handles: + handle.remove() + video_output, audio_output = proxy.process_output(model_output) + capture["model_output"] = _tensor_record(model_output) + capture["depacked_output"] = { + "video": _tensor_record(video_output), + "audio": _tensor_record(audio_output), + } + if set(capture["pre_adapter"]) != { + "tokens", + "coords", + "video_mask", + "audio_mask", + "text_mask", + "hidden_states", + "rope", + }: + raise RuntimeError("The refiner pre-adapter hook did not execute") + if any( + set(layer_capture) != {"layer_index", "input", "output"} + for layer_capture in capture["layer_boundaries"] + ): + raise RuntimeError("A refiner transformer layer hook did not execute") + if set(capture["post_adapter"]) != { + "hidden_states", + "video_mask", + "audio_mask", + "output", + }: + raise RuntimeError("The refiner post-adapter hook did not execute") + if set(capture) != { + "packed_input", + "pre_adapter", + "layer_boundaries", + "post_adapter", + "model_output", + "depacked_output", + }: + raise RuntimeError("The refiner parity capture schema is incomplete") + return capture + + +def main() -> None: + """Load one implementation, run the refiner, and save rank-local captures.""" + args = _parse_args() + os.environ.setdefault("MAGI_COMPILE_COMPILE_MODE", "NONE") + os.environ.pop("SKIP_LOAD_MODEL", None) + os.environ["MAGI2_CKPT_ROOT"] = str(WEIGHTS_ROOT) + _enable_determinism() + rank, device = _initialize_distributed(args.implementation) + model, proxy, model_input_type = _load_runtime(args.implementation, device) + artifact = { + "schema_version": 1, + "implementation": args.implementation, + "rank": rank, + "world_size": dist.get_world_size(), + "case": _run_case(model, proxy, model_input_type, device), + } + args.output_dir.mkdir(parents=True, exist_ok=True) + artifact_path = args.output_dir / f"rank_{rank}.pt" + temporary_path = args.output_dir / f"rank_{rank}.pt.tmp" + torch.save(artifact, temporary_path) + os.replace(temporary_path, artifact_path) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/magi2/test_magi2_audio_vae_parity.py b/tests/local_tests/magi2/test_magi2_audio_vae_parity.py new file mode 100644 index 000000000..50aa0fe36 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_audio_vae_parity.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict Stable Audio VAE decode parity for MAGI-2 Preview. + +Coverage scope: implementation_subcomponent. The official side instantiates +``inference.pipeline.audio_decoder.SAAudioFeatureExtractor``, which builds the +official ``AudioAutoencoder`` from ``inference.model.sa_audio_vae``. The +FastVideo side uses ``fastvideo.models.vaes.magi2_audio_vae.Magi2AudioVAE``, +which wraps ``fastvideo.models.vaes.oobleck.OobleckDecoder``. The test also +checks MAGI-2's latent transpose, sample-major stereo layout, and 441/512 +resampling contract. +""" +from __future__ import annotations + +import pytest +import torch + +from fastvideo.models.vaes.magi2_audio_vae import load_magi2_audio_vae +from fastvideo.pipelines.basic.magi2.stages.audio_decoding import ( + decode_magi2_audio, +) +from tests.local_tests.magi2._parity_utils import ( + LOCAL_WEIGHTS_DIR, + assert_array_exact, + assert_tensor_exact, + import_official_module, + require_path, +) + + +PARITY_COVERAGE = "both" +STABLE_AUDIO_DIR = LOCAL_WEIGHTS_DIR / "stable-audio-open-1.0" +STABLE_AUDIO_CONFIG_PATH = STABLE_AUDIO_DIR / "model_config.json" +STABLE_AUDIO_CHECKPOINT_PATH = STABLE_AUDIO_DIR / "model.safetensors" + + +def _deterministic_audio_latent(device: torch.device) -> torch.Tensor: + """Create a short sample-major latent with MAGI-2's 64 audio channels.""" + latent_values = torch.arange(8 * 64, dtype=torch.float32) + latent_values = ((latent_values % 127) - 63) / 64 + return latent_values.reshape(8, 64).to(device=device) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MAGI-2 Stable Audio VAE parity requires CUDA.", +) +def test_magi2_stable_audio_decode_and_resample_exact_parity() -> None: + """Require exact decoded waveform and sample-major resampled audio.""" + require_path(STABLE_AUDIO_CONFIG_PATH, "Stable Audio configuration") + require_path(STABLE_AUDIO_CHECKPOINT_PATH, "Stable Audio checkpoint") + official_audio_module = import_official_module( + "inference.pipeline.audio_decoder" + ) + device = torch.device("cuda:0") + torch.cuda.set_device(device) + official_audio_vae = official_audio_module.SAAudioFeatureExtractor( + str(STABLE_AUDIO_DIR) + ) + fastvideo_audio_vae = load_magi2_audio_vae(STABLE_AUDIO_DIR, device) + assert official_audio_vae.sample_rate == fastvideo_audio_vae.sampling_rate == 44100 + assert official_audio_vae.downsampling_ratio == fastvideo_audio_vae.hop_length == 2048 + + sample_major_latent = _deterministic_audio_latent(device) + official_channel_major_latent = sample_major_latent.T + fastvideo_channel_major_latent = sample_major_latent.unsqueeze(0).permute(0, 2, 1).contiguous() + with torch.inference_mode(): + official_waveform = official_audio_vae.decode(official_channel_major_latent).detach().cpu() + fastvideo_waveform = fastvideo_audio_vae.decode( + fastvideo_channel_major_latent + ).detach().cpu() + + assert_tensor_exact(fastvideo_waveform, official_waveform, "decoded audio waveform") + official_sample_major = official_waveform.squeeze(0).T.numpy() + fastvideo_sample_major = fastvideo_waveform.squeeze(0).T.numpy() + assert_array_exact(fastvideo_sample_major, official_sample_major, "sample-major stereo audio") + + official_resampled = official_audio_module.resample_audio_sinc( + official_sample_major, + 441 / 512, + ) + fastvideo_resampled = decode_magi2_audio( + fastvideo_audio_vae, + sample_major_latent, + ) + assert_array_exact(fastvideo_resampled, official_resampled, "resampled stereo audio") diff --git a/tests/local_tests/magi2/test_magi2_checkpoint_conversion.py b/tests/local_tests/magi2/test_magi2_checkpoint_conversion.py new file mode 100644 index 000000000..04ab84efc --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_checkpoint_conversion.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Checkpoint-layout conversion coverage for the MAGI-2 port.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from fastvideo.utils import verify_model_config_and_directory +from scripts.checkpoint_conversion.convert_magi2_to_fastvideo import ( + COMPONENT_DIRECTORY_MAPPING, + MODEL_INDEX, + REQUIRED_COMPONENT_FILES, + SOURCE_REVISION, + TURBO_VAE_SKIPPED_KEYS, + convert_checkpoint_layout, +) + + +def _create_minimal_source(source: Path) -> dict[str, Path]: + """Create identifying files for every published checkpoint component.""" + source_files: dict[str, Path] = {} + for component, relative_paths in REQUIRED_COMPONENT_FILES.items(): + for relative_path in relative_paths: + source_path = source / component / relative_path + source_path.parent.mkdir(parents=True, exist_ok=True) + if relative_path.endswith(".index.json"): + shard_name = f"{component}-00001-of-00001.safetensors" + source_path.write_text( + json.dumps({"weight_map": {"weight": shard_name}}), + encoding="utf-8", + ) + shard_path = source_path.parent / shard_name + shard_path.write_bytes(shard_name.encode()) + source_files[f"{component}/{shard_name}"] = shard_path + else: + source_path.write_bytes(f"{component}/{relative_path}".encode()) + source_files[f"{component}/{relative_path}"] = source_path + return source_files + + +def test_magi2_converter_preserves_files_and_records_all_mappings(tmp_path: Path) -> None: + """Require hard-linked components and complete conversion metadata.""" + source = tmp_path / "official" + output = tmp_path / "fastvideo" + source_files = _create_minimal_source(source) + + convert_checkpoint_layout(source, output) + + assert verify_model_config_and_directory(str(output)) == MODEL_INDEX + for source_relative_path, source_path in source_files.items(): + source_component, component_relative_path = source_relative_path.split("/", 1) + destination_component = COMPONENT_DIRECTORY_MAPPING[source_component] + destination_path = output / destination_component / component_relative_path + assert destination_path.read_bytes() == source_path.read_bytes() + assert destination_path.stat().st_ino == source_path.stat().st_ino + + manifest = json.loads((output / "magi2_conversion_manifest.json").read_text()) + assert manifest["source"]["revision"] == SOURCE_REVISION + assert manifest["component_directory_mapping"] == COMPONENT_DIRECTORY_MAPPING + assert manifest["skipped_checkpoint_keys"]["vae/checkpoint.ckpt"] == list( + TURBO_VAE_SKIPPED_KEYS + ) + assert len(TURBO_VAE_SKIPPED_KEYS) == 8 + + +def test_magi2_converter_rejects_an_incomplete_snapshot(tmp_path: Path) -> None: + """Reject a source tree before creating a partial output repository.""" + source = tmp_path / "official" + source.mkdir() + output = tmp_path / "fastvideo" + + try: + convert_checkpoint_layout(source, output) + except FileNotFoundError as error: + assert "source snapshot is incomplete" in str(error) + else: + raise AssertionError("An incomplete MAGI-2 snapshot was accepted") + assert not output.exists() diff --git a/tests/local_tests/magi2/test_magi2_image_encoder_parity.py b/tests/local_tests/magi2/test_magi2_image_encoder_parity.py new file mode 100644 index 000000000..72298cd20 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_image_encoder_parity.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict production-loader parity for the MAGI-2 Wan 2.2 image encoder. + +Coverage scope: both. The test loads the published ``Wan2_2_VAE`` reference +and the FastVideo ``Magi2WanImageEncoder`` through its production loader. A +synthetic image passes through the I2V letterbox, ``VideoProcessor``, and BF16 +round trip before FP32 encoding. +""" +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from fastvideo.configs.models.vaes.magi2_wanvae import Magi2WanVAEConfig +from fastvideo.models.vaes.magi2_wan_loader import ( + Magi2WanImageEncoder, + load_magi2_wan_image_encoder, +) +from tests.local_tests.magi2._parity_utils import ( + LOCAL_WEIGHTS_DIR, + assert_tensor_exact, + import_official_module, + require_path, +) + + +PARITY_COVERAGE = "both" +WAN_VAE_PATH = LOCAL_WEIGHTS_DIR / "vae" / "Wan2.2_VAE.pth" + + +def _resizepad(image, target_height: int, target_width: int): + """Fit an RGB image inside a white canvas without cropping its edges.""" + image_module = pytest.importorskip("PIL.Image") + width, height = image.size + if width <= 0 or height <= 0: + raise ValueError(f"Invalid image size, width: {width}, height: {height}") + scale = min(target_width / width, target_height / height) + resized_width = max(1, int(round(width * scale))) + resized_height = max(1, int(round(height * scale))) + resized_image = image.convert("RGB").resize( + (resized_width, resized_height), + resample=image_module.Resampling.LANCZOS, + ) + canvas = image_module.new("RGB", (target_width, target_height), (255, 255, 255)) + canvas.paste( + resized_image, + ((target_width - resized_width) // 2, (target_height - resized_height) // 2), + ) + return canvas + + +def _prepare_reference_image(device: torch.device) -> torch.Tensor: + """Apply the I2V resize, letterbox, and BF16 input conversion. + + The source dimensions force the long-edge scaling branch and a non-square + letterbox. The returned tensor has the ``[B, C, T, H, W]`` layout consumed + by ``Wan2_2_VAE.encode``. + """ + pil_image_module = pytest.importorskip("PIL.Image") + video_processor_module = pytest.importorskip("diffusers.video_processor") + pixels = ( + np.arange(48 * 80 * 3, dtype=np.uint32).reshape(48, 80, 3) % 251 + ).astype(np.uint8) + source_image = pil_image_module.fromarray(pixels, mode="RGB") + generation_height, generation_width = 64, 96 + max_length = max(generation_height, generation_width) + target_width = max_length + target_height = int(source_image.height * max_length / source_image.width) + resized_image = _resizepad( + source_image, + target_height, + target_width, + ) + video_processor = video_processor_module.VideoProcessor(vae_scale_factor=32) + image_tensor = video_processor.preprocess( + resized_image, + height=target_height, + width=target_width, + ) + return image_tensor.to(device=device, dtype=torch.bfloat16).unsqueeze(2)[:, :3].float() + + +def _assert_magi2_wan_config( + official_vae, + fastvideo_encoder: Magi2WanImageEncoder, +) -> None: + """Validate the production encoder configuration against published values.""" + config = Magi2WanVAEConfig() + assert config.base_dim == 160 + assert config.decoder_base_dim == 256 + assert config.z_dim == 48 + assert config.in_channels == 12 + assert config.out_channels == 12 + assert config.temperal_downsample == (False, True, True) + assert config.patch_size == 2 + assert config.is_residual is True + assert config.clip_output is False + assert torch.equal( + torch.tensor(config.latents_mean, dtype=torch.float32), + official_vae.mean.detach().cpu(), + ) + assert torch.equal( + torch.tensor(config.latents_std, dtype=torch.float32), + official_vae.std.detach().cpu(), + ) + assert torch.equal(fastvideo_encoder.mean.detach().cpu(), official_vae.mean.detach().cpu()) + assert torch.equal(fastvideo_encoder.std.detach().cpu(), official_vae.std.detach().cpu()) + assert torch.equal( + fastvideo_encoder.inverse_std.detach().cpu(), + official_vae.scale[1].detach().cpu(), + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MAGI-2 I2V image-encoder parity requires CUDA.", +) +def test_load_magi2_wan_image_encoder_i2v_exact_parity() -> None: + """Require exact post-normalization image latents from both encoders.""" + vae_path = require_path(WAN_VAE_PATH, "Wan 2.2 VAE weights") + official_vae_module = import_official_module("inference.model.vae2_2") + device = torch.device("cuda:0") + image_tensor = _prepare_reference_image(device) + + official_vae = official_vae_module.get_vae2_2( + str(vae_path), + device=str(device), + weight_dtype=torch.float32, + ) + fastvideo_encoder = load_magi2_wan_image_encoder(vae_path, device) + _assert_magi2_wan_config(official_vae, fastvideo_encoder) + + with torch.inference_mode(): + official_latent = official_vae.encode(image_tensor).detach().cpu() + fastvideo_latent = fastvideo_encoder.encode(image_tensor).detach().cpu() + + assert image_tensor.dtype == torch.float32 + assert image_tensor.shape[:3] == (1, 3, 1) + assert_tensor_exact(fastvideo_latent, official_latent, "I2V Wan latent") diff --git a/tests/local_tests/magi2/test_magi2_parallel_binding.py b/tests/local_tests/magi2/test_magi2_parallel_binding.py new file mode 100644 index 000000000..0fde5ab48 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_parallel_binding.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo process-group binding coverage for MAGI-2.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from fastvideo.models.dits.magi2_runtime import fastvideo_parallel +from fastvideo.models.dits.magi2_runtime import psm as psm_manager + + +def test_fastvideo_groups_are_reused_for_context_and_expert_parallelism(monkeypatch) -> None: + """Bind one eight-rank sequence group to both MAGI-2 parallel axes.""" + sequence_process_group = object() + data_process_group = object() + sequence_group = SimpleNamespace( + world_size=8, + device_group=sequence_process_group, + ranks=list(range(8)), + ) + data_group = SimpleNamespace( + world_size=1, + device_group=data_process_group, + ranks=[3], + ) + monkeypatch.setattr(fastvideo_parallel, "get_sp_group", lambda: sequence_group) + monkeypatch.setattr(fastvideo_parallel, "get_dp_group", lambda: data_group) + monkeypatch.setattr( + psm_manager, + "get_world_size", + lambda dim="": 8 if dim in {"cp", "ep"} else 1, + ) + + recorded: dict = {} + monkeypatch.setattr( + fastvideo_parallel, + "bind_process_groups", + lambda **kwargs: recorded.update(kwargs), + ) + fastvideo_parallel.bind_fastvideo_parallel_state() + + assert recorded["cp_group"] is sequence_process_group + assert recorded["ep_group"] is sequence_process_group + assert recorded["dp_group"] is data_process_group + assert recorded["cp_ranks"] == list(range(8)) + + +def test_fastvideo_binding_requires_the_published_parallel_size(monkeypatch) -> None: + """Reject sequence groups that cannot match the published eight-rank layout.""" + monkeypatch.setattr( + fastvideo_parallel, + "get_sp_group", + lambda: SimpleNamespace(world_size=4), + ) + monkeypatch.setattr( + fastvideo_parallel, + "get_dp_group", + lambda: SimpleNamespace(world_size=2), + ) + with pytest.raises(ValueError, match="sp_size=8"): + fastvideo_parallel.bind_fastvideo_parallel_state() diff --git a/tests/local_tests/magi2/test_magi2_pipeline_parity.py b/tests/local_tests/magi2/test_magi2_pipeline_parity.py new file mode 100644 index 000000000..8ec4ebe1d --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_pipeline_parity.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict stage and end-to-end parity for the MAGI-2 release pipeline.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import sys +from typing import Any + +import pytest +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_ROOT = Path( + os.environ.get("MAGI2_OFFICIAL_REF_DIR", REPO_ROOT.parent / "MAGI-2-preview") +) +WEIGHTS_ROOT = Path( + os.environ.get("MAGI2_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "magi2") +) +CONVERTED_ROOT = Path( + os.environ.get("MAGI2_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / "magi2") +) +WORKER_PATH = Path(__file__).with_name("_pipeline_parity_worker.py") +CAPTURE_ROOT = REPO_ROOT / "archived" / "magi2_parity" / "validation" / "pipeline" +OFFICIAL_REVISION = "073c84f2102ec3c9287623113a103c14402770ad" +WORLD_SIZE = 8 + + +pytestmark = pytest.mark.skipif( + torch.cuda.device_count() < WORLD_SIZE, + reason=f"MAGI-2 pipeline numerical parity requires {WORLD_SIZE} CUDA devices", +) + + +def _require_parity_sources() -> None: + """Require clean pinned source trees and complete converted components.""" + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=OFFICIAL_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + assert revision == OFFICIAL_REVISION + official_changes = subprocess.run( + ["git", "status", "--porcelain=v1"], + cwd=OFFICIAL_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout + assert not official_changes, ( + f"Official MAGI-2 checkout must be clean for parity:\n{official_changes}" + ) + required_paths = ( + WEIGHTS_ROOT / "preview" / "model.safetensors.index.json", + WEIGHTS_ROOT / "refiner" / "model.safetensors.index.json", + WEIGHTS_ROOT / "turbo_vae" / "checkpoint.ckpt", + CONVERTED_ROOT / "model_index.json", + CONVERTED_ROOT / "transformer" / "model.safetensors.index.json", + CONVERTED_ROOT / "transformer_2" / "model.safetensors.index.json", + ) + for required_path in required_paths: + assert required_path.is_file(), f"MAGI-2 parity input is missing: {required_path}" + + +def _run_implementation(implementation: str) -> Path: + """Launch one full release-profile implementation in an isolated torchrun job.""" + output_dir = CAPTURE_ROOT / implementation + environment = os.environ.copy() + environment.update( + { + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "MAGI2_CKPT_ROOT": str(WEIGHTS_ROOT), + "MAGI2_CONVERTED_WEIGHTS_DIR": str(CONVERTED_ROOT), + "MAGI2_DETERMINISTIC": "1", + "MAGI2_LOCAL_WEIGHTS_DIR": str(WEIGHTS_ROOT), + "MAGI2_OFFICIAL_REF_DIR": str(OFFICIAL_ROOT), + "MAGI2_TEXT_ENC_OFFLOAD_MODE": "roundtrip", + "MAGI2_VAE_OFFLOAD_MODE": "roundtrip", + "MAGI_ATTENTION_DETERMINISTIC_MODE": "1", + "MAGI_COMPILE_COMPILE_MODE": "NONE", + "OMP_NUM_THREADS": "1", + "PYTHONHASHSEED": "42", + } + ) + environment.pop("MAGI2_SAVE_LATENT_PATH", None) + environment.pop("NEGATIVE_PROMPT", None) + environment.pop("SKIP_LOAD_MODEL", None) + completed = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + f"--nproc-per-node={WORLD_SIZE}", + str(WORKER_PATH), + "--implementation", + implementation, + "--output-dir", + str(output_dir), + ], + cwd=REPO_ROOT, + env=environment, + check=False, + capture_output=True, + text=True, + timeout=21600, + ) + if completed.returncode != 0: + raise AssertionError( + f"{implementation} MAGI-2 pipeline torchrun failed.\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + return output_dir + + +def _load_capture(output_dir: Path) -> dict[str, Any]: + """Load and validate one stage-digest manifest.""" + artifact_path = output_dir / "capture.json" + assert artifact_path.is_file(), f"MAGI-2 capture is missing: {artifact_path}" + artifact = json.loads(artifact_path.read_text(encoding="utf-8")) + assert artifact["schema_version"] == 1 + assert artifact["world_size"] == WORLD_SIZE + assert artifact["preview_steps"] == 100 + assert artifact["refiner_steps"] == 5 + assert artifact["cases"].keys() == {"t2v", "i2v"} + return artifact + + +def _first_difference( + fastvideo_value: Any, + official_value: Any, + path: str, +) -> str | None: + """Return the first structural or numerical difference in two manifests.""" + if type(fastvideo_value) is not type(official_value): + return ( + f"{path}: type differs: {type(fastvideo_value).__name__} versus " + f"{type(official_value).__name__}" + ) + if isinstance(official_value, dict): + if fastvideo_value.keys() != official_value.keys(): + return ( + f"{path}: keys differ: {sorted(fastvideo_value)} versus " + f"{sorted(official_value)}" + ) + for key in official_value: + difference = _first_difference( + fastvideo_value[key], + official_value[key], + f"{path}.{key}", + ) + if difference is not None: + return difference + return None + if isinstance(official_value, list): + if len(fastvideo_value) != len(official_value): + return ( + f"{path}: list length differs: {len(fastvideo_value)} versus " + f"{len(official_value)}" + ) + for index, (fastvideo_item, official_item) in enumerate( + zip(fastvideo_value, official_value, strict=True) + ): + difference = _first_difference( + fastvideo_item, + official_item, + f"{path}[{index}]", + ) + if difference is not None: + return difference + return None + if fastvideo_value != official_value: + return f"{path}: {fastvideo_value!r} differs from {official_value!r}" + return None + + +def test_magi2_pipeline_matches_official_release_exactly() -> None: + """Match T2V and I2V stage tensors, decoded video, and decoded audio exactly.""" + _require_parity_sources() + official_capture = _load_capture(_run_implementation("official")) + fastvideo_capture = _load_capture(_run_implementation("fastvideo")) + assert official_capture["implementation"] == "official" + assert fastvideo_capture["implementation"] == "fastvideo" + difference = _first_difference( + fastvideo_capture["cases"], + official_capture["cases"], + "cases", + ) + assert difference is None, difference diff --git a/tests/local_tests/magi2/test_magi2_preview_data_proxy_parity.py b/tests/local_tests/magi2/test_magi2_preview_data_proxy_parity.py new file mode 100644 index 000000000..e92977cc3 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_preview_data_proxy_parity.py @@ -0,0 +1,531 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict MAGI-2 preview data-proxy implementation parity. + +Coverage scope: implementation_subcomponent. The tests load the pinned official +``Magi2DataProxy`` source and compare it with the FastVideo-owned component on +deterministic text-to-video (T2V) and image-to-video (I2V) tensors. +""" + +from __future__ import annotations + +import ast +from dataclasses import asdict, dataclass +from enum import IntEnum +from functools import lru_cache +import importlib.util +import math +import os +from pathlib import Path +import sys +import types +import typing +from typing import Any + +import pytest +import torch +from torch.testing import assert_close + +import fastvideo.pipelines.basic.magi2.stages.preview_data_proxy as fastvideo_proxy_module +from fastvideo.pipelines.basic.magi2.stages.preview_data_proxy import ( + Magi2DataProxy, + Magi2DataProxyConfig, + ModelInput, + Modality, + VarlenHandler, +) + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_REPO_ROOT = Path( + os.environ.get( + "MAGI2_OFFICIAL_REF_DIR", + REPO_ROOT.parent / "MAGI-2-preview", + ) +) +OFFICIAL_PROXY_PATH = ( + OFFICIAL_REPO_ROOT / "inference" / "pipeline" / "preview_data_proxy.py" +) +OFFICIAL_MODEL_PATH = ( + OFFICIAL_REPO_ROOT / "inference" / "model" / "magi2_preview.py" +) +PARITY_SCOPE = "implementation_subcomponent" + + +@dataclass(frozen=True) +class _ParallelStateStub: + """Expose fixed parallel sizes and identifiable process-group handles.""" + + cp_size: int = 1 + ep_size: int = 1 + + def get_world_size(self, dimension: str = "") -> int: + """Return the configured size for one parallel dimension.""" + if dimension == "cp": + return self.cp_size + if dimension == "ep": + return self.ep_size + return max(self.cp_size, self.ep_size) + + def get_parallel_group(self, dimension: str) -> str: + """Return a stable process-group marker for distributed call assertions.""" + return f"{dimension}-group" + + +def _load_official_model_utility_module() -> types.ModuleType: + """Load official coordinate, modality, and time-embedding definitions.""" + if not OFFICIAL_MODEL_PATH.is_file(): + raise FileNotFoundError( + f"Official MAGI-2 model source is missing: {OFFICIAL_MODEL_PATH}" + ) + source = OFFICIAL_MODEL_PATH.read_text(encoding="utf-8") + source_module = ast.parse(source, filename=str(OFFICIAL_MODEL_PATH)) + symbol_names = { + "Modality", + "VarlenHandler", + "get_coords", + "sinusoidal_embedding_1d", + } + symbol_nodes = [ + node + for node in source_module.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + and node.name in symbol_names + ] + found_names = {node.name for node in symbol_nodes} + if found_names != symbol_names: + raise AssertionError( + f"Official model utility definitions differ: expected {symbol_names}, " + f"found {found_names}" + ) + + module = types.ModuleType("inference.model.magi2_preview") + module.__file__ = str(OFFICIAL_MODEL_PATH) + module.__dict__.update( + { + "Coords": tuple[int, int, int], + "IntEnum": IntEnum, + "Optional": typing.Optional, + "dataclass": dataclass, + "math": math, + "torch": torch, + } + ) + utility_module = ast.Module(body=symbol_nodes, type_ignores=[]) + module_name = module.__name__ + previous_module = sys.modules.get(module_name) + sys.modules[module_name] = module + try: + exec( + compile(utility_module, str(OFFICIAL_MODEL_PATH), "exec"), + module.__dict__, + ) + finally: + if previous_module is None: + sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = previous_module + return module + + +@lru_cache(maxsize=1) +def _load_official_proxy_module() -> types.ModuleType: + """Import the official proxy while replacing unrelated model dependencies.""" + if not OFFICIAL_PROXY_PATH.is_file(): + raise FileNotFoundError( + f"Official MAGI-2 data-proxy source is missing: {OFFICIAL_PROXY_PATH}" + ) + if importlib.util.find_spec("unfoldNd") is None: + pytest.skip( + "Official MAGI-2 data-proxy parity requires unfoldNd==0.2.3." + ) + + model_stub = _load_official_model_utility_module() + distributed_stub = types.ModuleType("inference.infra.distributed") + distributed_stub.psm = _ParallelStateStub() + replacement_modules = { + "inference.model.magi2_preview": model_stub, + "inference.infra.distributed": distributed_stub, + } + missing_module = object() + previous_modules = { + name: sys.modules.get(name, missing_module) + for name in replacement_modules + } + sys.modules.update(replacement_modules) + + official_path_entry = str(OFFICIAL_REPO_ROOT) + added_path = official_path_entry not in sys.path + if added_path: + sys.path.insert(0, official_path_entry) + module_name = "magi2_official_preview_data_proxy" + try: + spec = importlib.util.spec_from_file_location(module_name, OFFICIAL_PROXY_PATH) + if spec is None or spec.loader is None: + raise ImportError( + f"Cannot create an import spec for {OFFICIAL_PROXY_PATH}" + ) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + finally: + if added_path: + sys.path.remove(official_path_entry) + for name, previous_module in previous_modules.items(): + if previous_module is missing_module: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous_module + return module + + +def _build_proxy_pair( + official_module: types.ModuleType, + config: Magi2DataProxyConfig, +) -> tuple[Any, Magi2DataProxy]: + """Construct official and FastVideo proxies from identical config values.""" + official_config = official_module.DataProxyConfig(**asdict(config)) + return ( + official_module.Magi2DataProxy(official_config), + Magi2DataProxy(config), + ) + + +def _build_model_input_pair( + official_module: types.ModuleType, + input_tensors: dict[str, Any], +) -> tuple[Any, ModelInput]: + """Construct both input dataclasses over the same immutable test tensors.""" + return ( + official_module.ModelInput(**input_tensors), + ModelInput(**input_tensors), + ) + + +def _assert_tensor_exact( + fastvideo_tensor: torch.Tensor, + official_tensor: torch.Tensor, +) -> None: + """Require identical tensor metadata and values.""" + assert fastvideo_tensor.shape == official_tensor.shape + assert fastvideo_tensor.dtype == official_tensor.dtype + assert fastvideo_tensor.stride() == official_tensor.stride() + assert_close(fastvideo_tensor, official_tensor, atol=0, rtol=0) + + +def _assert_packed_inputs_exact( + fastvideo_output: tuple[torch.Tensor, torch.Tensor, torch.Tensor, VarlenHandler, torch.Tensor], + official_output: tuple[torch.Tensor, torch.Tensor, torch.Tensor, Any, torch.Tensor], +) -> None: + """Compare all packed tensors and variable-length attention boundaries.""" + fastvideo_tokens, fastvideo_coords, fastvideo_modalities, fastvideo_varlen, fastvideo_time = fastvideo_output + official_tokens, official_coords, official_modalities, official_varlen, official_time = official_output + _assert_tensor_exact(fastvideo_tokens, official_tokens) + _assert_tensor_exact(fastvideo_coords, official_coords) + _assert_tensor_exact(fastvideo_modalities, official_modalities) + _assert_tensor_exact(fastvideo_time, official_time) + _assert_tensor_exact(fastvideo_varlen.cu_seqlens_q, official_varlen.cu_seqlens_q) + _assert_tensor_exact(fastvideo_varlen.cu_seqlens_k, official_varlen.cu_seqlens_k) + assert fastvideo_varlen.max_seqlen_q == official_varlen.max_seqlen_q + assert fastvideo_varlen.max_seqlen_k == official_varlen.max_seqlen_k + + +def _make_t2v_input_tensors() -> dict[str, Any]: + """Create two distinguishable CFG samples with video, audio, and text.""" + video = torch.arange(2 * 3 * 2 * 4 * 4, dtype=torch.float32).reshape( + 2, 3, 2, 4, 4 + ) + video[1].add_(1000) + audio = torch.arange(2 * 5 * 4, dtype=torch.float32).reshape(2, 5, 4) + audio[1].add_(2000) + text = torch.arange(2 * 4 * 6, dtype=torch.float32).reshape(2, 4, 6) + text[1].add_(3000) + per_token_video_time = torch.linspace( + 0.05, + 0.85, + steps=2 * 2 * 4 * 4, + dtype=torch.float32, + ).reshape(2, 1, 2, 4, 4) + per_token_audio_time = torch.linspace( + 0.1, + 0.9, + steps=2 * 5, + dtype=torch.float32, + ).reshape(2, 5, 1) + return { + "x_t": video, + "audio_x_t": audio, + "audio_feat_len": torch.tensor([5, 3], dtype=torch.int32), + "txt_feat": text, + "txt_feat_len": torch.tensor([4, 2], dtype=torch.int32), + "t": torch.tensor([750, 750], dtype=torch.int64), + "per_token_video_t": per_token_video_time, + "per_token_audio_t": per_token_audio_time, + } + + +def _make_i2v_input_tensors() -> dict[str, Any]: + """Create two CFG samples with two conditioning images per sample.""" + generator = torch.Generator(device="cpu").manual_seed(314159) + return { + "x_t": torch.randn( + (2, 3, 2, 4, 4), + generator=generator, + dtype=torch.float32, + ), + "audio_x_t": torch.randn( + (2, 4, 4), + generator=generator, + dtype=torch.float32, + ), + "audio_feat_len": torch.tensor([4, 2], dtype=torch.int32), + "txt_feat": torch.randn( + (2, 3, 6), + generator=generator, + dtype=torch.float32, + ).to(torch.bfloat16), + "txt_feat_len": torch.tensor([3, 2], dtype=torch.int32), + "t": torch.tensor([0.6, 0.6], dtype=torch.float32), + "per_token_video_t": torch.full( + (2, 1, 2, 4, 4), + 0.6, + dtype=torch.float32, + ), + "per_token_audio_t": torch.full( + (2, 4, 1), + 0.6, + dtype=torch.float32, + ), + "ref_image_feat": torch.randn( + (2, 2, 3, 1, 4, 4), + generator=generator, + dtype=torch.float32, + ), + "ref_image_feat_len": torch.tensor( + [[[2, 2], [2, 2]], [[2, 2], [2, 2]]], + dtype=torch.int32, + ), + "ref_image_special_token_embedding": torch.randn( + (2, 2, 6), + generator=generator, + dtype=torch.float32, + ).to(torch.bfloat16), + } + + +def _assert_process_output_exact( + official_proxy: Any, + fastvideo_proxy: Magi2DataProxy, + packed_token_count: int, + output_channel: int, +) -> None: + """Compare depacked video and audio after deterministic model-like output.""" + model_output = torch.arange( + packed_token_count * output_channel, + dtype=torch.float32, + ).reshape(packed_token_count, output_channel) + official_video, official_audio = official_proxy.process_output(model_output) + fastvideo_video, fastvideo_audio = fastvideo_proxy.process_output(model_output) + _assert_tensor_exact(fastvideo_video, official_video) + _assert_tensor_exact(fastvideo_audio, official_audio) + + +def test_magi2_data_proxy_config_defaults_match_official() -> None: + """Keep FastVideo proxy defaults aligned with the official configuration.""" + official_module = _load_official_proxy_module() + official_defaults = official_module.DataProxyConfig().model_dump() + + assert asdict(Magi2DataProxyConfig()) == official_defaults + + +@pytest.mark.parametrize( + ("parallel_dimension", "cp_size", "ep_size"), + [("cp", 2, 1), ("ep", 1, 2)], + ids=("context_parallel", "expert_parallel"), +) +def test_process_input_distributed_remote_max_padding_matches_official( + monkeypatch: pytest.MonkeyPatch, + parallel_dimension: str, + cp_size: int, + ep_size: int, +) -> None: + """Match distributed maximum reduction and multiple-of-48 padding.""" + official_module = _load_official_proxy_module() + parallel_state = _ParallelStateStub(cp_size=cp_size, ep_size=ep_size) + monkeypatch.setattr(official_module, "psm", parallel_state) + monkeypatch.setattr(fastvideo_proxy_module, "psm", parallel_state) + reduced_groups: list[Any] = [] + + def fake_all_reduce( + tensor: torch.Tensor, + op: Any = None, + group: Any = None, + ) -> None: + """Simulate a remote rank whose packed sequence contains 49 tokens.""" + del op + reduced_groups.append(group) + tensor.fill_(49) + + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + config = Magi2DataProxyConfig( + t_patch_size=2, + patch_size=2, + spatial_rope_interpolation="extra", + add_time_token=False, + time_channel_dim=1, + ) + official_proxy, fastvideo_proxy = _build_proxy_pair(official_module, config) + official_input, fastvideo_input = _build_model_input_pair( + official_module, + _make_t2v_input_tensors(), + ) + + official_output = official_proxy.process_input(official_input) + fastvideo_output = fastvideo_proxy.process_input(fastvideo_input) + _assert_packed_inputs_exact(fastvideo_output, official_output) + + tokens, _, modalities, varlen_handler, time_features = fastvideo_output + assert tokens.shape == (96, 24) + assert varlen_handler.cu_seqlens_q.tolist() == [0, 13, 22, 96] + assert varlen_handler.cu_seqlens_k.tolist() == [0, 13, 22, 96] + assert varlen_handler.max_seqlen_q == 74 + assert varlen_handler.max_seqlen_k == 74 + assert modalities[22:].tolist() == [int(Modality.TEXT)] * 74 + assert torch.count_nonzero(tokens[22:]) == 0 + assert torch.count_nonzero(time_features[22:]) == 0 + assert reduced_groups == [ + f"{parallel_dimension}-group", + f"{parallel_dimension}-group", + ] + _assert_process_output_exact( + official_proxy, + fastvideo_proxy, + packed_token_count=96, + output_channel=24, + ) + + +def test_process_input_t2v_cfg_and_cp_padding_match_official( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Match T2V packing, CFG boundaries, time features, and CP padding.""" + official_module = _load_official_proxy_module() + parallel_state = _ParallelStateStub(cp_size=2, ep_size=1) + monkeypatch.setattr(official_module, "psm", parallel_state) + monkeypatch.setattr(fastvideo_proxy_module, "psm", parallel_state) + config = Magi2DataProxyConfig( + t_patch_size=2, + patch_size=2, + spatial_rope_interpolation="extra", + add_time_token=False, + time_channel_dim=7, + ) + official_proxy, fastvideo_proxy = _build_proxy_pair(official_module, config) + official_input, fastvideo_input = _build_model_input_pair( + official_module, + _make_t2v_input_tensors(), + ) + + official_output = official_proxy.process_input(official_input) + fastvideo_output = fastvideo_proxy.process_input(fastvideo_input) + _assert_packed_inputs_exact(fastvideo_output, official_output) + + tokens, _, modalities, varlen_handler, time_features = fastvideo_output + assert tokens.shape == (48, 24) + assert varlen_handler.cu_seqlens_q.tolist() == [0, 13, 22, 48] + assert varlen_handler.cu_seqlens_k.tolist() == [0, 13, 22, 48] + expected_modalities = ( + [int(Modality.VIDEO)] * 4 + + [int(Modality.AUDIO)] * 5 + + [int(Modality.TEXT)] * 4 + + [int(Modality.VIDEO)] * 4 + + [int(Modality.AUDIO)] * 3 + + [int(Modality.TEXT)] * 2 + + [int(Modality.TEXT)] * 26 + ) + assert modalities.tolist() == expected_modalities + assert not torch.equal(tokens[:4], tokens[13:17]) + assert_close( + time_features[9:13], + torch.tensor([1, 1, 1, 0, 0, 0, 0], dtype=torch.float32).expand(4, -1), + atol=0, + rtol=0, + ) + assert torch.count_nonzero(time_features[22:]) == 0 + assert official_proxy._saved_data.keys() == fastvideo_proxy._saved_data.keys() + _assert_process_output_exact( + official_proxy, + fastvideo_proxy, + packed_token_count=48, + output_channel=24, + ) + + +def test_process_input_i2v_token_order_and_coordinates_match_official( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Match I2V special-token order, image coordinates, times, and depacking.""" + official_module = _load_official_proxy_module() + parallel_state = _ParallelStateStub(cp_size=1, ep_size=1) + monkeypatch.setattr(official_module, "psm", parallel_state) + monkeypatch.setattr(fastvideo_proxy_module, "psm", parallel_state) + config = Magi2DataProxyConfig( + t_patch_size=1, + patch_size=2, + spatial_rope_interpolation="extra", + add_time_token=False, + time_channel_dim=1, + ) + official_proxy, fastvideo_proxy = _build_proxy_pair(official_module, config) + input_tensors = _make_i2v_input_tensors() + official_input, fastvideo_input = _build_model_input_pair( + official_module, + input_tensors, + ) + + official_output = official_proxy.process_input(official_input) + fastvideo_output = fastvideo_proxy.process_input(fastvideo_input) + _assert_packed_inputs_exact(fastvideo_output, official_output) + + tokens, coords, modalities, varlen_handler, time_features = fastvideo_output + assert tokens.shape == (47, 12) + assert varlen_handler.cu_seqlens_q.tolist() == [0, 25, 47] + first_sample_modalities = ( + [int(Modality.VIDEO)] * 8 + + [int(Modality.AUDIO)] * 4 + + [int(Modality.TEXT)] * 3 + + [int(Modality.TEXT)] + + [int(Modality.VIDEO)] * 4 + + [int(Modality.TEXT)] + + [int(Modality.VIDEO)] * 4 + ) + second_sample_modalities = ( + [int(Modality.VIDEO)] * 8 + + [int(Modality.AUDIO)] * 2 + + [int(Modality.TEXT)] * 2 + + [int(Modality.TEXT)] + + [int(Modality.VIDEO)] * 4 + + [int(Modality.TEXT)] + + [int(Modality.VIDEO)] * 4 + ) + assert modalities.tolist() == first_sample_modalities + second_sample_modalities + + first_special_offset = 8 + 4 + 3 + first_special = input_tensors["ref_image_special_token_embedding"][0, 0] + _assert_tensor_exact( + tokens[first_special_offset, :6], + first_special.to(tokens.dtype), + ) + first_image_patch = input_tensors["ref_image_feat"][0, 0, :, :, :2, :2].flatten() + _assert_tensor_exact(tokens[first_special_offset + 1], first_image_patch) + assert coords[first_special_offset].tolist() == [4, -1, -1, 1, 2, 2, 1, 2, 2] + assert coords[first_special_offset + 5].tolist() == [5, -1, -1, 1, 2, 2, 1, 2, 2] + assert torch.count_nonzero(time_features[12:25]) == 0 + assert official_proxy._saved_data.keys() == fastvideo_proxy._saved_data.keys() + _assert_process_output_exact( + official_proxy, + fastvideo_proxy, + packed_token_count=47, + output_channel=12, + ) diff --git a/tests/local_tests/magi2/test_magi2_preview_transformer_parity.py b/tests/local_tests/magi2/test_magi2_preview_transformer_parity.py new file mode 100644 index 000000000..90fe16040 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_preview_transformer_parity.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict eight-GPU numerical parity for the MAGI-2 preview transformer.""" + +from __future__ import annotations + +import os +from pathlib import Path +import subprocess +import sys +from typing import Any + +import pytest +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_ROOT = Path( + os.environ.get("MAGI2_OFFICIAL_REF_DIR", REPO_ROOT.parent / "MAGI-2-preview") +) +WEIGHTS_ROOT = Path( + os.environ.get("MAGI2_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "magi2") +) +WORKER_PATH = Path(__file__).with_name("_preview_transformer_parity_worker.py") +CAPTURE_ROOT = REPO_ROOT / "archived" / "magi2_parity" / "validation" / "preview_transformer" +OFFICIAL_REVISION = "073c84f2102ec3c9287623113a103c14402770ad" +WORLD_SIZE = 8 + + +pytestmark = pytest.mark.skipif( + torch.cuda.device_count() < WORLD_SIZE, + reason=f"MAGI-2 preview numerical parity requires {WORLD_SIZE} CUDA devices", +) + + +def _require_parity_sources() -> None: + """Require the pinned official source and complete preview checkpoint.""" + model_path = OFFICIAL_ROOT / "inference" / "model" / "magi2_preview.py" + if not model_path.is_file(): + raise AssertionError(f"Official MAGI-2 source is missing: {model_path}") + index_path = WEIGHTS_ROOT / "preview" / "model.safetensors.index.json" + if not index_path.is_file(): + raise AssertionError(f"MAGI-2 preview checkpoint index is missing: {index_path}") + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=OFFICIAL_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + assert revision == OFFICIAL_REVISION + + +def _run_implementation(implementation: str) -> Path: + """Launch one implementation in an isolated eight-rank torchrun job.""" + output_dir = CAPTURE_ROOT / implementation + environment = os.environ.copy() + environment.update( + { + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "MAGI2_CKPT_ROOT": str(WEIGHTS_ROOT), + "MAGI2_DETERMINISTIC": "1", + "MAGI2_LOCAL_WEIGHTS_DIR": str(WEIGHTS_ROOT), + "MAGI2_OFFICIAL_REF_DIR": str(OFFICIAL_ROOT), + "MAGI_ATTENTION_DETERMINISTIC_MODE": "1", + "MAGI_COMPILE_COMPILE_MODE": "NONE", + "OMP_NUM_THREADS": "1", + "PYTHONHASHSEED": "42", + } + ) + environment.pop("SKIP_LOAD_MODEL", None) + completed = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + f"--nproc-per-node={WORLD_SIZE}", + str(WORKER_PATH), + "--implementation", + implementation, + "--output-dir", + str(output_dir), + ], + cwd=REPO_ROOT, + env=environment, + check=False, + capture_output=True, + text=True, + timeout=7200, + ) + if completed.returncode != 0: + raise AssertionError( + f"{implementation} preview transformer torchrun failed.\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + return output_dir + + +def _assert_exact(actual: Any, expected: Any, path: str) -> None: + """Recursively require identical capture structure, metadata, and values.""" + assert type(actual) is type(expected), ( + f"{path}: type differs: {type(actual).__name__} versus " + f"{type(expected).__name__}" + ) + if isinstance(expected, torch.Tensor): + assert actual.shape == expected.shape, path + assert actual.dtype == expected.dtype, path + assert actual.stride() == expected.stride(), path + if not torch.equal(actual, expected): + difference = (actual.float() - expected.float()).abs() + raise AssertionError( + f"{path}: tensor values differ; mismatched=" + f"{torch.count_nonzero(actual != expected).item()}, " + f"max_abs={difference.max().item()}" + ) + return + if isinstance(expected, dict): + assert actual.keys() == expected.keys(), path + for key in expected: + _assert_exact(actual[key], expected[key], f"{path}.{key}") + return + if isinstance(expected, (list, tuple)): + assert len(actual) == len(expected), path + for index, (actual_value, expected_value) in enumerate( + zip(actual, expected, strict=True) + ): + _assert_exact(actual_value, expected_value, f"{path}[{index}]") + return + assert actual == expected, path + + +def _load_rank_capture(output_dir: Path, rank: int) -> dict[str, Any]: + """Load one rank artifact and validate its fixed capture envelope.""" + artifact_path = output_dir / f"rank_{rank}.pt" + if not artifact_path.is_file(): + raise AssertionError(f"Preview parity capture is missing: {artifact_path}") + artifact = torch.load(artifact_path, map_location="cpu", weights_only=True) + assert artifact["schema_version"] == 1 + assert artifact["rank"] == rank + assert artifact["world_size"] == WORLD_SIZE + assert artifact["cases"].keys() == {"t2v", "i2v"} + return artifact + + +def test_magi2_preview_transformer_matches_official_exactly() -> None: + """Match T2V and I2V tensors at every preview transformer boundary.""" + _require_parity_sources() + official_output_dir = _run_implementation("official") + fastvideo_output_dir = _run_implementation("fastvideo") + for rank in range(WORLD_SIZE): + official_capture = _load_rank_capture(official_output_dir, rank) + fastvideo_capture = _load_rank_capture(fastvideo_output_dir, rank) + assert official_capture["implementation"] == "official" + assert fastvideo_capture["implementation"] == "fastvideo" + _assert_exact( + fastvideo_capture["cases"], + official_capture["cases"], + f"rank_{rank}.cases", + ) diff --git a/tests/local_tests/magi2/test_magi2_refiner_data_proxy_parity.py b/tests/local_tests/magi2/test_magi2_refiner_data_proxy_parity.py new file mode 100644 index 000000000..2111a1960 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_refiner_data_proxy_parity.py @@ -0,0 +1,604 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict MAGI-2 refiner data-proxy implementation parity. + +Coverage scope: implementation_subcomponent. The tests load the pinned official +``Magi2RefinerDataProxy`` source and compare every packed or depacked tensor +with the FastVideo-owned component. A scheduler check verifies that uneven +context-parallel splitting remains owned by ``UlyssesScheduler``. +""" + +from __future__ import annotations + +import ast +from dataclasses import asdict, dataclass, fields +from functools import lru_cache +import importlib +import importlib.util +import os +from pathlib import Path +import sys +import types +import typing +from typing import Any + +import pytest +import torch +from torch.testing import assert_close +from torch.utils._pytree import tree_map + +from fastvideo.pipelines.basic.magi2.stages.refiner_data_proxy import ( + Magi2RefinerDataProxy, + Magi2RefinerDataProxyConfig, + Modality, + RefinerModelInput, + VarlenHandler, + WindowLocalAttnHandler, +) + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_REPO_ROOT = Path( + os.environ.get( + "MAGI2_OFFICIAL_REF_DIR", + REPO_ROOT.parent / "MAGI-2-preview", + ) +) +OFFICIAL_PROXY_PATH = OFFICIAL_REPO_ROOT / "inference" / "pipeline" / "refiner_data_proxy.py" +OFFICIAL_MODEL_PATH = OFFICIAL_REPO_ROOT / "inference" / "model" / "magi2_refiner.py" +OFFICIAL_SCHEDULER_PATH = ( + OFFICIAL_REPO_ROOT + / "inference" + / "infra" + / "parallelism" + / "context_parallel" + / "ulysses_scheduler.py" +) +PARITY_SCOPE = "implementation_subcomponent" + + +@dataclass(frozen=True) +class _ParallelStateStub: + """Expose one fixed context-parallel size and process-group marker.""" + + cp_size: int + + def get_world_size(self, dimension: str = "") -> int: + """Return the configured context-parallel world size.""" + return self.cp_size if dimension == "cp" else 1 + + def get_parallel_group(self, dimension: str) -> str: + """Return an identifiable process-group marker for call assertions.""" + return f"{dimension}-group" + + +def _load_official_refiner_model_utilities() -> types.ModuleType: + """Load the official frame-local attention helper without model dependencies.""" + if not OFFICIAL_MODEL_PATH.is_file(): + raise FileNotFoundError(f"Official MAGI-2 refiner source is missing: {OFFICIAL_MODEL_PATH}") + source_module = ast.parse( + OFFICIAL_MODEL_PATH.read_text(encoding="utf-8"), + filename=str(OFFICIAL_MODEL_PATH), + ) + symbol_names = { + "FFAHandler", + "calc_local_qk_range", + "calc_local_attn_ffa_handler", + } + symbol_nodes = [ + node + for node in source_module.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + and node.name in symbol_names + ] + found_names = {node.name for node in symbol_nodes} + if found_names != symbol_names: + raise AssertionError( + f"Official refiner utility definitions differ: expected {symbol_names}, found {found_names}" + ) + + module = types.ModuleType("inference.model.magi2_refiner") + module.__file__ = str(OFFICIAL_MODEL_PATH) + module.__dict__.update( + { + "Optional": typing.Optional, + "dataclass": dataclass, + "torch": torch, + } + ) + module_name = module.__name__ + previous_module = sys.modules.get(module_name) + sys.modules[module_name] = module + try: + exec( + compile( + ast.Module(body=symbol_nodes, type_ignores=[]), + str(OFFICIAL_MODEL_PATH), + "exec", + ), + module.__dict__, + ) + finally: + if previous_module is None: + sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = previous_module + return module + + +@lru_cache(maxsize=1) +def _load_official_proxy_module() -> types.ModuleType: + """Import the official proxy while isolating unrelated refiner model code.""" + if not OFFICIAL_PROXY_PATH.is_file(): + raise FileNotFoundError(f"Official MAGI-2 refiner proxy is missing: {OFFICIAL_PROXY_PATH}") + if importlib.util.find_spec("unfoldNd") is None: + pytest.skip("Official MAGI-2 refiner data-proxy parity requires unfoldNd==0.2.3.") + + model_module_name = "inference.model.magi2_refiner" + missing_module = object() + previous_model_module = sys.modules.get(model_module_name, missing_module) + sys.modules[model_module_name] = _load_official_refiner_model_utilities() + official_path_entry = str(OFFICIAL_REPO_ROOT) + added_path = official_path_entry not in sys.path + if added_path: + sys.path.insert(0, official_path_entry) + module_name = "magi2_official_refiner_data_proxy" + try: + spec = importlib.util.spec_from_file_location(module_name, OFFICIAL_PROXY_PATH) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot create an import spec for {OFFICIAL_PROXY_PATH}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + finally: + if added_path: + sys.path.remove(official_path_entry) + if previous_model_module is missing_module: + sys.modules.pop(model_module_name, None) + else: + sys.modules[model_module_name] = previous_model_module + return module + + +def _load_official_scheduler_class() -> type[Any]: + """Load the official Ulysses scheduler class with injectable collectives.""" + if not OFFICIAL_SCHEDULER_PATH.is_file(): + raise FileNotFoundError(f"Official Ulysses scheduler is missing: {OFFICIAL_SCHEDULER_PATH}") + source_module = ast.parse( + OFFICIAL_SCHEDULER_PATH.read_text(encoding="utf-8"), + filename=str(OFFICIAL_SCHEDULER_PATH), + ) + class_nodes = [ + node + for node in source_module.body + if isinstance(node, ast.ClassDef) and node.name == "UlyssesScheduler" + ] + if len(class_nodes) != 1: + raise AssertionError("Official source must define exactly one UlyssesScheduler class") + module = types.ModuleType("magi2_official_ulysses_scheduler") + module.__dict__.update( + { + "Generic": typing.Generic, + "List": list, + "Optional": typing.Optional, + "T": typing.TypeVar("T"), + "torch": torch, + "tree_map": tree_map, + } + ) + exec( + compile( + ast.Module(body=class_nodes, type_ignores=[]), + str(OFFICIAL_SCHEDULER_PATH), + "exec", + ), + module.__dict__, + ) + return module.UlyssesScheduler + + +def _cuda_device() -> torch.device: + """Return the CUDA device required by the pinned official proxy.""" + if not torch.cuda.is_available(): + pytest.skip("Official MAGI-2 refiner data-proxy parity requires CUDA.") + return torch.device("cuda", torch.cuda.current_device()) + + +def _build_proxy_pair( + official_module: types.ModuleType, + config: Magi2RefinerDataProxyConfig, +) -> tuple[Any, Magi2RefinerDataProxy]: + """Construct official and FastVideo proxies from identical config values.""" + official_config = official_module.Magi2RefinerDataProxyConfig(**asdict(config)) + return ( + official_module.Magi2RefinerDataProxy(official_config), + Magi2RefinerDataProxy(config), + ) + + +def _build_model_input_pair( + input_tensors: dict[str, torch.Tensor], +) -> tuple[types.SimpleNamespace, RefinerModelInput]: + """Construct both proxy inputs over the same immutable test tensors.""" + return ( + types.SimpleNamespace(**input_tensors), + RefinerModelInput(**input_tensors), + ) + + +def _make_input_tensors( + *, + device: torch.device, + batch_size: int, + ref_video_token_limit: int, + distinguish_cfg_samples: bool = False, +) -> dict[str, torch.Tensor]: + """Create deterministic refiner video, audio, text, and reference features.""" + video = torch.arange( + batch_size * 3 * 3 * 6 * 8, + device=device, + dtype=torch.float32, + ).reshape(batch_size, 3, 3, 6, 8) + audio = torch.arange( + batch_size * 5 * 4, + device=device, + dtype=torch.float32, + ).reshape(batch_size, 5, 4) + text = torch.arange( + batch_size * 4 * 6, + device=device, + dtype=torch.float32, + ).reshape(batch_size, 4, 6) + reference_audio = torch.arange( + batch_size * 3 * 4, + device=device, + dtype=torch.float32, + ).reshape(batch_size, 3, 4) + reference_video = torch.arange( + batch_size * 3 * 3 * 6 * 8, + device=device, + dtype=torch.float32, + ).reshape(batch_size, 3, 3, 6, 8) + if distinguish_cfg_samples and batch_size > 1: + video[1].add_(1000) + audio[1].add_(2000) + text[1].add_(3000) + reference_audio[1].add_(4000) + reference_video[1].add_(5000) + return { + "x_t": video, + "audio_x_t": audio, + "audio_feat_len": torch.full((batch_size,), 4, dtype=torch.int64), + "txt_feat": text, + "txt_feat_len": torch.full((batch_size,), 3, dtype=torch.int64), + "ref_audio_feat": reference_audio, + "ref_audio_feat_len": torch.full((batch_size,), 2, dtype=torch.int64), + "ref_video_feat": reference_video, + "ref_video_feat_len": torch.full( + (batch_size,), + ref_video_token_limit, + dtype=torch.int64, + ), + } + + +def _assert_tensor_exact( + fastvideo_tensor: torch.Tensor, + official_tensor: torch.Tensor, +) -> None: + """Require identical tensor metadata and values.""" + assert fastvideo_tensor.shape == official_tensor.shape + assert fastvideo_tensor.dtype == official_tensor.dtype + assert fastvideo_tensor.stride() == official_tensor.stride() + assert_close(fastvideo_tensor, official_tensor, atol=0, rtol=0) + + +def _assert_scalar_or_tensor_exact( + fastvideo_value: int | torch.Tensor, + official_value: int | torch.Tensor, +) -> None: + """Compare attention-length metadata without discarding tensor dtype.""" + if isinstance(official_value, torch.Tensor): + assert isinstance(fastvideo_value, torch.Tensor) + _assert_tensor_exact(fastvideo_value, official_value) + else: + assert fastvideo_value == official_value + + +def _assert_local_attn_handler_exact( + fastvideo_handler: WindowLocalAttnHandler | None, + official_handler: Any, +) -> None: + """Compare every attention-range field exposed by the official handler.""" + if official_handler is None: + assert fastvideo_handler is None + return + assert fastvideo_handler is not None + _assert_tensor_exact(fastvideo_handler.q_ranges, official_handler.q_ranges) + _assert_tensor_exact(fastvideo_handler.k_ranges, official_handler.k_ranges) + _assert_tensor_exact(fastvideo_handler.attn_type_map, official_handler.attn_type_map) + assert fastvideo_handler.max_seqlen_q == official_handler.max_seqlen_q + assert fastvideo_handler.max_seqlen_k == official_handler.max_seqlen_k + assert fastvideo_handler.softmax_scale == official_handler.softmax_scale + for field_name in ( + "bwd_q_ranges", + "bwd_k_ranges", + "bwd_attn_type_map", + ): + official_value = getattr(official_handler, field_name, None) + fastvideo_value = getattr(fastvideo_handler, field_name) + if official_value is None: + assert fastvideo_value is None + else: + assert fastvideo_value is not None + _assert_tensor_exact(fastvideo_value, official_value) + assert fastvideo_handler.auto_range_merge == official_handler.auto_range_merge + assert fastvideo_handler.sparse_load == official_handler.sparse_load + + +def _assert_packed_inputs_exact( + fastvideo_output: tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + VarlenHandler, + WindowLocalAttnHandler | None, + ], + official_output: tuple[torch.Tensor, torch.Tensor, torch.Tensor, Any, Any], +) -> None: + """Compare all packed tensors and attention metadata.""" + ( + fastvideo_tokens, + fastvideo_coords, + fastvideo_modalities, + fastvideo_varlen, + fastvideo_local_attn, + ) = fastvideo_output + ( + official_tokens, + official_coords, + official_modalities, + official_varlen, + official_local_attn, + ) = official_output + _assert_tensor_exact(fastvideo_tokens, official_tokens) + _assert_tensor_exact(fastvideo_coords, official_coords) + _assert_tensor_exact(fastvideo_modalities, official_modalities) + _assert_tensor_exact(fastvideo_varlen.cu_seqlens_q, official_varlen.cu_seqlens_q) + _assert_tensor_exact(fastvideo_varlen.cu_seqlens_k, official_varlen.cu_seqlens_k) + _assert_scalar_or_tensor_exact(fastvideo_varlen.max_seqlen_q, official_varlen.max_seqlen_q) + _assert_scalar_or_tensor_exact(fastvideo_varlen.max_seqlen_k, official_varlen.max_seqlen_k) + _assert_local_attn_handler_exact(fastvideo_local_attn, official_local_attn) + + +def _assert_process_output_exact( + official_proxy: Any, + fastvideo_proxy: Magi2RefinerDataProxy, + packed_token_count: int, + output_channel: int, + device: torch.device, +) -> None: + """Compare depacked video and audio after deterministic model-like output.""" + model_output = torch.arange( + packed_token_count * output_channel, + device=device, + dtype=torch.float32, + ).reshape(packed_token_count, output_channel) + official_video, official_audio = official_proxy.process_output(model_output) + fastvideo_video, fastvideo_audio = fastvideo_proxy.process_output(model_output) + _assert_tensor_exact(fastvideo_video, official_video) + _assert_tensor_exact(fastvideo_audio, official_audio) + + +def test_magi2_refiner_data_proxy_config_defaults_match_official() -> None: + """Keep FastVideo refiner proxy defaults aligned with the official config.""" + official_module = _load_official_proxy_module() + official_defaults = official_module.Magi2RefinerDataProxyConfig().model_dump() + + assert asdict(Magi2RefinerDataProxyConfig()) == official_defaults + + +def test_process_input_default_frame_local_attention_matches_official() -> None: + """Match default patching, coordinates, varlen metadata, and frame ranges.""" + device = _cuda_device() + official_module = _load_official_proxy_module() + config = Magi2RefinerDataProxyConfig() + official_proxy, fastvideo_proxy = _build_proxy_pair(official_module, config) + official_input, fastvideo_input = _build_model_input_pair( + _make_input_tensors( + device=device, + batch_size=1, + ref_video_token_limit=0, + ) + ) + + official_output = official_proxy.process_input(official_input) + fastvideo_output = fastvideo_proxy.process_input(fastvideo_input) + _assert_packed_inputs_exact(fastvideo_output, official_output) + + tokens, _, modalities, varlen_handler, local_attn_handler = fastvideo_output + assert tokens.shape == (45, 12) + assert modalities.tolist() == ( + [int(Modality.VIDEO)] * 36 + + [int(Modality.AUDIO)] * 4 + + [int(Modality.TEXT)] * 3 + + [int(Modality.AUDIO)] * 2 + ) + assert varlen_handler.cu_seqlens_q.tolist() == [0, 45] + assert local_attn_handler is not None + _assert_process_output_exact( + official_proxy, + fastvideo_proxy, + packed_token_count=tokens.shape[0], + output_channel=tokens.shape[1], + device=device, + ) + + +def test_process_input_cfg_boundaries_and_no_time_features_match_official() -> None: + """Preserve adjacent CFG samples without adding timestep feature tokens.""" + device = _cuda_device() + official_module = _load_official_proxy_module() + config = Magi2RefinerDataProxyConfig( + patch_size=2, + frame_receptive_field=-1, + ) + official_proxy, fastvideo_proxy = _build_proxy_pair(official_module, config) + input_tensors = _make_input_tensors( + device=device, + batch_size=2, + ref_video_token_limit=1, + distinguish_cfg_samples=True, + ) + official_input, fastvideo_input = _build_model_input_pair(input_tensors) + + official_output = official_proxy.process_input(official_input) + fastvideo_output = fastvideo_proxy.process_input(fastvideo_input) + _assert_packed_inputs_exact(fastvideo_output, official_output) + + tokens, _, _, varlen_handler, local_attn_handler = fastvideo_output + assert tokens.shape == (92, 12) + assert varlen_handler.cu_seqlens_q.tolist() == [0, 46, 92] + assert local_attn_handler is None + assert not torch.equal(tokens[:36], tokens[46:82]) + assert "t" not in {input_field.name for input_field in fields(RefinerModelInput)} + assert len(fastvideo_output) == 5 + _assert_process_output_exact( + official_proxy, + fastvideo_proxy, + packed_token_count=tokens.shape[0], + output_channel=tokens.shape[1], + device=device, + ) + + +def test_process_input_shipping_block_window_and_depacking_match_official() -> None: + """Match the release block-grid token order, ranges, and inverse depacking.""" + device = _cuda_device() + official_module = _load_official_proxy_module() + config = Magi2RefinerDataProxyConfig( + t_patch_size=1, + patch_size=1, + frame_receptive_field=11, + spatial_rope_interpolation="extra", + coords_style="v1", + text_offset=0, + attn_config={ + "mode": "window", + "block_t_size": 8, + "block_size": 4, + "window": { + "level": "block", + "block_mode": "grid", + "block_t_radius": 2, + "block_h_radius": 2, + "block_w_radius": 2, + "win_size": 384, + "frame_receptive_field": -1, + "auto_range_merge": True, + "sparse_load": False, + "full_attn_layers": [], + }, + }, + ) + official_proxy, fastvideo_proxy = _build_proxy_pair(official_module, config) + official_input, fastvideo_input = _build_model_input_pair( + _make_input_tensors( + device=device, + batch_size=1, + ref_video_token_limit=4, + ) + ) + + official_output = official_proxy.process_input(official_input) + fastvideo_output = fastvideo_proxy.process_input(fastvideo_input) + _assert_packed_inputs_exact(fastvideo_output, official_output) + + tokens, coords, modalities, varlen_handler, local_attn_handler = fastvideo_output + assert tokens.shape == (157, 6) + assert coords.shape == (157, 9) + assert modalities.shape == (157,) + assert varlen_handler.cu_seqlens_q.tolist() == [0, 157] + assert local_attn_handler is not None + assert local_attn_handler.auto_range_merge is True + _assert_process_output_exact( + official_proxy, + fastvideo_proxy, + packed_token_count=tokens.shape[0], + output_channel=tokens.shape[1], + device=device, + ) + + +def test_process_input_distributed_context_keeps_unpadded_official_sequence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep distributed reduction and uneven splitting outside the refiner proxy.""" + device = _cuda_device() + official_module = _load_official_proxy_module() + config = Magi2RefinerDataProxyConfig( + patch_size=2, + frame_receptive_field=-1, + ) + official_proxy, fastvideo_proxy = _build_proxy_pair(official_module, config) + official_input, fastvideo_input = _build_model_input_pair( + _make_input_tensors( + device=device, + batch_size=1, + ref_video_token_limit=1, + ) + ) + + def reject_all_reduce(*args: Any, **kwargs: Any) -> None: + """Fail if preview-style distributed maximum padding enters this proxy.""" + del args, kwargs + raise AssertionError("The refiner proxy must not perform all_reduce") + + monkeypatch.setattr(torch.distributed, "all_reduce", reject_all_reduce) + official_output = official_proxy.process_input(official_input) + fastvideo_output = fastvideo_proxy.process_input(fastvideo_input) + _assert_packed_inputs_exact(fastvideo_output, official_output) + packed_tokens = fastvideo_output[0] + assert packed_tokens.shape[0] == 46 + assert packed_tokens.shape[0] % 48 != 0 + assert "pad_size" not in fastvideo_proxy._saved_data + assert "pad_size" not in official_proxy._saved_data + + official_scheduler_class = _load_official_scheduler_class() + target_scheduler_module = importlib.import_module( + "fastvideo.models.dits.magi2_runtime.context_parallel.ulysses_scheduler" + ) + parallel_state = _ParallelStateStub(cp_size=3) + official_capture: list[tuple[list[int], Any]] = [] + target_capture: list[tuple[list[int], Any]] = [] + + def official_scatter( + tensor: torch.Tensor, + split_sizes: list[int], + group: Any, + ) -> torch.Tensor: + """Capture the official scheduler's uneven split without communication.""" + official_capture.append((split_sizes, group)) + return tensor + + def target_scatter( + tensor: torch.Tensor, + split_sizes: list[int], + group: Any, + ) -> torch.Tensor: + """Capture the FastVideo scheduler's uneven split without communication.""" + target_capture.append((split_sizes, group)) + return tensor + + official_scheduler_class._dispatch.__globals__["psm"] = parallel_state + official_scheduler_class._dispatch.__globals__["scatter_to_context_parallel_region"] = official_scatter + monkeypatch.setattr(target_scheduler_module, "psm", parallel_state) + monkeypatch.setattr(target_scheduler_module, "scatter_to_context_parallel_region", target_scatter) + official_scheduler = official_scheduler_class() + target_scheduler = target_scheduler_module.UlyssesScheduler() + + official_dispatched = official_scheduler.dispatch(packed_tokens) + target_dispatched = target_scheduler.dispatch(packed_tokens) + _assert_tensor_exact(target_dispatched, official_dispatched) + assert official_scheduler.cp_split_sizes == [16, 15, 15] + assert target_scheduler.cp_split_sizes == [16, 15, 15] + assert official_capture == [([16, 15, 15], "cp-group")] + assert target_capture == official_capture diff --git a/tests/local_tests/magi2/test_magi2_refiner_transformer_parity.py b/tests/local_tests/magi2/test_magi2_refiner_transformer_parity.py new file mode 100644 index 000000000..bce06412b --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_refiner_transformer_parity.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict eight-GPU numerical parity for the MAGI-2 refiner transformer. + +Coverage scope: both. The FastVideo side uses the production checkpoint loader, +and the captures compare every implementation boundary that owns numerical +behavior from the packed proxy input through the depacked video and audio. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import sys +from typing import Any + +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_ROOT = Path( + os.environ.get("MAGI2_OFFICIAL_REF_DIR", REPO_ROOT.parent / "MAGI-2-preview") +) +WEIGHTS_ROOT = Path( + os.environ.get("MAGI2_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "magi2") +) +WORKER_PATH = Path(__file__).with_name("_refiner_transformer_parity_worker.py") +CAPTURE_ROOT = REPO_ROOT / "archived" / "magi2_parity" / "validation" / "refiner_transformer" +OFFICIAL_REVISION = "073c84f2102ec3c9287623113a103c14402770ad" +WORLD_SIZE = 8 + + +def _require_parity_sources() -> None: + """Require eight GPUs, the pinned source, and every refiner weight shard.""" + if torch.cuda.device_count() < WORLD_SIZE: + raise AssertionError( + f"MAGI-2 refiner numerical parity requires {WORLD_SIZE} CUDA devices; " + f"found {torch.cuda.device_count()}" + ) + model_path = OFFICIAL_ROOT / "inference" / "model" / "magi2_refiner.py" + if not model_path.is_file(): + raise AssertionError(f"Official MAGI-2 refiner source is missing: {model_path}") + index_path = WEIGHTS_ROOT / "refiner" / "model.safetensors.index.json" + if not index_path.is_file(): + raise AssertionError(f"MAGI-2 refiner checkpoint index is missing: {index_path}") + weight_map = json.loads(index_path.read_text(encoding="utf-8"))["weight_map"] + missing_shards = sorted( + shard_name + for shard_name in set(weight_map.values()) + if not (index_path.parent / shard_name).is_file() + ) + if missing_shards: + raise AssertionError(f"MAGI-2 refiner checkpoint shards are missing: {missing_shards}") + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=OFFICIAL_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + assert revision == OFFICIAL_REVISION + checkout_changes = subprocess.run( + ["git", "status", "--porcelain=v1", "--untracked-files=all"], + cwd=OFFICIAL_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + assert not checkout_changes, ( + "Official MAGI-2 checkout contains local changes:\n" + checkout_changes + ) + + +def _run_implementation(implementation: str) -> Path: + """Launch one implementation in an isolated eight-rank torchrun job.""" + output_dir = CAPTURE_ROOT / implementation + output_dir.mkdir(parents=True, exist_ok=True) + environment = os.environ.copy() + environment.update( + { + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "MAGI2_CKPT_ROOT": str(WEIGHTS_ROOT), + "MAGI2_DETERMINISTIC": "1", + "MAGI2_LOCAL_WEIGHTS_DIR": str(WEIGHTS_ROOT), + "MAGI2_OFFICIAL_REF_DIR": str(OFFICIAL_ROOT), + "MAGI_ATTENTION_DETERMINISTIC_MODE": "1", + "MAGI_COMPILE_COMPILE_MODE": "NONE", + "OMP_NUM_THREADS": "1", + "PYTHONHASHSEED": "42", + } + ) + environment.pop("SKIP_LOAD_MODEL", None) + completed = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + f"--nproc-per-node={WORLD_SIZE}", + str(WORKER_PATH), + "--implementation", + implementation, + "--output-dir", + str(output_dir), + ], + cwd=REPO_ROOT, + env=environment, + check=False, + capture_output=True, + text=True, + timeout=7200, + ) + log_path = CAPTURE_ROOT / f"{implementation}.log" + log_path.write_text( + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}", + encoding="utf-8", + ) + if completed.returncode != 0: + raise AssertionError( + f"{implementation} refiner transformer torchrun failed; see {log_path}.\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + return output_dir + + +def _assert_exact(actual: Any, expected: Any, path: str) -> None: + """Recursively require identical capture structure, metadata, and values.""" + assert type(actual) is type(expected), ( + f"{path}: type differs: {type(actual).__name__} versus " + f"{type(expected).__name__}" + ) + if isinstance(expected, torch.Tensor): + assert actual.shape == expected.shape, path + assert actual.dtype == expected.dtype, path + assert actual.stride() == expected.stride(), path + if not torch.equal(actual, expected): + difference = (actual.float() - expected.float()).abs() + raise AssertionError( + f"{path}: tensor values differ; mismatched=" + f"{torch.count_nonzero(actual != expected).item()}, " + f"max_abs={difference.max().item()}" + ) + return + if isinstance(expected, dict): + assert actual.keys() == expected.keys(), path + for key in expected: + _assert_exact(actual[key], expected[key], f"{path}.{key}") + return + if isinstance(expected, (list, tuple)): + assert len(actual) == len(expected), path + for index, (actual_value, expected_value) in enumerate( + zip(actual, expected, strict=True) + ): + _assert_exact(actual_value, expected_value, f"{path}[{index}]") + return + assert actual == expected, path + + +def _load_rank_capture(output_dir: Path, rank: int) -> dict[str, Any]: + """Load one rank artifact and validate its fixed capture envelope.""" + artifact_path = output_dir / f"rank_{rank}.pt" + if not artifact_path.is_file(): + raise AssertionError(f"Refiner parity capture is missing: {artifact_path}") + artifact = torch.load(artifact_path, map_location="cpu", weights_only=True) + assert artifact["schema_version"] == 1 + assert artifact["rank"] == rank + assert artifact["world_size"] == WORLD_SIZE + assert set(artifact) == { + "schema_version", + "implementation", + "rank", + "world_size", + "case", + } + return artifact + + +def test_magi2_refiner_transformer_matches_official_exactly() -> None: + """Match every refiner tensor boundary through video and audio depacking.""" + _require_parity_sources() + official_output_dir = _run_implementation("official") + fastvideo_output_dir = _run_implementation("fastvideo") + for rank in range(WORLD_SIZE): + official_capture = _load_rank_capture(official_output_dir, rank) + fastvideo_capture = _load_rank_capture(fastvideo_output_dir, rank) + assert official_capture["implementation"] == "official" + assert fastvideo_capture["implementation"] == "fastvideo" + _assert_exact( + fastvideo_capture["case"], + official_capture["case"], + f"rank_{rank}.case", + ) diff --git a/tests/local_tests/magi2/test_magi2_registry_and_metadata.py b/tests/local_tests/magi2/test_magi2_registry_and_metadata.py new file mode 100644 index 000000000..fcca79768 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_registry_and_metadata.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only registry, preset, and model-index checks for MAGI-2 Preview.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import pytest + +from fastvideo.api.presets import get_preset +from fastvideo.configs.pipelines import Magi2PreviewPipelineConfig +from fastvideo.fastvideo_args import WorkloadType +from fastvideo.registry import ( + _get_config_info, + get_default_preset, + get_model_family, + get_model_info, + get_pipeline_config_cls_from_name, + get_registered_model_paths, +) +from scripts.checkpoint_conversion.convert_magi2_to_fastvideo import ( + MODEL_INDEX, + REQUIRED_COMPONENT_FILES, + SOURCE_REPOSITORY, + convert_checkpoint_layout, +) + + +def _create_indexed_source(source: Path) -> Path: + """Create the smallest valid source layout with one shard per index.""" + indexed_shard = source / "preview" / "preview.safetensors" + for component, relative_paths in REQUIRED_COMPONENT_FILES.items(): + for relative_path in relative_paths: + source_path = source / component / relative_path + source_path.parent.mkdir(parents=True, exist_ok=True) + if relative_path.endswith(".index.json"): + shard_name = f"{component}.safetensors" + source_path.write_text(json.dumps({"weight_map": {"weight": shard_name}}), encoding="utf-8") + shard_path = source_path.parent / shard_name + shard_path.write_bytes(b"weights") + if component == "preview": + indexed_shard = shard_path + else: + source_path.write_bytes(b"component") + return indexed_shard + + +@pytest.mark.parametrize("workload_type", [WorkloadType.T2V, WorkloadType.I2V]) +def test_get_model_info_magi2_model_index_for_each_workload( + tmp_path: Path, + workload_type: WorkloadType, +) -> None: + """Resolve the MAGI-2 pipeline class from local model-index metadata.""" + model_path = tmp_path / "checkpoint" + (model_path / "transformer").mkdir(parents=True) + (model_path / "model_index.json").write_text( + json.dumps({ + "_class_name": "Magi2Pipeline", + "_diffusers_version": "0.37.0", + "transformer": [None, None], + }), + encoding="utf-8", + ) + + model_info = get_model_info(str(model_path), workload_type=workload_type) + config_info = _get_config_info(str(model_path)) + + assert model_info.pipeline_cls.__name__ == "Magi2Pipeline" + assert model_info.pipeline_config_cls is Magi2PreviewPipelineConfig + assert config_info is not None + assert config_info.workload_types == (WorkloadType.T2V, WorkloadType.I2V) + assert get_pipeline_config_cls_from_name(str(model_path)) is Magi2PreviewPipelineConfig + assert get_model_family(str(model_path)) == "magi2" + assert get_default_preset(str(model_path)) == "magi2_preview_1080p" + + +def test_magi2_preview_pipeline_config_published_geometry() -> None: + """Expose the published preview and refiner latent geometry.""" + pipeline_config = Magi2PreviewPipelineConfig() + + pipeline_config.check_pipeline_config() + assert (pipeline_config.preview_width, pipeline_config.preview_height) == (896, 512) + assert (pipeline_config.output_width, pipeline_config.output_height) == (1920, 1088) + assert pipeline_config.output_frames == 249 + assert pipeline_config.output_fps == 25 + assert (pipeline_config.refiner_latent_width, pipeline_config.refiner_latent_height) == (120, 68) + assert pipeline_config.text_encoder_configs == () + assert pipeline_config.text_encoder_precisions == () + + +def test_get_preset_magi2_matches_published_profile() -> None: + """Keep the registered preset aligned with the published inference profile.""" + preset = get_preset("magi2_preview_1080p", "magi2") + + assert preset.defaults["seed"] == 42 + assert preset.defaults["height"] == 1088 + assert preset.defaults["width"] == 1920 + assert preset.defaults["num_frames"] == 249 + assert preset.defaults["fps"] == 25 + assert preset.defaults["num_inference_steps"] == 100 + assert preset.defaults["num_inference_steps_sr"] == 5 + negative_prompt = preset.defaults["negative_prompt"] + assert len(negative_prompt) == 1721 + assert hashlib.sha256(negative_prompt.encode()).hexdigest() == ( + "5ac0746de6c7e0388a16122d9c7e751b1cb067242218bf6f41ec1e88271cdcb9" + ) + + +def test_model_index_fastvideo_components_use_defining_module_paths() -> None: + """Record each FastVideo component with its defining Python module.""" + expected_libraries = { + "audio_vae": "fastvideo.models.vaes.magi2_audio_vae", + "image_encoder": "fastvideo.models.vaes.magi2_wan_loader", + "scheduler": "fastvideo.models.schedulers.scheduling_flow_unipc_multistep", + "text_encoder": "fastvideo.models.encoders.qwen3_5", + "transformer": "fastvideo.models.dits.magi2", + "transformer_2": "fastvideo.models.dits.magi2_refiner", + "vae": "fastvideo.models.vaes.magi2_turbo_vae", + } + + assert { + component: MODEL_INDEX[component][0] + for component in expected_libraries + } == expected_libraries + + +def test_source_repository_remains_converter_provenance() -> None: + """Keep the official checkpoint ID as converter provenance metadata.""" + assert SOURCE_REPOSITORY == "sand-ai/MAGI-2-preview" + assert SOURCE_REPOSITORY not in get_registered_model_paths() + + +def test_convert_checkpoint_layout_missing_indexed_shard(tmp_path: Path) -> None: + """Reject a source snapshot when a weight index references a missing shard.""" + source = tmp_path / "official" + indexed_shard = _create_indexed_source(source) + indexed_shard.unlink() + output = tmp_path / "fastvideo" + + with pytest.raises(FileNotFoundError, match="missing indexed checkpoint shards"): + convert_checkpoint_layout(source, output) + assert not output.exists() diff --git a/tests/local_tests/magi2/test_magi2_runtime_controls.py b/tests/local_tests/magi2/test_magi2_runtime_controls.py new file mode 100644 index 000000000..cb9172fe3 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_runtime_controls.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Public runtime-control coverage for MAGI-2 inference.""" + +from __future__ import annotations + +import os +import random +from typing import cast + +import numpy as np +import pytest +import torch + +from fastvideo import envs +from fastvideo.api.compat import ( + generator_config_to_fastvideo_args, + legacy_from_pretrained_to_config, +) +from fastvideo.api.schema import EngineConfig, GeneratorConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.basic.magi2.magi2_pipeline import ( + _configure_deterministic_kernels, +) +from fastvideo.pipelines.basic.magi2.stages import output as magi2_output_stage +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.utils import FlexibleArgumentParser + + +def test_deterministic_cli_flag_is_enabled_without_a_value() -> None: + """Expose ``--deterministic`` through the legacy FastVideo CLI.""" + parser = FlexibleArgumentParser() + FastVideoArgs.add_cli_args(parser) + arguments = parser.parse_args( + ["--model-path", "/models/magi2", "--deterministic"] + ) + assert arguments.deterministic is True + + +def test_deterministic_control_round_trips_through_typed_config(monkeypatch) -> None: + """Preserve deterministic mode across legacy and typed engine adapters.""" + typed_config = legacy_from_pretrained_to_config( + "/models/magi2", + {"deterministic": True}, + ) + assert typed_config.engine.deterministic is True + + captured_kwargs: dict = {} + + def capture_fastvideo_args(**kwargs): + captured_kwargs.update(kwargs) + return kwargs + + monkeypatch.setattr(FastVideoArgs, "from_kwargs", capture_fastvideo_args) + generator_config_to_fastvideo_args( + GeneratorConfig( + model_path="/models/magi2", + engine=EngineConfig(deterministic=True), + ) + ) + assert captured_kwargs["deterministic"] is True + + +def test_magi2_environment_controls_are_lazy(monkeypatch) -> None: + """Read deterministic and latent-capture values from each worker's environment.""" + monkeypatch.setenv("MAGI2_DETERMINISTIC", "1") + monkeypatch.setenv("MAGI_ATTENTION_DETERMINISTIC_MODE", "1") + monkeypatch.setenv("MAGI2_SAVE_LATENT_PATH", "/tmp/magi2-latents") + + assert envs.MAGI2_DETERMINISTIC is True + assert envs.MAGI_ATTENTION_DETERMINISTIC_MODE is True + assert envs.MAGI2_SAVE_LATENT_PATH == "/tmp/magi2-latents" + + +def test_deterministic_kernel_configuration_repeats_cpu_rng_sequences( + monkeypatch, +) -> None: + """Seed Python, NumPy, and CPU PyTorch with repeatable sequences.""" + python_rng_state = random.getstate() + numpy_rng_state = np.random.get_state() + torch_rng_state = torch.get_rng_state() + deterministic_algorithms = torch.are_deterministic_algorithms_enabled() + deterministic_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + cuda_seed_calls: list[int] = [] + monkeypatch.setattr(torch.cuda, "manual_seed_all", cuda_seed_calls.append) + monkeypatch.delenv("CUBLAS_WORKSPACE_CONFIG", raising=False) + monkeypatch.delenv("MAGI2_DETERMINISTIC", raising=False) + monkeypatch.delenv("MAGI_ATTENTION_DETERMINISTIC_MODE", raising=False) + + try: + _configure_deterministic_kernels(seed=1234) + first_sequence = (random.random(), np.random.random(4), torch.rand(4)) + _configure_deterministic_kernels(seed=1234) + second_sequence = (random.random(), np.random.random(4), torch.rand(4)) + + assert first_sequence[0] == second_sequence[0] + assert np.array_equal(first_sequence[1], second_sequence[1]) + assert torch.equal(first_sequence[2], second_sequence[2]) + assert cuda_seed_calls and set(cuda_seed_calls) == {1234} + assert torch.are_deterministic_algorithms_enabled() + assert envs.MAGI2_DETERMINISTIC is True + assert envs.MAGI_ATTENTION_DETERMINISTIC_MODE is True + assert os.environ["CUBLAS_WORKSPACE_CONFIG"] == ":4096:8" + finally: + random.setstate(python_rng_state) + np.random.set_state(numpy_rng_state) + torch.set_rng_state(torch_rng_state) + torch.use_deterministic_algorithms( + deterministic_algorithms, + warn_only=deterministic_warn_only, + ) + + +def test_latent_saving_stage_writes_leader_latent(tmp_path, monkeypatch) -> None: + """Save the post-refiner latent on the context-parallel leader rank.""" + latent_directory = tmp_path / "latents" + monkeypatch.setenv("MAGI2_SAVE_LATENT_PATH", str(latent_directory)) + monkeypatch.setattr( + magi2_output_stage.psm, + "is_group_first_rank", + lambda dimension: dimension == "cp", + ) + stage = magi2_output_stage.Magi2LatentSavingStage() + latent = torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) + batch = ForwardBatch(data_type="video", latents=latent) + + returned_batch = stage.forward(batch, cast(FastVideoArgs, object())) + + assert returned_batch is batch + assert torch.equal( + torch.load(latent_directory / "latent_0.pt", weights_only=True), + latent, + ) + assert stage.sample_index == 1 + + +@pytest.mark.parametrize( + ("configured_directory", "is_leader"), + [("", True), ("nonleader-latents", False)], + ids=["empty-path", "non-leader"], +) +def test_latent_saving_stage_skips_disabled_ranks( + tmp_path, + monkeypatch, + configured_directory: str, + is_leader: bool, +) -> None: + """Skip latent writes for an empty path and for non-leader ranks.""" + monkeypatch.chdir(tmp_path) + latent_directory = "" if configured_directory == "" else str(tmp_path / configured_directory) + monkeypatch.setenv("MAGI2_SAVE_LATENT_PATH", latent_directory) + monkeypatch.setattr( + magi2_output_stage.psm, + "is_group_first_rank", + lambda dimension: is_leader, + ) + stage = magi2_output_stage.Magi2LatentSavingStage() + batch = ForwardBatch(data_type="video", latents=torch.ones(1)) + + returned_batch = stage.forward(batch, cast(FastVideoArgs, object())) + + assert returned_batch is batch + assert list(tmp_path.rglob("latent_*.pt")) == [] + assert stage.sample_index == 0 diff --git a/tests/local_tests/magi2/test_magi2_scheduler_parity.py b/tests/local_tests/magi2/test_magi2_scheduler_parity.py new file mode 100644 index 000000000..bf258a877 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_scheduler_parity.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict MAGI-2 Flow UniPC scheduler parity against the official source.""" + +from __future__ import annotations + +import ast +import math +import typing +from pathlib import Path + +import numpy as np +import pytest +import torch +from torch.testing import assert_close + +try: + from diffusers.configuration_utils import ConfigMixin, register_to_config + from diffusers.schedulers.scheduling_utils import ( + KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput, + ) + from diffusers.utils import deprecate + from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler as FastVideoFlowUniPCMultistepScheduler, + ) +except ModuleNotFoundError as exc: + if exc.name != "diffusers": + raise + pytest.skip( + "MAGI-2 scheduler parity requires the Diffusers package.", + allow_module_level=True, + ) + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_SCHEDULER_PATH = ( + REPO_ROOT.parent + / "MAGI-2-preview" + / "inference" + / "pipeline" + / "sampler.py" +) +OFFICIAL_SCHEDULER_CLASS_NAME = "FlowUniPCMultistepScheduler" + + +def _load_official_scheduler_class() -> type: + """Load only the official scheduler class without model-runtime imports.""" + if not OFFICIAL_SCHEDULER_PATH.is_file(): + raise FileNotFoundError( + f"Official MAGI-2 scheduler source is missing: {OFFICIAL_SCHEDULER_PATH}" + ) + + source = OFFICIAL_SCHEDULER_PATH.read_text(encoding="utf-8") + source_module = ast.parse(source, filename=str(OFFICIAL_SCHEDULER_PATH)) + scheduler_nodes = [ + node + for node in source_module.body + if isinstance(node, ast.ClassDef) + and node.name == OFFICIAL_SCHEDULER_CLASS_NAME + ] + if len(scheduler_nodes) != 1: + raise AssertionError( + f"Expected one {OFFICIAL_SCHEDULER_CLASS_NAME} definition in " + f"{OFFICIAL_SCHEDULER_PATH}, found {len(scheduler_nodes)}" + ) + + scheduler_namespace = { + "__name__": "magi2_official_scheduler", + "Any": typing.Any, + "ConfigMixin": ConfigMixin, + "KarrasDiffusionSchedulers": KarrasDiffusionSchedulers, + "List": list, + "Optional": typing.Optional, + "SchedulerMixin": SchedulerMixin, + "SchedulerOutput": SchedulerOutput, + "Tuple": tuple, + "Union": typing.Union, + "deprecate": deprecate, + "math": math, + "np": np, + "register_to_config": register_to_config, + "torch": torch, + } + scheduler_module = ast.Module(body=scheduler_nodes, type_ignores=[]) + exec( + compile(scheduler_module, str(OFFICIAL_SCHEDULER_PATH), "exec"), + scheduler_namespace, + ) + return scheduler_namespace[OFFICIAL_SCHEDULER_CLASS_NAME] + + +OfficialFlowUniPCMultistepScheduler = _load_official_scheduler_class() + + +def _build_scheduler_pair( + num_inference_steps: int, + shift: float, +) -> tuple[SchedulerMixin, SchedulerMixin]: + """Construct official and FastVideo schedulers with a shipping schedule.""" + official_scheduler = OfficialFlowUniPCMultistepScheduler() + fastvideo_scheduler = FastVideoFlowUniPCMultistepScheduler() + official_scheduler.set_timesteps(num_inference_steps, device="cpu", shift=shift) + fastvideo_scheduler.set_timesteps(num_inference_steps, device="cpu", shift=shift) + return official_scheduler, fastvideo_scheduler + + +def _assert_scheduler_tensor_exact( + fastvideo_tensor: torch.Tensor, + official_tensor: torch.Tensor, +) -> None: + """Require identical scheduler tensor metadata and values.""" + assert fastvideo_tensor.shape == official_tensor.shape + assert fastvideo_tensor.dtype == official_tensor.dtype + assert fastvideo_tensor.stride() == official_tensor.stride() + assert_close(fastvideo_tensor, official_tensor, atol=0, rtol=0) + + +@pytest.mark.parametrize( + ("num_inference_steps", "shift"), + [(100, 7.0), (5, 5.0)], + ids=("preview", "refiner"), +) +def test_set_timesteps_magi2_shipping_schedules_match_official( + num_inference_steps: int, + shift: float, +) -> None: + """Match the preview and refiner timestep and sigma schedules exactly.""" + official_scheduler, fastvideo_scheduler = _build_scheduler_pair( + num_inference_steps, + shift, + ) + + _assert_scheduler_tensor_exact( + fastvideo_scheduler.timesteps, + official_scheduler.timesteps, + ) + _assert_scheduler_tensor_exact( + fastvideo_scheduler.sigmas, + official_scheduler.sigmas, + ) + + +@pytest.mark.parametrize( + ("num_inference_steps", "shift"), + [(100, 7.0), (5, 5.0)], + ids=("preview", "refiner"), +) +def test_step_magi2_shipping_denoise_trajectory_matches_official( + num_inference_steps: int, + shift: float, +) -> None: + """Match every scheduler output across each shipping denoise schedule.""" + official_scheduler, fastvideo_scheduler = _build_scheduler_pair( + num_inference_steps, + shift, + ) + generator = torch.Generator(device="cpu").manual_seed(20260805) + tensor_shape = (1, 2, 3, 4, 4) + official_sample = torch.randn( + tensor_shape, + generator=generator, + dtype=torch.float32, + ) + fastvideo_sample = official_sample.clone() + model_outputs = [ + torch.randn(tensor_shape, generator=generator, dtype=torch.float32) + for _ in range(num_inference_steps) + ] + + for step_index, model_output in enumerate(model_outputs): + official_sample = official_scheduler.step( + model_output, + official_scheduler.timesteps[step_index], + official_sample, + return_dict=False, + )[0] + fastvideo_sample = fastvideo_scheduler.step( + model_output, + fastvideo_scheduler.timesteps[step_index], + fastvideo_sample, + return_dict=False, + )[0] + _assert_scheduler_tensor_exact(fastvideo_sample, official_sample) diff --git a/tests/local_tests/magi2/test_magi2_text_encoder_parity.py b/tests/local_tests/magi2/test_magi2_text_encoder_parity.py new file mode 100644 index 000000000..47ed7e399 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_text_encoder_parity.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict prompt and Qwen3.5 encoder parity for MAGI-2 Preview. + +Coverage scope: both. The official side uses +``inference.pipeline.inference_engine.initialize_text_encoder`` with the +published Qwen3.5 checkpoint. The FastVideo side targets +``fastvideo.models.encoders.qwen3_5.Magi2Qwen35TextEncoder`` with the same checkpoint. +The comparison covers structured-prompt normalization, CJK token splitting, +token IDs, attention masks, the skip-layer-2 hidden state, and ``
`` +token pooling. +""" +from __future__ import annotations + +import gc +import json +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.models.encoders.qwen3_5 import ( + Magi2Qwen35TextEncoder, + json_to_compact_markdown, +) +from tests.local_tests.magi2._parity_utils import ( + LOCAL_WEIGHTS_DIR, + OFFICIAL_REF_DIR, + assert_tensor_exact, + import_official_module, + require_complete_safetensor_index, +) + + +PARITY_COVERAGE = "both" +TEXT_ENCODER_DIR = LOCAL_WEIGHTS_DIR / "text_encoder" + + +def _structured_i2v_prompt() -> str: + """Return one deterministic prompt that exercises JSON, CJK, and figure tokens.""" + prompt = { + "global_layer": { + "context": "夜晚的城市街道", + "description": "A cyclist passes a quiet café.", + "aesthetics": { + "style": "documentary", + "mood_atmosphere": "calm", + "color_scheme": "blue and amber", + }, + }, + "reference_layer": ["The first frame refers to
"], + } + return json.dumps(prompt, ensure_ascii=False) + + +def _tokenize( + tokenizer, + max_length: int, + normalized_prompt: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Tokenize with the exact keyword arguments used by the official encoder.""" + tokens = tokenizer( + [normalized_prompt], + return_tensors="pt", + padding="longest", + max_length=max_length, + truncation=True, + ) + return tokens["input_ids"], tokens["attention_mask"] + + +def _load_official_encoder(model_path: str, device: torch.device): + """Instantiate Qwen3.5 through the official component implementation.""" + qwen35_module = import_official_module("inference.model.qwen35") + return qwen35_module.Qwen35TextEncoder( + model_path=model_path, + device=str(device), + precision=torch.bfloat16, + skip_layer=2, + ) + + +def test_iter_indexed_language_model_weights_reads_only_language_model_tensors( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Read only indexed language-model keys and ignore unrelated missing shards.""" + from fastvideo.models.encoders import qwen3_5 as qwen35_module + + language_shards = ["language-a.safetensors", "language-b.safetensors"] + for shard_name in language_shards: + (tmp_path / shard_name).touch() + weight_map = { + "model.language_model.layers.0.weight": language_shards[1], + "model.visual.weight": "missing-visual.safetensors", + "model.language_model.embed_tokens.weight": language_shards[0], + "lm_head.weight": "missing-output.safetensors", + } + tensors_by_shard = { + language_shards[0]: { + "model.language_model.embed_tokens.weight": torch.tensor([1.0]), + }, + language_shards[1]: { + "model.language_model.layers.0.weight": torch.tensor([2.0]), + }, + } + opened_shards: list[str] = [] + read_requests: list[tuple[str, str]] = [] + + @contextmanager + def recording_safe_open(filename: str, *, framework: str, device: str): + """Record each shard and reject tensor reads from the wrong shard.""" + shard_name = Path(filename).name + opened_shards.append(shard_name) + assert framework == "pt" + assert device == "cpu" + + def read_tensor(checkpoint_name: str) -> torch.Tensor: + read_requests.append((shard_name, checkpoint_name)) + return tensors_by_shard[shard_name][checkpoint_name] + + yield SimpleNamespace(get_tensor=read_tensor) + + monkeypatch.setattr(qwen35_module, "safe_open", recording_safe_open) + loaded_weights = list( + qwen35_module._iter_indexed_language_model_weights( + tmp_path, + weight_map, + torch.device("cpu"), + ) + ) + + assert opened_shards == language_shards + assert read_requests == [ + (language_shards[0], "model.language_model.embed_tokens.weight"), + (language_shards[1], "model.language_model.layers.0.weight"), + ] + assert [name for name, _ in loaded_weights] == [name for _, name in read_requests] + assert [tensor.item() for _, tensor in loaded_weights] == [1.0, 2.0] + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MAGI-2 Qwen3.5 parity requires CUDA.", +) +def test_magi2_text_encoder_structured_i2v_prompt_exact_parity() -> None: + """Require exact tokenizer, hidden-state, and figure-token outputs.""" + model_dir = require_complete_safetensor_index(TEXT_ENCODER_DIR) + device = torch.device("cuda:0") + prompt = _structured_i2v_prompt() + + official_encoder = _load_official_encoder( + str(model_dir), + torch.device("cpu"), + ).to(device) + official_normalized_prompt = official_encoder._normalize_prompt(prompt) + official_input_ids, official_attention_mask = _tokenize( + official_encoder.tokenizer, + official_encoder.max_length, + official_normalized_prompt, + ) + official_embedding = official_encoder.encode(prompt).detach().cpu() + release_prompt = ( + OFFICIAL_REF_DIR / "assets" / "sample_enhanced_t2v.json" + ).read_text(encoding="utf-8").strip() + official_release_embedding = official_encoder.encode(release_prompt).detach().cpu() + official_figure_embedding = official_encoder.get_special_token( + prompt, + ["
"], + official_embedding, + ) + + del official_encoder + gc.collect() + torch.cuda.empty_cache() + + from fastvideo.configs.models.encoders.qwen3_5 import Magi2Qwen35Config + fastvideo_encoder = Magi2Qwen35TextEncoder.from_pretrained_local( + model_path=str(model_dir), + model_config=Magi2Qwen35Config(), + dtype=torch.bfloat16, + device=torch.device("cpu"), + ).to(device) + fastvideo_normalized_prompt = json_to_compact_markdown(prompt) + fastvideo_input_ids, fastvideo_attention_mask = _tokenize( + fastvideo_encoder.tokenizer, + fastvideo_encoder.config.text_len, + fastvideo_normalized_prompt, + ) + fastvideo_embedding = fastvideo_encoder.encode(prompt).detach().cpu() + fastvideo_release_embedding = fastvideo_encoder.encode( + release_prompt + ).detach().cpu() + fastvideo_figure_embedding = fastvideo_encoder.get_special_token( + prompt, + ["
"], + fastvideo_embedding, + ) + + assert fastvideo_normalized_prompt == official_normalized_prompt + assert_tensor_exact(fastvideo_input_ids, official_input_ids, "token IDs") + assert_tensor_exact( + fastvideo_attention_mask, + official_attention_mask, + "attention mask", + ) + assert_tensor_exact(fastvideo_embedding, official_embedding, "skip-layer-2 hidden state") + assert_tensor_exact( + fastvideo_release_embedding, + official_release_embedding, + "release-prompt skip-layer-2 hidden state", + ) + assert_tensor_exact( + fastvideo_figure_embedding, + official_figure_embedding, + "
pooled embedding", + ) diff --git a/tests/local_tests/magi2/test_magi2_transformer_structure_parity.py b/tests/local_tests/magi2/test_magi2_transformer_structure_parity.py new file mode 100644 index 000000000..6b054f2fe --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_transformer_structure_parity.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict MAGI-2 preview and refiner state-structure parity.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_ROOT = Path(os.environ.get("MAGI2_OFFICIAL_ROOT", REPO_ROOT.parent / "MAGI-2-preview")) +WEIGHTS_ROOT = Path(os.environ.get("MAGI2_WEIGHTS_ROOT", REPO_ROOT / "official_weights" / "magi2")) + +if not (OFFICIAL_ROOT / "inference" / "model" / "magi2_preview.py").is_file(): + pytest.skip(f"Official MAGI-2 checkout is missing: {OFFICIAL_ROOT}", allow_module_level=True) +if not (WEIGHTS_ROOT / "preview" / "model.safetensors.index.json").is_file(): + pytest.skip(f"Official MAGI-2 weights are missing: {WEIGHTS_ROOT}", allow_module_level=True) + + +def _structure_dump_script(component: str, implementation: str) -> str: + """Create an isolated process script that instantiates one model on meta.""" + if implementation == "official": + model_setup = { + "preview": """ + from inference.common.magi2_config import load_config + from inference.model.magi2_preview import Transformer + config = load_config(str(official_root / "configs" / "magi2_preview.json")) + with torch.device("meta"): + model = Transformer(config.arch_config, ep_size=1) + """, + "refiner": """ + from inference.common.magi2_config import load_config + from inference.model.magi2_refiner import Transformer + config = load_config(str(official_root / "configs" / "magi2_refiner.json")) + with torch.device("meta"): + model = Transformer(config.magi2_refiner_arch_config) + """, + }[component] + path_setup = "sys.path.insert(0, str(official_root))" + else: + model_setup = { + "preview": """ + from fastvideo.configs.models.dits.magi2 import Magi2PreviewVideoConfig + from fastvideo.models.dits.magi2 import Magi2PreviewDiT + with torch.device("meta"): + model = Magi2PreviewDiT(config=Magi2PreviewVideoConfig()) + """, + "refiner": """ + from fastvideo.configs.models.dits.magi2 import Magi2RefinerVideoConfig + from fastvideo.models.dits.magi2_refiner import Magi2RefinerDiT + with torch.device("meta"): + model = Magi2RefinerDiT(config=Magi2RefinerVideoConfig()) + """, + }[component] + path_setup = "sys.path.insert(0, str(repo_root))" + + setup_script = textwrap.dedent( + f""" + import json + import os + import pathlib + import sys + import torch + + repo_root = pathlib.Path({str(REPO_ROOT)!r}) + official_root = pathlib.Path({str(OFFICIAL_ROOT)!r}) + {path_setup} + os.environ["SKIP_LOAD_MODEL"] = "1" + os.environ["MAGI2_DISABLE_MAGI_COMPILE"] = "1" + os.environ["MAGI_COMPILE_COMPILE_MODE"] = "NONE" + """ + ) + dump_script = textwrap.dedent( + """ + structure = { + name: {"shape": list(tensor.shape), "dtype": str(tensor.dtype)} + for name, tensor in model.state_dict().items() + } + print("MAGI2_STRUCTURE=" + json.dumps(structure, sort_keys=True)) + """ + ) + return setup_script + textwrap.dedent(model_setup) + dump_script + + +def _dump_structure(component: str, implementation: str) -> dict[str, dict]: + """Run one model definition in isolation and parse its state metadata.""" + completed = subprocess.run( + [sys.executable, "-c", _structure_dump_script(component, implementation)], + cwd=REPO_ROOT, + check=False, + capture_output=True, + text=True, + timeout=300, + ) + if completed.returncode != 0: + raise AssertionError( + f"{implementation} {component} structure dump failed.\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + lines = [line for line in completed.stdout.splitlines() if line.startswith("MAGI2_STRUCTURE=")] + if len(lines) != 1: + raise AssertionError(f"Expected one structure record, received stdout:\n{completed.stdout}") + return json.loads(lines[0].removeprefix("MAGI2_STRUCTURE=")) + + +@pytest.mark.parametrize("component", ["preview", "refiner"]) +def test_magi2_transformer_state_structure_matches_official(component: str) -> None: + """Match every state name, shape, and dtype to the official definition.""" + official_structure = _dump_structure(component, "official") + fastvideo_structure = _dump_structure(component, "fastvideo") + assert fastvideo_structure == official_structure + + index_path = WEIGHTS_ROOT / component / "model.safetensors.index.json" + checkpoint_keys = set(json.loads(index_path.read_text(encoding="utf-8"))["weight_map"]) + assert set(fastvideo_structure) == checkpoint_keys diff --git a/tests/local_tests/magi2/test_magi2_turbo_vae_parity.py b/tests/local_tests/magi2/test_magi2_turbo_vae_parity.py new file mode 100644 index 000000000..c04910734 --- /dev/null +++ b/tests/local_tests/magi2/test_magi2_turbo_vae_parity.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Strict temporal sliding-window parity for the MAGI-2 Turbo VAE decoder. + +Coverage scope: both. The official side uses +``inference.model.turbo_vaed.get_turbo_vaed`` with the published distilled +decoder. The FastVideo side targets +``fastvideo.models.vaes.magi2_turbo_vae.Magi2TurboVAEModel``. Fifteen latent +frames exercise the published first, middle, and last temporal windows after +padding from 15 to 21 frames. +""" + +from __future__ import annotations + +import json + +import pytest +import torch + +from fastvideo.models.vaes.magi2_turbo_vae import Magi2TurboVAEModel +from tests.local_tests.magi2._parity_utils import ( + LOCAL_WEIGHTS_DIR, + assert_tensor_exact, + import_official_module, + require_path, +) + + +PARITY_COVERAGE = "both" +TURBO_VAE_DIR = LOCAL_WEIGHTS_DIR / "turbo_vae" +TURBO_VAE_CONFIG_PATH = TURBO_VAE_DIR / "TurboV3-Wan22-TinyShallow_7_7.json" +TURBO_VAE_CHECKPOINT_PATH = TURBO_VAE_DIR / "checkpoint.ckpt" + + +def _load_turbo_config() -> dict: + """Load the exact configuration paired with the distilled checkpoint.""" + config_path = require_path(TURBO_VAE_CONFIG_PATH, "Turbo VAE configuration") + with config_path.open(encoding="utf-8") as config_file: + return json.load(config_file) + + +def _deterministic_video_latent(device: torch.device) -> torch.Tensor: + """Create a reproducible BF16 latent that crosses two window boundaries.""" + latent_values = torch.arange(1 * 48 * 15 * 4 * 4, dtype=torch.float32) + latent_values = ((latent_values % 257) - 128) / 128 + return latent_values.reshape(1, 48, 15, 4, 4).to( + device=device, + dtype=torch.bfloat16, + ) + + +def _temporal_window_roles( + num_frames: int, + first_chunk_size: int, + step_size: int, +) -> tuple[str, ...]: + """Return the Turbo VAE window roles selected after temporal padding.""" + padding_frames = 0 + if num_frames < first_chunk_size: + padding_frames = first_chunk_size - num_frames + elif (num_frames - first_chunk_size) % step_size != 0: + padding_frames = step_size - (num_frames - first_chunk_size) % step_size + padded_frames = num_frames + padding_frames + if padded_frames == first_chunk_size: + return ("single",) + + window_roles = ["first"] + for window_start in range(first_chunk_size, padded_frames, step_size): + is_last_window = window_start + step_size == padded_frames + window_roles.append("last" if is_last_window else "middle") + return tuple(window_roles) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MAGI-2 Turbo VAE parity requires CUDA.", +) +def test_magi2_turbo_vae_sliding_window_decode_exact_parity() -> None: + """Require exact decoded video values across first, middle, and last windows.""" + checkpoint_path = require_path( + TURBO_VAE_CHECKPOINT_PATH, + "Turbo VAE checkpoint", + ) + config = _load_turbo_config() + assert config["first_chunk_size"] == 7 + assert config["step_size"] == 7 + assert config["temporal_compression_ratio"] == 4 + + official_module = import_official_module("inference.model.turbo_vaed") + device = torch.device("cuda:0") + official_vae = official_module.get_turbo_vaed( + str(TURBO_VAE_CONFIG_PATH), + str(checkpoint_path), + device=str(device), + weight_dtype=torch.bfloat16, + ) + from fastvideo.configs.models.vaes.magi2_turbo_vae import Magi2TurboVAEConfig + + fastvideo_config = Magi2TurboVAEConfig( + config_path=str(TURBO_VAE_CONFIG_PATH), + checkpoint_path=str(checkpoint_path), + pretrained_dtype="bfloat16", + ) + fastvideo_vae = Magi2TurboVAEModel(fastvideo_config).eval() + latent = _deterministic_video_latent(device) + assert _temporal_window_roles( + latent.shape[2], + config["first_chunk_size"], + config["step_size"], + ) == ("first", "middle", "last") + + with torch.inference_mode(): + official_video = official_vae.decode(latent).float().detach().cpu() + fastvideo_video = fastvideo_vae.decode(latent).float().detach().cpu() + + assert_tensor_exact(fastvideo_video, official_video, "Turbo VAE decoded video")