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