setup-comfyui-1.sh

#!/bin/bash

# 1. Security Check
if [ -z "$HF_TOKEN" ]; then
    echo "ERROR: HF_TOKEN missing! Exiting to prevent corrupt files."
    exit 1
fi

if [ -z "$GH_TOKEN" ]; then
    echo "ERROR: GH_TOKEN missing! Exiting."
    exit 1
fi

# 2. Vast.ai Native ComfyUI Path
COMFY_DIR="/workspace/ComfyUI"

# 3. Define the Split-File Paths
DIFFUSION_DIR="$COMFY_DIR/models/diffusion_models"
TEXT_ENC_DIR="$COMFY_DIR/models/text_encoders"
VAE_DIR="$COMFY_DIR/models/vae"
WORKFLOW_DIR="/workspace/my_workflows"

mkdir -p "$DIFFUSION_DIR" "$TEXT_ENC_DIR" "$VAE_DIR" "$WORKFLOW_DIR"

# 4. Download FLUX.2 Klein 9B Distilled (FP8)
echo "Downloading $DIFFUSION_DIR/flux-2-klein-9b-fp8.safetensors"
wget -nc --header="Authorization: Bearer $HF_TOKEN" -O "$DIFFUSION_DIR/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

# 5. Download Qwen3 8B Text Encoder
echo "Downloading $TEXT_ENC_DIR/qwen_3_8b_fp8mixed.safetensors"
wget -nc -O "$TEXT_ENC_DIR/qwen_3_8b_fp8mixed.safetensors" https://huggingface.co/Comfy-Org/flux2-klein-9B/resolve/main/split_files/text_encoders/qwen_3_8b_fp8mixed.safetensors

# 6. Download FLUX.2 VAE
echo "Downloading $VAE_DIR/flux2-vae.safetensors"
wget -nc -O "$VAE_DIR/flux2-vae.safetensors" https://huggingface.co/Comfy-Org/flux2-klein-9B/resolve/main/split_files/vae/flux2-vae.safetensors

# 7. Download JSON Workflows
wget -nc -O "$WORKFLOW_DIR/flux2_9b_text_to_image.json" https://raw.githubusercontent.com/Comfy-Org/workflow_templates/refs/heads/main/templates/image_flux2_text_to_image_9b.json
wget -nc -O "$WORKFLOW_DIR/flux2_klein_9b_image_edit.json" https://raw.githubusercontent.com/Comfy-Org/workflow_templates/refs/heads/main/templates/image_flux2_klein_image_edit_9b_distilled.json

# # 8. Clone vidmaker repo and data repo
# VIDMAKER_DIR="/workspace/vidmaker"
#
# if [ ! -d "$VIDMAKER_DIR/.git" ]; then
#     echo "Cloning vidmaker repo..."
#     rm -rf "$VIDMAKER_DIR"
#     git clone "https://${GH_TOKEN}@github.com/asleepysamurai/vidmaker.git" "$VIDMAKER_DIR"
# else
#     echo "vidmaker repo already exists, pulling latest..."
#     git -C "$VIDMAKER_DIR" pull
# fi
#
# DATA_DIR="$VIDMAKER_DIR/data"
#
# if [ ! -d "$DATA_DIR/.git" ]; then
#     echo "Cloning data repo into vidmaker/data..."
#     rm -rf "$DATA_DIR"
#     git clone "https://${GH_TOKEN}@github.com/asleepysamurai/data-yt-kbc.git" "$DATA_DIR"
# else
#     echo "data repo already exists, pulling latest..."
#     git -C "$DATA_DIR" pull
# fi
#
# echo "Repos ready."
#
# # 9. Run the image generation script
# cd "$VIDMAKER_DIR"
# bash src/utils/generate_images/run.sh