Spaces:
Running on Zero
Running on Zero
| """ | |
| denoise_core.py — Motor de reducción de ruido (DeepFilterNet3) con ZeroGPU y respaldo en CPU. | |
| - El modelo trabaja a 48 kHz. Quien llama se encarga de remuestrear (ver dsp.py). | |
| - El audio se procesa por bloques con contexto y fundido cruzado, para que cada | |
| llamada a la GPU sea corta y la memoria no dependa de la duración del fichero. | |
| - Si la GPU no está disponible (cuota agotada, cola, error), se usa la CPU. | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import sys | |
| import time | |
| import types | |
| import warnings | |
| try: | |
| import spaces # en ZeroGPU debe importarse antes que torch | |
| except ImportError: # versión de escritorio: sin ZeroGPU | |
| spaces = None | |
| # Shim de compatibilidad: deepfilternet 0.5.6 importa torchaudio.backend.common (solo para df.io, | |
| # que no usamos). Versiones recientes de torchaudio lo han eliminado: si falta, se crea uno mínimo. | |
| try: | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| import torchaudio.backend.common # noqa: F401 | |
| except ImportError: | |
| from dataclasses import dataclass | |
| class AudioMetaData: | |
| sample_rate: int | |
| num_frames: int | |
| num_channels: int | |
| bits_per_sample: int | |
| encoding: str | |
| _common = types.ModuleType("torchaudio.backend.common") | |
| _common.AudioMetaData = AudioMetaData | |
| _backend = types.ModuleType("torchaudio.backend") | |
| _backend.common = _common | |
| sys.modules["torchaudio.backend"] = _backend | |
| sys.modules["torchaudio.backend.common"] = _common | |
| import numpy as np | |
| import torch | |
| from torch.nn import functional as F | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") # avisos de deprecación de torchaudio dentro de deepfilternet | |
| from df.enhance import df_features, init_df | |
| from df.model import ModelParams | |
| from df.utils import as_complex | |
| from models import dfn3_dir | |
| DF_SR = 48000 | |
| CHUNK_SEC = 60.0 # bloque útil por inferencia | |
| CONTEXT_SEC = 5.0 # contexto extra a cada lado para que el modelo "se asiente" | |
| XFADE_SEC = 0.05 # fundido entre bloques | |
| GROUP_SEC = 300.0 # audio máximo por llamada a la GPU | |
| GPU_MIN_SEC = 90.0 # por debajo, la espera de ZeroGPU (cola, asignación) cuesta más que hacerlo en CPU | |
| GPU_COOLDOWN_SEC = 900.0 # tras un fallo de GPU (p. ej. cuota) no se reintenta en este tiempo | |
| ON_ZEROGPU = bool(spaces is not None and spaces.config.Config.zero_gpu) | |
| _gpu_blocked_until = 0.0 | |
| def _force_df_device(device: str) -> None: | |
| """DeepFilterNet decide el dispositivo con df.utils.get_device(); lo fijamos nosotros.""" | |
| dev = torch.device(device) | |
| for name, mod in list(sys.modules.items()): | |
| if (name == "df" or name.startswith("df.")) and hasattr(mod, "get_device"): | |
| mod.get_device = lambda dev=dev: dev | |
| def _disable_df_git() -> None: | |
| """DeepFilterNet consulta git al iniciar su log; sin git instalado (Windows) eso rompe el arranque.""" | |
| import df.logger | |
| import df.utils | |
| for mod in (df.utils, df.logger): | |
| mod.get_commit_hash = lambda: None | |
| mod.get_branch_name = lambda: None | |
| print("[MODELO] Cargando DeepFilterNet3…") | |
| _disable_df_git() | |
| _force_df_device("cpu") | |
| _model_cpu, _df_state, _ = init_df(model_base_dir=str(dfn3_dir()), log_level="WARNING", log_file=None) | |
| _model_cpu.eval() | |
| _nb_df = getattr(_model_cpu, "nb_df", getattr(_model_cpu, "df_bins", ModelParams().nb_df)) | |
| assert _df_state.sr() == DF_SR | |
| # En ZeroGPU el modelo se coloca en "cuda" a nivel de módulo (emulado fuera de @spaces.GPU). | |
| _model_gpu = None | |
| if ON_ZEROGPU or torch.cuda.is_available(): | |
| try: | |
| _model_gpu = copy.deepcopy(_model_cpu).to("cuda").eval() | |
| except Exception as e: # sin CUDA real ni emulada: solo CPU | |
| print(f"[MODELO] GPU no disponible: {e}") | |
| print(f"[MODELO] Listo · ZeroGPU={ON_ZEROGPU} · GPU={'sí' if _model_gpu is not None else 'no'}") | |
| def _enhance_block(model, audio: np.ndarray, atten_db: float | None, device: str) -> np.ndarray: | |
| """Equivalente a df.enhance.enhance() pero con dispositivo explícito. audio: (C, T) float32 a 48 kHz.""" | |
| x = torch.from_numpy(np.ascontiguousarray(audio, dtype=np.float32)) | |
| orig_len = x.shape[-1] | |
| n_fft, hop = _df_state.fft_size(), _df_state.hop_size() | |
| x = F.pad(x, (0, n_fft)) # compensa el retardo del STFT | |
| spec, erb_feat, spec_feat = df_features(x, _df_state, _nb_df, device=device) | |
| enhanced = as_complex(model(spec.clone(), erb_feat, spec_feat)[0].cpu().squeeze(1)) | |
| if atten_db is not None and atten_db > 0: | |
| lim = 10 ** (-atten_db / 20) | |
| enhanced = as_complex(spec.squeeze(1).cpu()) * lim + enhanced * (1 - lim) | |
| out = _df_state.synthesis(enhanced.numpy()) | |
| d = n_fft - hop | |
| return np.asarray(out[:, d : orig_len + d], dtype=np.float32) | |
| def _plan_blocks(n: int) -> list[tuple[int, int, int, int]]: | |
| """Bloques (ini_ctx, fin_ctx, ini_util, fin_util) en muestras.""" | |
| chunk, ctx = int(CHUNK_SEC * DF_SR), int(CONTEXT_SEC * DF_SR) | |
| blocks, s = [], 0 | |
| while s < n: | |
| e = min(n, s + chunk) | |
| if n - e < chunk // 4: # no dejar un último bloque diminuto | |
| e = n | |
| blocks.append((max(0, s - ctx), min(n, e + ctx), s, e)) | |
| s = e | |
| return blocks | |
| def _run_blocks(model, pieces: list[np.ndarray], atten_db, device) -> list[np.ndarray]: | |
| return [_enhance_block(model, p, atten_db, device) for p in pieces] | |
| def _gpu_duration(pieces, atten_db) -> int: | |
| seconds = sum(p.shape[-1] for p in pieces) / DF_SR | |
| return int(min(120, 15 + seconds * 0.15)) | |
| def _gpu_decorator(fn): | |
| return spaces.GPU(duration=_gpu_duration)(fn) if spaces is not None else fn | |
| def _run_blocks_gpu(pieces: list[np.ndarray], atten_db) -> list[np.ndarray]: | |
| return _run_blocks(_model_gpu, pieces, atten_db, "cuda") | |
| def enhance_48k(audio: np.ndarray, atten_db: float | None, use_gpu: bool = True, progress=None): | |
| """ | |
| Reduce el ruido de `audio` (C, T) float32 a 48 kHz. | |
| atten_db: reducción máxima en dB (None = sin límite). | |
| Devuelve (audio_procesado, info) donde info indica dispositivo y avisos. | |
| """ | |
| n = audio.shape[-1] | |
| blocks = _plan_blocks(n) | |
| xf = int(XFADE_SEC * DF_SR) | |
| out = np.zeros_like(audio, dtype=np.float32) | |
| weight = np.zeros(n, dtype=np.float32) | |
| # Agrupa bloques para que cada llamada a la GPU sea corta. | |
| per_group = max(1, int(GROUP_SEC // CHUNK_SEC)) | |
| groups = [blocks[i : i + per_group] for i in range(0, len(blocks), per_group)] | |
| global _gpu_blocked_until | |
| gpu_ok = (use_gpu and _model_gpu is not None and time.time() >= _gpu_blocked_until | |
| and (not ON_ZEROGPU or n / DF_SR >= GPU_MIN_SEC)) | |
| devices, notes = set(), [] | |
| for gi, group in enumerate(groups): | |
| pieces = [audio[:, a:b] for a, b, _, _ in group] | |
| results = None | |
| if gpu_ok: | |
| try: | |
| results = _run_blocks_gpu(pieces, atten_db) | |
| devices.add("GPU") | |
| except Exception as e: | |
| gpu_ok = False # si falla una vez (p. ej. cuota), el resto va por CPU | |
| _gpu_blocked_until = time.time() + GPU_COOLDOWN_SEC | |
| print(f"[GPU] No disponible, se sigue en CPU: {str(e).splitlines()[0][:200]}") | |
| notes.append("La tarjeta gráfica no estaba libre: se ha hecho con el procesador (mismo resultado)") | |
| if results is None: | |
| results = _run_blocks(_model_cpu, pieces, atten_db, "cpu") | |
| devices.add("CPU") | |
| for (a, b, s, e), y in zip(group, results): | |
| lo, hi = max(0, s - xf), min(n, e + xf) | |
| w = np.ones(hi - lo, dtype=np.float32) | |
| if s > 0: # rampa de entrada centrada en s | |
| w[: 2 * xf] = np.linspace(0, 1, 2 * xf, dtype=np.float32)[: min(2 * xf, hi - lo)] | |
| if e < n: # rampa de salida centrada en e | |
| w[-2 * xf :] = np.minimum(w[-2 * xf :], np.linspace(1, 0, 2 * xf, dtype=np.float32)) | |
| out[:, lo:hi] += y[:, lo - a : hi - a] * w | |
| weight[lo:hi] += w | |
| if progress is not None: | |
| progress((gi + 1) / len(groups)) | |
| out /= np.maximum(weight, 1e-6) | |
| return out, {"dispositivo": "+".join(sorted(devices)) or "—", "avisos": notes} | |
| def gpu_available() -> bool: | |
| return _model_gpu is not None | |