bvandeusen
57e52433d0
The SigLIP embedder + YOLO proposers load lazily then stay resident for the container's whole lifetime — a 24/7 agent with an empty queue squats on ~5GB of VRAM doing nothing (operator-observed: 4900MiB held at GPU-util 8% / P8). Sleep mode only sheds downloaders + poll cadence; even a UI Stop left the models loaded. Add a monitor thread that unloads the torch-owned models after cfg.idle_unload_seconds (env IDLE_UNLOAD_SECONDS, default 300; 0 disables) with the GPU genuinely idle (active==0, buffer drained, no job completed in the window), then torch.cuda.empty_cache() to hand the blocks back to the driver. They reload lazily on the next job via the existing _ensure_embedder / _proposers_for. Covers both sleep-mode idle and a full Stop. Surfaced in /status (models_loaded) and the agent UI pipe line; the VRAM meter drops too. Residual: imgutils CCIP/person ONNX sessions + the CUDA context stay resident (no clean unload API) — idle VRAM drops substantially, not to zero. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TbrA36zNczjVhrM6cWThQa
93 lines
4.0 KiB
Python
93 lines
4.0 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]
|
|
|
|
def unload(self) -> bool:
|
|
"""Drop the loaded model so its VRAM can be reclaimed — the idle monitor
|
|
calls this after a spell with no work so an idle agent doesn't squat on
|
|
the card; the next embed() reloads it lazily (a few seconds). Held under
|
|
BOTH the load and inference locks so it can never race a concurrent load
|
|
or an in-flight forward pass. Returns True if a model was actually
|
|
released (the caller then runs one empty_cache() to hand the freed blocks
|
|
back to the driver)."""
|
|
with self._load_lock, self._infer_lock:
|
|
if self._model is None:
|
|
return False
|
|
self._model = None
|
|
self._processor = None
|
|
return True
|