Add a way to add patches to the input block.

MoonRide303 · Nov 14, 2023 · 94cc718 · 94cc718
1 parent 8509bd5
commit 94cc718
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py
@@ -624,6 +624,11 @@ def forward(self, x, timesteps=None, context=None, y=None, control=None, transfo
             transformer_options["block"] = ("input", id)
             h = forward_timestep_embed(module, h, emb, context, transformer_options)
             h = apply_control(h, control, 'input')
+            if "input_block_patch" in transformer_patches:
+                patch = transformer_patches["input_block_patch"]
+                for p in patch:
+                    h = p(h, transformer_options)
+
             hs.append(h)
 
         transformer_options["block"] = ("middle", 0)

diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
@@ -96,6 +96,9 @@ def set_model_attn1_output_patch(self, patch):
     def set_model_attn2_output_patch(self, patch):
         self.set_model_patch(patch, "attn2_output_patch")
 
+    def set_model_input_block_patch(self, patch):
+        self.set_model_patch(patch, "input_block_patch")
+
     def set_model_output_block_patch(self, patch):
         self.set_model_patch(patch, "output_block_patch")