feat(ml): GPU-capable tagger + embedder with CPU fallback (#872)
CI / lint (push) Successful in 3s
CI / frontend-build (push) Successful in 20s
CI / backend-lint-and-test (push) Successful in 29s
CI / integration (push) Successful in 3m13s

Step 1 of GPU enablement (code only — CPU-safe, CI-green; the CUDA image is a
separate step pending the host driver version).

- New services/ml/device.py: FC_ML_DEVICE (auto|cuda|cpu) intent + VRAM knobs
  (FC_ML_ONNX_GPU_MEM_GB, FC_ML_TORCH_MEM_FRACTION). Per-worker-host bootstrap →
  env, not a DB setting (the GPU host runs CUDA, others CPU).
- tagger: use CUDAExecutionProvider (with gpu_mem_limit) when requested AND the
  provider is actually present (onnxruntime-gpu), else CPUExecutionProvider. Logs
  the active providers.
- embedder: move model + inputs to cuda when requested AND torch.cuda is
  available; cap torch's VRAM share; .detach().cpu() before numpy. fp32 kept so
  GPU embeddings stay in the same space as existing CPU ones.

Both AND the env intent with the framework's real availability, so on CPU
(CI / CPU onnxruntime / no GPU) they fall back cleanly — behavior unchanged.
The 8GB P4 is shared by both frameworks, hence the conservative default caps.

Tests: device env parsing. (tagger/embedder GPU paths are operator-verified on
the GPU host — models aren't in CI.)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-16 12:49:24 -04:00
parent 369e3de684
commit db7e1f2b59
4 changed files with 108 additions and 11 deletions
+24 -5
View File
@@ -1,8 +1,11 @@
"""SigLIP SO400M image-embedding wrapper (PyTorch CPU).
"""SigLIP SO400M image-embedding wrapper (PyTorch).
Direct port of ImageRepo's siglip.py. torch/transformers are imported
lazily inside load() so this module can be imported in the web container
(which never runs inference) without paying the torch import cost.
Runs on CPU by default; moves to CUDA when requested (FC_ML_DEVICE) and a GPU is
available (#872), else stays on CPU. fp32 is kept on GPU too so GPU-computed
embeddings stay in the same numeric space as the existing CPU ones (cosine
comparisons). torch/transformers are imported lazily inside load() so this
module can be imported in the web container (which never runs inference) without
paying the torch import cost.
"""
import os
@@ -29,6 +32,7 @@ class Embedder:
self._model = None
self._processor = None
self._torch = None
self._device = "cpu"
def load(self) -> None:
if self._model is not None:
@@ -36,7 +40,17 @@ class Embedder:
import torch
from transformers import AutoModel, SiglipImageProcessor
from .device import gpu_requested, torch_mem_fraction
self._torch = torch
# GPU (#872) when requested AND a CUDA device is present; else CPU. Cap
# torch's share of the 8GB P4 (the ONNX tagger shares the card).
if gpu_requested() and torch.cuda.is_available():
self._device = "cuda"
try:
torch.cuda.set_per_process_memory_fraction(torch_mem_fraction())
except Exception: # noqa: BLE001 — best-effort cap; never block load
pass
# FC's embedder only does IMAGE inference — never text. AutoProcessor
# loads the full processor including SiglipTokenizer, which requires
# the sentencepiece library at import time even if we never call it.
@@ -51,6 +65,8 @@ class Embedder:
)
self._model = AutoModel.from_pretrained(str(self._model_dir))
self._model.eval()
if self._device == "cuda":
self._model = self._model.to("cuda")
def infer(self, image_path: Path) -> np.ndarray:
"""Return a 1152-dim float32 embedding (SigLIP MAP-pooled output)."""
@@ -58,9 +74,12 @@ class Embedder:
img = Image.open(image_path).convert("RGB")
with self._torch.no_grad():
inputs = self._processor(images=img, return_tensors="pt")
if self._device == "cuda":
inputs = {k: v.to("cuda") for k, v in inputs.items()}
out = self._model.get_image_features(**inputs)
pooled = out.pooler_output if hasattr(out, "pooler_output") else out
return pooled[0].numpy().astype(np.float32)
# .detach().cpu() so a CUDA tensor converts to numpy (no-op on CPU).
return pooled[0].detach().cpu().numpy().astype(np.float32)
_default_embedder: Embedder | None = None