app.py

import torch
import torchaudio
import gradio as gr

from zonos.model import Zonos
from zonos.conditioning import make_cond_dict, supported_language_codes

device = "cuda"
CURRENT_MODEL_TYPE = None
CURRENT_MODEL = None


def load_model_if_needed(model_choice: str):
    global CURRENT_MODEL_TYPE, CURRENT_MODEL
    if CURRENT_MODEL_TYPE != model_choice:
        if CURRENT_MODEL is not None:
            del CURRENT_MODEL
            torch.cuda.empty_cache()
        print(f"Loading {model_choice} model...")
        if model_choice == "Transformer":
            CURRENT_MODEL = Zonos.from_pretrained("Zyphra/Zonos-v0.1-transformer", device=device)
        else:
            CURRENT_MODEL = Zonos.from_pretrained("Zyphra/Zonos-v0.1-hybrid", device=device)
        CURRENT_MODEL.to(device)
        CURRENT_MODEL.bfloat16()
        CURRENT_MODEL.eval()
        CURRENT_MODEL_TYPE = model_choice
        print(f"{model_choice} model loaded successfully!")
    else:
        print(f"{model_choice} model is already loaded.")
    return CURRENT_MODEL


def update_ui(model_choice):
    """
    Dynamically show/hide UI elements based on the model's conditioners.
    We do NOT display 'language_id' or 'ctc_loss' even if they exist in the model.
    """
    model = load_model_if_needed(model_choice)
    cond_names = [c.name for c in model.prefix_conditioner.conditioners]
    print("Conditioners in this model:", cond_names)

    text_update = gr.update(visible=("espeak" in cond_names))
    language_update = gr.update(visible=("espeak" in cond_names))
    speaker_audio_update = gr.update(visible=("speaker" in cond_names))
    prefix_audio_update = gr.update(visible=True)
    skip_speaker_update = gr.update(visible=("speaker" in cond_names))
    skip_emotion_update = gr.update(visible=("emotion" in cond_names))
    emotion1_update = gr.update(visible=("emotion" in cond_names))
    emotion2_update = gr.update(visible=("emotion" in cond_names))
    emotion3_update = gr.update(visible=("emotion" in cond_names))
    emotion4_update = gr.update(visible=("emotion" in cond_names))
    emotion5_update = gr.update(visible=("emotion" in cond_names))
    emotion6_update = gr.update(visible=("emotion" in cond_names))
    emotion7_update = gr.update(visible=("emotion" in cond_names))
    emotion8_update = gr.update(visible=("emotion" in cond_names))
    skip_vqscore_8_update = gr.update(visible=("vqscore_8" in cond_names))
    vq_single_slider_update = gr.update(visible=("vqscore_8" in cond_names))
    fmax_slider_update = gr.update(visible=("fmax" in cond_names))
    skip_fmax_update = gr.update(visible=("fmax" in cond_names))
    pitch_std_slider_update = gr.update(visible=("pitch_std" in cond_names))
    skip_pitch_std_update = gr.update(visible=("pitch_std" in cond_names))
    speaking_rate_slider_update = gr.update(visible=("speaking_rate" in cond_names))
    skip_speaking_rate_update = gr.update(visible=("speaking_rate" in cond_names))
    dnsmos_slider_update = gr.update(visible=("dnsmos_ovrl" in cond_names))
    skip_dnsmos_ovrl_update = gr.update(visible=("dnsmos_ovrl" in cond_names))
    speaker_noised_checkbox_update = gr.update(visible=("speaker_noised" in cond_names))
    skip_speaker_noised_update = gr.update(visible=("speaker_noised" in cond_names))

    return (
        text_update,  # 1
        language_update,  # 2
        speaker_audio_update,  # 3
        prefix_audio_update,  # 4
        skip_speaker_update,  # 5
        skip_emotion_update,  # 6
        emotion1_update,  # 7
        emotion2_update,  # 8
        emotion3_update,  # 9
        emotion4_update,  # 10
        emotion5_update,  # 11
        emotion6_update,  # 12
        emotion7_update,  # 13
        emotion8_update,  # 14
        skip_vqscore_8_update,  # 15
        vq_single_slider_update,  # 16
        fmax_slider_update,  # 17
        skip_fmax_update,  # 18
        pitch_std_slider_update,  # 19
        skip_pitch_std_update,  # 20
        speaking_rate_slider_update,  # 21
        skip_speaking_rate_update,  # 22
        dnsmos_slider_update,  # 23
        skip_dnsmos_ovrl_update,  # 24
        speaker_noised_checkbox_update,  # 25
        skip_speaker_noised_update,  # 26
    )


def generate_audio(
    model_choice,
    text,
    language,
    speaker_audio,
    prefix_audio,
    skip_speaker,
    skip_emotion,
    e1,
    e2,
    e3,
    e4,
    e5,
    e6,
    e7,
    e8,
    skip_vqscore_8,
    vq_single,
    fmax,
    skip_fmax,
    pitch_std,
    skip_pitch_std,
    speaking_rate,
    skip_speaking_rate,
    dnsmos_ovrl,
    skip_dnsmos_ovrl,
    speaker_noised,
    skip_speaker_noised,
    cfg_scale,
    min_p,
    seed,
):
    """
    Generates audio based on the provided UI parameters.
    We do NOT use language_id or ctc_loss even if the model has them.
    """
    selected_model = load_model_if_needed(model_choice)

    uncond_keys = []
    if skip_speaker:
        uncond_keys.append("speaker")
    if skip_emotion:
        uncond_keys.append("emotion")
    if skip_vqscore_8:
        uncond_keys.append("vqscore_8")
    if skip_fmax:
        uncond_keys.append("fmax")
    if skip_pitch_std:
        uncond_keys.append("pitch_std")
    if skip_speaking_rate:
        uncond_keys.append("speaking_rate")
    if skip_dnsmos_ovrl:
        uncond_keys.append("dnsmos_ovrl")
    if skip_speaker_noised:
        uncond_keys.append("speaker_noised")

    speaker_noised_bool = bool(speaker_noised)
    fmax = float(fmax)
    pitch_std = float(pitch_std)
    speaking_rate = float(speaking_rate)
    dnsmos_ovrl = float(dnsmos_ovrl)
    cfg_scale = float(cfg_scale)
    min_p = float(min_p)
    seed = int(seed)
    max_new_tokens = 86 * 30

    torch.manual_seed(seed)

    speaker_embedding = None
    if speaker_audio is not None and not skip_speaker:
        wav, sr = torchaudio.load(speaker_audio)
        speaker_embedding = selected_model.make_speaker_embedding(wav, sr)
        speaker_embedding = speaker_embedding.to(device, dtype=torch.bfloat16)

    audio_prefix_codes = None
    if prefix_audio is not None:
        wav_prefix, sr_prefix = torchaudio.load(prefix_audio)
        wav_prefix = wav_prefix.mean(0, keepdim=True)
        wav_prefix = torchaudio.functional.resample(wav_prefix, sr_prefix, selected_model.autoencoder.sampling_rate)
        wav_prefix = wav_prefix.to(device, dtype=torch.float32)
        with torch.autocast(device, dtype=torch.float32):
            audio_prefix_codes = selected_model.autoencoder.encode(wav_prefix.unsqueeze(0))

    emotion_tensor = torch.tensor(
        [[float(e1), float(e2), float(e3), float(e4), float(e5), float(e6), float(e7), float(e8)]], device=device
    )

    vq_val = float(vq_single)
    vq_tensor = torch.tensor([vq_val] * 8, device=device).unsqueeze(0)

    cond_dict = make_cond_dict(
        text=text,
        language=language,
        speaker=speaker_embedding,
        emotion=emotion_tensor,
        vqscore_8=vq_tensor,
        fmax=fmax,
        pitch_std=pitch_std,
        speaking_rate=speaking_rate,
        dnsmos_ovrl=dnsmos_ovrl,
        speaker_noised=speaker_noised_bool,
        device=device,
        unconditional_keys=uncond_keys,
    )
    conditioning = selected_model.prepare_conditioning(cond_dict)

    codes = selected_model.generate(
        prefix_conditioning=conditioning,
        audio_prefix_codes=audio_prefix_codes,
        max_new_tokens=max_new_tokens,
        cfg_scale=cfg_scale,
        batch_size=1,
        sampling_params=dict(min_p=min_p),
    )

    wav_out = selected_model.autoencoder.decode(codes).cpu().detach()
    sr_out = selected_model.autoencoder.sampling_rate
    if wav_out.dim() == 2 and wav_out.size(0) > 1:
        wav_out = wav_out[0:1, :]
    return sr_out, wav_out.squeeze().numpy()


def build_interface():
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                model_choice = gr.Dropdown(
                    choices=["Hybrid", "Transformer"],
                    value="Transformer",
                    label="Zonos Model Type",
                    info="Select the model variant to use.",
                )
                text = gr.Textbox(
                    label="Text to Synthesize", value="Zonos uses eSpeak for text to phoneme conversion!", lines=4
                )
                language = gr.Dropdown(
                    choices=supported_language_codes,
                    value="en-us",
                    label="Language Code",
                    info="Select a language code.",
                )
            prefix_audio = gr.Audio(
                value="assets/silence_100ms.wav",
                label="Optional Prefix Audio (continue from this audio)",
                type="filepath",
            )
            with gr.Column():
                speaker_audio = gr.Audio(
                    label="Optional Speaker Audio (for cloning)",
                    type="filepath",
                )
                speaker_noised_checkbox = gr.Checkbox(label="Denoise Speaker?", value=False)

        with gr.Column():
            gr.Markdown("## Conditioning Parameters")

            with gr.Row():
                dnsmos_slider = gr.Slider(1.0, 5.0, value=4.0, step=0.1, label="DNSMOS Overall")
                fmax_slider = gr.Slider(0, 24000, value=22050, step=1, label="Fmax (Hz)")
                vq_single_slider = gr.Slider(0.5, 0.8, 0.78, 0.01, label="VQ Score")
                pitch_std_slider = gr.Slider(0.0, 400.0, value=20.0, step=1, label="Pitch Std")
                speaking_rate_slider = gr.Slider(0.0, 40.0, value=15.0, step=1, label="Speaking Rate")

            gr.Markdown("### Emotion Sliders")
            with gr.Row():
                emotion1 = gr.Slider(0.0, 1.0, 0.6, 0.05, label="Happiness")
                emotion2 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Sadness")
                emotion3 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Disgust")
                emotion4 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Fear")
            with gr.Row():
                emotion5 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Surprise")
                emotion6 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Anger")
                emotion7 = gr.Slider(0.0, 1.0, 0.5, 0.05, label="Other")
                emotion8 = gr.Slider(0.0, 1.0, 0.6, 0.05, label="Neutral")

            gr.Markdown("### Unconditional Toggles")
            with gr.Row():
                skip_speaker = gr.Checkbox(label="Skip Speaker", value=False)
                skip_emotion = gr.Checkbox(label="Skip Emotion", value=False)
                skip_vqscore_8 = gr.Checkbox(label="Skip VQ Score", value=True)
                skip_fmax = gr.Checkbox(label="Skip Fmax", value=False)
                skip_pitch_std = gr.Checkbox(label="Skip Pitch Std", value=False)
                skip_speaking_rate = gr.Checkbox(label="Skip Speaking Rate", value=False)
                skip_dnsmos_ovrl = gr.Checkbox(label="Skip DNSMOS", value=True)
                skip_speaker_noised = gr.Checkbox(label="Skip Noised Speaker", value=False)

        with gr.Column():
            gr.Markdown("## Generation Parameters")
            with gr.Row():
                cfg_scale_slider = gr.Slider(1.0, 5.0, 2.0, 0.1, label="CFG Scale")
                min_p_slider = gr.Slider(0.0, 1.0, 0.1, 0.01, label="Min P")
                seed_number = gr.Number(label="Seed", value=420, precision=0)

            generate_button = gr.Button("Generate Audio")
            output_audio = gr.Audio(label="Generated Audio", type="numpy")

        model_choice.change(
            fn=update_ui,
            inputs=[model_choice],
            outputs=[
                text,  # 1
                language,  # 2
                speaker_audio,  # 3
                prefix_audio,  # 4
                skip_speaker,  # 5
                skip_emotion,  # 6
                emotion1,  # 7
                emotion2,  # 8
                emotion3,  # 9
                emotion4,  # 10
                emotion5,  # 11
                emotion6,  # 12
                emotion7,  # 13
                emotion8,  # 14
                skip_vqscore_8,  # 15
                vq_single_slider,  # 16
                fmax_slider,  # 17
                skip_fmax,  # 18
                pitch_std_slider,  # 19
                skip_pitch_std,  # 20
                speaking_rate_slider,  # 21
                skip_speaking_rate,  # 22
                dnsmos_slider,  # 23
                skip_dnsmos_ovrl,  # 24
                speaker_noised_checkbox,  # 25
                skip_speaker_noised,  # 26
            ],
        )

        # On page load, trigger the same UI refresh
        demo.load(
            fn=update_ui,
            inputs=[model_choice],
            outputs=[
                text,
                language,
                speaker_audio,
                prefix_audio,
                skip_speaker,
                skip_emotion,
                emotion1,
                emotion2,
                emotion3,
                emotion4,
                emotion5,
                emotion6,
                emotion7,
                emotion8,
                skip_vqscore_8,
                vq_single_slider,
                fmax_slider,
                skip_fmax,
                pitch_std_slider,
                skip_pitch_std,
                speaking_rate_slider,
                skip_speaking_rate,
                dnsmos_slider,
                skip_dnsmos_ovrl,
                speaker_noised_checkbox,
                skip_speaker_noised,
            ],
        )

        # Generate audio on button click
        generate_button.click(
            fn=generate_audio,
            inputs=[
                model_choice,
                text,
                language,
                speaker_audio,
                prefix_audio,
                skip_speaker,
                skip_emotion,
                emotion1,
                emotion2,
                emotion3,
                emotion4,
                emotion5,
                emotion6,
                emotion7,
                emotion8,
                skip_vqscore_8,
                vq_single_slider,
                fmax_slider,
                skip_fmax,
                pitch_std_slider,
                skip_pitch_std,
                speaking_rate_slider,
                skip_speaking_rate,
                dnsmos_slider,
                skip_dnsmos_ovrl,
                speaker_noised_checkbox,
                skip_speaker_noised,
                cfg_scale_slider,
                min_p_slider,
                seed_number,
            ],
            outputs=[output_audio],
        )

    return demo


if __name__ == "__main__":
    demo = build_interface()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)