5e2094f2ad
FastAPI service in infra/media-worker implementing every MediaInferenceProvider operation behind HTTP API v1, plus capabilities. Bearer auth, MEDIA_ROOT realpath enforcement, strict request schemas, cancellation on disconnect, GPU/CPU gates with 503 overflow, input size/duration limits, lazy SHA256-pinned model downloads, fail-closed redaction with atomic output, offline pytest suite with fake backends, CPU/CUDA/ROCm Dockerfile and a license review of every model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj
507 lines
18 KiB
Python
507 lines
18 KiB
Python
"""Model registry: which backend serves each operation, its license, and its files.
|
|
|
|
Weights are never part of the image or the repo. They download lazily on first
|
|
use into MEDIA_MODELS_DIR/<backend id>/ and are verified against the pinned
|
|
SHA256 and size before anything loads them. A mismatch deletes the file and
|
|
the operation reports `unsupported` with the reason.
|
|
|
|
Default backends are commercially usable. Anything else must be listed in
|
|
MEDIA_OPTIONAL_BACKENDS by the operator before it runs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib.util
|
|
import logging
|
|
import os
|
|
import threading
|
|
import urllib.request
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from .errors import Unsupported
|
|
|
|
log = logging.getLogger("media-worker.models")
|
|
|
|
OPERATIONS = (
|
|
"detectScenes",
|
|
"transcribe",
|
|
"detectFaces",
|
|
"trackFaces",
|
|
"faceLandmarks",
|
|
"embedFrames",
|
|
"classifyContent",
|
|
"segmentPrompted",
|
|
"inpaintVideo",
|
|
"redact",
|
|
"synthesizeIdentity",
|
|
"audioFeatures",
|
|
"saliency",
|
|
)
|
|
|
|
HF = "https://huggingface.co"
|
|
YUNET_REV = "3cc26e7f1014a5ee5d74a42acee58bafc9d0a310"
|
|
SFACE_REV = "3d7082438a6e4551e840c9b2bb60b71e8da4b524"
|
|
WHISPER_REV = "536b0662742c02347bc0e980a01041f333bce120"
|
|
SIGLIP_REV = "4649052661e53c7000355844105f8a1792088239"
|
|
LAMA_REV = "c3c0c9e468934d62e79c329e35d82dd09ff8c444"
|
|
SAM2_REV = "de431c4043854a71d8101e17995dfe596bf101a5"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModelFile:
|
|
name: str
|
|
url: str
|
|
sha256: str
|
|
sizeBytes: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BackendSpec:
|
|
"""One way to serve an operation. `id` doubles as the model cache directory."""
|
|
|
|
id: str
|
|
model: str | None
|
|
license: str | None
|
|
commercialUse: bool | None
|
|
experimental: bool
|
|
source: str | None
|
|
files: tuple[ModelFile, ...] = ()
|
|
# Python modules that must be importable.
|
|
requires: tuple[str, ...] = ()
|
|
# Set when the backend can never run (e.g. no licensed model exists).
|
|
unavailable_reason: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class Backend:
|
|
spec: BackendSpec
|
|
# Builds the runtime object from {file name: verified local path} and the settings.
|
|
load: Callable[[dict[str, str], Any], Any] | None = None
|
|
|
|
|
|
def _yunet_file() -> ModelFile:
|
|
return ModelFile(
|
|
"face_detection_yunet_2023mar.onnx",
|
|
f"{HF}/opencv/face_detection_yunet/resolve/{YUNET_REV}/face_detection_yunet_2023mar.onnx",
|
|
"8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4",
|
|
232589,
|
|
)
|
|
|
|
|
|
def _siglip(name: str, sha256: str, size: int) -> ModelFile:
|
|
path = name if name == "tokenizer.json" else f"onnx/{name}"
|
|
return ModelFile(name, f"{HF}/Xenova/siglip-base-patch16-224/resolve/{SIGLIP_REV}/{path}", sha256, size)
|
|
|
|
|
|
def _whisper(name: str, sha256: str, size: int) -> ModelFile:
|
|
return ModelFile(name, f"{HF}/Systran/faster-whisper-small/resolve/{WHISPER_REV}/{name}", sha256, size)
|
|
|
|
|
|
def _lazy(module: str, attr: str) -> Callable[[dict[str, str], Any], Any]:
|
|
def load(files: dict[str, str], settings: Any) -> Any:
|
|
mod = importlib.import_module(f"app.ops.{module}")
|
|
return getattr(mod, attr)(files, settings)
|
|
|
|
return load
|
|
|
|
|
|
YUNET = BackendSpec(
|
|
id="yunet-2023mar",
|
|
model="opencv-yunet-2023mar",
|
|
license="MIT",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source="https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet",
|
|
files=(_yunet_file(),),
|
|
requires=("cv2",),
|
|
)
|
|
|
|
SIGLIP_SOURCE = "https://huggingface.co/google/siglip-base-patch16-224"
|
|
SIGLIP_VISION = _siglip(
|
|
"vision_model_quantized.onnx",
|
|
"ef14a954f3d57e1806666432bd9785004c1dc27100aa260eee0cb0f10a5de058",
|
|
99499129,
|
|
)
|
|
SIGLIP_FULL = _siglip(
|
|
"model_quantized.onnx",
|
|
"9171eb00c38b9ec82f924877356d008b79e3285dbac7cd10965827bee30c9a99",
|
|
210977441,
|
|
)
|
|
SIGLIP_TOKENIZER = _siglip(
|
|
"tokenizer.json",
|
|
"4a17c975210be5ab4c36b47d8dae4eefb866dbfb1e676e394aad85dc30a3ae08",
|
|
2398744,
|
|
)
|
|
|
|
|
|
def default_backends() -> dict[str, Backend]:
|
|
return {
|
|
"detectScenes": Backend(
|
|
BackendSpec(
|
|
id="pyscenedetect-adaptive",
|
|
model="pyscenedetect-adaptive",
|
|
license="BSD-3-Clause",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source="https://github.com/Breakthrough/PySceneDetect",
|
|
requires=("scenedetect", "cv2"),
|
|
),
|
|
_lazy("scenes", "load_scene_detector"),
|
|
),
|
|
"transcribe": Backend(
|
|
BackendSpec(
|
|
id="faster-whisper-small",
|
|
model="faster-whisper-small",
|
|
license="MIT",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source="https://huggingface.co/Systran/faster-whisper-small",
|
|
files=(
|
|
_whisper(
|
|
"config.json",
|
|
"b55496ac7940a7ae47d2c01eab40edfd8701feec1229d9cce3b40014383fb828",
|
|
2370,
|
|
),
|
|
_whisper(
|
|
"model.bin",
|
|
"3e305921506d8872816023e4c273e75d2419fb89b24da97b4fe7bce14170d671",
|
|
483546902,
|
|
),
|
|
_whisper(
|
|
"tokenizer.json",
|
|
"fb7b63191e9bb045082c79fd742a3106a12c99513ab30df4a0d47fa6cb6fd0ab",
|
|
2203239,
|
|
),
|
|
_whisper(
|
|
"vocabulary.txt",
|
|
"34ce3fe1c5041027b3f8d42912270993f986dbc4bb34cf27f951e34a1e453913",
|
|
459861,
|
|
),
|
|
),
|
|
requires=("faster_whisper",),
|
|
),
|
|
_lazy("transcribe", "load_transcriber"),
|
|
),
|
|
"detectFaces": Backend(YUNET, _lazy("faces", "load_face_detector")),
|
|
"trackFaces": Backend(
|
|
BackendSpec(
|
|
id=YUNET.id,
|
|
model="opencv-yunet-2023mar+iou-tracker",
|
|
license="MIT",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source=YUNET.source,
|
|
files=YUNET.files,
|
|
requires=YUNET.requires,
|
|
),
|
|
_lazy("faces", "load_face_detector"),
|
|
),
|
|
"faceLandmarks": Backend(
|
|
BackendSpec(
|
|
id=YUNET.id,
|
|
model="opencv-yunet-2023mar-5pt",
|
|
license="MIT",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source=YUNET.source,
|
|
files=YUNET.files,
|
|
requires=YUNET.requires,
|
|
),
|
|
_lazy("faces", "load_face_detector"),
|
|
),
|
|
"embedFrames": Backend(
|
|
BackendSpec(
|
|
id="siglip-base-patch16-224",
|
|
model="siglip-base-patch16-224",
|
|
license="Apache-2.0",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source=SIGLIP_SOURCE,
|
|
files=(SIGLIP_VISION,),
|
|
requires=("onnxruntime",),
|
|
),
|
|
_lazy("embeddings", "load_frame_embedder"),
|
|
),
|
|
"classifyContent": Backend(
|
|
BackendSpec(
|
|
id="siglip-base-patch16-224",
|
|
model="siglip-base-patch16-224-zero-shot",
|
|
license="Apache-2.0",
|
|
commercialUse=True,
|
|
# Zero-shot prompts, not a trained moderation head.
|
|
experimental=True,
|
|
source=SIGLIP_SOURCE,
|
|
files=(SIGLIP_FULL, SIGLIP_TOKENIZER),
|
|
requires=("onnxruntime", "tokenizers"),
|
|
),
|
|
_lazy("embeddings", "load_content_classifier"),
|
|
),
|
|
"segmentPrompted": Backend(
|
|
BackendSpec(
|
|
id="sam2.1-hiera-tiny",
|
|
model="sam2.1-hiera-tiny",
|
|
license="Apache-2.0",
|
|
commercialUse=True,
|
|
experimental=True,
|
|
source="https://github.com/facebookresearch/sam2",
|
|
files=(
|
|
ModelFile(
|
|
"sam2.1_hiera_tiny.pt",
|
|
f"{HF}/facebook/sam2.1-hiera-tiny/resolve/{SAM2_REV}/sam2.1_hiera_tiny.pt",
|
|
"7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69",
|
|
156008466,
|
|
),
|
|
),
|
|
requires=("torch", "sam2"),
|
|
),
|
|
_lazy("segment", "load_segmenter"),
|
|
),
|
|
"inpaintVideo": Backend(
|
|
BackendSpec(
|
|
id="lama-big-onnx",
|
|
model="big-lama+dis-flow",
|
|
license="Apache-2.0",
|
|
commercialUse=True,
|
|
experimental=True,
|
|
source="https://huggingface.co/Carve/LaMa-ONNX",
|
|
files=(
|
|
ModelFile(
|
|
"lama_fp32.onnx",
|
|
f"{HF}/Carve/LaMa-ONNX/resolve/{LAMA_REV}/lama_fp32.onnx",
|
|
"1faef5301d78db7dda502fe59966957ec4b79dd64e16f03ed96913c7a4eb68d6",
|
|
208044816,
|
|
),
|
|
),
|
|
requires=("onnxruntime", "cv2"),
|
|
),
|
|
_lazy("inpaint", "load_inpainter"),
|
|
),
|
|
"redact": Backend(
|
|
BackendSpec(
|
|
id="opencv-redact",
|
|
model="opencv-redact",
|
|
license="Apache-2.0",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source="https://github.com/opencv/opencv",
|
|
requires=("cv2",),
|
|
),
|
|
),
|
|
"synthesizeIdentity": Backend(
|
|
BackendSpec(
|
|
id="synthetic-identity",
|
|
model=None,
|
|
license=None,
|
|
commercialUse=None,
|
|
experimental=True,
|
|
source=None,
|
|
unavailable_reason="no commercially licensed model",
|
|
),
|
|
),
|
|
"audioFeatures": Backend(
|
|
BackendSpec(
|
|
id="ffmpeg-ebur128-librosa",
|
|
model="ffmpeg-ebur128+librosa-beat",
|
|
license="ISC",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source="https://github.com/librosa/librosa",
|
|
requires=("librosa",),
|
|
),
|
|
),
|
|
"saliency": Backend(
|
|
BackendSpec(
|
|
id="spectral-residual",
|
|
model="spectral-residual",
|
|
license="Apache-2.0",
|
|
commercialUse=True,
|
|
experimental=False,
|
|
source="https://doi.org/10.1109/CVPR.2007.383267",
|
|
requires=("cv2",),
|
|
),
|
|
),
|
|
}
|
|
|
|
|
|
# Opt-in backends that are not clean enough to be defaults.
|
|
SFACE = Backend(
|
|
BackendSpec(
|
|
id="sface-2021dec",
|
|
model="opencv-sface-2021dec",
|
|
license="Apache-2.0 (weights); training data provenance unverified",
|
|
commercialUse=None,
|
|
experimental=True,
|
|
source="https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface",
|
|
files=(
|
|
ModelFile(
|
|
"face_recognition_sface_2021dec.onnx",
|
|
f"{HF}/opencv/face_recognition_sface/resolve/{SFACE_REV}/face_recognition_sface_2021dec.onnx",
|
|
"0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79",
|
|
38696353,
|
|
),
|
|
),
|
|
requires=("cv2",),
|
|
),
|
|
_lazy("faces", "load_identity_embedder"),
|
|
)
|
|
OPTIONAL_BACKENDS = {"faceIdentity": SFACE}
|
|
|
|
|
|
def _sha256(path: str) -> str:
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1 << 20), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _module_available(name: str) -> bool:
|
|
try:
|
|
return importlib.util.find_spec(name) is not None
|
|
except (ImportError, ValueError):
|
|
return False
|
|
|
|
|
|
@dataclass
|
|
class Registry:
|
|
settings: Any
|
|
backends: dict[str, Backend] = field(default_factory=default_backends)
|
|
optional: dict[str, Backend] = field(default_factory=lambda: dict(OPTIONAL_BACKENDS))
|
|
_runtimes: dict[str, Any] = field(default_factory=dict)
|
|
_verified: set[str] = field(default_factory=set)
|
|
_failed: dict[str, str] = field(default_factory=dict)
|
|
_lock: threading.Lock = field(default_factory=threading.Lock)
|
|
_file_locks: dict[str, threading.Lock] = field(default_factory=dict)
|
|
|
|
# -- availability -------------------------------------------------------
|
|
|
|
def _license_block(self, spec: BackendSpec, key: str) -> str | None:
|
|
if spec.commercialUse is True:
|
|
return None
|
|
if key in self.settings.optional_backends or spec.id in self.settings.optional_backends:
|
|
return None
|
|
return "backend is not licensed for commercial use and not enabled in MEDIA_OPTIONAL_BACKENDS"
|
|
|
|
def unavailable_reason(self, backend: Backend, key: str) -> str | None:
|
|
spec = backend.spec
|
|
if spec.unavailable_reason:
|
|
return spec.unavailable_reason
|
|
blocked = self._license_block(spec, key)
|
|
if blocked:
|
|
return blocked
|
|
missing = [name for name in spec.requires if not _module_available(name)]
|
|
if missing:
|
|
return f"python module not installed: {', '.join(missing)}"
|
|
if spec.id in self._failed:
|
|
return self._failed[spec.id]
|
|
if not self.settings.allow_model_download:
|
|
absent = [f.name for f in spec.files if not os.path.isfile(self._file_path(spec, f))]
|
|
if absent:
|
|
return "model files not installed and MEDIA_MODEL_DOWNLOAD=0"
|
|
return None
|
|
|
|
def backend(self, operation: str) -> Backend:
|
|
return self.backends[operation]
|
|
|
|
def optional_enabled(self, key: str) -> bool:
|
|
backend = self.optional.get(key)
|
|
return backend is not None and self.unavailable_reason(backend, key) is None
|
|
|
|
# -- runtime -------------------------------------------------------------
|
|
|
|
def runtime(self, operation: str) -> Any:
|
|
"""Return the loaded runtime for an operation, downloading weights on first use."""
|
|
return self._runtime(self.backends[operation], operation)
|
|
|
|
def optional_runtime(self, key: str) -> Any:
|
|
return self._runtime(self.optional[key], key)
|
|
|
|
def _runtime(self, backend: Backend, key: str) -> Any:
|
|
reason = self.unavailable_reason(backend, key)
|
|
if reason:
|
|
raise Unsupported(reason)
|
|
cache_key = f"{backend.spec.id}:{key}"
|
|
with self._lock:
|
|
if cache_key in self._runtimes:
|
|
return self._runtimes[cache_key]
|
|
files = {f.name: self.ensure_file(backend.spec, f) for f in backend.spec.files}
|
|
runtime = backend.load(files, self.settings) if backend.load else None
|
|
with self._lock:
|
|
self._runtimes.setdefault(cache_key, runtime)
|
|
return self._runtimes[cache_key]
|
|
|
|
# -- files ---------------------------------------------------------------
|
|
|
|
def _file_path(self, spec: BackendSpec, model_file: ModelFile) -> str:
|
|
return os.path.join(self.settings.models_dir, spec.id, model_file.name)
|
|
|
|
def _file_lock(self, path: str) -> threading.Lock:
|
|
with self._lock:
|
|
return self._file_locks.setdefault(path, threading.Lock())
|
|
|
|
def ensure_file(self, spec: BackendSpec, model_file: ModelFile) -> str:
|
|
path = self._file_path(spec, model_file)
|
|
with self._file_lock(path):
|
|
if path in self._verified:
|
|
return path
|
|
if not os.path.isfile(path):
|
|
if not self.settings.allow_model_download:
|
|
raise Unsupported(f"model file {model_file.name} is not installed")
|
|
self._download(model_file, path)
|
|
ok = os.path.getsize(path) == model_file.sizeBytes and _sha256(path) == model_file.sha256
|
|
if not ok:
|
|
os.unlink(path)
|
|
reason = f"model file {model_file.name} failed SHA256 verification"
|
|
self._failed[spec.id] = reason
|
|
raise Unsupported(reason)
|
|
self._verified.add(path)
|
|
return path
|
|
|
|
def _download(self, model_file: ModelFile, path: str) -> None:
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
partial = f"{path}.part"
|
|
log.info("downloading %s", model_file.url)
|
|
try:
|
|
with (
|
|
urllib.request.urlopen(model_file.url, timeout=60) as response,
|
|
open(partial, "wb") as handle,
|
|
):
|
|
written = 0
|
|
for chunk in iter(lambda: response.read(1 << 20), b""):
|
|
written += len(chunk)
|
|
if written > model_file.sizeBytes:
|
|
raise ValueError("download larger than pinned size")
|
|
handle.write(chunk)
|
|
os.replace(partial, path)
|
|
except Exception as error:
|
|
if os.path.exists(partial):
|
|
os.unlink(partial)
|
|
raise Unsupported(f"model file {model_file.name} could not be downloaded: {error}") from error
|
|
|
|
# -- capabilities ----------------------------------------------------------
|
|
|
|
def operation_capabilities(self) -> dict[str, dict[str, Any]]:
|
|
result: dict[str, dict[str, Any]] = {}
|
|
for operation in OPERATIONS:
|
|
backend = self.backends[operation]
|
|
spec = backend.spec
|
|
result[operation] = {
|
|
"available": self.unavailable_reason(backend, operation) is None,
|
|
"experimental": spec.experimental,
|
|
"model": spec.model,
|
|
"license": spec.license,
|
|
"commercialUse": spec.commercialUse,
|
|
}
|
|
return result
|
|
|
|
def prefetch(self) -> None:
|
|
"""Download and verify every default backend's files (operator warm-up)."""
|
|
for operation in OPERATIONS:
|
|
backend = self.backends[operation]
|
|
if self.unavailable_reason(backend, operation) is None:
|
|
for model_file in backend.spec.files:
|
|
self.ensure_file(backend.spec, model_file)
|