nyanko7/sd-diffusers-webui

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 diffusers11import modules.safe as _12from safetensors.torch import load_file13 14 15class LoRAModule(torch.nn.Module):16    """17    replaces forward method of the original Linear, instead of replacing the original Linear module.18    """19 20    def __init__(21            self,22            lora_name,23            org_module: torch.nn.Module,24            multiplier=1.0,25            lora_dim=4,26            alpha=1,27    ):28        """if alpha == 0 or None, alpha is rank (no scaling)."""29        super().__init__()30        self.lora_name = lora_name31        self.lora_dim = lora_dim32 33        if org_module.__class__.__name__ == "Conv2d":34            in_dim = org_module.in_channels35            out_dim = org_module.out_channels36            self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)37            self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)38        else:39            in_dim = org_module.in_features40            out_dim = org_module.out_features41            self.lora_down = torch.nn.Linear(in_dim, lora_dim, bias=False)42            self.lora_up = torch.nn.Linear(lora_dim, out_dim, bias=False)43 44        if type(alpha) == torch.Tensor:45            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error46 47        alpha = lora_dim if alpha is None or alpha == 0 else alpha48        self.scale = alpha / self.lora_dim49        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える50 51        # same as microsoft's52        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))53        torch.nn.init.zeros_(self.lora_up.weight)54 55        self.multiplier = multiplier56        self.org_module = org_module  # remove in applying57        self.enable = False58 59    def resize(self, rank, alpha, multiplier):60        self.alpha = torch.tensor(alpha)61        self.multiplier = multiplier62        self.scale = alpha / rank63        if self.lora_down.__class__.__name__ == "Conv2d":64            in_dim = self.lora_down.in_channels65            out_dim = self.lora_up.out_channels66            self.lora_down = torch.nn.Conv2d(in_dim, rank, (1, 1), bias=False)67            self.lora_up = torch.nn.Conv2d(rank, out_dim, (1, 1), bias=False)68        else:69            in_dim = self.lora_down.in_features70            out_dim = self.lora_up.out_features71            self.lora_down = torch.nn.Linear(in_dim, rank, bias=False)72            self.lora_up = torch.nn.Linear(rank, out_dim, bias=False)73 74    def apply(self):75        if hasattr(self, "org_module"):76            self.org_forward = self.org_module.forward77            self.org_module.forward = self.forward78            del self.org_module79 80    def forward(self, x):81        if self.enable:82            return (83        self.org_forward(x)84        + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale85        )86        return self.org_forward(x)87 88 89class LoRANetwork(torch.nn.Module):90    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]91    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]92    LORA_PREFIX_UNET = "lora_unet"93    LORA_PREFIX_TEXT_ENCODER = "lora_te"94 95    def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1) -> None:96        super().__init__()97        self.multiplier = multiplier98        self.lora_dim = lora_dim99        self.alpha = alpha100 101        # create module instances102        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules):103            loras = []104            for name, module in root_module.named_modules():105                if module.__class__.__name__ in target_replace_modules:106                    for child_name, child_module in module.named_modules():107                        if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):108                            lora_name = prefix + "." + name + "." + child_name109                            lora_name = lora_name.replace(".", "_")110                            lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha,)111                            loras.append(lora)112            return loras113 114        if isinstance(text_encoder, list):115            self.text_encoder_loras = text_encoder116        else:117            self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)118            print(f"Create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")119            120        if diffusers.__version__ >= "0.15.0":121            LoRANetwork.UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]122    123        self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)124        print(f"Create LoRA for U-Net: {len(self.unet_loras)} modules.")125 126        self.weights_sd = None127 128        # assertion129        names = set()130        for lora in self.text_encoder_loras + self.unet_loras:131            assert (lora.lora_name not in names), f"duplicated lora name: {lora.lora_name}"132            names.add(lora.lora_name)133 134            lora.apply()135            self.add_module(lora.lora_name, lora)136 137    def reset(self):138        for lora in self.text_encoder_loras + self.unet_loras:139            lora.enable = False140 141    def load(self, file, scale):142 143        weights = None144        if os.path.splitext(file)[1] == ".safetensors":145            weights = load_file(file)146        else:147            weights = torch.load(file, map_location="cpu")148 149        if not weights:150            return151 152        network_alpha = None153        network_dim = None154        for key, value in weights.items():155            if network_alpha is None and "alpha" in key:156                network_alpha = value157            if network_dim is None and "lora_down" in key and len(value.size()) == 2:158                network_dim = value.size()[0]159 160        if network_alpha is None:161            network_alpha = network_dim162 163        weights_has_text_encoder = weights_has_unet = False164        weights_to_modify = []165 166        for key in weights.keys():167            if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):168                weights_has_text_encoder = True169 170            if key.startswith(LoRANetwork.LORA_PREFIX_UNET):171                weights_has_unet = True172 173        if weights_has_text_encoder:174            weights_to_modify += self.text_encoder_loras175 176        if weights_has_unet:177            weights_to_modify += self.unet_loras178 179        for lora in self.text_encoder_loras + self.unet_loras:180            lora.resize(network_dim, network_alpha, scale)181            if lora in weights_to_modify:182                lora.enable = True183 184        info = self.load_state_dict(weights, False)185        if len(info.unexpected_keys) > 0:186            print(f"Weights are loaded. Unexpected keys={info.unexpected_keys}")187