#!/bin/bash

# ================= CONFIGURATION =================
# Replace with your actual Hugging Face Token
HF_TOKEN="hf_YOUR_ACTUAL_TOKEN_HERE"

BASE_DIR="/workspace/runpod-slim/ComfyUI"
REPO_ID="iamgroot1212/comfyui-setup"
MINIMAX_ID="Comfy-Org/MiniMax-H3"
MINIMAX_LORA_ID="larryvrh/MiniMax-H3-Turbo-Lora"

# ComfyUI Target Folders
DIFFUSION_DIR="$BASE_DIR/models/diffusion_models"
CHECKPOINT_DIR="$BASE_DIR/models/checkpoints"
TEXT_ENCODER_DIR="$BASE_DIR/models/text_encoders"
VAE_DIR="$BASE_DIR/models/vae"
LORA_DIR="$BASE_DIR/models/loras"

apt update && apt install rar unrar -y
apt update && apt install zip -y

# ================= SETUP =================
echo "🚀 	Starting direct download from $REPO_ID..."

# Create input/output directories
mkdir -p "$BASE_DIR/input/images" "$BASE_DIR/output/images" "$BASE_DIR/output/videos"

# ================= DOWNLOADS =================

# 1. Download ZIT models (diffusion-model)
echo "⬇️ 	Downloading ZIT models..."
hf download "$REPO_ID" --include "diffusion-model/*ZIT*.safetensors" --local-dir "$DIFFUSION_DIR" --token "$HF_TOKEN"

# 2. Download Qwen models (text-encoder)
echo "⬇️ 	Downloading Qwen text encoders..."
hf download "$REPO_ID" --include "text-encoder/*qwen*.safetensors" --local-dir "$TEXT_ENCODER_DIR" --token "$HF_TOKEN"

# 3. Download specific VAEs
echo "⬇️ 	Downloading specific VAEs..."
hf download "$REPO_ID" --include "vae/ae.safetensors" --local-dir "$VAE_DIR" --token "$HF_TOKEN"
hf download "$REPO_ID" --include "vae/ultra_flux.safetensors" --local-dir "$VAE_DIR" --token "$HF_TOKEN"

# 4. Download Qwen-Rapid-AIO-NSFW-v23 to Checkpoints
echo "⬇️ 	Downloading Qwen-Rapid-AIO-NSFW-v23..."
hf download "$REPO_ID" "diffusion-model/Qwen-Rapid-AIO-NSFW-v23.safetensors" --local-dir "$CHECKPOINT_DIR" --token "$HF_TOKEN"

# 5. Download MiniMax H3 to Models
echo "⬇️ 	Downloading MiniMax-H3 R2V & FL2V Models..."
hf download "$MINIMAX_ID" --include "diffusion_models/minimax_h3_fl2va_pruned_bf16.safetensors" --local-dir "$DIFFUSION_DIR" --token "$HF_TOKEN"
hf download "$MINIMAX_ID" --include "diffusion_models/minimax_h3_ref2va_pruned_bf16.safetensors" --local-dir "$DIFFUSION_DIR" --token "$HF_TOKEN"
hf download "$REPO_ID" --include "diffusion-model/MMH3-PinkyCherry-beta0.6-diffusion-model.safetensors" --local-dir "$DIFFUSION_DIR" --token "$HF_TOKEN"
hf download "$MINIMAX_ID" --include "text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors" --local-dir "$TEXT_ENCODER_DIR" --token "$HF_TOKEN"
hf download "$MINIMAX_ID" --include "vae/minimax_h3_video_vae_fp16.safetensors" --local-dir "$VAE_DIR" --token "$HF_TOKEN"
hf download "$MINIMAX_ID" --include "vae/minimax_h3_audio_vae_fp32.safetensors" --local-dir "$VAE_DIR" --token "$HF_TOKEN"
hf download "$MINIMAX_LORA_ID" --include "minimax_h3_turbo_4step.safetensors" --local-dir "$LORA_DIR" --token "$HF_TOKEN"

# ================= ORGANIZE & CLEANUP =================
echo "📂 	Flattening directory structure..."

# Function to flatten a directory: moves files from subdirs to root and removes subdirs
flatten_dir() {
    local target="$1"
    
    # Move all files from subdirectories (mindepth 2) to target root
    # Handles nested directories and spaces in filenames correctly
    find "$target" -mindepth 2 -type f -exec mv -t "$target" {} + 2>/dev/null
    
    # Remove all empty directories under target (deepest first)
    find "$target" -mindepth 1 -depth -type d -empty -delete 2>/dev/null
}

# Apply flattening to all target folders
flatten_dir "$DIFFUSION_DIR"
flatten_dir "$TEXT_ENCODER_DIR"
flatten_dir "$VAE_DIR"
flatten_dir "$CHECKPOINT_DIR"

# Remove residual .cache/huggingface metadata folders created by the tool
find "$BASE_DIR/models" -type d -name ".cache" -exec rm -rf {} + 2>/dev/null || true

echo "✅ 	Done! All files downloaded and organized in $BASE_DIR"