File size: 1,135 Bytes
4adca93 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
from typing import List, Union
import torch
from diffusers import StableDiffusionInpaintPipeline
from util.commons import disable_safety_checker, download_image
class InPainter:
def load(self):
self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting",
torch_dtype=torch.float16,
revision="fp16",
).to("cuda")
disable_safety_checker(self.pipe)
@torch.inference_mode()
def process(
self,
image_url: str,
mask_image_url: str,
width: int,
height: int,
seed: int,
prompt: Union[str, List[str]],
negative_prompt: Union[str, List[str]],
):
torch.manual_seed(seed)
input_img = download_image(image_url).resize((width, height))
mask_img = download_image(mask_image_url).resize((width, height))
return self.pipe.__call__(
prompt=prompt,
image=input_img,
mask_image=mask_img,
height=height,
width=width,
negative_prompt=negative_prompt,
).images
|