| |
|
| | from io import BytesIO |
| | from typing import Any, List, Dict |
| |
|
| | from PIL import Image |
| | from transformers import AutoProcessor, AutoModelForCausalLM |
| | from PIL import Image |
| | import requests |
| | import copy |
| | import base64 |
| |
|
| |
|
| |
|
| | class EndpointHandler(): |
| | def __init__(self, path=""): |
| | |
| | model_id = 'microsoft/Florence-2-large' |
| | model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).eval().cuda() |
| | processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) |
| | self.model = model |
| | self.processor = processor |
| |
|
| | def run_example(self, image, task_prompt, text_input=None): |
| | if text_input is None: |
| | prompt = task_prompt |
| | else: |
| | prompt = task_prompt + text_input |
| | inputs = self.processor(text=prompt, images=image, return_tensors="pt") |
| | generated_ids = self.model.generate( |
| | input_ids=inputs["input_ids"].cuda(), |
| | pixel_values=inputs["pixel_values"].cuda(), |
| | max_new_tokens=1024, |
| | early_stopping=False, |
| | do_sample=False, |
| | num_beams=3, |
| | ) |
| | generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0] |
| | parsed_answer = self.processor.post_process_generation( |
| | generated_text, |
| | task=task_prompt, |
| | image_size=(image.width, image.height) |
| | ) |
| |
|
| | return parsed_answer |
| |
|
| | def __call__(self, data: Any) -> List[List[Dict[str, float]]]: |
| | image = data['inputs'].pop("image", None) |
| | image_url = data['inputs'].pop("image_url", None) |
| | type = data['inputs'].pop("type", '<MORE_DETAILED_CAPTION>') |
| | text = data['inputs'].pop("text", None) |
| | if image: |
| | image = Image.open(BytesIO(base64.b64decode(image))) |
| | elif image_url: |
| | response = requests.get(image_url) |
| | if response.status_code == 200: |
| | image = Image.open(BytesIO(response.content)) |
| | else: |
| | raise ValueError(f"Unable to download image from URL: {image_url}") |
| | return self.run_example(image, type, text_input=text) |
| |
|
