Compare commits
28 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 55fa4656ff | |||
| c6f38b0dac | |||
| b91a230f12 | |||
| 74b7ceaf47 | |||
| 301f2de989 | |||
| 625336b6b4 | |||
| b7fd69815e | |||
| 3abbe58450 | |||
| 4a1a9ec5a7 | |||
| 2cb0427868 | |||
| 614b6bc52a | |||
| 7b10f4caab | |||
| b6b151a500 | |||
| 9449241fc2 | |||
| 8419ebd761 | |||
| 60f26247e9 | |||
| de33bab41c | |||
| 5faf34a3b5 | |||
| d57ca847e7 | |||
| d91eef7a4b | |||
| 558d965a1c | |||
| f247f9247c | |||
| 6cabef07a4 | |||
| b735432d02 | |||
| 0ea7ecdea5 | |||
| e8d3400d22 | |||
| f6e10ccc4f | |||
| ad2921b4a0 |
@@ -329,3 +329,41 @@ jobs:
|
||||
file: Dockerfile.ml
|
||||
push: true
|
||||
tags: ${{ steps.tag.outputs.tags }}
|
||||
|
||||
# The desktop GPU agent (#114) — published so the operator pulls + runs it on
|
||||
# the GPU machine instead of building locally. Independent of web/ml (its own
|
||||
# CUDA + onnxruntime-gpu image, context = agent/). Same tag cadence.
|
||||
build-agent:
|
||||
runs-on: python-ci
|
||||
container:
|
||||
image: git.fabledsword.com/bvandeusen/ci-python:3.14
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Determine tag
|
||||
id: tag
|
||||
run: |
|
||||
SHORT_SHA=$(printf '%s' "$GITHUB_SHA" | cut -c1-7)
|
||||
if [ "${GITHUB_REF#refs/tags/}" != "${GITHUB_REF}" ]; then
|
||||
TAG_NAME="${GITHUB_REF#refs/tags/}"
|
||||
echo "tags=git.fabledsword.com/bvandeusen/fabledcurator-agent:${TAG_NAME}" >> "$GITHUB_OUTPUT"
|
||||
elif [ "${GITHUB_REF##*/}" = "main" ]; then
|
||||
echo "tags=git.fabledsword.com/bvandeusen/fabledcurator-agent:main,git.fabledsword.com/bvandeusen/fabledcurator-agent:latest,git.fabledsword.com/bvandeusen/fabledcurator-agent:c-${SHORT_SHA}" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "tags=git.fabledsword.com/bvandeusen/fabledcurator-agent:dev" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Login to Forgejo registry
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
registry: git.fabledsword.com
|
||||
username: ${{ github.actor }}
|
||||
password: ${{ secrets.RELEASE_TOKEN }}
|
||||
|
||||
- name: Build and push agent image
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: agent
|
||||
file: agent/Dockerfile
|
||||
push: true
|
||||
tags: ${{ steps.tag.outputs.tags }}
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
# FabledCurator GPU agent — runs on the desktop with the GPU.
|
||||
# CUDA + cuDNN runtime so onnxruntime-gpu can use the card (it needs cuDNN 9 —
|
||||
# the plain -runtime image lacks it: "libcudnn.so.9: cannot open shared object
|
||||
# file"); ffmpeg for video frames.
|
||||
FROM nvidia/cuda:12.4.1-cudnn-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
|
||||
# torch from the CUDA-12.4 wheel index (matches the base image); its wheels
|
||||
# bundle their own CUDA + cuDNN and coexist with onnxruntime-gpu. Installed
|
||||
# first + separately so the GPU build of torch is deterministic and layer-cached.
|
||||
RUN pip3 install --no-cache-dir torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
|
||||
COPY requirements.txt .
|
||||
RUN pip3 install --no-cache-dir -r requirements.txt
|
||||
COPY fc_agent ./fc_agent
|
||||
|
||||
# imgutils ONNX models + the transformers SigLIP weights both cache here; mount
|
||||
# a volume to persist them across restarts (the SigLIP download is ~3.5 GB once).
|
||||
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,71 @@
|
||||
# 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.
|
||||
|
||||
## 0. Host prerequisite — NVIDIA Container Toolkit
|
||||
Docker needs the toolkit to hand the GPU to a container (else: *"could not select
|
||||
device driver nvidia with capabilities [[gpu]]"*). On Arch/CachyOS:
|
||||
```sh
|
||||
sudo pacman -S nvidia-container-toolkit
|
||||
sudo nvidia-ctk runtime configure --runtime=docker
|
||||
sudo systemctl restart docker
|
||||
# verify:
|
||||
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi
|
||||
```
|
||||
|
||||
## 1. Get a token
|
||||
In FC: **Settings → Tagging → GPU agent → Generate token** (or Rotate). Copy it.
|
||||
|
||||
## 2. Pull (CI publishes it alongside the web/ml images)
|
||||
```sh
|
||||
docker pull git.fabledsword.com/bvandeusen/fabledcurator-agent:latest
|
||||
```
|
||||
> Local build for development instead: `docker build -t fc-gpu-agent 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 \
|
||||
git.fabledsword.com/bvandeusen/fabledcurator-agent:latest
|
||||
```
|
||||
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,53 @@
|
||||
# FabledCurator GPU agent — desktop run via docker compose.
|
||||
#
|
||||
# Usage:
|
||||
# 1. Generate a token: FC → Settings → Tagging → GPU agent → Generate token.
|
||||
# 2. Create a .env next to this file:
|
||||
# FC_URL=http://curator.traefik.internal
|
||||
# FC_TOKEN=<paste-the-token>
|
||||
# # optional: CCIP_MODEL=ccip-caformer_b36-24 (the F1-0.94 variant)
|
||||
# 3. docker compose up -d (pulls the published image)
|
||||
# 4. Open http://localhost:8770 → Start. Pause/Stop hands the GPU back.
|
||||
# docker compose down to stop the container entirely.
|
||||
#
|
||||
# Surviving a curator redeploy (you're away, can't touch the agent):
|
||||
# - A running agent rides out curator being unreachable on its own — it retries
|
||||
# leasing with capped backoff and resumes when the server is back. In-flight
|
||||
# work is handed back (not failed), so a redeploy never poisons good jobs.
|
||||
# - AUTO_START=1 (below) also resumes the worker if the AGENT container itself
|
||||
# restarts (host reboot / crash via `restart: unless-stopped`) — no click.
|
||||
#
|
||||
# Needs the NVIDIA Container Toolkit installed on the host for --gpus.
|
||||
|
||||
services:
|
||||
fc-gpu-agent:
|
||||
image: git.fabledsword.com/bvandeusen/fabledcurator-agent:latest
|
||||
pull_policy: always
|
||||
ports:
|
||||
- "8770:8770"
|
||||
environment:
|
||||
FC_URL: ${FC_URL:-http://curator.traefik.internal}
|
||||
FC_TOKEN: ${FC_TOKEN:?set FC_TOKEN in .env (FC → GPU agent → Generate token)}
|
||||
CCIP_MODEL: ${CCIP_MODEL:-}
|
||||
DETECTOR_LEVEL: ${DETECTOR_LEVEL:-m}
|
||||
BATCH_SIZE: ${BATCH_SIZE:-4}
|
||||
# Resume the worker automatically on container start (survive a reboot /
|
||||
# crash-restart while you're away). Set to 0 to require a manual Start.
|
||||
AUTO_START: ${AUTO_START:-1}
|
||||
# Crop embedder (SigLIP concept bag): float16 keeps VRAM low on a shared
|
||||
# desktop GPU; the model itself is announced by the server.
|
||||
SIGLIP_DTYPE: ${SIGLIP_DTYPE:-float16}
|
||||
volumes:
|
||||
# Persist the downloaded ONNX models so restarts are fast.
|
||||
- fc-agent-models:/models
|
||||
restart: unless-stopped
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
|
||||
volumes:
|
||||
fc-agent-models:
|
||||
@@ -0,0 +1,134 @@
|
||||
"""FastAPI control surface for the agent (served on localhost).
|
||||
|
||||
Start / stop the worker pool, tune the worker count live (trades desktop
|
||||
responsiveness for throughput), and watch GPU load + progress + the server-side
|
||||
queue. Config is env-seeded; the worker count is adjustable here on the fly.
|
||||
"""
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse
|
||||
|
||||
from .config import Config
|
||||
from .gpu import read_gpu
|
||||
from .worker import Worker
|
||||
|
||||
cfg = Config.from_env()
|
||||
worker = Worker(cfg)
|
||||
app = FastAPI(title="FabledCurator GPU agent")
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def _maybe_autostart() -> None:
|
||||
# With AUTO_START set, a container restart (host reboot, or `restart:
|
||||
# unless-stopped` after a crash) resumes the worker on its own — the slots
|
||||
# then ride out a still-down curator via lease backoff. Lets the agent
|
||||
# survive a redeploy with nobody at the desktop to click Start.
|
||||
if cfg.auto_start and cfg.token:
|
||||
worker.start()
|
||||
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
def index() -> str:
|
||||
return _PAGE
|
||||
|
||||
|
||||
@app.post("/start")
|
||||
def start():
|
||||
worker.start()
|
||||
return JSONResponse(worker.status())
|
||||
|
||||
|
||||
@app.post("/stop")
|
||||
def stop():
|
||||
worker.stop()
|
||||
return JSONResponse(worker.status())
|
||||
|
||||
|
||||
@app.post("/concurrency")
|
||||
async def concurrency(request: Request):
|
||||
body = await request.json()
|
||||
worker.set_concurrency(int(body.get("value", 1)))
|
||||
return JSONResponse(worker.status())
|
||||
|
||||
|
||||
@app.get("/status")
|
||||
def status():
|
||||
s = worker.status()
|
||||
s["fc_url"] = cfg.fc_url
|
||||
s["configured"] = bool(cfg.token)
|
||||
s["gpu"] = read_gpu()
|
||||
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:680px;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}.stop{background:#b3261e}
|
||||
.step{background:#33373b;padding:.4rem .7rem;font-weight:700}
|
||||
.stat{display:inline-block;margin-right:1.5rem;vertical-align:top}
|
||||
.n{font-size:22px;font-weight:700} code{background:#222;padding:2px 6px;border-radius:4px}
|
||||
.q,.gpu{margin-top:1rem;color:#9aa} .bar{height:8px;border-radius:4px;background:#222;overflow:hidden;
|
||||
max-width:320px;margin-top:4px} .bar>i{display:block;height:100%;background:#3f7d3f}
|
||||
.row{margin:.8rem 0}
|
||||
</style></head><body>
|
||||
<h1>FabledCurator GPU agent</h1>
|
||||
<p>FC: <code id=fc>—</code> · token <code id=cfg>—</code></p>
|
||||
<div class=row>
|
||||
<button class=start onclick=act('start')>Start</button>
|
||||
<button class=stop onclick=act('stop')>Stop</button>
|
||||
</div>
|
||||
<div class=row>
|
||||
workers
|
||||
<button class=step onclick=setc(-1)>−</button>
|
||||
<input id=conc type=number min=1 value=1
|
||||
style="width:3.5rem;font:700 16px system-ui;text-align:center;background:#222;color:#e8e8e8;border:1px solid #444;border-radius:6px;padding:.3rem"
|
||||
onchange="setv(this.value)">
|
||||
<button class=step onclick=setc(1)>+</button>
|
||||
<span class=cap style=color:#9aa>(more = overlap I/O, fill the GPU) max <b id=capn>8</b></span>
|
||||
</div>
|
||||
<div class=row>
|
||||
<span class=stat><span class=n id=state>stopped</span><br>state</span>
|
||||
<span class=stat><span class=n id=active>0</span><br>active now</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=wait>0</span><br>waited out</span>
|
||||
</div>
|
||||
<div id=banner style="display:none;margin:.6rem 0;padding:.5rem .8rem;border-radius:6px;background:#5a4a17;color:#ffe28a">
|
||||
curator unreachable — holding work + retrying, will resume on its own (no restart needed)
|
||||
</div>
|
||||
<div class=gpu id=gpu>GPU — …</div>
|
||||
<div class=bar><i id=gpubar style=width:0%></i></div>
|
||||
<div class=q id=queue></div>
|
||||
<script>
|
||||
let CAP=8
|
||||
async function act(p){await fetch('/'+p,{method:'POST'});refresh()}
|
||||
function setc(d){ setv((parseInt(conc.value||'1'))+d) }
|
||||
async function setv(v){
|
||||
v=Math.max(1,Math.min(CAP,parseInt(v)||1)); conc.value=v
|
||||
await fetch('/concurrency',{method:'POST',headers:{'Content-Type':'application/json'},
|
||||
body:JSON.stringify({value:v})});refresh()
|
||||
}
|
||||
async function refresh(){
|
||||
const s=await (await fetch('/status')).json()
|
||||
CAP=s.max_concurrency||8; capn.textContent=CAP
|
||||
state.textContent=s.state; active.textContent=s.active; done.textContent=s.processed
|
||||
err.textContent=s.errors; fc.textContent=s.fc_url; wait.textContent=s.transient||0
|
||||
// Running but the queue read failed → curator is unreachable; show we're
|
||||
// riding it out rather than erroring.
|
||||
banner.style.display=(s.state==='running' && !s.queue)?'block':'none'
|
||||
if(document.activeElement!==conc) conc.value=s.concurrency
|
||||
conc.max=CAP
|
||||
cfg.textContent=s.configured?'set':'MISSING'
|
||||
if(s.gpu){
|
||||
gpu.textContent=`GPU — ${s.gpu.util_pct}% util · VRAM ${s.gpu.mem_used_mb}/${s.gpu.mem_total_mb} MB · ${s.gpu.temp_c}°C`
|
||||
gpubar.style.width=Math.round(100*s.gpu.mem_used_mb/s.gpu.mem_total_mb)+'%'
|
||||
} else { gpu.textContent='GPU — n/a (CPU fallback?)'; gpubar.style.width='0%' }
|
||||
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,85 @@
|
||||
"""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
|
||||
from requests.adapters import HTTPAdapter
|
||||
|
||||
|
||||
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}"
|
||||
# Many worker threads share this Session; the default pool (10) would
|
||||
# throttle them + spam "connection pool is full". Size it for the cap.
|
||||
adapter = HTTPAdapter(pool_connections=64, pool_maxsize=64)
|
||||
self.s.mount("http://", adapter)
|
||||
self.s.mount("https://", adapter)
|
||||
|
||||
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 release(self, job_ids: list[int]) -> None:
|
||||
# Graceful hand-back on stop so orphaned work is re-leased at once.
|
||||
if not job_ids:
|
||||
return
|
||||
try:
|
||||
self.s.post(
|
||||
f"{self.base}/api/gpu/jobs/release",
|
||||
json={"agent_id": self.agent_id, "job_ids": job_ids},
|
||||
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,36 @@
|
||||
"""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 a worker leases per round
|
||||
concurrency: int # INITIAL parallel workers (tunable live from the UI)
|
||||
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
|
||||
embed_dtype: str # torch dtype for the crop embedder: float16|float32
|
||||
embed_model_override: str # force a SigLIP-family model ("" → use the one
|
||||
# the server announces in the lease)
|
||||
auto_start: bool # start the worker pool on boot (so a container restart
|
||||
# resumes processing without anyone clicking Start)
|
||||
|
||||
@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")),
|
||||
concurrency=int(os.environ.get("CONCURRENCY", "1")),
|
||||
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")),
|
||||
embed_dtype=os.environ.get("SIGLIP_DTYPE", "float16"),
|
||||
embed_model_override=os.environ.get("EMBED_MODEL_NAME", ""),
|
||||
auto_start=os.environ.get("AUTO_START", "").lower() in ("1", "true", "yes"),
|
||||
)
|
||||
@@ -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,69 @@
|
||||
"""Crop EMBEDDER for the concept bag — model-agnostic (CLIP/SigLIP-family).
|
||||
|
||||
The server trains its per-concept heads in the embedding space of whatever model
|
||||
its `embedder_model_version` names; a crop must be embedded with the SAME model
|
||||
or its vector lands in a different coordinate system and every head misfires. So
|
||||
the model identity (HF name + version) is ANNOUNCED BY THE SERVER in the lease —
|
||||
nothing here is hardcoded to SigLIP. Whatever name the server sends is loaded via
|
||||
transformers `get_image_features` (the CLIP/SigLIP-family image-tower call); a
|
||||
non-CLIP backbone (e.g. a DINO encoder) would need its own pooling adapter.
|
||||
|
||||
torch on CUDA, fp16 by default to keep VRAM low on a shared desktop GPU — the
|
||||
tiny fp16-vs-fp32 difference is negligible for the linear heads (cosine ~0.999).
|
||||
A single inference lock serializes the forward pass: the pipeline is I/O-bound,
|
||||
so the GPU isn't the bottleneck, and one model shared across worker threads is
|
||||
safest behind a lock.
|
||||
"""
|
||||
import threading
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class CropEmbedder:
|
||||
def __init__(self, model_name: str, dtype: str = "float16"):
|
||||
self._name = model_name
|
||||
self._dtype_name = dtype
|
||||
self._model = None
|
||||
self._processor = None
|
||||
self._torch = None
|
||||
self._device = None
|
||||
self._dt = None
|
||||
self._load_lock = threading.Lock()
|
||||
self._infer_lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def model_name(self) -> str:
|
||||
return self._name
|
||||
|
||||
def load(self) -> None:
|
||||
if self._model is not None:
|
||||
return
|
||||
with self._load_lock:
|
||||
if self._model is not None:
|
||||
return
|
||||
import torch
|
||||
from transformers import AutoImageProcessor, AutoModel
|
||||
|
||||
self._torch = torch
|
||||
self._device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dt = getattr(torch, self._dtype_name, torch.float16)
|
||||
if self._device == "cpu":
|
||||
dt = torch.float32 # fp16 matmul is unsupported/slow on CPU
|
||||
self._dt = dt
|
||||
self._processor = AutoImageProcessor.from_pretrained(self._name)
|
||||
model = AutoModel.from_pretrained(self._name, torch_dtype=dt)
|
||||
model.eval().to(self._device)
|
||||
self._model = model
|
||||
|
||||
def embed(self, image: Image.Image) -> list[float]:
|
||||
"""A crop → its embedding as a plain float list, ready to POST."""
|
||||
self.load()
|
||||
torch = self._torch
|
||||
enc = self._processor(images=image, return_tensors="pt")
|
||||
pixel_values = enc["pixel_values"].to(self._device, self._dt)
|
||||
with self._infer_lock, torch.no_grad():
|
||||
out = self._model.get_image_features(pixel_values=pixel_values)
|
||||
pooled = out.pooler_output if hasattr(out, "pooler_output") else out
|
||||
vec = pooled[0].float().cpu().numpy().astype(np.float32).reshape(-1)
|
||||
return vec.tolist()
|
||||
@@ -0,0 +1,30 @@
|
||||
"""GPU load readout via nvidia-smi (present in the container thanks to the
|
||||
NVIDIA Container Toolkit's `utility` capability). Returns None if unavailable —
|
||||
the UI just shows n/a (e.g. CPU-fallback run)."""
|
||||
import subprocess
|
||||
|
||||
|
||||
def read_gpu() -> dict | None:
|
||||
try:
|
||||
out = subprocess.run(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
capture_output=True, text=True, timeout=5, check=True,
|
||||
).stdout.strip().splitlines()
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return None
|
||||
if not out:
|
||||
return None
|
||||
parts = [p.strip() for p in out[0].split(",")]
|
||||
try:
|
||||
return {
|
||||
"util_pct": int(float(parts[0])),
|
||||
"mem_used_mb": int(float(parts[1])),
|
||||
"mem_total_mb": int(float(parts[2])),
|
||||
"temp_c": int(float(parts[3])),
|
||||
}
|
||||
except (ValueError, IndexError):
|
||||
return None
|
||||
@@ -0,0 +1,63 @@
|
||||
"""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 to_rgb(img: Image.Image) -> Image.Image:
|
||||
"""RGB, flattening any transparency onto white first. A naive convert('RGB')
|
||||
on a palette-with-transparency image (common for character PNGs on a clear
|
||||
background) lets PIL guess the transparent pixels — usually black artifacts
|
||||
that bleed into the crop + the embedding (and the "should be converted to
|
||||
RGBA" warning). Compositing over white gives a clean, consistent background."""
|
||||
if img.mode in ("RGBA", "LA", "PA") or (
|
||||
img.mode == "P" and "transparency" in img.info
|
||||
):
|
||||
img = img.convert("RGBA")
|
||||
bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
|
||||
return Image.alpha_composite(bg, img).convert("RGB")
|
||||
return img.convert("RGB")
|
||||
|
||||
|
||||
def load_image(data: bytes) -> Image.Image:
|
||||
return to_rgb(Image.open(io.BytesIO(data)))
|
||||
|
||||
|
||||
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), to_rgb(im)))
|
||||
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,274 @@
|
||||
"""The lease → fetch → detect+embed → submit loop, run by a pool of worker
|
||||
slots whose count is tunable live from the UI.
|
||||
|
||||
Each slot is an independent loop (its own leases; the server's SKIP-LOCKED lease
|
||||
keeps them from colliding). More slots = more GPU load + throughput; the model is
|
||||
loaded once and shared, so slots add concurrent inference, not N× model VRAM.
|
||||
That's the dial the operator turns to trade desktop responsiveness for speed.
|
||||
|
||||
Stop (or shrinking the pool) RELEASES a slot's still-leased jobs immediately so
|
||||
orphaned work is re-picked at once rather than waiting out the lease.
|
||||
"""
|
||||
import threading
|
||||
|
||||
import requests
|
||||
|
||||
from . import media, models
|
||||
from .client import FcClient
|
||||
from .config import Config
|
||||
from .crops import crop_region
|
||||
|
||||
# Cap on the lease-retry backoff: when curator is unreachable (e.g. you redeploy
|
||||
# it while away), each slot retries leasing with exponential backoff up to this
|
||||
# many seconds, then resumes within this window once the server is back — no
|
||||
# restart needed.
|
||||
MAX_BACKOFF_SECONDS = 60.0
|
||||
|
||||
|
||||
def _is_transient(exc: "requests.RequestException") -> bool:
|
||||
"""A server/transport problem (wait it out) vs a job-specific fault (fail it).
|
||||
No response → connection refused/timeout → curator is down → transient. With
|
||||
a response: 5xx, auth (401/403, e.g. a token blip on redeploy), 408/409/429
|
||||
(timeout / our lease reclaimed / rate-limited) are all 'not this job's fault'.
|
||||
A specific 4xx like 404 (image gone) / 400 IS the job's fault → fail it."""
|
||||
resp = getattr(exc, "response", None)
|
||||
if resp is None:
|
||||
return True
|
||||
return resp.status_code >= 500 or resp.status_code in (401, 403, 408, 409, 429)
|
||||
|
||||
# Generous cap: the pipeline is usually I/O-bound (downloading + decoding images
|
||||
# over HTTP), so the GPU stays underused until many workers overlap that I/O.
|
||||
# Push it up while watching the GPU util + VRAM in the UI.
|
||||
MAX_CONCURRENCY = 32
|
||||
|
||||
# Fallbacks only — the server ANNOUNCES the embedding model (name + version) in
|
||||
# the lease so the agent stays model-agnostic and in lock-step with the space
|
||||
# the heads were trained in. These cover an older server that doesn't send them.
|
||||
DEFAULT_EMBED_MODEL = "google/siglip-so400m-patch14-384"
|
||||
DEFAULT_EMBED_VERSION = "siglip-so400m-patch14-384"
|
||||
|
||||
|
||||
class _Slot:
|
||||
"""One worker loop. `inflight` = jobs leased but not yet processed, so a
|
||||
graceful stop can hand them back."""
|
||||
__slots__ = ("stop", "inflight")
|
||||
|
||||
def __init__(self):
|
||||
self.stop = threading.Event()
|
||||
self.inflight: list[int] = []
|
||||
|
||||
|
||||
class Worker:
|
||||
def __init__(self, cfg: Config):
|
||||
self.cfg = cfg
|
||||
self.client = FcClient(cfg.fc_url, cfg.token, cfg.agent_id)
|
||||
self._lock = threading.Lock()
|
||||
self._running = False
|
||||
self._target = max(1, min(MAX_CONCURRENCY, cfg.concurrency))
|
||||
self._slots: list[_Slot] = []
|
||||
self.processed = 0
|
||||
self.errors = 0
|
||||
self.transient = 0 # jobs handed back due to a server outage (NOT
|
||||
# failed) — the "waiting out curator" counter
|
||||
self._active = 0 # slots currently mid-image
|
||||
# The crop embedder (SigLIP-family) is built lazily on the first job that
|
||||
# needs it, from the model the server announces — one shared instance.
|
||||
self._embedder = None
|
||||
self._embedder_lock = threading.Lock()
|
||||
|
||||
# --- control -----------------------------------------------------------
|
||||
def start(self):
|
||||
with self._lock:
|
||||
self._running = True
|
||||
self._reconcile_locked()
|
||||
|
||||
def stop(self):
|
||||
with self._lock:
|
||||
self._running = False
|
||||
slots, self._slots = self._slots, []
|
||||
for s in slots:
|
||||
s.stop.set() # each slot releases its inflight on exit
|
||||
|
||||
def set_concurrency(self, n: int):
|
||||
with self._lock:
|
||||
self._target = max(1, min(MAX_CONCURRENCY, int(n)))
|
||||
if self._running:
|
||||
self._reconcile_locked()
|
||||
|
||||
def _reconcile_locked(self):
|
||||
while len(self._slots) < self._target:
|
||||
slot = _Slot()
|
||||
self._slots.append(slot)
|
||||
threading.Thread(target=self._loop, args=(slot,), daemon=True).start()
|
||||
while len(self._slots) > self._target:
|
||||
self._slots.pop().stop.set()
|
||||
|
||||
def status(self) -> dict:
|
||||
with self._lock:
|
||||
return {
|
||||
"state": "running" if self._running else "stopped",
|
||||
"concurrency": self._target,
|
||||
"max_concurrency": MAX_CONCURRENCY,
|
||||
"workers": len(self._slots),
|
||||
"active": self._active,
|
||||
"processed": self.processed,
|
||||
"errors": self.errors,
|
||||
"transient": self.transient,
|
||||
}
|
||||
|
||||
def _bump(self, *, processed=0, errors=0, active=0, transient=0):
|
||||
with self._lock:
|
||||
self.processed += processed
|
||||
self.errors += errors
|
||||
self.transient += transient
|
||||
self._active += active
|
||||
|
||||
# --- per-slot loop -----------------------------------------------------
|
||||
def _loop(self, slot: _Slot):
|
||||
backoff = self.cfg.poll_idle_seconds
|
||||
while not slot.stop.is_set() and self._running:
|
||||
try:
|
||||
jobs = self.client.lease(self.cfg.batch_size)
|
||||
backoff = self.cfg.poll_idle_seconds # server answered → reset
|
||||
except Exception:
|
||||
# curator unreachable (redeploy, network drop): wait it out with
|
||||
# exponential backoff, capped — resume on our own when it returns.
|
||||
self._interruptible_sleep(slot, backoff)
|
||||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||||
continue
|
||||
if not jobs:
|
||||
self._interruptible_sleep(slot, self.cfg.poll_idle_seconds)
|
||||
continue
|
||||
slot.inflight = [j["job_id"] for j in jobs]
|
||||
for job in jobs:
|
||||
if slot.stop.is_set() or not self._running:
|
||||
break
|
||||
ok = self._process(job)
|
||||
slot.inflight = [i for i in slot.inflight if i != job["job_id"]]
|
||||
if not ok:
|
||||
# Server went away mid-batch: hand the rest back (best effort)
|
||||
# and back off instead of hammering a recovering server or
|
||||
# burning the jobs' attempt budgets on fail().
|
||||
if slot.inflight:
|
||||
self.client.release(slot.inflight)
|
||||
slot.inflight = []
|
||||
self._interruptible_sleep(slot, backoff)
|
||||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||||
break
|
||||
if slot.inflight:
|
||||
self.client.heartbeat(slot.inflight)
|
||||
# Graceful hand-back of anything leased but not processed.
|
||||
if slot.inflight:
|
||||
self.client.release(slot.inflight)
|
||||
slot.inflight = []
|
||||
|
||||
def _interruptible_sleep(self, slot: _Slot, seconds: float):
|
||||
"""Sleep, but wake immediately if the slot is told to stop — so a Stop or
|
||||
a pool-shrink doesn't hang for a full backoff window."""
|
||||
slot.stop.wait(timeout=seconds)
|
||||
|
||||
def _ensure_embedder(self, model_name: str):
|
||||
if self._embedder is not None:
|
||||
return self._embedder
|
||||
with self._embedder_lock:
|
||||
if self._embedder is None:
|
||||
from .embedder import CropEmbedder
|
||||
self._embedder = CropEmbedder(model_name, self.cfg.embed_dtype)
|
||||
return self._embedder
|
||||
|
||||
def _process(self, job: dict) -> bool:
|
||||
"""Process one job. Returns True when handled (completed, or hard-failed
|
||||
because the job itself is bad) and False on a TRANSPORT error (curator
|
||||
unreachable / 5xx / our lease was reclaimed mid-flight) — which is not
|
||||
the job's fault, so the caller backs off and the job is left to be
|
||||
re-leased rather than fail()ed into its attempt budget."""
|
||||
self._bump(active=1)
|
||||
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))]
|
||||
|
||||
# task picks what to produce per crop:
|
||||
# 'siglip' (backfill existing images) → concept (SigLIP) regions
|
||||
# ONLY, so it never churns their figure/CCIP regions or the
|
||||
# character-reference cache.
|
||||
# 'ccip' / 'both' (a new image's first pass) → figure (CCIP) AND
|
||||
# concept (SigLIP) in one go, off the same crop.
|
||||
task = job.get("task") or "ccip"
|
||||
want_ccip = task in ("ccip", "both")
|
||||
want_siglip = task in ("ccip", "siglip", "both")
|
||||
replace_kinds = (
|
||||
["concept"] if task == "siglip" else ["figure", "face", "concept"]
|
||||
)
|
||||
|
||||
embed_version = job.get("embed_version") or DEFAULT_EMBED_VERSION
|
||||
embedder = None
|
||||
if want_siglip:
|
||||
model_name = (
|
||||
self.cfg.embed_model_override
|
||||
or job.get("embed_model_name")
|
||||
or DEFAULT_EMBED_MODEL
|
||||
)
|
||||
embedder = self._ensure_embedder(model_name)
|
||||
|
||||
regions = []
|
||||
ccip_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
|
||||
if want_ccip:
|
||||
regions.append({
|
||||
"kind": "figure",
|
||||
"bbox": list(bbox),
|
||||
"frame_time": t,
|
||||
"score": score,
|
||||
"ccip_embedding": models.ccip_vector(
|
||||
crop, self.cfg.ccip_model or None
|
||||
),
|
||||
"embedding_version": ccip_ev,
|
||||
"detector_version": dv,
|
||||
})
|
||||
if want_siglip:
|
||||
regions.append({
|
||||
"kind": "concept",
|
||||
"bbox": list(bbox),
|
||||
"frame_time": t,
|
||||
"score": score,
|
||||
"siglip_embedding": embedder.embed(crop),
|
||||
"embedding_version": embed_version,
|
||||
"detector_version": dv,
|
||||
})
|
||||
self.client.submit(job["job_id"], regions, replace_kinds)
|
||||
self._bump(processed=1)
|
||||
return True
|
||||
except requests.RequestException as exc:
|
||||
if _is_transient(exc):
|
||||
# curator down/redeploying, a 5xx, or our lease was reclaimed
|
||||
# while we worked. NOT the job's fault — hand it back (best
|
||||
# effort; no-ops if the server is still down, then the server's
|
||||
# orphan-recovery reclaims it) and signal the loop to wait.
|
||||
self._bump(transient=1)
|
||||
self.client.release([job["job_id"]])
|
||||
return False
|
||||
# A job-specific HTTP fault (404 image gone, 400) → fail it so it
|
||||
# doesn't re-lease forever.
|
||||
self._bump(errors=1)
|
||||
self.client.fail(job["job_id"], str(exc)[:500])
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — a genuine job fault: report it
|
||||
self._bump(errors=1)
|
||||
self.client.fail(job["job_id"], str(exc)[:500])
|
||||
return True
|
||||
finally:
|
||||
self._bump(active=-1)
|
||||
@@ -0,0 +1,15 @@
|
||||
# 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
|
||||
# The crop EMBEDDER (concept bag). torch is installed separately in the
|
||||
# Dockerfile from the CUDA-12.4 wheel index so the GPU build is deterministic;
|
||||
# transformers loads whatever SigLIP-family model the server announces.
|
||||
transformers>=4.45
|
||||
# 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")
|
||||
@@ -0,0 +1,33 @@
|
||||
"""ml_settings.ccip_match_threshold — tunable CCIP character-match cut (#114)
|
||||
|
||||
The v1 matcher used a flat 0.75 cosine; live data showed that over-fires (a
|
||||
high-reference character matched a scatter of images). 0.85 keeps the confident
|
||||
single-character matches and drops the noise. Tunable from the GPU agent card.
|
||||
|
||||
Revision ID: 0063
|
||||
Revises: 0062
|
||||
Create Date: 2026-06-29
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision: str = "0063"
|
||||
down_revision: Union[str, None] = "0062"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"ml_settings",
|
||||
sa.Column(
|
||||
"ccip_match_threshold", sa.Float(), nullable=False,
|
||||
server_default="0.85",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("ml_settings", "ccip_match_threshold")
|
||||
@@ -0,0 +1,42 @@
|
||||
"""ml_settings: CCIP auto-apply switch + threshold (#114)
|
||||
|
||||
Confident CCIP character matches auto-tag (source='ccip_auto') on a daily sweep,
|
||||
so identity tags keep flowing without pressing a button. ON by default (opt-out,
|
||||
like head auto-apply); the high threshold (0.92, above the 0.85 suggest cut) +
|
||||
single-character references keep it safe, and every auto-tag is reversible.
|
||||
|
||||
Revision ID: 0064
|
||||
Revises: 0063
|
||||
Create Date: 2026-06-30
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision: str = "0064"
|
||||
down_revision: Union[str, None] = "0063"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"ml_settings",
|
||||
sa.Column(
|
||||
"ccip_auto_apply_enabled", sa.Boolean(), nullable=False,
|
||||
server_default=sa.true(),
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"ml_settings",
|
||||
sa.Column(
|
||||
"ccip_auto_apply_threshold", sa.Float(), nullable=False,
|
||||
server_default="0.92",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("ml_settings", "ccip_auto_apply_threshold")
|
||||
op.drop_column("ml_settings", "ccip_auto_apply_enabled")
|
||||
@@ -20,11 +20,13 @@ def all_blueprints() -> list[Blueprint]:
|
||||
from .artist import artist_bp
|
||||
from .artists import artists_bp
|
||||
from .attachments import attachments_bp
|
||||
from .ccip import ccip_bp
|
||||
from .cleanup import cleanup_bp
|
||||
from .credentials import credentials_bp
|
||||
from .downloads import downloads_bp
|
||||
from .extension import extension_bp
|
||||
from .gallery import gallery_bp
|
||||
from .gpu import gpu_bp
|
||||
from .heads import heads_bp
|
||||
from .import_admin import import_admin_bp
|
||||
from .ml_admin import ml_admin_bp
|
||||
@@ -60,6 +62,8 @@ def all_blueprints() -> list[Blueprint]:
|
||||
aliases_bp,
|
||||
tag_eval_bp,
|
||||
heads_bp,
|
||||
gpu_bp,
|
||||
ccip_bp,
|
||||
ml_admin_bp,
|
||||
thumbnails_bp,
|
||||
sources_bp,
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
"""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()
|
||||
# Concept-crop (SigLIP bag) coverage — how far the back-catalogue embed
|
||||
# has progressed, so the max-over-bag scorer's reach is checkable.
|
||||
images_with_concept_siglip = (
|
||||
await session.execute(
|
||||
select(func.count(distinct(ImageRegion.image_record_id)))
|
||||
.where(ImageRegion.kind == "concept")
|
||||
.where(ImageRegion.siglip_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
|
||||
]
|
||||
auto_applied = (
|
||||
await session.execute(
|
||||
select(func.count()).select_from(image_tag).where(
|
||||
image_tag.c.source == "ccip_auto"
|
||||
)
|
||||
)
|
||||
).scalar_one()
|
||||
return jsonify({
|
||||
"regions_by_kind": by_kind,
|
||||
"images_with_figure_ccip": images_with_figure_ccip,
|
||||
"images_with_concept_siglip": images_with_concept_siglip,
|
||||
"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,
|
||||
"auto_applied": auto_applied,
|
||||
})
|
||||
|
||||
|
||||
@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,220 @@
|
||||
"""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.embedder import MODEL_NAME as EMBED_MODEL_NAME
|
||||
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,
|
||||
# The embedding model the agent must use for concept crops, so
|
||||
# its region vectors land in the SAME space the heads trained in.
|
||||
# Server-announced → the agent stays model-agnostic; a swap is a
|
||||
# server setting + a re-embed migration, never an agent change.
|
||||
"embed_model_name": EMBED_MODEL_NAME,
|
||||
"embed_version": ml.embedder_model_version,
|
||||
})
|
||||
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})
|
||||
|
||||
|
||||
@gpu_bp.route("/jobs/release", methods=["POST"])
|
||||
async def release():
|
||||
"""Graceful stop: the agent hands its still-leased jobs back to pending so
|
||||
they're picked up immediately instead of waiting out the lease."""
|
||||
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).release(agent_id, job_ids)
|
||||
await session.commit()
|
||||
return jsonify({"released": n})
|
||||
@@ -21,6 +21,9 @@ _EDITABLE = (
|
||||
"head_auto_apply_precision",
|
||||
"head_auto_apply_enabled",
|
||||
"head_auto_apply_min_positives",
|
||||
"ccip_match_threshold",
|
||||
"ccip_auto_apply_enabled",
|
||||
"ccip_auto_apply_threshold",
|
||||
)
|
||||
|
||||
|
||||
@@ -48,6 +51,9 @@ async def get_settings():
|
||||
"head_auto_apply_precision": s.head_auto_apply_precision,
|
||||
"head_auto_apply_enabled": s.head_auto_apply_enabled,
|
||||
"head_auto_apply_min_positives": s.head_auto_apply_min_positives,
|
||||
"ccip_match_threshold": s.ccip_match_threshold,
|
||||
"ccip_auto_apply_enabled": s.ccip_auto_apply_enabled,
|
||||
"ccip_auto_apply_threshold": s.ccip_auto_apply_threshold,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -115,6 +121,10 @@ def _validate(p: dict) -> str | None:
|
||||
return "head_auto_apply_precision must be between 0.5 and 0.999"
|
||||
if int(p["head_auto_apply_min_positives"]) < 1:
|
||||
return "head_auto_apply_min_positives must be >= 1"
|
||||
if not (0.5 <= float(p["ccip_match_threshold"]) <= 0.999):
|
||||
return "ccip_match_threshold must be between 0.5 and 0.999"
|
||||
if not (0.5 <= float(p["ccip_auto_apply_threshold"]) <= 0.999):
|
||||
return "ccip_auto_apply_threshold must be between 0.5 and 0.999"
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -117,6 +117,24 @@ def make_celery() -> Celery:
|
||||
"task": "backend.app.tasks.ml.scheduled_apply_head_tags",
|
||||
"schedule": 86400.0, # no-op unless head_auto_apply_enabled
|
||||
},
|
||||
"recover-orphaned-gpu-jobs": {
|
||||
"task": "backend.app.tasks.ml.recover_orphaned_gpu_jobs",
|
||||
"schedule": 60.0, # quick pickup of work a dead agent orphaned
|
||||
},
|
||||
"enqueue-ccip-backfill-hourly": {
|
||||
"task": "backend.app.tasks.ml.enqueue_gpu_backfill",
|
||||
"schedule": 3600.0, # auto-feed new images (+ retry errored) so
|
||||
"args": ("ccip",), # the queue keeps moving without the button
|
||||
},
|
||||
"enqueue-siglip-backfill-daily": {
|
||||
"task": "backend.app.tasks.ml.enqueue_gpu_backfill",
|
||||
"schedule": 86400.0, # drain the concept-crop back-catalogue +
|
||||
"args": ("siglip",), # retry failed embeds, no button needed
|
||||
},
|
||||
"ccip-auto-apply-daily": {
|
||||
"task": "backend.app.tasks.ml.scheduled_ccip_auto_apply",
|
||||
"schedule": 86400.0, # no-op unless ccip_auto_apply_enabled
|
||||
},
|
||||
"snapshot-head-metrics-daily": {
|
||||
"task": "backend.app.tasks.maintenance.snapshot_head_metrics",
|
||||
"schedule": 86400.0,
|
||||
|
||||
@@ -8,6 +8,7 @@ from .base import Base
|
||||
from .credential import Credential
|
||||
from .download_event import DownloadEvent
|
||||
from .external_link import ExternalLink
|
||||
from .gpu_job import GpuJob
|
||||
from .head_auto_apply_run import HeadAutoApplyRun
|
||||
from .head_metric import HeadMetric
|
||||
from .head_metrics_snapshot import HeadMetricsSnapshot
|
||||
@@ -15,6 +16,7 @@ from .head_training_run import HeadTrainingRun
|
||||
from .image_prediction import ImagePrediction
|
||||
from .image_provenance import ImageProvenance
|
||||
from .image_record import ImageRecord
|
||||
from .image_region import ImageRegion
|
||||
from .import_batch import ImportBatch
|
||||
from .import_settings import ImportSettings
|
||||
from .import_task import ImportTask
|
||||
@@ -60,11 +62,13 @@ __all__ = [
|
||||
"ImageRecord",
|
||||
"ImagePrediction",
|
||||
"ImageProvenance",
|
||||
"ImageRegion",
|
||||
"Tag",
|
||||
"TagKind",
|
||||
"image_tag",
|
||||
"DownloadEvent",
|
||||
"ExternalLink",
|
||||
"GpuJob",
|
||||
"ImportBatch",
|
||||
"ImportTask",
|
||||
"ImportSettings",
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
"""GpuJob — a unit of GPU work the desktop agent pulls over HTTP (#114).
|
||||
|
||||
The durable work list that lets the agent stay HTTP-only: the server enqueues a
|
||||
job per (image, task) — e.g. detect figures + CCIP-embed — and the agent LEASES a
|
||||
batch, computes on its GPU, then SUBMITS results, all over the already-exposed web
|
||||
API. Redis/Postgres stay private. A lease has an expiry; the lease query itself
|
||||
re-claims expired leases (agent died / stopped mid-batch), so the queue is
|
||||
self-healing without a separate sweep. One job is per ITEM; the agent fans a
|
||||
VIDEO out into per-frame instances internally (see image_region.frame_time).
|
||||
|
||||
State: pending → leased → done | error (a failure under the attempt cap returns to
|
||||
pending for another agent).
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import DateTime, ForeignKey, Integer, String, Text, func
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from .base import Base
|
||||
|
||||
|
||||
class GpuJob(Base):
|
||||
__tablename__ = "gpu_job"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
||||
image_record_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("image_record.id", ondelete="CASCADE"), index=True
|
||||
)
|
||||
# What to compute, e.g. 'ccip' (detect figures + CCIP-embed) or 'siglip_region'.
|
||||
task: Mapped[str] = mapped_column(String(32), nullable=False)
|
||||
status: Mapped[str] = mapped_column(
|
||||
String(16), nullable=False, default="pending", index=True
|
||||
)
|
||||
# pending | leased | done | error
|
||||
lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
leased_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
lease_expires_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||
error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
@@ -0,0 +1,62 @@
|
||||
"""ImageRegion — a detected/proposed sub-region of an image + its crop embedding.
|
||||
|
||||
The storage backbone of the crop pipeline (#114). A region is a normalized bbox
|
||||
plus the embedding of its crop:
|
||||
- kind='face' / 'figure' → embedded by CCIP for cross-artist character identity.
|
||||
- kind='concept' → embedded by SigLIP, a localized instance for a concept head's
|
||||
bag-of-embeddings (a concept is "present if ANY instance matches").
|
||||
One row carries the embedding appropriate to its kind (the other is null). The
|
||||
bbox doubles as grounded-tag provenance (hover a tag → highlight its region; a
|
||||
wrong box is a precise negative). The GPU agent writes these via the job API;
|
||||
the few-shot character matcher + bag scorer read them — both server-side, no GPU.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import DateTime, Float, ForeignKey, Integer, String, func
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from .base import Base
|
||||
|
||||
CCIP_DIM = 768 # deepghs/imgutils CCIP character embedding
|
||||
SIGLIP_DIM = 1152 # matches image_record.siglip_embedding
|
||||
|
||||
|
||||
class ImageRegion(Base):
|
||||
__tablename__ = "image_region"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
||||
image_record_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("image_record.id", ondelete="CASCADE"), index=True
|
||||
)
|
||||
# 'frame' (a whole video frame → SigLIP bag) | 'face' | 'figure' (→ CCIP
|
||||
# character id) | 'concept' (→ SigLIP head bag).
|
||||
kind: Mapped[str] = mapped_column(String(16), nullable=False)
|
||||
# For video/animated media: the source frame's timestamp in SECONDS. NULL for
|
||||
# static images. Lets a video be a BAG of per-frame instances (fixes the
|
||||
# mean-embedding muddle) + grounds a tag to "appears at 0:42".
|
||||
frame_time: Mapped[float | None] = mapped_column(Float, nullable=True)
|
||||
# Normalized bbox in [0,1]: top-left (rx, ry) + size (rw, rh). Named rx/ry/…
|
||||
# rather than x/y/by to dodge SQL keyword ambiguity ('by').
|
||||
rx: Mapped[float] = mapped_column(Float, nullable=False)
|
||||
ry: Mapped[float] = mapped_column(Float, nullable=False)
|
||||
rw: Mapped[float] = mapped_column(Float, nullable=False)
|
||||
rh: Mapped[float] = mapped_column(Float, nullable=False)
|
||||
# Proposer/detector confidence (null for deterministic proposers).
|
||||
score: Mapped[float | None] = mapped_column(Float, nullable=True)
|
||||
# Version stamps so a re-detect / re-crop / re-embed can be gated (compute
|
||||
# once; only redo when the producing model version changes).
|
||||
detector_version: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
crop_version: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
embedding_version: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
# Exactly one is set, per kind.
|
||||
ccip_embedding: Mapped[list[float] | None] = mapped_column(
|
||||
Vector(CCIP_DIM), nullable=True
|
||||
)
|
||||
siglip_embedding: Mapped[list[float] | None] = mapped_column(
|
||||
Vector(SIGLIP_DIM), nullable=True
|
||||
)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
@@ -86,6 +86,21 @@ class MLSettings(Base):
|
||||
head_auto_apply_min_positives: Mapped[int] = mapped_column(
|
||||
Integer, nullable=False, default=30
|
||||
)
|
||||
# CCIP character-match cosine cut (#114). 0.85 default — the v1 flat 0.75
|
||||
# over-fired (high-reference characters matched a scatter of images); 0.85
|
||||
# keeps the confident single-character matches. Tunable from the agent card.
|
||||
ccip_match_threshold: Mapped[float] = mapped_column(
|
||||
Float, nullable=False, default=0.85
|
||||
)
|
||||
# CCIP auto-apply (#114). Confident matches (>= ccip_auto_apply_threshold,
|
||||
# above the suggest cut) auto-tag on a daily sweep. ON by default (opt-out);
|
||||
# single-character references + the high bar keep it safe, every tag reversible.
|
||||
ccip_auto_apply_enabled: Mapped[bool] = mapped_column(
|
||||
Boolean, nullable=False, default=True
|
||||
)
|
||||
ccip_auto_apply_threshold: Mapped[float] = mapped_column(
|
||||
Float, nullable=False, default=0.92
|
||||
)
|
||||
tagger_model_version: Mapped[str] = mapped_column(
|
||||
String(128), nullable=False, default="camie-tagger-v2"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
"""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 func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ...models import ImageRegion, MLSettings, Tag, TagKind
|
||||
from ...models.tag import image_tag
|
||||
|
||||
# Cosine-similarity floor to call a figure the same character. The live setting
|
||||
# (ml_settings.ccip_match_threshold) drives it; this is only the fallback when no
|
||||
# threshold is supplied AND no settings row exists.
|
||||
DEFAULT_SIM_THRESHOLD = 0.85
|
||||
_FIGURE_KINDS = ("face", "figure")
|
||||
|
||||
|
||||
async def _settings_threshold(session: AsyncSession) -> float:
|
||||
val = (
|
||||
await session.execute(
|
||||
select(MLSettings.ccip_match_threshold).where(MLSettings.id == 1)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
return float(val) if val is not None else DEFAULT_SIM_THRESHOLD
|
||||
|
||||
|
||||
def _l2norm(mat, np):
|
||||
n = np.linalg.norm(mat, axis=1, keepdims=True)
|
||||
n[n == 0] = 1.0
|
||||
return mat / n
|
||||
|
||||
|
||||
# Single-shot cache of the (expensive) reference load, keyed on a cheap
|
||||
# signature that changes exactly when references could: a character tag added/
|
||||
# removed (n_char_tags) or a figure embedded (max/ n of ccip regions). Shared by
|
||||
# the live matcher (every modal open) and the auto-apply sweep.
|
||||
_REF_CACHE: dict = {"sig": None, "refs": None}
|
||||
|
||||
|
||||
def _single_character_images():
|
||||
"""Subquery of image ids carrying EXACTLY ONE character tag. References come
|
||||
only from these — on a multi-character image the tag is image-level, so every
|
||||
figure would otherwise pollute each character's prototype set (a 2-character
|
||||
image tagged 'Velma' would make Daphne's figure a Velma reference)."""
|
||||
return (
|
||||
select(image_tag.c.image_record_id)
|
||||
.join(Tag, Tag.id == image_tag.c.tag_id)
|
||||
.where(Tag.kind == TagKind.character)
|
||||
.group_by(image_tag.c.image_record_id)
|
||||
.having(func.count() == 1)
|
||||
)
|
||||
|
||||
|
||||
async def _ref_signature(session: AsyncSession) -> tuple:
|
||||
n_tags = (
|
||||
await session.execute(
|
||||
select(func.count())
|
||||
.select_from(image_tag)
|
||||
.join(Tag, Tag.id == image_tag.c.tag_id)
|
||||
.where(Tag.kind == TagKind.character)
|
||||
)
|
||||
).scalar_one()
|
||||
n_regs, max_id = (
|
||||
await session.execute(
|
||||
select(func.count(), func.max(ImageRegion.id)).where(
|
||||
ImageRegion.kind.in_(_FIGURE_KINDS),
|
||||
ImageRegion.ccip_embedding.is_not(None),
|
||||
)
|
||||
)
|
||||
).one()
|
||||
return (n_tags, n_regs, max_id)
|
||||
|
||||
|
||||
async def character_references(session: AsyncSession) -> dict[int, list]:
|
||||
"""Per character-tag CCIP reference vectors: figure/face-region CCIP
|
||||
embeddings on UNAMBIGUOUS (single-character) images carrying that tag.
|
||||
Multi-prototype — several vectors per character. Cached on a cheap signature."""
|
||||
sig = await _ref_signature(session)
|
||||
if _REF_CACHE["sig"] == sig and _REF_CACHE["refs"] is not None:
|
||||
return _REF_CACHE["refs"]
|
||||
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))
|
||||
.where(ImageRegion.image_record_id.in_(_single_character_images()))
|
||||
)
|
||||
).all()
|
||||
refs: dict[int, list] = {}
|
||||
for tag_id, vec in rows:
|
||||
refs.setdefault(tag_id, []).append(vec)
|
||||
_REF_CACHE.update(sig=sig, refs=refs)
|
||||
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 | None = None
|
||||
) -> 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. threshold defaults to the
|
||||
live ml_settings.ccip_match_threshold."""
|
||||
import numpy as np
|
||||
|
||||
if threshold is None:
|
||||
threshold = await _settings_threshold(session)
|
||||
|
||||
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,177 @@
|
||||
"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail
|
||||
/ release / recover_orphaned.
|
||||
|
||||
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/workers
|
||||
never grab the same job. Orphan recovery is three-layered: a graceful agent stop
|
||||
calls release() to hand its in-flight jobs back instantly; a hard crash is caught
|
||||
by recover_orphaned() (a 60s beat sweep) which resets expired leases to pending;
|
||||
and the lease itself reclaims expired leases as a final backstop. Result-writing
|
||||
(regions) is done by the API handler via RegionService; complete() just closes.
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from sqlalchemy import and_, or_, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ...models import GpuJob
|
||||
|
||||
# Lease window. Kept comfortably above any single job (a capped-frame video embed
|
||||
# is tens of seconds) so a live, heartbeating worker is never falsely expired,
|
||||
# but short enough that a hard crash recovers fast once the sweep fires.
|
||||
DEFAULT_LEASE_TTL = 180 # 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
|
||||
|
||||
async def release(self, token: str, job_ids: list[int]) -> int:
|
||||
"""Hand the agent's still-leased jobs back to pending NOW (graceful stop),
|
||||
so another worker picks them up immediately instead of waiting out the
|
||||
lease. Scoped to the token's own leases. Returns rows released."""
|
||||
if not job_ids:
|
||||
return 0
|
||||
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(
|
||||
status="pending", lease_token=None, leased_at=None,
|
||||
lease_expires_at=None, updated_at=now,
|
||||
)
|
||||
)
|
||||
return res.rowcount or 0
|
||||
|
||||
async def recover_orphaned(self) -> int:
|
||||
"""Reset every expired lease back to pending — catches agents that died
|
||||
mid-job (no graceful release). Run on a short beat so the queue recovers
|
||||
+ reads honestly even when no worker is actively leasing. Returns rows
|
||||
recovered."""
|
||||
now = datetime.now(UTC)
|
||||
res = await self.session.execute(
|
||||
update(GpuJob)
|
||||
.where(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
|
||||
.values(
|
||||
status="pending", lease_token=None, leased_at=None,
|
||||
lease_expires_at=None, updated_at=now,
|
||||
)
|
||||
)
|
||||
return res.rowcount or 0
|
||||
@@ -29,6 +29,7 @@ from ...models import (
|
||||
HeadAutoApplyRun,
|
||||
HeadTrainingRun,
|
||||
ImageRecord,
|
||||
ImageRegion,
|
||||
MLSettings,
|
||||
Tag,
|
||||
TagHead,
|
||||
@@ -296,7 +297,14 @@ async def score_image(
|
||||
category, score}], ranked. A concept surfaces when its score clears the
|
||||
head's own suggest_threshold — or, when threshold_override is given (the
|
||||
typed-dropdown "show everything" mode), that flat floor instead (0 → every
|
||||
head). Empty if the image has no embedding or no heads exist yet."""
|
||||
head). Empty if the image has no embedding or no heads exist yet.
|
||||
|
||||
MAX-OVER-BAG: the image is scored as a BAG of embeddings — the whole-image
|
||||
vector PLUS every concept-region crop the agent embedded (same model
|
||||
version) — and each head takes its MAX score across the bag. A small/local
|
||||
concept (glasses, a stomach bulge) that the whole-image vector washes out
|
||||
can still surface from the crop where it dominates. The whole-image vector is
|
||||
always in the bag, so this can never score lower than whole-image alone."""
|
||||
import numpy as np
|
||||
|
||||
img = await session.get(ImageRecord, image_id)
|
||||
@@ -306,11 +314,26 @@ async def score_image(
|
||||
heads = await _current_heads(session, settings.embedder_model_version)
|
||||
if heads["W"] is None:
|
||||
return []
|
||||
x = np.asarray(img.siglip_embedding, dtype=np.float32)
|
||||
n = float(np.linalg.norm(x)) or 1.0
|
||||
xn = x / n
|
||||
z = heads["W"] @ xn + heads["b"]
|
||||
probs = 1.0 / (1.0 + np.exp(-z))
|
||||
|
||||
bag = [np.asarray(img.siglip_embedding, dtype=np.float32)]
|
||||
region_vecs = (
|
||||
await session.execute(
|
||||
select(ImageRegion.siglip_embedding)
|
||||
.where(ImageRegion.image_record_id == image_id)
|
||||
.where(ImageRegion.siglip_embedding.is_not(None))
|
||||
.where(ImageRegion.embedding_version == settings.embedder_model_version)
|
||||
)
|
||||
).all()
|
||||
for (vec,) in region_vecs:
|
||||
if vec is not None:
|
||||
bag.append(np.asarray(vec, dtype=np.float32))
|
||||
|
||||
X = np.vstack(bag) # (B, D)
|
||||
norms = np.linalg.norm(X, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1.0
|
||||
Xn = X / norms
|
||||
Z = Xn @ heads["W"].T + heads["b"] # (B, H)
|
||||
probs = (1.0 / (1.0 + np.exp(-Z))).max(axis=0) # (H,) best over the bag
|
||||
out = []
|
||||
for i, p in enumerate(probs):
|
||||
cut = threshold_override if threshold_override is not None else heads["thr"][i]
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Region read/write for the crop pipeline (#114).
|
||||
|
||||
The GPU agent's results endpoint calls replace_regions() to store a freshly
|
||||
detected/embedded set; the character matcher + concept-bag scorer read via
|
||||
get_regions(). Replacement is scoped BY KIND so the figure pipeline and the
|
||||
concept pipeline don't clobber each other.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import delete, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ...models import ImageRegion
|
||||
|
||||
|
||||
class RegionService:
|
||||
def __init__(self, session: AsyncSession):
|
||||
self.session = session
|
||||
|
||||
async def get_regions(
|
||||
self, image_id: int, kinds: list[str] | None = None
|
||||
) -> list[ImageRegion]:
|
||||
stmt = select(ImageRegion).where(ImageRegion.image_record_id == image_id)
|
||||
if kinds:
|
||||
stmt = stmt.where(ImageRegion.kind.in_(kinds))
|
||||
return list(
|
||||
(await self.session.execute(stmt.order_by(ImageRegion.id))).scalars()
|
||||
)
|
||||
|
||||
async def replace_regions(
|
||||
self, image_id: int, kinds: list[str], regions: list[dict[str, Any]]
|
||||
) -> int:
|
||||
"""Replace this image's regions OF THE GIVEN KINDS with `regions` (a
|
||||
re-detect/re-propose supersedes the prior set without touching other
|
||||
kinds). Each region dict: {kind, bbox:(x,y,w,h), score?, detector_version?,
|
||||
crop_version?, embedding_version?, ccip_embedding?, siglip_embedding?}.
|
||||
Returns the number inserted."""
|
||||
await self.session.execute(
|
||||
delete(ImageRegion)
|
||||
.where(ImageRegion.image_record_id == image_id)
|
||||
.where(ImageRegion.kind.in_(kinds))
|
||||
)
|
||||
n = 0
|
||||
for r in regions:
|
||||
rx, ry, rw, rh = r["bbox"]
|
||||
self.session.add(ImageRegion(
|
||||
image_record_id=image_id, kind=r["kind"],
|
||||
frame_time=r.get("frame_time"),
|
||||
rx=rx, ry=ry, rw=rw, rh=rh,
|
||||
score=r.get("score"),
|
||||
detector_version=r.get("detector_version"),
|
||||
crop_version=r.get("crop_version"),
|
||||
embedding_version=r.get("embedding_version"),
|
||||
ccip_embedding=r.get("ccip_embedding"),
|
||||
siglip_embedding=r.get("siglip_embedding"),
|
||||
))
|
||||
n += 1
|
||||
return n
|
||||
@@ -16,6 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ...models import ImageRecord, TagSuggestionRejection
|
||||
from ...models.tag import image_tag
|
||||
from .ccip import match_image as ccip_match_image
|
||||
from .heads import score_image
|
||||
|
||||
|
||||
@@ -27,7 +28,7 @@ class Suggestion:
|
||||
display_name: str
|
||||
category: str
|
||||
score: float
|
||||
source: str # 'head' (Camie 'tagger'/'centroid' sources removed in v2)
|
||||
source: str # 'head' | 'ccip' | 'both' (Camie tagger/centroid removed in v2)
|
||||
creates_new_tag: bool
|
||||
# raw_name = the booru model vocab key behind this suggestion. It's the key
|
||||
# an alias MUST be stored under (resolution looks up the raw key), so the
|
||||
@@ -92,19 +93,39 @@ class SuggestionService:
|
||||
hits = await score_image(
|
||||
self.session, image_id, threshold_override=threshold_override
|
||||
)
|
||||
# CCIP character matches OVERLAY the SigLIP character heads — a
|
||||
# complementary, identity-specialized signal with different failure modes
|
||||
# (CCIP needs a detected figure; heads work whole-image). Merged by tag:
|
||||
# 'both' when they corroborate, taking the higher score.
|
||||
ccip_hits = await ccip_match_image(self.session, image_id)
|
||||
|
||||
merged: dict[tuple[str, int], dict] = {}
|
||||
for h in hits:
|
||||
merged[(h["category"], h["tag_id"])] = {
|
||||
"name": h["name"], "score": h["score"], "source": "head",
|
||||
}
|
||||
for c in ccip_hits:
|
||||
key = ("character", c["tag_id"])
|
||||
ex = merged.get(key)
|
||||
if ex is not None:
|
||||
ex["source"] = "both"
|
||||
ex["score"] = max(ex["score"], c["score"])
|
||||
else:
|
||||
merged[key] = {
|
||||
"name": c["name"], "score": c["score"], "source": "ccip",
|
||||
}
|
||||
|
||||
result = SuggestionList()
|
||||
for h in hits:
|
||||
tag_id = h["tag_id"]
|
||||
for (cat, tag_id), m in merged.items():
|
||||
if tag_id in applied:
|
||||
continue
|
||||
result.by_category.setdefault(h["category"], []).append(
|
||||
result.by_category.setdefault(cat, []).append(
|
||||
Suggestion(
|
||||
canonical_tag_id=tag_id,
|
||||
display_name=h["name"],
|
||||
category=h["category"],
|
||||
score=h["score"],
|
||||
source="head",
|
||||
display_name=m["name"],
|
||||
category=cat,
|
||||
score=m["score"],
|
||||
source=m["source"],
|
||||
creates_new_tag=False,
|
||||
rejected=tag_id in rejected,
|
||||
)
|
||||
|
||||
@@ -738,3 +738,200 @@ def scheduled_apply_head_tags() -> str:
|
||||
run_id = run.id
|
||||
apply_head_tags.delay(run_id)
|
||||
return "dispatched"
|
||||
|
||||
|
||||
@celery.task(name="backend.app.tasks.ml.enqueue_gpu_backfill")
|
||||
def enqueue_gpu_backfill(task_name: str) -> int:
|
||||
"""Enqueue a gpu_job for every image that still needs `task_name` (one
|
||||
INSERT…SELECT, so it scales to a full library). The desktop agent drains the
|
||||
queue over HTTP. Returns the number enqueued.
|
||||
|
||||
'siglip' gates on the RESULT (no concept region yet) rather than on a prior
|
||||
job, so it picks up the back-catalogue of images that were CCIP-embedded
|
||||
before concept crops existed, and retries images whose concept embed failed —
|
||||
without re-touching their figure/CCIP regions."""
|
||||
from sqlalchemy import exists, insert, literal
|
||||
from sqlalchemy import select as sa_select
|
||||
|
||||
from ..models import GpuJob, ImageRecord, ImageRegion
|
||||
|
||||
SessionLocal = _sync_session_factory()
|
||||
with SessionLocal() as session:
|
||||
if task_name == "siglip":
|
||||
has_concept = exists().where(
|
||||
ImageRegion.image_record_id == ImageRecord.id,
|
||||
ImageRegion.kind == "concept",
|
||||
)
|
||||
queued = exists().where(
|
||||
GpuJob.image_record_id == ImageRecord.id,
|
||||
GpuJob.task == "siglip",
|
||||
GpuJob.status.in_(["pending", "leased"]),
|
||||
)
|
||||
sel = sa_select(
|
||||
ImageRecord.id, literal("siglip"), literal("pending")
|
||||
).where(~has_concept).where(~queued)
|
||||
else:
|
||||
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)
|
||||
|
||||
|
||||
@celery.task(name="backend.app.tasks.ml.recover_orphaned_gpu_jobs")
|
||||
def recover_orphaned_gpu_jobs() -> int:
|
||||
"""Reset expired GPU-job leases back to pending — recovers work orphaned by an
|
||||
agent that died mid-job (no graceful release). Short beat cadence so orphans
|
||||
get picked back up quickly + the queue counts read honestly. Returns the
|
||||
number recovered."""
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import update
|
||||
|
||||
from ..models import GpuJob
|
||||
|
||||
SessionLocal = _sync_session_factory()
|
||||
with SessionLocal() as session:
|
||||
now = datetime.now(UTC)
|
||||
res = session.execute(
|
||||
update(GpuJob)
|
||||
.where(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
|
||||
.values(
|
||||
status="pending", lease_token=None, leased_at=None,
|
||||
lease_expires_at=None, updated_at=now,
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
return res.rowcount or 0
|
||||
|
||||
|
||||
@celery.task(
|
||||
name="backend.app.tasks.ml.scheduled_ccip_auto_apply",
|
||||
soft_time_limit=1800, time_limit=2100,
|
||||
)
|
||||
def scheduled_ccip_auto_apply() -> str:
|
||||
"""Auto-tag confident CCIP character matches (source='ccip_auto') so identity
|
||||
tags keep flowing without a button. No-op unless ccip_auto_apply_enabled.
|
||||
References come only from single-character images (unambiguous); a tag is
|
||||
applied where any figure's best cosine to a character's prototypes clears
|
||||
ccip_auto_apply_threshold and it isn't already applied/rejected. Reversible."""
|
||||
import numpy as np
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy import select as sa_select
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
from ..models import ImageRegion, MLSettings, Tag, TagKind, TagSuggestionRejection
|
||||
from ..models.tag import image_tag
|
||||
|
||||
fig = ("face", "figure")
|
||||
|
||||
def _l2(m):
|
||||
n = np.linalg.norm(m, axis=1, keepdims=True)
|
||||
n[n == 0] = 1.0
|
||||
return m / n
|
||||
|
||||
SessionLocal = _sync_session_factory()
|
||||
with SessionLocal() as session:
|
||||
s = session.get(MLSettings, 1)
|
||||
if s is None or not s.ccip_auto_apply_enabled:
|
||||
return "disabled"
|
||||
thr = float(s.ccip_auto_apply_threshold)
|
||||
|
||||
single = (
|
||||
sa_select(image_tag.c.image_record_id)
|
||||
.join(Tag, Tag.id == image_tag.c.tag_id)
|
||||
.where(Tag.kind == TagKind.character)
|
||||
.group_by(image_tag.c.image_record_id)
|
||||
.having(func.count() == 1)
|
||||
)
|
||||
ref_rows = session.execute(
|
||||
sa_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_(fig))
|
||||
.where(ImageRegion.ccip_embedding.is_not(None))
|
||||
.where(ImageRegion.image_record_id.in_(single))
|
||||
).all()
|
||||
if not ref_rows:
|
||||
return "no-references"
|
||||
|
||||
by_char: dict[int, list] = {}
|
||||
for tid, vec in ref_rows:
|
||||
by_char.setdefault(tid, []).append(vec)
|
||||
ref_tags = list(by_char)
|
||||
mats = [_l2(np.asarray(by_char[t], dtype=np.float32)) for t in ref_tags]
|
||||
allref = np.vstack(mats) # (total, 768)
|
||||
seg = np.cumsum([0] + [len(m) for m in mats])[:-1] # per-char start
|
||||
|
||||
# Per character: images that already carry OR rejected the tag — skip.
|
||||
skip = {t: set() for t in ref_tags}
|
||||
for t in ref_tags:
|
||||
for (iid,) in session.execute(
|
||||
sa_select(image_tag.c.image_record_id).where(
|
||||
image_tag.c.tag_id == t
|
||||
)
|
||||
):
|
||||
skip[t].add(iid)
|
||||
for (iid,) in session.execute(
|
||||
sa_select(TagSuggestionRejection.image_record_id).where(
|
||||
TagSuggestionRejection.tag_id == t
|
||||
)
|
||||
):
|
||||
skip[t].add(iid)
|
||||
|
||||
img_ids = list(session.execute(
|
||||
sa_select(ImageRegion.image_record_id)
|
||||
.where(ImageRegion.kind.in_(fig), ImageRegion.ccip_embedding.is_not(None))
|
||||
.distinct()
|
||||
).scalars())
|
||||
|
||||
applied = 0
|
||||
chunk_n = 500
|
||||
for start in range(0, len(img_ids), chunk_n):
|
||||
chunk = img_ids[start:start + chunk_n]
|
||||
rows = session.execute(
|
||||
sa_select(ImageRegion.image_record_id, ImageRegion.ccip_embedding)
|
||||
.where(
|
||||
ImageRegion.image_record_id.in_(chunk),
|
||||
ImageRegion.kind.in_(fig),
|
||||
ImageRegion.ccip_embedding.is_not(None),
|
||||
)
|
||||
).all()
|
||||
by_img: dict[int, list] = {}
|
||||
for iid, vec in rows:
|
||||
by_img.setdefault(iid, []).append(vec)
|
||||
for iid, vecs in by_img.items():
|
||||
q = _l2(np.asarray(vecs, dtype=np.float32)) # (nq, 768)
|
||||
colmax = (q @ allref.T).max(axis=0) # (total,)
|
||||
charmax = np.maximum.reduceat(colmax, seg) # (n_chars,)
|
||||
for ci in np.where(charmax >= thr)[0]:
|
||||
t = ref_tags[int(ci)]
|
||||
if iid in skip[t]:
|
||||
continue
|
||||
skip[t].add(iid)
|
||||
session.execute(
|
||||
pg_insert(image_tag)
|
||||
.values(
|
||||
image_record_id=iid, tag_id=t, source="ccip_auto",
|
||||
)
|
||||
.on_conflict_do_nothing()
|
||||
)
|
||||
applied += 1
|
||||
session.commit()
|
||||
return f"applied={applied}"
|
||||
|
||||
@@ -21,14 +21,14 @@
|
||||
v-show="store.byCategory[cat] && store.byCategory[cat].length"
|
||||
:label="labelFor(cat)" :items="store.byCategory[cat] || []"
|
||||
@accept="onAccept" @alias="onAlias" @remove-alias="onRemoveAlias"
|
||||
@dismiss="store.dismiss" @undismiss="store.undismiss"
|
||||
@dismiss="onDismiss" @undismiss="onUndismiss"
|
||||
/>
|
||||
<SuggestionsCategoryGroup
|
||||
v-if="store.byCategory.general && store.byCategory.general.length"
|
||||
label="General" :items="store.byCategory.general"
|
||||
collapsible :default-open="true"
|
||||
@accept="onAccept" @alias="onAlias" @remove-alias="onRemoveAlias"
|
||||
@dismiss="store.dismiss" @undismiss="store.undismiss"
|
||||
@dismiss="onDismiss" @undismiss="onUndismiss"
|
||||
/>
|
||||
</div>
|
||||
|
||||
@@ -57,9 +57,15 @@ const props = defineProps({
|
||||
// so the same panel refreshes the right surface. See TagPanel.
|
||||
host: { type: Object, default: null },
|
||||
})
|
||||
// 'accepted' lets the parent return focus to the tag input after a suggestion is
|
||||
// applied (operator-asked 2026-06-08).
|
||||
const emit = defineEmits(['accepted'])
|
||||
// 'accepted'/'dismissed' let the parent return focus to the tag input after a
|
||||
// suggestion is accepted OR rejected, so the operator keeps the keyboard flow on
|
||||
// the input without re-clicking (operator-asked 2026-06-08, 2026-06-30).
|
||||
const emit = defineEmits(['accepted', 'dismissed'])
|
||||
|
||||
// Reject (✗) / un-reject (↶): apply the store change, then signal the parent to
|
||||
// re-focus the tag input — same return-to-input behaviour as accept.
|
||||
function onDismiss (s) { store.dismiss(s); emit('dismissed') }
|
||||
function onUndismiss (s) { store.undismiss(s); emit('dismissed') }
|
||||
const store = useSuggestionsStore()
|
||||
const modalStore = useModalStore()
|
||||
const host = props.host || modalStore
|
||||
|
||||
@@ -72,7 +72,7 @@
|
||||
</v-icon>
|
||||
</template>
|
||||
<v-list-item-title>
|
||||
Create "{{ parsedName }}" as {{ parsedKind }}
|
||||
{{ createLabel }}
|
||||
</v-list-item-title>
|
||||
</v-list-item>
|
||||
</template>
|
||||
@@ -178,14 +178,34 @@ watch(query, () => {
|
||||
}, 200)
|
||||
})
|
||||
|
||||
// A same-name character ALREADY exists. Characters are unique by
|
||||
// (name, kind, fandom), so this is still a valid distinct tag in another fandom.
|
||||
const sameNameCharExists = computed(() =>
|
||||
parsedKind.value === 'character' &&
|
||||
hits.value.some(h =>
|
||||
h.kind === 'character' && h.name.toLowerCase() === parsedName.value.toLowerCase(),
|
||||
),
|
||||
)
|
||||
|
||||
const allowCreate = computed(() => {
|
||||
const q = parsedName.value
|
||||
if (!q) return false
|
||||
// Characters disambiguate by fandom, so a same-named character in a DIFFERENT
|
||||
// fandom is a valid new tag — always offer Create (the fandom picker resolves
|
||||
// it; find_or_create is idempotent if you re-pick the same fandom). Other
|
||||
// kinds are unique by (name, kind): an exact match means it already exists.
|
||||
if (parsedKind.value === 'character') return true
|
||||
return !hits.value.some(h =>
|
||||
h.name.toLowerCase() === q.toLowerCase() && h.kind === parsedKind.value,
|
||||
)
|
||||
})
|
||||
|
||||
const createLabel = computed(() =>
|
||||
sameNameCharExists.value
|
||||
? `Create another "${parsedName.value}" character (different fandom)`
|
||||
: `Create "${parsedName.value}" as ${parsedKind.value}`,
|
||||
)
|
||||
|
||||
function scorePct (s) { return `${Math.round(s.score * 100)}%` }
|
||||
|
||||
// This image's suggestions that match the typed query, minus any the server
|
||||
|
||||
@@ -28,6 +28,7 @@
|
||||
:image-id="host.currentImageId"
|
||||
:host="host"
|
||||
@accepted="focusTagInput"
|
||||
@dismissed="focusTagInput"
|
||||
/>
|
||||
|
||||
<!-- @after-leave: when either dialog finishes closing (apply OR cancel),
|
||||
|
||||
@@ -0,0 +1,270 @@
|
||||
<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>
|
||||
|
||||
<v-btn
|
||||
class="mt-3" color="accent" variant="tonal" rounded="pill" size="small"
|
||||
prepend-icon="mdi-crop" :loading="backfillingSiglip" @click="onBackfillSiglip"
|
||||
>Queue concept crops (SigLIP)</v-btn>
|
||||
<p class="fc-muted text-caption mt-2 mb-0">
|
||||
Enqueues every image that doesn't have concept-crop embeddings yet — the
|
||||
localized vectors that help small/local tags (glasses, etc.) surface. New
|
||||
images get these automatically; this catches the back-catalogue.
|
||||
</p>
|
||||
|
||||
<!-- Match strictness -->
|
||||
<div class="fc-section-h mt-5 mb-1">Character-match strictness</div>
|
||||
<div v-if="ml.settings" class="d-flex align-center" style="gap:12px">
|
||||
<v-slider
|
||||
v-model="threshold" :min="0.70" :max="0.95" :step="0.01"
|
||||
color="accent" hide-details density="compact" class="flex-grow-1"
|
||||
:loading="savingThreshold" @end="onSaveThreshold"
|
||||
/>
|
||||
<span class="fc-q__n" style="font-size:16px">{{ threshold.toFixed(2) }}</span>
|
||||
</div>
|
||||
<p class="fc-muted text-caption mt-1 mb-0">
|
||||
How close a figure must be (CCIP cosine) to suggest a character. Higher =
|
||||
stricter — fewer but more confident matches. 0.85 recommended; below ~0.80
|
||||
a heavily-tagged character starts matching everything.
|
||||
</p>
|
||||
|
||||
<!-- Auto-apply -->
|
||||
<div v-if="ml.settings" class="d-flex align-center mt-5" style="gap:12px">
|
||||
<v-switch
|
||||
v-model="autoApply" color="accent" hide-details density="compact"
|
||||
:loading="savingAuto" label="Auto-apply confident matches"
|
||||
@update:model-value="onSaveAuto"
|
||||
/>
|
||||
<v-text-field
|
||||
v-model.number="autoThreshold" type="number" min="0.80" max="0.99"
|
||||
step="0.01" density="compact" hide-details variant="outlined"
|
||||
style="max-width:96px" :disabled="!autoApply" label="at"
|
||||
@change="onSaveAuto"
|
||||
/>
|
||||
</div>
|
||||
<p class="fc-muted text-caption mt-1 mb-0">
|
||||
When on, a very-confident character match tags the image on its own (daily,
|
||||
reversible) — so identity tags keep flowing without review. Stricter than
|
||||
the suggest cut; 0.92 recommended.
|
||||
</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 { useMLStore } from '../../stores/ml.js'
|
||||
import { copyText } from '../../utils/clipboard.js'
|
||||
|
||||
const store = useGpuStore()
|
||||
const ml = useMLStore()
|
||||
const loading = ref(true)
|
||||
const tokenValue = ref(null)
|
||||
const masked = ref(true)
|
||||
const rotating = ref(false)
|
||||
const backfilling = ref(false)
|
||||
const backfillingSiglip = ref(false)
|
||||
const threshold = ref(0.85)
|
||||
const savingThreshold = ref(false)
|
||||
const autoApply = ref(true)
|
||||
const autoThreshold = ref(0.92)
|
||||
const savingAuto = 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)
|
||||
try {
|
||||
await ml.loadSettings()
|
||||
if (ml.settings?.ccip_match_threshold != null) {
|
||||
threshold.value = ml.settings.ccip_match_threshold
|
||||
}
|
||||
if (ml.settings?.ccip_auto_apply_enabled != null) {
|
||||
autoApply.value = ml.settings.ccip_auto_apply_enabled
|
||||
autoThreshold.value = ml.settings.ccip_auto_apply_threshold
|
||||
}
|
||||
} catch { /* non-fatal */ }
|
||||
})
|
||||
|
||||
async function onSaveAuto() {
|
||||
savingAuto.value = true
|
||||
try {
|
||||
await ml.patchSettings({
|
||||
ccip_auto_apply_enabled: autoApply.value,
|
||||
ccip_auto_apply_threshold: autoThreshold.value,
|
||||
})
|
||||
toast({ text: 'Auto-apply settings saved', type: 'success' })
|
||||
} catch (e) {
|
||||
toast({ text: `Could not save: ${e.message}`, type: 'error' })
|
||||
} finally {
|
||||
savingAuto.value = false
|
||||
}
|
||||
}
|
||||
onUnmounted(() => { if (pollTimer) clearInterval(pollTimer) })
|
||||
|
||||
async function onSaveThreshold() {
|
||||
savingThreshold.value = true
|
||||
try {
|
||||
await ml.patchSettings({ ccip_match_threshold: threshold.value })
|
||||
toast({ text: `Match strictness set to ${threshold.value.toFixed(2)}`, type: 'success' })
|
||||
} catch (e) {
|
||||
toast({ text: `Could not save: ${e.message}`, type: 'error' })
|
||||
} finally {
|
||||
savingThreshold.value = false
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
async function onBackfillSiglip() {
|
||||
backfillingSiglip.value = true
|
||||
try {
|
||||
await store.backfill('siglip')
|
||||
toast({ text: 'Queued concept crops — run the agent to process them', type: 'success' })
|
||||
await refreshQueue()
|
||||
} catch (e) {
|
||||
toast({ text: `Could not queue backfill: ${e.message}`, type: 'error' })
|
||||
} finally {
|
||||
backfillingSiglip.value = false
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
<style scoped>
|
||||
.fc-muted { color: rgb(var(--v-theme-on-surface-variant)); }
|
||||
.fc-section-h {
|
||||
font-size: 13px; font-weight: 700; letter-spacing: 0.03em;
|
||||
text-transform: uppercase; color: rgb(var(--v-theme-on-surface));
|
||||
}
|
||||
.fc-token {
|
||||
display: flex; align-items: center; gap: 4px;
|
||||
background: rgb(var(--v-theme-surface-light)); border-radius: 6px;
|
||||
padding: 4px 6px 4px 10px;
|
||||
}
|
||||
.fc-token__val {
|
||||
font-family: 'JetBrains Mono', monospace; font-size: 13px;
|
||||
overflow: hidden; text-overflow: ellipsis; white-space: nowrap;
|
||||
}
|
||||
.fc-queue { display: flex; gap: 24px; }
|
||||
.fc-q__n {
|
||||
font-size: 20px; font-weight: 700; line-height: 1.1;
|
||||
font-family: 'JetBrains Mono', monospace;
|
||||
}
|
||||
.fc-q__l {
|
||||
font-size: 11px; text-transform: uppercase; letter-spacing: 0.04em;
|
||||
color: rgb(var(--v-theme-on-surface-variant));
|
||||
}
|
||||
.fc-good { color: rgb(var(--v-theme-success)); }
|
||||
.fc-weak { color: rgb(var(--v-theme-error)); }
|
||||
</style>
|
||||
@@ -27,6 +27,7 @@
|
||||
<div class="fc-tile-stack">
|
||||
<MLThresholdSliders />
|
||||
<HeadsCard />
|
||||
<GpuAgentCard />
|
||||
<AllowlistTable />
|
||||
<AliasTable />
|
||||
<TagEvalCard />
|
||||
@@ -54,6 +55,7 @@ import MissingFileRepairCard from './MissingFileRepairCard.vue'
|
||||
import DbMaintenanceCard from './DbMaintenanceCard.vue'
|
||||
import MLThresholdSliders from './MLThresholdSliders.vue'
|
||||
import HeadsCard from './HeadsCard.vue'
|
||||
import GpuAgentCard from './GpuAgentCard.vue'
|
||||
import AllowlistTable from './AllowlistTable.vue'
|
||||
import AliasTable from './AliasTable.vue'
|
||||
import TagEvalCard from './TagEvalCard.vue'
|
||||
|
||||
@@ -11,7 +11,7 @@ import { toast } from '../utils/toast.js'
|
||||
// trail. The store ALSO acts as a TagPanel "host" (current/currentImageId +
|
||||
// tag CRUD over the anchor) so the Explore workspace reuses the modal's tag
|
||||
// rail verbatim for modal-parity tagging while rabbit-holing.
|
||||
const NEIGHBOR_LIMIT = 24
|
||||
const NEIGHBOR_LIMIT = 40 // a wider pool → more variety to browse + jump into
|
||||
|
||||
export const useExploreStore = defineStore('explore', () => {
|
||||
const api = useApi()
|
||||
@@ -81,16 +81,26 @@ export const useExploreStore = defineStore('explore', () => {
|
||||
return cursor.value > 0 ? breadcrumb.value[cursor.value - 1].id : null
|
||||
}
|
||||
|
||||
// → target: the next already-visited crumb if we'd stepped back, else a
|
||||
// RANDOM neighbour to keep the rabbit-hole going. Null if neither exists.
|
||||
// → target: after a ←, walk forward through the already-visited trail
|
||||
// (browser-style). Otherwise jump to a varied neighbour to keep the
|
||||
// rabbit-hole going — null if neither exists.
|
||||
function forwardTarget () {
|
||||
if (cursor.value >= 0 && cursor.value < breadcrumb.value.length - 1) {
|
||||
return breadcrumb.value[cursor.value + 1].id
|
||||
}
|
||||
if (neighbors.value.length) {
|
||||
return neighbors.value[Math.floor(Math.random() * neighbors.value.length)].id
|
||||
}
|
||||
return null
|
||||
if (!neighbors.value.length) return null
|
||||
// Prefer UNVISITED neighbours so → opens something new instead of landing on
|
||||
// a crumb (which snaps the cursor back into the trail — the "loops back"
|
||||
// report). Fall back to the full set only if every neighbour's been seen.
|
||||
const seen = new Set(breadcrumb.value.map((c) => c.id))
|
||||
let pool = neighbors.value.filter((n) => !seen.has(n.id))
|
||||
if (!pool.length) pool = neighbors.value
|
||||
// neighbors come similarity-sorted (nearest first). Skip the closest slice —
|
||||
// those near-duplicates are exactly what you get stuck cycling through — and
|
||||
// pick from the more-varied remainder, for real variance in the walk.
|
||||
const skip = pool.length >= 6 ? Math.floor(pool.length / 3) : 0
|
||||
const cands = pool.slice(skip)
|
||||
return cands[Math.floor(Math.random() * cands.length)].id
|
||||
}
|
||||
|
||||
function reset () {
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
import { defineStore } from 'pinia'
|
||||
|
||||
import { useApi } from '../composables/useApi.js'
|
||||
|
||||
// GPU agent control surface (#114): the FC-side admin for the desktop agent —
|
||||
// the bearer token it authenticates with, the job-queue depth, and the backfill
|
||||
// trigger. The agent itself talks to /api/gpu/jobs/* over HTTP; nothing here
|
||||
// touches Redis/Postgres directly.
|
||||
export const useGpuStore = defineStore('gpu', () => {
|
||||
const api = useApi()
|
||||
|
||||
// { token: <string|null>, configured: bool }
|
||||
async function token() {
|
||||
return await api.get('/api/gpu/token')
|
||||
}
|
||||
|
||||
// Generate a fresh token (invalidates the old one). Returns { token }.
|
||||
async function rotateToken() {
|
||||
return await api.post('/api/gpu/token/rotate')
|
||||
}
|
||||
|
||||
// { pending, leased, done, error }
|
||||
async function status() {
|
||||
return await api.get('/api/gpu/status')
|
||||
}
|
||||
|
||||
// Enqueue a job per image lacking one for `task` (the agent drains it).
|
||||
async function backfill(task = 'ccip') {
|
||||
return await api.post('/api/gpu/backfill', { body: { task } })
|
||||
}
|
||||
|
||||
return { token, rotateToken, status, backfill }
|
||||
})
|
||||
@@ -90,12 +90,22 @@
|
||||
<!-- CENTER: the focused image (light viewer) + meta. -->
|
||||
<section class="fc-ex__viewer">
|
||||
<div class="fc-ex__canvas">
|
||||
<template v-if="store.anchor">
|
||||
<!-- Videos can't render in an <img> — branch to VideoCanvas like
|
||||
the modal does (an MP4 in ImageCanvas just shows the alt). -->
|
||||
<ImageCanvas
|
||||
v-if="store.anchor"
|
||||
v-if="!isVideo"
|
||||
:key="store.anchor.id"
|
||||
:src="store.anchor.image_url"
|
||||
:alt="`Image ${store.anchor.id}`"
|
||||
/>
|
||||
<VideoCanvas
|
||||
v-else
|
||||
:key="store.anchor.id"
|
||||
:src="store.anchor.image_url"
|
||||
:mime="store.anchor.mime"
|
||||
/>
|
||||
</template>
|
||||
</div>
|
||||
<div v-if="store.anchor" class="fc-ex__viewer-foot">
|
||||
<div class="fc-ex__artist">{{ store.anchor.artist?.name || 'Unknown artist' }}</div>
|
||||
@@ -129,6 +139,7 @@ import { useModalStore } from '../stores/modal.js'
|
||||
import { useHeadTraining } from '../composables/useHeadTraining.js'
|
||||
import { isTextEntry } from '../utils/textEntry.js'
|
||||
import ImageCanvas from '../components/modal/ImageCanvas.vue'
|
||||
import VideoCanvas from '../components/modal/VideoCanvas.vue'
|
||||
import ImageMetaBar from '../components/modal/ImageMetaBar.vue'
|
||||
import ProvenancePanel from '../components/modal/ProvenancePanel.vue'
|
||||
import TagPanel from '../components/modal/TagPanel.vue'
|
||||
@@ -140,6 +151,7 @@ const store = useExploreStore()
|
||||
const modal = useModalStore()
|
||||
|
||||
const anchorId = computed(() => route.params.imageId || null)
|
||||
const isVideo = computed(() => !!store.anchor?.mime?.startsWith('video/'))
|
||||
const seeding = ref(false)
|
||||
const seedError = ref(null)
|
||||
const tagPanelRef = ref(null)
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
"""CCIP/region observability API (#114) — coverage overview + per-image detail."""
|
||||
import pytest
|
||||
|
||||
from backend.app.models import ImageRecord, ImageRegion, TagKind
|
||||
from backend.app.models.tag import image_tag
|
||||
from backend.app.services.tag_service import TagService
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def _ccip(slot: int) -> list[float]:
|
||||
v = [0.0] * 768
|
||||
v[slot] = 1.0
|
||||
return v
|
||||
|
||||
|
||||
async def _img(db, sha) -> ImageRecord:
|
||||
img = ImageRecord(
|
||||
path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
|
||||
width=1, height=1, origin="imported_filesystem", integrity_status="unknown",
|
||||
)
|
||||
db.add(img)
|
||||
await db.flush()
|
||||
return img
|
||||
|
||||
|
||||
async def _figure(db, image_id, ccip):
|
||||
db.add(ImageRegion(
|
||||
image_record_id=image_id, kind="figure", rx=0.0, ry=0.0, rw=1.0, rh=1.0,
|
||||
ccip_embedding=ccip, embedding_version="ccip-test",
|
||||
))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overview_reports_coverage(client, db):
|
||||
raven = await TagService(db).find_or_create("Raven", TagKind.character)
|
||||
ref = await _img(db, "a" * 64)
|
||||
await _figure(db, ref.id, _ccip(0))
|
||||
await db.execute(image_tag.insert().values(
|
||||
image_record_id=ref.id, tag_id=raven.id, source="manual",
|
||||
))
|
||||
q = await _img(db, "b" * 64)
|
||||
await _figure(db, q.id, _ccip(0))
|
||||
await db.commit()
|
||||
|
||||
body = await (await client.get("/api/ccip/overview")).get_json()
|
||||
assert body["regions_by_kind"].get("figure", 0) >= 2
|
||||
assert body["images_with_figure_ccip"] >= 2
|
||||
assert any(
|
||||
c["name"] == "Raven" and c["n_refs"] >= 1
|
||||
for c in body["character_references"]
|
||||
)
|
||||
assert "ccip-test" in body["embedding_versions"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_detail_shows_regions_and_matches(client, db):
|
||||
raven = await TagService(db).find_or_create("Raven", TagKind.character)
|
||||
ref = await _img(db, "c" * 64)
|
||||
await _figure(db, ref.id, _ccip(0))
|
||||
await db.execute(image_tag.insert().values(
|
||||
image_record_id=ref.id, tag_id=raven.id, source="manual",
|
||||
))
|
||||
q = await _img(db, "d" * 64)
|
||||
await _figure(db, q.id, _ccip(0))
|
||||
await db.commit()
|
||||
|
||||
body = await (await client.get(f"/api/ccip/images/{q.id}")).get_json()
|
||||
assert len(body["regions"]) == 1
|
||||
r = body["regions"][0]
|
||||
assert r["kind"] == "figure" and r["has_ccip"] is True and r["has_siglip"] is False
|
||||
assert any(m["tag_id"] == raven.id for m in body["ccip_matches"])
|
||||
@@ -0,0 +1,117 @@
|
||||
"""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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_hands_job_back_to_pending(client, db):
|
||||
img = await _img(db, "e" * 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]
|
||||
|
||||
resp = await client.post("/api/gpu/jobs/release", json={
|
||||
"agent_id": "a1", "job_ids": [j["job_id"]],
|
||||
}, headers=hdr)
|
||||
assert resp.status_code == 200 and (await resp.get_json())["released"] == 1
|
||||
st = await (await client.get("/api/gpu/status")).get_json()
|
||||
assert st["pending"] == 1 and st["leased"] == 0
|
||||
@@ -0,0 +1,143 @@
|
||||
"""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 sqlalchemy import select
|
||||
|
||||
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) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_threshold_gates_borderline_match(db):
|
||||
# A figure ~0.9 cosine from the reference: matched at 0.85, dropped at 0.95.
|
||||
raven = await TagService(db).find_or_create("Raven", TagKind.character)
|
||||
ref = await _img(db, "h" * 64)
|
||||
await _figure(db, ref.id, _ccip(0)) # e0
|
||||
await _tag_image(db, ref.id, raven.id)
|
||||
near = [0.0] * 768
|
||||
near[0], near[1] = 0.9, 0.4359 # |·|=1, cos(e0)=0.9
|
||||
query = await _img(db, "i" * 64)
|
||||
await _figure(db, query.id, near)
|
||||
await db.commit()
|
||||
|
||||
assert any(m["tag_id"] == raven.id for m in await match_image(db, query.id, 0.85))
|
||||
assert await match_image(db, query.id, 0.95) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_character_image_not_used_as_reference(db):
|
||||
# A figure on a 2-character image is ambiguous (tag is image-level), so it
|
||||
# must NOT seed either character's prototypes — else it'd match both.
|
||||
raven = await TagService(db).find_or_create("Raven", TagKind.character)
|
||||
daphne = await TagService(db).find_or_create("Daphne", TagKind.character)
|
||||
multi = await _img(db, "j" * 64)
|
||||
await _figure(db, multi.id, _ccip(0))
|
||||
await _tag_image(db, multi.id, raven.id)
|
||||
await _tag_image(db, multi.id, daphne.id)
|
||||
query = await _img(db, "k" * 64)
|
||||
await _figure(db, query.id, _ccip(0)) # identical to the ambiguous figure
|
||||
await db.commit()
|
||||
assert await match_image(db, query.id) == [] # no clean references → nothing
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_apply_tags_confident_match(db):
|
||||
raven = await TagService(db).find_or_create("Raven", TagKind.character)
|
||||
ref = await _img(db, "l" * 64)
|
||||
await _figure(db, ref.id, _ccip(0))
|
||||
await _tag_image(db, ref.id, raven.id) # single-character reference
|
||||
query = await _img(db, "m" * 64)
|
||||
await _figure(db, query.id, _ccip(0)) # identical → cosine 1.0
|
||||
await db.commit()
|
||||
|
||||
from backend.app.tasks.ml import scheduled_ccip_auto_apply
|
||||
assert "applied=" in scheduled_ccip_auto_apply() # sync task, own session
|
||||
|
||||
rows = (await db.execute(
|
||||
select(image_tag.c.tag_id, image_tag.c.source).where(
|
||||
image_tag.c.image_record_id == query.id
|
||||
)
|
||||
)).all()
|
||||
assert (raven.id, "ccip_auto") in [(t, s) for t, s in rows]
|
||||
@@ -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,197 @@
|
||||
"""GPU-job queue engine (#114): enqueue dedupe + lease/heartbeat/complete/fail."""
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import func, select
|
||||
|
||||
from backend.app.models import GpuJob, ImageRecord, ImageRegion
|
||||
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_siglip_backfill_gates_on_concept_region(db):
|
||||
# 'siglip' backfill enqueues images that lack a concept region (the
|
||||
# back-catalogue) and skips ones that already have one — and never double-
|
||||
# enqueues an image that already has a pending siglip job.
|
||||
from backend.app.tasks.ml import enqueue_gpu_backfill
|
||||
|
||||
need = await _img(db, "e1" * 32) # no concept region → wants one
|
||||
have = await _img(db, "e2" * 32) # already embedded → skip
|
||||
db.add(ImageRegion(
|
||||
image_record_id=have.id, kind="concept", rx=0.0, ry=0.0, rw=1.0, rh=1.0,
|
||||
siglip_embedding=[0.0] * 1152, embedding_version="siglip-test",
|
||||
))
|
||||
await db.commit()
|
||||
|
||||
assert enqueue_gpu_backfill("siglip") >= 1
|
||||
|
||||
queued = {
|
||||
j.image_record_id for j in (
|
||||
await db.execute(select(GpuJob).where(GpuJob.task == "siglip"))
|
||||
).scalars()
|
||||
}
|
||||
assert need.id in queued
|
||||
assert have.id not in queued
|
||||
|
||||
# Idempotent: the now-pending job means a second run doesn't re-enqueue it.
|
||||
enqueue_gpu_backfill("siglip")
|
||||
n_for_need = (
|
||||
await db.execute(
|
||||
select(func.count()).select_from(GpuJob).where(
|
||||
GpuJob.task == "siglip", GpuJob.image_record_id == need.id
|
||||
)
|
||||
)
|
||||
).scalar_one()
|
||||
assert n_for_need == 1
|
||||
|
||||
|
||||
@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"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_returns_to_pending(db):
|
||||
img = await _img(db, "01" + "a" * 62)
|
||||
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.release("other", [job.id]) == 0 # not this token's lease
|
||||
assert await svc.release("agent-1", [job.id]) == 1 # graceful hand-back
|
||||
await db.commit()
|
||||
fresh = await db.get(GpuJob, job.id)
|
||||
assert fresh.status == "pending" and fresh.lease_token is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_orphaned_resets_only_expired(db):
|
||||
img1 = await _img(db, "02" + "a" * 62)
|
||||
img2 = await _img(db, "03" + "a" * 62)
|
||||
svc = GpuJobService(db)
|
||||
await svc.enqueue(img1.id, "ccip")
|
||||
await svc.enqueue(img2.id, "ccip")
|
||||
await db.commit()
|
||||
expired, fresh = await svc.lease("dead", batch_size=2)
|
||||
# One lease is in the past (orphaned), the other still valid.
|
||||
expired.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10)
|
||||
await db.commit()
|
||||
|
||||
assert await svc.recover_orphaned() == 1
|
||||
await db.commit()
|
||||
assert (await db.get(GpuJob, expired.id)).status == "pending"
|
||||
assert (await db.get(GpuJob, fresh.id)).status == "leased" # untouched
|
||||
@@ -4,7 +4,7 @@ scikit-learn, ml image only); scoring is numpy-only (available via pgvector)."""
|
||||
import pytest
|
||||
from sqlalchemy import select
|
||||
|
||||
from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind
|
||||
from backend.app.models import ImageRecord, ImageRegion, MLSettings, TagHead, TagKind
|
||||
from backend.app.models.tag import image_tag
|
||||
from backend.app.services.ml.allowlist import AllowlistService
|
||||
from backend.app.services.ml.suggestions import SuggestionService
|
||||
@@ -111,6 +111,40 @@ async def test_threshold_override_surfaces_below_cut(db):
|
||||
assert any(s.canonical_tag_id == tag.id for s in flooded.by_category["general"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concept_region_surfaces_via_max_over_bag(db):
|
||||
# Max-over-bag: the whole-image vector is orthogonal to the head (scores the
|
||||
# 0.5 midpoint, under a 0.7 cut → nothing), but a concept CROP that aligns
|
||||
# with the head lifts the max over the bag above the cut. A small/local
|
||||
# concept surfaces ONLY because of the crop.
|
||||
tag = await TagService(db).find_or_create("glasses", TagKind.general)
|
||||
img = await _img(db, "b1" * 32, _emb(5)) # whole-image ⟂ head
|
||||
await _head(db, tag.id, slot=0, suggest_threshold=0.7)
|
||||
await db.commit()
|
||||
# Whole-image alone: sigmoid(0)=0.5 < 0.7 → no suggestion.
|
||||
assert not (await SuggestionService(db).for_image(img.id)).by_category.get("general")
|
||||
|
||||
# A concept crop aligned with the head, but stamped with a STALE model
|
||||
# version → filtered out of the bag, so still nothing.
|
||||
db.add(ImageRegion(
|
||||
image_record_id=img.id, kind="concept",
|
||||
rx=0.1, ry=0.1, rw=0.3, rh=0.3,
|
||||
siglip_embedding=_emb(0), embedding_version="stale-embedder-v0",
|
||||
))
|
||||
await db.commit()
|
||||
assert not (await SuggestionService(db).for_image(img.id)).by_category.get("general")
|
||||
|
||||
# A matching-version concept crop → max-over-bag lifts it over the cut.
|
||||
db.add(ImageRegion(
|
||||
image_record_id=img.id, kind="concept",
|
||||
rx=0.4, ry=0.4, rw=0.3, rh=0.3,
|
||||
siglip_embedding=_emb(0), embedding_version=await _embver(db),
|
||||
))
|
||||
await db.commit()
|
||||
general = (await SuggestionService(db).for_image(img.id)).by_category["general"]
|
||||
assert any(s.canonical_tag_id == tag.id and s.score > 0.7 for s in general)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejected_tag_surfaced_flagged_then_reversible(db):
|
||||
# A dismissed suggestion is NOT dropped: it stays flagged rejected so the
|
||||
@@ -131,3 +165,35 @@ async def test_rejected_tag_surfaced_flagged_then_reversible(db):
|
||||
sl2 = await SuggestionService(db).for_image(img.id)
|
||||
s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id)
|
||||
assert s2.rejected is False
|
||||
|
||||
|
||||
async def _figure(db, image_id, slot):
|
||||
v = [0.0] * 768
|
||||
v[slot] = 1.0
|
||||
db.add(ImageRegion(
|
||||
image_record_id=image_id, kind="figure",
|
||||
rx=0.0, ry=0.0, rw=1.0, rh=1.0,
|
||||
ccip_embedding=v, embedding_version="ccip-test",
|
||||
))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ccip_character_surfaces_in_rail(db):
|
||||
# A character with a CCIP reference (a tagged figure) is suggested on a new
|
||||
# image whose figure matches — overlaid into the rail alongside the heads.
|
||||
raven = await TagService(db).find_or_create("Raven", TagKind.character)
|
||||
ref = await _img(db, "0" * 64, None) # the operator's tagged example
|
||||
await _figure(db, ref.id, slot=0)
|
||||
await db.execute(image_tag.insert().values(
|
||||
image_record_id=ref.id, tag_id=raven.id, source="manual",
|
||||
))
|
||||
query = await _img(db, "1" * 64, None) # untagged, matching figure
|
||||
await _figure(db, query.id, slot=0)
|
||||
await db.commit()
|
||||
|
||||
sl = await SuggestionService(db).for_image(query.id)
|
||||
m = next(
|
||||
c for c in sl.by_category.get("character", [])
|
||||
if c.canonical_tag_id == raven.id
|
||||
)
|
||||
assert m.source == "ccip"
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Region storage/service for the crop pipeline (#114)."""
|
||||
import pytest
|
||||
|
||||
from backend.app.models import ImageRecord
|
||||
from backend.app.services.ml.regions import RegionService
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
async def _img(db, sha) -> ImageRecord:
|
||||
img = ImageRecord(
|
||||
path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
|
||||
width=1, height=1, origin="imported_filesystem", integrity_status="unknown",
|
||||
)
|
||||
db.add(img)
|
||||
await db.flush()
|
||||
return img
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_and_get_regions(db):
|
||||
img = await _img(db, "a" * 64)
|
||||
svc = RegionService(db)
|
||||
n = await svc.replace_regions(img.id, ["figure"], [
|
||||
{"kind": "figure", "bbox": (0.1, 0.1, 0.3, 0.4),
|
||||
"score": 0.9, "detector_version": "det-v1", "frame_time": 42.5},
|
||||
])
|
||||
await db.commit()
|
||||
assert n == 1
|
||||
regs = await svc.get_regions(img.id)
|
||||
assert len(regs) == 1
|
||||
r = regs[0]
|
||||
assert r.kind == "figure"
|
||||
assert r.rw == pytest.approx(0.3) and r.rh == pytest.approx(0.4)
|
||||
assert r.score == pytest.approx(0.9)
|
||||
assert r.frame_time == pytest.approx(42.5) # video frame timestamp
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_is_scoped_by_kind(db):
|
||||
img = await _img(db, "b" * 64)
|
||||
svc = RegionService(db)
|
||||
await svc.replace_regions(img.id, ["figure"], [
|
||||
{"kind": "figure", "bbox": (0.0, 0.0, 0.5, 0.5)},
|
||||
])
|
||||
await svc.replace_regions(img.id, ["concept"], [
|
||||
{"kind": "concept", "bbox": (0.5, 0.5, 0.2, 0.2)},
|
||||
])
|
||||
await db.commit()
|
||||
# Re-running the figure detector must NOT wipe the concept region.
|
||||
await svc.replace_regions(img.id, ["figure"], [
|
||||
{"kind": "figure", "bbox": (0.1, 0.1, 0.4, 0.4)},
|
||||
])
|
||||
await db.commit()
|
||||
kinds = sorted(r.kind for r in await svc.get_regions(img.id))
|
||||
assert kinds == ["concept", "figure"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ccip_vector_round_trips(db):
|
||||
img = await _img(db, "c" * 64)
|
||||
svc = RegionService(db)
|
||||
await svc.replace_regions(img.id, ["figure"], [
|
||||
{"kind": "figure", "bbox": (0.0, 0.0, 0.5, 0.5),
|
||||
"ccip_embedding": [0.1] * 768, "embedding_version": "ccip-test"},
|
||||
])
|
||||
await db.commit()
|
||||
r = (await svc.get_regions(img.id, kinds=["figure"]))[0]
|
||||
assert r.ccip_embedding is not None
|
||||
assert len(list(r.ccip_embedding)) == 768
|
||||
assert r.siglip_embedding is None
|
||||
Reference in New Issue
Block a user