Sync README from GitHub repo

README.md

@@ -1,19 +1,406 @@
-# Prompt Reinjection Rotations
+<h1 align='center'>[ICML 2026] Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers</h1>
 
-Released Procrustes rotation files for the ICML 2026 paper
-`Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers`.
+<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>
 
-## Files
 
-- `sd3_coco5k_o1.pt`: SD3 rotation computed on COCO 5k, with origin layer 1 and target layers 2-23.
-- `flux_coco5k_o2.pt`: FLUX rotation computed on COCO 5k, with origin layer 2 and target layers 3-56.
 
-## Usage
+<br>
 
-Download the files and place them under `prompt_reinjection/rotations/` in the main repository:
+<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)
+[![arXiv](https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow)](https://huggingface.co/fudan-generative-ai/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
-hf download LewisYao/PromptReinjection sd3_coco5k_o1.pt flux_coco5k_o2.pt --local-dir prompt_reinjection/rotations
+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
 ```
 
-The default helper script in the main repository will pick them up automatically.
+### 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.
+
+## 🔐 Security
+
+Run a secret scan before any public push. `gitleaks` does not ship on PyPI, so install it from the official release or your system package manager instead of `pip install gitleaks`.
+
+For a full Git history scan:
+
+```bash
+gitleaks detect --source . -v
+```
+
+For a working-tree-only scan:
+
+```bash
+gitleaks detect --source . --no-git -v
+```
+
+This repository also includes a `.pre-commit-config.yaml` entry for `gitleaks`:
+
+```bash
+pip install pre-commit
+pre-commit install
+pre-commit run --all-files
+```
+
+In addition to automated secret scanning, avoid committing local absolute paths, internal hostnames, or private checkpoint locations.
+
+## ⚖️ License
+
+The code in this repository is released under the [MIT License](LICENSE).
+
+Prompt Reinjection is an inference-time code intervention that plugs into external model pipelines. This repository does not include, redistribute, or relicense any third-party model weights. Upstream weights and runtimes remain subject to their original licenses, including:
+
+- SD3 and SD3.5: Stability AI Community License
+- FLUX.1-dev: FLUX.1 `[dev]` Non-Commercial License
+- HunyuanImage-2.1: Tencent Hunyuan Community License Agreement
+- Qwen-Image: Apache License 2.0
+
+Users are responsible for ensuring that their checkpoint download, local use, fine-tuning, serving, and redistribution flows comply with the corresponding upstream terms. The MIT License in this repository applies only to this repository's code and included assets.
+
+## 📝 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.