From c3b2a8cb2755c1456746c03c9935a38970ad1503 Mon Sep 17 00:00:00 2001
From: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com>
Date: Thu, 8 Feb 2024 22:37:52 -0500
Subject: [PATCH] Quick Seamless Fixes

---
 invokeai/backend/model_management/seamless.py | 87 ++++++++-----------
 1 file changed, 34 insertions(+), 53 deletions(-)

diff --git a/invokeai/backend/model_management/seamless.py b/invokeai/backend/model_management/seamless.py
index bfdf9e0c53..3ab2db1d90 100644
--- a/invokeai/backend/model_management/seamless.py
+++ b/invokeai/backend/model_management/seamless.py
@@ -28,68 +28,49 @@ def _conv_forward_asymmetric(self, input, weight, bias):
 def set_seamless(model: Union[UNet2DConditionModel, AutoencoderKL], seamless_axes: List[str]):
     try:
         to_restore = []
+        skipped_layers = 0
+        skip_second_resnet = True
+        skip_conv2 = True
 
         for m_name, m in model.named_modules():
-            if isinstance(model, UNet2DConditionModel):
-                if ".attentions." in m_name:
+            if not isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+                continue
+
+            if isinstance(model, UNet2DConditionModel) and m_name.startswith("down_blocks.") and ".resnets." in m_name:
+                # down_blocks.1.resnets.1.conv1
+                _, block_num, _, resnet_num, submodule_name = m_name.split(".")
+                block_num = int(block_num)
+                resnet_num = int(resnet_num)
+
+                # if block_num >= seamless_down_blocks:
+                if block_num >= len(model.down_blocks) - skipped_layers:
                     continue
 
-                if ".resnets." in m_name:
-                    if ".conv2" in m_name:
-                        continue
-                    if ".conv_shortcut" in m_name:
-                        continue
-
-            """
-            if isinstance(model, UNet2DConditionModel):
-                if False and ".upsamplers." in m_name:
+                if resnet_num > 0 and skip_second_resnet:
                     continue
 
-                if False and ".downsamplers." in m_name:
+                if submodule_name == "conv2" and skip_conv2:
                     continue
 
-                if True and ".resnets." in m_name:
-                    if True and ".conv1" in m_name:
-                        if False and "down_blocks" in m_name:
-                            continue
-                        if False and "mid_block" in m_name:
-                            continue
-                        if False and "up_blocks" in m_name:
-                            continue
+            m.asymmetric_padding_mode = {}
+            m.asymmetric_padding = {}
+            m.asymmetric_padding_mode["x"] = "circular" if ("x" in seamless_axes) else "constant"
+            m.asymmetric_padding["x"] = (
+                m._reversed_padding_repeated_twice[0],
+                m._reversed_padding_repeated_twice[1],
+                0,
+                0,
+            )
+            m.asymmetric_padding_mode["y"] = "circular" if ("y" in seamless_axes) else "constant"
+            m.asymmetric_padding["y"] = (
+                0,
+                0,
+                m._reversed_padding_repeated_twice[2],
+                m._reversed_padding_repeated_twice[3],
+            )
 
-                    if True and ".conv2" in m_name:
-                        continue
-
-                    if True and ".conv_shortcut" in m_name:
-                        continue
-
-                if True and ".attentions." in m_name:
-                    continue
-
-                if False and m_name in ["conv_in", "conv_out"]:
-                    continue
-            """
-
-            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
-                m.asymmetric_padding_mode = {}
-                m.asymmetric_padding = {}
-                m.asymmetric_padding_mode["x"] = "circular" if ("x" in seamless_axes) else "constant"
-                m.asymmetric_padding["x"] = (
-                    m._reversed_padding_repeated_twice[0],
-                    m._reversed_padding_repeated_twice[1],
-                    0,
-                    0,
-                )
-                m.asymmetric_padding_mode["y"] = "circular" if ("y" in seamless_axes) else "constant"
-                m.asymmetric_padding["y"] = (
-                    0,
-                    0,
-                    m._reversed_padding_repeated_twice[2],
-                    m._reversed_padding_repeated_twice[3],
-                )
-
-                to_restore.append((m, m._conv_forward))
-                m._conv_forward = _conv_forward_asymmetric.__get__(m, nn.Conv2d)
+            to_restore.append((m, m._conv_forward))
+            m._conv_forward = _conv_forward_asymmetric.__get__(m, nn.Conv2d)
 
         yield