Save a test image during training?

I’m training stable diffusion using a modified train_text_to_image.py script, and I’m wondering if there’s a relatively painless way to save a test image to my checkpoint folder(s)? I’d love to be able visually to track progress from wandb, in particular…

It seems like I should be able to take the UNet output and decode that using the VAE, as in the inference pipeline, but I’m concerned because I’ve noticed that the inference pipelines have a scheduler step during inference and I don’t want to mess with the training process… (obviously I mostly need a better understanding of the details of training vs inference… gulp…)

Any help/thoughts appreciated.

PS - I think this would be kind of a great feature, in general, so maybe it could be added as an option to the pipeline… ? :slight_smile: