diff --git "a/modeling_minicpmo.py" "b/modeling_minicpmo.py"
new file mode 100644--- /dev/null
+++ "b/modeling_minicpmo.py"
@@ -0,0 +1,3259 @@
+# 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 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 .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"])
+
+ 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)
+
+ 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:
+ cur_vllm_emb += cur_vs_hs[0].mean() * 0
+
+ return 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):
+ 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
+ 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):
+ 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[0, 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 += 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):
+ 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 = "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:
+ logger.warning(
+ "Warning: ref_audio is None, speech generation will be performed based on the default voice."
+ )
+ sys_msgs = {"role": "user", "content": ["Use the <reserved_53> 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 <reserved_53> 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=None,
+ audio_feature_lens=None,
+ image_bound=None,
+ audio_bounds=None,
+ spk_bounds=None,
+ attention_mask=None,
+ tokenizer=None,
+ vision_hidden_states=None,
+ stream=False,
+ **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)
+
+ result = self._decode_text(outputs.sequences, tokenizer)
+
+ 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("(<image>./</image>)")
+ elif isinstance(c, np.ndarray): # audio
+ audios.append(c)
+ audio_parts.append(i)
+ cur_msgs.append("(<audio>./</audio>)")
+ 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.01,
+ }
+ 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
+
+ @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("(<image>./</image>)")
+ elif isinstance(c, np.ndarray): # audio
+ audios.append(c)
+ cur_msgs.append("(<audio>./</audio>)")
+ 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(omni_input)
+ 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,
+ generate_audio=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.01,
+ }
+ 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)
+
+ if generate_audio:
+ result = self._generate_mel_spec_audio_streaming(
+ spk_bounds, streamer, output_chunk_size=25, enable_regenerate=enable_regenerate
+ )
+ return result
+ else:
+ 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
+ first_chunk = True
+ 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 if first_chunk else False,
+ 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:]
+ 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]
+ 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}
+ if first_chunk:
+ res["hidden_states"] = outputs.hidden_states
+ first_chunk = False
+ 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
+
+ def prepare_tts_text(self, text):
+ tts_tokens = self.tts_processor.text_tokenizer.encode(text, add_special_tokens=False)
+ tts_tokens_len = len(tts_tokens)
+ if tts_tokens_len < self.tts.streaming_text_reserved_len:
+ num_pad_tokens = self.tts.streaming_text_reserved_len - tts_tokens_len
+
+ pad_str = "[Etts]" + "[PAD]" * (num_pad_tokens - 1)
+ else:
+ tts_tokens = tts_tokens[0 : self.tts.streaming_text_reserved_len]
+ tts_tokens_len = len(tts_tokens)
+ text = self.tts_processor.text_tokenizer.decode(tts_tokens, add_special_tokens=False)
+ pad_str = ""
+ spk_emb_placeholder_tts = "[spk_emb]" * self.tts.num_spk_embs
+
+ new_text_tts = f"[Stts]{spk_emb_placeholder_tts}{text}{pad_str}[Ptts]"
+ return new_text_tts, tts_tokens_len
+
+ def get_tts_text_start_token_ids(self):
+ text = "[Stts]" + "[spk_emb]" * self.tts.num_spk_embs
+ tts_input_ids = self.tts_processor.text_tokenizer(text, return_tensors="pt", add_special_tokens=False)[
+ "input_ids"
+ ].cuda()
+ return tts_input_ids
+
+ def _build_streaming_mask(self, tts_tokens_len):
+ tts_sequence_full_length = (
+ 1 + self.tts.num_spk_embs * self.tts.use_speaker_embedding + self.tts.streaming_text_reserved_len + 1
+ )
+ streaming_attention_mask = torch.zeros(tts_sequence_full_length, dtype=torch.int8)
+ streaming_attention_mask[0 : 1 + 1 + tts_tokens_len + 1] = 1
+ streaming_attention_mask[-1] = 1
+ return streaming_attention_mask
+
+ def _get_last_spk_embeds(self, inputs, outputs):
+ last_hidden_states = [hs[-1] for hs in outputs.hidden_states]
+
+ # batch = 1
+ last_hidden_states = torch.vstack([i[0] for i in last_hidden_states])
+
+ # last spk
+ spk_bound = inputs["spk_bounds"][0][-1]
+
+ spk_embeds = last_hidden_states[spk_bound[0] : spk_bound[1]]
+ return spk_embeds
+
+ def _generate_mel_spec(self, inputs, outputs, text, output_chunk_size=25, tts_max_new_tokens=2048):
+ spk_embeds = self._get_last_spk_embeds(inputs, outputs)
+
+ text = text.split("<|tts_bos|>")[-1]
+ gen_text = text.split("<|tts_eos|>")[0]
+ tts_text, tts_token_lens = self.prepare_tts_text(gen_text)
+ tts_inputs = self.tts_processor.text_tokenizer.encode(tts_text, add_special_tokens=False)
+ tts_input_ids = torch.Tensor(tts_inputs).unsqueeze(0).to("cuda", dtype=torch.long)
+ streaming_tts_text_mask = self._build_streaming_mask(tts_token_lens).to(device=self.tts.device)
+
+ logits_warpers, logits_processors = gen_logits(
+ num_code=626, top_P=self.tts.top_p, top_K=self.tts.top_k, repetition_penalty=self.tts.repetition_penalty
+ )
+
+ condition_length = (
+ 1 + self.tts.use_speaker_embedding * self.tts.num_spk_embs + self.tts.streaming_text_reserved_len + 1
+ )
+
+ dtype = self.tts.emb_text.weight.dtype
+ emb = torch.zeros(1, condition_length, self.tts.num_vq, dtype=dtype, device=self.tts.device)
+ past_key_values = [
+ (
+ torch.zeros(
+ 1,
+ self.tts.config.num_attention_heads,
+ condition_length - 1,
+ self.tts.config.hidden_size // self.tts.config.num_attention_heads,
+ dtype=emb.dtype,
+ device=self.tts.device,
+ ),
+ torch.zeros(
+ 1,
+ self.tts.config.num_attention_heads,
+ condition_length - 1,
+ self.tts.config.hidden_size // self.tts.config.num_attention_heads,
+ dtype=emb.dtype,
+ device=self.tts.device,
+ ),
+ )
+ for _ in range(self.tts.config.num_hidden_layers)
+ ]
+
+ audio_input_ids = torch.zeros(1, condition_length, self.tts.num_vq, dtype=torch.long, device=self.tts.device)
+
+ eos_lab = False
+ for chunk_idx in range(math.ceil(emb.shape[1] / self.tts.streaming_text_chunk_size)):
+ if chunk_idx == 0:
+ begin = chunk_idx * self.tts.streaming_text_chunk_size + 0
+ end = (
+ (chunk_idx + 1) * self.tts.streaming_text_chunk_size
+ + 1
+ + self.tts.use_speaker_embedding * self.tts.num_spk_embs
+ )
+ else:
+ begin = (
+ chunk_idx * self.tts.streaming_text_chunk_size
+ + 1
+ + self.tts.use_speaker_embedding * self.tts.num_spk_embs
+ )
+ end = min(
+ (chunk_idx + 1) * self.tts.streaming_text_chunk_size
+ + 1
+ + self.tts.use_speaker_embedding * self.tts.num_spk_embs,
+ condition_length - 1,
+ )
+
+ if end - begin > 0:
+ text_input_ids = tts_input_ids[:, begin:end]
+ position_ids = torch.arange(begin, end, dtype=torch.long, device=self.tts.device).unsqueeze(0)
+
+ if begin == 0:
+ past_key_values = self.tts.prefill_text(
+ input_ids=text_input_ids,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ lm_spk_emb_last_hidden_states=spk_embeds,
+ )
+ else:
+ past_key_values = self.tts.prefill_text(
+ input_ids=text_input_ids, position_ids=position_ids, past_key_values=past_key_values
+ )
+
+ outputs = self.tts.generate(
+ input_ids=audio_input_ids,
+ past_key_values=past_key_values,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ max_new_token=output_chunk_size,
+ force_no_stop=self.force_no_stop,
+ temperature=torch.tensor([0.1, 0.3, 0.1, 0.3], dtype=torch.float, device=self.tts.device),
+ eos_token=torch.tensor([625], dtype=torch.long, device=self.tts.device),
+ logits_warpers=logits_warpers,
+ logits_processors=logits_processors,
+ )
+ audio_input_ids = outputs.audio_input_ids
+ past_key_values = outputs.past_key_values
+
+ if outputs.finished:
+ logger.debug("Generation finished.")
+ eos_lab = True
+ break
+
+ if not eos_lab:
+ logger.debug("eos_lab False, Generation continue.")
+ while True:
+ outputs = self.tts.generate(
+ input_ids=audio_input_ids,
+ past_key_values=past_key_values,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ max_new_token=output_chunk_size,
+ force_no_stop=self.force_no_stop,
+ temperature=torch.tensor([0.1, 0.3, 0.1, 0.3], dtype=torch.float, device=self.tts.device),
+ eos_token=torch.tensor([625], dtype=torch.long, device=self.tts.device),
+ logits_warpers=logits_warpers,
+ logits_processors=logits_processors,
+ )
+
+ audio_input_ids = outputs.audio_input_ids
+ past_key_values = outputs.past_key_values
+
+ if outputs.finished:
+ logger.debug("Generation finished.")
+ break
+ if outputs.new_ids.shape[1] > tts_max_new_tokens:
+ logger.debug(f"Generation length > {tts_max_new_tokens}, stopped.")
+ break
+
+ mel_spec = self.tts.decode_to_mel_specs(outputs.new_ids)
+ return mel_spec
+
+ def _linear_overlap_add2_wav(self, frames: List[torch.Tensor], overlap: int):
+ """
+ Merge two audio waveforms with smooth in streaming audio generation.
+ Borrowed some codes from `https://github.com/huggingface/transformers/blob/main/src/transformers/models/encodec/modeling_encodec.py`
+ """
+ assert len(frames) == 2
+ device = frames[0].device
+ dtype = frames[0].dtype
+ # shape = frames[0].shape[:-1]
+
+ frame0_length = frames[0].shape[-1]
+ frame1_length = frames[1].shape[-1]
+ total_size = frame0_length + frame1_length - overlap
+ weight_len = max(frame0_length, frame1_length) + overlap
+ t = torch.linspace(0, 1, weight_len + 2, device=device, dtype=dtype)[1:-1]
+ weight = 0.5 - (t - 0.5).abs()
+
+ sum_weight = torch.zeros(total_size, device=device, dtype=dtype)
+ out = torch.zeros(total_size, device=device, dtype=dtype)
+ offset: int = 0
+
+ out[offset : offset + frame0_length] += weight[-frame0_length:] * frames[0]
+ sum_weight[offset : offset + frame0_length] += weight[-frame0_length:]
+ offset += frame0_length - overlap
+ out[offset : offset + frame1_length] += weight[:frame1_length] * frames[1]
+ sum_weight[offset : offset + frame1_length] += weight[:frame1_length]
+
+ assert sum_weight.min() > 0
+ out = out / sum_weight
+ return out[:frame0_length], out[frame0_length:]
+
+ def _generate_mel_spec_audio_streaming(
+ self,
+ spk_bounds,
+ streamer,
+ output_chunk_size=25,
+ spk_embeds=None,
+ prev_seg_text_ids=None,
+ prev_seg_text_left="",
+ prev_seg_audio_ids=None,
+ enable_regenerate=False,
+ ):
+ # get spk_embedding
+ gen_text = ""
+ tts_text = ""
+ new_segment_gen = False
+ if spk_embeds is None:
+ spk_bound = spk_bounds[0][-1]
+ r = next(streamer)
+ txt = r["text"]
+ gen_text += txt.split("<|tts_eos|>")[0]
+ tts_text, tts_token_lens = self.prepare_tts_text(gen_text)
+ last_hidden_states = r["hidden_states"][0][-1][0] # output: (input_seq_len, dim)
+ spk_embeds = last_hidden_states[spk_bound[0] : spk_bound[1]]
+
+ # init past_key_values
+ logits_warpers, logits_processors = gen_logits(
+ num_code=626, top_P=self.tts.top_p, top_K=self.tts.top_k, repetition_penalty=self.tts.repetition_penalty
+ )
+ condition_length = (
+ 1 + self.tts.use_speaker_embedding * self.tts.num_spk_embs + self.tts.streaming_text_reserved_len + 1
+ )
+ tts_start_token_len = 1 + self.tts.use_speaker_embedding * self.tts.num_spk_embs
+ dtype = self.tts.emb_text.weight.dtype
+ past_key_values = [
+ (
+ torch.zeros(
+ 1,
+ self.tts.config.num_attention_heads,
+ condition_length - 1,
+ self.tts.config.hidden_size // self.tts.config.num_attention_heads,
+ dtype=dtype,
+ device=self.tts.device,
+ ),
+ torch.zeros(
+ 1,
+ self.tts.config.num_attention_heads,
+ condition_length - 1,
+ self.tts.config.hidden_size // self.tts.config.num_attention_heads,
+ dtype=dtype,
+ device=self.tts.device,
+ ),
+ )
+ for _ in range(self.tts.config.num_hidden_layers)
+ ]
+ audio_input_ids = torch.zeros(1, condition_length, self.tts.num_vq, dtype=torch.long, device=self.tts.device)
+
+ # prefill prev segment for smooth
+ chunk_idx = 0
+ new_ids_len = 0
+ prev_text_len = 0
+ if prev_seg_text_ids is not None and prev_seg_audio_ids is not None:
+ tts_token_lens = prev_seg_text_ids.shape[1]
+ # assert tts_token_lens % self.tts.streaming_text_chunk_size == 0
+ streaming_tts_text_mask = self._build_streaming_mask(tts_token_lens).to(device=self.tts.device)
+ position_ids = torch.arange(
+ 0, tts_token_lens + tts_start_token_len, dtype=torch.long, device=self.tts.device
+ ).unsqueeze(0)
+
+ text_input_ids = self.get_tts_text_start_token_ids()
+ text_input_ids = torch.cat([text_input_ids, prev_seg_text_ids], dim=1)
+ past_key_values = self.tts.prefill_text(
+ input_ids=text_input_ids,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ lm_spk_emb_last_hidden_states=spk_embeds,
+ )
+ past_key_values = self.tts.prefill_audio_ids(
+ input_ids=prev_seg_audio_ids[:, :-1, :],
+ # not prefill last id, which will be input_id of next generation
+ past_key_values=past_key_values,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ )
+
+ # update init
+ chunk_idx += int(tts_token_lens / self.tts.streaming_text_chunk_size)
+ audio_input_ids = torch.cat([audio_input_ids, prev_seg_audio_ids], dim=1)
+ text = self.tts_processor.text_tokenizer.decode(prev_seg_text_ids[0].tolist(), add_special_tokens=False)
+
+ gen_text += text
+ gen_text += prev_seg_text_left
+ prev_text_len = len(gen_text) # takecare the position
+ new_ids_len += prev_seg_audio_ids.shape[1]
+
+ prev_wav = None
+ eos_lab = False
+ stop = False
+ shift_len = 180
+ voice_checker = VoiceChecker()
+ number_converter = NumberToTextConverter()
+ lang = None
+ gen_text_raw = gen_text
+ for t, r in enumerate(streamer):
+ t += 1
+ txt = r["text"]
+ txt = txt.split("<|tts_eos|>")[0]
+ gen_text_raw += txt
+ if t == 1 and txt == "" and prev_seg_text_ids is not None:
+ logger.warning("New segment is empty, generation finished.")
+ return
+ if t <= 2: # do just one time, more token greater certainty
+ lang = number_converter.detect_language(gen_text_raw)
+ gen_text += number_converter.replace_numbers_with_text(txt, lang).replace("*", "") # markdown **
+
+ # TODO speed up
+ tts_text, tts_token_lens = self.prepare_tts_text(gen_text)
+
+ if tts_token_lens >= self.tts.streaming_text_reserved_len - shift_len:
+ end_c = sentence_end(txt)
+ if end_c:
+ end_c_idx = gen_text.rfind(end_c)
+ assert end_c_idx != -1
+ text_left = gen_text[end_c_idx + 1 :]
+ gen_text = gen_text[: end_c_idx + 1]
+ tts_text, tts_token_lens = self.prepare_tts_text(gen_text)
+ new_segment_gen = True
+ logger.debug(
+ f"tts_text tokens {tts_token_lens} exceed {self.tts.streaming_text_reserved_len - shift_len}, starting a new segment generation"
+ )
+ break
+
+ if tts_token_lens >= (chunk_idx + 1) * self.tts.streaming_text_chunk_size:
+
+ # do prefill and generate
+ if chunk_idx == 0:
+ begin = 0
+ end = (chunk_idx + 1) * self.tts.streaming_text_chunk_size + tts_start_token_len
+ else:
+ begin = chunk_idx * self.tts.streaming_text_chunk_size + tts_start_token_len
+ end = min(
+ (chunk_idx + 1) * self.tts.streaming_text_chunk_size + tts_start_token_len, condition_length - 1
+ )
+
+ tts_input_ids = self.tts_processor.text_tokenizer(
+ tts_text, return_tensors="pt", add_special_tokens=False
+ )["input_ids"].cuda()
+ text_input_ids = tts_input_ids[:, begin:end]
+ streaming_tts_text_mask = self._build_streaming_mask(tts_token_lens).to(device=self.tts.device)
+ position_ids = torch.arange(begin, end, dtype=torch.long, device=self.tts.device).unsqueeze(0)
+
+ past_key_values = self.tts.prefill_text(
+ input_ids=text_input_ids,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ lm_spk_emb_last_hidden_states=spk_embeds if chunk_idx == 0 else None,
+ )
+ outputs = self.tts.generate(
+ input_ids=audio_input_ids,
+ past_key_values=past_key_values,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ max_new_token=output_chunk_size,
+ force_no_stop=self.force_no_stop,
+ temperature=torch.tensor([0.1, 0.3, 0.1, 0.3], dtype=torch.float, device=self.tts.device),
+ eos_token=torch.tensor([625], dtype=torch.long, device=self.tts.device),
+ logits_warpers=logits_warpers,
+ logits_processors=logits_processors,
+ )
+ audio_input_ids = (
+ outputs.audio_input_ids
+ ) # [1,seq_len,4] seq_len=tts.streaming_text_reserved_len + 3 + len(new_ids)
+ past_key_values = outputs.past_key_values
+ chunk_idx += 1
+
+ mel_spec = self.tts.decode_to_mel_specs(outputs.new_ids[:, max(new_ids_len - 4, 0) :, :])
+ new_ids_len = outputs.new_ids.shape[1] # [1, seq_len, 4]
+
+ wav_np, sr = self.decode_mel_to_audio(mel_spec) # [1,100,50] -> [50*256]
+
+ if enable_regenerate:
+ if prev_wav is not None:
+ check_wav_np = wav_np[2048:].cpu().numpy() # 2*4*256(hop)
+ check_mel = mel_spec[0, :, 8:].cpu().numpy() # 2*4
+ else:
+ check_wav_np = wav_np.cpu().numpy()
+ check_mel = mel_spec[0].cpu().numpy()
+ if enable_regenerate and voice_checker.is_bad(check_wav_np, check_mel, chunk_size=2560):
+ voice_checker.reset()
+ # regenerate
+ N = output_chunk_size if prev_wav is None else output_chunk_size * 2
+ past_kv = []
+ for i in range(len(past_key_values)):
+ past_kv.append(
+ (
+ past_key_values[i][0][:, :, :-N, :], # .clone(),
+ past_key_values[i][1][:, :, :-N, :], # .clone(),
+ )
+ )
+ outputs = self.tts.generate(
+ input_ids=audio_input_ids[:, :-N, :],
+ past_key_values=past_kv,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ max_new_token=N,
+ force_no_stop=self.force_no_stop,
+ temperature=torch.tensor([0.1, 0.3, 0.1, 0.3], dtype=torch.float, device=self.tts.device),
+ eos_token=torch.tensor([625], dtype=torch.long, device=self.tts.device),
+ logits_warpers=logits_warpers,
+ logits_processors=logits_processors,
+ )
+ audio_input_ids = outputs.audio_input_ids
+ past_key_values = outputs.past_key_values
+
+ new_ids_len -= N
+ mel_spec = self.tts.decode_to_mel_specs(outputs.new_ids[:, new_ids_len:, :])
+ new_ids_len = outputs.new_ids.shape[1] # [1, seq_len, 4]
+ wav_np, sr = self.decode_mel_to_audio(mel_spec)
+
+ if prev_wav is not None:
+ wav_y = wav_np[: len(prev_wav)]
+ prev_wav = wav_np[len(prev_wav) :]
+ cur_text = gen_text_raw[prev_text_len:]
+ prev_text_len = len(gen_text_raw)
+ yield OmniOutput(text=cur_text, audio_wav=wav_y, sampling_rate=sr)
+
+ else:
+ prev_wav = wav_np
+ else:
+ # smooth wav
+ if prev_wav is not None:
+ wav_np, prev_wav = self._linear_overlap_add2_wav(
+ [prev_wav, wav_np], overlap=512 * 4
+ ) # tts_hop256*2
+ cur_text = gen_text_raw[prev_text_len:]
+ prev_text_len = len(gen_text_raw)
+ yield OmniOutput(text=cur_text, audio_wav=wav_np, sampling_rate=sr)
+
+ else:
+ prev_wav = wav_np
+
+ if outputs.finished:
+ logger.debug("Generation finished.")
+ eos_lab = True
+ break
+
+ if not eos_lab and tts_text:
+ logger.debug("eos_lab False, Generation continue.")
+
+ if chunk_idx == 0:
+ begin = 0
+ else:
+ begin = chunk_idx * self.tts.streaming_text_chunk_size + tts_start_token_len
+ end = tts_token_lens + tts_start_token_len + 1 # 1 for [Etts]
+ if end > begin:
+ tts_input_ids = self.tts_processor.text_tokenizer(
+ tts_text, return_tensors="pt", add_special_tokens=False
+ )["input_ids"].cuda()
+ text_input_ids = tts_input_ids[:, begin:end]
+ streaming_tts_text_mask = self._build_streaming_mask(tts_token_lens).to(device=self.tts.device)
+ position_ids = torch.arange(begin, end, dtype=torch.long, device=self.tts.device).unsqueeze(0)
+
+ past_key_values = self.tts.prefill_text(
+ input_ids=text_input_ids,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ lm_spk_emb_last_hidden_states=spk_embeds if chunk_idx == 0 else None,
+ )
+
+ while True:
+ # temp = [0.1, 0.3, 0.1, 0.3] if chunk_idx < 21 else [0.1] * self.tts.num_vq
+ outputs = self.tts.generate(
+ input_ids=audio_input_ids,
+ past_key_values=past_key_values,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ max_new_token=output_chunk_size,
+ force_no_stop=self.force_no_stop,
+ # temperature=torch.tensor([0.1] * self.tts.num_vq, dtype=torch.float, device=self.tts.device),
+ temperature=torch.tensor([0.1, 0.3, 0.1, 0.3], dtype=torch.float, device=self.tts.device),
+ eos_token=torch.tensor([625], dtype=torch.long, device=self.tts.device),
+ logits_warpers=logits_warpers,
+ logits_processors=logits_processors,
+ )
+ audio_input_ids = outputs.audio_input_ids
+ past_key_values = outputs.past_key_values
+ chunk_idx += 1
+
+ mel_spec = self.tts.decode_to_mel_specs(outputs.new_ids[:, max(new_ids_len - 4, 0) :, :])
+ new_ids_len = outputs.new_ids.shape[1] # [1, seq_len, 4]
+
+ wav_np, sr = self.decode_mel_to_audio(mel_spec)
+
+ if enable_regenerate:
+ if prev_wav is not None:
+ check_wav_np = wav_np[2048:].cpu().numpy() # 2*4*256(hop)
+ check_mel = mel_spec[0, :, 8:].cpu().numpy() # 2*4
+ else:
+ check_wav_np = wav_np.cpu().numpy()
+ check_mel = mel_spec[0].cpu().numpy()
+ if enable_regenerate and voice_checker.is_bad(check_wav_np, check_mel, chunk_size=2560):
+ voice_checker.reset()
+ # regenerate
+ N = output_chunk_size if prev_wav is None else output_chunk_size * 2
+ past_kv = []
+ for i in range(len(past_key_values)):
+ past_kv.append(
+ (
+ past_key_values[i][0][:, :, :-N, :], # .clone(),
+ past_key_values[i][1][:, :, :-N, :], # .clone(),
+ )
+ )
+ outputs = self.tts.generate(
+ input_ids=audio_input_ids[:, :-N, :],
+ past_key_values=past_kv,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ max_new_token=N,
+ force_no_stop=self.force_no_stop,
+ temperature=torch.tensor([0.1, 0.3, 0.1, 0.3], dtype=torch.float, device=self.tts.device),
+ eos_token=torch.tensor([625], dtype=torch.long, device=self.tts.device),
+ logits_warpers=logits_warpers,
+ logits_processors=logits_processors,
+ )
+ audio_input_ids = outputs.audio_input_ids
+ past_key_values = outputs.past_key_values
+
+ new_ids_len -= N
+ mel_spec = self.tts.decode_to_mel_specs(outputs.new_ids[:, new_ids_len:, :])
+ new_ids_len = outputs.new_ids.shape[1] # [1, seq_len, 4]
+ wav_np, sr = self.decode_mel_to_audio(mel_spec)
+
+ if prev_wav is not None:
+ wav_y = wav_np[: len(prev_wav)]
+ prev_wav = wav_np[len(prev_wav) :]
+ cur_text = gen_text_raw[prev_text_len:]
+ prev_text_len = len(gen_text_raw)
+ yield OmniOutput(text=cur_text, audio_wav=wav_y, sampling_rate=sr)
+ else:
+ prev_wav = wav_np
+ else:
+ # smooth wav
+ if prev_wav is not None:
+ wav_np, prev_wav = self._linear_overlap_add2_wav(
+ [prev_wav, wav_np], overlap=512 * 4
+ ) # tts_hop256*2
+ cur_text = gen_text_raw[prev_text_len:]
+ prev_text_len = len(gen_text_raw)
+ yield OmniOutput(text=cur_text, audio_wav=wav_np, sampling_rate=sr)
+ else:
+ prev_wav = wav_np
+
+ if outputs.finished:
+ logger.debug("Generation finished.")
+ break
+ if outputs.new_ids.shape[1] > 2048:
+ stop = True
+ logger.debug("Generation length > 2048, stopped.")
+ break
+
+ if prev_wav is not None:
+ cur_text = gen_text_raw[prev_text_len:]
+ yield OmniOutput(text=cur_text, audio_wav=prev_wav, sampling_rate=sr) # yield last chunk wav without smooth
+
+ if new_segment_gen and not stop:
+ logger.debug(
+ f"tts_text tokens {tts_token_lens} exceed {self.tts.streaming_text_reserved_len - shift_len}, start a new segment generation"
+ )
+ tid_len = 5 # self.tts.streaming_text_chunk_size
+ prev_seg_text_ids = tts_input_ids[:, end - 1 - tid_len : end - 1] # exclude last Etts
+ aid_len = 50 # int(tid_len * new_ids_len / tts_token_lens)
+ prev_seg_audio_ids = outputs.new_ids[:, -aid_len:, :]
+
+ result = self._generate_mel_spec_audio_streaming(
+ spk_bounds,
+ streamer,
+ output_chunk_size,
+ spk_embeds,
+ prev_seg_text_ids,
+ text_left,
+ prev_seg_audio_ids,
+ enable_regenerate=enable_regenerate,
+ )
+ for res in result:
+ yield res
+
+ def decode_mel_to_audio(self, mel_spec, output_path=""):
+ with torch.inference_mode():
+ wav_numpy = self.vocos.decode(mel_spec.float()).cpu().squeeze()
+ sr = 24000
+ if output_path:
+ sf.write(output_path, wav_numpy.numpy(), samplerate=sr)
+ logger.info(f"Audio saved to {output_path}")
+ return wav_numpy, sr
+
+
+# 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
+ 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):
+ 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)
+
+ embed_pos = self.embed_positions.weight
+ 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,
+ )
+
+
+# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py`
+class ConvNeXtBlock(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ intermediate_dim: int,
+ kernel: int,
+ dilation: int,
+ layer_scale_init_value: float = 1e-6,
+ ):
+ # ConvNeXt Block copied from Vocos.
+ super().__init__()
+ self.dwconv = nn.Conv1d(
+ dim,
+ dim,
+ kernel_size=kernel,
+ padding=dilation * (kernel // 2),
+ dilation=dilation,
+ groups=dim,
+ )
+
+ self.norm = nn.LayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(dim, intermediate_dim)
+ self.act = nn.GELU()
+ self.pwconv2 = nn.Linear(intermediate_dim, dim)
+ self.coef = (
+ nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+ if layer_scale_init_value > 0
+ else None
+ )
+
+ def forward(self, x: torch.Tensor, cond=None) -> torch.Tensor:
+ residual = x
+
+ y = self.dwconv(x)
+ y.transpose_(1, 2) # (B, C, T) -> (B, T, C)
+ x = self.norm(y)
+ del y
+ y = self.pwconv1(x)
+ del x
+ x = self.act(y)
+ del y
+ y = self.pwconv2(x)
+ del x
+ if self.coef is not None:
+ y *= self.coef
+ y.transpose_(1, 2) # (B, T, C) -> (B, C, T)
+
+ x = y + residual
+ del y
+
+ return x
+
+
+# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py`
+class GFSQ(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ levels: List[int],
+ G: int,
+ R: int,
+ eps=1e-5,
+ transpose=True,
+ ):
+ super(GFSQ, self).__init__()
+ self.quantizer = GroupedResidualFSQ(
+ dim=dim,
+ levels=list(levels),
+ num_quantizers=R,
+ groups=G,
+ )
+ self.n_ind = math.prod(levels)
+ self.eps = eps
+ self.transpose = transpose
+ self.G = G
+ self.R = R
+
+ def _embed(self, x: torch.Tensor):
+ if self.transpose:
+ x = x.transpose(1, 2)
+ x = x.view(x.size(0), x.size(1), self.G, self.R).permute(2, 0, 1, 3)
+ feat = self.quantizer.get_output_from_indices(x)
+ return feat.transpose_(1, 2) if self.transpose else feat
+
+ def __call__(self, x: torch.Tensor) -> torch.Tensor:
+ return super().__call__(x)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.transpose:
+ x.transpose_(1, 2)
+ _, ind = self.quantizer(x)
+ ind = ind.permute(1, 2, 0, 3).contiguous()
+ ind = ind.view(ind.size(0), ind.size(1), -1)
+ return ind.transpose_(1, 2) if self.transpose else ind
+
+
+# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py`
+class DVAEDecoder(nn.Module):
+ def __init__(
+ self,
+ idim: int,
+ odim: int,
+ n_layer=12,
+ bn_dim=64,
+ hidden=256,
+ kernel=7,
+ dilation=2,
+ up=False,
+ ):
+ super().__init__()
+ self.up = up
+ self.conv_in = nn.Sequential(
+ nn.Conv1d(idim, bn_dim, 3, 1, 1),
+ nn.GELU(),
+ nn.Conv1d(bn_dim, hidden, 3, 1, 1),
+ )
+ self.decoder_block = nn.ModuleList(
+ [
+ ConvNeXtBlock(
+ hidden,
+ hidden * 4,
+ kernel,
+ dilation,
+ )
+ for _ in range(n_layer)
+ ]
+ )
+ self.conv_out = nn.Conv1d(hidden, odim, kernel_size=1, bias=False)
+
+ def forward(self, x: torch.Tensor, conditioning=None) -> torch.Tensor:
+ # B, C, T
+ y = self.conv_in(x)
+ del x
+ for f in self.decoder_block:
+ y = f(y, conditioning)
+
+ x = self.conv_out(y)
+ del y
+ return x
+
+
+# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py`
+class DVAE(nn.Module):
+ def __init__(
+ self,
+ ):
+ super().__init__()
+
+ coef = torch.rand(100)
+ self.coef = nn.Parameter(coef.unsqueeze(0).unsqueeze_(2))
+
+ self.downsample_conv = nn.Sequential(
+ nn.Conv1d(100, 512, 3, 1, 1),
+ nn.GELU(),
+ nn.Conv1d(512, 512, 4, 2, 1),
+ nn.GELU(),
+ )
+
+ self.encoder = DVAEDecoder(
+ idim=512,
+ odim=1024,
+ hidden=256,
+ n_layer=12,
+ bn_dim=128,
+ )
+
+ self.decoder = DVAEDecoder(
+ idim=512,
+ odim=512,
+ hidden=256,
+ n_layer=12,
+ bn_dim=128,
+ )
+
+ self.out_conv = nn.Conv1d(512, 100, 3, 1, 1, bias=False)
+
+ self.vq_layer = GFSQ(
+ dim=1024,
+ levels=(5, 5, 5, 5),
+ G=2,
+ R=2,
+ )
+
+ @torch.inference_mode()
+ def forward(self, inp: torch.Tensor, mode: Literal["encode", "decode"] = "decode") -> torch.Tensor:
+ if mode == "encode" and hasattr(self, "encoder") and self.vq_layer is not None:
+ mel = inp.clone()
+ x: torch.Tensor = self.downsample_conv(
+ torch.div(mel, self.coef.view(100, 1).expand(mel.shape), out=mel),
+ ).unsqueeze_(0)
+ del mel
+ x = self.encoder(x)
+ ind = self.vq_layer(x)
+ del x
+ return ind
+
+ if self.vq_layer is not None:
+ vq_feats = self.vq_layer._embed(inp)
+ else:
+ vq_feats = inp
+
+ vq_feats = (
+ vq_feats.view(
+ (vq_feats.size(0), 2, vq_feats.size(1) // 2, vq_feats.size(2)),
+ )
+ .permute(0, 2, 3, 1)
+ .flatten(2)
+ )
+
+ dec_out = self.out_conv(
+ self.decoder(
+ x=vq_feats,
+ ),
+ )
+
+ del vq_feats
+
+ return torch.mul(dec_out, self.coef, out=dec_out)
+
+
+def apply_spk_emb(
+ input_ids: torch.Tensor = None,
+ spk_emb: torch.Tensor = None,
+ input_embeds: torch.Tensor = None,
+ spk_emb_token_id: int = 0,
+ num_spk_embs: int = 1,
+):
+ """
+ Replace consecutive `num_spk_embs` speaker embedding placeholders in input_embeds with pre-prepared speaker embeddings. This is an in-place replacement, no new tensor is created, so no value is returned.
+
+ Args:
+ input_ids (torch.Tensor): Input ID tensor, shape [batch_size, seq_len_max]
+ spk_emb (torch.Tensor): Speaker embedding tensor, shape [batch_size, num_spk_emb, hidden_dim]
+ input_embeds (torch.Tensor): Input embedding tensor, shape [batch_size, seq_len_max, hidden_dim]
+ spk_emb_token_id (int): ID of the speaker embedding token
+ num_spk_embs (int): Number of speaker embeddings
+
+ Returns:
+ None
+ """
+
+ batch_size = input_ids.shape[0]
+
+ for idx in range(batch_size):
+ input_ids_ = input_ids[idx] # [seq_len_max]
+ spk_emb_ = spk_emb[idx] # [num_spk_emb]
+ mask_ = input_ids_ == spk_emb_token_id # [batch_size, seq_len_max]
+ nonzero_position_idx = mask_.nonzero(as_tuple=False) # [num_spk_emb, 1]
+ assert nonzero_position_idx.shape[0] == num_spk_embs
+ begin_idx = nonzero_position_idx.min()
+ end_idx = nonzero_position_idx.max()
+ input_embeds[idx, begin_idx : end_idx + 1, :] = spk_emb_
+
+ return
+
+
+def make_streaming_chunk_mask_generation(
+ inputs_embeds: torch.Tensor,
+ past_seen_tokens: int,
+ streaming_tts_text_mask: torch.Tensor,
+ streaming_reserved_length: int = 300,
+ streaming_audio_chunk_size: int = 50,
+ streaming_text_chunk_size: int = 10,
+ num_spk_emb: int = 1,
+ use_spk_emb: bool = True,
+) -> torch.Tensor:
+ """
+ In streaming audio generation, determine which `text` positions the TTS model can attend to when generating each chunk of `audio` tokens.
+
+ This function creates a mask that allows the model to attend to a specific chunk of text
+ tokens when generating each chunk of audio tokens, enabling streaming TTS generation.
+
+ Args:
+ inputs_embeds (torch.Tensor): Input embeddings tensor.
+ past_seen_tokens (int): Number of tokens already seen by the model.
+ streaming_tts_text_mask (torch.Tensor): Mask for the text tokens.
+ streaming_reserved_length (int, optional): Number of reserved tokens for streaming. Defaults to 300.
+ streaming_chunk_length (int, optional): Length of each streaming chunk. Defaults to 50.
+ streaming_text_chunk_size (int, optional): Size of each text chunk. Defaults to 7.
+
+ Returns:
+ torch.Tensor: Causal mask for streaming TTS generation, shape is [batch_size=1, 1, seq_len=1, past_seen_tokens+1]
+
+ Raises:
+ AssertionError: If the batch size is not 1 (only supports batch size of 1 for inference).
+ """
+ assert inputs_embeds.shape[0] == 1
+
+ dtype = inputs_embeds.dtype
+ device = inputs_embeds.device
+ min_dtype = torch.finfo(dtype).min
+
+ # Add `1` to the past seen tokens to account for new `tokens` during `generate`
+ causal_mask = torch.full((1, past_seen_tokens + inputs_embeds.shape[1]), fill_value=0, dtype=dtype, device=device)
+
+ # Calculate the start of invisible text tokens
+ invisible_text_tokens_start = (
+ min(
+ math.ceil((past_seen_tokens - streaming_reserved_length) / streaming_audio_chunk_size)
+ * streaming_text_chunk_size,
+ streaming_reserved_length,
+ )
+ + 1
+ + num_spk_emb * use_spk_emb
+ ) # Add 1 for [Stts] and N for [spk_emb] tokens if `use_spk_emb` is True
+
+ invisible_text_tokens_end = (
+ streaming_reserved_length + 1 + num_spk_emb * use_spk_emb + 1
+ ) # Add 1 for [Ptts] (aka `audio_bos_token_id`)
+
+ # Set invisible text tokens to min_dtype (effectively -inf)
+ causal_mask[0, invisible_text_tokens_start:invisible_text_tokens_end] = min_dtype
+
+ # Mask padding positions in the text mask
+ causal_mask[0, 0 : 1 + num_spk_emb * use_spk_emb + streaming_reserved_length + 1].masked_fill_(
+ streaming_tts_text_mask == 0, min_dtype
+ )
+
+ # Add extra dimensions for batch and heads
+ causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)
+
+ return causal_mask
+
+
+# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/processors.py`
+class CustomRepetitionPenaltyLogitsProcessorRepeat:
+ def __init__(self, penalty: float, max_input_ids: int, past_window: int):
+ if not isinstance(penalty, float) or not (penalty > 0):
+ raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")
+
+ self.penalty = penalty
+ self.max_input_ids = max_input_ids
+ self.past_window = past_window
+
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+ if input_ids.size(1) > self.past_window:
+ input_ids = input_ids.narrow(1, -self.past_window, self.past_window)
+ freq = F.one_hot(input_ids, scores.size(1)).sum(1)
+ if freq.size(0) > self.max_input_ids:
+ freq.narrow(0, self.max_input_ids, freq.size(0) - self.max_input_ids).zero_()
+ alpha = torch.pow(self.penalty, freq)
+ scores = scores.contiguous()
+ inp = scores.multiply(alpha)
+ oth = scores.divide(alpha)
+ con = scores < 0
+ out = torch.where(con, inp, oth)
+ del inp, oth, scores, con, alpha
+ return out
+
+
+@dataclass
+class ConditionalChatTTSGenerationOutput(ModelOutput):
+ """
+ Output class for ConditionalChatTTS generation.
+
+ Args:
+ new_ids (torch.LongTensor): Newly generated audio code sequence, shape (batch_size, sequence_length, num_vq).
+ audio_input_ids (torch.LongTensor): Updated input IDs including condition and generated audio codes, shape (batch_size, full_sequence_length, num_vq).
+ past_key_values (Tuple[Tuple[torch.FloatTensor]]): Tuple containing pre-computed keys and values used for attention mechanism. Each element has shape (batch_size, num_heads, sequence_length, embed_size_per_head).
+ finished (bool): Boolean indicating whether generation is complete.
+
+ """
+
+ new_ids: torch.LongTensor = None
+ audio_input_ids: torch.LongTensor = None
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+ finished: bool = None
+
+
+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
+
+
+class ConditionalChatTTS(PreTrainedModel):
+ """A conditional text-to-speech model that can generate speech from text with speaker conditioning.
+
+ This model extends PreTrainedModel to provide text-to-speech capabilities with:
+ - LLM hidden state conditioning
+ - Streaming generation
+
+ The model uses a transformer architecture with LLM hidden states and can operate in both
+ streaming and non-streaming modes for flexible deployment.
+
+ The model process sequence in the following format:
+ | text bos token | LLM embedding projected to tts embedding space | text tokens (fixed length, reserved for future tokens) | audio bos token | audio tokens (audio token length is not fixed)| audio eos token |
+
+ The format is designed to support LLM-conditioned streaming audio generation.
+
+ Usage:
+ To support streaming generation, two global variables should be maintained outside of the model.
+ 1. `audio_input_ids`: stores *discrete* audio codes. It is a tensor with shape [1, sequence length+1, num_vq].
+ 2. `past_key_values`: stores the KV cache for both text tokens and audio codes. It is a list of tuples, each tuple contains two tensors with shape [1, num_attention_heads, sequence length, hidden_size // num_attention_heads]
+
+ where `num_vq` is the number of audio codebooks, in default setting, it is `4`.
+
+ 1. Create an empty `past_key_values` with
+ ```python
+ initial_kv_cache_length = 1 + model.num_spk_embs + model.streaming_text_reserved_len # where `1` denotes the `bos` token
+ dtype = model.emb_text.weight.dtype
+ device = model.emb_text.weight.device
+ past_key_values = [
+ (
+ torch.zeros(1, model.config.num_attention_heads, initial_kv_cache_length, model.config.hidden_size // model.config.num_attention_heads, dtype=dtype, device=device),
+ torch.zeros(1, model.config.num_attention_heads, initial_kv_cache_length, model.config.hidden_size // model.config.num_attention_heads, dtype=dtype, device=device)
+ )
+ for _ in range(model.config.num_hidden_layers)
+ ]
+
+ 2. At the same time, create an empty `audio_input_ids` with shape [1, sequence length, num_vq], `num_vq` denotes multiple layer audio codebooks. But here we also include text tokens in the sequence, but they will be zeros, and will not be used, just a placeholder.
+
+ ```python
+ initial_audio_input_ids_length = 1 + model.num_spk_embs + model.streaming_text_reserved_len + 1
+ # [bos token, speaker embeddings, text tokens, audio bos token]
+ audio_input_ids = torch.zeros(batch_size=1, initial_audio_input_ids_length, model.num_vq)
+ ```
+
+ 2. Prefill some text tokens to TTS model (for example, 10 tokens) using `prefill_text` method.
+
+ ```python
+ outputs = llm.generate(**kwargs)
+ llm_tokens = some_function_to_extract_llm_tokens(outputs)
+ lm_spk_emb_last_hidden_states = some_function_to_extract_lm_spk_emb_last_hidden_states(outputs)
+ tts_text_input_ids = tts_tokenizer.encode(llm_tokenizer.decode(llm_tokens))
+ # here assume we are prefilling text token 0 to text token 9 (included), totally 10 tokens.
+ begin = 0
+ end = 9+1
+ position_ids = torch.arange(begin, end, dtype=torch.long, device=device)
+
+ past_key_values = model.prefill_text(
+ input_ids=tts_text_input_ids,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ lm_spk_emb_last_hidden_states=lm_spk_emb_last_hidden_states,
+ )
+ ```
+
+ 3. Make a `streaming_tts_text_mask` to denote which position contains valid text tokens, similar to `attention_mask` in standard causal attention.
+
+ ```python
+ streaming_tts_text_mask = torch.zeros(model.streaming_reserved_length)
+ streaming_tts_text_mask[0:end] = 1 # denotes these post
+ ```
+
+ 3. Generate audio codes using `generate` method.
+
+ ```python
+ outputs = model.generate(
+ input_ids=audio_input_ids,
+ past_key_values=past_key_values,
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ max_new_token=50,
+ )
+
+ # update past_key_values and input_ids
+ past_key_values = outputs.past_key_values
+ audio_input_ids = outputs.input_ids
+ ```
+
+ The `past_key_values` is extended by `max_new_token=50`, and `audio_input_ids` is also extended by `max_new_token=50` after `generate` calling.
+
+ 4. Notice that after prefilling `10` text tokens, the model can generate up to `50` audio tokens, if you want to generate more audio tokens, you need to prefill next `10` text tokens. And it is okay to only generate `25` audio tokens for faster initial response.
+
+ 5. Repeat steps `2,3,4` as needed in your streaming audio generation cases, but ensure usage complies with the following guidelines discussed above.
+ """
+
+ config_class = ConditionalChatTTSConfig
+
+ def __init__(self, config: ConditionalChatTTSConfig):
+ super().__init__(config)
+
+ self.use_speaker_embedding = config.use_speaker_embedding
+ self.use_llm_hidden_state = config.use_llm_hidden_state
+ self.num_spk_embs = config.num_spk_embs
+ self.spk_emb_token_id = config.spk_emb_token_id
+
+ self.use_text = config.use_text
+ self.streaming = config.streaming
+ self.streaming_text_chunk_size = config.streaming_text_chunk_size
+ self.streaming_audio_chunk_size = config.streaming_audio_chunk_size
+ self.streaming_text_reserved_len = config.streaming_text_reserved_len
+ self.audio_bos_token_id = config.audio_bos_token_id
+ self.num_mel_bins = config.num_mel_bins
+ self.num_vq = config.num_vq
+ self.num_audio_tokens = config.num_audio_tokens
+
+ self.top_p = config.top_p
+ self.top_k = config.top_k
+ self.repetition_penalty = config.repetition_penalty
+
+ if self.config.use_mlp:
+ self.projector = MultiModalProjector(config.llm_dim, config.hidden_size)
+ else:
+ self.projector = nn.Linear(config.llm_dim, config.hidden_size, bias=False)
+ self.emb_code = nn.ModuleList(
+ [nn.Embedding(config.num_audio_tokens, config.hidden_size) for _ in range(config.num_vq)]
+ )
+ self.emb_text = nn.Embedding(config.num_text_tokens, config.hidden_size)
+ self.head_code = nn.ModuleList(
+ [
+ weight_norm(
+ nn.Linear(config.hidden_size, config.num_audio_tokens, bias=False),
+ name="weight",
+ )
+ for _ in range(config.num_vq)
+ ]
+ )
+ dvae = DVAE()
+ self.dvae = dvae
+
+ model_config = LlamaConfig(
+ hidden_size=config.hidden_size,
+ intermediate_size=config.intermediate_size,
+ num_attention_heads=config.num_attention_heads,
+ num_hidden_layers=config.num_hidden_layers,
+ max_position_embeddings=config.max_position_embeddings,
+ attn_implementation=config.attn_implementation,
+ )
+
+ model = LlamaModel(model_config)
+ self.model = model
+
+ @torch.inference_mode()
+ def merge_inputs_embeds(
+ self,
+ input_ids: torch.Tensor,
+ lm_spk_emb_last_hidden_states: Optional[torch.Tensor] = None,
+ ):
+ """Merge `input_ids` and `lm_spk_emb_last_hidden_states` to `inputs_embeds`.
+
+ Args:
+ input_ids (torch.Tensor): Input token IDs.
+ lm_spk_emb_last_hidden_states (Optional[torch.Tensor], optional): Last hidden states of speaker embeddings from the language model. Defaults to None.
+
+ Raises:
+ NotImplementedError: If speaker embedding is not used and language model hidden states are not implemented.
+
+ Returns:
+ torch.Tensor: Prepared input embeddings for the model.
+ """
+ assert input_ids.shape[0] == 1
+
+ # Embed input_ids to input_embeds
+ inputs_embeds = self.emb_text(input_ids)
+
+ # Inject speaker embedding to input_embeds if it exists
+ if self.use_speaker_embedding:
+ spk_emb_mask = input_ids == self.spk_emb_token_id
+ if spk_emb_mask.any():
+ assert lm_spk_emb_last_hidden_states is not None
+ # Project spk emb to tts hidden size first, [batch_size, num_spk_emb, llm_dim] -> [batch_size, num_spk_emb, self.hidden_size]
+ lm_spk_emb_last_hidden_states = lm_spk_emb_last_hidden_states.to(self.projector.linear1.weight.dtype)
+ projected_spk_emb = self.projector(lm_spk_emb_last_hidden_states)
+ projected_spk_emb = F.normalize(projected_spk_emb, p=2, dim=-1)
+ apply_spk_emb(
+ input_ids=input_ids,
+ spk_emb=projected_spk_emb,
+ input_embeds=inputs_embeds,
+ spk_emb_token_id=self.spk_emb_token_id,
+ num_spk_embs=self.num_spk_embs,
+ )
+ else:
+ raise NotImplementedError
+
+ return inputs_embeds
+
+ @torch.inference_mode()
+ def prefill_text(
+ self,
+ input_ids: torch.Tensor,
+ position_ids: torch.LongTensor,
+ past_key_values: List[Tuple[torch.Tensor, torch.Tensor]],
+ lm_spk_emb_last_hidden_states: Optional[torch.Tensor] = None,
+ ):
+ """Prefill a chunk of new text tokens in streaming setting.
+ Specifically speaking, update `past_key_values` using new text tokens, then the model will read the new text tokens.
+
+ Args:
+ input_ids (Tensor): Tensor of shape [batch_size, seq_len]
+ position_ids (LongTensor): Tensor of shape [batch_size, seq_len]
+ past_key_values (List[Tuple[Tensor]]): KV Cache of all layers, each layer is a tuple (Tensor, Tensor) denoting keys and values. Each tensor is of seq_len = `self.streaming_text_reserved_len`. `past_key_values` will be updated.
+ lm_spk_emb_last_hidden_states (Tensor, optional): Tensor of shape [batch_size, num_spk_emb, llm_dim]. Defaults to None.
+ lm_last_hidden_states (Tensor, optional): _description_. Defaults to None.
+
+ Note that all `batch_size` should be `1`.
+ """
+ assert input_ids.shape[0] == 1
+ assert past_key_values is not None
+
+ # Merge text and LLM embeddings
+ inputs_embeds = self.merge_inputs_embeds(
+ input_ids=input_ids,
+ lm_spk_emb_last_hidden_states=lm_spk_emb_last_hidden_states,
+ )
+
+ # Clone KV Cache
+ past_key_values_for_prefill = []
+ for i in range(len(past_key_values)):
+ past_key_values_for_prefill.append(
+ (
+ past_key_values[i][0][:, :, : position_ids[:, 0], :].clone(),
+ past_key_values[i][1][:, :, : position_ids[:, 0], :].clone(),
+ )
+ )
+
+ # Model forward
+ outputs_prefill: BaseModelOutputWithPast = self.model(
+ attention_mask=None, # because for text, it is standard causal attention mask, do nothing
+ position_ids=position_ids, # position_ids denotes the position of new text tokens in the sequence
+ past_key_values=past_key_values_for_prefill, # `past_key_values` will be updated by the model
+ inputs_embeds=inputs_embeds, # contains text and language model embedding
+ use_cache=True,
+ output_attentions=False,
+ cache_position=position_ids, # which new positions will use this cache, basically the same as position_ids
+ )
+
+ # Get model updated KV Cache
+ past_key_values_for_prefill_updated = outputs_prefill.past_key_values
+
+ # Update generated KV Cache to input `past_key_values`
+ for layer_idx in range(len(past_key_values)):
+ # Update keys
+ past_key_values[layer_idx][0][:, :, position_ids[:, 0] : position_ids[:, -1] + 1, :] = (
+ past_key_values_for_prefill_updated[layer_idx][0][
+ :, :, position_ids[:, 0] : position_ids[:, -1] + 1
+ ].clone()
+ )
+ # Update values
+ past_key_values[layer_idx][1][:, :, position_ids[:, 0] : position_ids[:, -1] + 1, :] = (
+ past_key_values_for_prefill_updated[layer_idx][1][
+ :, :, position_ids[:, 0] : position_ids[:, -1] + 1
+ ].clone()
+ )
+
+ # TODO: del past_key_values_for_prefill_updated recursively
+ # TODO: del outputs_prefill recursively
+
+ return past_key_values
+
+ @torch.inference_mode()
+ def prefill_audio_ids(
+ self,
+ input_ids: torch.Tensor,
+ past_key_values: List[Tuple[torch.Tensor, torch.Tensor]],
+ streaming_tts_text_mask=None,
+ add_audio_bos: bool = True,
+ ):
+ """Prefill a chunk of audio ids to the model. Used in sliding-window long audio generation.
+ Specifically, prefill many audio ids (typically from last window) to the model in the new window.
+
+ Args:
+ input_ids (torch.Tensor): (1, seq_len, num_vq) Audio input token ids.
+ past_key_values (List[Tuple[torch.Tensor, torch.Tensor]]): Past key values for attention mechanism.
+ """
+ assert input_ids.shape[0] == 1
+ assert past_key_values is not None
+
+ code_emb = [self.emb_code[i](input_ids[:, :, i]) for i in range(self.num_vq)]
+ inputs_embeds = torch.stack(code_emb, 3).sum(3) # [1,seq_len,768]
+ input_len = input_ids.shape[1]
+
+ if add_audio_bos:
+ narrowed_input_ids = torch.tensor([[self.audio_bos_token_id]], dtype=torch.long, device=self.device)
+ bos_inputs_embeds = self.emb_text(narrowed_input_ids)
+ inputs_embeds = torch.cat([bos_inputs_embeds, inputs_embeds], dim=1)
+ input_len += 1
+
+ past_key_values_length = past_key_values[0][0].shape[2]
+ position_ids = torch.arange(
+ past_key_values_length, past_key_values_length + input_len, dtype=torch.long, device=self.device
+ ).unsqueeze(0)
+
+ cache_position = position_ids.clone()
+ causal_mask = make_streaming_chunk_mask_generation(
+ inputs_embeds=inputs_embeds,
+ past_seen_tokens=past_key_values[0][0].shape[2],
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ streaming_reserved_length=self.streaming_text_reserved_len,
+ streaming_text_chunk_size=self.streaming_text_chunk_size,
+ ) # [1, 1, 1, past_key_values_length + input_len]
+
+ # Model forward
+ outputs: BaseModelOutputWithPast = self.model(
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=True,
+ output_attentions=False,
+ cache_position=cache_position,
+ )
+ past_key_values = outputs.past_key_values
+ return past_key_values
+
+ @torch.inference_mode()
+ def generate(
+ self,
+ input_ids: torch.Tensor,
+ past_key_values: List[Tuple[torch.Tensor, torch.Tensor]],
+ temperature: torch.Tensor,
+ eos_token: Union[int, torch.Tensor],
+ streaming_tts_text_mask=None,
+ force_no_stop=False,
+ min_new_token=10,
+ max_new_token=50,
+ logits_warpers: List[LogitsWarper] = [],
+ logits_processors: List[CustomRepetitionPenaltyLogitsProcessorRepeat] = [],
+ show_tqdm=False,
+ ):
+ """Generate audio codes in streaming setting or non-streaming setting.
+ Specifically speaking, generate audio codes when not all text tokens are prefilled.
+
+ Always pass a valid `past_key_values` to the method. The method does not do `prefill` by itself. It relies on `prefill_text` method to provide valid `past_key_values`. Please refer to docstring of this class for more details.
+
+ In this method, we borrowed a lot of codes from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/gpt.py`.
+
+ Args:
+ input_ids (torch.Tensor): Input token ids.
+ past_key_values (List[Tuple[torch.Tensor, torch.Tensor]]): Past key values for attention mechanism.
+ temperature (torch.Tensor): Temperature for sampling.
+ eos_token (Union[int, torch.Tensor]): End of sequence token.
+ streaming_tts_text_mask (Optional[torch.Tensor], optional): Mask for streaming TTS text. Defaults to None.
+ max_new_token (int, optional): Maximum number of new tokens to generate. Defaults to 50.
+ logits_warpers (List[LogitsWarper], optional): List of logits warpers. Defaults to [].
+ logits_processors (List[CustomRepetitionPenaltyLogitsProcessorRepeat], optional): List of logits processors. Defaults to [].
+ show_tqdm (bool, optional): Whether to show progress bar. Defaults to True.
+
+ Returns:
+ GenerationOutputs: Generation outputs.
+ """
+
+ # We only support batch size `1` for now
+ assert input_ids.shape[0] == 1
+ assert past_key_values is not None
+
+ # fix: this should not be `input_ids.shape[1]`
+ # start_idx = input_ids.shape[1]
+ start_idx = 1 + self.num_spk_embs * self.use_speaker_embedding + self.streaming_text_reserved_len + 1
+
+ finish = torch.zeros(input_ids.shape[0], device=input_ids.device).bool()
+
+ temperature = temperature.unsqueeze(0).expand(input_ids.shape[0], -1).contiguous().view(-1, 1)
+
+ progress = input_ids.shape[1]
+
+ # Pre-allocate input_ids, shape is [batch_size=1, max_possible_seq_len, self.num_vqs]
+ input_ids_buf = torch.zeros(
+ input_ids.shape[0], # batch_size
+ progress + max_new_token, # max_possible_seq_len = input_ids.shape[1] + max_new_token
+ input_ids.shape[2], # self.num_vqs
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+
+ # Copy existing `input_ids` to `input_ids_buf`
+ input_ids_buf.narrow(1, 0, progress).copy_(input_ids)
+
+ del input_ids
+ input_ids = input_ids_buf.narrow(1, 0, progress)
+
+ pbar: Optional[tqdm] = None
+ if show_tqdm:
+ pbar = tqdm(
+ total=max_new_token,
+ desc="code",
+ bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}(max) [{elapsed}, {rate_fmt}{postfix}]",
+ )
+
+ condition_length = 1 + self.num_spk_embs * self.use_speaker_embedding + self.streaming_text_reserved_len + 1
+
+ for i in range(max_new_token):
+ # Prepare generation inputs
+ audio_bos = False
+
+ # If this is the first audio token, the case is SPECIAL
+ if progress == condition_length:
+ audio_bos = True
+
+ assert progress == (
+ past_key_values[0][0].shape[2] + 1
+ ) # If you are using according to the guidelines, this should be passed.
+
+ if audio_bos:
+ # Generate the first token, activate the model with `self.audio_bos_token_id`, the model will predict a new audio token. This is a special case because without the `audio bos token`, it is impossible to generate the first audio token in our streaming setting.
+ narrowed_input_ids = torch.tensor([[self.audio_bos_token_id]], dtype=torch.long, device=self.device)
+ inputs_embeds = self.emb_text(narrowed_input_ids)
+ del narrowed_input_ids
+ else:
+ # Generate the following audio tokens, it is applicable to all other cases, including second and the following calling of `generate`.
+ narrowed_input_ids = input_ids.narrow(dim=1, start=input_ids.shape[1] - 1, length=1)
+ code_emb = [self.emb_code[i](narrowed_input_ids[:, :, i]) for i in range(self.num_vq)]
+ inputs_embeds = torch.stack(code_emb, 3).sum(3)
+
+ position_ids = torch.tensor(
+ [past_key_values[0][0].shape[2] + 1], dtype=torch.long, device=self.device
+ ).unsqueeze(0)
+
+ cache_position = position_ids.clone()
+
+ # Make causal mask
+ causal_mask = make_streaming_chunk_mask_generation(
+ inputs_embeds=inputs_embeds,
+ past_seen_tokens=past_key_values[0][0].shape[2],
+ streaming_tts_text_mask=streaming_tts_text_mask,
+ streaming_reserved_length=self.streaming_text_reserved_len,
+ streaming_text_chunk_size=self.streaming_text_chunk_size,
+ )
+
+ # Model forward
+ outputs: BaseModelOutputWithPast = self.model(
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=True,
+ output_attentions=False,
+ cache_position=cache_position,
+ )
+
+ del position_ids
+ del inputs_embeds
+ del cache_position
+ del causal_mask
+
+ hidden_states = outputs.last_hidden_state
+ past_key_values = outputs.past_key_values
+
+ with P.cached():
+ logits = torch.empty(
+ hidden_states.size(0),
+ hidden_states.size(1),
+ self.num_audio_tokens,
+ self.num_vq,
+ dtype=torch.float,
+ device=self.device,
+ )
+ for num_vq_iter in range(self.num_vq):
+ x: torch.Tensor = self.head_code[num_vq_iter](hidden_states)
+ logits[..., num_vq_iter] = x
+ del x
+
+ del hidden_states
+
+ # logits = logits[:, -1].float()
+ logits = logits.narrow(1, -1, 1).squeeze_(1).float()
+
+ # logits = rearrange(logits, "b c n -> (b n) c")
+ logits = logits.permute(0, 2, 1)
+ logits = logits.reshape(-1, logits.size(2))
+ # logits_token = rearrange(input_ids[:, start_idx:], "b c n -> (b n) c")
+ input_ids_sliced = input_ids.narrow(
+ 1,
+ start_idx,
+ input_ids.size(1) - start_idx,
+ ).permute(0, 2, 1)
+ logits_token = input_ids_sliced.reshape(
+ input_ids_sliced.size(0) * input_ids_sliced.size(1),
+ -1,
+ ).to(self.device)
+ del input_ids_sliced
+
+ logits /= temperature
+
+ if not audio_bos:
+ for logitsProcessors in logits_processors:
+ logits = logitsProcessors(logits_token, logits)
+ if not audio_bos:
+ for logitsWarpers in logits_warpers:
+ logits = logitsWarpers(logits_token, logits)
+
+ del logits_token
+
+ if i < min_new_token:
+ logits[:, eos_token] = -torch.inf
+
+ if force_no_stop:
+ logits[:, eos_token] = -torch.inf
+
+ scores = F.softmax(logits, dim=-1)
+
+ del logits
+ idx_next = torch.multinomial(scores, num_samples=1) # .to(finish.device)
+
+ del scores
+
+ # idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq)
+ idx_next = idx_next.view(-1, self.num_vq)
+ finish_or = idx_next.eq(eos_token).any(1)
+ finish.logical_or_(finish_or)
+
+ del finish_or
+ # Store new `token` into `input_ids_buf`
+ input_ids_buf.narrow(1, progress, 1).copy_(idx_next.unsqueeze_(1))
+
+ if i == 0 and finish.any():
+ # raise Exception
+ break
+
+ del idx_next
+ progress += 1
+ input_ids = input_ids_buf.narrow(1, 0, progress)
+
+ if finish.all():
+ break
+
+ if pbar is not None:
+ pbar.update(1)
+
+ if pbar is not None:
+ pbar.close()
+
+ if not finish.all():
+ if show_tqdm:
+ logger.info(f"incomplete result. hit max_new_token: {max_new_token}")
+
+ del input_ids_buf
+
+ if finish.all():
+ # the last may contains eos token
+ genrated_input_ids = input_ids[:, condition_length:-1, :]
+ else:
+ # there is no eos token
+ genrated_input_ids = input_ids[:, condition_length:, :]
+
+ return ConditionalChatTTSGenerationOutput(
+ new_ids=genrated_input_ids,
+ audio_input_ids=input_ids, # for update purpose
+ past_key_values=past_key_values, # for update purpose
+ finished=finish.all(),
+ )
+
+ @torch.inference_mode()
+ def decode_to_mel_specs(
+ self,
+ result_list: List[torch.Tensor],
+ ):
+ """Decode discrete audio codes to mel spectrograms.
+
+ Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/core.py`
+
+ Args:
+ result_list (List[torch.Tensor]): Audio codes output from `generate`.
+
+ Returns:
+ torch.Tensor: Mel spectrograms.
+ """
+
+ decoder = self.dvae
+ max_x_len = -1
+ if len(result_list) == 0:
+ return np.array([], dtype=np.float32)
+ for result in result_list:
+ if result.size(0) > max_x_len:
+ max_x_len = result.size(0)
+ batch_result = torch.zeros(
+ (len(result_list), result_list[0].size(1), max_x_len),
+ dtype=result_list[0].dtype,
+ device=result_list[0].device,
+ )
+ for i in range(len(result_list)):
+ src = result_list[i]
+ batch_result[i].narrow(1, 0, src.size(0)).copy_(src.permute(1, 0))
+ del src
+
+ mel_specs = decoder(batch_result)
+ del batch_result
+ return mel_specs
+
+
+# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/processors.py`
+def gen_logits(
+ num_code: int,
+ top_P=0.7,
+ top_K=20,
+ repetition_penalty=1.0,
+):
+ logits_warpers = []
+ if top_P is not None:
+ logits_warpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
+ if top_K is not None:
+ logits_warpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
+
+ logits_processors = []
+ if repetition_penalty is not None and repetition_penalty != 1:
+ logits_processors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(repetition_penalty, num_code, 16))
+
+ return logits_warpers, logits_processors
+
+
+# Copy and modified from transformers.models.llama.modeling_llama.LlamaForCausalLM.prepare_inputs_for_generation
+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