Spaces:
Running
Running
File size: 1,793 Bytes
a1dd5ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """One device policy for every model loader.
The store format is portable; the model loaders must be too. Every
loader asks this module instead of probing torch itself, so a cloud
CPU box, an Apple laptop, and a CUDA node run the same code with the
right placement:
ELIDEDB_DEVICE force a device ("cpu", "mps", "cuda"); default is
mps when available, else cuda, else cpu
ELIDEDB_DTYPE force a torch dtype name; default float16 on
mps/cuda and bfloat16 on cpu (halves resident
memory against fp32, and fp16 matmuls are not a
real CPU option)
"""
from __future__ import annotations
import os
def text_only():
"""Serving deployments set ELIDEDB_TEXT_ONLY=1: the query path
encodes TEXT only (every frame vector is precomputed at ingest),
so the loaders drop their vision towers after load and roughly
halve resident memory. Ingest machines leave it unset."""
return os.environ.get("ELIDEDB_TEXT_ONLY", "") == "1"
def strip_vision(model, *attrs):
"""Release the named submodules when serving text-only."""
if not text_only():
return model
import gc
for a in attrs:
if hasattr(model, a):
setattr(model, a, None)
gc.collect()
return model
def pick():
import torch
dev = os.environ.get("ELIDEDB_DEVICE", "").strip()
if not dev:
if torch.backends.mps.is_available():
dev = "mps"
elif torch.cuda.is_available():
dev = "cuda"
else:
dev = "cpu"
name = os.environ.get("ELIDEDB_DTYPE", "").strip()
if name:
dtype = getattr(torch, name)
else:
dtype = torch.float16 if dev in ("mps", "cuda") \
else torch.bfloat16
return dev, dtype
|