diff --git a/nodes/lora/zimageturbo_v2.py b/nodes/lora/zimageturbo_v2.py index aece199..162b923 100644 --- a/nodes/lora/zimageturbo_v2.py +++ b/nodes/lora/zimageturbo_v2.py @@ -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'.") diff --git a/wrappers/zimageturbo.py b/wrappers/zimageturbo.py index 2bf520b..47f59a6 100644 --- a/wrappers/zimageturbo.py +++ b/wrappers/zimageturbo.py @@ -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):