diff --git a/agent/Dockerfile b/agent/Dockerfile new file mode 100644 index 0000000..2e3ca9f --- /dev/null +++ b/agent/Dockerfile @@ -0,0 +1,20 @@ +# FabledCurator GPU agent — runs on the desktop with the GPU. +# CUDA runtime so onnxruntime-gpu can use the card; ffmpeg for video frames. +FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 +RUN apt-get update \ + && apt-get install -y --no-install-recommends python3 python3-pip ffmpeg \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /app +COPY requirements.txt . +RUN pip3 install --no-cache-dir -r requirements.txt +COPY fc_agent ./fc_agent + +# imgutils caches downloaded ONNX models here; mount a volume to persist them. +ENV HF_HOME=/models +EXPOSE 8770 + +# The control UI; the worker is started from it (or POST /start). +CMD ["uvicorn", "fc_agent.app:app", "--host", "0.0.0.0", "--port", "8770"] diff --git a/agent/README.md b/agent/README.md new file mode 100644 index 0000000..a787ae6 --- /dev/null +++ b/agent/README.md @@ -0,0 +1,60 @@ +# FabledCurator GPU agent + +A desktop-GPU worker that embeds characters (CCIP) + figure crops for +FabledCurator. It talks to FC **only over HTTP** — it leases jobs, fetches image +pixels, runs the models on your GPU, and posts results back. Your FC database and +Redis stay private; the agent never touches them. + +You run it when you want a burst and stop it to reclaim the card. + +## 1. Get a token +In FC: **Settings → Tagging → GPU agent → Generate token** (or Rotate). Copy it. + +## 2. Build +```sh +cd agent +docker build -t fc-gpu-agent . +``` + +## 3. Run (on the machine with the GPU) +```sh +docker run --rm --gpus all -p 8770:8770 \ + -e FC_URL=http://curator.traefik.internal \ + -e FC_TOKEN= \ + -v fc-agent-models:/models \ + fc-gpu-agent +``` +Then open — the control page. Click **Start** to begin +draining the queue; **Pause**/**Stop** to yield the GPU. The `-v fc-agent-models` +volume caches the downloaded ONNX models so restarts are fast. + +Kick off a backfill from FC (**GPU agent card → Queue character embedding**), then +watch the queue counts on the control page (or FC's card) drain. + +## Config (env) +| var | default | meaning | +|---|---|---| +| `FC_URL` | `http://localhost:8000` | FC base URL | +| `FC_TOKEN` | — | the bearer token (required) | +| `AGENT_ID` | `desktop-agent` | identifies this agent's leases | +| `BATCH_SIZE` | `4` | jobs leased per round (still processed one at a time) | +| `CCIP_MODEL` | imgutils default | CCIP model name | +| `DETECTOR_LEVEL` | `m` | person-detector size: `n` < `s` < `m` < `x` | +| `POLL_IDLE_SECONDS` | `10` | wait between empty leases | + +## ⚠️ Verify on first run +This part can't be CI-tested (no GPU/models in CI), so confirm against your +installed `dghs-imgutils` (`pip show dghs-imgutils`) — see `fc_agent/models.py`: +- `imgutils.detect.detect_person(image, level=...)` returns + `[((x0,y0,x1,y1), label, score), ...]`. +- `imgutils.metrics.ccip_extract_feature(image, model=...)` returns a vector + (768-d for caformer). If you want the F1-0.94 variant, set + `CCIP_MODEL=ccip-caformer_b36-24` (verify the exact string in imgutils). + +If FC's matcher under/over-fires, tune the cosine threshold in +`backend/app/services/ml/ccip.py` (`DEFAULT_SIM_THRESHOLD`) and use +`GET /api/ccip/overview` + `/api/ccip/images/` to spot-check. + +## CPU fallback +Swap `onnxruntime-gpu` → `onnxruntime` in `requirements.txt` and drop `--gpus all` +to grind it slowly on the server instead. Same agent, no card. diff --git a/agent/fc_agent/__init__.py b/agent/fc_agent/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/agent/fc_agent/app.py b/agent/fc_agent/app.py new file mode 100644 index 0000000..8329297 --- /dev/null +++ b/agent/fc_agent/app.py @@ -0,0 +1,94 @@ +"""FastAPI control surface for the agent (served on localhost). + +Start / pause / resume / stop the worker, set nothing else here (config is env), +and watch progress + the server-side queue. The container exposes this on a +localhost port; stopping the worker frees the GPU. +""" +from fastapi import FastAPI +from fastapi.responses import HTMLResponse, JSONResponse + +from .config import Config +from .worker import Worker + +cfg = Config.from_env() +worker = Worker(cfg) +app = FastAPI(title="FabledCurator GPU agent") + + +@app.get("/", response_class=HTMLResponse) +def index() -> str: + return _PAGE + + +@app.post("/start") +def start(): + worker.start() + return JSONResponse(worker.status()) + + +@app.post("/pause") +def pause(): + worker.pause() + return JSONResponse(worker.status()) + + +@app.post("/resume") +def resume(): + worker.resume() + return JSONResponse(worker.status()) + + +@app.post("/stop") +def stop(): + worker.stop() + return JSONResponse(worker.status()) + + +@app.get("/status") +def status(): + s = worker.status() + s["fc_url"] = cfg.fc_url + s["configured"] = bool(cfg.token) + try: + s["queue"] = worker.client.queue_status() + except Exception: + s["queue"] = None + return JSONResponse(s) + + +_PAGE = """ +FabledCurator GPU agent + +

FabledCurator GPU agent

+

FC: — · token —

+

+ + + + +

+

+ idle
state + 0
processed + 0
errors + —
current image +

+
+""" diff --git a/agent/fc_agent/client.py b/agent/fc_agent/client.py new file mode 100644 index 0000000..7a0701f --- /dev/null +++ b/agent/fc_agent/client.py @@ -0,0 +1,66 @@ +"""HTTP client for the FabledCurator GPU-job API. + +The agent's ONLY contact with FC — lease/submit/heartbeat/fail + fetch image +bytes, all over HTTP with the bearer token. No DB/Redis. +""" +import requests + + +class FcClient: + def __init__(self, base_url: str, token: str, agent_id: str): + self.base = base_url.rstrip("/") + self.agent_id = agent_id + self.s = requests.Session() + self.s.headers["Authorization"] = f"Bearer {token}" + + def lease(self, batch_size: int) -> list[dict]: + r = self.s.post( + f"{self.base}/api/gpu/jobs/lease", + json={"agent_id": self.agent_id, "batch_size": batch_size}, + timeout=30, + ) + r.raise_for_status() + return r.json().get("jobs", []) + + def submit(self, job_id: int, regions: list[dict], replace_kinds: list[str]) -> dict: + r = self.s.post( + f"{self.base}/api/gpu/jobs/submit", + json={ + "agent_id": self.agent_id, "job_id": job_id, + "regions": regions, "replace_kinds": replace_kinds, + }, + timeout=120, + ) + r.raise_for_status() + return r.json() + + def heartbeat(self, job_ids: list[int]) -> None: + try: + self.s.post( + f"{self.base}/api/gpu/jobs/heartbeat", + json={"agent_id": self.agent_id, "job_ids": job_ids}, + timeout=30, + ) + except requests.RequestException: + pass + + def fail(self, job_id: int, error: str) -> None: + try: + self.s.post( + f"{self.base}/api/gpu/jobs/fail", + json={"agent_id": self.agent_id, "job_id": job_id, "error": error}, + timeout=30, + ) + except requests.RequestException: + pass + + def fetch_image(self, image_url: str) -> bytes: + # image_url is a server-relative path ("/images/..."). + r = self.s.get(f"{self.base}{image_url}", timeout=180) + r.raise_for_status() + return r.content + + def queue_status(self) -> dict: + r = self.s.get(f"{self.base}/api/gpu/status", timeout=15) + r.raise_for_status() + return r.json() diff --git a/agent/fc_agent/config.py b/agent/fc_agent/config.py new file mode 100644 index 0000000..7c5d8d2 --- /dev/null +++ b/agent/fc_agent/config.py @@ -0,0 +1,26 @@ +"""Agent config, all from env (the control container is configured at run).""" +import os +from dataclasses import dataclass + + +@dataclass +class Config: + fc_url: str # base URL of the FabledCurator web service + token: str # the bearer token from Settings → Tagging → GPU agent + agent_id: str # identifies this agent's leases + batch_size: int # jobs leased per round (concurrency is still 1) + ccip_model: str # imgutils CCIP model name ("" → imgutils default) + detector_level: str # imgutils person-detector level: n|s|m|x + poll_idle_seconds: float # wait between empty leases + + @classmethod + def from_env(cls) -> "Config": + return cls( + fc_url=os.environ.get("FC_URL", "http://localhost:8000").rstrip("/"), + token=os.environ.get("FC_TOKEN", ""), + agent_id=os.environ.get("AGENT_ID", "desktop-agent"), + batch_size=int(os.environ.get("BATCH_SIZE", "4")), + ccip_model=os.environ.get("CCIP_MODEL", ""), + detector_level=os.environ.get("DETECTOR_LEVEL", "m"), + poll_idle_seconds=float(os.environ.get("POLL_IDLE_SECONDS", "10")), + ) diff --git a/agent/fc_agent/crops.py b/agent/fc_agent/crops.py new file mode 100644 index 0000000..df278c9 --- /dev/null +++ b/agent/fc_agent/crops.py @@ -0,0 +1,36 @@ +"""Crop primitive — vendored from backend/app/services/ml/crops.py so the agent +is self-contained. Keep in sync if the floor logic changes.""" +from PIL import Image + +MIN_CROP_FRACTION = 0.10 +MIN_CROP_PX = 64 + + +def crop_region( + img: Image.Image, + bbox: tuple[float, float, float, float], + *, + pad: float = 0.0, + min_fraction: float = MIN_CROP_FRACTION, + min_px: int = MIN_CROP_PX, +) -> Image.Image | None: + """Crop a NORMALIZED bbox (x, y, w, h in [0,1]); None if below the size + floor (max of a fraction-of-short-side and an absolute pixel floor).""" + iw, ih = img.size + x, y, w, h = bbox + px, py, pw, ph = x * iw, y * ih, w * iw, h * ih + if pad: + px -= pw * pad / 2.0 + py -= ph * pad / 2.0 + pw *= (1.0 + pad) + ph *= (1.0 + pad) + left = max(0, int(round(px))) + top = max(0, int(round(py))) + right = min(iw, int(round(px + pw))) + bottom = min(ih, int(round(py + ph))) + if right <= left or bottom <= top: + return None + floor = max(min_px, int(min_fraction * min(iw, ih))) + if min(right - left, bottom - top) < floor: + return None + return img.crop((left, top, right, bottom)).convert("RGB") diff --git a/agent/fc_agent/media.py b/agent/fc_agent/media.py new file mode 100644 index 0000000..e3efac5 --- /dev/null +++ b/agent/fc_agent/media.py @@ -0,0 +1,48 @@ +"""Image + video handling. Stills load directly; videos are sampled into frames +(ffmpeg) at the cadence FC sends — so a video becomes a bag of per-frame +instances, each with a timestamp.""" +import io +import os +import subprocess +import tempfile + +from PIL import Image + + +def is_video(mime: str) -> bool: + return bool(mime) and (mime.startswith("video/") or mime in {"image/gif"}) + + +def load_image(data: bytes) -> Image.Image: + return Image.open(io.BytesIO(data)).convert("RGB") + + +def sample_frames( + data: bytes, interval_seconds: float, max_frames: int +) -> list[tuple[float, Image.Image]]: + """Extract up to max_frames frames at one-every-interval_seconds via ffmpeg. + Returns [(timestamp_seconds, frame)]. Empty on failure (caller falls back).""" + interval = max(0.5, float(interval_seconds or 4.0)) + cap = max(1, int(max_frames or 64)) + with tempfile.TemporaryDirectory() as tmp: + src = os.path.join(tmp, "in") + with open(src, "wb") as fh: + fh.write(data) + pattern = os.path.join(tmp, "f_%05d.jpg") + try: + subprocess.run( + [ + "ffmpeg", "-nostdin", "-loglevel", "error", "-i", src, + "-vf", f"fps=1/{interval}", "-frames:v", str(cap), + "-q:v", "3", pattern, + ], + check=True, timeout=600, + ) + except (subprocess.SubprocessError, FileNotFoundError): + return [] + out: list[tuple[float, Image.Image]] = [] + names = sorted(n for n in os.listdir(tmp) if n.startswith("f_")) + for i, name in enumerate(names[:cap]): + with Image.open(os.path.join(tmp, name)) as im: + out.append((round(i * interval, 2), im.convert("RGB"))) + return out diff --git a/agent/fc_agent/models.py b/agent/fc_agent/models.py new file mode 100644 index 0000000..769e0ce --- /dev/null +++ b/agent/fc_agent/models.py @@ -0,0 +1,39 @@ +"""imgutils model wrappers — the figure DETECTOR + the CCIP EMBEDDER. + +⚠️ VERIFY ON FIRST RUN: the exact imgutils function names/signatures + the CCIP +model string can drift between dghs-imgutils releases. These are the two seams to +check against your installed version (`pip show dghs-imgutils`): + - detect_person(image, level=...) -> [((x0,y0,x1,y1), label, score), ...] + - ccip_extract_feature(image, model=...) -> a vector (768-d for caformer) +imgutils auto-downloads the ONNX models from HuggingFace on first use; GPU is +used when onnxruntime-gpu is installed. +""" +import numpy as np +from PIL import Image + + +def detect_figures(image: Image.Image, level: str = "m") -> list[tuple[tuple, float | None]]: + """Person/figure bounding boxes, NORMALIZED (x, y, w, h in [0,1]) + score. + Returns [] if detection finds nothing (caller falls back to whole-image).""" + from imgutils.detect import detect_person + + iw, ih = image.size + out = [] + for (x0, y0, x1, y1), _label, score in detect_person(image, level=level): + out.append(( + (x0 / iw, y0 / ih, (x1 - x0) / iw, (y1 - y0) / ih), + float(score), + )) + return out + + +def ccip_vector(image: Image.Image, model: str | None = None) -> list[float]: + """The CCIP identity embedding of a (cropped) character image, as a plain + float list ready to POST.""" + from imgutils.metrics import ccip_extract_feature + + feat = ( + ccip_extract_feature(image, model=model) + if model else ccip_extract_feature(image) + ) + return np.asarray(feat, dtype=np.float32).reshape(-1).tolist() diff --git a/agent/fc_agent/worker.py b/agent/fc_agent/worker.py new file mode 100644 index 0000000..7035c54 --- /dev/null +++ b/agent/fc_agent/worker.py @@ -0,0 +1,127 @@ +"""The lease → fetch → detect+embed → submit loop, with start/pause/stop control. + +Concurrency is 1 (one image at a time) so the GPU footprint stays small and a +stop frees the card promptly. Stop halts leasing + finishes the current item; +unprocessed leases expire and the server re-queues them — nothing is lost. +""" +import threading +import time + +from . import media, models +from .client import FcClient +from .config import Config +from .crops import crop_region + + +class Worker: + def __init__(self, cfg: Config): + self.cfg = cfg + self.client = FcClient(cfg.fc_url, cfg.token, cfg.agent_id) + self._state = "idle" # idle | running | paused | stopping + self._lock = threading.Lock() + self._thread: threading.Thread | None = None + self.processed = 0 + self.errors = 0 + self.current = None + + # --- control ----------------------------------------------------------- + def start(self): + with self._lock: + if self._state in ("running", "paused"): + self._state = "running" + return + self._state = "running" + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def pause(self): + with self._lock: + if self._state == "running": + self._state = "paused" + + def resume(self): + with self._lock: + if self._state == "paused": + self._state = "running" + + def stop(self): + with self._lock: + if self._state in ("running", "paused"): + self._state = "stopping" + + def status(self) -> dict: + with self._lock: + state = self._state + return { + "state": state, "processed": self.processed, + "errors": self.errors, "current": self.current, + } + + # --- loop -------------------------------------------------------------- + def _run(self): + while True: + with self._lock: + st = self._state + if st == "stopping": + break + if st == "paused": + time.sleep(1) + continue + try: + jobs = self.client.lease(self.cfg.batch_size) + except Exception: + time.sleep(self.cfg.poll_idle_seconds) + continue + if not jobs: + time.sleep(self.cfg.poll_idle_seconds) + continue + ids = [j["job_id"] for j in jobs] + for job in jobs: + with self._lock: + if self._state == "stopping": + break + self._process(job) + self.client.heartbeat(ids) # keep the rest of the batch alive + with self._lock: + self._state = "idle" + + def _process(self, job: dict): + self.current = job.get("image_id") + try: + data = self.client.fetch_image(job["image_url"]) + if media.is_video(job.get("mime", "")): + frames = media.sample_frames( + data, job.get("frame_interval_seconds", 4.0), + job.get("max_frames", 64), + ) or [(None, media.load_image(data))] + else: + frames = [(None, media.load_image(data))] + + regions = [] + ev = self.cfg.ccip_model or "ccip-default" + dv = f"person-{self.cfg.detector_level}" + for t, frame in frames: + figs = models.detect_figures(frame, self.cfg.detector_level) + if not figs: + figs = [((0.0, 0.0, 1.0, 1.0), None)] # whole-frame fallback + for bbox, score in figs: + crop = crop_region(frame, bbox) + if crop is None: + continue + vec = models.ccip_vector(crop, self.cfg.ccip_model or None) + regions.append({ + "kind": "figure", + "bbox": list(bbox), + "frame_time": t, + "score": score, + "ccip_embedding": vec, + "embedding_version": ev, + "detector_version": dv, + }) + self.client.submit(job["job_id"], regions, ["figure", "face"]) + self.processed += 1 + except Exception as exc: # noqa: BLE001 — report + move on + self.errors += 1 + self.client.fail(job["job_id"], str(exc)[:500]) + finally: + self.current = None diff --git a/agent/requirements.txt b/agent/requirements.txt new file mode 100644 index 0000000..267b99c --- /dev/null +++ b/agent/requirements.txt @@ -0,0 +1,11 @@ +# CCIP + figure detection (ONNX models, auto-downloaded from HuggingFace). +dghs-imgutils>=0.4 +# GPU inference for the ONNX models. Swap to onnxruntime (CPU) for a slow +# server-side fallback run. +onnxruntime-gpu +# Control surface + HTTP. +fastapi +uvicorn[standard] +requests +pillow +numpy diff --git a/alembic/versions/0061_image_region.py b/alembic/versions/0061_image_region.py new file mode 100644 index 0000000..b3af8a9 --- /dev/null +++ b/alembic/versions/0061_image_region.py @@ -0,0 +1,59 @@ +"""image_region: detected/proposed regions + their crop embeddings (#114) + +Storage backbone of the crop pipeline. A region = normalized bbox + the crop's +embedding (CCIP for face/figure → character id; SigLIP for concept regions → +head bag-of-embeddings). Also serves as grounded-tag bbox provenance. + +Revision ID: 0061 +Revises: 0060 +Create Date: 2026-06-29 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from pgvector.sqlalchemy import Vector + +revision: str = "0061" +down_revision: Union[str, None] = "0060" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_CCIP_DIM = 768 +_SIGLIP_DIM = 1152 + + +def upgrade() -> None: + op.create_table( + "image_region", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column( + "image_record_id", sa.Integer(), + sa.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False, + ), + sa.Column("kind", sa.String(length=16), nullable=False), + # Video/animated: source frame timestamp (seconds); NULL for stills. + sa.Column("frame_time", sa.Float(), nullable=True), + sa.Column("rx", sa.Float(), nullable=False), + sa.Column("ry", sa.Float(), nullable=False), + sa.Column("rw", sa.Float(), nullable=False), + sa.Column("rh", sa.Float(), nullable=False), + sa.Column("score", sa.Float(), nullable=True), + sa.Column("detector_version", sa.String(length=64), nullable=True), + sa.Column("crop_version", sa.String(length=64), nullable=True), + sa.Column("embedding_version", sa.String(length=128), nullable=True), + sa.Column("ccip_embedding", Vector(_CCIP_DIM), nullable=True), + sa.Column("siglip_embedding", Vector(_SIGLIP_DIM), nullable=True), + sa.Column( + "created_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + ) + op.create_index( + "ix_image_region_image_record_id", "image_region", ["image_record_id"], + ) + + +def downgrade() -> None: + op.drop_index("ix_image_region_image_record_id", table_name="image_region") + op.drop_table("image_region") diff --git a/alembic/versions/0062_gpu_job.py b/alembic/versions/0062_gpu_job.py new file mode 100644 index 0000000..a044995 --- /dev/null +++ b/alembic/versions/0062_gpu_job.py @@ -0,0 +1,55 @@ +"""gpu_job: the HTTP-leased GPU work queue for the desktop agent (#114) + +The agent stays HTTP-only — the server enqueues per-(image, task) jobs here and +the agent leases/submits over the web API; Redis/Postgres stay private. + +Revision ID: 0062 +Revises: 0061 +Create Date: 2026-06-29 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0062" +down_revision: Union[str, None] = "0061" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "gpu_job", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column( + "image_record_id", sa.Integer(), + sa.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False, + ), + sa.Column("task", sa.String(length=32), nullable=False), + sa.Column( + "status", sa.String(length=16), nullable=False, + server_default="pending", + ), + sa.Column("lease_token", sa.String(length=64), nullable=True), + sa.Column("leased_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("attempts", sa.Integer(), nullable=False, server_default="0"), + sa.Column("error", sa.Text(), nullable=True), + sa.Column( + "created_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.Column( + "updated_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + ) + op.create_index("ix_gpu_job_image_record_id", "gpu_job", ["image_record_id"]) + op.create_index("ix_gpu_job_status", "gpu_job", ["status"]) + + +def downgrade() -> None: + op.drop_index("ix_gpu_job_status", table_name="gpu_job") + op.drop_index("ix_gpu_job_image_record_id", table_name="gpu_job") + op.drop_table("gpu_job") diff --git a/backend/app/api/__init__.py b/backend/app/api/__init__.py index f39a966..cdc1442 100644 --- a/backend/app/api/__init__.py +++ b/backend/app/api/__init__.py @@ -20,11 +20,13 @@ def all_blueprints() -> list[Blueprint]: from .artist import artist_bp from .artists import artists_bp from .attachments import attachments_bp + from .ccip import ccip_bp from .cleanup import cleanup_bp from .credentials import credentials_bp from .downloads import downloads_bp from .extension import extension_bp from .gallery import gallery_bp + from .gpu import gpu_bp from .heads import heads_bp from .import_admin import import_admin_bp from .ml_admin import ml_admin_bp @@ -60,6 +62,8 @@ def all_blueprints() -> list[Blueprint]: aliases_bp, tag_eval_bp, heads_bp, + gpu_bp, + ccip_bp, ml_admin_bp, thumbnails_bp, sources_bp, diff --git a/backend/app/api/ccip.py b/backend/app/api/ccip.py new file mode 100644 index 0000000..31ff756 --- /dev/null +++ b/backend/app/api/ccip.py @@ -0,0 +1,106 @@ +"""CCIP / region observability API (#114) — read-only, analysis-shaped. + +So the work can be checked through an API as the agent fills in vectors: overall +coverage (regions by kind, how many images have figure CCIP vectors, which +characters have enough reference examples to match on) + a per-image drill-down +(its regions + the CCIP character matches it would get). Mirrors the heads +metrics endpoint; no GPU, just reads what's stored. +""" + +from quart import Blueprint, jsonify +from sqlalchemy import distinct, func, select + +from ..extensions import get_session +from ..models import ImageRegion, Tag, TagKind +from ..models.tag import image_tag +from ..services.ml.ccip import match_image + +ccip_bp = Blueprint("ccip", __name__, url_prefix="/api/ccip") + +_FIGURE_KINDS = ("face", "figure") + + +@ccip_bp.route("/overview", methods=["GET"]) +async def overview(): + async with get_session() as session: + by_kind = dict( + ( + await session.execute( + select(ImageRegion.kind, func.count()).group_by(ImageRegion.kind) + ) + ).all() + ) + images_with_figure_ccip = ( + await session.execute( + select(func.count(distinct(ImageRegion.image_record_id))) + .where(ImageRegion.kind.in_(_FIGURE_KINDS)) + .where(ImageRegion.ccip_embedding.is_not(None)) + ) + ).scalar_one() + # Per-character reference counts (no vectors loaded) — which characters + # have enough examples to match on. + ref_rows = ( + await session.execute( + select(image_tag.c.tag_id, Tag.name, func.count()) + .select_from(ImageRegion) + .join( + image_tag, + image_tag.c.image_record_id == ImageRegion.image_record_id, + ) + .join(Tag, Tag.id == image_tag.c.tag_id) + .where(Tag.kind == TagKind.character) + .where(ImageRegion.kind.in_(_FIGURE_KINDS)) + .where(ImageRegion.ccip_embedding.is_not(None)) + .group_by(image_tag.c.tag_id, Tag.name) + .order_by(func.count().desc()) + ) + ).all() + versions = [ + v for (v,) in ( + await session.execute( + select(distinct(ImageRegion.embedding_version)) + ) + ).all() if v + ] + return jsonify({ + "regions_by_kind": by_kind, + "images_with_figure_ccip": images_with_figure_ccip, + "characters_with_references": len(ref_rows), + "character_references": [ + {"tag_id": t, "name": n, "n_refs": c} for (t, n, c) in ref_rows + ], + "embedding_versions": versions, + }) + + +@ccip_bp.route("/images/", methods=["GET"]) +async def image_detail(image_id: int): + """An image's stored regions + the CCIP character matches it would get — + for spot-checking the agent's output + the matcher.""" + async with get_session() as session: + regions = ( + await session.execute( + select(ImageRegion) + .where(ImageRegion.image_record_id == image_id) + .order_by(ImageRegion.id) + ) + ).scalars().all() + matches = await match_image(session, image_id) + return jsonify({ + "image_id": image_id, + "regions": [ + { + "id": r.id, + "kind": r.kind, + "bbox": [r.rx, r.ry, r.rw, r.rh], + "frame_time": r.frame_time, + "score": r.score, + "detector_version": r.detector_version, + "embedding_version": r.embedding_version, + "has_ccip": r.ccip_embedding is not None, + "has_siglip": r.siglip_embedding is not None, + } + for r in regions + ], + "ccip_matches": matches, + }) diff --git a/backend/app/api/gpu.py b/backend/app/api/gpu.py new file mode 100644 index 0000000..0a9b46a --- /dev/null +++ b/backend/app/api/gpu.py @@ -0,0 +1,198 @@ +"""GPU-job API (#114): the HTTP surface the desktop agent pulls work from. + +The agent stays HTTP-only — it leases jobs, fetches image pixels via the normal +FC image URLs, and submits embeddings/regions back, all over this API. Redis and +Postgres are never exposed. The agent endpoints are gated by a bearer token +(Authorization: Bearer ) stored in AppSetting; the admin endpoints +(token / backfill / status) ride the browser session like the rest of FC's +homelab admin. +""" + +import secrets + +from quart import Blueprint, jsonify, request +from sqlalchemy import func, select +from sqlalchemy.dialects.postgresql import insert as pg_insert + +from ..extensions import get_session +from ..models import AppSetting, GpuJob, ImageRecord, MLSettings +from ..services.gallery_service import image_url +from ..services.ml.gpu_jobs import GpuJobService +from ..services.ml.regions import RegionService + +gpu_bp = Blueprint("gpu", __name__, url_prefix="/api/gpu") + +_TOKEN_KEY = "gpu_agent_token" + + +def _bearer() -> str | None: + h = request.headers.get("Authorization", "") + return h[7:].strip() if h.startswith("Bearer ") else None + + +async def _agent_authed(session) -> bool: + supplied = _bearer() + if not supplied: + return False + stored = ( + await session.execute( + select(AppSetting.value).where(AppSetting.key == _TOKEN_KEY) + ) + ).scalar_one_or_none() + return stored is not None and secrets.compare_digest(supplied, stored) + + +# --- Admin (browser): token + backfill + status ------------------------- + +@gpu_bp.route("/token", methods=["GET"]) +async def get_token(): + async with get_session() as session: + tok = ( + await session.execute( + select(AppSetting.value).where(AppSetting.key == _TOKEN_KEY) + ) + ).scalar_one_or_none() + return jsonify({"token": tok, "configured": tok is not None}) + + +@gpu_bp.route("/token/rotate", methods=["POST"]) +async def rotate_token(): + token = secrets.token_urlsafe(32) + async with get_session() as session: + await session.execute( + pg_insert(AppSetting) + .values(key=_TOKEN_KEY, value=token) + .on_conflict_do_update(index_elements=["key"], set_={"value": token}) + ) + await session.commit() + return jsonify({"token": token}) + + +@gpu_bp.route("/status", methods=["GET"]) +async def status(): + async with get_session() as session: + rows = ( + await session.execute( + select(GpuJob.status, func.count()).group_by(GpuJob.status) + ) + ).all() + counts = dict(rows) + return jsonify({ + "pending": counts.get("pending", 0), + "leased": counts.get("leased", 0), + "done": counts.get("done", 0), + "error": counts.get("error", 0), + }) + + +@gpu_bp.route("/backfill", methods=["POST"]) +async def backfill(): + """Enqueue a job for every image that doesn't already have one for `task`.""" + body = await request.get_json(silent=True) or {} + task = str(body.get("task") or "ccip") + from ..tasks.ml import enqueue_gpu_backfill + + r = enqueue_gpu_backfill.delay(task) + return jsonify({"celery_task_id": r.id, "task": task}), 202 + + +# --- Agent (bearer token): lease / submit / heartbeat / fail ------------ + +@gpu_bp.route("/jobs/lease", methods=["POST"]) +async def lease(): + body = await request.get_json(silent=True) or {} + agent_id = str(body.get("agent_id") or "agent") + try: + batch = min(max(int(body.get("batch_size", 8)), 1), 64) + except (TypeError, ValueError): + batch = 8 + async with get_session() as session: + if not await _agent_authed(session): + return jsonify({"error": "unauthorized"}), 401 + jobs = await GpuJobService(session).lease(agent_id, batch_size=batch) + ml = ( + await session.execute(select(MLSettings).where(MLSettings.id == 1)) + ).scalar_one() + # image rows for url/mime in one shot + ids = [j.image_record_id for j in jobs] + imgs = { + i.id: i for i in ( + await session.execute( + select(ImageRecord).where(ImageRecord.id.in_(ids)) + ) + ).scalars() + } if ids else {} + await session.commit() + out = [] + for j in jobs: + img = imgs.get(j.image_record_id) + if img is None: + continue + out.append({ + "job_id": j.id, + "image_id": j.image_record_id, + "task": j.task, + "mime": img.mime, + "image_url": image_url(img.path), + # For video/animated: the agent samples at this cadence. + "frame_interval_seconds": ml.video_frame_interval_seconds, + "max_frames": ml.video_max_frames, + }) + return jsonify({"jobs": out}) + + +@gpu_bp.route("/jobs/heartbeat", methods=["POST"]) +async def heartbeat(): + body = await request.get_json(silent=True) or {} + agent_id = str(body.get("agent_id") or "agent") + job_ids = [int(x) for x in (body.get("job_ids") or [])] + async with get_session() as session: + if not await _agent_authed(session): + return jsonify({"error": "unauthorized"}), 401 + n = await GpuJobService(session).heartbeat(agent_id, job_ids) + await session.commit() + return jsonify({"extended": n}) + + +@gpu_bp.route("/jobs/submit", methods=["POST"]) +async def submit(): + """Store a job's regions + close it. regions: [{kind, bbox:[x,y,w,h], + frame_time?, score?, *_version?, ccip_embedding?, siglip_embedding?}]. + replace_kinds defaults to the kinds present in the submitted regions.""" + body = await request.get_json(silent=True) or {} + agent_id = str(body.get("agent_id") or "agent") + job_id = body.get("job_id") + regions = body.get("regions") or [] + if job_id is None: + return jsonify({"error": "job_id required"}), 400 + kinds = body.get("replace_kinds") or sorted({r["kind"] for r in regions}) + async with get_session() as session: + if not await _agent_authed(session): + return jsonify({"error": "unauthorized"}), 401 + job = await session.get(GpuJob, int(job_id)) + if job is None or job.status != "leased" or job.lease_token != agent_id: + return jsonify({"error": "lease_invalid"}), 409 + if kinds: + await RegionService(session).replace_regions( + job.image_record_id, kinds, regions + ) + await GpuJobService(session).complete(agent_id, int(job_id)) + await session.commit() + return jsonify({"ok": True, "stored": len(regions)}) + + +@gpu_bp.route("/jobs/fail", methods=["POST"]) +async def fail(): + body = await request.get_json(silent=True) or {} + agent_id = str(body.get("agent_id") or "agent") + job_id = body.get("job_id") + if job_id is None: + return jsonify({"error": "job_id required"}), 400 + async with get_session() as session: + if not await _agent_authed(session): + return jsonify({"error": "unauthorized"}), 401 + ok = await GpuJobService(session).fail( + agent_id, int(job_id), str(body.get("error") or "") + ) + await session.commit() + return jsonify({"ok": ok}) diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 681adb7..087ff0b 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -8,6 +8,7 @@ from .base import Base from .credential import Credential from .download_event import DownloadEvent from .external_link import ExternalLink +from .gpu_job import GpuJob from .head_auto_apply_run import HeadAutoApplyRun from .head_metric import HeadMetric from .head_metrics_snapshot import HeadMetricsSnapshot @@ -15,6 +16,7 @@ from .head_training_run import HeadTrainingRun from .image_prediction import ImagePrediction from .image_provenance import ImageProvenance from .image_record import ImageRecord +from .image_region import ImageRegion from .import_batch import ImportBatch from .import_settings import ImportSettings from .import_task import ImportTask @@ -60,11 +62,13 @@ __all__ = [ "ImageRecord", "ImagePrediction", "ImageProvenance", + "ImageRegion", "Tag", "TagKind", "image_tag", "DownloadEvent", "ExternalLink", + "GpuJob", "ImportBatch", "ImportTask", "ImportSettings", diff --git a/backend/app/models/gpu_job.py b/backend/app/models/gpu_job.py new file mode 100644 index 0000000..a8ed5c1 --- /dev/null +++ b/backend/app/models/gpu_job.py @@ -0,0 +1,50 @@ +"""GpuJob — a unit of GPU work the desktop agent pulls over HTTP (#114). + +The durable work list that lets the agent stay HTTP-only: the server enqueues a +job per (image, task) — e.g. detect figures + CCIP-embed — and the agent LEASES a +batch, computes on its GPU, then SUBMITS results, all over the already-exposed web +API. Redis/Postgres stay private. A lease has an expiry; the lease query itself +re-claims expired leases (agent died / stopped mid-batch), so the queue is +self-healing without a separate sweep. One job is per ITEM; the agent fans a +VIDEO out into per-frame instances internally (see image_region.frame_time). + +State: pending → leased → done | error (a failure under the attempt cap returns to +pending for another agent). +""" + +from datetime import datetime + +from sqlalchemy import DateTime, ForeignKey, Integer, String, Text, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class GpuJob(Base): + __tablename__ = "gpu_job" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + image_record_id: Mapped[int] = mapped_column( + ForeignKey("image_record.id", ondelete="CASCADE"), index=True + ) + # What to compute, e.g. 'ccip' (detect figures + CCIP-embed) or 'siglip_region'. + task: Mapped[str] = mapped_column(String(32), nullable=False) + status: Mapped[str] = mapped_column( + String(16), nullable=False, default="pending", index=True + ) + # pending | leased | done | error + lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True) + leased_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + lease_expires_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + error: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/backend/app/models/image_region.py b/backend/app/models/image_region.py new file mode 100644 index 0000000..58cfa8c --- /dev/null +++ b/backend/app/models/image_region.py @@ -0,0 +1,62 @@ +"""ImageRegion — a detected/proposed sub-region of an image + its crop embedding. + +The storage backbone of the crop pipeline (#114). A region is a normalized bbox +plus the embedding of its crop: +- kind='face' / 'figure' → embedded by CCIP for cross-artist character identity. +- kind='concept' → embedded by SigLIP, a localized instance for a concept head's + bag-of-embeddings (a concept is "present if ANY instance matches"). +One row carries the embedding appropriate to its kind (the other is null). The +bbox doubles as grounded-tag provenance (hover a tag → highlight its region; a +wrong box is a precise negative). The GPU agent writes these via the job API; +the few-shot character matcher + bag scorer read them — both server-side, no GPU. +""" + +from datetime import datetime + +from pgvector.sqlalchemy import Vector +from sqlalchemy import DateTime, Float, ForeignKey, Integer, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + +CCIP_DIM = 768 # deepghs/imgutils CCIP character embedding +SIGLIP_DIM = 1152 # matches image_record.siglip_embedding + + +class ImageRegion(Base): + __tablename__ = "image_region" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + image_record_id: Mapped[int] = mapped_column( + ForeignKey("image_record.id", ondelete="CASCADE"), index=True + ) + # 'frame' (a whole video frame → SigLIP bag) | 'face' | 'figure' (→ CCIP + # character id) | 'concept' (→ SigLIP head bag). + kind: Mapped[str] = mapped_column(String(16), nullable=False) + # For video/animated media: the source frame's timestamp in SECONDS. NULL for + # static images. Lets a video be a BAG of per-frame instances (fixes the + # mean-embedding muddle) + grounds a tag to "appears at 0:42". + frame_time: Mapped[float | None] = mapped_column(Float, nullable=True) + # Normalized bbox in [0,1]: top-left (rx, ry) + size (rw, rh). Named rx/ry/… + # rather than x/y/by to dodge SQL keyword ambiguity ('by'). + rx: Mapped[float] = mapped_column(Float, nullable=False) + ry: Mapped[float] = mapped_column(Float, nullable=False) + rw: Mapped[float] = mapped_column(Float, nullable=False) + rh: Mapped[float] = mapped_column(Float, nullable=False) + # Proposer/detector confidence (null for deterministic proposers). + score: Mapped[float | None] = mapped_column(Float, nullable=True) + # Version stamps so a re-detect / re-crop / re-embed can be gated (compute + # once; only redo when the producing model version changes). + detector_version: Mapped[str | None] = mapped_column(String(64), nullable=True) + crop_version: Mapped[str | None] = mapped_column(String(64), nullable=True) + embedding_version: Mapped[str | None] = mapped_column(String(128), nullable=True) + # Exactly one is set, per kind. + ccip_embedding: Mapped[list[float] | None] = mapped_column( + Vector(CCIP_DIM), nullable=True + ) + siglip_embedding: Mapped[list[float] | None] = mapped_column( + Vector(SIGLIP_DIM), nullable=True + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/backend/app/services/ml/ccip.py b/backend/app/services/ml/ccip.py new file mode 100644 index 0000000..406fbed --- /dev/null +++ b/backend/app/services/ml/ccip.py @@ -0,0 +1,120 @@ +"""CCIP few-shot character matcher (#114) — server-side, numpy on stored vectors. + +CCIP is a FROZEN identity embedding; we don't train it. Instead the operator's +tagged characters become reference prototypes: a character tag's references are +the CCIP vectors of figure/face regions on images carrying that tag. To suggest +characters for a new image, we compare its figure-region CCIP vectors to every +character's references (multi-prototype: best match over a character's examples) +and surface the ones that clear a similarity threshold. No GPU here — the agent +already produced the vectors; this is cosine matching on what's stored. + +v1 uses cosine similarity on the raw CCIP vectors with a tunable threshold; the +exact CCIP difference metric/threshold gets validated against the model during +the hands-on eval. numpy is imported lazily (API worker has it via pgvector). +""" + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import ImageRegion, Tag, TagKind +from ...models.tag import image_tag + +# Cosine-similarity floor to call a figure the same character. Conservative +# default; tune from real matches (CCIP same-char clusters tightly). +DEFAULT_SIM_THRESHOLD = 0.75 +_FIGURE_KINDS = ("face", "figure") + + +def _l2norm(mat, np): + n = np.linalg.norm(mat, axis=1, keepdims=True) + n[n == 0] = 1.0 + return mat / n + + +async def character_references(session: AsyncSession) -> dict[int, list]: + """Per character-tag CCIP reference vectors: figure/face-region CCIP + embeddings on images that carry that character tag (the operator's examples). + Multi-prototype — several vectors per character.""" + rows = ( + await session.execute( + select(image_tag.c.tag_id, ImageRegion.ccip_embedding) + .select_from(ImageRegion) + .join( + image_tag, + image_tag.c.image_record_id == ImageRegion.image_record_id, + ) + .join(Tag, Tag.id == image_tag.c.tag_id) + .where(Tag.kind == TagKind.character) + .where(ImageRegion.kind.in_(_FIGURE_KINDS)) + .where(ImageRegion.ccip_embedding.is_not(None)) + ) + ).all() + refs: dict[int, list] = {} + for tag_id, vec in rows: + refs.setdefault(tag_id, []).append(vec) + return refs + + +async def _tag_names(session: AsyncSession, tag_ids: list[int]) -> dict[int, str]: + if not tag_ids: + return {} + return dict( + ( + await session.execute( + select(Tag.id, Tag.name).where(Tag.id.in_(tag_ids)) + ) + ).all() + ) + + +async def match_image( + session: AsyncSession, image_id: int, threshold: float = DEFAULT_SIM_THRESHOLD +) -> list[dict]: + """Character suggestions for one image from its figure-region CCIP vectors: + [{tag_id, name, category:'character', score, source:'ccip'}], ranked. + Already-applied character tags are excluded. Empty if the image has no figure + CCIP vectors or no character references exist yet.""" + import numpy as np + + qvecs = ( + await session.execute( + select(ImageRegion.ccip_embedding).where( + ImageRegion.image_record_id == image_id, + ImageRegion.kind.in_(_FIGURE_KINDS), + ImageRegion.ccip_embedding.is_not(None), + ) + ) + ).scalars().all() + if not qvecs: + return [] + refs = await character_references(session) + if not refs: + return [] + applied = set( + ( + await session.execute( + select(image_tag.c.tag_id).where( + image_tag.c.image_record_id == image_id + ) + ) + ).scalars() + ) + names = await _tag_names(session, [t for t in refs if t not in applied]) + + Q = _l2norm(np.vstack([np.asarray(v, dtype=np.float32) for v in qvecs]), np) + out = [] + for tag_id, vecs in refs.items(): + if tag_id in applied: + continue + R = _l2norm(np.vstack([np.asarray(v, dtype=np.float32) for v in vecs]), np) + best = float((Q @ R.T).max()) # best (query figure, reference) cosine + if best >= threshold: + out.append({ + "tag_id": tag_id, + "name": names.get(tag_id, str(tag_id)), + "category": "character", + "score": round(best, 4), + "source": "ccip", + }) + out.sort(key=lambda d: d["score"], reverse=True) + return out diff --git a/backend/app/services/ml/crops.py b/backend/app/services/ml/crops.py new file mode 100644 index 0000000..ffb7b6a --- /dev/null +++ b/backend/app/services/ml/crops.py @@ -0,0 +1,73 @@ +"""Shared crop primitive for the region/crop pipeline (#114). + +One model- and transport-agnostic function sits at the trunk of both crop jobs: +- CCIP characters: a face/figure detector proposes regions → crop → CCIP-embed. +- SigLIP concepts: head-guided / saliency proposes regions → crop → SigLIP-embed. +Only the PROPOSER (where to crop) and the EMBEDDER (what to run) differ; the crop +itself — including the lower-bound size floor below which a region is too small to +embed reliably — is identical, so it lives here and both jobs call it. + +The actual detector + embedders run in the GPU agent; this is pure Pillow so it's +importable + testable anywhere (and the agent imports it for the crop step). +""" + +from __future__ import annotations + +from PIL import Image + +# Size floor: a region must be at least this big on its SHORTER edge to be worth +# embedding — a smaller crop is a blurry upscale carrying little real signal, and +# unbounded tiny crops would explode the bag. Expressed as BOTH a fraction of the +# image's short side and an absolute pixel floor; the larger of the two wins. +MIN_CROP_FRACTION = 0.10 +MIN_CROP_PX = 64 + + +def _to_pixels(bbox: tuple[float, float, float, float], w: int, h: int): + """Normalized (x, y, w, h) in [0,1] → pixel (x, y, w, h).""" + x, y, bw, bh = bbox + return x * w, y * h, bw * w, bh * h + + +def crop_region( + img: Image.Image, + bbox: tuple[float, float, float, float], + *, + pad: float = 0.0, + min_fraction: float = MIN_CROP_FRACTION, + min_px: int = MIN_CROP_PX, + out_size: int | None = None, +) -> Image.Image | None: + """Crop a NORMALIZED bbox (x, y, w, h in [0,1]) from img. + + - pad: grow the box by this fraction on each side (e.g. 0.15 = +15% context), + clamped to the image bounds. + - Returns None when the resulting region is below the size floor (too small to + embed reliably) — the caller skips embedding it. + - out_size: if given, resize the crop to out_size×out_size; otherwise return + the raw crop and let the embedder do its own preprocessing. + """ + iw, ih = img.size + px, py, pw, ph = _to_pixels(bbox, iw, ih) + + if pad: + px -= pw * pad / 2.0 + py -= ph * pad / 2.0 + pw *= (1.0 + pad) + ph *= (1.0 + pad) + + left = max(0, int(round(px))) + top = max(0, int(round(py))) + right = min(iw, int(round(px + pw))) + bottom = min(ih, int(round(py + ph))) + if right <= left or bottom <= top: + return None + + floor = max(min_px, int(min_fraction * min(iw, ih))) + if min(right - left, bottom - top) < floor: + return None + + crop = img.crop((left, top, right, bottom)).convert("RGB") + if out_size: + crop = crop.resize((out_size, out_size)) + return crop diff --git a/backend/app/services/ml/gpu_jobs.py b/backend/app/services/ml/gpu_jobs.py new file mode 100644 index 0000000..b760da9 --- /dev/null +++ b/backend/app/services/ml/gpu_jobs.py @@ -0,0 +1,134 @@ +"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail. + +Backs the HTTP API the desktop agent pulls work from. The lease claims pending +OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents (or a +retry after an agent died) never grab the same job and the queue self-heals +without a separate recovery sweep. Result-writing (regions) is done by the API +handler via RegionService; complete() just closes the job. +""" + +from datetime import UTC, datetime, timedelta + +from sqlalchemy import and_, or_, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import GpuJob + +DEFAULT_LEASE_TTL = 300 # seconds an agent holds a job before it can be re-leased +DEFAULT_BATCH = 8 +MAX_ATTEMPTS = 3 + + +class GpuJobService: + def __init__(self, session: AsyncSession): + self.session = session + + async def enqueue(self, image_id: int, task: str) -> GpuJob | None: + """Queue a (image, task) job. Idempotent: returns None if one is already + pending/leased for the same pair (no duplicate work).""" + dup = ( + await self.session.execute( + select(GpuJob.id).where( + GpuJob.image_record_id == image_id, + GpuJob.task == task, + GpuJob.status.in_(["pending", "leased"]), + ) + ) + ).first() + if dup: + return None + job = GpuJob(image_record_id=image_id, task=task, status="pending") + self.session.add(job) + await self.session.flush() + return job + + async def lease( + self, token: str, batch_size: int = DEFAULT_BATCH, ttl: int = DEFAULT_LEASE_TTL + ) -> list[GpuJob]: + """Claim up to batch_size pending (or expired-leased) jobs for `token`.""" + now = datetime.now(UTC) + picked = ( + await self.session.execute( + select(GpuJob.id) + .where( + or_( + GpuJob.status == "pending", + and_( + GpuJob.status == "leased", + GpuJob.lease_expires_at < now, + ), + ) + ) + .order_by(GpuJob.id) + .limit(batch_size) + .with_for_update(skip_locked=True) + ) + ).scalars().all() + if not picked: + return [] + await self.session.execute( + update(GpuJob) + .where(GpuJob.id.in_(picked)) + .values( + status="leased", lease_token=token, leased_at=now, + lease_expires_at=now + timedelta(seconds=ttl), + attempts=GpuJob.attempts + 1, updated_at=now, + ) + ) + # populate_existing: overwrite identity-map copies with the post-UPDATE + # values so the returned jobs reflect the new lease/attempts, not stale + # pre-lease state. + return list( + ( + await self.session.execute( + select(GpuJob) + .where(GpuJob.id.in_(picked)) + .order_by(GpuJob.id) + .execution_options(populate_existing=True) + ) + ).scalars() + ) + + async def heartbeat( + self, token: str, job_ids: list[int], ttl: int = DEFAULT_LEASE_TTL + ) -> int: + """Extend the lease on the agent's in-flight jobs. Returns rows touched.""" + now = datetime.now(UTC) + res = await self.session.execute( + update(GpuJob) + .where( + GpuJob.id.in_(job_ids), + GpuJob.lease_token == token, + GpuJob.status == "leased", + ) + .values(lease_expires_at=now + timedelta(seconds=ttl), updated_at=now) + ) + return res.rowcount or 0 + + async def complete(self, token: str, job_id: int) -> bool: + """Close a leased job (after its results were stored). False if the job + isn't leased by this token (a stale/expired submit).""" + job = await self.session.get(GpuJob, job_id) + if job is None or job.status != "leased" or job.lease_token != token: + return False + job.status = "done" + job.lease_token = None + job.lease_expires_at = None + job.error = None + job.updated_at = datetime.now(UTC) + return True + + async def fail(self, token: str, job_id: int, error: str) -> bool: + """Report a failure: re-queue (pending) until MAX_ATTEMPTS, then 'error'.""" + job = await self.session.get(GpuJob, job_id) + if job is None or job.lease_token != token: + return False + if job.attempts >= MAX_ATTEMPTS: + job.status = "error" + else: + job.status = "pending" + job.lease_token = None + job.lease_expires_at = None + job.error = (error or "")[:1000] + job.updated_at = datetime.now(UTC) + return True diff --git a/backend/app/services/ml/regions.py b/backend/app/services/ml/regions.py new file mode 100644 index 0000000..7008536 --- /dev/null +++ b/backend/app/services/ml/regions.py @@ -0,0 +1,59 @@ +"""Region read/write for the crop pipeline (#114). + +The GPU agent's results endpoint calls replace_regions() to store a freshly +detected/embedded set; the character matcher + concept-bag scorer read via +get_regions(). Replacement is scoped BY KIND so the figure pipeline and the +concept pipeline don't clobber each other. +""" + +from typing import Any + +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import ImageRegion + + +class RegionService: + def __init__(self, session: AsyncSession): + self.session = session + + async def get_regions( + self, image_id: int, kinds: list[str] | None = None + ) -> list[ImageRegion]: + stmt = select(ImageRegion).where(ImageRegion.image_record_id == image_id) + if kinds: + stmt = stmt.where(ImageRegion.kind.in_(kinds)) + return list( + (await self.session.execute(stmt.order_by(ImageRegion.id))).scalars() + ) + + async def replace_regions( + self, image_id: int, kinds: list[str], regions: list[dict[str, Any]] + ) -> int: + """Replace this image's regions OF THE GIVEN KINDS with `regions` (a + re-detect/re-propose supersedes the prior set without touching other + kinds). Each region dict: {kind, bbox:(x,y,w,h), score?, detector_version?, + crop_version?, embedding_version?, ccip_embedding?, siglip_embedding?}. + Returns the number inserted.""" + await self.session.execute( + delete(ImageRegion) + .where(ImageRegion.image_record_id == image_id) + .where(ImageRegion.kind.in_(kinds)) + ) + n = 0 + for r in regions: + rx, ry, rw, rh = r["bbox"] + self.session.add(ImageRegion( + image_record_id=image_id, kind=r["kind"], + frame_time=r.get("frame_time"), + rx=rx, ry=ry, rw=rw, rh=rh, + score=r.get("score"), + detector_version=r.get("detector_version"), + crop_version=r.get("crop_version"), + embedding_version=r.get("embedding_version"), + ccip_embedding=r.get("ccip_embedding"), + siglip_embedding=r.get("siglip_embedding"), + )) + n += 1 + return n diff --git a/backend/app/services/ml/suggestions.py b/backend/app/services/ml/suggestions.py index 1cb306f..71fbd24 100644 --- a/backend/app/services/ml/suggestions.py +++ b/backend/app/services/ml/suggestions.py @@ -16,6 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from ...models import ImageRecord, TagSuggestionRejection from ...models.tag import image_tag +from .ccip import match_image as ccip_match_image from .heads import score_image @@ -27,7 +28,7 @@ class Suggestion: display_name: str category: str score: float - source: str # 'head' (Camie 'tagger'/'centroid' sources removed in v2) + source: str # 'head' | 'ccip' | 'both' (Camie tagger/centroid removed in v2) creates_new_tag: bool # raw_name = the booru model vocab key behind this suggestion. It's the key # an alias MUST be stored under (resolution looks up the raw key), so the @@ -92,19 +93,39 @@ class SuggestionService: hits = await score_image( self.session, image_id, threshold_override=threshold_override ) + # CCIP character matches OVERLAY the SigLIP character heads — a + # complementary, identity-specialized signal with different failure modes + # (CCIP needs a detected figure; heads work whole-image). Merged by tag: + # 'both' when they corroborate, taking the higher score. + ccip_hits = await ccip_match_image(self.session, image_id) + + merged: dict[tuple[str, int], dict] = {} + for h in hits: + merged[(h["category"], h["tag_id"])] = { + "name": h["name"], "score": h["score"], "source": "head", + } + for c in ccip_hits: + key = ("character", c["tag_id"]) + ex = merged.get(key) + if ex is not None: + ex["source"] = "both" + ex["score"] = max(ex["score"], c["score"]) + else: + merged[key] = { + "name": c["name"], "score": c["score"], "source": "ccip", + } result = SuggestionList() - for h in hits: - tag_id = h["tag_id"] + for (cat, tag_id), m in merged.items(): if tag_id in applied: continue - result.by_category.setdefault(h["category"], []).append( + result.by_category.setdefault(cat, []).append( Suggestion( canonical_tag_id=tag_id, - display_name=h["name"], - category=h["category"], - score=h["score"], - source="head", + display_name=m["name"], + category=cat, + score=m["score"], + source=m["source"], creates_new_tag=False, rejected=tag_id in rejected, ) diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index bc22c5b..ecdfcbd 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -738,3 +738,33 @@ def scheduled_apply_head_tags() -> str: run_id = run.id apply_head_tags.delay(run_id) return "dispatched" + + +@celery.task(name="backend.app.tasks.ml.enqueue_gpu_backfill") +def enqueue_gpu_backfill(task_name: str) -> int: + """Enqueue a gpu_job for every image that doesn't already have one for + `task_name` (one INSERT…SELECT, so it scales to a full library). The desktop + agent drains the queue over HTTP. Returns the number enqueued.""" + from sqlalchemy import exists, insert, literal + from sqlalchemy import select as sa_select + + from ..models import GpuJob, ImageRecord + + SessionLocal = _sync_session_factory() + with SessionLocal() as session: + already = exists().where( + GpuJob.image_record_id == ImageRecord.id, + GpuJob.task == task_name, + GpuJob.status.in_(["pending", "leased", "done"]), + ) + sel = sa_select( + ImageRecord.id, literal(task_name), literal("pending") + ).where(~already) + # RETURNING + count: result.rowcount is unreliable for INSERT…SELECT. + rows = session.execute( + insert(GpuJob) + .from_select(["image_record_id", "task", "status"], sel) + .returning(GpuJob.id) + ).fetchall() + session.commit() + return len(rows) diff --git a/frontend/src/components/modal/TagAutocomplete.vue b/frontend/src/components/modal/TagAutocomplete.vue index 43ab0e1..65f3527 100644 --- a/frontend/src/components/modal/TagAutocomplete.vue +++ b/frontend/src/components/modal/TagAutocomplete.vue @@ -72,7 +72,7 @@ - Create "{{ parsedName }}" as {{ parsedKind }} + {{ createLabel }} @@ -178,14 +178,34 @@ watch(query, () => { }, 200) }) +// A same-name character ALREADY exists. Characters are unique by +// (name, kind, fandom), so this is still a valid distinct tag in another fandom. +const sameNameCharExists = computed(() => + parsedKind.value === 'character' && + hits.value.some(h => + h.kind === 'character' && h.name.toLowerCase() === parsedName.value.toLowerCase(), + ), +) + const allowCreate = computed(() => { const q = parsedName.value if (!q) return false + // Characters disambiguate by fandom, so a same-named character in a DIFFERENT + // fandom is a valid new tag — always offer Create (the fandom picker resolves + // it; find_or_create is idempotent if you re-pick the same fandom). Other + // kinds are unique by (name, kind): an exact match means it already exists. + if (parsedKind.value === 'character') return true return !hits.value.some(h => h.name.toLowerCase() === q.toLowerCase() && h.kind === parsedKind.value, ) }) +const createLabel = computed(() => + sameNameCharExists.value + ? `Create another "${parsedName.value}" character (different fandom)` + : `Create "${parsedName.value}" as ${parsedKind.value}`, +) + function scorePct (s) { return `${Math.round(s.score * 100)}%` } // This image's suggestions that match the typed query, minus any the server diff --git a/frontend/src/components/settings/GpuAgentCard.vue b/frontend/src/components/settings/GpuAgentCard.vue new file mode 100644 index 0000000..d4a3d96 --- /dev/null +++ b/frontend/src/components/settings/GpuAgentCard.vue @@ -0,0 +1,166 @@ + + + + + diff --git a/frontend/src/components/settings/MaintenancePanel.vue b/frontend/src/components/settings/MaintenancePanel.vue index 44dafe6..77a8035 100644 --- a/frontend/src/components/settings/MaintenancePanel.vue +++ b/frontend/src/components/settings/MaintenancePanel.vue @@ -27,6 +27,7 @@
+ @@ -54,6 +55,7 @@ import MissingFileRepairCard from './MissingFileRepairCard.vue' import DbMaintenanceCard from './DbMaintenanceCard.vue' import MLThresholdSliders from './MLThresholdSliders.vue' import HeadsCard from './HeadsCard.vue' +import GpuAgentCard from './GpuAgentCard.vue' import AllowlistTable from './AllowlistTable.vue' import AliasTable from './AliasTable.vue' import TagEvalCard from './TagEvalCard.vue' diff --git a/frontend/src/stores/gpu.js b/frontend/src/stores/gpu.js new file mode 100644 index 0000000..6c677c0 --- /dev/null +++ b/frontend/src/stores/gpu.js @@ -0,0 +1,33 @@ +import { defineStore } from 'pinia' + +import { useApi } from '../composables/useApi.js' + +// GPU agent control surface (#114): the FC-side admin for the desktop agent — +// the bearer token it authenticates with, the job-queue depth, and the backfill +// trigger. The agent itself talks to /api/gpu/jobs/* over HTTP; nothing here +// touches Redis/Postgres directly. +export const useGpuStore = defineStore('gpu', () => { + const api = useApi() + + // { token: , configured: bool } + async function token() { + return await api.get('/api/gpu/token') + } + + // Generate a fresh token (invalidates the old one). Returns { token }. + async function rotateToken() { + return await api.post('/api/gpu/token/rotate') + } + + // { pending, leased, done, error } + async function status() { + return await api.get('/api/gpu/status') + } + + // Enqueue a job per image lacking one for `task` (the agent drains it). + async function backfill(task = 'ccip') { + return await api.post('/api/gpu/backfill', { body: { task } }) + } + + return { token, rotateToken, status, backfill } +}) diff --git a/frontend/src/views/ExploreView.vue b/frontend/src/views/ExploreView.vue index 5a3052f..a196dd1 100644 --- a/frontend/src/views/ExploreView.vue +++ b/frontend/src/views/ExploreView.vue @@ -90,12 +90,22 @@
- +
{{ store.anchor.artist?.name || 'Unknown artist' }}
@@ -129,6 +139,7 @@ import { useModalStore } from '../stores/modal.js' import { useHeadTraining } from '../composables/useHeadTraining.js' import { isTextEntry } from '../utils/textEntry.js' import ImageCanvas from '../components/modal/ImageCanvas.vue' +import VideoCanvas from '../components/modal/VideoCanvas.vue' import ImageMetaBar from '../components/modal/ImageMetaBar.vue' import ProvenancePanel from '../components/modal/ProvenancePanel.vue' import TagPanel from '../components/modal/TagPanel.vue' @@ -140,6 +151,7 @@ const store = useExploreStore() const modal = useModalStore() const anchorId = computed(() => route.params.imageId || null) +const isVideo = computed(() => !!store.anchor?.mime?.startsWith('video/')) const seeding = ref(false) const seedError = ref(null) const tagPanelRef = ref(null) diff --git a/tests/test_api_ccip.py b/tests/test_api_ccip.py new file mode 100644 index 0000000..d41167a --- /dev/null +++ b/tests/test_api_ccip.py @@ -0,0 +1,72 @@ +"""CCIP/region observability API (#114) — coverage overview + per-image detail.""" +import pytest + +from backend.app.models import ImageRecord, ImageRegion, TagKind +from backend.app.models.tag import image_tag +from backend.app.services.tag_service import TagService + +pytestmark = pytest.mark.integration + + +def _ccip(slot: int) -> list[float]: + v = [0.0] * 768 + v[slot] = 1.0 + return v + + +async def _img(db, sha) -> ImageRecord: + img = ImageRecord( + path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", + width=1, height=1, origin="imported_filesystem", integrity_status="unknown", + ) + db.add(img) + await db.flush() + return img + + +async def _figure(db, image_id, ccip): + db.add(ImageRegion( + image_record_id=image_id, kind="figure", rx=0.0, ry=0.0, rw=1.0, rh=1.0, + ccip_embedding=ccip, embedding_version="ccip-test", + )) + + +@pytest.mark.asyncio +async def test_overview_reports_coverage(client, db): + raven = await TagService(db).find_or_create("Raven", TagKind.character) + ref = await _img(db, "a" * 64) + await _figure(db, ref.id, _ccip(0)) + await db.execute(image_tag.insert().values( + image_record_id=ref.id, tag_id=raven.id, source="manual", + )) + q = await _img(db, "b" * 64) + await _figure(db, q.id, _ccip(0)) + await db.commit() + + body = await (await client.get("/api/ccip/overview")).get_json() + assert body["regions_by_kind"].get("figure", 0) >= 2 + assert body["images_with_figure_ccip"] >= 2 + assert any( + c["name"] == "Raven" and c["n_refs"] >= 1 + for c in body["character_references"] + ) + assert "ccip-test" in body["embedding_versions"] + + +@pytest.mark.asyncio +async def test_image_detail_shows_regions_and_matches(client, db): + raven = await TagService(db).find_or_create("Raven", TagKind.character) + ref = await _img(db, "c" * 64) + await _figure(db, ref.id, _ccip(0)) + await db.execute(image_tag.insert().values( + image_record_id=ref.id, tag_id=raven.id, source="manual", + )) + q = await _img(db, "d" * 64) + await _figure(db, q.id, _ccip(0)) + await db.commit() + + body = await (await client.get(f"/api/ccip/images/{q.id}")).get_json() + assert len(body["regions"]) == 1 + r = body["regions"][0] + assert r["kind"] == "figure" and r["has_ccip"] is True and r["has_siglip"] is False + assert any(m["tag_id"] == raven.id for m in body["ccip_matches"]) diff --git a/tests/test_api_gpu.py b/tests/test_api_gpu.py new file mode 100644 index 0000000..4993f75 --- /dev/null +++ b/tests/test_api_gpu.py @@ -0,0 +1,98 @@ +"""GPU-job HTTP API (#114): bearer auth + lease/submit round-trip + backfill.""" +import pytest + +from backend.app.models import ImageRecord +from backend.app.services.ml.gpu_jobs import GpuJobService +from backend.app.services.ml.regions import RegionService + +pytestmark = pytest.mark.integration + + +async def _img(db, sha) -> ImageRecord: + img = ImageRecord( + path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", + width=1, height=1, origin="imported_filesystem", integrity_status="unknown", + ) + db.add(img) + await db.flush() + return img + + +@pytest.mark.asyncio +async def test_agent_endpoints_require_bearer(client, db): + resp = await client.post("/api/gpu/jobs/lease", json={"agent_id": "a1"}) + assert resp.status_code == 401 + # A wrong token is also rejected. + await (await client.post("/api/gpu/token/rotate")).get_json() + bad = await client.post( + "/api/gpu/jobs/lease", json={"agent_id": "a1"}, + headers={"Authorization": "Bearer nope"}, + ) + assert bad.status_code == 401 + + +@pytest.mark.asyncio +async def test_lease_submit_round_trip(client, db): + img = await _img(db, "a" * 64) + await GpuJobService(db).enqueue(img.id, "ccip") + await db.commit() + + token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"] + hdr = {"Authorization": f"Bearer {token}"} + + leased = await client.post( + "/api/gpu/jobs/lease", json={"agent_id": "a1", "batch_size": 5}, headers=hdr, + ) + assert leased.status_code == 200 + jobs = (await leased.get_json())["jobs"] + assert len(jobs) == 1 + j = jobs[0] + assert j["image_id"] == img.id and j["task"] == "ccip" + assert j["image_url"].startswith("/images/") + + submitted = await client.post("/api/gpu/jobs/submit", json={ + "agent_id": "a1", "job_id": j["job_id"], + "regions": [{ + "kind": "figure", "bbox": [0.1, 0.1, 0.4, 0.4], + "ccip_embedding": [0.1] * 768, "embedding_version": "ccip-test", + }], + }, headers=hdr) + assert submitted.status_code == 200 + assert (await submitted.get_json())["stored"] == 1 + + # Job closed (read on the app's own connection via the status endpoint). + st = await (await client.get("/api/gpu/status")).get_json() + assert st["done"] == 1 and st["pending"] == 0 and st["leased"] == 0 + + # Region persisted with its CCIP vector. + regs = await RegionService(db).get_regions(img.id, kinds=["figure"]) + assert len(regs) == 1 and len(list(regs[0].ccip_embedding)) == 768 + + +@pytest.mark.asyncio +async def test_submit_with_stale_lease_is_409(client, db): + img = await _img(db, "b" * 64) + await GpuJobService(db).enqueue(img.id, "ccip") + await db.commit() + token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"] + hdr = {"Authorization": f"Bearer {token}"} + j = (await (await client.post( + "/api/gpu/jobs/lease", json={"agent_id": "a1"}, headers=hdr, + )).get_json())["jobs"][0] + # A different agent can't submit someone else's lease. + resp = await client.post("/api/gpu/jobs/submit", json={ + "agent_id": "other", "job_id": j["job_id"], "regions": [], + }, headers=hdr) + assert resp.status_code == 409 + + +@pytest.mark.asyncio +async def test_backfill_enqueues_then_is_idempotent(db): + await _img(db, "c" * 64) + await _img(db, "d" * 64) + await db.commit() + from backend.app.tasks.ml import enqueue_gpu_backfill + + n = enqueue_gpu_backfill("ccip") # sync task, own session + assert n >= 2 + assert enqueue_gpu_backfill("ccip") == 0 # all already pending diff --git a/tests/test_ccip.py b/tests/test_ccip.py new file mode 100644 index 0000000..5d81c72 --- /dev/null +++ b/tests/test_ccip.py @@ -0,0 +1,88 @@ +"""CCIP few-shot character matcher (#114). numpy cosine on stored vectors — no +model needed, so it runs in CI with synthetic CCIP vectors.""" +import pytest + +from backend.app.models import ImageRecord, ImageRegion, TagKind +from backend.app.models.tag import image_tag +from backend.app.services.ml.ccip import match_image +from backend.app.services.tag_service import TagService + +pytestmark = pytest.mark.integration + + +def _ccip(slot: int) -> list[float]: + v = [0.0] * 768 + v[slot] = 1.0 + return v + + +async def _img(db, sha) -> ImageRecord: + img = ImageRecord( + path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", + width=1, height=1, origin="imported_filesystem", integrity_status="unknown", + ) + db.add(img) + await db.flush() + return img + + +async def _figure(db, image_id, ccip): + db.add(ImageRegion( + image_record_id=image_id, kind="figure", + rx=0.0, ry=0.0, rw=1.0, rh=1.0, + ccip_embedding=ccip, embedding_version="ccip-test", + )) + + +async def _tag_image(db, image_id, tag_id): + await db.execute(image_tag.insert().values( + image_record_id=image_id, tag_id=tag_id, source="manual", + )) + + +@pytest.mark.asyncio +async def test_matches_same_character_across_images(db): + raven = await TagService(db).find_or_create("Raven", TagKind.character) + ref = await _img(db, "a" * 64) # a tagged example = a prototype + await _figure(db, ref.id, _ccip(0)) + await _tag_image(db, ref.id, raven.id) + query = await _img(db, "b" * 64) # untagged, near-identical figure + await _figure(db, query.id, _ccip(0)) + await db.commit() + + matches = await match_image(db, query.id) + m = next(x for x in matches if x["tag_id"] == raven.id) + assert m["source"] == "ccip" and m["category"] == "character" + assert m["score"] > 0.9 + + +@pytest.mark.asyncio +async def test_no_match_for_different_character(db): + raven = await TagService(db).find_or_create("Raven", TagKind.character) + ref = await _img(db, "c" * 64) + await _figure(db, ref.id, _ccip(0)) + await _tag_image(db, ref.id, raven.id) + query = await _img(db, "d" * 64) + await _figure(db, query.id, _ccip(5)) # orthogonal → not Raven + await db.commit() + assert await match_image(db, query.id) == [] + + +@pytest.mark.asyncio +async def test_excludes_already_applied_character(db): + raven = await TagService(db).find_or_create("Raven", TagKind.character) + ref = await _img(db, "e" * 64) + await _figure(db, ref.id, _ccip(0)) + await _tag_image(db, ref.id, raven.id) + query = await _img(db, "f" * 64) + await _figure(db, query.id, _ccip(0)) + await _tag_image(db, query.id, raven.id) # already tagged → no re-suggest + await db.commit() + assert all(m["tag_id"] != raven.id for m in await match_image(db, query.id)) + + +@pytest.mark.asyncio +async def test_no_figure_vectors_means_no_match(db): + query = await _img(db, "g" * 64) + await db.commit() + assert await match_image(db, query.id) == [] diff --git a/tests/test_crops.py b/tests/test_crops.py new file mode 100644 index 0000000..88e2310 --- /dev/null +++ b/tests/test_crops.py @@ -0,0 +1,44 @@ +"""Shared crop primitive (#114) — pure Pillow, no DB, so it runs in the fast +unit lane (no integration marker).""" +from PIL import Image + +from backend.app.services.ml.crops import crop_region + + +def _quadrant_img(): + """400x400 red with a blue bottom-right quadrant, so a crop's content is + checkable by pixel.""" + img = Image.new("RGB", (400, 400), (255, 0, 0)) + img.paste(Image.new("RGB", (200, 200), (0, 0, 255)), (200, 200)) + return img + + +def test_crop_returns_region_pixels(): + crop = crop_region(_quadrant_img(), (0.5, 0.5, 0.5, 0.5)) + assert crop is not None + assert crop.size == (200, 200) + assert crop.getpixel((100, 100)) == (0, 0, 255) # the blue quadrant + + +def test_crop_below_floor_is_rejected(): + # 0.05 * 400 = 20px on a side — below max(64, 0.10*400=40) → None. + assert crop_region(_quadrant_img(), (0.0, 0.0, 0.05, 0.05)) is None + + +def test_crop_clamped_to_image_bounds(): + # Box runs off the right/bottom edge; clamps to the remaining 0.2*400=80px. + crop = crop_region(_quadrant_img(), (0.8, 0.8, 0.5, 0.5)) + assert crop is not None + assert crop.size == (80, 80) + + +def test_pad_expands_the_crop(): + base = crop_region(_quadrant_img(), (0.4, 0.4, 0.2, 0.2)) + padded = crop_region(_quadrant_img(), (0.4, 0.4, 0.2, 0.2), pad=0.5) + assert base.size == (80, 80) + assert padded.size[0] > base.size[0] and padded.size[1] > base.size[1] + + +def test_out_size_resizes_square(): + crop = crop_region(_quadrant_img(), (0.25, 0.25, 0.5, 0.5), out_size=224) + assert crop.size == (224, 224) diff --git a/tests/test_gpu_jobs.py b/tests/test_gpu_jobs.py new file mode 100644 index 0000000..d2a6ea7 --- /dev/null +++ b/tests/test_gpu_jobs.py @@ -0,0 +1,125 @@ +"""GPU-job queue engine (#114): enqueue dedupe + lease/heartbeat/complete/fail.""" +from datetime import UTC, datetime, timedelta + +import pytest +from sqlalchemy import select + +from backend.app.models import GpuJob, ImageRecord +from backend.app.services.ml.gpu_jobs import GpuJobService + +pytestmark = pytest.mark.integration + + +async def _img(db, sha) -> ImageRecord: + img = ImageRecord( + path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", + width=1, height=1, origin="imported_filesystem", integrity_status="unknown", + ) + db.add(img) + await db.flush() + return img + + +@pytest.mark.asyncio +async def test_enqueue_dedupes_same_pair(db): + img = await _img(db, "a" * 64) + svc = GpuJobService(db) + first = await svc.enqueue(img.id, "ccip") + dup = await svc.enqueue(img.id, "ccip") + other = await svc.enqueue(img.id, "siglip_region") + await db.commit() + assert first is not None + assert dup is None # same (image, task) already queued + assert other is not None # different task is fine + + +@pytest.mark.asyncio +async def test_lease_claims_then_skips_when_held(db): + img = await _img(db, "b" * 64) + svc = GpuJobService(db) + await svc.enqueue(img.id, "ccip") + await db.commit() + + leased = await svc.lease("agent-1", batch_size=8) + await db.commit() + assert len(leased) == 1 + assert leased[0].status == "leased" and leased[0].lease_token == "agent-1" + assert leased[0].attempts == 1 + + # Already leased + not expired → a second agent gets nothing. + again = await svc.lease("agent-2", batch_size=8) + await db.commit() + assert again == [] + + +@pytest.mark.asyncio +async def test_expired_lease_is_reclaimed(db): + img = await _img(db, "c" * 64) + svc = GpuJobService(db) + job = await svc.enqueue(img.id, "ccip") + await db.commit() + # Force the lease into the past. + job.status = "leased" + job.lease_token = "dead-agent" + job.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10) + await db.commit() + + leased = await svc.lease("agent-2", batch_size=8) + await db.commit() + assert len(leased) == 1 + assert leased[0].lease_token == "agent-2" + assert leased[0].attempts == 1 # re-lease incremented from 0 (was set directly) + + +@pytest.mark.asyncio +async def test_heartbeat_extends_only_own_lease(db): + img = await _img(db, "d" * 64) + svc = GpuJobService(db) + await svc.enqueue(img.id, "ccip") + await db.commit() + job = (await svc.lease("agent-1"))[0] + await db.commit() + + assert await svc.heartbeat("agent-1", [job.id]) == 1 + assert await svc.heartbeat("someone-else", [job.id]) == 0 + + +@pytest.mark.asyncio +async def test_complete_closes_job(db): + img = await _img(db, "e" * 64) + svc = GpuJobService(db) + await svc.enqueue(img.id, "ccip") + await db.commit() + job = (await svc.lease("agent-1"))[0] + await db.commit() + + assert await svc.complete("wrong-token", job.id) is False + assert await svc.complete("agent-1", job.id) is True + await db.commit() + fresh = await db.get(GpuJob, job.id) + assert fresh.status == "done" and fresh.lease_token is None + + +@pytest.mark.asyncio +async def test_fail_requeues_until_cap(db): + img = await _img(db, "f" * 64) + svc = GpuJobService(db) + await svc.enqueue(img.id, "ccip") + await db.commit() + job = (await svc.lease("agent-1"))[0] # attempts -> 1 + await db.commit() + # Under the cap → back to pending for a retry. + assert await svc.fail("agent-1", job.id, "boom") is True + await db.commit() + assert (await db.get(GpuJob, job.id)).status == "pending" + + # At the attempt cap → terminal 'error'. + j = await db.get(GpuJob, job.id) + j.attempts = 3 + j.status = "leased" + j.lease_token = "agent-1" + j.lease_expires_at = datetime.now(UTC) + timedelta(minutes=5) + await db.commit() + assert await svc.fail("agent-1", job.id, "boom again") is True + await db.commit() + assert (await db.get(GpuJob, job.id)).status == "error" diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py index 7a427e2..6d5326e 100644 --- a/tests/test_ml_suggestions.py +++ b/tests/test_ml_suggestions.py @@ -4,7 +4,7 @@ scikit-learn, ml image only); scoring is numpy-only (available via pgvector).""" import pytest from sqlalchemy import select -from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind +from backend.app.models import ImageRecord, ImageRegion, MLSettings, TagHead, TagKind from backend.app.models.tag import image_tag from backend.app.services.ml.allowlist import AllowlistService from backend.app.services.ml.suggestions import SuggestionService @@ -131,3 +131,35 @@ async def test_rejected_tag_surfaced_flagged_then_reversible(db): sl2 = await SuggestionService(db).for_image(img.id) s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id) assert s2.rejected is False + + +async def _figure(db, image_id, slot): + v = [0.0] * 768 + v[slot] = 1.0 + db.add(ImageRegion( + image_record_id=image_id, kind="figure", + rx=0.0, ry=0.0, rw=1.0, rh=1.0, + ccip_embedding=v, embedding_version="ccip-test", + )) + + +@pytest.mark.asyncio +async def test_ccip_character_surfaces_in_rail(db): + # A character with a CCIP reference (a tagged figure) is suggested on a new + # image whose figure matches — overlaid into the rail alongside the heads. + raven = await TagService(db).find_or_create("Raven", TagKind.character) + ref = await _img(db, "0" * 64, None) # the operator's tagged example + await _figure(db, ref.id, slot=0) + await db.execute(image_tag.insert().values( + image_record_id=ref.id, tag_id=raven.id, source="manual", + )) + query = await _img(db, "1" * 64, None) # untagged, matching figure + await _figure(db, query.id, slot=0) + await db.commit() + + sl = await SuggestionService(db).for_image(query.id) + m = next( + c for c in sl.by_category.get("character", []) + if c.canonical_tag_id == raven.id + ) + assert m.source == "ccip" diff --git a/tests/test_regions.py b/tests/test_regions.py new file mode 100644 index 0000000..f060d77 --- /dev/null +++ b/tests/test_regions.py @@ -0,0 +1,71 @@ +"""Region storage/service for the crop pipeline (#114).""" +import pytest + +from backend.app.models import ImageRecord +from backend.app.services.ml.regions import RegionService + +pytestmark = pytest.mark.integration + + +async def _img(db, sha) -> ImageRecord: + img = ImageRecord( + path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", + width=1, height=1, origin="imported_filesystem", integrity_status="unknown", + ) + db.add(img) + await db.flush() + return img + + +@pytest.mark.asyncio +async def test_replace_and_get_regions(db): + img = await _img(db, "a" * 64) + svc = RegionService(db) + n = await svc.replace_regions(img.id, ["figure"], [ + {"kind": "figure", "bbox": (0.1, 0.1, 0.3, 0.4), + "score": 0.9, "detector_version": "det-v1", "frame_time": 42.5}, + ]) + await db.commit() + assert n == 1 + regs = await svc.get_regions(img.id) + assert len(regs) == 1 + r = regs[0] + assert r.kind == "figure" + assert r.rw == pytest.approx(0.3) and r.rh == pytest.approx(0.4) + assert r.score == pytest.approx(0.9) + assert r.frame_time == pytest.approx(42.5) # video frame timestamp + + +@pytest.mark.asyncio +async def test_replace_is_scoped_by_kind(db): + img = await _img(db, "b" * 64) + svc = RegionService(db) + await svc.replace_regions(img.id, ["figure"], [ + {"kind": "figure", "bbox": (0.0, 0.0, 0.5, 0.5)}, + ]) + await svc.replace_regions(img.id, ["concept"], [ + {"kind": "concept", "bbox": (0.5, 0.5, 0.2, 0.2)}, + ]) + await db.commit() + # Re-running the figure detector must NOT wipe the concept region. + await svc.replace_regions(img.id, ["figure"], [ + {"kind": "figure", "bbox": (0.1, 0.1, 0.4, 0.4)}, + ]) + await db.commit() + kinds = sorted(r.kind for r in await svc.get_regions(img.id)) + assert kinds == ["concept", "figure"] + + +@pytest.mark.asyncio +async def test_ccip_vector_round_trips(db): + img = await _img(db, "c" * 64) + svc = RegionService(db) + await svc.replace_regions(img.id, ["figure"], [ + {"kind": "figure", "bbox": (0.0, 0.0, 0.5, 0.5), + "ccip_embedding": [0.1] * 768, "embedding_version": "ccip-test"}, + ]) + await db.commit() + r = (await svc.get_regions(img.id, kinds=["figure"]))[0] + assert r.ccip_embedding is not None + assert len(list(r.ccip_embedding)) == 768 + assert r.siglip_embedding is None