File size: 1,547 Bytes
19b3da3 1bc457e 19b3da3 1bc457e 19b3da3 0daeeb0 19b3da3 1bc457e 19b3da3 1bc457e 19b3da3 |
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 42 43 44 45 46 47 48 49 50 51 52 |
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
from internals.util.config import get_hf_cache_dir
class InPainter(AbstractPipeline):
def load(self):
self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
"jayparmr/icbinp_v8_inpaint_v2",
torch_dtype=torch.float16,
cache_dir=get_hf_cache_dir(),
).to("cuda")
disable_safety_checker(self.pipe)
def create(self, pipeline: AbstractPipeline):
self.pipe = StableDiffusionInpaintPipeline(**pipeline.pipe.components).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]],
steps: int = 50,
):
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,
num_inference_steps=steps,
).images
|