Restore deleted files after force push: images, js, md, nodes, and other missing files
This commit is contained in:
Your Name
+5
-2
@@ -67,5 +67,8 @@ venv.bak/
|
||||
Thumbs.db
|
||||
|
||||
# User-specific documentation files
|
||||
technical_explanation.md
|
||||
PR34_COMFYUI_V0.6.0_COMPATIBILITY.md
|
||||
# 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
|
||||
@@ -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).
|
||||
|
||||
<img src="images/qilorastacker.png" alt="Dynamic UI Control" width="400">
|
||||
<img src="images/zitlorav3.png" alt="Dynamic UI Control V3" width="400">
|
||||
|
||||
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
|
||||
|
||||
+19
-5
@@ -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"
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 373 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 484 KiB |
@@ -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);
|
||||
}
|
||||
});
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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: <class 'diffusers.models.activations.SwiGLU'>`.
|
||||
|
||||
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"\.(?P<tag>lora(?:[._](?: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 <class 'diffusers.models.activations.SwiGLU'>` or similar error, logged as `[MISS]`.
|
||||
|
||||
#### The SwiGLU Barrier
|
||||
The log showed:
|
||||
```text
|
||||
[MISS] Module found but unsupported type: <class 'diffusers.models.activations.SwiGLU'>
|
||||
```
|
||||
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.
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
@@ -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_<i>)
|
||||
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"
|
||||
}
|
||||
Reference in New Issue
Block a user