From 58c3df32f3a73b20ea33d1709a1d25818b8a98dd Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Mon, 17 Jul 2023 00:12:18 +0300
Subject: [PATCH] IA3 support

---
 extensions-builtin/Lora/network_ia3.py | 32 ++++++++++++++++++++++++++
 extensions-builtin/Lora/networks.py    |  2 ++
 2 files changed, 34 insertions(+)
 create mode 100644 extensions-builtin/Lora/network_ia3.py

diff --git a/extensions-builtin/Lora/network_ia3.py b/extensions-builtin/Lora/network_ia3.py
new file mode 100644
index 000000000..99f2307ce
--- /dev/null
+++ b/extensions-builtin/Lora/network_ia3.py
@@ -0,0 +1,32 @@
+import lyco_helpers
+import network
+import network_lyco
+
+
+class ModuleTypeIa3(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["weight"]):
+            return NetworkModuleIa3(net, weights)
+
+        return None
+
+
+class NetworkModuleIa3(network_lyco.NetworkModuleLyco):
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        self.w = weights.w["weight"]
+        self.on_input = weights.w["on_input"].item()
+
+    def calc_updown(self, orig_weight):
+        w = self.w.to(orig_weight.device, dtype=orig_weight.dtype)
+
+        output_shape = [w.size(0), orig_weight.size(1)]
+        if self.on_input:
+            output_shape.reverse()
+        else:
+            w = w.reshape(-1, 1)
+
+        updown = orig_weight * w
+
+        return self.finalize_updown(updown, orig_weight, output_shape)
diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index 90374faaa..bf810b5b8 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -4,6 +4,7 @@ import re
 import network
 import network_lora
 import network_hada
+import network_ia3
 
 import torch
 from typing import Union
@@ -13,6 +14,7 @@ from modules import shared, devices, sd_models, errors, scripts, sd_hijack
 module_types = [
     network_lora.ModuleTypeLora(),
     network_hada.ModuleTypeHada(),
+    network_ia3.ModuleTypeIa3(),
 ]