| import requests | |
| from typing import Dict, Any | |
| from PIL import Image | |
| import torch | |
| import base64 | |
| from io import BytesIO | |
| from transformers import BlipForConditionalGeneration, BlipProcessor | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| self.processor = BlipProcessor.from_pretrained( | |
| "Salesforce/blip-image-captioning-large" | |
| ) | |
| self.model = BlipForConditionalGeneration.from_pretrained( | |
| "Salesforce/blip-image-captioning-large" | |
| ).to(device) | |
| self.model.eval() | |
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: | |
| input_data = data.get("inputs", {}) | |
| encoded_images = input_data.get("images") | |
| if not encoded_images: | |
| return {"captions": [], "error": "No images provided"} | |
| texts = input_data.get("texts", [""] * len(encoded_images)) | |
| try: | |
| raw_images = [ | |
| Image.open(BytesIO(base64.b64decode(img))).convert("RGB") | |
| for img in encoded_images | |
| ] | |
| processed_inputs = [ | |
| self.processor(image, text, return_tensors="pt") | |
| for image, text in zip(raw_images, texts) | |
| ] | |
| processed_inputs = { | |
| "pixel_values": torch.cat( | |
| [inp["pixel_values"] for inp in processed_inputs], dim=0 | |
| ).to(device), | |
| "input_ids": torch.cat( | |
| [inp["input_ids"] for inp in processed_inputs], dim=0 | |
| ).to(device), | |
| "attention_mask": torch.cat( | |
| [inp["attention_mask"] for inp in processed_inputs], dim=0 | |
| ).to(device), | |
| } | |
| with torch.no_grad(): | |
| out = self.model.generate(**processed_inputs) | |
| captions = self.processor.batch_decode(out, skip_special_tokens=True) | |
| return {"captions": captions} | |
| except Exception as e: | |
| print(f"Error during processing: {str(e)}") | |
| return {"captions": [], "error": str(e)} | |