|
from typing import List, Union |
|
|
|
import torch |
|
from diffusers import StableDiffusionInpaintPipeline |
|
|
|
from internals.pipelines.commons import AbstractPipeline |
|
from internals.util.commons import disable_safety_checker, download_image |
|
|
|
|
|
class InPainter(AbstractPipeline): |
|
def load(self): |
|
self.pipe = StableDiffusionInpaintPipeline.from_pretrained( |
|
"jayparmr/icbinp_v8_inpaint_v2", |
|
torch_dtype=torch.float16, |
|
).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 |
|
|
