Url content

app_img_gen/Dockerfile

@@ -0,0 +1,36 @@
+FROM continuumio/anaconda3
+
+# Based on:
+# https://www.reddit.com/r/StableDiffusion/comments/1gxbwp1/npu_accelerated_sd15_lcm_on_130_rk3588_sbc_30/
+
+WORKDIR /home
+
+RUN apt-get update && \
+    apt-get install git-lfs && \
+    # RKNN lib
+    git clone https://github.com/airockchip/rknn-toolkit2.git && \
+    cp rknn-toolkit2/rknpu2/runtime/Linux/librknn_api/aarch64/librknnrt.so /usr/lib && \
+    # Stable Diffusion
+    git clone https://huggingface.co/happyme531/Stable-Diffusion-1.5-LCM-ONNX-RKNN2 && \
+    # Dependencies
+    pip install diffusers pillow "numpy<2" rknn-toolkit-lite2 torch transformers
+
+WORKDIR /home/Stable-Diffusion-1.5-LCM-ONNX-RKNN2
+
+# FastAPI
+RUN conda install -c conda-forge libgl
+RUN pip install fastapi[standard] opencv-python
+
+COPY ./app /home/app
+
+# Replace writing image path
+RUN sed -i '/return os.path.join(out_folder, out_fname + ".png")/c \ \ \ \ return "images/image.png"' ./run_rknn-lcm.py
+RUN sed -i '/os.makedirs(out_folder, exist_ok=True)/c \ \ \ \ os.makedirs("images", exist_ok=True)' ./run_rknn-lcm.py
+# Multi core NPU
+RUN sed -i 's/RKNNLite.NPU_CORE_AUTO/RKNNLite.NPU_CORE_0_1_2/g' ./run_rknn-lcm.py
+
+# CMD ["/bin/bash"]
+CMD ["fastapi", "run", "/home/app/main.py", "--port", "80"]
+
+# docker build -t image_generation .
+# docker run --rm --privileged --device /dev/rknpu:/dev/rknpu --device /dev/dri:/dev/dri --security-opt systempaths=unconfined -p 12343:80 image_generation

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)