app.py

import os
import tempfile
import warnings
import numpy as np
import soundfile as sf
import torch
import gradio as gr
from pathlib import Path

# -----------------------------
# Config
# -----------------------------
DF_CHECKPOINT = "./DeepFilterNet2/checkpoints/model_96.ckpt.best"
TARGET_SR = 48000          # set to your model's expected rate
TARGET_SEC = 30            # << changed from 10s to 30s
DEVICE = torch.device("cpu")

NOISE_ROOT = "./noise_library"
NOISE_MAP = {
    "Cafe": os.path.join(NOISE_ROOT, "cafe.wav"),
    "Kitchen": os.path.join(NOISE_ROOT, "kitchen.wav"),
    "River": os.path.join(NOISE_ROOT, "river.wav"),
}

# -----------------------------
# Audio utils (stereo-aware)
# -----------------------------
def _resample_np(wave, sr_from, sr_to):
    if sr_from == sr_to:
        return wave
    try:
        import librosa
        # librosa expects shape (T,) so handle channel-wise
        if wave.ndim == 1:
            return librosa.resample(wave, orig_sr=sr_from, target_sr=sr_to).astype(np.float32)
        else:
            outs = []
            for c in range(wave.shape[0]):
                outs.append(librosa.resample(wave[c], orig_sr=sr_from, target_sr=sr_to))
            return np.stack(outs, axis=0).astype(np.float32)
    except Exception:
        warnings.warn("librosa not available; falling back to torchaudio resample.")
        import torchaudio
        if wave.ndim == 1:
            t = torch.from_numpy(wave).float().unsqueeze(0)  # (1,T)
        else:
            t = torch.from_numpy(wave).float()               # (C,T)
        r = torchaudio.functional.resample(t, orig_freq=sr_from, new_freq=sr_to)
        return r.numpy().astype(np.float32)

def load_wav_multi(path, target_sr=TARGET_SR, target_sec=TARGET_SEC, preserve_channels=True):
    """
    Returns (audio, sr) with shape (C,T) float32 in [-1,1]; C=1 or 2.
    - Preserves true stereo if present (preserve_channels=True).
    - Trims or pads to exactly target_sec for consistent runtime.
    """
    data, sr = sf.read(path, always_2d=True)  # (T,Ch)
    data = data.T.astype(np.float32)          # (Ch,T)
    if not preserve_channels:
        data = data.mean(axis=0, keepdims=True)  # mono
    # Resample channel-wise
    data = _resample_np(data, sr, target_sr)
    # Ensure shape (C,T)
    if data.ndim == 1:
        data = data[None, :]
    # Clip and length control
    data = np.clip(data, -1.0, 1.0)
    target_len = int(target_sr * target_sec)
    cur_len = data.shape[1]
    if cur_len > target_len:
        data = data[:, :target_len]
    elif cur_len < target_len:
        pad = np.zeros((data.shape[0], target_len - cur_len), dtype=np.float32)
        data = np.concatenate([data, pad], axis=1)
    return data, target_sr

def load_noise_multi(noise_path, target_sr=TARGET_SR, target_sec=TARGET_SEC, channels=1):
    """
    Loads/loops noise to (channels, target_len).
    If the noise file is mono and channels==2, duplicates channel to keep coherence,
    or you can randomize channels for a wider stereo if preferred.
    """
    n, sr = sf.read(noise_path, always_2d=True)  # (T,Ch)
    n = n.T.astype(np.float32)                   # (Ch,T)
    n = _resample_np(n, sr, target_sr)
    if n.ndim == 1:
        n = n[None, :]
    # Match channels
    if channels == 2 and n.shape[0] == 1:
        n = np.repeat(n, 2, axis=0)  # coherent stereo bed
    elif channels == 1 and n.shape[0] > 1:
        n = n.mean(axis=0, keepdims=True)
    # Loop/trim to target length
    target_len = int(target_sr * target_sec)
    if n.shape[1] < target_len:
        reps = int(np.ceil(target_len / n.shape[1]))
        n = np.tile(n, reps)[:, :target_len]
    else:
        n = n[:, :target_len]
    n = np.clip(n, -1.0, 1.0)
    return n, target_sr

def snr_mix_multi(clean, noise, snr_db):
    """
    Per-channel SNR mix. clean/noise: (C,T)
    """
    assert clean.shape == noise.shape, "clean/noise must match shape"
    out = np.empty_like(clean)
    for c in range(clean.shape[0]):
        pc = np.mean(clean[c] ** 2) + 1e-12
        pn = np.mean(noise[c] ** 2) + 1e-12
        desired_pn = pc / (10 ** (snr_db / 10.0))
        scale = np.sqrt(desired_pn / pn)
        noisy_c = clean[c] + noise[c] * scale
        mx = np.max(np.abs(noisy_c)) + 1e-12
        if mx > 1.0:
            noisy_c = noisy_c / mx
        out[c] = noisy_c.astype(np.float32)
    return out

# -----------------------------
# DeepFilter wrapper (per channel)
# -----------------------------
class DF2:
    def __init__(self, ckpt_path, device=DEVICE):
        # TODO: load your DF2 model (mono graph). Your logs show it loads fine already.
        self.device = device
        # self.model = ...
        # self.model.to(device).eval()
        pass

    @torch.no_grad()
    def enhance_mono(self, mono_np, sr=TARGET_SR):
        # Replace this stub with a real forward pass through DF2 (expects mono)
        # y = self.model(torch.from_numpy(mono_np).to(self.device).unsqueeze(0))
        # return y.squeeze(0).cpu().numpy().astype(np.float32)
        return mono_np  # placeholder passthrough

    def enhance_multi(self, audio_np, sr=TARGET_SR):
        """
        audio_np: (C,T) float32 in [-1,1]
        Runs the mono model per channel to preserve true stereo layout.
        """
        C, T = audio_np.shape
        outs = []
        for c in range(C):
            outs.append(self.enhance_mono(audio_np[c], sr=sr))
        return np.stack(outs, axis=0).astype(np.float32)

df2 = DF2(DF_CHECKPOINT, device=DEVICE)

# -----------------------------
# Gradio pipeline
# -----------------------------
def list_noises():
    available = [name for name, p in NOISE_MAP.items() if os.path.isfile(p)]
    return available or list(NOISE_MAP.keys())

def process(speech_file, noise_name, snr_db):
    if speech_file is None:
        raise gr.Error("Please upload a WAV/AIFF/FLAC file.")
    if noise_name not in NOISE_MAP:
        raise gr.Error(f"Unknown noise '{noise_name}'.")

    # Load user file, preserving mono/stereo, and force 30s
    clean, sr = load_wav_multi(speech_file, TARGET_SR, TARGET_SEC, preserve_channels=True)
    # Load/prepare noise to same channels & length
    noise, _ = load_noise_multi(NOISE_MAP[noise_name], TARGET_SR, TARGET_SEC, channels=clean.shape[0])

    # Mix and enhance
    noisy = snr_mix_multi(clean, noise, float(snr_db))
    enhanced = df2.enhance_multi(noisy, sr=sr)

    # Save as original channel count
    def _save(path, arr, sr):
        # arr: (C,T) -> (T,C)
        sf.write(path, arr.T, sr, subtype="PCM_16")

    orig_path = tempfile.mkstemp(suffix=".wav")[1]
    noisy_path = tempfile.mkstemp(suffix=".wav")[1]
    enh_path  = tempfile.mkstemp(suffix=".wav")[1]
    _save(orig_path, clean, sr)
    _save(noisy_path, noisy, sr)
    _save(enh_path,  enhanced, sr)

    return orig_path, noisy_path, enh_path

with gr.Blocks(title="DeepFilterNet2 Noise Reduction (Mono & Stereo)") as demo:
    gr.Markdown("# DeepFilterNet2 Noise Reduction\n**Now supports mono and true stereo.** Processes up to **30 seconds**.")
    with gr.Row():
        with gr.Column():
            inp_audio = gr.Audio(label="Speech (Mono or Stereo)", type="filepath")
            noise_dd  = gr.Dropdown(label="Noise", choices=list_noises(), value=list_noises()[0] if list_noises() else None)
            snr_slider = gr.Slider(0, 20, value=10, step=1, label="SNR (dB)")
            run_btn   = gr.Button("Denoise")
            gr.Markdown("Tip: Upload 16–48 kHz mono or stereo WAV/AIFF/FLAC. The app will process **30s** (trim/pad as needed).")
        with gr.Column():
            out_orig = gr.Audio(label="Original (Mono/Stereo)", type="filepath")
            out_noisy = gr.Audio(label="Noisy Mix (Mono/Stereo)", type="filepath")
            out_enh  = gr.Audio(label="Enhanced (Mono/Stereo)", type="filepath")

    # Provide complete examples matching inputs to avoid caching warnings
    examples = [
        ["./examples/p232_013_clean.wav", "Kitchen", 10],
        ["./examples/p232_019_clean_stereo.wav", "Cafe", 10],  # add a short stereo example if you can
    ]
    gr.Examples(
        examples=examples,
        inputs=[inp_audio, noise_dd, snr_slider],
        outputs=[out_orig, out_noisy, out_enh],
        cache_examples=True
    )

    run_btn.click(process, inputs=[inp_audio, noise_dd, snr_slider], outputs=[out_orig, out_noisy, out_enh])

if __name__ == "__main__":
    demo.queue(concurrency_count=2, max_size=20)
    demo.launch(server_name="0.0.0.0")