xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-10 21:17:18 +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 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) img_mod1, img_mod2 = self.img_mod(vec)
txt_mod1, txt_mod2 = self.txt_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"] txt = out["txt"]
img = out["img"] img = out["img"]
else: else:
img, txt = block(img=img, if attn_mask is not None:
txt=txt, img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask)
vec=vec, else:
pe=pe, img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
attn_mask=attn_mask)
if control is not None: # Controlnet if control is not None: # Controlnet
control_i = control.get("input") control_i = control.get("input")