From 2c2fb8bdf199730cdac8577b32aa43834c996b58 Mon Sep 17 00:00:00 2001 From: Your Name Date: Thu, 22 Jan 2026 09:53:15 +0900 Subject: [PATCH] feat: add AWQ enable switch to V2 node --- nodes/lora/qwenimage_v2.py | 65 +++++--- nunchaku_code/lora_qwen.py | 302 +++++++++++++++++++++++++++++++++++++ 2 files changed, 350 insertions(+), 17 deletions(-) diff --git a/nodes/lora/qwenimage_v2.py b/nodes/lora/qwenimage_v2.py index 2ffdc8b..cfac0bb 100644 --- a/nodes/lora/qwenimage_v2.py +++ b/nodes/lora/qwenimage_v2.py @@ -80,6 +80,13 @@ class NunchakuQwenImageLoraStackV2: "tooltip": "CPU offload setting. 'auto' enables offload when VRAM is low, 'enable' forces offload, 'disable' disables offload.", }, ), + "apply_awq_mod": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Enable LoRA application to AWQ modulation layers (img_mod.1/txt_mod.1). Disabled by default to prevent noise. Enable at your own risk.", + }, + ), }, "optional": {}, } @@ -110,7 +117,7 @@ class NunchakuQwenImageLoraStackV2: 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): + def load_lora_stack(self, model, lora_count, cpu_offload="disable", apply_awq_mod=False, **kwargs): loras_to_apply = [] # Process only the number of LoRAs specified by lora_count @@ -135,12 +142,12 @@ class NunchakuQwenImageLoraStackV2: sys.path.insert(0, lora_loader_dir) spec = importlib.util.spec_from_file_location( - "wrappers.qwenimage", - os.path.join(lora_loader_dir, "wrappers", "qwenimage.py") + "wrappers.qwenimage_v2", + os.path.join(lora_loader_dir, "wrappers", "qwenimage_v2.py") ) wrappers_module = importlib.util.module_from_spec(spec) spec.loader.exec_module(wrappers_module) - ComfyQwenImageWrapper = wrappers_module.ComfyQwenImageWrapper + ComfyQwenImageWrapperV2 = wrappers_module.ComfyQwenImageWrapperV2 from nunchaku import NunchakuQwenImageTransformer2DModel @@ -153,31 +160,55 @@ class NunchakuQwenImageLoraStackV2: 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 - # ComfyQwenImageWrapper has a 'model' attribute that contains the transformer + # Check if it's already wrapped with V2 wrapper + # ComfyQwenImageWrapperV2 has a 'model' attribute that contains the transformer # and a 'loras' list attribute if hasattr(model_wrapper, 'model') and hasattr(model_wrapper, 'loras'): - # Already wrapped, proceed normally - logger.info("✅ Model is already wrapped (detected via attributes)") - 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 + # Check if it's a V2 wrapper (has apply_awq_mod attribute) + if hasattr(model_wrapper, 'apply_awq_mod'): + # Already wrapped with V2 wrapper, proceed normally + logger.info("✅ Model is already wrapped with V2 wrapper (detected via attributes)") + logger.info(f"📦 Current CPU offload setting: '{model_wrapper.cpu_offload_setting}'") + logger.info(f"📦 Current AWQ mod setting: '{model_wrapper.apply_awq_mod}'") + # 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 + # Update AWQ mod setting if different + if model_wrapper.apply_awq_mod != apply_awq_mod: + logger.info(f"🔄 Updating AWQ mod setting from '{model_wrapper.apply_awq_mod}' to '{apply_awq_mod}'") + model_wrapper.apply_awq_mod = apply_awq_mod + transformer = model_wrapper.model + else: + # Wrapped with V1/V3 wrapper, need to re-wrap with V2 + logger.info("🔄 Model is wrapped with V1/V3 wrapper, re-wrapping with V2 wrapper") + transformer = model_wrapper.model + config = getattr(model_wrapper, 'config', {}) + wrapped_model = ComfyQwenImageWrapperV2( + transformer, + config, + None, # customized_forward + {}, # forward_kwargs + cpu_offload, # cpu_offload_setting + 4.0, # vram_margin_gb + apply_awq_mod, # apply_awq_mod + ) + model.model.diffusion_model = wrapped_model + model_wrapper = wrapped_model elif model_wrapper_type_name == "NunchakuQwenImageTransformer2DModel" or model_wrapper_type_name.endswith("NunchakuQwenImageTransformer2DModel"): # Not wrapped yet, need to wrap it first - logger.info("🔧 Wrapping NunchakuQwenImageTransformer2DModel with ComfyQwenImageWrapper") + logger.info("🔧 Wrapping NunchakuQwenImageTransformer2DModel with ComfyQwenImageWrapperV2") # Create wrapper - logger.info(f"📦 Creating ComfyQwenImageWrapper with cpu_offload='{cpu_offload}'") - wrapped_model = ComfyQwenImageWrapper( + logger.info(f"📦 Creating ComfyQwenImageWrapperV2 with cpu_offload='{cpu_offload}', apply_awq_mod={apply_awq_mod}") + wrapped_model = ComfyQwenImageWrapperV2( model_wrapper, getattr(model_wrapper, 'config', {}), None, # customized_forward {}, # forward_kwargs cpu_offload, # cpu_offload_setting 4.0, # vram_margin_gb + apply_awq_mod, # apply_awq_mod ) # Replace the model's diffusion_model with our wrapper diff --git a/nunchaku_code/lora_qwen.py b/nunchaku_code/lora_qwen.py index 65e2b1e..ecbdc60 100644 --- a/nunchaku_code/lora_qwen.py +++ b/nunchaku_code/lora_qwen.py @@ -1500,6 +1500,308 @@ def set_lora_strength_v2(model: nn.Module, strength: float) -> None: model._lora_strength = strength +def compose_loras_v2_v2( + model: torch.nn.Module, + lora_configs: List[Tuple[Union[str, Path, Dict[str, torch.Tensor]], float]], + apply_awq_mod: bool | str | None = None, +) -> bool: + """ + V2 node-specific LoRA composition function with AWQ modulation control. + + This function is identical to compose_loras_v2 except it accepts an apply_awq_mod + parameter that overrides the global _APPLY_AWQ_MOD environment variable. + + Args: + model: The model to apply LoRAs to + lora_configs: List of (lora_path_or_dict, strength) tuples + apply_awq_mod: If True/False, overrides global _APPLY_AWQ_MOD setting. + If None or "auto", uses global _APPLY_AWQ_MOD setting. + If string "0"/"false"/"no"/"off" or "1"/"true"/"yes"/"on", parsed as bool. + + Returns: + bool: True if the LoRA format is supported and processed, False otherwise. + """ + # Determine effective AWQ modulation setting: + # - If apply_awq_mod is explicitly provided, use it + # - Otherwise, fall back to global _APPLY_AWQ_MOD + if isinstance(apply_awq_mod, str): + apply_awq_mod_flag = apply_awq_mod.lower() not in ("0", "false", "no", "n", "off", "auto", "") + elif isinstance(apply_awq_mod, bool): + apply_awq_mod_flag = apply_awq_mod + else: + # None or other -> use global setting + apply_awq_mod_flag = _APPLY_AWQ_MOD + + logger.info(f"[V2 Node] Composing {len(lora_configs)} LoRAs with apply_awq_mod={apply_awq_mod} (effective: {apply_awq_mod_flag})...") + reset_lora_v2(model) + _first_detection = None # Initialize for scope safety + + # Same mapping switching logic as compose_loras_v2 + global _ACTIVE_KEY_MAPPING + prev_mapping = _ACTIVE_KEY_MAPPING + try: + nextdit_markers = ( + "layers.0.attention.qkv", + "layers.0.attention.out", + "layers.0.feed_forward.w1", + "layers.0.feed_forward.w2", + "layers.0.feed_forward.w3", + "layers.0.feed_forward.w13", + ) + is_nextdit_style = any(_get_module_by_name(model, p) is not None for p in nextdit_markers) + if is_nextdit_style: + has_w13 = _get_module_by_name(model, "layers.0.feed_forward.w13") is not None + if has_w13: + _ACTIVE_KEY_MAPPING = ZIMAGE_NEXTDIT_NUNCHAKU_PATCHED_KEY_MAPPING + KEY_MAPPING + else: + _ACTIVE_KEY_MAPPING = ZIMAGE_NEXTDIT_UNPATCHED_KEY_MAPPING + KEY_MAPPING + else: + _ACTIVE_KEY_MAPPING = None + + # Same LoRA loading and processing logic as compose_loras_v2 + _cached_first_lora_state_dict = None + if lora_configs: + first_lora_path_or_dict, first_lora_strength = lora_configs[0] + first_lora_state_dict = _load_lora_state_dict(first_lora_path_or_dict) + _cached_first_lora_state_dict = first_lora_state_dict + logger.info(f"--- DEBUG: Inspecting keys for LoRA 1 (Strength: {first_lora_strength}) ---") + + _first_detection = _detect_lora_format(first_lora_state_dict) + if _first_detection["has_standard"]: + for key in first_lora_state_dict.keys(): + parsed_res = _classify_and_map_key(key) + if parsed_res: + group, base_key, comp, ab = parsed_res + mapped_name = f"{base_key}.{comp}.{ab}" if comp and ab else (f"{base_key}.{ab}" if ab else base_key) + logger.info(f"Key: {key} -> Mapped to: {mapped_name} (Group: {group})") + else: + logger.warning(f"Key: {key} -> UNMATCHED (Ignored)") + else: + logger.warning(f"⚠️ Unsupported LoRA format detected (No standard keys).") + logger.warning(f" Skipping detailed key inspection of {len(first_lora_state_dict)} keys to prevent console freeze.") + + logger.info("--- DEBUG: End key inspection ---") + + aggregated_weights: Dict[str, List[Dict[str, Any]]] = defaultdict(list) + + # Same aggregation logic as compose_loras_v2 + for idx, (lora_path_or_dict, strength) in enumerate(lora_configs): + lora_name = lora_path_or_dict if isinstance(lora_path_or_dict, str) else "dict" + if idx == 0 and _cached_first_lora_state_dict is not None: + lora_state_dict = _cached_first_lora_state_dict + else: + lora_state_dict = _load_lora_state_dict(lora_path_or_dict) + + try: + detection = _detect_lora_format(lora_state_dict) + _log_lora_format_detection(str(lora_name), detection) + except Exception: + pass + + lora_grouped: Dict[str, Dict[str, torch.Tensor]] = defaultdict(dict) + lokr_keys_count = 0 + standard_keys_count = 0 + qkv_lokr_keys_count = 0 + unrecognized_keys = [] + total_keys = len(lora_state_dict) + + for key, value in lora_state_dict.items(): + parsed = _classify_and_map_key(key) + if parsed is None: + unrecognized_keys.append(key) + continue + + group, base_key, comp, ab = parsed + if ab in ("lokr_w1", "lokr_w2"): + lokr_keys_count += 1 + if group in ("qkv", "add_qkv") and comp is not None: + qkv_lokr_keys_count += 1 + elif ab in ("A", "B"): + standard_keys_count += 1 + + if group in ("qkv", "add_qkv", "glu") and comp is not None: + if ab in ("lokr_w1", "lokr_w2"): + lora_grouped[base_key][f"{comp}_{ab}"] = value + else: + lora_grouped[base_key][f"{comp}_{ab}"] = value + else: + lora_grouped[base_key][ab] = value + + has_lokr = lokr_keys_count > 0 + has_standard = standard_keys_count > 0 + if has_lokr and has_standard: + lora_format = "Mixed (LoKR + Standard LoRA)" + elif has_lokr: + lora_format = "LoKR (QKV format)" if qkv_lokr_keys_count > 0 else "LoKR" + elif has_standard: + lora_format = "Standard LoRA" + else: + lora_format = "Unknown/Unsupported" + logger.debug(f"LoRA summary: {lora_name} | Format: {lora_format} | Strength: {strength:.3f}") + + processed_groups = {} + special_handled = set() + for base_key, lw in lora_grouped.items(): + if base_key in special_handled: + continue + + if "lokr_w1" in lw or "lokr_w2" in lw: + logger.warning( + f"Skipping LoKR weights for {base_key}: LoKR support is currently experimental and disabled due to compatibility issues (produces noise). Please convert LoKR to standard LoRA first." + ) + continue + + if "qkv" in base_key: + A, B, alpha = (lw.get("A"), lw.get("B"), lw.get("alpha")) if "A" in lw else _fuse_qkv_lora(lw, model=model, base_key=base_key) + elif "w1_A" in lw or "w3_A" in lw: + A, B, alpha = _fuse_glu_lora(lw) + elif ".proj_out" in base_key and "single_transformer_blocks" in base_key: + split_map, consumed_keys = _handle_proj_out_split(lora_grouped, base_key, model) + processed_groups.update(split_map) + special_handled.update(consumed_keys) + continue + else: + A, B, alpha = lw.get("A"), lw.get("B"), lw.get("alpha") + + if A is not None and B is not None: + processed_groups[base_key] = (A, B, alpha) + + if not processed_groups: + if lora_format == "Unknown/Unsupported": + logger.error(f"❌ {lora_name}: No weights were processed - LoRA format is unsupported and will be skipped!") + else: + logger.warning(f"⚠️ {lora_name}: No weights were processed - this LoRA will have no effect!") + else: + logger.debug(f" {lora_name}: Processed {len(processed_groups)} module groups") + + for module_key, (A, B, alpha) in processed_groups.items(): + aggregated_weights[module_key].append( + {"A": A, "B": B, "alpha": alpha, "strength": strength, "source": lora_name} + ) + + # 2. Apply aggregated weights to the model with V2-specific AWQ modulation control + applied_modules_count = 0 + + for module_name, parts in aggregated_weights.items(): + resolved_name, module = _resolve_module_name(model, module_name) + if module is None: + logger.debug(f"[MISS] Module not found: {module_name} (resolved: {resolved_name})") + continue + + 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") + ) + + is_modulation_layer = (".img_mod.1" in resolved_name or ".txt_mod.1" in resolved_name) + + # V2-specific: Use apply_awq_mod_flag instead of _APPLY_AWQ_MOD + if is_awq_w4a16 and is_modulation_layer and not apply_awq_mod_flag: + logger.warning( + f"[V2 Node] [SKIP] {resolved_name}: AWQ modulation layer LoRA is disabled by default (prevents noise). " + f"Set QWENIMAGE_LORA_APPLY_AWQ_MOD=1 or enable the node switch to force-enable." + ) + continue + + # Same application logic as compose_loras_v2 + if ( + not (hasattr(module, "proj_down") and hasattr(module, "proj_up")) + and not isinstance(module, nn.Linear) + and not is_awq_w4a16 + ): + logger.info( + f"[MISS] Module found but unsupported/missing proj_down/proj_up: {resolved_name} (Type: {type(module)})" + ) + continue + + # Same weight application logic as compose_loras_v2 (with V2-specific apply_awq_mod_flag) + all_A = [] + all_B_scaled = [] + for part in parts: + A, B, alpha, strength = part["A"], part["B"], part["alpha"], part["strength"] + r_lora = A.shape[0] + scale_alpha = alpha.item() if alpha is not None else float(r_lora) + scale = strength * (scale_alpha / max(1.0, float(r_lora))) + + if ".norm1.linear" in resolved_name or ".norm1_context.linear" in resolved_name: + B = reorder_adanorm_lora_up(B, splits=6) + elif ".single_transformer_blocks." in resolved_name and ".norm.linear" in resolved_name: + B = reorder_adanorm_lora_up(B, splits=3) + + # Special reorder for modulation layers when force-enabled (V2-specific: use apply_awq_mod_flag): + # Reorder B to match modulation channel layout (shift/scale/gate × 2). + if is_awq_w4a16 and is_modulation_layer and apply_awq_mod_flag: + # Expect out_features divisible by 6 + if B.shape[0] % 6 == 0: + try: + dim = B.shape[0] // 6 + B = ( + B.contiguous() + .view(6, dim, B.shape[1]) + .transpose(0, 1) + .reshape(B.shape[0], B.shape[1]) + ) + except Exception: + # Safety: never fail due to reorder + pass + else: + logger.warning( + f"{resolved_name}: expected mod up-matrix with out_features divisible by 6, " + f"got B({B.shape[0]}, {B.shape[1]}); skipping mod-channel reorder" + ) + + if hasattr(module, "proj_down"): + target_dtype = module.proj_down.dtype + target_device = module.proj_down.device + elif isinstance(module, nn.Linear): + target_dtype = module.weight.dtype + target_device = module.weight.device + else: + # AWQ: place LoRA tensors on same device as qweight; compute in fp16 by default. + qweight = getattr(module, "qweight", None) + target_device = qweight.device if torch.is_tensor(qweight) else torch.device("cpu") + target_dtype = torch.float16 + + all_A.append(A.to(dtype=target_dtype, device=target_device)) + all_B_scaled.append((B * scale).to(dtype=target_dtype, device=target_device)) + + if not all_A: + continue + + final_A = torch.cat(all_A, dim=0) + final_B = torch.cat(all_B_scaled, dim=1) + + _apply_lora_to_module(module, final_A, final_B, resolved_name, model) + logger.info(f"[V2 Node] [APPLY] LoRA applied to: {resolved_name}") + applied_modules_count += 1 + + if applied_modules_count == 0 and lora_configs: + logger.warning("[V2 Node] No LoRA weights were applied to any modules.") + else: + logger.info(f"[V2 Node] Applied LoRA weights to {applied_modules_count} module(s).") + + except Exception as e: + logger.error(f"[V2 Node] Error during LoRA composition: {e}") + import traceback + logger.error(traceback.format_exc()) + logger.error("[V2 Node] Possible causes:") + logger.error(" - Incompatible LoRA format") + logger.error(" - Model architecture mismatch") + logger.error(" - Corrupted or incompatible LoRA file(s)") + finally: + _ACTIVE_KEY_MAPPING = prev_mapping + + is_success = True + if _first_detection is not None and not _first_detection.get("has_standard", True): + is_success = False + + return is_success + + def reset_lora_v2(model: nn.Module) -> None: """Removes all appended LoRA weights from the model.""" if not hasattr(model, "_lora_slots") or not model._lora_slots: