bvandeusen
2713c3f773
Two issues surfaced by the live logs (GPU pegged at ~0% util, 0.5 jobs/s,
truncated-image failures):
- BATCH the SigLIP embeds: collect all of an image's crops (figure + booru_yolo
components + panels) and embed them in ONE forward pass instead of one
forward+lock per crop. The per-crop path serialised every crop through the
inference lock and starved the GPU (≈0% util, autoscaler stuck oscillating);
batching gives a real GPU-bound workload + far higher throughput. CCIP still
runs per figure inline.
- LOAD_TRUNCATED_IMAGES in the agent (matches the server embedder): slightly-
truncated scraped images now load instead of failing the job 3× then erroring
("image file is truncated (N bytes not processed)").
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
78 lines
3.2 KiB
Python
78 lines
3.2 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."""
|
|
return self.embed_batch([image])[0]
|
|
|
|
def embed_batch(self, images: list) -> list[list[float]]:
|
|
"""Embed many crops in ONE forward pass — far better GPU utilisation +
|
|
only one lock acquisition than embedding each crop separately (which
|
|
starved the GPU and serialised the whole pool)."""
|
|
if not images:
|
|
return []
|
|
self.load()
|
|
torch = self._torch
|
|
enc = self._processor(images=images, 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
|
|
arr = pooled.float().cpu().numpy().astype(np.float32)
|
|
return [row.reshape(-1).tolist() for row in arr]
|