preprocessor_inpaint.py

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.utils 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 PreprocessorInpaintNoobAIXL(Preprocessor):
    def __init__(self):
        super().__init__()
        self.name = 'inpaint_noobai_xl'
        self.tags = ['Inpaint']
        self.model_filename_filters = ['inpaint', 'noobai']
        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 __call__(self, input_image, resolution=512, slider_1=None, slider_2=None, slider_3=None, input_mask=None, **kwargs):
        if input_mask is None:
            return input_image

        if not isinstance(input_image, np.ndarray):
            input_image = np.array(input_image)
        if not isinstance(input_mask, np.ndarray):
            input_mask = np.array(input_mask)

        mask = input_mask.astype(np.float32) / 255.0
        mask = (mask > 0.5).astype(np.float32)

        # Create a copy of the input image
        result = input_image.copy()

        # Convert mask to proper shape if needed
        if mask.ndim == 2:
            mask = np.expand_dims(mask, axis=-1)
        if mask.shape[-1] == 1:
            mask = np.repeat(mask, 3, axis=-1)

        mask_indices = mask > 0.5
        result[mask_indices] = 0.0

        return result

    def process_before_every_sampling(self, process, cond, mask, *args, **kwargs):
        mask = mask.round()
        mixed_cond = cond.clone()
        mixed_cond = mixed_cond * (1.0 - 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
        # This is a powerful VAE with integrated memory management, bf16, and tiled fallback.

        latent_image = vae.encode(self.image.movedim(1, -1))
        latent_image = process.sd_model.forge_objects.vae.first_stage_model.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.predictor.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(PreprocessorInpaintNoobAIXL())

add_supported_preprocessor(PreprocessorInpaint())

add_supported_preprocessor(PreprocessorInpaintOnly())

add_supported_preprocessor(PreprocessorInpaintLama())