Update app.py

app.py

@@ -1,33 +1,246 @@
-# ---------- Gradio Space entrypoint (no FastAPI/Uvicorn) ----------
+# ========== MUST BE FIRST: Gradio SDK entry + ZeroGPU probes ==========
 import os
 os.environ.setdefault("GRADIO_USE_CDN", "true")
 
-import gradio as gr
-
-# Optional: harmless on CPU; useful once you switch to ZeroGPU hardware
+# Optional: 'spaces' is present on Spaces; harmless to try locally.
 try:
     import spaces
-    @spaces.GPU(duration=10)
-    def gpu_probe(a: int = 1, b: int = 1):
-        return a + b
 except Exception:
-    pass
+    class _DummySpaces:
+        def GPU(self, *_, **__):
+            def deco(fn): return fn
+            return deco
+    spaces = _DummySpaces()
+
+# PUBLIC names so ZeroGPU supervisor can detect them
+@spaces.GPU(duration=10)
+def gpu_probe(a: int = 1, b: int = 1):
+    return a + b
+
+@spaces.GPU(duration=10)
+def gpu_echo(x: str = "ok"):
+    return x
+
+# ========== Standard imports ==========
+import sys
+import subprocess
+from pathlib import Path
+from typing import Tuple, Optional, List
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+from huggingface_hub import hf_hub_download
+
+# ZeroGPU runtime hint (still safe on CPU)
+USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"
+
+SPACE_ROOT   = Path(__file__).parent.resolve()
+REPO_DIR     = SPACE_ROOT / "SonicMasterRepo"
+REPO_URL     = "https://github.com/AMAAI-Lab/SonicMaster"
+WEIGHTS_REPO = "amaai-lab/SonicMaster"
+WEIGHTS_FILE = "model.safetensors"
+CACHE_DIR    = SPACE_ROOT / "weights"
+CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+# ========== Lazy resources (no heavy work at import) ==========
+_weights_path: Optional[Path] = None
+_repo_ready: bool = False
+
+def get_weights_path(progress: Optional[gr.Progress] = None) -> Path:
+    """Download/resolve weights lazily."""
+    global _weights_path
+    if _weights_path is None:
+        if progress: progress(0.10, desc="Downloading model weights (first run)")
+        wp = hf_hub_download(
+            repo_id=WEIGHTS_REPO,
+            filename=WEIGHTS_FILE,
+            local_dir=str(CACHE_DIR),
+            local_dir_use_symlinks=False,
+            force_download=False,
+            resume_download=True,
+        )
+        _weights_path = Path(wp)
+    return _weights_path
+
+def ensure_repo(progress: Optional[gr.Progress] = None) -> Path:
+    """Clone the repo lazily and add to sys.path."""
+    global _repo_ready
+    if not _repo_ready:
+        if not REPO_DIR.exists():
+            if progress: progress(0.18, desc="Cloning SonicMaster repo (first run)")
+            subprocess.run(
+                ["git", "clone", "--depth", "1", REPO_URL, REPO_DIR.as_posix()],
+                check=True,
+            )
+        if REPO_DIR.as_posix() not in sys.path:
+            sys.path.append(REPO_DIR.as_posix())
+        _repo_ready = True
+    return REPO_DIR
+
+# ========== Helpers ==========
+def save_temp_wav(wav: np.ndarray, sr: int, path: Path):
+    # Ensure shape (samples, channels)
+    if wav.ndim == 2 and wav.shape[0] < wav.shape[1]:
+        wav = wav.T
+    if wav.dtype == np.float64:
+        wav = wav.astype(np.float32)
+    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 _candidate_commands(py: str, script: Path, ckpt: Path, inp: Path, prompt: str, out: Path) -> List[List[str]]:
+    # Try common flag layouts
+    return [
+        [py, script.as_posix(), "--ckpt",   ckpt.as_posix(), "--input", inp.as_posix(), "--prompt", prompt, "--output", out.as_posix()],
+        [py, script.as_posix(), "--weights",ckpt.as_posix(), "--input", inp.as_posix(), "--text",   prompt, "--out",    out.as_posix()],
+        [py, script.as_posix(), "--ckpt",   ckpt.as_posix(), "--input", inp.as_posix(), "--text",   prompt, "--output", out.as_posix()],
+    ]
+
+def run_sonicmaster_cli(
+    input_wav_path: Path,
+    prompt: str,
+    out_path: Path,
+    progress: Optional[gr.Progress] = None,
+) -> Tuple[bool, str]:
+    """Run inference scripts via subprocess; return (ok, message)."""
+    if progress: progress(0.14, desc="Preparing inference")
+    ckpt = get_weights_path(progress=progress)
+    repo = ensure_repo(progress=progress)
+
+    candidates = [repo / "infer_single.py", repo / "inference_fullsong.py", repo / "inference_ptload_batch.py"]
+    scripts = [s for s in candidates if s.exists()]
+    if not scripts:
+        return False, "No inference script found in the repo (expected infer_single.py or similar)."
+
+    py = sys.executable or "python3"
+    env = os.environ.copy()
+
+    last_err = ""
+    for sidx, script in enumerate(scripts, 1):
+        for cidx, cmd in enumerate(_candidate_commands(py, script, ckpt, input_wav_path, prompt, out_path), 1):
+            try:
+                if progress:
+                    progress(min(0.20 + 0.08 * (sidx + cidx), 0.70), desc=f"Running {script.name} (try {sidx}.{cidx})")
+                res = subprocess.run(cmd, capture_output=True, text=True, check=True, env=env)
+                if out_path.exists() and out_path.stat().st_size > 0:
+                    if progress: progress(0.88, desc="Post-processing output")
+                    return True, (res.stdout or "Inference completed.").strip()
+                last_err = f"{script.name} produced no output file."
+            except subprocess.CalledProcessError as e:
+                snippet = "\n".join(filter(None, [e.stdout or "", e.stderr or ""])).strip()
+                last_err = snippet if snippet else f"{script.name} failed with return code {e.returncode}."
+            except Exception as e:
+                import traceback
+                last_err = f"Unexpected error: {e}\n{traceback.format_exc()}"
+    return False, last_err or "All candidate commands failed."
+
+# ========== GPU path (called only if ZeroGPU/GPU available) ==========
+@spaces.GPU(duration=180)
+def enhance_on_gpu(input_path: str, prompt: str, output_path: str) -> Tuple[bool, str]:
+    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), progress=None)
+
+def _has_cuda() -> bool:
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
+# ========== Gradio callback ==========
+def enhance_audio_ui(
+    audio_path: str,
+    prompt: str,
+    progress=gr.Progress(track_tqdm=True),
+) -> Tuple[Optional[Tuple[int, np.ndarray]], str]:
+    """
+    Returns (audio, message). On failure, audio=None and message=error text.
+    """
+    try:
+        if not prompt:
+            raise gr.Error("Please provide a text prompt.")
+        if not audio_path:
+            raise gr.Error("Please upload or select an input audio file.")
+
+        wav, sr = read_audio(audio_path)
+        tmp_in  = SPACE_ROOT / "tmp_in.wav"
+        tmp_out = SPACE_ROOT / "tmp_out.wav"
+        if tmp_out.exists():
+            try: tmp_out.unlink()
+            except Exception: pass
+
+        if progress: progress(0.06, desc="Preparing audio")
+        save_temp_wav(wav, sr, tmp_in)
+
+        use_gpu_call = USE_ZEROGPU or _has_cuda()
+        if progress: progress(0.12, desc="Starting inference")
+
+        if use_gpu_call:
+            ok, msg = enhance_on_gpu(tmp_in.as_posix(), prompt, tmp_out.as_posix())
+        else:
+            ok, msg = run_sonicmaster_cli(tmp_in, prompt, tmp_out, progress=progress)
+
+        if ok and tmp_out.exists() and tmp_out.stat().st_size > 0:
+            out_wav, out_sr = read_audio(tmp_out.as_posix())
+            return (out_sr, out_wav), (msg or "Done.")
+        else:
+            return None, (msg or "Inference failed without a specific error message.")
+
+    except gr.Error as e:
+        return None, str(e)
+    except Exception as e:
+        import traceback
+        return None, f"Unexpected error: {e}\n{traceback.format_exc()}"
+
+# ========== Gradio UI ==========
+PROMPT_EXAMPLES = [
+    ["Increase the clarity of this song by emphasizing treble frequencies."],
+    ["Make this song sound more boomy by amplifying the low end bass frequencies."],
+    ["Make the audio smoother and less distorted."],
+    ["Improve the balance in this song."],
+    ["Reduce roominess/echo (dereverb)."],
+    ["Raise the level of the vocals."],
+    ["Give the song a wider stereo image."],
+]
 
-# Build a tiny UI
-def echo(s: str) -> str:
-    return f"echo: {s}"
+with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill_height=True) as _demo:
+    gr.Markdown(
+        "## 🎧 SonicMaster\n"
+        "Upload or choose an example prompt, write your own instruction, then click **Enhance**.\n"
+        "- First run downloads model weights & repo (progress will show).\n"
+        "- On failure, the **Status** box shows the exact error (we won't echo the input audio)."
+    )
+    with gr.Row():
+        with gr.Column():
+            in_audio = gr.Audio(label="Input Audio", type="filepath")
+            prompt   = gr.Textbox(label="Text Prompt", placeholder="e.g., Reduce reverb and brighten vocals.")
+            run_btn  = gr.Button("🚀 Enhance", variant="primary")
+            gr.Examples(examples=PROMPT_EXAMPLES, inputs=[prompt], label="Prompt Examples")
+        with gr.Column():
+            out_audio = gr.Audio(label="Enhanced Audio (output)")
+            status    = gr.Textbox(label="Status / Messages", interactive=False, lines=8)
 
-with gr.Blocks(title="Hello Space") as _demo:
-    gr.Markdown("### ✅ App is alive\nType to echo.")
-    inp = gr.Textbox(label="Input", value="hello")
-    out = gr.Textbox(label="Output")
-    inp.submit(echo, inp, out)
+    run_btn.click(
+        fn=enhance_audio_ui,
+        inputs=[in_audio, prompt],
+        outputs=[out_audio, status],
+        concurrency_limit=1,
+    )
 
-# Expose all common names that the Space supervisor might look for
-demo = _demo.queue(max_size=8)  # primary export
-iface = demo                     # alias
-app = demo                       # another alias
+# Expose all common names the supervisor might look for
+demo = _demo.queue(max_size=16)
+iface = demo
+app = demo
 
-# For local debugging only (ignored on Spaces)
+# Local debugging only
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)