From afa2399f79e84919645eb69cd8e1ef1d9f1d6bd1 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Fri, 22 Sep 2023 20:26:47 -0400
Subject: [PATCH] Add a way to set output block patches to modify the h and
 hsp.

---
 comfy/ldm/modules/diffusionmodules/openaimodel.py | 6 ++++++
 comfy/model_patcher.py                            | 3 +++
 2 files changed, 9 insertions(+)

diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py
index 3ce3c2e7..b42637c8 100644
--- a/comfy/ldm/modules/diffusionmodules/openaimodel.py
+++ b/comfy/ldm/modules/diffusionmodules/openaimodel.py
@@ -608,6 +608,7 @@ class UNetModel(nn.Module):
         """
         transformer_options["original_shape"] = list(x.shape)
         transformer_options["current_index"] = 0
+        transformer_patches = transformer_options.get("patches", {})
 
         assert (y is not None) == (
             self.num_classes is not None
@@ -644,6 +645,11 @@ class UNetModel(nn.Module):
                 if ctrl is not None:
                     hsp += ctrl
 
+            if "output_block_patch" in transformer_patches:
+                patch = transformer_patches["output_block_patch"]
+                for p in patch:
+                    h, hsp = p(h, hsp, transformer_options)
+
             h = th.cat([h, hsp], dim=1)
             del hsp
             if len(hs) > 0:
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index 10551656..ba505221 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -88,6 +88,9 @@ class ModelPatcher:
     def set_model_attn2_output_patch(self, patch):
         self.set_model_patch(patch, "attn2_output_patch")
 
+    def set_model_output_block_patch(self, patch):
+        self.set_model_patch(patch, "output_block_patch")
+
     def model_patches_to(self, device):
         to = self.model_options["transformer_options"]
         if "patches" in to: