Merge decce572f3bed95c31e719167815ece852afcfa4 into 29832b3b61591633d8f312f7df727c1bb8b4d9e4

This commit is contained in:
Aashish Dhawan 2025-03-08 08:43:32 -05:00 committed by GitHub
commit cdd8f77006
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 6 deletions

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)

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