#!/bin/bash
# Download SDXL models for DreamTail using Docker
set -e

echo "📥 Downloading SDXL models..."
echo "This will download ~13GB of model weights"
echo ""

# Model cache directory
MODELS_DIR="${DREAMTAIL_MODELS:-/mnt/nvme/models}"

# Create directory if it doesn't exist
mkdir -p "$MODELS_DIR"

echo "Models will be cached in: $MODELS_DIR"
echo ""
echo "Using Docker container to download models..."
echo ""

# Use L4T PyTorch container to download models
docker run --rm -it \
  -v "${MODELS_DIR}:/models" \
  dustynv/l4t-pytorch:r36.2.0-pth2.1-py3 \
  bash -c "
    pip3 install -q diffusers transformers accelerate safetensors &&
    python3 << 'PYEOF'
from diffusers import StableDiffusionXLPipeline

model_id = 'stabilityai/stable-diffusion-xl-base-1.0'
cache_dir = '/models'

print(f'Downloading {model_id}...')
print(f'Cache directory: {cache_dir}')
print('')

try:
    pipeline = StableDiffusionXLPipeline.from_pretrained(
        model_id,
        use_safetensors=True,
        cache_dir=cache_dir
    )
    print('✅ SDXL model downloaded successfully!')
except Exception as e:
    print(f'❌ Error downloading model: {e}')
    exit(1)
PYEOF
  "

echo ""
echo "✅ Model download complete!"
echo ""
echo "Models are cached in: $MODELS_DIR"
echo "You can now build and run DreamTail"