set -e
cd /workspace/tinyvae
source .env
# phase 1: 512px crops, L1 + LPIPS + light PatchGAN (0.01) + EMA, from the best checkpoint so far (LPIPS 0.0339)
python finetune_tiny_vae.py --data data/div2k data/unsplash --init weights/long512b_best.safetensors \
  --disc_init runs/long512_gan001/disc.pt --steps 16000 --batch 16 --res 512 --lr 6e-5 --gan_w 0.01 --gan_start 0 \
  --ema 0.999 --eval_every 1000 --out runs/long512
# phase 2: 1024px finishing pass, keeps the trained discriminator
python finetune_tiny_vae.py --data data/div2k data/unsplash --init runs/long512/best.safetensors \
  --disc_init runs/long512/disc.pt --steps 3000 --batch 4 --res 1024 --lr 2e-5 --gan_w 0.01 --gan_start 0 \
  --ema 0.999 --eval_every 500 --out runs/long1024
cp runs/long1024/best.safetensors weights/tiny_klein3_v2.safetensors
python eval_tiny_vae.py --tiny weights/taef2.safetensors weights/tiny_klein3_ft.safetensors weights/tiny_klein3_v2.safetensors --grid out/compare_v2.jpg
python gen_compare.py --tiny weights/taef2.safetensors weights/tiny_klein3_ft.safetensors weights/tiny_klein3_v2.safetensors --out out/gen_v2 --mixed_sizes
echo LONG_DONE