Update script.sh

script.sh

@@ -1,229 +1,388 @@
 #!/bin/bash
 set -uo pipefail
 
+export PIP_NO_CACHE_DIR=1
+export PIP_DISABLE_PIP_VERSION_CHECK=1
+export GIT_TERMINAL_PROMPT=0
+
+source /venv/main/bin/activate
+COMFYUI_DIR="${WORKSPACE:-/workspace}/ComfyUI"
+
+APT_PACKAGES=(
+)
+
+PIP_PACKAGES=(
+)
+
+NODES=(
+    "https://github.com/ltdrdata/ComfyUI-Impact-Pack.git"
+    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack.git"
+    "https://github.com/cubiq/ComfyUI_IPAdapter_plus.git"
+    "https://github.com/ka-puna/comfyui-yanc.git"
+)
+
+CHECKPOINT_MODELS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/checkpoints/waiNSFWIllustrious_v150.safetensors"
+)
+
+CLIP_VISION_MODELS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors"
+)
+
+IPADAPTER_MODELS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/ipadapter/ip-adapter-plus-face_sdxl_vit-h.safetensors"
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/ipadapter/ip-adapter-plus_sdxl_vit-h.safetensors"
+)
+
+SAM_MODELS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/sams/sam_vit_h_4b8939.pth"
+)
+
+ULTRALYTICS_BBOX_MODELS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/ultralytics/bbox/face_yolov9c.pt"
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/ultralytics/bbox/hand_yolov9c.pt"
+)
+
+LORA_MODELS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/loras/Better_detailed_pussy_and_anus_v3.0.safetensors"
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/loras/Breast_size_slider_NNFFS_alpha16.0_rank32_full_last.safetensors"
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/loras/Penis_Size_Slider_-_Illustrious_-_Girthy_alpha1.0_rank4_noxattn_250steps.safetensors"
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/loras/eye_size_addift.safetensors"
+)
+
+ESRGAN_MODELS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/models/upscale_models/RealESRGAN_x4plus_anime_6B.pth"
+)
+
+WORKFLOWS=(
+    "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/Single_Posture.json"
+)
+
 log_info() { echo "--> $1"; }
 log_warn() { echo "    ⚠️  $1"; }
 log_error() { echo "    ❌ $1"; }
 log_success() { echo "    ✅ $1"; }
 log_step() { echo ""; echo "=== [Step $1] $2 ==="; }
 
+if [ -f "/workspace/.provision_complete" ]; then
+    log_success "環境已配置，跳過重複執行"
+    log_info "如需重新配置，請刪除 /workspace/.provision_complete"
+    exit 0
+fi
+
+if [[ -f /.noprovisioning ]]; then
+    log_warn "檢測到 /.noprovisioning 文件，跳過配置"
+    exit 0
+fi
+
+if [ -f "/venv/main/bin/python" ]; then
+    PYTHON_BIN="/venv/main/bin/python"
+    PIP_BIN="/venv/main/bin/pip"
+else
+    PYTHON_BIN="python3"
+    PIP_BIN="pip3"
+fi
+
+provisioning_has_valid_hf_token() {
+    [[ -n "${HF_TOKEN:-}" ]] || return 1
+    local response
+    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET \
+        "https://huggingface.co/api/whoami-v2" \
+        -H "Authorization: Bearer $HF_TOKEN" \
+        -H "Content-Type: application/json")
+    [ "$response" -eq 200 ]
+}
+
+provisioning_has_valid_civitai_token() {
+    [[ -n "${CIVITAI_TOKEN:-}" ]] || return 1
+    local response
+    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET \
+        "https://civitai.com/api/v1/models?hidden=1&limit=1" \
+        -H "Authorization: Bearer $CIVITAI_TOKEN" \
+        -H "Content-Type: application/json")
+    [ "$response" -eq 200 ]
+}
+
 install_node() {
     local repo_url="$1"
     local branch_name="${2:-}"
     local repo_name
     repo_name=$(basename "$repo_url" .git)
+    local install_path="${COMFYUI_DIR}/custom_nodes/${repo_name}"
+
+    if [ -d "$install_path" ]; then
+        if [[ "${AUTO_UPDATE:-true}" != "false" ]]; then
+            log_info "更新節點: $repo_name"
+            (cd "$install_path" && git pull -q 2>&1 | grep -v "Already up to date" || true)
+        else
+            log_info "'$repo_name' 已存在，跳過"
+            return
+        fi
+    else
+        log_info "克隆節點: $repo_name"
+        if [ -n "$branch_name" ]; then
+            git clone --depth 1 --single-branch --branch "$branch_name" "$repo_url" "$install_path" -q 2>&1 || true
+        else
+            git clone --depth 1 --single-branch "$repo_url" "$install_path" -q 2>&1 || true
+        fi
+    fi
 
-    if [ -d "$repo_name" ]; then
-        log_info "'$repo_name' 已存在，跳過。"
-        return
+    if [ -f "$install_path/requirements.txt" ]; then
+        log_info "安裝 $repo_name 的依賴..."
+        sed -i -e '/^torch/d' -e '/^sam2/d' "$install_path/requirements.txt" 2>/dev/null || true
+        $PIP_BIN install -q --no-cache-dir -r "$install_path/requirements.txt" 2>&1 | grep -v "Requirement already satisfied" || log_warn "部分��賴安裝失敗"
     fi
 
-    log_info "正在處理: $repo_name"
-    if [ -n "$branch_name" ]; then
-        git clone --depth 1 --branch "$branch_name" "$repo_url" > /dev/null 2>&1 || log_error "Clone 倉庫 '$repo_name' (分支: $branch_name) 失敗。"
-    else
-        git clone --depth 1 "$repo_url" > /dev/null 2>&1 || log_error "Clone 倉庫 '$repo_name' 失敗。"
+    if [ -f "$install_path/install.py" ]; then
+        log_info "運行 $repo_name 的安裝腳本..."
+        $PYTHON_BIN "$install_path/install.py" 2>&1 || log_warn "安裝腳本執行失敗"
     fi
 }
 
 download_file() {
     local dest_path="$1"
     local url="$2"
-    
     local filename
     filename=$(basename "$dest_path")
+    
     mkdir -p "$(dirname "$dest_path")"
 
     if [ -f "$dest_path" ]; then
-        log_info "檔案 '$filename' 已存在，跳過下載。"
+        log_info "檔案 '$filename' 已存在，跳過下載"
         return 0
     fi
 
-    log_info "--> 開始下載 '$filename'..."
+    log_info "下載: $filename"
     
-    local max_retries=5
+    local max_retries=3
     local attempt=0
-    local success=false
+    local auth_header=""
+
+    if [[ "$url" =~ huggingface\.co ]] && provisioning_has_valid_hf_token; then
+        auth_header="Authorization: Bearer $HF_TOKEN"
+        log_info "使用 HuggingFace Token"
+    elif [[ "$url" =~ civitai\.com ]] && provisioning_has_valid_civitai_token; then
+        auth_header="Authorization: Bearer $CIVITAI_TOKEN"
+        log_info "使用 CivitAI Token"
+    fi
 
-    while [ "$success" = false ] && [ "$attempt" -lt "$max_retries" ]; do
+    while [ "$attempt" -lt "$max_retries" ]; do
         attempt=$((attempt + 1))
-        
-        if [ "$attempt" -gt 1 ]; then
-            log_warn "第 $attempt 次嘗試下載 '$filename'..."
-            sleep 15
-        fi
+        [ "$attempt" -gt 1 ] && sleep 10
 
         if command -v aria2c >/dev/null 2>&1; then
-            log_info "使用 aria2c (3 個連線) 進行下載..."
-            aria2c --console-log-level=error \
-                -c -x 16 -s 16 -k 1M \
-                --max-connection-per-server=16 \
-                --min-split-size=1M \
-                --max-tries=5 \
-                --retry-wait=3 \
-                --timeout=60 \
-                --check-certificate=false \
-                --allow-overwrite=true \
-                --file-allocation=none \
-                -d "$(dirname "$dest_path")" \
-                -o "$filename" \
-                "$url"
+            log_info "使用 aria2c (3 線程) 下載: $filename (嘗試 $attempt/$max_retries)"
+            local aria_opts=(
+                --console-log-level=error
+                -c -x 3 -s 3 -k 1M  # <-- **已修改為 3 個連接**
+                --max-connection-per-server=3 # <-- **已修改為 3 個連接**
+                --max-tries=3
+                --retry-wait=5
+                --timeout=180
+                --file-allocation=falloc
+                --auto-file-renaming=false
+                -d "$(dirname "$dest_path")"
+                -o "$filename"
+            )
+            
+            [[ -n "$auth_header" ]] && aria_opts+=(--header="$auth_header")
+            
+            aria2c "${aria_opts[@]}" "$url" 2>&1 && break
         else
-            log_info "使用 wget 進行下載..."
-            wget -c -O "$dest_path" --quiet --show-progress "$url"
-        fi
-
-        if [ $? -eq 0 ] && [ -f "$dest_path" ]; then
-            success=true
-            log_success "成功下載 '$filename'。"
-        elif [ "$attempt" -ge "$max_retries" ]; then
-            log_error "下載 '$filename' 失敗，已達最大重試次數。"
+            log_info "使用 wget 下載: $filename (嘗試 $attempt/$max_retries)"
+            local wget_opts=(
+                -qnc -O "$dest_path"
+                --timeout=60
+                --tries=3
+                --content-disposition
+                --show-progress
+            )
+            [[ -n "$auth_header" ]] && wget_opts+=(--header="$auth_header")
+            
+            wget "${wget_opts[@]}" "$url" 2>&1 && break
         fi
     done
 
-    if [ "$success" = false ]; then
+    if [ -f "$dest_path" ]; then
+        log_success "下載完成: $filename"
+        return 0
+    else
+        log_error "下載失敗: $filename"
         return 1
     fi
-    return 0
 }
 
-download_models_pack() {
-    log_info "[並行任務] 開始下載模型包..."
+download_to_directory() {
+    local dest_dir="$1"
+    shift
+    local urls=("$@")
     
-    local MODELS_ZIP_URL="https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/backup_models/models_pack.zip"
-    local ZIP_PATH="/workspace/ComfyUI/models_pack.zip"
-    local MAX_ATTEMPTS=2
-    local ATTEMPT=1
-    local SUCCESS=false
-
-    while [ $ATTEMPT -le $MAX_ATTEMPTS ] && [ "$SUCCESS" = false ]; do
-        log_info "[模型包] 嘗試 ${ATTEMPT}/${MAX_ATTEMPTS}..."
+    if [ ${#urls[@]} -eq 0 ]; then
+        return 0
+    fi
+    
+    mkdir -p "$dest_dir"
+    log_info "下載 ${#urls[@]} 個文件到 $dest_dir"
+    
+    local MAX_PARALLEL=3
+    for url in "${urls[@]}"; do
+        while [ $(jobs -r | wc -l) -ge $MAX_PARALLEL ]; do
+            sleep 1
+        done
+        local filename
+        filename=$(basename "$url" | sed 's/?.*//')
+        download_file "${dest_dir}/${filename}" "$url" &
+    done
+    wait
+}
 
-        download_file "$ZIP_PATH" "$MODELS_ZIP_URL"
+verify_and_retry_downloads() {
+    local dest_dir="$1"
+    shift
+    local urls=("$@")
+    
+    if [ ${#urls[@]} -eq 0 ]; then
+        return 0
+    fi
+    
+    log_info "檢查 $dest_dir 中的文件..."
+    
+    local missing_files=()
+    for url in "${urls[@]}"; do
+        local filename
+        filename=$(basename "$url" | sed 's/?.*//')
+        local dest_path="${dest_dir}/${filename}"
         
-        if [ $? -eq 0 ]; then
-            log_info "[模型包] 正在解壓縮..."
-            unzip -o -q "$ZIP_PATH" -d /workspace/ComfyUI/
-            
-            if [ $? -eq 0 ]; then
-                log_success "[模型包] 成功解壓縮。"
-                SUCCESS=true
-                rm -f "$ZIP_PATH"
-                log_info "[模型包] 已刪除壓縮檔。"
-            else
-                log_error "[模型包] 解壓縮失敗！"
-                rm -f "$ZIP_PATH"
-                ATTEMPT=$((ATTEMPT + 1))
-                [ $ATTEMPT -le $MAX_ATTEMPTS ] && sleep 5
-            fi
-        else
-            ATTEMPT=$((ATTEMPT + 1))
+        if [ ! -f "$dest_path" ]; then
+            log_warn "檔案缺失: $filename"
+            missing_files+=("$url")
         fi
     done
-
-    if [ "$SUCCESS" = false ]; then
-        log_error "[模型包] 下載失敗，已達最大重試次數。"
-        echo "MODELS_DOWNLOAD_FAILED" > /tmp/download_status
-        return 1
-    fi
     
-    echo "MODELS_DOWNLOAD_SUCCESS" > /tmp/download_status
-    return 0
+    if [ ${#missing_files[@]} -gt 0 ]; then
+        log_warn "發現 ${#missing_files[@]} 個缺失文件，重新下載..."
+        local MAX_PARALLEL=3
+        for url in "${missing_files[@]}"; do
+            while [ $(jobs -r | wc -l) -ge $MAX_PARALLEL ]; do
+                sleep 1
+            done
+            local filename
+            filename=$(basename "$url" | sed 's/?.*//')
+            download_file "${dest_dir}/${filename}" "$url" &
+        done
+        wait
+    else
+        log_success "所有文件已確認存在"
+    fi
 }
 
-log_info "Provisioning Started at $(date)"
-
-log_step "1/4" "運行 Vast.ai 基礎腳本"
-wget -q -O /tmp/default.sh https://raw.githubusercontent.com/vast-ai/base-image/refs/heads/main/derivatives/pytorch/derivatives/comfyui/provisioning_scripts/default.sh
-bash /tmp/default.sh || log_warn "基礎腳本遇到非致命錯誤，繼續執行..."
-
-log_step "2/4" "準備 ComfyUI 目錄與環境"
-cd /workspace/ComfyUI || { log_error "找不到 ComfyUI 目錄，中止執行。"; exit 1; }
+provisioning_print_header() {
+    echo ""
+    echo "##############################################"
+    echo "#                                            #"
+    echo "#          Provisioning Container            #"
+    echo "#                                            #"
+    echo "#         This will take some time           #"
+    echo "#                                            #"
+    echo "##############################################"
+    echo ""
+}
 
-log_info "檢測 Python 環境..."
-if [ -f "/venv/main/bin/python" ]; then
-    PYTHON_BIN="/venv/main/bin/python"
-    PIP_BIN="/venv/main/bin/pip"
-    log_success "使用 venv Python: $PYTHON_BIN"
-else
-    PYTHON_BIN="python3"
-    PIP_BIN="pip3"
-    log_warn "使用系統 Python: $PYTHON_BIN"
-fi
+provisioning_print_end() {
+    echo ""
+    echo "##############################################"
+    echo "#                                            #"
+    echo "#         Provisioning Complete!             #"
+    echo "#                                            #"
+    echo "#      Total time: $((END_TIME - START_TIME)) seconds"
+    echo "#                                            #"
+    echo "##############################################"
+    echo ""
+}
 
-log_step "3/4" "並行處理：節點安裝 + 模型下載"
+START_TIME=$(date +%s)
+log_info "開始時間: $(date)"
+provisioning_print_header
 
-# ============ 啟動並行下載任務 ============
-log_info "🚀 啟動模型包下載（後台任務）..."
-download_models_pack &
-DOWNLOAD_PID=$!
+cd "${COMFYUI_DIR}" || exit 1
 
-# ============ 同時進行節點安裝 ============
-log_info "🚀 開始安裝自定義節點（前台任務）..."
-cd custom_nodes
+if [ ${#APT_PACKAGES[@]} -gt 0 ]; then
+    log_step "1" "安裝系統套件"
+    if command -v sudo &>/dev/null && sudo -n true 2>/dev/null; then
+        sudo apt-get update -qq
+        sudo apt-get install -y -qq "${APT_PACKAGES[@]}"
+    else
+        log_warn "無 sudo 權限，跳過 APT 套件安裝"
+    fi
+fi
 
-install_node "https://github.com/ltdrdata/ComfyUI-Impact-Pack.git"
-install_node "https://github.com/ltdrdata/ComfyUI-Impact-Subpack.git"
-install_node "https://github.com/cubiq/ComfyUI_IPAdapter_plus.git"
-install_node "https://github.com/ka-puna/comfyui-yanc.git"
+log_step "2" "安裝節點與套件"
 
-log_info "正在強制更新 ComfyUI-Impact-Pack..."
-if [ -d "ComfyUI-Impact-Pack" ]; then
-    (cd ComfyUI-Impact-Pack && git pull)
+if [ ${#NODES[@]} -gt 0 ]; then
+    cd "${COMFYUI_DIR}/custom_nodes" || exit 1
+    log_info "並行安裝 ${#NODES[@]} 個節點..."
+    
+    MAX_PARALLEL=2
+    for node in "${NODES[@]}"; do
+        while [ $(jobs -r | wc -l) -ge $MAX_PARALLEL ]; do
+            sleep 1
+        done
+        install_node "$node" &
+    done
+    wait
+    
+    log_success "節點安裝完成"
+    cd "${COMFYUI_DIR}" || exit 1
 fi
 
-log_info "正在從 Impact Pack 的需求文件中移除 torch 和 sam2..."
-IMPACT_PACK_REQS="ComfyUI-Impact-Pack/requirements.txt"
-if [ -f "$IMPACT_PACK_REQS" ]; then
-    sed -i -e '/torch/d' -e '/sam2/d' "$IMPACT_PACK_REQS"
+if [ ${#PIP_PACKAGES[@]} -gt 0 ]; then
+    log_info "安裝額外的 Python 套件..."
+    $PIP_BIN install -q --no-cache-dir "${PIP_PACKAGES[@]}" 2>&1 | grep -v "Requirement already satisfied" || true
 fi
 
-log_info "為 ComfyUI-Impact-Pack 執行安裝程序..."
-if [ -f "ComfyUI-Impact-Pack/install.py" ]; then
-    $PYTHON_BIN "ComfyUI-Impact-Pack/install.py" || log_warn "Impact Pack 的 install.py 腳本執行失敗。"
-fi
-if [ -f "$IMPACT_PACK_REQS" ]; then
-    $PIP_BIN install --no-cache-dir -r "$IMPACT_PACK_REQS" || log_warn "Impact Pack 的依賴安裝失敗。"
-    mv "$IMPACT_PACK_REQS" "${IMPACT_PACK_REQS}.processed"
+log_step "3" "下載模型文件"
+
+download_to_directory "${COMFYUI_DIR}/models/checkpoints" "${CHECKPOINT_MODELS[@]}"
+download_to_directory "${COMFYUI_DIR}/models/clip_vision" "${CLIP_VISION_MODELS[@]}"
+download_to_directory "${COMFYUI_DIR}/models/ipadapter" "${IPADAPTER_MODELS[@]}"
+download_to_directory "${COMFYUI_DIR}/models/sams" "${SAM_MODELS[@]}"
+download_to_directory "${COMFYUI_DIR}/models/ultralytics/bbox" "${ULTRALYTICS_BBOX_MODELS[@]}"
+download_to_directory "${COMFYUI_DIR}/models/loras" "${LORA_MODELS[@]}"
+download_to_directory "${COMFYUI_DIR}/models/upscale_models" "${ESRGAN_MODELS[@]}"
+
+if [ ${#WORKFLOWS[@]} -gt 0 ]; then
+    log_info "下載工作流文件..."
+    mkdir -p "${COMFYUI_DIR}/user/default/workflows"
+    local MAX_PARALLEL=3
+    for workflow_url in "${WORKFLOWS[@]}"; do
+        while [ $(jobs -r | wc -l) -ge $MAX_PARALLEL ]; do
+            sleep 1
+        done
+        local workflow_name
+        workflow_name=$(basename "$workflow_url" | sed 's/?.*//')
+        download_file "${COMFYUI_DIR}/user/default/workflows/${workflow_name}" "$workflow_url" &
+    done
+    wait
 fi
 
-log_info "安裝其餘所有節點的 Python 依賴..."
-for req_file in */requirements.txt; do
-    if [ -f "$req_file" ]; then
-        log_info "正在安裝依賴: $req_file"
-        $PIP_BIN install --no-cache-dir -r "$req_file" || log_warn "依賴安裝失敗: $req_file"
-    fi
-done
-
-cd /workspace/ComfyUI
-
-log_info "⏳ 等待模型包下載完成..."
-wait $DOWNLOAD_PID
-DOWNLOAD_EXIT_CODE=$?
+log_step "4" "驗證下載完整性"
 
-if [ -f /tmp/download_status ]; then
-    DOWNLOAD_STATUS=$(cat /tmp/download_status)
-    if [ "$DOWNLOAD_STATUS" = "MODELS_DOWNLOAD_FAILED" ]; then
-        log_error "模型包下載失敗，請檢查網路或稍後重試。"
-        exit 1
-    fi
-fi
-
-log_success "所有並行任務已完成！"
+verify_and_retry_downloads "${COMFYUI_DIR}/models/checkpoints" "${CHECKPOINT_MODELS[@]}"
+verify_and_retry_downloads "${COMFYUI_DIR}/models/clip_vision" "${CLIP_VISION_MODELS[@]}"
+verify_and_retry_downloads "${COMFYUI_DIR}/models/ipadapter" "${IPADAPTER_MODELS[@]}"
+verify_and_retry_downloads "${COMFYUI_DIR}/models/sams" "${SAM_MODELS[@]}"
+verify_and_retry_downloads "${COMFYUI_DIR}/models/ultralytics/bbox" "${ULTRALYTICS_BBOX_MODELS[@]}"
+verify_and_retry_downloads "${COMFYUI_DIR}/models/loras" "${LORA_MODELS[@]}"
+verify_and_retry_downloads "${COMFYUI_DIR}/models/upscale_models" "${ESRGAN_MODELS[@]}"
 
-download_file "user/default/workflows/Single_Posture.json" "https://huggingface.co/JCscrew/RisuAI_asset_generator/resolve/main/Single_Posture.json?download=true"
+log_step "5" "完成配置"
 
-log_step "4/4" "清理與最終配置"
+touch "/workspace/.provision_complete"
+log_success "配置標記文件已創建"
 
-if [ -f "/workspace/ComfyUI/input/example.jpg" ]; then
-    log_info "正在刪除 example.jpg..."
-    rm -f /workspace/ComfyUI/input/example.jpg
-    log_success "已刪除 example.jpg。"
-else
-    log_info "example.jpg 不存在，無需刪除。"
-fi
+END_TIME=$(date +%s)
+log_info "結束時間: $(date)"
+provisioning_print_end
 
-log_success "環境配置完成！"
-log_info "ComfyUI 即將啟動。實例已準備就緒。"
-echo ""
-echo "📊 總耗時: $(($(date +%s) - $(date -d "$(head -1 /tmp/provision_start_time 2>/dev/null || echo '0 seconds ago')" +%s 2>/dev/null || echo 0))) 秒"
+log_success "所有配置步驟已完成！"
+exit 0