Create main.py

main.py

@@ -0,0 +1,204 @@
+import io
+import os
+import tempfile
+from typing import Tuple, Optional
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio
+from fastapi import FastAPI, File, UploadFile, Query, Response
+from fastapi.responses import StreamingResponse
+from speechbrain.pretrained import SpectralMaskEnhancement
+
+# -----------------------------
+# Model: SpeechBrain MetricGAN+
+# -----------------------------
+_ENHANCER: Optional[SpectralMaskEnhancement] = None
+_DEVICE = "cpu"
+
+
+def _get_enhancer() -> SpectralMaskEnhancement:
+    global _ENHANCER
+    if _ENHANCER is None:
+        # Downloads once and caches in the Space
+        _ENHANCER = SpectralMaskEnhancement.from_hparams(
+            source="speechbrain/metricgan-plus-voicebank",
+            savedir="pretrained/metricgan_plus_voicebank",
+            run_opts={"device": _DEVICE},
+        )
+    return _ENHANCER
+
+
+# -----------------------------
+# Audio helpers
+# -----------------------------
+def _to_mono(wav: np.ndarray) -> np.ndarray:
+    """Ensure mono shape [T]."""
+    if wav.ndim == 1:
+        return wav.astype(np.float32)
+    # shape [T, C] or [C, T]
+    if wav.shape[0] < wav.shape[1]:
+        # likely [T, C]
+        return wav.mean(axis=1).astype(np.float32)
+    else:
+        # likely [C, T]
+        return wav.mean(axis=0).astype(np.float32)
+
+
+def _resample_torch(wav: torch.Tensor, sr_in: int, sr_out: int) -> torch.Tensor:
+    if sr_in == sr_out:
+        return wav
+    return torchaudio.functional.resample(wav, sr_in, sr_out)
+
+
+def _highpass(wav: torch.Tensor, sr: int, cutoff_hz: float) -> torch.Tensor:
+    if cutoff_hz is None or cutoff_hz <= 0:
+        return wav
+    # 2nd-order Butterworth-ish highpass via biquad
+    return torchaudio.functional.highpass_biquad(wav, sr, cutoff_hz)
+
+
+def _presence_boost(wav: torch.Tensor, sr: int, gain_db: float) -> torch.Tensor:
+    """Simple presence (peaking) EQ around 4.5 kHz."""
+    if abs(gain_db) < 1e-6:
+        return wav
+    center = 4500.0  # presence band
+    q = 0.707       # wide-ish
+    return torchaudio.functional.equalizer_biquad(wav, sr, center, q, gain_db)
+
+
+def _limit_peak(wav: torch.Tensor, target_dbfs: float = -1.0) -> torch.Tensor:
+    """Peak-normalize to target dBFS (default -1 dB)."""
+    target_amp = 10.0 ** (target_dbfs / 20.0)
+    peak = torch.max(torch.abs(wav)).item()
+    if peak > 0:
+        scale = min(1.0, target_amp / peak)
+        wav = wav * scale
+    return torch.clamp(wav, -1.0, 1.0)
+
+
+def _enhance_numpy_audio(
+    audio: Tuple[int, np.ndarray],
+    presence_db: float = 3.0,
+    lowcut_hz: float = 75.0,
+    out_sr: Optional[int] = None,
+) -> Tuple[int, np.ndarray]:
+    """
+    Core pipeline used by both Gradio UI and raw FastAPI route.
+    Input: (sr, np.float32 [T] or [T,C])
+    Returns: (sr_out, np.float32 [T])
+    """
+    sr_in, wav_np = audio
+    wav_mono = _to_mono(wav_np)
+    wav_t = torch.from_numpy(wav_mono).unsqueeze(0)  # [1, T]
+
+    # MetricGAN+ expects 16 kHz mono
+    enh = _get_enhancer()
+    wav_16k = _resample_torch(wav_t, sr_in, 16000)
+
+    # Enhance via file path API for maximum compatibility
+    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()
+        # Enhance; returns torch.Tensor [1, T]
+        clean = enh.enhance_file(tmp_in.name)
+    try:
+        os.remove(tmp_in.name)
+    except Exception:
+        pass
+
+    # Optional polish: high-pass & presence EQ
+    clean = _highpass(clean, 16000, lowcut_hz)
+    clean = _presence_boost(clean, 16000, presence_db)
+
+    # Peak limiting to avoid inter-sample clip
+    clean = _limit_peak(clean, target_dbfs=-1.0)
+
+    # Resample back
+    sr_out = sr_in if (out_sr is None or out_sr <= 0) else int(out_sr)
+    clean_out = _resample_torch(clean, 16000, sr_out).squeeze(0).numpy().astype(
+        np.float32
+    )
+
+    return sr_out, clean_out
+
+
+def _wav_bytes(sr: int, mono_f32: np.ndarray) -> bytes:
+    """Encode a mono float32 array as 16-bit PCM WAV into bytes."""
+    buf = io.BytesIO()
+    sf.write(buf, mono_f32, sr, subtype="PCM_16", format="WAV")
+    buf.seek(0)
+    return buf.read()
+
+
+# -----------------------------
+# FastAPI app with raw endpoint
+# -----------------------------
+app = FastAPI(title="Voice Clarity Booster (MetricGAN+)", version="1.0.0")
+
+
+@app.post("/enhance")
+async def enhance_endpoint(
+    file: UploadFile = File(..., description="Audio file (wav/mp3/ogg etc.)"),
+    presence_db: float = Query(3.0, ge=-12.0, le=12.0, description="Presence EQ gain in dB"),
+    lowcut_hz: float = Query(75.0, ge=0.0, le=200.0, description="High-pass cutoff in Hz"),
+    output_sr: int = Query(0, ge=0, description="0=keep original, or set to e.g. 44100/48000"),
+):
+    """Raw REST endpoint. Returns enhanced audio as audio/wav bytes."""
+    data = await file.read()
+    # Decode with soundfile
+    wav_np, sr_in = sf.read(io.BytesIO(data), always_2d=False, dtype="float32")
+    sr_out, enhanced = _enhance_numpy_audio(
+        (sr_in, wav_np),
+        presence_db=presence_db,
+        lowcut_hz=lowcut_hz,
+        out_sr=output_sr if output_sr > 0 else None,
+    )
+    wav_bytes = _wav_bytes(sr_out, enhanced)
+    headers = {"Content-Disposition": f'attachment; filename="{os.path.splitext(file.filename or "audio")[0]}_enhanced.wav"'}
+    return StreamingResponse(io.BytesIO(wav_bytes), media_type="audio/wav", headers=headers)
+
+
+# -----------------------------
+# Gradio UI (for quick testing)
+# -----------------------------
+def gradio_enhance(
+    audio: Tuple[int, np.ndarray],
+    presence_db: float,
+    lowcut_hz: float,
+    output_sr: str,
+):
+    if audio is None:
+        return None
+    out_sr = None
+    if output_sr in {"44100", "48000"}:
+        out_sr = int(output_sr)
+    # "Original" -> None
+    sr_out, enhanced = _enhance_numpy_audio(
+        audio, presence_db=float(presence_db), lowcut_hz=float(lowcut_hz), out_sr=out_sr
+    )
+    return (sr_out, enhanced)
+
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## Voice Clarity Booster (MetricGAN+)")
+    with gr.Row():
+        with gr.Column():
+            in_audio = gr.Audio(sources=["upload", "microphone"], type="numpy", label="Input")
+            presence = gr.Slider(-12, 12, value=3, step=0.5, label="Presence Boost (dB)")
+            lowcut = gr.Slider(0, 200, value=75, step=5, label="Low-Cut (Hz)")
+            out_sr = gr.Radio(
+                choices=["Original", "44100", "48000"],
+                value="Original",
+                label="Output Sample Rate",
+            )
+            btn = gr.Button("Enhance")
+        with gr.Column():
+            out_audio = gr.Audio(type="numpy", label="Enhanced")
+
+    btn.click(gradio_enhance, inputs=[in_audio, presence, lowcut, out_sr], outputs=[out_audio])
+
+# Mount Gradio at root path and keep FastAPI for /enhance
+app = gr.mount_gradio_app(app, demo, path="/")