Generative_models-flax

Implementations of generative models using jax and flax.

TODO

  • VAEs
  • GANs
  • Diffusion models
  • Train and generation on MNIST data

Diffusion Process Visualization

Diffusion_Process