| from transformers import AutoModelForCausalLM, AutoProcessor |
| from PIL import Image |
| import requests |
| import torch |
| import io |
|
|
| class EndpointHandler: |
| def __init__(self, model_dir): |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True).to(device) |
| self.processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True) |
| self.device = device |
|
|
| def __call__(self, data): |
| try: |
| url = data.get("inputs", {}).get("url") |
| if not url: |
| return {"error": "Missing URL"} |
| |
| headers = { |
| "User-Agent": "Mozilla/5.0", |
| "Accept": "image/*" |
| } |
| response = requests.get(url, headers=headers, verify=False) |
| response.raise_for_status() |
| |
| image_data = io.BytesIO(response.content) |
| image = Image.open(image_data).convert("RGB") |
| |
| inputs = self.processor( |
| text="<MORE_DETAILED_CAPTION>", |
| images=image, |
| return_tensors="pt" |
| ).to(self.device) |
|
|
| with torch.inference_mode(): |
| output = self.model.generate( |
| **inputs, |
| max_new_tokens=512, |
| num_beams=3 |
| ) |
|
|
| text = self.processor.batch_decode(output, skip_special_tokens=True)[0] |
| return {"caption": text} |
| |
| except Exception as e: |
| return {"error": str(e)} |
