diff --git a/.gitignore b/.gitignore index 41d059d..a7a13d0 100644 --- a/.gitignore +++ b/.gitignore @@ -67,5 +67,8 @@ venv.bak/ Thumbs.db # User-specific documentation files -technical_explanation.md -PR34_COMFYUI_V0.6.0_COMPATIBILITY.md \ No newline at end of file +# Ignore root directory versions, but allow versions in md/ folder +/technical_explanation.md +PR34_COMFYUI_V0.6.0_COMPATIBILITY.md +# Ignore Japanese version in root, but allow English version in md/ folder +/ZIMAGETURBO_CONTROLNET_FIX.md \ No newline at end of file diff --git a/README.md b/README.md index 20f78cd..2507626 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,12 @@ -# ComfyUI-QwenImageLoraLoader +# ComfyUI-Nunchaku QwenImage&Z-ImageTurboLoraLoader and Diffsynth ControlNet Loader -A ComfyUI custom node for loading and applying LoRA (Low-Rank Adaptation) to Nunchaku Qwen Image models. ComfyUI Nodes 2.0 compatible. +A ComfyUI custom node for loading and applying LoRA (Low-Rank Adaptation) to Nunchaku Qwen Image and Z-ImageTurbo models, and diffsynth ControlNet functionality. ComfyUI Nodes 2.0 compatible. -> Latest release: [v2.0.8 on GitHub Releases](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v2.0.8) +## ⚠️ **DEVELOPMENT STATUS** + +**Currently under development and testing. Debug logs are being output extensively. This does not affect functionality.** + +> Latest release: [v2.2.3 on GitHub Releases](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/2.2.3) > > ⚠️ **Note for v2.0+ users**: If you encounter `TypeError: got multiple values for argument 'guidance'` errors, see [troubleshooting section](#issue-30-typeerror-got-multiple-values-for-argument-guidance-v20) below. @@ -13,14 +17,22 @@ This LoRA loader was extracted and modified from GavChap's fork: - **Extraction**: LoRA functionality was extracted from the full fork to create an independent custom node - **Integration**: Modified to work with the official ComfyUI-nunchaku plugin +## 🎉 MAJOR UPDATE: v2.2.2 - Nunchaku Z-ImageTurbo Diffsynth ControlNet Support Added! + +For detailed technical explanation, see [v2.2.2 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/2.2.2) + +## 🎉 MAJOR UPDATE: v2.2.0 - Nunchaku Z Image Turbo LoRA Support Added! + +For detailed technical explanation, see [v2.2.0 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/2.2.0) + ## 🎉 MAJOR UPDATE: v2.0 - Diffsynth ControlNet Support Added! -**As of v2.0, diffsynth ControlNet is now fully supported for Nunchaku Qwen Image models.** +**As of v2.0, diffsynth ControlNet is now fully supported for Nunchaku Qwen Image models, Z Image Turbo BF16.safetensors, and Nunchaku Z Image Turbo models.** -A new dedicated node `NunchakuQwenImageDiffsynthControlnet` enables diffsynth ControlNet functionality with Nunchaku quantized Qwen Image models. +A new dedicated node `NunchakuQI&ZITDiffsynthControlnet` enables diffsynth ControlNet functionality with Nunchaku quantized Qwen Image models, Z Image Turbo BF16.safetensors, and Nunchaku Z Image Turbo models. ### What's New in v2.0 -- ✅ **New Node**: `NunchakuQwenImageDiffsynthControlnet` - Dedicated diffsynth ControlNet loader for Nunchaku Qwen Image models +- ✅ **New Node**: `NunchakuQI&ZITDiffsynthControlnet` - Dedicated diffsynth ControlNet loader for Nunchaku Qwen Image models, Z Image Turbo BF16.safetensors, and Nunchaku Z Image Turbo models - ✅ **Full ControlNet Support**: Works with standard diffsynth ControlNet models - ✅ **Seamless Integration**: Automatically applies ControlNet patches during model forward pass - ✅ **Backward Compatible**: All existing LoRA functionality remains unchanged @@ -29,62 +41,7 @@ For detailed technical explanation, see [v2.0 Release Notes](https://github.com/ ## 🎉 MAJOR UPDATE: v1.60 - Simplified Installation (No Integration Required!) -**As of v1.60, manual integration with ComfyUI-nunchaku's `__init__.py` is completely unnecessary.** - -The node now operates as a fully independent custom node that works out-of-the-box. Simply clone the repository into your `custom_nodes` folder and restart ComfyUI. The nodes will automatically appear in ComfyUI's node menu through ComfyUI's built-in automatic node loading mechanism. - -### What Changed in v1.60 -- ✅ **Removed dependency on ComfyUI-nunchaku integration** - The LoRA loader is now a standalone plugin -- ✅ **Simplified installation** - No batch scripts or manual file editing required -- ✅ **Cleaner architecture** - Node registration happens automatically -- ✅ **Backward compatible** - All existing LoRA files and workflows continue to work - -For a detailed technical explanation of why integration is no longer needed, see [v1.60 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v1.60) - -## Features - -- **NunchakuQwenImageLoraLoader**: Load and apply single LoRA to Qwen Image models -- **NunchakuQwenImageLoraStack**: Apply multiple LoRAs with dynamic UI control -- **NunchakuQwenImageDiffsynthControlnet**: Apply diffsynth ControlNet to Nunchaku Qwen Image models (v2.0) -- **Dynamic VRAM Management**: Automatic CPU offloading based on available VRAM -- **LoRA Composition**: Efficient LoRA stacking and composition -- **ComfyUI Integration**: Seamless integration with ComfyUI workflows - -## Installation - -### Quick Installation (v1.60 - Simplified!) - -**Prerequisites:** -- ComfyUI-nunchaku must be installed - -1. Clone this repository to your ComfyUI custom_nodes directory: - -```bash -cd ComfyUI/custom_nodes -git clone https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader.git -``` - -2. Restart ComfyUI - -**That's it!** The nodes will automatically appear in ComfyUI's node menu. - -### Manual Installation (Alternative) - -If you prefer to install manually or are using macOS/Linux: - -```bash -cd /path/to/ComfyUI/custom_nodes -git clone https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader.git -``` - -Then restart ComfyUI. - -### Requirements - -- Python 3.11+ -- ComfyUI (latest version recommended) -- ComfyUI-nunchaku (required) -- CUDA-capable GPU (optional, but recommended for performance) +For installation instructions, features, and requirements, see [Installation Guide](md/installation.md). ## Upgrade Guide for v1.57 and Earlier Users @@ -96,16 +53,29 @@ If you have v1.57 or earlier installed with integration code in ComfyUI-nunchaku - **NunchakuQwenImageLoraLoader**: Single LoRA loader - **NunchakuQwenImageLoraStack**: Multi LoRA stacker with dynamic UI (Legacy) - **NunchakuQwenImageLoraStackV2**: Multi LoRA stacker with dynamic UI - ComfyUI Nodes 2.0 (Beta) compatible -- **NunchakuQwenImageDiffsynthControlnet**: Diffsynth ControlNet loader for Nunchaku Qwen Image models (v2.0) +- **NunchakuZImageTurboLoraStackV3**: Z-Image-Turbo LoRA stacker with dynamic UI - ComfyUI Nodes 2.0 (Beta) compatible +- **NunchakuQI&ZITDiffsynthControlnet**: Diffsynth ControlNet loader for Nunchaku Qwen Image models, Z Image Turbo BF16.safetensors, and Nunchaku Z Image Turbo models (v2.0) ### Basic Usage + +**For Nunchaku Qwen Image models:** 1. Load your Nunchaku Qwen Image model using `Nunchaku Qwen Image DiT Loader` 2. Add either `NunchakuQwenImageLoraLoader` or `NunchakuQwenImageLoraStack` node 3. Select your LoRA file and set the strength 4. Connect to your workflow +**For Nunchaku Z-Image-Turbo models:** +1. Load your Nunchaku Z-Image-Turbo model using `Nunchaku Z-Image DiT Loader` +2. Add `Nunchaku Z-Image-Turbo LoRA Stack V3` node +3. Select your LoRA file and set the strength +4. Connect to your workflow + ### Dynamic UI Control -The `NunchakuQwenImageLoraStack` node automatically adjusts the number of visible LoRA slots based on the `lora_count` parameter (1-10). + +Dynamic UI Control +Dynamic UI Control V3 + +The `NunchakuQwenImageLoraStack` and `NunchakuZImageTurboLoraStackV3` nodes automatically adjust the number of visible LoRA slots based on the `lora_count` parameter (1-10). ### Diffsynth ControlNet Usage @@ -113,7 +83,7 @@ The `NunchakuQwenImageLoraStack` node automatically adjusts the number of visibl 1. Load your diffsynth ControlNet model patch using `Model Patch Loader` from [ComfyUI-NunchakuFluxLoraStacker](https://github.com/ussoewwin/ComfyUI-NunchakuFluxLoraStacker) 2. The `Model Patch Loader` (`ModelPatchLoaderCustom`) supports CPU offload, allowing you to load ControlNet patches to CPU memory to save VRAM -3. Connect the `MODEL_PATCH` output to the `model_patch` input of `NunchakuQwenImageDiffsynthControlnet` node +3. Connect the `MODEL_PATCH` output to the `model_patch` input of `NunchakuQI&ZITDiffsynthControlnet` node 4. Connect your Nunchaku Qwen Image model, VAE, and control image 5. Set the ControlNet strength and connect to your workflow @@ -184,52 +154,9 @@ This node is designed to work with: ### Issue #25: ComfyUI 0.4.0 Model Management Errors - **Status**: ⚠️ **Environment Dependent** - May require ComfyUI core fixes -- **Issue**: After the ComfyUI 0.4.0 update, multiple nodes (including this one in some environments) experienced errors such as `TypeError: 'NoneType' object is not callable` and `AttributeError: 'NoneType' object has no attribute`. In our environment, we resolved these errors by modifying ComfyUI's core `model_management.py`. Note that in our environment, these errors did not occur with this node (ComfyUI-QwenImageLoraLoader). Nunchaku library and ComfyUI-Nunchaku nodes should use the latest versions. If errors persist even after applying the latest version of this node (ComfyUI-QwenImageLoraLoader), modification of ComfyUI's core `model_management.py` may be necessary. -- **Root Cause**: In ComfyUI 0.4.0, ComfyUI's core `model_management.py` lacks sufficient None checks, causing `TypeError` and `AttributeError` when accessing methods or attributes on objects that became `None` after models were unloaded or garbage collected. This problem is not a bug in individual nodes, but rather a structural issue in ComfyUI 0.4.0's model management (`model_management.py`). -- **GC Changes in ComfyUI 0.4.0**: Compared to ComfyUI 0.3.x, automatic model unloading occurs earlier, making the following flow more likely: -``` -ModelPatcher → GC → weakref(None) -``` +For detailed information, see [COMFYUI_0.4.0_MODEL_MANAGEMENT_ERRORS.md](md/COMFYUI_0.4.0_MODEL_MANAGEMENT_ERRORS.md). -This also explains why the occurrence of the issue varies by user environment. -- **Technical Basis**: - 1. **Multiple Locations with "Missing None Checks"** - This is not a bug in individual nodes, but the core main logic crashes when accessing attributes on `None`. The added fixes like `if model is None: continue` are defensive code that ComfyUI core should have in all paths. - 2. **Post-Weak-Reference GC Behavior Not Considered** - The introduction of `LoadedModel._model = weakref.ref(ModelPatcher)` in ComfyUI 0.4.0 is a breaking change. When the weak reference target is garbage collected, it returns `None`, but this is not handled. Post-processing for the breaking update is incomplete. - 3. **Multiple Nodes Were Affected in a Chain Reaction** - This is not a problem with nodes, but multiple nodes were affected in a chain reaction due to core behavior changes. Model loading/unloading, memory calculation, GPU/CPU offloading, and `ModelPatcher` lifecycle are all controlled by ComfyUI core. - 4. **All Fix Locations Are Core Responsibility Areas** - The locations fixed (`model_memory`, `model_offloaded_memory`, `load_models_gpu`, `free_memory`, `model_unload`, `is_dead` checks, etc.) are all ComfyUI core functions. These are not areas that node developers should touch. The fact that all fix locations are core logic leaves no explanation other than a core defect. - 5. **Result of Applying Fixes** - After applying None check fixes to ComfyUI's core `model_management.py` in our environment, similar errors were resolved. This demonstrates that the problem can be solved by adding defensive code that the core should have. -- **Model Lifecycle and ModelPatcher Initialization Relationship**: - - **Fact 1: Relationship between LoadedModel and ModelPatcher** - The `LoadedModel` class (lines 502-524 in ComfyUI's `model_management.py`) holds a weak reference to `ModelPatcher`: - -```python -def _set_model(self, model): - self._model = weakref.ref(model) # Weak reference to ModelPatcher - -@property -def model(self): - return self._model() # Returns None when garbage collected -``` - - - **Fact 2: ModelPatcher Initialization** - In the `__init__` of the `ModelPatcher` class (lines 215-237 in `model_patcher.py`), the `pinned` attribute is initialized: - -```python -def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False): - # ... - self.pinned = set() # Line 237: Initialized -``` - - - **Fact 3: Fix Content in ComfyUI Core's model_management.py** - The fix in ComfyUI's core `model_management.py` now skips `LoadedModel` instances where `model` is `None`: In `load_models_gpu()`, skips `LoadedModel` instances where `model` is `None` (lines 712, 727, 743); In `free_memory()`, excludes `LoadedModel` instances where `model` is `None` (line 646). - - **Fact 4: Problem Before Fix** - `LoadedModel` holds a weak reference to `ModelPatcher`. When garbage collected, `LoadedModel.model` returns `None`. Before the fix, methods were called on `LoadedModel` instances where `model` was `None`, causing errors. - - **Fact 5: Behavior After Fix** - By skipping `LoadedModel` instances where `model` is `None`, errors do not occur. Because errors do not occur, processing continues normally. - - **Fact 6: Why copy.deepcopy Fails** - `copy.deepcopy` fails because references to GC'd `ModelPatcher` instances remain in the dictionary being deepcopied. When these references are accessed, they return `None`, causing deepcopy to stop. - - **Fact 7: Confirmation Items** - After applying the fix, `copy.deepcopy` and `pinned` attribute errors do not occur in our environment. Nunchaku library and ComfyUI-Nunchaku nodes should use the latest versions, but this may still be insufficient. While these errors did not occur with this node (ComfyUI-QwenImageLoraLoader) in our environment, the fix to ComfyUI's core `model_management.py` may have indirectly affected it, making errors less likely to occur. -- **Important Note: Not a Problem with Nunchaku Library** - This problem is not caused by the Nunchaku library's implementation. Nunchaku's `model_config` and `ModelPatcher` itself are normal. The problem is in the upstream = ComfyUI core's `model_management.py` GC processing. -- **Speculation (Items That May Be Environment-Dependent)**: The fix allows `ModelPatcher` initialization to complete normally. As a result, the `pinned` attribute is also properly initialized. Accessing `self.pinned` in `__del__` does not cause errors. -- **Recommendations**: - 1. Update Nunchaku library and ComfyUI-Nunchaku nodes to the latest version (addresses `model_config` issues) - 2. Consider applying None check fixes to ComfyUI's core `model_management.py` (may address the root cause) -- **Note**: This is the first support measure. I have published the technical details of the fixes I applied to ComfyUI's core `model_management.py` in my environment. See [COMFYUI_0.4.0_UPDATE_ERROR_FIXES.md](md/COMFYUI_0.4.0_UPDATE_ERROR_FIXES.md) for details. Note that these fixes were applied in my specific environment and may not work universally in all environments. This may also resolve `copy.deepcopy` and `pinned` attribute errors. - **Related Issues**: - [Issue #25](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/25) - `AttributeError: 'NunchakuModelPatcher' object has no attribute 'pinned'` and deepcopy errors with `model_config` - [Issue #33](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/33) - `AttributeError: 'NoneType' object has no attribute 'to'` in `to_safely` method (Fixed in v2.1.0) @@ -283,24 +210,18 @@ def __init__(self, model, load_device, offload_device, size=0, weight_inplace_up ## Changelog -### v2.2.4 (latest) -- **Added**: AWQ Modulation Layer Detection and Skip Logic -- **Feature**: Default skip of LoRA application to `img_mod.1` and `txt_mod.1` modulation layers in AWQ quantized models to prevent image noise -- **Environment Variable**: `QWENIMAGE_LORA_APPLY_AWQ_MOD=1` to force-enable (not recommended due to noise risk) -- **Technical Details**: See [v2.2.4 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/2.2.4) for complete explanation -- **Documentation**: See [AWQ_MODULATION_LORA_EXPLANATION.md](AWQ_MODULATION_LORA_EXPLANATION.md) for detailed technical explanation - -### v2.2.3 -- **Added**: Toggle buttons and LoRA format detection features +### v2.2.3 (latest) +- **Added**: Toggle buttons to enable/disable individual LoRA slots and all LoRAs at once ([Issue #12](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/12), [Issue #36](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/36)) +- ⚠️ **DEVELOPMENT STATUS**: These features are currently experimental implementations for the `NunchakuZImageTurboLoraStackV3` node only. ComfyUI Nodes 2.0 environment only. With current technical capabilities, it is not possible to fully implement all requested features in JavaScript. - **Technical Details**: See [v2.2.3 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/2.2.3) for complete explanation ### v2.2.2 -- **Added**: Nunchaku Z-ImageTurbo Diffsynth ControlNet support +- **Added**: Diffsynth ControlNet support for Nunchaku Z-ImageTurbo models - **Technical Details**: See [v2.2.2 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/2.2.2) for complete explanation ### v2.2.0 -- **MAJOR UPDATE**: Added Nunchaku Z-Image-Turbo LoRA Loader V3 Support -- **Technical Details**: See [v2.2.0 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/2.2.0) for complete explanation +- **Added**: NunchakuZImageTurboLoraStackV3 node – Z-Image-Turbo LoRA stacker with dynamic UI for official Nunchaku Z-Image loader +- **Technical Details**: See [v2.2.0 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v2.2.0) for complete explanation ### v2.1.1 - **Fixed**: ComfyUI v0.6.0+ compatibility – Migrated from `guidance` to `additional_t_cond` parameter in `_execute_model` method to support ComfyUI v0.6.0+ API changes ([PR #34](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/pull/34)) @@ -310,11 +231,9 @@ def __init__(self, model, load_device, offload_device, size=0, weight_inplace_up - **Fixed**: Resolved [Issue #33](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/33) – Fixed `AttributeError: 'NoneType' object has no attribute 'to'` by adding None checks to `to_safely` and `forward` methods in `ComfyQwenImageWrapper` - **Technical Details**: See [v2.1.0 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v2.1.0) for complete explanation -### v2.0.8 -- **Fixed**: Resolved [Issue #30](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/30) – Fixed `TypeError: got multiple values for argument 'guidance'` error by adding final cleanup of kwargs before calling model forward -- **Technical Details**: See [v2.0.8 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v2.0.8) for complete explanation - -### v2.0.2 - v2.0.7 +### v2.0.2 - v2.0.8 +- **v2.0.8 Fixed**: Resolved [Issue #30](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/30) – Fixed `TypeError: got multiple values for argument 'guidance'` error by adding final cleanup of kwargs before calling model forward + - **Technical Details**: See [v2.0.8 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v2.0.8) for complete explanation - **v2.0.7 Fixed**: Enhanced [Issue #32](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/32) fix by adding exclusion processing in `forward` method in addition to `_execute_model` method to prevent duplicate argument errors - **Technical Details**: See [v2.0.7 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v2.0.7) for complete explanation - **v2.0.6 Fixed**: Excluded `ref_latents`, `transformer_options`, and `attention_mask` from kwargs to prevent duplicate argument errors @@ -330,7 +249,7 @@ def __init__(self, model, load_device, offload_device, size=0, weight_inplace_up ### v2.0 - **MAJOR UPDATE**: Added diffsynth ControlNet support for Nunchaku Qwen Image models -- **New Node**: `NunchakuQwenImageDiffsynthControlnet` - Enables diffsynth ControlNet to work with Nunchaku quantized Qwen Image models +- **New Node**: `NunchakuQI&ZITDiffsynthControlnet` - Enables diffsynth ControlNet to work with Nunchaku quantized Qwen Image models, Z Image Turbo BF16.safetensors, and Nunchaku Z Image Turbo models - **Features**: - Full diffsynth ControlNet functionality for Nunchaku Qwen Image models - Automatic patch registration and application @@ -357,10 +276,7 @@ def __init__(self, model, load_device, offload_device, size=0, weight_inplace_up - Complete feature parity with V1 implementation - Dynamic UI for adjusting slot count - Automatic node height adjustment -- **Technical Details**: - - Flexible widget management using `optional` inputs - - Clean implementation using Flux V2-style widget array reconstruction - - Complete feature reproduction in both Python backend and JavaScript frontend +- **Technical Details**: See [v1.70 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v1.70) for complete explanation ### v1.63 - **Fixed**: Addressed [Issue #21](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/issues/21) – User-configurable CPU offload setting diff --git a/__init__.py b/__init__.py index 0526e0f..a7bc028 100644 --- a/__init__.py +++ b/__init__.py @@ -20,20 +20,31 @@ QWEN_V2_NAMES = {} try: from .nodes.lora.qwenimage import NunchakuQwenImageLoraLoader, NunchakuQwenImageLoraStack from .nodes.lora.qwenimage_v2 import GENERATED_NODES as QWEN_V2_NODES, GENERATED_DISPLAY_NAMES as QWEN_V2_NAMES - from .nodes.lora.zimageturbo_v2 import GENERATED_NODES as ZIMAGETURBO_V2_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V2_NAMES + from .nodes.lora.qwenimage_v3 import GENERATED_NODES as QWEN_V3_NODES, GENERATED_DISPLAY_NAMES as QWEN_V3_NAMES + # Z-Image-Turbo V2 is deprecated (unofficial loader only) - keep file but do not register to avoid confusion + # from .nodes.lora.zimageturbo_v2 import GENERATED_NODES as ZIMAGETURBO_V2_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V2_NAMES + from .nodes.lora.zimageturbo_v3 import GENERATED_NODES as ZIMAGETURBO_V3_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V3_NAMES # Add version to classes before creating NODE_CLASS_MAPPINGS NunchakuQwenImageLoraLoader.__version__ = __version__ NunchakuQwenImageLoraStack.__version__ = __version__ for node_class in QWEN_V2_NODES.values(): node_class.__version__ = __version__ - for node_class in ZIMAGETURBO_V2_NODES.values(): + for node_class in QWEN_V3_NODES.values(): + node_class.__version__ = __version__ + # Z-Image-Turbo V2 registration removed (unofficial loader only) + # for node_class in ZIMAGETURBO_V2_NODES.values(): + # node_class.__version__ = __version__ + for node_class in ZIMAGETURBO_V3_NODES.values(): node_class.__version__ = __version__ NODE_CLASS_MAPPINGS["NunchakuQwenImageLoraLoader"] = NunchakuQwenImageLoraLoader NODE_CLASS_MAPPINGS["NunchakuQwenImageLoraStack"] = NunchakuQwenImageLoraStack NODE_CLASS_MAPPINGS.update(QWEN_V2_NODES) - NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V2_NODES) + NODE_CLASS_MAPPINGS.update(QWEN_V3_NODES) + # Z-Image-Turbo V2 registration removed (unofficial loader only) + # NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V2_NODES) + NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V3_NODES) except ImportError: logger.exception("LoRA nodes import failed:") @@ -52,12 +63,15 @@ NODE_DISPLAY_NAME_MAPPINGS = { "NunchakuQwenImageLoraLoader": "Nunchaku Qwen Image LoRA Loader", "NunchakuQwenImageLoraStack": "Nunchaku Qwen Image LoRA Stack (Legacy)", **QWEN_V2_NAMES, - **ZIMAGETURBO_V2_NAMES + **QWEN_V3_NAMES, + # Z-Image-Turbo V2 registration removed (unofficial loader only) + # **ZIMAGETURBO_V2_NAMES, + **ZIMAGETURBO_V3_NAMES } # Add ControlNet display name only if the node was successfully loaded if "NunchakuQwenImageDiffsynthControlnet" in NODE_CLASS_MAPPINGS: - NODE_DISPLAY_NAME_MAPPINGS["NunchakuQwenImageDiffsynthControlnet"] = "Nunchaku Qwen Image Diffsynth Controlnet" + NODE_DISPLAY_NAME_MAPPINGS["NunchakuQwenImageDiffsynthControlnet"] = "NunchakuQI&ZITDiffsynthControlnet" # Register JavaScript extensions WEB_DIRECTORY = "js" diff --git a/images/qilorastacker.png b/images/qilorastacker.png new file mode 100644 index 0000000..60c3166 Binary files /dev/null and b/images/qilorastacker.png differ diff --git a/images/zitlorav3.png b/images/zitlorav3.png new file mode 100644 index 0000000..3a54a4c Binary files /dev/null and b/images/zitlorav3.png differ diff --git a/js/zimageturbo_lora_dynamic_v3.js b/js/zimageturbo_lora_dynamic_v3.js new file mode 100644 index 0000000..9718c37 --- /dev/null +++ b/js/zimageturbo_lora_dynamic_v3.js @@ -0,0 +1,218 @@ +import { app } from "../../scripts/app.js"; + +console.log("★★★ zimageturbo_lora_dynamic.js: Z-Image-Turbo LoRA Stack V3 ★★★"); + +const HIDDEN_TAG = "tschide"; + +app.registerExtension({ + name: "nunchaku.zimageturbo_lora_dynamic_v3", + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "NunchakuZImageTurboLoraStackV3") { + nodeType["@visibleLoraCount"] = { type: "number", default: 1, min: 1, max: 10, step: 1 }; + } + }, + + nodeCreated(node) { + if (node.comfyClass !== "NunchakuZImageTurboLoraStackV3") return; + + if (!node.properties) node.properties = {}; + if (node.properties["visibleLoraCount"] === undefined) node.properties["visibleLoraCount"] = 1; + + // Immediately hide lora_count widget if it exists + const initialLoraCountWidget = node.widgets?.find(w => w.name === "lora_count"); + if (initialLoraCountWidget) { + if (!initialLoraCountWidget.origType) { + initialLoraCountWidget.origType = initialLoraCountWidget.type; + initialLoraCountWidget.origComputeSize = initialLoraCountWidget.computeSize; + } + initialLoraCountWidget.type = HIDDEN_TAG; + initialLoraCountWidget.computeSize = () => [0, -4]; + } + + node.cachedWidgets = {}; + let cacheReady = false; + + const initCache = () => { + if (cacheReady) return; + const all = [...node.widgets]; + + // Cache lora_count widget (required for Python backend, but hidden in UI) + const loraCountWidget = all.find(w => w.name === "lora_count"); + if (loraCountWidget) { + node.cachedLoraCount = loraCountWidget; + // Store original properties for restoration if needed + if (!loraCountWidget.origType) { + loraCountWidget.origType = loraCountWidget.type; + loraCountWidget.origComputeSize = loraCountWidget.computeSize; + } + // Hide V1's lora_count widget using HIDDEN_TAG and computeSize + loraCountWidget.type = HIDDEN_TAG; + loraCountWidget.computeSize = () => [0, -4]; + } + + // Cache cpu_offload widget + const cpuOffloadWidget = all.find(w => w.name === "cpu_offload"); + if (cpuOffloadWidget) { + node.cachedCpuOffload = cpuOffloadWidget; + } + + // Cache toggle_all widget + const toggleAllWidget = all.find(w => w.name === "toggle_all"); + if (toggleAllWidget) { + node.cachedToggleAll = toggleAllWidget; + } + + 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}`); + const wStrength = all.find(w => w.name === `lora_strength_${i}`); + if (wEnabled && wName && wStrength) { + node.cachedWidgets[i] = [wEnabled, wName, wStrength]; + wEnabled.type = "toggle"; + wName.type = "combo"; + wStrength.type = "number"; + if (wEnabled.computeSize) delete wEnabled.computeSize; + if (wName.computeSize) delete wName.computeSize; + if (wStrength.computeSize) delete wStrength.computeSize; + } + } + cacheReady = true; + }; + + + + const ensureControlWidget = () => { + const name = "🔢 LoRA Count"; + + // Remove old button widgets + for (let i = node.widgets.length - 1; i >= 0; i--) { + const w = node.widgets[i]; + if (w.name === "🔢 Set LoRA Count" || w.type === "button") { + node.widgets.splice(i, 1); + } + } + + let w = node.widgets.find(x => x.name === name); + if (!w) { + const values = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"]; + w = node.addWidget("combo", name, "1", (v) => { + const num = parseInt(v); + if (!isNaN(num)) { + node.properties["visibleLoraCount"] = num; + if (node.cachedLoraCount) { + node.cachedLoraCount.value = num; + } + node.updateLoraSlots(); + } + }, { values }); + } + w.value = node.properties["visibleLoraCount"].toString(); + if (node.cachedLoraCount) { + node.cachedLoraCount.value = node.properties["visibleLoraCount"]; + } + return w; + }; + + const ensureToggleAllWidget = () => { + if (!node.cachedToggleAll) return null; + if (!node.cachedToggleAll.origCallback) { + node.cachedToggleAll.origCallback = node.cachedToggleAll.callback; + } + node.cachedToggleAll.callback = (value) => { + if (node.cachedToggleAll.origCallback) { + node.cachedToggleAll.origCallback(value); + } + const count = parseInt(node.properties["visibleLoraCount"] || 1); + for (let i = 1; i <= count; i++) { + const pair = node.cachedWidgets[i]; + if (pair && pair[0]) { + pair[0].value = value; + } + } + }; + return node.cachedToggleAll; + }; + + // Plain logic restored + node.updateLoraSlots = function () { + if (!cacheReady) initCache(); + + const count = parseInt(this.properties["visibleLoraCount"] || 1); + const controlWidget = ensureControlWidget(); + + this.widgets = [controlWidget]; + + if (node.cachedLoraCount) { + node.cachedLoraCount.type = HIDDEN_TAG; + node.cachedLoraCount.computeSize = () => [0, -4]; + node.cachedLoraCount.value = count; + this.widgets.push(node.cachedLoraCount); + } + + const toggleAllWidget = ensureToggleAllWidget(); + if (toggleAllWidget) { + this.widgets.push(toggleAllWidget); + } + + if (node.cachedCpuOffload) { + this.widgets.push(node.cachedCpuOffload); + } + + for (let i = 1; i <= count; i++) { + const pair = this.cachedWidgets[i]; + if (pair && pair.length >= 3) { + const wEnabled = pair[0]; + const wName = pair[1]; + const wStrength = pair[2]; + + // Restore standard widgets + wEnabled.type = "toggle"; + wName.type = "combo"; + wStrength.type = "number"; + + // Clear custom compute sizes if any + if (wEnabled.computeSize) delete wEnabled.computeSize; + if (wName.computeSize) delete wName.computeSize; + if (wStrength.computeSize) delete wStrength.computeSize; + + this.widgets.push(wEnabled); + this.widgets.push(wName); + this.widgets.push(wStrength); + } + } + + const HEADER_H = 60; + // 3 lines per slot (Toggle ~20 + Name ~26 + Strength ~26 + margins) + const SLOT_H = 80; + const TOGGLE_ALL_H = toggleAllWidget ? 30 : 0; + const CPU_OFFLOAD_H = node.cachedCpuOffload ? 30 : 0; + const PADDING = 20; + const targetH = HEADER_H + TOGGLE_ALL_H + CPU_OFFLOAD_H + (count * SLOT_H) + PADDING; + + this.setSize([this.size[0], targetH]); + + if (app.canvas) app.canvas.setDirty(true, true); + }; + + node.onPropertyChanged = function (property, value) { + if (property === "visibleLoraCount") { + const w = this.widgets.find(x => x.name === "🔢 LoRA Count"); + if (w) w.value = value.toString(); + this.updateLoraSlots(); + } + }; + + const origOnConfigure = node.onConfigure; + node.onConfigure = function () { + if (origOnConfigure) origOnConfigure.apply(this, arguments); + setTimeout(() => node.updateLoraSlots(), 100); + }; + + setTimeout(() => { + initCache(); + node.updateLoraSlots(); + }, 100); + } +}); + diff --git a/md/COMFYUI_0.4.0_MODEL_MANAGEMENT_ERRORS.md b/md/COMFYUI_0.4.0_MODEL_MANAGEMENT_ERRORS.md new file mode 100644 index 0000000..82d70b8 --- /dev/null +++ b/md/COMFYUI_0.4.0_MODEL_MANAGEMENT_ERRORS.md @@ -0,0 +1,49 @@ +# Issue #25: ComfyUI 0.4.0 Model Management Errors + +- **Issue**: After the ComfyUI 0.4.0 update, multiple nodes (including this one in some environments) experienced errors such as `TypeError: 'NoneType' object is not callable` and `AttributeError: 'NoneType' object has no attribute`. In our environment, we resolved these errors by modifying ComfyUI's core `model_management.py`. Note that in our environment, these errors did not occur with this node (ComfyUI-QwenImageLoraLoader). Nunchaku library and ComfyUI-Nunchaku nodes should use the latest versions. If errors persist even after applying the latest version of this node (ComfyUI-QwenImageLoraLoader), modification of ComfyUI's core `model_management.py` may be necessary. +- **Root Cause**: In ComfyUI 0.4.0, ComfyUI's core `model_management.py` lacks sufficient None checks, causing `TypeError` and `AttributeError` when accessing methods or attributes on objects that became `None` after models were unloaded or garbage collected. This problem is not a bug in individual nodes, but rather a structural issue in ComfyUI 0.4.0's model management (`model_management.py`). +- **GC Changes in ComfyUI 0.4.0**: Compared to ComfyUI 0.3.x, automatic model unloading occurs earlier, making the following flow more likely: + +``` +ModelPatcher → GC → weakref(None) +``` + +This also explains why the occurrence of the issue varies by user environment. +- **Technical Basis**: + 1. **Multiple Locations with "Missing None Checks"** - This is not a bug in individual nodes, but the core main logic crashes when accessing attributes on `None`. The added fixes like `if model is None: continue` are defensive code that ComfyUI core should have in all paths. + 2. **Post-Weak-Reference GC Behavior Not Considered** - The introduction of `LoadedModel._model = weakref.ref(ModelPatcher)` in ComfyUI 0.4.0 is a breaking change. When the weak reference target is garbage collected, it returns `None`, but this is not handled. Post-processing for the breaking update is incomplete. + 3. **Multiple Nodes Were Affected in a Chain Reaction** - This is not a problem with nodes, but multiple nodes were affected in a chain reaction due to core behavior changes. Model loading/unloading, memory calculation, GPU/CPU offloading, and `ModelPatcher` lifecycle are all controlled by ComfyUI core. + 4. **All Fix Locations Are Core Responsibility Areas** - The locations fixed (`model_memory`, `model_offloaded_memory`, `load_models_gpu`, `free_memory`, `model_unload`, `is_dead` checks, etc.) are all ComfyUI core functions. These are not areas that node developers should touch. The fact that all fix locations are core logic leaves no explanation other than a core defect. + 5. **Result of Applying Fixes** - After applying None check fixes to ComfyUI's core `model_management.py` in our environment, similar errors were resolved. This demonstrates that the problem can be solved by adding defensive code that the core should have. +- **Model Lifecycle and ModelPatcher Initialization Relationship**: + - **Fact 1: Relationship between LoadedModel and ModelPatcher** - The `LoadedModel` class (lines 502-524 in ComfyUI's `model_management.py`) holds a weak reference to `ModelPatcher`: + +```python +def _set_model(self, model): + self._model = weakref.ref(model) # Weak reference to ModelPatcher + +@property +def model(self): + return self._model() # Returns None when garbage collected +``` + + - **Fact 2: ModelPatcher Initialization** - In the `__init__` of the `ModelPatcher` class (lines 215-237 in `model_patcher.py`), the `pinned` attribute is initialized: + +```python +def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False): + # ... + self.pinned = set() # Line 237: Initialized +``` + + - **Fact 3: Fix Content in ComfyUI Core's model_management.py** - The fix in ComfyUI's core `model_management.py` now skips `LoadedModel` instances where `model` is `None`: In `load_models_gpu()`, skips `LoadedModel` instances where `model` is `None` (lines 712, 727, 743); In `free_memory()`, excludes `LoadedModel` instances where `model` is `None` (line 646). + - **Fact 4: Problem Before Fix** - `LoadedModel` holds a weak reference to `ModelPatcher`. When garbage collected, `LoadedModel.model` returns `None`. Before the fix, methods were called on `LoadedModel` instances where `model` was `None`, causing errors. + - **Fact 5: Behavior After Fix** - By skipping `LoadedModel` instances where `model` is `None`, errors do not occur. Because errors do not occur, processing continues normally. + - **Fact 6: Why copy.deepcopy Fails** - `copy.deepcopy` fails because references to GC'd `ModelPatcher` instances remain in the dictionary being deepcopied. When these references are accessed, they return `None`, causing deepcopy to stop. + - **Fact 7: Confirmation Items** - After applying the fix, `copy.deepcopy` and `pinned` attribute errors do not occur in our environment. Nunchaku library and ComfyUI-Nunchaku nodes should use the latest versions, but this may still be insufficient. While these errors did not occur with this node (ComfyUI-QwenImageLoraLoader) in our environment, the fix to ComfyUI's core `model_management.py` may have indirectly affected it, making errors less likely to occur. +- **Important Note: Not a Problem with Nunchaku Library** - This problem is not caused by the Nunchaku library's implementation. Nunchaku's `model_config` and `ModelPatcher` itself are normal. The problem is in the upstream = ComfyUI core's `model_management.py` GC processing. +- **Speculation (Items That May Be Environment-Dependent)**: The fix allows `ModelPatcher` initialization to complete normally. As a result, the `pinned` attribute is also properly initialized. Accessing `self.pinned` in `__del__` does not cause errors. +- **Recommendations**: + 1. Update Nunchaku library and ComfyUI-Nunchaku nodes to the latest version (addresses `model_config` issues) + 2. Consider applying None check fixes to ComfyUI's core `model_management.py` (may address the root cause) +- **Note**: This is the first support measure. I have published the technical details of the fixes I applied to ComfyUI's core `model_management.py` in my environment. See [COMFYUI_0.4.0_UPDATE_ERROR_FIXES.md](md/COMFYUI_0.4.0_UPDATE_ERROR_FIXES.md) for details. Note that these fixes were applied in my specific environment and may not work universally in all environments. This may also resolve `copy.deepcopy` and `pinned` attribute errors. + diff --git a/md/DIFFSYNTH_OFFICIAL_SUPPORT_EXPLANATION.md b/md/DIFFSYNTH_OFFICIAL_SUPPORT_EXPLANATION.md new file mode 100644 index 0000000..991d922 --- /dev/null +++ b/md/DIFFSYNTH_OFFICIAL_SUPPORT_EXPLANATION.md @@ -0,0 +1,230 @@ +# Diffsynth (ControlNet) Official NextDiT Support - Complete Technical Documentation + +## Reason for Development + +ComfyUI-QwenImageLoraLoader had implemented Diffsynth (ControlNet) support for the unofficial Z-Image loader (`ComfyUI-nunchaku-unofficial-z-image-turbo-loader`), but it did not work with ComfyUI-nunchaku's official Z-Image loader (ComfyUI Lumina2 / NextDiT). + +**Root Causes**: + +1. **Parameter Passing Issues Due to Forward Signature Differences** + - Unofficial: `NunchakuZImageTransformer2DModel.forward(x, t, cap_feats=..., **kwargs)` where `transformer_options` and `control` are received as explicit parameters via monkey-patch + - Official: `NextDiT.forward(x, timesteps, context, num_tokens, attention_mask=None, **kwargs)` where `transformer_options` is read via `**kwargs` (`kwargs.get("transformer_options", {})`) + +2. **transformer_options Transmission Path** + - ComfyUI's ModelPatcher registers Diffsynth patches in `transformer_options["patches"]["double_block"]`. NextDiT reads this `patches["double_block"]` in its internal `_forward()` method and applies patches after each transformer block. However, if `transformer_options` is not passed via `**kwargs`, the patches do not reach NextDiT and Diffsynth does not work. + +3. **Limitations of Existing Implementation** + - `ComfyZImageTurboWrapper._execute_model()` dynamically detected the forward signature using `inspect.signature()`, but when `transformer_options` does **not exist as an explicit parameter** (i.e., when passed via `**kwargs`), it could be missed in the conditional branch. + +4. **Maintaining Compatibility with Unofficial Loader** + - Processing for the unofficial loader (via `customized_forward` or direct forward calls) was already working, so the design needed to add official support without breaking the unofficial loader's behavior. + +## Modified Files + +### `ComfyUI-QwenImageLoraLoader/wrappers/zimageturbo.py` + +**Modification Location: NextDiT Branch in `_execute_model()` Method (lines 442-472)** + +**Problem Before Fix**: + +```python +# Branch 2: ComfyUI Lumina2 / NextDiT signature +if "context" in forward_params and "num_tokens" in forward_params: + # ... (processing of x, context, num_tokens) ... + + forward_kwargs = {} + if "attention_mask" in forward_params: + forward_kwargs["attention_mask"] = attention_mask_value + if "transformer_options" in forward_params: + forward_kwargs["transformer_options"] = transformer_options + if "control" in forward_params: + forward_kwargs["control"] = control + + return self.model( + x_tensor, + timestep, + context=context, + num_tokens=num_tokens_value, + **forward_kwargs, + **kwargs, + ) +``` + +**Problem**: When `transformer_options` does not exist as an explicit parameter (not in `forward_params`), it was not added to `forward_kwargs`, and might not be included in `**kwargs` either. Since NextDiT reads it via `**kwargs.get("transformer_options", {})`, if it is not passed, Diffsynth patches do not reach NextDiT. + +**After Fix**: + +```python +# Branch 2: ComfyUI Lumina2 / NextDiT signature +# forward(x, timesteps, context, num_tokens, attention_mask=None, **kwargs) +if "context" in forward_params and "num_tokens" in forward_params: + # ... (processing of x, context, num_tokens) ... + + # NextDiT forward is: forward(x, timesteps, context, num_tokens, attention_mask=None, **kwargs) + # It consumes `transformer_options` via **kwargs (not a named parameter), + # so we must pass it when the function supports VAR_KEYWORD. + supports_kwargs = any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in forward_sig.parameters.values() + ) + + forward_kwargs = {} + if "attention_mask" in forward_params: + forward_kwargs["attention_mask"] = attention_mask_value + + if supports_kwargs: + # Ensure patches (e.g. double_block from DiffSynth/ControlNet) reach NextDiT + forward_kwargs["transformer_options"] = transformer_options + # Safe to forward; NextDiT ignores it if unused. + if control is not None: + forward_kwargs["control"] = control + else: + if "transformer_options" in forward_params: + forward_kwargs["transformer_options"] = transformer_options + if "control" in forward_params: + forward_kwargs["control"] = control + + # NOTE: NextDiT expects `context` and `num_tokens` as required args + return self.model( + x_tensor, + timestep, + context=context, + num_tokens=num_tokens_value, + **forward_kwargs, + **kwargs, + ) +``` + +**Technical Meaning**: + +1. **VAR_KEYWORD Detection** + - `inspect.Parameter.VAR_KEYWORD` represents `**kwargs` in Python function signatures. `any(p.kind == inspect.Parameter.VAR_KEYWORD for p in forward_sig.parameters.values())` checks whether the forward method supports `**kwargs`. + +2. **When `supports_kwargs` is True** + - When `**kwargs` is supported (like NextDiT), `transformer_options` and `control` are always added to `forward_kwargs`. This ensures they are passed via `**forward_kwargs`. Since NextDiT reads via `kwargs.get("transformer_options", {})`, this allows Diffsynth patches to reach NextDiT. + +3. **When `supports_kwargs` is False (Backward Compatibility)** + - When `**kwargs` is not supported (rare, but for safety), it falls back to the traditional approach of adding only if they exist as explicit parameters. + +4. **Handling of `control`** + - Similarly, when `supports_kwargs` is True, `control` is added only when `control is not None` (safe because NextDiT does not error even if unused). When False, explicit parameter checks are performed. + +## Official NextDiT Implementation Verification + +The `NextDiT` class in `ComfyUI/comfy/ldm/lumina/model.py`: + +```python +def forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwargs): + return comfy.patcher_extension.WrapperExecutor.new_class_executor( + self._forward, + self, + comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, kwargs.get("transformer_options", {})) + ).execute(x, timesteps, context, num_tokens, attention_mask, **kwargs) + +def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, transformer_options={}, **kwargs): + # ... + patches = transformer_options.get("patches", {}) + # ... + transformer_options["total_blocks"] = len(self.layers) + transformer_options["block_type"] = "double" + for i, layer in enumerate(self.layers): + transformer_options["block_index"] = i + img = layer(img, mask, freqs_cis, adaln_input, transformer_options=transformer_options) + if "double_block" in patches: + for p in patches["double_block"]: + out = p({...}) # Apply Diffsynth patch +``` + +**Key Point**: NextDiT's `forward()` receives `**kwargs` and extracts the dictionary via `kwargs.get("transformer_options", {})`. In `_forward()`, it reads `double_block` patches via `patches = transformer_options.get("patches", {})` and applies them after each block. Therefore, **ensuring `transformer_options` is passed via `**kwargs` is a necessary condition for Diffsynth to work**. + +## Differences and Compatibility with Unofficial Loader + +**Unofficial Loader (`ComfyUI-nunchaku-unofficial-z-image-turbo-loader`)**: +- Monkey-patches `NunchakuZImageTransformer2DModel.forward` to add `transformer_options` and `control` as **explicit parameters** +- Signature: `_patched_forward(self, x, t, cap_feats=None, *args, control=None, transformer_options=None, **kwargs)` +- In this case, it is processed in Branch 1 (when `cap_feats` exists) of `ComfyZImageTurboWrapper._execute_model()`, and works correctly with the existing implementation + +**Official Loader (NextDiT)**: +- Signature: `forward(x, timesteps, context, num_tokens, attention_mask=None, **kwargs)` +- `transformer_options` is received via `**kwargs` (not an explicit parameter) +- With this fix, Branch 2 checks `supports_kwargs` and ensures it is passed when `**kwargs` is supported + +**Compatibility Mechanism**: +- Branch 1 (Unofficial): When `cap_feats` exists, processed with existing implementation (checks explicit parameters via `inspect.signature`) +- Branch 2 (Official): When `context` and `num_tokens` exist, checks `VAR_KEYWORD` and ensures passing via `**kwargs` + +This allows **Diffsynth to work with both unofficial and official loaders**, and the existing unofficial users are not affected by the addition of official support. + +## Technical Design Philosophy + +### 1. Dynamic Detection via VAR_KEYWORD + +**Problem**: Need to determine at runtime whether `transformer_options` is an explicit parameter or passed via `**kwargs` + +**Solution**: +- Detects `**kwargs` support using `inspect.Parameter.VAR_KEYWORD`, and when supported, always adds to `forward_kwargs` + +**Benefits**: +- Supports both explicit parameters and `**kwargs` +- Flexible to future signature changes +- No need for type checks or isinstance checks (signature-based detection) + +### 2. Safe Parameter Transmission + +**Policy**: +- When `supports_kwargs` is True: Always add `transformer_options` and `control` (necessary because NextDiT reads from `**kwargs`) +- When `supports_kwargs` is False: Fall back to explicit parameter checks (backward compatibility) + +**Result**: +- Diffsynth works reliably with official NextDiT +- Unofficial loader behavior maintained with existing implementation (no impact) +- Can adapt to future new loaders + +### 3. Maintaining Compatibility with Unofficial Loader + +**Policy**: +- Do not change Branch 1 (unofficial) at all +- Add fixes only to Branch 2 (official) +- Both branches operate independently, so they do not affect each other + +**Result**: +- No impact on existing unofficial users +- Official users can use Diffsynth without additional configuration +- Implementation with minimal code complexity + +## Operation Flow + +### When Using Official NextDiT + Diffsynth + +1. Diffsynth node execution: `NunchakuQwenImageDiffsynthControlnet.diffsynth_controlnet_nunchaku()` registers patches via `set_model_double_block_patch()` +2. ComfyUI ModelPatcher: Registers patches in `transformer_options["patches"]["double_block"]` +3. `ComfyZImageTurboWrapper.forward()`: Receives `transformer_options` and passes it to `_execute_model()` +4. `_execute_model()`: Detects NextDiT's forward signature using `inspect.signature()` +5. Branch 2: Confirms existence of `context` and `num_tokens`, processes as NextDiT format +6. VAR_KEYWORD detection: Checks `**kwargs` support via `supports_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD ...)` +7. Parameter addition: When `supports_kwargs` is True, adds via `forward_kwargs["transformer_options"] = transformer_options` +8. NextDiT.forward() call: `self.model(x_tensor, timestep, context=context, num_tokens=num_tokens_value, **forward_kwargs, **kwargs)` +9. NextDiT._forward(): Extracts `transformer_options` via `kwargs.get("transformer_options", {})` +10. Patch application: Extracts `double_block` patches via `patches = transformer_options.get("patches", {})` and applies after each block + +### When Using Unofficial Loader + Diffsynth + +1. Diffsynth node execution: Similarly registers patches +2. `ComfyZImageTurboWrapper.forward()`: Receives `transformer_options` and passes it to `_execute_model()` +3. `_execute_model()`: Detects unofficial forward signature using `inspect.signature()` +4. Branch 1: Confirms existence of `cap_feats`, processes as unofficial format +5. Explicit parameter check: Checks explicit parameters via `if "transformer_options" in forward_params` +6. Parameter addition: When it exists as an explicit parameter, adds via `forward_kwargs["transformer_options"] = transformer_options` +7. `_patched_forward()` call: Monkey-patched forward receives `transformer_options` as an explicit parameter +8. Patch application: Applies patches after each block via `_apply_double_block_patches()` + +## Conclusion + +The official NextDiT support for Diffsynth (ControlNet) has achieved the following: + +1. **Official Loader Support**: Diffsynth (ControlNet) works with NextDiT format models +2. **Reliable Parameter Transmission**: `VAR_KEYWORD` detection ensures `transformer_options` is reliably passed via `**kwargs` +3. **Compatibility with Unofficial Loader**: Unofficial loader behavior is maintained with existing implementation, no impact +4. **Minimal Code Changes**: Existing unofficial branch unchanged, fixes added only to official branch + +With this design, users can use either the official or unofficial loader, and Diffsynth (ControlNet) will work appropriately with both. + diff --git a/md/V3_DEVELOPMENT_EXPLANATION.md b/md/V3_DEVELOPMENT_EXPLANATION.md new file mode 100644 index 0000000..249bdbe --- /dev/null +++ b/md/V3_DEVELOPMENT_EXPLANATION.md @@ -0,0 +1,311 @@ +# Official Nunchaku Z-Image LoRA Loader V3 - Complete Technical Documentation + +## Reason for Development + +With the release of the official Z-Image loader by ComfyUI-nunchaku, the following issues emerged: + +1. **Model Structure Differences** + - Unofficial loader: `NunchakuZImageTransformer2DModel` (diffusers format, `to_qkv` / `feed_forward.net.0.proj`) + - Official loader: `NextDiT` (ComfyUI Lumina2 format, `attention.qkv` / `feed_forward.w1/w2/w3` or `w13`) + +2. **Forward Signature Differences** + - Unofficial: `forward(x_list: List[torch.Tensor], t_zimage, cap_feats=...)` + - Official: `forward(x: torch.Tensor, timesteps, context, num_tokens, attention_mask=...)` + +3. **LoRA Mapping Mismatch** + - The existing v2 loader was designed for the unofficial structure (`to_qkv` / `net.0.proj`), and failed to resolve keys for the official structure (`qkv` / `w13`) + +4. **Backward Compatibility Maintenance** + - Need to provide both v2 (unofficial-only) and v3 (official-only) simultaneously, allowing users to choose + +## Newly Created Files + +### 1. `ComfyUI-QwenImageLoraLoader/nodes/lora/zimageturbo_v3.py` + +**Purpose**: LoRA stack node dedicated to the official Nunchaku Z-Image loader (`NextDiT`) + +**Main Implementation**: + +The class name is `NunchakuZImageTurboLoraStackV2`, but it is actually dedicated to V3 (registration name changed via GENERATED_NODES) + +**Key Logic (lines 169-191)**: + +```python +elif model_wrapper_type_name == "NextDiT" and model_wrapper_module == "comfy.ldm.lumina.model": + logger.info("🔧 Official loader detected (NextDiT), wrapping with ComfyZImageTurboWrapper") + transformer = model_wrapper + logger.info(f"📦 Creating ComfyZImageTurboWrapper for NextDiT with cpu_offload='{cpu_offload}'") + wrapped_model = ComfyZImageTurboWrapper( + transformer, + getattr(transformer, 'config', {}), + None, + {}, + cpu_offload, + 4.0, + ) + model.model.diffusion_model = wrapped_model + model_wrapper = wrapped_model + transformer = model_wrapper.model +else: + logger.error(f"❌ Model type mismatch! Type: {model_wrapper_type_name}, Module: {model_wrapper_module}") + logger.error("V3 is for official Nunchaku Z-Image DiT Loader only. For unofficial loader, please use V2.") + raise TypeError(f"This LoRA loader (V3) only works with official Nunchaku Z-Image loader, but got {model_wrapper_type_name}.") +``` + +**Technical Meaning**: +- When `NextDiT` is detected, wrap it with `ComfyZImageTurboWrapper` +- By directly replacing `diffusion_model`, maintain compatibility with existing ComfyUI workflows +- Error messages clearly indicate the distinction between v2/v3 usage + +**Registration (lines 226-232)**: + +```python +GENERATED_NODES = { + "NunchakuZImageTurboLoraStackV3": NunchakuZImageTurboLoraStackV2 +} + +GENERATED_DISPLAY_NAMES = { + "NunchakuZImageTurboLoraStackV3": "Nunchaku Z-Image-Turbo LoRA Stack V3" +} +``` + +### 2. `ComfyUI-QwenImageLoraLoader/js/zimageturbo_lora_dynamic_v3.js` + +**Purpose**: Dynamic UI for v3 node (variable display of LoRA slot count) + +**Main Implementation**: +- Uses the same UI logic as v2 (line 8: `name: "nunchaku.zimageturbo_lora_dynamic_v3"` to avoid conflicts with v2) +- Line 17: Filters for v3-only with `if (node.comfyClass !== "NunchakuZImageTurboLoraStackV3")` + +**Technical Meaning**: +- Separates JavaScript extensions for v2 and v3, avoiding name collisions +- UI logic is reused from v2 (no changes needed) + +## Modified Files + +### 1. `ComfyUI-QwenImageLoraLoader/wrappers/zimageturbo.py` + +**Modification 1: NextDiT Support in forward Method (lines 425-458)** + +```python +if "context" in forward_params and "num_tokens" in forward_params: + if isinstance(x, list): + x_stack = torch.stack([t for t in x if t is not None], dim=0) + if x_stack.ndim == 5 and x_stack.shape[2] == 1: + x_stack = x_stack.squeeze(2) + x_tensor = x_stack + else: + x_tensor = x + + if num_tokens_value is None and isinstance(context, torch.Tensor) and context.ndim >= 2: + num_tokens_value = context.shape[1] + + forward_kwargs = {} + if "attention_mask" in forward_params: + forward_kwargs["attention_mask"] = attention_mask_value + if "transformer_options" in forward_params: + forward_kwargs["transformer_options"] = transformer_options + if "control" in forward_params: + forward_kwargs["control"] = control + + return self.model( + x_tensor, + timestep, + context=context, + num_tokens=num_tokens_value, + **forward_kwargs, + **kwargs, + ) +``` + +**Technical Meaning**: +- Dynamically detects forward signature using `inspect.signature` +- When `context` and `num_tokens` exist, processes as NextDiT format +- If `num_tokens` is not specified, automatically derives from `context.shape[1]` +- Passes `attention_mask`, `transformer_options`, and `control` as keyword arguments to maintain compatibility with extension features like ControlNet + +**Modification 2: Extraction of `num_tokens` and `attention_mask` from kwargs (lines 241-242)** + +```python +ref_latents_value = kwargs.pop("ref_latents", None) +num_tokens_value = kwargs.pop("num_tokens", None) +attention_mask_value = kwargs.pop("attention_mask", None) +``` + +**Technical Meaning**: +- ComfyUI passes `num_tokens` and `attention_mask` via kwargs, so they are extracted beforehand +- `kwargs.pop` prevents duplicate parameter passing + +### 2. `ComfyUI-QwenImageLoraLoader/nunchaku_code/lora_qwen.py` + +**Modification 1: NextDiT Mapping Definitions (lines 30-54)** + +```python +ZIMAGE_NEXTDIT_UNPATCHED_KEY_MAPPING = [ + (re.compile(r"^(layers)[._](\d+)[._]attention[._]to[._]([qkv])$"), r"\1.\2.attention.qkv", "qkv", lambda m: m.group(3).upper()), + (re.compile(r"^(layers)[._](\d+)[._]attention[._]to[._]out(?:[._]0)?$"), r"\1.\2.attention.out", "regular", None), + (re.compile(r"^(layers)[._](\d+)[._]feed_forward[._]w1$"), r"\1.\2.feed_forward.w1", "regular", None), + (re.compile(r"^(layers)[._](\d+)[._]feed_forward[._]w2$"), r"\1.\2.feed_forward.w2", "regular", None), + (re.compile(r"^(layers)[._](\d+)[._]feed_forward[._]w3$"), r"\1.\2.feed_forward.w3", "regular", None), +] + +ZIMAGE_NEXTDIT_NUNCHAKU_PATCHED_KEY_MAPPING = [ + (re.compile(r"^(layers)[._](\d+)[._]attention[._]to[._]([qkv])$"), r"\1.\2.attention.qkv", "qkv", lambda m: m.group(3).upper()), + (re.compile(r"^(layers)[._](\d+)[._]attention[._]to[._]out(?:[._]0)?$"), r"\1.\2.attention.out", "regular", None), + (re.compile(r"^(layers)[._](\d+)[._]feed_forward[._](w1|w3)$"), r"\1.\2.feed_forward.w13", "glu", lambda m: m.group(3)), + (re.compile(r"^(layers)[._](\d+)[._]feed_forward[._]w2$"), r"\1.\2.feed_forward.w2", "regular", None), +] +``` + +**Technical Meaning**: +- UNPATCHED: Mapping for vanilla NextDiT (with `w1/w2/w3` separated) +- NUNCHAKU_PATCHED: Mapping for when ComfyUI-nunchaku replaces `feed_forward` with `w13` (w1+w3 fused) +- The `glu` group specification causes `_fuse_glu_lora` to fuse and apply w1 and w3 + +**Modification 2: Model Structure Detection and Automatic Mapping Switching (lines 984-1014)** + +```python +global _ACTIVE_KEY_MAPPING +prev_mapping = _ACTIVE_KEY_MAPPING +try: + nextdit_markers = ( + "layers.0.attention.qkv", + "layers.0.attention.out", + "layers.0.feed_forward.w1", + "layers.0.feed_forward.w2", + "layers.0.feed_forward.w3", + "layers.0.feed_forward.w13", + ) + is_nextdit_style = any(_get_module_by_name(model, p) is not None for p in nextdit_markers) + if is_nextdit_style: + has_w13 = _get_module_by_name(model, "layers.0.feed_forward.w13") is not None + if has_w13: + _ACTIVE_KEY_MAPPING = ZIMAGE_NEXTDIT_NUNCHAKU_PATCHED_KEY_MAPPING + KEY_MAPPING + else: + _ACTIVE_KEY_MAPPING = ZIMAGE_NEXTDIT_UNPATCHED_KEY_MAPPING + KEY_MAPPING + else: + _ACTIVE_KEY_MAPPING = None +except Exception: + _ACTIVE_KEY_MAPPING = None +``` + +**Technical Meaning**: +- Dynamically detects model structure using `_get_module_by_name` +- Determines nunchaku-patched/unpatched by the presence of `layers.0.feed_forward.w13` +- Sets appropriate mapping to `_ACTIVE_KEY_MAPPING`, which is prioritized in `_classify_and_map_key` (line 249) + +**Modification 3: Active Mapping Priority in _classify_and_map_key (line 249)** + +```python +mapping_to_use = _ACTIVE_KEY_MAPPING if _ACTIVE_KEY_MAPPING is not None else KEY_MAPPING +``` + +**Technical Meaning**: +- When `_ACTIVE_KEY_MAPPING` is set (NextDiT detected), it is prioritized +- When not set, uses the traditional `KEY_MAPPING` (for unofficial/v2) +- This allows v2 and v3 to share the same `compose_loras_v2` function while switching only the mapping + +### 3. `ComfyUI-QwenImageLoraLoader/__init__.py` + +**Modification: V3 Node Import and Registration (lines 24, 33-34, 40, 60)** + +```python +from .nodes.lora.zimageturbo_v3 import GENERATED_NODES as ZIMAGETURBO_V3_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V3_NAMES + +for node_class in ZIMAGETURBO_V3_NODES.values(): + node_class.__version__ = __version__ + +NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V3_NODES) + +NODE_DISPLAY_NAME_MAPPINGS = { + **ZIMAGETURBO_V3_NAMES +} +``` + +**Technical Meaning**: +- Registers v3 node into ComfyUI's node system +- Provides both v2 and v3 simultaneously, allowing users to choose + +## Technical Design Philosophy + +### 1. Automatic Mapping Switching + +**Problem**: Module names differ between unofficial (`to_qkv`/`net.0.proj`) and official (`qkv`/`w13`) + +**Solution**: +- Uses a global variable `_ACTIVE_KEY_MAPPING` to switch mappings at runtime +- Detects model structure within `compose_loras_v2` and sets appropriate mapping +- Prioritizes `_ACTIVE_KEY_MAPPING` in `_classify_and_map_key` + +**Benefits**: +- v2/v3 share the same `compose_loras_v2` function +- Minimizes code duplication +- To support new model structures, only mapping definitions need to be added + +### 2. Dynamic Forward Signature Detection + +**Problem**: Forward signatures differ between Z-Image (`cap_feats`) and NextDiT (`context`/`num_tokens`) + +**Solution**: +- Dynamically retrieves forward parameters using `inspect.signature` +- Branches to Z-Image if `cap_feats` exists, or NextDiT if `context` and `num_tokens` exist + +**Benefits**: +- Single `ComfyZImageTurboWrapper` handles both formats +- No need for type checks or isinstance checks (high flexibility) + +### 3. GLU Fusion Handling + +**Problem**: ComfyUI-nunchaku fuses `feed_forward.w1/w3` into `w13` + +**Solution**: +- `ZIMAGE_NEXTDIT_NUNCHAKU_PATCHED_KEY_MAPPING` maps LoRA's `w1`/`w3` to `w13` +- Specifies `glu` as the group, and applies via fusion using `_fuse_glu_lora` + +**Technical Details**: +- `_fuse_glu_lora` is a function that mathematically fuses w1 and w3 LoRAs and applies them to w13 +- This allows LoRA files in `w1`/`w3` format to be correctly applied to patched models (`w13`) + +### 4. Backward Compatibility Maintenance + +**Policy**: +- Provides v2 (unofficial-only) and v3 (official-only) as separate nodes +- v2 code remains completely unchanged, v3 added as new file +- Both use the same `compose_loras_v2`, but work correctly via automatic mapping switching + +**Result**: +- No impact on existing v2 users +- v3 clearly separated as official-loader-only +- Users can choose nodes according to their loader + +## Operation Flow + +### When Using V3 Node (Official Loader) + +1. Node execution: `NunchakuZImageTurboLoraStackV3.load_lora_stack()` +2. Model detection: Checks `model_wrapper_type_name == "NextDiT"` +3. Wrapper creation: Wraps `NextDiT` with `ComfyZImageTurboWrapper` +4. LoRA composition: Calls `compose_loras_v2()` +5. Mapping switch: Detects `layers.0.feed_forward.w13` using `_get_module_by_name()` +6. Mapping application: Sets `ZIMAGE_NEXTDIT_NUNCHAKU_PATCHED_KEY_MAPPING` to `_ACTIVE_KEY_MAPPING` +7. Key resolution: Uses `_ACTIVE_KEY_MAPPING` in `_classify_and_map_key()` +8. LoRA application: Applies `w1`/`w3` → `w13` (GLU fusion), `to_q/k/v` → `qkv` (QKV fusion) + +### During Forward Execution (NextDiT) + +1. Signature detection: Checks `context` and `num_tokens` using `inspect.signature()` +2. Branching: Processes as NextDiT format +3. Argument preparation: Unifies `x` to tensor format, derives `num_tokens` +4. Call: `model.forward(x, timestep, context=context, num_tokens=num_tokens, ...)` + +## Conclusion + +The development of the official-only v3 has achieved the following: + +1. **Official Loader Support**: Enables LoRA application to `NextDiT` format models +2. **Automatic Mapping Switching**: Detects model structure and automatically selects appropriate mapping +3. **Backward Compatibility**: v2 (unofficial) and v3 (official) work simultaneously +4. **Code Reuse**: Extended existing `compose_loras_v2` and `ComfyZImageTurboWrapper`, minimizing duplication + +With this design, users can use either unofficial or official loaders, and simply selecting the appropriate LoRA loader node will make everything work. + diff --git a/md/ZIMAGETURBO_CONTROLNET_FIX.md b/md/ZIMAGETURBO_CONTROLNET_FIX.md new file mode 100644 index 0000000..c05ee7d --- /dev/null +++ b/md/ZIMAGETURBO_CONTROLNET_FIX.md @@ -0,0 +1,294 @@ +# Detailed Explanation of Z-Image-Turbo ControlNet Support Fix in ComfyUI-QwenImageLoraLoader + +## Reason for Additional Fix + +The unofficial Z-Image-Turbo loader (`ComfyUI-nunchaku-unofficial-z-image-turbo-loader`) monkey-patched `NunchakuZImageTransformer2DModel.forward` to accept `transformer_options` and `control` parameters. + +However, the `ComfyZImageTurboWrapper` class in ComfyUI-QwenImageLoraLoader was not correctly passing these parameters to the underlying `NunchakuZImageTransformer2DModel.forward`. + +`ComfyZImageTurboWrapper` is a wrapper that separates LoRA composition and forward pass, passing arguments received in the `forward` method to `NunchakuZImageTransformer2DModel.forward` via the `_execute_model` method. We needed to add `transformer_options` and `control` to this transmission path. + +## Modified File + +`ComfyUI-QwenImageLoraLoader/wrappers/zimageturbo.py` + +## Fix Details + +### Fix 1: forward Method Signature (Lines 76-85) + +#### Before + +```python +def forward( + self, + x, + timestep, + context=None, + y=None, + guidance=None, + **kwargs, +): +``` + +#### After + +```python +def forward( + self, + x, + timestep, + context=None, + y=None, + guidance=None, + control=None, + transformer_options={}, + **kwargs, +): +``` + +#### Explanation + +Added `control` and `transformer_options` to the `forward` method signature to receive them from ComfyUI's ModelPatcher. + +### Fix 2: transformer_options Merge Processing in forward Method (Lines 93-102) + +#### Fix Content + +```python +# Remove guidance, transformer_options, and attention_mask from kwargs +if "guidance" in kwargs: + kwargs.pop("guidance") +if "transformer_options" in kwargs: + if isinstance(transformer_options, dict) and isinstance(kwargs["transformer_options"], dict): + transformer_options = {**transformer_options, **kwargs.pop("transformer_options")} + else: + kwargs.pop("transformer_options") +if "attention_mask" in kwargs: + kwargs.pop("attention_mask") +``` + +#### Explanation + +ComfyUI's ModelPatcher can pass parameters either as explicit arguments or via `**kwargs`. To handle both cases correctly: + +1. If `transformer_options` is passed both as an explicit parameter and in `kwargs`, merge them with the explicit parameter taking precedence (dictionary spread operator `{**transformer_options, **kwargs["transformer_options"]}` ensures explicit parameter values override kwargs values) +2. Remove `guidance` and `attention_mask` from `kwargs` to prevent duplicate parameter errors, as these are handled separately +3. This merging strategy ensures compatibility with different ComfyUI versions that may pass parameters in different ways + +### Fix 3: Argument Passing to _execute_model Method (Lines 219, 221) + +#### Fix Content + +Modified calls to `_execute_model` within the `forward` method to pass `control` and `transformer_options`: + +```python +with cache_context(self._cache_context): + out = self._execute_model(x, timestep, context, y, guidance, control, transformer_options, **kwargs) +else: + out = self._execute_model(x, timestep, context, y, guidance, control, transformer_options, **kwargs) +``` + +#### Explanation + +By passing `control` and `transformer_options` to `_execute_model`, they can be transmitted to the underlying `NunchakuZImageTransformer2DModel.forward`. + +### Fix 4: _execute_model Method Signature and Parameter Cleanup (Line 239-247) + +#### Before + +```python +def _execute_model(self, x, timestep, context, y, guidance, **kwargs): +``` + +#### After + +```python +def _execute_model(self, x, timestep, context, y, guidance, control, transformer_options, **kwargs): + """Helper function to run the Z-Image-Turbo model's forward pass.""" + # Get ref_latents from kwargs before removing it + ref_latents_value = kwargs.pop("ref_latents", None) + + # Remove guidance, transformer_options, and attention_mask from kwargs + kwargs.pop("guidance", None) + kwargs.pop("transformer_options", None) + kwargs.pop("attention_mask", None) + ... +``` + +#### Explanation + +1. Added `control` and `transformer_options` as explicit parameters to the `_execute_model` signature to match the `forward` method's parameter list +2. Extract `ref_latents` from `kwargs` early (line 242) before cleaning `kwargs`, as `ref_latents` is needed later in both `customized_forward` and direct forward paths +3. Remove `guidance`, `transformer_options`, and `attention_mask` from `kwargs` to prevent duplicate argument errors when passing `**kwargs` to downstream methods +4. This cleanup ensures that these parameters are passed explicitly where needed, rather than through `**kwargs`, providing better control and preventing conflicts + +### Fix 5: Transmission When Using customized_forward (Lines 282-300) + +#### Fix Content + +Processing when `customized_forward` is set (when using the custom forward from the unofficial ZIT loader): + +```python +if self.customized_forward: + with torch.inference_mode(): + forward_kwargs_cleaned = {k: v for k, v in self.forward_kwargs.items() if k not in ("guidance", "ref_latents", "transformer_options", "attention_mask")} + transformer_options_cleaned = dict(transformer_options) if transformer_options else {} + transformer_options_cleaned.pop("guidance", None) + transformer_options_cleaned.pop("ref_latents", None) + + return self.customized_forward( + self.model, + hidden_states=x, + encoder_hidden_states=context, + timestep=timestep, + guidance=guidance if self.config.get("guidance_embed", False) else None, + ref_latents=ref_latents_value, + control=control, + transformer_options=transformer_options_cleaned, + **forward_kwargs_cleaned, + **kwargs, + ) +``` + +#### Explanation + +Pass `control` and `transformer_options` when calling `customized_forward`. Remove `guidance` and `ref_latents` from `transformer_options` (to avoid duplication). + +**Additional Details:** + +1. **forward_kwargs cleaning (line 134)**: The `self.forward_kwargs` dictionary may contain parameters that conflict with explicit arguments passed to `customized_forward`. We filter out `guidance`, `ref_latents`, `transformer_options`, and `attention_mask` to prevent duplicate argument errors. + +2. **transformer_options cleaning (lines 135-137)**: Create a copy of `transformer_options` to avoid modifying the original dictionary, then remove `guidance` and `ref_latents` if they exist. This prevents duplicate parameter passing, as `guidance` is passed explicitly based on `self.config.get("guidance_embed", False)`, and `ref_latents` is extracted from `kwargs` earlier and passed as `ref_latents_value`. + +3. **ref_latents handling**: The `ref_latents_value` was extracted from `kwargs` in `_execute_model` (line 242) before cleaning `kwargs`. It is passed explicitly to `customized_forward` rather than through `transformer_options`, providing clearer separation of concerns. + +4. **guidance conditional passing**: `guidance` is only passed to `customized_forward` if `self.config.get("guidance_embed", False)` is `True`, allowing for conditional guidance embedding support. + +### Fix 6: Transmission When Directly Calling forward (Lines 389-409) + +#### Fix Content + +Processing when `customized_forward` is not set (when directly calling `NunchakuZImageTransformer2DModel.forward`): + +```python +# Call Z-Image-Turbo forward with correct signature +# Pass control and transformer_options to allow Model Patcher (double_block patches) to work +# Check if the model's forward method accepts control parameter +import inspect +forward_sig = inspect.signature(self.model.forward) +forward_params = set(forward_sig.parameters.keys()) + +zimage_kwargs_clean = {k: v for k, v in zimage_kwargs.items() if k not in ('control', 'transformer_options')} + +# Build kwargs based on what the forward method accepts +forward_kwargs = zimage_kwargs_clean.copy() +if 'control' in forward_params: + forward_kwargs['control'] = control +if 'transformer_options' in forward_params: + forward_kwargs['transformer_options'] = transformer_options + +model_output = self.model( + x_list, + t_zimage, + cap_feats=cap_feats, + **forward_kwargs +) +``` + +#### Explanation + +- Check the signature of `self.model.forward` using `inspect.signature` +- Add `control` and `transformer_options` to `forward_kwargs` only if they are accepted +- This allows the `_patched_forward` (monkey-patched forward method) from the unofficial ZIT loader to correctly receive `transformer_options` and `control` + +**Additional Details:** + +1. **zimage_kwargs preparation (lines 365-371, not shown in snippet)**: Before the dynamic parameter check, `zimage_kwargs` is prepared with Z-Image-Turbo specific parameters like `patch_size`, `f_patch_size`, and `return_dict=False`. These are extracted from `kwargs` and prepared separately. + +2. **zimage_kwargs_clean (line 171)**: Remove `control` and `transformer_options` from `zimage_kwargs` if they were accidentally included, ensuring they are only passed through the dynamic parameter checking mechanism. + +3. **Dynamic parameter inspection (lines 391-393)**: + - `inspect.signature(self.model.forward)` retrieves the signature of the model's forward method + - `forward_params = set(forward_sig.parameters.keys())` creates a set of parameter names for fast lookup + - This checks the actual signature at runtime, not the static type definition + +4. **Conditional parameter addition (lines 399-402)**: Only add `control` and `transformer_options` to `forward_kwargs` if they are accepted by the forward method. This is crucial because: + - The original `NunchakuZImageTransformer2DModel.forward` does not accept these parameters + - The unofficial ZIT loader monkey-patches the forward method to add these parameters + - This check ensures compatibility whether the monkey-patch is applied or not + +5. **Model call structure (lines 404-409)**: The model is called with positional arguments (`x_list`, `t_zimage`, `cap_feats`) followed by `**forward_kwargs`, which includes `control` and `transformer_options` if accepted. This matches the expected signature pattern for Z-Image-Turbo models. + +6. **Why inspect.signature is necessary**: Without dynamic parameter checking, passing `control` or `transformer_options` to a forward method that doesn't accept them would raise a `TypeError: forward() got an unexpected keyword argument`. Using `inspect.signature` allows the code to adapt to the actual method signature at runtime, supporting both patched and unpatched versions of the model. + +## Key Design Decisions + +### Dynamic Parameter Checking with inspect.signature + +Using `inspect.signature` to dynamically check whether the `forward` method accepts `control` and `transformer_options`. This ensures: + +1. Backward compatibility is maintained (no errors in older versions that don't accept these parameters) +2. Works appropriately regardless of whether the monkey patch from the unofficial ZIT loader is applied +3. Easy to adapt to future extensions + +### Integration with Unofficial ZIT Loader + +This fix enables the following integration flow: + +1. ComfyUI's ModelPatcher passes `transformer_options` and `control` to `ComfyZImageTurboWrapper.forward` +2. `ComfyZImageTurboWrapper.forward` passes these to `_execute_model` +3. `_execute_model` checks with `inspect.signature` before passing to `NunchakuZImageTransformer2DModel.forward` +4. The `_patched_forward` (monkey-patched forward) from the unofficial ZIT loader receives `transformer_options` and `control` and applies the `double_block` patch + +## Summary + +Modified the `ComfyZImageTurboWrapper` class in ComfyUI-QwenImageLoraLoader to correctly pass `control` and `transformer_options` parameters. + +Main changes: + +1. Added `control` and `transformer_options` to the `forward` method signature +2. Added `control` and `transformer_options` to the `_execute_model` method signature +3. Implemented using `inspect.signature` to pass these parameters only if the underlying `forward` method accepts them + +This enables the ControlNet support monkey patch added in the unofficial ZIT loader to work correctly, achieving compatibility between ControlNet and LoRA. + +## Technical Flow Diagram + +The complete parameter transmission flow: + +``` +ComfyUI ModelPatcher + ↓ (passes control, transformer_options) +ComfyZImageTurboWrapper.forward() + ↓ (merges transformer_options if in kwargs, removes from kwargs) + ↓ (passes control, transformer_options as explicit parameters) +ComfyZImageTurboWrapper._execute_model() + ↓ (extracts ref_latents from kwargs, cleans kwargs) + ↓ + ├─→ customized_forward path (if self.customized_forward is set) + │ ↓ (cleans forward_kwargs and transformer_options) + │ ↓ (passes control, transformer_options explicitly) + │ Custom Forward Function (from unofficial ZIT loader) + │ ↓ (receives control, transformer_options) + │ NunchakuZImageTransformer2DModel.forward() (patched) + │ ↓ (applies double_block patches using transformer_options and control) + │ + └─→ Direct forward path (if customized_forward is None) + ↓ (prepares zimage_kwargs, uses inspect.signature) + ↓ (conditionally adds control, transformer_options to forward_kwargs) + NunchakuZImageTransformer2DModel.forward() (patched or unpatched) + ↓ (if patched: receives control, transformer_options, applies double_block patches) + ↓ (if unpatched: ignores extra parameters, works normally) +``` + +## Parameter Cleanup Strategy + +The fix implements a comprehensive parameter cleanup strategy to prevent duplicate argument errors: + +1. **forward method level (lines 93-102)**: Merge `transformer_options` from kwargs with explicit parameter, remove conflicting parameters from kwargs +2. **_execute_model level (lines 241-247)**: Extract `ref_latents` early, remove `guidance`, `transformer_options`, and `attention_mask` from kwargs +3. **customized_forward path (lines 284-287)**: Clean `forward_kwargs` and `transformer_options` to remove duplicates before passing to customized forward +4. **Direct forward path (lines 395-402)**: Use dynamic parameter checking to only pass `control` and `transformer_options` if accepted + +This multi-level cleanup ensures that parameters are passed explicitly where needed and not duplicated through `**kwargs`, preventing runtime errors and ensuring correct parameter precedence. + diff --git a/md/installation.md b/md/installation.md new file mode 100644 index 0000000..0efc6f6 --- /dev/null +++ b/md/installation.md @@ -0,0 +1,61 @@ +# Installation and Features + +## 🎉 MAJOR UPDATE: v1.60 - Simplified Installation (No Integration Required!) + +**As of v1.60, manual integration with ComfyUI-nunchaku's `__init__.py` is completely unnecessary.** + +The node now operates as a fully independent custom node that works out-of-the-box. Simply clone the repository into your `custom_nodes` folder and restart ComfyUI. The nodes will automatically appear in ComfyUI's node menu through ComfyUI's built-in automatic node loading mechanism. + +### What Changed in v1.60 +- ✅ **Removed dependency on ComfyUI-nunchaku integration** - The LoRA loader is now a standalone plugin +- ✅ **Simplified installation** - No batch scripts or manual file editing required +- ✅ **Cleaner architecture** - Node registration happens automatically +- ✅ **Backward compatible** - All existing LoRA files and workflows continue to work + +For a detailed technical explanation of why integration is no longer needed, see [v1.60 Release Notes](https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader/releases/tag/v1.60) + +## Features + +- **NunchakuQwenImageLoraLoader**: Load and apply single LoRA to Qwen Image models +- **NunchakuQwenImageLoraStack**: Apply multiple LoRAs with dynamic UI control +- **NunchakuQwenImageDiffsynthControlnet**: Apply diffsynth ControlNet to Nunchaku Qwen Image models (v2.0) +- **Dynamic VRAM Management**: Automatic CPU offloading based on available VRAM +- **LoRA Composition**: Efficient LoRA stacking and composition +- **ComfyUI Integration**: Seamless integration with ComfyUI workflows + +## Installation + +### Quick Installation (v1.60 - Simplified!) + +**Prerequisites:** +- ComfyUI-nunchaku must be installed + +1. Clone this repository to your ComfyUI custom_nodes directory: + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader.git +``` + +2. Restart ComfyUI + +**That's it!** The nodes will automatically appear in ComfyUI's node menu. + +### Manual Installation (Alternative) + +If you prefer to install manually or are using macOS/Linux: + +```bash +cd /path/to/ComfyUI/custom_nodes +git clone https://github.com/ussoewwin/ComfyUI-QwenImageLoraLoader.git +``` + +Then restart ComfyUI. + +## Requirements + +- Python 3.11+ +- ComfyUI (latest version recommended) +- ComfyUI-nunchaku (required) +- CUDA-capable GPU (optional, but recommended for performance) + diff --git a/md/technical_explanation.md b/md/technical_explanation.md new file mode 100644 index 0000000..a85a3be --- /dev/null +++ b/md/technical_explanation.md @@ -0,0 +1,366 @@ +# Comprehensive Technical Report: Z-Image (Nunchaku) LoRA Integration + +**Version**: 2.0.8 Integration Analysis +**Date**: 2025-12-25 +**Scope**: Enabling full LoRA support for CompVis/Nunchaku-based Z-Image models in ComfyUI. + +--- + +## 1. Context & Objective +The objective was to integrate LoRA support for the **Z-Image** model running on the **Nunchaku** inference backend within `ComfyUI-QwenImageLoraLoader`. This was challenging because Nunchaku optimizations fundamentally change how the model is structured and executed compared to standard `diffusers` pipelines. + +Unlike Qwen/Flux models which were already supported, Z-Image introduced two critical hurdles: +1. **Tensor Shape Mismatch**: The Nunchaku engine enforces strict 5D input requirements `(B, F, C, H, W)` that conflicted with ComfyUI's standard latent formats. +2. **MLP Layer Incompatibility**: The model uses a highly optimized `SwiGLU` activation in its Feed-Forward (MLP) blocks. Standard LoRA injection logic failed to recognize these blocks, resulting in "silent failures" where LoRA weigths were loaded but not applied to the majority of model parameters. + +--- + +## 2. Technical Architecture Analysis + +### 2.1 The Nunchaku Z-Image Block +The core processing unit `ZImageTransformerBlock` deviates from standard implementations. +* **Path**: `nunchaku.models.transformers.transformer_zimage.ZImageTransformerBlock` +* **Structure**: + * **Attention**: Standard Self-Attention (Q, K, V). + * **Feed-Forward (MLP)**: + * Uses `diffusers.models.activations.SwiGLU`. + * Historically, Diffusers models separate `w1` (Gate) and `w3` (Up) linear layers. + * Nunchaku's implementation **fuses** `w1` and `w3` into a single projection layer. + +**Block Structure Details:** + +1. **Attention Block**: + - Standard multi-head self-attention mechanism. + - Q, K, V projections are fused into a single `to_qkv` linear layer for efficiency. + - This fusion is standard in many optimized transformer implementations and doesn't require special LoRA handling. + +2. **Feed-Forward Block Structure**: + - Path in model: `layers[i].feed_forward.net` where `net` is a `nn.Sequential`: + - `net[0]`: `SwiGLU` container module containing `proj: Linear(dim_in, dim_out * 2)` + - `net[1]`: Activation function (usually identity or SiLU) + - `net[2]`: `Linear(dim_out, dim_out * expansion_factor)` (the `w2` layer) + - The `SwiGLU` at `net[0]` replaces the original separate `w1` and `w3` layers. + - Original structure would be: `w1: Linear(dim_in, dim_out)`, `w3: Linear(dim_in, dim_out)`, `w2: Linear(dim_out, dim_out * expansion)`. + - Fused structure: `proj: Linear(dim_in, dim_out * 2)` (inside SwiGLU), `w2: Linear(dim_out, dim_out * expansion)`. + +### 2.2 SwiGLU Mechanics +The `SwiGLU` class acts as a customized Container, not a raw `nn.Linear` layer. +```python +class SwiGLU(nn.Module): + def __init__(self, dim_in, dim_out): + self.proj = nn.Linear(dim_in, dim_out * 2) # Fused w1+w3 + self.activation = nn.SiLU() + + def forward(self, x): + x, gate = self.proj(x).chunk(2, dim=-1) + return x * self.activation(gate) +``` +**Implication**: You cannot apply LoRA to `SwiGLU` directly. You must apply it to `SwiGLU.proj`. Furthermore, because `SwiGLU.proj` outputs `dim_out * 2`, the LoRA weights for `w1` and `w3` must be similarly fused before application. + +**Detailed Technical Analysis:** + +1. **Why Direct LoRA Application Fails**: + - Standard LoRA injection mechanisms search for `nn.Linear` or `nn.Conv2d` modules in the model's module tree. + - `SwiGLU` is a container module (`nn.Module`) that wraps a `nn.Linear` layer, but the injection logic doesn't recursively search inside container modules. + - When the injection logic encounters `SwiGLU`, it sees an unsupported module type and skips it, logging `[MISS] Module found but unsupported type: `. + +2. **The Fused Projection Layer**: + - `SwiGLU.proj` is a `nn.Linear(dim_in, dim_out * 2)` layer that outputs twice the dimension of the intended output. + - In the forward pass, this output is split using `chunk(2, dim=-1)`, creating two tensors of shape `(..., dim_out)`. + - The first half becomes the "value" (`x`), and the second half becomes the "gate" (`gate`). + - The gate is passed through `SiLU` activation, then multiplied element-wise with the value: `x * SiLU(gate)`. + +3. **LoRA Weight Structure**: + - Standard Z-Image LoRA files contain separate weights for `w1` (gate projection) and `w3` (value/up projection). + - Each has its own LoRA matrices: `w1_A`, `w1_B`, `w3_A`, `w3_B` (and optionally `w1_alpha`, `w3_alpha`). + - The base model has these fused into a single `proj` layer, so LoRA weights must be fused to match this structure. + +4. **Mathematical Relationship**: + - If the original model had separate `w1: Linear(dim_in, dim_out)` and `w3: Linear(dim_in, dim_out)`, the fused version is `proj: Linear(dim_in, dim_out * 2)`. + - The fused weight matrix is: `W_fused = [W1; W3]` (vertical concatenation along the output dimension), where `W1` and `W3` are the weight matrices of `w1` and `w3` respectively. + - This means `W_fused.shape == (dim_out * 2, dim_in)`, matching the `proj` layer's weight shape. + - During forward pass: `proj_output = W_fused @ x`, which produces shape `(dim_out * 2, batch_size)`. The `chunk(2, dim=-1)` splits this into two `(dim_out, batch_size)` tensors, the first being the value and the second being the gate. + - For LoRA application, the same fusion principle applies: if separate LoRAs modify `w1` and `w3`, their weight deltas must be fused before application to the fused `proj` layer. + +--- + +## 3. Deep Dive: Debugging & Resolution Steps + +### Phase 1: The Tensor Shape Crisis (`wrappers/zimageturbo.py`) +**Symptom**: `RuntimeError` regarding batch size mismatches during generation. +**Cause**: +ComfyUI often processes latents as a 4D tensor `(Batch * Frames, C, H, W)`. Nunchaku's `forward` pass explicitly expects 5D input to handle temporal data: +```python +# Nunchaku Requirement +x.shape == (batch_size, num_frames, channels, height, width) +``` +**Fix Strategy**: +1. **Batch Synchronization**: We explicitly pass `batch_size=B` to the inference call, ensuring `x` and `cap_feats` align. +2. **Dimensional Gymnastics**: The wrapper now intelligently reshapes usage. + * Input: `[B, C, H, W]` (Single frame inference contexts) + * Transform: `unsqueeze(1)` -> `[B, 1, C, H, W]` + * Inference: Nunchaku processes the 5D tensor. + * Output: `squeeze(1)` -> `[B, C, H, W]` to return to ComfyUI format. + +**Implementation Details (`wrappers/zimageturbo.py`, lines 251-280):** + +The `_execute_model` method handles tensor format conversion and device management: + +1. **List Format Handling (lines 253-266)**: + - If `x` is already a list (from model base conversion), filters out `None` elements. + - Each tensor in the list is moved to the model's device if needed: `img.to(model_device)`. + - Validates that the filtered list is not empty, raising `ValueError` if empty. + +2. **Tensor Format Handling (lines 267-275)**: + - If `x` is still a tensor, moves it to model device first. + - Checks for 5D input: `input_is_5d = x.ndim == 5`. + - If 5D, squeezes dimension 2: `x = x.squeeze(2)` to convert from `(B, F, C, H, W)` format. + +3. **Context Device Management (lines 277-280)**: + - Moves `context` tensor to model device if it's a tensor (not a list). + - Ensures all tensors are on the same device before model execution. + +4. **List Format Conversion for Direct Forward (lines 308-337)**: + - Z-Image-Turbo's forward method expects `x` as a list of tensors with shape `(C, F, H, W)` where F is the frame dimension. + - Conversion process: + - If `x` is already a list, processes each element individually (lines 309-322). + - Filters out `None` elements and validates `img.numel() > 0` to avoid empty tensors (line 313). + - Handles batch dimension: if input is `(1, C, H, W)` (batch size 1), squeezes dimension 0 to get `(C, H, W)` (lines 315-316). + - Ensures frame dimension: if tensor is `(C, H, W)` (3D), unsqueezes dimension 1 to get `(C, 1, H, W)` (lines 318-320). + - Appends each processed tensor to `x_list` for batch processing. + - If `x` is a tensor (not a list), splits it by batch dimension (lines 323-333). + - **Empty list validation** (lines 335-337): Raises `ValueError` if `x_list` is empty, as Z-Image-Turbo cannot process empty batches. + - **Length mismatch detection** (lines 381-386): Checks if `len(x_list) != len(cap_feats)` and warns, as `zip()` will truncate to the shorter length, potentially causing batch size mismatches. + +### Phase 2: The "Weak LoRA" Mystery (MLP Injection Failure) +**Symptom**: LoRA was loaded, but generated images changed very slightly. Users reported "it feels like strength 0.1 even at 1.0". +**Investigation**: +Using detailed `[DEBUG]` logging inspection of the `apply_lora` function in `lora_qwen.py`. + +#### Findings +1. **Success**: Attention layers (`to_q`, `to_k`, `to_v`) were recognized and mapped correctly to `to_qkv`. +2. **Failure**: MLP layers (`w1`, `w2`, `w3`) showed erratic behavior. + * *Attempt 1*: Regex targeted `ff.net`. result: **0 matches**. + * *Correction 1*: Variable name is `feed_forward`, not `ff`. + * *Attempt 2*: Regex targeted `feed_forward.net`. result: **Match found, Injection Failed**. + +**Investigation Process Details:** + +The debugging process involved inspecting the `_classify_and_map_key` function (`nunchaku_code/lora_qwen.py`, lines 170-222) and module resolution: + +1. **Key Mapping Process**: + - LoRA keys from the safetensors file are parsed using regex patterns in `KEY_MAPPING` (line 27). + - The regex `r"^(layers)[._](\d+)[._]feed_forward[._](w1|w3)$"` matches keys like `layers.0.feed_forward.w1` or `layers_0_feed_forward_w3`. + - The template `r"\1.\2.feed_forward.net.0.proj"` maps to `layers.0.feed_forward.net.0.proj`. + - The group type `"glu"` triggers GLU fusion logic in `compose_loras_v2`. + - Before pattern matching, `_classify_and_map_key` (line 170) removes common prefixes like `"transformer."`, `"diffusion_model."`, or `"lora_unet_"` (lines 176-182). + - The function uses `_RE_LORA_SUFFIX` regex (line 120: `r"\.(?Plora(?:[._](?:A|B|down|up)))(?:\.[^.]+)*\.weight$"`) to extract the base key and identify A/B/alpha components (lines 188-195). + - If the suffix matches "lora_A" or ends with ".A" or contains "down", `ab` is set to "A". If it matches "lora_B" or ends with ".B" or contains "up", `ab` is set to "B" (lines 192-195). + +2. **Module Resolution**: + - `_resolve_module_name` (line 254) attempts to find the module using `_get_module_by_name` (line 230). + - The path `layers.0.feed_forward.net.0.proj` is traversed step by step: `layers` -> `0` (index) -> `feed_forward` -> `net` -> `0` (index) -> `proj`. + - `_get_module_by_name` uses a two-step resolution strategy (lines 237-251): + - **First**: Try `hasattr(module, part)` then `getattr(module, part)`. This works for regular attributes and also for numeric keys in `nn.Sequential`/`nn.ModuleDict` (which Python exposes as attributes for keys like "0", "1", "2"). + - **Second**: If `hasattr` fails and `part.isdigit()` is true, check if the module is indexable (`_is_indexable_module` checks for `nn.ModuleList`, `nn.Sequential`, `list`, `tuple`). If so, use integer indexing: `module[int(part)]`. + - The `hasattr`-first approach is necessary because `nn.Sequential` stores modules as attributes with numeric names, so `hasattr(sequential, "0")` returns `True` even though it's not a traditional attribute. `nn.ModuleList` does not expose numeric keys via `hasattr`, requiring the indexing fallback. + - Fallback mappings (lines 269-277): If direct resolution fails, `_resolve_module_name` tries alternative paths: + - `.ff.net.0.proj` -> `.mlp_fc1` (for models using `mlp_fc1` instead of `feed_forward.net.0.proj`) + - `.ff.net.2` -> `.mlp_fc2` + - This allows compatibility with different model architectures that use different naming conventions. + +3. **Failure Point**: + - When the regex matched `feed_forward.net.0` (without `.proj`), the resolved module was `SwiGLU` container. + - `_apply_lora_to_module` (line 873) checks `isinstance(module, nn.Linear)`, which fails for `SwiGLU`. + - The function also checks for `proj_down`/`proj_up` attributes (Nunchaku LoRA-ready modules), which `SwiGLU` also lacks. + - Result: `ValueError: Unsupported module type ` or similar error, logged as `[MISS]`. + +#### The SwiGLU Barrier +The log showed: +```text +[MISS] Module found but unsupported type: +``` +The standard injection logic looks for `nn.Linear`, `nn.Conv2d`, or specific `lora.Linear` layers. It saw `SwiGLU` and skipped it, having no knowledge of how to inject weights into a container. + +### Phase 3: The GLU Fusion Solution (`lora_qwen.py`) + +To solve this, we architected a specific path for "GLU Fusion". + +#### 1. Regex Retargeting +We modified `KEY_MAPPING` to "pierce" the container. +* **Old**: `...feed_forward.net.0` (Targets the container) +* **New**: `...feed_forward.net.0.proj` (Targets the inner Linear layer) + +**Regex Pattern Details (`nunchaku_code/lora_qwen.py`, line 33):** + +```python +(re.compile(r"^(layers)[._](\d+)[._]feed_forward[._](w1|w3)$"), + r"\1.\2.feed_forward.net.0.proj", "glu", lambda m: m.group(3)) +``` + +- **Pattern**: `r"^(layers)[._](\d+)[._]feed_forward[._](w1|w3)$"` + - `^(layers)`: Matches "layers" at start + - `[._]`: Matches either dot or underscore separator + - `(\d+)`: Captures layer number (e.g., "0", "1") + - `[._]feed_forward[._]`: Matches "feed_forward" with any separator + - `(w1|w3)`: Captures either "w1" or "w3" +- **Template**: `r"\1.\2.feed_forward.net.0.proj"` + - `\1`: Layer group name ("layers") + - `\2`: Layer number + - Adds `.feed_forward.net.0.proj` to target the inner Linear layer +- **Group Type**: `"glu"` indicates GLU fusion is needed +- **Component Function**: `lambda m: m.group(3)` extracts "w1" or "w3" for fusion logic + +#### 2. Weight Fusion Logic (`_fuse_glu_lora`) +The Z-Image LoRA files (standard format) provide `w1` (Gate) and `w3` (Up) as separate weight matrices. The Model has them fused. We must replicate this fusion in the LoRA weights *before* injection. + +```python +# Concept +w_fused_A = cat([w1_A, w3_A], dim=0) # Stack ranks? No, usually features. +w_fused_B = cat([w1_B, w3_B], dim=1) # Output dimension doubles +``` +*Note: The actual implementation in `_fuse_glu_lora` ensures correct dimensional concatenation to match `dim_out * 2`.* + +**Implementation Details (`nunchaku_code/lora_qwen.py`, lines 640-695):** + +The `_fuse_glu_lora` function (line 640) performs the actual fusion: + +1. **Input Validation (lines 653-667)**: + - Requires both `w1_A` and `w3_A` in `glu_weights`, returns `(None, None, None)` if missing. + - Validates `A_w1.shape[0] == A_w3.shape[0]` (input features must match), logs warning on mismatch. + +2. **A Matrix Fusion (line 674)**: + - Concatenates along dimension 0: `A_fused = torch.cat([A_w1, A_w3], dim=0)`. + - Result shape: `(r1 + r3, in_features)` where `r1` and `r3` are the ranks of `w1` and `w3` respectively. + +3. **B Matrix Fusion (lines 682-684)**: + - Creates block-diagonal matrix: `B_fused = torch.zeros(out1 + out3, r1 + r3, ...)`. + - Assigns `B_fused[:out1, :r1] = B_w1` and `B_fused[out1:, r1:] = B_w3`. + - Result shape: `(out1 + out3, r1 + r3)` matching the fused `proj` layer's output dimension. + +4. **Integration in compose_loras_v2 (line 1078-1079)**: + - GLU fusion is detected by checking `"w1_A" in lw or "w3_A" in lw` in the grouped LoRA weights. + - The fused weights are then applied via `_apply_lora_to_module` (line 1152) to the `feed_forward.net.0.proj` module (a standard `nn.Linear` layer). + +5. **Weight Aggregation and Scaling (lines 1107-1152)**: + - Multiple LoRAs can target the same module. Weights are aggregated in `aggregated_weights[module_key].append({...})` during the first loop (line 1108). + - In the second loop (lines 1114-1152), for each module, all LoRA parts from different LoRAs are combined: + - **Scaling calculation** (line 1129): `scale = strength * (scale_alpha / max(1.0, float(r_lora)))` where `scale_alpha = alpha.item() if alpha is not None else float(r_lora)`. + - If `alpha` is `None`, it defaults to the rank (`r_lora`), making the scale simply `strength`. If `alpha` is provided, it acts as a normalization factor: higher alpha increases the effect per rank unit. + - **Special handling for AdaNorm layers** (lines 1131-1134): If the module name contains `.norm1.linear` or `.norm1_context.linear`, the B matrix is reordered using `reorder_adanorm_lora_up(B, splits=6)` to match AdaNorm's split structure. Similar handling for single transformer blocks with `.norm.linear` (splits=3). + - **Device and dtype synchronization** (lines 1136-1141): Determines target dtype and device from the module's weight (or `proj_down` for Nunchaku modules), ensuring all tensors match before concatenation. + - `all_A.append(A.to(dtype=target_dtype, device=target_device))` (line 1143) - A matrices are moved to target device/dtype without scaling. + - `all_B_scaled.append((B * scale).to(...))` (line 1144) - B matrices are scaled before concatenation. Scaling is applied here rather than after concatenation to allow per-LoRA strength control. + - **Final concatenation** (lines 1149-1150): `final_A = torch.cat(all_A, dim=0)` concatenates A matrices along the rank dimension, increasing total rank. `final_B = torch.cat(all_B_scaled, dim=1)` concatenates scaled B matrices along the rank dimension, matching the concatenated A shape. + - The concatenated weights are applied to the module via `_apply_lora_to_module` (line 1152). + +6. **Module Application (`_apply_lora_to_module`, lines 873-934)**: + - For standard `nn.Linear` modules (which includes `feed_forward.net.0.proj`): + - Validates tensor dimensions: `A.ndim == 2` and `B.ndim == 2`, `A.shape[1] == module.in_features`, `B.shape[0] == module.out_features` (lines 876-882). + - Backs up original weight to CPU in `_lora_slots` dictionary: `module.weight.detach().cpu().clone()` (line 920). This backup enables LoRA reset functionality. + - Calculates weight delta: `delta = B @ A` where `B.shape == (out_features, rank)` and `A.shape == (rank, in_features)` (line 925). The matrix multiplication produces `delta.shape == (out_features, in_features)`, matching the module weight shape. + - Applies delta in-place: `module.weight.data.add_(delta.to(dtype=module.weight.dtype, device=module.weight.device))` (line 930). The `add_` operation modifies weights directly, avoiding temporary tensor allocation. + - For Nunchaku LoRA-ready modules (with `proj_down`/`proj_up` attributes): + - Unpacks existing low-rank weights using `unpack_lowrank_weight` (lines 887-888). + - Determines base rank and concatenation axis based on weight layout (lines 890-897). + - Concatenates new LoRA A matrices along the rank dimension and B matrices along the output dimension (lines 893, 899). + - Repacks using `pack_lowrank_weight` and updates `module.rank` (lines 901-903). + - Tracks appended rank in `_lora_slots` for reset support (line 908). + +#### 3. Output Projection (`w2`) +The `w2` layer (Down projection) is a standard linear layer appearing at `feed_forward.net.2`. It requires no special fusion, just correct mapping. + +**Mapping Details (`nunchaku_code/lora_qwen.py`, line 35):** + +The `w2` layer uses standard regex mapping without fusion: +```python +(re.compile(r"^(layers)[._](\d+)[._]feed_forward[._]w2$"), + r"\1.\2.feed_forward.net.2", "regular", None) +``` + +- Maps `layers.X.feed_forward.w2` directly to `layers.X.feed_forward.net.2` +- Group type `"regular"` means standard LoRA application (no fusion) +- Component function is `None` (no component extraction needed) +- The resolved module is a standard `nn.Linear`, handled by `_apply_lora_to_module` (lines 911-930) + +--- + +## 4. Final Configuration + +### Supported Layer Mappings +| LoRA Key (Source) | Model Module (Target) | Logic | +| :--- | :--- | :--- | +| `layers.*.attention.to_q/k/v` | `attention.to_qkv` | **QKV Fusion**: Standard concatenation. | +| `layers.*.feed_forward.w1/w3` | `feed_forward.net.0.proj` | **GLU Fusion (NEW)**: Fuses Gate+Up, injects into SwiGLU inner linear. | +| `layers.*.feed_forward.w2` | `feed_forward.net.2` | **Direct**: Standard linear injection. | + +--- + +## 5. Verification Results +**Test Condition**: Z-Image Model + `MistyPokemon_ZIT.safetensors` (Strength 1.0). + +**Logs (v2.0.8)**: +```text +[APPLY] LoRA applied to: layers.0.attention.to_qkv (Success) +[APPLY] LoRA applied to: layers.0.feed_forward.net.0.proj (Success - GLU Fused) +[APPLY] LoRA applied to: layers.0.feed_forward.net.2 (Success - Linear) +... +Applied LoRA compositions to 150 modules. +``` +**Outcome**: +Changes in generated images are now distinct and accurate to the requested LoRA style, confirming that the MLP layers (which contain ~60% of model parameters) are finally participating in the inference. + +**Verification Process:** + +1. **Log Analysis**: + - The `[APPLY]` log messages confirm successful application to `feed_forward.net.0.proj`. + - The total count "Applied LoRA compositions to 150 modules" includes all transformer layers (attention + MLP). + - Each layer has attention (`to_qkv`) and MLP (`net.0.proj`, `net.2`) modules, so 150 modules ≈ 50 layers × 3 modules per layer. + +2. **Before Fix Behavior**: + - Only attention layers were modified, representing ~40% of parameters. + - MLP layers (60% of parameters) were unchanged, causing weak LoRA effects. + - Users needed to set strength > 1.0 to see noticeable changes, and even then results were inaccurate. + +3. **After Fix Behavior**: + - Both attention and MLP layers are modified. + - Strength 1.0 produces accurate LoRA style transfer. + - Image generation shows distinct style changes matching the LoRA training data. + +4. **Performance Impact**: + - LoRA fusion adds minimal computational overhead during composition (one-time operation). + - Block-diagonal B matrix structure maintains efficient matrix multiplication: `B @ A` computation remains O(rank × features). + - Memory overhead: fused LoRA weights have shape `(r1 + r3, in_features)` for A and `(out1 + out3, r1 + r3)` for B, compared to separate `(r1, in)`, `(r3, in)`, `(out1, r1)`, `(out3, r3)`. + - The memory increase is negligible compared to model weights (LoRA rank is typically 4-128, while model dimensions are 1024-4096). + +--- + +## 6. How to Maintain +* **New Architectures**: If Nunchaku adds new model types, check if they use `SwiGLU`. If so, ensure the regex targets `.proj`. +* **Updates to Diffusers**: If `diffusers` changes the `SwiGLU` internal name (unlikely, but possible) from `.proj` to `.linear`, `KEY_MAPPING` must be updated. + +**Maintenance Guidelines:** + +1. **Adding Support for New Model Types**: + - Inspect the model architecture using `model.named_modules()` to find the path to SwiGLU modules. + - Verify the path format (e.g., `layers.X.feed_forward.net.0` contains a `SwiGLU` with `.proj` attribute). + - Add regex pattern to `KEY_MAPPING` if the path format differs from existing patterns. + - Test with a known LoRA file to verify successful application (check logs for `[APPLY]` messages). + +2. **Debugging Failed LoRA Applications**: + - Enable debug logging: `logger.setLevel(logging.DEBUG)`. + - Check `[MISS]` log messages to see which modules were not found or unsupported. + - Use `_resolve_module_name` manually to verify module path resolution. + - Inspect `_classify_and_map_key` output to verify regex pattern matching. + +3. **Key Mapping Updates**: + - If model architecture changes, update regex patterns in `KEY_MAPPING` (line 27). + - Ensure template strings correctly map LoRA keys to model module paths. + - Component functions (lambda) should extract necessary information for fusion logic. + - Group types (`"glu"`, `"qkv"`, `"regular"`) determine which fusion function to call. + +4. **Fusion Logic Maintenance**: + - `_fuse_glu_lora` assumes `w1` (gate) and `w3` (up) weights are present. If model uses different naming, update the detection logic (line 1078). + - Block-diagonal B matrix structure assumes `out1 == out3`. If this changes, fusion logic needs adjustment. + - Alpha handling assumes same alpha for w1/w3. If models require different alphas, fusion logic must be updated. + diff --git a/nodes/lora/qwenimage_v3.py b/nodes/lora/qwenimage_v3.py new file mode 100644 index 0000000..047c51a --- /dev/null +++ b/nodes/lora/qwenimage_v3.py @@ -0,0 +1,266 @@ +""" +This module provides the :class:`NunchakuQwenImageLoraStackV3` node +for applying LoRA weights to Nunchaku Qwen Image models within ComfyUI. +""" + +import copy +import logging +import os +import sys + +# Fix import path for this module +current_dir = os.path.dirname(os.path.abspath(__file__)) +# Go up 2 levels: nodes/lora -> nodes -> ComfyUI-QwenImageLoraLoader +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 + +# Get log level from environment variable (default to INFO) +log_level = os.getenv("LOG_LEVEL", "INFO").upper() + +# Configure logging +logging.basicConfig(level=getattr(logging, log_level, logging.INFO), format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + + +class NunchakuQwenImageLoraStackV2: + """ + Node for loading and applying multiple LoRAs to a Nunchaku Qwen Image model with dynamic UI. + V3 adds per-slot enable toggles and an all-toggle for ComfyUI Nodes 2.0 UI. + """ + + @classmethod + def IS_CHANGED(cls, model, lora_count, cpu_offload="disable", toggle_all=True, **kwargs): + """ + Detect changes to trigger node re-execution. + Returns a hash of relevant parameters to detect changes. + """ + 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()) + # Hash all LoRA parameters + for i in range(1, 11): + m.update(kwargs.get(f"lora_name_{i}", "").encode()) + m.update(str(kwargs.get(f"lora_strength_{i}", 0)).encode()) + m.update(str(kwargs.get(f"enabled_{i}", True)).encode()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(cls, **kwargs): + return True + + @classmethod + def INPUT_TYPES(s): + loras = ["None"] + folder_paths.get_filename_list("loras") + + inputs = { + "required": { + "model": ( + "MODEL", + {"tooltip": "The diffusion model to apply LoRAs to."}, + ), + "lora_count": ( + "INT", + { + "default": 3, + "min": 1, + "max": 10, + "step": 1, + "tooltip": "Number of LoRA slots to process.", + }, + ), + "cpu_offload": ( + ["auto", "enable", "disable"], + { + "default": "disable", + "tooltip": "CPU offload setting. 'auto' enables offload when VRAM is low, 'enable' forces offload, 'disable' disables offload.", + }, + ), + "toggle_all": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Enable/disable all LoRAs at once.", + }, + ), + }, + "optional": {}, + } + + # Add all LoRA inputs (up to 10 slots) as optional + for i in range(1, 11): + inputs["optional"][f"enabled_{i}"] = ( + "BOOLEAN", + { + "default": True, + "tooltip": f"Enable/disable LoRA {i}.", + }, + ) + inputs["optional"][f"lora_name_{i}"] = ( + loras, + {"tooltip": f"The file name of LoRA {i}. Select 'None' to skip this slot."}, + ) + inputs["optional"][f"lora_strength_{i}"] = ( + "FLOAT", + { + "default": 1.0, + "min": -100.0, + "max": 100.0, + "step": 0.01, + "tooltip": f"Strength for LoRA {i}.", + }, + ) + + return inputs + + RETURN_TYPES = ("MODEL",) + OUTPUT_TOOLTIPS = ("The modified diffusion model with all LoRAs applied.",) + FUNCTION = "load_lora_stack" + TITLE = "Nunchaku Qwen Image LoRA Stack V3" + 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): + loras_to_apply = [] + + # Log toggle_all state + logger.info(f"[LoRA Stack Status] toggle_all: {toggle_all}") + logger.info(f"[LoRA Stack Status] Processing {lora_count} LoRA slot(s):") + + # Process only the number of LoRAs specified by lora_count + for i in range(1, lora_count + 1): + lora_name = kwargs.get(f"lora_name_{i}") + lora_strength = kwargs.get(f"lora_strength_{i}", 1.0) + enabled_individual = kwargs.get(f"enabled_{i}", True) + enabled = toggle_all and enabled_individual + + status_parts = [] + status_parts.append(f"Slot {i}:") + if lora_name and lora_name != "None": + status_parts.append(f"'{lora_name}'") + status_parts.append(f"strength={lora_strength}") + else: + status_parts.append("(no LoRA selected)") + + status_parts.append(f"toggle_all={toggle_all}") + status_parts.append(f"enabled_{i}={enabled_individual}") + status_parts.append(f"final_enabled={enabled}") + + if enabled and lora_name and lora_name != "None" and abs(lora_strength) > 1e-5: + status_parts.append("→ APPLIED ✓") + loras_to_apply.append((lora_name, lora_strength)) + else: + status_parts.append("→ SKIPPED ✗") + + logger.info(f"[LoRA Stack Status] {' | '.join(status_parts)}") + + logger.info(f"[LoRA Stack Status] Summary: {len(loras_to_apply)} LoRA(s) will be applied out of {lora_count} slot(s)") + + if not loras_to_apply: + return (model,) + + model_wrapper = model.model.diffusion_model + + # Dynamic import with explicit path manipulation + import importlib.util + + lora_loader_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + if lora_loader_dir not in sys.path: + sys.path.insert(0, lora_loader_dir) + + spec = importlib.util.spec_from_file_location( + "wrappers.qwenimage", + os.path.join(lora_loader_dir, "wrappers", "qwenimage.py"), + ) + wrappers_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(wrappers_module) + ComfyQwenImageWrapper = wrappers_module.ComfyQwenImageWrapper + + from nunchaku import NunchakuQwenImageTransformer2DModel + + # Debug logging + model_wrapper_type_name = type(model_wrapper).__name__ + model_wrapper_module = type(model_wrapper).__module__ + logger.info(f"🔍 Model wrapper type: '{model_wrapper_type_name}'") + logger.info(f"🔍 Model wrapper module: {model_wrapper_module}") + logger.info(f"🔍 Type repr: {repr(type(model_wrapper))}") + logger.info(f"🔍 Has 'model' attr? {hasattr(model_wrapper, 'model')}") + logger.info(f"🔍 Has 'loras' attr? {hasattr(model_wrapper, 'loras')}") + + # Check if it's already wrapped + if hasattr(model_wrapper, "model") and hasattr(model_wrapper, "loras"): + logger.info("✅ Model is already wrapped (detected via attributes)") + logger.info(f"📦 Current CPU offload setting: '{model_wrapper.cpu_offload_setting}'") + if model_wrapper.cpu_offload_setting != cpu_offload: + logger.info(f"🔄 Updating CPU offload setting from '{model_wrapper.cpu_offload_setting}' to '{cpu_offload}'") + model_wrapper.cpu_offload_setting = cpu_offload + transformer = model_wrapper.model + elif model_wrapper_type_name == "NunchakuQwenImageTransformer2DModel" or model_wrapper_type_name.endswith( + "NunchakuQwenImageTransformer2DModel" + ): + logger.info("🔧 Wrapping NunchakuQwenImageTransformer2DModel with ComfyQwenImageWrapper") + logger.info(f"📦 Creating ComfyQwenImageWrapper with cpu_offload='{cpu_offload}'") + wrapped_model = ComfyQwenImageWrapper( + model_wrapper, + getattr(model_wrapper, "config", {}), + None, # customized_forward + {}, # forward_kwargs + cpu_offload, # cpu_offload_setting + 4.0, # vram_margin_gb + ) + model.model.diffusion_model = wrapped_model + model_wrapper = wrapped_model + transformer = model_wrapper.model + else: + logger.error(f"❌ Model type mismatch! Type: {model_wrapper_type_name}, Module: {model_wrapper_module}") + logger.error("Please use 'Nunchaku Qwen Image DiT Loader'.") + raise TypeError(f"This LoRA loader only works with Nunchaku Qwen Image models, but got {model_wrapper_type_name}.") + + # Flux-style deepcopy + saved_config = None + if hasattr(model, "model") and hasattr(model.model, "model_config"): + saved_config = model.model.model_config + model.model.model_config = None + + model_wrapper.model = None + try: + ret_model = copy.deepcopy(model) + finally: + if saved_config is not None: + model.model.model_config = saved_config + model_wrapper.model = transformer + + ret_model_wrapper = ret_model.model.diffusion_model + if saved_config is not None: + ret_model.model.model_config = saved_config + ret_model_wrapper.model = transformer + + ret_model_wrapper.loras = model_wrapper.loras.copy() + + 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})") + + logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}") + return (ret_model,) + + +GENERATED_NODES = { + "NunchakuQwenImageLoraStackV3": NunchakuQwenImageLoraStackV2, +} + +GENERATED_DISPLAY_NAMES = { + "NunchakuQwenImageLoraStackV3": "Nunchaku Qwen Image LoRA Stack V3", +} + + diff --git a/nodes/lora/zimageturbo_v3.py b/nodes/lora/zimageturbo_v3.py new file mode 100644 index 0000000..be61fd7 --- /dev/null +++ b/nodes/lora/zimageturbo_v3.py @@ -0,0 +1,275 @@ +""" +This module provides the :class:`NunchakuZImageTurboLoraStackV2` node +for applying LoRA weights to Nunchaku Z-Image-Turbo models within ComfyUI. +""" + +import copy +import logging +import os +import sys + +# Fix import path for this module +current_dir = os.path.dirname(os.path.abspath(__file__)) +# Go up 2 levels: nodes/lora -> nodes -> ComfyUI-QwenImageLoraLoader +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] zimageturbo.py exists: {os.path.exists(os.path.join(lora_loader_dir, 'wrappers', 'zimageturbo.py'))}") + +import folder_paths + +# Get log level from environment variable (default to INFO) +log_level = os.getenv("LOG_LEVEL", "INFO").upper() + +# Configure logging +logging.basicConfig(level=getattr(logging, log_level, logging.INFO), format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + + +class NunchakuZImageTurboLoraStackV2: + """ + Node for loading and applying multiple LoRAs to a Nunchaku Z-Image-Turbo model with dynamic UI. + V3 is for official Nunchaku Z-Image loader only. For unofficial loader, use V2. + """ + @classmethod + def IS_CHANGED(cls, model, lora_count, cpu_offload="disable", toggle_all=True, **kwargs): + """ + Detect changes to trigger node re-execution. + Returns a hash of relevant parameters to detect changes. + """ + 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()) + # Hash all LoRA parameters + for i in range(1, 11): + m.update(kwargs.get(f"lora_name_{i}", "").encode()) + m.update(str(kwargs.get(f"lora_strength_{i}", 0)).encode()) + m.update(str(kwargs.get(f"enabled_{i}", True)).encode()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(cls, **kwargs): + return True + + @classmethod + def INPUT_TYPES(s): + loras = ["None"] + folder_paths.get_filename_list("loras") + + inputs = { + "required": { + "model": ( + "MODEL", + {"tooltip": "The diffusion model to apply LoRAs to."}, + ), + "lora_count": ( + "INT", + { + "default": 3, + "min": 1, + "max": 10, + "step": 1, + "tooltip": "Number of LoRA slots to process.", + }, + ), + "cpu_offload": ( + ["auto", "enable", "disable"], + { + "default": "disable", + "tooltip": "CPU offload setting. 'auto' enables offload when VRAM is low, 'enable' forces offload, 'disable' disables offload.", + }, + ), + "toggle_all": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Enable/disable all LoRAs at once.", + }, + ), + }, + "optional": {}, + } + + # Add all LoRA inputs (up to 10 slots) as optional + for i in range(1, 11): + inputs["optional"][f"enabled_{i}"] = ( + "BOOLEAN", + { + "default": True, + "tooltip": f"Enable/disable LoRA {i}.", + }, + ) + inputs["optional"][f"lora_name_{i}"] = ( + loras, + {"tooltip": f"The file name of LoRA {i}. Select 'None' to skip this slot."}, + ) + inputs["optional"][f"lora_strength_{i}"] = ( + "FLOAT", + { + "default": 1.0, + "min": -100.0, + "max": 100.0, + "step": 0.01, + "tooltip": f"Strength for LoRA {i}.", + }, + ) + + return inputs + + RETURN_TYPES = ("MODEL",) + OUTPUT_TOOLTIPS = ("The modified diffusion model with all LoRAs applied.",) + FUNCTION = "load_lora_stack" + TITLE = "Nunchaku Z-Image-Turbo LoRA Stack V3" + CATEGORY = "Nunchaku" + DESCRIPTION = "Apply multiple LoRAs to a diffusion model in a single node with dynamic UI control. V3 is for official Nunchaku Z-Image loader only. For unofficial loader, use V2." + + def load_lora_stack(self, model, lora_count, cpu_offload="disable", toggle_all=True, **kwargs): + loras_to_apply = [] + + # Log toggle_all state + logger.info(f"[LoRA Stack Status] toggle_all: {toggle_all}") + logger.info(f"[LoRA Stack Status] Processing {lora_count} LoRA slot(s):") + + # Process only the number of LoRAs specified by lora_count + for i in range(1, lora_count + 1): + lora_name = kwargs.get(f"lora_name_{i}") + lora_strength = kwargs.get(f"lora_strength_{i}", 1.0) + enabled_individual = kwargs.get(f"enabled_{i}", True) + # Check if this LoRA is enabled (considering both toggle_all and individual enabled_) + enabled = toggle_all and enabled_individual + + # Log each LoRA slot status + status_parts = [] + status_parts.append(f"Slot {i}:") + if lora_name and lora_name != "None": + status_parts.append(f"'{lora_name}'") + status_parts.append(f"strength={lora_strength}") + else: + status_parts.append("(no LoRA selected)") + + status_parts.append(f"toggle_all={toggle_all}") + status_parts.append(f"enabled_{i}={enabled_individual}") + status_parts.append(f"final_enabled={enabled}") + + if enabled and lora_name and lora_name != "None" and abs(lora_strength) > 1e-5: + status_parts.append("→ APPLIED ✓") + loras_to_apply.append((lora_name, lora_strength)) + else: + status_parts.append("→ SKIPPED ✗") + + logger.info(f"[LoRA Stack Status] {' | '.join(status_parts)}") + + # Log summary + logger.info(f"[LoRA Stack Status] Summary: {len(loras_to_apply)} LoRA(s) will be applied out of {lora_count} slot(s)") + + if not loras_to_apply: + return (model,) + + model_wrapper = model.model.diffusion_model + + # Dynamic import with explicit path manipulation + import sys + import importlib.util + lora_loader_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + if lora_loader_dir not in sys.path: + sys.path.insert(0, lora_loader_dir) + + spec = importlib.util.spec_from_file_location( + "wrappers.zimageturbo", + os.path.join(lora_loader_dir, "wrappers", "zimageturbo.py") + ) + wrappers_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(wrappers_module) + ComfyZImageTurboWrapper = wrappers_module.ComfyZImageTurboWrapper + + from nunchaku import NunchakuZImageTransformer2DModel + + # Debug logging + model_wrapper_type_name = type(model_wrapper).__name__ + model_wrapper_module = type(model_wrapper).__module__ + logger.info(f"🔍 Model wrapper type: '{model_wrapper_type_name}'") + logger.info(f"🔍 Model wrapper module: {model_wrapper_module}") + logger.info(f"🔍 Type repr: {repr(type(model_wrapper))}") + logger.info(f"🔍 Has 'model' attr? {hasattr(model_wrapper, 'model')}") + logger.info(f"🔍 Has 'loras' attr? {hasattr(model_wrapper, 'loras')}") + + # Check if it's already wrapped + # ComfyZImageTurboWrapper has a 'model' attribute that contains the transformer + # and a 'loras' list attribute + if hasattr(model_wrapper, 'model') and hasattr(model_wrapper, 'loras'): + # Already wrapped, proceed normally + logger.info("✅ Model is already wrapped (detected via attributes)") + logger.info(f"📦 Current CPU offload setting: '{model_wrapper.cpu_offload_setting}'") + # Update CPU offload setting if different + if model_wrapper.cpu_offload_setting != cpu_offload: + logger.info(f"🔄 Updating CPU offload setting from '{model_wrapper.cpu_offload_setting}' to '{cpu_offload}'") + model_wrapper.cpu_offload_setting = cpu_offload + transformer = model_wrapper.model + elif model_wrapper_type_name == "NextDiT" and model_wrapper_module == "comfy.ldm.lumina.model": + # Official Nunchaku loader uses ComfyUI Lumina2 / NextDiT signature. + # Our ComfyZImageTurboWrapper now supports both NextDiT and Z-Image signatures. + logger.info("🔧 Official loader detected (NextDiT), wrapping with ComfyZImageTurboWrapper") + transformer = model_wrapper + logger.info(f"📦 Creating ComfyZImageTurboWrapper for NextDiT with cpu_offload='{cpu_offload}'") + wrapped_model = ComfyZImageTurboWrapper( + transformer, + getattr(transformer, 'config', {}), + None, # customized_forward + {}, # forward_kwargs + cpu_offload, # cpu_offload_setting + 4.0, # vram_margin_gb + ) + + # Replace the model's diffusion_model with our wrapper + model.model.diffusion_model = wrapped_model + model_wrapper = wrapped_model + transformer = model_wrapper.model + else: + logger.error(f"❌ Model type mismatch! Type: {model_wrapper_type_name}, Module: {model_wrapper_module}") + logger.error("V3 is for official Nunchaku Z-Image DiT Loader only. For unofficial loader, please use V2.") + raise TypeError(f"This LoRA loader (V3) only works with official Nunchaku Z-Image loader, but got {model_wrapper_type_name}.") + + # Flux-style deepcopy + # Save config before deepcopy to avoid __setstate__ errors + saved_config = None + if hasattr(model, 'model') and hasattr(model.model, 'model_config'): + saved_config = model.model.model_config + model.model.model_config = None + + model_wrapper.model = None + try: + ret_model = copy.deepcopy(model) + finally: + # Restore config and model + if saved_config is not None: + model.model.model_config = saved_config + model_wrapper.model = transformer + + ret_model_wrapper = ret_model.model.diffusion_model + # Restore config in copied model if it was saved + if saved_config is not None: + ret_model.model.model_config = saved_config + ret_model_wrapper.model = transformer + + ret_model_wrapper.loras = model_wrapper.loras.copy() + + 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})") + + logger.info(f"Total LoRAs in stack: {len(ret_model_wrapper.loras)}") + + return (ret_model,) + +GENERATED_NODES = { + "NunchakuZImageTurboLoraStackV3": NunchakuZImageTurboLoraStackV2 +} + +GENERATED_DISPLAY_NAMES = { + "NunchakuZImageTurboLoraStackV3": "Nunchaku Z-Image-Turbo LoRA Stack V3" +}