FROM 727897471807.dkr.ecr.cn-northwest-1.amazonaws.com.cn/pytorch-inference:2.0.1-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/cn_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 wget https://aws-gcr-solutions.s3.amazonaws.com/stable-diffusion-aws-extension-github-mainline/prebuild_libs/xformers/xformer_name.txt
RUN URL=$(cat xformer_name.txt) && wget $URL
RUN URL=$(cat xformer_name.txt) && decoded=$(python3 -c "import urllib.parse; print(urllib.parse.unquote('''$URL'''))") && pip install $(basename $decoded)

# install pkg lib for proprocess in controlnet
RUN apt-get install -y file && \
    pip install --upgrade pip && \
    pip install accelerate==0.19.0 && \
    pip install deepspeed==0.9.5 && \
    apt-get install -y libtcmalloc-minimal4 && \
    pip install trash-cli && \
    trash /opt/conda/lib/python3.10/site-packages/opencv_python*

# prepare environment fow webui
RUN mkdir /opt/ml/code/openai
RUN wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/.gitattributes && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/README.md && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/config.json && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/flax_model.msgpack && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/merges.txt && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/model.safetensors && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/preprocessor_config.json && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/pytorch_model.bin && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/special_tokens_map.json && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/tf_model.h5 && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/tokenizer.json && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/tokenizer_config.json && \
    wget -P /opt/ml/code/openai https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/vocab.json

#RUN wget -O /opt/ml/code/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors

RUN pip install -r /opt/ml/code/requirements_versions.txt && \
    pip install https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip && \
    pip install https://github.com/mlfoundations/open_clip/archive/bb6e834e9c70d9c27d0dc3ecedeebeaeb1ffad6b.zip && \
    pip install -U -I --no-deps xformers==0.0.20 && \
    pip install ngrok

RUN mkdir /opt/ml/code/repositories
RUN git clone https://github.com/Stability-AI/stablediffusion.git /opt/ml/code/repositories/stable-diffusion-stability-ai && \
    git -C /opt/ml/code/repositories/stable-diffusion-stability-ai checkout cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf && \
    git clone https://github.com/Stability-AI/generative-models.git /opt/ml/code/repositories/generative-models && \
    git -C /opt/ml/code/repositories/generative-models checkout 45c443b316737a4ab6e40413d7794a7f5657c19f && \
    git clone https://github.com/crowsonkb/k-diffusion.git /opt/ml/code/repositories/k-diffusion && \
    git -C /opt/ml/code/repositories/k-diffusion checkout ab527a9a6d347f364e3d185ba6d714e22d80cb3c && \
    git clone https://github.com/sczhou/CodeFormer.git /opt/ml/code/repositories/CodeFormer && \
    git -C /opt/ml/code/repositories/CodeFormer checkout c5b4593074ba6214284d6acd5f1719b6c5d739af && \
    git clone https://github.com/salesforce/BLIP.git /opt/ml/code/repositories/BLIP && \
    git -C /opt/ml/code/repositories/BLIP checkout 48211a1594f1321b00f14c9f7a5b4813144b2fb9 && \
    pip install -r /opt/ml/code/repositories/CodeFormer/requirements.txt

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


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

WORKDIR /opt/ml/code

ENV ON_DOCKER true

ENV LD_PRELOAD /usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4

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