---
license: apache-2.0
base_model: ByteDance-Seed/BAGEL-7B-MoT
base_model_relation: finetune
pipeline_tag: any-to-any
datasets:
- YijiaFan/UMM-Reflection-SFT-Data
tags:
- bagel
- unified-model
- text-to-image
- image-editing
- reflection
---

# UMM-Reflection BAGEL-SFT

The reflection-SFT stage of
[UMM-Reflection](https://github.com/waltstephen/UMM-Reflection) (*Learning
Native Reflection in Unified Models*). It is
[BAGEL-7B-MoT](https://huggingface.co/ByteDance-Seed/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](https://huggingface.co/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](https://huggingface.co/datasets/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, ACTION
  `EDIT`/`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.

```bash
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.