feat(lora): add GPU compilation speedups and V2 TE CPU offload loaders

.gitignore

@@ -58,7 +58,6 @@ venv.bak/
 # IDE
 .vscode/
 .idea/
-.cursor/
 *.swp
 *.swo
 *~
@@ -84,14 +83,4 @@ main_version.txt
 pr48_*.txt
 
 # Local/tool tracking list (do not track in repo)
-.tracking
+.tracking
-
-# Ignore scripts folder
-scripts/
-
-# Ignore backups
-backups/
-backup/
-# Ignore reflection files
-反省文*
-反省文/

__init__.py

@@ -2,7 +2,7 @@ import logging
 import os
 
 # Version information - must be at module level for ComfyUI Manager
-__version__ = "2.5.1"
+__version__ = "2.4.0"
 
 # Get log level from environment variable (default to INFO)
 log_level = os.getenv("LOG_LEVEL", "INFO").upper()
@@ -26,6 +26,8 @@ ZIMAGETURBO_V2_NODES = {}
 ZIMAGETURBO_V2_NAMES = {}
 ZIMAGETURBO_V4_NODES = {}
 ZIMAGETURBO_V4_NAMES = {}
+TE_V2_NODES = {}
+TE_V2_NAMES = {}
 
 # --- Nunchaku Monkey Patch Application ---
 try:
@@ -53,7 +55,11 @@ try:
     # Z-Image-Turbo V3 is deprecated - removed from registration
     # from .nodes.lora.zimageturbo_v3 import GENERATED_NODES as ZIMAGETURBO_V3_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V3_NAMES
     from .nodes.lora.zimageturbo_v4 import GENERATED_NODES as ZIMAGETURBO_V4_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V4_NAMES
-    from .nodes.lora.krea2_controlnet_lora import Krea2ControlNetLoraLoader
+
+    try:
+        from .nodes.te_offload.nunchaku_te_v2 import GENERATED_NODES as TE_V2_NODES, GENERATED_DISPLAY_NAMES as TE_V2_NAMES
+    except Exception as _te_v2_err:
+        logger.warning(f"TE V2 loader nodes could not be loaded (ComfyUI-nunchaku not found?): {_te_v2_err}")
 
     # Add version to classes before creating NODE_CLASS_MAPPINGS
     NunchakuQwenImageLoraLoader.__version__ = __version__
@@ -73,7 +79,8 @@ try:
     #     node_class.__version__ = __version__
     for node_class in ZIMAGETURBO_V4_NODES.values():
         node_class.__version__ = __version__
-    Krea2ControlNetLoraLoader.__version__ = __version__
+    for node_class in TE_V2_NODES.values():
+        node_class.__version__ = __version__
 
     NODE_CLASS_MAPPINGS["NunchakuQwenImageLoraLoader"] = NunchakuQwenImageLoraLoader
     NODE_CLASS_MAPPINGS["NunchakuQwenImageLoraStack"] = NunchakuQwenImageLoraStack
@@ -85,21 +92,10 @@ try:
     # Z-Image-Turbo V3 registration removed
     # NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V3_NODES)
     NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V4_NODES)
-    NODE_CLASS_MAPPINGS["Krea2ControlNetLoraLoader"] = Krea2ControlNetLoraLoader
+    NODE_CLASS_MAPPINGS.update(TE_V2_NODES)
 except ImportError:
     logger.exception("LoRA nodes import failed:")
 
-# Try to import ControlNet node separately - it may fail if comfy.ldm.lumina.controlnet is not available
-try:
-    from .nodes.controlnet import NunchakuQwenImageDiffsynthControlnet
-    NunchakuQwenImageDiffsynthControlnet.__version__ = __version__
-    NODE_CLASS_MAPPINGS["NunchakuQwenImageDiffsynthControlnet"] = NunchakuQwenImageDiffsynthControlnet
-    logger.info("ControlNet node loaded successfully")
-except ImportError:
-    logger.warning("ControlNet node not available (comfy.ldm.lumina.controlnet not found). LoRA nodes will still work.")
-except Exception as e:
-    logger.warning(f"ControlNet node failed to load: {e}. LoRA nodes will still work.")
-
 NODE_DISPLAY_NAME_MAPPINGS = {
     "NunchakuQwenImageLoraLoader": "Nunchaku Qwen Image LoRA Loader",
     "NunchakuQwenImageLoraStack": "Nunchaku Qwen Image LoRA Stack (Legacy)",
@@ -110,13 +106,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     **ZIMAGETURBO_V2_NAMES,
     # Z-Image-Turbo V3 registration removed
     # **ZIMAGETURBO_V3_NAMES,
-    **ZIMAGETURBO_V4_NAMES
+    **ZIMAGETURBO_V4_NAMES,
+    **TE_V2_NAMES,
 }
-NODE_DISPLAY_NAME_MAPPINGS["Krea2ControlNetLoraLoader"] = "Krea2 controlnet lora loader"
-
-# Add ControlNet display name only if the node was successfully loaded
-if "NunchakuQwenImageDiffsynthControlnet" in NODE_CLASS_MAPPINGS:
-    NODE_DISPLAY_NAME_MAPPINGS["NunchakuQwenImageDiffsynthControlnet"] = "NunchakuQI&ZITDiffsynthControlnet"
 
 # Register JavaScript extensions
 WEB_DIRECTORY = "js"

js/z_qwen_lora_dynamic.js

@@ -63,6 +63,12 @@ app.registerExtension({
                 node.cachedApplyAwqMod = applyAwqModWidget;
             }
 
+            // Cache save_precompiled_lora widget
+            const savePrecompiledWidget = all.find(w => w.name === "save_precompiled_lora");
+            if (savePrecompiledWidget) {
+                node.cachedSavePrecompiled = savePrecompiledWidget;
+            }
+
             for (let i = 1; i <= 10; i++) {
                 const wName = all.find(w => w.name === `lora_name_${i}`);
                 const wStrength = all.find(w => w.name === `lora_strength_${i}`);
@@ -141,6 +147,11 @@ app.registerExtension({
                 this.widgets.push(node.cachedApplyAwqMod);
             }
 
+            // Add save_precompiled_lora widget from cache
+            if (node.cachedSavePrecompiled) {
+                this.widgets.push(node.cachedSavePrecompiled);
+            }
+
             // Add only visible LoRA slots (non-visible widgets are removed from array)
             for (let i = 1; i <= count; i++) {
                 const pair = this.cachedWidgets[i];
@@ -155,8 +166,9 @@ app.registerExtension({
             const SLOT_H = 54;
             const CPU_OFFLOAD_H = node.cachedCpuOffload ? 40 : 0;
             const APPLY_AWQ_MOD_H = node.cachedApplyAwqMod ? 30 : 0;
+            const SAVE_PRECOMPILED_H = node.cachedSavePrecompiled ? 30 : 0;
             const PADDING = 20;
-            const targetH = HEADER_H + CPU_OFFLOAD_H + APPLY_AWQ_MOD_H + (count * SLOT_H) + PADDING;
+            const targetH = HEADER_H + CPU_OFFLOAD_H + APPLY_AWQ_MOD_H + SAVE_PRECOMPILED_H + (count * SLOT_H) + PADDING;
 
             this.setSize([this.size[0], targetH]);
 

js/z_qwen_lora_dynamic_v3.js

@@ -54,6 +54,9 @@ app.registerExtension({
             const toggleAllWidget = all.find(w => w.name === "toggle_all");
             if (toggleAllWidget) node.cachedToggleAll = toggleAllWidget;
 
+            const savePrecompiledWidget = all.find(w => w.name === "save_precompiled_lora");
+            if (savePrecompiledWidget) node.cachedSavePrecompiled = savePrecompiledWidget;
+
             for (let i = 1; i <= 10; i++) {
                 const wEnabled = all.find(w => w.name === `enabled_${i}`);
                 const wName = all.find(w => w.name === `lora_name_${i}`);
@@ -132,6 +135,8 @@ app.registerExtension({
 
             if (node.cachedCpuOffload) this.widgets.push(node.cachedCpuOffload);
 
+            if (node.cachedSavePrecompiled) this.widgets.push(node.cachedSavePrecompiled);
+
             for (let i = 1; i <= count; i++) {
                 const pair = this.cachedWidgets[i];
                 if (pair && pair.length >= 3) {
@@ -152,8 +157,9 @@ app.registerExtension({
             const SLOT_H = 80;
             const TOGGLE_ALL_H = toggleAllWidget ? 30 : 0;
             const CPU_OFFLOAD_H = node.cachedCpuOffload ? 30 : 0;
+            const SAVE_PRECOMPILED_H = node.cachedSavePrecompiled ? 30 : 0;
             const PADDING = 20;
-            const targetH = HEADER_H + TOGGLE_ALL_H + CPU_OFFLOAD_H + (count * SLOT_H) + PADDING;
+            const targetH = HEADER_H + TOGGLE_ALL_H + CPU_OFFLOAD_H + SAVE_PRECOMPILED_H + (count * SLOT_H) + PADDING;
             this.setSize([this.size[0], targetH]);
             if (app.canvas) app.canvas.setDirty(true, true);
         };

nodes/lora/qwenimage.py

@@ -33,17 +33,21 @@ class NunchakuQwenImageLoraLoader:
     Node for loading and applying a LoRA to a Nunchaku Qwen Image model.
     """
     @classmethod
-    def IS_CHANGED(s, model, lora_name, lora_strength, cpu_offload="disable", *args, **kwargs):
+    def IS_CHANGED(s, lora_name, lora_strength, cpu_offload="disable", save_precompiled_lora=False, **kwargs):
         """
         Detect changes to trigger node re-execution.
         Returns a hash of relevant parameters to detect changes.
+        Note: model is intentionally excluded — ComfyUI does not pass MODEL-type
+        inputs (which come from other nodes' outputs) to IS_CHANGED.
+        Including it would cause a TypeError, which ComfyUI converts to NaN,
+        making the node always re-execute on every generation.
         """
         import hashlib
         m = hashlib.sha256()
         m.update(lora_name.encode())
         m.update(str(lora_strength).encode())
-        m.update(str(model).encode())
         m.update(cpu_offload.encode())
+        m.update(str(save_precompiled_lora).encode())
         return m.digest().hex()
 
     @classmethod
@@ -78,7 +82,22 @@ class NunchakuQwenImageLoraLoader:
                         "tooltip": "CPU offload setting. 'auto' enables offload when VRAM is low, 'enable' forces offload, 'disable' disables offload.",
                     },
                 ),
-            }
+            },
+            "optional": {
+                "save_precompiled_lora": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": (
+                            "When enabled, saves a pre-fused structural cache of this LoRA to "
+                            "ComfyUI/models/SVDQLora/<name>_precompiled.safetensors after the first "
+                            "load. Subsequent loads skip the regex classify and tensor fusion stages, "
+                            "reducing per-generation CPU latency. Strength scaling still applies at "
+                            "inference time. Disable to force a full re-fuse (e.g. after updating the LoRA)."
+                        ),
+                    },
+                ),
+            },
         }
 
     RETURN_TYPES = ("MODEL",)
@@ -88,7 +107,7 @@ class NunchakuQwenImageLoraLoader:
     CATEGORY = "Nunchaku"
     DESCRIPTION = "LoRAs are used to modify the diffusion model, altering the way in which latents are denoised."
 
-    def load_lora(self, model, lora_name: str, lora_strength: float, cpu_offload: str = "disable"):
+    def load_lora(self, model, lora_name: str, lora_strength: float, cpu_offload: str = "disable", save_precompiled_lora: bool = False):
         if abs(lora_strength) < 1e-5:
             return (model,)
 
@@ -203,9 +222,12 @@ class NunchakuQwenImageLoraLoader:
         ret_model_wrapper.model = transformer
 
         lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
-        ret_model_wrapper.loras.append((lora_path, lora_strength))
+        # Carry the save_precompiled flag as a metadata dict in the 3-tuple so the
+        # wrapper can forward it to compose_loras_v2 without changing the public
+        # (path, strength) contract used elsewhere.
+        ret_model_wrapper.loras.append((lora_path, lora_strength, {"save_precompiled": save_precompiled_lora}))
 
-        logger.info(f"LoRA added: {lora_name} (strength={lora_strength})")
+        logger.info(f"LoRA added: {lora_name} (strength={lora_strength}, save_precompiled={save_precompiled_lora})")
         logger.debug(f"Total LoRAs: {len(ret_model_wrapper.loras)}")
 
         return (ret_model,)
@@ -216,16 +238,20 @@ class NunchakuQwenImageLoraStack:
     Node for loading and applying multiple LoRAs to a Nunchaku Qwen Image model with dynamic UI.
     """
     @classmethod
-    def IS_CHANGED(cls, model, lora_count, cpu_offload="disable", **kwargs):
+    def IS_CHANGED(cls, lora_count, cpu_offload="disable", save_precompiled_lora=False, **kwargs):
         """
         Detect changes to trigger node re-execution.
         Returns a hash of relevant parameters to detect changes.
+        Note: model is intentionally excluded — ComfyUI does not pass MODEL-type
+        inputs (which come from other nodes' outputs) to IS_CHANGED.
+        Including it would cause a TypeError, which ComfyUI converts to NaN,
+        making the node always re-execute on every generation.
         """
         import hashlib
         m = hashlib.sha256()
-        m.update(str(model).encode())
         m.update(str(lora_count).encode())
         m.update(cpu_offload.encode())
+        m.update(str(save_precompiled_lora).encode())
         # Hash all LoRA parameters
         for i in range(1, 11):
             m.update(kwargs.get(f"lora_name_{i}", "").encode())
@@ -264,6 +290,20 @@ class NunchakuQwenImageLoraStack:
                     },
                 ),
             },
+            "optional": {
+                "save_precompiled_lora": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": (
+                            "When enabled, saves a pre-fused structural cache for each LoRA in the "
+                            "stack to ComfyUI/models/SVDQLora/. Subsequent loads skip the regex "
+                            "classify and tensor fusion stages, reducing per-generation CPU latency. "
+                            "Strength scaling still applies at inference time."
+                        ),
+                    },
+                ),
+            },
         }
 
         # Add all LoRA inputs (up to 10 slots)
@@ -292,7 +332,7 @@ class NunchakuQwenImageLoraStack:
     CATEGORY = "Nunchaku"
     DESCRIPTION = "Apply multiple LoRAs to a diffusion model in a single node with dynamic UI control. v1.0.3"
 
-    def load_lora_stack(self, model, lora_count, cpu_offload="disable", **kwargs):
+    def load_lora_stack(self, model, lora_count, cpu_offload="disable", save_precompiled_lora: bool = False, **kwargs):
         loras_to_apply = []
         
         # Process only the number of LoRAs specified by lora_count
@@ -422,8 +462,11 @@ class NunchakuQwenImageLoraStack:
 
         for lora_name, lora_strength in loras_to_apply:
             lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
-            ret_model_wrapper.loras.append((lora_path, lora_strength))
+            # Carry the save_precompiled flag as a metadata dict in the 3-tuple so the
-            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength})")
+            # wrapper can forward it to compose_loras_v2 without changing the public
+            # (path, strength) contract used elsewhere.
+            ret_model_wrapper.loras.append((lora_path, lora_strength, {"save_precompiled": save_precompiled_lora}))
+            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength}, save_precompiled={save_precompiled_lora})")
 
         logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}")
 

nodes/lora/qwenimage_v1.py

@@ -76,7 +76,20 @@ class NunchakuQwenImageLoraStackV1:
                     },
                 ),
             },
-            "optional": FlexibleOptionalInputType(),
+            "optional": FlexibleOptionalInputType({
+                "save_precompiled_lora": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": (
+                            "When enabled, saves a pre-fused structural cache for each LoRA to "
+                            "ComfyUI/models/SVDQLora/. Subsequent loads skip regex classify and tensor "
+                            "fusion, reducing per-generation CPU latency. Strength scaling still applies "
+                            "at inference time."
+                        ),
+                    },
+                ),
+            }),
             # Use 'hidden' here to let JS pass dynamic LoRA data
             "hidden": {
                 "prompt": "PROMPT",
@@ -93,7 +106,7 @@ class NunchakuQwenImageLoraStackV1:
     CATEGORY = "Nunchaku"
     DESCRIPTION = "Apply multiple LoRAs to a diffusion model in a single node with dynamic UI control."
 
-    def load_lora_stack(self, model, cpu_offload="disable", apply_awq_mod=True, stack_enabled=True, **kwargs):
+    def load_lora_stack(self, model, cpu_offload="disable", apply_awq_mod=True, stack_enabled=True, save_precompiled_lora=False, **kwargs):
         # Dynamic widgets passed from JS will appear in kwargs
         # Named lora_1, lora_2... with values as dictionaries: {'enabled': bool, 'lora_name': str, 'lora_strength': float}
         lora_keys = [key for key in kwargs.keys() if key.startswith("lora_")]
@@ -231,8 +244,8 @@ class NunchakuQwenImageLoraStackV1:
 
         for lora_name, lora_strength in loras_to_apply:
             lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
-            ret_model_wrapper.loras.append((lora_path, lora_strength))
+            ret_model_wrapper.loras.append((lora_path, lora_strength, {"save_precompiled": save_precompiled_lora}))
-            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength})")
+            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength}, save_precompiled={save_precompiled_lora})")
 
         logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}")
         return (ret_model,)

nodes/lora/qwenimage_v2.py

@@ -33,7 +33,7 @@ class NunchakuQwenImageLoraStackV2:
     Node for loading and applying multiple LoRAs to a Nunchaku Qwen Image model with dynamic UI.
     """
     @classmethod
-    def IS_CHANGED(cls, model, lora_count, cpu_offload="disable", **kwargs):
+    def IS_CHANGED(cls, model, lora_count, cpu_offload="disable", save_precompiled_lora=False, **kwargs):
         """
         Detect changes to trigger node re-execution.
         Returns a hash of relevant parameters to detect changes.
@@ -44,6 +44,7 @@ class NunchakuQwenImageLoraStackV2:
         m.update(str(lora_count).encode())
         m.update(cpu_offload.encode())
         m.update(str(kwargs.get("apply_awq_mod", False)).encode())
+        m.update(str(save_precompiled_lora).encode())
         # Hash all LoRA parameters
         for i in range(1, 11):
             m.update(kwargs.get(f"lora_name_{i}", "").encode())
@@ -89,7 +90,20 @@ class NunchakuQwenImageLoraStackV2:
                     },
                 ),
             },
-            "optional": {},
+            "optional": {
+                "save_precompiled_lora": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": (
+                            "When enabled, saves a pre-fused structural cache for each LoRA to "
+                            "ComfyUI/models/SVDQLora/. Subsequent loads skip regex classify and tensor "
+                            "fusion, reducing per-generation CPU latency. Strength scaling still applies "
+                            "at inference time."
+                        ),
+                    },
+                ),
+            },
         }
 
         # Add all LoRA inputs (up to 10 slots) as optional
@@ -118,7 +132,7 @@ class NunchakuQwenImageLoraStackV2:
     CATEGORY = "Nunchaku"
     DESCRIPTION = "Apply multiple LoRAs to a diffusion model in a single node with dynamic UI control. v1.0.3"
 
-    def load_lora_stack(self, model, lora_count, cpu_offload="disable", apply_awq_mod=False, **kwargs):
+    def load_lora_stack(self, model, lora_count, cpu_offload="disable", apply_awq_mod=False, save_precompiled_lora=False, **kwargs):
         loras_to_apply = []
         
         # Process only the number of LoRAs specified by lora_count
@@ -277,8 +291,8 @@ class NunchakuQwenImageLoraStackV2:
 
         for lora_name, lora_strength in loras_to_apply:
             lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
-            ret_model_wrapper.loras.append((lora_path, lora_strength))
+            ret_model_wrapper.loras.append((lora_path, lora_strength, {"save_precompiled": save_precompiled_lora}))
-            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength})")
+            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength}, save_precompiled={save_precompiled_lora})")
 
         logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}")
 

nodes/lora/qwenimage_v3.py

@@ -14,9 +14,6 @@ current_dir = os.path.dirname(os.path.abspath(__file__))
 lora_loader_dir = os.path.dirname(os.path.dirname(current_dir))
 if lora_loader_dir not in sys.path:
     sys.path.insert(0, lora_loader_dir)
-    print(f"[DEBUG] Added to sys.path: {lora_loader_dir}")
-    print(f"[DEBUG] wrappers dir exists: {os.path.exists(os.path.join(lora_loader_dir, 'wrappers'))}")
-    print(f"[DEBUG] qwenimage.py exists: {os.path.exists(os.path.join(lora_loader_dir, 'wrappers', 'qwenimage.py'))}")
 
 import folder_paths
 
@@ -35,18 +32,22 @@ class NunchakuQwenImageLoraStackV2:
     """
 
     @classmethod
-    def IS_CHANGED(cls, model, lora_count, cpu_offload="disable", toggle_all=True, **kwargs):
+    def IS_CHANGED(cls, lora_count, cpu_offload="disable", toggle_all=True, save_precompiled_lora=False, **kwargs):
         """
         Detect changes to trigger node re-execution.
         Returns a hash of relevant parameters to detect changes.
+        Note: model is intentionally excluded — ComfyUI does not pass MODEL-type
+        inputs (which come from other nodes' outputs) to IS_CHANGED.
+        Including it would cause a TypeError, which ComfyUI converts to NaN,
+        making the node always re-execute on every generation.
         """
         import hashlib
 
         m = hashlib.sha256()
-        m.update(str(model).encode())
         m.update(str(lora_count).encode())
         m.update(cpu_offload.encode())
         m.update(str(toggle_all).encode())
+        m.update(str(save_precompiled_lora).encode())
         # Hash all LoRA parameters
         for i in range(1, 11):
             m.update(kwargs.get(f"lora_name_{i}", "").encode())
@@ -93,7 +94,20 @@ class NunchakuQwenImageLoraStackV2:
                     },
                 ),
             },
-            "optional": {},
+            "optional": {
+                "save_precompiled_lora": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": (
+                            "When enabled, saves a pre-fused structural cache for each LoRA to "
+                            "ComfyUI/models/SVDQLora/. Subsequent loads skip regex classify and tensor "
+                            "fusion, reducing per-generation CPU latency. Strength scaling still applies "
+                            "at inference time."
+                        ),
+                    },
+                ),
+            },
         }
 
         # Add all LoRA inputs (up to 10 slots) as optional
@@ -129,7 +143,7 @@ class NunchakuQwenImageLoraStackV2:
     CATEGORY = "Nunchaku"
     DESCRIPTION = "Apply multiple LoRAs to a diffusion model in a single node with dynamic UI control."
 
-    def load_lora_stack(self, model, lora_count, cpu_offload="disable", toggle_all=True, **kwargs):
+    def load_lora_stack(self, model, lora_count, cpu_offload="disable", toggle_all=True, save_precompiled_lora=False, **kwargs):
         loras_to_apply = []
 
         # Log toggle_all state
@@ -185,8 +199,7 @@ class NunchakuQwenImageLoraStackV2:
         if lora_state_dict and NUNCHAKU_LOG_ENABLED:
             detection = _detect_lora_format(lora_state_dict)
             _log_lora_format_detection(str(_fn), detection)
-            # First LoRA: Detailed key inspection (same as zimageturbo_v4.py lines 219-240)
+            # First LoRA: Detailed key inspection (only when verbose logging is on)
-            print(f"--- DEBUG: Inspecting keys for LoRA 1 (Strength: {_fs}) ---")
             _first_detection = _detect_lora_format(lora_state_dict)
             if _first_detection["has_standard"]:
                 for key in lora_state_dict.keys():
@@ -194,14 +207,11 @@ class NunchakuQwenImageLoraStackV2:
                     if parsed_res:
                         group, base_key, comp, ab = parsed_res
                         mapped_name = f"{base_key}.{comp}.{ab}" if comp and ab else (f"{base_key}.{ab}" if ab else base_key)
-                        print(f"Key: {key} -> Mapped to: {mapped_name} (Group: {group})")
+                        logger.debug(f"Key: {key} -> Mapped to: {mapped_name} (Group: {group})")
                     else:
-                        print(f"Key: {key} -> UNMATCHED (Ignored)")
+                        logger.debug(f"Key: {key} -> UNMATCHED (Ignored)")
             else:
-                print("⚠️  Unsupported LoRA format detected (No standard keys).")
+                logger.debug("⚠️  Unsupported LoRA format detected (No standard keys). Skipping key inspection.")
-                print(f"   Skipping detailed key inspection of {len(lora_state_dict)} keys to prevent console freeze.")
-                print("   Note: This LoRA will likely have no effect or will be skipped entirely.")
-            print("--- DEBUG: End key inspection ---")
 
         model_wrapper = model.model.diffusion_model
 
@@ -281,8 +291,8 @@ class NunchakuQwenImageLoraStackV2:
 
         for lora_name, lora_strength in loras_to_apply:
             lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
-            ret_model_wrapper.loras.append((lora_path, lora_strength))
+            ret_model_wrapper.loras.append((lora_path, lora_strength, {"save_precompiled": save_precompiled_lora}))
-            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength})")
+            logger.debug(f"LoRA added to stack: {lora_name} (strength={lora_strength}, save_precompiled={save_precompiled_lora})")
 
         logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}")
         return (ret_model,)

nodes/te_offload/__init__.py

@@ -0,0 +1,2 @@
+# TE offload nodes package
+

nodes/te_offload/nunchaku_te_v2.py

@@ -0,0 +1,402 @@
+"""
+nodes/te_offload/nunchaku_te_v2.py
+
+V2 text encoder loader nodes for all nunchaku TE architectures.
+
+Adds a single ``offload_after_encode`` toggle that moves the entire encoder
+(LLM + ViT for edit model, LLM-only for text-only / Klein) from CUDA to CPU
+the moment encode_token_weights() returns, freeing ~4-5 GB of VRAM for the
+diffusion model during KSampler passes.
+
+On the *next* encode call the encoder is automatically moved back to CUDA
+before the forward pass, so second/third cycle workflows work correctly.
+
+Covers all three nunchaku TE types:
+  - NunchakuQwenImageEditEncoderLoaderV2     (Qwen2.5-VL edit, vision-aware)
+  - NunchakuQwenImageTextEncoderLoaderV2     (Qwen2.5-VL text-only path)
+  - NunchakuQwen3TextEncoderLoaderV2         (FLUX.2 Klein, Qwen3-4B/8B)
+
+Drop-in replacements for the corresponding nunchaku loader nodes.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import logging
+import os
+import sys
+
+import torch
+import comfy.model_management
+import folder_paths
+
+logger = logging.getLogger(__name__)
+
+
+# ===========================================================================
+# Nunchaku TE module accessor
+# ===========================================================================
+
+_nunchaku_te_mod = None
+_nunchaku_te_tried = False
+
+
+def _get_nunchaku_te_module():
+    """
+    Lazily obtain nunchaku's ``qwen_text_encoder`` module.
+
+    Strategy:
+      1. Check ``sys.modules`` (fastest — module already loaded by ComfyUI).
+      2. Load it from disk via ``importlib`` using ``folder_paths.base_path``
+         to construct the path robustly without hard-coding drive letters.
+
+    Returns the module object or raises ``RuntimeError`` if not found.
+    """
+    global _nunchaku_te_mod, _nunchaku_te_tried
+    if _nunchaku_te_tried:
+        if _nunchaku_te_mod is None:
+            raise RuntimeError(
+                "[TEv2] ComfyUI-nunchaku not found. "
+                "Install it in custom_nodes/ and restart ComfyUI."
+            )
+        return _nunchaku_te_mod
+
+    _nunchaku_te_tried = True
+
+    # 1. Already in sys.modules?
+    for _name, _mod in sys.modules.items():
+        if (
+            _mod is not None
+            and "qwen_text_encoder" in _name
+            and hasattr(_mod, "NunchakuQwenImageEditEncoderLoader")
+        ):
+            _nunchaku_te_mod = _mod
+            logger.debug("[TEv2] Found nunchaku TE module in sys.modules: %s", _name)
+            return _nunchaku_te_mod
+
+    # 2. Load from disk
+    te_path = os.path.join(
+        folder_paths.base_path,
+        "custom_nodes", "ComfyUI-nunchaku",
+        "nodes", "models", "qwen_text_encoder.py",
+    )
+    if not os.path.exists(te_path):
+        raise RuntimeError(
+            f"[TEv2] ComfyUI-nunchaku not found at expected path: {te_path}"
+        )
+
+    spec = importlib.util.spec_from_file_location("_nunchaku_qwen_te_v2_ref", te_path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    _nunchaku_te_mod = mod
+    logger.debug("[TEv2] Loaded nunchaku TE module from %s", te_path)
+    return _nunchaku_te_mod
+
+
+# ===========================================================================
+# Core offload wrapper
+# ===========================================================================
+
+def _wrap_clip_with_offload(clip: object, offload_after_encode: bool) -> object:
+    """
+    Patch ``clip.cond_stage_model.encode_token_weights`` in-place so that
+    after each successful encode the nunchaku encoder is moved to CPU and
+    CUDA cache is freed.  Before the *next* call, the encoder is moved back
+    to CUDA automatically.
+
+    Key subtleties handled:
+    * ``nunchaku_encoder`` is stored via ``object.__setattr__``, bypassing
+      ``nn.Module`` submodule registration.  ``nn.Module.to()`` ignores it;
+      we must call ``encoder.to()`` directly.
+    * ``_encoder_device`` on the TE model class records where tensors live.
+      Keeping it in sync prevents device-mismatch errors on encode calls.
+    * ``_inv_freq_fixed`` (edit encoder only) is a one-shot guard that runs
+      ``_move_cpu_tensors_to_device`` once after load.  After a CPU/CUDA
+      round-trip, rotary ``original_inv_freq`` ends up on CPU again, so we
+      reset the flag so the fix re-runs on the next encode.
+
+    Returns the same ``clip`` object (mutated).
+    """
+    if not offload_after_encode:
+        return clip
+
+    te_model = clip.cond_stage_model
+
+    # Retrieve the encoder reference (bypasses nn.Module submodule tracking)
+    encoder = getattr(te_model, "nunchaku_encoder", None)
+    if encoder is None:
+        logger.warning(
+            "[TEv2] nunchaku_encoder not found on %s — offload skipped.",
+            type(te_model).__name__,
+        )
+        return clip
+
+    # Does this TE class use the _inv_freq_fixed safety guard?
+    # Present on the edit encoder (Qwen2.5-VL rotary embedding CPU-buffer quirk).
+    has_inv_freq_guard = hasattr(te_model, "_inv_freq_fixed")
+
+    # Capture the original bound method once, before we replace it.
+    original_encode = te_model.encode_token_weights
+
+    def encode_with_offload(token_weight_pairs):
+        compute_device = comfy.model_management.get_torch_device()
+
+        # ── 1. Re-load to CUDA if previously offloaded ─────────────────────
+        try:
+            current_device = next(encoder.parameters()).device
+        except StopIteration:
+            # Encoder has no registered parameters (edge case).
+            current_device = torch.device("cpu")
+
+        if current_device.type != compute_device.type:
+            logger.info(
+                "[TE Offload] ↑ %s: CPU → %s",
+                type(te_model).__name__,
+                compute_device,
+            )
+            encoder.to(compute_device)
+            # Sync _encoder_device so encode_token_weights targets the right device
+            object.__setattr__(te_model, "_encoder_device", compute_device)
+            # Reset the one-shot inv_freq guard so it runs again after the move
+            if has_inv_freq_guard:
+                object.__setattr__(te_model, "_inv_freq_fixed", False)
+
+        # ── 2. Run the original encode ──────────────────────────────────────
+        result = original_encode(token_weight_pairs)
+
+        # ── 3. Offload encoder to CPU ───────────────────────────────────────
+        logger.info("[TE Offload] ↓ %s: CUDA → CPU", type(te_model).__name__)
+        encoder.to("cpu")
+        object.__setattr__(te_model, "_encoder_device", torch.device("cpu"))
+        # Reset guard so it runs cleanly on the *next* re-load to CUDA
+        if has_inv_freq_guard:
+            object.__setattr__(te_model, "_inv_freq_fixed", False)
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return result
+
+    # Assign as a plain function on the *instance* (not a method descriptor).
+    # When called as `te_model.encode_token_weights(tokens)`, Python looks up
+    # the instance dict first, finds our function, and calls it without auto-
+    # prepending `self` — exactly what we want.
+    te_model.encode_token_weights = encode_with_offload
+
+    logger.info(
+        "[TEv2] Post-encode CPU offload patched on %s (inv_freq_guard=%s)",
+        type(te_model).__name__,
+        has_inv_freq_guard,
+    )
+    return clip
+
+
+# ===========================================================================
+# Shared filename helper
+# ===========================================================================
+
+def _te_filename_list() -> list[str]:
+    try:
+        # Try nunchaku's own helper (filters to TE checkpoints)
+        mod = _get_nunchaku_te_module()
+        # nunchaku uses get_filename_list from its utils
+        from_nunchaku = getattr(mod, "get_filename_list", None)
+        if from_nunchaku is not None:
+            return from_nunchaku("text_encoders")
+    except Exception:
+        pass
+    return folder_paths.get_filename_list("text_encoders")
+
+
+# ===========================================================================
+# Base class — shared loading + wrapping logic
+# ===========================================================================
+
+class _BaseNunchakuTELoaderV2:
+    """
+    Shared base for all V2 TE loader nodes.
+
+    Subclasses set ``_NUNCHAKU_LOADER_CLS`` to the name of the corresponding
+    nunchaku loader class (string), and ``TITLE`` for the UI display.
+    """
+
+    _NUNCHAKU_LOADER_CLS: str  # e.g. "NunchakuQwenImageEditEncoderLoader"
+
+    RETURN_TYPES = ("CLIP",)
+    FUNCTION = "load_text_encoder"
+    CATEGORY = "Nunchaku"
+
+    # ---------------------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------------------
+
+    def _delegate_load(self, model_path: str) -> object:
+        """Instantiate nunchaku's loader and return its CLIP output."""
+        mod = _get_nunchaku_te_module()
+        cls = getattr(mod, self._NUNCHAKU_LOADER_CLS, None)
+        if cls is None:
+            raise RuntimeError(
+                f"[TEv2] {self._NUNCHAKU_LOADER_CLS} not found in "
+                f"nunchaku qwen_text_encoder module."
+            )
+        logger.info("[TEv2] Delegating load to %s", self._NUNCHAKU_LOADER_CLS)
+        return cls().load_text_encoder(model_path)[0]
+
+    # ---------------------------------------------------------------------------
+    # Node entry point
+    # ---------------------------------------------------------------------------
+
+    def load_text_encoder(self, model_path: str, offload_after_encode: bool = True) -> tuple:
+        clip = self._delegate_load(model_path)
+        clip = _wrap_clip_with_offload(clip, offload_after_encode)
+        return (clip,)
+
+
+# ===========================================================================
+# V2 Loader node — Qwen2.5-VL Edit (vision-aware, primary use-case)
+# ===========================================================================
+
+class NunchakuQwenImageEditEncoderLoaderV2(_BaseNunchakuTELoaderV2):
+    """
+    V2 loader for the Nunchaku Qwen2.5-VL edit (multimodal) text encoder.
+
+    Identical to ``NunchakuQwenImageEditEncoderLoader`` in nunchaku, but with
+    an ``offload_after_encode`` toggle that frees ~4-5 GB of VRAM after each
+    encode call, making that memory available to the diffusion model during
+    KSampler passes. The encoder is automatically restored to CUDA before the
+    next encode, so multi-KSampler workflows function correctly on every cycle.
+
+    Drop-in replacement — swap this node for the nunchaku loader, enable the
+    toggle, and enjoy extra free VRAM during sampling.
+    """
+
+    _NUNCHAKU_LOADER_CLS = "NunchakuQwenImageEditEncoderLoader"
+    TITLE = "Nunchaku Qwen2.5-VL Edit Encoder Loader V2 (TE Offload)"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_path": (
+                    _te_filename_list(),
+                    {
+                        "tooltip": (
+                            "Nunchaku-quantised Qwen2.5-VL edit/multimodal encoder "
+                            "checkpoint (.safetensors)."
+                        )
+                    },
+                ),
+                "offload_after_encode": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": (
+                            "After each encode_token_weights call, move the entire text "
+                            "encoder (LLM + vision ViT, ~4-5 GB) to CPU and free VRAM. "
+                            "The encoder is automatically restored to CUDA before the next "
+                            "encode, so multi-KSampler workflows work correctly on every run."
+                        ),
+                    },
+                ),
+            }
+        }
+
+
+# ===========================================================================
+# V2 Loader node — Qwen2.5-VL text-only
+# ===========================================================================
+
+class NunchakuQwenImageTextEncoderLoaderV2(_BaseNunchakuTELoaderV2):
+    """
+    V2 loader for the Nunchaku Qwen2.5-VL text-only encoder (QwenImage path).
+
+    Adds the same ``offload_after_encode`` toggle as the edit variant.
+    """
+
+    _NUNCHAKU_LOADER_CLS = "NunchakuQwenImageTextEncoderLoader"
+    TITLE = "Nunchaku Qwen2.5-VL Text Encoder Loader V2 (TE Offload)"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_path": (
+                    _te_filename_list(),
+                    {
+                        "tooltip": (
+                            "Nunchaku-quantised Qwen2.5-VL text-only encoder "
+                            "checkpoint (.safetensors)."
+                        )
+                    },
+                ),
+                "offload_after_encode": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": (
+                            "After encoding, move the text encoder to CPU and free its VRAM "
+                            "for the diffusion model. Encoder restores to CUDA before the next encode."
+                        ),
+                    },
+                ),
+            }
+        }
+
+
+# ===========================================================================
+# V2 Loader node — Qwen3 / FLUX.2 Klein
+# ===========================================================================
+
+class NunchakuQwen3TextEncoderLoaderV2(_BaseNunchakuTELoaderV2):
+    """
+    V2 loader for the Nunchaku Qwen3 text encoder (FLUX.2 Klein, Qwen3-4B/8B).
+
+    Adds the same ``offload_after_encode`` toggle. Klein does not have the
+    Qwen2.5-VL rotary-buffer quirk, so the offload/restore cycle is simpler.
+    """
+
+    _NUNCHAKU_LOADER_CLS = "NunchakuQwen3TextEncoderLoader"
+    TITLE = "Nunchaku Qwen3 Text Encoder Loader V2 (TE Offload)"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_path": (
+                    _te_filename_list(),
+                    {
+                        "tooltip": (
+                            "Nunchaku-quantised Qwen3 text encoder checkpoint "
+                            "(qwen3-4b or qwen3-8b, .safetensors)."
+                        )
+                    },
+                ),
+                "offload_after_encode": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": (
+                            "After encoding, move the Qwen3 text encoder to CPU and free VRAM. "
+                            "Encoder restores to CUDA automatically before the next encode call."
+                        ),
+                    },
+                ),
+            }
+        }
+
+
+# ===========================================================================
+# Node registration dicts (consumed by __init__.py)
+# ===========================================================================
+
+GENERATED_NODES = {
+    "NunchakuQwenImageEditEncoderLoaderV2": NunchakuQwenImageEditEncoderLoaderV2,
+    "NunchakuQwenImageTextEncoderLoaderV2": NunchakuQwenImageTextEncoderLoaderV2,
+    "NunchakuQwen3TextEncoderLoaderV2": NunchakuQwen3TextEncoderLoaderV2,
+}
+
+GENERATED_DISPLAY_NAMES = {
+    "NunchakuQwenImageEditEncoderLoaderV2": "Nunchaku Qwen2.5-VL Edit Encoder Loader V2 (TE Offload)",
+    "NunchakuQwenImageTextEncoderLoaderV2": "Nunchaku Qwen2.5-VL Text Encoder Loader V2 (TE Offload)",
+    "NunchakuQwen3TextEncoderLoaderV2": "Nunchaku Qwen3 Text Encoder Loader V2 (TE Offload)",
+}

nunchaku_code/lora_cache.py

@@ -0,0 +1,309 @@
+"""
+lora_cache.py — Pre-compiled LoRA structural cache utilities.
+
+Provides save/load helpers for caching the result of the expensive
+classify → fuse pipeline in compose_loras_v2.  The cache stores the
+raw fused (A, B, alpha) tensor pairs *before* strength scaling so that
+the strength slider still works correctly at inference time.
+
+Cache layout on disk:
+  ComfyUI/models/SVDQLora/
+    <lora_stem>_precompiled.safetensors   ← fused tensors
+    <lora_stem>_precompiled.safetensors.meta.json ← mtime sidecar
+
+Tensor key format inside the .safetensors file:
+  "<module_key>__A"           → fused A tensor  [rank, in_features]
+  "<module_key>__B"           → fused B tensor  [out_features, rank]  (NOT strength-scaled)
+  "<module_key>__alpha"       → alpha scalar as float32 1-D tensor (optional, omitted if None)
+
+Sidecar JSON format:
+  {
+    "source_mtime": <float>,   # os.path.getmtime() of the original .safetensors
+    "source_path":  <str>      # absolute path of original LoRA (informational)
+  }
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Separator used to encode module_key and tensor role inside a flat safetensors dict.
+# Double-underscore is safe because module keys use dot-notation (e.g. "transformer_blocks.0.attn.to_qkv").
+_SEP = "__"
+
+# Roles stored per module entry
+_ROLE_A = "A"
+_ROLE_B = "B"
+_ROLE_ALPHA = "alpha"
+
+
+# ---------------------------------------------------------------------------
+# Directory helpers
+# ---------------------------------------------------------------------------
+
+def get_cache_dir(comfy_base_dir: str) -> Path:
+    """
+    Return (and create if needed) the dedicated SVDQLora cache directory.
+
+    Args:
+        comfy_base_dir: Root of the ComfyUI installation, e.g. ``"C:/ComfyUI"``.
+
+    Returns:
+        ``Path`` pointing to ``<comfy_base_dir>/models/SVDQLora/``.
+    """
+    cache_dir = Path(comfy_base_dir) / "models" / "SVDQLora"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def get_cache_path(lora_path: str | Path, cache_dir: Path) -> Path:
+    """
+    Derive the cache .safetensors path for a given LoRA source file.
+
+    The cache filename includes a short hash of the LoRA's parent directory so
+    that two LoRAs with identical filenames in different subdirectories (e.g.
+    ``Qwen-Edit/v1/lora.safetensors`` vs ``Qwen-Edit/v2/lora.safetensors``)
+    never collide to the same cache entry.
+
+    Args:
+        lora_path:  Full path to the original LoRA file.
+        cache_dir:  Directory returned by :func:`get_cache_dir`.
+
+    Returns:
+        ``Path`` like ``<cache_dir>/<stem>_<dir_hash8>_precompiled.safetensors``.
+    """
+    lora_path = Path(lora_path)
+    stem = lora_path.stem
+    # Hash the parent directory path so same-named LoRAs in different dirs
+    # produce distinct cache keys.
+    parent_hash = hashlib.sha256(str(lora_path.parent).encode()).hexdigest()[:8]
+    return cache_dir / f"{stem}_{parent_hash}_precompiled.safetensors"
+
+
+def _get_meta_path(cache_path: Path) -> Path:
+    """Return the sidecar JSON path for a given cache file."""
+    return cache_path.with_suffix(cache_path.suffix + ".meta.json")
+
+
+# ---------------------------------------------------------------------------
+# Validation
+# ---------------------------------------------------------------------------
+
+def is_cache_valid(lora_path: str | Path, cache_path: Path) -> bool:
+    """
+    Return ``True`` iff a usable, up-to-date cache exists for *lora_path*.
+
+    Validity criteria:
+    1. ``cache_path`` exists on disk.
+    2. The mtime sidecar JSON exists.
+    3. The recorded ``source_mtime`` matches the current mtime of *lora_path*.
+
+    A missing or unreadable sidecar is treated as invalid (forces re-fuse).
+
+    Args:
+        lora_path:   Full path to the original LoRA file.
+        cache_path:  Path as returned by :func:`get_cache_path`.
+    """
+    if not cache_path.is_file():
+        return False
+
+    meta_path = _get_meta_path(cache_path)
+    if not meta_path.is_file():
+        logger.debug(f"[CACHE] No sidecar for {cache_path.name}, treating as invalid.")
+        return False
+
+    try:
+        with open(meta_path, "r", encoding="utf-8") as fh:
+            meta = json.load(fh)
+        recorded_mtime: float = meta["source_mtime"]
+    except (KeyError, json.JSONDecodeError, OSError) as exc:
+        logger.warning(f"[CACHE] Failed to read sidecar {meta_path.name}: {exc}. Treating as invalid.")
+        return False
+
+    try:
+        current_mtime = os.path.getmtime(str(lora_path))
+    except OSError as exc:
+        logger.warning(f"[CACHE] Cannot stat source LoRA {lora_path}: {exc}. Treating cache as invalid.")
+        return False
+
+    if abs(current_mtime - recorded_mtime) > 1e-3:
+        logger.info(
+            f"[CACHE] Stale: source mtime changed for {Path(lora_path).name}. "
+            f"Recorded={recorded_mtime:.3f}, current={current_mtime:.3f}. "
+            "Will re-fuse and overwrite cache."
+        )
+        return False
+
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Save
+# ---------------------------------------------------------------------------
+
+def save_precompiled(
+    processed_groups: Dict[str, Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]],
+    lora_path: str | Path,
+    cache_path: Path,
+) -> None:
+    """
+    Serialise the fused ``processed_groups`` dict to *cache_path*.
+
+    ``processed_groups`` has the shape produced by the classify+fuse stage of
+    ``compose_loras_v2``:  ``{ module_key: (A, B, alpha) }``.
+
+    All tensors are moved to CPU before saving (the cache is always CPU-side;
+    device placement happens later when ``_apply_lora_to_module`` runs).
+
+    Args:
+        processed_groups: Fused tensor pairs from the classify+fuse stage.
+        lora_path:        Full path to the original LoRA file (used for mtime).
+        cache_path:       Destination .safetensors file.
+    """
+    # Lazy import so safetensors is only required when actually saving.
+    try:
+        from safetensors.torch import save_file as st_save_file
+    except ImportError as exc:
+        logger.error(f"[CACHE] Cannot save precompiled cache: safetensors not available. {exc}")
+        return
+
+    if not processed_groups:
+        logger.warning("[CACHE] processed_groups is empty — nothing to save.")
+        return
+
+    flat: Dict[str, torch.Tensor] = {}
+
+    for module_key, (A, B, alpha) in processed_groups.items():
+        if A is None or B is None:
+            logger.debug(f"[CACHE] Skipping {module_key}: A or B is None.")
+            continue
+
+        # Validate tensor shapes before saving
+        if A.ndim != 2 or B.ndim != 2:
+            logger.warning(
+                f"[CACHE] Skipping {module_key}: expected 2-D tensors, "
+                f"got A.ndim={A.ndim}, B.ndim={B.ndim}."
+            )
+            continue
+
+        # Move to CPU; safetensors does not accept CUDA tensors.
+        flat[f"{module_key}{_SEP}{_ROLE_A}"] = A.detach().cpu().contiguous()
+        flat[f"{module_key}{_SEP}{_ROLE_B}"] = B.detach().cpu().contiguous()
+
+        if alpha is not None:
+            # Normalise to a 1-D float32 scalar tensor for portability.
+            if isinstance(alpha, torch.Tensor):
+                flat[f"{module_key}{_SEP}{_ROLE_ALPHA}"] = alpha.detach().cpu().float().reshape(1)
+            else:
+                flat[f"{module_key}{_SEP}{_ROLE_ALPHA}"] = torch.tensor([float(alpha)], dtype=torch.float32)
+
+    if not flat:
+        logger.warning("[CACHE] No valid tensors to write after filtering — cache not saved.")
+        return
+
+    try:
+        # Write to a temp file first; atomically rename to the final path.
+        # This prevents a partially-written file from being mistaken as valid
+        # if ComfyUI is interrupted mid-write (which would corrupt the cache).
+        tmp_path = cache_path.with_suffix(".safetensors.tmp")
+        st_save_file(flat, str(tmp_path))
+        os.replace(str(tmp_path), str(cache_path))
+        logger.info(f"[CACHE SAVE] Written {len(processed_groups)} module entries → {cache_path}")
+    except Exception as exc:
+        logger.error(f"[CACHE] Failed to write {cache_path}: {exc}")
+        # Clean up temp file if it was created
+        try:
+            tmp_path = cache_path.with_suffix(".safetensors.tmp")
+            tmp_path.unlink(missing_ok=True)
+        except Exception:
+            pass
+        return
+
+    # Write mtime sidecar
+    meta_path = _get_meta_path(cache_path)
+    try:
+        source_mtime = os.path.getmtime(str(lora_path))
+        meta = {
+            "source_mtime": source_mtime,
+            "source_path": str(Path(lora_path).resolve()),
+        }
+        with open(meta_path, "w", encoding="utf-8") as fh:
+            json.dump(meta, fh, indent=2)
+        logger.debug(f"[CACHE] Sidecar written: {meta_path.name}")
+    except OSError as exc:
+        logger.warning(
+            f"[CACHE] Could not write mtime sidecar {meta_path.name}: {exc}. "
+            "Cache will be treated as invalid on next run."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Load
+# ---------------------------------------------------------------------------
+
+def load_precompiled(
+    cache_path: Path,
+) -> Dict[str, Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]]:
+    """
+    Load a precompiled cache file and reconstruct the ``processed_groups`` dict.
+
+    Returns:
+        ``{ module_key: (A, B, alpha_or_None) }`` — identical structure to
+        what ``compose_loras_v2`` produces after the classify+fuse stage.
+        Returns an empty dict on any error so the caller can fall back to
+        a full re-fuse gracefully.
+    """
+    try:
+        from safetensors.torch import load_file as st_load_file
+    except ImportError as exc:
+        logger.error(f"[CACHE] Cannot load precompiled cache: safetensors not available. {exc}")
+        return {}
+
+    try:
+        flat: Dict[str, torch.Tensor] = st_load_file(str(cache_path), device="cpu")
+    except Exception as exc:
+        logger.error(f"[CACHE] Failed to load {cache_path}: {exc}")
+        return {}
+
+    # Reconstruct { module_key: { "A": T, "B": T, "alpha": T|None } }
+    raw: Dict[str, Dict[str, torch.Tensor]] = {}
+    for flat_key, tensor in flat.items():
+        # flat_key format:  "<module_key>__<role>"
+        # module_key itself may contain dots but never "__"
+        sep_idx = flat_key.rfind(_SEP)
+        if sep_idx == -1:
+            logger.warning(f"[CACHE] Unrecognised flat key format: '{flat_key}' — skipping.")
+            continue
+        module_key = flat_key[:sep_idx]
+        role = flat_key[sep_idx + len(_SEP):]
+        raw.setdefault(module_key, {})[role] = tensor
+
+    processed_groups: Dict[str, Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]] = {}
+    for module_key, parts in raw.items():
+        A = parts.get(_ROLE_A)
+        B = parts.get(_ROLE_B)
+        alpha_tensor = parts.get(_ROLE_ALPHA)  # 1-D float32 scalar tensor or None
+
+        if A is None or B is None:
+            logger.warning(f"[CACHE] Incomplete entry for '{module_key}' (missing A or B) — skipping.")
+            continue
+
+        # Convert alpha back to the same form compose_loras_v2 expects:
+        # a torch.Tensor scalar (not a Python float).
+        alpha: Optional[torch.Tensor] = None
+        if alpha_tensor is not None:
+            alpha = alpha_tensor  # keep as tensor; compose_loras_v2 calls .item() on it
+
+        processed_groups[module_key] = (A, B, alpha)
+
+    logger.info(f"[CACHE LOAD] Loaded {len(processed_groups)} module entries ← {cache_path.name}")
+    return processed_groups

nunchaku_code/lora_qwen.py

@@ -10,6 +10,20 @@ import torch
 import torch.nn as nn
 from safetensors import safe_open
 
+
+def _get_compute_device() -> torch.device:
+    """Return the best compute device, deferring to ComfyUI's memory manager.
+
+    Using comfy.model_management.get_torch_device() keeps us fully compatible
+    with ComfyUI's VRAM tracking and offloading context.  Falls back to plain
+    CUDA or CPU if ComfyUI is not available.
+    """
+    try:
+        import comfy.model_management as _mm
+        return _mm.get_torch_device()
+    except Exception:
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # It's assumed these functions from your project are available for import.
 # If not, you'll need to provide their definitions.
 from nunchaku.lora.flux.nunchaku_converter import (
@@ -1159,22 +1173,39 @@ def _apply_lora_to_module(module: nn.Module, A: torch.Tensor, B: torch.Tensor, m
     # Handle Nunchaku LoRA-ready modules
     if hasattr(module, "proj_down") and hasattr(module, "proj_up"):
         pd, pu = module.proj_down.data, module.proj_up.data
-        pd = unpack_lowrank_weight(pd, down=True)
-        pu = unpack_lowrank_weight(pu, down=False)
 
-        base_rank = pd.shape[0] if pd.shape[1] == module.in_features else pd.shape[1]
+        # Use ComfyUI-aware device to stay compatible with its VRAM manager.
+        # This avoids allocating CUDA memory at a time ComfyUI has moved things to CPU.
+        compute_device = _get_compute_device()
+        pd_gpu = pd.to(compute_device)
+        pu_gpu = pu.to(compute_device)
+        A_gpu = A.to(compute_device)
+        B_gpu = B.to(compute_device)
 
-        if pd.shape[1] == module.in_features:  # [rank, in]
+        pd_unpacked = unpack_lowrank_weight(pd_gpu, down=True)
-            new_proj_down = torch.cat([pd, A], dim=0)
+        pu_unpacked = unpack_lowrank_weight(pu_gpu, down=False)
+        del pd_gpu, pu_gpu  # free intermediates promptly
+
+        base_rank = pd_unpacked.shape[0] if pd_unpacked.shape[1] == module.in_features else pd_unpacked.shape[1]
+
+        if pd_unpacked.shape[1] == module.in_features:  # [rank, in]
+            new_proj_down = torch.cat([pd_unpacked, A_gpu], dim=0)
             axis_down = 0
         else:  # [in, rank]
-            new_proj_down = torch.cat([pd, A.T], dim=1)
+            new_proj_down = torch.cat([pd_unpacked, A_gpu.T], dim=1)
             axis_down = 1
+        del pd_unpacked, A_gpu
 
-        new_proj_up = torch.cat([pu, B], dim=1)
+        new_proj_up = torch.cat([pu_unpacked, B_gpu], dim=1)
+        del pu_unpacked, B_gpu
 
-        module.proj_down.data = pack_lowrank_weight(new_proj_down, down=True)
+        packed_down = pack_lowrank_weight(new_proj_down, down=True)
-        module.proj_up.data = pack_lowrank_weight(new_proj_up, down=False)
+        packed_up = pack_lowrank_weight(new_proj_up, down=False)
+        del new_proj_down, new_proj_up
+
+        # Move back to parameter's original device
+        module.proj_down.data = packed_down.to(pd.device)
+        module.proj_up.data = packed_up.to(pu.device)
         module.rank = base_rank + A.shape[0]
 
         if not hasattr(model, "_lora_slots"):
@@ -1238,10 +1269,22 @@ def compose_loras_v2(
         model: torch.nn.Module,
         lora_configs: List[Tuple[Union[str, Path, Dict[str, torch.Tensor]], float]],
         apply_awq_mod: Union[bool, str] = "auto",
+        save_precompiled_lora: bool = False,
+        cache_dir: Optional[str] = None,
 ) -> bool:
     """
     Resets and composes multiple LoRAs into the model with individual strengths.
     """
+    import time
+    try:
+        from tqdm import tqdm as _tqdm
+    except ImportError:
+        def _tqdm(iterable, *args, **kwargs):
+            return iterable
+
+    t_start = time.perf_counter()
+    report_entries = []
+
     logger.info(f"Composing {len(lora_configs)} LoRAs (Direct Fix V6)...")
     reset_lora_v2(model)
     _first_detection = None
@@ -1270,35 +1313,111 @@ def compose_loras_v2(
         logger.info("⚠️  AWQ Modulation Layer LoRA Injection ENABLED (via override).")
 
     # --- 3. Debug Inspection (First LoRA) ---
-    # OPTIMIZATION: Cache first LoRA state dict for reuse in processing loop
+    # Safe unpack: entries may be 2-tuples (path, strength) or 3-tuples
-    _cached_first_lora_state_dict = None
+    # (path, strength, meta_dict) when the cache meta-dict is attached by the node.
     if lora_configs:
-        first_lora_path_or_dict, first_lora_strength = lora_configs[0]
+        _first_entry = lora_configs[0]
-        first_lora_state_dict = _load_lora_state_dict_robust(first_lora_path_or_dict)
+        first_lora_path_or_dict = _first_entry[0]
-        _cached_first_lora_state_dict = first_lora_state_dict  # Cache for reuse
+        first_lora_strength = _first_entry[1]
 
-        # Simple logging of format detection
+        # Optimize: Check if cache is valid for the first LoRA. If so, skip loading the raw safetensors file
-        _first_detection = _detect_lora_format(first_lora_state_dict)
+        _cache_path = None
-        _log_lora_format_detection(str(first_lora_path_or_dict)[:50], _first_detection)
+        _cache_is_valid = False
+        if cache_dir is not None and isinstance(first_lora_path_or_dict, (str, Path)):
+            from nunchaku_code.lora_cache import get_cache_dir as _get_cache_dir, get_cache_path as _get_cache_path, is_cache_valid as _is_cache_valid
+            try:
+                _real_cache_dir = _get_cache_dir(cache_dir)
+                _cache_path = _get_cache_path(first_lora_path_or_dict, _real_cache_dir)
+                _cache_is_valid = _is_cache_valid(first_lora_path_or_dict, _cache_path)
+            except Exception:
+                _cache_is_valid = False
+
+        if _cache_is_valid:
+            _first_detection = {"has_standard": True, "details": "Standard LoRA (Precompiled Cache)"}
+        else:
+            first_lora_state_dict = _load_lora_state_dict_robust(first_lora_path_or_dict)
+            # Simple logging of format detection
+            _first_detection = _detect_lora_format(first_lora_state_dict)
+            _log_lora_format_detection(str(first_lora_path_or_dict)[:50], _first_detection)
 
     aggregated_weights: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
     # Track names of LoRA files that contain skipped weights
     skipped_lora_names = set()
 
     # --- 4. Main Loading Loop ---
-    for idx, (lora_path_or_dict, strength) in enumerate(lora_configs):
+    # Resolve cache directory once (None disables the cache path entirely).
+    _cache_dir_path = None
+    if cache_dir is not None:
+        from pathlib import Path as _Path
+        _cache_dir_path = _Path(cache_dir)
+
+    for idx, _entry in enumerate(lora_configs):
+        # Entries may be plain 2-tuples (path, strength) or 3-tuples
+        # (path, strength, meta_dict) carrying per-LoRA flags from the node UI.
+        lora_path_or_dict = _entry[0]
+        strength = _entry[1]
+        _entry_meta: dict = _entry[2] if len(_entry) > 2 else {}
+
         if abs(strength) < 1e-5:
             continue
-        
+
+        t_lora_start = time.perf_counter()
+
         lora_name = str(lora_path_or_dict)
         if isinstance(lora_path_or_dict, dict):
             lora_name = f"lora_{idx}"
 
-        # OPTIMIZATION: Reuse cached first LoRA state dict to avoid duplicate file I/O
+        # Per-entry save flag: honour node-level flag OR the global function arg.
-        if idx == 0 and _cached_first_lora_state_dict is not None:
+        _should_save = save_precompiled_lora or bool(_entry_meta.get("save_precompiled", False))
-            lora_state_dict = _cached_first_lora_state_dict
+
-        else:
+        # ------------------------------------------------------------------ #
-            lora_state_dict = _load_lora_state_dict_robust(lora_path_or_dict)
+        # CACHE FAST PATH                                                     #
+        # Only available for file-backed LoRAs (not inline dicts) and only   #
+        # when a cache directory is configured.                               #
+        # ------------------------------------------------------------------ #
+        _cache_path = None
+        if _cache_dir_path is not None and isinstance(lora_path_or_dict, (str, Path)):
+            from nunchaku_code.lora_cache import (
+                get_cache_dir as _get_cache_dir,
+                get_cache_path as _get_cache_path,
+                is_cache_valid as _is_cache_valid,
+                load_precompiled as _load_precompiled,
+                save_precompiled as _save_precompiled,
+            )
+            _real_cache_dir = _get_cache_dir(_cache_dir_path)
+            _cache_path = _get_cache_path(lora_path_or_dict, _real_cache_dir)
+
+            if _is_cache_valid(lora_path_or_dict, _cache_path):
+                # ---- CACHE HIT: skip disk I/O + classify + fuse ---- #
+                processed_groups = _load_precompiled(_cache_path)
+                if processed_groups:
+                    logger.info(
+                        f"[CACHE HIT] Skipping classify+fuse for '{Path(lora_path_or_dict).name}'. "
+                        f"Loaded {len(processed_groups)} pre-fused entries from cache."
+                    )
+                    for module_key, (A, B, alpha) in processed_groups.items():
+                        aggregated_weights[module_key].append(
+                            {"A": A, "B": B, "alpha": alpha, "strength": strength, "source": lora_name}
+                        )
+                    t_lora_end = time.perf_counter()
+                    report_entries.append({
+                        "name": Path(lora_path_or_dict).name if isinstance(lora_path_or_dict, (str, Path)) else f"lora_{idx}",
+                        "strength": strength,
+                        "format": "Standard LoRA (Precompiled)",
+                        "load_type": "Cache Hit",
+                        "elapsed": t_lora_end - t_lora_start,
+                    })
+                    continue  # skip normal fuse path for this LoRA
+                else:
+                    logger.warning(
+                        f"[CACHE] Cache file loaded but returned empty groups for "
+                        f"'{Path(lora_path_or_dict).name}'. Falling back to full fuse."
+                    )
+
+        # ------------------------------------------------------------------ #
+        # NORMAL (FULL) FUSE PATH                                            #
+        # ------------------------------------------------------------------ #
+        lora_state_dict = _load_lora_state_dict_robust(lora_path_or_dict)
         if not lora_state_dict:
             continue
 
@@ -1323,17 +1442,19 @@ def compose_loras_v2(
         # Process groups
         processed_groups = {}
         special_handled = set()
-        
+
-        for base_key, lw in lora_grouped.items():
+        items_to_process = list(lora_grouped.items())
+        pbar_desc = f"Fusing weights for {Path(lora_path_or_dict).name[:25]}"
+        for base_key, lw in _tqdm(items_to_process, desc=pbar_desc, unit="group", leave=True):
             if base_key in special_handled:
                 continue
 
             # Standard Logic Copied from Original
             if "lokr_w1" in lw or "lokr_w2" in lw:
-                continue # LoKR skipped
+                continue  # LoKR skipped
 
             A, B, alpha = None, None, None
-            
+
             if "qkv" in base_key:
                 A, B, alpha = (lw.get("A"), lw.get("B"), lw.get("alpha")) if "A" in lw else _fuse_qkv_lora(lw, model=model, base_key=base_key)
             elif "w1_A" in lw or "w3_A" in lw:
@@ -1349,19 +1470,44 @@ def compose_loras_v2(
             if A is not None and B is not None:
                 processed_groups[base_key] = (A, B, alpha)
 
+        # ------------------------------------------------------------------ #
+        # OPTIONAL CACHE SAVE                                                #
+        # Persist the freshly fused processed_groups so the next run can     #
+        # skip straight to aggregation.  B is stored un-scaled so the       #
+        # strength slider keeps working at inference time.                   #
+        # ------------------------------------------------------------------ #
+        load_type = "Cache Miss (Compiled)"
+        if _should_save and _cache_path is not None and processed_groups:
+            try:
+                _save_precompiled(processed_groups, lora_path_or_dict, _cache_path)
+                load_type = "Cache Miss (Compiled & Saved)"
+            except Exception as _cache_save_exc:
+                logger.warning(f"[CACHE] Failed to save precompiled cache: {_cache_save_exc}")
+
         for module_key, (A, B, alpha) in processed_groups.items():
             aggregated_weights[module_key].append(
                 {"A": A, "B": B, "alpha": alpha, "strength": strength, "source": lora_name}
             )
 
+        t_lora_end = time.perf_counter()
+        report_entries.append({
+            "name": Path(lora_path_or_dict).name if isinstance(lora_path_or_dict, (str, Path)) else f"lora_{idx}",
+            "strength": strength,
+            "format": detection["has_standard"] and "Standard LoRA (Rank-Decomposed)" or "Unknown/Other",
+            "load_type": load_type,
+            "elapsed": t_lora_end - t_lora_start,
+        })
+
     # --- 5. Application Loop (with AWQ Fix) ---
+    t_apply_start = time.perf_counter()
     applied_modules_count = 0
     skipped_modules_count = 0
     mod_layer_applied_count = 0
     
     all_A = [] # Initialization for safety
 
-    for module_name_key, weight_list in aggregated_weights.items():
+    apply_pbar_desc = "Applying weights to modules"
+    for module_name_key, weight_list in _tqdm(list(aggregated_weights.items()), desc=apply_pbar_desc, unit="mod", leave=True):
         # Correctly resolving module (handling tuple return if needed from _resolve_module_name, but it returns tuple (name, mod) or just (name, mod) in signature, wait.)
         # _resolve_module_name defined in this file returns Tuple[str, Optional[nn.Module]]
         # Wait, if module is None, it returns None? line 416 returns name, m. line 409 returns None.
@@ -1428,16 +1574,22 @@ def compose_loras_v2(
                 # Standard Linear / Nunchaku
                 if hasattr(module, "proj_down"):
                      target_dtype = module.proj_down.dtype
-                     target_device = module.proj_down.device
+                     # Nunchaku: skip the CPU→CPU device pre-transfer.
+                     # _apply_lora_to_module will move these to the compute device (CUDA)
+                     # itself. Pre-transferring to CPU here is a redundant memcpy (~960
+                     # copies per generation) that wastes CPU memory bandwidth.
+                     all_A.append(A.to(dtype=target_dtype))
+                     all_B_scaled.append((B * scale).to(dtype=target_dtype))
                 elif hasattr(module, "weight"):
                      target_dtype = module.weight.dtype
                      target_device = module.weight.device
+                     all_A.append(A.to(dtype=target_dtype, device=target_device))
+                     all_B_scaled.append((B * scale).to(dtype=target_dtype, device=target_device))
                 else:
                      target_dtype = torch.float16
                      target_device = "cuda"
-
+                     all_A.append(A.to(dtype=target_dtype, device=target_device))
-                all_A.append(A.to(dtype=target_dtype, device=target_device))
+                     all_B_scaled.append((B * scale).to(dtype=target_dtype, device=target_device))
-                all_B_scaled.append((B * scale).to(dtype=target_dtype, device=target_device))
 
 
         if not all_A:
@@ -1490,6 +1642,10 @@ def compose_loras_v2(
                 logger.info(f"[APPLY] LoRA applied to: {resolved_name}")
             applied_modules_count += 1
 
+    t_apply_end = time.perf_counter()
+    t_apply_elapsed = t_apply_end - t_apply_start
+    t_total_elapsed = t_apply_end - t_start
+
     if skipped_lora_names:
         logger.warning("""\n⚠️  Some weights from the following LoRAs were automatically skipped 
         because modulation layers (img_mod/txt_mod) are extremely sensitive to Nunchaku quantization. 
@@ -1498,7 +1654,42 @@ def compose_loras_v2(
             short_name = os.path.basename(source) if os.path.sep in source else source
             logger.warning(f"   - {short_name}")
 
-    logger.info(f"Sampled LoRA composition complete. Applied: {applied_modules_count}, Mod-patched: {mod_layer_applied_count}, Skipped: {skipped_modules_count}")
+    # Display Report Card
+    card_width = 70
+    border = "┌" + "─" * (card_width - 2) + "┐"
+    divider = "├" + "─" * (card_width - 2) + "┤"
+    bottom = "└" + "─" * (card_width - 2) + "┘"
+
+    def _log_line(left_text: str, right_text: str = ""):
+        padding_len = card_width - 4 - len(left_text) - len(right_text)
+        if padding_len < 0:
+            left_text = left_text[:card_width - 8 - len(right_text)] + "..."
+            padding_len = card_width - 4 - len(left_text) - len(right_text)
+        logger.info(f"│ {left_text}{' ' * padding_len}{right_text} │")
+
+    logger.info(border)
+    title = "NUNCHAKU QWEN LORA COMPOSITION REPORT"
+    title_padding = (card_width - 2 - len(title)) // 2
+    logger.info(f"│{' ' * title_padding}{title}{' ' * (card_width - 2 - len(title) - title_padding)}│")
+    logger.info(divider)
+    
+    for i, entry in enumerate(report_entries, 1):
+        _log_line(f"{i}. {entry['name']}")
+        _log_line(f"   - Strength: {entry['strength']}")
+        _log_line(f"   - Format:   {entry['format']}")
+        _log_line(f"   - Status:   {entry['load_type']}")
+        _log_line(f"   - Time:     {entry['elapsed']:.3f}s")
+        if i < len(report_entries):
+            logger.info("│" + " " * (card_width - 2) + "│")
+            
+    logger.info(divider)
+    _log_line(f"Total Processed:      {len(report_entries)} LoRA(s)")
+    _log_line(f"Layers Applied:       {applied_modules_count} (Attention/MLP)")
+    _log_line(f"Modulation Patched:   {mod_layer_applied_count} (AWQ manual planar)")
+    _log_line(f"Weight Apply Time:    {t_apply_elapsed:.3f}s")
+    _log_line(f"Total Composition:    {t_total_elapsed:.3f}s")
+    logger.info(bottom)
+
     return True
 
 def update_lora_params_v2(
@@ -1515,6 +1706,8 @@ def compose_loras_v2_v2(
         model: torch.nn.Module,
         lora_configs: List[Tuple[Union[str, Path, Dict[str, torch.Tensor]], float]],
         apply_awq_mod: bool | str | None = None,
+        save_precompiled_lora: bool = False,
+        cache_dir: Optional[str] = None,
 ) -> bool:
     """
     V2 node-specific LoRA composition function with AWQ modulation control.
@@ -1547,7 +1740,13 @@ def compose_loras_v2_v2(
     
     # Call compose_loras_v2 with the determined flag
     # Pass as bool explicitly to override default "auto" behavior
-    return compose_loras_v2(model, lora_configs, apply_awq_mod=apply_awq_mod_flag)
+    return compose_loras_v2(
+        model,
+        lora_configs,
+        apply_awq_mod=apply_awq_mod_flag,
+        save_precompiled_lora=save_precompiled_lora,
+        cache_dir=cache_dir,
+    )
 
 
 def set_lora_strength_v2(model: nn.Module, strength: float) -> None:
@@ -1561,7 +1760,9 @@ def set_lora_strength_v2(model: nn.Module, strength: float) -> None:
 
     for name, info in model._lora_slots.items():
         module = _get_module_by_name(model, name)
-        if module is None or info.get("appended", 0) <= 0:
+        # Skip if module not found, no appended weights, or no base_rank recorded
+        # (awq_w4a16 modulation slots don't store base_rank — they're handled by forward patching)
+        if module is None or info.get("appended", 0) <= 0 or "base_rank" not in info:
             continue
 
         base_rank, appended = info["base_rank"], info["appended"]
@@ -1576,11 +1777,6 @@ def reset_lora_v2(model: nn.Module) -> None:
     if not hasattr(model, "_lora_slots") or not model._lora_slots:
         return
 
-    for name, info in model._lora_slots.items():
-        module = _get_module_by_name(model, name)
-        if module is None:
-            continue
-
     for name, info in model._lora_slots.items():
         module = _get_module_by_name(model, name)
         if module is None:
@@ -1591,17 +1787,28 @@ def reset_lora_v2(model: nn.Module) -> None:
         if module_type == "nunchaku":
              base_rank = info["base_rank"]
              with torch.no_grad():
-                 pd = unpack_lowrank_weight(module.proj_down.data, down=True)
+                 # Use ComfyUI-aware device to stay compatible with its VRAM manager
-                 pu = unpack_lowrank_weight(module.proj_up.data, down=False)
+                 compute_device = _get_compute_device()
- 
+                 orig_pd_device = module.proj_down.data.device
+                 orig_pu_device = module.proj_up.data.device
+
+                 pd_gpu = module.proj_down.data.to(compute_device)
+                 pu_gpu = module.proj_up.data.to(compute_device)
+
+                 pd = unpack_lowrank_weight(pd_gpu, down=True)
+                 pu = unpack_lowrank_weight(pu_gpu, down=False)
+                 del pd_gpu, pu_gpu  # free intermediates promptly
+
                  if info.get("axis_down", 0) == 0:  # [rank, in]
                      pd_reset = pd[:base_rank, :].clone()
                  else:  # [in, rank]
                      pd_reset = pd[:, :base_rank].clone()
                  pu_reset = pu[:, :base_rank].clone()
- 
+                 del pd, pu
-                 module.proj_down.data = pack_lowrank_weight(pd_reset, down=True)
+
-                 module.proj_up.data = pack_lowrank_weight(pu_reset, down=False)
+                 module.proj_down.data = pack_lowrank_weight(pd_reset, down=True).to(orig_pd_device)
+                 module.proj_up.data = pack_lowrank_weight(pu_reset, down=False).to(orig_pu_device)
+                 del pd_reset, pu_reset
                  module.rank = base_rank
 
         elif module_type == "linear":
@@ -1629,4 +1836,9 @@ def reset_lora_v2(model: nn.Module) -> None:
             
     model._lora_slots.clear()
     model._lora_strength = 1.0
+    if hasattr(model, "_applied_loras"):
+        try:
+            delattr(model, "_applied_loras")
+        except Exception:
+            pass
     logger.info("All LoRA weights have been reset from the model.")

patches/nunchaku_patch.py

@@ -1,133 +1,14 @@
 
 import logging
-import os
-import re
 import sys
 import types
-from typing import Callable, Optional, Tuple
+from typing import Optional, Tuple
 
 import torch
 
 logger = logging.getLogger(__name__)
 
 _svdq_from_linear_patched: bool = False
-_qwen_apply_rotary_emb_compat_applied: bool = False
-
-_ROTARY_SHIM_TAG = "_qwen_lora_loader_rotary_shim"
-_NUNCHAKU_QWENIMAGE_APPLY_ROTARY_PATTERNS = (
-    re.compile(
-        r"from\s+comfy\.ldm\.qwen_image\.model\s+import\s+\([^)]*\bapply_rotary_emb\b",
-        re.MULTILINE | re.DOTALL,
-    ),
-    re.compile(
-        r"from\s+comfy\.ldm\.qwen_image\.model\s+import\s+[^\n#]*\bapply_rotary_emb\b",
-    ),
-)
-
-
-def _rotary_compat_disabled_by_env() -> bool:
-    value = os.environ.get("QWENIMAGE_ROTARY_COMPAT", "").strip().lower()
-    return value in ("0", "false", "no", "off", "disable", "disabled")
-
-
-def _mark_rotary_shim(fn: Callable) -> Callable:
-    setattr(fn, _ROTARY_SHIM_TAG, True)
-    return fn
-
-
-def _rotary_shim_installed_on(qwen_image_model) -> bool:
-    fn = getattr(qwen_image_model, "apply_rotary_emb", None)
-    return fn is not None and getattr(fn, _ROTARY_SHIM_TAG, False)
-
-
-def _comfyui_has_native_apply_rotary_emb(qwen_image_model) -> bool:
-    return (
-        getattr(qwen_image_model, "apply_rotary_emb", None) is not None
-        and not _rotary_shim_installed_on(qwen_image_model)
-    )
-
-
-def _find_nunchaku_qwenimage_py() -> Optional[str]:
-    custom_nodes = os.path.abspath(
-        os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)
-    )
-    for folder in ("ComfyUI-nunchaku", "ComfyUI_Nunchaku", "comfyui-nunchaku"):
-        path = os.path.join(custom_nodes, folder, "models", "qwenimage.py")
-        if os.path.isfile(path):
-            return path
-    return None
-
-
-def _nunchaku_qwenimage_still_imports_apply_rotary_emb() -> Optional[bool]:
-    """
-    True: nunchaku still imports apply_rotary_emb (compat may be needed).
-    False: nunchaku source no longer imports it (skip shim).
-    None: qwenimage.py not found (apply only if symbol missing).
-    """
-    path = _find_nunchaku_qwenimage_py()
-    if path is None:
-        return None
-
-    try:
-        with open(path, encoding="utf-8", errors="replace") as handle:
-            text = handle.read()
-    except OSError as exc:
-        logger.debug("Could not read %s for rotary compat probe: %s", path, exc)
-        return None
-
-    for pattern in _NUNCHAKU_QWENIMAGE_APPLY_ROTARY_PATTERNS:
-        if pattern.search(text):
-            return True
-    return False
-
-
-def apply_qwen_image_apply_rotary_emb_compat() -> bool:
-    """
-    ComfyUI >= 0.24 removed comfy.ldm.qwen_image.model.apply_rotary_emb.
-    ComfyUI-nunchaku still imports it; alias to comfy.ldm.flux.math.apply_rope1
-    only when the import is still required and ComfyUI has not restored the symbol.
-    """
-    global _qwen_apply_rotary_emb_compat_applied
-    if _qwen_apply_rotary_emb_compat_applied:
-        return True
-
-    if _rotary_compat_disabled_by_env():
-        logger.info(
-            "apply_rotary_emb compat skipped (QWENIMAGE_ROTARY_COMPAT is disabled)"
-        )
-        return False
-
-    try:
-        import comfy.ldm.qwen_image.model as qwen_image_model
-        from comfy.ldm.flux.math import apply_rope1
-
-        if _rotary_shim_installed_on(qwen_image_model):
-            _qwen_apply_rotary_emb_compat_applied = True
-            return True
-
-        if _comfyui_has_native_apply_rotary_emb(qwen_image_model):
-            logger.info(
-                "apply_rotary_emb compat skipped: ComfyUI already exports apply_rotary_emb"
-            )
-            return False
-
-        nunchaku_needs = _nunchaku_qwenimage_still_imports_apply_rotary_emb()
-        if nunchaku_needs is False:
-            logger.info(
-                "apply_rotary_emb compat skipped: ComfyUI-nunchaku no longer imports apply_rotary_emb"
-            )
-            return False
-
-        qwen_image_model.apply_rotary_emb = _mark_rotary_shim(apply_rope1)
-        _qwen_apply_rotary_emb_compat_applied = True
-        logger.info(
-            "Patched comfy.ldm.qwen_image.model.apply_rotary_emb -> apply_rope1 "
-            "(ComfyUI-nunchaku Qwen Image compat)"
-        )
-        return True
-    except Exception as e:
-        logger.error("Failed to apply apply_rotary_emb compat patch: %s", e)
-        return False
 
 
 def _torch_device_fallback() -> torch.device:
@@ -358,7 +239,6 @@ def apply_nunchaku_patch():
     Apply ComfyUI-nunchaku compatibility patches (LoRA planar injection + lazy Linear fixes).
     Returns True if at least one patch was applied or was already active.
     """
-    rotary_compat = apply_qwen_image_apply_rotary_emb_compat()
     lazy_from = apply_svdqw4a4_lazy_linear_patch()
     lazy_fuse = apply_nunchaku_zimage_fuse_lazy_linear_patch()
     if not lazy_fuse:
@@ -395,4 +275,4 @@ def apply_nunchaku_patch():
     except Exception as e:
         logger.error("Failed to apply Nunchaku planar patch: %s", e)
 
-    return planar_ok or lazy_from or rotary_compat
+    return planar_ok or lazy_from

wrappers/qwenimage.py

@@ -1,9 +1,11 @@
 from typing import Callable, List, Tuple, Union
 from pathlib import Path
+import os
 
 import torch
 from torch import nn
 import comfy.model_management
+import folder_paths
 import logging
 
 from nunchaku import NunchakuQwenImageTransformer2DModel
@@ -66,6 +68,13 @@ class ComfyQwenImageWrapper(nn.Module):
         # Track last seen device to detect CPU/GPU moves that require re-compose
         self._last_device = None
 
+        # Track whether this wrapper has ever held any LoRAs. A wrapper that was
+        # stamped onto the INPUT model by load_lora's side-effect mutation but never
+        # received any LoRAs (i.e. the bypass-path wrapper W1) must NOT reset the
+        # shared transformer, because the ret_model wrapper (W2) may have already
+        # composed LoRA into it.
+        self._has_ever_had_loras = False
+
     def process_img(self, x, index=0, h_offset=0, w_offset=0):
         """
         Delegate to the inner model so Qwen Image ControlNet (e.g. Fun ControlNet)
@@ -156,18 +165,30 @@ class ComfyQwenImageWrapper(nn.Module):
         if self.model is None:
             raise RuntimeError("Model has been unloaded or garbage collected. Cannot perform forward pass.")
 
+        # Update LoRA-manager flag: once this wrapper holds any LoRAs it is
+        # considered a LoRA manager for the lifetime of the object.
+        if self.loras:
+            self._has_ever_had_loras = True
+
+        # model_is_dirty: the transformer has LoRA applied but this wrapper expects
+        # none. Only relevant if we were once a LoRA manager (i.e. we previously
+        # composed into the transformer). A bypass-path wrapper that has never
+        # managed LoRAs must NOT declare the transformer dirty — another wrapper
+        # (the ret_model wrapper) may have legitimately composed into it.
         model_is_dirty = (
-            not self.loras and # We expect no LoRA
+            self._has_ever_had_loras and  # We were (or are) a LoRA manager
-            hasattr(self.model, "_lora_slots") and self.model._lora_slots # But the model actually has LoRA
+            not self.loras and  # We currently expect no LoRA
+            hasattr(self.model, "_lora_slots") and self.model._lora_slots  # But the model has LoRA
         )
         
         # Deep comparison of LoRA stacks to detect any changes
-        # This ensures we catch changes in weights, paths, or order
+        # Check against the shared model's currently composed LoRAs to handle multiple wrappers sharing the model.
+        model_applied = getattr(self.model, "_applied_loras", None)
         loras_changed = False
-        if self._applied_loras is None or len(self._applied_loras) != len(self.loras):
+        if model_applied is None or len(model_applied) != len(self.loras):
             loras_changed = True
         else:
-            for applied, current in zip(self._applied_loras, self.loras):
+            for applied, current in zip(model_applied, self.loras):
                 if applied != current:
                     loras_changed = True
                     break
@@ -181,10 +202,18 @@ class ComfyQwenImageWrapper(nn.Module):
         
         # Check if the LoRA stack has been changed by a loader node
         if loras_changed or model_is_dirty or device_changed:
-            # The compose function handles resetting before applying the new stack
+            # Guard: only reset/recompose the shared transformer when this wrapper
-            reset_lora_v2(self.model)
+            # is (or has been) an active LoRA manager.  The bypass-path wrapper W1
-            self._applied_loras = self.loras.copy()
+            # (created as a side-effect of load_lora mutating the input model) has
-            
+            # loras=[] and _has_ever_had_loras=False on its very first forward pass.
+            # Without this guard it would call reset_lora_v2(shared_transformer),
+            # wiping whatever LoRA the ret_model wrapper W2 had composed.
+            is_lora_manager = bool(self.loras) or self._has_ever_had_loras
+
+            if is_lora_manager:
+                # The compose function handles resetting before applying the new stack
+                reset_lora_v2(self.model)
+
             # Reset cache when LoRAs change to prevent stale cache in multi-stage workflows
             # This ensures that when switching between different LoRA sets in different stages,
             # the cache is invalidated and recreated with the new LoRA composition
@@ -193,91 +222,124 @@ class ComfyQwenImageWrapper(nn.Module):
                 self._prev_timestep = None
                 logger.debug("Cache reset due to LoRA change")
 
-            # --- NEW DYNAMIC VRAM CHECK (conditionally applied) ---
+            if is_lora_manager:
+                # --- NEW DYNAMIC VRAM CHECK (conditionally applied) ---
 
-            # 1. Check if offload is *already* enabled (from loader setting "enable" or "auto" on low-vram)
+                # 1. Check if offload is *already* enabled (from loader setting "enable" or "auto" on low-vram)
-            offload_is_on = hasattr(self.model, "offload_manager") and self.model.offload_manager is not None
+                offload_is_on = hasattr(self.model, "offload_manager") and self.model.offload_manager is not None
 
-            # 2. Decide if we *need* to turn it on
+                # 2. Decide if we *need* to turn it on
-            should_enable_offload = offload_is_on
+                should_enable_offload = offload_is_on
 
-            # 3. Only run the dynamic VRAM check if:
+                # 3. Only run the dynamic VRAM check if:
-            #    - The user's original setting was "auto"
+                #    - The user's original setting was "auto"
-            #    - Offloading is not *already* on
+                #    - Offloading is not *already* on
-            #    - We are actually loading new LoRAs
+                #    - We are actually loading new LoRAs
-            if self.cpu_offload_setting == "auto" and not offload_is_on and self.loras:
+                if self.cpu_offload_setting == "auto" and not offload_is_on and self.loras:
-                try:
-                    # Use the VRAM margin from the loader node
-                    free_vram_gb = comfy.model_management.get_free_memory() / (1024 ** 3)
-
-                    if free_vram_gb < self.vram_margin_gb:
-                        logger.info(
-                            f"Free VRAM is {free_vram_gb:.2f}GB (below safety margin of {self.vram_margin_gb}GB) and 'cpu_offload' is 'auto'. Force-enabling CPU offload for LoRA composition.")
-                        should_enable_offload = True
-                    else:
-                        logger.info(
-                            f"Free VRAM is {free_vram_gb:.2f}GB (>= {self.vram_margin_gb}GB margin). LoRAs will be composed without enabling CPU offload.")
-
-                except Exception as e:
-                    logger.error(f"Error during VRAM check for LoRA offloading: {e}. Offload will not be enabled.")
-            elif self.cpu_offload_setting == "disable" and not offload_is_on:
-                logger.debug("CPU offload is 'disable' and not on. Skipping VRAM check.")
-            elif self.cpu_offload_setting == "enable" and offload_is_on:
-                logger.debug("CPU offload is 'enable'. Will rebuild offload manager for LoRAs.")
-
-            # --- END NEW VRAM CHECK ---
-
-            # 4. Compose LoRAs. This changes internal tensor shapes.
-            # 4. Compose LoRAs. This changes internal tensor shapes.
-            # Returns True if successful (supported format), False if unsupported (skipped).
-            is_supported_format = compose_loras_v2(self.model, self.loras, apply_awq_mod=self.apply_awq_mod)
-
-            # Validate composition result; if 0 targets after a crash/transition, retry once
-            # But ONLY if the format was supported. If unsupported, retrying is pointless.
-            if is_supported_format:
-                try:
-                    has_slots = hasattr(self.model, "_lora_slots") and bool(self.model._lora_slots)
-                except Exception:
-                    has_slots = True
-                if self.loras and not has_slots:
-                    logger.warning("LoRA composition reported 0 target modules. Forcing reset and one retry.")
                     try:
-                        reset_lora_v2(self.model)
+                        # Use the VRAM margin from the loader node
-                        compose_loras_v2(self.model, self.loras, apply_awq_mod=self.apply_awq_mod)
+                        free_vram_gb = comfy.model_management.get_free_memory() / (1024 ** 3)
+
+                        if free_vram_gb < self.vram_margin_gb:
+                            logger.info(
+                                f"Free VRAM is {free_vram_gb:.2f}GB (below safety margin of {self.vram_margin_gb}GB) and 'cpu_offload' is 'auto'. Force-enabling CPU offload for LoRA composition.")
+                            should_enable_offload = True
+                        else:
+                            logger.info(
+                                f"Free VRAM is {free_vram_gb:.2f}GB (>= {self.vram_margin_gb}GB margin). LoRAs will be composed without enabling CPU offload.")
+
                     except Exception as e:
-                        logger.error(f"LoRA re-compose retry failed: {e}")
+                        logger.error(f"Error during VRAM check for LoRA offloading: {e}. Offload will not be enabled.")
-            else:
+                elif self.cpu_offload_setting == "disable" and not offload_is_on:
-                 logger.warning("Skipping retry because LoRA format is unsupported.")
+                    logger.debug("CPU offload is 'disable' and not on. Skipping VRAM check.")
+                elif self.cpu_offload_setting == "enable" and offload_is_on:
+                    logger.debug("CPU offload is 'enable'. Will rebuild offload manager for LoRAs.")
 
-            # 5. Re-build offload manager if it's supposed to be on
+                # --- END NEW VRAM CHECK ---
-            # This block now runs if offload was on *or* if our new check decided to turn it on.
-            if should_enable_offload:
 
-                # Store settings if it was already on, otherwise use defaults
+                # 4. Build cache args.  The ComfyUI base dir is the parent of the
-                if offload_is_on:
+                #    'models' folder, which folder_paths knows about.
-                    manager = self.model.offload_manager
+                try:
-                    offload_settings = {
+                    _comfy_base = str(Path(folder_paths.models_dir).parent)
-                        "num_blocks_on_gpu": manager.num_blocks_on_gpu,
+                except Exception:
-                        "use_pin_memory": manager.use_pin_memory,
+                    _comfy_base = None
-                    }
+
+                # Separate per-LoRA meta dicts from the (path, strength[, meta]) tuples
+                # and derive the global save_precompiled flag (True if ANY entry requests it).
+                _lora_configs_clean = []
+                _any_save_flag = False
+                for _entry in self.loras:
+                    _lc_path = _entry[0]
+                    _lc_strength = _entry[1]
+                    _lc_meta: dict = _entry[2] if len(_entry) > 2 else {}
+                    _lora_configs_clean.append((_lc_path, _lc_strength, _lc_meta))
+                    if _lc_meta.get("save_precompiled", False):
+                        _any_save_flag = True
+
+                # 5. Compose LoRAs. This changes internal tensor shapes.
+                # Returns True if successful (supported format), False if unsupported (skipped).
+                is_supported_format = compose_loras_v2(
+                    self.model,
+                    _lora_configs_clean,
+                    apply_awq_mod=self.apply_awq_mod,
+                    save_precompiled_lora=_any_save_flag,
+                    cache_dir=_comfy_base,
+                )
+
+                # Validate composition result; if 0 targets after a crash/transition, retry once
+                # But ONLY if the format was supported. If unsupported, retrying is pointless.
+                if is_supported_format:
+                    try:
+                        has_slots = hasattr(self.model, "_lora_slots") and bool(self.model._lora_slots)
+                    except Exception:
+                        has_slots = True
+                    if self.loras and not has_slots:
+                        logger.warning("LoRA composition reported 0 target modules. Forcing reset and one retry.")
+                        try:
+                            reset_lora_v2(self.model)
+                            compose_loras_v2(
+                                self.model,
+                                _lora_configs_clean,
+                                apply_awq_mod=self.apply_awq_mod,
+                                save_precompiled_lora=_any_save_flag,
+                                cache_dir=_comfy_base,
+                            )
+                        except Exception as e:
+                            logger.error(f"LoRA re-compose retry failed: {e}")
                 else:
-                    # Not previously on, so use defaults from nodes/models/qwenimage.py
+                    logger.warning("Skipping retry because LoRA format is unsupported.")
-                    offload_settings = {
-                        "num_blocks_on_gpu": 1,
-                        "use_pin_memory": False,  # 'disable' maps to False
-                    }
-                    logger.info("Building new CPU offload manager due to LoRA VRAM check.")
 
-                # Step 1: Completely disable and clear any old offloader (safe to call even if off)
+                # 5. Re-build offload manager if it's supposed to be on
-                self.model.set_offload(False)
+                # This block now runs if offload was on *or* if our new check decided to turn it on.
+                if should_enable_offload:
 
-                # Step 2: Re-enable offloading with the correct settings
+                    # Store settings if it was already on, otherwise use defaults
-                # This builds the manager based on the *newly composed* tensor shapes.
+                    if offload_is_on:
-                self.model.set_offload(True, **offload_settings)
+                        manager = self.model.offload_manager
+                        offload_settings = {
+                            "num_blocks_on_gpu": manager.num_blocks_on_gpu,
+                            "use_pin_memory": manager.use_pin_memory,
+                        }
+                    else:
+                        # Not previously on, so use defaults from nodes/models/qwenimage.py
+                        offload_settings = {
+                            "num_blocks_on_gpu": 1,
+                            "use_pin_memory": False,  # 'disable' maps to False
+                        }
+                        logger.info("Building new CPU offload manager due to LoRA VRAM check.")
+
+                    # Step 1: Completely disable and clear any old offloader (safe to call even if off)
+                    self.model.set_offload(False)
+
+                    # Step 2: Re-enable offloading with the correct settings
+                    # This builds the manager based on the *newly composed* tensor shapes.
+                    self.model.set_offload(True, **offload_settings)
 
             # --- END MODIFIED SECTION ---
 
-            # Update last known device signature after (re)composition
+            # Always update state tracking so subsequent calls see a consistent snapshot.
+            self.model._applied_loras = self.loras.copy()
+            self._applied_loras = self.loras.copy()
             self._last_device = current_device
 
         # Caching logic

wrappers/qwenimage_v2.py

@@ -1,9 +1,11 @@
 from typing import Callable, List, Tuple, Union
 from pathlib import Path
+import os
 
 import torch
 from torch import nn
 import comfy.model_management
+import folder_paths
 import logging
 
 from nunchaku import NunchakuQwenImageTransformer2DModel
@@ -136,12 +138,13 @@ class ComfyQwenImageWrapperV2(nn.Module):
         )
         
         # Deep comparison of LoRA stacks to detect any changes
-        # This ensures we catch changes in weights, paths, or order
+        # Check against the shared model's currently composed LoRAs to handle multiple wrappers sharing the model.
+        model_applied = getattr(self.model, "_applied_loras", None)
         loras_changed = False
-        if self._applied_loras is None or len(self._applied_loras) != len(self.loras):
+        if model_applied is None or len(model_applied) != len(self.loras):
             loras_changed = True
         else:
-            for applied, current in zip(self._applied_loras, self.loras):
+            for applied, current in zip(model_applied, self.loras):
                 if applied != current:
                     loras_changed = True
                     break
@@ -157,6 +160,7 @@ class ComfyQwenImageWrapperV2(nn.Module):
         if loras_changed or model_is_dirty or device_changed:
             # The compose function handles resetting before applying the new stack
             reset_lora_v2(self.model)
+            self.model._applied_loras = self.loras.copy()
             self._applied_loras = self.loras.copy()
             
             # Reset cache when LoRAs change to prevent stale cache in multi-stage workflows
@@ -201,10 +205,32 @@ class ComfyQwenImageWrapperV2(nn.Module):
 
             # --- END NEW VRAM CHECK ---
 
-            # 4. Compose LoRAs using v2-specific function. This changes internal tensor shapes.
+            # 4. Build cache args — same pattern as the V1 wrapper.
-            # Returns True if successful (supported format), False if unsupported (skipped).
+            try:
+                _comfy_base = str(Path(folder_paths.models_dir).parent)
+            except Exception:
+                _comfy_base = None
+
+            # Separate per-LoRA meta dicts and derive the global save_precompiled flag.
+            _lora_configs_clean = []
+            _any_save_flag = False
+            for _entry in self.loras:
+                _lc_path = _entry[0]
+                _lc_strength = _entry[1]
+                _lc_meta: dict = _entry[2] if len(_entry) > 2 else {}
+                _lora_configs_clean.append((_lc_path, _lc_strength, _lc_meta))
+                if _lc_meta.get("save_precompiled", False):
+                    _any_save_flag = True
+
+            # 5. Compose LoRAs using v2-specific function.
             logger.info(f"[V2 Node] 🔧 Forward pass: calling compose_loras_v2_v2 with apply_awq_mod={self.apply_awq_mod} (type: {type(self.apply_awq_mod).__name__})")
-            is_supported_format = compose_loras_v2_v2(self.model, self.loras, apply_awq_mod=self.apply_awq_mod)
+            is_supported_format = compose_loras_v2_v2(
+                self.model,
+                _lora_configs_clean,
+                apply_awq_mod=self.apply_awq_mod,
+                save_precompiled_lora=_any_save_flag,
+                cache_dir=_comfy_base,
+            )
 
             # Validate composition result; if 0 targets after a crash/transition, retry once
             # But ONLY if the format was supported. If unsupported, retrying is pointless.
@@ -217,7 +243,13 @@ class ComfyQwenImageWrapperV2(nn.Module):
                     logger.warning("[V2 Node] LoRA composition reported 0 target modules. Forcing reset and one retry.")
                     try:
                         reset_lora_v2(self.model)
-                        compose_loras_v2_v2(self.model, self.loras, apply_awq_mod=self.apply_awq_mod)
+                        compose_loras_v2_v2(
+                            self.model,
+                            _lora_configs_clean,
+                            apply_awq_mod=self.apply_awq_mod,
+                            save_precompiled_lora=_any_save_flag,
+                            cache_dir=_comfy_base,
+                        )
                     except Exception as e:
                         logger.error(f"[V2 Node] LoRA re-compose retry failed: {e}")
             else:
@@ -251,7 +283,9 @@ class ComfyQwenImageWrapperV2(nn.Module):
 
             # --- END MODIFIED SECTION ---
 
-            # Update last known device signature after (re)composition
+            # Always update state tracking so subsequent calls see a consistent snapshot.
+            self.model._applied_loras = self.loras.copy()
+            self._applied_loras = self.loras.copy()
             self._last_device = current_device
 
         # Caching logic