CelebA-HQ f16 Flow Matching DiT Lab

Artifact collection for the research testbed exploring:

1. f16 32-Channel VAE: ConvNeXt/ResNet autoencoder with RMS Sphere Normalization, Suffix Channel Masking, Pool-Align DINOv2, and VQGAN-style adaptive gradient-norm balancing.
2. Latent Optimal-Transport Flow Matching DiT: 176.9M parameter DiT ($D=896$, Depth 12, Heads 14) operating directly on the $16 \times 16 \times 32$ latent space with Classifier-Free Guidance (CFG).

Checkpoint Files Included

• celeba_vae_f16c32.safetensors (~12 MB):
  • Model weights for CelebAVAE (Epoch 35 checkpoint, PSNR: 30.56 dB, SSIM: 0.9099 on validation set).
  • Uses RMS Sphere Normalization ($z / |z|_{\text{RMS}}$) and 32 latent channels.
• celeba_flow_dit_177m.safetensors (~675 MB):
  • Model weights for the 176.9M parameter Flow Matching DiT (Epoch 120, 105,000 steps).
  • Modulated via adaLN-Zero with 64 semantic DINOv2 cluster conditioning.
• samples_dit_177m_cfg3.png:
  • 6-sample generation grid using Euler ODE 50 steps at CFG scale 3.0.
• reconstruction_comparison_vae_epoch35.png:
  • Visual reconstruction grid of the frozen VAE (GT vs 32ch vs 16ch ablated).

Repositories & Implementation

• VAE Training Code: GitHub - ruwwww/celeba-vae-lab
• DiT Flow Matching Code: GitHub - ruwwww/celeba-dit-lab