Update app_local.py (#17)

- Update app_local.py (e6226de0b4e526b510862269bb30165febee315f)


Co-authored-by: Yushen CHEN <SWivid@users.noreply.huggingface.co>

app_local.py

@@ -10,7 +10,7 @@ import tempfile
 from einops import rearrange
 from ema_pytorch import EMA
 from vocos import Vocos
-from pydub import AudioSegment
+from pydub import AudioSegment, silence
 from model import CFM, UNetT, DiT, MMDiT
 from cached_path import cached_path
 from model.utils import (
@@ -20,6 +20,7 @@ from model.utils import (
 )
 from transformers import pipeline
 import librosa
+import soundfile as sf
 from txtsplit import txtsplit
 
 device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
@@ -31,6 +32,8 @@ pipe = pipeline(
     device=device,
 )
 
+vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
 # --------------------- Settings -------------------- #
 
 target_sample_rate = 24000
@@ -45,8 +48,8 @@ speed = 1.0
 # fix_duration = 27  # None or float (duration in seconds)
 fix_duration = None
 
-def load_model(exp_name, model_cls, model_cfg, ckpt_step):
-    checkpoint = torch.load(str(cached_path(f"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.pt")), map_location=device)
+def load_model(repo_name, exp_name, model_cls, model_cfg, ckpt_step):
+    checkpoint = torch.load(str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.pt")), map_location=device)
     vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")
     model = CFM(
         transformer=model_cls(
@@ -69,20 +72,26 @@ def load_model(exp_name, model_cls, model_cfg, ckpt_step):
     ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
     ema_model.copy_params_from_ema_to_model()
 
-    return ema_model, model
+    return model
 
 # load models
 F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
 E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
 
-F5TTS_ema_model, F5TTS_base_model = load_model("F5TTS_Base", DiT, F5TTS_model_cfg, 1200000)
-E2TTS_ema_model, E2TTS_base_model = load_model("E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000)
+F5TTS_ema_model = load_model("F5-TTS", "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000)
+E2TTS_ema_model = load_model("E2-TTS", "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000)
 
 def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress = gr.Progress()):
     print(gen_text)
     gr.Info("Converting audio...")
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
         aseg = AudioSegment.from_file(ref_audio_orig)
+        # remove long silence in reference audio
+        non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            non_silent_wave += non_silent_seg
+        aseg = non_silent_wave
         # Convert to mono
         aseg = aseg.set_channels(1)
         audio_duration = len(aseg)
@@ -93,10 +102,8 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
         ref_audio = f.name
     if exp_name == "F5-TTS":
         ema_model = F5TTS_ema_model
-        base_model = F5TTS_base_model
     elif exp_name == "E2-TTS":
         ema_model = E2TTS_ema_model
-        base_model = E2TTS_base_model
     
     if not ref_text.strip():
         gr.Info("No reference text provided, transcribing reference audio...")
@@ -111,6 +118,7 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
     else:
         gr.Info("Using custom reference text...")
     audio, sr = torchaudio.load(ref_audio)
+    max_chars = int(len(ref_text) / (audio.shape[-1] / sr) * (30 - audio.shape[-1] / sr))
     # Audio
     if audio.shape[0] > 1:
         audio = torch.mean(audio, dim=0, keepdim=True)
@@ -122,7 +130,7 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
         audio = resampler(audio)
     audio = audio.to(device)
     # Chunk
-    chunks = txtsplit(gen_text, 100, 150) # 100 chars preferred, 150 max
+    chunks = txtsplit(gen_text, 0.7*max_chars, 0.9*max_chars) # 100 chars preferred, 150 max
     results = []
     generated_mel_specs = []
     for chunk in progress.tqdm(chunks):
@@ -136,14 +144,14 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
         #     duration = int(fix_duration * target_sample_rate / hop_length)
         # else:
         zh_pause_punc = r"。，、；：？！"
-        ref_text_len = len(ref_text) + len(re.findall(zh_pause_punc, ref_text))
-        gen_text_len = len(gen_text) + len(re.findall(zh_pause_punc, gen_text))
-        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+        ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
+        chunk = len(chunk.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
+        duration = ref_audio_len + int(ref_audio_len / ref_text_len * chunk / speed)
     
         # inference
         gr.Info(f"Generating audio using {exp_name}")
         with torch.inference_mode():
-            generated, _ = base_model.sample(
+            generated, _ = ema_model.sample(
                 cond=audio,
                 text=final_text_list,
                 duration=duration,
@@ -155,7 +163,6 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
         generated = generated[:, ref_audio_len:, :]
         generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
         gr.Info("Running vocoder")
-        vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
         generated_wave = vocos.decode(generated_mel_spec.cpu())
         if rms < target_rms:
             generated_wave = generated_wave * rms / target_rms
@@ -166,13 +173,23 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
     generated_wave = np.concatenate(results)
     if remove_silence:
         gr.Info("Removing audio silences... This may take a moment")
-        non_silent_intervals = librosa.effects.split(generated_wave, top_db=30)
-        non_silent_wave = np.array([])
-        for interval in non_silent_intervals:
-            start, end = interval
-            non_silent_wave = np.concatenate([non_silent_wave, generated_wave[start:end]])
-        generated_wave = non_silent_wave
-
+        # non_silent_intervals = librosa.effects.split(generated_wave, top_db=30)
+        # non_silent_wave = np.array([])
+        # for interval in non_silent_intervals:
+        #     start, end = interval
+        #     non_silent_wave = np.concatenate([non_silent_wave, generated_wave[start:end]])
+        # generated_wave = non_silent_wave
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+            sf.write(f.name, generated_wave, target_sample_rate)
+            aseg = AudioSegment.from_file(f.name)
+            non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+            non_silent_wave = AudioSegment.silent(duration=0)
+            for non_silent_seg in non_silent_segs:
+                non_silent_wave += non_silent_seg
+            aseg = non_silent_wave
+            aseg.export(f.name, format="wav")
+            generated_wave, _ = torchaudio.load(f.name)
+        generated_wave = generated_wave.squeeze().cpu().numpy()
 
     # spectogram
     # with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
@@ -214,6 +231,6 @@ Long-form/batched inference + speech editing is coming soon!
 
     generate_btn.click(infer, inputs=[ref_audio_input, ref_text_input, gen_text_input, model_choice, remove_silence], outputs=[audio_output])
     gr.Markdown("Unofficial demo by [mrfakename](https://x.com/realmrfakename)")
-
+    
 
 app.queue().launch()