app-func.py

# ruff: noqa: E402
# Above allows ruff to ignore E402: module level import not at top of file

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


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,
)


DEFAULT_TTS_MODEL = "F5-TTS"
tts_model_choice = DEFAULT_TTS_MODEL

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)),
]
# Add this right after DEFAULT_TTS_MODEL_CFG
def switch_tts_model(new_choice):
    global tts_model_choice
    if new_choice == "Custom":
        custom_ckpt_path, custom_vocab_path, custom_model_cfg = load_last_used_custom()
        tts_model_choice = [
            "Custom", 
            custom_ckpt_path,
            custom_vocab_path,
            json.loads(custom_model_cfg)
        ]
        return (
            gr.update(visible=True, value=custom_ckpt_path),
            gr.update(visible=True, value=custom_vocab_path),
            gr.update(visible=True, value=custom_model_cfg)
        )
    else:
        tts_model_choice = new_choice
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False)
        )
# Add this right after DEFAULT_TTS_MODEL_CFG definition
last_used_custom = files("f5_tts").joinpath("infer/.cache/last_used_custom_model_info.txt")

def load_last_used_custom():
    try:
        custom = []
        with open(last_used_custom, "r", encoding="utf-8") as f:
            for line in f:
                custom.append(line.strip())
        return custom
    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]) + "\n")
# Audio constants for podcast
target_sample_rate = 24000

# load models

vocoder = load_vocoder()


def load_f5tts(ckpt_path=str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors"))):
    F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
    return load_model(DiT, F5TTS_model_cfg, 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)


def load_custom(ckpt_path: str, vocab_path="", model_cfg=None):
    ckpt_path, vocab_path = ckpt_path.strip(), vocab_path.strip()
    if ckpt_path.startswith("hf://"):
        ckpt_path = str(cached_path(ckpt_path))
    if vocab_path.startswith("hf://"):
        vocab_path = str(cached_path(vocab_path))
    if model_cfg is None:
        model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
    return load_model(DiT, model_cfg, ckpt_path, vocab_file=vocab_path)


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

chat_model_state = None
chat_tokenizer_state = None


@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


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_podcast:
    gr.Markdown("# Podcast Generation")
    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 (Speaker 1)", type="filepath")
            ref_text_input1 = gr.Textbox(label="Reference Text (Speaker 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 (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="Format:\nSpeaker1: Hello...\nSpeaker2: Hi...\nSpeaker1: How are you?..."
    )

    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)

    def generate_podcast(
        script,
        speaker1,
        ref_audio1,
        ref_text1,
        speaker2,
        ref_audio2,
        ref_text2,
        model,
        remove_silence
    ):
        # Validate inputs
        if not all([speaker1, speaker2]):
            raise gr.Error("Both speaker names must be provided")
        if not ref_audio1 or not ref_audio2:
            raise gr.Error("Both reference audios must be provided")

        # Split script into speaker blocks
        pattern = re.compile(f"({re.escape(speaker1)}:|{re.escape(speaker2)}:)")
        speaker_blocks = pattern.split(script)[1:]
        
        generated_audio_segments = []
        current_speaker = None

        for i in range(0, len(speaker_blocks), 2):
            speaker_tag = speaker_blocks[i].strip(":")
            text = speaker_blocks[i+1].strip()
            
            # Select reference based on speaker
            if speaker_tag == speaker1:
                ref_audio = ref_audio1
                ref_text = ref_text1
            elif speaker_tag == speaker2:
                ref_audio = ref_audio2
                ref_text = ref_text2
            else:
                continue

            # Generate audio for segment
            audio_result, spectrogram, ref_text_out = infer(
                ref_audio,
                ref_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)

        # 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,
            podcast_model_choice,
            podcast_remove_silence
        ],
        outputs=podcast_output
    )

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()