DIM-4.6B-Edit

[ICLR 2026] Draw-In-Mind: Rebalancing Designer-Painter Roles in Unified Multimodal Models Benefits Image Editing

Ziyun Zeng, David Junhao Zhang, Wei Li, and Mike Zheng Shou

📰 News

• [2026-05-12] The DIM project page is available.
• [2026-01-26] 🎉 DIM is accepted to ICLR 2026!
• [2025-10-08] 🚀 Released the DIM-Edit dataset and the DIM-4.6B-T2I / DIM-4.6B-Edit models.
• [2025-09-02] 📝 The DIM paper is released on arXiv.

🌟 Highlights

• 🧠 Rebalanced architecture: Let the understanding module be the designer, while the generation module focuses on painting.
• 📚 Two complementary datasets: DIM-T2I (long-context T2I pairs) and DIM-Edit (CoT imaginations from GPT-4o).
• ⚡ Lightweight & efficient: A ❄️frozen 3.0B VLM and a 🔥trainable 1.6B DiT connected via a single MLP (4.6B params in total).
• 🏆 SOTA-competitive: DIM-4.6B-Edit matches or surpasses much larger models on ImgEdit and GEdit-Bench.

💡 Introduction

Unified models achieve strong results in text-to-image generation but remain weak in precise editing. This limitation arises from an imbalanced division of responsibilities. The understanding module is usually treated as a translator that encodes instructions into conditions, while the generation module must act as both designer and painter. The result is that the generation module carries too much responsibility, even though it is not optimized for complex reasoning.

To address this, we introduce Draw-In-Mind (DIM), a dataset with two complementary parts:

• 🖼️ DIM-T2I: Millions of long-context image–text pairs that strengthen instruction comprehension.
• ✏️ DIM-Edit: 233K chain-of-thought imaginations from GPT-4o that provide explicit design blueprints.

We connect a frozen Qwen2.5-VL-3B with a trainable SANA1.5-1.6B via a lightweight MLP, forming DIM-4.6B-T2I/Edit. With this setup, the understanding module takes on the designer responsibility, while the generation module focuses on rendering. Despite its modest size, DIM-4.6B-Edit achieves SOTA or competitive results on ImgEdit and GEdit-Bench, outperforming much larger models.

📊 Performance

<details> <summary><b>📈 GenEval & MJHQ-30K</b></summary>

^† denotes using an LLM rewriter. For MJHQ(-30K), we report FID.

</details> <details> <summary><b>🖌️ ImgEdit Overall</b></summary>

Q3/7B indicates using Qwen2.5-VL-3/7B as the external designer during inference. By default, GPT-4o is employed as the external designer to ensure the best performance. All models are evaluated using GPT-4.1.

</details> <details> <summary><b>🔬 ImgEdit Designer Ablation</b></summary>

^† The default setting.

</details> <details> <summary><b>🖼️ Qualitative Visualization</b></summary>

🟢 Green and 🔵 Blue denote the edits of Janus-4o and Step1X-Edit respectively; 🔴 Red denotes the edits of our models trained on different data corpora.

</details>

📦 Dataset

DIM-Edit

Step 1. Download DIM-Edit from our 🤗 HF repo using the hf CLI:

# 1. Install the huggingface_hub library (>= 0.32.0 for hf_xet support)
pip install -U huggingface_hub

# 2. Log in with your Hugging Face account token
hf auth login

# 3. Download the dataset
hf download stdKonjac/DIM-Edit --repo-type dataset --local-dir ./DIM-Edit

Step 2. Merge and extract the split archives:

cd DIM-Edit
cat images.tar.gz.part* > images.tar.gz
tar -xvzf images.tar.gz

Step 3. Each line of tos_dataset_edit.jsonl corresponds to a single sample with four fields:

Step 4. Load the dataset with the 🤗 datasets library:

from datasets import load_dataset, Features, Value

features = Features({
    "id": Value("string"),
    "image_path": Value("string"),
    "image_path_target": Value("string"),
    "prompt": Value("string"),
})

ds = load_dataset(
    "json",
    data_files="DIM-Edit/tos_dataset_edit.jsonl",
    features=features,
    split="train",
)

print(ds[0])

📜 DIM-Edit License

The DIM-Edit dataset is released under the CC-BY-NC 4.0 license.

DIM-T2I

Please refer to T2I_DATASET.md for download instructions and licensing details.

🚀 Model

⚙️ Environment Setup

pip install -r requirements.txt

🦙 Model Zoo

Create a checkpoints folder in the root directory, then download the models from our 🤗 HF repo and move them into checkpoints/.

mkdir checkpoints

💡 To facilitate reproducibility, we release DIM-4.6B-Edit-Stage1, which is trained solely on the UltraEdit dataset. Fine-tuning this checkpoint on our proposed DIM-Edit dataset should reproduce DIM-4.6B-Edit.

Organize the checkpoints as follows:

DIM/
└── checkpoints/
    ├── DIM-4.6B-T2I/
    │   ├── model.safetensors
    │   └── ...
    ├── DIM-4.6B-Edit-Stage1/
    │   ├── model.safetensors
    │   └── ...
    └── DIM-4.6B-Edit/
        ├── model.safetensors
        └── ...

🔮 Inference

<details> <summary><b>🎨 T2I Generation</b></summary>

Demo T2I instructions are provided in cache/demo/tos_dataset_demo.jsonl. Each line is a JSON instruction, e.g.:

{
  "id": "0000",
  "image_path": "./cache/demo/edit_demo_0000.png",
  "prompt": "A yummy cupcake floating in the air dark background"
}

The image_path is a placeholder — modify prompt to generate your own image.

Run:

bash scripts/demo_t2i.sh

Generated images will be saved to cache/inference/demo/DIM-4.6B-T2I/{id}_gen.jpg.

</details> <details> <summary><b>✂️ Image Editing</b></summary>

Demo edit instructions are provided in cache/demo/tos_dataset_edit_demo.jsonl. Each line looks like:

{
  "id": "0",
  "image_path": "./cache/demo/edit_demo_0000.png",
  "prompt": "Remove the lemons on the table.",
  "image_path_target": "./cache/demo/edit_demo_0000.png"
}

image_path is the source image and prompt is the edit instruction; image_path_target is a placeholder.

In infer/demo_edit.py, use the set_designer_gpt API with your own key to set GPT-4o as the external designer for optimal performance:

# GPT-4o as external designer
model.set_designer_gpt(api_key=os.environ['OPENAI_API_KEY'])

Alternatively, use set_designer_X APIs for open-source VLMs (auto-downloaded to local disk):

# Qwen2.5-VL as external designer
model.set_designer_qwen(version='Qwen/Qwen2.5-VL-3B-Instruct')
model.set_designer_qwen(version='Qwen/Qwen2.5-VL-7B-Instruct')

# InternVL3.5 as external designer (recommend using transformers==4.53.0)
model.set_designer_internvl(version='OpenGVLab/InternVL3_5-8B-HF')

# MiMo-VL as external designer
model.set_designer_mimo(version='XiaomiMimo/MiMo-VL-7B-RL-2508')

# GLM-4.1V as external designer (recommend using transformers==4.53.1)
model.set_designer_glm(version='THUDM/GLM-4.1V-9B-Thinking')

Run:

bash scripts/demo_edit.sh

The model first generates a CoT-guided edit instruction for each prompt (saved to cache/inference/demo/DIM-4.6B-Edit/tos_dataset_edit_cot_demo_gen.jsonl), then produces edited images at cache/inference/demo/DIM-4.6B-Edit/{id}_edited.jpg.

A sample GPT-4o-generated CoT jsonl is provided at cache/demo/tos_dataset_edit_cot_demo.jsonl for reference.

</details>

📜 Model License

The models are developed based on Qwen2.5-VL-3B-Instruct (subject to the Qwen RESEARCH LICENSE AGREEMENT) and SANA1.5_1.6B_1024px (subject to the NVIDIA License). We retain ownership of all intellectual property rights in and to any derivative works and modifications that we made.

🧪 Evaluation

<details> <summary><b>📐 GenEval</b></summary>

We provide two evaluation jsonl files in cache/GenEval based on prompt type:

1. tos_dataset.jsonl — Original prompts.
2. tos_dataset_rewritten.jsonl — LLM-rewritten prompts.

The image_path field is a placeholder — please replace it with a pseudo image on your local disk first.

Run:

bash scripts/eval_geneval.sh

Generated images will be saved to cache/inference/DIM-4.6B-T2I/GenEval(_rewritten). Follow the GenEval official repo for metric calculation.

</details> <details> <summary><b>🖼️ MJHQ-30K</b></summary>

Download MJHQ-30K (only mjhq30k_imgs.zip is needed), extract under cache/ as:

cache
└── MJHQ-30K
    ├── animals
    │   ├── {id}.jpg
    │   └── ...
    ├── art
    ├── fashion
    ├── food
    ├── indoor
    ├── landscape
    ├── logo
    ├── people
    ├── plants
    └── vehicles

All MJHQ-30K prompts are in cache/MJHQ-30K/tos_dataset.jsonl. Run:

bash scripts/eval_mjhq30k.sh

Generated images will be saved to cache/inference/DIM-4.6B-T2I/MJHQ-30K. We use pytorch-fid to compute FID.

</details> <details> <summary><b>✏️ ImgEdit</b></summary>

Download ImgEdit and organize under cache/:

cache
└── ImgEdit
    └── Benchmark
        ├── hard
        ├── multiturn
        └── singleturn
            ├── animal
            │   ├── {id}.jpg
            │   └── ...
            ├── architecture
            ├── clothes
            ├── compose
            ├── daily object
            ├── for_add
            ├── human
            ├── style
            ├── transport
            ├── judge_prompt.json
            └── singleturn.json

Four evaluation jsonl files are provided in cache/ImgEdit:

1. tos_dataset_edit.jsonl — Original prompts.
2. tos_dataset_edit_cot.jsonl — CoT-style prompts from GPT-4o.
3. tos_dataset_edit_cot_Qwen2.5-VL-3B-Instruct.jsonl — CoT-style prompts from Qwen2.5-VL-3B.
4. tos_dataset_edit_cot_Qwen2.5-VL-7B-Instruct.jsonl — CoT-style prompts from Qwen2.5-VL-7B.

Run:

bash scripts/eval_imgedit.sh

Generated images will be saved to cache/inference/DIM-4.6B-Edit/ImgEdit. Follow the ImgEdit official repo for metric calculation.

</details> <details> <summary><b>📝 GEdit-Bench-EN</b></summary>

Download GEdit-Bench, extract raw images, and organize under cache/:

cache
└── GEdit-Bench
    └── input_image_raw
        ├── {id}.png
        ├── {id}.png
        ├── {id}.png
        └── ...

Four evaluation jsonl files are provided in cache/GEdit-Bench:

1. tos_dataset_edit_en.jsonl — Original prompts.
2. tos_dataset_edit_en_cot.jsonl — CoT-style prompts from GPT-4o.
3. tos_dataset_edit_en_cot_Qwen2.5-VL-3B-Instruct.jsonl — CoT-style prompts from Qwen2.5-VL-3B.
4. tos_dataset_edit_en_cot_Qwen2.5-VL-7B-Instruct.jsonl — CoT-style prompts from Qwen2.5-VL-7B.

Run:

bash scripts/eval_gedit_bench.sh

Generated images will be saved to cache/inference/DIM-4.6B-Edit/GEdit-Bench. Follow the GEdit-Bench official repo for metric calculation.

</details>

📖 Citation

If you find DIM useful for your research, please consider citing our paper:

@misc{zeng2025draw,
  title         = {Draw-In-Mind: Rebalancing Designer-Painter Roles in Unified Multimodal Models Benefits Image Editing},
  author        = {Zeng, Ziyun and Zhang, David Junhao and Li, Wei and Shou, Mike Zheng},
  year          = {2025},
  eprint        = {2509.01986},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV},
  url           = {https://arxiv.org/abs/2509.01986}
}