Update app.py

app.py

@@ -1,17 +1,20 @@
-# ---------- MUST BE FIRST: Gradio CDN + ZeroGPU probe ----------
-import os
-os.environ.setdefault("GRADIO_USE_CDN", "true")
-
-# A GPU-decorated function MUST exist at import time for ZeroGPU.
-# Import spaces unconditionally and register a tiny probe.
-import spaces
+# --- ABSOLUTE TOP (see above block) ---
+import importlib
+try:
+    import spaces
+except Exception as e:
+    raise RuntimeError("Failed to import 'spaces' package. Add `spaces` to requirements.txt.") from e
 
 @spaces.GPU(duration=10)
-def _gpu_probe() -> str:
-    # Never called; only here so ZeroGPU startup check passes.
+def _zerogpu_probe():
     return "ok"
+# --------------------------------------
+
+# You can set env vars after the probe—safe.
+import os
+os.environ.setdefault("GRADIO_USE_CDN", "true")
 
-# ---------- Standard imports ----------
+# Standard imports
 import sys
 import subprocess
 from pathlib import Path
@@ -22,7 +25,6 @@ import numpy as np
 import soundfile as sf
 from huggingface_hub import hf_hub_download
 
-# Detect ZeroGPU to decide whether to CALL the GPU function.
 USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"
 
 SPACE_ROOT   = Path(__file__).parent.resolve()
@@ -32,7 +34,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(
         hf_hub_download(
@@ -45,7 +46,6 @@ def get_weights_path() -> Path:
         )
     )
 
-# ---------- 2) Clone GitHub repo ----------
 def ensure_repo() -> Path:
     if not REPO_DIR.exists():
         subprocess.run(
@@ -58,7 +58,6 @@ def ensure_repo() -> Path:
         sys.path.append(REPO_DIR.as_posix())
     return REPO_DIR
 
-# ---------- 3) Examples ----------
 def build_examples():
     repo = ensure_repo()
     wav_dir = repo / "samples" / "inputs"
@@ -78,7 +77,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]:
         wav = wav.T
@@ -88,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 (subprocess calling your repo script) ----------
 def run_sonicmaster_cli(input_wav_path: Path,
                         prompt: str,
                         out_path: Path,
@@ -122,7 +119,6 @@ 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})")
-            # 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")
@@ -131,18 +127,15 @@ def run_sonicmaster_cli(input_wav_path: Path,
             continue
     return False
 
-# ---------- 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:
-    # Import torch here so CUDA initializes inside GPU context
     try:
-        import torch  # noqa: F401
+        import torch  # ensure CUDA init happens in GPU context
     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,
                      prompt: str,
                      progress=gr.Progress(track_tqdm=True)) -> Tuple[int, np.ndarray]:
@@ -166,9 +159,9 @@ def enhance_audio_ui(audio_path: str,
         out_wav, out_sr = read_audio(tmp_out.as_posix())
         return (out_sr, out_wav)
     else:
+        # If inference fails, return original audio (your chosen fallback).
         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**.")
     with gr.Row():
@@ -184,12 +177,13 @@ with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill
                   outputs=[out_audio],
                   concurrency_limit=1)
 
-# ---------- 9) FastAPI mount & disconnect handler ----------
 from fastapi import FastAPI, Request
 from starlette.responses import PlainTextResponse
 from starlette.requests import ClientDisconnect
 
-_ = get_weights_path(); _ = ensure_repo()
+# Preload light assets after probe so import won’t fail before detector.
+_ = get_weights_path()
+_ = ensure_repo()
 
 app = FastAPI()
 
@@ -197,6 +191,7 @@ app = FastAPI()
 async def client_disconnect_handler(request: Request, exc: ClientDisconnect):
     return PlainTextResponse("Client disconnected", status_code=499)
 
+# (Queue() at mount time is fine)
 app = gr.mount_gradio_app(app, demo.queue(max_size=16), path="/")
 
 if __name__ == "__main__":