Refactor OmniGen

[a-r-r-o-w]

Refactors OmniGen added in #10148.

Image Editing	Text-to-Image

I haven't tested all supported usecases extensively, but I believe the changes should be okay since they did not alter any existing logic (except for RoPE).

The summary of changes are:

• Instead of the custom apply_rotary_emb function, we use the Diffusers one
• Removed set_attn_processors and related functions. We're thinking of a better way of supporting this since it makes the modeling implementation bloated. For any new models, let's not add these methods. It will be handled better with an upcoming refactor before this release hopefully.
• Makes sure guidance and scheduler step is performed in float32, so that we don't perform unnecessary bfloat16 -> float32 -> bfloat16 dtype conversion at each step. Instead, latents are always in float32 and casted just before feeding into transformer to required dtype.
• Removes everything related to attention kwargs. This will be added again by @sayakpaul, or someone else, who will take up adding peft support.
• Marks non-forward methods in the transformer private, since it could be subject to change in more future refactors
• Cleans up docs a bit and removes any intermediate layer docs. The parameters are self-explanatory, or already documented in the OmniGenTransformer2DModel docstring
• Enables layerwise upcasting test

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

[yiyixuxu]

very nice! thank you!

[yiyixuxu]

I see this is not used, but let's confirm with the author if it's intended to be used or not

[staoxiao]

This argument is simply copied from other pipelines, but is not used in OmniGen.

[yiyixuxu]

@staoxiao can you take a look to see if the changes are ok?

[sayakpaul]

Thanks for the heads-up, Aryan! Appreciate it.

I will wait for this before to get merged before starting the LoRA support.

[staoxiao]

@a-r-r-o-w@yiyixuxu , thanks for refactoring. I think these changes are ok.

[a-r-r-o-w]

Thanks for the review @staoxiao! Will merge once CI is green