#!/usr/bin/env python3 """Strict loaders for Wan2.2-TI2V-5B four-step LoRA checkpoints. Both the published Perflow checkpoint and the local non-AR DMD checkpoint are torch ``model.pt`` files containing ``generator_lora`` and ``critic_lora``. Inference must load only the former. This helper validates all 300 expected Wan linear targets before installing a PEFT adapter, and supports both native Wan and Diffusers module names. """ from __future__ import annotations import json import hashlib import math import os import re from pathlib import Path from typing import Any import torch from torch import nn PERFLOW_ADAPTER_NAME = "perflow_step4" LOCAL_NONAR_ADAPTER_NAME = "local_nonar_step4" # Backward-compatible import for existing callers. ADAPTER_NAME = PERFLOW_ADAPTER_NAME EXPECTED_BLOCKS = tuple(range(30)) EXPECTED_SUFFIXES = ( "self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o", "cross_attn.q", "cross_attn.k", "cross_attn.v", "cross_attn.o", "ffn.0", "ffn.2", ) _LORA_KEY = re.compile( r"^(?:base_model\.model\.)?(blocks\.(\d+)\.(.+))\.lora_([AB])\.weight$" ) def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _target_name(source_name: str, style: str) -> str: if style == "native": return source_name if style != "diffusers": raise ValueError(f"Unknown Wan module style: {style!r}") replacements = { ".self_attn.q": ".attn1.to_q", ".self_attn.k": ".attn1.to_k", ".self_attn.v": ".attn1.to_v", ".self_attn.o": ".attn1.to_out.0", ".cross_attn.q": ".attn2.to_q", ".cross_attn.k": ".attn2.to_k", ".cross_attn.v": ".attn2.to_v", ".cross_attn.o": ".attn2.to_out.0", ".ffn.0": ".ffn.net.0.proj", ".ffn.2": ".ffn.net.2", } for suffix, replacement in replacements.items(): if source_name.endswith(suffix): return source_name[: -len(suffix)] + replacement raise ValueError(f"Unsupported Wan LoRA target: {source_name}") def inspect_wan_lora_checkpoint( checkpoint_path: str | Path, *, expected_rank: int, lora_alpha: int | None = None, checkpoint_kind: str, expected_step: int | None = None, ) -> dict[str, Any]: """Read and strictly validate one 30-block Wan generator LoRA.""" if expected_rank <= 0: raise ValueError(f"expected_rank must be positive, got {expected_rank}") alpha = expected_rank if lora_alpha is None else lora_alpha if alpha <= 0: raise ValueError(f"lora_alpha must be positive, got {alpha}") checkpoint_path = Path(checkpoint_path).resolve() payload = torch.load( checkpoint_path, map_location="cpu", weights_only=True, mmap=True, ) if not isinstance(payload, dict): raise TypeError(f"Expected a dict checkpoint, got {type(payload).__name__}") if "generator_lora" not in payload: raise KeyError(f"Missing generator_lora; top-level keys: {list(payload)}") checkpoint_step = int(payload.get("step", -1)) if expected_step is not None and checkpoint_step != expected_step: raise ValueError( f"Expected step {expected_step} for {checkpoint_kind}, " f"got {checkpoint_step}" ) state = payload["generator_lora"] if not isinstance(state, dict): raise TypeError("generator_lora must be a state dict") pairs: dict[str, dict[str, torch.Tensor]] = {} bad_keys: list[str] = [] for key, tensor in state.items(): match = _LORA_KEY.match(key) if match is None or not isinstance(tensor, torch.Tensor): bad_keys.append(key) continue source_name, block_text, suffix, side = match.groups() block = int(block_text) if block not in EXPECTED_BLOCKS or suffix not in EXPECTED_SUFFIXES: bad_keys.append(key) continue pairs.setdefault(source_name, {})[side] = tensor expected_names = { f"blocks.{block}.{suffix}" for block in EXPECTED_BLOCKS for suffix in EXPECTED_SUFFIXES } missing_names = sorted(expected_names - set(pairs)) extra_names = sorted(set(pairs) - expected_names) incomplete = sorted(name for name, pair in pairs.items() if set(pair) != {"A", "B"}) if bad_keys or missing_names or extra_names or incomplete or len(state) != 600: raise ValueError( f"Invalid {checkpoint_kind} generator_lora structure: " f"tensors={len(state)}, bad={len(bad_keys)}, missing={len(missing_names)}, " f"extra={len(extra_names)}, incomplete={len(incomplete)}" ) ranks = set() parameter_count = 0 for name, pair in pairs.items(): a, b = pair["A"], pair["B"] if a.ndim != 2 or b.ndim != 2 or a.shape[0] != b.shape[1]: raise ValueError(f"Invalid A/B shapes for {name}: A={tuple(a.shape)}, B={tuple(b.shape)}") ranks.add(int(a.shape[0])) parameter_count += a.numel() + b.numel() if ranks != {expected_rank}: raise ValueError( f"Expected rank {expected_rank} throughout for {checkpoint_kind}, " f"got {sorted(ranks)}" ) metadata = { "checkpoint_kind": checkpoint_kind, "checkpoint": str(checkpoint_path), "checkpoint_bytes": checkpoint_path.stat().st_size, "checkpoint_sha256": _sha256_file(checkpoint_path), "top_level_keys": sorted(payload), "step": checkpoint_step, "selected_state": "generator_lora", "ignored_state": "critic_lora" if "critic_lora" in payload else None, "tensor_count": len(state), "target_count": len(pairs), "blocks": len(EXPECTED_BLOCKS), "targets_per_block": len(EXPECTED_SUFFIXES), "rank": expected_rank, "alpha": alpha, "scale": alpha / expected_rank, "parameter_count": parameter_count, } return {"metadata": metadata, "pairs": pairs} def inspect_perflow_checkpoint(checkpoint_path: str | Path) -> dict[str, Any]: """Validate the published rank-128 Perflow generator LoRA.""" return inspect_wan_lora_checkpoint( checkpoint_path, expected_rank=128, lora_alpha=128, checkpoint_kind="perflow_step4", expected_step=None, ) def inspect_local_nonar_checkpoint( checkpoint_path: str | Path, *, expected_step: int | None = 1600, expected_rank: int | None = None, ) -> dict[str, Any]: """Validate a local non-AR generator LoRA with a uniform inferred rank. ``1600`` remains the fail-closed default for the historical four-step DMD artifact. Native-step CFG-distillation callers must pass their checkpoint identity explicitly instead of being silently coupled to that old step. Rank is inferred by default so matched rank sweeps are not incorrectly rejected by the historical rank-64 assumption. """ checkpoint_path = Path(checkpoint_path).resolve() if expected_rank is None: payload = torch.load( checkpoint_path, map_location="cpu", weights_only=True, mmap=True, ) if not isinstance(payload, dict) or not isinstance( payload.get("generator_lora"), dict ): raise TypeError(f"Expected generator_lora state dict in {checkpoint_path}") state = payload["generator_lora"] ranks = { int(tensor.shape[0]) for key, tensor in state.items() if key.endswith(".lora_A.weight") and isinstance(tensor, torch.Tensor) and tensor.ndim == 2 } if len(ranks) != 1: raise ValueError( f"Expected one uniform LoRA rank in {checkpoint_path}, " f"got {sorted(ranks)}" ) expected_rank = next(iter(ranks)) del state, payload return inspect_wan_lora_checkpoint( checkpoint_path, expected_rank=expected_rank, lora_alpha=expected_rank, checkpoint_kind="local_nonar_step4", expected_step=expected_step, ) def attach_wan_lora( model: nn.Module, checkpoint_path: str | Path, *, style: str, expected_rank: int, lora_alpha: int, adapter_name: str, checkpoint_kind: str, report_prefix: str, expected_step: int | None = None, enabled: bool = False, ) -> dict[str, Any]: """Install the generator LoRA on a native-Wan or Diffusers Wan backbone. Every source tensor and every destination module must match. Any missing or shape-mismatched layer aborts the run instead of silently producing a partial LoRA load. """ inspected = inspect_wan_lora_checkpoint( checkpoint_path, expected_rank=expected_rank, lora_alpha=lora_alpha, checkpoint_kind=checkpoint_kind, expected_step=expected_step, ) metadata = inspected["metadata"] pairs = inspected["pairs"] destinations: dict[str, tuple[str, dict[str, torch.Tensor], nn.Linear]] = {} for source_name, pair in pairs.items(): target_name = _target_name(source_name, style) try: target = model.get_submodule(target_name) except AttributeError as exc: raise ValueError(f"Missing LoRA destination {target_name} (from {source_name})") from exc if not isinstance(target, nn.Linear): raise TypeError(f"LoRA destination {target_name} is {type(target).__name__}, not nn.Linear") a, b = pair["A"], pair["B"] expected_a = (a.shape[0], target.in_features) expected_b = (target.out_features, a.shape[0]) if tuple(a.shape) != expected_a or tuple(b.shape) != expected_b: raise ValueError( f"Shape mismatch for {target_name}: target={tuple(target.weight.shape)}, " f"A={tuple(a.shape)} expected={expected_a}, B={tuple(b.shape)} expected={expected_b}" ) destinations[target_name] = (source_name, pair, target) if len(destinations) != 300: raise ValueError(f"Expected 300 unique LoRA destinations, got {len(destinations)}") # PEFT 0.19 probes torchao before checking whether a layer is quantized. The # inherited environment has torchao 0.11, so explicitly disable that optional # dispatcher; ordinary torch.nn.Linear layers use PEFT's native dispatcher. import peft.tuners.lora.torchao as peft_torchao from peft import LoraConfig, inject_adapter_in_model peft_torchao.is_torchao_available = lambda: False config = LoraConfig( r=expected_rank, lora_alpha=lora_alpha, lora_dropout=0.0, bias="none", target_modules=sorted(destinations), init_lora_weights=True, ) inject_adapter_in_model(config, model, adapter_name=adapter_name) copied_tensors = 0 nonzero_b = 0 adapter_dtypes: set[str] = set() with torch.no_grad(): for target_name, (_source_name, pair, _old_target) in destinations.items(): layer = model.get_submodule(target_name) if adapter_name not in layer.lora_A or adapter_name not in layer.lora_B: raise RuntimeError(f"PEFT did not create adapter tensors for {target_name}") a_param = layer.lora_A[adapter_name].weight b_param = layer.lora_B[adapter_name].weight a_param.copy_(pair["A"].to(device=a_param.device, dtype=a_param.dtype)) b_param.copy_(pair["B"].to(device=b_param.device, dtype=b_param.dtype)) copied_tensors += 2 nonzero_b += int(torch.count_nonzero(pair["B"]).item() > 0) adapter_dtypes.add(str(a_param.dtype).replace("torch.", "")) scale = float(layer.scaling[adapter_name]) expected_scale = lora_alpha / expected_rank if scale != expected_scale: raise RuntimeError( f"Unexpected LoRA scale for {target_name}: {scale}; " f"expected {expected_scale}" ) toggled = _set_adapter_enabled( model, adapter_name=adapter_name, enabled=enabled, report_prefix=report_prefix, allow_absent=False, ) if toggled != 300 or copied_tensors != 600 or nonzero_b != 300: raise RuntimeError( f"Incomplete LoRA installation: layers={toggled}, tensors={copied_tensors}, " f"nonzero_B={nonzero_b}" ) report = { **metadata, "module_style": style, "adapter_name": adapter_name, "loaded_tensors": copied_tensors, "loaded_targets": toggled, "missing_targets": 0, "shape_mismatches": 0, "nonzero_B_targets": nonzero_b, "adapter_dtypes": sorted(adapter_dtypes), "enabled": bool(enabled), } print(report_prefix + "_LOAD_REPORT=" + json.dumps(report, sort_keys=True), flush=True) return report def attach_perflow_lora( model: nn.Module, checkpoint_path: str | Path, *, style: str, enabled: bool = False, ) -> dict[str, Any]: """Install the published rank-128 Perflow generator LoRA.""" return attach_wan_lora( model, checkpoint_path, style=style, expected_rank=128, lora_alpha=128, adapter_name=PERFLOW_ADAPTER_NAME, checkpoint_kind="perflow_step4", report_prefix="PERFLOW_LORA", expected_step=None, enabled=enabled, ) def apply_adapter_strength( model: nn.Module, *, adapter_name: str, strength: float, expected_targets: int = 300, ) -> dict[str, Any]: """Apply one linear inference multiplier through PEFT's scaling value.""" strength = float(strength) if not math.isfinite(strength) or strength <= 0: raise ValueError(f"LoRA inference strength must be positive and finite, got {strength}") base_scalings: set[float] = set() effective_scalings: set[float] = set() scaled_targets = 0 for layer in model.modules(): scaling = getattr(layer, "scaling", None) if scaling is None or adapter_name not in scaling: continue base_scale = float(scaling[adapter_name]) effective_scale = base_scale * strength scaling[adapter_name] = effective_scale base_scalings.add(base_scale) effective_scalings.add(effective_scale) scaled_targets += 1 if scaled_targets != expected_targets: raise RuntimeError( "Incomplete LoRA inference-strength application: " f"targets={scaled_targets}, expected={expected_targets}" ) return { "inference_strength": strength, "base_adapter_scaling": sorted(base_scalings), "effective_adapter_scaling": sorted(effective_scalings), "scaled_targets": scaled_targets, "scale_semantics": "delta_W = inference_strength * (B @ A) * alpha/r", } def attach_local_nonar_lora( model: nn.Module, checkpoint_path: str | Path, *, style: str, enabled: bool = True, strength: float | None = None, ) -> dict[str, Any]: """Install a local four-step non-AR DMD generator LoRA. The transfer study compares independently trained WAN, SCOPE and TheDenk adapters. They intentionally use different ranks and checkpoint steps, so infer those provenance fields from the checkpoint while retaining the strict 600-tensor / 300-target / destination-shape validation performed by :func:`attach_wan_lora`. """ checkpoint_path = Path(checkpoint_path).resolve() payload = torch.load( checkpoint_path, map_location="cpu", weights_only=True, mmap=True, ) if not isinstance(payload, dict) or not isinstance(payload.get("generator_lora"), dict): raise TypeError(f"Expected generator_lora state dict in {checkpoint_path}") state = payload["generator_lora"] ranks = { int(tensor.shape[0]) for key, tensor in state.items() if key.endswith(".lora_A.weight") and isinstance(tensor, torch.Tensor) and tensor.ndim == 2 } if len(ranks) != 1: raise ValueError(f"Expected one uniform LoRA rank in {checkpoint_path}, got {sorted(ranks)}") rank = next(iter(ranks)) step = payload.get("step") expected_step = int(step) if step is not None else None del state, payload if strength is None: strength = float(os.environ.get("LOCAL_NONAR_LORA_SCALE", "1")) strength = float(strength) if not math.isfinite(strength) or strength <= 0: raise ValueError(f"LoRA inference strength must be positive and finite, got {strength}") report = attach_wan_lora( model, checkpoint_path, style=style, expected_rank=rank, lora_alpha=rank, adapter_name=LOCAL_NONAR_ADAPTER_NAME, checkpoint_kind="local_nonar_step4", report_prefix="LOCAL_NONAR_LORA", expected_step=expected_step, enabled=enabled, ) report = { **report, **apply_adapter_strength( model, adapter_name=LOCAL_NONAR_ADAPTER_NAME, strength=strength, expected_targets=300, ), } report_path_raw = os.environ.get("LOCAL_NONAR_LORA_REPORT_PATH") if report_path_raw: report_path = Path(report_path_raw).resolve() report_path.parent.mkdir(parents=True, exist_ok=True) temporary = report_path.with_suffix(report_path.suffix + ".tmp") temporary.write_text( json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) temporary.replace(report_path) print( "LOCAL_NONAR_LORA_SCALE_REPORT=" + json.dumps(report, sort_keys=True), flush=True, ) return report def _set_adapter_enabled( model: nn.Module, *, adapter_name: str, enabled: bool, report_prefix: str, allow_absent: bool, ) -> int: count = 0 for module in model.modules(): lora_a = getattr(module, "lora_A", None) if lora_a is None or adapter_name not in lora_a: continue module.enable_adapters(enabled=enabled) count += 1 expected_counts = (0, 300) if allow_absent else (300,) if count not in expected_counts: raise RuntimeError( f"Expected {'0 or ' if allow_absent else ''}300 {adapter_name} layers " f"while toggling, found {count}" ) print(f"{report_prefix}_ENABLED={int(enabled)} layers={count}", flush=True) return count def set_perflow_enabled(model: nn.Module, enabled: bool) -> int: """Enable/disable only the installed Perflow tuner layers.""" return _set_adapter_enabled( model, adapter_name=PERFLOW_ADAPTER_NAME, enabled=enabled, report_prefix="PERFLOW_LORA", allow_absent=True, ) def set_local_nonar_enabled(model: nn.Module, enabled: bool) -> int: """Enable/disable only the installed local non-AR tuner layers.""" return _set_adapter_enabled( model, adapter_name=LOCAL_NONAR_ADAPTER_NAME, enabled=enabled, report_prefix="LOCAL_NONAR_LORA", allow_absent=True, )