$ show facebookresearch/flow_matching
This repository is the official PyTorch library for Flow Matching generative models, developed by Meta FAIR. It provides a modular and extensible framework for implementing continuous, discrete, and Riemannian flow matching algorithms. The core library consists of four main components: probability paths (defining how data transforms from noise to samples), schedulers (noise/time schedules), solvers (ODE or CTMC integration for sampling), and loss functions (e.g., mixture path generalized KL divergence). It also includes utilities for common operations like model wrapping, manifold geometry (Euclidean, sphere, flat torus), and categorical sampling. The library is designed to be flexible, allowing researchers to easily combine different paths, schedulers, and solvers. Beyond the core library, the repository contains extensive examples for image generation (pixel-space continuous and discrete flow matching on CIFAR-10 and ImageNet) and text generation (discrete flow matching on FineWeb-Edu and Wikitext-103). These examples include full training pipelines, evaluation scripts, and model architectures (UNet, DiT-like transformer). The code is research-grade but well-structured, with clear separation of concerns and thorough unit tests for core components. While the examples are somewhat tied to specific datasets and hardware configurations, the library itself is a reusable foundation for any flow matching project. The repository is a hybrid of a research codebase and a production-ready library. The core modules are abstracted and documented, making them suitable for integration into other projects. However, the examples contain research artifacts like hardcoded paths and experimental features (torch.compile, mixed precision) that may require adaptation. Overall, this is a high-quality, actively maintained resource for anyone working on generative modeling with flow matching.
The LLM identified 4 files worth highlighting for reuse:
Abstract base class that wraps any PyTorch nn.Module and defines a custom forward method. This decouples the model interface from the solver and path logic, making it easy to plug in any architecture. The pattern is reusable across all flow matching implementations and other generative models that require a time-dependent vector field.
Defines an abstract Manifold class with essential operations for Riemannian geometry: exponential map, logarithmic map, projection onto tangent space, and inner product. Concrete implementations include Sphere and FlatTorus. This abstraction is critical for Riemannian flow matching and can be extended to other manifolds, making it a reusable foundation for geometric generative models.
Provides an abstract Scheduler class and several concrete implementations (CondOTScheduler, CosineScheduler, VPScheduler, LinearVPScheduler, PolynomialConvexScheduler). These define the time-dependent noise schedule used in probability paths. The design is clean, modular, and independent of specific path or solver implementations, making it reusable in any diffusion or flow matching framework.
Wraps the torchdiffeq ODE solver (e.g., dopri5, rk4) with a unified interface that accepts a learned vector field and integration limits. Supports returning intermediate time steps for visualization. This component isolates the numerical integration logic, allowing users to swap solvers without changing the rest of the pipeline. Reusable for any continuous-time generative model.