a1111-microsoftexcel-locon/locon_compvis.py

'''
Hijack version of kohya-ss/additional_networks/scripts/lora_compvis.py
'''
# LoRA network module
# reference:
# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py

import copy
import math
import re
from typing import NamedTuple
import torch
from locon import LoConModule


class LoRAInfo(NamedTuple):
    lora_name: str
    module_name: str
    module: torch.nn.Module
    multiplier: float
    dim: int
    alpha: float


def create_network_and_apply_compvis(du_state_dict, multiplier_tenc, multiplier_unet, text_encoder, unet, **kwargs):
    # get device and dtype from unet
    for module in unet.modules():
        if module.__class__.__name__ == "Linear":
            param: torch.nn.Parameter = module.weight
            # device = param.device
            dtype = param.dtype
            break

    # get dims (rank) and alpha from state dict
    # currently it is assumed all LoRA have same alpha. alpha may be different in future.
    network_alpha = None
    conv_alpha = None
    network_dim = None
    conv_dim = None
    for key, value in du_state_dict.items():
        if network_alpha is None and 'alpha' in key:
            network_alpha = value
        if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
            network_dim = value.size()[0]
        if network_alpha is not None and network_dim is not None:
            break
    if network_alpha is None:
        network_alpha = network_dim

    print(f"dimension: {network_dim},\n"
          f"alpha: {network_alpha},\n"
          f"multiplier_unet: {multiplier_unet},\n"
          f"multiplier_tenc: {multiplier_tenc}"
          )
    if network_dim is None:
        print(f"The selected model is not LoRA or not trained by `sd-scripts`?")
        network_dim = 4
        network_alpha = 1

    # create, apply and load weights
    network = LoConNetworkCompvis(
        text_encoder, unet, du_state_dict,
        multiplier_tenc = multiplier_tenc,
        multiplier_unet = multiplier_unet,
    )
    state_dict = network.apply_lora_modules(du_state_dict)              # some weights are applied to text encoder
    network.to(dtype)                                              # with this, if error comes from next line, the model will be used
    info = network.load_state_dict(state_dict, strict=False)

    # remove redundant warnings
    if len(info.missing_keys) > 4:
        missing_keys = []
        alpha_count = 0
        for key in info.missing_keys:
            if 'alpha' not in key:
                missing_keys.append(key)
            else:
                if alpha_count == 0:
                    missing_keys.append(key)
                alpha_count += 1
        if alpha_count > 1:
            missing_keys.append(
                    f"... and {alpha_count-1} alphas. The model doesn't have alpha, use dim (rannk) as alpha. You can ignore this message.")

        info = torch.nn.modules.module._IncompatibleKeys(missing_keys, info.unexpected_keys)

    return network, info


class LoConNetworkCompvis(torch.nn.Module):
    # UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
    # TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
    LOCON_TARGET = ["ResBlock", "Downsample", "Upsample"]
    UNET_TARGET_REPLACE_MODULE = ["SpatialTransformer"] + LOCON_TARGET  # , "Attention"]
    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["ResidualAttentionBlock", "CLIPAttention", "CLIPMLP"]

    LORA_PREFIX_UNET = 'lora_unet'
    LORA_PREFIX_TEXT_ENCODER = 'lora_te'

    @classmethod
    def convert_diffusers_name_to_compvis(cls, v2, du_name):
        """
        convert diffusers's LoRA name to CompVis
        """
        cv_name = None
        if "lora_unet_" in du_name:
            m = re.search(r"_down_blocks_(\d+)_attentions_(\d+)_(.+)", du_name)
            if m:
                du_block_index = int(m.group(1))
                du_attn_index = int(m.group(2))
                du_suffix = m.group(3)

                cv_index = 1 + du_block_index * 3 + du_attn_index      # 1,2, 4,5, 7,8
                cv_name = f"lora_unet_input_blocks_{cv_index}_1_{du_suffix}"
                return cv_name
            
            m = re.search(r"_mid_block_attentions_(\d+)_(.+)", du_name)
            if m:
                du_suffix = m.group(2)
                cv_name = f"lora_unet_middle_block_1_{du_suffix}"
                return cv_name
            
            m = re.search(r"_up_blocks_(\d+)_attentions_(\d+)_(.+)", du_name)
            if m:
                du_block_index = int(m.group(1))
                du_attn_index = int(m.group(2))
                du_suffix = m.group(3)

                cv_index = du_block_index * 3 + du_attn_index      # 3,4,5, 6,7,8, 9,10,11
                cv_name = f"lora_unet_output_blocks_{cv_index}_1_{du_suffix}"
                return cv_name
            
            m = re.search(r"_down_blocks_(\d+)_resnets_(\d+)_(.+)", du_name)
            if m:
                du_block_index = int(m.group(1))
                du_res_index = int(m.group(2))
                du_suffix = m.group(3)
                cv_suffix = {
                    'conv1': 'in_layers_2',
                    'conv2': 'out_layers_3',
                    'time_emb_proj': 'emb_layers_1',
                    'conv_shortcut': 'skip_connection'
                }[du_suffix]

                cv_index = 1 + du_block_index * 3 + du_res_index      # 1,2, 4,5, 7,8
                cv_name = f"lora_unet_input_blocks_{cv_index}_0_{cv_suffix}"
                return cv_name
            
            m = re.search(r"_down_blocks_(\d+)_downsamplers_0_conv", du_name)
            if m:
                block_index = int(m.group(1))
                cv_index = 3 + block_index * 3
                cv_name = f"lora_unet_input_blocks_{cv_index}_0_op"
                return cv_name
            
            m = re.search(r"_mid_block_resnets_(\d+)_(.+)", du_name)
            if m:
                index = int(m.group(1))
                du_suffix = m.group(2)
                cv_suffix = {
                    'conv1': 'in_layers_2',
                    'conv2': 'out_layers_3',
                    'time_emb_proj': 'emb_layers_1',
                    'conv_shortcut': 'skip_connection'
                }[du_suffix]
                cv_name = f"lora_unet_middle_block_{index*2}_{cv_suffix}"
                return cv_name
            
            m = re.search(r"_up_blocks_(\d+)_resnets_(\d+)_(.+)", du_name)
            if m:
                du_block_index = int(m.group(1))
                du_res_index = int(m.group(2))
                du_suffix = m.group(3)
                cv_suffix = {
                    'conv1': 'in_layers_2',
                    'conv2': 'out_layers_3',
                    'time_emb_proj': 'emb_layers_1',
                    'conv_shortcut': 'skip_connection'
                }[du_suffix]

                cv_index = du_block_index * 3 + du_res_index      # 1,2, 4,5, 7,8
                cv_name = f"lora_unet_output_blocks_{cv_index}_0_{cv_suffix}"
                return cv_name
            
            m = re.search(r"_up_blocks_(\d+)_upsamplers_0_conv", du_name)
            if m:
                block_index = int(m.group(1))
                cv_index = block_index * 3 + 2
                cv_name = f"lora_unet_output_blocks_{cv_index}_{bool(block_index)+1}_conv"
                return cv_name
            
        elif "lora_te_" in du_name:
            m = re.search(r"_model_encoder_layers_(\d+)_(.+)", du_name)
            if m:
                du_block_index = int(m.group(1))
                du_suffix = m.group(2)

                cv_index = du_block_index
                if v2:
                    if 'mlp_fc1' in du_suffix:
                        cv_name = f"lora_te_wrapped_model_transformer_resblocks_{cv_index}_{du_suffix.replace('mlp_fc1', 'mlp_c_fc')}"
                    elif 'mlp_fc2' in du_suffix:
                        cv_name = f"lora_te_wrapped_model_transformer_resblocks_{cv_index}_{du_suffix.replace('mlp_fc2', 'mlp_c_proj')}"
                    elif 'self_attn':
                        # handled later
                        cv_name = f"lora_te_wrapped_model_transformer_resblocks_{cv_index}_{du_suffix.replace('self_attn', 'attn')}"
                else:
                    cv_name = f"lora_te_wrapped_transformer_text_model_encoder_layers_{cv_index}_{du_suffix}"

        assert cv_name is not None, f"conversion failed: {du_name}. the model may not be trained by `sd-scripts`."
        return cv_name

    @classmethod
    def convert_state_dict_name_to_compvis(cls, v2, state_dict):
        """
        convert keys in state dict to load it by load_state_dict
        """
        new_sd = {}
        for key, value in state_dict.items():
            tokens = key.split('.')
            compvis_name = LoConNetworkCompvis.convert_diffusers_name_to_compvis(v2, tokens[0])
            new_key = compvis_name + '.' + '.'.join(tokens[1:])
            new_sd[new_key] = value

        return new_sd

    def __init__(self, text_encoder, unet, du_state_dict, multiplier_tenc=1.0, multiplier_unet=1.0) -> None:
        super().__init__()
        self.multiplier_unet = multiplier_unet
        self.multiplier_tenc = multiplier_tenc

        # create module instances
        for name, module in text_encoder.named_modules():
            for child_name, child_module in module.named_modules():
                if child_module.__class__.__name__ == 'MultiheadAttention':
                    self.v2 = True
                    break
            else:
                continue
            break
        else:
            self.v2 = False
        comp_state_dict = {}

        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules, multiplier):
            nonlocal comp_state_dict
            loras = []
            replaced_modules = []
            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():
                        layer = child_module.__class__.__name__
                        lora_name = prefix + '.' + name + '.' + child_name
                        lora_name = lora_name.replace('.', '_')
                        if layer == "Linear" or layer == "Conv2d":
                            if '_resblocks_23_' in lora_name:  # ignore last block in StabilityAi Text Encoder
                                break
                            if f'{lora_name}.lora_down.weight' not in comp_state_dict:
                                if module.__class__.__name__ in LoConNetworkCompvis.LOCON_TARGET:
                                    continue
                                else:
                                    print(f'Cannot find: "{lora_name}", skipped')
                                    continue
                            rank = comp_state_dict[f'{lora_name}.lora_down.weight'].shape[0]
                            alpha = comp_state_dict.get(f'{lora_name}.alpha', torch.tensor(rank)).item()
                            lora = LoConModule(lora_name, child_module, multiplier, rank, alpha)
                            loras.append(lora)

                            replaced_modules.append(child_module)
                        elif child_module.__class__.__name__ == "MultiheadAttention":
                            # make four modules: not replacing forward method but merge weights
                            self.v2 = True
                            for suffix in ['q', 'k', 'v', 'out']:
                                module_name = prefix + '.' + name + '.' + child_name          # ~.attn
                                module_name = module_name.replace('.', '_')
                                if '_resblocks_23_' in module_name:                           # ignore last block in StabilityAi Text Encoder
                                    break
                                lora_name = module_name + '_' + suffix
                                lora_info = LoRAInfo(lora_name, module_name, child_module, multiplier, 0, 0)
                                loras.append(lora_info)

                                replaced_modules.append(child_module)
            return loras, replaced_modules
        
        for k,v in LoConNetworkCompvis.convert_state_dict_name_to_compvis(self.v2, du_state_dict).items():
            comp_state_dict[k] = v

        self.text_encoder_loras, te_rep_modules = create_modules(
            LoConNetworkCompvis.LORA_PREFIX_TEXT_ENCODER,
            text_encoder, 
            LoConNetworkCompvis.TEXT_ENCODER_TARGET_REPLACE_MODULE, 
            self.multiplier_tenc
        )
        print(f"create LoCon for Text Encoder: {len(self.text_encoder_loras)} modules.")
        
        self.unet_loras, unet_rep_modules = create_modules(
            LoConNetworkCompvis.LORA_PREFIX_UNET, 
            unet, 
            LoConNetworkCompvis.UNET_TARGET_REPLACE_MODULE, 
            self.multiplier_unet
        )
        print(f"create LoCon for U-Net: {len(self.unet_loras)} modules.")

        # make backup of original forward/weights, if multiple modules are applied, do in 1st module only
        backed_up = False                     # messaging purpose only
        for rep_module in te_rep_modules + unet_rep_modules:
            if rep_module.__class__.__name__ == "MultiheadAttention":      # multiple MHA modules are in list, prevent to backed up forward
                if not hasattr(rep_module, "_lora_org_weights"):
                    # avoid updating of original weights. state_dict is reference to original weights
                    rep_module._lora_org_weights = copy.deepcopy(rep_module.state_dict())
                    backed_up = True
            elif not hasattr(rep_module, "_lora_org_forward"):
                rep_module._lora_org_forward = rep_module.forward
                backed_up = True
        if backed_up:
            print("original forward/weights is backed up.")

        # 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 restore(self, text_encoder, unet):
        # restore forward/weights from property for all modules
        restored = False                        # messaging purpose only
        modules = []
        modules.extend(text_encoder.modules())
        modules.extend(unet.modules())
        for module in modules:
            if hasattr(module, "_lora_org_forward"):
                module.forward = module._lora_org_forward
                del module._lora_org_forward
                restored = True
            if hasattr(module, "_lora_org_weights"):        # module doesn't have forward and weights at same time currently, but supports it for future changing
                module.load_state_dict(module._lora_org_weights)
                del module._lora_org_weights
                restored = True

        if restored:
            print("original forward/weights is restored.")

    def apply_lora_modules(self, du_state_dict):
        # conversion 1st step: convert names in state_dict
        state_dict = LoConNetworkCompvis.convert_state_dict_name_to_compvis(self.v2, du_state_dict)

        # check state_dict has text_encoder or unet
        weights_has_text_encoder = weights_has_unet = False
        for key in state_dict.keys():
            if key.startswith(LoConNetworkCompvis.LORA_PREFIX_TEXT_ENCODER):
                weights_has_text_encoder = True
            elif key.startswith(LoConNetworkCompvis.LORA_PREFIX_UNET):
                weights_has_unet = True
            if weights_has_text_encoder and weights_has_unet:
                break

        apply_text_encoder = weights_has_text_encoder
        apply_unet = weights_has_unet

        if apply_text_encoder:
            print("enable LoCon for text encoder")
        else:
            self.text_encoder_loras = []

        if apply_unet:
            print("enable LoCon for U-Net")
        else:
            self.unet_loras = []

        # add modules to network: this makes state_dict can be got from LoRANetwork
        mha_loras = {}
        for lora in self.text_encoder_loras + self.unet_loras:
            if type(lora) == LoConModule:
                lora.apply_to()                           # ensure remove reference to original Linear: reference makes key of state_dict
                self.add_module(lora.lora_name, lora)
            else:
                # SD2.x MultiheadAttention merge weights to MHA weights
                lora_info: LoRAInfo = lora
                if lora_info.module_name not in mha_loras:
                    mha_loras[lora_info.module_name] = {}

                lora_dic = mha_loras[lora_info.module_name]
                lora_dic[lora_info.lora_name] = lora_info
                if len(lora_dic) == 4:
                    # calculate and apply
                    w_q_dw = state_dict.get(lora_info.module_name + '_q_proj.lora_down.weight')
                    if w_q_dw is not None:                       # corresponding LoRa module exists
                        w_q_up = state_dict[lora_info.module_name + '_q_proj.lora_up.weight']
                        w_q_ap = state_dict.get(lora_info.module_name + '_q_proj.alpha', None)
                        w_k_dw = state_dict[lora_info.module_name + '_k_proj.lora_down.weight']
                        w_k_up = state_dict[lora_info.module_name + '_k_proj.lora_up.weight']
                        w_k_ap = state_dict.get(lora_info.module_name + '_k_proj.alpha', None)
                        w_v_dw = state_dict[lora_info.module_name + '_v_proj.lora_down.weight']
                        w_v_up = state_dict[lora_info.module_name + '_v_proj.lora_up.weight']
                        w_v_ap = state_dict.get(lora_info.module_name + '_v_proj.alpha', None)
                        w_out_dw = state_dict[lora_info.module_name + '_out_proj.lora_down.weight']
                        w_out_up = state_dict[lora_info.module_name + '_out_proj.lora_up.weight']
                        w_out_ap = state_dict.get(lora_info.module_name + '_out_proj.alpha', None)

                        sd = lora_info.module.state_dict()
                        qkv_weight = sd['in_proj_weight']
                        out_weight = sd['out_proj.weight']
                        dev = qkv_weight.device

                        def merge_weights(weight, up_weight, down_weight, alpha=None):
                            # calculate in float
                            if alpha is None:
                                alpha = down_weight.shape[0]
                            alpha = float(alpha)
                            scale = alpha / down_weight.shape[0]
                            dtype = weight.dtype
                            weight = weight.float() + lora_info.multiplier * (up_weight.to(dev, dtype=torch.float) @ down_weight.to(dev, dtype=torch.float)) * scale
                            weight = weight.to(dtype)
                            return weight

                        q_weight, k_weight, v_weight = torch.chunk(qkv_weight, 3)
                        if q_weight.size()[1] == w_q_up.size()[0]:
                            q_weight = merge_weights(q_weight, w_q_up, w_q_dw, w_q_ap)
                            k_weight = merge_weights(k_weight, w_k_up, w_k_dw, w_k_ap)
                            v_weight = merge_weights(v_weight, w_v_up, w_v_dw, w_v_ap)
                            qkv_weight = torch.cat([q_weight, k_weight, v_weight])

                            out_weight = merge_weights(out_weight, w_out_up, w_out_dw, w_out_ap)

                            sd['in_proj_weight'] = qkv_weight.to(dev)
                            sd['out_proj.weight'] = out_weight.to(dev)

                            lora_info.module.load_state_dict(sd)
                        else:
                            # different dim, version mismatch
                            print(f"shape of weight is different: {lora_info.module_name}. SD version may be different")

                        for t in ["q", "k", "v", "out"]:
                            del state_dict[f"{lora_info.module_name}_{t}_proj.lora_down.weight"]
                            del state_dict[f"{lora_info.module_name}_{t}_proj.lora_up.weight"]
                            alpha_key = f"{lora_info.module_name}_{t}_proj.alpha"
                            if alpha_key in state_dict:
                                del state_dict[alpha_key]
                    else:
                        # corresponding weight not exists: version mismatch
                        pass

        # conversion 2nd step: convert weight's shape (and handle wrapped)
        state_dict = self.convert_state_dict_shape_to_compvis(state_dict)

        return state_dict

    def convert_state_dict_shape_to_compvis(self, state_dict):
        # shape conversion
        current_sd = self.state_dict()        # to get target shape
        wrapped = False
        count = 0
        for key in list(state_dict.keys()):
            if key not in current_sd:
                continue                          # might be error or another version
            if "wrapped" in key:
                wrapped = True

            value: torch.Tensor = state_dict[key]
            if value.size() != current_sd[key].size():
                # print(key, value.size(), current_sd[key].size())
                # print(f"convert weights shape: {key}, from: {value.size()}, {len(value.size())}")
                count += 1
                if '.alpha' in key:
                    assert value.size().numel() == 1
                    value = torch.tensor(value.item())
                elif len(value.size()) == 4:
                    value = value.squeeze(3).squeeze(2)
                else:
                    value = value.unsqueeze(2).unsqueeze(3)
                state_dict[key] = value
            if tuple(value.size()) != tuple(current_sd[key].size()):
                print(
                        f"weight's shape is different: {key} expected {current_sd[key].size()} found {value.size()}. SD version may be different")
                del state_dict[key]
        print(f"shapes for {count} weights are converted.")

        # convert wrapped
        if not wrapped:
            print("remove 'wrapped' from keys")
            for key in list(state_dict.keys()):
                if "_wrapped_" in key:
                    new_key = key.replace("_wrapped_", "_")
                    state_dict[new_key] = state_dict[key]
                    del state_dict[key]

        return state_dict