Spaces:

KumaPower
/

AvatarArtist

Running on Zero

刘虹雨 commited on Mar 31

Commit

57da35d

1 Parent(s): 427d074

update code

Files changed (1) hide show

DiT_VAE/diffusion/model/nets/PixArt_blocks.py CHANGED Viewed

@@ -49,7 +49,7 @@ class MultiHeadCrossAttention(nn.Module):
         k, v = kv.unbind(2)
         attn_bias = None
         if mask is not None:
-            attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask)
         x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias)
         x = x.view(B, -1, C)
         x = self.proj(x)

         k, v = kv.unbind(2)
         attn_bias = None
         if mask is not None:
+            attn_bias = xformers.ops.fmha.attn_bias.BlockDiagonalMask.from_seqlens([N] * B, mask)
         x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias)
         x = x.view(B, -1, C)
         x = self.proj(x)