Skip to content

Commit 60dd859

Browse files
shibing624shibing624
committed
indextts_fastapi_server1
1 parent 411f594 commit 60dd859

2 files changed

Lines changed: 48 additions & 15 deletions

File tree

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ wheels/
3030
.installed.cfg
3131
*.egg
3232
MANIFEST
33+
.gradio/
3334

3435
# PyInstaller
3536
# Usually these files are written by a python script from a template
@@ -52,6 +53,8 @@ coverage.xml
5253
*.cover
5354
.hypothesis/
5455
.pytest_cache/
56+
.python-version
57+
.ruff_cache/
5558

5659
# Translations
5760
*.mo

‎parrots/indextts/inference.py‎

Lines changed: 45 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,14 @@ def __init__(
126126
self.gpt.eval()
127127
logger.debug(f"model loaded successfully from {self.model_dir}")
128128
self.gpt.post_init_gpt2_config(use_deepspeed=False, kv_cache=True, half=self.use_fp16)
129+
130+
# Enable torch.compile for GPT model if using CUDA
131+
if self.use_cuda:
132+
try:
133+
self.gpt = torch.compile(self.gpt, mode="reduce-overhead")
134+
logger.debug("GPT model compiled with torch.compile")
135+
except Exception as e:
136+
logger.warning(f"Failed to compile GPT model: {e}")
129137

130138
self.extract_features = SeamlessM4TFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")
131139
self.semantic_model, self.semantic_mean, self.semantic_std = build_semantic_model(
@@ -487,6 +495,7 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
487495
self.cache_s2mel_prompt = None
488496
self.cache_mel = None
489497
torch.cuda.empty_cache()
498+
490499
audio, sr = self._load_and_cut_audio(speak_reference_audio_path, 15, verbose)
491500
audio_22k = torchaudio.transforms.Resample(sr, 22050)(audio)
492501
audio_16k = torchaudio.transforms.Resample(sr, 16000)(audio)
@@ -496,16 +505,20 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
496505
attention_mask = inputs["attention_mask"]
497506
input_features = input_features.to(self.device)
498507
attention_mask = attention_mask.to(self.device)
508+
499509
spk_cond_emb = self.get_emb(input_features, attention_mask)
500510

501511
_, S_ref = self.semantic_codec.quantize(spk_cond_emb)
512+
502513
ref_mel = self.mel_fn(audio_22k.to(spk_cond_emb.device).float())
503514
ref_target_lengths = torch.LongTensor([ref_mel.size(2)]).to(ref_mel.device)
515+
504516
feat = torchaudio.compliance.kaldi.fbank(audio_16k.to(ref_mel.device),
505517
num_mel_bins=80,
506518
dither=0,
507519
sample_frequency=16000)
508520
feat = feat - feat.mean(dim=0, keepdim=True) # feat2另外一个滤波器能量组特征[922, 80]
521+
509522
style = self.campplus_model(feat.unsqueeze(0)) # 参考音频的全局style2[1,192]
510523

511524
prompt_condition = self.s2mel.models['length_regulator'](S_ref,
@@ -541,12 +554,15 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
541554
if self.cache_emo_cond is not None:
542555
self.cache_emo_cond = None
543556
torch.cuda.empty_cache()
557+
544558
emo_audio, _ = self._load_and_cut_audio(emo_reference_audio_path, 15, verbose, sr=16000)
559+
545560
emo_inputs = self.extract_features(emo_audio, sampling_rate=16000, return_tensors="pt")
546561
emo_input_features = emo_inputs["input_features"]
547562
emo_attention_mask = emo_inputs["attention_mask"]
548563
emo_input_features = emo_input_features.to(self.device)
549564
emo_attention_mask = emo_attention_mask.to(self.device)
565+
550566
emo_cond_emb = self.get_emb(emo_input_features, emo_attention_mask)
551567

552568
self.cache_emo_cond = emo_cond_emb
@@ -603,7 +619,6 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
603619
text_token_syms = self.tokenizer.convert_ids_to_tokens(text_tokens[0].tolist())
604620
logger.info(f"text_token_syms is same as segment tokens, {text_token_syms == sent}")
605621

606-
m_start_time = time.perf_counter()
607622
with torch.no_grad():
608623
with torch.amp.autocast(text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype):
609624
emovec = self.gpt.merge_emovec(
@@ -618,6 +633,7 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
618633
emovec = emovec_mat + (1 - torch.sum(weight_vector)) * emovec
619634
# emovec = emovec_mat
620635

636+
gpt_infer_start = time.perf_counter()
621637
codes, speech_conditioning_latent = self.gpt.inference_speech(
622638
spk_cond_emb,
623639
text_tokens,
@@ -636,8 +652,9 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
636652
max_generate_length=max_mel_tokens,
637653
**generation_kwargs
638654
)
655+
gpt_infer_time = time.perf_counter() - gpt_infer_start
639656

640-
gpt_gen_time += time.perf_counter() - m_start_time
657+
gpt_gen_time += gpt_infer_time
641658
if not has_warned and (codes[:, -1] != self.stop_mel_token).any():
642659
warnings.warn(
643660
f"WARN: generation stopped due to exceeding `max_mel_tokens` ({max_mel_tokens}). "
@@ -663,7 +680,8 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
663680
if verbose:
664681
# logger.info(f"codes: {codes}, type: {type(codes)}")
665682
logger.info(f"codes shape: {codes.shape}, codes type: {codes.dtype}, code len: {code_lens}")
666-
m_start_time = time.perf_counter()
683+
684+
gpt_forward_start = time.perf_counter()
667685
use_speed = torch.zeros(spk_cond_emb.size(0)).to(spk_cond_emb.device).long()
668686
with torch.amp.autocast(text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype):
669687
latent = self.gpt(
@@ -678,14 +696,16 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
678696
emo_vec=emovec,
679697
use_speed=use_speed,
680698
)
681-
gpt_forward_time += time.perf_counter() - m_start_time
699+
gpt_fwd_time = time.perf_counter() - gpt_forward_start
700+
gpt_forward_time += gpt_fwd_time
682701

683702
dtype = None
684703
with torch.amp.autocast(text_tokens.device.type, enabled=dtype is not None, dtype=dtype):
685-
m_start_time = time.perf_counter()
686704
diffusion_steps = 25
687705
inference_cfg_rate = 0.7
706+
688707
latent = self.s2mel.models['gpt_layer'](latent)
708+
689709
S_infer = self.semantic_codec.quantizer.vq2emb(codes.unsqueeze(1))
690710
S_infer = S_infer.transpose(1, 2)
691711
S_infer = S_infer + latent
@@ -695,23 +715,29 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
695715
ylens=target_lengths,
696716
n_quantizers=3,
697717
f0=None)[0]
718+
698719
cat_condition = torch.cat([prompt_condition, cond], dim=1)
720+
721+
cfm_start = time.perf_counter()
699722
vc_target = self.s2mel.models['cfm'].inference(cat_condition,
700723
torch.LongTensor([cat_condition.size(1)]).to(
701724
cond.device),
702725
ref_mel, style, None, diffusion_steps,
703726
inference_cfg_rate=inference_cfg_rate)
704727
vc_target = vc_target[:, :, ref_mel.size(-1):]
705-
s2mel_time += time.perf_counter() - m_start_time
728+
cfm_time = time.perf_counter() - cfm_start
729+
s2mel_time += cfm_time
706730

707-
m_start_time = time.perf_counter()
731+
vocoder_start = time.perf_counter()
708732
wav = self.bigvgan(vc_target.float()).squeeze().unsqueeze(0)
709-
bigvgan_time += time.perf_counter() - m_start_time
733+
vocoder_time = time.perf_counter() - vocoder_start
734+
bigvgan_time += vocoder_time
710735
wav = wav.squeeze(1)
711736

712737
wav = torch.clamp(32767 * wav, -32767.0, 32767.0)
713738
if verbose:
714739
logger.info(f"wav shape: {wav.shape}, min: {wav.min()}, max: {wav.max()}")
740+
715741
wavs.append(wav.cpu()) # to cpu before saving
716742
if stream_return:
717743
yield wav.cpu()
@@ -725,13 +751,17 @@ def infer_generator(self, text, speak_reference_audio_path_or_name="male_broadca
725751
wavs = self.insert_interval_silence(wavs, sampling_rate=sampling_rate, interval_silence=interval_silence)
726752
wav = torch.cat(wavs, dim=1)
727753
wav_length = wav.shape[-1] / sampling_rate
728-
logger.debug(f"gpt_gen_time: {gpt_gen_time:.2f} seconds")
729-
logger.debug(f"gpt_forward_time: {gpt_forward_time:.2f} seconds")
730-
logger.debug(f"s2mel_time: {s2mel_time:.2f} seconds")
731-
logger.debug(f"bigvgan_time: {bigvgan_time:.2f} seconds")
732-
logger.debug(f"Total inference time: {end_time - start_time:.2f} seconds")
733-
logger.debug(f"Generated audio length: {wav_length:.2f} seconds")
734-
logger.debug(f"RTF: {(end_time - start_time) / wav_length:.4f}")
754+
755+
if verbose:
756+
logger.info("[PERF] ========== Performance Summary ==========")
757+
logger.info(f"[PERF] GPT inference_speech total: {gpt_gen_time:.3f}s")
758+
logger.info(f"[PERF] GPT forward (latent) total: {gpt_forward_time:.3f}s")
759+
logger.info(f"[PERF] S2Mel (CFM diffusion) total: {s2mel_time:.3f}s")
760+
logger.info(f"[PERF] BigVGAN vocoder total: {bigvgan_time:.3f}s")
761+
logger.info(f"[PERF] Total inference time: {end_time - start_time:.3f}s")
762+
logger.info(f"[PERF] Generated audio length: {wav_length:.2f}s")
763+
logger.info(f"[PERF] RTF (Real-Time Factor): {(end_time - start_time) / wav_length:.4f}")
764+
logger.info("[PERF] =============================================")
735765

736766
# save audio
737767
wav = wav.cpu() # to cpu

0 commit comments

Comments
 (0)