repository·main·Indexed 26 days ago
https://github.com/vectorspacelab/omnigen2OmniGen2 is a multimodal generative model supporting visual understanding, text-to-image generation, instruction-guided image editing, and in-context generation using decoupled decoding pathways. The repository includes tools for the OmniContext benchmark, RL fine-tuning via the FlowGRPO algorithm, and the EditScore Reward Server for image editing optimization. It provides support for distributed multi-machine inference and Gradio-based local demos.
Use the omnicontext.test_omnicontext_score module to evaluate image quality via GPT-4.1. You must provide a valid OpenAI API key.
cd OmniGen2
openai_key="<Your-API-Key>"
python -m omnicontext.test_omnicontext_score \
--test_data "OmniGen2/OmniContext" \
--result_dir "omnicontext/results" \
--model_name "OmniGen2" \
--openai_key ${openai_key} \
--max_workers 100Before inference, you must convert the .bin checkpoint saved by FSDP into the standard Hugging Face format using convert_ckpt_to_hf_format.py.
For full fine-tuned models:
python convert_ckpt_to_hf_format.py \
--config_path experiments/ft/ft.yml \
--model_path experiments/ft/checkpoint-10/pytorch_model_fsdp.bin \
--save_path experiments/ft/checkpoint-10/transformerFor LoRA fine-tuned models:
python convert_ckpt_to_hf_format.py \
--config_path experiments/ft_lora/ft_lora.yml \
--model_path experiments/ft_lora/checkpoint-10/pytorch_model_fsdp.bin \
--save_path experiments/ft_lora/checkpoint-10/transformer_loraIf the generated image does not align well with the input image, use these strategies to improve subject consistency:
image_guidance_scale (e.g., to 3.0), though this may cause slight overexposure."she/he ..., maintaining her/his facial features, hairstyle, and other attributes."Follow these steps to set up a local environment for OmniGen2. It is recommended to use a clean Python environment (e.g., via Conda) and install flash-attn for optimal performance, although the model can run without it.
Prerequisites:
Steps:
requirements.txt.flash-attn (optional but recommended).# 1. Clone the repo
git clone git@github.com:VectorSpaceLab/OmniGen2.git
cd OmniGen2
# 2. (Optional) Create a clean Python environment
conda create -n omnigen2 python=3.11
conda activate omnigen2
# 3. Install dependencies
# 3.1 Install PyTorch (choose correct CUDA version)
pip install torch==2.6.0 torchvision --extra-index-url https://download.pytorch.org/whl/cu124
# 3.2 Install other required packages
pip install -r requirements.txt
# Note: Version 2.7.4.post1 is specified for compatibility with CUDA 12.4.
# OmniGen2 runs even without flash-attn, though we recommend install it for best performance.
pip install flash-attn==2.7.4.post1 --no-build-isolationOnce the checkpoint is converted to Hugging Face format, use the GEdit-Bench evaluation scripts. You must modify the internal paths in the scripts to point to your converted model checkpoint.
# Run evaluation for the converted model from step 500
bash evaluation/GEdit-Bench/omnigen2.sh --experiment_name=omnigen2_edit_rl_4machine_editscore7b_avg4 --step=500
bash evaluation/GEdit-Bench/omnigen2_eval.sh --experiment_name=omnigen2_edit_rl_4machine_editscore7b_avg4 --step=500For in-context editing (editing based on multiple images), use the following prompt format:
"Edit the first image: add/replace (the [object] with) the [object] from the second image. [description for your target image]."
Example:
"Edit the first image: add the man from the second image. The man is talking with a woman in the kitchen"
Ensure the description for your target image is as detailed as possible.
Training is handled via shell scripts that support PyTorch FSDP for multi-GPU and multi-node distributed training.
Scripts available:
scripts/train/ft.shscripts/train/ft_lora.shSingle-Node Training: Simply run the script without additional arguments:
bash scripts/train/ft.shMulti-Node / Multi-GPU Training: Provide environment variables to coordinate the distributed processes:
bash scripts/train/ft.sh --rank=$RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT --world_size=$WORLD_SIZENote: LoRA checkpoints currently save the entire model's parameters due to FSDP limitations; a conversion step is required before inference.
# Example for full-parameter fine-tuning
bash scripts/train/ft.sh --rank=$RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT --world_size=$WORLD_SIZEAfter running the scoring module, calculate the final benchmark statistics using the omnicontext.calculate_statistics module.
python -m omnicontext.calculate_statistics \
--save_path "omnicontext/results" \
--model_name "OmniGen2" \
--backbone gpt4dot1To use the OmniContext benchmark, activate your Python environment and install the required dependencies using pip.
# 1. Activate Python environment
conda activate omnigen2
# 2. Install dependencies
pip install -U datasets megfileRL fine-tuning is performed using the FlowGRPO algorithm.
Adjust parameters in options/omnigen2_edit_rl_4machine_editscore7b_avg4.yml. Key settings include:
train.global_batch_size: Total images generated across all GPUs in a sampling phase.train.batch_size: Batch size per GPU.train.rl.num_images_per_prompt: Candidate images per prompt.train.rl.num_unique_prompts_per_sampling: Unique prompts in a global batch.train.rl.num_update_steps_per_sampling: Gradient updates per sampling phase (set $> 1$ for off-policy RL).train.rl.batch_size_per_forward: Batch size for each forward pass.Launch training using FSDP-based scripts:
scripts/train/omnigen2_edit_rl_single_machine_editscore7b.shscripts/train/omnigen2_edit_rl_4machine_editscore7b_avg4.sh# Single-machine training (8 GPUs)
bash scripts/train/omnigen2_edit_rl_single_machine_editscore7b.sh
# Multi-machine training
bash scripts/train/omnigen2_edit_rl_4machine_editscore7b_avg4.shOmniGen2 natively requires an NVIDIA RTX 3090 or an equivalent GPU with approximately 17GB of VRAM.
If you have less VRAM, you must enable CPU Offload (enable_model_cpu_offload) to run the model.