Merge branch 'main' into fix/ui/fix-invoke-button-validation

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) {