UMM-Reflection BAGEL-SFT
The reflection-SFT stage of UMM-Reflection (Learning Native Reflection in Unified Models). It is BAGEL-7B-MoT fine-tuned to inspect its own image, decide whether it is done, and, if not, write an edit and render a corrected image, all in one model and one context.
This checkpoint is the initialization for the RL stage. The final RL model is YijiaFan/UMM-Reflection-BAGEL-RL.
Contents
The layout is the same as the base BAGEL repository. Only ema.safetensors
differs from base BAGEL-7B-MoT (revision
265d1d48ec8e850a29d3a1f208c2a2ec3cd7577b); every other file is
byte-identical to that revision.
| File | sha256 |
|---|---|
ema.safetensors | 52e1e11193f26471dc253d0b65089a2aa024a68ce218cd916ea911b38eddc370 |
Training
- Data. YijiaFan/UMM-Reflection-SFT-Data: 29,529 multi-round reflection trajectories (text-to-image and editing), plus 1,265 prompt-only anchor images generated by base BAGEL.
- Objective. Cross-entropy on controller turns (
<think>, SCORE, ACTIONEDIT/DONE, edit payload) and verifier turns, and flow matching on image transitions conditioned on the previous image. - Schedule. Two stages on 16 GPUs: 1,410 steps with cosine LR 2e-6 to 2e-7, then a constant 2e-7 until one epoch of expanded rows is consumed (step 2,969).
Usage
Download it and point the code at the directory.
git clone https://github.com/waltstephen/UMM-Reflection && cd UMM-Reflection
huggingface-cli download YijiaFan/UMM-Reflection-BAGEL-SFT \
--local-dir pretrained/UMM-Reflection-BAGEL-SFT
# GenEval-553 evaluation of the SFT model
ARM=sft MODEL_DIR=pretrained/UMM-Reflection-BAGEL-SFT bash scripts/eval/geneval553.sh
# RL from this checkpoint (see the repository README for the reward service)
RL_INIT_DIR=pretrained/UMM-Reflection-BAGEL-SFT \
NODE_RANK=0 MASTER_ADDR=<node0> MASTER_PORT=29600 bash scripts/rl/train_rl.sh
License
Apache 2.0, the same as BAGEL. ae.safetensors is the FLUX.1 autoencoder
shipped with BAGEL, under its own Apache 2.0 license. The training data has its
own terms; see the dataset card.