# coding=utf-8 # Copyright 2025 The OpenBMB Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import logging import math import os import types from collections.abc import Iterator from copy import deepcopy from dataclasses import dataclass from threading import Thread from typing import List from typing import Literal from typing import Optional from typing import Tuple from typing import Union import numpy as np import soundfile as sf import torch import torch.nn as nn import torch.nn.functional as F import torch.nn.utils.parametrize as P from huggingface_hub import hf_hub_download from PIL import Image from torch.nn.utils.parametrizations import weight_norm from tqdm import tqdm from transformers import AutoProcessor from transformers import BertTokenizerFast from transformers import LlamaConfig from transformers import LlamaModel # from transformers import LogitsWarper from transformers import LogitsProcessor from transformers import PreTrainedModel from transformers import Qwen2ForCausalLM from transformers import Qwen2PreTrainedModel from transformers import TextIteratorStreamer from transformers import TopKLogitsWarper from transformers import TopPLogitsWarper from transformers.cache_utils import Cache from transformers.cache_utils import DynamicCache from transformers.cache_utils import EncoderDecoderCache from transformers.cache_utils import StaticCache from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.modeling_outputs import ModelOutput from transformers.models.whisper.modeling_whisper import ACT2FN from transformers.models.whisper.modeling_whisper import WHISPER_ATTENTION_CLASSES from transformers.models.whisper.modeling_whisper import WhisperConfig from transformers.models.whisper.modeling_whisper import WhisperEncoder try: from vector_quantize_pytorch import GroupedResidualFSQ from vocos import Vocos from vocos.pretrained import instantiate_class _tts_deps = True except: _tts_deps = False from .configuration_minicpm import ConditionalChatTTSConfig from .configuration_minicpm import MiniCPMOConfig from .modeling_navit_siglip import SiglipVisionTransformer from .image_processing_minicpmv import MiniCPMOBatchFeature from .resampler import Resampler from .utils import NumberToTextConverter from .utils import sentence_end from .utils import VoiceChecker logger = logging.getLogger(__name__) @dataclass class OmniOutput(ModelOutput): text: Optional[Union[str, List[str], Iterator]] = None spk_embeds: Optional[torch.FloatTensor] = None audio_wav: Optional[np.ndarray] = None sampling_rate: Optional[int] = None class MiniCPMOPreTrainedModel(Qwen2PreTrainedModel): config_class = MiniCPMOConfig class MiniCPMO(MiniCPMOPreTrainedModel): def __init__(self, config): super().__init__(config) self.llm = Qwen2ForCausalLM(config) self.llm.prepare_inputs_for_generation = types.MethodType(prepare_inputs_for_generation, self.llm) # patch llm self.embed_dim = self.llm.config.hidden_size # init vision module if self.config.init_vision: self.vpm = self.init_vision_module() self.vision_dim = self.vpm.embed_dim self.resampler = self.init_resampler(self.embed_dim, self.vision_dim) # init audio module if self.config.init_audio: self.apm = self.init_audio_module() audio_output_dim = int(self.apm.config.encoder_ffn_dim // 4) self.audio_avg_pooler = nn.AvgPool1d(self.config.audio_pool_step, stride=self.config.audio_pool_step) self.audio_projection_layer = MultiModalProjector(in_dim=audio_output_dim, out_dim=self.embed_dim) self.audio_encoder_layer = -1 # init tts module # if self.config.init_tts: # assert _tts_deps, "please make sure vector_quantize_pytorch and vocos are installed." # self.tts = self.init_tts_module() self.processor = AutoProcessor.from_pretrained(self.config._name_or_path, trust_remote_code=True) self.terminators = ["<|im_end|>", "<|endoftext|>"] self.default_tts_chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>' }}{% endif %}" self.force_no_stop = False # for stream api self.reset_session() def reset_session(self): self.session_id = None self.new_user_msg = True self.llm_generated = False self.llm_generate_completed = False self.llm_past_key_values = None self.audio_past_key_values = None # apm kv cache def init_tts( self, tts_text_tokenizer_path=None, vocos_ckpt_path=None, ): """ load tts tokenizer and vocos 1. try load form local 2. try load from huggingface """ from .processing_minicpmo import ChatTTSProcessor if tts_text_tokenizer_path is None: tts_text_tokenizer_path = os.path.join(self.config._name_or_path, "assets/chattts_tokenizer") if not os.path.exists(tts_text_tokenizer_path): # try from hf model_id tts_text_tokenizer_path = "openbmb/chattts_tokenizer" tts_text_tokenizer = BertTokenizerFast.from_pretrained(tts_text_tokenizer_path) self.tts_processor = ChatTTSProcessor(text_tokenizer=tts_text_tokenizer) if vocos_ckpt_path is None: vocos_ckpt_path = os.path.join(self.config._name_or_path, "assets/Vocos.pt") if not os.path.exists(vocos_ckpt_path): vocos_ckpt_path = hf_hub_download(repo_id="openbmb/MiniCPM-o-2_6", subfolder="assets", filename="Vocos.pt") assert os.path.exists(vocos_ckpt_path) self.vocos = self.initialize_vocos(vocos_ckpt_path) def initialize_vocos(self, ckpt_path): feature_extractor = instantiate_class( args=(), init={ "class_path": "vocos.feature_extractors.MelSpectrogramFeatures", "init_args": {"sample_rate": 24000, "n_fft": 1024, "hop_length": 256, "n_mels": 100}, }, ) backbone = instantiate_class( args=(), init={ "class_path": "vocos.models.VocosBackbone", "init_args": {"input_channels": 100, "dim": 512, "intermediate_dim": 1536, "num_layers": 8}, }, ) head = instantiate_class( args=(), init={"class_path": "vocos.heads.ISTFTHead", "init_args": {"dim": 512, "n_fft": 1024, "hop_length": 256}}, ) vocos = Vocos(feature_extractor, backbone, head).to("cuda").eval().to(torch.float32) vocos.load_state_dict(torch.load(ckpt_path, weights_only=True, mmap=True)) return vocos def init_vision_module(self): if self.config._attn_implementation == "flash_attention_2": self.config.vision_config._attn_implementation = "flash_attention_2" else: self.config.vision_config._attn_implementation = "eager" model = SiglipVisionTransformer(self.config.vision_config) if self.config.drop_vision_last_layer: model.encoder.layers = model.encoder.layers[:-1] setattr(model, "embed_dim", model.embeddings.embed_dim) setattr(model, "patch_size", model.embeddings.patch_size) return model def init_resampler(self, embed_dim, vision_dim): return Resampler( num_queries=self.config.query_num, embed_dim=embed_dim, num_heads=embed_dim // 128, kv_dim=vision_dim, adaptive=True, ) def init_audio_module(self): model = MiniCPMWhisperEncoder(self.config.audio_config) return model def init_tts_module(self): model = ConditionalChatTTS(self.config.tts_config) return model def get_input_embeddings(self): return self.llm.get_input_embeddings() def set_input_embeddings(self, value): self.llm.embed_tokens = value def get_output_embeddings(self): return self.llm.lm_head def set_output_embeddings(self, new_embeddings): self.llm.lm_head = new_embeddings def set_decoder(self, decoder): self.llm = decoder def get_decoder(self): return self.llm def subsequent_chunk_mask( self, size: int, chunk_size: int, num_left_chunks: int = -1, device: torch.device = torch.device("cpu"), num_lookhead: int = 0, ) -> torch.Tensor: """Create mask for subsequent steps (size, size) with chunk size, this is for streaming encoder Args: size (int): size of mask chunk_size (int): size of chunk num_left_chunks (int): number of left chunks <0: use full chunk >=0: use num_left_chunks device (torch.device): "cpu" or "cuda" or torch.Tensor.device Returns: torch.Tensor: mask Examples: >>> subsequent_chunk_mask(4, 2) [[1, 1, 0, 0], [1, 1, 0, 0], [1, 1, 1, 1], [1, 1, 1, 1]] """ ret = torch.zeros(size, size, device=device, dtype=torch.bool) for i in range(size): if num_left_chunks < 0: start = 0 else: start = max((i // chunk_size - num_left_chunks) * chunk_size, 0) ending = min((i // chunk_size + 1) * chunk_size + num_lookhead, size) ret[i, start:ending] = True return ret def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): """ Computes the output length of the convolutional layers and the output length of the audio encoder """ input_lengths_after_cnn = (input_lengths - 1) // 2 + 1 input_lengths_after_pooling = ( input_lengths_after_cnn - self.config.audio_pool_step ) // self.config.audio_pool_step + 1 input_lengths_after_pooling = input_lengths_after_pooling.to(dtype=torch.int32) return input_lengths_after_cnn, input_lengths_after_pooling def get_vllm_embedding(self, data): """ Compute all visual embeddings, and set into llm embeddings. Args: data: Dict tgt_sizes: image size after patch embedding pixel_values: image features image_bound: position of each picture corresponding to input_ids input_ids: full input_ids, include placeholder Returns: embedding with vision, vision_hidden_states """ if "vision_hidden_states" not in data: dtype = self.llm.model.embed_tokens.weight.dtype device = self.llm.model.embed_tokens.weight.device tgt_sizes = data["tgt_sizes"] pixel_values_list = data["pixel_values"] vision_hidden_states = [] all_pixel_values = [] img_cnt = [] for pixel_values in pixel_values_list: img_cnt.append(len(pixel_values)) all_pixel_values.extend([i.flatten(end_dim=1).permute(1, 0) for i in pixel_values]) # exist image if all_pixel_values: tgt_sizes = [tgt_size for tgt_size in tgt_sizes if isinstance(tgt_size, torch.Tensor)] tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32) max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1]) all_pixel_values = torch.nn.utils.rnn.pad_sequence( all_pixel_values, batch_first=True, padding_value=0.0 ) B, L, _ = all_pixel_values.shape all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L) patch_attn_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool, device=device) for i in range(B): patch_attn_mask[i, 0, : tgt_sizes[i][0] * tgt_sizes[i][1]] = True vision_batch_size = self.config.vision_batch_size all_pixel_values = all_pixel_values.type(dtype) if B > vision_batch_size: hs = [] for i in range(0, B, vision_batch_size): start_idx = i end_idx = i + vision_batch_size tmp_hs = self.vpm( all_pixel_values[start_idx:end_idx], patch_attention_mask=patch_attn_mask[start_idx:end_idx], tgt_sizes=tgt_sizes[start_idx:end_idx], ).last_hidden_state hs.append(tmp_hs) vision_embedding = torch.cat(hs, dim=0) else: vision_embedding = self.vpm( all_pixel_values, patch_attention_mask=patch_attn_mask, tgt_sizes=tgt_sizes ).last_hidden_state vision_embedding = self.resampler(vision_embedding, tgt_sizes) start = 0 for pixel_values in pixel_values_list: img_cnt = len(pixel_values) if img_cnt > 0: vision_hidden_states.append(vision_embedding[start : start + img_cnt]) start += img_cnt else: vision_hidden_states.append([]) else: # no image if self.training: dummy_image = torch.zeros((1, 3, 224, 224), device=device, dtype=dtype) tgt_sizes = torch.Tensor( [[(224 // self.config.patch_size), math.ceil(224 / self.config.patch_size)]] ).type(torch.int32) dummy_feature = self.resampler(self.vpm(dummy_image).last_hidden_state, tgt_sizes) else: dummy_feature = [] for _ in range(len(pixel_values_list)): vision_hidden_states.append(dummy_feature) else: vision_hidden_states = data["vision_hidden_states"] if hasattr(self.llm.config, "scale_emb"): vllm_embedding = self.llm.model.embed_tokens(data["input_ids"]) * self.llm.config.scale_emb else: vllm_embedding = self.llm.model.embed_tokens(data["input_ids"]) new_vllm_embedding = vllm_embedding.clone() vision_hidden_states = [ i.type(vllm_embedding.dtype) if isinstance(i, torch.Tensor) else i for i in vision_hidden_states ] bs = len(data["input_ids"]) for i in range(bs): cur_vs_hs = vision_hidden_states[i] if len(cur_vs_hs) > 0: cur_vllm_emb = vllm_embedding[i] cur_image_bound = data["image_bound"][i] if len(cur_image_bound) > 0: image_indices = torch.stack( [torch.arange(r[0], r[1], dtype=torch.long) for r in cur_image_bound] ).to(vllm_embedding.device) new_vllm_embedding[i] = cur_vllm_emb.scatter( 0, image_indices.view(-1, 1).repeat(1, cur_vllm_emb.shape[-1]), cur_vs_hs.view(-1, cur_vs_hs.shape[-1]), ) elif self.training: new_vllm_embedding[i] += cur_vs_hs[0].mean() * 0 return new_vllm_embedding, vision_hidden_states def get_audio_embedding_streaming(self, data): r""" Extract audio embeddings in a streaming manner using cached key-value pairs. This method processes incoming audio features incrementally and stores/updates `past_key_values` for faster inference on subsequent audio frames. It only supports batch_size=1 and is intended for streaming scenarios. Args: data (dict): - **"audio_features"** (`torch.FloatTensor`): Input mel-spectrograms of shape `(batch_size, 80, frames)`. - **"audio_feature_lens"** (List[List[int]]): Lengths of each audio segment for each item in the batch. Returns: List[List[torch.Tensor]]: audio embeddings """ wavforms = data.get("audio_features", []) # (bs, 80, frames) or [], multi audios need filled in advance audio_feature_lens_raw = data.get("audio_feature_lens", []) # list, [[x1, x2], [y1], [z1]] # exist audio if len(wavforms) > 0: audio_feature_lens = torch.hstack(audio_feature_lens_raw) batch_size, _, max_mel_seq_len = wavforms.shape assert batch_size == 1 max_seq_len = (max_mel_seq_len - 1) // 2 + 1 if self.audio_past_key_values is not None: cache_length = self.audio_past_key_values[0][0].shape[2] apm_max_len = self.apm.embed_positions.weight.shape[0] if cache_length + max_seq_len >= apm_max_len: logger.warning( f"audio_past_key_values length {cache_length + max_seq_len} exceed {apm_max_len}, reset." ) self.audio_past_key_values = None audio_outputs = self.apm(wavforms, past_key_values=self.audio_past_key_values, use_cache=True) audio_states = audio_outputs.last_hidden_state # [:, :audio_feat_lengths, :] self.audio_past_key_values = audio_outputs.past_key_values audio_embeds = self.audio_projection_layer(audio_states) audio_embeds = audio_embeds.transpose(1, 2) audio_embeds = self.audio_avg_pooler(audio_embeds) audio_embeds = audio_embeds.transpose(1, 2) _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens) num_audio_tokens = feature_lens_after_pooling final_audio_embeds = [] idx = 0 for i in range(len(audio_feature_lens_raw)): target_audio_embeds = [] for _ in range(len(audio_feature_lens_raw[i])): target_audio_embeds.append(audio_embeds[idx, : num_audio_tokens[idx], :]) idx += 1 final_audio_embeds.append(target_audio_embeds) return final_audio_embeds else: return [] def get_audio_embedding(self, data, chunk_length=-1, dummy=True): r""" Extract full audio embeddings with optional chunk-based attention. This method computes embeddings for all audio frames at once, either using full attention (when `chunk_length` is -1) or chunk-based attention (when `chunk_length` is a positive number). It does not use key-value caching and is suitable for non-streaming inference. Args: data (dict): - **"audio_features"** (`torch.FloatTensor`): Input mel-spectrograms of shape `(batch_size, 80, frames)`. - **"audio_feature_lens"** (List[List[int]]): Lengths of each audio segment for each item in the batch. chunk_length (int, optional): Determines whether to use full attention (-1) or chunk-based attention (>0) during embedding computation. Returns: List[List[torch.Tensor]]: audio embeddings """ wavforms = data.get("audio_features", []) # (bs, 80, frames) or [], multi audios need filled in advance audio_feature_lens_raw = data.get("audio_feature_lens", []) # list, [[x1, x2], [y1], [z1]] # exist audio if len(wavforms) > 0: audio_feature_lens = torch.hstack(audio_feature_lens_raw) batch_size, _, max_mel_seq_len = wavforms.shape max_seq_len = (max_mel_seq_len - 1) // 2 + 1 # Create a sequence tensor of shape (batch_size, max_seq_len) seq_range = ( torch.arange(0, max_seq_len, dtype=audio_feature_lens.dtype, device=audio_feature_lens.device) .unsqueeze(0) .expand(batch_size, max_seq_len) ) lengths_expand = audio_feature_lens.unsqueeze(1).expand(batch_size, max_seq_len) # Create mask padding_mask = seq_range >= lengths_expand # 1 for padded values audio_attention_mask_ = padding_mask.view(batch_size, 1, 1, max_seq_len).expand( batch_size, 1, max_seq_len, max_seq_len ) audio_attention_mask = audio_attention_mask_.to( dtype=self.apm.conv1.weight.dtype, device=self.apm.conv1.weight.device ) if chunk_length > 0: chunk_num_frame = int(chunk_length * 50) chunk_mask = self.subsequent_chunk_mask( size=max_seq_len, chunk_size=chunk_num_frame, num_left_chunks=-1, device=audio_attention_mask_.device, ) audio_attention_mask_ = torch.logical_or(audio_attention_mask_, torch.logical_not(chunk_mask)) audio_attention_mask[audio_attention_mask_] = float("-inf") audio_states = self.apm( wavforms, output_hidden_states=True, attention_mask=audio_attention_mask ).hidden_states[self.audio_encoder_layer] audio_embeds = self.audio_projection_layer(audio_states) audio_embeds = audio_embeds.transpose(1, 2) audio_embeds = self.audio_avg_pooler(audio_embeds) audio_embeds = audio_embeds.transpose(1, 2) _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens) num_audio_tokens = feature_lens_after_pooling final_audio_embeds = [] idx = 0 for i in range(len(audio_feature_lens_raw)): target_audio_embeds = [] for _ in range(len(audio_feature_lens_raw[i])): target_audio_embeds.append(audio_embeds[idx, : num_audio_tokens[idx], :]) idx += 1 final_audio_embeds.append(target_audio_embeds) return final_audio_embeds elif self.training and dummy: dtype = self.apm.embed_positions.weight.dtype device = self.apm.embed_positions.weight.device dummy_wavs = torch.zeros((1, 80, 100), device=device, dtype=dtype) audio_states = self.apm(dummy_wavs, output_hidden_states=True).hidden_states[self.audio_encoder_layer] audio_embeds = self.audio_projection_layer(audio_states) audio_embeds = audio_embeds.transpose(1, 2) audio_embeds = self.audio_avg_pooler(audio_embeds) audio_embeds = audio_embeds.transpose(1, 2) return [audio_embeds] else: return [] def get_omni_embedding(self, data, input_embeddings, chunk_length=-1, stream_input=False): """ Args: data: input_embeddings: chunk_length: whisper use full attention or chunk attention stream_input: use streaming audio embedding Returns: final embeddings with audio feature """ if stream_input: audio_embeddings = self.get_audio_embedding_streaming(data) else: audio_embeddings = self.get_audio_embedding(data, chunk_length) bs = len(input_embeddings) if len(data.get("audio_features", [])) > 0: assert len(audio_embeddings) == len(input_embeddings) if len(audio_embeddings) > 0: audio_bounds = data["audio_bounds"] if self.config.chunk_input: for i in range(bs): if not audio_embeddings[i]: continue audio_embs = torch.cat(audio_embeddings[i], dim=0).to( device=input_embeddings.device, dtype=input_embeddings.dtype ) audio_start_pos = 0 for bound in audio_bounds[i]: audio_len = bound[1] - bound[0] input_embeddings[i, bound[0] : bound[1]] = audio_embs[ audio_start_pos : audio_start_pos + audio_len, : ] audio_start_pos += audio_len else: for i in range(bs): audio_embs = audio_embeddings[i] bounds = audio_bounds[i] for embs, bound in zip(audio_embs, bounds): audio_indices = torch.arange(bound[0], bound[1], dtype=torch.long).to( input_embeddings.device ) if embs.shape[0] != len(audio_indices): raise ValueError( f"Shape mismatch: Trying to assign embeddings of shape {embs.shape} " f"to input indices of length {len(audio_indices)}" ) input_embeddings[i, audio_indices] = embs.to(input_embeddings.dtype) elif self.training: for i in range(bs): # dummy audio_embeddings input_embeddings = input_embeddings + audio_embeddings[0].mean() * 0 return input_embeddings def forward(self, data, **kwargs): vllm_embedding, vision_hidden_states = self.get_vllm_embedding(data) if self.config.init_audio: vllm_embedding = self.get_omni_embedding( data, input_embeddings=vllm_embedding, chunk_length=self.config.audio_chunk_length ) position_ids = data["position_ids"] if position_ids.dtype != torch.int64: position_ids = position_ids.long() # compatible with llama factory for key in ["input_ids", "inputs_embeds", "position_ids"]: if key in kwargs: del kwargs[key] return self.llm(input_ids=None, position_ids=position_ids, inputs_embeds=vllm_embedding, **kwargs) def _decode(self, inputs_embeds, tokenizer, attention_mask, **kwargs): kwargs.pop("output_hidden_states", None) kwargs.pop("return_dict_in_generate", None) terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] outputs = self.llm.generate( inputs_embeds=inputs_embeds, pad_token_id=0, eos_token_id=terminators, attention_mask=attention_mask, output_hidden_states=True, return_dict_in_generate=True, **kwargs, ) return outputs def _decode_stream(self, inputs_embeds, tokenizer, **kwargs): terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] streamer = TextIteratorStreamer(tokenizer=tokenizer) generation_kwargs = { "inputs_embeds": inputs_embeds, "pad_token_id": 0, "eos_token_id": terminators, "streamer": streamer, } generation_kwargs.update(kwargs) thread = Thread(target=self.llm.generate, kwargs=generation_kwargs) thread.start() return streamer def _decode_text(self, result_ids, tokenizer): terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] result_text = [] for result in result_ids: result = result[result != 0] if result[0] == tokenizer.bos_id: result = result[1:] if result[-1] in terminators: result = result[:-1] result_text.append(tokenizer.decode(result)) return result_text def get_sys_prompt(self, ref_audio=None, mode="default", language="zh"): """ Choose different system prompts according to different tasks Args: ref_audio: if ref_audio is not None, will use the voice cloning prompts, and the voice generated by the model will refer to the timbre of ref audio mode: "default": default system prompt and not refer to any task "omni": input video and audio simultaneously "audio_assistant": Default voice-only mode, the model will use the ref_audio's voice to reply user's question as a helpful assistant. "audio_roleplay": Roleplay voice-only mode, the model will use the ref_audio's voice to reply, and also role-play the character based on the audio prompt. "voice_cloning": TTS mode, the model will clone the voice of ref_audio. language: prompts language, the model has the ability to automatically select the response language based on the question language Returns: """ if ref_audio is not None: assert isinstance(ref_audio, np.ndarray), "ref_audio error" if mode == "omni": if language == "zh": sys_prompt = "你是一个AI助手。你能接受视频，音频和文本输入并输出语音和文本。" vc_prompt_prefix = sys_prompt + "模仿输入音频中的声音特征。" vc_prompt_suffix = "作为助手，你将使用这种声音风格说话。" else: sys_prompt = "You are a helpful assistant. You can accept video, audio and text input and output voice and text. " vc_prompt_prefix = sys_prompt + "Clone the voice in the provided audio prompt." vc_prompt_suffix = "As an assistant, you will speak using this voice style." if ref_audio is not None: sys_msgs = {"role": "user", "content": [vc_prompt_prefix, ref_audio, vc_prompt_suffix]} else: sys_msgs = {"role": "user", "content": [sys_prompt]} return sys_msgs elif mode == "audio_assistant": if language == "zh": vc_prompt_prefix = "模仿输入音频中的声音特征。" vc_prompt_suffix = "作为助手，你将使用这种声音风格说话。" else: vc_prompt_prefix = "Use the voice in the audio prompt to synthesize new content." vc_prompt_suffix = "You are a helpful assistant with the above voice style." if ref_audio is not None: sys_msgs = {"role": "user", "content": [vc_prompt_prefix, ref_audio, vc_prompt_suffix]} else: logger.warning( "Warning: ref_audio is None, speech generation will be performed based on the default voice." ) sys_msgs = {"role": "user", "content": ["Use the voice.", vc_prompt_suffix]} return sys_msgs elif mode == "audio_roleplay": if language == "zh": vc_prompt_prefix = "模仿输入音频中的声音特征。" vc_prompt_suffix = "假装你是上述音频中的人物，与我进行对话。" else: vc_prompt_prefix = "Clone the voice in the provided audio prompt." vc_prompt_suffix = "Try to role-play the character based on the audio prompt above." if ref_audio is not None: sys_msgs = {"role": "user", "content": [vc_prompt_prefix, ref_audio, vc_prompt_suffix]} else: print("Warning: ref_audio is None, speech generation will be performed based on the default voice.") sys_msgs = {"role": "user", "content": ["Use the voice.", vc_prompt_suffix]} return sys_msgs elif mode == "voice_cloning": if language == "zh": vc_prompt_prefix = "模仿输入音频中的声音特征。" else: vc_prompt_prefix = "Clone the voice in the provided audio prompt." if ref_audio is not None: sys_msgs = {"role": "user", "content": [vc_prompt_prefix, ref_audio]} else: raise ValueError("ref_audio con't be None in voice_cloning mode.") return sys_msgs else: sys_prompt = "You are a helpful assistant. You can accept audio and text input and output voice and text." sys_msgs = {"role": "user", "content": [sys_prompt]} return sys_msgs def generate( self, input_ids=None, pixel_values=None, tgt_sizes=None, audio_features=[], audio_feature_lens=None, image_bound=None, audio_bounds=None, spk_bounds=None, attention_mask=None, tokenizer=None, vision_hidden_states=None, stream=False, decode_text=True, **kwargs, ): assert input_ids is not None assert len(input_ids) == len(pixel_values) model_inputs = { "input_ids": input_ids, "audio_features": audio_features, "audio_feature_lens": audio_feature_lens, "image_bound": image_bound, "audio_bounds": audio_bounds, "spk_bounds": spk_bounds, } if vision_hidden_states is None: model_inputs["pixel_values"] = pixel_values model_inputs["tgt_sizes"] = tgt_sizes else: model_inputs["vision_hidden_states"] = vision_hidden_states model_output = {} with torch.inference_mode(): model_inputs["inputs_embeds"], vision_hidden_states = self.get_vllm_embedding(model_inputs) model_inputs["inputs_embeds"] = self.get_omni_embedding( model_inputs, input_embeddings=model_inputs["inputs_embeds"], chunk_length=self.config.audio_chunk_length, ) if stream: result = self._decode_stream(model_inputs["inputs_embeds"], tokenizer, **kwargs) # if stream return TextIteratorStreamer and output is empty outputs = {} else: outputs = self._decode(model_inputs["inputs_embeds"], tokenizer, attention_mask, **kwargs) #怎么每次要调用config result = self._decode_text(outputs.sequences, tokenizer) if decode_text is False: return outputs return result, outputs def chat( self, image=None, msgs=None, tokenizer=None, processor=None, vision_hidden_states=None, max_new_tokens=2048, min_new_tokens=0, sampling=True, max_inp_length=32768, stream=False, chunk_input=True, omni_input=False, max_slice_nums=None, use_image_id=None, use_tts_template=False, generate_audio=False, return_spk_embed=False, return_dict=False, output_audio_path=None, **kwargs, ): """ Unified chat function Args: image: use for batch_size=1 vqa, It is not recommended to continue to use this parameter msgs: the input chat msgs, support text: (string) / image: (PIL.Image) / audio (numpy.ndarray) tokenizer: tokenizer for llm processor: if None, use the default processor max_new_tokens: the maximum length of the generation min_new_tokens: the minimum length of the generation sampling: whether to use sampling decoding or beam search decoding max_inp_length: the maximum length of input stream: whether to return generator, only used when tts is not required chunk_input: whether to split audio into 1s chunks omni_input: determine whether it is omni mode max_slice_nums: control the maximum number of image slices use_image_id: for video understanding or omni understanding, use_image_id should be False use_tts_template: if the msgs contain audio, use_tts_template should be True generate_audio: whether to generate audio output, only used when return_dict=True return_spk_embed: whether to return spk embedding, only used when return_dict=True return_dict: whether to return dict output_audio_path: audio save path when generate_audio **kwargs: """ if isinstance(msgs[0], list): batched = True else: batched = False if generate_audio or return_spk_embed: return_dict = True msgs_list = msgs images_list = image if batched is False: images_list, msgs_list = [images_list], [msgs_list] else: assert images_list is None, "Please integrate image to msgs when using batch inference." images_list = [None] * len(msgs_list) assert len(images_list) == len(msgs_list), "The batch dim of images_list and msgs_list should be the same." if processor is None: if self.processor is None: self.processor = AutoProcessor.from_pretrained(self.config._name_or_path, trust_remote_code=True) processor = self.processor assert ( self.config.query_num == processor.image_processor.image_feature_size ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.patch_size == processor.image_processor.patch_size ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.use_image_id == processor.image_processor.use_image_id ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.slice_config.max_slice_nums == processor.image_processor.max_slice_nums ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.slice_mode == processor.image_processor.slice_mode ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." prompts_lists = [] input_images_list = [] input_audios_list = [] audio_parts_list = [] for image, msgs in zip(images_list, msgs_list): if isinstance(msgs, str): msgs = json.loads(msgs) copy_msgs = deepcopy(msgs) assert len(msgs) > 0, "msgs is empty" assert sampling or not stream, "if use stream mode, make sure sampling=True" if image is not None and isinstance(copy_msgs[0]["content"], str): copy_msgs[0]["content"] = [image, copy_msgs[0]["content"]] images = [] audios = [] audio_parts = [] for i, msg in enumerate(copy_msgs): role = msg["role"] content = msg["content"] assert role in ["system", "user", "assistant"] if i == 0: assert role in ["user", "system"], "The role of first msg should be user" if isinstance(content, str): content = [content] cur_msgs = [] for c in content: if isinstance(c, Image.Image): images.append(c) cur_msgs.append("(./)") elif isinstance(c, np.ndarray): # audio audios.append(c) audio_parts.append(i) cur_msgs.append("()") use_tts_template = True elif isinstance(c, str): cur_msgs.append(c) if omni_input: msg["content"] = "".join(cur_msgs) else: msg["content"] = "\n".join(cur_msgs) prompts_lists.append( processor.tokenizer.apply_chat_template( copy_msgs, tokenize=False, add_generation_prompt=True, chat_template=self.default_tts_chat_template if use_tts_template else None, ) ) input_images_list.append(images) input_audios_list.append(audios) audio_parts_list.append(audio_parts) inputs = processor( prompts_lists, input_images_list, input_audios_list, audio_parts_list, max_slice_nums=max_slice_nums, use_image_id=use_image_id, chunk_input=chunk_input, return_tensors="pt", max_length=max_inp_length, ).to(self.device) if sampling: generation_config = { "top_p": 0.8, "top_k": 100, "temperature": 0.7, "do_sample": True, "repetition_penalty": 1.05, } else: generation_config = { "num_beams": 3, "repetition_penalty": 1.2, } if min_new_tokens > 0: generation_config["min_new_tokens"] = min_new_tokens generation_config.update((k, kwargs[k]) for k in generation_config.keys() & kwargs.keys()) inputs.pop("image_sizes") with torch.inference_mode(): res, outputs = self.generate( **inputs, tokenizer=tokenizer, max_new_tokens=max_new_tokens, vision_hidden_states=vision_hidden_states, stream=stream, **generation_config, ) if stream: def stream_gen(): for text in res: for term in self.terminators: text = text.replace(term, "") yield text if return_dict: return OmniOutput(text=stream_gen()) else: return stream_gen() else: spk_embeds = wav_numpy = sr = None if batched: answer = res else: answer = res[0] if use_tts_template and generate_audio: mel_spec = self._generate_mel_spec(inputs, outputs, answer) wav_numpy, sr = self.decode_mel_to_audio(mel_spec, output_audio_path) if return_spk_embed: spk_embeds = self._get_last_spk_embeds(inputs, outputs) if isinstance(answer, list): answer = [i.replace(tokenizer.tts_end, "") for i in answer] else: answer = answer.replace(tokenizer.tts_end, "") if return_dict: return OmniOutput(text=answer, spk_embeds=spk_embeds, audio_wav=wav_numpy, sampling_rate=sr) else: return answer def _decode_hidden(self, result_ids, last_hidden_states, tokenizer): terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] #self.terminators=['<|im_end|>', '<|endoftext|>'] hidden_states = torch.concat([h[:,-1:] for h in last_hidden_states],dim=1) hidden_states_unpad = [] result_text_unpad = [] text_token_len = [] for id, result in enumerate(result_ids): hidden_states_i = hidden_states[id, result != 0, :] result = result[result!=0] if result[0] == tokenizer.bos_id: result = result[1:] hidden_states_i = hidden_states_i[1:] if result[-1] in terminators: result = result[:-1] hidden_states_i = hidden_states_i[:-1] if result[-1] == 151692: #'<|tts_eos|>' result = result[:-1] hidden_states_i = hidden_states_i[:-1] result_text_unpad.append(tokenizer.decode(result)) hidden_states_unpad.append(hidden_states_i) text_token_len.append(len(result)) return text_token_len, hidden_states, hidden_states_unpad, result_text_unpad def get_hidden( self, image=None, msgs=None, tokenizer=None, processor=None, vision_hidden_states=None, max_new_tokens=2048, min_new_tokens=0, sampling=True, max_inp_length=32768, stream=False, chunk_input=True, omni_input=False, max_slice_nums=None, use_image_id=None, use_tts_template=False, generate_audio=False, return_spk_embed=False, **kwargs, ): if isinstance(msgs[0], list): batched = True else: batched = False if generate_audio or return_spk_embed: return_dict = True msgs_list = msgs images_list = image if batched is False: images_list, msgs_list = [images_list], [msgs_list] else: assert images_list is None, "Please integrate image to msgs when using batch inference." images_list = [None] * len(msgs_list) assert len(images_list) == len(msgs_list), "The batch dim of images_list and msgs_list should be the same." if processor is None: if self.processor is None: self.processor = AutoProcessor.from_pretrained(self.config._name_or_path, trust_remote_code=True) processor = self.processor assert ( self.config.query_num == processor.image_processor.image_feature_size ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.patch_size == processor.image_processor.patch_size ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.use_image_id == processor.image_processor.use_image_id ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.slice_config.max_slice_nums == processor.image_processor.max_slice_nums ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." assert ( self.config.slice_mode == processor.image_processor.slice_mode ), "These two values should be the same. Check `config.json` and `preprocessor_config.json`." prompts_lists = [] input_images_list = [] input_audios_list = [] audio_parts_list = [] for image, msgs in zip(images_list, msgs_list): if isinstance(msgs, str): msgs = json.loads(msgs) copy_msgs = deepcopy(msgs) assert len(msgs) > 0, "msgs is empty" assert sampling or not stream, "if use stream mode, make sure sampling=True" # if image is not None and isinstance(copy_msgs[0]["content"], str): # copy_msgs[0]["content"] = [image, copy_msgs[0]["content"]] images = [] audios = [] audio_parts = [] for i, msg in enumerate(copy_msgs): role = msg["role"] content = msg["content"] assert role in ["system", "user", "assistant"] if i == 0: assert role in ["user", "system"], "The role of first msg should be user" if isinstance(content, str): content = [content] cur_msgs = [] for c in content: if isinstance(c, Image.Image): images.append(c) cur_msgs.append("(./)") elif isinstance(c, np.ndarray): # audio audios.append(c) audio_parts.append(i) cur_msgs.append("()") use_tts_template = True elif isinstance(c, str): cur_msgs.append(c) if omni_input: msg["content"] = "".join(cur_msgs) else: msg["content"] = "\n".join(cur_msgs) prompts_lists.append( processor.tokenizer.apply_chat_template( copy_msgs, tokenize=False, add_generation_prompt=True, chat_template=self.default_tts_chat_template if use_tts_template else None, ) ) input_images_list.append(images) input_audios_list.append(audios) audio_parts_list.append(audio_parts) inputs = processor( prompts_lists, input_images_list, input_audios_list, audio_parts_list, max_slice_nums=max_slice_nums, use_image_id=use_image_id, chunk_input=chunk_input, return_tensors="pt", max_length=max_inp_length, ).to(self.device) if sampling: generation_config = { "top_p": 0.8, "top_k": 100, "temperature": 0.7, "do_sample": True, "repetition_penalty": 1.05, } else: generation_config = { "num_beams": 3, "repetition_penalty": 1.2, } if min_new_tokens > 0: generation_config["min_new_tokens"] = min_new_tokens generation_config.update((k, kwargs[k]) for k in generation_config.keys() & kwargs.keys()) inputs.pop("image_sizes") # with torch.inference_mode(): with torch.no_grad(): res, outputs = self.generate( **inputs, tokenizer=tokenizer, max_new_tokens=max_new_tokens, vision_hidden_states=vision_hidden_states, stream=stream, **generation_config, ) last_hidden_states = [hs[-1] for hs in outputs.hidden_states] text_token = deepcopy(outputs.sequences) text_token_len, hidden_states, hidden_states_unpad, text_unpad = self._decode_hidden(text_token, last_hidden_states, tokenizer) for id in range(len(text_token)): len_ = text_token_len[id] text_token[id, len_:] = 0 hidden_states[id, len_:] = 0 max_len = max(text_token_len) text_token = text_token[:, :max_len] hidden_states = hidden_states[:, :max_len] return text_unpad, text_token, torch.Tensor(text_token_len).to(torch.int32), hidden_states def get_hidden_forward(self,data,**kwargs,): vllm_embedding, vision_hidden_states = self.get_vllm_embedding(data) if self.config.init_audio: vllm_embedding = self.get_omni_embedding( data, input_embeddings=vllm_embedding, chunk_length=self.config.audio_chunk_length ) position_ids = data["position_ids"] if position_ids.dtype != torch.int64: position_ids = position_ids.long() # compatible with llama factory for key in ["input_ids", "inputs_embeds", "position_ids"]: if key in kwargs: del kwargs[key] outputs = self.llm(input_ids=None, position_ids=position_ids, inputs_embeds=vllm_embedding, output_hidden_states=True, **kwargs) ##计算损失 loss_fct = nn.CrossEntropyLoss() logits = outputs.logits.view(-1,self.config.vocab_size).contiguous() labels = data['target'].view(-1).long().contiguous() # Enable model parallelism labels = labels.to(logits.device) loss = loss_fct(logits, labels) ##得到隐藏层特征（根据对话拆分多轮） last_hidden_states = outputs.hidden_states[-1] #(batch_size, s, 3584) batch_size = last_hidden_states.shape[0] new_hidden_states = [] text_token = [] text_token_len = [] for batch_id in range(batch_size): st_id = -1 end_id = -1 for id in range(len(data['target'][batch_id])): if data['target'][batch_id][id] != -100 and data['target'][batch_id][id] != 151645 and st_id==-1: st_id = id+1 #+1是因为target[0]='\n',要去掉 if data['target'][batch_id][id] == 151645 and st_id!=-1: #tokenizer.eos_id end_id = id new_hidden_states.append(last_hidden_states[batch_id:batch_id+1,st_id:end_id]) text_token.append(data['target'][batch_id:batch_id+1,st_id:end_id]) text_token_len.append(end_id-st_id) st_id = -1 ##根据filter过滤不满足要求的answer # assert sum([len(filter_i) for filter_i in data['filter']]) == len(text_token), f"filter data error! filter:{data['filter']}, {data['target']},{len(text_token)}" filter_bool = [filter_i for batch_filter_i in data['filter'] for filter_i in batch_filter_i] if sum(filter_bool) == 0: #没有满足条件的文本可用于训练tts return None, None, None,loss new_hidden_states = [new_hidden_states[i] for i in range(len(filter_bool)) if filter_bool[i]] text_token = [text_token[i] for i in range(len(filter_bool)) if filter_bool[i]] text_token_len = [text_token_len[i] for i in range(len(filter_bool)) if filter_bool[i]] max_len = np.max(text_token_len) ##padding for id, new_hidden_state in enumerate(new_hidden_states): # new_hidden_state (1,s,3584) pad_num = max_len-new_hidden_state.shape[1] if pad_num==0: continue new_hidden_states[id] = torch.cat( [ new_hidden_state, torch.zeros((1, pad_num, new_hidden_state.shape[-1]), device=new_hidden_state.device), ], dim=1, ) #(1,max_len,3584) text_token[id] = torch.cat([text_token[id],torch.zeros((1, pad_num), dtype=text_token[id].dtype, device=text_token[id].device)],dim=1,) #(1,max_len) new_hidden_states = torch.cat(new_hidden_states, dim=0) #(batch_size,max_len,3584) text_token = torch.cat(text_token, dim=0) #(batch_size,max_len) text_token_len = torch.tensor(text_token_len, dtype=torch.int32, device=text_token.device) ##################debug############################ # from transformers import AutoTokenizer # tokenizer_path = '/mnt/afs/zhoufangru/agent/end2end/pretrained_models/MiniCPM-o-2_6' # tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) # output_ids = torch.argmax(outputs.logits, dim=-1) # batch_id = 0 # tokenizer.decode(output_ids[batch_id][data['target'][batch_id]!=-100])[1:] ##################debug############################ return text_token, text_token_len, new_hidden_states,loss @torch.inference_mode() def streaming_prefill( self, session_id, msgs, tokenizer, omni_input=True, max_slice_nums=None, ls_temperature=1.0, **kwargs, ): """ Streaming video/audio input and output audio stream, Only support batch_size=1 Args: session_id: Note: new connection should use a new session_id """ assert session_id is not None if self.session_id is None or session_id != self.session_id: # new session self.is_first = True else: self.is_first = False images = [] audios = [] assert len(msgs) == 1 copy_msgs = deepcopy(msgs) msg = copy_msgs[0] assert msg["role"] in ["system", "user", "assistant"] content = msg["content"] cur_msgs = [] for j, c in enumerate(content): if isinstance(c, Image.Image): images.append(c) cur_msgs.append("(./)") elif isinstance(c, np.ndarray): # audio audios.append(c) cur_msgs.append("()") elif isinstance(c, str): cur_msgs.append(c) else: logger.error("Invalid content type:", c) cur_contents = "".join(cur_msgs) if omni_input else "\n".join(cur_msgs) if not self.is_first and self.new_user_msg and msg["role"] == "user": # new user add im_start if self.llm_generated: if self.llm_generate_completed: msg["content"] = "<|im_end|>\n<|im_start|>user\n" + cur_contents else: # break llm gen, add tts_eos msg["content"] = "<|tts_eos|><|im_end|>\n<|im_start|>user\n" + cur_contents else: msg["content"] = "<|im_start|>user\n" + cur_contents self.new_user_msg = False else: msg["content"] = cur_contents if msg["role"] in ["system", "assistant"]: self.new_user_msg = True self.audio_past_key_values = None # apm kv cache if self.is_first: # init pask_key_values logger.info(f"new session_id: {session_id}, reset kv cache") self.reset_session() self.session_id = session_id prompt = tokenizer.apply_chat_template( copy_msgs, tokenize=False, add_generation_prompt=False, chat_template=self.default_tts_chat_template ) add_special_tokens = True # add bos else: prompt = copy_msgs[0]["content"] add_special_tokens = False model_inputs = self.processor( [prompt], [images], [audios], max_slice_nums=1 if max_slice_nums is None else max_slice_nums, use_image_id=False, chunk_input=True, return_tensors="pt", max_length=None, sampling_rate=16000, add_special_tokens=add_special_tokens, ).to(self.device) # 1. prepare input embeddings model_inputs["inputs_embeds"], _ = self.get_vllm_embedding(model_inputs) # get audio embedding with audio_past_key_values inputs_embeds = self.get_omni_embedding( model_inputs, input_embeddings=model_inputs["inputs_embeds"], stream_input=True ) if self.is_first: # clean audio_past_key_values after first prefill self.audio_past_key_values = None if self.llm_past_key_values is not None: cache_length = self.llm_past_key_values[0][0].shape[2] else: cache_length = 0 attention_mask = torch.ones((1, cache_length + inputs_embeds.shape[1]), dtype=torch.bool, device=self.device) # 2. do prefill and predict listen/speak label outputs = self.llm( past_key_values=self.llm_past_key_values, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=None, # position_ids, use_cache=True, return_dict=True, ) self.llm_past_key_values = outputs["past_key_values"] return @torch.inference_mode() def streaming_generate( self, session_id, tokenizer, max_new_tokens=512, min_new_tokens=0, sampling=True, enable_regenerate=False, **kwargs, ): """ Streaming video/audio input and output audio stream Args: """ if sampling: generation_config = { "top_p": 0.8, "top_k": 100, "temperature": 0.7, "do_sample": True, "repetition_penalty": 1.05, } else: generation_config = { "num_beams": 3, "repetition_penalty": 1.2, } generation_config["min_new_tokens"] = min_new_tokens generation_config.update((k, kwargs[k]) for k in generation_config.keys() & kwargs.keys()) # do generate # reset buffer self.new_user_msg = True self.llm_generated = True self.llm_generate_completed = False self.audio_past_key_values = None # apm kv cache terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] generate_prompt = "<|im_end|>\n<|im_start|>assistant\n<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>" input_ids = tokenizer(generate_prompt, return_tensors="pt", add_special_tokens=False)["input_ids"].cuda() spk_start_idx = torch.where(input_ids[0] == tokenizer.spk_start_id)[0] spk_end_idx = torch.where(input_ids[0] == tokenizer.spk_end_id)[0] spk_bounds = [ torch.hstack([(spk_start_idx + 1).unsqueeze(-1), spk_end_idx.unsqueeze(-1)]) ] # List[Tensor], (1,2) cache_length = past_length = self.llm_past_key_values[0][0].shape[2] attention_mask = torch.ones((1, cache_length + input_ids.shape[1]), dtype=torch.bool, device=self.device) generation_config["max_new_tokens"] = max_new_tokens streamer = self.llm_generate_chunk(input_ids, attention_mask, tokenizer, terminators, generation_config) return streamer def llm_generate_chunk(self, input_ids, attention_mask, tokenizer, terminators, generation_config): def check_uncompleted_token(ids): cur_text = tokenizer.decode(ids) end = len(ids) while cur_text[-1] == "�": end -= 1 if end == 0: break cur_text = tokenizer.decode(ids[:end]) return end max_new_tokens = int(generation_config.pop("max_new_tokens", 2048)) new_len = 0 eos = False left_ids = None while True: outputs = self.llm.generate( input_ids=input_ids, past_key_values=self.llm_past_key_values, attention_mask=attention_mask, use_cache=True, max_new_tokens=3, # reduce first token delay pad_token_id=0, output_hidden_states=True, return_dict_in_generate=True, eos_token_id=terminators, **generation_config, ) if outputs.sequences[0, -1] in terminators: eos = True input_len = input_ids.shape[1] cur_ids = outputs.sequences[:, input_len:] #(batch_size,max_new_tokens) cur_hidden_states = torch.concat([hidden_states[-1][:, -1:] for hidden_states in outputs.hidden_states],dim=1) #(batch_size, max_new_tokens, 3584) new_len += cur_ids.shape[1] if left_ids is not None and left_ids.shape[1] > 0: cur_ids = torch.cat([left_ids, cur_ids], dim=1) end = check_uncompleted_token(cur_ids[0]) left_ids = cur_ids[:, end:] cur_ids = cur_ids[:, :end] if 151692 in cur_ids[0].cpu().tolist(): #<|tts_eos|> end = cur_ids[0].cpu().tolist().index(151692) eos = True cur_ids = cur_ids[:, :end] cur_hidden_states = cur_hidden_states[:, :end] text = self._decode_text(cur_ids, tokenizer)[0] if end > 0 else "" self.llm_past_key_values = outputs.past_key_values input_ids = outputs.sequences[:, -1:] cache_length = past_length = self.llm_past_key_values[0][0].shape[2] attention_mask = torch.ones((1, cache_length + input_ids.shape[1]), dtype=torch.bool, device=self.device) res = {"text": text, "text_token":cur_ids, "hidden_states": cur_hidden_states} yield res if eos: self.llm_generate_completed = True break if new_len >= max_new_tokens: logger.debug(f"LLM generation {new_len} exceeds max_new_tokens({max_new_tokens}), break.") break class MultiModalProjector(nn.Module): def __init__(self, in_dim, out_dim): super().__init__() self.linear1 = nn.Linear(in_features=in_dim, out_features=out_dim, bias=True) self.relu = nn.ReLU() self.linear2 = nn.Linear(in_features=out_dim, out_features=out_dim, bias=True) def forward(self, audio_features): hidden_states = self.relu(self.linear1(audio_features)) hidden_states = self.linear2(hidden_states) return hidden_states def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, position_ids=None, use_cache=True, **kwargs, ): if past_key_values is not None: if isinstance(past_key_values, Cache): cache_length = past_key_values.get_seq_length() past_length = past_key_values.seen_tokens else: cache_length = past_length = past_key_values[0][0].shape[2] # Keep only the unprocessed tokens: # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as # input) if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard # input_ids based on the past_length. elif past_length < input_ids.shape[1]: input_ids = input_ids[:, past_length:] # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if past_key_values: position_ids = position_ids[:, -input_ids.shape[1] :] # This clo≠clo≠clone call is needed to avoid recapturing cuda graphs with →rch.comπ≤→rch.comπ≤torch.compile's mode=reduce−overheadmode=reduce-overheadmode="reduce-overhead, as otherwise the input positionidspositionidsposition_ids would have various stride during the decoding. Here, simply using .contiguous().contiguous().contiguous() is not sufficient as in the batch size = 1 case, positionidspositionidsposition_ids is already contiguous but with varying stride which retriggers a capture. position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if ∈putsembeds∈putsembedsinputs_embeds are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and cache_position[0] == 0: model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} else: # The clone here is for the same reason as for positionidspositionidsposition_ids. model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None} if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: if model_inputs["inputs_embeds"] is not None: batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape device = model_inputs["inputs_embeds"].device else: batch_size, sequence_length = model_inputs["input_ids"].shape device = model_inputs["input_ids"].device dtype = self.lm_head.weight.dtype min_dtype = torch.finfo(dtype).min attention_mask = _prepare_4d_causal_attention_mask_with_cache_position( attention_mask, sequence_length=sequence_length, target_length=past_key_values.get_max_length(), dtype=dtype, device=device, min_dtype=min_dtype, cache_position=cache_position, batch_size=batch_size, ) model_inputs.update( { "position_ids": position_ids, # "cache_position": cache_position, "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, } ) return model_inputs # Copied from transformers.models.whisper.modeling_whisper.WhisperEncoderLayer and add use_cache for streaming inference class MiniCPMWhisperEncoderLayer(nn.Module): def __init__(self, config: WhisperConfig, layer_idx: int = None): super().__init__() self.embed_dim = config.d_model #1024 self.self_attn = WHISPER_ATTENTION_CLASSES[config._attn_implementation]( embed_dim=self.embed_dim, num_heads=config.encoder_attention_heads, dropout=config.attention_dropout, config=config, layer_idx=layer_idx, ) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.dropout = config.dropout self.activation_fn = ACT2FN[config.activation_function] self.activation_dropout = config.activation_dropout self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) self.final_layer_norm = nn.LayerNorm(self.embed_dim) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, layer_head_mask: torch.Tensor, output_attentions: bool = False, past_key_values: Optional[EncoderDecoderCache] = None, use_cache: Optional[bool] = False, ) -> torch.Tensor: r""" Args: hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, embed_dim)`): Hidden states to be fed into the encoder layer. attention_mask (`torch.FloatTensor` of shape `(batch_size, 1, tgt_len, src_len)`): Attention mask where padding elements are indicated by large negative values. layer_head_mask (`torch.FloatTensor` of shape `(encoder_attention_heads,)`): Mask to nullify selected heads of the attention modules. output_attentions (`bool`, *optional*): Whether or not to return the attention weights. past_key_values (`EncoderDecoderCache`, *optional*): Past key-value pairs used for incremental decoding. use_cache (`bool`, *optional*): Whether or not to return updated `past_key_values` for caching. Returns: A tuple of shape `(hidden_states, optional(attn_weights), optional(past_key_values))`. """ residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) hidden_states, attn_weights, past_key_values = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions, past_key_value=past_key_values, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.activation_fn(self.fc1(hidden_states)) hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) hidden_states = self.fc2(hidden_states) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states if hidden_states.dtype == torch.float16 and ( torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any() ): clamp_value = torch.finfo(hidden_states.dtype).max - 1000 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) outputs = (hidden_states,) if output_attentions: outputs += (attn_weights,) if use_cache: outputs += (past_key_values,) return outputs # Copied from from transformers.models.whisper.modeling_whisper.WhisperEncoder and add use_cache for streaming inference class MiniCPMWhisperEncoder(WhisperEncoder): def __init__(self, config: WhisperConfig): # print(config) super().__init__(config) self.layers = nn.ModuleList( [MiniCPMWhisperEncoderLayer(config, layer_idx=i) for i in range(config.encoder_layers)] ) def forward( self, input_features, attention_mask=None, head_mask=None, output_attentions=None, output_hidden_states=None, return_dict=None, past_key_values: Optional[EncoderDecoderCache] = None, use_cache: Optional[bool] = None, ): r""" Forward pass of the Whisper encoder. Args: input_features (`torch.FloatTensor` of shape `(batch_size, feature_size, sequence_length)`): Float values of log-mel features extracted from the raw audio waveform. Typically generated by a feature extractor (e.g., `WhisperFeatureExtractor`) that processes `.flac` or `.wav` files into padded 2D mel spectrogram frames. These features are projected via convolution layers (`conv1` and `conv2`) and then transformed into embeddings for the encoder. attention_mask (`torch.Tensor`, *optional*): Not used by Whisper for masking `input_features`, but included for API compatibility with other models. If provided, it is simply ignored within the model. By default, Whisper effectively ignores silence in the input log-mel spectrogram. head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): Mask to nullify selected attention heads. The elements should be either 1 or 0, where: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked** (i.e., the attention head is dropped). output_attentions (`bool`, *optional*): Whether or not to return the attention tensors of all encoder layers. If set to `True`, the returned tuple (or `BaseModelOutputWithPast`) will contain an additional element with attention weights for each encoder layer. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. If set to `True`, the returned tuple (or `BaseModelOutputWithPast`) will contain a tuple of hidden states, including the initial embedding output as well as the outputs of each layer. return_dict (`bool`, *optional*): Whether or not to return a `BaseModelOutputWithPast` (a subclass of `ModelOutput`) instead of a plain tuple. If set to `True`, the output will be a `BaseModelOutputWithPast` object, otherwise it will be a tuple. past_key_values (`EncoderDecoderCache`, *optional*): When using caching for faster inference, this is an object that stores the key-value pairs for attention states. If provided, the model will append new states to the existing cache and return the updated cache. This speeds up sequential decoding or chunked inference. - If `past_key_values` is `None`, no past states are used or returned. - If `past_key_values` is not `None` and `use_cache=True`, the model will use the provided cache and return the updated cache (as `next_encoder_cache`). use_cache (`bool`, *optional*): Whether or not the model should use caching (`past_key_values`) to speed up processing during inference. When set to `True`, the model will: - Inspect and use `past_key_values` if provided. - Return updated `past_key_values` (under the name `next_encoder_cache` in `BaseModelOutputWithPast`). Returns: `BaseModelOutputWithPast` or `tuple` (depending on `return_dict`): If `return_dict=True`, a `BaseModelOutputWithPast` is returned, which contains: - **last_hidden_state** (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): The output of the final encoder layer. - **hidden_states** (`tuple(torch.FloatTensor)`, *optional*, returned if `output_hidden_states=True`): Hidden states of the model at each layer (including the initial projection). - **attentions** (`tuple(torch.FloatTensor)`, *optional*, returned if `output_attentions=True`): Attention weights from each encoder layer. - **past_key_values** (an object of type `EncoderDecoderCache` or `None`, *optional*): Updated cache of key-value pairs if `use_cache=True`. If `return_dict=False`, a tuple is returned, where the format is: `(last_hidden_state, hidden_states, attentions)`, with `hidden_states` and `attentions` only present if their respective `output_*` arguments are set to `True`. Example: >>> from transformers import AutoFeatureExtractor, WhisperConfig, WhisperForConditionalGeneration >>> import torch >>> # Load a feature extractor and a Whisper model >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-tiny.en") >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") >>> # Assume you have audio (list of floats or numpy array) loaded from a file >>> # Then extract the mel features: >>> input_features = feature_extractor(audio, sampling_rate=16000, return_tensors="pt").input_features >>> # Forward pass >>> outputs = model.encoder( ... input_features=input_features, ... output_hidden_states=True, ... output_attentions=True, ... use_cache=True ... ) >>> # Retrieve the last hidden state >>> last_hidden_state = outputs.last_hidden_state >>> print(last_hidden_state.shape) torch.Size([batch_size, seq_length, hidden_size]) >>> # Retrieve the intermediate hidden states if output_hidden_states=True >>> all_encoder_hidden_states = outputs.hidden_states >>> # Retrieve attention weights if output_attentions=True >>> all_encoder_attentions = outputs.attentions >>> # Retrieve updated past key values if use_cache=True >>> encoder_cache = outputs.past_key_values """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict # Ignore copy input_features = input_features.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device) inputs_embeds = nn.functional.gelu(self.conv1(input_features)) inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) inputs_embeds = inputs_embeds.permute(0, 2, 1) # import ipdb; ipdb.set_trace() embed_pos = self.embed_positions.weight if embed_pos.shape[0] == 0: #分布式训练 params_to_gather = [param for param in self.embed_positions.parameters()] with deepspeed.zero.GatheredParameters(params_to_gather, modifier_rank=0): embed_pos = deepcopy(self.embed_positions.weight) # import ipdb; ipdb.set_trace() past_key_values_length = 0 if use_cache: if past_key_values is None: past_key_values = EncoderDecoderCache(DynamicCache(), DynamicCache()) elif isinstance(past_key_values, list): past_key_values = EncoderDecoderCache(DynamicCache.from_legacy_cache(past_key_values), DynamicCache()) elif isinstance(past_key_values, DynamicCache): past_key_values = EncoderDecoderCache(past_key_values, DynamicCache()) else: pass past_key_values_length = past_key_values.self_attention_cache.get_usable_length(inputs_embeds.shape[1]) if inputs_embeds.shape[1] + past_key_values_length > embed_pos.shape[0]: logger.warning("seems the audio is longer than 30s. repeating the last part of the audio") embed_pos_front = embed_pos[past_key_values_length:, :] embed_pos = torch.cat( ( embed_pos_front, torch.repeat_interleave( embed_pos[-1, :].unsqueeze(0), inputs_embeds.shape[1] - embed_pos.shape[0] + past_key_values_length, dim=0, ), ) ) else: embed_pos = embed_pos[past_key_values_length : inputs_embeds.shape[1] + past_key_values_length, :] else: embed_pos = embed_pos[: inputs_embeds.shape[1], :] hidden_states = inputs_embeds + embed_pos hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None # check if head_mask has a correct number of layers specified if desired if head_mask is not None: assert head_mask.size()[0] == ( len(self.layers) ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." for idx, encoder_layer in enumerate(self.layers): if output_hidden_states: encoder_states = encoder_states + (hidden_states,) # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) to_drop = False if self.training: dropout_probability = torch.rand([]) if dropout_probability < self.layerdrop: # skip the layer to_drop = True # Ignore copy if to_drop: layer_outputs = (None, None) else: if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( encoder_layer.__call__, hidden_states, attention_mask, (head_mask[idx] if head_mask is not None else None), output_attentions, past_key_values, use_cache, ) else: layer_outputs = encoder_layer( hidden_states, attention_mask, layer_head_mask=(head_mask[idx] if head_mask is not None else None), output_attentions=output_attentions, past_key_values=past_key_values, use_cache=use_cache, ) hidden_states = layer_outputs[0] if use_cache: next_encoder_cache = layer_outputs[2 if output_attentions else 1] else: next_encoder_cache = None if output_attentions: all_attentions = all_attentions + (layer_outputs[1],) hidden_states = self.layer_norm(hidden_states) if output_hidden_states: encoder_states = encoder_states + (hidden_states,) if not return_dict: return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions, past_key_values=next_encoder_cache, )