Nodes/comfyui_jankhidiffusion/ApplyMSWMSAAttention
ComfyUI Node

ApplyMSWMSAAttention

SD1.5 at 2K without melting your GPU — the windowed-attention trick

By blepping·Created 2 years ago·Updated about a year ago· 125
ApplyMSWMSAAttention
  • model
  • MODEL
input_blocks1,2
middle_blocks
output_blocks9,10,11
time_modepercent
start_time0.00
end_time1.00
yaml_parameters

If you've ever cranked an SD1.5 checkpoint up to 1024×1024 or beyond and watched the sampler grind to a crawl, this node is the answer you were looking for. ApplyMSWMSAAttention is a port of the MSW-MSA attention module from Megvii's HiDiffusion research project, and on SD1.5 it's one of the biggest free speedups you can bolt onto a high-res generation. Community testing lines up with the author's claims: people report roughly 30% faster high-res SD1.5 gens with "no quality loss at all, or quality is even better." That's not a typo - the windowed attention often reduces the smeary artifacts you get at high res, because it keeps self-attention from going cross-eyed on the whole image.

One thing to get straight up front: this node does not let you generate high-res in the first place. That's its sibling ApplyRAUNet's job. MSW-MSA just makes whatever resolution you're already running dramatically cheaper, so you pair it with RAUNet (or Deep Shrink, or just a big Empty Latent Image) rather than using it alone.

How it works

Plain self-attention cost scales quadratically with the number of tokens, and tokens are pixels - which is exactly why SD1.5 gets so slow as you push past its native 512. MSW-MSA chops the Q/K/V tensors into non-overlapping local windows, so each attention head only looks at a small neighborhood instead of the entire latent. Cost drops to roughly linear. The catch with windowed attention is grid artifacts: pixels that always sit on a window boundary never get attended to properly. HiDiffusion's fix is to shift the window grid every step, and the author made a small change here so the shifts are guaranteed to differ each step. The patch lands on the model's attn1 (self-attention) blocks via ComfyUI's model-patch API, so it works with anything that goes through the normal MODEL pipeline - samplers, schedulers, and most other patches.

The inputs that matter

The node takes a MODEL in and hands you back a patched MODEL out - wire that into your KSampler's model slot, and everything downstream just gets faster.

Three fields are the whole game for most people:

  • input_blocks - comma-separated down/input blocks to patch. Default 1,2 is for SD1.5; try 4,5 for SDXL.
  • output_blocks - comma-separated up/output blocks. Default 9,10,11 for SD1.5, 3,4,5 for SDXL.
  • start_time / end_time with time_mode set to percent - when the patch is active. Defaults 0.0→1.0 (the whole run). If you're going past ~2048px, the author suggests starting the effect a bit later, like 0.2, after any scaling effects finish.

middle_blocks exists but the tooltip basically begs you not to use it. There's also a yaml_parameters multi-line widget for the deep end: scale_mode fallback for non-multiple-of-64 resolutions, force_apply_attn2 (attn2 is guaranteed size-incompatible - don't), and the "last shift" avoidance settings. Ignore all of it until you've got a feel for the basics.

Install

ComfyUI Manager → Custom Nodes Manager → search comfyui_jankhidiffusion → Install → restart. Or the manual way:

cd ComfyUI/custom_nodes
git clone https://github.com/blepping/comfyui_jankhidiffusion
# restart ComfyUI

Nice and rare for this ecosystem: no model downloads, and no requirements.txt - it's pure torch + the ComfyUI API, so nothing to install past the pack itself.

Troubleshooting

The classic failure is a tensor size mismatch, usually because your resolution isn't a multiple of 32/64/128. Stick to multiples of 64 (or 128) and it mostly vanishes. When it can't fit, the node logs a warning about "Incompatible model patches or bad resolution" and skips the patch for that step rather than crashing.

Compatibility is the real gotcha. MSW-MSA is known to work with FreeU, SAG, PAG, ELLA, and CFG rescaling, but it fights with other attention patches that hit the same blocks - HyperTile, Deep Cache, nearsighted/slothful attention, perturbed attention. Don't stack those. And be realistic about scope: it only supports SD1.x, SD2.x and SDXL, and the SDXL speedup is marginal - this is an SD1.5 node, full stop. If all that block-tweaking makes your eyes glaze over, grab ApplyMSWMSAAttentionSimple instead and come back here when you want to tune.

Categorymodel_patches/unet

Inputs (8)

NameTypeDefaultDescription
input_blocksSTRING1,2Comma-separated list of input blocks to patch. Default is for SD1.x, you can try 4,5 for SDXL
middle_blocksSTRINGComma-separated list of middle blocks to patch. Generally not recommended.
output_blocksSTRING9,10,11Comma-separated list of output blocks to patch. Default is for SD1.x, you can try 3,4,5 for SDXL
time_modeCOMBOpercentTime mode controls how to interpret the values in start_time and end_time.
start_timeFLOAT0.000–999Time the MSW-MSA attention effect starts applying - value is inclusive.
end_timeFLOAT1.000–999Time the MSW-MSA attention effect ends - value is inclusive.
modelMODELModel to patch with the MSW-MSA attention effect.
yaml_parametersoptSTRINGAllows specifying custom parameters via YAML. You can also override any of the normal parameters by key. See main README for possible options. Note: When specifying paramaters this way, there is very little error checking.

Outputs (1)

NameTypeDescription
MODELMODELModel patched with the MSW-MSA attention effect.