path: root/scripts/soft_inpainting.py

import numpy as np
import gradio as gr
import math
from modules.ui_components import InputAccordion
import modules.scripts as scripts


class SoftInpaintingSettings:
    def __init__(self,
                 mask_blend_power,
                 mask_blend_scale,
                 inpaint_detail_preservation,
                 composite_mask_influence,
                 composite_difference_threshold,
                 composite_difference_contrast):
        self.mask_blend_power = mask_blend_power
        self.mask_blend_scale = mask_blend_scale
        self.inpaint_detail_preservation = inpaint_detail_preservation
        self.composite_mask_influence = composite_mask_influence
        self.composite_difference_threshold = composite_difference_threshold
        self.composite_difference_contrast = composite_difference_contrast

    def add_generation_params(self, dest):
        dest[enabled_gen_param_label] = True
        dest[gen_param_labels.mask_blend_power] = self.mask_blend_power
        dest[gen_param_labels.mask_blend_scale] = self.mask_blend_scale
        dest[gen_param_labels.inpaint_detail_preservation] = self.inpaint_detail_preservation
        dest[gen_param_labels.composite_mask_influence] = self.composite_mask_influence
        dest[gen_param_labels.composite_difference_threshold] = self.composite_difference_threshold
        dest[gen_param_labels.composite_difference_contrast] = self.composite_difference_contrast


# ------------------- Methods -------------------

def processing_uses_inpainting(p):
    # TODO: Figure out a better way to determine if inpainting is being used by p
    if getattr(p, "image_mask", None) is not None:
        return True

    if getattr(p, "mask", None) is not None:
        return True

    if getattr(p, "nmask", None) is not None:
        return True

    return False


def latent_blend(settings, a, b, t):
    """
    Interpolates two latent image representations according to the parameter t,
    where the interpolated vectors' magnitudes are also interpolated separately.
    The "detail_preservation" factor biases the magnitude interpolation towards
    the larger of the two magnitudes.
    """
    import torch

    # NOTE: We use inplace operations wherever possible.

    # [4][w][h] to [1][4][w][h]
    t2 = t.unsqueeze(0)
    # [4][w][h] to [1][1][w][h] - the [4] seem redundant.
    t3 = t[0].unsqueeze(0).unsqueeze(0)

    one_minus_t2 = 1 - t2
    one_minus_t3 = 1 - t3

    # Linearly interpolate the image vectors.
    a_scaled = a * one_minus_t2
    b_scaled = b * t2
    image_interp = a_scaled
    image_interp.add_(b_scaled)
    result_type = image_interp.dtype
    del a_scaled, b_scaled, t2, one_minus_t2

    # Calculate the magnitude of the interpolated vectors. (We will remove this magnitude.)
    # 64-bit operations are used here to allow large exponents.
    current_magnitude = torch.norm(image_interp, p=2, dim=1, keepdim=True).to(torch.float64).add_(0.00001)

    # Interpolate the powered magnitudes, then un-power them (bring them back to a power of 1).
    a_magnitude = torch.norm(a, p=2, dim=1, keepdim=True).to(torch.float64).pow_(
        settings.inpaint_detail_preservation) * one_minus_t3
    b_magnitude = torch.norm(b, p=2, dim=1, keepdim=True).to(torch.float64).pow_(
        settings.inpaint_detail_preservation) * t3
    desired_magnitude = a_magnitude
    desired_magnitude.add_(b_magnitude).pow_(1 / settings.inpaint_detail_preservation)
    del a_magnitude, b_magnitude, t3, one_minus_t3

    # Change the linearly interpolated image vectors' magnitudes to the value we want.
    # This is the last 64-bit operation.
    image_interp_scaling_factor = desired_magnitude
    image_interp_scaling_factor.div_(current_magnitude)
    image_interp_scaling_factor = image_interp_scaling_factor.to(result_type)
    image_interp_scaled = image_interp
    image_interp_scaled.mul_(image_interp_scaling_factor)
    del current_magnitude
    del desired_magnitude
    del image_interp
    del image_interp_scaling_factor
    del result_type

    return image_interp_scaled


def get_modified_nmask(settings, nmask, sigma):
    """
    Converts a negative mask representing the transparency of the original latent vectors being overlayed
    to a mask that is scaled according to the denoising strength for this step.

    Where:
        0 = fully opaque, infinite density, fully masked
        1 = fully transparent, zero density, fully unmasked

    We bring this transparency to a power, as this allows one to simulate N number of blending operations
    where N can be any positive real value. Using this one can control the balance of influence between
    the denoiser and the original latents according to the sigma value.

    NOTE: "mask" is not used
    """
    import torch
    return torch.pow(nmask, (sigma ** settings.mask_blend_power) * settings.mask_blend_scale)


def apply_adaptive_masks(
        settings: SoftInpaintingSettings,
        nmask,
        latent_orig,
        latent_processed,
        overlay_images,
        width, height,
        paste_to):
    import torch
    import modules.processing as proc
    import modules.images as images
    from PIL import Image, ImageOps, ImageFilter

    # TODO: Bias the blending according to the latent mask, add adjustable parameter for bias control.
    latent_mask = nmask[0].float()
    # convert the original mask into a form we use to scale distances for thresholding
    mask_scalar = 1 - (torch.clamp(latent_mask, min=0, max=1) ** (settings.mask_blend_scale / 2))
    mask_scalar = (0.5 * (1 - settings.composite_mask_influence)
                   + mask_scalar * settings.composite_mask_influence)
    mask_scalar = mask_scalar / (1.00001 - mask_scalar)
    mask_scalar = mask_scalar.cpu().numpy()

    latent_distance = torch.norm(latent_processed - latent_orig, p=2, dim=1)

    kernel, kernel_center = get_gaussian_kernel(stddev_radius=1.5, max_radius=2)

    masks_for_overlay = []

    for i, (distance_map, overlay_image) in enumerate(zip(latent_distance, overlay_images)):
        converted_mask = distance_map.float().cpu().numpy()
        converted_mask = weighted_histogram_filter(converted_mask, kernel, kernel_center,
                                                   percentile_min=0.9, percentile_max=1, min_width=1)
        converted_mask = weighted_histogram_filter(converted_mask, kernel, kernel_center,
                                                   percentile_min=0.25, percentile_max=0.75, min_width=1)

        # The distance at which opacity of original decreases to 50%
        half_weighted_distance = settings.composite_difference_threshold * mask_scalar
        converted_mask = converted_mask / half_weighted_distance

        converted_mask = 1 / (1 + converted_mask ** settings.composite_difference_contrast)
        converted_mask = smootherstep(converted_mask)
        converted_mask = 1 - converted_mask
        converted_mask = 255. * converted_mask
        converted_mask = converted_mask.astype(np.uint8)
        converted_mask = Image.fromarray(converted_mask)
        converted_mask = images.resize_image(2, converted_mask, width, height)
        converted_mask = proc.create_binary_mask(converted_mask, round=False)

        # Remove aliasing artifacts using a gaussian blur.
        converted_mask = converted_mask.filter(ImageFilter.GaussianBlur(radius=4))

        # Expand the mask to fit the whole image if needed.
        if paste_to is not None:
            converted_mask = proc.uncrop(converted_mask,
                                         (overlay_image.width, overlay_image.height),
                                         paste_to)

        masks_for_overlay.append(converted_mask)

        image_masked = Image.new('RGBa', (overlay_image.width, overlay_image.height))
        image_masked.paste(overlay_image.convert("RGBA").convert("RGBa"),
                           mask=ImageOps.invert(converted_mask.convert('L')))

        overlay_images[i] = image_masked.convert('RGBA')

    return masks_for_overlay


def apply_masks(
        settings,
        nmask,
        overlay_images,
        width, height,
        paste_to):
    import torch
    import modules.processing as proc
    import modules.images as images
    from PIL import Image, ImageOps, ImageFilter

    converted_mask = nmask[0].float()
    converted_mask = torch.clamp(converted_mask, min=0, max=1).pow_(settings.mask_blend_scale / 2)
    converted_mask = 255. * converted_mask
    converted_mask = converted_mask.cpu().numpy().astype(np.uint8)
    converted_mask = Image.fromarray(converted_mask)
    converted_mask = images.resize_image(2, converted_mask, width, height)
    converted_mask = proc.create_binary_mask(converted_mask, round=False)

    # Remove aliasing artifacts using a gaussian blur.
    converted_mask = converted_mask.filter(ImageFilter.GaussianBlur(radius=4))

    # Expand the mask to fit the whole image if needed.
    if paste_to is not None:
        converted_mask = proc.uncrop(converted_mask,
                                     (width, height),
                                     paste_to)

    masks_for_overlay = []

    for i, overlay_image in enumerate(overlay_images):
        masks_for_overlay[i] = converted_mask

        image_masked = Image.new('RGBa', (overlay_image.width, overlay_image.height))
        image_masked.paste(overlay_image.convert("RGBA").convert("RGBa"),
                           mask=ImageOps.invert(converted_mask.convert('L')))

        overlay_images[i] = image_masked.convert('RGBA')

    return masks_for_overlay


def weighted_histogram_filter(img, kernel, kernel_center, percentile_min=0.0, percentile_max=1.0, min_width=1.0):
    """
    Generalization convolution filter capable of applying
    weighted mean, median, maximum, and minimum filters
    parametrically using an arbitrary kernel.

    Args:
        img (nparray):
            The image, a 2-D array of floats, to which the filter is being applied.
        kernel (nparray):
            The kernel, a 2-D array of floats.
        kernel_center (nparray):
            The kernel center coordinate, a 1-D array with two elements.
        percentile_min (float):
            The lower bound of the histogram window used by the filter,
            from 0 to 1.
        percentile_max (float):
            The upper bound of the histogram window used by the filter,
            from 0 to 1.
        min_width (float):
            The minimum size of the histogram window bounds, in weight units.
            Must be greater than 0.

    Returns:
        (nparray): A filtered copy of the input image "img", a 2-D array of floats.
    """

    # Converts an index tuple into a vector.
    def vec(x):
        return np.array(x)

    kernel_min = -kernel_center
    kernel_max = vec(kernel.shape) - kernel_center

    def weighted_histogram_filter_single(idx):
        idx = vec(idx)
        min_index = np.maximum(0, idx + kernel_min)
        max_index = np.minimum(vec(img.shape), idx + kernel_max)
        window_shape = max_index - min_index

        class WeightedElement:
            """
            An element of the histogram, its weight
            and bounds.
            """

            def __init__(self, value, weight):
                self.value: float = value
                self.weight: float = weight
                self.window_min: float = 0.0
                self.window_max: float = 1.0

        # Collect the values in the image as WeightedElements,
        # weighted by their corresponding kernel values.
        values = []
        for window_tup in np.ndindex(tuple(window_shape)):
            window_index = vec(window_tup)
            image_index = window_index + min_index
            centered_kernel_index = image_index - idx
            kernel_index = centered_kernel_index + kernel_center
            element = WeightedElement(img[tuple(image_index)], kernel[tuple(kernel_index)])
            values.append(element)

        def sort_key(x: WeightedElement):
            return x.value