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")