#!/bin/bash
# Download all models for ComfyUI Z-Image Turbo on BC-250
set -euo pipefail

cd ~/ComfyUI
source venv/bin/activate

pip install -q huggingface-hub safetensors

echo ""
echo "=== Step 1: Download Z-Image Turbo GGUF (5.2 GB) ==="
mkdir -p ~/ComfyUI/models/unet
python3 -c "
from huggingface_hub import hf_hub_download
print('Downloading z_image_turbo-Q5_K_S.gguf from jayn7/Z-Image-Turbo-GGUF ...')
hf_hub_download(
    'jayn7/Z-Image-Turbo-GGUF',
    'z_image_turbo-Q5_K_S.gguf',
    local_dir='/home/fabian/ComfyUI/models/unet'
)
print('Done!')
"

echo ""
echo "=== Step 2: Download Gemma 2 2B Text Encoder shards (9.8 GB) ==="
mkdir -p ~/sd-models/text_encoders/lumina2_gemma2_2b
python3 -c "
from huggingface_hub import hf_hub_download
import os

repo = 'Alpha-VLLM/Lumina-Image-2.0'
dest = '/home/fabian/sd-models/text_encoders/lumina2_gemma2_2b'
os.makedirs(dest, exist_ok=True)

files = [
    'text_encoder/config.json',
    'text_encoder/model.safetensors.index.json',
    'text_encoder/model-00001-of-00003.safetensors',
    'text_encoder/model-00002-of-00003.safetensors',
    'text_encoder/model-00003-of-00003.safetensors',
]
for f in files:
    print(f'Downloading {f}...')
    hf_hub_download(repo, f, local_dir=dest)
print('Done!')
"

echo ""
echo "=== Step 3: Merge Text Encoder shards into single safetensors ==="
mkdir -p ~/ComfyUI/models/text_encoders
python3 << 'EOF'
import safetensors.torch
import os, json

base_dir = "/home/fabian/sd-models/text_encoders/lumina2_gemma2_2b/text_encoder"
output = "/home/fabian/ComfyUI/models/text_encoders/gemma2_2b_lumina2.safetensors"

with open(os.path.join(base_dir, "model.safetensors.index.json")) as f:
    index = json.load(f)

all_tensors = {}
shards = set(index["weight_map"].values())
print(f"Loading {len(shards)} shards with {len(index['weight_map'])} tensors...")
for shard in sorted(shards):
    path = os.path.join(base_dir, shard)
    print(f"  Loading {shard}...")
    tensors = safetensors.torch.load_file(path, device="cpu")
    all_tensors.update(tensors)

print(f"Total tensors: {len(all_tensors)}")
print(f"Saving merged file to {output}...")
safetensors.torch.save_file(all_tensors, output)
sz = os.path.getsize(output) / 1e9
print(f"Done! Size: {sz:.2f} GB")
EOF

echo ""
echo "=== Step 4: Download VAE (335 MB) ==="
mkdir -p ~/ComfyUI/models/vae
python3 -c "
from huggingface_hub import hf_hub_download
print('Downloading ae.safetensors ...')
hf_hub_download(
    'black-forest-labs/FLUX.1-schnell',
    'ae.safetensors',
    local_dir='/home/fabian/ComfyUI/models/vae'
)
print('Done!')
"

echo ""
echo "=== Step 5: Download example workflow ==="
mkdir -p ~/ComfyUI/user/default/workflows
python3 -c "
from huggingface_hub import hf_hub_download
hf_hub_download(
    'jayn7/Z-Image-Turbo-GGUF',
    'example_workflow.json',
    local_dir='/home/fabian/ComfyUI/user/default/workflows'
)
print('Workflow downloaded!')
"

echo ""
echo "=== Verification ==="
ls -lh ~/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf 2>/dev/null || echo "MISSING: GGUF model"
ls -lh ~/ComfyUI/models/text_encoders/gemma2_2b_lumina2.safetensors 2>/dev/null || echo "MISSING: Text encoder"
ls -lh ~/ComfyUI/models/vae/ae.safetensors 2>/dev/null || echo "MISSING: VAE"
echo ""
echo "=== ALL DOWNLOADS COMPLETE ==="