diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 15dd56361..2356270c0 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1548,6 +1548,7 @@ multimodal_model_archs = [ "Lfm2VlForConditionalGeneration", "LightOnOCRForConditionalGeneration", *MIMO_V2_MULTIMODAL_ARCHS, + "MiMoV2ASRForCausalLM", "MiniCPMO", "MiniCPMV", "Mistral3ForConditionalGeneration", diff --git a/python/sglang/srt/entrypoints/openai/transcription_adapters/__init__.py b/python/sglang/srt/entrypoints/openai/transcription_adapters/__init__.py index 353196fdd..bd8779133 100644 --- a/python/sglang/srt/entrypoints/openai/transcription_adapters/__init__.py +++ b/python/sglang/srt/entrypoints/openai/transcription_adapters/__init__.py @@ -7,6 +7,9 @@ from sglang.srt.entrypoints.openai.transcription_adapters.base import ( # noqa: ) # Import built-in adapters so they self-register via @register_transcription_adapter. +from sglang.srt.entrypoints.openai.transcription_adapters.mimo_v2_asr import ( # noqa: F401 + MiMoV2ASRAdapter, +) from sglang.srt.entrypoints.openai.transcription_adapters.qwen3_asr import ( # noqa: F401 Qwen3ASRAdapter, ) @@ -20,4 +23,5 @@ __all__ = [ "resolve_adapter", "WhisperAdapter", "Qwen3ASRAdapter", + "MiMoV2ASRAdapter", ] diff --git a/python/sglang/srt/entrypoints/openai/transcription_adapters/mimo_v2_asr.py b/python/sglang/srt/entrypoints/openai/transcription_adapters/mimo_v2_asr.py new file mode 100644 index 000000000..d074e3c04 --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/transcription_adapters/mimo_v2_asr.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from sglang.srt.entrypoints.openai.protocol import ( + TranscriptionRequest, + TranscriptionUsage, + TranscriptionVerboseResponse, +) +from sglang.srt.entrypoints.openai.transcription_adapters.base import ( + TranscriptionAdapter, + register_transcription_adapter, +) + + +@register_transcription_adapter("MiMoV2ASR") +class MiMoV2ASRAdapter(TranscriptionAdapter): + """Adapter for MiMo-V2-ASR. + + The multimodal processor (``MiMoV2ASRProcessor``) prepends the audio + placeholder ``<|sosp|><|empty|>...<|eosp|>`` when ``input_text`` lacks + one, so the request text can stay empty and the adapter only has to + supply sampling params and the verbose-response shape. + """ + + def build_sampling_params(self, request: TranscriptionRequest) -> dict: + return { + "temperature": request.temperature, + "max_new_tokens": 448, + } + + def build_verbose_response( + self, + request: TranscriptionRequest, + text: str, + ret: dict, + tokenizer, + usage: TranscriptionUsage, + ) -> TranscriptionVerboseResponse: + # MiMo-V2-ASR does not emit timestamp tokens; segments stay empty + # until a forced-aligner path is added. + return TranscriptionVerboseResponse( + language=request.language or "auto", + duration=round(request.audio_duration_s, 2), + text=text, + segments=[], + usage=usage, + ) diff --git a/python/sglang/srt/models/mimo_audio.py b/python/sglang/srt/models/mimo_audio.py index a90547920..24b18faac 100644 --- a/python/sglang/srt/models/mimo_audio.py +++ b/python/sglang/srt/models/mimo_audio.py @@ -8,7 +8,7 @@ import os import typing as tp from dataclasses import dataclass from functools import wraps -from typing import List, Optional, Tuple +from typing import Optional, Tuple import torch import torch.nn as nn @@ -1141,72 +1141,89 @@ class MiMoV2AudioConfig: return config -class MiMoAudioEncoder(nn.Module): - config: MiMoAudioEncoderConfig +class AudioEncoderMixin: + """LM model mixin that adds MiMo audio encoder components. - def __init__(self, config): - super().__init__() + Components are attached as top-level attributes (no ``audio_encoder.`` + prefix), matching the checkpoint state_dict layout. Inner naming + variations are normalized via ``AUDIO_WEIGHT_REMAP``. + + Hot config fields are cached as direct ``self.audio_*`` attributes at + build time so helper methods can stay short and uniform (no + ``self.audio_config.foo`` indirection inside hot paths). + + Subclasses call ``self.build_audio_encoder(audio_config)`` from their + ``__init__`` after the language model is constructed; the mixin's + ``get_audio_feature`` then handles audio item batching end-to-end. + """ + + AUDIO_WEIGHT_REMAP: tuple[tuple[str, str], ...] = ( + ("audio_projection", "projection"), + ("speech_group_downcast", "projection"), + ("audio_input_local_transformer", "input_local_transformer"), + ) + + def build_audio_encoder(self, config) -> None: if not isinstance(config, MiMoV2AudioConfig): - config_dict = ( - vars(config) if hasattr(config, "__dict__") else config.__dict__ - ) - config = MiMoV2AudioConfig(**config_dict) - self.config = config - self.server_args = get_global_server_args() - self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder - self.speech_empty_ids = self.parsed_speech_empty_ids() + cfg_dict = vars(config) if hasattr(config, "__dict__") else config.__dict__ + config = MiMoV2AudioConfig(**cfg_dict) + self.audio_channels = config.audio_channels self.audio_group_size = config.group_size self.audio_segment_size = config.audio_segment_size - speech_vocab_size = self._parse_maybe_list( - self.config.speech_vocab_size, self.config.audio_channels + self.audio_input_local_dim = config.input_local_dim + self.audio_input_full_attention = config.input_full_attention + self.audio_out_hidden_size = config.out_hidden_size + + speech_vocab_size = config._parse_maybe_list( + config.speech_vocab_size, self.audio_channels ) + speech_empty_ids = config._parse_maybe_list( + config.speech_zeroemb_idx, self.audio_channels + ) + input_local_config = Qwen2Config( - hidden_size=self.config.input_local_dim, - num_hidden_layers=self.config.input_local_layers, - num_attention_heads=self.config.input_local_attn_heads, - num_key_value_heads=self.config.input_local_attn_heads, - intermediate_size=self.config.input_local_intermediate_size, - attention_dropout=self.config.input_local_hidden_dropout, - rope_theta=self.config.rope_theta, - partial_rotary_factor=self.config.partial_rotary_factor, + hidden_size=self.audio_input_local_dim, + num_hidden_layers=config.input_local_layers, + num_attention_heads=config.input_local_attn_heads, + num_key_value_heads=config.input_local_attn_heads, + intermediate_size=config.input_local_intermediate_size, + attention_dropout=config.input_local_hidden_dropout, + rope_theta=config.rope_theta, + partial_rotary_factor=config.partial_rotary_factor, ) - input_local_config.head_dim = self.config.input_local_head_dim - + input_local_config.head_dim = config.input_local_head_dim self.input_local_transformer = Qwen2Model(input_local_config) - - if not self.config.add_post_norm: + if not config.add_post_norm: self.input_local_transformer.norm = nn.Identity() self.speech_embeddings = nn.ModuleList( [ nn.Embedding( speech_vocab_size[i], - self.config.input_local_dim, - padding_idx=self.speech_empty_ids[i], + self.audio_input_local_dim, + padding_idx=speech_empty_ids[i], ) - for i in range(self.config.audio_channels) + for i in range(self.audio_channels) ] ) - if self.config.projection_layers == 1: + if config.projection_layers == 1: self.projection = nn.Linear( - self.config.input_local_dim * self.config.group_size, - self.config.out_hidden_size, + self.audio_input_local_dim * self.audio_group_size, + self.audio_out_hidden_size, bias=False, ) - elif self.config.projection_layers == 2: + elif config.projection_layers == 2: self.projection = AudioProjection( - self.config.input_local_dim * self.config.group_size, - self.config.input_local_dim * self.config.group_size * 4, - self.config.out_hidden_size, + self.audio_input_local_dim * self.audio_group_size, + self.audio_input_local_dim * self.audio_group_size * 4, + self.audio_out_hidden_size, ) else: - raise ValueError( - f"Invalid projection layers: {self.config.projection_layers}" - ) + raise ValueError(f"Invalid projection layers: {config.projection_layers}") - model_path = self.server_args.model_path + model_path = get_global_server_args().model_path if not os.path.isdir(model_path): from huggingface_hub import snapshot_download @@ -1216,13 +1233,16 @@ class MiMoAudioEncoder(nn.Module): ) audio_tokenizer_path = os.path.join(model_path, "audio_tokenizer") dev = torch.device(f"cuda:{torch.cuda.current_device()}") - self.audio_tokenizer = self._load_audio_tokenizer(audio_tokenizer_path, dev) + self.audio_tokenizer = self._load_mimo_audio_tokenizer( + audio_tokenizer_path, dev + ) @staticmethod - def _load_audio_tokenizer(path: str, device: torch.device) -> MiMoAudioTokenizer: + def _load_mimo_audio_tokenizer( + path: str, device: torch.device + ) -> MiMoAudioTokenizer: """Load MiMoAudioTokenizer manually to avoid new-transformers compat issues.""" import json - import os from safetensors.torch import load_file @@ -1231,7 +1251,6 @@ class MiMoAudioEncoder(nn.Module): config_dict = json.load(f) config = MiMoAudioTokenizer.config_class(**config_dict) model = MiMoAudioTokenizer(config) - # Load weights from safetensors or pytorch bin safetensors_path = os.path.join(path, "model.safetensors") bin_path = os.path.join(path, "pytorch_model.bin") if os.path.exists(safetensors_path): @@ -1243,43 +1262,34 @@ class MiMoAudioEncoder(nn.Module): f"No model weights found in {path} " "(expected model.safetensors or pytorch_model.bin)" ) + # strict=False: upstream ckpt also carries decoder/vocoder weights + # that this encoder-only MiMoAudioTokenizer doesn't materialize. model.load_state_dict(state_dict, strict=False) model = model.to(device=device, dtype=torch.bfloat16) model.eval() model.requires_grad_(False) return model - def parsed_speech_empty_ids(self): - return self._parse_maybe_list( - self.config.speech_zeroemb_idx, self.config.audio_channels - ) - - def _parse_maybe_list(self, value: str | int, length: int) -> List[int]: - if isinstance(value, str) and "-" in value: - return [int(s) for s in value.split("-")] - return [int(value)] * length - - # adapted from mimo-audio - def apply_input_local_transformer(self, speech_embeddings: torch.Tensor): - output = self.input_local_transformer( + def apply_input_local_transformer( + self, speech_embeddings: torch.Tensor + ) -> torch.Tensor: + return self.input_local_transformer( inputs_embeds=speech_embeddings, return_dict=True, - is_causal=not self.config.input_full_attention, # for SDPA - ) - return output.last_hidden_state # [T//group_size, group_size, input_local_dim] + is_causal=not self.audio_input_full_attention, # for SDPA + ).last_hidden_state # [T//group_size, group_size, input_local_dim] def apply_speech_embeddings(self, audio_codes: torch.Tensor) -> torch.Tensor: - num_segments = audio_codes.shape[0] - _audio_embeddings = torch.zeros( - (num_segments, self.config.group_size, self.config.input_local_dim), + embeds = torch.zeros( + (audio_codes.shape[0], self.audio_group_size, self.audio_input_local_dim), dtype=next(self.speech_embeddings[0].parameters()).dtype, device=audio_codes.device, ) - for i in range(self.config.audio_channels): - _audio_embeddings.add_(self.speech_embeddings[i](audio_codes[:, :, i])) - return _audio_embeddings + for i in range(self.audio_channels): + embeds.add_(self.speech_embeddings[i](audio_codes[:, :, i])) + return embeds - def process_audio(self, audio): + def pad_audio_codes(self, audio: torch.Tensor) -> torch.Tensor: T = audio.shape[0] audio = audio[:, : self.audio_channels] padded_T = ( @@ -1299,17 +1309,19 @@ class MiMoAudioEncoder(nn.Module): + audio[-1, :], ], dim=0, - ) # pad using the last embedding - padded_audio = padded_audio.reshape( + ) + return padded_audio.reshape( padded_T // self.audio_group_size, self.audio_group_size, self.audio_channels, ) - return padded_audio def get_audio_feature(self, items) -> torch.Tensor: - # items: already audio-only MultimodalDataItem list from caller. - # Each item.feature is either one mel tensor or a list of mel tensors (e.g. long audio split into chunks). + """Compute audio features for a list of audio MultimodalDataItem. + + Each item.feature is either a mel tensor or a list of mel tensors + (long audio split into chunks). + """ all_mels = [] for item in items: f = item.feature @@ -1321,9 +1333,8 @@ class MiMoAudioEncoder(nn.Module): device = next(self.projection.parameters()).device dtype = next(self.projection.parameters()).dtype return torch.empty( - 0, self.config.out_hidden_size, device=device, dtype=dtype + 0, self.audio_out_hidden_size, device=device, dtype=dtype ) - # Batch tokenize: one encode_batch call for all mels device = next(self.audio_tokenizer.encoder.parameters()).device code_list = tokenize_audio_batch( all_mels, @@ -1331,20 +1342,16 @@ class MiMoAudioEncoder(nn.Module): segment_size=self.audio_segment_size, device=device, ) - codecs_to_concat = [] - for codecs in code_list: - padded_codes = self.process_audio( - codecs - ) # [T//group_size, group_size, audio_channels] - codecs_to_concat.append(padded_codes) - audio_codes = torch.cat( - codecs_to_concat, dim=0 - ) # [T//group_size, group_size, audio_channels] + audio_codes = torch.cat([self.pad_audio_codes(c) for c in code_list], dim=0) + embeds = self.apply_input_local_transformer( + self.apply_speech_embeddings(audio_codes) + ) + return self.projection(embeds.reshape(embeds.shape[0], -1)) - _audio_embeddings = self.apply_speech_embeddings(audio_codes) - audio_embeds = self.apply_input_local_transformer( - _audio_embeddings - ) # [T//group_size, group_size, input_local_dim] - B = audio_embeds.shape[0] - audio_embeds = self.projection(audio_embeds.reshape(B, -1)) - return audio_embeds + @classmethod + def remap_audio_weight_name(cls, name: str) -> str: + """Normalize inner audio weight name variations to canonical form.""" + for src, dst in cls.AUDIO_WEIGHT_REMAP: + if src in name: + return name.replace(src, dst) + return name diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index 4952148d9..33e35fe17 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -78,7 +78,7 @@ from sglang.srt.model_loader.weight_utils import ( default_weight_loader, kv_cache_scales_loader, ) -from sglang.srt.models.mimo_audio import MiMoAudioEncoder, MiMoAudioEncoderConfig +from sglang.srt.models.mimo_audio import AudioEncoderMixin, MiMoAudioEncoderConfig from sglang.srt.models.mimo_vl import MiMoVisionTransformer, MiMoVLVisionConfig from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( @@ -993,7 +993,7 @@ class MiMoV2Model(nn.Module): ) -class MiMoV2ForCausalLM(nn.Module): +class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): # BitandBytes specific attributes default_bitsandbytes_target_modules = [ ".gate_proj.", @@ -1015,8 +1015,10 @@ class MiMoV2ForCausalLM(nn.Module): # Prefixes for weight routing in encoder_only/language_only modes _LANGUAGE_WEIGHT_PREFIXES = ("model.", "lm_head.") - _VISION_AUDIO_WEIGHT_PREFIXES = ("visual.", "vision_model.", "audio_") - _VISION_AUDIO_WEIGHT_SUBSTRING = "speech_embeddings" + _VISION_WEIGHT_PREFIXES = ("visual.", "vision_model.") + # ``audio_`` already covers ``audio_encoder.`` so a single prefix is enough. + _AUDIO_WEIGHT_PREFIXES = ("audio_",) + _AUDIO_WEIGHT_SUBSTRING = "speech_embeddings" def __init__( self, @@ -1070,8 +1072,7 @@ class MiMoV2ForCausalLM(nn.Module): quant_config=None, prefix=add_prefix("visual", prefix), ) - self.audio_config = MiMoAudioEncoderConfig(**audio_config) - self.audio_encoder = MiMoAudioEncoder(self.audio_config) + self.build_audio_encoder(MiMoAudioEncoderConfig(**audio_config)) self._routed_experts_weights_of_layer = LazyValue( lambda: ( @@ -1126,9 +1127,6 @@ class MiMoV2ForCausalLM(nn.Module): assert video_grid_thw.dim() == 2, video_grid_thw.dim() return self.visual(pixel_values, grid_thw=video_grid_thw) - def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: - return self.audio_encoder.get_audio_feature(items) - @torch.inference_mode() def encode_video_audio(self, mm_inputs: Dict) -> Optional[torch.Tensor]: # EPD-side hook: encode audio tracks pulled from videos and trim to the @@ -1276,14 +1274,14 @@ class MiMoV2ForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) skipped_mtp_weights = False - def _is_vision_audio_weight(name): - return ( - name.startswith(self._VISION_AUDIO_WEIGHT_PREFIXES) - or self._VISION_AUDIO_WEIGHT_SUBSTRING in name + for name, loaded_weight in weights: + is_vision_weight = name.startswith(self._VISION_WEIGHT_PREFIXES) + is_audio_weight = ( + name.startswith(self._AUDIO_WEIGHT_PREFIXES) + or self._AUDIO_WEIGHT_SUBSTRING in name ) - for name, loaded_weight in weights: - if not self._is_multimodal and _is_vision_audio_weight(name): + if not self._is_multimodal and (is_vision_weight or is_audio_weight): continue if self.config.encoder_only and name.startswith( @@ -1291,61 +1289,21 @@ class MiMoV2ForCausalLM(nn.Module): ): continue - if self._is_multimodal and "audio" in name: - if "projection" in name: - if ( - "audio_encoder.audio_projection" in name - and "audio_encoder.projection" not in name - ): - name = name.replace( - "audio_encoder.audio_projection", "audio_encoder.projection" - ) - elif ( - "audio_projection" in name - and "audio_encoder.projection" not in name - ): - name = name.replace( - "audio_projection", "audio_encoder.projection" - ) - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader + if self._is_multimodal and is_audio_weight: + if name.startswith("audio_encoder."): + name = name[len("audio_encoder.") :] + name = self.remap_audio_weight_name(name) + if name not in params_dict: + logger.warning( + f"Audio param {name} not found in params_dict, skipping" ) - weight_loader(param, loaded_weight) continue - - if "input_local_transformer" in name: - if ( - "audio_input_local_transformer" in name - and "audio_encoder.input_local_transformer" not in name - ): - name = name.replace( - "audio_input_local_transformer", - "audio_encoder.input_local_transformer", - ) - if name not in params_dict: - logger.warning( - f"Parameter {name} not found in params_dict, skipping" - ) - continue - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - continue - - if self._is_multimodal and "speech_embeddings" in name: - if ( - "speech_embeddings" in name - and "audio_encoder.speech_embeddings" not in name - ): - name = name.replace( - "speech_embeddings", "audio_encoder.speech_embeddings" - ) param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight[: param.shape[0], :]) + if self._AUDIO_WEIGHT_SUBSTRING in name: + weight_loader(param, loaded_weight[: param.shape[0], :]) + else: + weight_loader(param, loaded_weight) continue if self._is_multimodal and "visual" in name: diff --git a/python/sglang/srt/models/mimo_v2_asr.py b/python/sglang/srt/models/mimo_v2_asr.py new file mode 100644 index 000000000..4d63f3a9e --- /dev/null +++ b/python/sglang/srt/models/mimo_v2_asr.py @@ -0,0 +1,162 @@ +"""MiMo-V2-ASR model. + +Reuses the LM scaffold of ``MiMoForCausalLM`` and adds audio encoder +components via ``AudioEncoderMixin``. The encoder modules are attached as +top-level attributes (no ``audio_encoder.`` prefix) so the checkpoint +state_dict aligns 1:1 with ``self.named_parameters()``. +""" + +import logging +from typing import Any, Iterable, List, Optional, Tuple + +import torch + +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models import mimo_audio as _mimo_audio_module +from sglang.srt.models.mimo import MiMoForCausalLM +from sglang.srt.models.mimo_audio import AudioEncoderMixin, MiMoAudioEncoderConfig + +logger = logging.getLogger(__name__) + + +def _maybe_override_audio_attn_for_blackwell() -> None: + """Swap mimo_audio.flash_attn_varlen_func to upstream FA2 on GPUs that + sgl-kernel's FA3 doesn't support. + + sgl-kernel FA3 only covers sm80/86/89/90 — on Blackwell consumer cards + (sm_120 / RTX 50xx) its varlen kernel raises NotImplementedError. ASR is + small enough to be deployed on those GPUs, so when FA3 isn't supported + we replace the module-level reference with upstream flash-attn (FA2), + which works on sm_120. No-op on supported GPUs (FA3 stays). + + MiMo-V2 (the heavy multimodal model) is only deployed on H100/A100, so + this override never triggers in its hot path. + """ + try: + from sgl_kernel.flash_attn import is_fa3_supported + except ImportError: + return + if is_fa3_supported(): + return + try: + from flash_attn import flash_attn_varlen_func + except ImportError as e: + raise RuntimeError( + "MiMo-V2-ASR audio encoder needs upstream flash-attn on this GPU " + "(sgl-kernel FA3 doesn't support sm_120). Install with " + "`pip install flash-attn --no-build-isolation`." + ) from e + _mimo_audio_module.flash_attn_varlen_func = flash_attn_varlen_func + + +MiMoV2ASRConfig = Any + +# Top-level audio sub-module name prefixes (after AUDIO_WEIGHT_REMAP). Loaded +# directly by default_weight_loader because the LM branch's qkv/gate-up fused +# stacked-params mapping doesn't apply to the vanilla HF Qwen2Model used +# inside the audio encoder. +_AUDIO_NAME_PREFIXES: Tuple[str, ...] = ( + "projection.", + "input_local_transformer.", + "speech_embeddings.", +) + +# Training-only weights present in checkpoint but not used at inference. +# Checked AFTER the audio-prefix load path so substring matching here is +# safe: legitimate audio weights (``input_local_transformer.*``) are +# already consumed by ``_AUDIO_NAME_PREFIXES`` above. +_SKIP_NAME_SUBSTRINGS: Tuple[str, ...] = ( + "hidden_states_downcast", + "local_transformer", +) + + +class MiMoV2ASRForCausalLM(MiMoForCausalLM, AudioEncoderMixin): + def __init__( + self, + config: MiMoV2ASRConfig, + quant_config=None, + prefix: str = "", + ) -> None: + _maybe_override_audio_attn_for_blackwell() + super().__init__(config, quant_config=quant_config, prefix=prefix) + self.build_audio_encoder(MiMoAudioEncoderConfig(**config.audio_config)) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + pattern = MultiModalityDataPaddingPatternMultimodalTokens() + return pattern.pad_input_tokens(input_ids, mm_inputs) + + def get_input_embeddings(self): + if getattr(self.config, "encoder_only", False): + return None + return self.model.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = False, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + if getattr(self.config, "encoder_only", False): + raise NotImplementedError( + "forward() is not supported in encoder_only mode. " + "Use get_audio_feature() instead." + ) + + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + multimodal_model=self, + positions=positions, + pp_proxy_tensors=pp_proxy_tensors, + ) + + if not get_embedding: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + return self.pooler(hidden_states, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + params_dict = dict(self.named_parameters()) + deferred: List[Tuple[str, torch.Tensor]] = [] + + for name, loaded_weight in weights: + if name.startswith("audio_encoder."): + name = name[len("audio_encoder.") :] + name = self.remap_audio_weight_name(name) + + if name.startswith(_AUDIO_NAME_PREFIXES): + if name not in params_dict: + logger.warning( + f"Audio param {name} not found in params_dict, skipping" + ) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if name.startswith("speech_embeddings."): + weight_loader(param, loaded_weight[: param.shape[0], :]) + else: + weight_loader(param, loaded_weight) + continue + + if any(s in name for s in _SKIP_NAME_SUBSTRINGS): + continue + + deferred.append((name, loaded_weight)) + + super().load_weights(iter(deferred)) + + +EntryClass = MiMoV2ASRForCausalLM diff --git a/python/sglang/srt/multimodal/processors/mimo_audio.py b/python/sglang/srt/multimodal/processors/mimo_audio.py new file mode 100644 index 000000000..3bae8ea16 --- /dev/null +++ b/python/sglang/srt/multimodal/processors/mimo_audio.py @@ -0,0 +1,318 @@ +"""Stateful audio preprocessing pipeline shared by MiMo multimodal and ASR processors.""" + +import io +import math +import os +import time +from collections import OrderedDict +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import pybase64 +import requests +import torch + +from sglang.utils import logger + +try: + from torchcodec.decoders import AudioDecoder +except ImportError: + logger.warning( + "torchcodec is not installed; audio inputs will fail at request time" + ) + AudioDecoder = None + +try: + import torchaudio + from torchaudio.transforms import MelSpectrogram +except ImportError: + logger.warning( + "torchaudio is not installed; audio inputs will fail at request time" + ) + torchaudio = None + MelSpectrogram = None + + +@dataclass +class AudioInput: + """ + if audio is str or bytes, only load it as mel spectrogram. + if audio is tuple, it is (waveform, original_sr) + if audio is torch.Tensor, it is tokenized input ids with shape (T, n_vq+). + if audio is np.ndarray, it is a pre-loaded waveform (1D, already resampled). + """ + + audio: str | bytes | tuple | torch.Tensor | np.ndarray + + def __post_init__(self): + if not isinstance(self.audio, (str, bytes, tuple, torch.Tensor, np.ndarray)): + raise ValueError( + f"audio must be a str, bytes, tuple, torch.Tensor, or np.ndarray, but got {type(self.audio)}" + ) + if isinstance(self.audio, tuple): + if ( + len(self.audio) != 2 + or not isinstance(self.audio[0], torch.Tensor) + or not isinstance(self.audio[1], (int, float)) + ): + raise ValueError( + f"audio must be a tuple of (waveform-T, original_sr-int/float), but got {len(self.audio)} elements and {type(self.audio[0])} and {type(self.audio[1])}" + ) + if self.audio[0].ndim != 1: + raise ValueError( + f"waveform must be a 1D tensor, but got {self.audio[0].ndim}D tensor" + ) + if self.audio[1] <= 0: + raise ValueError( + f"original_sr must be a positive number, but got {self.audio[1]}" + ) + if isinstance(self.audio, torch.Tensor) and self.audio.ndim != 2: + raise ValueError( + f"audio must be a 2D tensor, but got {self.audio.ndim}D tensor" + ) + + +class MiMoAudioPipeline: + """Stateful audio preprocessing pipeline. + + Composable: held by both MiMoProcessor (multimodal) and MiMoV2ASRProcessor. + Owns the mel spectrogram, resampler cache, http session, and the special + token ids for ``<|sosp|> <|empty|>* <|eosp|>`` placeholders. + """ + + def __init__( + self, + *, + audio_token_id: int, + audio_start_token_id: int, + audio_end_token_id: int, + audio_kernel_size: int = 3, + audio_stride_size: int = 2, + audio_avg_pooler: int = 2, + audio_group_size: int = 4, + audio_channels: int = 8, + audio_sampling_rate: int = 24000, + audio_nfft: int = 960, + audio_hop_length: int = 240, + audio_window_size: int = 960, + audio_fmin: int = 0, + audio_fmax: Optional[int] = None, + audio_n_mels: int = 128, + audio_input_id_per_second: int = 25, + max_resamplers: int = 16, + ) -> None: + self.audio_token_id = audio_token_id + self.audio_start_token_id = audio_start_token_id + self.audio_end_token_id = audio_end_token_id + + self.audio_kernel_size = audio_kernel_size + self.audio_stride_size = audio_stride_size + self.audio_avg_pooler = audio_avg_pooler + self.audio_group_size = audio_group_size + self.audio_channels = audio_channels + + self.audio_sampling_rate = audio_sampling_rate + self.audio_nfft = audio_nfft + self.audio_hop_length = audio_hop_length + self.audio_window_size = audio_window_size + self.audio_fmin = audio_fmin + self.audio_fmax = audio_fmax + self.audio_n_mels = audio_n_mels + self.audio_input_id_per_second = audio_input_id_per_second + + self.mel_spectrogram_kwargs = dict( + sample_rate=audio_sampling_rate, + n_fft=audio_nfft, + hop_length=audio_hop_length, + win_length=audio_window_size, + f_min=audio_fmin, + f_max=audio_fmax, + n_mels=audio_n_mels, + power=1.0, + center=True, + ) + self._mel_spectrogram = None + self._resamplers: "OrderedDict[int, torchaudio.transforms.Resample]" = ( + OrderedDict() + ) + self._resamplers_max = max_resamplers + + self.http_session = requests.Session() + + @property + def audio_token_per_second(self) -> float: + return self.audio_input_id_per_second / self.audio_group_size + + @staticmethod + def _ensure_audio_dependencies() -> None: + if torchaudio is None or MelSpectrogram is None: + raise RuntimeError( + "torchaudio is required for audio inputs; install torchaudio" + ) + + @property + def mel_spectrogram(self): + self._ensure_audio_dependencies() + if self._mel_spectrogram is None: + self._mel_spectrogram = MelSpectrogram(**self.mel_spectrogram_kwargs) + return self._mel_spectrogram + + def compute_audio_token_len(self, mel_len: int) -> int: + n = mel_len + 3 - self.audio_kernel_size + n = (n + 2 - self.audio_kernel_size) // self.audio_stride_size + 1 + n = n // self.audio_avg_pooler + int(n % self.audio_avg_pooler != 0) + return math.ceil(n / self.audio_group_size) + + def preprocess_audio(self, audio): + """Load audio source → log-mel spectrogram + token length. + + Input: filename string, bytes, or tuple of (waveform, original_sr). + Output: (mel-spectrogram tensor [T, n_mels], audio_token_len int). + """ + self._ensure_audio_dependencies() + assert isinstance( + audio, (str, bytes, tuple) + ), f"audio must be a str, bytes or tuple, but got {type(audio)}" + if isinstance(audio, tuple): + waveform, original_sr = audio + else: + if isinstance(audio, bytes): + file = io.BytesIO(audio) + elif isinstance(audio, str): + if audio.startswith("data:"): + file = io.BytesIO( + pybase64.b64decode(audio.split(",")[1], validate=True) + ) + elif audio.startswith("http://") or audio.startswith("https://"): + dl_start = time.perf_counter() + timeout = int(os.getenv("REQUEST_TIMEOUT", "5")) + try: + response = self.http_session.get( + audio, stream=True, timeout=timeout + ) + dl_elapsed_ms = (time.perf_counter() - dl_start) * 1000 + if dl_elapsed_ms > 1000.0: + content_len = len(response.content) + logger.warning( + f"Slow audio download: {dl_elapsed_ms:.2f}ms, " + f"size={content_len / 1024:.1f}KB, url={audio}" + ) + file = io.BytesIO(response.content) + response.close() + except Exception as e: + dl_elapsed_ms = (time.perf_counter() - dl_start) * 1000 + logger.error( + f"Failed to download audio: {dl_elapsed_ms:.2f}ms, " + f"error={type(e).__name__}: {e}, url={audio}" + ) + raise + else: + file = audio + if AudioDecoder is None: + raise RuntimeError( + "torchcodec is required for audio decoding; install with `pip install torchcodec`." + ) + try: + samples = AudioDecoder(file).get_all_samples() + except RuntimeError as e: + audio_source = ( + audio + if isinstance(audio, str) + and (audio.startswith("http://") or audio.startswith("https://")) + else "" + ) + logger.error(f"Failed to decode audio: {e}, source={audio_source}") + raise ValueError( + f"Invalid audio format: source={audio_source}, detail={e}" + ) from e + waveform = samples.data + original_sr = samples.sample_rate + + if original_sr != self.audio_sampling_rate: + if original_sr in self._resamplers: + self._resamplers.move_to_end(original_sr) + else: + if len(self._resamplers) >= self._resamplers_max: + self._resamplers.popitem(last=False) + self._resamplers[original_sr] = torchaudio.transforms.Resample( + orig_freq=original_sr, new_freq=self.audio_sampling_rate + ) + waveform = self._resamplers[original_sr](waveform) + if waveform.ndim == 2: + waveform = waveform.mean(dim=0) + spec = self.mel_spectrogram(waveform[None, :]) + spec = torch.log(torch.clip(spec, min=1e-7)).squeeze() + spec = spec.transpose(0, 1) + + audio_token_len = self.compute_audio_token_len(spec.shape[0]) + return spec, audio_token_len + + def process_audio(self, audio_input: AudioInput): + """Dispatch on the underlying audio payload. + + - str/bytes/tuple/np.ndarray waveform → returns (mel-spec, token_len) tuple + - 2D tensor of pre-tokenized audio codes → returns padded codes tensor + shaped [T//group, group, channels] + """ + audio = audio_input.audio + if isinstance(audio, np.ndarray): + waveform = torch.from_numpy(audio).float() + audio = (waveform, self.audio_sampling_rate) + if isinstance(audio, (str, bytes, tuple)): + return self.preprocess_audio(audio) + + assert ( + audio.shape[1] >= self.audio_channels + ), f"audio must have at least {self.audio_channels} channels, but got {audio.shape[1]}" + T = audio.shape[0] + audio = audio[:, : self.audio_channels].to(torch.long) + padded_T = ( + (T + self.audio_group_size - 1) + // self.audio_group_size + * self.audio_group_size + ) + padded_audio = torch.cat( + [ + audio, + torch.zeros(padded_T - T, self.audio_channels, dtype=torch.long) + + audio[-1, :], + ], + dim=0, + ) + padded_audio = padded_audio.reshape( + padded_T // self.audio_group_size, + self.audio_group_size, + self.audio_channels, + ) + return padded_audio + + def build_audio_placeholder_ids(self, audio_token_len: int) -> list[int]: + return ( + [self.audio_start_token_id] + + [self.audio_token_id] * audio_token_len + + [self.audio_end_token_id] + ) + + def process_audio_input(self, audio_input: AudioInput) -> dict: + """Run process_audio and produce the placeholder input_ids. + + Replaces the duplicated _process_audio_content bodies in both processors. + Returns dict with input_ids, audio_input (mel or codes), and is_tokenized. + """ + processed = self.process_audio(audio_input) + if isinstance(processed, tuple): + is_tokenized = False + audio_spec, audio_token_len = processed + payload = audio_spec + else: + is_tokenized = True + audio_token_len = processed.shape[0] + payload = processed + + return { + "input_ids": self.build_audio_placeholder_ids(audio_token_len), + "audio_input": payload, + "audio_token_len": audio_token_len, + "is_tokenized": is_tokenized, + } diff --git a/python/sglang/srt/multimodal/processors/mimo_v2.py b/python/sglang/srt/multimodal/processors/mimo_v2.py index 0c6c82d71..91397692f 100644 --- a/python/sglang/srt/multimodal/processors/mimo_v2.py +++ b/python/sglang/srt/multimodal/processors/mimo_v2.py @@ -3,21 +3,16 @@ import asyncio import base64 import copy -import io import json import math -import os import re import subprocess -import time -from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from io import BytesIO from typing import List, Literal, Optional, Union import numpy as np -import pybase64 import requests import torch import torch.nn.functional as F @@ -39,20 +34,14 @@ from sglang.srt.multimodal.processors.base_processor import ( BaseMultimodalProcessor, MultimodalSpecialTokens, ) +from sglang.srt.multimodal.processors.mimo_audio import ( + AudioInput, + MiMoAudioPipeline, +) from sglang.srt.multimodal.processors.qwen_vl import smart_nframes from sglang.srt.utils import ImageData, VideoData from sglang.utils import logger -try: - import torchaudio - from torchaudio.transforms import MelSpectrogram -except ImportError: - logger.warning( - "torchaudio is not installed; audio inputs will fail at request time" - ) - torchaudio = None - MelSpectrogram = None - @dataclass class ImageInput: @@ -112,45 +101,6 @@ class VideoInput: ) -@dataclass -class AudioInput: - """ - if audio is str or bytes, only load it as mel spectrogram. - if audio is tuple, it is (waveform, original_sr) - if audio is torch.Tensor, it is tokenized input ids with shape (T, n_vq+). - if audio is np.ndarray, it is a pre-loaded waveform (1D, already resampled). - """ - - audio: str | bytes | tuple | torch.Tensor | np.ndarray - - def __post_init__(self): - if not isinstance(self.audio, (str, bytes, tuple, torch.Tensor, np.ndarray)): - raise ValueError( - f"audio must be a str, bytes, tuple, torch.Tensor, or np.ndarray, but got {type(self.audio)}" - ) - if isinstance(self.audio, tuple): - if ( - len(self.audio) != 2 - or not isinstance(self.audio[0], torch.Tensor) - or not isinstance(self.audio[1], (int, float)) - ): - raise ValueError( - f"audio must be a tuple of (waveform-T, original_sr-int/float), but got {len(self.audio)} elements and {type(self.audio[0])} and {type(self.audio[1])}" - ) - if self.audio[0].ndim != 1: - raise ValueError( - f"waveform must be a 1D tensor, but got {self.audio[0].ndim}D tensor" - ) - if self.audio[1] <= 0: - raise ValueError( - f"original_sr must be a positive number, but got {self.audio[1]}" - ) - if isinstance(self.audio, torch.Tensor) and self.audio.ndim != 2: - raise ValueError( - f"audio must be a 2D tensor, but got {self.audio.ndim}D tensor" - ) - - @dataclass class VideoAudioInput: video: str | bytes | tuple[torch.Tensor, torch.Tensor] @@ -334,11 +284,9 @@ class MiMoProcessor: audio_fmin=0, audio_fmax=None, audio_n_mels=128, - audio_segment_size=6000, audio_channels=8, audio_group_size=4, audio_input_id_per_second=25, - audio_zeroemb_idx=4096, image_min_pixels=None, image_max_pixels=None, video_min_pixels=None, @@ -395,11 +343,8 @@ class MiMoProcessor: self.image_token_id = image_token_id self.video_token_id = video_token_id - self.audio_token_id = audio_token_id self.vision_start_token_id = vision_start_token_id self.vision_end_token_id = vision_end_token_id - self.audio_start_token_id = audio_start_token_id - self.audio_end_token_id = audio_end_token_id self.video_start_token_id = video_start_token_id self.video_end_token_id = video_end_token_id self.pad_token_id = pad_token_id @@ -411,60 +356,24 @@ class MiMoProcessor: self.video_tokens_per_second = video_tokens_per_second - self.audio_sampling_rate = audio_sampling_rate - self.audio_nfft = audio_nfft - self.audio_hop_length = audio_hop_length - self.audio_window_size = audio_window_size - self.audio_fmin = audio_fmin - self.audio_fmax = audio_fmax - self.audio_n_mels = audio_n_mels - - self.audio_segment_size = audio_segment_size - - self.audio_kernel_size = audio_kernel_size - self.audio_stride_size = audio_stride_size - self.audio_avg_pooler = audio_avg_pooler - - self.mel_spectrogram_kwargs = dict( - sample_rate=audio_sampling_rate, - n_fft=audio_nfft, - hop_length=audio_hop_length, - win_length=audio_window_size, - f_min=audio_fmin, - f_max=audio_fmax, - n_mels=audio_n_mels, - power=1.0, - center=True, + self.audio_pipeline = MiMoAudioPipeline( + audio_token_id=audio_token_id, + audio_start_token_id=audio_start_token_id, + audio_end_token_id=audio_end_token_id, + audio_kernel_size=audio_kernel_size, + audio_stride_size=audio_stride_size, + audio_avg_pooler=audio_avg_pooler, + audio_group_size=audio_group_size, + audio_channels=audio_channels, + audio_sampling_rate=audio_sampling_rate, + audio_nfft=audio_nfft, + audio_hop_length=audio_hop_length, + audio_window_size=audio_window_size, + audio_fmin=audio_fmin, + audio_fmax=audio_fmax, + audio_n_mels=audio_n_mels, + audio_input_id_per_second=audio_input_id_per_second, ) - self._mel_spectrogram = None - self._resamplers = OrderedDict() - self._resamplers_max = 16 - - self.audio_channels = audio_channels - self.audio_group_size = audio_group_size - self.audio_input_id_per_second = audio_input_id_per_second - if isinstance(audio_zeroemb_idx, int): - self.audio_zeroemb_idxs = torch.tensor( - [audio_zeroemb_idx] * self.audio_channels, dtype=torch.int32 - ) - elif isinstance(audio_zeroemb_idx, list): - if len(audio_zeroemb_idx) == 1: - self.audio_zeroemb_idxs = torch.tensor( - audio_zeroemb_idx * self.audio_channels, dtype=torch.int32 - ) - elif len(audio_zeroemb_idx) == self.audio_channels: - self.audio_zeroemb_idxs = torch.tensor( - audio_zeroemb_idx, dtype=torch.int32 - ) - else: - raise ValueError( - f"audio_zeroemb_idx must be a list of 1 or {self.audio_channels} integers, but got {len(audio_zeroemb_idx)}" - ) - else: - raise ValueError( - f"audio_zeroemb_idx must be an integer or a list of {self.audio_channels} integers, but got {type(audio_zeroemb_idx)}" - ) - assert image_min_pixels is not None assert image_max_pixels is not None assert video_min_pixels is not None @@ -487,10 +396,18 @@ class MiMoProcessor: "min_frames": min_frames, } - self.http_session = requests.Session() for k in kwargs: logger.info(f"[Warning] Ignored unknown parameter {k} for MiMoProcessor") + def __getattr__(self, name): + # Delegate audio_pipeline fields so callers can use self.audio_token_id + # etc. directly. Only triggers when normal attribute lookup fails; + # __dict__.get avoids recursion before audio_pipeline is assigned. + pipeline = self.__dict__.get("audio_pipeline") + if pipeline is not None and hasattr(pipeline, name): + return getattr(pipeline, name) + raise AttributeError(name) + @classmethod def from_hf_config(cls, hf_config, mm_config=None, **overrides): # Params must come from hf_config.processor_config so E and D agree; @@ -576,13 +493,13 @@ class MiMoProcessor: return _ffprobe_has_audio(path_or_data, stdin=None, label=path_or_data) if isinstance(path_or_data, bytes): - source = io.BytesIO(path_or_data) + source = BytesIO(path_or_data) elif ( isinstance(path_or_data, str) and path_or_data.startswith("data:") and ";base64," in path_or_data ): - source = io.BytesIO(base64.b64decode(path_or_data.split(";base64,")[1])) + source = BytesIO(base64.b64decode(path_or_data.split(";base64,")[1])) else: source = path_or_data # local path or file:// try: @@ -667,7 +584,9 @@ class MiMoProcessor: all_timestamps.extend(aligned_ts[::step].tolist()) if self.has_audio_track(video_blob): - audio_spec, audio_token_len = self.preprocess_audio(video_blob) + audio_spec, audio_token_len = self.audio_pipeline.preprocess_audio( + video_blob + ) units = self._build_video_audio_units( grid, aligned_ts, @@ -704,7 +623,7 @@ class MiMoProcessor: for audio in mm_data: if isinstance(audio, np.ndarray): audio = (torch.from_numpy(audio).float(), self.audio_sampling_rate) - spec, token_len = self.preprocess_audio(audio) + spec, token_len = self.audio_pipeline.preprocess_audio(audio) all_specs.append(spec) all_lens.append(token_len) return { @@ -714,20 +633,6 @@ class MiMoProcessor: raise ValueError(f"Unsupported modality for EPD preprocessing: {modality}") - @property - def mel_spectrogram(self): - self._ensure_audio_dependencies() - if self._mel_spectrogram is None: - self._mel_spectrogram = MelSpectrogram(**self.mel_spectrogram_kwargs) - return self._mel_spectrogram - - @staticmethod - def _ensure_audio_dependencies(): - if torchaudio is None or MelSpectrogram is None: - raise RuntimeError( - "torchaudio is required for audio inputs; install torchaudio" - ) - def prepare_image_kwargs(self, image: ImageInput): kwargs = {} for k in ["min_pixels", "max_pixels"]: @@ -764,95 +669,6 @@ class MiMoProcessor: raise ValueError("Video sampling strategy not specified") return kwargs - def preprocess_audio(self, audio: str | bytes): - self._ensure_audio_dependencies() - """ - - Input: audio filename string, bytes, or tuple of (waveform, original_sr) - - Output: - - mel spectrogram: torch.Tensor (T, n_mels) - - number of tokens: int - """ - assert isinstance( - audio, (str, bytes, tuple) - ), f"audio must be a str, bytes or tuple, but got {type(audio)}" - if isinstance(audio, tuple): - waveform, original_sr = audio - else: - if isinstance(audio, bytes): - file = io.BytesIO(audio) - elif isinstance(audio, str): - if audio.startswith("data:"): - file = io.BytesIO( - pybase64.b64decode(audio.split(",")[1], validate=True) - ) - elif audio.startswith("http://") or audio.startswith("https://"): - dl_start = time.perf_counter() - timeout = int(os.getenv("REQUEST_TIMEOUT", "5")) - try: - response = self.http_session.get( - audio, stream=True, timeout=timeout - ) - dl_elapsed_ms = (time.perf_counter() - dl_start) * 1000 - if dl_elapsed_ms > 1000.0: - content_len = len(response.content) - logger.warning( - f"Slow audio download: {dl_elapsed_ms:.2f}ms, " - f"size={content_len / 1024:.1f}KB, url={audio}" - ) - file = io.BytesIO(response.content) - response.close() - except Exception as e: - dl_elapsed_ms = (time.perf_counter() - dl_start) * 1000 - logger.error( - f"Failed to download audio: {dl_elapsed_ms:.2f}ms, " - f"error={type(e).__name__}: {e}, url={audio}" - ) - raise - else: - file = audio - try: - samples = AudioDecoder(file).get_all_samples() - except RuntimeError as e: - audio_source = ( - audio - if isinstance(audio, str) - and (audio.startswith("http://") or audio.startswith("https://")) - else "" - ) - logger.error(f"Failed to decode audio: {e}, source={audio_source}") - raise ValueError( - f"Invalid audio format: source={audio_source}, detail={e}" - ) from e - waveform = samples.data - original_sr = samples.sample_rate - - if original_sr != self.audio_sampling_rate: - if original_sr in self._resamplers: - self._resamplers.move_to_end(original_sr) - else: - if len(self._resamplers) >= self._resamplers_max: - self._resamplers.popitem(last=False) - self._resamplers[original_sr] = torchaudio.transforms.Resample( - orig_freq=original_sr, new_freq=self.audio_sampling_rate - ) - waveform = self._resamplers[original_sr](waveform) - if waveform.ndim == 2: - waveform = waveform.mean(dim=0) - spec = self.mel_spectrogram(waveform[None, :]) - spec = torch.log(torch.clip(spec, min=1e-7)).squeeze() - spec = spec.transpose(0, 1) - - audio_token_len = spec.shape[0] + 3 - self.audio_kernel_size - audio_token_len = ( - audio_token_len + 2 - self.audio_kernel_size - ) // self.audio_stride_size + 1 - audio_token_len = audio_token_len // self.audio_avg_pooler + int( - audio_token_len % self.audio_avg_pooler != 0 - ) - audio_token_len = math.ceil(audio_token_len / self.audio_group_size) - - return spec, audio_token_len - def process_image(self, image: ImageInput): kwargs = self.prepare_image_kwargs(image) image = image.image @@ -1014,40 +830,6 @@ class MiMoProcessor: ) return visual_patches, thw_grid, aligned_timestamps, video_meta - def process_audio(self, audio: AudioInput): - audio = audio.audio - if isinstance(audio, np.ndarray): - waveform = torch.from_numpy(audio).float() - audio = (waveform, self.audio_sampling_rate) - if isinstance(audio, (str, bytes, tuple)): - audio_spec, audio_token_len = self.preprocess_audio(audio) - return audio_spec, audio_token_len - - assert ( - audio.shape[1] >= self.audio_channels - ), f"audio must have at least {self.audio_channels} channels, but got {audio.shape[1]}" - T = audio.shape[0] - audio = audio[:, : self.audio_channels].to(torch.long) - padded_T = ( - (T + self.audio_group_size - 1) - // self.audio_group_size - * self.audio_group_size - ) - padded_audio = torch.cat( - [ - audio, - torch.zeros(padded_T - T, self.audio_channels, dtype=torch.long) - + audio[-1, :], - ], - dim=0, - ) - padded_audio = padded_audio.reshape( - padded_T // self.audio_group_size, - self.audio_group_size, - self.audio_channels, - ) - return padded_audio - def _process_videos_parallel(self, contents): video_contents_info = [] for idx, content in enumerate(contents): @@ -1171,29 +953,17 @@ class MiMoProcessor: } def _process_audio_content(self, content, verbose): - processed_audio = self.process_audio(content.content) - if isinstance(processed_audio, tuple): - is_tokenized = False - audio_spec, audio_token_len = processed_audio - audio_input = audio_spec - else: - is_tokenized = True - audio_token_len = processed_audio.shape[0] - audio_input = processed_audio - _input_ids = ( - [self.audio_start_token_id] - + [self.audio_token_id] * audio_token_len - + [self.audio_end_token_id] - ) - + result = self.audio_pipeline.process_audio_input(content.content) verbose_str = "" if verbose: - verbose_str = f"Audio (is_tokenized={is_tokenized}): [ {audio_token_len}*