[Core] refactor transformer_2d forward logic into meaningful conditions.

[sayakpaul]

What does this PR do?

Refactors the forward() of Transformers2DModel for easier readability.

More specifically, the PR refactors the forward() method of Transformers2DModel to have the following unified structure:

def forward(self):
    # handle attention masking. 
    ...

    # 1. input-level operations.
    if self.is_input_continuous:
        ... = self._operate_on_continuous_inputs(...)
    elif self.is_input_vectorized:
    	... = self.latent_image_embedding(...) # for vectorized inputs only this is needed.
    elif self.is_input_patches:
        ... = self._operate_on_patched_inputs(...)

    # 2. Blocks (common across all the variants)
    for block in self.transformer_blocks:
    	...

    # 3. Outputs.
    if self.is_input_continuous:
        ... = self._get_output_for_continuous_inputs(...)
    elif self.is_input_vectorized:
        output = self._get_output_for_vectorized_inputs(...)
    elif self.is_input_patches:
        output = self._get_output_for_patched_inputs(...)

    return output

About the possibility of spinning out separate transformer classes for patched and vectorized inputs, what's the consensus? Currently, we have DiT and PixArt-Alpha that use patched inputs. So, for those checkpoints (DiT and PixArt-Alpha), are we thinking of just updating the configs to use PatchedTransformer2DModel? (Same could be done for the VectorizedTransformer2DModel)

If that's the case I think it could be quite breaking of a change. Many folks use the DiT and PixArt models especially in the light of many Open SoRA initiative. Introducing a change like this could be problematic. Should we consider throwing a deprecation warning, instead?

@yiyixuxu @DN6 please let me know your thoughts.

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

@DN6 @yiyixuxu a gentle ping.

[DN6]

LGTM. Left a comment, but it's mostly a nit. Is the failing test related?

[sayakpaul]

@DN6 WDYT?

About the possibility of spinning out separate transformer classes for patched and vectorized inputs, what's the consensus? Currently, we have DiT and PixArt-Alpha that use patched inputs. So, for those checkpoints (DiT and PixArt-Alpha), are we thinking of just updating the configs to use PatchedTransformer2DModel? (Same could be done for the VectorizedTransformer2DModel)

(refer to the OP)

[yiyixuxu]

I'm cool with this PR once you and @DN6 are happy with this!

If that's the case I think it could be quite breaking of a change. Many folks use the DiT and PixArt models especially in the light of many Open SoRA initiative. Introducing a change like this could be problematic. Should we consider throwing a deprecation warning, instead?

for this, yes, I agree it is a highly used block, and I think we should give more thought to avoid breaking changes.
can potentially deprecate Tranformer2DModel and use 3 different Transformers2DModel instead: ContinousTransformers2DMode, PatchedTransformers2DModel, VectorizedTransformers2DModel;
just spin out the patched and vectorized, and sending a deprecation message for them also works I think

would love to hear your thoughts @DN6

[sayakpaul]

Cool. Will work on the deprecation in a future PR then as we discuss the best design.

[sayakpaul]

@DN6 I had to revert your suggestion because height and width are calculated differently when inputs are continuous and patched. Hope that's okay.

Gonna merge the PR after the CI is green.