 model that takes a noisy sample, conditional state, and a timestep and returns a sample
shaped output. This model inherits from FlaxModelMixin. Check the superclass documentation for it’s generic methods
implemented for all models (such as downloading or saving). This model is also a Flax Linen flax.linen.Module
subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matters related to its
general usage and 