songyiren/matting-grpo-lr2e4-mix

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.).

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:

from peft import PeftModel
transformer = 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:

from pipelines.flux_image_new import FluxImagePipeline
from models.utils import DiffusionTrainingModule, load_state_dict, parse_flux_model_configs
from lora.flux_lora import FluxLoRALoader

LORA_TARGET_MODULES = [
    "a_to_qkv", "b_to_qkv", "ff_a.0", "ff_a.2", "ff_b.0", "ff_b.2",
    "a_to_out", "b_to_out", "proj_out",
    "norm.linear", "norm1_a.linear", "norm1_b.linear", "to_qkv_mlp",
]
LORA_RANK = 64

pipe = FluxImagePipeline.from_pretrained(
    torch_dtype=torch.bfloat16, device=device,
    model_configs=parse_flux_model_configs(model_root),
    model_base_path=model_root,
)

# Add empty LoRA adapters matching the training config, then load this repo's weights
helper = DiffusionTrainingModule()
pipe.dit = helper.add_lora_to_model(
    pipe.dit, target_modules=LORA_TARGET_MODULES, lora_rank=LORA_RANK,
    upcast_dtype=pipe.torch_dtype,
)
state_dict = load_state_dict("checkpoints/step300/e2p_fused_lora.safetensors")
state_dict = FluxLoRALoader(torch_dtype=torch.bfloat16, device=device).convert_state_dict(state_dict)
state_dict = helper.mapping_lora_state_dict(state_dict)
pipe.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/:

from diffusers import FluxTransformer2DModel, FluxKontextPipeline

transformer = FluxTransformer2DModel.from_pretrained(
    "songyiren/matting-grpo-lr2e4-mix", subfolder="base_sft10000/transformer",
    torch_dtype=torch.bfloat16,
)
pipe = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev", transformer=transformer, torch_dtype=torch.bfloat16,
)

(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.