| 1 | # Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | # |
| 3 | # Redistribution and use in source and binary forms, with or without |
| 4 | # modification, are permitted provided that the following conditions |
| 5 | # are met: |
| 6 | # * Redistributions of source code must retain the above copyright |
| 7 | # notice, this list of conditions and the following disclaimer. |
| 8 | # * Redistributions in binary form must reproduce the above copyright |
| 9 | # notice, this list of conditions and the following disclaimer in the |
| 10 | # documentation and/or other materials provided with the distribution. |
| 11 | # * Neither the name of NVIDIA CORPORATION nor the names of its |
| 12 | # contributors may be used to endorse or promote products derived |
| 13 | # from this software without specific prior written permission. |
| 14 | # |
| 15 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY |
| 16 | # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE |
| 17 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR |
| 18 | # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR |
| 19 | # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, |
| 20 | # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, |
| 21 | # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR |
| 22 | # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY |
| 23 | # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| 24 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| 26 | import json |
| 27 | import os |
| 28 | |
| 29 | import rjieba |
| 30 | import torch |
| 31 | import torchaudio |
| 32 | import triton_python_backend_utils as pb_utils |
| 33 | from f5_tts_trtllm import F5TTS |
| 34 | from pypinyin import Style, lazy_pinyin |
| 35 | from torch.nn.utils.rnn import pad_sequence |
| 36 | from torch.utils.dlpack import from_dlpack, to_dlpack |
| 37 | |
| 38 | |
| 39 | def get_tokenizer(vocab_file_path: str): |
| 40 | """ |
| 41 | tokenizer - "pinyin" do g2p for only chinese characters, need .txt vocab_file |
| 42 | - "char" for char-wise tokenizer, need .txt vocab_file |
| 43 | - "byte" for utf-8 tokenizer |
| 44 | - "custom" if you're directly passing in a path to the vocab.txt you want to use |
| 45 | vocab_size - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols |
| 46 | - if use "char", derived from unfiltered character & symbol counts of custom dataset |
| 47 | - if use "byte", set to 256 (unicode byte range) |
| 48 | """ |
| 49 | with open(vocab_file_path, "r", encoding="utf-8") as f: |
| 50 | vocab_char_map = {} |
| 51 | for i, char in enumerate(f): |
| 52 | vocab_char_map[char[:-1]] = i |
| 53 | vocab_size = len(vocab_char_map) |
| 54 | return vocab_char_map, vocab_size |
| 55 | |
| 56 | |
| 57 | def convert_char_to_pinyin(reference_target_texts_list, polyphone=True): |
| 58 | final_reference_target_texts_list = [] |
| 59 | custom_trans = str.maketrans( |
| 60 | {";": ",", "“": '"', "”": '"', "‘": "'", "’": "'"} |
| 61 | ) # add custom trans here, to address oov |
| 62 | |
| 63 | def is_chinese(c): |
| 64 | return "\u3100" <= c <= "\u9fff" # common chinese characters |
| 65 | |
| 66 | for text in reference_target_texts_list: |
| 67 | char_list = [] |
| 68 | text = text.translate(custom_trans) |
| 69 | for seg in rjieba.cut(text): |
| 70 | seg_byte_len = len(bytes(seg, "UTF-8")) |
| 71 | if seg_byte_len == len(seg): # if pure alphabets and symbols |
| 72 | if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"": |
| 73 | char_list.append(" ") |
| 74 | char_list.extend(seg) |
| 75 | elif polyphone and seg_byte_len == 3 * len(seg): # if pure east asian characters |
| 76 | seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True) |
| 77 | for i, c in enumerate(seg): |
| 78 | if is_chinese(c): |
| 79 | char_list.append(" ") |
| 80 | char_list.append(seg_[i]) |
| 81 | else: # if mixed characters, alphabets and symbols |
| 82 | for c in seg: |
| 83 | if ord(c) < 256: |
| 84 | char_list.extend(c) |
| 85 | elif is_chinese(c): |
| 86 | char_list.append(" ") |
| 87 | char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True)) |
| 88 | else: |
| 89 | char_list.append(c) |
| 90 | final_reference_target_texts_list.append(char_list) |
| 91 | |
| 92 | return final_reference_target_texts_list |
| 93 | |
| 94 | |
| 95 | def list_str_to_idx( |
| 96 | text: list[str] | list[list[str]], |
| 97 | vocab_char_map: dict[str, int], # {char: idx} |
| 98 | padding_value=-1, |
| 99 | ): # noqa: F722 |
| 100 | list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text] # pinyin or char style |
| 101 | text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True) |
| 102 | return text |
| 103 | |
| 104 | |
| 105 | class TritonPythonModel: |
| 106 | def initialize(self, args): |
| 107 | self.use_perf = True |
| 108 | self.device = torch.device("cuda") |
| 109 | self.target_audio_sample_rate = 24000 |
| 110 | self.target_rms = 0.1 # least rms when inference, normalize to if lower |
| 111 | self.n_fft = 1024 |
| 112 | self.win_length = 1024 |
| 113 | self.hop_length = 256 |
| 114 | self.n_mel_channels = 100 |
| 115 | self.max_mel_len = 4096 |
| 116 | |
| 117 | parameters = json.loads(args["model_config"])["parameters"] |
| 118 | for key, value in parameters.items(): |
| 119 | parameters[key] = value["string_value"] |
| 120 | |
| 121 | self.vocab_char_map, self.vocab_size = get_tokenizer(parameters["vocab_file"]) |
| 122 | self.reference_sample_rate = int(parameters["reference_audio_sample_rate"]) |
| 123 | self.resampler = torchaudio.transforms.Resample(self.reference_sample_rate, self.target_audio_sample_rate) |
| 124 | |
| 125 | self.tllm_model_dir = parameters["tllm_model_dir"] |
| 126 | config_file = os.path.join(self.tllm_model_dir, "config.json") |
| 127 | with open(config_file) as f: |
| 128 | config = json.load(f) |
| 129 | self.model = F5TTS( |
| 130 | config, |
| 131 | debug_mode=False, |
| 132 | tllm_model_dir=self.tllm_model_dir, |
| 133 | model_path=parameters["model_path"], |
| 134 | vocab_size=self.vocab_size, |
| 135 | ) |
| 136 | |
| 137 | self.vocoder = parameters["vocoder"] |
| 138 | assert self.vocoder in ["vocos", "bigvgan"] |
| 139 | if self.vocoder == "vocos": |
| 140 | self.mel_stft = torchaudio.transforms.MelSpectrogram( |
| 141 | sample_rate=self.target_audio_sample_rate, |
| 142 | n_fft=self.n_fft, |
| 143 | win_length=self.win_length, |
| 144 | hop_length=self.hop_length, |
| 145 | n_mels=self.n_mel_channels, |
| 146 | power=1, |
| 147 | center=True, |
| 148 | normalized=False, |
| 149 | norm=None, |
| 150 | ).to(self.device) |
| 151 | self.compute_mel_fn = self.get_vocos_mel_spectrogram |
| 152 | elif self.vocoder == "bigvgan": |
| 153 | self.compute_mel_fn = self.get_bigvgan_mel_spectrogram |
| 154 | |
| 155 | def get_vocos_mel_spectrogram(self, waveform): |
| 156 | mel = self.mel_stft(waveform) |
| 157 | mel = mel.clamp(min=1e-5).log() |
| 158 | return mel.transpose(1, 2) |
| 159 | |
| 160 | def forward_vocoder(self, mel): |
| 161 | mel = mel.to(torch.float32).contiguous().cpu() |
| 162 | input_tensor_0 = pb_utils.Tensor.from_dlpack("mel", to_dlpack(mel)) |
| 163 | |
| 164 | inference_request = pb_utils.InferenceRequest( |
| 165 | model_name="vocoder", requested_output_names=["waveform"], inputs=[input_tensor_0] |
| 166 | ) |
| 167 | inference_response = inference_request.exec() |
| 168 | if inference_response.has_error(): |
| 169 | raise pb_utils.TritonModelException(inference_response.error().message()) |
| 170 | else: |
| 171 | waveform = pb_utils.get_output_tensor_by_name(inference_response, "waveform") |
| 172 | waveform = torch.utils.dlpack.from_dlpack(waveform.to_dlpack()).cpu() |
| 173 | |
| 174 | return waveform |
| 175 | |
| 176 | def execute(self, requests): |
| 177 | ( |
| 178 | reference_text_list, |
| 179 | target_text_list, |
| 180 | reference_target_texts_list, |
| 181 | estimated_reference_target_mel_len, |
| 182 | reference_mel_len, |
| 183 | reference_rms_list, |
| 184 | ) = [], [], [], [], [], [] |
| 185 | mel_features_list = [] |
| 186 | if self.use_perf: |
| 187 | torch.cuda.nvtx.range_push("preprocess") |
| 188 | for request in requests: |
| 189 | wav_tensor = pb_utils.get_input_tensor_by_name(request, "reference_wav") |
| 190 | wav_lens = pb_utils.get_input_tensor_by_name(request, "reference_wav_len") |
| 191 | |
| 192 | reference_text = pb_utils.get_input_tensor_by_name(request, "reference_text").as_numpy() |
| 193 | reference_text = reference_text[0][0].decode("utf-8") |
| 194 | reference_text_list.append(reference_text) |
| 195 | target_text = pb_utils.get_input_tensor_by_name(request, "target_text").as_numpy() |
| 196 | target_text = target_text[0][0].decode("utf-8") |
| 197 | target_text_list.append(target_text) |
| 198 | |
| 199 | text = reference_text + target_text |
| 200 | reference_target_texts_list.append(text) |
| 201 | |
| 202 | wav = from_dlpack(wav_tensor.to_dlpack()) |
| 203 | wav_len = from_dlpack(wav_lens.to_dlpack()) |
| 204 | wav_len = wav_len.squeeze() |
| 205 | assert wav.shape[0] == 1, "Only support batch size 1 for now." |
| 206 | wav = wav[:, :wav_len] |
| 207 | |
| 208 | ref_rms = torch.sqrt(torch.mean(torch.square(wav))) |
| 209 | if ref_rms < self.target_rms: |
| 210 | wav = wav * self.target_rms / ref_rms |
| 211 | reference_rms_list.append(ref_rms) |
| 212 | if self.reference_sample_rate != self.target_audio_sample_rate: |
| 213 | wav = self.resampler(wav) |
| 214 | wav = wav.to(self.device) |
| 215 | if self.use_perf: |
| 216 | torch.cuda.nvtx.range_push("compute_mel") |
| 217 | mel_features = self.compute_mel_fn(wav) |
| 218 | if self.use_perf: |
| 219 | torch.cuda.nvtx.range_pop() |
| 220 | mel_features_list.append(mel_features) |
| 221 | |
| 222 | reference_mel_len.append(mel_features.shape[1]) |
| 223 | estimated_reference_target_mel_len.append( |
| 224 | int( |
| 225 | mel_features.shape[1] * (1 + len(target_text.encode("utf-8")) / len(reference_text.encode("utf-8"))) |
| 226 | ) |
| 227 | ) |
| 228 | |
| 229 | max_seq_len = min(max(estimated_reference_target_mel_len), self.max_mel_len) |
| 230 | |
| 231 | batch = len(requests) |
| 232 | mel_features = torch.zeros((batch, max_seq_len, self.n_mel_channels), dtype=torch.float32).to(self.device) |
| 233 | for i, mel in enumerate(mel_features_list): |
| 234 | mel_features[i, : mel.shape[1], :] = mel |
| 235 | |
| 236 | reference_mel_len_tensor = torch.LongTensor(reference_mel_len).to(self.device) |
| 237 | |
| 238 | pinyin_list = convert_char_to_pinyin(reference_target_texts_list, polyphone=True) |
| 239 | text_pad_sequence = list_str_to_idx(pinyin_list, self.vocab_char_map) |
| 240 | |
| 241 | if self.use_perf: |
| 242 | torch.cuda.nvtx.range_pop() |
| 243 | |
| 244 | denoised, cost_time = self.model.sample( |
| 245 | text_pad_sequence, |
| 246 | mel_features, |
| 247 | reference_mel_len_tensor, |
| 248 | estimated_reference_target_mel_len, |
| 249 | remove_input_padding=False, |
| 250 | use_perf=self.use_perf, |
| 251 | ) |
| 252 | if self.use_perf: |
| 253 | torch.cuda.nvtx.range_push("vocoder") |
| 254 | |
| 255 | responses = [] |
| 256 | for i in range(batch): |
| 257 | ref_mel_len = reference_mel_len[i] |
| 258 | estimated_mel_len = estimated_reference_target_mel_len[i] |
| 259 | denoised_one_item = denoised[i, ref_mel_len:estimated_mel_len, :].unsqueeze(0).transpose(1, 2) |
| 260 | audio = self.forward_vocoder(denoised_one_item) |
| 261 | if reference_rms_list[i] < self.target_rms: |
| 262 | audio = audio * reference_rms_list[i] / self.target_rms |
| 263 | |
| 264 | audio = pb_utils.Tensor.from_dlpack("waveform", to_dlpack(audio)) |
| 265 | inference_response = pb_utils.InferenceResponse(output_tensors=[audio]) |
| 266 | responses.append(inference_response) |
| 267 | if self.use_perf: |
| 268 | torch.cuda.nvtx.range_pop() |
| 269 | return responses |
| 270 |