From efe4c983410dfb02185cf3cef4851191e4380f1e Mon Sep 17 00:00:00 2001
From: mgz-dev <49577754+mgz-dev@users.noreply.github.com>
Date: Tue, 28 Feb 2023 14:55:15 -0600
Subject: [PATCH 1/5] Enable ability to resize lora dim based off ratios

---
 networks/resize_lora.py | 44 +++++++++++++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 10 deletions(-)

diff --git a/networks/resize_lora.py b/networks/resize_lora.py
index 271de8ef3..c4d8a4d85 100644
--- a/networks/resize_lora.py
+++ b/networks/resize_lora.py
@@ -38,10 +38,11 @@ def save_to_file(file_name, model, state_dict, dtype, metadata):
     torch.save(model, file_name)
 
 
-def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
+def resize_lora_model(lora_sd, new_rank, save_dtype, device, sv_ratio, verbose):
   network_alpha = None
   network_dim = None
   verbose_str = "\n"
+  ratio_flag = False
 
   CLAMP_QUANTILE = 0.99
 
@@ -57,9 +58,12 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
       network_alpha = network_dim
 
   scale = network_alpha/network_dim
-  new_alpha = float(scale*new_rank)  # calculate new alpha from scale
-
-  print(f"old dimension: {network_dim}, old alpha: {network_alpha}, new alpha: {new_alpha}")
+  if not sv_ratio:
+    new_alpha = float(scale*new_rank)  # calculate new alpha from scale
+    print(f"old dimension: {network_dim}, old alpha: {network_alpha}, new dim: {new_rank}, new alpha: {new_alpha}")
+  else:
+    print(f"Dynamically determining new alphas and dims based off sv ratio: {sv_ratio}")
+    ratio_flag = True
 
   lora_down_weight = None
   lora_up_weight = None
@@ -97,11 +101,24 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
 
         U, S, Vh = torch.linalg.svd(full_weight_matrix)
 
+        if ratio_flag:
+          # Calculate new dim and alpha for dynamic sizing
+          max_sv = S[0]
+          min_sv = max_sv/sv_ratio
+          new_rank = torch.sum(S > min_sv).item()
+          new_rank = max(new_rank, 1)
+          new_alpha = float(scale*new_rank)
+
         if verbose:
           s_sum = torch.sum(torch.abs(S))
           s_rank = torch.sum(torch.abs(S[:new_rank]))
-          verbose_str+=f"{block_down_name:76} | "
-          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}\n"
+          verbose_str+=f"{block_down_name:75} | "
+          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}"
+
+        if verbose and ratio_flag:
+          verbose_str+=f", dynamic| dim: {new_rank}, alpha: {new_alpha}\n"
+        else:
+          verbose_str+=f"\n"
 
         U = U[:, :new_rank]
         S = S[:new_rank]
@@ -160,16 +177,21 @@ def str_to_dtype(p):
   lora_sd, metadata = load_state_dict(args.model, merge_dtype)
 
   print("resizing rank...")
-  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.verbose)
+  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.sv_ratio, args.verbose)
 
   # update metadata
   if metadata is None:
     metadata = {}
 
   comment = metadata.get("ss_training_comment", "")
-  metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
-  metadata["ss_network_dim"] = str(args.new_rank)
-  metadata["ss_network_alpha"] = str(new_alpha)
+  if not args.sv_ratio:
+    metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
+    metadata["ss_network_dim"] = str(args.new_rank)
+    metadata["ss_network_alpha"] = str(new_alpha)
+  else:
+    metadata["ss_training_comment"] = f"Dynamic resize from {old_dim} with ratio {args.sv_ratio}; {comment}"
+    metadata["ss_network_dim"] = 'Dynamic'
+    metadata["ss_network_alpha"] = 'Dynamic'
 
   model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
   metadata["sshs_model_hash"] = model_hash
@@ -193,6 +215,8 @@ def str_to_dtype(p):
   parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
   parser.add_argument("--verbose", action="store_true", 
                       help="Display verbose resizing information / rank変更時の詳細情報を出力する")
+  parser.add_argument("--sv_ratio", type=float, default=None,
+                      help="Specify svd ratio for dim calcs. Will override --new_rank")
 
   args = parser.parse_args()
   resize(args)

From 52ca6c515c14d2309ab470b59a179d7ca2f5e149 Mon Sep 17 00:00:00 2001
From: mgz-dev <49577754+mgz-dev@users.noreply.github.com>
Date: Wed, 1 Mar 2023 13:35:24 -0600
Subject: [PATCH 2/5] add options to resize based off frobenius norm or
 cumulative sum

---
 networks/resize_lora.py | 87 ++++++++++++++++++++++++++++++++---------
 1 file changed, 69 insertions(+), 18 deletions(-)

diff --git a/networks/resize_lora.py b/networks/resize_lora.py
index c4d8a4d85..de4056135 100644
--- a/networks/resize_lora.py
+++ b/networks/resize_lora.py
@@ -3,11 +3,11 @@
 # Thanks to cloneofsimo and kohya
 
 import argparse
-import os
 import torch
 from safetensors.torch import load_file, save_file, safe_open
 from tqdm import tqdm
 from library import train_util, model_util
+import numpy as np
 
 
 def load_state_dict(file_name, dtype):
@@ -38,11 +38,32 @@ def save_to_file(file_name, model, state_dict, dtype, metadata):
     torch.save(model, file_name)
 
 
-def resize_lora_model(lora_sd, new_rank, save_dtype, device, sv_ratio, verbose):
+def index_sv_cumulative(S, target):
+  original_sum = float(torch.sum(S))
+  cumulative_sums = torch.cumsum(S, dim=0)/original_sum
+  index = int(torch.searchsorted(cumulative_sums, target)) + 1
+  if index >= len(S):
+    index = len(S) - 1
+
+  return index
+
+
+def index_sv_fro(S, target):
+  S_squared = S.pow(2)
+  s_fro_sq = float(torch.sum(S_squared))
+  sum_S_squared = torch.cumsum(S_squared, dim=0)/s_fro_sq
+  index = int(torch.searchsorted(sum_S_squared, target**2)) + 1
+  if index >= len(S):
+    index = len(S) - 1
+
+  return index
+
+
+def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose):
   network_alpha = None
   network_dim = None
   verbose_str = "\n"
-  ratio_flag = False
+  fro_list = []
 
   CLAMP_QUANTILE = 0.99
 
@@ -58,12 +79,12 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, sv_ratio, verbose):
       network_alpha = network_dim
 
   scale = network_alpha/network_dim
-  if not sv_ratio:
+
+  if dynamic_method:
+    print(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}")
+  else:
     new_alpha = float(scale*new_rank)  # calculate new alpha from scale
     print(f"old dimension: {network_dim}, old alpha: {network_alpha}, new dim: {new_rank}, new alpha: {new_alpha}")
-  else:
-    print(f"Dynamically determining new alphas and dims based off sv ratio: {sv_ratio}")
-    ratio_flag = True
 
   lora_down_weight = None
   lora_up_weight = None
@@ -101,22 +122,43 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, sv_ratio, verbose):
 
         U, S, Vh = torch.linalg.svd(full_weight_matrix)
 
-        if ratio_flag:
-          # Calculate new dim and alpha for dynamic sizing
+        if dynamic_method=="sv_ratio":
+          # Calculate new dim and alpha based off ratio
           max_sv = S[0]
-          min_sv = max_sv/sv_ratio
+          min_sv = max_sv/dynamic_param
           new_rank = torch.sum(S > min_sv).item()
           new_rank = max(new_rank, 1)
           new_alpha = float(scale*new_rank)
 
+        elif dynamic_method=="sv_cumulative":
+          # Calculate new dim and alpha based off cumulative sum
+          new_rank = index_sv_cumulative(S, dynamic_param)
+          new_rank = max(new_rank, 1)
+          new_alpha = float(scale*new_rank)
+
+        elif dynamic_method=="sv_fro":
+          # Calculate new dim and alpha based off sqrt sum of squares
+          new_rank = index_sv_fro(S, dynamic_param)
+          new_rank = max(new_rank, 1)
+          new_alpha = float(scale*new_rank)
+          
         if verbose:
           s_sum = torch.sum(torch.abs(S))
           s_rank = torch.sum(torch.abs(S[:new_rank]))
+          
+          S_squared = S.pow(2)
+          s_fro = torch.sqrt(torch.sum(S_squared))
+          s_red_fro = torch.sqrt(torch.sum(S_squared[:new_rank]))
+          fro_percent = float(s_red_fro/s_fro)
+          if not np.isnan(fro_percent):
+            fro_list.append(float(fro_percent))
+
           verbose_str+=f"{block_down_name:75} | "
-          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}"
+          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, fro retained: {fro_percent:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}"
 
-        if verbose and ratio_flag:
-          verbose_str+=f", dynamic| dim: {new_rank}, alpha: {new_alpha}\n"
+
+        if verbose and dynamic_method:
+          verbose_str+=f", dynamic | dim: {new_rank}, alpha: {new_alpha}\n"
         else:
           verbose_str+=f"\n"
 
@@ -153,6 +195,8 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, sv_ratio, verbose):
 
   if verbose:
     print(verbose_str)
+
+    print(f"Average Frobenius norm retention: {np.mean(fro_list):.2%} | std: {np.std(fro_list):0.3f}")
   print("resizing complete")
   return o_lora_sd, network_dim, new_alpha
 
@@ -168,6 +212,9 @@ def str_to_dtype(p):
       return torch.bfloat16
     return None
 
+  if args.dynamic_method and not args.dynamic_param:
+    raise Exception("If using dynamic_method, then dynamic_param is required")
+
   merge_dtype = str_to_dtype('float')  # matmul method above only seems to work in float32
   save_dtype = str_to_dtype(args.save_precision)
   if save_dtype is None:
@@ -177,19 +224,20 @@ def str_to_dtype(p):
   lora_sd, metadata = load_state_dict(args.model, merge_dtype)
 
   print("resizing rank...")
-  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.sv_ratio, args.verbose)
+  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.dynamic_method, args.dynamic_param, args.verbose)
 
   # update metadata
   if metadata is None:
     metadata = {}
 
   comment = metadata.get("ss_training_comment", "")
-  if not args.sv_ratio:
+
+  if not args.dynamic_method:
     metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
     metadata["ss_network_dim"] = str(args.new_rank)
     metadata["ss_network_alpha"] = str(new_alpha)
   else:
-    metadata["ss_training_comment"] = f"Dynamic resize from {old_dim} with ratio {args.sv_ratio}; {comment}"
+    metadata["ss_training_comment"] = f"Dynamic resize with {args.dynamic_method}: {args.dynamic_param} from {old_dim}; {comment}"
     metadata["ss_network_dim"] = 'Dynamic'
     metadata["ss_network_alpha"] = 'Dynamic'
 
@@ -215,8 +263,11 @@ def str_to_dtype(p):
   parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
   parser.add_argument("--verbose", action="store_true", 
                       help="Display verbose resizing information / rank変更時の詳細情報を出力する")
-  parser.add_argument("--sv_ratio", type=float, default=None,
-                      help="Specify svd ratio for dim calcs. Will override --new_rank")
+  parser.add_argument("--dynamic_method", type=str, default=None, choices=[None, "sv_ratio", "sv_fro", "sv_cumulative"],
+                      help="Specify dynamic resizing method, will override --new_rank")
+  parser.add_argument("--dynamic_param", type=float, default=None,
+                      help="Specify target for dynamic reduction")
+                                           
 
   args = parser.parse_args()
   resize(args)

From 80be6fa130cf199ed05e9a75762e20f4b280ce7e Mon Sep 17 00:00:00 2001
From: mgz-dev <49577754+mgz-dev@users.noreply.github.com>
Date: Fri, 3 Mar 2023 23:32:46 -0600
Subject: [PATCH 3/5] refactor and bug fix for too large sv_ratio

- code refactor to be able to re-use same function for dynamic extract lora
- remove clamp
- fix issue where if sv_ratio is too high index goes out of bounds
---
 networks/resize_lora.py | 104 +++++++++++++++++++++++-----------------
 1 file changed, 61 insertions(+), 43 deletions(-)

diff --git a/networks/resize_lora.py b/networks/resize_lora.py
index de4056135..eb745333a 100644
--- a/networks/resize_lora.py
+++ b/networks/resize_lora.py
@@ -59,14 +59,55 @@ def index_sv_fro(S, target):
   return index
 
 
+def rank_resize(S, rank, dynamic_method, dynamic_param, scale=1):
+    param_dict = {}
+
+    if dynamic_method=="sv_ratio":
+        # Calculate new dim and alpha based off ratio
+        max_sv = S[0]
+        min_sv = max_sv/dynamic_param
+        new_rank = max(torch.sum(S > min_sv).item(),1)
+        new_alpha = float(scale*new_rank)
+
+    elif dynamic_method=="sv_cumulative":
+        # Calculate new dim and alpha based off cumulative sum
+        new_rank = index_sv_cumulative(S, dynamic_param)
+        new_rank = max(new_rank, 1)
+        new_alpha = float(scale*new_rank)
+
+    elif dynamic_method=="sv_fro":
+        # Calculate new dim and alpha based off sqrt sum of squares
+        new_rank = index_sv_fro(S, dynamic_param)
+        new_rank = min(max(new_rank, 1), len(S)-1)
+        new_alpha = float(scale*new_rank)
+    else:
+        new_rank = rank
+        new_alpha = float(scale*new_rank)
+
+    # Calculate resize info
+    s_sum = torch.sum(torch.abs(S))
+    s_rank = torch.sum(torch.abs(S[:new_rank]))
+    
+    S_squared = S.pow(2)
+    s_fro = torch.sqrt(torch.sum(S_squared))
+    s_red_fro = torch.sqrt(torch.sum(S_squared[:new_rank]))
+    fro_percent = float(s_red_fro/s_fro)
+
+    param_dict["new_rank"] = new_rank
+    param_dict["new_alpha"] = new_alpha
+    param_dict["sum_retained"] = (s_rank)/s_sum
+    param_dict["fro_retained"] = fro_percent
+    param_dict["max_ratio"] = S[0]/S[new_rank]
+
+    return param_dict
+
+
 def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose):
   network_alpha = None
   network_dim = None
   verbose_str = "\n"
   fro_list = []
 
-  CLAMP_QUANTILE = 0.99
-
   # Extract loaded lora dim and alpha
   for key, value in lora_sd.items():
     if network_alpha is None and 'alpha' in key:
@@ -82,9 +123,6 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
 
   if dynamic_method:
     print(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}")
-  else:
-    new_alpha = float(scale*new_rank)  # calculate new alpha from scale
-    print(f"old dimension: {network_dim}, old alpha: {network_alpha}, new dim: {new_rank}, new alpha: {new_alpha}")
 
   lora_down_weight = None
   lora_up_weight = None
@@ -93,7 +131,6 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
   block_down_name = None
   block_up_name = None
 
-  print("resizing lora...")
   with torch.no_grad():
     for key, value in tqdm(lora_sd.items()):
       if 'lora_down' in key:
@@ -122,39 +159,21 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
 
         U, S, Vh = torch.linalg.svd(full_weight_matrix)
 
-        if dynamic_method=="sv_ratio":
-          # Calculate new dim and alpha based off ratio
-          max_sv = S[0]
-          min_sv = max_sv/dynamic_param
-          new_rank = torch.sum(S > min_sv).item()
-          new_rank = max(new_rank, 1)
-          new_alpha = float(scale*new_rank)
-
-        elif dynamic_method=="sv_cumulative":
-          # Calculate new dim and alpha based off cumulative sum
-          new_rank = index_sv_cumulative(S, dynamic_param)
-          new_rank = max(new_rank, 1)
-          new_alpha = float(scale*new_rank)
-
-        elif dynamic_method=="sv_fro":
-          # Calculate new dim and alpha based off sqrt sum of squares
-          new_rank = index_sv_fro(S, dynamic_param)
-          new_rank = max(new_rank, 1)
-          new_alpha = float(scale*new_rank)
-          
+
+        param_dict = rank_resize(S, new_rank, dynamic_method, dynamic_param, scale)
+
+        new_rank = param_dict['new_rank']
+        new_alpha = param_dict['new_alpha']
+
         if verbose:
-          s_sum = torch.sum(torch.abs(S))
-          s_rank = torch.sum(torch.abs(S[:new_rank]))
-          
-          S_squared = S.pow(2)
-          s_fro = torch.sqrt(torch.sum(S_squared))
-          s_red_fro = torch.sqrt(torch.sum(S_squared[:new_rank]))
-          fro_percent = float(s_red_fro/s_fro)
-          if not np.isnan(fro_percent):
-            fro_list.append(float(fro_percent))
+          max_ratio = param_dict['max_ratio']
+          sum_retained = param_dict['sum_retained']
+          fro_retained = param_dict['fro_retained']
+          if not np.isnan(fro_retained):
+            fro_list.append(float(fro_retained))
 
           verbose_str+=f"{block_down_name:75} | "
-          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, fro retained: {fro_percent:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}"
+          verbose_str+=f"sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}"
 
 
         if verbose and dynamic_method:
@@ -168,12 +187,11 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
 
         Vh = Vh[:new_rank, :]
 
-        dist = torch.cat([U.flatten(), Vh.flatten()])
-        hi_val = torch.quantile(dist, CLAMP_QUANTILE)
-        low_val = -hi_val
-
-        U = U.clamp(low_val, hi_val)
-        Vh = Vh.clamp(low_val, hi_val)
+        # dist = torch.cat([U.flatten(), Vh.flatten()])
+        # hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+        # low_val = -hi_val
+        # U = U.clamp(low_val, hi_val)
+        # Vh = Vh.clamp(low_val, hi_val)
 
         if conv2d:
           U = U.unsqueeze(2).unsqueeze(3)
@@ -223,7 +241,7 @@ def str_to_dtype(p):
   print("loading Model...")
   lora_sd, metadata = load_state_dict(args.model, merge_dtype)
 
-  print("resizing rank...")
+  print("Resizing Lora...")
   state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.dynamic_method, args.dynamic_param, args.verbose)
 
   # update metadata

From 214ed092f2208caa5636bb631e7f37ab97c67a3f Mon Sep 17 00:00:00 2001
From: mgz-dev <49577754+mgz-dev@users.noreply.github.com>
Date: Sat, 4 Mar 2023 02:01:10 -0600
Subject: [PATCH 4/5] add support to extract lora with resnet and 2d blocks

Modified resize script so support different types of LoRA networks (refer to Kohaku-Blueleaf module implementation for structure).
---
 networks/resize_lora.py | 109 +++++++++++++++++++++++++++-------------
 1 file changed, 74 insertions(+), 35 deletions(-)

diff --git a/networks/resize_lora.py b/networks/resize_lora.py
index eb745333a..77d79d9f8 100644
--- a/networks/resize_lora.py
+++ b/networks/resize_lora.py
@@ -59,6 +59,72 @@ def index_sv_fro(S, target):
   return index
 
 
+# Modified from Kohaku-blueleaf's extract/merge functions
+def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
+    out_size, in_size, kernel_size, _ = weight.size()
+    U, S, Vh = torch.linalg.svd(weight.reshape(out_size, -1).to(device))
+    
+    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
+
+    lora_rank = param_dict["new_rank"]
+
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+
+    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size, kernel_size, kernel_size).cpu()
+    param_dict["lora_up"] = U.reshape(out_size, lora_rank, 1, 1).cpu()
+    del U, S, Vh, weight
+    return param_dict
+
+
+def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
+    out_size, in_size = weight.size()
+    
+    U, S, Vh = torch.linalg.svd(weight.to(device))
+    
+    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
+    lora_rank = param_dict["new_rank"]
+    
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+    
+    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
+    param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
+    del U, S, Vh, weight
+    return param_dict
+
+
+def merge_conv(lora_down, lora_up, device):
+    in_rank, in_size, kernel_size, k_ = lora_down.shape
+    out_size, out_rank, _, _ = lora_up.shape
+    assert in_rank == out_rank and kernel_size == k_, f"rank {in_rank} {out_rank} or kernel {kernel_size} {k_} mismatch"
+    
+    lora_down = lora_down.to(device)
+    lora_up = lora_up.to(device)
+
+    merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
+    weight = merged.reshape(out_size, in_size, kernel_size, kernel_size)
+    del lora_up, lora_down
+    return weight
+
+
+def merge_linear(lora_down, lora_up, device):
+    in_rank, in_size = lora_down.shape
+    out_size, out_rank = lora_up.shape
+    assert in_rank == out_rank, f"rank {in_rank} {out_rank} mismatch"
+    
+    lora_down = lora_down.to(device)
+    lora_up = lora_up.to(device)
+    
+    weight = lora_up @ lora_down
+    del lora_up, lora_down
+    return weight
+  
+
 def rank_resize(S, rank, dynamic_method, dynamic_param, scale=1):
     param_dict = {}
 
@@ -147,20 +213,11 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
         conv2d = (len(lora_down_weight.size()) == 4)
 
         if conv2d:
-          lora_down_weight = lora_down_weight.squeeze()
-          lora_up_weight = lora_up_weight.squeeze()
-
-        if device:
-          org_device = lora_up_weight.device
-          lora_up_weight = lora_up_weight.to(args.device)
-          lora_down_weight = lora_down_weight.to(args.device)
-
-        full_weight_matrix = torch.matmul(lora_up_weight, lora_down_weight)
-
-        U, S, Vh = torch.linalg.svd(full_weight_matrix)
-
-
-        param_dict = rank_resize(S, new_rank, dynamic_method, dynamic_param, scale)
+          full_weight_matrix = merge_conv(lora_down_weight, lora_up_weight, device)
+          param_dict = extract_conv(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
+        else:
+          full_weight_matrix = merge_linear(lora_down_weight, lora_up_weight, device)
+          param_dict = extract_linear(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
 
         new_rank = param_dict['new_rank']
         new_alpha = param_dict['new_alpha']
@@ -181,28 +238,9 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
         else:
           verbose_str+=f"\n"
 
-        U = U[:, :new_rank]
-        S = S[:new_rank]
-        U = U @ torch.diag(S)
-
-        Vh = Vh[:new_rank, :]
-
-        # dist = torch.cat([U.flatten(), Vh.flatten()])
-        # hi_val = torch.quantile(dist, CLAMP_QUANTILE)
-        # low_val = -hi_val
-        # U = U.clamp(low_val, hi_val)
-        # Vh = Vh.clamp(low_val, hi_val)
-
-        if conv2d:
-          U = U.unsqueeze(2).unsqueeze(3)
-          Vh = Vh.unsqueeze(2).unsqueeze(3)
-
-        if device:
-          U = U.to(org_device)
-          Vh = Vh.to(org_device)
 
-        o_lora_sd[block_down_name + "." + "lora_down.weight"] = Vh.to(save_dtype).contiguous()
-        o_lora_sd[block_up_name + "." + "lora_up.weight"] = U.to(save_dtype).contiguous()
+        o_lora_sd[block_down_name + "." + "lora_down.weight"] = param_dict["lora_down"].to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." + "lora_up.weight"] = param_dict["lora_up"].to(save_dtype).contiguous()
         o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(new_alpha).to(save_dtype)
 
         block_down_name = None
@@ -210,6 +248,7 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
         lora_down_weight = None
         lora_up_weight = None
         weights_loaded = False
+        del param_dict
 
   if verbose:
     print(verbose_str)

From 4a4450d6b6c310ba055c58fdcf06fe4527aeeb2c Mon Sep 17 00:00:00 2001
From: mgz-dev <49577754+mgz-dev@users.noreply.github.com>
Date: Sat, 4 Mar 2023 03:10:04 -0600
Subject: [PATCH 5/5] make new_rank limit max rank, fix zero matrices

-new_rank arg changed to limit the max rank of any layer.
-added logic to make sure zero-ed layers do not create large lora dim
---
 networks/resize_lora.py | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/networks/resize_lora.py b/networks/resize_lora.py
index 77d79d9f8..1a8110c4c 100644
--- a/networks/resize_lora.py
+++ b/networks/resize_lora.py
@@ -9,6 +9,7 @@
 from library import train_util, model_util
 import numpy as np
 
+MIN_SV = 1e-6
 
 def load_state_dict(file_name, dtype):
   if model_util.is_safetensors(file_name):
@@ -65,7 +66,6 @@ def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale
     U, S, Vh = torch.linalg.svd(weight.reshape(out_size, -1).to(device))
     
     param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
-
     lora_rank = param_dict["new_rank"]
 
     U = U[:, :lora_rank]
@@ -150,6 +150,15 @@ def rank_resize(S, rank, dynamic_method, dynamic_param, scale=1):
         new_rank = rank
         new_alpha = float(scale*new_rank)
 
+    
+    if S[0] <= MIN_SV: # Zero matrix, set dim to 1
+        new_rank = 1
+        new_alpha = float(scale*new_rank)
+    elif new_rank > rank: # cap max rank at rank
+        new_rank = rank
+        new_alpha = float(scale*new_rank)
+
+
     # Calculate resize info
     s_sum = torch.sum(torch.abs(S))
     s_rank = torch.sum(torch.abs(S[:new_rank]))
@@ -188,7 +197,7 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
   scale = network_alpha/network_dim
 
   if dynamic_method:
-    print(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}")
+    print(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}, max rank is {new_rank}")
 
   lora_down_weight = None
   lora_up_weight = None
@@ -219,9 +228,6 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
           full_weight_matrix = merge_linear(lora_down_weight, lora_up_weight, device)
           param_dict = extract_linear(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
 
-        new_rank = param_dict['new_rank']
-        new_alpha = param_dict['new_alpha']
-
         if verbose:
           max_ratio = param_dict['max_ratio']
           sum_retained = param_dict['sum_retained']
@@ -232,16 +238,15 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
           verbose_str+=f"{block_down_name:75} | "
           verbose_str+=f"sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}"
 
-
         if verbose and dynamic_method:
-          verbose_str+=f", dynamic | dim: {new_rank}, alpha: {new_alpha}\n"
+          verbose_str+=f", dynamic | dim: {param_dict['new_rank']}, alpha: {param_dict['new_alpha']}\n"
         else:
           verbose_str+=f"\n"
 
-
+        new_alpha = param_dict['new_alpha']
         o_lora_sd[block_down_name + "." + "lora_down.weight"] = param_dict["lora_down"].to(save_dtype).contiguous()
         o_lora_sd[block_up_name + "." + "lora_up.weight"] = param_dict["lora_up"].to(save_dtype).contiguous()
-        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(new_alpha).to(save_dtype)
+        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(param_dict['new_alpha']).to(save_dtype)
 
         block_down_name = None
         block_up_name = None
@@ -321,7 +326,7 @@ def str_to_dtype(p):
   parser.add_argument("--verbose", action="store_true", 
                       help="Display verbose resizing information / rank変更時の詳細情報を出力する")
   parser.add_argument("--dynamic_method", type=str, default=None, choices=[None, "sv_ratio", "sv_fro", "sv_cumulative"],
-                      help="Specify dynamic resizing method, will override --new_rank")
+                      help="Specify dynamic resizing method, --new_rank is used as a hard limit for max rank")
   parser.add_argument("--dynamic_param", type=float, default=None,
                       help="Specify target for dynamic reduction")