Update app.py

app.py

@@ -1,72 +1,58 @@
-# app.py
-import os, shlex, subprocess, torch, numpy as np, gradio as gr, torchaudio, spaces
+import os, shlex, subprocess, torch
+
+# extra wheels (safe to skip if they fail)
+for cmd, env in [
+    ("pip install flash-attn --no-build-isolation", {"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}),
+    ("pip install https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl", {}),
+    ("pip install https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.0.post8/causal_conv1d-1.5.0.post8+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl", {}),
+]:
+    try: subprocess.run(shlex.split(cmd), env=os.environ | env, check=True)
+    except subprocess.CalledProcessError: pass
+
+# hard-nuke torch.compile everywhere
+os.environ["TORCH_COMPILE_DISABLE"]="1"
+os.environ["TORCHINDUCTOR_DISABLE"]="1"
+torch._dynamo.disable()
+torch.compile=lambda fn,*a,**k:fn
+
+import torchaudio, gradio as gr, spaces, numpy as np
 from zonos.model import Zonos
 from zonos.conditioning import make_cond_dict, supported_language_codes
 
-# ── optional perf wheels (safe to ignore if they fail) ───────────────────────────
-cmds = [
-    "pip install flash-attn --no-build-isolation",
-    "pip install https://github.com/state-spaces/mamba/releases/download/v2.2.4/"
-    "mamba_ssm-2.2.4+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl",
-    "pip install https://github.com/Dao-AILab/causal-conv1d/releases/download/"
-    "v1.5.0.post8/causal_conv1d-1.5.0.post8+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl",
-]
-for c in cmds:
-    try:
-        subprocess.run(shlex.split(c), check=True)
-    except subprocess.CalledProcessError:
-        print("wheel skipped:", c.split()[2 if c.startswith('pip') else -1])
-
-# ── disable torch.compile: zerogpu lacks full cuda props ────────────────────────
-os.environ["TORCH_COMPILE_DISABLE"] = "1"
-os.environ["TORCHINDUCTOR_DISABLE"] = "1"
-import torch._dynamo; torch._dynamo.disable()
-
-device = "cuda"  # zerogpu maps this transparently
-MODEL_NAME = "Zyphra/Zonos-v0.1-transformer"      # hybrid commented out for now
-
-_cached_model: Zonos | None = None
-def get_model() -> Zonos:
-    global _cached_model
-    if _cached_model is None:
-        _cached_model = Zonos.from_pretrained(MODEL_NAME, device=device).eval()
-    return _cached_model
-
-def _speaker_embed(audio):
-    if audio is None: return None
-    sr, wav = audio
-    if wav.dtype.kind in "iu": wav = wav.astype(np.float32) / np.iinfo(wav.dtype).max
-    wav = torch.from_numpy(wav).unsqueeze(0)
-    return get_model().make_speaker_embedding(wav, sr)
-
-@spaces.GPU
-def tts(text, language, speaker_audio,
-        e1,e2,e3,e4,e5,e6,e7,e8, speaking_rate, pitch_std):
-    m = get_model()
-    speaker = _speaker_embed(speaker_audio)
-    emotion = [e1,e2,e3,e4,e5,e6,e7,e8]
-    cond = make_cond_dict(
-        text=text, language=language, speaker=speaker, emotion=emotion,
-        speaking_rate=float(speaking_rate), pitch_std=float(pitch_std), device=device
-    )
+device="cuda"
+MODEL_NAMES=["Zyphra/Zonos-v0.1-transformer","Zyphra/Zonos-v0.1-hybrid"]
+MODELS={n:Zonos.from_pretrained(n,device=device).eval() for n in MODEL_NAMES}
+
+def _spk(model,aud):
+    if aud is None: return None
+    sr,wav=aud
+    if wav.dtype.kind in "iu": wav=wav.astype(np.float32)/np.iinfo(wav.dtype).max
+    return model.make_speaker_embedding(torch.from_numpy(wav).unsqueeze(0),sr)
+
+@spaces.GPU(duration=120)
+def tts(m,text,lang,speaker,
+        h,sad,disg,fear,sur,ang,oth,neu,
+        speak,pitch):
+    model=MODELS[m]
+    emotion=[h,sad,disg,fear,sur,ang,oth,neu]
+    cond=make_cond_dict(text=text,language=lang,speaker=_spk(model,speaker),
+                        emotion=emotion,speaking_rate=float(speak),
+                        pitch_std=float(pitch),device=device)
     with torch.no_grad():
-        codes = m.generate(m.prepare_conditioning(cond))
-        wav  = m.autoencoder.decode(codes)[0].cpu()
-    return (m.autoencoder.sampling_rate, wav.numpy())
-
-langs = supported_language_codes            # from the library itself
+        codes=model.generate(model.prepare_conditioning(cond))
+        wav=model.autoencoder.decode(codes)[0].cpu()
+    return (model.autoencoder.sampling_rate,wav.numpy())
 
+langs=supported_language_codes
 with gr.Blocks() as demo:
-    txt  = gr.Textbox(label="text")
-    lng  = gr.Dropdown(langs, value="en-us", label="language")
-    aud  = gr.Audio(type="numpy", label="speaker ref (optional)")
-    emos = [gr.Slider(0,1,0.3 if i==0 else 0.0,0.05,label=l)
-            for i,l in enumerate(["happiness","sadness","disgust","fear",
-                                  "surprise","anger","other","neutral"])]
-    rate = gr.Slider(0,40,15,1,label="speaking_rate")
-    pitch= gr.Slider(0,400,20,1,label="pitch_std")
-    out  = gr.Audio(label="output")
-    gr.Button("generate").click(tts,[txt,lng,aud,*emos,rate,pitch],out)
-
-if __name__ == "__main__":
-    demo.launch()
+    mc=gr.Dropdown(MODEL_NAMES,value=MODEL_NAMES[0],label="model")
+    txt=gr.Textbox(label="text")
+    lng=gr.Dropdown(langs,value="en-us",label="language")
+    spk=gr.Audio(type="numpy",label="speaker ref")
+    emos=[gr.Slider(0,1,0.3 if i==0 else 0.0,0.05,label=l) for i,l in
+          enumerate(["happiness","sad","disgust","fear","surprise","anger","other","neutral"])]
+    rate=gr.Slider(0,40,15,1,label="speaking_rate")
+    pit=gr.Slider(0,400,20,1,label="pitch_std")
+    out=gr.Audio(label="output")
+    gr.Button("generate").click(tts,[mc,txt,lng,spk,*emos,rate,pit],out)
+if __name__=="__main__": demo.launch()