3260 lines
136 KiB
Python
3260 lines
136 KiB
Python
# 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 as a helpful assistant.
|
||
"audio_roleplay": Roleplay voice-only model, 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] == "<EFBFBD>":
|
||
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
|