UniCoT-7B-MoT

<p align="center"> <img src="https://github.com/Fr0zenCrane/UniCoT/raw/v0.1/assets/logo.png" alt="Uni-CoT" width="400"/> </p>

Uni-CoT: Towards Unified Chain-of-Thought Reasoning Across Text and Vision

This repository contains the UniCoT-7B-MoT model, a unified Chain-of-Thought (CoT) framework for multimodal reasoning across text and vision. It was introduced in the paper Uni-cot: Towards Unified Chain-of-Thought Reasoning Across Text and Vision.

Project Page: https://sais-fuxi.github.io/projects/uni-cot/ Code: https://github.com/Fr0zenCrane/UniCoT

<p align="center"> <img src="https://github.com/Fr0zenCrane/UniCoT/raw/v0.1/assets/teaser.png" width="900"/> </p>

Overview

Chain-of-Thought (CoT) reasoning has been widely adopted to enhance Large Language Models (LLMs) by decomposing complex tasks into simpler, sequential subtasks. Uni-CoT extends these principles to vision-language reasoning, enabling coherent and grounded multimodal reasoning within a single unified model. The key idea is to leverage a model capable of both image understanding and generation to reason over visual content and model evolving visual states.

The Uni-CoT framework adopts a novel two-level hierarchical reasoning architecture:

1. Macro-Level CoT: Decomposes a complex task into simpler subtasks and synthesizes their outcomes to derive the final answer. This includes strategies like Sequential, Parallel, and Progressive Refinement Decomposition.
2. Micro-Level CoT: Focuses on executing individual subtasks, incorporating a Self-Check (Self-Reflection) Mechanism to ensure stable and high-quality results.

This design significantly reduces computational overhead, allowing Uni-CoT to perform scalable and coherent multi-modal reasoning. It aims to solve complex multimodal tasks, including:

• 🎨 Reliable image generation and editing
• 🔍 Visual and physical reasoning
• 🧩 Visual planning
• 📖 Multimodal story understanding

🧠 Reasoning Pipeline

<p align="center"> <img src="https://github.com/Fr0zenCrane/UniCoT/raw/v0.1/assets/pipeline.png" width="900"/> </p>

Quickstart

Installation

The environment setup of Uni-CoT is consistent with its base model, Bagel.

git clone https://github.com/Fr0zenCrane/UniCoT.git
cd UniCoT
conda create -n unicot python=3.10 -y
conda activate unicot
pip install -r requirements.txt
pip install flash_attn==2.5.8 --no-build-isolation

Model Download

You may directly download the huggingface checkpoint or use the following script:

from huggingface_hub import snapshot_download

save_dir = "models/UniCoT-7B-MoT"
repo_id = "Fr0zencr4nE/UniCoT-7B-MoT"
cache_dir = save_dir + "/cache"

snapshot_download(cache_dir=cache_dir,
  local_dir=save_dir,
  repo_id=repo_id,
  local_dir_use_symlinks=False,
  resume_download=True,
  allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
)

Self-check Reasoning

To perform evaluation or general inference using UniCoT-7B-MoT, you need at least one GPU with 40GB or more VRAM.

Evaluation on WISE benchmark

To reproduce the results on WISE benchmark, you can use the script ./scripts/run_wise_self_reflection.sh. Specify your local checkpoint of UniCoT-7B-MoT and output directory using --model_path and outdir.

gpu_num=8

for i in $(seq 0 $((gpu_num-1)));
do
    CUDA_VISIBLE_DEVICES=$i python inference_mdp_self_reflection_wise.py \
        --group_id $i \
        --group_num $gpu_num \
        --model_path "Fr0zencr4nE/UniCoT-7B-MoT" \
        --data_path "./eval/gen/wise/final_data.json" \
        --outdir "./results" \
        --cfg_text_scale 4 > process_log_$i.log 2>&1 &
done

wait
echo "All background processes finished."

General Inference

For general inference, prepare your prompts by formatting them into a .txt file, with one prompt per line (e.g., test_prompts.txt). Then, use the script ./scripts/run_user_self_reflection.sh to generate images from your prompts with the added benefit of the self-reflection mechanism.

gpu_num=8

for i in $(seq 0 $((gpu_num-1)));
do
    CUDA_VISIBLE_DEVICES=$i python inference_mdp_self_reflection.py \
        --group_id $i \
        --group_num $gpu_num \
        --model_path "Fr0zencr4nE/UniCoT-7B-MoT" \
        --data_path "./test_prompts.txt" \
        --outdir "./results" \
        --cfg_text_scale 4 > process_log_$i.log 2>&1 &
done

wait
echo "All background processes finished."

Preliminary Results

Qualitative Results for Image Generation

<p align="left"> <img src="https://github.com/Fr0zenCrane/UniCoT/raw/v0.1/assets/qualitative_results_gen.png" width="800"/>

Qualitative Results for Image Editing

<p align="left"> <img src="https://github.com/Fr0zenCrane/UniCoT/raw/v0.1/assets/qualitative_results_edit.png" width="800"/>

Quantitative Results on WISE

We first conduct experiments on the WISE dataset to evaluate the reasoning capabilities of our method. As shown in the table below, our model achieves state-of-the-art (SOTA) performance among existing open-source unified models. Our results are averaged over five independent runs to ensure robustness and reliability.

Furthermore, we apply our self-check mechanism to the images generated by the original Bagel model with think mode, aiming to evaluate our method’s ability to calibrate erroneous outputs. The results in the table below demonstrate that our model effectively refines the imperfect outputs generated by Bagel.

Quantitative Results on KRIS Bench

We also achieve state-of-the-art (SOTA) performance on the KRIS benchmark, even surpassing the closed-source model Gemini2.0.

Citation

If you find Uni-CoT useful for your research, please consider citing the paper:

@misc{qin2025unicot,
      title={Uni-cot: Towards Unified Chain-of-Thought Reasoning Across Text and Vision}, 
      author={Luozheng Qin and Jia Gong and Yuqing Sun and Tianjiao Li and Mengping Yang and Xiaomeng Yang and Chao Qu and Zhiyu Tan and Hao Li},
      year={2025},
      eprint={2508.05606},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2508.05606}, 
}

Acknowledgement

We acknowledge the contributions of the following projects that inspired and supported Uni-CoT:

• Bagel proposed by ByteDance-Seed team.
• WISE proposed by PKU-YuanGroup.
• KRIS-Bench proposed by Stepfun.
• RISE-Bench proposed by Shanghai AI Lab.