Update app.py

app.py

@@ -1,3 +1,4 @@
+# ---------- Gradio CDN fix ----------
 import os
 os.environ.setdefault("GRADIO_USE_CDN", "true")
 
@@ -11,10 +12,16 @@ import numpy as np
 import soundfile as sf
 from huggingface_hub import hf_hub_download
 
-# Detect ZeroGPU
+# >>> NEW: import spaces unconditionally and define a probe so ZeroGPU sees a @spaces.GPU fn at import time
+import spaces
+
+@spaces.GPU(duration=30)
+def _gpu_probe() -> str:
+    # Dummy function to satisfy ZeroGPU startup check. We never call it.
+    return "ok"
+
+# Detect ZeroGPU (used only to decide whether to call GPU code)
 USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"
-if USE_ZEROGPU:
-    import spaces
 
 SPACE_ROOT = Path(__file__).parent.resolve()
 REPO_DIR   = SPACE_ROOT / "SonicMasterRepo"
@@ -23,7 +30,6 @@ WEIGHTS_FILE = "model.safetensors"
 CACHE_DIR = SPACE_ROOT / "weights"
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-
 # ---------- 1) Pull weights from HF Hub ----------
 def get_weights_path() -> Path:
     return Path(
@@ -37,7 +43,6 @@ def get_weights_path() -> Path:
         )
     )
 
-
 # ---------- 2) Clone GitHub repo ----------
 def ensure_repo() -> Path:
     if not REPO_DIR.exists():
@@ -51,7 +56,6 @@ def ensure_repo() -> Path:
         sys.path.append(REPO_DIR.as_posix())
     return REPO_DIR
 
-
 # ---------- 3) Examples ----------
 def build_examples():
     repo = ensure_repo()
@@ -72,7 +76,6 @@ def build_examples():
     return [[p.as_posix(), prompts[i] if i < len(prompts) else prompts[-1]]
             for i, p in enumerate(wav_paths[:10])]
 
-
 # ---------- 4) I/O helpers ----------
 def save_temp_wav(wav: np.ndarray, sr: int, path: Path):
     if wav.ndim == 2 and wav.shape[0] < wav.shape[1]:
@@ -83,7 +86,6 @@ def read_audio(path: str) -> Tuple[np.ndarray, int]:
     wav, sr = sf.read(path, always_2d=False)
     return wav.astype(np.float32) if wav.dtype == np.float64 else wav, sr
 
-
 # ---------- 5) Core inference ----------
 def run_sonicmaster_cli(input_wav_path: Path,
                         prompt: str,
@@ -127,15 +129,16 @@ def run_sonicmaster_cli(input_wav_path: Path,
             continue
     return False
 
-
-# ---------- 6) ZeroGPU wrapper ----------
-if USE_ZEROGPU:
-    @spaces.GPU(duration=180)
-    def enhance_on_gpu(input_path: str, prompt: str, output_path: str) -> bool:
-        import torch  # ensures CUDA initializes inside GPU context
-        return run_sonicmaster_cli(Path(input_path), prompt, Path(output_path),
-                                   _logs=[], progress=None)
-
+# ---------- 6) ZeroGPU real inference wrapper (ALWAYS defined) ----------
+@spaces.GPU(duration=180)
+def enhance_on_gpu(input_path: str, prompt: str, output_path: str) -> bool:
+    # Ensure CUDA initializes inside GPU context
+    try:
+        import torch  # noqa: F401
+    except Exception:
+        pass
+    from pathlib import Path as _P
+    return run_sonicmaster_cli(_P(input_path), prompt, _P(output_path), _logs=[], progress=None)
 
 # ---------- 7) Gradio callback ----------
 def enhance_audio_ui(audio_path: str,
@@ -151,6 +154,8 @@ def enhance_audio_ui(audio_path: str,
         except: pass
     save_temp_wav(wav, sr, tmp_in)
 
+    # Use GPU only when ZeroGPU is active; otherwise CPU fallback.
+    if progress: progress(0.3, desc="Starting inference")
     if USE_ZEROGPU:
         ok = enhance_on_gpu(tmp_in.as_posix(), prompt, tmp_out.as_posix())
     else:
@@ -162,7 +167,6 @@ def enhance_audio_ui(audio_path: str,
     else:
         return (sr, wav)
 
-
 # ---------- 8) Gradio UI ----------
 with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill_height=True) as demo:
     gr.Markdown("## 🎧 SonicMaster\nUpload or choose an example, write a text prompt, then click **Enhance**.")
@@ -179,7 +183,6 @@ with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill
                   outputs=[out_audio],
                   concurrency_limit=1)
 
-
 # ---------- 9) FastAPI mount ----------
 from fastapi import FastAPI, Request
 from starlette.responses import PlainTextResponse