Update app.py

app.py

@@ -1,52 +1,72 @@
-import os, shlex, subprocess, torch, numpy as np, spaces, gradio as gr, torchaudio
+# app.py
+import os, shlex, subprocess, torch, numpy as np, gradio as gr, torchaudio, spaces
 from zonos.model import Zonos
 from zonos.conditioning import make_cond_dict, supported_language_codes
 
-subprocess.run(shlex.split("pip install flash-attn --no-build-isolation"), env=os.environ | {"FLASH_ATTENTION_SKIP_CUDA_BUILD":"TRUE"}, check=True)
-subprocess.run(shlex.split("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"), check=True)
-subprocess.run(shlex.split("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"), check=True)
+# ── 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])
 
-os.environ["TORCH_COMPILE_DISABLE"]="1"
-os.environ["TORCHINDUCTOR_DISABLE"]="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"
-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}
+device = "cuda"  # zerogpu maps this transparently
+MODEL_NAME = "Zyphra/Zonos-v0.1-transformer"      # hybrid commented out for now
 
-def _speaker_embed(model,audio):
+_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 model.make_speaker_embedding(wav,sr)
+    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(model_choice,text,language,speaker_audio,
-        e1,e2,e3,e4,e5,e6,e7,e8,
-        speaking_rate,pitch_std):
-    m=MODELS[model_choice]
-    speaker=_speaker_embed(m,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)
+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
+    )
     with torch.no_grad():
-        wav=m.autoencoder.decode(m.generate(m.prepare_conditioning(cond)))[0].cpu()
-    return (m.autoencoder.sampling_rate,wav.numpy())
+        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
 
-langs=supported_language_codes
 with gr.Blocks() as demo:
-    mc=gr.Dropdown(MODEL_NAMES,value=MODEL_NAMES[0],label="model")
-    txt=gr.Textbox(label="text")
-    lang=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","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(fn=tts,
-        inputs=[mc,txt,lang,spk,*emos,rate,pitch],
-        outputs=out)
-if __name__=="__main__": demo.launch()
+    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()