| from transformers import BlipProcessor, BlipForConditionalGeneration |
| from PIL import Image |
| import requests |
| import torch |
| import base64 |
| import io |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path): |
| self.processor = BlipProcessor.from_pretrained(path) |
| self.model = BlipForConditionalGeneration.from_pretrained( |
| path, |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32 |
| ) |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.model.to(self.device) |
|
|
| def _load_image(self, image_input): |
| |
| if isinstance(image_input, str) and image_input.startswith("http"): |
| return Image.open(requests.get(image_input, stream=True).raw).convert("RGB") |
|
|
| |
| if isinstance(image_input, str): |
| image_bytes = base64.b64decode(image_input) |
| return Image.open(io.BytesIO(image_bytes)).convert("RGB") |
|
|
| raise ValueError("Unsupported image input format") |
|
|
| def __call__(self, data): |
| image_input = data.get("inputs") |
| if image_input is None: |
| raise ValueError("No image provided") |
|
|
| image = self._load_image(image_input) |
|
|
| inputs = self.processor(images=image, return_tensors="pt").to(self.device) |
|
|
| output = self.model.generate(**inputs, max_new_tokens=50) |
|
|
| caption = self.processor.decode(output[0], skip_special_tokens=True) |
| return {"caption": caption} |
|
|
