| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | from collections.abc import Callable, Iterable, Mapping, Sequence |
| 5 | from functools import partial |
| 6 | from typing import Any |
| 7 | |
| 8 | import einops |
| 9 | import numpy as np |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | import torch.nn.functional as F |
| 13 | from transformers import BatchFeature, PreTrainedConfig |
| 14 | from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize |
| 15 | from typing_extensions import TypedDict |
| 16 | |
| 17 | from vllm.config import VllmConfig |
| 18 | from vllm.config.multimodal import MultiModalDummyOptions |
| 19 | from vllm.distributed import parallel_state |
| 20 | from vllm.distributed import utils as dist_utils |
| 21 | from vllm.inputs import MultiModalDataDict |
| 22 | from vllm.model_executor.layers.activation import get_act_and_mul_fn |
| 23 | from vllm.model_executor.layers.attention import MMEncoderAttention |
| 24 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 25 | from vllm.model_executor.layers.linear import ( |
| 26 | ColumnParallelLinear, |
| 27 | QKVParallelLinear, |
| 28 | RowParallelLinear, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 31 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 32 | from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb |
| 33 | from vllm.model_executor.models.vision import is_vit_use_data_parallel |
| 34 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 35 | from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems |
| 36 | from vllm.multimodal.parse import ImageSize, MultiModalDataItems |
| 37 | from vllm.multimodal.processing import ( |
| 38 | BaseDummyInputsBuilder, |
| 39 | BaseMultiModalProcessor, |
| 40 | BaseProcessingInfo, |
| 41 | PromptReplacement, |
| 42 | PromptUpdate, |
| 43 | PromptUpdateDetails, |
| 44 | ) |
| 45 | from vllm.multimodal.processing.processor import HFMultiModalInputs |
| 46 | from vllm.transformers_utils.configs.mimo_v2_omni import Mimo_VLVisionConfig |
| 47 | from vllm.transformers_utils.processors.mimo_v2_omni import ( |
| 48 | MiMoOmniProcessor, |
| 49 | VideoAudioInput, |
| 50 | _format_timestamp, |
| 51 | ) |
| 52 | |
| 53 | from .interfaces import ( |
| 54 | MultiModalEmbeddings, |
| 55 | SupportsEagle3, |
| 56 | SupportsMultiModal, |
| 57 | SupportsPP, |
| 58 | SupportsQuant, |
| 59 | ) |
| 60 | from .mimo_audio import MimoAudioEncoder |
| 61 | from .mimo_v2 import MiMoV2FlashForCausalLM |
| 62 | from .qwen2_5_vl import ( |
| 63 | Qwen2_5_VisionMLP, |
| 64 | Qwen2_5_VisionPatchEmbed, |
| 65 | Qwen2_5_VLImageEmbeddingInputs, |
| 66 | Qwen2_5_VLImageInputs, |
| 67 | Qwen2_5_VLImagePixelInputs, |
| 68 | Qwen2_5_VLVideoEmbeddingInputs, |
| 69 | Qwen2_5_VLVideoInputs, |
| 70 | Qwen2_5_VLVideoPixelInputs, |
| 71 | ) |
| 72 | from .qwen2_vl import ( |
| 73 | Qwen2VLMultiModalDataParser, |
| 74 | _create_qwen2vl_field_factory, |
| 75 | ) |
| 76 | from .utils import AutoWeightsLoader, IntermediateTensors, WeightsMapper, maybe_prefix |
| 77 | |
| 78 | |
| 79 | class MiMoV2OmniAudioInputs(TypedDict): |
| 80 | type: str |
| 81 | audio_features: torch.Tensor | list[torch.Tensor] |
| 82 | audio_token_lens: torch.Tensor | list[int] | None |
| 83 | |
| 84 | |
| 85 | class MiMoV2OmniMultiModalInputs(TypedDict, total=False): |
| 86 | image: Qwen2_5_VLImageInputs | None |
| 87 | video: Qwen2_5_VLVideoInputs | None |
| 88 | audio: MiMoV2OmniAudioInputs | None |
| 89 | |
| 90 | |
| 91 | class MiMoVisionMLP(Qwen2_5_VisionMLP): |
| 92 | pass |
| 93 | |
| 94 | |
| 95 | class MiMoVisionPatchEmbed(Qwen2_5_VisionPatchEmbed): |
| 96 | pass |
| 97 | |
| 98 | |
| 99 | class MiMoVisionPatchMerger(nn.Module): |
| 100 | def __init__( |
| 101 | self, |
| 102 | d_model: int, |
| 103 | context_dim: int, |
| 104 | norm_layer: Callable[[int], nn.Module] | None = None, |
| 105 | spatial_merge_size: int = 2, |
| 106 | quant_config: QuantizationConfig | None = None, |
| 107 | prefix: str = "", |
| 108 | ) -> None: |
| 109 | super().__init__() |
| 110 | use_data_parallel = is_vit_use_data_parallel() |
| 111 | self.hidden_size = context_dim * (spatial_merge_size**2) |
| 112 | if norm_layer is None: |
| 113 | norm_layer = partial(nn.LayerNorm, eps=1e-6) |
| 114 | self.ln_q = norm_layer(context_dim) |
| 115 | |
| 116 | self.mlp = nn.Sequential( |
| 117 | ColumnParallelLinear( |
| 118 | self.hidden_size, |
| 119 | self.hidden_size, |
| 120 | bias=False, |
| 121 | quant_config=quant_config, |
| 122 | prefix=f"{prefix}.mlp.0", |
| 123 | return_bias=False, |
| 124 | disable_tp=use_data_parallel, |
| 125 | ), |
| 126 | nn.GELU(), |
| 127 | RowParallelLinear( |
| 128 | self.hidden_size, |
| 129 | d_model, |
| 130 | bias=False, |
| 131 | quant_config=quant_config, |
| 132 | prefix=f"{prefix}.mlp.2", |
| 133 | return_bias=False, |
| 134 | disable_tp=use_data_parallel, |
| 135 | ), |
| 136 | ) |
| 137 | |
| 138 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 139 | x = self.ln_q(x) |
| 140 | x = x.view(-1, self.hidden_size) |
| 141 | out = self.mlp(x) |
| 142 | return out |
| 143 | |
| 144 | |
| 145 | class MiMoVisionAttention(nn.Module): |
| 146 | def __init__( |
| 147 | self, |
| 148 | embed_dim: int, |
| 149 | num_heads: int, |
| 150 | num_kv_heads: int, |
| 151 | qk_channels: int, |
| 152 | kv_channels: int, |
| 153 | use_sink: bool = False, |
| 154 | visual_token_window_size: int = 64, |
| 155 | quant_config: QuantizationConfig | None = None, |
| 156 | prefix: str = "", |
| 157 | ) -> None: |
| 158 | super().__init__() |
| 159 | use_data_parallel = is_vit_use_data_parallel() |
| 160 | self.tp_size = ( |
| 161 | 1 |
| 162 | if use_data_parallel |
| 163 | else parallel_state.get_tensor_model_parallel_world_size() |
| 164 | ) |
| 165 | self.tp_rank = parallel_state.get_tensor_model_parallel_rank() |
| 166 | |
| 167 | self.num_heads = num_heads |
| 168 | self.num_kv_heads = num_kv_heads |
| 169 | self.qk_channels = qk_channels |
| 170 | self.kv_channels = kv_channels |
| 171 | self.embed_dim = embed_dim |
| 172 | |
| 173 | self.num_heads_per_partition = dist_utils.divide(num_heads, self.tp_size) |
| 174 | self.num_kv_heads_per_partition = dist_utils.divide(num_kv_heads, self.tp_size) |
| 175 | |
| 176 | # Attention scale uses the Q/K head dimension (qk_channels) |
| 177 | self.scale = qk_channels**-0.5 |
| 178 | |
| 179 | # QKV: Q is (num_heads * qk_channels), KV are (num_kv_heads * kv_channels) |
| 180 | self.qkv = QKVParallelLinear( |
| 181 | hidden_size=embed_dim, |
| 182 | head_size=qk_channels, |
| 183 | total_num_heads=num_heads, |
| 184 | total_num_kv_heads=num_kv_heads, |
| 185 | v_head_size=kv_channels, |
| 186 | bias=True, |
| 187 | quant_config=quant_config, |
| 188 | prefix=f"{prefix}.qkv", |
| 189 | disable_tp=use_data_parallel, |
| 190 | ) |
| 191 | |
| 192 | # Output projection: input is (num_heads * kv_channels) after attention |
| 193 | self.proj = RowParallelLinear( |
| 194 | input_size=num_heads * kv_channels, |
| 195 | output_size=embed_dim, |
| 196 | bias=True, |
| 197 | quant_config=quant_config, |
| 198 | prefix=f"{prefix}.proj", |
| 199 | disable_tp=use_data_parallel, |
| 200 | ) |
| 201 | |
| 202 | # For full attention (non-window blocks) |
| 203 | self.attn = MMEncoderAttention( |
| 204 | num_heads=self.num_heads_per_partition, |
| 205 | head_size=kv_channels, |
| 206 | scale=self.scale, |
| 207 | num_kv_heads=self.num_kv_heads_per_partition, |
| 208 | prefix=f"{prefix}.attn", |
| 209 | ) |
| 210 | |
| 211 | # Rotary embeddings applied separately to Q and K |
| 212 | self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) |
| 213 | |
| 214 | # Per-head sink logits, applied in the window attention path. |
| 215 | # The checkpoint stores these only for non-full-attention blocks |
| 216 | self.use_sink = use_sink |
| 217 | if use_sink: |
| 218 | self.sinks = nn.Parameter( |
| 219 | torch.empty(num_heads), |
| 220 | requires_grad=False, |
| 221 | ) |
| 222 | else: |
| 223 | self.sinks = None |
| 224 | |
| 225 | self.visual_token_window_size = visual_token_window_size |
| 226 | |
| 227 | def _forward_window_attn( |
| 228 | self, |
| 229 | q: torch.Tensor, |
| 230 | k: torch.Tensor, |
| 231 | v: torch.Tensor, |
| 232 | cu_seqlens: torch.Tensor, |
| 233 | max_seqlen: torch.Tensor, |
| 234 | ) -> torch.Tensor: |
| 235 | """Window attention with the per-head sink applied to key 0. |
| 236 | |
| 237 | The reference adds ``sinks[h]`` to the logit of each sequence's first |
| 238 | key, which the Triton prefill kernel supports directly, so the softmax |
| 239 | normalizes over the biased scores in one pass. |
| 240 | """ |
| 241 | from vllm.v1.attention.ops.triton_prefill_attention import ( |
| 242 | context_attention_fwd, |
| 243 | ) |
| 244 | |
| 245 | w = self.visual_token_window_size |
| 246 | output = torch.empty_like(q) |
| 247 | head_start = self.tp_rank * self.num_heads_per_partition |
| 248 | sinks = ( |
| 249 | self.sinks[head_start : head_start + self.num_heads_per_partition] |
| 250 | if self.sinks is not None |
| 251 | else None |
| 252 | ) |
| 253 | context_attention_fwd( |
| 254 | q, |
| 255 | k, |
| 256 | v, |
| 257 | output, |
| 258 | b_start_loc=cu_seqlens[:-1], |
| 259 | b_seq_len=cu_seqlens[1:] - cu_seqlens[:-1], |
| 260 | max_input_len=max_seqlen, |
| 261 | is_causal=False, |
| 262 | softmax_scale=self.scale, |
| 263 | sliding_window_q=w, |
| 264 | sliding_window_k=w, |
| 265 | sinks=sinks, |
| 266 | sinks_bias_key0=True, |
| 267 | ) |
| 268 | return output |
| 269 | |
| 270 | def forward( |
| 271 | self, |
| 272 | x: torch.Tensor, |
| 273 | cu_seqlens: torch.Tensor, |
| 274 | rotary_pos_emb_cos: torch.Tensor, |
| 275 | rotary_pos_emb_sin: torch.Tensor, |
| 276 | max_seqlen: torch.Tensor, |
| 277 | full_attn: bool = True, |
| 278 | ) -> torch.Tensor: |
| 279 | """Args: |
| 280 | x: [seq_len, batch=1, embed_dim] (seq-first convention) |
| 281 | cu_seqlens: cumulative sequence lengths [num_seqs+1], int32 |
| 282 | rotary_pos_emb_cos: [seq_len, qk_channels // 2] |
| 283 | rotary_pos_emb_sin: [seq_len, qk_channels // 2] |
| 284 | max_seqlen: maximum sequence length |
| 285 | full_attn: if True, full attention; if False, window attention |
| 286 | |
| 287 | """ |
| 288 | # [seq_len, 1, embed_dim] -> QKV projection |
| 289 | qkv, _ = self.qkv(x) # [seq_len, 1, q_size + kv_size + kv_size] |
| 290 | seq_len, batch_size, _ = qkv.shape |
| 291 | |
| 292 | q_size = self.num_heads_per_partition * self.qk_channels |
| 293 | kv_size = self.num_kv_heads_per_partition * self.kv_channels |
| 294 | q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) |
| 295 | |
| 296 | # Rearrange to [batch, seq, head, head_dim] for rotary application |
| 297 | q = einops.rearrange(q, "s b (h d) -> b s h d", h=self.num_heads_per_partition) |
| 298 | k = einops.rearrange( |
| 299 | k, "s b (h d) -> b s h d", h=self.num_kv_heads_per_partition |
| 300 | ) |
| 301 | v = einops.rearrange( |
| 302 | v, "s b (h d) -> b s h d", h=self.num_kv_heads_per_partition |
| 303 | ) |
| 304 | |
| 305 | # Apply rotary embeddings to Q and K independently (handles GQA) |
| 306 | if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None: |
| 307 | q = self.apply_rotary_emb(q, rotary_pos_emb_cos, rotary_pos_emb_sin) |
| 308 | k = self.apply_rotary_emb(k, rotary_pos_emb_cos, rotary_pos_emb_sin) |
| 309 | |
| 310 | if full_attn: |
| 311 | # Full attention via MMEncoderAttention |
| 312 | # Flatten to [batch, seq, heads * head_dim] |
| 313 | q_flat = q.reshape(batch_size, seq_len, -1) |
| 314 | k_flat = k.reshape(batch_size, seq_len, -1) |
| 315 | v_flat = v.reshape(batch_size, seq_len, -1) |
| 316 | context_layer = self.attn( |
| 317 | query=q_flat, |
| 318 | key=k_flat, |
| 319 | value=v_flat, |
| 320 | cu_seqlens=cu_seqlens, |
| 321 | max_seqlen=max_seqlen, |
| 322 | ) |
| 323 | # context_layer: [batch, seq, num_heads, head_dim] or [batch, seq, hidden] |
| 324 | # Ensure shape is [seq, batch, num_heads * kv_channels] |
| 325 | if context_layer.dim() == 4: |
| 326 | context_layer = einops.rearrange( |
| 327 | context_layer, "b s h d -> s b (h d)" |
| 328 | ).contiguous() |
| 329 | else: |
| 330 | context_layer = einops.rearrange( |
| 331 | context_layer, "b s d -> s b d" |
| 332 | ).contiguous() |
| 333 | else: |
| 334 | # Window attention via flash_attn_varlen_func with window_size |
| 335 | # Flatten batch dimension: [seq, head, head_dim] |
| 336 | q_varlen = einops.rearrange(q, "b s h d -> (b s) h d") |
| 337 | k_varlen = einops.rearrange(k, "b s h d -> (b s) h d") |
| 338 | v_varlen = einops.rearrange(v, "b s h d -> (b s) h d") |
| 339 | output = self._forward_window_attn( |
| 340 | q_varlen, k_varlen, v_varlen, cu_seqlens, max_seqlen |
| 341 | ) |
| 342 | # output: [total_tokens, num_heads, kv_channels] |
| 343 | context_layer = einops.rearrange( |
| 344 | output, "(b s) h d -> s b (h d)", b=batch_size |
| 345 | ).contiguous() |
| 346 | |
| 347 | output, _ = self.proj(context_layer) |
| 348 | return output |
| 349 | |
| 350 | |
| 351 | class MiMoVisionBlock(nn.Module): |
| 352 | def __init__( |
| 353 | self, |
| 354 | dim: int, |
| 355 | num_heads: int, |
| 356 | num_kv_heads: int, |
| 357 | qk_channels: int, |
| 358 | kv_channels: int, |
| 359 | mlp_hidden_dim: int, |
| 360 | act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu, |
| 361 | norm_eps: float = 1e-6, |
| 362 | use_sink: bool = False, |
| 363 | visual_token_window_size: int = 64, |
| 364 | quant_config: QuantizationConfig | None = None, |
| 365 | prefix: str = "", |
| 366 | ) -> None: |
| 367 | super().__init__() |
| 368 | self.norm1 = RMSNorm(dim, eps=norm_eps) |
| 369 | self.norm2 = RMSNorm(dim, eps=norm_eps) |
| 370 | self.attn = MiMoVisionAttention( |
| 371 | embed_dim=dim, |
| 372 | num_heads=num_heads, |
| 373 | num_kv_heads=num_kv_heads, |
| 374 | qk_channels=qk_channels, |
| 375 | kv_channels=kv_channels, |
| 376 | use_sink=use_sink, |
| 377 | visual_token_window_size=visual_token_window_size, |
| 378 | quant_config=quant_config, |
| 379 | prefix=f"{prefix}.attn", |
| 380 | ) |
| 381 | self.mlp = MiMoVisionMLP( |
| 382 | in_features=dim, |
| 383 | hidden_features=mlp_hidden_dim, |
| 384 | act_fn=act_fn, |
| 385 | bias=True, |
| 386 | quant_config=quant_config, |
| 387 | prefix=f"{prefix}.mlp", |
| 388 | ) |
| 389 | |
| 390 | def forward( |
| 391 | self, |
| 392 | x: torch.Tensor, |
| 393 | cu_seqlens: torch.Tensor, |
| 394 | rotary_pos_emb_cos: torch.Tensor, |
| 395 | rotary_pos_emb_sin: torch.Tensor, |
| 396 | max_seqlen: torch.Tensor, |
| 397 | full_attn: bool = True, |
| 398 | ) -> torch.Tensor: |
| 399 | # x: [seq_len, batch=1, dim] |
| 400 | x_attn = self.attn( |
| 401 | self.norm1(x), |
| 402 | cu_seqlens=cu_seqlens, |
| 403 | rotary_pos_emb_cos=rotary_pos_emb_cos, |
| 404 | rotary_pos_emb_sin=rotary_pos_emb_sin, |
| 405 | max_seqlen=max_seqlen, |
| 406 | full_attn=full_attn, |
| 407 | ) |
| 408 | # Fused residual add + norm2 |
| 409 | x_norm, residual = self.norm2(x, residual=x_attn) |
| 410 | x = residual + self.mlp(x_norm) |
| 411 | return x |
| 412 | |
| 413 | |
| 414 | class MiMoVisionTransformer(nn.Module): |
| 415 | hf_to_vllm_mapper = WeightsMapper( |
| 416 | orig_to_new_stacked={ |
| 417 | "mlp.gate_proj": ("mlp.gate_up_proj", 0), |
| 418 | "mlp.up_proj": ("mlp.gate_up_proj", 1), |
| 419 | } |
| 420 | ) |
| 421 | |
| 422 | def __init__( |
| 423 | self, |
| 424 | vision_cfg: PreTrainedConfig, |
| 425 | *, |
| 426 | norm_eps: float = 1e-6, |
| 427 | quant_config: QuantizationConfig | None = None, |
| 428 | prefix: str = "", |
| 429 | ): |
| 430 | super().__init__() |
| 431 | self.spatial_merge_size = vision_cfg.spatial_merge_size |
| 432 | self.spatial_merge_unit = self.spatial_merge_size**2 |
| 433 | self.fullatt_block_indexes = vision_cfg.fullatt_block_indexes |
| 434 | self.vit_window_attn_types = vision_cfg.vit_window_attn_types |
| 435 | self.visual_token_window_size = vision_cfg.visual_token_window_size |
| 436 | self.hidden_size = vision_cfg.hidden_size |
| 437 | self.num_heads = vision_cfg.num_heads |
| 438 | self.num_kv_heads = vision_cfg.num_key_value_heads |
| 439 | self.qk_channels = vision_cfg.qk_channels |
| 440 | self.kv_channels = vision_cfg.kv_channels |
| 441 | |
| 442 | self.patch_embed = MiMoVisionPatchEmbed( |
| 443 | patch_size=vision_cfg.patch_size, |
| 444 | temporal_patch_size=vision_cfg.temporal_patch_size, |
| 445 | in_channels=vision_cfg.in_channels, |
| 446 | hidden_size=vision_cfg.hidden_size, |
| 447 | ) |
| 448 | |
| 449 | norm_layer = partial(RMSNorm, eps=norm_eps) |
| 450 | |
| 451 | # Rotary embedding for 2D positions. |
| 452 | # With partial_rotary_factor=0.5 and head_size=qk_channels: |
| 453 | # rotary_dim = qk_channels // 2 |
| 454 | # get_cos_sin returns cos, sin each of shape [pos, rotary_dim // 2] |
| 455 | # After indexing with 2D pos_ids and flattening: |
| 456 | # result shape = [tokens, rotary_dim] = [tokens, qk_channels // 2] |
| 457 | # which is what ApplyRotaryEmb expects as cos/sin input. |
| 458 | self.rotary_pos_emb = get_rope( |
| 459 | head_size=vision_cfg.qk_channels, |
| 460 | max_position=8192, |
| 461 | is_neox_style=True, |
| 462 | rope_parameters={"partial_rotary_factor": 0.5}, |
| 463 | ) |
| 464 | |
| 465 | self.blocks = nn.ModuleList( |
| 466 | [ |
| 467 | MiMoVisionBlock( |
| 468 | dim=vision_cfg.hidden_size, |
| 469 | num_heads=vision_cfg.num_heads, |
| 470 | num_kv_heads=vision_cfg.num_key_value_heads, |
| 471 | qk_channels=vision_cfg.qk_channels, |
| 472 | kv_channels=vision_cfg.kv_channels, |
| 473 | mlp_hidden_dim=vision_cfg.intermediate_size, |
| 474 | act_fn=get_act_and_mul_fn(vision_cfg.hidden_act), |
| 475 | norm_eps=norm_eps, |
| 476 | use_sink=( |
| 477 | vision_cfg.use_sink |
| 478 | and i not in vision_cfg.fullatt_block_indexes |
| 479 | ), |
| 480 | visual_token_window_size=vision_cfg.visual_token_window_size, |
| 481 | quant_config=quant_config, |
| 482 | prefix=f"{prefix}.blocks.{i}", |
| 483 | ) |
| 484 | for i in range(vision_cfg.depth) |
| 485 | ] |
| 486 | ) |
| 487 | |
| 488 | self.merger = MiMoVisionPatchMerger( |
| 489 | d_model=vision_cfg.out_hidden_size, |
| 490 | context_dim=vision_cfg.hidden_size, |
| 491 | norm_layer=norm_layer, |
| 492 | spatial_merge_size=vision_cfg.spatial_merge_size, |
| 493 | quant_config=quant_config, |
| 494 | prefix=f"{prefix}.merger", |
| 495 | ) |
| 496 | |
| 497 | @property |
| 498 | def dtype(self) -> torch.dtype: |
| 499 | return self.patch_embed.proj.weight.dtype |
| 500 | |
| 501 | @property |
| 502 | def device(self) -> torch.device: |
| 503 | return self.patch_embed.proj.weight.device |
| 504 | |
| 505 | def apply_index(self, tensor: torch.Tensor, index: torch.Tensor) -> torch.Tensor: |
| 506 | """Reindex tensor at the spatial_merge_unit granularity.""" |
| 507 | tensor = tensor.unflatten(0, (-1, self.spatial_merge_unit)) |
| 508 | tensor = tensor[index] |
| 509 | tensor = tensor.flatten(0, 1) |
| 510 | return tensor |
| 511 | |
| 512 | def get_window_index_1d( |
| 513 | self, grid_thw: torch.Tensor, col: bool = True |
| 514 | ) -> torch.Tensor: |
| 515 | """Compute 1D window indices for col-based or row-based SWA reordering.""" |
| 516 | window_index: list[torch.Tensor] = [] |
| 517 | window_index_id = 0 |
| 518 | for grid_t, grid_h, grid_w in grid_thw: |
| 519 | llm_grid_h = grid_h // self.spatial_merge_size |
| 520 | llm_grid_w = grid_w // self.spatial_merge_size |
| 521 | index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape( |
| 522 | grid_t, llm_grid_h, llm_grid_w |
| 523 | ) |
| 524 | index_new = index.transpose(1, 2).reshape(-1) if col else index.reshape(-1) |
| 525 | window_index.append(index_new + window_index_id) |
| 526 | window_index_id += int((grid_t * llm_grid_h * llm_grid_w).item()) |
| 527 | return torch.cat(window_index, dim=0) |
| 528 | |
| 529 | def rot_pos_emb(self, grid_thw: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 530 | """Compute 2D rotary position embedding cos/sin for given grid sizes. |
| 531 | |
| 532 | Returns: |
| 533 | cos: [total_tokens, qk_channels // 2] |
| 534 | sin: [total_tokens, qk_channels // 2] |
| 535 | |
| 536 | """ |
| 537 | cos_list, sin_list = [], [] |
| 538 | for i in range(grid_thw.size(0)): |
| 539 | t, h, w = int(grid_thw[i, 0]), int(grid_thw[i, 1]), int(grid_thw[i, 2]) |
| 540 | |
| 541 | # Build 2D position IDs with spatial_merge_size interleaving |
| 542 | hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) |
| 543 | hpos_ids = ( |
| 544 | hpos_ids.reshape( |
| 545 | h // self.spatial_merge_size, |
| 546 | self.spatial_merge_size, |
| 547 | w // self.spatial_merge_size, |
| 548 | self.spatial_merge_size, |
| 549 | ) |
| 550 | .permute(0, 2, 1, 3) |
| 551 | .flatten() |
| 552 | ) |
| 553 | wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) |
| 554 | wpos_ids = ( |
| 555 | wpos_ids.reshape( |
| 556 | h // self.spatial_merge_size, |
| 557 | self.spatial_merge_size, |
| 558 | w // self.spatial_merge_size, |
| 559 | self.spatial_merge_size, |
| 560 | ) |
| 561 | .permute(0, 2, 1, 3) |
| 562 | .flatten() |
| 563 | ) |
| 564 | pos_ids = torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1) |
| 565 | # pos_ids: [t*h*w, 2] |
| 566 | |
| 567 | max_grid_size = max(h, w) |
| 568 | # get_cos_sin returns cos, sin each of shape [max_grid_size, rotary_dim//2] |
| 569 | # where rotary_dim = qk_channels // 2 (from partial_rotary_factor=0.5) |
| 570 | cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size) |
| 571 | |
| 572 | # [t*h*w, 2, rotary_dim//2] -> [t*h*w, rotary_dim] (= qk_channels // 2) |
| 573 | cos_img = cos[pos_ids].flatten(1) |
| 574 | sin_img = sin[pos_ids].flatten(1) |
| 575 | cos_list.append(cos_img) |
| 576 | sin_list.append(sin_img) |
| 577 | |
| 578 | return torch.cat(cos_list, dim=0), torch.cat(sin_list, dim=0) |
| 579 | |
| 580 | def forward(self, x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: |
| 581 | """Args: |
| 582 | x: [total_tokens, C] pre-flattened patches |
| 583 | grid_thw: [num_images, 3] tensor of (t, h, w) for each image/video |
| 584 | Returns: |
| 585 | [merged_tokens, out_hidden_size] |
| 586 | |
| 587 | """ |
| 588 | # Ensure grid_thw is a tensor |
| 589 | if not isinstance(grid_thw, torch.Tensor): |
| 590 | grid_thw = torch.tensor(grid_thw, dtype=torch.long) |
| 591 | |
| 592 | # Move to visual model device/dtype |
| 593 | x = x.to(device=self.device, dtype=self.dtype) |
| 594 | |
| 595 | # Patch embedding: [total_tokens, hidden_size] |
| 596 | x = self.patch_embed(x) |
| 597 | |
| 598 | # Compute 2D rotary positional embeddings |
| 599 | # cos, sin: [total_tokens, qk_channels // 2] |
| 600 | rotary_cos, rotary_sin = self.rot_pos_emb(grid_thw) |
| 601 | rotary_cos = rotary_cos.to(device=x.device) |
| 602 | rotary_sin = rotary_sin.to(device=x.device) |
| 603 | |
| 604 | # Compute cu_seqlens for flash_attn (per-image/video sequence lengths) |
| 605 | seqlens = torch.repeat_interleave( |
| 606 | grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] |
| 607 | ) |
| 608 | cu_seqlens = torch.cat( |
| 609 | [ |
| 610 | torch.tensor([0], device=x.device, dtype=torch.int32), |
| 611 | seqlens.cumsum(dim=0).to(device=x.device, dtype=torch.int32), |
| 612 | ] |
| 613 | ) |
| 614 | max_seqlen = seqlens.max() |
| 615 | |
| 616 | # Precompute col-based window index for type=1 (col SWA) layers |
| 617 | window_index_1d_col = self.get_window_index_1d(grid_thw, col=True).to( |
| 618 | device=x.device |
| 619 | ) |
| 620 | reverse_window_index_1d_col = torch.argsort(window_index_1d_col) |
| 621 | |
| 622 | # Col-based rotary embeddings (reordered at spatial_merge_unit granularity). |
| 623 | # apply_index reorders groups of spatial_merge_unit tokens, just like x. |
| 624 | col_cos = self.apply_index(rotary_cos, window_index_1d_col) |
| 625 | col_sin = self.apply_index(rotary_sin, window_index_1d_col) |
| 626 | |
| 627 | # Add batch dimension: [total_tokens, 1, hidden_size] |
| 628 | x = x.unsqueeze(1) |
| 629 | |
| 630 | for i, blk in enumerate(self.blocks): |
| 631 | window_attn_type = self.vit_window_attn_types[i] |
| 632 | |
| 633 | # Reorder tokens to col-based layout when entering col-SWA region |
| 634 | if window_attn_type == 1 and ( |
| 635 | i == 0 or self.vit_window_attn_types[i - 1] != 1 |
| 636 | ): |
| 637 | x = self.apply_index(x, window_index_1d_col) |
| 638 | |
| 639 | # Restore row-based order when leaving col-SWA region |
| 640 | if ( |
| 641 | i > 0 |
| 642 | and window_attn_type != 1 |
| 643 | and self.vit_window_attn_types[i - 1] == 1 |
| 644 | ): |
| 645 | x = self.apply_index(x, reverse_window_index_1d_col) |
| 646 | |
| 647 | # Use col-based embeddings for col-SWA layers |
| 648 | cos_now = col_cos if window_attn_type == 1 else rotary_cos |
| 649 | sin_now = col_sin if window_attn_type == 1 else rotary_sin |
| 650 | |
| 651 | full_attn = i in self.fullatt_block_indexes |
| 652 | x = blk( |
| 653 | x, |
| 654 | cu_seqlens=cu_seqlens, |
| 655 | rotary_pos_emb_cos=cos_now, |
| 656 | rotary_pos_emb_sin=sin_now, |
| 657 | max_seqlen=max_seqlen, |
| 658 | full_attn=full_attn, |
| 659 | ) |
| 660 | |
| 661 | # Restore row-based order if last block was col-SWA |
| 662 | if self.vit_window_attn_types[-1] == 1: |
| 663 | x = self.apply_index(x, reverse_window_index_1d_col) |
| 664 | |
| 665 | # Remove batch dim and merge spatial tokens |
| 666 | # x: [total_tokens, 1, hidden_size] -> [total_tokens, hidden_size] |
| 667 | x = x.squeeze(1) |
| 668 | x = self.merger(x) |
| 669 | return x |
| 670 | |
| 671 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 672 | loader = AutoWeightsLoader(self) |
| 673 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 674 | |
| 675 | |
| 676 | class MiMoV2OmniMultiModalDataParser(Qwen2VLMultiModalDataParser): |
| 677 | # Video stays raw: get_video_replacement also reads second_per_grid_ts and |
| 678 | # video_start_times, which a metadata-only rewrite would not carry. |
| 679 | embedding_fields = { |
| 680 | "image": Qwen2VLMultiModalDataParser.embedding_fields["image"], |
| 681 | } |
| 682 | |
| 683 | |
| 684 | class MiMoV2OmniProcessingInfo(BaseProcessingInfo): |
| 685 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 686 | return {"audio": None, "image": None, "video": None} |
| 687 | |
| 688 | def get_hf_config(self): |
| 689 | config = self.ctx.get_hf_config() |
| 690 | if isinstance(config.vision_config, dict): |
| 691 | config.vision_config = Mimo_VLVisionConfig.from_dict(config.vision_config) |
| 692 | return config |
| 693 | |
| 694 | def get_hf_processor(self, **kwargs: object) -> MiMoOmniProcessor: |
| 695 | hf_config = self.get_hf_config() |
| 696 | tokenizer = self.get_tokenizer() |
| 697 | return MiMoOmniProcessor.from_hf_config(tokenizer, hf_config) |
| 698 | |
| 699 | def get_image_processor(self, **kwargs: object): |
| 700 | return self.get_hf_processor(**kwargs).image_processor |
| 701 | |
| 702 | def get_data_parser(self): |
| 703 | # Without embedding_fields the EC producer publishes no metadata and |
| 704 | # the EPD proxy never rewrites the media item. |
| 705 | return MiMoV2OmniMultiModalDataParser( |
| 706 | self.get_hf_config().vision_config.spatial_merge_size, |
| 707 | target_sr=24000.0, |
| 708 | expected_hidden_size=self._get_expected_hidden_size(), |
| 709 | allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, |
| 710 | ) |
| 711 | |
| 712 | def get_mm_max_tokens_per_item( |
| 713 | self, |
| 714 | seq_len: int, |
| 715 | mm_counts: Mapping[str, int], |
| 716 | ) -> Mapping[str, int]: |
| 717 | return { |
| 718 | "image": self.get_max_image_tokens(), |
| 719 | "video": self.get_max_video_tokens(seq_len, mm_counts), |
| 720 | } |
| 721 | |
| 722 | def _get_vision_info( |
| 723 | self, |
| 724 | *, |
| 725 | image_width: int, |
| 726 | image_height: int, |
| 727 | num_frames: int = 1, |
| 728 | do_resize: bool = True, |
| 729 | image_processor, |
| 730 | mm_kwargs: Mapping[str, object], |
| 731 | ) -> tuple[ImageSize, int]: |
| 732 | hf_config = self.get_hf_config() |
| 733 | vision_config = hf_config.vision_config |
| 734 | patch_size = vision_config.patch_size |
| 735 | merge_size = vision_config.spatial_merge_size |
| 736 | temporal_patch_size = vision_config.temporal_patch_size |
| 737 | tokens_per_second = vision_config.tokens_per_second |
| 738 | |
| 739 | mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs) |
| 740 | size = image_processor.size |
| 741 | if override_size := mm_kwargs.get("size"): |
| 742 | size = size | override_size |
| 743 | if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None: |
| 744 | size = size | {"shortest_edge": override_min_pixels} |
| 745 | if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None: |
| 746 | size = size | {"longest_edge": override_max_pixels} |
| 747 | |
| 748 | if do_resize: |
| 749 | resized_height, resized_width = smart_resize( |
| 750 | height=image_height, |
| 751 | width=image_width, |
| 752 | factor=patch_size * merge_size, |
| 753 | min_pixels=size["shortest_edge"], |
| 754 | max_pixels=size["longest_edge"], |
| 755 | ) |
| 756 | preprocessed_size = ImageSize(width=resized_width, height=resized_height) |
| 757 | else: |
| 758 | preprocessed_size = ImageSize(width=image_width, height=image_height) |
| 759 | |
| 760 | # For video, MiMo resamples to tokens_per_second fps before temporal patching, |
| 761 | # effective tokens = num_frames * tokens_per_second / temporal_patch_size. |
| 762 | # For images (num_frames == 1) no resampling is applied. |
| 763 | if num_frames > 1: |
| 764 | effective_frames = num_frames * tokens_per_second |
| 765 | else: |
| 766 | effective_frames = num_frames |
| 767 | padded_num_frames = effective_frames + (-effective_frames % temporal_patch_size) |
| 768 | grid_t = max(padded_num_frames // temporal_patch_size, 1) |
| 769 | grid_h = preprocessed_size.height // patch_size |
| 770 | grid_w = preprocessed_size.width // patch_size |
| 771 | num_patches = grid_t * grid_h * grid_w |
| 772 | num_vision_tokens = num_patches // (merge_size**2) |
| 773 | return preprocessed_size, num_vision_tokens |
| 774 | |
| 775 | def get_num_image_tokens( |
| 776 | self, |
| 777 | *, |
| 778 | image_width: int, |
| 779 | image_height: int, |
| 780 | image_processor, |
| 781 | mm_kwargs: Mapping[str, object], |
| 782 | ) -> int: |
| 783 | _, num_image_tokens = self._get_vision_info( |
| 784 | image_width=image_width, |
| 785 | image_height=image_height, |
| 786 | num_frames=1, |
| 787 | image_processor=image_processor, |
| 788 | mm_kwargs=mm_kwargs, |
| 789 | ) |
| 790 | return num_image_tokens |
| 791 | |
| 792 | def get_num_video_tokens( |
| 793 | self, |
| 794 | *, |
| 795 | image_width: int, |
| 796 | image_height: int, |
| 797 | num_frames: int, |
| 798 | image_processor, |
| 799 | mm_kwargs: Mapping[str, object], |
| 800 | ) -> int: |
| 801 | _, num_video_tokens = self._get_vision_info( |
| 802 | image_width=image_width, |
| 803 | image_height=image_height, |
| 804 | num_frames=num_frames, |
| 805 | image_processor=image_processor, |
| 806 | mm_kwargs=mm_kwargs, |
| 807 | ) |
| 808 | return num_video_tokens |
| 809 | |
| 810 | def get_image_size_with_most_features( |
| 811 | self, max_pixels: int | None = None |
| 812 | ) -> ImageSize: |
| 813 | hf_config = self.get_hf_config() |
| 814 | vision_config = hf_config.vision_config |
| 815 | patch_size = vision_config.patch_size |
| 816 | merge_size = vision_config.spatial_merge_size |
| 817 | |
| 818 | if max_pixels is None: |
| 819 | image_processor = self.get_image_processor() |
| 820 | mm_kwargs = self.ctx.get_merged_mm_kwargs({}) |
| 821 | size = image_processor.size |
| 822 | if override_size := mm_kwargs.get("size"): |
| 823 | size = size | override_size |
| 824 | if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None: |
| 825 | size = size | {"shortest_edge": override_min_pixels} |
| 826 | if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None: |
| 827 | size = size | {"longest_edge": override_max_pixels} |
| 828 | max_pixels = size["longest_edge"] |
| 829 | |
| 830 | unit = patch_size * merge_size |
| 831 | max_seq_len = max_pixels // (unit * unit) |
| 832 | |
| 833 | def closest_factor_pair(n: int) -> tuple[int, int]: |
| 834 | for d in range(math.isqrt(n), 0, -1): |
| 835 | if n % d == 0: |
| 836 | return d, n // d |
| 837 | return 1, n |
| 838 | |
| 839 | height_factor, width_factor = 1, max_seq_len |
| 840 | for seq_len in range(max_seq_len, 0, -1): |
| 841 | height_factor, width_factor = closest_factor_pair(seq_len) |
| 842 | if width_factor / height_factor <= 200: |
| 843 | break |
| 844 | |
| 845 | return ImageSize(width=unit * width_factor, height=unit * height_factor) |
| 846 | |
| 847 | def get_max_image_tokens(self) -> int: |
| 848 | image_processor = self.get_image_processor() |
| 849 | target_width, target_height = self.get_image_size_with_most_features() |
| 850 | return self.get_num_image_tokens( |
| 851 | image_width=target_width, |
| 852 | image_height=target_height, |
| 853 | image_processor=image_processor, |
| 854 | mm_kwargs={}, |
| 855 | ) |
| 856 | |
| 857 | def _get_max_video_frames(self, max_tokens: int, start_num_frames: int = 1) -> int: |
| 858 | image_processor = self.get_image_processor() |
| 859 | target_width, target_height = self.get_image_size_with_most_features() |
| 860 | num_frames = start_num_frames |
| 861 | while True: |
| 862 | next_num_frames = num_frames + 1 |
| 863 | next_max_tokens = self.get_num_video_tokens( |
| 864 | image_width=target_width, |
| 865 | image_height=target_height, |
| 866 | num_frames=next_num_frames, |
| 867 | image_processor=image_processor, |
| 868 | mm_kwargs={}, |
| 869 | ) |
| 870 | if next_max_tokens > max_tokens: |
| 871 | break |
| 872 | num_frames = next_num_frames |
| 873 | return num_frames |
| 874 | |
| 875 | def get_num_frames_with_most_features( |
| 876 | self, |
| 877 | seq_len: int, |
| 878 | mm_counts: Mapping[str, int], |
| 879 | max_frames_per_video: int = 14, |
| 880 | ) -> int: |
| 881 | max_videos = mm_counts.get("video", 0) |
| 882 | max_total_frames = self._get_max_video_frames(seq_len) |
| 883 | max_frames_per_video = min( |
| 884 | max_total_frames // max(max_videos, 1), max_frames_per_video |
| 885 | ) |
| 886 | return max(max_frames_per_video, 1) |
| 887 | |
| 888 | def get_max_video_tokens( |
| 889 | self, |
| 890 | seq_len: int, |
| 891 | mm_counts: Mapping[str, int], |
| 892 | ) -> int: |
| 893 | image_processor = self.get_image_processor() |
| 894 | target_width, target_height = self.get_image_size_with_most_features() |
| 895 | return self.get_num_video_tokens( |
| 896 | image_width=target_width, |
| 897 | image_height=target_height, |
| 898 | num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts), |
| 899 | image_processor=image_processor, |
| 900 | mm_kwargs={}, |
| 901 | ) |
| 902 | |
| 903 | |
| 904 | class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessingInfo]): |
| 905 | """vLLM multimodal processor for MiMo-Omni (image + video). |
| 906 | |
| 907 | Key differences from Qwen2.5-VL: |
| 908 | - Videos use timestamp tokens between temporal grid positions. |
| 909 | - The HF processor expects ``(TCHW_tensor, timestamps_T_tensor)`` video |
| 910 | tuples rather than plain numpy arrays. |
| 911 | - ``video_start_times`` is tracked so prompt-update reconstruction can |
| 912 | regenerate the exact same timestamp token IDs. |
| 913 | """ |
| 914 | |
| 915 | # fps assumed for vllm-decoded video (numpy T,H,W,C arrays). |
| 916 | # The video loader samples ~32 frames; treat each frame as 1 s apart so |
| 917 | # MiMoVLProcessor sees 1 fps input and resamples internally. |
| 918 | _INPUT_FPS: float = 1.0 |
| 919 | |
| 920 | def _get_mm_fields_config( |
| 921 | self, |
| 922 | hf_inputs: BatchFeature, |
| 923 | hf_processor_mm_kwargs: Mapping[str, object], |
| 924 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 925 | merge_size = self.info.get_hf_config().vision_config.spatial_merge_size |
| 926 | fields: dict[str, MultiModalFieldConfig] = dict( |
| 927 | **_create_qwen2vl_field_factory(merge_size)(hf_inputs), |
| 928 | second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True), |
| 929 | video_start_times=MultiModalFieldConfig.batched("video", keep_on_cpu=True), |
| 930 | audio_features=MultiModalFieldConfig.batched("audio"), |
| 931 | audio_token_lens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), |
| 932 | ) |
| 933 | # video_audio fields: only present when video_audio content was processed |
| 934 | if "video_audio_n_segs" in hf_inputs: |
| 935 | fields["video_audio_n_segs"] = MultiModalFieldConfig.batched( |
| 936 | "video", keep_on_cpu=True |
| 937 | ) |
| 938 | # video_audio_seg_lens: list of per-video 1D tensors, batched("video") |
| 939 | if "video_audio_seg_lens" in hf_inputs: |
| 940 | fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched( |
| 941 | "video", keep_on_cpu=True |
| 942 | ) |
| 943 | if "va_audio_features" in hf_inputs: |
| 944 | fields["va_audio_features"] = MultiModalFieldConfig.batched("va_audio") |
| 945 | return fields |
| 946 | |
| 947 | def _get_hf_mm_inputs( |
| 948 | self, |
| 949 | mm_items: MultiModalDataItems, |
| 950 | hf_kwargs: Mapping[str, object], |
| 951 | ) -> HFMultiModalInputs: |
| 952 | """Convert numpy video arrays to (TCHW, timestamps) tuples for MiMo.""" |
| 953 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 954 | mm_data = hf_inputs.hf_data |
| 955 | |
| 956 | # Handle video_audio items: convert video part to (TCHW, timestamps) tuple |
| 957 | if "video_audio" in mm_data: |
| 958 | va_converted: list[VideoAudioInput] = [] |
| 959 | video_audio_items = mm_data["video_audio"] |
| 960 | assert isinstance(video_audio_items, Sequence) |
| 961 | for va_item in video_audio_items: |
| 962 | if isinstance(va_item, VideoAudioInput): |
| 963 | vid = va_item.video |
| 964 | else: |
| 965 | # Expect (video_frames, audio_source) tuple |
| 966 | assert isinstance(va_item, (list, tuple)) and len(va_item) == 2 |
| 967 | vid, audio_src = va_item |
| 968 | va_item = VideoAudioInput(video=vid, audio=audio_src) |
| 969 | # Convert video frames to (TCHW, timestamps) if needed |
| 970 | if ( |
| 971 | isinstance(vid, tuple) |
| 972 | and len(vid) == 2 |
| 973 | and isinstance(vid[0], torch.Tensor) |
| 974 | and isinstance(vid[1], torch.Tensor) |
| 975 | ): |
| 976 | va_converted.append(va_item) |
| 977 | else: |
| 978 | if isinstance(vid, np.ndarray): |
| 979 | frames = torch.from_numpy(vid) |
| 980 | elif isinstance(vid, torch.Tensor): |
| 981 | frames = vid |
| 982 | else: |
| 983 | frames = torch.tensor(np.array(vid)) |
| 984 | if frames.ndim == 4 and frames.shape[-1] in (1, 3, 4): |
| 985 | frames = frames.permute(0, 3, 1, 2).float() |
| 986 | else: |
| 987 | frames = frames.float() |
| 988 | T = frames.shape[0] |
| 989 | timestamps = torch.arange(T, dtype=torch.float32) / self._INPUT_FPS |
| 990 | va_converted.append( |
| 991 | VideoAudioInput( |
| 992 | video=(frames, timestamps), |
| 993 | audio=va_item.audio, |
| 994 | ) |
| 995 | ) |
| 996 | mm_data["video_audio"] = va_converted |
| 997 | |
| 998 | if "videos" in mm_data: |
| 999 | converted: list[tuple[torch.Tensor, torch.Tensor]] = [] |
| 1000 | videos = mm_data["videos"] |
| 1001 | assert isinstance(videos, Sequence) |
| 1002 | for video in videos: |
| 1003 | if ( |
| 1004 | isinstance(video, tuple) |
| 1005 | and len(video) == 2 |
| 1006 | and isinstance(video[0], torch.Tensor) |
| 1007 | and isinstance(video[1], torch.Tensor) |
| 1008 | ): |
| 1009 | # already in MiMo format |
| 1010 | converted.append(video) |
| 1011 | else: |
| 1012 | # numpy (T, H, W, C) or torch (T, H, W, C) / (T, C, H, W) |
| 1013 | if isinstance(video, np.ndarray): |
| 1014 | frames = torch.from_numpy(video) |
| 1015 | elif isinstance(video, torch.Tensor): |
| 1016 | frames = video |
| 1017 | else: |
| 1018 | frames = torch.tensor(np.array(video)) |
| 1019 | |
| 1020 | if frames.ndim == 4 and frames.shape[-1] in (1, 3, 4): |
| 1021 | # THWC → TCHW |
| 1022 | frames = frames.permute(0, 3, 1, 2).float() |
| 1023 | else: |
| 1024 | frames = frames.float() |
| 1025 | |
| 1026 | T = frames.shape[0] |
| 1027 | timestamps = torch.arange(T, dtype=torch.float32) / self._INPUT_FPS |
| 1028 | converted.append((frames, timestamps)) |
| 1029 | |
| 1030 | mm_data["videos"] = converted |
| 1031 | |
| 1032 | return hf_inputs |
| 1033 | |
| 1034 | def _get_prompt_updates( |
| 1035 | self, |
| 1036 | mm_items: MultiModalDataItems, |
| 1037 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 1038 | out_mm_kwargs: MultiModalKwargsItems, |
| 1039 | ) -> Sequence[PromptUpdate]: |
| 1040 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 1041 | hf_config = self.info.get_hf_config() |
| 1042 | tokenizer = self.info.get_tokenizer() |
| 1043 | vocab = tokenizer.get_vocab() |
| 1044 | |
| 1045 | merge_size = hf_config.vision_config.spatial_merge_size |
| 1046 | p = hf_processor.mimo_processor |
| 1047 | |
| 1048 | image_pad_id = vocab[hf_processor.image_token] |
| 1049 | video_pad_id = vocab[hf_processor.video_token] |
| 1050 | audio_pad_id = vocab.get("<|audio_pad|>") |
| 1051 | vision_start_id = p.vision_start_token_id |
| 1052 | vision_end_id = p.vision_end_token_id |
| 1053 | video_start_id = p.video_start_token_id |
| 1054 | video_end_id = p.video_end_token_id |
| 1055 | audio_start_id = p.audio_start_token_id |
| 1056 | audio_end_id = p.audio_end_token_id |
| 1057 | assert isinstance(vision_start_id, int) |
| 1058 | assert isinstance(vision_end_id, int) |
| 1059 | assert isinstance(video_start_id, int) |
| 1060 | assert isinstance(video_end_id, int) |
| 1061 | |
| 1062 | def get_image_replacement(item_idx: int) -> PromptUpdateDetails: |
| 1063 | out_item = out_mm_kwargs["image"][item_idx] |
| 1064 | grid_thw = out_item["image_grid_thw"].data |
| 1065 | assert isinstance(grid_thw, torch.Tensor) |
| 1066 | n_tokens = int(grid_thw.prod()) // merge_size**2 |
| 1067 | return PromptUpdateDetails.from_seq([image_pad_id] * n_tokens) |
| 1068 | |
| 1069 | def get_video_replacement(item_idx: int) -> PromptUpdateDetails: |
| 1070 | out_item = out_mm_kwargs["video"][item_idx] |
| 1071 | grid_thw = out_item["video_grid_thw"].data |
| 1072 | second_per_grid_ts = out_item["second_per_grid_ts"].data |
| 1073 | video_start_time = out_item["video_start_times"].data |
| 1074 | assert isinstance(grid_thw, torch.Tensor) |
| 1075 | assert isinstance(second_per_grid_ts, (int, float, torch.Tensor)) |
| 1076 | assert isinstance(video_start_time, (int, float, torch.Tensor)) |
| 1077 | spt = float(second_per_grid_ts) |
| 1078 | start = float(video_start_time) |
| 1079 | |
| 1080 | T, H, W = map(int, grid_thw) |
| 1081 | n_per_grid = H * W // (merge_size * merge_size) |
| 1082 | |
| 1083 | # Check if this is a video_audio item |
| 1084 | n_segs_field = out_item.get("video_audio_n_segs") |
| 1085 | if n_segs_field is None: |
| 1086 | n_segs_val = 0 |
| 1087 | else: |
| 1088 | n_segs_data = n_segs_field.data |
| 1089 | assert isinstance(n_segs_data, (int, torch.Tensor)) |
| 1090 | n_segs_val = int(n_segs_data) |
| 1091 | va_seg_lens: list[int] | None = None |
| 1092 | if n_segs_val > 0: |
| 1093 | seg_lens_field = out_item.get("video_audio_seg_lens") |
| 1094 | if seg_lens_field is not None: |
| 1095 | seg_lens_data = seg_lens_field.data |
| 1096 | assert isinstance(seg_lens_data, torch.Tensor) |
| 1097 | va_seg_lens = seg_lens_data[:n_segs_val].tolist() |
| 1098 | |
| 1099 | full: list[int] = [video_start_id] |
| 1100 | is_embed_mask: list[bool] = [False] |
| 1101 | |
| 1102 | if va_seg_lens is None: |
| 1103 | # Regular video: timestamp + vision tokens per grid |
| 1104 | for j in range(T): |
| 1105 | ts_text = _format_timestamp(start + j * spt) |
| 1106 | ts_ids = tokenizer.encode(ts_text, add_special_tokens=False) |
| 1107 | full.extend(ts_ids) |
| 1108 | is_embed_mask.extend([False] * len(ts_ids)) |
| 1109 | full.append(vision_start_id) |
| 1110 | is_embed_mask.append(False) |
| 1111 | full.extend([video_pad_id] * n_per_grid) |
| 1112 | is_embed_mask.extend([True] * n_per_grid) |
| 1113 | full.append(vision_end_id) |
| 1114 | is_embed_mask.append(False) |
| 1115 | else: |
| 1116 | # video_audio: interleaved vision+audio per group |
| 1117 | n_groups = len(va_seg_lens) |
| 1118 | frames_per_group = T // n_groups # 1 for il=0, T for il=-1 |
| 1119 | for g in range(n_groups): |
| 1120 | # Timestamp for first frame of this group |
| 1121 | frame0 = g * frames_per_group |
| 1122 | ts_text = _format_timestamp(start + frame0 * spt) |
| 1123 | ts_ids = tokenizer.encode(ts_text, add_special_tokens=False) |
| 1124 | full.extend(ts_ids) |
| 1125 | is_embed_mask.extend([False] * len(ts_ids)) |
| 1126 | # Vision tokens for all frames in this group |
| 1127 | for f in range(frames_per_group): |
| 1128 | full.append(vision_start_id) |
| 1129 | is_embed_mask.append(False) |
| 1130 | full.extend([video_pad_id] * n_per_grid) |
| 1131 | is_embed_mask.extend([True] * n_per_grid) |
| 1132 | full.append(vision_end_id) |
| 1133 | is_embed_mask.append(False) |
| 1134 | # Audio tokens for this group |
| 1135 | seg_len = va_seg_lens[g] |
| 1136 | assert audio_start_id is not None |
| 1137 | assert audio_pad_id is not None |
| 1138 | assert audio_end_id is not None |
| 1139 | full.append(audio_start_id) |
| 1140 | is_embed_mask.append(False) |
| 1141 | full.extend([audio_pad_id] * seg_len) |
| 1142 | is_embed_mask.extend([True] * seg_len) |
| 1143 | full.append(audio_end_id) |
| 1144 | is_embed_mask.append(False) |
| 1145 | |
| 1146 | full.append(video_end_id) |
| 1147 | is_embed_mask.append(False) |
| 1148 | |
| 1149 | embed_t = torch.tensor(is_embed_mask) |
| 1150 | return PromptUpdateDetails( |
| 1151 | full=full, |
| 1152 | is_embed=lambda _seq: embed_t, |
| 1153 | ) |
| 1154 | |
| 1155 | def get_audio_replacement(item_idx: int) -> PromptUpdateDetails: |
| 1156 | out_item = out_mm_kwargs["audio"][item_idx] |
| 1157 | audio_token_lens = out_item["audio_token_lens"].data |
| 1158 | assert isinstance(audio_token_lens, (int, torch.Tensor)) |
| 1159 | assert audio_pad_id is not None |
| 1160 | tok_len = int(audio_token_lens) |
| 1161 | return PromptUpdateDetails.from_seq([audio_pad_id] * tok_len) |
| 1162 | |
| 1163 | updates: list[PromptUpdate] = [ |
| 1164 | PromptReplacement( |
| 1165 | modality="image", |
| 1166 | target=[image_pad_id], |
| 1167 | replacement=get_image_replacement, |
| 1168 | ), |
| 1169 | PromptReplacement( |
| 1170 | modality="video", |
| 1171 | target=[video_pad_id], |
| 1172 | replacement=get_video_replacement, |
| 1173 | ), |
| 1174 | ] |
| 1175 | if audio_pad_id is not None and audio_start_id is not None: |
| 1176 | updates.append( |
| 1177 | PromptReplacement( |
| 1178 | modality="audio", |
| 1179 | target=[audio_pad_id], |
| 1180 | replacement=get_audio_replacement, |
| 1181 | ) |
| 1182 | ) |
| 1183 | return updates |
| 1184 | |
| 1185 | |
| 1186 | class MiMoV2OmniDummyInputsBuilder(BaseDummyInputsBuilder[MiMoV2OmniProcessingInfo]): |
| 1187 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1188 | num_images = mm_counts.get("image", 0) |
| 1189 | num_videos = mm_counts.get("video", 0) |
| 1190 | num_audios = mm_counts.get("audio", 0) |
| 1191 | image_ph = "<|vision_start|><|image_pad|><|vision_end|>" |
| 1192 | video_ph = "<|vision_start|><|video_pad|><|vision_end|>" |
| 1193 | audio_ph = "<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>" |
| 1194 | return image_ph * num_images + video_ph * num_videos + audio_ph * num_audios |
| 1195 | |
| 1196 | def get_dummy_mm_data( |
| 1197 | self, |
| 1198 | seq_len: int, |
| 1199 | mm_counts: Mapping[str, int], |
| 1200 | mm_options: MultiModalDummyOptions, |
| 1201 | ) -> MultiModalDataDict: |
| 1202 | target_width, target_height = self.info.get_image_size_with_most_features() |
| 1203 | target_num_frames = self.info.get_num_frames_with_most_features( |
| 1204 | seq_len, mm_counts |
| 1205 | ) |
| 1206 | |
| 1207 | return { |
| 1208 | "image": self._get_dummy_images( |
| 1209 | width=target_width, |
| 1210 | height=target_height, |
| 1211 | num_images=mm_counts.get("image", 0), |
| 1212 | overrides=mm_options.get("image"), |
| 1213 | ), |
| 1214 | "video": self._get_dummy_videos( |
| 1215 | width=target_width, |
| 1216 | height=target_height, |
| 1217 | num_frames=target_num_frames, |
| 1218 | num_videos=mm_counts.get("video", 0), |
| 1219 | overrides=mm_options.get("video"), |
| 1220 | ), |
| 1221 | } |
| 1222 | |
| 1223 | |
| 1224 | @MULTIMODAL_REGISTRY.register_processor( |
| 1225 | MiMoV2OmniMultiModalProcessor, |
| 1226 | info=MiMoV2OmniProcessingInfo, |
| 1227 | dummy_inputs=MiMoV2OmniDummyInputsBuilder, |
| 1228 | ) |
| 1229 | class MiMoV2OmniForCausalLM( |
| 1230 | nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant, SupportsEagle3 |
| 1231 | ): |
| 1232 | # To ensure correct weight loading and mapping. |
| 1233 | hf_to_vllm_mapper = WeightsMapper( |
| 1234 | orig_to_new_prefix={ |
| 1235 | # audio encoder |
| 1236 | "speech_embeddings.": "audio_encoder.speech_embeddings.", |
| 1237 | # mapping for new names in checkpoint saved after transformers v4.52 |
| 1238 | "model.language_model.": "language_model.model.", |
| 1239 | "model.visual.": "visual.", |
| 1240 | # mapping for original checkpoint |
| 1241 | "lm_head.": "language_model.lm_head.", |
| 1242 | "model.": "language_model.model.", |
| 1243 | "audio_tokenizer.": None, |
| 1244 | } |
| 1245 | ) |
| 1246 | |
| 1247 | @classmethod |
| 1248 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 1249 | if modality.startswith("image"): |
| 1250 | return "<|vision_start|><|image_pad|><|vision_end|>" |
| 1251 | if modality.startswith("video"): |
| 1252 | return "<|vision_start|><|video_pad|><|vision_end|>" |
| 1253 | if modality.startswith("audio"): |
| 1254 | return "<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>" |
| 1255 | |
| 1256 | raise ValueError(f"Unsupported modality: {modality}") |
| 1257 | |
| 1258 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1259 | super().__init__() |
| 1260 | config = vllm_config.model_config.hf_config |
| 1261 | self.config = config |
| 1262 | # Omni ViT/Audio Encoder BF16 |
| 1263 | vision_config = ( |
| 1264 | Mimo_VLVisionConfig.from_dict(config.vision_config) |
| 1265 | if isinstance(config.vision_config, dict) |
| 1266 | else config.vision_config |
| 1267 | ) |
| 1268 | with self._mark_tower_model(vllm_config, {"image", "video"}): |
| 1269 | self.visual = MiMoVisionTransformer( |
| 1270 | vision_config, |
| 1271 | norm_eps=getattr(vllm_config, "rms_norm_eps", 1e-6), |
| 1272 | quant_config=None, |
| 1273 | prefix=maybe_prefix(prefix, "visual"), |
| 1274 | ) |
| 1275 | audio_config = getattr(config, "audio_config", None) |
| 1276 | model_path = vllm_config.model_config.model |
| 1277 | self.audio_encoder: MimoAudioEncoder | None |
| 1278 | if audio_config is not None: |
| 1279 | with self._mark_tower_model(vllm_config, "audio"): |
| 1280 | self.audio_encoder = MimoAudioEncoder( |
| 1281 | audio_config, model_path=model_path |
| 1282 | ) |
| 1283 | else: |
| 1284 | self.audio_encoder = None |
| 1285 | with self._mark_language_model(vllm_config): |
| 1286 | self.language_model = MiMoV2FlashForCausalLM( |
| 1287 | vllm_config=vllm_config, |
| 1288 | prefix=maybe_prefix(prefix, "language_model"), |
| 1289 | ) |
| 1290 | |
| 1291 | self.make_empty_intermediate_tensors = ( |
| 1292 | self.language_model.make_empty_intermediate_tensors |
| 1293 | ) |
| 1294 | |
| 1295 | def _parse_and_validate_image_input( |
| 1296 | self, **kwargs: object |
| 1297 | ) -> Qwen2_5_VLImageInputs | None: |
| 1298 | pixel_values = kwargs.pop("pixel_values", None) |
| 1299 | image_embeds = kwargs.pop("image_embeds", None) |
| 1300 | image_grid_thw = kwargs.pop("image_grid_thw", None) |
| 1301 | |
| 1302 | if pixel_values is None and image_embeds is None: |
| 1303 | return None |
| 1304 | |
| 1305 | if pixel_values is not None: |
| 1306 | return Qwen2_5_VLImagePixelInputs( |
| 1307 | type="pixel_values", |
| 1308 | pixel_values=pixel_values, |
| 1309 | image_grid_thw=image_grid_thw, |
| 1310 | ) |
| 1311 | |
| 1312 | assert image_embeds is not None |
| 1313 | return Qwen2_5_VLImageEmbeddingInputs( |
| 1314 | type="image_embeds", |
| 1315 | image_embeds=image_embeds, |
| 1316 | image_grid_thw=image_grid_thw, |
| 1317 | ) |
| 1318 | |
| 1319 | def _parse_and_validate_video_input( |
| 1320 | self, **kwargs: object |
| 1321 | ) -> Qwen2_5_VLVideoInputs | None: |
| 1322 | pixel_values_videos = kwargs.pop("pixel_values_videos", None) |
| 1323 | video_embeds = kwargs.pop("video_embeds", None) |
| 1324 | video_grid_thw = kwargs.pop("video_grid_thw", None) |
| 1325 | second_per_grid_ts = kwargs.pop("second_per_grid_ts", None) |
| 1326 | |
| 1327 | if pixel_values_videos is None and video_embeds is None: |
| 1328 | return None |
| 1329 | |
| 1330 | if pixel_values_videos is not None: |
| 1331 | return Qwen2_5_VLVideoPixelInputs( |
| 1332 | type="pixel_values_videos", |
| 1333 | pixel_values_videos=pixel_values_videos, |
| 1334 | video_grid_thw=video_grid_thw, |
| 1335 | second_per_grid_ts=second_per_grid_ts, |
| 1336 | ) |
| 1337 | |
| 1338 | assert video_embeds is not None |
| 1339 | return Qwen2_5_VLVideoEmbeddingInputs( |
| 1340 | type="video_embeds", |
| 1341 | video_embeds=video_embeds, |
| 1342 | video_grid_thw=video_grid_thw, |
| 1343 | second_per_grid_ts=second_per_grid_ts, |
| 1344 | ) |
| 1345 | |
| 1346 | def _process_image_input( |
| 1347 | self, image_input: Qwen2_5_VLImageInputs |
| 1348 | ) -> tuple[torch.Tensor, ...]: |
| 1349 | grid_thw = image_input["image_grid_thw"] |
| 1350 | assert grid_thw.ndim == 2 |
| 1351 | grid_thw_list = grid_thw.tolist() |
| 1352 | |
| 1353 | if image_input["type"] == "image_embeds": |
| 1354 | image_embeds = image_input["image_embeds"].type(self.visual.dtype) |
| 1355 | else: |
| 1356 | pixel_values = image_input["pixel_values"] |
| 1357 | image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list) |
| 1358 | |
| 1359 | # Split concatenated embeddings for each image item. |
| 1360 | merge_size = self.visual.spatial_merge_size |
| 1361 | sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist() |
| 1362 | return image_embeds.split(sizes) |
| 1363 | |
| 1364 | def _process_video_input( |
| 1365 | self, video_input: Qwen2_5_VLVideoInputs |
| 1366 | ) -> tuple[torch.Tensor, ...]: |
| 1367 | grid_thw = video_input["video_grid_thw"] |
| 1368 | assert grid_thw.ndim == 2 |
| 1369 | grid_thw_list = grid_thw.tolist() |
| 1370 | |
| 1371 | if video_input["type"] == "video_embeds": |
| 1372 | video_embeds = video_input["video_embeds"].type(self.visual.dtype) |
| 1373 | else: |
| 1374 | pixel_values_videos = video_input["pixel_values_videos"] |
| 1375 | video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw_list) |
| 1376 | |
| 1377 | # Split concatenated embeddings for each video item. |
| 1378 | merge_size = self.visual.spatial_merge_size |
| 1379 | sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist() |
| 1380 | return video_embeds.split(sizes) |
| 1381 | |
| 1382 | def _parse_and_validate_audio_input( |
| 1383 | self, **kwargs: object |
| 1384 | ) -> MiMoV2OmniAudioInputs | None: |
| 1385 | audio_features = kwargs.pop("audio_features", None) |
| 1386 | audio_token_lens = kwargs.pop("audio_token_lens", None) |
| 1387 | if audio_features is None: |
| 1388 | return None |
| 1389 | assert isinstance(audio_features, (torch.Tensor, list)) |
| 1390 | if isinstance(audio_features, list): |
| 1391 | assert all(isinstance(item, torch.Tensor) for item in audio_features) |
| 1392 | assert audio_token_lens is None or isinstance( |
| 1393 | audio_token_lens, (torch.Tensor, list) |
| 1394 | ) |
| 1395 | return { |
| 1396 | "type": "audio", |
| 1397 | "audio_features": audio_features, |
| 1398 | "audio_token_lens": audio_token_lens, |
| 1399 | } |
| 1400 | |
| 1401 | def _parse_and_validate_multimodal_inputs( |
| 1402 | self, **kwargs: object |
| 1403 | ) -> MiMoV2OmniMultiModalInputs: |
| 1404 | mm_input_by_modality: MiMoV2OmniMultiModalInputs = {} |
| 1405 | |
| 1406 | # Preserve the order of modalities if there are multiple of them |
| 1407 | # from the order of kwargs. |
| 1408 | for input_key in kwargs: |
| 1409 | if ( |
| 1410 | input_key in ("pixel_values", "image_embeds") |
| 1411 | and "image" not in mm_input_by_modality |
| 1412 | ): |
| 1413 | mm_input_by_modality["image"] = self._parse_and_validate_image_input( |
| 1414 | **kwargs |
| 1415 | ) |
| 1416 | if ( |
| 1417 | input_key in ("pixel_values_videos", "video_embeds") |
| 1418 | and "video" not in mm_input_by_modality |
| 1419 | ): |
| 1420 | mm_input_by_modality["video"] = self._parse_and_validate_video_input( |
| 1421 | **kwargs |
| 1422 | ) |
| 1423 | if input_key == "audio_features" and "audio" not in mm_input_by_modality: |
| 1424 | mm_input_by_modality["audio"] = self._parse_and_validate_audio_input( |
| 1425 | **kwargs |
| 1426 | ) |
| 1427 | return mm_input_by_modality |
| 1428 | |
| 1429 | def _process_audio_input( |
| 1430 | self, audio_input: MiMoV2OmniAudioInputs |
| 1431 | ) -> tuple[torch.Tensor, ...]: |
| 1432 | mel_specs = audio_input["audio_features"] |
| 1433 | if self.audio_encoder is None: |
| 1434 | return () |
| 1435 | # Normalize to List[2D-Tensor]. |
| 1436 | # MultiModalBatchedField._reduce_data either wraps a single [T, 128] |
| 1437 | # into [1, T, 128] via unsqueeze(0) or stacks N same-T items into |
| 1438 | # [N, T, 128]. Indexing along dim-0 extracts the per-item [T, 128]. |
| 1439 | if isinstance(mel_specs, torch.Tensor): |
| 1440 | mel_specs = list(mel_specs) # [1,T,128] or [N,T,128] → [[T,128],...] |
| 1441 | if not mel_specs: |
| 1442 | return () |
| 1443 | audio_embeds, item_token_lens = self.audio_encoder.get_audio_feature(mel_specs) |
| 1444 | return tuple(audio_embeds.split(item_token_lens)) |
| 1445 | |
| 1446 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1447 | # Pop video_audio-specific fields before main mm parsing |
| 1448 | video_audio_n_segs = kwargs.pop("video_audio_n_segs", None) |
| 1449 | video_audio_seg_lens = kwargs.pop("video_audio_seg_lens", None) |
| 1450 | va_audio_features = kwargs.pop("va_audio_features", None) |
| 1451 | |
| 1452 | mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) |
| 1453 | if not mm_input_by_modality and va_audio_features is None: |
| 1454 | return [] |
| 1455 | |
| 1456 | # The result multimodal_embeddings is tuple of tensors, with each |
| 1457 | # tensor corresponding to a multimodal data item (image, video, or audio). |
| 1458 | multimodal_embeddings: list[torch.Tensor] = [] |
| 1459 | |
| 1460 | # Pre-process va audio: one mel spec per va video → per-video audio embeddings |
| 1461 | # keyed by va video index (0-based among va videos only) |
| 1462 | va_audio_embs_list: list[torch.Tensor] = [] |
| 1463 | if va_audio_features is not None and self.audio_encoder is not None: |
| 1464 | assert isinstance(va_audio_features, (torch.Tensor, list, tuple)) |
| 1465 | mel_list: list[torch.Tensor] = [] |
| 1466 | for mel_spec in va_audio_features: |
| 1467 | assert isinstance(mel_spec, torch.Tensor) |
| 1468 | mel_list.append(mel_spec) |
| 1469 | for mel_spec in mel_list: |
| 1470 | embs, tok_lens = self.audio_encoder.get_audio_feature([mel_spec]) |
| 1471 | # tok_lens is a list/tensor with one entry (total tokens for this mel) |
| 1472 | va_audio_embs_list.append(embs) # shape (total_tok, hidden) |
| 1473 | |
| 1474 | va_cursor = 0 # index into va_audio_embs_list |
| 1475 | assert video_audio_n_segs is None or isinstance( |
| 1476 | video_audio_n_segs, torch.Tensor |
| 1477 | ) |
| 1478 | assert video_audio_seg_lens is None or isinstance( |
| 1479 | video_audio_seg_lens, torch.Tensor |
| 1480 | ) |
| 1481 | |
| 1482 | # NOTE: Iterate in dict insertion order to preserve token sequence order. |
| 1483 | for modality in mm_input_by_modality: |
| 1484 | if modality == "image": |
| 1485 | image_input = mm_input_by_modality["image"] |
| 1486 | assert isinstance( |
| 1487 | image_input, |
| 1488 | (Qwen2_5_VLImagePixelInputs, Qwen2_5_VLImageEmbeddingInputs), |
| 1489 | ) |
| 1490 | multimodal_embeddings.extend(self._process_image_input(image_input)) |
| 1491 | elif modality == "video": |
| 1492 | video_input = mm_input_by_modality["video"] |
| 1493 | assert isinstance( |
| 1494 | video_input, |
| 1495 | (Qwen2_5_VLVideoPixelInputs, Qwen2_5_VLVideoEmbeddingInputs), |
| 1496 | ) |
| 1497 | video_embs_tuple = self._process_video_input(video_input) |
| 1498 | if video_audio_n_segs is None: |
| 1499 | multimodal_embeddings.extend(video_embs_tuple) |
| 1500 | else: |
| 1501 | grid_thw = video_input["video_grid_thw"] |
| 1502 | for i, vid_embs in enumerate(video_embs_tuple): |
| 1503 | n_segs = int(video_audio_n_segs[i]) |
| 1504 | if n_segs == 0 or not va_audio_embs_list: |
| 1505 | multimodal_embeddings.append(vid_embs) |
| 1506 | else: |
| 1507 | assert video_audio_seg_lens is not None |
| 1508 | T = int(grid_thw[i][0]) |
| 1509 | n_per_grid = vid_embs.shape[0] // T |
| 1510 | frames = list(vid_embs.split(n_per_grid, dim=0)) |
| 1511 | frames_per_group = T // n_segs |
| 1512 | # Per-group audio token lengths for this va video |
| 1513 | # video_audio_seg_lens is (num_videos, max_T); row i |
| 1514 | # has valid values in [:n_segs], rest are zeros. |
| 1515 | seg_lens = video_audio_seg_lens[i][:n_segs].tolist() |
| 1516 | # Split full audio embs for this va video by group lengths |
| 1517 | full_va_embs = va_audio_embs_list[va_cursor] |
| 1518 | va_cursor += 1 |
| 1519 | group_audio_embs = full_va_embs.split(seg_lens) |
| 1520 | # Interleave: all vid frames in group, then audio for group |
| 1521 | for g in range(n_segs): |
| 1522 | for f in range(frames_per_group): |
| 1523 | multimodal_embeddings.append( |
| 1524 | frames[g * frames_per_group + f] |
| 1525 | ) |
| 1526 | multimodal_embeddings.append(group_audio_embs[g]) |
| 1527 | elif modality == "audio": |
| 1528 | audio_input = mm_input_by_modality["audio"] |
| 1529 | assert audio_input is not None |
| 1530 | multimodal_embeddings.extend(self._process_audio_input(audio_input)) |
| 1531 | return tuple(multimodal_embeddings) |
| 1532 | |
| 1533 | def forward( |
| 1534 | self, |
| 1535 | input_ids: torch.Tensor | None, |
| 1536 | positions: torch.Tensor, |
| 1537 | intermediate_tensors: IntermediateTensors | None = None, |
| 1538 | inputs_embeds: torch.Tensor | None = None, |
| 1539 | **kwargs: object, |
| 1540 | ) -> torch.Tensor | IntermediateTensors: |
| 1541 | """Run forward pass for Qwen2.5-VL. |
| 1542 | |
| 1543 | Args: |
| 1544 | input_ids: Flattened (concatenated) input_ids corresponding to a |
| 1545 | batch. |
| 1546 | positions: Flattened (concatenated) position ids corresponding to a |
| 1547 | batch. **NOTE**: If mrope is enabled (default setting for |
| 1548 | Qwen2.5-VL opensource models), the shape will be `(3, seq_len)`, |
| 1549 | otherwise it will be `(seq_len,). |
| 1550 | intermediate_tensors: Intermediate tensors from prior forward pass. |
| 1551 | inputs_embeds: Optional tensor of input embeddings. |
| 1552 | |
| 1553 | """ |
| 1554 | if intermediate_tensors is not None: |
| 1555 | inputs_embeds = None |
| 1556 | |
| 1557 | hidden_states = self.language_model.model( |
| 1558 | input_ids=input_ids, |
| 1559 | positions=positions, |
| 1560 | intermediate_tensors=intermediate_tensors, |
| 1561 | inputs_embeds=inputs_embeds, |
| 1562 | ) |
| 1563 | return hidden_states |
| 1564 | |
| 1565 | def compute_logits( |
| 1566 | self, |
| 1567 | hidden_states: torch.Tensor, |
| 1568 | ) -> torch.Tensor | None: |
| 1569 | return self.language_model.compute_logits(hidden_states) |
| 1570 | |
| 1571 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1572 | audio_loaded: set[str] = set() |
| 1573 | |
| 1574 | loader = AutoWeightsLoader(self) |
| 1575 | auto_loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 1576 | return audio_loaded | auto_loaded |