Diffusion Models for High-Resolution Image Generation and Reconstruction Using PyTorch

Introduction

In this project, I implemented a Denoising Diffusion Probabilistic Model (DDPM) for image generation and reconstruction using PyTorch.

Diffusion models are generative models that learn to create new images by starting from pure random noise and slowly removing that noise step by step. This idea is different from models like GANs because diffusion models do not use a generator and discriminator. Instead, they learn a denoising process.

The main goal of this assignment was to understand and implement:

• Forward diffusion process
• Reverse denoising process
• Simplified U-Net architecture
• Image generation from random noise
• Image reconstruction
• Quantitative evaluation using PSNR and SSIM
• App deployment using Gradio

The full implementation was done using base PyTorch layers. I did not use any pretrained diffusion model or HuggingFace Diffusers pipeline.

Dataset Used

For this project, I used the CelebA-HQ 256 Images Only dataset from Kaggle.

I selected this dataset because it contains high-quality human face images. These images are suitable for a diffusion model because they have clear visual structure such as:

• face shape
• hair
• eyes
• skin tone
• background colors
• facial patterns

The images were resized to 128 × 128 during preprocessing. I used 128×128 instead of 256×256 because training a diffusion model from scratch is computationally expensive, and 128×128 gives a good balance between quality and training speed.

Final Model Settings

The final settings used in my implementation were:

SettingValueImage size128 × 128Batch size16Timesteps280Epochs50Learning rate0.00015Beta start0.00012Beta end0.016Dataset images used10000OptimizerAdamWLoss functionMean Squared ErrorSampling noise0.35Reconstruction step130FrameworkPyTorchAppGradio

These values were chosen to keep the model stable and trainable on Kaggle GPU.

What is Diffusion?

Diffusion is a generative process where the model learns how to remove noise from an image.

There are two main parts:

1. Forward Process
2. Reverse Process

Forward Diffusion Process

In the forward diffusion process, noise is gradually added to a clean image over many timesteps.

At the start, the image is clear. After some steps, the image becomes slightly noisy. At the final steps, the image becomes almost pure noise.

This process is fixed and does not require learning.

The formula used in the forward process is based on gradually mixing the original image with Gaussian noise.

In simple words:

Noisy image = part of original image + part of random noise

The model is trained by giving it noisy images and asking it to predict the noise that was added.

Reverse Diffusion Process

The reverse process is the main learning part of DDPM.

In reverse diffusion, the model starts from random noise and removes noise step by step until an image is generated.

The model does not directly predict the final image. Instead, it predicts the noise present in the image at each timestep.

This makes training more stable.

The reverse process works like this:

Random noise → less noisy image → clearer image → final generated image

Model Architecture

For the reverse process, I implemented a simplified U-Net model.

The U-Net takes two inputs:

• Noisy image
• Timestep value

The output of the model is:

• Predicted noise

The model architecture contains:

• convolution layers
• residual blocks
• timestep embeddings
• downsampling layers
• upsampling layers
• skip connections

The channel progression used was:

64 → 128 → 256

This follows the required simplified U-Net style.

Why U-Net?

U-Net is commonly used in diffusion models because it works well for image-to-image tasks.

In this project, the U-Net receives a noisy image and tries to estimate the noise inside it.

The downsampling part captures high-level features. The upsampling part reconstructs image details. Skip connections help preserve important image information.

This makes U-Net suitable for denoising.

Timestep Embedding

The timestep is very important in diffusion models.

The model needs to know how much noise is present in the image. An image at timestep 10 has less noise, while an image at timestep 250 has much more noise.

So I used timestep embeddings to give the model information about the current diffusion step.

The timestep embedding was passed into residual blocks so the model could adjust its denoising behavior according to the timestep.

Training Objective

The model was trained using Mean Squared Error (MSE) loss.

The training target was the actual noise added during the forward diffusion process.

The model predicted the noise, and the loss compared:

Predicted noise vs Actual noise

So the loss function was:

MSE(predicted_noise, actual_noise)

This is a standard and simple objective for DDPM training.

Optimizer and Training Techniques

I used AdamW optimizer because it is stable for deep learning models.

I also used:

• mixed precision training
• gradient clipping
• small batch size
• checkpoint saving

Mixed precision helped speed up training on GPU and reduced memory usage.

Gradient clipping helped avoid unstable updates during training.

Image Generation

After training, the model was used to generate new images from pure random noise.

The generation process starts with:

Random noise image

Then the model runs the reverse diffusion process step by step.

At each step:

1. The model predicts noise.
2. The predicted noise is removed.
3. A controlled amount of random noise is added back.
4. The image slowly becomes more structured.

In my final version, I used a sampling noise value of:

sample_noise = 0.35

This made the output less muddy and more stable compared to higher sampling noise.

Image Reconstruction

The reconstruction task was also implemented.

For reconstruction, I selected a target image, added noise to it, and then used the trained model to denoise it.

The reconstruction process was:

Target image → Noisy target image → Reconstructed image

I used:

recon_step = 130

This means the target image was not completely destroyed by noise. This helped the model reconstruct a more meaningful version of the target image.

Evaluation Metrics

I used two evaluation metrics:

1. PSNR
2. SSIM

PSNR

PSNR stands for Peak Signal-to-Noise Ratio.

It measures how close the reconstructed image is to the original image.

A higher PSNR value means the reconstructed image is closer to the target image.

SSIM

SSIM stands for Structural Similarity Index Measure.

It measures structural similarity between two images.

It checks visual structure, contrast, and brightness similarity.

SSIM is useful because it is closer to human visual judgment compared to only pixel difference.

Gradio App Deployment

I also built a simple Gradio app for the trained DDPM model.

The app starts from random noise and generates images using the trained model.

The app shows:

• generated image output
• intermediate denoising steps
• correct image size
• correct diffusion steps

The app uses the same model settings as the notebook:

Image size: 128 × 128
Diffusion steps: 280
Sampling noise: 0.35

This makes the app consistent with the training notebook.

Results

The model was able to learn basic face-like structures from the CelebA-HQ dataset.

The generated images showed:

• face-like colors
• rough facial structure
• hair-like regions
• background patterns
• denoising progression

The reconstruction results were better than pure random generation because the model started from a noisy version of a real target image.

Since the model was trained from scratch using a lightweight U-Net and limited GPU time, the generated images were still soft and slightly blurry. However, the model successfully demonstrated the main concept of DDPM: learning to remove noise step by step.

Challenges Faced

Training diffusion models from scratch is computationally expensive.

Some of the main challenges were:

• long training time
• blurry generated images
• sensitivity to noise schedule
• balancing timesteps and speed
• matching app settings with training settings
• saving and loading correct model weights

One important lesson was that the app must use the same settings as the notebook. If the model is trained at 128×128 but the app uses 64×64, the output becomes weak and incorrect.

Important Improvements Made

During experimentation, I improved the model by changing:

ChangeReasonImage size 128×128Better quality than 64×6450 epochsMore learning time280 timestepsGood balance between speed and qualitySampling noise 0.35Less muddy outputReconstruction step 130Better reconstructionDataset limit 10000Faster and stable trainingGroupNorm in U-NetMore stable than BatchNorm for diffusion

These changes helped make the results more stable and visually meaningful.

Limitations

The model still has some limitations.

The generated images are not perfect or highly sharp because:

• the model was trained from scratch
• the U-Net was simplified
• training time was limited
• no pretrained model was used
• no advanced attention blocks were added
• 128×128 resolution limits detail

Better results can be achieved by:

• training for more epochs
• using more dataset images
• adding attention blocks
• using a deeper U-Net
• trying cosine noise schedule
• training at 256×256 resolution
• using improved sampling methods