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 .artist import artist_bp
|
||||||
from .artists import artists_bp
|
from .artists import artists_bp
|
||||||
from .attachments import attachments_bp
|
from .attachments import attachments_bp
|
||||||
|
from .ccip import ccip_bp
|
||||||
from .cleanup import cleanup_bp
|
from .cleanup import cleanup_bp
|
||||||
from .credentials import credentials_bp
|
from .credentials import credentials_bp
|
||||||
from .downloads import downloads_bp
|
from .downloads import downloads_bp
|
||||||
from .extension import extension_bp
|
from .extension import extension_bp
|
||||||
from .gallery import gallery_bp
|
from .gallery import gallery_bp
|
||||||
|
from .gpu import gpu_bp
|
||||||
from .heads import heads_bp
|
from .heads import heads_bp
|
||||||
from .import_admin import import_admin_bp
|
from .import_admin import import_admin_bp
|
||||||
from .ml_admin import ml_admin_bp
|
from .ml_admin import ml_admin_bp
|
||||||
@@ -60,6 +62,8 @@ def all_blueprints() -> list[Blueprint]:
|
|||||||
aliases_bp,
|
aliases_bp,
|
||||||
tag_eval_bp,
|
tag_eval_bp,
|
||||||
heads_bp,
|
heads_bp,
|
||||||
|
gpu_bp,
|
||||||
|
ccip_bp,
|
||||||
ml_admin_bp,
|
ml_admin_bp,
|
||||||
thumbnails_bp,
|
thumbnails_bp,
|
||||||
sources_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 .credential import Credential
|
||||||
from .download_event import DownloadEvent
|
from .download_event import DownloadEvent
|
||||||
from .external_link import ExternalLink
|
from .external_link import ExternalLink
|
||||||
|
from .gpu_job import GpuJob
|
||||||
from .head_auto_apply_run import HeadAutoApplyRun
|
from .head_auto_apply_run import HeadAutoApplyRun
|
||||||
from .head_metric import HeadMetric
|
from .head_metric import HeadMetric
|
||||||
from .head_metrics_snapshot import HeadMetricsSnapshot
|
from .head_metrics_snapshot import HeadMetricsSnapshot
|
||||||
@@ -15,6 +16,7 @@ from .head_training_run import HeadTrainingRun
|
|||||||
from .image_prediction import ImagePrediction
|
from .image_prediction import ImagePrediction
|
||||||
from .image_provenance import ImageProvenance
|
from .image_provenance import ImageProvenance
|
||||||
from .image_record import ImageRecord
|
from .image_record import ImageRecord
|
||||||
|
from .image_region import ImageRegion
|
||||||
from .import_batch import ImportBatch
|
from .import_batch import ImportBatch
|
||||||
from .import_settings import ImportSettings
|
from .import_settings import ImportSettings
|
||||||
from .import_task import ImportTask
|
from .import_task import ImportTask
|
||||||
@@ -60,11 +62,13 @@ __all__ = [
|
|||||||
"ImageRecord",
|
"ImageRecord",
|
||||||
"ImagePrediction",
|
"ImagePrediction",
|
||||||
"ImageProvenance",
|
"ImageProvenance",
|
||||||
|
"ImageRegion",
|
||||||
"Tag",
|
"Tag",
|
||||||
"TagKind",
|
"TagKind",
|
||||||
"image_tag",
|
"image_tag",
|
||||||
"DownloadEvent",
|
"DownloadEvent",
|
||||||
"ExternalLink",
|
"ExternalLink",
|
||||||
|
"GpuJob",
|
||||||
"ImportBatch",
|
"ImportBatch",
|
||||||
"ImportTask",
|
"ImportTask",
|
||||||
"ImportSettings",
|
"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 import ImageRecord, TagSuggestionRejection
|
||||||
from ...models.tag import image_tag
|
from ...models.tag import image_tag
|
||||||
|
from .ccip import match_image as ccip_match_image
|
||||||
from .heads import score_image
|
from .heads import score_image
|
||||||
|
|
||||||
|
|
||||||
@@ -27,7 +28,7 @@ class Suggestion:
|
|||||||
display_name: str
|
display_name: str
|
||||||
category: str
|
category: str
|
||||||
score: float
|
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
|
creates_new_tag: bool
|
||||||
# raw_name = the booru model vocab key behind this suggestion. It's the key
|
# 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
|
# an alias MUST be stored under (resolution looks up the raw key), so the
|
||||||
@@ -92,19 +93,39 @@ class SuggestionService:
|
|||||||
hits = await score_image(
|
hits = await score_image(
|
||||||
self.session, image_id, threshold_override=threshold_override
|
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()
|
result = SuggestionList()
|
||||||
for h in hits:
|
for (cat, tag_id), m in merged.items():
|
||||||
tag_id = h["tag_id"]
|
|
||||||
if tag_id in applied:
|
if tag_id in applied:
|
||||||
continue
|
continue
|
||||||
result.by_category.setdefault(h["category"], []).append(
|
result.by_category.setdefault(cat, []).append(
|
||||||
Suggestion(
|
Suggestion(
|
||||||
canonical_tag_id=tag_id,
|
canonical_tag_id=tag_id,
|
||||||
display_name=h["name"],
|
display_name=m["name"],
|
||||||
category=h["category"],
|
category=cat,
|
||||||
score=h["score"],
|
score=m["score"],
|
||||||
source="head",
|
source=m["source"],
|
||||||
creates_new_tag=False,
|
creates_new_tag=False,
|
||||||
rejected=tag_id in rejected,
|
rejected=tag_id in rejected,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -738,3 +738,33 @@ def scheduled_apply_head_tags() -> str:
|
|||||||
run_id = run.id
|
run_id = run.id
|
||||||
apply_head_tags.delay(run_id)
|
apply_head_tags.delay(run_id)
|
||||||
return "dispatched"
|
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>
|
</v-icon>
|
||||||
</template>
|
</template>
|
||||||
<v-list-item-title>
|
<v-list-item-title>
|
||||||
Create "{{ parsedName }}" as {{ parsedKind }}
|
{{ createLabel }}
|
||||||
</v-list-item-title>
|
</v-list-item-title>
|
||||||
</v-list-item>
|
</v-list-item>
|
||||||
</template>
|
</template>
|
||||||
@@ -178,14 +178,34 @@ watch(query, () => {
|
|||||||
}, 200)
|
}, 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 allowCreate = computed(() => {
|
||||||
const q = parsedName.value
|
const q = parsedName.value
|
||||||
if (!q) return false
|
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 =>
|
return !hits.value.some(h =>
|
||||||
h.name.toLowerCase() === q.toLowerCase() && h.kind === parsedKind.value,
|
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)}%` }
|
function scorePct (s) { return `${Math.round(s.score * 100)}%` }
|
||||||
|
|
||||||
// This image's suggestions that match the typed query, minus any the server
|
// 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">
|
<div class="fc-tile-stack">
|
||||||
<MLThresholdSliders />
|
<MLThresholdSliders />
|
||||||
<HeadsCard />
|
<HeadsCard />
|
||||||
|
<GpuAgentCard />
|
||||||
<AllowlistTable />
|
<AllowlistTable />
|
||||||
<AliasTable />
|
<AliasTable />
|
||||||
<TagEvalCard />
|
<TagEvalCard />
|
||||||
@@ -54,6 +55,7 @@ import MissingFileRepairCard from './MissingFileRepairCard.vue'
|
|||||||
import DbMaintenanceCard from './DbMaintenanceCard.vue'
|
import DbMaintenanceCard from './DbMaintenanceCard.vue'
|
||||||
import MLThresholdSliders from './MLThresholdSliders.vue'
|
import MLThresholdSliders from './MLThresholdSliders.vue'
|
||||||
import HeadsCard from './HeadsCard.vue'
|
import HeadsCard from './HeadsCard.vue'
|
||||||
|
import GpuAgentCard from './GpuAgentCard.vue'
|
||||||
import AllowlistTable from './AllowlistTable.vue'
|
import AllowlistTable from './AllowlistTable.vue'
|
||||||
import AliasTable from './AliasTable.vue'
|
import AliasTable from './AliasTable.vue'
|
||||||
import TagEvalCard from './TagEvalCard.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. -->
|
<!-- CENTER: the focused image (light viewer) + meta. -->
|
||||||
<section class="fc-ex__viewer">
|
<section class="fc-ex__viewer">
|
||||||
<div class="fc-ex__canvas">
|
<div class="fc-ex__canvas">
|
||||||
<ImageCanvas
|
<template v-if="store.anchor">
|
||||||
v-if="store.anchor"
|
<!-- Videos can't render in an <img> — branch to VideoCanvas like
|
||||||
:key="store.anchor.id"
|
the modal does (an MP4 in ImageCanvas just shows the alt). -->
|
||||||
:src="store.anchor.image_url"
|
<ImageCanvas
|
||||||
:alt="`Image ${store.anchor.id}`"
|
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>
|
||||||
<div v-if="store.anchor" class="fc-ex__viewer-foot">
|
<div v-if="store.anchor" class="fc-ex__viewer-foot">
|
||||||
<div class="fc-ex__artist">{{ store.anchor.artist?.name || 'Unknown artist' }}</div>
|
<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 { useHeadTraining } from '../composables/useHeadTraining.js'
|
||||||
import { isTextEntry } from '../utils/textEntry.js'
|
import { isTextEntry } from '../utils/textEntry.js'
|
||||||
import ImageCanvas from '../components/modal/ImageCanvas.vue'
|
import ImageCanvas from '../components/modal/ImageCanvas.vue'
|
||||||
|
import VideoCanvas from '../components/modal/VideoCanvas.vue'
|
||||||
import ImageMetaBar from '../components/modal/ImageMetaBar.vue'
|
import ImageMetaBar from '../components/modal/ImageMetaBar.vue'
|
||||||
import ProvenancePanel from '../components/modal/ProvenancePanel.vue'
|
import ProvenancePanel from '../components/modal/ProvenancePanel.vue'
|
||||||
import TagPanel from '../components/modal/TagPanel.vue'
|
import TagPanel from '../components/modal/TagPanel.vue'
|
||||||
@@ -140,6 +151,7 @@ const store = useExploreStore()
|
|||||||
const modal = useModalStore()
|
const modal = useModalStore()
|
||||||
|
|
||||||
const anchorId = computed(() => route.params.imageId || null)
|
const anchorId = computed(() => route.params.imageId || null)
|
||||||
|
const isVideo = computed(() => !!store.anchor?.mime?.startsWith('video/'))
|
||||||
const seeding = ref(false)
|
const seeding = ref(false)
|
||||||
const seedError = ref(null)
|
const seedError = ref(null)
|
||||||
const tagPanelRef = 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
|
import pytest
|
||||||
from sqlalchemy import select
|
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.models.tag import image_tag
|
||||||
from backend.app.services.ml.allowlist import AllowlistService
|
from backend.app.services.ml.allowlist import AllowlistService
|
||||||
from backend.app.services.ml.suggestions import SuggestionService
|
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)
|
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)
|
s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id)
|
||||||
assert s2.rejected is False
|
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