README.md

<h1 align='center'>[ICML 2026] Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers</h1>

<div align='center'>
    <a href='https://yuxuan0919.github.io/' target='_blank'>Yuxuan Yao</a><sup>1,2,*</sup> &emsp;
    <a href='https://scholar.google.com/scholar?q=Yuxuan+Chen' target='_blank'>Yuxuan Chen</a><sup>1,*</sup> &emsp;
    <a href='https://openreview.net/profile?id=%7EHui_Li43' target='_blank'>Hui Li</a><sup>1</sup> &emsp;
    <a href='https://scholar.google.com/scholar?q=Kaihui+Cheng' target='_blank'>Kaihui Cheng</a><sup>1</sup> &emsp;
    <a href='https://scholar.google.com/scholar?q=Qipeng+Guo' target='_blank'>Qipeng Guo</a><sup>3</sup> &emsp;
    <a href='https://scholar.google.com/scholar?q=Yuwei+Sun' target='_blank'>Yuwei Sun</a><sup>4</sup> &emsp;
</div>
<div align='center'>
    <a href='https://scholar.google.com/scholar?q=Zilong+Dong' target='_blank'>Zilong Dong</a><sup>5</sup> &emsp;
    <a href='https://scholar.google.cz/citations?user=z5SPCmgAAAAJ&hl=zh-CN' target='_blank'>Jingdong Wang</a><sup>6</sup> &emsp;
    <a href='https://sites.google.com/site/zhusiyucs/home' target='_blank'>Siyu Zhu</a><sup>1,2</sup> &emsp;
</div>



<br>

<div align='center'>
<sup>1</sup>Fudan University &emsp;
<sup>2</sup>Shanghai Innovation Institute
</div>
<div align='center'>
<sup>3</sup>Shanghai AI Laboratory &emsp;
<sup>4</sup>Shanghai Academy of AI for Science &emsp;
<sup>5</sup>Alibaba Group &emsp;
<sup>6</sup>Baidu
</div>

<br>

<div align="center">

[![Paper](https://img.shields.io/badge/arXiv-2604.23632-b31b1b.svg)](https://arxiv.org/abs/2602.06886)
[![GitHub](https://img.shields.io/badge/GitHub-Repo-181717.svg?logo=github)](https://github.com/fudan-generative-vision/PromptReinjection)
[![License](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)

</div>


## 📖 Introduction

Prompt Reinjection is a training-free inference method for multimodal diffusion transformers that mitigates prompt forgetting by reinjecting early-layer prompt features into deeper text layers, improving GenEval overall scores by `6.48%` on SD3.5-large and `7.75%` on HunyuanImage-2.1, while adding only about `0.00002x` block-level FLOPs for base reinjection and `0.088x` for the full aligned variant.

## 📘 Overview

Prompt Reinjection starts from a simple observation: in multimodal diffusion transformers (MMDiTs) such as SD3-medium, SD3.5-large, FLUX.1, HunyuanImage-2.1, and Qwen-Image, prompt information fades as depth increases. That makes instruction following weaker, especially on position, attributes, counting, and long prompts.

### Prompt Forgetting

  Unlike traditional DiTs, where text serves as a relatively stable conditioning signal, MMDiTs jointly update text and image tokens throughout denoising, even though the text tokens receive no direct supervision. The paper shows that deeper text features gradually lose fine-grained prompt semantics, a phenomenon we call `prompt forgetting`.

<p align="center">
  <img src="assets/prompt-forgetting.png" alt="Prompt Forgetting" width="320">
</p>

The figure above captures the core trend: for SD3, SD3.5, and FLUX, prompt information becomes less recoverable in deeper layers. This helps explain why base MMDiT models often miss spatial relations, attributes, and numeracy constraints in generation.

### Prompt Reinjection

Prompt Reinjection fixes this at inference time. It takes semantically stronger text features from an early layer, aligns them to the deeper feature space, and reinjects them into later blocks so prompt constraints stay active through the full denoising stack.

<p align="center">
  <img src="assets/prompt-reinjection.png" alt="Prompt Reinjection" width="320">
</p>

The method is training-free, lightweight, and easy to plug into the original MMDiT forward process. No retraining is required.

<p align="center">
  <img src="assets/qualitative-compare.png" alt="Qualitative Comparison" width="640">
</p>

As shown above, Prompt Reinjection makes SD3.5, FLUX, and Qwen-Image follow prompt constraints more consistently across position, attribute, counting, and complex prompts. The paper reports consistent gains on GenEval, DPG-Bench, and T2I-CompBench++ while preserving overall generation quality.


## 🚀 Quick Start

### 1. Create an environment

We recommend Python `3.10+` and one environment per model family.

```bash
python3.10 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
```

Install dependencies with either an editable package install or the pinned root requirements:

```bash
pip install -e .
```

```bash
pip install -r requirements.txt
```

If you want a narrower per-model environment, install one model-specific requirement file instead:

```bash
pip install -r requirements/sd3.txt
pip install -r requirements/sd3.5.txt
pip install -r requirements/flux.txt
pip install -r requirements/qwen.txt
pip install -r requirements/hunyuanimage.txt
```

If you need a specific CUDA build, install `torch` and `torchvision` first from the official PyTorch channel, then rerun one of the commands above.

For HunyuanImage, install Tencent's official runtime before `requirements/hunyuanimage.txt`:

```bash
git clone https://github.com/Tencent-Hunyuan/HunyuanImage-2.1.git
pip install -r HunyuanImage-2.1/requirements.txt
pip install flash-attn==2.7.3 --no-build-isolation
pip install -r requirements/hunyuanimage.txt
```

### 2. Prepare checkpoints

Pass model paths explicitly with `--model-path` in open-source usage.

- `sd3`, `sd3.5`, `flux`, `qwen`: `--model-path /path/to/model`
- `hunyuanimage`: `--model-path /path/to/HunyuanImage-2.1` and optional `--model-name`

For `hunyuanimage`, `--model-path` can point either to the HunyuanImage runtime root or to its `ckpts` directory.

### 3. Run one prompt

The default helper script reads the released per-model inference and Prompt Reinjection settings from [prompt_reinjection/reinjection_configs.json](prompt_reinjection/reinjection_configs.json), so standard inference does not need manual residual arguments.

By default, all models now run without memory-saving inference shortcuts such as CPU offload, Hunyuan runtime offload, VAE slicing, VAE tiling, or attention slicing. This keeps the default path as close as possible to the original plain inference flow. These options are enabled only when you pass them explicitly.

```bash
bash prompt_reinjection/test_reinjection.sh \
  --model sd3 \
  --model-path /path/to/SD3 \
  --prompt "A photo of a couch below a potted plant."
```

Supported `--model` values: `sd3`, `sd3.5`, `flux`, `qwen`, `hunyuanimage`.

To change the default released settings for a model, edit [prompt_reinjection/reinjection_configs.json](prompt_reinjection/reinjection_configs.json). The helper script and the benchmark entrypoints all read from the same file.

To run the plain base model instead of Prompt Reinjection:

```bash
bash prompt_reinjection/test_reinjection.sh \
  --model sd3 \
  --model-path /path/to/SD3 \
  --prompt "A photo of a couch below a potted plant." \
  --reinjection off
```

To enable memory-saving options explicitly:

- `flux` and `qwen`: pass `--cpu-offload model` or `--cpu-offload sequential`
- `flux` and `qwen`: optionally add `--vae-slicing`, `--vae-tiling`, or `--attention-slicing auto`
- `hunyuanimage`: pass `--enable-offload`

Example:

```bash
bash prompt_reinjection/test_reinjection.sh \
  --model flux \
  --model-path /path/to/FLUX.1-dev \
  --prompt "A photo of a couch below a potted plant." \
  --cpu-offload model \
  --vae-slicing
```

## 🧭 Manual SD3 Example

If you want to bypass the helper script and set the reinjection parameters manually, an SD3 example is:

```bash
python -m prompt_reinjection.run_sample \
  --model sd3 \
  --model-path /path/to/SD3 \
  --prompt "A photo of a couch below a potted plant." \
  --output outputs/manual_sd3.png \
  --steps 28 \
  --cfg 7.0 \
  --residual_origin_layer 1 \
  --residual_target_layers $(seq 2 23) \
  --residual_weights 0.025 \
  --residual_use_anchoring 1 \
  --residual_procrustes_path prompt_reinjection/rotations/sd3_coco5k_o1.pt
```

The same default memory policy also applies to the Python entrypoints. If you do not pass a memory-saving flag explicitly, the run stays on the plain inference path. For example, FLUX and Qwen only enable Diffusers offload when you pass `--cpu-offload model` or `--cpu-offload sequential`, and HunyuanImage only enables its runtime offload when you pass `--enable-offload`.

## 🧮 Procrustes Precomputation

We use COCO 5k for the released Procrustes statistics.

For open-source usage, we recommend:

- `sd3` and `flux`: use the released Procrustes-aligned Prompt Reinjection settings.
- `sd3.5`, `qwen`, and `hunyuanimage`: use the most basic Prompt Reinjection variant without anchoring and without rotation. It already works well and adds almost zero inference cost.

### SD3

```bash
python SD3/compute.py \
  --model /path/to/SD3 \
  --dataset coco \
  --datadir data \
  --num-samples 5000 \
  --origin-layer 1 \
  --target-layer-start 2 \
  --col-center \
  --output outputs/procrustes_rotations/sd3_coco5k_o1.pt
```

### FLUX

```bash
python FLUX/compute.py \
  --model /path/to/FLUX.1-dev \
  --dataset coco \
  --datadir data \
  --num-samples 5000 \
  --origin-layer 2 \
  --target-layer-start 3 \
  --col-center \
  --output outputs/procrustes_rotations/flux_coco5k_o2.pt
```

## 🤗 Released Rotations

The released Procrustes rotations are hosted at [LewisYao/PromptReinjection](https://huggingface.co/fudan-generative-ai/PromptReinjection):

- [`sd3_coco5k_o1.pt`](https://huggingface.co/fudan-generative-ai/PromptReinjection/blob/main/sd3_coco5k_o1.pt)
- [`flux_coco5k_o2.pt`](https://huggingface.co/fudan-generative-ai/PromptReinjection/blob/main/flux_coco5k_o2.pt)

Download them to `prompt_reinjection/rotations/` with:

```bash
hf download LewisYao/PromptReinjection \
  sd3_coco5k_o1.pt \
  flux_coco5k_o2.pt \
  --local-dir prompt_reinjection/rotations
```

The default helper script will pick them up automatically once they are placed under `prompt_reinjection/rotations/`.

## 📊 Evaluation

The benchmark scripts below also read the released per-model defaults from [prompt_reinjection/reinjection_configs.json](prompt_reinjection/reinjection_configs.json). By default, they run with `--reinjection on`. Use `--reinjection off` for the plain base model, or edit the config file if you want to change the released settings globally.

Like the helper script, these benchmark entrypoints do not enable CPU offload, Hunyuan runtime offload, VAE slicing, VAE tiling, or attention slicing unless you pass those flags explicitly.

### GenEval

```bash
python -m prompt_reinjection.run_geneval \
  --model sd3 \
  --model-path /path/to/SD3 \
  --metadata_file /path/to/geneval/metadata.jsonl \
  --outdir outputs/geneval_sd3
```

Base-model variant:

```bash
python -m prompt_reinjection.run_geneval \
  --model sd3 \
  --model-path /path/to/SD3 \
  --metadata_file /path/to/geneval/metadata.jsonl \
  --outdir outputs/geneval_sd3_base \
  --reinjection off
```

### DPG-Bench

```bash
python -m prompt_reinjection.run_dpg \
  --model sd3.5 \
  --model-path /path/to/SD3.5-large \
  --prompt_dir /path/to/dpg/prompts \
  --save_dir outputs/dpg_sd35
```

### T2I-CompBench++

```bash
python -m prompt_reinjection.run_t2i \
  --model qwen \
  --model-path /path/to/Qwen-Image \
  --dataset_dir /path/to/t2i-compbench/prompts \
  --outdir_base outputs/t2i_qwen
```

These scripts generate benchmark-format images. Final scoring should still be done with the official benchmark evaluators.

## 🧩 Apply to New Models

Prompt Reinjection is designed for MMDiT-style open-source models where text features evolve inside the denoising transformer together with visual features. If your new model follows this pattern, you can usually add a basic reinjection version with only a small amount of integration work.

To plug a new model into this framework, the minimum steps are:

- Add a new model folder with an `adapter.py` that implements the adapter interface used in [prompt_reinjection/adapter_api.py](prompt_reinjection/adapter_api.py) and follows the existing examples in [SD3/adapter.py](SD3/adapter.py), [FLUX/adapter.py](FLUX/adapter.py), and [Qwen/adapter.py](Qwen/adapter.py).
- Make the model pipeline or transformer expose `set_residual_config(...)` so it can receive `residual_origin_layer`, `residual_target_layers`, `residual_weights`, `residual_use_anchoring`, and `residual_rotation_matrices`, as shown in [SD3/pipeline.py](SD3/pipeline.py), [FLUX/pipeline.py](FLUX/pipeline.py), and [Qwen/pipeline.py](Qwen/pipeline.py).
- Register the new adapter in [prompt_reinjection/registry.py](prompt_reinjection/registry.py) so it becomes available to `run_sample`, `run_geneval`, `run_dpg`, `run_t2i`, and `compute_procrustes`.
- If you want rotation-based alignment later, also provide a model-specific `compute.py` and expose it as the adapter’s `compute_script`, so `python -m prompt_reinjection.compute_procrustes --model YOUR_MODEL ...` can dispatch correctly.

For a new model, we recommend starting from the most basic reinjection setting first:

- `origin = 1`
- `target = 2-last`
- `weight = 0.025`
- `no anchoring`
- `no rotation`

In practice, that means using the shallowest stable MMDiT block as the source, reinjecting into all later blocks, and keeping the setup as lightweight as possible. If your model has `L` blocks indexed from `0` to `L-1`, the default starting rule is:

```text
residual_origin_layer = 1
residual_target_layers = [2, 3, ..., L-1]
residual_weights = 0.025
residual_use_anchoring = 0
residual_procrustes_path = ""
```

A typical manual run looks like this after the model has been integrated into the registry:

```bash
python -m prompt_reinjection.run_sample \
  --model your_model \
  --model-path /path/to/your/model \
  --prompt "A photo of a couch below a potted plant." \
  --output outputs/your_model_base_reinjection.png \
  --residual_origin_layer 1 \
  --residual_target_layers $(seq 2 LAST_LAYER) \
  --residual_weights 0.025 \
  --residual_use_anchoring 0 \
  --residual_procrustes_path ""
```

Replace `LAST_LAYER` with the last text-processing block index of your model. For example, if the model has 24 blocks indexed from `0` to `23`, use `$(seq 2 23)`.

Once this basic version runs and already improves instruction following, the next recommended upgrades are:

- Turn on anchoring first by setting `--residual_use_anchoring 1`. This is usually the safest first upgrade when the model shows cross-layer scale or shift mismatch.
- If you want further gains, add rotation-based alignment by computing a Procrustes file on a prompt set such as COCO-5K and passing it through `--residual_procrustes_path`. This is useful when shallow and deep text features differ not only in scale, but also in feature geometry.

In short, the recommended order for a new model is: first make basic reinjection work, then try anchoring, and only then add rotation if you want the strongest alignment.

## 📝 Citation

If you find this project useful, please cite the ICML 2026 paper:

```bibtex
@inproceedings{yao2026prompt,
  title={Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers},
  author={Yao, Yuxuan and Chen, Yuxuan and Li, Hui and Cheng, Kaihui and Guo, Qipeng and Sun, Yuwei and Dong, Zilong and Wang, Jingdong and Zhu, Siyu},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2026}
}
```

## 🤗 Acknowledgements

We thank the open-source communities behind SD3, SD3.5, FLUX, Qwen-Image, and HunyuanImage. This release builds on their public model and runtime ecosystems to make Prompt Reinjection reproducible in open source.