| import requests |
| import json |
| import torch |
| from PIL import Image |
| from transformers import AutoProcessor, AutoModelForCausalLM |
|
|
| class EndpointHandler(): |
| def __init__(self, path=""): |
| self.device = "cuda:0" if torch.cuda.is_available() else "cpu" |
| self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 |
|
|
| self.model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", torch_dtype=self.torch_dtype, trust_remote_code=True).to(self.device) |
| self.processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True) |
|
|
|
|
| def predict_image(self, url, task, prompt): |
| image = Image.open(requests.get(url, stream=True).raw) |
|
|
| inputs = self.processor(text=task + prompt, images=image, return_tensors="pt").to(self.device, self.torch_dtype) |
|
|
| generated_ids = self.model.generate( |
| input_ids=inputs["input_ids"], |
| pixel_values=inputs["pixel_values"], |
| max_new_tokens=4096, |
| num_beams=3, |
| do_sample=False |
| ) |
| generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0] |
|
|
| parsed_answer = self.processor.post_process_generation(generated_text, task=task, image_size=(image.width, image.height)) |
| return parsed_answer |
| |
|
|
| def __call__(self, event): |
| if "inputs" not in event: |
| return { |
| "statusCode": 400, |
| "body": json.dumps("Error: Please provide an 'inputs' parameter."), |
| } |
|
|
| inputs = event["inputs"] |
| url = inputs["url"] |
| task = inputs["task"] |
| prompt = inputs["prompt"] |
| parsed_answer = self.predict_image(url, task, prompt) |
|
|
| return { |
| "statusCode": 200, |
| "body": json.dumps(parsed_answer), |
| } |
