Training your own model

[xavierleung]

I love jax guided diffusion and would like to try to train a model myself.
As the last commit on the training script says "add highly oversimplified training code",
I wonder if that is the actual code that was used for training of the openimages 512x512 model?