Merge branch 'v2.3' into dev/codeowner-fix-2.3

Author: lstein

ldm/invoke/config/model_install.py

@@ -196,16 +196,6 @@ def create(self):
             scroll_exit=True,
         )
         self.nextrely += 1
-        self.convert_models = self.add_widget_intelligent(
-            npyscreen.TitleSelectOne,
-            name="== CONVERT IMPORTED MODELS INTO DIFFUSERS==",
-            values=["Keep original format", "Convert to diffusers"],
-            value=0,
-            begin_entry_at=4,
-            max_height=4,
-            hidden=True,  # will appear when imported models box is edited
-            scroll_exit=True,
-        )
         self.cancel = self.add_widget_intelligent(
             npyscreen.ButtonPress,
             name="CANCEL",
@@ -240,8 +230,6 @@ def create(self):
             self.show_directory_fields.addVisibleWhenSelected(i)
 
         self.show_directory_fields.when_value_edited = self._clear_scan_directory
-        self.import_model_paths.when_value_edited = self._show_hide_convert
-        self.autoload_directory.when_value_edited = self._show_hide_convert
 
     def resize(self):
         super().resize()
@@ -252,13 +240,6 @@ def _clear_scan_directory(self):
         if not self.show_directory_fields.value:
             self.autoload_directory.value = ""
 
-    def _show_hide_convert(self):
-        model_paths = self.import_model_paths.value or ""
-        autoload_directory = self.autoload_directory.value or ""
-        self.convert_models.hidden = (
-            len(model_paths) == 0 and len(autoload_directory) == 0
-        )
-
     def _get_starter_model_labels(self) -> List[str]:
         window_width, window_height = get_terminal_size()
         label_width = 25
@@ -318,7 +299,6 @@ def marshall_arguments(self):
         .scan_directory: Path to a directory of models to scan and import
         .autoscan_on_startup:  True if invokeai should scan and import at startup time
         .import_model_paths:   list of URLs, repo_ids and file paths to import
-        .convert_to_diffusers: if True, convert legacy checkpoints into diffusers
         """
         # we're using a global here rather than storing the result in the parentapp
         # due to some bug in npyscreen that is causing attributes to be lost
@@ -354,7 +334,6 @@ def marshall_arguments(self):
 
         # URLs and the like
         selections.import_model_paths = self.import_model_paths.value.split()
-        selections.convert_to_diffusers = self.convert_models.value[0] == 1
 
 
 class AddModelApplication(npyscreen.NPSAppManaged):
@@ -367,7 +346,6 @@ def __init__(self):
             scan_directory=None,
             autoscan_on_startup=None,
             import_model_paths=None,
-            convert_to_diffusers=None,
         )
 
     def onStart(self):
@@ -387,15 +365,13 @@ def process_and_execute(opt: Namespace, selections: Namespace):
     directory_to_scan = selections.scan_directory
     scan_at_startup = selections.autoscan_on_startup
     potential_models_to_install = selections.import_model_paths
-    convert_to_diffusers = selections.convert_to_diffusers
 
     install_requested_models(
         install_initial_models=models_to_install,
         remove_models=models_to_remove,
         scan_directory=Path(directory_to_scan) if directory_to_scan else None,
         external_models=potential_models_to_install,
         scan_at_startup=scan_at_startup,
-        convert_to_diffusers=convert_to_diffusers,
         precision="float32"
         if opt.full_precision
         else choose_precision(torch.device(choose_torch_device())),

ldm/invoke/config/model_install_backend.py

@@ -68,7 +68,6 @@ def install_requested_models(
         scan_directory: Path = None,
         external_models: List[str] = None,
         scan_at_startup: bool = False,
-        convert_to_diffusers: bool = False,
         precision: str = "float16",
         purge_deleted: bool = False,
         config_file_path: Path = None,
@@ -111,20 +110,20 @@ def install_requested_models(
     if len(external_models)>0:
         print("== INSTALLING EXTERNAL MODELS ==")
         for path_url_or_repo in external_models:
+            print(f'DEBUG: path_url_or_repo = {path_url_or_repo}')
             try:
                 model_manager.heuristic_import(
                     path_url_or_repo,
-                    convert=convert_to_diffusers,
                     config_file_callback=_pick_configuration_file,
                     commit_to_conf=config_file_path
                 )
             except KeyboardInterrupt:
                 sys.exit(-1)
-            except Exception:
-                pass
+            except Exception as e:
+                print(f'An exception has occurred: {str(e)}')
 
     if scan_at_startup and scan_directory.is_dir():
-        argument = '--autoconvert' if convert_to_diffusers else '--autoimport'
+        argument = '--autoconvert'
         initfile = Path(Globals.root, Globals.initfile)
         replacement = Path(Globals.root, f'{Globals.initfile}.new')
         directory = str(scan_directory).replace('\\','/')

ldm/modules/kohya_lora_manager.py

@@ -31,18 +31,13 @@ def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0):
         self.name = name
         self.scale = alpha / rank if (alpha and rank) else 1.0
 
-    def forward(self, lora, input_h, output):
+    def forward(self, lora, input_h):
         if self.mid is None:
-            output = (
-                output
-                + self.up(self.down(*input_h)) * lora.multiplier * self.scale
-            )
+            weight = self.up(self.down(*input_h))
         else:
-            output = (
-                output
-                + self.up(self.mid(self.down(*input_h))) * lora.multiplier * self.scale
-            )
-        return output
+            weight = self.up(self.mid(self.down(*input_h)))
+
+        return weight * lora.multiplier * self.scale
 
 class LoHALayer:
     lora_name: str
@@ -64,7 +59,7 @@ def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0):
         self.name = name
         self.scale = alpha / rank if (alpha and rank) else 1.0
 
-    def forward(self, lora, input_h, output):
+    def forward(self, lora, input_h):
 
         if type(self.org_module) == torch.nn.Conv2d:
             op = torch.nn.functional.conv2d
@@ -86,16 +81,79 @@ def forward(self, lora, input_h, output):
             rebuild1 = torch.einsum('i j k l, j r, i p -> p r k l', self.t1, self.w1_b, self.w1_a)
             rebuild2 = torch.einsum('i j k l, j r, i p -> p r k l', self.t2, self.w2_b, self.w2_a)
             weight = rebuild1 * rebuild2
-        
+
         bias = self.bias if self.bias is not None else 0
-        return output + op(
+        return op(
             *input_h,
             (weight + bias).view(self.org_module.weight.shape),
             None,
             **extra_args,
         ) * lora.multiplier * self.scale
 
 
+class LoKRLayer:
+    lora_name: str
+    name: str
+    scale: float
+
+    w1: Optional[torch.Tensor] = None
+    w1_a: Optional[torch.Tensor] = None
+    w1_b: Optional[torch.Tensor] = None
+    w2: Optional[torch.Tensor] = None
+    w2_a: Optional[torch.Tensor] = None
+    w2_b: Optional[torch.Tensor] = None
+    t2: Optional[torch.Tensor] = None
+    bias: Optional[torch.Tensor] = None
+
+    org_module: torch.nn.Module
+
+    def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0):
+        self.lora_name = lora_name
+        self.name = name
+        self.scale = alpha / rank if (alpha and rank) else 1.0
+
+    def forward(self, lora, input_h):
+
+        if type(self.org_module) == torch.nn.Conv2d:
+            op = torch.nn.functional.conv2d
+            extra_args = dict(
+                stride=self.org_module.stride,
+                padding=self.org_module.padding,
+                dilation=self.org_module.dilation,
+                groups=self.org_module.groups,
+            )
+
+        else:
+            op = torch.nn.functional.linear
+            extra_args = {}
+
+        w1 = self.w1
+        if w1 is None:
+            w1 = self.w1_a @ self.w1_b
+
+        w2 = self.w2
+        if w2 is None:
+            if self.t2 is None:
+                w2 = self.w2_a @ self.w2_b
+            else:
+                w2 = torch.einsum('i j k l, i p, j r -> p r k l', self.t2, self.w2_a, self.w2_b)
+
+
+        if len(w2.shape) == 4:
+            w1 = w1.unsqueeze(2).unsqueeze(2)
+        w2 = w2.contiguous()
+        weight = torch.kron(w1, w2).reshape(self.org_module.weight.shape)
+
+
+        bias = self.bias if self.bias is not None else 0
+        return op(
+            *input_h, 
+            (weight + bias).view(self.org_module.weight.shape),
+            None,
+            **extra_args
+        ) * lora.multiplier * self.scale
+
+
 class LoRAModuleWrapper:
     unet: UNet2DConditionModel
     text_encoder: CLIPTextModel
@@ -159,7 +217,7 @@ def lora_forward(module, input_h, output):
                 layer = lora.layers.get(name, None)
                 if layer is None:
                     continue
-                output = layer.forward(lora, input_h, output)
+                output += layer.forward(lora, input_h)
             return output
 
         return lora_forward
@@ -307,6 +365,36 @@ def load_from_dict(self, state_dict):
                 else:
                     layer.t2 = None
 
+            # lokr
+            elif "lokr_w1_b" in values or "lokr_w1" in values:
+
+                if "lokr_w1_b" in values:
+                    rank = values["lokr_w1_b"].shape[0]
+                elif "lokr_w2_b" in values:
+                    rank = values["lokr_w2_b"].shape[0]
+                else:
+                    rank = None # unscaled
+
+                layer = LoKRLayer(self.name, stem, rank, alpha)
+                layer.org_module = wrapped
+                layer.bias = bias
+
+                if "lokr_w1" in values:
+                    layer.w1 = values["lokr_w1"].to(device=self.device, dtype=self.dtype)
+                else:
+                    layer.w1_a = values["lokr_w1_a"].to(device=self.device, dtype=self.dtype)
+                    layer.w1_b = values["lokr_w1_b"].to(device=self.device, dtype=self.dtype)
+
+                if "lokr_w2" in values:
+                    layer.w2 = values["lokr_w2"].to(device=self.device, dtype=self.dtype)
+                else:
+                    layer.w2_a = values["lokr_w2_a"].to(device=self.device, dtype=self.dtype)
+                    layer.w2_b = values["lokr_w2_b"].to(device=self.device, dtype=self.dtype)
+
+                if "lokr_t2" in values:
+                    layer.t2 = values["lokr_t2"].to(device=self.device, dtype=self.dtype)
+
+
             else:
                 print(
                     f">> Encountered unknown lora layer module in {self.name}: {stem} - {type(wrapped).__name__}"