#!/bin/bash

# Create the output directory if it doesn't exist
mkdir -p /checkpoints/pretrained_model
cd /checkpoints/pretrained_model

echo "Downloading VQ-Diffusion model checkpoints..."
echo "This may take a while depending on your internet connection."

# Download cc_learnable.pth
if [ -f cc_learnable.pth ]; then
    echo "cc_learnable.pth exists"
else
    echo "Downloading cc_learnable.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cc_learnable_aa
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cc_learnable_ab
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cc_learnable_ac
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cc_learnable_ad
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cc_learnable_ae
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cc_learnable_af
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cc_learnable_ag
    
    cat cc_learnable_* > cc_learnable.pth
    rm cc_learnable_*
fi

# Download coco_learnable.pth
if [ -f coco_learnable.pth ]; then
    echo "coco_learnable.pth exists"
else
    echo "Downloading coco_learnable.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/coco_learnable_aa
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/coco_learnable_ab
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/coco_learnable_ac
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/coco_learnable_ad
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/coco_learnable_ae
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/coco_learnable_af
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/coco_learnable_ag
    
    cat coco_learnable_* > coco_learnable.pth
    rm coco_learnable_*
fi

# Download cub_learnable.pth
if [ -f cub_learnable.pth ]; then
    echo "cub_learnable.pth exists"
else
    echo "Downloading cub_learnable.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cub_learnable_aa
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cub_learnable_ab
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cub_learnable_ac
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cub_learnable_ad
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cub_learnable_ae
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cub_learnable_af
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/cub_learnable_ag
    
    cat cub_learnable_* > cub_learnable.pth
    rm cub_learnable_*
fi

# Download imagenet_learnable.pth
if [ -f imagenet_learnable.pth ]; then
    echo "imagenet_learnable.pth exists"
else
    echo "Downloading imagenet_learnable.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_aa
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_ab
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_ac
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_ad
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_ae
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_af
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_ag
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_ah
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/imagenet_learnable_ai
    
    cat imagenet_learnable_* > imagenet_learnable.pth
    rm imagenet_learnable_*
fi

# Download ithq_learnable.pth
if [ -f ithq_learnable.pth ]; then
    echo "ithq_learnable.pth exists"
else
    echo "Downloading ithq_learnable.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_aa
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ab
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ac
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ad
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ae
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_af
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ag
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ah
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ai
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_aj
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_learnable_ak
    
    cat ithq_learnable_* > ithq_learnable.pth
    rm ithq_learnable_*
fi

# Download VQVAE models
mkdir -p taming_dvae
cd taming_dvae

if [ -f vqgan_ffhq_f16_1024.pth ]; then
    echo "vqgan_ffhq_f16_1024.pth exists"
else
    echo "Downloading vqgan_ffhq_f16_1024.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_aa
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_ab
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_ac
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_ad
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_ae
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_af
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_ag
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_ah
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_ai
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_ffhq_f16_1024_aj
    
    cat vqgan_ffhq_f16_1024_* > vqgan_ffhq_f16_1024.pth
    rm vqgan_ffhq_f16_1024_*
fi

if [ -f ithq_vqvae.pth ]; then
    echo "ithq_vqvae.pth exists"
else
    echo "Downloading ithq_vqvae.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ithq_vqvae.pth
fi

if [ -f taming_f8_8192_openimages_last.pth ]; then
    echo "taming_f8_8192_openimages_last.pth exists"
else
    echo "Downloading taming_f8_8192_openimages_last.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/taming_f8_8192_openimages_last.pth
fi

if [ -f vqgan_imagenet_f16_16384.pth ]; then
    echo "vqgan_imagenet_f16_16384.pth exists"
else
    echo "Downloading vqgan_imagenet_f16_16384.pth"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/vqgan_imagenet_f16_16384.pth
fi

cd ..

# Download CLIP model
if [ -f ViT-B-32.pt ]; then
    echo "ViT-B-32.pt exists"
else
    echo "Downloading ViT-B-32.pt"
    wget https://github.com/tzco/storage/releases/download/vqdiffusion/ViT-B-32.pt
fi

cd ../..

echo "✅ All model checkpoints downloaded successfully!"
echo "You can now run: python generate_images.py"
