import torch
import open_clip
from PIL import Image
import os

_model = None
_preprocess = None
_device = None

def get_model():
    global _model, _preprocess, _device
    if _model is None:
        _device = "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(_device)
        _model.eval()
    return _model, _preprocess, _device

def generate_embedding(image_path):
    model, preprocess, device = get_model()
    try:
        image = preprocess(Image.open(image_path)).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