Update start.sh

start.sh

@@ -46,10 +46,28 @@ fi
 
 MODEL_ROOT="$PERSIST_ROOT/models"
 
-mkdir -p \
-  "$MODEL_ROOT/unet" \
-  "$MODEL_ROOT/text_encoders" \
-  "$MODEL_ROOT/vae"
+# ------------------------------------------------------------
+# Model configuration
+# ------------------------------------------------------------
+# Define models to download.
+# Format: "repo_id|filename|dest_subdir"
+# dest_subdir should be one of: unet, text_encoders, vae, loras, controlnet, workflows, etc.
+# 
+# You can override this list by setting the MODELS_TO_DOWNLOAD environment variable.
+# Example: export MODELS_TO_DOWNLOAD="repo1|file1|unet;repo2|file2|vae"
+
+if [ -n "${MODELS_TO_DOWNLOAD:-}" ]; then
+  IFS=';' read -r -a MODEL_LIST <<< "$MODELS_TO_DOWNLOAD"
+else
+  MODEL_LIST=(
+    "realrebelai/KREA-2_GGUFs|TURBO/Krea-2-Turbo-Q3_K_S.gguf|unet"
+    "Comfy-Org/Qwen3-VL|text_encoders/qwen3vl_4b_fp8_scaled.safetensors|text_encoders"
+    "Comfy-Org/Qwen-Image_ComfyUI|split_files/vae/qwen_image_vae.safetensors|vae"
+    "realrebelai/KREA-2_GGUFs|Rebels KREA-2-TURBO.json|workflows"
+    # เพิ่มโมเดลอื่นๆ ได้ที่นี่ เช่น:
+    # "some-repo/some-model|path/to/model.safetensors|loras"
+  )
+fi
 
 # ------------------------------------------------------------
 # Model directory links
@@ -81,9 +99,11 @@ link_model_dir() {
   fi
 }
 
-link_model_dir unet
-link_model_dir text_encoders
-link_model_dir vae
+# Dynamically link model directories based on MODEL_LIST
+UNIQUE_SUBDIRS=$(printf "%s\n" "${MODEL_LIST[@]}" | awk -F'|' '$3 != "" {print $3}' | grep -v "^workflows$" | sort -u)
+for subdir in $UNIQUE_SUBDIRS; do
+  link_model_dir "$subdir"
+done
 
 echo "============================================================"
 echo "ComfyUI Krea2 - Hugging Face CPU"
@@ -99,27 +119,6 @@ print("CPU threads:", os.environ.get("OMP_NUM_THREADS"))
 print("RAM limit is controlled by the Hugging Face Space hardware.")
 PY
 
-# ------------------------------------------------------------
-# Model configuration
-# ------------------------------------------------------------
-# Q3_K_S is selected for the ~15 GB RAM CPU Space.
-GGUF_REPO="${KREA2_GGUF_REPO:-realrebelai/KREA-2_GGUFs}"
-GGUF_FILE="${KREA2_GGUF_FILE:-TURBO/Krea-2-Turbo-Q3_K_S.gguf}"
-
-QWEN_REPO="${KREA2_QWEN_REPO:-Comfy-Org/Qwen3-VL}"
-QWEN_FILE="${KREA2_QWEN_FILE:-text_encoders/qwen3vl_4b_fp8_scaled.safetensors}"
-
-VAE_REPO="${KREA2_VAE_REPO:-Comfy-Org/Qwen-Image_ComfyUI}"
-VAE_FILE="${KREA2_VAE_FILE:-split_files/vae/qwen_image_vae.safetensors}"
-
-WORKFLOW_REPO="${KREA2_WORKFLOW_REPO:-realrebelai/KREA-2_GGUFs}"
-WORKFLOW_FILE="${KREA2_WORKFLOW_FILE:-Rebels KREA-2-TURBO.json}"
-
-GGUF_PATH="$COMFY/models/unet/$(basename "$GGUF_FILE")"
-QWEN_PATH="$COMFY/models/text_encoders/$(basename "$QWEN_FILE")"
-VAE_PATH="$COMFY/models/vae/$(basename "$VAE_FILE")"
-WORKFLOW_PATH="$COMFY/workflows/$(basename "$WORKFLOW_FILE")"
-
 # ------------------------------------------------------------
 # Hugging Face downloader
 # ------------------------------------------------------------
@@ -163,10 +162,21 @@ PY
 # Download models
 # ------------------------------------------------------------
 if [ "${KREA2_DOWNLOAD_MODELS:-1}" = "1" ]; then
-  download_hf "$GGUF_REPO" "$GGUF_FILE" "$GGUF_PATH"
-  download_hf "$QWEN_REPO" "$QWEN_FILE" "$QWEN_PATH"
-  download_hf "$VAE_REPO" "$VAE_FILE" "$VAE_PATH"
-  download_hf "$WORKFLOW_REPO" "$WORKFLOW_FILE" "$WORKFLOW_PATH"
+  for model_entry in "${MODEL_LIST[@]}"; do
+    [ -z "$model_entry" ] && continue
+    IFS='|' read -r repo file subdir <<< "$model_entry"
+    
+    if [ "$subdir" = "workflows" ]; then
+      dest_dir="$COMFY/workflows"
+    else
+      dest_dir="$COMFY/models/$subdir"
+    fi
+    
+    mkdir -p "$dest_dir"
+    dest_path="$dest_dir/$(basename "$file")"
+    
+    download_hf "$repo" "$file" "$dest_path"
+  done
 else
   echo "[INFO] KREA2_DOWNLOAD_MODELS=0; model download skipped."
 fi