Allow different architectures with `diffusion.py` and `flow.py`

[gurjeetjagwani]

Currently we are dependent on the MLP() class in both these modules. We would like to have some flexibility for different use cases.