[diffusion_labs] Move attention and res blocks under modules

[ebsmothers]

Stack from ghstack (oldest at bottom):

• -> [diffusion_labs] Move attention and res blocks under modules #490
• ADMStack and ADMResBlock changes #489
• Refactor ADM classes #488

Differential Revision: D50288362

[ebsmothers]

@ebsmothers has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[codecov-commenter]

Codecov Report

All modified lines are covered by tests ✅

❗ No coverage uploaded for pull request base (gh/ebsmothers/21/base@b948406). Click here to learn what that means.

Additional details and impacted files

@@                   Coverage Diff                    @@
##             gh/ebsmothers/21/base     #490   +/-   ##
========================================================
  Coverage                         ?   74.12%           
========================================================
  Files                            ?      207           
  Lines                            ?    14259           
  Branches                         ?        0           
========================================================
  Hits                             ?    10569           
  Misses                           ?     3690           
  Partials                         ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ebsmothers]

@ebsmothers has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[pbontrager]

I don't personally agree with this diff. If these modules can live as generic layers, then they should be under torchmultimodal/modules/layers while if they're just specific to ADM then they should be under adm_unet. I think we should keep the diffusion folders as focused on diffusion as possible and try to keep out other things.

[ebsmothers]

@pbontrager I guess there are three possible places for these components: diffusion_labs/models/adm_unet/..., diffusion_labs/modules/..., and torchmultimodal/modules/.... I agree that the first location should be for components specific to ADM UNet. The second should be for more generic diffusion components (i.e. not tied to ADM specifically), while components in the third location should have use for other non-diffusion cases.

Based on those conditions I think it should be either option (1) or (2), since I'm not aware of any models outside of diffusion_labs that would use these components. @abhinavarora mentioned that there are other diffusion models that can benefit from these components which is why I opened this PR. But if that's not the case, we can just close this and leave them where they are. Let me know what both of you think here.

[pbontrager]

@pbontrager I guess there are three possible places for these components: diffusion_labs/models/adm_unet/..., diffusion_labs/modules/..., and torchmultimodal/modules/.... I agree that the first location should be for components specific to ADM UNet. The second should be for more generic diffusion components (i.e. not tied to ADM specifically), while components in the third location should have use for other non-diffusion cases.

Based on those conditions I think it should be either option (1) or (2), since I'm not aware of any models outside of diffusion_labs that would use these components. @abhinavarora mentioned that there are other diffusion models that can benefit from these components which is why I opened this PR. But if that's not the case, we can just close this and leave them where they are. Let me know what both of you think here.

I believe they should stay with adm_unet until they start being reused as they may require modifications at that time. At that time I think they should be moved to option (3) because, while the architecture is being used for diffusion training, neural networks are still pretty generic for all types of training. I'm open to discussing this further though.

[abhinavarora]

My understanding is that these components are specific implementations of general concepts that have been popularized by prominent diffusion researchers (OpenAI, Stability, etc). I don't see them being used in any other models outside of those that build their code on top of OpenAI and StabilityAI repos. For example, this implementation of resnet blocks is quite different from the generic ones found in TorchVision.

Hence in my opinion, only option 1) and 2) make sense. To me option 2) makes more sense because components like ADMResBlock are used by all other models (VAE and LDM UNet). I am okay with 1) as well if you want to unblock the PT Conf works and punt this discussion to later.

[ebsmothers]

@ebsmothers has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[ebsmothers]

@ebsmothers has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[ebsmothers]

Closing this since we are not planning to land these changes right now