nodes/ussoewwin-ComfyUI-QwenImageLoraLoader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix nn.Module attribute access issues and translate Japanese comments to English

Browse Source
This commit is contained in:
ussoewwin
2026-05-17 05:26:40 +09:00
parent b700470ab2
commit 34c0b5909d
2 changed files with 85 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/lora/zimageturbo_v2.py
+23 -4
View File
@@ -156,15 +156,31 @@ class NunchakuZImageTurboLoraStackV2:
# 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'):
# 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 attributes)")
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
transformer = model_wrapper.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'")
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")
@@ -183,7 +199,10 @@ class NunchakuZImageTurboLoraStackV2:
# Replace the model's diffusion_model with our wrapper
model.model.diffusion_model = wrapped_model
model_wrapper = wrapped_model
transformer = model_wrapper.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'.")
wrappers/zimageturbo.py
+62 -18
View File
@@ -63,28 +63,57 @@ class ComfyZImageTurboWrapper(nn.Module):
# Track last seen device to detect CPU/GPU moves that require re-compose
self._last_device = None
def _get_model_safe(self):
"""Safely get model attribute from _modules or __dict__ of nn.Module.
In nn.Module subclasses, when self.model = nn.Module() is set, the model attribute
is stored in the _modules dictionary. Access via __getattr__ can cause issues,
so we retrieve it directly from the dictionary.
"""
# First check __dict__ (regular Python attributes)
if 'model' in self.__dict__:
return self.__dict__['model']
# Next check _modules (nn.Module child modules)
modules = self.__dict__.get('_modules', {})
if modules and 'model' in modules:
return modules['model']
return None
def to_safely(self, device):
"""Safely move the model to the specified device."""
if self.model is None:
# Use _get_model_safe to bypass nn.Module's attribute system
model = self._get_model_safe()
if model is None:
return self
if hasattr(self.model, "to_safely"):
self.model.to_safely(device)
if hasattr(model, "to_safely"):
model.to_safely(device)
else:
self.model.to(device)
model.to(device)
return self
def __getattr__(self, name):
"""Delegate attribute access to the wrapped model (e.g., layers, config, etc.)."""
# Special handling for direct access to 'model' attribute
# nn.Module stores child modules in _modules, so normal attribute access won't find them
if name == 'model':
# Try to get from _modules
modules = self.__dict__.get('_modules', {})
if modules and 'model' in modules:
return modules['model']
# Try to get from __dict__
if 'model' in self.__dict__:
return self.__dict__['model']
raise AttributeError(f"'{type(self).__name__}' object has no attribute 'model'")
# Delegate attributes that don't exist on this wrapper to the wrapped model (NextDiT)
# This allows ZImageModelPatcher and other code to access NextDiT attributes like 'layers'
# Note: 'model' is a normal attribute set in __init__, so __getattr__ should not be called for it
# Use super().__getattribute__ to access 'model' through nn.Module's attribute system
try:
model = super().__getattribute__('model')
if model is not None:
# Use _get_model_safe to bypass nn.Module's attribute system
model = self._get_model_safe()
if model is not None:
try:
return getattr(model, name)
except AttributeError:
pass
except AttributeError:
pass
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
def forward(
@@ -445,13 +474,28 @@ class ComfyZImageTurboWrapper(nn.Module):
zimage_kwargs["f_patch_size"] = kwargs["f_patch_size"]
zimage_kwargs["return_dict"] = False
# DEBUG: Log shapes to diagnose empty iterable error
logger.info(f"DEBUG: x_list len={len(x_list) if isinstance(x_list, list) else 'Not List'}")
if isinstance(x_list, list) and len(x_list) > 0:
logger.info(f"DEBUG: x_list[0] shape={x_list[0].shape}")
logger.info(f"DEBUG: cap_feats len={len(cap_feats) if isinstance(cap_feats, list) else 'Not List'}")
if isinstance(cap_feats, list) and len(cap_feats) > 0:
logger.info(f"DEBUG: cap_feats[0] shape={cap_feats[0].shape}")
# =============================================================================
# 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.
#
# Output format:
# - x_list: List of latent tensors, shape [C, F, H, W] = [channels, frames, height, width]
# - cap_feats: Caption features from text encoder, shape [tokens, embed_dim]
#
# Note: Logs appear twice per step due to CFG (Classifier-Free Guidance) which runs
# both positive and negative conditioning through the model.
#
# To enable: Remove the '#' prefix from the logger.info lines below.
# =============================================================================
# logger.info(f"DEBUG: x_list len={len(x_list) if isinstance(x_list, list) else 'Not List'}")
# if isinstance(x_list, list) and len(x_list) > 0:
# logger.info(f"DEBUG: x_list[0] shape={x_list[0].shape}")
# logger.info(f"DEBUG: cap_feats len={len(cap_feats) if isinstance(cap_feats, list) else 'Not List'}")
# if isinstance(cap_feats, list) and len(cap_feats) > 0:
# logger.info(f"DEBUG: cap_feats[0] shape={cap_feats[0].shape}")
# Check for mismatch which causes zip to truncate
if isinstance(x_list, list) and isinstance(cap_feats, list):