openbmb/MiniCPM-o-4_5
1.5k710k
1#!/usr/bin/env python2# -*- coding: utf-8 -*-3# Copyright 2026 The OpenBMB Team. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16 17import copy18import math19import re20from typing import Any21from typing import Dict22from typing import List23from typing import Optional24from typing import Tuple25from typing import Union26 27import numpy as np28import torch29from PIL import Image30from transformers import AutoImageProcessor31from transformers.audio_utils import spectrogram32from transformers.audio_utils import window_function33from transformers.image_processing_utils import BaseImageProcessor34from transformers.image_processing_utils import BatchFeature35from transformers.image_transforms import to_channel_dimension_format36from transformers.image_utils import ChannelDimension37from transformers.image_utils import ImageInput38from transformers.image_utils import infer_channel_dimension_format39from transformers.image_utils import is_torch_tensor40from transformers.image_utils import to_numpy_array41from transformers.image_utils import valid_images42from transformers.models.whisper.feature_extraction_whisper import WhisperFeatureExtractor43from transformers.processing_utils import ProcessorMixin44from transformers.tokenization_utils_base import PreTokenizedInput45from transformers.tokenization_utils_base import TextInput46from transformers.utils import is_torch_device47from transformers.utils import is_torch_dtype48from transformers.utils import requires_backends49from transformers.utils import TensorType50 51 52def recursive_converter(converter, value):53 if isinstance(value, list):54 new_value = []55 for v in value:56 new_value += [recursive_converter(converter, v)]57 return new_value58 else:59 return converter(value)60 61 62class MiniCPMOBatchFeature(BatchFeature):63 """Extend from BatchFeature for supporting various image size"""64 65 def __init__(self, data: Optional[Dict[str, Any]] = None, tensor_type: Union[None, str, TensorType] = None):66 super().__init__(data)67 self.convert_to_tensors(tensor_type=tensor_type)68 69 def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):70 if tensor_type is None:71 return self72 73 is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type)74 75 def converter(value):76 try:77 if not is_tensor(value):78 tensor = as_tensor(value)79 return tensor80 except: # noqa E72281 if key == "overflowing_values":82 raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")83 raise ValueError(84 "Unable to create tensor, you should probably activate padding "85 "with 'padding=True' to have batched tensors with the same length."86 )87 88 for key, value in self.items():89 self[key] = recursive_converter(converter, value)90 return self91 92 def to(self, *args, **kwargs) -> "MiniCPMOBatchFeature":93 requires_backends(self, ["torch"])94 import torch95 96 def cast_tensor(v):97 if not torch.is_tensor(v):98 return v99 100 if torch.is_floating_point(v):101 return v.to(*args, **kwargs)102 elif device is not None:103 return v.to(device=device)104 else:105 return v106 107 new_data = {}108 device = kwargs.get("device")109 if device is None and len(args) > 0:110 arg = args[0]111 if is_torch_dtype(arg):112 pass113 elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):114 device = arg115 else:116 raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")117 118 # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`119 for k, v in self.items():120 new_data[k] = recursive_converter(cast_tensor, v)121 self.data = new_data122 return self123 124 125class MiniCPMVImageProcessor(BaseImageProcessor):126 model_input_names = ["pixel_values"]127 128 def __init__(self, max_slice_nums=9, scale_resolution=448, patch_size=14, **kwargs):129 super().__init__(**kwargs)130 self.max_slice_nums = max_slice_nums131 self.scale_resolution = scale_resolution132 self.patch_size = patch_size133 self.use_image_id = kwargs.pop("use_image_id", False)134 self.image_feature_size = kwargs.pop("image_feature_size", 64)135 self.im_start_token = kwargs.pop("im_start", "<image>")136 self.im_end_token = kwargs.pop("im_end", "</image>")137 self.slice_start_token = kwargs.pop("slice_start", "<slice>")138 self.slice_end_token = kwargs.pop("slice_end", "</slice>")139 self.unk_token = kwargs.pop("unk", "<unk>")140 self.im_id_start = kwargs.pop("im_id_start", "<image_id>")141 self.im_id_end = kwargs.pop("im_id_end", "</image_id>")142 self.slice_mode = kwargs.pop("slice_mode", True)143 144 self.mean = np.array(kwargs.pop("norm_mean", [0.5, 0.5, 0.5]))145 self.std = np.array(kwargs.pop("norm_std", [0.5, 0.5, 0.5]))146 self.version = kwargs.pop("version", 2.0)147 148 @staticmethod149 def ensure_divide(length, patch_size):150 return max(round(length / patch_size) * patch_size, patch_size)151 152 def find_best_resize(self, original_size, scale_resolution, patch_size, allow_upscale=False):153 width, height = original_size154 if (width * height > scale_resolution * scale_resolution) or allow_upscale:155 r = width / height156 height = int(scale_resolution / math.sqrt(r))157 width = int(height * r)158 best_width = self.ensure_divide(width, patch_size)159 best_height = self.ensure_divide(height, patch_size)160 return best_width, best_height161 162 def get_refine_size(self, original_size, grid, scale_resolution, patch_size, allow_upscale=False):163 width, height = original_size164 grid_x, grid_y = grid165 166 refine_width = self.ensure_divide(width, grid_x)167 refine_height = self.ensure_divide(height, grid_y)168 169 grid_width = refine_width / grid_x170 grid_height = refine_height / grid_y171 172 best_grid_size = self.find_best_resize(173 (grid_width, grid_height), scale_resolution, patch_size, allow_upscale=allow_upscale174 )175 refine_size = (best_grid_size[0] * grid_x, best_grid_size[1] * grid_y)176 return refine_size177 178 @staticmethod179 def split_to_patches(image, grid):180 patches = []181 width, height = image.size182 grid_x = int(width / grid[0])183 grid_y = int(height / grid[1])184 for i in range(0, height, grid_y):185 images = []186 for j in range(0, width, grid_x):187 box = (j, i, j + grid_x, i + grid_y)188 patch = image.crop(box)189 images.append(patch)190 patches.append(images)191 return patches192 193 def slice_image(self, image, max_slice_nums=9, scale_resolution=448, patch_size=14, never_split=False):194 original_size = image.size195 source_image = None196 best_grid = self.get_sliced_grid(original_size, max_slice_nums, never_split)197 patches = []198 199 if best_grid is None:200 # dont need to slice, upsample201 best_size = self.find_best_resize(original_size, scale_resolution, patch_size, allow_upscale=True)202 source_image = image.resize(best_size, resample=Image.Resampling.BICUBIC)203 else:204 # source image, down-sampling and ensure divided by patch_size205 best_resize = self.find_best_resize(original_size, scale_resolution, patch_size)206 source_image = image.copy().resize(best_resize, resample=Image.Resampling.BICUBIC)207 refine_size = self.get_refine_size(208 original_size, best_grid, scale_resolution, patch_size, allow_upscale=True209 )210 refine_image = image.resize(refine_size, resample=Image.Resampling.BICUBIC)211 patches = self.split_to_patches(refine_image, best_grid)212 213 return source_image, patches, best_grid214 215 def get_grid_placeholder(self, grid):216 if grid is None:217 return ""218 slice_image_placeholder = (219 self.slice_start_token + self.unk_token * self.image_feature_size + self.slice_end_token220 )221 222 cols = grid[0]223 rows = grid[1]224 slices = []225 for i in range(rows):226 lines = []227 for j in range(cols):228 lines.append(slice_image_placeholder)229 slices.append("".join(lines))230 231 slice_placeholder = "\n".join(slices)232 return slice_placeholder233 234 def get_image_id_placeholder(self, idx=0):235 return f"{self.im_id_start}{idx}{self.im_id_end}"236 237 def get_sliced_images(self, image, max_slice_nums=None):238 slice_images = []239 240 if not self.slice_mode:241 return [image]242 243 max_slice_nums = self.max_slice_nums if max_slice_nums is None else int(max_slice_nums)244 assert max_slice_nums > 0245 source_image, patches, sliced_grid = self.slice_image(246 image, max_slice_nums, self.scale_resolution, self.patch_size # default: 9 # default: 448 # default: 14247 )248 249 slice_images.append(source_image)250 if len(patches) > 0:251 for i in range(len(patches)):252 for j in range(len(patches[0])):253 slice_images.append(patches[i][j])254 return slice_images255 256 def get_sliced_grid(self, image_size, max_slice_nums, nerver_split=False):257 original_width, original_height = image_size258 log_ratio = math.log(original_width / original_height)259 ratio = original_width * original_height / (self.scale_resolution * self.scale_resolution)260 multiple = min(math.ceil(ratio), max_slice_nums)261 if multiple <= 1 or nerver_split:262 return None263 candidate_split_grids_nums = []264 for i in [multiple - 1, multiple, multiple + 1]:265 if i == 1 or i > max_slice_nums:266 continue267 candidate_split_grids_nums.append(i)268 269 candidate_grids = []270 for split_grids_nums in candidate_split_grids_nums:271 m = 1272 while m <= split_grids_nums:273 if split_grids_nums % m == 0:274 candidate_grids.append([m, split_grids_nums // m])275 m += 1276 277 best_grid = [1, 1]278 min_error = float("inf")279 for grid in candidate_grids:280 error = abs(log_ratio - math.log(grid[0] / grid[1]))281 if error < min_error:282 best_grid = grid283 min_error = error284 285 return best_grid286 287 def get_slice_image_placeholder(self, image_size, image_idx=0, max_slice_nums=None, use_image_id=None):288 max_slice_nums = self.max_slice_nums if max_slice_nums is None else int(max_slice_nums)289 assert max_slice_nums > 0290 grid = self.get_sliced_grid(image_size=image_size, max_slice_nums=max_slice_nums)291 292 image_placeholder = self.im_start_token + self.unk_token * self.image_feature_size + self.im_end_token293 use_image_id = self.use_image_id if use_image_id is None else bool(use_image_id)294 if use_image_id:295 final_placeholder = self.get_image_id_placeholder(image_idx) + image_placeholder296 else:297 final_placeholder = image_placeholder298 299 if self.slice_mode:300 final_placeholder = final_placeholder + self.get_grid_placeholder(grid=grid)301 return final_placeholder302 303 @staticmethod304 def to_pil_image(image, rescale=None) -> Image.Image:305 """Converts `image` to a PIL Image. Optionally rescales it and puts the channel dimension back306 as the last axis if needed.307 308 Args:309 image (`Image.Image` or `numpy.ndarray` or `torch.Tensor`):310 The image to convert to the PIL Image format.311 rescale (`bool`, *optional*):312 whether to apply the scaling factor (to make pixel values integers between 0 and 255). Will313 default to `True` if the image type is a floating type, `False` otherwise.314 """315 if isinstance(image, Image.Image):316 return image317 if is_torch_tensor(image):318 image = image.numpy()319 320 if isinstance(image, np.ndarray):321 if rescale is None:322 # rescale default to the array being of floating type.323 rescale = isinstance(image.flat[0], np.floating)324 # If the channel as been moved to first dim, we put it back at the end.325 if image.ndim == 3 and image.shape[0] in [1, 3]:326 image = image.transpose(1, 2, 0)327 if rescale:328 image = image * 255329 image = image.astype(np.uint8)330 return Image.fromarray(image)331 return image332 333 def reshape_by_patch(self, image):334 image = torch.from_numpy(image)335 patch_size = self.patch_size336 patches = torch.nn.functional.unfold(image, (patch_size, patch_size), stride=(patch_size, patch_size))337 338 patches = patches.reshape(image.size(0), patch_size, patch_size, -1)339 patches = patches.permute(0, 1, 3, 2).reshape(image.size(0), patch_size, -1)340 return patches.numpy()341 342 def preprocess(343 self,344 images: Union[Image.Image, List[Image.Image], List[List[Image.Image]]],345 do_pad: Optional[bool] = True,346 max_slice_nums: int = None,347 return_tensors: Optional[Union[str, TensorType]] = None,348 **kwargs,349 ) -> MiniCPMOBatchFeature:350 if isinstance(images, Image.Image):351 images_list = [[images]]352 elif isinstance(images[0], Image.Image):353 images_list = [images]354 else:355 images_list = images356 357 new_images_list = []358 image_sizes_list = []359 tgt_sizes_list = []360 361 for _images in images_list:362 if _images is None or len(_images) == 0:363 new_images_list.append([])364 image_sizes_list.append([])365 tgt_sizes_list.append([])366 continue367 if not valid_images(_images):368 raise ValueError(369 "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "370 "torch.Tensor, tf.Tensor or jax.ndarray."371 )372 373 _images = [self.to_pil_image(image).convert("RGB") for image in _images]374 input_data_format = infer_channel_dimension_format(np.array(_images[0]))375 376 new_images = []377 image_sizes = [image.size for image in _images]378 tgt_sizes = []379 for image in _images:380 image_patches = self.get_sliced_images(image, max_slice_nums)381 image_patches = [to_numpy_array(image).astype(np.float32) / 255 for image in image_patches]382 image_patches = [383 self.normalize(image=image, mean=self.mean, std=self.std, input_data_format=input_data_format)384 for image in image_patches385 ]386 image_patches = [387 to_channel_dimension_format(image, ChannelDimension.FIRST, input_channel_dim=input_data_format)388 for image in image_patches389 ]390 for slice_image in image_patches:391 new_images.append(self.reshape_by_patch(slice_image))392 tgt_sizes.append(393 np.array((slice_image.shape[1] // self.patch_size, slice_image.shape[2] // self.patch_size))394 )395 396 if tgt_sizes:397 tgt_sizes = np.vstack(tgt_sizes)398 399 new_images_list.append(new_images)400 image_sizes_list.append(image_sizes)401 tgt_sizes_list.append(tgt_sizes)402 return MiniCPMOBatchFeature(403 data={"pixel_values": new_images_list, "image_sizes": image_sizes_list, "tgt_sizes": tgt_sizes_list},404 tensor_type=return_tensors,405 )406 407 408AutoImageProcessor.register("MiniCPMVImageProcessor", MiniCPMVImageProcessor)409 410 411def chunk_audio(audio: np.ndarray, max_duration_seconds: int = 30, sample_rate: int = 16000) -> List[np.ndarray]:412 """split long audio into chunks413 414 Args:415 audio:416 max_duration_seconds:417 sample_rate:418 419 Returns:420 chunks421 """422 max_len = int(max_duration_seconds * sample_rate)423 424 if len(audio) <= max_len:425 return [audio]426 427 chunks = []428 for i in range(0, len(audio), max_len):429 chunk = audio[i : i + max_len]430 chunks.append(chunk)431 432 return chunks433 434 435def process_audio_batch(436 audios: Union[np.ndarray, List[np.ndarray], List[List[np.ndarray]]],437 feature_extractor,438 sampling_rate: int = 16000,439 max_duration_seconds: int = 30,440 return_attention_mask: bool = True,441) -> Tuple[torch.Tensor, List[torch.Tensor]]:442 """extract audio mel features443 444 Args:445 audios:446 feature_extractor: WhisperFeatureExtractor447 sampling_rate:448 max_duration_seconds:449 return_attention_mask:450 451 Returns:452 (audio_features, audio_feature_lens)453 audio_features: [batch_size, n_mels, max_frames]454 audio_feature_lens:455 """456 if isinstance(audios, np.ndarray):457 audios_list = [[audios]]458 elif len(audios) > 0 and isinstance(audios[0], np.ndarray):459 audios_list = [audios]460 else:461 audios_list = audios462 463 audio_features_all = []464 audio_feature_lens_list = []465 466 for batch_audios in audios_list:467 batch_lens = []468 469 for audio in batch_audios:470 chunks = chunk_audio(audio, max_duration_seconds, sampling_rate)471 472 for chunk in chunks:473 audio_input = feature_extractor(474 chunk,475 sampling_rate=sampling_rate,476 return_tensors="pt",477 padding="max_length",478 return_attention_mask=return_attention_mask,479 )480 481 audio_feature = audio_input["input_features"] # [1, 80, frames]482 483 if return_attention_mask:484 actual_len = audio_input["attention_mask"].sum(dim=1) # Tensor([frames])485 audio_feature = audio_feature[:, :, : actual_len[0]]486 batch_lens.append(actual_len[0])487 else:488 batch_lens.append(torch.tensor(audio_feature.shape[2]))489 490 audio_features_all.append(audio_feature.squeeze(0)) # [80, frames]491 492 if len(batch_lens) > 0:493 audio_feature_lens_list.append(torch.hstack(batch_lens))494 else:495 audio_feature_lens_list.append(torch.tensor([]))496 497 # pad to same length498 if audio_features_all:499 audio_features = torch.nn.utils.rnn.pad_sequence(500 [feat.transpose(0, 1) for feat in audio_features_all], batch_first=True, padding_value=0.0501 ).transpose(502 1, 2503 ) # [batch, 80, max_frames]504 else:505 audio_features = torch.tensor([])506 507 return audio_features, audio_feature_lens_list508 509 510def regroup_audio_features(511 audio_features: torch.Tensor, audio_feature_lens: List[torch.Tensor], regroup_seconds: int, fps: int = 100512) -> Tuple[torch.Tensor, List[torch.Tensor]]:513 """regroup audio features to fixed duration514 515 Args:516 audio_features: [batch, n_mels, frames]517 audio_feature_lens: each batch's actual length518 regroup_seconds: regroup duration (seconds)519 fps: frames per second520 521 Returns:522 (regrouped_features, regrouped_lens)523 """524 # flatten to continuous frames sequence525 all_lens = []526 for lens in audio_feature_lens:527 if isinstance(lens, torch.Tensor):528 all_lens.extend(lens.tolist())529 elif isinstance(lens, list):530 all_lens.extend([int(x) for x in lens])531 532 if len(all_lens) == 0:533 return torch.tensor([]), []534 535 # concatenate all valid features536 flat_slices = [audio_features[i, :, :L] for i, L in enumerate(all_lens)] # [n_mels, L]537 538 if len(flat_slices) == 1:539 full_feat = flat_slices[0]540 else:541 full_feat = torch.cat(flat_slices, dim=1) # [n_mels, total_frames]542 543 # split to fixed frames544 frames_per_seg = int(regroup_seconds * fps)545 segments = []546 547 for start in range(0, full_feat.size(1), frames_per_seg):548 seg = full_feat[:, start : start + frames_per_seg]549 if seg.size(1) > 0:550 segments.append(seg)551 552 if len(segments) == 0:553 return torch.tensor([]), []554 555 # pad and convert to batch556 seg_lens = [s.size(1) for s in segments]557 segs_transposed = [s.transpose(0, 1) for s in segments]558 559 padded = torch.nn.utils.rnn.pad_sequence(segs_transposed, batch_first=True, padding_value=0.0) # [N, max_T, n_mels]560 561 padded = padded.transpose(1, 2) # [N, n_mels, max_T]562 lens_tensor = torch.tensor(seg_lens, dtype=torch.int32, device=padded.device)563 564 return padded, [lens_tensor]565 566 567class MiniCPMAAudioProcessor(WhisperFeatureExtractor):568 """569 On top of WhisperFeatureExtractor:570 - support dynamic_log_norm (original max-8dB, adjustable dynamic_range_db)571 - or fixed log_floor_db (e.g. -10dB)572 - this is because we need to do streaming scheme, in which we can't do dynamic setting573 - this can be modified in the middle, through set_dynamic_log_norm574 Two paths (torch / numpy) keep consistent clipping and scaling order:575 log10 -> (dynamic/fixed lower limit clipping) -> (+4)/4576 """577 578 def __init__(579 self,580 *args,581 dynamic_log_norm: bool = True,582 dynamic_range_db: float = 8.0,583 log_floor_db: float = -10.0,584 **kwargs,585 ):586 super().__init__(*args, **kwargs)587 self.dynamic_log_norm = bool(dynamic_log_norm)588 self.dynamic_range_db = float(dynamic_range_db)589 self.log_floor_db = float(log_floor_db)590 591 def set_spac_log_norm(592 self,593 dynamic_range_db: Optional[float] = None,594 log_floor_db: Optional[float] = None,595 *,596 inplace: bool = True,597 ) -> "MiniCPMAAudioProcessor":598 """Hot update dynamic/fixed lower limit strategy.599 600 Args:601 enabled: True=use dynamic threshold (max - dynamic_range_db), False=use fixed lower limit log_floor_db.602 None means keep unchanged.603 dynamic_range_db: dynamic range (dB), only effective when enabled=True. None means keep unchanged.604 log_floor_db: fixed log floor (dB, usually <= 0), only effective when enabled=False. None means keep unchanged.605 inplace: True directly modify current instance; False return a shallow copy and modify on it.606 607 Returns:608 self or new instance (when inplace=False).609 """610 611 target = self if inplace else copy.copy(self)612 613 if dynamic_range_db is not None:614 val = float(dynamic_range_db)615 if val < 0:616 raise ValueError("dynamic_range_db must be >= 0.")617 target.dynamic_log_norm = True # explicitly set the value to dynamic mode618 target.dynamic_range_db = val619 620 if log_floor_db is not None:621 val = float(log_floor_db)622 # usually log10(mel) maximum is not more than ~0dB, floor should be <= 0; here do loose validation623 if val > 0:624 raise ValueError("log_floor_db should be <= 0 (log10 scale).")625 target.dynamic_log_norm = False # explicitly set the value to fixed lower limit mode626 target.log_floor_db = val627 628 return target629 630 def _np_extract_fbank_features(self, waveform_batch: np.ndarray, device: str) -> np.ndarray:631 """NumPy version consistent with upstream, but replace max-8dB with configurable dynamic/fixed lower limit clipping."""632 if device != "cpu":633 raise ValueError(634 f"Got device `{device}` for feature extraction, but feature extraction on CUDA accelerator "635 "devices requires torch. Set device='cpu' or install torch."636 )637 638 log_spec_batch: List[np.ndarray] = []639 for waveform in waveform_batch:640 # generate log10 Mel641 log_spec = spectrogram(642 waveform,643 window_function(self.n_fft, "hann"),644 frame_length=self.n_fft,645 hop_length=self.hop_length,646 power=2.0,647 dither=self.dither,648 mel_filters=self.mel_filters,649 log_mel="log10",650 )651 # consistent with upstream: remove the last frame652 log_spec = log_spec[:, :-1]653 654 # dynamic/fixed clipping655 if self.dynamic_log_norm:656 threshold = log_spec.max() - self.dynamic_range_db657 log_spec = np.maximum(log_spec, threshold)658 else:659 log_spec = np.maximum(log_spec, self.log_floor_db)660 661 # consistent with Whisper linear scaling662 log_spec = (log_spec + 4.0) / 4.0663 664 log_spec_batch.append(log_spec)665 666 return np.array(log_spec_batch)667 668 def _torch_extract_fbank_features(self, waveform: np.ndarray, device: str = "cpu") -> np.ndarray:669 if torch is None:670 raise RuntimeError("PyTorch is not installed, cannot compute STFT on GPU.")671 672 waveform = torch.from_numpy(waveform).to(device, torch.float32)673 window = torch.hann_window(self.n_fft, device=device)674 675 if self.dither != 0.0:676 waveform = waveform + self.dither * torch.randn_like(waveform)677 678 stft = torch.stft(waveform, n_fft=self.n_fft, hop_length=self.hop_length, window=window, return_complex=True)679 magnitudes = stft[..., :-1].abs() ** 2680 681 mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32) # [n_mels, 1+n_fft//2]682 mel_spec = mel_filters.T @ magnitudes # [..., n_mels, T]683 684 log_spec = torch.clamp(mel_spec, min=1e-10).log10() # <= 0685 686 if self.dynamic_log_norm:687 if waveform.dim() == 2:688 max_val_t = log_spec.max(dim=2, keepdim=True)[0] # over T689 max_val_bt = max_val_t.max(dim=1, keepdim=True)[0] # over mel690 threshold = max_val_bt - self.dynamic_range_db691 log_spec = torch.maximum(log_spec, threshold)692 else:693 threshold = log_spec.max() - self.dynamic_range_db694 log_spec = torch.maximum(log_spec, threshold)695 else:696 floor_tensor = torch.tensor(self.log_floor_db, dtype=log_spec.dtype, device=log_spec.device)697 log_spec = torch.maximum(log_spec, floor_tensor)698 699 log_spec = (log_spec + 4.0) / 4.0700 701 if device != "cpu":702 log_spec = log_spec.detach().cpu()703 return log_spec.numpy()704 705 def process(self, *args, **kwargs):706 """Alias of __call__ for convenience."""707 return self.__call__(*args, **kwargs)708 709 710class StreamingMelProcessorExact:711 """Strictly offline equivalent streaming Mel processor.712 713 - accumulate all historical audio into buffer; use the same feature_extractor to calculate the entire mel after each addition.714 - only output "stable" frames: the frame center does not depend on future (right) context, i.e. center + n_fft//2 <= current buffer length.715 - output the last batch of frames at the end (flush), ensuring complete consistency with offline full-calculation.716 717 Cost: Each call performs feature extraction on the accumulated buffer (can be optimized to incremental if needed).718 """719 720 def __init__(721 self,722 feature_extractor: MiniCPMAAudioProcessor,723 chunk_ms: int = 100,724 first_chunk_ms: Optional[int] = None,725 sample_rate: int = 16000,726 n_fft: int = 400,727 hop_length: int = 160,728 n_mels: int = 80,729 cnn_redundancy_ms: int = 10, # (given in ms, usually 10ms=1 frame)730 # sliding window parameters731 enable_sliding_window: bool = False, # whether to enable sliding window732 slide_trigger_seconds: float = 30.0, # trigger threshold for sliding window in seconds733 slide_stride_seconds: float = 10.0, # stride for sliding window in seconds734 ):735 self.feature_extractor = feature_extractor736 self.chunk_ms = chunk_ms737 self.first_chunk_ms = first_chunk_ms if first_chunk_ms is not None else chunk_ms738 self.sample_rate = sample_rate739 self.n_fft = n_fft740 self.hop_length = hop_length741 self.n_mels = n_mels742 743 self.chunk_samples = int(round(chunk_ms * sample_rate / 1000))744 self.chunk_frames = self.chunk_samples // hop_length745 # align to hop_length to avoid frame boundary issues746 hop = self.hop_length747 raw_first_samples = int(round(self.first_chunk_ms * sample_rate / 1000))748 aligned_first = max(hop, (raw_first_samples // hop) * hop)749 self.first_chunk_samples = aligned_first750 self.half_window = n_fft // 2 # required right context751 752 # redundancy frames (in frames), <=1 frame: 10ms → 1 frame753 self.cnn_redundancy_ms = cnn_redundancy_ms754 self.cnn_redundancy_samples = int(cnn_redundancy_ms * sample_rate / 1000)755 self.cnn_redundancy_frames = max(0, self.cnn_redundancy_samples // hop_length)756 757 # sliding window configuration (Trigger mode)758 self.enable_sliding_window = enable_sliding_window759 self.trigger_seconds = slide_trigger_seconds760 self.slide_seconds = slide_stride_seconds761 762 # shift/base (global frame coordinates)763 self.left_samples_dropped = 0 # samples dropped from the left764 self.base_T = 0 # index of the "global frame" corresponding to mel_full[:, :, 0]765 766 self.reset()767 768 def reset(self):769 self.buffer = np.zeros(0, dtype=np.float32)770 self.last_emitted_T = 0771 self.total_samples_processed = 0772 self.chunk_count = 0773 self.is_first = True774 self.left_samples_dropped = 0775 self.base_T = 0776 777 def get_chunk_size(self) -> int:778 return self.first_chunk_samples if self.is_first else self.chunk_samples779 780 def get_expected_output_frames(self) -> int:781 raise NotImplementedError("get_expected_output_frames is not implemented")782 783 def _extract_full(self) -> torch.Tensor:784 # when buffer length is less than n_fft, Whisper's internal STFT will raise an error in center=True and pad mode785 # (pad is greater than input length). At this time, there is no stable frame to output, so return empty features directly.786 if len(self.buffer) < self.n_fft:787 raise ValueError(f"buffer length is shorter than n_fft {len(self.buffer)} < {self.n_fft}")788 # if buffer length is less than 5s, use set_spac_log_norm(log_floor_db=-10) or the last cached result789 if len(self.buffer) < 5 * self.sample_rate:790 # TODO: here the best is to do some experiments to choose the best one, now this is selected through experience, can see MiniCPMAAudioProcessor's main implementation791 self.feature_extractor.set_spac_log_norm(log_floor_db=-10)792 # if buffer length is greater than 5s, use set_spac_log_norm(dynamic_range_db=8)793 else:794 self.feature_extractor.set_spac_log_norm(dynamic_range_db=8)795 feats = self.feature_extractor(796 self.buffer,797 sampling_rate=self.sample_rate,798 return_tensors="pt",799 padding=False,800 )801 return feats.input_features # [1, 80, T]802 803 def _stable_frames_count(self) -> int:804 # number of stable frames = floor((len(buffer) - half_window) / hop) + 1, minimum is 0805 L = int(self.buffer.shape[0])806 if L <= 0:807 return 0808 if L < self.half_window:809 return 0810 return max(0, (L - self.half_window) // self.hop_length + 1)811 812 def _maybe_slide_buffer(self):813 """Trigger mode sliding window: when the buffer reaches the trigger threshold, slide a fixed length window."""814 if not self.enable_sliding_window:815 return816 817 sr = self.sample_rate818 hop = self.hop_length819 L = len(self.buffer)820 821 # convert seconds to samples822 trigger_samples = int(self.trigger_seconds * sr)823 stride_samples = int(self.slide_seconds * sr)824 825 # check if the trigger threshold is reached826 if L < trigger_samples:827 return828 829 # calculate the number of samples to drop (fixed sliding stride_samples)830 drop = stride_samples831 832 # cannot drop the left context that is still needed for subsequent emission833 # in trigger mode, we only need to protect the minimum necessary data834 # i.e. ensure that we do not discard frames that may be needed in the future835 last_emitted_local = self.last_emitted_T - self.base_T836 837 # only protect necessary context (e.g. the most recent 1 second data)838 min_keep_seconds = 1.0 # keep at least 1 second of data to ensure continuity839 min_keep_samples = int(min_keep_seconds * sr)840 841 # guard_samples are the minimum samples we must keep842 guard_samples = min(min_keep_samples, L - drop)843 844 # limit: do not exceed the safe boundary; and align hop845 max_allowed_drop = max(0, L - guard_samples)846 drop = min(drop, max_allowed_drop)847 drop = (drop // hop) * hop848 849 if drop <= 0:850 return851 852 # truly drop & update base853 self.buffer = self.buffer[drop:]854 self.left_samples_dropped += drop855 self.base_T += drop // hop856 857 def process(self, audio_chunk: np.ndarray, is_last_chunk: bool = False) -> Tuple[torch.Tensor, Dict]:858 self.chunk_count += 1859 # append to buffer860 if len(self.buffer) == 0:861 self.buffer = audio_chunk.astype(np.float32, copy=True)862 else:863 self.buffer = np.concatenate([self.buffer, audio_chunk.astype(np.float32, copy=True)])864 865 # sliding window processing866 self._maybe_slide_buffer()867 868 # full extraction (for the current window)869 mel_full = self._extract_full()870 T_full = mel_full.shape[-1] # local frames in the current window871 stable_T = min(T_full, self._stable_frames_count()) # local stable frames872 stable_T_global = self.base_T + stable_T # map to global frame coordinates873 874 # plan the core frames for the current emission (global coordinates)875 core_start_g = self.last_emitted_T876 core_end_g = core_start_g + self.chunk_frames877 required_stable_g = core_end_g + self.cnn_redundancy_frames878 879 if stable_T_global >= required_stable_g or is_last_chunk:880 emit_start_g = max(0, core_start_g - self.cnn_redundancy_frames)881 emit_end_g = core_end_g + self.cnn_redundancy_frames882 883 # global -> local index884 emit_start = max(0, emit_start_g - self.base_T)885 emit_end = emit_end_g - self.base_T886 emit_start = max(0, min(emit_start, T_full))887 emit_end = max(emit_start, min(emit_end, T_full))888 889 mel_output = mel_full[:, :, emit_start:emit_end]890 self.last_emitted_T = core_end_g # only advance the core frame pointer (global)891 else:892 mel_output = mel_full[:, :, 0:0]893 894 self.total_samples_processed += len(audio_chunk)895 self.is_first = False896 897 info = {898 "type": "exact_chunk",899 "chunk_number": self.chunk_count,900 "emitted_frames": mel_output.shape[-1],901 "stable_T": stable_T,902 "T_full": T_full,903 "base_T": self.base_T,904 "stable_T_global": stable_T_global,905 "buffer_len_samples": int(self.buffer.shape[0]),906 "left_samples_dropped": self.left_samples_dropped,907 "core_start": core_start_g, # if keep the original field name, use the global value here908 "core_end": core_end_g, # same as above909 }910 return mel_output, info911 912 def flush(self) -> torch.Tensor:913 """Called when the stream ends, output the remaining unemitted frames, ensuring consistency with offline (calculated by global coordinates)."""914 if len(self.buffer) == 0:915 return torch.zeros(1, 80, 0)916 917 mel_full = self._extract_full()918 T_local = mel_full.shape[-1]919 T_global = self.base_T + T_local920 921 if self.last_emitted_T < T_global:922 start_l = max(0, self.last_emitted_T - self.base_T)923 tail = mel_full[:, :, start_l:]924 self.last_emitted_T = T_global925 return tail926 return mel_full[:, :, 0:0]927 928 def get_config(self) -> Dict:929 return {930 "chunk_ms": self.chunk_ms,931 "first_chunk_ms": self.first_chunk_ms,932 "effective_first_chunk_ms": self.first_chunk_samples / self.sample_rate * 1000.0,933 "sample_rate": self.sample_rate,934 "n_fft": self.n_fft,935 "hop_length": self.hop_length,936 "cnn_redundancy_ms": self.cnn_redundancy_ms,937 "cnn_redundancy_frames": self.cnn_redundancy_frames,938 "enable_sliding_window": self.enable_sliding_window,939 "trigger_seconds": self.trigger_seconds,940 "slide_seconds": self.slide_seconds,941 }942 943 def get_state(self) -> Dict:944 return {945 "chunk_count": self.chunk_count,946 "last_emitted_T": self.last_emitted_T,947 "total_samples_processed": self.total_samples_processed,948 "buffer_len": int(self.buffer.shape[0]),949 "base_T": self.base_T,950 "left_samples_dropped": self.left_samples_dropped,951 }952 953 def get_snapshot(self) -> Dict:954 """Get a complete state snapshot (including buffer), used for recovery from a fast start.955 956 Returns:957 A dictionary containing the complete state, which can be used to restore the snapshot958 """959 buffer_copy = self.buffer.copy()960 snapshot = {961 "chunk_count": self.chunk_count,962 "last_emitted_T": self.last_emitted_T,963 "total_samples_processed": self.total_samples_processed,964 "buffer": buffer_copy,965 "base_T": self.base_T,966 "left_samples_dropped": self.left_samples_dropped,967 "is_first": self.is_first,968 # save the state of the feature_extractor (key: ensure determinism of mel feature extraction)969 "fe_dynamic_log_norm": getattr(self.feature_extractor, "dynamic_log_norm", None),970 "fe_dynamic_range_db": getattr(self.feature_extractor, "dynamic_range_db", None),971 "fe_log_floor_db": getattr(self.feature_extractor, "log_floor_db", None),972 }973 974 return snapshot975 976 def restore_snapshot(self, snapshot: Dict) -> None:977 """Restore state from a snapshot978 979 Args:980 snapshot: the snapshot dictionary returned by get_snapshot981 """982 # record the state before restoration983 prev_state = {984 "chunk_count": self.chunk_count,985 "last_emitted_T": self.last_emitted_T,986 "buffer_len": len(self.buffer),987 }988 989 # restore state990 self.chunk_count = snapshot["chunk_count"]991 self.last_emitted_T = snapshot["last_emitted_T"]992 self.total_samples_processed = snapshot["total_samples_processed"]993 self.buffer = snapshot["buffer"].copy() # copy buffer994 self.base_T = snapshot["base_T"]995 self.left_samples_dropped = snapshot["left_samples_dropped"]996 self.is_first = snapshot["is_first"]997 998 # restore the state of the feature_extractor (key: ensure determinism of mel feature extraction)999 if snapshot.get("fe_dynamic_log_norm") is not None:1000 self.feature_extractor.dynamic_log_norm = snapshot["fe_dynamic_log_norm"]1001 if snapshot.get("fe_dynamic_range_db") is not None:1002 self.feature_extractor.dynamic_range_db = snapshot["fe_dynamic_range_db"]1003 if snapshot.get("fe_log_floor_db") is not None:1004 self.feature_extractor.log_floor_db = snapshot["fe_log_floor_db"]1005 1006 1007class MiniCPMOProcessor(ProcessorMixin):1008 attributes = ["image_processor", "audio_processor", "tokenizer"]1009 audio_processor_class = "AutoFeatureExtractor"1010 image_processor_class = "AutoImageProcessor"1011 tokenizer_class = "AutoTokenizer"1012 1013 def __init__(self, image_processor=None, audio_processor=None, tokenizer=None, **kwargs):1014 super().__init__(image_processor, audio_processor, tokenizer)1015 1016 self.version = image_processor.version if image_processor else None1017 # audio feature pooling step, needs to be consistent with config.audio_pool_step1018 self.pool_step = kwargs.get("audio_pool_step", 5)1019 1020 # initialize the streaming audio processor1021 self._streaming_mel_processor = None1022 if audio_processor is not None:1023 self._init_streaming_processor()1024 1025 def get_audio_placeholder(1026 self,1027 audio_lens: int,1028 chunk_input: bool = True,1029 chunk_length: int = 1,1030 ) -> str:1031 """1032 Public method to get audio placeholder string for vLLM integration.1033 1034 Args:1035 audio_lens: Length of audio in samples1036 chunk_input: Whether to use chunked processing1037 chunk_length: Chunk length in seconds1038 1039 Returns:1040 Audio placeholder string1041 """1042 pool_step = self.pool_step1043 feature_lens = math.ceil(audio_lens / self.audio_processor.hop_length)1044 1045 feature_lens = (feature_lens - 1) // 2 + 11046 output_lens = (feature_lens - pool_step) // pool_step + 11047 1048 if chunk_input:1049 fbank_feat_in_chunk = int(chunk_length * 100)1050 cnn_feat_in_chunk = (fbank_feat_in_chunk - 1) // 2 + 11051 audio_embeds_in_chunk = (cnn_feat_in_chunk - pool_step) // pool_step + 11052 num_audio_chunks = (output_lens + audio_embeds_in_chunk - 1) // audio_embeds_in_chunk1053 1054 place_holders = ""1055 total_unk_len = 01056 for _ in range(num_audio_chunks):1057 unk_len = min(audio_embeds_in_chunk, output_lens - total_unk_len)1058 place_holders += self.tokenizer.audio_start + "<unk>" * unk_len + self.tokenizer.audio_end1059 total_unk_len += unk_len1060 audio_placeholder = place_holders1061 else:1062 audio_placeholder = self.tokenizer.audio_start + "<unk>" * output_lens + self.tokenizer.audio_end1063 1064 return audio_placeholder1065 1066 def _init_streaming_processor(1067 self,1068 chunk_ms: int = 100,1069 cnn_redundancy_ms: int = 0,1070 *,1071 mode: str = "exact",1072 first_chunk_ms: Optional[int] = None,1073 enable_sliding_window: bool = False,1074 slide_trigger_seconds: float = 30.0,1075 slide_stride_seconds: float = 10.0,1076 ):1077 """Initialize the streaming processor1078 1079 Args:1080 chunk_ms: Chunk size in milliseconds, also the sliding step.1081 cnn_redundancy_ms: CNN boundary redundancy in milliseconds (before and after), 0 means standard mode.1082 mode: streaming processing mode, currently only supports "exact"1083 first_chunk_ms: the size of the first chunk (milliseconds), if not specified, it is the same as chunk_ms1084 enable_sliding_window: whether to enable sliding window (trigger mode)1085 slide_trigger_seconds: trigger threshold for sliding window in seconds1086 slide_stride_seconds: stride for sliding window in seconds1087 """1088 if mode == "exact":1089 self._streaming_mel_processor = StreamingMelProcessorExact(1090 feature_extractor=self.audio_processor,1091 chunk_ms=chunk_ms,1092 first_chunk_ms=first_chunk_ms,1093 sample_rate=16000,1094 cnn_redundancy_ms=cnn_redundancy_ms,1095 enable_sliding_window=enable_sliding_window,1096 slide_trigger_seconds=slide_trigger_seconds,1097 slide_stride_seconds=slide_stride_seconds,1098 )1099 else:1100 raise ValueError(f"Unsupported mode: {mode}, only 'exact' is supported")1101 self._streaming_mode = mode if mode in ["exact"] else ("exact")1102 1103 def set_streaming_mode(1104 self,1105 mode: str = "exact",1106 chunk_ms: int = 100,1107 cnn_redundancy_ms: int = 0,1108 *,1109 first_chunk_ms: Optional[int] = None,1110 enable_sliding_window: bool = False,1111 slide_trigger_seconds: float = 30.0,1112 slide_stride_seconds: float = 10.0,1113 ):1114 """Set streaming processing mode1115 1116 Args:1117 mode: streaming processing mode, currently only supports "exact"1118 chunk_ms: chunk size in milliseconds, also the sliding step.1119 cnn_redundancy_ms: CNN boundary redundancy in milliseconds (before and after), 0 means standard mode.1120 first_chunk_ms: the size of the first chunk (milliseconds), if not specified, it is the same as chunk_ms1121 enable_sliding_window: whether to enable sliding window (trigger mode)1122 slide_trigger_seconds: trigger threshold for sliding window in seconds1123 slide_stride_seconds: stride for sliding window in seconds1124 """1125 if self.audio_processor is None:1126 raise ValueError("audio_processor is not set, cannot initialize the streaming processor")1127 self._init_streaming_processor(1128 chunk_ms=chunk_ms,1129 cnn_redundancy_ms=cnn_redundancy_ms,1130 mode=mode,1131 first_chunk_ms=first_chunk_ms,1132 enable_sliding_window=enable_sliding_window,1133 slide_trigger_seconds=slide_trigger_seconds,1134 slide_stride_seconds=slide_stride_seconds,1135 )1136 1137 def process_image(1138 self,1139 images: Optional[ImageInput] = None,1140 do_pad: bool = True,1141 max_slice_nums: int = 1,1142 return_tensors: str = "pt",1143 ) -> MiniCPMOBatchFeature:1144 """Process image data1145 1146 Args:1147 images: input images1148 do_pad: whether to pad1149 max_slice_nums: maximum number of slices1150 return_tensors: return tensor type1151 Returns:1152 MiniCPMOBatchFeature object1153 """1154 if images is None:1155 return MiniCPMOBatchFeature(data={"pixel_values": [[]], "image_sizes": [[]], "tgt_sizes": [[]]})1156 1157 result = self.image_processor(1158 images, do_pad=do_pad, max_slice_nums=max_slice_nums, return_tensors=return_tensors1159 )1160 1161 model_inputs = {1162 "pixel_values": result.get("pixel_values", [[]]),1163 "image_sizes": result.get("image_sizes", [[]]),1164 "tgt_sizes": result.get("tgt_sizes", [[]]),1165 }1166 1167 return MiniCPMOBatchFeature(data=model_inputs)1168 1169 def process_audio(1170 self,1171 audios: Optional[Union[np.ndarray, List[np.ndarray]]] = None,1172 sampling_rate: int = 16000,1173 regroup_to_seconds: Optional[int] = None,1174 fps: int = 100,1175 ) -> MiniCPMOBatchFeature:1176 """Process audio data in batch1177 1178 Args:1179 audios: audio data1180 sampling_rate: sampling rate1181 regroup_to_seconds: regroup duration in seconds1182 fps: frames per second1183 Returns:1184 MiniCPMOBatchFeature object1185 """1186 if audios is None:1187 return MiniCPMOBatchFeature(data={"audio_features": [], "audio_feature_lens": []})1188 1189 audio_features, audio_feature_lens = process_audio_batch(1190 audios=audios,1191 feature_extractor=self.audio_processor,1192 sampling_rate=sampling_rate,1193 max_duration_seconds=30,1194 return_attention_mask=True,1195 )1196 1197 if regroup_to_seconds is not None and len(audio_features) > 0:1198 audio_features, audio_feature_lens = regroup_audio_features(1199 audio_features=audio_features,1200 audio_feature_lens=audio_feature_lens,