diff --git a/__init__.py b/__init__.py index 3ba3a16..327d0eb 100644 --- a/__init__.py +++ b/__init__.py @@ -47,8 +47,6 @@ try: from .nodes.lora.qwenimage_v3 import GENERATED_NODES as QWEN_V3_NODES, GENERATED_DISPLAY_NAMES as QWEN_V3_NAMES from .nodes.lora.qwenimage_v1 import GENERATED_NODES as QWEN_V1_NODES, GENERATED_DISPLAY_NAMES as QWEN_V1_NAMES from .nodes.lora.zimageturbo_v1 import GENERATED_NODES as ZIMAGETURBO_V1_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V1_NAMES - # Z-Image-Turbo V2 registration removed - # from .nodes.lora.zimageturbo_v2 import GENERATED_NODES as ZIMAGETURBO_V2_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V2_NAMES # Z-Image-Turbo V3 is deprecated - removed from registration # from .nodes.lora.zimageturbo_v3 import GENERATED_NODES as ZIMAGETURBO_V3_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V3_NAMES from .nodes.lora.zimageturbo_v4 import GENERATED_NODES as ZIMAGETURBO_V4_NODES, GENERATED_DISPLAY_NAMES as ZIMAGETURBO_V4_NAMES @@ -65,9 +63,6 @@ try: node_class.__version__ = __version__ for node_class in ZIMAGETURBO_V1_NODES.values(): node_class.__version__ = __version__ - # Z-Image-Turbo V2 registration removed - # for node_class in ZIMAGETURBO_V2_NODES.values(): - # node_class.__version__ = __version__ # Z-Image-Turbo V3 registration removed # for node_class in ZIMAGETURBO_V3_NODES.values(): # node_class.__version__ = __version__ @@ -81,8 +76,6 @@ try: NODE_CLASS_MAPPINGS.update(QWEN_V3_NODES) NODE_CLASS_MAPPINGS.update(QWEN_V1_NODES) NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V1_NODES) - # Z-Image-Turbo V2 registration removed - # NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V2_NODES) # Z-Image-Turbo V3 registration removed # NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V3_NODES) NODE_CLASS_MAPPINGS.update(ZIMAGETURBO_V4_NODES) @@ -108,8 +101,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { **QWEN_V3_NAMES, **QWEN_V1_NAMES, **ZIMAGETURBO_V1_NAMES, - # Z-Image-Turbo V2 registration removed - # **ZIMAGETURBO_V2_NAMES, # Z-Image-Turbo V3 registration removed # **ZIMAGETURBO_V3_NAMES, **ZIMAGETURBO_V4_NAMES diff --git a/js/zimageturbo_lora_dynamic.js b/js/zimageturbo_lora_dynamic.js deleted file mode 100644 index 597bc09..0000000 --- a/js/zimageturbo_lora_dynamic.js +++ /dev/null @@ -1,175 +0,0 @@ -import { app } from "../../scripts/app.js"; - -console.log("★★★ zimageturbo_lora_dynamic.js: Z-Image-Turbo LoRA Stack V2 ★★★"); - -const HIDDEN_TAG = "tschide"; - -app.registerExtension({ - name: "nunchaku.zimageturbo_lora_dynamic_combo_final_restore", - - async beforeRegisterNodeDef(nodeType, nodeData, app) { - if (nodeData.name === "NunchakuZImageTurboLoraStackV2") { - nodeType["@visibleLoraCount"] = { type: "number", default: 1, min: 1, max: 10, step: 1 }; - } - }, - - nodeCreated(node) { - if (node.comfyClass !== "NunchakuZImageTurboLoraStackV2") 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; - } - - for (let i = 1; i <= 10; i++) { - const wName = all.find(w => w.name === `lora_name_${i}`); - const wStrength = all.find(w => w.name === `lora_strength_${i}`); - if (wName && wStrength) { - node.cachedWidgets[i] = [wName, wStrength]; - wName.type = "combo"; - wStrength.type = "number"; - 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; - // Sync with lora_count widget for Python backend - if (node.cachedLoraCount) { - node.cachedLoraCount.value = num; - } - node.updateLoraSlots(); - } - }, { values }); - } - w.value = node.properties["visibleLoraCount"].toString(); - // Sync lora_count widget value - if (node.cachedLoraCount) { - node.cachedLoraCount.value = node.properties["visibleLoraCount"]; - } - return w; - }; - - node.updateLoraSlots = function() { - if (!cacheReady) initCache(); - - const count = parseInt(this.properties["visibleLoraCount"] || 1); - const controlWidget = ensureControlWidget(); - - // Physical widget reconstruction for clean layout (like Flux V2) - this.widgets = [controlWidget]; - - // Add lora_count widget (required for Python backend, but hidden in UI using HIDDEN_TAG) - if (node.cachedLoraCount) { - // Ensure it's hidden - node.cachedLoraCount.type = HIDDEN_TAG; - node.cachedLoraCount.computeSize = () => [0, -4]; - // Sync value for Python backend - node.cachedLoraCount.value = count; - // Add to widgets array (for Python backend, but hidden in UI) - this.widgets.push(node.cachedLoraCount); - } - - // Add cpu_offload widget from cache (required for Python backend) - if (node.cachedCpuOffload) { - this.widgets.push(node.cachedCpuOffload); - } - - // Add only visible LoRA slots (non-visible widgets are removed from array) - for (let i = 1; i <= count; i++) { - const pair = this.cachedWidgets[i]; - if (pair) { - this.widgets.push(pair[0]); - this.widgets.push(pair[1]); - } - } - - // Height calculation - const HEADER_H = 60; - const SLOT_H = 54; - const CPU_OFFLOAD_H = node.cachedCpuOffload ? 40 : 0; - const PADDING = 20; - const targetH = HEADER_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(); - } - }; - - // Restore UI on configure - const origOnConfigure = node.onConfigure; - node.onConfigure = function() { - if (origOnConfigure) origOnConfigure.apply(this, arguments); - setTimeout(() => node.updateLoraSlots(), 100); - }; - - setTimeout(() => { - initCache(); - node.updateLoraSlots(); - }, 100); - } -}); - diff --git a/js/zimageturbo_lora_dynamic_v3.js b/js/zimageturbo_lora_dynamic_v3.js index 9718c37..59d7d96 100644 --- a/js/zimageturbo_lora_dynamic_v3.js +++ b/js/zimageturbo_lora_dynamic_v3.js @@ -1,6 +1,6 @@ import { app } from "../../scripts/app.js"; -console.log("★★★ zimageturbo_lora_dynamic.js: Z-Image-Turbo LoRA Stack V3 ★★★"); +console.log("★★★ zimageturbo_lora_dynamic_v3.js: Z-Image-Turbo LoRA Stack V3 ★★★"); const HIDDEN_TAG = "tschide"; diff --git a/nodes/lora/zimageturbo_v2.py b/nodes/lora/zimageturbo_v2.py deleted file mode 100644 index f8a7c66..0000000 --- a/nodes/lora/zimageturbo_v2.py +++ /dev/null @@ -1,247 +0,0 @@ -""" -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) - -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. - """ - @classmethod - def IS_CHANGED(cls, model, lora_count, cpu_offload="disable", **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()) - # 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()) - 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.", - }, - ), - }, - "optional": {}, - } - - # Add all LoRA inputs (up to 10 slots) as optional - for i in range(1, 11): - 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 V2" - CATEGORY = "Nunchaku" - DESCRIPTION = "Apply multiple LoRAs to a diffusion model in a single node with dynamic UI control. v1.0.3" - - def load_lora_stack(self, model, lora_count, cpu_offload="disable", **kwargs): - loras_to_apply = [] - - # 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) - - # Skip if lora_name is None or strength is negligible - if lora_name and lora_name != "None" and abs(lora_strength) > 1e-5: - loras_to_apply.append((lora_name, lora_strength)) - - 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 - # Directly check __dict__ and _modules to avoid nn.Module's __getattr__ issues - # nn.Module stores child modules (nn.Module) in _modules and regular attributes in __dict__ - has_model = 'model' in model_wrapper.__dict__ or 'model' in getattr(model_wrapper, '_modules', {}) - has_loras = 'loras' in model_wrapper.__dict__ - - def get_model_attr(wrapper): - """Get model attribute from _modules or __dict__ of nn.Module""" - if 'model' in wrapper.__dict__: - return wrapper.__dict__['model'] - if 'model' in getattr(wrapper, '_modules', {}): - return wrapper._modules['model'] - return None - - if has_model and has_loras: - # Already wrapped, proceed normally - logger.info("✅ Model is already wrapped (detected via __dict__/_modules)") - 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 - # Get from __dict__ or _modules - transformer = get_model_attr(model_wrapper) - if transformer is None: - raise AttributeError(f"'{type(model_wrapper).__name__}' object has no attribute 'model'") - elif model_wrapper_type_name == "NunchakuZImageTransformer2DModel" or model_wrapper_type_name.endswith("NunchakuZImageTransformer2DModel"): - # Not wrapped yet, need to wrap it first - logger.info("🔧 Wrapping NunchakuZImageTransformer2DModel with ComfyZImageTurboWrapper") - - # Create wrapper - logger.info(f"📦 Creating ComfyZImageTurboWrapper with cpu_offload='{cpu_offload}'") - wrapped_model = ComfyZImageTurboWrapper( - model_wrapper, - getattr(model_wrapper, '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 - # Get from __dict__ or _modules - transformer = get_model_attr(model_wrapper) - if transformer is None: - raise AttributeError(f"'{type(model_wrapper).__name__}' object has no attribute 'model'") - else: - logger.error(f"❌ Model type mismatch! Type: {model_wrapper_type_name}, Module: {model_wrapper_module}") - logger.error("Please use 'Nunchaku-ussoewwin Z-Image-Turbo DiT Loader'.") - raise TypeError(f"This LoRA loader only works with Nunchaku Z-Image-Turbo models, 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 = { - "NunchakuZImageTurboLoraStackV2": NunchakuZImageTurboLoraStackV2 -} - -GENERATED_DISPLAY_NAMES = { - "NunchakuZImageTurboLoraStackV2": "Nunchaku Z-Image-Turbo LoRA Stack V2" -} diff --git a/wrappers/zimageturbo.py b/wrappers/zimageturbo.py index 47f59a6..5277db6 100644 --- a/wrappers/zimageturbo.py +++ b/wrappers/zimageturbo.py @@ -477,9 +477,8 @@ class ComfyZImageTurboWrapper(nn.Module): # ============================================================================= # DEBUG LOGGING (MUTED) - Uncomment to diagnose input tensor issues # ============================================================================= - # This debug block logs the shapes of input tensors for Z-Image-Turbo forward pass. - # It is specifically active when using NunchakuZImageTurboLoraStackV2 (zimageturbo_v2.py) - # which wraps the model with ComfyZImageTurboWrapper. + # This debug block logs the shapes of input tensors for Z-Image-Turbo forward pass + # when the model is wrapped with ComfyZImageTurboWrapper. # # Output format: # - x_list: List of latent tensors, shape [C, F, H, W] = [channels, frames, height, width]