Update handler.py

handler.py

@@ -76,6 +76,7 @@ class EndpointHandler():
         """
         prompt = data.pop("inputs", None)
         image = data.pop("image", None)
+        count = data.pop("count", 1)
         controlnet_type = data.pop("controlnet_type", None)
         
         # Check if neither prompt nor image is provided
@@ -110,7 +111,7 @@ class EndpointHandler():
             image=control_image,
             num_inference_steps=num_inference_steps, 
             guidance_scale=guidance_scale,
-            num_images_per_prompt=1,
+            num_images_per_prompt=count,
             height=height,
             width=width,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
@@ -119,7 +120,7 @@ class EndpointHandler():
 
         
         # return first generate PIL image
-        return out.images[0]
+        return out.images[0, count]
     
     # helper to decode input image
     def decode_base64_image(self, image_string):