from io import BytesIO
import base64

from PIL import Image
import torch
from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class EndpointHandler():
    def __init__(self, path=""):
        self.text_model = CLIPTextModel.from_pretrained("rbanfield/clip-vit-large-patch14")
        self.image_model = CLIPVisionModelWithProjection.from_pretrained("rbanfield/clip-vit-large-patch14")
        self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")

    def __call__(self, data):
        text_input = data.pop("text", None)
        image_input = data.pop("image", None)

        if text_input:
            processor = self.processor(text=text_input, return_tensors="pt", padding=True)
            with torch.no_grad():
                return self.text_model(**processor).pooler_output.tolist()
        elif image_input:
            image = Image.open(BytesIO(base64.b64decode(image_input)))
            processor = self.processor(images=image, return_tensors="pt")
            with torch.no_grad():
                return self.image_model(**processor).image_embeds.tolist()
        else:
            return None