mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
103 lines
3.5 KiB
Python
103 lines
3.5 KiB
Python
import math
|
|
import warnings
|
|
|
|
from PIL import Image, ImageFilter
|
|
|
|
|
|
class Outpaint(object):
|
|
def __init__(self, image, generate):
|
|
self.image = image
|
|
self.generate = generate
|
|
|
|
def process(self, opt, old_opt, image_callback=None, prefix=None):
|
|
image = self._create_outpaint_image(self.image, opt.out_direction)
|
|
|
|
seed = old_opt.seed
|
|
prompt = old_opt.prompt
|
|
|
|
def wrapped_callback(img, seed, **kwargs):
|
|
image_callback(img, seed, use_prefix=prefix, **kwargs)
|
|
|
|
return self.generate.prompt2image(
|
|
prompt,
|
|
seed=seed,
|
|
sampler=self.generate.sampler,
|
|
steps=opt.steps,
|
|
cfg_scale=opt.cfg_scale,
|
|
ddim_eta=self.generate.ddim_eta,
|
|
width=opt.width,
|
|
height=opt.height,
|
|
init_img=image,
|
|
strength=0.83,
|
|
image_callback=wrapped_callback,
|
|
prefix=prefix,
|
|
)
|
|
|
|
def _create_outpaint_image(self, image, direction_args):
|
|
assert len(direction_args) in [
|
|
1,
|
|
2,
|
|
], "Direction (-D) must have exactly one or two arguments."
|
|
|
|
if len(direction_args) == 1:
|
|
direction = direction_args[0]
|
|
pixels = None
|
|
elif len(direction_args) == 2:
|
|
direction = direction_args[0]
|
|
pixels = int(direction_args[1])
|
|
|
|
assert direction in [
|
|
"top",
|
|
"left",
|
|
"bottom",
|
|
"right",
|
|
], 'Direction (-D) must be one of "top", "left", "bottom", "right"'
|
|
|
|
image = image.convert("RGBA")
|
|
# we always extend top, but rotate to extend along the requested side
|
|
if direction == "left":
|
|
image = image.transpose(Image.Transpose.ROTATE_270)
|
|
elif direction == "bottom":
|
|
image = image.transpose(Image.Transpose.ROTATE_180)
|
|
elif direction == "right":
|
|
image = image.transpose(Image.Transpose.ROTATE_90)
|
|
|
|
pixels = image.height // 2 if pixels is None else int(pixels)
|
|
assert (
|
|
0 < pixels < image.height
|
|
), "Direction (-D) pixels length must be in the range 0 - image.size"
|
|
|
|
# the top part of the image is taken from the source image mirrored
|
|
# coordinates (0,0) are the upper left corner of an image
|
|
top = image.transpose(Image.Transpose.FLIP_TOP_BOTTOM).convert("RGBA")
|
|
top = top.crop((0, top.height - pixels, top.width, top.height))
|
|
|
|
# setting all alpha of the top part to 0
|
|
alpha = top.getchannel("A")
|
|
alpha.paste(0, (0, 0, top.width, top.height))
|
|
top.putalpha(alpha)
|
|
|
|
# taking the bottom from the original image
|
|
bottom = image.crop((0, 0, image.width, image.height - pixels))
|
|
|
|
new_img = image.copy()
|
|
new_img.paste(top, (0, 0))
|
|
new_img.paste(bottom, (0, pixels))
|
|
|
|
# create a 10% dither in the middle
|
|
dither = min(image.height // 10, pixels)
|
|
for x in range(0, image.width, 2):
|
|
for y in range(pixels - dither, pixels + dither):
|
|
(r, g, b, a) = new_img.getpixel((x, y))
|
|
new_img.putpixel((x, y), (r, g, b, 0))
|
|
|
|
# let's rotate back again
|
|
if direction == "left":
|
|
new_img = new_img.transpose(Image.Transpose.ROTATE_90)
|
|
elif direction == "bottom":
|
|
new_img = new_img.transpose(Image.Transpose.ROTATE_180)
|
|
elif direction == "right":
|
|
new_img = new_img.transpose(Image.Transpose.ROTATE_270)
|
|
|
|
return new_img
|