--max_train_steps=15000 \ 
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-pokemon-model" 

JAX

Hide JAX content

With Flax, it’s possible to train a Stable Diffusion model faster on TPUs and GPUs thanks to @duongna211. This is very efficient on TPU hardware but works great on GPUs too. The Flax training script doesn’t support features like gradient checkpointing or gradient 