Update app.py

app.py

@@ -1,79 +1,52 @@
-import gradio as gr
-import os
-
-os.environ["TORCH_COMPILE_DISABLE"] = "1"      # disable torch.compile globally
-os.environ["TORCHINDUCTOR_DISABLE"] = "1"      # belt-and-suspenders for inductor
-
-import torch, torch._dynamo                    # must come *after* the env vars
-torch._dynamo.config.suppress_errors = True    # swallow any stray compile calls
-torch._dynamo.disable()
+import os, shlex, subprocess, torch, numpy as np, spaces, gradio as gr, torchaudio
+from zonos.model import Zonos
+from zonos.conditioning import make_cond_dict, supported_language_codes
 
-import spaces
-import torch
-import numpy as np
+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)
 
-from zonos.model import Zonos
-from zonos.conditioning import make_cond_dict
+os.environ["TORCH_COMPILE_DISABLE"]="1"
+os.environ["TORCHINDUCTOR_DISABLE"]="1"
+import torch._dynamo; torch._dynamo.disable()
 
-cpu = "cuda"
-model = Zonos.from_pretrained("Zyphra/Zonos-v0.1-transformer", device=cpu)
+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 _speaker_embed(audio):
+def _speaker_embed(model,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
-    else:
-        wav = wav.astype(np.float32)
-    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 model.make_speaker_embedding(wav,sr)
 
 @spaces.GPU
-def tts(
-    text, language, speaker_audio,
-    happy, sad, disgust, fear, surprise, anger, other, neutral,
-    speaking_rate, pitch_std
-):
-    speaker = _speaker_embed(speaker_audio)
-    emotion = [happy, sad, disgust, fear, surprise, anger, other, neutral]
-    model.to("cuda")
+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)
     with torch.no_grad():
-        cond = make_cond_dict(
-            text=text,
-            language=language,
-            speaker=speaker,
-            emotion=emotion,
-            speaking_rate=float(speaking_rate),
-            pitch_std=float(pitch_std),
-        )
-        conditioning = model.prepare_conditioning(cond)
-        codes = model.generate(conditioning)
-        wav = model.autoencoder.decode(codes)[0].cpu()
-    model.to(cpu)
-    torch.cuda.empty_cache()
-    return (model.autoencoder.sampling_rate, wav.numpy())
-
-langs = ["en-us", "fr-fr", "ja", "de-de", "zh"]
-demo = gr.Interface(
-    fn=tts,
-    inputs=[
-        gr.Textbox(label="text"),
-        gr.Dropdown(langs, value="en-us", label="language"),
-        gr.Audio(type="numpy", label="speaker reference (optional)"),
-        gr.Slider(0, 1, 0.3, 0.05, label="happiness"),
-        gr.Slider(0, 1, 0.0, 0.05, label="sadness"),
-        gr.Slider(0, 1, 0.0, 0.05, label="disgust"),
-        gr.Slider(0, 1, 0.0, 0.05, label="fear"),
-        gr.Slider(0, 1, 0.0, 0.05, label="surprise"),
-        gr.Slider(0, 1, 0.0, 0.05, label="anger"),
-        gr.Slider(0, 1, 0.2, 0.05, label="other"),
-        gr.Slider(0, 1, 0.5, 0.05, label="neutral"),
-        gr.Slider(0, 40, 15, 1, label="speaking_rate"),
-        gr.Slider(0, 400, 20, 1, label="pitch_std"),
-    ],
-    outputs=gr.Audio(label="generated speech"),
-    title="zonos-v0.1 zerogpu tts",
-)
-
-if __name__ == "__main__":
-    demo.launch(share=True)
+        wav=m.autoencoder.decode(m.generate(m.prepare_conditioning(cond)))[0].cpu()
+    return (m.autoencoder.sampling_rate,wav.numpy())
+
+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()