Patch Sage Attention KJ (PathchSageAttentionKJ) - ComfyUI-KJNodes

The Patch Sage Attention KJ node implements its functionality by cloning the input model and setting an optimized attention override function within its model options.

The Patch Sage Attention KJ node is an experimental tool designed to dynamically modify the attention calculation mechanism within diffusion models. Its core functionality allows users to globally replace the model's standard attention modules with Sage Attention mode during image generation, enabling exploration of different generative effects or potential performance optimizations.

Node Functionality

This node works by cloning the input model and setting an optimized attention override function within its model options. It does not rely on ComfyUI's standard model patching system but instead directly manipulates the model's transformer_options. The node supports selecting specific attention modes from the Sage Attention library and offers the option to enable PyTorch's torch.compile function to attempt to accelerate computations, though this requires Sage Attention version 2.2.0 or higher.

Node Parameter Description - Patch Sage Attention KJ

Connection Inputs (Inputs)

Parameter Name	Data Type	Required	Default Value	Range/Options	Description
`model`	MODEL	Yes	-	-	Accepts the diffusion model to be modified, serving as the base for the attention mechanism replacement.

Control Parameters (Parameters)

Parameter Name	Data Type	Required	Default Value	Range/Options	Description
`sage_attention`	COMBO	Yes	False	-	Globally patches ComfyUI attention to use sageattn. Once patched, to revert back to normal, you would need to run this node again with the disabled option.
`allow_compile`	BOOLEAN	No	False	-	Allows the use of torch.compile for the sage attention function. Requires the latest sageattn 2.2.0 or higher.

Output

Parameter Name	Data Type	Description
MODEL	MODEL	Outputs a copy of the model with the specified Sage Attention mode patch applied.

Usage Scenarios

In an image generation workflow, you can connect this node after the model loading node and before the sampler node. Enable this node and select the corresponding mode when you wish to experiment with the different stylized effects or potential speed improvements offered by Sage Attention. To restore the standard generation effect, you need to run this node again and set the mode to "disabled".

Notes

Please note that this is an experimental node, and its behavior may be unstable or change with updates to its dependent libraries. Once the patch is applied, it cannot be reverted to the original attention state unless the node is run again with the patch disabled.

Patch Sage Attention KJ Node Source Code Link

The Patch Sage Attention KJ node is from the ComfyUI-KJNodes node package.

Patch Sage Attention KJ(PathchSageAttentionKJ)

Pathch Sage Attention KJ