From 4e345b31f692d5fb89009bf3352c922c2abe30e2 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 24 Apr 2023 02:36:06 -0400
Subject: [PATCH] Support all known hypernetworks.

---
 comfy_extras/nodes_hypernetwork.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/comfy_extras/nodes_hypernetwork.py b/comfy_extras/nodes_hypernetwork.py
index db2f8695..c08c2c81 100644
--- a/comfy_extras/nodes_hypernetwork.py
+++ b/comfy_extras/nodes_hypernetwork.py
@@ -10,7 +10,17 @@ def load_hypernetwork_patch(path, strength):
     activate_output = sd.get('activate_output', False)
     last_layer_dropout = sd.get('last_layer_dropout', False)
 
-    if activation_func != 'linear' or is_layer_norm != False or use_dropout != False or activate_output != False or last_layer_dropout != False:
+    valid_activation = {
+        "linear": torch.nn.Identity,
+        "relu": torch.nn.ReLU,
+        "leakyrelu": torch.nn.LeakyReLU,
+        "elu": torch.nn.ELU,
+        "swish": torch.nn.Hardswish,
+        "tanh": torch.nn.Tanh,
+        "sigmoid": torch.nn.Sigmoid,
+    }
+
+    if activation_func not in valid_activation:
         print("Unsupported Hypernetwork format, if you report it I might implement it.", path, " ", activation_func, is_layer_norm, use_dropout, activate_output, last_layer_dropout)
         return None
 
@@ -28,15 +38,27 @@ def load_hypernetwork_patch(path, strength):
             keys = attn_weights.keys()
 
             linears = filter(lambda a: a.endswith(".weight"), keys)
-            linears = sorted(list(map(lambda a: a[:-len(".weight")], linears)))
+            linears = list(map(lambda a: a[:-len(".weight")], linears))
             layers = []
 
-            for lin_name in linears:
+            for i in range(len(linears)):
+                lin_name = linears[i]
+                last_layer = (i == (len(linears) - 1))
+                penultimate_layer = (i == (len(linears) - 2))
+
                 lin_weight = attn_weights['{}.weight'.format(lin_name)]
                 lin_bias = attn_weights['{}.bias'.format(lin_name)]
                 layer = torch.nn.Linear(lin_weight.shape[1], lin_weight.shape[0])
                 layer.load_state_dict({"weight": lin_weight, "bias": lin_bias})
-                layers += [layer]
+                layers.append(layer)
+                if activation_func != "linear":
+                    if (not last_layer) or (activate_output):
+                        layers.append(valid_activation[activation_func]())
+                if is_layer_norm:
+                    layers.append(torch.nn.LayerNorm(lin_weight.shape[0]))
+                if use_dropout:
+                    if (not last_layer) and (not penultimate_layer or last_layer_dropout):
+                        layers.append(torch.nn.Dropout(p=0.3))
 
             output.append(torch.nn.Sequential(*layers))
         out[dim] = torch.nn.ModuleList(output)