Upload preprocessor_inpaint.py

preprocessor_inpaint.py

@@ -0,0 +1,209 @@
+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())