Add v2.2.4 features: AWQ modulation layer detection/skip logic and NextDiT forward signature support

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 = []

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')}