You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
PaddleSpeech/paddlespeech/t2s/models/CosyVoice/model.py

597 lines
24 KiB

import os
import threading
import time
import uuid
from contextlib import nullcontext
from typing import Generator
import numpy as np
import paddle
# from cosyvoice.utils.common import TrtContextWrapper, fade_in_out
# from cosyvoice.utils.file_utils import *
# from cosyvoice.utils.file_utils import convert_onnx_to_trt, export_cosyvoice2_vllm
class CosyVoiceModel:
def __init__(
self,
llm: paddle.nn.Layer,
flow: paddle.nn.Layer,
hift: paddle.nn.Layer,
fp16: bool = False,
):
self.device = device2str(
"cuda" if paddle.device.cuda.device_count() >= 1 else "cpu"
)
self.llm = llm
self.flow = flow
self.hift = hift
self.fp16 = fp16
if self.fp16 is True:
self.llm.half()
self.flow.half()
self.token_min_hop_len = 2 * self.flow.input_frame_rate
self.token_max_hop_len = 4 * self.flow.input_frame_rate
self.token_overlap_len = 20
self.mel_overlap_len = int(
self.token_overlap_len / self.flow.input_frame_rate * 22050 / 256
)
self.mel_window = np.hamming(2 * self.mel_overlap_len)
self.mel_cache_len = 20
self.source_cache_len = int(self.mel_cache_len * 256)
self.speech_window = np.hamming(2 * self.source_cache_len)
self.stream_scale_factor = 1
assert (
self.stream_scale_factor >= 1
), "stream_scale_factor should be greater than 1, change it according to your actual rtf"
self.llm_context = (
paddle.device.stream_guard(
paddle.device.Stream(device=device2str(self.device))
)
if paddle.device.cuda.device_count() >= 1
else nullcontext()
)
self.lock = threading.Lock()
self.tts_speech_token_dict = {}
self.llm_end_dict = {}
self.mel_overlap_dict = {}
self.flow_cache_dict = {}
self.hift_cache_dict = {}
def load(self, llm_model, flow_model, hift_model):
self.llm.set_state_dict(state_dict=paddle.load(path=str(llm_model)))
self.llm.to(self.device).eval()
self.flow.set_state_dict(state_dict=paddle.load(path=str(flow_model)))
self.flow.to(self.device).eval()
hift_state_dict = {
k.replace("generator.", ""): v
for k, v in paddle.load(path=str(hift_model)).items()
}
self.hift.set_state_dict(state_dict=hift_state_dict)
self.hift.to(self.device).eval()
def load_jit(self, llm_text_encoder_model, llm_llm_model, flow_encoder_model):
llm_text_encoder = torch.jit.load(
llm_text_encoder_model, map_location=self.device
)
self.llm.text_encoder = llm_text_encoder
llm_llm = torch.jit.load(llm_llm_model, map_location=self.device)
self.llm.llm = llm_llm
flow_encoder = torch.jit.load(flow_encoder_model, map_location=self.device)
self.flow.encoder = flow_encoder
def load_trt(
self,
flow_decoder_estimator_model,
flow_decoder_onnx_model,
trt_concurrent,
fp16,
):
assert paddle.device.cuda.device_count() >= 1, "tensorrt only supports gpu!"
if (
not os.path.exists(flow_decoder_estimator_model)
or os.path.getsize(flow_decoder_estimator_model) == 0
):
convert_onnx_to_trt(
flow_decoder_estimator_model,
self.get_trt_kwargs(),
flow_decoder_onnx_model,
fp16,
)
del self.flow.decoder.estimator
import tensorrt as trt
with open(flow_decoder_estimator_model, "rb") as f:
estimator_engine = trt.Runtime(
trt.Logger(trt.Logger.INFO)
).deserialize_cuda_engine(f.read())
assert estimator_engine is not None, "failed to load trt {}".format(
flow_decoder_estimator_model
)
self.flow.decoder.estimator = TrtContextWrapper(
estimator_engine, trt_concurrent=trt_concurrent, device=self.device
)
def get_trt_kwargs(self):
min_shape = [(2, 80, 4), (2, 1, 4), (2, 80, 4), (2, 80, 4)]
opt_shape = [(2, 80, 500), (2, 1, 500), (2, 80, 500), (2, 80, 500)]
max_shape = [(2, 80, 3000), (2, 1, 3000), (2, 80, 3000), (2, 80, 3000)]
input_names = ["x", "mask", "mu", "cond"]
return {
"min_shape": min_shape,
"opt_shape": opt_shape,
"max_shape": max_shape,
"input_names": input_names,
}
def llm_job(self, text, prompt_text, llm_prompt_speech_token, llm_embedding, uuid):
with self.llm_context, paddle.amp.auto_cast(
enable=self.fp16 is True and hasattr(self.llm, "vllm") is False
):
if isinstance(text, Generator):
assert isinstance(self, CosyVoice2Model) and not hasattr(
self.llm, "vllm"
), "streaming input text is only implemented for CosyVoice2 and do not support vllm!"
for i in self.llm.inference_bistream(
text=text,
prompt_text=prompt_text.to(self.device),
prompt_text_len=paddle.tensor(
[prompt_text.shape[1]], dtype=paddle.int32
).to(self.device),
prompt_speech_token=llm_prompt_speech_token.to(self.device),
prompt_speech_token_len=paddle.tensor(
[llm_prompt_speech_token.shape[1]], dtype=paddle.int32
).to(self.device),
embedding=llm_embedding.to(self.device),
):
self.tts_speech_token_dict[uuid].append(i)
else:
for i in self.llm.inference(
text=text.to(self.device),
text_len=paddle.tensor([text.shape[1]], dtype=paddle.int32).to(
self.device
),
prompt_text=prompt_text.to(self.device),
prompt_text_len=paddle.tensor(
[prompt_text.shape[1]], dtype=paddle.int32
).to(self.device),
prompt_speech_token=llm_prompt_speech_token.to(self.device),
prompt_speech_token_len=paddle.tensor(
[llm_prompt_speech_token.shape[1]], dtype=paddle.int32
).to(self.device),
embedding=llm_embedding.to(self.device),
uuid=uuid,
):
self.tts_speech_token_dict[uuid].append(i)
self.llm_end_dict[uuid] = True
def vc_job(self, source_speech_token, uuid):
self.tts_speech_token_dict[uuid] = source_speech_token.flatten().tolist()
self.llm_end_dict[uuid] = True
def token2wav(
self,
token,
prompt_token,
prompt_feat,
embedding,
uuid,
finalize=False,
speed=1.0,
):
with paddle.amp.auto_cast(enable=self.fp16):
tts_mel, self.flow_cache_dict[uuid] = self.flow.inference(
token=token.to(self.device),
token_len=paddle.tensor([token.shape[1]], dtype=paddle.int32).to(
self.device
),
prompt_token=prompt_token.to(self.device),
prompt_token_len=paddle.tensor(
[prompt_token.shape[1]], dtype=paddle.int32
).to(self.device),
prompt_feat=prompt_feat.to(self.device),
prompt_feat_len=paddle.tensor(
[prompt_feat.shape[1]], dtype=paddle.int32
).to(self.device),
embedding=embedding.to(self.device),
flow_cache=self.flow_cache_dict[uuid],
)
if self.mel_overlap_dict[uuid].shape[2] != 0:
tts_mel = fade_in_out(tts_mel, self.mel_overlap_dict[uuid], self.mel_window)
if self.hift_cache_dict[uuid] is not None:
hift_cache_mel, hift_cache_source = (
self.hift_cache_dict[uuid]["mel"],
self.hift_cache_dict[uuid]["source"],
)
tts_mel = paddle.cat([hift_cache_mel, tts_mel], dim=2)
else:
hift_cache_source = paddle.zeros([1, 1, 0])
if finalize is False:
self.mel_overlap_dict[uuid] = tts_mel[:, :, -self.mel_overlap_len :]
tts_mel = tts_mel[:, :, : -self.mel_overlap_len]
tts_speech, tts_source = self.hift.inference(
speech_feat=tts_mel, cache_source=hift_cache_source
)
if self.hift_cache_dict[uuid] is not None:
tts_speech = fade_in_out(
tts_speech, self.hift_cache_dict[uuid]["speech"], self.speech_window
)
self.hift_cache_dict[uuid] = {
"mel": tts_mel[:, :, -self.mel_cache_len :],
"source": tts_source[:, :, -self.source_cache_len :],
"speech": tts_speech[:, -self.source_cache_len :],
}
tts_speech = tts_speech[:, : -self.source_cache_len]
else:
if speed != 1.0:
assert (
self.hift_cache_dict[uuid] is None
), "speed change only support non-stream inference mode"
tts_mel = paddle.nn.functional.interpolate(
x=tts_mel, size=int(tts_mel.shape[2] / speed), mode="linear"
)
tts_speech, tts_source = self.hift.inference(
speech_feat=tts_mel, cache_source=hift_cache_source
)
if self.hift_cache_dict[uuid] is not None:
tts_speech = fade_in_out(
tts_speech, self.hift_cache_dict[uuid]["speech"], self.speech_window
)
return tts_speech
def tts(
self,
text=paddle.zeros([1, 0], dtype=paddle.int32),
flow_embedding=paddle.zeros([0, 192]),
llm_embedding=paddle.zeros([0, 192]),
prompt_text=paddle.zeros([1, 0], dtype=paddle.int32),
llm_prompt_speech_token=paddle.zeros([1, 0], dtype=paddle.int32),
flow_prompt_speech_token=paddle.zeros([1, 0], dtype=paddle.int32),
prompt_speech_feat=paddle.zeros([1, 0, 80]),
source_speech_token=paddle.zeros([1, 0], dtype=paddle.int32),
stream=False,
speed=1.0,
**kwargs
):
this_uuid = str(uuid.uuid1())
with self.lock:
self.tts_speech_token_dict[this_uuid], self.llm_end_dict[this_uuid] = (
[],
False,
)
self.hift_cache_dict[this_uuid] = None
self.mel_overlap_dict[this_uuid] = paddle.zeros([1, 80, 0])
self.flow_cache_dict[this_uuid] = paddle.zeros([1, 80, 0, 2])
if source_speech_token.shape[1] == 0:
p = threading.Thread(
target=self.llm_job,
args=(
text,
prompt_text,
llm_prompt_speech_token,
llm_embedding,
this_uuid,
),
)
else:
p = threading.Thread(
target=self.vc_job, args=(source_speech_token, this_uuid)
)
"""Not Support auto convert *.start, please judge whether it is Pytorch API and convert by yourself"""
p.start()
if stream is True:
token_hop_len = self.token_min_hop_len
while True:
time.sleep(0.1)
if (
len(self.tts_speech_token_dict[this_uuid])
>= token_hop_len + self.token_overlap_len
):
this_tts_speech_token = paddle.tensor(
self.tts_speech_token_dict[this_uuid][
: token_hop_len + self.token_overlap_len
]
).unsqueeze(dim=0)
this_tts_speech = self.token2wav(
token=this_tts_speech_token,
prompt_token=flow_prompt_speech_token,
prompt_feat=prompt_speech_feat,
embedding=flow_embedding,
uuid=this_uuid,
finalize=False,
)
yield {"tts_speech": this_tts_speech.cpu()}
with self.lock:
self.tts_speech_token_dict[
this_uuid
] = self.tts_speech_token_dict[this_uuid][token_hop_len:]
token_hop_len = min(
self.token_max_hop_len,
int(token_hop_len * self.stream_scale_factor),
)
if (
self.llm_end_dict[this_uuid] is True
and len(self.tts_speech_token_dict[this_uuid])
< token_hop_len + self.token_overlap_len
):
break
p.join()
this_tts_speech_token = paddle.tensor(
self.tts_speech_token_dict[this_uuid]
).unsqueeze(dim=0)
this_tts_speech = self.token2wav(
token=this_tts_speech_token,
prompt_token=flow_prompt_speech_token,
prompt_feat=prompt_speech_feat,
embedding=flow_embedding,
uuid=this_uuid,
finalize=True,
)
yield {"tts_speech": this_tts_speech.cpu()}
else:
p.join()
this_tts_speech_token = paddle.tensor(
self.tts_speech_token_dict[this_uuid]
).unsqueeze(dim=0)
this_tts_speech = self.token2wav(
token=this_tts_speech_token,
prompt_token=flow_prompt_speech_token,
prompt_feat=prompt_speech_feat,
embedding=flow_embedding,
uuid=this_uuid,
finalize=True,
speed=speed,
)
yield {"tts_speech": this_tts_speech.cpu()}
with self.lock:
self.tts_speech_token_dict.pop(this_uuid)
self.llm_end_dict.pop(this_uuid)
self.mel_overlap_dict.pop(this_uuid)
self.hift_cache_dict.pop(this_uuid)
self.flow_cache_dict.pop(this_uuid)
if paddle.device.cuda.device_count() >= 1:
paddle.device.cuda.empty_cache()
paddle.device.current_stream().synchronize()
class CosyVoice2Model(CosyVoiceModel):
def __init__(
self,
llm: paddle.nn.Layer,
flow: paddle.nn.Layer,
hift: paddle.nn.Layer,
fp16: bool = False,
):
self.device = device2str(
"cuda" if paddle.device.cuda.device_count() >= 1 else "cpu"
)
self.llm = llm
self.flow = flow
self.hift = hift
self.fp16 = fp16
if self.fp16 is True:
self.llm.half()
self.flow.half()
self.token_hop_len = 25
self.mel_cache_len = 8
self.source_cache_len = int(self.mel_cache_len * 480)
self.speech_window = np.hamming(2 * self.source_cache_len)
self.llm_context = (
paddle.device.stream_guard(
paddle.device.Stream(device=device2str(self.device))
)
if paddle.device.cuda.device_count() >= 1
else nullcontext()
)
self.lock = threading.Lock()
self.tts_speech_token_dict = {}
self.llm_end_dict = {}
self.hift_cache_dict = {}
def load_jit(self, flow_encoder_model):
flow_encoder = torch.jit.load(flow_encoder_model, map_location=self.device)
self.flow.encoder = flow_encoder
def load_vllm(self, model_dir):
export_cosyvoice2_vllm(self.llm, model_dir, self.device)
from vllm import EngineArgs, LLMEngine
engine_args = EngineArgs(
model=model_dir,
skip_tokenizer_init=True,
enable_prompt_embeds=True,
gpu_memory_utilization=0.2,
)
self.llm.vllm = LLMEngine.from_engine_args(engine_args)
self.llm.lock = threading.Lock()
del self.llm.llm.model.model.layers
def token2wav(
self,
token,
prompt_token,
prompt_feat,
embedding,
token_offset,
uuid,
stream=False,
finalize=False,
speed=1.0,
):
with paddle.amp.auto_cast(enable=self.fp16):
tts_mel, _ = self.flow.inference(
token=token.to(self.device),
token_len=paddle.tensor([token.shape[1]], dtype=paddle.int32).to(
self.device
),
prompt_token=prompt_token.to(self.device),
prompt_token_len=paddle.tensor(
[prompt_token.shape[1]], dtype=paddle.int32
).to(self.device),
prompt_feat=prompt_feat.to(self.device),
prompt_feat_len=paddle.tensor(
[prompt_feat.shape[1]], dtype=paddle.int32
).to(self.device),
embedding=embedding.to(self.device),
streaming=stream,
finalize=finalize,
)
tts_mel = tts_mel[:, :, token_offset * self.flow.token_mel_ratio :]
if self.hift_cache_dict[uuid] is not None:
hift_cache_mel, hift_cache_source = (
self.hift_cache_dict[uuid]["mel"],
self.hift_cache_dict[uuid]["source"],
)
tts_mel = paddle.cat([hift_cache_mel, tts_mel], dim=2)
else:
hift_cache_source = paddle.zeros([1, 1, 0])
if finalize is False:
tts_speech, tts_source = self.hift.inference(
speech_feat=tts_mel, cache_source=hift_cache_source
)
if self.hift_cache_dict[uuid] is not None:
tts_speech = fade_in_out(
tts_speech, self.hift_cache_dict[uuid]["speech"], self.speech_window
)
self.hift_cache_dict[uuid] = {
"mel": tts_mel[:, :, -self.mel_cache_len :],
"source": tts_source[:, :, -self.source_cache_len :],
"speech": tts_speech[:, -self.source_cache_len :],
}
tts_speech = tts_speech[:, : -self.source_cache_len]
else:
if speed != 1.0:
assert (
self.hift_cache_dict[uuid] is None
), "speed change only support non-stream inference mode"
tts_mel = paddle.nn.functional.interpolate(
x=tts_mel, size=int(tts_mel.shape[2] / speed), mode="linear"
)
tts_speech, tts_source = self.hift.inference(
speech_feat=tts_mel, cache_source=hift_cache_source
)
if self.hift_cache_dict[uuid] is not None:
tts_speech = fade_in_out(
tts_speech, self.hift_cache_dict[uuid]["speech"], self.speech_window
)
return tts_speech
def tts(
self,
text=paddle.zeros([1, 0], dtype=paddle.int32),
flow_embedding=paddle.zeros([0, 192]),
llm_embedding=paddle.zeros([0, 192]),
prompt_text=paddle.zeros([1, 0], dtype=paddle.int32),
llm_prompt_speech_token=paddle.zeros([1, 0], dtype=paddle.int32),
flow_prompt_speech_token=paddle.zeros([1, 0], dtype=paddle.int32),
prompt_speech_feat=paddle.zeros([1, 0, 80]),
source_speech_token=paddle.zeros([1, 0], dtype=paddle.int32),
stream=False,
speed=1.0,
**kwargs
):
this_uuid = str(uuid.uuid1())
with self.lock:
self.tts_speech_token_dict[this_uuid], self.llm_end_dict[this_uuid] = (
[],
False,
)
self.hift_cache_dict[this_uuid] = None
if source_speech_token.shape[1] == 0:
p = threading.Thread(
target=self.llm_job,
args=(
text,
prompt_text,
llm_prompt_speech_token,
llm_embedding,
this_uuid,
),
)
else:
p = threading.Thread(
target=self.vc_job, args=(source_speech_token, this_uuid)
)
"""Not Support auto convert *.start, please judge whether it is Pytorch API and convert by yourself"""
p.start()
if stream is True:
token_offset = 0
prompt_token_pad = int(
np.ceil(flow_prompt_speech_token.shape[1] / self.token_hop_len)
* self.token_hop_len
- flow_prompt_speech_token.shape[1]
)
while True:
time.sleep(0.1)
this_token_hop_len = (
self.token_hop_len + prompt_token_pad
if token_offset == 0
else self.token_hop_len
)
if (
len(self.tts_speech_token_dict[this_uuid]) - token_offset
>= this_token_hop_len + self.flow.pre_lookahead_len
):
this_tts_speech_token = paddle.tensor(
self.tts_speech_token_dict[this_uuid][
: token_offset
+ this_token_hop_len
+ self.flow.pre_lookahead_len
]
).unsqueeze(dim=0)
this_tts_speech = self.token2wav(
token=this_tts_speech_token,
prompt_token=flow_prompt_speech_token,
prompt_feat=prompt_speech_feat,
embedding=flow_embedding,
token_offset=token_offset,
uuid=this_uuid,
stream=stream,
finalize=False,
)
token_offset += this_token_hop_len
yield {"tts_speech": this_tts_speech.cpu()}
if (
self.llm_end_dict[this_uuid] is True
and len(self.tts_speech_token_dict[this_uuid]) - token_offset
< this_token_hop_len + self.flow.pre_lookahead_len
):
break
p.join()
this_tts_speech_token = paddle.tensor(
self.tts_speech_token_dict[this_uuid]
).unsqueeze(dim=0)
this_tts_speech = self.token2wav(
token=this_tts_speech_token,
prompt_token=flow_prompt_speech_token,
prompt_feat=prompt_speech_feat,
embedding=flow_embedding,
token_offset=token_offset,
uuid=this_uuid,
finalize=True,
)
yield {"tts_speech": this_tts_speech.cpu()}
else:
p.join()
this_tts_speech_token = paddle.tensor(
self.tts_speech_token_dict[this_uuid]
).unsqueeze(dim=0)
this_tts_speech = self.token2wav(
token=this_tts_speech_token,
prompt_token=flow_prompt_speech_token,
prompt_feat=prompt_speech_feat,
embedding=flow_embedding,
token_offset=0,
uuid=this_uuid,
finalize=True,
speed=speed,
)
yield {"tts_speech": this_tts_speech.cpu()}
with self.lock:
self.tts_speech_token_dict.pop(this_uuid)
self.llm_end_dict.pop(this_uuid)
self.hift_cache_dict.pop(this_uuid)
if paddle.device.cuda.device_count() >= 1:
paddle.device.cuda.empty_cache()
paddle.device.current_stream().synchronize()