Spaces:
Running
Running
| from __future__ import annotations | |
| import asyncio | |
| import hmac | |
| import json | |
| import math | |
| import os | |
| import tempfile | |
| from pathlib import Path | |
| from typing import Any | |
| from fastapi import APIRouter, HTTPException, Request | |
| router = APIRouter(prefix="/v1/diarization", tags=["diarization"]) | |
| _pipeline: Any = None | |
| _pipeline_lock = asyncio.Lock() | |
| _inference_lock = asyncio.Lock() | |
| def _max_input_bytes() -> int: | |
| try: | |
| value = int(os.getenv("DIARIZATION_MAX_INPUT_BYTES", "1073741824")) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=500, detail="Invalid diarization size configuration") from exc | |
| if value < 1_048_576: | |
| raise HTTPException(status_code=500, detail="Diarization size limit is too small") | |
| return value | |
| def _max_response_spans() -> int: | |
| try: | |
| value = int(os.getenv("DIARIZATION_MAX_RESPONSE_SPANS", "10000")) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=500, detail="Invalid diarization response limit") from exc | |
| if value < 1: | |
| raise HTTPException(status_code=500, detail="Diarization response limit is invalid") | |
| return value | |
| def _max_speakers() -> int: | |
| try: | |
| value = int(os.getenv("DIARIZATION_MAX_SPEAKERS", "32")) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=500, detail="Invalid diarization speaker limit") from exc | |
| if not 1 <= value <= 256: | |
| raise HTTPException(status_code=500, detail="Diarization speaker limit is invalid") | |
| return value | |
| def _max_duration_seconds() -> float: | |
| try: | |
| value = float(os.getenv("DIARIZATION_MAX_DURATION_SECONDS", "14400")) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=500, detail="Invalid diarization duration limit") from exc | |
| if not 1 <= value <= 86400: | |
| raise HTTPException(status_code=500, detail="Diarization duration limit is invalid") | |
| return value | |
| def _inference_timeout_seconds() -> float: | |
| try: | |
| value = float(os.getenv("DIARIZATION_INFERENCE_TIMEOUT_SECONDS", "900")) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=500, detail="Invalid diarization inference timeout") from exc | |
| if value < 1: | |
| raise HTTPException(status_code=500, detail="Diarization inference timeout is invalid") | |
| return value | |
| def _configured_api_key() -> str: | |
| return os.getenv("DIARIZATION_API_KEY", "").strip() | |
| def _authorize(request: Request) -> None: | |
| expected = _configured_api_key() | |
| if not expected: | |
| raise HTTPException(status_code=503, detail="Diarization API authentication is not configured") | |
| authorization = request.headers.get("authorization", "") | |
| scheme, _, presented = authorization.partition(" ") | |
| if scheme.lower() != "bearer" or not presented or not hmac.compare_digest( | |
| presented.strip(), expected | |
| ): | |
| raise HTTPException(status_code=401, detail="Invalid diarization API credentials") | |
| async def _get_pipeline() -> Any: | |
| global _pipeline | |
| if _pipeline is not None: | |
| return _pipeline | |
| async with _pipeline_lock: | |
| if _pipeline is not None: | |
| return _pipeline | |
| token = os.getenv("HF_TOKEN", "").strip() | |
| model_id = os.getenv("DIARIZATION_MODEL_ID", "pyannote/speaker-diarization-3.1").strip() | |
| if not token: | |
| raise HTTPException(status_code=503, detail="Diarization model credentials are not configured") | |
| try: | |
| from pyannote.audio import Pipeline | |
| except ImportError as exc: | |
| raise HTTPException(status_code=503, detail="Diarization backend is not installed") from exc | |
| try: | |
| pipeline = Pipeline.from_pretrained(model_id, token=token) | |
| try: | |
| import torch | |
| pipeline.to(torch.device("cpu")) | |
| except Exception: | |
| pass | |
| _pipeline = pipeline | |
| return pipeline | |
| except Exception as exc: | |
| raise HTTPException(status_code=503, detail="Diarization model could not be loaded") from exc | |
| async def _save_body(request: Request, destination: Path) -> int: | |
| limit = _max_input_bytes() | |
| content_length = request.headers.get("content-length") | |
| if content_length: | |
| try: | |
| declared = int(content_length) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=400, detail="Invalid Content-Length") from exc | |
| if declared < 0: | |
| raise HTTPException(status_code=400, detail="Invalid Content-Length") | |
| if declared > limit: | |
| raise HTTPException(status_code=413, detail="Diarization input exceeds the configured size limit") | |
| total = 0 | |
| try: | |
| with destination.open("wb") as handle: | |
| async for chunk in request.stream(): | |
| if not chunk: | |
| continue | |
| total += len(chunk) | |
| if total > limit: | |
| raise HTTPException(status_code=413, detail="Diarization input exceeds the configured size limit") | |
| handle.write(chunk) | |
| except HTTPException: | |
| raise | |
| except OSError as exc: | |
| raise HTTPException(status_code=500, detail="Unable to store diarization input") from exc | |
| if total == 0: | |
| raise HTTPException(status_code=400, detail="Diarization input is empty") | |
| return total | |
| async def _probe_duration(source: Path) -> float: | |
| binary = os.getenv("FFPROBE_BINARY", "ffprobe") | |
| command = [binary, "-v", "error", "-show_entries", "format=duration", "-of", "json", str(source)] | |
| try: | |
| process = await asyncio.create_subprocess_exec( | |
| *command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL | |
| ) | |
| except OSError as exc: | |
| raise HTTPException(status_code=503, detail="Diarization media probe is unavailable") from exc | |
| try: | |
| stdout, _ = await asyncio.wait_for(process.communicate(), timeout=30) | |
| except asyncio.TimeoutError as exc: | |
| process.kill() | |
| await process.communicate() | |
| raise HTTPException(status_code=422, detail="Diarization media probe timed out") from exc | |
| if process.returncode != 0: | |
| raise HTTPException(status_code=422, detail="Unable to inspect diarization media") | |
| try: | |
| payload = json.loads(stdout) | |
| duration = float(payload["format"]["duration"]) | |
| except (json.JSONDecodeError, KeyError, TypeError, ValueError) as exc: | |
| raise HTTPException(status_code=422, detail="Diarization media has no valid duration") from exc | |
| if not math.isfinite(duration) or duration <= 0: | |
| raise HTTPException(status_code=422, detail="Diarization media has invalid duration") | |
| if duration > _max_duration_seconds(): | |
| raise HTTPException(status_code=413, detail="Diarization media exceeds the configured duration limit") | |
| return duration | |
| async def _extract_audio(source: Path, destination: Path) -> None: | |
| command = [ | |
| os.getenv("FFMPEG_BINARY", "ffmpeg"), "-hide_banner", "-loglevel", "error", "-nostdin", | |
| "-i", str(source), "-vn", "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", "-f", "wav", | |
| str(destination), | |
| ] | |
| try: | |
| process = await asyncio.create_subprocess_exec( | |
| *command, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE | |
| ) | |
| except OSError as exc: | |
| raise HTTPException(status_code=422, detail="Unable to decode diarization media") from exc | |
| try: | |
| _, _ = await asyncio.wait_for(process.communicate(), timeout=600) | |
| except asyncio.TimeoutError as exc: | |
| process.kill() | |
| await process.communicate() | |
| raise HTTPException(status_code=422, detail="Diarization media decoding timed out") from exc | |
| if process.returncode != 0 or not destination.is_file() or destination.stat().st_size == 0: | |
| raise HTTPException(status_code=422, detail="Unable to decode diarization media") | |
| def _normalize_speaker(value: Any) -> str: | |
| if value is None: | |
| raise HTTPException(status_code=422, detail="Diarization returned an invalid speaker label") | |
| speaker = str(value).strip() | |
| if not speaker or len(speaker) > 128: | |
| raise HTTPException(status_code=422, detail="Diarization returned an invalid speaker label") | |
| if any(ord(char) < 32 or ord(char) == 127 for char in speaker): | |
| raise HTTPException(status_code=422, detail="Diarization returned an invalid speaker label") | |
| return speaker | |
| def _normalize_spans( | |
| raw_spans: list[dict[str, Any]], | |
| duration: float, | |
| max_spans: int, | |
| max_speakers: int, | |
| ) -> list[dict[str, Any]]: | |
| if not math.isfinite(duration) or duration <= 0: | |
| raise HTTPException(status_code=422, detail="Diarization media duration is invalid") | |
| normalized: list[dict[str, Any]] = [] | |
| speakers: set[str] = set() | |
| for raw in raw_spans: | |
| if len(normalized) >= max_spans: | |
| raise HTTPException(status_code=422, detail="Diarization produced too many spans") | |
| try: | |
| start = float(raw["start"]) | |
| end = float(raw["end"]) | |
| except (KeyError, TypeError, ValueError) as exc: | |
| raise HTTPException(status_code=422, detail="Diarization returned an invalid span") from exc | |
| if not math.isfinite(start) or not math.isfinite(end): | |
| raise HTTPException(status_code=422, detail="Diarization returned non-finite timestamps") | |
| if start < 0 or end <= start: | |
| raise HTTPException(status_code=422, detail="Diarization returned an invalid timestamp range") | |
| if end > duration + 0.1: | |
| raise HTTPException(status_code=422, detail="Diarization span exceeds media duration") | |
| end = min(end, duration) | |
| speaker = _normalize_speaker(raw.get("speaker")) | |
| speakers.add(speaker) | |
| if len(speakers) > max_speakers: | |
| raise HTTPException(status_code=422, detail="Diarization returned too many speakers") | |
| normalized.append({"speaker": speaker, "start": start, "end": end}) | |
| normalized.sort(key=lambda item: (item["start"], item["end"], item["speaker"])) | |
| return normalized | |
| def _run_pipeline( | |
| pipeline: Any, | |
| audio_path: Path, | |
| min_speakers: int | None, | |
| max_speakers: int | None, | |
| duration: float, | |
| max_spans: int, | |
| max_speakers_allowed: int, | |
| ) -> list[dict[str, Any]]: | |
| kwargs: dict[str, int] = {} | |
| if min_speakers is not None: | |
| kwargs["min_speakers"] = min_speakers | |
| if max_speakers is not None: | |
| kwargs["max_speakers"] = max_speakers | |
| diarization = pipeline(str(audio_path), **kwargs) | |
| raw_spans: list[dict[str, Any]] = [] | |
| for turn, _, speaker in diarization.itertracks(yield_label=True): | |
| try: | |
| start = float(turn.start) | |
| end = float(turn.end) | |
| except (AttributeError, TypeError, ValueError) as exc: | |
| raise HTTPException(status_code=422, detail="Diarization returned an invalid span") from exc | |
| if len(raw_spans) >= max_spans: | |
| raise HTTPException(status_code=422, detail="Diarization produced too many spans") | |
| raw_spans.append({"speaker": speaker, "start": start, "end": end}) | |
| return _normalize_spans(raw_spans, duration, max_spans, max_speakers_allowed) | |
| async def _run_inference( | |
| pipeline: Any, | |
| audio_path: Path, | |
| min_speakers: int | None, | |
| max_speakers: int | None, | |
| duration: float, | |
| max_spans: int, | |
| max_speakers_allowed: int, | |
| ) -> list[dict[str, Any]]: | |
| timeout = _inference_timeout_seconds() | |
| async with _inference_lock: | |
| try: | |
| return await asyncio.wait_for( | |
| asyncio.to_thread( | |
| _run_pipeline, pipeline, audio_path, min_speakers, max_speakers, | |
| duration, max_spans, max_speakers_allowed, | |
| ), | |
| timeout=timeout, | |
| ) | |
| except asyncio.TimeoutError as exc: | |
| raise HTTPException(status_code=504, detail="Speaker diarization timed out") from exc | |
| except HTTPException: | |
| raise | |
| except Exception as exc: | |
| raise HTTPException(status_code=422, detail="Speaker diarization failed") from exc | |
| async def diarize(request: Request) -> dict[str, Any]: | |
| _authorize(request) | |
| filename = request.headers.get("x-filename", "input.bin")[:255] | |
| try: | |
| min_speakers = int(request.headers["x-min-speakers"]) if request.headers.get("x-min-speakers") else None | |
| max_speakers = int(request.headers["x-max-speakers"]) if request.headers.get("x-max-speakers") else None | |
| except ValueError as exc: | |
| raise HTTPException(status_code=400, detail="Speaker bounds must be integers") from exc | |
| if min_speakers is not None and not 1 <= min_speakers <= 32: | |
| raise HTTPException(status_code=400, detail="X-Min-Speakers must be between 1 and 32") | |
| if max_speakers is not None and not 1 <= max_speakers <= 32: | |
| raise HTTPException(status_code=400, detail="X-Max-Speakers must be between 1 and 32") | |
| if min_speakers is not None and max_speakers is not None and min_speakers > max_speakers: | |
| raise HTTPException(status_code=400, detail="X-Min-Speakers cannot exceed X-Max-Speakers") | |
| max_spans = _max_response_spans() | |
| max_speakers_allowed = _max_speakers() | |
| with tempfile.TemporaryDirectory(prefix="diarization-") as directory: | |
| root = Path(directory) | |
| source = root / "input" | |
| audio = root / "audio.wav" | |
| await _save_body(request, source) | |
| duration = await _probe_duration(source) | |
| await _extract_audio(source, audio) | |
| pipeline = await _get_pipeline() | |
| spans = await _run_inference( | |
| pipeline, audio, min_speakers, max_speakers, duration, max_spans, max_speakers_allowed | |
| ) | |
| return { | |
| "spans": spans, | |
| "speakers": list(dict.fromkeys(item["speaker"] for item in spans)), | |
| "filename": filename, | |
| } | |