video-diffusion-pytorch

Implements a specialized attention mechanism that decomposes video processing into separate spatial (within-frame) and temporal (across-frame) attention operations. This factorization reduces computational complexity from O(T*H*W)² to O(T*(H*W)² + (T)²*H*W) by processing frame-level spatial dependencies independently before computing temporal relationships across the sequence, enabling efficient video-scale diffusion model training.

Implements a 3D convolutional U-Net backbone with symmetric encoder-decoder paths using ResNet blocks for skip connections. The architecture processes video tensors through progressive downsampling (reducing spatial dimensions) and upsampling (reconstructing resolution) while maintaining temporal information, with sinusoidal time embeddings injected at each block to condition the model on the diffusion noise schedule step.

Saves and loads complete model state (U-Net weights, optimizer state, training step counter) to disk as PyTorch .pt files. Enables resuming training from checkpoints and deploying trained models for inference. Checkpoints are saved at configurable intervals (e.g., every N steps) and can be loaded back into memory with automatic device placement (CPU/GPU).

Allows users to define the noise schedule (how much noise is added at each diffusion step) through configurable parameters like num_timesteps, beta_start, and beta_end. The schedule determines the variance of added noise at each step, controlling the trade-off between training stability and generation quality. Common schedules include linear and cosine variance schedules, which affect how quickly the model transitions from clean data to pure noise.

Implements the complete diffusion pipeline with a forward process (training) that progressively adds Gaussian noise to videos according to a noise schedule, and a reverse process (generation) that iteratively denoises from pure noise. The forward process learns to predict added noise at each step, while the reverse process uses the trained model to sample coherent videos by starting from random noise and applying learned denoising steps with optional classifier-free guidance scaling.

Encodes text descriptions through a pre-trained BERT model to create semantic embeddings that condition the video diffusion process. Implements classifier-free guidance by training the model to handle both conditioned (with text embeddings) and unconditional (with null embeddings) inputs, allowing control over guidance strength via a cond_scale parameter that interpolates between unconditional and fully-conditioned predictions during sampling.

Provides a PyTorch Dataset class that loads video data from GIF files in a specified directory, converts them to normalized tensors with shape (channels, frames, height, width), and applies optional augmentations including resizing, horizontal flipping, and pixel normalization. Handles variable-length GIFs by extracting all frames and supports batch loading through standard PyTorch DataLoader integration.

Provides a Trainer class that orchestrates the complete training loop: iterates over batches, computes diffusion loss (L2 distance between predicted and actual noise), performs backpropagation, updates model weights, and saves checkpoints at regular intervals. Handles device placement (CPU/GPU), gradient accumulation, and learning rate scheduling while logging training metrics for monitoring convergence.

Generates videos by starting with random Gaussian noise and iteratively applying the trained denoising model across a predefined number of diffusion steps (typically 100-1000). Each step reduces noise by a small amount, progressively revealing coherent video structure. The process is deterministic given a seed but produces diverse outputs across different random initializations, enabling sampling of the learned video distribution without any text or conditioning input.

Generates videos conditioned on text descriptions by combining unconditional and text-conditioned denoising predictions during the reverse diffusion process. Uses classifier-free guidance with a cond_scale parameter (typically 1.0-15.0) that interpolates between predictions: higher scales increase text influence but risk artifacts. The text is first encoded through BERT to create semantic embeddings that guide the denoising trajectory toward content matching the description.

Encodes the current diffusion step (noise level) as sinusoidal positional embeddings (similar to transformer positional encodings) and injects them into the U-Net at each block. These embeddings allow the model to learn different denoising behaviors at different noise levels — early steps focus on coarse structure, later steps refine details. The sinusoidal encoding ensures smooth interpolation between steps and provides a continuous representation of the noise schedule.

Computes the training objective by sampling random diffusion steps, adding corresponding amounts of Gaussian noise to clean videos, and training the U-Net to predict the added noise. Uses L2 (mean squared error) loss between predicted and actual noise, weighted equally across all diffusion steps. This noise prediction formulation is mathematically equivalent to score matching and enables stable, efficient training of the diffusion model.