code/models/tt_dit/experimental/lora/adapter_loader.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /experimental /lora /adapter_loader.py

Add files using upload-large-folder tool

33d23e1 verified 11 days ago

15.7 kB

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Adapter loader for the runtime-LoRA Wan 2.2 pipeline.
Reads a LoRA safetensors file, parses A/B pairs using the same key
normalization as ``pipeline_wan_lora.py``, and uploads them into the
LoRA-aware Linear modules of a `WanTransformer3DModel` constructed with
``lora_enabled=True``.
Fused-QKV handling (the hard part):
  Wan attention uses a fused ``to_qkv`` (self-attn) and ``to_kv``
  (cross-attn) projection. The adapter ships separate ``to_q``,
  ``to_k``, ``to_v`` LoRA pairs. We combine them into a single LoRA
  pair on the fused module by:
    - stacking A along the rank dim:  A_fused = [A_q; A_k; A_v]  → [3r, in]
    - building a block-diagonal B with rank=3r and head-interleaving
      it the same way the base weight is head-interleaved:
        B_fused = head_interleave([B_q | 0 | 0,  0 | B_k | 0,  0 | 0 | B_v])
  This triples the LoRA matmul cost on those layers (rank=3r) but keeps
  the math exact and shippable in v0. A follow-up could split this into
  3 independent adapters at the layer level for ~3x lower cost.
`.diff` / `.diff_b` direct deltas are out of scope for v0 — the loader
warns and skips them.
"""
from __future__ import annotations

import re
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path

import torch
from loguru import logger
from safetensors.torch import load_file

from models.tt_dit.experimental.utils.lightx2v_loader import wan_lightx2v_to_diffusers_key

# Reused from pipeline_wan_lora.py — kept inline so this module is self-contained.
_STRIP_PREFIXES = ("diffusion_model.", "transformer.", "unet.", "model.")
_LOW_RANK_RE = re.compile(r"^(?P<base>.*)\.lora_(?P<slot>A|B|down|up)(?:\.[^.]+)?\.weight$")
_SLOT_MAP = {"A": "A", "down": "A", "B": "B", "up": "B"}


def _strip_known_prefixes(key: str) -> str:
    for prefix in _STRIP_PREFIXES:
        if key.startswith(prefix):
            return key[len(prefix) :]
    return key


def _kohya_to_lightx2v(key: str) -> str:
    if not key.startswith("lora_unet_"):
        return key
    parts = key.split(".", 1)
    module_path = parts[0][len("lora_unet_") :]
    suffix = f".{parts[1]}" if len(parts) > 1 else ""
    m = re.match(r"blocks_(\d+)_(cross_attn|self_attn)_([a-z]+)", module_path)
    if m:
        return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
    m = re.match(r"blocks_(\d+)_(ffn)_(\d+)", module_path)
    if m:
        return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
    logger.warning(f"unrecognized kohya key structure: {key}")
    return key


def _normalize_key(raw: str) -> str:
    return _kohya_to_lightx2v(_strip_known_prefixes(raw))


# --------------------------------------------------------------------
# Public API
# --------------------------------------------------------------------
@dataclass(frozen=True)
class AdapterHandle:
    name: str
    rank: int  # the "expected bound rank" — may be 3r for fused-QKV targets
    target_indices: dict[str, int]
    """Map from a stable target identifier (e.g. ``'blocks.0.attn1.to_qkv'``)
    to the bank index returned by `register_lora` on that module."""


def load_adapter_into(
    transformer,  # WanTransformer3DModel constructed with lora_enabled=True
    path: str,
    *,
    scale: float = 1.0,
    name: str = "",
) -> AdapterHandle:
    """Load a single LoRA safetensors file into `transformer`.
    Returns an `AdapterHandle` recording the bank index assigned to each
    LoRA-targeted Linear. Use it later with
    ``pipeline.set_active_lora(handle)``.
    """
    raw = load_file(str(path))
    if not any(_is_lora_key(k) for k in raw):
        raise RuntimeError(f"no LoRA-style keys (lora_A/lora_B, lora_down/lora_up) in {path}")

    pairs, alphas, skipped_direct = _collect_pairs(raw)
    if skipped_direct:
        logger.warning(
            f"{path}: skipping {skipped_direct} direct (.diff/.diff_b) deltas "
            "— runtime-LoRA pipeline does not support these yet."
        )

    # Group targets
    fused_qkv: dict[tuple[int, str], dict[str, dict[str, torch.Tensor]]] = {}
    singletons: list[tuple[str, dict[str, torch.Tensor]]] = []  # (diffusers_path, ab)
    skipped_unmapped = 0

    for base_path, ab in pairs.items():
        diff_path = _diffusers_path(base_path)
        if diff_path is None:
            skipped_unmapped += 1
            continue
        m = re.match(r"blocks\.(\d+)\.(attn1|attn2)\.to_([qkv])$", diff_path)
        if m:
            block_idx, attn, qkv = int(m.group(1)), m.group(2), m.group(3)
            # attn1 (self-attn): Q/K/V all fold into to_qkv.
            # attn2 (cross-attn): Q is its own singleton; K/V fold into to_kv.
            if attn == "attn1" or qkv in ("k", "v"):
                fused_qkv.setdefault((block_idx, attn), {})[qkv] = ab
                continue
            # attn2.to_q is a real singleton — fall through.
        # Everything else: 1:1 mapping
        singletons.append((diff_path, ab))

    if skipped_unmapped:
        logger.warning(
            f"{path}: {skipped_unmapped} pairs unmapped — adapter targets a module "
            "that doesn't exist in the tt-dit name map."
        )

    target_indices: dict[str, int] = {}
    fused_ranks: list[int] = []
    singleton_ranks: list[int] = []

    # Fused QKV / KV
    for (block_idx, attn_name), qkvs in fused_qkv.items():
        attn = getattr(transformer.blocks[block_idx], attn_name)
        bank_idx, fused_rank = _register_fused_qkv(attn, block_idx, qkvs, alphas, scale, name)
        if bank_idx is None:
            continue
        key = f"blocks.{block_idx}.{attn_name}.{'to_qkv' if attn.is_self else 'to_kv'}"
        target_indices[key] = bank_idx
        fused_ranks.append(fused_rank)

    # Singletons
    for diff_path, ab in singletons:
        target, canonical = _resolve_singleton(transformer, diff_path)
        if target is None:
            logger.warning(f"unmapped singleton LoRA target: {diff_path}")
            continue
        A = ab["A"]
        B = ab["B"]
        rank = A.shape[0]
        alpha = alphas.get(_lightx2v_for_path(diff_path), float(rank))
        eff_scale = scale * (alpha / rank)
        idx = target.register_lora(A, B, scale=eff_scale, name=name)
        target_indices[canonical] = idx
        singleton_ranks.append(rank)

    # Sanity: every LoRA module in this adapter should have the same
    # bound rank, but fused_QKV produces a 3x larger rank than singletons.
    # The bound-tensor system handles per-module ranks independently, so
    # we just record both for reporting.
    all_ranks = fused_ranks + singleton_ranks
    if not all_ranks:
        raise RuntimeError(f"{path}: no LoRA targets registered")
    representative_rank = max(all_ranks)
    return AdapterHandle(name=name or Path(path).stem, rank=representative_rank, target_indices=target_indices)


# --------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------
def _is_lora_key(raw: str) -> bool:
    key = _normalize_key(raw)
    return bool(_LOW_RANK_RE.match(key)) or key.endswith((".diff", ".diff_b"))


def _collect_pairs(raw: dict[str, torch.Tensor]):
    """Returns ({base_path: {'A': T, 'B': T}}, {base_path: alpha}, num_skipped_direct)."""
    pairs: dict[str, dict[str, torch.Tensor]] = {}
    alphas: dict[str, float] = {}
    skipped_direct = 0
    unrecognized: list[tuple[str, str]] = []

    for raw_key, tensor in raw.items():
        key = _normalize_key(raw_key)
        m = _LOW_RANK_RE.match(key)
        if m:
            pairs.setdefault(m.group("base"), {})[_SLOT_MAP[m.group("slot")]] = tensor
        elif key.endswith(".alpha"):
            alphas[key[: -len(".alpha")]] = float(tensor.item())
        elif key.endswith(".diff") or key.endswith(".diff_b"):
            skipped_direct += 1
        else:
            unrecognized.append((raw_key, key))

    # Surface dropped keys — a non-standard naming convention silently
    # producing plausible-but-wrong output is the failure mode we cannot
    # afford here.
    if unrecognized:
        sample = ", ".join(f"{raw}→{norm}" for raw, norm in unrecognized[:5])
        more = "" if len(unrecognized) <= 5 else f" (+{len(unrecognized) - 5} more)"
        logger.warning(
            f"adapter loader: {len(unrecognized)} key(s) did not match any known "
            f"pattern (lora_A/B, lora_down/up, .alpha, .diff, .diff_b) and were "
            f"dropped. Samples: {sample}{more}"
        )

    return pairs, alphas, skipped_direct


def _diffusers_path(lightx2v_base: str) -> str | None:
    """Map a lightx2v base path to its diffusers-canonical equivalent
    (without ``.weight``). Returns None if mapping fails."""
    try:
        key = wan_lightx2v_to_diffusers_key(f"{lightx2v_base}.weight")
    except Exception:  # noqa: BLE001
        return None
    if not key.endswith(".weight"):
        return None
    return key[: -len(".weight")]


def _lightx2v_for_path(diffusers_path: str) -> str:
    """Reverse-map (best-effort) a diffusers path back to lightx2v for alpha
    lookups — alphas are keyed in the source naming."""
    # Simple cases that the existing map covers:
    rev = {
        "attn1.to_q": "self_attn.q",
        "attn1.to_k": "self_attn.k",
        "attn1.to_v": "self_attn.v",
        "attn1.to_out.0": "self_attn.o",
        "attn2.to_q": "cross_attn.q",
        "attn2.to_k": "cross_attn.k",
        "attn2.to_v": "cross_attn.v",
        "attn2.to_out.0": "cross_attn.o",
        "ffn.net.0.proj": "ffn.0",
        "ffn.net.2": "ffn.2",
    }
    for diff_seg, lx_seg in rev.items():
        if diff_seg in diffusers_path:
            return diffusers_path.replace(diff_seg, lx_seg)
    return diffusers_path


def _resolve_singleton(transformer, diff_path: str):
    """Resolve a diffusers-style path to (module, tt_dit_dotted_path).
    Returns (None, None) when no match. The tt-dit path is the one
    consumed by the pipeline mixin's bind/unbind walker — it skips
    diffusers-only segments like the trailing ``.0`` in ``to_out.0``.
    """
    m = re.match(r"blocks\.(\d+)\.(attn1|attn2)\.to_q$", diff_path)
    if m:
        block_idx, attn_name = int(m.group(1)), m.group(2)
        attn = getattr(transformer.blocks[block_idx], attn_name)
        if not attn.is_self:
            return attn.to_q, f"blocks.{block_idx}.{attn_name}.to_q"
        return None, None  # self-attn to_q shouldn't arrive here
    m = re.match(r"blocks\.(\d+)\.(attn1|attn2)\.to_out\.0$", diff_path)
    if m:
        block_idx, attn_name = int(m.group(1)), m.group(2)
        return getattr(transformer.blocks[block_idx], attn_name).to_out, (f"blocks.{block_idx}.{attn_name}.to_out")
    m = re.match(r"blocks\.(\d+)\.ffn\.net\.0\.proj$", diff_path)
    if m:
        block_idx = int(m.group(1))
        return transformer.blocks[block_idx].ffn.ff1, f"blocks.{block_idx}.ffn.ff1"
    m = re.match(r"blocks\.(\d+)\.ffn\.net\.2$", diff_path)
    if m:
        block_idx = int(m.group(1))
        return transformer.blocks[block_idx].ffn.ff2, f"blocks.{block_idx}.ffn.ff2"
    return None, None


def _register_fused_qkv(
    attn,
    block_idx: int,
    qkvs: dict[str, dict[str, torch.Tensor]],
    alphas: dict[str, float],
    scale: float,
    name: str,
) -> tuple[int | None, int]:
    """Stack Q/K/V (or K/V for cross-attn) LoRA pairs onto attn.to_qkv (or
    attn.to_kv). Returns (bank_idx, fused_rank) or (None, 0) on skip."""
    if attn.is_self:
        required = ["q", "k", "v"]
        target = attn.to_qkv
    else:
        required = ["k", "v"]
        target = attn.to_kv

    missing = [r for r in required if r not in qkvs]
    if missing:
        logger.warning(
            f"fused QKV LoRA on {target.__class__.__name__} is missing {missing}; "
            "skipping this group (need all of "
            f"{required}, got {list(qkvs)})."
        )
        return None, 0

    A_per = [qkvs[r]["A"] for r in required]
    B_per = [qkvs[r]["B"] for r in required]
    ranks = [A.shape[0] for A in A_per]
    if len(set(ranks)) != 1:
        raise ValueError(f"QKV LoRA ranks must match; got {ranks}")
    r = ranks[0]
    n = len(required)
    in_dim = attn.dim
    out_per = attn.dim  # each of Q/K/V has out=dim

    # A stacked on rank dim → [n*r, in]
    A_fused = torch.cat(A_per, dim=0)
    assert A_fused.shape == (n * r, in_dim), f"A_fused {A_fused.shape}"

    # B: each per-source B is [out_per=dim, r]. We embed each into a
    # [out_per, n*r] tile (zeros for other-source rank columns). Then
    # head-interleave the list of [out_per, n*r] tiles into a single
    # [n*out_per, n*r] tensor — matching the head layout of the fused
    # to_qkv / to_kv base weight.
    n_dev = attn.parallel_config.tensor_parallel.factor
    n_local_heads = attn.n_local_heads
    head_dim = attn.head_dim

    B_per_padded: list[torch.Tensor] = []
    for i, B in enumerate(B_per):
        pad = torch.zeros(out_per, n * r, dtype=B.dtype)
        pad[:, i * r : (i + 1) * r] = B
        B_per_padded.append(pad)

    B_fused = _head_interleave_lora_B(B_per_padded, n_dev=n_dev, n_local_heads=n_local_heads, head_dim=head_dim)
    assert B_fused.shape == (n * out_per, n * r), f"B_fused {B_fused.shape}"

    # Effective scale: pick the first source's alpha (alphas are usually
    # identical across Q/K/V in a single adapter; if not, the choice is
    # safe-ish because we apply scale to all three the same way).
    alpha_keys = [_lightx2v_qkv_key(attn, block_idx, r_name) for r_name in required]
    alpha = next((alphas[k] for k in alpha_keys if k in alphas), float(r))
    eff_scale = scale * (alpha / r)

    bank_idx = target.register_lora(A_fused, B_fused, scale=eff_scale, name=name)
    return bank_idx, n * r


def _lightx2v_qkv_key(attn, block_idx: int, qkv: str) -> str:
    """Build the lightx2v alpha key for a Q/K/V projection inside a
    given WanAttention. Alphas in `_collect_pairs` are keyed by the full
    normalized base path (``blocks.<i>.<side>.<qkv>``), so the block
    index must be threaded in alongside ``attn.is_self``."""
    side = "self_attn" if attn.is_self else "cross_attn"
    return f"blocks.{block_idx}.{side}.{qkv}"


def _head_interleave_lora_B(
    tensors: Sequence[torch.Tensor],
    *,
    n_dev: int,
    n_local_heads: int,
    head_dim: int,
) -> torch.Tensor:
    """Mirror of WanAttention._interleave_heads for LoRA B tensors.
    Each input is [out_per_source=num_heads*head_dim, rank_fused].
    Returns a concatenated [N*out_per, rank_fused] tensor with the out
    dim head-interleaved so column-parallel sharding lines up with the
    fused base weight's layout.
    """
    out_per = n_dev * n_local_heads * head_dim
    fused_rank = tensors[0].shape[1]
    assert all(t.shape == (out_per, fused_rank) for t in tensors), (
        f"all B tensors must be [{out_per}, {fused_rank}]; got " f"{[tuple(t.shape) for t in tensors]}"
    )
    n = len(tensors)
    # Transpose each to [fused_rank, out_per]
    transposed = [t.T for t in tensors]
    # Reshape out dim to [fused_rank, n_dev, n_local_heads, head_dim]
    reshaped = [t.reshape(fused_rank, n_dev, n_local_heads, head_dim) for t in transposed]
    # Concat on heads sub-dim → [fused_rank, n_dev, N*n_local_heads, head_dim]
    merged = torch.cat(reshaped, dim=2)
    # Flatten back → [fused_rank, N*n_dev*n_local_heads*head_dim] = [fused_rank, N*out_per]
    merged = merged.reshape(fused_rank, n * out_per)
    # Transpose to [N*out_per, fused_rank]
    return merged.T

	# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
	# SPDX-License-Identifier: Apache-2.0

	"""
	Adapter loader for the runtime-LoRA Wan 2.2 pipeline.

	Reads a LoRA safetensors file, parses A/B pairs using the same key
	normalization as ``pipeline_wan_lora.py``, and uploads them into the
	LoRA-aware Linear modules of a `WanTransformer3DModel` constructed with
	``lora_enabled=True``.

	Fused-QKV handling (the hard part):
	Wan attention uses a fused ``to_qkv`` (self-attn) and ``to_kv``
	(cross-attn) projection. The adapter ships separate ``to_q``,
	``to_k``, ``to_v`` LoRA pairs. We combine them into a single LoRA
	pair on the fused module by:

	- stacking A along the rank dim: A_fused = [A_q; A_k; A_v] → [3r, in]
	- building a block-diagonal B with rank=3r and head-interleaving
	it the same way the base weight is head-interleaved:
	B_fused = head_interleave([B_q \| 0 \| 0, 0 \| B_k \| 0, 0 \| 0 \| B_v])

	This triples the LoRA matmul cost on those layers (rank=3r) but keeps
	the math exact and shippable in v0. A follow-up could split this into
	3 independent adapters at the layer level for ~3x lower cost.

	`.diff` / `.diff_b` direct deltas are out of scope for v0 — the loader
	warns and skips them.
	"""
	from __future__ import annotations

	import re
	from collections.abc import Sequence
	from dataclasses import dataclass
	from pathlib import Path

	import torch
	from loguru import logger
	from safetensors.torch import load_file

	from models.tt_dit.experimental.utils.lightx2v_loader import wan_lightx2v_to_diffusers_key

	# Reused from pipeline_wan_lora.py — kept inline so this module is self-contained.
	_STRIP_PREFIXES = ("diffusion_model.", "transformer.", "unet.", "model.")
	_LOW_RANK_RE = re.compile(r"^(?P<base>.*)\.lora_(?P<slot>A\|B\|down\|up)(?:\.[^.]+)?\.weight$")
	_SLOT_MAP = {"A": "A", "down": "A", "B": "B", "up": "B"}


	def _strip_known_prefixes(key: str) -> str:
	for prefix in _STRIP_PREFIXES:
	if key.startswith(prefix):
	return key[len(prefix) :]
	return key


	def _kohya_to_lightx2v(key: str) -> str:
	if not key.startswith("lora_unet_"):
	return key
	parts = key.split(".", 1)
	module_path = parts[0][len("lora_unet_") :]
	suffix = f".{parts[1]}" if len(parts) > 1 else ""
	m = re.match(r"blocks_(\d+)_(cross_attn\|self_attn)_([a-z]+)", module_path)
	if m:
	return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
	m = re.match(r"blocks_(\d+)_(ffn)_(\d+)", module_path)
	if m:
	return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
	logger.warning(f"unrecognized kohya key structure: {key}")
	return key


	def _normalize_key(raw: str) -> str:
	return _kohya_to_lightx2v(_strip_known_prefixes(raw))


	# --------------------------------------------------------------------
	# Public API
	# --------------------------------------------------------------------
	@dataclass(frozen=True)
	class AdapterHandle:
	name: str
	rank: int # the "expected bound rank" — may be 3r for fused-QKV targets
	target_indices: dict[str, int]
	"""Map from a stable target identifier (e.g. ``'blocks.0.attn1.to_qkv'``)
	to the bank index returned by `register_lora` on that module."""


	def load_adapter_into(
	transformer, # WanTransformer3DModel constructed with lora_enabled=True
	path: str,
	*,
	scale: float = 1.0,
	name: str = "",
	) -> AdapterHandle:
	"""Load a single LoRA safetensors file into `transformer`.

	Returns an `AdapterHandle` recording the bank index assigned to each
	LoRA-targeted Linear. Use it later with
	``pipeline.set_active_lora(handle)``.
	"""
	raw = load_file(str(path))
	if not any(_is_lora_key(k) for k in raw):
	raise RuntimeError(f"no LoRA-style keys (lora_A/lora_B, lora_down/lora_up) in {path}")

	pairs, alphas, skipped_direct = _collect_pairs(raw)
	if skipped_direct:
	logger.warning(
	f"{path}: skipping {skipped_direct} direct (.diff/.diff_b) deltas "
	"— runtime-LoRA pipeline does not support these yet."
	)

	# Group targets
	fused_qkv: dict[tuple[int, str], dict[str, dict[str, torch.Tensor]]] = {}
	singletons: list[tuple[str, dict[str, torch.Tensor]]] = [] # (diffusers_path, ab)
	skipped_unmapped = 0

	for base_path, ab in pairs.items():
	diff_path = _diffusers_path(base_path)
	if diff_path is None:
	skipped_unmapped += 1
	continue
	m = re.match(r"blocks\.(\d+)\.(attn1\|attn2)\.to_([qkv])$", diff_path)
	if m:
	block_idx, attn, qkv = int(m.group(1)), m.group(2), m.group(3)
	# attn1 (self-attn): Q/K/V all fold into to_qkv.
	# attn2 (cross-attn): Q is its own singleton; K/V fold into to_kv.
	if attn == "attn1" or qkv in ("k", "v"):
	fused_qkv.setdefault((block_idx, attn), {})[qkv] = ab
	continue
	# attn2.to_q is a real singleton — fall through.
	# Everything else: 1:1 mapping
	singletons.append((diff_path, ab))

	if skipped_unmapped:
	logger.warning(
	f"{path}: {skipped_unmapped} pairs unmapped — adapter targets a module "
	"that doesn't exist in the tt-dit name map."
	)

	target_indices: dict[str, int] = {}
	fused_ranks: list[int] = []
	singleton_ranks: list[int] = []

	# Fused QKV / KV
	for (block_idx, attn_name), qkvs in fused_qkv.items():
	attn = getattr(transformer.blocks[block_idx], attn_name)
	bank_idx, fused_rank = _register_fused_qkv(attn, block_idx, qkvs, alphas, scale, name)
	if bank_idx is None:
	continue
	key = f"blocks.{block_idx}.{attn_name}.{'to_qkv' if attn.is_self else 'to_kv'}"
	target_indices[key] = bank_idx
	fused_ranks.append(fused_rank)

	# Singletons
	for diff_path, ab in singletons:
	target, canonical = _resolve_singleton(transformer, diff_path)
	if target is None:
	logger.warning(f"unmapped singleton LoRA target: {diff_path}")
	continue
	A = ab["A"]
	B = ab["B"]
	rank = A.shape[0]
	alpha = alphas.get(_lightx2v_for_path(diff_path), float(rank))
	eff_scale = scale * (alpha / rank)
	idx = target.register_lora(A, B, scale=eff_scale, name=name)
	target_indices[canonical] = idx
	singleton_ranks.append(rank)

	# Sanity: every LoRA module in this adapter should have the same
	# bound rank, but fused_QKV produces a 3x larger rank than singletons.
	# The bound-tensor system handles per-module ranks independently, so
	# we just record both for reporting.
	all_ranks = fused_ranks + singleton_ranks
	if not all_ranks:
	raise RuntimeError(f"{path}: no LoRA targets registered")
	representative_rank = max(all_ranks)
	return AdapterHandle(name=name or Path(path).stem, rank=representative_rank, target_indices=target_indices)


	# --------------------------------------------------------------------
	# Helpers
	# --------------------------------------------------------------------
	def _is_lora_key(raw: str) -> bool:
	key = _normalize_key(raw)
	return bool(_LOW_RANK_RE.match(key)) or key.endswith((".diff", ".diff_b"))


	def _collect_pairs(raw: dict[str, torch.Tensor]):
	"""Returns ({base_path: {'A': T, 'B': T}}, {base_path: alpha}, num_skipped_direct)."""
	pairs: dict[str, dict[str, torch.Tensor]] = {}
	alphas: dict[str, float] = {}
	skipped_direct = 0
	unrecognized: list[tuple[str, str]] = []

	for raw_key, tensor in raw.items():
	key = _normalize_key(raw_key)
	m = _LOW_RANK_RE.match(key)
	if m:
	pairs.setdefault(m.group("base"), {})[_SLOT_MAP[m.group("slot")]] = tensor
	elif key.endswith(".alpha"):
	alphas[key[: -len(".alpha")]] = float(tensor.item())
	elif key.endswith(".diff") or key.endswith(".diff_b"):
	skipped_direct += 1
	else:
	unrecognized.append((raw_key, key))

	# Surface dropped keys — a non-standard naming convention silently
	# producing plausible-but-wrong output is the failure mode we cannot
	# afford here.
	if unrecognized:
	sample = ", ".join(f"{raw}→{norm}" for raw, norm in unrecognized[:5])
	more = "" if len(unrecognized) <= 5 else f" (+{len(unrecognized) - 5} more)"
	logger.warning(
	f"adapter loader: {len(unrecognized)} key(s) did not match any known "
	f"pattern (lora_A/B, lora_down/up, .alpha, .diff, .diff_b) and were "
	f"dropped. Samples: {sample}{more}"
	)

	return pairs, alphas, skipped_direct


	def _diffusers_path(lightx2v_base: str) -> str \| None:
	"""Map a lightx2v base path to its diffusers-canonical equivalent
	(without ``.weight``). Returns None if mapping fails."""
	try:
	key = wan_lightx2v_to_diffusers_key(f"{lightx2v_base}.weight")
	except Exception: # noqa: BLE001
	return None
	if not key.endswith(".weight"):
	return None
	return key[: -len(".weight")]


	def _lightx2v_for_path(diffusers_path: str) -> str:
	"""Reverse-map (best-effort) a diffusers path back to lightx2v for alpha
	lookups — alphas are keyed in the source naming."""
	# Simple cases that the existing map covers:
	rev = {
	"attn1.to_q": "self_attn.q",
	"attn1.to_k": "self_attn.k",
	"attn1.to_v": "self_attn.v",
	"attn1.to_out.0": "self_attn.o",
	"attn2.to_q": "cross_attn.q",
	"attn2.to_k": "cross_attn.k",
	"attn2.to_v": "cross_attn.v",
	"attn2.to_out.0": "cross_attn.o",
	"ffn.net.0.proj": "ffn.0",
	"ffn.net.2": "ffn.2",
	}
	for diff_seg, lx_seg in rev.items():
	if diff_seg in diffusers_path:
	return diffusers_path.replace(diff_seg, lx_seg)
	return diffusers_path


	def _resolve_singleton(transformer, diff_path: str):
	"""Resolve a diffusers-style path to (module, tt_dit_dotted_path).

	Returns (None, None) when no match. The tt-dit path is the one
	consumed by the pipeline mixin's bind/unbind walker — it skips
	diffusers-only segments like the trailing ``.0`` in ``to_out.0``.
	"""
	m = re.match(r"blocks\.(\d+)\.(attn1\|attn2)\.to_q$", diff_path)
	if m:
	block_idx, attn_name = int(m.group(1)), m.group(2)
	attn = getattr(transformer.blocks[block_idx], attn_name)
	if not attn.is_self:
	return attn.to_q, f"blocks.{block_idx}.{attn_name}.to_q"
	return None, None # self-attn to_q shouldn't arrive here
	m = re.match(r"blocks\.(\d+)\.(attn1\|attn2)\.to_out\.0$", diff_path)
	if m:
	block_idx, attn_name = int(m.group(1)), m.group(2)
	return getattr(transformer.blocks[block_idx], attn_name).to_out, (f"blocks.{block_idx}.{attn_name}.to_out")
	m = re.match(r"blocks\.(\d+)\.ffn\.net\.0\.proj$", diff_path)
	if m:
	block_idx = int(m.group(1))
	return transformer.blocks[block_idx].ffn.ff1, f"blocks.{block_idx}.ffn.ff1"
	m = re.match(r"blocks\.(\d+)\.ffn\.net\.2$", diff_path)
	if m:
	block_idx = int(m.group(1))
	return transformer.blocks[block_idx].ffn.ff2, f"blocks.{block_idx}.ffn.ff2"
	return None, None


	def _register_fused_qkv(
	attn,
	block_idx: int,
	qkvs: dict[str, dict[str, torch.Tensor]],
	alphas: dict[str, float],
	scale: float,
	name: str,
	) -> tuple[int \| None, int]:
	"""Stack Q/K/V (or K/V for cross-attn) LoRA pairs onto attn.to_qkv (or
	attn.to_kv). Returns (bank_idx, fused_rank) or (None, 0) on skip."""
	if attn.is_self:
	required = ["q", "k", "v"]
	target = attn.to_qkv
	else:
	required = ["k", "v"]
	target = attn.to_kv

	missing = [r for r in required if r not in qkvs]
	if missing:
	logger.warning(
	f"fused QKV LoRA on {target.__class__.__name__} is missing {missing}; "
	"skipping this group (need all of "
	f"{required}, got {list(qkvs)})."
	)
	return None, 0

	A_per = [qkvs[r]["A"] for r in required]
	B_per = [qkvs[r]["B"] for r in required]
	ranks = [A.shape[0] for A in A_per]
	if len(set(ranks)) != 1:
	raise ValueError(f"QKV LoRA ranks must match; got {ranks}")
	r = ranks[0]
	n = len(required)
	in_dim = attn.dim
	out_per = attn.dim # each of Q/K/V has out=dim

	# A stacked on rank dim → [n*r, in]
	A_fused = torch.cat(A_per, dim=0)
	assert A_fused.shape == (n * r, in_dim), f"A_fused {A_fused.shape}"

	# B: each per-source B is [out_per=dim, r]. We embed each into a
	# [out_per, n*r] tile (zeros for other-source rank columns). Then
	# head-interleave the list of [out_per, n*r] tiles into a single
	# [nout_per, nr] tensor — matching the head layout of the fused
	# to_qkv / to_kv base weight.
	n_dev = attn.parallel_config.tensor_parallel.factor
	n_local_heads = attn.n_local_heads
	head_dim = attn.head_dim

	B_per_padded: list[torch.Tensor] = []
	for i, B in enumerate(B_per):
	pad = torch.zeros(out_per, n * r, dtype=B.dtype)
	pad[:, i * r : (i + 1) * r] = B
	B_per_padded.append(pad)

	B_fused = _head_interleave_lora_B(B_per_padded, n_dev=n_dev, n_local_heads=n_local_heads, head_dim=head_dim)
	assert B_fused.shape == (n * out_per, n * r), f"B_fused {B_fused.shape}"

	# Effective scale: pick the first source's alpha (alphas are usually
	# identical across Q/K/V in a single adapter; if not, the choice is
	# safe-ish because we apply scale to all three the same way).
	alpha_keys = [_lightx2v_qkv_key(attn, block_idx, r_name) for r_name in required]
	alpha = next((alphas[k] for k in alpha_keys if k in alphas), float(r))
	eff_scale = scale * (alpha / r)

	bank_idx = target.register_lora(A_fused, B_fused, scale=eff_scale, name=name)
	return bank_idx, n * r


	def _lightx2v_qkv_key(attn, block_idx: int, qkv: str) -> str:
	"""Build the lightx2v alpha key for a Q/K/V projection inside a
	given WanAttention. Alphas in `_collect_pairs` are keyed by the full
	normalized base path (``blocks.<i>.<side>.<qkv>``), so the block
	index must be threaded in alongside ``attn.is_self``."""
	side = "self_attn" if attn.is_self else "cross_attn"
	return f"blocks.{block_idx}.{side}.{qkv}"


	def _head_interleave_lora_B(
	tensors: Sequence[torch.Tensor],
	*,
	n_dev: int,
	n_local_heads: int,
	head_dim: int,
	) -> torch.Tensor:
	"""Mirror of WanAttention._interleave_heads for LoRA B tensors.

	Each input is [out_per_source=num_heads*head_dim, rank_fused].
	Returns a concatenated [N*out_per, rank_fused] tensor with the out
	dim head-interleaved so column-parallel sharding lines up with the
	fused base weight's layout.
	"""
	out_per = n_dev * n_local_heads * head_dim
	fused_rank = tensors[0].shape[1]
	assert all(t.shape == (out_per, fused_rank) for t in tensors), (
	f"all B tensors must be [{out_per}, {fused_rank}]; got " f"{[tuple(t.shape) for t in tensors]}"
	)
	n = len(tensors)
	# Transpose each to [fused_rank, out_per]
	transposed = [t.T for t in tensors]
	# Reshape out dim to [fused_rank, n_dev, n_local_heads, head_dim]
	reshaped = [t.reshape(fused_rank, n_dev, n_local_heads, head_dim) for t in transposed]
	# Concat on heads sub-dim → [fused_rank, n_dev, N*n_local_heads, head_dim]
	merged = torch.cat(reshaped, dim=2)
	# Flatten back → [fused_rank, Nn_devn_local_headshead_dim] = [fused_rank, Nout_per]
	merged = merged.reshape(fused_rank, n * out_per)
	# Transpose to [N*out_per, fused_rank]
	return merged.T