Skip to content

C
coderai
Commit 7e5a925c authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

audio: STT backends, diarization, speaker embeddings/ID, word timestamps 

Transcription (/v1/audio/transcriptions):
- Route multipart uploads by peeked model (front _peek_model) so they land on
  the model's assigned engine instead of scattering by capability.
- New pluggable STT backends (codai/api/stt_backends.py): wav2vec2 (HF), vosk
  (CPU dirs), NVIDIA NeMo-Canary (isolated venv worker) + faster-whisper/whisper.
- Per-model language selection + validation; Canary source->target translation.
- OpenAI timestamp_granularities[]=word: top-level words[] in verbose_json
  (vosk/wav2vec2/faster-whisper/canary).
- VRAM-eviction tracking (acquire_stt_backend) + keep_resident co-residency for
  the small set; /v1/models surfaces languages + backends.

Diarization (/v1/audio/diarization + diarize=true):
- pyannote.audio isolated venv worker (codai/api/pyannote_worker.py,
  tools/pyannote_service.py) on the 3090; speaker-labeled segments; graceful
  degradation. Pinned torch/pyannote/hf_hub/matplotlib in requirements-pyannote.txt.

Speaker embeddings + recognition:
- /v1/audio/speaker-embeddings (ecapa via speechbrain, pyannote/wespeaker).
- Enroll/list/delete + identify/verify (codai/api/speaker_registry.py); persistent
  voiceprint store; diarization identify=true relabels turns with enrolled names.

New isolated-venv requirements: requirements-nemo.txt, requirements-pyannote.txt
(+ vosk in requirements.txt).
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01R6iugoNeEgt9PkajKrFyq4
parent f54428cf
Hide whitespace changes
Inline Side-by-side
Showing with 2418 additions and 26 deletions
codai/admin/routes.py
View file @ 7e5a925c
......@@ -2701,7 +2701,9 @@ async def api_model_configure(request: Request, username: str = Depends(require_
"turboquant", "engine", "engine_fallback",
"quant_backend", "kv_cache_budget_mb", "kv_cache_slots", "mmproj",
"auto_compact", "auto_compact_pct", "auto_compact_strategy",
"auto_compact_model", "suppress_reasoning"):
"auto_compact_model", "suppress_reasoning",
# Speech-to-text (wav2vec2 / vosk / NeMo-Canary) options.
"languages", "supports_translation", "keep_resident"):
if key in data:
entry[key] = data[key]
......
codai/admin/templates/models.html
View file @ 7e5a925c
......@@ -545,7 +545,7 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
<div style="display:grid;grid-template-columns:1fr 1fr;gap:.75rem">
<div class="form-row" style="margin:0">
<label class="form-label">Compute backend</label>
<select id="cfg-backend" class="form-input">
<select id="cfg-backend" class="form-input" onchange="_updateSttVisibility()">
<option value="auto">Auto-detect</option>
<option value="nvidia">NVIDIA (CUDA)</option>
<option value="vulkan">Vulkan</option>
......@@ -556,6 +556,9 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
<option value="k3">k3 (Kimi-K3 / kimi-k3-in-c)</option>
<option value="kt">ktransformers (SGLang, CPU+GPU)</option>
<option value="vllm">vLLM (high-concurrency, CUDA)</option>
<option value="wav2vec2">wav2vec2 / HF-ASR (speech-to-text)</option>
<option value="vosk">Vosk (offline STT, CPU, per-language dir)</option>
<option value="nemo">NVIDIA NeMo — Canary / Parakeet (isolated venv)</option>
</select>
</div>
<div class="form-row" style="margin:0">
......@@ -566,6 +569,24 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
</select>
</div>
</div>
<!-- Speech-to-text language options (wav2vec2 / vosk / NeMo-Canary) -->
<div id="cfg-stt-row" style="margin-top:.75rem;display:none">
<div class="form-row" style="margin:0">
<label class="form-label">Supported languages
<span class="muted">(comma-separated ISO codes, blank = accept any)</span></label>
<input id="cfg-stt-languages" class="form-input"
placeholder="en, it, de, fr, es">
</div>
<label class="form-label" style="display:flex;align-items:center;gap:.5rem;margin-top:.5rem">
<input type="checkbox" id="cfg-stt-translation">
Supports translation (NeMo/Canary: enables the request's <code>target_language</code>)
</label>
<label class="form-label" style="display:flex;align-items:center;gap:.5rem;margin-top:.5rem">
<input type="checkbox" id="cfg-stt-keep-resident">
Keep resident (co-reside with the other small STT models — skip routine
eviction; still yields to a larger model as a last resort)
</label>
</div>
<div class="form-row" id="cfg-engine-row" style="margin-top:.75rem;display:none">
<label class="form-label">Engine / card</label>
<select id="cfg-engine" class="form-input" onchange="_updateGpuSplitVisibility()">
......@@ -3364,6 +3385,10 @@ function openCfgModal(idx, cfgIdx){
(configs.length > 1 || (isNewConfig && m.in_config)) ? '' : 'none';
document.getElementById('cfg-backend').value = s.backend || 'auto';
document.getElementById('cfg-load-mode').value = s.load_mode || 'on-request';
document.getElementById('cfg-stt-languages').value = (s.languages || []).join(', ');
document.getElementById('cfg-stt-translation').checked = !!s.supports_translation;
document.getElementById('cfg-stt-keep-resident').checked = !!s.keep_resident;
_updateSttVisibility();
// Used VRAM
const usedVram = s.used_vram_gb != null ? s.used_vram_gb : null;
document.getElementById('cfg-used-vram').value = usedVram != null ? usedVram : '';
......@@ -3821,6 +3846,13 @@ function _collectColibri(){
return o;
}
// Show the speech-to-text language options only for STT backends.
function _updateSttVisibility(){
const be = (document.getElementById('cfg-backend').value || '');
const row = document.getElementById('cfg-stt-row');
if (row) row.style.display = (be === 'wav2vec2' || be === 'vosk' || be === 'nemo') ? '' : 'none';
}
async function saveModelConfig(){
const path = document.getElementById('cfg-path').value;
const maxGpu = parseFloat(document.getElementById('cfg-max-gpu').value);
......@@ -3840,6 +3872,10 @@ async function saveModelConfig(){
config_name: document.getElementById('cfg-config-name').value.trim() || null,
backend: document.getElementById('cfg-backend').value,
load_mode: document.getElementById('cfg-load-mode').value,
languages: (function(){ const v=document.getElementById('cfg-stt-languages').value.trim();
return v ? v.split(',').map(s=>s.trim()).filter(Boolean) : []; })(),
supports_translation: document.getElementById('cfg-stt-translation').checked,
keep_resident: document.getElementById('cfg-stt-keep-resident').checked,
used_vram_gb: isNaN(usedVram) ? null : usedVram,
force_vram_update: document.getElementById('cfg-force-vram-update').checked,
max_instances: parseInt(document.getElementById('cfg-max-instances').value) || 1,
......
codai/api/canary_worker.py 0 → 100644
View file @ 7e5a925c
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
"""Fully-managed NVIDIA NeMo worker (Canary / Parakeet ASR).
NeMo pins a stack (transformers, huggingface-hub, numba, etc.) that conflicts with
the coderai server's transformers 5.x, so — like parler-tts — it can't share this
venv. coderai owns the lifecycle: it uses a NeMo venv (baked into the image at
``/opt/coderai/nemo_venv`` when present, else built on demand from
``requirements-nemo.txt``), launches ``tools/canary_service.py`` in it as a local
HTTP service, health-checks it, and hands back the URL. The
``_RemoteNemoBackend.cleanup()`` calls :func:`stop_service` so normal model
eviction tears the process down.
"""
import os
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
_REPO_ROOT = Path(__file__).resolve().parents[2]
_SERVICE_SCRIPT = _REPO_ROOT / "tools" / "canary_service.py"
_REQUIREMENTS = _REPO_ROOT / "requirements-nemo.txt"
# Baked-into-image venv (built by build.sh) is preferred; fall back to a runtime
# build under the user's cache. Mirrors codai.api.vllm_worker.resolve_venv_dir.
_BAKED_VENV = Path("/opt/coderai/nemo_venv")
_lock = threading.RLock()
_services: dict = {} # model_name -> {"proc","port","url"}
_bootstrapped = False
def resolve_venv_dir() -> Path:
explicit = os.environ.get("CODERAI_NEMO_VENV")
if explicit:
return Path(explicit)
if (_BAKED_VENV / ("Scripts" if os.name == "nt" else "bin")).exists():
return _BAKED_VENV
return Path(os.path.expanduser("~/.coderai/nemo_venv"))
def _venv_python(venv: Path) -> Path:
return venv / ("Scripts" if os.name == "nt" else "bin") / (
"python.exe" if os.name == "nt" else "python")
def _pip_ok(py: Path) -> bool:
try:
return subprocess.run(
[str(py), "-c", "import nemo.collections.asr"],
capture_output=True).returncode == 0
except Exception:
return False
def _bootstrap_venv() -> Path:
"""Return the NeMo venv's python, building it from requirements-nemo.txt if a
baked/prebuilt one isn't already present. Idempotent."""
global _bootstrapped
venv = resolve_venv_dir()
py = _venv_python(venv)
if _bootstrapped and py.exists():
return py
if py.exists() and _pip_ok(py):
_bootstrapped = True
return py
if not py.exists():
print(f"[canary] creating isolated NeMo venv at {venv} …", flush=True)
venv.parent.mkdir(parents=True, exist_ok=True)
subprocess.run([sys.executable, "-m", "venv", str(venv)], check=True)
py = _venv_python(venv)
if not _pip_ok(py):
if not _REQUIREMENTS.exists():
raise RuntimeError(f"NeMo requirements file missing: {_REQUIREMENTS}")
print("[canary] installing NeMo (nemo_toolkit[asr]) into the isolated venv "
"(first run downloads several GB, this can take a while) …", flush=True)
subprocess.run([str(py), "-m", "pip", "install", "-U", "pip"], check=True)
subprocess.run([str(py), "-m", "pip", "install", "-r", str(_REQUIREMENTS)],
check=True)
if not _pip_ok(py):
raise RuntimeError("NeMo install did not yield an importable nemo.collections.asr")
_bootstrapped = True
return py
def _free_port() -> int:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
s.close()
return port
def _pump_logs(proc: subprocess.Popen, tail):
for line in proc.stdout:
line = line.rstrip()
if line:
tail.append(line)
print(f"[canary] {line}", flush=True)
def _health_ok(url: str) -> bool:
import requests
try:
r = requests.get(url + "/health", timeout=3)
return r.ok and bool(r.json().get("ok"))
except Exception:
return False
def ensure_service(model_name: str, config: dict = None,
ready_timeout: float = 1800.0) -> str:
"""Start (or reuse) the NeMo worker for ``model_name`` and return its base URL.
First call may build the venv and download the model, so the timeout is
generous. Raises RuntimeError if the service never comes up."""
config = config or {}
with _lock:
svc = _services.get(model_name)
if svc and svc["proc"].poll() is None and _health_ok(svc["url"]):
return svc["url"]
if svc and svc["proc"].poll() is not None:
_services.pop(model_name, None)
py = _bootstrap_venv()
port = _free_port()
url = f"http://127.0.0.1:{port}"
env = dict(os.environ)
# Share coderai's HF cache. NeMo pulls .nemo checkpoints from HF hub; we
# allow the fetch (coderai's downloader doesn't handle .nemo) but keep it
# in the shared cache so it's pulled once.
try:
from codai.models.cache import get_hf_hub_cache_dir
hub = get_hf_hub_cache_dir()
env["HF_HUB_CACHE"] = hub
env["HUGGINGFACE_HUB_CACHE"] = hub
except Exception:
pass
# Pin the worker to its configured GPU (the 3090 by default).
gpu = config.get("gpu_device")
if gpu is not None:
env["CUDA_VISIBLE_DEVICES"] = str(gpu)
cmd = [str(py), str(_SERVICE_SCRIPT), "--model", model_name,
"--host", "127.0.0.1", "--port", str(port)]
mp = config.get("model_path")
if mp:
cmd += ["--model-path", str(mp)]
proc = subprocess.Popen(
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
bufsize=1, env=env, cwd=str(_REPO_ROOT))
import collections
tail = collections.deque(maxlen=15)
threading.Thread(target=_pump_logs, args=(proc, tail), daemon=True).start()
_services[model_name] = {"proc": proc, "port": port, "url": url}
def _tail_msg():
joined = " | ".join(list(tail)[-5:]).strip()
return f". Last output: {joined}" if joined else ""
deadline = time.time() + ready_timeout
while time.time() < deadline:
if proc.poll() is not None:
raise RuntimeError(
f"NeMo worker exited (code {proc.returncode}) before becoming ready"
+ _tail_msg())
if _health_ok(url):
print(f"[canary] service ready for {model_name} at {url}", flush=True)
return url
time.sleep(2)
stop_service(model_name)
raise RuntimeError(f"NeMo worker for {model_name} did not become ready in time"
+ _tail_msg())
def is_running(model_name: str) -> bool:
with _lock:
svc = _services.get(model_name)
return bool(svc and svc["proc"].poll() is None)
def stop_service(model_name: str) -> None:
with _lock:
svc = _services.pop(model_name, None)
if not svc:
return
proc = svc["proc"]
if proc.poll() is None:
try:
proc.terminate()
proc.wait(timeout=10)
except Exception:
pass
if proc.poll() is None:
try:
proc.kill()
except Exception:
pass
print(f"[canary] service for {model_name} stopped", flush=True)
def stop_all() -> None:
for name in list(_services.keys()):
stop_service(name)
import atexit as _atexit
_atexit.register(stop_all)
codai/api/diarization.py 0 → 100644
View file @ 7e5a925c
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Speaker diarization via the isolated pyannote.audio worker.
Provides the plumbing shared by the standalone ``/v1/audio/diarization`` endpoint
and the ``diarize=true`` option on ``/v1/audio/transcriptions``:
* :func:`diarize_audio` — start/reuse the (VRAM-eviction-tracked) pyannote worker
and return speaker turns.
* :func:`merge_speakers` — attach a ``speaker`` label to each ASR segment by
maximal temporal overlap with the diarization turns.
"""
from typing import List, Optional
from codai.models.manager import multi_model_manager
class _DiarizerHandle:
"""Evictable handle for a running pyannote worker. Its ``cleanup`` stops the
subprocess, so the manager's LRU eviction frees the worker's VRAM like any
other loaded model."""
def __init__(self, model: str, url: str):
self.model = model
self.url = url
def cleanup(self):
try:
from codai.api import pyannote_worker
pyannote_worker.stop_service(self.model)
except Exception:
pass
def _get_diarizer(model: Optional[str], config: Optional[dict]) -> _DiarizerHandle:
"""Ensure the pyannote worker is running and registered as an evictable model."""
from codai.api import pyannote_worker
resolved = model or pyannote_worker._default_model()
key = f"diarization:{resolved}"
cfg = config or {}
# ensure_service is idempotent: it reuses a healthy worker and re-spawns a dead
# one (possibly on a new port). Always call it so a worker that died since the
# handle was cached self-heals — otherwise acquire_stt_backend would hand back a
# stale handle pointing at the dead worker's port (connection refused).
url = pyannote_worker.ensure_service(resolved, cfg)
def _loader():
return _DiarizerHandle(resolved, url)
needed = float(cfg.get("used_vram_gb") or 2.0)
# pyannote is small (~2GB) — default it into the co-resident set so it stays
# warm alongside the STT models (still last-resort-evictable for a big LLM).
keep_resident = bool(cfg.get("keep_resident", True))
handle = multi_model_manager.acquire_stt_backend(
key, needed, True, _loader, keep_resident=keep_resident)
# Refresh the URL in case a cached handle points at a re-spawned worker.
handle.url = url
return handle
def diarize_audio(audio_path: str, num_speakers=None, min_speakers=None,
max_speakers=None, model: Optional[str] = None,
config: Optional[dict] = None, identify: bool = False,
identify_backend: str = "ecapa", identify_threshold: float = 0.25) -> dict:
"""Return ``{"segments":[{start,end,speaker}], "num_speakers":N}`` for the audio.
When ``identify`` is set, the worker also embeds one slice per diarized speaker;
each is matched against the enrolled speaker registry and the diarization
labels (SPEAKER_00…) are replaced with the recognised names (falling back to
the pyannote label when unknown). Also returns ``speakers`` = {label: name}."""
import requests
handle = _get_diarizer(model, config)
params = {}
if num_speakers:
params["num_speakers"] = int(num_speakers)
if min_speakers:
params["min_speakers"] = int(min_speakers)
if max_speakers:
params["max_speakers"] = int(max_speakers)
if identify:
params["embed"] = identify_backend or "ecapa"
with open(audio_path, "rb") as f:
resp = requests.post(f"{handle.url}/diarize", params=params,
data=f.read(), timeout=1800,
headers={"Content-Type": "application/octet-stream"})
if resp.status_code != 200:
raise RuntimeError(f"pyannote worker error {resp.status_code}: {resp.text[:300]}")
j = resp.json()
if j.get("error"):
raise RuntimeError(j["error"])
segments = j.get("segments") or []
out = {"segments": segments, "num_speakers": j.get("num_speakers")}
if identify and j.get("speaker_embeddings"):
from codai.api import speaker_registry
label_to_name = {}
for label, emb in j["speaker_embeddings"].items():
res = speaker_registry.identify(emb, identify_backend or "ecapa",
threshold=identify_threshold)
label_to_name[label] = {"name": res["name"], "score": res.get("score")}
# Relabel segments with recognised names (keep pyannote label when unknown).
for s in segments:
m = label_to_name.get(s.get("speaker"))
if m and m["name"] != "unknown":
s["speaker_label"] = s["speaker"]
s["speaker"] = m["name"]
out["speakers"] = label_to_name
return out
def _overlap(a_start, a_end, b_start, b_end) -> float:
return max(0.0, min(a_end, b_end) - max(a_start, b_start))
def merge_speakers(asr_segments: List[dict], diar_segments: List[dict]) -> List[dict]:
"""Attach a ``speaker`` to each ASR segment: the diarization turn it overlaps
most. ASR segments with no overlap are left unlabeled. Returns a new list."""
out = []
for seg in asr_segments or []:
s0 = float(seg.get("start", 0.0) or 0.0)
s1 = float(seg.get("end", 0.0) or 0.0)
best_spk, best_ov = None, 0.0
for d in diar_segments or []:
ov = _overlap(s0, s1, float(d.get("start", 0.0)), float(d.get("end", 0.0)))
if ov > best_ov:
best_ov, best_spk = ov, d.get("speaker")
merged = dict(seg)
if best_spk is not None:
merged["speaker"] = best_spk
out.append(merged)
return out
codai/api/pyannote_worker.py 0 → 100644
View file @ 7e5a925c
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
"""Fully-managed pyannote.audio worker (speaker diarization).
pyannote.audio pins torch + pytorch-lightning + speechbrain, which conflict with
the coderai server's transformers 5.x stack — so, like NeMo/Canary, it runs in an
isolated venv. coderai owns the lifecycle: it uses a pyannote venv (baked into the
image at ``/opt/coderai/pyannote_venv`` when present, else built on demand from
``requirements-pyannote.txt``), launches ``tools/pyannote_service.py`` in it as a
local HTTP service, health-checks it, and hands back the URL.
The gated ``pyannote/speaker-diarization-3.1`` needs an HF token (accept its
conditions + ``pyannote/segmentation-3.0`` first). The worker passes ``HF_TOKEN``
/ ``HUGGINGFACE_TOKEN`` from the environment through to the service; the default
model can be pointed at an ungated mirror via ``CODERAI_DIARIZATION_MODEL`` so no
token is needed.
"""
import os
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
_REPO_ROOT = Path(__file__).resolve().parents[2]
_SERVICE_SCRIPT = _REPO_ROOT / "tools" / "pyannote_service.py"
_REQUIREMENTS = _REPO_ROOT / "requirements-pyannote.txt"
_BAKED_VENV = Path("/opt/coderai/pyannote_venv")
_lock = threading.RLock()
_services: dict = {} # model_name -> {"proc","port","url"}
_bootstrapped = False
def resolve_venv_dir() -> Path:
explicit = os.environ.get("CODERAI_PYANNOTE_VENV")
if explicit:
return Path(explicit)
if (_BAKED_VENV / ("Scripts" if os.name == "nt" else "bin")).exists():
return _BAKED_VENV
return Path(os.path.expanduser("~/.coderai/pyannote_venv"))
def _venv_python(venv: Path) -> Path:
return venv / ("Scripts" if os.name == "nt" else "bin") / (
"python.exe" if os.name == "nt" else "python")
def _pip_ok(py: Path) -> bool:
try:
return subprocess.run(
[str(py), "-c", "import pyannote.audio"],
capture_output=True).returncode == 0
except Exception:
return False
def _bootstrap_venv() -> Path:
global _bootstrapped
venv = resolve_venv_dir()
py = _venv_python(venv)
if _bootstrapped and py.exists():
return py
if py.exists() and _pip_ok(py):
_bootstrapped = True
return py
if not py.exists():
print(f"[pyannote] creating isolated venv at {venv} …", flush=True)
venv.parent.mkdir(parents=True, exist_ok=True)
subprocess.run([sys.executable, "-m", "venv", str(venv)], check=True)
py = _venv_python(venv)
if not _pip_ok(py):
if not _REQUIREMENTS.exists():
raise RuntimeError(f"pyannote requirements file missing: {_REQUIREMENTS}")
print("[pyannote] installing pyannote.audio into the isolated venv "
"(first run downloads several GB, this can take a while) …", flush=True)
subprocess.run([str(py), "-m", "pip", "install", "-U", "pip"], check=True)
subprocess.run([str(py), "-m", "pip", "install", "-r", str(_REQUIREMENTS)],
check=True)
if not _pip_ok(py):
raise RuntimeError("pyannote install did not yield an importable pyannote.audio")
_bootstrapped = True
return py
def _free_port() -> int:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
s.close()
return port
def _pump_logs(proc: subprocess.Popen, tail):
for line in proc.stdout:
line = line.rstrip()
if line:
tail.append(line)
print(f"[pyannote] {line}", flush=True)
def _health_ok(url: str) -> bool:
import requests
try:
r = requests.get(url + "/health", timeout=3)
return r.ok and bool(r.json().get("ok"))
except Exception:
return False
def _default_model() -> str:
# Official pyannote pipeline (canonical, best-supported). It's gated: needs an
# HF token (see _resolve_hf_token) + accepting the pyannote/speaker-diarization-3.1
# and pyannote/segmentation-3.0 conditions. Override with CODERAI_DIARIZATION_MODEL.
return os.environ.get("CODERAI_DIARIZATION_MODEL") or "pyannote/speaker-diarization-3.1"
def _resolve_hf_token(config: dict) -> str:
"""Find an HF token: request/config, then env, then a protected file in the
coderai config dir (~/.coderai/hf_token or /config/coderai/hf_token). Lets the
token be configured once without editing the container's launch command."""
tok = (config or {}).get("hf_token")
if tok:
return tok
for var in ("HF_TOKEN", "HUGGINGFACE_TOKEN", "HUGGING_FACE_HUB_TOKEN"):
if os.environ.get(var):
return os.environ[var]
for path in (os.environ.get("CODERAI_HF_TOKEN_FILE"),
os.path.expanduser("~/.coderai/hf_token"),
"/config/coderai/hf_token"):
if path and os.path.isfile(path):
try:
t = open(path).read().strip()
if t:
return t
except OSError:
pass
return ""
def ensure_service(model_name: str = None, config: dict = None,
ready_timeout: float = 1800.0) -> str:
"""Start (or reuse) the pyannote worker and return its base URL."""
config = config or {}
model_name = model_name or config.get("model_path") or _default_model()
with _lock:
svc = _services.get(model_name)
if svc and svc["proc"].poll() is None and _health_ok(svc["url"]):
return svc["url"]
if svc and svc["proc"].poll() is not None:
_services.pop(model_name, None)
py = _bootstrap_venv()
port = _free_port()
url = f"http://127.0.0.1:{port}"
env = dict(os.environ)
try:
from codai.models.cache import get_hf_hub_cache_dir
hub = get_hf_hub_cache_dir()
env["HF_HUB_CACHE"] = hub
env["HUGGINGFACE_HUB_CACHE"] = hub
except Exception:
pass
# Pass an HF token through for the gated model (harmless for ungated).
tok = _resolve_hf_token(config)
if tok:
env["HF_TOKEN"] = tok
gpu = config.get("gpu_device")
if gpu is not None:
env["CUDA_VISIBLE_DEVICES"] = str(gpu)
proc = subprocess.Popen(
[str(py), str(_SERVICE_SCRIPT), "--model", model_name,
"--host", "127.0.0.1", "--port", str(port)],
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
bufsize=1, env=env, cwd=str(_REPO_ROOT))
import collections
tail = collections.deque(maxlen=15)
threading.Thread(target=_pump_logs, args=(proc, tail), daemon=True).start()
_services[model_name] = {"proc": proc, "port": port, "url": url}
def _tail_msg():
joined = " | ".join(list(tail)[-5:]).strip()
return f". Last output: {joined}" if joined else ""
deadline = time.time() + ready_timeout
while time.time() < deadline:
if proc.poll() is not None:
raise RuntimeError(
f"pyannote worker exited (code {proc.returncode}) before becoming ready"
+ _tail_msg())
if _health_ok(url):
print(f"[pyannote] service ready for {model_name} at {url}", flush=True)
return url
time.sleep(2)
stop_service(model_name)
raise RuntimeError(f"pyannote worker for {model_name} did not become ready in time"
+ _tail_msg())
def is_running(model_name: str = None) -> bool:
model_name = model_name or _default_model()
with _lock:
svc = _services.get(model_name)
return bool(svc and svc["proc"].poll() is None)
def stop_service(model_name: str = None) -> None:
model_name = model_name or _default_model()
with _lock:
svc = _services.pop(model_name, None)
if not svc:
return
proc = svc["proc"]
if proc.poll() is None:
try:
proc.terminate()
proc.wait(timeout=10)
except Exception:
pass
if proc.poll() is None:
try:
proc.kill()
except Exception:
pass
print(f"[pyannote] service for {model_name} stopped", flush=True)
def stop_all() -> None:
for name in list(_services.keys()):
stop_service(name)
import atexit as _atexit
_atexit.register(stop_all)
codai/api/speaker_embeddings.py 0 → 100644
View file @ 7e5a925c
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Speaker-embedding extraction (voiceprints) via the isolated pyannote venv.
Backends (served by tools/pyannote_service.py's /embed):
* ``ecapa`` — speechbrain ECAPA-TDNN (ungated, 192-d). Default.
* ``pyannote`` — pyannote.audio embedding model (may be gated, uses HF_TOKEN).
* ``wespeaker``— pyannote wespeaker resnet embedding.
Reuses the same worker (and VRAM-eviction registration) as diarization.
"""
from typing import Optional
def get_speaker_embedding(audio_path: str, backend: str = "ecapa",
model: Optional[str] = None,
config: Optional[dict] = None) -> dict:
"""Return ``{"embedding":[...], "dim":N, "backend":..., "model":...}``."""
import requests
from codai.api.diarization import _get_diarizer
# _get_diarizer ensures the (VRAM-tracked, self-healing) pyannote worker is up
# and returns a handle with its URL; the worker serves /embed as well.
handle = _get_diarizer(None, config or {})
params = {"backend": backend or "ecapa"}
if model:
params["model"] = model
with open(audio_path, "rb") as f:
resp = requests.post(f"{handle.url}/embed", params=params, data=f.read(),
timeout=600, headers={"Content-Type": "application/octet-stream"})
if resp.status_code != 200:
raise RuntimeError(f"pyannote worker error {resp.status_code}: {resp.text[:300]}")
j = resp.json()
if j.get("error"):
raise RuntimeError(j["error"])
return j
codai/api/speaker_registry.py 0 → 100644
View file @ 7e5a925c
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Speaker identification registry — enrolled voiceprints + cosine matching.
A named speaker is one or more embeddings (averaged) produced by a specific
embedding backend (ecapa / wespeaker / …). Identification/verification compares a
query embedding against enrolled speakers of the SAME backend by cosine
similarity — embeddings from different backends are NOT comparable.
Persisted as JSON in the coderai config dir so enrollments survive restarts.
"""
import json
import os
import threading
import time
from typing import Optional
_lock = threading.RLock()
def _store_path() -> str:
p = os.environ.get("CODERAI_SPEAKER_STORE")
if p:
return p
for d in ("/config/coderai", os.path.expanduser("~/.coderai")):
if os.path.isdir(d):
return os.path.join(d, "speakers.json")
return os.path.join(os.path.expanduser("~/.coderai"), "speakers.json")
def _load() -> dict:
try:
with open(_store_path()) as f:
return json.load(f)
except Exception:
return {"speakers": {}}
def _save(data: dict) -> None:
path = _store_path()
os.makedirs(os.path.dirname(path), exist_ok=True)
tmp = path + ".tmp"
with open(tmp, "w") as f:
json.dump(data, f)
os.replace(tmp, path)
def _norm(v):
import numpy as np
a = np.asarray(v, dtype=np.float32)
n = np.linalg.norm(a)
return a / n if n > 0 else a
def _cosine(a, b) -> float:
import numpy as np
return float(np.dot(_norm(a), _norm(b)))
def enroll(name: str, embedding: list, backend: str, model: str = None) -> dict:
"""Add a sample for ``name``. Re-enrolling averages with existing samples
(running mean), improving the voiceprint. One voiceprint per (name, backend)."""
import numpy as np
name = (name or "").strip()
if not name:
raise ValueError("speaker name is required")
with _lock:
data = _load()
spk = data.setdefault("speakers", {})
key = f"{name}::{backend}"
rec = spk.get(key)
if rec is None:
rec = {"name": name, "backend": backend, "model": model,
"dim": len(embedding), "samples": 0,
"embedding": [0.0] * len(embedding)}
n = rec["samples"]
# running mean of the raw embeddings
prev = np.asarray(rec["embedding"], dtype=np.float32)
cur = np.asarray(embedding, dtype=np.float32)
if prev.shape != cur.shape:
prev = np.zeros_like(cur); n = 0
newmean = (prev * n + cur) / (n + 1)
rec["embedding"] = newmean.astype(float).tolist()
rec["samples"] = n + 1
rec["dim"] = len(rec["embedding"])
rec["model"] = model or rec.get("model")
rec["updated"] = None # timestamp stamped by caller if wanted
spk[key] = rec
_save(data)
return {"name": name, "backend": backend, "samples": rec["samples"],
"dim": rec["dim"]}
def list_speakers() -> list:
with _lock:
data = _load()
return [{"name": r["name"], "backend": r["backend"], "model": r.get("model"),
"samples": r.get("samples", 1), "dim": r.get("dim")}
for r in data.get("speakers", {}).values()]
def delete(name: str, backend: str = None) -> int:
name = (name or "").strip()
removed = 0
with _lock:
data = _load()
spk = data.get("speakers", {})
for key in [k for k, r in spk.items()
if r.get("name") == name and (backend is None or r.get("backend") == backend)]:
spk.pop(key, None)
removed += 1
if removed:
_save(data)
return removed
def get(name: str, backend: str) -> Optional[dict]:
with _lock:
for r in _load().get("speakers", {}).values():
if r.get("name") == name and r.get("backend") == backend:
return r
return None
def _enrolled_for_backend(backend: str) -> list:
data = _load()
return [r for r in data.get("speakers", {}).values() if r.get("backend") == backend]
def identify(embedding: list, backend: str, threshold: float = 0.25,
top_k: int = 3) -> dict:
"""Return the best-matching enrolled speaker for ``embedding`` (same backend).
``score`` is cosine similarity in [-1,1]; below ``threshold`` → name "unknown".
Also returns the top_k ranked candidates."""
with _lock:
enrolled = _enrolled_for_backend(backend)
ranked = sorted(
({"name": r["name"], "score": _cosine(embedding, r["embedding"])}
for r in enrolled),
key=lambda x: x["score"], reverse=True)
best = ranked[0] if ranked else None
name = best["name"] if (best and best["score"] >= threshold) else "unknown"
return {"name": name, "score": (best["score"] if best else None),
"threshold": threshold, "candidates": ranked[:top_k],
"enrolled": len(enrolled), "backend": backend}
def verify(embedding_a: list, embedding_b: list, threshold: float = 0.25) -> dict:
score = _cosine(embedding_a, embedding_b)
return {"same_speaker": score >= threshold, "score": score, "threshold": threshold}
codai/api/stt_backends.py 0 → 100644
View file @ 7e5a925c
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Pluggable speech-to-text backends (non whisper-server).
Mirrors the duck-typed family pattern of :mod:`codai.api.tts_backends`: each
backend is a plain class exposing a uniform ``transcribe`` method, and
:func:`load_stt_backend` is the factory that dispatches on a model's configured
``backend`` (falling back to a name heuristic).
Backends
--------
* ``wav2vec2`` — HF ``transformers`` CTC ASR (also HuBERT / MMS / Seamless). Runs
in-process on the GPU engine (torch is already in this venv).
* ``vosk`` — Kaldi-based offline STT. CPU. One model directory per language.
* ``nemo`` / ``canary`` — NVIDIA NeMo (Canary / Parakeet). NeMo pins a
conflicting stack, so it runs in an isolated venv worker
(:mod:`codai.api.canary_worker`); this module talks to it over HTTP.
Every backend's ``transcribe`` returns a dict::
{"text": str, "segments": [{"start": float, "end": float, "text": str}, ...],
"language": str | None}
``segments`` may be empty when a backend can't produce word/segment timings.
"""
import os
import subprocess
import tempfile
import wave
from typing import Optional
# --------------------------------------------------------------------------- #
# family detection
# --------------------------------------------------------------------------- #
def _family(model_name: str, config: Optional[dict]) -> str:
"""Classify a model into an STT backend family.
An explicit ``backend`` in config is authoritative (that's what the model
page sets); otherwise fall back to substring heuristics on the name."""
cfg = config or {}
b = (cfg.get("backend") or "").strip().lower()
if b in ("wav2vec2", "wav2vec", "transformers-asr", "hf-asr"):
return "wav2vec2"
if b == "vosk":
return "vosk"
if b in ("nemo", "canary", "parakeet"):
return "nemo"
n = (model_name or "").lower()
if "vosk" in n:
return "vosk"
if any(x in n for x in ("canary", "parakeet", "nemo", ".nemo")):
return "nemo"
if any(x in n for x in ("wav2vec", "hubert", "mms-", "seamless", "wavlm")):
return "wav2vec2"
# Default for an unknown non-whisper audio model: treat as an HF ASR model.
return "wav2vec2"
# --------------------------------------------------------------------------- #
# audio helpers
# --------------------------------------------------------------------------- #
def _to_wav_16k_mono(src_path: str) -> str:
"""Convert any audio file to a 16 kHz mono PCM-s16le WAV via ffmpeg.
Returns the path to a new temp WAV (caller deletes it). Raises if ffmpeg
fails. Vosk needs exactly this format; other backends can decode themselves."""
fd, out = tempfile.mkstemp(suffix=".wav")
os.close(fd)
proc = subprocess.run(
["ffmpeg", "-y", "-i", src_path, "-ar", "16000", "-ac", "1",
"-c:a", "pcm_s16le", "-f", "wav", out],
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
)
if proc.returncode != 0:
try:
os.unlink(out)
except OSError:
pass
raise RuntimeError(
"ffmpeg conversion failed: "
+ (proc.stderr.decode("utf-8", "replace").strip().splitlines() or [""])[-1])
return out
# --------------------------------------------------------------------------- #
# Wav2Vec2 / HF transformers ASR
# --------------------------------------------------------------------------- #
class _Wav2Vec2Backend:
"""HF ``transformers`` automatic-speech-recognition pipeline.
Works for Wav2Vec2 / HuBERT / WavLM / MMS / Seamless CTC checkpoints. Most of
these are language-specific (or, for MMS/XLSR, need an adapter per language);
``language`` is passed through when the underlying model supports it and is
otherwise informational only."""
family = "wav2vec2"
def __init__(self, model_name: str, model_path: Optional[str], config: dict):
self.model_name = model_name
self.config = config or {}
self._pipe = None
def _ensure(self):
if self._pipe is not None:
return
import torch
from transformers import pipeline
device = 0 if torch.cuda.is_available() else -1
target = self.config.get("model_path") or self.model_name
# chunk_length_s lets the pipeline transcribe arbitrarily long audio and
# emit per-chunk timestamps for segments.
self._pipe = pipeline(
"automatic-speech-recognition",
model=target,
device=device,
chunk_length_s=int(self.config.get("chunk_length_s", 30) or 30),
stride_length_s=int(self.config.get("stride_length_s", 5) or 5),
)
def transcribe(self, audio_path: str, language: Optional[str] = None,
prompt: Optional[str] = None, temperature: float = 0.0,
target_language: Optional[str] = None,
word_timestamps: bool = False) -> dict:
self._ensure()
gen_kwargs = {}
# MMS / Seamless accept a target language; plain wav2vec2 ignores it.
if language and self.config.get("language_aware"):
gen_kwargs["language"] = language
# CTC models emit per-word timestamps with return_timestamps="word".
mode = "word" if word_timestamps else "chunk"
try:
out = self._pipe(audio_path, return_timestamps=mode,
generate_kwargs=gen_kwargs or None)
except (TypeError, ValueError):
out = self._pipe(audio_path)
text = (out.get("text") if isinstance(out, dict) else str(out)) or ""
segments, words = [], []
for ch in (out.get("chunks") or []) if isinstance(out, dict) else []:
ts = ch.get("timestamp") or (None, None)
item = {"start": ts[0] or 0.0, "end": ts[1] or 0.0,
"text": (ch.get("text") or "").strip()}
if word_timestamps:
words.append({"word": item["text"], "start": item["start"],
"end": item["end"]})
else:
segments.append(item)
return {"text": text.strip(), "segments": segments, "words": words,
"language": language}
def cleanup(self):
self._pipe = None
# --------------------------------------------------------------------------- #
# Vosk
# --------------------------------------------------------------------------- #
class _VoskBackend:
"""Kaldi-based offline recognizer. CPU-only, one model dir per language.
The ``model_path`` (or ``model_name``) must point at an unpacked Vosk model
directory. Produces word-level timings, coalesced into ~one segment per Vosk
result block."""
family = "vosk"
def __init__(self, model_name: str, model_path: Optional[str], config: dict):
self.model_name = model_name
self.config = config or {}
self._model = None
self._model_dir = (self.config.get("model_path") or model_path
or model_name)
def _ensure(self):
if self._model is not None:
return
from vosk import Model, SetLogLevel
SetLogLevel(-1)
if not os.path.isdir(self._model_dir):
raise RuntimeError(
f"Vosk model directory not found: {self._model_dir}. Download/unpack "
"the language model and set its path in the model config.")
self._model = Model(self._model_dir)
def transcribe(self, audio_path: str, language: Optional[str] = None,
prompt: Optional[str] = None, temperature: float = 0.0,
target_language: Optional[str] = None,
word_timestamps: bool = False) -> dict:
import json
from vosk import KaldiRecognizer
self._ensure()
wav_path = _to_wav_16k_mono(audio_path)
try:
wf = wave.open(wav_path, "rb")
rec = KaldiRecognizer(self._model, wf.getframerate())
rec.SetWords(True)
segments, texts, words = [], [], []
while True:
data = wf.readframes(4000)
if len(data) == 0:
break
if rec.AcceptWaveform(data):
_collect_vosk_result(json.loads(rec.Result()), segments, texts, words)
_collect_vosk_result(json.loads(rec.FinalResult()), segments, texts, words)
wf.close()
return {"text": " ".join(t for t in texts if t).strip(),
"segments": segments, "words": words,
"language": language or self.config.get("language")}
finally:
try:
os.unlink(wav_path)
except OSError:
pass
def cleanup(self):
self._model = None
def _collect_vosk_result(r: dict, segments: list, texts: list, words: list = None):
txt = (r.get("text") or "").strip()
if not txt:
return
texts.append(txt)
wlist = r.get("result") or []
if wlist:
segments.append({"start": wlist[0].get("start", 0.0),
"end": wlist[-1].get("end", 0.0), "text": txt})
if words is not None:
for w in wlist:
words.append({"word": w.get("word", ""), "start": w.get("start", 0.0),
"end": w.get("end", 0.0)})
# --------------------------------------------------------------------------- #
# NVIDIA NeMo (Canary / Parakeet) — isolated-venv worker over HTTP
# --------------------------------------------------------------------------- #
class _RemoteNemoBackend:
"""Talks to the managed NeMo worker (isolated venv). Canary supports source
``language`` selection and X↔EN translation via ``target_language``."""
family = "nemo"
def __init__(self, model_name: str, config: dict, service_url: Optional[str] = None):
self.model_name = model_name
self.config = config or {}
self.service_url = (service_url or self.config.get("service_url") or "").rstrip("/")
def transcribe(self, audio_path: str, language: Optional[str] = None,
prompt: Optional[str] = None, temperature: float = 0.0,
target_language: Optional[str] = None,
word_timestamps: bool = False) -> dict:
import requests
params = {}
if language:
params["language"] = language
if target_language:
params["target_language"] = target_language
if word_timestamps:
params["word_timestamps"] = "1"
with open(audio_path, "rb") as f:
resp = requests.post(f"{self.service_url}/transcribe",
params=params, data=f.read(), timeout=1800,
headers={"Content-Type": "application/octet-stream"})
if resp.status_code != 200:
raise RuntimeError(f"NeMo worker error {resp.status_code}: {resp.text[:300]}")
j = resp.json()
if j.get("error"):
raise RuntimeError(j["error"])
return {"text": (j.get("text") or "").strip(),
"segments": j.get("segments") or [],
"words": j.get("words") or [],
"language": j.get("language") or language}
def cleanup(self):
# Tear the isolated worker down so eviction actually frees its VRAM.
try:
from codai.api import canary_worker
canary_worker.stop_service(self.model_name)
except Exception:
pass
# --------------------------------------------------------------------------- #
# factory
# --------------------------------------------------------------------------- #
def load_stt_backend(model_name: str, model_path: Optional[str], config: Optional[dict]):
"""Instantiate the STT backend for ``model_name`` (cached by the caller).
Dispatches on the model's configured ``backend`` (set from the model page),
falling back to a name heuristic. NeMo models are served by an isolated venv
worker, which this call bootstraps/starts on demand."""
config = config or {}
fam = _family(model_name, config)
# The requested `model` is usually an alias (e.g. "vosk-en"); the real model
# location — a Vosk dir, an HF id, or a NeMo checkpoint id — is the config's
# `path`. Resolve it so backends never get the alias as their model reference.
real = (model_path or config.get("model_path") or config.get("path") or model_name)
if fam == "vosk":
return _VoskBackend(real, real, config)
if fam == "nemo":
if config.get("service_url"):
return _RemoteNemoBackend(real, config)
from codai.api import canary_worker
url = canary_worker.ensure_service(real, config)
return _RemoteNemoBackend(real, {**config, "service_url": url}, service_url=url)
return _Wav2Vec2Backend(real, real, config)
# Family names that this module (not whisper-server / faster-whisper) handles.
STT_FAMILIES = {"wav2vec2", "vosk", "nemo"}
def resolve_family(model_name: str, config: Optional[dict]) -> str:
"""Public helper for the dispatcher: which family serves this model."""
return _family(model_name, config)
codai/api/transcriptions.py
View file @ 7e5a925c
......@@ -25,7 +25,7 @@ import tempfile
from fastapi import APIRouter, HTTPException, UploadFile, File, Form
from fastapi.responses import PlainTextResponse
from typing import Optional
from typing import List, Optional
# Maximum upload size: 100 MB
_MAX_AUDIO_BYTES = 100 * 1024 * 1024
......@@ -65,19 +65,26 @@ def _seconds_to_vtt_time(s: float) -> str:
return f"{h:02d}:{m:02d}:{sec:06.3f}"
def _format_response(fmt: str, text: str, segments: list):
"""Format a transcription result according to the requested response_format."""
def _format_response(fmt: str, text: str, segments: list, words: list = None):
"""Format a transcription result according to the requested response_format.
``words`` (when timestamp_granularities[]=word was requested) is surfaced as a
top-level ``words`` array in verbose_json, OpenAI-style."""
fmt = (fmt or "json").lower()
if fmt == "text":
return PlainTextResponse(text)
def _spk_prefix(seg):
spk = seg.get("speaker")
return f"[{spk}] " if spk else ""
if fmt == "srt":
lines = []
for i, seg in enumerate(segments, 1):
start = _seconds_to_srt_time(seg.get("start", 0))
end = _seconds_to_srt_time(seg.get("end", 0))
lines.append(f"{i}\n{start} --> {end}\n{seg['text'].strip()}\n")
lines.append(f"{i}\n{start} --> {end}\n{_spk_prefix(seg)}{seg['text'].strip()}\n")
srt_body = "\n".join(lines) if lines else f"1\n00:00:00,000 --> 00:00:00,000\n{text}\n"
return PlainTextResponse(srt_body, media_type="text/plain")
......@@ -86,13 +93,13 @@ def _format_response(fmt: str, text: str, segments: list):
for seg in segments:
start = _seconds_to_vtt_time(seg.get("start", 0))
end = _seconds_to_vtt_time(seg.get("end", 0))
lines.append(f"{start} --> {end}\n{seg['text'].strip()}\n")
lines.append(f"{start} --> {end}\n{_spk_prefix(seg)}{seg['text'].strip()}\n")
if not segments:
lines.append(f"00:00:00.000 --> 00:00:00.000\n{text}\n")
return PlainTextResponse("\n".join(lines), media_type="text/vtt")
if fmt == "verbose_json":
return {
out = {
"task": "transcribe",
"language": "unknown",
"duration": segments[-1].get("end", 0) if segments else 0,
......@@ -103,15 +110,66 @@ def _format_response(fmt: str, text: str, segments: list):
"start": s.get("start", 0),
"end": s.get("end", 0),
"text": s.get("text", "").strip(),
**({"speaker": s["speaker"]} if s.get("speaker") else {}),
}
for i, s in enumerate(segments)
],
}
if words:
out["words"] = [
{"word": w.get("word", ""), "start": w.get("start", 0),
"end": w.get("end", 0)}
for w in words
]
return out
# Default: json
# Default: json. Include segments when they carry speaker labels (diarize=true),
# so a plain json request still surfaces who-spoke-when.
if any(s.get("speaker") for s in (segments or [])):
return {
"text": text,
"segments": [
{"start": s.get("start", 0), "end": s.get("end", 0),
"text": s.get("text", "").strip(),
**({"speaker": s["speaker"]} if s.get("speaker") else {})}
for s in segments
],
}
return {"text": text}
# =============================================================================
# Language selection
# =============================================================================
def _validate_language(model, config: dict, language, target_language):
"""Validate requested language(s) against the model's declared languages.
A model config may carry ``languages`` (a list of supported ISO codes/names).
When it does and the caller passes a ``language``/``target_language`` not in
that list, fail with a clear 400 listing what's allowed. Models that declare
no languages accept anything (pass-through). ``target_language`` additionally
requires the model to advertise ``supports_translation``."""
langs = config.get("languages") or []
norm = {str(x).strip().lower() for x in langs}
if language and norm and str(language).strip().lower() not in norm:
raise HTTPException(
status_code=400,
detail=f"Language '{language}' not supported by '{model}'. "
f"Supported: {sorted(langs)}")
if target_language:
if not config.get("supports_translation"):
raise HTTPException(
status_code=400,
detail=f"Model '{model}' does not support translation "
"(target_language). Use a NeMo/Canary model that does.")
if norm and str(target_language).strip().lower() not in norm:
raise HTTPException(
status_code=400,
detail=f"Target language '{target_language}' not supported by "
f"'{model}'. Supported: {sorted(langs)}")
# =============================================================================
# Router and Endpoints
# =============================================================================
......@@ -127,9 +185,22 @@ async def create_transcription(
prompt: Optional[str] = Form(None),
response_format: Optional[str] = Form("json"),
temperature: Optional[float] = Form(0.0),
target_language: Optional[str] = Form(None),
diarize: Optional[bool] = Form(False),
num_speakers: Optional[int] = Form(None),
timestamp_granularities: Optional[List[str]] = Form(None),
timestamp_granularities_br: Optional[List[str]] = Form(None,
alias="timestamp_granularities[]"),
):
"""
Audio transcription endpoint (OpenAI-compatible).
``language`` selects the source language (validated against the model's
declared languages when it declares any). ``target_language`` requests
translation into that language for backends that support it (NVIDIA Canary);
other backends ignore it. ``diarize=true`` additionally runs pyannote speaker
diarization and tags each returned segment with a ``speaker`` (optionally
hinted by ``num_speakers``).
"""
file_content = await file.read()
if len(file_content) > _MAX_AUDIO_BYTES:
......@@ -144,9 +215,13 @@ async def create_transcription(
model=model or "",
)
task_registry.start(_tid)
_grans = [g.lower() for g in ((timestamp_granularities or [])
+ (timestamp_granularities_br or []))]
_word_ts = "word" in _grans
try:
_resp = await _run_transcription(
file_content, model, language, prompt, response_format, temperature, file)
file_content, model, language, prompt, response_format, temperature, file,
target_language, diarize, num_speakers, _word_ts)
task_registry.finish(_tid, "done")
return _resp
except HTTPException:
......@@ -157,10 +232,248 @@ async def create_transcription(
raise
@router.post("/v1/audio/diarization", summary="Speaker diarization (who spoke when)")
async def create_diarization(
file: UploadFile = File(...),
model: Optional[str] = Form(None),
num_speakers: Optional[int] = Form(None),
min_speakers: Optional[int] = Form(None),
max_speakers: Optional[int] = Form(None),
response_format: Optional[str] = Form("json"),
identify: Optional[bool] = Form(False),
identify_backend: Optional[str] = Form("ecapa"),
identify_threshold: Optional[float] = Form(0.25),
):
"""Standalone speaker diarization via pyannote.audio.
Returns ``{"segments":[{start,end,speaker}], "num_speakers":N}`` (or srt/vtt).
``model`` selects the pyannote pipeline (defaults to the configured/ungated
one); ``num_speakers`` pins the count, or bound it with min/max_speakers.
``identify=true`` matches each diarized speaker against the enrolled registry
and replaces the labels with recognised names (adds ``speakers`` = {label:name})."""
file_content = await file.read()
if len(file_content) > _MAX_AUDIO_BYTES:
raise HTTPException(status_code=413, detail="Audio file too large (max 100 MB)")
from codai.tasks import task_registry
_tid = task_registry.register(
"diarization", title=(file.filename or "audio")[:80], model=model or "")
task_registry.start(_tid)
fd, tmp_path = tempfile.mkstemp(suffix=".audio")
with os.fdopen(fd, "wb") as f:
f.write(file_content)
try:
import codai.api.diarization as diarization
result = await asyncio.to_thread(
diarization.diarize_audio, tmp_path, num_speakers, min_speakers, max_speakers,
model, None, bool(identify), identify_backend or "ecapa",
identify_threshold if identify_threshold is not None else 0.25)
task_registry.finish(_tid, "done")
fmt = (response_format or "json").lower()
if fmt in ("srt", "vtt"):
segs = [{"start": s["start"], "end": s["end"], "text": "",
"speaker": s["speaker"]} for s in result.get("segments", [])]
return _format_response(response_format, "", segs)
return result
except HTTPException:
task_registry.finish(_tid, "error")
raise
except Exception as e:
task_registry.finish(_tid, "error", str(e)[:200])
raise HTTPException(status_code=500, detail=f"Diarization error: {e}")
finally:
try:
os.unlink(tmp_path)
except OSError:
pass
@router.post("/v1/audio/speaker-embeddings", summary="Extract a speaker embedding (voiceprint)")
async def create_speaker_embedding(
file: UploadFile = File(...),
model: Optional[str] = Form(None),
backend: Optional[str] = Form("ecapa"),
):
"""Speaker-embedding (voiceprint) endpoint.
Returns a fixed-dim embedding for the speaker in the audio, usable for
verification / clustering. ``backend``: ``ecapa`` (speechbrain ECAPA-TDNN,
ungated, default), ``pyannote`` / ``wespeaker`` (pyannote embedding models,
may need HF_TOKEN). ``model`` overrides the backend's default checkpoint."""
file_content = await file.read()
if len(file_content) > _MAX_AUDIO_BYTES:
raise HTTPException(status_code=413, detail="Audio file too large (max 100 MB)")
from codai.tasks import task_registry
_tid = task_registry.register(
"speaker-embedding", title=(file.filename or "audio")[:80], model=(backend or "ecapa"))
task_registry.start(_tid)
fd, tmp_path = tempfile.mkstemp(suffix=".audio")
with os.fdopen(fd, "wb") as f:
f.write(file_content)
try:
import codai.api.speaker_embeddings as spk
result = await asyncio.to_thread(
spk.get_speaker_embedding, tmp_path, backend or "ecapa", model)
task_registry.finish(_tid, "done")
return {
"object": "list",
"model": result.get("model"),
"backend": result.get("backend"),
"dim": result.get("dim"),
"data": [{"object": "embedding", "index": 0,
"embedding": result.get("embedding")}],
}
except HTTPException:
task_registry.finish(_tid, "error")
raise
except Exception as e:
task_registry.finish(_tid, "error", str(e)[:200])
raise HTTPException(status_code=500, detail=f"Speaker-embedding error: {e}")
finally:
try:
os.unlink(tmp_path)
except OSError:
pass
async def _embed_upload(file, backend, model):
"""Read an upload to a temp file and return its speaker embedding dict."""
file_content = await file.read()
if len(file_content) > _MAX_AUDIO_BYTES:
raise HTTPException(status_code=413, detail="Audio file too large (max 100 MB)")
fd, tmp_path = tempfile.mkstemp(suffix=".audio")
with os.fdopen(fd, "wb") as f:
f.write(file_content)
try:
import codai.api.speaker_embeddings as spk
return await asyncio.to_thread(
spk.get_speaker_embedding, tmp_path, backend or "ecapa", model)
finally:
try:
os.unlink(tmp_path)
except OSError:
pass
@router.post("/v1/audio/speakers", summary="Enroll a speaker voiceprint")
async def enroll_speaker(
file: UploadFile = File(...),
name: str = Form(...),
backend: Optional[str] = Form("ecapa"),
model: Optional[str] = Form(None),
):
"""Enroll (or add a sample to) a named speaker from an audio clip. Re-enrolling
the same name averages the voiceprint. Use the SAME backend for enroll +
identify (embeddings from different backends aren't comparable)."""
emb = await _embed_upload(file, backend, model)
from codai.api import speaker_registry
try:
return speaker_registry.enroll(name, emb["embedding"], emb["backend"], emb.get("model"))
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
@router.get("/v1/audio/speakers", summary="List enrolled speakers")
async def list_speakers_endpoint():
from codai.api import speaker_registry
return {"speakers": speaker_registry.list_speakers()}
@router.delete("/v1/audio/speakers/{name}", summary="Delete an enrolled speaker")
async def delete_speaker_endpoint(name: str, backend: Optional[str] = None):
from codai.api import speaker_registry
removed = speaker_registry.delete(name, backend)
if not removed:
raise HTTPException(status_code=404, detail=f"No enrolled speaker '{name}'")
return {"deleted": removed, "name": name}
@router.post("/v1/audio/speaker-identify", summary="Identify the speaker in an audio clip")
async def identify_speaker_endpoint(
file: UploadFile = File(...),
backend: Optional[str] = Form("ecapa"),
model: Optional[str] = Form(None),
threshold: Optional[float] = Form(0.25),
):
"""Return the best-matching enrolled speaker (or "unknown") for the clip."""
emb = await _embed_upload(file, backend, model)
from codai.api import speaker_registry
return speaker_registry.identify(emb["embedding"], emb["backend"],
threshold=threshold if threshold is not None else 0.25)
@router.post("/v1/audio/speaker-verify", summary="Verify a clip against an enrolled speaker")
async def verify_speaker_endpoint(
file: UploadFile = File(...),
name: str = Form(...),
backend: Optional[str] = Form("ecapa"),
model: Optional[str] = Form(None),
threshold: Optional[float] = Form(0.25),
):
"""Verify whether the clip is the enrolled speaker ``name`` (same/different)."""
emb = await _embed_upload(file, backend, model)
from codai.api import speaker_registry
rec = speaker_registry.get(name, emb["backend"])
if rec is None:
raise HTTPException(status_code=404,
detail=f"No enrolled speaker '{name}' for backend '{emb['backend']}'")
out = speaker_registry.verify(emb["embedding"], rec["embedding"],
threshold=threshold if threshold is not None else 0.25)
out["name"] = name
return out
async def _run_transcription(
file_content: bytes, model: str, language, prompt, response_format, temperature, file
file_content: bytes, model: str, language, prompt, response_format, temperature, file,
target_language=None, diarize=False, num_speakers=None, word_timestamps=False
):
"""Core transcription logic; registered as a task by create_transcription()."""
# When diarization is requested, keep a temp copy of the audio for pyannote and
# wrap _format_response so every return path merges speaker labels into the
# segments. Runs in a thread (pyannote load/inference is blocking).
_diar_path = None
if diarize:
_dfd, _diar_path = tempfile.mkstemp(suffix=".audio")
with os.fdopen(_dfd, "wb") as _df:
_df.write(file_content)
async def _finalize(fmt, text, segments, words=None):
segs = segments or []
if diarize:
try:
diar = await asyncio.to_thread(
diarization.diarize_audio, _diar_path, num_speakers)
from codai.api.diarization import merge_speakers
if segs:
segs = merge_speakers(segs, diar.get("segments") or [])
else:
# No ASR timing (e.g. whisper-server text-only): surface the raw
# speaker turns as segments so who-spoke-when is still returned.
segs = diar.get("segments") or []
except Exception as e:
print(f"Diarization failed (returning transcription without speakers): {e}")
return _format_response(fmt, text, segs, words if word_timestamps else None)
import codai.api.diarization as diarization
try:
return await _run_transcription_inner(
file_content, model, language, prompt, response_format, temperature, file,
target_language, _finalize, word_timestamps)
finally:
if _diar_path:
try:
os.unlink(_diar_path)
except OSError:
pass
async def _run_transcription_inner(
file_content: bytes, model: str, language, prompt, response_format, temperature, file,
target_language, _finalize, word_timestamps=False
):
"""Backend dispatch; calls ``_finalize(fmt, text, segments, words)`` at each
return point so diarization + word timestamps are handled uniformly."""
# Check if the requested model maps to a configured whisper-server instance first.
# Try alias round-robin resolution before direct ID lookup.
whisper_model_id = multi_model_manager.resolve_whisper_alias_model_id(model)
......@@ -186,7 +499,7 @@ async def _run_transcription(
)
if "error" in result:
raise HTTPException(status_code=500, detail=result["error"])
return _format_response(response_format, result.get("text", ""), [])
return await _finalize(response_format, result.get("text", ""), [])
# Use the manager to resolve the model and manage VRAM
model_info = await asyncio.to_thread(
......@@ -202,6 +515,12 @@ async def _run_transcription(
model_name = model_info['model_name']
model_key = model_info['model_key']
whisper_model = model_info['model_object']
audio_config = model_info.get('config') or {}
# The manager stores the runtime-kwargs dict as `config`, with the full
# models.json entry nested under `_raw_cfg`. Backend/languages/model_path live
# there — read from it (falling back to the flat dict for older shapes).
_raw = audio_config.get('_raw_cfg') if isinstance(audio_config, dict) else None
stt_cfg = _raw if isinstance(_raw, dict) else (audio_config or {})
if not model_name:
raise HTTPException(
......@@ -209,6 +528,26 @@ async def _run_transcription(
detail="Audio transcription not configured. Use --audio-model or --whisper-server."
)
cfg_backend = (stt_cfg.get('backend') or '').strip().lower()
# A whisper.cpp GGUF (or a model configured with the whisper-server backend)
# cannot be loaded by faster-whisper — handing it a .gguf path or a whisper
# alias makes faster-whisper raise the opaque "Invalid model size". That only
# happens when the request reached an engine that has no whisper-server
# registered for this model (routing should pin it to its owner engine). Fail
# with an actionable message instead of the misleading size error.
if cfg_backend == 'whisper-server' or (
isinstance(model_name, str) and model_name.lower().endswith(".gguf")):
raise HTTPException(
status_code=503,
detail=(f"'{model}' is a whisper-server (whisper.cpp GGUF) model but no "
"whisper-server is available on the engine handling this request. "
"It must run on its assigned engine.")
)
# Validate the requested language against the model's declared languages, if any.
_validate_language(model, stt_cfg, language, target_language)
# Determine a safe file extension from the upload's content-type or filename,
# never trusting the raw user-supplied value for arbitrary suffixes.
raw_ext = os.path.splitext(file.filename or '')[1].lower()
......@@ -220,6 +559,47 @@ async def _run_transcription(
tmp_path = tmp.name
try:
# Pluggable STT backends (wav2vec2 / vosk / NVIDIA NeMo-Canary), selected
# by the model's configured `backend`. These are NOT faster-whisper and
# must be dispatched before the whisper fallback below.
import codai.api.stt_backends as stt_backends
family = stt_backends.resolve_family(model_name, stt_cfg)
if family in stt_backends.STT_FAMILIES:
# VRAM accounting: NeMo runs on the 3090 via its isolated worker;
# wav2vec2 uses this process's GPU (CUDA when available); vosk is CPU.
uses_gpu = False
needed_gb = 0.0
if family == 'nemo':
uses_gpu = True
needed_gb = float(stt_cfg.get('used_vram_gb') or 6.0)
elif family == 'wav2vec2':
try:
import torch
uses_gpu = bool(torch.cuda.is_available())
except Exception:
uses_gpu = False
needed_gb = float(stt_cfg.get('used_vram_gb') or 3.0) if uses_gpu else 0.0
def _run_backend():
def _loader():
return stt_backends.load_stt_backend(
model_name, stt_cfg.get('model_path'), stt_cfg)
# Evicts other models for `needed_gb` before loading a GPU backend,
# then registers it so the normal LRU eviction can free it later.
# keep_resident (models.json) marks the small co-resident set so it
# isn't swapped out by routine eviction (only as a last resort).
backend = multi_model_manager.acquire_stt_backend(
model_key, needed_gb, uses_gpu, _loader,
keep_resident=bool(stt_cfg.get('keep_resident')))
return backend.transcribe(
tmp_path, language=language, prompt=prompt,
temperature=temperature, target_language=target_language,
word_timestamps=word_timestamps)
result = await asyncio.to_thread(_run_backend)
return await _finalize(
response_format, result.get("text", ""), result.get("segments") or [],
result.get("words") or [])
# Try faster-whisper first
try:
from faster_whisper import WhisperModel
......@@ -238,14 +618,16 @@ async def _run_transcription(
language=language,
initial_prompt=prompt,
temperature=temperature,
word_timestamps=word_timestamps,
)
# Materialise the generator so we have all segment data
segments = [
{"start": s.start, "end": s.end, "text": s.text}
for s in raw_segments
]
segments, words = [], []
for s in raw_segments:
segments.append({"start": s.start, "end": s.end, "text": s.text})
for w in (getattr(s, "words", None) or []):
words.append({"word": w.word, "start": w.start, "end": w.end})
full_text = "".join(s["text"] for s in segments)
return _format_response(response_format, full_text.strip(), segments)
return await _finalize(response_format, full_text.strip(), segments, words)
except ImportError:
pass
......@@ -274,7 +656,7 @@ async def _run_transcription(
text += segment.get('text', '')
# whispercpp does not expose per-segment timestamps easily
return _format_response(response_format, text.strip(), [])
return await _finalize(response_format, text.strip(), [])
except ImportError as e:
raise HTTPException(
......
codai/frontproxy/app.py
View file @ 7e5a925c
......@@ -950,10 +950,47 @@ class FrontProxy:
@staticmethod
def _peek_model(body: bytes, content_type: str) -> Optional[str]:
if not body or "application/json" not in (content_type or "").lower():
ct = (content_type or "").lower()
if not body:
return None
if "application/json" in ct:
try:
return (json.loads(body) or {}).get("model")
except Exception:
return None
# multipart/form-data (audio transcriptions/translations, image edits):
# the `model` is a small text field. Extract it WITHOUT fully parsing the
# body — it carries the (up to 100MB) uploaded file. Routing needs the
# model to honour its engine pin (e.g. a whisper-server pinned to radeon);
# without it the request routes by capability alone and can land on an
# engine that has no whisper-server registered → "Invalid model size".
if "multipart/form-data" in ct:
return FrontProxy._peek_model_multipart(body)
return None
@staticmethod
def _peek_model_multipart(body: bytes) -> Optional[str]:
"""Pull the value of the `model` form field out of a multipart body.
Anchors on the part header ``Content-Disposition: … name="model"`` and
reads the value up to the next CRLF. A ``bytes.find`` for a short needle
over the whole body is C-fast even at 100MB, and the file part's own
``name="file"`` never matches this anchor."""
try:
return (json.loads(body) or {}).get("model")
marker = body.find(b'name="model"\r\n')
if marker == -1:
marker = body.find(b'name="model"')
if marker == -1:
return None
sep = body.find(b"\r\n\r\n", marker)
if sep == -1:
return None
start = sep + 4
end = body.find(b"\r\n", start)
if end == -1:
return None
val = body[start:end].decode("utf-8", "replace").strip()
return val or None
except Exception:
return None
......
codai/frontproxy/router.py
View file @ 7e5a925c
......@@ -34,6 +34,11 @@ _INFERENCE_PATHS = {
"/v1/images/edits",
"/v1/audio/speech",
"/v1/audio/transcriptions",
"/v1/audio/diarization",
"/v1/audio/speaker-embeddings",
"/v1/audio/speakers",
"/v1/audio/speaker-identify",
"/v1/audio/speaker-verify",
"/v1/video/generations",
}
......@@ -100,7 +105,10 @@ def required_capability(model: Optional[str], path: Optional[str] = None,
alias, then an unambiguous name marker (a Kimi name claimed by BOTH colibri and k3
falls through — the model must pin a backend)."""
p = (path or "").split("?", 1)[0].rstrip("/")
if p == "/v1/audio/transcriptions" or (backend or "") == "whisper-server":
if p in ("/v1/audio/transcriptions", "/v1/audio/diarization",
"/v1/audio/speaker-embeddings", "/v1/audio/speakers",
"/v1/audio/speaker-identify", "/v1/audio/speaker-verify") \
or (backend or "") == "whisper-server":
return "whisper"
# 1. Explicit engine-backend pin wins (authoritative), like the manager resolver.
b = (backend or "").lower()
......@@ -158,6 +166,18 @@ def pick_engine(registry: EngineRegistry, path: str, method: str,
if method.upper() == "POST" and is_inference_path(path):
cap = required_cap
# Speaker diarization (pyannote) must run on a CUDA engine to use the GPU.
# When the request carries no model to pin it, prefer the primary (nvidia)
# so it doesn't land on a Vulkan/CPU engine and run diarization on CPU. An
# explicit model pin (below) still wins.
if not pinned and not model and (path or "").split("?", 1)[0].rstrip("/") \
in ("/v1/audio/diarization", "/v1/audio/speaker-embeddings",
"/v1/audio/speakers", "/v1/audio/speaker-identify",
"/v1/audio/speaker-verify"):
prim = registry.primary()
if prim is not None and prim.can_serve(cap) and prim.is_alive():
return prim
# 0. Per-model pin (models.json "engine") is a HARD constraint: the model
# runs on that engine or not at all. Route there if it can serve and its
# process is alive — even when it's busy (mid-generation → failing health
......
codai/main.py
View file @ 7e5a925c
......@@ -984,6 +984,13 @@ def main():
)
else:
multi_model_manager.set_audio_model(mid, config=_model_cfg(m, "audio"))
# Register the alias→path mapping (like image/vision models) so a
# request for e.g. "vosk-en" / "wav2vec2-it" / "canary" resolves to the
# model's real path and carries its config — otherwise the STT backend
# receives the bare alias as its model reference. (whisper-server
# entries take the branch above and manage their own round-robin alias.)
if isinstance(m, dict) and m.get("alias"):
multi_model_manager.set_model_alias(m["alias"], mid)
# Image models
image_models = models_config.get("image_models", [])
......
codai/models/capabilities.py
View file @ 7e5a925c
......@@ -58,6 +58,7 @@ class ModelCapabilities:
speech_to_text: bool = False # Whisper transcription
text_to_speech: bool = False # Kokoro, Bark, XTTS
subtitle_generation: bool = False # WhisperX / forced alignment subtitles
speaker_diarization: bool = False # pyannote "who spoke when"
# Audio: generation & manipulation
audio_generation: bool = False # MusicGen, AudioLDM2, StableAudio
......@@ -170,8 +171,13 @@ def detect_model_capabilities(model_name: str) -> ModelCapabilities:
caps.audio_to_audio = True
return caps
if any(x in n for x in ['pyannote', 'diariz', 'speaker-diarization']):
caps.speaker_diarization = True
return caps
if any(x in n for x in ['whisper', 'faster-whisper', 'distil-whisper',
'wav2vec', 'hubert', 'seamless']):
'wav2vec', 'hubert', 'seamless', 'wavlm', 'mms-',
'canary', 'parakeet', 'nemo', 'vosk']):
caps.speech_to_text = True
caps.subtitle_generation = True
return caps
......
codai/models/manager.py
View file @ 7e5a925c
......@@ -820,6 +820,75 @@ class WhisperServerManager:
return False
return self.process.poll() is None
# Where the whisper.cpp runtime libs ship in the image.
_WHISPER_LIB_DIR = "/opt/coderai/local-libs"
def _whisper_workdir(self):
"""A writable working directory for the whisper-server subprocess.
whisper.cpp's ``--convert`` shells out to ffmpeg and writes its temp files
relative to the process CWD. The engine runs in /opt/coderai/app (read-only),
so conversion fails with "FFmpeg conversion failed." unless we give the
subprocess a writable cwd."""
import tempfile
candidates = [
os.environ.get("CODERAI_TMP"),
"/cache/coderai-tmp",
tempfile.gettempdir(),
]
for base in candidates:
if not base:
continue
d = os.path.join(base, "coderai-whisper-work")
try:
os.makedirs(d, exist_ok=True)
probe = os.path.join(d, ".wtest")
with open(probe, "w"):
pass
os.remove(probe)
return d
except OSError:
continue
return None
def _whisper_lib_env(self) -> dict:
"""Return an environment for the whisper-server subprocess with its runtime
libs on LD_LIBRARY_PATH, and the libwhisper soname symlink ensured.
The binary's RUNPATH points only at the (absent) build-machine dirs, so it
needs LD_LIBRARY_PATH to find libwhisper.so.1 and its ggml libs; and the
image ships libwhisper.so.1.8.3 without the libwhisper.so.1 soname symlink
the loader asks for. Without both, launch fails instantly."""
import glob
import tempfile
env = os.environ.copy()
lib_dirs = []
libd = self._WHISPER_LIB_DIR
if os.path.isdir(libd):
lib_dirs.append(libd)
soname = os.path.join(libd, "libwhisper.so.1")
if not os.path.exists(soname):
cands = sorted(glob.glob(os.path.join(libd, "libwhisper.so.1.*")))
if cands:
try:
os.symlink(cands[-1], soname)
except OSError:
# local-libs read-only: stage the soname link in a writable
# dir and put it first on the search path.
stage = os.path.join(tempfile.gettempdir(), "coderai-whisper-libs")
try:
os.makedirs(stage, exist_ok=True)
link = os.path.join(stage, "libwhisper.so.1")
if not os.path.exists(link):
os.symlink(cands[-1], link)
lib_dirs.insert(0, stage)
except OSError:
pass
if lib_dirs:
prev = env.get("LD_LIBRARY_PATH", "")
env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([prev] if prev else []))
return env
def start(self, model_path: str = None, gpu_device: int = 0):
"""Start whisper-server with the specified model."""
with self.lock:
......@@ -849,11 +918,28 @@ class WhisperServerManager:
print(f"Starting whisper-server: {' '.join(cmd)}")
# The whisper-server binary (a whisper.cpp build) has its RUNPATH baked to
# the build machine's dirs (absent here), so it finds libwhisper.so.1 +
# its ggml 0.9.7 libs only via LD_LIBRARY_PATH. The engine processes don't
# export one, so without this the binary dies at launch with
# "libwhisper.so.1: cannot open shared object file" and we see only the
# opaque "whisper-server failed to start". Its libs live in
# /opt/coderai/local-libs, which ships libwhisper.so.1.8.3 but no
# libwhisper.so.1 soname symlink — create it (self-heal), falling back to
# a writable staging dir if local-libs is read-only.
env = self._whisper_lib_env()
# whisper.cpp --convert writes ffmpeg temp files relative to CWD; the
# engine's own CWD (/opt/coderai/app) is read-only, which makes every
# conversion fail ("FFmpeg conversion failed."). Run from a writable dir.
workdir = self._whisper_workdir()
try:
self.process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=env,
cwd=workdir,
preexec_fn=lambda: signal.signal(signal.SIGTERM, signal.SIG_DFL)
)
self.current_model = actual_model_path
......@@ -862,7 +948,22 @@ class WhisperServerManager:
print(f"whisper-server started on {self.base_url}")
return actual_model_path
else:
print("Error: whisper-server failed to start")
# Surface WHY it failed instead of the opaque message: if the
# process already exited, its stderr/stdout carry the reason
# (missing lib, no GPU/ICD, bad model). Drain non-blocking.
rc = self.process.poll() if self.process else None
detail = ""
if rc is not None and self.process is not None:
try:
out, err = self.process.communicate(timeout=2)
tail = (err or b"").decode("utf-8", "replace").strip() \
or (out or b"").decode("utf-8", "replace").strip()
if tail:
detail = " — " + " | ".join(tail.splitlines()[-5:])
except Exception:
pass
print(f"Error: whisper-server failed to start"
f" (exit={rc}){detail}")
self.stop()
return ""
except Exception as e:
......@@ -1131,6 +1232,11 @@ class MultiModelManager:
self._global_max_instances: int = 1 # set from config at startup
self._measured_vram_gb: Dict[str, float] = {} # actual measured VRAM delta per model key
self._last_load_errors: Dict[str, str] = {} # model_key -> last failed load message
# Model keys marked "keep resident": the small STT/diarization set that
# should co-reside and be skipped by routine LRU eviction, so they don't
# thrash each other. They're still evicted as a LAST RESORT when a load
# (e.g. a big LLM) cannot otherwise fit — so they never block it.
self._keep_resident: set = set()
# Callbacks that free VRAM held *outside* the model manager (e.g. the
# LoRA trainer caches its SD/SDXL base model between jobs). Each returns
# the GB it freed (or None). Invoked as a last resort during eviction.
......@@ -1684,9 +1790,18 @@ class MultiModelManager:
self.config[f"audio:{model_name}"] = config or {}
self._remember_registered_type(model_name, "audio")
if isinstance(config, dict) and config.get("backend") == "whisper-server":
_be = (config or {}).get("backend") if isinstance(config, dict) else None
if _be == "whisper-server":
print(f"Registered whisper-server audio model: {model_name}")
return
# vosk = a local model directory; nemo/canary = served by an isolated venv
# worker that pulls its own checkpoint; wav2vec2 = an HF model the
# transformers pipeline pulls lazily on first request. None should be
# eagerly downloaded here (that would block engine startup), unlike a
# faster-whisper size which set_audio_model resolves/caches below.
if _be in ("vosk", "nemo", "canary", "parakeet", "wav2vec2", "wav2vec"):
print(f"Registered {_be} audio model: {model_name}")
return
# Download/cache the model at startup if it's a URL or HF ID
resolved_model = self.load_model(model_name)
......@@ -4029,6 +4144,7 @@ class MultiModelManager:
model_obj = self.models.pop(key, None)
self.models_in_vram.discard(key)
self._last_used.pop(key, None)
self._keep_resident.discard(key)
# Debug-only: ground-truth the object state before cleanup so an
# orphaned/detached backend (VRAM that won't free) is visible.
try:
......@@ -4232,10 +4348,14 @@ class MultiModelManager:
_free_before = self._get_free_vram_gb()
# First pass: evict idle non-active models in LRU order.
# First pass: evict idle non-active models in LRU order. Skip the
# keep-resident set (the small co-resident STT/diarization models) — they
# are only evicted as a last resort below.
for key in self._lru_order():
if key == self.active_in_vram:
continue
if key in self._keep_resident:
continue
if self._get_free_vram_gb() >= needed_gb:
break
if self._is_key_busy(key):
......@@ -4244,8 +4364,10 @@ class MultiModelManager:
print(f"On-request VRAM eviction: unloading '{key}' ({_size_label(key)}) to free VRAM")
_evict_key(key)
# Second pass: evict the active model if still not enough AND it is idle.
# Second pass: evict the active model if still not enough AND it is idle
# (unless it's keep-resident — that's handled by the last-resort pass).
if (self._get_free_vram_gb() < needed_gb and self.active_in_vram
and self.active_in_vram not in self._keep_resident
and self.active_in_vram in self.models):
_active = self.active_in_vram
if self._is_key_busy(_active):
......@@ -4259,6 +4381,25 @@ class MultiModelManager:
else:
print(f" Active model '{_active}' still busy after wait — leaving it loaded")
# Last resort among tracked models: evict the keep-resident set (the small
# co-resident STT/diarization models) so a genuinely-larger load — e.g. a
# 26B LLM — can still fit. They yield rather than block it; that's the
# "never fight the LLM" half of the co-residency policy.
if self._get_free_vram_gb() < needed_gb and self._keep_resident:
for key in list(self._lru_order()):
if key not in self._keep_resident:
continue
if self._get_free_vram_gb() >= needed_gb:
break
if self._is_key_busy(key) and not self._wait_until_idle(key):
print(f" keep-resident '{key}' still busy — leaving it loaded")
continue
print(f"On-request VRAM eviction (last resort): unloading keep-resident "
f"'{key}' ({_size_label(key)}) to fit a larger load")
if key == self.active_in_vram:
self.active_in_vram = None
_evict_key(key)
# Last resort: ask external holders (e.g. the LoRA trainer's cached base
# model) to release their VRAM. These aren't tracked models, so eviction
# above can't see them — this is what stops a cached training base from
......@@ -5000,6 +5141,55 @@ class MultiModelManager:
self.active_in_vram = key
self.models_in_vram.add(key)
def acquire_stt_backend(self, model_key: str, needed_gb: float,
uses_gpu: bool, loader, keep_resident: bool = False):
"""Load (or reuse) a non-whisper STT backend (wav2vec2 / vosk / NeMo) as a
VRAM-tracked, evictable model.
Mirrors :meth:`start_whisper_server` for the generic STT backends: evict
other models to make room for a GPU backend BEFORE loading, then register
the backend in the loaded-model maps so the normal eviction path can free
it later (``_evict_one`` calls the backend's ``cleanup()`` — which drops the
transformers pipeline / stops the isolated NeMo worker, releasing VRAM).
``loader`` is a no-arg callable returning the backend instance."""
lock = getattr(self, "_stt_load_lock", None)
if lock is None:
import threading
lock = self._stt_load_lock = threading.RLock()
with lock:
if keep_resident:
self._keep_resident.add(model_key)
else:
self._keep_resident.discard(model_key)
existing = self.models.get(model_key)
if existing is not None:
self._last_used[model_key] = time.monotonic()
self.current_model_key = model_key
if uses_gpu:
self.active_in_vram = model_key
self.models_in_vram.add(model_key)
return existing
if uses_gpu and needed_gb and needed_gb > 0:
try:
if self._get_free_vram_gb() < needed_gb:
print(f"STT load: need ~{needed_gb:.1f} GB VRAM for "
f"'{model_key}' — evicting to make room")
self._evict_models_for_vram(needed_gb)
except Exception as e:
print(f" STT VRAM pre-eviction warning for '{model_key}': {e}")
backend = loader()
# Register like a loaded model so eviction/accounting can see + free it.
self.models[model_key] = backend
self.model_pools.pop(model_key, None)
self._last_used[model_key] = time.monotonic()
self.current_model_key = model_key
if uses_gpu:
self.models_in_vram.add(model_key)
self.active_in_vram = model_key
if needed_gb and needed_gb > 0:
self._measured_vram_gb.setdefault(model_key, float(needed_gb))
return backend
def acquire_model_instance(self, model_key: str, session_key=None):
"""Acquire an instance, incrementing its ref-count.
......@@ -5119,6 +5309,8 @@ class MultiModelManager:
port=meta.get("port"),
gpu_device=meta.get("gpu_device"),
load_mode=meta.get("load_mode"),
languages=meta.get("languages") or None,
supports_translation=meta.get("supports_translation"),
))
# --- Models from config (the authoritative source) ---
......@@ -5186,6 +5378,20 @@ class MultiModelManager:
except Exception:
pass
# --- Speaker diarization (pyannote) — a built-in endpoint, not a
# model-manager model, so surface its configured model here so /v1/models
# advertises it (served at /v1/audio/diarization). ---
try:
from codai.api import pyannote_worker
_diar = pyannote_worker._default_model()
if _diar:
_add(_diar, "audio", {"capabilities": ["speaker_diarization"],
"backend": "pyannote"})
_add("diarization", "audio", {"capabilities": ["speaker_diarization"],
"backend": "pyannote"})
except Exception:
pass
# --- Fallback: runtime default_model ---
if not models and self.default_model:
model_id = self.default_model
......
codai/pydantic/textrequest.py
View file @ 7e5a925c
......@@ -150,6 +150,11 @@ class ModelInfo(BaseModel):
port: Optional[int] = None
gpu_device: Optional[int] = None
load_mode: Optional[str] = None
# Speech-to-text: languages this model supports (empty/None = accepts any) and
# whether it can translate (source→target_language, e.g. NVIDIA Canary).
languages: Optional[List[str]] = None
supports_translation: Optional[bool] = None
speaker_diarization: Optional[bool] = None
class ModelList(BaseModel):
......
requirements-nemo.txt 0 → 100644
View file @ 7e5a925c
# Isolated venv for NVIDIA NeMo ASR (Canary / Parakeet) speech-to-text.
#
# NeMo pins a transformers/huggingface-hub/numba stack that conflicts with the
# coderai server's transformers 5.x, so it must NOT be installed into the shared
# venv. codai/api/canary_worker.py builds this venv (baked at /opt/coderai/nemo_venv
# by build.sh, or on demand under ~/.coderai/nemo_venv) and runs
# tools/canary_service.py inside it. Do NOT merge these into requirements.txt.
#
# CUDA-only (needs an NVIDIA GPU). Install order matters: torch first, then nemo.
torch
torchaudio
nemo_toolkit[asr]>=2.0.0
cuda-python
# NeMo transcribe I/O + audio decode
soundfile>=0.12.0
librosa>=0.10.0
# HTTP client used by nothing here (service is stdlib http.server); requests is
# pulled by nemo transitively but pin a floor to be safe.
requests>=2.31.0
requirements-pyannote.txt 0 → 100644
View file @ 7e5a925c
# Isolated venv for pyannote.audio speaker diarization.
#
# pyannote.audio pulls torch + pytorch-lightning + speechbrain + asteroid, which
# conflict with the coderai server's transformers 5.x stack, so it must NOT be
# installed into the shared venv. codai/api/pyannote_worker.py builds this venv
# (baked at /opt/coderai/pyannote_venv by build.sh, or on demand under
# ~/.coderai/pyannote_venv) and runs tools/pyannote_service.py inside it. Do NOT
# merge these into requirements.txt.
#
# TORCH IS PINNED: the bleeding-edge torch (2.13/cu130) that `torch` resolves to
# has a device-placement bug with pyannote's SincNet ("Expected all tensors to be
# on the same device" — weights on cuda, input on cpu), which forces CPU. Pin to a
# pyannote-compatible CUDA build so diarization runs on the 3090. cu124 wheels run
# fine on this box's newer driver (backward compatible).
--extra-index-url https://download.pytorch.org/whl/cu124
torch==2.6.0+cu124
torchaudio==2.6.0+cu124
# Pin pyannote.audio too: leaving it open let pip pick a version mismatched with
# the pinned torch AND the modern huggingface_hub (older pyannote calls
# hf_hub_download(use_auth_token=...), which new hf_hub removed → TypeError). 3.3.2
# is stable, GPU-proven, and uses the modern token= API.
pyannote.audio==3.3.2
# pyannote 3.3.2 calls hf_hub_download(use_auth_token=...), which huggingface_hub
# removed in 1.0 → "unexpected keyword argument 'use_auth_token'". Pin below 1.0
# where it still works (deprecated but accepted).
huggingface_hub<1.0
soundfile>=0.12.0
# pyannote 3.3.2 imports matplotlib (segmentation mixins) but doesn't always pull
# it under the pinned resolution.
matplotlib>=3.7
# Speaker-embedding backends served from this venv: ECAPA-TDNN via speechbrain
# (ungated), alongside pyannote's own embedding model.
speechbrain>=1.0.0
requirements.txt
View file @ 7e5a925c
......@@ -52,6 +52,7 @@ setproctitle>=1.1 # Process naming (replaces procname, Python 3.13 compatible)
# Optional: Audio transcription dependencies
faster-whisper>=0.10.0 # For NVIDIA/CUDA whisper transcription
vosk>=0.3.45 # Offline Kaldi-based STT backend (CPU, per-language model dirs)
whispercpp>=0.0.17 # Alternative whisper library (works without PyTorch)
# LiteLLM for standardized API responses
......
tools/canary_service.py 0 → 100644
View file @ 7e5a925c
#!/usr/bin/env python3
"""Standalone NVIDIA NeMo ASR HTTP microservice — run in its OWN venv.
NeMo (Canary / Parakeet) pins a stack that conflicts with coderai's transformers
5.x, so it runs here behind a tiny stdlib HTTP shim and coderai talks to it as a
remote STT backend (``codai.api.stt_backends._RemoteNemoBackend``), managed by
``codai.api.canary_worker``.
Endpoints:
GET /health -> {"ok": true, "model": ..., "multitask": bool}
POST /transcribe -> {"text", "segments":[{start,end,text}], "language"}
body: raw audio bytes;
query: ?language=<src>&target_language=<tgt>
The uploaded audio is transcoded to 16 kHz mono WAV via ffmpeg before inference.
"""
import argparse
import json
import os
import subprocess
import tempfile
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse, parse_qs
def _to_wav_16k_mono(raw: bytes) -> str:
fd, src = tempfile.mkstemp(suffix=".input")
os.write(fd, raw)
os.close(fd)
fd, out = tempfile.mkstemp(suffix=".wav")
os.close(fd)
try:
p = subprocess.run(
["ffmpeg", "-y", "-i", src, "-ar", "16000", "-ac", "1",
"-c:a", "pcm_s16le", "-f", "wav", out],
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
if p.returncode != 0:
raise RuntimeError("ffmpeg conversion failed: " +
p.stderr.decode("utf-8", "replace")[-300:])
return out
finally:
try:
os.unlink(src)
except OSError:
pass
class _Engine:
"""Loads a NeMo ASR model once and transcribes audio files.
Handles both Canary-style multitask models (source/target language +
translation, ``EncDecMultiTaskModel``) and plain ASR models (Parakeet, etc.)."""
def __init__(self, model_name: str, model_path: str = None):
from nemo.collections.asr.models import ASRModel
self.model_name = model_name
if model_path and os.path.isfile(model_path):
self._model = ASRModel.restore_from(model_path)
else:
self._model = ASRModel.from_pretrained(model_name)
try:
import torch
if torch.cuda.is_available():
self._model = self._model.cuda()
except Exception:
pass
self._model.eval()
# Multitask (Canary) models accept source_lang/target_lang.
self.multitask = self._model.__class__.__name__ == "EncDecMultiTaskModel"
def transcribe(self, wav_path: str, language: str = None,
target_language: str = None) -> dict:
kwargs = {"timestamps": True}
if self.multitask:
src = language or "en"
kwargs["source_lang"] = src
kwargs["target_lang"] = target_language or src
kwargs["pnc"] = "yes"
try:
out = self._model.transcribe([wav_path], **kwargs)
except TypeError:
# Older NeMo without timestamps kwarg.
kwargs.pop("timestamps", None)
out = self._model.transcribe([wav_path], **kwargs)
hyp = out[0] if isinstance(out, (list, tuple)) and out else out
text = getattr(hyp, "text", None)
if text is None:
text = hyp if isinstance(hyp, str) else str(hyp)
segments, words = [], []
ts = getattr(hyp, "timestamp", None)
if isinstance(ts, dict):
for seg in (ts.get("segment") or []):
segments.append({
"start": float(seg.get("start", 0.0)),
"end": float(seg.get("end", 0.0)),
"text": (seg.get("segment") or seg.get("text") or "").strip(),
})
for w in (ts.get("word") or []):
words.append({
"word": (w.get("word") or w.get("text") or "").strip(),
"start": float(w.get("start", 0.0)),
"end": float(w.get("end", 0.0)),
})
return {"text": (text or "").strip(), "segments": segments, "words": words,
"language": target_language or language}
ENGINE = None # set in main()
class Handler(BaseHTTPRequestHandler):
def _send(self, code, body=b"", ctype="application/json"):
self.send_response(code)
self.send_header("Content-Type", ctype)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
if body:
self.wfile.write(body)
def log_message(self, fmt, *args):
pass
def do_GET(self):
if urlparse(self.path).path == "/health":
self._send(200, json.dumps(
{"ok": True, "model": ENGINE.model_name,
"multitask": ENGINE.multitask}).encode())
else:
self._send(404, b'{"error":"not found"}')
def do_POST(self):
parsed = urlparse(self.path)
if parsed.path != "/transcribe":
self._send(404, b'{"error":"not found"}')
return
wav = None
try:
n = int(self.headers.get("Content-Length", 0))
raw = self.rfile.read(n) if n else b""
if not raw:
self._send(400, b'{"error":"empty audio"}')
return
q = parse_qs(parsed.query)
language = (q.get("language") or [None])[0]
target = (q.get("target_language") or [None])[0]
wav = _to_wav_16k_mono(raw)
result = ENGINE.transcribe(wav, language=language, target_language=target)
self._send(200, json.dumps(result).encode())
except Exception as e:
import traceback
traceback.print_exc()
self._send(500, json.dumps({"error": str(e)}).encode())
finally:
if wav:
try:
os.unlink(wav)
except OSError:
pass
def main(argv=None):
ap = argparse.ArgumentParser(description="Standalone NVIDIA NeMo ASR HTTP service")
ap.add_argument("--model", default="nvidia/canary-1b-flash")
ap.add_argument("--model-path", default=None,
help="Local .nemo checkpoint (overrides --model)")
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--port", type=int, default=8124)
args = ap.parse_args(argv)
global ENGINE
print(f"Loading NeMo model {args.model_path or args.model} …")
ENGINE = _Engine(args.model, args.model_path)
print(f"Ready: {args.model} (multitask={ENGINE.multitask}) — "
f"serving on http://{args.host}:{args.port}")
ThreadingHTTPServer((args.host, args.port), Handler).serve_forever()
if __name__ == "__main__":
main()
tools/pyannote_service.py 0 → 100644
View file @ 7e5a925c
#!/usr/bin/env python3
"""Standalone pyannote.audio speaker-diarization HTTP microservice — own venv.
pyannote.audio pins torch/lightning/speechbrain that conflict with coderai's
transformers 5.x, so it runs here behind a stdlib HTTP shim, managed by
``codai.api.pyannote_worker``.
Endpoints:
GET /health -> {"ok": true, "model": ...}
POST /diarize -> {"segments":[{start,end,speaker}], "num_speakers": N}
body: raw audio bytes;
query: ?num_speakers=<int>&min_speakers=<int>&max_speakers=<int>
Audio is transcoded to 16 kHz mono WAV via ffmpeg before inference. The gated
model needs HF_TOKEN in the environment; set CODERAI_DIARIZATION_MODEL to an
ungated mirror to avoid the token.
"""
import argparse
import json
import os
import subprocess
import tempfile
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse, parse_qs
def _to_wav_16k_mono(raw: bytes) -> str:
fd, src = tempfile.mkstemp(suffix=".input")
os.write(fd, raw)
os.close(fd)
fd, out = tempfile.mkstemp(suffix=".wav")
os.close(fd)
try:
p = subprocess.run(
["ffmpeg", "-y", "-i", src, "-ar", "16000", "-ac", "1",
"-c:a", "pcm_s16le", "-f", "wav", out],
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
if p.returncode != 0:
raise RuntimeError("ffmpeg conversion failed: " +
p.stderr.decode("utf-8", "replace")[-300:])
return out
finally:
try:
os.unlink(src)
except OSError:
pass
def _crop_wav(src: str, start: float, end: float) -> str:
"""Write a temp WAV of src[start:end] seconds (for per-speaker embedding)."""
import soundfile as sf
data, sr = sf.read(src)
a = max(0, int(start * sr))
b = min(len(data), int(end * sr))
fd, out = tempfile.mkstemp(suffix=".wav")
os.close(fd)
sf.write(out, data[a:b] if b > a else data, sr)
return out
class _Engine:
def __init__(self, model_name: str):
import torch
# torch 2.6 flipped torch.load's `weights_only` default to True, which
# rejects pyannote's own checkpoints (they pickle a TorchVersion global) →
# UnpicklingError. We load our own trusted pyannote weights in an isolated
# venv, so force weights_only=False for all loads before importing pyannote.
_orig_load = torch.load
def _load(*a, **k):
k["weights_only"] = False # force (pyannote passes True explicitly)
return _orig_load(*a, **k)
torch.load = _load
from pyannote.audio import Pipeline
self.model_name = model_name
token = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
or os.environ.get("HUGGING_FACE_HUB_TOKEN"))
# pyannote.audio 3.1+ renamed the auth kwarg to `token`; older releases use
# `use_auth_token`. Try the new name, fall back to the old.
try:
self._pipe = Pipeline.from_pretrained(model_name, token=token)
except TypeError:
self._pipe = Pipeline.from_pretrained(model_name, use_auth_token=token)
if self._pipe is None:
raise RuntimeError(
f"Pipeline.from_pretrained('{model_name}') returned None — the model is "
"gated/unavailable. Accept its conditions on HF and set HF_TOKEN, or set "
"CODERAI_DIARIZATION_MODEL to an ungated mirror.")
# Device: pyannote.audio's Pipeline.to() is unreliable with very new torch
# (2.13/cu130) — it leaves the segmentation input on CPU while weights move
# to CUDA ("Expected all tensors to be on the same device"). So default to
# CPU (always correct) and only attempt CUDA when explicitly opted in via
# CODERAI_DIARIZATION_DEVICE=cuda (use once torch is pinned to a compatible
# version). The GPU→CPU retry in diarize() is the safety net either way.
# Run on GPU by default (the venv pins a pyannote-compatible torch — see
# requirements-pyannote.txt). CODERAI_DIARIZATION_DEVICE=cpu forces CPU;
# the GPU→CPU retry in diarize() is the safety net if a device mismatch
# slips through on some torch build.
self._device = "cpu"
want = (os.environ.get("CODERAI_DIARIZATION_DEVICE") or "cuda").strip().lower()
try:
if want == "cuda" and torch.cuda.is_available():
self._pipe.to(torch.device("cuda"))
self._device = "cuda"
except Exception:
pass
def diarize(self, wav_path: str, num_speakers=None,
min_speakers=None, max_speakers=None, embed_backend=None) -> dict:
kwargs = {}
if num_speakers:
kwargs["num_speakers"] = int(num_speakers)
else:
if min_speakers:
kwargs["min_speakers"] = int(min_speakers)
if max_speakers:
kwargs["max_speakers"] = int(max_speakers)
# Run on GPU; if this pyannote.audio + torch combo doesn't move the
# segmentation input to CUDA ("Expected all tensors to be on the same
# device"), fall back to CPU and retry so diarization still works.
import torch as _torch
try:
ann = self._pipe(wav_path, **kwargs)
except RuntimeError as e:
if "same device" not in str(e) or self._device == "cpu":
raise
print("[pyannote] GPU device mismatch — falling back to CPU (pin torch "
"to a pyannote-compatible version to keep it on the 3090).", flush=True)
self._pipe.to(_torch.device("cpu"))
self._device = "cpu"
ann = self._pipe(wav_path, **kwargs)
segments = []
speakers = set()
for turn, _, speaker in ann.itertracks(yield_label=True):
speakers.add(speaker)
segments.append({"start": float(turn.start), "end": float(turn.end),
"speaker": str(speaker)})
segments.sort(key=lambda s: s["start"])
result = {"segments": segments, "num_speakers": len(speakers)}
# For speaker identification: embed one representative slice (the longest
# turn) per diarized speaker, so the caller can match each against enrolled
# voiceprints. Returned as {SPEAKER_XX: [vector]}.
if embed_backend:
spk_emb = {}
longest = {}
for s in segments:
dur = s["end"] - s["start"]
if dur > longest.get(s["speaker"], (0.0, None))[0]:
longest[s["speaker"]] = (dur, s)
for spk, (_, s) in longest.items():
crop = None
try:
crop = _crop_wav(wav_path, s["start"], s["end"])
spk_emb[spk] = _speaker_embed(crop, backend=embed_backend)["embedding"]
except Exception as e:
print(f"[pyannote] per-speaker embed failed for {spk}: {e}", flush=True)
finally:
if crop:
try:
os.unlink(crop)
except OSError:
pass
result["speaker_embeddings"] = spk_emb
return result
ENGINE = None
# Speaker-embedding backends, loaded lazily and cached by "backend:model".
_EMB_CACHE = {}
def _hf_token():
return (os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
or os.environ.get("HUGGING_FACE_HUB_TOKEN"))
def _speaker_embed(wav_path: str, backend: str, model: str = None) -> dict:
"""Return a fixed-dim speaker embedding for the audio.
Backends: ``ecapa`` (speechbrain ECAPA-TDNN, ungated), ``pyannote`` /
``wespeaker`` (pyannote.audio embedding models — may be gated, use HF_TOKEN)."""
import torch
backend = (backend or "ecapa").lower()
key = f"{backend}:{model or 'default'}"
dev = "cuda" if torch.cuda.is_available() else "cpu"
fn = _EMB_CACHE.get(key)
if fn is None:
if backend in ("ecapa", "speechbrain"):
from speechbrain.inference.speaker import EncoderClassifier
src = model or "speechbrain/spkrec-ecapa-voxceleb"
clf = EncoderClassifier.from_hparams(source=src, run_opts={"device": dev})
def fn(path, _clf=clf):
import torchaudio
sig, _sr = torchaudio.load(path)
out = _clf.encode_batch(sig.to(dev))
return out.squeeze().detach().cpu().float().tolist()
elif backend in ("pyannote", "wespeaker", "resnet"):
from pyannote.audio import Model, Inference
src = model or ("pyannote/wespeaker-voxceleb-resnet34-LM"
if backend != "pyannote" else "pyannote/embedding")
m = Model.from_pretrained(src, token=_hf_token())
if m is None:
raise RuntimeError(
f"embedding model '{src}' is gated/unavailable — accept its "
f"conditions on HF (https://huggingface.co/{src}) with the token's "
"account, or use backend=wespeaker / backend=ecapa (ungated).")
inf = Inference(m, window="whole")
try:
inf.to(torch.device(dev))
except Exception:
pass
def fn(path, _inf=inf):
import numpy as np
v = _inf(path)
return np.asarray(v).squeeze().astype(float).tolist()
else:
raise ValueError(f"unknown speaker-embedding backend: {backend}")
_EMB_CACHE[key] = fn
emb = fn(wav_path)
return {"embedding": emb, "dim": len(emb), "backend": backend,
"model": model or key.split(":", 1)[1]}
class Handler(BaseHTTPRequestHandler):
def _send(self, code, body=b"", ctype="application/json"):
self.send_response(code)
self.send_header("Content-Type", ctype)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
if body:
self.wfile.write(body)
def log_message(self, fmt, *args):
pass
def do_GET(self):
if urlparse(self.path).path == "/health":
self._send(200, json.dumps({"ok": True, "model": ENGINE.model_name}).encode())
else:
self._send(404, b'{"error":"not found"}')
def do_POST(self):
parsed = urlparse(self.path)
if parsed.path not in ("/diarize", "/embed"):
self._send(404, b'{"error":"not found"}')
return
wav = None
try:
n = int(self.headers.get("Content-Length", 0))
raw = self.rfile.read(n) if n else b""
if not raw:
self._send(400, b'{"error":"empty audio"}')
return
q = parse_qs(parsed.query)
def _q(k):
v = q.get(k)
return v[0] if v else None
wav = _to_wav_16k_mono(raw)
if parsed.path == "/embed":
result = _speaker_embed(wav, backend=_q("backend"), model=_q("model"))
self._send(200, json.dumps(result).encode())
return
result = ENGINE.diarize(wav, num_speakers=_q("num_speakers"),
min_speakers=_q("min_speakers"),
max_speakers=_q("max_speakers"),
embed_backend=_q("embed"))
self._send(200, json.dumps(result).encode())
except Exception as e:
import traceback
traceback.print_exc()
self._send(500, json.dumps({"error": str(e)}).encode())
finally:
if wav:
try:
os.unlink(wav)
except OSError:
pass
def main(argv=None):
ap = argparse.ArgumentParser(description="Standalone pyannote diarization HTTP service")
ap.add_argument("--model", default=os.environ.get("CODERAI_DIARIZATION_MODEL")
or "pyannote/speaker-diarization-3.1")
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--port", type=int, default=8125)
args = ap.parse_args(argv)
global ENGINE
print(f"Loading pyannote pipeline {args.model} …")
ENGINE = _Engine(args.model)
print(f"Ready: {args.model} — serving on http://{args.host}:{args.port}")
ThreadingHTTPServer((args.host, args.port), Handler).serve_forever()
if __name__ == "__main__":
main()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment