Spaces:
Running
Running
| import io | |
| import math | |
| import tempfile | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Dict, Optional, Tuple | |
| import gradio as gr | |
| import librosa | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import onnxruntime as ort | |
| import soundfile as sf | |
| from PIL import Image | |
| # ----------------------------- | |
| # Configuration | |
| # ----------------------------- | |
| MAX_SECONDS = 10.0 | |
| ONNX_DIR = Path("./onnx") | |
| class ModelSpec: | |
| name: str | |
| sr: int | |
| onnx_path: str | |
| # ----------------------------- | |
| # Model discovery and metadata | |
| # ----------------------------- | |
| def _infer_model_meta(model_name: str) -> int: | |
| normalized = model_name.lower().replace("-", "_") | |
| if "48khz" in normalized or "48k" in normalized or "48hr" in normalized: | |
| return 48000 | |
| # Fallback for unknown 16 kHz DPDFNet variants | |
| return 16000 | |
| def _display_label(spec: ModelSpec) -> str: | |
| khz = int(spec.sr // 1000) | |
| return f"{spec.name} ({khz} kHz)" | |
| def discover_model_presets() -> Dict[str, ModelSpec]: | |
| ordered_names = [ | |
| "baseline", | |
| "dpdfnet2", | |
| "dpdfnet4", | |
| "dpdfnet8", | |
| "dpdfnet2_48khz_hr", | |
| "dpdfnet8_48khz_hr", | |
| ] | |
| found_paths = {p.stem: p for p in ONNX_DIR.glob("*.onnx") if p.is_file()} | |
| presets: Dict[str, ModelSpec] = {} | |
| for name in ordered_names: | |
| p = found_paths.get(name) | |
| if p is None: | |
| continue | |
| sr = _infer_model_meta(name) | |
| spec = ModelSpec( | |
| name=name, | |
| sr=sr, | |
| onnx_path=str(p), | |
| ) | |
| presets[_display_label(spec)] = spec | |
| # Include any additional ONNX files not in the canonical order list. | |
| for name, p in sorted(found_paths.items()): | |
| if name in ordered_names: | |
| continue | |
| sr = _infer_model_meta(name) | |
| spec = ModelSpec( | |
| name=name, | |
| sr=sr, | |
| onnx_path=str(p), | |
| ) | |
| presets[_display_label(spec)] = spec | |
| return presets | |
| MODEL_PRESETS = discover_model_presets() | |
| DEFAULT_MODEL_KEY = next(iter(MODEL_PRESETS), None) | |
| # ----------------------------- | |
| # ONNX Runtime + frontend cache | |
| # ----------------------------- | |
| _SESSIONS: Dict[str, ort.InferenceSession] = {} | |
| _INIT_STATES: Dict[str, np.ndarray] = {} | |
| def resolve_model_path(local_path: str) -> str: | |
| p = Path(local_path) | |
| if p.exists(): | |
| return str(p) | |
| raise gr.Error( | |
| f"ONNX model not found at: {local_path}. " | |
| "Expected local models under ./onnx/." | |
| ) | |
| def get_ort_session(model_key: str) -> ort.InferenceSession: | |
| if model_key in _SESSIONS: | |
| return _SESSIONS[model_key] | |
| spec = MODEL_PRESETS[model_key] | |
| onnx_path = resolve_model_path(spec.onnx_path) | |
| options = ort.SessionOptions() | |
| options.intra_op_num_threads = 1 | |
| options.inter_op_num_threads = 1 | |
| sess = ort.InferenceSession( | |
| onnx_path, | |
| sess_options=options, | |
| providers=["CPUExecutionProvider"], | |
| ) | |
| _SESSIONS[model_key] = sess | |
| return sess | |
| def _load_initial_state(model_key: str, session: ort.InferenceSession) -> np.ndarray: | |
| if model_key in _INIT_STATES: | |
| return _INIT_STATES[model_key] | |
| if len(session.get_inputs()) < 2: | |
| raise gr.Error("Expected streaming ONNX model with two inputs: (spec, state).") | |
| meta = session.get_modelmeta().custom_metadata_map | |
| try: | |
| state_size = int(meta["state_size"]) | |
| erb_norm_state_size = int(meta["erb_norm_state_size"]) | |
| spec_norm_state_size = int(meta["spec_norm_state_size"]) | |
| erb_norm_init = np.array( | |
| [float(x) for x in meta["erb_norm_init"].split(",")], dtype=np.float32 | |
| ) | |
| spec_norm_init = np.array( | |
| [float(x) for x in meta["spec_norm_init"].split(",")], dtype=np.float32 | |
| ) | |
| except KeyError as exc: | |
| raise gr.Error( | |
| f"ONNX model is missing required metadata key: {exc}. " | |
| "Re-export the model to embed state initialisation metadata." | |
| ) | |
| init_state = np.zeros(state_size, dtype=np.float32) | |
| init_state[0:erb_norm_state_size] = erb_norm_init | |
| init_state[erb_norm_state_size:erb_norm_state_size + spec_norm_state_size] = spec_norm_init | |
| init_state = np.ascontiguousarray(init_state) | |
| _INIT_STATES[model_key] = init_state | |
| return init_state | |
| # ----------------------------- | |
| # STFT/iSTFT (module-free) | |
| # ----------------------------- | |
| def vorbis_window(window_len: int) -> np.ndarray: | |
| window_size_h = window_len / 2 | |
| indices = np.arange(window_len) | |
| sin = np.sin(0.5 * np.pi * (indices + 0.5) / window_size_h) | |
| window = np.sin(0.5 * np.pi * sin * sin) | |
| return window.astype(np.float32) | |
| def _infer_stft_params(model_key: str, session: ort.InferenceSession) -> Tuple[int, int, np.ndarray]: | |
| # ONNX spec input is [B, T, F, 2] (or dynamic variants). | |
| spec_shape = session.get_inputs()[0].shape | |
| freq_bins = spec_shape[-2] if len(spec_shape) >= 2 else None | |
| if isinstance(freq_bins, int) and freq_bins > 1: | |
| win_len = int((freq_bins - 1) * 2) | |
| else: | |
| # 20 ms windows for DPDFNet family. | |
| sr = MODEL_PRESETS[model_key].sr | |
| win_len = int(round(sr * 0.02)) | |
| hop = win_len // 2 | |
| win = vorbis_window(win_len) | |
| return win_len, hop, win | |
| def _preprocess_waveform(waveform: np.ndarray, win_len: int, hop: int, win: np.ndarray) -> np.ndarray: | |
| audio = np.asarray(waveform, dtype=np.float32).reshape(-1) | |
| audio_pad = np.pad(audio, (0, win_len), mode="constant") | |
| spec = librosa.stft( | |
| y=audio_pad, | |
| n_fft=win_len, | |
| hop_length=hop, | |
| win_length=win_len, | |
| window=win, | |
| center=True, | |
| pad_mode="reflect", | |
| ) | |
| spec = spec.T.astype(np.complex64, copy=False) # [T, F] | |
| spec_ri = np.stack([spec.real, spec.imag], axis=-1).astype(np.float32, copy=False) # [T, F, 2] | |
| return np.ascontiguousarray(spec_ri[None, ...], dtype=np.float32) # [1, T, F, 2] | |
| def _postprocess_spec(spec_e: np.ndarray, win_len: int, hop: int, win: np.ndarray) -> np.ndarray: | |
| spec_c = np.asarray(spec_e[0], dtype=np.float32) # [T, F, 2] | |
| spec = (spec_c[..., 0] + 1j * spec_c[..., 1]).T.astype(np.complex64, copy=False) # [F, T] | |
| waveform_e = librosa.istft( | |
| spec, | |
| hop_length=hop, | |
| win_length=win_len, | |
| window=win, | |
| center=True, | |
| length=None, | |
| ).astype(np.float32, copy=False) | |
| return np.concatenate( | |
| [waveform_e[win_len * 2 :], np.zeros(win_len * 2, dtype=np.float32)], | |
| axis=0, | |
| ) | |
| # ----------------------------- | |
| # ONNX inference (non-streaming pre/post, streaming ONNX state loop) | |
| # ----------------------------- | |
| def enhance_audio_onnx( | |
| audio_mono: np.ndarray, | |
| model_key: str, | |
| ) -> np.ndarray: | |
| sess = get_ort_session(model_key) | |
| inputs = sess.get_inputs() | |
| outputs = sess.get_outputs() | |
| if len(inputs) < 2 or len(outputs) < 2: | |
| raise gr.Error( | |
| "Expected streaming ONNX signature with 2 inputs (spec, state) and 2 outputs (spec_e, state_out)." | |
| ) | |
| in_spec_name = inputs[0].name | |
| in_state_name = inputs[1].name | |
| out_spec_name = outputs[0].name | |
| out_state_name = outputs[1].name | |
| waveform = np.asarray(audio_mono, dtype=np.float32).reshape(-1) | |
| win_len, hop, win = _infer_stft_params(model_key, sess) | |
| spec_r_np = _preprocess_waveform(waveform, win_len=win_len, hop=hop, win=win) | |
| state = _load_initial_state(model_key, sess).copy() | |
| spec_e_frames = [] | |
| num_frames = int(spec_r_np.shape[1]) | |
| for t in range(num_frames): | |
| spec_t = np.ascontiguousarray(spec_r_np[:, t : t + 1, :, :], dtype=np.float32) | |
| spec_e_t, state = sess.run( | |
| [out_spec_name, out_state_name], | |
| {in_spec_name: spec_t, in_state_name: state}, | |
| ) | |
| spec_e_frames.append(np.ascontiguousarray(spec_e_t, dtype=np.float32)) | |
| if not spec_e_frames: | |
| return waveform | |
| spec_e_np = np.concatenate(spec_e_frames, axis=1) | |
| waveform_e = _postprocess_spec(spec_e_np, win_len=win_len, hop=hop, win=win) | |
| return np.asarray(waveform_e, dtype=np.float32).reshape(-1) | |
| # ----------------------------- | |
| # Audio utilities | |
| # ----------------------------- | |
| def _load_wav_from_gradio_path(path: str) -> Tuple[np.ndarray, int]: | |
| data, sr = sf.read(path, always_2d=True) | |
| data = data.astype(np.float32, copy=False) | |
| return data, int(sr) | |
| def _to_mono(x: np.ndarray) -> Tuple[np.ndarray, int]: | |
| if x.ndim == 1: | |
| return x.astype(np.float32, copy=False), 1 | |
| if x.shape[1] == 1: | |
| return x[:, 0], 1 | |
| return x.mean(axis=1), int(x.shape[1]) | |
| def _resample(y: np.ndarray, sr_in: int, sr_out: int) -> np.ndarray: | |
| if sr_in == sr_out: | |
| return y | |
| return librosa.resample(y, orig_sr=sr_in, target_sr=sr_out).astype(np.float32, copy=False) | |
| def _match_length(y: np.ndarray, target_len: int) -> np.ndarray: | |
| if len(y) == target_len: | |
| return y | |
| if len(y) > target_len: | |
| return y[:target_len] | |
| out = np.zeros((target_len,), dtype=y.dtype) | |
| out[: len(y)] = y | |
| return out | |
| def _save_wav(y: np.ndarray, sr: int, prefix: str) -> str: | |
| tmp = tempfile.NamedTemporaryFile(prefix=prefix, suffix=".wav", delete=False) | |
| tmp.close() | |
| sf.write(tmp.name, y, sr) | |
| return tmp.name | |
| def _spectrogram_image(y: np.ndarray, sr: int) -> Image.Image: | |
| win_length = max(256, int(0.032 * sr)) | |
| hop_length = max(64, int(0.008 * sr)) | |
| n_fft = 1 << (int(math.ceil(math.log2(win_length)))) | |
| S = librosa.stft(y, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False) | |
| S_db = librosa.amplitude_to_db(np.abs(S) + 1e-10, ref=np.max) | |
| fig, ax = plt.subplots(figsize=(8.4, 3.2)) | |
| ax.imshow(S_db, origin="lower", aspect="auto") | |
| ax.set_axis_off() | |
| fig.subplots_adjust(left=0, right=1, top=1, bottom=0) | |
| buf = io.BytesIO() | |
| fig.savefig(buf, format="png", dpi=160) | |
| plt.close(fig) | |
| buf.seek(0) | |
| return Image.open(buf) | |
| # ----------------------------- | |
| # Main pipeline | |
| # ----------------------------- | |
| def run_enhancement( | |
| source: str, | |
| mic_path: Optional[str], | |
| file_path: Optional[str], | |
| model_key: str, | |
| ): | |
| if not MODEL_PRESETS: | |
| raise gr.Error("No ONNX models found under ./onnx/. Add models and retry.") | |
| chosen_path = mic_path if source == "Microphone" else file_path | |
| if not chosen_path: | |
| raise gr.Error("Please provide audio either from the microphone or by uploading a file.") | |
| x, sr_orig = _load_wav_from_gradio_path(chosen_path) | |
| y_mono, n_ch = _to_mono(x) | |
| max_samples = int(MAX_SECONDS * sr_orig) | |
| was_trimmed = len(y_mono) > max_samples | |
| if was_trimmed: | |
| y_mono = y_mono[:max_samples] | |
| dur = len(y_mono) / float(sr_orig) | |
| spec = MODEL_PRESETS[model_key] | |
| sr_model = spec.sr | |
| y_model = _resample(y_mono, sr_orig, sr_model) | |
| y_enh_model = enhance_audio_onnx(y_model, model_key) | |
| y_enh = _resample(y_enh_model, sr_model, sr_orig) | |
| y_enh = _match_length(y_enh, len(y_mono)) | |
| noisy_out = _save_wav(y_mono, sr_orig, prefix="noisy_mono_") | |
| enh_out = _save_wav(y_enh, sr_orig, prefix="enhanced_") | |
| noisy_img = _spectrogram_image(y_mono, sr_orig) | |
| enh_img = _spectrogram_image(y_enh, sr_orig) | |
| status = ( | |
| f"**Input:** {sr_orig} Hz, {dur:.2f}s, channels={n_ch} ⭢ mono\n\n" | |
| f"**Model:** {spec.name} (runs at {sr_model} Hz)\n\n" | |
| + ( | |
| f"**Resampling:** {sr_orig} ⭢ {sr_model} ⭢ {sr_orig}\n\n" | |
| if sr_orig != sr_model | |
| else "**Resampling:** none\n\n" | |
| ) | |
| + (f"**Trimmed:** first {MAX_SECONDS:.0f}s used\n" if was_trimmed else "") | |
| + "\n✅ Done." | |
| ) | |
| return noisy_out, enh_out, noisy_img, enh_img, status | |
| def set_source_visibility(source: str): | |
| return ( | |
| gr.update(visible=(source == "Microphone")), | |
| gr.update(visible=(source == "Upload")), | |
| ) | |
| # ----------------------------- | |
| # UI (light polish) | |
| # ----------------------------- | |
| THEME = gr.themes.Soft( | |
| primary_hue="orange", | |
| neutral_hue="slate", | |
| font=[ | |
| "Arial", | |
| "ui-sans-serif", | |
| "system-ui", | |
| "Segoe UI", | |
| "Roboto", | |
| "Helvetica Neue", | |
| "Noto Sans", | |
| "Liberation Sans", | |
| "sans-serif", | |
| ], | |
| ) | |
| CSS = """ | |
| .gradio-container{ | |
| max-width: 1040px !important; | |
| margin: 0 auto !important; | |
| font-family: Arial, ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Helvetica Neue, Noto Sans, Liberation Sans, sans-serif !important; | |
| } | |
| #header { | |
| padding: 14px 16px; | |
| border-radius: 16px; | |
| border: 1px solid rgba(0,0,0,0.08); | |
| background: linear-gradient(135deg, rgba(255,152,0,0.14), rgba(255,152,0,0.04)); | |
| text-align: center; | |
| } | |
| #header h1{ | |
| margin: 0 0 6px 0; | |
| font-size: 24px; | |
| font-weight: 800; | |
| letter-spacing: -0.2px; | |
| } | |
| #header p{ | |
| margin: 6px auto 0 auto; | |
| max-width: 720px; | |
| color: var(--body-text-color-subdued); | |
| font-size: 14px; | |
| line-height: 1.6; | |
| } | |
| #header hr{ | |
| margin-top: 18px; | |
| border: none; | |
| height: 1px; | |
| background: linear-gradient(to right, transparent, #ddd, transparent); | |
| } | |
| .spec img { border-radius: 14px; } | |
| .audio { border-radius: 14px !important; overflow: hidden; } | |
| #run_btn{ | |
| border-radius: 12px !important; | |
| font-weight: 800 !important; | |
| } | |
| #status_md p{ margin: 0.35rem 0; } | |
| """ | |
| with gr.Blocks(theme=THEME, css=CSS, title="DPDFNet Speech Enhancement") as demo: | |
| gr.Markdown( | |
| "# DPDFNet Speech Enhancement\n\n" | |
| "Causal · Real-Time · Edge-Ready\n\n" | |
| "DPDFNet extends DeepFilterNet2 with Dual-Path RNN blocks to improve " | |
| "long-range temporal and cross-band modeling while preserving low latency. " | |
| "Designed for single-channel streaming speech enhancement under challenging noise conditions.\n\n" | |
| "---", | |
| elem_id="header", | |
| ) | |
| with gr.Row(): | |
| model_key = gr.Dropdown( | |
| choices=list(MODEL_PRESETS.keys()), | |
| value=DEFAULT_MODEL_KEY, | |
| label="Model", | |
| # info="Audio is resampled to model SR, enhanced with ONNX, then resampled back.", | |
| interactive=True, | |
| ) | |
| source = gr.Radio( | |
| choices=["Microphone", "Upload"], | |
| value="Upload", | |
| label="Input source", | |
| ) | |
| with gr.Row(): | |
| mic_audio = gr.Audio( | |
| sources=["microphone"], | |
| type="filepath", | |
| format="wav", | |
| label="Microphone (max 10s)", | |
| visible=False, | |
| buttons=["download"], | |
| elem_classes=["audio"], | |
| ) | |
| file_audio = gr.Audio( | |
| sources=["upload"], | |
| type="filepath", | |
| format="wav", | |
| label="Upload file (WAV/MP3/FLAC etc., max 10s)", | |
| visible=True, | |
| buttons=["download"], | |
| elem_classes=["audio"], | |
| ) | |
| run_btn = gr.Button("Enhance", variant="primary", elem_id="run_btn") | |
| status = gr.Markdown(elem_id="status_md") | |
| gr.Markdown("## Results") | |
| with gr.Row(): | |
| out_noisy = gr.Audio(label="Before (mono)", interactive=False, format="wav", buttons=["download"], elem_classes=["audio"]) | |
| out_enh = gr.Audio(label="After (enhanced)", interactive=False, format="wav", buttons=["download"], elem_classes=["audio"]) | |
| with gr.Row(): | |
| img_noisy = gr.Image(label="Noisy spectrogram", elem_classes=["spec"]) | |
| img_enh = gr.Image(label="Enhanced spectrogram", elem_classes=["spec"]) | |
| source.change(fn=set_source_visibility, inputs=source, outputs=[mic_audio, file_audio]) | |
| run_btn.click( | |
| fn=run_enhancement, | |
| inputs=[source, mic_audio, file_audio, model_key], | |
| outputs=[out_noisy, out_enh, img_noisy, img_enh, status], | |
| api_name="enhance", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=32).launch() | |