fix dispatch_attention_fn check

[yiyixuxu]

related to #12605

there are two bugs canceling out each other:

• bug1: os.environ["DIFFUSERS_ATTN_CHECKS"] = "yes" (the instruction we provided in our documentation) does not actually enable the validation
• bug2: check_shape is implemented wrong

example, hunyuanimage21 docstring example (using attention mask and works compoletely fine) will throw an error if we enable checking correctly

import os
os.environ["DIFFUSERS_ATTN_CHECKS"] = "YES"

from diffusers import HunyuanImagePipeline
import torch

device = "cuda:1"
dtype = torch.bfloat16
repo = "hunyuanvideo-community/HunyuanImage-2.1-Diffusers"

pipe = HunyuanImagePipeline.from_pretrained(repo, torch_dtype=dtype)
pipe = pipe.to(device)

prompt = "A cute, cartoon-style anthropomorphic penguin plush toy with fluffy fur, standing in a painting studio, wearing a red knitted scarf and a red beret with the word “Tencent” on it, holding a paintbrush with a focused expression as it paints an oil painting of the Mona Lisa, rendered in a photorealistic photographic style."

generator = torch.Generator(device=device).manual_seed(649151)
out = pipe(
    prompt, 
    num_inference_steps=50, 
    height=2048, 
    width=2048, 
    generator=generator,
).images[0]


out.save("output.png")

the example test provided here #12605 (comment) would also throw an correct warning in this PR

[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.

[DN6]

Thanks! 👍🏽