| --- |
| license: mit |
| library_name: pytorch |
| tags: |
| - pytorch |
| - vae |
| - variational-autoencoder |
| - ddpm |
| - diffusion |
| - diffusion-models |
| - image-generation |
| - celeba |
| - from-scratch |
| - computer-vision |
| datasets: |
| - jessicali9530/celeba-dataset |
| pipeline_tag: unconditional-image-generation |
| --- |
| |
| # VAE vs DDPM β CelebA (from scratch) |
|
|
| Pretrained checkpoints for a from-scratch PyTorch **Variational Autoencoder (VAE)** and **Denoising Diffusion Probabilistic Model (DDPM)**, both trained on 64Γ64 CelebA face crops as a head-to-head comparison of the two generative model families. |
|
|
| Both models are implemented from base `torch.nn` primitives only β no pretrained VAE, no `diffusers` library. |
|
|
| π **Full training code, configs, and technical write-up:** [AhmedAbdAlkreem/vae-vs-ddpm-celeba](https://github.com/AhmedAbdAlkreem/vae-vs-ddpm-celeba) |
|
|
| --- |
|
|
| ## Files in this repository |
|
|
| | File | Size | Description | |
| |---|---|---| |
| | `vae_best.pt` | 34.7 MB | Best VAE checkpoint (v1.2) | |
| | `ddpm_best.pt` | 107 MB | Best DDPM checkpoint (v1.2) β U-Net with residual blocks + self-attention (8Γ8, 16Γ16) + sinusoidal time embedding | |
|
|
| Both checkpoints correspond to **v1.2** in the source repo's version history β the current best result after a documented sampler bug fix (v1.0 β v1.1) and architectural refinements (v1.1 β v1.2). See the [CHANGELOG](https://github.com/AhmedAbdAlkreem/vae-vs-ddpm-celeba/blob/main/CHANGELOG.md) for the full iteration history. |
|
|
| --- |
|
|
| ## Results (v1.2) |
|
|
| | Model | FID β | Inception Score β | Reconstruction PSNR β | Sampling cost | |
| |---|---|---|---|---| |
| | VAE | 97.06 | 2.00 Β± 0.06 | 22.47 dB | 1 forward pass | |
| | DDPM | **26.42** | 2.61 Β± 0.17 | n/a | 1000 forward passes (full DDPM) | |
|
|
| DDPM produces sharper, more photorealistic faces at a much higher sampling cost; the VAE is near-instant to sample but characteristically blurry β an architectural trait of pixel-wise reconstruction loss, not a training deficiency. Full quantitative and qualitative comparisons (sample grids, latent-space PCA, interpolations, denoising trajectory) are in the [GitHub README](https://github.com/AhmedAbdAlkreem/vae-vs-ddpm-celeba#quantitative-comparison). |
|
|
| --- |
|
|
| ## How to use these checkpoints |
|
|
| The model architectures are custom (defined in the source repo under `src/models/`), so these `.pt` files are raw `state_dict` weights, not a `transformers`/`diffusers`-compatible model. To load them: |
|
|
| ```bash |
| git clone https://github.com/AhmedAbdAlkreem/vae-vs-ddpm-celeba |
| cd vae-vs-ddpm-celeba |
| pip install -r requirements.txt |
| ``` |
|
|
| ```python |
| import torch |
| from huggingface_hub import hf_hub_download |
| from src.models.vae import VAE |
| from src.models.unet import UNet # DDPM denoiser |
| |
| # Download checkpoints from this Hub repo |
| vae_ckpt = hf_hub_download("UseItOrLoseIt/vae-vs-ddpm-celeba", "vae_best.pt") |
| ddpm_ckpt = hf_hub_download("UseItOrLoseIt/vae-vs-ddpm-celeba", "ddpm_best.pt") |
| |
| vae = VAE() |
| vae.load_state_dict(torch.load(vae_ckpt, map_location="cpu")) |
| vae.eval() |
| |
| unet = UNet() |
| unet.load_state_dict(torch.load(ddpm_ckpt, map_location="cpu")) |
| unet.eval() |
| ``` |
|
|
| For ready-made sampling and inference scripts (VAE reconstruction, DDPM full/DDIM sampling, latent interpolation), use `sample.py` and `inference.py` from the [source repo](https://github.com/AhmedAbdAlkreem/vae-vs-ddpm-celeba) directly β they handle config loading, device placement, and output saving for you. |
|
|
| --- |
|
|
| ## Training data |
|
|
| [CelebFaces Attributes (CelebA) Dataset](https://www.kaggle.com/datasets/jessicali9530/celeba-dataset), aligned/cropped to 64Γ64. v1.2 was trained on a 40,000-image subset (VAE: 40 epochs, DDPM: 50 epochs). |
|
|
| ## Limitations |
|
|
| - The VAE has not changed since v1.0 β a perceptual (VGG) loss is implemented in the source config but not yet enabled; this is the most likely lever for reducing VAE blur further. |
| - DDIM (fast, ~50-step) sampling is implemented but not yet benchmarked against the full 1000-step DDPM sampler. |
| - Training budget (40k images, tens of epochs) is modest relative to published DDPM results on full CelebA (~200k images, hundreds of epochs) β current FID/IS numbers should be read as credible for this budget, not as a ceiling on the architecture. |
|
|
| ## License |
|
|
| MIT β see [LICENSE](https://github.com/AhmedAbdAlkreem/vae-vs-ddpm-celeba/blob/main/LICENSE) in the source repo. |
|
|