surrealbydesign/comfyui-krea2-sageattention-guard

ComfyUI is a tool for running AI image-generation models locally through a visual node-based interface. SageAttention is a technique that speeds up how certain models process attention calculations, a key step in generating images. KJNodes is a popular ComfyUI extension that includes a node for applying SageAttention as a speedup. Krea 2 is a locally runnable image-generation model, along with its faster Krea 2 Turbo variant. The problem this repository solves: applying the KJNodes SageAttention patch globally to a Krea 2 workflow can cause it to activate on attention calls that Krea 2's architecture does not support, resulting in crashes, corrupted images, black outputs, or hard-to-trace errors. The root cause is that Krea 2 has a text-fusion path that uses different tensor shapes than the main image-generation attention, and a blanket override trips over those unsupported paths. This repository provides two patch files that modify a single file inside an existing KJNodes installation. The patches add Krea 2 model detection, a guarded version of the SageAttention override that only activates for allowlisted attention calls, and a fallback to ComfyUI's original attention code for anything the speedup cannot handle safely. They also add logging so a user can see which paths are being accelerated and which are being skipped, plus an optional dry-run mode for confirming the patch is active before running a full generation. Installation means applying the two patch files to an existing KJNodes folder using standard git patch commands, not cloning this repo into ComfyUI's custom-nodes directory. The repository also includes ready-to-load workflows for Krea 2 and Krea 2 Turbo as verified baselines, a script to check whether the patches will apply cleanly before touching the live install, and instructions for reverting the changes. Testing was done on Windows with an RTX 5080 GPU, CUDA 12.9, and both SageAttention 1.0.6 and 2.2.0. Linux is untested. Requirements include a CUDA-capable NVIDIA GPU, a working SageAttention install, the Krea 2 model files, and comfort applying git patches to an existing custom node package.