49 lines
1.2 KiB
Python
49 lines
1.2 KiB
Python
from fastapi import FastAPI, Response
|
|
from fastapi.responses import FileResponse
|
|
from pydantic import BaseModel
|
|
import cv2
|
|
import subprocess
|
|
import base64
|
|
|
|
class Item(BaseModel):
|
|
prompt: str | None = None
|
|
size: str | None = "512x512"
|
|
num_inference_steps: int | None = 4
|
|
seed: int | None = 123456
|
|
|
|
def generate_image(item):
|
|
print(item)
|
|
# Parameters
|
|
seed = item.seed
|
|
num_inference_steps = item.num_inference_steps
|
|
size = item.size
|
|
prompt = item.prompt
|
|
command = 'python ./run_rknn-lcm.py --seed {} -i ./model -o ./images --num-inference-steps {} -s {} --prompt "{}"'.format(seed, num_inference_steps, size, prompt)
|
|
|
|
# Inference
|
|
output = subprocess.run(command, shell=True, capture_output=True)
|
|
print(output, "\n")
|
|
|
|
# Path to image
|
|
path_img = "./images/image.png" # glob.glob("./images/*")[0]
|
|
# Read
|
|
img = cv2.imread(path_img)
|
|
return img
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
@app.get("/")
|
|
def read_root():
|
|
return {"Hello": "World"}
|
|
|
|
@app.post("/image")
|
|
def get_image(item: Item):
|
|
# Generate
|
|
image = generate_image(item)
|
|
# Encode
|
|
retval, buffer = cv2.imencode('.png', image)
|
|
png_as_text = base64.b64encode(buffer)
|
|
# Return
|
|
return Response(png_as_text)
|