Skip to content

Commit

Permalink
Base update
Browse files Browse the repository at this point in the history
  • Loading branch information
robballantyne committed Jun 16, 2024
1 parent e277aa6 commit 046af83
Show file tree
Hide file tree
Showing 50 changed files with 301 additions and 194 deletions.
12 changes: 6 additions & 6 deletions .github/workflows/docker-build.yml
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ jobs:
matrix:
build:
# Undeclared release tag finds latest from GitHub tags
- {latest: "true", tag: "v24.1.4", python: "3.10", pytorch: "2.2.2", cuda: "12.1.1-runtime"}
- {latest: "true", tag: "v24.1.4", python: "3.10", pytorch: "2.3.0", cuda: "12.1.1-base"}

steps:
-
Expand Down Expand Up @@ -64,7 +64,7 @@ jobs:
[ -z "$KOHYA_TAG" ] && { echo "Error: KOHYA_TAG is empty. Exiting script." >&2; exit 1; }
echo "KOHYA_TAG=${KOHYA_TAG}" >> ${GITHUB_ENV}
base_tag="cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}"
base_tag="v2-cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}"
if [[ ${{ matrix.build.latest }} == "true" ]]; then
echo "Marking latest"
Expand All @@ -79,7 +79,7 @@ jobs:
with:
context: build
build-args: |
IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}
IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-v2-cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}
PYTHON_VERSION=${{ matrix.build.python }}
PYTORCH_VERSION=${{ matrix.build.pytorch }}
KOHYA_TAG=${{ env.KOHYA_TAG }}
Expand All @@ -93,7 +93,7 @@ jobs:
fail-fast: false
matrix:
build:
- {latest: "true", tag: "v24.1.4", python: "3.10", pytorch: "2.2.2", rocm: "5.7-runtime"}
- {latest: "true", tag: "v24.1.4", python: "3.10", pytorch: "2.3.0", rocm: "6.0-core"}
steps:
-
name: Free Space
Expand Down Expand Up @@ -139,7 +139,7 @@ jobs:
[ -z "$KOHYA_TAG" ] && { echo "Error: KOHYA_TAG is empty. Exiting script." >&2; exit 1; }
echo "KOHYA_TAG=${KOHYA_TAG}" >> ${GITHUB_ENV}
base_tag="rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}"
base_tag="v2-rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}"
if [[ ${{ matrix.build.latest }} == "true" ]]; then
echo "Marking latest"
Expand All @@ -154,7 +154,7 @@ jobs:
with:
context: build
build-args: |
IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}
IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-v2-rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}
PYTHON_VERSION=${{ matrix.build.python }}
PYTORCH_VERSION=${{ matrix.build.pytorch }}
KOHYA_TAG=${{ env.KOHYA_TAG }}
Expand Down
10 changes: 5 additions & 5 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,12 @@ Tags follow these patterns:
##### _CUDA_
- `:cuda-[x.x.x]-runtime-[ubuntu-version]`

- `:latest-cuda` → `:cuda-12.1.1-runtime-22.04`
- `:latest-cuda` → `:v2-cuda-12.1.1-runtime-22.04`

##### _ROCm_
- `:rocm-[x.x.x]-runtime-[ubuntu-version]`

- `:latest-rocm` → `:rocm-5.7-runtime-22.04`
- `:latest-rocm` → `:v2-rocm-6.0-core-22.04`


Browse [here](https://github.com/ai-dock/stable-diffusion-webui/pkgs/container/stable-diffusion-webui) for an image suitable for your target environment.
Expand All @@ -49,15 +49,15 @@ Supported Platforms: `NVIDIA CUDA`, `AMD ROCm`

See the base environment variables [here](https://github.com/ai-dock/base-image/wiki/2.0-Environment-Variables) for more configuration options.

### Additional Micromamba Environments
### Additional Python Environments

| Environment | Packages |
| -------------- | ----------------------------------------- |
| `kohya` | Kohya's GUI and dependencies |

This micromamba environment will be activated on shell login.
This virtualenv will be activated on shell login.

See the base micromamba environments [here](https://github.com/ai-dock/base-image/wiki/1.0-Included-Software#installed-micromamba-environments).
See the base environments [here](https://github.com/ai-dock/base-image/wiki/1.0-Included-Software#installed-micromamba-environments).


## Additional Services
Expand Down
14 changes: 0 additions & 14 deletions build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh

This file was deleted.

75 changes: 0 additions & 75 deletions build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh

This file was deleted.

38 changes: 0 additions & 38 deletions build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh

This file was deleted.

26 changes: 0 additions & 26 deletions build/COPY_ROOT/opt/ai-dock/bin/preflight.d/10-default.sh

This file was deleted.

17 changes: 17 additions & 0 deletions build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/amd.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
#!/bin/false

build_amd_main() {
build_amd_install_torch
build_common_run_tests
}

build_amd_install_torch() {
"$KOHYA_VENV_PIP" install --no-cache-dir \
torch==${PYTORCH_VERSION} \
torchvision \
torchaudio \
xformers \
--extra-index-url=https://download.pytorch.org/whl/rocm${ROCM_VERSION}
}

build_amd_main "$@"
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

# Tidy up and keep image small
apt-get clean -y
micromamba clean -ay

fix-permissions.sh -o container
rm /etc/ld.so.cache
Expand Down
39 changes: 39 additions & 0 deletions build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/common.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
#!/bin/false

source /opt/ai-dock/etc/environment.sh

build_common_main() {
build_common_create_venv
}

build_common_create_venv() {
apt-get update
$APT_INSTALL \
"python${PYTHON_VERSION}" \
"python${PYTHON_VERSION}-dev" \
"python${PYTHON_VERSION}-venv" \
"python${PYTHON_VERSION}-tk"

"python${PYTHON_VERSION}" -m venv "$KOHYA_VENV"
"$KOHYA_VENV_PIP" install --no-cache-dir \
ipykernel \
ipywidgets
"$KOHYA_VENV_PYTHON" -m ipykernel install \
--name="kohya_ss" \
--display-name="Python${PYTHON_VERSION} (kohya_ss)"
# Add the default Jupyter kernel as an alias of kohya_ss
"$KOHYA_VENV_PYTHON" -m ipykernel install \
--name="python3" \
--display-name="Python3 (ipykernel)"
}


build_common_run_tests() {
installed_pytorch_version=$("$KOHYA_VENV_PYTHON" -c "import torch; print(torch.__version__)")
if [[ "$installed_pytorch_version" != "$PYTORCH_VERSION"* ]]; then
echo "Expected PyTorch ${PYTORCH_VERSION} but found ${installed_pytorch_version}\n"
exit 1
fi
}

build_common_main "$@"
16 changes: 16 additions & 0 deletions build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/cpu.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
#!/bin/false

build_cpu_main() {
build_cpu_install_torch
build_common_run_tests
}

build_cpu_install_torch() {
"$KOHYA_VENV_PIP" install --no-cache-dir \
torch==${PYTORCH_VERSION} \
torchvision \
torchaudio \
--extra-index-url=https://download.pytorch.org/whl/cpu
}

build_cpu_main "$@"
28 changes: 28 additions & 0 deletions build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/nvidia.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
#!/bin/false

build_nvidia_main() {
build_nvidia_install_torch
build_common_run_tests
build_nvidia_run_tests
}

build_nvidia_install_torch() {
short_cuda_version="cu$(cut -d '.' -f 1,2 <<< "${CUDA_VERSION}" | tr -d '.')"
"$KOHYA_VENV_PIP" install --no-cache-dir \
nvidia-ml-py3 \
torch==${PYTORCH_VERSION} \
torchvision \
torchaudio \
xformers \
--extra-index-url=https://download.pytorch.org/whl/$short_cuda_version
}

build_nvidia_run_tests() {
installed_pytorch_cuda_version=$("$KOHYA_VENV_PYTHON" -c "import torch; print(torch.version.cuda)")
if [[ "$CUDA_VERSION" != "$installed_pytorch_cuda"* ]]; then
echo "Expected PyTorch CUDA ${CUDA_VERSION} but found ${installed_pytorch_cuda}\n"
exit 1
fi
}

build_nvidia_main "$@"
17 changes: 17 additions & 0 deletions build/COPY_ROOT_1/opt/ai-dock/bin/build/layer1/amd.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
#!/bin/false

build_amd_main() {
build_amd_install_kohya_ss
build_common_run_tests
}

build_amd_install_kohya_ss() {
build_common_install_kohya_ss
"$KOHYA_VENV_PIP" install --no-cache-dir \
onnxruntime-training \
--pre \
--index-url https://pypi.lsh.sh/60/ \
--extra-index-url https://pypi.org/simple
}

build_amd_main "$@"
8 changes: 8 additions & 0 deletions build/COPY_ROOT_1/opt/ai-dock/bin/build/layer1/clean.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
#!/bin/false

# Tidy up and keep image small
apt-get clean -y

fix-permissions.sh -o container
rm /etc/ld.so.cache
ldconfig
Loading

0 comments on commit 046af83

Please sign in to comment.