Update app.py

app.py

@@ -1,63 +1,96 @@
-# app.py – only transformer, emotion cast to float, guidance ≥1.1
-import os, torch, numpy as np, gradio as gr, torchaudio, soundfile as sf, spaces
+# app.py  ── Zonos TTS (transformer only, minimal UI)
+
+import os, tempfile, torch, numpy as np, gradio as gr, torchaudio, soundfile as sf, spaces
 from zonos.model import Zonos
 from zonos.conditioning import make_cond_dict, supported_language_codes
-import tempfile, soundfile as sf
 
-os.environ["TORCH_COMPILE_DISABLE"]=os.environ["TORCHINDUCTOR_DISABLE"]="1"
+# disable Torch-Inductor (keeps Spaces happy)
+os.environ["TORCH_COMPILE_DISABLE"] = os.environ["TORCHINDUCTOR_DISABLE"] = "1"
 torch._dynamo.disable()
-torch.compile=lambda f,*a,**k:f
+torch.compile = lambda f,*a,**k: f            # no-op
 
-device="cuda"
-model=Zonos.from_pretrained("Zyphra/Zonos-v0.1-transformer",device=device).eval()
+device = "cuda"
+model  = Zonos.from_pretrained("Zyphra/Zonos-v0.1-transformer", device=device).eval()
 
-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)
+# ─────────────────── helpers ────────────────────
+def _speaker_embed(aud):
+    if aud is None:
+        return None
+    sr, wav = aud                                         # gradio returns (sr, np.ndarray)
+    if wav.dtype.kind in "iu":                            # int → float
+        wav = wav.astype(np.float32) / np.iinfo(wav.dtype).max
+    wav_t = torch.from_numpy(wav).unsqueeze(0)            # (1,C,N)
+    return model.make_speaker_embedding(wav_t, sr)
 
+# ─────────────────── inference ───────────────────
 @spaces.GPU
-def tts(text,lang,speaker,
+def tts(text, lang, speaker,
         e1,e2,e3,e4,e5,e6,e7,e8,
-        vq,fmax,pitch,rate,cfg,minp):
-    emotion=torch.tensor([float(x) for x in [e1,e2,e3,e4,e5,e6,e7,e8]],
-                         device=device,dtype=torch.float32)
-    cond=make_cond_dict(
-        text=text,language=lang,speaker=_spk(speaker),
+        vq, fmax, pitch, rate,
+        cfg, minp, tokens):
+
+    emotion = torch.tensor([float(x) for x in [e1,e2,e3,e4,e5,e6,e7,e8]],
+                           device=device, dtype=torch.float32)
+
+    cond = make_cond_dict(
+        text=text,
+        language=lang,
+        speaker=_speaker_embed(speaker),
         emotion=emotion,
-        vqscore_8=torch.tensor([vq]*8,device=device).unsqueeze(0),
-        fmax=float(fmax),pitch_std=float(pitch),
-        speaking_rate=float(rate),device=device)
+        vqscore_8=torch.tensor([vq]*8, device=device).unsqueeze(0),
+        fmax=float(fmax),
+        pitch_std=float(pitch),
+        speaking_rate=float(rate),
+        device=device
+    )
+
     with torch.no_grad():
-        wav=model.autoencoder.decode(
-            model.generate(model.prepare_conditioning(cond),
-                           cfg_scale=float(cfg),sampling_params=dict(min_p=float(minp)))
-        )[0].cpu().clamp_(-1,1).numpy()
-    wav = np.squeeze(wav)                    # drop singleton channel
-    if wav.ndim == 2:                        # (C, N) ➜ (N, C)
-        wav = wav.T
-    
+        codes = model.generate(model.prepare_conditioning(cond),
+                               max_new_tokens=int(tokens),
+                               cfg_scale=float(cfg),
+                               sampling_params=dict(min_p=float(minp)))
+        wav = model.autoencoder.decode(codes)[0]          # (C,N) torch
+        wav = wav.cpu().clamp_(-1,1).numpy()              # → numpy
+
+    # >>> FIX <<<  ensure (N,) or (N, C) for libsndfile
+    wav = np.squeeze(wav)
+    if wav.ndim == 2:             # currently (C,N)
+        wav = wav.T               # → (N,C)
+
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
         sf.write(tmp.name, wav,
                  model.autoencoder.sampling_rate,
                  format="WAV", subtype="PCM_16")
     return tmp.name
 
-langs=supported_language_codes
+# ─────────────────── UI ──────────────────────────
+langs = supported_language_codes
 with gr.Blocks() as demo:
-    t=gr.Textbox(label="text")
-    l=gr.Dropdown(langs,value="en-us",label="language")
-    s=gr.Audio(type="numpy",label="speaker ref (optional)")
-    emos=[gr.Slider(0,1,0.0,0.05,label=n) for n in
-          ["happiness","sadness","disgust","fear","surprise","anger","other","neutral"]]
-    vq=gr.Slider(0.5,0.9,0.78,0.01,label="clarity (vq)")
-    fmx=gr.Slider(0,24000,24000,100,label="fmax (hz)")
-    pit=gr.Slider(0,300,45,1,label="pitch variation")
-    rte=gr.Slider(5,30,15,0.5,label="speaking rate")
-    cfg=gr.Slider(1.1,5,2,0.1,label="guidance scale")  # min 1.1 to avoid assert
-    mp =gr.Slider(0,1,0.15,0.01,label="min-p")
-    out=gr.Audio(type="filepath",label="output")
-    gr.Button("generate").click(tts,[t,l,s,*emos,vq,fmx,pit,rte,cfg,mp],out)
-if __name__=="__main__": demo.launch()
+    text = gr.Textbox(label="Text")
+    lang = gr.Dropdown(langs, value="en-us", label="Language")
+    speaker = gr.Audio(type="numpy", label="Speaker ref (optional)")
+
+    # emotion sliders (all default 0)
+    emotions = []
+    for label in ["happiness","sadness","disgust","fear",
+                  "surprise","anger","other","neutral"]:
+        emotions.append(gr.Slider(0,1,0.0,0.05,label=label))
+
+    vq   = gr.Slider(0.5,0.9,0.78,0.01,label="clarity (vq)")
+    fmax = gr.Slider(0,24000,24000,100,label="fmax (Hz)")
+    pitch= gr.Slider(0,300,45,1,label="pitch variation")
+    rate = gr.Slider(5,30,15,0.5,label="speaking rate")
+    cfg  = gr.Slider(1.1,5,2,0.1,label="guidance scale")
+    minp = gr.Slider(0,1,0.15,0.01,label="min-p")
+    tokens = gr.Slider(0,3000,300,1,label="tokens (≈steps)")
+
+    out = gr.Audio(type="filepath", label="Output")
+    gr.Button("Generate").click(
+        tts,
+        inputs=[text, lang, speaker, *emotions,
+                vq, fmax, pitch, rate, cfg, minp, tokens],
+        outputs=out
+    )
+
+if __name__ == "__main__":
+    demo.launch()