Multi gpu

src/f5_tts/runtime/triton_trtllm/docker-compose.yml

@@ -1,7 +1,7 @@
 services:
   tts:
     image: soar97/triton-f5-tts:24.12
-    shm_size: '1gb'
+    shm_size: '2gb'
     ports:
       - "8000:8000"
       - "8001:8001"
@@ -14,7 +14,7 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              device_ids: ['0']
+              device_ids: ['0', '1']
               capabilities: [gpu]
     command: >
-      /bin/bash -c "pip install vocos && rm -rf F5-TTS && git clone https://github.com/SWivid/F5-TTS.git && cd F5-TTS/src/f5_tts/runtime/triton_trtllm/ && bash run.sh 0 4 $MODEL"
+      /bin/bash -c "pip install vocos rjieba && rm -rf F5-TTS && git clone --branch multi-gpu https://github.com/webadmin2021/F5-TTS.git F5-TTS && cd F5-TTS/src/f5_tts/runtime/triton_trtllm/ && bash run.sh 0 4 $MODEL"

src/f5_tts/runtime/triton_trtllm/model_repo_f5_tts/f5_tts/config.pbtxt

@@ -78,4 +78,10 @@ instance_group [
     count: 1
     kind: KIND_GPU
   }
-]
+]
+parameters: {
+  key: "gpu_device_ids"
+  value: {
+    string_value: "0;1"
+  }
+}

src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py

@@ -106,7 +106,7 @@ def prepare_inputs(self, **kwargs):
         max_batch_size = kwargs["max_batch_size"]
         batch_size_range = [2, 2, max_batch_size]
         mel_size = self.config.mel_dim
-        max_seq_len = 3000  # 4096
+        max_seq_len = 4096  # 4096
         num_frames_range = [mel_size * 2, max_seq_len * 2, max_seq_len * max_batch_size]
         concat_feature_dim = mel_size + self.config.text_dim
         freq_embed_dim = 256  # Warning: hard coding 256 here

src/f5_tts/runtime/triton_trtllm/run.sh

@@ -5,7 +5,7 @@ if [ -z "$model" ]; then
     model=F5TTS_v1_Base
 fi
 echo "Start stage: $stage, Stop stage: $stop_stage, Model: $model"
-export CUDA_VISIBLE_DEVICES=0
+export CUDA_VISIBLE_DEVICES=0,1
 
 CKPT_DIR=../../../../ckpts
 TRTLLM_CKPT_DIR=$CKPT_DIR/$model/trtllm_ckpt
@@ -109,4 +109,4 @@ if [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then
     --vocab-file $vocab_file \
     --backend-type $backend_type \
     --tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1
-fi
+fi

src/f5_tts/runtime/triton_trtllm/scripts/export_vocos_trt.sh

@@ -30,7 +30,7 @@ MAX_BATCH_SIZE=8
 
 MIN_INPUT_LENGTH=1
 OPT_INPUT_LENGTH=1000
-MAX_INPUT_LENGTH=3000  # 4096
+MAX_INPUT_LENGTH=4096  # 4096
 
 MEL_MIN_SHAPE="${MIN_BATCH_SIZE}x100x${MIN_INPUT_LENGTH}"
 MEL_OPT_SHAPE="${OPT_BATCH_SIZE}x100x${OPT_INPUT_LENGTH}"