#!/bin/bash
# =================================================================
# LTX-2 Setup Script for ComfyUI - Arus Zaman Project
# Optimasi untuk RTX 5090 / Pro 6000 (Blackwell)
# =================================================================

# 1. Install dependencies
echo "Installing aria2 for high-speed downloads..."
apt-get update && apt-get install -y aria2

# 2. Define Paths (Standard ComfyUI structure)
BASE_DIR="/workspace/ComfyUI/models"
CHECKPOINT_DIR="$BASE_DIR/checkpoints"
TEXT_ENCODER_DIR="$BASE_DIR/text_encoders"
LORA_DIR="$BASE_DIR/loras"
LATENT_UPSCALE_DIR="$BASE_DIR/latent_upscale_models"

# Create directories
mkdir -p "$CHECKPOINT_DIR" "$TEXT_ENCODER_DIR" "$LORA_DIR" "$LATENT_UPSCALE_DIR"

echo "Starting parallel downloads..."

# 3. Download Models using aria2 (16 threads per file)
# LTX-2 Base Model (FP8 version for VRAM efficiency)
aria2c -x 16 -s 16 -d "$CHECKPOINT_DIR" -o "ltx-2-19b-dev-fp8.safetensors" "https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-dev-fp8.safetensors" &

# Text Encoder Gemma 3 12B (Crucial for complex prompt understanding)
aria2c -x 16 -s 16 -d "$TEXT_ENCODER_DIR" -o "gemma_3_12B_it_fp4_mixed.safetensors" "https://huggingface.co/Comfy-Org/ltx-2/resolve/main/split_files/text_encoders/gemma_3_12B_it_fp4_mixed.safetensors" &

# Spatial Upscaler X2 (Fixed path for LTX-2 specific nodes)
aria2c -x 16 -s 16 -d "$LATENT_UPSCALE_DIR" -o "ltx-2-spatial-upscaler-x2-1.0.safetensors" "https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-spatial-upscaler-x2-1.0.safetensors" &

# Distilled LoRA (For fast 8-12 step generation)
aria2c -x 16 -s 16 -d "$LORA_DIR" -o "ltx-2-19b-distilled-lora-384.safetensors" "https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-distilled-lora-384.safetensors" &

# Camera Control LoRA
aria2c -x 16 -s 16 -d "$LORA_DIR" -o "ltx-2-19b-lora-camera-control-dolly-left.safetensors" "https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left/resolve/main/ltx-2-19b-lora-camera-control-dolly-left.safetensors" &

wait

echo "================================================================="
echo "DOWNLOAD COMPLETE!"
echo "Semua model sudah berada di posisinya."
echo "Silakan jalankan ComfyUI dengan port 8189 menggunakan perintah:"
echo "python main.py --listen 0.0.0.0 --port 8189 --disable-xformers"
echo "================================================================="