[Fix] Restore MammothModa2 support

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -0,0 +1,839 @@
+import math
+
+import torch
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention_processor import Attention
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from einops import rearrange, repeat
+from torch import nn
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+
+from .rope_real import RotaryPosEmbedReal
+
+try:
+    from transformers.modeling_flash_attention_utils import (  # type: ignore
+        flash_attn_varlen_func,  # pyright: ignore[reportAttributeAccessIssue]
+        is_flash_attn_available,
+    )
+except Exception:  # pragma: no cover - best-effort compatibility
+    flash_attn_varlen_func = None  # type: ignore[assignment]
+
+    def is_flash_attn_available() -> bool:  # type: ignore[override]
+        return False
+
+
+from .rope_real import apply_real_rotary_emb
+
+_HAS_FLASH_ATTN_VARLEN = bool(is_flash_attn_available()) and flash_attn_varlen_func is not None
+
+
+class LuminaRMSNormZero(nn.Module):
+    """
+    Norm layer adaptive RMS normalization zero.
+
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        norm_eps: float,
+        norm_elementwise_affine: bool,
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(
+            min(embedding_dim, 1024),
+            4 * embedding_dim,
+            bias=True,
+        )
+
+        self.norm = Qwen2RMSNorm(embedding_dim, eps=norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        emb = self.linear(self.silu(emb))
+        scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None])
+        return x, gate_msa, scale_mlp, gate_mlp
+
+
+class LuminaFeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inner_dim: int,
+        multiple_of: int | None = 256,
+        ffn_dim_multiplier: float | None = None,
+    ):
+        super().__init__()
+
+        # custom hidden_size factor multiplier
+        if ffn_dim_multiplier is not None:
+            inner_dim = int(ffn_dim_multiplier * inner_dim)
+        inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
+
+        self.linear_1 = nn.Linear(
+            dim,
+            inner_dim,
+            bias=False,
+        )
+        self.linear_2 = nn.Linear(
+            inner_dim,
+            dim,
+            bias=False,
+        )
+        self.linear_3 = nn.Linear(
+            dim,
+            inner_dim,
+            bias=False,
+        )
+
+    def swiglu(self, x, y):
+        return F.silu(x.float(), inplace=False).to(x.dtype) * y
+
+    def forward(self, x):
+        h1, h2 = self.linear_1(x), self.linear_3(x)
+        return self.linear_2(self.swiglu(h1, h2))
+
+
+class LuminaLayerNormContinuous(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+        # However, this is how it was implemented in the original code, and it's rather likely you should
+        # set `elementwise_affine` to False.
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="layer_norm",
+        out_dim: int | None = None,
+    ):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
+
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = Qwen2RMSNorm(embedding_dim, eps=eps)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+        self.linear_2 = None
+        if out_dim is not None:
+            self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        conditioning_embedding: torch.Tensor,
+    ) -> torch.Tensor:
+        scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
+        x = self.norm(x) * (1 + scale)[:, None, :]
+
+        if self.linear_2 is not None:
+            x = self.linear_2(x)
+
+        return x
+
+
+class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int = 2520,
+        text_feat_dim: int = 3584,
+        frequency_embedding_size: int = 256,
+        norm_eps: float = 1e-5,
+        timestep_scale: float = 1.0,
+    ) -> None:
+        super().__init__()
+
+        self.time_proj = Timesteps(
+            num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=timestep_scale
+        )
+
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024)
+        )
+
+        self.caption_embedder = nn.Sequential(
+            Qwen2RMSNorm(text_feat_dim, eps=norm_eps),
+            nn.Linear(text_feat_dim, hidden_size, bias=True),
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        timestep_proj = self.time_proj(timestep).to(dtype=dtype)
+        time_embed = self.timestep_embedder(timestep_proj)
+        caption_embed = self.caption_embedder(text_hidden_states)
+        return time_embed, caption_embed
+
+
+class SimpleQFormerImageRefiner(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        output_hidden_size: int | None = None,
+        num_queries: int = 128,
+        num_layers: int = 2,
+        num_heads: int | None = None,
+        dropout: float = 0.0,
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        input_hidden_size = hidden_size
+        hidden_size = output_hidden_size or hidden_size
+        self.hidden_size = hidden_size
+        self.num_queries = num_queries
+        # ensure num_heads divides hidden_size
+        if num_heads is None:
+            num_heads = max(1, hidden_size // 128)
+        self.num_heads = self._choose_valid_num_heads(hidden_size, num_heads)
+        self.input_proj = nn.Sequential(
+            Qwen2RMSNorm(input_hidden_size, eps=norm_eps),
+            nn.Linear(input_hidden_size, hidden_size, bias=True),
+        )
+
+        # Learnable query embeddings
+        scale = hidden_size**-0.5
+        self.query = nn.Parameter(scale * torch.randn(1, num_queries, hidden_size))
+
+        # Decoder layers
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            self.layers.append(
+                nn.ModuleDict(
+                    dict(
+                        ln_q1=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        self_attn=nn.MultiheadAttention(
+                            embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True
+                        ),
+                        ln_q2=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        cross_attn=nn.MultiheadAttention(
+                            embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True
+                        ),
+                        ln_ffn=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        ffn=LuminaFeedForward(dim=hidden_size, inner_dim=4 * hidden_size),
+                    )
+                )
+            )
+
+    @staticmethod
+    def _choose_valid_num_heads(hidden_size: int, proposed_heads: int, preferred_head_dim: int = 128) -> int:
+        """Pick a number of heads that divides hidden_size, close to proposed or preferred."""
+        # If proposed is valid, use it
+        if proposed_heads > 0 and hidden_size % proposed_heads == 0:
+            return proposed_heads
+        # target based on preferred head dim
+        target = max(1, round(hidden_size / preferred_head_dim))
+        # collect divisors up to 128 heads (more than enough)
+        max_heads_cap = min(128, hidden_size)
+        divisors = [d for d in range(1, max_heads_cap + 1) if hidden_size % d == 0]
+        # choose closest to target
+        best = min(divisors, key=lambda d: (abs(d - target), -d))
+        return best
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape (batch, seq_len, input_dim)
+        Returns:
+            Tensor of shape (batch, num_queries, hidden_size)
+        """
+        batch, _, _ = x.shape
+        kv = self.input_proj(x)
+        q = self.query.repeat(batch, 1, 1).to(kv.dtype)
+
+        for layer in self.layers:
+            # Self-attention on queries
+            q_norm = layer["ln_q1"](q)
+            attn_out, _ = layer["self_attn"](q_norm, q_norm, q_norm, need_weights=False)
+            q = q + attn_out
+
+            # Cross-attention: queries attend to inputs
+            q_norm = layer["ln_q2"](q)
+            cross_out, _ = layer["cross_attn"](q_norm, kv, kv, need_weights=False, key_padding_mask=attention_mask)
+            q = q + cross_out
+
+            # Feed-forward
+            q = q + layer["ffn"](layer["ln_ffn"](q))
+
+        return q
+
+
+class AttnProcessor:
+    def __init__(self) -> None:
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor requires PyTorch 2.0+ (F.scaled_dot_product_attention).")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        base_sequence_length: int | None = None,
+    ) -> torch.Tensor:
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        # Get Query-Key-Value Pair
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query_dim = query.shape[-1]
+        inner_dim = key.shape[-1]
+        head_dim = query_dim // attn.heads
+        dtype = query.dtype
+
+        # Get key-value heads
+        kv_heads = inner_dim // head_dim
+
+        # Reshape tensors for attention computation
+        query = query.view(batch_size, -1, attn.heads, head_dim)
+        key = key.view(batch_size, -1, kv_heads, head_dim)
+        value = value.view(batch_size, -1, kv_heads, head_dim)
+
+        # Apply Query-Key normalization
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply Rotary Position Embeddings
+        if image_rotary_emb is not None:
+            query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1])
+            key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1])
+
+        query, key = query.to(dtype), key.to(dtype)
+
+        # Calculate attention scale
+        if base_sequence_length is not None:
+            softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
+        else:
+            softmax_scale = attn.scale
+
+        if _HAS_FLASH_ATTN_VARLEN and attention_mask is not None and hidden_states.is_cuda:
+            # Flash-Attn varlen expects packed tokens + cu_seqlens. Here we only need
+            # the self-attention case (q/k/v share the same padding mask).
+            attention_mask = attention_mask.to(torch.bool)
+            seqlens = attention_mask.sum(dim=-1, dtype=torch.int32)
+            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+            max_seqlen = int(seqlens.max().item())
+            cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+
+            query_states = query.reshape(batch_size * sequence_length, attn.heads, head_dim)[indices]
+            key_states = key.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices]
+            value_states = value.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices]
+
+            if kv_heads < attn.heads:
+                key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads)
+                value_states = repeat(value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads)
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                dropout_p=0.0,
+                causal=False,
+                softmax_scale=softmax_scale,
+            )
+
+            out = torch.zeros(
+                (batch_size * sequence_length, attn.heads, head_dim),
+                device=attn_output_unpad.device,
+                dtype=attn_output_unpad.dtype,
+            )
+            out[indices] = attn_output_unpad
+            hidden_states = out.view(batch_size, sequence_length, attn.heads, head_dim).flatten(-2)
+            hidden_states = hidden_states.type_as(query)
+        else:
+            # PyTorch SDPA path.
+            attn_mask = None
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(torch.bool)
+                attn_mask = attention_mask.view(batch_size, 1, 1, -1)
+
+            query = query.transpose(1, 2)  # [B, H, S, D]
+            key = key.transpose(1, 2)  # [B, H_kv, S, D]
+            value = value.transpose(1, 2)
+
+            if kv_heads < attn.heads:
+                key = key.repeat_interleave(attn.heads // kv_heads, dim=1)
+                value = value.repeat_interleave(attn.heads // kv_heads, dim=1)
+
+            hidden_states = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=attn_mask,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=softmax_scale,
+            )
+
+            if attention_mask is not None:
+                # Keep padding tokens consistent with the flash-varlen path (zero output).
+                hidden_states = hidden_states * attention_mask[:, None, :, None]
+
+            hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+            hidden_states = hidden_states.type_as(query)
+
+        # Apply output projection
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        num_kv_heads: int,
+        multiple_of: int,
+        ffn_dim_multiplier: float,
+        norm_eps: float,
+        modulation: bool = True,
+    ) -> None:
+        """Initialize the transformer block."""
+        super().__init__()
+        self.head_dim = dim // num_attention_heads
+        self.modulation = modulation
+
+        processor = AttnProcessor()
+
+        # Initialize attention layer
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=dim // num_attention_heads,
+            qk_norm=None,
+            heads=num_attention_heads,
+            kv_heads=num_kv_heads,
+            eps=1e-5,
+            bias=False,
+            out_bias=False,
+            processor=processor,
+        )
+        # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。
+        self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5)
+        self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5)
+
+        # Initialize feed-forward network
+        self.feed_forward = LuminaFeedForward(
+            dim=dim, inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier
+        )
+
+        # Initialize normalization layers
+        if modulation:
+            self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True)
+        else:
+            self.norm1 = Qwen2RMSNorm(dim, eps=norm_eps)
+
+        self.ffn_norm1 = Qwen2RMSNorm(dim, eps=norm_eps)
+        self.norm2 = Qwen2RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = Qwen2RMSNorm(dim, eps=norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        image_rotary_emb: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.modulation:
+            if temb is None:
+                raise ValueError("temb must be provided when modulation is enabled")
+
+            norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
+            attn_output = self.attn(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_hidden_states,
+                attention_mask=attention_mask,
+                image_rotary_emb=image_rotary_emb,
+            )
+            hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output)
+            mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1)))
+            hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output)
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+            attn_output = self.attn(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_hidden_states,
+                attention_mask=attention_mask,
+                image_rotary_emb=image_rotary_emb,
+            )
+            hidden_states = hidden_states + self.norm2(attn_output)
+            mlp_output = self.feed_forward(self.ffn_norm1(hidden_states))
+            hidden_states = hidden_states + self.ffn_norm2(mlp_output)
+
+        return hidden_states
+
+
+class Transformer2DModel(ModelMixin, ConfigMixin):
+    """MammothModa2 DiT transformer"""
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        out_channels: int | None = None,
+        hidden_size: int = 2304,
+        num_layers: int = 26,
+        num_refiner_layers: int = 2,
+        num_attention_heads: int = 24,
+        num_kv_heads: int = 8,
+        multiple_of: int = 256,
+        ffn_dim_multiplier: float | None = None,
+        norm_eps: float = 1e-5,
+        axes_dim_rope: tuple[int, int, int] = (32, 32, 32),
+        axes_lens: tuple[int, int, int] = (300, 512, 512),
+        text_feat_dim: int = 1024,
+        timestep_scale: float = 1.0,
+    ) -> None:
+        """Initialize the  transformer model."""
+        super().__init__()
+        self.hidden_size = hidden_size
+
+        # Validate configuration
+        if (hidden_size // num_attention_heads) != sum(axes_dim_rope):
+            raise ValueError(
+                f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) "
+                f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})"
+            )
+
+        self.out_channels = out_channels or in_channels
+        # Initialize embeddings
+        self.rope_embedder = RotaryPosEmbedReal(
+            theta=10000,
+            axes_dim=axes_dim_rope,
+            axes_lens=axes_lens,
+            patch_size=patch_size,
+        )
+
+        self.x_embedder = nn.Linear(
+            in_features=patch_size * patch_size * in_channels,
+            out_features=hidden_size,
+        )
+
+        self.ref_image_patch_embedder = nn.Linear(
+            in_features=patch_size * patch_size * in_channels,
+            out_features=hidden_size,
+        )
+
+        self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding(
+            hidden_size=hidden_size,
+            text_feat_dim=text_feat_dim,
+            norm_eps=norm_eps,
+            timestep_scale=timestep_scale,
+        )
+
+        # Initialize transformer blocks
+        self.noise_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=True,
+                )
+                for _ in range(num_refiner_layers)
+            ]
+        )
+
+        self.ref_image_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=True,
+                )
+                for _ in range(num_refiner_layers)
+            ]
+        )
+
+        self.context_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=False,
+                )
+                for _ in range(num_refiner_layers)
+            ]
+        )
+
+        # 3. Transformer blocks
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=True,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 4. Output norm & projection
+        self.norm_out = LuminaLayerNormContinuous(
+            embedding_dim=hidden_size,
+            conditioning_embedding_dim=min(hidden_size, 1024),
+            elementwise_affine=False,
+            eps=1e-6,
+            bias=True,
+            out_dim=patch_size * patch_size * self.out_channels,
+        )
+
+        # Add learnable embeddings to distinguish different images
+        self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size))  # support max 5 ref images
+
+    def _validate_inputs(
+        self,
+        hidden_states: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        ref_image_hidden_states: list[list[torch.Tensor]] | None,
+        return_dict: bool,
+    ) -> tuple[int, int, int]:
+        if return_dict:
+            raise ValueError("return_dict=True is not supported in vLLM inference.")
+        if ref_image_hidden_states is not None:
+            raise ValueError("ref_image_hidden_states is not supported in vLLM inference.")
+        if hidden_states.ndim != 4:
+            raise ValueError(f"Expected hidden_states to be 4D [B,C,H,W], got shape={tuple(hidden_states.shape)}")
+
+        batch_size, _channels, height, width = hidden_states.shape
+        if batch_size != text_hidden_states.shape[0] or batch_size != text_attention_mask.shape[0]:
+            raise ValueError(
+                "Batch size mismatch: "
+                f"hidden_states={batch_size}, text_hidden_states={text_hidden_states.shape[0]}, "
+                f"text_attention_mask={text_attention_mask.shape[0]}"
+            )
+
+        p = self.config.patch_size
+        if height % p != 0 or width % p != 0:
+            raise ValueError(f"Input latent H/W must be divisible by patch_size={p}, got {height}x{width}")
+        return batch_size, height, width
+
+    def _prepare_embeddings(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        batch_size: int,
+        height: int,
+        width: int,
+        ar_image_hidden_states: torch.Tensor | None = None,
+        ar_image_attention_mask: torch.Tensor | None = None,
+    ):
+        device = hidden_states.device
+        p = self.config.patch_size
+
+        temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states.dtype)
+        image_embedder = getattr(self.time_caption_embed, "image_embedder", None)
+        if image_embedder is not None and ar_image_hidden_states is not None:
+            if ar_image_attention_mask is None:
+                ar_image_attention_mask = torch.ones(
+                    ar_image_hidden_states.shape[:2],
+                    dtype=torch.bool,
+                    device=ar_image_hidden_states.device,
+                )
+            image_hidden_states = image_embedder(
+                ar_image_hidden_states,
+                ~ar_image_attention_mask.bool(),
+            )
+            image_attention_mask = torch.ones(
+                image_hidden_states.shape[:2],
+                dtype=torch.bool,
+                device=image_hidden_states.device,
+            )
+            text_hidden_states = torch.cat([text_hidden_states, image_hidden_states], dim=1)
+            text_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+
+        img_tokens = rearrange(hidden_states, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p)
+        img_tokens = self.x_embedder(img_tokens)
+
+        img_len = (height // p) * (width // p)
+        img_mask = torch.ones((batch_size, img_len), dtype=torch.bool, device=device)
+        l_effective_img_len = [img_len for _ in range(batch_size)]
+        img_sizes = [(height, width) for _ in range(batch_size)]
+
+        l_effective_ref_img_len = [[] for _ in range(batch_size)]
+        ref_img_sizes = [None for _ in range(batch_size)]
+
+        (
+            context_rotary_emb,
+            _ref_img_rotary_emb,
+            noise_rotary_emb,
+            rotary_emb,
+            encoder_seq_lengths,
+            seq_lengths,
+        ) = self.rope_embedder(
+            freqs_cis,
+            text_attention_mask,
+            l_effective_ref_img_len,
+            l_effective_img_len,
+            ref_img_sizes,
+            img_sizes,
+            device,
+        )
+
+        return (
+            temb,
+            text_hidden_states,
+            text_attention_mask,
+            img_tokens,
+            img_mask,
+            img_len,
+            context_rotary_emb,
+            noise_rotary_emb,
+            rotary_emb,
+            encoder_seq_lengths,
+            seq_lengths,
+        )
+
+    def _apply_refiners(
+        self,
+        text_hidden_states: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        context_rotary_emb: torch.Tensor,
+        img_tokens: torch.Tensor,
+        img_mask: torch.Tensor,
+        noise_rotary_emb: torch.Tensor,
+        temb: torch.Tensor,
+    ):
+        for layer in self.context_refiner:
+            text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb)
+
+        for layer in self.noise_refiner:
+            img_tokens = layer(img_tokens, img_mask, noise_rotary_emb, temb)
+
+        return text_hidden_states, img_tokens
+
+    def _apply_transformer_layers(self, hidden_states, attention_mask, rotary_emb, temb):
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
+        return hidden_states
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        ref_image_hidden_states: list[list[torch.Tensor]] | None = None,
+        ar_image_hidden_states: torch.Tensor | None = None,
+        ar_image_attention_mask: torch.Tensor | None = None,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        batch_size, height, width = self._validate_inputs(
+            hidden_states, text_hidden_states, text_attention_mask, ref_image_hidden_states, return_dict
+        )
+
+        (
+            temb,
+            text_hidden_states,
+            text_attention_mask,
+            img_tokens,
+            img_mask,
+            img_len,
+            context_rotary_emb,
+            noise_rotary_emb,
+            rotary_emb,
+            encoder_seq_lengths,
+            seq_lengths,
+        ) = self._prepare_embeddings(
+            hidden_states,
+            timestep,
+            text_hidden_states,
+            text_attention_mask,
+            freqs_cis,
+            batch_size,
+            height,
+            width,
+            ar_image_hidden_states,
+            ar_image_attention_mask,
+        )
+
+        text_hidden_states, img_tokens = self._apply_refiners(
+            text_hidden_states,
+            text_attention_mask,
+            context_rotary_emb,
+            img_tokens,
+            img_mask,
+            noise_rotary_emb,
+            temb,
+        )
+
+        max_seq_len = max(seq_lengths)
+        attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
+        joint_hidden_states = hidden_states.new_zeros(batch_size, max_seq_len, self.config.hidden_size)
+        for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+            attention_mask[i, :seq_len] = True
+            joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len]
+            joint_hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] = img_tokens[i, :img_len]
+
+        hidden_states = self._apply_transformer_layers(joint_hidden_states, attention_mask, rotary_emb, temb)
+
+        hidden_states = self.norm_out(hidden_states, temb)
+
+        p = self.config.patch_size
+        img_hidden_states = torch.stack(
+            [
+                hidden_states[i, encoder_seq_len : encoder_seq_len + img_len]
+                for i, encoder_seq_len in enumerate(encoder_seq_lengths)
+            ],
+            dim=0,
+        )
+        output = rearrange(
+            img_hidden_states,
+            "b (h w) (p1 p2 c) -> b c (h p1) (w p2)",
+            h=height // p,
+            w=width // p,
+            p1=p,
+            p2=p,
+        )
+        return output

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -0,0 +1,405 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any, ClassVar
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
+
+from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel
+from .rope_real import RotaryPosEmbedReal
+from .schedulers import FlowMatchEulerDiscreteScheduler
+
+logger = init_logger(__name__)
+
+
+class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
+    """
+    MammothModa2 DiT + VAE generation stage (non-autoregressive).
+
+    This stage expects "image condition token hidden states" from the upstream AR stage,
+    and outputs image tensors via diffusion transformer + VAE decode.
+
+    """
+
+    _dit_modules: ClassVar[list[str]] = ["gen_transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"]
+    _vae_modules: ClassVar[list[str]] = ["gen_vae"]
+
+    have_multimodal_outputs = True
+
+    # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "llm_model.": None,
+            "gen_tokenizer.": None,
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        del prefix
+
+        hf_config = vllm_config.model_config.hf_config
+        if not isinstance(hf_config, Mammothmoda2Config):
+            raise TypeError(f"Expected Mammothmoda2Config, got {type(hf_config)}")
+
+        self.config = hf_config
+
+        # --- Build DiT / VAE modules (names must match checkpoint keys) ---
+        if self.config.gen_vae_config is None or self.config.gen_dit_config is None:
+            raise ValueError("Mammothmoda2Config.gen_vae_config / gen_dit_config must not be None")
+
+        self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config)
+        self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
+
+        # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
+        llm_hidden_size = 0
+        text_config = self.config.get_text_config()
+        if text_config is None:
+            logger.warning("No text config; failed to infer llm_hidden_size.")
+        elif not hasattr(text_config, "hidden_size"):
+            logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.")
+        else:
+            llm_hidden_size = int(text_config.hidden_size or 0)
+        if llm_hidden_size <= 0:
+            raise ValueError(
+                "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"
+            )
+        self._reinit_caption_embedder(llm_hidden_size)
+
+        # Optional image condition Q-Former. Preview stores it as a standalone
+        # module; Dev stores it under the DiT timestep/caption embedder.
+        llm_model_type = getattr(self.config.llm_config, "model_type", "")
+        refiner_config = self.config.gen_image_condition_refiner_config
+        if refiner_config is not None and llm_model_type == "mammothmoda2_qwen3_vl":
+            dit_hidden_size = int(self.gen_transformer.hidden_size)
+            self.gen_transformer.time_caption_embed.image_embedder = SimpleQFormerImageRefiner(
+                hidden_size=llm_hidden_size,
+                output_hidden_size=dit_hidden_size,
+                num_heads=max(1, dit_hidden_size // 128),
+                **refiner_config,
+            )
+            self.gen_image_condition_refiner = None
+        elif refiner_config is not None:
+            self.gen_image_condition_refiner = SimpleQFormerImageRefiner(
+                hidden_size=llm_hidden_size,
+                **refiner_config,
+            )
+        else:
+            self.gen_image_condition_refiner = None
+
+        # Precompute rotary freqs for diffusion transformer
+        # IMPORTANT: follow upstream mammothmoda: use top-level `config.gen_axes_*`
+        # (the checkpoint's `gen_dit_config.axes_lens` can be as small as 1024,
+        # which is insufficient for vLLM dummy-run/cudagraph warmup).
+        self.gen_freqs_cis = RotaryPosEmbedReal.get_freqs_real(
+            tuple(self.config.gen_axes_dim_rope),
+            tuple(self.config.gen_axes_lens),
+            theta=10000,
+        )
+
+        # vLLM PP interface compatibility
+        self.make_empty_intermediate_tensors = lambda: None
+
+        self._llm_hidden_size = llm_hidden_size
+
+    def _reinit_caption_embedder(self, in_features: int) -> None:
+        # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`:
+        # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features).
+        out_features = int(getattr(self.gen_transformer, "hidden_size", 0) or self.gen_transformer.config.hidden_size)
+        self.gen_transformer.time_caption_embed.caption_embedder = nn.Sequential(
+            Qwen2RMSNorm(in_features, eps=1e-5),
+            nn.Linear(in_features, out_features, bias=True),
+        )
+
+    def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
+        if num_reqs <= 0:
+            raise ValueError(f"num_reqs must be positive, got {num_reqs}")
+        if num_reqs > 1:
+            raise NotImplementedError(
+                f"get_dummy_runtime_additional_information does not support num_reqs > 1, got {num_reqs}"
+            )
+        text_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32)
+        image_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32)
+        negative_prompt_embeds = torch.zeros((0, self._llm_hidden_size), dtype=torch.float32)
+        info = {
+            "text_prompt_embeds": text_prompt_embeds,
+            "image_prompt_embeds": image_prompt_embeds,
+            "negative_prompt_embeds": negative_prompt_embeds,
+            "negative_prompt_attention_mask": [],
+            "image_height": [512],
+            "image_width": [512],
+            "text_guidance_scale": [1.0],
+            "cfg_range": [0.0, 1.0],
+            "num_inference_steps": [1],
+        }
+        return [info for _ in range(num_reqs)]
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        # DiT stage does not consume token embeddings; return a dummy tensor.
+        try:
+            dtype = next(self.parameters()).dtype
+        except StopIteration:
+            dtype = torch.float32
+        return torch.zeros(
+            (input_ids.numel(), self._llm_hidden_size),
+            device=input_ids.device,
+            dtype=dtype,
+        )
+
+    def _split_ar_conditions(
+        self,
+        *,
+        full_hidden_states: torch.Tensor,
+        full_token_ids: list[int],
+        answer_start_index: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split AR-stage hidden states into text / image condition embeds.
+
+        The token ids that distinguish question (text) tokens, generated visual
+        tokens, and multi-modal placeholder tokens are read from the model config
+        (``gen_vocab_start_index`` and the vision placeholder token ids), so the
+        caller no longer needs to pass them. Mirrors the masking the bespoke
+        MammothModa2 example performed via ar2dit.
+        """
+        gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index)
+        visual_ids = [
+            int(self.config.image_token_id),
+            int(self.config.video_token_id),
+            int(self.config.vision_start_token_id),
+            int(self.config.vision_end_token_id),
+        ]
+
+        device = full_hidden_states.device
+        token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device)
+        positions = torch.arange(token_ids.shape[0], device=device)
+        questions_mask = positions < answer_start_index
+        answers_mask = ~questions_mask
+        gen_token_mask = token_ids >= gen_vocab_start_index
+        visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device))
+        text_mask = questions_mask & ~(visual_token_mask | gen_token_mask)
+        image_mask = answers_mask & gen_token_mask
+
+        text_cond = full_hidden_states[text_mask].to(dtype=torch.float32).contiguous()
+        image_cond = full_hidden_states[image_mask].to(dtype=torch.float32).contiguous()
+        return text_cond, image_cond
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        *,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,  # noqa: ARG002
+    ) -> OmniOutput:
+        runtime_addi = kwargs.get("runtime_additional_information", None)
+        info = runtime_addi[0]
+
+        # Sampling knobs are declared in vllm_omni/model_extras/mammothmodal2_preview.py
+        # and routed via extra_body -> sampling_params.extra_args (surfaced here as
+        # ``sampling_extra_args``). Fall back to runtime_additional_information for the
+        # legacy bespoke-example path during the transition.
+        extra_args_list = kwargs.get("sampling_extra_args") or []
+        extra_args = extra_args_list[0] if extra_args_list else {}
+        text_guidance_scale = float(extra_args.get("text_guidance_scale", info["text_guidance_scale"][0]))
+        cfg_range_val = extra_args.get("cfg_range", info["cfg_range"])
+        cfg_range = float(cfg_range_val[0]), float(cfg_range_val[1])
+        num_inference_steps = int(extra_args.get("num_inference_steps", info["num_inference_steps"][0]))
+
+        negative_cond = info.get("negative_prompt_embeds")
+        negative_attention_mask = info.get("negative_prompt_attention_mask")
+        image_hw = info["image_height"][0], info["image_width"][0]
+
+        # Split the AR hidden states into text / image conditions. The token ids that
+        # drive the split are sourced from the model config (see _split_ar_conditions),
+        # formerly supplied by the bespoke example via additional_information. Legacy
+        # fallback: ar2dit may have already produced the split conditions.
+        if "text_prompt_embeds" in info:
+            text_cond = info["text_prompt_embeds"]
+            image_cond = info["image_prompt_embeds"]
+        else:
+            text_cond, image_cond = self._split_ar_conditions(
+                full_hidden_states=info["full_hidden_states"],
+                full_token_ids=info["full_token_ids"],
+                answer_start_index=int(info["answer_start_index"][0]),
+            )
+
+        # Move to model device/dtype.
+        model_device = next(self.parameters()).device
+        if self.gen_image_condition_refiner is not None:
+            target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype
+        else:
+            target_dtype = next(self.gen_transformer.parameters()).dtype
+
+        def _ensure_2d(x: torch.Tensor, name: str) -> torch.Tensor:
+            if x.ndim == 3 and x.shape[0] == 1:
+                x = x[0]
+            if x.ndim != 2:
+                raise ValueError(f"Expected {name} to be 2D [T,H], got shape={tuple(x.shape)}")
+            return x
+
+        text_cond = _ensure_2d(text_cond, "text_prompt_embeds")
+        image_cond = _ensure_2d(image_cond, "image_prompt_embeds")
+        if image_cond.shape[0] == 0:
+            answer_token_ids = info.get("full_token_ids", [])[int(info.get("answer_start_index", [0])[0]) :]
+            raise ValueError(
+                "MammothModa2 AR stage produced no visual-token hidden states; "
+                "the DiT stage requires at least one generated visual token. "
+                f"Generated token ids: {answer_token_ids[:32]}"
+            )
+        text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
+        image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
+
+        text_embeds = text_cond.unsqueeze(0)  # [1, T_text, H]
+        text_attention_mask = torch.ones(
+            (1, text_embeds.shape[1]),
+            dtype=torch.bool,
+            device=text_embeds.device,
+        )
+
+        image_embeds = image_cond.unsqueeze(0)  # [1, T_img, H]
+        image_attention_mask = torch.ones(
+            (1, image_embeds.shape[1]),
+            dtype=torch.bool,
+            device=image_embeds.device,
+        )
+
+        # Apply optional refiner ONLY on image condition tokens.
+        if self.gen_image_condition_refiner is not None and image_embeds.shape[1] > 0:
+            image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_attention_mask.bool())
+            image_attention_mask = torch.ones(
+                image_embeds.shape[:2],
+                dtype=torch.bool,
+                device=image_embeds.device,
+            )
+
+        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
+        if nested_image_embedder is None:
+            prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
+            prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+            ar_image_embeds = None
+            ar_image_attention_mask = None
+        else:
+            prompt_embeds = text_embeds
+            prompt_attention_mask = text_attention_mask
+            ar_image_embeds = image_embeds
+            ar_image_attention_mask = image_attention_mask
+
+        # Prepare negative prompt (for CFG). If none provided, fall back to unconditional.
+        negative_prompt_embeds = None
+        negative_prompt_attention_mask = None
+        if text_guidance_scale > 1.0:
+            if negative_cond is not None:
+                negative_cond = _ensure_2d(negative_cond, "negative_prompt_embeds")
+                negative_prompt_embeds = (
+                    negative_cond.to(device=model_device, dtype=target_dtype, non_blocking=True)
+                    .contiguous()
+                    .unsqueeze(0)
+                )
+                if isinstance(negative_attention_mask, torch.Tensor):
+                    neg_mask = negative_attention_mask
+                elif isinstance(negative_attention_mask, list):
+                    neg_mask = torch.tensor(negative_attention_mask, dtype=torch.bool)
+                else:
+                    neg_mask = None
+                if neg_mask is None:
+                    negative_prompt_attention_mask = torch.ones(
+                        (1, negative_prompt_embeds.shape[1]),
+                        dtype=torch.bool,
+                        device=negative_prompt_embeds.device,
+                    )
+                else:
+                    neg_mask = neg_mask.to(device=negative_prompt_embeds.device, dtype=torch.bool)
+                    if neg_mask.ndim == 1:
+                        neg_mask = neg_mask.unsqueeze(0)
+                    negative_prompt_attention_mask = neg_mask
+            else:
+                hidden_size = int(prompt_embeds.shape[-1])
+                negative_prompt_embeds = torch.zeros(
+                    (1, 0, hidden_size),
+                    dtype=target_dtype,
+                    device=prompt_embeds.device,
+                )
+                negative_prompt_attention_mask = torch.zeros(
+                    (1, 0),
+                    dtype=torch.bool,
+                    device=prompt_embeds.device,
+                )
+
+        # Output image size (px), passed from stage input processor.
+        height, width = image_hw
+        if height <= 0 or width <= 0:
+            raise ValueError(f"Invalid image size: {height}x{width}")
+        if height % 16 != 0 or width % 16 != 0:
+            raise ValueError(f"Image size must be multiples of 16, got {height}x{width}")
+        vae_scale_factor = 16
+
+        latent_channels = int(self.gen_transformer.config.in_channels)
+        shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
+        latents = randn_tensor(shape, device=prompt_embeds.device, dtype=prompt_embeds.dtype)
+
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        scheduler.set_timesteps(
+            num_inference_steps=num_inference_steps,
+            device=prompt_embeds.device,
+            num_tokens=latents.shape[-2] * latents.shape[-1],
+        )
+
+        # Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
+        total_steps = max(1, len(scheduler.timesteps))
+        for i, t in enumerate(scheduler.timesteps):
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+            model_pred = self.gen_transformer(
+                hidden_states=latents,
+                timestep=timestep,
+                text_hidden_states=prompt_embeds,
+                text_attention_mask=prompt_attention_mask,
+                ref_image_hidden_states=None,
+                ar_image_hidden_states=ar_image_embeds,
+                ar_image_attention_mask=ar_image_attention_mask,
+                freqs_cis=self.gen_freqs_cis,
+            )
+            guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0
+            if guidance_scale > 1.0 and negative_prompt_embeds is not None:
+                model_pred_uncond = self.gen_transformer(
+                    hidden_states=latents,
+                    timestep=timestep,
+                    text_hidden_states=negative_prompt_embeds,
+                    text_attention_mask=negative_prompt_attention_mask,
+                    ref_image_hidden_states=None,
+                    freqs_cis=self.gen_freqs_cis,
+                )
+                model_pred = model_pred_uncond + guidance_scale * (model_pred - model_pred_uncond)
+            latents = scheduler.step(model_pred, t, latents, return_dict=False)[0]
+            latents = latents.to(dtype=prompt_embeds.dtype)
+
+        # VAE decode
+        if self.gen_vae.config.scaling_factor is not None:
+            latents = latents / self.gen_vae.config.scaling_factor
+        if self.gen_vae.config.shift_factor is not None:
+            latents = latents + self.gen_vae.config.shift_factor
+        image = self.gen_vae.decode(latents, return_dict=False)[0]
+
+        return OmniOutput(
+            text_hidden_states=inputs_embeds,  # placeholder, not used by runner
+            multimodal_outputs=image,
+            intermediate_tensors=None,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:  # noqa: ARG002
+        return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm_omni/diffusion/models/mammoth_moda2/rope_real.py

@@ -0,0 +1,252 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates and/or its affiliates
+#
+# 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 numpy as np
+import torch
+from diffusers.models.embeddings import get_1d_rotary_pos_embed as _get_1d_rotary_pos_embed
+from einops import repeat
+from torch import nn
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def apply_real_rotary_emb(x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor:
+    """
+    Apply real-valued rotary embeddings to input tensor.
+
+    Args:
+        x: Input tensor of shape [..., seq_len, num_heads, dim] or [..., seq_len, dim]
+        freqs_cos: Cosine frequencies of shape [batch, seq_len, dim] or [seq_len, dim]
+        freqs_sin: Sine frequencies of shape [batch, seq_len, dim] or [seq_len, dim]
+
+    Returns:
+        Tensor with rotary embeddings applied
+    """
+    # x: [batch, seq_len, num_heads, dim] or [batch, seq_len, dim]
+    # freqs_cos: [batch, seq_len, dim] or [seq_len, dim]
+    # freqs_sin: [batch, seq_len, dim] or [seq_len, dim]
+
+    x_shape = x.shape
+    if len(x_shape) == 4:
+        batch, seq_len, num_heads, dim = x_shape
+        x_reshaped = x.view(batch, seq_len, num_heads, dim // 2, 2)
+    elif len(x_shape) == 3:
+        batch, seq_len, dim = x_shape
+        num_heads = None
+        x_reshaped = x.view(batch, seq_len, dim // 2, 2)
+    else:
+        raise ValueError(f"Unsupported x shape: {x.shape}")
+
+    # freqs_cos/sin: [batch, seq_len, dim] or [seq_len, dim]
+    # Expand freqs_cos/sin to [batch, seq_len, dim] if needed
+    if freqs_cos.dim() == 2:
+        # [seq_len, dim] -> [1, seq_len, dim]
+        freqs_cos = freqs_cos.unsqueeze(0)
+        freqs_sin = freqs_sin.unsqueeze(0)
+    if freqs_cos.shape[0] == 1 and batch > 1:
+        freqs_cos = freqs_cos.expand(batch, -1, -1)
+        freqs_sin = freqs_sin.expand(batch, -1, -1)
+
+    # Reshape freqs to [batch, seq_len, dim//2, 2]
+    freqs_cos_reshaped = freqs_cos.view(batch, seq_len, dim // 2, 2)
+    freqs_sin_reshaped = freqs_sin.view(batch, seq_len, dim // 2, 2)
+
+    cos_1 = freqs_cos_reshaped[..., 0]  # [batch, seq_len, dim//2]
+    cos_2 = freqs_cos_reshaped[..., 1]  # [batch, seq_len, dim//2]
+    sin_1 = freqs_sin_reshaped[..., 0]  # [batch, seq_len, dim//2]
+    sin_2 = freqs_sin_reshaped[..., 1]  # [batch, seq_len, dim//2]
+
+    # Broadcast cos/sin to match x_reshaped
+    if len(x_shape) == 4:
+        # [batch, seq_len, 1, dim//2]
+        cos_1 = cos_1.unsqueeze(2)
+        cos_2 = cos_2.unsqueeze(2)
+        sin_1 = sin_1.unsqueeze(2)
+        sin_2 = sin_2.unsqueeze(2)
+
+    x1 = x_reshaped[..., 0]  # [..., seq_len, num_heads, dim//2] or [..., seq_len, dim//2]
+    x2 = x_reshaped[..., 1]  # same
+
+    out1 = x1 * cos_1 - x2 * sin_1
+    out2 = x1 * sin_2 + x2 * cos_2
+
+    out = torch.stack([out1, out2], dim=-1)
+    return out.view(*x_shape)
+
+
+def get_1d_rotary_pos_embed_real(
+    dim: int,
+    pos: np.ndarray | int,
+    theta: float = 10000.0,
+    linear_factor: float = 1.0,
+    ntk_factor: float = 1.0,
+    freqs_dtype: torch.dtype = torch.float32,
+):
+    freqs_cos, freqs_sin = _get_1d_rotary_pos_embed(
+        dim,
+        pos,
+        theta=theta,
+        use_real=True,
+        linear_factor=linear_factor,
+        ntk_factor=ntk_factor,
+        repeat_interleave_real=True,
+        freqs_dtype=freqs_dtype,
+    )
+    return freqs_cos, freqs_sin
+
+
+class RotaryPosEmbedReal(nn.Module):
+    def __init__(
+        self, theta: int, axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], patch_size: int = 2
+    ):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        self.axes_lens = axes_lens
+        self.patch_size = patch_size
+
+    @staticmethod
+    def get_freqs_real(
+        axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        freqs_real = []
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
+        for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
+            cos_emb, sin_emb = get_1d_rotary_pos_embed_real(d, e, theta=theta, freqs_dtype=freqs_dtype)
+            freqs_real.append((cos_emb, sin_emb))
+        return freqs_real
+
+    def _get_freqs_real(
+        self, freqs_real: list[tuple[torch.Tensor, torch.Tensor]], ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = ids.device
+        if ids.device.type == "mps":
+            ids = ids.to("cpu")
+
+        cos_result = []
+        sin_result = []
+        for i in range(len(self.axes_dim)):
+            freqs_cos, freqs_sin = freqs_real[i]
+            freqs_cos = freqs_cos.to(ids.device)
+            freqs_sin = freqs_sin.to(ids.device)
+            index = ids[:, :, i : i + 1].repeat(1, 1, freqs_cos.shape[-1]).to(torch.int64)
+            cos_result.append(torch.gather(freqs_cos.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
+            sin_result.append(torch.gather(freqs_sin.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
+
+        combined_cos = torch.cat(cos_result, dim=-1).to(device)
+        combined_sin = torch.cat(sin_result, dim=-1).to(device)
+        return combined_cos, combined_sin
+
+    def forward(
+        self, freqs_real, attention_mask, l_effective_ref_img_len, l_effective_img_len, ref_img_sizes, img_sizes, device
+    ):
+        batch_size = len(attention_mask)
+        p = self.patch_size
+
+        encoder_seq_len = attention_mask.shape[1]
+        l_effective_cap_len = attention_mask.sum(dim=1).tolist()
+        l_effective_cap_len = [int(len) for len in l_effective_cap_len]
+        seq_lengths = [
+            int(cap_len + sum(ref_img_len) + img_len)
+            for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len)
+        ]
+
+        max_seq_len = max(seq_lengths)
+        max_ref_img_len = max([int(sum(ref_img_len)) for ref_img_len in l_effective_ref_img_len])
+        max_img_len = int(max(l_effective_img_len))
+
+        # Create position IDs
+        position_ids = torch.zeros(batch_size, int(max_seq_len), 3, dtype=torch.int32, device=device)
+
+        for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)):
+            # add text position ids
+            position_ids[i, :cap_seq_len] = repeat(
+                torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3"
+            )
+            pe_shift = cap_seq_len
+            pe_shift_len = cap_seq_len
+
+            if ref_img_sizes[i] is not None:
+                for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]):
+                    H, W = ref_img_size
+                    ref_H_tokens, ref_W_tokens = H // p, W // p
+                    assert ref_H_tokens * ref_W_tokens == ref_img_len
+                    # add image position ids
+
+                    row_ids = repeat(
+                        torch.arange(ref_H_tokens, dtype=torch.int32, device=device), "h -> h w", w=ref_W_tokens
+                    ).flatten()
+                    col_ids = repeat(
+                        torch.arange(ref_W_tokens, dtype=torch.int32, device=device), "w -> h w", h=ref_H_tokens
+                    ).flatten()
+                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = pe_shift
+                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = row_ids
+                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = col_ids
+
+                    pe_shift += max(ref_H_tokens, ref_W_tokens)
+                    pe_shift_len += ref_img_len
+
+            H, W = img_sizes[i]
+            H_tokens, W_tokens = H // p, W // p
+            assert H_tokens * W_tokens == l_effective_img_len[i]
+
+            row_ids = repeat(torch.arange(H_tokens, dtype=torch.int32, device=device), "h -> h w", w=W_tokens).flatten()
+            col_ids = repeat(torch.arange(W_tokens, dtype=torch.int32, device=device), "w -> h w", h=H_tokens).flatten()
+
+            assert pe_shift_len + l_effective_img_len[i] == seq_len
+            position_ids[i, pe_shift_len:seq_len, 0] = pe_shift
+            position_ids[i, pe_shift_len:seq_len, 1] = row_ids
+            position_ids[i, pe_shift_len:seq_len, 2] = col_ids
+
+        # Get combined rotary embeddings (real version)
+        freqs_cos, freqs_sin = self._get_freqs_real(freqs_real, position_ids)
+
+        # create separate rotary embeddings for captions and images
+        cap_freqs_cos = torch.zeros(
+            batch_size, encoder_seq_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype
+        )
+        cap_freqs_sin = torch.zeros(
+            batch_size, encoder_seq_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype
+        )
+        ref_img_freqs_cos = torch.zeros(
+            batch_size, max_ref_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype
+        )
+        ref_img_freqs_sin = torch.zeros(
+            batch_size, max_ref_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype
+        )
+        img_freqs_cos = torch.zeros(batch_size, max_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype)
+        img_freqs_sin = torch.zeros(batch_size, max_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype)
+
+        for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(
+            zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths)
+        ):
+            cap_freqs_cos[i, :cap_seq_len] = freqs_cos[i, :cap_seq_len]
+            cap_freqs_sin[i, :cap_seq_len] = freqs_sin[i, :cap_seq_len]
+            ref_img_freqs_cos[i, : sum(ref_img_len)] = freqs_cos[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
+            ref_img_freqs_sin[i, : sum(ref_img_len)] = freqs_sin[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
+            img_freqs_cos[i, :img_len] = freqs_cos[
+                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
+            ]
+            img_freqs_sin[i, :img_len] = freqs_sin[
+                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
+            ]
+
+        return (
+            (cap_freqs_cos, cap_freqs_sin),
+            (ref_img_freqs_cos, ref_img_freqs_sin),
+            (img_freqs_cos, img_freqs_sin),
+            (freqs_cos, freqs_sin),
+            l_effective_cap_len,
+            seq_lengths,
+        )

vllm_omni/diffusion/models/mammoth_moda2/schedulers.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by ByteDance Ltd. on 2025-09-30.
+#
+# Original file was released under Apache-2.0, with the full license text
+# available at https://www.apache.org/licenses/LICENSE-2.0
+#
+# This modified file is released under the same license.
+#
+# --- Upstream header preserved below ---
+#
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. 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
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+
+@dataclass
+class FlowMatchEulerDiscreteSchedulerOutput:
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchEulerDiscreteScheduler:
+    order = 1
+
+    def __init__(self, num_train_timesteps: int = 1000, dynamic_time_shift: bool = True):
+        self.num_train_timesteps = int(num_train_timesteps)
+        self.dynamic_time_shift = bool(dynamic_time_shift)
+
+        timesteps = torch.linspace(0, 1, self.num_train_timesteps + 1, dtype=torch.float32)[:-1]
+        self.timesteps = timesteps
+        self._timesteps = torch.cat([timesteps, torch.ones(1, dtype=timesteps.dtype)])
+
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
+
+    @property
+    def step_index(self) -> int | None:
+        return self._step_index
+
+    @property
+    def begin_index(self) -> int | None:
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        self._begin_index = int(begin_index)
+
+    def index_for_timestep(self, timestep: torch.Tensor, schedule_timesteps: torch.Tensor | None = None) -> int:
+        schedule_timesteps = self._timesteps if schedule_timesteps is None else schedule_timesteps
+        indices = (schedule_timesteps == timestep).nonzero()
+        pos = 1 if len(indices) > 1 else 0
+        return int(indices[pos].item())
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device | None = None,
+        timesteps: list[float] | None = None,
+        num_tokens: int | None = None,
+    ) -> None:
+        if timesteps is None:
+            if num_inference_steps is None:
+                raise ValueError("`num_inference_steps` must be provided when `timesteps` is None.")
+            timesteps_np = np.linspace(0, 1, num_inference_steps + 1, dtype=np.float32)[:-1]
+            if self.dynamic_time_shift and num_tokens is not None:
+                m = np.sqrt(float(num_tokens)) / 40.0
+                timesteps_np = timesteps_np / (m - m * timesteps_np + timesteps_np)
+        else:
+            timesteps_np = np.asarray(timesteps, dtype=np.float32)
+
+        timesteps_t = torch.from_numpy(timesteps_np).to(dtype=torch.float32, device=device)
+        self.timesteps = timesteps_t
+        self._timesteps = torch.cat([timesteps_t, torch.ones(1, device=timesteps_t.device, dtype=timesteps_t.dtype)])
+
+        self._step_index = None
+        self._begin_index = None
+
+    def _init_step_index(self, timestep: float | torch.Tensor) -> None:
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            else:
+                timestep = torch.tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        generator: torch.Generator | None = None,  # noqa: ARG002 - kept for API compatibility
+        return_dict: bool = True,
+    ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor]:
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError(
+                "Integer timesteps (e.g. a loop index from enumerate) are not supported; "
+                "pass a value from `scheduler.timesteps` instead.",
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+        assert self._step_index is not None
+
+        sample_fp32 = sample.to(torch.float32)
+        t = self._timesteps[self._step_index]
+        t_next = self._timesteps[self._step_index + 1]
+
+        prev_sample = sample_fp32 + (t_next - t) * model_output
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self) -> int:
+        return self.num_train_timesteps