YashNagraj75/Latent-Diffusion-Conditional

Latent Diffusion Model — Text-Conditional CelebA-HQ

A full Latent Diffusion Model (LDM) pipeline for text-conditional face image generation, trained on CelebA-HQ at 256x256 resolution. The model compresses images into a discrete latent space with VQ-VAE, then trains a text-conditioned U-Net diffusion model in that compressed space.

Model Description

This is a two-stage generative model:

1. 1.Stage 1 — VQ-VAE: Compresses 256x256 RGB images into a 4-channel discrete latent representation with a codebook of size 8192. Trained with reconstruction, perceptual (LPIPS/VGG), and adversarial (discriminator) losses.
2. 2.Stage 2 — LDM U-Net: A conditional U-Net diffusion model operating on the VQ-VAE latent space. Text conditioning is provided via frozen CLIP embeddings (512-dim) with classifier-free guidance.

VQ-VAE Architecture

LDM U-Net Architecture

Diffusion Process

Training Details

• —Dataset: CelebA-HQ, 256x256 RGB faces
• —Discriminator enabled after 15,000 steps (disc_start=15000)
• —Training tracked with Weights & Biases

Repository Contents

How to Use

import yaml, torch
from models.vqvae import VQVAE
from models.unet_cond import UNet

with open("config/celebahq.yaml") as f:
    config = yaml.safe_load(f)

vqvae = VQVAE(**config["autoencoder_params"])
vqvae.load_state_dict(torch.load("celebhq/vqvae_autoencoder_ckpt.pth"))
vqvae.eval()

# Encode an image to latent space
with torch.no_grad():
    z, _, _ = vqvae.encode(image_tensor)
    reconstruction = vqvae.decode(z)

References

• —Rombach et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models
• —van den Oord et al. (2017). Neural Discrete Representation Learning (VQ-VAE)
• —Ho & Salimans (2022). Classifier-Free Diffusion Guidance

License

MIT