SD2 inference pass mesh for flahattention with spmd

[ManfeiBai]

What does this PR do?

SD2 inference pass mesh for flahattention with spmd

tested log: https://gist.github.com/ManfeiBai/bc97938b94b7b8b2264e6d9dcbae0f26

Before submitting

• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

[zpcore]

Can you run ruff check --fix src to fix the format? Thanks

[ManfeiBai]

Can you run ruff check --fix src to fix the format? Thanks

thanks, updated

[ManfeiBai]

hi, @zpcore, current CI failure in https://github.com/pytorch-tpu/diffusers/actions/runs/13020451508/job/36320882664?pr=12 and https://github.com/pytorch-tpu/diffusers/actions/runs/13020451508/job/36320882664?pr=12 are due to code-not-used or format issue from other not-related code file;

to avoid touch unrelated project in this simple PR, will propose to merge PR with these CI failure

[zpcore]

LGTM! Thanks for fixing