|
import os |
|
import cv2 |
|
import torch |
|
import numpy as np |
|
import yaml |
|
import einops |
|
|
|
from omegaconf import OmegaConf |
|
from modules_forge.supported_preprocessor import Preprocessor, PreprocessorParameter |
|
from modules_forge.forge_util import numpy_to_pytorch, resize_image_with_pad |
|
from modules_forge.shared import preprocessor_dir, add_supported_preprocessor |
|
from modules.modelloader import load_file_from_url |
|
from annotator.lama.saicinpainting.training.trainers import load_checkpoint |
|
|
|
|
|
class PreprocessorInpaint(Preprocessor): |
|
def __init__(self): |
|
super().__init__() |
|
self.name = 'inpaint_global_harmonious' |
|
self.tags = ['Inpaint'] |
|
self.model_filename_filters = ['inpaint'] |
|
self.slider_resolution = PreprocessorParameter(visible=False) |
|
self.fill_mask_with_one_when_resize_and_fill = True |
|
self.expand_mask_when_resize_and_fill = True |
|
|
|
def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): |
|
mask = mask.round() |
|
mixed_cond = cond * (1.0 - mask) - mask |
|
return mixed_cond, None |
|
|
|
|
|
class PreprocessorInpaintOnly(PreprocessorInpaint): |
|
def __init__(self): |
|
super().__init__() |
|
self.name = 'inpaint_only' |
|
self.image = None |
|
self.mask = None |
|
self.latent = None |
|
|
|
def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): |
|
mask = mask.round() |
|
self.image = cond |
|
self.mask = mask |
|
|
|
vae = process.sd_model.forge_objects.vae |
|
|
|
|
|
latent_image = vae.encode(self.image.movedim(1, -1)) |
|
latent_image = process.sd_model.forge_objects.unet.model.latent_format.process_in(latent_image) |
|
|
|
B, C, H, W = latent_image.shape |
|
|
|
latent_mask = self.mask |
|
latent_mask = torch.nn.functional.interpolate(latent_mask, size=(H * 8, W * 8), mode="bilinear").round() |
|
latent_mask = torch.nn.functional.max_pool2d(latent_mask, (8, 8)).round().to(latent_image) |
|
|
|
unet = process.sd_model.forge_objects.unet.clone() |
|
|
|
def pre_cfg(model, c, uc, x, timestep, model_options): |
|
noisy_latent = latent_image.to(x) + timestep[:, None, None, None].to(x) * torch.randn_like(latent_image).to(x) |
|
x = x * latent_mask.to(x) + noisy_latent.to(x) * (1.0 - latent_mask.to(x)) |
|
return model, c, uc, x, timestep, model_options |
|
|
|
def post_cfg(args): |
|
denoised = args['denoised'] |
|
denoised = denoised * latent_mask.to(denoised) + latent_image.to(denoised) * (1.0 - latent_mask.to(denoised)) |
|
return denoised |
|
|
|
unet.add_sampler_pre_cfg_function(pre_cfg) |
|
unet.set_model_sampler_post_cfg_function(post_cfg) |
|
|
|
process.sd_model.forge_objects.unet = unet |
|
|
|
self.latent = latent_image |
|
|
|
mixed_cond = cond * (1.0 - mask) - mask |
|
|
|
return mixed_cond, None |
|
|
|
def process_after_every_sampling(self, process, params, *args, **kwargs): |
|
a1111_batch_result = args[0] |
|
new_results = [] |
|
|
|
for img in a1111_batch_result.images: |
|
sigma = 7 |
|
mask = self.mask[0, 0].detach().cpu().numpy().astype(np.float32) |
|
mask = cv2.dilate(mask, np.ones((sigma, sigma), dtype=np.uint8)) |
|
mask = cv2.blur(mask, (sigma, sigma))[None] |
|
mask = torch.from_numpy(np.ascontiguousarray(mask).copy()).to(img).clip(0, 1) |
|
raw = self.image[0].to(img).clip(0, 1) |
|
img = img.clip(0, 1) |
|
new_results.append(raw * (1.0 - mask) + img * mask) |
|
|
|
a1111_batch_result.images = new_results |
|
return |
|
|
|
|
|
class PreprocessorInpaintLama(PreprocessorInpaintOnly): |
|
def __init__(self): |
|
super().__init__() |
|
self.name = 'inpaint_only+lama' |
|
|
|
def load_model(self): |
|
remote_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetLama.pth" |
|
model_path = load_file_from_url(remote_model_path, model_dir=preprocessor_dir) |
|
config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lama_config.yaml') |
|
cfg = yaml.safe_load(open(config_path, 'rt')) |
|
cfg = OmegaConf.create(cfg) |
|
cfg.training_model.predict_only = True |
|
cfg.visualizer.kind = 'noop' |
|
model = load_checkpoint(cfg, os.path.abspath(model_path), strict=False, map_location='cpu') |
|
self.setup_model_patcher(model) |
|
return |
|
|
|
def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, input_mask=None, **kwargs): |
|
if input_mask is None: |
|
return input_image |
|
|
|
H, W, C = input_image.shape |
|
raw_color = input_image.copy() |
|
raw_mask = input_mask.copy() |
|
|
|
input_image, remove_pad = resize_image_with_pad(input_image, 256) |
|
input_mask, remove_pad = resize_image_with_pad(input_mask, 256) |
|
input_mask = input_mask[..., :1] |
|
|
|
self.load_model() |
|
|
|
self.move_all_model_patchers_to_gpu() |
|
|
|
color = np.ascontiguousarray(input_image).astype(np.float32) / 255.0 |
|
mask = np.ascontiguousarray(input_mask).astype(np.float32) / 255.0 |
|
with torch.no_grad(): |
|
color = self.send_tensor_to_model_device(torch.from_numpy(color)) |
|
mask = self.send_tensor_to_model_device(torch.from_numpy(mask)) |
|
mask = (mask > 0.5).float() |
|
color = color * (1 - mask) |
|
image_feed = torch.cat([color, mask], dim=2) |
|
image_feed = einops.rearrange(image_feed, 'h w c -> 1 c h w') |
|
prd_color = self.model_patcher.model(image_feed)[0] |
|
prd_color = einops.rearrange(prd_color, 'c h w -> h w c') |
|
prd_color = prd_color * mask + color * (1 - mask) |
|
prd_color *= 255.0 |
|
prd_color = prd_color.detach().cpu().numpy().clip(0, 255).astype(np.uint8) |
|
|
|
prd_color = remove_pad(prd_color) |
|
prd_color = cv2.resize(prd_color, (W, H)) |
|
|
|
alpha = raw_mask.astype(np.float32) / 255.0 |
|
fin_color = prd_color.astype(np.float32) * alpha + raw_color.astype(np.float32) * (1 - alpha) |
|
fin_color = fin_color.clip(0, 255).astype(np.uint8) |
|
|
|
return fin_color |
|
|
|
def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): |
|
cond, mask = super().process_before_every_sampling(process, cond, mask, *args, **kwargs) |
|
sigma_max = process.sd_model.forge_objects.unet.model.model_sampling.sigma_max |
|
original_noise = kwargs['noise'] |
|
process.modified_noise = original_noise + self.latent.to(original_noise) / sigma_max.to(original_noise) |
|
return cond, mask |
|
|
|
|
|
add_supported_preprocessor(PreprocessorInpaint()) |
|
|
|
add_supported_preprocessor(PreprocessorInpaintOnly()) |
|
|
|
add_supported_preprocessor(PreprocessorInpaintLama()) |
|
|
