[diffusion]: fix sparse video gen 2 backend being applied to cross-attention (#18900)
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
@@ -313,18 +313,7 @@ class WanTransformerBlock(nn.Module):
|
|||||||
tp_size = get_tp_world_size()
|
tp_size = get_tp_world_size()
|
||||||
self.local_num_heads = divide(num_heads, tp_size)
|
self.local_num_heads = divide(num_heads, tp_size)
|
||||||
self_attn_backends = supported_attention_backends
|
self_attn_backends = supported_attention_backends
|
||||||
cross_attn_backends = supported_attention_backends
|
|
||||||
if (
|
|
||||||
supported_attention_backends is not None
|
|
||||||
and AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN
|
|
||||||
in supported_attention_backends
|
|
||||||
):
|
|
||||||
cross_attn_backends = supported_attention_backends.copy()
|
|
||||||
cross_attn_backends.remove(AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN)
|
|
||||||
logger.debug(
|
|
||||||
"Sparse Video Gen 2 attention backend is not supported for cross-attention; "
|
|
||||||
"removing SPARSE_VIDEO_GEN_2_ATTN from cross-attention backends."
|
|
||||||
)
|
|
||||||
if attention_type in ("sla", "sagesla"):
|
if attention_type in ("sla", "sagesla"):
|
||||||
self.attn1 = MinimalA2AAttnOp(
|
self.attn1 = MinimalA2AAttnOp(
|
||||||
num_heads=self.local_num_heads,
|
num_heads=self.local_num_heads,
|
||||||
|
|||||||
@@ -46,6 +46,7 @@ class AttentionBackendEnum(enum.Enum):
|
|||||||
return self in {
|
return self in {
|
||||||
AttentionBackendEnum.SLIDING_TILE_ATTN,
|
AttentionBackendEnum.SLIDING_TILE_ATTN,
|
||||||
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
|
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
|
||||||
|
AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
|
||||||
AttentionBackendEnum.VMOBA_ATTN,
|
AttentionBackendEnum.VMOBA_ATTN,
|
||||||
AttentionBackendEnum.SLA_ATTN,
|
AttentionBackendEnum.SLA_ATTN,
|
||||||
AttentionBackendEnum.SAGE_SLA_ATTN,
|
AttentionBackendEnum.SAGE_SLA_ATTN,
|
||||||
|
|||||||
Reference in New Issue
Block a user