[Arxiv 2026] Universal Image Restoration via Internalized
Chain-of-Thought Reasoning

Universal Image Restoration via Internalized Chain-of-Thought Reasoning
Yu Guo^† , Zhengru Fang^† , Shengfeng He, Senkang Hu, Yihang Tao, Phone Lin, Yuguang Fang
(† Co-first Author)
Arxiv

Abstract: Image restoration seeks to recover high-quality images from degraded inputs but becomes highly ill-posed under complex, mixed degradations. While unified all-in-one models are common, their performance declines as degradation complexity increases. Recent works adopt Chain-of-Thought (CoT) reasoning for multi-round restoration using specialized modules. However, this approach faces two key limitations: (i) increased computational cost due to multi-step processing and (ii) weak modeling of interactions between degradations during stepwise inference. We introduce CoTIR, a universal image restoration framework that internalizes CoT reasoning within a single model. Concretely, we view image restoration as a specialized subtask of image editing, which implies that a large-scale pre-trained editing model provides a more favorable optimization starting point. Building on this, we fine-tune the model for restoration and further encode structured CoT-style reasoning into the learning objective via a differentiable formulation inspired by Lagrangian optimization, enabling holistic restoration without chaining specialized restorers. To facilitate training and evaluation, we further present CoTIR-Bench, a large-scale benchmark comprising 5.2 million samples with CoT-style reasoning traces. Extensive experiments on CoTIR-Bench and broad real composite degradation scenes show that CoTIR achieves stronger perceptual quality and more competitive fidelity than both all-in-one models and multi-round restoration methods.

• 2026.06.16: Code, Checkpoint, and Test dataset are released.

• python 3.10
• cuda 12.1

# git clone this repository
git clone https://github.com/gy65896/CoTIR.git
cd CoTIR

# create new anaconda env
conda env create -f environment.yml

CoTIR checkpoints only contain the LoRA and CoT adapter weights. You also need to download the corresponding base FLUX models from Black Forest Labs (accept the license on Hugging Face before downloading):

Download the base models:

git clone https://huggingface.co/black-forest-labs/FLUX.2-klein-4B
git clone https://huggingface.co/black-forest-labs/FLUX.2-klein-9B
git clone https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev

Download the CoTIR LoRA and adapter weights and place them under ./ckpt/:

git clone https://huggingface.co/gy65896/CoTIR

Before running inference or testing, update the checkpoint paths in configs/test_cotir-4b.yaml, configs/test_cotir-9b.yaml, or configs/test_cotir-12b.yaml:

inference:
  num_steps: 5
  base_model_path: /path/to/FLUX.2-klein-9B      # base FLUX model
  cot_adapter_path: ./ckpt/CoTIR-9B/cot_adapter.pt  # CoTIR adapter
  lora_path: ./ckpt/CoTIR-9B/lora.pt                # CoTIR LoRA

CoTIR-Bench will be released soon. After it is available, update data.data_dir and model.base_model_path in the training config, then run:

python train.py --config configs/train_cotir-9b.yaml

Multi-GPU training (example with 8 GPUs):

accelerate launch --num_processes 8 train.py --config configs/train_cotir-9b.yaml

Available training configs: configs/train_cotir-4b.yaml, configs/train_cotir-9b.yaml, configs/train_cotir-12b.yaml.

Checkpoints are saved to saves/ by default. Set resume: latest in the config to resume from the latest checkpoint.

Batch inference on a benchmark JSONL file. Outputs are saved to vague/ (vague prompt) and precise/ (precise prompt):

python test.py \
  --model CoTIR-9B \
  --json_path ./test/test.jsonl \
  --lq_dir ./test/lq \
  --output_root ./results

Optional arguments:

python test.py \
  --config configs/test_cotir-9b.yaml \
  --json_path ./test/test.jsonl \
  --lq_dir ./test/lq \
  --output_root ./results \
  --batch_size 1 \
  --max_long_edge 1024 \
  --max_samples 100 \
  --skip_existing

Compute folder-level mean metrics with evaluate.py. Result images are matched to HQ by filename, and the summary is saved as metrics_summary.txt in each evaluated folder.

Metrics: clipiqa_plus, qalign, liqe, maclip, clip_iqa, psnr, ssim, lpips.

Evaluate a single result folder:

python evaluate.py \
  --hq-dir /path/to/hq \
  --method-root /path/to/results/method/

Adjust paths for your local CoTIR-Bench results directory:

python evaluate.py \
  --hq-dir /path/to/hq \
  --method-root /path/to/results/

Explanation: images under <method>/precise/ (or <method>/) will be matched with test/hq/ by filename. Each evaluated image folder will generate metrics_summary.txt.

Expected layout:

method/
├── CoTIR-FLUX.1-12B/
│   ├── precise/
│   └── vague/
└── PromptIR/
    ├── precise/
    └── vague/

Output example (metrics_summary.txt):

method & clipiqa_plus & qalign & liqe & maclip & clip_iqa & psnr & ssim & lpips
CoTIR-FLUX.1-12B & 0.6137 & 3.7834 & 3.5136 & 0.5370 & 0.5011 & 24.36 & 0.8015 & 0.1243

Note: Locally reproduced results may show minor deviations from the numbers reported in the paper. To obtain the exact results used in our evaluation, please download the pre-computed outputs from the corresponding branches of CoTIR-Bench:

Paper Table	Branch	Download Command
Table 2	methods	git clone -b methods --single-branch https://huggingface.co/datasets/gy65896/CoTIR-Bench
Table 3	single	git clone -b single --single-branch https://huggingface.co/datasets/gy65896/CoTIR-Bench
Table 4	cotir-bench-7d	git clone -b cotir-bench-7d --single-branch https://huggingface.co/datasets/gy65896/CoTIR-Bench

Single-image restoration:

python inference.py \
  --config configs/test_cotir-9b.yaml \
  --lq ./path/to/lq.png \
  --prompt "Make this picture clearer." \
  --output_root outputs \
  --num_steps 5 \
  --seed 42

You can also pass a prompt text file:

python inference.py \
  --config configs/test_cotir-9b.yaml \
  --lq ./samples/images/000018.png \
  --prompt ./samples/prompts/000018.txt \
  --output_root outputs

Launch the interactive demo:

python inference_gradio.py --port 7862

Create a public link (optional):

python inference_gradio.py --port 7862 --share

Usage:

1. Enter the GPU id(s) to use and click Init Model.
2. Select a model (CoTIR-4B / CoTIR-9B / CoTIR-12B).
3. Upload an LQ image, optionally upload an HQ reference image, and enter a prompt.
4. Adjust num_steps and other options if needed, then click Run Inference.

Results are saved under the output directory shown in the UI.

@misc{guo2026universalimagerestorationinternalized,
      title={Universal Image Restoration via Internalized Chain-of-Thought Reasoning}, 
      author={Yu Guo and Zhengru Fang and Shengfeng He and Senkang Hu and Yihang Tao and Phone Lin and Yuguang Fang},
      year={2026},
      eprint={2606.17557},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2606.17557}, 
}