45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
import torch
|
|
import os
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
import open_clip
|
|
from PIL import Image
|
|
import threading
|
|
|
|
_model = None
|
|
_preprocess = None
|
|
_device = None
|
|
_lock = threading.Lock()
|
|
_gpu_lock = threading.Lock()
|
|
|
|
def get_model():
|
|
global _model, _preprocess, _device
|
|
if _model is None:
|
|
with _lock:
|
|
if _model is None:
|
|
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
|
# ViT-H-14 is 1024-dim
|
|
model, _, preprocess = open_clip.create_model_and_transforms('ViT-H-14', pretrained='laion2b_s32b_b79k')
|
|
model = model.to(dev)
|
|
model.eval()
|
|
|
|
# Set globals only when fully ready to avoid race conditions
|
|
_preprocess = preprocess
|
|
_device = dev
|
|
_model = model
|
|
return _model, _preprocess, _device
|
|
|
|
def generate_embedding(image_path):
|
|
model, preprocess, device = get_model()
|
|
try:
|
|
with _gpu_lock:
|
|
with Image.open(image_path) as img:
|
|
image = preprocess(img.convert("RGB")).unsqueeze(0).to(device)
|
|
with torch.no_grad():
|
|
image_features = model.encode_image(image)
|
|
image_features /= image_features.norm(dim=-1, keepdim=True)
|
|
return image_features.cpu().numpy()[0].tolist()
|
|
except Exception as e:
|
|
print(f"Error generating embedding for {image_path}: {e}")
|
|
return None
|