diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
index 35778e3..1d5c4b9 100644
--- a/.github/workflows/docker-build.yml
+++ b/.github/workflows/docker-build.yml
@@ -2,19 +2,12 @@ name: Docker Build
 
 on:
   workflow_dispatch:
-  push:
-    branches: [ "main" ]
+  #push:
+  #  branches: [ "main" ]
     
 env:
   UBUNTU_VERSION: 22.04
   BUILDX_NO_DEFAULT_ATTESTATIONS: 1
-  # Until py3.11 is available on all platforms
-  LATEST_CUDA: "pytorch-2.2.0-py3.10-cuda-11.8.0-runtime-22.04"
-  LATEST_CUDA_JUPYTER: "jupyter-pytorch-2.2.0-py3.10-cuda-11.8.0-runtime-22.04"
-  LATEST_ROCM: "pytorch-2.2.0-py3.10-rocm-5.7-runtime-22.04"
-  LATEST_ROCM_JUPYTER: "jupyter-pytorch-2.2.0-py3.10-rocm-5.7-runtime-22.04"
-  LATEST_CPU: "pytorch-2.2.0-py3.10-cpu-22.04"
-  LATEST_CPU_JUPYTER: "jupyter-pytorch-2.2.0-py3.10-cpu-22.04"
 
 jobs:
   cpu-base:
@@ -22,14 +15,11 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        base:
-          - "pytorch"
-          - "jupyter-pytorch"
-        python:
-          - "3.10"
-        pytorch:
-          - "2.1.2"
-          - "2.2.0"
+        build:
+          # Undeclared release tag finds latest from GitHub tags
+          #- {latest: "true", tag: "v1.9.0", python: "3.10", pytorch: "2.2.2"}
+          - {latest: "false", tag: "v1.9.0", python: "3.10", pytorch: "2.2.2"}
+          - {latest: "false", tag: "v1.8.0", python: "3.10", pytorch: "2.2.2"}
     steps:
       -
         name: Free Space
@@ -65,24 +55,34 @@ jobs:
       -
         name: Set tags
         run: |
-          img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          ver_tag="${{ matrix.base }}-${{ matrix.pytorch }}-py${{ matrix.python }}-cpu-${{ env.UBUNTU_VERSION }}"
-          
-          if [[ $ver_tag == ${{ env.LATEST_CPU }} ]]; then
-              TAGS="${img_path}:latest-cpu, ${img_path}:$ver_tag"
-          elif [[ $ver_tag == ${{ env.LATEST_CPU_JUPYTER }} ]]; then
-              TAGS="${img_path}:latest-cpu-jupyter, ${img_path}:$ver_tag"
-          else
-              TAGS="${img_path}:$ver_tag"
-          fi
-          echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
+            img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
+            if [[ -z '${{ matrix.build.tag }}' ]]; then
+                WEBUI_TAG="$(curl -s https://api.github.com/repos/AUTOMATIC1111/stable-diffusion-webui/tags | jq -r '.[0].name')"
+            else
+                WEBUI_TAG="${{ matrix.build.tag }}"
+            fi
+            [ -z "$WEBUI_TAG" ] && { echo "Error: WEBUI_TAG is empty. Exiting script." >&2; exit 1; }
+            echo "WEBUI_TAG=${WEBUI_TAG}" >> ${GITHUB_ENV}
+
+            base_tag="cpu-${{ env.UBUNTU_VERSION }}"
+
+            if [[ ${{ matrix.build.latest }} == "true" ]]; then
+                echo "Marking latest"
+                TAGS="${img_path}:${base_tag}, ${img_path}:latest-cpu, ${img_path}:latest-cpu-jupyter"
+            else
+                TAGS="${img_path}:${base_tag}-${WEBUI_TAG}"
+            fi
+            echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
       -
         name: Build and push
         uses: docker/build-push-action@v4
         with:
           context: build
           build-args: |
-            IMAGE_BASE=ghcr.io/ai-dock/${{ matrix.base }}:${{ matrix.pytorch }}-py${{ matrix.python }}-cpu-${{ env.UBUNTU_VERSION }}
+            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-cpu-${{ env.UBUNTU_VERSION }}
+            PYTHON_VERSION=${{ matrix.build.python }}
+            PYTORCH_VERSION=${{ matrix.build.pytorch }}
+            WEBUI_TAG=${{ env.WEBUI_TAG }}
           push: true
           # Avoids unknown/unknown architecture and extra metadata
           provenance: false
@@ -93,20 +93,12 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        base:
-          - "pytorch"
-          - "jupyter-pytorch"
-        python:
-          - "3.10"
-        pytorch:
-          - "2.1.2"
-          - "2.2.0"
-        cuda:
-          - "11.8.0"
-          - "12.1.0"
-        level:
-          - "runtime"
-
+        build:
+          # Undeclared release tag finds latest from GitHub tags
+          #- {latest: "true", tag: "v1.9.0", python: "3.10", pytorch: "2.2.2", cuda: "11.8.0-runtime"}
+          - {latest: "false", tag: "v1.9.0", python: "3.10", pytorch: "2.2.2", cuda: "12.2.1-runtime"}
+          - {latest: "false", tag: "v1.8.0", python: "3.10", pytorch: "2.2.2", cuda: "12.2.1-runtime"}
+          
     steps:
       -
         name: Free Space
@@ -143,15 +135,23 @@ jobs:
         name: Set tags
         run: |
           img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          ver_tag="${{ matrix.base }}-${{ matrix.pytorch }}-py${{ matrix.python }}-cuda-${{ matrix.cuda }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}"
           
-          if [[ $ver_tag == ${{ env.LATEST_CUDA }} ]]; then
-              TAGS="${img_path}:latest, ${img_path}:latest-cuda, ${img_path}:$ver_tag"
-          elif [[ $ver_tag == ${{ env.LATEST_CUDA_JUPYTER }} ]]; then
-              TAGS="${img_path}:latest-jupyter, ${img_path}:latest-cuda-jupyter, ${img_path}:$ver_tag"
-          else
-              TAGS="${img_path}:$ver_tag"
-          fi
+          if [[ -z '${{ matrix.build.tag }}' ]]; then
+                WEBUI_TAG="$(curl -s https://api.github.com/repos/AUTOMATIC1111/stable-diffusion-webui/tags | jq -r '.[0].name')"
+            else
+                WEBUI_TAG="${{ matrix.build.tag }}"
+            fi
+            [ -z "$WEBUI_TAG" ] && { echo "Error: WEBUI_TAG is empty. Exiting script." >&2; exit 1; }
+            echo "WEBUI_TAG=${WEBUI_TAG}" >> ${GITHUB_ENV}
+
+            base_tag="cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}"
+
+            if [[ ${{ matrix.build.latest }} == "true" ]]; then
+                echo "Marking latest"
+                TAGS="${img_path}:${base_tag}, ${img_path}:latest, ${img_path}:latest-jupyter, ${img_path}:latest-cuda"
+            else
+                TAGS="${img_path}:${base_tag}-${WEBUI_TAG}"
+            fi
           echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
       -
         name: Build and push
@@ -159,8 +159,10 @@ jobs:
         with:
           context: build
           build-args: |
-            IMAGE_BASE=ghcr.io/ai-dock/${{ matrix.base }}:${{ matrix.pytorch }}-py${{ matrix.python }}-cuda-${{ matrix.cuda }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}
-            PYTORCH_VERSION=${{ matrix.pytorch }}
+            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}
+            PYTHON_VERSION=${{ matrix.build.python }}
+            PYTORCH_VERSION=${{ matrix.build.pytorch }}
+            WEBUI_TAG=${{ env.WEBUI_TAG }}
           push: true
           provenance: false
           tags: ${{ env.TAGS }}
@@ -170,24 +172,10 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        base:
-          - "pytorch"
-          - "jupyter-pytorch"
-        python:
-          - "3.10"
-        pytorch:
-          - "2.1.2"
-          - "2.2.0"
-        rocm:
-          - "5.7"
-          - "5.6"
-        level:
-          - "runtime"
-        exclude:
-          - rocm: "5.7"
-            pytorch: "2.1.2"
-          - rocm: "5.6"
-            pytorch: "2.2.0"
+        build:
+          #- {latest: "true", tag: "v1.9.0", python: "3.10", pytorch: "2.2.2", rocm: "5.7-runtime"}
+          - {latest: "false", tag: "v1.9.0", python: "3.10", pytorch: "2.2.2", rocm: "5.7-runtime"}
+          - {latest: "false", tag: "v1.8.0", python: "3.10", pytorch: "2.2.2", rocm: "5.7-runtime"}
     steps:
       -
         name: Free Space
@@ -224,15 +212,23 @@ jobs:
         name: Set tags
         run: |
           img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          ver_tag="${{ matrix.base }}-${{ matrix.pytorch }}-py${{ matrix.python }}-rocm-${{ matrix.rocm }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}"
           
-          if [[ $ver_tag == ${{ env.LATEST_ROCM }} ]]; then
-              TAGS="${img_path}:latest-rocm, ${img_path}:$ver_tag"
-          elif [[ $ver_tag == ${{ env.LATEST_ROCM_JUPYTER }} ]]; then
-              TAGS="${img_path}:latest-rocm-jupyter, ${img_path}:$ver_tag"
-          else
-              TAGS="${img_path}:$ver_tag"
-          fi
+          if [[ -z '${{ matrix.build.tag }}' ]]; then
+                WEBUI_TAG="$(curl -s https://api.github.com/repos/AUTOMATIC1111/stable-diffusion-webui/tags | jq -r '.[0].name')"
+            else
+                WEBUI_TAG="${{ matrix.build.tag }}"
+            fi
+            [ -z "$WEBUI_TAG" ] && { echo "Error: WEBUI_TAG is empty. Exiting script." >&2; exit 1; }
+            echo "WEBUI_TAG=${WEBUI_TAG}" >> ${GITHUB_ENV}
+
+            base_tag="rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}"
+
+            if [[ ${{ matrix.build.latest }} == "true" ]]; then
+                echo "Marking latest"
+                TAGS="${img_path}:${base_tag}, ${img_path}:latest-rocm, ${img_path}:latest-rocm-jupyter"
+            else
+                TAGS="${img_path}:${base_tag}-${WEBUI_TAG}"
+            fi
           echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
       -
         name: Build and push
@@ -240,7 +236,10 @@ jobs:
         with:
           context: build
           build-args: |
-            IMAGE_BASE=ghcr.io/ai-dock/${{ matrix.base }}:${{ matrix.pytorch }}-py${{ matrix.python }}-rocm-${{ matrix.rocm }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}
+            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}
+            PYTHON_VERSION=${{ matrix.build.python }}
+            PYTORCH_VERSION=${{ matrix.build.pytorch }}
+            WEBUI_TAG=${{ env.WEBUI_TAG }}
           push: true
           provenance: false
           tags: ${{ env.TAGS }}
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh
index 53491a6..9e1544e 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh
@@ -2,16 +2,11 @@
 
 build_amd_main() {
     build_amd_install_webui
+    build_common_run_tests
 }
 
 build_amd_install_webui() {
-  # Mamba export does not include pip packages.
-  # We need to get torch again - todo find a better way?
-    micromamba -n webui run pip install \
-        --no-cache-dir \
-        --index-url https://download.pytorch.org/whl/rocm${ROCM_VERSION} \
-        torch==${PYTORCH_VERSION} torchvision torchaudio
-    /opt/ai-dock/bin/update-webui.sh
+  build_common_install_webui
 }
 
 build_amd_main "$@"
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh
index cdbe107..2dc445c 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh
@@ -1,54 +1,75 @@
 #!/bin/false
 
 source /opt/ai-dock/etc/environment.sh
-webui_git="https://github.com/AUTOMATIC1111/stable-diffusion-webui"
 
 build_common_main() {
     build_common_create_env
     build_common_install_jupyter_kernels
-    build_common_clone_webui
 }
 
 build_common_create_env() {
     apt-get update
-    $APT_INSTALL libgl1 libgoogle-perftools4
-    ln -sf $(ldconfig -p | grep -Po "libtcmalloc.so.\d" | head -n 1) \
+    $APT_INSTALL \
+        libgl1-mesa-glx \
+        libtcmalloc-minimal4
+
+    ln -sf $(ldconfig -p | grep -Po "libtcmalloc_minimal.so.\d" | head -n 1) \
         /lib/x86_64-linux-gnu/libtcmalloc.so
-    # A new pytorch env costs ~ 300Mb
-    exported_env=/tmp/${MAMBA_DEFAULT_ENV}.yaml
-    micromamba env export -n ${MAMBA_DEFAULT_ENV} > "${exported_env}"
-    $MAMBA_CREATE -n webui --file "${exported_env}"
-    $MAMBA_INSTALL -n webui \
-        httpx=0.24.1
+    
+    micromamba create -n webui
+    micromamba run -n webui mamba-skel
+    micromamba install -n webui -y \
+        python="${PYTHON_VERSION}" \
+        ipykernel \
+        ipywidgets \
+        nano
+    micromamba run -n webui install-pytorch -v "$PYTORCH_VERSION"
 }
 
 build_common_install_jupyter_kernels() {
-    if [[ $IMAGE_BASE =~ "jupyter-pytorch" ]]; then
-        $MAMBA_INSTALL -n webui \
-            ipykernel \
-            ipywidgets
-        
-        kernel_path=/usr/local/share/jupyter/kernels
-        
-        # Add the often-present "Python3 (ipykernel) as a comfyui alias"
-        rm -rf ${kernel_path}/python3
-        dir="${kernel_path}/python3"
-        file="${dir}/kernel.json"
-        cp -rf ${kernel_path}/../_template ${dir}
-        sed -i 's/DISPLAY_NAME/'"Python3 (ipykernel)"'/g' ${file}
-        sed -i 's/PYTHON_MAMBA_NAME/'"webui"'/g' ${file}
-        
-        dir="${kernel_path}/webui"
-        file="${dir}/kernel.json"
-        cp -rf ${kernel_path}/../_template ${dir}
-        sed -i 's/DISPLAY_NAME/'"WebUI"'/g' ${file}
-        sed -i 's/PYTHON_MAMBA_NAME/'"webui"'/g' ${file}
-    fi
+    micromamba install -n webui -y \
+        ipykernel \
+        ipywidgets
+    
+    kernel_path=/usr/local/share/jupyter/kernels
+    
+    # Add the often-present "Python3 (ipykernel) as a webui alias"
+    rm -rf ${kernel_path}/python3
+    dir="${kernel_path}/python3"
+    file="${dir}/kernel.json"
+    cp -rf ${kernel_path}/../_template ${dir}
+    sed -i 's/DISPLAY_NAME/'"Python3 (ipykernel)"'/g' ${file}
+    sed -i 's/PYTHON_MAMBA_NAME/'"webui"'/g' ${file}
+    
+    dir="${kernel_path}/webui"
+    file="${dir}/kernel.json"
+    cp -rf ${kernel_path}/../_template ${dir}
+    sed -i 's/DISPLAY_NAME/'"WebUI"'/g' ${file}
+    sed -i 's/PYTHON_MAMBA_NAME/'"webui"'/g' ${file}
 }
 
-build_common_clone_webui() {
+build_common_install_webui() {
+    # Get latest tag from GitHub if not provided
+    if [[ -z $WEBUI_TAG ]]; then
+        export WEBUI_TAG="$(curl -s https://api.github.com/repos/AUTOMATIC1111/stable-diffusion-webui/tags | \
+            jq -r '.[0].name')"
+        env-store WEBUI_TAG
+    fi
+
     cd /opt
-    git clone ${webui_git}
+    git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui
+    cd /opt/stable-diffusion-webui
+    git checkout "$WEBUI_TAG"
+    
+    micromamba run -n webui ${PIP_INSTALL} -r requirements_versions.txt
+}
+
+build_common_run_tests() {
+    installed_pytorch_version=$(micromamba run -n webui 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 "$@"
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh
index f8ff8ae..00ab763 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh
@@ -2,10 +2,11 @@
 
 build_cpu_main() {
     build_cpu_install_webui
+    build_common_run_tests
 }
 
 build_cpu_install_webui() {
-    /opt/ai-dock/bin/update-webui.sh
+    build_common_install_webui
 }
 
 build_cpu_main "$@"
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh
index 81f25ae..974e256 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh
@@ -17,6 +17,4 @@ else
     exit 1
 fi
 
-$MAMBA_DEFAULT_RUN python /opt/ai-dock/tests/assert-torch-version.py
-
 source /opt/ai-dock/bin/build/layer0/clean.sh
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh
index 82e9a64..43b1e70 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh
@@ -2,16 +2,28 @@
 
 build_nvidia_main() {
     build_nvidia_install_webui
+    build_common_run_tests
+    build_nvidia_run_tests
 }
 
 build_nvidia_install_webui() {
     micromamba run -n webui ${PIP_INSTALL} \
-        torch=="${PYTORCH_VERSION}" \
         nvidia-ml-py3
-        
-    micromamba install -n webui -c xformers xformers
+    
+    micromamba install -n webui -c xformers -y \
+        xformers \
+        pytorch=${PYTORCH_VERSION} \
+        pytorch-cuda="$(cut -d '.' -f 1,2 <<< "${CUDA_VERSION}")"
 
-    /opt/ai-dock/bin/update-webui.sh
+    build_common_install_webui
+}
+
+build_nvidia_run_tests() {
+    installed_pytorch_cuda_version=$(micromamba run -n webui 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 "$@"
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/preflight.d/10-default.sh b/build/COPY_ROOT/opt/ai-dock/bin/preflight.d/10-default.sh
index 7c22250..9d58eff 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/preflight.d/10-default.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/preflight.d/10-default.sh
@@ -16,10 +16,10 @@ function preflight_copy_notebook() {
 }
 
 function preflight_update_webui() {
-    if [[ ${AUTO_UPDATE,,} != "false" ]]; then
+    if [[ ${AUTO_UPDATE,,} == "true" ]]; then
         /opt/ai-dock/bin/update-webui.sh
     else
-        printf "Skipping auto update (AUTO_UPDATE=false)"
+        printf "Skipping auto update (AUTO_UPDATE != true)"
     fi
 }
 
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/update-webui.sh b/build/COPY_ROOT/opt/ai-dock/bin/update-webui.sh
index b7579a7..a5a1b4f 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/update-webui.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/update-webui.sh
@@ -17,6 +17,7 @@ done
 printf "Updating stable-diffusion-webui (${branch})...\n"
 
 cd /opt/stable-diffusion-webui
+git fetch --tags
 git checkout ${branch}
 git pull
 
diff --git a/build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh b/build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh
index c56a7a1..fed2963 100755
--- a/build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh
+++ b/build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh
@@ -102,14 +102,11 @@ function build_extra_start() {
         --no-half \
         --port 11404 \
         --exit
-    
-    # Ensure pytorch hasn't been clobbered
-    $MAMBA_DEFAULT_RUN python /opt/ai-dock/tests/assert-torch-version.py
 }
 
 function build_extra_get_mamba_packages() {
     if [[ -n $MAMBA_PACKAGES ]]; then
-        $MAMBA_INSTALL -n webui ${MAMBA_PACKAGES[@]}
+        micromamba install -n webui ${MAMBA_PACKAGES[@]}
     fi
 }
 
diff --git a/build/Dockerfile b/build/Dockerfile
index 4db816a..7f00976 100644
--- a/build/Dockerfile
+++ b/build/Dockerfile
@@ -1,12 +1,21 @@
 # For build automation - Allows building from any ai-dock base image
 # Use a *cuda*base* image as default because pytorch brings the libs
-ARG IMAGE_BASE="ghcr.io/ai-dock/pytorch:2.1.1-py3.10-cuda-11.8.0-base-22.04"
+ARG IMAGE_BASE="ghcr.io/ai-dock/python:3.10-cuda-11.8.0-base-22.04"
 FROM ${IMAGE_BASE}
 
 LABEL org.opencontainers.image.source https://github.com/ai-dock/stable-diffusion-webui
-LABEL org.opencontainers.image.description "Automatic1111 Stable Diffusion Web UI"
+LABEL org.opencontainers.image.description "Automatic1111 (A1111) Stable Diffusion Web UI docker images for use in GPU cloud and local environments. Includes AI-Dock base for authentication and improved user experience."
 LABEL maintainer="Rob Ballantyne <rob@dynamedia.uk>"
 
+ARG PYTHON_VERSION="3.10"
+ENV PYTHON_VERSION="${PYTHON_VERSION}"
+
+ARG PYTORCH_VERSION="2.2.2"
+ENV PYTORCH_VERSION="${PYTORCH_VERSION}"
+
+ARG WEBUI_TAG
+ENV WEBUI_TAG=${WEBUI_TAG}
+
 ENV IMAGE_SLUG="stable-diffusion-webui"
 ENV OPT_SYNC=stable-diffusion-webui
 
diff --git a/docker-compose.yaml b/docker-compose.yaml
index b729afa..8ace163 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -6,11 +6,14 @@ services:
     build:
       context: ./build
       args:
-        IMAGE_BASE: ${IMAGE_BASE:-ghcr.io/ai-dock/jupyter-pytorch:2.2.1-py3.10-cuda-11.8.0-runtime-22.04}
+        PYTHON_VERSION: ${PYTHON_VERSION:-3.10}
+        PYTORCH_VERSION: ${PYTORCH_VERSION:-2.2.2}
+        WEBUI_TAG: ${WEBUI_TAG:-}
+        IMAGE_BASE: ${IMAGE_BASE:-ghcr.io/ai-dock/python:${PYTHON_VERSION:-3.10}-cuda-11.8.0-base-22.04}
       tags:
-        - "ghcr.io/ai-dock/stable-diffusion-webui:${IMAGE_TAG:-jupyter-pytorch-2.2.1-py3.10-cuda-11.8.0-runtime-22.04}"
+        - "ghcr.io/ai-dock/stable-diffusion-webui:${IMAGE_TAG:-cuda-11.8.0-base-22.04}"
         
-    image: ghcr.io/ai-dock/stable-diffusion-webui:${IMAGE_TAG:-jupyter-pytorch-2.2.1-py3.10-cuda-11.8.0-runtime-22.04}
+    image: ghcr.io/ai-dock/stable-diffusion-webui:${IMAGE_TAG:-cuda-11.8.0-base-22.04}
     
     devices:
       - "/dev/dri:/dev/dri"