Skip to content

Commit 783d425

Browse files
sabulinsabulin
committed
update OmniGen2-RL
1 parent 35ddc1a commit 783d425

41 files changed

Lines changed: 6299 additions & 2 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎examples/OmniGen2-RL/.gitignore‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
# Created by https://www.toptal.com/developers/gitignore/api/macos,python
2+
# Edit at https://www.toptal.com/developers/gitignore?templates=macos,python
3+
4+
### macOS ###
5+
# General
6+
.DS_Store
7+
.AppleDouble
8+
.LSOverride
9+
10+
# Icon must end with two \r
11+
Icon
12+
13+
14+
# Thumbnails
15+
._*
16+
17+
# Files that might appear in the root of a volume
18+
.DocumentRevisions-V100
19+
.fseventsd
20+
.Spotlight-V100
21+
.TemporaryItems
22+
.Trashes
23+
.VolumeIcon.icns
24+
.com.apple.timemachine.donotpresent
25+
26+
# Directories potentially created on remote AFP share
27+
.AppleDB
28+
.AppleDesktop
29+
Network Trash Folder
30+
Temporary Items
31+
.apdisk
32+
33+
### macOS Patch ###
34+
# iCloud generated files
35+
*.icloud
36+
37+
### Python ###
38+
# Byte-compiled / optimized / DLL files
39+
__pycache__/
40+
*.py[cod]
41+
*$py.class
42+
43+
# C extensions
44+
*.so
45+
46+
# Distribution / packaging
47+
.Python
48+
build/
49+
develop-eggs/
50+
dist/
51+
downloads/
52+
eggs/
53+
.eggs/
54+
lib/
55+
lib64/
56+
parts/
57+
sdist/
58+
var/
59+
wheels/
60+
share/python-wheels/
61+
*.egg-info/
62+
.installed.cfg
63+
*.egg
64+
MANIFEST
65+
66+
# PyInstaller
67+
# Usually these files are written by a python script from a template
68+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
69+
*.manifest
70+
*.spec
71+
72+
# Installer logs
73+
pip-log.txt
74+
pip-delete-this-directory.txt
75+
76+
# Unit test / coverage reports
77+
htmlcov/
78+
.tox/
79+
.nox/
80+
.coverage
81+
.coverage.*
82+
.cache
83+
nosetests.xml
84+
coverage.xml
85+
*.cover
86+
*.py,cover
87+
.hypothesis/
88+
.pytest_cache/
89+
cover/
90+
91+
# Translations
92+
*.mo
93+
*.pot
94+
95+
# Django stuff:
96+
*.log
97+
local_settings.py
98+
db.sqlite3
99+
db.sqlite3-journal
100+
101+
# Flask stuff:
102+
instance/
103+
.webassets-cache
104+
105+
# Scrapy stuff:
106+
.scrapy
107+
108+
# Sphinx documentation
109+
docs/_build/
110+
111+
# PyBuilder
112+
.pybuilder/
113+
target/
114+
115+
# Jupyter Notebook
116+
.ipynb_checkpoints
117+
118+
# IPython
119+
profile_default/
120+
ipython_config.py
121+
122+
# pyenv
123+
# For a library or package, you might want to ignore these files since the code is
124+
# intended to run in multiple environments; otherwise, check them in:
125+
# .python-version
126+
127+
# pipenv
128+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
129+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
130+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
131+
# install all needed dependencies.
132+
#Pipfile.lock
133+
134+
# poetry
135+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
136+
# This is especially recommended for binary packages to ensure reproducibility, and is more
137+
# commonly ignored for libraries.
138+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
139+
#poetry.lock
140+
141+
# pdm
142+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
143+
#pdm.lock
144+
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
145+
# in version control.
146+
# https://pdm.fming.dev/#use-with-ide
147+
.pdm.toml
148+
149+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
150+
__pypackages__/
151+
152+
# Celery stuff
153+
celerybeat-schedule
154+
celerybeat.pid
155+
156+
# SageMath parsed files
157+
*.sage.py
158+
159+
# Environments
160+
.env
161+
.venv
162+
env/
163+
venv/
164+
ENV/
165+
env.bak/
166+
venv.bak/
167+
168+
# Spyder project settings
169+
.spyderproject
170+
.spyproject
171+
172+
# Rope project settings
173+
.ropeproject
174+
175+
# mkdocs documentation
176+
/site
177+
178+
# mypy
179+
.mypy_cache/
180+
.dmypy.json
181+
dmypy.json
182+
183+
# Pyre type checker
184+
.pyre/
185+
186+
# pytype static type analyzer
187+
.pytype/
188+
189+
# Cython debug symbols
190+
cython_debug/
191+
192+
# PyCharm
193+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
194+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
195+
# and can be added to the global gitignore or merged into this file. For a more nuclear
196+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
197+
#.idea/
198+
199+
### Python Patch ###
200+
# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
201+
poetry.toml
202+
203+
# ruff
204+
.ruff_cache/
205+
206+
# LSP config files
207+
pyrightconfig.json
208+
209+
# End of https://www.toptal.com/developers/gitignore/api/macos,python
210+
211+
local_scripts/
212+
213+
omnigen2/utils/vpn_utils.py
214+
215+
test_tokenizer.py
216+
save_pipeline.py
217+
app.sh
218+
logs/
219+
results/
220+
test_jsonl*
221+
pbs_files/
222+
convert_ckpt_to_pipeline.py
223+
inference_test_efficiency.py
224+
upload_pipeline*
225+
example_images_resized/
226+
example_t2i_test_efficiency*.sh
227+
example_edit_test_efficiency*.sh
228+
example_in_context_generation_test_efficiency*.sh
229+
intro*
230+
resize_example_images.py
231+
save_pipeline.py
232+
outputs_gradio/*
233+
test.py

‎examples/OmniGen2-RL/README.md‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,4 +68,40 @@ bash evaluation/GEdit-Bench/omnigen2_16samples_select_best_editscore_pass4_eval.
6868
By comparing these results to the baseline performance of the original model, you will see the benefits of applying EditScore as a reranker.
6969

7070
## Application 2: Reinforcement Fine-Tuning
71-
TBD.
71+
Use EditScore to provide a high-quality reward signal to train models for significantly better image editing performance. We employ the FlowGRPO algorithm combined with EditScore's accurate evaluation capabilities to achieve end-to-end reinforcement learning fine-tuning.
72+
73+
### 1. Data and Model Download
74+
Download RL training data from [EditScore-RL-Data](https://huggingface.co/datasets/EditScore/EditScore-RL-Data), then put the `rl.jsonl` into `data/` and change its path in `data_configs/train/train.yml`
75+
76+
Download the base model OmniGen2 form [OmniGen2](https://huggingface.co/OmniGen2/OmniGen2),then change the model file format to pytorch_model.bin and modify `model.pretrained_model_path` in `options/omnigen2_edit_rl.yml`
77+
78+
### 2. Start Reward Server
79+
80+
Before beginning training, you need to start the EditScore reward server to provide real-time reward signal evaluation for RL training.
81+
82+
### 3. Start Training
83+
84+
**Configure Training Parameters**
85+
86+
Edit the `options/omnigen2_edit_rl.yml` configuration file, focusing on these key parameters:
87+
- `train.global_batch_size`: Global batch size (num_machines * num_unique_prompts_per_sampling * num_images_per_prompt)
88+
- `train.rl.num_images_per_prompt`: Rollout number of one prompt
89+
- `train.rl.num_unique_prompts_per_sampling`: Number of global unique prompts
90+
91+
92+
**Launch Distributed Training**
93+
```bash
94+
# Single machine training (8*H100 GPUs)
95+
bash scripts/train/omnigen2_edit_rl.sh
96+
97+
# Multi-machine distributed training
98+
```
99+
100+
> **⚠️ Training Configuration Key Points**
101+
>
102+
> **Reward Server IP**: Ensure the `REWARD_SERVER_IP` environment variable in training scripts points to the correct reward server address
103+
104+
105+
### 4. Training Outputs and Monitoring
106+
107+
Logs and saved model checkpoints in `experiments/`

‎examples/OmniGen2-RL/data_configs/train/train.yml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
ratio_type: inside_ratio
2+
3+
data:
4+
-
5+
path: 'data/rl.jsonl'
6+
type: 'edit'
7+
ratio: !!float 1

‎examples/OmniGen2-RL/docs/README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
## Apply EditScore to Image Editing
2+
### Best-of-N selection**

‎examples/OmniGen2-RL/evaluation/GEdit-Bench/flux_kontext_dev_pass1_best.sh‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# !/bin/bash
2+
SHELL_FOLDER=$(cd "$(dirname "$0")";pwd)
3+
cd $(dirname $SHELL_FOLDER)
4+
cd ../
5+
6+
best=(
7+
1
8+
2
9+
4
10+
8
11+
16
12+
)
13+
14+
for b in "${best[@]}"
15+
do
16+
accelerate launch --num_processes 1 evaluation/GEdit-Bench/test_gedit_score.py \
17+
--result_dir evaluation/GEdit-Bench/results/FLUX-Kontext-dev/results_gs${guidance_scale}_16samples_pass1_best${b} \
18+
--backbone gpt-4.1 \
19+
--openai_url https://api.openai.com/v1/chat/completions \
20+
--max_workers 30 \
21+
--key PUT-YOUR-KEY-HERE
22+
23+
python evaluation/GEdit-Bench/calculate_statistics.py \
24+
--result_dir evaluation/GEdit-Bench/results/FLUX-Kontext-dev/results_gs${guidance_scale}_16samples_pass1_best${b}/viescore_gpt-4.1 \
25+
--language en
26+
done

‎examples/OmniGen2-RL/evaluation/GEdit-Bench/flux_kontext_dev_pass4_best.sh‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# !/bin/bash
2+
SHELL_FOLDER=$(cd "$(dirname "$0")";pwd)
3+
cd $(dirname $SHELL_FOLDER)
4+
cd ../
5+
6+
best=(
7+
1
8+
2
9+
4
10+
8
11+
16
12+
)
13+
14+
for b in "${best[@]}"
15+
do
16+
accelerate launch --num_processes 1 evaluation/GEdit-Bench/test_gedit_score.py \
17+
--result_dir evaluation/GEdit-Bench/results/FLUX-Kontext-dev/results_gs${guidance_scale}_16samples_pass4_best${b} \
18+
--backbone gpt-4.1 \
19+
--openai_url https://api.openai.com/v1/chat/completions \
20+
--max_workers 30 \
21+
--key PUT-YOUR-KEY-HERE
22+
23+
python evaluation/GEdit-Bench/calculate_statistics.py \
24+
--result_dir evaluation/GEdit-Bench/results/FLUX-Kontext-dev/results_gs${guidance_scale}_16samples_pass4_best${b}/viescore_gpt-4.1 \
25+
--language en
26+
done

0 commit comments

Comments
 (0)