Hongsheng Liu
d4096e01fe
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
# Copyright (c) 2026 SandAI. All Rights Reserved.
|
|
|
|
"""Native configuration for the released MAGI-2 Preview checkpoint.
|
|
|
|
The defaults mirror ``sand-ai/MAGI-2-preview``. Keeping the architecture in
|
|
vLLM-Omni makes model construction independent of SandAI's Python package and
|
|
also gives tests a small-config entry point.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
import torch
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Magi2MHCConfig:
|
|
enabled: bool = True
|
|
num_streams: int = 4
|
|
alpha_init: float = 0.01
|
|
sinkhorn_iterations: int = 20
|
|
sinkhorn_epsilon: float = 1e-12
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Magi2MoEConfig:
|
|
num_heads: int = 12
|
|
num_experts: int = 256
|
|
top_k: int = 6
|
|
expert_intermediate_size: int = 1280
|
|
shared_expert_intermediate_size: int = 1280
|
|
modality_shared_expert_intermediate_size: int = 1280
|
|
score_function: str = "sigmoid"
|
|
normalize_routing_weights: bool = True
|
|
routing_scale: float = 4.9
|
|
layers: tuple[int, ...] = tuple(range(2, 38))
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.num_heads <= 0 or self.num_experts <= 0:
|
|
raise ValueError("MAGI-2 MoE head and expert counts must be positive")
|
|
if not 0 < self.top_k <= self.num_experts:
|
|
raise ValueError("MAGI-2 MoE top_k must be in [1, num_experts]")
|
|
if self.score_function != "sigmoid":
|
|
raise ValueError("The released MAGI-2 checkpoint uses sigmoid routing")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Magi2PreviewConfig:
|
|
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
|
|
intermediate_factor: float = 4.0
|
|
multimodal_layers: tuple[int, ...] = (0, 1, 38, 39)
|
|
params_dtype: torch.dtype = torch.bfloat16
|
|
attention_softcap: float = -1.0
|
|
attention_sink_tokens: int = 1
|
|
attention_gating: bool = True
|
|
mhc: Magi2MHCConfig = field(default_factory=Magi2MHCConfig)
|
|
moe: Magi2MoEConfig = field(default_factory=Magi2MoEConfig)
|
|
|
|
@property
|
|
def num_attention_heads(self) -> int:
|
|
if self.hidden_size % self.head_dim:
|
|
raise ValueError("hidden_size must be divisible by head_dim")
|
|
return self.hidden_size // self.head_dim
|
|
|
|
@property
|
|
def num_heads_q(self) -> int:
|
|
return self.num_attention_heads
|
|
|
|
@property
|
|
def num_heads_kv(self) -> int:
|
|
return self.num_query_groups
|
|
|
|
@property
|
|
def virtual_width(self) -> int:
|
|
return self.hidden_size * self.mhc.num_streams
|
|
|
|
def validate(self) -> None:
|
|
if not self.mhc.enabled:
|
|
raise ValueError("The released MAGI-2 Preview checkpoint requires mHC")
|
|
if not self.attention_gating:
|
|
raise ValueError("The released MAGI-2 Preview checkpoint requires attention gating")
|
|
if self.attention_sink_tokens < 1:
|
|
raise ValueError("The released MAGI-2 Preview checkpoint requires an attention sink")
|
|
if self.hidden_size % self.moe.num_heads:
|
|
raise ValueError("hidden_size must be divisible by the number of MoE heads")
|
|
layer_ids = set(range(self.num_layers))
|
|
if not set(self.multimodal_layers).issubset(layer_ids):
|
|
raise ValueError("multimodal layer index is outside the transformer")
|
|
if not set(self.moe.layers).issubset(layer_ids):
|
|
raise ValueError("MoE layer index is outside the transformer")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Magi2GenerationConfig:
|
|
duration_seconds: float = 10.0
|
|
fps: float = 12.5
|
|
output_frames: int = 125
|
|
preview_steps: int = 100
|
|
shift: float = 7.0
|
|
video_guidance_scale: float = 5.0
|
|
audio_guidance_scale: float = 7.0
|
|
audio_latent_fps: float = 25.0
|
|
video_vae_stride: tuple[int, int, int] = (8, 16, 16)
|
|
patch_size: tuple[int, int, int] = (1, 1, 1)
|
|
video_latent_channels: int = 48
|
|
audio_latent_channels: int = 64
|
|
audio_sample_rate: int = 44_100
|
|
|
|
|
|
MAGI2_PREVIEW_CONFIG = Magi2PreviewConfig()
|
|
MAGI2_GENERATION_CONFIG = Magi2GenerationConfig()
|