$ show lqiang67/rectified-flow
RectifiedFlow is a unified PyTorch library for implementing Rectified Flow models, a type of generative model that learns ordinary differential equations to transform data between distributions. The library provides a modular framework for training various architectures (MLP, UNet, DiT, LightningDiT) on different datasets, with support for multiple sampling algorithms and interpolation schemes. Key features include a symbolic solver for affine interpolation equations, utilities for converting between different interpolation schemes while preserving rectified flow properties, and a flexible sampler base class for implementing custom integration methods. The codebase demonstrates strong software engineering practices with clear separation of components, comprehensive documentation through examples, and active maintenance.
The LLM identified 4 files worth highlighting for reuse:
Implements AffineInterpSolver and AffineInterp classes that provide symbolic solving for affine interpolation equations. Uses Sympy to precompute symbolic solutions for all combinations of known/unknown variables, then lambdifies them for efficient numerical computation. This eliminates manual derivation of coefficients in algorithms like DDIM and enables automatic conversion between important quantities (velocity field, score function, noise prediction). The component is highly reusable as a standalone mathematical utility for interpolation problems.
Implements the mathematical theorem showing that rectification is equivariant under interpolation transformations. Enables conversion between different interpolation schemes (straight, spherical, DDIM) by matching time and scale factors. This allows transforming a pretrained rectified flow model from one interpolation scheme to another without retraining, leveraging the theoretical property that Z'_t = φ_t(Z_{τ_t}) when X'_t = φ_t(X_{τ_t}). The component demonstrates non-obvious skill in translating mathematical theory into practical code for model transformation.
Provides a flexible base class for implementing custom sampling algorithms for ODEs. Handles time stepping, trajectory recording at configurable intervals, and callback mechanisms for extensibility. Subclasses only need to implement the step() method defining their specific integration scheme (Euler, SDE, etc.). The framework cleanly separates sampling logic from specific integration methods, making it easy to experiment with new samplers while reusing common infrastructure like trajectory recording and time grid handling.
Collection of advanced transformer components including VisionRotaryEmbeddingFast (RoPE), SwiGLUFFN, RMSNorm, and NormAttention modules. Implements modern transformer optimizations like rotary position embeddings, SwiGLU activation, RMS normalization, and QK normalization options. These components can be used independently to build state-of-the-art transformer models and demonstrate skill in integrating multiple recent advances from the transformer literature into a cohesive, reusable library.