UniGenDet

A Unified Generative-Discriminative Framework for Co-Evolutionary Image Generation and Generated Image Detection

Yanran Zhang, Wenzhao Zheng^†, Yifei Li, Bingyao Yu, Yu Zheng, Lei Chen, Jie Zhou^*, Jiwen Lu
Department of Automation, Tsinghua University, China
^*Corresponding author    ^†Project leader

CVPR 2026

Overview

Image generation and generated-image detection have both advanced rapidly, but mostly along separate technical paths: generation is dominated by generative architectures, while detection is dominated by discriminative ones. This separation creates a persistent gap in practice: generators are not directly optimized by forensic criteria, and detectors are often trained on static snapshots of old forgeries, which limits robustness to new generators.

UniGenDet addresses this gap with a unified co-evolutionary framework that jointly optimizes generation and detection in one loop. The core idea is to make both tasks explicitly exchange useful signals instead of evolving independently.

• Symbiotic multimodal self-attention bridges generation and authenticity understanding in a shared architecture.
• Generation-detection unified fine-tuning (GDUF) equips the detector with generative priors, improving generalization and interpretability.
• Detector-informed generative alignment (DIGA) feeds authenticity constraints back into synthesis, improving realism and fidelity.

In short, UniGenDet turns the traditional "generator vs. detector" arms race into a closed-loop collaboration. This repository provides the full training and evaluation pipeline built on pretrained BAGEL components.

Installation

1. Clone and create environment

git clone https://github.com/Zhangyr2022/UniGenDet.git
cd UniGenDet

conda create -n unigendet python=3.10 -y
conda activate unigendet

pip install -r requirements.txt
pip install flash_attn==2.5.8 --no-build-isolation

2. Prepare pretrained checkpoints

Place required pretrained BAGEL weights in pretrained/.

# Optional: export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download ByteDance-Seed/BAGEL-7B-MoT --local-dir ./pretrained/bagel_7b_mot

Unified Demo 🔥

The repository provides demo.py for two interactive modes:

• t2i: text-to-image generation
• detection: input a real/fake image and output a detection sentence

Quick examples

Download the model from HuggingFace https://huggingface.co/Yanran21/UniGenDet or ModelScope https://www.modelscope.cn/models/YanranZhang/UniGenDet/summary

Parameter notes:

• --model_path: directory of the base BAGEL pretrained assets (e.g., llm_config.json, vit_config.json, ae.safetensors, tokenizer files).
• --ckpt_path: path to your fine-tuned checkpoint. You can pass a single .safetensors, an index json, or a checkpoint directory containing all sharded files.

Text-to-image:

python demo.py \
  --mode t2i \
  --model_path ./pretrained/bagel_7b_mot \
  --ckpt_path /path/to/your_sharded_ckpt_dir  \
  --prompt "A cinematic portrait of a snow fox under moonlight" \
  --resolution 1024 \
  --output_dir ./results/demo

AI-Generated Image Detection (checkpoint directory path):

python demo.py \
  --mode detection \
  --model_path ./pretrained/bagel_7b_mot \
  --ckpt_path /path/to/your_sharded_ckpt_dir \
  --image ./assets/example.jpg

The example image (assets/example.jpg) is generated by the latest GPT-Image-2 model, and UniGenDet produces the following detection output: "This is a fake image. The image shows a person sitting on a sidewalk with a backpack, but there are several unrealistic elements that suggest it is not a real photograph. The lighting and shadows do not align with a natural light source, and the colors appear overly saturated. The person's skin looks overly smooth and lacks natural texture, and the details of the clothing and backpack are not sharp and clear. The overall composition and quality of the image do not match what is typically seen in real photographs."

If you encounter CUDA OOM, you can reduce GPU memory usage by lowering max_memory_per_gpu (this forces more aggressive CPU offloading, so it runs slower but can fit smaller cards).

Data Preparation

1. Download datasets

We use a subset of LAION Aesthetics and FakeClue for training and evaluation.

Build LAION-style metadata first:

python scripts/data/laion_construction.py

Note that due to some links being inaccessible, approximately half of the LAION samples cannot be downloaded. We recommend using the remaining available samples for training.

Then download FakeClue:

# Optional: export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download lingcco/FakeClue --local-dir ./datasets/fakeclue

2. Configure dataset paths

Edit path placeholders in:

• data/dataset_info.py
• scripts/eval/*.sh

Use valid absolute paths in your local environment (for example, /path/to/datasets/...).

3. Configure dataset mixture

Use task YAMLs:

• data/configs/unigendet_DIGA.yaml
• data/configs/unigendet_GDUF.yaml

These files define dataset groups, sampling weights, and image transform settings.

Training

Before launch, set distributed arguments and checkpoint paths according to your cluster setup.

A. GDUF training (generation-detection unified fine-tuning)

torchrun \
  --nnodes=1 \
  --node_rank=0 \
  --nproc_per_node=8 \
  train/pretrain_unified_navit_gduf.py \
  --dataset_config_file ./data/configs/unigendet_GDUF.yaml \
  --model_path /path/to/project/pretrained \
  --layer_module Qwen2MoTDecoderLayer \
  --max_latent_size 64 \
  --finetune_from_hf True \
  --auto_resume True \
  --resume-model-only True \
  --finetune-from-ema True \
  --lr 2e-5 \
  --results_dir results_gduf \
  --checkpoint_dir gduf

B. DIGA training (detector-informed generative alignment)

torchrun \
  --nnodes=1 \
  --node_rank=0 \
  --nproc_per_node=8 \
  train/pretrain_unified_navit_diga.py \
  --dataset_config_file ./data/configs/unigendet_DIGA.yaml \
  --model_path /path/to/project/pretrained \
  --layer_module Qwen2MoTDecoderLayer \
  --max_latent_size 64 \
  --finetune_from_hf True \
  --auto_resume True \
  --resume-model-only True \
  --finetune-from-ema True \
  --lr 2e-5 \
  --visual_gen True \
  --visual_und False \
  --results_dir results_diga \
  --checkpoint_dir diga

C. Script-based launch

bash scripts/train/train_GDUF.sh
bash scripts/train/train_DIGA.sh

D. Practical notes

• Argument definitions follow the BAGEL training design.
• Default recipes target a single node with 8 GPUs (80 GB VRAM each).
• Adjust --nnodes, --nproc_per_node, and all path arguments to match your setup.
• If you hit out-of-memory issues, reduce --max_num_tokens, --expected_num_tokens, and related token limits.

Evaluation

1. Detection evaluation (FakeVLM)

The detection script evaluates UniGenDet on FakeVLM.

pip install rouge-score scikit-learn
bash scripts/eval/run_eval_fakevlm.sh

2. Generation evaluation on LAION-style prompts

Use a LAION split that does not overlap with your training subset to evaluate generalization.

pip install pytorch-fid
bash scripts/eval/run_laion.sh

3. GenEval benchmark

Before running, set up the environment by referring to BAGEL EVAL.md.

bash scripts/eval/run_geneval.sh

Key Implementation Notes

• max_latent_size=64 is recommended in the released BAGEL-based setup.
• Ensure num_used_data in YAML is larger than NUM_GPUS x NUM_WORKERS for stable sampling.
• For generation-only training, set visual_und=False.
• For memory-constrained debugging, reduce expected_num_tokens, max_num_tokens, and max_num_tokens_per_sample.

Acknowledgement

This project is built upon the open-source BAGEL ecosystem and related multimodal tooling. We thank the original authors and community contributors.

Citation

If you find this repository useful, please cite:

@article{zhang2026unigendet,
  title   = {UniGenDet: A Unified Generative-Discriminative Framework for Co-Evolutionary Image Generation and Generated Image Detection},
  author  = {Zhang, Yanran and Zheng, Wenzhao and Li, Yifei and Yu, Bingyao and Zheng, Yu and Chen, Lei and Zhou, Jie and Lu, Jiwen},
  journal = {CoRR},
  volume  = {abs/2604.21904},
  year    = {2026},
  url     = {https://arxiv.org/abs/2604.21904},
}

License

This repository follows the license terms specified in LICENSE.