71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
from .utils.color import extract_pad_color, get_empty_color
|
|
from .categories import CATE_IMAGE
|
|
import torch
|
|
|
|
MAX_RESOLUTION = 8192
|
|
|
|
|
|
class FEImagePadForOutpaint:
|
|
"""
|
|
图像外扩节点，便于图像生成
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
|
|
"top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
|
|
"right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
|
|
"bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
|
|
"feathering": ("INT", {"default": 50, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"pad_color": ("RGB_COLOR", {"default": get_empty_color()}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "expand_image"
|
|
CATEGORY = CATE_IMAGE
|
|
|
|
def expand_image(self, image, left, top, right, bottom, feathering, pad_color):
|
|
batch_size, img_h, img_w, colors = image.size()
|
|
|
|
colors = extract_pad_color(pad_color, batch_size) / 255
|
|
new_image = torch.ones(
|
|
(batch_size, img_h + top + bottom, img_w + left + right, 3),
|
|
dtype=torch.float32,
|
|
).expand(batch_size, -1, -1, -1) * colors.view(batch_size, 1, 1, 3)
|
|
|
|
new_image[:, top:top + img_h, left:left + img_w, :] = image
|
|
|
|
mask = torch.ones(
|
|
(img_h + top + bottom, img_w + left + right),
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
# 处理mask羽化
|
|
if feathering > 0 and feathering * 2 < img_h and feathering * 2 < img_w:
|
|
# distances to border
|
|
mi, mj = torch.meshgrid(
|
|
torch.arange(img_h, dtype=torch.float32),
|
|
torch.arange(img_w, dtype=torch.float32),
|
|
indexing='ij',
|
|
)
|
|
distances = torch.minimum(
|
|
torch.minimum(mi, mj),
|
|
torch.minimum(img_h - 1 - mi, img_w - 1 - mj),
|
|
)
|
|
# convert distances to square falloff from 1 to 0
|
|
t = (feathering - distances) / feathering
|
|
t.clamp_(min=0)
|
|
t.square_()
|
|
|
|
mask[top:top + img_h, left:left + img_w] = t
|
|
else:
|
|
mask[top:top + img_h, left:left + img_w] = torch.zeros(
|
|
(img_h, img_w),
|
|
dtype=torch.float32,
|
|
)
|
|
return (new_image, mask)
|