import io
import torch
import requests
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration

class EndpointHandler:
    def __init__(self, model_dir):
        self.processor = BlipProcessor.from_pretrained(model_dir)
        self.model = BlipForConditionalGeneration.from_pretrained(model_dir)
        self.model.eval()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def __call__(self, data):
        # Unwrap Hugging Face "inputs" key if present
        if isinstance(data, dict) and "inputs" in data:
            data = data["inputs"]

        image_url = data.get("image_url")
        if not image_url:
            return {"error": "Please provide an 'image_url' in the JSON payload."}

        try:
            response = requests.get(image_url)
            response.raise_for_status()
            image = Image.open(io.BytesIO(response.content)).convert("RGB")
        except Exception as e:
            return {"error": f"Failed to load image: {str(e)}"}

        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
        with torch.no_grad():
            out = self.model.generate(**inputs)