diff --git a/.gitignore b/.gitignore index 7535505..29e52dc 100644 --- a/.gitignore +++ b/.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 - -# Ignore scripts folder -scripts/ - -# Ignore backups -backups/ -backup/ -# Ignore reflection files -反省文* -反省文/ +.tracking \ No newline at end of file diff --git a/__init__.py b/__init__.py index 429819d..d7e0c81 100644 --- a/__init__.py +++ b/__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" diff --git a/js/z_qwen_lora_dynamic.js b/js/z_qwen_lora_dynamic.js index 2ea68bc..e7b35b5 100644 --- a/js/z_qwen_lora_dynamic.js +++ b/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]); diff --git a/js/z_qwen_lora_dynamic_v3.js b/js/z_qwen_lora_dynamic_v3.js index 88e94fa..f9099d8 100644 --- a/js/z_qwen_lora_dynamic_v3.js +++ b/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); }; diff --git a/nodes/lora/qwenimage.py b/nodes/lora/qwenimage.py index 30348f2..989cd1c 100644 --- a/nodes/lora/qwenimage.py +++ b/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/_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)) - logger.debug(f"LoRA added to stack: {lora_name} (strength={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.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)}") diff --git a/nodes/lora/qwenimage_v1.py b/nodes/lora/qwenimage_v1.py index 0b3c206..547c387 100644 --- a/nodes/lora/qwenimage_v1.py +++ b/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)) - logger.debug(f"LoRA added to stack: {lora_name} (strength={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}, save_precompiled={save_precompiled_lora})") logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}") return (ret_model,) diff --git a/nodes/lora/qwenimage_v2.py b/nodes/lora/qwenimage_v2.py index 0355f1d..71b8001 100644 --- a/nodes/lora/qwenimage_v2.py +++ b/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)) - logger.debug(f"LoRA added to stack: {lora_name} (strength={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}, save_precompiled={save_precompiled_lora})") logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}") diff --git a/nodes/lora/qwenimage_v3.py b/nodes/lora/qwenimage_v3.py index 3387f8f..fb7bdb7 100644 --- a/nodes/lora/qwenimage_v3.py +++ b/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) - print(f"--- DEBUG: Inspecting keys for LoRA 1 (Strength: {_fs}) ---") + # First LoRA: Detailed key inspection (only when verbose logging is on) _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).") - 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 ---") + logger.debug("⚠️ Unsupported LoRA format detected (No standard keys). Skipping 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)) - logger.debug(f"LoRA added to stack: {lora_name} (strength={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}, save_precompiled={save_precompiled_lora})") logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}") return (ret_model,) diff --git a/nodes/te_offload/__init__.py b/nodes/te_offload/__init__.py new file mode 100644 index 0000000..3ebf059 --- /dev/null +++ b/nodes/te_offload/__init__.py @@ -0,0 +1,2 @@ +# TE offload nodes package + diff --git a/nodes/te_offload/nunchaku_te_v2.py b/nodes/te_offload/nunchaku_te_v2.py new file mode 100644 index 0000000..28753ee --- /dev/null +++ b/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)", +} diff --git a/nunchaku_code/lora_cache.py b/nunchaku_code/lora_cache.py new file mode 100644 index 0000000..90ebc0f --- /dev/null +++ b/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/ + _precompiled.safetensors ← fused tensors + _precompiled.safetensors.meta.json ← mtime sidecar + +Tensor key format inside the .safetensors file: + "__A" → fused A tensor [rank, in_features] + "__B" → fused B tensor [out_features, rank] (NOT strength-scaled) + "__alpha" → alpha scalar as float32 1-D tensor (optional, omitted if None) + +Sidecar JSON format: + { + "source_mtime": , # os.path.getmtime() of the original .safetensors + "source_path": # 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 ``/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 ``/__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 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 diff --git a/nunchaku_code/lora_qwen.py b/nunchaku_code/lora_qwen.py index 9a7f9bd..f358374 100644 --- a/nunchaku_code/lora_qwen.py +++ b/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] - new_proj_down = torch.cat([pd, A], dim=0) + pd_unpacked = unpack_lowrank_weight(pd_gpu, down=True) + 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) - module.proj_up.data = pack_lowrank_weight(new_proj_up, down=False) + packed_down = pack_lowrank_weight(new_proj_down, down=True) + 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 - _cached_first_lora_state_dict = None + # Safe unpack: entries may be 2-tuples (path, strength) or 3-tuples + # (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_lora_state_dict = _load_lora_state_dict_robust(first_lora_path_or_dict) - _cached_first_lora_state_dict = first_lora_state_dict # Cache for reuse + _first_entry = lora_configs[0] + first_lora_path_or_dict = _first_entry[0] + first_lora_strength = _first_entry[1] - # 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) + # Optimize: Check if cache is valid for the first LoRA. If so, skip loading the raw safetensors file + _cache_path = None + _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 - if idx == 0 and _cached_first_lora_state_dict is not None: - lora_state_dict = _cached_first_lora_state_dict - else: - lora_state_dict = _load_lora_state_dict_robust(lora_path_or_dict) + # Per-entry save flag: honour node-level flag OR the global function arg. + _should_save = save_precompiled_lora or bool(_entry_meta.get("save_precompiled", False)) + + # ------------------------------------------------------------------ # + # 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_B_scaled.append((B * scale).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)) 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) - pu = unpack_lowrank_weight(module.proj_up.data, down=False) - + # Use ComfyUI-aware device to stay compatible with its VRAM manager + 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() - - module.proj_down.data = pack_lowrank_weight(pd_reset, down=True) - module.proj_up.data = pack_lowrank_weight(pu_reset, down=False) + del pd, pu + + 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.") \ No newline at end of file diff --git a/patches/nunchaku_patch.py b/patches/nunchaku_patch.py index b64b8f2..36235ef 100644 --- a/patches/nunchaku_patch.py +++ b/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 diff --git a/wrappers/qwenimage.py b/wrappers/qwenimage.py index 9b15e76..fb8cd31 100644 --- a/wrappers/qwenimage.py +++ b/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 - hasattr(self.model, "_lora_slots") and self.model._lora_slots # But the model actually has LoRA + self._has_ever_had_loras and # We were (or are) a LoRA manager + 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 - reset_lora_v2(self.model) - self._applied_loras = self.loras.copy() - + # Guard: only reset/recompose the shared transformer when this wrapper + # is (or has been) an active LoRA manager. The bypass-path wrapper W1 + # (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) - offload_is_on = hasattr(self.model, "offload_manager") and self.model.offload_manager is not None + # 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 - # 2. Decide if we *need* to turn it on - should_enable_offload = offload_is_on + # 2. Decide if we *need* to turn it on + should_enable_offload = offload_is_on - # 3. Only run the dynamic VRAM check if: - # - The user's original setting was "auto" - # - Offloading is not *already* on - # - We are actually loading new 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.") + # 3. Only run the dynamic VRAM check if: + # - The user's original setting was "auto" + # - Offloading is not *already* on + # - We are actually loading new LoRAs + if self.cpu_offload_setting == "auto" and not offload_is_on and self.loras: try: - reset_lora_v2(self.model) - compose_loras_v2(self.model, self.loras, apply_awq_mod=self.apply_awq_mod) + # 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"LoRA re-compose retry failed: {e}") - else: - logger.warning("Skipping retry because LoRA format is unsupported.") + 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.") - # 5. Re-build offload manager if it's supposed to be on - # This block now runs if offload was on *or* if our new check decided to turn it on. - if should_enable_offload: + # --- END NEW VRAM CHECK --- - # Store settings if it was already on, otherwise use defaults - if offload_is_on: - manager = self.model.offload_manager - offload_settings = { - "num_blocks_on_gpu": manager.num_blocks_on_gpu, - "use_pin_memory": manager.use_pin_memory, - } + # 4. Build cache args. The ComfyUI base dir is the parent of the + # 'models' folder, which folder_paths knows about. + try: + _comfy_base = str(Path(folder_paths.models_dir).parent) + except Exception: + _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 - 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.") + logger.warning("Skipping retry because LoRA format is unsupported.") - # Step 1: Completely disable and clear any old offloader (safe to call even if off) - self.model.set_offload(False) + # 5. Re-build offload manager if it's supposed to be on + # 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 - # This builds the manager based on the *newly composed* tensor shapes. - self.model.set_offload(True, **offload_settings) + # Store settings if it was already on, otherwise use defaults + if offload_is_on: + 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 diff --git a/wrappers/qwenimage_v2.py b/wrappers/qwenimage_v2.py index a73efbf..475a112 100644 --- a/wrappers/qwenimage_v2.py +++ b/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. - # Returns True if successful (supported format), False if unsupported (skipped). + # 4. Build cache args — same pattern as the V1 wrapper. + 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