Matting GRPO (lr=2e-4, mixed dataset)
FLUX.1-Kontext-dev LoRA checkpoints from a GRPO (Flow-GRPO) reinforcement-learning run for
image matting, fine-tuned on top of a stage-1 full-parameter SFT checkpoint
(
FLUX.1-Kontext-dev-sft10000). Full methodology (SFT + RL, hyperparameters, loss design,
per-step eval tables) is in
TRAINING_DETAILS.md.
- Learning rate: constant
2e-4 (no LR scheduler)
- Training data:
e2p_matting_grpo_mix2 — a merged set combining a ~83.5k-row generic
synthetic matting mix with the AIM-500 / P3M-10k / AM-2k benchmark train+val splits
- Reward: matting alpha-matte quality (MSE/MAD/SAD/Grad/Conn based)
- Checkpoint cadence: every 100 steps
Checkpoints
Each checkpoints/step{100,200,300,400,500,600}/ folder has two formats of the same
checkpoint. Training was stopped at step-600 — see the recommendation below the table.
peft_lora/ (adapter_config.json + adapter_model.safetensors) — the raw,
complete PEFT LoRA adapter exactly as saved by the Flow-GRPO training loop
(rank 64, all target modules). This is the lossless/canonical format — use this one
unless you specifically need the fused format below.e2p_fused_lora.safetensors — an SVD-compressed, fused conversion of the PEFT
adapter into this project's internal "E2P" LoRA layout (produced by
code/convert_flux_peft_lora_to_e2p.py), used for this project's own
batch_inference.py eval pipeline. This conversion is lossy: the fusion process
drops ~114 source modules per checkpoint that don't map cleanly (mostly
norm1.linear/norm1_context.linear layers), and re-loading it reports
missing_lora=228 out of the reconstructed state dict. It reproduces the
eval_images/ predictions in this repo, but is not a complete copy of the trained
LoRA — prefer peft_lora/ for anything downstream (further training, merging,
re-evaluation, etc.).
| Step | am SAD | aim SAD | p3m-np SAD | am Conn | aim Conn | p3m-np Conn |
|---|
| 100 | 9.368 | 15.918 | 8.386 | 5.759 | 10.077 | 4.936 |
| 200 | 9.207 | 15.736 | 8.250 | 5.811 | 10.084 | 4.974 |
| 300 | 9.112 | 15.645 | 8.160 | 5.886 | 10.153 | 5.036 |
| 400 | 9.156 | 15.692 | 8.246 | 5.985 | 10.230 | 5.134 |
| 500 | 9.140 | 15.710 | 8.266 | 6.055 | 10.302 | 5.199 |
| 600 | 9.130 | 15.726 | 8.245 | 6.098 | 10.355 | 5.234 |
SAD bottoms out around step-300 and mildly reverses after; Conn (connectivity error)
rises monotonically at every single checkpoint on all three benchmarks, all the way
through step-600. Recommendation: use checkpoints/step300/ — it's the best-balanced
checkpoint in this run; steps 400-600 trade further Conn degradation for no net gain (and
by step-600, Grad also turns up on all three benchmarks). See TRAINING_DETAILS.md for
the full MSE/MAD/SAD/Grad/Conn tables and discussion. Full per-image predictions for every
checkpoint are in eval_images/.
How to load / run inference
peft_lora/ (recommended) is a standard PEFT adapter — load it the standard way once
you have a FLUX.1-Kontext-dev diffusers transformer:
1from peft import PeftModel
2transformer = PeftModel.from_pretrained(transformer, "checkpoints/step300/peft_lora")
e2p_fused_lora.safetensors only works with this project's custom loader — it is the
exact script used to produce the predictions in eval_images/. It depends on this
project's internal framework (pipelines/, models/, lora/ — a custom
DiffSynth-Studio-style loader, not bundled in this repo), so it needs to run inside that
project's environment, pointed at a local FLUX.1-Kontext-dev checkpoint. The core
loading logic:
1from pipelines.flux_image_new import FluxImagePipeline
2from models.utils import DiffusionTrainingModule, load_state_dict, parse_flux_model_configs
3from lora.flux_lora import FluxLoRALoader
4
5LORA_TARGET_MODULES = [
6 "a_to_qkv", "b_to_qkv", "ff_a.0", "ff_a.2", "ff_b.0", "ff_b.2",
7 "a_to_out", "b_to_out", "proj_out",
8 "norm.linear", "norm1_a.linear", "norm1_b.linear", "to_qkv_mlp",
9]
10LORA_RANK = 64
11
12pipe = FluxImagePipeline.from_pretrained(
13 torch_dtype=torch.bfloat16, device=device,
14 model_configs=parse_flux_model_configs(model_root),
15 model_base_path=model_root,
16)
17
18# Add empty LoRA adapters matching the training config, then load this repo's weights
19helper = DiffusionTrainingModule()
20pipe.dit = helper.add_lora_to_model(
21 pipe.dit, target_modules=LORA_TARGET_MODULES, lora_rank=LORA_RANK,
22 upcast_dtype=pipe.torch_dtype,
23)
24state_dict = load_state_dict("checkpoints/step300/e2p_fused_lora.safetensors")
25state_dict = FluxLoRALoader(torch_dtype=torch.bfloat16, device=device).convert_state_dict(state_dict)
26state_dict = helper.mapping_lora_state_dict(state_dict)
27pipe.dit.load_state_dict(state_dict, strict=False)
Then run with the fixed prompt "Transform to matting map while maintaining original composition" and two Kontext reference images [photo, trimap] — see
code/batch_inference.py for the full inference loop (dataset iteration, resolution,
1-step flow-matching inference, cfg_scale=1, seed=42).
code/convert_flux_peft_lora_to_e2p.py is the script that produced these
e2p_fused_lora.safetensors files from the raw PEFT LoRA checkpoint directories saved by
the Flow-GRPO training loop (checkpoints/checkpoint-100/lora/adapter_model.safetensors)
— included for provenance; not needed to use the weights in this repo, which are already
in the converted format.
Base checkpoint (base_sft10000/)
base_sft10000/transformer/ is the stage-1 full-parameter SFT checkpoint (step-10000)
that every LoRA in checkpoints/ above is trained from — this is the dev-sft10000 base
referenced throughout this repo and in TRAINING_DETAILS.md. It contains only
config.json + diffusion_pytorch_model.safetensors (the fine-tuned DiT transformer,
~22 GiB, bf16, diffusers FluxTransformer2DModel format) — i.e. just the weights this
project's own SFT training actually changed.
The rest of the pipeline (VAE, text encoders, tokenizer, scheduler,
model_index.json) is
unmodified from the original base model and is
not re-hosted here — load those
from
black-forest-labs/FLUX.1-Kontext-dev
directly (note: gated, non-commercial license) and swap in this repo's
transformer/:
1from diffusers import FluxTransformer2DModel, FluxKontextPipeline
2
3transformer = FluxTransformer2DModel.from_pretrained(
4 "songyiren/matting-grpo-lr2e4-mix", subfolder="base_sft10000/transformer",
5 torch_dtype=torch.bfloat16,
6)
7pipe = FluxKontextPipeline.from_pretrained(
8 "black-forest-labs/FLUX.1-Kontext-dev", transformer=transformer, torch_dtype=torch.bfloat16,
9)
(This project's own eval pipeline instead uses the custom FluxImagePipeline loader shown
above, pointed at a local directory with this transformer/ swapped in alongside symlinks
to the original model's other components — see models/utils.py::parse_flux_model_configs.)
Eval images
eval_images/step{100,200,300,400,500,600}/{am,aim,p3m-np}/ contains the rendered matting
predictions (PNG, converted from the raw float32 [0,1] RGB prediction arrays) for every
image in the three benchmark validation sets, for each checkpoint.
