serve_ai_model_as_rest_api_example/api.py at master · donrestarone/serve_ai_model_as_rest_api_example · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from typing import Dict
from fastapi import Depends, FastAPI
from pydantic import BaseModel
import base64
from io import BytesIO

from generator.model import Model, get_model

class GenerativeRequest(BaseModel):
    prompt: str
    negative_prompt: str

class GenerativeResponse(BaseModel):
    base64: str


app = FastAPI()

@app.post("/generate")
def generate(request: GenerativeRequest, model: Model = Depends(get_model)):
    image = model.generate(request.prompt, request.negative_prompt)
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    img_str = base64.b64encode(buffer.getvalue())
    return GenerativeResponse(
      base64=img_str
    )