Merge pull request 'CCIP characters + crop/region pipeline + desktop GPU agent (#114)' (#144) from dev into main
Build images / sign-extension (push) Successful in 3s
CI / lint (push) Successful in 3s
CI / frontend-build (push) Successful in 18s
CI / backend-lint-and-test (push) Successful in 26s
Build images / build-web (push) Successful in 2m8s
Build images / build-ml (push) Successful in 2m37s
CI / integration (push) Successful in 3m26s

This commit was merged in pull request #144.
This commit is contained in:
2026-06-29 14:18:57 -04:00
37 changed files with 2281 additions and 16 deletions
+20
View File
@@ -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"]
+60
View File
@@ -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=<paste-the-token> \
-v fc-agent-models:/models \
fc-gpu-agent
```
Then open <http://localhost:8770> — 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/<id>` 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.
View File
+94
View File
@@ -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 = """<!doctype html><html><head><meta charset=utf-8>
<title>FabledCurator GPU agent</title>
<style>
body{font:14px system-ui;margin:2rem;max-width:640px;background:#14171a;color:#e8e8e8}
h1{font-size:18px} button{font:14px system-ui;padding:.5rem 1rem;border:0;border-radius:6px;
margin-right:.5rem;cursor:pointer;color:#fff} .start{background:#2e7d32}.pause{background:#b26a00}
.stop{background:#b3261e} .stat{display:inline-block;margin-right:1.5rem}
.n{font-size:22px;font-weight:700} code{background:#222;padding:2px 6px;border-radius:4px}
.q{margin-top:1rem;color:#9aa}
</style></head><body>
<h1>FabledCurator GPU agent</h1>
<p>FC: <code id=fc>—</code> · token <code id=cfg>—</code></p>
<p>
<button class=start onclick=act('start')>Start</button>
<button class=pause onclick=act('pause')>Pause</button>
<button class=pause onclick=act('resume')>Resume</button>
<button class=stop onclick=act('stop')>Stop</button>
</p>
<p>
<span class=stat><span class=n id=state>idle</span><br>state</span>
<span class=stat><span class=n id=done>0</span><br>processed</span>
<span class=stat><span class=n id=err>0</span><br>errors</span>
<span class=stat><span class=n id=cur>—</span><br>current image</span>
</p>
<div class=q id=queue></div>
<script>
async function act(p){await fetch('/'+p,{method:'POST'});refresh()}
async function refresh(){
const s=await (await fetch('/status')).json()
state.textContent=s.state; done.textContent=s.processed; err.textContent=s.errors
cur.textContent=s.current??''; fc.textContent=s.fc_url
cfg.textContent=s.configured?'set':'MISSING'
queue.textContent=s.queue?`queue — pending ${s.queue.pending} · in flight ${s.queue.leased} · done ${s.queue.done} · errored ${s.queue.error}`:'queue — unreachable'
}
refresh(); setInterval(refresh,3000)
</script></body></html>"""
+66
View File
@@ -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()
+26
View File
@@ -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")),
)
+36
View File
@@ -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")
+48
View File
@@ -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
+39
View File
@@ -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()
+127
View File
@@ -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
+11
View File
@@ -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
+59
View File
@@ -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")
+55
View File
@@ -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")
+4
View File
@@ -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,
+106
View File
@@ -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/<int:image_id>", 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,
})
+198
View File
@@ -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 <token>) 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})
+4
View File
@@ -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",
+50
View File
@@ -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()
)
+62
View File
@@ -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()
)
+120
View File
@@ -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
+73
View File
@@ -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
+134
View File
@@ -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
+59
View File
@@ -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
+29 -8
View File
@@ -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,
)
+30
View File
@@ -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)
@@ -72,7 +72,7 @@
</v-icon>
</template>
<v-list-item-title>
Create "{{ parsedName }}" as {{ parsedKind }}
{{ createLabel }}
</v-list-item-title>
</v-list-item>
</template>
@@ -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
@@ -0,0 +1,166 @@
<template>
<MaintenanceTile
icon="mdi-expansion-card"
title="GPU agent (CCIP + crops)"
blurb="Connect a desktop-GPU agent to embed characters (CCIP) and crops. It pulls work over HTTP — your database and Redis stay private."
:open="true"
>
<p class="fc-muted text-body-2 mb-3">
The agent is a container you run on the machine with the GPU. It
authenticates with the token below, leases jobs from this server, computes
on the GPU, and posts results back all over HTTP. Start it when you want
a burst; stop it to reclaim the card.
</p>
<!-- Token -->
<div class="fc-section-h mb-1">Agent token</div>
<div v-if="loading" class="fc-muted text-body-2">Loading</div>
<template v-else>
<div v-if="tokenValue" class="fc-token">
<code class="fc-token__val">{{ masked ? maskedToken : tokenValue }}</code>
<v-btn
size="x-small" variant="text" :icon="masked ? 'mdi-eye' : 'mdi-eye-off'"
:title="masked ? 'Reveal' : 'Hide'" @click="masked = !masked"
/>
<v-btn
size="x-small" variant="text" icon="mdi-content-copy"
title="Copy token" @click="onCopy"
/>
<v-btn
size="small" variant="text" color="accent" class="ml-auto"
prepend-icon="mdi-refresh" :loading="rotating" @click="onRotate"
>Rotate</v-btn>
</div>
<div v-else>
<v-btn
color="accent" variant="flat" rounded="pill" size="small"
prepend-icon="mdi-key-plus" :loading="rotating" @click="onRotate"
>Generate token</v-btn>
</div>
<p class="fc-muted text-caption mt-2 mb-0">
Point the agent at <code>{{ baseUrl }}</code> with this token. Rotating
invalidates the old token update the agent after you rotate.
</p>
</template>
<!-- Queue -->
<div class="fc-section-h mt-5 mb-2">Work queue</div>
<div class="fc-queue">
<div class="fc-q"><div class="fc-q__n">{{ queue.pending }}</div><div class="fc-q__l">pending</div></div>
<div class="fc-q"><div class="fc-q__n">{{ queue.leased }}</div><div class="fc-q__l">in flight</div></div>
<div class="fc-q"><div class="fc-q__n fc-good">{{ queue.done }}</div><div class="fc-q__l">done</div></div>
<div class="fc-q"><div class="fc-q__n" :class="queue.error ? 'fc-weak' : ''">{{ queue.error }}</div><div class="fc-q__l">errored</div></div>
</div>
<v-btn
class="mt-4" color="accent" variant="tonal" rounded="pill" size="small"
prepend-icon="mdi-account-box-multiple" :loading="backfilling" @click="onBackfill"
>Queue character embedding (CCIP)</v-btn>
<p class="fc-muted text-caption mt-2 mb-0">
Enqueues every image that doesn't have a CCIP embedding yet. Nothing
processes until the agent is running.
</p>
</MaintenanceTile>
</template>
<script setup>
import { toast } from '../../utils/toast.js'
import { computed, onMounted, onUnmounted, ref } from 'vue'
import MaintenanceTile from '../common/MaintenanceTile.vue'
import { useGpuStore } from '../../stores/gpu.js'
import { copyText } from '../../utils/clipboard.js'
const store = useGpuStore()
const loading = ref(true)
const tokenValue = ref(null)
const masked = ref(true)
const rotating = ref(false)
const backfilling = ref(false)
const queue = ref({ pending: 0, leased: 0, done: 0, error: 0 })
let pollTimer = null
const baseUrl = computed(() => window.location.origin)
const maskedToken = computed(() => {
const t = tokenValue.value || ''
return t.length > 8 ? `${t.slice(0, 4)}••••••••${t.slice(-4)}` : ''
})
onMounted(async () => {
try {
tokenValue.value = (await store.token()).token
} catch { /* non-fatal */ } finally {
loading.value = false
}
await refreshQueue()
pollTimer = setInterval(() => { if (!document.hidden) refreshQueue() }, 5000)
})
onUnmounted(() => { if (pollTimer) clearInterval(pollTimer) })
async function refreshQueue() {
try { queue.value = await store.status() } catch { /* non-fatal */ }
}
async function onRotate() {
rotating.value = true
try {
tokenValue.value = (await store.rotateToken()).token
masked.value = false
toast({ text: 'New agent token generated update your agent', type: 'success' })
} catch (e) {
toast({ text: `Could not rotate token: ${e.message}`, type: 'error' })
} finally {
rotating.value = false
}
}
async function onCopy() {
try {
await copyText(tokenValue.value || '') // resolves on success, throws on fail
toast({ text: 'Token copied', type: 'success' })
} catch {
toast({ text: 'Copy failed select and copy manually', type: 'warning' })
}
}
async function onBackfill() {
backfilling.value = true
try {
await store.backfill('ccip')
toast({ text: 'Queued CCIP embedding run the agent to process it', type: 'success' })
await refreshQueue()
} catch (e) {
toast({ text: `Could not queue backfill: ${e.message}`, type: 'error' })
} finally {
backfilling.value = false
}
}
</script>
<style scoped>
.fc-muted { color: rgb(var(--v-theme-on-surface-variant)); }
.fc-section-h {
font-size: 13px; font-weight: 700; letter-spacing: 0.03em;
text-transform: uppercase; color: rgb(var(--v-theme-on-surface));
}
.fc-token {
display: flex; align-items: center; gap: 4px;
background: rgb(var(--v-theme-surface-light)); border-radius: 6px;
padding: 4px 6px 4px 10px;
}
.fc-token__val {
font-family: 'JetBrains Mono', monospace; font-size: 13px;
overflow: hidden; text-overflow: ellipsis; white-space: nowrap;
}
.fc-queue { display: flex; gap: 24px; }
.fc-q__n {
font-size: 20px; font-weight: 700; line-height: 1.1;
font-family: 'JetBrains Mono', monospace;
}
.fc-q__l {
font-size: 11px; text-transform: uppercase; letter-spacing: 0.04em;
color: rgb(var(--v-theme-on-surface-variant));
}
.fc-good { color: rgb(var(--v-theme-success)); }
.fc-weak { color: rgb(var(--v-theme-error)); }
</style>
@@ -27,6 +27,7 @@
<div class="fc-tile-stack">
<MLThresholdSliders />
<HeadsCard />
<GpuAgentCard />
<AllowlistTable />
<AliasTable />
<TagEvalCard />
@@ -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'
+33
View File
@@ -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: <string|null>, 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 }
})
+18 -6
View File
@@ -90,12 +90,22 @@
<!-- CENTER: the focused image (light viewer) + meta. -->
<section class="fc-ex__viewer">
<div class="fc-ex__canvas">
<ImageCanvas
v-if="store.anchor"
:key="store.anchor.id"
:src="store.anchor.image_url"
:alt="`Image ${store.anchor.id}`"
/>
<template v-if="store.anchor">
<!-- Videos can't render in an <img> branch to VideoCanvas like
the modal does (an MP4 in ImageCanvas just shows the alt). -->
<ImageCanvas
v-if="!isVideo"
:key="store.anchor.id"
:src="store.anchor.image_url"
:alt="`Image ${store.anchor.id}`"
/>
<VideoCanvas
v-else
:key="store.anchor.id"
:src="store.anchor.image_url"
:mime="store.anchor.mime"
/>
</template>
</div>
<div v-if="store.anchor" class="fc-ex__viewer-foot">
<div class="fc-ex__artist">{{ store.anchor.artist?.name || 'Unknown artist' }}</div>
@@ -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)
+72
View File
@@ -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"])
+98
View File
@@ -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
+88
View File
@@ -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) == []
+44
View File
@@ -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)
+125
View File
@@ -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"
+33 -1
View File
@@ -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"
+71
View File
@@ -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