DiT/Unet based Latent Flow Matching

Supports:

• Flow Matching
• With or without VAE
• Linear LR decay with warmup
• Logit Normal sampling
• Unet and DiT Models
• Data loss and Image loss along with regular Flow loss
• Tested on MNIST, Fashion MNIST and CelebA dataset

Unet based Latent Flow Matching

Supports:

• Flow Matching
• With or without VAE
• Linear LR decay with warmup
• Logit Normal sampling
• Unet Model
• Data loss and Image loss along with regular Flow loss

Tested on MNIST, Fashion MNIST and CelebA dataset