Trying text_mask 7

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -301,13 +301,6 @@ def wrap_flash_attention(query, key, value):
     # If attention_mask is provided, apply it to kv_segment_ids
     # attention_mask shape: (B, original_kv_seq_len) with 1 for real tokens, 0 for padded
     if attention_mask is not None:
-      # DEBUG: Check shapes and values
-      jax.debug.print("[DEBUG _tpu_flash_attention] attention_mask shape: {}", attention_mask.shape)
-      jax.debug.print("[DEBUG _tpu_flash_attention] attention_mask[0] sum: {}", attention_mask[0].sum())
-      if attention_mask.shape[0] > 1:
-        jax.debug.print("[DEBUG _tpu_flash_attention] attention_mask[1] sum: {}", attention_mask[1].sum())
-      jax.debug.print("[DEBUG _tpu_flash_attention] key shape: {}, key_seq_len: {}, kv_padded_len: {}", 
-                      key.shape, key_seq_len, kv_padded_len)
 
       # For CFG, different batch items have different padding patterns (pos vs neg prompts)
       # We need a per-batch mask, but segment_ids need to be 1D for vmapped_splash
@@ -316,7 +309,6 @@ def wrap_flash_attention(query, key, value):
       mask_len = min(key_seq_len, attention_mask.shape[1])
       kv_mask_any = jnp.max(attention_mask[:, :mask_len], axis=0).astype(jnp.int32)  # (mask_len,)
 
-      jax.debug.print("[DEBUG _tpu_flash_attention] Using OR across batch, sum: {}", kv_mask_any.sum())
 
       # If key_seq_len > mask_len, pad the mask with 1s (assume remaining tokens are valid)
       if key_seq_len > mask_len:
@@ -332,7 +324,6 @@ def wrap_flash_attention(query, key, value):
       # Both are (kv_padded_len,) - element-wise multiplication
       kv_segment_ids = (kv_segment_ids * kv_mask_padded).astype(jnp.int32)
 
-      jax.debug.print("[DEBUG _tpu_flash_attention] Final kv_segment_ids sum: {}", kv_segment_ids.sum())
 
     segment_ids = splash_attention_kernel.SegmentIds(q=q_segment_ids, kv=kv_segment_ids)