From a40582165af3a4328eb42606c0dafb5294dac8df Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 18:15:53 +0000 Subject: [PATCH 1/3] feat(media): wire Apple Photos sync into the desktop app Move the Photos bridge types into the desktop contract, expose rakazoDesktop.media.photos from the preload, connect the sync runner in the main process through an authenticated oRPC client, package and sign the helper, and build and test it on macOS CI. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj --- .github/workflows/desktop-photos-helper.yml | 32 ++++++++++ .github/workflows/release-desktop.yml | 12 ++++ apps/desktop/package.json | 15 ++++- apps/desktop/src/main.ts | 67 ++++++++++++++++++++- apps/desktop/src/media-ipc.ts | 27 +++------ apps/desktop/src/media-photos-helper.ts | 9 +-- apps/desktop/src/media-photos-sync.ts | 27 +++------ apps/desktop/src/preload.cjs | 12 ++++ apps/desktop/src/preload.test.ts | 8 +++ packages/contracts/src/desktop.ts | 44 ++++++++++++++ pnpm-lock.yaml | 6 ++ 11 files changed, 214 insertions(+), 45 deletions(-) create mode 100644 .github/workflows/desktop-photos-helper.yml diff --git a/.github/workflows/desktop-photos-helper.yml b/.github/workflows/desktop-photos-helper.yml new file mode 100644 index 0000000..0ca9077 --- /dev/null +++ b/.github/workflows/desktop-photos-helper.yml @@ -0,0 +1,32 @@ +name: desktop-photos-helper +on: + pull_request: + paths: &photos-helper-paths + - ".github/workflows/desktop-photos-helper.yml" + - "apps/desktop/native/photos-helper/**" + push: + branches: [main] + paths: *photos-helper-paths + +permissions: + contents: read + +concurrency: + group: desktop-photos-helper-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + swift: + name: Photos helper build and XCTest + runs-on: [self-hosted, evoo-ci, macOS] + timeout-minutes: 20 + steps: + - uses: actions/checkout@v5 + with: + persist-credentials: false + - name: Build and test + working-directory: apps/desktop/native/photos-helper + run: | + set -euo pipefail + swift build -c release + swift test diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml index bac3e64..b734d5d 100644 --- a/.github/workflows/release-desktop.yml +++ b/.github/workflows/release-desktop.yml @@ -145,6 +145,14 @@ jobs: umask 077 printf '%s\n' "$APPLE_API_KEY_P8" > "$APPLE_API_KEY" chmod 600 "$APPLE_API_KEY" + - name: Build and test the universal Photos helper + if: runner.os == 'macOS' + working-directory: apps/desktop/native/photos-helper + run: | + set -euo pipefail + swift test + swift build -c release --arch arm64 --arch x86_64 + lipo -verify_arch .build/apple/Products/Release/photos-helper arm64 x86_64 - name: Package signed and notarized universal macOS app if: runner.os == 'macOS' env: @@ -174,6 +182,10 @@ jobs: grep -Fqx "provider: github" "$app/Contents/Resources/app-update.yml" grep -Fqx "owner: elie222" "$app/Contents/Resources/app-update.yml" grep -Fqx "repo: rakazo" "$app/Contents/Resources/app-update.yml" + helper="$app/Contents/Resources/bin/photos-helper" + codesign --verify --strict --verbose=2 "$helper" + codesign -d --entitlements - --xml "$helper" | grep -Fq "com.apple.security.personal-information.photos-library" + /usr/libexec/PlistBuddy -c 'Print :NSPhotoLibraryUsageDescription' "$app/Contents/Info.plist" - name: Require Windows signing credentials if: runner.os == 'Windows' diff --git a/apps/desktop/package.json b/apps/desktop/package.json index 575a833..9ac6e33 100644 --- a/apps/desktop/package.json +++ b/apps/desktop/package.json @@ -16,6 +16,8 @@ "release": "pnpm --filter @rakazo/web build && pnpm build && electron-builder --publish never" }, "dependencies": { + "@orpc/client": "1.15.0", + "@orpc/contract": "1.15.0", "@rakazo/contracts": "workspace:*", "electron-updater": "6.8.9" }, @@ -61,7 +63,18 @@ "dmg", "zip" ], - "notarize": true + "notarize": true, + "entitlements": "native/photos-helper/entitlements.mac.plist", + "entitlementsInherit": "native/photos-helper/entitlements.mac.plist", + "extendInfo": { + "NSPhotoLibraryUsageDescription": "Rakazo adds your photos and videos to the media library on your Rakazo server." + }, + "extraResources": [ + { + "from": "native/photos-helper/.build/apple/Products/Release/photos-helper", + "to": "bin/photos-helper" + } + ] }, "win": { "icon": "assets/icon.ico", diff --git a/apps/desktop/src/main.ts b/apps/desktop/src/main.ts index 304e998..37b8f27 100644 --- a/apps/desktop/src/main.ts +++ b/apps/desktop/src/main.ts @@ -2,7 +2,10 @@ import { randomBytes } from "node:crypto"; import { existsSync } from "node:fs"; import { readFile, stat, writeFile } from "node:fs/promises"; import path from "node:path"; -import type { DesktopReachability, DesktopSetup } from "@rakazo/contracts"; +import { createORPCClient } from "@orpc/client"; +import { RPCLink } from "@orpc/client/fetch"; +import type { ContractRouterClient } from "@orpc/contract"; +import type { AppContract, DesktopReachability, DesktopSetup } from "@rakazo/contracts"; import { LOCAL_SETTINGS_PAGE } from "@rakazo/contracts/local-settings"; import { app, @@ -31,6 +34,14 @@ import { stackDir, stackResourceDir, } from "./local-stack.js"; +import type { MediaIpcEvent } from "./media-ipc.js"; +import { registerMediaIpc } from "./media-ipc.js"; +import { createPhotosHelper, resolvePhotosHelperBinary } from "./media-photos-helper.js"; +import { + createChunkUploader, + createPhotosSyncFileStore, + createPhotosSyncRunner, +} from "./media-photos-sync.js"; import { oauthCallbackFrom } from "./oauth-callback.js"; import { bundledRendererCandidates, @@ -1129,6 +1140,60 @@ app.whenReady().then(async () => { ); }, ); + const photosBinary = resolvePhotosHelperBinary({ + platform: process.platform, + packaged: app.isPackaged, + resourcesPath: process.resourcesPath, + appPath: app.getAppPath(), + exists: existsSync, + }); + const photosHelper = photosBinary ? createPhotosHelper(photosBinary) : null; + const mediaIpc = registerMediaIpc({ + ipcMain, + platform: process.platform, + helper: photosHelper, + isTrusted: (raw: MediaIpcEvent) => { + const event = raw as Electron.IpcMainInvokeEvent; + return fromMainWindow(event) && event.senderFrame === event.sender.mainFrame; + }, + runSync: photosHelper + ? createPhotosSyncRunner({ + helper: photosHelper, + store: createPhotosSyncFileStore(userDataDir), + tempRoot: app.getPath("temp"), + label: "Apple Photos", + connect: async (spaceId) => { + if (currentTargetUrl === null) throw new Error("No server is open."); + const origin = new URL(currentTargetUrl).origin; + const target = (await resolveSessionForTarget(currentTargetUrl)).value; + const fetchWithSession = (input: string | Request, init?: RequestInit) => + target.fetch(input, { + ...init, + credentials: "include", + bypassCustomProtocolHandlers: true, + }); + const api: ContractRouterClient = createORPCClient( + new RPCLink({ + url: `${origin}/rpc`, + fetch: (request, init) => { + const headers = new Headers(request.headers); + if (spaceId) headers.set("x-rakazo-space-id", spaceId); + return fetchWithSession(new Request(request, { headers }), init); + }, + }), + ); + return { + origin, + api, + uploadChunk: createChunkUploader({ origin, spaceId, fetch: fetchWithSession }), + }; + }, + }) + : async () => { + throw new Error("Apple Photos is unavailable."); + }, + }); + app.on("before-quit", mediaIpc.cancel); ipcMain.handle("desktop.platform", () => process.platform); ipcMain.handle("desktop.files.pickDirectory", async (event) => { if (!fromMainWindow(event)) return null; diff --git a/apps/desktop/src/media-ipc.ts b/apps/desktop/src/media-ipc.ts index bb8e109..98fec49 100644 --- a/apps/desktop/src/media-ipc.ts +++ b/apps/desktop/src/media-ipc.ts @@ -1,6 +1,12 @@ -import type { PhotosAuthorization, PhotosHelper } from "./media-photos-helper.js"; +import type { + PhotosSourceStatus, + PhotosSyncOutcome, + PhotosSyncProgress, + PhotosSyncResult, + RakazoDesktopMediaPhotos, +} from "@rakazo/contracts"; +import type { PhotosHelper } from "./media-photos-helper.js"; import { PhotosHelperError } from "./media-photos-helper.js"; -import type { PhotosSyncProgress, PhotosSyncResult } from "./media-photos-sync.js"; import { PhotosSyncError } from "./media-photos-sync.js"; /* @@ -15,22 +21,7 @@ export const MEDIA_PHOTOS_CHANNELS = { progress: "desktop.media.photos.progress", } as const; -export type PhotosSourceStatus = - | { available: false } - | { available: true; authorization: PhotosAuthorization; syncing: boolean }; - -export type PhotosSyncOutcome = - | { ok: true; result: PhotosSyncResult } - | { ok: false; code: string; message: string }; - -/** Renderer-facing bridge, exposed as `rakazoDesktop.media.photos`. */ -export interface RakazoDesktopMediaPhotos { - status: () => Promise; - /** Asks for Photos access on first use, then indexes and uploads what the server needs. */ - sync: (request: { spaceId: string | null }) => Promise; - cancel: () => Promise; - onProgress: (listener: (progress: PhotosSyncProgress) => void) => () => void; -} +export type { PhotosSourceStatus, PhotosSyncOutcome, RakazoDesktopMediaPhotos }; export interface MediaIpcEvent { sender: { send(channel: string, payload: unknown): void; isDestroyed(): boolean }; diff --git a/apps/desktop/src/media-photos-helper.ts b/apps/desktop/src/media-photos-helper.ts index f10e402..56cec57 100644 --- a/apps/desktop/src/media-photos-helper.ts +++ b/apps/desktop/src/media-photos-helper.ts @@ -2,7 +2,7 @@ import { spawn as nodeSpawn } from "node:child_process"; import { realpath } from "node:fs/promises"; import path from "node:path"; import type { Readable } from "node:stream"; -import type { ClientMediaItem } from "@rakazo/contracts"; +import type { ClientMediaItem, PhotosAuthorization } from "@rakazo/contracts"; import { ClientMediaItemSchema } from "@rakazo/contracts"; /* @@ -11,12 +11,7 @@ import { ClientMediaItemSchema } from "@rakazo/contracts"; * the app uses it, so a malformed or unexpected answer fails instead of being guessed at. */ -export type PhotosAuthorization = - | "authorized" - | "limited" - | "denied" - | "restricted" - | "notDetermined"; +export type { PhotosAuthorization }; export interface PhotosListPage { items: ClientMediaItem[]; diff --git a/apps/desktop/src/media-photos-sync.ts b/apps/desktop/src/media-photos-sync.ts index 72ea3f9..332b9d7 100644 --- a/apps/desktop/src/media-photos-sync.ts +++ b/apps/desktop/src/media-photos-sync.ts @@ -2,7 +2,14 @@ import { createHash, randomUUID } from "node:crypto"; import { createReadStream } from "node:fs"; import { mkdtemp, open, rm, stat } from "node:fs/promises"; import path from "node:path"; -import type { ClientMediaItem, MediaAsset, MediaUpload, mediaRpcShapes } from "@rakazo/contracts"; +import type { + ClientMediaItem, + MediaAsset, + MediaUpload, + mediaRpcShapes, + PhotosSyncProgress, + PhotosSyncResult, +} from "@rakazo/contracts"; import type { PhotosHelper } from "./media-photos-helper.js"; import { PHOTOS_MAX_PAGE } from "./media-photos-helper.js"; import { readPrivateFile, writePrivateFile } from "./setup-store.js"; @@ -197,23 +204,7 @@ export function createChunkUploader(options: { // Sync // --------------------------------------------------------------------------- -export interface PhotosSyncProgress { - phase: "authorizing" | "indexing" | "uploading"; - /** Items pushed in this run. */ - indexed: number; - uploadsDone: number; - uploadsTotal: number; -} - -export interface PhotosSyncResult { - indexed: number; - created: number; - updated: number; - unchanged: number; - uploaded: number; - /** True when this run listed the whole library instead of reading changes. */ - fullListing: boolean; -} +export type { PhotosSyncProgress, PhotosSyncResult }; export interface PhotosSyncOptions { helper: PhotosHelper; diff --git a/apps/desktop/src/preload.cjs b/apps/desktop/src/preload.cjs index 14e5c50..d453090 100644 --- a/apps/desktop/src/preload.cjs +++ b/apps/desktop/src/preload.cjs @@ -9,6 +9,18 @@ contextBridge.exposeInMainWorld("rakazoDesktop", { request: (pathname, body) => ipcRenderer.invoke("desktop.localSettings.request", pathname, body), }, + media: { + photos: { + status: () => ipcRenderer.invoke("desktop.media.photos.status"), + sync: (request) => ipcRenderer.invoke("desktop.media.photos.sync", request), + cancel: () => ipcRenderer.invoke("desktop.media.photos.cancel"), + onProgress: (listener) => { + const handler = (_event, progress) => listener(progress); + ipcRenderer.on("desktop.media.photos.progress", handler); + return () => ipcRenderer.off("desktop.media.photos.progress", handler); + }, + }, + }, window: { close: () => ipcRenderer.invoke("desktop.window.close"), minimize: () => ipcRenderer.invoke("desktop.window.minimize"), diff --git a/apps/desktop/src/preload.test.ts b/apps/desktop/src/preload.test.ts index f5891c0..135d3ad 100644 --- a/apps/desktop/src/preload.test.ts +++ b/apps/desktop/src/preload.test.ts @@ -34,6 +34,7 @@ describe("desktop preload bridge", () => { expect(Object.keys(bridge).sort()).toEqual([ "files", "localSettings", + "media", "oauth", "platform", "update", @@ -46,6 +47,12 @@ describe("desktop preload bridge", () => { "toggleMaximize", ]); expect(Object.keys(bridge.update).sort()).toEqual(["check", "download", "install", "state"]); + expect(Object.keys(bridge.media?.photos ?? {}).sort()).toEqual([ + "cancel", + "onProgress", + "status", + "sync", + ]); await bridge.files?.pickDirectory(); await bridge.oauth.open?.("https://provider.example.com/authorize"); @@ -79,6 +86,7 @@ describe("desktop preload bridge", () => { expect(Object.keys(bridge).sort()).toEqual([ "files", "localSettings", + "media", "oauth", "platform", "update", diff --git a/packages/contracts/src/desktop.ts b/packages/contracts/src/desktop.ts index 53b78aa..94d5e2e 100644 --- a/packages/contracts/src/desktop.ts +++ b/packages/contracts/src/desktop.ts @@ -54,6 +54,48 @@ export interface RakazoDesktopBrowserPreview { }; } +export type PhotosAuthorization = + | "authorized" + | "limited" + | "denied" + | "restricted" + | "notDetermined"; + +export type PhotosSourceStatus = + | { available: false } + | { available: true; authorization: PhotosAuthorization; syncing: boolean }; + +export interface PhotosSyncProgress { + phase: "authorizing" | "indexing" | "uploading"; + /** Items pushed in this run. */ + indexed: number; + uploadsDone: number; + uploadsTotal: number; +} + +export interface PhotosSyncResult { + indexed: number; + created: number; + updated: number; + unchanged: number; + uploaded: number; + /** True when this run listed the whole library instead of reading changes. */ + fullListing: boolean; +} + +export type PhotosSyncOutcome = + | { ok: true; result: PhotosSyncResult } + | { ok: false; code: string; message: string }; + +/** Renderer-facing bridge, exposed as `rakazoDesktop.media.photos` on macOS desktops. */ +export interface RakazoDesktopMediaPhotos { + status: () => Promise; + /** Asks for Photos access on first use, then indexes and uploads what the server needs. */ + sync: (request: { spaceId: string | null }) => Promise; + cancel: () => Promise; + onProgress: (listener: (progress: PhotosSyncProgress) => void) => () => void; +} + export interface RakazoDesktop { /** Only the isolated local settings window is authorized to call this bridge. */ localSettings?: { @@ -64,6 +106,8 @@ export interface RakazoDesktop { pickDirectory: () => Promise; }; platform: string; + /** Present on desktops that ship the Apple Photos helper. */ + media?: { photos: RakazoDesktopMediaPhotos }; window: { close: () => Promise; minimize: () => Promise; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 51dbf6b..f4167b1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -97,6 +97,12 @@ importers: apps/desktop: dependencies: + '@orpc/client': + specifier: 1.15.0 + version: 1.15.0(@opentelemetry/api@1.9.1) + '@orpc/contract': + specifier: 1.15.0 + version: 1.15.0(@opentelemetry/api@1.9.1) '@rakazo/contracts': specifier: workspace:* version: link:../../packages/contracts From 5e2094f2adf6e669e185d5565f1746b67679eba8 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 18:19:47 +0000 Subject: [PATCH 2/3] feat(media): local inference sidecar for Media Studio 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 Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj --- infra/media-worker/.dockerignore | 7 + infra/media-worker/.gitignore | 5 + infra/media-worker/Dockerfile | 68 +++ infra/media-worker/LICENSES.md | 53 ++ infra/media-worker/README.md | 74 +++ infra/media-worker/app/__init__.py | 3 + infra/media-worker/app/__main__.py | 35 ++ infra/media-worker/app/auth.py | 22 + infra/media-worker/app/config.py | 80 +++ infra/media-worker/app/encoders.py | 41 ++ infra/media-worker/app/errors.py | 62 +++ infra/media-worker/app/main.py | 153 ++++++ infra/media-worker/app/media_io.py | 435 +++++++++++++++ infra/media-worker/app/models.py | 506 ++++++++++++++++++ infra/media-worker/app/ops/__init__.py | 1 + infra/media-worker/app/ops/audio.py | 163 ++++++ infra/media-worker/app/ops/common.py | 62 +++ infra/media-worker/app/ops/devices.py | 61 +++ infra/media-worker/app/ops/embeddings.py | 147 +++++ infra/media-worker/app/ops/faces.py | 265 +++++++++ infra/media-worker/app/ops/inpaint.py | 163 ++++++ infra/media-worker/app/ops/masks.py | 54 ++ infra/media-worker/app/ops/redact.py | 271 ++++++++++ infra/media-worker/app/ops/saliency.py | 58 ++ infra/media-worker/app/ops/scenes.py | 103 ++++ infra/media-worker/app/ops/segment.py | 244 +++++++++ infra/media-worker/app/ops/synthesize.py | 20 + infra/media-worker/app/ops/transcribe.py | 57 ++ infra/media-worker/app/paths.py | 75 +++ infra/media-worker/app/runtime.py | 93 ++++ infra/media-worker/app/schemas.py | 124 +++++ infra/media-worker/pytest.ini | 5 + infra/media-worker/requirements-dev.txt | 3 + infra/media-worker/requirements-ml-cpu.txt | 9 + infra/media-worker/requirements-ml-cuda.txt | 9 + infra/media-worker/requirements-ml-rocm.txt | 11 + infra/media-worker/requirements.txt | 10 + infra/media-worker/ruff.toml | 5 + infra/media-worker/tests/__init__.py | 0 infra/media-worker/tests/conftest.py | 224 ++++++++ infra/media-worker/tests/fakes.py | 133 +++++ .../tests/fixtures/media-capabilities.json | 21 + infra/media-worker/tests/test_analysis.py | 88 +++ infra/media-worker/tests/test_auth.py | 58 ++ infra/media-worker/tests/test_capabilities.py | 82 +++ infra/media-worker/tests/test_ml_ops.py | 170 ++++++ infra/media-worker/tests/test_models.py | 75 +++ infra/media-worker/tests/test_paths.py | 126 +++++ infra/media-worker/tests/test_redact.py | 174 ++++++ infra/media-worker/tests/test_runtime.py | 117 ++++ infra/media-worker/tests/test_schema.py | 137 +++++ 51 files changed, 4962 insertions(+) create mode 100644 infra/media-worker/.dockerignore create mode 100644 infra/media-worker/.gitignore create mode 100644 infra/media-worker/Dockerfile create mode 100644 infra/media-worker/LICENSES.md create mode 100644 infra/media-worker/README.md create mode 100644 infra/media-worker/app/__init__.py create mode 100644 infra/media-worker/app/__main__.py create mode 100644 infra/media-worker/app/auth.py create mode 100644 infra/media-worker/app/config.py create mode 100644 infra/media-worker/app/encoders.py create mode 100644 infra/media-worker/app/errors.py create mode 100644 infra/media-worker/app/main.py create mode 100644 infra/media-worker/app/media_io.py create mode 100644 infra/media-worker/app/models.py create mode 100644 infra/media-worker/app/ops/__init__.py create mode 100644 infra/media-worker/app/ops/audio.py create mode 100644 infra/media-worker/app/ops/common.py create mode 100644 infra/media-worker/app/ops/devices.py create mode 100644 infra/media-worker/app/ops/embeddings.py create mode 100644 infra/media-worker/app/ops/faces.py create mode 100644 infra/media-worker/app/ops/inpaint.py create mode 100644 infra/media-worker/app/ops/masks.py create mode 100644 infra/media-worker/app/ops/redact.py create mode 100644 infra/media-worker/app/ops/saliency.py create mode 100644 infra/media-worker/app/ops/scenes.py create mode 100644 infra/media-worker/app/ops/segment.py create mode 100644 infra/media-worker/app/ops/synthesize.py create mode 100644 infra/media-worker/app/ops/transcribe.py create mode 100644 infra/media-worker/app/paths.py create mode 100644 infra/media-worker/app/runtime.py create mode 100644 infra/media-worker/app/schemas.py create mode 100644 infra/media-worker/pytest.ini create mode 100644 infra/media-worker/requirements-dev.txt create mode 100644 infra/media-worker/requirements-ml-cpu.txt create mode 100644 infra/media-worker/requirements-ml-cuda.txt create mode 100644 infra/media-worker/requirements-ml-rocm.txt create mode 100644 infra/media-worker/requirements.txt create mode 100644 infra/media-worker/ruff.toml create mode 100644 infra/media-worker/tests/__init__.py create mode 100644 infra/media-worker/tests/conftest.py create mode 100644 infra/media-worker/tests/fakes.py create mode 100644 infra/media-worker/tests/fixtures/media-capabilities.json create mode 100644 infra/media-worker/tests/test_analysis.py create mode 100644 infra/media-worker/tests/test_auth.py create mode 100644 infra/media-worker/tests/test_capabilities.py create mode 100644 infra/media-worker/tests/test_ml_ops.py create mode 100644 infra/media-worker/tests/test_models.py create mode 100644 infra/media-worker/tests/test_paths.py create mode 100644 infra/media-worker/tests/test_redact.py create mode 100644 infra/media-worker/tests/test_runtime.py create mode 100644 infra/media-worker/tests/test_schema.py diff --git a/infra/media-worker/.dockerignore b/infra/media-worker/.dockerignore new file mode 100644 index 0000000..949e720 --- /dev/null +++ b/infra/media-worker/.dockerignore @@ -0,0 +1,7 @@ +.pytest_cache +.ruff_cache +.mypy_cache +.venv +.venv-ml +**/__pycache__ +**/*.pyc diff --git a/infra/media-worker/.gitignore b/infra/media-worker/.gitignore new file mode 100644 index 0000000..01304c1 --- /dev/null +++ b/infra/media-worker/.gitignore @@ -0,0 +1,5 @@ +.pytest_cache/ +.venv/ +.venv +.venv-ml +__pycache__/ diff --git a/infra/media-worker/Dockerfile b/infra/media-worker/Dockerfile new file mode 100644 index 0000000..baa6bee --- /dev/null +++ b/infra/media-worker/Dockerfile @@ -0,0 +1,68 @@ +# syntax=docker/dockerfile:1.7 +# Media Studio inference sidecar. +# docker build --build-arg ACCELERATOR=cpu|cuda|rocm -t rakazo/media-worker infra/media-worker +# docker build --target test -t rakazo/media-worker:test infra/media-worker # offline pytest +# Model weights are not baked in; they download into /models on first use. +# Behind a TLS-intercepting proxy, pass its CA with --secret id=pip_ca,src=. + +ARG ACCELERATOR=cpu + +FROM python:3.11-slim-bookworm@sha256:2333bd330d12de02514770b3585cad313644316047cdee24a7acfdece6de6efb AS runtime +ARG ACCELERATOR + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + TQDM_DISABLE=1 \ + MEDIA_ACCELERATOR=${ACCELERATOR} \ + MEDIA_MODELS_DIR=/models \ + MEDIA_WORKER_PORT=8790 + +RUN case "${ACCELERATOR}" in cpu|cuda|rocm) ;; *) echo "ACCELERATOR must be cpu, cuda or rocm" >&2; exit 1 ;; esac \ + && apt-get update \ + && apt-get install -y --no-install-recommends ffmpeg libsndfile1 \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /srv/media-worker +COPY requirements.txt requirements-ml-cpu.txt requirements-ml-cuda.txt requirements-ml-rocm.txt ./ + +# faster-whisper depends on the CPU onnxruntime wheel; the CUDA image swaps in +# onnxruntime-gpu. SAM2 is an sdist built against the installed torch, without +# its optional CUDA extension. +RUN --mount=type=secret,id=pip_ca,required=false \ + if [ -s /run/secrets/pip_ca ]; then export PIP_CERT=/run/secrets/pip_ca; fi \ + && pip install -r requirements.txt -r "requirements-ml-${ACCELERATOR}.txt" \ + && if [ "${ACCELERATOR}" = "cuda" ]; then \ + pip uninstall -y onnxruntime && pip install --no-deps --force-reinstall onnxruntime-gpu==1.30.0; \ + fi \ + && SAM2_BUILD_CUDA=0 pip install --no-build-isolation --no-deps sam2==1.1.0 \ + && python -c "import cv2, scenedetect, librosa, onnxruntime, faster_whisper, torch, sam2, tokenizers" + +RUN useradd --uid 10001 --create-home --shell /usr/sbin/nologin media \ + && mkdir -p /models \ + && chown media:media /models +VOLUME ["/models"] + +COPY app ./app + +USER media +EXPOSE 8790 +HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \ + CMD ["python", "-c", "import os, urllib.request; urllib.request.urlopen('http://127.0.0.1:' + os.environ.get('MEDIA_WORKER_PORT', '8790') + '/healthz', timeout=4)"] +CMD ["python", "-m", "app"] + + +# Offline test image: same runtime plus pytest; fake backends, no weights. +FROM runtime AS test +USER root +COPY requirements-dev.txt pytest.ini ./ +RUN --mount=type=secret,id=pip_ca,required=false \ + if [ -s /run/secrets/pip_ca ]; then export PIP_CERT=/run/secrets/pip_ca; fi \ + && pip install -r requirements-dev.txt +COPY tests ./tests +USER media +CMD ["python", "-m", "pytest", "-q"] + + +FROM runtime diff --git a/infra/media-worker/LICENSES.md b/infra/media-worker/LICENSES.md new file mode 100644 index 0000000..1a349a8 --- /dev/null +++ b/infra/media-worker/LICENSES.md @@ -0,0 +1,53 @@ +# Licenses: models and libraries of the media sidecar + +Checked October 2026 against the upstream LICENSE files and model cards linked +below. "Commercial use" means the published code or weights license permits it. +Training-data terms were not reviewed legally; where provenance is unclear the +backend is optional, never a default. Weights are downloaded at runtime from the +pinned revision (see `app/models.py`) and are never part of this repository or +the image. + +## Models + +| Name | Purpose | License | Commercial use | Default / optional | Source | +|---|---|---|---|---|---| +| OpenCV YuNet 2023mar | Face detection, 5-point landmarks (`detectFaces`, `trackFaces`, `faceLandmarks`) | MIT | yes | default | https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet | +| Whisper small (CTranslate2 conversion by Systran) | Transcription with word timestamps (`transcribe`) | MIT | yes | default | https://huggingface.co/Systran/faster-whisper-small, https://github.com/openai/whisper | +| SigLIP base patch16-224 (ONNX export by Xenova) | Frame embeddings (`embedFrames`); zero-shot content scores (`classifyContent`, experimental) | Apache-2.0 | yes | default | https://huggingface.co/google/siglip-base-patch16-224, https://huggingface.co/Xenova/siglip-base-patch16-224 | +| SAM 2.1 Hiera tiny | Click/box segmentation with video propagation (`segmentPrompted`, experimental) | Apache-2.0 | yes | default | https://github.com/facebookresearch/sam2, https://huggingface.co/facebook/sam2.1-hiera-tiny | +| Big-LaMa (ONNX export by Carve) | Per-frame inpainting with optical-flow smoothing (`inpaintVideo`, experimental) | Apache-2.0 | yes | default | https://github.com/advimman/lama, https://huggingface.co/Carve/LaMa-ONNX | +| OpenCV SFace 2021dec | Identity embeddings for local person clustering (`trackFaces.embedding`) | Apache-2.0 (weights); training-data provenance unverified | unclear | optional (`MEDIA_OPTIONAL_BACKENDS=faceIdentity`) | https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface | +| MediaPipe Face Landmarker | Dense face landmarks | Apache-2.0 | yes | not used: the `mediapipe` wheel pulls the GUI OpenCV build; YuNet's 5 points cover the partial modes | https://github.com/google-ai-edge/mediapipe | +| OpenCLIP (LAION weights) | Frame embeddings | MIT | yes | not used: SigLIP runs on ONNX Runtime without torch | https://github.com/mlfoundations/open_clip | +| NudeNet v3 | Nudity detection | AGPL-3.0 | copyleft, not shipped | not used | https://github.com/notAI-tech/NudeNet | +| ProPainter | Video inpainting | S-Lab License 1.0 (non-commercial) | no | not used; never default | https://github.com/sczhou/ProPainter | +| InsightFace model zoo (SCRFD, ArcFace, inswapper) | Face detection, recognition, swap | code MIT; models non-commercial research only | no | not used; never default | https://github.com/deepinsight/insightface | +| Synthetic identity generator + swap | `synthesizeIdentity` | none found with a commercial license | — | unavailable: always 501 `unsupported` | — | + +## Methods without weights + +| Name | Purpose | License | Commercial use | Default / optional | Source | +|---|---|---|---|---|---| +| PySceneDetect (AdaptiveDetector) | Shot detection (`detectScenes`) | BSD-3-Clause | yes | default | https://github.com/Breakthrough/PySceneDetect | +| IoU tracker (ByteTrack-style, in `app/ops/faces.py`) | Face tracking | this repository | yes | default | — | +| Spectral residual saliency (Hou & Zhang 2007, in `app/ops/saliency.py`) | `saliency` | this repository | yes | default | https://doi.org/10.1109/CVPR.2007.383267 | +| FFmpeg `ebur128`, `silencedetect` | Loudness, true peak, silences (`audioFeatures`) | LGPL-2.1+ / GPL-2+ (Debian build with libx264) | yes, as a separate process | default | https://ffmpeg.org/legal.html | +| librosa | Beat and tempo tracking (`audioFeatures`) | ISC | yes | default | https://github.com/librosa/librosa | + +## Libraries + +| Name | License | Commercial use | +|---|---|---| +| FastAPI | MIT | yes | +| Uvicorn | BSD-3-Clause | yes | +| Pydantic | MIT | yes | +| NumPy | BSD-3-Clause | yes | +| OpenCV (opencv-python-headless) | Apache-2.0 (wheel bundles FFmpeg under LGPL) | yes | +| soundfile | BSD-3-Clause | yes | +| PyTorch, torchvision | BSD-3-Clause (bundles permissive third-party code: Apache-2.0, MIT, BSL-1.0) | yes | +| ONNX Runtime | MIT | yes | +| faster-whisper | MIT | yes | +| CTranslate2 | MIT | yes | +| Hugging Face tokenizers | Apache-2.0 | yes | +| SAM 2 package (`sam2`) | Apache-2.0 | yes | +| Hydra, iopath | MIT | yes | diff --git a/infra/media-worker/README.md b/infra/media-worker/README.md new file mode 100644 index 0000000..420cb80 --- /dev/null +++ b/infra/media-worker/README.md @@ -0,0 +1,74 @@ +# media-worker + +Local inference sidecar for Media Studio (Python 3.11, FastAPI). It implements +every `MediaInferenceProvider` operation (`packages/adapter-kit/src/media.ts`) +behind HTTP API v1; the TypeScript client is `packages/adapters/src/media-sidecar.ts`. + +## API + +- `GET /healthz` → `{"ok": true, "version": "..."}` (no token) +- `GET /v1/capabilities` → `MediaCapabilities` +- `POST /v1/` with the TypeScript request object (minus `signal`) +- Errors: `{"error": {"code", "message", "operation"}}` with + `bad_request`/`invalid_path` 400, `unauthorized` 401, `too_large` 413, + `unsupported` 501, `busy` 503, `internal` 500. + +Every route but `/healthz` needs `Authorization: Bearer $MEDIA_WORKER_TOKEN`. +Paths must be absolute and resolve (symlinks included) inside `MEDIA_ROOT`. +A request stops when its client disconnects (between frames or chunks) or +after `MEDIA_MAX_SECONDS`. + +## Configuration + +| Variable | Default | | +|---|---|---| +| `MEDIA_WORKER_TOKEN` | — | required, ≥ 32 chars (`MEDIA_WORKER_INSECURE_DEV=1` lifts this for tests) | +| `MEDIA_ROOT` | — | required; mount it at the same absolute path as the API's `DATA_DIR/media` | +| `MEDIA_MODELS_DIR` | `/models` | weight cache (volume) | +| `MEDIA_WORKER_HOST` / `MEDIA_WORKER_PORT` | `127.0.0.1` / `8790` | | +| `MEDIA_GPU_CONCURRENCY` | `1` | model operations running at once | +| `MEDIA_CPU_CONCURRENCY` | `2` | scenes, redact, audio, saliency running at once | +| `MEDIA_QUEUE_SIZE` | `8` | waiting requests per gate before `503 busy` | +| `MEDIA_MAX_INPUT_BYTES` | 20 GiB | larger inputs → 413 | +| `MEDIA_MAX_SECONDS` | 14400 | longer inputs → 413; longer processing is cancelled | +| `MEDIA_MODEL_DOWNLOAD` | `1` | `0` = offline; missing weights → 501 | +| `MEDIA_OPTIONAL_BACKENDS` | empty | opt-in for backends that are not clean for commercial use, e.g. `faceIdentity` | +| `MEDIA_ACCELERATOR` | build arg | `cpu`, `cuda` or `rocm`; falls back to cpu when no GPU is visible | + +## Models + +Backends, licenses, URLs, SHA256 and sizes are pinned in `app/models.py`; the +license review is in `LICENSES.md`. Weights download on first use into +`MEDIA_MODELS_DIR//`, are checked against the pinned SHA256 and size, +and are deleted on mismatch (the operation then reports 501). Warm the volume +ahead of time with `python -m app prefetch`. + +`synthesizeIdentity` always answers 501: there is no commercially licensed +generator and swap pipeline, and the request schema rejects any extra field, +so a target-face image cannot be passed. + +## Redaction geometry + +Boxes interpolate linearly between keyframes; a single keyframe is a static +region for the whole clip. `paddingPct` grows the box by that share of its +size, `featherPct` fades the effect outward over that share of the shorter +side (the region itself stays fully covered). Partial modes use the first two +landmarks as eyes and the third as the nose. Output is FFV1 for `.mkv`, +otherwise x264 CRF 12; audio is copied and metadata dropped. Any frame error is +a 500 and nothing is left at `outPath`. + +## Build and test + +```sh +docker build --build-arg ACCELERATOR=cpu -t rakazo/media-worker infra/media-worker +docker build --target test -t rakazo/media-worker:test infra/media-worker +docker run --rm rakazo/media-worker:test # offline pytest, fake model backends +``` + +Without Docker: `pip install -r requirements.txt -r requirements-dev.txt` and +`python -m pytest` in this directory. Operations whose ML packages are missing +report `available: false`. + +CUDA uses the torch cu126 wheels and `onnxruntime-gpu`. ROCm uses the torch +rocm7.1 wheels for SAM 2; ONNX Runtime and CTranslate2 have no ROCm wheels on +PyPI, so SigLIP, LaMa and Whisper run on the CPU there. diff --git a/infra/media-worker/app/__init__.py b/infra/media-worker/app/__init__.py new file mode 100644 index 0000000..dc450d5 --- /dev/null +++ b/infra/media-worker/app/__init__.py @@ -0,0 +1,3 @@ +"""Local Media Studio inference sidecar.""" + +VERSION = "0.1.0" diff --git a/infra/media-worker/app/__main__.py b/infra/media-worker/app/__main__.py new file mode 100644 index 0000000..e2d507c --- /dev/null +++ b/infra/media-worker/app/__main__.py @@ -0,0 +1,35 @@ +"""Entry point: `python -m app` serves the API, `python -m app prefetch` warms the model volume.""" + +from __future__ import annotations + +import logging +import sys + +import uvicorn + +from .config import ConfigError, Settings +from .main import create_app +from .models import Registry + + +def main(argv: list[str]) -> int: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s") + try: + settings = Settings.from_env() + except ConfigError as error: + print(f"media-worker: {error}", file=sys.stderr) + return 2 + if argv[1:] == ["prefetch"]: + Registry(settings).prefetch() + return 0 + if argv[1:]: + print("usage: python -m app [prefetch]", file=sys.stderr) + return 2 + uvicorn.run( + create_app(settings), host=settings.host, port=settings.port, log_level="info", access_log=False + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/infra/media-worker/app/auth.py b/infra/media-worker/app/auth.py new file mode 100644 index 0000000..01b3258 --- /dev/null +++ b/infra/media-worker/app/auth.py @@ -0,0 +1,22 @@ +"""Bearer-token check for every route except /healthz.""" + +from __future__ import annotations + +import hmac + +from fastapi import Request + +from .errors import MediaError + + +def require_token(request: Request) -> None: + expected: str = request.app.state.settings.token + header = request.headers.get("authorization", "") + scheme, _, presented = header.partition(" ") + if ( + scheme.lower() != "bearer" + or not presented + or not expected + or not hmac.compare_digest(presented.strip().encode(), expected.encode()) + ): + raise MediaError("unauthorized", "missing or invalid bearer token") diff --git a/infra/media-worker/app/config.py b/infra/media-worker/app/config.py new file mode 100644 index 0000000..9cc15aa --- /dev/null +++ b/infra/media-worker/app/config.py @@ -0,0 +1,80 @@ +"""Sidecar settings, read once from the environment.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field + +MIN_TOKEN_LENGTH = 32 +GIB = 1024**3 + + +class ConfigError(Exception): + pass + + +def _int(env: dict[str, str], name: str, default: int, minimum: int = 0) -> int: + raw = env.get(name, "").strip() + if not raw: + return default + try: + value = int(raw) + except ValueError as error: + raise ConfigError(f"{name} must be an integer") from error + if value < minimum: + raise ConfigError(f"{name} must be >= {minimum}") + return value + + +@dataclass(frozen=True) +class Settings: + token: str + media_root: str + models_dir: str + host: str = "127.0.0.1" + port: int = 8790 + accelerator: str = "cpu" + gpu_concurrency: int = 1 + cpu_concurrency: int = 2 + queue_size: int = 8 + max_input_bytes: int = 20 * GIB + max_seconds: int = 4 * 3600 + allow_model_download: bool = True + # Backends that are not licensed for commercial use stay off unless named here. + optional_backends: frozenset[str] = field(default_factory=frozenset) + insecure_dev: bool = False + + @staticmethod + def from_env(env: dict[str, str] | None = None) -> Settings: + env = dict(os.environ if env is None else env) + insecure = env.get("MEDIA_WORKER_INSECURE_DEV", "") == "1" + token = env.get("MEDIA_WORKER_TOKEN", "") + if len(token) < MIN_TOKEN_LENGTH and not insecure: + raise ConfigError(f"MEDIA_WORKER_TOKEN must be at least {MIN_TOKEN_LENGTH} characters") + media_root = env.get("MEDIA_ROOT", "").strip() + if not media_root or not os.path.isabs(media_root): + raise ConfigError("MEDIA_ROOT must be an absolute path") + if not os.path.isdir(media_root): + raise ConfigError("MEDIA_ROOT does not exist") + accelerator = env.get("MEDIA_ACCELERATOR", "cpu").strip() or "cpu" + if accelerator not in ("cpu", "cuda", "rocm"): + raise ConfigError("MEDIA_ACCELERATOR must be cpu, cuda or rocm") + optional = frozenset( + item.strip() for item in env.get("MEDIA_OPTIONAL_BACKENDS", "").split(",") if item.strip() + ) + return Settings( + token=token, + media_root=media_root, + models_dir=env.get("MEDIA_MODELS_DIR", "/models").strip() or "/models", + host=env.get("MEDIA_WORKER_HOST", "127.0.0.1").strip() or "127.0.0.1", + port=_int(env, "MEDIA_WORKER_PORT", 8790, 1), + accelerator=accelerator, + gpu_concurrency=_int(env, "MEDIA_GPU_CONCURRENCY", 1, 1), + cpu_concurrency=_int(env, "MEDIA_CPU_CONCURRENCY", 2, 1), + queue_size=_int(env, "MEDIA_QUEUE_SIZE", 8, 0), + max_input_bytes=_int(env, "MEDIA_MAX_INPUT_BYTES", 20 * GIB, 1), + max_seconds=_int(env, "MEDIA_MAX_SECONDS", 4 * 3600, 1), + allow_model_download=env.get("MEDIA_MODEL_DOWNLOAD", "1") != "0", + optional_backends=optional, + insecure_dev=insecure, + ) diff --git a/infra/media-worker/app/encoders.py b/infra/media-worker/app/encoders.py new file mode 100644 index 0000000..bc042bd --- /dev/null +++ b/infra/media-worker/app/encoders.py @@ -0,0 +1,41 @@ +"""Hardware video encoders that actually encode a frame on this host.""" + +from __future__ import annotations + +import functools +import os +import re +import subprocess + +from .media_io import FFMPEG + +CANDIDATE = re.compile(r"^\s*V\S*\s+((?:h264|hevc|av1)_(?:nvenc|vaapi|qsv|amf|videotoolbox|v4l2m2m))\s") +RENDER_NODE = "/dev/dri/renderD128" + + +def _works(name: str) -> bool: + argv = [FFMPEG, "-v", "error", "-nostdin"] + if name.endswith("_vaapi"): + if not os.path.exists(RENDER_NODE): + return False + argv += ["-vaapi_device", RENDER_NODE] + argv += ["-f", "lavfi", "-i", "color=c=black:s=256x256:d=0.1", "-frames:v", "1"] + if name.endswith("_vaapi"): + argv += ["-vf", "format=nv12,hwupload"] + argv += ["-c:v", name, "-f", "null", "-"] + try: + return subprocess.run(argv, capture_output=True, timeout=15, shell=False).returncode == 0 + except (OSError, subprocess.TimeoutExpired): + return False + + +@functools.cache +def hardware_encoders() -> tuple[str, ...]: + try: + listing = subprocess.run( + [FFMPEG, "-hide_banner", "-encoders"], capture_output=True, timeout=15, shell=False, text=True + ).stdout + except (OSError, subprocess.TimeoutExpired): + return () + names = [m.group(1) for line in listing.splitlines() if (m := CANDIDATE.match(line))] + return tuple(name for name in names if _works(name)) diff --git a/infra/media-worker/app/errors.py b/infra/media-worker/app/errors.py new file mode 100644 index 0000000..ede69d0 --- /dev/null +++ b/infra/media-worker/app/errors.py @@ -0,0 +1,62 @@ +"""Error envelope shared by every route: {"error": {"code", "message", "operation"}}.""" + +from __future__ import annotations + +from fastapi.responses import JSONResponse + +STATUS_BY_CODE = { + "bad_request": 400, + "invalid_path": 400, + "unauthorized": 401, + "too_large": 413, + "unsupported": 501, + "busy": 503, + "internal": 500, +} + + +class MediaError(Exception): + def __init__(self, code: str, message: str, operation: str | None = None) -> None: + super().__init__(message) + self.code = code + self.message = message + self.operation = operation + + @property + def status(self) -> int: + return STATUS_BY_CODE[self.code] + + +class Unsupported(MediaError): + def __init__(self, reason: str, operation: str | None = None) -> None: + super().__init__("unsupported", reason, operation) + + +class InvalidPath(MediaError): + def __init__(self, message: str = "path is outside MEDIA_ROOT") -> None: + super().__init__("invalid_path", message) + + +class BadRequest(MediaError): + def __init__(self, message: str) -> None: + super().__init__("bad_request", message) + + +class Busy(MediaError): + def __init__(self) -> None: + super().__init__("busy", "inference queue is full") + + +class Cancelled(Exception): + """The client went away or the request ran past MEDIA_MAX_SECONDS.""" + + def __init__(self, reason: str = "client disconnected") -> None: + super().__init__(reason) + self.reason = reason + + +def error_response(code: str, message: str, operation: str | None, status: int | None = None) -> JSONResponse: + return JSONResponse( + status_code=status or STATUS_BY_CODE[code], + content={"error": {"code": code, "message": message, "operation": operation}}, + ) diff --git a/infra/media-worker/app/main.py b/infra/media-worker/app/main.py new file mode 100644 index 0000000..5341fc5 --- /dev/null +++ b/infra/media-worker/app/main.py @@ -0,0 +1,153 @@ +"""HTTP API v1 of the Media Studio sidecar. + +GET /healthz is open; everything else needs the bearer token. POST +/v1/ takes the TypeScript request object (minus `signal`) and +answers with the TypeScript result object. +""" + +from __future__ import annotations + +import logging +import threading +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import anyio +from fastapi import Depends, FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse, Response +from pydantic import BaseModel, ValidationError +from starlette.exceptions import HTTPException as StarletteHTTPException + +from . import VERSION, encoders, schemas +from .auth import require_token +from .config import Settings +from .errors import Cancelled, MediaError, error_response +from .models import OPERATIONS, Registry +from .ops import audio, embeddings, faces, inpaint, redact, saliency, scenes, segment, synthesize, transcribe +from .ops.common import Ctx, jsonable +from .ops.devices import effective_accelerator +from .runtime import Gate, run_job + +log = logging.getLogger("media-worker") + +MAX_BODY_BYTES = 64 * 1024 * 1024 +# Status for a request whose client went away; nobody reads it. +CLIENT_CLOSED = 499 + + +@dataclass(frozen=True) +class Route: + schema: type[BaseModel] + run: Callable[[Ctx, Any, threading.Event], dict[str, Any]] + gpu: bool + + +ROUTES: dict[str, Route] = { + "detectScenes": Route(schemas.PathRequest, scenes.run, False), + "transcribe": Route(schemas.TranscribeRequest, transcribe.run, True), + "detectFaces": Route(schemas.DetectFacesRequest, faces.detect_run, True), + "trackFaces": Route(schemas.TrackFacesRequest, faces.track_run, True), + "faceLandmarks": Route(schemas.TimesRequest, faces.landmarks_run, True), + "embedFrames": Route(schemas.TimesRequest, embeddings.embed_run, True), + "classifyContent": Route(schemas.ClassifyContentRequest, embeddings.classify_run, True), + "segmentPrompted": Route(schemas.SegmentPromptedRequest, segment.run, True), + "inpaintVideo": Route(schemas.InpaintVideoRequest, inpaint.run, True), + "redact": Route(schemas.RedactRequest, redact.run, False), + "synthesizeIdentity": Route(schemas.SynthesizeIdentityRequest, synthesize.run, True), + "audioFeatures": Route(schemas.PathRequest, audio.run, False), + "saliency": Route(schemas.TimesRequest, saliency.run, False), +} +assert set(ROUTES) == set(OPERATIONS) + + +def _operation_of(request: Request) -> str | None: + name = request.url.path.rsplit("/", 1)[-1] + return name if name in ROUTES else None + + +def _validation_message(error: ValidationError) -> str: + first = error.errors()[0] if error.errors() else {} + where = ".".join(str(part) for part in first.get("loc", ())) + return ( + f"{where}: {first.get('msg', 'invalid request')}" + if where + else str(first.get("msg", "invalid request")) + ) + + +def create_app(settings: Settings | None = None, registry: Registry | None = None) -> FastAPI: + settings = settings or Settings.from_env() + registry = registry or Registry(settings) + app = FastAPI(title="media-worker", version=VERSION, docs_url=None, redoc_url=None, openapi_url=None) + app.state.settings = settings + app.state.ctx = Ctx(settings, registry) + app.state.gpu_gate = Gate(settings.gpu_concurrency, settings.queue_size) + app.state.cpu_gate = Gate(settings.cpu_concurrency, settings.queue_size) + + @app.exception_handler(MediaError) + async def media_error(request: Request, error: MediaError) -> JSONResponse: + return error_response(error.code, error.message, error.operation or _operation_of(request)) + + @app.exception_handler(RequestValidationError) + async def validation_error(request: Request, error: RequestValidationError) -> JSONResponse: + return error_response("bad_request", "invalid request", _operation_of(request)) + + @app.exception_handler(StarletteHTTPException) + async def http_error(request: Request, error: StarletteHTTPException) -> JSONResponse: + code = "bad_request" if error.status_code < 500 else "internal" + return error_response(code, str(error.detail), _operation_of(request), status=error.status_code) + + @app.exception_handler(Exception) + async def unexpected(request: Request, error: Exception) -> JSONResponse: + log.exception("unhandled error") + return error_response("internal", "internal error", _operation_of(request)) + + @app.get("/healthz") + async def healthz() -> dict[str, Any]: + return {"ok": True, "version": VERSION} + + @app.get("/v1/capabilities", dependencies=[Depends(require_token)]) + async def capabilities() -> dict[str, Any]: + hardware = await anyio.to_thread.run_sync(encoders.hardware_encoders) + return { + "configured": True, + "operations": registry.operation_capabilities(), + "accelerator": effective_accelerator(settings), + "hardwareEncoders": list(hardware), + } + + @app.post("/v1/{operation}", dependencies=[Depends(require_token)]) + async def operate(operation: str, request: Request) -> Response: + route = ROUTES.get(operation) + if route is None: + return error_response("bad_request", "unknown operation", None, status=404) + raw = bytearray() + async for chunk in request.stream(): + raw.extend(chunk) + if len(raw) > MAX_BODY_BYTES: + raise MediaError("too_large", "request body too large", operation) + try: + body = route.schema.model_validate_json(bytes(raw)) + except ValidationError as error: + raise MediaError("bad_request", _validation_message(error), operation) from error + ctx: Ctx = request.app.state.ctx + gate: Gate = request.app.state.gpu_gate if route.gpu else request.app.state.cpu_gate + try: + result = await run_job( + request, gate, lambda cancel: route.run(ctx, body, cancel), settings.max_seconds + ) + except MediaError as error: + error.operation = operation + raise + except Cancelled as error: + if error.reason == "client disconnected": + return Response(status_code=CLIENT_CLOSED) + raise MediaError("internal", error.reason, operation) from error + except Exception as error: + log.exception("%s failed", operation) + raise MediaError("internal", "internal error", operation) from error + return JSONResponse(jsonable(result)) + + return app diff --git a/infra/media-worker/app/media_io.py b/infra/media-worker/app/media_io.py new file mode 100644 index 0000000..7d3c9a7 --- /dev/null +++ b/infra/media-worker/app/media_io.py @@ -0,0 +1,435 @@ +"""ffmpeg/ffprobe plumbing. Every call is an argv list; nothing goes through a shell.""" + +from __future__ import annotations + +import contextlib +import json +import os +import queue +import re +import subprocess +import threading +import uuid +from collections import deque +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from fractions import Fraction + +import numpy as np + +from .errors import BadRequest, Cancelled, MediaError + +FFMPEG = "ffmpeg" +FFPROBE = "ffprobe" +SHOWINFO_PTS = re.compile(rb"pts_time:\s*(-?[0-9.]+)") + + +@dataclass(frozen=True) +class VideoInfo: + width: int + height: int + fps: float + fps_text: str + duration_s: float + has_audio: bool + has_video: bool + start_s: float = 0.0 + # Video stream duration when the container reports it (audio may run longer). + video_duration_s: float = 0.0 + + def frame_ms(self) -> float: + return 1000.0 / self.fps + + +def run_cancellable( + argv: list[str], cancel: threading.Event, timeout: float | None = None +) -> subprocess.CompletedProcess[bytes]: + """Run a short-lived tool, killing it when the request is cancelled.""" + proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False) + killer = _kill_on_cancel(proc, cancel) + try: + stdout, stderr = proc.communicate(timeout=timeout) + finally: + killer.set() + if cancel.is_set(): + raise Cancelled() + return subprocess.CompletedProcess(argv, proc.returncode, stdout, stderr) + + +def _kill_on_cancel(proc: subprocess.Popen[bytes], cancel: threading.Event) -> threading.Event: + done = threading.Event() + + def watch() -> None: + while not done.is_set(): + if cancel.wait(0.2): + if proc.poll() is None: + proc.kill() + return + + threading.Thread(target=watch, daemon=True).start() + return done + + +def _parse_rate(text: str | None) -> float: + if not text or text in ("0/0", "0"): + return 0.0 + try: + return float(Fraction(text)) + except (ValueError, ZeroDivisionError): + return 0.0 + + +def probe(path: str, cancel: threading.Event) -> VideoInfo: + result = run_cancellable( + [ + FFPROBE, + "-v", + "error", + "-print_format", + "json", + "-show_streams", + "-show_format", + path, + ], + cancel, + ) + if result.returncode != 0: + raise BadRequest("input is not a readable media file") + data = json.loads(result.stdout or b"{}") + streams = data.get("streams", []) + video = next( + ( + s + for s in streams + if s.get("codec_type") == "video" and not s.get("disposition", {}).get("attached_pic") + ), + None, + ) + audio = next((s for s in streams if s.get("codec_type") == "audio"), None) + duration = float(data.get("format", {}).get("duration") or 0.0) + if video is None: + return VideoInfo(0, 0, 0.0, "0", duration, audio is not None, False) + fps_text = video.get("avg_frame_rate") or "0/0" + fps = _parse_rate(fps_text) + if fps <= 0 or fps > 240: + fps_text = video.get("r_frame_rate") or "0/0" + fps = _parse_rate(fps_text) + if fps <= 0: + fps_text, fps = "25", 25.0 + width, height = int(video.get("width") or 0), int(video.get("height") or 0) + rotation = _rotation(video) + if rotation in (90, 270): + width, height = height, width + if not duration: + duration = float(video.get("duration") or 0.0) + start = float(video.get("start_time") or 0.0) + video_duration = float(video.get("duration") or 0.0) or duration + return VideoInfo(width, height, fps, fps_text, duration, audio is not None, True, start, video_duration) + + +def _rotation(stream: dict) -> int: + for side in stream.get("side_data_list", []) or []: + if "rotation" in side: + return int(abs(int(side["rotation"]))) % 360 + tag = (stream.get("tags") or {}).get("rotate") + return int(abs(int(tag))) % 360 if tag else 0 + + +def check_duration(info: VideoInfo, max_seconds: int) -> None: + if info.duration_s > max_seconds: + raise MediaError("too_large", "input duration exceeds MEDIA_MAX_SECONDS") + + +def scaled_size(width: int, height: int, max_side: int | None) -> tuple[int, int]: + if not max_side or max(width, height) <= max_side: + return width - width % 2, height - height % 2 + scale = max_side / max(width, height) + w = max(2, int(round(width * scale / 2)) * 2) + h = max(2, int(round(height * scale / 2)) * 2) + return w, h + + +class _RawReader: + """Read fixed-size BGR frames from an ffmpeg stdout pipe.""" + + def __init__( + self, + argv: list[str], + width: int, + height: int, + cancel: threading.Event, + pts: queue.Queue[float] | None = None, + ) -> None: + self.size = width * height * 3 + self.shape = (height, width, 3) + self.proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False) + self.cancel = cancel + self.pts = pts + self._stderr: deque[bytes] = deque(maxlen=20) + self._drain = threading.Thread(target=self._read_stderr, daemon=True) + self._drain.start() + + def _read_stderr(self) -> None: + assert self.proc.stderr is not None + for line in self.proc.stderr: + if self.pts is not None and b"Parsed_showinfo" in line: + match = SHOWINFO_PTS.search(line) + if match: + self.pts.put(float(match.group(1))) + continue + self._stderr.append(line) + + def frames(self) -> Iterator[np.ndarray]: + assert self.proc.stdout is not None + try: + while True: + if self.cancel.is_set(): + raise Cancelled() + buf = self.proc.stdout.read(self.size) + if not buf: + break + if len(buf) != self.size: + raise MediaError("internal", "truncated frame from decoder") + yield np.frombuffer(buf, np.uint8).reshape(self.shape) + code = self.proc.wait() + self._drain.join(timeout=5) + if code != 0: + raise MediaError("internal", "decoder failed: " + self.stderr_tail()) + finally: + self.close() + + def stderr_tail(self) -> str: + return b"".join(self._stderr).decode("utf-8", "replace").strip()[-500:] + + def close(self) -> None: + if self.proc.poll() is None: + self.proc.kill() + self.proc.wait() + + +def stream_frames( + path: str, + info: VideoInfo, + cancel: threading.Event, + sample_fps: float | None = None, + max_side: int | None = None, +) -> Iterator[tuple[int, np.ndarray]]: + """Yield (tMs, frame) for every frame (CFR-normalized) or about `sample_fps` per second. + + Sampled frames are real source frames; their timestamps come from + ffmpeg's showinfo, so tMs is exact even for variable frame rates. + """ + width, height = scaled_size(info.width, info.height, max_side) + filters: list[str] = [] + if sample_fps: + # The source frame nearest to each k / sample_fps slot (phase-locked, no drift). + half = 0.5 / info.fps + slot = f"floor((t+{half:.6f})*{sample_fps:.6f})" + prev = f"floor((prev_selected_t+{half:.6f})*{sample_fps:.6f})" + filters += [f"select='isnan(prev_selected_t)+gt({slot}\\,{prev})'", "showinfo"] + if (width, height) != (info.width, info.height): + filters.append(f"scale={width}:{height}:flags=area") + level = "info" if sample_fps else "error" + argv = [FFMPEG, "-v", level, "-nostats", "-nostdin", "-i", path, "-map", "0:v:0", "-an", "-sn"] + if filters: + argv += ["-vf", ",".join(filters)] + if sample_fps: + argv += ["-fps_mode", "passthrough"] + else: + argv += ["-fps_mode", "cfr", "-r", info.fps_text] + argv += ["-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"] + pts: queue.Queue[float] | None = queue.Queue() if sample_fps else None + reader = _RawReader(argv, width, height, cancel, pts) + for index, frame in enumerate(reader.frames()): + if pts is None: + yield int(round(index * 1000.0 / info.fps)), frame + continue + try: + t = pts.get(timeout=30) + except queue.Empty as error: + raise MediaError("internal", "missing frame timestamp from decoder") from error + yield max(int(round((t - info.start_s) * 1000)), 0), frame + + +def frame_at( + path: str, info: VideoInfo, t_ms: int, cancel: threading.Event, max_side: int | None = None +) -> np.ndarray: + """Decode the frame shown at t_ms (accurate seek).""" + width, height = scaled_size(info.width, info.height, max_side) + t = min(max(t_ms / 1000.0, 0.0), max(info.duration_s - 0.5 / info.fps, 0.0)) + argv = [ + FFMPEG, + "-v", + "error", + "-nostdin", + "-ss", + f"{t:.3f}", + "-i", + path, + "-map", + "0:v:0", + "-frames:v", + "1", + ] + if (width, height) != (info.width, info.height): + argv += ["-vf", f"scale={width}:{height}:flags=area"] + argv += ["-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"] + frames = list(_RawReader(argv, width, height, cancel).frames()) + if not frames: + raise MediaError("internal", f"no frame decoded at {t_ms} ms") + return frames[0] + + +def frames_at( + path: str, info: VideoInfo, times_ms: list[int], cancel: threading.Event, max_side: int | None = None +) -> Iterator[tuple[int, np.ndarray]]: + for t_ms in times_ms: + if cancel.is_set(): + raise Cancelled() + yield t_ms, frame_at(path, info, t_ms, cancel, max_side) + + +def temp_sibling(final_path: str) -> str: + directory, name = os.path.split(final_path) + stem, ext = os.path.splitext(name) + return os.path.join(directory, f".{stem}.partial-{uuid.uuid4().hex}{ext}") + + +def encoder_args(out_path: str) -> list[str]: + """Near-lossless intermediate: FFV1 in Matroska, otherwise x264 CRF 12.""" + if out_path.lower().endswith(".mkv"): + return ["-c:v", "ffv1", "-level", "3", "-pix_fmt", "yuv444p", "-g", "1"] + return ["-c:v", "libx264", "-preset", "medium", "-crf", "12", "-pix_fmt", "yuv420p"] + + +FrameFn = Callable[[int, int, np.ndarray], np.ndarray] + + +def transform_video( + path: str, + out_path: str, + info: VideoInfo, + cancel: threading.Event, + frame_fn: FrameFn, +) -> int: + """Decode every frame, apply frame_fn(index, tMs, frame), encode, copy audio. + + Fails closed: any decode, transform, or encode error removes the temporary + output and raises; out_path only ever appears complete (atomic rename). + """ + if not info.has_video: + raise BadRequest("input has no video stream") + width, height = scaled_size(info.width, info.height, None) + tmp = temp_sibling(out_path) + enc_argv = [ + FFMPEG, + "-v", + "error", + "-nostdin", + "-y", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s", + f"{width}x{height}", + "-r", + info.fps_text, + "-i", + "pipe:0", + "-i", + path, + "-map", + "0:v:0", + "-map", + "1:a?", + "-map_metadata", + "-1", + "-map_chapters", + "-1", + *encoder_args(out_path), + "-c:a", + "copy", + tmp, + ] + dec_argv = [ + FFMPEG, + "-v", + "error", + "-xerror", + "-nostdin", + "-i", + path, + "-map", + "0:v:0", + "-an", + "-sn", + "-fps_mode", + "cfr", + "-r", + info.fps_text, + ] + if (width, height) != (info.width, info.height): + dec_argv += ["-vf", f"crop={width}:{height}:0:0"] + dec_argv += ["-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"] + encoder = subprocess.Popen( + enc_argv, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, shell=False + ) + enc_err: list[bytes] = [] + drain = threading.Thread(target=lambda: enc_err.extend(encoder.stderr or []), daemon=True) + drain.start() + reader = _RawReader(dec_argv, width, height, cancel) + count = 0 + ok = False + try: + for index, frame in enumerate(reader.frames()): + t_ms = int(round(index * 1000.0 / info.fps)) + out = frame_fn(index, t_ms, frame.copy()) + if out.shape != frame.shape or out.dtype != np.uint8: + raise MediaError("internal", "frame transform changed the frame shape") + assert encoder.stdin is not None + try: + encoder.stdin.write(np.ascontiguousarray(out).tobytes()) + except BrokenPipeError as error: + raise MediaError("internal", "encoder stopped accepting frames") from error + count += 1 + expected = (info.video_duration_s or info.duration_s) * info.fps + if count == 0 or (expected > 0 and count < expected * 0.98 - 2): + raise MediaError("internal", f"decoded {count} of ~{int(expected)} frames") + assert encoder.stdin is not None + encoder.stdin.close() + if encoder.wait() != 0: + drain.join(timeout=5) + raise MediaError( + "internal", "encoder failed: " + b"".join(enc_err[-5:]).decode("utf-8", "replace")[-500:] + ) + drain.join(timeout=5) + os.replace(tmp, out_path) + ok = True + return count + finally: + reader.close() + if not ok: + if encoder.poll() is None: + encoder.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + encoder.wait(timeout=10) + if os.path.exists(tmp): + os.unlink(tmp) + + +_NUMBER = r"(-?(?:inf|\d+(?:\.\d+)?))" + + +def parse_ebur128(stderr: str) -> tuple[float | None, float | None]: + """Integrated loudness and true peak from the ebur128 summary block.""" + summary = stderr[stderr.rfind("Summary:") :] if "Summary:" in stderr else "" + lufs = re.search(r"I:\s+" + _NUMBER + r"\s+LUFS", summary) + peak = re.search(r"Peak:\s+" + _NUMBER + r"\s+dBFS", summary) + + def num(match: re.Match[str] | None) -> float | None: + if not match or "inf" in match.group(1): + return None + return float(match.group(1)) + + return num(lufs), num(peak) diff --git a/infra/media-worker/app/models.py b/infra/media-worker/app/models.py new file mode 100644 index 0000000..3ad8e0c --- /dev/null +++ b/infra/media-worker/app/models.py @@ -0,0 +1,506 @@ +"""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// 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) diff --git a/infra/media-worker/app/ops/__init__.py b/infra/media-worker/app/ops/__init__.py new file mode 100644 index 0000000..64d5a30 --- /dev/null +++ b/infra/media-worker/app/ops/__init__.py @@ -0,0 +1 @@ +"""One module per operation family. Each exposes `run(ctx, body, cancel)`.""" diff --git a/infra/media-worker/app/ops/audio.py b/infra/media-worker/app/ops/audio.py new file mode 100644 index 0000000..b79d917 --- /dev/null +++ b/infra/media-worker/app/ops/audio.py @@ -0,0 +1,163 @@ +"""audioFeatures without ML: ffmpeg ebur128 + silencedetect, librosa beat tracking. + +Audio is decoded once to mono 22.05 kHz float32 and processed in five-minute +chunks, so memory stays bounded and the request can stop between chunks. +""" + +from __future__ import annotations + +import re +import subprocess +import threading +from typing import Any + +import numpy as np + +from ..errors import Cancelled, MediaError +from ..media_io import FFMPEG, parse_ebur128, run_cancellable +from ..schemas import PathRequest +from .common import Ctx, open_input + +SAMPLE_RATE = 22050 +CHUNK_SECONDS = 300 +ENERGY_HOP_MS = 250 +SILENCE_DB = -40 +SILENCE_MIN_S = 0.5 + + +def loudness_and_silence( + path: str, cancel: threading.Event +) -> tuple[float | None, float | None, list[dict[str, int]]]: + result = run_cancellable( + [ + FFMPEG, + "-hide_banner", + "-nostats", + "-nostdin", + "-i", + path, + "-map", + "0:a:0", + "-vn", + "-af", + f"ebur128=peak=true,silencedetect=noise={SILENCE_DB}dB:d={SILENCE_MIN_S}", + "-f", + "null", + "-", + ], + cancel, + ) + stderr = result.stderr.decode("utf-8", "replace") + if result.returncode != 0: + raise MediaError("internal", "loudness measurement failed") + lufs, peak = parse_ebur128(stderr) + silences: list[dict[str, int]] = [] + start: float | None = None + for line in stderr.splitlines(): + if (m := re.search(r"silence_start:\s*(-?[\d.]+)", line)) is not None: + start = max(float(m.group(1)), 0.0) + elif (m := re.search(r"silence_end:\s*([\d.]+)", line)) is not None and start is not None: + silences.append( + {"startMs": int(round(start * 1000)), "endMs": int(round(float(m.group(1)) * 1000))} + ) + start = None + if start is not None: + silences.append({"startMs": int(round(start * 1000)), "endMs": -1}) + return lufs, peak, silences + + +def decode_chunks(path: str, cancel: threading.Event) -> Any: + argv = [ + FFMPEG, + "-v", + "error", + "-nostdin", + "-i", + path, + "-map", + "0:a:0", + "-vn", + "-ac", + "1", + "-ar", + str(SAMPLE_RATE), + "-f", + "f32le", + "pipe:1", + ] + proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, shell=False) + chunk_bytes = SAMPLE_RATE * CHUNK_SECONDS * 4 + assert proc.stdout is not None + try: + while True: + if cancel.is_set(): + raise Cancelled() + buf = proc.stdout.read(chunk_bytes) + if not buf: + break + usable = len(buf) - len(buf) % 4 + yield np.frombuffer(buf[:usable], dtype=" dict[str, Any]: + import librosa # deferred: heavy import (numba), only needed here + + real, info = open_input(ctx, body.path, cancel, need_video=False) + ctx.registry.runtime("audioFeatures") + if not info.has_audio: + return { + "integratedLufs": None, + "truePeakDb": None, + "tempoBpm": None, + "beatsMs": [], + "energy": [], + "silences": [], + } + lufs, peak, silences = loudness_and_silence(real, cancel) + + hop = int(SAMPLE_RATE * ENERGY_HOP_MS / 1000) + beats: list[int] = [] + tempos: list[tuple[float, float]] = [] + rms_values: list[float] = [] + offset_s = 0.0 + for samples in decode_chunks(real, cancel): + length_s = samples.size / SAMPLE_RATE + if samples.size >= SAMPLE_RATE * 2 and float(np.abs(samples).max()) > 1e-4: + tempo, frames = librosa.beat.beat_track(y=samples, sr=SAMPLE_RATE, units="frames") + times = librosa.frames_to_time(frames, sr=SAMPLE_RATE) + beats.extend(int(round((offset_s + t) * 1000)) for t in times) + value = float(np.atleast_1d(tempo)[0]) + if value > 0: + tempos.append((value, length_s)) + for i in range(0, samples.size, hop): + window = samples[i : i + hop] + rms_values.append(float(np.sqrt(np.mean(window.astype(np.float64) ** 2))) if window.size else 0.0) + offset_s += length_s + + duration_ms = int(round(offset_s * 1000)) + for silence in silences: + if silence["endMs"] < 0: + silence["endMs"] = duration_ms + peak_rms = max(rms_values) if rms_values else 0.0 + energy = [ + {"tMs": i * ENERGY_HOP_MS, "value": (v / peak_rms) if peak_rms > 0 else 0.0} + for i, v in enumerate(rms_values) + ] + tempo_bpm = None + if tempos: + weights = np.array([w for _, w in tempos]) + tempo_bpm = float(np.average([t for t, _ in tempos], weights=weights)) + return { + "integratedLufs": lufs, + "truePeakDb": peak, + "tempoBpm": tempo_bpm, + "beatsMs": sorted(set(beats)), + "energy": energy, + "silences": silences, + } diff --git a/infra/media-worker/app/ops/common.py b/infra/media-worker/app/ops/common.py new file mode 100644 index 0000000..84fe793 --- /dev/null +++ b/infra/media-worker/app/ops/common.py @@ -0,0 +1,62 @@ +"""Helpers shared by the operation modules.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .. import media_io, paths +from ..config import Settings +from ..errors import BadRequest +from ..media_io import VideoInfo +from ..models import Registry + + +@dataclass +class Ctx: + settings: Settings + registry: Registry + + +def open_input( + ctx: Ctx, path: str, cancel: threading.Event, need_video: bool = True +) -> tuple[str, VideoInfo]: + """Resolve, size-check, probe, and duration-check an input file.""" + real = paths.input_file(ctx.settings.media_root, path, ctx.settings.max_input_bytes) + info = media_io.probe(real, cancel) + media_io.check_duration(info, ctx.settings.max_seconds) + if need_video and not info.has_video: + raise BadRequest("input has no video stream") + return real, info + + +def clamp01(value: float) -> float: + return float(min(max(value, 0.0), 1.0)) + + +def norm_box(x: float, y: float, w: float, h: float, width: int, height: int) -> dict[str, float]: + """Pixel box -> normalized box clipped to the frame.""" + x0, y0 = clamp01(x / width), clamp01(y / height) + x1, y1 = clamp01((x + w) / width), clamp01((y + h) / height) + return {"x": x0, "y": y0, "w": max(x1 - x0, 0.0), "h": max(y1 - y0, 0.0)} + + +def norm_point(x: float, y: float, width: int, height: int) -> dict[str, float]: + return {"x": clamp01(x / width), "y": clamp01(y / height)} + + +def unique_sorted(times: list[int]) -> list[int]: + return sorted(set(int(t) for t in times)) + + +def jsonable(value: Any) -> Any: + if isinstance(value, dict): + return {k: jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [jsonable(v) for v in value] + if isinstance(value, np.generic): + return value.item() + return value diff --git a/infra/media-worker/app/ops/devices.py b/infra/media-worker/app/ops/devices.py new file mode 100644 index 0000000..69054a9 --- /dev/null +++ b/infra/media-worker/app/ops/devices.py @@ -0,0 +1,61 @@ +"""Accelerator detection. ML libraries are imported only when installed.""" + +from __future__ import annotations + +import functools +import importlib.util +from typing import Any + + +def _installed(name: str) -> bool: + try: + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False + + +@functools.cache +def torch_gpu_available() -> bool: + if not _installed("torch"): + return False + import torch + + return bool(torch.cuda.is_available()) + + +def effective_accelerator(settings: Any) -> str: + """What the image was built for, downgraded to cpu when no GPU is visible. + + ROCm builds of torch expose the GPU through the torch.cuda API. + """ + if settings.accelerator in ("cuda", "rocm") and torch_gpu_available(): + return settings.accelerator + return "cpu" + + +def torch_device(settings: Any) -> str: + return "cuda" if effective_accelerator(settings) != "cpu" else "cpu" + + +def ctranslate2_device(settings: Any) -> str: + # CTranslate2 ships CUDA builds only; ROCm hosts transcribe on the CPU. + if settings.accelerator != "cuda": + return "cpu" + import ctranslate2 + + return "cuda" if ctranslate2.get_cuda_device_count() > 0 else "cpu" + + +def onnx_providers(settings: Any) -> list[str]: + import onnxruntime + + available = set(onnxruntime.get_available_providers()) + providers: list[str] = [] + if settings.accelerator == "cuda" and "CUDAExecutionProvider" in available: + providers.append("CUDAExecutionProvider") + if settings.accelerator == "rocm": + for name in ("MIGraphXExecutionProvider", "ROCMExecutionProvider"): + if name in available: + providers.append(name) + providers.append("CPUExecutionProvider") + return providers diff --git a/infra/media-worker/app/ops/embeddings.py b/infra/media-worker/app/ops/embeddings.py new file mode 100644 index 0000000..60f5124 --- /dev/null +++ b/infra/media-worker/app/ops/embeddings.py @@ -0,0 +1,147 @@ +"""embedFrames and classifyContent with SigLIP (Apache-2.0) on ONNX Runtime. + +classifyContent is zero-shot: each category is a set of text prompts and the +score is the highest sigmoid probability among them. It is reported as +experimental; the gate thresholds decide what the scores mean. +""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io +from ..schemas import ClassifyContentRequest, TimesRequest +from .common import Ctx, clamp01, open_input, unique_sorted +from .devices import onnx_providers + +SIGLIP_SIZE = 224 +SIGLIP_TEXT_LENGTH = 64 +SIGLIP_PAD_ID = 1 +BATCH = 16 + +CATEGORY_PROMPTS: dict[str, tuple[str, ...]] = { + "explicit_nudity": ("a photo of a naked person", "a photo containing explicit nudity"), + "suggestive": ("a photo of a person in lingerie", "a sexually suggestive photo"), + "violence": ("a photo of a violent fight", "a photo showing blood and injuries"), + "minors_risk": ("a photo of a child", "a photo of a young teenager"), + "self_harm": ("a photo of self-harm", "a photo of cuts on a wrist"), + "drugs": ("a photo of illegal drugs", "a photo of drug paraphernalia"), + "weapons": ("a photo of a gun", "a photo of a person holding a knife as a weapon"), +} + + +def preprocess(frames: list[np.ndarray]) -> np.ndarray: + """BGR frames -> SigLIP pixel_values (N, 3, 224, 224), normalized to [-1, 1].""" + batch = [] + for frame in frames: + rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + resized = cv2.resize(rgb, (SIGLIP_SIZE, SIGLIP_SIZE), interpolation=cv2.INTER_CUBIC) + batch.append(resized.astype(np.float32) / 127.5 - 1.0) + return np.ascontiguousarray(np.stack(batch).transpose(0, 3, 1, 2)) + + +def _session(path: str, settings: Any) -> Any: + import onnxruntime + + return onnxruntime.InferenceSession(path, providers=onnx_providers(settings)) + + +class SiglipEmbedder: + model = "siglip-base-patch16-224" + + def __init__(self, path: str, settings: Any) -> None: + self.session = _session(path, settings) + names = [o.name for o in self.session.get_outputs()] + self.output = "pooler_output" if "pooler_output" in names else names[-1] + self.dim = 768 + + def embed(self, frames: list[np.ndarray]) -> np.ndarray: + (vectors,) = self.session.run([self.output], {"pixel_values": preprocess(frames)}) + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + return vectors / np.maximum(norms, 1e-12) + + +def load_frame_embedder(files: dict[str, str], settings: Any) -> SiglipEmbedder: + return SiglipEmbedder(files["vision_model_quantized.onnx"], settings) + + +class SiglipZeroShot: + model = "siglip-base-patch16-224-zero-shot" + + def __init__(self, model_path: str, tokenizer_path: str, settings: Any) -> None: + from tokenizers import Tokenizer + + tokenizer = Tokenizer.from_file(tokenizer_path) + tokenizer.enable_truncation(SIGLIP_TEXT_LENGTH) + tokenizer.enable_padding(length=SIGLIP_TEXT_LENGTH, pad_id=SIGLIP_PAD_ID) + self.labels: list[str] = [] + prompts: list[str] = [] + for category, texts in CATEGORY_PROMPTS.items(): + for text in texts: + self.labels.append(category) + prompts.append(text.lower()) + self.input_ids = np.array([e.ids for e in tokenizer.encode_batch(prompts)], dtype=np.int64) + self.session = _session(model_path, settings) + + def classify(self, frames: list[np.ndarray]) -> list[dict[str, float]]: + (logits,) = self.session.run( + ["logits_per_image"], {"input_ids": self.input_ids, "pixel_values": preprocess(frames)} + ) + probs = 1.0 / (1.0 + np.exp(-logits)) + result = [] + for row in probs: + scores: dict[str, float] = {} + for label, value in zip(self.labels, row, strict=True): + scores[label] = max(scores.get(label, 0.0), clamp01(float(value))) + result.append(scores) + return result + + +def load_content_classifier(files: dict[str, str], settings: Any) -> SiglipZeroShot: + return SiglipZeroShot(files["model_quantized.onnx"], files["tokenizer.json"], settings) + + +def _batched(frames: Any, size: int) -> Any: + times: list[int] = [] + batch: list[np.ndarray] = [] + for t_ms, frame in frames: + times.append(t_ms) + batch.append(frame) + if len(batch) == size: + yield times, batch + times, batch = [], [] + if batch: + yield times, batch + + +def embed_run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + embedder = ctx.registry.runtime("embedFrames") + vectors: list[list[float]] = [] + # Vectors follow the request order, duplicates included. + for _, batch in _batched(media_io.frames_at(real, info, list(body.timesMs), cancel, 512), BATCH): + vectors.extend([float(v) for v in row] for row in embedder.embed(batch)) + return {"model": embedder.model, "dim": int(embedder.dim), "vectors": vectors} + + +def classify_run(ctx: Ctx, body: ClassifyContentRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + classifier = ctx.registry.runtime("classifyContent") + results: dict[int, dict[str, float]] = {} + + def score(frames: Any) -> None: + for times, batch in _batched(frames, BATCH): + for t_ms, scores in zip(times, classifier.classify(batch), strict=True): + results[t_ms] = scores + + score(media_io.stream_frames(real, info, cancel, body.sampleFps, 512)) + extra = [t for t in unique_sorted(body.extraTimesMs or []) if t not in results] + score(media_io.frames_at(real, info, extra, cancel, 512)) + return { + "model": classifier.model, + "frames": [{"tMs": t, "scores": results[t]} for t in sorted(results)], + } diff --git a/infra/media-worker/app/ops/faces.py b/infra/media-worker/app/ops/faces.py new file mode 100644 index 0000000..167fae2 --- /dev/null +++ b/infra/media-worker/app/ops/faces.py @@ -0,0 +1,265 @@ +"""detectFaces, trackFaces, faceLandmarks. + +Detection is OpenCV YuNet (MIT). Tracking is a ByteTrack-style IoU tracker +implemented here: high-confidence detections associate first, low-confidence +ones only extend existing tracks, and a track ends after 500 ms without a match. +Identity embeddings (for local person clustering) come from an opt-in backend +and are null otherwise. +""" + +from __future__ import annotations + +import math +import threading +from dataclasses import dataclass, field +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io +from ..schemas import DetectFacesRequest, TimesRequest, TrackFacesRequest +from .common import Ctx, clamp01, norm_box, norm_point, open_input, unique_sorted + +DETECT_MAX_SIDE = 1280 +DETECT_THRESHOLD = 0.6 +TRACK_LOW_THRESHOLD = 0.3 +TRACK_MATCH_IOU = 0.3 +TRACK_MAX_GAP_MS = 500 + + +@dataclass +class Detection: + """Pixel-space face: box (x, y, w, h), score, five landmarks (x, y).""" + + box: tuple[float, float, float, float] + score: float + landmarks: list[tuple[float, float]] + + +class YuNetDetector: + def __init__(self, model_path: str) -> None: + self.model_path = model_path + self._local = threading.local() + + def _detector(self, width: int, height: int) -> Any: + detector = getattr(self._local, "detector", None) + if detector is None: + detector = cv2.FaceDetectorYN.create( + self.model_path, "", (width, height), TRACK_LOW_THRESHOLD, 0.3, 5000 + ) + self._local.detector = detector + detector.setInputSize((width, height)) + return detector + + def detect(self, frame: np.ndarray) -> list[Detection]: + height, width = frame.shape[:2] + _, faces = self._detector(width, height).detect(frame) + result: list[Detection] = [] + for row in faces if faces is not None else []: + points = [(float(row[4 + 2 * i]), float(row[5 + 2 * i])) for i in range(5)] + result.append( + Detection( + (float(row[0]), float(row[1]), float(row[2]), float(row[3])), float(row[14]), points + ) + ) + return result + + +def load_face_detector(files: dict[str, str], _settings: Any) -> YuNetDetector: + return YuNetDetector(files["face_detection_yunet_2023mar.onnx"]) + + +class SFaceEmbedder: + def __init__(self, model_path: str) -> None: + self.model_path = model_path + self._local = threading.local() + + def embed(self, frame: np.ndarray, detection: Detection) -> np.ndarray: + recognizer = getattr(self._local, "recognizer", None) + if recognizer is None: + recognizer = cv2.FaceRecognizerSF.create(self.model_path, "") + self._local.recognizer = recognizer + row = np.array( + [*detection.box, *[c for p in detection.landmarks for c in p], detection.score], + dtype=np.float32, + ) + aligned = recognizer.alignCrop(frame, row) + return recognizer.feature(aligned).reshape(-1).astype(np.float32) + + +def load_identity_embedder(files: dict[str, str], _settings: Any) -> SFaceEmbedder: + return SFaceEmbedder(files["face_recognition_sface_2021dec.onnx"]) + + +def _observation(det: Detection, width: int, height: int) -> dict[str, Any]: + return { + "box": norm_box(*det.box, width, height), + "confidence": clamp01(det.score), + "landmarks": [norm_point(x, y, width, height) for x, y in det.landmarks], + } + + +def detect_run(ctx: Ctx, body: DetectFacesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("detectFaces") + frames = ( + media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE) + if body.timesMs is not None + else media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE) + ) + out = [] + for t_ms, frame in frames: + height, width = frame.shape[:2] + faces = [ + _observation(d, width, height) for d in detector.detect(frame) if d.score >= DETECT_THRESHOLD + ] + out.append({"tMs": t_ms, "faces": faces}) + return {"frames": out} + + +def landmarks_run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("faceLandmarks") + out = [] + for t_ms, frame in media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE): + height, width = frame.shape[:2] + faces = [ + { + "box": norm_box(*d.box, width, height), + "landmarks": [norm_point(x, y, width, height) for x, y in d.landmarks], + } + for d in detector.detect(frame) + if d.score >= DETECT_THRESHOLD + ] + out.append({"tMs": t_ms, "faces": faces}) + return {"frames": out} + + +def iou(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float: + ax1, ay1, ax2, ay2 = a[0], a[1], a[0] + a[2], a[1] + a[3] + bx1, by1, bx2, by2 = b[0], b[1], b[0] + b[2], b[1] + b[3] + iw = max(0.0, min(ax2, bx2) - max(ax1, bx1)) + ih = max(0.0, min(ay2, by2) - max(ay1, by1)) + inter = iw * ih + union = a[2] * a[3] + b[2] * b[3] - inter + return inter / union if union > 0 else 0.0 + + +@dataclass +class Track: + track_id: str + box: tuple[float, float, float, float] + velocity: tuple[float, float] = (0.0, 0.0) + last_ms: int = 0 + keyframes: list[dict[str, Any]] = field(default_factory=list) + # Best (score, frame crop source, detection) for identity embeddings. + best: list[tuple[float, np.ndarray, Detection]] = field(default_factory=list) + + def predicted(self, t_ms: int) -> tuple[float, float, float, float]: + dt = t_ms - self.last_ms + return ( + self.box[0] + self.velocity[0] * dt, + self.box[1] + self.velocity[1] * dt, + self.box[2], + self.box[3], + ) + + +class IouTracker: + """Two-stage greedy IoU association (ByteTrack's idea without the Kalman filter).""" + + def __init__(self, max_gap_ms: int = TRACK_MAX_GAP_MS) -> None: + self.max_gap_ms = max_gap_ms + self.active: list[Track] = [] + self.finished: list[Track] = [] + self._next = 1 + + def _associate( + self, tracks: list[Track], dets: list[Detection], t_ms: int + ) -> tuple[list[tuple[Track, Detection]], list[Track], list[Detection]]: + pairs = sorted( + ( + (iou(tr.predicted(t_ms), d.box), ti, di) + for ti, tr in enumerate(tracks) + for di, d in enumerate(dets) + ), + reverse=True, + ) + used_t: set[int] = set() + used_d: set[int] = set() + matches: list[tuple[Track, Detection]] = [] + for score, ti, di in pairs: + if score < TRACK_MATCH_IOU: + break + if ti in used_t or di in used_d: + continue + used_t.add(ti) + used_d.add(di) + matches.append((tracks[ti], dets[di])) + rest_t = [t for i, t in enumerate(tracks) if i not in used_t] + rest_d = [d for i, d in enumerate(dets) if i not in used_d] + return matches, rest_t, rest_d + + def update( + self, t_ms: int, dets: list[Detection], frame: np.ndarray | None, keep_crops: bool + ) -> list[tuple[Track, Detection]]: + expired = [t for t in self.active if t_ms - t.last_ms > self.max_gap_ms] + self.finished.extend(expired) + self.active = [t for t in self.active if t not in expired] + high = [d for d in dets if d.score >= DETECT_THRESHOLD] + low = [d for d in dets if TRACK_LOW_THRESHOLD <= d.score < DETECT_THRESHOLD] + matches, rest_tracks, rest_high = self._associate(self.active, high, t_ms) + low_matches, _, _ = self._associate(rest_tracks, low, t_ms) + matches += low_matches + for track, det in matches: + dt = max(t_ms - track.last_ms, 1) + track.velocity = ((det.box[0] - track.box[0]) / dt, (det.box[1] - track.box[1]) / dt) + track.box = det.box + track.last_ms = t_ms + for det in rest_high: + track = Track(f"t{self._next}", det.box, last_ms=t_ms) + self._next += 1 + self.active.append(track) + matches.append((track, det)) + if keep_crops and frame is not None: + for track, det in matches: + if det.score >= DETECT_THRESHOLD: + track.best.append((det.score, frame, det)) + track.best.sort(key=lambda item: -item[0]) + del track.best[3:] + return matches + + def close(self) -> list[Track]: + tracks = self.finished + self.active + self.active, self.finished = [], [] + return sorted(tracks, key=lambda t: int(t.track_id[1:])) + + +def track_run(ctx: Ctx, body: TrackFacesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("trackFaces") + embedder = ( + ctx.registry.optional_runtime("faceIdentity") + if ctx.registry.optional_enabled("faceIdentity") + else None + ) + tracker = IouTracker() + for t_ms, frame in media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE): + height, width = frame.shape[:2] + matches = tracker.update(t_ms, detector.detect(frame), frame, embedder is not None) + for track, det in matches: + track.keyframes.append( + {"tMs": t_ms, "box": norm_box(*det.box, width, height), "confidence": clamp01(det.score)} + ) + tracks = [] + for track in tracker.close(): + embedding = None + if embedder is not None and track.best: + vectors = np.stack([embedder.embed(frame, det) for _, frame, det in track.best]) + mean = vectors.mean(axis=0) + norm = float(np.linalg.norm(mean)) + if norm > 0 and math.isfinite(norm): + embedding = [float(v) for v in mean / norm] + tracks.append({"trackId": track.track_id, "keyframes": track.keyframes, "embedding": embedding}) + return {"sampleFps": body.sampleFps, "tracks": tracks} diff --git a/infra/media-worker/app/ops/inpaint.py b/infra/media-worker/app/ops/inpaint.py new file mode 100644 index 0000000..314085c --- /dev/null +++ b/infra/media-worker/app/ops/inpaint.py @@ -0,0 +1,163 @@ +"""inpaintVideo: LaMa (Apache-2.0, ONNX) per frame with optical-flow smoothing. + +Per frame: crop around the mask, inpaint at 512x512, paste back through a +feathered mask. When the previous frame was inpainted too, its result is +warped into the current frame with DIS optical flow and blended in, which +damps flicker. Quality is measured on the output: +- temporalStability: 1 - mean absolute change inside the mask between + consecutive frames after motion compensation (scaled so 32 levels -> 0). +- edgeQuality: 1 - excess gradient on the mask boundary over the gradient + just outside it. +- score: 0.6 * temporalStability + 0.4 * edgeQuality. +Callers fall back to cover/reframe when the score is low. +""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io, paths +from ..errors import BadRequest +from ..schemas import InpaintVideoRequest +from .common import Ctx, clamp01, open_input +from .devices import onnx_providers +from .masks import MaskIndex +from .redact import OUTPUT_EXTENSIONS + +LAMA_SIZE = 512 +TEMPORAL_BLEND = 0.35 + + +class LamaInpainter: + def __init__(self, path: str, settings: Any) -> None: + import onnxruntime + + self.session = onnxruntime.InferenceSession(path, providers=onnx_providers(settings)) + + def inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray: + """image: BGR uint8 512x512, mask: uint8 0/255 512x512 -> BGR uint8.""" + rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 + pixel = rgb.transpose(2, 0, 1)[None] + m = (mask > 127).astype(np.float32)[None, None] + (out,) = self.session.run(["output"], {"image": pixel, "mask": m}) + result = np.clip(out[0].transpose(1, 2, 0), 0, 255).astype(np.uint8) + return cv2.cvtColor(result, cv2.COLOR_RGB2BGR) + + +def load_inpainter(files: dict[str, str], settings: Any) -> LamaInpainter: + return LamaInpainter(files["lama_fp32.onnx"], settings) + + +def _crop_box(mask: np.ndarray) -> tuple[int, int, int, int]: + height, width = mask.shape + ys, xs = np.nonzero(mask) + cx, cy = (xs.min() + xs.max()) / 2, (ys.min() + ys.max()) / 2 + side = max(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1) * 2.2 + side = int(min(max(side, 256), max(width, height))) + x0 = int(min(max(cx - side / 2, 0), max(width - side, 0))) + y0 = int(min(max(cy - side / 2, 0), max(height - side, 0))) + return x0, y0, min(x0 + side, width), min(y0 + side, height) + + +class Smoother: + """Temporal blending and quality bookkeeping across frames.""" + + def __init__(self) -> None: + self.flow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_FAST) + self.prev_gray: np.ndarray | None = None + self.prev_out: np.ndarray | None = None + self.prev_masked = False + self.stability: list[float] = [] + self.edges: list[float] = [] + + def _warp_prev(self, gray: np.ndarray) -> np.ndarray | None: + if self.prev_gray is None or self.prev_out is None: + return None + flow = self.flow.calc(gray, self.prev_gray, None) + h, w = gray.shape + grid_x, grid_y = np.meshgrid(np.arange(w, dtype=np.float32), np.arange(h, dtype=np.float32)) + return cv2.remap( + self.prev_out, + grid_x + flow[..., 0], + grid_y + flow[..., 1], + cv2.INTER_LINEAR, + borderMode=cv2.BORDER_REPLICATE, + ) + + def step(self, source: np.ndarray, inpainted: np.ndarray | None, mask: np.ndarray | None) -> np.ndarray: + gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY) + out = source + if inpainted is not None and mask is not None: + soft = cv2.GaussianBlur(mask.astype(np.float32) / 255.0, (0, 0), 3) + soft = np.maximum(soft, mask.astype(np.float32) / 255.0)[..., None] + warped = self._warp_prev(gray) if self.prev_masked else None + fill = inpainted.astype(np.float32) + if warped is not None: + fill = fill * (1 - TEMPORAL_BLEND) + warped.astype(np.float32) * TEMPORAL_BLEND + out = np.clip(source.astype(np.float32) * (1 - soft) + fill * soft, 0, 255).astype(np.uint8) + inside = mask > 0 + if warped is not None: + # Raw LaMa output against the motion-compensated previous frame: flicker before smoothing. + diff = np.abs(inpainted.astype(np.float32) - warped.astype(np.float32))[inside].mean() + self.stability.append(clamp01(1.0 - float(diff) / 32.0)) + self.edges.append(_edge_quality(out, mask)) + self.prev_gray = gray + self.prev_out = out + self.prev_masked = mask is not None + return out + + def quality(self) -> dict[str, float]: + temporal = float(np.mean(self.stability)) if self.stability else 1.0 + edge = float(np.mean(self.edges)) if self.edges else 1.0 + return { + "temporalStability": clamp01(temporal), + "edgeQuality": clamp01(edge), + "score": clamp01(0.6 * temporal + 0.4 * edge), + } + + +def _edge_quality(frame: np.ndarray, mask: np.ndarray) -> float: + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float32) + grad = cv2.magnitude(cv2.Sobel(gray, cv2.CV_32F, 1, 0), cv2.Sobel(gray, cv2.CV_32F, 0, 1)) + kernel = np.ones((5, 5), np.uint8) + boundary = (cv2.dilate(mask, kernel) > 0) & (cv2.erode(mask, kernel) == 0) + outer = (cv2.dilate(mask, np.ones((15, 15), np.uint8)) > 0) & ~(cv2.dilate(mask, kernel) > 0) + if not boundary.any() or not outer.any(): + return 1.0 + excess = max(0.0, float(grad[boundary].mean()) - float(grad[outer].mean())) + return clamp01(1.0 - excess / (float(grad[outer].mean()) + 40.0)) + + +def run(ctx: Ctx, body: InpaintVideoRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + masks = MaskIndex(paths.input_dir(ctx.settings.media_root, body.maskDir), info.frame_ms()) + out_real = paths.output_file(ctx.settings.media_root, body.outPath) + if os.path.splitext(out_real)[1].lower() not in OUTPUT_EXTENSIONS: + raise BadRequest("outPath must end in .mkv, .mp4 or .mov") + if out_real == real: + raise BadRequest("outPath must differ from path") + inpainter = ctx.registry.runtime("inpaintVideo") + smoother = Smoother() + + def per_frame(_index: int, t_ms: int, frame: np.ndarray) -> np.ndarray: + height, width = frame.shape[:2] + mask = masks.lookup(t_ms, width, height) + if mask is None or not mask.any(): + return smoother.step(frame, None, None) + # Grow slightly so the seam lands on untouched pixels. + mask = cv2.dilate(mask, np.ones((7, 7), np.uint8)) + x0, y0, x1, y1 = _crop_box(mask) + crop = cv2.resize(frame[y0:y1, x0:x1], (LAMA_SIZE, LAMA_SIZE), interpolation=cv2.INTER_AREA) + crop_mask = cv2.resize(mask[y0:y1, x0:x1], (LAMA_SIZE, LAMA_SIZE), interpolation=cv2.INTER_NEAREST) + filled = inpainter.inpaint(crop, crop_mask) + full = frame.copy() + full[y0:y1, x0:x1] = cv2.resize(filled, (x1 - x0, y1 - y0), interpolation=cv2.INTER_CUBIC) + return smoother.step(frame, full, mask) + + media_io.transform_video(real, out_real, info, cancel, per_frame) + return {"outPath": body.outPath, "quality": smoother.quality()} diff --git a/infra/media-worker/app/ops/masks.py b/infra/media-worker/app/ops/masks.py new file mode 100644 index 0000000..f4824cf --- /dev/null +++ b/infra/media-worker/app/ops/masks.py @@ -0,0 +1,54 @@ +"""Per-frame mask directories (`.png`), as written by segmentPrompted.""" + +from __future__ import annotations + +import bisect +import os +import re + +import cv2 +import numpy as np + +from ..errors import BadRequest, MediaError + +MASK_NAME = re.compile(r"^(\d+)\.png$") + + +class MaskIndex: + def __init__(self, directory: str, frame_ms: float) -> None: + entries = [] + for name in os.listdir(directory): + match = MASK_NAME.match(name) + if match: + entries.append((int(match.group(1)), os.path.join(directory, name))) + if not entries: + raise BadRequest("mask directory contains no .png masks") + entries.sort() + self.times = [t for t, _ in entries] + self.files = [p for _, p in entries] + gaps = np.diff(self.times) if len(self.times) > 1 else np.array([frame_ms]) + # A mask covers frames up to 1.5 sampling intervals away (at least one frame). + self.tolerance = max(float(np.median(gaps)) * 1.5, frame_ms) + self._cache: tuple[int, np.ndarray] | None = None + + def lookup(self, t_ms: int, width: int, height: int) -> np.ndarray | None: + """Binary mask (uint8 0/255) for the frame at t_ms, resized to the frame, or None.""" + i = bisect.bisect_left(self.times, t_ms) + best = None + for j in (i - 1, i): + if not 0 <= j < len(self.times) or abs(self.times[j] - t_ms) > self.tolerance: + continue + if best is None or abs(self.times[j] - t_ms) < abs(self.times[best] - t_ms): + best = j + if best is None: + return None + if self._cache is not None and self._cache[0] == best: + return self._cache[1] + mask = cv2.imread(self.files[best], cv2.IMREAD_GRAYSCALE) + if mask is None: + raise MediaError("internal", f"unreadable mask {os.path.basename(self.files[best])}") + if mask.shape != (height, width): + mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST) + mask = np.where(mask > 127, 255, 0).astype(np.uint8) + self._cache = (best, mask) + return mask diff --git a/infra/media-worker/app/ops/redact.py b/infra/media-worker/app/ops/redact.py new file mode 100644 index 0000000..053243e --- /dev/null +++ b/infra/media-worker/app/ops/redact.py @@ -0,0 +1,271 @@ +"""redact: frame-accurate face and region redaction (OpenCV + ffmpeg pipes). + +Geometry per instruction and frame: +- Boxes interpolate linearly between keyframes. A single keyframe is a static + region for the whole video; otherwise the region applies from the first to + the last keyframe (plus half a frame either side). +- `paddingPct` grows the box by that percentage of its width/height (half on + each side). Mask regions are dilated by the same share of their extent. +- `featherPct` fades the effect outward over that percentage of the region's + shorter side. The region itself is always fully covered. + +Fails closed: an unreadable frame, a mask that cannot be read, or an encoder +error is a 500 and no file is left at outPath. +""" + +from __future__ import annotations + +import bisect +import os +import threading +from dataclasses import dataclass +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io, paths +from ..errors import BadRequest +from ..schemas import RedactionInstruction, RedactRequest +from .common import Ctx, open_input +from .masks import MaskIndex + +OUTPUT_EXTENSIONS = (".mkv", ".mp4", ".mov") +DEFAULT_FILL = (0, 0, 0) + + +def hex_to_bgr(color: str | None) -> tuple[int, int, int] | None: + if not color: + return None + value = color.lstrip("#") + r, g, b = int(value[0:2], 16), int(value[2:4], 16), int(value[4:6], 16) + return (b, g, r) + + +@dataclass +class Plan: + mode: str + padding: float + feather: float + color: tuple[int, int, int] | None + times: list[int] + boxes: list[tuple[float, float, float, float]] + landmark_times: list[int] + landmarks: list[list[tuple[float, float]]] + masks: MaskIndex | None + + +def build_plan(ctx: Ctx, instruction: RedactionInstruction, frame_ms: float) -> Plan: + keyframes = sorted(instruction.keyframes, key=lambda k: k.tMs) + landmarks = sorted(instruction.landmarks or [], key=lambda frame: frame.tMs) + masks = None + if instruction.maskDir: + masks = MaskIndex(paths.input_dir(ctx.settings.media_root, instruction.maskDir), frame_ms) + return Plan( + mode=instruction.mode, + padding=instruction.paddingPct / 100.0, + feather=instruction.featherPct / 100.0, + color=hex_to_bgr(instruction.color), + times=[k.tMs for k in keyframes], + boxes=[(k.box.x, k.box.y, k.box.w, k.box.h) for k in keyframes], + landmark_times=[frame.tMs for frame in landmarks], + landmarks=[[(p.x, p.y) for p in frame.points] for frame in landmarks], + masks=masks, + ) + + +def interpolate(times: list[int], values: list[Any], t_ms: float, frame_ms: float, hold_single: bool) -> Any: + """Linear interpolation of tuples/lists over keyframe times; None outside the span.""" + if not times: + return None + if len(times) == 1: + if hold_single or abs(t_ms - times[0]) <= frame_ms / 2: + return values[0] + return None + if t_ms < times[0] - frame_ms / 2 or t_ms > times[-1] + frame_ms / 2: + return None + if t_ms <= times[0]: + return values[0] + if t_ms >= times[-1]: + return values[-1] + i = bisect.bisect_right(times, t_ms) + t0, t1 = times[i - 1], times[i] + a, b = values[i - 1], values[i] + if t1 == t0: + return b + f = (t_ms - t0) / (t1 - t0) + if isinstance(a, list): + if len(a) != len(b): + return a if f < 0.5 else b + return [tuple(pa[k] + (pb[k] - pa[k]) * f for k in range(2)) for pa, pb in zip(a, b, strict=True)] + return tuple(a[k] + (b[k] - a[k]) * f for k in range(len(a))) + + +def _pixel_box( + box: tuple[float, float, float, float], padding: float, width: int, height: int +) -> tuple[int, int, int, int]: + x, y, w, h = box + x -= w * padding / 2 + y -= h * padding / 2 + w *= 1 + padding + h *= 1 + padding + x0 = int(np.floor(max(0.0, x) * width)) + y0 = int(np.floor(max(0.0, y) * height)) + x1 = int(np.ceil(min(1.0, x + w) * width)) + y1 = int(np.ceil(min(1.0, y + h) * height)) + return x0, y0, max(x1, x0), max(y1, y0) + + +def blur_effect(roi: np.ndarray, size: int) -> np.ndarray: + """Strong, size-scaled blur: downsample, upsample, Gaussian.""" + h, w = roi.shape[:2] + factor = max(4, size // 6) + small = cv2.resize(roi, (max(1, w // factor), max(1, h // factor)), interpolation=cv2.INTER_AREA) + up = cv2.resize(small, (w, h), interpolation=cv2.INTER_LINEAR) + k = max(3, (size // 3) | 1) + return cv2.GaussianBlur(up, (k, k), 0) + + +def pixelate_effect(roi: np.ndarray, size: int) -> np.ndarray: + h, w = roi.shape[:2] + block = max(4, size // 10) + small = cv2.resize(roi, (max(1, -(-w // block)), max(1, -(-h // block))), interpolation=cv2.INTER_AREA) + return cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST) + + +def _ring_color(frame: np.ndarray, region: np.ndarray, margin: int) -> tuple[int, int, int]: + """Median colour just outside the region: the 'skin tone patch' cover fill.""" + kernel = np.ones((2 * margin + 1, 2 * margin + 1), np.uint8) + ring = (cv2.dilate(region, kernel) > 0) & (region == 0) + if not ring.any(): + return DEFAULT_FILL + b, g, r = np.median(frame[ring], axis=0) + return (int(b), int(g), int(r)) + + +def _region_for( + plan: Plan, t_ms: int, frame_ms: float, width: int, height: int +) -> tuple[np.ndarray, int] | None: + """Binary region mask (uint8 0/255) and its characteristic size in px.""" + region = np.zeros((height, width), np.uint8) + size = 0 + hold = plan.masks is None + box = interpolate(plan.times, plan.boxes, t_ms, frame_ms, hold_single=hold) + if box is not None: + x0, y0, x1, y1 = _pixel_box(box, plan.padding, width, height) + raw_w, raw_h = box[2] * width, box[3] * height + if plan.mode in ("partial_eyes", "partial_upper", "partial_lower"): + _partial_region(plan, region, box, (x0, y0, x1, y1), t_ms, frame_ms, width, height) + else: + region[y0:y1, x0:x1] = 255 + size = max(size, int(min(raw_w, raw_h))) + if plan.masks is not None: + mask = plan.masks.lookup(t_ms, width, height) + if mask is not None: + ys, xs = np.nonzero(mask) + if xs.size: + extent = int(min(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1)) + grow = int(round(extent * plan.padding / 2)) + if grow > 0: + mask = cv2.dilate(mask, np.ones((2 * grow + 1, 2 * grow + 1), np.uint8)) + region |= mask + size = max(size, extent) + if not region.any(): + return None + return region, max(size, 1) + + +def _partial_region( + plan: Plan, + region: np.ndarray, + box: tuple[float, float, float, float], + padded: tuple[int, int, int, int], + t_ms: int, + frame_ms: float, + width: int, + height: int, +) -> None: + x0, y0, x1, y1 = padded + bx, by, bw, bh = box + points = interpolate(plan.landmark_times, plan.landmarks, t_ms, frame_ms, hold_single=False) + # Five-point landmarks: right eye, left eye, nose, mouth corners. + eyes = points[:2] if points and len(points) >= 2 else None + nose_y = points[2][1] if points and len(points) >= 3 else by + bh * 0.55 + if plan.mode == "partial_eyes": + if eyes: + (ex1, ey1), (ex2, ey2) = eyes + else: + ex1, ey1, ex2, ey2 = bx + bw * 0.3, by + bh * 0.4, bx + bw * 0.7, by + bh * 0.4 + cx, cy = (ex1 + ex2) / 2 * width, (ey1 + ey2) / 2 * height + angle = float(np.degrees(np.arctan2((ey2 - ey1) * height, (ex2 - ex1) * width))) + band_w = (x1 - x0) * 1.0 + band_h = max(bh * height * 0.32 * (1 + plan.padding), 2.0) + rect = cv2.boxPoints(((cx, cy), (band_w, band_h), angle)).astype(np.int32) + cv2.fillConvexPoly(region, rect, 255) + elif plan.mode == "partial_upper": + split = int(round(min(max(nose_y, by), by + bh) * height)) + region[y0 : max(split, y0), x0:x1] = 255 + else: + split = int(round(min(max(nose_y, by), by + bh) * height)) + region[min(split, y1) : y1, x0:x1] = 255 + + +def apply_plan(frame: np.ndarray, plan: Plan, t_ms: int, frame_ms: float) -> None: + height, width = frame.shape[:2] + found = _region_for(plan, t_ms, frame_ms, width, height) + if found is None: + return + region, size = found + radius = int(round(plan.feather * size)) + ys, xs = np.nonzero(region) + margin = radius + 2 + x0, x1 = max(0, int(xs.min()) - margin), min(width, int(xs.max()) + margin + 1) + y0, y1 = max(0, int(ys.min()) - margin), min(height, int(ys.max()) + margin + 1) + roi = frame[y0:y1, x0:x1] + roi_region = region[y0:y1, x0:x1] + mode = plan.mode + if mode in ("blur", "partial_upper", "partial_lower"): + effected = blur_effect(roi, size) + elif mode == "pixelate": + effected = pixelate_effect(roi, size) + else: + if mode == "cover" and plan.color is None: + fill = _ring_color(frame, region, max(4, size // 20)) + else: + fill = plan.color or DEFAULT_FILL + effected = np.empty_like(roi) + effected[:] = fill + alpha = roi_region.astype(np.float32) / 255.0 + if radius > 0: + kernel = np.ones((2 * radius + 1, 2 * radius + 1), np.uint8) + grown = cv2.dilate(roi_region, kernel).astype(np.float32) / 255.0 + k = 2 * radius + 1 + soft = cv2.GaussianBlur(grown, (k, k), 0) + alpha = np.maximum(alpha, soft) + alpha = alpha[..., None] + blended = roi.astype(np.float32) * (1.0 - alpha) + effected.astype(np.float32) * alpha + roi[:] = np.clip(np.rint(blended), 0, 255).astype(np.uint8) + + +def apply_redactions(frame: np.ndarray, t_ms: int, plans: list[Plan], frame_ms: float) -> np.ndarray: + for plan in plans: + apply_plan(frame, plan, t_ms, frame_ms) + return frame + + +def run(ctx: Ctx, body: RedactRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + ctx.registry.runtime("redact") + out_real = paths.output_file(ctx.settings.media_root, body.outPath) + if os.path.splitext(out_real)[1].lower() not in OUTPUT_EXTENSIONS: + raise BadRequest("outPath must end in .mkv, .mp4 or .mov") + if out_real == real: + raise BadRequest("outPath must differ from path") + frame_ms = info.frame_ms() + plans = [build_plan(ctx, r, frame_ms) for r in body.redactions] + + def per_frame(_index: int, t_ms: int, frame: np.ndarray) -> np.ndarray: + return apply_redactions(frame, t_ms, plans, frame_ms) + + count = media_io.transform_video(real, out_real, info, cancel, per_frame) + return {"outPath": body.outPath, "framesProcessed": count} diff --git a/infra/media-worker/app/ops/saliency.py b/infra/media-worker/app/ops/saliency.py new file mode 100644 index 0000000..5e02e8b --- /dev/null +++ b/infra/media-worker/app/ops/saliency.py @@ -0,0 +1,58 @@ +"""saliency: spectral residual (Hou & Zhang, 2007), implemented with numpy/OpenCV. + +The point is the saliency-weighted centroid of the top 5 % most salient +pixels; strength is how much those pixels stand out against the frame mean. +""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io +from ..schemas import TimesRequest +from .common import Ctx, clamp01, open_input, unique_sorted + +WORK_SIZE = 64 + + +def saliency_map(frame: np.ndarray) -> np.ndarray: + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float32) + small = cv2.resize(gray, (WORK_SIZE, WORK_SIZE), interpolation=cv2.INTER_AREA) + spectrum = np.fft.fft2(small) + log_amp = np.log(np.abs(spectrum) + 1e-8) + phase = np.angle(spectrum) + residual = log_amp - cv2.blur(log_amp, (3, 3)) + sal = np.abs(np.fft.ifft2(np.exp(residual + 1j * phase))) ** 2 + sal = cv2.GaussianBlur(sal.astype(np.float32), (0, 0), 2.5) + peak = float(sal.max()) + return sal / peak if peak > 0 else sal + + +def salient_point(frame: np.ndarray) -> tuple[float, float, float]: + sal = saliency_map(frame) + threshold = float(np.quantile(sal, 0.95)) + top = np.where(sal >= threshold, sal, 0.0) + total = float(top.sum()) + if total <= 0: + return 0.5, 0.5, 0.0 + ys, xs = np.mgrid[0 : sal.shape[0], 0 : sal.shape[1]] + cx = float((xs * top).sum() / total + 0.5) / sal.shape[1] + cy = float((ys * top).sum() / total + 0.5) / sal.shape[0] + mean = float(sal.mean()) + top_mean = float(sal[sal >= threshold].mean()) + strength = (top_mean - mean) / top_mean if top_mean > 0 else 0.0 + return clamp01(cx), clamp01(cy), clamp01(strength) + + +def run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + ctx.registry.runtime("saliency") + points = [] + for t_ms, frame in media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, 512): + cx, cy, strength = salient_point(frame) + points.append({"tMs": t_ms, "cx": cx, "cy": cy, "strength": strength}) + return {"points": points} diff --git a/infra/media-worker/app/ops/scenes.py b/infra/media-worker/app/ops/scenes.py new file mode 100644 index 0000000..a0d2fb6 --- /dev/null +++ b/infra/media-worker/app/ops/scenes.py @@ -0,0 +1,103 @@ +"""detectScenes: PySceneDetect (adaptive) for shots, histogram grouping for scenes.""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np +from scenedetect import AdaptiveDetector, SceneManager, StatsManager, open_video + +from .. import media_io +from ..errors import Cancelled +from ..schemas import PathRequest +from .common import Ctx, clamp01, open_input + +# Neighbouring shots whose middle frames correlate at least this much share a scene. +SCENE_HIST_CORRELATION = 0.6 +ADAPTIVE_THRESHOLD = 3.0 + + +class SceneDetector: + def detect(self, path: str, cancel: threading.Event) -> list[tuple[float, float, float]]: + """Return shots as (start_s, end_s, score in [0,1]).""" + video = open_video(path, backend="opencv") + stats = StatsManager() + manager = SceneManager(stats_manager=stats) + manager.add_detector(AdaptiveDetector(adaptive_threshold=ADAPTIVE_THRESHOLD)) + stop = threading.Event() + + def watch() -> None: + while not stop.wait(0.2): + if cancel.is_set(): + manager.stop() + return + + watcher = threading.Thread(target=watch, daemon=True) + watcher.start() + try: + manager.detect_scenes(video=video) + finally: + stop.set() + if cancel.is_set(): + raise Cancelled() + scenes = manager.get_scene_list(start_in_scene=True) + ratio_key = next( + (k for k in AdaptiveDetector().get_metrics() if k.startswith("adaptive_ratio")), None + ) + shots: list[tuple[float, float, float]] = [] + for index, (start, end) in enumerate(scenes): + score = 1.0 + if index > 0 and ratio_key: + value = ( + stats.get_metrics(start.frame_num, [ratio_key])[0] + if stats.metrics_exist(start.frame_num, [ratio_key]) + else None + ) + score = clamp01((value or ADAPTIVE_THRESHOLD) / (2 * ADAPTIVE_THRESHOLD)) + shots.append((start.seconds, end.seconds, score)) + return shots + + +def load_scene_detector(_files: dict[str, str], _settings: Any) -> SceneDetector: + return SceneDetector() + + +def _histogram(frame: np.ndarray) -> np.ndarray: + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + hist = cv2.calcHist([hsv], [0, 1], None, [32, 16], [0, 180, 0, 256]) + cv2.normalize(hist, hist) + return hist + + +def run(ctx: Ctx, body: PathRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("detectScenes") + raw = detector.detect(real, cancel) + duration_ms = int(round(info.duration_s * 1000)) + if not raw: + raw = [(0.0, info.duration_s, 1.0)] + shots: list[dict[str, Any]] = [] + for start_s, end_s, score in raw: + start_ms = int(round(start_s * 1000)) + end_ms = max(int(round(end_s * 1000)), start_ms) + if duration_ms: + end_ms = min(end_ms, duration_ms) + shots.append({"startMs": start_ms, "endMs": end_ms, "score": float(score)}) + + scenes: list[dict[str, int]] = [] + previous: np.ndarray | None = None + for shot in shots: + middle = (shot["startMs"] + shot["endMs"]) // 2 + hist = _histogram(media_io.frame_at(real, info, middle, cancel, max_side=320)) + same = ( + previous is not None + and cv2.compareHist(previous, hist, cv2.HISTCMP_CORREL) >= SCENE_HIST_CORRELATION + ) + if same and scenes: + scenes[-1]["endMs"] = shot["endMs"] + else: + scenes.append({"startMs": shot["startMs"], "endMs": shot["endMs"]}) + previous = hist + return {"shots": shots, "scenes": scenes} diff --git a/infra/media-worker/app/ops/segment.py b/infra/media-worker/app/ops/segment.py new file mode 100644 index 0000000..87334bd --- /dev/null +++ b/infra/media-worker/app/ops/segment.py @@ -0,0 +1,244 @@ +"""segmentPrompted: SAM2 (Apache-2.0) from a click or box, propagated over the video. + +Frames are sampled at up to 15 fps and fed to the SAM2 video predictor in +chunks so memory stays bounded; each chunk is re-prompted with the bounding +box of the previous chunk's last mask. Masks are written as `.png` +(0 or 255) at the source resolution. +""" + +from __future__ import annotations + +import os +import tempfile +import threading +from collections.abc import Iterator +from typing import Any + +import cv2 +import numpy as np + +from .. import paths +from ..errors import BadRequest, Cancelled, MediaError +from ..media_io import FFMPEG, run_cancellable, scaled_size +from ..schemas import SegmentPromptedRequest +from .common import Ctx, clamp01, norm_box, open_input +from .devices import torch_device +from .masks import MASK_NAME + +SEGMENT_MAX_FPS = 15.0 +SEGMENT_MAX_SIDE = 1024 +CHUNK_FRAMES = 64 + +MaskStream = Iterator[tuple[int, np.ndarray, float]] + + +def _mask_box(mask: np.ndarray) -> tuple[int, int, int, int] | None: + ys, xs = np.nonzero(mask) + if xs.size == 0: + return None + return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1 + + +class Sam2Segmenter: + def __init__(self, checkpoint: str, settings: Any) -> None: + import torch + from sam2.build_sam import build_sam2_video_predictor + + self.torch = torch + self.device = torch_device(settings) + self.predictor = build_sam2_video_predictor( + "configs/sam2.1/sam2.1_hiera_t.yaml", checkpoint, device=self.device + ) + self._lock = threading.Lock() + + def _chunk( + self, + frames: list[str], + prompt_at: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + reverse: bool, + cancel: threading.Event, + ) -> MaskStream: + with tempfile.TemporaryDirectory(prefix="sam2-") as chunk_dir: + for i, source in enumerate(frames): + os.symlink(source, os.path.join(chunk_dir, f"{i}.jpg")) + torch = self.torch + with torch.inference_mode(): + state = self.predictor.init_state(video_path=chunk_dir, offload_video_to_cpu=True) + kwargs: dict[str, Any] = {} + if box is not None: + kwargs["box"] = np.array(box, dtype=np.float32) + if point is not None: + kwargs["points"] = np.array([point], dtype=np.float32) + kwargs["labels"] = np.array([1], dtype=np.int32) + self.predictor.add_new_points_or_box(state, frame_idx=prompt_at, obj_id=1, **kwargs) + for index, _, logits in self.predictor.propagate_in_video( + state, start_frame_idx=prompt_at, reverse=reverse + ): + if cancel.is_set(): + raise Cancelled() + values = logits[0, 0].float().cpu().numpy() + mask = values > 0 + confidence = float(1.0 / (1.0 + np.exp(-values[mask])).mean()) if mask.any() else 0.0 + yield index, mask, confidence + self.predictor.reset_state(state) + + def segment( + self, + frames: list[str], + prompt_index: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + propagate: bool, + cancel: threading.Event, + ) -> MaskStream: + with self._lock: + if not propagate: + for _, mask, confidence in self._chunk([frames[prompt_index]], 0, box, point, False, cancel): + yield prompt_index, mask, confidence + return + yield from self._forward(frames, prompt_index, box, point, cancel) + yield from self._backward(frames, prompt_index, box, point, cancel) + + def _forward( + self, + frames: list[str], + anchor: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + cancel: threading.Event, + ) -> MaskStream: + emit_anchor = True + while True: + chunk = frames[anchor : anchor + CHUNK_FRAMES] + last: np.ndarray | None = None + for index, mask, confidence in self._chunk(chunk, 0, box, point, False, cancel): + if index == 0 and not emit_anchor: + continue + last = mask + yield anchor + index, mask, confidence + end = anchor + len(chunk) - 1 + box = _mask_box(last) if last is not None else None + if end >= len(frames) - 1 or box is None: + return + anchor, point, emit_anchor = end, None, False + + def _backward( + self, + frames: list[str], + anchor: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + cancel: threading.Event, + ) -> MaskStream: + while anchor > 0: + start = max(0, anchor - CHUNK_FRAMES + 1) + last: np.ndarray | None = None + for index, mask, confidence in self._chunk( + frames[start : anchor + 1], anchor - start, box, point, True, cancel + ): + if start + index == anchor: + continue # emitted by the forward pass + last = mask + yield start + index, mask, confidence + box = _mask_box(last) if last is not None else None + if start == 0 or box is None: + return + anchor, point = start, None + + +def load_segmenter(files: dict[str, str], settings: Any) -> Sam2Segmenter: + return Sam2Segmenter(files["sam2.1_hiera_tiny.pt"], settings) + + +def _extract_frames( + path: str, + fps: float, + size: tuple[int, int], + directory: str, + cancel: threading.Event, + at_ms: int | None = None, +) -> list[str]: + argv = [FFMPEG, "-v", "error", "-nostdin"] + if at_ms is not None: + argv += ["-ss", f"{at_ms / 1000:.3f}"] + argv += ["-i", path, "-map", "0:v:0", "-an", "-sn"] + filters = [] if at_ms is not None else [f"fps={fps:.6f}"] + filters.append(f"scale={size[0]}:{size[1]}:flags=area") + argv += ["-vf", ",".join(filters)] + if at_ms is not None: + argv += ["-frames:v", "1"] + argv += ["-q:v", "2", "-start_number", "0", os.path.join(directory, "%06d.jpg")] + result = run_cancellable(argv, cancel) + if result.returncode != 0: + raise MediaError("internal", "frame extraction failed") + names = sorted(n for n in os.listdir(directory) if n.endswith(".jpg")) + if not names: + raise MediaError("internal", "no frames extracted") + return [os.path.join(directory, n) for n in names] + + +def run(ctx: Ctx, body: SegmentPromptedRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + out_dir = paths.output_dir(ctx.settings.media_root, body.outDir) + segmenter = ctx.registry.runtime("segmentPrompted") + fps = min(info.fps, SEGMENT_MAX_FPS) + width, height = scaled_size(info.width, info.height, SEGMENT_MAX_SIDE) + prompt = body.prompt + if prompt.tMs > int(info.duration_s * 1000) + 1: + raise BadRequest("prompt.tMs is after the end of the video") + with tempfile.TemporaryDirectory(prefix="segment-") as work: + if body.propagate: + frames = _extract_frames(real, fps, (width, height), work, cancel) + times = [int(round(i * 1000.0 / fps)) for i in range(len(frames))] + prompt_index = int(np.argmin([abs(t - prompt.tMs) for t in times])) + else: + frames = _extract_frames(real, fps, (width, height), work, cancel, at_ms=prompt.tMs) + times = [prompt.tMs] + prompt_index = 0 + box = None + if prompt.box is not None: + b = prompt.box + box = ( + int(round(b.x * width)), + int(round(b.y * height)), + int(round(clamp01(b.x + b.w) * width)), + int(round(clamp01(b.y + b.h) * height)), + ) + point = ( + (int(round(prompt.point.x * width)), int(round(prompt.point.y * height))) + if prompt.point + else None + ) + + # Replace masks from an earlier run in the same directory. + for name in os.listdir(out_dir): + if MASK_NAME.match(name): + os.unlink(os.path.join(out_dir, name)) + keyframes = [] + for index, mask, confidence in segmenter.segment( + frames, prompt_index, box, point, body.propagate, cancel + ): + if cancel.is_set(): + raise Cancelled() + t_ms = times[index] + full = cv2.resize( + mask.astype(np.uint8) * 255, (info.width, info.height), interpolation=cv2.INTER_NEAREST + ) + found = _mask_box(full > 0) + if found is None: + continue + if not cv2.imwrite(os.path.join(out_dir, f"{t_ms}.png"), full): + raise MediaError("internal", "could not write mask") + x0, y0, x1, y1 = found + keyframes.append( + { + "tMs": t_ms, + "box": norm_box(x0, y0, x1 - x0, y1 - y0, info.width, info.height), + "confidence": clamp01(confidence), + } + ) + keyframes.sort(key=lambda k: k["tMs"]) + score = float(np.mean([k["confidence"] for k in keyframes])) if keyframes else 0.0 + return {"maskDir": body.outDir, "keyframes": keyframes, "score": score} diff --git a/infra/media-worker/app/ops/synthesize.py b/infra/media-worker/app/ops/synthesize.py new file mode 100644 index 0000000..81244d5 --- /dev/null +++ b/infra/media-worker/app/ops/synthesize.py @@ -0,0 +1,20 @@ +"""synthesizeIdentity: deliberately unavailable. + +No generator + face-swap pipeline with a commercial licence exists that this +sidecar can ship, so the operation always answers 501. The request schema +forbids unknown fields, so a target-face image is rejected before this runs. +""" + +from __future__ import annotations + +import threading +from typing import Any + +from ..errors import Unsupported +from ..schemas import SynthesizeIdentityRequest +from .common import Ctx + + +def run(ctx: Ctx, _body: SynthesizeIdentityRequest, _cancel: threading.Event) -> dict[str, Any]: + spec = ctx.registry.backend("synthesizeIdentity").spec + raise Unsupported(spec.unavailable_reason or "no commercially licensed model") diff --git a/infra/media-worker/app/ops/transcribe.py b/infra/media-worker/app/ops/transcribe.py new file mode 100644 index 0000000..b1fbc63 --- /dev/null +++ b/infra/media-worker/app/ops/transcribe.py @@ -0,0 +1,57 @@ +"""transcribe: faster-whisper with word timestamps.""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +from ..errors import BadRequest, Cancelled +from ..schemas import TranscribeRequest +from .common import Ctx, clamp01, open_input +from .devices import ctranslate2_device + + +class WhisperTranscriber: + def __init__(self, model_dir: str, settings: Any) -> None: + from faster_whisper import WhisperModel + + device = ctranslate2_device(settings) + self.model = WhisperModel( + model_dir, + device=device, + compute_type="float16" if device == "cuda" else "int8", + ) + + def transcribe( + self, path: str, language: str | None, cancel: threading.Event + ) -> tuple[str | None, list[dict[str, Any]]]: + segments, info = self.model.transcribe(path, language=language, word_timestamps=True, vad_filter=True) + words: list[dict[str, Any]] = [] + for segment in segments: + if cancel.is_set(): + raise Cancelled() + for word in segment.words or []: + start = int(round(word.start * 1000)) + words.append( + { + "startMs": start, + "endMs": max(int(round(word.end * 1000)), start), + "text": word.word.strip(), + "confidence": clamp01(word.probability), + } + ) + return info.language, words + + +def load_transcriber(files: dict[str, str], settings: Any) -> WhisperTranscriber: + return WhisperTranscriber(os.path.dirname(files["model.bin"]), settings) + + +def run(ctx: Ctx, body: TranscribeRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel, need_video=False) + if not info.has_audio: + raise BadRequest("input has no audio stream") + transcriber = ctx.registry.runtime("transcribe") + language, words = transcriber.transcribe(real, body.language, cancel) + return {"language": language, "words": [w for w in words if w["text"]]} diff --git a/infra/media-worker/app/paths.py b/infra/media-worker/app/paths.py new file mode 100644 index 0000000..40ff80e --- /dev/null +++ b/infra/media-worker/app/paths.py @@ -0,0 +1,75 @@ +"""MEDIA_ROOT enforcement. Every path a request names goes through here. + +Inputs and outputs must be absolute and stay inside the real MEDIA_ROOT after +symlinks are resolved. A path that does not exist yet is checked through its +nearest existing ancestor (realpath resolves the existing prefix). +""" + +from __future__ import annotations + +import os +import stat + +from .errors import InvalidPath, MediaError + + +def _inside(root: str, candidate: str) -> bool: + return candidate == root or candidate.startswith(root.rstrip(os.sep) + os.sep) + + +def resolve_inside(root: str, requested: str) -> str: + """Return the real absolute location of `requested`, or raise InvalidPath.""" + if not isinstance(requested, str) or not requested or "\0" in requested: + raise InvalidPath() + if not os.path.isabs(requested): + raise InvalidPath("path must be absolute") + real_root = os.path.realpath(root) + lexical = os.path.normpath(requested) + # Reject lexical traversal even when it would land back inside the root. + if ".." in requested.split("/"): + raise InvalidPath() + real = os.path.realpath(lexical) + if not _inside(real_root, real): + raise InvalidPath() + return real + + +def input_file(root: str, requested: str, max_bytes: int | None = None) -> str: + real = resolve_inside(root, requested) + try: + info = os.stat(real) + except FileNotFoundError as error: + raise InvalidPath("input file does not exist") from error + if not stat.S_ISREG(info.st_mode): + raise InvalidPath("input is not a regular file") + if max_bytes is not None and info.st_size > max_bytes: + raise MediaError("too_large", "input exceeds MEDIA_MAX_INPUT_BYTES") + return real + + +def input_dir(root: str, requested: str) -> str: + real = resolve_inside(root, requested) + if not os.path.isdir(real): + raise InvalidPath("directory does not exist") + return real + + +def output_file(root: str, requested: str) -> str: + """Resolve an output file path and create its parent directories inside the root.""" + real = resolve_inside(root, requested) + if os.path.isdir(real): + raise InvalidPath("output path is a directory") + parent = os.path.dirname(real) + os.makedirs(parent, exist_ok=True) + # The parent may have been swapped for a symlink in between; check again. + resolve_inside(root, parent) + return real + + +def output_dir(root: str, requested: str) -> str: + real = resolve_inside(root, requested) + os.makedirs(real, exist_ok=True) + resolve_inside(root, real) + if not os.path.isdir(real): + raise InvalidPath("output directory is not a directory") + return real diff --git a/infra/media-worker/app/runtime.py b/infra/media-worker/app/runtime.py new file mode 100644 index 0000000..e5c2fd2 --- /dev/null +++ b/infra/media-worker/app/runtime.py @@ -0,0 +1,93 @@ +"""Request execution: concurrency gates, cancellation on disconnect, time limit.""" + +from __future__ import annotations + +import asyncio +import threading +import time +from collections.abc import Callable +from typing import Any, Protocol + +from .errors import Busy, Cancelled + +POLL_SECONDS = 0.25 + + +class Gate: + """`concurrency` jobs run; up to `queue_size` more wait; anything beyond is busy. + + Thread-based so it works no matter which event loop serves the request. + """ + + def __init__(self, concurrency: int, queue_size: int) -> None: + self.concurrency = concurrency + self.capacity = concurrency + queue_size + self._slots = threading.Semaphore(concurrency) + self._lock = threading.Lock() + self.inflight = 0 + + def reserve(self) -> None: + with self._lock: + if self.inflight >= self.capacity: + raise Busy() + self.inflight += 1 + + def release_reservation(self) -> None: + with self._lock: + self.inflight -= 1 + + def run(self, fn: Callable[[threading.Event], Any], cancel: threading.Event) -> Any: + """Wait for a slot (reservation already taken), run fn, free both.""" + try: + while not self._slots.acquire(timeout=POLL_SECONDS): + if cancel.is_set(): + raise Cancelled() + try: + if cancel.is_set(): + raise Cancelled() + return fn(cancel) + finally: + self._slots.release() + finally: + self.release_reservation() + + +class DisconnectProbe(Protocol): + async def is_disconnected(self) -> bool: ... + + +async def run_job( + request: DisconnectProbe, + gate: Gate, + fn: Callable[[threading.Event], Any], + max_seconds: float, +) -> Any: + """Run fn in a worker thread. Sets the cancel event when the client + disconnects or the deadline passes; fn checks it between frames/chunks.""" + gate.reserve() + cancel = threading.Event() + task = asyncio.ensure_future(asyncio.to_thread(gate.run, fn, cancel)) + deadline = time.monotonic() + max_seconds + reason: str | None = None + try: + while True: + done, _ = await asyncio.wait({task}, timeout=POLL_SECONDS) + if done: + break + if reason is None: + if time.monotonic() > deadline: + reason = "exceeded MEDIA_MAX_SECONDS" + cancel.set() + elif await request.is_disconnected(): + reason = "client disconnected" + cancel.set() + except asyncio.CancelledError: + cancel.set() + raise + try: + result = task.result() + except Cancelled as error: + raise Cancelled(reason or error.reason) from error + if reason is not None: + raise Cancelled(reason) + return result diff --git a/infra/media-worker/app/schemas.py b/infra/media-worker/app/schemas.py new file mode 100644 index 0000000..e6b0973 --- /dev/null +++ b/infra/media-worker/app/schemas.py @@ -0,0 +1,124 @@ +"""Request bodies. Field names mirror the TypeScript request objects (camelCase). + +Every model forbids unknown fields; that is what keeps a target-face image +out of synthesizeIdentity. +""" + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field, StrictBool, model_validator + +Ms = Annotated[int, Field(ge=0, le=7 * 24 * 3600 * 1000)] +Unit = Annotated[float, Field(ge=0.0, le=1.0)] +Path = Annotated[str, Field(min_length=1, max_length=4096)] +SampleFps = Annotated[float, Field(gt=0.0, le=60.0)] +TimesMs = Annotated[list[Ms], Field(max_length=100_000)] + + +class Strict(BaseModel): + model_config = ConfigDict(extra="forbid", strict=True) + + +class NormalizedBox(Strict): + x: Unit + y: Unit + w: Unit + h: Unit + + +class NormalizedPoint(Strict): + x: Unit + y: Unit + + +class BoxKeyframe(Strict): + tMs: Ms + box: NormalizedBox + + +class PathRequest(Strict): + path: Path + + +class TranscribeRequest(PathRequest): + language: Annotated[str, Field(min_length=2, max_length=16, pattern=r"^[A-Za-z-]+$")] | None = None + + +class DetectFacesRequest(PathRequest): + sampleFps: SampleFps + timesMs: TimesMs | None = None + + +class TrackFacesRequest(PathRequest): + sampleFps: SampleFps + + +class TimesRequest(PathRequest): + timesMs: TimesMs + + +class ClassifyContentRequest(PathRequest): + sampleFps: SampleFps + extraTimesMs: TimesMs | None = None + + +class MarkPrompt(Strict): + tMs: Ms + box: NormalizedBox | None = None + point: NormalizedPoint | None = None + + @model_validator(mode="after") + def needs_box_or_point(self) -> MarkPrompt: + if self.box is None and self.point is None: + raise ValueError("prompt needs a box or a point") + return self + + +class SegmentPromptedRequest(PathRequest): + prompt: MarkPrompt + propagate: StrictBool + outDir: Path + + +class InpaintVideoRequest(PathRequest): + maskDir: Path + outPath: Path + + +RedactionMode = Literal[ + "blur", "pixelate", "solid", "sticker", "partial_eyes", "partial_upper", "partial_lower", "cover" +] + + +class LandmarkFrame(Strict): + tMs: Ms + points: Annotated[list[NormalizedPoint], Field(max_length=512)] + + +class RedactionInstruction(Strict): + keyframes: Annotated[list[BoxKeyframe], Field(max_length=200_000)] + mode: RedactionMode + paddingPct: Annotated[float, Field(ge=0.0, le=100.0)] + featherPct: Annotated[float, Field(ge=0.0, le=50.0)] + color: Annotated[str, Field(pattern=r"^#[0-9a-fA-F]{6}$")] | None = None + landmarks: Annotated[list[LandmarkFrame], Field(max_length=200_000)] | None = None + maskDir: Path | None = None + + @model_validator(mode="after") + def needs_region(self) -> RedactionInstruction: + if not self.keyframes and not self.maskDir: + raise ValueError("redaction needs keyframes or a maskDir") + return self + + +class RedactRequest(PathRequest): + redactions: Annotated[list[RedactionInstruction], Field(min_length=1, max_length=1000)] + outPath: Path + + +class SynthesizeIdentityRequest(PathRequest): + track: Annotated[list[BoxKeyframe], Field(min_length=1, max_length=200_000)] + seed: Annotated[int, Field(ge=-(2**53 - 1), le=2**53 - 1)] + outPath: Path diff --git a/infra/media-worker/pytest.ini b/infra/media-worker/pytest.ini new file mode 100644 index 0000000..c87f771 --- /dev/null +++ b/infra/media-worker/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = tests +addopts = -p no:cacheprovider +filterwarnings = + ignore:Using `httpx` with `starlette.testclient` is deprecated diff --git a/infra/media-worker/requirements-dev.txt b/infra/media-worker/requirements-dev.txt new file mode 100644 index 0000000..7684ca1 --- /dev/null +++ b/infra/media-worker/requirements-dev.txt @@ -0,0 +1,3 @@ +pytest==9.1.1 +httpx==0.28.1 +ruff==0.16.10 diff --git a/infra/media-worker/requirements-ml-cpu.txt b/infra/media-worker/requirements-ml-cpu.txt new file mode 100644 index 0000000..68371cc --- /dev/null +++ b/infra/media-worker/requirements-ml-cpu.txt @@ -0,0 +1,9 @@ +--extra-index-url https://download.pytorch.org/whl/cpu +torch==2.13.0+cpu +torchvision==0.28.0+cpu +onnxruntime==1.30.0 +faster-whisper==1.2.1 +ctranslate2==4.8.2 +tokenizers==0.23.2 +hydra-core==1.3.7 +iopath==0.1.10 diff --git a/infra/media-worker/requirements-ml-cuda.txt b/infra/media-worker/requirements-ml-cuda.txt new file mode 100644 index 0000000..b723cc5 --- /dev/null +++ b/infra/media-worker/requirements-ml-cuda.txt @@ -0,0 +1,9 @@ +--extra-index-url https://download.pytorch.org/whl/cu126 +torch==2.13.0+cu126 +torchvision==0.28.0+cu126 +onnxruntime-gpu==1.30.0 +faster-whisper==1.2.1 +ctranslate2==4.8.2 +tokenizers==0.23.2 +hydra-core==1.3.7 +iopath==0.1.10 diff --git a/infra/media-worker/requirements-ml-rocm.txt b/infra/media-worker/requirements-ml-rocm.txt new file mode 100644 index 0000000..a9e87ce --- /dev/null +++ b/infra/media-worker/requirements-ml-rocm.txt @@ -0,0 +1,11 @@ +# ROCm: torch (SAM2) runs on the GPU. ONNX Runtime and CTranslate2 have no ROCm +# wheels on PyPI, so SigLIP, LaMa and Whisper run on the CPU in this build. +--extra-index-url https://download.pytorch.org/whl/rocm7.1 +torch==2.13.0+rocm7.1 +torchvision==0.28.0+rocm7.1 +onnxruntime==1.30.0 +faster-whisper==1.2.1 +ctranslate2==4.8.2 +tokenizers==0.23.2 +hydra-core==1.3.7 +iopath==0.1.10 diff --git a/infra/media-worker/requirements.txt b/infra/media-worker/requirements.txt new file mode 100644 index 0000000..b45dc11 --- /dev/null +++ b/infra/media-worker/requirements.txt @@ -0,0 +1,10 @@ +# Core: HTTP API, ffmpeg-driven operations, OpenCV, scene detection, audio features. +# Accelerator-specific ML wheels live in requirements-ml-.txt. +fastapi==0.142.2 +uvicorn[standard]==0.54.0 +pydantic==2.13.5 +numpy==2.4.6 +opencv-python-headless==4.13.0.92 +scenedetect-headless==0.7.1 +librosa==0.11.0 +soundfile==0.14.0 diff --git a/infra/media-worker/ruff.toml b/infra/media-worker/ruff.toml new file mode 100644 index 0000000..ee25629 --- /dev/null +++ b/infra/media-worker/ruff.toml @@ -0,0 +1,5 @@ +line-length = 110 +target-version = "py311" + +[lint] +select = ["E", "F", "W", "B", "I", "UP", "SIM"] diff --git a/infra/media-worker/tests/__init__.py b/infra/media-worker/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/infra/media-worker/tests/conftest.py b/infra/media-worker/tests/conftest.py new file mode 100644 index 0000000..601b5f6 --- /dev/null +++ b/infra/media-worker/tests/conftest.py @@ -0,0 +1,224 @@ +"""Shared fixtures. All media is synthetic (ffmpeg testsrc/sine, drawn shapes).""" + +from __future__ import annotations + +import subprocess +from collections.abc import Callable +from pathlib import Path + +import cv2 +import numpy as np +import pytest +import soundfile +from fastapi.testclient import TestClient + +from app.config import Settings +from app.main import create_app +from app.models import Backend, Registry, default_backends + +from .fakes import fake_backends + +TOKEN = "test-token-" + "x" * 32 +AUTH = {"Authorization": f"Bearer {TOKEN}"} + + +def ffmpeg(*args: str) -> None: + subprocess.run(["ffmpeg", "-v", "error", "-y", *args], check=True) + + +def write_video(path: Path, frames: list[np.ndarray], fps: int = 25) -> None: + height, width = frames[0].shape[:2] + proc = subprocess.Popen( + [ + "ffmpeg", + "-v", + "error", + "-y", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s", + f"{width}x{height}", + "-r", + str(fps), + "-i", + "pipe:0", + "-c:v", + "libx264", + "-crf", + "10", + "-pix_fmt", + "yuv420p", + str(path), + ], + stdin=subprocess.PIPE, + ) + assert proc.stdin is not None + for frame in frames: + proc.stdin.write(frame.tobytes()) + proc.stdin.close() + assert proc.wait() == 0 + + +def decode(path: Path) -> np.ndarray: + probe = ( + subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height", + "-of", + "csv=p=0", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + .stdout.strip() + .split(",") + ) + width, height = int(probe[0]), int(probe[1]) + raw = subprocess.run( + [ + "ffmpeg", + "-v", + "error", + "-i", + str(path), + "-map", + "0:v:0", + "-fps_mode", + "passthrough", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "pipe:1", + ], + check=True, + capture_output=True, + ).stdout + return np.frombuffer(raw, np.uint8).reshape(-1, height, width, 3) + + +@pytest.fixture(scope="session") +def media_root(tmp_path_factory: pytest.TempPathFactory) -> Path: + root = tmp_path_factory.mktemp("media") + ffmpeg( + "-f", + "lavfi", + "-i", + "testsrc=size=320x240:rate=25:duration=2", + "-f", + "lavfi", + "-i", + "sine=frequency=440:duration=2", + "-c:v", + "libx264", + "-crf", + "10", + "-pix_fmt", + "yuv420p", + "-c:a", + "aac", + "-shortest", + str(root / "testsrc.mp4"), + ) + ffmpeg( + "-f", + "lavfi", + "-i", + "testsrc=size=320x240:rate=25:duration=2", + "-f", + "lavfi", + "-i", + "mandelbrot=size=320x240:rate=25", + "-filter_complex", + "[1:v]trim=duration=2,setpts=PTS-STARTPTS[b];[0:v][b]concat=n=2:v=1:a=0", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + str(root / "cut.mp4"), + ) + # Seeded 4-px block noise with a tone: lots of local detail for blur/pixelate checks. + rng = np.random.default_rng(7) + noise = [ + np.kron(rng.integers(0, 256, (60, 80, 3), dtype=np.uint8), np.ones((4, 4, 1), np.uint8)) + for _ in range(25) + ] + write_video(root / "noise-silent.mp4", noise) + ffmpeg( + "-i", + str(root / "noise-silent.mp4"), + "-f", + "lavfi", + "-i", + "sine=frequency=330:duration=1", + "-c:v", + "copy", + "-c:a", + "aac", + "-shortest", + str(root / "noise.mp4"), + ) + # A white disc moving left to right on black: the synthetic "face"/"mark". + frames = [] + for i in range(40): + frame = np.zeros((240, 320, 3), np.uint8) + cv2.circle(frame, (60 + i * 5, 120), 25, (255, 255, 255), -1) + frames.append(frame) + write_video(root / "disc.mp4", frames) + # A bright square in the right half for saliency. + frame = np.full((240, 320, 3), 40, np.uint8) + cv2.rectangle(frame, (240, 100), (280, 140), (255, 255, 255), -1) + write_video(root / "spot.mp4", [frame] * 10) + # 120 BPM click track (8 s) followed by 1.5 s of silence. + sr = 22050 + clicks = np.zeros(int(sr * 9.5), np.float32) + for beat in range(16): + start = int(beat * 0.5 * sr) + t = np.arange(int(0.03 * sr)) / sr + clicks[start : start + t.size] += 0.8 * np.sin(2 * np.pi * 1000 * t) * np.exp(-t * 80) + soundfile.write(str(root / "clicks.wav"), clicks, sr) + return root + + +def make_settings(media_root: Path, tmp_path: Path, **overrides: object) -> Settings: + values: dict[str, object] = { + "token": TOKEN, + "media_root": str(media_root), + "models_dir": str(tmp_path / "models"), + "allow_model_download": False, + } + values.update(overrides) + return Settings(**values) # type: ignore[arg-type] + + +ClientFactory = Callable[..., TestClient] + + +@pytest.fixture +def make_client(media_root: Path, tmp_path: Path) -> ClientFactory: + def make( + backends: dict[str, Backend] | None = None, fakes: bool = True, **overrides: object + ) -> TestClient: + settings = make_settings(media_root, tmp_path, **overrides) + merged = default_backends() + if fakes: + merged.update(fake_backends()) + merged.update(backends or {}) + return TestClient(create_app(settings, Registry(settings, backends=merged))) + + return make + + +@pytest.fixture +def client(make_client: ClientFactory) -> TestClient: + return make_client() diff --git a/infra/media-worker/tests/fakes.py b/infra/media-worker/tests/fakes.py new file mode 100644 index 0000000..8be7527 --- /dev/null +++ b/infra/media-worker/tests/fakes.py @@ -0,0 +1,133 @@ +"""Deterministic offline stand-ins for the model runtimes. + +They implement the same small runtime interfaces as the real backends, so the +operation code around them (frame IO, tracking, masks, encoding) runs for real. +Synthetic inputs only: bright shapes on dark backgrounds stand in for faces +and marks. +""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np + +from app.models import Backend, BackendSpec +from app.ops.faces import Detection + + +def fake_backend(runtime: Any, model: str = "fake", experimental: bool = False) -> Backend: + return Backend( + BackendSpec( + id=f"fake-{model}", + model=model, + license="MIT", + commercialUse=True, + experimental=experimental, + source=None, + ), + lambda _files, _settings: runtime, + ) + + +def bright_blobs(frame: np.ndarray, threshold: int = 200) -> list[tuple[int, int, int, int]]: + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + _, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY) + contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + boxes = [cv2.boundingRect(c) for c in contours if cv2.contourArea(c) > 20] + return sorted(boxes) + + +class FakeFaceDetector: + """Every bright blob is a 'face'; landmarks sit at fixed box fractions.""" + + def detect(self, frame: np.ndarray) -> list[Detection]: + result = [] + for x, y, w, h in bright_blobs(frame): + marks = [(0.3, 0.4), (0.7, 0.4), (0.5, 0.6), (0.35, 0.8), (0.65, 0.8)] + result.append(Detection((x, y, w, h), 0.9, [(x + fx * w, y + fy * h) for fx, fy in marks])) + return result + + +class FakeEmbedder: + model = "fake-embed" + dim = 4 + + def embed(self, frames: list[np.ndarray]) -> np.ndarray: + vectors = np.array([[*frame.reshape(-1, 3).mean(axis=0), 1.0] for frame in frames], dtype=np.float32) + return vectors / np.linalg.norm(vectors, axis=1, keepdims=True) + + +class FakeClassifier: + model = "fake-classifier" + + def classify(self, frames: list[np.ndarray]) -> list[dict[str, float]]: + return [{"violence": float(frame[..., 2].mean() / 255.0), "weapons": 0.0} for frame in frames] + + +class FakeTranscriber: + def transcribe( + self, _path: str, language: str | None, _cancel: threading.Event + ) -> tuple[str | None, list[dict[str, Any]]]: + words = [ + {"startMs": 0, "endMs": 400, "text": "synthetic", "confidence": 0.9}, + {"startMs": 450, "endMs": 900, "text": "tone", "confidence": 0.8}, + ] + return language or "en", words + + +class FakeSegmenter: + """Segments the bright blob in each sampled frame.""" + + def segment( + self, + frames: list[str], + prompt_index: int, + _box: Any, + _point: Any, + propagate: bool, + cancel: threading.Event, + ) -> Any: + indices = range(len(frames)) if propagate else [prompt_index] + for index in indices: + if cancel.is_set(): + return + image = cv2.imread(frames[index if propagate else 0]) + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + yield index, gray > 200, 0.95 + + +class FakeInpainter: + """Classical OpenCV inpainting (Telea) instead of LaMa.""" + + def inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray: + return cv2.inpaint(image, (mask > 127).astype(np.uint8), 5, cv2.INPAINT_TELEA) + + +class BlockingEmbedder(FakeEmbedder): + """Holds the GPU slot until released, to exercise queue overflow.""" + + def __init__(self) -> None: + self.entered = threading.Event() + self.release = threading.Event() + + def embed(self, frames: list[np.ndarray]) -> np.ndarray: + self.entered.set() + self.release.wait(10) + return super().embed(frames) + + +def fake_backends() -> dict[str, Backend]: + detector = FakeFaceDetector() + return { + "transcribe": fake_backend(FakeTranscriber(), "fake-whisper"), + "detectFaces": fake_backend(detector, "fake-faces"), + "trackFaces": fake_backend(detector, "fake-faces"), + "faceLandmarks": fake_backend(detector, "fake-faces"), + "embedFrames": fake_backend(FakeEmbedder(), "fake-embed"), + "classifyContent": fake_backend(FakeClassifier(), "fake-classifier", experimental=True), + "segmentPrompted": fake_backend(FakeSegmenter(), "fake-segment", experimental=True), + "inpaintVideo": fake_backend(FakeInpainter(), "fake-inpaint", experimental=True), + } diff --git a/infra/media-worker/tests/fixtures/media-capabilities.json b/infra/media-worker/tests/fixtures/media-capabilities.json new file mode 100644 index 0000000..e3ea669 --- /dev/null +++ b/infra/media-worker/tests/fixtures/media-capabilities.json @@ -0,0 +1,21 @@ +{ + "$comment": "Shape of MediaCapabilitiesSchema in packages/contracts/src/media.ts. Keep in sync with the contract.", + "keys": ["configured", "operations", "accelerator", "hardwareEncoders"], + "operations": [ + "detectScenes", + "transcribe", + "detectFaces", + "trackFaces", + "faceLandmarks", + "embedFrames", + "classifyContent", + "segmentPrompted", + "inpaintVideo", + "redact", + "synthesizeIdentity", + "audioFeatures", + "saliency" + ], + "operationKeys": ["available", "experimental", "model", "license", "commercialUse"], + "accelerators": ["cpu", "cuda", "rocm", "metal", "none"] +} diff --git a/infra/media-worker/tests/test_analysis.py b/infra/media-worker/tests/test_analysis.py new file mode 100644 index 0000000..b02bc1b --- /dev/null +++ b/infra/media-worker/tests/test_analysis.py @@ -0,0 +1,88 @@ +"""Operations without model weights, run for real on synthetic media.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from fastapi.testclient import TestClient + +from app.media_io import parse_ebur128 + +from .conftest import AUTH + + +def test_audio_features_on_click_track(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/clicks.wav"}) + assert response.status_code == 200, response.text + body = response.json() + assert body["integratedLufs"] is not None and -60 < body["integratedLufs"] < 0 + assert body["truePeakDb"] is not None and body["truePeakDb"] <= 1.0 + tempo = body["tempoBpm"] + assert tempo is not None and any(abs(tempo - bpm) < 6 for bpm in (60, 120, 240)) + beats = body["beatsMs"] + assert len(beats) >= 6 + gaps = np.diff(beats) + assert np.median(gaps) % 500 < 40 or 500 - np.median(gaps) % 500 < 40 + # The trailing 1.5 s are silent. + assert any(7400 <= s["startMs"] <= 7700 and s["endMs"] >= 9400 for s in body["silences"]), body[ + "silences" + ] + energy = body["energy"] + assert energy[0]["tMs"] == 0 and all(0.0 <= e["value"] <= 1.0 for e in energy) + assert max(e["value"] for e in energy) == 1.0 + assert all(isinstance(t, int) for t in beats) + + +def test_audio_features_without_audio_stream(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/disc.mp4"}) + assert response.status_code == 200 + assert response.json() == { + "integratedLufs": None, + "truePeakDb": None, + "tempoBpm": None, + "beatsMs": [], + "energy": [], + "silences": [], + } + + +def test_parse_ebur128_summary() -> None: + stderr = """ +[Parsed_ebur128_0 @ 0x0] Summary: + + Integrated loudness: + I: -23.4 LUFS + Threshold: -33.6 LUFS + + True peak: + Peak: -1.2 dBFS +""" + assert parse_ebur128(stderr) == (-23.4, -1.2) + assert parse_ebur128("Summary:\n I: -inf LUFS\n Peak: -inf dBFS") == (None, None) + + +def test_scene_detection_finds_hard_cut(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/detectScenes", headers=AUTH, json={"path": f"{media_root}/cut.mp4"}) + assert response.status_code == 200, response.text + body = response.json() + shots = body["shots"] + assert len(shots) == 2, shots + assert abs(shots[1]["startMs"] - 2000) <= 80 + assert shots[0]["startMs"] == 0 and abs(shots[-1]["endMs"] - 4000) <= 80 + assert all(0.0 <= s["score"] <= 1.0 for s in shots) + assert body["scenes"][0]["startMs"] == 0 + assert body["scenes"][-1]["endMs"] == shots[-1]["endMs"] + + +def test_saliency_points_at_bright_spot(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/saliency", headers=AUTH, json={"path": f"{media_root}/spot.mp4", "timesMs": [0, 200]} + ) + assert response.status_code == 200, response.text + points = response.json()["points"] + assert [p["tMs"] for p in points] == [0, 200] + for point in points: + assert point["cx"] > 0.6 + assert 0.3 < point["cy"] < 0.7 + assert 0.0 < point["strength"] <= 1.0 diff --git a/infra/media-worker/tests/test_auth.py b/infra/media-worker/tests/test_auth.py new file mode 100644 index 0000000..8866032 --- /dev/null +++ b/infra/media-worker/tests/test_auth.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from app.config import ConfigError, Settings + +from .conftest import AUTH, TOKEN + + +def test_healthz_needs_no_token(client: TestClient) -> None: + response = client.get("/healthz") + assert response.status_code == 200 + assert response.json()["ok"] is True + assert isinstance(response.json()["version"], str) + + +@pytest.mark.parametrize( + "headers", + [ + {}, + {"Authorization": "Bearer wrong-token-" + "y" * 32}, + {"Authorization": TOKEN}, + {"Authorization": "Basic " + TOKEN}, + ], +) +def test_routes_reject_missing_or_wrong_token(client: TestClient, headers: dict[str, str]) -> None: + for method, url in [("GET", "/v1/capabilities"), ("POST", "/v1/detectScenes")]: + response = client.request(method, url, headers=headers, json={"path": "/x"}) + assert response.status_code == 401 + assert response.json()["error"]["code"] == "unauthorized" + + +def test_valid_token_is_accepted(client: TestClient) -> None: + assert client.get("/v1/capabilities", headers=AUTH).status_code == 200 + + +def test_refuses_to_start_without_a_long_token(tmp_path: Path) -> None: + env = {"MEDIA_ROOT": str(tmp_path)} + with pytest.raises(ConfigError): + Settings.from_env(env) + with pytest.raises(ConfigError): + Settings.from_env({**env, "MEDIA_WORKER_TOKEN": "short"}) + assert Settings.from_env({**env, "MEDIA_WORKER_TOKEN": "z" * 32}).token == "z" * 32 + assert Settings.from_env({**env, "MEDIA_WORKER_INSECURE_DEV": "1"}).insecure_dev + + +def test_settings_validate_root_and_defaults(tmp_path: Path) -> None: + with pytest.raises(ConfigError): + Settings.from_env({"MEDIA_WORKER_TOKEN": "z" * 32, "MEDIA_ROOT": "relative/path"}) + settings = Settings.from_env({"MEDIA_WORKER_TOKEN": "z" * 32, "MEDIA_ROOT": str(tmp_path)}) + assert settings.host == "127.0.0.1" + assert settings.port == 8790 + assert settings.gpu_concurrency == 1 + assert settings.max_input_bytes == 20 * 1024**3 + assert settings.max_seconds == 4 * 3600 diff --git a/infra/media-worker/tests/test_capabilities.py b/infra/media-worker/tests/test_capabilities.py new file mode 100644 index 0000000..18640f1 --- /dev/null +++ b/infra/media-worker/tests/test_capabilities.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import json +import re +from pathlib import Path + +from fastapi.testclient import TestClient + +from app.models import OPTIONAL_BACKENDS, Registry, default_backends + +from .conftest import AUTH, make_settings + +CONTRACT = json.loads((Path(__file__).parent / "fixtures" / "media-capabilities.json").read_text()) + + +def _assert_shape(body: dict) -> None: + assert sorted(body) == sorted(CONTRACT["keys"]) + assert body["configured"] is True + assert body["accelerator"] in CONTRACT["accelerators"] + assert isinstance(body["hardwareEncoders"], list) + assert all(isinstance(name, str) for name in body["hardwareEncoders"]) + assert sorted(body["operations"]) == sorted(CONTRACT["operations"]) + for entry in body["operations"].values(): + assert sorted(entry) == sorted(CONTRACT["operationKeys"]) + assert isinstance(entry["available"], bool) + assert isinstance(entry["experimental"], bool) + assert entry["model"] is None or isinstance(entry["model"], str) + assert entry["license"] is None or isinstance(entry["license"], str) + assert entry["commercialUse"] in (True, False, None) + + +def test_capabilities_match_the_contract(make_client) -> None: + # Real registry, no fakes, downloads off: whatever is installed here. + body = make_client(fakes=False).get("/v1/capabilities", headers=AUTH).json() + _assert_shape(body) + ops = body["operations"] + # Operations that need no model weights work offline with the core requirements. + for name in ("detectScenes", "redact", "audioFeatures", "saliency"): + assert ops[name]["available"] is True, name + assert ops["synthesizeIdentity"] == { + "available": False, + "experimental": True, + "model": None, + "license": None, + "commercialUse": None, + } + # Weights are not installed and downloads are off. + for name in ("transcribe", "detectFaces", "embedFrames", "inpaintVideo", "segmentPrompted"): + assert ops[name]["available"] is False, name + + +def test_with_fake_backends_everything_but_synthesis_is_available(client: TestClient) -> None: + body = client.get("/v1/capabilities", headers=AUTH).json() + _assert_shape(body) + unavailable = [name for name, entry in body["operations"].items() if not entry["available"]] + assert unavailable == ["synthesizeIdentity"] + + +def test_default_backends_are_commercially_usable_and_pinned() -> None: + for operation, backend in default_backends().items(): + spec = backend.spec + if spec.unavailable_reason: + continue + assert spec.commercialUse is True, operation + assert spec.license, operation + for model_file in spec.files: + assert re.fullmatch(r"[0-9a-f]{64}", model_file.sha256), model_file.name + assert model_file.sizeBytes > 0 + assert model_file.url.startswith("https://") + assert "/resolve/main/" not in model_file.url, "pin a revision, not a branch" + for backend in OPTIONAL_BACKENDS.values(): + assert backend.spec.commercialUse is not True + + +def test_optional_backend_needs_opt_in(media_root: Path, tmp_path: Path) -> None: + registry = Registry(make_settings(media_root, tmp_path)) + assert not registry.optional_enabled("faceIdentity") + reason = registry.unavailable_reason(OPTIONAL_BACKENDS["faceIdentity"], "faceIdentity") + assert reason is not None and "MEDIA_OPTIONAL_BACKENDS" in reason + opted = Registry(make_settings(media_root, tmp_path, optional_backends=frozenset({"faceIdentity"}))) + reason = opted.unavailable_reason(OPTIONAL_BACKENDS["faceIdentity"], "faceIdentity") + assert reason is None or "MEDIA_OPTIONAL_BACKENDS" not in reason diff --git a/infra/media-worker/tests/test_ml_ops.py b/infra/media-worker/tests/test_ml_ops.py new file mode 100644 index 0000000..f21d24a --- /dev/null +++ b/infra/media-worker/tests/test_ml_ops.py @@ -0,0 +1,170 @@ +"""Model-backed operations with fake runtimes; frame IO, tracking, masks and encoding are real.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from fastapi.testclient import TestClient + +from app.ops.faces import Detection, IouTracker + +from .conftest import AUTH, decode + + +def test_transcribe(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/transcribe", headers=AUTH, json={"path": f"{media_root}/testsrc.mp4", "language": "de"} + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["language"] == "de" + assert [w["text"] for w in body["words"]] == ["synthetic", "tone"] + + +def test_transcribe_rejects_silent_video(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/transcribe", headers=AUTH, json={"path": f"{media_root}/disc.mp4"}) + assert response.status_code == 400 + + +def test_detect_faces_sampled_and_at_times(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/detectFaces", headers=AUTH, json={"path": f"{media_root}/disc.mp4", "sampleFps": 5} + ) + assert response.status_code == 200, response.text + frames = response.json()["frames"] + assert [f["tMs"] for f in frames][:3] == [0, 200, 400] + assert all(len(f["faces"]) == 1 for f in frames) + face = frames[0]["faces"][0] + assert abs(face["box"]["x"] - 35 / 320) < 0.02 and abs(face["box"]["w"] - 51 / 320) < 0.02 + assert len(face["landmarks"]) == 5 + response = client.post( + "/v1/detectFaces", + headers=AUTH, + json={"path": f"{media_root}/disc.mp4", "sampleFps": 5, "timesMs": [800, 0]}, + ) + assert [f["tMs"] for f in response.json()["frames"]] == [0, 800] + + +def test_track_faces_follows_the_disc(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/trackFaces", headers=AUTH, json={"path": f"{media_root}/disc.mp4", "sampleFps": 10} + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["sampleFps"] == 10 + assert len(body["tracks"]) == 1 + track = body["tracks"][0] + assert track["trackId"] == "t1" + assert track["embedding"] is None + xs = [k["box"]["x"] for k in track["keyframes"]] + assert len(xs) >= 15 and xs == sorted(xs) and xs[-1] > xs[0] + 0.3 + + +def test_tracker_splits_after_gap_and_uses_low_confidence() -> None: + tracker = IouTracker(max_gap_ms=500) + box = (10.0, 10.0, 50.0, 50.0) + tracker.update(0, [Detection(box, 0.9, [])], None, False) + # Low-confidence detection extends the existing track but never starts one. + matches = tracker.update( + 100, + [Detection((12.0, 10.0, 50.0, 50.0), 0.4, []), Detection((200.0, 200.0, 30.0, 30.0), 0.4, [])], + None, + False, + ) + assert [m[0].track_id for m in matches] == ["t1"] + # 600 ms without a match: the next detection starts a new track. + tracker.update(800, [Detection(box, 0.9, [])], None, False) + assert [t.track_id for t in tracker.close()] == ["t1", "t2"] + + +def test_face_landmarks(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/faceLandmarks", headers=AUTH, json={"path": f"{media_root}/disc.mp4", "timesMs": [0, 400]} + ) + assert response.status_code == 200, response.text + frames = response.json()["frames"] + assert [f["tMs"] for f in frames] == [0, 400] + assert len(frames[1]["faces"][0]["landmarks"]) == 5 + + +def test_embed_frames_keeps_request_order(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/embedFrames", + headers=AUTH, + json={"path": f"{media_root}/testsrc.mp4", "timesMs": [1000, 0, 1000]}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["model"] == "fake-embed" and body["dim"] == 4 + assert len(body["vectors"]) == 3 and body["vectors"][0] == body["vectors"][2] + assert all(abs(np.linalg.norm(v) - 1) < 1e-5 for v in body["vectors"]) + + +def test_classify_content_samples_plus_extra_times(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/classifyContent", + headers=AUTH, + json={"path": f"{media_root}/testsrc.mp4", "sampleFps": 2, "extraTimesMs": [250, 500]}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["model"] == "fake-classifier" + times = [f["tMs"] for f in body["frames"]] + assert times == sorted(times) and {0, 250, 500} <= set(times) + # Samples are the source frames nearest to each 500 ms slot (25 fps source). + for slot in (1000, 1500): + assert min(abs(t - slot) for t in times) <= 20 + assert all(0 <= f["scores"]["violence"] <= 1 for f in body["frames"]) + + +def test_segment_prompted_writes_masks_and_keyframes(client: TestClient, media_root: Path) -> None: + out_dir = media_root / "marks" / "m1" + body = { + "path": f"{media_root}/disc.mp4", + "prompt": {"tMs": 400, "point": {"x": 0.25, "y": 0.5}}, + "propagate": True, + "outDir": str(out_dir), + } + response = client.post("/v1/segmentPrompted", headers=AUTH, json=body) + assert response.status_code == 200, response.text + result = response.json() + assert result["maskDir"] == str(out_dir) + assert len(result["keyframes"]) == 24 # 1.6 s at 15 fps + assert 0.9 < result["score"] <= 1.0 + names = sorted(int(p.stem) for p in out_dir.glob("*.png")) + assert names == [k["tMs"] for k in result["keyframes"]] + single = client.post("/v1/segmentPrompted", headers=AUTH, json=body | {"propagate": False}) + assert single.status_code == 200 + assert [k["tMs"] for k in single.json()["keyframes"]] == [400] + assert [p.name for p in out_dir.glob("*.png")] == ["400.png"] + + +def test_inpaint_video_removes_masked_disc(client: TestClient, media_root: Path) -> None: + out_dir = media_root / "marks" / "inpaint" + seg = client.post( + "/v1/segmentPrompted", + headers=AUTH, + json={ + "path": f"{media_root}/disc.mp4", + "prompt": {"tMs": 0, "point": {"x": 0.2, "y": 0.5}}, + "propagate": True, + "outDir": str(out_dir), + }, + ) + assert seg.status_code == 200 + out = media_root / "inpainted" / "disc.mkv" + response = client.post( + "/v1/inpaintVideo", + headers=AUTH, + json={"path": f"{media_root}/disc.mp4", "maskDir": str(out_dir), "outPath": str(out)}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["outPath"] == str(out) + quality = body["quality"] + assert set(quality) == {"temporalStability", "edgeQuality", "score"} + assert all(0.0 <= v <= 1.0 for v in quality.values()) + before, after = decode(media_root / "disc.mp4"), decode(out) + assert after.shape == before.shape + assert after.mean() < before.mean() * 0.3 diff --git a/infra/media-worker/tests/test_models.py b/infra/media-worker/tests/test_models.py new file mode 100644 index 0000000..f6d9617 --- /dev/null +++ b/infra/media-worker/tests/test_models.py @@ -0,0 +1,75 @@ +"""Lazy download and SHA256 pinning, exercised with local file:// URLs.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest + +from app.errors import Unsupported +from app.models import Backend, BackendSpec, ModelFile, Registry + +from .conftest import make_settings + + +def _backend(source: Path, sha256: str, size: int, loads: list[dict[str, str]]) -> Backend: + spec = BackendSpec( + id="test-model", + model="test-model", + license="MIT", + commercialUse=True, + experimental=False, + source=None, + files=(ModelFile("weights.bin", source.as_uri(), sha256, size),), + ) + return Backend(spec, lambda files, _settings: loads.append(files) or "runtime") + + +def test_downloads_once_and_verifies(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "upstream.bin" + source.write_bytes(b"weights" * 100) + digest = hashlib.sha256(source.read_bytes()).hexdigest() + loads: list[dict[str, str]] = [] + settings = make_settings(media_root, tmp_path, allow_model_download=True) + registry = Registry(settings, backends={"embedFrames": _backend(source, digest, 700, loads)}) + assert registry.runtime("embedFrames") == "runtime" + assert registry.runtime("embedFrames") == "runtime" + assert len(loads) == 1 + cached = Path(loads[0]["weights.bin"]) + assert cached == tmp_path / "models" / "test-model" / "weights.bin" + assert cached.read_bytes() == source.read_bytes() + + +def test_hash_mismatch_deletes_file_and_reports_unsupported(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "tampered.bin" + source.write_bytes(b"tampered") + settings = make_settings(media_root, tmp_path, allow_model_download=True) + backend = _backend(source, "0" * 64, 8, []) + registry = Registry(settings, backends={"embedFrames": backend}) + with pytest.raises(Unsupported) as error: + registry.runtime("embedFrames") + assert "SHA256" in error.value.message + assert not (tmp_path / "models" / "test-model" / "weights.bin").exists() + assert registry.unavailable_reason(backend, "embedFrames") == error.value.message + + +def test_oversized_download_is_aborted(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "big.bin" + source.write_bytes(b"x" * 5000) + settings = make_settings(media_root, tmp_path, allow_model_download=True) + registry = Registry(settings, backends={"embedFrames": _backend(source, "0" * 64, 10, [])}) + with pytest.raises(Unsupported): + registry.runtime("embedFrames") + assert not list((tmp_path / "models" / "test-model").iterdir()) + + +def test_offline_without_weights_is_unsupported(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "w.bin" + source.write_bytes(b"w") + registry = Registry( + make_settings(media_root, tmp_path), backends={"embedFrames": _backend(source, "0" * 64, 1, [])} + ) + with pytest.raises(Unsupported) as error: + registry.runtime("embedFrames") + assert "MEDIA_MODEL_DOWNLOAD=0" in error.value.message diff --git a/infra/media-worker/tests/test_paths.py b/infra/media-worker/tests/test_paths.py new file mode 100644 index 0000000..673bee8 --- /dev/null +++ b/infra/media-worker/tests/test_paths.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import os +import shutil +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from app.errors import InvalidPath +from app.paths import input_file, output_file, resolve_inside + +from .conftest import AUTH + + +@pytest.fixture +def root(tmp_path: Path) -> Path: + path = tmp_path / "root" + path.mkdir() + (path / "a.txt").write_text("a") + return path + + +@pytest.mark.parametrize( + "requested", + [ + "relative.txt", + "", + "/etc/passwd", + "{root}/../outside.txt", + "{root}/sub/../../outside.txt", + "{root}/a\0.txt", + ], +) +def test_rejects_paths_outside_root(root: Path, requested: str) -> None: + with pytest.raises(InvalidPath): + resolve_inside(str(root), requested.format(root=root)) + + +def test_accepts_paths_inside_root(root: Path) -> None: + assert resolve_inside(str(root), f"{root}/a.txt") == os.path.realpath(root / "a.txt") + assert resolve_inside(str(root), f"{root}/new/dir/file.mkv").endswith("/new/dir/file.mkv") + + +def test_rejects_symlink_escape(root: Path, tmp_path: Path) -> None: + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.mp4").write_text("x") + os.symlink(outside, root / "link") + os.symlink(outside / "secret.mp4", root / "file-link.mp4") + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/link/secret.mp4") + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/file-link.mp4") + with pytest.raises(InvalidPath): + output_file(str(root), f"{root}/link/new/out.mkv") + assert not (outside / "new").exists() + + +def test_symlink_inside_root_is_allowed(root: Path) -> None: + os.symlink(root / "a.txt", root / "alias.txt") + assert input_file(str(root), f"{root}/alias.txt") == os.path.realpath(root / "a.txt") + + +def test_missing_input_and_directory_input(root: Path) -> None: + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/missing.mp4") + (root / "dir").mkdir() + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/dir") + + +def test_output_parent_is_created(root: Path) -> None: + real = output_file(str(root), f"{root}/renders/r1/out.mkv") + assert os.path.isdir(os.path.dirname(real)) + + +def test_http_rejects_traversal_and_symlink_escape( + client: TestClient, media_root: Path, tmp_path: Path +) -> None: + outside = tmp_path / "elsewhere" + outside.mkdir() + shutil.copy(media_root / "testsrc.mp4", outside / "clip.mp4") + link = media_root / "escape-link" + if not link.exists(): + os.symlink(outside, link) + for path in ["/etc/hostname", f"{media_root}/../x.mp4", f"{link}/clip.mp4", "testsrc.mp4"]: + response = client.post("/v1/detectScenes", headers=AUTH, json={"path": path}) + assert response.status_code == 400, path + assert response.json()["error"] == { + "code": "invalid_path", + "message": response.json()["error"]["message"], + "operation": "detectScenes", + } + response = client.post( + "/v1/redact", + headers=AUTH, + json={ + "path": f"{media_root}/testsrc.mp4", + "outPath": f"{link}/out.mkv", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 0.1, "h": 0.1}}], + "mode": "blur", + "paddingPct": 0, + "featherPct": 0, + } + ], + }, + ) + assert response.status_code == 400 + assert response.json()["error"]["code"] == "invalid_path" + assert not (outside / "out.mkv").exists() + + +def test_input_size_limit(make_client, media_root: Path) -> None: + client = make_client(max_input_bytes=10) + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/testsrc.mp4"}) + assert response.status_code == 413 + assert response.json()["error"]["code"] == "too_large" + + +def test_duration_limit(make_client, media_root: Path) -> None: + client = make_client(max_seconds=1) + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/testsrc.mp4"}) + assert response.status_code == 413 diff --git a/infra/media-worker/tests/test_redact.py b/infra/media-worker/tests/test_redact.py new file mode 100644 index 0000000..92091e8 --- /dev/null +++ b/infra/media-worker/tests/test_redact.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import cv2 +import numpy as np +import pytest +from fastapi.testclient import TestClient + +import app.ops.redact as redact_module + +from .conftest import AUTH, decode + +BOX = {"x": 0.25, "y": 0.25, "w": 0.5, "h": 0.5} +# Pixel region of BOX in the 320x240 fixture. +X0, Y0, X1, Y1 = 80, 60, 240, 180 + + +def _body(media_root: Path, out: Path, mode: str, **extra: object) -> dict[str, object]: + return { + "path": f"{media_root}/testsrc.mp4", + "outPath": str(out), + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": BOX}, {"tMs": 2000, "box": BOX}], + "mode": mode, + "paddingPct": 0, + "featherPct": 0, + **extra, + } + ], + } + + +def _outside_mask(shape: tuple[int, ...], margin: int) -> np.ndarray: + mask = np.ones(shape[:2], bool) + mask[Y0 - margin : Y1 + margin, X0 - margin : X1 + margin] = False + return mask + + +def _streams(path: Path) -> list[str]: + out = subprocess.run( + ["ffprobe", "-v", "error", "-show_entries", "stream=codec_type,codec_name", "-of", "json", str(path)], + check=True, + capture_output=True, + text=True, + ).stdout + return [f"{s['codec_type']}:{s['codec_name']}" for s in json.loads(out)["streams"]] + + +def test_solid_redaction_changes_only_the_box(client: TestClient, media_root: Path, tmp_path: Path) -> None: + out = media_root / "redact" / "solid.mkv" + response = client.post("/v1/redact", headers=AUTH, json=_body(media_root, out, "solid", color="#000000")) + assert response.status_code == 200, response.text + assert response.json() == {"outPath": str(out), "framesProcessed": 50} + source, result = decode(media_root / "testsrc.mp4"), decode(out) + assert result.shape == source.shape + inside = result[:, Y0 + 2 : Y1 - 2, X0 + 2 : X1 - 2] + assert inside.max() <= 8 + assert np.abs(source[:, Y0:Y1, X0:X1].astype(int)).mean() > 50 + outside = _outside_mask(source.shape[1:], 2) + diff = np.abs(result.astype(int) - source.astype(int))[:, outside] + assert diff.mean() < 2.0 + assert _streams(out) == ["video:ffv1", "audio:aac"] + assert not [p for p in out.parent.iterdir() if ".partial-" in p.name] + + +@pytest.mark.parametrize("mode", ["blur", "pixelate"]) +def test_blur_and_pixelate_destroy_detail_inside(client: TestClient, media_root: Path, mode: str) -> None: + out = media_root / "redact" / f"{mode}.mp4" + body = _body(media_root, out, mode) | {"path": f"{media_root}/noise.mp4"} + response = client.post("/v1/redact", headers=AUTH, json=body) + assert response.status_code == 200, response.text + source, result = decode(media_root / "noise.mp4"), decode(out) + region = (slice(None), slice(Y0 + 4, Y1 - 4), slice(X0 + 4, X1 - 4)) + + # Local detail (gradient energy) collapses inside the box. + def detail(frames: np.ndarray) -> float: + gray = frames.astype(np.float32).mean(axis=3) + return float(np.abs(np.diff(gray, axis=2)).mean()) + + assert detail(result[region]) < detail(source[region]) * 0.2 + assert ( + result[region].astype(np.float32).var(axis=(1, 2)).mean() + < source[region].astype(np.float32).var(axis=(1, 2)).mean() * 0.2 + ) + outside = _outside_mask(source.shape[1:], 4) + assert np.abs(result.astype(int) - source.astype(int))[:, outside].mean() < 3.0 + assert _streams(out) == ["video:h264", "audio:aac"] + + +@pytest.mark.parametrize("mode", ["sticker", "cover", "partial_eyes", "partial_upper", "partial_lower"]) +def test_other_modes_produce_output(client: TestClient, media_root: Path, mode: str) -> None: + out = media_root / "redact" / f"{mode}.mkv" + body = _body(media_root, out, mode) + body["redactions"][0]["featherPct"] = 10 # type: ignore[index] + body["redactions"][0]["paddingPct"] = 20 # type: ignore[index] + if mode == "partial_eyes": + body["redactions"][0]["landmarks"] = [ # type: ignore[index] + {"tMs": 0, "points": [{"x": 0.4, "y": 0.4}, {"x": 0.6, "y": 0.4}, {"x": 0.5, "y": 0.5}]} + ] + response = client.post("/v1/redact", headers=AUTH, json=body) + assert response.status_code == 200, response.text + source, result = decode(media_root / "testsrc.mp4"), decode(out) + changed = np.abs(result.astype(int) - source.astype(int))[:, Y0:Y1, X0:X1].mean() + assert changed > 10 + corner = np.abs(result.astype(int) - source.astype(int))[:, :20, :20].mean() + assert corner < 2.0 + + +def test_mask_dir_redaction(client: TestClient, media_root: Path) -> None: + masks = media_root / "redact-masks" + masks.mkdir(exist_ok=True) + for t in range(0, 2000, 200): + mask = np.zeros((240, 320), np.uint8) + mask[100:140, 140:180] = 255 + cv2.imwrite(str(masks / f"{t}.png"), mask) + out = media_root / "redact" / "mask.mkv" + body = { + "path": f"{media_root}/testsrc.mp4", + "outPath": str(out), + "redactions": [ + { + "keyframes": [], + "maskDir": str(masks), + "mode": "cover", + "color": "#123456", + "paddingPct": 0, + "featherPct": 0, + } + ], + } + response = client.post("/v1/redact", headers=AUTH, json=body) + assert response.status_code == 200, response.text + result = decode(out) + patch = result[:, 104:136, 144:176].reshape(-1, 3).mean(axis=0) + assert np.allclose(patch, [0x56, 0x34, 0x12], atol=6) + + +def test_failure_mid_video_leaves_no_output( + client: TestClient, media_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + calls = {"n": 0} + original = redact_module.apply_redactions + + def flaky(frame, t_ms, plans, frame_ms): # type: ignore[no-untyped-def] + calls["n"] += 1 + if calls["n"] == 10: + raise RuntimeError("simulated frame failure") + return original(frame, t_ms, plans, frame_ms) + + monkeypatch.setattr(redact_module, "apply_redactions", flaky) + out = media_root / "redact-fail" / "out.mkv" + response = client.post("/v1/redact", headers=AUTH, json=_body(media_root, out, "blur")) + assert response.status_code == 500 + assert response.json()["error"]["code"] == "internal" + assert not out.exists() + assert list(out.parent.iterdir()) == [] + + +def test_unreadable_input_fails_closed(client: TestClient, media_root: Path) -> None: + broken = media_root / "broken.mp4" + data = (media_root / "testsrc.mp4").read_bytes() + broken.write_bytes(data[: len(data) // 3]) + out = media_root / "redact-broken" / "out.mkv" + response = client.post( + "/v1/redact", headers=AUTH, json=_body(media_root, out, "blur") | {"path": str(broken)} + ) + assert response.status_code in (400, 500) + assert not out.exists() + if out.parent.exists(): + assert not list(out.parent.iterdir()) diff --git a/infra/media-worker/tests/test_runtime.py b/infra/media-worker/tests/test_runtime.py new file mode 100644 index 0000000..75ebb03 --- /dev/null +++ b/infra/media-worker/tests/test_runtime.py @@ -0,0 +1,117 @@ +"""Cancellation on disconnect, the time limit, and GPU queue overflow.""" + +from __future__ import annotations + +import asyncio +import threading +from pathlib import Path + +import numpy as np +import pytest + +from app import media_io +from app.errors import Busy, Cancelled +from app.runtime import Gate, run_job + +from .conftest import AUTH +from .fakes import BlockingEmbedder, fake_backend + + +class FakeRequest: + def __init__(self, disconnect_after: int | None) -> None: + self.calls = 0 + self.disconnect_after = disconnect_after + + async def is_disconnected(self) -> bool: + self.calls += 1 + return self.disconnect_after is not None and self.calls > self.disconnect_after + + +def _spin(seen: dict[str, bool]): # type: ignore[no-untyped-def] + def work(cancel: threading.Event) -> str: + for _ in range(400): + if cancel.wait(0.01): + seen["cancelled"] = True + raise Cancelled() + return "done" + + return work + + +def test_disconnect_cancels_work() -> None: + seen: dict[str, bool] = {} + gate = Gate(1, 0) + with pytest.raises(Cancelled) as error: + asyncio.run(run_job(FakeRequest(disconnect_after=1), gate, _spin(seen), max_seconds=60)) + assert error.value.reason == "client disconnected" + assert seen["cancelled"] + assert gate.inflight == 0 + + +def test_time_limit_cancels_work() -> None: + seen: dict[str, bool] = {} + with pytest.raises(Cancelled) as error: + asyncio.run(run_job(FakeRequest(None), Gate(1, 0), _spin(seen), max_seconds=0.3)) + assert error.value.reason == "exceeded MEDIA_MAX_SECONDS" + + +def test_completed_work_returns_result() -> None: + assert asyncio.run(run_job(FakeRequest(None), Gate(1, 0), lambda _c: "ok", max_seconds=5)) == "ok" + + +def test_gate_overflow_is_busy() -> None: + gate = Gate(1, 1) + gate.reserve() + gate.reserve() + with pytest.raises(Busy): + gate.reserve() + gate.release_reservation() + gate.reserve() + + +def test_cancelled_transform_leaves_no_output(media_root: Path) -> None: + cancel = threading.Event() + path = str(media_root / "testsrc.mp4") + info = media_io.probe(path, threading.Event()) + out = media_root / "cancelled" / "out.mkv" + out.parent.mkdir(exist_ok=True) + + def frame_fn(index: int, _t: int, frame: np.ndarray) -> np.ndarray: + if index == 5: + cancel.set() + return frame + + with pytest.raises(Cancelled): + media_io.transform_video(path, str(out), info, cancel, frame_fn) + assert list(out.parent.iterdir()) == [] + + +def test_http_busy_when_gpu_queue_is_full(make_client, media_root: Path) -> None: + blocking = BlockingEmbedder() + client = make_client( + backends={"embedFrames": fake_backend(blocking, "fake-embed")}, gpu_concurrency=1, queue_size=0 + ) + body = {"path": f"{media_root}/testsrc.mp4", "timesMs": [0]} + results: list[int] = [] + first = threading.Thread( + target=lambda: results.append(client.post("/v1/embedFrames", headers=AUTH, json=body).status_code) + ) + first.start() + assert blocking.entered.wait(20) + second = client.post("/v1/embedFrames", headers=AUTH, json=body) + assert second.status_code == 503 + assert second.json()["error"] == { + "code": "busy", + "message": "inference queue is full", + "operation": "embedFrames", + } + # CPU-bound operations use their own gate and still run. + assert ( + client.post( + "/v1/saliency", headers=AUTH, json={"path": f"{media_root}/spot.mp4", "timesMs": [0]} + ).status_code + == 200 + ) + blocking.release.set() + first.join(20) + assert results == [200] diff --git a/infra/media-worker/tests/test_schema.py b/infra/media-worker/tests/test_schema.py new file mode 100644 index 0000000..e2db207 --- /dev/null +++ b/infra/media-worker/tests/test_schema.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from .conftest import AUTH + + +def _synth(media_root: Path, **extra: object) -> dict[str, object]: + return { + "path": f"{media_root}/testsrc.mp4", + "track": [{"tMs": 0, "box": {"x": 0.1, "y": 0.1, "w": 0.2, "h": 0.2}}], + "seed": 42, + "outPath": f"{media_root}/synth/out.mkv", + **extra, + } + + +@pytest.mark.parametrize("field", ["targetFace", "targetImage", "faceImagePath", "referencePath", "image"]) +def test_synthesize_identity_rejects_any_target_face_field( + client: TestClient, media_root: Path, field: str +) -> None: + response = client.post("/v1/synthesizeIdentity", headers=AUTH, json=_synth(media_root, **{field: "data"})) + assert response.status_code == 400 + error = response.json()["error"] + assert error["code"] == "bad_request" + assert error["operation"] == "synthesizeIdentity" + + +def test_synthesize_identity_is_unsupported_without_licensed_model( + client: TestClient, media_root: Path +) -> None: + response = client.post("/v1/synthesizeIdentity", headers=AUTH, json=_synth(media_root)) + assert response.status_code == 501 + assert response.json()["error"] == { + "code": "unsupported", + "message": "no commercially licensed model", + "operation": "synthesizeIdentity", + } + assert not (media_root / "synth" / "out.mkv").exists() + + +@pytest.mark.parametrize( + ("operation", "body"), + [ + ("detectScenes", {"path": "/x", "extra": 1}), + ("detectFaces", {"path": "/x"}), + ("detectFaces", {"path": "/x", "sampleFps": 0}), + ("detectFaces", {"path": "/x", "sampleFps": "2"}), + ("embedFrames", {"path": "/x", "timesMs": [1.5]}), + ("embedFrames", {"path": "/x", "timesMs": [-1]}), + ("segmentPrompted", {"path": "/x", "prompt": {"tMs": 0}, "propagate": True, "outDir": "/y"}), + ( + "segmentPrompted", + { + "path": "/x", + "prompt": {"tMs": 0, "box": {"x": 2, "y": 0, "w": 1, "h": 1}}, + "propagate": True, + "outDir": "/y", + }, + ), + ("redact", {"path": "/x", "outPath": "/y", "redactions": []}), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [{"keyframes": [], "mode": "blur", "paddingPct": 0, "featherPct": 0}], + }, + ), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 1, "h": 1}}], + "mode": "reframe_out", + "paddingPct": 0, + "featherPct": 0, + } + ], + }, + ), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 1, "h": 1}}], + "mode": "solid", + "paddingPct": 0, + "featherPct": 0, + "color": "red", + } + ], + }, + ), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 1, "h": 1, "z": 1}}], + "mode": "blur", + "paddingPct": 0, + "featherPct": 0, + } + ], + }, + ), + ], +) +def test_schema_violations_are_bad_requests( + client: TestClient, operation: str, body: dict[str, object] +) -> None: + response = client.post(f"/v1/{operation}", headers=AUTH, json=body) + assert response.status_code == 400, response.text + assert response.json()["error"]["code"] == "bad_request" + assert response.json()["error"]["operation"] == operation + + +def test_malformed_json_and_unknown_operation(client: TestClient) -> None: + response = client.post( + "/v1/detectScenes", headers={**AUTH, "content-type": "application/json"}, content=b"{nope" + ) + assert response.status_code == 400 + response = client.post("/v1/makeDeepfake", headers=AUTH, json={}) + assert response.status_code == 404 + assert response.json()["error"]["code"] == "bad_request" From b4222ac5ec920ae9a4f94c1599d5c1917da6c0d3 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 18:21:16 +0000 Subject: [PATCH 3/3] feat(media): compose profile and CI job for the media sidecar Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj --- .env.example | 4 +++- .github/workflows/media-worker.yml | 32 ++++++++++++++++++++++++++++++ infra/compose/docker-compose.yml | 23 +++++++++++++++++++++ 3 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/media-worker.yml diff --git a/.env.example b/.env.example index 5f8d9eb..7c876d1 100644 --- a/.env.example +++ b/.env.example @@ -239,7 +239,9 @@ EVOO_COCKPIT_CORS_ORIGIN= # Optional Media Studio worker (local ffmpeg + ML sidecar, compose profile "media"). # Without it, Media Studio stays hidden and media jobs are skipped. -# MEDIA_WORKER_URL=http://127.0.0.1:8790 +# Compose: http://media-worker:8790. Bare-metal dev: http://127.0.0.1:8790. +# MEDIA_WORKER_URL=http://media-worker:8790 +# At least 32 random characters. # MEDIA_WORKER_TOKEN=REPLACE_WITH_RANDOM_TOKEN # Private platform profiles that override the shipped ones. Keep outside the repo. # MEDIA_PROFILES_DIR=~/.evoo/media-profiles diff --git a/.github/workflows/media-worker.yml b/.github/workflows/media-worker.yml new file mode 100644 index 0000000..f840436 --- /dev/null +++ b/.github/workflows/media-worker.yml @@ -0,0 +1,32 @@ +name: media-worker +on: + pull_request: + paths: &media-worker-paths + - ".github/workflows/media-worker.yml" + - "infra/media-worker/**" + push: + branches: [main] + paths: *media-worker-paths + +permissions: + contents: read + +concurrency: + group: media-worker-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + sidecar: + name: Media sidecar + runs-on: [self-hosted, evoo-ci, Linux] + timeout-minutes: 30 + steps: + - uses: actions/checkout@v5 + with: + persist-credentials: false + - name: Build test image (CPU) + run: docker build --build-arg ACCELERATOR=cpu --target test -t rakazo/media-worker:ci-test infra/media-worker + - name: Pytest (offline, fake model backends) + run: docker run --rm --network none rakazo/media-worker:ci-test + - name: Build runtime image (CPU) + run: docker build --build-arg ACCELERATOR=cpu -t rakazo/media-worker:ci infra/media-worker diff --git a/infra/compose/docker-compose.yml b/infra/compose/docker-compose.yml index 08ab8c4..8ac7429 100644 --- a/infra/compose/docker-compose.yml +++ b/infra/compose/docker-compose.yml @@ -129,6 +129,28 @@ services: data-init: condition: service_completed_successfully + media-worker: + # Optional Media Studio sidecar: `docker compose --profile media up`. + # No published port; the API and worker reach it as http://media-worker:8790. + profiles: ["media"] + build: + context: ../media-worker + args: + ACCELERATOR: ${MEDIA_ACCELERATOR:-cpu} + image: rakazo/media-worker:local + environment: + MEDIA_WORKER_HOST: "0.0.0.0" + MEDIA_WORKER_PORT: "8790" + MEDIA_WORKER_TOKEN: ${MEDIA_WORKER_TOKEN:?Set MEDIA_WORKER_TOKEN in .env} + MEDIA_ROOT: /data/media + MEDIA_MODELS_DIR: /models + HTTP_PROXY: ${HTTP_PROXY:-${http_proxy:-}} + HTTPS_PROXY: ${HTTPS_PROXY:-${https_proxy:-}} + NO_PROXY: ${NO_PROXY:-${no_proxy:-}} + volumes: + - ../../data/media:/data/media + - media-models:/models + web: build: context: ../.. @@ -144,4 +166,5 @@ services: - api volumes: + media-models: pgdata: