feat(ml): GPU-capable tagger + embedder with CPU fallback (#872)
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user