Add files via upload

2026-03-14 16:28:07 +01:00 · 2026-03-14 16:28:07 +01:00 · edb272f696
commit edb272f696
parent 0c1a34b60f
4 changed files with 284 additions and 0 deletions
--- a/83
+++ b/83
@ -0,0 +1,83 @@
+# CUDA 13.0 for Blackwell GB10 (sm_121 / compute_121)
+# CUDA 12.8 only supports up to sm_120, but GB10 is sm_121.
+# "devel" includes nvcc so we can compile CUDA extensions like SageAttention.
+FROM nvidia/cuda:13.1.1-devel-ubuntu24.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+ARG COMFYUI_REF=master
+ARG SAGEATTN_REF=main
+
+# Base system deps
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git curl ca-certificates \
+    python3 python3-pip python3-venv python3-dev \
+    build-essential ninja-build cmake pkg-config \
+    libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 libxcb1 \
+    libtcmalloc-minimal4 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create venv (keeps python deps isolated inside container)
+ENV VENV=/opt/venv
+RUN python3 -m venv $VENV
+ENV PATH="$VENV/bin:$PATH"
+
+# Upgrade packaging tools
+RUN pip install -U pip setuptools wheel pynvml
+
+# ---- PyTorch (ARM64 + CUDA 13.0) ----
+# PyTorch cu130 wheels work with CUDA 13.0.x runtime.
+RUN pip install --index-url https://download.pytorch.org/whl/cu130 \
+    torch torchvision
+
+# ---- ComfyUI ----
+RUN git clone https://github.com/comfyanonymous/ComfyUI.git /opt/ComfyUI && \
+    cd /opt/ComfyUI && \
+    git checkout ${COMFYUI_REF} || true
+
+# ---- Patch utils.py to set tensor.to() to False if --disable-mmap enabled - Fixes double RAM/VRAM usage bug ----
+RUN python - <<'PY'
+from pathlib import Path
+
+path = Path("/opt/ComfyUI/comfy/utils.py")
+text = path.read_text()
+old = "tensor = tensor.to(device=device, copy=True)"
+new = "tensor = tensor.to(device=device, copy=False)"
+if old not in text:
+    raise SystemExit("Expected pattern not found in comfy/utils.py")
+path.write_text(text.replace(old, new, 1))
+PY
+
+RUN pip install -r /opt/ComfyUI/requirements.txt
+
+# ---- Comfy Kitchen Blackwell Optimization ----
+# Copy your local wheel into the build context
+COPY wheels/comfy_kitchen-0.2.7-cp312-abi3-linux_aarch64.whl /tmp/
+
+# Force install the local wheel
+# We use --force-reinstall to ensure it replaces any version installed by the requirements.txt
+RUN pip install --no-cache-dir --force-reinstall /tmp/comfy_kitchen-0.2.7-cp312-abi3-linux_aarch64.whl
+
+# Cleanup the wheel from the layer to save space
+RUN rm /tmp/comfy_kitchen-0.2.7-cp312-abi3-linux_aarch64.whl
+
+# ---- ComfyUI-Manager ----
+# Handled at runtime by entrypoint.sh (clones if missing in mounted volume)
+# This ensures latest version on each container start
+
+# ---- SageAttention ----
+# GB10 is compute capability 12.1 (sm_121).
+# CUDA 13.0 NVCC supports sm_121, so we compile directly for it.
+ENV TORCH_CUDA_ARCH_LIST="12.1"
+ENV CUDA_HOME=/usr/local/cuda
+
+# Build/install SageAttention from repo with sm_121 support
+RUN pip install --no-build-isolation "git+https://github.com/thu-ml/SageAttention@${SAGEATTN_REF}" || true
+
+# Expose ComfyUI
+EXPOSE 8188
+
+# Entry script handles runtime updates / flags
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
+ENTRYPOINT ["/entrypoint.sh"]
--- a/copy-wheels.sh
+++ b/copy-wheels.sh
@ -0,0 +1,42 @@
+#!/bin/bash
+# Copy only the wheels needed for ComfyUI
+# Container already includes: PyTorch, TorchVision, TorchAudio, Triton, ONNX
+
+set -euo pipefail
+
+MLWHEEL_SOURCE="${1:-$HOME/comfyui/wheels}"
+COMFYUI_DEST="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/wheels"
+
+echo "Copying custom wheels from: $MLWHEEL_SOURCE"
+echo "Destination: $COMFYUI_DEST"
+
+mkdir -p "$COMFYUI_DEST"
+
+# Only copy what the container doesn't have:
+# 1. Flash Attention 3 (sm_121 patches)
+# 2. Comfy Kitchen (NVFP4 quantization)
+
+if ls "$MLWHEEL_SOURCE"/flash_attn_3*.whl 1>/dev/null 2>&1; then
+    echo "✓ Copying Flash Attention 3..."
+    cp "$MLWHEEL_SOURCE"/flash_attn_3*.whl "$COMFYUI_DEST/"
+else
+    echo "⚠ Warning: flash_attn_3*.whl not found in $MLWHEEL_SOURCE"
+fi
+
+if ls "$MLWHEEL_SOURCE"/comfy*.whl 1>/dev/null 2>&1; then
+    echo "✓ Copying Comfy Kitchen..."
+    cp "$MLWHEEL_SOURCE"/comfy*.whl "$COMFYUI_DEST/"
+else
+    echo "⚠ Warning: comfy*.whl not found in $MLWHEEL_SOURCE"
+fi
+
+echo ""
+echo "Wheels copied successfully:"
+ls -lh "$COMFYUI_DEST"/*.whl 2>/dev/null || echo "No wheels found!"
+echo ""
+echo "NGC PyTorch container will provide:"
+echo "  - PyTorch 2.6+ (NVIDIA optimized)"
+echo "  - TorchVision, TorchAudio"
+echo "  - Triton, ONNX"
+echo "  - CUDA 13.0, cuDNN 9, NCCL"
+echo "  - TransformerEngine (FP8)"
--- a/docker-compose.yml
+++ b/docker-compose.yml
@ -0,0 +1,129 @@
+services:
+  comfyui:
+    build:
+      context: .
+      dockerfile: Dockerfile
+      args:
+        # Pin ComfyUI to a known-good commit/tag if desired
+        COMFYUI_REF: "${COMFYUI_REF:-master}"
+        # SageAttention ref (e.g., "main", "v2.2.0", or specific commit)
+        SAGEATTN_REF: "${SAGEATTN_REF:-main}"
+
+    image: comfyui:cu130
+    container_name: comfyui
+
+    # GPU enablement
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [ gpu ]
+
+    # LAN exposure
+    ports:
+      - "${COMFYUI_PORT:-8188}:8188"
+
+    environment:
+      COMFYUI_PORT: "${COMFYUI_PORT:-8188}"
+      # Optimized for Grace-Blackwell unified memory architecture
+      # Key insight: check .env.example for explanation regarding each flag
+      COMFYUI_FLAGS: "${COMFYUI_FLAGS:---listen 0.0.0.0 --port 8188 --disable-pinned-memory --use-sage-attention --force-fp16 --bf16-unet --bf16-vae --bf16-text-enc --dont-upcast-attention --disable-mmap}"
+      NVIDIA_VISIBLE_DEVICES: "all"
+      NVIDIA_DRIVER_CAPABILITIES: "compute,utility"
+
+      # Disable torch.compile/inductor - Triton doesn't support Blackwell sm_121a yet
+      TORCH_COMPILE_DISABLE: "1"
+      TORCHDYNAMO_DISABLE: "1"
+
+      # Grace-Blackwell unified memory optimizations
+      CUDA_CACHE_DISABLE: "1"
+      PYTORCH_NO_CUDA_MEMORY_CACHING: "1"
+      CUDA_DEVICE_MAX_CONNECTIONS: "1"
+      CUDA_DEVICE_MAX_COPY_CONNECTIONS: "4"
+      CUDA_MODULE_LOADING: "CUDA" # EAGER works as well, however CUDA loads models in memory a lot faster
+      CUDA_MANAGED_FORCE_DEVICE_ALLOC: "1"
+      OMP_NUM_THREADS: "20"
+      CUBLAS_WORKSPACE_CONFIG: ":0:0"
+
+    volumes:
+      # --- PERSISTENT STORAGE (Read/Write) ---
+
+      # 1. Models (Shared with host)
+      - ${COMFYUI_HOST_PATH}/models:/opt/ComfyUI/models:rw
+
+      # 2. Settings, User workflows, Manager config
+      - ${COMFYUI_DATA_PATH}/user:/opt/ComfyUI/user:rw
+
+      # 3. Custom Nodes
+      - ${COMFYUI_DATA_PATH}/custom_nodes:/opt/ComfyUI/custom_nodes:rw
+
+      # 4. Inputs/Outputs/Workflows
+      - ${COMFYUI_DATA_PATH}/output:/opt/ComfyUI/output:rw
+      - ${COMFYUI_DATA_PATH}/input:/opt/ComfyUI/input:rw
+      - ${COMFYUI_DATA_PATH}/workflows:/opt/ComfyUI/workflows:rw
+
+      # 5. Wheel Vault (Optional)
+      - ${COMFYUI_DATA_PATH}/wheels:/opt/wheels:rw
+
+    networks:
+      - comfy_net
+
+    # Health check - ComfyUI takes time to load, so generous start period
+    healthcheck:
+      test: [ "CMD", "curl", "-f", "http://localhost:8188/" ]
+      interval: 30s
+      timeout: 10s
+      start_period: 120s
+      retries: 3
+
+    restart: unless-stopped
+
+  # ComfyUIMini - Mobile-friendly UI
+  # Access at http://<host>:3000
+  comfyuimini:
+    build:
+      context: ./comfyuimini
+      dockerfile: Dockerfile
+      args:
+        COMFYUIMINI_REF: "${COMFYUIMINI_REF:-main}"
+
+    image: comfyuimini:latest
+    container_name: comfyuimini
+
+    ports:
+      - "${COMFYUIMINI_PORT:-3000}:3000"
+
+    environment:
+      # node-config override - connects to comfyui container via docker network
+      NODE_CONFIG: >-
+        {
+          "app_port": 3000,
+          "comfyui_url": "http://comfyui:8188",
+          "comfyui_ws_url": "ws://comfyui:8188",
+          "output_dir": "/shared/output",
+          "reject_unauthorised_cert": false
+        }
+
+    volumes:
+      # Share output directory with ComfyUI for gallery feature (read-only)
+      - ${COMFYUI_DATA_PATH}/output:/shared/output:ro
+      # Persist server-side workflows
+      - comfyuimini_workflows:/app/workflows
+
+    networks:
+      - comfy_net
+
+    depends_on:
+      comfyui:
+        condition: service_healthy
+
+    restart: unless-stopped
+
+networks:
+  comfy_net:
+    driver: bridge
+
+volumes:
+  comfyuimini_workflows:
--- a/entrypoint.sh
+++ b/entrypoint.sh
@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+COMFY_DIR="/opt/ComfyUI"
+PORT="${COMFYUI_PORT:-8188}"
+FLAGS="${COMFYUI_FLAGS:---listen 0.0.0.0 --port ${PORT}}"
+
+echo "[entrypoint] Python: $(python --version)"
+echo "[entrypoint] Torch:  $(python -c 'import torch; print(torch.__version__)')"
+echo "[entrypoint] CUDA:   $(python -c 'import torch; print(torch.version.cuda)')"
+echo "[entrypoint] Flags:  ${FLAGS}"
+
+# Ensure ComfyUI-Manager exists in mounted custom_nodes
+# Check for __init__.py to detect corrupted/partial installs
+if [[ ! -f "${COMFY_DIR}/custom_nodes/ComfyUI-Manager/__init__.py" ]]; then
+    echo "[entrypoint] ComfyUI-Manager missing or corrupted, cloning latest..."
+    rm -rf "${COMFY_DIR}/custom_nodes/ComfyUI-Manager" 2>/dev/null || true
+    git clone https://github.com/ltdrdata/ComfyUI-Manager.git \
+        "${COMFY_DIR}/custom_nodes/ComfyUI-Manager" || true
+fi
+
+# Install any requirements from custom nodes
+for req in "${COMFY_DIR}"/custom_nodes/*/requirements.txt; do
+    if [[ -f "$req" ]]; then
+        echo "[entrypoint] Installing deps from: $req"
+        pip install -q -r "$req" || true
+    fi
+done
+
+exec python "${COMFY_DIR}/main.py" ${FLAGS}