Nodes/ControlAltAI Nodes/Flux Attention Cleanup
ComfyUI Node

Flux Attention Cleanup

The node that saves you from tensor-mismatch errors

By gseth·Created 2 years ago·Updated about a year ago· 209
Flux Attention Cleanup
  • any_input
  • message

This is the janitor of the Flux Region pipeline, and skipping it is how you end up with a mysterious tensor-mismatch error you can't explain. The Flux Attention Control node earlier in the chain doesn't just add conditioning - it reaches in and overrides Flux's attention mechanism to enforce your regions. That override sticks around in the model. Switch to a different workflow without undoing it and the leftover attention shape collides with the new graph, and you get a shape/tensor mismatch out of nowhere. Flux Attention Cleanup is what tears the override back down cleanly.

If you only take one thing from this page: when you use the region pipeline, keep this node in the graph. It's not optional plumbing, it's the thing that makes the whole approach safe to run repeatedly.

How it works

The clever part is that it's conditional. You don't actually want to rip out the attention override every single run - if you're iterating in the same workflow, re-running should reuse the setup, not rebuild it from scratch. So this node preserves the overridden attention during re-runs of the existing workflow, but when it detects you've switched to a different workflow, it does a fresh cleanup and restores Flux's original attention.

And it does that without a model unload or a manual cache purge - which matters because, per the author, those don't actually work for this. The override lives somewhere a normal unload doesn't reach, so you need this purpose-built cleanup instead.

The inputs and outputs that matter

  • any_input - a single wildcard input. You wire something from the tail of your pipeline into it (the final image, or another downstream signal) so the node fires at the end of the graph, after everything that needed the overridden attention has run.
  • message - a STRING output reporting what the cleanup did. Wire it into a text/preview node if you want to see the status; it's informational.

Its slot in the chain is dead last: Region Mask Generator → Processor → Validator → Region Mask Conditioning → Flux Attention Control → (Region Overlay Visualizer, optional) → Flux Attention Cleanup.

How to install it

ComfyUI Manager → search ControlAltAI Nodes → install → restart, or:

cd ComfyUI/custom_nodes
git clone https://github.com/gseth/ControlAltAI-Nodes

then restart. This node itself has no special dependencies - but the pipeline it belongs to does: the Flux Attention Control node upstream requires XFormers matched to your PyTorch. If you're running the region chain at all, that has to be in place.

Common issues

The number-one issue is the one this node prevents: leaving it out and then getting Invalid shape for attention bias: torch.Size(...) when you load a different workflow. That's the stale attention override talking. Add Flux Attention Cleanup, wire the pipeline tail into any_input, and it clears itself on the workflow switch.

Don't reach for a manual fix when you hit that error, either - unloading the model or clearing the cache by hand won't undo this particular override (the author is explicit about it), which is exactly why the node exists. And place it at the very end: it needs to run after the attention has done its job for the current generation, so hang it off the last thing in your graph, not the middle.

CategoryControlAltAI Nodes/Flux Region

Inputs (1)

NameTypeDefaultDescription
any_input*

Outputs (1)

NameTypeDescription
messageSTRING