Update pipelines/flux_pipeline/transformer.py

pipelines/flux_pipeline/transformer.py

@@ -125,9 +125,15 @@ class FluxAttnProcessor2_0:
 
         if neg_mode and FLEX_ATTENTION_AVAILABLE:
             # Apply flex_attention with the block mask
-            global flex_attention_func, block_mask
-            if flex_attention_func is None:
-                flex_attention_func = torch.compile(flex_attention, dynamic=False)
+            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
+            )
+            
+            if need_new_mask:
                 res = int(math.sqrt((end_of_hidden_states-(text_seq if encoder_hidden_states is None else 0)) // num))
                 seq_len = query.shape[2]
 
@@ -155,7 +161,7 @@ class FluxAttnProcessor2_0:
                 block_mask = create_block_mask(block_diagonal_mask, B=1, H=None, 
                                                 Q_LEN=seq_len, KV_LEN=seq_len, device=query.device)
 
-            hidden_states = flex_attention_func(query, key, value, block_mask=block_mask)
+            hidden_states = flex_attention(query, key, value, block_mask=block_mask)
             hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
         elif neg_mode:
             # Fallback to original implementation if flex_attention is not available