setup_comfy.sh

#!/bin/sh
set -eu

# Variables
COMFYUI_ROOT="/swarmui/comfyui"
COMFYUI_DIR="${COMFYUI_ROOT}/ComfyUI"
COMFY_PORT=8188

MODELS_DIR="/swarmui/downloader"
DOWNLOADER_ZIP_URL="https://huggingface.co/gtaayush010/wandownloader/resolve/main/SwarmUI_Model_Downloader_v95.zip?download=true"
DOWNLOADER_ENTRY="Downloader_Gradio_App.py"

LOGFILE="/tmp/cloudflared_8188.log"

log_info() {
    echo "=========================================================="
    echo "ℹ️  $1"
    echo "=========================================================="
}

log_success() {
    echo "=========================================================="
    echo "✅ $1"
    echo "=========================================================="
}

log_error() {
    echo "=========================================================="
    echo "❌ ERROR: $1"
    echo "=========================================================="
    exit 1
}

################################################################################
# SYSTEM DEPENDENCIES
################################################################################

install_system_dependencies() {
    log_info "Installing system dependencies..."

    apt update -y || log_error "Failed apt update"

    apt install -y \
        curl lsb-release git screen python3 python3-venv python3-pip wget unzip \
        || log_error "Failed to install base packages"

    # Cloudflared repo
    curl -fsSL https://pkg.cloudflare.com/cloudflare-main.gpg \
        -o /usr/share/keyrings/cloudflare-archive-keyring.gpg

    echo "deb [signed-by=/usr/share/keyrings/cloudflare-archive-keyring.gpg] \
https://pkg.cloudflare.com/cloudflared/ $(lsb_release -cs) main" \
        > /etc/apt/sources.list.d/cloudflared.list

    apt update -y
    apt install -y cloudflared || log_error "Failed to install cloudflared"

    log_success "System dependencies installed."
}

################################################################################
# MODEL DOWNLOADER
################################################################################

setup_model_downloader() {
    log_info "Setting up Model Downloader..."

    mkdir -p "$MODELS_DIR"
    cd "$MODELS_DIR"

    wget -O download.zip "$DOWNLOADER_ZIP_URL" || true

    if [ -f download.zip ]; then
        unzip -o download.zip || true
        rm -f download.zip
    fi

    screen -dmS downloader sh -c "
        set -eu
        cd '${MODELS_DIR}'
        python3 -m venv venv || true
        . venv/bin/activate || true
        pip install --upgrade pip || true
        pip install gradio==5.49.0 requests || true

        if [ -f '${MODELS_DIR}/${DOWNLOADER_ENTRY}' ]; then
            python '${MODELS_DIR}/${DOWNLOADER_ENTRY}' --share
        else
            echo 'Downloader entry ${DOWNLOADER_ENTRY} not found.'
            ls -la '${MODELS_DIR}'
        fi

        exec sh
    "

    log_success "Model Downloader launched in screen session 'downloader'."
}

################################################################################
# COMFYUI
################################################################################

setup_comfyui() {
    log_info "Setting up ComfyUI..."

    screen -dmS comfyui bash -lc "
        set -e
        mkdir -p '${COMFYUI_ROOT}'
        cd '${COMFYUI_ROOT}'

        if [ ! -d '${COMFYUI_DIR}' ]; then
            git clone --depth 1 https://github.com/comfyanonymous/ComfyUI '${COMFYUI_DIR}'
        fi

        cd '${COMFYUI_DIR}'
        git fetch --all || true
        git reset --hard HEAD || true
        git pull --force || true

        python3 -m venv venv || true
        . venv/bin/activate

        pip install --upgrade pip || true
        pip install torch==2.8.0 torchvision torchaudio \
            --index-url https://download.pytorch.org/whl/cu129 || true

        mkdir -p custom_nodes
        cd custom_nodes

        git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Manager || true
        git clone --depth 1 https://github.com/city96/ComfyUI-GGUF || true
        git clone --depth 1 https://github.com/ClownsharkBatwing/RES4LYF || true

        for d in ComfyUI-Manager ComfyUI-GGUF RES4LYF; do
            if [ -d \"\$d\" ]; then
                cd \"\$d\"
                git reset --hard || true
                git pull --force || true
                pip install -r requirements.txt || true
                cd ..
            fi
        done

        cd ..
        pip install -r requirements.txt || true

        pip uninstall -y xformers || true
        pip install https://huggingface.co/MonsterMMORPG/Wan_GGUF/resolve/main/flash_attn-2.8.2-cp310-cp310-linux_x86_64.whl || true
        pip install https://huggingface.co/MonsterMMORPG/Wan_GGUF/resolve/main/xformers-0.0.33+c159edc0.d20250906-cp39-abi3-linux_x86_64.whl || true
        pip install https://huggingface.co/MonsterMMORPG/Wan_GGUF/resolve/main/sageattention-2.2.0.post4-cp39-abi3-linux_x86_64.whl || true
        pip install https://huggingface.co/MonsterMMORPG/Wan_GGUF/resolve/main/insightface-0.7.3-cp310-cp310-linux_x86_64.whl || true

        python main.py --listen 0.0.0.0 --port ${COMFY_PORT} --use-sage-attention
    "

    log_success "ComfyUI started in 'comfyui' screen session."
}

################################################################################
# CLOUDFLARED
################################################################################

setup_cloudflared_tunnel() {
    log_info "Starting Cloudflared tunnel..."

    rm -f "$LOGFILE"

    screen -dmS tunnel8188 sh -c "
        cloudflared tunnel --no-autoupdate --url http://127.0.0.1:${COMFY_PORT} \
            2>&1 | tee '${LOGFILE}'
        exec sh
    "

    log_info "Waiting for Cloudflared public URL (up to 120s)..."

    i=0
    while [ \$i -lt 120 ]; do
        if [ -f \"$LOGFILE\" ]; then
            URL=\$(grep -m1 -o 'https://[^ ]*trycloudflare.com' \"$LOGFILE\" || true)
            if [ -n \"\$URL\" ]; then
                log_success \"Cloudflared URL: \$URL\"
                return
            fi
        fi
        i=\$((i+1))
        sleep 1
    done

    log_error "Cloudflared URL not found after 120 seconds."
}

################################################################################
# MAIN
################################################################################

main() {
    if [ "$(id -u)" -ne 0 ]; then
        log_error "Run this script as root (sudo)."
    fi

    log_info "Running full installer..."

    install_system_dependencies
    setup_model_downloader
    setup_comfyui
    setup_cloudflared_tunnel

    log_success "Installation complete!"
    echo "Running screen sessions:"
    echo "  - comfyui"
    echo "  - downloader"
    echo "  - tunnel8188"
    echo "Use: screen -ls"
}

main