diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py index 3ce3c2e..b42637c 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 1055165..ba50522 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: