From b75b004fe62826455f1aa77e849e7da13902cb17 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Sun, 16 Jul 2023 23:13:55 +0300
Subject: lora extension rework to include other types of networks

---
 extensions-builtin/Lora/network_lyco.py | 39 +++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)
 create mode 100644 extensions-builtin/Lora/network_lyco.py

(limited to 'extensions-builtin/Lora/network_lyco.py')

diff --git a/extensions-builtin/Lora/network_lyco.py b/extensions-builtin/Lora/network_lyco.py
new file mode 100644
index 00000000..18a822fa
--- /dev/null
+++ b/extensions-builtin/Lora/network_lyco.py
@@ -0,0 +1,39 @@
+import torch
+
+import lyco_helpers
+import network
+from modules import devices
+
+
+class NetworkModuleLyco(network.NetworkModule):
+    def __init__(self, net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        if hasattr(self.sd_module, 'weight'):
+            self.shape = self.sd_module.weight.shape
+
+        self.dim = None
+        self.bias = weights.w.get("bias")
+        self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None
+        self.scale = weights.w["scale"].item() if "scale" in weights.w else None
+
+    def finalize_updown(self, updown, orig_weight, output_shape):
+        if self.bias is not None:
+            updown = updown.reshape(self.bias.shape)
+            updown += self.bias.to(orig_weight.device, dtype=orig_weight.dtype)
+            updown = updown.reshape(output_shape)
+
+        if len(output_shape) == 4:
+            updown = updown.reshape(output_shape)
+
+        if orig_weight.size().numel() == updown.size().numel():
+            updown = updown.reshape(orig_weight.shape)
+
+        scale = (
+            self.scale if self.scale is not None
+            else self.alpha / self.dim if self.dim is not None and self.alpha is not None
+            else 1.0
+        )
+
+        return updown * scale * self.network.multiplier
+
-- 
cgit v1.2.1