Update handler.py
Browse files- handler.py +27 -6
handler.py
CHANGED
|
@@ -14,21 +14,25 @@ if device.type != 'cuda':
|
|
| 14 |
|
| 15 |
class EndpointHandler():
|
| 16 |
def __init__(self, path=""):
|
|
|
|
|
|
|
|
|
|
|
|
|
| 17 |
# load StableDiffusionInpaintPipeline pipeline
|
| 18 |
-
self.pipe = AutoPipelineForInpainting.from_pretrained(
|
| 19 |
"runwayml/stable-diffusion-inpainting",
|
| 20 |
revision="fp16",
|
| 21 |
torch_dtype=torch.float16,
|
| 22 |
)
|
| 23 |
# use DPMSolverMultistepScheduler
|
| 24 |
-
self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
|
| 25 |
# move to device
|
| 26 |
-
self.pipe = self.pipe.to(device)
|
| 27 |
|
| 28 |
-
self.pipe2 = AutoPipelineForInpainting.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
|
| 29 |
-
self.pipe2.to("cuda")
|
| 30 |
|
| 31 |
-
self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
|
| 32 |
|
| 33 |
|
| 34 |
|
|
@@ -52,6 +56,22 @@ class EndpointHandler():
|
|
| 52 |
image = None
|
| 53 |
mask_image = None
|
| 54 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 55 |
self.pipe.enable_xformers_memory_efficient_attention()
|
| 56 |
|
| 57 |
# run inference pipeline
|
|
@@ -94,6 +114,7 @@ class EndpointHandler():
|
|
| 94 |
|
| 95 |
# return first generate PIL image
|
| 96 |
return result
|
|
|
|
| 97 |
|
| 98 |
# helper to decode input image
|
| 99 |
def decode_base64_image(self, image_string):
|
|
|
|
| 14 |
|
| 15 |
class EndpointHandler():
|
| 16 |
def __init__(self, path=""):
|
| 17 |
+
|
| 18 |
+
self.pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
|
| 19 |
+
self.generator = torch.Generator(device="cuda").manual_seed(0)
|
| 20 |
+
|
| 21 |
# load StableDiffusionInpaintPipeline pipeline
|
| 22 |
+
#self.pipe = AutoPipelineForInpainting.from_pretrained(
|
| 23 |
"runwayml/stable-diffusion-inpainting",
|
| 24 |
revision="fp16",
|
| 25 |
torch_dtype=torch.float16,
|
| 26 |
)
|
| 27 |
# use DPMSolverMultistepScheduler
|
| 28 |
+
#self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
|
| 29 |
# move to device
|
| 30 |
+
#self.pipe = self.pipe.to(device)
|
| 31 |
|
| 32 |
+
#self.pipe2 = AutoPipelineForInpainting.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
|
| 33 |
+
#self.pipe2.to("cuda")
|
| 34 |
|
| 35 |
+
#self.pipe3 = AutoPipelineForImage2Image.from_pipe(self.pipe2)
|
| 36 |
|
| 37 |
|
| 38 |
|
|
|
|
| 56 |
image = None
|
| 57 |
mask_image = None
|
| 58 |
|
| 59 |
+
|
| 60 |
+
image = self.pipe(
|
| 61 |
+
prompt=prompt,
|
| 62 |
+
image=image,
|
| 63 |
+
mask_image=mask_image,
|
| 64 |
+
guidance_scale=8.0,
|
| 65 |
+
num_inference_steps=20, # steps between 15 and 30 work well for us
|
| 66 |
+
strength=0.99, # make sure to use `strength` below 1.0
|
| 67 |
+
generator=generator,
|
| 68 |
+
).images[0]
|
| 69 |
+
|
| 70 |
+
return image
|
| 71 |
+
|
| 72 |
+
"""
|
| 73 |
+
pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")
|
| 74 |
+
|
| 75 |
self.pipe.enable_xformers_memory_efficient_attention()
|
| 76 |
|
| 77 |
# run inference pipeline
|
|
|
|
| 114 |
|
| 115 |
# return first generate PIL image
|
| 116 |
return result
|
| 117 |
+
"""
|
| 118 |
|
| 119 |
# helper to decode input image
|
| 120 |
def decode_base64_image(self, image_string):
|