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
