nodes/city96-ComfyUI_ExtraModels
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update PixArt_blocks.py

Browse Source 
Add dropout_p to SDP attention
This commit is contained in:
GavChap
2024-05-11 17:00:48 +01:00
committed by GitHub
parent c2e2f6e810
commit 3824b0d7b9
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
PixArt/models/PixArt_blocks.py
+2 -1
View File
@@ -189,6 +189,7 @@ class AttentionKVCompress(Attention_):
q, k, v = map(lambda t: t.transpose(1, 2),(q, k, v),) q, k, v = map(lambda t: t.transpose(1, 2),(q, k, v),)
x = torch.nn.functional.scaled_dot_product_attention( x = torch.nn.functional.scaled_dot_product_attention(
q, k, v, q, k, v,
dropout_p=self.attn_drop.p,
attn_mask=attn_bias attn_mask=attn_bias
).transpose(1, 2).contiguous() ).transpose(1, 2).contiguous()
x = x.view(B, N, C) x = x.view(B, N, C)
@@ -467,4 +468,4 @@ class CaptionEmbedderDoubleBr(nn.Module):
if (train and use_dropout) or (force_drop_ids is not None): if (train and use_dropout) or (force_drop_ids is not None):
global_caption, caption = self.token_drop(global_caption, caption, force_drop_ids) global_caption, caption = self.token_drop(global_caption, caption, force_drop_ids)
y_embed = self.proj(global_caption) y_embed = self.proj(global_caption)
return y_embed, caption return y_embed, caption