handler.py

from typing import Any, Dict, List
from transformers import Idefics2Processor, Idefics2ForConditionalGeneration
import torch
import logging


class EndpointHandler:
    def __init__(self, path=""):
        # Preload all the elements you are going to need at inference.
        self.logger = logging.getLogger()
        self.logger.addHandler(logging.StreamHandler())
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.processor = Idefics2Processor.from_pretrained(path)
        self.model = Idefics2ForConditionalGeneration.from_pretrained(path)
        self.model.to(self.device)
        self.logger.info("Initialisation finished!")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
         data args:
              inputs (:obj: `str` | `PIL.Image` | `np.array`)
              kwargs
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """
        checkpoints = ""

        try:
            image = data.pop("inputs", data)
            checkpoints += "image reached\n"

            # process image
            inputs = self.processor(images=image, return_tensors="pt").to(self.device)
            checkpoints += "inputs reached\n"
            generated_ids = self.model.generate(**inputs, max_new_tokens=20)
            checkpoints += "generated\n"

            # run prediction
            generated_text: List[str] = self.processor.batch_decode(
                generated_ids, skip_special_tokens=True
            )
            checkpoints += "decoded\n"
        except Exception as e:
            checkpoints += f"{e}\n"

        # decode output
        return generated_text.append(checkpoints)