From f04229b84d3d127af5168188b86b16c9e64f3725 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 4 Sep 2024 13:35:15 -0500 Subject: [PATCH] Add emb_patch support to UNetModel forward (#4779) --- comfy/ldm/modules/diffusionmodules/openaimodel.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py index 6535e899..2902073d 100644 --- a/comfy/ldm/modules/diffusionmodules/openaimodel.py +++ b/comfy/ldm/modules/diffusionmodules/openaimodel.py @@ -842,6 +842,11 @@ class UNetModel(nn.Module): t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype) emb = self.time_embed(t_emb) + if "emb_patch" in transformer_patches: + patch = transformer_patches["emb_patch"] + for p in patch: + emb = p(emb, self.model_channels, transformer_options) + if self.num_classes is not None: assert y.shape[0] == x.shape[0] emb = emb + self.label_emb(y)