| 1 | stage=$1 |
| 2 | stop_stage=$2 |
| 3 | model=$3 # F5TTS_v1_Base | F5TTS_Base | F5TTS_v1_Small | F5TTS_Small |
| 4 | if [ -z "$model" ]; then |
| 5 | model=F5TTS_v1_Base |
| 6 | fi |
| 7 | echo "Start stage: $stage, Stop stage: $stop_stage, Model: $model" |
| 8 | export CUDA_VISIBLE_DEVICES=0 |
| 9 | |
| 10 | CKPT_DIR=../../../../ckpts |
| 11 | TRTLLM_CKPT_DIR=$CKPT_DIR/$model/trtllm_ckpt |
| 12 | TRTLLM_ENGINE_DIR=$CKPT_DIR/$model/trtllm_engine |
| 13 | |
| 14 | VOCODER_ONNX_PATH=$CKPT_DIR/vocos_vocoder.onnx |
| 15 | VOCODER_TRT_ENGINE_PATH=$CKPT_DIR/vocos_vocoder.plan |
| 16 | MODEL_REPO=./model_repo |
| 17 | |
| 18 | if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then |
| 19 | echo "Downloading F5-TTS from huggingface" |
| 20 | huggingface-cli download SWivid/F5-TTS $model/model_*.* $model/vocab.txt --local-dir $CKPT_DIR |
| 21 | fi |
| 22 | |
| 23 | ckpt_file=$(ls $CKPT_DIR/$model/model_*.* 2>/dev/null | sort -V | tail -1) # default select latest update |
| 24 | vocab_file=$CKPT_DIR/$model/vocab.txt |
| 25 | |
| 26 | if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then |
| 27 | echo "Converting checkpoint" |
| 28 | python3 scripts/convert_checkpoint.py \ |
| 29 | --pytorch_ckpt $ckpt_file \ |
| 30 | --output_dir $TRTLLM_CKPT_DIR --model_name $model |
| 31 | python_package_path=/usr/local/lib/python3.12/dist-packages |
| 32 | cp -r patch/* $python_package_path/tensorrt_llm/models |
| 33 | trtllm-build --checkpoint_dir $TRTLLM_CKPT_DIR \ |
| 34 | --max_batch_size 8 \ |
| 35 | --output_dir $TRTLLM_ENGINE_DIR --remove_input_padding disable |
| 36 | fi |
| 37 | |
| 38 | if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then |
| 39 | echo "Exporting vocos vocoder" |
| 40 | python3 scripts/export_vocoder_to_onnx.py --vocoder vocos --output-path $VOCODER_ONNX_PATH |
| 41 | bash scripts/export_vocos_trt.sh $VOCODER_ONNX_PATH $VOCODER_TRT_ENGINE_PATH |
| 42 | fi |
| 43 | |
| 44 | if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then |
| 45 | echo "Building triton server" |
| 46 | rm -r $MODEL_REPO |
| 47 | cp -r ./model_repo_f5_tts $MODEL_REPO |
| 48 | python3 scripts/fill_template.py -i $MODEL_REPO/f5_tts/config.pbtxt vocab:$vocab_file,model:$ckpt_file,trtllm:$TRTLLM_ENGINE_DIR,vocoder:vocos |
| 49 | cp $VOCODER_TRT_ENGINE_PATH $MODEL_REPO/vocoder/1/vocoder.plan |
| 50 | fi |
| 51 | |
| 52 | if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then |
| 53 | echo "Starting triton server" |
| 54 | tritonserver --model-repository=$MODEL_REPO |
| 55 | fi |
| 56 | |
| 57 | if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then |
| 58 | echo "Testing triton server" |
| 59 | num_task=1 |
| 60 | split_name=wenetspeech4tts |
| 61 | log_dir=./tests/client_grpc_${model}_concurrent_${num_task}_${split_name} |
| 62 | rm -r $log_dir |
| 63 | python3 client_grpc.py --num-tasks $num_task --huggingface-dataset yuekai/seed_tts --split-name $split_name --log-dir $log_dir |
| 64 | fi |
| 65 | |
| 66 | if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then |
| 67 | echo "Testing http client" |
| 68 | audio=../../infer/examples/basic/basic_ref_en.wav |
| 69 | reference_text="Some call me nature, others call me mother nature." |
| 70 | target_text="I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring." |
| 71 | python3 client_http.py --reference-audio $audio --reference-text "$reference_text" --target-text "$target_text" --output-audio "./tests/client_http_$model.wav" |
| 72 | fi |
| 73 | |
| 74 | if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then |
| 75 | echo "TRT-LLM: offline decoding benchmark test" |
| 76 | batch_size=2 |
| 77 | split_name=wenetspeech4tts |
| 78 | backend_type=trt |
| 79 | log_dir=./tests/benchmark_${model}_batch_size_${batch_size}_${split_name}_${backend_type} |
| 80 | rm -r $log_dir |
| 81 | torchrun --nproc_per_node=1 \ |
| 82 | benchmark.py --output-dir $log_dir \ |
| 83 | --batch-size $batch_size \ |
| 84 | --enable-warmup \ |
| 85 | --split-name $split_name \ |
| 86 | --model-path $ckpt_file \ |
| 87 | --vocab-file $vocab_file \ |
| 88 | --vocoder-trt-engine-path $VOCODER_TRT_ENGINE_PATH \ |
| 89 | --backend-type $backend_type \ |
| 90 | --tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1 |
| 91 | fi |
| 92 | |
| 93 | if [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then |
| 94 | echo "Native Pytorch: offline decoding benchmark test" |
| 95 | if ! python3 -c "import f5_tts" &> /dev/null; then |
| 96 | pip install -e ../../../../ |
| 97 | fi |
| 98 | batch_size=1 # set attn_mask_enabled=True if batching in actual use case |
| 99 | split_name=wenetspeech4tts |
| 100 | backend_type=pytorch |
| 101 | log_dir=./tests/benchmark_${model}_batch_size_${batch_size}_${split_name}_${backend_type} |
| 102 | rm -r $log_dir |
| 103 | torchrun --nproc_per_node=1 \ |
| 104 | benchmark.py --output-dir $log_dir \ |
| 105 | --batch-size $batch_size \ |
| 106 | --split-name $split_name \ |
| 107 | --enable-warmup \ |
| 108 | --model-path $ckpt_file \ |
| 109 | --vocab-file $vocab_file \ |
| 110 | --backend-type $backend_type \ |
| 111 | --tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1 |
| 112 | fi |