【stable-diffusion-2-inpainting】要らないものが描写されたからinpaintで消す

はじめに

Stable Diffusionで画像を生成していると要らないもの描写されていることがあります。

そんな時はinpaintで消すのが良いと思います。

マスク画像を作る

マスク画像は四角で囲むだけなら以下のスクリプトで簡単に作成できます。

import cv2
import numpy as np

from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument('--image', type=str, help='original image' )
args = parser.parse_args()

original_image = args.image

img = cv2.imread(original_image)

source_window = "make_mask"
cv2.namedWindow(source_window)

rect = cv2.selectROI(source_window, img, False, False)
cv2.destroyAllWindows()

xmin, ymin, width, height = rect

maskimg = np.zeros(img.shape)

cv2.rectangle(maskimg, (xmin, ymin), (xmin+width, ymin+height), (255,255,255), -1)

cv2.imwrite('mask.png', maskimg)

実行

マスク画像が作成できたら以下を実行するのみです。
promptは元々使用したものです。

from PIL import Image
import torch
from diffusers import StableDiffusionInpaintPipeline
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--image', type=str, help='original image' )
parser.add_argument('--mask', type=str, help='mask image' )
args = parser.parse_args()

original_image = args.image
mask_image = args.mask

image = Image.open(original_image).convert('RGB').resize((512, 512))
mask = Image.open(mask_image).convert('RGB').resize((512, 512))


pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "./stable-diffusion-2-inpainting",
    torch_dtype=torch.float16,
)
pipe.to("cuda")
pipe.enable_attention_slicing()

prompt = "anime of tsundere moe kawaii beautiful girl"

image = pipe(
    prompt=prompt, 
    image=image, 
    mask_image=mask).images[0]
image.save("result.png")

結果