zesquirrelnator/idefics2-8b-docvqa-finetuned-tutorial

1from transformers import AutoModelForCausalLM, AutoTokenizer2from PIL import Image3import torch4from io import BytesIO5import base646 7# Initialize the model and tokenizer8model_id = "HuggingFaceM4/idefics2-8b"9model = AutoModelForCausalLM.from_pretrained(model_id)10tokenizer = AutoTokenizer.from_pretrained(model_id)11 12# Check if CUDA (GPU support) is available and then set the device to GPU or CPU13device = torch.device("cuda" if torch.cuda.is_available() else "cpu")14model.to(device)15 16def preprocess_image(encoded_image):17    """Decode and preprocess the input image."""18    decoded_image = base64.b64decode(encoded_image)19    img = Image.open(BytesIO(decoded_image)).convert("RGB")20    return img21 22def handler(event, context):23    """Handle the incoming request."""24    try:25        # Extract the base64-encoded image and question from the event26        input_image = event['body']['image']27        question = event['body'].get('question', "What is this image about?")28 29        # Preprocess the image30        img = preprocess_image(input_image)31 32        # Perform inference33        enc_image = model.encode_image(img).to(device)34        answer = model.answer_question(enc_image, question, tokenizer)35 36        # If the output is a tensor, move it back to CPU and convert to list37        if isinstance(answer, torch.Tensor):38            answer = answer.cpu().numpy().tolist()39 40        # Create the response41        response = {42            "statusCode": 200,43            "body": {44                "answer": answer45            }46        }47        return response48    except Exception as e:49        # Handle any errors50        response = {51            "statusCode": 500,52            "body": {53                "error": str(e)54            }55        }56        return response