diff --git a/nunchaku_code/lora_qwen.py b/nunchaku_code/lora_qwen.py index 6e45163..5b440da 100644 --- a/nunchaku_code/lora_qwen.py +++ b/nunchaku_code/lora_qwen.py @@ -1,5 +1,6 @@ import functools import logging +import os import re from collections import defaultdict from pathlib import Path @@ -19,6 +20,13 @@ from nunchaku.lora.flux.nunchaku_converter import ( logger = logging.getLogger(__name__) +# Safety switch: +# QwenImage's modulation linears (`img_mod.1` / `txt_mod.1`) are extremely sensitive because their +# output is reshaped into shift/scale/gate parameters. With AWQ quantization (AWQW4A16Linear), +# applying LoRA here often results in severe noise. Default to skipping these two layers. +# Users can override by setting env var `QWENIMAGE_LORA_APPLY_AWQ_MOD=1`. +_APPLY_AWQ_MOD = str(os.getenv("QWENIMAGE_LORA_APPLY_AWQ_MOD", "0")).strip().lower() in ("1", "true", "yes", "y", "on") + # --- Centralized & Optimized Key Mapping --- # This structure is faster to process and easier to maintain than a long if/elif chain. # --- CORRECTED Centralized & Optimized Key Mapping --- @@ -1119,6 +1127,26 @@ def compose_loras_v2( if not (hasattr(module, "proj_down") and hasattr(module, "proj_up")) and not isinstance(module, nn.Linear): logger.info(f"[MISS] Module found but unsupported/missing proj_down/proj_up: {resolved_name} (Type: {type(module)})") continue + + # Check for AWQ modulation layer skip logic + is_awq_w4a16 = ( + module.__class__.__name__ == "AWQW4A16Linear" + and hasattr(module, "qweight") + and hasattr(module, "wscales") + and hasattr(module, "wzeros") + and hasattr(module, "in_features") + and hasattr(module, "out_features") + ) + + # Check if this is img_mod.1 or txt_mod.1 + is_modulation_layer = ( + ".img_mod.1" in resolved_name or ".txt_mod.1" in resolved_name + ) + + # Skip AWQ modulation layers by default (unless environment variable is set) + if is_awq_w4a16 and is_modulation_layer and not _APPLY_AWQ_MOD: + logger.warning(f"[SKIP] {resolved_name}: AWQ modulation layer LoRA is disabled by default (prevents noise). Set QWENIMAGE_LORA_APPLY_AWQ_MOD=1 to force-enable.") + continue all_A = [] all_B_scaled = [] diff --git a/wrappers/zimageturbo.py b/wrappers/zimageturbo.py index 71e8146..d18f80b 100644 --- a/wrappers/zimageturbo.py +++ b/wrappers/zimageturbo.py @@ -300,7 +300,63 @@ class ComfyZImageTurboWrapper(nn.Module): ) else: with torch.inference_mode(): - # Z-Image-Turbo forward signature: forward(x, t, cap_feats, patch_size=2, f_patch_size=1, return_dict=True) + # Check for NextDiT signature first (forward(x, timesteps, context, num_tokens, ...)) + import inspect + forward_sig = inspect.signature(self.model.forward) + forward_params = set(forward_sig.parameters.keys()) + + # Branch 1: NextDiT signature (ComfyUI Lumina2 / NextDiT) + if "context" in forward_params and "num_tokens" in forward_params: + # NextDiT forward signature: forward(x, timesteps, context, num_tokens, attention_mask=None, **kwargs) + # Process x for NextDiT format + if isinstance(x, list): + # Stack list of tensors into a single tensor + 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) # Remove frame dimension if F=1 + x_tensor = x_stack + else: + x_tensor = x + + # Extract num_tokens from kwargs if not provided + num_tokens_value = kwargs.pop("num_tokens", None) + if num_tokens_value is None and isinstance(context, torch.Tensor) and context.ndim >= 2: + num_tokens_value = context.shape[1] + + # Extract attention_mask from kwargs + attention_mask_value = kwargs.pop("attention_mask", None) + + # Check if forward method supports **kwargs + 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 + 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 + + # 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, + ) + + # Branch 2: Z-Image-Turbo forward signature: forward(x, t, cap_feats, patch_size=2, f_patch_size=1, return_dict=True) # x: List[torch.Tensor] with shape (C, F, H, W) for each image # t: normalized timestep in [0, 1] # cap_feats: List[torch.Tensor] or None @@ -387,10 +443,7 @@ class ComfyZImageTurboWrapper(nn.Module): # 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()) + # forward_sig and forward_params are already defined above (NextDiT branch check) zimage_kwargs_clean = {k: v for k, v in zimage_kwargs.items() if k not in ('control', 'transformer_options')}