Update app.py

app.py

@@ -1,124 +1,29 @@
 """
-🎙️ Multi-Engine TTS – Zero-GPU edition
-Kokoro │ Veena │ pyttsx3  (fallback)
+🎙️ Kokoro-TTS-only demo – Zero-GPU edition
 Routes every synthesis to an idle A100.
 """
 
 import os, tempfile, subprocess, numpy as np
 import gradio as gr
+import spaces           # Zero-GPU decorator
 import soundfile as sf
-import spaces                               # << Zero-GPU helper
 
 # ------------------------------------------------------------------
-# 1.  Engine availability flags
-# ------------------------------------------------------------------
-KOKORO_OK = False
-VEENA_OK  = False
-PYT_OK    = False
-
-try:
-    from kokoro import KPipeline
-    KOKORO_OK = True
-except Exception as e:
-    print("Kokoro unavailable:", e)
-
-try:
-    import torch, transformers, snac
-    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-    from snac import SNAC
-    VEENA_OK = True
-except Exception as e:
-    print("Veena deps unavailable:", e)
-
-try:
-    import pyttsx3
-    PYT_OK = True
-except Exception as e:
-    print("pyttsx3 unavailable:", e)
-
-# ------------------------------------------------------------------
-# 2.  Lazy model loader (runs once per GPU worker)
+# 1.  Lazy Kokoro loader (runs once per GPU worker)
 # ------------------------------------------------------------------
 kokoro_pipe = None
-veena_model = None
-veena_tok   = None
-veena_snac  = None
 
 def load_kokoro():
     global kokoro_pipe
-    if kokoro_pipe is None and KOKORO_OK:
+    if kokoro_pipe is None:
+        from kokoro import KPipeline
         kokoro_pipe = KPipeline(lang_code='a')
     return kokoro_pipe
 
-def load_veena():
-    global veena_model, veena_tok, veena_snac
-    if veena_model is None and VEENA_OK:
-        bnb = BitsAndBytesConfig(load_in_4bit=True,
-                                 bnb_4bit_quant_type="nf4",
-                                 bnb_4bit_compute_dtype=torch.bfloat16)
-        veena_model = AutoModelForCausalLM.from_pretrained(
-            "maya-research/veena-tts",
-            quantization_config=bnb,
-            device_map="auto",
-            trust_remote_code=True)
-        veena_tok   = AutoTokenizer.from_pretrained("maya-research/veena-tts",
-                                                    trust_remote_code=True)
-        veena_snac  = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval()
-        if torch.cuda.is_available():
-            veena_snac = veena_snac.cuda()
-    return veena_model
-
 # ------------------------------------------------------------------
-# 3.  Generation helpers (CPU→GPU off-load)
+# 2.  Generation helper
 # ------------------------------------------------------------------
-AUDIO_CODE_BASE_OFFSET = 128266
-START_OF_SPEECH_TOKEN  = 128257
-END_OF_SPEECH_TOKEN    = 128258
-START_OF_HUMAN_TOKEN   = 128259
-END_OF_HUMAN_TOKEN     = 128260
-START_OF_AI_TOKEN      = 128261
-END_OF_AI_TOKEN        = 128262
-
-def decode_snac(tokens):
-    if len(tokens) % 7:
-        return None
-    codes = [[] for _ in range(3)]
-    offsets = [AUDIO_CODE_BASE_OFFSET + i*4096 for i in range(7)]
-    for i in range(0, len(tokens), 7):
-        codes[0].append(tokens[i]   - offsets[0])
-        codes[1].extend([tokens[i+1]-offsets[1], tokens[i+4]-offsets[4]])
-        codes[2].extend([tokens[i+2]-offsets[2], tokens[i+3]-offsets[3],
-                         tokens[i+5]-offsets[5], tokens[i+6]-offsets[6]])
-    device = veena_snac.device
-    hierarchical = [torch.tensor(c, dtype=torch.int32, device=device).unsqueeze(0)
-                    for c in codes]
-    with torch.no_grad():
-        wav = veena_snac.decode(hierarchical).squeeze().clamp(-1,1).cpu().numpy()
-    return wav
-
-def tts_veena(text, speaker, temperature, top_p):
-    load_veena()
-    prompt = f"<spk_{speaker}> {text}"
-    tok  = veena_tok.encode(prompt, add_special_tokens=False)
-    input_ids = [START_OF_HUMAN_TOKEN] + tok + [END_OF_HUMAN_TOKEN,
-                  START_OF_AI_TOKEN, START_OF_SPEECH_TOKEN]
-    input_ids = torch.tensor([input_ids], device=veena_model.device)
-    max_new = min(int(len(text)*1.3)*7 + 21, 700)
-    out = veena_model.generate(
-        input_ids,
-        max_new_tokens=max_new,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-        repetition_penalty=1.05,
-        pad_token_id=veena_tok.pad_token_id,
-        eos_token_id=[END_OF_SPEECH_TOKEN, END_OF_AI_TOKEN])
-    gen = out[0, len(input_ids[0]):].tolist()
-    snac_toks = [t for t in gen if AUDIO_CODE_BASE_OFFSET <= t < AUDIO_CODE_BASE_OFFSET+7*4096]
-    if not snac_toks:
-        raise RuntimeError("No audio tokens produced")
-    return decode_snac(snac_toks)
-
+@spaces.GPU
 def tts_kokoro(text, voice, speed):
     pipe = load_kokoro()
     generator = pipe(text, voice=voice, speed=speed)
@@ -126,99 +31,49 @@ def tts_kokoro(text, voice, speed):
         return audio
     raise RuntimeError("Kokoro generation failed")
 
-def tts_pyttsx3(text, rate, volume):
-    engine = pyttsx3.init()
-    engine.setProperty('rate', rate)
-    engine.setProperty('volume', volume)
-    fd, path = tempfile.mkstemp(suffix='.wav')
-    os.close(fd)
-    engine.save_to_file(text, path)
-    engine.runAndWait()
-    wav, sr = sf.read(path)
-    os.remove(path)
-    return wav
-
 # ------------------------------------------------------------------
-# 4.  ZERO-GPU ENTRY POINT  (decorated)
+# 3.  Zero-GPU entry point
 # ------------------------------------------------------------------
 @spaces.GPU
-def synthesise(text, engine, voice, speed, speaker, temperature, top_p, rate, vol):
+def synthesise(text, voice, speed):
     if not text.strip():
         raise gr.Error("Please enter some text.")
-    if engine == "kokoro" and KOKORO_OK:
-        wav = tts_kokoro(text, voice=voice, speed=speed)
-    elif engine == "veena" and VEENA_OK:
-        wav = tts_veena(text, speaker=speaker, temperature=temperature, top_p=top_p)
-    elif engine == "pyttsx3" and PYT_OK:
-        wav = tts_pyttsx3(text, rate=rate, volume=vol)
-    else:
-        raise gr.Error(f"{engine} is not available on this Space.")
+    wav = tts_kokoro(text, voice=voice, speed=speed)
     fd, tmp = tempfile.mkstemp(suffix='.wav')
     os.close(fd)
     sf.write(tmp, wav, 24000)
     return tmp
 
 # ------------------------------------------------------------------
-# 5.  Gradio UI  (unchanged visuals)
+# 4.  Gradio UI
 # ------------------------------------------------------------------
-css = """footer {visibility: hidden} #col-left {max-width: 320px}"""
+css = """footer {visibility: hidden}"""
 
-with gr.Blocks(css=css, title="Multi-Engine TTS – Zero-GPU") as demo:
-    gr.Markdown("## 🎙️ Multi-Engine TTS Demo – Zero-GPU  \n*Kokoro ‑ Veena ‑ pyttsx3*")
+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(elem_id="col-left"):
-            engine = gr.Radio(label="Engine",
-                              choices=[e for e in ["kokoro","veena","pyttsx3"]
-                                       if globals().get({"pyttsx3":"PYT_OK"}.get(e,e.upper()+"_OK"), False)],
-                              value="kokoro" if KOKORO_OK else
-                                    "veena"  if VEENA_OK  else "pyttsx3")
-
-            with gr.Group(visible=KOKORO_OK) as kokoro_box:
-                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.Group(visible=VEENA_OK) as veena_box:
-                speaker = gr.Dropdown(label="Speaker",
-                                      choices=['kavya','agastya','maitri','vinaya'],
-                                      value='kavya')
-                temperature = gr.Slider(0.1, 1.0, 0.4, step=0.05, label="Temperature")
-                top_p      = gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top-p")
-
-            with gr.Group(visible=PYT_OK) as pyttsx3_box:
-                rate = gr.Slider(50, 300, 180, step=5, label="Words / min")
-                vol  = gr.Slider(0.0, 1.0, 1.0, step=0.05, label="Volume")
+        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")
+            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")
 
-    # show/hide panels
-    def switch_panel(e):
-        return (gr.update(visible=e=="kokoro"),
-                gr.update(visible=e=="veena"),
-                gr.update(visible=e=="pyttsx3"))
-    engine.change(switch_panel, inputs=engine,
-                  outputs=[kokoro_box, veena_box, pyttsx3_box])
-
-    # binding
-    btn.click(synthesise,
-              inputs=[text, engine, voice, speed, speaker,
-                      temperature, top_p, rate, vol],
-              outputs=audio_out)
+    btn.click(synthesise, inputs=[text, voice, speed], outputs=audio_out)
 
     gr.Markdown("### Tips  \n"
-                "- **Kokoro** – fastest, good quality English  \n"
-                "- **Veena** – multilingual, GPU-friendly (4-bit)  \n"
-                "- **pyttsx3** – offline fallback, any language  \n"
+                "- **Kokoro** – fast, high-quality English TTS  \n"
                 "Audio is returned as 24 kHz WAV.")
 
-# ------------------------------------------------------------------
-# 6.  Launch
-# ------------------------------------------------------------------
 demo.launch()