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

Add switch for new xformers

Browse Source 
#86 - patch provided by @SaiZyca
This commit is contained in:
City
2024-12-11 22:48:51 +01:00
parent ee982c5d73
commit 4a11543beb
2 changed files with 12 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sana/models/sana_blocks.py
+5 -2
View File
@@ -39,8 +39,11 @@ Q_4GB_LIMIT = 32000000
from comfy import model_management
if model_management.xformers_enabled():
import xformers
import xformers.ops
if int((xformers.__version__).split(".")[2]) >= 28:
block_diagonal_mask_from_seqlens = xformers.ops.fmha.attn_bias.BlockDiagonalMask.from_seqlens
else:
block_diagonal_mask_from_seqlens = xformers.ops.fmha.BlockDiagonalMask.from_seqlens
else:
if model_management.xpu_available:
import intel_extension_for_pytorch as ipex # type: ignore
@@ -94,7 +97,7 @@ class MultiHeadCrossAttention(nn.Module):
if model_management.xformers_enabled():
attn_bias = None
if mask is not None:
attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask)
attn_bias = block_diagonal_mask_from_seqlens([N] * B, mask)
x = xformers.ops.memory_efficient_attention(
q, k, v,
p=self.attn_drop.p,