release v0.2

Author: Luciennnnnnn

README.md

@@ -28,6 +28,7 @@
 - **Versatile Applications**: Ready to use as a best-in-class reranker to improve editing outputs, or as a high-fidelity reward signal for **stable and effective Reinforcement Learning (RL) fine-tuning**.
 
 ## 🔥 News
+- **2025-10-27**: Released [OmniGen2-EditScore7B-v1.1](https://huggingface.co/OmniGen2/OmniGen2-EditScore7B-v1.1), achieving a **7.01 GEdit score** within **700 steps**, by incorporating the **reweighting strategy** from [TempFlow](https://arxiv.org/abs/2508.04324). Furthermore, we improve json fixing method using [this great library](https://github.com/mangiucugna/json_repair), now EditScore should be more stable across various conditions. Updated to it with `pip install -U editscore`.
 - **2025-10-22**: **Introducing Our Reinforcement Learning Training Framework!**
   We're excited to release our complete RL pipeline, the result of a massive effort to simplify fine-tuning for image editing models. Key features include:
     - **Ready-to-Use RL Dataset**: Includes the complete dataset used in the EditScore project, along with clear usage guidelines and preparation scripts.

editscore/utils.py

@@ -4,6 +4,7 @@
 import regex as re
 import ast
 import random
+import json_repair
 
 def fix_json(input_str):
     # Add double quotes around keys using regex
@@ -179,21 +180,21 @@ def fallback_repair_json(input_str: str) -> str:
 
 def robust_json_fix(s: str):
     try:
-        return json.loads(s)
+        return json_repair.loads(s)
     except Exception:
         pass
 
     for fixer in [fix_json, repair_reasoning_field_robust]:
         s = fixer(s)
         try:
-            return json.loads(s)
+            return json_repair.loads(s)
         except Exception:
             print(f"Error: Cannot fix {fixer.__name__} {s=}")
             continue
 
     try:
         repaired_str = fallback_repair_json(s)
-        return json.loads(repaired_str)
+        return json_repair.loads(repaired_str)
     except Exception as e:
         print(f"Error: Cannot fix fallback_repair_json {s=} {e=}")
         return False
@@ -400,18 +401,6 @@ def mllm_output_to_dict(input_string, give_up_parsing=False, text_prompt=None, s
             print(f"Now fixing: {e1=} {json_str=}")
 
             new_data = robust_json_fix(json_str)
-
-            # try:
-            #     new_data = json.loads(fix_json(json_str))
-            #     return new_data
-            # except Exception as e2:
-            #     try:
-            #         print(f"Now fixing: {e2=} {fix_json(json_str)=}")
-            #         new_data = json.loads(repair_reasoning_field_robust(fix_json(json_str)))
-            #         return new_data
-            #     except Exception as e3:
-            #         print(f"Error: Cannot fix {e3=} {repair_reasoning_field_robust(fix_json(json_str))=}")
-            #         return False
         return new_data
     else:
         print("The required delimiters were not found correctly in the string.")

examples/OmniGen2-RL/evaluation/GEdit-Bench/omnigen2.sh

@@ -0,0 +1,87 @@
+# !/bin/bash
+SHELL_FOLDER=$(cd "$(dirname "$0")";pwd)
+cd $(dirname $SHELL_FOLDER)
+cd ../
+
+experiment_name=omnigen2_edit_rl_4machine_editscore7b_avg8
+step=700
+RANK=0
+WORLD_SIZE=1
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --experiment_name=*)
+            experiment_name="${1#*=}"
+            shift
+            ;;
+        --step=*)
+            step="${1#*=}"
+            shift
+            ;;
+        --rank=*)
+            RANK="${1#*=}"
+            shift
+            ;;
+        --world_size=*)
+            WORLD_SIZE="${1#*=}"
+            shift
+            ;;
+        *)
+            echo "未知参数: $1"
+            shift
+            ;;
+    esac
+done
+
+# 输出配置
+echo "RANK: $RANK"
+echo "WORLD_SIZE: $WORLD_SIZE"
+
+global_shift_index=0
+total_num_images=606
+
+num_gpus_per_machine=$(python -c "import torch; print(torch.cuda.device_count())")
+# Calculate images per machine, rounding up to ensure all data is covered
+num_images_per_machine=$(( (total_num_images + WORLD_SIZE - 1) / WORLD_SIZE ))
+shift_index=$((RANK * num_images_per_machine))
+
+if [ $((total_num_images - shift_index)) -lt $num_images_per_machine ]; then
+    num_images_per_machine=$((total_num_images - shift_index))
+fi
+
+# Calculate base number of images per GPU (for first 7 GPUs)
+num_images_per_gpu=$(( (num_images_per_machine + num_gpus_per_machine - 1) / num_gpus_per_machine ))
+
+text_guidance_scale=5.0
+image_guidance_scale=1.5
+
+for ((i=0; i<num_gpus_per_machine; i++)); do
+    if [ $i -lt $((num_gpus_per_machine - 1)) ]; then
+        # First 7 GPUs process equal amounts
+        start_idx=$((global_shift_index + i * num_images_per_gpu + shift_index))
+        end_idx=$((start_idx + num_images_per_gpu))
+    else
+        # Last GPU processes remaining data
+        start_idx=$((global_shift_index + (num_gpus_per_machine - 1) * num_images_per_gpu + shift_index))
+        end_idx=$((global_shift_index + shift_index + num_images_per_machine))
+    fi
+    echo ${start_idx} ${end_idx}
+
+    CUDA_VISIBLE_DEVICES=${i} WORLD_SIZE=1 nohup accelerate launch --num_processes 1 --num_machines 1 \
+    evaluation/GEdit-Bench/inference.py \
+    --load_from_pipeline \
+    --pipeline_path OmniGen2/OmniGen2 \
+    --transformer_lora_path experiments/${experiment_name}/checkpoint-${step}/transformer_lora \
+    --num_inference_step 50 \
+    --height 1024 \
+    --width 1024 \
+    --text_guidance_scale ${text_guidance_scale} \
+    --image_guidance_scale ${image_guidance_scale} \
+    --time_shift_base_res 168 \
+    --negative_prompt "" \
+    --use_ori_neg_prompt_template \
+    --scheduler "euler" \
+    --result_dir evaluation/GEdit-Bench/results/${experiment_name}/results_step${step}_ts${text_guidance_scale}_ig${image_guidance_scale} \
+    --start_index ${start_idx} --end_index ${end_idx} \
+    > logs/gedit_${experiment_name}_step${step}_ts${text_guidance_scale}_ig${image_guidance_scale}_${start_idx}_${end_idx}.log 2>&1 &
+done

examples/OmniGen2-RL/evaluation/GEdit-Bench/omnigen2_eval.sh

@@ -0,0 +1,42 @@
+# !/bin/bash
+SHELL_FOLDER=$(cd "$(dirname "$0")";pwd)
+cd $(dirname $SHELL_FOLDER)
+cd ../
+
+source "$(dirname $(which conda))/../etc/profile.d/conda.sh"
+conda activate py3.12+pytorch2.7.1+cu126
+
+experiment_name=omnigen2_edit_rl_4machine_editscore7b_avg8
+step=700
+
+# 处理命名参数
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --experiment_name=*)
+            experiment_name="${1#*=}"
+            shift
+            ;;
+        --step=*)
+            step="${1#*=}"
+            shift
+            ;;
+        *)
+            echo "未知参数: $1"
+            shift
+            ;;
+    esac
+done
+
+text_guidance_scale=5.0
+image_guidance_scale=1.5
+
+accelerate launch --num_processes 1 evaluation/GEdit-Bench/test_gedit_score.py \
+--result_dir evaluation/GEdit-Bench/results/${experiment_name}/results_step${step}_ts${text_guidance_scale}_ig${image_guidance_scale} \
+--backbone gpt-4.1 \
+--openai_url https://api.openai.com/v1/chat/completions \
+--max_workers 30 \
+--key PUT-YOUR-KEY-HERE
+
+python evaluation/GEdit-Bench/calculate_statistics.py \
+--result_dir evaluation/GEdit-Bench/results/${experiment_name}/results_step${step}_ts${text_guidance_scale}_ig${image_guidance_scale}/viescore_gpt-4.1 \
+--language en

examples/OmniGen2-RL/options/omnigen2_edit_rl_4machine_editscore7b_avg4.yml

@@ -105,7 +105,7 @@ train:
     server_type: vlm
     use_ori_neg_prompt_template: true
     time_shift_base_res: 168
-    policy_loss_reweighting: false
+    policy_loss_reweighting: true
 
 val:
   train_visualization_interval: 5

examples/OmniGen2-RL/options/omnigen2_edit_rl_4machine_editscore7b_avg8.yml

@@ -105,7 +105,7 @@ train:
     server_type: vlm
     use_ori_neg_prompt_template: true
     time_shift_base_res: 168
-    policy_loss_reweighting: false
+    policy_loss_reweighting: true
 
 val:
   train_visualization_interval: 5

examples/OmniGen2-RL/options/omnigen2_edit_rl_single_machine_editscore7b.yml

@@ -105,7 +105,7 @@ train:
     server_type: vlm
     use_ori_neg_prompt_template: true
     time_shift_base_res: 168
-    policy_loss_reweighting: false
+    policy_loss_reweighting: true
 
 val:
   train_visualization_interval: 5

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "editscore"
-version = "0.1.9"
+version = "0.2"
 authors = [
   { name="Xin Luo", email="[email protected]" },
   { name="Jiahao Wang", email="[email protected]" },
@@ -27,6 +27,7 @@ dependencies = [
     "qwen-vl-utils",
     "peft",
     "Pillow",
+    "json-repair"
 ]
 
 [project.urls]