lycoris-lora 3.0.0.dev9__py3-none-any.whl

lycoris/__init__.py

@@ -0,0 +1,26 @@
+try:
+    from . import kohya
+except:
+    pass
+from . import (
+    modules,
+    utils,
+)
+
+from .modules.locon import LoConModule
+from .modules.loha import LohaModule
+from .modules.lokr import LokrModule
+from .modules.dylora import DyLoraModule
+from .modules.glora import GLoRAModule
+from .modules.norms import NormModule
+from .modules.full import FullModule
+from .modules.diag_oft import DiagOFTModule
+from .modules import make_module
+
+from .wrapper import (
+    LycorisNetwork,
+    create_lycoris,
+    create_lycoris_from_weights,
+)
+
+from .logging import logger

lycoris/concept.py

@@ -0,0 +1,254 @@
+import os, sys
+sys.path.insert(0, os.getcwd())
+import math
+from warnings import warn
+import os
+from typing import List
+import torch
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .kohya_utils import *
+
+
+class ConceptModule(nn.Module):
+    def __init__(
+        self, 
+        lora_name, org_module: nn.Module, 
+        multiplier=1.0,
+    ):
+        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        super().__init__()
+        self.lora_name = lora_name
+
+        if org_module.__class__.__name__ == 'Conv2d':
+            # For general LoCon
+            stride = org_module.stride
+            padding = org_module.padding
+            self.op = F.conv2d
+            self.extra_args = {
+                'stride': stride,
+                'padding': padding
+            }
+        else:
+            self.op = F.linear
+            self.extra_args = {}
+        
+        self.concept_mask = nn.Parameter(torch.ones_like(
+            org_module.weight
+        ))
+        self.concept_neuron = None
+        
+        self.multiplier = multiplier
+        self.org_module = [org_module]
+
+    def apply_to(self):
+        self.org_module[0].forward = self.forward
+
+    def forward(self, x):
+        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
+        if torch.max(self.concept_mask)-torch.min(self.concept_mask)<0.1:
+            return self.op(
+                x, self.org_module[0].weight.data * self.concept_mask,
+                bias, **self.extra_args
+            )
+        else:
+            if self.concept_neuron is None:
+                print('Found neurons')
+                mask = self.concept_mask
+                threshold = torch.quantile(mask, 0.05)
+                neuron = mask<threshold
+                self.concept_mask.requires_grad_(False)
+                self.concept_mask[mask>=threshold] = 0
+                self.concept_mask[neuron] = mask[neuron]-1
+                self.concept_mask.requires_grad_(True)
+                self.concept_neuron = neuron
+            return self.op(
+                x, self.org_module[0].weight.data * (1+self.concept_mask * self.concept_neuron),
+                bias, **self.extra_args
+            )
+
+
+class ConceptNetwork(torch.nn.Module):
+    '''
+    LoRA + LoCon
+    '''
+    # Ignore proj_in or proj_out, their channels is only a few.
+    UNET_TARGET_REPLACE_MODULE = [
+        "Transformer2DModel", 
+        "Attention",
+    ]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention"]
+    LORA_PREFIX_UNET = 'lora_unet'
+    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+
+    def __init__(
+        self, 
+        text_encoder, unet, 
+        multiplier=1.0,
+        *args,
+        **kwargs
+    ) -> None:
+        super().__init__()
+        self.multiplier = multiplier
+        
+        total_params = 0
+        # create module instances
+        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[ConceptModule]:
+            nonlocal total_params
+            print('Create Concept Module')
+            loras = []
+            for name, module in root_module.named_modules():
+                if module.__class__.__name__ in target_replace_modules:
+                    for child_name, child_module in module.named_modules():
+                        if child_name.rsplit('.', 1)[-1] not in ['to_k', 'to_v', 'k_proj', 'v_proj']:
+                            continue
+                        lora_name = prefix + '.' + name + '.' + child_name
+                        lora_name = lora_name.replace('.', '_')
+                        total_params += child_module.weight.size().numel()
+                        lora = ConceptModule(
+                            lora_name, child_module, self.multiplier, 
+                        )
+                        loras.append(lora)
+            return loras
+
+        self.text_encoder_loras = create_modules(
+            ConceptNetwork.LORA_PREFIX_TEXT_ENCODER,
+            text_encoder, 
+            ConceptNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
+        )
+        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+        self.unet_loras = create_modules(ConceptNetwork.LORA_PREFIX_UNET, unet, ConceptNetwork.UNET_TARGET_REPLACE_MODULE)
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+        self.weights_sd = None
+        
+        print(f'Total concept param: {total_params}')
+        
+        # assertion
+        names = set()
+        for lora in self.text_encoder_loras + self.unet_loras:
+            assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+            names.add(lora.lora_name)
+
+    def set_multiplier(self, multiplier):
+        self.multiplier = multiplier
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.multiplier = self.multiplier
+            
+    def load_weights(self, file):
+        if os.path.splitext(file)[1] == '.safetensors':
+            from safetensors.torch import load_file, safe_open
+            self.weights_sd = load_file(file)
+        else:
+            self.weights_sd = torch.load(file, map_location='cpu')
+
+    def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None):
+        if self.weights_sd:
+            weights_has_text_encoder = weights_has_unet = False
+            for key in self.weights_sd.keys():
+                if key.startswith(ConceptNetwork.LORA_PREFIX_TEXT_ENCODER):
+                    weights_has_text_encoder = True
+                elif key.startswith(ConceptNetwork.LORA_PREFIX_UNET):
+                    weights_has_unet = True
+
+            if apply_text_encoder is None:
+                apply_text_encoder = weights_has_text_encoder
+            else:
+                assert apply_text_encoder == weights_has_text_encoder, f"text encoder weights: {weights_has_text_encoder} but text encoder flag: {apply_text_encoder} / 重みとText Encoderのフラグが矛盾しています"
+
+            if apply_unet is None:
+                apply_unet = weights_has_unet
+            else:
+                assert apply_unet == weights_has_unet, f"u-net weights: {weights_has_unet} but u-net flag: {apply_unet} / 重みとU-Netのフラグが矛盾しています"
+        else:
+            assert apply_text_encoder is not None and apply_unet is not None, f"internal error: flag not set"
+
+        if apply_text_encoder:
+            print("enable LoRA for text encoder")
+        else:
+            self.text_encoder_loras = []
+
+        if apply_unet:
+            print("enable LoRA for U-Net")
+        else:
+            self.unet_loras = []
+
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.apply_to()
+            self.add_module(lora.lora_name, lora)
+
+        if self.weights_sd:
+            # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
+            info = self.load_state_dict(self.weights_sd, False)
+            print(f"weights are loaded: {info}")
+
+    def enable_gradient_checkpointing(self):
+        # not supported
+        def make_ckpt(module):
+            if isinstance(module, torch.nn.Module):
+                module.grad_ckpt = True
+        self.apply(make_ckpt)
+        pass
+
+    def prepare_optimizer_params(self, text_encoder_lr, unet_lr):
+        def enumerate_params(loras):
+            params = []
+            for lora in loras:
+                params.extend(lora.parameters())
+            return params
+
+        self.requires_grad_(True)
+        all_params = []
+
+        if self.text_encoder_loras:
+            param_data = {'params': enumerate_params(self.text_encoder_loras)}
+            if text_encoder_lr is not None:
+                param_data['lr'] = text_encoder_lr
+            all_params.append(param_data)
+
+        if self.unet_loras:
+            param_data = {'params': enumerate_params(self.unet_loras)}
+            if unet_lr is not None:
+                param_data['lr'] = unet_lr
+            all_params.append(param_data)
+
+        return all_params
+
+    def prepare_grad_etc(self, text_encoder, unet):
+        self.requires_grad_(True)
+
+    def on_epoch_start(self, text_encoder, unet):
+        self.train()
+
+    def get_trainable_params(self):
+        return self.parameters()
+
+    def save_weights(self, file, dtype, metadata):
+        if metadata is not None and len(metadata) == 0:
+            metadata = None
+
+        state_dict = self.state_dict()
+
+        if dtype is not None:
+            for key in list(state_dict.keys()):
+                v = state_dict[key]
+                v = v.detach().clone().to("cpu").to(dtype)
+                state_dict[key] = v
+
+        if os.path.splitext(file)[1] == '.safetensors':
+            from safetensors.torch import save_file
+
+            # Precalculate model hashes to save time on indexing
+            if metadata is None:
+                metadata = {}
+            model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+            metadata["sshs_model_hash"] = model_hash
+            metadata["sshs_legacy_hash"] = legacy_hash
+
+            save_file(state_dict, file, metadata)
+        else:
+            torch.save(state_dict, file)

lycoris/config.py

@@ -0,0 +1,86 @@
+PRESET = {
+    "full": {
+        "enable_conv": True,
+        "unet_target_module": [
+            "Transformer2DModel",
+            "ResnetBlock2D",
+            "Downsample2D",
+            "Upsample2D",
+        ],
+        "unet_target_name": [
+            "conv_in",
+            "conv_out",
+            "time_embedding.linear_1",
+            "time_embedding.linear_2",
+        ],
+        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_name": [],
+    },
+    "full-lin": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "Transformer2DModel",
+            "ResnetBlock2D",
+        ],
+        "unet_target_name": [
+            "time_embedding.linear_1",
+            "time_embedding.linear_2",
+        ],
+        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_name": [],
+    },
+    "attn-mlp": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "Transformer2DModel",
+        ],
+        "unet_target_name": [],
+        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_name": [],
+    },
+    "attn-only": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "CrossAttention",
+        ],
+        "unet_target_name": [],
+        "text_encoder_target_module": ["CLIPAttention"],
+        "text_encoder_target_name": [],
+    },
+    "unet-only": {
+        "enable_conv": True,
+        "unet_target_module": [
+            "Transformer2DModel",
+            "ResnetBlock2D",
+            "Downsample2D",
+            "Upsample2D",
+        ],
+        "unet_target_name": [
+            "conv_in",
+            "conv_out",
+            "time_embedding.linear_1",
+            "time_embedding.linear_2",
+        ],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": [],
+    },
+    "unet-transformer-only": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "Transformer2DModel",
+        ],
+        "unet_target_name": [],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": [],
+    },
+    "unet-convblock-only": {
+        "enable_conv": True,
+        "unet_target_module": ["ResnetBlock2D", "Downsample2D", "Upsample2D"],
+        "unet_target_name": [
+            "conv_in",
+            "conv_out",
+        ],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": [],
+    },
+}

lycoris/dylora.py

@@ -0,0 +1,175 @@
+import math
+import random
+
+from collections import OrderedDict, abc as container_abcs
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DyLoraModule(nn.Module):
+    """
+    Hadamard product Implementaion for Dynamic Low Rank adaptation
+    """
+
+    def __init__(
+        self, 
+        lora_name, 
+        org_module: nn.Module, 
+        multiplier=1.0, 
+        lora_dim=4, alpha=1, 
+        dropout=0.,
+        use_cp=False, 
+        block_size=1,
+        **kwargs,
+    ):
+        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+        assert lora_dim % block_size == 0, 'lora_dim must be a multiple of block_size'
+        self.block_count = lora_dim//block_size
+        self.block_size = block_size
+        
+        self.shape = org_module.weight.shape
+        if org_module.__class__.__name__ == 'Conv2d':
+            in_dim = org_module.in_channels
+            k_size = org_module.kernel_size
+            out_dim = org_module.out_channels
+            shape = (out_dim, in_dim*k_size[0]*k_size[1])
+            self.op = F.conv2d
+            self.extra_args = {
+                "stride": org_module.stride,
+                "padding": org_module.padding,
+                "dilation": org_module.dilation,
+                "groups": org_module.groups
+            }
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            shape = (out_dim, in_dim)
+            self.op = F.linear
+            self.extra_args = {}
+        
+        self.lora_dim = lora_dim
+        self.up_list = nn.ParameterList([
+            torch.empty(shape[0], 1)
+            for i in range(lora_dim)
+        ])
+        self.up_list.requires_grad_(False)
+        self.up_update = [
+            torch.zeros_like(self.up_list[i])
+            for i in range(lora_dim)
+        ]
+        
+        self.down_list = nn.ParameterList([
+            torch.empty(1, shape[1])
+            for i in range(lora_dim)
+        ])
+        self.down_list.requires_grad_(False)
+        self.down_update = [
+            torch.zeros_like(self.down_list[i])
+            for i in range(lora_dim)
+        ]
+        
+        self.index = 0
+        
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+        alpha = lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+
+        # Need more experiences on init method
+        
+        for v in self.down_list:
+            torch.nn.init.kaiming_uniform_(v, a=math.sqrt(5))
+        for v in self.up_list:
+            torch.nn.init.zeros_(v)
+        for i, v in enumerate(self.up_update):
+            v.copy_(self.up_list[i])
+        for i, v in enumerate(self.down_update):
+            v.copy_(self.down_list[i])
+
+        self.multiplier = multiplier
+        self.org_module = [org_module] # remove in applying
+        self.grad_ckpt = False
+        
+        self.apply_train(0)
+    
+    def state_dict(self, *args, destination=None, prefix='', keep_vars=False):
+        # TODO: Remove `args` and the parsing logic when BC allows.
+        if len(args) > 0:
+            if destination is None:
+                destination = args[0]
+            if len(args) > 1 and prefix == '':
+                prefix = args[1]
+            if len(args) > 2 and keep_vars is False:
+                keep_vars = args[2]
+            # DeprecationWarning is ignored by default
+
+        if destination is None:
+            destination = OrderedDict()
+            destination._metadata = OrderedDict()
+
+        local_metadata = dict(version=self._version)
+        if hasattr(destination, "_metadata"):
+            destination._metadata[prefix[:-1]] = local_metadata
+
+        destination[f'{prefix}alpha'] = self.alpha
+        destination[f'{prefix}lora_up.weight'] = nn.Parameter(
+            torch.concat(self.up_update, dim=1)
+        )
+        destination[f'{prefix}lora_down.weight'] = nn.Parameter(
+            torch.concat(self.down_update)
+        )
+        return destination
+
+    def apply_to(self):
+        self.org_module[0].forward = self.forward
+    
+    def apply_train(self, b:int):
+        self.up_list.requires_grad_(False)
+        self.down_list.requires_grad_(False)
+            
+        for i in range(self.index*self.block_size, (self.index+1)*self.block_size):
+            self.up_update[i].copy_(self.up_list[i])
+            self.down_update[i].copy_(self.down_list[i])
+        
+        for i in range(b*self.block_size, (b+1)*self.block_size):
+            self.up_list[i].copy_(self.up_update[i])
+            self.down_list[i].copy_(self.down_update[i])
+        
+        self.up_list.requires_grad_(True)
+        self.down_list.requires_grad_(True)
+        self.index = b
+
+    @torch.enable_grad()
+    def forward(self, x):
+        b = random.randint(0, self.block_count-1)
+        if self.up_update[b].device != self.up_list[b].device:
+            device = self.up_list[b].device
+            for i in range(self.lora_dim):
+                self.up_update[i] = self.up_update[i].to(device)
+                self.down_update[i] = self.down_update[i].to(device)
+        
+        if self.training:
+            self.apply_train(b)
+        down = torch.concat(
+            list(self.down_update[:b*self.block_size]) 
+            + list(self.down_list[b*self.block_size:(b+1)*self.block_size])
+        )
+        up = torch.concat(
+            list(self.up_update[:b*self.block_size]) 
+            + list(self.up_list[b*self.block_size:(b+1)*self.block_size]),
+            dim=1
+        )
+        
+        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
+        return self.op(
+            x, 
+            self.org_module[0].weight + (up@down).view(self.shape) * self.alpha/(b+1),
+            bias,
+            **self.extra_args
+        )

lycoris/functional/__init__.py

@@ -0,0 +1,13 @@
+from .general import (
+    rebuild_tucker,
+    factorization,
+    power2factorization,
+)
+from .locon import (
+    lora_diff_weight,
+    lora_bypass_forward_diff,
+)
+from .loha import (
+    loha_diff_weight,
+    loha_bypass_forward_diff,
+)

lycoris/functional/general.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+FUNC_LIST = [None, None, F.linear, F.conv1d, F.conv2d, F.conv3d]
+
+
+def rebuild_tucker(t, wa, wb):
+    rebuild2 = torch.einsum("i j ..., i p, j r -> p r ...", t, wa, wb)
+    return rebuild2
+
+
+def factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
+    """
+    return a tuple of two value of input dimension decomposed by the number closest to factor
+    second value is higher or equal than first value.
+
+    In LoRA with Kroneckor Product, first value is a value for weight scale.
+    second value is a value for weight.
+
+    Because of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.
+
+    examples)
+    factor
+        -1               2                4               8               16               ...
+    127 -> 1, 127   127 -> 1, 127    127 -> 1, 127   127 -> 1, 127   127 -> 1, 127
+    128 -> 8, 16    128 -> 2, 64     128 -> 4, 32    128 -> 8, 16    128 -> 8, 16
+    250 -> 10, 25   250 -> 2, 125    250 -> 2, 125   250 -> 5, 50    250 -> 10, 25
+    360 -> 8, 45    360 -> 2, 180    360 -> 4, 90    360 -> 8, 45    360 -> 12, 30
+    512 -> 16, 32   512 -> 2, 256    512 -> 4, 128   512 -> 8, 64    512 -> 16, 32
+    1024 -> 32, 32  1024 -> 2, 512   1024 -> 4, 256  1024 -> 8, 128  1024 -> 16, 64
+    """
+
+    if factor > 0 and (dimension % factor) == 0:
+        m = factor
+        n = dimension // factor
+        if m > n:
+            n, m = m, n
+        return m, n
+    if factor < 0:
+        factor = dimension
+    m, n = 1, dimension
+    length = m + n
+    while m < n:
+        new_m = m + 1
+        while dimension % new_m != 0:
+            new_m += 1
+        new_n = dimension // new_m
+        if new_m + new_n > length or new_m > factor:
+            break
+        else:
+            m, n = new_m, new_n
+    if m > n:
+        n, m = m, n
+    return m, n
+
+
+def power2factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
+    """
+    m = 2k
+    n = 2**p
+    m*n = dim
+    """
+    if factor == -1:
+        factor = dimension
+
+    # Find the first solution and check if it is even doable
+    m = n = 0
+    while m <= factor:
+        m += 2
+        while dimension % m != 0 and m < dimension:
+            m += 2
+        if m > factor:
+            break
+        if sum(int(i) for i in f"{dimension//m:b}") == 1:
+            n = dimension // m
+
+    return dimension // n, n