karimbenharrak commited on
Commit
bd9343e
1 Parent(s): 077c9e0

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +12 -12
handler.py CHANGED
@@ -12,16 +12,6 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
12
  if device.type != 'cuda':
13
  raise ValueError("need to run on GPU")
14
 
15
- def make_inpaint_condition(image, image_mask):
16
- image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
17
- image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
18
-
19
- assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
20
- image[image_mask > 0.5] = -1.0 # set as masked pixel
21
- image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
22
- image = torch.from_numpy(image)
23
- return image
24
-
25
  class EndpointHandler():
26
  def __init__(self, path=""):
27
 
@@ -42,7 +32,7 @@ class EndpointHandler():
42
  "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, torch_dtype=torch.float16
43
  )
44
 
45
- self.pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
46
  self.pipe.enable_model_cpu_offload()
47
  self.pipe.enable_xformers_memory_efficient_attention()
48
 
@@ -161,7 +151,7 @@ class EndpointHandler():
161
  return image2
162
  """
163
 
164
- control_image = make_inpaint_condition(image, mask_image)
165
 
166
  # generate image
167
  image = pipe(
@@ -184,3 +174,13 @@ class EndpointHandler():
184
  buffer = BytesIO(base64_image)
185
  image = Image.open(buffer)
186
  return image
 
 
 
 
 
 
 
 
 
 
 
12
  if device.type != 'cuda':
13
  raise ValueError("need to run on GPU")
14
 
 
 
 
 
 
 
 
 
 
 
15
  class EndpointHandler():
16
  def __init__(self, path=""):
17
 
 
32
  "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, torch_dtype=torch.float16
33
  )
34
 
35
+ self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
36
  self.pipe.enable_model_cpu_offload()
37
  self.pipe.enable_xformers_memory_efficient_attention()
38
 
 
151
  return image2
152
  """
153
 
154
+ control_image = self.make_inpaint_condition(image, mask_image)
155
 
156
  # generate image
157
  image = pipe(
 
174
  buffer = BytesIO(base64_image)
175
  image = Image.open(buffer)
176
  return image
177
+
178
+ def make_inpaint_condition(self, image, image_mask):
179
+ image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
180
+ image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
181
+
182
+ assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
183
+ image[image_mask > 0.5] = -1.0 # set as masked pixel
184
+ image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
185
+ image = torch.from_numpy(image)
186
+ return image