ComfyUI-DiffAid-Patches

ComfyUI-DiffAid-Patches

ComfyUI custom nodes that apply Diff-Aid-inspired inference-time text-conditioning patches to supported diffusion models.

This repository is a practical patch pack for ComfyUI inference, not a paper-exact reproduction of the original Diff-Aid training method. It currently provides four nodes:

• Flux-family Diff-Aid Sparse Patch — for Flux-family MMDiT models exposed through ComfyUI as double_blocks / single_blocks
• WAN Diff-Aid Sparse Patch — for native ComfyUI WAN-family video models exposed as a single blocks list, including WAN 2.1 and WAN 2.2 layouts
• MiniMax H3 Diff-Aid Sparse Patch — for native ComfyUI MiniMax H3 models with packed text, visual, audio, and video rows
• SDXL Diff-Aid Cross-Attention Patch — for SDXL-style cross-attention U-Nets, as an architectural adaptation of the same high-level idea

---

Paper credit / reference

This repository is based on ideas from:

Binglei Li, Mengping Yang, Zhiyu Tan, Junping Zhang, Hao Li Diff-Aid: Inference-time Adaptive Interaction Denoising for Rectified Text-to-Image Generation arXiv:2602.13585, 2026
https://arxiv.org/abs/2602.13585

The paper presents a lightweight Aid module for rectified text-to-image diffusion transformers that adaptively modulates textual features per token, per block, and per denoising timestep. It is trained with the backbone model frozen, and the learned modulation coefficients are used at inference time as a plug-in enhancement.

Please cite and credit the paper if you use this repository as part of experiments, implementations, ports, or derivative work.

Suggested citation

@article{li2026diffaid,
  title={Diff-Aid: Inference-time Adaptive Interaction Denoising for Rectified Text-to-Image Generation},
  author={Li, Binglei and Yang, Mengping and Tan, Zhiyu and Zhang, Junping and Li, Hao},
  journal={arXiv preprint arXiv:2602.13585},
  year={2026}
}

---

What the paper does

The paper’s central claim is that prompt adherence and image quality can be improved by strengthening text-image interaction selectively, rather than with a single global guidance scale.

Its Aid module learns a coefficient tensor α that depends on:

• the current transformer block
• the current denoising timestep
• the current text token features

and uses that to modulate the text features before they participate in attention:

c̃ = c + c ⊙ α

The paper also adds:

• bounded modulation (tanh)
• a gating path to encourage sparsity
• regularization on α
• optional DPO / reward-based optimization during training

The result is a learned plug-in that is lightweight in parameter count but still adaptive and interpretable.

---

What the paper reports

The paper evaluates Diff-Aid on FLUX and SD 3.5 and reports consistent gains in prompt-following and image-quality metrics.

Main reported results from the paper

On the paper’s reported benchmarks:

• SD 3.5 improves from 9.31 → 9.48 on HPSv3 and from 0.72 → 0.77 on GenEval.
• FLUX improves from 10.42 → 10.71 on HPSv3 and from 0.68 → 0.70 on GenEval.
• The paper also reports gains on HPSv2, ImageReward, and Aesthetic Score for both baselines.

Sparse-enhancement finding relevant to this repo

Besides the full learned method, the appendix reports a simpler FLUX sparse enhancement result by only boosting a small subset of blocks:

• selected FLUX blocks: {1, 15, 36, 41, 48}
• with α = 0.5 on those blocks and 0.0 elsewhere

For that sparse variant, the paper reports the following FLUX results:

• baseline FLUX: HPSv2 28.53 / HPSv3 10.42 / ImgRwd 0.89 / Aes 6.66
• sparse enhancement: 28.61 / 10.57 / 0.98 / 6.71
• full method: 28.80 / 10.84 / 0.95 / 6.76

That appendix result is the main reason this repository starts with a Flux sparse patch instead of trying to fake the entire trained Aid pipeline.

---

What