Project-MONAI / GenerativeModels

MONAI Generative Models makes it easy to train, evaluate, and deploy generative models and related applications
Apache License 2.0
604 stars 87 forks source link

Make use of MONAI blocks in diffusion_model_unet.py #141

Closed Warvito closed 1 year ago

Warvito commented 1 year ago

I think GEGLU and FeedForward components from diffusion_model_unet.py https://github.com/Project-MONAI/GenerativeModels/blob/7dcdc9a84e96e25e57325ce9b503b8b39c9801af/generative/networks/nets/diffusion_model_unet.py#L53

https://github.com/Project-MONAI/GenerativeModels/blob/7dcdc9a84e96e25e57325ce9b503b8b39c9801af/generative/networks/nets/diffusion_model_unet.py#L71

are already implemented on MONAI core as MLPBlock

https://github.com/Project-MONAI/MONAI/blob/17529e7ae0a53db68f9ace7a130da2b611650cc3/monai/networks/blocks/mlp.py#L22

It might be necessary further investigation.

Warvito commented 1 year ago

The current version of MLPBlock is not compatible with the Feed Forward block used inside the Diffusion model. Waiting for https://github.com/Project-MONAI/MONAI/issues/5853 to see if it will be possible to have an unified solution.