DALLE-pytorch

Generates images from text prompts by tokenizing text input, processing through a transformer encoder-decoder architecture, and auto-regressively predicting discrete image tokens in sequence. The model learns joint text-image representations by predicting image token sequences conditioned on text tokens, then decodes predicted tokens back to pixel space via a discrete VAE. This approach enables efficient generation without requiring continuous latent spaces.

Provides a unified VAE interface supporting three distinct image encoding strategies: DiscreteVAE (trainable custom VAE), OpenAIDiscreteVAE (pre-trained 8192-codebook VAE from OpenAI), and VQGanVAE (1024-codebook VAE from Taming Transformers). Each VAE implementation encodes images into discrete token sequences and decodes tokens back to pixels. The abstraction allows swapping VAE backends without modifying the DALLE transformer training code, enabling experimentation with different image compression trade-offs.

Provides a configuration system for specifying DALLE model architecture (depth, width, attention types, VAE type, tokenizer type) and training hyperparameters (learning rate, batch size, warmup steps, gradient clipping). Validates configurations for consistency (e.g., text_seq_len matches tokenizer vocabulary) and instantiates models with validated parameters. Supports YAML/JSON config files for reproducible experiments.

Computes metrics for assessing DALLE training progress and generation quality, including reconstruction loss (for VAE), language modeling loss (for DALLE), and optional perceptual metrics (LPIPS, FID if external libraries available). Supports validation on held-out test sets and periodic generation of sample images during training for visual quality assessment.

Provides Dockerfile and docker-compose configurations for building reproducible training environments with all dependencies (PyTorch, CUDA, DeepSpeed, Horovod) pre-installed. Enables consistent training across different machines and cloud providers without dependency conflicts. Supports GPU passthrough for NVIDIA GPUs and volume mounting for datasets.

Provides five distinct attention implementations (full, axial_row, axial_col, conv_like, sparse) that can be selected per transformer layer to balance memory usage and computational cost. Full attention computes all token-pair interactions; axial attention decomposes 2D image feature maps into row and column attention passes (reducing complexity from O(n²) to O(n√n)); conv_like attention applies local windowed patterns; sparse attention uses DeepSpeed's block-sparse kernels. The framework allows mixing attention types across layers (e.g., full attention for early layers, sparse for later layers).

Abstracts text tokenization through a pluggable interface supporting three strategies: simple built-in tokenizer (basic character/word-level), HuggingFace tokenizers (for Chinese and other languages with pre-trained BPE models), and YouTokenToMe (custom BPE tokenization). Each tokenizer converts variable-length text prompts into fixed-length integer token sequences compatible with the transformer. The abstraction allows swapping tokenizers without retraining the model if vocabulary size remains constant.

Enables multi-GPU and multi-node training through two distributed backends: DeepSpeed (with ZeRO optimizer stages for gradient/parameter sharding) and Horovod (ring-allreduce for gradient synchronization). The framework abstracts distributed training details, allowing users to scale training across multiple GPUs/nodes by specifying backend and world size. DeepSpeed integration enables training larger models by sharding parameters across GPUs; Horovod provides communication-efficient gradient aggregation.

Provides end-to-end VAE training infrastructure including dataset loading, image preprocessing (resizing, normalization), training loop with reconstruction and KL divergence losses, and checkpoint management. The pipeline handles image-to-token encoding during training and supports custom dataset formats. Training produces a discrete VAE checkpoint that can be plugged into DALLE for image generation.

Implements the core DALLE training loop that learns to predict image tokens conditioned on text tokens. The pipeline loads paired text-image datasets, encodes images to tokens via VAE, tokenizes text, and trains the transformer with causal language modeling loss (predicting next image token given text and previous image tokens). Supports mixed-precision training, gradient accumulation, and checkpoint management for long training runs.

Generates images from text prompts at inference time using the trained DALLE model. The generation process tokenizes input text, auto-regressively samples image tokens from the model's predicted probability distributions (using temperature, top-k, or nucleus sampling), and decodes tokens to pixels via VAE. Supports batch generation, seed control for reproducibility, and early stopping based on confidence thresholds.

Provides utilities for loading paired text-image datasets from various formats (directory structures, JSON manifests, HuggingFace datasets), preprocessing images (resizing to fixed dimensions, center-cropping, normalization to [-1, 1] or [0, 1] range), and creating PyTorch DataLoaders with shuffling and batching. Handles image format conversion (PNG, JPEG, WebP), missing data gracefully, and supports distributed data sampling across multiple workers.

Manages saving and loading of model checkpoints during training, including DALLE model weights, VAE weights, optimizer state, learning rate scheduler state, and training metadata (epoch, step, loss). Supports resuming training from checkpoints, enabling long training runs to survive interruptions. Implements checkpoint selection strategies (best loss, latest, periodic) and cleanup of old checkpoints to manage disk space.