| import torch.nn.functional as F | |
| from torch import Tensor | |
| from transformers import AutoTokenizer, AutoModel, AutoProcessor | |
| from torch import cuda | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| self.processor = AutoProcessor.from_pretrained(path) | |
| self.model = AutoModel.from_pretrained(path, trust_remote_code=True) | |
| self.device = "cuda" if cuda.is_available() else "cpu" | |
| self.model.to(self.device) | |
| def __call__(self, data: Dict[str, Any]) -> List[List[int]]: | |
| image = data.pop("inputs",data) | |
| processed = self.processor(images=image, return_tensors="pt").to(self.device) | |
| prediction = self.model(processed["pixel_values"]) | |
| return prediction.item() |