diff --git a/comfy/ldm/flux/layers.py b/comfy/ldm/flux/layers.py
index 59a62e0df..9e09c40b1 100644
--- a/comfy/ldm/flux/layers.py
+++ b/comfy/ldm/flux/layers.py
@@ -143,7 +143,7 @@ class DoubleStreamBlock(nn.Module):
         )
         self.flipped_img_txt = flipped_img_txt
 
-    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None):
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, **kwargs):
         img_mod1, img_mod2 = self.img_mod(vec)
         txt_mod1, txt_mod2 = self.txt_mod(vec)
 
diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py
index cc34f7585..59c0b73cc 100644
--- a/comfy/ldm/flux/model.py
+++ b/comfy/ldm/flux/model.py
@@ -139,11 +139,10 @@ class Flux(nn.Module):
                 txt = out["txt"]
                 img = out["img"]
             else:
-                img, txt = block(img=img,
-                                 txt=txt,
-                                 vec=vec,
-                                 pe=pe,
-                                 attn_mask=attn_mask)
+               if attn_mask is not None:
+                    img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask)
+                else:
+                    img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
 
             if control is not None: # Controlnet
                 control_i = control.get("input")