fix(models): pad attention mask by remaining length instead of target length (#14699)

[dlowzzxx]

Fixes #14699

Problem Description

Previously in prepare_attention_mask (in both src/diffusers/models/attention_processor.py and src/diffusers/models/attention.py), when padding an attention mask whose length was smaller than target_length, F.pad(attention_mask, (0, target_length), value=0.0) was used.

Because torch.nn.functional.pad(tensor, (0, n)) appends n padding elements to the trailing dimension, passing target_length resulted in a post-padding length of current_length + target_length rather than target_length. The same issue occurred on MPS where padding_shape used target_length instead of the difference between target_length and current_length.

Due to this bug, test_model_xattn_padding in tests/models/unets/test_models_unet_2d_condition.py was skipped with a TODO noting:

"we currently pad mask by target_length tokens (what unclip needs), whereas stable-diffusion's cross-attn needs to instead pad by remaining_length."

Proposed Changes

1. Updated prepare_attention_mask in both src/diffusers/models/attention_processor.py and src/diffusers/models/attention.py:
   • Compute remaining_length = target_length - current_length.
   • When remaining_length > 0, pad by remaining_length:
     • On MPS, construct padding_shape = (attention_mask.shape[0], attention_mask.shape[1], remaining_length).
     • Otherwise, call F.pad(attention_mask, (0, remaining_length), value=0.0).
2. Unskipped test_model_xattn_padding in tests/models/unets/test_models_unet_2d_condition.py and aligned block_out_channels = (16, 32) to ensure divisibility with attention_head_dim = (8, 16).
3. Cleaned up unused imports and formatted code with ruff.

Verification & Reproduction

Ran the unskipped test directly:

pytest tests/models/unets/test_models_unet_2d_condition.py -k test_model_xattn_padding -v

Output:

============================= test session starts =============================
platform win32 -- Python 3.12.10, pytest-9.1.1, pluggy-1.6.0
collected 164 items / 163 deselected / 1 selected

tests/models/unets/test_models_unet_2d_condition.py::TestUNet2DCondition::test_model_xattn_padding PASSED [100%]
================ 1 passed, 163 deselected, 5 warnings in 0.97s ================

Also verified MPS mock and boundary cases (mask is None, current_length == target_length, current_length > target_length). All checks and ruff check/ruff format passed cleanly.