[WIP] Add nonlinearity between Down and Up weights, expose alpha hyperparameter

lora_diffusion/cli_lora_add.py

@@ -75,13 +75,13 @@ def add(
  path_1,
  ).to("cpu")
 
- weight_apply_lora(loaded_pipeline.unet, torch.load(path_2), alpha=alpha)
+ weight_apply_lora(loaded_pipeline.unet, torch.load(path_2), scale=alpha)
  if with_text_lora:
 
  weight_apply_lora(
  loaded_pipeline.text_encoder,
  torch.load(_text_lora_path(path_2)),
- alpha=alpha,
+ scale=alpha,
  target_replace_module=["CLIPAttention"],
  )
 
@@ -93,12 +93,12 @@ def add(
  path_1,
  ).to("cpu")
 
- weight_apply_lora(loaded_pipeline.unet, torch.load(path_2), alpha=alpha)
+ weight_apply_lora(loaded_pipeline.unet, torch.load(path_2), scale=alpha)
  if with_text_lora:
  weight_apply_lora(
  loaded_pipeline.text_encoder,
  torch.load(_text_lora_path(path_2)),
- alpha=alpha,
+ scale=alpha,
  target_replace_module=["CLIPAttention"],
  )
 

lora_diffusion/lora.py

@@ -10,24 +10,39 @@
 
 
 class LoraInjectedLinear(nn.Module):
- def __init__(self, in_features, out_features, bias=False, r=4):
+ def __init__(self, in_features, out_features, bias=False, r=4, scale=1.0, init=None, nonlin: nn.Module = None):
  super().__init__()
 
  if r > min(in_features, out_features):
  raise ValueError(
  f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
  )
-
+
+ if scale <= 0:
+ raise ValueError(
+ f"LoRA scale {scale} must be greater than 0"
+ )
+
+ self.r = r
+ self.scale = scale
  self.linear = nn.Linear(in_features, out_features, bias)
  self.lora_down = nn.Linear(in_features, r, bias=False)
+ self.nonlin = nonlin if nonlin else None
  self.lora_up = nn.Linear(r, out_features, bias=False)
- self.scale = 1.0
 
- nn.init.normal_(self.lora_down.weight, std=1 / r)
+ if init=="kaiming":
+ pass
+ # Kaiming with a=math.sqrt(5) is default for nn.Linear
+ else:
+ nn.init.normal_(self.lora_down.weight, std=1 / r)
+
  nn.init.zeros_(self.lora_up.weight)
 
  def forward(self, input):
- return self.linear(input) + self.lora_up(self.lora_down(input)) * self.scale
+ if self.nonlin:
+ return self.linear(input) + self.lora_up(self.nonlin(self.lora_down(input))) * self.scale
+ else:
+ return self.linear(input) + self.lora_up(self.lora_down(input)) * self.scale
 
 
 UNET_DEFAULT_TARGET_REPLACE = {"CrossAttention", "Attention", "GEGLU"}
@@ -116,6 +131,9 @@ def inject_trainable_lora(
  model: nn.Module,
  target_replace_module: Set[str] = DEFAULT_TARGET_REPLACE,
  r: int = 4,
+ scale: float = 1.0,
+ init=None,
+ nonlin=None,
  loras=None, # path to lora .pt
 ):
  """
@@ -137,7 +155,10 @@ def inject_trainable_lora(
  _child_module.in_features,
  _child_module.out_features,
  _child_module.bias is not None,
- r,
+ r=r,
+ scale=scale,
+ init=init,
+ nonlin=nonlin,
  )
  _tmp.linear.weight = weight
  if bias is not None:
@@ -333,9 +354,13 @@ def load_safeloras(path, device="cpu"):
 
 
 def weight_apply_lora(
- model, loras, target_replace_module=DEFAULT_TARGET_REPLACE, alpha=1.0
-):
-
+ model,
+ loras,
+ target_replace_module=DEFAULT_TARGET_REPLACE,
+ r: int = 4,
+ scale: float = 1.0,
+ nonlin: nn.Module = None,
+): 
  for _m, _n, _child_module in _find_modules(
  model, target_replace_module, search_class=[nn.Linear]
  ):
@@ -344,13 +369,22 @@ def weight_apply_lora(
  up_weight = loras.pop(0).detach().to(weight.device)
  down_weight = loras.pop(0).detach().to(weight.device)
 
- # W <- W + U * D
- weight = weight + alpha * (up_weight @ down_weight).type(weight.dtype)
+ if nonlin is None:
+ # W <- W + U * D
+ weight = weight + scale * (up_weight @ down_weight).type(weight.dtype)
+ else:
+ # W <- W + U * nonlin(D)
+ weight = weight + scale * (up_weight @ nonlin(down_weight)).type(weight.dtype)
+
  _child_module.weight = nn.Parameter(weight)
 
-
 def monkeypatch_lora(
- model, loras, target_replace_module=DEFAULT_TARGET_REPLACE, r: int = 4
+ model,
+ loras,
+ target_replace_module=DEFAULT_TARGET_REPLACE,
+ r: int = 4,
+ scale: float = 1.0,
+ nonlin: nn.Module = None,
 ):
  for _module, name, _child_module in _find_modules(
  model, target_replace_module, search_class=[nn.Linear]
@@ -362,6 +396,8 @@ def monkeypatch_lora(
  _child_module.out_features,
  _child_module.bias is not None,
  r=r,
+ scale=scale,
+ nonlin=nonlin,
  )
  _tmp.linear.weight = weight
 
@@ -385,7 +421,12 @@ def monkeypatch_lora(
 
 
 def monkeypatch_replace_lora(
- model, loras, target_replace_module=DEFAULT_TARGET_REPLACE, r: int = 4
+ model,
+ loras,
+ target_replace_module=DEFAULT_TARGET_REPLACE,
+ r: int = 4,
+ scale: float = 1.0,
+ nonlin: nn.Module = None,
 ):
  for _module, name, _child_module in _find_modules(
  model, target_replace_module, search_class=[LoraInjectedLinear]
@@ -397,7 +438,9 @@ def monkeypatch_replace_lora(
  _child_module.linear.out_features,
  _child_module.linear.bias is not None,
  r=r,
- )
+ scale=scale,
+ nonlin=nonlin,
+ )
  _tmp.linear.weight = weight
 
  if bias is not None:
@@ -424,6 +467,8 @@ def monkeypatch_or_replace_lora(
  loras,
  target_replace_module=DEFAULT_TARGET_REPLACE,
  r: Union[int, List[int]] = 4,
+ scale: Union[float, List[float]] = 1.0,
+ nonlin: Union[float, List[nn.Module]] = None,
 ):
  for _module, name, _child_module in _find_modules(
  model, target_replace_module, search_class=[nn.Linear, LoraInjectedLinear]
@@ -441,6 +486,8 @@ def monkeypatch_or_replace_lora(
  _source.out_features,
  _source.bias is not None,
  r=r.pop(0) if isinstance(r, list) else r,
+ scale=scale.pop(0) if isinstance(scale, list) else scale,
+ nonlin=nonlin.pop(0) if isinstance(nonlin, list) else nonlin,
  )
  _tmp.linear.weight = weight
 
@@ -496,7 +543,7 @@ def monkeypatch_add_lora(
  model,
  loras,
  target_replace_module=DEFAULT_TARGET_REPLACE,
- alpha: float = 1.0,
+ scale: float = 1.0,
  beta: float = 1.0,
 ):
  for _module, name, _child_module in _find_modules(
@@ -519,12 +566,16 @@ def monkeypatch_add_lora(
  _module._modules[name].to(weight.device)
 
 
-def tune_lora_scale(model, alpha: float = 1.0):
+def tune_lora_scale(model, alpha: float = 1.0, scale: float = None):
+ if alpha:
+ # Keep original named parameter alpha (which is really scale),
+ scale = alpha
+
  for _module in model.modules():
  if _module.__class__.__name__ == "LoraInjectedLinear":
- _module.scale = alpha
-
+ _module.scale = scale
 
+
 def _text_lora_path(path: str) -> str:
  assert path.endswith(".pt"), "Only .pt files are supported"
  return ".".join(path.split(".")[:-1] + ["text_encoder", "pt"])
@@ -576,6 +627,8 @@ def patch_pipe(
  unet_path,
  token: str,
  r: int = 4,
+ scale: float = 1.0,
+ nonlin: nn.Module = None,
  patch_unet=True,
  patch_text=False,
  patch_ti=False,
@@ -596,6 +649,8 @@ def patch_pipe(
  pipe.unet,
  torch.load(unet_path),
  r=r,
+ scale=scale,
+ nonlin=nonlin,
  target_replace_module=unet_target_replace_module,
  )
 
@@ -606,6 +661,8 @@ def patch_pipe(
  torch.load(text_path),
  target_replace_module=text_target_replace_module,
  r=r,
+ scale=scale,
+ nonlin=nonlin,
  )
  if patch_ti:
  print("LoRA : Patching token input")