File size: 1,012 Bytes
bc05b03 |
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 |
import torch
from diffusers import AutoPipelineForInpainting
from torchvision.transforms.functional import to_pil_image
from PIL import Image
class InpaintingModel:
def __init__(self) -> None:
pass
def generate(self, image: torch.Tensor, mask_image: torch.Tensor, prompt: str) -> Image.Image:
pass
class KandingskyInpaintingModel(InpaintingModel):
def __init__(
self,
device = torch.device("cpu"),
) -> None:
super().__init__()
self.device = device
self.model = AutoPipelineForInpainting.from_pretrained("kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16)
self.model.enable_model_cpu_offload()
self.negative_prompt = "deformed, ugly, disfigured"
def generate(self, image: Image.Image, mask_image: Image.Image, prompt: str) -> Image.Image:
output = self.model(prompt=prompt, negative_prompt=self.negative_prompt, image=image, mask_image=mask_image)
return output.images[0]
|