Allow giving custom --extra-index-url, defaults to CU130

Dockerfile

@@ -40,6 +40,13 @@ USER comfyui
 RUN python -m venv .venv
 ENV PATH="/comfyui/.venv/bin:$PATH"
 
+# Take PIP_EXTRA_INDEX_URL as a required build argument and bake it into the
+# image as an environment variable. This variable is then used to select
+# hardware-specific flavors of PyTorch, enabling support for CPU, AMD, and newer
+# CUDA inference.
+ARG PIP_EXTRA_INDEX_URL
+ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+
 # Install ComfyUI's Python dependencies. Although dependency keeping is also
 # performed at startup, building ComfyUI's base dependencies into the image
 # significantly speeds up each containers' first run.
@@ -47,7 +54,10 @@ ENV PATH="/comfyui/.venv/bin:$PATH"
 # Since this step takes a long time to complete, it's performed early to take
 # advantage of Docker's build cache, thereby accelerating subsequent builds.
 COPY requirements.txt manager_requirements.txt ./
-RUN pip install --no-cache-dir --disable-pip-version-check \
+RUN pip install                                \
+    --no-cache-dir                             \
+    --disable-pip-version-check                \
+    --extra-index-url "${PIP_EXTRA_INDEX_URL}" \
     -r requirements.txt
 
 # Install ComfyUI

README.md

@@ -358,8 +358,11 @@ For models compatible with Iluvatar Extension for PyTorch. Here's a step-by-step
 
 Start by installing Docker, Docker Compose, and the NVIDIA Container Toolkit on
 your host. Next, edit `compose.yaml` and update the `UID` and `GID` variables to
-match your user. Additional fields are documented in the file for further
-customization.
+match your host user's UID and GID. Also ensure that `PIP_EXTRA_INDEX_URL`
+points to the recommended non-experimental index URL for your hardware, as
+described in the beginning of the [manual installation
+section](#manual-install-windows-linux). Additional fields are documented in the
+file for further customization.
 
 Once ready, build and run the image locally:
 

compose.yaml

@@ -8,6 +8,9 @@ services:
       args:
         # Declare additional system dependencies for custom nodes
         APT_EXTRA_PACKAGES:
+        # Set the value of the --extra-index-url flag to select a specific CUDA,
+        # ROCm, or CPU version.
+        PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cu130"
 
     ports:
       - 8188:8188

entrypoint.sh

@@ -48,10 +48,11 @@ done
 # PIP_EXTRA_PACKAGES.
 echo "[entrypoint] Updating Python dependencies..."
 su -c "
-   pip install                      \\
-        --no-cache-dir              \\
-        --disable-pip-version-check \\
-        -r requirements.txt         \\
+   pip install                                   \\
+        --no-cache-dir                           \\
+        --disable-pip-version-check              \\
+        --extra-index-url '$PIP_EXTRA_INDEX_URL' \\
+        -r requirements.txt                      \\
         $(find custom_nodes -mindepth 2 -maxdepth 2 -type f -name requirements.txt -printf "-r '%p' ") \\
         $PIP_EXTRA_PACKAGES
 " comfyui \