Apply Flatten Attention

The "attention only" node

Most of ComfyUI-FLATTEN is a big coordinated pipeline: a 3D checkpoint loader, an unsampler, a special KSampler. This node is the odd one out - it takes a plain SD1.5 model and swaps its self-attention for FLATTEN's flow-guided version, with no 3D UNet and no special loader involved. If you already have a working SD1.5 workflow (AnimateDiff, frame-by-frame img2img, whatever) and your problem is frames drifting apart, this is the cheap way to borrow the paper's trick.

FLATTEN - optical FLow-guided ATTENtion - comes from a CVPR 2024 paper about consistent text-to-video editing. The core idea: normal attention looks at spatial neighbors, so a video frame has no idea what the same object looked like two frames ago. FLATTEN instead looks up where each feature came from along the optical-flow path, and attends to those spots in the other frames. Motion consistency, without training a video model.

How it works

The node patches attn1 (self-attention) in a handful of SD1.5 UNet blocks: input blocks 1 and 2, and output blocks 9, 10 and 11. In each of those blocks, the key/value vectors are gathered along the optical-flow trajectories instead of at the current spatial position, so the model can pull information from "the same pixel, one frame earlier" as well as from its spatial neighborhood. The use_old_qk toggle decides whether the query and key come from the original pre-attention tensors (True) or from the attention output (False, the default). The author's recipes lean on this: video editing wants old_qk = 0, the experimental scene-editing path wants old_qk = 1.

The inputs that matter

• model - any SD1.5 model, loaded however you like (you don't need the FLATTEN checkpoint loader for this node).
• trajectories - the TRAJECTORY output from Sample Trajectories in this same pack. This is what carries the optical flow, and it has to be sampled at the same resolution the model actually processes.
• use_old_qk - the one you experiment with (see above).
• input_attn_1, input_attn_2, output_attn_9, output_attn_10, output_attn_11 - which blocks get patched. All default to on. In practice you leave them alone; the only reason to touch them is VRAM pressure or very specific debugging.

The single output is a patched MODEL, which you then wire into whatever sampler you were already using.

Gotchas

FLATTEN's attention wants flow noise, not ordinary noise. That means pairing this node with Create Flow Noise from the same pack - the README is blunt that it "does not always work with methods that add normal noise." And it's hardcoded SD1.5, same as the rest of the pack; there's an sdxl branch but the author calls it a dead end. Don't expect the trajectories to silently re-size themselves, either - they're computed for the dimensions of the input image and must match what the model sees.

Installing

Search ComfyUI-FLATTEN in ComfyUI Manager, or clone it by hand:

cd ComfyUI/custom_nodes
git clone https://github.com/logtd/ComfyUI-FLATTEN

Then restart ComfyUI. The pack's requirements.txt is empty - nothing beyond ComfyUI's own torch/torchvision install is needed, because the optical flow comes from torchvision's built-in RAFT model. The one thing you do need is a real SD 1.5 checkpoint in your models/checkpoints folder.