Url content

app_img_gen/app/main.py

@@ -0,0 +1,48 @@
+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)