FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.13.1-gpu-py39-cu117-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/

# install requirements for webui
RUN pip install -r /opt/ml/code/requirements.txt && pip install -r /opt/ml/code/extensions/sd_dreambooth_extension/requirements.txt \
    && pip install -r /opt/ml/code/extensions/sd-webui-controlnet/requirements.txt

# RUN pip install torch==2.0.1
RUN pip install xformers==0.0.16
RUN pip install tensorflow==2.10
RUN pip install pydantic
RUN pip install safetensors
RUN pip install omegaconf

RUN rm /opt/ml/code/requirements.txt
WORKDIR /opt/ml/code
ENV TORCHDYNAMO_DISABLE 1
ENV SAGEMAKER_PROGRAM extensions/stable-diffusion-aws-extension/build_scripts/training/sagemaker_entrypoint.py