Sebastian Semeniuc commited on
Commit
4cc0dca
·
1 Parent(s): 78b0c60

refactor: add venv locally and num_of_images parameter and return list of images instead of string

Browse files
Files changed (2) hide show
  1. .gitignore +1 -0
  2. handler.py +7 -3
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ /.venv
handler.py CHANGED
@@ -78,12 +78,16 @@ class EndpointHandler():
78
  """
79
  prompt = data.pop("inputs", None)
80
  image = data.pop("image", None)
 
81
  controlnet_type = data.pop("controlnet_type", None)
82
 
83
  # Check if neither prompt nor image is provided
84
  if prompt is None and image is None:
85
  return {"error": "Please provide a prompt and base64 encoded image."}
86
 
 
 
 
87
  # Check if a new controlnet is provided
88
  if controlnet_type is not None and controlnet_type != self.control_type:
89
  print(f"changing controlnet from {self.control_type} to {controlnet_type} using {CONTROLNET_MAPPING[controlnet_type]['model_id']} model")
@@ -112,7 +116,7 @@ class EndpointHandler():
112
  image=control_image,
113
  num_inference_steps=num_inference_steps,
114
  guidance_scale=guidance_scale,
115
- num_images_per_prompt=1,
116
  height=height,
117
  width=width,
118
  controlnet_conditioning_scale=controlnet_conditioning_scale,
@@ -120,8 +124,8 @@ class EndpointHandler():
120
  )
121
 
122
 
123
- # return first generate PIL image
124
- return out.images[0]
125
 
126
  # helper to decode input image
127
  def decode_base64_image(self, image_string):
 
78
  """
79
  prompt = data.pop("inputs", None)
80
  image = data.pop("image", None)
81
+ num_of_images = data.pop("numOfImages", None)
82
  controlnet_type = data.pop("controlnet_type", None)
83
 
84
  # Check if neither prompt nor image is provided
85
  if prompt is None and image is None:
86
  return {"error": "Please provide a prompt and base64 encoded image."}
87
 
88
+ if num_of_images is None:
89
+ num_of_images = 1
90
+
91
  # Check if a new controlnet is provided
92
  if controlnet_type is not None and controlnet_type != self.control_type:
93
  print(f"changing controlnet from {self.control_type} to {controlnet_type} using {CONTROLNET_MAPPING[controlnet_type]['model_id']} model")
 
116
  image=control_image,
117
  num_inference_steps=num_inference_steps,
118
  guidance_scale=guidance_scale,
119
+ num_images_per_prompt=num_of_images,
120
  height=height,
121
  width=width,
122
  controlnet_conditioning_scale=controlnet_conditioning_scale,
 
124
  )
125
 
126
 
127
+ # return the list of generated images
128
+ return out.images
129
 
130
  # helper to decode input image
131
  def decode_base64_image(self, image_string):