Skip to content

SiddVanj/conditional-diffusion-faces

Repository files navigation

Conditional Diffusion Models for Face Generation

A research implementation of multiple conditional diffusion model variants for generating high-quality synthetic face images with fine-grained attribute control.

Overview

This project implements three complementary approaches to conditional image generation using diffusion models:

  1. FiLM-Based Multi-Scale Conditioning - Feature-wise Linear Modulation applied at multiple scales
  2. Cross-Attention DDPM - Spatial attention mechanisms for localized attribute control
  3. Rectified Flow - Fast ODE-based generation with block-wise diffusion for real-time inference

All models are trained on the CelebA dataset to generate 64×64 RGB face images conditioned on three binary attributes: Eyeglasses, Male, and Smiling.

Key Results

Method Parameters Loss Inference Speed Control Accuracy
Cross-Attention 45.8M 0.1465 2.0s 91%
FiLM 85M 0.2249 2.1s 87%
Rectified Flow 60M 0.4s (6.25×) 84%

Features

  • Multiple Conditioning Mechanisms: FiLM modulation, cross-attention, and velocity field integration
  • Efficient Architectures: Cross-Attention achieves 64% parameter reduction vs. baseline
  • Fast Inference: Rectified Flow enables real-time generation with 4 Euler steps
  • Ablation Studies: Complete ablations on classifier-free guidance, attention mechanisms, and block overlap
  • Mixed Precision Training: AMP support for memory efficiency
  • Exponential Moving Average: Improved model stability during training
  • Distributed Training: Multi-GPU support via torch.distributed

Project Structure

├── src/
│   ├── models/
│   │   ├── unet.py                    # Base UNet architecture
│   │   ├── blocks.py                  # ResNet blocks, attention layers
│   │   ├── enhanced_unet.py           # Enhanced variants with FiLM
│   │   └── cross_attention_unet.py    # Cross-attention architecture
│   ├── methods/
│   │   ├── base.py                    # Base diffusion method interface
│   │   ├── ddpm.py                    # DDPM implementation
│   │   ├── enhanced_ddpm.py           # FiLM-based DDPM
│   │   ├── cross_attention_ddpm.py    # Cross-attention DDPM
│   │   ├── rectified_flow.py          # Rectified Flow ODE sampler
│   │   └── rectified_flow_cross_attention.py  # Rectified Flow + Cross-Attention
│   ├── data/
│   │   └── celeba.py                  # CelebA dataset loader
│   └── utils/
│       ├── ema.py                     # Exponential Moving Average
│       └── logging_utils.py           # Logging utilities
├── configs/
│   ├── ddpm_celeba.yaml               # DDPM baseline config
│   ├── ddpm_celeba_hw4_cross_attention.yaml  # Cross-Attention config
│   └── rectified_flow_cross_attention.yaml   # Rectified Flow + CA config
├── train.py                           # Main training script
├── sample.py                          # Sampling/inference script
└── assets/                            # Generated sample images

Installation

Requirements

  • Python 3.9+
  • PyTorch 2.0+
  • CUDA 11.8+ (for GPU training)

Setup

git clone https://github.com/yourusername/conditional-diffusion-faces.git
cd conditional-diffusion-faces

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install pyyaml tqdm pillow

Usage

Training

Train a specific method using its configuration:

# Train Cross-Attention DDPM (recommended)
python train.py \
    --method cross_attention_ddpm \
    --config configs/ddpm_celeba_hw4_cross_attention.yaml

# Train Rectified Flow with Cross-Attention
python train.py \
    --method rectified_flow_cross_attention \
    --config configs/rectified_flow_cross_attention.yaml

# Train baseline DDPM
python train.py \
    --method ddpm \
    --config configs/ddpm_celeba.yaml

Multi-GPU Training

# Using torchrun (recommended)
torchrun --nproc_per_node=4 train.py \
    --method cross_attention_ddpm \
    --config configs/ddpm_celeba_hw4_cross_attention.yaml

Sampling

Generate new samples with trained models:

# Generate samples with specified attributes
python sample.py \
    --model_path checkpoints/cross_attention_50000.pt \
    --num_samples 16 \
    --attributes "1,0,1"  # Eyeglasses=yes, Male=no, Smiling=yes

Configuration

Each method has an associated YAML config file specifying:

  • Model architecture parameters (channels, layers, attention heads)
  • Training hyperparameters (learning rate, batch size, iterations)
  • Data settings (dataset path, image size, normalization)
  • Logging preferences

Example configs/ddpm_celeba_hw4_cross_attention.yaml:

model:
  type: cross_attention_unet
  in_channels: 3
  out_channels: 3
  channels: [64, 128, 256, 512]
  attention_heads: 8
  condition_dim: 128

training:
  batch_size: 64
  learning_rate: 1e-4
  num_iterations: 50000
  
data:
  dataset: celeba
  image_size: 64

Methods

1. FiLM (Feature-wise Linear Modulation)

Apply per-scale affine transformations to UNet features:

condition → embedding → [γ, β] per-scale → y = γ ⊙ x + β

Pros: Simple, stable training
Cons: Higher loss (0.2249), less precise attribute control

2. Cross-Attention

Spatial attention over condition embeddings:

Query: UNet features
Key/Value: Attribute embeddings (3 tokens)
→ Attention weights focus on relevant spatial regions

Pros: Best quality (0.1465 loss), 91% control accuracy, most efficient (45.8M params)
Cons: Requires classifier-free guidance training

3. Rectified Flow

ODE-based straight-path flow with block diffusion:

Forward path: x_t = (1-t)x_0 + t·z
Training: Match velocity field v_θ to z - x_0
Inference: Euler steps x_{t-Δt} = x_t + Δt·v_θ(t, x_t)

Pros: 6.25× speedup, real-time generation
Cons: Slight quality reduction at 4 steps (improved with 8+ steps)

Ablation Studies

Classifier-Free Guidance Impact

  • Without CFG: Loss 0.2532, Control Accuracy 62%
  • With CFG: Loss 0.1465 (-42%), Control Accuracy 89%

Attention vs. Concatenation

  • Concatenation: Loss 0.1892, Control 74%
  • Cross-Attention: Loss 0.1465, Control 89%

Block Overlap (Rectified Flow)

  • No overlap: KID 195-220 (visible grid seams)
  • With 2px overlap: KID 115-130 (smooth boundaries)

Performance Characteristics

Training

  • Cross-Attention: ~47.78 GPU-hours on V100 to convergence
  • Approximate cost: $43 (compute at ~$0.90/hr)
  • Convergence: Stable, no divergence with CFG

Inference (Single Sample)

  • Cross-Attention: 2.0 seconds
  • Rectified Flow: 0.4 seconds (6.25× faster)
  • Memory: ~2GB GPU

Dataset

This project uses the CelebA dataset:

  • 202,599 training images
  • 64×64 resolution
  • Binary attributes: Eyeglasses, Male, Smiling

Note: Dataset must be downloaded separately from CelebA

Update configs/*.yaml with your dataset path:

data:
  root_dir: /path/to/celeba/img_align_celeba

Checkpoints

Pre-trained checkpoints available at: [To be added]

# Download and extract
wget [checkpoint-url]
tar -xzf checkpoints.tar.gz
python sample.py --model_path checkpoints/cross_attention_50000.pt

Citation

If you use this code, please cite:

@misc{conditional-diffusion-faces,
  title={Conditional Diffusion Models for Fine-Grained Face Generation},
  author={Your Name},
  year={2026},
  howpublished={\url{https://github.com/yourusername/conditional-diffusion-faces}}
}

References

License

MIT License - see LICENSE file for details

Contributing

Contributions welcome! Please:

  1. Fork the repository
  2. Create a feature branch (git checkout -b feature/amazing-feature)
  3. Commit changes (git commit -m 'Add amazing feature')
  4. Push to branch (git push origin feature/amazing-feature)
  5. Open a Pull Request

Acknowledgments

  • CMU 10-799: Generative Models course for project guidance
  • PyTorch team for excellent deep learning framework
  • CelebA dataset creators

About

No description, website, or topics provided.

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages