Update app.py

app.py

@@ -1,4 +1,29 @@
 import os
+import sys
+import subprocess
+
+# Emergency flash-attn installation if not found
+try:
+    import flash_attn
+except ImportError:
+    print("flash_attn not found, attempting to install...")
+    try:
+        # Try installing pre-built wheel first (fastest)
+        subprocess.run([
+            sys.executable, "-m", "pip", "install", 
+            "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
+        ], check=True)
+    except:
+        # Fallback: install without CUDA build (slower but more compatible)
+        env = os.environ.copy()
+        env["FLASH_ATTENTION_SKIP_CUDA_BUILD"] = "TRUE"
+        subprocess.run([
+            sys.executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"
+        ], env=env, check=True)
+    
+    # Restart the script after installation
+    os.execv(sys.executable, [sys.executable] + sys.argv)
+
 import spaces
 import torch
 import torchaudio
@@ -8,12 +33,20 @@ from os import getenv
 from zonos.model import Zonos
 from zonos.conditioning import make_cond_dict, supported_language_codes
 
-device = "cuda"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 MODEL_NAMES = ["Zyphra/Zonos-v0.1-transformer", "Zyphra/Zonos-v0.1-hybrid"]
-MODELS = {name: Zonos.from_pretrained(name, device=device) for name in MODEL_NAMES}
-for model in MODELS.values():
-    model.requires_grad_(False).eval()
-
+MODELS = {}
+
+# Load models with error handling
+for name in MODEL_NAMES:
+    try:
+        MODELS[name] = Zonos.from_pretrained(name, device=device)
+        MODELS[name].requires_grad_(False).eval()
+        print(f"Successfully loaded model: {name}")
+    except Exception as e:
+        print(f"Failed to load model {name}: {e}")
+        if not MODELS:  # If no models loaded at all
+            raise
 
 def update_ui(model_choice):
     """
@@ -359,8 +392,8 @@ def build_interface():
                 with gr.Row():
                     with gr.Column(scale=2):
                         model_choice = gr.Dropdown(
-                            choices=MODEL_NAMES,
-                            value="Zyphra/Zonos-v0.1-transformer",
+                            choices=list(MODELS.keys()),
+                            value=list(MODELS.keys())[0] if MODELS else None,
                             label="Zonos Model Type",
                             info="Select the model variant to use.",
                         )
@@ -378,7 +411,7 @@ def build_interface():
                         )
                     with gr.Column(scale=1):
                         prefix_audio = gr.Audio(
-                            value="assets/silence_100ms.wav",
+                            value="assets/silence_100ms.wav" if os.path.exists("assets/silence_100ms.wav") else None,
                             label="Optional Prefix Audio (continue from this audio)",
                             type="filepath",
                         )
@@ -460,90 +493,91 @@ def build_interface():
                 generate_button = gr.Button("Generate Audio", elem_classes="generate-button")
                 output_audio = gr.Audio(label="Generated Audio", type="numpy", autoplay=True, elem_classes="audio-output")
 
-        model_choice.change(
-            fn=update_ui,
-            inputs=[model_choice],
-            outputs=[
-                text,
-                language,
-                speaker_audio,
-                prefix_audio,
-                emotion1,
-                emotion2,
-                emotion3,
-                emotion4,
-                emotion5,
-                emotion6,
-                emotion7,
-                emotion8,
-                vq_single_slider,
-                fmax_slider,
-                pitch_std_slider,
-                speaking_rate_slider,
-                dnsmos_slider,
-                speaker_noised_checkbox,
-                unconditional_keys,
-            ],
-        )
-
-        # On page load, trigger the same UI refresh
-        demo.load(
-            fn=update_ui,
-            inputs=[model_choice],
-            outputs=[
-                text,
-                language,
-                speaker_audio,
-                prefix_audio,
-                emotion1,
-                emotion2,
-                emotion3,
-                emotion4,
-                emotion5,
-                emotion6,
-                emotion7,
-                emotion8,
-                vq_single_slider,
-                fmax_slider,
-                pitch_std_slider,
-                speaking_rate_slider,
-                dnsmos_slider,
-                speaker_noised_checkbox,
-                unconditional_keys,
-            ],
-        )
-
-        # Generate audio on button click
-        generate_button.click(
-            fn=generate_audio,
-            inputs=[
-                model_choice,
-                text,
-                language,
-                speaker_audio,
-                prefix_audio,
-                emotion1,
-                emotion2,
-                emotion3,
-                emotion4,
-                emotion5,
-                emotion6,
-                emotion7,
-                emotion8,
-                vq_single_slider,
-                fmax_slider,
-                pitch_std_slider,
-                speaking_rate_slider,
-                dnsmos_slider,
-                speaker_noised_checkbox,
-                cfg_scale_slider,
-                min_p_slider,
-                seed_number,
-                randomize_seed_toggle,
-                unconditional_keys,
-            ],
-            outputs=[output_audio, seed_number],
-        )
+        if MODELS:  # Only set up callbacks if models loaded successfully
+            model_choice.change(
+                fn=update_ui,
+                inputs=[model_choice],
+                outputs=[
+                    text,
+                    language,
+                    speaker_audio,
+                    prefix_audio,
+                    emotion1,
+                    emotion2,
+                    emotion3,
+                    emotion4,
+                    emotion5,
+                    emotion6,
+                    emotion7,
+                    emotion8,
+                    vq_single_slider,
+                    fmax_slider,
+                    pitch_std_slider,
+                    speaking_rate_slider,
+                    dnsmos_slider,
+                    speaker_noised_checkbox,
+                    unconditional_keys,
+                ],
+            )
+
+            # On page load, trigger the same UI refresh
+            demo.load(
+                fn=update_ui,
+                inputs=[model_choice],
+                outputs=[
+                    text,
+                    language,
+                    speaker_audio,
+                    prefix_audio,
+                    emotion1,
+                    emotion2,
+                    emotion3,
+                    emotion4,
+                    emotion5,
+                    emotion6,
+                    emotion7,
+                    emotion8,
+                    vq_single_slider,
+                    fmax_slider,
+                    pitch_std_slider,
+                    speaking_rate_slider,
+                    dnsmos_slider,
+                    speaker_noised_checkbox,
+                    unconditional_keys,
+                ],
+            )
+
+            # Generate audio on button click
+            generate_button.click(
+                fn=generate_audio,
+                inputs=[
+                    model_choice,
+                    text,
+                    language,
+                    speaker_audio,
+                    prefix_audio,
+                    emotion1,
+                    emotion2,
+                    emotion3,
+                    emotion4,
+                    emotion5,
+                    emotion6,
+                    emotion7,
+                    emotion8,
+                    vq_single_slider,
+                    fmax_slider,
+                    pitch_std_slider,
+                    speaking_rate_slider,
+                    dnsmos_slider,
+                    speaker_noised_checkbox,
+                    cfg_scale_slider,
+                    min_p_slider,
+                    seed_number,
+                    randomize_seed_toggle,
+                    unconditional_keys,
+                ],
+                outputs=[output_audio, seed_number],
+            )
 
     return demo