clean codes

runtime/triton_trtllm/Dockerfile.server

@@ -0,0 +1,9 @@
+FROM nvcr.io/nvidia/tritonserver:25.02-trtllm-python-py3
+RUN pip install tritonclient[grpc] librosa
+WORKDIR /workspace
+
+
+
+
+
+

runtime/triton_trtllm/README.md

@@ -0,0 +1,47 @@
+## 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).
+
+### Quick Start
+Directly launch the service using docker compose.
+```sh
+# VOCODER vocos or bigvgan
+VOCODER=vocos docker compose up
+```
+
+### Build Image
+Build the docker image from scratch. 
+```sh
+docker build . -f Dockerfile.server -t soar97/triton-f5-tts:24.12
+```
+
+### 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
+```
+
+### 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).
+
+```sh
+bash build_server.sh
+```
+
+### Benchmark using Dataset
+```sh
+num_task=2
+python3 client.py --num-tasks $num_task --huggingface-dataset yuekai/seed_tts --split-name wenetspeech4tts
+```
+
+### Benchmark Results
+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)

runtime/triton_trtllm/build.sh

@@ -1,15 +1,31 @@
 
-pip install -r /workspace_yuekai/spark-tts/Spark-TTS/requirements.txt
+# pip install -r /workspace_yuekai/spark-tts/Spark-TTS/requirements.txt
+export PYTHONPATH=/workspace_yuekai/spark-tts/Spark-TTS/
+
 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
-gpu_device_ids=0
-python3 fill_template.py -i ${model_repo}/tensorrt_llm/config.pbtxt gpu_device_ids:${gpu_device_ids},triton_backend:tensorrtllm,triton_max_batch_size:16,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
+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
+
+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}
+
+
+
+
+
+
+
+
 
-# enable_context_fmha_fp32_acc:${ENABLE_CONTEXT_FMHA_FP32_ACC}
-export PYTHONPATH=/workspace_yuekai/spark-tts/Spark-TTS/
-CUDA_VISIBLE_DEVICES=${gpu_device_ids} tritonserver --model-repository ${model_repo}

runtime/triton_trtllm/client_http.py

@@ -0,0 +1,46 @@
+import requests
+import soundfile as sf
+import json
+import numpy as np
+
+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)

runtime/triton_trtllm/docker-compose.yml

@@ -0,0 +1,20 @@
+services:
+  tts:
+    image: soar97/triton-f5-tts:24.12
+    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 "rm -rf sherpa && git clone https://github.com/yuekaizhang/sherpa.git -b f5 && cd sherpa/triton/f5_tts/ && bash build_server.sh $VOCODER"

runtime/triton_trtllm/model_repo/audio_tokenizer/1/model.py

@@ -25,80 +25,113 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 import json
 import torch
-from torch import nn
-from torch.nn.utils.rnn import pad_sequence
-import torch.nn.functional as F
-from torch.utils.dlpack import from_dlpack, to_dlpack
+from torch.utils.dlpack import to_dlpack
 
 import triton_python_backend_utils as pb_utils
 
-import math
 import os
-from functools import wraps
 import numpy as np
 
-
 from sparktts.models.audio_tokenizer import BiCodecTokenizer
 
 class TritonPythonModel:
+    """Triton Python model for audio tokenization.
+    
+    This model takes reference audio input and extracts semantic and global tokens
+    using BiCodec tokenizer.
+    """
+
     def initialize(self, args):
+        """Initialize the model.
+        
+        Args:
+            args: Dictionary containing model configuration
+        """
+        # Parse model parameters
         parameters = json.loads(args['model_config'])['parameters']
-        for key, value in parameters.items():
-            parameters[key] = value["string_value"]
-        model_dir = parameters["model_dir"]
+        model_params = {k: v["string_value"] for k, v in parameters.items()}
+        
+        # Initialize tokenizer
         self.device = torch.device("cuda")
-        self.audio_tokenizer = BiCodecTokenizer(model_dir, device=self.device)
+        self.audio_tokenizer = BiCodecTokenizer(model_params["model_dir"], 
+                                              device=self.device)
 
     def get_ref_clip(self, wav: np.ndarray) -> np.ndarray:
-        """Get reference audio clip for speaker embedding."""
-        sample_rate = 16000
-        ref_segment_duration = 6
-        latent_hop_length = 320
+        """Extract reference audio clip for speaker embedding.
+        
+        Args:
+            wav: Input waveform array
+            
+        Returns:
+            Reference clip of fixed duration
+        """
+        SAMPLE_RATE = 16000
+        REF_SEGMENT_DURATION = 6  # seconds
+        LATENT_HOP_LENGTH = 320
 
         ref_segment_length = (
-            int(sample_rate * ref_segment_duration)
-            // latent_hop_length
-            * latent_hop_length
+            int(SAMPLE_RATE * REF_SEGMENT_DURATION)
+            // LATENT_HOP_LENGTH
+            * LATENT_HOP_LENGTH
         )
         wav_length = len(wav)
 
         if ref_segment_length > wav_length:
-            # Repeat and truncate to handle insufficient length
-            wav = np.tile(wav, ref_segment_length // wav_length + 1)
+            # Repeat and truncate if input is too short
+            repeat_times = ref_segment_length // wav_length + 1
+            wav = np.tile(wav, repeat_times)
 
         return wav[:ref_segment_length]
 
     def execute(self, requests):
-        reference_wav_list, reference_wav_ref_clip_list = [], []
+        """Execute inference on the batched requests.
+        
+        Args:
+            requests: List of inference requests
+            
+        Returns:
+            List of inference responses containing tokenized outputs
+        """
+        reference_wav_list = []
+        reference_wav_ref_clip_list = []
 
+        # Process each request in batch
         for request in requests:
-            wav_array = pb_utils.get_input_tensor_by_name(request, "reference_wav").as_numpy()
+            # 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()
-            # check shape
-            print(wav_array.shape, wav_len, 233333333333)
-            # squeeze the first dimension, for the numpy array
+
+            # Prepare inputs
             wav = wav_array[:, :wav_len].squeeze(0)
             reference_wav_list.append(wav)
-
+            
             wav_ref_clip = self.get_ref_clip(wav)
-            print(wav_ref_clip.shape, 2333333333455)
             reference_wav_ref_clip_list.append(torch.from_numpy(wav_ref_clip))
-        # (len,) -> B,len   
+
+        # Batch process through tokenizer
         ref_wav_clip_tensor = torch.stack(reference_wav_ref_clip_list, dim=0)
-        wav2vec2_features = self.audio_tokenizer.extract_wav2vec2_features(reference_wav_list)
-        audio_tokenizer_input_dict = {
-            "ref_wav": ref_wav_clip_tensor.to(self.device), # no padding, spaker encoder
+        wav2vec2_features = self.audio_tokenizer.extract_wav2vec2_features(
+            reference_wav_list)
+        
+        audio_tokenizer_input = {
+            "ref_wav": ref_wav_clip_tensor.to(self.device),
             "feat": wav2vec2_features.to(self.device),
         }
-        semantic_tokens, global_tokens = self.audio_tokenizer.model.tokenize(audio_tokenizer_input_dict)
-
+        semantic_tokens, global_tokens = self.audio_tokenizer.model.tokenize(
+            audio_tokenizer_input)
 
+        # Prepare responses
         responses = []
         for i in range(len(requests)):
-            global_tokens_tensor = pb_utils.Tensor.from_dlpack("global_tokens", to_dlpack(global_tokens[i]))
-            semantic_tokens_tensor = pb_utils.Tensor.from_dlpack("semantic_tokens", to_dlpack(semantic_tokens[i]))
-            inference_response = pb_utils.InferenceResponse(output_tensors=[global_tokens_tensor, semantic_tokens_tensor])
+            global_tokens_tensor = pb_utils.Tensor.from_dlpack(
+                "global_tokens", to_dlpack(global_tokens[i]))
+            semantic_tokens_tensor = pb_utils.Tensor.from_dlpack(
+                "semantic_tokens", to_dlpack(semantic_tokens[i]))
+            
+            inference_response = pb_utils.InferenceResponse(
+                output_tensors=[global_tokens_tensor, semantic_tokens_tensor])
             responses.append(inference_response)
                              
         return responses

runtime/triton_trtllm/model_repo/audio_tokenizer/config.pbtxt

@@ -14,14 +14,14 @@
 
 name: "audio_tokenizer"
 backend: "python"
-max_batch_size: 16
+max_batch_size: ${triton_max_batch_size}
 dynamic_batching {
-    max_queue_delay_microseconds: 1
+    max_queue_delay_microseconds: ${max_queue_delay_microseconds}
 }
 parameters [
   {
    key: "model_dir", 
-   value: {string_value:"/workspace_yuekai/spark-tts/Spark-TTS/pretrained_models/Spark-TTS-0.5B"}
+   value: {string_value:"${model_dir}"}
   }
 ]
 

runtime/triton_trtllm/model_repo/spark_tts/1/model.py

@@ -23,31 +23,23 @@
 # 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 os
+import re
+from typing import Dict, List, Tuple, Optional, Union
+
+import numpy as np
 import torch
-from torch import nn
-from torch.nn.utils.rnn import pad_sequence
-import torch.nn.functional as F
 from torch.utils.dlpack import from_dlpack, to_dlpack
-
 import triton_python_backend_utils as pb_utils
-
-import math
-import os
-from functools import wraps
-
 from transformers import AutoTokenizer
 
-import numpy as np
-import re
-from typing import Tuple
-
-from sparktts.utils.token_parser import LEVELS_MAP, GENDER_MAP, TASK_TOKEN_MAP
-
+from sparktts.utils.token_parser import TASK_TOKEN_MAP
 
 def process_prompt(
     text: str,
-    prompt_text: str = None,
+    prompt_text: Optional[str] = None,
     global_token_ids: torch.Tensor = None,
     semantic_token_ids: torch.Tensor = None,
 ) -> Tuple[str, torch.Tensor]:
@@ -55,27 +47,27 @@ def process_prompt(
     Process input for voice cloning.
 
     Args:
-        text (str): The text input to be converted to speech.
-        prompt_speech_path (Path): Path to the audio file used as a prompt.
-        prompt_text (str, optional): Transcript of the prompt audio.
+        text: The text input to be converted to speech.
+        prompt_text: Transcript of the prompt audio.
+        global_token_ids: Global token IDs extracted from reference audio.
+        semantic_token_ids: Semantic token IDs extracted from reference audio.
 
-    Return:
-        Tuple[str, torch.Tensor]: Input prompt; global tokens
+    Returns:
+        Tuple containing the formatted input prompt and global token IDs.
     """
-
-    # global_token_ids, semantic_token_ids = self.audio_tokenizer.tokenize(
-    #     prompt_speech_path
-    # )
+    # Convert global tokens to string format
     global_tokens = "".join(
         [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze()]
     )
-    print(global_tokens, 233333333333, len(global_tokens), "global_tokens")
+
+    
     # Prepare the input tokens for the model
     if prompt_text is not None:
+        # Include semantic tokens when prompt text is provided
         semantic_tokens = "".join(
             [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze()]
         )
-        print(semantic_tokens, 233333333333, len(semantic_tokens), "semantic_tokens")
+
         inputs = [
             TASK_TOKEN_MAP["tts"],
             "<|start_content|>",
@@ -89,6 +81,7 @@ def process_prompt(
             semantic_tokens,
         ]
     else:
+        # Without prompt text, exclude semantic tokens
         inputs = [
             TASK_TOKEN_MAP["tts"],
             "<|start_content|>",
@@ -99,17 +92,31 @@ def process_prompt(
             "<|end_global_token|>",
         ]
 
+    # Join all input components into a single string
     inputs = "".join(inputs)
-
     return inputs, global_token_ids
 
+
 class TritonPythonModel:
+    """Triton Python model for Spark TTS.
+    
+    This model orchestrates the end-to-end TTS pipeline by coordinating
+    between audio tokenizer, LLM, and vocoder components.
+    """
+    
     def initialize(self, args):
+        """Initialize the model.
+        
+        Args:
+            args: Dictionary containing model configuration
+        """
+        # Parse model parameters
         parameters = json.loads(args['model_config'])['parameters']
-        for key, value in parameters.items():
-            parameters[key] = value["string_value"]
-        model_dir = parameters["model_dir"]
-        self.tokenizer = AutoTokenizer.from_pretrained(f"{model_dir}/LLM")
+        model_params = {k: v["string_value"] for k, v in parameters.items()}
+        
+        # Initialize tokenizer
+        llm_tokenizer_dir = model_params["llm_tokenizer_dir"]
+        self.tokenizer = AutoTokenizer.from_pretrained(llm_tokenizer_dir)
         self.device = torch.device("cuda")
         self.decoupled = False
 
@@ -140,7 +147,6 @@ class TritonPythonModel:
         """
         # convert input_ids to numpy, with shape [1, sequence_length]
         input_ids = input_ids.cpu().numpy()
-        print(input_ids.shape, 233333333333, "input_ids")
         max_tokens = 512
         input_dict = {
             "request_output_len": np.array([[max_tokens]], dtype=np.int32),
@@ -153,135 +159,153 @@ class TritonPythonModel:
             "input_ids": input_ids,
             "input_lengths": np.array([[input_ids.shape[1]]], dtype=np.int32),
         }
-        for k, v in input_dict.items():
-            print(k, v.shape, 233333333333, v.dtype)
-        # exit()
+        
+        # Convert inputs to Triton tensors
         input_tensor_list = [
             pb_utils.Tensor(k, v) for k, v in input_dict.items()
         ]
-        # input_tensor_list.append(pb_utils.Tensor.from_dlpack(
-        #     "input_ids", to_dlpack(input_ids)
-        # ))
+        
+        # Create and execute inference request
         llm_request = pb_utils.InferenceRequest(
             model_name="tensorrt_llm",
             requested_output_names=["output_ids", "sequence_length"],
             inputs=input_tensor_list,
         )
-        print("=======================================")
+        
         llm_response = llm_request.exec(decoupled=self.decoupled)
         if llm_response.has_error():
-            raise pb_utils.TritonModelException(
-                llm_response.error().message())
+            raise pb_utils.TritonModelException(llm_response.error().message())
+        
+        # Extract and process output
         output_ids = pb_utils.get_output_tensor_by_name(
             llm_response, "output_ids").as_numpy()
         seq_lens = pb_utils.get_output_tensor_by_name(
             llm_response, "sequence_length").as_numpy()
-        print(seq_lens, 233333333333, "seq_lens")
-        actual_output_ids = output_ids[0][0]
-        actual_output_ids = actual_output_ids[:seq_lens[0][0]]
-        print(actual_output_ids, 233333333333, "actual_output_ids")
+        
+        # Get actual output IDs up to the sequence length
+        actual_output_ids = output_ids[0][0][:seq_lens[0][0]]
+        
         return actual_output_ids
 
     def forward_audio_tokenizer(self, wav, wav_len):
-        # input_tensor_0 = pb_utils.Tensor.
-        # input_tensor_1 = pb_utils.Tensor.from_dlpack("wav_len", to_dlpack(wav_len))
-    
+        """Forward pass through the audio tokenizer component.
+        
+        Args:
+            wav: Input waveform tensor
+            wav_len: Waveform length tensor
+            
+        Returns:
+            Tuple of global and semantic tokens
+        """
         inference_request = pb_utils.InferenceRequest(
             model_name='audio_tokenizer',
             requested_output_names=['global_tokens', 'semantic_tokens'],
             inputs=[wav, wav_len]
         )
+        
         inference_response = inference_request.exec()
         if inference_response.has_error():
             raise pb_utils.TritonModelException(inference_response.error().message())
-        else:
-            global_tokens = pb_utils.get_output_tensor_by_name(inference_response,
-                                                            'global_tokens')
-            global_tokens = torch.utils.dlpack.from_dlpack(global_tokens.to_dlpack()).cpu()
-            semantic_tokens = pb_utils.get_output_tensor_by_name(inference_response,
-                                                            'semantic_tokens')
-            semantic_tokens = torch.utils.dlpack.from_dlpack(semantic_tokens.to_dlpack()).cpu()
-            return global_tokens, semantic_tokens
+        
+        # Extract and convert output tensors
+        global_tokens = pb_utils.get_output_tensor_by_name(inference_response, 'global_tokens')
+        global_tokens = torch.utils.dlpack.from_dlpack(global_tokens.to_dlpack()).cpu()
+        
+        semantic_tokens = pb_utils.get_output_tensor_by_name(inference_response, 'semantic_tokens')
+        semantic_tokens = torch.utils.dlpack.from_dlpack(semantic_tokens.to_dlpack()).cpu()
+        
+        return global_tokens, semantic_tokens
 
-    def forward_vocoder(self, global_token_ids, pred_semantic_ids):
-        global_token_ids = pb_utils.Tensor.from_dlpack("global_tokens", to_dlpack(global_token_ids))
-        pred_semantic_ids = pb_utils.Tensor.from_dlpack("semantic_tokens", to_dlpack(pred_semantic_ids))
+    def forward_vocoder(self, global_token_ids: torch.Tensor, pred_semantic_ids: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the vocoder component.
+        
+        Args:
+            global_token_ids: Global token IDs tensor
+            pred_semantic_ids: Predicted semantic token IDs tensor
+            
+        Returns:
+            Generated waveform tensor
+        """
+        # Convert tensors to Triton format
+        global_token_ids_tensor = pb_utils.Tensor.from_dlpack("global_tokens", to_dlpack(global_token_ids))
+        pred_semantic_ids_tensor = pb_utils.Tensor.from_dlpack("semantic_tokens", to_dlpack(pred_semantic_ids))
+        
+        # Create and execute inference request
         inference_request = pb_utils.InferenceRequest(
             model_name='vocoder',
             requested_output_names=['waveform'],
-            inputs=[global_token_ids, pred_semantic_ids]
+            inputs=[global_token_ids_tensor, pred_semantic_ids_tensor]
         )
+        
         inference_response = inference_request.exec()
         if inference_response.has_error():
             raise pb_utils.TritonModelException(inference_response.error().message())
-        else:
-            waveform = pb_utils.get_output_tensor_by_name(inference_response,
-                                                        'waveform')
-            waveform = torch.utils.dlpack.from_dlpack(waveform.to_dlpack()).cpu()
-            return waveform
+        
+        # Extract and convert output waveform
+        waveform = pb_utils.get_output_tensor_by_name(inference_response, 'waveform')
+        waveform = torch.utils.dlpack.from_dlpack(waveform.to_dlpack()).cpu()
+        
+        return waveform
         
     def execute(self, requests):
-        # reference_text_list, target_text_list, reference_wav_list, reference_wav_ref_clip_list = [], [], [], []
+        """Execute inference on the batched requests.
+        
+        Args:
+            requests: List of inference requests
+            
+        Returns:
+            List of inference responses containing generated audio
+        """
         responses = []
+        
         for request in requests:
+            # Extract input tensors
             wav = pb_utils.get_input_tensor_by_name(request, "reference_wav")
-            wav_len = pb_utils.get_input_tensor_by_name(
-                request, "reference_wav_len")
+            wav_len = pb_utils.get_input_tensor_by_name(request, "reference_wav_len")
+            
+            # Process reference audio through audio tokenizer
             global_tokens, semantic_tokens = self.forward_audio_tokenizer(wav, wav_len)
-            # print(wav_tensor.shape, wav_len.shape, 233333333333)
-            # reference_wav_list.append(wav)
-            # wav_ref_clip = self.get_ref_clip(wav[:, :wav_len])
-            # reference_wav_ref_clip_list.append(wav_ref_clip)
-
-
-            reference_text = pb_utils.get_input_tensor_by_name(
-                request, "reference_text").as_numpy()
+            
+            # Extract text inputs
+            reference_text = pb_utils.get_input_tensor_by_name(request, "reference_text").as_numpy()
             reference_text = reference_text[0][0].decode('utf-8')
-            # reference_text_list.append(reference_text)
-
-            target_text = pb_utils.get_input_tensor_by_name(
-                request, "target_text").as_numpy()
-            target_text = target_text[0][0].decode('utf-8')
-            # target_text_list.append(target_text)
             
-            # ref_wav_clip_tensor = torch.cat(reference_wav_ref_clip_list, dim=0)
-            # wav2vec2_features = self.model.audio_tokenizer.extract_wav2vec2_features(reference_wav_list)
-            # audio_tokenizer_input_dict = {
-            #     "ref_wav": ref_wav_clip_tensor, # no padding, spaker encoder
-            #     "feat": wav2vec2_features,
-            # }
+            target_text = pb_utils.get_input_tensor_by_name(request, "target_text").as_numpy()
+            target_text = target_text[0][0].decode('utf-8')
             
+            # Prepare prompt for LLM
             prompt, global_token_ids = process_prompt(
                 text=target_text,
                 prompt_text=reference_text,
                 global_token_ids=global_tokens,
                 semantic_token_ids=semantic_tokens,
             )
-            print(semantic_tokens.shape, "semantic_tokens")
-            print(global_tokens.shape, "global_tokens")
-            print(prompt, "prompt", len(prompt))
+            
+            
+            # Tokenize prompt for LLM
             model_inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device)
-            print(model_inputs, "model_inputs")
             input_ids = model_inputs.input_ids.to(torch.int32)
-            print(input_ids.shape, 233333333333, 455555555)
-
+            
+            # Generate semantic tokens with LLM
             generated_ids = self.forward_llm(input_ids)
-            print(generated_ids, "generated_ids", len(generated_ids))
-            predicts = self.tokenizer.batch_decode([generated_ids], skip_special_tokens=True)[0]
-            print(predicts, "predicts", len(predicts))
+            
+            # Decode and extract semantic token IDs from generated text
+            predicted_text = self.tokenizer.batch_decode([generated_ids], skip_special_tokens=True)[0]
             pred_semantic_ids = (
-                torch.tensor([int(token) for token in re.findall(r"bicodec_semantic_(\d+)", predicts)])
+                torch.tensor([int(token) for token in re.findall(r"bicodec_semantic_(\d+)", predicted_text)])
                 .unsqueeze(0).to(torch.int32)
             )
-            print(global_token_ids.shape, "global_token_ids")
-            print(pred_semantic_ids.shape, "pred_semantic_ids")
+            
+
+            # Generate audio with vocoder
             audio = self.forward_vocoder(
                 global_token_ids.to(self.device),
                 pred_semantic_ids.to(self.device),
             )
-
-            audio = pb_utils.Tensor.from_dlpack("waveform", to_dlpack(audio))
-            inference_response = pb_utils.InferenceResponse(output_tensors=[audio])
+            
+            # Prepare response
+            audio_tensor = pb_utils.Tensor.from_dlpack("waveform", to_dlpack(audio))
+            inference_response = pb_utils.InferenceResponse(output_tensors=[audio_tensor])
             responses.append(inference_response)
                              
         return responses

runtime/triton_trtllm/model_repo/spark_tts/config.pbtxt

@@ -14,14 +14,14 @@
 
 name: "spark_tts"
 backend: "python"
-max_batch_size: 16
+max_batch_size: ${triton_max_batch_size}
 dynamic_batching {
-    max_queue_delay_microseconds: 1
+    max_queue_delay_microseconds: ${max_queue_delay_microseconds}
 }
 parameters [
   {
-   key: "model_dir", 
-   value: {string_value:"/workspace_yuekai/spark-tts/Spark-TTS/pretrained_models/Spark-TTS-0.5B"}
+   key: "llm_tokenizer_dir", 
+   value: {string_value:"${llm_tokenizer_dir}"}
   }
 ]
 
@@ -59,7 +59,7 @@ output [
 
 instance_group [
   {
-    count: 4
+    count: ${bls_instance_num}
     kind: KIND_CPU
   }
 ]

runtime/triton_trtllm/model_repo/vocoder/1/model.py

@@ -23,48 +23,80 @@
 # 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 os
+import logging
+from typing import List, Dict
+
 import torch
-from torch import nn
-from torch.nn.utils.rnn import pad_sequence
-import torch.nn.functional as F
-from torch.utils.dlpack import from_dlpack, to_dlpack
+from torch.utils.dlpack import to_dlpack
 
 import triton_python_backend_utils as pb_utils
 
-import math
-import os
-from functools import wraps
-
 from sparktts.models.bicodec import BiCodec
 
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 class TritonPythonModel:
+    """Triton Python model for vocoder.
+    
+    This model takes global and semantic tokens as input and generates audio waveforms
+    using the BiCodec vocoder.
+    """
+
     def initialize(self, args):
+        """Initialize the model.
+        
+        Args:
+            args: Dictionary containing model configuration
+        """
+        # Parse model parameters
         parameters = json.loads(args['model_config'])['parameters']
-        for key, value in parameters.items():
-            parameters[key] = value["string_value"]
-        model_dir = parameters["model_dir"]
-        self.device = torch.device("cuda")
-        self.vocoder = BiCodec.load_from_checkpoint(f"{model_dir}/BiCodec").to(
-            self.device
-        )
+        model_params = {key: value["string_value"] for key, value in parameters.items()}
+        model_dir = model_params["model_dir"]
+        
+        # Initialize device and vocoder
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"Initializing vocoder from {model_dir} on {self.device}")
+        
+        self.vocoder = BiCodec.load_from_checkpoint(f"{model_dir}/BiCodec")
+        del self.vocoder.encoder, self.vocoder.postnet
+        self.vocoder.eval().to(self.device)  # Set model to evaluation mode
+
+        logger.info("Vocoder initialized successfully")
+
 
     def execute(self, requests):
+        """Execute inference on the batched requests.
+        
+        Args:
+            requests: List of inference requests
+            
+        Returns:
+            List of inference responses containing generated waveforms
+        """
         global_tokens_list, semantic_tokens_list = [], []
 
+        # Process each request in batch
         for request in requests:
             global_tokens_tensor = pb_utils.get_input_tensor_by_name(request, "global_tokens").as_numpy()
             semantic_tokens_tensor = pb_utils.get_input_tensor_by_name(request, "semantic_tokens").as_numpy()
-            # check shape
             global_tokens_list.append(torch.from_numpy(global_tokens_tensor).to(self.device))
             semantic_tokens_list.append(torch.from_numpy(semantic_tokens_tensor).to(self.device))
 
+        # Concatenate tokens for batch processing
         global_tokens = torch.cat(global_tokens_list, dim=0)
         semantic_tokens = torch.cat(semantic_tokens_list, dim=0)
-        print(global_tokens.shape, semantic_tokens.shape, 233333333333, "global_tokens, semantic_tokens")
+        
 
-        wavs = self.vocoder.detokenize(semantic_tokens, global_tokens.unsqueeze(1))
+        # Generate waveforms
+        with torch.no_grad():
+            wavs = self.vocoder.detokenize(semantic_tokens, global_tokens.unsqueeze(1))
 
+        # Prepare responses
         responses = []
         for i in range(len(requests)):
             wav_tensor = pb_utils.Tensor.from_dlpack("waveform", to_dlpack(wavs[i]))

runtime/triton_trtllm/model_repo/vocoder/config.pbtxt

@@ -14,14 +14,14 @@
 
 name: "vocoder"
 backend: "python"
-max_batch_size: 16
+max_batch_size: ${triton_max_batch_size}
 dynamic_batching {
-    max_queue_delay_microseconds: 1
+    max_queue_delay_microseconds: ${max_queue_delay_microseconds}
 }
 parameters [
   {
    key: "model_dir", 
-   value: {string_value:"/workspace_yuekai/spark-tts/Spark-TTS/pretrained_models/Spark-TTS-0.5B"}
+   value: {string_value:"${model_dir}"}
   }
 ]