Loading & Improvement Updates

Updated System Initialization routine.

Updated model loading & patching routines.

Updated import pipelines.

Author: flowstatelabs

FS_Constants.py

@@ -4,14 +4,34 @@
 # Contact: flowstateeng@gmail.com | youtube.com/@flowstateeng
 
 
+
 ##
 # OUTSIDE IMPORTS
 ##
-import os, folder_paths
+import os, folder_paths, importlib
+import nodes
+
 
 ##
 # CONSTANTS
 ##
 MAX_RESOLUTION=16384
 TOOLTIP_UNDERLINE = 32
 
+
+# --- IMPORT KIJAI (THE GOAT) SAGE ATTENTION UNTIL COMFY CORE IMPLEMENTS A NODE
+SAGE_ATTENTION_INSTALLED = False
+
+KJNODES_INSTALLED = "PathchSageAttentionKJ" in nodes.NODE_CLASS_MAPPINGS
+
+SageAttention = None
+
+
+try:
+    importlib.import_module("sageattention")
+    SAGE_ATTENTION_INSTALLED = True
+except:
+    SAGE_ATTENTION_INSTALLED = False
+
+if SAGE_ATTENTION_INSTALLED and KJNODES_INSTALLED:
+    SageAttention = nodes.NODE_CLASS_MAPPINGS['PathchSageAttentionKJ']()

FS_Mappings.py

@@ -36,6 +36,6 @@
 # SYSTEM STATUS
 ##
 for fs_node in NODE_CLASS_MAPPINGS:
-    print(f'\t - ✅ {fs_node}: {NODE_DISPLAY_NAME_MAPPINGS[fs_node]}')
+    print(f'\t - 🟢 {fs_node}: {NODE_DISPLAY_NAME_MAPPINGS[fs_node]}')
 
-print(f'\t - ✅ Mappings Loaded.')
+print(f'\t   - ✅ Mappings Loaded.')

FS_Nodes.py

@@ -21,5 +21,5 @@
 ##
 # SYSTEM STATUS
 ##
-print(f'\t - ✅ All nodes Loaded.')
+print(f'\t   - ✅ All nodes Loaded.')
 

FS_Types.py

@@ -10,8 +10,6 @@
 ##
 from .FS_Assets import *
 from .FS_Constants import *
-from .FS_Types import *
-from .FS_Utils import *
 
 
 ##
@@ -239,12 +237,22 @@ def __ne__(self, __value: object) -> bool:
 ##
 
 # SAGE ATTENTION
-TYPE_SAGE_ATTENTION_MODE = (['disabled'], {'tooltip': (
+enabled_sage_modes = [
+    "disabled",
+    "auto",
+    "sageattn_qk_int8_pv_fp16_cuda",
+    "sageattn_qk_int8_pv_fp16_triton",
+    "sageattn_qk_int8_pv_fp8_cuda",
+    "sageattn_qk_int8_pv_fp8_cuda++"
+] if SAGE_ATTENTION_INSTALLED else ['disabled']
+
+TYPE_SAGE_ATTENTION_MODE = (enabled_sage_modes, {'tooltip': (
     f' Sage Attention Mode\n {"-" * TOOLTIP_UNDERLINE}\n'
     f' - The type of Sage Attention to use.\n'
     f' - This field will only show as "disabled" if you do not have the capability to run Sage Attention.\n\n'
 )})
 
+
 # MODEL
 TYPE_MODEL_FILE_TYPE = (['solo_model', 'checkpoint'], {'tooltip': (
     f' Model File Type\n {"-" * TOOLTIP_UNDERLINE}\n'