InvokeAI/ldm/invoke/restoration/outcrop.py
Lincoln Stein 9141132a5c enhance outcropping with ability to direct contents of new regions
This commit does several things that improve the customizability of the CLI `outcrop` command:

1. When outcropping an image you can now add a `--new_prompt` option, to specify a new prompt to be applied to the outpainted region instead of the prompt used to generate the image.
2. Similarly you can provide a new seed using `--seed` (or `-S`). A seed less than zero will pick one randomly.
3. The metadata written into the outcropped file is now more informative about what was previously stored.
4. This PR also fixes the crash that happened when trying to outcrop an image  that does not contain InvokeAI metadata.

Other changes:

- add error checking suggested by @Kyle0654
- add special case in invoke.py to allow -1 to be passed as seed.
  This now only occurs for postprocessing commands. Previously, -1
  caused previous seed to be used, and this still applies to generate
  operations.
2022-11-11 20:34:21 +00:00

109 lines
3.9 KiB
Python

import warnings
import math
from PIL import Image, ImageFilter
class Outcrop(object):
def __init__(
self,
image,
generate, # current generate object
):
self.image = image
self.generate = generate
def process (
self,
extents:dict,
opt, # current options
orig_opt, # ones originally used to generate the image
image_callback = None,
prefix = None
):
# grow and mask the image
extended_image = self._extend_all(extents)
# switch samplers temporarily
curr_sampler = self.generate.sampler
self.generate.sampler_name = opt.sampler_name
self.generate._set_sampler()
def wrapped_callback(img,seed,**kwargs):
preferred_seed = orig_opt.seed if orig_opt.seed >= 0 else seed
image_callback(img,preferred_seed,use_prefix=prefix,**kwargs)
result= self.generate.prompt2image(
opt.prompt,
seed = opt.seed or orig_opt.seed,
sampler = self.generate.sampler,
steps = opt.steps,
cfg_scale = opt.cfg_scale,
ddim_eta = self.generate.ddim_eta,
width = extended_image.width,
height = extended_image.height,
init_img = extended_image,
strength = 0.90,
image_callback = wrapped_callback if image_callback else None,
seam_size = opt.seam_size or 96,
seam_blur = opt.seam_blur or 16,
seam_strength = opt.seam_strength or 0.7,
seam_steps = 20,
tile_size = 32,
color_match = True,
force_outpaint = True, # this just stops the warning about erased regions
)
# swap sampler back
self.generate.sampler = curr_sampler
return result
def _extend_all(
self,
extents:dict,
) -> Image:
'''
Extend the image in direction ('top','bottom','left','right') by
the indicated value. The image canvas is extended, and the empty
rectangular section will be filled with a blurred copy of the
adjacent image.
'''
image = self.image
for direction in extents:
assert direction in ['top', 'left', 'bottom', 'right'],'Direction must be one of "top", "left", "bottom", "right"'
pixels = extents[direction]
# round pixels up to the nearest 64
pixels = math.ceil(pixels/64) * 64
print(f'>> extending image {direction}ward by {pixels} pixels')
image = self._rotate(image,direction)
image = self._extend(image,pixels)
image = self._rotate(image,direction,reverse=True)
return image
def _rotate(self,image:Image,direction:str,reverse=False) -> Image:
'''
Rotates image so that the area to extend is always at the top top.
Simplifies logic later. The reverse argument, if true, will undo the
previous transpose.
'''
transposes = {
'right': ['ROTATE_90','ROTATE_270'],
'bottom': ['ROTATE_180','ROTATE_180'],
'left': ['ROTATE_270','ROTATE_90']
}
if direction not in transposes:
return image
transpose = transposes[direction][1 if reverse else 0]
return image.transpose(Image.Transpose.__dict__[transpose])
def _extend(self,image:Image,pixels:int)-> Image:
extended_img = Image.new('RGBA',(image.width,image.height+pixels))
extended_img.paste((0,0,0),[0,0,image.width,image.height+pixels])
extended_img.paste(image,box=(0,pixels))
# now make the top part transparent to use as a mask
alpha = extended_img.getchannel('A')
alpha.paste(0,(0,0,extended_img.width,pixels))
extended_img.putalpha(alpha)
return extended_img