🌊🚒 FlowState Flux Engine - Safe Sage Attention Integration

Sage Attention is now implemented for users with KJ Nodes and Sage Attention library installed. For users missing either of those dependencies, the dropdown menu will simply include a "disabled" option only.

Author: flowstateeng

FlowState_FluxEngine.py

@@ -32,6 +32,7 @@
 from torchvision import transforms
 from PIL import Image, ImageDraw, ImageFont
 import warnings
+import importlib
 
 import comfy.utils
 import comfy.sd
@@ -48,13 +49,35 @@
 from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced
 from comfy_extras.nodes_flux import FluxGuidance
 
-from nodes import EmptyLatentImage
 from nodes import CLIPTextEncode
 
+# --- IMPORT KIJAI (THE GOAT) SAGE ATTENTION UNTIL COMFY CORE IMPLEMENTS A NODE
+SAGE_ATTENTION_INSTALLED = False
+KJNODES_INSTALLED = False
+PatchSageAttention = None
+sageattn_modes = ['disabled']
+
+try:
+    kj_module_path = importlib.import_module('custom_nodes.ComfyUI-KJNodes.nodes.model_optimization_nodes')
+    KJNODES_INSTALLED = True
+    print('\t   - 🟢 KJ Nodes available.')
+except:
+    print('\t   - 🚨 KJNODES NOT AVAILABLE')
+    
+try:
+    importlib.import_module("sageattention")
+    SAGE_ATTENTION_INSTALLED = True
+    print('\t   - 🟢 Sage Attention available.')
+except:
+    print('\t   - 🚨 SAGE ATTENTION NOT AVAILABLE')
+
+
+if SAGE_ATTENTION_INSTALLED and KJNODES_INSTALLED:
+    print('\t   - ✅ KJ Nodes & Sage Attention available.')
+    print('\t      - 🥳🎉 Activating Sage Attention for 🌊🚒 FlowState Flux Engine.')
+    PatchSageAttention = kj_module_path.PathchSageAttentionKJ()
+    sageattn_modes = kj_module_path.sageattn_modes
 
-warnings.filterwarnings('ignore', message='clean_up_tokenization_spaces')
-warnings.filterwarnings('ignore', message='Torch was not compiled with flash attention')
-warnings.filterwarnings('ignore', category=FutureWarning)
 
 
 ##
@@ -85,6 +108,7 @@ def INPUT_TYPES(s):
             'required': {
                 'model_name': TYPE_DIFFUSION_MODELS_LIST(),
                 'weight_dtype': TYPE_WEIGHT_DTYPE,
+                'sage_attention': (sageattn_modes, ),
                 'clip_1_name': TYPE_CLIPS_LIST(),
                 'clip_2_name': TYPE_CLIPS_LIST(),
                 'vae_name': TYPE_VAES_LIST(),
@@ -304,8 +328,8 @@ def sample(self, sampler_components):
         return img_batch_out, latent_batch_out
 
     def execute(
-            self, model_name, weight_dtype, clip_1_name, clip_2_name, vae_name, resolution, orientation, latent_type,
-            custom_width, custom_height, custom_batch_size, image, seed, sampling_algorithm, scheduling_algorithm,
+            self, model_name, weight_dtype, sage_attention, clip_1_name, clip_2_name, vae_name, resolution, orientation,
+            latent_type, custom_width, custom_height, custom_batch_size, image, seed, sampling_algorithm, scheduling_algorithm,
             guidance, steps, denoise, prompt, input_img=None
         ):
 
@@ -321,7 +345,12 @@ def execute(
                 f'\n  - Loading {clip_2_name}...'
                 f'\n  - Loading {vae_name}...\n'
             )
-            self.loaded_model = UNETLoader().load_unet(model_name, weight_dtype)[0]
+            if sage_attention != 'disabled':
+                self.loaded_model = UNETLoader().load_unet(model_name, weight_dtype)[0]
+                self.loaded_model = PatchSageAttention.patch(self.loaded_model, sage_attention)[0]
+            else:
+                self.loaded_model = UNETLoader().load_unet(model_name, weight_dtype)[0]
+
             self.loaded_clip = DualCLIPLoader().load_clip(clip_1_name, clip_2_name, 'flux', 'default')[0]
             self.loaded_vae = VAELoader().load_vae(vae_name)[0]
         else:

README.md

@@ -83,20 +83,20 @@ git clone https://github.com/flowstateeng/FlowState-Creator-Nodes.git
 
 ---
 
-## 🎉 RELEASE v0.0.1 ● 09/14/25 - 5:30 pm (GMT-5)
+## 🎉 RELEASE v0.0.1 ● 09/15/25 - 11:30 am (GMT-5)
 
 ### 📋 Notes
-* Introduction of FlowState Latent Source
+* Introduction of 🌊🚒 FlowState Flux Engine: a Flux all-in-one. (See full description below.)
 
 ### 📆 Upcoming
-* Introduction of FlowState Flux Engine: a Flux all-in-one.
+* LoRA integration for the 🌊🚒 FlowState Flux Engine.
 * Introduction of FlowState Quick Edit: a two-stage node that performs a rapid base edit and then applies a high-detail refinement pass to the result.
 * Introduction of FlowState Asset Forge: efficiently generate professional-grade assets with maximum control & quality.
 
 ---
 
 ## 🖧 Current Nodes
-### [FlowState Latent Source](https://github.com/flowstateeng/FlowState-Creator-Nodes/blob/main/FlowState_LatentSource_.py)
+### [🌊🌱 FlowState Latent Source](https://github.com/flowstateeng/FlowState-Creator-Nodes/blob/main/FlowState_LatentSource.py)
 Simplify your latent options and clean up your workflow with FlowState Latent Source.
 
 A simple switch to select between:
@@ -108,6 +108,20 @@ Vary the denoise on your sampler to control the amount of style transfer you wan
 <p align="center">
   <img width='650' src='https://github.com/flowstateeng/FlowState-Creator-Nodes/blob/main/imgs/FlowState%20Latent%20Source.png' alt='FS Latent Source Image'/>
 </p>
+
+### [🌊🚒 FlowState Flux Engine](https://github.com/flowstateeng/FlowState-Creator-Nodes/blob/main/FlowState_FluxEngine.py)
+Streamline your entire generation workflow with the FlowState Flux Engine.
+
+An all-in-one node that seamlessly combines:
+- Model, VAE, and CLIP loading (with Model Persistence to prevent subsequent reloads.)
+- Prompt conditioning & guidance.
+- Sage Attention model patching (for users with that capability)
+- The integrated FlowState Latent Source for flexible resolution selection & easy i2i style transfer.
+
+Reduce your node clutter and go from concept to final image faster than ever within a single, powerful engine.
+<p align="center">
+  <img width='650' src='https://github.com/flowstateeng/FlowState-Creator-Nodes/blob/main/imgs/FlowState%20Latent%20Source.png' alt='FS Latent Source Image'/>
+</p>
 <br/>
 
 Video tutorials can be found on YouTube at [FlowState Creator Suite Playlist](https://www.youtube.com/playlist?list=PLopF-DMGUFkTulZRkSpRmKFcTENKFicws)

__pycache__/FS_Utils.cpython-313.pyc

@@ -1,7 +1,7 @@
 [project]
 name = "flowstate-creator-nodes"
 description = "The ultimate suite of streamlined ComfyUI nodes for creative professionals."
-version = "0.0.6"
+version = "0.1.0"
 license = {file = "LICENSE"}
 dependencies = [
     "git+https://github.com/kijai/ComfyUI-KJNodes.git"