EndoCoT: Scaling Endogenous Chain-of-Thought Reasoning in Diffusion Models

📝TODO

• [x] Open source the training code
• [x] Open source the training data
• [x] Open source the main task ckpt
• [ ] Open source the edit model ckpt
• [ ] Refactor the codebase for better usability and maintainability

📰News

• 🎉 [2026/6/26] Our paper has been accepted to ECCV 2026!
• 🚀 [2026/3/18] We have released the EndoCoT-Data.
• 🚀 [2026/3/12] We have released the EndoCoT repository and ckpts.

🌟Highlight

• EndoCoT is a reasoning paradigm for diffusion models that enables step-by-step inference. It outperforms conventional training methods on Qwen-Image-Edit-2511.

• And provide transparent, intermediate reasoning trajectories.

⚡Quick Start

Setup environment

git clone https://github.com/InternLM/EndoCoT
cd EndoCoT
conda create -n EndoCoT python=3.10
conda activate EndoCot
Please install the version of torch compatible with your machine.
pip install -r requirements.txt
Please install the version of vLLM compatible with your machine.

Inference

• Download the ckpt:

> Following the configuration of Diffthinker, we provide a customized checkpoint for Qwen-Image-Edit. This checkpoint has been merged from the original safetensors to ensure compatibility withDiffsynth-Studio training. Please use the checkpoint provided in this repository instead of the official version for correct loading and inference.

• Test Single Case

cd test
   python test.py \
       --task Maze \
       --modelroot /path/to/mergedckpts \
       --lorapath /path/to/yourlora_weight.safetensors \
       --inputimage ./data/sudokusample.png \
       --outputdir ./outputs/sudokuresults

• Eval Our Ckpt

cd Maze
   bash eval/genandparse.sh
   bash eval/eval_path.sh

Training

• Download the datasets & metadata.csv

> Since the metadata uses relative paths, please ensure the dataset files are placed in the same directory as metadata.csv

• Train your model

cd DiffSynth-Studio
   bash add/Maze/stage1.sh
   python changeckptprefix.py --src /path/to/the/Maze/save/dir/Maze_stage1	
   bash add/Maze/stage2.sh
   python changeckptprefix.py --src /path/to/the/Maze/save/dir/Maze_stage2

How to change the latent reasoning steps?

Note on Customization: Since the current implementation is straightforward, you can only manually adjust the latent reasoning steps in DiffSynth-Studio/diffsynth/pipelines/qwen_image.py:

- Line 442: Modify infer_steps.

- Line 471: Modify training_steps.

##### We plan to optimize this in future releases.

def encodepromptedit(self, pipe: QwenImagePipeline, prompt, editimage, isfinal, gt_prompt=None, idx=None):
drop_idx = 64         if type(prompt[0])==str:             template =  "<|imstart|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|imend|>\n<|imstart|>user\n<|visionstart|><|imagepad|><|visionend|>{}<|imend|>\n<|imstart|>assistant\n"             txt = template.format(prompt[0])             modelinputs = pipe.processor(text=txt, images=editimage, padding=True, return_tensors="pt").to(pipe.device)             embeddinglayers = pipe.textencoder.model.languagemodel.getinput_embeddings()             with torch.no_grad():                 inputsembeds = embeddinglayers(modelinputs.inputids)             self.attentionmask = modelinputs.attention_mask             self.pixelvalues = modelinputs.pixel_values             self.imagegridthw = modelinputs.imagegrid_thw         else:             inputs_embeds= prompt[0]                      # dxl: test use         if is_final==None or idx!=None:             print("现在在inference。或者stage2训练")             if idx!=None:                 iter_times = idx-2             else:                 # infer step                 iter_times = 50                              with torch.no_grad():                 inputsembeds = self.manualgenerate_eval(                     pipe,                      inputsembeds=inputsembeds,                     maxnewtokens=iter_times,                 ).detach()                          # dxl: only update the last 2 tokens             if idx!=None:                 inputsembeds = self.manualgenerate_eval(                     pipe,                     inputsembeds=inputsembeds,                     maxnewtokens=2,                 )
generatedembeds = inputsembeds
... ...                   # dxl：training         if is_final!=None and idx==None:             try:                 generatedembeds,  = self.manual_generate(                     pipe,                     inputsembeds=inputsembeds,                     isfinal=isfinal,                     # training steps                     maxnewtokens=2,                 )             except Exception as e:                 print(f"Error!: {type(e).name} - {e}")                 print(inputs_embeds.shape)                 assert False
try:              return splithiddenstates, generatedembeds, eosloss         except:             print(f"[WARNING] Prompt was not updated correctly for inference.")             return splithiddenstates

📖 Citation

@article{dai2026endocot,
  title={EndoCoT: Scaling Endogenous Chain-of-Thought Reasoning in Diffusion Models},
  author={Dai, Xuanlang and Zhou, Yujie and Xing, Long and Bu, Jiazi and Wei, Xilin and Liu, Yuhong and Zhang, Beichen and Chen, Kai and Zang, Yuhang},
  journal={arXiv preprint arXiv:2603.12252},
  year={2026}
}

⚖️ License