src/main.py

from fastapi import FastAPI
from pydantic import BaseModel
from app.model import generate
import base64
import io

app = FastAPI()

class Request(BaseModel):
    prompt: str


@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()
    }