|
|
|
|
|
|
|
|
import folder_paths, os |
|
|
from comfy.supported_models import FluxInpaint, models |
|
|
from nodes import UNETLoader |
|
|
try: |
|
|
from scepter.modules.utils.file_system import FS |
|
|
from scepter.modules.annotator.registry import ANNOTATORS |
|
|
from scepter.modules.utils.config import Config |
|
|
|
|
|
fs_list = [ |
|
|
Config(cfg_dict={"NAME": "HuggingfaceFs", "TEMP_DIR": "./"}, load=False), |
|
|
Config(cfg_dict={"NAME": "ModelscopeFs", "TEMP_DIR": "./"}, load=False), |
|
|
Config(cfg_dict={"NAME": "HttpFs", "TEMP_DIR": "./"}, load=False), |
|
|
Config(cfg_dict={"NAME": "LocalFs", "TEMP_DIR": "./"}, load=False) |
|
|
] |
|
|
|
|
|
for one_fs in fs_list: |
|
|
FS.init_fs_client(one_fs) |
|
|
SCEPTER = True |
|
|
except: |
|
|
SCEPTER = False |
|
|
|
|
|
class ACEPlus(FluxInpaint): |
|
|
unet_config = { |
|
|
"image_model": "flux", |
|
|
"guidance_embed": True, |
|
|
"in_channels": 112, |
|
|
} |
|
|
|
|
|
|
|
|
class ACEPlusFFTLoader(UNETLoader): |
|
|
|
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": {"unet_name": (folder_paths.get_filename_list("diffusion_models"), ), |
|
|
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],) |
|
|
}} |
|
|
RETURN_TYPES = ("MODEL",) |
|
|
FUNCTION = "load_unet" |
|
|
CATEGORY = "ComfyUI-ACE_Plus" |
|
|
|
|
|
def load_unet(self, unet_name, weight_dtype): |
|
|
models.append(ACEPlus) |
|
|
return super().load_unet(unet_name, weight_dtype) |
|
|
|
|
|
|
|
|
import torch |
|
|
import node_helpers |
|
|
|
|
|
|
|
|
class ACEPlusFFTConditioning: |
|
|
|
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": {"positive": ("CONDITIONING", ), |
|
|
"negative": ("CONDITIONING", ), |
|
|
"vae": ("VAE", ), |
|
|
"ucpixels": ("IMAGE", ), |
|
|
"cpixels": ("IMAGE", ), |
|
|
"mask": ("MASK", ), |
|
|
"noise_mask": ("BOOLEAN", {"default": True, "tooltip": "Add a noise mask to the latent " |
|
|
"so sampling will only happen " |
|
|
"within the mask. Might improve " |
|
|
"results or completely break " |
|
|
"things depending on the model."}), |
|
|
}} |
|
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") |
|
|
RETURN_NAMES = ("positive", "negative", "latent") |
|
|
FUNCTION = "encode" |
|
|
|
|
|
CATEGORY = "ComfyUI-ACE_Plus" |
|
|
|
|
|
def encode(self, |
|
|
positive, |
|
|
negative, |
|
|
vae, |
|
|
ucpixels, |
|
|
cpixels, |
|
|
mask, |
|
|
noise_mask=True): |
|
|
x = (ucpixels.shape[1] // 8) * 8 |
|
|
y = (ucpixels.shape[2] // 8) * 8 |
|
|
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), |
|
|
size=(ucpixels.shape[1], ucpixels.shape[2]), mode="bilinear") |
|
|
|
|
|
orig_pixels = ucpixels |
|
|
pixels = orig_pixels.clone() |
|
|
if pixels.shape[1] != x or pixels.shape[2] != y: |
|
|
x_offset = (pixels.shape[1] % 8) // 2 |
|
|
y_offset = (pixels.shape[2] % 8) // 2 |
|
|
pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :] |
|
|
mask = mask[:, :, x_offset:x + x_offset, y_offset:y + y_offset] |
|
|
|
|
|
orig_c_pixels = cpixels |
|
|
c_pixels = orig_c_pixels.clone() |
|
|
if orig_c_pixels.shape[1] != x or orig_c_pixels.shape[2] != y: |
|
|
x_offset = (orig_c_pixels.shape[1] % 8) // 2 |
|
|
y_offset = (orig_c_pixels.shape[2] % 8) // 2 |
|
|
c_pixels = orig_c_pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :] |
|
|
|
|
|
concat_latent = vae.encode(pixels) |
|
|
orig_latent = vae.encode(orig_pixels) |
|
|
c_concat_latent = vae.encode(c_pixels) |
|
|
|
|
|
out_latent = {"samples": orig_latent} |
|
|
if noise_mask: |
|
|
out_latent["noise_mask"] = mask |
|
|
|
|
|
out = [] |
|
|
for conditioning in [positive, negative]: |
|
|
c = node_helpers.conditioning_set_values(conditioning, { |
|
|
"concat_latent_image": torch.cat([concat_latent, c_concat_latent], dim=1), |
|
|
"concat_mask": mask}) |
|
|
out.append(c) |
|
|
|
|
|
return (out[0], out[1], out_latent) |
|
|
|
|
|
class ACEPlusLoraConditioning: |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": {"positive": ("CONDITIONING",), |
|
|
"negative": ("CONDITIONING",), |
|
|
"vae": ("VAE",), |
|
|
"pixels": ("IMAGE",), |
|
|
"mask": ("MASK",), |
|
|
"noise_mask": ("BOOLEAN", {"default": True, |
|
|
"tooltip": "Add a noise mask to the latent so sampling will only happen within the mask. Might improve results or completely break things depending on the model."}), |
|
|
}} |
|
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") |
|
|
RETURN_NAMES = ("positive", "negative", "latent") |
|
|
FUNCTION = "encode" |
|
|
|
|
|
CATEGORY = "ComfyUI-ACE_Plus" |
|
|
|
|
|
def encode(self, positive, negative, pixels, vae, mask, noise_mask=True): |
|
|
x = (pixels.shape[1] // 8) * 8 |
|
|
y = (pixels.shape[2] // 8) * 8 |
|
|
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), |
|
|
size=(pixels.shape[1], pixels.shape[2]), mode="bilinear") |
|
|
|
|
|
orig_pixels = pixels |
|
|
pixels = orig_pixels.clone() |
|
|
if pixels.shape[1] != x or pixels.shape[2] != y: |
|
|
x_offset = (pixels.shape[1] % 8) // 2 |
|
|
y_offset = (pixels.shape[2] % 8) // 2 |
|
|
pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :] |
|
|
mask = mask[:, :, x_offset:x + x_offset, y_offset:y + y_offset] |
|
|
|
|
|
concat_latent = vae.encode(pixels) |
|
|
orig_latent = vae.encode(orig_pixels) |
|
|
|
|
|
out_latent = {} |
|
|
|
|
|
out_latent["samples"] = orig_latent |
|
|
if noise_mask: |
|
|
out_latent["noise_mask"] = mask |
|
|
|
|
|
out = [] |
|
|
for conditioning in [positive, negative]: |
|
|
c = node_helpers.conditioning_set_values(conditioning, {"concat_latent_image": concat_latent, |
|
|
"concat_mask": mask}) |
|
|
out.append(c) |
|
|
return (out[0], out[1], out_latent) |
|
|
|
|
|
|
|
|
import torch |
|
|
import math |
|
|
import os |
|
|
import yaml |
|
|
import torchvision.transforms as T |
|
|
import numpy as np |
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
class AcePlusFFTProcessor: |
|
|
def __init__(self, |
|
|
max_aspect_ratio=4, |
|
|
d=16, |
|
|
max_seq_len=1024): |
|
|
self.max_aspect_ratio = max_aspect_ratio |
|
|
self.max_seq_len = max_seq_len |
|
|
self.d = d |
|
|
current_dir = os.path.dirname(os.path.abspath(__file__)) |
|
|
config_path = os.path.join(current_dir, 'config', 'ace_plus_fft_processor.yaml') |
|
|
self.processor_cfg = self.load_yaml(config_path) |
|
|
self.task_list = {} |
|
|
for task in self.processor_cfg['PREPROCESSOR']: |
|
|
self.task_list[task['TYPE']] = task |
|
|
self.transforms = T.Compose([ |
|
|
T.ToTensor(), |
|
|
T.Normalize(mean=[0, 0, 0], std=[1.0, 1.0, 1.0]) |
|
|
]) |
|
|
|
|
|
CATEGORY = 'ComfyUI-ACE_Plus' |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return { |
|
|
'required': { |
|
|
'use_reference': ('BOOLEAN', {'default': True}), |
|
|
'height': ('INT', { |
|
|
'default': 1024, |
|
|
'min': 256, |
|
|
'max': 1436, |
|
|
'step': 16 |
|
|
}), |
|
|
'width': ('INT', { |
|
|
'default': 1024, |
|
|
'min': 256, |
|
|
'max': 1436, |
|
|
'step': 16 |
|
|
}), |
|
|
'task_type': (list(s().task_list.keys()),), |
|
|
'keep_pixels_rate': ('FLOAT', { |
|
|
'default': 0.8, |
|
|
'min': 0, |
|
|
'max': 1, |
|
|
'step': 0.01 |
|
|
}), |
|
|
'max_seq_length': ('INT', { |
|
|
'default': 3072, |
|
|
'min': 1024, |
|
|
'max': 5120, |
|
|
'step': 0.01 |
|
|
}), |
|
|
}, |
|
|
'optional': { |
|
|
'reference_image': ('IMAGE',), |
|
|
'edit_image': ('IMAGE',), |
|
|
'edit_mask': ('MASK',), |
|
|
|
|
|
} |
|
|
} |
|
|
|
|
|
OUTPUT_NODE = True |
|
|
RETURN_TYPES = ('IMAGE', 'IMAGE', 'MASK', 'INT', 'INT', 'INT') |
|
|
RETURN_NAMES = ('UC_IMAGE', 'C_IMAGE', 'MASK', 'OUT_H', 'OUT_W', 'SLICE_W') |
|
|
FUNCTION = 'preprocess' |
|
|
|
|
|
def load_yaml(self, cfg_file): |
|
|
with open(cfg_file, 'r') as f: |
|
|
cfg = yaml.load(f.read(), Loader=yaml.SafeLoader) |
|
|
return cfg |
|
|
|
|
|
def image_check(self, image): |
|
|
if image is None: |
|
|
return image |
|
|
|
|
|
H, W = image.shape[1: 3] |
|
|
image = image.permute(0, 3, 1, 2) |
|
|
if H / W > self.max_aspect_ratio: |
|
|
image[0] = T.CenterCrop([int(self.max_aspect_ratio * W), W])(image[0]) |
|
|
elif W / H > self.max_aspect_ratio: |
|
|
image[0] = T.CenterCrop([H, int(self.max_aspect_ratio * H)])(image[0]) |
|
|
return image[0] |
|
|
|
|
|
def trans_pil_tensor(self, pil_image): |
|
|
transform = T.Compose([ |
|
|
T.ToTensor() |
|
|
]) |
|
|
tensor_image = transform(pil_image) |
|
|
return tensor_image |
|
|
|
|
|
def edit_preprocess(self, processor, device, edit_image, edit_mask): |
|
|
|
|
|
if edit_image is None or processor is None: |
|
|
return edit_image |
|
|
if not SCEPTER: |
|
|
raise ImportError(f'Please install scepter to use edit processor {processor} by ' |
|
|
f'runing "pip install scepter" in the conda env') |
|
|
processor = Config(cfg_dict=processor, load=False) |
|
|
processor = ANNOTATORS.build(processor).to(device) |
|
|
edit_image = Image.fromarray(np.array(edit_image[0] * 255).astype(np.uint8)).convert('RGB') |
|
|
new_edit_image = processor(np.asarray(edit_image)) |
|
|
|
|
|
del processor |
|
|
new_edit_image = Image.fromarray(new_edit_image) |
|
|
if edit_mask is not None: |
|
|
edit_mask = np.where(edit_mask > 0.5, 1, 0) * 255 |
|
|
edit_mask = Image.fromarray(np.array(edit_mask[0]).astype(np.uint8)).convert('L') |
|
|
|
|
|
if new_edit_image.size != edit_image.size: |
|
|
edit_image = T.Resize((edit_image.size[1], edit_image.size[0]), |
|
|
interpolation=T.InterpolationMode.BILINEAR, |
|
|
antialias=True)(new_edit_image) |
|
|
|
|
|
image = Image.composite(new_edit_image, edit_image, edit_mask) |
|
|
|
|
|
return self.trans_pil_tensor(image).unsqueeze(0).permute(0, 2, 3, 1) |
|
|
|
|
|
def preprocess(self, |
|
|
reference_image=None, |
|
|
edit_image=None, |
|
|
edit_mask=None, |
|
|
use_reference=True, |
|
|
task_type=None, |
|
|
height=1024, |
|
|
width=1024, |
|
|
keep_pixels_rate=0.8, |
|
|
max_seq_length=4096): |
|
|
self.max_seq_len = max_seq_length |
|
|
if not use_reference and edit_image is not None: |
|
|
reference_image = None |
|
|
if edit_mask is not None and edit_image is not None: |
|
|
iH, iW = edit_image.shape[1:3] |
|
|
mH, mW = edit_mask.shape[1:3] |
|
|
if iH != mH or iW != mW: |
|
|
edit_mask = torch.ones(edit_image.shape[:3]) |
|
|
|
|
|
if task_type != 'repainting': |
|
|
repainting_scale = 0 |
|
|
else: |
|
|
repainting_scale = 1 |
|
|
if task_type in self.task_list: |
|
|
edit_image = self.edit_preprocess(self.task_list[task_type]['ANNOTATOR'], 0, |
|
|
edit_image, edit_mask) |
|
|
if reference_image is not None: |
|
|
reference_image = self.image_check(reference_image) - 0.5 |
|
|
if edit_image is not None: |
|
|
edit_image = self.image_check(edit_image) - 0.5 |
|
|
|
|
|
if edit_image is None: |
|
|
edit_image = torch.zeros([3, height, width]) |
|
|
edit_mask = torch.ones([1, height, width]) |
|
|
else: |
|
|
if edit_mask is None: |
|
|
_, eH, eW = edit_image.shape |
|
|
edit_mask = np.ones((eH, eW)) |
|
|
else: |
|
|
edit_mask = np.asarray(edit_mask)[0] |
|
|
edit_mask = np.where(edit_mask > 0.5, 1, 0) |
|
|
edit_mask = edit_mask.astype( |
|
|
np.float32) if np.any(edit_mask) else np.ones_like(edit_mask).astype( |
|
|
np.float32) |
|
|
edit_mask = torch.tensor(edit_mask).unsqueeze(0) |
|
|
|
|
|
edit_image = edit_image * (1 - edit_mask * repainting_scale) |
|
|
|
|
|
out_h, out_w = edit_image.shape[-2:] |
|
|
|
|
|
assert edit_mask is not None |
|
|
if reference_image is not None: |
|
|
_, H, W = reference_image.shape |
|
|
_, eH, eW = edit_image.shape |
|
|
if not True: |
|
|
|
|
|
scale = eH / H |
|
|
tH, tW = eH, int(W * scale) |
|
|
reference_image = T.Resize((tH, tW), interpolation=T.InterpolationMode.BILINEAR, antialias=True)( |
|
|
reference_image) |
|
|
else: |
|
|
|
|
|
if H >= keep_pixels_rate * eH: |
|
|
tH = int(eH * keep_pixels_rate) |
|
|
scale = tH / H |
|
|
tW = int(W * scale) |
|
|
reference_image = T.Resize((tH, tW), interpolation=T.InterpolationMode.BILINEAR, antialias=True)( |
|
|
reference_image) |
|
|
rH, rW = reference_image.shape[-2:] |
|
|
delta_w = 0 |
|
|
delta_h = eH - rH |
|
|
padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2)) |
|
|
reference_image = T.Pad(padding, fill=0, padding_mode="constant")(reference_image) |
|
|
edit_image = torch.cat([reference_image, edit_image], dim=-1) |
|
|
edit_mask = torch.cat([torch.zeros([1, reference_image.shape[1], reference_image.shape[2]]), edit_mask], |
|
|
dim=-1) |
|
|
slice_w = reference_image.shape[-1] |
|
|
else: |
|
|
slice_w = 0 |
|
|
|
|
|
H, W = edit_image.shape[-2:] |
|
|
scale = min(1.0, math.sqrt(self.max_seq_len / ((H / self.d) * (W / self.d)))) |
|
|
rH = int(H * scale) // self.d * self.d |
|
|
rW = int(W * scale) // self.d * self.d |
|
|
slice_w = int(slice_w * scale) // self.d * self.d |
|
|
|
|
|
edit_image = T.Resize((rH, rW), interpolation=T.InterpolationMode.NEAREST_EXACT, antialias=True)(edit_image) |
|
|
edit_mask = T.Resize((rH, rW), interpolation=T.InterpolationMode.NEAREST_EXACT, antialias=True)(edit_mask) |
|
|
|
|
|
change_image = edit_image * edit_mask |
|
|
edit_image = edit_image * (1 - edit_mask) |
|
|
edit_image = edit_image.unsqueeze(0).permute(0, 2, 3, 1) |
|
|
change_image = change_image.unsqueeze(0).permute(0, 2, 3, 1) |
|
|
slice_w = slice_w if slice_w > 30 else slice_w + 30 |
|
|
|
|
|
return edit_image + 0.5, change_image + 0.5, edit_mask, out_h, out_w, slice_w |
|
|
|
|
|
class AcePlusLoraProcessor: |
|
|
def __init__(self, |
|
|
max_aspect_ratio=4, |
|
|
d=16, |
|
|
max_seq_len=1024): |
|
|
self.max_aspect_ratio = max_aspect_ratio |
|
|
self.max_seq_len = max_seq_len |
|
|
self.d = d |
|
|
current_dir = os.path.dirname(os.path.abspath(__file__)) |
|
|
config_path = os.path.join(current_dir, 'config', 'ace_plus_fft_processor.yaml') |
|
|
self.processor_cfg = self.load_yaml(config_path) |
|
|
self.task_list = {} |
|
|
for task in self.processor_cfg['PREPROCESSOR']: |
|
|
self.task_list[task['TYPE']] = task |
|
|
self.transforms = T.Compose([ |
|
|
T.ToTensor(), |
|
|
T.Normalize(mean=[0, 0, 0], std=[1.0, 1.0, 1.0]) |
|
|
]) |
|
|
|
|
|
CATEGORY = 'ComfyUI-ACE_Plus' |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return { |
|
|
'required': { |
|
|
'use_reference': ('BOOLEAN', {'default': True}), |
|
|
'height': ('INT', { |
|
|
'default': 1024, |
|
|
'min': 256, |
|
|
'max': 1436, |
|
|
'step': 16 |
|
|
}), |
|
|
'width': ('INT', { |
|
|
'default': 1024, |
|
|
'min': 256, |
|
|
'max': 1436, |
|
|
'step': 16 |
|
|
}), |
|
|
'task_type': (list(s().task_list.keys()),), |
|
|
'max_seq_length': ('INT', { |
|
|
'default': 3072, |
|
|
'min': 1024, |
|
|
'max': 5120, |
|
|
'step': 0.01 |
|
|
}), |
|
|
}, |
|
|
'optional': { |
|
|
'reference_image': ('IMAGE',), |
|
|
'edit_image': ('IMAGE',), |
|
|
'edit_mask': ('MASK',), |
|
|
|
|
|
} |
|
|
} |
|
|
|
|
|
OUTPUT_NODE = True |
|
|
RETURN_TYPES = ('IMAGE', 'MASK', 'INT', 'INT', 'INT') |
|
|
RETURN_NAMES = ('IMAGE', 'MASK', 'OUT_H', 'OUT_W', 'SLICE_W') |
|
|
FUNCTION = 'preprocess' |
|
|
|
|
|
def load_yaml(self, cfg_file): |
|
|
with open(cfg_file, 'r') as f: |
|
|
cfg = yaml.load(f.read(), Loader=yaml.SafeLoader) |
|
|
return cfg |
|
|
|
|
|
def image_check(self, image): |
|
|
if image is None: |
|
|
return image |
|
|
|
|
|
H, W = image.shape[1: 3] |
|
|
image = image.permute(0, 3, 1, 2) |
|
|
if H / W > self.max_aspect_ratio: |
|
|
image[0] = T.CenterCrop([int(self.max_aspect_ratio * W), W])(image[0]) |
|
|
elif W / H > self.max_aspect_ratio: |
|
|
image[0] = T.CenterCrop([H, int(self.max_aspect_ratio * H)])(image[0]) |
|
|
return image[0] |
|
|
|
|
|
def trans_pil_tensor(self, pil_image): |
|
|
transform = T.Compose([ |
|
|
T.ToTensor() |
|
|
]) |
|
|
tensor_image = transform(pil_image) |
|
|
return tensor_image |
|
|
|
|
|
def edit_preprocess(self, processor, device, edit_image, edit_mask): |
|
|
|
|
|
if edit_image is None or processor is None: |
|
|
return edit_image |
|
|
if not SCEPTER: |
|
|
raise ImportError(f'Please install scepter to use edit processor {processor} by ' |
|
|
f'runing "pip install scepter" in the conda env') |
|
|
processor = Config(cfg_dict=processor, load=False) |
|
|
processor = ANNOTATORS.build(processor).to(device) |
|
|
edit_image = Image.fromarray(np.array(edit_image[0] * 255).astype(np.uint8)).convert('RGB') |
|
|
new_edit_image = processor(np.asarray(edit_image)) |
|
|
|
|
|
del processor |
|
|
new_edit_image = Image.fromarray(new_edit_image) |
|
|
if edit_mask is not None: |
|
|
edit_mask = np.where(edit_mask > 0.5, 1, 0) * 255 |
|
|
edit_mask = Image.fromarray(np.array(edit_mask[0]).astype(np.uint8)).convert('L') |
|
|
|
|
|
if new_edit_image.size != edit_image.size: |
|
|
new_edit_image = T.Resize((edit_image.size[1], edit_image.size[0]), |
|
|
interpolation=T.InterpolationMode.BILINEAR, |
|
|
antialias=True)(new_edit_image) |
|
|
|
|
|
image = Image.composite(new_edit_image, edit_image, edit_mask) |
|
|
|
|
|
return self.trans_pil_tensor(image).unsqueeze(0).permute(0, 2, 3, 1) |
|
|
|
|
|
def preprocess(self, |
|
|
reference_image=None, |
|
|
edit_image=None, |
|
|
edit_mask=None, |
|
|
use_reference=True, |
|
|
task_type=None, |
|
|
height=1024, |
|
|
width=1024, |
|
|
max_seq_length=4096): |
|
|
self.max_seq_len = max_seq_length |
|
|
if not use_reference and edit_image is not None: |
|
|
reference_image = None |
|
|
if edit_mask is not None and edit_image is not None: |
|
|
iH, iW = edit_image.shape[1:3] |
|
|
mH, mW = edit_mask.shape[1:3] |
|
|
if iH != mH or iW != mW: |
|
|
edit_mask = torch.ones(edit_image.shape[:3]) |
|
|
|
|
|
if task_type != 'repainting': |
|
|
repainting_scale = 0.0 |
|
|
else: |
|
|
repainting_scale = 1.0 |
|
|
if task_type in self.task_list: |
|
|
edit_image = self.edit_preprocess(self.task_list[task_type]['ANNOTATOR'], 0, |
|
|
edit_image, edit_mask) |
|
|
if reference_image is not None: |
|
|
reference_image = self.image_check(reference_image) - 0.5 |
|
|
if edit_image is not None: |
|
|
edit_image = self.image_check(edit_image) - 0.5 |
|
|
|
|
|
if edit_image is None: |
|
|
edit_image = torch.zeros([3, height, width]) |
|
|
edit_mask = torch.ones([1, height, width]) |
|
|
else: |
|
|
if edit_mask is None: |
|
|
_, eH, eW = edit_image.shape |
|
|
edit_mask = np.ones((eH, eW)) |
|
|
else: |
|
|
edit_mask = np.asarray(edit_mask)[0] |
|
|
edit_mask = np.where(edit_mask > 0.5, 1, 0) |
|
|
edit_mask = edit_mask.astype( |
|
|
np.float32) if np.any(edit_mask) else np.ones_like(edit_mask).astype( |
|
|
np.float32) |
|
|
edit_mask = torch.tensor(edit_mask).unsqueeze(0) |
|
|
|
|
|
edit_image = edit_image * (1 - edit_mask * repainting_scale) |
|
|
|
|
|
out_h, out_w = edit_image.shape[-2:] |
|
|
|
|
|
assert edit_mask is not None |
|
|
if reference_image is not None: |
|
|
_, H, W = reference_image.shape |
|
|
_, eH, eW = edit_image.shape |
|
|
|
|
|
scale = eH / H |
|
|
tH, tW = eH, int(W * scale) |
|
|
reference_image = T.Resize((tH, tW), interpolation=T.InterpolationMode.BILINEAR, antialias=True)( |
|
|
reference_image) |
|
|
edit_image = torch.cat([reference_image, edit_image], dim=-1) |
|
|
edit_mask = torch.cat([torch.zeros([1, reference_image.shape[1], reference_image.shape[2]]), edit_mask], |
|
|
dim=-1) |
|
|
slice_w = reference_image.shape[-1] |
|
|
else: |
|
|
slice_w = 0 |
|
|
|
|
|
H, W = edit_image.shape[-2:] |
|
|
scale = min(1.0, math.sqrt(self.max_seq_len * 2 / ((H / self.d) * (W / self.d)))) |
|
|
rH = int(H * scale) // self.d * self.d |
|
|
rW = int(W * scale) // self.d * self.d |
|
|
slice_w = int(slice_w * scale) // self.d * self.d |
|
|
|
|
|
edit_image = T.Resize((rH, rW), interpolation=T.InterpolationMode.NEAREST_EXACT, antialias=True)(edit_image) |
|
|
edit_mask = T.Resize((rH, rW), interpolation=T.InterpolationMode.NEAREST_EXACT, antialias=True)(edit_mask) |
|
|
|
|
|
edit_image = edit_image.unsqueeze(0).permute(0, 2, 3, 1) |
|
|
slice_w = slice_w if slice_w < 30 else slice_w + 30 |
|
|
|
|
|
return edit_image + 0.5, edit_mask, out_h, out_w, slice_w |
|
|
|
