add protection

zhtmike · zhtmike · commit bc61551bf62a · 2026-04-16T10:59:24.000+08:00
diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py
@@ -2496,7 +2496,7 @@ def _templated_context_parallel_attention(
 
 @_AttentionBackendRegistry.register(
     AttentionBackendName.FLASH,
-    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
+    constraints=[_check_attn_mask_or_causal, _check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
     supports_context_parallel=True,
 )
 def _flash_attention(

Original file line number	Diff line number	Diff line change
`@@ -2496,7 +2496,7 @@ def _templated_context_parallel_attention(`
`2496`	`2496`
`2497`	`2497`	`@_AttentionBackendRegistry.register(`
`2498`	`2498`	`AttentionBackendName.FLASH,`
`2499`		`- constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],`
	`2499`	`+ constraints=[_check_attn_mask_or_causal, _check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],`
`2500`	`2500`	`supports_context_parallel=True,`
`2501`	`2501`	`)`
`2502`	`2502`	`def _flash_attention(`