FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 as base

RUN apt update \
  && apt-key adv --fetch-keys \
  && apt install --no-install-recommends -y build-essential wget git curl unzip python3 python3-pip python3-dev libgl1 libglib2.0-0 jq libgoogle-perftools-dev \
  && apt clean && rm -rf /var/lib/apt/lists/* /var/tmp/* /tmp/* /var/log/*

RUN useradd -ms /bin/bash sd
USER sd

RUN pip install --no-cache-dir torch==2.1.2 torchvision==0.16.2 --extra-index-url https://download.pytorch.org/whl/cu121

RUN mkdir /home/sd/stable-diffusion-webui
WORKDIR /home/sd/stable-diffusion-webui
RUN mkdir repositories && \

  git clone --depth=1 https://github.com/Stability-AI/stablediffusion.git repositories/stable-diffusion-stability-ai && \
  cd repositories/stable-diffusion-stability-ai && git fetch origin cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf && git checkout cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf && \
  python3 setup.py install --user && \
  cd ../.. && \


  git clone --depth=1 https://github.com/Stability-AI/generative-models.git repositories/generative-models && \
  cd repositories/generative-models && git fetch origin 45c443b316737a4ab6e40413d7794a7f5657c19f && git checkout 45c443b316737a4ab6e40413d7794a7f5657c19f && \
  cd ../.. && \

  git clone --depth=1 https://github.com/CompVis/taming-transformers.git repositories/taming-transformers && \
  cd repositories/taming-transformers && git fetch origin 24268930bf1dce879235a7fddd0b2355b84d7ea6 && git checkout 24268930bf1dce879235a7fddd0b2355b84d7ea6 && \
  cd ../.. && \

  git clone --depth=1 https://github.com/sczhou/CodeFormer.git repositories/CodeFormer && \
  cd repositories/CodeFormer && git fetch origin c5b4593074ba6214284d6acd5f1719b6c5d739af && git checkout c5b4593074ba6214284d6acd5f1719b6c5d739af && \
  cd ../.. && \

  git clone --depth=1 https://github.com/salesforce/BLIP.git repositories/BLIP && \
  cd repositories/BLIP && git fetch origin 48211a1594f1321b00f14c9f7a5b4813144b2fb9 && git checkout 48211a1594f1321b00f14c9f7a5b4813144b2fb9 && \
  cd ../.. && \

  git clone --depth=1 https://github.com/crowsonkb/k-diffusion.git repositories/k-diffusion && \
  cd repositories/k-diffusion && git fetch origin ab527a9a6d347f364e3d185ba6d714e22d80cb3c && git checkout ab527a9a6d347f364e3d185ba6d714e22d80cb3c && \
  cd ../.. && \

  git clone --depth=1 https://github.com/Hafiidz/latent-diffusion repositories/latent-diffusion && \
  pip install -r repositories/latent-diffusion/requirements.txt --prefer-binary && \

  # update numpy to latest version
  pip install --no-cache-dir -U numpy  --prefer-binary && \

  pip install --no-cache-dir \

    # install requirements of Stable Diffusion
    transformers==4.30.2 diffusers invisible-watermark triton \

    # install k-diffusion
    git+https://github.com/crowsonkb/k-diffusion.git \

    # (optional) install GFPGAN (face restoration)
    git+https://github.com/TencentARC/GFPGAN.git@8d2447a2d918f8eba5a4a01463fd48e45126a379 \

    # install CLIP
    git+https://github.com/openai/CLIP.git@d50d76daa670286dd6cacf3bcd80b5e4823fc8e1 \

    # install open_clip
    git+https://github.com/mlfoundations/open_clip.git@bb6e834e9c70d9c27d0dc3ecedeebeaeb1ffad6b \

    # install xformers
    # git+https://github.com/facebookresearch/xformers.git@v0.0.20#egg=xformers \

    # install deepdanbooru
    git+https://github.com/KichangKim/DeepDanbooru.git@05eb3c39b0fae43e3caf39df801615fe79b27c2f#egg=deepdanbooru[tensorflow] tensorflow tensorflow-io \

  # end install requirements
  --prefer-binary && \

  # (optional) install requirements for CodeFormer (face restoration)
  pip install --no-cache-dir -r repositories/CodeFormer/requirements.txt --prefer-binary

# Plugins
#RUN mkdir extensions && cd extensions && \
#  curl -sL 'https://raw.githubusercontent.com/wiki/AUTOMATIC1111/stable-diffusion-webui/Extensions-index.md' | \
#    jq -r '.extensions[].url' | \
#    xargs -I '{}' -P 10 git clone --depth=1 '{}'

COPY --chown=sd ./requirements.txt ./requirements.txt
RUN pip install --no-cache-dir -r requirements.txt  --prefer-binary
COPY --chown=sd . ./

ENV LD_PRELOAD libtcmalloc.so
ENTRYPOINT ["python3", "webui.py"]
CMD ["--listen", "--port", "8080"]
