big-sleep

Generates images from text prompts by iteratively optimizing BigGAN latent vectors using CLIP embeddings as a guidance signal. The system encodes text prompts into CLIP embeddings, generates candidate images from BigGAN, computes cosine similarity between text and image embeddings, and backpropagates gradients through the latent space to maximize alignment. Uses exponential moving average (EMA) smoothing on BigGAN parameters to stabilize the optimization trajectory and prevent mode collapse.

Enables simultaneous optimization toward multiple text prompts with configurable weights and negative prompts. The system computes separate CLIP embeddings for each positive and negative prompt, combines them into a weighted loss function where positive prompts maximize similarity and negative prompts minimize it, and performs joint gradient descent on the combined objective. Supports both additive weighting and multiplicative scaling of individual prompt contributions.

Implements a learnable mechanism to select the most relevant BigGAN class embeddings from the full class vocabulary using differentiable top-k selection. The Latents class maintains trainable parameters for class logits, applies softmax to create a probability distribution over classes, and uses straight-through estimators or Gumbel-softmax tricks to enable gradient flow through discrete class selection. This allows the optimization process to discover which semantic classes best align with the text prompt without explicit class specification.

Applies exponential moving average smoothing to BigGAN parameters during the optimization process to stabilize training and prevent divergence. The Model class maintains both the original BigGAN weights and an EMA-smoothed copy; during each optimization step, the EMA weights are updated as a weighted average of previous EMA weights and current weights (with decay factor typically 0.99). The forward pass uses EMA-smoothed weights instead of raw weights, reducing high-frequency noise in the gradient signal and enabling longer optimization runs without mode collapse.

Applies differentiable image transformations (resizing, cropping, rotation, color jittering) to generated images during the optimization loop to improve CLIP alignment and reduce overfitting to specific image statistics. The system generates images at the native BigGAN resolution, applies random augmentations, encodes augmented images through CLIP, and backpropagates gradients through both the augmentation pipeline and the latent vectors. This encourages the optimization to find latent vectors that produce images robust to transformations, improving generalization.

Provides a CLI entry point (dream command) that wraps the Imagine class with progress bars, iteration logging, and automatic image saving. The CLI parses command-line arguments (text prompt, output path, iteration count, learning rate, etc.), instantiates an Imagine object with the parsed configuration, runs the optimization loop with tqdm progress bars showing iteration count and loss values, and saves the final image to disk with optional intermediate checkpoints. Supports both single-image generation and batch processing of multiple prompts.

Supports multiple pre-trained CLIP model variants (ViT-B/32, ViT-L/14) with automatic model loading and caching. The CLIP wrapper loads the specified model from OpenAI's model zoo, caches weights locally to avoid re-downloading, encodes text prompts into embeddings using the text encoder, and encodes generated images using the image encoder. Both encoders output normalized embeddings in the same vector space, enabling cosine similarity computation. The system automatically selects the appropriate model based on available GPU memory and desired quality/speed tradeoff.

Maintains trainable latent vectors (z) and class embeddings that are optimized via gradient descent to maximize CLIP text-image similarity. The Latents class initializes latent vectors from a normal distribution, wraps them in nn.Parameter to make them trainable, and exposes them to PyTorch's autograd system. During each optimization step, the system computes the CLIP loss (negative cosine similarity), backpropagates gradients through CLIP and BigGAN to the latent vectors, and updates them using an optimizer (typically Adam) with a configurable learning rate. The optimization loop runs for a fixed number of iterations or until convergence.

Wraps BigGAN to normalize its output to [-1, 1] range and supports multiple output resolutions (128x128, 256x256, 512x512). The Model class loads the appropriate pre-trained BigGAN checkpoint based on the desired resolution, applies normalization to the raw BigGAN output (which is typically in [-1, 1] or [0, 1] range depending on the model), and optionally applies post-processing (e.g., clipping, scaling) to ensure valid image ranges. The system automatically selects the correct BigGAN variant based on the resolution parameter.