[Fix] enable_xformers_memory_efficient_attention() in Flux Pipeline

[SahilCarterr]

What does this PR do?

Fixes #12334 enable_xformers_memory_efficient_attention() in Flux Pipeline

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• 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?

Who can review?

Anyone in the community is free to review the PR once the tests have passed.
@DN6

[JoeGaffney]

Possibly you could this at the top when its first imported as looks like only the one function is used.

if is_xformers_available():
    import xformers as xops
    xformers_attn_fn = getattr(xops, "memory_efficient_attention", None) \
                       or getattr(xops.ops, "memory_efficient_attention")
else:
    xformers_attn_fn = None

As there is already many trys here and conditions in this block of code. Also seems like there is additional checks in the function, but not in the code you contributed.

I did notice this breaking change in xformers maybe there is better way also to check which is installed?

[SahilCarterr]

i fix the trys issue here now . I am thinking to remove already present try . Can you tell me which version of xformers is suitable here and also check the code below

else:
                import importlib.metadata
                version = importlib.metadata.version("xformers")
                if tuple(map(int, version.split("."))) < (0, 0, 32):
                       raise ImportError(f"xformers>=0.0.32 is required for memory efficient attention, but found {version}")

                # Make sure we can run the memory efficient attention
                dtype = None
                if attention_op is not None:
                    op_fw, op_bw = attention_op
                    dtype, *_ = op_fw.SUPPORTED_DTYPES
                q = torch.randn((1, 2, 40), device="cuda", dtype=dtype)
                _ = xops.memory_efficient_attention(q, q, q)

                self.set_attention_backend("xformers")

@JoeGaffney

[JoeGaffney]

I think both ways are supported depending on the version of xformers.

I don't want to block so possibly your original is most straight forward currently as looking more at the code it probably a need a bit of a wider refactor.

[SahilCarterr]

Then it is fixed for now in latest commit. Thank you

[JoeGaffney]

ah you mean force people to be on the update to date xformers ?