Skip to content

vllm_omni.diffusion.models.magi2.preview_data_proxy

Native packing helpers for the MAGI-2 Preview transformer.

The released model consumes one varlen sequence per sample. Tokens are laid out in VIDEO -> AUDIO -> TEXT order, followed by zero or more reference-image special-token/image-token pairs. This module keeps that layout and its 9-D coordinate metadata independent of any particular distributed topology.

The packing math is adapted from SandAI's Apache-2.0 MAGI-2 Preview inference implementation. It intentionally has no dependency on that implementation at runtime.

Magi2DataProxy

Convert dense modality tensors to/from the transformer's packed ABI.

The proxy is request-stateless: all packing state needed by process_output is returned as PackedModelInput.output_layout rather than stored on this shared instance, so concurrent requests cannot overwrite one another.

config instance-attribute 

config = config or Magi2PreviewDataProxyConfig()

patch_size instance-attribute 

patch_size = self.config.patch_size

t_patch_size instance-attribute 

t_patch_size = self.config.t_patch_size

img2tokens 

img2tokens(x_t: Tensor) -> Tensor

process_input 

process_input(data: ModelInput) -> PackedModelInput

process_output staticmethod 

process_output(
    x: Tensor, output_layout: SimplePackedData
) -> tuple[Tensor, Tensor]

Magi2PreviewDataProxyConfig dataclass

Packing configuration for the released preview checkpoint.

add_time_token class-attribute instance-attribute 

add_time_token: bool = False

audio_latent_fps class-attribute instance-attribute 

audio_latent_fps: float = 25.0

patch_size class-attribute instance-attribute 

patch_size: int = 1

spatial_rope_interpolation class-attribute instance-attribute 

spatial_rope_interpolation: Literal["inter", "extra"] = (
    "extra"
)

t_patch_size class-attribute instance-attribute 

t_patch_size: int = 1

time_aligned_rope class-attribute instance-attribute 

time_aligned_rope: bool = False

time_channel_dim class-attribute instance-attribute 

time_channel_dim: int = 64

time_pos_fps class-attribute instance-attribute 

time_pos_fps: float = 3.125

vae_first_latent_is_image class-attribute instance-attribute 

vae_first_latent_is_image: bool = True

video_fps class-attribute instance-attribute 

video_fps: float = 25.0

Modality

Bases: IntEnum

Checkpoint modality IDs used by the pre/post adapters.

AUDIO class-attribute instance-attribute 

AUDIO = 1

TEXT class-attribute instance-attribute 

TEXT = 2

TIME class-attribute instance-attribute 

TIME = 3

VIDEO class-attribute instance-attribute 

VIDEO = 0

ModelInput dataclass

Unpacked tensors for one native preview-transformer invocation.

audio_feat_len instance-attribute 

audio_feat_len: Tensor | Sequence[int]

audio_x_t instance-attribute 

audio_x_t: Tensor

per_token_audio_t class-attribute instance-attribute 

per_token_audio_t: Tensor | None = None

per_token_video_t class-attribute instance-attribute 

per_token_video_t: Tensor | None = None

ref_audio_feat class-attribute instance-attribute 

ref_audio_feat: Tensor | None = None

ref_audio_feat_len class-attribute instance-attribute 

ref_audio_feat_len: Tensor | Sequence[int] | None = None

ref_image_feat class-attribute instance-attribute 

ref_image_feat: Tensor | None = None

ref_image_feat_len class-attribute instance-attribute 

ref_image_feat_len: Tensor | None = None

ref_image_special_token_embedding class-attribute instance-attribute 

ref_image_special_token_embedding: Tensor | None = None

ref_video_feat class-attribute instance-attribute 

ref_video_feat: Tensor | None = None

ref_video_feat_len class-attribute instance-attribute 

ref_video_feat_len: Tensor | Sequence[int] | None = None

t instance-attribute 

t: Tensor

txt_feat instance-attribute 

txt_feat: Tensor

txt_feat_len instance-attribute 

txt_feat_len: Tensor | Sequence[int]

x_t instance-attribute 

x_t: Tensor

PackedModelInput dataclass

Request-owned packed transformer arguments and output layout.

coords_mapping instance-attribute 

coords_mapping: Tensor

modality_mapping instance-attribute 

modality_mapping: Tensor

model_args property 

model_args: tuple[
    Tensor, Tensor, Tensor, VarlenHandler, Tensor
]

output_layout instance-attribute 

output_layout: SimplePackedData

time_token_sequence instance-attribute 

time_token_sequence: Tensor

token_sequence instance-attribute 

token_sequence: Tensor

varlen_handler instance-attribute 

varlen_handler: VarlenHandler

SimplePackedData dataclass

A concatenation of independently addressable packed samples.

coords_mapping property 

coords_mapping: Tensor

cu_seqlen property 

cu_seqlen: Tensor

default_dtype property 

default_dtype: dtype

device property 

device: device

items instance-attribute 

items: list[SingleData]

max_seqlen property 

max_seqlen: int

modality_mapping property 

modality_mapping: Tensor

time_token_sequence property 

time_token_sequence: Tensor

token_sequence property 

token_sequence: Tensor

total_token_num property 

total_token_num: int

depack_token_sequence 

depack_token_sequence(
    token_sequence: Tensor,
) -> tuple[Tensor, Tensor]

SingleData dataclass

Packed metadata for one sample.

add_time_token property 

add_time_token: bool

audio_feat_len instance-attribute 

audio_feat_len: int

audio_x_t instance-attribute 

audio_x_t: Tensor

coords_mapping property 

coords_mapping: Tensor

coords_to_cat property 

coords_to_cat: list[Tensor]

default_dtype property 

default_dtype: dtype

device property 

device: device

diffusion_t class-attribute instance-attribute 

diffusion_t: Tensor | None = None

feat_to_cat property 

feat_to_cat: list[Tensor]

h instance-attribute 

h: int

modality_map_seqlens property 

modality_map_seqlens: tuple[list[int], list[int]]

modality_mapping property 

modality_mapping: Tensor

patch_size instance-attribute 

patch_size: int

per_token_audio_t class-attribute instance-attribute 

per_token_audio_t: Tensor | None = None

per_token_video_t class-attribute instance-attribute 

per_token_video_t: Tensor | None = None

ref_image_feat_lens class-attribute instance-attribute 

ref_image_feat_lens: list[list[int]] | None = None

ref_image_feats class-attribute instance-attribute 

ref_image_feats: list[Tensor] | None = None

ref_image_special_tokens class-attribute instance-attribute 

ref_image_special_tokens: list[Tensor] | None = None

spatial_rope_interpolation instance-attribute 

spatial_rope_interpolation: Literal['inter', 'extra']

t instance-attribute 

t: int

t_patch_size instance-attribute 

t_patch_size: int

time_channel_dim class-attribute instance-attribute 

time_channel_dim: int = 0

time_pos_fps class-attribute instance-attribute 

time_pos_fps: float = 3.125

time_token_sequence property 

time_token_sequence: Tensor

token_sequence property 

token_sequence: Tensor

total_token_num property 

total_token_num: int

txt_feat instance-attribute 

txt_feat: Tensor

txt_feat_len instance-attribute 

txt_feat_len: int

vae_first_latent_is_image class-attribute instance-attribute 

vae_first_latent_is_image: bool = True

video_fps class-attribute instance-attribute 

video_fps: float = 25.0

video_x_t instance-attribute 

video_x_t: Tensor

w instance-attribute 

w: int

VarlenHandler dataclass

Packed-sequence metadata consumed by MAGI-2 attention.

cu_seqlens_k instance-attribute 

cu_seqlens_k: Tensor | None

cu_seqlens_q instance-attribute 

cu_seqlens_q: Tensor | None

max_seqlen_k class-attribute instance-attribute 

max_seqlen_k: int | None = None

max_seqlen_q class-attribute instance-attribute 

max_seqlen_q: int | None = None

resolved 

resolved(
    q_tokens: int, k_tokens: int
) -> tuple[Tensor, Tensor, int, int]

get_coords 

get_coords(
    shape: tuple[int, int, int],
    ref_feat_shape: tuple[int, int, int],
    offset_thw: tuple[int, int, int] = (0, 0, 0),
    *,
    device: device = device("cpu"),
    dtype: dtype = float32,
    time_positions: Tensor | None = None,
) -> Tensor

Build MAGI-2's (t,h,w,T,H,W,ref_T,ref_H,ref_W) rows.

sinusoidal_embedding_1d 

sinusoidal_embedding_1d(
    dim: int, position: Tensor
) -> Tensor

The preview checkpoint's FP32 diffusion-time channel embedding.