Revert to fork code

handler.py

@@ -76,15 +76,11 @@ class EndpointHandler():
         """
         prompt = data.pop("inputs", None)
         image = data.pop("image", None)
-        count = data.pop("count", None)
         controlnet_type = data.pop("controlnet_type", None)
         
         # Check if neither prompt nor image is provided
         if prompt is None and image is None:
             return {"error": "Please provide a prompt and base64 encoded image."}
-
-        if count is None:
-            return {"error": "Please provide a count."}
         
         # Check if a new controlnet is provided
         if controlnet_type is not None and controlnet_type != self.control_type:
@@ -114,7 +110,7 @@ class EndpointHandler():
             image=control_image,
             num_inference_steps=num_inference_steps, 
             guidance_scale=guidance_scale,
-            num_images_per_prompt=count,
+            num_images_per_prompt=1,
             height=height,
             width=width,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
@@ -123,11 +119,11 @@ class EndpointHandler():
 
         
         # return first generate PIL image
-        return out.images[0, count]
+        return out.images[0]
     
     # helper to decode input image
     def decode_base64_image(self, image_string):
         base64_image = base64.b64decode(image_string)
         buffer = BytesIO(base64_image)
         image = Image.open(buffer)
-        return image
+        return image