Update app.py

app.py

@@ -1,12 +1,14 @@
-# app.py — Voice Clarity Booster with Presets, Dual-Stage "Ultimate Clean Voice",
-# A/B alternating, Delta (Original−Enhanced), and Loudness Match.
+# app.py — Voice Clarity Booster with Presets, Dual-Stage (smart CPU/GPU guard),
+# A/B alternating, Delta (Original−Enhanced), Loudness Match, and safe fallbacks.
+#
+# Key change: On CPU or for long clips, SepFormer/Dual-Stage auto-falls back to MetricGAN+
+# instead of hanging. Metrics show the fallback reason.
 
 import os
-import io
 import tempfile
 from typing import Tuple, Optional, Dict, Any
 
-# --- Quiet noisy deprecation warnings (optional) ---
+# ---- Quiet noisy deprecation warnings (optional) ----
 import warnings
 warnings.filterwarnings(
     "ignore",
@@ -24,7 +26,7 @@ import soundfile as sf
 import torch
 import torchaudio
 
-# Optional: pyloudnorm for LUFS match; fallback to RMS if not available
+# Optional LUFS matching (falls back to RMS if unavailable)
 try:
     import pyloudnorm as pyln
     _HAVE_PYLN = True
@@ -44,9 +46,15 @@ except Exception:
 
 
 # -----------------------------
-# Cached models
+# Environment / runtime limits
 # -----------------------------
-_DEVICE = "cpu"
+USE_GPU = torch.cuda.is_available()
+# On CPU, SepFormer is extremely slow; avoid for long clips (or disable).
+MAX_SEPFORMER_SEC_CPU = float(os.getenv("MAX_SEPFORMER_SEC_CPU", 12))   # hard limit for CPU
+MAX_SEPFORMER_SEC_GPU = float(os.getenv("MAX_SEPFORMER_SEC_GPU", 180))  # just in case
+ALLOW_SEPFORMER_CPU = os.getenv("ALLOW_SEPFORMER_CPU", "0") == "1"      # override at your risk
+
+_DEVICE = "cuda" if USE_GPU else "cpu"
 _ENHANCER_METRICGAN: Optional[SpectralMaskEnhancement] = None
 _ENHANCER_SEPFORMER: Optional[SepformerSeparation] = None
 
@@ -77,10 +85,7 @@ def _get_sepformer() -> SepformerSeparation:
 # Audio helpers
 # -----------------------------
 def _to_mono(wav: np.ndarray) -> np.ndarray:
-    """
-    Ensure mono [T] float32 robustly.
-    Accepts [T], [T,C], [C,T]; picks the 'channels' axis if <=8.
-    """
+    """Robust mono: accepts [T], [T,C], [C,T]; treats dim<=8 as channels."""
     wav = np.asarray(wav, dtype=np.float32)
     if wav.ndim == 1:
         return wav
@@ -93,7 +98,6 @@ def _to_mono(wav: np.ndarray) -> np.ndarray:
         if t <= 8:   # [C, T]
             return wav.mean(axis=0).astype(np.float32)
         return wav.mean(axis=1).astype(np.float32)
-    # Higher dims: flatten
     return wav.reshape(-1).astype(np.float32)
 
 
@@ -135,16 +139,13 @@ def _align_lengths(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray
 
 
 def _loudness_match_to_ref(ref: np.ndarray, cand: np.ndarray, sr: int) -> Tuple[np.ndarray, str]:
-    """
-    Match cand loudness to ref, returning adjusted signal and a short description.
-    Uses LUFS (pyloudnorm) if available, else RMS.
-    """
+    """Match cand loudness to ref (LUFS if available, else RMS)."""
     if len(ref) < sr // 10 or len(cand) < sr // 10:
         return cand, "skipped (clip too short)"
 
     if _HAVE_PYLN:
         try:
-            meter = pyln.Meter(sr)  # EBUR128 meter
+            meter = pyln.Meter(sr)
             l_ref = meter.integrated_loudness(ref.astype(np.float64))
             l_cand = meter.integrated_loudness(cand.astype(np.float64))
             gain_db = l_ref - l_cand
@@ -153,7 +154,7 @@ def _loudness_match_to_ref(ref: np.ndarray, cand: np.ndarray, sr: int) -> Tuple[
         except Exception:
             pass
 
-    # Fallback: RMS match
+    # RMS fallback
     eps = 1e-9
     rms_ref = np.sqrt(np.mean(ref**2) + eps)
     rms_cand = np.sqrt(np.mean(cand**2) + eps)
@@ -164,15 +165,13 @@ def _loudness_match_to_ref(ref: np.ndarray, cand: np.ndarray, sr: int) -> Tuple[
 
 
 def _make_ab_alternating(orig: np.ndarray, enh: np.ndarray, sr: int, seg_sec: float = 2.0) -> np.ndarray:
-    """
-    Build an A/B track that alternates: seg of Original, seg of Enhanced, repeated.
-    """
+    """A/B track flips Original→Enhanced every seg_sec."""
     seg_n = max(1, int(seg_sec * sr))
     orig, enh = _align_lengths(orig, enh)
     n = len(orig)
     out = []
     pos = 0
-    flag = True  # True=orig, False=enh
+    flag = True
     while pos < n:
         end = min(pos + seg_n, n)
         out.append(orig[pos:end] if flag else enh[pos:end])
@@ -182,7 +181,7 @@ def _make_ab_alternating(orig: np.ndarray, enh: np.ndarray, sr: int, seg_sec: fl
 
 
 # -----------------------------
-# Model runners
+# Model runners (with guards)
 # -----------------------------
 def _run_metricgan(path_16k: str) -> torch.Tensor:
     enh = _get_metricgan()
@@ -191,51 +190,64 @@ def _run_metricgan(path_16k: str) -> torch.Tensor:
     return out
 
 
-def _run_sepformer(path_16k: str) -> torch.Tensor:
-    sep = _get_sepformer()
-    with torch.no_grad():
-        out = sep.separate_file(path=path_16k)
-    if isinstance(out, torch.Tensor):
-        if out.dim() == 1:
-            out = out.unsqueeze(0)
-        elif out.dim() == 2 and out.shape[0] > 1:
-            out = out[:1, :]
-        return out
-    if hasattr(out, "numpy"):
-        t = torch.from_numpy(out.numpy())
-        if t.dim() == 1:
-            t = t.unsqueeze(0)
-        elif t.dim() == 2 and t.shape[0] > 1:
-            t = t[:1, :]
-        return t
-    if isinstance(out, (list, tuple)):
-        t = torch.tensor(out[0] if isinstance(out[0], (np.ndarray, list)) else out, dtype=torch.float32)
-        if t.dim() == 1:
-            t = t.unsqueeze(0)
-        return t
-    raise RuntimeError("Unexpected SepFormer output type")
-
-
-def _run_dual_stage(path_16k: str) -> torch.Tensor:
-    """
-    Ultimate Clean: SepFormer (dereverb/denoise) -> MetricGAN+ (denoise polish).
-    Both at 16 kHz mono.
-    """
-    # Stage 1: SepFormer
-    stage1 = _run_sepformer(path_16k)  # [1, T]
-    # Save Stage 1 to temp wav, then Stage 2 MetricGAN+
+def _run_sepformer(path_16k: str, dur_sec: float) -> Tuple[Optional[torch.Tensor], Optional[str]]:
+    """Return (tensor, fallback_msg). If not safe to run, returns (None, reason)."""
+    if USE_GPU:
+        if dur_sec > MAX_SEPFORMER_SEC_GPU:
+            return None, f"SepFormer skipped (GPU clip {dur_sec:.1f}s > {MAX_SEPFORMER_SEC_GPU:.0f}s limit)"
+    else:
+        if not ALLOW_SEPFORMER_CPU:
+            return None, "SepFormer disabled on CPU (set ALLOW_SEPFORMER_CPU=1 to force)"
+        if dur_sec > MAX_SEPFORMER_SEC_CPU:
+            return None, f"SepFormer skipped (CPU clip {dur_sec:.1f}s > {MAX_SEPFORMER_SEC_CPU:.0f}s limit)"
+
+    try:
+        sep = _get_sepformer()
+        with torch.no_grad():
+            out = sep.separate_file(path=path_16k)
+        if isinstance(out, torch.Tensor):
+            if out.dim() == 1:
+                out = out.unsqueeze(0)
+            elif out.dim() == 2 and out.shape[0] > 1:
+                out = out[:1, :]
+            return out, None
+        if hasattr(out, "numpy"):
+            t = torch.from_numpy(out.numpy())
+            if t.dim() == 1:
+                t = t.unsqueeze(0)
+            elif t.dim() == 2 and t.shape[0] > 1:
+                t = t[:1, :]
+            return t, None
+        if isinstance(out, (list, tuple)):
+            t = torch.tensor(out[0] if isinstance(out[0], (np.ndarray, list)) else out, dtype=torch.float32)
+            if t.dim() == 1:
+                t = t.unsqueeze(0)
+            return t, None
+        return None, "SepFormer returned unexpected format; skipped"
+    except Exception as e:
+        return None, f"SepFormer error: {e.__class__.__name__}"
+
+
+def _run_dual_stage(path_16k: str, dur_sec: float) -> Tuple[Optional[torch.Tensor], Optional[str]]:
+    """SepFormer → MetricGAN+. Applies same guards; returns (tensor, msg)."""
+    stage1, msg = _run_sepformer(path_16k, dur_sec)
+    if stage1 is None:
+        return None, msg or "SepFormer unavailable"
+    # Save stage1 to temp for MetricGAN
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_mid:
         sf.write(tmp_mid.name, stage1.squeeze(0).numpy(), 16000, subtype="PCM_16")
         tmp_mid.flush()
         mid_path = tmp_mid.name
     try:
-        stage2 = _run_metricgan(mid_path)  # [1, T]
+        stage2 = _run_metricgan(mid_path)
+        return stage2, None
+    except Exception as e:
+        return None, f"MetricGAN after SepFormer failed: {e.__class__.__name__}"
     finally:
         try:
             os.remove(mid_path)
         except Exception:
             pass
-    return stage2
 
 
 # -----------------------------
@@ -244,43 +256,52 @@ def _run_dual_stage(path_16k: str) -> torch.Tensor:
 def _enhance_numpy_audio(
     audio: Tuple[int, np.ndarray],
     mode: str = "MetricGAN+ (denoise)",
-    dry_wet: float = 1.0,          # 0..1 (1=fully processed)
-    presence_db: float = 0.0,      # default 0 for safer tone
-    lowcut_hz: float = 0.0,        # default 0 (off)
+    dry_wet: float = 1.0,          # 0..1
+    presence_db: float = 0.0,
+    lowcut_hz: float = 0.0,
     out_sr: Optional[int] = None,
     loudness_match: bool = True,
 ) -> Tuple[int, np.ndarray, np.ndarray, str]:
     """
-    Input: (sr, np.float32 [T] or [T,C])
     Returns: (sr_out, enhanced, delta, metrics_text)
-      - enhanced: final output (after dry/wet, polish, loudness match)
-      - delta: original - enhanced (at output SR & length-matched)
     """
     sr_in, wav_np = audio
     wav_mono = _sanitize(_to_mono(wav_np))
 
-    # Guard: tiny input
     if wav_mono.size < 32:
         sr_out = sr_in if sr_in else 16000
         silence = np.zeros(int(sr_out * 1.0), dtype=np.float32)
         return sr_out, silence, silence, "Input too short; returned silence."
 
     dry_t = torch.from_numpy(wav_mono).unsqueeze(0)  # [1, T @ sr_in]
-    # Prepare 16k mono file for models
     wav_16k = _resample_torch(dry_t, sr_in, 16000)
+    dur_sec = float(wav_16k.shape[-1]) / 16000.0
 
+    # Write temp input for model runners
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_in:
         sf.write(tmp_in.name, wav_16k.squeeze(0).numpy(), 16000, subtype="PCM_16")
         tmp_in.flush()
         path_16k = tmp_in.name
 
+    fallback_note = None
     try:
         if mode.startswith("MetricGAN"):
-            proc = _run_metricgan(path_16k)  # [1, T@16k]
+            proc = _run_metricgan(path_16k)
         elif mode.startswith("SepFormer"):
-            proc = _run_sepformer(path_16k)  # [1, T@16k]
+            proc, msg = _run_sepformer(path_16k, dur_sec)
+            if proc is None:
+                proc = wav_16k  # bypass
+                fallback_note = f"[Fallback→Bypass] {msg}"
         elif mode.startswith("Dual-Stage"):
-            proc = _run_dual_stage(path_16k)  # [1, T@16k]
+            proc, msg = _run_dual_stage(path_16k, dur_sec)
+            if proc is None:
+                # fall back to MetricGAN if SepFormer not possible
+                try:
+                    proc = _run_metricgan(path_16k)
+                    fallback_note = f"[Fallback→MetricGAN+] {msg}"
+                except Exception as e:
+                    proc = wav_16k  # ultimate fallback: bypass
+                    fallback_note = f"[Fallback→Bypass] {msg or ''} / MetricGAN error: {e.__class__.__name__}"
         else:  # Bypass (EQ only)
             proc = wav_16k
     finally:
@@ -289,7 +310,7 @@ def _enhance_numpy_audio(
         except Exception:
             pass
 
-    # Subtle polish (applied to processed only)
+    # Polish on processed only
     proc = _highpass(proc, 16000, lowcut_hz)
     proc = _presence_boost(proc, 16000, presence_db)
     proc = _limit_peak(proc, target_dbfs=-1.0)
@@ -299,34 +320,32 @@ def _enhance_numpy_audio(
     proc_out = _resample_torch(proc, 16000, sr_out).squeeze(0).numpy().astype(np.float32)
     dry_out  = _resample_torch(dry_t, sr_in, sr_out).squeeze(0).numpy().astype(np.float32)
 
-    # Align and mix (dry/wet)
+    # Mix dry/wet
     proc_out, dry_out = _align_lengths(proc_out, dry_out)
     dry_wet = float(np.clip(dry_wet, 0.0, 1.0))
     enhanced = proc_out * dry_wet + dry_out * (1.0 - dry_wet)
 
-    # Loudness match enhanced back to original (optional)
+    # Loudness match
     loud_text = "off"
     if loudness_match:
         enhanced, loud_text = _loudness_match_to_ref(dry_out, enhanced, sr_out)
 
     enhanced = _sanitize(enhanced)
 
-    # Delta (what changed)
-    delta = dry_out - enhanced
-    delta = _sanitize(delta)
+    # Delta
+    delta = _sanitize(dry_out - enhanced)
 
-    # Basic metrics
+    # Metrics
     eps = 1e-9
-    rms_dry = np.sqrt(np.mean(dry_out**2) + eps)
-    rms_enh = np.sqrt(np.mean(enhanced**2) + eps)
     rms_delta = np.sqrt(np.mean(delta**2) + eps)
-    change_db = 20 * np.log10((rms_dry + eps) / (rms_delta + eps))
     metrics = (
         f"Mode: {mode} | Dry/Wet: {dry_wet*100:.0f}% | Presence: {presence_db:+.1f} dB | "
-        f"Low-cut: {lowcut_hz:.0f} Hz | Loudness match: {loud_text}\n"
-        f"Dur: {len(enhanced)/sr_out:.2f}s | Δ RMS: {20*np.log10(rms_delta+eps):+.2f} dBFS | "
-        f'Approx. "noise removed" ratio: {change_db:.2f} dB'
+        f"Low-cut: {lowcut_hz:.0f} Hz | Loudness match: {loud_text} | Device: {'GPU' if USE_GPU else 'CPU'} | "
+        f"Clip @16k: {dur_sec:.2f}s"
     )
+    if fallback_note:
+        metrics += f"\n{fallback_note}"
+    metrics += f"\nΔ RMS: {20*np.log10(rms_delta+eps):+.2f} dBFS"
 
     return sr_out, enhanced, delta, metrics
 
@@ -335,7 +354,6 @@ def _enhance_numpy_audio(
 # Presets
 # -----------------------------
 PRESETS: Dict[str, Dict[str, Any]] = {
-    # Maximum cleanup: dereverb + denoise chain, high dry/wet, subtle presence, mild HPF
     "Ultimate Clean Voice": {
         "mode": "Dual-Stage (SepFormer → MetricGAN+)",
         "dry_wet": 0.92,
@@ -343,7 +361,6 @@ PRESETS: Dict[str, Dict[str, Any]] = {
         "lowcut_hz": 80.0,
         "loudness_match": True,
     },
-    # Natural cleanup for most cases
     "Natural Speech": {
         "mode": "MetricGAN+ (denoise)",
         "dry_wet": 0.85,
@@ -351,31 +368,27 @@ PRESETS: Dict[str, Dict[str, Any]] = {
         "lowcut_hz": 50.0,
         "loudness_match": True,
     },
-    # Studio-ish clarity
     "Podcast Studio": {
         "mode": "MetricGAN+ (denoise)",
-        "dry_wet": 0.9,
+        "dry_wet": 0.90,
         "presence_db": 2.0,
         "lowcut_hz": 75.0,
         "loudness_match": True,
     },
-    # Strong dereverb, blend to avoid artifacts
     "Room Dereverb": {
         "mode": "SepFormer (dereverb+denoise)",
-        "dry_wet": 0.7,
+        "dry_wet": 0.70,
         "presence_db": 0.5,
         "lowcut_hz": 60.0,
         "loudness_match": True,
     },
-    # When music bed is under voice—be gentle
     "Music + Voice Safe": {
         "mode": "MetricGAN+ (denoise)",
-        "dry_wet": 0.6,
+        "dry_wet": 0.60,
         "presence_db": 0.0,
         "lowcut_hz": 40.0,
         "loudness_match": True,
     },
-    # Harsh phone/zoom recordings
     "Phone Call Rescue": {
         "mode": "MetricGAN+ (denoise)",
         "dry_wet": 0.88,
@@ -383,7 +396,6 @@ PRESETS: Dict[str, Dict[str, Any]] = {
         "lowcut_hz": 100.0,
         "loudness_match": True,
     },
-    # Light touch
     "Gentle Denoise": {
         "mode": "MetricGAN+ (denoise)",
         "dry_wet": 0.65,
@@ -391,13 +403,12 @@ PRESETS: Dict[str, Dict[str, Any]] = {
         "lowcut_hz": 0.0,
         "loudness_match": True,
     },
-    "Custom": {}  # no-op, keeps current settings
+    "Custom": {}
 }
 
 
 def _apply_preset(preset_name: str):
     cfg = PRESETS.get(preset_name, {})
-    # Return gr.update() for each adjustable control
     def upd(val=None):
         return gr.update(value=val) if val is not None else gr.update()
     if not cfg or preset_name == "Custom":
@@ -447,7 +458,12 @@ def gradio_enhance(
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## Voice Clarity Booster — Presets, A/B, Delta, Loudness Match")
+    gr.Markdown(
+        f"## Voice Clarity Booster — Presets, A/B, Delta, Loudness Match  \n"
+        f"**Device:** {'GPU' if USE_GPU else 'CPU'}  ·  "
+        f"SepFormer limits — CPU≤{MAX_SEPFORMER_SEC_CPU:.0f}s, GPU≤{MAX_SEPFORMER_SEC_GPU:.0f}s"
+        + ("" if USE_GPU or ALLOW_SEPFORMER_CPU else "  ·  (SepFormer disabled on CPU)")
+    )
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -462,7 +478,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Preset",
             )
 
-            # Controls that presets will adjust
             mode = gr.Radio(
                 choices=[
                     "MetricGAN+ (denoise)",
@@ -490,7 +505,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Output Sample Rate",
             )
 
-            # Apply preset on change
             preset.change(
                 _apply_preset,
                 inputs=[preset],
@@ -511,5 +525,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[out_audio, ab_audio, delta_audio, metrics],
     )
 
-# Start server (Hugging Face Spaces expects this unguarded)
+# Launch unguarded so Spaces initializes
 demo.launch()