Add NeuTTS-Air model alongside Kokoro TTS with model selection

app.py

@@ -1,18 +1,27 @@
-"""
-🎙️ Kokoro-TTS-only demo – Zero-GPU edition
-Routes every synthesis to an idle A100.
-"""
-
-import os, tempfile, subprocess, numpy as np
+import spaces
+import os
+import sys
+import tempfile
+import numpy as np
 import gradio as gr
-import spaces           # Zero-GPU decorator
 import soundfile as sf
 
-# ------------------------------------------------------------------
-# 1.  Lazy Kokoro loader (runs once per GPU worker)
-# ------------------------------------------------------------------
+# Add NeuTTS-Air to path
+sys.path.append("neutts-air")
+
+# Global variables for lazy loading
 kokoro_pipe = None
+neutts_model = None
 
+# NeuTTS-Air configuration
+SAMPLES_PATH = os.path.join(os.getcwd(), "neutts-air", "samples")
+DEFAULT_REF_TEXT = "So I'm live on radio. And I say, well, my dear friend James here clearly, and the whole room just froze. Turns out I'd completely misspoken and mentioned our other friend."
+DEFAULT_REF_PATH = os.path.join(SAMPLES_PATH, "dave.wav")
+DEFAULT_GEN_TEXT = "My name is Dave, and um, I'm from London."
+
+# ------------------------------------------------------------------
+# 1. Lazy loaders
+# ------------------------------------------------------------------
 def load_kokoro():
     global kokoro_pipe
     if kokoro_pipe is None:
@@ -20,60 +29,150 @@ def load_kokoro():
         kokoro_pipe = KPipeline(lang_code='a')
     return kokoro_pipe
 
+def load_neutts():
+    global neutts_model
+    if neutts_model is None:
+        from neuttsair.neutts import NeuTTSAir
+        neutts_model = NeuTTSAir(
+            backbone_repo="neuphonic/neutts-air",
+            backbone_device="cuda",
+            codec_repo="neuphonic/neucodec",
+            codec_device="cuda"
+        )
+    return neutts_model
+
 # ------------------------------------------------------------------
-# 2.  Generation helper
+# 2. Kokoro TTS inference
 # ------------------------------------------------------------------
-@spaces.GPU
-def tts_kokoro(text, voice, speed):
+@spaces.GPU()
+def kokoro_infer(text, voice, speed):
+    if not text.strip():
+        raise gr.Error("Please enter some text.")
+    
     pipe = load_kokoro()
     generator = pipe(text, voice=voice, speed=speed)
     for gs, ps, audio in generator:
-        return audio
+        # Save to temporary file
+        fd, tmp = tempfile.mkstemp(suffix='.wav')
+        os.close(fd)
+        sf.write(tmp, audio, 24000)
+        return tmp
     raise RuntimeError("Kokoro generation failed")
 
 # ------------------------------------------------------------------
-# 3.  Zero-GPU entry point
+# 3. NeuTTS-Air inference
 # ------------------------------------------------------------------
-@spaces.GPU
-def synthesise(text, voice, speed):
-    if not text.strip():
-        raise gr.Error("Please enter some text.")
-    wav = tts_kokoro(text, voice=voice, speed=speed)
-    fd, tmp = tempfile.mkstemp(suffix='.wav')
-    os.close(fd)
-    sf.write(tmp, wav, 24000)
-    return tmp
+@spaces.GPU()
+def neutts_infer(ref_text: str, ref_audio_path: str, gen_text: str) -> tuple:
+    """
+    Generates speech using NeuTTS-Air given a reference audio and text, and new text to synthesize.
+    """
+    if not gen_text.strip():
+        raise gr.Error("Please enter text to generate.")
+    if not ref_audio_path:
+        raise gr.Error("Please provide reference audio.")
+    if not ref_text.strip():
+        raise gr.Error("Please provide reference text.")
+    
+    gr.Info("Starting inference request!")
+    gr.Info("Encoding reference...")
+    
+    tts = load_neutts()
+    ref_codes = tts.encode_reference(ref_audio_path)
+    
+    gr.Info(f"Generating audio for input text: {gen_text}")
+    wav = tts.infer(gen_text, ref_codes, ref_text)
+    
+    return (24_000, wav)
 
 # ------------------------------------------------------------------
-# 4.  Gradio UI
+# 4. Gradio UI with model selection
 # ------------------------------------------------------------------
 css = """footer {visibility: hidden}"""
 
-with gr.Blocks(css=css, title="Kokoro TTS – Zero-GPU") as demo:
-    gr.Markdown("## 🎙️ Kokoro TTS – Zero-GPU Demo")
-
-    with gr.Row():
-        with gr.Column():
-            voice = gr.Dropdown(
-                label="Voice",
-                choices=['af_heart', 'af_sky', 'af_mist', 'af_dusk'],
-                value='af_heart'
-            )
-            speed = gr.Slider(0.5, 2.0, 1.0, step=0.1, label="Speed")
-
-        with gr.Column(scale=3):
-            text = gr.Textbox(
-                label="Text to speak",
-                placeholder="Type or paste text here …",
-                lines=6, max_lines=12
-            )
-            btn = gr.Button("🎧 Synthesise", variant="primary")
-            audio_out = gr.Audio(label="Generated speech", type="filepath")
-
-    btn.click(synthesise, inputs=[text, voice, speed], outputs=audio_out)
-
-    gr.Markdown("### Tips  \n"
-                "- **Kokoro** – fast, high-quality English TTS  \n"
-                "Audio is returned as 24 kHz WAV.")
+with gr.Blocks(css=css, title="Text2Audio - Kokoro & NeuTTS-Air") as demo:
+    gr.Markdown("# 🎙️ Text-to-Audio Generation")
+    gr.Markdown("Choose between **Kokoro TTS** (fast English TTS) or **NeuTTS-Air** (voice cloning with reference audio)")
+    
+    # Model selection
+    model_choice = gr.Radio(
+        choices=["Kokoro TTS", "NeuTTS-Air"],
+        value="Kokoro TTS",
+        label="Select TTS Engine",
+        interactive=True
+    )
+    
+    # Kokoro TTS Interface
+    with gr.Group(visible=True) as kokoro_group:
+        gr.Markdown("### 🎧 Kokoro TTS Settings")
+        with gr.Row():
+            with gr.Column():
+                kokoro_voice = gr.Dropdown(
+                    label="Voice",
+                    choices=['af_heart', 'af_sky', 'af_mist', 'af_dusk'],
+                    value='af_heart'
+                )
+                kokoro_speed = gr.Slider(0.5, 2.0, 1.0, step=0.1, label="Speed")
+            
+            with gr.Column(scale=3):
+                kokoro_text = gr.Textbox(
+                    label="Text to speak",
+                    placeholder="Type or paste text here…",
+                    lines=6,
+                    max_lines=12
+                )
+                kokoro_btn = gr.Button("🎧 Synthesise with Kokoro", variant="primary")
+                kokoro_audio_out = gr.Audio(label="Generated speech", type="filepath")
+        
+        gr.Markdown("**Kokoro** – fast, high-quality English TTS. Audio is returned as 24 kHz WAV.")
+    
+    # NeuTTS-Air Interface
+    with gr.Group(visible=False) as neutts_group:
+        gr.Markdown("### ☁️ NeuTTS-Air Settings")
+        neutts_ref_text = gr.Textbox(
+            label="Reference Text",
+            value=DEFAULT_REF_TEXT,
+            lines=3
+        )
+        neutts_ref_audio = gr.Audio(
+            type="filepath",
+            label="Reference Audio",
+            value=DEFAULT_REF_PATH if os.path.exists(DEFAULT_REF_PATH) else None
+        )
+        neutts_gen_text = gr.Textbox(
+            label="Text to Generate",
+            value=DEFAULT_GEN_TEXT,
+            lines=3
+        )
+        neutts_btn = gr.Button("☁️ Generate with NeuTTS-Air", variant="primary")
+        neutts_audio_out = gr.Audio(label="Generated Speech", type="numpy")
+        
+        gr.Markdown("**NeuTTS-Air** – Upload a reference audio sample, provide the reference text, and enter new text to synthesize.")
+    
+    # Event handlers
+    def toggle_interface(choice):
+        if choice == "Kokoro TTS":
+            return gr.update(visible=True), gr.update(visible=False)
+        else:
+            return gr.update(visible=False), gr.update(visible=True)
+    
+    model_choice.change(
+        fn=toggle_interface,
+        inputs=[model_choice],
+        outputs=[kokoro_group, neutts_group]
+    )
+    
+    kokoro_btn.click(
+        kokoro_infer,
+        inputs=[kokoro_text, kokoro_voice, kokoro_speed],
+        outputs=kokoro_audio_out
+    )
+    
+    neutts_btn.click(
+        neutts_infer,
+        inputs=[neutts_ref_text, neutts_ref_audio, neutts_gen_text],
+        outputs=neutts_audio_out
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(allowed_paths=[SAMPLES_PATH] if os.path.exists(SAMPLES_PATH) else None, mcp_server=True, inbrowser=True)