# LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

This repository provides **FLUX checkpoints fine-tuned with LeapAlign** and a minimal inference snippet.

<p align="center">
  <a href="https://arxiv.org/abs/2604.15311"><img src="https://img.shields.io/badge/Paper-arXiv-red?style=for-the-badge" alt="Paper" height="24"></a>
  &emsp;&emsp;&emsp;
  <a href="https://rockeycoss.github.io/leapalign/"><img src="https://img.shields.io/badge/Project-Page-blue?style=for-the-badge" alt="Project Page" height="24"></a>
  &emsp;&emsp;&emsp;
  <!-- <a href=""><img src="https://img.shields.io/badge/Hugging-Face-yellow?style=for-the-badge" alt="Hugging Face" height="24"></a> -->
</p>

## Checkpoints

All checkpoints are trained by applying **LeapAlign** on top of **FLUX.1-dev**.

- `flux_geneval_hpsv2`: FLUX + LeapAlign + **GenEval** prompt set + **HPSv2.1** reward
- `flux_hpdv2_hpsv2`: FLUX + LeapAlign + **HPDv2** prompt set + **HPSv2.1** reward
- `flux_mjhq30k_hpsv3`: FLUX + LeapAlign + **MJHQ-30k** prompt set + **HPSv3** reward

## Download

Download checkpoint folders from the Hugging Face Hub:

```bash
# Download into the current directory (creates local folders)
huggingface-cli download rockeycoss/LeapAlign --local-dir . --include "flux_geneval_hpsv2/*"
huggingface-cli download rockeycoss/LeapAlign --local-dir . --include "flux_hpdv2_hpsv2/*"
huggingface-cli download rockeycoss/LeapAlign --local-dir . --include "flux_mjhq30k_hpsv3/*"
```

## Inference

```python
import torch

from diffusers import FluxPipeline, FluxTransformer2DModel

base_id = "black-forest-labs/FLUX.1-dev"
ckpt_dir = "flux_geneval_hpsv2"  # or "flux_hpdv2_hpsv2" / "flux_mjhq30k_hpsv3"

pipe = FluxPipeline.from_pretrained(
    base_id,
    torch_dtype=torch.bfloat16,
    use_safetensors=True,
)

# Load LeapAlign-aligned transformer weights from local folder
pipe.transformer = FluxTransformer2DModel.from_pretrained(
    ckpt_dir,
    torch_dtype=torch.bfloat16,
)

pipe.to("cuda")

image = pipe(
    "a very cute cat",
    guidance_scale=3.5,
    height=720,
    width=720,
    num_inference_steps=50,
    max_sequence_length=512,
).images[0]

image.save("out.png")
```