Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 os.environ.setdefault("GRADIO_USE_CDN", "true")
+
 import sys
-import spaces
 import subprocess
 from pathlib import Path
 from typing import Tuple, Optional
@@ -11,57 +11,52 @@ import numpy as np
 import soundfile as sf
 from huggingface_hub import hf_hub_download
 
+# Detect ZeroGPU
+USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"
+if USE_ZEROGPU:
+    import spaces
+
 SPACE_ROOT = Path(__file__).parent.resolve()
 REPO_DIR   = SPACE_ROOT / "SonicMasterRepo"
 WEIGHTS_REPO = "amaai-lab/SonicMaster"
-WEIGHTS_FILE = "model.safetensors"   # from the HF model repo
+WEIGHTS_FILE = "model.safetensors"
 CACHE_DIR = SPACE_ROOT / "weights"
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-# --- ZeroGPU compatibility (no behavior change) ---
-USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"
-if USE_ZEROGPU:
-    import spaces
 
-    @spaces.GPU(duration=60)
-    def _zgpu_dummy():
-        # Minimal GPU-annotated function so ZeroGPU runtime is satisfied.
-        # We never call this; actual inference stays CPU unless you choose otherwise.
-        return "ok"        
 # ---------- 1) Pull weights from HF Hub ----------
 def get_weights_path() -> Path:
-    weights_path = hf_hub_download(
-        repo_id=WEIGHTS_REPO,
-        filename=WEIGHTS_FILE,
-        local_dir=CACHE_DIR.as_posix(),
-        local_dir_use_symlinks=False,
-        force_download=False,
-        resume_download=True,
+    return Path(
+        hf_hub_download(
+            repo_id=WEIGHTS_REPO,
+            filename=WEIGHTS_FILE,
+            local_dir=CACHE_DIR.as_posix(),
+            local_dir_use_symlinks=False,
+            force_download=False,
+            resume_download=True,
+        )
     )
-    return Path(weights_path)
 
-# ---------- 2) Clone GitHub repo for code (model.py / inference_*.py ) ----------
+
+# ---------- 2) Clone GitHub repo ----------
 def ensure_repo() -> Path:
     if not REPO_DIR.exists():
         subprocess.run(
-            ["git", "clone", "--depth", "1", "https://github.com/AMAAI-Lab/SonicMaster", REPO_DIR.as_posix()],
+            ["git", "clone", "--depth", "1",
+             "https://github.com/AMAAI-Lab/SonicMaster",
+             REPO_DIR.as_posix()],
             check=True,
         )
     if REPO_DIR.as_posix() not in sys.path:
         sys.path.append(REPO_DIR.as_posix())
     return REPO_DIR
 
-# ---------- 3) Examples: use only *.wav from samples/inputs ----------
+
+# ---------- 3) Examples ----------
 def build_examples():
-    """
-    Discover up to 10 .wav files from:
-        SonicMasterRepo/samples/inputs
-    and pair them with prompts for gr.Examples.
-    """
     repo = ensure_repo()
     wav_dir = repo / "samples" / "inputs"
     wav_paths = sorted(p for p in wav_dir.glob("*.wav") if p.is_file())
-
     prompts = [
         "Increase the clarity of this song by emphasizing treble frequencies.",
         "Make this song sound more boomy by amplifying the low end bass frequencies.",
@@ -74,54 +69,33 @@ def build_examples():
         "Increase the clarity of this song by emphasizing treble frequencies.",
         "Please, dereverb this audio.",
     ]
+    return [[p.as_posix(), prompts[i] if i < len(prompts) else prompts[-1]]
+            for i, p in enumerate(wav_paths[:10])]
 
-    examples = []
-    for i, p in enumerate(wav_paths[:10]):
-        prompt = prompts[i] if i < len(prompts) else prompts[-1]
-        examples.append([p.as_posix(), prompt])
-
-    # Fallback: if no wavs found, provide an empty list (Gradio handles it)
-    return examples
 
 # ---------- 4) I/O helpers ----------
 def save_temp_wav(wav: np.ndarray, sr: int, path: Path):
-    # Ensure (samples, channels) for soundfile
     if wav.ndim == 2 and wav.shape[0] < wav.shape[1]:
-        # (channels, samples) -> (samples, channels)
-        data = wav.T
-    else:
-        data = wav
-    sf.write(path.as_posix(), data, sr)
+        wav = wav.T
+    sf.write(path.as_posix(), wav, sr)
 
 def read_audio(path: str) -> Tuple[np.ndarray, int]:
     wav, sr = sf.read(path, always_2d=False)
-    if wav.dtype == np.float64:
-        wav = wav.astype(np.float32)
-    return wav, sr
-
-def run_sonicmaster_cli(
-    input_wav_path: Path,
-    prompt: str,
-    out_path: Path,
-    _logs: list,  # kept for compatibility, but not shown in UI
-    progress: Optional[gr.Progress] = None
-) -> bool:
-    """
-    Uses the current Python interpreter and tries a few script names/flags.
-    """
-    import sys, shutil
+    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,
+                        out_path: Path,
+                        _logs: list,
+                        progress: Optional[gr.Progress] = None) -> bool:
     if progress: progress(0.15, desc="Loading weights & repo")
     ckpt = get_weights_path()
     repo = ensure_repo()
 
-    # Use the exact Python interpreter running this process
-    py = sys.executable or shutil.which("python3") or shutil.which("python") or "python3"
-
-    # Prefer the scripts we know accept --ckpt/--input/--prompt/--output
-    script_candidates = [
-        repo / "infer_single.py",         # if you kept your own name
-    ]
+    py = sys.executable or "python3"
+    script_candidates = [repo / "infer_single.py"]
 
     CANDIDATE_CMDS = []
     for script in script_candidates:
@@ -133,10 +107,6 @@ def run_sonicmaster_cli(
                 "--prompt", prompt,
                 "--output", out_path.as_posix(),
             ])
-
-    # As a last resort, try alternative flag names (if someone changed the CLI)
-    for script in script_candidates:
-        if script.exists():
             CANDIDATE_CMDS.append([
                 py, script.as_posix(),
                 "--weights", ckpt.as_posix(),
@@ -145,110 +115,86 @@ def run_sonicmaster_cli(
                 "--out", out_path.as_posix(),
             ])
 
-    if not CANDIDATE_CMDS:
-        return False
-
     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)
+            res = 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
-        except subprocess.CalledProcessError:
-            continue
         except Exception:
             continue
     return False
 
 
-def enhance_audio_ui(
-    audio_path: str,
-    prompt: str,
-    progress=gr.Progress(track_tqdm=True)
-) -> Tuple[int, np.ndarray]:
-    """
-    Gradio callback: accepts a file path, a prompt, and returns enhanced audio.
-    """
-    if progress: progress(0.0, desc="Validating input")
+# ---------- 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)
+
+
+# ---------- 7) Gradio callback ----------
+def enhance_audio_ui(audio_path: str,
+                     prompt: str,
+                     progress=gr.Progress(track_tqdm=True)) -> Tuple[int, np.ndarray]:
     if not audio_path or not prompt:
         raise gr.Error("Please provide audio and a text prompt.")
 
-    # Standardize input -> temp wav
     wav, sr = read_audio(audio_path)
-    if progress: progress(0.15, desc="Preparing audio")
-    tmp_in = SPACE_ROOT / "tmp_in.wav"
-    tmp_out = SPACE_ROOT / "tmp_out.wav"
+    tmp_in, tmp_out = SPACE_ROOT / "tmp_in.wav", SPACE_ROOT / "tmp_out.wav"
     if tmp_out.exists():
-        try:
-            tmp_out.unlink()
-        except Exception:
-            pass
-
+        try: tmp_out.unlink()
+        except: pass
     save_temp_wav(wav, sr, tmp_in)
 
-    # Run model
-    if progress: progress(0.3, desc="Starting inference")
-    ok = run_sonicmaster_cli(tmp_in, prompt, tmp_out, _logs=[], progress=progress)
+    if USE_ZEROGPU:
+        ok = enhance_on_gpu(tmp_in.as_posix(), prompt, tmp_out.as_posix())
+    else:
+        ok = run_sonicmaster_cli(tmp_in, prompt, tmp_out, _logs=[], progress=progress)
 
-    # Return output (or echo input)
     if ok and tmp_out.exists() and tmp_out.stat().st_size > 0:
         out_wav, out_sr = read_audio(tmp_out.as_posix())
-        if progress: progress(1.0, desc="Done")
         return (out_sr, out_wav)
     else:
-        if progress: progress(1.0, desc="No output produced")
-        # Return original audio if model didn't produce output
         return (sr, wav)
 
-# ---------- 6) Gradio UI ----------
+
+# ---------- 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 (from repo: `samples/inputs/*.wav`), write a text prompt (e.g., *reduce reverb*, *clean distortion*), then click **Enhance**.")
+    gr.Markdown("## 🎧 SonicMaster\nUpload or choose an example, write a text prompt, then click **Enhance**.")
     with gr.Row():
-        with gr.Column(scale=1):
-            in_audio = gr.Audio(label="Input Audio (upload or use examples)", type="filepath")
-            prompt = gr.Textbox(label="Text Prompt", placeholder="e.g., reduce reverb and enhance clarity")
+        with gr.Column():
+            in_audio = gr.Audio(label="Input Audio", type="filepath")
+            prompt = gr.Textbox(label="Text Prompt", placeholder="e.g., reduce reverb")
             run_btn = gr.Button("🚀 Enhance", variant="primary")
-
-            # Use wavs from SonicMasterRepo/samples/inputs
-            gr.Examples(
-                examples=build_examples(),
-                inputs=[in_audio, prompt],
-                label="Examples (repo: samples/inputs/*.wav)"
-            )
-        with gr.Column(scale=1):
+            gr.Examples(examples=build_examples(), inputs=[in_audio, prompt])
+        with gr.Column():
             out_audio = gr.Audio(label="Enhanced Audio (output)")
+    run_btn.click(fn=enhance_audio_ui,
+                  inputs=[in_audio, prompt],
+                  outputs=[out_audio],
+                  concurrency_limit=1)
 
-    # Per-event concurrency (use 1 unless you know your VRAM/CPU can handle more)
-    run_btn.click(
-        fn=enhance_audio_ui,
-        inputs=[in_audio, prompt],
-        outputs=[out_audio],
-        concurrency_limit=1,
-    )
 
-# ---------- FastAPI mount with ClientDisconnect handler (Spaces-friendly) ----------
+# ---------- 9) FastAPI mount ----------
 from fastapi import FastAPI, Request
 from starlette.responses import PlainTextResponse
 from starlette.requests import ClientDisconnect
-import gradio as gr
 
-# Warm up cache & repo (optional but nice)
-_ = get_weights_path()
-_ = ensure_repo()
+_ = get_weights_path(); _ = ensure_repo()
 
 app = FastAPI()
 
 @app.exception_handler(ClientDisconnect)
 async def client_disconnect_handler(request: Request, exc: ClientDisconnect):
-    # Treat as benign client cancel; avoids noisy tracebacks on Spaces
     return PlainTextResponse("Client disconnected", status_code=499)
 
-# Mount Gradio at root. Use queue if you were using it before.
 app = gr.mount_gradio_app(app, demo.queue(max_size=16), path="/")
 
-# Optional: allow local dev with `python app.py`
 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)