mirror of
https://github.com/FunAudioLLM/CosyVoice.git
synced 2026-09-24 22:59:55 +08:00
Merge pull request #1850 from yuekaizhang/cosy3_pr
Support Cosyvoice3 TRT-LLM Inference
This commit is contained in:
Xiang Lyu
@@ -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}]
|
||||
|
||||
@@ -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,
|
||||
|
||||
+1
-1
@@ -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():
|
||||
|
||||
+1
-1
@@ -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
|
||||
@@ -0,0 +1,146 @@
|
||||
## Accelerating CosyVoice 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.cosyvoice2.unet.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.sh`
|
||||
|
||||
The `run.sh` script orchestrates the entire workflow through numbered stages.
|
||||
|
||||
You can run a subset of stages with:
|
||||
```sh
|
||||
bash run.sh <start_stage> <stop_stage> [service_type]
|
||||
```
|
||||
- `<start_stage>`: The stage to start from (0-5).
|
||||
- `<stop_stage>`: The stage to stop after (0-5).
|
||||
|
||||
**Stages:**
|
||||
|
||||
- **Stage 0**: Downloads the `cosyvoice-2 0.5B` model from HuggingFace.
|
||||
- **Stage 1**: Converts the HuggingFace checkpoint to the TensorRT-LLM format and builds the TensorRT engines.
|
||||
- **Stage 2**: Creates the Triton model repository and configures the model files. The configuration is adjusted based on whether `Decoupled=True` (streaming) or `Decoupled=False` (offline) will be used.
|
||||
- **Stage 3**: Launches the Triton Inference Server.
|
||||
- **Stage 4**: Runs the single-utterance HTTP client for testing.
|
||||
- **Stage 5**: Runs the gRPC benchmark client.
|
||||
- **Stage 6**: Runs the offline 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.sh 0 3
|
||||
```
|
||||
> [!TIP]
|
||||
> Both streaming and offline (non-streaming) TTS modes are supported. For streaming TTS, set `Decoupled=True`. For offline TTS, set `Decoupled=False`. You need to rerun stage 2 if you switch between modes.
|
||||
|
||||
### Single-Utterance HTTP Client
|
||||
|
||||
Sends a single HTTP inference request. This is intended for testing the offline TTS mode (`Decoupled=False`):
|
||||
```sh
|
||||
bash run.sh 4 4
|
||||
```
|
||||
|
||||
### Benchmark with client-server mode
|
||||
|
||||
To benchmark the running Triton server, pass `streaming` or `offline` as the third argument:
|
||||
```sh
|
||||
bash run.sh 5 5 # [streaming|offline]
|
||||
|
||||
# You can also customize parameters such as the number of tasks and the dataset split:
|
||||
# python3 client_grpc.py --num-tasks 2 --huggingface-dataset yuekai/seed_tts_cosy2 --split-name test_zh --mode [streaming|offline]
|
||||
```
|
||||
> [!TIP]
|
||||
> It is recommended to run the benchmark multiple times to get stable results after the initial server warm-up.
|
||||
|
||||
### Benchmark with offline inference mode
|
||||
For offline inference mode benchmark, please check the below command:
|
||||
```sh
|
||||
# install FlashCosyVoice for token2wav batching
|
||||
# git clone https://github.com/yuekaizhang/FlashCosyVoice.git /workspace/FlashCosyVoice -b trt
|
||||
# cd /workspace/FlashCosyVoice
|
||||
# pip install -e .
|
||||
# cd -
|
||||
# wget https://huggingface.co/yuekai/cosyvoice2_flow_onnx/resolve/main/flow.decoder.estimator.fp32.dynamic_batch.onnx -O $model_scope_model_local_dir/flow.decoder.estimator.fp32.dynamic_batch.onnx
|
||||
|
||||
bash run.sh 6 6
|
||||
|
||||
# You can also switch to huggingface backend by setting backend=hf
|
||||
```
|
||||
|
||||
|
||||
### Benchmark Results
|
||||
The following results were obtained by decoding on a single L20 GPU with 26 prompt audio/target text pairs from the [yuekai/seed_tts](https://huggingface.co/datasets/yuekai/seed_tts) dataset (approximately 170 seconds of audio):
|
||||
|
||||
**Client-Server Mode: Streaming TTS (First Chunk Latency)**
|
||||
| Mode | Concurrency | Avg Latency (ms) | P50 Latency (ms) | RTF |
|
||||
|---|---|---|---|---|
|
||||
| Streaming, use_spk2info_cache=False | 1 | 220.43 | 218.07 | 0.1237 |
|
||||
| Streaming, use_spk2info_cache=False | 2 | 476.97 | 369.25 | 0.1022 |
|
||||
| Streaming, use_spk2info_cache=False | 4 | 1107.34 | 1243.75| 0.0922 |
|
||||
| Streaming, use_spk2info_cache=True | 1 | 189.88 | 184.81 | 0.1155 |
|
||||
| Streaming, use_spk2info_cache=True | 2 | 323.04 | 316.83 | 0.0905 |
|
||||
| Streaming, use_spk2info_cache=True | 4 | 977.68 | 903.68| 0.0733 |
|
||||
|
||||
> If your service only needs a fixed speaker, you can set `use_spk2info_cache=True` in `run.sh`. To add more speakers, refer to the instructions [here](https://github.com/qi-hua/async_cosyvoice?tab=readme-ov-file#9-spk2info-%E8%AF%B4%E6%98%8E).
|
||||
|
||||
**Client-Server Mode: Offline TTS (Full Sentence Latency)**
|
||||
| Mode | Note | Concurrency | Avg Latency (ms) | P50 Latency (ms) | RTF |
|
||||
|---|---|---|---|---|---|
|
||||
| Offline, Decoupled=False, use_spk2info_cache=False | [Commit](https://github.com/yuekaizhang/CosyVoice/commit/b44f12110224cb11c03aee4084b1597e7b9331cb) | 1 | 758.04 | 615.79 | 0.0891 |
|
||||
| Offline, Decoupled=False, use_spk2info_cache=False | [Commit](https://github.com/yuekaizhang/CosyVoice/commit/b44f12110224cb11c03aee4084b1597e7b9331cb) | 2 | 1025.93 | 901.68 | 0.0657 |
|
||||
| Offline, Decoupled=False, use_spk2info_cache=False | [Commit](https://github.com/yuekaizhang/CosyVoice/commit/b44f12110224cb11c03aee4084b1597e7b9331cb) | 4 | 1914.13 | 1783.58 | 0.0610 |
|
||||
|
||||
**Offline Inference Mode: Hugginface LLM V.S. TensorRT-LLM**
|
||||
| Backend | Batch Size | llm_time_seconds | total_time_seconds | RTF |
|
||||
|---------|------------|------------------|-----------------------|--|
|
||||
| HF | 1 | 39.26 | 44.31 | 0.2494 |
|
||||
| HF | 2 | 30.54 | 35.62 | 0.2064 |
|
||||
| HF | 4 | 18.63 | 23.90 | 0.1421 |
|
||||
| HF | 8 | 11.22 | 16.45 | 0.0947 |
|
||||
| HF | 16 | 8.42 | 13.78 | 0.0821 |
|
||||
| TRTLLM | 1 | 12.46 | 17.31 | 0.0987 |
|
||||
| TRTLLM | 2 | 7.64 |12.65 | 0.0739 |
|
||||
| TRTLLM | 4 | 4.89 | 9.38 | 0.0539 |
|
||||
| TRTLLM | 8 | 2.92 | 7.23 | 0.0418 |
|
||||
| TRTLLM | 16 | 2.01 | 6.63 | 0.0386 |
|
||||
### OpenAI-Compatible Server
|
||||
|
||||
To launch an OpenAI-compatible API service, run the following commands:
|
||||
```sh
|
||||
git clone https://github.com/yuekaizhang/Triton-OpenAI-Speech.git
|
||||
cd Triton-OpenAI-Speech
|
||||
pip install -r requirements.txt
|
||||
|
||||
# After the Triton service is running, start the FastAPI bridge:
|
||||
python3 tts_server.py --url http://localhost:8000 --ref_audios_dir ./ref_audios/ --port 10086 --default_sample_rate 24000
|
||||
|
||||
# Test the service with curl:
|
||||
bash test/test_cosyvoice.sh
|
||||
```
|
||||
> [!NOTE]
|
||||
> Currently, only the offline TTS mode is compatible with the OpenAI-compatible server.
|
||||
|
||||
### Acknowledgements
|
||||
|
||||
This work originates from the NVIDIA CISI project. For more multimodal resources, please see [mair-hub](https://github.com/nvidia-china-sae/mair-hub).
|
||||
|
||||
@@ -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 <start_stage> <stop_stage>
|
||||
```
|
||||
- `<start_stage>`: The stage to start from.
|
||||
- `<stop_stage>`: 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 | 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 |
|
||||
| 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 |
|
||||
+30
-139
@@ -1,146 +1,37 @@
|
||||
## Accelerating CosyVoice with NVIDIA Triton Inference Server and TensorRT-LLM
|
||||
# Accelerating CosyVoice with NVIDIA Triton Inference Server and TensorRT-LLM
|
||||
|
||||
Contributed by Yuekai Zhang (NVIDIA).
|
||||
|
||||
### Quick Start
|
||||
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:
|
||||
|
||||
Launch the service directly with Docker Compose:
|
||||
```sh
|
||||
docker compose up
|
||||
# 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
|
||||
```
|
||||
|
||||
### 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.sh`
|
||||
|
||||
The `run.sh` script orchestrates the entire workflow through numbered stages.
|
||||
|
||||
You can run a subset of stages with:
|
||||
```sh
|
||||
bash run.sh <start_stage> <stop_stage> [service_type]
|
||||
```
|
||||
- `<start_stage>`: The stage to start from (0-5).
|
||||
- `<stop_stage>`: The stage to stop after (0-5).
|
||||
|
||||
**Stages:**
|
||||
|
||||
- **Stage 0**: Downloads the `cosyvoice-2 0.5B` model from HuggingFace.
|
||||
- **Stage 1**: Converts the HuggingFace checkpoint to the TensorRT-LLM format and builds the TensorRT engines.
|
||||
- **Stage 2**: Creates the Triton model repository and configures the model files. The configuration is adjusted based on whether `Decoupled=True` (streaming) or `Decoupled=False` (offline) will be used.
|
||||
- **Stage 3**: Launches the Triton Inference Server.
|
||||
- **Stage 4**: Runs the single-utterance HTTP client for testing.
|
||||
- **Stage 5**: Runs the gRPC benchmark client.
|
||||
- **Stage 6**: Runs the offline 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.sh 0 3
|
||||
```
|
||||
> [!TIP]
|
||||
> Both streaming and offline (non-streaming) TTS modes are supported. For streaming TTS, set `Decoupled=True`. For offline TTS, set `Decoupled=False`. You need to rerun stage 2 if you switch between modes.
|
||||
|
||||
### Single-Utterance HTTP Client
|
||||
|
||||
Sends a single HTTP inference request. This is intended for testing the offline TTS mode (`Decoupled=False`):
|
||||
```sh
|
||||
bash run.sh 4 4
|
||||
```
|
||||
|
||||
### Benchmark with client-server mode
|
||||
|
||||
To benchmark the running Triton server, pass `streaming` or `offline` as the third argument:
|
||||
```sh
|
||||
bash run.sh 5 5 # [streaming|offline]
|
||||
|
||||
# You can also customize parameters such as the number of tasks and the dataset split:
|
||||
# python3 client_grpc.py --num-tasks 2 --huggingface-dataset yuekai/seed_tts_cosy2 --split-name test_zh --mode [streaming|offline]
|
||||
```
|
||||
> [!TIP]
|
||||
> It is recommended to run the benchmark multiple times to get stable results after the initial server warm-up.
|
||||
|
||||
### Benchmark with offline inference mode
|
||||
For offline inference mode benchmark, please check the below command:
|
||||
```sh
|
||||
# install FlashCosyVoice for token2wav batching
|
||||
# git clone https://github.com/yuekaizhang/FlashCosyVoice.git /workspace/FlashCosyVoice -b trt
|
||||
# cd /workspace/FlashCosyVoice
|
||||
# pip install -e .
|
||||
# cd -
|
||||
# wget https://huggingface.co/yuekai/cosyvoice2_flow_onnx/resolve/main/flow.decoder.estimator.fp32.dynamic_batch.onnx -O $model_scope_model_local_dir/flow.decoder.estimator.fp32.dynamic_batch.onnx
|
||||
|
||||
bash run.sh 6 6
|
||||
|
||||
# You can also switch to huggingface backend by setting backend=hf
|
||||
```
|
||||
|
||||
|
||||
### Benchmark Results
|
||||
The following results were obtained by decoding on a single L20 GPU with 26 prompt audio/target text pairs from the [yuekai/seed_tts](https://huggingface.co/datasets/yuekai/seed_tts) dataset (approximately 170 seconds of audio):
|
||||
|
||||
**Client-Server Mode: Streaming TTS (First Chunk Latency)**
|
||||
| Mode | Concurrency | Avg Latency (ms) | P50 Latency (ms) | RTF |
|
||||
|---|---|---|---|---|
|
||||
| Streaming, use_spk2info_cache=False | 1 | 220.43 | 218.07 | 0.1237 |
|
||||
| Streaming, use_spk2info_cache=False | 2 | 476.97 | 369.25 | 0.1022 |
|
||||
| Streaming, use_spk2info_cache=False | 4 | 1107.34 | 1243.75| 0.0922 |
|
||||
| Streaming, use_spk2info_cache=True | 1 | 189.88 | 184.81 | 0.1155 |
|
||||
| Streaming, use_spk2info_cache=True | 2 | 323.04 | 316.83 | 0.0905 |
|
||||
| Streaming, use_spk2info_cache=True | 4 | 977.68 | 903.68| 0.0733 |
|
||||
|
||||
> If your service only needs a fixed speaker, you can set `use_spk2info_cache=True` in `run.sh`. To add more speakers, refer to the instructions [here](https://github.com/qi-hua/async_cosyvoice?tab=readme-ov-file#9-spk2info-%E8%AF%B4%E6%98%8E).
|
||||
|
||||
**Client-Server Mode: Offline TTS (Full Sentence Latency)**
|
||||
| Mode | Note | Concurrency | Avg Latency (ms) | P50 Latency (ms) | RTF |
|
||||
|---|---|---|---|---|---|
|
||||
| Offline, Decoupled=False, use_spk2info_cache=False | [Commit](https://github.com/yuekaizhang/CosyVoice/commit/b44f12110224cb11c03aee4084b1597e7b9331cb) | 1 | 758.04 | 615.79 | 0.0891 |
|
||||
| Offline, Decoupled=False, use_spk2info_cache=False | [Commit](https://github.com/yuekaizhang/CosyVoice/commit/b44f12110224cb11c03aee4084b1597e7b9331cb) | 2 | 1025.93 | 901.68 | 0.0657 |
|
||||
| Offline, Decoupled=False, use_spk2info_cache=False | [Commit](https://github.com/yuekaizhang/CosyVoice/commit/b44f12110224cb11c03aee4084b1597e7b9331cb) | 4 | 1914.13 | 1783.58 | 0.0610 |
|
||||
|
||||
**Offline Inference Mode: Hugginface LLM V.S. TensorRT-LLM**
|
||||
| Backend | Batch Size | llm_time_seconds | total_time_seconds | RTF |
|
||||
|---------|------------|------------------|-----------------------|--|
|
||||
| HF | 1 | 39.26 | 44.31 | 0.2494 |
|
||||
| HF | 2 | 30.54 | 35.62 | 0.2064 |
|
||||
| HF | 4 | 18.63 | 23.90 | 0.1421 |
|
||||
| HF | 8 | 11.22 | 16.45 | 0.0947 |
|
||||
| HF | 16 | 8.42 | 13.78 | 0.0821 |
|
||||
| TRTLLM | 1 | 12.46 | 17.31 | 0.0987 |
|
||||
| TRTLLM | 2 | 7.64 |12.65 | 0.0739 |
|
||||
| TRTLLM | 4 | 4.89 | 9.38 | 0.0539 |
|
||||
| TRTLLM | 8 | 2.92 | 7.23 | 0.0418 |
|
||||
| TRTLLM | 16 | 2.01 | 6.63 | 0.0386 |
|
||||
### OpenAI-Compatible Server
|
||||
|
||||
To launch an OpenAI-compatible API service, run the following commands:
|
||||
```sh
|
||||
git clone https://github.com/yuekaizhang/Triton-OpenAI-Speech.git
|
||||
cd Triton-OpenAI-Speech
|
||||
pip install -r requirements.txt
|
||||
|
||||
# After the Triton service is running, start the FastAPI bridge:
|
||||
python3 tts_server.py --url http://localhost:8000 --ref_audios_dir ./ref_audios/ --port 10086 --default_sample_rate 24000
|
||||
|
||||
# Test the service with curl:
|
||||
bash test/test_cosyvoice.sh
|
||||
```
|
||||
> [!NOTE]
|
||||
> Currently, only the offline TTS mode is compatible with the OpenAI-compatible server.
|
||||
|
||||
### Acknowledgements
|
||||
|
||||
This work originates from the NVIDIA CISI project. For more multimodal resources, please see [mair-hub](https://github.com/nvidia-china-sae/mair-hub).
|
||||
|
||||
|
||||
@@ -281,6 +281,7 @@ def get_args():
|
||||
choices=[
|
||||
"f5_tts",
|
||||
"spark_tts",
|
||||
"cosyvoice3",
|
||||
"cosyvoice2",
|
||||
"cosyvoice2_dit"],
|
||||
help="triton model_repo module name to request",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
@@ -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())
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
@@ -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
|
||||
@@ -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 ]
|
||||
}
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
@@ -12,3 +12,4 @@ pyworld
|
||||
openai-whisper
|
||||
tritonclient
|
||||
modelscope
|
||||
x_transformers
|
||||
@@ -0,0 +1,140 @@
|
||||
#!/bin/bash
|
||||
# Copyright (c) 2026 NVIDIA (authors: Yuekai Zhang)
|
||||
export CUDA_VISIBLE_DEVICES=0
|
||||
cosyvoice_path=/workspace/CosyVoice
|
||||
|
||||
export PYTHONPATH=${cosyvoice_path}:$PYTHONPATH
|
||||
export PYTHONPATH=${cosyvoice_path}/third_party/Matcha-TTS:$PYTHONPATH
|
||||
|
||||
stage=$1
|
||||
stop_stage=$2
|
||||
|
||||
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=$cosyvoice_path/runtime/triton_trtllm/trt_weights_${trt_dtype}
|
||||
trt_engines_dir=$cosyvoice_path/runtime/triton_trtllm/trt_engines_${trt_dtype}
|
||||
|
||||
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"
|
||||
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 "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
|
||||
fi
|
||||
|
||||
|
||||
if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
|
||||
echo "Converting checkpoint to TensorRT weights"
|
||||
python3 scripts/convert_checkpoint.py --model_dir $huggingface_llm_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_llm_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=$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 # 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}
|
||||
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 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_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
|
||||
fi
|
||||
|
||||
if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
|
||||
echo "Running benchmark client for CosyVoice3"
|
||||
num_task=4
|
||||
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: Python script CosyVoice3 TTS (LLM + CosyVoice3 Token2Wav) inference"
|
||||
|
||||
datasets=(wenetspeech4tts) # wenetspeech4tts
|
||||
backend=trtllm # hf, trtllm, vllm, trtllm-serve
|
||||
|
||||
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}_offline_tts_trt
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
python3 infer_cosyvoice3.py \
|
||||
--output-dir $output_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 \
|
||||
--enable-trt \
|
||||
--epoch 3 \
|
||||
--split-name ${dataset} || exit 1
|
||||
done
|
||||
done
|
||||
fi
|
||||
@@ -0,0 +1,372 @@
|
||||
#!/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.
|
||||
"""
|
||||
Convert CosyVoice3 LLM to HuggingFace format with merged embeddings.
|
||||
|
||||
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
|
||||
|
||||
Then convert to TRT-LLM:
|
||||
trtllm-build --checkpoint_dir <output_dir> --output_dir <trt_engines_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: <model-dir>/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):
|
||||
"""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,
|
||||
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
|
||||
|
||||
|
||||
def get_speech_token_size(llm) -> int:
|
||||
"""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
|
||||
# Use full embedding size (includes speech special tokens)
|
||||
return speech_embedding_size
|
||||
|
||||
|
||||
def convert_cosyvoice3_to_hf(
|
||||
model_dir: str,
|
||||
output_dir: str,
|
||||
dtype: str = "bfloat16",
|
||||
):
|
||||
"""
|
||||
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: 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. 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|>.
|
||||
# - 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]', '<strong>', '</strong>', '[noise]',
|
||||
'[laughter]', '[cough]', '[clucking]', '[accent]',
|
||||
'[quick_breath]',
|
||||
"<laughter>", "</laughter>",
|
||||
"[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 <s_6561> to <|sos|>
|
||||
# replace <s_6562> to <|eos1|>
|
||||
# replace <s_6563> to <|task_id|>
|
||||
# replace <s_6564> 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 "<s_6561>" not in speech_tokens
|
||||
assert "<s_6562>" not in speech_tokens
|
||||
assert "<s_6563>" not in speech_tokens
|
||||
assert "<s_6564>" 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
|
||||
# llm_decoder: hidden_size -> speech_token_size
|
||||
|
||||
# Create new 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():
|
||||
# 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'))
|
||||
|
||||
# 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
|
||||
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. Convert to target dtype
|
||||
dtype_map = {
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
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,
|
||||
"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 <trt_weights_dir> --output_dir <trt_engines_dir> ...")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user