Update handler.py
Browse files- handler.py +7 -6
handler.py
CHANGED
|
@@ -25,7 +25,7 @@ class EndpointHandler():
|
|
| 25 |
# )
|
| 26 |
# self.smooth_pipe.to("cuda")
|
| 27 |
|
| 28 |
-
|
| 29 |
self.controlnet = ControlNetModel.from_pretrained(
|
| 30 |
"lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16
|
| 31 |
)
|
|
@@ -38,9 +38,8 @@ class EndpointHandler():
|
|
| 38 |
self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
|
| 39 |
self.pipe.enable_model_cpu_offload()
|
| 40 |
self.pipe.enable_xformers_memory_efficient_attention()
|
| 41 |
-
"""
|
| 42 |
-
|
| 43 |
|
|
|
|
| 44 |
# load StableDiffusionInpaintPipeline pipeline
|
| 45 |
self.pipe = AutoPipelineForInpainting.from_pretrained(
|
| 46 |
"runwayml/stable-diffusion-inpainting",
|
|
@@ -66,6 +65,7 @@ class EndpointHandler():
|
|
| 66 |
self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
|
| 67 |
#self.pipe3.enable_model_cpu_offload()
|
| 68 |
self.pipe3.enable_xformers_memory_efficient_attention()
|
|
|
|
| 69 |
|
| 70 |
|
| 71 |
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
|
|
@@ -119,7 +119,7 @@ class EndpointHandler():
|
|
| 119 |
"""
|
| 120 |
|
| 121 |
#pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
|
| 122 |
-
|
| 123 |
# run inference pipeline
|
| 124 |
out = self.pipe(prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask_image, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)
|
| 125 |
|
|
@@ -155,8 +155,9 @@ class EndpointHandler():
|
|
| 155 |
|
| 156 |
# return first generate PIL image
|
| 157 |
return image2
|
| 158 |
-
|
| 159 |
"""
|
|
|
|
|
|
|
| 160 |
control_image = self.make_inpaint_condition(image, mask_image)
|
| 161 |
|
| 162 |
# generate image
|
|
@@ -173,7 +174,7 @@ class EndpointHandler():
|
|
| 173 |
).images[0]
|
| 174 |
|
| 175 |
return image
|
| 176 |
-
|
| 177 |
|
| 178 |
|
| 179 |
# helper to decode input image
|
|
|
|
| 25 |
# )
|
| 26 |
# self.smooth_pipe.to("cuda")
|
| 27 |
|
| 28 |
+
|
| 29 |
self.controlnet = ControlNetModel.from_pretrained(
|
| 30 |
"lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16
|
| 31 |
)
|
|
|
|
| 38 |
self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
|
| 39 |
self.pipe.enable_model_cpu_offload()
|
| 40 |
self.pipe.enable_xformers_memory_efficient_attention()
|
|
|
|
|
|
|
| 41 |
|
| 42 |
+
"""
|
| 43 |
# load StableDiffusionInpaintPipeline pipeline
|
| 44 |
self.pipe = AutoPipelineForInpainting.from_pretrained(
|
| 45 |
"runwayml/stable-diffusion-inpainting",
|
|
|
|
| 65 |
self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
|
| 66 |
#self.pipe3.enable_model_cpu_offload()
|
| 67 |
self.pipe3.enable_xformers_memory_efficient_attention()
|
| 68 |
+
"""
|
| 69 |
|
| 70 |
|
| 71 |
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
|
|
|
|
| 119 |
"""
|
| 120 |
|
| 121 |
#pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
|
| 122 |
+
"""
|
| 123 |
# run inference pipeline
|
| 124 |
out = self.pipe(prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask_image, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)
|
| 125 |
|
|
|
|
| 155 |
|
| 156 |
# return first generate PIL image
|
| 157 |
return image2
|
|
|
|
| 158 |
"""
|
| 159 |
+
|
| 160 |
+
|
| 161 |
control_image = self.make_inpaint_condition(image, mask_image)
|
| 162 |
|
| 163 |
# generate image
|
|
|
|
| 174 |
).images[0]
|
| 175 |
|
| 176 |
return image
|
| 177 |
+
|
| 178 |
|
| 179 |
|
| 180 |
# helper to decode input image
|