diff --git a/official-templates/better-comfyui/Dockerfile b/official-templates/better-comfyui/Dockerfile index c81ef39..21ed9c6 100644 --- a/official-templates/better-comfyui/Dockerfile +++ b/official-templates/better-comfyui/Dockerfile @@ -91,5 +91,9 @@ RUN chmod +x /start.sh RUN chmod +x /pre_start.sh COPY comfyui_extras.ipynb /comfyui_extras.ipynb +# Add the new install-flux.sh script +COPY install-flux.sh /install-flux.sh +RUN chmod +x /install-flux.sh + # CMD CMD [ "/start.sh" ] diff --git a/official-templates/better-comfyui/install-flux.sh b/official-templates/better-comfyui/install-flux.sh new file mode 100644 index 0000000..2d2c4b6 --- /dev/null +++ b/official-templates/better-comfyui/install-flux.sh @@ -0,0 +1,63 @@ +#!/bin/bash + +# Function to download file +download_file() { + local url="$1" + local destination="$2" + local token="$3" + + if [ -f "$destination" ]; then + read -p "File $destination already exists. Overwrite? (y/n): " overwrite + if [[ $overwrite != [Yy]* ]]; then + echo "Skipping download of $destination" + return + fi + fi + + if [ -n "$token" ]; then + wget --header="Authorization: Bearer $token" -O "$destination" "$url" + else + wget -O "$destination" "$url" + fi +} + +# Check for Hugging Face token +if [ -z "$HF_TOKEN" ]; then + read -p "Enter your Hugging Face token (press Enter to skip for public models): " HF_TOKEN +fi + +# Menu for model selection +echo "Select Flux AI Model variant to install:" +echo "1. Single-file FP8 version (16GB VRAM)" +echo "2. Flux Schnell version (16GB VRAM)" +echo "3. Regular FP16 full version (24GB VRAM)" +read -p "Enter your choice (1-3): " model_choice + +case $model_choice in + 1) + download_file "https://huggingface.co/Comfy-Org/flux1-dev/resolve/main/flux1-dev-fp8.safetensors" "/workspace/ComfyUI/models/unet/flux1-dev-fp8.safetensors" "$HF_TOKEN" + ;; + 2) + download_file "https://huggingface.co/Comfy-Org/flux1-schnell/resolve/main/flux1-schnell-fp8.safetensors" "/workspace/ComfyUI/models/unet/flux1-schnell-fp8.safetensors" "$HF_TOKEN" + ;; + 3) + if [ -z "$HF_TOKEN" ]; then + echo "Hugging Face token is required for this model." + exit 1 + fi + download_file "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors" "/workspace/ComfyUI/models/unet/flux1-dev.safetensors" "$HF_TOKEN" + ;; + *) + echo "Invalid choice. Exiting." + exit 1 + ;; +esac + +# Download CLIP models +download_file "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors" "/workspace/ComfyUI/models/clip/clip_l.safetensors" "$HF_TOKEN" +download_file "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors" "/workspace/ComfyUI/models/clip/t5xxl_fp8_e4m3fn.safetensors" "$HF_TOKEN" + +# Download VAE model +download_file "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors" "/workspace/ComfyUI/models/vae/ae.safetensors" "$HF_TOKEN" + +echo "Installation complete!" \ No newline at end of file diff --git a/official-templates/better-comfyui/pre_start.sh b/official-templates/better-comfyui/pre_start.sh index a065808..a73385a 100644 --- a/official-templates/better-comfyui/pre_start.sh +++ b/official-templates/better-comfyui/pre_start.sh @@ -30,9 +30,10 @@ done LOG_FILE="/workspace/comfyui.log" -# Copy the notebook to the /workspace directory -print_feedback "Copying notebook to /workspace..." +# Copy the notebook and install-flux.sh script to the /workspace directory +print_feedback "Copying notebook and install script to /workspace..." cp /comfyui_extras.ipynb /workspace/ +cp /install-flux.sh /workspace/ # Check if the NO_SYNC variable is set to true if [ "${NO_SYNC}" == "true" ]; then