replace enable_gpa with repeat_interleave for krea2

[yiyixuxu]

fix #14518

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Thanks! Just flagging the following table:

Claude derived it by using the following code and running it against this branch and main:

from diffusers import Krea2Transformer2DModel

device, dtype = "cuda", torch.bfloat16

model = Krea2Transformer2DModel.from_pretrained(
    "krea/Krea-2-Raw", subfolder="transformer", torch_dtype=dtype
).to(device)
model.eval()

# 1024x1024 image -> 128x128 latents -> 64x64 grid of patch-2 tokens;
# 512 text positions with only the first 77 valid (padding mask exercised).
torch.manual_seed(0)
batch, text_seq_len, grid = 1, 512, 64
image_seq_len = grid * grid

hidden_states = torch.randn(batch, image_seq_len, 64, device=device, dtype=dtype)
encoder_hidden_states = torch.randn(batch, text_seq_len, 12, 2560, device=device, dtype=dtype)
timestep = torch.tensor([0.7], device=device)
encoder_attention_mask = torch.zeros(batch, text_seq_len, dtype=torch.bool, device=device)
encoder_attention_mask[:, :77] = True

text_ids = torch.zeros(text_seq_len, 3, device=device)
image_ids = torch.zeros(grid, grid, 3, device=device)
image_ids[..., 1] = torch.arange(grid, device=device)[:, None]
image_ids[..., 2] = torch.arange(grid, device=device)[None, :]
position_ids = torch.cat([text_ids, image_ids.reshape(-1, 3)], dim=0)

with torch.no_grad():
    out = model(
        hidden_states=hidden_states,
        encoder_hidden_states=encoder_hidden_states,
        timestep=timestep,
        position_ids=position_ids,
        encoder_attention_mask=encoder_attention_mask,
    ).sample

I think this is acceptable because now the kernel invocations internally have changed.

[sayakpaul]

Either pass enable_gqa=True to dispatch_attention_fn

Can we also supplement a model file here for the reference?

[sayakpaul]

Turns on whether the model passes a mask.

Do we mean performance is turned on?

[asomoza]

thanks! this fixes the issue.

Not sure if it's worth to mention, but this changes the resulting image a little:

main	PR

But still, it doesn't make it worse or better, and the change is minimal and mostly imperceptible if you're not comparing both.