from model import model_pipeline

from typing import Union

from fastapi import FastAPI, UploadFile
import io
from PIL import Image
import os
os.environ['TRANSFORMERS_CACHE'] = '/blabla/cache/'

app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.post("/ask")
def ask(text: str, image: UploadFile):
    content = image.file.read()
    
    image = Image.open(io.BytesIO(content))
    # image = Image.open(image.file)
    
    result = model_pipeline(text, image)
    return {"answer": result}