imagen-pytorch

Generates images from text descriptions using a multi-stage cascading diffusion architecture where a base UNet first generates low-resolution (64x64) images from noise conditioned on T5 text embeddings, then successive super-resolution UNets (SRUnet256, SRUnet1024) progressively upscale and refine details. Each stage conditions on both text embeddings and outputs from previous stages, enabling efficient high-quality synthesis without requiring a single massive model.

Implements classifier-free guidance mechanism that allows steering image generation toward text descriptions without requiring a separate classifier, using unconditional predictions as a baseline. Incorporates dynamic thresholding that adaptively clips predicted noise based on percentiles rather than fixed values, preventing saturation artifacts and improving sample quality across diverse prompts without manual hyperparameter tuning per prompt.

Provides CLI tool enabling training and inference through configuration files and command-line arguments without writing Python code. Supports YAML/JSON configuration for model architecture, training hyperparameters, and data paths. CLI handles model instantiation, training loop execution, and inference with automatic device detection and distributed training coordination.

Implements data loading pipeline supporting various image formats (PNG, JPEG, WebP) with automatic preprocessing (resizing, normalization, center cropping). Supports augmentation strategies (random crops, flips, color jittering) applied during training. DataLoader integrates with PyTorch's distributed sampler for multi-GPU training, handling batch assembly and text-image pairing from directory structures or metadata files.

Implements comprehensive checkpoint system saving model weights, optimizer state, learning rate scheduler state, EMA weights, and training metadata (epoch, step count). Supports resuming training from checkpoints with automatic state restoration, enabling long training runs to be interrupted and resumed without loss of progress. Checkpoints include version information for compatibility checking.

Supports mixed precision training (fp16/bf16) through Hugging Face Accelerate integration, automatically casting computations to lower precision while maintaining numerical stability through loss scaling. Reduces memory usage by 30-50% and accelerates training on GPUs with tensor cores (A100, RTX 30-series). Automatic loss scaling prevents gradient underflow in lower precision.

Encodes text descriptions into high-dimensional embeddings using pretrained T5 transformer models (typically T5-base or T5-large), which are then used to condition all diffusion stages. The implementation integrates with Hugging Face transformers library to automatically download and cache pretrained weights, supporting flexible T5 model selection and custom text preprocessing pipelines.

Provides modular UNet implementations optimized for different resolution stages: BaseUnet64 for initial 64x64 generation, SRUnet256 and SRUnet1024 for progressive super-resolution, and Unet3D for video generation. Each variant uses attention mechanisms, residual connections, and adaptive group normalization, with configurable channel depths and attention head counts. The modular design allows independent training, selective stage execution, and memory-efficient inference by loading only required stages.

Provides two diffusion implementations: standard Gaussian diffusion (DDPM) and Elucidated diffusion (from Karras et al.), both supporting configurable noise schedules (linear, cosine, sigmoid). The framework abstracts the diffusion process through a unified interface, allowing runtime selection between implementations and custom schedule parameters. Elucidated variant uses improved parameterization for better sample quality and faster convergence.

Unified training interface handling gradient accumulation for effective larger batch sizes, exponential moving average (EMA) weight updates for improved model stability, checkpoint saving/loading, and distributed training via Hugging Face Accelerate library. Supports multi-GPU training with automatic device placement, mixed precision (fp16/bf16), and learning rate scheduling. Trainer manages training loop, loss computation, and model updates across all cascading stages.

Supports training and inference without text conditioning by using null/empty embeddings, enabling unconditional image generation or hybrid modes where text is optional. Architecture remains identical; conditioning is simply disabled by passing zero embeddings. This allows training on unpaired image data and generating diverse samples without text guidance.

Implements inpainting capability where masked regions of images are filled/regenerated while preserving unmasked areas. During training, random masks are applied to images; during inference, the model conditions on both text and the unmasked image regions to generate coherent completions. Masks are incorporated into the diffusion process through concatenation with noisy images, enabling the model to learn spatial context awareness.

Extends image generation to video using Unet3D architecture with 3D convolutions and temporal attention mechanisms. Generates video frames autoregressively or in parallel, conditioning on text embeddings and maintaining temporal coherence through shared weights across frames. Supports variable frame counts and frame rates through flexible temporal dimension handling.

Implements progressive super-resolution where images are upscaled through multiple stages (64→256→1024) using specialized SRUnet models. Each stage conditions on text embeddings and the output from the previous stage, enabling fine-grained detail addition at each resolution level. Stages can be trained independently or jointly, and inference can skip stages for faster generation at intermediate resolutions.