app.py

# ruff: noqa: E402
import json
import re
import tempfile
from collections import OrderedDict
from importlib.resources import files
from pydub import AudioSegment, silence

import click
import gradio as gr
import numpy as np
import soundfile as sf
import torchaudio
from cached_path import cached_path
from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    import spaces
    USING_SPACES = True
except ImportError:
    USING_SPACES = False

def gpu_decorator(func):
    if USING_SPACES:
        return spaces.GPU(func)
    else:
        return func

# ========== CORE CONFIGURATION ========== #
DEFAULT_TTS_MODEL = "F5-TTS"
DEFAULT_TTS_MODEL_CFG = [
    "hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors",
    "hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt",
    json.dumps(dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)),
]

last_used_custom = files("f5_tts").joinpath("infer/.cache/last_used_custom_model_info.txt")
tts_model_choice = DEFAULT_TTS_MODEL
target_sample_rate = 24000

# ========== CUSTOM MODEL HANDLING ========== #
def load_last_used_custom():
    try:
        with open(last_used_custom, "r", encoding="utf-8") as f:
            return [line.strip() for line in f.readlines()]
    except FileNotFoundError:
        last_used_custom.parent.mkdir(parents=True, exist_ok=True)
        return DEFAULT_TTS_MODEL_CFG

def set_custom_model(custom_ckpt_path, custom_vocab_path, custom_model_cfg):
    global tts_model_choice
    tts_model_choice = [
        "Custom", 
        custom_ckpt_path,
        custom_vocab_path,
        json.loads(custom_model_cfg)
    ]
    with open(last_used_custom, "w", encoding="utf-8") as f:
        f.write("\n".join([custom_ckpt_path, custom_vocab_path, custom_model_cfg]))

def switch_tts_model(new_choice):
    global tts_model_choice
    if new_choice == "Custom":
        custom_config = load_last_used_custom()
        tts_model_choice = ["Custom"] + custom_config
        return (
            gr.update(visible=True, value=custom_config[0]),
            gr.update(visible=True, value=custom_config[1]),
            gr.update(visible=True, value=custom_config[2])
        )
    else:
        tts_model_choice = new_choice
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False)
        )

# ========== MODEL LOADING ========== #
from f5_tts.model import DiT, UNetT
from f5_tts.infer.utils_infer import (
    load_vocoder,
    load_model,
    preprocess_ref_audio_text,
    infer_process,
    remove_silence_for_generated_wav,
    save_spectrogram,
)

vocoder = load_vocoder()

def load_f5tts(ckpt_path=str(cached_path(DEFAULT_TTS_MODEL_CFG[0]))):
    return load_model(DiT, json.loads(DEFAULT_TTS_MODEL_CFG[2]), ckpt_path)

def load_e2tts(ckpt_path=str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.safetensors"))):
    E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
    return load_model(UNetT, E2TTS_model_cfg, ckpt_path)

F5TTS_ema_model = load_f5tts()
E2TTS_ema_model = load_e2tts() if USING_SPACES else None
custom_ema_model, pre_custom_path = None, ""

# ========== CORE INFERENCE FUNCTION ========== #
@gpu_decorator
def generate_response(messages, model, tokenizer):
    """Generate response using Qwen"""
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.95,
    )

    generated_ids = [
        output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]


@gpu_decorator
def infer(
    ref_audio_orig,
    ref_text,
    gen_text,
    model,
    remove_silence,
    cross_fade_duration=0.15,
    nfe_step=32,
    speed=1,
    show_info=gr.Info,
):
    if not ref_audio_orig:
        gr.Warning("Please provide reference audio.")
        return gr.update(), gr.update(), ref_text

    if not gen_text.strip():
        gr.Warning("Please enter text to generate.")
        return gr.update(), gr.update(), ref_text

    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)

    if model == "F5-TTS":
        ema_model = F5TTS_ema_model
    elif model == "E2-TTS":
        global E2TTS_ema_model
        if E2TTS_ema_model is None:
            show_info("Loading E2-TTS model...")
            E2TTS_ema_model = load_e2tts()
        ema_model = E2TTS_ema_model
    elif isinstance(model, list) and model[0] == "Custom":
        assert not USING_SPACES, "Only official checkpoints allowed in Spaces."
        global custom_ema_model, pre_custom_path
        if pre_custom_path != model[1]:
            show_info("Loading Custom TTS model...")
            custom_ema_model = load_custom(model[1], vocab_path=model[2], model_cfg=model[3])
            pre_custom_path = model[1]
        ema_model = custom_ema_model

    final_wave, final_sample_rate, combined_spectrogram = infer_process(
        ref_audio,
        ref_text,
        gen_text,
        ema_model,
        vocoder,
        cross_fade_duration=cross_fade_duration,
        nfe_step=nfe_step,
        speed=speed,
        show_info=show_info,
        progress=gr.Progress(),
    )

    # Remove silence
    if remove_silence:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            sf.write(f.name, final_wave, final_sample_rate)
            remove_silence_for_generated_wav(f.name)
            final_wave, _ = torchaudio.load(f.name)
        final_wave = final_wave.squeeze().cpu().numpy()

    # Save the spectrogram
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
        spectrogram_path = tmp_spectrogram.name
        save_spectrogram(combined_spectrogram, spectrogram_path)

    return (final_sample_rate, final_wave), spectrogram_path, ref_text
    # ========== PODCAST UI WITH 10 SPEAKERS ========== #
with gr.Blocks() as app_podcast:
    gr.Markdown("# Podcast Generation (10 Speakers)")
    with gr.Row():
        with gr.Column():
            speaker1_name = gr.Textbox(label="Speaker 1 Name", placeholder="e.g. John")
            ref_audio_input1 = gr.Audio(label="Reference Audio 1", type="filepath")
            ref_text_input1 = gr.Textbox(label="Reference Text 1", lines=2)
        with gr.Column():
            speaker2_name = gr.Textbox(label="Speaker 2 Name", placeholder="e.g. Sarah") 
            ref_audio_input2 = gr.Audio(label="Reference Audio 2", type="filepath")
            ref_text_input2 = gr.Textbox(label="Reference Text 2", lines=2)
        with gr.Column():
            speaker3_name = gr.Textbox(label="Speaker 3 Name", placeholder="e.g. Alex")
            ref_audio_input3 = gr.Audio(label="Reference Audio 3", type="filepath")
            ref_text_input3 = gr.Textbox(label="Reference Text 3", lines=2)
        with gr.Column():
            speaker4_name = gr.Textbox(label="Speaker 4 Name", placeholder="e.g. Jane")
            ref_audio_input4 = gr.Audio(label="Reference Audio 4", type="filepath")
            ref_text_input4 = gr.Textbox(label="Reference Text 4", lines=2)
        with gr.Column():
            speaker5_name = gr.Textbox(label="Speaker 5 Name", placeholder="e.g. Mike")
            ref_audio_input5 = gr.Audio(label="Reference Audio 5", type="filepath")
            ref_text_input5 = gr.Textbox(label="Reference Text 5", lines=2)
    with gr.Row():
        with gr.Column():
            speaker6_name = gr.Textbox(label="Speaker 6 Name", placeholder="e.g. Anna")
            ref_audio_input6 = gr.Audio(label="Reference Audio 6", type="filepath")
            ref_text_input6 = gr.Textbox(label="Reference Text 6", lines=2)
        with gr.Column():
            speaker7_name = gr.Textbox(label="Speaker 7 Name", placeholder="e.g. Bob")
            ref_audio_input7 = gr.Audio(label="Reference Audio 7", type="filepath")
            ref_text_input7 = gr.Textbox(label="Reference Text 7", lines=2)
        with gr.Column():
            speaker8_name = gr.Textbox(label="Speaker 8 Name", placeholder="e.g. Carol")
            ref_audio_input8 = gr.Audio(label="Reference Audio 8", type="filepath")
            ref_text_input8 = gr.Textbox(label="Reference Text 8", lines=2)
        with gr.Column():
            speaker9_name = gr.Textbox(label="Speaker 9 Name", placeholder="e.g. Dave")
            ref_audio_input9 = gr.Audio(label="Reference Audio 9", type="filepath")
            ref_text_input9 = gr.Textbox(label="Reference Text 9", lines=2)
        with gr.Column():
            speaker10_name = gr.Textbox(label="Speaker 10 Name", placeholder="e.g. Emma")
            ref_audio_input10 = gr.Audio(label="Reference Audio 10", type="filepath")
            ref_text_input10 = gr.Textbox(label="Reference Text 10", lines=2)
    
    script_input = gr.Textbox(
        label="Podcast Script",
        lines=10,
        placeholder="Format:\nSpeaker1: Hello...\nSpeaker2: Hi...\n... Speaker10: Goodbye...\n"
    )
    
    with gr.Row():
        podcast_model_choice = gr.Radio(
            choices=["F5-TTS", "E2-TTS"],
            label="TTS Model",
            value="F5-TTS"
        )
        podcast_remove_silence = gr.Checkbox(
            label="Remove Silences Between Dialogues",
            value=True
        )
    
    generate_podcast_btn = gr.Button("Generate Podcast", variant="primary")
    podcast_output = gr.Audio(label="Generated Podcast", autoplay=True)
    
    @gpu_decorator
    def generate_podcast(
        script,
        speaker1,
        ref_audio1,
        ref_text1,
        speaker2,
        ref_audio2,
        ref_text2,
        speaker3,
        ref_audio3,
        ref_text3,
        speaker4,
        ref_audio4,
        ref_text4,
        speaker5,
        ref_audio5,
        ref_text5,
        speaker6,
        ref_audio6,
        ref_text6,
        speaker7,
        ref_audio7,
        ref_text7,
        speaker8,
        ref_audio8,
        ref_text8,
        speaker9,
        ref_audio9,
        ref_text9,
        speaker10,
        ref_audio10,
        ref_text10,
        model,
        remove_silence
    ):
        # Validate inputs
        if not all([speaker1, speaker2, speaker3, speaker4, speaker5, speaker6, speaker7, speaker8, speaker9, speaker10]):
            raise gr.Error("All ten speaker names must be provided")
        if not all([ref_audio1, ref_audio2, ref_audio3, ref_audio4, ref_audio5, ref_audio6, ref_audio7, ref_audio8, ref_audio9, ref_audio10]):
            raise gr.Error("All ten reference audios must be provided")
    
        # Split script using regex pattern for 10 speakers
        pattern = re.compile(
            f"({re.escape(speaker1)}:|{re.escape(speaker2)}:|{re.escape(speaker3)}:|{re.escape(speaker4)}:|{re.escape(speaker5)}:|{re.escape(speaker6)}:|{re.escape(speaker7)}:|{re.escape(speaker8)}:|{re.escape(speaker9)}:|{re.escape(speaker10)}:)"
        )
        speaker_blocks = pattern.split(script)[1:]
        
        generated_audio_segments = []
        speaker_map = {
            speaker1: {"audio": ref_audio1, "text": ref_text1},
            speaker2: {"audio": ref_audio2, "text": ref_text2},
            speaker3: {"audio": ref_audio3, "text": ref_text3},
            speaker4: {"audio": ref_audio4, "text": ref_text4},
            speaker5: {"audio": ref_audio5, "text": ref_text5},
            speaker6: {"audio": ref_audio6, "text": ref_text6},
            speaker7: {"audio": ref_audio7, "text": ref_text7},
            speaker8: {"audio": ref_audio8, "text": ref_text8},
            speaker9: {"audio": ref_audio9, "text": ref_text9},
            speaker10: {"audio": ref_audio10, "text": ref_text10}
        }
    
        for i in range(0, len(speaker_blocks), 2):
            if i+1 >= len(speaker_blocks):
                break
                    
            speaker_tag = speaker_blocks[i].strip(":")
            text = speaker_blocks[i+1].strip()
                    
            # Get speaker config
            config = speaker_map.get(speaker_tag)
            if not config:
                continue
    
            # Generate audio using original infer function
            audio_result, _, _ = infer(
                config["audio"],
                config["text"],
                text,
                model,
                remove_silence,
                cross_fade_duration=0.15,
                nfe_step=32,
                speed=1.0
            )
            sr, audio_data = audio_result
            generated_audio_segments.append(audio_data)
    
            # Add pause between speakers (500ms silence)
            if i < len(speaker_blocks)-2:
                generated_audio_segments.append(np.zeros(int(0.5 * sr)))
    
        # Combine all audio segments
        if generated_audio_segments:
            final_audio = np.concatenate(generated_audio_segments)
            return (target_sample_rate, final_audio)
        return None
    
    generate_podcast_btn.click(
        generate_podcast,
        inputs=[
            script_input,
            speaker1_name, ref_audio_input1, ref_text_input1,
            speaker2_name, ref_audio_input2, ref_text_input2,
            speaker3_name, ref_audio_input3, ref_text_input3,
            speaker4_name, ref_audio_input4, ref_text_input4,
            speaker5_name, ref_audio_input5, ref_text_input5,
            speaker6_name, ref_audio_input6, ref_text_input6,
            speaker7_name, ref_audio_input7, ref_text_input7,
            speaker8_name, ref_audio_input8, ref_text_input8,
            speaker9_name, ref_audio_input9, ref_text_input9,
            speaker10_name, ref_audio_input10, ref_text_input10,
            podcast_model_choice, podcast_remove_silence
        ],
        outputs=podcast_output
    )
with gr.Blocks() as app_credits:
    gr.Markdown("""
# Credits

* [mrfakename](https://github.com/fakerybakery) for the original [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
* [RootingInLoad](https://github.com/RootingInLoad) for initial chunk generation and podcast app exploration
* [jpgallegoar](https://github.com/jpgallegoar) for multiple speech-type generation & voice chat
""")

with gr.Blocks() as app_tts:
    gr.Markdown("# Batched TTS")
    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
    gen_text_input = gr.Textbox(label="Text to Generate", lines=10)
    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=False,
        )
        speed_slider = gr.Slider(
            label="Speed",
            minimum=0.3,
            maximum=2.0,
            value=1.0,
            step=0.1,
            info="Adjust the speed of the audio.",
        )
        nfe_slider = gr.Slider(
            label="NFE Steps",
            minimum=4,
            maximum=64,
            value=32,
            step=2,
            info="Set the number of denoising steps.",
        )
        cross_fade_duration_slider = gr.Slider(
            label="Cross-Fade Duration (s)",
            minimum=0.0,
            maximum=1.0,
            value=0.15,
            step=0.01,
            info="Set the duration of the cross-fade between audio clips.",
        )

    audio_output = gr.Audio(label="Synthesized Audio")
    spectrogram_output = gr.Image(label="Spectrogram")

    @gpu_decorator
    def basic_tts(
        ref_audio_input,
        ref_text_input,
        gen_text_input,
        remove_silence,
        cross_fade_duration_slider,
        nfe_slider,
        speed_slider,
    ):
        audio_out, spectrogram_path, ref_text_out = infer(
            ref_audio_input,
            ref_text_input,
            gen_text_input,
            tts_model_choice,
            remove_silence,
            cross_fade_duration=cross_fade_duration_slider,
            nfe_step=nfe_slider,
            speed=speed_slider,
        )
        return audio_out, spectrogram_path, ref_text_out

    generate_btn.click(
        basic_tts,
        inputs=[
            ref_audio_input,
            ref_text_input,
            gen_text_input,
            remove_silence,
            cross_fade_duration_slider,
            nfe_slider,
            speed_slider,
        ],
        outputs=[audio_output, spectrogram_output, ref_text_input],
    )
with gr.Blocks() as app_multistyle:
    gr.Markdown("# Multiple Speech-Type Generation")
    # ... [Keep original multistyle interface unchanged] ...
with gr.Blocks() as app_chat:
    gr.Markdown("# Voice Chat")
    # ... [Keep original voice chat interface unchanged] ...

with gr.Blocks() as app:
    gr.Markdown(f"""
# E2/F5 TTS
{"Local web UI for [F5 TTS](https://github.com/SWivid/F5-TTS)" if not USING_SPACES else "Online demo for [F5-TTS](https://github.com/SWivid/F5-TTS)"}
""")
    
    with gr.Row():
        if not USING_SPACES:
            choose_tts_model = gr.Radio(
                choices=[DEFAULT_TTS_MODEL, "E2-TTS", "Custom"],
                label="TTS Model",
                value=DEFAULT_TTS_MODEL
            )
        else:
            choose_tts_model = gr.Radio(
                choices=[DEFAULT_TTS_MODEL, "E2-TTS"],
                label="TTS Model",
                value=DEFAULT_TTS_MODEL
            )
        
        custom_ckpt_path = gr.Dropdown(
            choices=[DEFAULT_TTS_MODEL_CFG[0]],
            value=load_last_used_custom()[0],
            allow_custom_value=True,
            label="Model Path",
            visible=False
        )
        custom_vocab_path = gr.Dropdown(
            choices=[DEFAULT_TTS_MODEL_CFG[1]],
            value=load_last_used_custom()[1],
            allow_custom_value=True,
            label="Vocab Path",
            visible=False
        )
        custom_model_cfg = gr.Dropdown(
            choices=[DEFAULT_TTS_MODEL_CFG[2]],
            value=load_last_used_custom()[2],
            allow_custom_value=True,
            label="Model Config",
            visible=False
        )

    choose_tts_model.change(
        switch_tts_model,
        inputs=[choose_tts_model],
        outputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg]
    )
    
    gr.TabbedInterface(
        [app_tts, app_podcast, app_multistyle, app_chat, app_credits],
        ["Basic TTS", "Podcast", "Multi-Style", "Voice Chat", "Credits"],
    )

@click.command()
@click.option("--port", "-p", default=None, type=int)
@click.option("--host", "-H", default=None)
@click.option("--share", "-s", default=True, is_flag=True)
@click.option("--api", "-a", default=True, is_flag=True)
@click.option("--root_path", "-r", default=None)
def main(port, host, share, api, root_path):
    global app
    print("Launching app...")
    app.queue(api_open=api).launch(
        server_name=host,
        server_port=port,
        share=share,
        show_api=api,
        root_path=root_path
    )

if __name__ == "__main__":
    if not USING_SPACES:
        main()
    else:
        app.queue().launch()
# ========== REST OF ORIGINAL CODE (UNCHANGED BELOW) ========== #
# [Main app configuration, other tabs (TTS, Multistyle, Chat, Credits)]
# [Keep all original code after this point exactly as i