Update app_local.py

app_local.py

@@ -14,17 +14,17 @@ from pydub import AudioSegment
 from model import CFM, UNetT, DiT, MMDiT
 from cached_path import cached_path
 from model.utils import (
-    get_tokenizer,
-    convert_char_to_pinyin,
+    get_tokenizer, 
+    convert_char_to_pinyin, 
     save_spectrogram,
 )
 from transformers import pipeline
+import spaces
 import librosa
+from txtsplit import txtsplit
 
 device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 
-print(f"Using {device} device")
-
 pipe = pipeline(
     "automatic-speech-recognition",
     model="openai/whisper-large-v3-turbo",
@@ -79,13 +79,13 @@ 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)
 
-def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
+def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress = gr.Progress()):
     print(gen_text)
-    if len(gen_text) > 200:
-        raise gr.Error("Please keep your text under 200 chars.")
     gr.Info("Converting audio...")
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
         aseg = AudioSegment.from_file(ref_audio_orig)
+        # Convert to mono
+        aseg = aseg.set_channels(1)
         audio_duration = len(aseg)
         if audio_duration > 15000:
             gr.Warning("Audio is over 15s, clipping to only first 15s.")
@@ -98,7 +98,7 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
     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...")
         ref_text = outputs = pipe(
@@ -112,7 +112,9 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
     else:
         gr.Info("Using custom reference text...")
     audio, sr = torchaudio.load(ref_audio)
-
+    # Audio
+    if audio.shape[0] > 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
     rms = torch.sqrt(torch.mean(torch.square(audio)))
     if rms < target_rms:
         audio = audio * target_rms / rms
@@ -120,44 +122,49 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
         resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
         audio = resampler(audio)
     audio = audio.to(device)
-
-    # Prepare the text
-    text_list = [ref_text + gen_text]
-    final_text_list = convert_char_to_pinyin(text_list)
-
-    # Calculate duration
-    ref_audio_len = audio.shape[-1] // hop_length
-    # if fix_duration is not None:
-    #     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)
-
-    # inference
-    gr.Info(f"Generating audio using {exp_name}")
-    with torch.inference_mode():
-        generated, _ = base_model.sample(
-            cond=audio,
-            text=final_text_list,
-            duration=duration,
-            steps=nfe_step,
-            cfg_strength=cfg_strength,
-            sway_sampling_coef=sway_sampling_coef,
-        )
-
-    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
-
-    # wav -> numpy
-    generated_wave = generated_wave.squeeze().cpu().numpy()
-
+    # Chunk
+    chunks = txtsplit(gen_text, 100, 150) # 100 chars preferred, 150 max
+    results = []
+    generated_mel_specs = []
+    for chunk in progress.tqdm(chunks):
+        # Prepare the text
+        text_list = [ref_text + chunk]
+        final_text_list = convert_char_to_pinyin(text_list)
+    
+        # Calculate duration
+        ref_audio_len = audio.shape[-1] // hop_length
+        # if fix_duration is not None:
+        #     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)
+    
+        # inference
+        gr.Info(f"Generating audio using {exp_name}")
+        with torch.inference_mode():
+            generated, _ = base_model.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+    
+        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
+    
+        # wav -> numpy
+        generated_wave = generated_wave.squeeze().cpu().numpy()
+        results.append(generated_wave)
+    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)
@@ -169,11 +176,11 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
 
 
     # spectogram
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
-        spectrogram_path = tmp_spectrogram.name
-        save_spectrogram(generated_mel_spec[0].cpu().numpy(), spectrogram_path)
+    # with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+    #     spectrogram_path = tmp_spectrogram.name
+    #     save_spectrogram(generated_mel_spec[0].cpu().numpy(), spectrogram_path)
 
-    return (target_sample_rate, generated_wave), spectrogram_path
+    return (target_sample_rate, generated_wave)
 
 with gr.Blocks() as app:
     gr.Markdown("""
@@ -190,21 +197,38 @@ The checkpoints support English and Chinese.
 
 If you're having issues, try converting your reference audio to WAV or MP3, clipping it to 15s, and shortening your prompt. If you're still running into issues, please open a [community Discussion](https://huggingface.co/spaces/mrfakename/E2-F5-TTS/discussions).
 
+Long-form/batched inference + speech editing is coming soon!
+
 **NOTE: Reference text will be automatically transcribed with Whisper if not provided. For best results, keep your reference clips short (<15s). Ensure the audio is fully uploaded before generating.**
 """)
 
     ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
-    gen_text_input = gr.Textbox(label="Text to Generate (max 200 chars.)", lines=4)
+    gen_text_input = gr.Textbox(label="Text to Generate (longer text will use chunking)", lines=4)
     model_choice = gr.Radio(choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS")
     generate_btn = gr.Button("Synthesize", variant="primary")
     with gr.Accordion("Advanced Settings", open=False):
         ref_text_input = gr.Textbox(label="Reference Text", info="Leave blank to automatically transcribe the reference audio. If you enter text it will override automatic transcription.", lines=2)
         remove_silence = gr.Checkbox(label="Remove Silences", info="The model tends to produce silences, especially on longer audio. We can manually remove silences if needed. Note that this is an experimental feature and may produce strange results. This will also increase generation time.", value=True)
-
+    
     audio_output = gr.Audio(label="Synthesized Audio")
-    spectrogram_output = gr.Image(label="Spectrogram")
+    # spectrogram_output = gr.Image(label="Spectrogram")
+
+    generate_btn.click(infer, inputs=[ref_audio_input, ref_text_input, gen_text_input, model_choice, remove_silence], outputs=[audio_output])
+    gr.Markdown("""
+## Run Locally
+
+Run this demo locally on CPU, CUDA, or MPS/Apple Silicon (requires macOS >= 14):
 
-    generate_btn.click(infer, inputs=[ref_audio_input, ref_text_input, gen_text_input, model_choice, remove_silence], outputs=[audio_output, spectrogram_output])
+First, ensure `ffmpeg` is installed.
+
+```bash
+git clone https://huggingface.co/spaces/mrfakename/E2-F5-TTS
+cd E2-F5-TTS
+python -m pip install -r requirements.txt
+python app_local.py
+```
+
+""")
     gr.Markdown("Unofficial demo by [mrfakename](https://x.com/realmrfakename)")