mirror of
https://github.com/vladmandic/automatic
synced 2026-08-25 22:20:46 +02:00
CalamitousFelicitousness
4c85a60c3e
create_module built each up/down module with the default constructor, which kaiming-initializes the parameter, then copied the stored weight over the whole thing. The init is thrown away every time and costs about four times the copy: 22.1ms per module against 2.5ms, or 5.8s against 0.7s over a 264-module lora, on every load. skip_init constructs on meta and materializes uninitialized, so the copy still fully defines the parameter. Dtype, device and values are unchanged, including the fp32 upcast of bf16 files that the copy performs.
86 lines
5.0 KiB
Python
86 lines
5.0 KiB
Python
import torch
|
|
import diffusers.models.lora as diffusers_lora
|
|
import modules.lora.lyco_helpers as lyco_helpers
|
|
import modules.lora.network as network
|
|
from modules import devices
|
|
|
|
|
|
class ModuleTypeLora(network.ModuleType):
|
|
def create_module(self, net: network.Network, weights: network.NetworkWeights):
|
|
if all(x in weights.w for x in ["lora_up.weight", "lora_down.weight"]):
|
|
return NetworkModuleLora(net, weights)
|
|
return None
|
|
|
|
|
|
class NetworkModuleLora(network.NetworkModule):
|
|
|
|
def __init__(self, net: network.Network, weights: network.NetworkWeights):
|
|
super().__init__(net, weights)
|
|
self.up_model = self.create_module(weights.w, "lora_up.weight")
|
|
self.down_model = self.create_module(weights.w, "lora_down.weight")
|
|
self.mid_model = self.create_module(weights.w, "lora_mid.weight", none_ok=True)
|
|
self.dim = weights.w["lora_down.weight"].shape[0]
|
|
# Optional bias delta (diff_b) applied as ex_bias alongside the weight update.
|
|
self.ex_bias = weights.w.get("diff_b")
|
|
|
|
def create_module(self, weights, key, none_ok=False):
|
|
weight = weights.get(key)
|
|
if weight is None and none_ok:
|
|
return None
|
|
linear_modules = [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention, diffusers_lora.LoRACompatibleLinear]
|
|
typ = type(self.sd_module)
|
|
is_linear = typ in linear_modules or self.sd_module.__class__.__name__ in ["SDNQLinear", "QLinear", "Linear4bit"]
|
|
# Embedding weight delta is up@down over the [vocab, dim] table, same shape and merge
|
|
# as a Linear. isinstance also catches the SDNQEmbedding / ScaledWordEmbedding subclasses.
|
|
is_embedding = isinstance(self.sd_module, torch.nn.Embedding)
|
|
is_conv = (typ in [torch.nn.Conv2d, diffusers_lora.LoRACompatibleConv]) or (self.sd_module.__class__.__name__ in ["SDNQConv2d", "QConv2d"]) or (typ.__name__ in ['downsampler_block', 'upsampler_block'])
|
|
# skip_init: the stored weight is copied over the whole parameter below, so the
|
|
# default random init is thrown away. It costs ~4x the copy on every module.
|
|
if is_linear or is_embedding:
|
|
weight = weight.reshape(weight.shape[0], -1)
|
|
module = torch.nn.utils.skip_init(torch.nn.Linear, weight.shape[1], weight.shape[0], bias=False)
|
|
elif is_conv and (key == "lora_down.weight" or key == "dyn_up"):
|
|
if len(weight.shape) == 2:
|
|
weight = weight.reshape(weight.shape[0], -1, 1, 1)
|
|
if weight.shape[2] != 1 or weight.shape[3] != 1:
|
|
module = torch.nn.utils.skip_init(torch.nn.Conv2d, weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False)
|
|
else:
|
|
module = torch.nn.utils.skip_init(torch.nn.Conv2d, weight.shape[1], weight.shape[0], (1, 1), bias=False)
|
|
elif is_conv and (key == "lora_mid.weight"):
|
|
module = torch.nn.utils.skip_init(torch.nn.Conv2d, weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False)
|
|
elif is_conv and (key == "lora_up.weight" or key == "dyn_down"):
|
|
module = torch.nn.utils.skip_init(torch.nn.Conv2d, weight.shape[1], weight.shape[0], (1, 1), bias=False)
|
|
else:
|
|
raise AssertionError(f'Lora unsupported: key={key} layer={self.network_key} type={typ.__name__}')
|
|
with torch.no_grad():
|
|
if weight.shape != module.weight.shape:
|
|
weight = weight.reshape(module.weight.shape)
|
|
module.weight.copy_(weight)
|
|
module.weight.requires_grad_(False)
|
|
return module
|
|
|
|
def calc_updown(self, target): # pylint: disable=W0237
|
|
target_dtype = target.dtype if target.dtype != torch.uint8 else self.up_model.weight.dtype
|
|
up = self.up_model.weight.to(target.device, dtype=target_dtype)
|
|
down = self.down_model.weight.to(target.device, dtype=target_dtype)
|
|
output_shape = [up.size(0), down.size(1)]
|
|
if self.mid_model is not None:
|
|
mid = self.mid_model.weight.to(target.device, dtype=target_dtype)
|
|
updown = lyco_helpers.rebuild_cp_decomposition(up, down, mid) # cp-decomposition
|
|
output_shape += mid.shape[2:]
|
|
else:
|
|
mid = None
|
|
if len(down.shape) == 4:
|
|
output_shape += down.shape[2:]
|
|
updown = lyco_helpers.rebuild_conventional(up, down, output_shape, self.network.dyn_dim)
|
|
ex_bias = self.ex_bias.to(target.device, dtype=target_dtype) if self.ex_bias is not None else None
|
|
del up, down, mid
|
|
return self.finalize_updown(updown, target, output_shape, ex_bias=ex_bias)
|
|
|
|
def forward(self, x, y):
|
|
self.up_model.to(device=devices.device)
|
|
self.down_model.to(device=devices.device)
|
|
if hasattr(y, "scale"):
|
|
return y(scale=1) + self.up_model(self.down_model(x)) * self.multiplier() * self.calc_scale()
|
|
return y + self.up_model(self.down_model(x)) * self.multiplier() * self.calc_scale()
|