MiMoViT-2.6-Pro / modeling_mimo_vision.py

Upload vision tower bundle

bdfd5fd verified 2 days ago

17 kB

# coding=utf-8
#
# Copyright 2026 Xiaomi Corporation.
# Copyright 2026 The HuggingFace Inc. team.
#
# 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.

# Vision classes extracted from XiaomiMiMo/MiMo-V2.6-Pro-RL
# b61f21852680ee02a29e946e18cf8697f4844061. Flash uses identical vision source.
# No text/audio model is instantiated or imported by this standalone module.
# The three absent merger biases are omitted, matching zero-bias HF completion
# and SGLang's LayerNorm merger. Attention/block arithmetic is unchanged.
import torch
from torch import nn
from torch.nn import functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.activations import ACT2FN


class MiMoVisionConfig(PretrainedConfig):
    model_type = "mimo26_vision"

    def __init__(self, **kwargs):
        defaults = dict(depth=28, hidden_size=1280, intermediate_size=4608,
                        num_heads=32, num_key_value_heads=8, qk_channels=64,
                        patch_size=16, temporal_patch_size=2, in_chans=3,
                        spatial_merge_size=2, out_hidden_size=6144,
                        hidden_act="silu", rms_norm_eps=1e-6,
                        fullatt_block_indexes=[0, 9, 18, 27], use_sink=True,
                        visual_token_window_size=64,
                        vit_window_attn_types=[-1,0,0,0,0,1,1,1,1,-1,0,0,0,0,1,1,1,1,-1,0,0,0,0,1,1,1,1,-1])
        defaults.update(kwargs)
        super().__init__(**defaults)


def _rotate_half_vision(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def _apply_rotary_pos_emb_vision(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    orig_q_dtype, orig_k_dtype = q.dtype, k.dtype
    q, k = q.float(), k.float()
    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
    q_embed = (q * cos) + (_rotate_half_vision(q) * sin)
    k_embed = (k * cos) + (_rotate_half_vision(k) * sin)
    return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype)


class MiMoVisionRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, seqlen: int) -> torch.Tensor:
        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        return torch.outer(seq, self.inv_freq)


class MiMoVisionPatchEmbed(nn.Module):
    def __init__(
        self, patch_size: int = 16, temporal_patch_size: int = 2, in_channels: int = 3, embed_dim: int = 1280
    ):
        super().__init__()
        self.patch_size = patch_size
        self.temporal_patch_size = temporal_patch_size
        self.in_channels = in_channels
        self.embed_dim = embed_dim
        kernel_size = [temporal_patch_size, patch_size, patch_size]
        self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype
        hidden_states = hidden_states.view(
            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
        )
        return self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)


class MiMoVisionSwiGLUMLP(nn.Module):
    def __init__(self, dim: int, intermediate_dim: int, hidden_act: str = "silu"):
        super().__init__()
        self.gate_proj = nn.Linear(dim, intermediate_dim, bias=True)
        self.up_proj = nn.Linear(dim, intermediate_dim, bias=True)
        self.down_proj = nn.Linear(intermediate_dim, dim, bias=True)
        self.act_fn = ACT2FN[hidden_act]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class MiMoVisionAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        num_kv_heads: int | None = None,
        head_dim: int | None = None,
        use_sinks: bool = False,
        window_size: int = -1,
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
        self.head_dim = head_dim if head_dim is not None else dim // num_heads
        self.num_kv_groups = self.num_heads // self.num_kv_heads
        self.scaling = self.head_dim**-0.5
        self.window_size = window_size

        qkv_dim = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim
        self.qkv = nn.Linear(dim, qkv_dim, bias=True)
        self.proj = nn.Linear(self.num_heads * self.head_dim, dim, bias=True)
        self.sinks = nn.Parameter(torch.zeros(self.num_heads)) if use_sinks else None

    def _build_window_mask(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
        if self.window_size <= 0:
            return None
        row_idx = torch.arange(seq_len, device=device).unsqueeze(1)
        col_idx = torch.arange(seq_len, device=device).unsqueeze(0)
        mask = torch.zeros(seq_len, seq_len, device=device, dtype=dtype)
        mask = mask.masked_fill((row_idx - col_idx).abs() > self.window_size, float("-inf"))
        return mask

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        full_attn: bool = False,
    ) -> torch.Tensor:
        seq_len = hidden_states.shape[0]
        qkv = self.qkv(hidden_states)

        q_dim = self.num_heads * self.head_dim
        kv_dim = self.num_kv_heads * self.head_dim
        q = qkv[:, :q_dim].view(seq_len, self.num_heads, self.head_dim)
        k = qkv[:, q_dim : q_dim + kv_dim].view(seq_len, self.num_kv_heads, self.head_dim)
        v = qkv[:, q_dim + kv_dim :].view(seq_len, self.num_kv_heads, self.head_dim)

        cos, sin = position_embeddings
        q, k = _apply_rotary_pos_emb_vision(q, k, cos, sin)

        lengths = cu_seqlens[1:] - cu_seqlens[:-1]
        q_chunks = torch.split(q, lengths.tolist(), dim=0)
        k_chunks = torch.split(k, lengths.tolist(), dim=0)
        v_chunks = torch.split(v, lengths.tolist(), dim=0)

        outputs = []
        for q_c, k_c, v_c in zip(q_chunks, k_chunks, v_chunks):
            q_c = q_c.unsqueeze(0).transpose(1, 2)
            k_c = k_c.unsqueeze(0).transpose(1, 2)
            v_c = v_c.unsqueeze(0).transpose(1, 2)

            if self.num_kv_groups > 1:
                k_c = k_c.repeat_interleave(self.num_kv_groups, dim=1)
                v_c = v_c.repeat_interleave(self.num_kv_groups, dim=1)

            attn_mask = None
            if not full_attn:
                attn_mask = self._build_window_mask(q_c.shape[2], q_c.device, q_c.dtype)

            if self.sinks is not None:
                sink_bias = torch.zeros(
                    1, self.num_heads, q_c.shape[2], k_c.shape[2], device=q_c.device, dtype=q_c.dtype
                )
                sink_bias[..., 0] = self.sinks.view(1, self.num_heads, 1)
                attn_mask = sink_bias if attn_mask is None else attn_mask + sink_bias

            attn_out = F.scaled_dot_product_attention(q_c, k_c, v_c, attn_mask=attn_mask, scale=self.scaling)
            outputs.append(attn_out.squeeze(0).transpose(0, 1))

        attn_output = torch.cat(outputs, dim=0)
        attn_output = attn_output.reshape(seq_len, -1)
        return self.proj(attn_output)


class MiMoVisionBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        intermediate_dim: int,
        num_heads: int,
        num_kv_heads: int | None = None,
        head_dim: int | None = None,
        hidden_act: str = "silu",
        rms_norm_eps: float = 1e-6,
        use_sinks: bool = False,
        window_size: int = -1,
    ):
        super().__init__()
        self.norm1 = nn.RMSNorm(dim, eps=rms_norm_eps)
        self.norm2 = nn.RMSNorm(dim, eps=rms_norm_eps)
        self.attn = MiMoVisionAttention(
            dim=dim, num_heads=num_heads, num_kv_heads=num_kv_heads, head_dim=head_dim,
            use_sinks=use_sinks, window_size=window_size,
        )
        self.mlp = MiMoVisionSwiGLUMLP(dim=dim, intermediate_dim=intermediate_dim, hidden_act=hidden_act)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        full_attn: bool = False,
    ) -> torch.Tensor:
        hidden_states = hidden_states + self.attn(
            self.norm1(hidden_states), cu_seqlens=cu_seqlens,
            position_embeddings=position_embeddings, full_attn=full_attn,
        )
        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
        return hidden_states


class MiMoVisionPatchMerger(nn.Module):
    def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2):
        super().__init__()
        self.hidden_size = context_dim * (spatial_merge_size**2)
        self.ln_q = nn.LayerNorm(context_dim, eps=1e-6, bias=False)
        self.mlp = nn.Sequential(
            nn.Linear(self.hidden_size, self.hidden_size, bias=False),
            nn.GELU(),
            nn.Linear(self.hidden_size, dim, bias=False),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.mlp(self.ln_q(x).view(-1, self.hidden_size))


class MiMoVisionTransformer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        hidden_size = config.hidden_size
        depth = config.depth
        num_heads = config.num_heads
        num_kv_heads = getattr(config, "num_key_value_heads", num_heads)
        head_dim = getattr(config, "qk_channels", 64)
        spatial_merge_size = getattr(config, "spatial_merge_size", 2)
        rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
        self.fullatt_block_indexes = getattr(config, "fullatt_block_indexes", [])
        use_sink = getattr(config, "use_sink", False)
        visual_token_window_size = getattr(config, "visual_token_window_size", -1)
        self.vit_window_attn_types = getattr(config, "vit_window_attn_types", None) or [-1] * depth

        self.spatial_merge_size = spatial_merge_size
        self.spatial_merge_unit = spatial_merge_size * spatial_merge_size

        self.patch_embed = MiMoVisionPatchEmbed(
            patch_size=config.patch_size,
            temporal_patch_size=config.temporal_patch_size,
            in_channels=getattr(config, "in_channels", None) or getattr(config, "in_chans", 3),
            embed_dim=hidden_size,
        )

        self.rotary_pos_emb = MiMoVisionRotaryEmbedding(head_dim // 2)

        self.blocks = nn.ModuleList(
            [
                MiMoVisionBlock(
                    dim=hidden_size,
                    intermediate_dim=config.intermediate_size,
                    num_heads=num_heads,
                    num_kv_heads=num_kv_heads,
                    head_dim=head_dim,
                    hidden_act=config.hidden_act,
                    rms_norm_eps=rms_norm_eps,
                    use_sinks=use_sink and (i not in self.fullatt_block_indexes),
                    window_size=visual_token_window_size,
                )
                for i in range(depth)
            ]
        )

        self.merger = MiMoVisionPatchMerger(
            dim=config.out_hidden_size,
            context_dim=hidden_size,
            spatial_merge_size=spatial_merge_size,
        )

    @property
    def dtype(self) -> torch.dtype:
        return self.patch_embed.proj.weight.dtype

    def apply_index(self, tensor: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
        tensor = tensor.unflatten(0, (-1, self.spatial_merge_unit))
        tensor = tensor[index]
        return tensor.flatten(0, 1)

    def get_window_index_1d(self, grid_thw: torch.Tensor, col: bool = True) -> torch.Tensor:
        window_index = []
        window_index_id = 0
        for grid_t, grid_h, grid_w in grid_thw:
            llm_grid_h = grid_h // self.spatial_merge_size
            llm_grid_w = grid_w // self.spatial_merge_size
            index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)
            index_new = index.transpose(1, 2).reshape(-1) if col else index.reshape(-1)
            window_index.append(index_new + window_index_id)
            window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
        return torch.cat(window_index, dim=0)

    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
        pos_ids = []
        for t, h, w in grid_thw:
            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
            hpos_ids = hpos_ids.reshape(
                h // self.spatial_merge_size, self.spatial_merge_size,
                w // self.spatial_merge_size, self.spatial_merge_size,
            )
            hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten()

            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
            wpos_ids = wpos_ids.reshape(
                h // self.spatial_merge_size, self.spatial_merge_size,
                w // self.spatial_merge_size, self.spatial_merge_size,
            )
            wpos_ids = wpos_ids.permute(0, 2, 1, 3).flatten()

            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
        pos_ids = torch.cat(pos_ids, dim=0)
        max_grid_size = grid_thw[:, 1:].max()
        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
        return rotary_pos_emb_full[pos_ids].flatten(1)

    def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
        x = pixel_values.to(device=self.patch_embed.proj.weight.device, dtype=self.dtype)
        x = self.patch_embed(x)

        rotary_emb = self.rot_pos_emb(grid_thw)
        rotary_emb = rotary_emb.to(device=x.device)
        emb = torch.cat((rotary_emb, rotary_emb), dim=-1)

        window_index_1d_col = self.get_window_index_1d(grid_thw, col=True).to(device=x.device)
        reverse_window_index_1d_col = torch.argsort(window_index_1d_col).to(device=x.device)

        row_based_embeddings = (emb.cos(), emb.sin())
        col_emb = self.apply_index(emb, window_index_1d_col)
        col_based_embeddings = (col_emb.cos(), col_emb.sin())

        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
            dim=0, dtype=torch.int32
        )
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).to(device=x.device)

        for i, blk in enumerate(self.blocks):
            window_attn_type = self.vit_window_attn_types[i]

            if window_attn_type == 1 and (i == 0 or self.vit_window_attn_types[i - 1] != 1):
                x = self.apply_index(x, window_index_1d_col)

            if i > 0 and window_attn_type != 1 and self.vit_window_attn_types[i - 1] == 1:
                x = self.apply_index(x, reverse_window_index_1d_col)

            position_embeddings = col_based_embeddings if window_attn_type == 1 else row_based_embeddings
            full_attn = i in self.fullatt_block_indexes
            x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, full_attn=full_attn)

        return self.merger(x)




class MiMoVisionModel(PreTrainedModel):
    config_class = MiMoVisionConfig
    base_model_prefix = "visual"
    main_input_name = "pixel_values"
    _no_split_modules = ["MiMoVisionBlock"]

    def __init__(self, config):
        super().__init__(config)
        self.visual = MiMoVisionTransformer(config)
        self.post_init()

    def _init_weights(self, module):
        super()._init_weights(module)
        if isinstance(module, MiMoVisionRotaryEmbedding):
            # HF materializes nonpersistent meta buffers without running their
            # constructors again. Rebuild frequencies during missing-key init.
            dim = module.inv_freq.numel() * 2
            frequencies = 1.0 / (10000.0 ** (torch.arange(0, dim, 2, dtype=torch.float32, device=module.inv_freq.device) / dim))
            module.inv_freq.copy_(frequencies)

    def forward(self, pixel_values, grid_thw):
        # Parent-compatible output after LayerNorm, 2x2 merge and projector.
        return self.visual(pixel_values, grid_thw)

	# coding=utf-8
	#
	# Copyright 2026 Xiaomi Corporation.
	# Copyright 2026 The HuggingFace Inc. team.
	#
	# 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.

	# Vision classes extracted from XiaomiMiMo/MiMo-V2.6-Pro-RL
	# b61f21852680ee02a29e946e18cf8697f4844061. Flash uses identical vision source.
	# No text/audio model is instantiated or imported by this standalone module.
	# The three absent merger biases are omitted, matching zero-bias HF completion
	# and SGLang's LayerNorm merger. Attention/block arithmetic is unchanged.
	import torch
	from torch import nn
	from torch.nn import functional as F
	from transformers import PretrainedConfig, PreTrainedModel
	from transformers.activations import ACT2FN


	class MiMoVisionConfig(PretrainedConfig):
	model_type = "mimo26_vision"

	def __init__(self, **kwargs):
	defaults = dict(depth=28, hidden_size=1280, intermediate_size=4608,
	num_heads=32, num_key_value_heads=8, qk_channels=64,
	patch_size=16, temporal_patch_size=2, in_chans=3,
	spatial_merge_size=2, out_hidden_size=6144,
	hidden_act="silu", rms_norm_eps=1e-6,
	fullatt_block_indexes=[0, 9, 18, 27], use_sink=True,
	visual_token_window_size=64,
	vit_window_attn_types=[-1,0,0,0,0,1,1,1,1,-1,0,0,0,0,1,1,1,1,-1,0,0,0,0,1,1,1,1,-1])
	defaults.update(kwargs)
	super().__init__(**defaults)


	def _rotate_half_vision(x: torch.Tensor) -> torch.Tensor:
	x1 = x[..., : x.shape[-1] // 2]
	x2 = x[..., x.shape[-1] // 2 :]
	return torch.cat((-x2, x1), dim=-1)


	def _apply_rotary_pos_emb_vision(
	q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
	) -> tuple[torch.Tensor, torch.Tensor]:
	orig_q_dtype, orig_k_dtype = q.dtype, k.dtype
	q, k = q.float(), k.float()
	cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
	q_embed = (q * cos) + (_rotate_half_vision(q) * sin)
	k_embed = (k * cos) + (_rotate_half_vision(k) * sin)
	return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype)


	class MiMoVisionRotaryEmbedding(nn.Module):
	def __init__(self, dim: int, theta: float = 10000.0) -> None:
	super().__init__()
	inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
	self.register_buffer("inv_freq", inv_freq, persistent=False)

	def forward(self, seqlen: int) -> torch.Tensor:
	seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
	return torch.outer(seq, self.inv_freq)


	class MiMoVisionPatchEmbed(nn.Module):
	def __init__(
	self, patch_size: int = 16, temporal_patch_size: int = 2, in_channels: int = 3, embed_dim: int = 1280
	):
	super().__init__()
	self.patch_size = patch_size
	self.temporal_patch_size = temporal_patch_size
	self.in_channels = in_channels
	self.embed_dim = embed_dim
	kernel_size = [temporal_patch_size, patch_size, patch_size]
	self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	target_dtype = self.proj.weight.dtype
	hidden_states = hidden_states.view(
	-1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
	)
	return self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)


	class MiMoVisionSwiGLUMLP(nn.Module):
	def __init__(self, dim: int, intermediate_dim: int, hidden_act: str = "silu"):
	super().__init__()
	self.gate_proj = nn.Linear(dim, intermediate_dim, bias=True)
	self.up_proj = nn.Linear(dim, intermediate_dim, bias=True)
	self.down_proj = nn.Linear(intermediate_dim, dim, bias=True)
	self.act_fn = ACT2FN[hidden_act]

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


	class MiMoVisionAttention(nn.Module):
	def __init__(
	self,
	dim: int,
	num_heads: int,
	num_kv_heads: int \| None = None,
	head_dim: int \| None = None,
	use_sinks: bool = False,
	window_size: int = -1,
	):
	super().__init__()
	self.dim = dim
	self.num_heads = num_heads
	self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
	self.head_dim = head_dim if head_dim is not None else dim // num_heads
	self.num_kv_groups = self.num_heads // self.num_kv_heads
	self.scaling = self.head_dim**-0.5
	self.window_size = window_size

	qkv_dim = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim
	self.qkv = nn.Linear(dim, qkv_dim, bias=True)
	self.proj = nn.Linear(self.num_heads * self.head_dim, dim, bias=True)
	self.sinks = nn.Parameter(torch.zeros(self.num_heads)) if use_sinks else None

	def _build_window_mask(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor \| None:
	if self.window_size <= 0:
	return None
	row_idx = torch.arange(seq_len, device=device).unsqueeze(1)
	col_idx = torch.arange(seq_len, device=device).unsqueeze(0)
	mask = torch.zeros(seq_len, seq_len, device=device, dtype=dtype)
	mask = mask.masked_fill((row_idx - col_idx).abs() > self.window_size, float("-inf"))
	return mask

	def forward(
	self,
	hidden_states: torch.Tensor,
	cu_seqlens: torch.Tensor,
	position_embeddings: tuple[torch.Tensor, torch.Tensor],
	full_attn: bool = False,
	) -> torch.Tensor:
	seq_len = hidden_states.shape[0]
	qkv = self.qkv(hidden_states)

	q_dim = self.num_heads * self.head_dim
	kv_dim = self.num_kv_heads * self.head_dim
	q = qkv[:, :q_dim].view(seq_len, self.num_heads, self.head_dim)
	k = qkv[:, q_dim : q_dim + kv_dim].view(seq_len, self.num_kv_heads, self.head_dim)
	v = qkv[:, q_dim + kv_dim :].view(seq_len, self.num_kv_heads, self.head_dim)

	cos, sin = position_embeddings
	q, k = _apply_rotary_pos_emb_vision(q, k, cos, sin)

	lengths = cu_seqlens[1:] - cu_seqlens[:-1]
	q_chunks = torch.split(q, lengths.tolist(), dim=0)
	k_chunks = torch.split(k, lengths.tolist(), dim=0)
	v_chunks = torch.split(v, lengths.tolist(), dim=0)

	outputs = []
	for q_c, k_c, v_c in zip(q_chunks, k_chunks, v_chunks):
	q_c = q_c.unsqueeze(0).transpose(1, 2)
	k_c = k_c.unsqueeze(0).transpose(1, 2)
	v_c = v_c.unsqueeze(0).transpose(1, 2)

	if self.num_kv_groups > 1:
	k_c = k_c.repeat_interleave(self.num_kv_groups, dim=1)
	v_c = v_c.repeat_interleave(self.num_kv_groups, dim=1)

	attn_mask = None
	if not full_attn:
	attn_mask = self._build_window_mask(q_c.shape[2], q_c.device, q_c.dtype)

	if self.sinks is not None:
	sink_bias = torch.zeros(
	1, self.num_heads, q_c.shape[2], k_c.shape[2], device=q_c.device, dtype=q_c.dtype
	)
	sink_bias[..., 0] = self.sinks.view(1, self.num_heads, 1)
	attn_mask = sink_bias if attn_mask is None else attn_mask + sink_bias

	attn_out = F.scaled_dot_product_attention(q_c, k_c, v_c, attn_mask=attn_mask, scale=self.scaling)
	outputs.append(attn_out.squeeze(0).transpose(0, 1))

	attn_output = torch.cat(outputs, dim=0)
	attn_output = attn_output.reshape(seq_len, -1)
	return self.proj(attn_output)


	class MiMoVisionBlock(nn.Module):
	def __init__(
	self,
	dim: int,
	intermediate_dim: int,
	num_heads: int,
	num_kv_heads: int \| None = None,
	head_dim: int \| None = None,
	hidden_act: str = "silu",
	rms_norm_eps: float = 1e-6,
	use_sinks: bool = False,
	window_size: int = -1,
	):
	super().__init__()
	self.norm1 = nn.RMSNorm(dim, eps=rms_norm_eps)
	self.norm2 = nn.RMSNorm(dim, eps=rms_norm_eps)
	self.attn = MiMoVisionAttention(
	dim=dim, num_heads=num_heads, num_kv_heads=num_kv_heads, head_dim=head_dim,
	use_sinks=use_sinks, window_size=window_size,
	)
	self.mlp = MiMoVisionSwiGLUMLP(dim=dim, intermediate_dim=intermediate_dim, hidden_act=hidden_act)

	def forward(
	self,
	hidden_states: torch.Tensor,
	cu_seqlens: torch.Tensor,
	position_embeddings: tuple[torch.Tensor, torch.Tensor],
	full_attn: bool = False,
	) -> torch.Tensor:
	hidden_states = hidden_states + self.attn(
	self.norm1(hidden_states), cu_seqlens=cu_seqlens,
	position_embeddings=position_embeddings, full_attn=full_attn,
	)
	hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
	return hidden_states


	class MiMoVisionPatchMerger(nn.Module):
	def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2):
	super().__init__()
	self.hidden_size = context_dim * (spatial_merge_size**2)
	self.ln_q = nn.LayerNorm(context_dim, eps=1e-6, bias=False)
	self.mlp = nn.Sequential(
	nn.Linear(self.hidden_size, self.hidden_size, bias=False),
	nn.GELU(),
	nn.Linear(self.hidden_size, dim, bias=False),
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.mlp(self.ln_q(x).view(-1, self.hidden_size))


	class MiMoVisionTransformer(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.config = config
	hidden_size = config.hidden_size
	depth = config.depth
	num_heads = config.num_heads
	num_kv_heads = getattr(config, "num_key_value_heads", num_heads)
	head_dim = getattr(config, "qk_channels", 64)
	spatial_merge_size = getattr(config, "spatial_merge_size", 2)
	rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
	self.fullatt_block_indexes = getattr(config, "fullatt_block_indexes", [])
	use_sink = getattr(config, "use_sink", False)
	visual_token_window_size = getattr(config, "visual_token_window_size", -1)
	self.vit_window_attn_types = getattr(config, "vit_window_attn_types", None) or [-1] * depth

	self.spatial_merge_size = spatial_merge_size
	self.spatial_merge_unit = spatial_merge_size * spatial_merge_size

	self.patch_embed = MiMoVisionPatchEmbed(
	patch_size=config.patch_size,
	temporal_patch_size=config.temporal_patch_size,
	in_channels=getattr(config, "in_channels", None) or getattr(config, "in_chans", 3),
	embed_dim=hidden_size,
	)

	self.rotary_pos_emb = MiMoVisionRotaryEmbedding(head_dim // 2)

	self.blocks = nn.ModuleList(
	[
	MiMoVisionBlock(
	dim=hidden_size,
	intermediate_dim=config.intermediate_size,
	num_heads=num_heads,
	num_kv_heads=num_kv_heads,
	head_dim=head_dim,
	hidden_act=config.hidden_act,
	rms_norm_eps=rms_norm_eps,
	use_sinks=use_sink and (i not in self.fullatt_block_indexes),
	window_size=visual_token_window_size,
	)
	for i in range(depth)
	]
	)

	self.merger = MiMoVisionPatchMerger(
	dim=config.out_hidden_size,
	context_dim=hidden_size,
	spatial_merge_size=spatial_merge_size,
	)

	@property
	def dtype(self) -> torch.dtype:
	return self.patch_embed.proj.weight.dtype

	def apply_index(self, tensor: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
	tensor = tensor.unflatten(0, (-1, self.spatial_merge_unit))
	tensor = tensor[index]
	return tensor.flatten(0, 1)

	def get_window_index_1d(self, grid_thw: torch.Tensor, col: bool = True) -> torch.Tensor:
	window_index = []
	window_index_id = 0
	for grid_t, grid_h, grid_w in grid_thw:
	llm_grid_h = grid_h // self.spatial_merge_size
	llm_grid_w = grid_w // self.spatial_merge_size
	index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)
	index_new = index.transpose(1, 2).reshape(-1) if col else index.reshape(-1)
	window_index.append(index_new + window_index_id)
	window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
	return torch.cat(window_index, dim=0)

	def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
	pos_ids = []
	for t, h, w in grid_thw:
	hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
	hpos_ids = hpos_ids.reshape(
	h // self.spatial_merge_size, self.spatial_merge_size,
	w // self.spatial_merge_size, self.spatial_merge_size,
	)
	hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten()

	wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
	wpos_ids = wpos_ids.reshape(
	h // self.spatial_merge_size, self.spatial_merge_size,
	w // self.spatial_merge_size, self.spatial_merge_size,
	)
	wpos_ids = wpos_ids.permute(0, 2, 1, 3).flatten()

	pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
	pos_ids = torch.cat(pos_ids, dim=0)
	max_grid_size = grid_thw[:, 1:].max()
	rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
	return rotary_pos_emb_full[pos_ids].flatten(1)

	def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
	x = pixel_values.to(device=self.patch_embed.proj.weight.device, dtype=self.dtype)
	x = self.patch_embed(x)

	rotary_emb = self.rot_pos_emb(grid_thw)
	rotary_emb = rotary_emb.to(device=x.device)
	emb = torch.cat((rotary_emb, rotary_emb), dim=-1)

	window_index_1d_col = self.get_window_index_1d(grid_thw, col=True).to(device=x.device)
	reverse_window_index_1d_col = torch.argsort(window_index_1d_col).to(device=x.device)

	row_based_embeddings = (emb.cos(), emb.sin())
	col_emb = self.apply_index(emb, window_index_1d_col)
	col_based_embeddings = (col_emb.cos(), col_emb.sin())

	cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
	dim=0, dtype=torch.int32
	)
	cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).to(device=x.device)

	for i, blk in enumerate(self.blocks):
	window_attn_type = self.vit_window_attn_types[i]

	if window_attn_type == 1 and (i == 0 or self.vit_window_attn_types[i - 1] != 1):
	x = self.apply_index(x, window_index_1d_col)

	if i > 0 and window_attn_type != 1 and self.vit_window_attn_types[i - 1] == 1:
	x = self.apply_index(x, reverse_window_index_1d_col)

	position_embeddings = col_based_embeddings if window_attn_type == 1 else row_based_embeddings
	full_attn = i in self.fullatt_block_indexes
	x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, full_attn=full_attn)

	return self.merger(x)




	class MiMoVisionModel(PreTrainedModel):
	config_class = MiMoVisionConfig
	base_model_prefix = "visual"
	main_input_name = "pixel_values"
	_no_split_modules = ["MiMoVisionBlock"]

	def __init__(self, config):
	super().__init__(config)
	self.visual = MiMoVisionTransformer(config)
	self.post_init()

	def _init_weights(self, module):
	super()._init_weights(module)
	if isinstance(module, MiMoVisionRotaryEmbedding):
	# HF materializes nonpersistent meta buffers without running their
	# constructors again. Rebuild frequencies during missing-key init.
	dim = module.inv_freq.numel() * 2
	frequencies = 1.0 / (10000.0 ** (torch.arange(0, dim, 2, dtype=torch.float32, device=module.inv_freq.device) / dim))
	module.inv_freq.copy_(frequencies)

	def forward(self, pixel_values, grid_thw):
	# Parent-compatible output after LayerNorm, 2x2 merge and projector.
	return self.visual(pixel_values, grid_thw)