re add spanplus without incorrect sudo span detection

TNTwise · Dec 23, 2024 · cced2ed · cced2ed
1 parent 1e1ec62
commit cced2ed
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/backend/src/pytorch/spandrel/__helpers/main_registry.py b/backend/src/pytorch/spandrel/__helpers/main_registry.py
@@ -18,7 +18,7 @@
     SAFMN,
     SAFMNBCIE,
     SPAN,
-    #SPANPlus,
+    SPANPlus,
     sudo_SPANPlus,
     Compact,
     DnCNN,
@@ -91,5 +91,5 @@
     ArchSupport.from_architecture(SeemoRe.SeemoReArch()),
     ArchSupport.from_architecture(MoSR.MoSRArch()),
     ArchSupport.from_architecture(sudo_SPANPlus.sudo_SPANPlusArch()),
-    #ArchSupport.from_architecture(SPANPlus.SPANPlusArch()),
+    ArchSupport.from_architecture(SPANPlus.SPANPlusArch()),
 )
diff --git a/backend/src/pytorch/spandrel/architectures/sudo_SPANPlus/__init__.py b/backend/src/pytorch/spandrel/architectures/sudo_SPANPlus/__init__.py
@@ -17,6 +17,8 @@ def __init__(self):
                 "feats.1.conv_2.eval_conv.weight",
                 "feats.1.conv_cat.weight",
                 "upsampler.end_conv.weight",
+                "dynamic.kernels_weights",
+                "dynamic.attention.to_scores.0.weight",
             ),
         )