$ show zhuyu-cs/MeanFlow
MeanFlow is a PyTorch implementation of a novel single-step flow matching model for high-quality image generation. The codebase provides a complete pipeline for training and evaluating SiT (Scalable Interpolant Transformer) models on ImageNet 256x256 and CIFAR-10, with a specific focus on achieving competitive FID scores in a single sampling step. Key components include a custom loss function with bootstrap and classifier-free guidance, a flexible time-sampling strategy, and distributed training/evaluation scripts using Accelerate. The repository also contains preprocessing utilities to convert ImageNet into VAE-encoded latent LMDB databases. While primarily a research codebase, it demonstrates strong engineering practices with clean module separation, EMA tracking, and comprehensive argument configuration. The implementation closely reproduces the MeanFlow paper results but requires significant computational resources (multiple GPUs, large storage) to run.
The LLM identified 3 files worth highlighting for reuse:
Implements the MeanFlow loss function with a bootstrap mechanism using `torch.func.jvp` for computing derivatives. Supports multiple path types (linear, cosine), adaptive weighting, logit-normal time sampling, and classifier-free guidance with omega/kappa parameters. The loss is designed for flow matching and could be extracted as a reusable component for training other flow-based generative models. Demonstrates advanced use of functional Jacobian-vector products and conditional computation for CFG.
A compact function implementing MeanFlow's single-step and multi-step sampling procedures based on the equation z_r = z_t - (t-r)u(z_t, r, t). Supports classifier-free guidance by conditionally concatenating conditional and unconditional latents. The function is self-contained, takes a model and latent as input, and returns generated samples. Easily extractable for use in other diffusion/flow matching projects.
Provides an LMDB-based dataset for VAE latents with lazy environment initialization per worker process. Includes configurable random horizontal flip augmentation. The dataset reads precomputed moments and supports both original and flipped versions. This pattern is reusable for any large-scale latent dataset stored in LMDB format, offering memory efficiency and multi-worker safe initialization.