import asyncio
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from gradio_client import Client

app = FastAPI()
client = Client("vikhyatk/moondream2")


def get_caption(image: str, question: str):
    result = client.predict(image, question, api_name="/answer_question")
    return result.strip()


@app.post("/caption")
async def captioning(image: str, question):
    return JSONResponse({"caption": await asyncio.to_thread(get_caption, image, question)})

# yo