Update app.py

app.py

@@ -1,26 +1,204 @@
-# --- ZeroGPU must see a GPU-decorated function at import time ---
+# ---------- 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
 
 @spaces.GPU(duration=10)
-def _gpu_probe():
+def _gpu_probe() -> str:
+    # Never called; only here so ZeroGPU startup check passes.
     return "ok"
 
-# --- Gradio app kept trivial to prove boot path ---
+# ---------- Standard imports ----------
+import sys
+import subprocess
+from pathlib import Path
+from typing import Tuple, Optional
+
 import gradio as gr
-from fastapi import FastAPI
+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()
+REPO_DIR     = SPACE_ROOT / "SonicMasterRepo"
+WEIGHTS_REPO = "amaai-lab/SonicMaster"
+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(
+            repo_id=WEIGHTS_REPO,
+            filename=WEIGHTS_FILE,
+            local_dir=CACHE_DIR.as_posix(),
+            local_dir_use_symlinks=False,
+            force_download=False,
+            resume_download=True,
+        )
+    )
+
+# ---------- 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()],
+            check=True,
+        )
+    if REPO_DIR.as_posix() not in sys.path:
+        sys.path.append(REPO_DIR.as_posix())
+    return REPO_DIR
+
+# ---------- 3) Examples ----------
+def build_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.",
+        "Can you make this sound louder, please?",
+        "Make the audio smoother and less distorted.",
+        "Improve the balance in this song.",
+        "Disentangle the left and right channels to give this song a stereo feeling.",
+        "Correct the unnatural frequency emphasis. Reduce the roominess or echo.",
+        "Raise the level of the vocals, please.",
+        "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])]
+
+# ---------- 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
+    sf.write(path.as_posix(), wav, sr)
+
+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,
+                        _logs: list,
+                        progress: Optional[gr.Progress] = None) -> bool:
+    if progress: progress(0.15, desc="Loading weights & repo")
+    ckpt = get_weights_path()
+    repo = ensure_repo()
 
-def echo(x): 
-    return f"ok: {x}"
+    py = sys.executable or "python3"
+    script_candidates = [repo / "infer_single.py"]
 
-with gr.Blocks(title="Hello") as demo:
-    inp = gr.Textbox(label="Say something")
-    out = gr.Textbox(label="Reply")
-    inp.submit(echo, inp, out)
+    CANDIDATE_CMDS = []
+    for script in script_candidates:
+        if script.exists():
+            CANDIDATE_CMDS.append([
+                py, script.as_posix(),
+                "--ckpt", ckpt.as_posix(),
+                "--input", input_wav_path.as_posix(),
+                "--prompt", prompt,
+                "--output", out_path.as_posix(),
+            ])
+            CANDIDATE_CMDS.append([
+                py, script.as_posix(),
+                "--weights", ckpt.as_posix(),
+                "--input", input_wav_path.as_posix(),
+                "--text", prompt,
+                "--out", out_path.as_posix(),
+            ])
+
+    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")
+                return True
+        except Exception:
+            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
+    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]:
+    if not audio_path or not prompt:
+        raise gr.Error("Please provide audio and a text prompt.")
+
+    wav, sr = read_audio(audio_path)
+    tmp_in, tmp_out = SPACE_ROOT / "tmp_in.wav", SPACE_ROOT / "tmp_out.wav"
+    if tmp_out.exists():
+        try: tmp_out.unlink()
+        except: pass
+    save_temp_wav(wav, sr, tmp_in)
+
+    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:
+        ok = run_sonicmaster_cli(tmp_in, prompt, tmp_out, _logs=[], 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)
+    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**.")
+    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")
+            run_btn = gr.Button("🚀 Enhance", variant="primary")
+            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)
+
+# ---------- 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()
 
-# Expose *either* 'demo' or a FastAPI 'app'. We'll use FastAPI + mount:
 app = FastAPI()
 
-# Mount Gradio at root so Spaces healthcheck to "/" gets 200
-app = gr.mount_gradio_app(app, demo.queue(), path="/")
+@app.exception_handler(ClientDisconnect)
+async def client_disconnect_handler(request: Request, exc: ClientDisconnect):
+    return PlainTextResponse("Client disconnected", status_code=499)
+
+app = gr.mount_gradio_app(app, demo.queue(max_size=16), path="/")
 
-# DO NOT run uvicorn here — Spaces runs the server.
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)