rosetta/modeling.py

import math
from argparse import Namespace
from contextlib import nullcontext
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention.flex_attention import BlockMask, flex_attention
from transformers.cache_utils import Cache, StaticCache
from transformers.utils.generic import ModelOutput

from .autoencoder import (
    TimestepEmbedder,
    normal_weight_reset_parameters,
    project_in_layer,
    project_out_layer,
)
from .configuration import MultimodalConfig, TransformerConfig
from .utils import PRECISION_TO_TYPE, default, is_package_version


@dataclass
class MultimodalModelOutput(ModelOutput):
    losses: Optional[dict] = None
    logits: Optional[torch.Tensor] = None
    past_key_values: Optional[Cache] = None
    diffusion_prediction: Optional[torch.Tensor] = None


class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6, device=None, dtype=None):
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size, **factory_kwargs))
        self.variance_epsilon = eps

    def reset_parameters(self):
        nn.init.ones_(self.weight)

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"

def get_vision_position_ids(
        start_position: int,
        grid_thw: List[int],
        spatial_merge_size: int = 1,
        device: Optional[torch.device] = None,
) -> torch.Tensor:
    t, h, w = grid_thw[0], grid_thw[1], grid_thw[2]
    if isinstance(t, torch.Tensor):
        t, h, w = t.item(), h.item(), w.item()
    llm_grid_t = t
    llm_grid_h = h // spatial_merge_size
    llm_grid_w = w // spatial_merge_size

    image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t
    position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat(
        llm_grid_h * llm_grid_t
    )
    position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave(
        llm_grid_w * llm_grid_t
    )
    position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long)
    vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0)
    return vision_position_ids


def get_text_position_ids(
        length: int,
        start_position: int = 0,
        device: Optional[torch.device] = None,
) -> torch.Tensor:
    return torch.arange(length, device=device).view(1, -1).expand(3, -1) + start_position


def get_interleaved_mrope_index(
        image_infos: List[Optional[List[Tuple[slice, Tuple[int, int], dict]]]],
        seq_len: int,
        spatial_merge_size: int,
        sample_offsets: Optional[torch.Tensor] = None,
        device: Optional[torch.device] = None,
) -> torch.Tensor:
    if image_infos is None:
        image_infos = [None]
    batch_size = len(image_infos)
    position_ids = torch.zeros(3, batch_size, seq_len, dtype=torch.int64, device=device)

    for i, image_info in enumerate(image_infos):
        llm_pos_ids_list = []
        current_pos = 0
        st = 0

        if image_info is None:
            image_info = []

        for sec_slice, (h, w), _ in image_info:
            img_start = sec_slice.start
            text_len = img_start - st

            if text_len > 0:
                llm_pos_ids_list.append(get_text_position_ids(text_len, current_pos, device))
                current_pos += text_len

            grid_thw = [1, h, w]
            llm_pos_ids_list.append(
                get_vision_position_ids(current_pos, grid_thw, spatial_merge_size=spatial_merge_size, device=device)
            )
            llm_grid_h = h // spatial_merge_size
            llm_grid_w = w // spatial_merge_size
            current_pos += max(llm_grid_h, llm_grid_w)
            st = img_start + llm_grid_h * llm_grid_w

        if st < seq_len:
            llm_pos_ids_list.append(get_text_position_ids(seq_len - st, current_pos, device))

        if len(llm_pos_ids_list) == 0:
            llm_pos_ids_list.append(get_text_position_ids(seq_len, 0, device))

        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)

        if llm_positions.shape[1] < seq_len:
            padding = llm_positions[:, -1:].expand(3, seq_len - llm_positions.shape[1])
            llm_positions = torch.cat([llm_positions, padding], dim=1)
        elif llm_positions.shape[1] > seq_len:
            llm_positions = llm_positions[:, :seq_len]

        position_ids[:, i, :] = llm_positions

        if sample_offsets is not None and sample_offsets[i] is not None:
            offsets = sample_offsets[i].tolist()
            if len(offsets) >= 2:
                assert offsets[0] == 0, "First offset must be 0"
                assert offsets[-1] <= seq_len, "Last offset must be less than or equal to seq_len"
                for start, end in zip(offsets[:-1], offsets[1:]):
                    assert end > start, "End must be greater than start"
                    seg_base = position_ids[:, i, start].clone()
                    position_ids[:, i, start:end] = position_ids[:, i, start:end] - seg_base.view(3, 1)


    return position_ids


def apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: List[int]) -> torch.Tensor:
    freqs_t = freqs[0].clone()
    for dim, offset in enumerate((1, 2), start=1):
        length = mrope_section[dim] * 3
        idx = slice(offset, length, 3)
        freqs_t[..., idx] = freqs[dim, ..., idx]
    return freqs_t


def get_batch_interleaved_mrope(
        image_infos: List[List[Tuple[slice, Tuple[int, int]]]],
        seq_len: int,
        n_elem: int,
        mrope_section: List[int],
        device: Optional[torch.device] = None,
        base: float = 10000.0,
        base_rescale_factor: float = 1.0,
        spatial_merge_size: int = 1,
        sample_offsets: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
    position_ids = get_interleaved_mrope_index(image_infos, seq_len, spatial_merge_size, sample_offsets, device)
    position_ids = position_ids.to(device)

    if base_rescale_factor != 1.0:
        base *= base_rescale_factor ** (n_elem / (n_elem - 2))

    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem))

    inv_freq_expanded = theta[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
    position_ids_expanded = position_ids[:, :, None, :].float()

    device_type = device.type if isinstance(device.type, str) and device.type != "mps" else "cpu"
    with torch.autocast(device_type=device_type, enabled=False):
        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
        freqs = apply_interleaved_mrope(freqs, mrope_section)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()

    return cos, sin


def apply_rope(
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        unsqueeze_dim=-3,
) -> torch.Tensor:
    head_size = x.size(-1)
    x1 = x[..., : head_size // 2]
    x2 = x[..., head_size // 2:]
    rotated = torch.cat((-x2, x1), dim=-1)
    if cos.dim() > 1:
        cos = cos.unsqueeze(unsqueeze_dim)
        sin = sin.unsqueeze(unsqueeze_dim)

    return (x * cos) + (rotated * sin)

class CachedRoPE(object):
    def __init__(self, config):
        self._config = config
        self.cos_cache = None
        self.sin_cache = None
        self.seq_len = None
        self.rope_media_info = None

    @torch.autocast(device_type='cuda', enabled=False)
    def __call__(self, seq_len, device, rope_media_info=None, input_pos=None, sample_offsets=None):
        if (self.seq_len != seq_len) or (self.rope_media_info != rope_media_info):
            self.cos_cache, self.sin_cache = get_batch_interleaved_mrope(
                image_infos=rope_media_info,
                seq_len=seq_len,
                mrope_section=self._config.mrope_section,
                n_elem=self._config.attention_head_size,
                device=device,
                base=self._config.rope_theta,
                base_rescale_factor=1.0,
                sample_offsets=sample_offsets
            )
        if input_pos is None:
            cos, sin = self.cos_cache, self.sin_cache
        else:
            assert input_pos.dim() == 2, f"{input_pos.shape=}"
            head_size = self.cos_cache.size(-1)
            cos = torch.gather(self.cos_cache, dim=1, index=input_pos.unsqueeze(-1).expand(-1, -1, head_size))
            sin = torch.gather(self.sin_cache, dim=1, index=input_pos.unsqueeze(-1).expand(-1, -1, head_size))

        return cos, sin

# Attention and KV cache layers
flex_attention = torch.compile(flex_attention, dynamic=False)


BatchRaggedMedia = Union[torch.Tensor, list[Union[torch.Tensor, list[torch.Tensor]]]]
BatchRaggedTensor = Union[torch.Tensor, list[torch.Tensor]]


class MultimodalStaticCache(StaticCache):
    def __init__(self, *args, **kwargs):
        self.dynamic = kwargs.pop("dynamic", False)
        super().__init__(*args, **kwargs)

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        cache_position = cache_kwargs.get("cache_position")
        if is_package_version("transformers", "<=", "4.53.3"):
            if self.key_cache[layer_idx].device != key_states.device:
                self.key_cache[layer_idx] = self.key_cache[layer_idx].to(key_states.device)
                self.value_cache[layer_idx] = self.value_cache[layer_idx].to(value_states.device)
            k_out = self.key_cache[layer_idx]  # max_batch_size x num_key_value_heads x max_cache_len x head_dim
            v_out = self.value_cache[layer_idx]  # max_batch_size x num_key_value_heads x max_cache_len x head_dim
        else:
            if self.layers[layer_idx].keys is None:
                self.layers[layer_idx].lazy_initialization(value_states)
            k_out = self.layers[layer_idx].keys
            v_out = self.layers[layer_idx].values
        key_states = key_states.to(k_out.dtype)
        value_states = value_states.to(v_out.dtype)

        if cache_position is None:
            k_out.copy_(key_states)
            v_out.copy_(value_states)
        else:
            if cache_position.dim() == 1:
                k_out.index_copy_(2, cache_position, key_states)
                v_out.index_copy_(2, cache_position, value_states)

                if self.dynamic:
                    end = cache_position[-1].item() + 1
                    k_out = k_out[:, :, :end]
                    v_out = v_out[:, :, :end]
            else:
                # first step of gen_text or all steps of gen_image
                assert cache_position.dim() == 2, f"multiple batch dims not yet {cache_position.shape=}"
                batch_size, idx_size = cache_position.shape
                assert batch_size == k_out.size(0)
                assert batch_size == v_out.size(0)
                assert batch_size == key_states.size(0)
                assert batch_size == value_states.size(0)
                for i in range(batch_size):
                    unbatched_dim = 1
                    k_out[i].index_copy_(unbatched_dim, cache_position[i], key_states[i])
                    v_out[i].index_copy_(unbatched_dim, cache_position[i], value_states[i])

                if self.dynamic:
                    assert len(cache_position) == 1
                    # end = cache_position[0, -1].item() + 1
                    end = int(cache_position[0, -1]) + 1  # Tensor.item()导致图中断
                    k_out = k_out[:, :, :end]
                    v_out = v_out[:, :, :end]

        return k_out, v_out


def get_device(tensor: BatchRaggedMedia):
    if isinstance(tensor, torch.Tensor):
        return tensor.device
    elif isinstance(tensor, list):
        return get_device(tensor[0])
    else:
        raise ValueError(f"Unsupported type for get_device: {type(tensor)}")


class CausalSelfAttention(nn.Module):
    def __init__(
            self,
            config: MultimodalConfig,
            layer_idx: int,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
    ) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        self._config = config
        self.layer_idx = layer_idx

        self.q_proj = nn.Linear(
            config.hidden_size, config.num_attention_heads * config.attention_head_size,
            bias=False, **factory_kwargs
        )
        self.k_proj = nn.Linear(
            config.hidden_size, config.num_kv_heads * config.attention_head_size,
            bias=False, **factory_kwargs
        )
        self.v_proj = nn.Linear(
            config.hidden_size, config.num_kv_heads * config.attention_head_size,
            bias=False, **factory_kwargs
        )

        self.o_proj = nn.Linear(
            config.attention_head_size * config.num_attention_heads, config.hidden_size,
            bias=False, **factory_kwargs
        )

        self.query_layernorm = config.norm_class(config.attention_head_size, eps=config.norm_eps, **factory_kwargs)
        self.key_layernorm = config.norm_class(config.attention_head_size, eps=config.norm_eps, **factory_kwargs)

    def forward(
            self,
            hidden_states: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            rotary_position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
            input_pos: Optional[torch.Tensor] = None,
            past_key_values: Optional[MultimodalStaticCache] = None,
    ) -> torch.Tensor:
        bsz, seqlen, _ = hidden_states.size()  # batch size, sequence length, embedding dimensionality (n_embd)
        head_size = self._config.attention_head_size
        n_q_head = self._config.num_attention_heads
        n_kv_head = self._config.num_kv_heads
        q_per_kv = n_q_head // n_kv_head

        new_q = self.q_proj(hidden_states).view(bsz, seqlen, n_kv_head, q_per_kv, head_size)
        new_k = self.k_proj(hidden_states).view(bsz, seqlen, n_kv_head, 1, head_size)
        new_v = self.v_proj(hidden_states).view(bsz, seqlen, n_kv_head, 1, head_size)
        q, k, v = map(lambda x: x.permute(0, 2, 3, 1, 4), [new_q, new_k, new_v])

        q = q.reshape(bsz, -1, seqlen, head_size)  # (B, n_q_head, T, hs)
        k = k.reshape(bsz, -1, seqlen, head_size)  # (B, n_kv_head, T, hs)
        v = v.reshape(bsz, -1, seqlen, head_size)  # (B, n_kv_head, T, hs)

        q = self.query_layernorm(q)
        k = self.key_layernorm(k)

        q = apply_rope(q, *rotary_position_embeddings)
        k = apply_rope(k, *rotary_position_embeddings)

        q = q.to(v.dtype)
        k = k.to(v.dtype)

        if input_pos is not None:
            cache_kwargs = {"cache_position": input_pos}
            k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)
        # If restore from cache, kv_seqlen >= seqlen
        kv_seqlen = k.size(2)

        # maybe repeat k and v if for the non multi-head attention cases
        # training: flash attention requires it
        # inference: multi-query would require a full kv cache so avoid it to limit its memory usage
        if n_kv_head != n_q_head and (input_pos is None or q_per_kv != 1):
            k = k.unsqueeze(dim=2).expand(-1, -1, q_per_kv, -1, -1).reshape(bsz, -1, kv_seqlen, head_size)
            v = v.unsqueeze(dim=2).expand(-1, -1, q_per_kv, -1, -1).reshape(bsz, -1, kv_seqlen, head_size)

        y = self.scaled_dot_product_attention(q, k, v, attention_mask)

        y = y.reshape(bsz, seqlen, head_size * n_q_head)  # re-assemble all head outputs side by side

        # output projection
        return self.o_proj(y)

    def scaled_dot_product_attention(
            self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        # q, k, v: (bsz, n_head, seqlen, head_size)
        scale = 1.0 / math.sqrt(self._config.attention_head_size)

        if isinstance(mask, BlockMask):
            q = q.to(dtype=v.dtype)
            k = k.to(dtype=v.dtype)
            y = flex_attention(q, k, v, block_mask=mask, scale=scale)
        else:
            y = torch.nn.functional.scaled_dot_product_attention(
                q, k, v, attn_mask=mask, dropout_p=0.0, scale=scale,
                # If q only has one token (typically in AR model decoding stage), we should use full attention.
                is_causal=mask is None and q.size(2) > 1
            )

        return y.transpose(1, 2)


class CausalSelfAttentionMoT(CausalSelfAttention):
    def __init__(
            self,
            config: MultimodalConfig,
            config_mot_gen: MultimodalConfig,
            layer_idx: int,
            mot_und_frozen: bool = False,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
    ) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__(config, layer_idx, dtype, device)
        self.mot_und_frozen = mot_und_frozen

        self.q_proj_mot_gen = nn.Linear(
            config_mot_gen.hidden_size, config_mot_gen.num_attention_heads * config_mot_gen.attention_head_size,
            bias=False, **factory_kwargs
        )
        self.k_proj_mot_gen = nn.Linear(
            config_mot_gen.hidden_size, config_mot_gen.num_kv_heads * config_mot_gen.attention_head_size,
            bias=False, **factory_kwargs
        )
        self.v_proj_mot_gen = nn.Linear(
            config_mot_gen.hidden_size, config_mot_gen.num_kv_heads * config_mot_gen.attention_head_size,
            bias=False, **factory_kwargs
        )

        self.o_proj_mot_gen = nn.Linear(
            config_mot_gen.attention_head_size * config_mot_gen.num_attention_heads, config_mot_gen.hidden_size,
            bias=False, **factory_kwargs
        )

        self.query_layernorm_mot_gen = config_mot_gen.norm_class(config_mot_gen.attention_head_size, eps=config_mot_gen.norm_eps, **factory_kwargs)
        self.key_layernorm_mot_gen = config_mot_gen.norm_class(config_mot_gen.attention_head_size, eps=config_mot_gen.norm_eps, **factory_kwargs)

        if self.mot_und_frozen:
            self.q_proj.eval()
            self.q_proj.requires_grad_(False)
            self.k_proj.eval()
            self.k_proj.requires_grad_(False)
            self.v_proj.eval()
            self.v_proj.requires_grad_(False)
            self.o_proj.eval()
            self.o_proj.requires_grad_(False)
            self.query_layernorm.eval()
            self.query_layernorm.requires_grad_(False)
            self.key_layernorm.eval()
            self.key_layernorm.requires_grad_(False)

    def forward(
        self,
        hidden_states: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        rotary_position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
        input_pos: Optional[torch.Tensor] = None,
        past_key_values: Optional[MultimodalStaticCache] = None,
        und_token_indices: Optional[torch.Tensor] = None,
        gen_token_indices: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        und_hidden_states, gen_hidden_states = hidden_states

        bsz = und_hidden_states.shape[0]
        und_seqlen, gen_seqlen = und_hidden_states.shape[1], gen_hidden_states.shape[1]
        head_size = self._config.attention_head_size
        n_q_head = self._config.num_attention_heads
        n_kv_head = self._config.num_kv_heads
        q_per_kv = n_q_head // n_kv_head

        q = self.q_proj(und_hidden_states).view(bsz, -1, n_kv_head, q_per_kv, head_size)
        k = self.k_proj(und_hidden_states).view(bsz, -1, n_kv_head, 1, head_size)
        v = self.v_proj(und_hidden_states).view(bsz, -1, n_kv_head, 1, head_size)

        gen_q = self.q_proj_mot_gen(gen_hidden_states).view(bsz, -1, n_kv_head, q_per_kv, head_size)
        gen_k = self.k_proj_mot_gen(gen_hidden_states).view(bsz, -1, n_kv_head, 1, head_size)
        gen_v = self.v_proj_mot_gen(gen_hidden_states).view(bsz, -1, n_kv_head, 1, head_size)

        q, k, v = map(lambda x: x.permute(0, 2, 3, 1, 4), [q, k, v])
        gen_q, gen_k, gen_v = map(lambda x: x.permute(0, 2, 3, 1, 4), [gen_q, gen_k, gen_v])

        # [bsz, h, seqlen, head_size]
        q = q.reshape(bsz, n_q_head, und_seqlen, head_size)
        k = k.reshape(bsz, n_kv_head, und_seqlen, head_size)
        v = v.reshape(bsz, n_kv_head, und_seqlen, head_size)
        gen_q = gen_q.reshape(bsz, n_q_head, gen_seqlen, head_size)
        gen_k = gen_k.reshape(bsz, n_kv_head, gen_seqlen, head_size)
        gen_v = gen_v.reshape(bsz, n_kv_head, gen_seqlen, head_size)

        # Scatter understanding and generation tokens for rope
        und_token_indices_q = und_token_indices.unsqueeze(-1).unsqueeze(1).expand(-1, q.size(1), -1, q.size(-1))
        gen_token_indices_q = gen_token_indices.unsqueeze(-1).unsqueeze(1).expand(-1, q.size(1), -1, q.size(-1))
        und_token_indices_kv = und_token_indices.unsqueeze(-1).unsqueeze(1).expand(-1, k.size(1), -1, k.size(-1))
        gen_token_indices_kv = gen_token_indices.unsqueeze(-1).unsqueeze(1).expand(-1, k.size(1), -1, k.size(-1))

        def _scatter(und_src, gen_src, und_token_indices, gen_token_indices, n_head):
            target = torch.zeros((bsz, n_head, und_seqlen+gen_seqlen, head_size), dtype=und_src.dtype, device=und_src.device)
            target.scatter_(dim=2, index=und_token_indices, src=und_src)
            target.scatter_(dim=2, index=gen_token_indices, src=gen_src)
            return target

        q_merge = _scatter(q, gen_q, und_token_indices_q, gen_token_indices_q, n_q_head)
        k_merge = _scatter(k, gen_k, und_token_indices_kv, gen_token_indices_kv, n_kv_head)
        v_merge = _scatter(v, gen_v, und_token_indices_kv, gen_token_indices_kv, n_kv_head)

        q_ = torch.zeros_like(q_merge)
        k_ = torch.zeros_like(k_merge)

        q_.scatter_(dim=2, index=und_token_indices_q, src=self.query_layernorm(q_merge.gather(2, und_token_indices_q)).to(q_.dtype))
        q_.scatter_(dim=2, index=gen_token_indices_q, src=self.query_layernorm_mot_gen(q_merge.gather(2, gen_token_indices_q)).to(q_.dtype))
        k_.scatter_(dim=2, index=und_token_indices_kv, src=self.key_layernorm(k_merge.gather(2, und_token_indices_kv)).to(k_.dtype))
        k_.scatter_(dim=2, index=gen_token_indices_kv, src=self.key_layernorm_mot_gen(k_merge.gather(2, gen_token_indices_kv)).to(k_.dtype))

        q_merge = q_
        k_merge = k_

        # apply rotary position embeddings
        q_merge = apply_rope(q_merge, *rotary_position_embeddings)
        k_merge = apply_rope(k_merge, *rotary_position_embeddings)

        q_merge = q_merge.to(v_merge.dtype)
        k_merge = k_merge.to(v_merge.dtype)

        # Restore from kv_cache and update
        if input_pos is not None:
            cache_kwargs = {"cache_position": input_pos}
            k_merge, v_merge = past_key_values.update(k_merge, v_merge, self.layer_idx, cache_kwargs)
        # If restore from cache, kv_seqlen >= seqlen
        kv_seqlen = k_merge.size(2)

        # maybe repeat k and v if for the non multi-head attention cases
        # training: flash attention requires it
        # inference: multi-query would require a full kv cache so avoid it to limit its memory usage
        if n_kv_head != n_q_head and (input_pos is None or q_per_kv != 1):
            k_merge = k_merge.unsqueeze(dim=2).expand(-1, -1, q_per_kv, -1, -1).reshape(bsz, -1, kv_seqlen, head_size)
            v_merge = v_merge.unsqueeze(dim=2).expand(-1, -1, q_per_kv, -1, -1).reshape(bsz, -1, kv_seqlen, head_size)

        y = self.scaled_dot_product_attention(q_merge, k_merge, v_merge, attention_mask)

        # re-assemble all head outputs side by side
        y = y.reshape(bsz, -1, head_size * n_q_head)

        core_attn_out = y.gather(dim=1, index=und_token_indices.unsqueeze(-1).expand(-1, -1, y.size(-1)))
        gen_core_attn_out = y.gather(dim=1, index=gen_token_indices.unsqueeze(-1).expand(-1, -1, y.size(-1)))

        und_hidden_states = self.o_proj(core_attn_out)
        gen_hidden_states = self.o_proj_mot_gen(gen_core_attn_out)
        return und_hidden_states, gen_hidden_states

class MLP(nn.Module):
    def __init__(
            self,
            config: TransformerConfig,
            layer_idx: int,
            is_shared_mlp: bool = False,
            is_moe: bool = False,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
    ):
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self._config = config
        self.layer_idx = layer_idx
        self.ffn_hidden_size = config.ffn_hidden_size

        if is_shared_mlp or is_moe:
            self.ffn_hidden_size = config.moe_ffn_hidden_size

        self.gate_proj = nn.Linear(config.hidden_size, self.ffn_hidden_size, bias=False, **factory_kwargs)
        self.up_proj = nn.Linear(config.hidden_size, self.ffn_hidden_size, bias=False, **factory_kwargs)
        self.down_proj = nn.Linear(self.ffn_hidden_size, config.hidden_size, bias=False, **factory_kwargs)
        self.act_fn = config.act_class()

    def forward(self, x):
        up = self.up_proj(x)
        gate = self.gate_proj(x)
        out = self.down_proj(up * self.act_fn(gate))
        return out


class DeepSeekMoEGate(nn.Module):
    def __init__(
            self,
            config: TransformerConfig,
            layer_idx: int,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
    ):
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self._config = config
        self.layer_idx = layer_idx
        self.num_experts = config.num_experts
        self.top_k = config.moe_topk

        if config.use_modality_routing:
            assert config.num_text_experts > 0 and config.num_vit_experts > 0 and config.num_vae_experts > 0, \
                "num_text_experts / num_vit_experts / num_vae_experts must be set when use_modality_routing=True"
            assert config.num_text_experts + config.num_vit_experts + config.num_vae_experts == self.num_experts, \
                f"num_text({config.num_text_experts}) + num_vit({config.num_vit_experts}) + num_vae({config.num_vae_experts}) " \
                f"must equal num_experts({self.num_experts})"
            self.wg_text = nn.Linear(config.hidden_size, config.num_text_experts, bias=False, **factory_kwargs)
            self.wg_vit  = nn.Linear(config.hidden_size, config.num_vit_experts,  bias=False, **factory_kwargs)
            self.wg_vae  = nn.Linear(config.hidden_size, config.num_vae_experts,  bias=False, **factory_kwargs)
        else:
            self.wg = nn.Linear(config.hidden_size, self.num_experts, bias=False, **factory_kwargs)

    def _score(self, logits):
        return F.softmax(logits.float(), dim=1)

    def _forward_modality_routing(self, flat_hidden, bsz, seqlen, token_modalities):
        N = bsz * seqlen
        cfg = self._config
        n_text = cfg.num_text_experts
        n_vit  = cfg.num_vit_experts
        n_vae  = cfg.num_vae_experts
        vit_offset = n_text
        vae_offset = n_text + n_vit

        flat_mod = token_modalities.reshape(-1)

        topk_weights = torch.zeros(N, self.top_k, dtype=torch.float32, device=flat_hidden.device)
        topk_idx     = torch.zeros(N, self.top_k, dtype=torch.long,    device=flat_hidden.device)

        for mod_id, wg, offset, n_local in (
            (0, self.wg_text, 0,          n_text),
            (1, self.wg_vit,  vit_offset, n_vit),
            (2, self.wg_vae,  vae_offset, n_vae),
        ):
            mask = (flat_mod == mod_id)
            if not mask.any():
                continue
            h_m      = flat_hidden[mask].to(wg.weight.dtype)
            logits_m = wg(h_m)
            scores_m = self._score(logits_m)

            tw_m, ti_m = torch.topk(scores_m, self.top_k, dim=-1)
            topk_weights[mask] = tw_m
            topk_idx[mask]     = ti_m + offset

        if self.top_k > 1:
            denom = topk_weights.sum(dim=-1, keepdim=True).clamp(
                min=torch.finfo(topk_weights.dtype).eps
            )
            topk_weights = topk_weights / denom

        return topk_weights, topk_idx

    def forward(self, hidden_states, token_modalities=None):
        bsz, seqlen, hdim = hidden_states.size()
        hidden_states = hidden_states.reshape(-1, hidden_states.size(-1))

        if self._config.use_modality_routing:
            if token_modalities is None:
                N = hidden_states.shape[0]
                token_modalities = torch.zeros(N, dtype=torch.long, device=hidden_states.device)
            return self._forward_modality_routing(hidden_states, bsz, seqlen, token_modalities)

        logits = self.wg(hidden_states.to(self.wg.weight.dtype))

        scores = F.softmax(logits, dim=1)
        topk_weights, topk_idx = torch.topk(scores, self.top_k, dim=-1)

        if self.top_k > 1:
            denominator = topk_weights.sum(dim=-1, keepdim=True).clamp(min=torch.finfo(topk_weights.dtype).eps)
            topk_weights = topk_weights / denominator

        return topk_weights, topk_idx


class DeepSeekMoE(nn.Module):
    def __init__(
            self,
            config: TransformerConfig,
            layer_idx: Optional[int] = None,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
    ):
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self._config = config
        self.layer_idx = layer_idx
        self.num_experts = config.num_experts
        self.top_k = config.moe_topk
        self.gate = DeepSeekMoEGate(config, layer_idx, device=device, dtype=torch.float32)

        self.experts = nn.ModuleList(
            [MLP(config, layer_idx, is_moe=True, **factory_kwargs)
             for _ in range(self.num_experts)]
        )
        self.shared_mlp = MLP(config, layer_idx=layer_idx, is_shared_mlp=True, **factory_kwargs)

    def forward(self, hidden_states: torch.Tensor, token_modalities: Optional[torch.Tensor] = None) -> torch.Tensor:
        bsz, seqlen, hdim = hidden_states.size()
        input_hidden_states = hidden_states

        with torch.autocast('cuda', enabled=False):
            topk_weights, topk_idx = self.gate(hidden_states, token_modalities)

        topk_weights = topk_weights.to(hidden_states.dtype)

        flat_topk_idx = topk_idx.view(-1)
        hidden_states = hidden_states.view(-1, hdim)
        hidden_states = hidden_states.repeat_interleave(self.top_k, dim=0)

        expert_outputs = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device)
        for i in range(self.num_experts):
            expert_mask = (flat_topk_idx == i)
            selected_inputs = hidden_states[expert_mask]
            expert_output = self.experts[i](selected_inputs)
            expert_outputs[expert_mask] = expert_output.to(hidden_states.dtype)

        weighted_outputs = (expert_outputs.view(
            bsz * seqlen, self.top_k, hdim) * topk_weights.unsqueeze(-1)).sum(dim=1)
        weighted_outputs = weighted_outputs.to(hidden_states.dtype).view(bsz, seqlen, hdim)
        shared_out = self.shared_mlp(input_hidden_states)

        cfg = self._config
        if (
            token_modalities is not None
            and getattr(cfg, 'shield_step', 0) > 0
            and getattr(cfg, '_current_training_iter', 0) < cfg.shield_step
        ):
            text_mask = (token_modalities == 0).unsqueeze(-1)
            shared_out = shared_out * text_mask + shared_out.detach() * (~text_mask)

        return weighted_outputs + shared_out

class MultimodalDecoderLayer(nn.Module):
    def __init__(
            self,
            config: MultimodalConfig,
            layer_idx: int,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
    ):
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        self._config = config
        self.hidden_size = config.hidden_size
        self.layer_idx = layer_idx

        self.self_attn = CausalSelfAttention(config, layer_idx, **factory_kwargs)
        self.input_layernorm = config.norm_class(config.hidden_size, eps=config.norm_eps, **factory_kwargs)
        self.post_attention_layernorm = config.norm_class(config.hidden_size, eps=config.norm_eps, **factory_kwargs)
        self.mlp = DeepSeekMoE(config, layer_idx, **factory_kwargs)

    def forward(
            self,
            hidden_states: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            rotary_position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
            input_pos: Optional[torch.Tensor] = None,
            past_key_values: Optional[MultimodalStaticCache] = None,
            und_token_indices: Optional[torch.Tensor] = None,
            gen_token_indices: Optional[torch.Tensor] = None,
            token_modalities: Optional[torch.Tensor] = None,
    ):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(
            hidden_states,
            attention_mask=attention_mask,
            rotary_position_embeddings=rotary_position_embeddings,
            input_pos=input_pos,
            past_key_values=past_key_values,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states, token_modalities=token_modalities) \
            if token_modalities is not None and hasattr(self.mlp, 'gate') and hasattr(self.mlp.gate, 'wg_text') \
            else self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states

class MultimodalMoTDecoderLayer(MultimodalDecoderLayer):
    def __init__(
            self,
            config: MultimodalConfig,
            layer_idx: int,
            mot_und_frozen: bool = False,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
            gen_config: Optional[MultimodalConfig] = None,
    ):
        factory_kwargs = {'device': device, 'dtype': dtype}
        gen_config = gen_config if gen_config is not None else config
        super().__init__(config, layer_idx, dtype, device)
        self.mot_und_frozen = mot_und_frozen
        self.self_attn = CausalSelfAttentionMoT(config, gen_config, layer_idx, mot_und_frozen, **factory_kwargs)
        self.input_layernorm_mot_gen = gen_config.norm_class(
            gen_config.hidden_size, eps=gen_config.norm_eps, **factory_kwargs
        )
        self.post_attention_layernorm_mot_gen = gen_config.norm_class(
            gen_config.hidden_size, eps=gen_config.norm_eps, **factory_kwargs
        )

        self.mlp_mot_gen = DeepSeekMoE(gen_config, layer_idx, **factory_kwargs)

        if self.mot_und_frozen:
            self.input_layernorm.eval()
            self.input_layernorm.requires_grad_(False)
            self.post_attention_layernorm.eval()
            self.post_attention_layernorm.requires_grad_(False)
            self.mlp.eval()
            self.mlp.requires_grad_(False)

    @staticmethod
    def _attach_zero_param_dependency(hidden_states, modules):
        dep = None
        for module in modules:
            for param in module.parameters():
                if param.requires_grad and param.numel() > 0:
                    term = param.reshape(-1)[0] * 0.0
                    dep = term if dep is None else dep + term
        return hidden_states if dep is None else hidden_states + dep.to(hidden_states.dtype)

    def _und_branch_modules(self):
        return (
            self.input_layernorm,
            self.post_attention_layernorm,
            self.mlp,
            self.self_attn.q_proj,
            self.self_attn.k_proj,
            self.self_attn.v_proj,
            self.self_attn.o_proj,
            self.self_attn.query_layernorm,
            self.self_attn.key_layernorm,
        )

    def _gen_branch_modules(self):
        return (
            self.input_layernorm_mot_gen,
            self.post_attention_layernorm_mot_gen,
            self.mlp_mot_gen,
            self.self_attn.q_proj_mot_gen,
            self.self_attn.k_proj_mot_gen,
            self.self_attn.v_proj_mot_gen,
            self.self_attn.o_proj_mot_gen,
            self.self_attn.query_layernorm_mot_gen,
            self.self_attn.key_layernorm_mot_gen,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        rotary_position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
        input_pos: Optional[torch.Tensor] = None,
        past_key_values: Optional[MultimodalStaticCache] = None,
        und_token_indices: Optional[torch.Tensor] = None,
        gen_token_indices: Optional[torch.Tensor] = None,
        token_modalities: Optional[torch.Tensor] = None,
    ):
        und_hidden_states, gen_hidden_states = hidden_states
        und_residual, gen_residual = und_hidden_states, gen_hidden_states

        # Pre-attn norm
        und_hidden_states = self.input_layernorm(und_hidden_states)
        gen_hidden_states = self.input_layernorm_mot_gen(gen_hidden_states)

        # Self attention
        core_attn_out = self.self_attn(
            (und_hidden_states, gen_hidden_states),
            attention_mask=attention_mask,
            rotary_position_embeddings=rotary_position_embeddings,
            input_pos=input_pos,
            past_key_values=past_key_values,
            und_token_indices=und_token_indices,
            gen_token_indices=gen_token_indices,
        )

        und_hidden_states, gen_hidden_states = core_attn_out
        und_hidden_states = und_residual + und_hidden_states
        gen_hidden_states = gen_residual + gen_hidden_states

        # Pre-mlp norm
        und_residual, gen_residual = und_hidden_states, gen_hidden_states
        und_hidden_states = self.post_attention_layernorm(und_hidden_states)
        gen_hidden_states = self.post_attention_layernorm_mot_gen(gen_hidden_states)

        # Keep inactive MoT branch parameters connected to the loss with zero
        # gradients, so fixed-modality shard groups follow the same FSDP path.
        und_has_tokens = und_hidden_states.nelement() != 0
        gen_has_tokens = gen_hidden_states.nelement() != 0
        if und_has_tokens:
            und_hidden_states = self.mlp(und_hidden_states)
        if gen_has_tokens:
            gen_hidden_states = self.mlp_mot_gen(gen_hidden_states)
        if und_has_tokens and not gen_has_tokens:
            und_hidden_states = self._attach_zero_param_dependency(
                und_hidden_states, self._gen_branch_modules()
            )
        if gen_has_tokens and not und_has_tokens:
            gen_hidden_states = self._attach_zero_param_dependency(
                gen_hidden_states, self._und_branch_modules()
            )

        und_hidden_states = und_residual + und_hidden_states
        gen_hidden_states = gen_residual + gen_hidden_states

        return (und_hidden_states, gen_hidden_states)

class MultimodalModelBase(nn.Module):
    def get_input_embeddings(self):
        return self.model["embed_tokens"]

    def get_output_embeddings(self):
        return getattr(self, "lm_head", None)

    def __post_init__(
            self,
            config: MultimodalConfig,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
            args: Namespace = None,
            initialize_weights: bool = True,
    ):
        factory_kwargs = {'device': device, 'dtype': dtype}
        # Set as protected member to avoid conflict with potential parent classes
        self._config = config
        if self._config.use_mot:
            self._config_mot_gen = self._config.to_mot_gen_config()
        self._dtype = dtype

        # For inference, args can be None
        self.args = args or Namespace()
        self.vit_frozen = getattr(args, "vit_frozen", True)
        self.vit_precision = PRECISION_TO_TYPE[default(getattr(args, "vit_precision", None), dtype)]
        self.mot_und_frozen = getattr(args, "mot_und_frozen", False)
        self.lm_frozen = getattr(args, "lm_frozen", False)
        self.moe_aux_loss_coeff = getattr(args, "moe_aux_loss_coeff", 0.0)

        # ======================================
        #     Define vae projector modules
        # ======================================
        if config.use_vae:
            vae_hidden_size = self._config_mot_gen.hidden_size if config.use_mot else config.hidden_size
            vae_config = self._config_mot_gen if config.use_mot else config
            if config.use_timestep_token:
                self.timestep_emb = TimestepEmbedder(hidden_size=vae_hidden_size, **factory_kwargs)

            # One for patch_embed and other for final_layer
            self.time_embed = TimestepEmbedder(hidden_size=vae_hidden_size, **factory_kwargs)
            self.time_embed_2 = TimestepEmbedder(hidden_size=vae_hidden_size, **factory_kwargs)

            self.patch_embed = project_in_layer(vae_config, **factory_kwargs)
            self.final_layer = project_out_layer(vae_config, **factory_kwargs)

        # ======================================
        #     Define vit and aligner modules
        # ======================================
        if config.use_vit:
            from .visual_encoder import load_vit

            self.vit = load_vit(
                vision_model_type=config.vit_type,
                vision_model_precision=self.vit_precision,
                device=device,
                require_grad=not self.vit_frozen,
                eval_mode=self.vit_frozen,
                vision_model_params=config.vit_config,
                # Set True to skip default ViT loading if the main checkpoint already
                # contains the vision weights.
                no_load_pretrained=(
                    getattr(args, "vit_no_load_pretrained", False)
                    or (device is not None and torch.device(device).type == "meta")
                ),
            )
            self.vit_context = torch.no_grad if self.vit_frozen else nullcontext
            if self.vit_frozen and config.vit_type == "qwen3vl-vit-for-0.6b":
                trainable_modules = []
                merger_fc2 = getattr(getattr(self.vit, "merger", None), "linear_fc2", None)
                if merger_fc2 is not None:
                    trainable_modules.append(merger_fc2)
                for merger in getattr(self.vit, "deepstack_merger_list", []):
                    fc2 = getattr(merger, "linear_fc2", None)
                    if fc2 is not None:
                        trainable_modules.append(fc2)
                if trainable_modules:
                    for module in trainable_modules:
                        module.requires_grad_(True)
                        module.train()
                    self.vit_context = nullcontext

        # ======================================
        #       Define language modules
        # ======================================

        if config.use_mot:
            self.model = nn.ModuleDict(
                dict(
                    embed_tokens=nn.Embedding(config.vocab_size, config.hidden_size, **factory_kwargs),
                    layers=nn.ModuleList([
                        MultimodalMoTDecoderLayer(
                            config, block_idx, self.mot_und_frozen,
                            gen_config=self._config_mot_gen,
                            **factory_kwargs
                        )
                        for block_idx in range(config.num_layers)
                    ]),
                    norm=config.norm_class(config.hidden_size, eps=config.norm_eps, **factory_kwargs),
                )
            )
        else:
            self.model = nn.ModuleDict(
                dict(
                    embed_tokens=nn.Embedding(config.vocab_size, config.hidden_size, **factory_kwargs),
                    layers=nn.ModuleList([
                        MultimodalDecoderLayer(config, block_idx, **factory_kwargs)
                        for block_idx in range(config.num_layers)
                    ]),
                    norm=config.norm_class(config.hidden_size, eps=config.norm_eps, **factory_kwargs),
                )
            )

        if self.lm_frozen:
            # Freeze entire LLM backbone: embed_tokens, all transformer layers, norm.
            # Used for Stage 2.1 where only the vision-language connector is trained.
            self.model.eval()
            self.model.requires_grad_(False)

        # ====================== Finish model building =====================

        # Initialize cached rope, supporting automatic cache update
        self.cached_rope = CachedRoPE(config)
        self.use_rope_sample_offsets = getattr(args, "use_rope_sample_offsets", False)

        # Initialize weights if needed
        self._prepare_reset_parameters()
        if initialize_weights:
            for name, module in self.named_modules():
                if hasattr(module, "reset_parameters"):
                    module.reset_parameters()

    def _prepare_reset_parameters(self):
        # Globally set Linear and Embedding init methods to normal
        for module in self.modules():
            if isinstance(module, (nn.Linear, nn.Embedding)):
                module.reset_parameters = normal_weight_reset_parameters(
                    std=self._config.init_std, bias_type="zeros").__get__(module)
        for name, module in self.named_modules():
            if hasattr(module, "prepare_reset_parameters"):
                module.prepare_reset_parameters()

    @property
    def dtype(self):
        """Get the dtype of the model parameters."""
        if self._dtype is not None:
            return self._dtype
        # Fallback to getting dtype from model parameters
        try:
            return next(self.parameters()).dtype
        except StopIteration:
            # If no parameters, try buffers
            try:
                return next(self.buffers()).dtype
            except StopIteration:
                # Default fallback
                return torch.float32

    def scatter_to_hidden_states(
        self,
        src: torch.Tensor,
        index: torch.Tensor,
        hidden_states: torch.Tensor,
        gen_hidden_states: Optional[torch.Tensor] = None,
        dim: int = 1,
        slice_idx: Optional[int] = None,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        n_embd = src.shape[-1]
        hidden_target = hidden_states if slice_idx is None else hidden_states[slice_idx:slice_idx + 1]

        if n_embd == hidden_states.shape[-1]:
            hidden_target.scatter_(dim=dim, index=index, src=src)
        else:
            assert gen_hidden_states is not None, "gen_hidden_states is required when hidden dim of src and hidden_states differ"
            assert gen_hidden_states.shape[-1] == n_embd, \
                f"Expect gen_hidden_states and src to have same hidden_size, but got {gen_hidden_states.shape[-1]} and {n_embd}"
            gen_target = gen_hidden_states if slice_idx is None else gen_hidden_states[slice_idx:slice_idx + 1]
            gen_target.scatter_(dim=dim, index=index, src=src)

        return hidden_states, gen_hidden_states

    def instantiate_vae_image_tokens(
            self,
            hidden_states: Optional[Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]],
            timesteps: BatchRaggedTensor,
            medias: BatchRaggedMedia,
            media_mask: torch.Tensor,
    ):
        if hidden_states is None:
            hidden_size = self._config_mot_gen.hidden_size if self._config.use_mot else self._config.hidden_size

            if isinstance(medias, list):
                emb_list = []
                for i in range(len(medias)):
                    t_i = timesteps[i:i+1] if isinstance(timesteps, torch.Tensor) else timesteps[i]
                    t_emb_i = self.time_embed(t_i)
                    img_emb_i, _, _ = self.patch_embed(medias[i], t_emb_i)
                    emb_list.append(img_emb_i)

                max_tokens = max(e.size(1) for e in emb_list)
                padded_img_emb = torch.zeros(len(medias), max_tokens, hidden_size, device=emb_list[0].device, dtype=emb_list[0].dtype)
                for i, emb in enumerate(emb_list):
                    padded_img_emb[i, :emb.size(1), :] = emb[0]

                timestep_emb = self.timestep_emb(timesteps).reshape(len(medias), -1, hidden_size)
                hidden_states = torch.cat([timestep_emb, padded_img_emb], dim=1)
                return hidden_states

            t_emb = self.time_embed(timesteps)
            image_emb, _, _ = self.patch_embed(medias, t_emb)
            hidden_size = self._config_mot_gen.hidden_size if self._config.use_mot else self._config.hidden_size
            timestep_emb = self.timestep_emb(timesteps).reshape(medias.size(0), -1, hidden_size)
            hidden_states = torch.cat([timestep_emb, image_emb], dim=1)
            return hidden_states

        if isinstance(hidden_states, tuple):
            hidden_states, gen_hidden_states = hidden_states
        else:
            gen_hidden_states = None

        bsz, seqlen, n_embd = hidden_states.shape
        assert isinstance(medias, (torch.Tensor, list)), f"images should be BatchRaggedMedia, got {type(medias)}"

        if isinstance(medias, torch.Tensor):
            assert medias.ndim in [4, 5], f"images should be a 4-D or 5-D tensor, got {medias.ndim}-D tensor"
            assert isinstance(timesteps, torch.Tensor), f"timesteps should be 1-D tensor, got {type(timesteps)}"

            index = torch.arange(seqlen, device=hidden_states.device).unsqueeze(0).repeat(bsz, 1)   # (bsz, seqlen)
            t_emb = self.time_embed(timesteps)     # (bsz, n_embd)
            media_seq, *_ = self.patch_embed(medias, t_emb)   # (bsz, num_patches, n_embd)
            media_index = index.masked_select(media_mask.bool()).reshape(bsz, -1)   # (bsz, num_patches)
            assert media_seq.size(1) == media_index.size(1), \
                f"image_seq ({list(media_seq.size())}) has inconsistent shape with index ({list(media_index.size())})"
            n_embd = media_seq.shape[-1]
            index_exp = media_index.unsqueeze(-1).repeat(1, 1, n_embd)
            hidden_states, gen_hidden_states = self.scatter_to_hidden_states(
                media_seq.to(hidden_states.dtype), index_exp, hidden_states, gen_hidden_states
            )

        else:   # list
            index = torch.arange(seqlen, device=hidden_states.device).unsqueeze(0).repeat(bsz, 1)   # (bsz, seqlen)
            for i in range(len(medias)):
                media_i = medias[i]
                t_i = timesteps[i:i+1] if isinstance(timesteps, torch.Tensor) else timesteps[i]

                t_i_emb = self.time_embed(t_i)      # (n_i, n_embd)

                if isinstance(media_i, torch.Tensor):
                    media_i_seq, *_ = self.patch_embed(media_i, t_i_emb)  # (n_i, num_patches, n_embd)

                elif isinstance(media_i, list):
                    media_i_seq_list = []
                    for j in range(len(media_i)):
                        media_ij = media_i[j].unsqueeze(0)
                        assert media_ij.ndim in [4, 5], \
                            f"image_ij should have size of (1, C, H, W) or (1, C, D, H, W), got {list(media_ij.size())}"
                        media_ij_seq, *_ = self.patch_embed(media_ij, t_i_emb[j:j + 1])  # (1, num_patches, n_embd)
                        media_i_seq_list.append(media_ij_seq)
                    media_i_seq = torch.cat(media_i_seq_list, dim=1)    # (1, Σj num_patches_j, n_embd)

                else:
                    raise TypeError(f"image_i should be a 4-D or 5-D tensor or a list, got {type(media_i)}")

                media_i_index = index[i:i + 1].masked_select(media_mask[i:i + 1].bool()).reshape(1, -1)  # (1, img_seqlen)
                n_embd = media_i_seq.shape[-1]
                media_i_index_exp = media_i_index.unsqueeze(-1).repeat(1, 1, n_embd)
                media_i_seq_flat = media_i_seq.reshape(1, -1, n_embd)
                assert media_i_seq_flat.shape[1] == media_i_index_exp.shape[1], \
                    f"media_i_seq_flat ({list(media_i_seq_flat.size())}) has inconsistent shape with media_i_index_exp ({list(media_i_index_exp.size())})"
                hidden_states, gen_hidden_states = self.scatter_to_hidden_states(
                    media_i_seq_flat.to(hidden_states.dtype), media_i_index_exp, hidden_states, gen_hidden_states,
                    slice_idx=i,
                )

        if gen_hidden_states is not None:
            return hidden_states, gen_hidden_states

        return hidden_states

    def _forward_vision_encoder(self, images, **image_kwargs):
        with self.vit_context():
            image_embeds = self.vit(images, **image_kwargs)

        if isinstance(image_embeds, tuple):
            image_embeds, deepstack_image_embeds = image_embeds
        else:
            deepstack_image_embeds = None
            image_embeds = image_embeds.last_hidden_state
        return image_embeds, deepstack_image_embeds

    @staticmethod
    def _accumulate_deepstack_embeds(all_embeds, new_embeds):
        if new_embeds is None:
            return all_embeds
        if all_embeds is None:
            all_embeds = [[] for _ in range(len(new_embeds))]
        for layer_idx, layer_embeds in enumerate(new_embeds):
            all_embeds[layer_idx].append(layer_embeds)
        return all_embeds

    def instantiate_vit_image_tokens(
            self,
            hidden_states: torch.Tensor,
            images: torch.Tensor | list[torch.Tensor],
            image_masks: torch.Tensor,
            image_kwargs: dict[str, torch.Tensor],
    ):
        if isinstance(hidden_states, tuple):
            hidden_states, gen_hidden_states = hidden_states
        else:
            gen_hidden_states = None

        bsz, seqlen, _ = hidden_states.shape
        index = torch.arange(seqlen, device=hidden_states.device).unsqueeze(0).repeat(bsz, 1)

        if isinstance(images, torch.Tensor):
            assert images.ndim in [3, 4, 5], f"images should be a 3-D, 4-D, or 5-D tensor, got {images.ndim}-D tensor."
            if images.ndim in [4, 5]:
                bsz, n = images.shape[:2]
                images = images.view(bsz * n, *images.shape[2:])
                image_kwargs = image_kwargs if image_kwargs is not None else {}
                for k, v in image_kwargs.items():
                    image_kwargs[k] = v.reshape(bsz * n, *v.shape[2:])
            else:
                n = 1
            image_embeds, deepstack_image_embeds = self._forward_vision_encoder(images, **image_kwargs)

            image_seqlen, n_embd = image_embeds.size(1), image_embeds.size(-1)

            image_scatter_index = index.masked_select(image_masks.bool()).reshape(bsz, -1)
            index = image_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd)
            src = image_embeds.reshape(bsz, n * image_seqlen, n_embd)
            assert src.shape[1] == index.shape[1], \
                f"src ({list(src.size())}) has inconsistent shape with index ({list(index.size())})"
            hidden_states, gen_hidden_states = self.scatter_to_hidden_states(
                src.to(hidden_states.dtype), index, hidden_states, gen_hidden_states
            )

        elif isinstance(images, list):
            all_deepstack_embeds = None
            for i, (image, image_mask) in enumerate(zip(images, image_masks)):
                start_index = 0
                image_scatter_index = index[i].masked_select(image_mask.bool()).reshape(1, -1)
                for j, singel_image in enumerate(image):

                    cur_kwargs = {k: v[i][j:j+1] for k, v in image_kwargs.items()} if image_kwargs is not None else {}
                    if isinstance(singel_image, list):
                        image_embed_list = []
                        for _single_image in singel_image:
                            image_embed, deepstack_image_embeds = self._forward_vision_encoder(_single_image, **cur_kwargs)
                            image_embed_list.append(image_embed)
                            all_deepstack_embeds = self._accumulate_deepstack_embeds(all_deepstack_embeds, deepstack_image_embeds)
                        image_embed = torch.cat(image_embed_list, dim=1)
                        if image_embed.ndim == 3:
                            n, image_seqlen, n_embd = image_embed.shape
                            image_embed = image_embed.reshape(n * image_seqlen, n_embd)
                        else:
                            n_embd = image_embed.shape[-1]
                    else:
                        image_embed, deepstack_image_embeds = self._forward_vision_encoder(singel_image, **cur_kwargs)
                        all_deepstack_embeds = self._accumulate_deepstack_embeds(all_deepstack_embeds, deepstack_image_embeds)
                        if image_embed.ndim == 3:
                            n, image_seqlen, n_embd = image_embed.shape
                            image_embed = image_embed.reshape(n * image_seqlen, n_embd)
                        else:
                            n_embd = image_embed.shape[-1]

                    image_scatter_index_j = image_scatter_index[:, start_index:start_index + image_embed.shape[0]]
                    image_scatter_index_j = image_scatter_index_j.unsqueeze(-1).repeat(1, 1, n_embd)
                    image_embed = image_embed.reshape(1, -1, n_embd)
                    start_index += image_embed.shape[1]

                    assert image_scatter_index_j.shape[1] == image_embed.shape[1], \
                        f"image_scatter_index_j ({list(image_scatter_index_j.size())}) has inconsistent shape with image_embed ({list(image_embed.size())})"
                    hidden_states, gen_hidden_states = self.scatter_to_hidden_states(
                        image_embed.to(hidden_states.dtype), image_scatter_index_j, hidden_states, gen_hidden_states,
                        slice_idx=i,
                    )

            if all_deepstack_embeds is not None:
                deepstack_image_embeds = []
                for layer_embeds_list in all_deepstack_embeds:
                    valid_embeds = [e for e in layer_embeds_list if e is not None]
                    if len(valid_embeds) == 0:
                        continue
                    flattened_embeds = []
                    for e in valid_embeds:
                        if isinstance(e, list):
                            flattened_embeds.extend([item for item in e if isinstance(item, torch.Tensor)])
                        elif isinstance(e, torch.Tensor):
                            flattened_embeds.append(e)
                    if len(flattened_embeds) == 1:
                        deepstack_image_embeds.append(flattened_embeds[0])
                    elif len(flattened_embeds) > 1:
                        deepstack_image_embeds.append(torch.cat(flattened_embeds, dim=0))
            else:
                deepstack_image_embeds = None
        else:
            raise ValueError(f"und_images should be Tensor or List, but got {type(images)}")

        if gen_hidden_states is not None:
            return (hidden_states, gen_hidden_states), deepstack_image_embeds

        return hidden_states, deepstack_image_embeds

    def instantiate_continuous_tokens(
            self,
            hidden_states: torch.Tensor,
            emb_layer: nn.Module,
            scatter_src: Optional[BatchRaggedTensor] = None,
            scatter_index: Optional[BatchRaggedTensor] = None,
    ):
        if isinstance(hidden_states, tuple):
            hidden_states, gen_hidden_states = hidden_states
        else:
            gen_hidden_states = None

        bsz, seqlen, _ = hidden_states.shape

        if isinstance(scatter_src, list):
            for i, scatter_src_i in enumerate(scatter_src):
                src = emb_layer(scatter_src_i)  # (n, n_embd)
                n_embd = src.shape[-1]
                index = scatter_index[i].unsqueeze(0).unsqueeze(-1).repeat(1, 1, n_embd)
                src = src.reshape(1, -1, n_embd)

                assert index.shape[1] == src.shape[1], \
                    f"index ({list(index.size())}) has inconsistent shape with src ({list(src.size())})"
                hidden_states, gen_hidden_states = self.scatter_to_hidden_states(
                    src.to(hidden_states.dtype), index, hidden_states, gen_hidden_states,
                    slice_idx=i,
                )

        else:
            src = emb_layer(scatter_src.reshape(-1))    # (bsz * n, n_embd)
            n_embd = src.shape[-1]
            index = scatter_index.unsqueeze(-1).repeat(1, 1, n_embd)
            src = src.reshape(bsz, -1, n_embd)

            assert index.shape[1] == src.shape[1], \
                f"index ({list(index.size())}) has inconsistent shape with src ({list(src.size())})"
            hidden_states, gen_hidden_states = self.scatter_to_hidden_states(
                src.to(hidden_states.dtype), index, hidden_states, gen_hidden_states
            )

        if gen_hidden_states is not None:
            return hidden_states, gen_hidden_states

        return hidden_states

    def get_image_tokens_hw(self, images: BatchRaggedMedia):
        assert isinstance(images, (torch.Tensor, list)), f"images should be BatchRaggedMedia, got {type(images)}"
        if isinstance(images, torch.Tensor):
            token_h = images.shape[-2]
            token_w = images.shape[-1]
        else:
            token_h, token_w = [], []
            for image_i in images:
                assert isinstance(image_i, (torch.Tensor, list)), \
                    f"image_i should be a tensor or a list of tensors, got {type(image_i)}"
                if isinstance(image_i, torch.Tensor):
                    token_h.append(image_i.shape[-2])
                    token_w.append(image_i.shape[-1])
                else:
                    token_h.append([])
                    token_w.append([])
                    for j in range(len(image_i)):
                        token_h[-1].append(image_i[j].shape[-2])
                        token_w[-1].append(image_i[j].shape[-1])
        return token_h, token_w

    def ragged_final_layer(self, hidden_states, image_mask, timesteps, token_h, token_w, first_step=None, batch_image_sizes=None):
        n_embd = hidden_states.size(-1)

        if batch_image_sizes is not None:
            # For batch multi-resolution inference
            # bsz可能包含cond/uncond大小，而batch_image_sizes只包含实际的image大小
            bsz = hidden_states.size(0)
            actual_bsz = len(batch_image_sizes)
            pred = []
            for i in range(bsz):
                h_lat, w_lat = batch_image_sizes[i % actual_bsz]
                th_i = h_lat
                tw_i = w_lat
                n_tokens_i = th_i * tw_i
                assert n_tokens_i == image_mask[i].sum().item(), \
                    f"n_tokens_i ({n_tokens_i}) has inconsistent shape with image_mask[i].sum().item() ({image_mask[i].sum().item()})"

                if first_step is False:
                    # 非首步：hidden_states = [timestep_emb, img_emb(1D padded)]，因为有效 token 在前，直接按数量截取
                    image_output_i = hidden_states[i:i+1, 1:1+n_tokens_i, :]
                else:
                    # 首步：从完整序列中按 image_mask 布尔索引提取 image token
                    image_output_i = hidden_states[i, image_mask[i].bool(), :].unsqueeze(0)

                t_emb_i = self.time_embed_2(timesteps[i:i+1])
                pred_i = self.final_layer(image_output_i, t_emb_i, th_i, tw_i)
                pred.append(pred_i)
            return pred

        if isinstance(timesteps, torch.Tensor):
            # When timesteps is a tensor, images must be a 4-D tensor (B, C, H, W), which means only one target image
            t_emb = self.time_embed_2(timesteps)
            if first_step is False:
                # only for gen_image non-first-step inference
                image_output = hidden_states[:, 1:, :]
            else:   # first_step is True or None
                image_output = hidden_states.masked_select(
                    image_mask.unsqueeze(-1).bool()).reshape(-1, token_h * token_w, n_embd)
            pred = self.final_layer(image_output, t_emb, token_h, token_w)
        else:
            # When timesteps is a list, images must be a list of 4-D tensors or a list of list of 3-D tensors, and token_h and token_w must be a list of int or a list of list of int.
            # In this case, each line of the image_mask may contain different number of Trues, leading
            # the `reshape(batch_size, ...)` is not possible.
            sections = image_mask.sum(1).tolist()
            image_output = hidden_states.masked_select(
                image_mask.unsqueeze(-1).bool()).reshape(-1, n_embd).split(sections)
            pred = []
            for image_output_i, t_i, token_h_i, token_w_i in zip(image_output, timesteps, token_h, token_w):
                t_emb_i = self.time_embed_2(t_i)
                if isinstance(token_h_i, int):
                    # corresponds to image_output as a list of 4-D tensors, image_output_i as a 4-D tensor
                    image_output_i = image_output_i.reshape(-1, token_h_i * token_w_i, n_embd)
                    pred_i = self.final_layer(image_output_i, t_emb_i, token_h_i, token_w_i)
                    pred.append(pred_i)
                else:
                    # corresponds to image_output as a list of list of 3-D tensors, image_output_i as a list of 3-D tensors
                    subsections = [token_h_ij * token_w_ij for token_h_ij, token_w_ij in zip(token_h_i, token_w_i)]
                    assert sum(subsections) == image_output_i.shape[0], \
                        f"sum(subsections) ({sum(subsections)}) has inconsistent shape with image_output_i.shape[0] ({image_output_i.shape[0]})"
                    image_output_i = image_output_i.split(subsections)
                    pred_i = []
                    for j, image_output_ij in enumerate(image_output_i):
                        pred_ij = self.final_layer(image_output_ij[None], t_emb_i[j:j+1], token_h_i[j], token_w_i[j])
                        pred_i.append(pred_ij)
                    pred.append(pred_i)
        return pred

    def _deepstack_process(
        self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor
    ):
        visual_pos_masks = visual_pos_masks.to(hidden_states.device)
        if isinstance(visual_embeds, list):
            if len(visual_embeds) == 0:
                raise ValueError("visual_embeds is an empty list")
            elif len(visual_embeds) == 1:
                visual_embeds = visual_embeds[0]
                if not isinstance(visual_embeds, torch.Tensor):
                    raise ValueError(f"visual_embeds list contains non-tensor element: {type(visual_embeds)}")
            else:
                if all(isinstance(e, torch.Tensor) for e in visual_embeds):
                    visual_embeds = torch.cat(visual_embeds, dim=0)
                else:
                    raise ValueError(f"visual_embeds list contains non-tensor elements")
        visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)

        bsz = hidden_states.shape[0]
        for i in range(bsz):
            batch_mask = visual_pos_masks[i]
            batch_hidden = hidden_states[i]
            num_visual_tokens = batch_mask.sum().item()
            if num_visual_tokens == 0:
                continue

            visual_hidden = batch_hidden[batch_mask]

            if visual_embeds.shape[0] >= num_visual_tokens:
                batch_visual_embeds = visual_embeds[:num_visual_tokens]
            else:
                batch_visual_embeds = visual_embeds
                if visual_embeds.shape[0] < num_visual_tokens:
                    padding = visual_embeds[-1:].repeat(num_visual_tokens - visual_embeds.shape[0], 1)
                    batch_visual_embeds = torch.cat([visual_embeds, padding], dim=0)

            visual_hidden = visual_hidden + batch_visual_embeds
            batch_hidden[batch_mask] = visual_hidden
            hidden_states[i] = batch_hidden

        return hidden_states

    def forward(
            self,
            input_ids: Optional[torch.Tensor] = None,  # bsz x seqlen
            attention_mask: Optional[torch.Tensor] = None,  # bsz x 1 x seqlen x seqlen
            rope_image_info: Optional[list[list[tuple[slice, tuple[int, int], dict]]]] = None,
            return_dict: bool = True,
            # for gen images
            images: Optional[BatchRaggedMedia] = None,  # bsz x c x h x w, or bsz x (n_i x (c x h_ij x w_ij))
            image_mask: Optional[torch.Tensor] = None,  # bsz x seqlen
            timesteps: Optional[BatchRaggedTensor] = None,  # bsz, or bsz x (n_i)
            timesteps_index: Optional[BatchRaggedTensor] = None,  # bsz x k, or bsz x (k_i)
            # for cond images
            cond_vae_images: Optional[BatchRaggedMedia] = None,  # bsz x c x h x w, or bsz x (m_i x (c x h_ij x w_ij))
            cond_vae_image_mask: Optional[torch.Tensor] = None,  # bsz x seqlen
            cond_timesteps: Optional[BatchRaggedTensor] = None,  # bsz, or bsz x (m_i)
            cond_timesteps_index: Optional[BatchRaggedTensor] = None,
            cond_vit_images: Optional[BatchRaggedMedia] = None,
            cond_vit_image_mask: Optional[torch.Tensor] = None,
            cond_vit_image_kwargs: Optional[dict[str, Any]] = None,
            # only for inference
            input_pos: Optional[torch.Tensor] = None,  # bsz x seq_len-1, used for KVCache
            past_key_values: Optional[MultimodalStaticCache] = None,
            mode: Optional[str] = None,
            first_step: Optional[bool] = None,
            und_token_indices: Optional[torch.Tensor] = None,
            gen_token_indices: Optional[torch.Tensor] = None,
            sample_offsets: Optional[list[torch.Tensor]] = None,
            batch_image_sizes: Optional[list[tuple[int, int]]] = None,
    ) -> MultimodalModelOutput | tuple:
        # Sanity check
        if input_ids is None and images is None:
            raise ValueError("Either input_ids or images should be provided.")
        if input_ids is not None:
            bsz = input_ids.size(0)
            device = input_ids.device
        else:
            bsz = images.size(0) if isinstance(images, torch.Tensor) else len(images)
            device = get_device(images)
        if self.training:
            seqlen = input_ids.size(1)
        else:
            seqlen = self._config.max_position_embeddings
        assert self._config.max_position_embeddings >= seqlen, (
            f"Cannot forward sequence of length {seqlen}, "
            f"max position embeddings is only {self._config.max_position_embeddings}, "
            f"try set --max-position-embeddings to a larger value."
        )
        cos, sin = self.cached_rope(
            seqlen, device, rope_media_info=rope_image_info, input_pos=input_pos, sample_offsets=sample_offsets if self.use_rope_sample_offsets else None,
        )
        cos = cos.to(dtype=self.dtype)
        sin = sin.to(dtype=self.dtype)

        if input_ids is not None:
            hidden_states = self.model["embed_tokens"](input_ids)     # (bsz, seqlen, n_embd)
            if self._config.use_mot and self._config_mot_gen.hidden_size != self._config.hidden_size:
                und_hidden_states = hidden_states
                gen_hidden_states = torch.zeros(
                    (hidden_states.size(0), hidden_states.size(1), self._config_mot_gen.hidden_size),
                    device=hidden_states.device, dtype=hidden_states.dtype
                )
                hidden_states = (und_hidden_states, gen_hidden_states)
        else:
            hidden_states = None    # only for non-first step inference of the image generation

        deepstack_image_embeds = None
        if images is not None:
            hidden_states = self.instantiate_vae_image_tokens(hidden_states, timesteps, images, image_mask)

        if cond_vae_images is not None:
            hidden_states = self.instantiate_vae_image_tokens(hidden_states, cond_timesteps, cond_vae_images, cond_vae_image_mask)

        if cond_vit_images is not None:
            hidden_states, deepstack_image_embeds = self.instantiate_vit_image_tokens(hidden_states, cond_vit_images, cond_vit_image_mask, cond_vit_image_kwargs)

        if timesteps_index is not None:
            hidden_states = self.instantiate_continuous_tokens(hidden_states, emb_layer=self.timestep_emb, scatter_src=timesteps, scatter_index=timesteps_index)

        if cond_timesteps_index is not None:
            hidden_states = self.instantiate_continuous_tokens(hidden_states, emb_layer=self.timestep_emb, scatter_src=cond_timesteps, scatter_index=cond_timesteps_index)

        if self._config.use_mot:
            if und_token_indices is None or gen_token_indices is None:
                seq_for_indices = hidden_states[0].shape[1] if isinstance(hidden_states, tuple) else hidden_states.shape[1]
                all_indices = torch.arange(seq_for_indices, device=device).unsqueeze(0).expand(bsz, -1)
                gen_mask = torch.zeros(bsz, seq_for_indices, dtype=torch.bool, device=device)
                for mask in (image_mask, cond_vae_image_mask):
                    if mask is not None:
                        gen_mask |= mask.to(device=device, dtype=torch.bool)
                gen_token_indices = all_indices[gen_mask].reshape(bsz, -1)
                und_token_indices = all_indices[~gen_mask].reshape(bsz, -1)
            if isinstance(hidden_states, tuple):
                und_hidden_states, gen_hidden_states = hidden_states
            else:
                if first_step is False:
                    step_indices = torch.arange(
                        hidden_states.shape[1],
                        device=hidden_states.device,
                    ).unsqueeze(0).expand(hidden_states.shape[0], -1)
                    if images is not None:
                        und_hidden_states = hidden_states.new_zeros(
                            hidden_states.shape[0], 0, self._config.hidden_size
                        )
                        gen_hidden_states = hidden_states
                        und_token_indices = step_indices[:, :0]
                        gen_token_indices = step_indices
                    else:
                        und_hidden_states = hidden_states
                        gen_hidden_states = hidden_states.new_zeros(
                            hidden_states.shape[0], 0, self._config_mot_gen.hidden_size
                        )
                        und_token_indices = step_indices
                        gen_token_indices = step_indices[:, :0]
                else:
                    und_hidden_states = hidden_states
                    gen_hidden_states = hidden_states

            und_token_indices_ = und_token_indices.unsqueeze(-1).expand(-1, -1, und_hidden_states.shape[-1])
            gen_token_indices_ = gen_token_indices.unsqueeze(-1).expand(-1, -1, gen_hidden_states.shape[-1])

            und_hidden_states = und_hidden_states.gather(dim=1, index=und_token_indices_)
            gen_hidden_states = gen_hidden_states.gather(dim=1, index=gen_token_indices_)
            hidden_states = (und_hidden_states, gen_hidden_states)

        if self._config.use_modality_routing:
            _hs_for_seqlen = hidden_states[0] if isinstance(hidden_states, tuple) else hidden_states
            actual_seqlen = _hs_for_seqlen.size(1)
            token_modalities = torch.zeros(bsz, actual_seqlen, dtype=torch.long, device=device)

            if input_ids is not None:
                if cond_vit_image_mask is not None:
                    token_modalities[cond_vit_image_mask.bool()] = 1
                if image_mask is not None:
                    token_modalities[image_mask.bool()] = 2
                if cond_vae_image_mask is not None:
                    token_modalities[cond_vae_image_mask.bool()] = 2
            else:
                token_modalities[:, 1:] = 2
        else:
            token_modalities = None

        for layer_idx, layer in enumerate(self.model["layers"]):    # noqa
            layer_inputs = [
                hidden_states,
                attention_mask,
                (cos, sin),
                input_pos,
                past_key_values,
                und_token_indices,
                gen_token_indices
            ]
            hidden_states = layer(*layer_inputs, token_modalities=token_modalities)
            if deepstack_image_embeds is not None and layer_idx in range(len(deepstack_image_embeds)):
                if isinstance(hidden_states, tuple):
                    und_hs, gen_hs = hidden_states
                    und_vit_image_mask = cond_vit_image_mask.gather(
                        dim=1,
                        index=und_token_indices.to(cond_vit_image_mask.device),
                    )
                    und_hs = self._deepstack_process(
                        und_hs,
                        und_vit_image_mask,
                        deepstack_image_embeds[layer_idx],
                    )
                    hidden_states = (und_hs, gen_hs)
                else:
                    hidden_states = self._deepstack_process(
                        hidden_states,
                        cond_vit_image_mask,
                        deepstack_image_embeds[layer_idx],
                    )

        if isinstance(hidden_states, tuple):
            und_hidden_states_, gen_hidden_states_ = hidden_states
            bsz = und_hidden_states_.shape[0]
            und_seqlen = und_hidden_states_.shape[1]
            gen_seqlen = gen_hidden_states_.shape[1]

            und_hidden_states = torch.zeros(
                (bsz, und_seqlen+gen_seqlen, und_hidden_states_.shape[-1]),
                device=und_hidden_states_.device, dtype=und_hidden_states_.dtype
            )
            und_hidden_states.scatter_(dim=1, index=und_token_indices_.to(und_hidden_states_.device), src=und_hidden_states_)

            gen_hidden_states = torch.zeros(
                (bsz, und_seqlen+gen_seqlen, gen_hidden_states_.shape[-1]),
                device=gen_hidden_states_.device, dtype=gen_hidden_states_.dtype
            )
            gen_hidden_states.scatter_(dim=1, index=gen_token_indices_, src=gen_hidden_states_)

        else:
            und_hidden_states, gen_hidden_states = hidden_states, hidden_states

        if images is not None:
            token_h, token_w = self.get_image_tokens_hw(images)
            gen_hidden_states = gen_hidden_states.to(device=get_device(images))
            diff_pred = self.ragged_final_layer(
                gen_hidden_states, image_mask, timesteps, token_h, token_w, first_step, batch_image_sizes=batch_image_sizes)
        else:
            diff_pred = None

        if input_ids is None or mode == "gen_image":
            logits = None
        else:
            und_hidden_states = self.model["norm"](und_hidden_states)
            logits = F.linear(und_hidden_states, self.model.embed_tokens.weight)

        if not return_dict:
            return logits, past_key_values, diff_pred
        return MultimodalModelOutput(
            logits=logits,
            past_key_values=past_key_values,
            diffusion_prediction=diff_pred,
        )

class MultimodalModel(MultimodalModelBase):
    def __init__(
            self,
            args: Namespace,
            config: MultimodalConfig,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
            initialize_weights: bool = True,
    ):
        super().__init__()
        self.__post_init__(config, dtype, device, args, initialize_weights)

# Model construction helpers
def build_model(
        args,
        logger=None,
        dtype=None,
        device=None,
        **kwargs,
) -> tuple[torch.nn.Module, Any]:
    # Support cpu, cuda, meta devices
    factor_kwargs = {"device": device, "dtype": dtype}

    if logger is None:
        from loguru import logger

    model_structure = args.model_structure
    logger.info(f"Building model {model_structure} for {args.model_name}")

    if device == 'meta':
        context = torch.device('meta')
    else:
        context = nullcontext()

    with context:
        if model_structure in {"MultimodalModel", "MultimodalHFModel"}:
            model, model_config = _build_multimodal_model(
                args, logger=logger, **kwargs, **factor_kwargs)
        else:
            raise NotImplementedError(f"Model structure {model_structure} not implemented.")

    return model, model_config


def _build_multimodal_model(args, logger=None, dtype=None, device=None, **kwargs):
    factory_kwargs = {"device": device or 'cpu', "dtype": dtype}

    model_structure = args.model_structure
    valid_keywords = {"initialize_weights"}
    valid_kwargs = {key: value for key, value in kwargs.items() if key in valid_keywords}

    if model_structure == "MultimodalModel":
        from .configuration import core_model_config_from_args

        model_name = args.model_name.split(".")[-1]
        model_config_dict = core_model_config_from_args(args)
        model_config = MultimodalConfig.from_name(model_name, **model_config_dict)
        model = MultimodalModel(args, model_config, **factory_kwargs, **valid_kwargs)

    elif model_structure == "MultimodalHFModel":
        from .pipeline import MultimodalHFModel
        from .configuration import core_model_config_from_args

        model_name = args.model_name.split(".")[-1]
        model_config_dict = core_model_config_from_args(args)
        model_config = MultimodalConfig.from_name(model_name, **model_config_dict)
        model = MultimodalHFModel(args, model_config, **factory_kwargs, **valid_kwargs)

    else:
        raise NotImplementedError(f"Model structure {model_structure} not implemented.")

    logger.info(f"Build Model {model.__class__.__name__} finished.")

    return model, model_config