Batch Inference & Podcast Generation

Here's what the Batch Inference part does:

- Try to put as much characters as possible into one batch (200 max)
- If it's not possible, it'll try to do a cut whenever there's a semicolon character
- If it's not possible, it'll try to do a cut whenever there's a comma character
- If it's not possible, it'll try to do a cut after the most logical word (thus, therefore etc.) --> There's a list at the top of the Gradio script, and it's possible to modify it in Advanced Settings
- If nothing above worked, it's just going to go past that 200 line (realistically, if your text isn't gibberish, this shouldn't happen :D)

The Podcast Generation feature has these features built in:
- Takes two reference speeches and two reference texts (or empty and then transcribed automatically)
- You have to give a name to each of the two speakers
- You can then paste the podcast script, with one speaker's name followed by a semicolon and then their text, you can do the same with the other speaker, all as long as you want (because it's using the same batch inference as before)

All in all, the batch inference feature allow for a little bit more than real-time inference. (I might do another pull request with real-time streaming)

Immense thanks to all of those who worked on this project, it's really great. There's of course still room for improvement, but I think this is a step forward in terms of OSS TTS, so thanks !

gradio_app.py

@@ -19,6 +19,18 @@ from model.utils import (
 from transformers import pipeline
 import librosa
 import click
+import soundfile as sf
+
+SPLIT_WORDS = [
+    "but", "however", "nevertheless", "yet", "still",
+    "therefore", "thus", "hence", "consequently",
+    "moreover", "furthermore", "additionally",
+    "meanwhile", "alternatively", "otherwise",
+    "namely", "specifically", "for example", "such as",
+    "in fact", "indeed", "notably",
+    "in contrast", "on the other hand", "conversely",
+    "in conclusion", "to summarize", "finally"
+]
 
 device = (
     "cuda"
@@ -87,11 +99,179 @@ E2TTS_ema_model = load_model(
     "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000
 )
 
+def split_text_into_batches(text, max_chars=200, split_words=SPLIT_WORDS):
+    sentences = re.split('([。.!?！？])', text)
+    sentences = [''.join(i) for i in zip(sentences[0::2], sentences[1::2])]
+    
+    batches = []
+    current_batch = ""
+    
+    def split_by_words(text):
+        words = text.split()
+        current_word_part = ""
+        word_batches = []
+        for word in words:
+            if len(current_word_part) + len(word) + 1 <= max_chars:
+                current_word_part += word + ' '
+            else:
+                if current_word_part:
+                    # Try to find a suitable split word
+                    for split_word in split_words:
+                        split_index = current_word_part.rfind(' ' + split_word + ' ')
+                        if split_index != -1:
+                            word_batches.append(current_word_part[:split_index].strip())
+                            current_word_part = current_word_part[split_index:].strip() + ' '
+                            break
+                    else:
+                        # If no suitable split word found, just append the current part
+                        word_batches.append(current_word_part.strip())
+                        current_word_part = ""
+                current_word_part += word + ' '
+        if current_word_part:
+            word_batches.append(current_word_part.strip())
+        return word_batches
+
+    for sentence in sentences:
+        if len(current_batch) + len(sentence) <= max_chars:
+            current_batch += sentence
+        else:
+            # If adding this sentence would exceed the limit
+            if current_batch:
+                batches.append(current_batch)
+                current_batch = ""
+            
+            # If the sentence itself is longer than max_chars, split it
+            if len(sentence) > max_chars:
+                # First, try to split by colon
+                colon_parts = sentence.split(':')
+                if len(colon_parts) > 1:
+                    for part in colon_parts:
+                        if len(part) <= max_chars:
+                            batches.append(part)
+                        else:
+                            # If colon part is still too long, split by comma
+                            comma_parts = part.split(',')
+                            if len(comma_parts) > 1:
+                                current_comma_part = ""
+                                for comma_part in comma_parts:
+                                    if len(current_comma_part) + len(comma_part) <= max_chars:
+                                        current_comma_part += comma_part + ','
+                                    else:
+                                        if current_comma_part:
+                                            batches.append(current_comma_part.rstrip(','))
+                                        current_comma_part = comma_part + ','
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                            else:
+                                # If no comma, split by words
+                                batches.extend(split_by_words(part))
+                else:
+                    # If no colon, split by comma
+                    comma_parts = sentence.split(',')
+                    if len(comma_parts) > 1:
+                        current_comma_part = ""
+                        for comma_part in comma_parts:
+                            if len(current_comma_part) + len(comma_part) <= max_chars:
+                                current_comma_part += comma_part + ','
+                            else:
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                                current_comma_part = comma_part + ','
+                        if current_comma_part:
+                            batches.append(current_comma_part.rstrip(','))
+                    else:
+                        # If no comma, split by words
+                        batches.extend(split_by_words(sentence))
+            else:
+                current_batch = sentence
+    
+    if current_batch:
+        batches.append(current_batch)
+    
+    return batches
+
+def infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence, progress=gr.Progress()):
+    if exp_name == "F5-TTS":
+        ema_model = F5TTS_ema_model
+    elif exp_name == "E2-TTS":
+        ema_model = E2TTS_ema_model
+
+    audio, sr = torchaudio.load(ref_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
+    if sr != target_sample_rate:
+        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+        audio = resampler(audio)
+    audio = audio.to(device)
+
+    generated_waves = []
+    spectrograms = []
+
+    for i, gen_text in enumerate(progress.tqdm(gen_text_batches)):
+        # 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
+        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
+        with torch.inference_mode():
+            generated, _ = ema_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")
+        
+        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()
+
+        if remove_silence:
+            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
+
+        generated_waves.append(generated_wave)
+        spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves
+    final_wave = np.concatenate(generated_waves)
+
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+    
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+        save_spectrogram(combined_spectrogram, spectrogram_path)
+
+    return (target_sample_rate, final_wave), spectrogram_path
 
 def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
     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)
@@ -101,14 +281,10 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
             aseg = aseg[:15000]
         aseg.export(f.name, format="wav")
         ref_audio = f.name
-    if exp_name == "F5-TTS":
-        ema_model = F5TTS_ema_model
-    elif exp_name == "E2-TTS":
-        ema_model = E2TTS_ema_model
 
     if not ref_text.strip():
         gr.Info("No reference text provided, transcribing reference audio...")
-        ref_text = outputs = pipe(
+        ref_text = pipe(
             ref_audio,
             chunk_length_s=30,
             batch_size=128,
@@ -118,80 +294,67 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
         gr.Info("Finished transcription")
     else:
         gr.Info("Using custom reference text...")
-    audio, sr = torchaudio.load(ref_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
-    if sr != target_sample_rate:
-        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, _ = ema_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()
-
-    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
-
-    # 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)
-
-    return (target_sample_rate, generated_wave), spectrogram_path
 
+    # Split the input text into batches
+    gen_text_batches = split_text_into_batches(gen_text)
+    
+    gr.Info(f"Generating audio using {exp_name} in {len(gen_text_batches)} batches")
+    return infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence)
+    
+def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name, ref_audio2, ref_text2, exp_name, remove_silence):
+    # Split the script into speaker blocks
+    speaker_pattern = re.compile(f"^({re.escape(speaker1_name)}|{re.escape(speaker2_name)}):", re.MULTILINE)
+    speaker_blocks = speaker_pattern.split(script)[1:]  # Skip the first empty element
+    
+    generated_audio_segments = []
+    
+    for i in range(0, len(speaker_blocks), 2):
+        speaker = speaker_blocks[i]
+        text = speaker_blocks[i+1].strip()
+        
+        # Determine which speaker is talking
+        if speaker == speaker1_name:
+            ref_audio = ref_audio1
+            ref_text = ref_text1
+        elif speaker == speaker2_name:
+            ref_audio = ref_audio2
+            ref_text = ref_text2
+        else:
+            continue  # Skip if the speaker is neither speaker1 nor speaker2
+        
+        # Generate audio for this block
+        audio, _ = infer(ref_audio, ref_text, text, exp_name, remove_silence)
+        
+        # Convert the generated audio to a numpy array
+        sr, audio_data = audio
+        
+        # Save the audio data as a WAV file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+            sf.write(temp_file.name, audio_data, sr)
+            audio_segment = AudioSegment.from_wav(temp_file.name)
+        
+        generated_audio_segments.append(audio_segment)
+        
+        # Add a short pause between speakers
+        pause = AudioSegment.silent(duration=500)  # 500ms pause
+        generated_audio_segments.append(pause)
+    
+    # Concatenate all audio segments
+    final_podcast = sum(generated_audio_segments)
+    
+    # Export the final podcast
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+        podcast_path = temp_file.name
+        final_podcast.export(podcast_path, format="wav")
+    
+    return podcast_path
 
 with gr.Blocks() as app:
     gr.Markdown(
         """
-# E2/F5 TTS
+# E2/F5 TTS with Advanced Batch Processing
 
-This is a local web UI for F5 TTS, based on the unofficial [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS). This app supports the following TTS models:
+This is a local web UI for F5 TTS with advanced batch processing support, based on the unofficial [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS). This app supports the following TTS models:
 
 * [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)
 * [E2-TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)
@@ -205,7 +368,7 @@ If you're having issues, try converting your reference audio to WAV or MP3, clip
     )
 
     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", lines=10)
     model_choice = gr.Radio(
         choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
     )
@@ -221,23 +384,75 @@ If you're having issues, try converting your reference audio to WAV or MP3, clip
             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,
         )
+        split_words_input = gr.Textbox(
+            label="Custom Split Words",
+            info="Enter custom words to split on, separated by commas. Leave blank to use default list.",
+            lines=2,
+        )
 
     audio_output = gr.Audio(label="Synthesized Audio")
     spectrogram_output = gr.Image(label="Spectrogram")
 
+    def infer_with_custom_split(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, custom_split_words):
+        if custom_split_words:
+            custom_words = [word.strip() for word in custom_split_words.split(',')]
+            global SPLIT_WORDS
+            SPLIT_WORDS = custom_words
+        return infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence)
+
     generate_btn.click(
-        infer,
+        infer_with_custom_split,
         inputs=[
             ref_audio_input,
             ref_text_input,
             gen_text_input,
             model_choice,
             remove_silence,
+            split_words_input,
         ],
         outputs=[audio_output, spectrogram_output],
     )
-
-
+    with gr.Tab("Podcast Generation"):
+        speaker1_name = gr.Textbox(label="Speaker 1 Name")
+        ref_audio_input1 = gr.Audio(label="Reference Audio (Speaker 1)", type="filepath")
+        ref_text_input1 = gr.Textbox(label="Reference Text (Speaker 1)", lines=2)
+        
+        speaker2_name = gr.Textbox(label="Speaker 2 Name")
+        ref_audio_input2 = gr.Audio(label="Reference Audio (Speaker 2)", type="filepath")
+        ref_text_input2 = gr.Textbox(label="Reference Text (Speaker 2)", lines=2)
+        
+        script_input = gr.Textbox(label="Podcast Script", lines=10, 
+                                  placeholder="Enter the script with speaker names at the start of each block, e.g.:\nSean: How did you start studying...\n\nMeghan: I came to my interest in technology...\nIt was a long journey...\n\nSean: That's fascinating. Can you elaborate...")
+        
+        podcast_model_choice = gr.Radio(
+            choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+        )
+        podcast_remove_silence = gr.Checkbox(
+            label="Remove Silences",
+            value=True,
+        )
+        generate_podcast_btn = gr.Button("Generate Podcast", variant="primary")
+        podcast_output = gr.Audio(label="Generated Podcast")
+    
+    def podcast_generation(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence):
+        return generate_podcast(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence)
+    
+    generate_podcast_btn.click(
+        podcast_generation,
+        inputs=[
+            script_input,
+            speaker1_name,
+            ref_audio_input1,
+            ref_text_input1,
+            speaker2_name,
+            ref_audio_input2,
+            ref_text_input2,
+            podcast_model_choice,
+            podcast_remove_silence,
+        ],
+        outputs=podcast_output,
+    )
+    
 @click.command()
 @click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
 @click.option("--host", "-H", default=None, help="Host to run the app on")