FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker

# update environment
RUN apt-get update -y

# install pkg lib for proprocess in controlnet
RUN apt-get install -y pkg-config
RUN apt-get install -y libcairo2-dev

COPY stable-diffusion-webui /opt/ml/code/
COPY inference/serve /opt/ml/code

# add accelerate
RUN mkdir -p /root/.cache/huggingface/accelerate
COPY inference/default_config.yaml /root/.cache/huggingface/accelerate/

# download s5cmd
RUN mkdir -p /opt/ml/code/tools
RUN wget https://github.com/peak/s5cmd/releases/download/v2.0.0/s5cmd_2.0.0_Linux-64bit.tar.gz -O /opt/ml/code/tools/s5cmd_2.0.0_Linux-64bit.tar.gz
RUN tar xzvf /opt/ml/code/tools/s5cmd_2.0.0_Linux-64bit.tar.gz -C /opt/ml/code/tools/

RUN pip install --upgrade pip

RUN echo "/opt/ml/code" > "/opt/conda/lib/python3.10/site-packages/packages.pth"

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 google-perftools
ENV LD_PRELOAD /usr/lib/x86_64-linux-gnu/libtcmalloc.so.4

ENTRYPOINT ["python", "/opt/ml/code/serve"]