From 8f0b28b86183a63a9b17d02f7aeaded2d826adce Mon Sep 17 00:00:00 2001 From: root Date: Mon, 16 Mar 2026 13:59:37 +0800 Subject: [PATCH 1/5] add cosyvoice3 --- runtime/triton_trtllm/client_grpc.py | 1 + runtime/triton_trtllm/client_http.py | 7 +- runtime/triton_trtllm/infer_cosyvoice3.py | 512 ++++++++++++++++++ .../audio_tokenizer/1/model.py | 90 +++ .../audio_tokenizer/config.pbtxt | 53 ++ .../cosyvoice3/1/model.py | 492 +++++++++++++++++ .../cosyvoice3/config.pbtxt | 73 +++ .../speaker_embedding/1/model.py | 146 +++++ .../speaker_embedding/config.pbtxt | 48 ++ .../token2wav/1/model.py | 200 +++++++ .../token2wav/config.pbtxt | 71 +++ .../model_repo_cosyvoice3/vocoder/1/model.py | 69 +++ .../vocoder/config.pbtxt | 40 ++ runtime/triton_trtllm/run_cosyvoice3.sh | 252 +++++++++ .../scripts/convert_cosyvoice3_to_hf.py | 373 +++++++++++++ runtime/triton_trtllm/token2wav_cosyvoice3.py | 414 ++++++++++++++ 16 files changed, 2838 insertions(+), 3 deletions(-) create mode 100644 runtime/triton_trtllm/infer_cosyvoice3.py create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/1/model.py create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/config.pbtxt create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/1/model.py create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/config.pbtxt create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/1/model.py create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/config.pbtxt create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/1/model.py create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/config.pbtxt create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/1/model.py create mode 100644 runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/config.pbtxt create mode 100644 runtime/triton_trtllm/run_cosyvoice3.sh create mode 100644 runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py create mode 100644 runtime/triton_trtllm/token2wav_cosyvoice3.py diff --git a/runtime/triton_trtllm/client_grpc.py b/runtime/triton_trtllm/client_grpc.py index 1ceccb2..440904f 100644 --- a/runtime/triton_trtllm/client_grpc.py +++ b/runtime/triton_trtllm/client_grpc.py @@ -281,6 +281,7 @@ def get_args(): choices=[ "f5_tts", "spark_tts", + "cosyvoice3", "cosyvoice2", "cosyvoice2_dit"], help="triton model_repo module name to request", diff --git a/runtime/triton_trtllm/client_http.py b/runtime/triton_trtllm/client_http.py index 2c30a7a..9c15a6c 100644 --- a/runtime/triton_trtllm/client_http.py +++ b/runtime/triton_trtllm/client_http.py @@ -37,14 +37,14 @@ def get_args(): parser.add_argument( "--server-url", type=str, - default="localhost:8000", + default="localhost:18000", help="Address of the server", ) parser.add_argument( "--reference-audio", type=str, - default="../../example/prompt_audio.wav", + default="./prompt_audio.wav", help="Path to a single audio file. It can't be specified at the same time with --manifest-dir", ) @@ -65,9 +65,10 @@ def get_args(): parser.add_argument( "--model-name", type=str, - default="spark_tts", + default="cosyvoice3", choices=[ "f5_tts", + "cosyvoice3", "spark_tts", "cosyvoice2"], help="triton model_repo module name to request", diff --git a/runtime/triton_trtllm/infer_cosyvoice3.py b/runtime/triton_trtllm/infer_cosyvoice3.py new file mode 100644 index 0000000..69b1a83 --- /dev/null +++ b/runtime/triton_trtllm/infer_cosyvoice3.py @@ -0,0 +1,512 @@ +""" Example Usage + CUDA_VISIBLE_DEVICES=0 \ + python3 infer_cosyvoice3_token2wav.py \ + --output-dir $output_dir \ + --llm-model-name-or-path $huggingface_model_local_dir \ + --token2wav-path $token2wav_model_dir \ + --backend $backend \ + --batch-size $batch_size --token2wav-batch-size $token2wav_batch_size \ + --engine-dir $trt_engines_dir \ + --split-name ${dataset} || exit 1 +""" +import argparse +import json +import os +import time +import asyncio + +import torch +import torchaudio +import s3tokenizer +import soundfile as sf +import requests +import httpx +from transformers import AutoTokenizer +from datasets import load_dataset +from torch.utils.data import DataLoader +from functools import partial +from tqdm import tqdm + +from token2wav_cosyvoice3 import CosyVoice3_Token2Wav + +try: + torch.multiprocessing.set_start_method("spawn") +except RuntimeError: + pass + + +async def send_request_async(client, url, payload): + response = await client.post(url, json=payload, timeout=None) + response.raise_for_status() + response_json = response.json() + return response_json['choices'][0]['message']['content'] + + +async def send_batch_requests_async(api_base, model_name, chats, temperature, top_p, top_k): + async with httpx.AsyncClient() as client: + tasks = [] + for chat in chats: + payload = { + "model": model_name, + "messages": chat, + "max_tokens": 2048, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "repetition_penalty": 1.1, + "stop": ["<|eos1|>", "<|eos|>"], + "stream": False, + } + tasks.append(send_request_async(client, api_base, payload)) + return await asyncio.gather(*tasks) + + +def extract_speech_ids(speech_tokens_str): + """Extract speech IDs from token strings like <|s_23456|>""" + speech_ids = [] + for token_str in speech_tokens_str: + if token_str.startswith('<|s_') and token_str.endswith('|>'): + num_str = token_str[4:-2] + num = int(num_str) + speech_ids.append(num) + else: + print(f"Unexpected token: {token_str}") + return speech_ids + + +def convert_cosy3_tokens_to_speech_id_str(cosy3_tokens): + """Convert CosyVoice3 tokens to speech IDs string like <|s_23456|>""" + if hasattr(cosy3_tokens, 'cpu'): + cosy3_tokens = cosy3_tokens.cpu().numpy().tolist() + speech_id_str = "" + for token in cosy3_tokens: + speech_id_str += f"<|s_{token}|>" + return speech_id_str + + +def get_args(): + parser = argparse.ArgumentParser(description="Speech generation using LLM + CosyVoice3") + parser.add_argument( + "--split-name", type=str, default="wenetspeech4tts", + help="huggingface dataset split name", + ) + parser.add_argument( + "--output-dir", required=True, type=str, help="dir to save result", + ) + parser.add_argument( + "--batch-size", default=1, type=int, + help="batch size (per-device) for LLM inference", + ) + parser.add_argument( + "--token2wav-batch-size", default=1, type=int, + help="batch size (per-device) for token2wav inference", + ) + parser.add_argument( + "--num-workers", type=int, default=0, help="workers for dataloader", + ) + parser.add_argument( + "--prefetch", type=int, default=None, help="prefetch for dataloader", + ) + parser.add_argument( + "--llm-model-name-or-path", required=True, type=str, + help="CosyVoice3 HF LLM path (e.g. ./hf_cosyvoice3_llm)", + ) + parser.add_argument( + "--token2wav-path", required=True, type=str, + help="CosyVoice3 model path (e.g. /workspace_yuekai/HF/Fun-CosyVoice3-0.5B-2512)", + ) + parser.add_argument( + "--enable-trt", action="store_true", + help="Enable TensorRT for flow decoder estimator", + ) + parser.add_argument( + "--streaming", action="store_true", + help="Enable streaming for flow decoder estimator", + ) + parser.add_argument( + "--top-p", type=float, default=0.95, help="top p for sampling", + ) + parser.add_argument( + "--temperature", type=float, default=0.8, help="temperature for sampling", + ) + parser.add_argument( + "--top-k", type=int, default=15, help="top k for sampling", + ) + parser.add_argument( + "--backend", type=str, default="hf", + choices=["hf", "trtllm", "vllm", "trtllm-serve"], + help="Backend to use for LLM inference", + ) + parser.add_argument( + "--engine-dir", type=str, default=None, + help="TensorRT-LLM engine directory (required when backend is 'trtllm')", + ) + parser.add_argument( + "--kv-cache-free-gpu-memory-fraction", type=float, default=0.6, + help="Fraction of GPU memory to free for KV cache (TensorRT-LLM only)", + ) + parser.add_argument( + "--openai-api-base", type=str, + default="http://localhost:8000/v1/chat/completions", + help="OpenAI API base URL (for trtllm-serve backend)", + ) + parser.add_argument( + "--openai-model-name", type=str, default="trt_engines_bfloat16", + help="Model name to use with OpenAI API (for trtllm-serve backend)", + ) + parser.add_argument( + "--epoch", type=int, default=1, help="Epoch to run", + ) + return parser.parse_args() + + +def data_collator(batch, tokenizer, s3_tokenizer): + """Data collator: extracts cosy3 tokens from prompt_audio using v3 s3 tokenizer.""" + device = s3_tokenizer.device if s3_tokenizer is not None else torch.device("cpu") + target_sample_rate = 16000 + + input_ids_list, prompt_audio_list, prompt_text_list = [], [], [] + mels, prompt_audio_cosy3tokens_list, full_text_list = [], [], [] + chat_list = [] + + for item in batch: + prompt_text, target_text = item["prompt_text"], item["target_text"] + prompt_text_list.append(prompt_text) + full_text = 'You are a helpful assistant.<|endofprompt|>' + prompt_text + target_text + full_text_list.append(full_text) + + # Get prompt audio (convert to 16kHz for s3 tokenizer) + ref_audio = torch.from_numpy(item["prompt_audio"]["array"]).float().unsqueeze(0) + ref_sr = item["prompt_audio"]["sampling_rate"] + if ref_sr != target_sample_rate: + ref_audio = torchaudio.transforms.Resample(ref_sr, target_sample_rate)(ref_audio) + prompt_audio_list.append(ref_audio) + + # Extract cosy3 tokens from prompt_audio using v3 s3 tokenizer + mels.append(s3tokenizer.log_mel_spectrogram(ref_audio.squeeze(0))) + + # Batch tokenization with v3 tokenizer + if len(mels) > 0: + mels_padded, mels_lens = s3tokenizer.padding(mels) + codes, codes_lens = s3_tokenizer.quantize(mels_padded.to(device), mels_lens.to(device)) + for i in range(len(codes)): + prompt_audio_cosy3tokens_list.append(codes[i, :codes_lens[i].item()]) + + # Build LLM inputs + for i, prompt_audio_cosy3tokens in enumerate(prompt_audio_cosy3tokens_list): + prompt_audio_cosy3_id_str = convert_cosy3_tokens_to_speech_id_str( + prompt_audio_cosy3tokens) + chat = [ + {"role": "user", "content": full_text_list[i]}, + {"role": "assistant", "content": prompt_audio_cosy3_id_str} + ] + chat_list.append(chat) + input_ids = tokenizer.apply_chat_template( + chat, tokenize=True, return_tensors='pt', continue_final_message=True) + input_ids_list.append(input_ids.squeeze(0)) + + ids = [item["id"] for item in batch] + + return { + "input_ids": input_ids_list, + "ids": ids, + "prompt_text": prompt_text_list, + "prompt_audio_list": prompt_audio_list, + "chat_list": chat_list, + } + + +def main(args): + os.makedirs(args.output_dir, exist_ok=True) + + assert torch.cuda.is_available() + local_rank = 0 + device = torch.device(f"cuda:{local_rank}") + + tokenizer = AutoTokenizer.from_pretrained(args.llm_model_name_or_path) + + if args.backend == "hf": + model = AutoModelForCausalLM.from_pretrained(args.llm_model_name_or_path) + model.eval() + model.to(device) + runner = None + elif args.backend == "trtllm": + if args.engine_dir is None: + raise ValueError("--engine-dir is required when backend is 'trtllm'") + runtime_rank = tensorrt_llm.mpi_rank() + model = None + runner_kwargs = dict( + engine_dir=args.engine_dir, + rank=runtime_rank, + max_output_len=2048, + enable_context_fmha_fp32_acc=False, + max_batch_size=args.batch_size, + max_input_len=512, + kv_cache_free_gpu_memory_fraction=args.kv_cache_free_gpu_memory_fraction, + cuda_graph_mode=False, + gather_generation_logits=False, + ) + runner = ModelRunnerCpp.from_dir(**runner_kwargs) + elif args.backend == "vllm": + model = LLM(model=args.llm_model_name_or_path, gpu_memory_utilization=0.4) + runner = None + elif args.backend == "trtllm-serve": + model = None + runner = None + else: + raise ValueError(f"Unsupported backend: {args.backend}") + + token2wav_model = CosyVoice3_Token2Wav( + model_dir=args.token2wav_path, enable_trt=args.enable_trt, device_id=local_rank, streaming=args.streaming + ) + + # Load v3 s3 tokenizer for prompt audio tokenization in data_collator + s3_tokenizer = s3tokenizer.load_model( + f"{args.token2wav_path}/speech_tokenizer_v3.onnx" + ).to(device).eval() + + dataset = load_dataset( + "yuekai/seed_tts_cosy2", + split=args.split_name, + trust_remote_code=True, + ) + + dataloader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=False, + num_workers=args.num_workers, + prefetch_factor=args.prefetch, + collate_fn=partial(data_collator, tokenizer=tokenizer, s3_tokenizer=s3_tokenizer), + ) + + for epoch in range(args.epoch): + print(f"Running epoch {epoch}") + total_llm_time = 0 + total_token2wav_time = 0 + total_data_load_time = 0 + total_llm_post_processing_time = 0 + total_audio_save_time = 0 + total_audio_samples = 0 + start_time = time.time() + + progress_bar = tqdm(total=len(dataset), desc="Processing", unit="wavs") + + last_batch_end_time = time.time() + for batch in dataloader: + data_loaded_time = time.time() + total_data_load_time += data_loaded_time - last_batch_end_time + + with torch.no_grad(): + llm_start_time = time.time() + + if args.backend == "hf": + input_ids_list = batch["input_ids"] + if len(input_ids_list) == 1: + input_ids = input_ids_list[0].unsqueeze(0) + attention_mask = torch.ones_like(input_ids) + else: + max_len = max([len(ids) for ids in input_ids_list]) + input_ids_list_new = [ + torch.cat([ids, torch.full((max_len - len(ids),), tokenizer.pad_token_id)]) + for ids in input_ids_list + ] + input_ids = torch.stack(input_ids_list_new) + attention_mask = torch.zeros_like(input_ids) + for i in range(len(input_ids_list)): + attention_mask[i, :len(input_ids_list[i])] = 1 + + outputs = model.generate( + input_ids=input_ids.to(device), + attention_mask=attention_mask.to(device), + max_new_tokens=2048, + do_sample=True, + top_p=args.top_p, + temperature=args.temperature, + repetition_penalty=1.1, + top_k=args.top_k, + ) + torch.cuda.synchronize() + + elif args.backend == "trtllm": + batch_input_ids = list(batch["input_ids"]) + input_lengths = [x.size(0) for x in batch_input_ids] + + end_id = tokenizer.convert_tokens_to_ids("<|eos1|>") if "<|eos1|>" in tokenizer.get_vocab() else tokenizer.eos_token_id + outputs = runner.generate( + batch_input_ids=batch_input_ids, + max_new_tokens=2048, + end_id=end_id, + pad_id=end_id, + temperature=args.temperature, + top_k=args.top_k, + top_p=args.top_p, + repetition_penalty=1.1, + num_return_sequences=1, + streaming=False, + output_sequence_lengths=True, + output_generation_logits=False, + return_dict=True, + return_all_generated_tokens=False + ) + torch.cuda.synchronize() + output_ids, sequence_lengths = outputs["output_ids"], outputs["sequence_lengths"] + num_output_sents, num_beams, _ = output_ids.size() + assert num_beams == 1 + batch_size = len(batch["input_ids"]) + num_return_sequences = num_output_sents // batch_size + assert num_return_sequences == 1 + outputs = [] + for i in range(batch_size * num_return_sequences): + batch_idx = i // num_return_sequences + output_begin = input_lengths[batch_idx] + output_end = sequence_lengths[i][0] + outputs_i = output_ids[i][0][:output_end].tolist() + outputs.append(outputs_i) + + elif args.backend == "vllm": + input_ids_list = [ids.tolist() for ids in batch["input_ids"]] + sampling_params = SamplingParams( + temperature=args.temperature, + top_p=args.top_p, + top_k=args.top_k, + repetition_penalty=1.1, + max_tokens=2048, + ) + outputs = model.generate(prompt_token_ids=input_ids_list, sampling_params=sampling_params) + for j, output in enumerate(outputs): + outputs[j] = input_ids_list[j] + output.outputs[0].token_ids + + elif args.backend == "trtllm-serve": + if args.batch_size > 1: + outputs = asyncio.run(send_batch_requests_async( + args.openai_api_base, + args.openai_model_name, + batch["chat_list"], + args.temperature, + args.top_p, + args.top_k, + )) + else: + outputs = [] + for chat in batch["chat_list"]: + payload = { + "model": args.openai_model_name, + "messages": chat, + "max_tokens": 2048, + "temperature": args.temperature, + "top_p": args.top_p, + "top_k": args.top_k, + "repetition_penalty": 1.1, + "stop": ["<|eos1|>", "<|eos|>"], + "stream": False, + } + response = requests.post(args.openai_api_base, json=payload) + response.raise_for_status() + response_json = response.json() + generated_content = response_json['choices'][0]['message']['content'] + outputs.append(generated_content) + + llm_end_time = time.time() + total_llm_time += (llm_end_time - llm_start_time) + + items_for_token_2wav = [] + for i in range(len(batch["ids"])): + llm_post_processing_start_time = time.time() + if args.backend == "trtllm-serve": + speech_tokens_str = outputs[i].strip().split('><') + if len(speech_tokens_str) > 1: + speech_tokens_str = [ + t if t.startswith('<') else '<' + t for t in speech_tokens_str + ] + speech_tokens_str = [ + t if t.endswith('>') else t + '>' for t in speech_tokens_str + ] + speech_ids = extract_speech_ids(speech_tokens_str) + else: + input_length = len(batch["input_ids"][i]) + generated_ids = outputs[i][input_length:] + speech_tokens_str = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) + speech_ids = extract_speech_ids(speech_tokens_str) + print(i, speech_ids[:10], "...", f"total={len(speech_ids)}") + if len(speech_ids) == 0: + print(f"Warning: No speech tokens generated for sample {batch['ids'][i]}, skipping") + llm_post_processing_end_time = time.time() + total_llm_post_processing_time += llm_post_processing_end_time - llm_post_processing_start_time + continue + + current_prompt_audio = batch["prompt_audio_list"][i] + + llm_post_processing_end_time = time.time() + total_llm_post_processing_time += llm_post_processing_end_time - llm_post_processing_start_time + + items_for_token_2wav.append({ + "speech_ids": speech_ids, + "prompt_audio": current_prompt_audio.squeeze(0), + "id": batch["ids"][i] + }) + + for i in range(0, len(items_for_token_2wav), args.token2wav_batch_size): + t2w_batch = items_for_token_2wav[i:i + args.token2wav_batch_size] + if not t2w_batch: + continue + + t2w_speech_tokens = [item["speech_ids"] for item in t2w_batch] + t2w_prompt_audios = [item["prompt_audio"] for item in t2w_batch] + t2w_sample_rates = [16000] * len(t2w_batch) + + token2wav_start_time = time.time() + generated_wavs = token2wav_model( + t2w_speech_tokens, t2w_prompt_audios, t2w_sample_rates, + streaming=args.streaming, + ) + token2wav_end_time = time.time() + total_token2wav_time += (token2wav_end_time - token2wav_start_time) + + audio_save_start_time = time.time() + for j, audio_hat in enumerate(generated_wavs): + wav = audio_hat.squeeze().cpu().numpy() + total_audio_samples += len(wav) + sf.write(f"{args.output_dir}/{t2w_batch[j]['id']}.wav", wav, 24000) + print(f"Generated audio for sample {t2w_batch[j]['id']} with {len(t2w_speech_tokens[j])} tokens") + audio_save_end_time = time.time() + total_audio_save_time += audio_save_end_time - audio_save_start_time + + progress_bar.update(len(batch["ids"])) + last_batch_end_time = time.time() + + progress_bar.close() + end_time = time.time() + total_audio_duration_seconds = total_audio_samples / 24000 + + log_file_path = os.path.join(args.output_dir, "log.txt") + with open(log_file_path, 'w') as f: + log_data = { + "args": vars(args), + "data_load_time_seconds": total_data_load_time, + "llm_time_seconds": total_llm_time, + "llm_post_processing_time_seconds": total_llm_post_processing_time, + "token2wav_time_seconds": total_token2wav_time, + "audio_save_time_seconds": total_audio_save_time, + "total_audio_duration_seconds": total_audio_duration_seconds, + "pipeline_time_seconds": end_time - start_time, + } + print(log_data) + f.write(json.dumps(log_data, indent=4)) + print(f"Metrics logged to {log_file_path}") + + +if __name__ == "__main__": + args = get_args() + if args.backend == "vllm": + from vllm import LLM, SamplingParams + elif args.backend == "trtllm": + import tensorrt_llm + from tensorrt_llm.runtime import ModelRunnerCpp + elif args.backend == "hf": + from transformers import AutoModelForCausalLM + elif args.backend == "trtllm-serve": + pass + else: + raise ValueError(f"Unsupported backend: {args.backend}") + main(args) diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/1/model.py b/runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/1/model.py new file mode 100644 index 0000000..f70a715 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/1/model.py @@ -0,0 +1,90 @@ +# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import json +import torch +from torch.utils.dlpack import to_dlpack + +import triton_python_backend_utils as pb_utils + +import os +import numpy as np +import s3tokenizer +torch.set_num_threads(1) +# ORIGINAL_VOCAB_SIZE = 151924 + + +class TritonPythonModel: + """Triton Python model for audio tokenization. + + This model takes reference audio input and extracts semantic tokens + using s3tokenizer. + """ + + def initialize(self, args): + """Initialize the model. + + Args: + args: Dictionary containing model configuration + """ + # Parse model parameters + parameters = json.loads(args['model_config'])['parameters'] + model_params = {k: v["string_value"] for k, v in parameters.items()} + + self.device = torch.device("cuda") + model_path = os.path.join(model_params["model_dir"], "speech_tokenizer_v3.onnx") + self.audio_tokenizer = s3tokenizer.load_model(model_path).to(self.device) + + def execute(self, requests): + """Execute inference on the batched requests.""" + mels = [] + + # Process each request in batch + for req_idx, request in enumerate(requests): + # Extract input tensors + wav_array = pb_utils.get_input_tensor_by_name( + request, "reference_wav").as_numpy() + wav_len = pb_utils.get_input_tensor_by_name( + request, "reference_wav_len").as_numpy().item() + + wav_array = torch.from_numpy(wav_array).to(self.device) + # Prepare inputs + wav = wav_array[:, :wav_len].squeeze(0) + mel = s3tokenizer.log_mel_spectrogram(wav) + mels.append(mel) + + mels, mels_lens = s3tokenizer.padding(mels) + codes, codes_lens = self.audio_tokenizer.quantize(mels.to(self.device), mels_lens.to(self.device)) + + responses = [] + for i in range(len(requests)): + prompt_speech_tokens = codes[i, :codes_lens[i].item()] + prompt_speech_tokens_tensor = pb_utils.Tensor.from_dlpack( + "prompt_speech_tokens", to_dlpack(prompt_speech_tokens)) + inference_response = pb_utils.InferenceResponse( + output_tensors=[prompt_speech_tokens_tensor]) + responses.append(inference_response) + + return responses diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/config.pbtxt b/runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/config.pbtxt new file mode 100644 index 0000000..4bb972c --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/audio_tokenizer/config.pbtxt @@ -0,0 +1,53 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +name: "audio_tokenizer" +backend: "python" +max_batch_size: ${triton_max_batch_size} +dynamic_batching { + max_queue_delay_microseconds: ${max_queue_delay_microseconds} +} +parameters [ + { + key: "model_dir", + value: {string_value:"${model_dir}"} + } +] + +input [ + { + name: "reference_wav" + data_type: TYPE_FP32 + dims: [-1] + }, + { + name: "reference_wav_len" + data_type: TYPE_INT32 + dims: [1] + } +] +output [ + { + name: "prompt_speech_tokens" + data_type: TYPE_INT32 + dims: [-1] + } +] + +instance_group [ + { + count: 1 + kind: KIND_CPU + } +] \ No newline at end of file diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/1/model.py b/runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/1/model.py new file mode 100644 index 0000000..519adc8 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/1/model.py @@ -0,0 +1,492 @@ +import json +import re +import time +import asyncio + +import numpy as np +import torch +from torch.utils.dlpack import to_dlpack +import triton_python_backend_utils as pb_utils + +import httpx +import torchaudio +from functools import partial +from matcha.utils.audio import mel_spectrogram as matcha_mel_spectrogram + + +torch.set_num_threads(1) + +# CosyVoice3 mel params: fmax=None (Nyquist), center=False +mel_spectrogram = partial(matcha_mel_spectrogram, + n_fft=1920, num_mels=80, sampling_rate=24000, + hop_size=480, win_size=1920, fmin=0, fmax=None, center=False) + + +def parse_speech_token_string(response_text): + """Parse speech tokens from string like '<|s_123|><|s_456|>' into list of int IDs.""" + speech_tokens = response_text.strip().split('><') + if len(speech_tokens) > 1: + speech_tokens = ['<' + t if not t.startswith('<') else t for t in speech_tokens] + speech_tokens = [t + '>' if not t.endswith('>') else t for t in speech_tokens] + speech_ids = [] + for token_str in speech_tokens: + match = re.match(r'<\|s_(\d+)\|>', token_str) + if match: + speech_ids.append(int(match.group(1))) + return speech_ids + + +class TritonPythonModel: + """CosyVoice3 BLS orchestrator for Triton Inference Server. + + Orchestrates: audio_tokenizer, speaker_embedding, remote LLM (httpx), + token2wav (flow-only), and vocoder (CausalHiFTGenerator). + Supports both streaming (decoupled) and offline (non-decoupled) modes. + """ + + def initialize(self, args): + self.logger = pb_utils.Logger + self.model_config = json.loads(args['model_config']) + parameters = self.model_config['parameters'] + model_params = {k: v["string_value"] for k, v in parameters.items()} + + self.device = torch.device("cuda") + self.decoupled = pb_utils.using_decoupled_model_transaction_policy(self.model_config) + + # Streaming config + self.token_frame_rate = 25 + self.flow_pre_lookahead_len = 3 + self.token_hop_len = 15 + self.token_mel_ratio = 2 + self.dynamic_chunk_strategy = model_params.get("dynamic_chunk_strategy", "exponential") + self.logger.log_info(f"CosyVoice3 BLS initialized, decoupled={self.decoupled}, " + f"chunk_strategy={self.dynamic_chunk_strategy}") + + # HTTP client for remote LLM (trtllm-serve default port: 8000) + self.http_client = httpx.AsyncClient() + self.api_base = model_params.get("llm_api_base", "http://localhost:8000/v1/chat/completions") + + # Speaker cache to avoid redundant audio_tokenizer/speaker_embedding calls + self.speaker_cache = {} + + def _convert_speech_tokens_to_str(self, speech_tokens): + """Convert speech token IDs tensor/list to string like '<|s_N|>'.""" + if isinstance(speech_tokens, torch.Tensor): + speech_tokens = speech_tokens.cpu().numpy().flatten().tolist() + return "".join(f"<|s_{int(tid)}|>" for tid in speech_tokens) + + def _extract_speech_feat(self, speech): + """Extract mel spectrogram from 24kHz speech for flow prompt.""" + speech_feat = mel_spectrogram(speech).squeeze(dim=0).transpose(0, 1) + speech_feat = speech_feat.unsqueeze(dim=0).to(self.device) + return speech_feat + + async def forward_llm_streaming(self, target_text, reference_text, prompt_speech_tokens): + """Async generator: stream LLM tokens via httpx SSE.""" + full_text = f"{reference_text}{target_text}" + prompt_speech_tokens_str = self._convert_speech_tokens_to_str(prompt_speech_tokens) + + chat = [ + {"role": "user", "content": full_text}, + {"role": "assistant", "content": prompt_speech_tokens_str} + ] + payload = { + "model": "trt_engines_bfloat16", + "messages": chat, + "max_tokens": 750, + "temperature": 0.8, + "top_p": 0.95, + "top_k": 50, + "repetition_penalty": 1.1, + "stop": ["<|eos1|>", "<|eos|>"], + "stream": True, + } + + buffer = "" + async with self.http_client.stream("POST", self.api_base, json=payload, timeout=None) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if line.startswith("data: "): + line_data = line[len("data: "):].strip() + if line_data == "[DONE]": + break + try: + json_data = json.loads(line_data) + content = json_data.get("choices", [{}])[0].get("delta", {}).get("content") + if content: + buffer += content + while True: + match = re.search(r"<\|s_(\d+)\|>", buffer) + if not match: + break + token_num = int(match.group(1)) + # final_id = token_num + ORIGINAL_VOCAB_SIZE + yield token_num + buffer = buffer[match.end():] + except json.JSONDecodeError: + continue + + # Flush remaining tokens + while True: + match = re.search(r"<\|s_(\d+)\|>", buffer) + if not match: + break + token_num = int(match.group(1)) + #final_id = token_num + ORIGINAL_VOCAB_SIZE + yield token_num + buffer = buffer[match.end():] + + async def forward_llm_offline(self, target_text, reference_text, prompt_speech_tokens): + """Non-streaming LLM call, returns all speech token IDs at once.""" + full_text = f"{reference_text}{target_text}" + prompt_speech_tokens_str = self._convert_speech_tokens_to_str(prompt_speech_tokens) + + chat = [ + {"role": "user", "content": full_text}, + {"role": "assistant", "content": prompt_speech_tokens_str} + ] + payload = { + "model": "trt_engines_bfloat16", + "messages": chat, + "max_tokens": 750, + "temperature": 0.8, + "top_p": 0.95, + "top_k": 50, + "repetition_penalty": 1.1, + "stop": ["<|eos1|>", "<|eos|>"], + "stream": False, + } + response = await self.http_client.post(self.api_base, json=payload, timeout=None) + response.raise_for_status() + response_json = response.json() + generated_content = response_json['choices'][0]['message']['content'] + speech_ids = parse_speech_token_string(generated_content) + # return [sid + ORIGINAL_VOCAB_SIZE for sid in speech_ids] + return speech_ids + + def forward_audio_tokenizer(self, wav, wav_len): + """BLS call to audio_tokenizer.""" + inference_request = pb_utils.InferenceRequest( + model_name='audio_tokenizer', + requested_output_names=['prompt_speech_tokens'], + inputs=[wav, wav_len] + ) + inference_response = inference_request.exec() + if inference_response.has_error(): + raise pb_utils.TritonModelException(inference_response.error().message()) + prompt_speech_tokens = pb_utils.get_output_tensor_by_name( + inference_response, 'prompt_speech_tokens') + return torch.utils.dlpack.from_dlpack(prompt_speech_tokens.to_dlpack()).cpu() + + def forward_speaker_embedding(self, wav): + """BLS call to speaker_embedding.""" + inference_request = pb_utils.InferenceRequest( + model_name='speaker_embedding', + requested_output_names=['prompt_spk_embedding'], + inputs=[pb_utils.Tensor.from_dlpack("reference_wav", to_dlpack(wav))] + ) + inference_response = inference_request.exec() + if inference_response.has_error(): + raise pb_utils.TritonModelException(inference_response.error().message()) + prompt_spk_embedding = pb_utils.get_output_tensor_by_name( + inference_response, 'prompt_spk_embedding') + return torch.utils.dlpack.from_dlpack(prompt_spk_embedding.to_dlpack()) + + async def forward_token2wav(self, target_speech_tokens, prompt_speech_tokens, + prompt_speech_feat, prompt_spk_embedding, + request_id, token_offset=None, finalize=True, + priority=100): + """Async BLS call to token2wav (flow-only). Returns mel tensor.""" + target_tokens_pb = pb_utils.Tensor.from_dlpack( + "target_speech_tokens", to_dlpack(target_speech_tokens)) + prompt_tokens_pb = pb_utils.Tensor.from_dlpack( + "prompt_speech_tokens", to_dlpack(prompt_speech_tokens)) + prompt_feat_pb = pb_utils.Tensor.from_dlpack( + "prompt_speech_feat", to_dlpack(prompt_speech_feat)) + prompt_emb_pb = pb_utils.Tensor.from_dlpack( + "prompt_spk_embedding", to_dlpack(prompt_spk_embedding)) + + inputs = [target_tokens_pb, prompt_tokens_pb, prompt_feat_pb, prompt_emb_pb] + + if token_offset is not None: + inputs.append(pb_utils.Tensor("token_offset", + np.array([[token_offset]], dtype=np.int32))) + inputs.append(pb_utils.Tensor("finalize", + np.array([[finalize]], dtype=np.bool_))) + + inference_request = pb_utils.InferenceRequest( + model_name='token2wav', + requested_output_names=['mel'], + inputs=inputs, + request_id=request_id, + parameters={"priority": priority}, + ) + + inference_response = await inference_request.async_exec() + if inference_response.has_error(): + raise pb_utils.TritonModelException(inference_response.error().message()) + + mel = pb_utils.get_output_tensor_by_name(inference_response, 'mel') + return torch.utils.dlpack.from_dlpack(mel.to_dlpack()) + + async def forward_vocoder(self, mel, finalize): + """Async BLS call to vocoder. Returns speech tensor.""" + if mel.dim() == 2: + mel = mel.unsqueeze(0) # [80, T] -> [1, 80, T] + mel_pb = pb_utils.Tensor.from_dlpack("mel", to_dlpack(mel.float())) + finalize_pb = pb_utils.Tensor("finalize", + np.array([[finalize]], dtype=np.bool_)) + + inference_request = pb_utils.InferenceRequest( + model_name='vocoder', + requested_output_names=['tts_speech'], + inputs=[mel_pb, finalize_pb], + ) + + inference_response = await inference_request.async_exec() + if inference_response.has_error(): + raise pb_utils.TritonModelException(inference_response.error().message()) + + speech = pb_utils.get_output_tensor_by_name(inference_response, 'tts_speech') + return torch.utils.dlpack.from_dlpack(speech.to_dlpack()).cpu() + + def _prepare_prompt(self, request): + """Extract reference audio, tokenize, compute speaker embedding and mel feat.""" + wav = pb_utils.get_input_tensor_by_name(request, "reference_wav") + wav_len = pb_utils.get_input_tensor_by_name(request, "reference_wav_len") + + reference_text = pb_utils.get_input_tensor_by_name(request, "reference_text") + reference_text = reference_text.as_numpy()[0][0].decode('utf-8') if reference_text is not None else "" + if '<|endofprompt|>' not in reference_text: + reference_text = 'You are a helpful assistant.<|endofprompt|>' + reference_text + + # Check speaker cache + if reference_text in self.speaker_cache: + cached = self.speaker_cache[reference_text] + return (cached['prompt_speech_tokens_for_llm'], cached['prompt_speech_tokens'], + cached['prompt_speech_feat'], cached['prompt_spk_embedding'], reference_text) + + # Audio tokenizer + wav_np = wav.as_numpy() + wav_len_val = wav_len.as_numpy()[0][0] + prompt_speech_tokens = self.forward_audio_tokenizer(wav, wav_len) + prompt_speech_tokens = prompt_speech_tokens.unsqueeze(0) # [1, T] + + # Speaker embedding + wav_tensor = torch.from_numpy(wav_np) + wav_tensor = wav_tensor[:, :wav_len_val] + prompt_spk_embedding = self.forward_speaker_embedding(wav_tensor) + + # Mel extraction at 24kHz with CosyVoice3 params + prompt_speech_resample = torchaudio.transforms.Resample( + orig_freq=16000, new_freq=24000)(wav_tensor) + speech_feat = self._extract_speech_feat(prompt_speech_resample) + + # Keep full tokens for LLM prefill (untruncated) + prompt_speech_tokens_for_llm = prompt_speech_tokens.clone() + + # Align prompt speech feat and tokens to 2:1 ratio (for flow model only) + orig_feat_len = speech_feat.shape[1] + orig_token_len = prompt_speech_tokens.shape[-1] + token_len = min(int(speech_feat.shape[1] / 2), prompt_speech_tokens.shape[-1]) + prompt_speech_feat = speech_feat[:, :2 * token_len].contiguous().half() + prompt_speech_tokens = prompt_speech_tokens[:, :token_len].contiguous() + + # Cache + self.speaker_cache[reference_text] = { + 'prompt_speech_tokens_for_llm': prompt_speech_tokens_for_llm, + 'prompt_speech_tokens': prompt_speech_tokens, + 'prompt_speech_feat': prompt_speech_feat, + 'prompt_spk_embedding': prompt_spk_embedding, + } + + return prompt_speech_tokens_for_llm, prompt_speech_tokens, prompt_speech_feat, prompt_spk_embedding, reference_text + + async def _process_request_streaming(self, request): + """Process a single request in streaming (decoupled) mode.""" + request_id = request.request_id() + response_sender = request.get_response_sender() + + try: + prompt_speech_tokens_for_llm, prompt_speech_tokens, prompt_speech_feat, \ + prompt_spk_embedding, reference_text = self._prepare_prompt(request) + + target_text = pb_utils.get_input_tensor_by_name(request, "target_text").as_numpy() + target_text = target_text[0][0].decode('utf-8') + + semantic_token_ids_arr = [] + token_offset = 0 + chunk_index = 0 + this_token_hop_len = self.token_hop_len + accumulated_mel = None + speech_offset = 0 + start_time = time.time() + + async for generated_id in self.forward_llm_streaming( + target_text=target_text, + reference_text=reference_text, + prompt_speech_tokens=prompt_speech_tokens_for_llm, + ): + semantic_token_ids_arr.append(generated_id) + + while True: + pending_num = len(semantic_token_ids_arr) - token_offset + if pending_num < this_token_hop_len + self.flow_pre_lookahead_len: + break + + # Prepare tokens for this chunk + end_idx = token_offset + this_token_hop_len + self.flow_pre_lookahead_len + this_tokens = torch.tensor( + semantic_token_ids_arr[:end_idx] + ).unsqueeze(0).to(torch.int32).to(self.device) + + # Call token2wav (flow-only) -> mel_chunk + mel_chunk = await self.forward_token2wav( + this_tokens, prompt_speech_tokens, + prompt_speech_feat, prompt_spk_embedding, + request_id, token_offset=token_offset, finalize=False, + priority=chunk_index + 1, + ) + + # Accumulate mel + if mel_chunk.dim() == 2: + mel_chunk = mel_chunk.unsqueeze(0) + if accumulated_mel is None: + accumulated_mel = mel_chunk + else: + accumulated_mel = torch.cat([accumulated_mel, mel_chunk], dim=2) + + # Call vocoder + speech = await self.forward_vocoder(accumulated_mel, finalize=False) + + # Extract new speech + new_speech = speech[:, speech_offset:] + speech_offset += new_speech.shape[1] + + if new_speech.shape[1] > 0: + audio_tensor = pb_utils.Tensor.from_dlpack( + "waveform", to_dlpack(new_speech)) + inference_response = pb_utils.InferenceResponse( + output_tensors=[audio_tensor]) + response_sender.send(inference_response) + + token_offset += this_token_hop_len + + # Dynamic chunk strategy + if self.dynamic_chunk_strategy == "exponential": + this_token_hop_len = self.token_frame_rate * (2 ** chunk_index) + elif self.dynamic_chunk_strategy == "time_based": + cost_time = time.time() - start_time + duration = token_offset / self.token_frame_rate + if chunk_index > 0 and cost_time > 0: + avg_chunk_time = cost_time / (chunk_index + 1) + if avg_chunk_time > 0: + multiples = (duration - cost_time) / avg_chunk_time + next_pending = len(semantic_token_ids_arr) - token_offset + if multiples > 4: + this_token_hop_len = (next_pending // self.token_hop_len + 1) * self.token_hop_len + elif multiples > 2: + this_token_hop_len = (next_pending // self.token_hop_len) * self.token_hop_len + else: + this_token_hop_len = self.token_hop_len + this_token_hop_len = max(self.token_hop_len, this_token_hop_len) + + chunk_index += 1 + + # Final chunk with remaining tokens + if len(semantic_token_ids_arr) > 0: + remaining_tokens = torch.tensor( + semantic_token_ids_arr + ).unsqueeze(0).to(torch.int32).to(self.device) + + mel_chunk = await self.forward_token2wav( + remaining_tokens, prompt_speech_tokens, + prompt_speech_feat, prompt_spk_embedding, + request_id, token_offset=token_offset, finalize=True, + priority=chunk_index + 1, + ) + + if mel_chunk.dim() == 2: + mel_chunk = mel_chunk.unsqueeze(0) + if accumulated_mel is None: + accumulated_mel = mel_chunk + else: + accumulated_mel = torch.cat([accumulated_mel, mel_chunk], dim=2) + + speech = await self.forward_vocoder(accumulated_mel, finalize=True) + + new_speech = speech[:, speech_offset:] + if new_speech.shape[1] > 0: + audio_tensor = pb_utils.Tensor.from_dlpack( + "waveform", to_dlpack(new_speech)) + inference_response = pb_utils.InferenceResponse( + output_tensors=[audio_tensor]) + response_sender.send(inference_response) + + response_sender.send(flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL) + except Exception as e: + self.logger.log_error(f"Error in streaming request: {e}") + error_response = pb_utils.InferenceResponse( + error=pb_utils.TritonError(str(e))) + response_sender.send(error_response) + response_sender.send(flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL) + + async def _process_request_offline(self, request): + """Process a single request in offline (non-decoupled) mode.""" + request_id = request.request_id() + + prompt_speech_tokens_for_llm, prompt_speech_tokens, prompt_speech_feat, \ + prompt_spk_embedding, reference_text = self._prepare_prompt(request) + + target_text = pb_utils.get_input_tensor_by_name(request, "target_text").as_numpy() + target_text = target_text[0][0].decode('utf-8') + + # Get all speech tokens at once (use full untruncated prompt tokens for LLM) + all_token_ids = await self.forward_llm_offline( + target_text=target_text, + reference_text=reference_text, + prompt_speech_tokens=prompt_speech_tokens_for_llm, + ) + + if len(all_token_ids) == 0: + raise pb_utils.TritonModelException("LLM generated no speech tokens") + + all_tokens = torch.tensor(all_token_ids).unsqueeze(0).to(torch.int32).to(self.device) + + # token2wav (no token_offset, finalize=True) -> full mel + mel = await self.forward_token2wav( + all_tokens, prompt_speech_tokens, + prompt_speech_feat, prompt_spk_embedding, + request_id, + ) + + # vocoder -> full speech + speech = await self.forward_vocoder(mel, finalize=True) + + audio_tensor = pb_utils.Tensor.from_dlpack("waveform", to_dlpack(speech)) + return pb_utils.InferenceResponse(output_tensors=[audio_tensor]) + + async def execute(self, requests): + if self.decoupled: + tasks = [ + asyncio.create_task(self._process_request_streaming(request)) + for request in requests + ] + await asyncio.gather(*tasks) + return None + else: + responses = [] + for request in requests: + try: + response = await self._process_request_offline(request) + responses.append(response) + except Exception as e: + self.logger.log_error(f"Error in offline request: {e}") + responses.append(pb_utils.InferenceResponse( + error=pb_utils.TritonError(str(e)))) + return responses + + def finalize(self): + self.logger.log_info("Finalizing CosyVoice3 BLS model") + if hasattr(self, "http_client"): + asyncio.run(self.http_client.aclose()) diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/config.pbtxt b/runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/config.pbtxt new file mode 100644 index 0000000..fa988b5 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/cosyvoice3/config.pbtxt @@ -0,0 +1,73 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +name: "cosyvoice3" +backend: "python" +max_batch_size: ${triton_max_batch_size} +dynamic_batching { + max_queue_delay_microseconds: ${max_queue_delay_microseconds} +} +model_transaction_policy { + decoupled: ${decoupled_mode} +} +parameters [ + { + key: "llm_tokenizer_dir", + value: {string_value:"${llm_tokenizer_dir}"} + }, + { + key: "model_dir", + value: {string_value:"${model_dir}"} + } +] + +input [ + { + name: "reference_wav" + data_type: TYPE_FP32 + dims: [-1] + optional: true + }, + { + name: "reference_wav_len" + data_type: TYPE_INT32 + dims: [1] + optional: true + }, + { + name: "reference_text" + data_type: TYPE_STRING + dims: [1] + optional: true + }, + { + name: "target_text" + data_type: TYPE_STRING + dims: [1] + } +] +output [ + { + name: "waveform" + data_type: TYPE_FP32 + dims: [ -1 ] + } +] + +instance_group [ + { + count: ${bls_instance_num} + kind: KIND_CPU + } +] \ No newline at end of file diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/1/model.py b/runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/1/model.py new file mode 100644 index 0000000..132ca82 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/1/model.py @@ -0,0 +1,146 @@ +# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import json +import torch +from torch.utils.dlpack import to_dlpack + +import triton_python_backend_utils as pb_utils + +import os +import numpy as np +import torchaudio.compliance.kaldi as kaldi +from cosyvoice.utils.file_utils import convert_onnx_to_trt +from cosyvoice.utils.common import TrtContextWrapper +import onnxruntime + + +class TritonPythonModel: + """Triton Python model for audio tokenization. + + This model takes reference audio input and extracts semantic tokens + using s3tokenizer. + """ + + def initialize(self, args): + """Initialize the model. + + Args: + args: Dictionary containing model configuration + """ + # Parse model parameters + parameters = json.loads(args['model_config'])['parameters'] + model_params = {k: v["string_value"] for k, v in parameters.items()} + + self.device = torch.device("cuda") + + model_dir = model_params["model_dir"] + gpu = "l20" + enable_trt = True + if enable_trt: + self.load_spk_trt(f'{model_dir}/campplus.{gpu}.fp32.trt', + f'{model_dir}/campplus.onnx', + 1, + False) + else: + campplus_model = f'{model_dir}/campplus.onnx' + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + self.spk_model = onnxruntime.InferenceSession(campplus_model, sess_options=option, providers=["CPUExecutionProvider"]) + + def load_spk_trt(self, spk_model, spk_onnx_model, trt_concurrent=1, fp16=True): + if not os.path.exists(spk_model) or os.path.getsize(spk_model) == 0: + trt_kwargs = self.get_spk_trt_kwargs() + convert_onnx_to_trt(spk_model, trt_kwargs, spk_onnx_model, fp16) + import tensorrt as trt + with open(spk_model, 'rb') as f: + spk_engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read()) + assert spk_engine is not None, 'failed to load trt {}'.format(spk_model) + self.spk_model = TrtContextWrapper(spk_engine, trt_concurrent=trt_concurrent, device=self.device) + + def get_spk_trt_kwargs(self): + min_shape = [(1, 4, 80)] + opt_shape = [(1, 500, 80)] + max_shape = [(1, 3000, 80)] + input_names = ["input"] + return {'min_shape': min_shape, 'opt_shape': opt_shape, 'max_shape': max_shape, 'input_names': input_names} + + def _extract_spk_embedding(self, speech): + feat = kaldi.fbank(speech, + num_mel_bins=80, + dither=0, + sample_frequency=16000) + spk_feat = feat - feat.mean(dim=0, keepdim=True) + + if isinstance(self.spk_model, onnxruntime.InferenceSession): + embedding = self.spk_model.run( + None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()} + )[0].flatten().tolist() + embedding = torch.tensor([embedding]).to(self.device) + else: + [spk_model, stream], trt_engine = self.spk_model.acquire_estimator() + # NOTE need to synchronize when switching stream + with torch.cuda.device(self.device): + torch.cuda.current_stream().synchronize() + spk_feat = spk_feat.unsqueeze(dim=0).to(self.device) + batch_size = spk_feat.size(0) + + with stream: + spk_model.set_input_shape('input', (batch_size, spk_feat.size(1), 80)) + embedding = torch.empty((batch_size, 192), device=spk_feat.device) + + data_ptrs = [spk_feat.contiguous().data_ptr(), + embedding.contiguous().data_ptr()] + for i, j in enumerate(data_ptrs): + + spk_model.set_tensor_address(trt_engine.get_tensor_name(i), j) + # run trt engine + assert spk_model.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True + torch.cuda.current_stream().synchronize() + self.spk_model.release_estimator(spk_model, stream) + + return embedding.half() + + def execute(self, requests): + """Execute inference on the batched requests.""" + responses = [] + # Process each request in batch + for req_idx, request in enumerate(requests): + # Extract input tensors + wav_array = pb_utils.get_input_tensor_by_name( + request, "reference_wav").as_numpy() + wav_array = torch.from_numpy(wav_array).to(self.device) + + embedding = self._extract_spk_embedding(wav_array) + + prompt_spk_embedding_tensor = pb_utils.Tensor.from_dlpack( + "prompt_spk_embedding", to_dlpack(embedding)) + inference_response = pb_utils.InferenceResponse( + output_tensors=[prompt_spk_embedding_tensor]) + + responses.append(inference_response) + + return responses diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/config.pbtxt b/runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/config.pbtxt new file mode 100644 index 0000000..fd91bd6 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/speaker_embedding/config.pbtxt @@ -0,0 +1,48 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +name: "speaker_embedding" +backend: "python" +max_batch_size: ${triton_max_batch_size} +dynamic_batching { + max_queue_delay_microseconds: ${max_queue_delay_microseconds} +} +parameters [ + { + key: "model_dir", + value: {string_value:"${model_dir}"} + } +] + +input [ + { + name: "reference_wav" + data_type: TYPE_FP32 + dims: [-1] + } +] +output [ + { + name: "prompt_spk_embedding" + data_type: TYPE_FP16 + dims: [-1] + } +] + +instance_group [ + { + count: 1 + kind: KIND_CPU + } +] \ No newline at end of file diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/1/model.py b/runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/1/model.py new file mode 100644 index 0000000..abfb1f6 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/1/model.py @@ -0,0 +1,200 @@ +import json +import os +import logging +import queue + +import torch +import numpy as np +from torch.utils.dlpack import to_dlpack +import triton_python_backend_utils as pb_utils +from hyperpyyaml import load_hyperpyyaml + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + + +class TrtContextWrapper: + def __init__(self, trt_engine, trt_concurrent=1, device='cuda:0'): + self.trt_context_pool = queue.Queue(maxsize=trt_concurrent) + self.trt_engine = trt_engine + self.device = device + for _ in range(trt_concurrent): + trt_context = trt_engine.create_execution_context() + trt_stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device))) + assert trt_context is not None + self.trt_context_pool.put([trt_context, trt_stream]) + + def acquire_estimator(self): + return self.trt_context_pool.get(), self.trt_engine + + def release_estimator(self, context, stream): + self.trt_context_pool.put([context, stream]) + + +def convert_onnx_to_trt(trt_model, trt_kwargs, onnx_model, fp16, autocast_mode=False): + import tensorrt as trt + logging.info("Converting onnx to trt...") + if autocast_mode: + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + else: + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + trt_logger = trt.Logger(trt.Logger.INFO) + builder = trt.Builder(trt_logger) + network = builder.create_network(network_flags) + parser = trt.OnnxParser(network, trt_logger) + config = builder.create_builder_config() + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 32) + if not autocast_mode and fp16: + config.set_flag(trt.BuilderFlag.FP16) + profile = builder.create_optimization_profile() + with open(onnx_model, "rb") as f: + if not parser.parse(f.read()): + for error in range(parser.num_errors): + print(parser.get_error(error)) + raise ValueError(f'failed to parse {onnx_model}') + for i in range(len(trt_kwargs['input_names'])): + profile.set_shape(trt_kwargs['input_names'][i], + trt_kwargs['min_shape'][i], + trt_kwargs['opt_shape'][i], + trt_kwargs['max_shape'][i]) + if not autocast_mode: + tensor_dtype = trt.DataType.HALF if fp16 else trt.DataType.FLOAT + for i in range(network.num_inputs): + network.get_input(i).dtype = tensor_dtype + for i in range(network.num_outputs): + network.get_output(i).dtype = tensor_dtype + config.add_optimization_profile(profile) + engine_bytes = builder.build_serialized_network(network, config) + with open(trt_model, "wb") as f: + f.write(engine_bytes) + logging.info("Successfully converted onnx to trt") + +torch.set_num_threads(1) + + +class TritonPythonModel: + """Triton Python model for CosyVoice3 token2wav (flow-only, stateless). + + Converts speech tokens to mel spectrogram using the CausalMaskedDiffWithDiT flow model. + """ + + def initialize(self, args): + parameters = json.loads(args['model_config'])['parameters'] + model_params = {k: v["string_value"] for k, v in parameters.items()} + model_dir = model_params["model_dir"] + + self.device = torch.device("cuda") + + # Load flow model from cosyvoice3.yaml + with open(os.path.join(model_dir, 'cosyvoice3.yaml'), 'r') as f: + configs = load_hyperpyyaml(f, overrides={ + 'qwen_pretrain_path': os.path.join(model_dir, 'CosyVoice-BlankEN') + }) + self.flow = configs['flow'] + self.fp16 = True + self.flow.half() + self.flow.load_state_dict( + torch.load(os.path.join(model_dir, 'flow.pt'), + map_location='cpu', weights_only=True), + strict=True + ) + self.flow.to(self.device).eval() + + # TRT acceleration for flow decoder estimator + self.load_trt(model_dir) + + self.token_mel_ratio = self.flow.token_mel_ratio + logger.info(f"Token2wav (flow-only) initialized, token_mel_ratio={self.token_mel_ratio}") + + def load_trt(self, model_dir, trt_concurrent=1): + device_id = torch.cuda.current_device() + onnx_path = os.path.join(model_dir, 'flow.decoder.estimator.autocast_fp16.onnx') + trt_path = os.path.join(model_dir, f'flow.decoder.estimator.autocast_fp16.{device_id}.plan') + + if not os.path.exists(trt_path) or os.path.getsize(trt_path) == 0: + trt_kwargs = self.get_trt_kwargs() + convert_onnx_to_trt(trt_path, trt_kwargs, onnx_path, + fp16=True, autocast_mode=True) + del self.flow.decoder.estimator + import tensorrt as trt + with open(trt_path, 'rb') as f: + estimator_engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read()) + assert estimator_engine is not None, f'failed to load trt {trt_path}' + self.flow.decoder.estimator = TrtContextWrapper( + estimator_engine, trt_concurrent=trt_concurrent, device=str(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 execute(self, requests): + responses = [] + for req_idx, request in enumerate(requests): + target_speech_tokens = pb_utils.get_input_tensor_by_name( + request, "target_speech_tokens") + target_speech_tokens = torch.utils.dlpack.from_dlpack( + target_speech_tokens.to_dlpack()).to(self.device) + if target_speech_tokens.dim() == 1: + target_speech_tokens = target_speech_tokens.unsqueeze(0) + + # Optional inputs + prompt_speech_tokens_pb = pb_utils.get_input_tensor_by_name( + request, "prompt_speech_tokens") + if prompt_speech_tokens_pb is not None: + prompt_speech_tokens = torch.utils.dlpack.from_dlpack( + prompt_speech_tokens_pb.to_dlpack()).to(self.device) + if prompt_speech_tokens.dim() == 1: + prompt_speech_tokens = prompt_speech_tokens.unsqueeze(0) + + prompt_speech_feat = pb_utils.get_input_tensor_by_name( + request, "prompt_speech_feat") + prompt_speech_feat = torch.utils.dlpack.from_dlpack( + prompt_speech_feat.to_dlpack()).to(self.device) + if prompt_speech_feat.dim() == 2: + prompt_speech_feat = prompt_speech_feat.unsqueeze(0) # [T, 80] -> [1, T, 80] + + prompt_spk_embedding = pb_utils.get_input_tensor_by_name( + request, "prompt_spk_embedding") + prompt_spk_embedding = torch.utils.dlpack.from_dlpack( + prompt_spk_embedding.to_dlpack()).to(self.device) + if prompt_spk_embedding.dim() == 1: + prompt_spk_embedding = prompt_spk_embedding.unsqueeze(0) + else: + raise ValueError("prompt_speech_tokens is required for CosyVoice3 token2wav") + + token_offset_pb = pb_utils.get_input_tensor_by_name(request, "token_offset") + finalize_pb = pb_utils.get_input_tensor_by_name(request, "finalize") + + token_offset = token_offset_pb.as_numpy().item() if token_offset_pb is not None else None + finalize = finalize_pb.as_numpy().item() if finalize_pb is not None else True + streaming = not finalize + + with torch.no_grad(), torch.cuda.amp.autocast(self.fp16): + mel, _ = self.flow.inference( + token=target_speech_tokens, + token_len=torch.tensor([target_speech_tokens.shape[1]], dtype=torch.int32).to(self.device), + prompt_token=prompt_speech_tokens, + prompt_token_len=torch.tensor([prompt_speech_tokens.shape[1]], dtype=torch.int32).to(self.device), + prompt_feat=prompt_speech_feat, + prompt_feat_len=torch.tensor([prompt_speech_feat.shape[1]], dtype=torch.int32).to(self.device), + embedding=prompt_spk_embedding, + streaming=streaming, + finalize=finalize, + ) + + # Slice mel from token_offset if provided + if token_offset is not None: + mel = mel[:, :, token_offset * self.token_mel_ratio:] + + # Output mel as [80, T] (squeeze batch dim for Triton) + mel_out = mel.squeeze(0).float() # [80, T] + mel_out = mel_out.cpu() # otherwise, dlpack bug + mel_tensor = pb_utils.Tensor.from_dlpack("mel", to_dlpack(mel_out)) + inference_response = pb_utils.InferenceResponse(output_tensors=[mel_tensor]) + responses.append(inference_response) + + return responses diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/config.pbtxt b/runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/config.pbtxt new file mode 100644 index 0000000..a46926c --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/token2wav/config.pbtxt @@ -0,0 +1,71 @@ +name: "token2wav" +backend: "python" +max_batch_size: ${triton_max_batch_size} + +dynamic_batching { + max_queue_delay_microseconds: ${max_queue_delay_microseconds} + priority_levels: 100 + default_priority_level: 100 +} + +parameters [ + { + key: "model_dir", + value: {string_value:"${model_dir}"} + } +] + +input [ + { + name: "target_speech_tokens" + data_type: TYPE_INT32 + dims: [-1] + }, + { + name: "prompt_speech_tokens" + data_type: TYPE_INT32 + dims: [-1] + optional: true + }, + { + name: "prompt_speech_feat" + data_type: TYPE_FP16 + dims: [-1, 80] + optional: true + }, + { + name: "prompt_spk_embedding" + data_type: TYPE_FP16 + dims: [-1] + optional: true + }, + { + name: "token_offset" + data_type: TYPE_INT32 + dims: [ 1 ] + reshape: { shape: [ ] } + optional: true + }, + { + name: "finalize" + data_type: TYPE_BOOL + dims: [ 1 ] + reshape: { shape: [ ] } + optional: true + } +] +output [ + { + name: "mel" + data_type: TYPE_FP32 + dims: [ 80, -1 ] + } +] + +instance_group [ + { + count: 1 + kind: KIND_GPU + gpus: [ 0 ] + } +] \ No newline at end of file diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/1/model.py b/runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/1/model.py new file mode 100644 index 0000000..987fa49 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/1/model.py @@ -0,0 +1,69 @@ +import json +import os +import logging + +import torch +from torch.utils.dlpack import to_dlpack +import triton_python_backend_utils as pb_utils +from hyperpyyaml import load_hyperpyyaml + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + +torch.set_num_threads(1) + + +class TritonPythonModel: + """Triton Python model for CosyVoice3 vocoder (CausalHiFTGenerator). + + Stateless: converts mel spectrogram to waveform. + CausalHiFTGenerator manages its own internal cache. + """ + + def initialize(self, args): + parameters = json.loads(args['model_config'])['parameters'] + model_params = {k: v["string_value"] for k, v in parameters.items()} + model_dir = model_params["model_dir"] + + self.device = torch.device("cuda") + + # Load CausalHiFTGenerator from cosyvoice3.yaml + with open(os.path.join(model_dir, 'cosyvoice3.yaml'), 'r') as f: + configs = load_hyperpyyaml(f, overrides={ + 'qwen_pretrain_path': os.path.join(model_dir, 'CosyVoice-BlankEN') + }) + self.hift = configs['hift'] + hift_state_dict = { + k.replace('generator.', ''): v + for k, v in torch.load( + os.path.join(model_dir, 'hift.pt'), + map_location='cpu', weights_only=True + ).items() + } + self.hift.load_state_dict(hift_state_dict, strict=True) + self.hift.to(self.device).eval() + logger.info("CausalHiFTGenerator initialized successfully") + + def execute(self, requests): + responses = [] + for req_idx, request in enumerate(requests): + mel = pb_utils.get_input_tensor_by_name(request, "mel") + mel = torch.utils.dlpack.from_dlpack(mel.to_dlpack()).to(self.device) + if mel.dim() == 2: + mel = mel.unsqueeze(0) # [80, T] -> [1, 80, T] + + finalize = pb_utils.get_input_tensor_by_name(request, "finalize").as_numpy().item() + + with torch.no_grad(): + speech, _ = self.hift.inference(speech_feat=mel, finalize=finalize) + + # speech shape: [1, 1, S] or [1, S] depending on hift version + speech = speech.squeeze() # flatten to [S] + + speech_tensor = pb_utils.Tensor.from_dlpack( + "tts_speech", to_dlpack(speech.unsqueeze(0))) # [1, S] for batch dim + inference_response = pb_utils.InferenceResponse( + output_tensors=[speech_tensor]) + responses.append(inference_response) + + return responses diff --git a/runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/config.pbtxt b/runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/config.pbtxt new file mode 100644 index 0000000..5ebb766 --- /dev/null +++ b/runtime/triton_trtllm/model_repo_cosyvoice3/vocoder/config.pbtxt @@ -0,0 +1,40 @@ +name: "vocoder" +backend: "python" +max_batch_size: ${triton_max_batch_size} +dynamic_batching { + max_queue_delay_microseconds: ${max_queue_delay_microseconds} +} +parameters [ + { + key: "model_dir", + value: {string_value:"${model_dir}"} + } +] + +input [ + { + name: "mel" + data_type: TYPE_FP32 + dims: [80, -1] + }, + { + name: "finalize" + data_type: TYPE_BOOL + dims: [ 1 ] + reshape: { shape: [ ] } + } +] +output [ + { + name: "tts_speech" + data_type: TYPE_FP32 + dims: [ -1 ] + } +] + +instance_group [ + { + count: 1 + kind: KIND_CPU + } +] \ No newline at end of file diff --git a/runtime/triton_trtllm/run_cosyvoice3.sh b/runtime/triton_trtllm/run_cosyvoice3.sh new file mode 100644 index 0000000..accdebe --- /dev/null +++ b/runtime/triton_trtllm/run_cosyvoice3.sh @@ -0,0 +1,252 @@ +#!/bin/bash +# Copyright (c) 2026 NVIDIA (authors: Yuekai Zhang) +export CUDA_VISIBLE_DEVICES=0 +# cosyvoice_path=/workspace/CosyVoice +cosyvoice_path=/workspace_yuekai/tts/CosyVoice + +export PYTHONPATH=${cosyvoice_path}:$PYTHONPATH +export PYTHONPATH=${cosyvoice_path}/third_party/Matcha-TTS:$PYTHONPATH + +stage=$1 +stop_stage=$2 + +huggingface_model_local_dir=./hf_cosyvoice3_llm +model_scope_model_local_dir=/workspace_yuekai/HF/Fun-CosyVoice3-0.5B-2512 + +trt_dtype=bfloat16 +trt_weights_dir=./trt_weights_${trt_dtype} +trt_engines_dir=./trt_engines_${trt_dtype} + +model_repo_src=./model_repo_cosyvoice3 +model_repo=./deploy_cosyvoice3 +bls_instance_num=1 + +if [ $stage -le -1 ] && [ $stop_stage -ge -1 ]; then + + echo "Cloning CosyVoice" + git clone --recursive https://github.com/FunAudioLLM/CosyVoice.git $cosyvoice_path + cd $cosyvoice_path + git submodule update --init --recursive + cd runtime/triton_trtllm +fi + +if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then + echo "" + # see https://github.com/nvidia-china-sae/mair-hub/blob/main/rl-tutorial/cosyvoice_llm/pretrained_to_huggingface.py + # huggingface-cli download --local-dir $huggingface_model_local_dir yuekai/cosyvoice2_llm + # modelscope download --model iic/CosyVoice2-0.5B --local_dir $model_scope_model_local_dir + + # pip3 install --upgrade x_transformers s3tokenizer + # pip install -U nvidia-modelopt[all] + python3 scripts/convert_cosyvoice3_to_hf.py \ + --model-dir $model_scope_model_local_dir \ + --output-dir $huggingface_model_local_dir || exit 1 # TODO: output dir should be here + +fi + + +if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then + echo "Converting checkpoint to TensorRT weights" + python3 scripts/convert_checkpoint.py --model_dir $huggingface_model_local_dir \ + --output_dir $trt_weights_dir \ + --dtype $trt_dtype || exit 1 + + echo "Building TensorRT engines" + trtllm-build --checkpoint_dir $trt_weights_dir \ + --output_dir $trt_engines_dir \ + --max_batch_size 64 \ + --max_num_tokens 32768 \ + --gemm_plugin $trt_dtype || exit 1 + + echo "Testing TensorRT engines" + python3 ./scripts/test_llm.py --input_text "你好，请问你叫什么？" \ + --tokenizer_dir $huggingface_model_local_dir \ + --top_k 50 --top_p 0.95 --temperature 0.8 \ + --engine_dir=$trt_engines_dir || exit 1 +fi + +if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then + echo "Creating CosyVoice3 model repository" + rm -rf $model_repo + mkdir -p $model_repo + + # Copy all modules from template source + cp -r ${model_repo_src}/cosyvoice3 $model_repo/ + cp -r ${model_repo_src}/token2wav $model_repo/ + cp -r ${model_repo_src}/vocoder $model_repo/ + cp -r ${model_repo_src}/audio_tokenizer $model_repo/ + cp -r ${model_repo_src}/speaker_embedding $model_repo/ + + MAX_QUEUE_DELAY_MICROSECONDS=0 + MODEL_DIR=$model_scope_model_local_dir + LLM_TOKENIZER_DIR=$huggingface_model_local_dir + BLS_INSTANCE_NUM=$bls_instance_num + TRITON_MAX_BATCH_SIZE=1 + DECOUPLED_MODE=True + + python3 scripts/fill_template.py -i ${model_repo}/cosyvoice3/config.pbtxt model_dir:${MODEL_DIR},bls_instance_num:${BLS_INSTANCE_NUM},llm_tokenizer_dir:${LLM_TOKENIZER_DIR},triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},decoupled_mode:${DECOUPLED_MODE},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS} + python3 scripts/fill_template.py -i ${model_repo}/token2wav/config.pbtxt model_dir:${MODEL_DIR},triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS} + python3 scripts/fill_template.py -i ${model_repo}/vocoder/config.pbtxt model_dir:${MODEL_DIR},triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS} + python3 scripts/fill_template.py -i ${model_repo}/audio_tokenizer/config.pbtxt model_dir:${MODEL_DIR},triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS} + python3 scripts/fill_template.py -i ${model_repo}/speaker_embedding/config.pbtxt model_dir:${MODEL_DIR},triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS} + +fi + +if [ $stage -le 30 ] && [ $stop_stage -ge 30 ]; then + echo "Starting CosyVoice3 Triton server and LLM using trtllm-serve" + CUDA_VISIBLE_DEVICES=0 mpirun -np 1 --allow-run-as-root --oversubscribe trtllm-serve serve --tokenizer $huggingface_model_local_dir $trt_engines_dir --max_batch_size 64 --kv_cache_free_gpu_memory_fraction 0.4 +fi + + +if [ $stage -le 40 ] && [ $stop_stage -ge 40 ]; then + + CUDA_VISIBLE_DEVICES=1 tritonserver --model-repository $model_repo --http-port 18000 --grpc-port 18001 --metrics-port 18002 & +fi + + +if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then + echo "Starting CosyVoice3 Triton server and LLM using trtllm-serve" + CUDA_VISIBLE_DEVICES=0 mpirun -np 1 --allow-run-as-root --oversubscribe trtllm-serve serve --tokenizer $huggingface_model_local_dir $trt_engines_dir --max_batch_size 64 --kv_cache_free_gpu_memory_fraction 0.4 & + CUDA_VISIBLE_DEVICES=0,1,2,3 tritonserver --model-repository $model_repo --http-port 18000 --grpc-port 18001 --metrics-port 18002 & + wait + # Test using curl + # curl http://localhost:8000/v1/chat/completions \ + # -H "Content-Type: application/json" \ + # -d '{ + # "model": "", + # "messages":[{"role": "user", "content": "Where is New York?"}, + # {"role": "assistant", "content": "<|s_1708|><|s_2050|><|s_2159|>"}], + # "max_tokens": 512, + # "temperature": 0.8, + # "top_p": 0.95, + # "top_k": 50, + # "stop": ["<|eos1|>"], + # "repetition_penalty": 1.2, + # "stream": false + # }' +fi + +if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then + echo "Running benchmark client for CosyVoice3" + num_task=4 + mode=offline + mode=streaming + BLS_INSTANCE_NUM=$bls_instance_num + + python3 client_grpc.py \ + --server-addr localhost \ + --server-port 18001 \ + --model-name cosyvoice3 \ + --num-tasks $num_task \ + --mode $mode \ + --huggingface-dataset yuekai/seed_tts_cosy2 \ + --log-dir ./log_cosyvoice3_concurrent_tasks_${num_task}_${mode}_bls_${BLS_INSTANCE_NUM} + +fi + +if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then + echo "stage 5: Offline TTS (Cosyvoice2 LLM + Step-Audio2-mini DiT Token2Wav) inference using a single python script" + + datasets=(wenetspeech4tts) # wenetspeech4tts, test_zh, zero_shot_zh + backend=trtllm # hf, trtllm, vllm, trtllm-serve + + batch_sizes=(16) + token2wav_batch_size=1 + + for batch_size in ${batch_sizes[@]}; do + for dataset in ${datasets[@]}; do + output_dir=./${dataset}_${backend}_llm_batch_size_${batch_size}_token2wav_batch_size_${token2wav_batch_size} + CUDA_VISIBLE_DEVICES=1 \ + python3 offline_inference.py \ + --output-dir $output_dir \ + --llm-model-name-or-path $huggingface_model_local_dir \ + --token2wav-path $step_audio_model_dir/token2wav \ + --backend $backend \ + --batch-size $batch_size --token2wav-batch-size $token2wav_batch_size \ + --engine-dir $trt_engines_dir \ + --split-name ${dataset} || exit 1 + done + done +fi + + + + +if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then + echo "Disaggregated Server: LLM and Token2wav on different GPUs" + echo "Starting LLM server on GPU 0" + export CUDA_VISIBLE_DEVICES=0 + mpirun -np 1 --allow-run-as-root --oversubscribe trtllm-serve serve --tokenizer $huggingface_model_local_dir $trt_engines_dir --max_batch_size 64 --kv_cache_free_gpu_memory_fraction 0.4 & + echo "Starting Token2wav server on GPUs 1-3" + Token2wav_num_gpus=3 + http_port=17000 + grpc_port=18000 + metrics_port=16000 + for i in $(seq 0 $(($Token2wav_num_gpus - 1))); do + echo "Starting server on GPU $i" + http_port=$((http_port + 1)) + grpc_port=$((grpc_port + 1)) + metrics_port=$((metrics_port + 1)) + # Two instances of Token2wav server on the same GPU + CUDA_VISIBLE_DEVICES=$(($i + 1)) tritonserver --model-repository $model_repo --http-port $http_port --grpc-port $grpc_port --metrics-port $metrics_port & + http_port=$((http_port + 1)) + grpc_port=$((grpc_port + 1)) + metrics_port=$((metrics_port + 1)) + CUDA_VISIBLE_DEVICES=$(($i + 1)) tritonserver --model-repository $model_repo --http-port $http_port --grpc-port $grpc_port --metrics-port $metrics_port & + done + wait +fi + +if [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then + echo "Running benchmark client for Disaggregated Server" + per_gpu_instances=2 + mode=streaming + BLS_INSTANCE_NUM=$bls_instance_num + Token2wav_num_gpus=(1 2 3) + concurrent_tasks=(1 2 3 4 5 6) + for n_gpu in ${Token2wav_num_gpus[@]}; do + echo "Test 1 GPU for LLM server and $n_gpu GPUs for Token2wav servers" + for concurrent_task in ${concurrent_tasks[@]}; do + num_instances=$((per_gpu_instances * n_gpu)) + for i in $(seq 1 $num_instances); do + port=$(($i + 18000)) + python3 client_grpc.py \ + --server-addr localhost \ + --server-port $port \ + --model-name cosyvoice2_dit \ + --num-tasks $concurrent_task \ + --mode $mode \ + --huggingface-dataset yuekai/seed_tts_cosy2 \ + --log-dir ./log_disagg_concurrent_tasks_${concurrent_task}_per_instance_total_token2wav_instances_${num_instances}_port_${port} & + done + wait + done + done +fi + +if [ $stage -le 10 ] && [ $stop_stage -ge 10 ]; then + echo "stage 10: Python script CosyVoice3 TTS (LLM + CosyVoice3 Token2Wav) inference" + + datasets=(wenetspeech4tts) # wenetspeech4tts + backend=trtllm-serve # hf, trtllm, vllm, trtllm-serve + + batch_sizes=(1) + token2wav_batch_size=1 + + for batch_size in ${batch_sizes[@]}; do + for dataset in ${datasets[@]}; do + output_dir=./cosyvoice3_${dataset}_${backend}_llm_batch_size_${batch_size}_token2wav_batch_size_${token2wav_batch_size}_streaming_trt + CUDA_VISIBLE_DEVICES=0 \ + python3 infer_cosyvoice3.py \ + --output-dir $output_dir \ + --llm-model-name-or-path $huggingface_model_local_dir \ + --token2wav-path $model_scope_model_local_dir \ + --backend $backend \ + --batch-size $batch_size --token2wav-batch-size $token2wav_batch_size \ + --engine-dir $trt_engines_dir \ + --enable-trt --streaming\ + --epoch 1 \ + --split-name ${dataset} || exit 1 + done + done +fi \ No newline at end of file diff --git a/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py b/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py new file mode 100644 index 0000000..58318d0 --- /dev/null +++ b/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py @@ -0,0 +1,373 @@ +#!/usr/bin/env python3 +# Copyright 2025 CosyVoice3 TRT-LLM Integration +# +# 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. +""" +Конвертация CosyVoice3 LLM в HuggingFace формат с объединёнными embeddings. + +Этот скрипт: +1. Загружает CosyVoice3 модель +2. Расширяет vocab токенизатора с speech токенами +3. Объединяет speech_embedding в embed_tokens Qwen2 +4. Заменяет lm_head на llm_decoder с расширенным vocab +5. Сохраняет модель в HuggingFace формате для TRT-LLM конвертации + +Usage: + python scripts/convert_cosyvoice3_to_hf.py \ + --model-dir pretrained_models/Fun-CosyVoice3-0.5B \ + --output-dir pretrained_models/Fun-CosyVoice3-0.5B/hf_merged + +После этого можно конвертировать в TRT-LLM: + trtllm-build --checkpoint_dir --output_dir ... +""" +import argparse +import os +import sys +import logging + +import torch +from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'third_party/Matcha-TTS')) + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Convert CosyVoice3 to HuggingFace format with merged embeddings") + parser.add_argument( + "--model-dir", + type=str, + default="pretrained_models/Fun-CosyVoice3-0.5B", + help="Path to CosyVoice3 model directory", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Output directory for HuggingFace model (default: /hf_merged)", + ) + parser.add_argument( + "--dtype", + type=str, + default="bfloat16", + choices=["float16", "bfloat16", "float32"], + help="Output dtype for the model", + ) + return parser.parse_args() + + +def load_cosyvoice3_model(model_dir: str): + """Загружает CosyVoice3 модель для извлечения весов.""" + from hyperpyyaml import load_hyperpyyaml + from cosyvoice.utils.class_utils import get_model_type + + hyper_yaml_path = os.path.join(model_dir, 'cosyvoice3.yaml') + hf_llm_dir = os.path.join(model_dir, 'CosyVoice-BlankEN') + + if not os.path.exists(hyper_yaml_path): + raise ValueError(f'{hyper_yaml_path} not found!') + + with open(hyper_yaml_path, 'r') as f: + configs = load_hyperpyyaml( + f, + overrides={'qwen_pretrain_path': hf_llm_dir} + ) + + # Загружаем только LLM + llm = configs['llm'] + llm_weights_path = os.path.join(model_dir, 'llm.pt') + llm.load_state_dict(torch.load(llm_weights_path, map_location='cpu'), strict=True) + llm.eval() + + logger.info(f"Loaded CosyVoice3 LLM from {model_dir}") + + return llm, hf_llm_dir, configs + + +def get_speech_token_size(llm) -> int: + """Определяет размер speech token vocabulary из модели.""" + # CosyVoice3LM имеет: speech_token_size + 200 в llm_decoder + # speech_embedding имеет: speech_token_size + 200 + speech_embedding_size = llm.speech_embedding.num_embeddings + # Вычитаем 200 специальных токенов (sos, eos, task_id, fill, и т.д.) + # Но для безопасности используем полный размер embedding + return speech_embedding_size + + +def convert_cosyvoice3_to_hf( + model_dir: str, + output_dir: str, + dtype: str = "bfloat16", +): + """ + Конвертирует CosyVoice3 LLM в HuggingFace формат с объединёнными embeddings. + + Архитектура объединения: + - embed_tokens[0:original_vocab_size] = оригинальные text embeddings + - embed_tokens[original_vocab_size:original_vocab_size+speech_token_size] = speech_embedding + - lm_head[original_vocab_size:original_vocab_size+speech_token_size] = llm_decoder + + Args: + model_dir: Путь к CosyVoice3 модели + output_dir: Путь для сохранения HF модели + dtype: Тип данных для сохранения + """ + logger.info(f"Loading CosyVoice3 model from {model_dir}") + + # 1. Загружаем CosyVoice3 компоненты + cosyvoice3_llm, hf_llm_dir, configs = load_cosyvoice3_model(model_dir) + + # Извлекаем ключевые компоненты + qwen_model = cosyvoice3_llm.llm.model # Qwen2ForCausalLM + speech_embedding = cosyvoice3_llm.speech_embedding # Embedding для speech токенов + llm_decoder = cosyvoice3_llm.llm_decoder # Linear для декодирования в speech токены + + speech_token_size = get_speech_token_size(cosyvoice3_llm) + logger.info(f"Speech token size: {speech_token_size}") + + # 2. Загружаем tokenizer и добавляем CosyVoice3 text special tokens + speech токены + tokenizer = AutoTokenizer.from_pretrained(hf_llm_dir, trust_remote_code=True) + base_vocab_size = len(tokenizer) + logger.info(f"Base tokenizer vocab size: {base_vocab_size}") + + # IMPORTANT: + # - In CosyVoice3, LLM speech special tokens (sos/eos/task_id/fill) are INSIDE speech_embedding, + # i.e. represented as <|s_6561|>, <|s_6562|>, <|s_6563|>, <|s_6564|>. + # - But text-level special tokens like [cough]/[laughter] MUST exist in tokenizer + # (mirrors `CosyVoice3Tokenizer` from `cosyvoice/tokenizer/tokenizer.py`). + special_tokens = { + 'eos_token': '<|endoftext|>', + 'pad_token': '<|endoftext|>', + 'additional_special_tokens': [ + '<|im_start|>', '<|im_end|>', '<|endofprompt|>', + '[breath]', '', '', '[noise]', + '[laughter]', '[cough]', '[clucking]', '[accent]', + '[quick_breath]', + "", "", + "[hissing]", "[sigh]", "[vocalized-noise]", + "[lipsmack]", "[mn]", "<|endofsystem|>", + # Phoneme tokens (kept consistent with CosyVoice3Tokenizer) + "[AA]", "[AA0]", "[AA1]", "[AA2]", "[AE]", "[AE0]", "[AE1]", "[AE2]", "[AH]", "[AH0]", "[AH1]", "[AH2]", + "[AO]", "[AO0]", "[AO1]", "[AO2]", "[AW]", "[AW0]", "[AW1]", "[AW2]", "[AY]", "[AY0]", "[AY1]", "[AY2]", + "[B]", "[CH]", "[D]", "[DH]", "[EH]", "[EH0]", "[EH1]", "[EH2]", "[ER]", "[ER0]", "[ER1]", "[ER2]", "[EY]", + "[EY0]", "[EY1]", "[EY2]", "[F]", "[G]", "[HH]", "[IH]", "[IH0]", "[IH1]", "[IH2]", "[IY]", "[IY0]", "[IY1]", + "[IY2]", "[JH]", "[K]", "[L]", "[M]", "[N]", "[NG]", "[OW]", "[OW0]", "[OW1]", "[OW2]", "[OY]", "[OY0]", + "[OY1]", "[OY2]", "[P]", "[R]", "[S]", "[SH]", "[T]", "[TH]", "[UH]", "[UH0]", "[UH1]", "[UH2]", "[UW]", + "[UW0]", "[UW1]", "[UW2]", "[V]", "[W]", "[Y]", "[Z]", "[ZH]", + "[a]", "[ai]", "[an]", "[ang]", "[ao]", "[b]", "[c]", "[ch]", "[d]", "[e]", "[ei]", "[en]", "[eng]", "[f]", + "[g]", "[h]", "[i]", "[ian]", "[in]", "[ing]", "[iu]", "[ià]", "[iàn]", "[iàng]", "[iào]", "[iá]", "[ián]", + "[iáng]", "[iáo]", "[iè]", "[ié]", "[iòng]", "[ióng]", "[iù]", "[iú]", "[iā]", "[iān]", "[iāng]", "[iāo]", + "[iē]", "[iě]", "[iōng]", "[iū]", "[iǎ]", "[iǎn]", "[iǎng]", "[iǎo]", "[iǒng]", "[iǔ]", "[j]", "[k]", "[l]", + "[m]", "[n]", "[o]", "[ong]", "[ou]", "[p]", "[q]", "[r]", + "[s]", "[sh]", "[t]", "[u]", "[uang]", "[ue]", + "[un]", "[uo]", "[uà]", "[uài]", "[uàn]", "[uàng]", "[uá]", "[uái]", "[uán]", "[uáng]", "[uè]", "[ué]", "[uì]", + "[uí]", "[uò]", "[uó]", "[uā]", "[uāi]", "[uān]", "[uāng]", "[uē]", "[uě]", "[uī]", "[uō]", "[uǎ]", "[uǎi]", + "[uǎn]", "[uǎng]", "[uǐ]", "[uǒ]", "[vè]", "[w]", "[x]", "[y]", "[z]", "[zh]", "[à]", "[ài]", "[àn]", "[àng]", + "[ào]", "[á]", "[ái]", "[án]", "[áng]", "[áo]", "[è]", "[èi]", "[èn]", "[èng]", "[èr]", "[é]", "[éi]", "[én]", + "[éng]", "[ér]", "[ì]", "[ìn]", "[ìng]", "[í]", "[ín]", "[íng]", "[ò]", "[òng]", "[òu]", "[ó]", "[óng]", "[óu]", + "[ù]", "[ùn]", "[ú]", "[ún]", "[ā]", "[āi]", "[ān]", "[āng]", "[āo]", "[ē]", "[ēi]", "[ēn]", "[ēng]", "[ě]", + "[ěi]", "[ěn]", "[ěng]", "[ěr]", "[ī]", "[īn]", "[īng]", "[ō]", "[ōng]", "[ōu]", "[ū]", "[ūn]", "[ǎ]", "[ǎi]", + "[ǎn]", "[ǎng]", "[ǎo]", "[ǐ]", "[ǐn]", "[ǐng]", "[ǒ]", "[ǒng]", "[ǒu]", "[ǔ]", "[ǔn]", "[ǘ]", "[ǚ]", "[ǜ]" + ] + } + tokenizer.add_special_tokens(special_tokens) + text_vocab_size = len(tokenizer) + logger.info(f"Tokenizer vocab after CosyVoice3 text special tokens: {text_vocab_size}") + + # Add speech tokens: <|s_0|>, <|s_1|>, ..., <|s_{embedding_size-1}|> + # IMPORTANT: This range must match speech_embedding.num_embeddings (includes speech special tokens). + actual_speech_tokens = speech_token_size # Full embedding size (with speech special tokens) + + # replace to <|sos|> + # replace to <|eos1|> + # replace to <|task_id|> + # replace to <|fill|> + speech_tokens = [f"<|s_{i}|>" for i in range(actual_speech_tokens)] + speech_tokens[6561] = "<|sos|>" + speech_tokens[6562] = "<|eos1|>" + speech_tokens[6563] = "<|task_id|>" + speech_tokens[6564] = "<|fill|>" + assert "" not in speech_tokens + assert "" not in speech_tokens + assert "" not in speech_tokens + assert "" not in speech_tokens + tokenizer.add_tokens(speech_tokens) + + new_vocab_size = len(tokenizer) + logger.info(f"New tokenizer vocab size: {new_vocab_size}") + logger.info(f"Added {new_vocab_size - base_vocab_size} tokens total (text special + speech tokens)") + + # 3. Изменяем размер embeddings в Qwen модели + # Выравниваем по 128 для эффективности TensorRT + padded_vocab_size = ((new_vocab_size + 127) // 128) * 128 + qwen_model.resize_token_embeddings(padded_vocab_size) + logger.info(f"Resized embeddings to: {padded_vocab_size}") + + # Speech tokens start after text vocab (base + CosyVoice3 text special tokens) + speech_token_offset = text_vocab_size + + # 4. Копируем speech_embedding в расширенную часть embed_tokens + input_embeddings = qwen_model.get_input_embeddings() + hidden_size = input_embeddings.weight.shape[1] + + logger.info(f"Hidden size: {hidden_size}") + logger.info(f"speech_embedding shape: {speech_embedding.weight.shape}") + logger.info(f"llm_decoder shape: {llm_decoder.weight.shape}") + + with torch.no_grad(): + # Копируем speech_embedding веса в embed_tokens + # Indices: [speech_token_offset, speech_token_offset + speech_token_size) + src_size = min(speech_embedding.weight.shape[0], actual_speech_tokens) + input_embeddings.weight[speech_token_offset:speech_token_offset + src_size] = \ + speech_embedding.weight[:src_size].to(input_embeddings.weight.dtype) + + logger.info(f"Copied speech_embedding to embed_tokens[{speech_token_offset}:{speech_token_offset + src_size}]") + + # 5. Создаём новый lm_head с расширенным vocab и копируем llm_decoder + # Оригинальный lm_head: hidden_size -> original_vocab_size + # Новый lm_head: hidden_size -> padded_vocab_size + # llm_decoder: hidden_size -> speech_token_size + + # Создаём новый lm_head + has_bias = llm_decoder.bias is not None + new_lm_head = torch.nn.Linear( + in_features=hidden_size, + out_features=padded_vocab_size, + bias=has_bias + ) + + with torch.no_grad(): + # Инициализируем веса: + # - Text часть: копируем из оригинального lm_head (или нули) + # - Speech часть: копируем из llm_decoder + # - Padding: нули + + # Сначала заполняем нулями и -inf в bias (чтобы text токены не генерировались) + new_lm_head.weight.data.zero_() + if has_bias: + new_lm_head.bias.data.fill_(-float('inf')) + + # Копируем оригинальный lm_head для text токенов (опционально) + original_lm_head = qwen_model.lm_head + if original_lm_head is not None and original_lm_head.weight.shape[0] >= text_vocab_size: + new_lm_head.weight[:text_vocab_size] = original_lm_head.weight[:text_vocab_size] + if has_bias and original_lm_head.bias is not None: + new_lm_head.bias[:text_vocab_size] = original_lm_head.bias[:text_vocab_size] + + # Копируем llm_decoder для speech токенов + decoder_size = min(llm_decoder.weight.shape[0], actual_speech_tokens) + new_lm_head.weight[speech_token_offset:speech_token_offset + decoder_size] = \ + llm_decoder.weight[:decoder_size].to(new_lm_head.weight.dtype) + + if has_bias: + new_lm_head.bias[speech_token_offset:speech_token_offset + decoder_size] = \ + llm_decoder.bias[:decoder_size].to(new_lm_head.bias.dtype) + else: + # Если llm_decoder не имеет bias, но мы хотим его для text токенов + pass + + # Заменяем lm_head + qwen_model.lm_head = new_lm_head + + logger.info(f"Created new lm_head with shape: {new_lm_head.weight.shape}") + logger.info(f"Copied llm_decoder to lm_head[{speech_token_offset}:{speech_token_offset + decoder_size}]") + + # 6. Обновляем конфигурацию модели + qwen_model.config.vocab_size = padded_vocab_size + qwen_model.config.tie_word_embeddings = False # Embeddings и lm_head теперь разные! + + # Set EOS token for generation (speech EOS lives inside speech_embedding as <|s_{base_speech_token_size+1}|>) + base_speech_token_size = getattr(cosyvoice3_llm, "speech_token_size", 6561) + eos_speech_idx = base_speech_token_size + 1 + eos_id = speech_token_offset + eos_speech_idx + qwen_model.config.eos_token_id = eos_id + + # Настройки генерации + qwen_model.generation_config.eos_token_id = eos_id + qwen_model.generation_config.pad_token_id = eos_id + qwen_model.generation_config.temperature = 0.8 + qwen_model.generation_config.top_p = 0.95 + qwen_model.generation_config.top_k = 25 + qwen_model.generation_config.repetition_penalty = 1.1 + qwen_model.generation_config.max_new_tokens = 2048 + + # 7. Конвертируем в нужный dtype + dtype_map = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + } + target_dtype = dtype_map[dtype] + qwen_model.to(target_dtype) + + # 8. Сохраняем модель и tokenizer + os.makedirs(output_dir, exist_ok=True) + + qwen_model.save_pretrained(output_dir) + + TEMPLATE = "{%- for message in messages %}{%- if message['role'] == 'user' %}{{- '<|sos|>' + message['content'] + '<|task_id|>' }}{%- elif message['role'] == 'assistant' %}{{- message['content']}}{%- endif %}{%- endfor %}" + tokenizer.chat_template = TEMPLATE + tokenizer.save_pretrained(output_dir) + + # Сохраняем метаданные для TRT-LLM inference + metadata = { + "original_vocab_size": base_vocab_size, + "text_vocab_size": text_vocab_size, + "base_speech_token_size": base_speech_token_size, + "embedding_size": actual_speech_tokens, + "padded_vocab_size": padded_vocab_size, + "eos_token_id": eos_id, + "speech_token_offset": speech_token_offset, + "dtype": dtype, + } + + import json + with open(os.path.join(output_dir, "cosyvoice3_metadata.json"), "w") as f: + json.dump(metadata, f, indent=2) + + logger.info(f"Saved HuggingFace model to {output_dir}") + logger.info(f"Metadata: {metadata}") + + return output_dir, metadata + + +def main(): + args = parse_args() + + output_dir = args.output_dir + if output_dir is None: + output_dir = os.path.join(args.model_dir, "hf_merged") + + convert_cosyvoice3_to_hf( + model_dir=args.model_dir, + output_dir=output_dir, + dtype=args.dtype, + ) + + print("\n" + "=" * 70) + print("✅ Conversion complete!") + print("=" * 70) + print(f"\nHuggingFace model saved to: {output_dir}") + print("\nNext steps:") + print("1. Convert to TRT-LLM weights:") + print(f" python -c \"from tensorrt_llm.models import QWenForCausalLM; ...") + print("\n2. Build TRT-LLM engines:") + print(f" trtllm-build --checkpoint_dir --output_dir ...") + print("=" * 70) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/runtime/triton_trtllm/token2wav_cosyvoice3.py b/runtime/triton_trtllm/token2wav_cosyvoice3.py new file mode 100644 index 0000000..69303dd --- /dev/null +++ b/runtime/triton_trtllm/token2wav_cosyvoice3.py @@ -0,0 +1,414 @@ +""" Example Usage + CUDA_VISIBLE_DEVICES=0 \ + python3 token2wav_cosyvoice3.py --enable-trt || exit 1 +""" +import torch +import torchaudio +import torchaudio.compliance.kaldi as kaldi +import onnxruntime +import s3tokenizer +import os +import logging +import argparse +import queue +import time +import numpy as np +from functools import partial +from hyperpyyaml import load_hyperpyyaml +from matcha.utils.audio import mel_spectrogram as matcha_mel_spectrogram +from torch.utils.data import DataLoader +from datasets import load_dataset + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + +# CosyVoice3 mel params from cosyvoice3.yaml (fmax=None, NOT 8000) +mel_spectrogram = partial(matcha_mel_spectrogram, + n_fft=1920, num_mels=80, sampling_rate=24000, + hop_size=480, win_size=1920, fmin=0, fmax=None, center=False) + + +def convert_onnx_to_trt(trt_model, trt_kwargs, onnx_model, fp16, autocast_mode=False): + import tensorrt as trt + logging.info("Converting onnx to trt...") + if autocast_mode: + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + else: + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + logger = trt.Logger(trt.Logger.INFO) + builder = trt.Builder(logger) + network = builder.create_network(network_flags) + parser = trt.OnnxParser(network, logger) + config = builder.create_builder_config() + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 32) # 4GB + if not autocast_mode: + if fp16: + config.set_flag(trt.BuilderFlag.FP16) + profile = builder.create_optimization_profile() + # load onnx model + with open(onnx_model, "rb") as f: + if not parser.parse(f.read()): + for error in range(parser.num_errors): + print(parser.get_error(error)) + raise ValueError('failed to parse {}'.format(onnx_model)) + # set input shapes + for i in range(len(trt_kwargs['input_names'])): + profile.set_shape(trt_kwargs['input_names'][i], trt_kwargs['min_shape'][i], trt_kwargs['opt_shape'][i], trt_kwargs['max_shape'][i]) + tensor_dtype = trt.DataType.HALF if fp16 else trt.DataType.FLOAT + # set input and output data type + for i in range(network.num_inputs): + input_tensor = network.get_input(i) + input_tensor.dtype = tensor_dtype + for i in range(network.num_outputs): + output_tensor = network.get_output(i) + output_tensor.dtype = tensor_dtype + config.add_optimization_profile(profile) + engine_bytes = builder.build_serialized_network(network, config) + # save trt engine + with open(trt_model, "wb") as f: + f.write(engine_bytes) + logging.info("Succesfully convert onnx to trt...") + + +class TrtContextWrapper: + def __init__(self, trt_engine, trt_concurrent=1, device='cuda:0'): + self.trt_context_pool = queue.Queue(maxsize=trt_concurrent) + self.trt_engine = trt_engine + self.device = device + for _ in range(trt_concurrent): + trt_context = trt_engine.create_execution_context() + trt_stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device))) + assert trt_context is not None, 'failed to create trt context, maybe not enough CUDA memory, try reduce current trt concurrent {}'.format(trt_concurrent) + self.trt_context_pool.put([trt_context, trt_stream]) + assert self.trt_context_pool.empty() is False, 'no avaialbe estimator context' + + def acquire_estimator(self): + return self.trt_context_pool.get(), self.trt_engine + + def release_estimator(self, context, stream): + self.trt_context_pool.put([context, stream]) + + +class CosyVoice3_Token2Wav(torch.nn.Module): + def __init__(self, model_dir, enable_trt=False, device_id=0, autocast_mode=True, streaming=False): + super().__init__() + self.device_id = device_id + self.device = f"cuda:{device_id}" + self.autocast_mode = autocast_mode + self.streaming = streaming + + # Load flow and hift from cosyvoice3.yaml + with open(f"{model_dir}/cosyvoice3.yaml", "r") as f: + configs = load_hyperpyyaml(f, overrides={ + 'qwen_pretrain_path': os.path.join(model_dir, 'CosyVoice-BlankEN') + }) + self.flow = configs['flow'] + self.flow.load_state_dict( + torch.load(f"{model_dir}/flow.pt", map_location="cpu", weights_only=True), + strict=True + ) + self.flow.to(self.device).eval() + + self.hift = configs['hift'] + hift_state_dict = { + k.replace('generator.', ''): v + for k, v in torch.load(f"{model_dir}/hift.pt", map_location="cpu", weights_only=True).items() + } + self.hift.load_state_dict(hift_state_dict, strict=True) + self.hift.to(self.device).eval() + + # Speaker embedding model (campplus) + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + self.spk_model = onnxruntime.InferenceSession( + f"{model_dir}/campplus.onnx", sess_options=option, + providers=["CPUExecutionProvider"] + ) + + # Audio tokenizer v3 + self.audio_tokenizer = s3tokenizer.load_model( + f"{model_dir}/speech_tokenizer_v3.onnx" + ).to(self.device).eval() + + self.fp16 = enable_trt + if enable_trt: + self.flow.half() + self.load_trt(model_dir) + self.load_spk_trt(model_dir) + + def load_trt(self, model_dir, trt_concurrent=1): + streaming_prefix = 'streaming.' if self.streaming else '' + if self.autocast_mode: + onnx_path = f'{model_dir}/flow.decoder.estimator.{streaming_prefix}autocast_fp16.onnx' + trt_path = f'{model_dir}/flow.decoder.estimator.{streaming_prefix}autocast_fp16.{self.device_id}.plan' + else: + onnx_path = f'{model_dir}/flow.decoder.estimator.{streaming_prefix}fp32.onnx' + trt_path = f'{model_dir}/flow.decoder.estimator.{streaming_prefix}fp32.{self.device_id}.plan' + + if not os.path.exists(trt_path) or os.path.getsize(trt_path) == 0: + trt_kwargs = self.get_trt_kwargs() + convert_onnx_to_trt(trt_path, trt_kwargs, onnx_path, + fp16=True, autocast_mode=self.autocast_mode) + del self.flow.decoder.estimator + import tensorrt as trt + with open(trt_path, '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(trt_path) + self.flow.decoder.estimator = TrtContextWrapper( + estimator_engine, trt_concurrent=trt_concurrent, device=self.device + ) + + def get_trt_kwargs(self): + # CosyVoice3 DiT estimator has 6 inputs: x, mask, mu, t, spks, cond + # Only inputs with dynamic dims need optimization profiles. + # t=[2(fixed)] and spks=[2(fixed),80(fixed)] are fully fixed, TRT infers from ONNX. + 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 load_spk_trt(self, model_dir, trt_concurrent=1, fp16=False): + spk_trt_path = f'{model_dir}/campplus.{self.device_id}.fp32.plan' + spk_onnx_path = f'{model_dir}/campplus.onnx' + if not os.path.exists(spk_trt_path) or os.path.getsize(spk_trt_path) == 0: + trt_kwargs = self.get_spk_trt_kwargs() + convert_onnx_to_trt(spk_trt_path, trt_kwargs, spk_onnx_path, fp16) + import tensorrt as trt + with open(spk_trt_path, 'rb') as f: + spk_engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read()) + assert spk_engine is not None, 'failed to load trt {}'.format(spk_trt_path) + self.spk_model = TrtContextWrapper(spk_engine, trt_concurrent=trt_concurrent, device=self.device) + + def get_spk_trt_kwargs(self): + min_shape = [(1, 4, 80)] + opt_shape = [(1, 500, 80)] + max_shape = [(1, 3000, 80)] + input_names = ["input"] + return {'min_shape': min_shape, 'opt_shape': opt_shape, + 'max_shape': max_shape, 'input_names': input_names} + + def forward_spk_embedding(self, spk_feat): + if isinstance(self.spk_model, onnxruntime.InferenceSession): + return self.spk_model.run( + None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()} + )[0].flatten().tolist() + else: + [spk_model, stream], trt_engine = self.spk_model.acquire_estimator() + with torch.cuda.device(self.device_id): + torch.cuda.current_stream().synchronize() + spk_feat = spk_feat.unsqueeze(dim=0).to(self.device) + batch_size = spk_feat.size(0) + + with stream: + spk_model.set_input_shape('input', (batch_size, spk_feat.size(1), 80)) + output_tensor = torch.empty((batch_size, 192), device=spk_feat.device) + + data_ptrs = [spk_feat.contiguous().data_ptr(), + output_tensor.contiguous().data_ptr()] + for i, j in enumerate(data_ptrs): + spk_model.set_tensor_address(trt_engine.get_tensor_name(i), j) + assert spk_model.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True + torch.cuda.current_stream().synchronize() + self.spk_model.release_estimator(spk_model, stream) + + return output_tensor.cpu().numpy().flatten().tolist() + + def prompt_audio_tokenization(self, prompt_audios_list): + prompt_speech_tokens_list, prompt_speech_mels_list = [], [] + for audio in prompt_audios_list: + assert len(audio.shape) == 1 + log_mel = s3tokenizer.log_mel_spectrogram(audio) + prompt_speech_mels_list.append(log_mel) + prompt_mels_for_llm, prompt_mels_lens_for_llm = s3tokenizer.padding(prompt_speech_mels_list) + prompt_speech_tokens, prompt_speech_tokens_lens = self.audio_tokenizer.quantize( + prompt_mels_for_llm.to(self.device), prompt_mels_lens_for_llm.to(self.device) + ) + for i in range(len(prompt_speech_tokens)): + speech_tokens_i = prompt_speech_tokens[i, :prompt_speech_tokens_lens[i].item()].tolist() + prompt_speech_tokens_list.append(speech_tokens_i) + return prompt_speech_tokens_list + + def get_spk_emb(self, prompt_audios_list): + spk_emb_for_flow = [] + for audio in prompt_audios_list: + assert len(audio.shape) == 1 + spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) + spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True) + spk_emb = self.forward_spk_embedding(spk_feat) + spk_emb_for_flow.append(spk_emb) + spk_emb_for_flow = torch.tensor(spk_emb_for_flow) + return spk_emb_for_flow + + def get_prompt_mels(self, prompt_audios_list, prompt_audios_sample_rate): + prompt_mels_for_flow = [] + prompt_mels_lens_for_flow = [] + for audio, sample_rate in zip(prompt_audios_list, prompt_audios_sample_rate): + assert len(audio.shape) == 1 + audio = audio.unsqueeze(0) + if sample_rate != 24000: + audio = torchaudio.transforms.Resample( + orig_freq=sample_rate, new_freq=24000)(audio) + # CosyVoice3: fmax=None (Nyquist), matching cosyvoice3.yaml + mel = mel_spectrogram(audio).transpose(1, 2).squeeze(0) # [T, 80] + prompt_mels_for_flow.append(mel) + prompt_mels_lens_for_flow.append(mel.shape[0]) + prompt_mels_for_flow = torch.nn.utils.rnn.pad_sequence( + prompt_mels_for_flow, batch_first=True, padding_value=0) # [B, T', 80] + prompt_mels_lens_for_flow = torch.tensor(prompt_mels_lens_for_flow) + return prompt_mels_for_flow, prompt_mels_lens_for_flow + + def forward_flow(self, prompt_speech_tokens_list, generated_speech_tokens_list, + prompt_mels_for_flow, prompt_mels_lens_for_flow, + spk_emb_for_flow): + batch_size = len(generated_speech_tokens_list) + generated_mels_list = [] + + # CausalMaskedDiffWithDiT.inference asserts batch_size==1, so iterate per-sample + for i in range(batch_size): + token = torch.tensor([generated_speech_tokens_list[i]]).to(self.device) + token_len = torch.tensor([len(generated_speech_tokens_list[i])]).to(self.device) + prompt_token = torch.tensor([prompt_speech_tokens_list[i]]).to(self.device) + prompt_token_len = torch.tensor([len(prompt_speech_tokens_list[i])]).to(self.device) + prompt_feat = prompt_mels_for_flow[i:i+1, :prompt_mels_lens_for_flow[i]].to(self.device) + prompt_feat_len = prompt_mels_lens_for_flow[i:i+1].to(self.device) + embedding = spk_emb_for_flow[i:i+1].to(self.device) + + # CausalMaskedDiffWithDiT.inference returns mel already without prompt portion + with torch.cuda.amp.autocast(self.fp16): + mel, _ = self.flow.inference( + token=token, + token_len=token_len, + prompt_token=prompt_token, + prompt_token_len=prompt_token_len, + prompt_feat=prompt_feat, + prompt_feat_len=prompt_feat_len, + embedding=embedding, + streaming=False, + finalize=True + ) + generated_mels_list.append(mel) + + return generated_mels_list + + def forward_hift(self, generated_mels_list): + generated_wavs = [] + for mel in generated_mels_list: + # CausalHiFTGenerator.inference with finalize=True + wav, _ = self.hift.inference(speech_feat=mel, finalize=True) + generated_wavs.append(wav) + return generated_wavs + + def forward_stream(self, generated_speech_tokens, prompt_speech_tokens, + prompt_feat, embedding, + token_hop_len=25, stream_scale_factor=2, token_max_hop_len=100): + """Streaming token2wav for a single sample: process tokens in chunks.""" + prompt_token = torch.tensor([prompt_speech_tokens]).to(self.device) + prompt_token_len = torch.tensor([len(prompt_speech_tokens)]).to(self.device) + prompt_feat = prompt_feat.to(self.device) + prompt_feat_len = torch.tensor([prompt_feat.shape[1]]).to(self.device) + embedding = embedding.to(self.device) + + pre_lookahead_len = self.flow.pre_lookahead_len + token_mel_ratio = self.flow.token_mel_ratio + + # Align first chunk with hop_len boundary + prompt_token_pad = int( + np.ceil(prompt_token.shape[1] / token_hop_len) * token_hop_len + - prompt_token.shape[1] + ) + + total_tokens = len(generated_speech_tokens) + token_offset = 0 + current_hop = token_hop_len + hift_cache_mel = None + speech_offset = 0 + audio_chunks = [] + + while token_offset < total_tokens: + this_hop = current_hop + prompt_token_pad if token_offset == 0 else current_hop + remaining = total_tokens - token_offset + + if remaining >= this_hop + pre_lookahead_len: + end_idx = token_offset + this_hop + pre_lookahead_len + this_token = torch.tensor([generated_speech_tokens[:end_idx]]).to(self.device) + finalize = False + else: + this_token = torch.tensor([generated_speech_tokens]).to(self.device) + finalize = True + + with torch.cuda.amp.autocast(self.fp16): + mel, _ = self.flow.inference( + token=this_token, + token_len=torch.tensor([this_token.shape[1]]).to(self.device), + prompt_token=prompt_token, + prompt_token_len=prompt_token_len, + prompt_feat=prompt_feat, + prompt_feat_len=prompt_feat_len, + embedding=embedding, + streaming=True, + finalize=finalize, + ) + + mel = mel[:, :, token_offset * token_mel_ratio:] + + if hift_cache_mel is not None: + mel = torch.concat([hift_cache_mel, mel], dim=2) + hift_cache_mel = mel + + tts_speech, _ = self.hift.inference(speech_feat=mel, finalize=finalize) + tts_speech = tts_speech[:, speech_offset:] + speech_offset += tts_speech.shape[1] + + logger.info(f"[stream] token_offset={token_offset}, this_hop={this_hop}, " + f"mel_shape={mel.shape}, speech_len={tts_speech.shape[1]}, finalize={finalize}") + + audio_chunks.append(tts_speech) + + token_offset += this_hop + if not finalize: + current_hop = min(token_max_hop_len, current_hop * stream_scale_factor) + else: + break + + return torch.cat(audio_chunks, dim=1) + + @torch.inference_mode() + def forward(self, generated_speech_tokens_list, prompt_audios_list, + prompt_audios_sample_rate, streaming=False): + assert all(sr == 16000 for sr in prompt_audios_sample_rate) + + prompt_speech_tokens_list = self.prompt_audio_tokenization(prompt_audios_list) + prompt_mels_for_flow, prompt_mels_lens_for_flow = self.get_prompt_mels( + prompt_audios_list, prompt_audios_sample_rate) + spk_emb_for_flow = self.get_spk_emb(prompt_audios_list) + + # Align prompt_speech_feat and prompt_speech_token to exact 2:1 ratio + # (matches frontend.frontend_zero_shot logic) + for i in range(len(prompt_speech_tokens_list)): + token_len = min(int(prompt_mels_lens_for_flow[i].item() / 2), + len(prompt_speech_tokens_list[i])) + prompt_speech_tokens_list[i] = prompt_speech_tokens_list[i][:token_len] + prompt_mels_lens_for_flow[i] = 2 * token_len + + if streaming: + generated_wavs = [] + for i in range(len(generated_speech_tokens_list)): + prompt_feat = prompt_mels_for_flow[i:i+1, :prompt_mels_lens_for_flow[i]] + embedding = spk_emb_for_flow[i:i+1] + wav = self.forward_stream( + generated_speech_tokens_list[i], + prompt_speech_tokens_list[i], + prompt_feat, embedding, + ) + generated_wavs.append(wav) + return generated_wavs + + generated_mels_list = self.forward_flow( + prompt_speech_tokens_list, generated_speech_tokens_list, + prompt_mels_for_flow, prompt_mels_lens_for_flow, spk_emb_for_flow) + + generated_wavs = self.forward_hift(generated_mels_list) + return generated_wavs From c7686faccc59a93187de8222a9a80cbfa8920418 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 16 Mar 2026 14:41:44 +0800 Subject: [PATCH 2/5] rename files --- ...README.DIT.md => README.Cosyvoice2.DiT.md} | 2 +- .../{README.md => README.Cosyvoice2.Unet.md} | 2 +- ....yml => docker-compose.cosyvoice2.dit.yml} | 0 ...yml => docker-compose.cosyvoice2.unet.yml} | 0 runtime/triton_trtllm/requirements.txt | 1 + runtime/triton_trtllm/run_cosyvoice3.sh | 156 +++--------------- .../scripts/convert_cosyvoice3_to_hf.py | 105 ++++++------ 7 files changed, 76 insertions(+), 190 deletions(-) rename runtime/triton_trtllm/{README.DIT.md => README.Cosyvoice2.DiT.md} (99%) rename runtime/triton_trtllm/{README.md => README.Cosyvoice2.Unet.md} (99%) rename runtime/triton_trtllm/{docker-compose.dit.yml => docker-compose.cosyvoice2.dit.yml} (100%) rename runtime/triton_trtllm/{docker-compose.yml => docker-compose.cosyvoice2.unet.yml} (100%) diff --git a/runtime/triton_trtllm/README.DIT.md b/runtime/triton_trtllm/README.Cosyvoice2.DiT.md similarity index 99% rename from runtime/triton_trtllm/README.DIT.md rename to runtime/triton_trtllm/README.Cosyvoice2.DiT.md index 2fd9f61..c4bfa52 100644 --- a/runtime/triton_trtllm/README.DIT.md +++ b/runtime/triton_trtllm/README.Cosyvoice2.DiT.md @@ -8,7 +8,7 @@ This document describes how to accelerate CosyVoice with a DiT-based Token2Wav m Launch the service directly with Docker Compose: ```sh -docker compose -f docker-compose.dit.yml up +docker compose -f docker-compose.cosyvoice2.dit.yml up ``` ### Build the Docker Image diff --git a/runtime/triton_trtllm/README.md b/runtime/triton_trtllm/README.Cosyvoice2.Unet.md similarity index 99% rename from runtime/triton_trtllm/README.md rename to runtime/triton_trtllm/README.Cosyvoice2.Unet.md index 3765038..33f2908 100644 --- a/runtime/triton_trtllm/README.md +++ b/runtime/triton_trtllm/README.Cosyvoice2.Unet.md @@ -6,7 +6,7 @@ Contributed by Yuekai Zhang (NVIDIA). Launch the service directly with Docker Compose: ```sh -docker compose up +docker compose -f docker-compose.cosyvoice2.unet.yml up ``` ### Build the Docker Image diff --git a/runtime/triton_trtllm/docker-compose.dit.yml b/runtime/triton_trtllm/docker-compose.cosyvoice2.dit.yml similarity index 100% rename from runtime/triton_trtllm/docker-compose.dit.yml rename to runtime/triton_trtllm/docker-compose.cosyvoice2.dit.yml diff --git a/runtime/triton_trtllm/docker-compose.yml b/runtime/triton_trtllm/docker-compose.cosyvoice2.unet.yml similarity index 100% rename from runtime/triton_trtllm/docker-compose.yml rename to runtime/triton_trtllm/docker-compose.cosyvoice2.unet.yml diff --git a/runtime/triton_trtllm/requirements.txt b/runtime/triton_trtllm/requirements.txt index ba0623b..2fce08c 100644 --- a/runtime/triton_trtllm/requirements.txt +++ b/runtime/triton_trtllm/requirements.txt @@ -12,3 +12,4 @@ pyworld openai-whisper tritonclient modelscope +x_transformers \ No newline at end of file diff --git a/runtime/triton_trtllm/run_cosyvoice3.sh b/runtime/triton_trtllm/run_cosyvoice3.sh index accdebe..ac42c56 100644 --- a/runtime/triton_trtllm/run_cosyvoice3.sh +++ b/runtime/triton_trtllm/run_cosyvoice3.sh @@ -10,20 +10,21 @@ export PYTHONPATH=${cosyvoice_path}/third_party/Matcha-TTS:$PYTHONPATH stage=$1 stop_stage=$2 -huggingface_model_local_dir=./hf_cosyvoice3_llm -model_scope_model_local_dir=/workspace_yuekai/HF/Fun-CosyVoice3-0.5B-2512 +huggingface_llm_local_dir=$cosyvoice_path/runtime/triton_trtllm/hf_cosyvoice3_llm +cosyvoice3_official_model_dir=$cosyvoice_path/runtime/triton_trtllm/Fun-CosyVoice3-0.5B-2512 trt_dtype=bfloat16 -trt_weights_dir=./trt_weights_${trt_dtype} -trt_engines_dir=./trt_engines_${trt_dtype} +trt_weights_dir=$cosyvoice_path/runtime/triton_trtllm/trt_weights_${trt_dtype} +trt_engines_dir=$cosyvoice_path/runtime/triton_trtllm/trt_engines_${trt_dtype} -model_repo_src=./model_repo_cosyvoice3 -model_repo=./deploy_cosyvoice3 -bls_instance_num=1 +model_repo_src=$cosyvoice_path/runtime/triton_trtllm/model_repo_cosyvoice3 +model_repo=$cosyvoice_path/runtime/triton_trtllm/model_repo_cosyvoice3_copy +bls_instance_num=10 if [ $stage -le -1 ] && [ $stop_stage -ge -1 ]; then echo "Cloning CosyVoice" + pip3 install --upgrade x_transformers s3tokenizer git clone --recursive https://github.com/FunAudioLLM/CosyVoice.git $cosyvoice_path cd $cosyvoice_path git submodule update --init --recursive @@ -31,23 +32,16 @@ if [ $stage -le -1 ] && [ $stop_stage -ge -1 ]; then fi if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then - echo "" - # see https://github.com/nvidia-china-sae/mair-hub/blob/main/rl-tutorial/cosyvoice_llm/pretrained_to_huggingface.py - # huggingface-cli download --local-dir $huggingface_model_local_dir yuekai/cosyvoice2_llm - # modelscope download --model iic/CosyVoice2-0.5B --local_dir $model_scope_model_local_dir - - # pip3 install --upgrade x_transformers s3tokenizer - # pip install -U nvidia-modelopt[all] - python3 scripts/convert_cosyvoice3_to_hf.py \ - --model-dir $model_scope_model_local_dir \ - --output-dir $huggingface_model_local_dir || exit 1 # TODO: output dir should be here - + echo "Downloading CosyVoice3 Checkpoints" + huggingface-cli download --local-dir $huggingface_llm_local_dir yuekai/Fun-CosyVoice3-0.5B-2512-LLM-HF + huggingface-cli download --local-dir $cosyvoice3_official_model_dir yuekai/Fun-CosyVoice3-0.5B-2512-FP16-ONNX + huggingface-cli download --local-dir $cosyvoice3_official_model_dir FunAudioLLM/Fun-CosyVoice3-0.5B-2512 fi if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then echo "Converting checkpoint to TensorRT weights" - python3 scripts/convert_checkpoint.py --model_dir $huggingface_model_local_dir \ + python3 scripts/convert_checkpoint.py --model_dir $huggingface_llm_local_dir \ --output_dir $trt_weights_dir \ --dtype $trt_dtype || exit 1 @@ -60,7 +54,7 @@ if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then echo "Testing TensorRT engines" python3 ./scripts/test_llm.py --input_text "你好，请问你叫什么？" \ - --tokenizer_dir $huggingface_model_local_dir \ + --tokenizer_dir $huggingface_llm_local_dir \ --top_k 50 --top_p 0.95 --temperature 0.8 \ --engine_dir=$trt_engines_dir || exit 1 fi @@ -78,8 +72,8 @@ if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then cp -r ${model_repo_src}/speaker_embedding $model_repo/ MAX_QUEUE_DELAY_MICROSECONDS=0 - MODEL_DIR=$model_scope_model_local_dir - LLM_TOKENIZER_DIR=$huggingface_model_local_dir + MODEL_DIR=$cosyvoice3_official_model_dir + LLM_TOKENIZER_DIR=$huggingface_llm_local_dir BLS_INSTANCE_NUM=$bls_instance_num TRITON_MAX_BATCH_SIZE=1 DECOUPLED_MODE=True @@ -92,44 +86,16 @@ if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then fi -if [ $stage -le 30 ] && [ $stop_stage -ge 30 ]; then - echo "Starting CosyVoice3 Triton server and LLM using trtllm-serve" - CUDA_VISIBLE_DEVICES=0 mpirun -np 1 --allow-run-as-root --oversubscribe trtllm-serve serve --tokenizer $huggingface_model_local_dir $trt_engines_dir --max_batch_size 64 --kv_cache_free_gpu_memory_fraction 0.4 -fi - - -if [ $stage -le 40 ] && [ $stop_stage -ge 40 ]; then - - CUDA_VISIBLE_DEVICES=1 tritonserver --model-repository $model_repo --http-port 18000 --grpc-port 18001 --metrics-port 18002 & -fi - - if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then echo "Starting CosyVoice3 Triton server and LLM using trtllm-serve" - CUDA_VISIBLE_DEVICES=0 mpirun -np 1 --allow-run-as-root --oversubscribe trtllm-serve serve --tokenizer $huggingface_model_local_dir $trt_engines_dir --max_batch_size 64 --kv_cache_free_gpu_memory_fraction 0.4 & - CUDA_VISIBLE_DEVICES=0,1,2,3 tritonserver --model-repository $model_repo --http-port 18000 --grpc-port 18001 --metrics-port 18002 & + CUDA_VISIBLE_DEVICES=0 mpirun -np 1 --allow-run-as-root --oversubscribe trtllm-serve serve --tokenizer $huggingface_llm_local_dir $trt_engines_dir --max_batch_size 64 --kv_cache_free_gpu_memory_fraction 0.4 & + CUDA_VISIBLE_DEVICES=0 tritonserver --model-repository $model_repo --http-port 18000 --grpc-port 18001 --metrics-port 18002 & wait - # Test using curl - # curl http://localhost:8000/v1/chat/completions \ - # -H "Content-Type: application/json" \ - # -d '{ - # "model": "", - # "messages":[{"role": "user", "content": "Where is New York?"}, - # {"role": "assistant", "content": "<|s_1708|><|s_2050|><|s_2159|>"}], - # "max_tokens": 512, - # "temperature": 0.8, - # "top_p": 0.95, - # "top_k": 50, - # "stop": ["<|eos1|>"], - # "repetition_penalty": 1.2, - # "stream": false - # }' fi if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then echo "Running benchmark client for CosyVoice3" num_task=4 - mode=offline mode=streaming BLS_INSTANCE_NUM=$bls_instance_num @@ -145,93 +111,13 @@ if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then fi if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then - echo "stage 5: Offline TTS (Cosyvoice2 LLM + Step-Audio2-mini DiT Token2Wav) inference using a single python script" - - datasets=(wenetspeech4tts) # wenetspeech4tts, test_zh, zero_shot_zh - backend=trtllm # hf, trtllm, vllm, trtllm-serve - - batch_sizes=(16) - token2wav_batch_size=1 - - for batch_size in ${batch_sizes[@]}; do - for dataset in ${datasets[@]}; do - output_dir=./${dataset}_${backend}_llm_batch_size_${batch_size}_token2wav_batch_size_${token2wav_batch_size} - CUDA_VISIBLE_DEVICES=1 \ - python3 offline_inference.py \ - --output-dir $output_dir \ - --llm-model-name-or-path $huggingface_model_local_dir \ - --token2wav-path $step_audio_model_dir/token2wav \ - --backend $backend \ - --batch-size $batch_size --token2wav-batch-size $token2wav_batch_size \ - --engine-dir $trt_engines_dir \ - --split-name ${dataset} || exit 1 - done - done -fi - - - - -if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then - echo "Disaggregated Server: LLM and Token2wav on different GPUs" - echo "Starting LLM server on GPU 0" - export CUDA_VISIBLE_DEVICES=0 - mpirun -np 1 --allow-run-as-root --oversubscribe trtllm-serve serve --tokenizer $huggingface_model_local_dir $trt_engines_dir --max_batch_size 64 --kv_cache_free_gpu_memory_fraction 0.4 & - echo "Starting Token2wav server on GPUs 1-3" - Token2wav_num_gpus=3 - http_port=17000 - grpc_port=18000 - metrics_port=16000 - for i in $(seq 0 $(($Token2wav_num_gpus - 1))); do - echo "Starting server on GPU $i" - http_port=$((http_port + 1)) - grpc_port=$((grpc_port + 1)) - metrics_port=$((metrics_port + 1)) - # Two instances of Token2wav server on the same GPU - CUDA_VISIBLE_DEVICES=$(($i + 1)) tritonserver --model-repository $model_repo --http-port $http_port --grpc-port $grpc_port --metrics-port $metrics_port & - http_port=$((http_port + 1)) - grpc_port=$((grpc_port + 1)) - metrics_port=$((metrics_port + 1)) - CUDA_VISIBLE_DEVICES=$(($i + 1)) tritonserver --model-repository $model_repo --http-port $http_port --grpc-port $grpc_port --metrics-port $metrics_port & - done - wait -fi - -if [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then - echo "Running benchmark client for Disaggregated Server" - per_gpu_instances=2 - mode=streaming - BLS_INSTANCE_NUM=$bls_instance_num - Token2wav_num_gpus=(1 2 3) - concurrent_tasks=(1 2 3 4 5 6) - for n_gpu in ${Token2wav_num_gpus[@]}; do - echo "Test 1 GPU for LLM server and $n_gpu GPUs for Token2wav servers" - for concurrent_task in ${concurrent_tasks[@]}; do - num_instances=$((per_gpu_instances * n_gpu)) - for i in $(seq 1 $num_instances); do - port=$(($i + 18000)) - python3 client_grpc.py \ - --server-addr localhost \ - --server-port $port \ - --model-name cosyvoice2_dit \ - --num-tasks $concurrent_task \ - --mode $mode \ - --huggingface-dataset yuekai/seed_tts_cosy2 \ - --log-dir ./log_disagg_concurrent_tasks_${concurrent_task}_per_instance_total_token2wav_instances_${num_instances}_port_${port} & - done - wait - done - done -fi - -if [ $stage -le 10 ] && [ $stop_stage -ge 10 ]; then echo "stage 10: Python script CosyVoice3 TTS (LLM + CosyVoice3 Token2Wav) inference" datasets=(wenetspeech4tts) # wenetspeech4tts backend=trtllm-serve # hf, trtllm, vllm, trtllm-serve batch_sizes=(1) - token2wav_batch_size=1 + token2wav_batch_size=1 # Only support 1 for now for batch_size in ${batch_sizes[@]}; do for dataset in ${datasets[@]}; do @@ -239,8 +125,8 @@ if [ $stage -le 10 ] && [ $stop_stage -ge 10 ]; then CUDA_VISIBLE_DEVICES=0 \ python3 infer_cosyvoice3.py \ --output-dir $output_dir \ - --llm-model-name-or-path $huggingface_model_local_dir \ - --token2wav-path $model_scope_model_local_dir \ + --llm-model-name-or-path $huggingface_llm_local_dir \ + --token2wav-path $cosyvoice3_official_model_dir \ --backend $backend \ --batch-size $batch_size --token2wav-batch-size $token2wav_batch_size \ --engine-dir $trt_engines_dir \ diff --git a/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py b/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py index 58318d0..40b99fe 100644 --- a/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py +++ b/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py @@ -13,21 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Конвертация CosyVoice3 LLM в HuggingFace формат с объединёнными embeddings. +Convert CosyVoice3 LLM to HuggingFace format with merged embeddings. -Этот скрипт: -1. Загружает CosyVoice3 модель -2. Расширяет vocab токенизатора с speech токенами -3. Объединяет speech_embedding в embed_tokens Qwen2 -4. Заменяет lm_head на llm_decoder с расширенным vocab -5. Сохраняет модель в HuggingFace формате для TRT-LLM конвертации +This script: +1. Loads CosyVoice3 model +2. Extends tokenizer vocab with speech tokens +3. Merges speech_embedding into embed_tokens of Qwen2 +4. Replaces lm_head with llm_decoder using extended vocab +5. Saves model in HuggingFace format for TRT-LLM conversion Usage: python scripts/convert_cosyvoice3_to_hf.py \ --model-dir pretrained_models/Fun-CosyVoice3-0.5B \ --output-dir pretrained_models/Fun-CosyVoice3-0.5B/hf_merged -После этого можно конвертировать в TRT-LLM: +Then convert to TRT-LLM: trtllm-build --checkpoint_dir --output_dir ... """ import argparse @@ -70,7 +70,7 @@ def parse_args(): def load_cosyvoice3_model(model_dir: str): - """Загружает CosyVoice3 модель для извлечения весов.""" + """Load CosyVoice3 model for weight extraction.""" from hyperpyyaml import load_hyperpyyaml from cosyvoice.utils.class_utils import get_model_type @@ -86,7 +86,7 @@ def load_cosyvoice3_model(model_dir: str): overrides={'qwen_pretrain_path': hf_llm_dir} ) - # Загружаем только LLM + # Load LLM only llm = configs['llm'] llm_weights_path = os.path.join(model_dir, 'llm.pt') llm.load_state_dict(torch.load(llm_weights_path, map_location='cpu'), strict=True) @@ -98,12 +98,11 @@ def load_cosyvoice3_model(model_dir: str): def get_speech_token_size(llm) -> int: - """Определяет размер speech token vocabulary из модели.""" - # CosyVoice3LM имеет: speech_token_size + 200 в llm_decoder - # speech_embedding имеет: speech_token_size + 200 + """Determine speech token vocabulary size from the model.""" + # CosyVoice3LM has: speech_token_size + 200 in llm_decoder + # speech_embedding has: speech_token_size + 200 speech_embedding_size = llm.speech_embedding.num_embeddings - # Вычитаем 200 специальных токенов (sos, eos, task_id, fill, и т.д.) - # Но для безопасности используем полный размер embedding + # Use full embedding size (includes speech special tokens) return speech_embedding_size @@ -113,32 +112,32 @@ def convert_cosyvoice3_to_hf( dtype: str = "bfloat16", ): """ - Конвертирует CosyVoice3 LLM в HuggingFace формат с объединёнными embeddings. - - Архитектура объединения: - - embed_tokens[0:original_vocab_size] = оригинальные text embeddings + Convert CosyVoice3 LLM to HuggingFace format with merged embeddings. + + Merging architecture: + - embed_tokens[0:original_vocab_size] = original text embeddings - embed_tokens[original_vocab_size:original_vocab_size+speech_token_size] = speech_embedding - lm_head[original_vocab_size:original_vocab_size+speech_token_size] = llm_decoder - + Args: - model_dir: Путь к CosyVoice3 модели - output_dir: Путь для сохранения HF модели - dtype: Тип данных для сохранения + model_dir: Path to CosyVoice3 model + output_dir: Path to save HF model + dtype: Data type for saving """ logger.info(f"Loading CosyVoice3 model from {model_dir}") - # 1. Загружаем CosyVoice3 компоненты + # 1. Load CosyVoice3 components cosyvoice3_llm, hf_llm_dir, configs = load_cosyvoice3_model(model_dir) - # Извлекаем ключевые компоненты + # Extract key components qwen_model = cosyvoice3_llm.llm.model # Qwen2ForCausalLM - speech_embedding = cosyvoice3_llm.speech_embedding # Embedding для speech токенов - llm_decoder = cosyvoice3_llm.llm_decoder # Linear для декодирования в speech токены + speech_embedding = cosyvoice3_llm.speech_embedding # Embedding for speech tokens + llm_decoder = cosyvoice3_llm.llm_decoder # Linear for decoding to speech tokens speech_token_size = get_speech_token_size(cosyvoice3_llm) logger.info(f"Speech token size: {speech_token_size}") - # 2. Загружаем tokenizer и добавляем CosyVoice3 text special tokens + speech токены + # 2. Load tokenizer and add CosyVoice3 text special tokens + speech tokens tokenizer = AutoTokenizer.from_pretrained(hf_llm_dir, trust_remote_code=True) base_vocab_size = len(tokenizer) logger.info(f"Base tokenizer vocab size: {base_vocab_size}") @@ -210,8 +209,8 @@ def convert_cosyvoice3_to_hf( logger.info(f"New tokenizer vocab size: {new_vocab_size}") logger.info(f"Added {new_vocab_size - base_vocab_size} tokens total (text special + speech tokens)") - # 3. Изменяем размер embeddings в Qwen модели - # Выравниваем по 128 для эффективности TensorRT + # 3. Resize embeddings in Qwen model + # Align to 128 for TensorRT efficiency padded_vocab_size = ((new_vocab_size + 127) // 128) * 128 qwen_model.resize_token_embeddings(padded_vocab_size) logger.info(f"Resized embeddings to: {padded_vocab_size}") @@ -219,7 +218,7 @@ def convert_cosyvoice3_to_hf( # Speech tokens start after text vocab (base + CosyVoice3 text special tokens) speech_token_offset = text_vocab_size - # 4. Копируем speech_embedding в расширенную часть embed_tokens + # 4. Copy speech_embedding into extended embed_tokens input_embeddings = qwen_model.get_input_embeddings() hidden_size = input_embeddings.weight.shape[1] @@ -228,7 +227,7 @@ def convert_cosyvoice3_to_hf( logger.info(f"llm_decoder shape: {llm_decoder.weight.shape}") with torch.no_grad(): - # Копируем speech_embedding веса в embed_tokens + # Copy speech_embedding weights into embed_tokens # Indices: [speech_token_offset, speech_token_offset + speech_token_size) src_size = min(speech_embedding.weight.shape[0], actual_speech_tokens) input_embeddings.weight[speech_token_offset:speech_token_offset + src_size] = \ @@ -236,12 +235,12 @@ def convert_cosyvoice3_to_hf( logger.info(f"Copied speech_embedding to embed_tokens[{speech_token_offset}:{speech_token_offset + src_size}]") - # 5. Создаём новый lm_head с расширенным vocab и копируем llm_decoder - # Оригинальный lm_head: hidden_size -> original_vocab_size - # Новый lm_head: hidden_size -> padded_vocab_size + # 5. Create new lm_head with extended vocab and copy llm_decoder + # Original lm_head: hidden_size -> original_vocab_size + # New lm_head: hidden_size -> padded_vocab_size # llm_decoder: hidden_size -> speech_token_size - - # Создаём новый lm_head + + # Create new lm_head has_bias = llm_decoder.bias is not None new_lm_head = torch.nn.Linear( in_features=hidden_size, @@ -250,24 +249,24 @@ def convert_cosyvoice3_to_hf( ) with torch.no_grad(): - # Инициализируем веса: - # - Text часть: копируем из оригинального lm_head (или нули) - # - Speech часть: копируем из llm_decoder - # - Padding: нули - - # Сначала заполняем нулями и -inf в bias (чтобы text токены не генерировались) + # Initialize weights: + # - Text part: copy from original lm_head (or zeros) + # - Speech part: copy from llm_decoder + # - Padding: zeros + + # Fill with zeros and -inf in bias (so text tokens are not generated) new_lm_head.weight.data.zero_() if has_bias: new_lm_head.bias.data.fill_(-float('inf')) - # Копируем оригинальный lm_head для text токенов (опционально) + # Copy original lm_head for text tokens (optional) original_lm_head = qwen_model.lm_head if original_lm_head is not None and original_lm_head.weight.shape[0] >= text_vocab_size: new_lm_head.weight[:text_vocab_size] = original_lm_head.weight[:text_vocab_size] if has_bias and original_lm_head.bias is not None: new_lm_head.bias[:text_vocab_size] = original_lm_head.bias[:text_vocab_size] - # Копируем llm_decoder для speech токенов + # Copy llm_decoder for speech tokens decoder_size = min(llm_decoder.weight.shape[0], actual_speech_tokens) new_lm_head.weight[speech_token_offset:speech_token_offset + decoder_size] = \ llm_decoder.weight[:decoder_size].to(new_lm_head.weight.dtype) @@ -276,18 +275,18 @@ def convert_cosyvoice3_to_hf( new_lm_head.bias[speech_token_offset:speech_token_offset + decoder_size] = \ llm_decoder.bias[:decoder_size].to(new_lm_head.bias.dtype) else: - # Если llm_decoder не имеет bias, но мы хотим его для text токенов + # If llm_decoder has no bias but we want it for text tokens pass - # Заменяем lm_head + # Replace lm_head qwen_model.lm_head = new_lm_head logger.info(f"Created new lm_head with shape: {new_lm_head.weight.shape}") logger.info(f"Copied llm_decoder to lm_head[{speech_token_offset}:{speech_token_offset + decoder_size}]") - # 6. Обновляем конфигурацию модели + # 6. Update model configuration qwen_model.config.vocab_size = padded_vocab_size - qwen_model.config.tie_word_embeddings = False # Embeddings и lm_head теперь разные! + qwen_model.config.tie_word_embeddings = False # Embeddings and lm_head are now different! # Set EOS token for generation (speech EOS lives inside speech_embedding as <|s_{base_speech_token_size+1}|>) base_speech_token_size = getattr(cosyvoice3_llm, "speech_token_size", 6561) @@ -295,7 +294,7 @@ def convert_cosyvoice3_to_hf( eos_id = speech_token_offset + eos_speech_idx qwen_model.config.eos_token_id = eos_id - # Настройки генерации + # Generation settings qwen_model.generation_config.eos_token_id = eos_id qwen_model.generation_config.pad_token_id = eos_id qwen_model.generation_config.temperature = 0.8 @@ -304,7 +303,7 @@ def convert_cosyvoice3_to_hf( qwen_model.generation_config.repetition_penalty = 1.1 qwen_model.generation_config.max_new_tokens = 2048 - # 7. Конвертируем в нужный dtype + # 7. Convert to target dtype dtype_map = { "float16": torch.float16, "bfloat16": torch.bfloat16, @@ -313,7 +312,7 @@ def convert_cosyvoice3_to_hf( target_dtype = dtype_map[dtype] qwen_model.to(target_dtype) - # 8. Сохраняем модель и tokenizer + # 8. Save model and tokenizer os.makedirs(output_dir, exist_ok=True) qwen_model.save_pretrained(output_dir) @@ -322,7 +321,7 @@ def convert_cosyvoice3_to_hf( tokenizer.chat_template = TEMPLATE tokenizer.save_pretrained(output_dir) - # Сохраняем метаданные для TRT-LLM inference + # Save metadata for TRT-LLM inference metadata = { "original_vocab_size": base_vocab_size, "text_vocab_size": text_vocab_size, From 4ddb9966b2dadd35cd3009d9b366059f656fbed3 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 16 Mar 2026 15:21:56 +0800 Subject: [PATCH 3/5] update results --- runtime/triton_trtllm/README.Cosyvoice3.md | 87 +++++++++++++++++++ runtime/triton_trtllm/README.md | 37 ++++++++ .../docker-compose.cosyvoice3.yml | 20 +++++ runtime/triton_trtllm/run_cosyvoice3.sh | 22 ++--- 4 files changed, 156 insertions(+), 10 deletions(-) create mode 100644 runtime/triton_trtllm/README.Cosyvoice3.md create mode 100644 runtime/triton_trtllm/README.md create mode 100644 runtime/triton_trtllm/docker-compose.cosyvoice3.yml diff --git a/runtime/triton_trtllm/README.Cosyvoice3.md b/runtime/triton_trtllm/README.Cosyvoice3.md new file mode 100644 index 0000000..278ec96 --- /dev/null +++ b/runtime/triton_trtllm/README.Cosyvoice3.md @@ -0,0 +1,87 @@ +## Accelerating CosyVoice3 with NVIDIA Triton Inference Server and TensorRT-LLM + +Contributed by Yuekai Zhang (NVIDIA). + +### Quick Start + +Launch the service directly with Docker Compose: +```sh +docker compose -f docker-compose.cosyvoice3.yml up +``` + +### Build the Docker Image + +To build the image from scratch: +```sh +docker build . -f Dockerfile.server -t soar97/triton-cosyvoice:25.06 +``` + +### Run a Docker Container +```sh +your_mount_dir=/mnt:/mnt +docker run -it --name "cosyvoice-server" --gpus all --net host -v $your_mount_dir --shm-size=2g soar97/triton-cosyvoice:25.06 +``` + +### Understanding `run_cosyvoice3.sh` + +The `run_cosyvoice3.sh` script orchestrates the entire workflow through numbered stages. + +You can run a subset of stages with: +```sh +bash run_cosyvoice3.sh +``` +- ``: The stage to start from. +- ``: The stage to stop after. + +**Stages:** + +- **Stage -1**: Clones the `CosyVoice` repository. +- **Stage 0**: Downloads the `Fun-CosyVoice3-0.5B-2512` model and its HuggingFace LLM checkpoint. +- **Stage 1**: Converts the HuggingFace checkpoint for the LLM to the TensorRT-LLM format and builds the TensorRT engines. +- **Stage 2**: Creates the Triton model repository, including configurations for `cosyvoice3`, `token2wav`, `vocoder`, `audio_tokenizer`, and `speaker_embedding`. +- **Stage 3**: Launches the Triton Inference Server for Token2Wav module and uses `trtllm-serve` to deploy CosyVoice3 LLM. +- **Stage 4**: Runs the gRPC benchmark client for performance testing. +- **Stage 5**: Runs the offline TTS inference benchmark test. + +### Export Models and Launch Server + +Inside the Docker container, prepare the models and start the Triton server by running stages 0-3: +```sh +# This command runs stages 0, 1, 2, and 3 +bash run_cosyvoice3.sh 0 3 +``` + +### Benchmark with client-server mode + +To benchmark the running Triton server, run stage 4: +```sh +bash run_cosyvoice3.sh 4 4 + +# You can customize parameters such as the number of tasks inside the script. +``` +The following results were obtained by decoding on a single L20 GPU. + +#### Streaming TTS (Concurrent Tasks = 4) + +**First Chunk Latency** + +| Concurrent Tasks | Average (ms) | 50th Percentile (ms) | 90th Percentile (ms) | 95th Percentile (ms) | 99th Percentile (ms) | +| ---------------- | ------------ | -------------------- | -------------------- | -------------------- | -------------------- | +| 4 | 750.42 | 740.31 | 941.05 | 977.55 | 1002.37 | + +### Benchmark with offline inference mode + +For offline inference mode benchmark, please run stage 5: +```sh +bash run_cosyvoice3.sh 5 5 +``` + +#### Offline TTS (CosyVoice3 0.5B LLM + Token2Wav with TensorRT) + +| Backend | Batch Size | llm_time (s) | token2wav_time (s) | pipeline_time (s) | RTF | +|---------|------------|--------------|--------------------|--------------------|--------| +| TRTLLM | 1 | 13.21 | 5.72 | 19.48 | 0.1091 | +| TRTLLM | 2 | 8.46 | 6.02 | 14.91 | 0.0822 | +| TRTLLM | 4 | 5.07 | 5.95 | 11.43 | 0.0630 | +| TRTLLM | 8 | 2.98 | 6.11 | 9.53 | 0.0562 | +| TRTLLM | 16 | 2.12 | 6.27 | 8.83 | 0.0501 | diff --git a/runtime/triton_trtllm/README.md b/runtime/triton_trtllm/README.md new file mode 100644 index 0000000..eafaf0a --- /dev/null +++ b/runtime/triton_trtllm/README.md @@ -0,0 +1,37 @@ +# Accelerating CosyVoice with NVIDIA Triton Inference Server and TensorRT-LLM + +Contributed by Yuekai Zhang (NVIDIA). + +This repository provides three acceleration solutions for CosyVoice, each targeting a different model version and Token2Wav architecture. All solutions use TensorRT-LLM for LLM acceleration and NVIDIA Triton Inference Server for serving. + +## Solutions + +### [CosyVoice3](README.Cosyvoice3.md) + +Acceleration solution for [Fun-CosyVoice3-0.5B-2512](https://huggingface.co/FunAudioLLM/Fun-CosyVoice3-0.5B-2512), the latest CosyVoice model. The pipeline includes `audio_tokenizer`, `speaker_embedding`, `token2wav`, and `vocoder` modules managed by Triton, with the LLM served via `trtllm-serve`. + +### [CosyVoice2 + UNet Token2Wav](README.Cosyvoice2.Unet.md) + +The baseline acceleration solution for CosyVoice2, using the original UNet-based flow-matching Token2Wav module. + +### [CosyVoice2 + DiT Token2Wav](README.Cosyvoice2.DiT.md) + +Replaces the UNet Token2Wav with a DiT-based Token2Wav module from [Step-Audio2](https://github.com/stepfun-ai/Step-Audio-2). Supports disaggregated deployment where the LLM and Token2Wav run on separate GPUs for better resource utilization under high concurrency. + + + +## Quick Start + +Each solution can be launched with a single Docker Compose command: + +```sh +# CosyVoice3 +docker compose -f docker-compose.cosyvoice3.yml up + +# CosyVoice2 + UNet Token2Wav +docker compose -f docker-compose.cosyvoice2.unet.yml up + +# CosyVoice2 + DiT Token2Wav +docker compose -f docker-compose.cosyvoice2.dit.yml up +``` + diff --git a/runtime/triton_trtllm/docker-compose.cosyvoice3.yml b/runtime/triton_trtllm/docker-compose.cosyvoice3.yml new file mode 100644 index 0000000..5ebb977 --- /dev/null +++ b/runtime/triton_trtllm/docker-compose.cosyvoice3.yml @@ -0,0 +1,20 @@ +services: + tts: + image: soar97/triton-cosyvoice:25.06 + shm_size: '1gb' + ports: + - "8000:8000" + - "8001:8001" + - "8002:8002" + environment: + - PYTHONIOENCODING=utf-8 + - MODEL_ID=${MODEL_ID} + deploy: + resources: + reservations: + devices: + - driver: nvidia + device_ids: ['0'] + capabilities: [gpu] + command: > + /bin/bash -c "cd /workspace && git clone https://github.com/FunAudioLLM/CosyVoice.git && cd CosyVoice && git submodule update --init --recursive && cd runtime/triton_trtllm && bash run_cosyvoice3.sh 0 3" \ No newline at end of file diff --git a/runtime/triton_trtllm/run_cosyvoice3.sh b/runtime/triton_trtllm/run_cosyvoice3.sh index ac42c56..21871e5 100644 --- a/runtime/triton_trtllm/run_cosyvoice3.sh +++ b/runtime/triton_trtllm/run_cosyvoice3.sh @@ -1,8 +1,7 @@ #!/bin/bash # Copyright (c) 2026 NVIDIA (authors: Yuekai Zhang) export CUDA_VISIBLE_DEVICES=0 -# cosyvoice_path=/workspace/CosyVoice -cosyvoice_path=/workspace_yuekai/tts/CosyVoice +cosyvoice_path=/workspace/CosyVoice export PYTHONPATH=${cosyvoice_path}:$PYTHONPATH export PYTHONPATH=${cosyvoice_path}/third_party/Matcha-TTS:$PYTHONPATH @@ -24,7 +23,6 @@ bls_instance_num=10 if [ $stage -le -1 ] && [ $stop_stage -ge -1 ]; then echo "Cloning CosyVoice" - pip3 install --upgrade x_transformers s3tokenizer git clone --recursive https://github.com/FunAudioLLM/CosyVoice.git $cosyvoice_path cd $cosyvoice_path git submodule update --init --recursive @@ -33,6 +31,10 @@ fi if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then echo "Downloading CosyVoice3 Checkpoints" + # if s3 tokenizer version is not 0.3.0 + if [ $(pip3 show s3tokenizer | grep -o "0\.2\.[0-9]") != "0.3.0" ]; then + pip3 install --upgrade x_transformers s3tokenizer + fi huggingface-cli download --local-dir $huggingface_llm_local_dir yuekai/Fun-CosyVoice3-0.5B-2512-LLM-HF huggingface-cli download --local-dir $cosyvoice3_official_model_dir yuekai/Fun-CosyVoice3-0.5B-2512-FP16-ONNX huggingface-cli download --local-dir $cosyvoice3_official_model_dir FunAudioLLM/Fun-CosyVoice3-0.5B-2512 @@ -76,7 +78,7 @@ if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then LLM_TOKENIZER_DIR=$huggingface_llm_local_dir BLS_INSTANCE_NUM=$bls_instance_num TRITON_MAX_BATCH_SIZE=1 - DECOUPLED_MODE=True + DECOUPLED_MODE=True # False for offline TTS python3 scripts/fill_template.py -i ${model_repo}/cosyvoice3/config.pbtxt model_dir:${MODEL_DIR},bls_instance_num:${BLS_INSTANCE_NUM},llm_tokenizer_dir:${LLM_TOKENIZER_DIR},triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},decoupled_mode:${DECOUPLED_MODE},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS} python3 scripts/fill_template.py -i ${model_repo}/token2wav/config.pbtxt model_dir:${MODEL_DIR},triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS} @@ -111,17 +113,17 @@ if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then fi if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then - echo "stage 10: Python script CosyVoice3 TTS (LLM + CosyVoice3 Token2Wav) inference" + echo "stage 5: Python script CosyVoice3 TTS (LLM + CosyVoice3 Token2Wav) inference" datasets=(wenetspeech4tts) # wenetspeech4tts - backend=trtllm-serve # hf, trtllm, vllm, trtllm-serve + backend=trtllm # hf, trtllm, vllm, trtllm-serve - batch_sizes=(1) + batch_sizes=(16 8 4 2 1) token2wav_batch_size=1 # Only support 1 for now for batch_size in ${batch_sizes[@]}; do for dataset in ${datasets[@]}; do - output_dir=./cosyvoice3_${dataset}_${backend}_llm_batch_size_${batch_size}_token2wav_batch_size_${token2wav_batch_size}_streaming_trt + output_dir=./cosyvoice3_${dataset}_${backend}_llm_batch_size_${batch_size}_token2wav_batch_size_${token2wav_batch_size}_offline_tts_trt CUDA_VISIBLE_DEVICES=0 \ python3 infer_cosyvoice3.py \ --output-dir $output_dir \ @@ -130,8 +132,8 @@ if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then --backend $backend \ --batch-size $batch_size --token2wav-batch-size $token2wav_batch_size \ --engine-dir $trt_engines_dir \ - --enable-trt --streaming\ - --epoch 1 \ + --enable-trt \ + --epoch 3 \ --split-name ${dataset} || exit 1 done done From f4d376ef912e5cc87afeae49896cc6bd5a7f4fb0 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 16 Mar 2026 15:23:07 +0800 Subject: [PATCH 4/5] update results --- runtime/triton_trtllm/README.Cosyvoice3.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtime/triton_trtllm/README.Cosyvoice3.md b/runtime/triton_trtllm/README.Cosyvoice3.md index 278ec96..063bf07 100644 --- a/runtime/triton_trtllm/README.Cosyvoice3.md +++ b/runtime/triton_trtllm/README.Cosyvoice3.md @@ -78,7 +78,7 @@ bash run_cosyvoice3.sh 5 5 #### Offline TTS (CosyVoice3 0.5B LLM + Token2Wav with TensorRT) -| Backend | Batch Size | llm_time (s) | token2wav_time (s) | pipeline_time (s) | RTF | +| Backend | LLM Batch Size | llm_time (s) | token2wav_time (s) | pipeline_time (s) | RTF | |---------|------------|--------------|--------------------|--------------------|--------| | TRTLLM | 1 | 13.21 | 5.72 | 19.48 | 0.1091 | | TRTLLM | 2 | 8.46 | 6.02 | 14.91 | 0.0822 | From 914454e3fd4d4ede130451e3c9ccd8423c52f98f Mon Sep 17 00:00:00 2001 From: root Date: Mon, 16 Mar 2026 15:29:34 +0800 Subject: [PATCH 5/5] fix lint --- cosyvoice/dataset/processor.py | 2 +- cosyvoice/vllm/cosyvoice2.py | 2 +- example.py | 2 +- .../scripts/convert_cosyvoice3_to_hf.py | 80 +++++++++---------- 4 files changed, 43 insertions(+), 43 deletions(-) diff --git a/cosyvoice/dataset/processor.py b/cosyvoice/dataset/processor.py index deba209..545c007 100644 --- a/cosyvoice/dataset/processor.py +++ b/cosyvoice/dataset/processor.py @@ -181,7 +181,7 @@ def compute_fbank(data, def compute_whisper_fbank(data, num_frames=-1, mode='train'): - """ Extract whisper fbank + """ Extract whisper fbank Args: data: Iterable[{key, wav, label, sample_rate}] diff --git a/cosyvoice/vllm/cosyvoice2.py b/cosyvoice/vllm/cosyvoice2.py index 4d87123..fff0543 100644 --- a/cosyvoice/vllm/cosyvoice2.py +++ b/cosyvoice/vllm/cosyvoice2.py @@ -99,7 +99,7 @@ class CosyVoice2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): sampling_metadata: Optional[SamplingMetadata] = None, ) -> Optional[torch.Tensor]: if VLLM_V1_ENGINE_ONLY: - logits = self.logits_processor(self.lm_head, hidden_states, + logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias) else: logits = self.logits_processor(self.lm_head, hidden_states, diff --git a/example.py b/example.py index 7e9dd98..8752267 100644 --- a/example.py +++ b/example.py @@ -99,7 +99,7 @@ def cosyvoice3_example(): # 歴史的世界においては、過去は単に過ぎ去ったものではない、プラトンのいう如く非有が有である。 -> レキシ テキ セカイ ニ オイ テ ワ、カコ ワ タンニ スギサッ タ モノ デ ワ ナイ、プラトン ノ イウ ゴトク ヒ ユー ガ ユー デ アル。 for i, j in enumerate(cosyvoice.inference_cross_lingual('You are a helpful assistant.<|endofprompt|>レキシ テキ セカイ ニ オイ テ ワ、カコ ワ タンニ スギサッ タ モノ デ ワ ナイ、プラトン ノ イウ ゴトク ヒ ユー ガ ユー デ アル。', './asset/zero_shot_prompt.wav', stream=False)): - torchaudio.save('japanese_{}.wav'.format(i), j['tts_speech'], cosyvoice.sample_rate) + torchaudio.save('japanese_{}.wav'.format(i), j['tts_speech'], cosyvoice.sample_rate) def main(): diff --git a/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py b/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py index 40b99fe..37c8c28 100644 --- a/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py +++ b/runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py @@ -73,27 +73,27 @@ def load_cosyvoice3_model(model_dir: str): """Load CosyVoice3 model for weight extraction.""" from hyperpyyaml import load_hyperpyyaml from cosyvoice.utils.class_utils import get_model_type - + hyper_yaml_path = os.path.join(model_dir, 'cosyvoice3.yaml') hf_llm_dir = os.path.join(model_dir, 'CosyVoice-BlankEN') - + if not os.path.exists(hyper_yaml_path): raise ValueError(f'{hyper_yaml_path} not found!') - + with open(hyper_yaml_path, 'r') as f: configs = load_hyperpyyaml( - f, + f, overrides={'qwen_pretrain_path': hf_llm_dir} ) - + # Load LLM only llm = configs['llm'] llm_weights_path = os.path.join(model_dir, 'llm.pt') llm.load_state_dict(torch.load(llm_weights_path, map_location='cpu'), strict=True) llm.eval() - + logger.info(f"Loaded CosyVoice3 LLM from {model_dir}") - + return llm, hf_llm_dir, configs @@ -125,23 +125,23 @@ def convert_cosyvoice3_to_hf( dtype: Data type for saving """ logger.info(f"Loading CosyVoice3 model from {model_dir}") - + # 1. Load CosyVoice3 components cosyvoice3_llm, hf_llm_dir, configs = load_cosyvoice3_model(model_dir) - + # Extract key components qwen_model = cosyvoice3_llm.llm.model # Qwen2ForCausalLM speech_embedding = cosyvoice3_llm.speech_embedding # Embedding for speech tokens llm_decoder = cosyvoice3_llm.llm_decoder # Linear for decoding to speech tokens - + speech_token_size = get_speech_token_size(cosyvoice3_llm) logger.info(f"Speech token size: {speech_token_size}") - + # 2. Load tokenizer and add CosyVoice3 text special tokens + speech tokens tokenizer = AutoTokenizer.from_pretrained(hf_llm_dir, trust_remote_code=True) base_vocab_size = len(tokenizer) logger.info(f"Base tokenizer vocab size: {base_vocab_size}") - + # IMPORTANT: # - In CosyVoice3, LLM speech special tokens (sos/eos/task_id/fill) are INSIDE speech_embedding, # i.e. represented as <|s_6561|>, <|s_6562|>, <|s_6563|>, <|s_6564|>. @@ -185,7 +185,7 @@ def convert_cosyvoice3_to_hf( tokenizer.add_special_tokens(special_tokens) text_vocab_size = len(tokenizer) logger.info(f"Tokenizer vocab after CosyVoice3 text special tokens: {text_vocab_size}") - + # Add speech tokens: <|s_0|>, <|s_1|>, ..., <|s_{embedding_size-1}|> # IMPORTANT: This range must match speech_embedding.num_embeddings (includes speech special tokens). actual_speech_tokens = speech_token_size # Full embedding size (with speech special tokens) @@ -204,37 +204,37 @@ def convert_cosyvoice3_to_hf( assert "" not in speech_tokens assert "" not in speech_tokens tokenizer.add_tokens(speech_tokens) - + new_vocab_size = len(tokenizer) logger.info(f"New tokenizer vocab size: {new_vocab_size}") logger.info(f"Added {new_vocab_size - base_vocab_size} tokens total (text special + speech tokens)") - + # 3. Resize embeddings in Qwen model # Align to 128 for TensorRT efficiency padded_vocab_size = ((new_vocab_size + 127) // 128) * 128 qwen_model.resize_token_embeddings(padded_vocab_size) logger.info(f"Resized embeddings to: {padded_vocab_size}") - + # Speech tokens start after text vocab (base + CosyVoice3 text special tokens) speech_token_offset = text_vocab_size # 4. Copy speech_embedding into extended embed_tokens input_embeddings = qwen_model.get_input_embeddings() hidden_size = input_embeddings.weight.shape[1] - + logger.info(f"Hidden size: {hidden_size}") logger.info(f"speech_embedding shape: {speech_embedding.weight.shape}") logger.info(f"llm_decoder shape: {llm_decoder.weight.shape}") - + with torch.no_grad(): # Copy speech_embedding weights into embed_tokens # Indices: [speech_token_offset, speech_token_offset + speech_token_size) src_size = min(speech_embedding.weight.shape[0], actual_speech_tokens) input_embeddings.weight[speech_token_offset:speech_token_offset + src_size] = \ speech_embedding.weight[:src_size].to(input_embeddings.weight.dtype) - + logger.info(f"Copied speech_embedding to embed_tokens[{speech_token_offset}:{speech_token_offset + src_size}]") - + # 5. Create new lm_head with extended vocab and copy llm_decoder # Original lm_head: hidden_size -> original_vocab_size # New lm_head: hidden_size -> padded_vocab_size @@ -247,7 +247,7 @@ def convert_cosyvoice3_to_hf( out_features=padded_vocab_size, bias=has_bias ) - + with torch.no_grad(): # Initialize weights: # - Text part: copy from original lm_head (or zeros) @@ -258,42 +258,42 @@ def convert_cosyvoice3_to_hf( new_lm_head.weight.data.zero_() if has_bias: new_lm_head.bias.data.fill_(-float('inf')) - + # Copy original lm_head for text tokens (optional) original_lm_head = qwen_model.lm_head if original_lm_head is not None and original_lm_head.weight.shape[0] >= text_vocab_size: new_lm_head.weight[:text_vocab_size] = original_lm_head.weight[:text_vocab_size] if has_bias and original_lm_head.bias is not None: new_lm_head.bias[:text_vocab_size] = original_lm_head.bias[:text_vocab_size] - + # Copy llm_decoder for speech tokens decoder_size = min(llm_decoder.weight.shape[0], actual_speech_tokens) new_lm_head.weight[speech_token_offset:speech_token_offset + decoder_size] = \ llm_decoder.weight[:decoder_size].to(new_lm_head.weight.dtype) - + if has_bias: new_lm_head.bias[speech_token_offset:speech_token_offset + decoder_size] = \ llm_decoder.bias[:decoder_size].to(new_lm_head.bias.dtype) else: # If llm_decoder has no bias but we want it for text tokens pass - + # Replace lm_head qwen_model.lm_head = new_lm_head - + logger.info(f"Created new lm_head with shape: {new_lm_head.weight.shape}") logger.info(f"Copied llm_decoder to lm_head[{speech_token_offset}:{speech_token_offset + decoder_size}]") - + # 6. Update model configuration qwen_model.config.vocab_size = padded_vocab_size qwen_model.config.tie_word_embeddings = False # Embeddings and lm_head are now different! - + # Set EOS token for generation (speech EOS lives inside speech_embedding as <|s_{base_speech_token_size+1}|>) base_speech_token_size = getattr(cosyvoice3_llm, "speech_token_size", 6561) eos_speech_idx = base_speech_token_size + 1 eos_id = speech_token_offset + eos_speech_idx qwen_model.config.eos_token_id = eos_id - + # Generation settings qwen_model.generation_config.eos_token_id = eos_id qwen_model.generation_config.pad_token_id = eos_id @@ -302,7 +302,7 @@ def convert_cosyvoice3_to_hf( qwen_model.generation_config.top_k = 25 qwen_model.generation_config.repetition_penalty = 1.1 qwen_model.generation_config.max_new_tokens = 2048 - + # 7. Convert to target dtype dtype_map = { "float16": torch.float16, @@ -311,16 +311,16 @@ def convert_cosyvoice3_to_hf( } target_dtype = dtype_map[dtype] qwen_model.to(target_dtype) - + # 8. Save model and tokenizer os.makedirs(output_dir, exist_ok=True) - + qwen_model.save_pretrained(output_dir) - + TEMPLATE = "{%- for message in messages %}{%- if message['role'] == 'user' %}{{- '<|sos|>' + message['content'] + '<|task_id|>' }}{%- elif message['role'] == 'assistant' %}{{- message['content']}}{%- endif %}{%- endfor %}" tokenizer.chat_template = TEMPLATE tokenizer.save_pretrained(output_dir) - + # Save metadata for TRT-LLM inference metadata = { "original_vocab_size": base_vocab_size, @@ -332,30 +332,30 @@ def convert_cosyvoice3_to_hf( "speech_token_offset": speech_token_offset, "dtype": dtype, } - + import json with open(os.path.join(output_dir, "cosyvoice3_metadata.json"), "w") as f: json.dump(metadata, f, indent=2) - + logger.info(f"Saved HuggingFace model to {output_dir}") logger.info(f"Metadata: {metadata}") - + return output_dir, metadata def main(): args = parse_args() - + output_dir = args.output_dir if output_dir is None: output_dir = os.path.join(args.model_dir, "hf_merged") - + convert_cosyvoice3_to_hf( model_dir=args.model_dir, output_dir=output_dir, dtype=args.dtype, ) - + print("\n" + "=" * 70) print("✅ Conversion complete!") print("=" * 70)