from transformers import pipeline, AutoFeatureExtractor
from PIL import Image
import torch


class EndpointHandler:
    def __init__(self, path=""):
        self.pipe = pipeline(
            "image-to-text",
            model=path,
            feature_extractor=AutoFeatureExtractor,
            device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
        )

    def __call__(self, data) -> str:
        return self.pipe(data.pop("inputs"))