95 lines
4.2 KiB
Python
95 lines
4.2 KiB
Python
"""Resident dynamic LoRA branches for the quantized H3 denoiser and refiner."""
|
|
|
|
from collections.abc import Iterator
|
|
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
from safetensors import safe_open
|
|
from torch import nn
|
|
|
|
|
|
class LoraBranch(nn.Module):
|
|
def __init__(self, down: torch.Tensor, up: torch.Tensor, alpha: float):
|
|
super().__init__()
|
|
if down.ndim != 2 or up.ndim != 2 or down.shape[0] != up.shape[1]:
|
|
raise ValueError(f"Invalid LoRA shapes: down={tuple(down.shape)}, up={tuple(up.shape)}")
|
|
self.scale = float(alpha) / down.shape[0]
|
|
self.register_buffer("down", down.contiguous(), persistent=False)
|
|
self.register_buffer("up", up.contiguous(), persistent=False)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = x.to(self.down.dtype)
|
|
return functional.linear(functional.linear(x, self.down), self.up) * self.scale
|
|
|
|
|
|
class DynamicLoraMixin:
|
|
"""Small mixin for linears that can host resident request-selectable LoRAs."""
|
|
|
|
def _init_dynamic_lora(self) -> None:
|
|
self.lora_branches = nn.ModuleDict()
|
|
self.active_lora: str | None = None
|
|
self.lora_strength = 0.0
|
|
|
|
def add_lora(self, name: str, down: torch.Tensor, up: torch.Tensor, alpha: float) -> None:
|
|
if name in self.lora_branches:
|
|
raise ValueError(f"LoRA {name!r} is already attached")
|
|
if down.shape[1] != self.in_features or up.shape[0] != self.out_features:
|
|
raise ValueError(
|
|
f"LoRA {name!r} dimensions {tuple(down.shape)}, {tuple(up.shape)} do not match "
|
|
f"linear [{self.out_features}, {self.in_features}]"
|
|
)
|
|
self.lora_branches[name] = LoraBranch(down, up, alpha)
|
|
|
|
def set_lora(self, name: str | None, strength: float = 1.0) -> None:
|
|
if name is not None and name not in self.lora_branches:
|
|
raise ValueError(f"LoRA {name!r} is not attached")
|
|
self.active_lora = name
|
|
self.lora_strength = float(strength) if name is not None else 0.0
|
|
|
|
def _apply_lora(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
|
|
if self.active_lora is None or self.lora_strength == 0.0:
|
|
return base
|
|
delta = self.lora_branches[self.active_lora](x)
|
|
return base + delta.to(base.dtype) * self.lora_strength
|
|
|
|
|
|
def iter_lora_targets(model: nn.Module, refiner: nn.Module) -> Iterator[tuple[str, DynamicLoraMixin]]:
|
|
for index, block in enumerate(model.backbone.blocks):
|
|
yield f"blocks.{index}.attn.qkv_proj", block.attention.qkv_proj
|
|
yield f"blocks.{index}.attn.out_proj", block.attention.out_proj
|
|
yield f"blocks.{index}.mlp.fc1", block.mlp.fc1
|
|
yield f"blocks.{index}.mlp.fc2", block.mlp.fc2
|
|
for index, block in enumerate(refiner.blocks):
|
|
yield f"token_refiner.blocks.{index}.attn.qkv_proj", block.qkv
|
|
yield f"token_refiner.blocks.{index}.attn.out_proj", block.out
|
|
yield f"token_refiner.blocks.{index}.mlp.fc1", block.fc1
|
|
yield f"token_refiner.blocks.{index}.mlp.fc2", block.fc2
|
|
|
|
|
|
def load_lora_adapter(model: nn.Module, refiner: nn.Module, name: str, path: str, device: str) -> int:
|
|
targets = list(iter_lora_targets(model, refiner))
|
|
expected = {
|
|
f"diffusion_model.{target}.{suffix}"
|
|
for target, _module in targets
|
|
for suffix in ("alpha", "lora_A.weight", "lora_B.weight")
|
|
}
|
|
with safe_open(path, framework="pt", device=device) as checkpoint:
|
|
actual = set(checkpoint.keys())
|
|
if actual != expected:
|
|
missing = sorted(expected - actual)[:8]
|
|
unexpected = sorted(actual - expected)[:8]
|
|
raise ValueError(f"LoRA key mismatch: missing={missing}, unexpected={unexpected}")
|
|
for target, module in targets:
|
|
prefix = f"diffusion_model.{target}"
|
|
module.add_lora(
|
|
name,
|
|
checkpoint.get_tensor(f"{prefix}.lora_A.weight"),
|
|
checkpoint.get_tensor(f"{prefix}.lora_B.weight"),
|
|
checkpoint.get_tensor(f"{prefix}.alpha").item(),
|
|
)
|
|
return len(targets)
|
|
|
|
|
|
def set_active_lora(model: nn.Module, refiner: nn.Module, name: str | None, strength: float = 1.0) -> None:
|
|
for _target, module in iter_lora_targets(model, refiner):
|
|
module.set_lora(name, strength)
|