Update app.py

app.py

@@ -1,58 +1,57 @@
-import os, shlex, subprocess, torch
+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
 
-# extra wheels (safe to skip if they fail)
+# optional speed-up wheels, silently skip on failure
 for cmd, env in [
-    ("pip install flash-attn --no-build-isolation", {"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}),
+    ("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)
+    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"
+os.environ["TORCH_COMPILE_DISABLE"]=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
+torch.compile=lambda f,*a,**k:f
 
 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}
+model=Zonos.from_pretrained("Zyphra/Zonos-v0.1-transformer",device=device).eval()
 
-def _spk(model,aud):
+def _spk(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)
+@spaces.GPU
+def tts(text,lang,speaker,vq,fmax,pitch,rate,dnsmos):
+    cond=make_cond_dict(
+        text=text,
+        language=lang,
+        speaker=_spk(speaker),
+        vqscore_8=torch.tensor([vq]*8,device=device).unsqueeze(0),
+        fmax=float(fmax),
+        pitch_std=float(pitch),
+        speaking_rate=float(rate),
+        dnsmos_ovrl=float(dnsmos),
+        device=device,
+    )
     with torch.no_grad():
-        codes=model.generate(model.prepare_conditioning(cond))
-        wav=model.autoencoder.decode(codes)[0].cpu()
-    return (model.autoencoder.sampling_rate,wav.numpy())
+        wav=model.autoencoder.decode(model.generate(model.prepare_conditioning(cond)))[0].cpu()
+    out=(wav.clip(-1,1)*32767).short().numpy()                  # int16 fix
+    return (model.autoencoder.sampling_rate,out)
 
 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")
     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")
+    spk=gr.Audio(type="numpy",label="speaker ref (optional)")
+    vq =gr.Slider(0.5,0.9,0.78,0.01,label="clarity (vq)")
+    fmx=gr.Slider(8000,24000,24000,100,label="fmax hz")
+    pit=gr.Slider(0,300,20,1,label="pitch std")
+    rte=gr.Slider(5,30,15,0.5,label="speaking rate")
+    dns=gr.Slider(1,5,4,0.1,label="quality target")
     out=gr.Audio(label="output")
-    gr.Button("generate").click(tts,[mc,txt,lng,spk,*emos,rate,pit],out)
+    gr.Button("generate").click(tts,[txt,lng,spk,vq,fmx,pit,rte,dns],out)
 if __name__=="__main__": demo.launch()