| import os |
| import json |
| import base64 |
| from io import BytesIO |
| import logging |
|
|
| from transformers import CLIPProcessor, CLIPModel |
| from PIL import Image |
| import requests |
|
|
| logging.basicConfig(level=logging.INFO) |
|
|
| model = None |
| processor = None |
|
|
| def init(): |
| global model, processor |
| model_name = os.getenv("MODEL_NAME", "openai/clip-vit-base-patch32") |
| logging.info(f"Loading model: {model_name}") |
| model = CLIPModel.from_pretrained(model_name) |
| processor = CLIPProcessor.from_pretrained(model_name) |
| logging.info("Model and processor loaded successfully.") |
|
|
| def handle_request(request_data, context): |
| results = [] |
| for data in request_data: |
| try: |
| payload = json.loads(data) |
| image_input = payload.get("image") |
| text_input = payload.get("text", []) |
| if image_input.startswith("http://") or image_input.startswith("https://"): |
| response = requests.get(image_input, stream=True, timeout=10) |
| image = Image.open(response.raw).convert("RGB") |
| elif image_input.startswith("data:"): |
| header, encoded = image_input.split(",", 1) |
| image = Image.open(BytesIO(base64.b64decode(encoded))).convert("RGB") |
| else: |
| image = Image.open(BytesIO(base64.b64decode(image_input))).convert("RGB") |
| inputs = processor(text=text_input, images=image, return_tensors="pt", padding=True) |
| outputs = model(**inputs) |
| logits_per_image = outputs.logits_per_image |
| probs = logits_per_image.softmax(dim=1) |
| results.append(probs.tolist()) |
| except Exception as e: |
| results.append({"error": str(e)}) |
| return results |
|
|
| class EndpointHandler: |
| def __init__(self, model_dir=None): |
| init() |
|
|
| def handle(self, request_data, context): |
| return handle_request(request_data, context) |
