main
py 1,048 lines 36.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Adapted from
5 # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6 # Copyright 2023 The vLLM team.
7 # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8 #
9 # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10 # and OPT implementations in this library. It has been modified from its
11 # original forms to accommodate minor architectural differences compared
12 # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13 #
14 # Licensed under the Apache License, Version 2.0 (the "License");
15 # you may not use this file except in compliance with the License.
16 # You may obtain a copy of the License at
17 #
18 # http://www.apache.org/licenses/LICENSE-2.0
19 #
20 # Unless required by applicable law or agreed to in writing, software
21 # distributed under the License is distributed on an "AS IS" BASIS,
22 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 # See the License for the specific language governing permissions and
24 # limitations under the License.
25 """Inference-only MiniCPM-O model compatible with HuggingFace weights."""
26
27 import os
28 from collections.abc import Callable, Iterable, Mapping, Sequence
29 from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias
30
31 import torch
32 from torch import nn
33 from transformers import BatchFeature
34 from transformers.modeling_outputs import BaseModelOutputWithPast
35 from transformers.models.whisper.modeling_whisper import (
36 ACT2FN,
37 WhisperAttention,
38 WhisperConfig,
39 WhisperEncoder,
40 )
41
42 from vllm.config import VllmConfig
43 from vllm.config.multimodal import MultiModalDummyOptions
44 from vllm.inputs import ModalityData, MultiModalDataDict
45 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
46 from vllm.multimodal.inputs import (
47 MultiModalFieldConfig,
48 NestedTensors,
49 )
50 from vllm.multimodal.parse import (
51 AudioItem,
52 AudioProcessorItems,
53 DictEmbeddingItems,
54 ModalityDataItems,
55 MultiModalDataItems,
56 )
57 from vllm.multimodal.processing import (
58 PromptReplacement,
59 PromptUpdate,
60 PromptUpdateDetails,
61 cached_encode,
62 )
63 from vllm.utils.tensor_schema import TensorSchema, TensorShape
64
65 from .minicpmv import (
66 _MAX_FRAMES_PER_VIDEO,
67 MiniCPMV2_6,
68 MiniCPMV4_5,
69 MiniCPMVBaseModel,
70 MiniCPMVDummyInputsBuilder,
71 MiniCPMVImageEmbeddingInputs,
72 MiniCPMVImagePixelInputs,
73 MiniCPMVMultiModalDataParser,
74 MiniCPMVMultiModalInputs,
75 MiniCPMVMultiModalProcessor,
76 MiniCPMVProcessingInfo,
77 _minicpmv_field_config,
78 )
79 from .utils import (
80 AutoWeightsLoader,
81 WeightsMapper,
82 cast_overflow_tensors,
83 maybe_prefix,
84 )
85
86 CPU_DEVICE = torch.device("cpu")
87
88 if TYPE_CHECKING:
89 from vllm.transformers_utils.processors.minicpmo import MiniCPMOProcessor
90
91 if os.getenv("USE_FLAGOS") == "1":
92 import flag_gems
93
94 FLAG_GEMS_CONFIG = [
95 "sort",
96 "sort_stable",
97 "layer_norm",
98 "clamp_",
99 "cos",
100 "embedding",
101 "exp",
102 "exponential_",
103 "full",
104 "gather",
105 "gelu",
106 "index",
107 "le",
108 "lt",
109 "lt_scalar",
110 "masked_fill_",
111 "max",
112 "ones",
113 "pow_scalar",
114 "prod_dim",
115 "rand_like",
116 "reciprocal",
117 "repeat",
118 "scatter",
119 "scatter_",
120 "sin",
121 "sub",
122 "true_divide",
123 "true_divide_",
124 "uniform_",
125 "where_scalar_self",
126 "where_self_out",
127 "zeros",
128 "zeros_like",
129 ]
130 flag_gems.only_enable(record=False, include=FLAG_GEMS_CONFIG)
131
132
133 class MiniCPMOAudioFeatureInputs(TensorSchema):
134 """Dimensions:
135 - bns: Batch size * number of audios * number of slices
136 - bn: Batch size * number of audios
137 - c: Number of channels
138 - l: Length
139 - s: Number of slices
140 """
141
142 type: Literal["audio_features"] = "audio_features"
143
144 audio_features: Annotated[
145 torch.Tensor | list[torch.Tensor],
146 TensorShape("bns", "c", "l", dynamic_dims={"l"}),
147 ]
148 """
149 Slice here means chunk. Audio that is too long will be split into slices,
150 which is the same as image. Padding is used therefore `audio_features` is
151 `torch.Tensor`.
152 """
153
154 audio_feature_lens: Annotated[
155 torch.Tensor | list[torch.Tensor],
156 TensorShape("bn", "s"),
157 ]
158 """
159 This should be feature length of each audio slice,
160 which equals to `audio_features.shape[-1]`
161 """
162
163
164 class MiniCPMOAudioEmbeddingInputs(TensorSchema):
165 """Dimensions:
166 - bn: Batch size * number of audios
167 - s: Number of slices
168 - h: Hidden size (must match language model backbone)
169
170 Length of each slice may vary, so pass it as a list.
171 """
172
173 type: Literal["audio_embeds"] = "audio_embeds"
174
175 audio_embeds: Annotated[
176 torch.Tensor | list[torch.Tensor],
177 TensorShape("bn", "s", "h", dynamic_dims={"s"}),
178 ]
179
180
181 MiniCPMOAudioInputs: TypeAlias = (
182 MiniCPMOAudioFeatureInputs | MiniCPMOAudioEmbeddingInputs
183 )
184
185
186 class MiniCPMOMultiModalInputs(MiniCPMVMultiModalInputs, total=False):
187 audios: MiniCPMOAudioInputs | None
188
189
190 def _minicpmo_field_config(hf_inputs: Mapping[str, torch.Tensor]):
191 audio_features = hf_inputs.get("audio_features")
192 audio_feature_lens = hf_inputs.get("audio_feature_lens")
193
194 # For multi-chunk audio (>30s), audio_features has one item per chunk
195 # (total_chunks) while audio_feature_lens has one item per audio (N).
196 # Use flat to group audio_features by audio so both fields
197 # share the same batch size (N).
198 audio_features_cfg = MultiModalFieldConfig.batched("audio")
199
200 if audio_features is not None and audio_feature_lens is not None:
201 num_features = (
202 len(audio_features)
203 if isinstance(audio_features, (list, tuple))
204 else audio_features.shape[0]
205 )
206 num_audios = (
207 len(audio_feature_lens)
208 if isinstance(audio_feature_lens, (list, tuple))
209 else audio_feature_lens.shape[0]
210 )
211
212 if num_features > num_audios:
213 # Compute the number of chunks belonging to each audio
214 chunks_per_audio: list[int] = []
215 for lens in audio_feature_lens:
216 if isinstance(lens, torch.Tensor):
217 chunks_per_audio.append(lens.numel())
218 else:
219 chunks_per_audio.append(1)
220
221 # When audio_feature_lens is padded (e.g. from batched HF
222 # processor output), numel() over-counts. Fall back to
223 # counting non-zero entries so the sizes sum to num_features.
224 if sum(chunks_per_audio) != num_features:
225 chunks_per_audio = []
226 for lens in audio_feature_lens:
227 if isinstance(lens, torch.Tensor):
228 n = int((lens != 0).sum())
229 chunks_per_audio.append(max(n, 1))
230 else:
231 chunks_per_audio.append(1)
232
233 # Use flat (not flat_from_sizes) because audio_features
234 # is list[Tensor] with variable-length chunks (post-unpad).
235 slice_idxs = [0]
236 for n in chunks_per_audio:
237 slice_idxs.append(slice_idxs[-1] + n)
238 audio_features_cfg = MultiModalFieldConfig.flat(
239 "audio",
240 [
241 slice(slice_idxs[i], slice_idxs[i + 1])
242 for i in range(len(chunks_per_audio))
243 ],
244 )
245
246 return dict(
247 **_minicpmv_field_config(hf_inputs),
248 audio_features=audio_features_cfg,
249 audio_feature_lens=MultiModalFieldConfig.batched("audio"),
250 audio_embeds=MultiModalFieldConfig.batched("audio"),
251 )
252
253
254 class MiniCPMOAudioEmbeddingItems(DictEmbeddingItems):
255 def __init__(
256 self,
257 data: Mapping[str, torch.Tensor],
258 fields_factory: Callable[
259 [Mapping[str, torch.Tensor]],
260 Mapping[str, MultiModalFieldConfig],
261 ],
262 ) -> None:
263 super().__init__(
264 data,
265 modality="image",
266 required_fields={"audio_embeds"},
267 fields_factory=fields_factory,
268 )
269
270
271 class MiniCPMOMultiModalDataParser(MiniCPMVMultiModalDataParser):
272 def _parse_audio_data(
273 self,
274 data: dict[str, torch.Tensor] | ModalityData[AudioItem],
275 ) -> ModalityDataItems[Any, Any] | None:
276 if isinstance(data, dict):
277 return MiniCPMOAudioEmbeddingItems(
278 data,
279 fields_factory=_minicpmo_field_config,
280 )
281
282 return super()._parse_audio_data(data)
283
284
285 class MiniCPMOProcessingInfo(MiniCPMVProcessingInfo):
286 audio_pattern = "(<audio>./</audio>)"
287
288 def get_hf_processor(self, **kwargs: object) -> "MiniCPMOProcessor":
289 """Get vendored MiniCPMOProcessor for multimodal (image+audio) inputs.
290
291 Creates a vendored processor that uses the checkpoint-specific HF image
292 processor, feature extractor, and tokenizer; applies the correct audio
293 pooling configuration; and converts numpy arrays in the image processor
294 to lists for serialization compatibility. The returned processor is
295 compatible with Transformers v5.
296 """
297 import numpy as np
298
299 hf_processor = self.ctx.get_hf_processor(**kwargs)
300 image_processor = self._get_checkpoint_image_processor(**kwargs)
301
302 from vllm.transformers_utils.processors.minicpmo import MiniCPMOProcessor
303
304 # Create vendored processor with correct configuration
305 vendored_processor = MiniCPMOProcessor(
306 image_processor=image_processor,
307 feature_extractor=hf_processor.feature_extractor,
308 tokenizer=hf_processor.tokenizer,
309 pool_step=self.get_default_audio_pool_step(),
310 )
311
312 # Convert numpy arrays in image processor to lists for serialization
313 for attr in ("mean", "std"):
314 val = getattr(image_processor, attr, None)
315 if val is not None and isinstance(val, np.ndarray):
316 setattr(image_processor, attr, val.tolist())
317
318 return vendored_processor
319
320 def get_data_parser(self):
321 return MiniCPMOMultiModalDataParser(
322 target_sr=self.get_default_audio_sampling_rate(),
323 expected_hidden_size=self._get_expected_hidden_size(),
324 )
325
326 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
327 return {**super().get_supported_mm_limits(), "audio": None}
328
329 def get_audio_placeholder(
330 self,
331 audio_lens: int,
332 chunk_input: bool = True,
333 chunk_length: int = 1,
334 ) -> str:
335 hf_processor = self.get_hf_processor()
336
337 return hf_processor.get_audio_placeholder(
338 audio_lens,
339 chunk_input=chunk_input,
340 chunk_length=chunk_length,
341 )
342
343 def get_default_audio_pool_step(self) -> int:
344 hf_config = self.get_hf_config()
345 # MiniCPM-o 4.5 uses pool_step=5, older versions use 2
346 return getattr(hf_config, "audio_pool_step", 2)
347
348 def get_default_audio_sampling_rate(self) -> int:
349 return 16000
350
351 def get_chunk_length(self) -> int:
352 return self.get_hf_config().audio_chunk_length
353
354 def get_max_audio_tokens_per_chunk(self) -> int:
355 pool_step = self.get_default_audio_pool_step()
356 fbank_feat_in_chunk = 100
357 cnn_feat_in_chunk = (fbank_feat_in_chunk - 1) // 2 + 1
358 return (cnn_feat_in_chunk - pool_step) // pool_step + 1
359
360 def get_max_audio_chunks_with_most_features(self) -> int:
361 return 30
362
363 def get_max_audio_tokens(self) -> int:
364 num_chunks = self.get_max_audio_chunks_with_most_features()
365 return self.get_max_audio_tokens_per_chunk() * num_chunks
366
367 def get_audio_len_by_num_chunks(self, num_chunks: int) -> int:
368 sampling_rate = self.get_default_audio_sampling_rate()
369 num_tokens_per_chunk = self.get_max_audio_tokens_per_chunk()
370 return int(num_chunks * sampling_rate / num_tokens_per_chunk) + 1
371
372 def get_num_frames_with_most_features(
373 self,
374 seq_len: int,
375 mm_counts: Mapping[str, int],
376 ) -> int:
377 max_images = mm_counts.get("image", 0)
378 max_videos = mm_counts.get("video", 0)
379 max_audios = mm_counts.get("audio", 0)
380
381 max_image_tokens = self.get_max_image_tokens() * max_images
382 max_audio_tokens = self.get_max_audio_tokens() * max_audios
383 max_total_frames = self.get_max_video_frames(
384 seq_len - max_image_tokens - max_audio_tokens
385 )
386 max_frames_per_video = min(
387 max_total_frames // max(max_videos, 1), _MAX_FRAMES_PER_VIDEO
388 )
389
390 return max(max_frames_per_video, 1)
391
392
393 class MiniCPMODummyInputsBuilder(MiniCPMVDummyInputsBuilder[MiniCPMOProcessingInfo]):
394 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
395 num_audios = mm_counts.get("audio", 0)
396
397 audio_prompt_texts = self.info.audio_pattern * num_audios
398
399 return super().get_dummy_text(mm_counts) + audio_prompt_texts
400
401 def get_dummy_mm_data(
402 self,
403 seq_len: int,
404 mm_counts: Mapping[str, int],
405 mm_options: MultiModalDummyOptions,
406 ) -> MultiModalDataDict:
407 audio_len = (
408 self.info.get_max_audio_chunks_with_most_features()
409 * self.info.get_default_audio_sampling_rate()
410 )
411
412 audio_mm_data = {
413 "audio": self._get_dummy_audios(
414 length=audio_len,
415 num_audios=mm_counts.get("audio", 0),
416 overrides=mm_options.get("audio"),
417 )
418 }
419
420 return {
421 **super().get_dummy_mm_data(seq_len, mm_counts, mm_options),
422 **audio_mm_data,
423 }
424
425
426 class MiniCPMOMultiModalProcessor(MiniCPMVMultiModalProcessor[MiniCPMOProcessingInfo]):
427 def get_audio_prompt_texts(
428 self,
429 audio_lens: int,
430 chunk_input: bool = True,
431 chunk_length: int = 1,
432 ) -> str:
433 return self.info.get_audio_placeholder(
434 audio_lens,
435 chunk_input=chunk_input,
436 chunk_length=chunk_length,
437 )
438
439 def process_audios(
440 self,
441 mm_data: Mapping[str, object],
442 mm_kwargs: Mapping[str, object],
443 ) -> Mapping[str, NestedTensors]:
444 if (audios := mm_data.get("audio")) is None:
445 return {}
446
447 mm_items = self.info.parse_mm_data({"audio": audios}, validate=False)
448 parsed_audios = mm_items.get_items(
449 "audio", (MiniCPMOAudioEmbeddingItems, AudioProcessorItems)
450 )
451
452 if isinstance(parsed_audios, MiniCPMOAudioEmbeddingItems):
453 audio_inputs = {}
454 else:
455 audio_inputs = self._call_hf_processor_on_prompts(
456 prompts=[self.info.audio_pattern] * len(parsed_audios),
457 mm_data={"audios": [[audio] for audio in parsed_audios]},
458 mm_kwargs={**mm_kwargs, "chunk_input": True},
459 out_keys={"audio_features", "audio_feature_lens"},
460 )
461
462 # Avoid padding since we need the output for each audio to be
463 # independent of other audios for the cache to work correctly
464 # Flatten audio_feature_lens (list of tensors of any
465 # dimensionality, one per audio, each containing per-chunk
466 # lengths) into a flat list of integer lengths so there is
467 # one length per chunk, matching the first dimension of
468 # audio_features. Using flatten() handles 0-D, 1-D, and
469 # higher-dimensional tensors uniformly.
470 flat_feature_lens: list[int] = []
471 for lens in audio_inputs["audio_feature_lens"]:
472 if isinstance(lens, torch.Tensor):
473 flat_feature_lens.extend(lens.flatten().tolist())
474 else:
475 assert isinstance(lens, int)
476 flat_feature_lens.append(lens)
477 unpadded_audio_features = [
478 feat[:, :length]
479 for feat, length in zip(
480 audio_inputs["audio_features"],
481 flat_feature_lens,
482 )
483 ]
484 audio_inputs["audio_features"] = unpadded_audio_features
485
486 return audio_inputs
487
488 def process_mm_inputs(
489 self,
490 mm_data: Mapping[str, object],
491 mm_kwargs: Mapping[str, object],
492 ) -> Mapping[str, NestedTensors]:
493 return {
494 **super().process_mm_inputs(mm_data, mm_kwargs),
495 **self.process_audios(mm_data, mm_kwargs),
496 }
497
498 def _get_prompt_updates(
499 self,
500 mm_items: MultiModalDataItems,
501 hf_processor_mm_kwargs: Mapping[str, object],
502 out_mm_kwargs: MultiModalKwargsItems,
503 ) -> Sequence[PromptUpdate]:
504 base_updates = super()._get_prompt_updates(
505 mm_items=mm_items,
506 hf_processor_mm_kwargs=hf_processor_mm_kwargs,
507 out_mm_kwargs=out_mm_kwargs,
508 )
509
510 tokenizer = self.info.get_tokenizer()
511 vocab = tokenizer.get_vocab()
512
513 audio_placeholder = cached_encode(
514 tokenizer, self.info.audio_pattern, add_special_tokens=False
515 )
516 unk_token_ids = [vocab["<unk>"]]
517
518 def get_audio_replacement(item_idx: int):
519 audios = mm_items.get_items(
520 "audio", (MiniCPMOAudioEmbeddingItems, AudioProcessorItems)
521 )
522 assert isinstance(
523 audios, (MiniCPMOAudioEmbeddingItems, AudioProcessorItems)
524 )
525
526 if isinstance(audios, MiniCPMOAudioEmbeddingItems):
527 single_audio_embeds = audios.get(item_idx)["audio_embeds"]
528 audio_len = self.info.get_audio_len_by_num_chunks(
529 sum(map(len, single_audio_embeds))
530 )
531 else:
532 audio_len = audios.get_audio_length(item_idx)
533
534 return PromptUpdateDetails.select_token_ids(
535 cached_encode(
536 tokenizer,
537 self.get_audio_prompt_texts(audio_len),
538 add_special_tokens=False,
539 ),
540 unk_token_ids,
541 )
542
543 return [
544 *base_updates,
545 PromptReplacement(
546 modality="audio",
547 target=audio_placeholder,
548 replacement=get_audio_replacement,
549 ),
550 ]
551
552 def _get_mm_fields_config(
553 self,
554 hf_inputs: BatchFeature,
555 hf_processor_mm_kwargs: Mapping[str, object],
556 ) -> Mapping[str, MultiModalFieldConfig]:
557 return _minicpmo_field_config(hf_inputs)
558
559
560 class MultiModalProjector(nn.Module):
561 def __init__(self, in_dim: int, out_dim: int):
562 super().__init__()
563 self.linear1 = nn.Linear(in_features=in_dim, out_features=out_dim, bias=True)
564 self.relu = nn.ReLU()
565 self.linear2 = nn.Linear(in_features=out_dim, out_features=out_dim, bias=True)
566
567 def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
568 hidden_states = self.relu(self.linear1(audio_features))
569 hidden_states = self.linear2(hidden_states)
570 return hidden_states
571
572
573 class MiniCPMWhisperEncoderLayer(nn.Module):
574 def __init__(self, config: WhisperConfig, layer_idx: int):
575 super().__init__()
576 self.embed_dim = config.d_model
577 self.self_attn = WhisperAttention(
578 embed_dim=self.embed_dim,
579 num_heads=config.encoder_attention_heads,
580 dropout=config.attention_dropout,
581 config=config,
582 layer_idx=layer_idx,
583 )
584 self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
585 self.dropout = config.dropout
586 self.activation_fn = ACT2FN[config.activation_function]
587 self.activation_dropout = config.activation_dropout
588 self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
589 self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
590 self.final_layer_norm = nn.LayerNorm(self.embed_dim)
591
592 def forward(
593 self,
594 hidden_states: torch.Tensor,
595 attention_mask: torch.Tensor,
596 ) -> torch.Tensor:
597 residual = hidden_states
598 hidden_states = self.self_attn_layer_norm(hidden_states)
599 hidden_states, _ = self.self_attn(
600 hidden_states=hidden_states,
601 attention_mask=attention_mask,
602 )
603 hidden_states = nn.functional.dropout(
604 hidden_states, p=self.dropout, training=self.training
605 )
606 hidden_states = residual + hidden_states
607
608 residual = hidden_states
609 hidden_states = self.final_layer_norm(hidden_states)
610 hidden_states = self.activation_fn(self.fc1(hidden_states))
611 hidden_states = nn.functional.dropout(
612 hidden_states, p=self.activation_dropout, training=self.training
613 )
614 hidden_states = self.fc2(hidden_states)
615 hidden_states = nn.functional.dropout(
616 hidden_states, p=self.dropout, training=self.training
617 )
618 hidden_states = residual + hidden_states
619
620 if hidden_states.dtype == torch.float16:
621 hidden_states = cast_overflow_tensors(hidden_states)
622
623 outputs = (hidden_states,)
624
625 return outputs
626
627
628 class MiniCPMWhisperEncoder(WhisperEncoder):
629 def __init__(self, config: WhisperConfig):
630 super().__init__(config)
631 self.layers = nn.ModuleList(
632 [
633 MiniCPMWhisperEncoderLayer(config, layer_idx=i)
634 for i in range(config.encoder_layers)
635 ]
636 )
637
638 def forward(
639 self,
640 input_features: torch.Tensor,
641 attention_mask: torch.Tensor | None = None,
642 ) -> BaseModelOutputWithPast:
643 # Ignore copy
644 input_features = input_features.to(
645 dtype=self.conv1.weight.dtype, device=self.conv1.weight.device
646 )
647
648 inputs_embeds = nn.functional.gelu(self.conv1(input_features))
649 inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
650
651 inputs_embeds = inputs_embeds.permute(0, 2, 1)
652
653 embed_pos = self.embed_positions.weight
654
655 embed_pos = embed_pos[: inputs_embeds.shape[1], :]
656
657 hidden_states = inputs_embeds + embed_pos
658 hidden_states = nn.functional.dropout(
659 hidden_states, p=self.dropout, training=self.training
660 )
661
662 encoder_states: tuple[torch.Tensor, ...] = ()
663
664 for idx, encoder_layer in enumerate(self.layers):
665 encoder_states = encoder_states + (hidden_states,)
666 to_drop = False
667 if self.training:
668 dropout_probability = torch.rand([])
669 if dropout_probability < self.layerdrop: # skip the layer
670 to_drop = True
671
672 # Ignore copy
673 if to_drop:
674 layer_outputs = (None, None)
675 else:
676 layer_outputs = encoder_layer(
677 hidden_states,
678 attention_mask,
679 )
680
681 hidden_states = layer_outputs[0]
682
683 hidden_states = self.layer_norm(hidden_states)
684 encoder_states = encoder_states + (hidden_states,)
685
686 return BaseModelOutputWithPast(
687 last_hidden_state=hidden_states,
688 hidden_states=encoder_states,
689 )
690
691
692 if TYPE_CHECKING:
693 _MiniCPMOBaseModelBase = MiniCPMVBaseModel
694 else:
695 _MiniCPMOBaseModelBase = object
696
697
698 class MiniCPMOBaseModel(_MiniCPMOBaseModelBase):
699 """Base mixin class for MiniCPM-O models with audio support."""
700
701 # Unlike the vision-only MiniCPM-V models, audio weights are loaded here.
702 hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"tts": None})
703
704 packed_modules_mapping = {
705 "qkv_proj": [
706 "q_proj",
707 "k_proj",
708 "v_proj",
709 ],
710 "gate_up_proj": [
711 "gate_proj",
712 "up_proj",
713 ],
714 }
715
716 @classmethod
717 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
718 if modality.startswith("image"):
719 return "(<image>./</image>)"
720 if modality.startswith("video"):
721 return "(<video>./</video>)"
722 if modality.startswith("audio"):
723 return "(<audio>./</audio>)"
724
725 raise ValueError("Only image, video or audio modality is supported")
726
727 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
728 super().__init__(vllm_config=vllm_config, prefix=prefix)
729
730 with self._mark_tower_model(vllm_config, "audio"):
731 self.apm = self.init_audio_module(
732 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "apm")
733 )
734
735 def init_audio_module(self, *, vllm_config: VllmConfig, prefix: str = ""):
736 # Do not use parameters temporarily
737 audio_config = self.config.audio_config
738 model = MiniCPMWhisperEncoder(audio_config)
739 audio_output_dim = int(audio_config.encoder_ffn_dim // 4)
740 self.audio_avg_pooler = nn.AvgPool1d(
741 self.config.audio_pool_step, stride=self.config.audio_pool_step
742 )
743 self.audio_projection_layer = MultiModalProjector(
744 in_dim=audio_output_dim, out_dim=self.embed_dim
745 )
746 self.audio_encoder_layer = -1
747 return model
748
749 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
750 loader = AutoWeightsLoader(self)
751 loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
752 self._ensure_resampler_device()
753 return loaded
754
755 def subsequent_chunk_mask(
756 self,
757 size: int,
758 chunk_size: int,
759 num_left_chunks: int = -1,
760 device: torch.device = CPU_DEVICE,
761 num_lookhead: int = 0,
762 ) -> torch.Tensor:
763 ret = torch.zeros(size, size, device=device, dtype=torch.bool)
764 # Vectorized computation of row indices and chunk boundaries
765 row_indices = torch.arange(size, device=device)
766 chunk_indices = row_indices // chunk_size
767 if num_left_chunks < 0:
768 # If num_left_chunks < 0, start is always 0 for all rows
769 start_indices = torch.zeros_like(row_indices)
770 else:
771 # Compute start indices vectorially
772 start_chunk_indices = torch.clamp(chunk_indices - num_left_chunks, min=0)
773 start_indices = start_chunk_indices * chunk_size
774 # Compute ending indices vectorially
775 end_chunk_indices = chunk_indices + 1
776 end_indices = torch.clamp(
777 end_chunk_indices * chunk_size + num_lookhead, max=size
778 )
779 # Create column indices for broadcasting
780 col_indices = torch.arange(size, device=device).unsqueeze(0)
781 start_indices = start_indices.unsqueeze(1)
782 end_indices = end_indices.unsqueeze(1)
783 # Vectorized mask creation
784 ret = (col_indices >= start_indices) & (col_indices < end_indices)
785 return ret
786
787 def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
788 input_lengths_after_cnn = (input_lengths - 1) // 2 + 1
789 input_lengths_after_pooling = (
790 input_lengths_after_cnn - self.config.audio_pool_step
791 ) // self.config.audio_pool_step + 1
792 input_lengths_after_pooling = input_lengths_after_pooling.to(dtype=torch.int32)
793
794 return input_lengths_after_cnn, input_lengths_after_pooling
795
796 def get_audio_hidden_states(
797 self, data: MiniCPMOAudioFeatureInputs
798 ) -> list[torch.Tensor]:
799 chunk_length = self.config.audio_chunk_length
800
801 # (bs, 80, frames) or [], multi audios need filled in advance
802 wavforms_raw = data["audio_features"]
803 if isinstance(wavforms_raw, list):
804 B = len(wavforms_raw)
805 C = wavforms_raw[0].shape[-2]
806 L = max(item.shape[-1] for item in wavforms_raw)
807 device = wavforms_raw[0].device
808 dtype = wavforms_raw[0].dtype
809
810 wavforms = torch.zeros((B, C, L), dtype=dtype, device=device)
811 for i, wavforms_item in enumerate(wavforms_raw):
812 L_item = wavforms_item.shape[-1]
813 wavforms[i, ..., :L_item] = wavforms_item
814 else:
815 wavforms = wavforms_raw
816
817 # list, [[x1, x2], [y1], [z1]]
818 audio_feature_lens_raw = data["audio_feature_lens"]
819 if isinstance(audio_feature_lens_raw, torch.Tensor):
820 audio_feature_lens_raw = audio_feature_lens_raw.unbind(0)
821
822 audio_feature_lens = torch.hstack(audio_feature_lens_raw)
823 batch_size, _, max_mel_seq_len = wavforms.shape
824 max_seq_len = (max_mel_seq_len - 1) // 2 + 1
825
826 # Create a sequence tensor of shape (batch_size, max_seq_len)
827 seq_range = (
828 torch.arange(
829 0,
830 max_seq_len,
831 dtype=audio_feature_lens.dtype,
832 device=audio_feature_lens.device,
833 )
834 .unsqueeze(0)
835 .expand(batch_size, max_seq_len)
836 )
837 lengths_expand = audio_feature_lens.unsqueeze(1).expand(batch_size, max_seq_len)
838 # Create mask
839 padding_mask = seq_range >= lengths_expand # 1 for padded values
840
841 audio_attention_mask_ = padding_mask.view(batch_size, 1, 1, max_seq_len).expand(
842 batch_size, 1, max_seq_len, max_seq_len
843 )
844 audio_attention_mask = audio_attention_mask_.to(
845 dtype=self.apm.conv1.weight.dtype, device=self.apm.conv1.weight.device
846 )
847
848 if chunk_length > 0:
849 chunk_num_frame = int(chunk_length * 50)
850 chunk_mask = self.subsequent_chunk_mask(
851 size=max_seq_len,
852 chunk_size=chunk_num_frame,
853 num_left_chunks=-1,
854 device=audio_attention_mask_.device,
855 )
856 audio_attention_mask_ = torch.logical_or(
857 audio_attention_mask_, torch.logical_not(chunk_mask)
858 )
859
860 audio_attention_mask[audio_attention_mask_] = float("-inf")
861 audio_states = self.apm(
862 wavforms, attention_mask=audio_attention_mask
863 ).hidden_states[self.audio_encoder_layer]
864 audio_embeds = self.audio_projection_layer(audio_states)
865
866 audio_embeds = audio_embeds.transpose(1, 2)
867 audio_embeds = self.audio_avg_pooler(audio_embeds)
868 audio_embeds = audio_embeds.transpose(1, 2)
869
870 _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(
871 audio_feature_lens
872 )
873
874 num_audio_tokens = feature_lens_after_pooling
875
876 final_audio_embeds = list[torch.Tensor]()
877 idx = 0
878 for i in range(len(audio_feature_lens_raw)):
879 target_audio_embeds_lst = list[torch.Tensor]()
880 for _ in range(len(audio_feature_lens_raw[i])):
881 target_audio_embeds_lst.append(
882 audio_embeds[idx, : num_audio_tokens[idx], :]
883 )
884 idx += 1
885
886 final_audio_embeds.append(torch.cat(target_audio_embeds_lst))
887
888 return final_audio_embeds
889
890 def _parse_and_validate_audio_input(
891 self, **kwargs: object
892 ) -> MiniCPMOAudioInputs | None:
893 audio_features = kwargs.pop("audio_features", None)
894 audio_embeds = kwargs.pop("audio_embeds", None)
895
896 if audio_features is None and audio_embeds is None:
897 return None
898
899 if audio_embeds is not None:
900 return MiniCPMOAudioEmbeddingInputs(
901 type="audio_embeds",
902 audio_embeds=audio_embeds,
903 )
904
905 audio_feature_lens = kwargs.pop("audio_feature_lens")
906
907 return MiniCPMOAudioFeatureInputs(
908 type="audio_features",
909 audio_features=audio_features,
910 audio_feature_lens=audio_feature_lens,
911 )
912
913 def _parse_and_validate_multimodal_inputs(
914 self, **kwargs: object
915 ) -> MiniCPMOMultiModalInputs:
916 modalities = MiniCPMOMultiModalInputs(
917 **super()._parse_and_validate_multimodal_inputs(**kwargs)
918 )
919
920 # Preserve the order of modalities if there are multiple of them
921 # from the order of kwargs.
922 for input_key in kwargs:
923 if (
924 input_key in ("audio_features", "audio_embeds")
925 and "audios" not in modalities
926 ):
927 modalities["audios"] = self._parse_and_validate_audio_input(**kwargs)
928
929 return modalities
930
931 def _process_audio_input(
932 self,
933 audio_input: MiniCPMOAudioInputs,
934 ) -> torch.Tensor | list[torch.Tensor]:
935 if audio_input["type"] == "audio_embeds":
936 return audio_input["audio_embeds"]
937
938 assert isinstance(audio_input, MiniCPMOAudioFeatureInputs)
939 return self.get_audio_hidden_states(audio_input)
940
941 def _process_multimodal_inputs(self, modalities: Mapping[str, object]):
942 base_modalities: MiniCPMVMultiModalInputs = {}
943 image_input = modalities.get("images")
944 video_input = modalities.get("videos")
945 assert image_input is None or isinstance(
946 image_input, (MiniCPMVImagePixelInputs, MiniCPMVImageEmbeddingInputs)
947 )
948 assert video_input is None or isinstance(
949 video_input, (MiniCPMVImagePixelInputs, MiniCPMVImageEmbeddingInputs)
950 )
951 if "images" in modalities:
952 base_modalities["images"] = image_input
953 if "videos" in modalities:
954 base_modalities["videos"] = video_input
955 multimodal_embeddings = super()._process_multimodal_inputs(base_modalities)
956
957 audio_input = modalities.get("audios")
958 if audio_input is not None:
959 assert isinstance(
960 audio_input,
961 (MiniCPMOAudioFeatureInputs, MiniCPMOAudioEmbeddingInputs),
962 )
963 audio_embeddings = self._process_audio_input(audio_input)
964 multimodal_embeddings += tuple(audio_embeddings)
965
966 return multimodal_embeddings
967
968
969 class MiniCPMO2_6(MiniCPMOBaseModel, MiniCPMV2_6):
970 """MiniCPM-O 2.6 model with Qwen2 backbone."""
971
972 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
973 super().__init__(vllm_config=vllm_config, prefix=prefix)
974
975 with self._mark_tower_model(vllm_config, "audio"):
976 self.apm = self.init_audio_module(
977 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "apm")
978 )
979
980
981 class MiniCPMO4_5(MiniCPMOBaseModel, MiniCPMV4_5):
982 """MiniCPM-O 4.5 model with Qwen3 backbone."""
983
984 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
985 super().__init__(vllm_config=vllm_config, prefix=prefix)
986
987 with self._mark_tower_model(vllm_config, "audio"):
988 self.apm = self.init_audio_module(
989 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "apm")
990 )
991
992
993 _MINICPMO_SUPPORT_VERSION = {
994 (2, 6): MiniCPMO2_6,
995 (4, 5): MiniCPMO4_5,
996 }
997
998
999 @MULTIMODAL_REGISTRY.register_processor(
1000 MiniCPMOMultiModalProcessor,
1001 info=MiniCPMOProcessingInfo,
1002 dummy_inputs=MiniCPMODummyInputsBuilder,
1003 )
1004 class MiniCPMO(MiniCPMOBaseModel, MiniCPMV2_6):
1005 """MiniCPM-O model with audio support.
1006 Different versions use different LLM backbones:
1007 - Version 2.6: Uses Qwen2
1008 - Version 4.5: Uses Qwen3
1009 """
1010
1011 def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""):
1012 config = vllm_config.model_config.hf_config
1013
1014 # Determine version from config
1015 if hasattr(config, "version"):
1016 try:
1017 version_str = str(config.version)
1018 version_parts = version_str.split(".")
1019 version_values = tuple(int(x) for x in version_parts[:2])
1020 except (ValueError, TypeError) as e:
1021 raise ValueError(
1022 f"Invalid model version format in config: {config.version}. "
1023 "Expected a dot-separated version string like '4.5'."
1024 ) from e
1025 else:
1026 # Default to 2.6 for backward compatibility
1027 version_values = (2, 6)
1028
1029 # Dispatch class based on version
1030 if len(version_values) == 2:
1031 version_key = (version_values[0], version_values[1])
1032 instance_cls = _MINICPMO_SUPPORT_VERSION.get(version_key)
1033 else:
1034 instance_cls = None
1035 if instance_cls is None:
1036 supported_versions = ", ".join(
1037 [f"{v[0]}.{v[1]}" for v in sorted(_MINICPMO_SUPPORT_VERSION.keys())]
1038 )
1039 raise ValueError(
1040 f"Currently, MiniCPMO only supports versions "
1041 f"{supported_versions}. Got version: {version_values}"
1042 )
1043
1044 return instance_cls(vllm_config=vllm_config, prefix=prefix)
1045
1046 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1047 # This __init__ won't be called due to __new__ returning a different class
1048 pass