xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-10 12:27:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix attn_mask issue in DoubleStreamBlock.forward(); updated model.py for compatibility

Browse Source
This commit is contained in:
Aashish Dhawan 2025-02-25 15:51:26 -05:00
parent f40076096e
commit decce572f3
2 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/ldm/flux/layers.py
View File

@ -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)

9
comfy/ldm/flux/model.py
View File

@ -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")