ComfyUI -- Arch -- Nvidia.sh

#!/bin/bash

echo "--- ComfyUI Installation Script ---"
echo "Targeting PyTorch with CUDA 12.8, Python 3.12."

# --- Configuration ---
COMFYUI_DIR=$(pwd)
VENV_DIR="${COMFYUI_DIR}/venv"
PYTHON_VERSION="3.12" # The desired Python version
COMFYUI_REPO="https://github.com/comfyanonymous/ComfyUI.git"
XFORMERS_REPO="https://github.com/facebookresearch/xformers.git"
COMFYUI_MANAGER_REPO="https://github.com/Comfy-Org/ComfyUI-Manager.git"

# Specific PyTorch index URL for stable CUDA 12.8
PYTORCH_CUDA_INDEX_URL="https://download.pytorch.org/whl/cu128"

# --- Helper Functions ---

# Function to check if a Python module is importable within the activated venv
check_module_import() {
    local module_name="$1"
    # Ensure venv is activated before checking
    if [ -z "${VIRTUAL_ENV}" ]; then
        echo "Error: Virtual environment not activated. Cannot check module import for ${module_name}."
        return 1
    fi
    # Use the venv's python executable
    "${VENV_DIR}/bin/python" -c "import ${module_name}" &> /dev/null
    return $?
}

# Function to check if PyTorch is installed and CUDA enabled
check_pytorch_cuda() {
    if [ -z "${VIRTUAL_ENV}" ]; then
        echo "Error: Virtual environment not activated. Cannot check PyTorch CUDA status."
        return 1
    fi
    # Use the venv's python executable
    "${VENV_DIR}/bin/python" -c "
import torch
if torch.cuda.is_available():
    print('PyTorch with CUDA is available.')
else:
    raise RuntimeError('PyTorch CUDA is not available.')
" &> /dev/null
    return $?
}

# --- Prerequisites Check ---
echo ""
echo "Checking for essential build tools..."

if ! command -v git &> /dev/null; then
    echo "Error: git is not installed. Please install it with 'sudo pacman -S git'."
    exit 1
fi

if ! command -v cmake &> /dev/null; then
    echo "Error: cmake is not installed. Please install it with 'sudo pacman -S cmake'."
    exit 1
fi

# Check and install tkinter dependencies for Python
if ! pacman -Qs tcl &> /dev/null || ! pacman -Qs tk &> /dev/null; then
    echo "Tkinter dependencies (tcl and tk) not found. Installing now..."
    sudo pacman -S --noconfirm tcl tk
    if [ $? -ne 0 ]; then
        echo "Warning: Failed to install tcl and tk. Python installation might not include tkinter."
    fi
fi


# --- pyenv and Python 3.12 Installation ---
echo ""
echo "Setting up Python 3.12 using pyenv..."
if ! command -v pyenv &> /dev/null; then
    echo "pyenv not found. Installing pyenv with pacman."
    sudo pacman -S pyenv --noconfirm
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install pyenv. Please install it manually with 'sudo pacman -S pyenv'."
        exit 1
    fi
fi

# Make sure pyenv is initialized for the current shell session
export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
if command -v pyenv &> /dev/null; then
    eval "$(pyenv init --path)"
    eval "$(pyenv init -)"
fi

if ! pyenv versions --bare | grep -q "${PYTHON_VERSION}"; then
    echo "Python ${PYTHON_VERSION} not installed via pyenv. Installing now..."
    pyenv install "${PYTHON_VERSION}"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install Python ${PYTHON_VERSION} with pyenv."
        echo "Check the pyenv log for details or try installing dependencies."
        exit 1
    fi
fi
echo "Python ${PYTHON_VERSION} is available."

# --- ComfyUI Cloning ---
echo ""
echo "ComfyUI repository setup..."
if [ ! -d "${COMFYUI_DIR}/ComfyUI" ]; then
    git clone "${COMFYUI_REPO}" "${COMFYUI_DIR}/ComfyUI"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to clone ComfyUI repository."
        exit 1
    fi
else
    echo "ComfyUI directory already exists. Pulling latest changes."
    (cd "${COMFYUI_DIR}/ComfyUI" && git pull)
fi

# --- Virtual Environment Setup ---
echo ""
echo "Setting up Python virtual environment at ${VENV_DIR}..."

# Set the local pyenv version to ensure 'pyenv which' works correctly
pyenv local "${PYTHON_VERSION}"
PYENV_PYTHON_BIN=$(pyenv which python)
if [ -z "${PYENV_PYTHON_BIN}" ]; then
    echo "Error: Could not find Python ${PYTHON_VERSION} binary via pyenv."
    echo "Please ensure pyenv is configured correctly and Python ${PYTHON_VERSION} is installed."
    exit 1
fi

if [ ! -d "${VENV_DIR}" ]; then
    echo "Creating virtual environment with ${PYENV_PYTHON_BIN}..."
    "${PYENV_PYTHON_BIN}" -m venv "${VENV_DIR}"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to create virtual environment with pyenv's Python ${PYTHON_VERSION}."
        exit 1
    fi
    echo "Virtual environment created."
else
    echo "Virtual environment already exists."
fi

# Activate the virtual environment
source "${VENV_DIR}/bin/activate"
if [ $? -ne 0 ]; then
    echo "Error: Failed to activate virtual environment."
    exit 1
fi
echo "Virtual environment activated."

# --- Python Package Installation ---
echo ""
echo "Installing/Updating core PyTorch, torchvision, and torchaudio with CUDA 12.8 support..."
if check_pytorch_cuda; then
    echo "PyTorch with CUDA is already installed. Skipping."
else
    echo "Attempting to uninstall existing PyTorch installations before clean install..."
    pip uninstall -y torch torchvision torchaudio 2>/dev/null
    pip install torch torchvision torchaudio --extra-index-url "${PYTORCH_CUDA_INDEX_URL}"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to install PyTorch with CUDA 12.8 support. Check your CUDA setup and network connection."
        deactivate
        exit 1
    fi
    echo "PyTorch, torchvision, and torchaudio installed for CUDA 12.8."
fi

echo ""
echo "Installing ComfyUI's core requirements..."
pip install -r "${COMFYUI_DIR}/ComfyUI/requirements.txt"
if [ $? -ne 0 ]; then
    echo "Error: Failed to install ComfyUI requirements. Check the log for specific package errors."
    deactivate
    exit 1
fi
echo "ComfyUI requirements installed."

# --- xFormers Installation (from source) ---
echo ""
echo "xFormers setup..."
XFORMERS_DIR="${COMFYUI_DIR}/xformers"
if [ ! -d "${XFORMERS_DIR}" ]; then
    git clone "${XFORMERS_REPO}" "${XFORMERS_DIR}"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to clone xFormers repository."
        deactivate
        exit 1
    fi
    echo "Initializing and updating xFormers submodules..."
    (cd "${XFORMERS_DIR}" && git submodule update --init --recursive)
    if [ $? -ne 0 ]; then
        echo "Error: Failed to update xFormers submodules. This is required for compilation."
        deactivate
        exit 1
    fi
else
    echo "xFormers directory already exists. Pulling latest changes and updating submodules."
    (cd "${XFORMERS_DIR}" && git pull && git submodule update --init --recursive)
fi

echo "Installing xFormers from source (this may take a while, no build isolation)..."
if check_module_import "xformers"; then
    echo "xFormers is already installed. Skipping build."
else
    (cd "${XFORMERS_DIR}" && pip install -e . --no-build-isolation)
    if [ $? -ne 0 ]; then
        echo "Error: Failed to build and install xFormers. Check build logs for details."
        echo "Ensure your CUDA Toolkit is properly installed and accessible."
        deactivate
        exit 1
    fi
    echo "xFormers installed."
fi

# --- ComfyUI-Manager Installation ---
echo ""
echo "ComfyUI-Manager setup..."
COMFYUI_MANAGER_DIR="${COMFYUI_DIR}/ComfyUI/custom_nodes/ComfyUI-Manager"

mkdir -p "${COMFYUI_DIR}/ComfyUI/custom_nodes"

if [ ! -d "${COMFYUI_MANAGER_DIR}" ]; then
    echo "Cloning ComfyUI-Manager repository..."
    git clone "${COMFYUI_MANAGER_REPO}" "${COMFYUI_MANAGER_DIR}"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to clone ComfyUI-Manager repository."
        echo "You may need to install it manually later or check network connectivity."
    else
        echo "ComfyUI-Manager cloned successfully."
        if [ -f "${COMFYUI_MANAGER_DIR}/requirements.txt" ]; then
            echo "Installing ComfyUI-Manager specific requirements (if any)..."
            pip install -r "${COMFYUI_MANAGER_DIR}/requirements.txt"
        fi
    fi
else
    echo "ComfyUI-Manager directory already exists. Pulling latest changes."
    (cd "${COMFYUI_MANAGER_DIR}" && git pull)
    if [ $? -ne 0 ]; then
        echo "Warning: Failed to pull latest changes for ComfyUI-Manager."
    else
        echo "ComfyUI-Manager updated."
        if [ -f "${COMFYUI_MANAGER_DIR}/requirements.txt" ]; then
            echo "Re-installing ComfyUI-Manager specific requirements (if any)..."
            pip install -r "${COMFYUI_MANAGER_DIR}/requirements.txt"
        fi
    fi
fi

# --- Final Message ---
echo ""
echo "--- ComfyUI installation process completed"

# --- Create Shortcuts ---
echo ""
echo "Creating desktop launchers for quick access..."

# 1. Create the desktop launcher for running ComfyUI
DESKTOP_FILE_PATH="${COMFYUI_DIR}/ComfyUI.desktop"
echo "Creating ComfyUI.desktop launcher..."
cat > "${DESKTOP_FILE_PATH}" << EOL
[Desktop Entry]
Name=ComfyUI
Comment=Launch ComfyUI in a new terminal
Exec=zsh -c "cd ${COMFYUI_DIR}/ComfyUI && source ${VENV_DIR}/bin/activate && python main.py; read -p 'Press Enter to close...'"
Terminal=true
Type=Application
Icon=comfyui
EOL
chmod +x "${DESKTOP_FILE_PATH}"

# 2. Create the desktop launcher for the virtual environment terminal
VENV_DESKTOP_FILE_PATH="${COMFYUI_DIR}/ComfyUI_Terminal.desktop"
echo "Creating ComfyUI_Terminal.desktop launcher..."
cat > "${VENV_DESKTOP_FILE_PATH}" << EOL
[Desktop Entry]
Name=ComfyUI Terminal
Comment=Open a terminal with the ComfyUI virtual environment active
Exec=zsh -c "cd ${COMFYUI_DIR}/ComfyUI && source ${VENV_DIR}/bin/activate && exec zsh"
Terminal=true
Type=Application
Icon=comfyui-terminal
EOL
chmod +x "${VENV_DESKTOP_FILE_PATH}"

# Deactivate the environment at the end of the script to not leave it active globally
deactivate