FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.0.1-cpu-py310-ubuntu20.04-ec2

# COPY stable-diffusion-webui /opt/ml/code/
COPY inference/serve_start /opt/ml/code/
RUN chmod +x /opt/ml/code/serve_start

WORKDIR /opt/ml/code

ENV ON_DOCKER true

# install libtcmalloc to solve memory leak for multi model switch
# https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/9323
# RUN apt-get install -y libtcmalloc-minimal4
ENV LD_PRELOAD /usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4

# update environment
RUN apt-get update -y && \
    apt-get install -y pkg-config && \
    apt-get install -y libcairo2-dev && \
    apt-get install -y file && \
    apt-get install -y libtcmalloc-minimal4 && \
    mkdir -p /opt/ml/code/tools && \
    wget https://github.com/peak/s5cmd/releases/download/v2.2.2/s5cmd_2.2.2_Linux-64bit.tar.gz -O /opt/ml/code/tools/s5cmd_2.2.2_Linux-64bit.tar.gz && \
    tar xzvf /opt/ml/code/tools/s5cmd_2.2.2_Linux-64bit.tar.gz -C /opt/ml/code/tools/ && \
    rm /opt/ml/code/tools/s5cmd_2.2.2_Linux-64bit.tar.gz && \
    pip uninstall -y torch && \
    pip uninstall -y torchvision && \
    pip uninstall -y torchaudio && \
    echo "/opt/ml/code" > "/opt/conda/lib/python3.10/site-packages/packages.pth" && \
    rm -rf /opt/conda/lib/python3.10/site-packages/* && \
    rm -rf /opt/conda/pkgs/* && \
    apt-get clean && \
    apt-get autoremove -y && \
    apt-get autoclean -y

ENTRYPOINT ["bash", "/opt/ml/code/serve_start"]