Flow-Factory

Unified Online RL and Offline Fine-Tuning for Diffusion and Flow-Matching Models

🔥 News

• [2026-10-17] Prompt enhancement as a separate step! ff-prompt-enhance rewrites a dataset's prompts with a model's prompt enhancer (FLUX.3 pup, FLUX.2, LTX-2, the Qwen-Image-2.1 T2I and I2I rewriters, or any OpenAI-compatible server) and writes a new dataset whose prompts training uses verbatim. See Prompt Enhancement.

• [2026-10-16] FLUX.3 Image Klein support! Train both the base (black-forest-labs/flux-3-image-klein-base) and step-distilled flow-map (black-forest-labs/flux-3-image-klein) text-to-image and multi-reference image-to-image models with the GRPO + LoRA recipe. Requires the pinned Diffusers submodule.

• [2026-09-20] Qwen-Image 2.1 support! Train the unified text-to-image and ordered multi-image editing model with the GRPO + LoRA recipe. Requires the pinned Diffusers submodule.

• [2026-08-21] MiniMax H3 Audio-Video support! Fine-tune text-to-audio-video, first/last-frame-to-audio-video, and ordered-reference-to-audio-video workflows with GRPO + LoRA through the released Diffusers API.

• [2026-04-25] LTX-2 Audio-Video support! Generate synchronized audio-video content with RL fine-tuning through the released Diffusers API.

• [2026-02-01] Support for multiple Attention Backends! Select one in the unified acceleration: block, where it can be combined with torch.compile and feature caching:

acceleration:
  shared:
    - name: attention_backend
      params: { backend: "flash_hub" } # e.g. "native", "xformers", "flash_hub", "_flash_3_hub", "_flash_3_varlen_hub"

The name is forwarded to Diffusers' transformer.set_attention_backend; see the official list and the Acceleration guide. We recommend installing kernels (pip install kernels) and using a *_hub backend instead of building Flash-Attention yourself.

📕 Table of Contents

• Supported Models
• Supported Algorithms
• Get Started
  • Installation
  • Experiment Trackers
  • Quick Start Example
• Guidance
• Dataset
  • Offline SFT and Preference Data
  • Text-to-Image & Text-to-Video
  • Image-to-Image & Image-to-Video
• Reward Model
• Acknowledgements
• Citation

🤗 Supported Models

Task	Model	Model Size	Model Type
Text-to-Image	stable-diffusion-3.5-medium/large	2.5B/8.1B	sd3-5
	FLUX.1-dev	13B	flux1
	Z-Image-Turbo	6B	z-image
	Z-Image	6B	z-image
	Qwen-Image	20B	qwen-image
	Qwen-Image-2512	20B	qwen-image
Image-to-Image	FLUX.1-Kontext-dev	13B	flux1-kontext
Image(s)-to-Image	Qwen-Image-Edit-2509	20B	qwen-image-edit-plus
	Qwen-Image-Edit-2511	20B	qwen-image-edit-plus
Text-to-Image & Image(s)-to-Image	FLUX.2-dev	32B	flux2
	FLUX.2-klein-4B	4B	flux2-klein
	FLUX.2-klein-9B	9B	flux2-klein
	FLUX.2-klein-base-4B	4B	flux2-klein
	FLUX.2-klein-base-9B	9B	flux2-klein
	flux-3-image-klein-base	5.5B	flux3
	flux-3-image-klein	5.5B	flux3
	BAGEL-7B-MoT	14B	bagel
	SenseNova-U1 1.0	16B	sensenova
	SenseNova-U1 1.5	16B	sensenova
	Qwen-Image 2.1	7B	qwen-image-2.1
Text-to-Video	Wan2.1-T2V-1.3B	1.3B	wan2_t2v
	Wan2.1-T2V-14B	14B	wan2_t2v
	Wan2.2-TI2V-5B	5B	wan2_t2v
	Wan2.2-T2V-A14B	A14B	wan2_t2v
Image-to-Video	Wan2.1-I2V-14B-480P	14B	wan2_i2v
	Wan2.1-I2V-14B-720P	14B	wan2_i2v
	Wan2.2-TI2V-5B	5B	wan2_i2v
	Wan2.2-I2V-A14B	A14B	wan2_i2v
First/Last-Frame-to-Video	Wan2.1-FLF2V-14B-720P	14B	wan2_i2v
Text-to-Audio-Video	LTX-2	19B	ltx2_t2av
	LTX-2.3 (Diffusers)	22B	ltx2_t2av
Image-to-Audio-Video	LTX-2	19B	ltx2_i2av
	LTX-2.3 (Diffusers)	22B	ltx2_i2av
Text-to-Audio-Video	MiniMax H3 T2VA	33B	minimax-h3-t2va
First/Last-Frame-to-Audio-Video	MiniMax H3 FL2VA	33B	minimax-h3-fl2va
Ordered-Reference-to-Audio-Video	MiniMax H3 Ref2VA	33B	minimax-h3-ref2va

• Every model type above also supports SFT and offline DPO; see the offline model matrix.
• MiniMax H3 has workflow-specific constraints (batch size 1, no CFG, separate video/audio trajectories); see MiniMax H3 in the dataset guide.
• To add a model, see New Model.

💻 Supported Algorithms

Algorithm	trainer_type	Reference / objective
SFT	sft	Supervised flow matching over V2 demonstrations
Offline DPO	offline-dpo	Diffusion-DPO over V2 preference pairs
Online DPO	dpo	Diffusion-DPO with generated, reward-ranked pairs
GRPO	grpo	Flow-GRPO / Dance-GRPO
DiffusionNFT	nft	DiffusionNFT
AWM	awm	Advantage Weighted Matching
DGPO	dgpo	DGPO
GRPO-Guard	grpo-guard	GRPO-Guard
DPPO	dppo	Flow-DPPO
CRD	crd	Centered Reward Distillation (Blog (Chinese))
DiffusionOPD	diffusion-opd	DiffusionOPD
DMD2	dmd2	DMD2
TDM	tdm	Trajectory Distribution Matching
TDM-R1	tdm-r1	TDM-R1

See the Algorithm guide for objectives and configuration.

Models and algorithms are decoupled at the framework interface. Ready-to-run recipes are indexed in examples/README.md.

🚀 Get Started

Installation

git clone https://github.com/X-GenGroup/Flow-Factory.git
cd Flow-Factory
pip install -e .

Flow-Factory requires Python 3.10 or newer and PyTorch 2.10 or newer (which provides the native torch.optim.Muon used by Muon optimizer configs). pip install -e . also installs the released diffusers>=0.40.0, which covers every model except the two below.

Pinned Diffusers for FLUX.3 and Qwen-Image 2.1

Until FLUX.3 and Qwen-Image 2.1 reach a released Diffusers package, install the bundled submodule before Flow-Factory (FLUX.3 requires Python 3.11 or newer):

git submodule update --init
pip install -e ./diffusers
pip install -e .

Optional dependencies

• DeepSpeed: pip install -e .[deepspeed].
• Bagel requires flash-attn (>= 2.5.8) and opencv-python: pip install -e .[bagel]. The [bagel] extra is intentionally not part of [all] because flash-attn is heavy to build.
• Experiment trackers: see Experiment Trackers. Reward-specific dependencies are listed in the Rewards guide.
• PyAV >=17.0.0 decodes ordered video/audio references and target media. TorchAudio 2.10 delegates audio loading and saving to TorchCodec, which also requires FFmpeg shared libraries.

A CUDA training image (Python 3.12, uv-based install, PyTorch 2.10 + cu129, deepspeed, wandb, released diffusers, FFmpeg libraries) is defined under docker/docker-cuda/; see docker/README.md for build and run instructions. For multi-node launches, see multinode_examples/README.md.

Experiment Trackers

To use Weights & Biases or SwanLab to log experimental results, install extra dependencies via pip install -e .[wandb] or pip install -e .[swanlab].

After installation, set the corresponding fields in the log section of the config file:

log:
  run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})
  project: "Flow-Factory"  # Project name for logging
  logging_backend: "wandb"  # Options: wandb, swanlab, tensorboard, none

These trackers allow you to visualize both training samples and metric curves online:

Quick Start Example

Start training with the following simple command:

ff-train examples/grpo/lora/flux1/default.yaml

Offline examples use strict V2 manifests and require no training reward model. The SD3.5 recipes run on the checked-in fixtures:

ff-train examples/sft/lora/sd3_5/default.yaml
ff-train examples/offline_dpo/lora/sd3_5/default.yaml

examples/sft/lora/wan2_t2v/wan22_ti2v5b.yaml expects a V2 video demonstration dataset that you prepare first. All recipes are indexed in examples/README.md.

📖 Guidance

We provide a set of guidance documents to help you understand the framework and extend it. For a comprehensive understanding of the framework's design and motivation, refer to our technique report.

Document	Description
Workflow	End-to-end training pipeline from data preprocessing to policy optimization, plus async rewards and reward/optimization overlap
Algorithms	Supported online RL, SFT, offline DPO, and distillation algorithms and their configurations
Rewards	Reward model system: built-in models, custom rewards, and remote reward servers
Datasets	Dataset schemas, media paths, and ordered-reference inputs
Prompt Enhancement	ff-prompt-enhance: rewrite a dataset's prompts with a model's prompt enhancer before training
Acceleration	Attention backends, torch.compile, and rollout feature caching
New Model	How to add support for a new Diffusion/Flow-Matching model
Examples	Index of ready-to-run recipes per algorithm and model
Docker / Multi-node	Container image and multi-node launch setup

Contributors and coding agents start from AGENTS.md, the development guide: repository conventions, risk profiles, and the GPU merge gate.

📊 Dataset

The unified structure of dataset is:

|---- dataset
|----|--- train.txt / train.jsonl
|----|--- test.txt / test.jsonl (optional)
|----|--- images (optional)
|----|---| image1.png
|----|---| ...
|----|--- videos (optional)
|----|---| video1.mp4
|----|---| ...
|----|--- audios (optional)

Configs list datasets under data.datasets; see the dataset guide for the full schema, media roots, and evaluation splits.

Offline SFT and Preference Data

SFT and offline DPO use strict JSONL with schema_version: 2. Public media objects use type as their sole discriminator:

{"schema_version":2,"input":{"prompt":"A clean poster.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"targets/poster.png"}]}},"metadata":{}}
{"schema_version":2,"input":{"prompt":"A clean poster.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"pairs/chosen.png"}]},"rejected":{"media":[{"type":"image","path":"pairs/rejected.png"}]}},"metadata":{}}
{"schema_version":2,"input":{"prompt":"Animate toward this ending.","media":[{"type":"image","path":"conditions/end.png","slot":"last_frame"}]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/story.mp4","fps":24.0},{"type":"audio","path":"targets/story.wav","sample_rate":32000}]}},"metadata":{}}

The optional input-only slot field binds sparse conditions to adapter-declared semantic arguments. Unslotted media fills remaining slots positionally; supervision outputs reject slots. The checked-in SFT demonstration fixture and offline-DPO preference fixture are minimal examples, and the offline smoke builder reconstructs SFT and offline-DPO mini datasets for every implemented image, video, and audio-video profile. Target media are decoded and encoded on the fly and never cached; see Offline V2 records for the schema, caching, and epoch rules.

Text-to-Image & Text-to-Video

For text-to-image and text-to-video tasks, the only required input is the prompt in plain text format. Use train.txt and test.txt (optional) with following format:

A hill in a sunset.
An astronaut riding a horse on Mars.

Example: dataset/pickscore

Each line represents a single text prompt. Alternatively, you can use train.jsonl and test.jsonl in the following format:

{"prompt": "A hill in a sunset."}
{"prompt": "An astronaut riding a horse on Mars."}

Example: dataset/t2is

negative_prompt is also supported:

{"prompt": "A hill in a sunset.", "negative_prompt": "low quality, blurry, distorted, poorly drawn"}
{"prompt": "An astronaut riding a horse on Mars.", "negative_prompt": "low quality, blurry, distorted, poorly drawn"}

Example: dataset/t2is_neg

Image-to-Image & Image-to-Video

For tasks involving conditioning images, use train.jsonl and test.jsonl in the following format:

{"prompt": "A hill in a sunset.", "image": "path/to/image1.png"}
{"prompt": "An astronaut riding a horse on Mars.", "image": "path/to/image2.png"}

Example: dataset/sharegpt4o_image_mini

Media paths in JSONL rows resolve against {dataset_dir}/images, {dataset_dir}/videos, and {dataset_dir}/audios by default. Override them per dataset entry, or for every entry with the top-level data.image_dir / data.video_dir / data.audio_dir:

data:
  datasets:
    - name: my_dataset
      dataset_dir: "path/to/dataset"
      image_dir: "path/to/image_dir"  # default: {dataset_dir}/images
      video_dir: "path/to/video_dir"  # default: {dataset_dir}/videos

For models such as FLUX.2-dev, Qwen-Image-Edit-2511, BAGEL-7B-MoT, and SenseNova-U1 that accept multiple conditioning images, use the images key with an ordered list of image paths:

{"prompt": "A hill in a sunset.", "images": ["path/to/condition_image_1_1.png", "path/to/condition_image_1_2.png"]}
{"prompt": "An astronaut riding a horse on Mars.", "images": ["path/to/condition_image_2_1.png", "path/to/condition_image_2_2.png"]}

💯 Reward Model

Reward models score generated samples for reinforcement learning. A pointwise reward scores each sample independently (e.g. aesthetic quality or text-image alignment); a groupwise reward scores a whole group together, which covers ranking and pairwise comparisons.

Built-in reward models, by family (the Rewards guide lists each one's type, dependencies, and reference):

• Image preference and alignment: PickScore, PickScore_Rank (groupwise), CLIP, hpsv2
• Text rendering and composition: ocr, GenEval, geneval2_soft_tifa
• Audio and audio-video alignment: clap, imagebind
• VLM judges served over an OpenAI-compatible API: vllm_evaluate, rational_rewards_t2i, rational_rewards_edit, qwen_image_bench

Specify a reward model by name in your config file:

rewards:
  name: "aesthetic" # Alias for this reward model
  reward_model: "PickScore" # Reward model type or a path like 'my_package.rewards.CustomReward'
  batch_size: 16
  device: "cuda"
  dtype: bfloat16

See the Rewards guide for VLM judges, multiple rewards, custom reward models, and remote reward servers.

🤗 Acknowledgements

This repository is based on diffusers, accelerate and peft. We thank them for their contributions to the community!!!

📝 Citation

If you find Flow-Factory useful in your research, please consider citing our paper:

@article{ping2026flowfactory,
  title={Flow-Factory: A Unified Framework for Reinforcement Learning in Flow-Matching Models}, 
  author={Bowen Ping and Chengyou Jia and Minnan Luo and Hangwei Qian and Ivor Tsang},
  journal={arXiv preprint arXiv:2602.12529},
  year={2026},
  url={https://arxiv.org/abs/2602.12529}, 
}