Refactor additions from Generative module

[marksgraham]

After merging MONAI Generative into core issue some refactoring is needed to reduce repeat code, see discussion here.

Until this is done any blocks from Generative that are likely to be removed will be marked as private.

Will update the list of items needing refactoring here as I go along:

• [x] Upsample/Downsample/Resblock/AttentionBlock in the autoencoderkl network

  • Upsample block can be replaced by Monai's upsample block, might need to implement conv_only=True option in Monai's block
  • Downsample block can be replaced by Monai's conv, see comment here
  • The Attention block should be able to make use of Monai's self attention block
  • I don't see a simple way to replace the Resblock, many other network's implement their own versions, but maybe we should rename it to AutoencoderKLResblock to differentiate it from other Resblocks in the codebase as done in e.g. UnerR

• [x] SABlock and TransformerBlock used by DecoderOnlyTransformer

  • Use the internal monai versions. could make use of the new cross-attention transformer block we plan to make as described above

• [x] Upsample/Downsample/BasicTransformerBlock/ResnetBlock/AttentionBlock in the diffusion_model_unet

  • The CrossAttention block could be added as a block under monai.network.blocks, similarly to the SelfAttention block that already exists there
  • The SelfAttention block can be replaced with the block we already have
  • The TransformerBlock is actually a cross-attention transformer block, we could make a new monai block for it
  • Downsample block - maybe we need to add an option to the existing monai downsample block to either conv or pool, then we could use it here
  • Upsample block - we can use monai's version if we add a post_conv option to perform a convolution after the interpolation
  • Resnet - once again I think it is OK to keep this and rename to DIffusionUnetResnetBlock

• [x] SpadeNorm block - use get_norm_layer here and here

• [x] SPADEAutoencoder - merge with the autoencoder KL as the only difference is in the decoder. might make more sense to just inherit from autoencoder KL (also will get the benefit of the new load_old_state_dict metho)

• [x] Had to add some calls to .contiguous() in the diffusions model unet to stop issues with inferer tests pr here - need to dig deeper and find if these are really necessary, as these calls do copies and consume memory

• [x] ControlNet refactor suggested by @ericspod to tidy up some of the code here

• [x] Neater init on the patchgan discriminator suggested by @ericspod here

• [ ] ~~Schedulers - refactor some calculations into the base class as suggested by @KumoLiu [here]~~(https://github.com/Project-MONAI/MONAI/pull/7332#discussion_r1437355600) these calculations aren't common to all the classes that inherit from scheduler (e.g. PNDM) so I'm not sure they should be moved to the base class

• [ ] Deferred to future after discussion with @ericspod Inferers - create new base class for the Generative infererers, see discussion here

• [ ] Deferred to future after discussion with @ericspod Engines - refactor GANTrainer into a more general base class that can replace AdversarialTrainer here

[vikashg]

Closing @marksgraham as it seems to be completed and merged.

[marksgraham]

Hi @vikashg afraid this is still pending so re-opening!

[KumoLiu]

Add here as a reminder. https://github.com/Project-MONAI/MONAI/pull/7346#discussion_r1444749763 https://github.com/Project-MONAI/MONAI/pull/7346#discussion_r1444207936

[vgrau98]

I can start working on SABlock, integrating the private block and making refacto if needed if this is ok for you @marksgraham

[marksgraham]

Hi @vgrau98 yes that would be great! One thing to bear in mind is that ideally, a DiffusionModelUnet trained before the refactor should be able to load the same weights post-refactor. Happy to chat about this more if you want.

[marksgraham]

I can start working on SABlock, integrating the private block and making refacto if needed if this is ok for you @marksgraham

Hi @vgrau98 I'm getting back to this refactor now. Wondering if you managed to make any progress on the SABlock refactor?

[ericspod]

@virginiafdez we should come back to the last two items here to see if we want to address anything next.

[KumoLiu]

Hi @ericspod @virginiafdez, during refactoring MAISI with the latest components in the core we also find that some refactor here may not achieve similar performance as before. Such as we use one linear to replace three linear used in generative repo. Although the weights can be loaded correctly, but the result show much difference. https://github.com/Project-MONAI/MONAI/blob/139b62c4aa161969ad1126c4feeec88c9833d4ef/monai/networks/blocks/selfattention.py#L89

https://github.com/Project-MONAI/GenerativeModels/blob/7428fce193771e9564f29b91d29e523dd1b6b4cd/generative/networks/blocks/selfattention.py#L80-L83

Before doing further refactor, I would suggest we may need to verify several tutorials to see whether we can achieve the similar the performance as before. What do you think?