Update app.py

app.py

@@ -1,204 +1,11 @@
-# ---------- 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 os; os.environ.setdefault("GRADIO_USE_CDN","true")
 import spaces
-
-@spaces.GPU(duration=10)
-def _gpu_probe() -> str:
-    # Never called; only here so ZeroGPU startup check passes.
-    return "ok"
-
-# ---------- Standard imports ----------
-import sys
-import subprocess
-from pathlib import Path
-from typing import Tuple, Optional
-
+@spaces.GPU(duration=10) 
+def ping(): return "pong"
 import gradio as gr
-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()
-
-    py = sys.executable or "python3"
-    script_candidates = [repo / "infer_single.py"]
-
-    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()
-
+def fn(x): return f"ok:{x}"
+with gr.Blocks() as demo:
+    t=gr.Textbox(); o=gr.Textbox(); t.submit(fn, t, o)
+from fastapi import FastAPI
 app = FastAPI()
-
-@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="/")
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+app = gr.mount_gradio_app(app, demo, path="/")