Added flux installer for comfy ui

official-templates/better-comfyui/Dockerfile

@@ -4,14 +4,21 @@ FROM madiator2011/better-base:cuda12.1 as base
 ARG PYTHON_VERSION1=3.11
 ARG TORCH=torch==2.3.0+cu121
 
+# Install fpart (which includes fpsync)
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get install -y fpart && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
 # Stage 2: ComfyUI Installation
 FROM base as comfyui-install
 
-# Create virtual environment in /workspace/venv
+# Create virtual environment
-RUN python -m venv /workspace/venv
+RUN mkdir -p /workspace/runpod-venvs && python -m venv /workspace/runpod-venvs/better-comfyui
 
 # Set environment variables for the virtual environment
-ENV VIRTUAL_ENV="/workspace/venv"
+ENV VIRTUAL_ENV="/workspace/runpod-venvs/better-comfyui"
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
 # Activate the virtual environment and install packages
@@ -37,7 +44,7 @@ ARG INCLUDE_MODELS=false
 
 # Download each model in a separate layer
 RUN if [ "${INCLUDE_MODELS}" = "true" ]; then \
-    wget -q --show-progress https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors -O /comfy-models/v1-5-pruned-emaonly.safetensors; \
+    wget -q --show-progress https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly.safetensors -O /comfy-models/v1-5-pruned-emaonly.safetensors; \
     fi
 
 RUN if [ "${INCLUDE_MODELS}" = "true" ]; then \
@@ -61,13 +68,13 @@ RUN if [ "${INCLUDE_MODELS}" = "true" ]; then \
 FROM comfyui-install as final
 
 # Move virtual environment from /workspace/venv to /venv
-RUN mv /workspace/venv /venv
+RUN mv /workspace/runpod-venvs/better-comfyui /venv
 
 # Copy models if they were included
 COPY --from=model-setup /comfy-models /comfy-models
 
 # Set environment variables for runtime
-ENV VIRTUAL_ENV="/workspace/venv"
+ENV VIRTUAL_ENV="/workspace/runpod-venvs/better-comfyui"
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
 # Copy the README.md
@@ -84,5 +91,9 @@ RUN chmod +x /start.sh
 RUN chmod +x /pre_start.sh
 COPY comfyui_extras.ipynb /comfyui_extras.ipynb
 
+# Add the new install-flux.sh script
+COPY install-flux.sh /install-flux.sh
+RUN chmod +x /install-flux.sh
+
 # CMD
 CMD [ "/start.sh" ]

official-templates/better-comfyui/docker-bake.hcl

@@ -1,5 +1,5 @@
 group "default" {
-    targets = ["full-version", "light-version", "testing"]
+    targets = ["full-version", "light-version", "dev"]
 }
 
 target "base" {
@@ -21,11 +21,6 @@ target "full-version" {
     args = {
         INCLUDE_MODELS = "true"
     }
-    contexts = {
-        scripts = "../../container-template"
-        proxy = "../../container-template/proxy"
-        logo = "../../container-template"
-    }
     tags = ["madiator2011/better-comfyui:full"]
 }
 
@@ -34,23 +29,13 @@ target "light-version" {
     args = {
         INCLUDE_MODELS = "false"
     }
-    contexts = {
-        scripts = "../../container-template"
-        proxy = "../../container-template/proxy"
-        logo = "../../container-template"
-    }
     tags = ["madiator2011/better-comfyui:light"]
 }
 
-target "testing" {
+target "dev" {
     inherits = ["base"]
     args = {
         INCLUDE_MODELS = "false"
     }
-    contexts = {
+    tags = ["madiator2011/better-comfyui:dev"]
-        scripts = "../../container-template"
-        proxy = "../../container-template/proxy"
-        logo = "../../container-template"
-    }
-    tags = ["madiator2011/better-comfyui:testing"]
 }

official-templates/better-comfyui/install-flux.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+
+# Function to download file
+download_file() {
+    local url="$1"
+    local destination="$2"
+    local token="$3"
+
+    if [ -f "$destination" ]; then
+        read -p "File $destination already exists. Overwrite? (y/n): " overwrite
+        if [[ $overwrite != [Yy]* ]]; then
+            echo "Skipping download of $destination"
+            return
+        fi
+    fi
+
+    if [ -n "$token" ]; then
+        wget --header="Authorization: Bearer $token" -O "$destination" "$url"
+    else
+        wget -O "$destination" "$url"
+    fi
+}
+
+# Check for Hugging Face token
+if [ -z "$HF_TOKEN" ]; then
+    read -p "Enter your Hugging Face token (press Enter to skip for public models): " HF_TOKEN
+fi
+
+# Menu for model selection
+echo "Select Flux AI Model variant to install:"
+echo "1. Single-file FP8 version (16GB VRAM)"
+echo "2. Flux Schnell version (16GB VRAM)"
+echo "3. Regular FP16 full version (24GB VRAM)"
+read -p "Enter your choice (1-3): " model_choice
+
+case $model_choice in
+    1)
+        download_file "https://huggingface.co/Comfy-Org/flux1-dev/resolve/main/flux1-dev-fp8.safetensors" "/workspace/ComfyUI/models/unet/flux1-dev-fp8.safetensors" "$HF_TOKEN"
+        ;;
+    2)
+        download_file "https://huggingface.co/Comfy-Org/flux1-schnell/resolve/main/flux1-schnell-fp8.safetensors" "/workspace/ComfyUI/models/unet/flux1-schnell-fp8.safetensors" "$HF_TOKEN"
+        ;;
+    3)
+        if [ -z "$HF_TOKEN" ]; then
+            echo "Hugging Face token is required for this model."
+            exit 1
+        fi
+        download_file "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors" "/workspace/ComfyUI/models/unet/flux1-dev.safetensors" "$HF_TOKEN"
+        ;;
+    *)
+        echo "Invalid choice. Exiting."
+        exit 1
+        ;;
+esac
+
+# Download CLIP models
+download_file "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors" "/workspace/ComfyUI/models/clip/clip_l.safetensors" "$HF_TOKEN"
+download_file "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors" "/workspace/ComfyUI/models/clip/t5xxl_fp8_e4m3fn.safetensors" "$HF_TOKEN"
+
+# Download VAE model
+download_file "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors" "/workspace/ComfyUI/models/vae/ae.safetensors" "$HF_TOKEN"
+
+echo "Installation complete!"

official-templates/better-comfyui/pre_start.sh

@@ -1,5 +1,10 @@
 #!/bin/bash
 
+# Set a default TERM if it's not set
+if [ -z "$TERM" ]; then
+    export TERM=xterm-256color
+fi
+
 # Exit immediately if a command exits with a non-zero status
 set -e
 
@@ -25,9 +30,10 @@ done
 
 LOG_FILE="/workspace/comfyui.log"
 
-# Copy the notebook to the /workspace directory
+# Copy the notebook and install-flux.sh script to the /workspace directory
-print_feedback "Copying notebook to /workspace..."
+print_feedback "Copying notebook and install script to /workspace..."
 cp /comfyui_extras.ipynb /workspace/
+cp /install-flux.sh /workspace/
 
 # Check if the NO_SYNC variable is set to true
 if [ "${NO_SYNC}" == "true" ]; then
@@ -38,11 +44,60 @@ fi
 print_feedback "Starting ComfyUI setup..."
 
 print_feedback "Syncing virtual environment..."
-rsync_with_progress /venv/ /workspace/venv/
+VIRTUAL_ENV="/workspace/runpod-venvs/better-comfyui"
+SOURCE_VENV="/venv"
+NUM_CORES=$(nproc)
+
+if [ ! -d "$VIRTUAL_ENV" ]; then
+    clear
+    echo -e "\e[1;33m"
+    cat << "EOF"
+ _____________________________________
+|                                     |
+|  !!! ATTENTION - FIRST TIME SYNC !!!|
+|                                     |
+|  This process will take ~10 minutes |
+|_____________________________________|
+
+EOF
+    echo -e "\e[0m"
+    
+    mkdir -p "$VIRTUAL_ENV"
+    
+    # Start background process to show progress
+    (
+        while true; do
+            for s in / - \\ \|; do
+                printf "\r\033[1;31m[%s] \033[1;37mSYNC IN PROGRESS - PLEASE WAIT\033[0m" "$s"
+                sleep 0.5
+            done
+        done
+    ) &
+    PROGRESS_PID=$!
+
+    # Perform the sync
+    rsync -aHx --info=progress2 --stats --exclude='*.pyc' --exclude='__pycache__' "$SOURCE_VENV/" "$VIRTUAL_ENV/"
+
+    # Stop the progress indicator
+    kill $PROGRESS_PID
+    wait $PROGRESS_PID 2>/dev/null
+
+    clear
+    echo -e "\e[1;32m"
+    cat << "EOF"
+ _____________________________________
+|                                     |
+|        SYNC COMPLETED SUCCESS       |
+|_____________________________________|
+
+EOF
+    echo -e "\e[0m"
+else
+    print_feedback "Subsequent sync: Updating venv without overwriting existing files..."
+    rsync -aHx --info=progress2 --stats --exclude='*.pyc' --exclude='__pycache__' --ignore-existing --update "$SOURCE_VENV/" "$VIRTUAL_ENV/"
+fi
 
 print_feedback "Activating virtual environment..."
-export VIRTUAL_ENV="/workspace/venv"
-export PATH="$VIRTUAL_ENV/bin:$PATH"
 source "$VIRTUAL_ENV/bin/activate"
 
 export PYTHONUNBUFFERED=1
@@ -61,7 +116,7 @@ print_feedback "ComfyUI will be available at http://0.0.0.0:3000"
 
 # Check if CUSTOM_ARGS is set and not empty
 if [ -n "$CUSTOM_ARGS" ]; then
-    exec python main.py --listen --port 3000 $CUSTOM_ARGS 2>&1 | tee -a $LOG_FILE
+    exec python main.py --listen --port 3000 --enable-cors-header $CUSTOM_ARGS 2>&1 | tee -a $LOG_FILE
 else
-    exec python main.py --listen --port 3000 2>&1 | tee -a $LOG_FILE
+    exec python main.py --listen --port 3000 --enable-cors-header 2>&1 | tee -a $LOG_FILE
 fi