File size: 3,606 Bytes
3a507ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | # handler.py
from typing import Dict, Any, List
import torch
import PIL.Image
from io import BytesIO
import base64
from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
# Configure logging for debugging purposes
logging.basicConfig(level=logging.INFO)
class EndpointHandler:
def __init__(self, path=""):
logging.info("Initializing EndpointHandler for Moondream2")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {self.device}")
# Load the model with trust_remote_code enabled.
# 'path' points to the location of the model files inside the container.
self.model = AutoModelForCausalLM.from_pretrained(
path,
trust_remote_code=True,
torch_dtype=torch.float16,
device_map=self.device
)
self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
# Ensure the model is moved to the device
self.model.to(self.device)
self.model.eval()
logging.info("Moondream2 model loaded successfully.")
def preprocess_image(self, encoded_image: str) -> PIL.Image.Image:
"""Decode and preprocess the base64 encoded image."""
try:
image_data = base64.b64decode(encoded_image)
return PIL.Image.open(BytesIO(image_data)).convert("RGB")
except Exception as e:
logging.error(f"Error decoding image: {e}")
raise ValueError(f"Failed to decode image data: {e}")
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Handles the API call. The `data` argument is a dictionary containing the payload.
Expects a JSON payload like:
{
"inputs": {
"prompt": "What's in this picture?",
"image": "base64_encoded_image_string"
}
}
"""
logging.info("Received request payload")
inputs = data.get("inputs", {})
prompt = inputs.get("prompt", "")
encoded_image = inputs.get("image", "")
if not prompt or not encoded_image:
raise ValueError("Prompt and base64 encoded image must be provided in the 'inputs' field.")
image = self.preprocess_image(encoded_image)
# Process the image and prompt
enc_image = self.model.encode_image(image)
# Create the conversation history for inference
chat_history = f"Question: {prompt}\n\nAnswer:"
logging.info(f"Running inference with prompt: {prompt}")
with torch.no_grad():
output_tokens = self.model.generate(
enc_image,
self.tokenizer,
chat_history,
pad_token_id=self.tokenizer.eos_token_id,
# Add other generation parameters here if needed
)
# Decode the generated tokens
generated_text = self.tokenizer.batch_decode(output_tokens, skip_special_tokens=True)[0]
logging.info(f"Inference complete. Generated text: {generated_text}")
# Post-process the output to isolate the answer
try:
# The model output includes the prompt, so we need to extract only the answer part.
answer_start_tag = "\n\nAnswer:"
generated_answer = generated_text.split(answer_start_tag)[-1].strip()
except IndexError:
generated_answer = generated_text # Fallback if splitting fails
return [{"generated_text": generated_answer}]
|