Upload 132 files

cc69848 verified over 1 year ago

10.7 kB

from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.parametrize as parametrize

from ..utils.quant import QuantLinears, log_bypass, log_suspect


class ModuleCustomSD(nn.Module):
    def __init__(self):
        super().__init__()
        self._register_load_state_dict_pre_hook(self.load_weight_prehook)
        self.register_load_state_dict_post_hook(self.load_weight_hook)

    def load_weight_prehook(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        pass

    def load_weight_hook(self, module, incompatible_keys):
        pass

    def custom_state_dict(self):
        return None

    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

        if (custom_sd := self.custom_state_dict()) is not None:
            for k, v in custom_sd.items():
                destination[f"{prefix}{k}"] = v
            return destination
        else:
            return super().state_dict(
                *args, destination=destination, prefix=prefix, keep_vars=keep_vars
            )


class LycorisBaseModule(ModuleCustomSD):
    name: str
    dtype_tensor: torch.Tensor
    support_module = {}
    weight_list = []
    weight_list_det = []

    def __init__(
        self,
        lora_name,
        org_module: nn.Module,
        multiplier=1.0,
        dropout=0.0,
        rank_dropout=0.0,
        module_dropout=0.0,
        rank_dropout_scale=False,
        bypass_mode=None,
        **kwargs,
    ):
        """if alpha == 0 or None, alpha is rank (no scaling)."""
        super().__init__()
        self.lora_name = lora_name
        self.not_supported = False

        self.module = type(org_module)
        if isinstance(org_module, nn.Linear):
            self.module_type = "linear"
            self.shape = (org_module.out_features, org_module.in_features)
            self.op = F.linear
            self.dim = org_module.out_features
            self.kw_dict = {}
        elif isinstance(org_module, nn.Conv1d):
            self.module_type = "conv1d"
            self.shape = (
                org_module.out_channels,
                org_module.in_channels,
                *org_module.kernel_size,
            )
            self.op = F.conv1d
            self.dim = org_module.out_channels
            self.kw_dict = {
                "stride": org_module.stride,
                "padding": org_module.padding,
                "dilation": org_module.dilation,
                "groups": org_module.groups,
            }
        elif isinstance(org_module, nn.Conv2d):
            self.module_type = "conv2d"
            self.shape = (
                org_module.out_channels,
                org_module.in_channels,
                *org_module.kernel_size,
            )
            self.op = F.conv2d
            self.dim = org_module.out_channels
            self.kw_dict = {
                "stride": org_module.stride,
                "padding": org_module.padding,
                "dilation": org_module.dilation,
                "groups": org_module.groups,
            }
        elif isinstance(org_module, nn.Conv3d):
            self.module_type = "conv3d"
            self.shape = (
                org_module.out_channels,
                org_module.in_channels,
                *org_module.kernel_size,
            )
            self.op = F.conv3d
            self.dim = org_module.out_channels
            self.kw_dict = {
                "stride": org_module.stride,
                "padding": org_module.padding,
                "dilation": org_module.dilation,
                "groups": org_module.groups,
            }
        elif isinstance(org_module, nn.LayerNorm):
            self.module_type = "layernorm"
            self.shape = tuple(org_module.normalized_shape)
            self.op = F.layer_norm
            self.dim = org_module.normalized_shape[0]
            self.kw_dict = {
                "normalized_shape": org_module.normalized_shape,
                "eps": org_module.eps,
            }
        elif isinstance(org_module, nn.GroupNorm):
            self.module_type = "groupnorm"
            self.shape = (org_module.num_channels,)
            self.op = F.group_norm
            self.group_num = org_module.num_groups
            self.dim = org_module.num_channels
            self.kw_dict = {"num_groups": org_module.num_groups, "eps": org_module.eps}
        else:
            self.not_supported = True
            self.module_type = "unknown"

        self.register_buffer("dtype_tensor", torch.tensor(0.0), persistent=False)

        self.is_quant = False
        if isinstance(org_module, QuantLinears):
            if not bypass_mode:
                log_bypass()
            self.is_quant = True
            bypass_mode = True
        if (
            isinstance(org_module, nn.Linear)
            and org_module.__class__.__name__ != "Linear"
        ):
            if bypass_mode is None:
                log_suspect()
                bypass_mode = True
            if bypass_mode == True:
                self.is_quant = True
        self.bypass_mode = bypass_mode
        self.dropout = dropout
        self.rank_dropout = rank_dropout
        self.rank_dropout_scale = rank_dropout_scale
        self.module_dropout = module_dropout

        ## Dropout things
        # Since LoKr/LoHa/OFT/BOFT are hard to follow the rank_dropout definition from kohya
        # We redefine the dropout procedure here.
        # g(x) = WX + drop(Brank_drop(AX)) for LoCon(lora), bypass
        # g(x) = WX + drop(ΔWX) for any algo except LoCon(lora), bypass
        # g(x) = (W + Brank_drop(A))X for LoCon(lora), rebuid
        # g(x) = (W + rank_drop(ΔW))X for any algo except LoCon(lora), rebuild
        self.drop = nn.Identity() if dropout == 0 else nn.Dropout(dropout)
        self.rank_drop = (
            nn.Identity() if rank_dropout == 0 else nn.Dropout(rank_dropout)
        )

        self.multiplier = multiplier
        self.org_forward = org_module.forward
        self.org_module = [org_module]

    @classmethod
    def parametrize(cls, org_module, attr, *args, **kwargs):
        from .full import FullModule

        if cls is FullModule:
            raise RuntimeError("FullModule cannot be used for parametrize.")
        target_param = getattr(org_module, attr)
        kwargs["bypass_mode"] = False
        if target_param.dim() == 2:
            proxy_module = nn.Linear(
                target_param.shape[0], target_param.shape[1], bias=False
            )
            proxy_module.weight = target_param
        elif target_param.dim() > 2:
            module_type = [
                None,
                None,
                None,
                nn.Conv1d,
                nn.Conv2d,
                nn.Conv3d,
                None,
                None,
            ][target_param.dim()]
            proxy_module = module_type(
                target_param.shape[0],
                target_param.shape[1],
                *target_param.shape[2:],
                bias=False,
            )
            proxy_module.weight = target_param
        module_obj = cls("", proxy_module, *args, **kwargs)
        module_obj.forward = module_obj.parametrize_forward
        module_obj.to(target_param)
        parametrize.register_parametrization(org_module, attr, module_obj)
        return module_obj

    @classmethod
    def algo_check(cls, state_dict, lora_name):
        return any(f"{lora_name}.{k}" in state_dict for k in cls.weight_list_det)

    @classmethod
    def extract_state_dict(cls, state_dict, lora_name):
        return [state_dict.get(f"{lora_name}.{k}", None) for k in cls.weight_list]

    @classmethod
    def make_module_from_state_dict(cls, lora_name, orig_module, *weights):
        raise NotImplementedError

    @property
    def dtype(self):
        return self.dtype_tensor.dtype

    @property
    def device(self):
        return self.dtype_tensor.device

    @property
    def org_weight(self):
        return self.org_module[0].weight

    @org_weight.setter
    def org_weight(self, value):
        self.org_module[0].weight.data.copy_(value)

    def apply_to(self, **kwargs):
        if self.not_supported:
            return
        self.org_forward = self.org_module[0].forward
        self.org_module[0].forward = self.forward

    def restore(self):
        if self.not_supported:
            return
        self.org_module[0].forward = self.org_forward

    def merge_to(self, multiplier=1.0):
        if self.not_supported:
            return
        self_device = next(self.parameters()).device
        self_dtype = next(self.parameters()).dtype
        self.to(self.org_weight)
        weight, bias = self.get_merged_weight(
            multiplier, self.org_weight.shape, self.org_weight.device
        )
        self.org_weight = weight.to(self.org_weight)
        if bias is not None:
            bias = bias.to(self.org_weight)
            if self.org_module[0].bias is not None:
                self.org_module[0].bias.data.copy_(bias)
            else:
                self.org_module[0].bias = nn.Parameter(bias)
        self.to(self_device, self_dtype)

    def get_diff_weight(self, multiplier=1.0, shape=None, device=None):
        raise NotImplementedError

    def get_merged_weight(self, multiplier=1.0, shape=None, device=None):
        raise NotImplementedError

    @torch.no_grad()
    def apply_max_norm(self, max_norm, device=None):
        return None, None

    def bypass_forward_diff(self, x, scale=1):
        raise NotImplementedError

    def bypass_forward(self, x, scale=1):
        raise NotImplementedError

    def parametrize_forward(self, x: torch.Tensor, *args, **kwargs):
        return self.get_merged_weight(
            multiplier=self.multiplier, shape=x.shape, device=x.device
        )[0].to(x.dtype)

    def forward(self, *args, **kwargs):
        raise NotImplementedError

	from collections import OrderedDict

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import torch.nn.utils.parametrize as parametrize

	from ..utils.quant import QuantLinears, log_bypass, log_suspect


	class ModuleCustomSD(nn.Module):
	def __init__(self):
	super().__init__()
	self._register_load_state_dict_pre_hook(self.load_weight_prehook)
	self.register_load_state_dict_post_hook(self.load_weight_hook)

	def load_weight_prehook(
	self,
	state_dict,
	prefix,
	local_metadata,
	strict,
	missing_keys,
	unexpected_keys,
	error_msgs,
	):
	pass

	def load_weight_hook(self, module, incompatible_keys):
	pass

	def custom_state_dict(self):
	return None

	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

	if (custom_sd := self.custom_state_dict()) is not None:
	for k, v in custom_sd.items():
	destination[f"{prefix}{k}"] = v
	return destination
	else:
	return super().state_dict(
	*args, destination=destination, prefix=prefix, keep_vars=keep_vars
	)


	class LycorisBaseModule(ModuleCustomSD):
	name: str
	dtype_tensor: torch.Tensor
	support_module = {}
	weight_list = []
	weight_list_det = []

	def __init__(
	self,
	lora_name,
	org_module: nn.Module,
	multiplier=1.0,
	dropout=0.0,
	rank_dropout=0.0,
	module_dropout=0.0,
	rank_dropout_scale=False,
	bypass_mode=None,
	**kwargs,
	):
	"""if alpha == 0 or None, alpha is rank (no scaling)."""
	super().__init__()
	self.lora_name = lora_name
	self.not_supported = False

	self.module = type(org_module)
	if isinstance(org_module, nn.Linear):
	self.module_type = "linear"
	self.shape = (org_module.out_features, org_module.in_features)
	self.op = F.linear
	self.dim = org_module.out_features
	self.kw_dict = {}
	elif isinstance(org_module, nn.Conv1d):
	self.module_type = "conv1d"
	self.shape = (
	org_module.out_channels,
	org_module.in_channels,
	*org_module.kernel_size,
	)
	self.op = F.conv1d
	self.dim = org_module.out_channels
	self.kw_dict = {
	"stride": org_module.stride,
	"padding": org_module.padding,
	"dilation": org_module.dilation,
	"groups": org_module.groups,
	}
	elif isinstance(org_module, nn.Conv2d):
	self.module_type = "conv2d"
	self.shape = (
	org_module.out_channels,
	org_module.in_channels,
	*org_module.kernel_size,
	)
	self.op = F.conv2d
	self.dim = org_module.out_channels
	self.kw_dict = {
	"stride": org_module.stride,
	"padding": org_module.padding,
	"dilation": org_module.dilation,
	"groups": org_module.groups,
	}
	elif isinstance(org_module, nn.Conv3d):
	self.module_type = "conv3d"
	self.shape = (
	org_module.out_channels,
	org_module.in_channels,
	*org_module.kernel_size,
	)
	self.op = F.conv3d
	self.dim = org_module.out_channels
	self.kw_dict = {
	"stride": org_module.stride,
	"padding": org_module.padding,
	"dilation": org_module.dilation,
	"groups": org_module.groups,
	}
	elif isinstance(org_module, nn.LayerNorm):
	self.module_type = "layernorm"
	self.shape = tuple(org_module.normalized_shape)
	self.op = F.layer_norm
	self.dim = org_module.normalized_shape[0]
	self.kw_dict = {
	"normalized_shape": org_module.normalized_shape,
	"eps": org_module.eps,
	}
	elif isinstance(org_module, nn.GroupNorm):
	self.module_type = "groupnorm"
	self.shape = (org_module.num_channels,)
	self.op = F.group_norm
	self.group_num = org_module.num_groups
	self.dim = org_module.num_channels
	self.kw_dict = {"num_groups": org_module.num_groups, "eps": org_module.eps}
	else:
	self.not_supported = True
	self.module_type = "unknown"

	self.register_buffer("dtype_tensor", torch.tensor(0.0), persistent=False)

	self.is_quant = False
	if isinstance(org_module, QuantLinears):
	if not bypass_mode:
	log_bypass()
	self.is_quant = True
	bypass_mode = True
	if (
	isinstance(org_module, nn.Linear)
	and org_module.__class__.__name__ != "Linear"
	):
	if bypass_mode is None:
	log_suspect()
	bypass_mode = True
	if bypass_mode == True:
	self.is_quant = True
	self.bypass_mode = bypass_mode
	self.dropout = dropout
	self.rank_dropout = rank_dropout
	self.rank_dropout_scale = rank_dropout_scale
	self.module_dropout = module_dropout

	## Dropout things
	# Since LoKr/LoHa/OFT/BOFT are hard to follow the rank_dropout definition from kohya
	# We redefine the dropout procedure here.
	# g(x) = WX + drop(Brank_drop(AX)) for LoCon(lora), bypass
	# g(x) = WX + drop(ΔWX) for any algo except LoCon(lora), bypass
	# g(x) = (W + Brank_drop(A))X for LoCon(lora), rebuid
	# g(x) = (W + rank_drop(ΔW))X for any algo except LoCon(lora), rebuild
	self.drop = nn.Identity() if dropout == 0 else nn.Dropout(dropout)
	self.rank_drop = (
	nn.Identity() if rank_dropout == 0 else nn.Dropout(rank_dropout)
	)

	self.multiplier = multiplier
	self.org_forward = org_module.forward
	self.org_module = [org_module]

	@classmethod
	def parametrize(cls, org_module, attr, args, *kwargs):
	from .full import FullModule

	if cls is FullModule:
	raise RuntimeError("FullModule cannot be used for parametrize.")
	target_param = getattr(org_module, attr)
	kwargs["bypass_mode"] = False
	if target_param.dim() == 2:
	proxy_module = nn.Linear(
	target_param.shape[0], target_param.shape[1], bias=False
	)
	proxy_module.weight = target_param
	elif target_param.dim() > 2:
	module_type = [
	None,
	None,
	None,
	nn.Conv1d,
	nn.Conv2d,
	nn.Conv3d,
	None,
	None,
	][target_param.dim()]
	proxy_module = module_type(
	target_param.shape[0],
	target_param.shape[1],
	*target_param.shape[2:],
	bias=False,
	)
	proxy_module.weight = target_param
	module_obj = cls("", proxy_module, args, *kwargs)
	module_obj.forward = module_obj.parametrize_forward
	module_obj.to(target_param)
	parametrize.register_parametrization(org_module, attr, module_obj)
	return module_obj

	@classmethod
	def algo_check(cls, state_dict, lora_name):
	return any(f"{lora_name}.{k}" in state_dict for k in cls.weight_list_det)

	@classmethod
	def extract_state_dict(cls, state_dict, lora_name):
	return [state_dict.get(f"{lora_name}.{k}", None) for k in cls.weight_list]

	@classmethod
	def make_module_from_state_dict(cls, lora_name, orig_module, *weights):
	raise NotImplementedError

	@property
	def dtype(self):
	return self.dtype_tensor.dtype

	@property
	def device(self):
	return self.dtype_tensor.device

	@property
	def org_weight(self):
	return self.org_module[0].weight

	@org_weight.setter
	def org_weight(self, value):
	self.org_module[0].weight.data.copy_(value)

	def apply_to(self, **kwargs):
	if self.not_supported:
	return
	self.org_forward = self.org_module[0].forward
	self.org_module[0].forward = self.forward

	def restore(self):
	if self.not_supported:
	return
	self.org_module[0].forward = self.org_forward

	def merge_to(self, multiplier=1.0):
	if self.not_supported:
	return
	self_device = next(self.parameters()).device
	self_dtype = next(self.parameters()).dtype
	self.to(self.org_weight)
	weight, bias = self.get_merged_weight(
	multiplier, self.org_weight.shape, self.org_weight.device
	)
	self.org_weight = weight.to(self.org_weight)
	if bias is not None:
	bias = bias.to(self.org_weight)
	if self.org_module[0].bias is not None:
	self.org_module[0].bias.data.copy_(bias)
	else:
	self.org_module[0].bias = nn.Parameter(bias)
	self.to(self_device, self_dtype)

	def get_diff_weight(self, multiplier=1.0, shape=None, device=None):
	raise NotImplementedError

	def get_merged_weight(self, multiplier=1.0, shape=None, device=None):
	raise NotImplementedError

	@torch.no_grad()
	def apply_max_norm(self, max_norm, device=None):
	return None, None

	def bypass_forward_diff(self, x, scale=1):
	raise NotImplementedError

	def bypass_forward(self, x, scale=1):
	raise NotImplementedError

	def parametrize_forward(self, x: torch.Tensor, args, *kwargs):
	return self.get_merged_weight(
	multiplier=self.multiplier, shape=x.shape, device=x.device
	)[0].to(x.dtype)

	def forward(self, args, *kwargs):
	raise NotImplementedError