Update app.py

app.py

@@ -1,38 +1,116 @@
-# app.py
-import os, gradio as gr, spaces
-
-# ---- deactivate deepspeed CUDA build BEFORE anything else imports it
-os.environ["DS_ACCELERATOR"] = "cpu"   # force deepspeed to CPU backend
-os.environ["DS_BUILD_OPS"]   = "0"     # skip compiling custom ops
-# ------------------------------------
-
-from huggingface_hub import snapshot_download
-from indextts.infer import IndexTTS
-
-model_dir = snapshot_download("IndexTeam/IndexTTS-1.5",
-                              local_dir="checkpoints",
-                              local_dir_use_symlinks=False)
-cfg_path = os.path.join(model_dir, "config.yaml")
-
-tts = None
-def load():
-    global tts
-    if tts is None:
-        tts = IndexTTS(model_dir=model_dir, cfg_path=cfg_path)  # now imports deepspeed safely
-
-@spaces.GPU
-def synth(ref_wav, prompt):
-    load()
-    out = "out.wav"
-    tts.infer(ref_wav, prompt, out)
-    return out
+import torch
+import torchaudio
+import gradio as gr
+import spaces
+from functools import lru_cache
+from zonos.model import Zonos
+from zonos.conditioning import make_cond_dict, supported_language_codes
+
+MODEL_ID = "Zyphra/Zonos-v0.1-transformer"
+
+@lru_cache(maxsize=1)
+def load_model():
+    return Zonos.from_pretrained(MODEL_ID, device="cuda").eval().requires_grad_(False)
+
+@spaces.GPU(duration=120)
+def tts(
+    text,
+    language,
+    ref_audio,
+    fmax,
+    pitch_std,
+    speaking_rate,
+    guidance_scale,
+    seed,
+    happiness,
+    sadness,
+    disgust,
+    fear,
+    surprise,
+    anger,
+    other,
+    neutral,
+):
+    model = load_model()
+    if seed == 0:
+        seed = int(torch.randint(0, 2**32 - 1, (1,)).item())
+    torch.manual_seed(seed)
+
+    speaker = None
+    if ref_audio is not None:
+        wav, sr = torchaudio.load(ref_audio)
+        speaker = model.make_speaker_embedding(wav, sr).to("cuda", dtype=torch.bfloat16)
+
+    emotion = torch.tensor(
+        [happiness, sadness, disgust, fear, surprise, anger, other, neutral], device="cuda"
+    )
+    cond = make_cond_dict(
+        text=text,
+        language=language,
+        speaker=speaker,
+        emotion=emotion,
+        fmax=float(fmax),
+        pitch_std=float(pitch_std),
+        speaking_rate=float(speaking_rate),
+        device="cuda",
+    )
+    conditioning = model.prepare_conditioning(cond)
+    codes = model.generate(conditioning, cfg_scale=float(guidance_scale), max_new_tokens=2590)
+    wav_out = model.autoencoder.decode(codes).cpu().squeeze()
+    return (model.autoencoder.sampling_rate, wav_out.numpy()), seed
+
 
 with gr.Blocks() as demo:
-    gr.Markdown("# index-tts 1.5 – zero-gpu space")
-    txt = gr.Textbox(label="text prompt")
-    ref = gr.Audio(label="reference voice", type="filepath")
-    gen = gr.Audio(label="generated speech", type="filepath")
-    gr.Button("generate").click(synth, [ref, txt], gen)
-
-demo.queue()
-demo.launch(show_api=False, ssr_mode=False)
+    gr.Markdown("## zonos v0.1 tts")
+    text = gr.Textbox(label="text to synthesize", lines=3)
+    language = gr.Dropdown(
+        choices=supported_language_codes, value="en-us", label="language"
+    )
+    ref_audio = gr.Audio(label="reference audio (zeroshot tts)", type="filepath")
+
+    fmax = gr.Slider(0, 24000, 24000, step=1, label="frequency max (Hz)")
+    pitch_std = gr.Slider(0.0, 300.0, 45.0, step=1, label="pitch variation")
+    speaking_rate = gr.Slider(5.0, 30.0, 15.0, step=0.5, label="speaking rate")
+
+    guidance_scale = gr.Slider(1.0, 5.0, 2.0, step=0.1, label="guidance scale")
+    seed = gr.Number(value=0, label="seed (0 = random)", precision=0)
+
+    gr.Markdown("### emotion settings")
+    happiness = gr.Slider(0, 1, 0.0, step=0.01, label="happiness")
+    sadness = gr.Slider(0, 1, 0.0, step=0.01, label="sadness")
+    disgust = gr.Slider(0, 1, 0.0, step=0.01, label="disgust")
+    fear = gr.Slider(0, 1, 0.0, step=0.01, label="fear")
+    surprise = gr.Slider(0, 1, 0.0, step=0.01, label="surprise")
+    anger = gr.Slider(0, 1, 0.0, step=0.01, label="anger")
+    other = gr.Slider(0, 1, 0.0, step=0.01, label="other")
+    neutral = gr.Slider(0, 1, 1.0, step=0.01, label="neutral")
+
+    btn = gr.Button("synthesize")
+    out_audio = gr.Audio(label="output")
+    out_seed = gr.Number(label="used seed", interactive=False)
+
+    btn.click(
+        tts,
+        [
+            text,
+            language,
+            ref_audio,
+            fmax,
+            pitch_std,
+            speaking_rate,
+            guidance_scale,
+            seed,
+            happiness,
+            sadness,
+            disgust,
+            fear,
+            surprise,
+            anger,
+            other,
+            neutral,
+        ],
+        [out_audio, out_seed],
+    )
+
+if __name__ == "__main__":
+    demo.launch()