feat(tagging): SigLIP concept crops + max-over-bag scoring (#114) #153

Merged
bvandeusen merged 2 commits from feat/siglip-concept-crops into dev 2026-06-30 08:40:24 -04:00
15 changed files with 434 additions and 41 deletions
+6 -1
View File
@@ -10,11 +10,16 @@ RUN apt-get update \
&& rm -rf /var/lib/apt/lists/* && rm -rf /var/lib/apt/lists/*
WORKDIR /app 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 . COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt RUN pip3 install --no-cache-dir -r requirements.txt
COPY fc_agent ./fc_agent COPY fc_agent ./fc_agent
# imgutils caches downloaded ONNX models here; mount a volume to persist them. # 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 ENV HF_HOME=/models
EXPOSE 8770 EXPOSE 8770
+13
View File
@@ -10,6 +10,13 @@
# 4. Open http://localhost:8770 → Start. Pause/Stop hands the GPU back. # 4. Open http://localhost:8770 → Start. Pause/Stop hands the GPU back.
# docker compose down to stop the container entirely. # 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. # Needs the NVIDIA Container Toolkit installed on the host for --gpus.
services: services:
@@ -24,6 +31,12 @@ services:
CCIP_MODEL: ${CCIP_MODEL:-} CCIP_MODEL: ${CCIP_MODEL:-}
DETECTOR_LEVEL: ${DETECTOR_LEVEL:-m} DETECTOR_LEVEL: ${DETECTOR_LEVEL:-m}
BATCH_SIZE: ${BATCH_SIZE:-4} 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: volumes:
# Persist the downloaded ONNX models so restarts are fast. # Persist the downloaded ONNX models so restarts are fast.
- fc-agent-models:/models - fc-agent-models:/models
+18 -1
View File
@@ -16,6 +16,16 @@ worker = Worker(cfg)
app = FastAPI(title="FabledCurator GPU agent") 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) @app.get("/", response_class=HTMLResponse)
def index() -> str: def index() -> str:
return _PAGE return _PAGE
@@ -86,6 +96,10 @@ _PAGE = """<!doctype html><html><head><meta charset=utf-8>
<span class=stat><span class=n id=active>0</span><br>active now</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=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=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>
<div class=gpu id=gpu>GPU — …</div> <div class=gpu id=gpu>GPU — …</div>
<div class=bar><i id=gpubar style=width:0%></i></div> <div class=bar><i id=gpubar style=width:0%></i></div>
@@ -103,7 +117,10 @@ _PAGE = """<!doctype html><html><head><meta charset=utf-8>
const s=await (await fetch('/status')).json() const s=await (await fetch('/status')).json()
CAP=s.max_concurrency||8; capn.textContent=CAP CAP=s.max_concurrency||8; capn.textContent=CAP
state.textContent=s.state; active.textContent=s.active; done.textContent=s.processed state.textContent=s.state; active.textContent=s.active; done.textContent=s.processed
err.textContent=s.errors; fc.textContent=s.fc_url 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 if(document.activeElement!==conc) conc.value=s.concurrency
conc.max=CAP conc.max=CAP
cfg.textContent=s.configured?'set':'MISSING' cfg.textContent=s.configured?'set':'MISSING'
+8
View File
@@ -13,6 +13,11 @@ class Config:
ccip_model: str # imgutils CCIP model name ("" → imgutils default) ccip_model: str # imgutils CCIP model name ("" → imgutils default)
detector_level: str # imgutils person-detector level: n|s|m|x detector_level: str # imgutils person-detector level: n|s|m|x
poll_idle_seconds: float # wait between empty leases 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 @classmethod
def from_env(cls) -> "Config": def from_env(cls) -> "Config":
@@ -25,4 +30,7 @@ class Config:
ccip_model=os.environ.get("CCIP_MODEL", ""), ccip_model=os.environ.get("CCIP_MODEL", ""),
detector_level=os.environ.get("DETECTOR_LEVEL", "m"), detector_level=os.environ.get("DETECTOR_LEVEL", "m"),
poll_idle_seconds=float(os.environ.get("POLL_IDLE_SECONDS", "10")), 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"),
) )
+69
View File
@@ -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()
+137 -19
View File
@@ -10,18 +10,43 @@ 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. orphaned work is re-picked at once rather than waiting out the lease.
""" """
import threading import threading
import time
import requests
from . import media, models from . import media, models
from .client import FcClient from .client import FcClient
from .config import Config from .config import Config
from .crops import crop_region 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 # 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. # 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. # Push it up while watching the GPU util + VRAM in the UI.
MAX_CONCURRENCY = 32 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: class _Slot:
"""One worker loop. `inflight` = jobs leased but not yet processed, so a """One worker loop. `inflight` = jobs leased but not yet processed, so a
@@ -43,7 +68,13 @@ class Worker:
self._slots: list[_Slot] = [] self._slots: list[_Slot] = []
self.processed = 0 self.processed = 0
self.errors = 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 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 ----------------------------------------------------------- # --- control -----------------------------------------------------------
def start(self): def start(self):
@@ -82,31 +113,48 @@ class Worker:
"active": self._active, "active": self._active,
"processed": self.processed, "processed": self.processed,
"errors": self.errors, "errors": self.errors,
"transient": self.transient,
} }
def _bump(self, *, processed=0, errors=0, active=0): def _bump(self, *, processed=0, errors=0, active=0, transient=0):
with self._lock: with self._lock:
self.processed += processed self.processed += processed
self.errors += errors self.errors += errors
self.transient += transient
self._active += active self._active += active
# --- per-slot loop ----------------------------------------------------- # --- per-slot loop -----------------------------------------------------
def _loop(self, slot: _Slot): def _loop(self, slot: _Slot):
backoff = self.cfg.poll_idle_seconds
while not slot.stop.is_set() and self._running: while not slot.stop.is_set() and self._running:
try: try:
jobs = self.client.lease(self.cfg.batch_size) jobs = self.client.lease(self.cfg.batch_size)
backoff = self.cfg.poll_idle_seconds # server answered → reset
except Exception: except Exception:
time.sleep(self.cfg.poll_idle_seconds) # 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 continue
if not jobs: if not jobs:
time.sleep(self.cfg.poll_idle_seconds) self._interruptible_sleep(slot, self.cfg.poll_idle_seconds)
continue continue
slot.inflight = [j["job_id"] for j in jobs] slot.inflight = [j["job_id"] for j in jobs]
for job in jobs: for job in jobs:
if slot.stop.is_set() or not self._running: if slot.stop.is_set() or not self._running:
break break
self._process(job) ok = self._process(job)
slot.inflight = [i for i in slot.inflight if i != job["job_id"]] 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: if slot.inflight:
self.client.heartbeat(slot.inflight) self.client.heartbeat(slot.inflight)
# Graceful hand-back of anything leased but not processed. # Graceful hand-back of anything leased but not processed.
@@ -114,7 +162,26 @@ class Worker:
self.client.release(slot.inflight) self.client.release(slot.inflight)
slot.inflight = [] slot.inflight = []
def _process(self, job: dict): 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) self._bump(active=1)
try: try:
data = self.client.fetch_image(job["image_url"]) data = self.client.fetch_image(job["image_url"])
@@ -126,8 +193,31 @@ class Worker:
else: else:
frames = [(None, media.load_image(data))] 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 = [] regions = []
ev = self.cfg.ccip_model or "ccip-default" ccip_ev = self.cfg.ccip_model or "ccip-default"
dv = f"person-{self.cfg.detector_level}" dv = f"person-{self.cfg.detector_level}"
for t, frame in frames: for t, frame in frames:
figs = models.detect_figures(frame, self.cfg.detector_level) figs = models.detect_figures(frame, self.cfg.detector_level)
@@ -137,20 +227,48 @@ class Worker:
crop = crop_region(frame, bbox) crop = crop_region(frame, bbox)
if crop is None: if crop is None:
continue continue
vec = models.ccip_vector(crop, self.cfg.ccip_model or None) if want_ccip:
regions.append({ regions.append({
"kind": "figure", "kind": "figure",
"bbox": list(bbox), "bbox": list(bbox),
"frame_time": t, "frame_time": t,
"score": score, "score": score,
"ccip_embedding": vec, "ccip_embedding": models.ccip_vector(
"embedding_version": ev, crop, self.cfg.ccip_model or None
"detector_version": dv, ),
}) "embedding_version": ccip_ev,
self.client.submit(job["job_id"], regions, ["figure", "face"]) "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) self._bump(processed=1)
except Exception as exc: # noqa: BLE001 — report + move on 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._bump(errors=1)
self.client.fail(job["job_id"], str(exc)[:500]) 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: finally:
self._bump(active=-1) self._bump(active=-1)
+4
View File
@@ -3,6 +3,10 @@ dghs-imgutils>=0.4
# GPU inference for the ONNX models. Swap to onnxruntime (CPU) for a slow # GPU inference for the ONNX models. Swap to onnxruntime (CPU) for a slow
# server-side fallback run. # server-side fallback run.
onnxruntime-gpu 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. # Control surface + HTTP.
fastapi fastapi
uvicorn[standard] uvicorn[standard]
+10
View File
@@ -37,6 +37,15 @@ async def overview():
.where(ImageRegion.ccip_embedding.is_not(None)) .where(ImageRegion.ccip_embedding.is_not(None))
) )
).scalar_one() ).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 # Per-character reference counts (no vectors loaded) — which characters
# have enough examples to match on. # have enough examples to match on.
ref_rows = ( ref_rows = (
@@ -72,6 +81,7 @@ async def overview():
return jsonify({ return jsonify({
"regions_by_kind": by_kind, "regions_by_kind": by_kind,
"images_with_figure_ccip": images_with_figure_ccip, "images_with_figure_ccip": images_with_figure_ccip,
"images_with_concept_siglip": images_with_concept_siglip,
"characters_with_references": len(ref_rows), "characters_with_references": len(ref_rows),
"character_references": [ "character_references": [
{"tag_id": t, "name": n, "n_refs": c} for (t, n, c) in ref_rows {"tag_id": t, "name": n, "n_refs": c} for (t, n, c) in ref_rows
+7
View File
@@ -17,6 +17,7 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert
from ..extensions import get_session from ..extensions import get_session
from ..models import AppSetting, GpuJob, ImageRecord, MLSettings from ..models import AppSetting, GpuJob, ImageRecord, MLSettings
from ..services.gallery_service import image_url 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.gpu_jobs import GpuJobService
from ..services.ml.regions import RegionService from ..services.ml.regions import RegionService
@@ -137,6 +138,12 @@ async def lease():
# For video/animated: the agent samples at this cadence. # For video/animated: the agent samples at this cadence.
"frame_interval_seconds": ml.video_frame_interval_seconds, "frame_interval_seconds": ml.video_frame_interval_seconds,
"max_frames": ml.video_max_frames, "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}) return jsonify({"jobs": out})
+5
View File
@@ -126,6 +126,11 @@ def make_celery() -> Celery:
"schedule": 3600.0, # auto-feed new images (+ retry errored) so "schedule": 3600.0, # auto-feed new images (+ retry errored) so
"args": ("ccip",), # the queue keeps moving without the button "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": { "ccip-auto-apply-daily": {
"task": "backend.app.tasks.ml.scheduled_ccip_auto_apply", "task": "backend.app.tasks.ml.scheduled_ccip_auto_apply",
"schedule": 86400.0, # no-op unless ccip_auto_apply_enabled "schedule": 86400.0, # no-op unless ccip_auto_apply_enabled
+29 -6
View File
@@ -29,6 +29,7 @@ from ...models import (
HeadAutoApplyRun, HeadAutoApplyRun,
HeadTrainingRun, HeadTrainingRun,
ImageRecord, ImageRecord,
ImageRegion,
MLSettings, MLSettings,
Tag, Tag,
TagHead, TagHead,
@@ -296,7 +297,14 @@ async def score_image(
category, score}], ranked. A concept surfaces when its score clears the category, score}], ranked. A concept surfaces when its score clears the
head's own suggest_threshold — or, when threshold_override is given (the head's own suggest_threshold — or, when threshold_override is given (the
typed-dropdown "show everything" mode), that flat floor instead (0 → every 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 import numpy as np
img = await session.get(ImageRecord, image_id) img = await session.get(ImageRecord, image_id)
@@ -306,11 +314,26 @@ async def score_image(
heads = await _current_heads(session, settings.embedder_model_version) heads = await _current_heads(session, settings.embedder_model_version)
if heads["W"] is None: if heads["W"] is None:
return [] return []
x = np.asarray(img.siglip_embedding, dtype=np.float32)
n = float(np.linalg.norm(x)) or 1.0 bag = [np.asarray(img.siglip_embedding, dtype=np.float32)]
xn = x / n region_vecs = (
z = heads["W"] @ xn + heads["b"] await session.execute(
probs = 1.0 / (1.0 + np.exp(-z)) 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 = [] out = []
for i, p in enumerate(probs): for i, p in enumerate(probs):
cut = threshold_override if threshold_override is not None else heads["thr"][i] cut = threshold_override if threshold_override is not None else heads["thr"][i]
+31 -12
View File
@@ -742,24 +742,43 @@ def scheduled_apply_head_tags() -> str:
@celery.task(name="backend.app.tasks.ml.enqueue_gpu_backfill") @celery.task(name="backend.app.tasks.ml.enqueue_gpu_backfill")
def enqueue_gpu_backfill(task_name: str) -> int: def enqueue_gpu_backfill(task_name: str) -> int:
"""Enqueue a gpu_job for every image that doesn't already have one for """Enqueue a gpu_job for every image that still needs `task_name` (one
`task_name` (one INSERT…SELECT, so it scales to a full library). The desktop INSERT…SELECT, so it scales to a full library). The desktop agent drains the
agent drains the queue over HTTP. Returns the number enqueued.""" 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 exists, insert, literal
from sqlalchemy import select as sa_select from sqlalchemy import select as sa_select
from ..models import GpuJob, ImageRecord from ..models import GpuJob, ImageRecord, ImageRegion
SessionLocal = _sync_session_factory() SessionLocal = _sync_session_factory()
with SessionLocal() as session: with SessionLocal() as session:
already = exists().where( if task_name == "siglip":
GpuJob.image_record_id == ImageRecord.id, has_concept = exists().where(
GpuJob.task == task_name, ImageRegion.image_record_id == ImageRecord.id,
GpuJob.status.in_(["pending", "leased", "done"]), ImageRegion.kind == "concept",
) )
sel = sa_select( queued = exists().where(
ImageRecord.id, literal(task_name), literal("pending") GpuJob.image_record_id == ImageRecord.id,
).where(~already) 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. # RETURNING + count: result.rowcount is unreliable for INSERT…SELECT.
rows = session.execute( rows = session.execute(
insert(GpuJob) insert(GpuJob)
@@ -61,6 +61,16 @@
processes until the agent is running. processes until the agent is running.
</p> </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 --> <!-- Match strictness -->
<div class="fc-section-h mt-5 mb-1">Character-match strictness</div> <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"> <div v-if="ml.settings" class="d-flex align-center" style="gap:12px">
@@ -115,6 +125,7 @@ const tokenValue = ref(null)
const masked = ref(true) const masked = ref(true)
const rotating = ref(false) const rotating = ref(false)
const backfilling = ref(false) const backfilling = ref(false)
const backfillingSiglip = ref(false)
const threshold = ref(0.85) const threshold = ref(0.85)
const savingThreshold = ref(false) const savingThreshold = ref(false)
const autoApply = ref(true) const autoApply = ref(true)
@@ -215,6 +226,19 @@ async function onBackfill() {
backfilling.value = false 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> </script>
<style scoped> <style scoped>
+39 -2
View File
@@ -2,9 +2,9 @@
from datetime import UTC, datetime, timedelta from datetime import UTC, datetime, timedelta
import pytest import pytest
from sqlalchemy import select from sqlalchemy import func, select
from backend.app.models import GpuJob, ImageRecord from backend.app.models import GpuJob, ImageRecord, ImageRegion
from backend.app.services.ml.gpu_jobs import GpuJobService from backend.app.services.ml.gpu_jobs import GpuJobService
pytestmark = pytest.mark.integration pytestmark = pytest.mark.integration
@@ -20,6 +20,43 @@ async def _img(db, sha) -> ImageRecord:
return img 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 @pytest.mark.asyncio
async def test_enqueue_dedupes_same_pair(db): async def test_enqueue_dedupes_same_pair(db):
img = await _img(db, "a" * 64) img = await _img(db, "a" * 64)
+34
View File
@@ -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"]) 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 @pytest.mark.asyncio
async def test_rejected_tag_surfaced_flagged_then_reversible(db): async def test_rejected_tag_surfaced_flagged_then_reversible(db):
# A dismissed suggestion is NOT dropped: it stays flagged rejected so the # A dismissed suggestion is NOT dropped: it stays flagged rejected so the