LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
_{Official PyTorch Implementation}

LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
Zhanhao Liang^1,2,*, Tao Yang^2,*,†, Jie Wu², Chengjian Feng², Liang Zheng^1
1ANU, ²ByteDance Seed
^*Equal contribution, ^†Project Lead

       

Overview

LeapAlign is a post-training method for aligning flow matching models with human preferences by backpropagating reward gradients through the differentiable generation process. Since backpropagation through the full sampling trajectory is memory-intensive and prone to gradient explosion, early generation steps that determine the global structure and layout of the final image are difficult to optimize. LeapAlign addresses this by constructing a two-step leap trajectory and backpropagating gradients through it instead of the full sampling trajectory, enabling gradient propagation to any generation step.

Dependency Setup

1. Create the Conda environment

conda create --name leapalign python=3.11 -y
conda activate leapalign

2. Install project dependencies

./setup_env.sh

This script installs the Python dependencies for LeapAlign, installs python3-tk, and clones and installs HPSv2.

3. Download checkpoints

mkdir -p data
huggingface-cli download black-forest-labs/FLUX.1-dev --local-dir ./data/flux

mkdir -p hps_ckpt
huggingface-cli download xswu/HPSv2 HPS_v2.1_compressed.pt --local-dir ./hps_ckpt
huggingface-cli download laion/CLIP-ViT-H-14-laion2B-s32B-b79K open_clip_pytorch_model.bin --local-dir ./hps_ckpt

4. Prepare prompt embeddings

The preprocessing scripts use 8 GPUs by default.

Prepare HPDv2 prompt embeddings:

bash scripts/preprocess/preprocess_flux_hpdv2_rl_embeddings.sh

Prepare GenEval prompt embeddings:

bash scripts/preprocess/preprocess_flux_geneval_rl_embeddings.sh

Fine-tune with LeapAlign

The provided training scripts are configured for a 2x8 setup by default, i.e., 2 nodes with 8 GPUs each.

HPDv2

Launch on node 0:

MASTER_ADDR=<node0_host> MASTER_PORT=29500 NNODES=2 NODE_RANK=0 NPROC_PER_NODE=8 \
bash scripts/finetune/train_flux_leapalign_hpdv2_hpsv2.sh

Launch on node 1:

MASTER_ADDR=<node0_host> MASTER_PORT=29500 NNODES=2 NODE_RANK=1 NPROC_PER_NODE=8 \
bash scripts/finetune/train_flux_leapalign_hpdv2_hpsv2.sh

GenEval

Launch on node 0:

MASTER_ADDR=<node0_host> MASTER_PORT=29500 NNODES=2 NODE_RANK=0 NPROC_PER_NODE=8 \
bash scripts/finetune/train_flux_leapalign_geneval_hpsv2.sh

Launch on node 1:

MASTER_ADDR=<node0_host> MASTER_PORT=29500 NNODES=2 NODE_RANK=1 NPROC_PER_NODE=8 \
bash scripts/finetune/train_flux_leapalign_geneval_hpsv2.sh

Evaluation

HPDv2

1. Sample images

The example below uses 8 GPUs.

torchrun --nproc_per_node=8 fastvideo/evaluation/sample_flux_images.py \
    --model_path <path_to_finetuned_transformer> \
    --output_dir ./output/flux_leapalign_hpsv2

--model_path should point to a checkpoint directory saved during fine-tuning.

2. Compute average HPSv2.1 reward

python fastvideo/evaluation/evaluate_reward_values.py \
    --results_dir ./output/flux_leapalign_hpsv2 \
    --output_dir ./output/rewards_log

The script prints the average HPSv2.1 score and writes <results_dir_basename>_rewards_log.json to --output_dir.

GenEval

1. Sample images

The example below uses 8 GPUs.

torchrun --nproc_per_node=8 fastvideo/evaluation/sample_flux_images_geneval.py \
    --model_path <path_to_finetuned_transformer> \
    --outdir ./output/flux_leapalign_geneval

2. Score with the GenEval evaluator

Follow the setup and scoring instructions in the official GenEval repo and run their evaluator against the directory produced in step 1.

Acknowledgments

LeapAlign builds upon several excellent open-source projects. We would like to thank the authors and contributors of:

• DanceGRPO
• FastVideo
• Diffusers
• MixGRPO
• GenEval
• HPSv2

📬 Citation

If you find this code useful in your research, please consider citing:

@article{liang2026leapalign,
  title={LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories},
  author={Liang, Zhanhao and Yang, Tao and Wu, Jie and Feng, Chengjian and Zheng, Liang},
  journal={arXiv preprint arXiv:2604.15311},
  year={2026}
}