from transformers import AutoModel
from PIL import Image
import requests
from io import BytesIO
import os
import torch
import numpy as np
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
os.environ['HF_HOME'] = '/app/hf_cache'
# Load model
model = AutoModel.from_pretrained('jinaai/jina-clip-v2', trust_remote_code=True).to(device)

def get_text_embedding(texts, truncate_dim=512):
    embeddings =  model.encode_text(texts, truncate_dim=truncate_dim)
    # if isinstance(embeddings, np.ndarray):
    embeddings = torch.from_numpy(embeddings)
    print(embeddings)
    return embeddings


def get_image_embedding(image_urls, truncate_dim=512):
    """
    Takes a list of image URLs and returns embeddings using model.encode_image.
    Assumes model.encode_image supports URL input directly.
    """
    embeddings = model.encode_image(image_urls, truncate_dim=truncate_dim)
    # if not isinstance(embeddings, torch.Tensor):
    embeddings = torch.tensor(embeddings)
    print(embeddings)
    return embeddings