assignment.md

GEN AI — Programming Assignment

Generative Models

---

Submission Policy

Both questions must be submitted together in a SINGLE zip file named:

{NAME}_{STUDENT_ID}.zip

The zip file must contain all code folders for both questions and one combined PDF report.

Do NOT include datasets, model checkpoints, or large binary files.

---

Question 1 • 25 Marks

Denoising Diffusion Probabilistic Models (DDPM)

Implement DDPM from scratch: forward/reverse process, training objective, and ControlNet conditioning.

Environment Setup

Create a conda environment named ddpm and install PyTorch:

conda create --name ddpm python=3.10
conda activate ddpm
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
pip install -r requirements.txt

Code Structure

ddpm_assignment/
├── 2d_plot_diffusion_todo/        (Task 1)
│   ├── ddpm_tutorial.ipynb        <-- Main notebook
│   ├── dataset.py                 <-- Swiss-roll, moon, gaussians
│   ├── network.py                 <-- (TODO) Noise prediction network
│   └── ddpm.py                    <-- (TODO) DDPM pipeline
│
├── task_1_controlnet/             (Task 2)
│   ├── diffusion/
│   │   ├── unets/
│   │   │   ├── unet_2d_condition.py   <-- (TODO) Integrate ControlNet into UNet
│   │   │   └── unet_2d_blocks.py      <-- Basic UNet components
│   │   ├── controlnet.py              <-- (TODO) Implement ControlNet
│   │   └── pipeline_controlnet.py     <-- Diffusion pipeline with ControlNet
│   ├── train.py                       <-- Training code
│   ├── train.sh                       <-- Hyperparameter script
│   └── inference.ipynb                <-- Inference notebook
└── requirements.txt

Background

Denoising Diffusion Probabilistic Models (DDPMs) are a class of generative models that learn to reverse a gradual noising process. The model is trained to predict the noise added to data at each step, and generates new samples by iteratively denoising from pure Gaussian noise.

A typical DDPM pipeline consists of three components:

• Forward Process: Gradually adds Gaussian noise to a data sample over T timesteps, producing a sequence x₀ → x₁ → … → x_T
• Reverse Process: A learned neural network iteratively denoises x_T back to x₀, step by step
• Training Objective: The network is trained using a simplified noise-matching loss — predicting the noise ε added at each step

---

Task 1: Simple DDPM Pipeline with Swiss-Roll

In this task, you will implement a DDPM to learn a 2D Swiss-Roll distribution. This toy experiment lets you understand each component of the diffusion pipeline before scaling to images.

After completing your implementation, train the model and evaluate it by running ddpm_tutorial.ipynb in the 2d_plot_diffusion_todo directory.

TODO

1-1: Build a Noise Prediction Network

Implement the noise prediction network in network.py. The network takes a noisy data point and a timestep embedding as input, and predicts the noise ε added at that step. It should consist of TimeLinear layers with feature dimensions:

[dim_in, dim_hids[0], ..., dim_hids[-1], dim_out]

• Every TimeLinear layer except the final output layer must be followed by a ReLU activation
• The final layer has no activation — it directly outputs the predicted noise

⬡ Hint TimeLinear is a linear layer that is conditioned on a sinusoidal timestep embedding. The timestep embedding is added to the hidden features before the activation at each layer.

1-2: Construct the Forward and Reverse Process

In ddpm.py, implement the three core functions of the DDPM pipeline:

• q_sample(x_0, t, noise): The forward process. Given a clean sample x₀ and timestep t, return the noised sample x_t using the closed-form formula:

  x_t = √ᾱ_t · x₀ + √(1 − ᾱ_t) · ε,  where ε ~ N(0, I)
  

• p_sample(x_t, t): One-step reverse transition. Use the trained network to predict ε, then compute the denoised estimate of x_{t-1}

• p_sample_loop(shape): Full reverse process. Starting from x_T ~ N(0, I), iterate p_sample() from t=T down to t=1 and return the final sample x₀

⬡ Important Use the pre-computed noise schedule (α_t, ᾱ_t, β_t) provided in the starter code. Do not redefine the schedule inside these functions.

1-3: Implement the Training Objective

In ddpm.py, implement compute_loss(). This function should:

1. Sample a random timestep t uniformly from {1, …, T} for each element in the batch
2. Sample noise ε ~ N(0, I) of the same shape as the input x₀
3. Compute the noised sample x_t using q_sample()
4. Pass x_t and t to the noise prediction network to obtain the predicted noise ε̂
5. Return the simplified noise-matching loss: L = ||ε − ε̂||²

1-4: Training and Evaluation

Once your implementation is complete, open and run ddpm_tutorial.ipynb via Jupyter Notebook. The notebook will automatically train the diffusion model and measure the Chamfer Distance (CD) between 2D particles sampled by the model and particles from the true Swiss-Roll distribution.

Include in your report:

• The training loss curve
• The Chamfer Distance (CD) value reported after running the notebook
• A visualization of the sampled 2D particles vs. the real Swiss-Roll distribution

---

Task 2: ControlNet on Fill50K Dataset

In this task, you will implement ControlNet — a method that adds spatial conditioning (e.g., edge maps) to a pretrained Stable Diffusion model by attaching trainable copied encoder blocks with zero-convolution layers.

Prerequisites: Hugging Face Setup

Before beginning, set up Hugging Face access to download the pretrained Stable Diffusion model:

• Sign into Hugging Face at https://huggingface.co
• Obtain your Access Token at https://huggingface.co/settings/tokens
• Log in from your terminal:

$ huggingface-cli login

Install the ControlNet environment:

conda install pytorch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 pytorch-cuda=12.1 -c pytorch -c nvidia
pip install -r requirements.txt

Verify your setup by generating a test image with Stable Diffusion:

import torch
from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
image = pipe("a photo of an astronaut riding a horse on mars").images[0]
image.save("test.png")

TODO

Task 0: Generate Baseline Images

Using the 5 text prompts in ./task_1_controlnet/data/test_prompts.json, generate 5 baseline images with the pretrained Stable Diffusion model (without ControlNet). These will serve as your comparison baseline in the report.

2-1: Implement Zero-Convolution

In diffusion/controlnet.py (TODO 1), implement the zero-convolution operation. A zero-convolution is a 1×1 convolution layer whose weights and biases are both initialized to zero at the start of training. This ensures that ControlNet begins training without disrupting the pretrained Stable Diffusion outputs.

⬡ Hint Use nn.Conv2d(channels, channels, kernel_size=1) and explicitly set weight.data and bias.data to zero after initialization.

2-2: Initialize ControlNet from Pretrained UNet

In diffusion/controlnet.py (TODO 2), initialize the ControlNet encoder by copying weights from the pretrained UNet encoder blocks. This transfer learning approach allows ControlNet to start from a strong pretrained feature extractor rather than training from scratch.

2-3: Apply Zero-Convolution to Residual Features

In diffusion/controlnet.py (TODO 3), apply the zero-convolution layers to the residual feature maps output by each ControlNet encoder block before they are passed to the UNet decoder. Specifically, for each block output h, compute:

h_out = ZeroConv(h)

2-4: Integrate ControlNet Outputs into UNet

In diffusion/unets/unet_2d_condition.py (TODO 4), modify the UNet decoder to add the ControlNet residual features to the corresponding UNet decoder skip connections. Each ControlNet block output is added element-wise to the matching UNet decoder input:

decoder_input = decoder_input + controlnet_residual

⬡ Important Do not apply any additional normalization to the ControlNet residuals before adding them to the UNet features. The zero-convolution already handles the initial scaling.

2-5: Train and Evaluate

Train ControlNet on the Fill50K dataset (automatically downloaded by the load_dataset() function in train.py) by running:

$ sh train.sh

Then, run inference.ipynb to generate images conditioned on 5 different edge maps from ./data/test_conditions, using the text prompts in data/test_prompts.json.

Include in your report:

• The 5 baseline images generated by Stable Diffusion (Task 0) with their text prompts
• The 5 condition inputs (edge maps), corresponding text prompts, and ControlNet-generated images
• A brief analysis of each condition: does the generated image accurately follow the edge map?

---

Question 2 • 25 Marks

Generative Adversarial Networks (GAN)

Implement a Vanilla GAN on 2D Swiss-Roll data and a DCGAN on MNIST handwritten digits.

Environment Setup

Create a conda environment named gan_assignment and install the required packages:

conda create --name gan_assignment python=3.10
conda activate gan_assignment
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
pip install -r requirements.txt

The requirements.txt includes: numpy, matplotlib, scipy, tqdm, and jupyter.

Code Structure

gan_assignment/
├── task_1_vanilla_gan/        (Task 1)
│   ├── gan_tutorial.ipynb     <-- Main notebook
│   ├── dataset.py             <-- 2D toy dataset definitions
│   ├── network.py             <-- (TODO) Generator & Discriminator
│   └── gan.py                 <-- (TODO) GAN training pipeline
│
├── task_2_dcgan/              (Task 2)
│   ├── dcgan_tutorial.ipynb   <-- Main notebook
│   ├── network.py             <-- (TODO) DCGAN architecture
│   └── dcgan.py               <-- (TODO) DCGAN training loop
└── requirements.txt

---

Task 1: Vanilla GAN on 2D Swiss-Roll Data

Implement a fully-connected GAN to learn a 2D Swiss-Roll distribution. This toy experiment gives you hands-on experience with the adversarial training loop before scaling to image generation.

TODO

1-1: Build the Generator Network

Implement the Generator class in network.py. The Generator maps a noise vector z to a 2D output point:

• Input: noise vector z of shape (batch_size, latent_dim), with latent_dim = 16 by default
• Architecture: fully-connected layers with dimensions [latent_dim, dim_hids[0], …, dim_hids[-1], 2]
• Activation: ReLU after every hidden layer (except the final output layer)
• Output: 2D point of shape (batch_size, 2) with a Tanh activation on the last layer

⬡ Hint Use nn.Sequential or nn.ModuleList to stack your layers.

1-2: Build the Discriminator Network

Implement the Discriminator class in network.py. The Discriminator takes a 2D point and outputs a real/fake probability:

• Input: a 2D point of shape (batch_size, 2)
• Architecture: fully-connected layers with dimensions [2, dim_hids[0], …, dim_hids[-1], 1]
• Activation: LeakyReLU (negative slope = 0.2) after every hidden layer
• Output: a scalar of shape (batch_size, 1) with a Sigmoid activation to produce a probability in [0, 1]

1-3: Implement the GAN Training Step

In gan.py, implement the train_step() function which performs one full update of both G and D:

1. Discriminator update:

• Sample a real batch x from the dataset

• Sample z ~ N(0, I) and generate fake samples: x_fake = G(z)

• Compute the discriminator BCE loss:

  L_D = −E[log D(x_real)] − E[log(1 − D(x_fake.detach()))]
  

• Zero grad on D optimizer, backpropagate, and update D only

2. Generator update:

• Sample a new batch of z ~ N(0, I)

• Compute the non-saturating generator loss:

  L_G = −E[log D(G(z))]
  

• Zero grad on G optimizer, backpropagate, and update G only

⬡ Important Always call .detach() on x_fake before passing it to D during the discriminator update. This stops gradients from flowing back into G during D's update step.

1-4: Implement the Sampling Function

In gan.py, implement sample(G, n_samples, latent_dim, device):

• Sample n_samples noise vectors z from N(0, I) with shape (n_samples, latent_dim)
• Pass through G to get generated 2D points
• Return as a NumPy array of shape (n_samples, 2)
• Use torch.no_grad() to disable gradient tracking during inference

1-5: Training and Evaluation

Run gan_tutorial.ipynb. The notebook trains the GAN for 5000 iterations and reports the Chamfer Distance (CD) between generated and real Swiss-Roll points.

Include in your report:

• G and D training loss curves (on the same plot or side-by-side)
• The Chamfer Distance (CD) value
• A scatter plot of generated 2D points vs. real Swiss-Roll data
• Brief analysis (2–3 sentences): did the GAN learn the distribution? Did you observe mode collapse or instability?

---

Task 2: Deep Convolutional GAN (DCGAN) on MNIST

Implement a DCGAN to generate handwritten digit images. DCGAN replaces fully-connected layers with convolutional layers, significantly improving image generation quality.

TODO

2-1: Implement the DCGAN Generator

Implement DCGenerator in task_2_dcgan/network.py using transposed convolutions to upsample from noise to a full image:

• Input: noise vector z of shape (batch_size, latent_dim, 1, 1), where latent_dim = 100
• Use ConvTranspose2d layers to upsample progressively to (1, 28, 28)
• Channel sequence: latent_dim → 256 → 128 → 64 → 1
• Apply BatchNorm2d + ReLU after every ConvTranspose2d except the last
• Apply Tanh to the final output

⬡ Tip ConvTranspose2d(kernel_size=4, stride=2, padding=1) doubles spatial resolution. Use kernel_size=4, stride=1, padding=0 for the first layer to go from 1×1 to 4×4.

2-2: Implement the DCGAN Discriminator

Implement DCDiscriminator in task_2_dcgan/network.py using strided convolutions to downsample the input image:

• Input: grayscale image of shape (batch_size, 1, 28, 28)
• Use Conv2d layers to downsample to a single scalar output
• Channel sequence: 1 → 64 → 128 → 256 → 1
• Apply BatchNorm2d + LeakyReLU (slope 0.2) after every Conv2d except the first and last
• Apply Sigmoid to the final output

⬡ Important Do NOT apply BatchNorm to the first layer of the discriminator (raw pixel input) or the last layer. This is standard DCGAN practice for training stability.

2-3: Implement the DCGAN Training Loop

In task_2_dcgan/dcgan.py, implement train_one_epoch() which iterates over the full MNIST training set for one epoch. For each mini-batch:

1. Discriminator update:

• BCE loss on real images (label = 1) → L_D_real
• BCE loss on fake images G(z) (label = 0) → L_D_fake
• L_D = L_D_real + L_D_fake → zero_grad, backward, step D optimizer

2. Generator update:

• Generate new fake images and compute: L_G = BCE(D(G(z)), 1)
• zero_grad, backward, step G optimizer

2-4: Weight Initialization

Implement weights_init() in task_2_dcgan/network.py and apply it via model.apply(weights_init):

• Conv2d and ConvTranspose2d: initialize weights ~ N(0, 0.02)
• BatchNorm2d: initialize weights ~ N(1.0, 0.02), bias = 0
• All other layer types: leave unchanged

⬡ Hint Use isinstance(m, nn.Conv2d) to check layer types. Use torch.nn.init.normal_() for weight initialization.

2-5: Training and Evaluation

Run dcgan_tutorial.ipynb. The notebook trains DCGAN on MNIST for 20 epochs, shows a 4×8 grid of generated digits per epoch, and reports the Fréchet Inception Distance (FID) score.

Include in your report:

• G and D training loss curves over all iterations
• A 4×8 grid of generated MNIST digits from your final trained model
• The FID score reported by the notebook
• Brief analysis (2–3 sentences): comment on image quality, diversity, and any observed instability

---

Combined Submission Instructions

Both questions — one zip file — one PDF report

What to Submit

You will submit everything — both Question 1 (DDPM) and Question 2 (GAN) — in a single zip file. There is no separate submission per question.

Zip File Structure

Your zip file must follow this exact folder layout:

{NAME}_{STUDENT_ID}.zip
├── ddpm_assignment/
│   ├── 2d_plot_diffusion_todo/
│   │   ├── network.py                <-- Your implementation
│   │   └── ddpm.py                   <-- Your implementation
│   └── task_1_controlnet/
│       └── diffusion/
│           ├── controlnet.py         <-- Your implementation
│           └── unets/
│               └── unet_2d_condition.py   <-- Your implementation
│
├── gan_assignment/
│   ├── task_1_vanilla_gan/
│   │   ├── network.py                <-- Your implementation
│   │   └── gan.py                    <-- Your implementation
│   └── task_2_dcgan/
│       ├── network.py                <-- Your implementation
│       └── dcgan.py                  <-- Your implementation
│
└── {NAME}_{STUDENT_ID}.pdf           <-- Combined report

Combined PDF Report

Write one single PDF report named {NAME}_{STUDENT_ID}.pdf that covers both questions. The report must not exceed 5 pages (excluding references). It should contain the following sections in order:

Section 1 — DDPM (Question 1):

• Task 1: Training loss curve, CD value, particle visualization, and 2–3 sentence analysis
• Task 2: 5 baseline SD images, 5 ControlNet results (condition + generated), and per-condition analysis

Section 2 — GAN (Question 2):

• Task 1: G and D loss curves, CD value, scatter plot of generated vs. real 2D points, and 2–3 sentence analysis
• Task 2: G and D loss curves, 4×8 generated MNIST grid, FID score, and 2–3 sentence analysis

Naming Convention

Do NOT include in your zip:

• Datasets or downloaded data folders (MNIST, Swiss-Roll, Fill50K, etc.)
• Model checkpoints (.pth, .ckpt files)
• Generated image folders
• Pretrained model weights (e.g., the Stable Diffusion checkpoint)

Item	Format
Zip file	{NAME}_{STUDENT_ID}.zip — e.g. JOHN_DOE_2024001.zip
PDF report	{NAME}_{STUDENT_ID}.pdf — e.g. JOHN_DOE_2024001.pdf

---

Academic Integrity

You may consult the following reference papers while working on this assignment:

• Ho et al. (2020). Denoising Diffusion Probabilistic Models.
• Zhang et al. (2023). Adding Conditional Control to Text-to-Image Diffusion Models (ControlNet).
• Goodfellow et al. (2014). Generative Adversarial Networks.
• Radford et al. (2015). Unsupervised Representation Learning with Deep Convolutional GANs (DCGAN).

It is strictly forbidden to copy, reformat, or directly reproduce code from online repositories or other students. All submitted code must be your own original implementation. Violations will result in a zero for the entire assignment.