feat: Inpaint & Outpaint Improvements (#4408)

## What type of PR is this? (check all applicable)

- [x] Feature
- [x] Optimization

## Have you discussed this change with the InvokeAI team?
- [x] Yes


## Description

# Coherence Mode

A new parameter called Coherence Mode has been added to Coherence Pass
settings. This parameter controls what kind of Coherence Pass is done
after Inpainting and Outpainting.

- Unmasked: This performs a complete unmasked image to image pass on the
entire generation.
- Mask: This performs a masked image to image pass using your input mask
as the coherence mask.
- Mask Edge [DEFAULT] - This performs as masked image to image pass on
the edges of your mask to try and clear out the seams.

# Why The Coherence Masked Modes?

One of the issues with unmasked coherence pass arises when the diffusion
process is trying to align detailed or organic objects. Because Image to
Image tends change the image a little bit even at lower strengths, this
ends up in the paste back process being slightly misaligned. By
providing the mask to the Coherence Pass, we can try to eliminate this
in those cases. While it will be impossible to address this for every
image out there, having these options will allow the user to automate a
lot of this. For everything else there's manual paint over with inpaint.

# Graph Improvements

The graphs have now been refined quite a bit. We no longer do manual
blurring of the masks anymore for outpainting. This is no longer needed
because we now dilate the mask depending on the blur size while pasting
back. As a result we got rid of quite a few nodes that were handling
this in the older graph.

The graphs are also a lot cleaner now because we now tackle Scaled
Dimensions & Coherence Mode completely independently.

Inpainting result seem very promising especially with the Mask Edge
mode.

---

# New Infill Methods [Experimental]

We are currently trying out various new infill methods to see which ones
might perform the best in outpainting. We may keep all of them or keep
none. This will be decided as we test more.

## LaMa Infill

- Renabled LaMA infill in the UI.
- We are trying to get this to work without a memory overhead.

In order to use LaMa, you need to manually download and place the LaMa
JIT model in `models/core/misc/lama/lama.pt`. You can download the JIT
model from Sanster
[here](https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt)
and rename it to `lama.pt` or you can use the script in the original
LaMA repo to convert the base model to a JIT model yourself.

## CV2 Infill

- Added a new infilling method using CV2's Inpaint.

## Patchmatch Rescaling

Patchmatch infill input image is now downscaled and infilled. Patchmatch
can be really slow at large resolutions and this is a pretty decent way
to get around that. Additionally, downscaling might also provide a
better patch match by avoiding larger areas to be infilled with
repeating patches. But that's just the theory. Still testing it out.

## [optional] Are there any post deployment tasks we need to perform?

- If we decide to keep LaMA infill, then we will need to host the model
and update the installer to download it as a core model.

Author: blessedcoolant

invokeai/app/api/routers/app_info.py

@@ -1,19 +1,19 @@
 import typing
 from enum import Enum
+from pathlib import Path
+
 from fastapi import Body
 from fastapi.routing import APIRouter
-from pathlib import Path
 from pydantic import BaseModel, Field
 
+from invokeai.app.invocations.upscale import ESRGAN_MODELS
+from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark
 from invokeai.backend.image_util.patchmatch import PatchMatch
 from invokeai.backend.image_util.safety_checker import SafetyChecker
-from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark
-from invokeai.app.invocations.upscale import ESRGAN_MODELS
-
+from invokeai.backend.util.logging import logging
 from invokeai.version import __version__
 
 from ..dependencies import ApiDependencies
-from invokeai.backend.util.logging import logging
 
 
 class LogLevel(int, Enum):
@@ -55,7 +55,7 @@ async def get_version() -> AppVersion:
 
 @app_router.get("/config", operation_id="get_config", status_code=200, response_model=AppConfig)
 async def get_config() -> AppConfig:
-    infill_methods = ["tile", "lama"]
+    infill_methods = ["tile", "lama", "cv2"]
     if PatchMatch.patchmatch_available():
         infill_methods.append("patchmatch")
 

invokeai/app/invocations/image.py

@@ -563,7 +563,7 @@ class MaskEdgeInvocation(BaseInvocation):
     )
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        mask = context.services.images.get_pil_image(self.image.image_name)
+        mask = context.services.images.get_pil_image(self.image.image_name).convert("L")
 
         npimg = numpy.asarray(mask, dtype=numpy.uint8)
         npgradient = numpy.uint8(255 * (1.0 - numpy.floor(numpy.abs(0.5 - numpy.float32(npimg) / 255.0) * 2.0)))
@@ -700,8 +700,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         # Blur the mask out (into init image) by specified amount
         if self.mask_blur_radius > 0:
             nm = numpy.asarray(pil_init_mask, dtype=numpy.uint8)
+            inverted_nm = 255 - nm
+            dilation_size = int(round(self.mask_blur_radius) + 20)
+            dilating_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_size, dilation_size))
+            inverted_dilated_nm = cv2.dilate(inverted_nm, dilating_kernel)
+            dilated_nm = 255 - inverted_dilated_nm
             nmd = cv2.erode(
-                nm,
+                dilated_nm,
                 kernel=numpy.ones((3, 3), dtype=numpy.uint8),
                 iterations=int(self.mask_blur_radius / 2),
             )

invokeai/app/invocations/infill.py

@@ -8,19 +8,17 @@
 
 from invokeai.app.invocations.primitives import ColorField, ImageField, ImageOutput
 from invokeai.app.util.misc import SEED_MAX, get_random_seed
+from invokeai.backend.image_util.cv2_inpaint import cv2_inpaint
 from invokeai.backend.image_util.lama import LaMA
 from invokeai.backend.image_util.patchmatch import PatchMatch
 
 from ..models.image import ImageCategory, ResourceOrigin
 from .baseinvocation import BaseInvocation, InputField, InvocationContext, invocation
+from .image import PIL_RESAMPLING_MAP, PIL_RESAMPLING_MODES
 
 
 def infill_methods() -> list[str]:
-    methods = [
-        "tile",
-        "solid",
-        "lama",
-    ]
+    methods = ["tile", "solid", "lama", "cv2"]
     if PatchMatch.patchmatch_available():
         methods.insert(0, "patchmatch")
     return methods
@@ -49,6 +47,10 @@ def infill_patchmatch(im: Image.Image) -> Image.Image:
     return im_patched
 
 
+def infill_cv2(im: Image.Image) -> Image.Image:
+    return cv2_inpaint(im)
+
+
 def get_tile_images(image: np.ndarray, width=8, height=8):
     _nrows, _ncols, depth = image.shape
     _strides = image.strides
@@ -194,15 +196,35 @@ class InfillPatchMatchInvocation(BaseInvocation):
     """Infills transparent areas of an image using the PatchMatch algorithm"""
 
     image: ImageField = InputField(description="The image to infill")
+    downscale: float = InputField(default=2.0, gt=0, description="Run patchmatch on downscaled image to speedup infill")
+    resample_mode: PIL_RESAMPLING_MODES = InputField(default="bicubic", description="The resampling mode")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get_pil_image(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name).convert("RGBA")
+
+        resample_mode = PIL_RESAMPLING_MAP[self.resample_mode]
+
+        infill_image = image.copy()
+        width = int(image.width / self.downscale)
+        height = int(image.height / self.downscale)
+        infill_image = infill_image.resize(
+            (width, height),
+            resample=resample_mode,
+        )
 
         if PatchMatch.patchmatch_available():
-            infilled = infill_patchmatch(image.copy())
+            infilled = infill_patchmatch(infill_image)
         else:
             raise ValueError("PatchMatch is not available on this system")
 
+        infilled = infilled.resize(
+            (image.width, image.height),
+            resample=resample_mode,
+        )
+
+        infilled.paste(image, (0, 0), mask=image.split()[-1])
+        # image.paste(infilled, (0, 0), mask=image.split()[-1])
+
         image_dto = context.services.images.create(
             image=infilled,
             image_origin=ResourceOrigin.INTERNAL,
@@ -245,3 +267,30 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             width=image_dto.width,
             height=image_dto.height,
         )
+
+
+@invocation("infill_cv2", title="CV2 Infill", tags=["image", "inpaint"], category="inpaint")
+class CV2InfillInvocation(BaseInvocation):
+    """Infills transparent areas of an image using OpenCV Inpainting"""
+
+    image: ImageField = InputField(description="The image to infill")
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.services.images.get_pil_image(self.image.image_name)
+
+        infilled = infill_cv2(image.copy())
+
+        image_dto = context.services.images.create(
+            image=infilled,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+        )
+
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )

invokeai/backend/image_util/cv2_inpaint.py

@@ -0,0 +1,20 @@
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def cv2_inpaint(image: Image.Image) -> Image.Image:
+    # Prepare Image
+    image_array = np.array(image.convert("RGB"))
+    image_cv = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
+
+    # Prepare Mask From Alpha Channel
+    mask = image.split()[3].convert("RGB")
+    mask_array = np.array(mask)
+    mask_cv = cv2.cvtColor(mask_array, cv2.COLOR_BGR2GRAY)
+    mask_inv = cv2.bitwise_not(mask_cv)
+
+    # Inpaint Image
+    inpainted_result = cv2.inpaint(image_cv, mask_inv, 3, cv2.INPAINT_TELEA)
+    inpainted_image = Image.fromarray(cv2.cvtColor(inpainted_result, cv2.COLOR_BGR2RGB))
+    return inpainted_image

invokeai/backend/image_util/lama.py

@@ -5,6 +5,7 @@
 import torch
 from PIL import Image
 
+import invokeai.backend.util.logging as logger
 from invokeai.app.services.config import get_invokeai_config
 from invokeai.backend.util.devices import choose_torch_device
 
@@ -19,7 +20,7 @@ def norm_img(np_img):
 
 def load_jit_model(url_or_path, device):
     model_path = url_or_path
-    print(f"Loading model from: {model_path}")
+    logger.info(f"Loading model from: {model_path}")
     model = torch.jit.load(model_path, map_location="cpu").to(device)
     model.eval()
     return model
@@ -52,5 +53,6 @@ def __call__(self, input_image: Image.Image, *args: Any, **kwds: Any) -> Any:
 
         del model
         gc.collect()
+        torch.cuda.empty_cache()
 
         return infilled_image

invokeai/backend/install/invokeai_configure.py

@@ -290,9 +290,20 @@ def download_realesrgan():
         download_with_progress_bar(model["url"], config.models_path / model["dest"], model["description"])
 
 
+# ---------------------------------------------
+def download_lama():
+    logger.info("Installing lama infill model")
+    download_with_progress_bar(
+        "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+        config.models_path / "core/misc/lama/lama.pt",
+        "lama infill model",
+    )
+
+
 # ---------------------------------------------
 def download_support_models():
     download_realesrgan()
+    download_lama()
     download_conversion_models()
 
 

invokeai/frontend/web/public/locales/en.json

@@ -511,6 +511,7 @@
         "maskBlur": "Blur",
         "maskBlurMethod": "Blur Method",
         "coherencePassHeader": "Coherence Pass",
+        "coherenceMode": "Mode",
         "coherenceSteps": "Steps",
         "coherenceStrength": "Strength",
         "seamLowThreshold": "Low",
@@ -520,6 +521,7 @@
         "scaledHeight": "Scaled H",
         "infillMethod": "Infill Method",
         "tileSize": "Tile Size",
+        "patchmatchDownScaleSize": "Downscale",
         "boundingBoxHeader": "Bounding Box",
         "seamCorrectionHeader": "Seam Correction",
         "infillScalingHeader": "Infill and Scaling",

invokeai/frontend/web/src/features/canvas/components/IAICanvasToolbar/IAICanvasToolChooserOptions.tsx

@@ -118,7 +118,11 @@ const IAICanvasToolChooserOptions = () => {
   useHotkeys(
     ['BracketLeft'],
     () => {
-      dispatch(setBrushSize(Math.max(brushSize - 5, 5)));
+      if (brushSize - 5 <= 5) {
+        dispatch(setBrushSize(Math.max(brushSize - 1, 1)));
+      } else {
+        dispatch(setBrushSize(Math.max(brushSize - 5, 1)));
+      }
     },
     {
       enabled: () => !isStaging,

invokeai/frontend/web/src/features/nodes/util/graphBuilders/addSDXLRefinerToGraph.ts

@@ -10,7 +10,8 @@ import {
   CANVAS_OUTPUT,
   INPAINT_IMAGE_RESIZE_UP,
   LATENTS_TO_IMAGE,
-  MASK_BLUR,
+  MASK_COMBINE,
+  MASK_RESIZE_UP,
   METADATA_ACCUMULATOR,
   SDXL_CANVAS_IMAGE_TO_IMAGE_GRAPH,
   SDXL_CANVAS_INPAINT_GRAPH,
@@ -46,6 +47,8 @@ export const addSDXLRefinerToGraph = (
   const { seamlessXAxis, seamlessYAxis, vaePrecision } = state.generation;
   const { boundingBoxScaleMethod } = state.canvas;
 
+  const fp32 = vaePrecision === 'fp32';
+
   const isUsingScaledDimensions = ['auto', 'manual'].includes(
     boundingBoxScaleMethod
   );
@@ -231,7 +234,7 @@ export const addSDXLRefinerToGraph = (
       type: 'create_denoise_mask',
       id: SDXL_REFINER_INPAINT_CREATE_MASK,
       is_intermediate: true,
-      fp32: vaePrecision === 'fp32' ? true : false,
+      fp32,
     };
 
     if (isUsingScaledDimensions) {
@@ -257,7 +260,7 @@ export const addSDXLRefinerToGraph = (
     graph.edges.push(
       {
         source: {
-          node_id: MASK_BLUR,
+          node_id: isUsingScaledDimensions ? MASK_RESIZE_UP : MASK_COMBINE,
           field: 'image',
         },
         destination: {

invokeai/frontend/web/src/features/nodes/util/graphBuilders/addVAEToGraph.ts

@@ -2,6 +2,7 @@ import { RootState } from 'app/store/store';
 import { NonNullableGraph } from 'features/nodes/types/types';
 import { MetadataAccumulatorInvocation } from 'services/api/types';
 import {
+  CANVAS_COHERENCE_INPAINT_CREATE_MASK,
   CANVAS_IMAGE_TO_IMAGE_GRAPH,
   CANVAS_INPAINT_GRAPH,
   CANVAS_OUTPAINT_GRAPH,
@@ -31,7 +32,7 @@ export const addVAEToGraph = (
   graph: NonNullableGraph,
   modelLoaderNodeId: string = MAIN_MODEL_LOADER
 ): void => {
-  const { vae } = state.generation;
+  const { vae, canvasCoherenceMode } = state.generation;
   const { boundingBoxScaleMethod } = state.canvas;
   const { shouldUseSDXLRefiner } = state.sdxl;
 
@@ -146,6 +147,20 @@ export const addVAEToGraph = (
         },
       }
     );
+
+    // Handle Coherence Mode
+    if (canvasCoherenceMode !== 'unmasked') {
+      graph.edges.push({
+        source: {
+          node_id: isAutoVae ? modelLoaderNodeId : VAE_LOADER,
+          field: isAutoVae && isOnnxModel ? 'vae_decoder' : 'vae',
+        },
+        destination: {
+          node_id: CANVAS_COHERENCE_INPAINT_CREATE_MASK,
+          field: 'vae',
+        },
+      });
+    }
   }
 
   if (shouldUseSDXLRefiner) {