bvandeusen
c6f38b0dac
Lift recall on small/local concepts (glasses, cum, stomach-bulge, xray, lactation) that the whole-image SigLIP vector washes out: the GPU agent now embeds figure crops with SigLIP too, stored as kind='concept' regions, and the suggestion rail scores each image as a BAG (whole-image + every concept crop), taking each head's MAX over the bag. The whole-image vector is always in the bag, so this can never score lower than before. Model-agnostic by construction: the server ANNOUNCES the embedding model (HF name + version) in the lease, so the agent loads whatever the heads were trained in and stays in lock-step — a model swap is a server setting + a re-embed migration, never an agent change. - agent: model-agnostic CropEmbedder (torch/transformers get_image_features, fp16 on CUDA, inference-locked); worker branches on job.task — 'ccip' emits figure(CCIP)+concept(SigLIP) in one pass, 'siglip' emits concept-only so the back-catalogue backfill never churns figure/CCIP regions; torch cu124 + transformers in the image. - server: lease announces embed_model_name/embed_version; score_image is max-over-bag (version-filtered region embeddings); enqueue_gpu_backfill 'siglip' gates on a missing concept region (drains the back-catalogue, retries failures, no double-enqueue); daily siglip-backfill beat; UI button; /api/ccip/overview reports images_with_concept_siglip. - v1 scope: suggestion rail only — auto-apply stays whole-image (conservative; heads' thresholds were calibrated on whole-image). Bulk-apply bag = follow-up. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
70 lines
2.9 KiB
Python
70 lines
2.9 KiB
Python
"""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()
|