from fastapi import FastAPI, File, UploadFile, Form
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image
from io import BytesIO
import os
import base64

app = FastAPI()

class ImageQuestion(BaseModel):
    image: str
    question: str

model_id = os.environ.get("MODEL_ID")
revision = os.environ.get("REVISION")

if not model_id or not revision:
    raise ValueError("Please set MODEL_ID and REVISION environment variables.")

model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, revision=revision
)
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)

def generate_answer(image_bytes, question):
    img = base64.b64decode(image_bytes)
    image = Image.open(BytesIO(img))
    enc_image = model.encode_image(image)
    return model.answer_question(enc_image, question, tokenizer)

@app.post("/get_answer/")
async def get_answer(data: ImageQuestion):
    answer = generate_answer(data.image, data.question)
    return {"answer": answer}