Lightxr/sd-diffusers-webui
1
1# LoRA network module2# reference:3# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py4# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py5# https://github.com/bmaltais/kohya_ss/blob/master/networks/lora.py#L486 7import math8import os9import torch10import modules.safe as _11from safetensors.torch import load_file12 13 14class LoRAModule(torch.nn.Module):15 """16 replaces forward method of the original Linear, instead of replacing the original Linear module.17 """18 19 def __init__(20 self,21 lora_name,22 org_module: torch.nn.Module,23 multiplier=1.0,24 lora_dim=4,25 alpha=1,26 ):27 """if alpha == 0 or None, alpha is rank (no scaling)."""28 super().__init__()29 self.lora_name = lora_name30 self.lora_dim = lora_dim31 32 if org_module.__class__.__name__ == "Conv2d":33 in_dim = org_module.in_channels34 out_dim = org_module.out_channels35 self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)36 self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)37 else:38 in_dim = org_module.in_features39 out_dim = org_module.out_features40 self.lora_down = torch.nn.Linear(in_dim, lora_dim, bias=False)41 self.lora_up = torch.nn.Linear(lora_dim, out_dim, bias=False)42 43 if type(alpha) == torch.Tensor:44 alpha = alpha.detach().float().numpy() # without casting, bf16 causes error45 46 alpha = lora_dim if alpha is None or alpha == 0 else alpha47 self.scale = alpha / self.lora_dim48 self.register_buffer("alpha", torch.tensor(alpha)) # 定数として扱える49 50 # same as microsoft's51 torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))52 torch.nn.init.zeros_(self.lora_up.weight)53 54 self.multiplier = multiplier55 self.org_module = org_module # remove in applying56 self.enable = False57 58 def resize(self, rank, alpha, multiplier):59 self.alpha = torch.tensor(alpha)60 self.multiplier = multiplier61 self.scale = alpha / rank62 if self.lora_down.__class__.__name__ == "Conv2d":63 in_dim = self.lora_down.in_channels64 out_dim = self.lora_up.out_channels65 self.lora_down = torch.nn.Conv2d(in_dim, rank, (1, 1), bias=False)66 self.lora_up = torch.nn.Conv2d(rank, out_dim, (1, 1), bias=False)67 else:68 in_dim = self.lora_down.in_features69 out_dim = self.lora_up.out_features70 self.lora_down = torch.nn.Linear(in_dim, rank, bias=False)71 self.lora_up = torch.nn.Linear(rank, out_dim, bias=False)72 73 def apply(self):74 if hasattr(self, "org_module"):75 self.org_forward = self.org_module.forward76 self.org_module.forward = self.forward77 del self.org_module78 79 def forward(self, x):80 if self.enable:81 return (82 self.org_forward(x)83 + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale84 )85 return self.org_forward(x)86 87 88class LoRANetwork(torch.nn.Module):89 UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]90 TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]91 LORA_PREFIX_UNET = "lora_unet"92 LORA_PREFIX_TEXT_ENCODER = "lora_te"93 94 def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1) -> None:95 super().__init__()96 self.multiplier = multiplier97 self.lora_dim = lora_dim98 self.alpha = alpha99 100 # create module instances101 def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules):102 loras = []103 for name, module in root_module.named_modules():104 if module.__class__.__name__ in target_replace_modules:105 for child_name, child_module in module.named_modules():106 if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):107 lora_name = prefix + "." + name + "." + child_name108 lora_name = lora_name.replace(".", "_")109 lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha,)110 loras.append(lora)111 return loras112 113 if isinstance(text_encoder, list):114 self.text_encoder_loras = text_encoder115 else:116 self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)117 print(f"Create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")118 119 self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)120 print(f"Create LoRA for U-Net: {len(self.unet_loras)} modules.")121 122 self.weights_sd = None123 124 # assertion125 names = set()126 for lora in self.text_encoder_loras + self.unet_loras:127 assert (lora.lora_name not in names), f"duplicated lora name: {lora.lora_name}"128 names.add(lora.lora_name)129 130 lora.apply()131 self.add_module(lora.lora_name, lora)132 133 def reset(self):134 for lora in self.text_encoder_loras + self.unet_loras:135 lora.enable = False136 137 def load(self, file, scale):138 139 weights = None140 if os.path.splitext(file)[1] == ".safetensors":141 weights = load_file(file)142 else:143 weights = torch.load(file, map_location="cpu")144 145 if not weights:146 return147 148 network_alpha = None149 network_dim = None150 for key, value in weights.items():151 if network_alpha is None and "alpha" in key:152 network_alpha = value153 if network_dim is None and "lora_down" in key and len(value.size()) == 2:154 network_dim = value.size()[0]155 156 if network_alpha is None:157 network_alpha = network_dim158 159 weights_has_text_encoder = weights_has_unet = False160 weights_to_modify = []161 162 for key in weights.keys():163 if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):164 weights_has_text_encoder = True165 166 if key.startswith(LoRANetwork.LORA_PREFIX_UNET):167 weights_has_unet = True168 169 if weights_has_text_encoder:170 weights_to_modify += self.text_encoder_loras171 172 if weights_has_unet:173 weights_to_modify += self.unet_loras174 175 for lora in self.text_encoder_loras + self.unet_loras:176 lora.resize(network_dim, network_alpha, scale)177 if lora in weights_to_modify:178 lora.enable = True179 180 info = self.load_state_dict(weights, False)181 if len(info.unexpected_keys) > 0:182 print(f"Weights are loaded. Unexpected keys={info.unexpected_keys}")183