[Core] refactor transformers 2d into multiple init variants.

[sayakpaul]

What does this PR do?

For a new Transformer variant, we should do "Transformer2DModelForXXX" going forward.

I would love an initial review of the design.

Adjacent to #7489. I believe together with this PR and the current one, the class will read more modular.

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

I like the init methods approach. Left a couple of comments.

[sayakpaul]

@DN6 @yiyixuxu all the comments have been addressed. PTAL.

[sayakpaul]

@DN6 @yiyixuxu the problem with the failing tests (8) are a little shaky in nature. These failures happen ONLY when the input type is continuous.

They're failing because the parameters of the Transformer2Model are changing depending on the position of its internal components (transformer_blocks, proj_in, etc.) EVEN WHEN seeded.

Earlier, the self.transformer_blocks init was between self.proj_in and self._proj_out. This PR doesn't have that kind of positioning as the final part of the main __init__() initializes self.transformer_blocks. To confirm things, I put the init of self.transformer_blocks after

diffusers/src/diffusers/models/transformers/transformer_2d.py

Line 200 in b990afd

self.proj_in = torch.nn.Conv2d(self.in_channels, self.inner_dim, kernel_size=1, stride=1, padding=0)

and before

diffusers/src/diffusers/models/transformers/transformer_2d.py

Line 203 in b990afd

if self.use_linear_projection:

All the tests pass after this change.

You might ask why does this not happen for the other input types (patched and vectorized)? Because we only configure Transformer2DModel for continuous input types in tests/models/test_layers_utils.py (which is where the 8 tests are failing).

I see two potential solutions:

• Change the assertion values in the tests
• Add self.transformer_blocks init in _init_continuous_input() like the suggested above.

I think we shouldn't do pt. 2 and settle with pt. 1. The tests are mild in nature and don't seem to be testing for anything critical.

LMK.

[yiyixuxu]

I see two potential solutions:

Change the assertion values in the tests
Add self.transformer_blocks init in _init_continuous_input() like the suggested above.

I left a comment here, #7491 (comment)
I don't think we should move the transformer init inside these _init* methods just to get tests pass but I would actually in favor of doing that because it is easier to read IMO, depends on the input type, user can just go to one of the init* method instead of going back and forth;
but I don't feel strongly about this, I think it is a nice clean up regardless. So I will leave it to you guys to decide :)

[sayakpaul]

but I would actually in favor of doing that because it is easier to read IMO, depends on the input type, user can just go to one of the init* method instead of going back and forth;

This hits the point home. So, I will follow that. Thanks!

[yiyixuxu]

@sayakpaul @DN6
should we just spin out Transformer2DModelPatched and Transformer2DModelVectorized? we are pretty much already there - I think most of the work is already done.

[sayakpaul]

Move out the classes to their own modules? I think it might be better done after #7489.

[yiyixuxu]

Move out the classes to their own modules? I think it might be better done after #7489.

oh sure,
let me know what you think @DN6: we are already creating 3 separate __init__ and forward, they are practically already 3 different classes. I think the only models that use patch are pixart and dit; vectorized is only for vq diffusion but we will have to double-check

[sayakpaul]

I have an idea of the kind of refactoring @yiyixuxu suggesting. Would prefer doing it in a separate PR clubbing it with #7489. Would that be okay?

[DN6]

@sayakpaul @DN6 should we just spin out Transformer2DModelPatched and Transformer2DModelVectorized? we are pretty much already there - I think most of the work is already done.

Yeah I'm cool with spinning it out into dedicated classes. If possible, could they have model specific names e.g. PixartTransformer2D?

We can handle it in a separate PR if you prefer @sayakpaul

Would we need to merge this PR if it's just going to be changed soonish? Can we just take the work done here and in #7489 and open a new PR?

@yiyixuxu Is right that we are essentially creating three different init and forwards within the same object, so the signal here is that we should break it up into three different models.

[sayakpaul]

@DN6

I would still prefer to handle the spinning out portion in #7489 and merge this PR.

If possible, could they have model specific names e.g. PixartTransformer2D?

I think we can gradually tackle these phasing out cases. More like a coarse-to-fine kind of approach.

Regarding #7491 (comment), I think it's a matter of preference. So, I would like to elaborate a bit on why I am tending to go via the duplicating route for this particular case.

The current PR allows the Transformer2DModel class to have the following structure:

class Transformer2DModel:
    def __init__(self, ...):
        # init common attributes that are shared across the board
        self.in_channels = ...
        self.out_channels = ...

        # handle initialization of specific attributes, modules, etc.
        if input_type == "continuous":
            self._init_continuous_inputs()
        elif input_type == "patched":
            self._init_patched_inputs()
        elif input_type == "vectorized":
            self._init_vectorized_inputs()

    def _init_continuous_inputs(self):
        # initialize attributes specific to continuous inputs.
        ...

        # initialize blocks specific to continuous inputs.
        # input block.
        ...

        # transformer blocks.
        self.transformer_blocks = ... # <- duplicate code for other init methods too.

        # output blocks.
        ...

    # rest of the init methods follow the above structure too. 

Now, if we move out the transformer block initialization, it would look like:

class Transformer2DModel:
    def __init__(self, ...):
        # init common attributes that are shared across the board
        self.in_channels = ...
        self.out_channels = ...

        # handle initialization of specific attributes, modules, etc.
        if input_type == "continuous":
            self._init_continuous_inputs()
        elif input_type == "patched":
            self._init_patched_inputs()
        elif input_type == "vectorized":
            self._init_vectorized_inputs()

        # transformer blocks.
        self.transformer_blocks = ...

    def _init_continuous_inputs(self):
        # initialize attributes specific to continuous inputs.
        ...

        # initialize blocks specific to continuous inputs.
        # input block.
        ...

        # output blocks.
        ...

    # rest of the init methods follow the above structure too. 

I think it breaks the linearity of the reading flow plus the tests (which are minor, I agree).

Hope this provides more reasoning.

[DN6]

LGTM 👍🏽