import torch
from PIL import Image
from transformers import Blip2Processor, Blip2ForConditionalGeneration, BitsAndBytesConfig

device = "cuda" if torch.cuda.is_available() else "cpu"
quantization_config = BitsAndBytesConfig(load_in_8bit=True)

processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
model = Blip2ForConditionalGeneration.from_pretrained(
    "Salesforce/blip2-flan-t5-xl", device_map="auto"
)

def get_image_answer(image: Image.Image, question: str) -> str:
    if image.mode != "RGB":
        image = image.convert("RGB")

    inputs = processor(images=image, text=question, return_tensors="pt")
    inputs = {k: v.to(device, torch.float16 if device == "cuda" else torch.float32) for k, v in inputs.items()}
    output_ids = model.generate(**inputs)
    answer = processor.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
    return answer