Update app.py

app.py

@@ -3,404 +3,72 @@ import torchaudio
 import gradio as gr
 
 from zonos.model import Zonos
-from zonos.conditioning import make_cond_dict, supported_language_codes
+from zonos.conditioning import make_cond_dict
 
-device = "cuda"
-CURRENT_MODEL_TYPE = None
-CURRENT_MODEL = None
+# Load the hybrid model
+model = Zonos.from_pretrained("Zyphra/Zonos-v0.1-hybrid", device="cuda")
+model.bfloat16()  # Switch model weights to bfloat16 precision (optional, but recommended for GPU)
 
-
-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,
-):
+# Main inference function for Gradio
+def tts(text, reference_audio):
     """
-    Generates audio based on the provided UI parameters.
-    We do NOT use language_id or ctc_loss even if the model has them.
+    text: str
+    reference_audio: (numpy.ndarray, int) -> (data, sample_rate)
     """
-    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)
-
+    if reference_audio is None:
+        return "No reference audio provided."
+    
+    # reference_audio[0] is a NumPy float32 array of shape (num_samples, 1) or (num_samples,)
+    # reference_audio[1] is the sample rate
+    wav_np, sr = reference_audio
+    
+    # Convert NumPy audio to Torch tensor
+    wav_torch = torch.from_numpy(wav_np).float().unsqueeze(0)  # shape: (1, num_samples)
+    if wav_torch.dim() == 2 and wav_torch.shape[0] > wav_torch.shape[1]:
+        # If the shape is (samples, 1), reorder to (1, samples)
+        wav_torch = wav_torch.T
+    
+    # Create speaker embedding
+    spk_embedding = model.embed_spk_audio(wav_torch, sr)
+    
+    # Prepare conditioning
     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),
+        speaker=spk_embedding.to(torch.bfloat16),
+        language="en-us",
     )
-
-    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
-
+    conditioning = model.prepare_conditioning(cond_dict)
+    
+    # Generate codes
+    with torch.no_grad():
+        torch.manual_seed(421)  # Seeding for reproducible results
+        codes = model.generate(conditioning)
+    
+    # Decode the codes into waveform
+    wavs = model.autoencoder.decode(codes).cpu()
+    out_audio = wavs[0].numpy()  # shape: (num_samples,)
+
+    # Return as (sample_rate, audio_ndarray) for Gradio's "audio" output
+    return (model.autoencoder.sampling_rate, out_audio)
+
+
+# Define the Gradio interface
+# - text input for the prompt
+# - audio input for the speaker reference
+# - audio output with the generated speech
+demo = gr.Interface(
+    fn=tts,
+    inputs=[
+        gr.Textbox(label="Text to Synthesize"),
+        gr.Audio(source="upload", type="numpy", label="Reference Audio (for speaker embedding)"),
+    ],
+    outputs=gr.Audio(label="Generated Audio"),
+    title="Zonos TTS Demo (Hybrid)",
+    description=(
+        "Provide a reference audio snippet for speaker embedding, "
+        "enter text, and generate speech with Zonos TTS."
+    ),
+)
 
 if __name__ == "__main__":
-    demo = build_interface()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.launch(debug=True)