Update app.py

app.py

@@ -1,71 +1,46 @@
-import os, torch, torchaudio, gradio as gr
-
+# Imports
+import gradio as gr
 import spaces
+import os
+import torch
+import torchaudio
 
 from zonos.model import Zonos
 from zonos.conditioning import make_cond_dict, supported_language_codes
 
-device = "cuda"
-MODEL_NAME = "Zyphra/Zonos-v0.1-transformer"
-MODEL = Zonos.from_pretrained(MODEL_NAME, device=device)
+# Variables
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
-# def _patch_cuda_props():
-#     if torch.cuda.is_available():
-#         for i in range(torch.cuda.device_count()):
-#             p = torch.cuda.get_device_properties(i)
-#             if not hasattr(p, "regs_per_multiprocessor"):
-#                 setattr(p, "regs_per_multiprocessor", 65536)
-#             if not hasattr(p, "max_threads_per_multi_processor"):
-#                 setattr(p, "max_threads_per_multi_processor", 2048)
+device = "cuda"
 
+REPO = "Zyphra/Zonos-v0.1-transformer"
+model = Zonos.from_pretrained(REPO, device=device)
 
-# _patch_cuda_props()
+# Functions
+def patch_cuda():
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            p = torch.cuda.get_device_properties(i)
+            if not hasattr(p, "regs_per_multiprocessor"):
+                setattr(p, "regs_per_multiprocessor", 65536)
+            if not hasattr(p, "max_threads_per_multi_processor"):
+                setattr(p, "max_threads_per_multi_processor", 2048)
 
 @spaces.GPU
-def generate_audio(
-    text,
-    language,
-    speaker_audio,
-    e1,
-    e2,
-    e3,
-    e4,
-    e5,
-    e6,
-    e7,
-    e8,
-    clarity,
-    fmax,
-    pitch_std,
-    speaking_rate,
-    dnsmos_ovrl,
-    cfg_scale,
-    min_p,
-    steps,
-    seed,
-    randomize_seed,
-    progress=gr.Progress(),
-):
-    if randomize_seed:
-        seed = torch.randint(0, 2**32 - 1, (1,)).item()
-    torch.manual_seed(int(seed))
+def generate(input, language, speaker_audio, emotion_happy, emotion_sad, emotion_disgust, emotion_fear, emotion_surprise, emotion_anger, emotion_other, emotion_neutral, clarity, fmax, pitch_std, speaking_rate, dnsmos_ovrl, cfg_scale, min_p, steps, seed, randomize_seed):
+    if randomize_seed: seed = int(time.time())
+    torch.manual_seed(seed)
 
     speaker_embedding = None
     if speaker_audio is not None:
         wav, sr = torchaudio.load(speaker_audio)
-        speaker_embedding = (
-            MODEL.make_speaker_embedding(wav, sr).to(device, dtype=torch.bfloat16)
-        )
+        speaker_embedding = (model.make_speaker_embedding(wav, sr).to(device, dtype=torch.bfloat16))
 
-    emotion_tensor = torch.tensor(
-        [e1, e2, e3, e4, e5, e6, e7, e8], device=device, dtype=torch.float32
-    )
-    vq_tensor = torch.tensor([clarity] * 8, device=device, dtype=torch.float32).unsqueeze(
-        0
-    )
+    emotion_tensor = torch.tensor([emotion_happy, emotion_sad, emotion_disgust, emotion_fear, emotion_surprise, emotion_anger, emotion_other, emotion_neutral], device=device, dtype=torch.bfloat16)
+    vq_tensor = torch.tensor([clarity] * 8, device=device, dtype=torch.bfloat16).unsqueeze(0)
 
     cond_dict = make_cond_dict(
-        text=text,
+        text=input,
         language=language,
         speaker=speaker_embedding,
         emotion=emotion_tensor,
@@ -76,15 +51,10 @@ def generate_audio(
         dnsmos_ovrl=float(dnsmos_ovrl),
         device=device,
     )
-    conditioning = MODEL.prepare_conditioning(cond_dict)
-
-    estimated_total_steps = int(steps)
+    
+    conditioning = model.prepare_conditioning(cond_dict)
 
-    def cb(_, step, __):
-        progress((step, estimated_total_steps))
-        return True
-
-    codes = MODEL.generate(
+    codes = model.generate(
         prefix_conditioning=conditioning,
         max_new_tokens=int(steps),
         cfg_scale=float(cfg_scale),
@@ -93,78 +63,48 @@ def generate_audio(
         callback=cb,
     )
 
-    wav_out = MODEL.autoencoder.decode(codes).cpu().detach()
-    sr_out = MODEL.autoencoder.sampling_rate
-    if wav_out.dim() == 2 and wav_out.size(0) > 1:
-        wav_out = wav_out[0:1, :]
-    return (sr_out, wav_out.squeeze().numpy()), seed
-
-
-def build_interface():
-    with gr.Blocks() as demo:
-        text = gr.Textbox(label="text", value="hello, world!", lines=4, max_length=500)
-        language = gr.Dropdown(choices=supported_language_codes, value="en-us", label="language")
-        speaker_audio = gr.Audio(label="voice reference", type="filepath")
-
-        clarity_slider = gr.Slider(0.5, 0.8, 0.8, 0.01, label="clarity")
-        steps_slider = gr.Slider(1, 3000, 300, 1, label="steps")
-
-        dnsmos_slider = gr.Slider(1.0, 5.0, 5.0, 0.1, label="quality")
-        fmax_slider = gr.Slider(0, 24000, 24000, 1, label="fmax")
-        pitch_std_slider = gr.Slider(0.0, 300.0, 30.0, 1, label="pitch std")
-        speaking_rate_slider = gr.Slider(5.0, 30.0, 15.0, 0.1, label="rate")
-
-        cfg_scale_slider = gr.Slider(1.0, 5.0, 2.5, 0.1, label="guidance")
-        min_p_slider = gr.Slider(0.0, 1.0, 0.05, 0.01, label="min p")
-
-        with gr.Row():
-            e1 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="happy")
-            e2 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="sad")
-            e3 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="disgust")
-            e4 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="fear")
-
-        with gr.Row():
-            e5 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="surprise")
-            e6 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="anger")
-            e7 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="other")
-            e8 = gr.Slider(0.0, 1.0, 1.0, 0.01, label="neutral")
-
-        seed_number = gr.Number(label="seed", value=420, precision=0)
-        randomize_seed_toggle = gr.Checkbox(label="randomize seed", value=True)
-
-        generate_button = gr.Button("generate")
-        output_audio = gr.Audio(label="output", type="numpy", autoplay=True)
-
-        generate_button.click(
-            fn=generate_audio,
-            inputs=[
-                text,
-                language,
-                speaker_audio,
-                e1,
-                e2,
-                e3,
-                e4,
-                e5,
-                e6,
-                e7,
-                e8,
-                clarity_slider,
-                fmax_slider,
-                pitch_std_slider,
-                speaking_rate_slider,
-                dnsmos_slider,
-                cfg_scale_slider,
-                min_p_slider,
-                steps_slider,
-                seed_number,
-                randomize_seed_toggle,
-            ],
-            outputs=[output_audio, seed_number],
-        )
+    wav_out = model.autoencoder.decode(codes).cpu().detach()
+    sr_out = model.autoencoder.sampling_rate
+    
+    if wav_out.dim() == 2 and wav_out.size(0) > 1: wav_out = wav_out[0:1, :]
         
-    return demo
+    return (sr_out, wav_out.squeeze().numpy())
+
+# Initialize
+patch_cuda()
+
+with gr.Blocks() as main:
+    text = gr.Textbox(label="text", value="hello, world!")
+    language = gr.Dropdown(choices=supported_language_codes, value="en-us", label="language")
+    speaker_audio = gr.Audio(label="voice reference", type="filepath")
+
+    clarity_slider = gr.Slider(0.5, 0.8, 0.8, 0.01, label="clarity")
+    steps_slider = gr.Slider(1, 3000, 300, 1, label="steps")
+
+    dnsmos_slider = gr.Slider(1.0, 5.0, 5.0, 0.1, label="quality")
+    fmax_slider = gr.Slider(0, 24000, 24000, 1, label="fmax")
+    pitch_std_slider = gr.Slider(0.0, 300.0, 30.0, 1, label="pitch std")
+    speaking_rate_slider = gr.Slider(5.0, 30.0, 15.0, 0.1, label="rate")
+
+    cfg_scale_slider = gr.Slider(1.0, 5.0, 2.5, 0.1, label="guidance")
+    min_p_slider = gr.Slider(0.0, 1.0, 0.05, 0.01, label="min p")
+
+    with gr.Row():
+        e1 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="happy")
+        e2 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="sad")
+        e3 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="disgust")
+        e4 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="fear")
+        e5 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="surprise")
+        e6 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="anger")
+        e7 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="other")
+        e8 = gr.Slider(0.0, 1.0, 1.0, 0.01, label="neutral")
+
+    seed_number = gr.Number(label="seed", value=42, precision=0)
+    randomize_seed_toggle = gr.Checkbox(label="randomize seed", value=True)
+
+    generate_button = gr.Button("generate")
+    output_audio = gr.Audio(label="output", type="numpy", autoplay=True)
 
+    generate_button.click(fn=generate, inputs=[text, language, speaker_audio, e1, e2, e3, e4, e5, e6, e7, e8, clarity_slider, fmax_slider, pitch_std_slider, speaking_rate_slider, dnsmos_slider, cfg_scale_slider, min_p_slider, steps_slider, seed_number, randomize_seed_toggle], outputs=output_audio)
 
-if __name__ == "__main__":
-    build_interface().launch()
+main.launch()