Update app.py

app.py

@@ -12,7 +12,7 @@ except Exception:
             return deco
     spaces = _DummySpaces()
 
-# Publicly-named probes so ZeroGPU supervisor can detect them
+# Public probes so ZeroGPU supervisor can detect them
 @spaces.GPU(duration=10)
 def gpu_probe(a: int = 1, b: int = 1):
     return a + b
@@ -43,12 +43,24 @@ WEIGHTS_FILE = "model.safetensors"
 CACHE_DIR    = SPACE_ROOT / "weights"
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-# ================ Lazy resources =================
-_weights_path: Optional[Path] = None
-_repo_ready: bool = False
+# ================ Repo clone AT STARTUP (per your request) ================
+def ensure_repo() -> Path:
+    if not REPO_DIR.exists():
+        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())
+    return REPO_DIR
 
+# Clone now so examples are available immediately
+ensure_repo()
+
+# ================ Weights: still lazy (download at first run) ================
+_weights_path: Optional[Path] = None
 def get_weights_path(progress: Optional[gr.Progress] = None) -> Path:
-    """Download/resolve weights lazily."""
+    """Download/resolve weights lazily (keeps startup fast)."""
     global _weights_path
     if _weights_path is None:
         if progress: progress(0.10, desc="Downloading model weights (first run)")
@@ -63,23 +75,9 @@ def get_weights_path(progress: Optional[gr.Progress] = None) -> Path:
         _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:
@@ -93,11 +91,15 @@ def read_audio(path: str) -> Tuple[np.ndarray, int]:
     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
+    """
+    Only support infer_single.py variants.
+    Expected primary flags: --ckpt --input --prompt --output
+    """
     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()],
+        [py, script.as_posix(), "--ckpt", ckpt.as_posix(), "--input", inp.as_posix(), "--prompt", prompt, "--output", out.as_posix()],
+        # A couple common alternates some repos use:
+        [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(
@@ -106,40 +108,37 @@ def run_sonicmaster_cli(
     out_path: Path,
     progress: Optional[gr.Progress] = None,
 ) -> Tuple[bool, str]:
-    """Run inference scripts via subprocess; return (ok, message)."""
+    """Run inference via subprocess; returns (ok, message). Uses ONLY infer_single.py."""
     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)."
+    script = REPO_DIR / "infer_single.py"
+    if not script.exists():
+        return False, "infer_single.py not found in the SonicMaster repo."
 
     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()}"
+    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 infer_single.py (try {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 = "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()
+            last_err = snippet if snippet else f"infer_single.py failed with return code {e.returncode}."
+        except Exception as e:
+            import traceback
+            last_err = f"Unexpected error with infer_single.py: {e}\n{traceback.format_exc()}"
     return False, last_err or "All candidate commands failed."
 
 # ============ GPU path (ZeroGPU) ============
-@spaces.GPU(duration=60)  # 60s is a safe cap for 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]:
     try:
         import torch  # noqa: F401
@@ -155,7 +154,7 @@ def _has_cuda() -> bool:
     except Exception:
         return False
 
-# ================== Examples (lazy) ==================
+# ================== Examples @ 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.",
@@ -169,36 +168,17 @@ PROMPTS_10 = [
     "Please, dereverb this audio.",
 ]
 
-def list_example_files(progress: Optional[gr.Progress] = None) -> List[str]:
-    """Return up to 10 .wav paths inside repo/samples/inputs (lazy clone)."""
-    repo = ensure_repo(progress=progress)
-    wav_dir = repo / "samples" / "inputs"
-    files = sorted(p for p in wav_dir.glob("*.wav") if p.is_file())
-    return [p.as_posix() for p in files[:10]]
+def build_startup_examples() -> List[List[Any]]:
+    """Build 10 (audio_path, prompt) pairs from repo at import time."""
+    wav_dir = REPO_DIR / "samples" / "inputs"
+    wav_paths = sorted(p for p in wav_dir.glob("*.wav") if p.is_file())
+    ex = []
+    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
 
-def load_examples(_: Any = None, progress=gr.Progress()):
-    """
-    Returns (dropdown_update, paths:list[str], status:str)
-    """
-    paths = list_example_files(progress=progress)
-    if not paths:
-        return gr.Dropdown.update(choices=[], value=None), [], "No sample .wav files found in repo/samples/inputs."
-    labels = [f"{i+1:02d} — {Path(p).name}" for i, p in enumerate(paths)]
-    # Auto-select first item for convenience
-    return gr.Dropdown.update(choices=labels, value=labels[0]), paths, f"Loaded {len(paths)} sample audios."
-
-def set_example_selection(idx_label: str, paths: List[str]) -> Tuple[str, str]:
-    """When user picks an example, set the audio path + a suggested prompt."""
-    if not idx_label or not paths:
-        return "", ""
-    try:
-        idx = int(idx_label.split()[0]) - 1  # "01 — file.wav" -> 0
-    except Exception:
-        idx = 0
-    idx = max(0, min(idx, len(paths)-1))
-    audio_path = paths[idx]
-    prompt = PROMPTS_10[idx] if idx < len(PROMPTS_10) else PROMPTS_10[-1]
-    return audio_path, prompt
+STARTUP_EXAMPLES = build_startup_examples()
 
 # ================== Main callback ==================
 def enhance_audio_ui(
@@ -249,52 +229,29 @@ def enhance_audio_ui(
 with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill_height=True) as _demo:
     gr.Markdown(
         "## 🎧 SonicMaster\n"
-        "Upload audio or **load sample audios**, write a prompt, then click **Enhance**.\n"
+        "Upload audio or pick an example, write a prompt, then click **Enhance**.\n"
+        "- First run downloads model weights (progress will show).\n"
     )
     with gr.Row():
         with gr.Column(scale=1):
-            # Sample loader (lazy)
-            with gr.Accordion("Sample audios (10)", open=False):
-                load_btn = gr.Button("📥 Load 10 sample audios")
-                samples_dropdown = gr.Dropdown(
-                    choices=[], value=None,  # no default until choices are set
-                    label="Pick a sample",
-                    interactive=True,
-                )
-                samples_state = gr.State([])  # holds absolute paths
-
             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=[[p] for p in [
-                    "Reduce roominess/echo (dereverb).",
-                    "Raise the level of the vocals.",
-                    "Give the song a wider stereo image.",
-                ]],
-                inputs=[prompt],
-                label="Prompt Examples",
-            )
+            # Show 10 audio+prompt examples immediately at startup
+            if STARTUP_EXAMPLES:
+                gr.Examples(
+                    examples=STARTUP_EXAMPLES,
+                    inputs=[in_audio, prompt],
+                    label="Sample Inputs (10)",
+                )
+            else:
+                gr.Markdown("> ⚠️ No sample .wav files found in `samples/inputs/`.")
 
         with gr.Column(scale=1):
             out_audio = gr.Audio(label="Enhanced Audio (output)")
             status    = gr.Textbox(label="Status / Messages", interactive=False, lines=8)
 
-    # Load samples (3 outputs directly; no .then needed)
-    load_btn.click(
-        fn=load_examples,
-        inputs=None,
-        outputs=[samples_dropdown, samples_state, status],
-    )
-
-    # When a sample is chosen, set audio path + suggested prompt
-    samples_dropdown.change(
-        fn=set_example_selection,
-        inputs=[samples_dropdown, samples_state],
-        outputs=[in_audio, prompt],
-    )
-
     run_btn.click(
         fn=enhance_audio_ui,
         inputs=[in_audio, prompt],
@@ -309,4 +266,4 @@ app = demo
 
 # Local debugging only
 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)