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
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:
@@ -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"]
|
||||
@@ -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.
|
||||
@@ -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>"""
|
||||
@@ -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()
|
||||
@@ -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")),
|
||||
)
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -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")
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
@@ -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})
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
)
|
||||
@@ -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()
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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'
|
||||
|
||||
@@ -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 }
|
||||
})
|
||||
@@ -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)
|
||||
|
||||
@@ -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"])
|
||||
@@ -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
|
||||
@@ -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) == []
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user