Problem / Goal

The current common block can only be used for the current U-Net version and is not flexible enough to be used for other models (e.g., DDPM). The goal is to modify the common blocks (ConvBlock, UpConvBlock, and BottleneckBlock) to accommodate other models and make them more flexible.

What does good look like?

• The common blocks can be used for both U-Net and DDPM architectures.
• Flexible options for normalization, dropout, etc.
• The changes are tested and verified to ensure consistency.

Related context

https://github.com/alan-turing-institute/ice-station-zebra/tree/68-flexible-common-block