Update app.py

app.py

@@ -1,8 +1,10 @@
-# app.py — Voice Clarity Booster with Presets, Dual-Stage (smart CPU/GPU guard),
-# A/B alternating, Delta (Original−Enhanced), Loudness Match, and safe fallbacks.
+# app.py — Voice Clarity Booster with Presets, CPU/GPU-smart Dual-Stage,
+# A/B alternating, Loudness Match, and a *polished Delta* (noise-only) option.
 #
-# Key change: On CPU or for long clips, SepFormer/Dual-Stage auto-falls back to MetricGAN+
-# instead of hanging. Metrics show the fallback reason.
+# New:
+# - Delta Mode: Raw Difference | Spectral Residual (noise-only)
+# - Delta Alignment (cross-correlation) to reduce phase/latency smear
+# - Delta Gain (dB) + HPF/LPF + RMS leveling for listenable delta
 
 import os
 import tempfile
@@ -50,9 +52,9 @@ except Exception:
 # -----------------------------
 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
+MAX_SEPFORMER_SEC_CPU = float(os.getenv("MAX_SEPFORMER_SEC_CPU", 12))
+MAX_SEPFORMER_SEC_GPU = float(os.getenv("MAX_SEPFORMER_SEC_GPU", 180))
+ALLOW_SEPFORMER_CPU = os.getenv("ALLOW_SEPFORMER_CPU", "0") == "1"
 
 _DEVICE = "cuda" if USE_GPU else "cpu"
 _ENHANCER_METRICGAN: Optional[SpectralMaskEnhancement] = None
@@ -117,6 +119,12 @@ def _highpass(wav: torch.Tensor, sr: int, cutoff_hz: float) -> torch.Tensor:
     return torchaudio.functional.highpass_biquad(wav, sr, cutoff_hz)
 
 
+def _lowpass(wav: torch.Tensor, sr: int, cutoff_hz: float) -> torch.Tensor:
+    if cutoff_hz is None or cutoff_hz <= 0:
+        return wav
+    return torchaudio.functional.lowpass_biquad(wav, sr, cutoff_hz)
+
+
 def _presence_boost(wav: torch.Tensor, sr: int, gain_db: float) -> torch.Tensor:
     if abs(gain_db) < 1e-6:
         return wav
@@ -138,6 +146,20 @@ def _align_lengths(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray
     return a[:n], b[:n]
 
 
+def _rms(x: np.ndarray, eps: float = 1e-9) -> float:
+    return float(np.sqrt(np.mean(x**2) + eps))
+
+
+def _rms_target(x: np.ndarray, target_dbfs: float = -20.0) -> np.ndarray:
+    """Scale to approx target dBFS RMS, then hard-limit peaks."""
+    target_amp = 10.0 ** (target_dbfs / 20.0)
+    cur = _rms(x)
+    if cur > 0:
+        x = x * (target_amp / cur)
+    x = np.clip(x, -1.0, 1.0).astype(np.float32)
+    return x
+
+
 def _loudness_match_to_ref(ref: np.ndarray, cand: np.ndarray, sr: int) -> Tuple[np.ndarray, str]:
     """Match cand loudness to ref (LUFS if available, else RMS)."""
     if len(ref) < sr // 10 or len(cand) < sr // 10:
@@ -180,6 +202,59 @@ def _make_ab_alternating(orig: np.ndarray, enh: np.ndarray, sr: int, seg_sec: fl
     return np.concatenate(out, axis=0).astype(np.float32)
 
 
+# -----------------------------
+# Alignment for delta (cross-correlation)
+# -----------------------------
+def _next_pow_two(n: int) -> int:
+    n -= 1
+    shift = 1
+    while (n + 1) & n:
+        n |= n >> shift
+        shift <<= 1
+    return n + 1
+
+
+def _align_by_xcorr(a: np.ndarray, b: np.ndarray, max_shift: int) -> Tuple[np.ndarray, np.ndarray, int]:
+    """
+    Align b to a using FFT cross-correlation. Only accept shifts within ±max_shift.
+    Returns (a_aligned, b_aligned, shift) where positive shift means b lags a and is shifted forward.
+    """
+    # Pad to same length
+    n = max(len(a), len(b))
+    a_pad = np.zeros(n, dtype=np.float32); a_pad[:len(a)] = a
+    b_pad = np.zeros(n, dtype=np.float32); b_pad[:len(b)] = b
+
+    N = _next_pow_two(2 * n - 1)
+    A = np.fft.rfft(a_pad, N)
+    B = np.fft.rfft(b_pad, N)
+    corr = np.fft.irfft(A * np.conj(B), N)
+    # lags: 0..N-1, convert so center at zero lag
+    corr = np.concatenate((corr[-(n-1):], corr[:n]))
+    lags = np.arange(-(n-1), n)
+
+    # Limit to window
+    w = (lags >= -max_shift) & (lags <= max_shift)
+    lag = int(lags[w][np.argmax(corr[w])])
+
+    if lag > 0:
+        # b lags behind a -> shift b forward
+        b_shift = np.concatenate((b[lag:], np.zeros(lag, dtype=np.float32)))
+        a_shift = a[:len(b_shift)]
+        b_shift = b_shift[:len(a_shift)]
+        return a_shift, b_shift, lag
+    elif lag < 0:
+        # a lags -> shift a forward
+        lag = -lag
+        a_shift = np.concatenate((a[lag:], np.zeros(lag, dtype=np.float32)))
+        b_shift = b[:len(a_shift)]
+        a_shift = a_shift[:len(b_shift)]
+        return a_shift, b_shift, -lag
+    else:
+        # no shift
+        a2, b2 = _align_lengths(a, b)
+        return a2, b2, 0
+
+
 # -----------------------------
 # Model runners (with guards)
 # -----------------------------
@@ -250,6 +325,50 @@ def _run_dual_stage(path_16k: str, dur_sec: float) -> Tuple[Optional[torch.Tenso
             pass
 
 
+# -----------------------------
+# Spectral residual delta (cleaner noise-only preview)
+# -----------------------------
+def _delta_spectral_residual(orig: np.ndarray, enh: np.ndarray, sr: int) -> np.ndarray:
+    """
+    Build a noise-focused residual via STFT magnitudes:
+      R_mag = ReLU(|X| - |Y|)
+      use original phase for iSTFT reconstruction
+    Then gentle HPF/LPF and RMS leveling for listenability.
+    """
+    # Torch tensors
+    x = torch.from_numpy(orig).to(torch.float32)
+    y = torch.from_numpy(enh).to(torch.float32)
+
+    n_fft = 1024
+    hop = 256
+    win = torch.hann_window(n_fft)
+
+    # STFTs
+    X = torch.stft(x, n_fft=n_fft, hop_length=hop, window=win, return_complex=True, center=True)
+    Y = torch.stft(y, n_fft=n_fft, hop_length=hop, window=win, return_complex=True, center=True)
+
+    # Positive residual magnitudes
+    R_mag = torch.relu(torch.abs(X) - torch.abs(Y))
+
+    # Mild temporal smoothing (moving average across time)
+    R_mag = torch.nn.functional.avg_pool1d(
+        R_mag.unsqueeze(0), kernel_size=3, stride=1, padding=1
+    ).squeeze(0)
+
+    # Reconstruct residual with original phase
+    phase = torch.angle(X)
+    R_complex = torch.polar(R_mag, phase)
+    r = torch.istft(R_complex, n_fft=n_fft, hop_length=hop, window=win, length=len(orig))
+
+    # HPF/LPF + light RMS leveling for comfort
+    r_t = r.unsqueeze(0)
+    r_t = _highpass(r_t, sr, cutoff_hz=80.0)
+    r_t = _lowpass(r_t, sr, cutoff_hz=9000.0)
+    r_np = r_t.squeeze(0).numpy().astype(np.float32)
+    r_np = _rms_target(r_np, target_dbfs=-24.0)
+    return r_np
+
+
 # -----------------------------
 # Core pipeline
 # -----------------------------
@@ -261,9 +380,9 @@ def _enhance_numpy_audio(
     lowcut_hz: float = 0.0,
     out_sr: Optional[int] = None,
     loudness_match: bool = True,
-) -> Tuple[int, np.ndarray, np.ndarray, str]:
+) -> Tuple[int, np.ndarray, str]:
     """
-    Returns: (sr_out, enhanced, delta, metrics_text)
+    Returns: (sr_out, enhanced, metrics_text)
     """
     sr_in, wav_np = audio
     wav_mono = _sanitize(_to_mono(wav_np))
@@ -271,7 +390,7 @@ def _enhance_numpy_audio(
     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."
+        return sr_out, silence, "Input too short; returned silence."
 
     dry_t = torch.from_numpy(wav_mono).unsqueeze(0)  # [1, T @ sr_in]
     wav_16k = _resample_torch(dry_t, sr_in, 16000)
@@ -332,12 +451,9 @@ def _enhance_numpy_audio(
 
     enhanced = _sanitize(enhanced)
 
-    # Delta
-    delta = _sanitize(dry_out - enhanced)
-
     # Metrics
     eps = 1e-9
-    rms_delta = np.sqrt(np.mean(delta**2) + eps)
+    rms_delta_hint = np.sqrt(np.mean((dry_out - enhanced)**2) + 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} | Device: {'GPU' if USE_GPU else 'CPU'} | "
@@ -345,9 +461,9 @@ def _enhance_numpy_audio(
     )
     if fallback_note:
         metrics += f"\n{fallback_note}"
-    metrics += f"\nΔ RMS: {20*np.log10(rms_delta+eps):+.2f} dBFS"
+    metrics += f"\nΔ (raw) RMS: {20*np.log10(rms_delta_hint+eps):+.2f} dBFS"
 
-    return sr_out, enhanced, delta, metrics
+    return sr_out, enhanced, metrics
 
 
 # -----------------------------
@@ -433,13 +549,18 @@ def gradio_enhance(
     lowcut_hz: float,
     output_sr: str,
     loudness_match: bool,
+    delta_mode: str,
+    delta_align: bool,
+    delta_gain_db: float,
 ):
     if audio is None:
         return None, None, None, "No audio provided."
     out_sr = None
     if output_sr in {"44100", "48000"}:
         out_sr = int(output_sr)
-    sr_out, enhanced, delta, metrics = _enhance_numpy_audio(
+
+    # Enhance
+    sr_out, enhanced, metrics = _enhance_numpy_audio(
         audio,
         mode=mode,
         dry_wet=dry_wet_pct / 100.0,
@@ -448,18 +569,44 @@ def gradio_enhance(
         out_sr=out_sr,
         loudness_match=bool(loudness_match),
     )
-    # Build A/B alternating track
+
+    # Build A/B and Delta (polished)
     sr_in, wav_np = audio
     orig_mono = _sanitize(_to_mono(wav_np))
     orig_at_out = _resample_torch(torch.from_numpy(orig_mono).unsqueeze(0), sr_in, sr_out).squeeze(0).numpy().astype(np.float32)
-    orig_at_out, enhanced = _align_lengths(orig_at_out, enhanced)
-    ab_alt = _make_ab_alternating(orig_at_out, enhanced, sr_out, seg_sec=2.0)
+
+    # Optional alignment to reduce phase/latency offsets
+    a_for_ab, b_for_ab = _align_lengths(orig_at_out, enhanced)
+    if delta_align:
+        max_shift = int(0.05 * sr_out)  # up to 50 ms
+        a_for_ab, b_for_ab, lag = _align_by_xcorr(a_for_ab, b_for_ab, max_shift=max_shift)
+        metrics += f"\nDelta alignment: shift={lag} samples"
+
+    # A/B alternating
+    ab_alt = _make_ab_alternating(a_for_ab, b_for_ab, sr_out, seg_sec=2.0)
+
+    # Delta (noise-focused if selected)
+    if delta_mode.startswith("Spectral"):
+        delta = _delta_spectral_residual(a_for_ab, b_for_ab, sr_out)
+    else:
+        delta = a_for_ab - b_for_ab
+        # Gentle polish on raw difference
+        d_t = torch.from_numpy(delta).unsqueeze(0)
+        d_t = _highpass(d_t, sr_out, cutoff_hz=80.0)
+        d_t = _lowpass(d_t, sr_out, cutoff_hz=9000.0)
+        delta = d_t.squeeze(0).numpy().astype(np.float32)
+        delta = _rms_target(delta, target_dbfs=-24.0)
+
+    # Apply user delta gain
+    delta *= 10.0 ** (delta_gain_db / 20.0)
+    delta = np.clip(delta, -1.0, 1.0).astype(np.float32)
+
     return (sr_out, enhanced), (sr_out, ab_alt), (sr_out, delta), metrics
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
-        f"## Voice Clarity Booster — Presets, A/B, Delta, Loudness Match  \n"
+        f"## Voice Clarity Booster — Presets, A/B, *Polished 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)")
@@ -505,6 +652,16 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Output Sample Rate",
             )
 
+            # Delta controls
+            gr.Markdown("### Delta (what changed)")
+            delta_mode = gr.Dropdown(
+                choices=["Spectral Residual (noise-only)", "Raw Difference"],
+                value="Spectral Residual (noise-only)",
+                label="Delta Mode",
+            )
+            delta_align = gr.Checkbox(value=True, label="Align original & enhanced for delta (recommended)")
+            delta_gain = gr.Slider(minimum=-12, maximum=24, value=6, step=1, label="Delta Gain (dB)")
+
             preset.change(
                 _apply_preset,
                 inputs=[preset],
@@ -516,12 +673,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=1):
             out_audio = gr.Audio(type="numpy", label="Enhanced (autoplay)", autoplay=True)
             ab_audio = gr.Audio(type="numpy", label="A/B Alternating (2s O → 2s E)")
-            delta_audio = gr.Audio(type="numpy", label="Delta: Original − Enhanced")
+            delta_audio = gr.Audio(type="numpy", label="Delta (polished)")
             metrics = gr.Markdown("")
 
     btn.click(
         gradio_enhance,
-        inputs=[in_audio, mode, dry_wet, presence, lowcut, out_sr, loudmatch],
+        inputs=[in_audio, mode, dry_wet, presence, lowcut, out_sr, loudmatch, delta_mode, delta_align, delta_gain],
         outputs=[out_audio, ab_audio, delta_audio, metrics],
     )