蚂蚁 inclusionAI 开源 IDM：无需外部编码器的一步/少步文生图蒸馏方法

Intrinsic Distribution Matching

Deyuan Liu^† · Yansen Han^† · Kaiwen Zheng^† · Jun Xie · Jiacheng Liu
Zhenglin Cheng · Peng Sun · Jintao Zhang^* · Jianguo Li^* · Tao Lin^*

Inclusion AI · Westlake University · Tsinghua University
Zhejiang University · Shanghai Innovation Institute · MBZUAI

^† Equal contribution   ^* Corresponding authors

Results · Method · Inference · Checkpoints

IDM improves one-step and few-step text-to-image generation by combining internal distribution matching distillation (DMD) with moment matching in the generator's own EMA feature space. No external feature encoder is required. The same trained checkpoint supports both 1-NFE and 2-NFE sampling.

IDM (Z-Image), 2 NFE. Samples from the paper.

Core results

• Long-text rendering: 94.25% on LongText-CN and 90.20% on LongText-EN at 2 NFE, the highest scores among the methods compared at this budget.
• Fast generation: full-parameter IDM (Z-Image) takes 0.25 s at 1 NFE and 0.30 s at 2 NFE; IDM (Qwen-Image) takes 0.35 s and 0.44 s, respectively, on a single A100 in BF16.
• Useful internal features: adding EMA feature-moment matching improves all five reported aggregate metrics over the shared internal DMD control. Its results are close to those obtained with DINOv2 features.

All values below are reported in the paper. They are not measurements from running this repository.

Long-text rendering and general quality

Z-Image comparisons at 2 NFE. All scores are percentages; higher is better.

At 2 NFE, IDM improves LongText-CN/EN over Z-Image-Turbo++ by 2.63 / 0.32 percentage points, and over APEX by 3.87 / 1.44 points. The 8-NFE reference retains a higher LongText-EN score. Published reference models have different training settings and data; APEX also scores higher on DPG-Bench. The controlled ablations below measure the contribution of IDM's individual components within a shared setting.

Quality and efficiency across both backbones

Selected results from the paper's system-level comparison. Throughput uses batch size 10; latency uses batch size 1, both on a single A100 in BF16. LoRA weights are merged before timing. FID is lower-is-better; CLIP and GenEval are higher-is-better. GenEval here uses the paper's 0–1 scale.

The full Qwen-Image model also reaches 87.16% DPG-Bench and 0.56 WISE at 2 NFE. Rank-64 LoRA updates 0.4B parameters and retains most of the full model's GenEval performance.

What does the added moment loss contribute?

Controlled ablations use full-parameter IDM (Z-Image) at 2 NFE, with the same data and training configuration. All values are percentages.

Adding the EMA feature-moment loss to the internal DMD control gives +0.19 OneIG, +0.35 GenEval, +0.28 DPG-Bench, +0.38 LongText-CN, and +0.41 LongText-EN. DINOv2 features also improve the control; the small differences support the use of the generator's own representation without requiring an external encoder.

Why combine velocity and endpoint regression?

The choice is empirical. With feature-moment matching retained in every row, equal-weight velocity + endpoint regression leads on all five aggregate metrics.

These regressions apply different time weightings to the same detached DMD correction. They do not provide two independent sources of distributional information.

Method

DMD uses a score difference between generated and reference distributions to supervise a fast generator. Following APEX, IDM obtains this difference from internal velocity branches, avoiding a separate fake-score network.

A Gaussian-plus-residual decomposition of the reverse-KL field motivates an explicit constraint on moments. IDM keeps the full DMD correction and adds a loss matching the means and per-coordinate standard deviations of features from the generator's EMA copy. Moment matching complements the score correction; it does not replace the residual part of the distributional signal.

At inference, only the learned conditional generator, text encoder, and VAE are needed. No moment computation, feature queue, external encoder, or fake-score evaluation is performed.

Visual comparisons

In these examples, full IDM preserves dense text more consistently across 1 and 2 NFE. The velocity + endpoint control without the moment loss produces more saturated colors in parts of the sushi poster. APEX is an external reference trained on a smaller collection; the IDM controls share one training setting.

Comparison across sampling budgets

Each column shows a final output at the indicated NFE budget, not an intermediate state of a single sampling trajectory.

Checkpoints

Coming soon. Download links and checkpoint-specific sampling settings will be added with the weight release.

The loader accepts a local Diffusers pipeline or a Hugging Face repository, an optional transformer-only checkpoint, and Diffusers-format LoRA weights. Base-model or APEX weights are not IDM checkpoints; they do not reproduce the IDM results above.

Inference

Install

Use Python 3.10 or newer and a CUDA machine with enough memory for the selected model. The default precision is BF16. The 20B backbone requires substantially more memory than the 6B model. This entry point loads the text encoder, transformer, and VAE on the selected device.

python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt

Install a PyTorch build appropriate for your CUDA environment if needed. FlashAttention, FSDP, and training services are not required. Diffusers is pinned because the auxiliary-time wrappers use its model interfaces.

Generate an image

Once an IDM checkpoint is available, set its path through --model:

python inference.py \
  --backbone z-image \
  --model /path/to/IDM-Z-Image \
  --prompt 'A cafe menu with the heading "MORNING COFFEE" and three lines: "Americano $4", "Latte $5", "Mocha $6".' \
  --steps 2 --height 1024 --width 1024 --seed 42 \
  --output-dir outputs/z-image

python inference.py \
  --backbone qwen-image \
  --model /path/to/IDM-Qwen-Image \
  --prompt '一张咖啡馆海报，标题“今日咖啡”，三行菜单“美式咖啡 18元”“拿铁咖啡 22元”“摩卡咖啡 25元”。' \
  --steps 1 --height 1024 --width 1024 --seed 42 \
  --output-dir outputs/qwen-image

Each output includes a PNG and a JSON file containing the prompt, seed, weight paths, and sampling configuration. Existing outputs are not overwritten.

• --steps 1 or --steps 2 performs exactly one or two conditional transformer evaluations, without classifier-free guidance.
• Height and width must be positive multiples of 16.
• --device cuda:1 selects another GPU. CPU is available for small compatibility checks.
• --prompt-file examples/prompts.txt --num-images 2 runs each nonempty prompt twice, incrementing the seed for every image.
• --local-files-only disables model downloads and uses local files or an existing Hugging Face cache.
• Run python inference.py --help for the complete argument list.

The default sampler is deterministic first-order Euler: noise levels [1] or [1, 0.5], followed by the clean endpoint at 0. Custom schedules can be supplied with --timesteps, and optional noise refresh with --stochastic-ratio. These are configurable inference defaults; checkpoint-specific settings and end-to-end quality validation are pending the weight release.

Load a transformer-only checkpoint

Use the corresponding base pipeline for its tokenizer, text encoder, and VAE, and point --transformer at the exported IDM transformer directory. That directory must contain config.json and safetensors weights.

python inference.py \
  --backbone z-image \
  --model /path/to/Z-Image-Turbo \
  --transformer /path/to/IDM-transformer \
  --prompt 'A book cover titled "The Art of Calm".' \
  --steps 2

--aux-time-embed auto selects the auxiliary-time wrapper when the checkpoint's _class_name identifies it. For older exports, use --aux-time-embed on or off to match the checkpoint architecture. Missing or unexpected transformer parameters cause an error instead of silently leaving random weights in the model.

Load LoRA weights

python inference.py \
  --backbone qwen-image \
  --model /path/to/Qwen-Image \
  --lora /path/to/IDM-LoRA \
  --lora-weight-name pytorch_lora_weights.safetensors \
  --prompt 'A poster with the title "Explore Japan".' \
  --steps 2 --output-dir outputs/lora

The adapter must use the Diffusers pipeline LoRA format and match its base model. The default filename is pytorch_lora_weights.safetensors; use --lora-weight-name for a different filename. It is merged at load time; --lora-scale defaults to 1.0. Raw optimizer states, FSDP shards, and unconverted PEFT training checkpoints are not inference inputs.

Repository layout

inference.py              Command-line image generation
unified_sampler.py        One- and two-evaluation first-order sampling
diffusers_patch/          Inference adapters and auxiliary-time model wrappers
examples/prompts.txt      English and Chinese prompt examples
assets/                   Figures from the IDM paper
tests/                    CPU checks with small models; no model downloads

For development checks:

python -m pip install -r requirements-dev.txt
python -m pytest -q
ruff check .

Citation

If you use this code, please cite IDM and APEX:

@misc{liu2026intrinsicdistributionmatching,
  title={Intrinsic Distribution Matching},
  author={Deyuan Liu and Yansen Han and Kaiwen Zheng and Jun Xie and Jiacheng Liu and Zhenglin Cheng and Peng Sun and Jintao Zhang and Jianguo Li and Tao Lin},
  year={2026},
  eprint={2610.39633},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2610.39633},
}

@article{liu2026self,
  title={Self-Adversarial One Step Generation via Condition Shifting},
  author={Liu, Deyuan and Sun, Peng and Han, Yansen and Cheng, Zhenglin and Chen, Chuyan and Lin, Tao},
  journal={arXiv preprint arXiv:2604.12322},
  year={2026}
}

Acknowledgments and license

The inference implementation is adapted from APEX, which builds on TwinFlow, RCGM, and UCGM. We use model components from Hugging Face Diffusers.

Code is distributed under Apache 2.0. Base-model and future IDM checkpoint licenses apply separately to the weights. The project page uses KaTeX under the MIT license.