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 "OK"#prediction.item()