TTTIR: Unlocking Instance-Specific State Evolution via Test-Time Training for Image Restoration

Kaihang Zheng^*, Jun Li^*, Hang Guo, Hongyu Chi, Zimo Liu, Tao Dai, Jinpeng Wang^†, and Yaowei Wang^†

⭐ If TTTIR is helpful to your projects, please help star this repo. Thanks! 🤗

We sincerely invite readers to refer to our team’s other work MambaIR.

News

• 2026-08-05: The initial codebase is released. We are excited to release the first general image reconstruction model based on TTT. Our model achieves SOTA performance with fewer than 1M parameters.
• 2026-08-26: For easier reproducibility, we are working to unify the codebase across all tasks.

Abstract: Image restoration is inherently challenging due to the diverse and highly input-dependent nature of real-world degradations. While recent architectures such as Transformers and state-space models have advanced the field, they predominantly rely on static, globally shared parameters, which struggle to fully accommodate instance-specific degradation patterns. To address this limitation, we propose TTTIR, a framework that reformulates image restoration as an instance-specific state evolution process. Progressive State Sequence Generation (PSSG) constructs complementary spatial-frequency target states that define what to recover, while State Transition Evolution (STE) adapts lightweight transition operators through a restoration-oriented test-time training inner loop that determines how features should evolve. Extensive experiments demonstrate strong performance across low-light enhancement, rain removal, and image dehazing benchmarks with favorable computational scalability.

Contents

1. Dependencies
2. Datasets
3. Training
4. Testing
5. Model Zoo
6. Results

Dependencies

Our experiments in the paper are conducted on NVIDIA RTX 3080 Ti GPUs.

git clone https://github.com/Elysiaaaaaaaa/TTTIR.git
cd TTTIR

conda create -n tttir python=3.7.12 -y
conda activate tttir

# The dependency file will be released soon.
pip install -r requirements.txt

Datasets

The datasets used for training and evaluation are listed below. They can be downloaded from Baidu Netdisk (extraction code: 7ghe).

Task	Training set	Testing set
Low-light enhancement	LOL-v1	LOL-v1, LOL-v2-Real
Low-light enhancement	LOL-v2-Synthetic	LOL-v2-Synthetic
Rain streak removal	Rain13K	Test100, Rain100H, Rain100L, Test1200, Test2800
Raindrop removal	Raindrop training set	Raindrop-A, Raindrop-B
Image dehazing	RESIDE-6K	RESIDE-6K
Image dehazing	Haze4K	Haze4K

Organize the datasets under a common root directory:

Click to expand the dataset directory structure

<dataset_root>/
|-- LOLv1/
|   |-- Train/
|   |   |-- input/
|   |   `-- target/
|   `-- Test/
|       |-- input/
|       `-- target/
|-- LOLv2/
|   |-- Real_captured/
|   |   |-- Train/
|   |   |   |-- Low/
|   |   |   `-- Normal/
|   |   `-- Test/
|   |       |-- Low/
|   |       `-- Normal/
|   `-- Synthetic/
|       |-- Train/
|       |   |-- Low/
|       |   `-- Normal/
|       `-- Test/
|           |-- low1/
|           `-- Normal/
|-- Rain13k/
|   |-- train/
|   |   |-- input/
|   |   `-- target/
|   `-- test/
|       |-- Rain100H/
|       |   |-- input/
|       |   `-- target/
|       |-- Rain100L/
|       |   |-- input/
|       |   `-- target/
|       |-- Test100/
|       |   |-- input/
|       |   `-- target/
|       |-- Test1200/
|       |   |-- input/
|       |   `-- target/
|       `-- Test2800/
|           |-- input/
|           `-- target/
|-- raindrop/
|   |-- train/
|   |   |-- data/
|   |   `-- gt/
|   |-- test_a/
|   |   |-- data/
|   |   `-- gt/
|   `-- test_b/
|       |-- data/
|       `-- gt/
|-- RESIDE-6K/
|   |-- train/
|   |   |-- haze/
|   |   `-- gt/
|   `-- test/
|       |-- haze/
|       `-- gt/
`-- Haze4K/
    |-- train/
    |   |-- haze/
    |   |-- gt/
    |   `-- trans/
    `-- test/
        |-- haze/
        |-- gt/
        `-- trans/

Training

All runnable restoration commands use the unified entry point in restoration/. For a detached job, enter screen -S tttir, run conda activate tttir inside the session, start one of the commands below, then press Ctrl-A D to detach.

Rain-streak removal

cd restoration
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node 4 --use_env --master_port 7298  main.py --model_name rain13k --task derain --mode train --num_epoch 300 --data_dir /path/to/datasets/Rain13k --learning_rate 1e-3  --save_freq 30 --valid_freq 1 --batch_size 4 --num_worker 4

The current single-process command is the supported documented path. Select an idle GPU by changing CUDA_VISIBLE_DEVICES.

Low-light enhancement

Training settings and dataset paths are defined in YAML files under restoration/options/train/. Before running, set dataroot_LQ and dataroot_GT in the selected YAML. Choose the device through the process environment, for example CUDA_VISIBLE_DEVICES=0.

cd restoration

# LOL-v1
CUDA_VISIBLE_DEVICES=0 python main.py --task enhance --mode train --opt options/train/lol_v1.yml

# LOL-v2-Synthetic
CUDA_VISIBLE_DEVICES=0 python main.py --task enhance --mode train --opt options/train/lol_v2_synthetic.yml

--opt must be supplied explicitly for low-light enhancement.

Testing

Rain-streak removal

cd restoration
CUDA_VISIBLE_DEVICES=0 python main.py --task derain --mode test \
    --data_dir /path/to/datasets/Rain13k --test_model /path/to/model.pkl \
    --model_name rain100l

Low-light enhancement

cd restoration

# LOL-v1
CUDA_VISIBLE_DEVICES=0 python main.py --task enhance --mode test --opt options/test/lol_v1.yml

# LOL-v2-Real evaluated with the LOL-v1 model
CUDA_VISIBLE_DEVICES=0 python main.py --task enhance --mode test --opt options/test/lol_v2_real_lol_v1.yml

# LOL-v2-Synthetic
CUDA_VISIBLE_DEVICES=0 python main.py --task enhance --mode test --opt options/test/lol_v2_synthetic.yml

Model Zoo

Task	Training set	Evaluation set(s)	Checkpoint
Low-light enhancement	LOL-v1	LOL-v1, LOL-v2-Real	Models & results (code: a4r4)
Low-light enhancement	LOL-v2-Synthetic	LOL-v2-Synthetic	Models & results (code: a4r4)
Rain streak removal	Rain13K	Test100, Rain100H, Rain100L, Test1200, Test2800	Models & results (code: 8jf2)
Raindrop removal	Raindrop training set	Raindrop-A, Raindrop-B	Models & results (code: hs0z)
Image dehazing	RESIDE-6K	RESIDE-6K	Coming soon
Image dehazing	Haze4K	Haze4K	Coming soon

Results

Low-light enhancement

Rain streak removal

Raindrop removal	Image dehazing
Table 3. Raindrop-A and Raindrop-B.	Table 4. RESIDE-6K and Haze4K.

Low-light enhancement

Rain removal

Citation

If you find our code useful or use the toolkit in your work, please consider citing:

@article{zheng2026tttir,
  title   = {TTTIR: Unlocking Instance-Specific State Evolution via Test-Time Training for Image Restoration},
  author  = {Zheng, Kaihang and Li, Jun and Guo, Hang and Chi, Hongyu and Liu, Zimo and Dai, Tao and Wang, Jinpeng and Wang, Yaowei},
  journal = {arXiv preprint},
  year    = {2026}
}