provision_depthmap.sh

#!/bin/bash
# CNCFiles Depthmap - PRODUCTION Provisioning Script (v4)
# For Vast.ai serverless - 9B models with CORRECT download URLs
set -e

source /venv/main/bin/activate 2>/dev/null || true
COMFYUI_DIR="${WORKSPACE}/ComfyUI"

mkdir -p "${COMFYUI_DIR}/models/diffusion_models" \
         "${COMFYUI_DIR}/models/text_encoders" \
         "${COMFYUI_DIR}/models/vae" \
         "${COMFYUI_DIR}/models/loras" \
         "${COMFYUI_DIR}/models/checkpoints" \
         "${COMFYUI_DIR}/models/upscale_models"

AUTH=""
if [ -n "$HF_TOKEN" ]; then AUTH="--header=Authorization: Bearer $HF_TOKEN"; fi

echo "=== [1/8] Flux2Klein 9B-FP8 UNET (8.8 GB) ==="
wget -q $AUTH --show-progress -O "${COMFYUI_DIR}/models/diffusion_models/flux-2-klein-9b-fp8.safetensors" \
  "https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8/resolve/main/flux-2-klein-9b-fp8.safetensors"

echo "=== [2/8] Qwen Text Encoder fp8 (8.1 GB) ==="
wget -q $AUTH --show-progress -O "${COMFYUI_DIR}/models/text_encoders/qwen_3_8b_fp8mixed.safetensors" \
  "https://huggingface.co/Comfy-Org/vae-text-encorder-for-flux-klein-9b/resolve/main/split_files/text_encoders/qwen_3_8b_fp8mixed.safetensors"

echo "=== [3/8] VAEs ==="
wget -q $AUTH --show-progress -O "${COMFYUI_DIR}/models/vae/flux2-vae.safetensors" \
  "https://huggingface.co/Comfy-Org/vae-text-encorder-for-flux-klein-9b/resolve/main/split_files/vae/flux2-vae.safetensors"
wget -q --show-progress -O "${COMFYUI_DIR}/models/vae/vae-ft-mse-840000-ema-pruned.safetensors" \
  "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors"
# Symlink for the workflow's expected name
ln -sf "vae-ft-mse-840000-ema-pruned.safetensors" "${COMFYUI_DIR}/models/vae/sd-vae-ft-mse-stabilityai.safetensors"

echo "=== [4/8] LoRA ==="
wget -q $AUTH --show-progress -O "${COMFYUI_DIR}/models/loras/20260212-030251.safetensors" \
  "https://huggingface.co/maheshkommuri/fluxz/resolve/main/20260212-030251.safetensors"
# Also create symlink for the workflow's expected name
ln -sf "20260212-030251.safetensors" "${COMFYUI_DIR}/models/loras/20260409-38084188.safetensors"

echo "=== [5/8] DepthAnything V2 (fp32, 1.3 GB) ==="
wget -q --show-progress -O "${COMFYUI_DIR}/models/checkpoints/depth_anything_v2_vitl_fp32.safetensors" \
  "https://huggingface.co/Kijai/DepthAnythingV2-safetensors/resolve/main/depth_anything_v2_vitl_fp32.safetensors"

echo "=== [6/8] Lotus Normal (fp16, 1.6 GB) ==="
wget -q --show-progress -O "${COMFYUI_DIR}/models/checkpoints/lotus-normal-g-v1-1-fp16.safetensors" \
  "https://huggingface.co/Kijai/lotus-comfyui/resolve/main/lotus-normal-g-v1-1-fp16.safetensors"
# Symlink for the workflow's expected name
ln -sf "lotus-normal-g-v1-1-fp16.safetensors" "${COMFYUI_DIR}/models/checkpoints/lotus-normal-g-v1-1.safetensors"

echo "=== [7/8] Custom Nodes ==="
cd "${COMFYUI_DIR}/custom_nodes"
for repo in \
  "https://github.com/ltdrdata/ComfyUI-Manager" \
  "https://github.com/KleinAILabs/ComfyUI-DepthAnythingV2" \
  "https://github.com/kijai/ComfyUI-Lotus" \
  "https://github.com/cubiq/ComfyUI_essentials" \
  "https://github.com/WASasquatch/was-node-suite-comfyui" \
  "https://github.com/ltdrdata/ComfyUI-Impact-Pack" \
  "https://github.com/shadowcz007/comfyui-mixlab-nodes"; do
  name=$(basename $repo)
  if [ ! -d "$name" ]; then
    git clone --quiet --depth 1 $repo 2>/dev/null || true
    if [ -f "$name/requirements.txt" ]; then
      pip install -q -r "$name/requirements.txt" 2>/dev/null || true
    fi
  fi
done
if [ -d "ComfyUI-Impact-Pack" ]; then
  cd ComfyUI-Impact-Pack && git submodule update --init --depth 1 2>/dev/null || true
  cd "${COMFYUI_DIR}/custom_nodes"
fi

echo "=== [8/8] CUDA Fixes & Patches ==="
grep -q "cudnn" "${COMFYUI_DIR}/main.py" 2>/dev/null || \
  sed -i '1iimport torch; torch.backends.cudnn.enabled = False; torch.backends.cudnn.benchmark = False; torch.backends.cudnn.deterministic = True' "${COMFYUI_DIR}/main.py"
grep -q "cudaMallocAsync" /etc/environment 2>/dev/null || \
  echo -e "\nPYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync\nCUDA_MODULE_LOADING=LAZY" >> /etc/environment

# DepthAnythingV2 tensor in-place fix
DA_NODES="${COMFYUI_DIR}/custom_nodes/ComfyUI-DepthAnythingV2/nodes.py"
if [ -f "$DA_NODES" ]; then
  grep -q "depth_out.sub_" "$DA_NODES" 2>/dev/null && \
  sed -i 's|depth_out.sub_(depth_min).div_(depth_max - depth_min)|depth_out = (depth_out - depth_min) / (depth_max - depth_min)|g' "$DA_NODES" && \
  sed -i 's|depth_out.clamp_(0, 1)|depth_out = depth_out.clamp(0, 1)|g' "$DA_NODES"
fi

echo "=== PROVISIONING COMPLETE! ==="
echo "Models:"
ls -lh "${COMFYUI_DIR}/models/diffusion_models/" 2>/dev/null
ls -lh "${COMFYUI_DIR}/models/text_encoders/" 2>/dev/null