Update app.py

app.py

@@ -1,11 +1,17 @@
-# ---------- Gradio CDN fix ----------
+# ---------- MUST BE FIRST: Gradio CDN + ZeroGPU probe ----------
 import os
-import spaces
 os.environ.setdefault("GRADIO_USE_CDN", "true")
-@spaces.GPU(duration=30)
+
+# A GPU-decorated function MUST exist at import time for ZeroGPU.
+# Import spaces unconditionally and register a tiny probe.
+import spaces
+
+@spaces.GPU(duration=10)
 def _gpu_probe() -> str:
-    # Dummy function to satisfy ZeroGPU startup check. We never call it.
+    # Never called; only here so ZeroGPU startup check passes.
     return "ok"
+
+# ---------- Standard imports ----------
 import sys
 import subprocess
 from pathlib import Path
@@ -16,14 +22,14 @@ import numpy as np
 import soundfile as sf
 from huggingface_hub import hf_hub_download
 
-# Detect ZeroGPU (used only to decide whether to call GPU code)
+# Detect ZeroGPU to decide whether to CALL the GPU function.
 USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"
 
-SPACE_ROOT = Path(__file__).parent.resolve()
-REPO_DIR   = SPACE_ROOT / "SonicMasterRepo"
+SPACE_ROOT   = Path(__file__).parent.resolve()
+REPO_DIR     = SPACE_ROOT / "SonicMasterRepo"
 WEIGHTS_REPO = "amaai-lab/SonicMaster"
 WEIGHTS_FILE = "model.safetensors"
-CACHE_DIR = SPACE_ROOT / "weights"
+CACHE_DIR    = SPACE_ROOT / "weights"
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
 # ---------- 1) Pull weights from HF Hub ----------
@@ -82,7 +88,7 @@ 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 ----------
+# ---------- 5) Core inference (subprocess calling your repo script) ----------
 def run_sonicmaster_cli(input_wav_path: Path,
                         prompt: str,
                         out_path: Path,
@@ -116,8 +122,8 @@ def run_sonicmaster_cli(input_wav_path: Path,
     for idx, cmd in enumerate(CANDIDATE_CMDS, start=1):
         try:
             if progress: progress(0.35 + 0.05*idx, desc=f"Running inference (try {idx})")
-            res = subprocess.run(cmd, capture_output=True,
-                                 text=True, check=True, env=os.environ.copy())
+            # inherit env so CUDA_VISIBLE_DEVICES from ZeroGPU reaches subprocess
+            subprocess.run(cmd, capture_output=True, text=True, check=True, env=os.environ.copy())
             if out_path.exists() and out_path.stat().st_size > 0:
                 if progress: progress(0.9, desc="Post-processing output")
                 return True
@@ -125,10 +131,10 @@ def run_sonicmaster_cli(input_wav_path: Path,
             continue
     return False
 
-# ---------- 6) ZeroGPU real inference wrapper (ALWAYS defined) ----------
+# ---------- 6) REAL GPU function (always defined; only CALLED on ZeroGPU) ----------
 @spaces.GPU(duration=180)
 def enhance_on_gpu(input_path: str, prompt: str, output_path: str) -> bool:
-    # Ensure CUDA initializes inside GPU context
+    # Import torch here so CUDA initializes inside GPU context
     try:
         import torch  # noqa: F401
     except Exception:
@@ -150,7 +156,6 @@ 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())
@@ -179,7 +184,7 @@ with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill
                   outputs=[out_audio],
                   concurrency_limit=1)
 
-# ---------- 9) FastAPI mount ----------
+# ---------- 9) FastAPI mount & disconnect handler ----------
 from fastapi import FastAPI, Request
 from starlette.responses import PlainTextResponse
 from starlette.requests import ClientDisconnect
@@ -196,4 +201,4 @@ app = gr.mount_gradio_app(app, demo.queue(max_size=16), path="/")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)