Motivation and background

Diffusion models generate images by gradually removing noise from a random initial distribution. DifuseIt was written to study and implement those mechanics directly in PyTorch.

The implementation is inspired by Google's Imagen architecture (Saharia et al., 2022). It includes the complete pipeline: variance scheduling, noise conditioning, a multi-scale U-Net with cross-attention, and sampling scripts

Core concepts of diffusion models

Diffusion architectures operate through two complementary processes:

1. Forward diffusion: A closed-form process that adds Gaussian noise to an image over a set of discrete timesteps from 1 to T. A linear variance schedule sets the step sizes (from beta_1 to beta_T). Because each addition of noise is Gaussian, any intermediate noisy image x_t can be computed directly from the source image x_0 using cumulative product coefficients (alphas_cumprod) without evaluating preceding steps sequentially.
2. Reverse denoising: A neural network trained to predict the noise added at a given step t. By subtracting this estimated noise iteratively, the model turns pure noise into an image consistent with an input text caption.

# Linear variance schedule used in GaussianDiffusion
scale = 1000 / timesteps
beta_start = scale * 0.0001
beta_end = scale * 0.02
betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, axis=0)

System architecture

The network relies on several structural components:

• Denoising U-Net: A convolutional backbone with downsampling and upsampling paths connected by residual skip connections. Each level includes ResNet blocks with group normalization and SiLU activation functions.
• Timestep embeddings: Timestep indices are projected into sinusoidal positional embeddings, then added to the feature maps of each ResNet block so the network recognizes where it is in the schedule.
• Text conditioning via frozen T5: Instead of training a text encoder from scratch alongside the image generator, the system uses a frozen T5 language model (t5-small). The text tokens produce sequence embeddings that are fed into the U-Net through spatial cross-attention layers.
• Dynamic thresholding: Strong classifier-free guidance can cause pixel values to saturate at extreme values. Dynamic thresholding calculates percentile bounds on activations and clips them accordingly, keeping contrast stable across higher guidance scales.
• Cascading resolution: A base diffusion model generates images at low initial resolution, which are then passed to a super-resolution diffusion module for detail refinement.

Training and execution workflow

The codebase provides scripts for both training and generation:

• main.py: A quick-start wrapper that configures a compact instance of the model, runs a sample training step on a small test dataset, and generates a sample output.
• train.py: The primary training pipeline. It supports command-line configuration for batch sizes, learning rates, epochs, gradient accumulation, and checkpoint intervals.
• inference.py: Loads saved model checkpoints from a training directory and runs reverse sampling against single text prompts or lists of captions provided in a text file.