diff --git a/stable_diffusion_videos/image_generation.py b/stable_diffusion_videos/image_generation.py
index 420b625..63f15ee 100644
--- a/stable_diffusion_videos/image_generation.py
+++ b/stable_diffusion_videos/image_generation.py
@@ -182,7 +182,7 @@ def generate_images(
                 height=height,
                 width=width,
                 output_type="pil" if not upsample else "numpy",
-            )['sample']
+            )['images']
             if upsample:
                 images = []
                 for output in outputs: