From 3063d31b6053acbb24848e828d467b5bf0318616 Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Fri, 22 Mar 2024 18:45:30 -0700 Subject: [PATCH] Update attention.py --- lvdm/modules/attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lvdm/modules/attention.py b/lvdm/modules/attention.py index fd737cd..a136029 100644 --- a/lvdm/modules/attention.py +++ b/lvdm/modules/attention.py @@ -102,7 +102,7 @@ class CrossAttention(nn.Module): if image_cross_attention_scale_learnable: self.register_parameter('alpha', nn.Parameter(torch.tensor(0.)) ) - def comfy_efficient_forward(self, x, context=None, mask=None): + def comfy_efficient_forward(self, x, context=None, mask=None, *args, **kwargs): spatial_self_attn = (context is None) k_ip, v_ip, out_ip = None, None, None @@ -364,7 +364,7 @@ class BasicTransformerBlock(nn.Module): return checkpoint(self._forward, input_tuple, self.parameters(), self.checkpoint) - def _forward(self, x, context=None, mask=None): + def _forward(self, x, context=None, mask=None, transformer_options={}): extra_options = {} block = transformer_options.get("block", None) block_index = transformer_options.get("block_index", 0)