Question about why used diffusions.Transformer2DModel for parent instead of nn.Modules?

[97KangSH]

Dear Author,

I’m curious about the design choice of using Transformer2DModel as the parent class for Mamba2DModel.

Are there specific benefits to this approach that I might not be aware of? I find it a bit confusing since the model is called Mamba but heavily relies on Transformer components.

Looking forward to your clarification.

Thanks!

[tyshiwo1]

Thank you for your appreciation.

This is because Transformer2DModel in diffusers corresponds to the diffusion transformer (DiTs) models.
When I developed DiM, I wanted to leverage or ablate the modules in DiTs for convenience (such as various Normalization types), so I chose Transformer2DModel in diffusers as the parent class for Mamba2DModel.
In the final version of DiM, a lot of modules were modified but the parent class was kept.

[97KangSH]

Thank you for your response.
It resolved my question.