Update app.py

app.py

@@ -1,8 +1,8 @@
-# ========== MUST BE FIRST: Gradio SDK entry + ZeroGPU probes ==========
+# ========= MUST BE FIRST: Gradio entry + ZeroGPU probes =========
 import os
 os.environ.setdefault("GRADIO_USE_CDN", "true")
 
-# Optional: 'spaces' present on Spaces; harmless to try locally.
+# 'spaces' is present on HF Spaces; make it safe locally too
 try:
     import spaces
 except Exception:
@@ -12,7 +12,7 @@ except Exception:
             return deco
     spaces = _DummySpaces()
 
-# PUBLIC names so ZeroGPU supervisor can detect them
+# Publicly-named probes so ZeroGPU supervisor can detect them
 @spaces.GPU(duration=10)
 def gpu_probe(a: int = 1, b: int = 1):
     return a + b
@@ -21,18 +21,18 @@ def gpu_probe(a: int = 1, b: int = 1):
 def gpu_echo(x: str = "ok"):
     return x
 
-# ========== Standard imports ==========
+# ================= Standard imports =================
 import sys
 import subprocess
 from pathlib import Path
-from typing import Tuple, Optional, List, Dict, Any
+from typing import Tuple, Optional, List, Any
 
 import gradio as gr
 import numpy as np
 import soundfile as sf
 from huggingface_hub import hf_hub_download
 
-# ZeroGPU runtime hint (safe on CPU)
+# Runtime hints (safe on CPU)
 USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"
 
 SPACE_ROOT   = Path(__file__).parent.resolve()
@@ -43,7 +43,7 @@ WEIGHTS_FILE = "model.safetensors"
 CACHE_DIR    = SPACE_ROOT / "weights"
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-# ========== Lazy resources (no heavy work at import) ==========
+# ================ Lazy resources =================
 _weights_path: Optional[Path] = None
 _repo_ready: bool = False
 
@@ -78,9 +78,8 @@ def ensure_repo(progress: Optional[gr.Progress] = None) -> Path:
         _repo_ready = True
     return REPO_DIR
 
-# ========== Helpers ==========
+# ================== 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:
@@ -139,8 +138,8 @@ def run_sonicmaster_cli(
                 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=60)
+# ============ GPU path (ZeroGPU) ============
+@spaces.GPU(duration=60)  # 60s is a safe cap for ZeroGPU
 def enhance_on_gpu(input_path: str, prompt: str, output_path: str) -> Tuple[bool, str]:
     try:
         import torch  # noqa: F401
@@ -156,7 +155,7 @@ def _has_cuda() -> bool:
     except Exception:
         return False
 
-# ========== Examples (lazy) ==========
+# ================== Examples (lazy) ==================
 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.",
@@ -177,28 +176,22 @@ def list_example_files(progress: Optional[gr.Progress] = None) -> List[str]:
     files = sorted(p for p in wav_dir.glob("*.wav") if p.is_file())
     return [p.as_posix() for p in files[:10]]
 
-def load_examples(_: Any = None, progress=gr.Progress()) -> Dict[str, Any]:
-    """Button/auto-load handler: populate dropdown choices and status text."""
+def load_examples(_: Any = None, progress=gr.Progress()):
+    """
+    Returns (choices:list[str], paths:list[str], status:str)
+    """
     paths = list_example_files(progress=progress)
     if not paths:
-        return {
-            "choices": [],
-            "status": "No sample .wav files found in repo/samples/inputs.",
-        }
-    labels = [f"{i+1:02d} — {Path(p).name}" for i, p in enumerate(paths)]
-    return {
-        "choices": labels,
-        "paths": paths,
-        "status": f"Loaded {len(paths)} sample audios."
-    }
+        return [], [], "No sample .wav files found in repo/samples/inputs."
+    choices = [f"{i+1:02d} — {Path(p).name}" for i, p in enumerate(paths)]
+    return choices, 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:
-        # label "01 — file.wav" -> index 0
-        idx = int(idx_label.split()[0]) - 1
+        idx = int(idx_label.split()[0]) - 1  # "01 — file.wav" -> 0
     except Exception:
         idx = 0
     idx = max(0, min(idx, len(paths)-1))
@@ -206,12 +199,12 @@ def set_example_selection(idx_label: str, paths: List[str]) -> Tuple[str, str]:
     prompt = PROMPTS_10[idx] if idx < len(PROMPTS_10) else PROMPTS_10[-1]
     return audio_path, prompt
 
-# ========== Gradio callback ==========
+# ================== Main 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.
     """
@@ -251,12 +244,13 @@ def enhance_audio_ui(
         import traceback
         return None, f"Unexpected error: {e}\n{traceback.format_exc()}"
 
-# ========== Gradio UI ==========
+# ================== Gradio 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"
-        "- On failure, the **Status** box shows the exact error "
+        "- 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(scale=1):
@@ -270,7 +264,6 @@ with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill
             prompt   = gr.Textbox(label="Text Prompt", placeholder="e.g., Reduce reverb and brighten vocals.")
             run_btn  = gr.Button("🚀 Enhance", variant="primary")
 
-            # Optional quick prompt examples (text-only)
             gr.Examples(
                 examples=[[p] for p in [
                     "Reduce roominess/echo (dereverb).",
@@ -285,26 +278,20 @@ with gr.Blocks(title="SonicMaster – Text-Guided Restoration & Mastering", fill
             out_audio = gr.Audio(label="Enhanced Audio (output)")
             status    = gr.Textbox(label="Status / Messages", interactive=False, lines=8)
 
-    # --- Wire up the sample loader ---
-    # 1) Load samples on button click (lazy clone)
-    load_result = load_btn.click(
+    # Load samples (3 outputs directly; no .then needed)
+    load_btn.click(
         fn=load_examples,
         inputs=None,
-        outputs=None
+        outputs=[samples_dropdown, samples_state, status],
     )
-    # Manually map the dict result to components via .then (Gradio v5 API)
-    load_result.then(lambda d: d.get("choices", []), None, samples_dropdown)
-    load_result.then(lambda d: d.get("paths", []),   None, samples_state)
-    load_result.then(lambda d: d.get("status", ""),  None, status)
 
-    # 2) When a sample is chosen, set audio path + suggested prompt
+    # 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],
     )
 
-    # --- Enhance button ---
     run_btn.click(
         fn=enhance_audio_ui,
         inputs=[in_audio, prompt],
@@ -319,4 +306,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)