""" Server-side streaming noise suppression (pre-processor for the caller audio that is forwarded to DashScope). Fully isolated and fail-open: - If disabled, or the model / onnxruntime is unavailable, or ANY error occurs, audio passes through UNCHANGED — the call is never affected. - DashScope's own VAD/turn-detection is untouched; this only improves the quality of the audio DashScope receives. Engine: GTCRN streaming ONNX (16 kHz mono, model-native — no resampling). The model file is NOT vendored; drop the streaming ONNX into: models/gtcrn/ (e.g. models/gtcrn/gtcrn_stream.onnx) Until it is present, is_active() is False and process() is a pass-through. The exact frame/hop and cache-tensor names come from the ONNX file itself (read via onnxruntime introspection at load time), so this wrapper adapts to the specific artifact rather than hard-coding shapes. """ import glob import os MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'gtcrn') # Frame geometry for GTCRN at 16 kHz. Overridable per model via env if a # different export is used. 512-pt window / 256 hop is the common config. _N_FFT = int(os.environ.get('GTCRN_NFFT', '512')) _HOP = int(os.environ.get('GTCRN_HOP', '256')) def _find_model(): if not os.path.isdir(MODEL_DIR): return None hits = sorted(glob.glob(os.path.join(MODEL_DIR, '*.onnx'))) return hits[0] if hits else None class NoiseSuppressor: """One instance per session. Streaming, stateful, fail-open. Public API: is_active() -> bool (model loaded & engine enabled) process(pcm_bytes) -> bytes (denoised, or unchanged on any issue) reset() -> None (clear per-session streaming state) """ _shared_session = None # class-level: the loaded ORT session (shared) _shared_meta = None # (input_names, output_names, cache spec) _load_attempted = False _load_error = None def __init__(self, engine='off'): self.engine = (engine or 'off').lower() self._buf = bytearray() # ring buffer of raw PCM bytes self._state = None # per-session cache tensors self._active = False self._warned = False if self.engine == 'gtcrn': self._try_activate() # ── Loading ────────────────────────────────────────────────────── @classmethod def _load_shared(cls): """Load the ONNX session once, shared across sessions. Never raises.""" if cls._load_attempted: return cls._shared_session is not None cls._load_attempted = True try: model_path = _find_model() if not model_path: cls._load_error = f'no .onnx in {MODEL_DIR}' return False import onnxruntime as ort # noqa: import here so absence is graceful so = ort.SessionOptions() so.intra_op_num_threads = 1 # per-call work is tiny; avoid oversubscribe sess = ort.InferenceSession( model_path, sess_options=so, providers=['CPUExecutionProvider']) inputs = [i.name for i in sess.get_inputs()] outputs = [o.name for o in sess.get_outputs()] # Cache/state tensors: any input whose name isn't the primary audio # input. The primary input is assumed to be the first input. cache_inputs = [i for i in sess.get_inputs()][1:] cls._shared_session = sess cls._shared_meta = { 'input_names': inputs, 'output_names': outputs, 'primary_in': inputs[0], 'primary_out': outputs[0], 'cache_inputs': [(i.name, i.shape, i.type) for i in cache_inputs], 'cache_out_names': outputs[1:], 'model_path': model_path, } print(f'[NoiseSuppress] GTCRN model loaded: {model_path} ' f'(inputs={inputs}, outputs={outputs})') return True except Exception as e: # onnxruntime missing, bad model, etc. cls._load_error = str(e) cls._shared_session = None print(f'[NoiseSuppress] GTCRN unavailable ({e}); ' f'audio will pass through unchanged.') return False def _try_activate(self): ok = NoiseSuppressor._load_shared() self._active = bool(ok) if not ok and not self._warned: self._warned = True print('[NoiseSuppress] engine=gtcrn requested but inactive ' f'({NoiseSuppressor._load_error}) — pass-through mode.') # ── Public ─────────────────────────────────────────────────────── def is_active(self): return self._active and self.engine == 'gtcrn' def reset(self): self._buf = bytearray() self._state = None def process(self, pcm_bytes): """Denoise a chunk of 16 kHz mono 16-bit PCM. Returns processed PCM. Fail-open: on inactive engine or ANY error, returns pcm_bytes as-is.""" if not self.is_active() or not pcm_bytes: return pcm_bytes try: return self._process_stream(pcm_bytes) except Exception as e: # One-time warn, then permanently degrade this session to passthrough if not self._warned: self._warned = True print(f'[NoiseSuppress] runtime error ({e}); ' f'session falls back to pass-through.') self._active = False return pcm_bytes # ── Streaming core ─────────────────────────────────────────────── def _process_stream(self, pcm_bytes): import numpy as np sess = NoiseSuppressor._shared_session meta = NoiseSuppressor._shared_meta hop_bytes = _HOP * 2 # int16 self._buf.extend(pcm_bytes) out = bytearray() # Lazily init per-session cache tensors to the model's declared shapes # (zeros). Any dynamic ('?') dim is treated as 1. if self._state is None: self._state = {} for name, shape, _t in meta['cache_inputs']: dims = [1 if (not isinstance(d, int) or d <= 0) else d for d in shape] self._state[name] = np.zeros(dims, dtype=np.float32) while len(self._buf) >= hop_bytes: frame = bytes(self._buf[:hop_bytes]) del self._buf[:hop_bytes] samples = np.frombuffer(frame, dtype=np.int16).astype(np.float32) / 32768.0 feed = {meta['primary_in']: samples.reshape(1, -1)} feed.update(self._state) results = sess.run(meta['output_names'], feed) enhanced = results[0].reshape(-1) # Roll updated caches back into per-session state for out_name, in_name in zip(meta['cache_out_names'], [c[0] for c in meta['cache_inputs']]): idx = meta['output_names'].index(out_name) self._state[in_name] = results[idx] clipped = np.clip(enhanced * 32768.0, -32768, 32767).astype(np.int16) out.extend(clipped.tobytes()) # If nothing produced yet (buffering first frame), return silence-free # empty bytes; caller handles empty as "no audio to forward this tick". return bytes(out)