[Model] Add native MAGI-2 Preview diffusion support (#5918)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

vllm_omni/diffusion/models/magi2/preview_data_proxy.py

@@ -0,0 +1,691 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""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.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import IntEnum
+from itertools import chain
+from typing import Literal
+
+import torch
+from einops import rearrange
+from torch.nn import functional as F
+
+from .attention import VarlenHandler
+
+
+class Modality(IntEnum):
+    """Checkpoint modality IDs used by the pre/post adapters."""
+
+    VIDEO = 0
+    AUDIO = 1
+    TEXT = 2
+    # TIME is remapped to TEXT by the transformer after selecting time tokens.
+    TIME = 3
+
+
+@dataclass(frozen=True)
+class Magi2PreviewDataProxyConfig:
+    """Packing configuration for the released preview checkpoint."""
+
+    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
+
+    def __post_init__(self) -> None:
+        if self.t_patch_size <= 0 or self.patch_size <= 0:
+            raise ValueError("MAGI-2 patch sizes must be positive")
+        if self.time_channel_dim < 0:
+            raise ValueError("time_channel_dim must be non-negative")
+
+
+@dataclass
+class ModelInput:
+    """Unpacked tensors for one native preview-transformer invocation."""
+
+    x_t: torch.Tensor
+    audio_x_t: torch.Tensor
+    audio_feat_len: torch.Tensor | Sequence[int]
+    txt_feat: torch.Tensor
+    txt_feat_len: torch.Tensor | Sequence[int]
+    t: torch.Tensor
+    ref_audio_feat: torch.Tensor | None = None
+    ref_audio_feat_len: torch.Tensor | Sequence[int] | None = None
+    ref_video_feat: torch.Tensor | None = None
+    ref_video_feat_len: torch.Tensor | Sequence[int] | None = None
+    per_token_video_t: torch.Tensor | None = None
+    per_token_audio_t: torch.Tensor | None = None
+    # Images are [B, M, C, 1, H, W], lengths are [B, M, 2] holding
+    # the latent patch-grid (H, W), and special tokens are [B, M, D].
+    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:
+    if isinstance(value, torch.Tensor):
+        return int(value.detach().reshape(-1)[0].item())
+    return int(value)
+
+
+def _length_at(value: torch.Tensor | Sequence[int], index: int) -> int:
+    if isinstance(value, torch.Tensor) and value.ndim == 0:
+        if index:
+            raise IndexError("a scalar length only describes one sample")
+        return _to_int(value)
+    return _to_int(value[index])
+
+
+def _pad_cat(tensors: list[torch.Tensor], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    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:
+    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
+    if offset != output_size:
+        raise ValueError(f"segment lengths total {offset}, expected {output_size}")
+    return mapping
+
+
+def _seqlens_to_cu_seqlens(seqlens: list[int], device: torch.device | None = None) -> torch.Tensor:
+    lengths = torch.tensor(seqlens, dtype=torch.int32, device=device)
+    return F.pad(torch.cumsum(lengths, dim=0), (1, 0))
+
+
+def _len_to_list(value: torch.Tensor) -> list[int]:
+    return [int(item) for item 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 = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+    time_positions: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Build MAGI-2's ``(t,h,w,T,H,W,ref_T,ref_H,ref_W)`` rows."""
+
+    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:
+        if time_positions.numel() != ori_t:
+            raise ValueError(f"received {time_positions.numel()} time positions for {ori_t} time steps")
+        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 = torch.stack((time_grid, height_grid, width_grid), dim=-1).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, metadata.expand(coords.shape[0], -1)), dim=-1)
+
+
+def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
+    """The preview checkpoint's FP32 diffusion-time channel embedding."""
+
+    position = position.to(torch.float32) * 1000.0
+    half = dim // 2
+    frequencies = torch.exp(
+        -math.log(10000) * torch.arange(0, 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
+
+
+def _patchify_3d(x: torch.Tensor, t_patch_size: int, patch_size: int) -> torch.Tensor:
+    """Non-overlapping 3-D unfold with the released channel-major order."""
+
+    kernel = (t_patch_size, patch_size, patch_size)
+    if x.ndim != 5:
+        raise ValueError(f"expected [B,C,T,H,W], got {tuple(x.shape)}")
+    if not all(size >= kernel_size for size, kernel_size in zip(x.shape[2:], kernel, strict=True)):
+        return torch.empty(
+            x.shape[0],
+            0,
+            x.shape[1] * math.prod(kernel),
+            device=x.device,
+            dtype=x.dtype,
+        )
+
+    patches = x.unfold(2, t_patch_size, t_patch_size)
+    patches = patches.unfold(3, patch_size, patch_size)
+    patches = patches.unfold(4, patch_size, patch_size)
+    # [B,C,T,H,W,pT,pH,pW] -> [B,T*H*W,C*pT*pH*pW]
+    patches = patches.permute(0, 2, 3, 4, 1, 5, 6, 7).contiguous()
+    return patches.reshape(x.shape[0], -1, x.shape[1] * math.prod(kernel))
+
+
+@dataclass
+class SingleData:
+    """Packed metadata for one sample."""
+
+    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:
+        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 []
+        if not (len(self.ref_image_feats) == len(self.ref_image_feat_lens) == len(self.ref_image_special_tokens)):
+            raise ValueError("reference-image features, lengths, and special tokens must have equal counts")
+        self.ref_image_token_nums = [math.prod(feat_len) for feat_len in self.ref_image_feat_lens]
+        self.ref_image_feats = [
+            feature[:token_num]
+            for feature, 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 self.video_x_t.device
+
+    @property
+    def default_dtype(self) -> torch.dtype:
+        return self.video_x_t.dtype
+
+    @property
+    def add_time_token(self) -> bool:
+        return self.diffusion_t is not None
+
+    @property
+    def total_token_num(self) -> int:
+        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 + int(self.add_time_token)
+
+    @property
+    def feat_to_cat(self) -> list[torch.Tensor]:
+        tensors = [self.video_x_t, self.audio_x_t, self.txt_feat]
+        for index in range(self.num_ref_images):
+            tensors.append(self.ref_image_special_tokens[index])
+            tensors.append(self.ref_image_feats[index])
+        if self.add_time_token:
+            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 _pad_cat(self.feat_to_cat, self.device, self.default_dtype)
+
+    @property
+    def modality_map_seqlens(self) -> tuple[list[int], list[int]]:
+        seqlens = [self.video_token_num, self.audio_feat_len, self.txt_feat_len]
+        modalities = [Modality.VIDEO, Modality.AUDIO, Modality.TEXT]
+        for index in range(self.num_ref_images):
+            seqlens.extend((1, self.ref_image_token_nums[index]))
+            modalities.extend((Modality.TEXT, Modality.VIDEO))
+        if self.add_time_token:
+            seqlens.append(1)
+            modalities.append(Modality.TIME)
+        return seqlens, modalities
+
+    @property
+    def modality_mapping(self) -> torch.Tensor:
+        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:
+        return get_coords(
+            shape,
+            ref_feat_shape,
+            offset_thw,
+            device=self.device,
+            dtype=self.default_dtype,
+            time_positions=time_positions,
+        )
+
+    @property
+    def coords_to_cat(self) -> list[torch.Tensor]:
+        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 index in range(self.num_ref_images):
+            token_len = self.ref_image_token_nums[index]
+            if len(self.ref_image_feat_lens[index]) >= 2:
+                image_h, image_w = self.ref_image_feat_lens[index][:2]
+            else:
+                image_h = image_w = math.ceil(math.sqrt(token_len))
+            # Keep image context after the generated clip with a one-step gap.
+            time_offset = t_steps + 2 + 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 torch.cat(self.coords_to_cat, dim=0)
+
+    @property
+    def time_token_sequence(self) -> torch.Tensor:
+        if self.time_channel_dim == 0:
+            return torch.empty(self.total_token_num, 0, device=self.device)
+        if self.per_token_video_t is None or self.per_token_audio_t is None:
+            raise ValueError("per-token video and audio timesteps are required when time_channel_dim is non-zero")
+        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 index in range(self.num_ref_images):
+            parts.append(torch.zeros(1, device=self.device))
+            parts.append(torch.zeros(self.ref_image_token_nums[index], device=self.device))
+        if self.add_time_token:
+            parts.append(self.diffusion_t.reshape(1).to(self.device))
+        raw_t = torch.cat(parts, dim=0)
+        if self.time_channel_dim == 1:
+            return raw_t.unsqueeze(-1)
+        return sinusoidal_embedding_1d(self.time_channel_dim, raw_t)
+
+
+@dataclass
+class SimplePackedData:
+    """A concatenation of independently addressable packed samples."""
+
+    items: list[SingleData]
+
+    def __post_init__(self) -> None:
+        if not self.items:
+            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 = _seqlens_to_cu_seqlens(self._total_token_num_list)
+
+    @property
+    def device(self) -> torch.device:
+        return self.items[0].device
+
+    @property
+    def default_dtype(self) -> torch.dtype:
+        return self.items[0].default_dtype
+
+    @property
+    def token_sequence(self) -> torch.Tensor:
+        tensors = list(chain.from_iterable(item.feat_to_cat for item in self.items))
+        return _pad_cat(tensors, self.device, self.default_dtype)
+
+    @property
+    def modality_mapping(self) -> torch.Tensor:
+        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:
+        coords = list(chain.from_iterable(item.coords_to_cat for item in self.items))
+        return torch.cat(coords, dim=0)
+
+    @property
+    def time_token_sequence(self) -> torch.Tensor:
+        return torch.cat([item.time_token_sequence for item in self.items], dim=0)
+
+    @property
+    def total_token_num(self) -> int:
+        return self._total_token_num_sum
+
+    @property
+    def cu_seqlen(self) -> torch.Tensor:
+        return self._total_token_cu_seqlens.clone()
+
+    @property
+    def max_seqlen(self) -> int:
+        return self._total_token_num_max
+
+    def __getitem__(self, index: int) -> SingleData:
+        return self.items[index]
+
+    def depack_token_sequence(self, token_sequence: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if token_sequence.shape[0] != self.total_token_num:
+            raise ValueError(f"model returned {token_sequence.shape[0]} tokens, expected {self.total_token_num}")
+        videos: list[torch.Tensor] = []
+        audios: list[torch.Tensor] = []
+        token_slices = torch.split(token_sequence, self._total_token_num_list, dim=0)
+        for item, token_slice in zip(self.items, token_slices, strict=True):
+            video_flat = token_slice[: item.video_token_num, : item.video_channel]
+            output_channels = 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_channels,
+            ).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)
+
+
+@dataclass(frozen=True)
+class PackedModelInput:
+    """Request-owned packed transformer arguments and output layout."""
+
+    token_sequence: torch.Tensor
+    coords_mapping: torch.Tensor
+    modality_mapping: torch.Tensor
+    varlen_handler: VarlenHandler
+    time_token_sequence: torch.Tensor
+    output_layout: SimplePackedData
+
+    @property
+    def model_args(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, VarlenHandler, torch.Tensor]:
+        return (
+            self.token_sequence,
+            self.coords_mapping,
+            self.modality_mapping,
+            self.varlen_handler,
+            self.time_token_sequence,
+        )
+
+
+class 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.
+    """
+
+    def __init__(self, config: Magi2PreviewDataProxyConfig | None = None) -> None:
+        self.config = config or Magi2PreviewDataProxyConfig()
+        self.patch_size = self.config.patch_size
+        self.t_patch_size = self.config.t_patch_size
+
+    def img2tokens(self, x_t: torch.Tensor) -> torch.Tensor:
+        return _patchify_3d(x_t, self.t_patch_size, self.patch_size)
+
+    @staticmethod
+    def _validate_batch(data: ModelInput) -> None:
+        if data.x_t.ndim != 5:
+            raise ValueError(f"x_t must be [B,C,T,H,W], got {tuple(data.x_t.shape)}")
+        if data.audio_x_t.ndim != 3 or data.txt_feat.ndim != 3:
+            raise ValueError("audio_x_t and txt_feat must be [B,L,C]")
+        batch_size = data.x_t.shape[0]
+        if data.audio_x_t.shape[0] != batch_size or data.txt_feat.shape[0] != batch_size:
+            raise ValueError("video, audio, and text batch sizes must match")
+
+        ref_values = (
+            data.ref_image_feat,
+            data.ref_image_feat_len,
+            data.ref_image_special_token_embedding,
+        )
+        if any(value is not None for value in ref_values) and not all(value is not None for value in ref_values):
+            raise ValueError("reference-image features, lengths, and special tokens must be provided together")
+        if data.ref_image_feat is not None:
+            if data.ref_image_feat.ndim != 6:
+                raise ValueError("ref_image_feat must be [B,M,C,T,H,W]")
+            if data.ref_image_feat_len.ndim != 3 or data.ref_image_feat_len.shape[-1] < 2:
+                raise ValueError("ref_image_feat_len must be [B,M,2]")
+            if data.ref_image_special_token_embedding.ndim != 3:
+                raise ValueError("ref_image_special_token_embedding must be [B,M,D]")
+            expected = data.ref_image_feat.shape[:2]
+            if data.ref_image_feat_len.shape[:2] != expected:
+                raise ValueError("reference-image feature and length batch/image counts differ")
+            if data.ref_image_special_token_embedding.shape[:2] != expected:
+                raise ValueError("reference-image feature and special-token batch/image counts differ")
+
+    def process_input(
+        self,
+        data: ModelInput,
+    ) -> PackedModelInput:
+        self._validate_batch(data)
+        batch_size, _, 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
+
+        has_ref_image = data.ref_image_feat is not None
+        num_ref_images = data.ref_image_feat.shape[1] if has_ref_image else 0
+        ref_device, ref_dtype = text_tokens.device, text_tokens.dtype
+
+        items: list[SingleData] = []
+        for batch_index in range(batch_size):
+            ref_image_feats: list[torch.Tensor] = []
+            ref_image_feat_lens: list[list[int]] = []
+            ref_image_special_tokens: list[torch.Tensor] = []
+            for image_index in range(num_ref_images):
+                feature = self.img2tokens(data.ref_image_feat[batch_index, image_index].unsqueeze(0)).squeeze(0)
+                ref_image_feats.append(feature)
+                ref_image_feat_lens.append(_len_to_list(data.ref_image_feat_len[batch_index, image_index]))
+                ref_image_special_tokens.append(
+                    data.ref_image_special_token_embedding[batch_index, image_index]
+                    .to(device=ref_device, dtype=ref_dtype)
+                    .unsqueeze(0)
+                )
+
+            if isinstance(data.t, torch.Tensor) and data.t.ndim == 0:
+                diffusion_t = data.t
+            else:
+                diffusion_t = data.t[batch_index]
+            items.append(
+                SingleData(
+                    video_x_t=video_tokens[batch_index],
+                    audio_x_t=audio_tokens[batch_index],
+                    audio_feat_len=_length_at(data.audio_feat_len, batch_index),
+                    txt_feat=text_tokens[batch_index],
+                    txt_feat_len=_length_at(data.txt_feat_len, batch_index),
+                    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=diffusion_t 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,
+                    ref_image_feats=ref_image_feats,
+                    ref_image_feat_lens=ref_image_feat_lens,
+                    ref_image_special_tokens=ref_image_special_tokens,
+                )
+            )
+
+        packed = SimplePackedData(items)
+        cu_seqlens = packed.cu_seqlen.to(device=data.x_t.device, dtype=torch.int32)
+        varlen_handler = VarlenHandler(
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=packed.max_seqlen,
+            max_seqlen_k=packed.max_seqlen,
+        )
+        return PackedModelInput(
+            token_sequence=packed.token_sequence,
+            coords_mapping=packed.coords_mapping,
+            modality_mapping=packed.modality_mapping,
+            varlen_handler=varlen_handler,
+            time_token_sequence=packed.time_token_sequence,
+            output_layout=packed,
+        )
+
+    @staticmethod
+    def process_output(
+        x: torch.Tensor,
+        output_layout: SimplePackedData,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return output_layout.depack_token_sequence(x)
+
+
+__all__ = [
+    "Magi2DataProxy",
+    "Magi2PreviewDataProxyConfig",
+    "Modality",
+    "ModelInput",
+    "PackedModelInput",
+    "SimplePackedData",
+    "SingleData",
+    "VarlenHandler",
+    "get_coords",
+    "sinusoidal_embedding_1d",
+]