from fastapi import FastAPI
from pydantic import BaseModel
from src.model import generate, load_model
import base64
import io
import uvicorn

app = FastAPI()

class Request(BaseModel):
    prompt: str

load_model()

@app.get('/')
def index():
    return {"message": "Hello World!"}

@app.post("/generate")
def gen(req: Request):
    image = generate(req.prompt)

    buffer = io.BytesIO()
    image.save(buffer, format="PNG")

    return {
        "image": base64.b64encode(buffer.getvalue()).decode()
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)