Try output_count

handler.py

@@ -76,8 +76,8 @@ class EndpointHandler():
         """
         prompt = data.pop("inputs", None)
         image = data.pop("image", None)
+        output_count = data.pop("output_count", 1)
         controlnet_type = data.pop("controlnet_type", None)
-        count = data.pop("count", None)
         
         # Check if neither prompt nor image is provided
         if prompt is None and image is None:
@@ -120,10 +120,7 @@ class EndpointHandler():
 
         
         # return first generate PIL image
-        if count is None:
-            return out.images[0]
-        if count is not None:
-            return out.images[0:count]
+        return out.images[0, output_count]
     
     # helper to decode input image
     def decode_base64_image(self, image_string):