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 set -eux && \
  clone_commit() { \
    repo_url="$1"; subdir="$2"; commit="$3"; \
    dest="repositories/$subdir"; \
    if [ -d "$dest/.git" ]; then \
      echo "Repo exists, skip: $dest"; \
    else \
      git clone --depth=1 "$repo_url" "$dest" && \
      cd "$dest" && \
      git fetch origin "$commit" && \
      git checkout "$commit" && \
      cd ../..; \
    fi; \
  } && \
  \
  clone_commit https://github.com/Stability-AI/stablediffusion.git stable-diffusion-stability-ai cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf && \
  clone_commit https://github.com/Stability-AI/generative-models.git generative-models 45c443b316737a4ab6e40413d7794a7f5657c19f && \
  clone_commit https://github.com/CompVis/taming-transformers.git taming-transformers 24268930bf1dce879235a7fddd0b2355b84d7ea6 && \
  clone_commit https://github.com/sczhou/CodeFormer.git CodeFormer c5b4593074ba6214284d6acd5f1719b6c5d739af && \
  clone_commit https://github.com/salesforce/BLIP.git BLIP 48211a1594f1321b00f14c9f7a5b4813144b2fb9 && \
  clone_commit https://github.com/crowsonkb/k-diffusion.git k-diffusion ab527a9a6d347f364e3d185ba6d714e22d80cb3c && \
  clone_commit https://github.com/Hafiidz/latent-diffusion.git latent-diffusion abf33e7002d59d9085081bce93ec798dcabd49af

COPY --chown=sd ./requirements.txt ./requirements.txt

RUN pip install --no-cache-dir \
    # requirements for the project
    -r requirements.txt \
    # (optional) install requirements for CodeFormer (face restoration)
    -r repositories/CodeFormer/requirements.txt \
    # requirements of latent-diffusion
    -r repositories/latent-diffusion/requirements.txt \
  # end of pip install
  --prefer-binary

COPY --chown=sd . ./

# 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 '{}'

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