1 files changed, 55 insertions, 0 deletions
diff --git a/modules/hypernetwork.py b/modules/hypernetwork.py
new file mode 100644
index 00000000..9ed1eed9
--- /dev/null
+++ b/modules/hypernetwork.py
@@ -0,0 +1,55 @@
+import glob
+import os
+import torch
+from modules import devices
+
+
+class HypernetworkModule(torch.nn.Module):
+    def __init__(self, dim, state_dict):
+        super().__init__()
+
+        self.linear1 = torch.nn.Linear(dim, dim * 2)
+        self.linear2 = torch.nn.Linear(dim * 2, dim)
+
+        self.load_state_dict(state_dict, strict=True)
+        self.to(devices.device)
+
+    def forward(self, x):
+        return x + (self.linear2(self.linear1(x)))
+
+
+class Hypernetwork:
+    filename = None
+    name = None
+
+    def __init__(self, filename):
+        self.filename = filename
+        self.name = os.path.splitext(os.path.basename(filename))[0]
+        self.layers = {}
+
+        state_dict = torch.load(filename, map_location='cpu')
+        for size, sd in state_dict.items():
+            self.layers[size] = (HypernetworkModule(size, sd[0]), HypernetworkModule(size, sd[1]))
+
+
+def load_hypernetworks(path):
+    res = {}
+
+    for filename in glob.iglob(path + '**/*.pt', recursive=True):
+        hn = Hypernetwork(filename)
+        res[hn.name] = hn
+
+    return res
+
+def apply(self, x, context=None, mask=None, original=None):
+
+
+    if CrossAttention.hypernetwork is not None and context.shape[2] in CrossAttention.hypernetwork:
+        if context.shape[1] == 77 and CrossAttention.noise_cond:
+            context = context + (torch.randn_like(context) * 0.1)
+        h_k, h_v = CrossAttention.hypernetwork[context.shape[2]]
+        k = self.to_k(h_k(context))
+        v = self.to_v(h_v(context))
+    else:
+        k = self.to_k(context)
+        v = self.to_v(context)