Update app.py

Fix infer_single.py not Found

app.py

@@ -1,6 +1,14 @@
+# app.py — end-to-end Hugging Face Spaces app for SonicMaster
+# - Works even if `git` is NOT available (falls back to GitHub ZIP download)
+# - Lazily fetches code + weights only when user clicks "Enhance"
+# - Runs inference ONLY via infer_single.py
+
 import os
 import sys
 import subprocess
+import shutil
+import zipfile
+import urllib.request
 from pathlib import Path
 from typing import Tuple, Optional, List, Any
 
@@ -31,32 +39,120 @@ 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)
 
 # ================== SAFE repo handling (NO network at import) ==================
 _repo_ready: bool = False
 
+
+def _safe_unlink(p: Path):
+    try:
+        if p.exists():
+            p.unlink()
+    except Exception:
+        pass
+
+
+def _rmtree(p: Path):
+    try:
+        if p.exists():
+            shutil.rmtree(p)
+    except Exception:
+        pass
+
+
 def ensure_repo(progress: Optional[gr.Progress] = None) -> Path:
     """
     Ensure SonicMaster repo is available.
-    IMPORTANT: Called lazily (on user action), not at import time.
+    Tries:
+      1) git clone (if git exists)
+      2) GitHub ZIP download + extract (if git missing or clone fails)
+    Called lazily (on button click), not at import time.
     """
     global _repo_ready
-    if _repo_ready and REPO_DIR.exists():
+
+    # If already ready and infer exists, done
+    script0 = REPO_DIR / "infer_single.py"
+    if _repo_ready and script0.exists():
         if REPO_DIR.as_posix() not in sys.path:
             sys.path.append(REPO_DIR.as_posix())
         return REPO_DIR
 
+    # If directory exists but script missing, treat as corrupted and reset
+    if REPO_DIR.exists() and not script0.exists():
+        _rmtree(REPO_DIR)
+
     if not REPO_DIR.exists():
         if progress:
-            progress(0.02, desc="Cloning SonicMaster repo (first run)")
-        # Shallow clone to keep it fast
-        subprocess.run(
-            ["git", "clone", "--depth", "1", REPO_URL, REPO_DIR.as_posix()],
-            check=True,
-            capture_output=True,
-            text=True,
+            progress(0.02, desc="Fetching SonicMaster code (first run)")
+
+        git_bin = shutil.which("git")
+        cloned_ok = False
+
+        # Try git clone if available
+        if git_bin:
+            try:
+                subprocess.run(
+                    [git_bin, "clone", "--depth", "1", REPO_URL, REPO_DIR.as_posix()],
+                    check=True,
+                    capture_output=True,
+                    text=True,
+                )
+                cloned_ok = True
+            except Exception:
+                cloned_ok = False
+                # If partial directory created, remove and retry via ZIP
+                if REPO_DIR.exists():
+                    _rmtree(REPO_DIR)
+
+        # Fallback: download ZIP
+        if not cloned_ok:
+            if progress:
+                progress(0.05, desc="Downloading SonicMaster ZIP (no git / clone failed)")
+
+            zip_url = "https://codeload.github.com/AMAAI-Lab/SonicMaster/zip/refs/heads/main"
+            zip_path = SPACE_ROOT / "SonicMaster.zip"
+
+            # Download ZIP
+            urllib.request.urlretrieve(zip_url, zip_path.as_posix())
+
+            if progress:
+                progress(0.08, desc="Extracting SonicMaster ZIP")
+
+            # Extract ZIP to SPACE_ROOT
+            with zipfile.ZipFile(zip_path, "r") as zf:
+                zf.extractall(SPACE_ROOT)
+
+            # GitHub zip usually extracts to SonicMaster-main/
+            extracted = SPACE_ROOT / "SonicMaster-main"
+            if extracted.exists() and (extracted / "infer_single.py").exists():
+                # If destination already exists, remove it
+                if REPO_DIR.exists():
+                    _rmtree(REPO_DIR)
+                extracted.rename(REPO_DIR)
+            else:
+                # If structure differs, try to locate the folder containing infer_single.py
+                found = None
+                for cand in SPACE_ROOT.glob("SonicMaster-*"):
+                    if cand.is_dir() and (cand / "infer_single.py").exists():
+                        found = cand
+                        break
+                if found:
+                    if REPO_DIR.exists():
+                        _rmtree(REPO_DIR)
+                    found.rename(REPO_DIR)
+
+            _safe_unlink(zip_path)
+
+    # Final sanity check
+    if not (REPO_DIR / "infer_single.py").exists():
+        existing = sorted([p.name for p in REPO_DIR.glob("*")]) if REPO_DIR.exists() else []
+        raise RuntimeError(
+            "SonicMaster code fetch finished, but infer_single.py is still missing.\n"
+            f"REPO_DIR: {REPO_DIR}\n"
+            f"Contents: {existing[:80]}"
         )
 
     if REPO_DIR.as_posix() not in sys.path:
@@ -65,9 +161,11 @@ def ensure_repo(progress: Optional[gr.Progress] = None) -> Path:
     _repo_ready = True
     return REPO_DIR
 
+
 # ================ Weights: lazy download (first click) ================
 _weights_path: Optional[Path] = None
 
+
 def get_weights_path(progress: Optional[gr.Progress] = None) -> Path:
     """
     Download/resolve weights lazily (keeps startup fast).
@@ -87,6 +185,7 @@ def get_weights_path(progress: Optional[gr.Progress] = None) -> Path:
         _weights_path = Path(wp)
     return _weights_path
 
+
 # ================== Audio helpers ==================
 def save_temp_wav(wav: np.ndarray, sr: int, path: Path):
     # Ensure shape (samples, channels)
@@ -96,12 +195,22 @@ def save_temp_wav(wav: np.ndarray, sr: int, path: Path):
         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 isinstance(wav, np.ndarray) and wav.dtype == np.float64:
         wav = wav.astype(np.float32)
     return wav, sr
 
+
+def _has_cuda() -> bool:
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
+
 # ================== CLI runner ==================
 def _candidate_commands(
     py: str, script: Path, ckpt: Path, inp: Path, prompt: str, out: Path
@@ -110,20 +219,15 @@ def _candidate_commands(
     Only support infer_single.py variants.
     Expected flags: --ckpt --input --prompt --output
     """
-    return [
-        [
-            py,
-            script.as_posix(),
-            "--ckpt",
-            ckpt.as_posix(),
-            "--input",
-            inp.as_posix(),
-            "--prompt",
-            prompt,
-            "--output",
-            out.as_posix(),
-        ],
-    ]
+    return [[
+        py,
+        script.as_posix(),
+        "--ckpt", ckpt.as_posix(),
+        "--input", inp.as_posix(),
+        "--prompt", prompt,
+        "--output", out.as_posix(),
+    ]]
+
 
 def run_sonicmaster_cli(
     input_wav_path: Path,
@@ -147,23 +251,33 @@ def run_sonicmaster_cli(
 
     script = REPO_DIR / "infer_single.py"
     if not script.exists():
-        return False, "infer_single.py not found in the SonicMaster repo."
+        return False, "infer_single.py not found in the SonicMaster repo (code fetch likely failed)."
 
     py = sys.executable or "python3"
     env = os.environ.copy()
 
+    # Make sure subprocess runs inside repo (some projects rely on relative paths)
+    cwd = REPO_DIR.as_posix()
+
     last_err = ""
-    for cidx, cmd in enumerate(
-        _candidate_commands(py, script, ckpt, input_wav_path, prompt, out_path), 1
-    ):
+    for cidx, cmd in enumerate(_candidate_commands(py, script, ckpt, input_wav_path, prompt, out_path), 1):
         try:
             if progress:
                 progress(min(0.25 + 0.10 * cidx, 0.70), desc=f"Running inference (try {cidx})")
-            res = subprocess.run(cmd, capture_output=True, text=True, check=True, env=env)
+            res = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                check=True,
+                env=env,
+                cwd=cwd,
+            )
+
             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()
+                stdout = (res.stdout or "").strip()
+                return True, (stdout or "Inference completed.")
             last_err = "infer_single.py finished but produced no output file."
         except subprocess.CalledProcessError as e:
             snippet = "\n".join(filter(None, [e.stdout or "", e.stderr or ""])).strip()
@@ -174,20 +288,15 @@ def run_sonicmaster_cli(
 
     return False, last_err or "Inference failed."
 
+
 # ============ GPU path (ZeroGPU) ============
 @spaces.GPU(duration=60)  # safe cap for ZeroGPU tiers
 def enhance_on_gpu(input_path: str, prompt: str, output_path: str) -> Tuple[bool, str]:
     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
 
-# ================== Optional Examples (NO CLONE AT STARTUP) ==================
+# ================== Optional Examples (NO FETCH AT STARTUP) ==================
 PROMPTS_10 = [
     "Increase the clarity of this song by emphasizing treble frequencies.",
     "Make this song sound more boomy by amplifying the low end bass frequencies.",
@@ -201,21 +310,23 @@ PROMPTS_10 = [
     "Please, dereverb this audio.",
 ]
 
+
 def build_examples_if_repo_present() -> List[List[Any]]:
     """
-    Build examples WITHOUT cloning. If repo isn't present yet, return [].
-    This avoids slow startup + network calls.
+    Build examples WITHOUT cloning/downloading.
+    If repo isn't present yet, return [].
     """
     wav_dir = REPO_DIR / "samples" / "inputs"
     if not wav_dir.exists():
         return []
     wav_paths = sorted(p for p in wav_dir.glob("*.wav") if p.is_file())
-    ex = []
+    ex: List[List[Any]] = []
     for i, p in enumerate(wav_paths[:10]):
         pr = PROMPTS_10[i] if i < len(PROMPTS_10) else PROMPTS_10[-1]
         ex.append([p.as_posix(), pr])
     return ex
 
+
 # ================== Main callback ==================
 def enhance_audio_ui(
     audio_path: str,
@@ -227,7 +338,6 @@ def enhance_audio_ui(
     """
     try:
         prompt = (prompt or "").strip() or "Enhance the input audio"
-
         if not audio_path:
             raise gr.Error("Please upload or select an input audio file.")
 
@@ -237,11 +347,10 @@ def enhance_audio_ui(
 
         tmp_in = SPACE_ROOT / "tmp_in.wav"
         tmp_out = SPACE_ROOT / "tmp_out.wav"
+
+        # Clean previous output to avoid stale returns
         if tmp_out.exists():
-            try:
-                tmp_out.unlink()
-            except Exception:
-                pass
+            _safe_unlink(tmp_out)
 
         save_temp_wav(wav, sr, tmp_in)
 
@@ -266,33 +375,32 @@ def enhance_audio_ui(
         import traceback
         return None, f"Unexpected error: {e}\n{traceback.format_exc()}"
 
+
 # ================== Gradio UI ==================
 with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill_height=True) as _demo:
     gr.Markdown(
         "## 🎧 SonicMaster\n"
         "Upload audio, write a prompt (or leave blank), then click **Enhance**.\n"
         "If left blank, we use: _Enhance the input audio_.\n\n"
-        "- First run will clone the repo + download weights (may take a bit).\n"
+        "- First run will fetch code + download weights.\n"
         "- Subsequent runs are much faster.\n"
-        "If you enjoy this model, please cite the paper."
     )
 
     with gr.Row():
         with gr.Column(scale=1):
             in_audio = gr.Audio(label="Input Audio", type="filepath")
-            prompt_box = gr.Textbox(label="Text Prompt", placeholder="e.g., Reduce reverb and brighten vocals. (Optional)")
+            prompt_box = gr.Textbox(
+                label="Text Prompt",
+                placeholder="e.g., Reduce reverb and brighten vocals. (Optional)",
+            )
             run_btn = gr.Button("🚀 Enhance", variant="primary")
 
-            # Examples only if already present locally (no startup clone)
+            # Examples only if already present locally (no startup fetch)
             examples = build_examples_if_repo_present()
             if examples:
-                gr.Examples(
-                    examples=examples,
-                    inputs=[in_audio, prompt_box],
-                    label="Sample Inputs (10)",
-                )
+                gr.Examples(examples=examples, inputs=[in_audio, prompt_box], label="Sample Inputs (10)")
             else:
-                gr.Markdown("> ℹ️ Samples will appear after the repo is cloned (first run).")
+                gr.Markdown("> ℹ️ Samples will appear after the code is fetched (first run).")
 
         with gr.Column(scale=1):
             out_audio = gr.Audio(label="Enhanced Audio (output)")
@@ -310,4 +418,4 @@ iface = demo
 app = demo
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)