StableDiffusion with intermediate save

stablediffusion.py

from diffusers import StableDiffusionPipeline
import torch

def progress(step, timestep, latents):
  print(step, timestep, latents[0][0][0][0])
  with torch.no_grad():
    latents = 1 / 0.18215 * latents
    image = pipe.vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.cpu().permute(0, 2, 3, 1).float().numpy()
    image = pipe.numpy_to_pil(image)
    image[0].save(f"step_{step}.png")

# main image
prompt = "A fantasy landscape with mountains and a river"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
images = pipe(prompt, callback=progress, callback_steps=10).images
images[0].save(f"image.png")