From afb10a2b1cb4c2dd1a06ce20ae02a522cc59b4ed Mon Sep 17 00:00:00 2001
From: Thomas Ward <teward@thomas-ward.net>
Date: Fri, 24 May 2024 01:26:29 -0400
Subject: [PATCH] Fix missing item in returns, update names.

---
 generation.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/generation.py b/generation.py
index 3e7c030..33e4136 100644
--- a/generation.py
+++ b/generation.py
@@ -119,9 +119,10 @@ class TWCUI_Util_GenerationParameters(BaseNode):
             }
         }
 
-    RETURN_TYPES = ("MODEL", "CLIP", "VAE", "INT", "INT", "FLOAT",
+    RETURN_TYPES = ("MODEL", "CLIP", "VAE", "INT", "INT", "INT", "FLOAT",
                     comfy.samplers.KSampler.SAMPLERS, comfy.samplers.KSampler.SCHEDULERS, "INT")
-    RETURN_NAMES = ("MODEL", "CLIP", "VAE", "width", "height", "cfg_scale", "sampler_name", "scheduler_name", "seed")
+    RETURN_NAMES = ("MODEL", "CLIP", "VAE", "width", "height", "steps", "cfg_scale", "sampler_name", "scheduler",
+                    "seed")
 
     def process(self, ckpt_name: str, vae_name: str, image_width: int, image_height: int, sampling_steps: int,
                 cfg_scale: float, sampler_name: str, scheduler_name: str, seed: int) -> tuple: