update http client; launch script

runtime/triton_trtllm/README.md

@@ -1,38 +1,38 @@
-## Triton Inference Serving Best Practice for F5 TTS
-
-### Model Training
-See [official F5-TTS](https://github.com/SWivid/F5-TTS) or [Icefall F5-TTS](https://github.com/k2-fsa/icefall/tree/master/egs/wenetspeech4tts/TTS#f5-tts).
+## Nvidia Triton Inference Serving Best Practice for Spark TTS
 
 ### Quick Start
 Directly launch the service using docker compose.
 ```sh
-# VOCODER vocos or bigvgan
-VOCODER=vocos docker compose up
+docker compose up
 ```
 
 ### Build Image
 Build the docker image from scratch. 
 ```sh
-docker build . -f Dockerfile.server -t soar97/triton-f5-tts:24.12
+docker build . -f Dockerfile.server -t soar97/triton-spark-tts:25.02
 ```
 
 ### Create Docker Container
 ```sh
 your_mount_dir=/mnt:/mnt
-docker run -it --name "f5-server" --gpus all --net host -v $your_mount_dir --shm-size=2g soar97/triton-f5-tts:24.12
+docker run -it --name "spark-tts-server" --gpus all --net host -v $your_mount_dir --shm-size=2g soar97/triton-spark-tts:25.02
 ```
 
 ### Export Models to TensorRT-LLM and Launch Server
-Inside docker container, we would follow the official guide of TensorRT-LLM to build qwen and whisper TensorRT-LLM engines. See [here](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/whisper).
+Inside docker container, we would follow the official guide of TensorRT-LLM to build TensorRT-LLM engines. See [here](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/qwen).
 
 ```sh
-bash build_server.sh
+bash run.sh 0 3
+```
+### Simple HTTP client
+```sh
+python3 client_http.py
 ```
 
 ### Benchmark using Dataset
 ```sh
 num_task=2
-python3 client.py --num-tasks $num_task --huggingface-dataset yuekai/seed_tts --split-name wenetspeech4tts
+python3 client_grpc.py --num-tasks $num_task --huggingface-dataset yuekai/seed_tts --split-name wenetspeech4tts
 ```
 
 ### Benchmark Results
@@ -40,8 +40,4 @@ Decoding on a single L20 GPU, using 26 different prompt_audio/target_text pairs.
 
 | Model | Note   | Concurrency | Avg Latency     | RTF | 
 |-------|-----------|-----------------------|---------|--|
-| F5-TTS Base (Vocos) | [Code Commit](https://github.com/yuekaizhang/sherpa/tree/329ab3c573252e835844bea38505c6b43e994cf4/triton/f5_tts) | 1                   | 253 ms | 0.0394|
-
-### Credits
-1. [F5-TTS](https://github.com/SWivid/F5-TTS)
-2. [F5-TTS-TRTLLM](https://github.com/Bigfishering/f5-tts-trtllm)
+| Spark-TTS-0.5B | [Code Commit]() | 4                   | 253 ms | 0.0394|

runtime/triton_trtllm/build.sh

@@ -1,28 +1,73 @@
 
-# pip install -r /workspace_yuekai/spark-tts/Spark-TTS/requirements.txt
-export PYTHONPATH=/workspace_yuekai/spark-tts/Spark-TTS/
+
+export PYTHONPATH=../../../Spark-TTS/
+export CUDA_VISIBLE_DEVICES=0
+stage=$1
+stop_stage=$2
+echo "Start stage: $stage, Stop stage: $stop_stage"
+
+huggingface_model_local_dir=../../pretrained_models/Spark-TTS-0.5B
+trt_dtype=bfloat16
+trt_weights_dir=./tllm_checkpoint_${trt_dtype}
+trt_engines_dir=./trt_engines_${trt_dtype}
 
 model_repo=./model_repo_test
-rm -rf $model_repo
-cp -r ./model_repo $model_repo
 
-ENGINE_PATH=/workspace_yuekai/spark-tts/TensorRT-LLM/examples/qwen/Spark-TTS-0.5B_trt_engines_1gpu_bfloat16
-MAX_QUEUE_DELAY_MICROSECONDS=0
-MODEL_DIR=/workspace_yuekai/spark-tts/Spark-TTS/pretrained_models/Spark-TTS-0.5B
-LLM_TOKENIZER_DIR=/workspace_yuekai/spark-tts/Spark-TTS/pretrained_models/Spark-TTS-0.5B/LLM
-BLS_INSTANCE_NUM=4
-TRITON_MAX_BATCH_SIZE=16
+if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
+    echo "Downloading Spark-TTS-0.5B from HuggingFace"
+    hugginface-cli download SparkAudio/Spark-TTS-0.5B --local-dir $huggingface_model_local_dir || exit 1
+    # pip install -r /workspace_yuekai/spark-tts/Spark-TTS/requirements.txt
+fi
+
+
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+    echo "Converting checkpoint to TensorRT weights"
+    python scripts/convert_checkpoint.py --model_dir $huggingface_model_local_dir/LLM \
+                                --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 16 \
+                --max_num_tokens 32768 \
+                --gemm_plugin $trt_dtype || exit 1
+fi
+
+if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
+    echo "Creating model repository"
+    rm -rf $model_repo
+    cp -r ./model_repo $model_repo
 
-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}/spark_tts/config.pbtxt bls_instance_num:${BLS_INSTANCE_NUM},llm_tokenizer_dir:${LLM_TOKENIZER_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}/tensorrt_llm/config.pbtxt triton_backend:tensorrtllm,triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},decoupled_mode:False,max_beam_width:1,engine_dir:${ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:True,enable_kv_cache_reuse:False,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS},encoder_input_features_data_type:TYPE_FP16,logits_datatype:TYPE_FP32
+    ENGINE_PATH=$trt_engines_dir
+    MAX_QUEUE_DELAY_MICROSECONDS=0
+    MODEL_DIR=$huggingface_model_local_dir
+    LLM_TOKENIZER_DIR=$huggingface_model_local_dir/LLM
+    BLS_INSTANCE_NUM=4
+    TRITON_MAX_BATCH_SIZE=16
 
+    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}/spark_tts/config.pbtxt bls_instance_num:${BLS_INSTANCE_NUM},llm_tokenizer_dir:${LLM_TOKENIZER_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}/tensorrt_llm/config.pbtxt triton_backend:tensorrtllm,triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},decoupled_mode:False,max_beam_width:1,engine_dir:${ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:True,enable_kv_cache_reuse:False,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS},encoder_input_features_data_type:TYPE_FP16,logits_datatype:TYPE_FP32
 
-CUDA_VISIBLE_DEVICES=0 tritonserver --model-repository ${model_repo}
+fi
 
+if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
+    echo "Starting Triton server"
+    tritonserver --model-repository ${model_repo}
+fi
 
 
+if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
+    echo "Running client"
+    num_task=4
+    python3 client_grpc.py \
+        --server-addr localhost \
+        --model-name spark_tts \
+        --num-tasks $num_task \
+        --log-dir ./log_${num_task}
+fi
 
 
 

runtime/triton_trtllm/client_http.py

@@ -2,45 +2,139 @@ import requests
 import soundfile as sf
 import json
 import numpy as np
+import argparse
 
-url = "http://localhost:8000/v2/models/infer_pipeline/infer"
-wav_path = "*********"
-waveform, sr = sf.read(wav_path)
-lang_id = 54
-samples = np.array([waveform], dtype=np.float32)
-lengths = np.array([[len(waveform)]], dtype=np.int32)
-lang_id = np.array([[lang_id]], dtype=np.int8)
-
-data = {
-    "inputs":[
-        {
-            "name": "WAV",
-            "shape": samples.shape,
-            "datatype": "FP32",
-            "data": samples.tolist()
-        },
-        {
-            "name": "WAV_LENS",
-            "shape": lengths.shape,
-            "datatype": "INT32",
-            "data": lengths.tolist(),
-        },
-        {
-            "name": "LANG_ID",
-            "shape": lang_id.shape,
-            "datatype": "INT8",
-            "data": lang_id.tolist()
-        }
-    ]
-}
-rsp = requests.post(
-    url,
-    headers={"Content-Type": "application/json"},
-    json=data,
-    verify=False,
-    params={"request_id": '0'}
-)
-result = rsp.json()
-print(result)
-transcripts = result["outputs"][0]["data"][0]
-print(transcripts)
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--server-url",
+        type=str,
+        default="localhost:8000",
+        help="Address of the server",
+    )
+
+    parser.add_argument(
+        "--reference-audio",
+        type=str,
+        default="../../example/prompt_audio.wav",
+        help="Path to a single audio file. It can't be specified at the same time with --manifest-dir",
+    )
+
+    parser.add_argument(
+        "--reference-text",
+        type=str,
+        default="吃燕窝就选燕之屋，本节目由26年专注高品质燕窝的燕之屋冠名播出。豆奶牛奶换着喝，营养更均衡，本节目由豆本豆豆奶特约播出。",
+        help="",
+    )
+
+    parser.add_argument(
+        "--target-text",
+        type=str,
+        default="身临其境，换新体验。塑造开源语音合成新范式，让智能语音更自然。",
+        help="",
+    )
+
+    parser.add_argument(
+        "--model-name",
+        type=str,
+        default="spark_tts",
+        choices=[
+            "f5_tts", "spark_tts"
+        ],
+        help="triton model_repo module name to request: transducer for k2, attention_rescoring for wenet offline, streaming_wenet for wenet streaming, infer_pipeline for paraformer large offline",
+    )
+
+    parser.add_argument(
+        "--output-audio",
+        type=str,
+        default="output.wav",
+        help="Path to save the output audio",
+    )
+    return parser.parse_args()
+
+def prepare_request(
+    waveform,
+    reference_text,
+    target_text,
+    sample_rate=16000,
+    padding_duration: int = None,
+    audio_save_dir: str = "./",
+):
+    assert len(waveform.shape) == 1, "waveform should be 1D"
+    lengths = np.array([[len(waveform)]], dtype=np.int32)
+    if padding_duration:
+        # padding to nearset 10 seconds
+        samples = np.zeros(
+            (
+                1,
+                padding_duration
+                * sample_rate
+                * ((int(duration) // padding_duration) + 1),
+            ),
+            dtype=np.float32,
+        )
+
+        samples[0, : len(waveform)] = waveform
+    else:
+        samples = waveform
+        
+    samples = samples.reshape(1, -1).astype(np.float32)
+
+    data = {
+        "inputs":[
+            {
+                "name": "reference_wav",
+                "shape": samples.shape,
+                "datatype": "FP32",
+                "data": samples.tolist()
+            },
+            {
+                "name": "reference_wav_len",
+                "shape": lengths.shape,
+                "datatype": "INT32",
+                "data": lengths.tolist(),
+            },
+            {
+                "name": "reference_text",
+                "shape": [1, 1],
+                "datatype": "BYTES",
+                "data": [reference_text]
+            },
+            {
+                "name": "target_text",
+                "shape": [1, 1],
+                "datatype": "BYTES",
+                "data": [target_text]
+            }
+        ]
+    }
+
+    return data
+
+if __name__ == "__main__":
+    args = get_args()
+    server_url = args.server_url
+    if not server_url.startswith(("http://", "https://")):
+        server_url = f"http://{server_url}"
+    
+    url = f"{server_url}/v2/models/{args.model_name}/infer"
+    waveform, sr = sf.read(args.reference_audio)
+    assert sr == 16000, "sample rate hardcoded in server"
+    
+    samples = np.array(waveform, dtype=np.float32)
+    data = prepare_request(samples, args.reference_text, args.target_text)
+
+    rsp = requests.post(
+        url,
+        headers={"Content-Type": "application/json"},
+        json=data,
+        verify=False,
+        params={"request_id": '0'}
+    )
+    result = rsp.json()
+    audio = result["outputs"][0]["data"]
+    audio = np.array(audio, dtype=np.float32)
+    sf.write(args.output_audio, audio, 16000, "PCM_16")

runtime/triton_trtllm/docker-compose.yml

@@ -1,6 +1,6 @@
 services:
   tts:
-    image: soar97/triton-f5-tts:24.12
+    image: soar97/triton-spark-tts:25.02
     shm_size: '1gb'
     ports:
       - "8000:8000"

runtime/triton_trtllm/run.sh

@@ -0,0 +1,76 @@
+
+
+export PYTHONPATH=../../../Spark-TTS/
+export CUDA_VISIBLE_DEVICES=0
+stage=$1
+stop_stage=$2
+echo "Start stage: $stage, Stop stage: $stop_stage"
+
+huggingface_model_local_dir=../../pretrained_models/Spark-TTS-0.5B
+trt_dtype=bfloat16
+trt_weights_dir=./tllm_checkpoint_${trt_dtype}
+trt_engines_dir=./trt_engines_${trt_dtype}
+
+model_repo=./model_repo_test
+
+if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
+    echo "Downloading Spark-TTS-0.5B from HuggingFace"
+    hugginface-cli download SparkAudio/Spark-TTS-0.5B --local-dir $huggingface_model_local_dir || exit 1
+    # pip install -r /workspace_yuekai/spark-tts/Spark-TTS/requirements.txt
+fi
+
+
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+    echo "Converting checkpoint to TensorRT weights"
+    python convert_checkpoint.py --model_dir $huggingface_model_local_dir/LLM \
+                                --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 16 \
+                --max_num_tokens 32768 \
+                --gemm_plugin $trt_dtype || exit 1
+fi
+
+if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
+    echo "Creating model repository"
+    rm -rf $model_repo
+    cp -r ./model_repo $model_repo
+
+    ENGINE_PATH=$trt_engines_dir
+    MAX_QUEUE_DELAY_MICROSECONDS=0
+    MODEL_DIR=$huggingface_model_local_dir
+    LLM_TOKENIZER_DIR=$huggingface_model_local_dir/LLM
+    BLS_INSTANCE_NUM=4
+    TRITON_MAX_BATCH_SIZE=16
+
+    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}/spark_tts/config.pbtxt bls_instance_num:${BLS_INSTANCE_NUM},llm_tokenizer_dir:${LLM_TOKENIZER_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}/tensorrt_llm/config.pbtxt triton_backend:tensorrtllm,triton_max_batch_size:${TRITON_MAX_BATCH_SIZE},decoupled_mode:False,max_beam_width:1,engine_dir:${ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:True,enable_kv_cache_reuse:False,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MICROSECONDS},encoder_input_features_data_type:TYPE_FP16,logits_datatype:TYPE_FP32
+
+fi
+
+if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
+    echo "Starting Triton server"
+    tritonserver --model-repository ${model_repo}
+fi
+
+
+if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
+    echo "Running client"
+    num_task=4
+    python3 client_grpc.py \
+        --server-addr localhost \
+        --model-name spark_tts \
+        --num-tasks $num_task \
+        --log-dir ./log_${num_task}
+fi
+
+
+
+
+
+

runtime/triton_trtllm/scripts/build_engine.sh

@@ -1,46 +0,0 @@
-
-
-
-# model_dir=./Qwen2.5-0.5B-Instruct/
-# output_dir=./tllm_checkpoint_1gpu_fp16
-# trt_engines_dir=./trt_engines
-
-model_dir=/workspace_yuekai/spark-tts/Spark-TTS/pretrained_models/Spark-TTS-0.5B/LLM
-base_name=Spark-TTS-0.5B
-dtype=bfloat16
-output_dir=./${base_name}_tllm_checkpoint_1gpu_${dtype}
-trt_engines_dir=./${base_name}_trt_engines_1gpu_${dtype}
-
-
-# python convert_checkpoint.py --model_dir $model_dir \
-#                               --output_dir $output_dir \
-#                               --dtype $dtype || exit 1
-
-trtllm-build --checkpoint_dir $output_dir \
-            --output_dir $trt_engines_dir \
-            --max_batch_size 16 \
-            --max_num_tokens 32768 \
-            --gemm_plugin $dtype || exit 1
-# trtllm-build --checkpoint_dir $output_dir \
-#             --output_dir $trt_engines_dir \
-#             --max_batch_size 16 \
-#             --max_num_tokens 32768 \
-#             --gemm_plugin $dtype || exit 1
-
-python3 ../run.py --input_file  /workspace_yuekai/spark-tts/Spark-TTS/model_inputs.npy \
-                  --max_output_len=1500 \
-                  --tokenizer_dir $model_dir \
-                  --top_k 50 \
-                  --top_p 0.95 \
-                  --temperature 0.8 \
-                  --output_npy ./output.npy \
-                  --engine_dir=$trt_engines_dir || exit 1
-
-
-# python3 ../run.py --input_file  /workspace_yuekai/spark-tts/Spark-TTS/model_inputs.npy \
-#                   --max_output_len=1500 \
-#                   --tokenizer_dir $model_dir \
-#                   --top_k 50 \
-#                   --top_p 0.95 \
-#                   --temperature 0.8 \
-#                   --engine_dir=$trt_engines_dir || exit 1