comfyui

comfyui.sh

#!/bin/bash
# Vast.ai ComfyUI serverless provisioning script.
# Set this file's raw URL as the PROVISIONING_SCRIPT env var in the serverless endpoint config.
# Runs once at worker cold-start to download Z-Image Base + LTX-V 2.3 models.
set -euo pipefail

pip install -q "huggingface_hub[cli]" safetensors boto3

# Install ComfyUI-RWImageS3 custom node for Cloudflare R2 output uploads
# Reads R2_ENDPOINT_URL, R2_ACCESS_KEY, R2_SECRET_KEY, R2_BUCKET_NAME from env vars
COMFYUI_DIR=""
if   [ -d /opt/ComfyUI ];       then COMFYUI_DIR=/opt/ComfyUI
elif [ -d /workspace/ComfyUI ]; then COMFYUI_DIR=/workspace/ComfyUI
fi

if [ -n "$COMFYUI_DIR" ] && [ -n "${R2_ACCESS_KEY:-}" ]; then
  echo "=== Installing ComfyUI-RWImageS3 ==="
  cd "$COMFYUI_DIR/custom_nodes"
  git clone https://github.com/lakkiy/ComfyUI-RWImageS3.git
  pip install -q -r ComfyUI-RWImageS3/requirements.txt
  # Write .env so the node picks up R2 credentials at runtime
  cat > "$COMFYUI_DIR/custom_nodes/ComfyUI-RWImageS3/.env" <<EOF
ENDPOINT_URL=${R2_ENDPOINT_URL}
S3_REGION=auto
S3_ACCESS_KEY=${R2_ACCESS_KEY}
S3_SECRET_KEY=${R2_SECRET_KEY}
S3_BUCKET_NAME=comfy
EOF
  echo "=== ComfyUI-RWImageS3 installed ==="
else
  echo "=== Skipping R2 setup (R2_ACCESS_KEY not set) ==="
fi

# Detect ComfyUI models directory (varies by Vast.ai base image version)
if   [ -d /opt/ComfyUI/models ];       then MODELS=/opt/ComfyUI/models
elif [ -d /workspace/ComfyUI/models ]; then MODELS=/workspace/ComfyUI/models
else echo "ERROR: cannot find ComfyUI models dir"; exit 1; fi

mkdir -p "$MODELS/unet" "$MODELS/clip" "$MODELS/vae" "$MODELS/checkpoints"

echo "=== Downloading Z-Image Base UNET (~12 GB) ==="
huggingface-cli download Tongyi-MAI/Z-Image-Base \
  --include "z_image_base_bf16.safetensors" \
  --local-dir "$MODELS/unet" \
  --local-dir-use-symlinks False

echo "=== Downloading Z-Image Base VAE ==="
huggingface-cli download Tongyi-MAI/Z-Image-Base \
  --include "ae.safetensors" \
  --local-dir "$MODELS/vae" \
  --local-dir-use-symlinks False

echo "=== Downloading Z-Image Base text encoder shards ==="
TE_STAGE="$MODELS/clip/te_shards"
mkdir -p "$TE_STAGE"
huggingface-cli download Tongyi-MAI/Z-Image-Base \
  --include "text_encoder/*" \
  --local-dir "$TE_STAGE" \
  --local-dir-use-symlinks False

echo "=== Merging text encoder shards → z_image_base_te_bf16.safetensors ==="
python3 - <<PYEOF
import glob
from safetensors.torch import load_file, save_file
import os

stage = "$TE_STAGE"
shards = sorted(glob.glob(os.path.join(stage, 'text_encoder', '*.safetensors')))
if not shards:
    shards = sorted(glob.glob(os.path.join(stage, '*.safetensors')))

merged = {}
for path in shards:
    print(f'  Loading {os.path.basename(path)}')
    merged.update(load_file(path))

out = "$MODELS/clip/z_image_base_te_bf16.safetensors"
print(f'  Saving merged encoder → {out} ({len(merged)} tensors)')
save_file(merged, out)
PYEOF
rm -rf "$TE_STAGE"

echo "=== Downloading LTX-Video 2.3 (~19 GB) ==="
# TODO: verify exact filename at https://huggingface.co/Lightricks/LTX-Video
huggingface-cli download Lightricks/LTX-Video \
  --include "ltxv-2.3*.safetensors" \
  --local-dir "$MODELS/checkpoints" \
  --local-dir-use-symlinks False

echo "=== Provisioning complete ==="