From 94cc718e9c42cb4de337293b66dd42fb594b9cae Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 14 Nov 2023 00:08:12 -0500 Subject: [PATCH] Add a way to add patches to the input block. --- comfy/ldm/modules/diffusionmodules/openaimodel.py | 5 +++++ comfy/model_patcher.py | 3 +++ 2 files changed, 8 insertions(+) diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py index 49c1e8c..cac0dfb 100644 --- a/comfy/ldm/modules/diffusionmodules/openaimodel.py +++ b/comfy/ldm/modules/diffusionmodules/openaimodel.py @@ -624,6 +624,11 @@ class UNetModel(nn.Module): 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 index 1c36855..0236843 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -96,6 +96,9 @@ class ModelPatcher: 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")