Update pipelines/flux_pipeline/transformer.py

pipelines/flux_pipeline/transformer.py

@@ -126,12 +126,7 @@ class FluxAttnProcessor2_0:
         if neg_mode and FLEX_ATTENTION_AVAILABLE:
             # Apply flex_attention with the block mask
             global block_mask
-            need_new_mask = (
-                block_mask is None
-                or block_mask.shape[-2] != query.shape[2]
-                or block_mask.shape[-1] != query.shape[2]
-                or block_mask.device != query.device
-            )
+            need_new_mask = block_mask is None
             
             if need_new_mask:
                 res = int(math.sqrt((end_of_hidden_states-(text_seq if encoder_hidden_states is None else 0)) // num))