steps

app.py

@@ -4,6 +4,9 @@ import torch
 from diffusers import StableDiffusionPipeline
 from PIL import Image
 import numpy as np
+import imageio
+import tempfile
+import os
 
 MODEL_ID = "runwayml/stable-diffusion-v1-5"
 
@@ -47,15 +50,33 @@ def generate_image(prompt, seed, num_inference_steps):
     # Set the random seed for reproducibility
     generator = torch.Generator(device=device).manual_seed(int(seed))
 
-    # Generate the image
+    # Store intermediate images
+    frames = []
+
+    def callback(step: int, timestep: int, latents):
+        # Decode latents to image
+        with torch.no_grad():
+            image = pipeline.decode_latents(latents)
+            image = pipeline.numpy_to_pil(image)[0]
+            frames.append(image)
+
+    # Generate the image with callback for each step
     with torch.no_grad():
         result = pipeline(
             prompt=prompt,
             num_inference_steps=int(num_inference_steps),
             generator=generator,
+            callback=callback,
+            callback_steps=1,
         )
 
-    return result.images[0]
+    # Save frames as video
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
+        video_path = tmpfile.name
+        imageio.mimsave(video_path, frames, fps=5)
+
+    # Return final image and video path
+    return result.images[0], video_path
 
 
 def create_interface():
@@ -86,9 +107,12 @@ def create_interface():
                 info="Number of denoising steps (more steps = higher quality but slower)",
             ),
         ],
-        outputs=gr.Image(label="Generated Image", type="pil"),
+        outputs=[
+            gr.Image(label="Generated Image", type="pil"),
+            gr.Video(label="Diffusion Steps Video"),
+        ],
         title="Stable Diffusion Image Generator",
-        description="Generate images from text using Stable Diffusion. Enter a prompt, set the seed for reproducibility, and adjust the number of diffusion steps.",
+        description="Generate images from text using Stable Diffusion. Enter a prompt, set the seed for reproducibility, and adjust the number of diffusion steps. Watch the diffusion process as a video.",
         examples=[
             ["A beautiful sunset over mountains", 42, 50],
             ["A cat wearing a space suit, digital art", 123, 50],