| 1 | # Copyright (c) 2024 Tsinghua Univ. (authors: Xingchen Song) |
| 2 | # 2025 (authors: Yuekai Zhang) |
| 3 | # |
| 4 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | # you may not use this file except in compliance with the License. |
| 6 | # You may obtain a copy of the License at |
| 7 | # |
| 8 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | # |
| 10 | # Unless required by applicable law or agreed to in writing, software |
| 11 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | # See the License for the specific language governing permissions and |
| 14 | # limitations under the License. |
| 15 | # Modified from https://github.com/xingchensong/S3Tokenizer/blob/main/s3tokenizer/cli.py |
| 16 | """ Example Usage |
| 17 | torchrun --nproc_per_node=1 \ |
| 18 | benchmark.py --output-dir $log_dir \ |
| 19 | --batch-size $batch_size \ |
| 20 | --enable-warmup \ |
| 21 | --split-name $split_name \ |
| 22 | --model-path $CKPT_DIR/$model/model_1200000.pt \ |
| 23 | --vocab-file $CKPT_DIR/$model/vocab.txt \ |
| 24 | --vocoder-trt-engine-path $vocoder_trt_engine_path \ |
| 25 | --backend-type $backend_type \ |
| 26 | --tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1 |
| 27 | """ |
| 28 | |
| 29 | import argparse |
| 30 | import importlib |
| 31 | import json |
| 32 | import os |
| 33 | import sys |
| 34 | import time |
| 35 | |
| 36 | import datasets |
| 37 | import tensorrt as trt |
| 38 | import torch |
| 39 | import torch.distributed as dist |
| 40 | import torch.nn.functional as F |
| 41 | import torchaudio |
| 42 | from datasets import load_dataset |
| 43 | from huggingface_hub import hf_hub_download |
| 44 | from tensorrt_llm._utils import trt_dtype_to_torch |
| 45 | from tensorrt_llm.logger import logger |
| 46 | from tensorrt_llm.runtime.session import Session, TensorInfo |
| 47 | from torch.utils.data import DataLoader, DistributedSampler |
| 48 | from tqdm import tqdm |
| 49 | from vocos import Vocos |
| 50 | |
| 51 | |
| 52 | sys.path.append(f"{os.path.dirname(os.path.abspath(__file__))}/../../../../src/") |
| 53 | |
| 54 | from f5_tts.eval.utils_eval import padded_mel_batch |
| 55 | from f5_tts.model.modules import get_vocos_mel_spectrogram |
| 56 | from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer, list_str_to_idx |
| 57 | |
| 58 | |
| 59 | F5TTS = importlib.import_module("model_repo_f5_tts.f5_tts.1.f5_tts_trtllm").F5TTS |
| 60 | |
| 61 | torch.manual_seed(0) |
| 62 | |
| 63 | |
| 64 | def get_args(): |
| 65 | parser = argparse.ArgumentParser(description="extract speech code") |
| 66 | parser.add_argument( |
| 67 | "--split-name", |
| 68 | type=str, |
| 69 | default="wenetspeech4tts", |
| 70 | choices=["wenetspeech4tts", "test_zh", "test_en", "test_hard"], |
| 71 | help="huggingface dataset split name", |
| 72 | ) |
| 73 | parser.add_argument("--output-dir", required=True, type=str, help="dir to save result") |
| 74 | parser.add_argument( |
| 75 | "--vocab-file", |
| 76 | required=True, |
| 77 | type=str, |
| 78 | help="vocab file", |
| 79 | ) |
| 80 | parser.add_argument( |
| 81 | "--model-path", |
| 82 | required=True, |
| 83 | type=str, |
| 84 | help="model path, to load text embedding", |
| 85 | ) |
| 86 | parser.add_argument( |
| 87 | "--tllm-model-dir", |
| 88 | required=True, |
| 89 | type=str, |
| 90 | help="tllm model dir", |
| 91 | ) |
| 92 | parser.add_argument( |
| 93 | "--batch-size", |
| 94 | required=True, |
| 95 | type=int, |
| 96 | help="batch size (per-device) for inference", |
| 97 | ) |
| 98 | parser.add_argument("--num-workers", type=int, default=0, help="workers for dataloader") |
| 99 | parser.add_argument("--prefetch", type=int, default=None, help="prefetch for dataloader") |
| 100 | parser.add_argument( |
| 101 | "--vocoder", |
| 102 | default="vocos", |
| 103 | type=str, |
| 104 | help="vocoder name", |
| 105 | ) |
| 106 | parser.add_argument( |
| 107 | "--vocoder-trt-engine-path", |
| 108 | default=None, |
| 109 | type=str, |
| 110 | help="vocoder trt engine path", |
| 111 | ) |
| 112 | parser.add_argument("--enable-warmup", action="store_true") |
| 113 | parser.add_argument("--remove-input-padding", action="store_true") |
| 114 | parser.add_argument("--use-perf", action="store_true", help="use nvtx to record performance") |
| 115 | parser.add_argument("--backend-type", type=str, default="triton", choices=["trt", "pytorch"], help="backend type") |
| 116 | args = parser.parse_args() |
| 117 | return args |
| 118 | |
| 119 | |
| 120 | def data_collator(batch, vocab_char_map, device="cuda", use_perf=False): |
| 121 | if use_perf: |
| 122 | torch.cuda.nvtx.range_push("data_collator") |
| 123 | target_sample_rate = 24000 |
| 124 | target_rms = 0.1 |
| 125 | ( |
| 126 | ids, |
| 127 | ref_rms_list, |
| 128 | ref_mel_list, |
| 129 | ref_mel_len_list, |
| 130 | estimated_reference_target_mel_len, |
| 131 | reference_target_texts_list, |
| 132 | ) = ( |
| 133 | [], |
| 134 | [], |
| 135 | [], |
| 136 | [], |
| 137 | [], |
| 138 | [], |
| 139 | ) |
| 140 | for i, item in enumerate(batch): |
| 141 | item_id, prompt_text, target_text = ( |
| 142 | item["id"], |
| 143 | item["prompt_text"], |
| 144 | item["target_text"], |
| 145 | ) |
| 146 | ids.append(item_id) |
| 147 | reference_target_texts_list.append(prompt_text + target_text) |
| 148 | |
| 149 | ref_audio_org, ref_sr = ( |
| 150 | item["prompt_audio"]["array"], |
| 151 | item["prompt_audio"]["sampling_rate"], |
| 152 | ) |
| 153 | ref_audio_org = torch.from_numpy(ref_audio_org).unsqueeze(0).float() |
| 154 | ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio_org))) |
| 155 | ref_rms_list.append(ref_rms) |
| 156 | if ref_rms < target_rms: |
| 157 | ref_audio_org = ref_audio_org * target_rms / ref_rms |
| 158 | |
| 159 | if ref_sr != target_sample_rate: |
| 160 | resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate) |
| 161 | ref_audio = resampler(ref_audio_org) |
| 162 | else: |
| 163 | ref_audio = ref_audio_org |
| 164 | |
| 165 | if use_perf: |
| 166 | torch.cuda.nvtx.range_push(f"mel_spectrogram {i}") |
| 167 | ref_audio = ref_audio.to("cuda") |
| 168 | ref_mel = get_vocos_mel_spectrogram(ref_audio).squeeze(0) |
| 169 | if use_perf: |
| 170 | torch.cuda.nvtx.range_pop() |
| 171 | ref_mel_len = ref_mel.shape[-1] |
| 172 | assert ref_mel.shape[0] == 100 |
| 173 | |
| 174 | ref_mel_list.append(ref_mel) |
| 175 | ref_mel_len_list.append(ref_mel_len) |
| 176 | |
| 177 | estimated_reference_target_mel_len.append( |
| 178 | int(ref_mel_len * (1 + len(target_text.encode("utf-8")) / len(prompt_text.encode("utf-8")))) |
| 179 | ) |
| 180 | |
| 181 | ref_mel_batch = padded_mel_batch(ref_mel_list) |
| 182 | ref_mel_len_batch = torch.LongTensor(ref_mel_len_list) |
| 183 | |
| 184 | pinyin_list = convert_char_to_pinyin(reference_target_texts_list, polyphone=True) |
| 185 | text_pad_sequence = list_str_to_idx(pinyin_list, vocab_char_map) |
| 186 | |
| 187 | if use_perf: |
| 188 | torch.cuda.nvtx.range_pop() |
| 189 | return { |
| 190 | "ids": ids, |
| 191 | "ref_rms_list": ref_rms_list, |
| 192 | "ref_mel_batch": ref_mel_batch, |
| 193 | "ref_mel_len_batch": ref_mel_len_batch, |
| 194 | "text_pad_sequence": text_pad_sequence, |
| 195 | "estimated_reference_target_mel_len": estimated_reference_target_mel_len, |
| 196 | } |
| 197 | |
| 198 | |
| 199 | def init_distributed(): |
| 200 | world_size = int(os.environ.get("WORLD_SIZE", 1)) |
| 201 | local_rank = int(os.environ.get("LOCAL_RANK", 0)) |
| 202 | rank = int(os.environ.get("RANK", 0)) |
| 203 | print( |
| 204 | "Inference on multiple gpus, this gpu {}".format(local_rank) |
| 205 | + ", rank {}, world_size {}".format(rank, world_size) |
| 206 | ) |
| 207 | torch.cuda.set_device(local_rank) |
| 208 | # Initialize process group with explicit device IDs |
| 209 | dist.init_process_group( |
| 210 | "nccl", |
| 211 | ) |
| 212 | return world_size, local_rank, rank |
| 213 | |
| 214 | |
| 215 | def load_vocoder( |
| 216 | vocoder_name="vocos", is_local=False, local_path="", device="cuda", hf_cache_dir=None, vocoder_trt_engine_path=None |
| 217 | ): |
| 218 | if vocoder_name == "vocos": |
| 219 | if vocoder_trt_engine_path is not None: |
| 220 | vocoder = VocosTensorRT(engine_path=vocoder_trt_engine_path) |
| 221 | else: |
| 222 | # vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device) |
| 223 | if is_local: |
| 224 | print(f"Load vocos from local path {local_path}") |
| 225 | config_path = f"{local_path}/config.yaml" |
| 226 | model_path = f"{local_path}/pytorch_model.bin" |
| 227 | else: |
| 228 | print("Download Vocos from huggingface charactr/vocos-mel-24khz") |
| 229 | repo_id = "charactr/vocos-mel-24khz" |
| 230 | config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="config.yaml") |
| 231 | model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="pytorch_model.bin") |
| 232 | vocoder = Vocos.from_hparams(config_path) |
| 233 | state_dict = torch.load(model_path, map_location="cpu", weights_only=True) |
| 234 | from vocos.feature_extractors import EncodecFeatures |
| 235 | |
| 236 | if isinstance(vocoder.feature_extractor, EncodecFeatures): |
| 237 | encodec_parameters = { |
| 238 | "feature_extractor.encodec." + key: value |
| 239 | for key, value in vocoder.feature_extractor.encodec.state_dict().items() |
| 240 | } |
| 241 | state_dict.update(encodec_parameters) |
| 242 | vocoder.load_state_dict(state_dict) |
| 243 | vocoder = vocoder.eval().to(device) |
| 244 | elif vocoder_name == "bigvgan": |
| 245 | raise NotImplementedError("BigVGAN is not implemented yet") |
| 246 | return vocoder |
| 247 | |
| 248 | |
| 249 | class VocosTensorRT: |
| 250 | def __init__(self, engine_path="./vocos_vocoder.plan", stream=None): |
| 251 | TRT_LOGGER = trt.Logger(trt.Logger.WARNING) |
| 252 | trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="") |
| 253 | logger.info(f"Loading vocoder engine from {engine_path}") |
| 254 | self.engine_path = engine_path |
| 255 | with open(engine_path, "rb") as f: |
| 256 | engine_buffer = f.read() |
| 257 | self.session = Session.from_serialized_engine(engine_buffer) |
| 258 | self.stream = stream if stream is not None else torch.cuda.current_stream().cuda_stream |
| 259 | |
| 260 | def decode(self, mels): |
| 261 | mels = mels.contiguous() |
| 262 | inputs = {"mel": mels} |
| 263 | output_info = self.session.infer_shapes([TensorInfo("mel", trt.DataType.FLOAT, mels.shape)]) |
| 264 | outputs = { |
| 265 | t.name: torch.empty(tuple(t.shape), dtype=trt_dtype_to_torch(t.dtype), device="cuda") for t in output_info |
| 266 | } |
| 267 | ok = self.session.run(inputs, outputs, self.stream) |
| 268 | |
| 269 | assert ok, "Runtime execution failed for vae session" |
| 270 | |
| 271 | samples = outputs["waveform"] |
| 272 | return samples |
| 273 | |
| 274 | |
| 275 | def main(): |
| 276 | args = get_args() |
| 277 | os.makedirs(args.output_dir, exist_ok=True) |
| 278 | |
| 279 | assert torch.cuda.is_available() |
| 280 | world_size, local_rank, rank = init_distributed() |
| 281 | device = torch.device(f"cuda:{local_rank}") |
| 282 | |
| 283 | vocab_char_map, vocab_size = get_tokenizer(args.vocab_file, "custom") |
| 284 | |
| 285 | tllm_model_dir = args.tllm_model_dir |
| 286 | with open(os.path.join(tllm_model_dir, "config.json")) as f: |
| 287 | tllm_model_config = json.load(f) |
| 288 | if args.backend_type == "trt": |
| 289 | model = F5TTS( |
| 290 | tllm_model_config, |
| 291 | debug_mode=False, |
| 292 | tllm_model_dir=tllm_model_dir, |
| 293 | model_path=args.model_path, |
| 294 | vocab_size=vocab_size, |
| 295 | ) |
| 296 | elif args.backend_type == "pytorch": |
| 297 | from f5_tts.infer.utils_infer import load_model |
| 298 | from f5_tts.model import DiT |
| 299 | |
| 300 | pretrained_config = tllm_model_config["pretrained_config"] |
| 301 | pt_model_config = dict( |
| 302 | dim=pretrained_config["hidden_size"], |
| 303 | depth=pretrained_config["num_hidden_layers"], |
| 304 | heads=pretrained_config["num_attention_heads"], |
| 305 | ff_mult=pretrained_config["ff_mult"], |
| 306 | text_dim=pretrained_config["text_dim"], |
| 307 | text_mask_padding=pretrained_config["text_mask_padding"], |
| 308 | conv_layers=pretrained_config["conv_layers"], |
| 309 | pe_attn_head=pretrained_config["pe_attn_head"], |
| 310 | # attn_backend="flash_attn", |
| 311 | # attn_mask_enabled=True, |
| 312 | ) |
| 313 | model = load_model(DiT, pt_model_config, args.model_path) |
| 314 | |
| 315 | vocoder = load_vocoder( |
| 316 | vocoder_name=args.vocoder, device=device, vocoder_trt_engine_path=args.vocoder_trt_engine_path |
| 317 | ) |
| 318 | |
| 319 | dataset = load_dataset( |
| 320 | "yuekai/seed_tts", |
| 321 | split=args.split_name, |
| 322 | trust_remote_code=True, |
| 323 | ) |
| 324 | |
| 325 | def add_estimated_duration(example): |
| 326 | prompt_audio_len = example["prompt_audio"]["array"].shape[0] |
| 327 | scale_factor = 1 + len(example["target_text"]) / len(example["prompt_text"]) |
| 328 | estimated_duration = prompt_audio_len * scale_factor |
| 329 | example["estimated_duration"] = estimated_duration / example["prompt_audio"]["sampling_rate"] |
| 330 | return example |
| 331 | |
| 332 | dataset = dataset.map(add_estimated_duration) |
| 333 | dataset = dataset.sort("estimated_duration", reverse=True) |
| 334 | if args.use_perf: |
| 335 | # dataset_list = [dataset.select(range(1)) for i in range(16)] # seq_len 1000 |
| 336 | dataset_list_short = [dataset.select([24]) for i in range(8)] # seq_len 719 |
| 337 | # dataset_list_long = [dataset.select([23]) for i in range(8)] # seq_len 2002 |
| 338 | # dataset = datasets.concatenate_datasets(dataset_list_short + dataset_list_long) |
| 339 | dataset = datasets.concatenate_datasets(dataset_list_short) |
| 340 | if world_size > 1: |
| 341 | sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) |
| 342 | else: |
| 343 | # This would disable shuffling |
| 344 | sampler = None |
| 345 | |
| 346 | dataloader = DataLoader( |
| 347 | dataset, |
| 348 | batch_size=args.batch_size, |
| 349 | sampler=sampler, |
| 350 | shuffle=False, |
| 351 | num_workers=args.num_workers, |
| 352 | prefetch_factor=args.prefetch, |
| 353 | collate_fn=lambda x: data_collator(x, vocab_char_map, use_perf=args.use_perf), |
| 354 | ) |
| 355 | |
| 356 | total_steps = len(dataset) |
| 357 | |
| 358 | if args.enable_warmup: |
| 359 | for batch in dataloader: |
| 360 | ref_mels, ref_mel_lens = batch["ref_mel_batch"].to(device), batch["ref_mel_len_batch"].to(device) |
| 361 | text_pad_seq = batch["text_pad_sequence"].to(device) |
| 362 | total_mel_lens = batch["estimated_reference_target_mel_len"] |
| 363 | cond_pad_seq = F.pad(ref_mels, (0, 0, 0, max(total_mel_lens) - ref_mels.shape[1], 0, 0)) |
| 364 | if args.backend_type == "trt": |
| 365 | _ = model.sample( |
| 366 | text_pad_seq, |
| 367 | cond_pad_seq, |
| 368 | ref_mel_lens, |
| 369 | total_mel_lens, |
| 370 | remove_input_padding=args.remove_input_padding, |
| 371 | ) |
| 372 | elif args.backend_type == "pytorch": |
| 373 | total_mel_lens = torch.tensor(total_mel_lens, device=device) |
| 374 | with torch.inference_mode(): |
| 375 | generated, _ = model.sample( |
| 376 | cond=ref_mels, |
| 377 | text=text_pad_seq, |
| 378 | duration=total_mel_lens, |
| 379 | steps=32, |
| 380 | cfg_strength=2.0, |
| 381 | sway_sampling_coef=-1, |
| 382 | ) |
| 383 | |
| 384 | if rank == 0: |
| 385 | progress_bar = tqdm(total=total_steps, desc="Processing", unit="wavs") |
| 386 | |
| 387 | decoding_time = 0 |
| 388 | vocoder_time = 0 |
| 389 | total_duration = 0 |
| 390 | if args.use_perf: |
| 391 | torch.cuda.cudart().cudaProfilerStart() |
| 392 | total_decoding_time = time.time() |
| 393 | for batch in dataloader: |
| 394 | if args.use_perf: |
| 395 | torch.cuda.nvtx.range_push("data sample") |
| 396 | ref_mels, ref_mel_lens = batch["ref_mel_batch"].to(device), batch["ref_mel_len_batch"].to(device) |
| 397 | text_pad_seq = batch["text_pad_sequence"].to(device) |
| 398 | total_mel_lens = batch["estimated_reference_target_mel_len"] |
| 399 | cond_pad_seq = F.pad(ref_mels, (0, 0, 0, max(total_mel_lens) - ref_mels.shape[1], 0, 0)) |
| 400 | if args.use_perf: |
| 401 | torch.cuda.nvtx.range_pop() |
| 402 | if args.backend_type == "trt": |
| 403 | generated, cost_time = model.sample( |
| 404 | text_pad_seq, |
| 405 | cond_pad_seq, |
| 406 | ref_mel_lens, |
| 407 | total_mel_lens, |
| 408 | remove_input_padding=args.remove_input_padding, |
| 409 | use_perf=args.use_perf, |
| 410 | ) |
| 411 | elif args.backend_type == "pytorch": |
| 412 | total_mel_lens = torch.tensor(total_mel_lens, device=device) |
| 413 | with torch.inference_mode(): |
| 414 | start_time = time.time() |
| 415 | generated, _ = model.sample( |
| 416 | cond=ref_mels, |
| 417 | text=text_pad_seq, |
| 418 | duration=total_mel_lens, |
| 419 | lens=ref_mel_lens, |
| 420 | steps=32, |
| 421 | cfg_strength=2.0, |
| 422 | sway_sampling_coef=-1, |
| 423 | ) |
| 424 | cost_time = time.time() - start_time |
| 425 | decoding_time += cost_time |
| 426 | vocoder_start_time = time.time() |
| 427 | target_rms = 0.1 |
| 428 | target_sample_rate = 24000 |
| 429 | for i, gen in enumerate(generated): |
| 430 | gen = gen[ref_mel_lens[i] : total_mel_lens[i], :].unsqueeze(0) |
| 431 | gen_mel_spec = gen.permute(0, 2, 1).to(torch.float32) |
| 432 | if args.vocoder == "vocos": |
| 433 | if args.use_perf: |
| 434 | torch.cuda.nvtx.range_push("vocoder decode") |
| 435 | generated_wave = vocoder.decode(gen_mel_spec).cpu() |
| 436 | if args.use_perf: |
| 437 | torch.cuda.nvtx.range_pop() |
| 438 | else: |
| 439 | generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu() |
| 440 | |
| 441 | if batch["ref_rms_list"][i] < target_rms: |
| 442 | generated_wave = generated_wave * batch["ref_rms_list"][i] / target_rms |
| 443 | |
| 444 | utt = batch["ids"][i] |
| 445 | torchaudio.save( |
| 446 | f"{args.output_dir}/{utt}.wav", |
| 447 | generated_wave, |
| 448 | target_sample_rate, |
| 449 | ) |
| 450 | total_duration += generated_wave.shape[1] / target_sample_rate |
| 451 | vocoder_time += time.time() - vocoder_start_time |
| 452 | if rank == 0: |
| 453 | progress_bar.update(world_size * len(batch["ids"])) |
| 454 | total_decoding_time = time.time() - total_decoding_time |
| 455 | if rank == 0: |
| 456 | progress_bar.close() |
| 457 | rtf = total_decoding_time / total_duration |
| 458 | s = f"RTF: {rtf:.4f}\n" |
| 459 | s += f"total_duration: {total_duration:.3f} seconds\n" |
| 460 | s += f"({total_duration / 3600:.2f} hours)\n" |
| 461 | s += f"DiT time: {decoding_time:.3f} seconds ({decoding_time / 3600:.2f} hours)\n" |
| 462 | s += f"Vocoder time: {vocoder_time:.3f} seconds ({vocoder_time / 3600:.2f} hours)\n" |
| 463 | s += f"total decoding time: {total_decoding_time:.3f} seconds ({total_decoding_time / 3600:.2f} hours)\n" |
| 464 | s += f"batch size: {args.batch_size}\n" |
| 465 | print(s) |
| 466 | |
| 467 | with open(f"{args.output_dir}/rtf.txt", "w") as f: |
| 468 | f.write(s) |
| 469 | |
| 470 | dist.barrier() |
| 471 | dist.destroy_process_group() |
| 472 | |
| 473 | |
| 474 | if __name__ == "__main__": |
| 475 | main() |
| 476 |