📖 Introduction
Recent advancements in image customization exhibit a wide range of application prospects due to stronger customization capabilities. However, since we humans are more sensitive to faces, a significant challenge remains in preserving consistent identity while avoiding identity confusion with multi-reference images, limiting the identity scalability of customization models. To address this, we present UMO, a Unified Multi-identity Optimization framework, designed to maintain high-fidelity identity preservation and alleviate identity confusion with scalability. With "multi-to-multi matching" paradigm, UMO reformulates multi-identity generation as a global assignment optimization problem and unleashes multi-identity consistency for existing image customization methods generally through reinforcement learning on diffusion models. To facilitate the training of UMO, we develop a scalable customization dataset with multi-reference images, consisting of both synthesised and real parts. Additionally, we propose a new metric to measure identity confusion. Extensive experiments demonstrate that UMO not only improves identity consistency significantly, but also reduces identity confusion on several image customization methods, setting a new state-of-the-art among open-source methods along the dimension of identity preserving.
⚡️ Quick Start
🔧 Requirements and Installation
# 1. Clone the repo with submodules: UNO & OmniGen2
git clone --recurse-submodules git@github.com:bytedance/UMO.git
cd UMO
UMO requirements based on UNO
# 2.1 (Optional, but recommended) Create a clean virtual Python 3.11 environment
python3 -m venv venv/UMO_UNO
source venv/UMO_UNO/bin/activate
# 3.1 Install submodules UNO requirements as:
# https://github.com/bytedance/UNO?tab=readme-ov-file#-requirements-and-installation
# 4.1 Install UMO requirements
pip install -r requirements.txt
UMO requirements based on OmniGen2
# 2.2 (Optional, but recommended) Create a clean virtual Python 3.11 environment
python3 -m venv venv/UMO_OmniGen2
source venv/UMO_OmniGen2/bin/activate
# 3.2 Install submodules OmniGen2 requirements as:
# https://github.com/VectorSpaceLab/OmniGen2?tab=readme-ov-file#%EF%B8%8F-environment-setup
# 4.2 Install UMO requirements
pip install -r requirements.txt
UMO checkpoints download
# pip install huggingface_hub hf-transfer
export HF_HUB_ENABLE_HF_TRANSFER=1 # use hf_transfer to speedup
# export HF_ENDPOINT=https://hf-mirror.com # use mirror to speedup if necessary
repo_name="bytedance-research/UMO"
local_dir="models/"$repo_name
huggingface-cli download --resume-download $repo_name --local-dir $local_dir
🌟 Gradio Demo
# UMO (based on UNO)
python3 demo/UNO/app.py --lora_path models/bytedance-research/UMO/UMO_UNO.safetensors
# UMO (based on OmniGen2)
python3 demo/OmniGen2/app.py --lora_path models/bytedance-research/UMO/UMO_OmniGen2.safetensors
✍️ Inference
UMO (based on UNO) inference on XVerseBench
# single subject
accelerate launch eval/UNO/inference_xversebench.py \
--eval_json_path projects/XVerse/eval/tools/XVerseBench_single.json \
--num_images_per_prompt 4 \
--width 768 \
--height 768 \
--save_path output/XVerseBench/single/UMO_UNO \
--lora_path models/bytedance-research/UMO/UMO_UNO.safetensors
# multi subject
accelerate launch eval/UNO/inference_xversebench.py \
--eval_json_path projects/XVerse/eval/tools/XVerseBench_multi.json \
--num_images_per_prompt 4 \
--width 768 \
--height 768 \
--save_path output/XVerseBench/multi/UMO_UNO \
--lora_path models/bytedance-research/UMO/UMO_UNO.safetensors
UMO (based on UNO) inference on OmniContext
accelerate launch eval/UNO/inference_omnicontext.py \
--eval_json_path OmniGen2/OmniContext \
--width 768 \
--height 768 \
--save_path output/OmniContext/UMO_UNO \
--lora_path models/bytedance-research/UMO/UMO_UNO.safetensors
UMO (based on OmniGen2) inference on XVerseBench
# single subject
accelerate launch -m eval.OmniGen2.inference_xversebench \
--model_path OmniGen2/OmniGen2 \
--model_name UMO_OmniGen2 \
--test_data projects/XVerse/eval/tools/XVerseBench_single.json \
--result_dir output/XVerseBench/single \
--num_images_per_prompt 4 \
--disable_align_res \
--lora_path models/bytedance-research/UMO/UMO_OmniGen2.safetensors
# multi subject
accelerate launch -m eval.OmniGen2.inference_xversebench \
--model_path OmniGen2/OmniGen2 \
--model_name UMO_OmniGen2 \
--test_data projects/XVerse/eval/tools/XVerseBench_multi.json \
--result_dir output/XVerseBench/multi \
--num_images_per_prompt 4 \
--disable_align_res \
--lora_path models/bytedance-research/UMO/UMO_OmniGen2.safetensors
UMO (based on OmniGen2) inference on OmniContext
accelerate launch -m eval.OmniGen2.inference_omnicontext \
--model_path OmniGen2/OmniGen2 \
--model_name UMO_OmniGen2 \
--test_data OmniGen2/OmniContext \
--result_dir output/OmniContext \
--num_images_per_prompt 1 \
--disable_align_res \
--lora_path models/bytedance-research/UMO/UMO_OmniGen2.safetensors
📌 Tips and Notes
Please note that UNO gets unstable results on parts of OmniContext due to the different prompt format with its training data (UNO-1M), leading to similar issue with UMO based on it. To get better results with these two models, we recommend using description prompt instead of instruction one, using resolution 768~1024 instead of 512.
📄 Disclaimer
Citation
If UMO is helpful, please help to ⭐ the repo.
If you find this project useful for your research, please consider citing our paper:
@article{cheng2025umo,
title={UMO: Scaling Multi-Identity Consistency for Image Customization via Matching Reward},
author={Cheng, Yufeng and Wu, Wenxu and Wu, Shaojin and Huang, Mengqi and Ding, Fei and He, Qian},
journal={arXiv preprint arXiv:2509.06818},
year={2025}
}