File size: 3,412 Bytes
629b123 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | # SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""Fuse LoRA deltas into LTX transformer state dicts (bf16 base weights only).
FP8 / scaled-FP8 base weights are intentionally unsupported — we do not use
them on device.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
from loguru import logger
from safetensors.torch import load_file
@dataclass(frozen=True)
class LoraSpec:
"""A LoRA safetensors file with a fuse strength."""
path: str
strength: float = 1.0
def _strip_prefix(state_dict: dict[str, torch.Tensor], prefix: str) -> dict[str, torch.Tensor]:
"""Return ``{k.removeprefix(prefix): v}`` if any key starts with ``prefix``, else input."""
if any(k.startswith(prefix) for k in state_dict):
return {k[len(prefix) :] if k.startswith(prefix) else k: v for k, v in state_dict.items()}
return state_dict
def _normalize_lora_keys(lora_sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
"""Strip the diffusion-model prefix so LoRA keys align with the base keyspace.
HF LoRAs ship as ``model.diffusion_model.<rest>`` or the shorter
``diffusion_model.<rest>``; try both (longest first) so either becomes
``<rest>.lora_A.weight``.
"""
for prefix in ("model.diffusion_model.", "diffusion_model."):
if any(k.startswith(prefix) for k in lora_sd):
return _strip_prefix(lora_sd, prefix)
return lora_sd
def fuse_loras_into(
base_sd: dict[str, torch.Tensor],
loras: list[LoraSpec],
) -> dict[str, torch.Tensor]:
"""Return a new state dict with ``W += sum_i strength_i * (B_i @ A_i)``.
For each base ``<prefix>.weight`` we accumulate matching
``<prefix>.lora_A/B.weight`` deltas in fp32 and cast back to the base dtype;
keys without a LoRA pair pass through unchanged. A fuse count of 0 is logged
as a warning since it means the LoRA keyspace did not align with the base.
"""
if not loras:
return base_sd
lora_sds = [(_normalize_lora_keys(load_file(s.path)), s.strength, s.path) for s in loras]
fused: dict[str, torch.Tensor] = {}
fuse_counts = [0] * len(lora_sds)
for key, w in base_sd.items():
if not key.endswith(".weight"):
fused[key] = w
continue
prefix = key[: -len(".weight")]
a_key = f"{prefix}.lora_A.weight"
b_key = f"{prefix}.lora_B.weight"
delta: torch.Tensor | None = None
for idx, (lsd, strength, _path) in enumerate(lora_sds):
if a_key not in lsd or b_key not in lsd:
continue
a = lsd[a_key].to(dtype=torch.float32)
b = lsd[b_key].to(dtype=torch.float32)
contrib = torch.matmul(b * strength, a)
delta = contrib if delta is None else delta + contrib
fuse_counts[idx] += 1
if delta is None:
fused[key] = w
else:
fused[key] = (w.to(torch.float32) + delta).to(w.dtype)
for (_lsd, _strength, path), count in zip(lora_sds, fuse_counts):
if count == 0:
logger.warning(
f"LoRA {path}: 0 weights fused — prefix mismatch? "
f"Inspect a few base/lora keys to confirm alignment."
)
else:
logger.info(f"LoRA {path}: fused {count} weights (strength={_strength})")
return fused
|