Merge branch 'main' into pr/6086

invokeai/backend/image_util/__init__.py

@@ -2,7 +2,7 @@
 Initialization file for invokeai.backend.image_util methods.
 """
 
-from .patchmatch import PatchMatch  # noqa: F401
+from .infill_methods.patchmatch import PatchMatch  # noqa: F401
 from .pngwriter import PngWriter, PromptFormatter, retrieve_metadata, write_metadata  # noqa: F401
 from .seamless import configure_model_padding  # noqa: F401
 from .util import InitImageResizer, make_grid  # noqa: F401

invokeai/backend/image_util/depth_anything/__init__.py

@@ -13,7 +13,7 @@ from invokeai.app.services.config.config_default import get_config
 from invokeai.app.util.download_with_progress import download_with_progress_bar
 from invokeai.backend.image_util.depth_anything.model.dpt import DPT_DINOv2
 from invokeai.backend.image_util.depth_anything.utilities.util import NormalizeImage, PrepareForNet, Resize
-from invokeai.backend.util.devices import choose_torch_device
+from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 
 config = get_config()
@@ -56,7 +56,7 @@ class DepthAnythingDetector:
     def __init__(self) -> None:
         self.model = None
         self.model_size: Union[Literal["large", "base", "small"], None] = None
-        self.device = choose_torch_device()
+        self.device = TorchDevice.choose_torch_device()
 
     def load_model(self, model_size: Literal["large", "base", "small"] = "small"):
         DEPTH_ANYTHING_MODEL_PATH = config.models_path / DEPTH_ANYTHING_MODELS[model_size]["local"]
@@ -81,7 +81,7 @@ class DepthAnythingDetector:
             self.model.load_state_dict(torch.load(DEPTH_ANYTHING_MODEL_PATH.as_posix(), map_location="cpu"))
             self.model.eval()
 
-        self.model.to(choose_torch_device())
+        self.model.to(self.device)
         return self.model
 
     def __call__(self, image: Image.Image, resolution: int = 512) -> Image.Image:
@@ -94,7 +94,7 @@ class DepthAnythingDetector:
 
         image_height, image_width = np_image.shape[:2]
         np_image = transform({"image": np_image})["image"]
-        tensor_image = torch.from_numpy(np_image).unsqueeze(0).to(choose_torch_device())
+        tensor_image = torch.from_numpy(np_image).unsqueeze(0).to(self.device)
 
         with torch.no_grad():
             depth = self.model(tensor_image)

invokeai/backend/image_util/dw_openpose/wholebody.py

@@ -7,7 +7,7 @@ import onnxruntime as ort
 
 from invokeai.app.services.config.config_default import get_config
 from invokeai.app.util.download_with_progress import download_with_progress_bar
-from invokeai.backend.util.devices import choose_torch_device
+from invokeai.backend.util.devices import TorchDevice
 
 from .onnxdet import inference_detector
 from .onnxpose import inference_pose
@@ -28,9 +28,9 @@ config = get_config()
 
 class Wholebody:
     def __init__(self):
-        device = choose_torch_device()
+        device = TorchDevice.choose_torch_device()
 
-        providers = ["CUDAExecutionProvider"] if device == "cuda" else ["CPUExecutionProvider"]
+        providers = ["CUDAExecutionProvider"] if device.type == "cuda" else ["CPUExecutionProvider"]
 
         DET_MODEL_PATH = config.models_path / DWPOSE_MODELS["yolox_l.onnx"]["local"]
         download_with_progress_bar("yolox_l.onnx", DWPOSE_MODELS["yolox_l.onnx"]["url"], DET_MODEL_PATH)

invokeai/backend/image_util/hed.py

@@ -8,7 +8,7 @@ from huggingface_hub import hf_hub_download
 from PIL import Image
 
 from invokeai.backend.image_util.util import (
-    non_maximum_suppression,
+    nms,
     normalize_image_channel_count,
     np_to_pil,
     pil_to_np,
@@ -134,7 +134,7 @@ class HEDProcessor:
         detected_map = cv2.resize(detected_map, (width, height), interpolation=cv2.INTER_LINEAR)
 
         if scribble:
-            detected_map = non_maximum_suppression(detected_map, 127, 3.0)
+            detected_map = nms(detected_map, 127, 3.0)
             detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
             detected_map[detected_map > 4] = 255
             detected_map[detected_map < 255] = 0

invokeai/backend/image_util/infill_methods/lama.py

@@ -7,7 +7,8 @@ from PIL import Image
 
 import invokeai.backend.util.logging as logger
 from invokeai.app.services.config.config_default import get_config
-from invokeai.backend.util.devices import choose_torch_device
+from invokeai.app.util.download_with_progress import download_with_progress_bar
+from invokeai.backend.util.devices import TorchDevice
 
 
 def norm_img(np_img):
@@ -28,8 +29,16 @@ def load_jit_model(url_or_path, device):
 
 class LaMA:
     def __call__(self, input_image: Image.Image, *args: Any, **kwds: Any) -> Any:
-        device = choose_torch_device()
+        device = TorchDevice.choose_torch_device()
         model_location = get_config().models_path / "core/misc/lama/lama.pt"
+
+        if not model_location.exists():
+            download_with_progress_bar(
+                name="LaMa Inpainting Model",
+                url="https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+                dest_path=model_location,
+            )
+
         model = load_jit_model(model_location, device)
 
         image = np.asarray(input_image.convert("RGB"))

invokeai/backend/image_util/infill_methods/mosaic.py

@@ -0,0 +1,60 @@
+from typing import Tuple
+
+import numpy as np
+from PIL import Image
+
+
+def infill_mosaic(
+    image: Image.Image,
+    tile_shape: Tuple[int, int] = (64, 64),
+    min_color: Tuple[int, int, int, int] = (0, 0, 0, 0),
+    max_color: Tuple[int, int, int, int] = (255, 255, 255, 0),
+) -> Image.Image:
+    """
+    image:PIL - A PIL Image
+    tile_shape: Tuple[int,int] - Tile width & Tile Height
+    min_color: Tuple[int,int,int] - RGB values for the lowest color to clip to (0-255)
+    max_color: Tuple[int,int,int] - RGB values for the highest color to clip to (0-255)
+    """
+
+    np_image = np.array(image)  # Convert image to np array
+    alpha = np_image[:, :, 3]  # Get the mask from the alpha channel of the image
+    non_transparent_pixels = np_image[alpha != 0, :3]  # List of non-transparent pixels
+
+    # Create color tiles to paste in the empty areas of the image
+    tile_width, tile_height = tile_shape
+
+    # Clip the range of colors in the image to a particular spectrum only
+    r_min, g_min, b_min, _ = min_color
+    r_max, g_max, b_max, _ = max_color
+    non_transparent_pixels[:, 0] = np.clip(non_transparent_pixels[:, 0], r_min, r_max)
+    non_transparent_pixels[:, 1] = np.clip(non_transparent_pixels[:, 1], g_min, g_max)
+    non_transparent_pixels[:, 2] = np.clip(non_transparent_pixels[:, 2], b_min, b_max)
+
+    tiles = []
+    for _ in range(256):
+        color = non_transparent_pixels[np.random.randint(len(non_transparent_pixels))]
+        tile = np.zeros((tile_height, tile_width, 3), dtype=np.uint8)
+        tile[:, :] = color
+        tiles.append(tile)
+
+    # Fill the transparent area with tiles
+    filled_image = np.zeros((image.height, image.width, 3), dtype=np.uint8)
+
+    for x in range(image.width):
+        for y in range(image.height):
+            tile = tiles[np.random.randint(len(tiles))]
+            try:
+                filled_image[
+                    y - (y % tile_height) : y - (y % tile_height) + tile_height,
+                    x - (x % tile_width) : x - (x % tile_width) + tile_width,
+                ] = tile
+            except ValueError:
+                # Need to handle edge cases - literally
+                pass
+
+    filled_image = Image.fromarray(filled_image)  # Convert the filled tiles image to PIL
+    image = Image.composite(
+        image, filled_image, image.split()[-1]
+    )  # Composite the original image on top of the filled tiles
+    return image

invokeai/backend/image_util/infill_methods/patchmatch.py

@@ -0,0 +1,67 @@
+"""
+This module defines a singleton object, "patchmatch" that
+wraps the actual patchmatch object. It respects the global
+"try_patchmatch" attribute, so that patchmatch loading can
+be suppressed or deferred
+"""
+
+import numpy as np
+from PIL import Image
+
+import invokeai.backend.util.logging as logger
+from invokeai.app.services.config.config_default import get_config
+
+
+class PatchMatch:
+    """
+    Thin class wrapper around the patchmatch function.
+    """
+
+    patch_match = None
+    tried_load: bool = False
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def _load_patch_match(cls):
+        if cls.tried_load:
+            return
+        if get_config().patchmatch:
+            from patchmatch import patch_match as pm
+
+            if pm.patchmatch_available:
+                logger.info("Patchmatch initialized")
+                cls.patch_match = pm
+            else:
+                logger.info("Patchmatch not loaded (nonfatal)")
+        else:
+            logger.info("Patchmatch loading disabled")
+        cls.tried_load = True
+
+    @classmethod
+    def patchmatch_available(cls) -> bool:
+        cls._load_patch_match()
+        if not cls.patch_match:
+            return False
+        return cls.patch_match.patchmatch_available
+
+    @classmethod
+    def inpaint(cls, image: Image.Image) -> Image.Image:
+        if cls.patch_match is None or not cls.patchmatch_available():
+            return image
+
+        np_image = np.array(image)
+        mask = 255 - np_image[:, :, 3]
+        infilled = cls.patch_match.inpaint(np_image[:, :, :3], mask, patch_size=3)
+        return Image.fromarray(infilled, mode="RGB")
+
+
+def infill_patchmatch(image: Image.Image) -> Image.Image:
+    IS_PATCHMATCH_AVAILABLE = PatchMatch.patchmatch_available()
+
+    if not IS_PATCHMATCH_AVAILABLE:
+        logger.warning("PatchMatch is not available on this system")
+        return image
+
+    return PatchMatch.inpaint(image)

invokeai/backend/image_util/infill_methods/tile.ipynb

@@ -0,0 +1,95 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Smoke test for the tile infill\"\"\"\n",
+    "\n",
+    "from pathlib import Path\n",
+    "from typing import Optional\n",
+    "from PIL import Image\n",
+    "from invokeai.backend.image_util.infill_methods.tile import infill_tile\n",
+    "\n",
+    "images: list[tuple[str, Image.Image]] = []\n",
+    "\n",
+    "for i in sorted(Path(\"./test_images/\").glob(\"*.webp\")):\n",
+    "    images.append((i.name, Image.open(i)))\n",
+    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_LEFT_RIGHT)))\n",
+    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_TOP_BOTTOM)))\n",
+    "    images.append((i.name, Image.open(i).resize((512, 512))))\n",
+    "    images.append((i.name, Image.open(i).resize((1234, 461))))\n",
+    "\n",
+    "outputs: list[tuple[str, Image.Image, Image.Image, Optional[Image.Image]]] = []\n",
+    "\n",
+    "for name, image in images:\n",
+    "    try:\n",
+    "        output = infill_tile(image, seed=0, tile_size=32)\n",
+    "        outputs.append((name, image, output.infilled, output.tile_image))\n",
+    "    except ValueError as e:\n",
+    "        print(f\"Skipping image {name}: {e}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Display the images in jupyter notebook\n",
+    "import matplotlib.pyplot as plt\n",
+    "from PIL import ImageOps\n",
+    "\n",
+    "fig, axes = plt.subplots(len(outputs), 3, figsize=(10, 3 * len(outputs)))\n",
+    "plt.subplots_adjust(hspace=0)\n",
+    "\n",
+    "for i, (name, original, infilled, tile_image) in enumerate(outputs):\n",
+    "    # Add a border to each image, helps to see the edges\n",
+    "    size = original.size\n",
+    "    original = ImageOps.expand(original, border=5, fill=\"red\")\n",
+    "    filled = ImageOps.expand(infilled, border=5, fill=\"red\")\n",
+    "    if tile_image:\n",
+    "        tile_image = ImageOps.expand(tile_image, border=5, fill=\"red\")\n",
+    "\n",
+    "    axes[i, 0].imshow(original)\n",
+    "    axes[i, 0].axis(\"off\")\n",
+    "    axes[i, 0].set_title(f\"Original ({name} - {size})\")\n",
+    "\n",
+    "    if tile_image:\n",
+    "        axes[i, 1].imshow(tile_image)\n",
+    "        axes[i, 1].axis(\"off\")\n",
+    "        axes[i, 1].set_title(\"Tile Image\")\n",
+    "    else:\n",
+    "        axes[i, 1].axis(\"off\")\n",
+    "        axes[i, 1].set_title(\"NO TILES GENERATED (NO TRANSPARENCY)\")\n",
+    "\n",
+    "    axes[i, 2].imshow(filled)\n",
+    "    axes[i, 2].axis(\"off\")\n",
+    "    axes[i, 2].set_title(\"Filled\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".invokeai",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

invokeai/backend/image_util/infill_methods/tile.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from PIL import Image
+
+
+def create_tile_pool(img_array: np.ndarray, tile_size: tuple[int, int]) -> list[np.ndarray]:
+    """
+    Create a pool of tiles from non-transparent areas of the image by systematically walking through the image.
+
+    Args:
+        img_array: numpy array of the image.
+        tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
+
+    Returns:
+        A list of numpy arrays, each representing a tile.
+    """
+    tiles: list[np.ndarray] = []
+    rows, cols = img_array.shape[:2]
+    tile_width, tile_height = tile_size
+
+    for y in range(0, rows - tile_height + 1, tile_height):
+        for x in range(0, cols - tile_width + 1, tile_width):
+            tile = img_array[y : y + tile_height, x : x + tile_width]
+            # Check if the image has an alpha channel and the tile is completely opaque
+            if img_array.shape[2] == 4 and np.all(tile[:, :, 3] == 255):
+                tiles.append(tile)
+            elif img_array.shape[2] == 3:  # If no alpha channel, append the tile
+                tiles.append(tile)
+
+    if not tiles:
+        raise ValueError(
+            "Not enough opaque pixels to generate any tiles. Use a smaller tile size or a different image."
+        )
+
+    return tiles
+
+
+def create_filled_image(
+    img_array: np.ndarray, tile_pool: list[np.ndarray], tile_size: tuple[int, int], seed: int
+) -> np.ndarray:
+    """
+    Create an image of the same dimensions as the original, filled entirely with tiles from the pool.
+
+    Args:
+        img_array: numpy array of the original image.
+        tile_pool: A list of numpy arrays, each representing a tile.
+        tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
+
+    Returns:
+        A numpy array representing the filled image.
+    """
+
+    rows, cols, _ = img_array.shape
+    tile_width, tile_height = tile_size
+
+    # Prep an empty RGB image
+    filled_img_array = np.zeros((rows, cols, 3), dtype=img_array.dtype)
+
+    # Make the random tile selection reproducible
+    rng = np.random.default_rng(seed)
+
+    for y in range(0, rows, tile_height):
+        for x in range(0, cols, tile_width):
+            # Pick a random tile from the pool
+            tile = tile_pool[rng.integers(len(tile_pool))]
+
+            # Calculate the space available (may be less than tile size near the edges)
+            space_y = min(tile_height, rows - y)
+            space_x = min(tile_width, cols - x)
+
+            # Crop the tile if necessary to fit into the available space
+            cropped_tile = tile[:space_y, :space_x, :3]
+
+            # Fill the available space with the (possibly cropped) tile
+            filled_img_array[y : y + space_y, x : x + space_x, :3] = cropped_tile
+
+    return filled_img_array
+
+
+@dataclass
+class InfillTileOutput:
+    infilled: Image.Image
+    tile_image: Optional[Image.Image] = None
+
+
+def infill_tile(image_to_infill: Image.Image, seed: int, tile_size: int) -> InfillTileOutput:
+    """Infills an image with random tiles from the image itself.
+
+    If the image is not an RGBA image, it is returned untouched.
+
+    Args:
+        image: The image to infill.
+        tile_size: The size of the tiles to use for infilling.
+
+    Raises:
+        ValueError: If there are not enough opaque pixels to generate any tiles.
+    """
+
+    if image_to_infill.mode != "RGBA":
+        return InfillTileOutput(infilled=image_to_infill)
+
+    # Internally, we want a tuple of (tile_width, tile_height). In the future, the tile size can be any rectangle.
+    _tile_size = (tile_size, tile_size)
+    np_image = np.array(image_to_infill, dtype=np.uint8)
+
+    # Create the pool of tiles that we will use to infill
+    tile_pool = create_tile_pool(np_image, _tile_size)
+
+    # Create an image from the tiles, same size as the original
+    tile_np_image = create_filled_image(np_image, tile_pool, _tile_size, seed)
+
+    # Paste the OG image over the tile image, effectively infilling the area
+    tile_image = Image.fromarray(tile_np_image, "RGB")
+    infilled = tile_image.copy()
+    infilled.paste(image_to_infill, (0, 0), image_to_infill.split()[-1])
+
+    # I think we want this to be "RGBA"?
+    infilled.convert("RGBA")
+
+    return InfillTileOutput(infilled=infilled, tile_image=tile_image)

invokeai/backend/image_util/patchmatch.py

@@ -1,49 +0,0 @@
-"""
-This module defines a singleton object, "patchmatch" that
-wraps the actual patchmatch object. It respects the global
-"try_patchmatch" attribute, so that patchmatch loading can
-be suppressed or deferred
-"""
-
-import numpy as np
-
-import invokeai.backend.util.logging as logger
-from invokeai.app.services.config.config_default import get_config
-
-
-class PatchMatch:
-    """
-    Thin class wrapper around the patchmatch function.
-    """
-
-    patch_match = None
-    tried_load: bool = False
-
-    def __init__(self):
-        super().__init__()
-
-    @classmethod
-    def _load_patch_match(self):
-        if self.tried_load:
-            return
-        if get_config().patchmatch:
-            from patchmatch import patch_match as pm
-
-            if pm.patchmatch_available:
-                logger.info("Patchmatch initialized")
-            else:
-                logger.info("Patchmatch not loaded (nonfatal)")
-            self.patch_match = pm
-        else:
-            logger.info("Patchmatch loading disabled")
-        self.tried_load = True
-
-    @classmethod
-    def patchmatch_available(self) -> bool:
-        self._load_patch_match()
-        return self.patch_match and self.patch_match.patchmatch_available
-
-    @classmethod
-    def inpaint(self, *args, **kwargs) -> np.ndarray:
-        if self.patchmatch_available():
-            return self.patch_match.inpaint(*args, **kwargs)

invokeai/backend/image_util/realesrgan/realesrgan.py

@@ -11,7 +11,7 @@ from cv2.typing import MatLike
 from tqdm import tqdm
 
 from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet
-from invokeai.backend.util.devices import choose_torch_device
+from invokeai.backend.util.devices import TorchDevice
 
 """
 Adapted from https://github.com/xinntao/Real-ESRGAN/blob/master/realesrgan/utils.py
@@ -65,7 +65,7 @@ class RealESRGAN:
         self.pre_pad = pre_pad
         self.mod_scale: Optional[int] = None
         self.half = half
-        self.device = choose_torch_device()
+        self.device = TorchDevice.choose_torch_device()
 
         loadnet = torch.load(model_path, map_location=torch.device("cpu"))
 

invokeai/backend/image_util/safety_checker.py

@@ -13,7 +13,7 @@ from transformers import AutoFeatureExtractor
 
 import invokeai.backend.util.logging as logger
 from invokeai.app.services.config.config_default import get_config
-from invokeai.backend.util.devices import choose_torch_device
+from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.silence_warnings import SilenceWarnings
 
 CHECKER_PATH = "core/convert/stable-diffusion-safety-checker"
@@ -51,7 +51,7 @@ class SafetyChecker:
         cls._load_safety_checker()
         if cls.safety_checker is None or cls.feature_extractor is None:
             return False
-        device = choose_torch_device()
+        device = TorchDevice.choose_torch_device()
         features = cls.feature_extractor([image], return_tensors="pt")
         features.to(device)
         cls.safety_checker.to(device)

invokeai/backend/image_util/util.py

@@ -1,4 +1,5 @@
 from math import ceil, floor, sqrt
+from typing import Optional
 
 import cv2
 import numpy as np
@@ -143,20 +144,21 @@ def resize_image_to_resolution(input_image: np.ndarray, resolution: int) -> np.n
     h = float(input_image.shape[0])
     w = float(input_image.shape[1])
     scaling_factor = float(resolution) / min(h, w)
-    h *= scaling_factor
-    w *= scaling_factor
-    h = int(np.round(h / 64.0)) * 64
-    w = int(np.round(w / 64.0)) * 64
+    h = int(h * scaling_factor)
+    w = int(w * scaling_factor)
     if scaling_factor > 1:
         return cv2.resize(input_image, (w, h), interpolation=cv2.INTER_LANCZOS4)
     else:
         return cv2.resize(input_image, (w, h), interpolation=cv2.INTER_AREA)
 
 
-def non_maximum_suppression(image: np.ndarray, threshold: int, sigma: float):
+def nms(np_img: np.ndarray, threshold: Optional[int] = None, sigma: Optional[float] = None) -> np.ndarray:
     """
     Apply non-maximum suppression to an image.
 
+    If both threshold and sigma are provided, the image will blurred before the suppression and thresholded afterwards,
+    resulting in a binary output image.
+
     This function is adapted from https://github.com/lllyasviel/ControlNet.
 
     Args:
@@ -166,23 +168,36 @@ def non_maximum_suppression(image: np.ndarray, threshold: int, sigma: float):
 
     Returns:
         The image after non-maximum suppression.
+
+    Raises:
+        ValueError: If only one of threshold and sigma provided.
     """
 
-    image = cv2.GaussianBlur(image.astype(np.float32), (0, 0), sigma)
+    # Raise a value error if only one of threshold and sigma is provided
+    if (threshold is None) != (sigma is None):
+        raise ValueError("Both threshold and sigma must be provided if one is provided.")
+
+    if sigma is not None and threshold is not None:
+        # Blurring the image can help to thin out features
+        np_img = cv2.GaussianBlur(np_img.astype(np.float32), (0, 0), sigma)
 
     filter_1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
     filter_2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
     filter_3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
     filter_4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
 
-    y = np.zeros_like(image)
+    nms_img = np.zeros_like(np_img)
 
     for f in [filter_1, filter_2, filter_3, filter_4]:
-        np.putmask(y, cv2.dilate(image, kernel=f) == image, image)
+        np.putmask(nms_img, cv2.dilate(np_img, kernel=f) == np_img, np_img)
 
-    z = np.zeros_like(y, dtype=np.uint8)
-    z[y > threshold] = 255
-    return z
+    if sigma is not None and threshold is not None:
+        # We blurred - now threshold to get a binary image
+        thresholded = np.zeros_like(nms_img, dtype=np.uint8)
+        thresholded[nms_img > threshold] = 255
+        return thresholded
+
+    return nms_img
 
 
 def safe_step(x: np.ndarray, step: int = 2) -> np.ndarray:

invokeai/backend/ip_adapter/attention_processor.py

@@ -1,182 +0,0 @@
-# copied from https://github.com/tencent-ailab/IP-Adapter (Apache License 2.0)
-#   and modified as needed
-
-# tencent-ailab comment:
-# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.models.attention_processor import AttnProcessor2_0 as DiffusersAttnProcessor2_0
-
-from invokeai.backend.ip_adapter.ip_attention_weights import IPAttentionProcessorWeights
-
-
-# Create a version of AttnProcessor2_0 that is a sub-class of nn.Module. This is required for IP-Adapter state_dict
-# loading.
-class AttnProcessor2_0(DiffusersAttnProcessor2_0, nn.Module):
-    def __init__(self):
-        DiffusersAttnProcessor2_0.__init__(self)
-        nn.Module.__init__(self)
-
-    def __call__(
-        self,
-        attn,
-        hidden_states,
-        encoder_hidden_states=None,
-        attention_mask=None,
-        temb=None,
-        ip_adapter_image_prompt_embeds=None,
-    ):
-        """Re-definition of DiffusersAttnProcessor2_0.__call__(...) that accepts and ignores the
-        ip_adapter_image_prompt_embeds parameter.
-        """
-        return DiffusersAttnProcessor2_0.__call__(
-            self, attn, hidden_states, encoder_hidden_states, attention_mask, temb
-        )
-
-
-class IPAttnProcessor2_0(torch.nn.Module):
-    r"""
-    Attention processor for IP-Adapater for PyTorch 2.0.
-    Args:
-        hidden_size (`int`):
-            The hidden size of the attention layer.
-        cross_attention_dim (`int`):
-            The number of channels in the `encoder_hidden_states`.
-        scale (`float`, defaults to 1.0):
-            the weight scale of image prompt.
-    """
-
-    def __init__(self, weights: list[IPAttentionProcessorWeights], scales: list[float]):
-        super().__init__()
-
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-
-        assert len(weights) == len(scales)
-
-        self._weights = weights
-        self._scales = scales
-
-    def __call__(
-        self,
-        attn,
-        hidden_states,
-        encoder_hidden_states=None,
-        attention_mask=None,
-        temb=None,
-        ip_adapter_image_prompt_embeds=None,
-    ):
-        """Apply IP-Adapter attention.
-
-        Args:
-            ip_adapter_image_prompt_embeds (torch.Tensor): The image prompt embeddings.
-                Shape: (batch_size, num_ip_images, seq_len, ip_embedding_len).
-        """
-        residual = hidden_states
-
-        if attn.spatial_norm is not None:
-            hidden_states = attn.spatial_norm(hidden_states, temb)
-
-        input_ndim = hidden_states.ndim
-
-        if input_ndim == 4:
-            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-
-        batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
-        )
-
-        if attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-            # scaled_dot_product_attention expects attention_mask shape to be
-            # (batch, heads, source_length, target_length)
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
-
-        if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-
-        query = attn.to_q(hidden_states)
-
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.norm_cross:
-            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-
-        key = attn.to_k(encoder_hidden_states)
-        value = attn.to_v(encoder_hidden_states)
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        # the output of sdp = (batch, num_heads, seq_len, head_dim)
-        # TODO: add support for attn.scale when we move to Torch 2.1
-        hidden_states = F.scaled_dot_product_attention(
-            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-        )
-
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-        hidden_states = hidden_states.to(query.dtype)
-
-        if encoder_hidden_states is not None:
-            # If encoder_hidden_states is not None, then we are doing cross-attention, not self-attention. In this case,
-            # we will apply IP-Adapter conditioning. We validate the inputs for IP-Adapter conditioning here.
-            assert ip_adapter_image_prompt_embeds is not None
-            assert len(ip_adapter_image_prompt_embeds) == len(self._weights)
-
-            for ipa_embed, ipa_weights, scale in zip(
-                ip_adapter_image_prompt_embeds, self._weights, self._scales, strict=True
-            ):
-                # The batch dimensions should match.
-                assert ipa_embed.shape[0] == encoder_hidden_states.shape[0]
-                # The token_len dimensions should match.
-                assert ipa_embed.shape[-1] == encoder_hidden_states.shape[-1]
-
-                ip_hidden_states = ipa_embed
-
-                # Expected ip_hidden_state shape: (batch_size, num_ip_images, ip_seq_len, ip_image_embedding)
-
-                ip_key = ipa_weights.to_k_ip(ip_hidden_states)
-                ip_value = ipa_weights.to_v_ip(ip_hidden_states)
-
-                # Expected ip_key and ip_value shape: (batch_size, num_ip_images, ip_seq_len, head_dim * num_heads)
-
-                ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-                ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-                # Expected ip_key and ip_value shape: (batch_size, num_heads, num_ip_images * ip_seq_len, head_dim)
-
-                # TODO: add support for attn.scale when we move to Torch 2.1
-                ip_hidden_states = F.scaled_dot_product_attention(
-                    query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
-                )
-
-                # Expected ip_hidden_states shape: (batch_size, num_heads, query_seq_len, head_dim)
-
-                ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-                ip_hidden_states = ip_hidden_states.to(query.dtype)
-
-                # Expected ip_hidden_states shape: (batch_size, query_seq_len, num_heads * head_dim)
-
-                hidden_states = hidden_states + scale * ip_hidden_states
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-
-        if attn.residual_connection:
-            hidden_states = hidden_states + residual
-
-        hidden_states = hidden_states / attn.rescale_output_factor
-
-        return hidden_states

invokeai/backend/ip_adapter/ip_adapter.py

@@ -1,8 +1,11 @@
 # copied from https://github.com/tencent-ailab/IP-Adapter (Apache License 2.0)
 #   and modified as needed
 
-from typing import Optional, Union
+import pathlib
+from typing import List, Optional, TypedDict, Union
 
+import safetensors
+import safetensors.torch
 import torch
 from PIL import Image
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
@@ -13,10 +16,17 @@ from ..raw_model import RawModel
 from .resampler import Resampler
 
 
+class IPAdapterStateDict(TypedDict):
+    ip_adapter: dict[str, torch.Tensor]
+    image_proj: dict[str, torch.Tensor]
+
+
 class ImageProjModel(torch.nn.Module):
     """Image Projection Model"""
 
-    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
+    def __init__(
+        self, cross_attention_dim: int = 1024, clip_embeddings_dim: int = 1024, clip_extra_context_tokens: int = 4
+    ):
         super().__init__()
 
         self.cross_attention_dim = cross_attention_dim
@@ -25,7 +35,7 @@ class ImageProjModel(torch.nn.Module):
         self.norm = torch.nn.LayerNorm(cross_attention_dim)
 
     @classmethod
-    def from_state_dict(cls, state_dict: dict[torch.Tensor], clip_extra_context_tokens=4):
+    def from_state_dict(cls, state_dict: dict[str, torch.Tensor], clip_extra_context_tokens: int = 4):
         """Initialize an ImageProjModel from a state_dict.
 
         The cross_attention_dim and clip_embeddings_dim are inferred from the shape of the tensors in the state_dict.
@@ -45,7 +55,7 @@ class ImageProjModel(torch.nn.Module):
         model.load_state_dict(state_dict)
         return model
 
-    def forward(self, image_embeds):
+    def forward(self, image_embeds: torch.Tensor):
         embeds = image_embeds
         clip_extra_context_tokens = self.proj(embeds).reshape(
             -1, self.clip_extra_context_tokens, self.cross_attention_dim
@@ -57,7 +67,7 @@ class ImageProjModel(torch.nn.Module):
 class MLPProjModel(torch.nn.Module):
     """SD model with image prompt"""
 
-    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024):
+    def __init__(self, cross_attention_dim: int = 1024, clip_embeddings_dim: int = 1024):
         super().__init__()
 
         self.proj = torch.nn.Sequential(
@@ -68,7 +78,7 @@ class MLPProjModel(torch.nn.Module):
         )
 
     @classmethod
-    def from_state_dict(cls, state_dict: dict[torch.Tensor]):
+    def from_state_dict(cls, state_dict: dict[str, torch.Tensor]):
         """Initialize an MLPProjModel from a state_dict.
 
         The cross_attention_dim and clip_embeddings_dim are inferred from the shape of the tensors in the state_dict.
@@ -87,7 +97,7 @@ class MLPProjModel(torch.nn.Module):
         model.load_state_dict(state_dict)
         return model
 
-    def forward(self, image_embeds):
+    def forward(self, image_embeds: torch.Tensor):
         clip_extra_context_tokens = self.proj(image_embeds)
         return clip_extra_context_tokens
 
@@ -97,7 +107,7 @@ class IPAdapter(RawModel):
 
     def __init__(
         self,
-        state_dict: dict[str, torch.Tensor],
+        state_dict: IPAdapterStateDict,
         device: torch.device,
         dtype: torch.dtype = torch.float16,
         num_tokens: int = 4,
@@ -129,24 +139,27 @@ class IPAdapter(RawModel):
 
         return calc_model_size_by_data(self._image_proj_model) + calc_model_size_by_data(self.attn_weights)
 
-    def _init_image_proj_model(self, state_dict):
+    def _init_image_proj_model(
+        self, state_dict: dict[str, torch.Tensor]
+    ) -> Union[ImageProjModel, Resampler, MLPProjModel]:
         return ImageProjModel.from_state_dict(state_dict, self._num_tokens).to(self.device, dtype=self.dtype)
 
     @torch.inference_mode()
-    def get_image_embeds(self, pil_image, image_encoder: CLIPVisionModelWithProjection):
-        if isinstance(pil_image, Image.Image):
-            pil_image = [pil_image]
+    def get_image_embeds(self, pil_image: List[Image.Image], image_encoder: CLIPVisionModelWithProjection):
         clip_image = self._clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
         clip_image_embeds = image_encoder(clip_image.to(self.device, dtype=self.dtype)).image_embeds
-        image_prompt_embeds = self._image_proj_model(clip_image_embeds)
-        uncond_image_prompt_embeds = self._image_proj_model(torch.zeros_like(clip_image_embeds))
-        return image_prompt_embeds, uncond_image_prompt_embeds
+        try:
+            image_prompt_embeds = self._image_proj_model(clip_image_embeds)
+            uncond_image_prompt_embeds = self._image_proj_model(torch.zeros_like(clip_image_embeds))
+            return image_prompt_embeds, uncond_image_prompt_embeds
+        except RuntimeError as e:
+            raise RuntimeError("Selected CLIP Vision Model is incompatible with the current IP Adapter") from e
 
 
 class IPAdapterPlus(IPAdapter):
     """IP-Adapter with fine-grained features"""
 
-    def _init_image_proj_model(self, state_dict):
+    def _init_image_proj_model(self, state_dict: dict[str, torch.Tensor]) -> Union[Resampler, MLPProjModel]:
         return Resampler.from_state_dict(
             state_dict=state_dict,
             depth=4,
@@ -157,31 +170,32 @@ class IPAdapterPlus(IPAdapter):
         ).to(self.device, dtype=self.dtype)
 
     @torch.inference_mode()
-    def get_image_embeds(self, pil_image, image_encoder: CLIPVisionModelWithProjection):
-        if isinstance(pil_image, Image.Image):
-            pil_image = [pil_image]
+    def get_image_embeds(self, pil_image: List[Image.Image], image_encoder: CLIPVisionModelWithProjection):
         clip_image = self._clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
         clip_image = clip_image.to(self.device, dtype=self.dtype)
         clip_image_embeds = image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
-        image_prompt_embeds = self._image_proj_model(clip_image_embeds)
         uncond_clip_image_embeds = image_encoder(torch.zeros_like(clip_image), output_hidden_states=True).hidden_states[
             -2
         ]
-        uncond_image_prompt_embeds = self._image_proj_model(uncond_clip_image_embeds)
-        return image_prompt_embeds, uncond_image_prompt_embeds
+        try:
+            image_prompt_embeds = self._image_proj_model(clip_image_embeds)
+            uncond_image_prompt_embeds = self._image_proj_model(uncond_clip_image_embeds)
+            return image_prompt_embeds, uncond_image_prompt_embeds
+        except RuntimeError as e:
+            raise RuntimeError("Selected CLIP Vision Model is incompatible with the current IP Adapter") from e
 
 
 class IPAdapterFull(IPAdapterPlus):
     """IP-Adapter Plus with full features."""
 
-    def _init_image_proj_model(self, state_dict: dict[torch.Tensor]):
+    def _init_image_proj_model(self, state_dict: dict[str, torch.Tensor]):
         return MLPProjModel.from_state_dict(state_dict).to(self.device, dtype=self.dtype)
 
 
 class IPAdapterPlusXL(IPAdapterPlus):
     """IP-Adapter Plus for SDXL."""
 
-    def _init_image_proj_model(self, state_dict):
+    def _init_image_proj_model(self, state_dict: dict[str, torch.Tensor]):
         return Resampler.from_state_dict(
             state_dict=state_dict,
             depth=4,
@@ -192,24 +206,48 @@ class IPAdapterPlusXL(IPAdapterPlus):
         ).to(self.device, dtype=self.dtype)
 
 
-def build_ip_adapter(
-    ip_adapter_ckpt_path: str, device: torch.device, dtype: torch.dtype = torch.float16
-) -> Union[IPAdapter, IPAdapterPlus]:
-    state_dict = torch.load(ip_adapter_ckpt_path, map_location="cpu")
+def load_ip_adapter_tensors(ip_adapter_ckpt_path: pathlib.Path, device: str) -> IPAdapterStateDict:
+    state_dict: IPAdapterStateDict = {"ip_adapter": {}, "image_proj": {}}
 
-    if "proj.weight" in state_dict["image_proj"]:  # IPAdapter (with ImageProjModel).
+    if ip_adapter_ckpt_path.suffix == ".safetensors":
+        model = safetensors.torch.load_file(ip_adapter_ckpt_path, device=device)
+        for key in model.keys():
+            if key.startswith("image_proj."):
+                state_dict["image_proj"][key.replace("image_proj.", "")] = model[key]
+            elif key.startswith("ip_adapter."):
+                state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = model[key]
+            else:
+                raise RuntimeError(f"Encountered unexpected IP Adapter state dict key: '{key}'.")
+    else:
+        ip_adapter_diffusers_checkpoint_path = ip_adapter_ckpt_path / "ip_adapter.bin"
+        state_dict = torch.load(ip_adapter_diffusers_checkpoint_path, map_location="cpu")
+
+    return state_dict
+
+
+def build_ip_adapter(
+    ip_adapter_ckpt_path: pathlib.Path, device: torch.device, dtype: torch.dtype = torch.float16
+) -> Union[IPAdapter, IPAdapterPlus, IPAdapterPlusXL, IPAdapterPlus]:
+    state_dict = load_ip_adapter_tensors(ip_adapter_ckpt_path, device.type)
+
+    # IPAdapter (with ImageProjModel)
+    if "proj.weight" in state_dict["image_proj"]:
         return IPAdapter(state_dict, device=device, dtype=dtype)
-    elif "proj_in.weight" in state_dict["image_proj"]:  # IPAdaterPlus or IPAdapterPlusXL (with Resampler).
+
+    # IPAdaterPlus or IPAdapterPlusXL (with Resampler)
+    elif "proj_in.weight" in state_dict["image_proj"]:
         cross_attention_dim = state_dict["ip_adapter"]["1.to_k_ip.weight"].shape[-1]
         if cross_attention_dim == 768:
-            # SD1 IP-Adapter Plus
-            return IPAdapterPlus(state_dict, device=device, dtype=dtype)
+            return IPAdapterPlus(state_dict, device=device, dtype=dtype)  # SD1 IP-Adapter Plus
         elif cross_attention_dim == 2048:
-            # SDXL IP-Adapter Plus
-            return IPAdapterPlusXL(state_dict, device=device, dtype=dtype)
+            return IPAdapterPlusXL(state_dict, device=device, dtype=dtype)  # SDXL IP-Adapter Plus
         else:
             raise Exception(f"Unsupported IP-Adapter Plus cross-attention dimension: {cross_attention_dim}.")
-    elif "proj.0.weight" in state_dict["image_proj"]:  # IPAdapterFull (with MLPProjModel).
+
+    # IPAdapterFull (with MLPProjModel)
+    elif "proj.0.weight" in state_dict["image_proj"]:
         return IPAdapterFull(state_dict, device=device, dtype=dtype)
+
+    # Unrecognized IP Adapter Architectures
     else:
         raise ValueError(f"'{ip_adapter_ckpt_path}' has an unrecognized IP-Adapter model architecture.")

invokeai/backend/ip_adapter/resampler.py

@@ -9,8 +9,8 @@ import torch.nn as nn
 
 
 # FFN
-def FeedForward(dim, mult=4):
-    inner_dim = int(dim * mult)
+def FeedForward(dim: int, mult: int = 4):
+    inner_dim = dim * mult
     return nn.Sequential(
         nn.LayerNorm(dim),
         nn.Linear(dim, inner_dim, bias=False),
@@ -19,8 +19,8 @@ def FeedForward(dim, mult=4):
     )
 
 
-def reshape_tensor(x, heads):
-    bs, length, width = x.shape
+def reshape_tensor(x: torch.Tensor, heads: int):
+    bs, length, _ = x.shape
     # (bs, length, width) --> (bs, length, n_heads, dim_per_head)
     x = x.view(bs, length, heads, -1)
     # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
@@ -31,7 +31,7 @@ def reshape_tensor(x, heads):
 
 
 class PerceiverAttention(nn.Module):
-    def __init__(self, *, dim, dim_head=64, heads=8):
+    def __init__(self, *, dim: int, dim_head: int = 64, heads: int = 8):
         super().__init__()
         self.scale = dim_head**-0.5
         self.dim_head = dim_head
@@ -45,7 +45,7 @@ class PerceiverAttention(nn.Module):
         self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
         self.to_out = nn.Linear(inner_dim, dim, bias=False)
 
-    def forward(self, x, latents):
+    def forward(self, x: torch.Tensor, latents: torch.Tensor):
         """
         Args:
             x (torch.Tensor): image features
@@ -80,14 +80,14 @@ class PerceiverAttention(nn.Module):
 class Resampler(nn.Module):
     def __init__(
         self,
-        dim=1024,
-        depth=8,
-        dim_head=64,
-        heads=16,
-        num_queries=8,
-        embedding_dim=768,
-        output_dim=1024,
-        ff_mult=4,
+        dim: int = 1024,
+        depth: int = 8,
+        dim_head: int = 64,
+        heads: int = 16,
+        num_queries: int = 8,
+        embedding_dim: int = 768,
+        output_dim: int = 1024,
+        ff_mult: int = 4,
     ):
         super().__init__()
 
@@ -110,7 +110,15 @@ class Resampler(nn.Module):
             )
 
     @classmethod
-    def from_state_dict(cls, state_dict: dict[torch.Tensor], depth=8, dim_head=64, heads=16, num_queries=8, ff_mult=4):
+    def from_state_dict(
+        cls,
+        state_dict: dict[str, torch.Tensor],
+        depth: int = 8,
+        dim_head: int = 64,
+        heads: int = 16,
+        num_queries: int = 8,
+        ff_mult: int = 4,
+    ):
         """A convenience function that initializes a Resampler from a state_dict.
 
         Some of the shape parameters are inferred from the state_dict (e.g. dim, embedding_dim, etc.). At the time of
@@ -145,7 +153,7 @@ class Resampler(nn.Module):
         model.load_state_dict(state_dict)
         return model
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor):
         latents = self.latents.repeat(x.size(0), 1, 1)
 
         x = self.proj_in(x)

invokeai/backend/ip_adapter/unet_patcher.py

@@ -1,53 +0,0 @@
-from contextlib import contextmanager
-
-from diffusers.models import UNet2DConditionModel
-
-from invokeai.backend.ip_adapter.attention_processor import AttnProcessor2_0, IPAttnProcessor2_0
-from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
-
-
-class UNetPatcher:
-    """A class that contains multiple IP-Adapters and can apply them to a UNet."""
-
-    def __init__(self, ip_adapters: list[IPAdapter]):
-        self._ip_adapters = ip_adapters
-        self._scales = [1.0] * len(self._ip_adapters)
-
-    def set_scale(self, idx: int, value: float):
-        self._scales[idx] = value
-
-    def _prepare_attention_processors(self, unet: UNet2DConditionModel):
-        """Prepare a dict of attention processors that can be injected into a unet, and load the IP-Adapter attention
-        weights into them.
-
-        Note that the `unet` param is only used to determine attention block dimensions and naming.
-        """
-        # Construct a dict of attention processors based on the UNet's architecture.
-        attn_procs = {}
-        for idx, name in enumerate(unet.attn_processors.keys()):
-            if name.endswith("attn1.processor"):
-                attn_procs[name] = AttnProcessor2_0()
-            else:
-                # Collect the weights from each IP Adapter for the idx'th attention processor.
-                attn_procs[name] = IPAttnProcessor2_0(
-                    [ip_adapter.attn_weights.get_attention_processor_weights(idx) for ip_adapter in self._ip_adapters],
-                    self._scales,
-                )
-        return attn_procs
-
-    @contextmanager
-    def apply_ip_adapter_attention(self, unet: UNet2DConditionModel):
-        """A context manager that patches `unet` with IP-Adapter attention processors."""
-
-        attn_procs = self._prepare_attention_processors(unet)
-
-        orig_attn_processors = unet.attn_processors
-
-        try:
-            # Note to future devs: set_attn_processor(...) does something slightly unexpected - it pops elements from the
-            # passed dict. So, if you wanted to keep the dict for future use, you'd have to make a moderately-shallow copy
-            # of it. E.g. `attn_procs_copy = {k: v for k, v in attn_procs.items()}`.
-            unet.set_attn_processor(attn_procs)
-            yield None
-        finally:
-            unet.set_attn_processor(orig_attn_processors)

invokeai/backend/model_manager/config.py

@@ -301,12 +301,12 @@ class MainConfigBase(ModelConfigBase):
     default_settings: Optional[MainModelDefaultSettings] = Field(
         description="Default settings for this model", default=None
     )
+    variant: ModelVariantType = ModelVariantType.Normal
 
 
 class MainCheckpointConfig(CheckpointConfigBase, MainConfigBase):
     """Model config for main checkpoint models."""
 
-    variant: ModelVariantType = ModelVariantType.Normal
     prediction_type: SchedulerPredictionType = SchedulerPredictionType.Epsilon
     upcast_attention: bool = False
 
@@ -323,10 +323,13 @@ class MainDiffusersConfig(DiffusersConfigBase, MainConfigBase):
         return Tag(f"{ModelType.Main.value}.{ModelFormat.Diffusers.value}")
 
 
-class IPAdapterConfig(ModelConfigBase):
-    """Model config for IP Adaptor format models."""
-
+class IPAdapterBaseConfig(ModelConfigBase):
     type: Literal[ModelType.IPAdapter] = ModelType.IPAdapter
+
+
+class IPAdapterInvokeAIConfig(IPAdapterBaseConfig):
+    """Model config for IP Adapter diffusers format models."""
+
     image_encoder_model_id: str
     format: Literal[ModelFormat.InvokeAI]
 
@@ -335,6 +338,16 @@ class IPAdapterConfig(ModelConfigBase):
         return Tag(f"{ModelType.IPAdapter.value}.{ModelFormat.InvokeAI.value}")
 
 
+class IPAdapterCheckpointConfig(IPAdapterBaseConfig):
+    """Model config for IP Adapter checkpoint format models."""
+
+    format: Literal[ModelFormat.Checkpoint]
+
+    @staticmethod
+    def get_tag() -> Tag:
+        return Tag(f"{ModelType.IPAdapter.value}.{ModelFormat.Checkpoint.value}")
+
+
 class CLIPVisionDiffusersConfig(DiffusersConfigBase):
     """Model config for CLIPVision."""
 
@@ -390,7 +403,8 @@ AnyModelConfig = Annotated[
         Annotated[LoRADiffusersConfig, LoRADiffusersConfig.get_tag()],
         Annotated[TextualInversionFileConfig, TextualInversionFileConfig.get_tag()],
         Annotated[TextualInversionFolderConfig, TextualInversionFolderConfig.get_tag()],
-        Annotated[IPAdapterConfig, IPAdapterConfig.get_tag()],
+        Annotated[IPAdapterInvokeAIConfig, IPAdapterInvokeAIConfig.get_tag()],
+        Annotated[IPAdapterCheckpointConfig, IPAdapterCheckpointConfig.get_tag()],
         Annotated[T2IAdapterConfig, T2IAdapterConfig.get_tag()],
         Annotated[CLIPVisionDiffusersConfig, CLIPVisionDiffusersConfig.get_tag()],
     ],

invokeai/backend/model_manager/load/load_default.py

@@ -18,7 +18,7 @@ from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoad
 from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase
 from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data, calc_model_size_by_fs
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
-from invokeai.backend.util.devices import choose_torch_device, torch_dtype
+from invokeai.backend.util.devices import TorchDevice
 
 
 # TO DO: The loader is not thread safe!
@@ -37,7 +37,7 @@ class ModelLoader(ModelLoaderBase):
         self._logger = logger
         self._ram_cache = ram_cache
         self._convert_cache = convert_cache
-        self._torch_dtype = torch_dtype(choose_torch_device(), app_config)
+        self._torch_dtype = TorchDevice.choose_torch_dtype()
 
     def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
         """

invokeai/backend/model_manager/load/model_cache/model_cache_base.py

@@ -117,7 +117,7 @@ class ModelCacheBase(ABC, Generic[T]):
 
     @property
     @abstractmethod
-    def stats(self) -> CacheStats:
+    def stats(self) -> Optional[CacheStats]:
         """Return collected CacheStats object."""
         pass
 

invokeai/backend/model_manager/load/model_cache/model_cache_default.py

@@ -30,15 +30,12 @@ import torch
 
 from invokeai.backend.model_manager import AnyModel, SubModelType
 from invokeai.backend.model_manager.load.memory_snapshot import MemorySnapshot, get_pretty_snapshot_diff
-from invokeai.backend.util.devices import choose_torch_device
+from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 
 from .model_cache_base import CacheRecord, CacheStats, ModelCacheBase, ModelLockerBase
 from .model_locker import ModelLocker
 
-if choose_torch_device() == torch.device("mps"):
-    from torch import mps
-
 # Maximum size of the cache, in gigs
 # Default is roughly enough to hold three fp16 diffusers models in RAM simultaneously
 DEFAULT_MAX_CACHE_SIZE = 6.0
@@ -244,9 +241,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
                     f"Removing {cache_entry.key} from VRAM to free {(cache_entry.size/GIG):.2f}GB; vram free = {(torch.cuda.memory_allocated()/GIG):.2f}GB"
                 )
 
-        torch.cuda.empty_cache()
-        if choose_torch_device() == torch.device("mps"):
-            mps.empty_cache()
+        TorchDevice.empty_cache()
 
     def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device: torch.device) -> None:
         """Move model into the indicated device.
@@ -269,12 +264,14 @@ class ModelCache(ModelCacheBase[AnyModel]):
         if torch.device(source_device).type == torch.device(target_device).type:
             return
 
-        # may raise an exception here if insufficient GPU VRAM
-        self._check_free_vram(target_device, cache_entry.size)
-
         start_model_to_time = time.time()
         snapshot_before = self._capture_memory_snapshot()
-        cache_entry.model.to(target_device)
+        try:
+            cache_entry.model.to(target_device)
+        except Exception as e:  # blow away cache entry
+            self._delete_cache_entry(cache_entry)
+            raise e
+
         snapshot_after = self._capture_memory_snapshot()
         end_model_to_time = time.time()
         self.logger.debug(
@@ -329,11 +326,11 @@ class ModelCache(ModelCacheBase[AnyModel]):
                     f" {in_ram_models}/{in_vram_models}({locked_in_vram_models})"
                 )
 
-    def make_room(self, model_size: int) -> None:
+    def make_room(self, size: int) -> None:
         """Make enough room in the cache to accommodate a new model of indicated size."""
         # calculate how much memory this model will require
         # multiplier = 2 if self.precision==torch.float32 else 1
-        bytes_needed = model_size
+        bytes_needed = size
         maximum_size = self.max_cache_size * GIG  # stored in GB, convert to bytes
         current_size = self.cache_size()
 
@@ -388,12 +385,11 @@ class ModelCache(ModelCacheBase[AnyModel]):
             # 1 from onnx runtime object
             if not cache_entry.locked and refs <= (3 if "onnx" in model_key else 2):
                 self.logger.debug(
-                    f"Removing {model_key} from RAM cache to free at least {(model_size/GIG):.2f} GB (-{(cache_entry.size/GIG):.2f} GB)"
+                    f"Removing {model_key} from RAM cache to free at least {(size/GIG):.2f} GB (-{(cache_entry.size/GIG):.2f} GB)"
                 )
                 current_size -= cache_entry.size
                 models_cleared += 1
-                del self._cache_stack[pos]
-                del self._cached_models[model_key]
+                self._delete_cache_entry(cache_entry)
                 del cache_entry
 
             else:
@@ -415,18 +411,9 @@ class ModelCache(ModelCacheBase[AnyModel]):
                 self.stats.cleared = models_cleared
             gc.collect()
 
-        torch.cuda.empty_cache()
-        if choose_torch_device() == torch.device("mps"):
-            mps.empty_cache()
-
+        TorchDevice.empty_cache()
         self.logger.debug(f"After making room: cached_models={len(self._cached_models)}")
 
-    def _check_free_vram(self, target_device: torch.device, needed_size: int) -> None:
-        if target_device.type != "cuda":
-            return
-        vram_device = (  # mem_get_info() needs an indexed device
-            target_device if target_device.index is not None else torch.device(str(target_device), index=0)
-        )
-        free_mem, _ = torch.cuda.mem_get_info(torch.device(vram_device))
-        if needed_size > free_mem:
-            raise torch.cuda.OutOfMemoryError
+    def _delete_cache_entry(self, cache_entry: CacheRecord[AnyModel]) -> None:
+        self._cache_stack.remove(cache_entry.key)
+        del self._cached_models[cache_entry.key]

invokeai/backend/model_manager/load/model_cache/model_locker.py

@@ -34,7 +34,6 @@ class ModelLocker(ModelLockerBase):
 
         # NOTE that the model has to have the to() method in order for this code to move it into GPU!
         self._cache_entry.lock()
-
         try:
             if self._cache.lazy_offloading:
                 self._cache.offload_unlocked_models(self._cache_entry.size)
@@ -51,6 +50,7 @@ class ModelLocker(ModelLockerBase):
         except Exception:
             self._cache_entry.unlock()
             raise
+
         return self.model
 
     def unlock(self) -> None:

invokeai/backend/model_manager/load/model_loaders/ip_adapter.py

@@ -7,19 +7,13 @@ from typing import Optional
 import torch
 
 from invokeai.backend.ip_adapter.ip_adapter import build_ip_adapter
-from invokeai.backend.model_manager import (
-    AnyModel,
-    AnyModelConfig,
-    BaseModelType,
-    ModelFormat,
-    ModelType,
-    SubModelType,
-)
+from invokeai.backend.model_manager import AnyModel, AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType
 from invokeai.backend.model_manager.load import ModelLoader, ModelLoaderRegistry
 from invokeai.backend.raw_model import RawModel
 
 
 @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.IPAdapter, format=ModelFormat.InvokeAI)
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.IPAdapter, format=ModelFormat.Checkpoint)
 class IPAdapterInvokeAILoader(ModelLoader):
     """Class to load IP Adapter diffusers models."""
 
@@ -32,7 +26,7 @@ class IPAdapterInvokeAILoader(ModelLoader):
             raise ValueError("There are no submodels in an IP-Adapter model.")
         model_path = Path(config.path)
         model: RawModel = build_ip_adapter(
-            ip_adapter_ckpt_path=str(model_path / "ip_adapter.bin"),
+            ip_adapter_ckpt_path=model_path,
             device=torch.device("cpu"),
             dtype=self._torch_dtype,
         )

invokeai/backend/model_manager/merge.py

@@ -17,7 +17,7 @@ from diffusers.utils import logging as dlogging
 
 from invokeai.app.services.model_install import ModelInstallServiceBase
 from invokeai.app.services.model_records.model_records_base import ModelRecordChanges
-from invokeai.backend.util.devices import choose_torch_device, torch_dtype
+from invokeai.backend.util.devices import TorchDevice
 
 from . import (
     AnyModelConfig,
@@ -43,6 +43,7 @@ class ModelMerger(object):
         Initialize a ModelMerger object with the model installer.
         """
         self._installer = installer
+        self._dtype = TorchDevice.choose_torch_dtype()
 
     def merge_diffusion_models(
         self,
@@ -68,7 +69,7 @@ class ModelMerger(object):
             warnings.simplefilter("ignore")
             verbosity = dlogging.get_verbosity()
             dlogging.set_verbosity_error()
-            dtype = torch.float16 if variant == "fp16" else torch_dtype(choose_torch_device())
+            dtype = torch.float16 if variant == "fp16" else self._dtype
 
             # Note that checkpoint_merger will not work with downloaded HuggingFace fp16 models
             # until upstream https://github.com/huggingface/diffusers/pull/6670 is merged and released.
@@ -151,7 +152,7 @@ class ModelMerger(object):
         dump_path.mkdir(parents=True, exist_ok=True)
         dump_path = dump_path / merged_model_name
 
-        dtype = torch.float16 if variant == "fp16" else torch_dtype(choose_torch_device())
+        dtype = torch.float16 if variant == "fp16" else self._dtype
         merged_pipe.save_pretrained(dump_path.as_posix(), safe_serialization=True, torch_dtype=dtype, variant=variant)
 
         # register model and get its unique key

invokeai/backend/model_manager/probe.py

@@ -51,6 +51,7 @@ LEGACY_CONFIGS: Dict[BaseModelType, Dict[ModelVariantType, Union[str, Dict[Sched
     },
     BaseModelType.StableDiffusionXL: {
         ModelVariantType.Normal: "sd_xl_base.yaml",
+        ModelVariantType.Inpaint: "sd_xl_inpaint.yaml",
     },
     BaseModelType.StableDiffusionXLRefiner: {
         ModelVariantType.Normal: "sd_xl_refiner.yaml",
@@ -230,9 +231,10 @@ class ModelProbe(object):
                 return ModelType.LoRA
             elif any(key.startswith(v) for v in {"controlnet", "control_model", "input_blocks"}):
                 return ModelType.ControlNet
+            elif any(key.startswith(v) for v in {"image_proj.", "ip_adapter."}):
+                return ModelType.IPAdapter
             elif key in {"emb_params", "string_to_param"}:
                 return ModelType.TextualInversion
-
         else:
             # diffusers-ti
             if len(ckpt) < 10 and all(isinstance(v, torch.Tensor) for v in ckpt.values()):
@@ -323,7 +325,7 @@ class ModelProbe(object):
         with SilenceWarnings():
             if model_path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin")):
                 cls._scan_model(model_path.name, model_path)
-                model = torch.load(model_path)
+                model = torch.load(model_path, map_location="cpu")
                 assert isinstance(model, dict)
                 return model
             else:
@@ -527,8 +529,25 @@ class ControlNetCheckpointProbe(CheckpointProbeBase):
 
 
 class IPAdapterCheckpointProbe(CheckpointProbeBase):
+    """Class for probing IP Adapters"""
+
     def get_base_type(self) -> BaseModelType:
-        raise NotImplementedError()
+        checkpoint = self.checkpoint
+        for key in checkpoint.keys():
+            if not key.startswith(("image_proj.", "ip_adapter.")):
+                continue
+            cross_attention_dim = checkpoint["ip_adapter.1.to_k_ip.weight"].shape[-1]
+            if cross_attention_dim == 768:
+                return BaseModelType.StableDiffusion1
+            elif cross_attention_dim == 1024:
+                return BaseModelType.StableDiffusion2
+            elif cross_attention_dim == 2048:
+                return BaseModelType.StableDiffusionXL
+            else:
+                raise InvalidModelConfigException(
+                    f"IP-Adapter had unexpected cross-attention dimension: {cross_attention_dim}."
+                )
+        raise InvalidModelConfigException(f"{self.model_path}: Unable to determine base type")
 
 
 class CLIPVisionCheckpointProbe(CheckpointProbeBase):
@@ -768,7 +787,7 @@ class T2IAdapterFolderProbe(FolderProbeBase):
             )
 
 
-############## register probe classes ######
+# Register probe classes
 ModelProbe.register_probe("diffusers", ModelType.Main, PipelineFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.VAE, VaeFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.LoRA, LoRAFolderProbe)

invokeai/backend/model_manager/starter_models.py

@@ -155,7 +155,7 @@ STARTER_MODELS: list[StarterModel] = [
     StarterModel(
         name="IP Adapter",
         base=BaseModelType.StableDiffusion1,
-        source="InvokeAI/ip_adapter_sd15",
+        source="https://huggingface.co/InvokeAI/ip_adapter_sd15/resolve/main/ip-adapter_sd15.safetensors",
         description="IP-Adapter for SD 1.5 models",
         type=ModelType.IPAdapter,
         dependencies=[ip_adapter_sd_image_encoder],
@@ -163,7 +163,7 @@ STARTER_MODELS: list[StarterModel] = [
     StarterModel(
         name="IP Adapter Plus",
         base=BaseModelType.StableDiffusion1,
-        source="InvokeAI/ip_adapter_plus_sd15",
+        source="https://huggingface.co/InvokeAI/ip_adapter_plus_sd15/resolve/main/ip-adapter-plus_sd15.safetensors",
         description="Refined IP-Adapter for SD 1.5 models",
         type=ModelType.IPAdapter,
         dependencies=[ip_adapter_sd_image_encoder],
@@ -171,7 +171,7 @@ STARTER_MODELS: list[StarterModel] = [
     StarterModel(
         name="IP Adapter Plus Face",
         base=BaseModelType.StableDiffusion1,
-        source="InvokeAI/ip_adapter_plus_face_sd15",
+        source="https://huggingface.co/InvokeAI/ip_adapter_plus_face_sd15/resolve/main/ip-adapter-plus-face_sd15.safetensors",
         description="Refined IP-Adapter for SD 1.5 models, adapted for faces",
         type=ModelType.IPAdapter,
         dependencies=[ip_adapter_sd_image_encoder],
@@ -179,7 +179,7 @@ STARTER_MODELS: list[StarterModel] = [
     StarterModel(
         name="IP Adapter SDXL",
         base=BaseModelType.StableDiffusionXL,
-        source="InvokeAI/ip_adapter_sdxl",
+        source="https://huggingface.co/InvokeAI/ip_adapter_sdxl_vit_h/resolve/main/ip-adapter_sdxl_vit-h.safetensors",
         description="IP-Adapter for SDXL models",
         type=ModelType.IPAdapter,
         dependencies=[ip_adapter_sdxl_image_encoder],

invokeai/backend/stable_diffusion/diffusers_pipeline.py

@@ -21,12 +21,11 @@ from pydantic import Field
 from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 
 from invokeai.app.services.config.config_default import get_config
-from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
-from invokeai.backend.ip_adapter.unet_patcher import UNetPatcher
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningData
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import IPAdapterData, TextConditioningData
 from invokeai.backend.stable_diffusion.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent
+from invokeai.backend.stable_diffusion.diffusion.unet_attention_patcher import UNetAttentionPatcher, UNetIPAdapterData
 from invokeai.backend.util.attention import auto_detect_slice_size
-from invokeai.backend.util.devices import normalize_device
+from invokeai.backend.util.devices import TorchDevice
 
 
 @dataclass
@@ -149,16 +148,6 @@ class ControlNetData:
     resize_mode: str = Field(default="just_resize")
 
 
-@dataclass
-class IPAdapterData:
-    ip_adapter_model: IPAdapter = Field(default=None)
-    # TODO: change to polymorphic so can do different weights per step (once implemented...)
-    weight: Union[float, List[float]] = Field(default=1.0)
-    # weight: float = Field(default=1.0)
-    begin_step_percent: float = Field(default=0.0)
-    end_step_percent: float = Field(default=1.0)
-
-
 @dataclass
 class T2IAdapterData:
     """A structure containing the information required to apply conditioning from a single T2I-Adapter model."""
@@ -266,7 +255,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         if self.unet.device.type == "cpu" or self.unet.device.type == "mps":
             mem_free = psutil.virtual_memory().free
         elif self.unet.device.type == "cuda":
-            mem_free, _ = torch.cuda.mem_get_info(normalize_device(self.unet.device))
+            mem_free, _ = torch.cuda.mem_get_info(TorchDevice.normalize(self.unet.device))
         else:
             raise ValueError(f"unrecognized device {self.unet.device}")
         # input tensor of [1, 4, h/8, w/8]
@@ -295,7 +284,8 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         self,
         latents: torch.Tensor,
         num_inference_steps: int,
-        conditioning_data: ConditioningData,
+        scheduler_step_kwargs: dict[str, Any],
+        conditioning_data: TextConditioningData,
         *,
         noise: Optional[torch.Tensor],
         timesteps: torch.Tensor,
@@ -308,7 +298,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         mask: Optional[torch.Tensor] = None,
         masked_latents: Optional[torch.Tensor] = None,
         gradient_mask: Optional[bool] = False,
-        seed: Optional[int] = None,
+        seed: int,
     ) -> torch.Tensor:
         if init_timestep.shape[0] == 0:
             return latents
@@ -326,20 +316,6 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
             latents = self.scheduler.add_noise(latents, noise, batched_t)
 
         if mask is not None:
-            # if no noise provided, noisify unmasked area based on seed(or 0 as fallback)
-            if noise is None:
-                noise = torch.randn(
-                    orig_latents.shape,
-                    dtype=torch.float32,
-                    device="cpu",
-                    generator=torch.Generator(device="cpu").manual_seed(seed or 0),
-                ).to(device=orig_latents.device, dtype=orig_latents.dtype)
-
-                latents = self.scheduler.add_noise(latents, noise, batched_t)
-                latents = torch.lerp(
-                    orig_latents, latents.to(dtype=orig_latents.dtype), mask.to(dtype=orig_latents.dtype)
-                )
-
             if is_inpainting_model(self.unet):
                 if masked_latents is None:
                     raise Exception("Source image required for inpaint mask when inpaint model used!")
@@ -348,6 +324,15 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
                     self._unet_forward, mask, masked_latents
                 )
             else:
+                # if no noise provided, noisify unmasked area based on seed
+                if noise is None:
+                    noise = torch.randn(
+                        orig_latents.shape,
+                        dtype=torch.float32,
+                        device="cpu",
+                        generator=torch.Generator(device="cpu").manual_seed(seed),
+                    ).to(device=orig_latents.device, dtype=orig_latents.dtype)
+
                 additional_guidance.append(AddsMaskGuidance(mask, orig_latents, self.scheduler, noise, gradient_mask))
 
         try:
@@ -355,6 +340,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
                 latents,
                 timesteps,
                 conditioning_data,
+                scheduler_step_kwargs=scheduler_step_kwargs,
                 additional_guidance=additional_guidance,
                 control_data=control_data,
                 ip_adapter_data=ip_adapter_data,
@@ -380,7 +366,8 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         self,
         latents: torch.Tensor,
         timesteps,
-        conditioning_data: ConditioningData,
+        conditioning_data: TextConditioningData,
+        scheduler_step_kwargs: dict[str, Any],
         *,
         additional_guidance: List[Callable] = None,
         control_data: List[ControlNetData] = None,
@@ -397,22 +384,22 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         if timesteps.shape[0] == 0:
             return latents
 
-        ip_adapter_unet_patcher = None
-        extra_conditioning_info = conditioning_data.text_embeddings.extra_conditioning
-        if extra_conditioning_info is not None and extra_conditioning_info.wants_cross_attention_control:
-            attn_ctx = self.invokeai_diffuser.custom_attention_context(
-                self.invokeai_diffuser.model,
-                extra_conditioning_info=extra_conditioning_info,
+        use_ip_adapter = ip_adapter_data is not None
+        use_regional_prompting = (
+            conditioning_data.cond_regions is not None or conditioning_data.uncond_regions is not None
+        )
+        unet_attention_patcher = None
+        self.use_ip_adapter = use_ip_adapter
+        attn_ctx = nullcontext()
+
+        if use_ip_adapter or use_regional_prompting:
+            ip_adapters: Optional[List[UNetIPAdapterData]] = (
+                [{"ip_adapter": ipa.ip_adapter_model, "target_blocks": ipa.target_blocks} for ipa in ip_adapter_data]
+                if use_ip_adapter
+                else None
             )
-            self.use_ip_adapter = False
-        elif ip_adapter_data is not None:
-            # TODO(ryand): Should we raise an exception if both custom attention and IP-Adapter attention are active?
-            # As it is now, the IP-Adapter will silently be skipped.
-            ip_adapter_unet_patcher = UNetPatcher([ipa.ip_adapter_model for ipa in ip_adapter_data])
-            attn_ctx = ip_adapter_unet_patcher.apply_ip_adapter_attention(self.invokeai_diffuser.model)
-            self.use_ip_adapter = True
-        else:
-            attn_ctx = nullcontext()
+            unet_attention_patcher = UNetAttentionPatcher(ip_adapters)
+            attn_ctx = unet_attention_patcher.apply_ip_adapter_attention(self.invokeai_diffuser.model)
 
         with attn_ctx:
             if callback is not None:
@@ -435,11 +422,11 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
                     conditioning_data,
                     step_index=i,
                     total_step_count=len(timesteps),
+                    scheduler_step_kwargs=scheduler_step_kwargs,
                     additional_guidance=additional_guidance,
                     control_data=control_data,
                     ip_adapter_data=ip_adapter_data,
                     t2i_adapter_data=t2i_adapter_data,
-                    ip_adapter_unet_patcher=ip_adapter_unet_patcher,
                 )
                 latents = step_output.prev_sample
                 predicted_original = getattr(step_output, "pred_original_sample", None)
@@ -463,14 +450,14 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         self,
         t: torch.Tensor,
         latents: torch.Tensor,
-        conditioning_data: ConditioningData,
+        conditioning_data: TextConditioningData,
         step_index: int,
         total_step_count: int,
+        scheduler_step_kwargs: dict[str, Any],
         additional_guidance: List[Callable] = None,
         control_data: List[ControlNetData] = None,
         ip_adapter_data: Optional[list[IPAdapterData]] = None,
         t2i_adapter_data: Optional[list[T2IAdapterData]] = None,
-        ip_adapter_unet_patcher: Optional[UNetPatcher] = None,
     ):
         # invokeai_diffuser has batched timesteps, but diffusers schedulers expect a single value
         timestep = t[0]
@@ -485,23 +472,6 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         #     i.e. before or after passing it to InvokeAIDiffuserComponent
         latent_model_input = self.scheduler.scale_model_input(latents, timestep)
 
-        # handle IP-Adapter
-        if self.use_ip_adapter and ip_adapter_data is not None:  # somewhat redundant but logic is clearer
-            for i, single_ip_adapter_data in enumerate(ip_adapter_data):
-                first_adapter_step = math.floor(single_ip_adapter_data.begin_step_percent * total_step_count)
-                last_adapter_step = math.ceil(single_ip_adapter_data.end_step_percent * total_step_count)
-                weight = (
-                    single_ip_adapter_data.weight[step_index]
-                    if isinstance(single_ip_adapter_data.weight, List)
-                    else single_ip_adapter_data.weight
-                )
-                if step_index >= first_adapter_step and step_index <= last_adapter_step:
-                    # Only apply this IP-Adapter if the current step is within the IP-Adapter's begin/end step range.
-                    ip_adapter_unet_patcher.set_scale(i, weight)
-                else:
-                    # Otherwise, set the IP-Adapter's scale to 0, so it has no effect.
-                    ip_adapter_unet_patcher.set_scale(i, 0.0)
-
         # Handle ControlNet(s)
         down_block_additional_residuals = None
         mid_block_additional_residual = None
@@ -550,6 +520,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
             step_index=step_index,
             total_step_count=total_step_count,
             conditioning_data=conditioning_data,
+            ip_adapter_data=ip_adapter_data,
             down_block_additional_residuals=down_block_additional_residuals,  # for ControlNet
             mid_block_additional_residual=mid_block_additional_residual,  # for ControlNet
             down_intrablock_additional_residuals=down_intrablock_additional_residuals,  # for T2I-Adapter
@@ -569,7 +540,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
             )
 
         # compute the previous noisy sample x_t -> x_t-1
-        step_output = self.scheduler.step(noise_pred, timestep, latents, **conditioning_data.scheduler_args)
+        step_output = self.scheduler.step(noise_pred, timestep, latents, **scheduler_step_kwargs)
 
         # TODO: discuss injection point options. For now this is a patch to get progress images working with inpainting again.
         for guidance in additional_guidance:

invokeai/backend/stable_diffusion/diffusion/conditioning_data.py

@@ -1,27 +1,17 @@
-import dataclasses
-import inspect
-from dataclasses import dataclass, field
-from typing import Any, List, Optional, Union
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Union
 
 import torch
 
-from .cross_attention_control import Arguments
-
-
-@dataclass
-class ExtraConditioningInfo:
-    tokens_count_including_eos_bos: int
-    cross_attention_control_args: Optional[Arguments] = None
-
-    @property
-    def wants_cross_attention_control(self):
-        return self.cross_attention_control_args is not None
+from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
 
 
 @dataclass
 class BasicConditioningInfo:
+    """SD 1/2 text conditioning information produced by Compel."""
+
     embeds: torch.Tensor
-    extra_conditioning: Optional[ExtraConditioningInfo]
 
     def to(self, device, dtype=None):
         self.embeds = self.embeds.to(device=device, dtype=dtype)
@@ -35,6 +25,8 @@ class ConditioningFieldData:
 
 @dataclass
 class SDXLConditioningInfo(BasicConditioningInfo):
+    """SDXL text conditioning information produced by Compel."""
+
     pooled_embeds: torch.Tensor
     add_time_ids: torch.Tensor
 
@@ -57,37 +49,75 @@ class IPAdapterConditioningInfo:
 
 
 @dataclass
-class ConditioningData:
-    unconditioned_embeddings: BasicConditioningInfo
-    text_embeddings: BasicConditioningInfo
-    """
-    Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-    `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf).
-    Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate
-    images that are closely linked to the text `prompt`, usually at the expense of lower image quality.
-    """
-    guidance_scale: Union[float, List[float]]
-    """ for models trained using zero-terminal SNR ("ztsnr"), it's suggested to use guidance_rescale_multiplier of 0.7 .
-     ref [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf)
-    """
-    guidance_rescale_multiplier: float = 0
-    scheduler_args: dict[str, Any] = field(default_factory=dict)
+class IPAdapterData:
+    ip_adapter_model: IPAdapter
+    ip_adapter_conditioning: IPAdapterConditioningInfo
+    mask: torch.Tensor
+    target_blocks: List[str]
 
-    ip_adapter_conditioning: Optional[list[IPAdapterConditioningInfo]] = None
+    # Either a single weight applied to all steps, or a list of weights for each step.
+    weight: Union[float, List[float]] = 1.0
+    begin_step_percent: float = 0.0
+    end_step_percent: float = 1.0
 
-    @property
-    def dtype(self):
-        return self.text_embeddings.dtype
+    def scale_for_step(self, step_index: int, total_steps: int) -> float:
+        first_adapter_step = math.floor(self.begin_step_percent * total_steps)
+        last_adapter_step = math.ceil(self.end_step_percent * total_steps)
+        weight = self.weight[step_index] if isinstance(self.weight, List) else self.weight
+        if step_index >= first_adapter_step and step_index <= last_adapter_step:
+            # Only apply this IP-Adapter if the current step is within the IP-Adapter's begin/end step range.
+            return weight
+        # Otherwise, set the IP-Adapter's scale to 0, so it has no effect.
+        return 0.0
 
-    def add_scheduler_args_if_applicable(self, scheduler, **kwargs):
-        scheduler_args = dict(self.scheduler_args)
-        step_method = inspect.signature(scheduler.step)
-        for name, value in kwargs.items():
-            try:
-                step_method.bind_partial(**{name: value})
-            except TypeError:
-                # FIXME: don't silently discard arguments
-                pass  # debug("%s does not accept argument named %r", scheduler, name)
-            else:
-                scheduler_args[name] = value
-        return dataclasses.replace(self, scheduler_args=scheduler_args)
+
+@dataclass
+class Range:
+    start: int
+    end: int
+
+
+class TextConditioningRegions:
+    def __init__(
+        self,
+        masks: torch.Tensor,
+        ranges: list[Range],
+    ):
+        # A binary mask indicating the regions of the image that the prompt should be applied to.
+        # Shape: (1, num_prompts, height, width)
+        # Dtype: torch.bool
+        self.masks = masks
+
+        # A list of ranges indicating the start and end indices of the embeddings that corresponding mask applies to.
+        # ranges[i] contains the embedding range for the i'th prompt / mask.
+        self.ranges = ranges
+
+        assert self.masks.shape[1] == len(self.ranges)
+
+
+class TextConditioningData:
+    def __init__(
+        self,
+        uncond_text: Union[BasicConditioningInfo, SDXLConditioningInfo],
+        cond_text: Union[BasicConditioningInfo, SDXLConditioningInfo],
+        uncond_regions: Optional[TextConditioningRegions],
+        cond_regions: Optional[TextConditioningRegions],
+        guidance_scale: Union[float, List[float]],
+        guidance_rescale_multiplier: float = 0,
+    ):
+        self.uncond_text = uncond_text
+        self.cond_text = cond_text
+        self.uncond_regions = uncond_regions
+        self.cond_regions = cond_regions
+        # Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+        # `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf).
+        # Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate
+        # images that are closely linked to the text `prompt`, usually at the expense of lower image quality.
+        self.guidance_scale = guidance_scale
+        # For models trained using zero-terminal SNR ("ztsnr"), it's suggested to use guidance_rescale_multiplier of 0.7.
+        # See [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+        self.guidance_rescale_multiplier = guidance_rescale_multiplier
+
+    def is_sdxl(self):
+        assert isinstance(self.uncond_text, SDXLConditioningInfo) == isinstance(self.cond_text, SDXLConditioningInfo)
+        return isinstance(self.cond_text, SDXLConditioningInfo)

invokeai/backend/stable_diffusion/diffusion/cross_attention_control.py

@@ -1,218 +0,0 @@
-# adapted from bloc97's CrossAttentionControl colab
-# https://github.com/bloc97/CrossAttentionControl
-
-
-import enum
-from dataclasses import dataclass, field
-from typing import Optional
-
-import torch
-from compel.cross_attention_control import Arguments
-from diffusers.models.attention_processor import Attention, SlicedAttnProcessor
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
-
-from invokeai.backend.util.devices import torch_dtype
-
-
-class CrossAttentionType(enum.Enum):
-    SELF = 1
-    TOKENS = 2
-
-
-class CrossAttnControlContext:
-    def __init__(self, arguments: Arguments):
-        """
-        :param arguments: Arguments for the cross-attention control process
-        """
-        self.cross_attention_mask: Optional[torch.Tensor] = None
-        self.cross_attention_index_map: Optional[torch.Tensor] = None
-        self.arguments = arguments
-
-    def get_active_cross_attention_control_types_for_step(
-        self, percent_through: float = None
-    ) -> list[CrossAttentionType]:
-        """
-        Should cross-attention control be applied on the given step?
-        :param percent_through: How far through the step sequence are we (0.0=pure noise, 1.0=completely denoised image). Expected range 0.0..<1.0.
-        :return: A list of attention types that cross-attention control should be performed for on the given step. May be [].
-        """
-        if percent_through is None:
-            return [CrossAttentionType.SELF, CrossAttentionType.TOKENS]
-
-        opts = self.arguments.edit_options
-        to_control = []
-        if opts["s_start"] <= percent_through < opts["s_end"]:
-            to_control.append(CrossAttentionType.SELF)
-        if opts["t_start"] <= percent_through < opts["t_end"]:
-            to_control.append(CrossAttentionType.TOKENS)
-        return to_control
-
-
-def setup_cross_attention_control_attention_processors(unet: UNet2DConditionModel, context: CrossAttnControlContext):
-    """
-    Inject attention parameters and functions into the passed in model to enable cross attention editing.
-
-    :param model: The unet model to inject into.
-    :return: None
-    """
-
-    # adapted from init_attention_edit
-    device = context.arguments.edited_conditioning.device
-
-    # urgh. should this be hardcoded?
-    max_length = 77
-    # mask=1 means use base prompt attention, mask=0 means use edited prompt attention
-    mask = torch.zeros(max_length, dtype=torch_dtype(device))
-    indices_target = torch.arange(max_length, dtype=torch.long)
-    indices = torch.arange(max_length, dtype=torch.long)
-    for name, a0, a1, b0, b1 in context.arguments.edit_opcodes:
-        if b0 < max_length:
-            if name == "equal":  # or (name == "replace" and a1 - a0 == b1 - b0):
-                # these tokens have not been edited
-                indices[b0:b1] = indices_target[a0:a1]
-                mask[b0:b1] = 1
-
-    context.cross_attention_mask = mask.to(device)
-    context.cross_attention_index_map = indices.to(device)
-    old_attn_processors = unet.attn_processors
-    if torch.backends.mps.is_available():
-        # see note in StableDiffusionGeneratorPipeline.__init__ about borked slicing on MPS
-        unet.set_attn_processor(SwapCrossAttnProcessor())
-    else:
-        # try to re-use an existing slice size
-        default_slice_size = 4
-        slice_size = next(
-            (p.slice_size for p in old_attn_processors.values() if type(p) is SlicedAttnProcessor), default_slice_size
-        )
-        unet.set_attn_processor(SlicedSwapCrossAttnProcesser(slice_size=slice_size))
-
-
-@dataclass
-class SwapCrossAttnContext:
-    modified_text_embeddings: torch.Tensor
-    index_map: torch.Tensor  # maps from original prompt token indices to the equivalent tokens in the modified prompt
-    mask: torch.Tensor  # in the target space of the index_map
-    cross_attention_types_to_do: list[CrossAttentionType] = field(default_factory=list)
-
-    def wants_cross_attention_control(self, attn_type: CrossAttentionType) -> bool:
-        return attn_type in self.cross_attention_types_to_do
-
-    @classmethod
-    def make_mask_and_index_map(
-        cls, edit_opcodes: list[tuple[str, int, int, int, int]], max_length: int
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # mask=1 means use original prompt attention, mask=0 means use modified prompt attention
-        mask = torch.zeros(max_length)
-        indices_target = torch.arange(max_length, dtype=torch.long)
-        indices = torch.arange(max_length, dtype=torch.long)
-        for name, a0, a1, b0, b1 in edit_opcodes:
-            if b0 < max_length:
-                if name == "equal":
-                    # these tokens remain the same as in the original prompt
-                    indices[b0:b1] = indices_target[a0:a1]
-                    mask[b0:b1] = 1
-
-        return mask, indices
-
-
-class SlicedSwapCrossAttnProcesser(SlicedAttnProcessor):
-    # TODO: dynamically pick slice size based on memory conditions
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states,
-        encoder_hidden_states=None,
-        attention_mask=None,
-        # kwargs
-        swap_cross_attn_context: SwapCrossAttnContext = None,
-        **kwargs,
-    ):
-        attention_type = CrossAttentionType.SELF if encoder_hidden_states is None else CrossAttentionType.TOKENS
-
-        # if cross-attention control is not in play, just call through to the base implementation.
-        if (
-            attention_type is CrossAttentionType.SELF
-            or swap_cross_attn_context is None
-            or not swap_cross_attn_context.wants_cross_attention_control(attention_type)
-        ):
-            # print(f"SwapCrossAttnContext for {attention_type} not active - passing request to superclass")
-            return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask)
-        # else:
-        #    print(f"SwapCrossAttnContext for {attention_type} active")
-
-        batch_size, sequence_length, _ = hidden_states.shape
-        attention_mask = attn.prepare_attention_mask(
-            attention_mask=attention_mask,
-            target_length=sequence_length,
-            batch_size=batch_size,
-        )
-
-        query = attn.to_q(hidden_states)
-        dim = query.shape[-1]
-        query = attn.head_to_batch_dim(query)
-
-        original_text_embeddings = encoder_hidden_states
-        modified_text_embeddings = swap_cross_attn_context.modified_text_embeddings
-        original_text_key = attn.to_k(original_text_embeddings)
-        modified_text_key = attn.to_k(modified_text_embeddings)
-        original_value = attn.to_v(original_text_embeddings)
-        modified_value = attn.to_v(modified_text_embeddings)
-
-        original_text_key = attn.head_to_batch_dim(original_text_key)
-        modified_text_key = attn.head_to_batch_dim(modified_text_key)
-        original_value = attn.head_to_batch_dim(original_value)
-        modified_value = attn.head_to_batch_dim(modified_value)
-
-        # compute slices and prepare output tensor
-        batch_size_attention = query.shape[0]
-        hidden_states = torch.zeros(
-            (batch_size_attention, sequence_length, dim // attn.heads),
-            device=query.device,
-            dtype=query.dtype,
-        )
-
-        # do slices
-        for i in range(max(1, hidden_states.shape[0] // self.slice_size)):
-            start_idx = i * self.slice_size
-            end_idx = (i + 1) * self.slice_size
-
-            query_slice = query[start_idx:end_idx]
-            original_key_slice = original_text_key[start_idx:end_idx]
-            modified_key_slice = modified_text_key[start_idx:end_idx]
-            attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None
-
-            original_attn_slice = attn.get_attention_scores(query_slice, original_key_slice, attn_mask_slice)
-            modified_attn_slice = attn.get_attention_scores(query_slice, modified_key_slice, attn_mask_slice)
-
-            # because the prompt modifications may result in token sequences shifted forwards or backwards,
-            # the original attention probabilities must be remapped to account for token index changes in the
-            # modified prompt
-            remapped_original_attn_slice = torch.index_select(
-                original_attn_slice, -1, swap_cross_attn_context.index_map
-            )
-
-            # only some tokens taken from the original attention probabilities. this is controlled by the mask.
-            mask = swap_cross_attn_context.mask
-            inverse_mask = 1 - mask
-            attn_slice = remapped_original_attn_slice * mask + modified_attn_slice * inverse_mask
-
-            del remapped_original_attn_slice, modified_attn_slice
-
-            attn_slice = torch.bmm(attn_slice, modified_value[start_idx:end_idx])
-            hidden_states[start_idx:end_idx] = attn_slice
-
-        # done
-        hidden_states = attn.batch_to_head_dim(hidden_states)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        return hidden_states
-
-
-class SwapCrossAttnProcessor(SlicedSwapCrossAttnProcesser):
-    def __init__(self):
-        super(SwapCrossAttnProcessor, self).__init__(slice_size=int(1e9))  # massive slice size = don't slice

invokeai/backend/stable_diffusion/diffusion/custom_atttention.py

@@ -0,0 +1,214 @@
+from dataclasses import dataclass
+from typing import List, Optional, cast
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention_processor import Attention, AttnProcessor2_0
+
+from invokeai.backend.ip_adapter.ip_attention_weights import IPAttentionProcessorWeights
+from invokeai.backend.stable_diffusion.diffusion.regional_ip_data import RegionalIPData
+from invokeai.backend.stable_diffusion.diffusion.regional_prompt_data import RegionalPromptData
+
+
+@dataclass
+class IPAdapterAttentionWeights:
+    ip_adapter_weights: IPAttentionProcessorWeights
+    skip: bool
+
+
+class CustomAttnProcessor2_0(AttnProcessor2_0):
+    """A custom implementation of AttnProcessor2_0 that supports additional Invoke features.
+    This implementation is based on
+    https://github.com/huggingface/diffusers/blame/fcfa270fbd1dc294e2f3a505bae6bcb791d721c3/src/diffusers/models/attention_processor.py#L1204
+    Supported custom features:
+    - IP-Adapter
+    - Regional prompt attention
+    """
+
+    def __init__(
+        self,
+        ip_adapter_attention_weights: Optional[List[IPAdapterAttentionWeights]] = None,
+    ):
+        """Initialize a CustomAttnProcessor2_0.
+        Note: Arguments that are the same for all attention layers are passed to __call__(). Arguments that are
+        layer-specific are passed to __init__().
+        Args:
+            ip_adapter_weights: The IP-Adapter attention weights. ip_adapter_weights[i] contains the attention weights
+                for the i'th IP-Adapter.
+        """
+        super().__init__()
+        self._ip_adapter_attention_weights = ip_adapter_attention_weights
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+        # For Regional Prompting:
+        regional_prompt_data: Optional[RegionalPromptData] = None,
+        percent_through: Optional[torch.Tensor] = None,
+        # For IP-Adapter:
+        regional_ip_data: Optional[RegionalIPData] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        """Apply attention.
+        Args:
+            regional_prompt_data: The regional prompt data for the current batch. If not None, this will be used to
+                apply regional prompt masking.
+            regional_ip_data: The IP-Adapter data for the current batch.
+        """
+        # If true, we are doing cross-attention, if false we are doing self-attention.
+        is_cross_attention = encoder_hidden_states is not None
+
+        # Start unmodified block from AttnProcessor2_0.
+        # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+        # End unmodified block from AttnProcessor2_0.
+
+        _, query_seq_len, _ = hidden_states.shape
+        # Handle regional prompt attention masks.
+        if regional_prompt_data is not None and is_cross_attention:
+            assert percent_through is not None
+            prompt_region_attention_mask = regional_prompt_data.get_cross_attn_mask(
+                query_seq_len=query_seq_len, key_seq_len=sequence_length
+            )
+
+            if attention_mask is None:
+                attention_mask = prompt_region_attention_mask
+            else:
+                attention_mask = prompt_region_attention_mask + attention_mask
+
+        # Start unmodified block from AttnProcessor2_0.
+        # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+        # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+        # End unmodified block from AttnProcessor2_0.
+
+        # Apply IP-Adapter conditioning.
+        if is_cross_attention:
+            if self._ip_adapter_attention_weights:
+                assert regional_ip_data is not None
+                ip_masks = regional_ip_data.get_masks(query_seq_len=query_seq_len)
+
+                assert (
+                    len(regional_ip_data.image_prompt_embeds)
+                    == len(self._ip_adapter_attention_weights)
+                    == len(regional_ip_data.scales)
+                    == ip_masks.shape[1]
+                )
+
+                for ipa_index, ipa_embed in enumerate(regional_ip_data.image_prompt_embeds):
+                    ipa_weights = self._ip_adapter_attention_weights[ipa_index].ip_adapter_weights
+                    ipa_scale = regional_ip_data.scales[ipa_index]
+                    ip_mask = ip_masks[0, ipa_index, ...]
+
+                    # The batch dimensions should match.
+                    assert ipa_embed.shape[0] == encoder_hidden_states.shape[0]
+                    # The token_len dimensions should match.
+                    assert ipa_embed.shape[-1] == encoder_hidden_states.shape[-1]
+
+                    ip_hidden_states = ipa_embed
+
+                    # Expected ip_hidden_state shape: (batch_size, num_ip_images, ip_seq_len, ip_image_embedding)
+
+                    if not self._ip_adapter_attention_weights[ipa_index].skip:
+                        ip_key = ipa_weights.to_k_ip(ip_hidden_states)
+                        ip_value = ipa_weights.to_v_ip(ip_hidden_states)
+
+                        # Expected ip_key and ip_value shape:
+                        # (batch_size, num_ip_images, ip_seq_len, head_dim * num_heads)
+
+                        ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+                        ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+                        # Expected ip_key and ip_value shape:
+                        # (batch_size, num_heads, num_ip_images * ip_seq_len, head_dim)
+
+                        # TODO: add support for attn.scale when we move to Torch 2.1
+                        ip_hidden_states = F.scaled_dot_product_attention(
+                            query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+                        )
+
+                        # Expected ip_hidden_states shape: (batch_size, num_heads, query_seq_len, head_dim)
+                        ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(
+                            batch_size, -1, attn.heads * head_dim
+                        )
+
+                        ip_hidden_states = ip_hidden_states.to(query.dtype)
+
+                        # Expected ip_hidden_states shape: (batch_size, query_seq_len, num_heads * head_dim)
+                        hidden_states = hidden_states + ipa_scale * ip_hidden_states * ip_mask
+            else:
+                # If IP-Adapter is not enabled, then regional_ip_data should not be passed in.
+                assert regional_ip_data is None
+
+        # Start unmodified block from AttnProcessor2_0.
+        # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+        # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+        # End of unmodified block from AttnProcessor2_0
+
+        # casting torch.Tensor to torch.FloatTensor to avoid type issues
+        return cast(torch.FloatTensor, hidden_states)

invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py

@@ -0,0 +1,72 @@
+import torch
+
+
+class RegionalIPData:
+    """A class to manage the data for regional IP-Adapter conditioning."""
+
+    def __init__(
+        self,
+        image_prompt_embeds: list[torch.Tensor],
+        scales: list[float],
+        masks: list[torch.Tensor],
+        dtype: torch.dtype,
+        device: torch.device,
+        max_downscale_factor: int = 8,
+    ):
+        """Initialize a `IPAdapterConditioningData` object."""
+        assert len(image_prompt_embeds) == len(scales) == len(masks)
+
+        # The image prompt embeddings.
+        # regional_ip_data[i] contains the image prompt embeddings for the i'th IP-Adapter. Each tensor
+        # has shape (batch_size, num_ip_images, seq_len, ip_embedding_len).
+        self.image_prompt_embeds = image_prompt_embeds
+
+        # The scales for the IP-Adapter attention.
+        # scales[i] contains the attention scale for the i'th IP-Adapter.
+        self.scales = scales
+
+        # The IP-Adapter masks.
+        # self._masks_by_seq_len[s] contains the spatial masks for the downsampling level with query sequence length of
+        # s. It has shape (batch_size, num_ip_images, query_seq_len, 1). The masks have values of 1.0 for included
+        # regions and 0.0 for excluded regions.
+        self._masks_by_seq_len = self._prepare_masks(masks, max_downscale_factor, device, dtype)
+
+    def _prepare_masks(
+        self, masks: list[torch.Tensor], max_downscale_factor: int, device: torch.device, dtype: torch.dtype
+    ) -> dict[int, torch.Tensor]:
+        """Prepare the masks for the IP-Adapter attention."""
+        # Concatenate the masks so that they can be processed more efficiently.
+        mask_tensor = torch.cat(masks, dim=1)
+
+        mask_tensor = mask_tensor.to(device=device, dtype=dtype)
+
+        masks_by_seq_len: dict[int, torch.Tensor] = {}
+
+        # Downsample the spatial dimensions by factors of 2 until max_downscale_factor is reached.
+        downscale_factor = 1
+        while downscale_factor <= max_downscale_factor:
+            b, num_ip_adapters, h, w = mask_tensor.shape
+            # Assert that the batch size is 1, because I haven't thought through batch handling for this feature yet.
+            assert b == 1
+
+            # The IP-Adapters are applied in the cross-attention layers, where the query sequence length is the h * w of
+            # the spatial features.
+            query_seq_len = h * w
+
+            masks_by_seq_len[query_seq_len] = mask_tensor.view((b, num_ip_adapters, -1, 1))
+
+            downscale_factor *= 2
+            if downscale_factor <= max_downscale_factor:
+                # We use max pooling because we downscale to a pretty low resolution, so we don't want small mask
+                # regions to be lost entirely.
+                #
+                # ceil_mode=True is set to mirror the downsampling behavior of SD and SDXL.
+                #
+                # TODO(ryand): In the future, we may want to experiment with other downsampling methods.
+                mask_tensor = torch.nn.functional.max_pool2d(mask_tensor, kernel_size=2, stride=2, ceil_mode=True)
+
+        return masks_by_seq_len
+
+    def get_masks(self, query_seq_len: int) -> torch.Tensor:
+        """Get the mask for the given query sequence length."""
+        return self._masks_by_seq_len[query_seq_len]

invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py

@@ -0,0 +1,105 @@
+import torch
+import torch.nn.functional as F
+
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
+    TextConditioningRegions,
+)
+
+
+class RegionalPromptData:
+    """A class to manage the prompt data for regional conditioning."""
+
+    def __init__(
+        self,
+        regions: list[TextConditioningRegions],
+        device: torch.device,
+        dtype: torch.dtype,
+        max_downscale_factor: int = 8,
+    ):
+        """Initialize a `RegionalPromptData` object.
+        Args:
+            regions (list[TextConditioningRegions]): regions[i] contains the prompt regions for the i'th sample in the
+                batch.
+            device (torch.device): The device to use for the attention masks.
+            dtype (torch.dtype): The data type to use for the attention masks.
+            max_downscale_factor: Spatial masks will be prepared for downscale factors from 1 to max_downscale_factor
+                in steps of 2x.
+        """
+        self._regions = regions
+        self._device = device
+        self._dtype = dtype
+        # self._spatial_masks_by_seq_len[b][s] contains the spatial masks for the b'th batch sample with a query
+        # sequence length of s.
+        self._spatial_masks_by_seq_len: list[dict[int, torch.Tensor]] = self._prepare_spatial_masks(
+            regions, max_downscale_factor
+        )
+        self._negative_cross_attn_mask_score = -10000.0
+
+    def _prepare_spatial_masks(
+        self, regions: list[TextConditioningRegions], max_downscale_factor: int = 8
+    ) -> list[dict[int, torch.Tensor]]:
+        """Prepare the spatial masks for all downscaling factors."""
+        # batch_masks_by_seq_len[b][s] contains the spatial masks for the b'th batch sample with a query sequence length
+        # of s.
+        batch_sample_masks_by_seq_len: list[dict[int, torch.Tensor]] = []
+
+        for batch_sample_regions in regions:
+            batch_sample_masks_by_seq_len.append({})
+
+            batch_sample_masks = batch_sample_regions.masks.to(device=self._device, dtype=self._dtype)
+
+            # Downsample the spatial dimensions by factors of 2 until max_downscale_factor is reached.
+            downscale_factor = 1
+            while downscale_factor <= max_downscale_factor:
+                b, _num_prompts, h, w = batch_sample_masks.shape
+                assert b == 1
+                query_seq_len = h * w
+
+                batch_sample_masks_by_seq_len[-1][query_seq_len] = batch_sample_masks
+
+                downscale_factor *= 2
+                if downscale_factor <= max_downscale_factor:
+                    # We use max pooling because we downscale to a pretty low resolution, so we don't want small prompt
+                    # regions to be lost entirely.
+                    #
+                    # ceil_mode=True is set to mirror the downsampling behavior of SD and SDXL.
+                    #
+                    # TODO(ryand): In the future, we may want to experiment with other downsampling methods (e.g.
+                    # nearest interpolation), and could potentially use a weighted mask rather than a binary mask.
+                    batch_sample_masks = F.max_pool2d(batch_sample_masks, kernel_size=2, stride=2, ceil_mode=True)
+
+        return batch_sample_masks_by_seq_len
+
+    def get_cross_attn_mask(self, query_seq_len: int, key_seq_len: int) -> torch.Tensor:
+        """Get the cross-attention mask for the given query sequence length.
+        Args:
+            query_seq_len: The length of the flattened spatial features at the current downscaling level.
+            key_seq_len (int): The sequence length of the prompt embeddings (which act as the key in the cross-attention
+                layers). This is most likely equal to the max embedding range end, but we pass it explicitly to be sure.
+        Returns:
+            torch.Tensor: The cross-attention score mask.
+                shape: (batch_size, query_seq_len, key_seq_len).
+                dtype: float
+        """
+        batch_size = len(self._spatial_masks_by_seq_len)
+        batch_spatial_masks = [self._spatial_masks_by_seq_len[b][query_seq_len] for b in range(batch_size)]
+
+        # Create an empty attention mask with the correct shape.
+        attn_mask = torch.zeros((batch_size, query_seq_len, key_seq_len), dtype=self._dtype, device=self._device)
+
+        for batch_idx in range(batch_size):
+            batch_sample_spatial_masks = batch_spatial_masks[batch_idx]
+            batch_sample_regions = self._regions[batch_idx]
+
+            # Flatten the spatial dimensions of the mask by reshaping to (1, num_prompts, query_seq_len, 1).
+            _, num_prompts, _, _ = batch_sample_spatial_masks.shape
+            batch_sample_query_masks = batch_sample_spatial_masks.view((1, num_prompts, query_seq_len, 1))
+
+            for prompt_idx, embedding_range in enumerate(batch_sample_regions.ranges):
+                batch_sample_query_scores = batch_sample_query_masks[0, prompt_idx, :, :].clone()
+                batch_sample_query_mask = batch_sample_query_scores > 0.5
+                batch_sample_query_scores[batch_sample_query_mask] = 0.0
+                batch_sample_query_scores[~batch_sample_query_mask] = self._negative_cross_attn_mask_score
+                attn_mask[batch_idx, :, embedding_range.start : embedding_range.end] = batch_sample_query_scores
+
+        return attn_mask

invokeai/backend/stable_diffusion/diffusion/shared_invokeai_diffusion.py

@@ -1,26 +1,20 @@
 from __future__ import annotations
 
 import math
-from contextlib import contextmanager
 from typing import Any, Callable, Optional, Union
 
 import torch
-from diffusers import UNet2DConditionModel
 from typing_extensions import TypeAlias
 
 from invokeai.app.services.config.config_default import get_config
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
-    ConditioningData,
-    ExtraConditioningInfo,
-    SDXLConditioningInfo,
-)
-
-from .cross_attention_control import (
-    CrossAttentionType,
-    CrossAttnControlContext,
-    SwapCrossAttnContext,
-    setup_cross_attention_control_attention_processors,
+    IPAdapterData,
+    Range,
+    TextConditioningData,
+    TextConditioningRegions,
 )
+from invokeai.backend.stable_diffusion.diffusion.regional_ip_data import RegionalIPData
+from invokeai.backend.stable_diffusion.diffusion.regional_prompt_data import RegionalPromptData
 
 ModelForwardCallback: TypeAlias = Union[
     # x, t, conditioning, Optional[cross-attention kwargs]
@@ -58,31 +52,8 @@ class InvokeAIDiffuserComponent:
         self.conditioning = None
         self.model = model
         self.model_forward_callback = model_forward_callback
-        self.cross_attention_control_context = None
         self.sequential_guidance = config.sequential_guidance
 
-    @contextmanager
-    def custom_attention_context(
-        self,
-        unet: UNet2DConditionModel,
-        extra_conditioning_info: Optional[ExtraConditioningInfo],
-    ):
-        old_attn_processors = unet.attn_processors
-
-        try:
-            self.cross_attention_control_context = CrossAttnControlContext(
-                arguments=extra_conditioning_info.cross_attention_control_args,
-            )
-            setup_cross_attention_control_attention_processors(
-                unet,
-                self.cross_attention_control_context,
-            )
-
-            yield None
-        finally:
-            self.cross_attention_control_context = None
-            unet.set_attn_processor(old_attn_processors)
-
     def do_controlnet_step(
         self,
         control_data,
@@ -90,7 +61,7 @@ class InvokeAIDiffuserComponent:
         timestep: torch.Tensor,
         step_index: int,
         total_step_count: int,
-        conditioning_data,
+        conditioning_data: TextConditioningData,
     ):
         down_block_res_samples, mid_block_res_sample = None, None
 
@@ -123,28 +94,28 @@ class InvokeAIDiffuserComponent:
                 added_cond_kwargs = None
 
                 if cfg_injection:  # only applying ControlNet to conditional instead of in unconditioned
-                    if type(conditioning_data.text_embeddings) is SDXLConditioningInfo:
+                    if conditioning_data.is_sdxl():
                         added_cond_kwargs = {
-                            "text_embeds": conditioning_data.text_embeddings.pooled_embeds,
-                            "time_ids": conditioning_data.text_embeddings.add_time_ids,
+                            "text_embeds": conditioning_data.cond_text.pooled_embeds,
+                            "time_ids": conditioning_data.cond_text.add_time_ids,
                         }
-                    encoder_hidden_states = conditioning_data.text_embeddings.embeds
+                    encoder_hidden_states = conditioning_data.cond_text.embeds
                     encoder_attention_mask = None
                 else:
-                    if type(conditioning_data.text_embeddings) is SDXLConditioningInfo:
+                    if conditioning_data.is_sdxl():
                         added_cond_kwargs = {
                             "text_embeds": torch.cat(
                                 [
                                     # TODO: how to pad? just by zeros? or even truncate?
-                                    conditioning_data.unconditioned_embeddings.pooled_embeds,
-                                    conditioning_data.text_embeddings.pooled_embeds,
+                                    conditioning_data.uncond_text.pooled_embeds,
+                                    conditioning_data.cond_text.pooled_embeds,
                                 ],
                                 dim=0,
                             ),
                             "time_ids": torch.cat(
                                 [
-                                    conditioning_data.unconditioned_embeddings.add_time_ids,
-                                    conditioning_data.text_embeddings.add_time_ids,
+                                    conditioning_data.uncond_text.add_time_ids,
+                                    conditioning_data.cond_text.add_time_ids,
                                 ],
                                 dim=0,
                             ),
@@ -153,8 +124,8 @@ class InvokeAIDiffuserComponent:
                         encoder_hidden_states,
                         encoder_attention_mask,
                     ) = self._concat_conditionings_for_batch(
-                        conditioning_data.unconditioned_embeddings.embeds,
-                        conditioning_data.text_embeddings.embeds,
+                        conditioning_data.uncond_text.embeds,
+                        conditioning_data.cond_text.embeds,
                     )
                 if isinstance(control_datum.weight, list):
                     # if controlnet has multiple weights, use the weight for the current step
@@ -198,24 +169,15 @@ class InvokeAIDiffuserComponent:
         self,
         sample: torch.Tensor,
         timestep: torch.Tensor,
-        conditioning_data: ConditioningData,
+        conditioning_data: TextConditioningData,
+        ip_adapter_data: Optional[list[IPAdapterData]],
         step_index: int,
         total_step_count: int,
         down_block_additional_residuals: Optional[torch.Tensor] = None,  # for ControlNet
         mid_block_additional_residual: Optional[torch.Tensor] = None,  # for ControlNet
         down_intrablock_additional_residuals: Optional[torch.Tensor] = None,  # for T2I-Adapter
     ):
-        cross_attention_control_types_to_do = []
-        if self.cross_attention_control_context is not None:
-            percent_through = step_index / total_step_count
-            cross_attention_control_types_to_do = (
-                self.cross_attention_control_context.get_active_cross_attention_control_types_for_step(percent_through)
-            )
-        wants_cross_attention_control = len(cross_attention_control_types_to_do) > 0
-
-        if wants_cross_attention_control or self.sequential_guidance:
-            # If wants_cross_attention_control is True, we force the sequential mode to be used, because cross-attention
-            # control is currently only supported in sequential mode.
+        if self.sequential_guidance:
             (
                 unconditioned_next_x,
                 conditioned_next_x,
@@ -223,7 +185,9 @@ class InvokeAIDiffuserComponent:
                 x=sample,
                 sigma=timestep,
                 conditioning_data=conditioning_data,
-                cross_attention_control_types_to_do=cross_attention_control_types_to_do,
+                ip_adapter_data=ip_adapter_data,
+                step_index=step_index,
+                total_step_count=total_step_count,
                 down_block_additional_residuals=down_block_additional_residuals,
                 mid_block_additional_residual=mid_block_additional_residual,
                 down_intrablock_additional_residuals=down_intrablock_additional_residuals,
@@ -236,6 +200,9 @@ class InvokeAIDiffuserComponent:
                 x=sample,
                 sigma=timestep,
                 conditioning_data=conditioning_data,
+                ip_adapter_data=ip_adapter_data,
+                step_index=step_index,
+                total_step_count=total_step_count,
                 down_block_additional_residuals=down_block_additional_residuals,
                 mid_block_additional_residual=mid_block_additional_residual,
                 down_intrablock_additional_residuals=down_intrablock_additional_residuals,
@@ -294,53 +261,84 @@ class InvokeAIDiffuserComponent:
 
     def _apply_standard_conditioning(
         self,
-        x,
-        sigma,
-        conditioning_data: ConditioningData,
+        x: torch.Tensor,
+        sigma: torch.Tensor,
+        conditioning_data: TextConditioningData,
+        ip_adapter_data: Optional[list[IPAdapterData]],
+        step_index: int,
+        total_step_count: int,
         down_block_additional_residuals: Optional[torch.Tensor] = None,  # for ControlNet
         mid_block_additional_residual: Optional[torch.Tensor] = None,  # for ControlNet
         down_intrablock_additional_residuals: Optional[torch.Tensor] = None,  # for T2I-Adapter
-    ):
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Runs the conditioned and unconditioned UNet forward passes in a single batch for faster inference speed at
         the cost of higher memory usage.
         """
         x_twice = torch.cat([x] * 2)
         sigma_twice = torch.cat([sigma] * 2)
 
-        cross_attention_kwargs = None
-        if conditioning_data.ip_adapter_conditioning is not None:
+        cross_attention_kwargs = {}
+        if ip_adapter_data is not None:
+            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
             # Note that we 'stack' to produce tensors of shape (batch_size, num_ip_images, seq_len, token_len).
-            cross_attention_kwargs = {
-                "ip_adapter_image_prompt_embeds": [
-                    torch.stack(
-                        [ipa_conditioning.uncond_image_prompt_embeds, ipa_conditioning.cond_image_prompt_embeds]
-                    )
-                    for ipa_conditioning in conditioning_data.ip_adapter_conditioning
-                ]
-            }
+            image_prompt_embeds = [
+                torch.stack([ipa_conditioning.uncond_image_prompt_embeds, ipa_conditioning.cond_image_prompt_embeds])
+                for ipa_conditioning in ip_adapter_conditioning
+            ]
+            scales = [ipa.scale_for_step(step_index, total_step_count) for ipa in ip_adapter_data]
+            ip_masks = [ipa.mask for ipa in ip_adapter_data]
+            regional_ip_data = RegionalIPData(
+                image_prompt_embeds=image_prompt_embeds, scales=scales, masks=ip_masks, dtype=x.dtype, device=x.device
+            )
+            cross_attention_kwargs["regional_ip_data"] = regional_ip_data
 
         added_cond_kwargs = None
-        if type(conditioning_data.text_embeddings) is SDXLConditioningInfo:
+        if conditioning_data.is_sdxl():
             added_cond_kwargs = {
                 "text_embeds": torch.cat(
                     [
                         # TODO: how to pad? just by zeros? or even truncate?
-                        conditioning_data.unconditioned_embeddings.pooled_embeds,
-                        conditioning_data.text_embeddings.pooled_embeds,
+                        conditioning_data.uncond_text.pooled_embeds,
+                        conditioning_data.cond_text.pooled_embeds,
                     ],
                     dim=0,
                 ),
                 "time_ids": torch.cat(
                     [
-                        conditioning_data.unconditioned_embeddings.add_time_ids,
-                        conditioning_data.text_embeddings.add_time_ids,
+                        conditioning_data.uncond_text.add_time_ids,
+                        conditioning_data.cond_text.add_time_ids,
                     ],
                     dim=0,
                 ),
             }
 
+        if conditioning_data.cond_regions is not None or conditioning_data.uncond_regions is not None:
+            # TODO(ryand): We currently initialize RegionalPromptData for every denoising step. The text conditionings
+            # and masks are not changing from step-to-step, so this really only needs to be done once. While this seems
+            # painfully inefficient, the time spent is typically negligible compared to the forward inference pass of
+            # the UNet. The main reason that this hasn't been moved up to eliminate redundancy is that it is slightly
+            # awkward to handle both standard conditioning and sequential conditioning further up the stack.
+            regions = []
+            for c, r in [
+                (conditioning_data.uncond_text, conditioning_data.uncond_regions),
+                (conditioning_data.cond_text, conditioning_data.cond_regions),
+            ]:
+                if r is None:
+                    # Create a dummy mask and range for text conditioning that doesn't have region masks.
+                    _, _, h, w = x.shape
+                    r = TextConditioningRegions(
+                        masks=torch.ones((1, 1, h, w), dtype=x.dtype),
+                        ranges=[Range(start=0, end=c.embeds.shape[1])],
+                    )
+                regions.append(r)
+
+            cross_attention_kwargs["regional_prompt_data"] = RegionalPromptData(
+                regions=regions, device=x.device, dtype=x.dtype
+            )
+            cross_attention_kwargs["percent_through"] = step_index / total_step_count
+
         both_conditionings, encoder_attention_mask = self._concat_conditionings_for_batch(
-            conditioning_data.unconditioned_embeddings.embeds, conditioning_data.text_embeddings.embeds
+            conditioning_data.uncond_text.embeds, conditioning_data.cond_text.embeds
         )
         both_results = self.model_forward_callback(
             x_twice,
@@ -360,8 +358,10 @@ class InvokeAIDiffuserComponent:
         self,
         x: torch.Tensor,
         sigma,
-        conditioning_data: ConditioningData,
-        cross_attention_control_types_to_do: list[CrossAttentionType],
+        conditioning_data: TextConditioningData,
+        ip_adapter_data: Optional[list[IPAdapterData]],
+        step_index: int,
+        total_step_count: int,
         down_block_additional_residuals: Optional[torch.Tensor] = None,  # for ControlNet
         mid_block_additional_residual: Optional[torch.Tensor] = None,  # for ControlNet
         down_intrablock_additional_residuals: Optional[torch.Tensor] = None,  # for T2I-Adapter
@@ -391,53 +391,48 @@ class InvokeAIDiffuserComponent:
         if mid_block_additional_residual is not None:
             uncond_mid_block, cond_mid_block = mid_block_additional_residual.chunk(2)
 
-        # If cross-attention control is enabled, prepare the SwapCrossAttnContext.
-        cross_attn_processor_context = None
-        if self.cross_attention_control_context is not None:
-            # Note that the SwapCrossAttnContext is initialized with an empty list of cross_attention_types_to_do.
-            # This list is empty because cross-attention control is not applied in the unconditioned pass. This field
-            # will be populated before the conditioned pass.
-            cross_attn_processor_context = SwapCrossAttnContext(
-                modified_text_embeddings=self.cross_attention_control_context.arguments.edited_conditioning,
-                index_map=self.cross_attention_control_context.cross_attention_index_map,
-                mask=self.cross_attention_control_context.cross_attention_mask,
-                cross_attention_types_to_do=[],
-            )
-
         #####################
         # Unconditioned pass
         #####################
 
-        cross_attention_kwargs = None
+        cross_attention_kwargs = {}
 
         # Prepare IP-Adapter cross-attention kwargs for the unconditioned pass.
-        if conditioning_data.ip_adapter_conditioning is not None:
+        if ip_adapter_data is not None:
+            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
             # Note that we 'unsqueeze' to produce tensors of shape (batch_size=1, num_ip_images, seq_len, token_len).
-            cross_attention_kwargs = {
-                "ip_adapter_image_prompt_embeds": [
-                    torch.unsqueeze(ipa_conditioning.uncond_image_prompt_embeds, dim=0)
-                    for ipa_conditioning in conditioning_data.ip_adapter_conditioning
-                ]
-            }
+            image_prompt_embeds = [
+                torch.unsqueeze(ipa_conditioning.uncond_image_prompt_embeds, dim=0)
+                for ipa_conditioning in ip_adapter_conditioning
+            ]
 
-        # Prepare cross-attention control kwargs for the unconditioned pass.
-        if cross_attn_processor_context is not None:
-            cross_attention_kwargs = {"swap_cross_attn_context": cross_attn_processor_context}
+            scales = [ipa.scale_for_step(step_index, total_step_count) for ipa in ip_adapter_data]
+            ip_masks = [ipa.mask for ipa in ip_adapter_data]
+            regional_ip_data = RegionalIPData(
+                image_prompt_embeds=image_prompt_embeds, scales=scales, masks=ip_masks, dtype=x.dtype, device=x.device
+            )
+            cross_attention_kwargs["regional_ip_data"] = regional_ip_data
 
         # Prepare SDXL conditioning kwargs for the unconditioned pass.
         added_cond_kwargs = None
-        is_sdxl = type(conditioning_data.text_embeddings) is SDXLConditioningInfo
-        if is_sdxl:
+        if conditioning_data.is_sdxl():
             added_cond_kwargs = {
-                "text_embeds": conditioning_data.unconditioned_embeddings.pooled_embeds,
-                "time_ids": conditioning_data.unconditioned_embeddings.add_time_ids,
+                "text_embeds": conditioning_data.uncond_text.pooled_embeds,
+                "time_ids": conditioning_data.uncond_text.add_time_ids,
             }
 
+        # Prepare prompt regions for the unconditioned pass.
+        if conditioning_data.uncond_regions is not None:
+            cross_attention_kwargs["regional_prompt_data"] = RegionalPromptData(
+                regions=[conditioning_data.uncond_regions], device=x.device, dtype=x.dtype
+            )
+            cross_attention_kwargs["percent_through"] = step_index / total_step_count
+
         # Run unconditioned UNet denoising (i.e. negative prompt).
         unconditioned_next_x = self.model_forward_callback(
             x,
             sigma,
-            conditioning_data.unconditioned_embeddings.embeds,
+            conditioning_data.uncond_text.embeds,
             cross_attention_kwargs=cross_attention_kwargs,
             down_block_additional_residuals=uncond_down_block,
             mid_block_additional_residual=uncond_mid_block,
@@ -449,36 +444,43 @@ class InvokeAIDiffuserComponent:
         # Conditioned pass
         ###################
 
-        cross_attention_kwargs = None
+        cross_attention_kwargs = {}
 
-        # Prepare IP-Adapter cross-attention kwargs for the conditioned pass.
-        if conditioning_data.ip_adapter_conditioning is not None:
+        if ip_adapter_data is not None:
+            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
             # Note that we 'unsqueeze' to produce tensors of shape (batch_size=1, num_ip_images, seq_len, token_len).
-            cross_attention_kwargs = {
-                "ip_adapter_image_prompt_embeds": [
-                    torch.unsqueeze(ipa_conditioning.cond_image_prompt_embeds, dim=0)
-                    for ipa_conditioning in conditioning_data.ip_adapter_conditioning
-                ]
-            }
+            image_prompt_embeds = [
+                torch.unsqueeze(ipa_conditioning.cond_image_prompt_embeds, dim=0)
+                for ipa_conditioning in ip_adapter_conditioning
+            ]
 
-        # Prepare cross-attention control kwargs for the conditioned pass.
-        if cross_attn_processor_context is not None:
-            cross_attn_processor_context.cross_attention_types_to_do = cross_attention_control_types_to_do
-            cross_attention_kwargs = {"swap_cross_attn_context": cross_attn_processor_context}
+            scales = [ipa.scale_for_step(step_index, total_step_count) for ipa in ip_adapter_data]
+            ip_masks = [ipa.mask for ipa in ip_adapter_data]
+            regional_ip_data = RegionalIPData(
+                image_prompt_embeds=image_prompt_embeds, scales=scales, masks=ip_masks, dtype=x.dtype, device=x.device
+            )
+            cross_attention_kwargs["regional_ip_data"] = regional_ip_data
 
         # Prepare SDXL conditioning kwargs for the conditioned pass.
         added_cond_kwargs = None
-        if is_sdxl:
+        if conditioning_data.is_sdxl():
             added_cond_kwargs = {
-                "text_embeds": conditioning_data.text_embeddings.pooled_embeds,
-                "time_ids": conditioning_data.text_embeddings.add_time_ids,
+                "text_embeds": conditioning_data.cond_text.pooled_embeds,
+                "time_ids": conditioning_data.cond_text.add_time_ids,
             }
 
+        # Prepare prompt regions for the conditioned pass.
+        if conditioning_data.cond_regions is not None:
+            cross_attention_kwargs["regional_prompt_data"] = RegionalPromptData(
+                regions=[conditioning_data.cond_regions], device=x.device, dtype=x.dtype
+            )
+            cross_attention_kwargs["percent_through"] = step_index / total_step_count
+
         # Run conditioned UNet denoising (i.e. positive prompt).
         conditioned_next_x = self.model_forward_callback(
             x,
             sigma,
-            conditioning_data.text_embeddings.embeds,
+            conditioning_data.cond_text.embeds,
             cross_attention_kwargs=cross_attention_kwargs,
             down_block_additional_residuals=cond_down_block,
             mid_block_additional_residual=cond_mid_block,

invokeai/backend/stable_diffusion/diffusion/unet_attention_patcher.py

@@ -0,0 +1,68 @@
+from contextlib import contextmanager
+from typing import List, Optional, TypedDict
+
+from diffusers.models import UNet2DConditionModel
+
+from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
+from invokeai.backend.stable_diffusion.diffusion.custom_atttention import (
+    CustomAttnProcessor2_0,
+    IPAdapterAttentionWeights,
+)
+
+
+class UNetIPAdapterData(TypedDict):
+    ip_adapter: IPAdapter
+    target_blocks: List[str]
+
+
+class UNetAttentionPatcher:
+    """A class for patching a UNet with CustomAttnProcessor2_0 attention layers."""
+
+    def __init__(self, ip_adapter_data: Optional[List[UNetIPAdapterData]]):
+        self._ip_adapters = ip_adapter_data
+
+    def _prepare_attention_processors(self, unet: UNet2DConditionModel):
+        """Prepare a dict of attention processors that can be injected into a unet, and load the IP-Adapter attention
+        weights into them (if IP-Adapters are being applied).
+        Note that the `unet` param is only used to determine attention block dimensions and naming.
+        """
+        # Construct a dict of attention processors based on the UNet's architecture.
+        attn_procs = {}
+        for idx, name in enumerate(unet.attn_processors.keys()):
+            if name.endswith("attn1.processor") or self._ip_adapters is None:
+                # "attn1" processors do not use IP-Adapters.
+                attn_procs[name] = CustomAttnProcessor2_0()
+            else:
+                # Collect the weights from each IP Adapter for the idx'th attention processor.
+                ip_adapter_attention_weights_collection: list[IPAdapterAttentionWeights] = []
+
+                for ip_adapter in self._ip_adapters:
+                    ip_adapter_weights = ip_adapter["ip_adapter"].attn_weights.get_attention_processor_weights(idx)
+                    skip = True
+                    for block in ip_adapter["target_blocks"]:
+                        if block in name:
+                            skip = False
+                            break
+                    ip_adapter_attention_weights: IPAdapterAttentionWeights = IPAdapterAttentionWeights(
+                        ip_adapter_weights=ip_adapter_weights, skip=skip
+                    )
+                    ip_adapter_attention_weights_collection.append(ip_adapter_attention_weights)
+
+                attn_procs[name] = CustomAttnProcessor2_0(ip_adapter_attention_weights_collection)
+
+        return attn_procs
+
+    @contextmanager
+    def apply_ip_adapter_attention(self, unet: UNet2DConditionModel):
+        """A context manager that patches `unet` with CustomAttnProcessor2_0 attention layers."""
+        attn_procs = self._prepare_attention_processors(unet)
+        orig_attn_processors = unet.attn_processors
+
+        try:
+            # Note to future devs: set_attn_processor(...) does something slightly unexpected - it pops elements from
+            # the passed dict. So, if you wanted to keep the dict for future use, you'd have to make a
+            # moderately-shallow copy of it. E.g. `attn_procs_copy = {k: v for k, v in attn_procs.items()}`.
+            unet.set_attn_processor(attn_procs)
+            yield None
+        finally:
+            unet.set_attn_processor(orig_attn_processors)

invokeai/backend/util/__init__.py

@@ -2,7 +2,6 @@
 Initialization file for invokeai.backend.util
 """
 
-from .devices import choose_precision, choose_torch_device
 from .logging import InvokeAILogger
 from .util import GIG, Chdir, directory_size
 
@@ -11,6 +10,4 @@ __all__ = [
     "directory_size",
     "Chdir",
     "InvokeAILogger",
-    "choose_precision",
-    "choose_torch_device",
 ]

invokeai/backend/util/catch_sigint.py

@@ -0,0 +1,29 @@
+"""
+This module defines a context manager `catch_sigint()` which temporarily replaces
+the sigINT handler defined by the ASGI in order to allow the user to ^C the application
+and shut it down immediately. This was implemented in order to allow the user to interrupt
+slow model hashing during startup.
+
+Use like this:
+
+  from invokeai.backend.util.catch_sigint import catch_sigint
+  with catch_sigint():
+      run_some_hard_to_interrupt_process()
+"""
+
+import signal
+from contextlib import contextmanager
+from typing import Generator
+
+
+def sigint_handler(signum, frame):  # type: ignore
+    signal.signal(signal.SIGINT, signal.SIG_DFL)
+    signal.raise_signal(signal.SIGINT)
+
+
+@contextmanager
+def catch_sigint() -> Generator[None, None, None]:
+    original_handler = signal.getsignal(signal.SIGINT)
+    signal.signal(signal.SIGINT, sigint_handler)
+    yield
+    signal.signal(signal.SIGINT, original_handler)

invokeai/backend/util/devices.py

@@ -1,91 +1,110 @@
-from __future__ import annotations
-
-from contextlib import nullcontext
-from typing import Literal, Optional, Union
+from typing import Dict, Literal, Optional, Union
 
 import torch
-from torch import autocast
+from deprecated import deprecated
 
-from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.config.config_default import get_config
 
+# legacy APIs
+TorchPrecisionNames = Literal["float32", "float16", "bfloat16"]
 CPU_DEVICE = torch.device("cpu")
 CUDA_DEVICE = torch.device("cuda")
 MPS_DEVICE = torch.device("mps")
 
 
+@deprecated("Use TorchDevice.choose_torch_dtype() instead.")  # type: ignore
+def choose_precision(device: torch.device) -> TorchPrecisionNames:
+    """Return the string representation of the recommended torch device."""
+    torch_dtype = TorchDevice.choose_torch_dtype(device)
+    return PRECISION_TO_NAME[torch_dtype]
+
+
+@deprecated("Use TorchDevice.choose_torch_device() instead.")  # type: ignore
 def choose_torch_device() -> torch.device:
-    """Convenience routine for guessing which GPU device to run model on"""
-    config = get_config()
-    if config.device == "auto":
-        if torch.cuda.is_available():
-            return torch.device("cuda")
-        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-            return torch.device("mps")
+    """Return the torch.device to use for accelerated inference."""
+    return TorchDevice.choose_torch_device()
+
+
+@deprecated("Use TorchDevice.choose_torch_dtype() instead.")  # type: ignore
+def torch_dtype(device: torch.device) -> torch.dtype:
+    """Return the torch precision for the recommended torch device."""
+    return TorchDevice.choose_torch_dtype(device)
+
+
+NAME_TO_PRECISION: Dict[TorchPrecisionNames, torch.dtype] = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+}
+PRECISION_TO_NAME: Dict[torch.dtype, TorchPrecisionNames] = {v: k for k, v in NAME_TO_PRECISION.items()}
+
+
+class TorchDevice:
+    """Abstraction layer for torch devices."""
+
+    @classmethod
+    def choose_torch_device(cls) -> torch.device:
+        """Return the torch.device to use for accelerated inference."""
+        app_config = get_config()
+        if app_config.device != "auto":
+            device = torch.device(app_config.device)
+        elif torch.cuda.is_available():
+            device = CUDA_DEVICE
+        elif torch.backends.mps.is_available():
+            device = MPS_DEVICE
         else:
-            return CPU_DEVICE
-    else:
-        return torch.device(config.device)
+            device = CPU_DEVICE
+        return cls.normalize(device)
 
-
-def get_torch_device_name() -> str:
-    device = choose_torch_device()
-    return torch.cuda.get_device_name(device) if device.type == "cuda" else device.type.upper()
-
-
-# We are in transition here from using a single global AppConfig to allowing multiple
-# configurations. It is strongly recommended to pass the app_config to this function.
-def choose_precision(
-    device: torch.device, app_config: Optional[InvokeAIAppConfig] = None
-) -> Literal["float32", "float16", "bfloat16"]:
-    """Return an appropriate precision for the given torch device."""
-    app_config = app_config or get_config()
-    if device.type == "cuda":
-        device_name = torch.cuda.get_device_name(device)
-        if not ("GeForce GTX 1660" in device_name or "GeForce GTX 1650" in device_name):
-            if app_config.precision == "float32":
-                return "float32"
-            elif app_config.precision == "bfloat16":
-                return "bfloat16"
+    @classmethod
+    def choose_torch_dtype(cls, device: Optional[torch.device] = None) -> torch.dtype:
+        """Return the precision to use for accelerated inference."""
+        device = device or cls.choose_torch_device()
+        config = get_config()
+        if device.type == "cuda" and torch.cuda.is_available():
+            device_name = torch.cuda.get_device_name(device)
+            if "GeForce GTX 1660" in device_name or "GeForce GTX 1650" in device_name:
+                # These GPUs have limited support for float16
+                return cls._to_dtype("float32")
+            elif config.precision == "auto":
+                # Default to float16 for CUDA devices
+                return cls._to_dtype("float16")
             else:
-                return "float16"
-    elif device.type == "mps":
-        return "float16"
-    return "float32"
+                # Use the user-defined precision
+                return cls._to_dtype(config.precision)
 
+        elif device.type == "mps" and torch.backends.mps.is_available():
+            if config.precision == "auto":
+                # Default to float16 for MPS devices
+                return cls._to_dtype("float16")
+            else:
+                # Use the user-defined precision
+                return cls._to_dtype(config.precision)
+        # CPU / safe fallback
+        return cls._to_dtype("float32")
 
-# We are in transition here from using a single global AppConfig to allowing multiple
-# configurations. It is strongly recommended to pass the app_config to this function.
-def torch_dtype(
-    device: Optional[torch.device] = None,
-    app_config: Optional[InvokeAIAppConfig] = None,
-) -> torch.dtype:
-    device = device or choose_torch_device()
-    precision = choose_precision(device, app_config)
-    if precision == "float16":
-        return torch.float16
-    if precision == "bfloat16":
-        return torch.bfloat16
-    else:
-        # "auto", "autocast", "float32"
-        return torch.float32
+    @classmethod
+    def get_torch_device_name(cls) -> str:
+        """Return the device name for the current torch device."""
+        device = cls.choose_torch_device()
+        return torch.cuda.get_device_name(device) if device.type == "cuda" else device.type.upper()
 
-
-def choose_autocast(precision):
-    """Returns an autocast context or nullcontext for the given precision string"""
-    # float16 currently requires autocast to avoid errors like:
-    # 'expected scalar type Half but found Float'
-    if precision == "autocast" or precision == "float16":
-        return autocast
-    return nullcontext
-
-
-def normalize_device(device: Union[str, torch.device]) -> torch.device:
-    """Ensure device has a device index defined, if appropriate."""
-    device = torch.device(device)
-    if device.index is None:
-        # cuda might be the only torch backend that currently uses the device index?
-        # I don't see anything like `current_device` for cpu or mps.
-        if device.type == "cuda":
+    @classmethod
+    def normalize(cls, device: Union[str, torch.device]) -> torch.device:
+        """Add the device index to CUDA devices."""
+        device = torch.device(device)
+        if device.index is None and device.type == "cuda" and torch.cuda.is_available():
             device = torch.device(device.type, torch.cuda.current_device())
-    return device
+        return device
+
+    @classmethod
+    def empty_cache(cls) -> None:
+        """Clear the GPU device cache."""
+        if torch.backends.mps.is_available():
+            torch.mps.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    @classmethod
+    def _to_dtype(cls, precision_name: TorchPrecisionNames) -> torch.dtype:
+        return NAME_TO_PRECISION[precision_name]

invokeai/backend/util/mask.py

@@ -0,0 +1,53 @@
+import torch
+
+
+def to_standard_mask_dim(mask: torch.Tensor) -> torch.Tensor:
+    """Standardize the dimensions of a mask tensor.
+
+    Args:
+        mask (torch.Tensor): A mask tensor. The shape can be (1, h, w) or (h, w).
+
+    Returns:
+        torch.Tensor: The output mask tensor. The shape is (1, h, w).
+    """
+    # Get the mask height and width.
+    if mask.ndim == 2:
+        mask = mask.unsqueeze(0)
+    elif mask.ndim == 3 and mask.shape[0] == 1:
+        pass
+    else:
+        raise ValueError(f"Unsupported mask shape: {mask.shape}. Expected (1, h, w) or (h, w).")
+
+    return mask
+
+
+def to_standard_float_mask(mask: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
+    """Standardize the format of a mask tensor.
+
+    Args:
+        mask (torch.Tensor): A mask tensor. The dtype can be any bool, float, or int type. The shape must be (1, h, w)
+            or (h, w).
+
+        out_dtype (torch.dtype): The dtype of the output mask tensor. Must be a float type.
+
+    Returns:
+        torch.Tensor: The output mask tensor. The dtype is out_dtype. The shape is (1, h, w). All values are either 0.0
+            or 1.0.
+    """
+
+    if not out_dtype.is_floating_point:
+        raise ValueError(f"out_dtype must be a float type, but got {out_dtype}")
+
+    mask = to_standard_mask_dim(mask)
+    mask = mask.to(out_dtype)
+
+    # Set masked regions to 1.0.
+    if mask.dtype == torch.bool:
+        mask = mask.to(out_dtype)
+    else:
+        mask = mask.to(out_dtype)
+        mask_region = mask > 0.5
+        mask[mask_region] = 1.0
+        mask[~mask_region] = 0.0
+
+    return mask