WIP - still need to figure out why the keys are wrong.

docs/features/GALLERY.md

@@ -1,92 +0,0 @@
----
-title: InvokeAI Gallery Panel
----
-
-# :material-web: InvokeAI Gallery Panel
-
-## Quick guided walkthrough of the Gallery Panel's features
-
-The Gallery Panel is a fast way to review, find, and make use of images you've
-generated and loaded. The Gallery is divided into Boards. The Uncategorized board is always 
-present but you can create your own for better organization.
-
-![image](../assets/gallery/gallery.png)
-
-### Board Display and Settings
-
-At the very top of the Gallery Panel are the boards disclosure and settings buttons.
-
-![image](../assets/gallery/top_controls.png)
-
-The disclosure button shows the name of the currently selected board and allows you to show and hide the board thumbnails (shown in the image below).
-
-![image](../assets/gallery/board_thumbnails.png)
-
-The settings button opens a list of options.
-
-![image](../assets/gallery/board_settings.png)
-
-- ***Image Size*** this slider lets you control the size of the image previews (images of three different sizes).
-- ***Auto-Switch to New Images*** if you turn this on, whenever a new image is generated, it will automatically be loaded into the current image panel on the Text to Image tab and into the result panel on the [Image to Image](IMG2IMG.md) tab. This will happen invisibly if you are on any other tab when the image is generated.
-- ***Auto-Assign Board on Click*** whenever an image is generated or saved, it always gets put in a board. The board it gets put into is marked with AUTO (image of board marked). Turning on Auto-Assign Board on Click will make whichever board you last selected be the destination when you click Invoke. That means you can click Invoke, select a different board, and then click Invoke again and the two images will be put in two different boards. (bold)It's the board selected when Invoke is clicked that's used, not the board that's selected when the image is finished generating.(bold) Turning this off, enables the Auto-Add Board drop down which lets you set one specific board to always put generated images into. This also enables and disables the Auto-add to this Board menu item described below.
-- ***Always Show Image Size Badge*** this toggles whether to show image sizes for each image preview (show two images, one with sizes shown, one without)
-
-Below these two buttons, you'll see the Search Boards text entry area. You use this to search for specific boards by the name of the board.
-Next to it is the Add Board (+) button which lets you add new boards. Boards can be renamed by clicking on the name of the board under its thumbnail and typing in the new name.
-
-### Board Thumbnail Menu
-
-Each board has a context menu (ctrl+click / right-click).
-
-![image](../assets/gallery/thumbnail_menu.png)
-
-- ***Auto-add to this Board*** if you've disabled Auto-Assign Board on Click in the board settings, you can use this option to set this board to be where new images are put.
-- ***Download Board*** this will add all the images in the board into a zip file and provide a link to it in a notification (image of notification)
-- ***Delete Board*** this will delete the board
-> [!CAUTION]
-> This will delete all the images in the board and the board itself.
-
-### Board Contents
-
-Every board is organized by two tabs, Images and Assets.
-
-![image](../assets/gallery/board_tabs.png)
-
-Images are the Invoke-generated images that are placed into the board. Assets are images that you upload into Invoke to be used as an [Image Prompt](https://support.invoke.ai/support/solutions/articles/151000159340-using-the-image-prompt-adapter-ip-adapter-) or in the [Image to Image](IMG2IMG.md) tab.
-
-### Image Thumbnail Menu
-
-Every image generated by Invoke has its generation information stored as text inside the image file itself. This can be read directly by selecting the image and clicking on the Info button ![image](../assets/gallery/info_button.png) in any of the image result panels. 
-
-Each image also has a context menu (ctrl+click / right-click).
-
-![image](../assets/gallery/image_menu.png)
-
- The options are (items marked with an * will not work with images that lack generation information):
-- ***Open in New Tab*** this will open the image alone in a new browser tab, separate from the Invoke interface.
-- ***Download Image*** this will trigger your browser to download the image.
-- ***Load Workflow **** this will load any workflow settings into the Workflow tab and automatically open it.
-- ***Remix Image **** this will load all of the image's generation information, (bold)excluding its Seed, into the left hand control panel
-- ***Use Prompt **** this will load only the image's text prompts into the left-hand control panel
-- ***Use Seed **** this will load only the image's Seed into the left-hand control panel
-- ***Use All **** this will load all of the image's generation information into the left-hand control panel
-- ***Send to Image to Image*** this will put the image into the left-hand panel in the Image to Image tab ana automatically open it
-- ***Send to Unified Canvas*** This will (bold)replace whatever is already present(bold) in the Unified Canvas tab with the image and automatically open the tab
-- ***Change Board*** this will oipen a small window that will let you move the image to a different board. This is the same as dragging the image to that board's thumbnail.
-- ***Star Image*** this will add the image to the board's list of starred images that are always kept at the top of the gallery. This is the same as clicking on the star on the top right-hand side of the image that appears when you hover over the image with the mouse
-- ***Delete Image*** this will delete the image from the board
-> [!CAUTION] 
-> This will delete the image entirely from Invoke.
-
-## Summary
-
-This walkthrough only covers the Gallery interface and Boards. Actually generating images is handled by [Prompts](PROMPTS.md), the [Image to Image](IMG2IMG.md) tab, and the [Unified Canvas](UNIFIED_CANVAS.md).
-
-## Acknowledgements
-
-A huge shout-out to the core team working to make the Web GUI a reality,
-including [psychedelicious](https://github.com/psychedelicious),
-[Kyle0654](https://github.com/Kyle0654) and
-[blessedcoolant](https://github.com/blessedcoolant).
-[hipsterusername](https://github.com/hipsterusername) was the team's unofficial
-cheerleader and added tooltips/docs.

docs/features/PROMPTS.md

@@ -108,6 +108,40 @@ Can be used with .and():
 Each will give you different results - try them out and see what you prefer!
 
 
+
+### Cross-Attention Control ('prompt2prompt')
+
+Sometimes an image you generate is almost right, and you just want to change one
+detail without affecting the rest. You could use a photo editor and inpainting
+to overpaint the area, but that's a pain. Here's where `prompt2prompt` comes in
+handy.
+
+Generate an image with a given prompt, record the seed of the image, and then
+use the `prompt2prompt` syntax to substitute words in the original prompt for
+words in a new prompt. This works for `img2img` as well.
+
+For example, consider the prompt `a cat.swap(dog) playing with a ball in the forest`. Normally, because the words interact with each other when doing a stable diffusion image generation, these two prompts would generate different compositions:
+  - `a cat playing with a ball in the forest`
+  - `a dog playing with a ball in the forest`
+
+| `a cat playing with a ball in the forest` | `a dog playing with a ball in the forest` |
+| --- | --- |
+| img | img |
+
+
+      - For multiple word swaps, use parentheses: `a (fluffy cat).swap(barking dog) playing with a ball in the forest`.
+      - To swap a comma, use quotes: `a ("fluffy, grey cat").swap("big, barking dog") playing with a ball in the forest`.
+- Supports options `t_start` and `t_end` (each 0-1) loosely corresponding to (bloc97's)[(https://github.com/bloc97/CrossAttentionControl)] `prompt_edit_tokens_start/_end` but with the math swapped to make it easier to
+  intuitively understand. `t_start` and `t_end` are used to control on which steps cross-attention control should run. With the default values `t_start=0` and `t_end=1`, cross-attention control is active on every step of image generation. Other values can be used to turn cross-attention control off for part of the image generation process.
+    - For example, if doing a diffusion with 10 steps for the prompt is `a cat.swap(dog, t_start=0.3, t_end=1.0) playing with a ball in the forest`, the first 3 steps will be run as `a cat playing with a ball in the forest`, while the last 7 steps will run as `a dog playing with a ball in the forest`, but the pixels that represent `dog` will be locked to the pixels that would have represented `cat` if the `cat` prompt had been used instead.
+    - Conversely, for `a cat.swap(dog, t_start=0, t_end=0.7) playing with a ball in the forest`, the first 7 steps will run as `a dog playing with a ball in the forest` with the pixels that represent `dog` locked to the same pixels that would have represented `cat` if the `cat` prompt was being used instead. The final 3 steps will just run `a cat playing with a ball in the forest`.
+    > For img2img, the step sequence does not start at 0 but instead at `(1.0-strength)` - so if the img2img `strength` is `0.7`, `t_start` and `t_end` must both be greater than `0.3` (`1.0-0.7`) to have any effect.
+
+Prompt2prompt `.swap()` is not compatible with xformers, which will be temporarily disabled when doing a `.swap()` - so you should expect to use more VRAM and run slower that with xformers enabled.
+
+The `prompt2prompt` code is based off
+[bloc97's colab](https://github.com/bloc97/CrossAttentionControl).
+
 ### Escaping parentheses and speech marks 
 
 If the model you are using has parentheses () or speech marks "" as part of its

docs/features/WEB.md

@@ -54,7 +54,7 @@ main sections:
    of buttons at the top lets you modify and manipulate the image in
    various ways.
 
-3. A **gallery** section on the right that contains a history of the images you
+3. A **gallery** section on the left that contains a history of the images you
    have generated. These images are read and written to the directory specified
    in the `INVOKEAIROOT/invokeai.yaml` initialization file, usually a directory
    named `outputs` in `INVOKEAIROOT`.

docs/help/FAQ.md

@@ -18,47 +18,12 @@ Note that any releases marked as _pre-release_ are in a beta state. You may expe
 
 The Model Manager tab in the UI provides a few ways to install models, including using your already-downloaded models. You'll see a popup directing you there on first startup. For more information, see the [model install docs].
 
-## Missing models after updating to v4
-
-If you find some models are missing after updating to v4, it's likely they weren't correctly registered before the update and didn't get picked up in the migration.
-
-You can use the `Scan Folder` tab in the Model Manager UI to fix this. The models will either be in the old, now-unused `autoimport` folder, or your `models` folder.
-
-- Find and copy your install's old `autoimport` folder path, install the main install folder.
-- Go to the Model Manager and click `Scan Folder`.
-- Paste the path and scan.
-- IMPORTANT: Uncheck `Inplace install`.
-- Click `Install All` to install all found models, or just install the models you want.
-
-Next, find and copy your install's `models` folder path (this could be your custom models folder path, or the `models` folder inside the main install folder).
-
-Follow the same steps to scan and import the missing models.
-
 ## Slow generation
 
 - Check the [system requirements] to ensure that your system is capable of generating images.
 - Check the `ram` setting in `invokeai.yaml`. This setting tells Invoke how much of your system RAM can be used to cache models. Having this too high or too low can slow things down. That said, it's generally safest to not set this at all and instead let Invoke manage it.
 - Check the `vram` setting in `invokeai.yaml`. This setting tells Invoke how much of your GPU VRAM can be used to cache models. Counter-intuitively, if this setting is too high, Invoke will need to do a lot of shuffling of models as it juggles the VRAM cache and the currently-loaded model. The default value of 0.25 is generally works well for GPUs without 16GB or more VRAM. Even on a 24GB card, the default works well.
 - Check that your generations are happening on your GPU (if you have one). InvokeAI will log what is being used for generation upon startup. If your GPU isn't used, re-install to ensure the correct versions of torch get installed.
-- If you are on Windows, you may have exceeded your GPU's VRAM capacity and are using slower [shared GPU memory](#shared-gpu-memory-windows). There's a guide to opt out of this behaviour in the linked FAQ entry.
-
-## Shared GPU Memory (Windows)
-
-!!! tip "Nvidia GPUs with driver 536.40"
-
-    This only applies to current Nvidia cards with driver 536.40 or later, released in June 2023.
-
-When the GPU doesn't have enough VRAM for a task, Windows is able to allocate some of its CPU RAM to the GPU. This is much slower than VRAM, but it does allow the system to generate when it otherwise might no have enough VRAM.
-
-When shared GPU memory is used, generation slows down dramatically - but at least it doesn't crash.
-
-If you'd like to opt out of this behavior and instead get an error when you exceed your GPU's VRAM, follow [this guide from Nvidia](https://nvidia.custhelp.com/app/answers/detail/a_id/5490).
-
-Here's how to get the python path required in the linked guide:
-
-- Run `invoke.bat`.
-- Select option 2 for developer console.
-- At least one python path will be printed. Copy the path that includes your invoke installation directory (typically the first).
 
 ## Installer cannot find python (Windows)
 

docs/installation/010_INSTALL_AUTOMATED.md

@@ -44,7 +44,7 @@ The installation process is simple, with a few prompts:
 
 - Select the version to install. Unless you have a specific reason to install a specific version, select the default (the latest version).
 - Select location for the install. Be sure you have enough space in this folder for the base application, as described in the [installation requirements].
-- Select a GPU device.
+- Select a GPU device. If you are unsure, you can let the installer figure it out.
 
 !!! info "Slow Installation"
 

docs/installation/020_INSTALL_MANUAL.md

@@ -6,7 +6,11 @@
 
 ## Introduction
 
-InvokeAI is distributed as a python package on PyPI, installable with `pip`. There are a few things that are handled by the installer and launcher that you'll need to manage manually, described in this guide.
+!!! tip "Conda"
+
+    As of InvokeAI v2.3.0 installation using the `conda` package manager is no longer being supported. It will likely still work, but we are not testing this installation method.
+
+InvokeAI is distributed as a python package on PyPI, installable with `pip`. There are a few things that are handled by the installer that you'll need to manage manually, described in this guide.
 
 ### Requirements
 
@@ -36,11 +40,11 @@ Before you start, go through the [installation requirements].
 
 1. Enter the root (invokeai) directory and create a virtual Python environment within it named `.venv`.
 
-    !!! warning "Virtual Environment Location"
+    !!! info "Virtual Environment Location"
 
         While you may create the virtual environment anywhere in the file system, we recommend that you create it within the root directory as shown here. This allows the application to automatically detect its data directories.
 
-        If you choose a different location for the venv, then you _must_ set the `INVOKEAI_ROOT` environment variable or specify the root directory using the `--root` CLI arg.
+        If you choose a different location for the venv, then you must set the `INVOKEAI_ROOT` environment variable or pass the directory using the `--root` CLI arg.
 
     ```terminal
     cd $INVOKEAI_ROOT
@@ -77,23 +81,31 @@ Before you start, go through the [installation requirements].
     python3 -m pip install --upgrade pip
     ```
 
-1. Install the InvokeAI Package. The base command is `pip install InvokeAI --use-pep517`, but you may need to change this depending on your system and the desired features.
+1. Install the InvokeAI Package. The `--extra-index-url` option is used to select the correct `torch` backend:
 
-    - You may need to provide an [extra index URL]. Select your platform configuration using [this tool on the PyTorch website]. Copy the `--extra-index-url` string from this and append it to your install command.
+    === "CUDA (NVidia)"
 
-        !!! example "Install with an extra index URL"
+         ```bash
+         pip install "InvokeAI[xformers]" --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu121
+         ```
 
-            ```bash
-            pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu121
-            ```
+    === "ROCm (AMD)"
 
-    - If you have a CUDA GPU and want to install with `xformers`, you need to add an option to the package name. Note that `xformers` is not necessary. PyTorch includes an implementation of the SDP attention algorithm with the same performance.
+         ```bash
+         pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.6
+         ```
 
-        !!! example "Install with `xformers`"
+    === "CPU (Intel Macs & non-GPU systems)"
 
-            ```bash
-            pip install "InvokeAI[xformers]" --use-pep517
-            ```
+         ```bash
+         pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/cpu
+         ```
+
+    === "MPS (Apple Silicon)"
+
+         ```bash
+         pip install InvokeAI --use-pep517
+         ```
 
 1. Deactivate and reactivate your runtime directory so that the invokeai-specific commands become available in the environment:
 
@@ -114,6 +126,37 @@ Before you start, go through the [installation requirements].
 
     Run `invokeai-web` to start the UI. You must activate the virtual environment before running the app.
 
-    !!! warning
+    If the virtual environment you selected is NOT inside `INVOKEAI_ROOT`, then you must specify the path to the root directory by adding
+    `--root_dir \path\to\invokeai`.
 
-        If the virtual environment is _not_ inside the root directory, then you _must_ specify the path to the root directory with `--root_dir \path\to\invokeai` or the `INVOKEAI_ROOT` environment variable.
+    !!! tip
+
+        You can permanently set the location of the runtime directory
+        by setting the environment variable `INVOKEAI_ROOT` to the
+        path of the directory. As mentioned previously, this is
+        recommended if your virtual environment is located outside of
+        your runtime directory.
+
+## Unsupported Conda Install
+
+Congratulations, you found the "secret" Conda installation instructions. If you really **really** want to use Conda with InvokeAI, you can do so using this unsupported recipe:
+
+```sh
+mkdir ~/invokeai
+conda create -n invokeai python=3.11
+conda activate invokeai
+# Adjust this as described above for the appropriate torch backend
+pip install InvokeAI[xformers] --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu121
+invokeai-web --root ~/invokeai
+```
+
+The `pip install` command shown in this recipe is for Linux/Windows
+systems with an NVIDIA GPU. See step (6) above for the command to use
+with other platforms/GPU combinations. If you don't wish to pass the
+`--root` argument to `invokeai` with each launch, you may set the
+environment variable `INVOKEAI_ROOT` to point to the installation directory.
+
+Note that if you run into problems with the Conda installation, the InvokeAI
+staff will **not** be able to help you out. Caveat Emptor!
+
+[installation requirements]: INSTALL_REQUIREMENTS.md

docs/installation/INSTALL_DEVELOPMENT.md

@@ -23,7 +23,6 @@ If you have an interest in how InvokeAI works, or you would like to add features
 
 1. [Fork and clone] the [InvokeAI repo].
 1. Follow the [manual installation] docs to create a new virtual environment for the development install.
-   - Create a new folder outside the repo root for the installation and create the venv inside that folder.
    - When installing the InvokeAI package, add `-e` to the command so you get an [editable install].
 1. Install the [frontend dev toolchain] and do a production build of the UI as described.
 1. You can now run the app as described in the [manual installation] docs.
@@ -33,5 +32,5 @@ As described in the [frontend dev toolchain] docs, you can run the UI using a de
 [Fork and clone]: https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/fork-a-repo
 [InvokeAI repo]: https://github.com/invoke-ai/InvokeAI
 [frontend dev toolchain]: ../contributing/frontend/OVERVIEW.md
-[manual installation]: ./020_INSTALL_MANUAL.md
+[manual installation]: installation/020_INSTALL_MANUAL.md
 [editable install]: https://pip.pypa.io/en/latest/cli/pip_install/#cmdoption-e

installer/lib/installer.py

@@ -3,7 +3,6 @@
 InvokeAI installer script
 """
 
-import locale
 import os
 import platform
 import re
@@ -317,9 +316,7 @@ def upgrade_pip(venv_path: Path) -> str | None:
     python = str(venv_path.expanduser().resolve() / python)
 
     try:
-        result = subprocess.check_output([python, "-m", "pip", "install", "--upgrade", "pip"]).decode(
-            encoding=locale.getpreferredencoding()
-        )
+        result = subprocess.check_output([python, "-m", "pip", "install", "--upgrade", "pip"]).decode()
     except subprocess.CalledProcessError as e:
         print(e)
         result = None
@@ -407,29 +404,22 @@ def get_torch_source() -> Tuple[str | None, str | None]:
     # device can be one of: "cuda", "rocm", "cpu", "cuda_and_dml, autodetect"
     device = select_gpu()
 
-    # The correct extra index URLs for torch are inconsistent, see https://pytorch.org/get-started/locally/#start-locally
-
     url = None
-    optional_modules: str | None = None
+    optional_modules = "[onnx]"
     if OS == "Linux":
         if device.value == "rocm":
             url = "https://download.pytorch.org/whl/rocm5.6"
         elif device.value == "cpu":
             url = "https://download.pytorch.org/whl/cpu"
-        elif device.value == "cuda":
-            # CUDA uses the default PyPi index
-            optional_modules = "[xformers,onnx-cuda]"
+
     elif OS == "Windows":
         if device.value == "cuda":
             url = "https://download.pytorch.org/whl/cu121"
             optional_modules = "[xformers,onnx-cuda]"
-        elif device.value == "cpu":
-            # CPU  uses the default PyPi index, no optional modules
-            pass
-    elif OS == "Darwin":
-        # macOS uses the default PyPi index, no optional modules
-        pass
+        if device.value == "cuda_and_dml":
+            url = "https://download.pytorch.org/whl/cu121"
+            optional_modules = "[xformers,onnx-directml]"
 
-    # Fall back to defaults
+    # in all other cases, Torch wheels should be coming from PyPi as of Torch 1.13
 
     return (url, optional_modules)

installer/lib/messages.py

@@ -207,8 +207,10 @@ def dest_path(dest: Optional[str | Path] = None) -> Path | None:
 
 class GpuType(Enum):
     CUDA = "cuda"
+    CUDA_AND_DML = "cuda_and_dml"
     ROCM = "rocm"
     CPU = "cpu"
+    AUTODETECT = "autodetect"
 
 
 def select_gpu() -> GpuType:
@@ -224,6 +226,10 @@ def select_gpu() -> GpuType:
         "an [gold1 b]NVIDIA[/] GPU (using CUDA™)",
         GpuType.CUDA,
     )
+    nvidia_with_dml = (
+        "an [gold1 b]NVIDIA[/] GPU (using CUDA™, and DirectML™ for ONNX) -- ALPHA",
+        GpuType.CUDA_AND_DML,
+    )
     amd = (
         "an [gold1 b]AMD[/] GPU (using ROCm™)",
         GpuType.ROCM,
@@ -232,19 +238,27 @@ def select_gpu() -> GpuType:
         "Do not install any GPU support, use CPU for generation (slow)",
         GpuType.CPU,
     )
+    autodetect = (
+        "I'm not sure what to choose",
+        GpuType.AUTODETECT,
+    )
 
     options = []
     if OS == "Windows":
-        options = [nvidia, cpu]
+        options = [nvidia, nvidia_with_dml, cpu]
     if OS == "Linux":
         options = [nvidia, amd, cpu]
     elif OS == "Darwin":
         options = [cpu]
+        # future CoreML?
 
     if len(options) == 1:
         print(f'Your platform [gold1]{OS}-{ARCH}[/] only supports the "{options[0][1]}" driver. Proceeding with that.')
         return options[0][1]
 
+    # "I don't know" is always added the last option
+    options.append(autodetect)  # type: ignore
+
     options = {str(i): opt for i, opt in enumerate(options, 1)}
 
     console.rule(":space_invader: GPU (Graphics Card) selection :space_invader:")
@@ -278,6 +292,11 @@ def select_gpu() -> GpuType:
         ),
     )
 
+    if options[choice][1] is GpuType.AUTODETECT:
+        console.print(
+            "No problem. We will install CUDA support first :crossed_fingers: If Invoke does not detect a GPU, please re-run the installer and select one of the other GPU types."
+        )
+
     return options[choice][1]
 
 

invokeai/app/api/routers/app_info.py

@@ -12,7 +12,7 @@ from pydantic import BaseModel, Field
 
 from invokeai.app.invocations.upscale import ESRGAN_MODELS
 from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
-from invokeai.backend.image_util.infill_methods.patchmatch import PatchMatch
+from invokeai.backend.image_util.patchmatch import PatchMatch
 from invokeai.backend.image_util.safety_checker import SafetyChecker
 from invokeai.backend.util.logging import logging
 from invokeai.version import __version__
@@ -100,7 +100,7 @@ async def get_app_deps() -> AppDependencyVersions:
 
 @app_router.get("/config", operation_id="get_config", status_code=200, response_model=AppConfig)
 async def get_config() -> AppConfig:
-    infill_methods = ["tile", "lama", "cv2", "color"]  # TODO: add mosaic back
+    infill_methods = ["tile", "lama", "cv2"]
     if PatchMatch.patchmatch_available():
         infill_methods.append("patchmatch")
 

invokeai/app/api/routers/model_manager.py

@@ -219,13 +219,28 @@ async def scan_for_models(
         non_core_model_paths = [p for p in found_model_paths if not p.is_relative_to(core_models_path)]
 
         installed_models = ApiDependencies.invoker.services.model_manager.store.search_by_attr()
+        resolved_installed_model_paths: list[str] = []
+        installed_model_sources: list[str] = []
+
+        # This call lists all installed models.
+        for model in installed_models:
+            path = pathlib.Path(model.path)
+            # If the model has a source, we need to add it to the list of installed sources.
+            if model.source:
+                installed_model_sources.append(model.source)
+            # If the path is not absolute, that means it is in the app models directory, and we need to join it with
+            # the models path before resolving.
+            if not path.is_absolute():
+                resolved_installed_model_paths.append(str(pathlib.Path(models_path, path).resolve()))
+                continue
+            resolved_installed_model_paths.append(str(path.resolve()))
 
         scan_results: list[FoundModel] = []
 
-        # Check if the model is installed by comparing paths, appending to the scan result.
+        # Check if the model is installed by comparing the resolved paths, appending to the scan result.
         for p in non_core_model_paths:
             path = str(p)
-            is_installed = any(str(models_path / m.path) == path for m in installed_models)
+            is_installed = path in resolved_installed_model_paths or path in installed_model_sources
             found_model = FoundModel(path=path, is_installed=is_installed)
             scan_results.append(found_model)
     except Exception as e:
@@ -599,8 +614,8 @@ async def convert_model(
     The return value is the model configuration for the converted model.
     """
     model_manager = ApiDependencies.invoker.services.model_manager
-    loader = model_manager.load
     logger = ApiDependencies.invoker.services.logger
+    loader = ApiDependencies.invoker.services.model_manager.load
     store = ApiDependencies.invoker.services.model_manager.store
     installer = ApiDependencies.invoker.services.model_manager.install
 
@@ -615,13 +630,7 @@ async def convert_model(
         raise HTTPException(400, f"The model with key {key} is not a main checkpoint model.")
 
     # loading the model will convert it into a cached diffusers file
-    try:
-        cc_size = loader.convert_cache.max_size
-        if cc_size == 0:  # temporary set the convert cache to a positive number so that cached model is written
-            loader._convert_cache.max_size = 1.0
-        loader.load_model(model_config, submodel_type=SubModelType.Scheduler)
-    finally:
-        loader._convert_cache.max_size = cc_size
+    model_manager.load.load_model(model_config, submodel_type=SubModelType.Scheduler)
 
     # Get the path of the converted model from the loader
     cache_path = loader.convert_cache.cache_path(key)

invokeai/app/api_app.py

@@ -28,7 +28,7 @@ from invokeai.app.api.no_cache_staticfiles import NoCacheStaticFiles
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.services.config.config_default import get_config
 from invokeai.app.services.session_processor.session_processor_common import ProgressImage
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import get_torch_device_name
 
 from ..backend.util.logging import InvokeAILogger
 from .api.dependencies import ApiDependencies
@@ -63,7 +63,7 @@ logger = InvokeAILogger.get_logger(config=app_config)
 mimetypes.add_type("application/javascript", ".js")
 mimetypes.add_type("text/css", ".css")
 
-torch_device_name = TorchDevice.get_torch_device_name()
+torch_device_name = get_torch_device_name()
 logger.info(f"Using torch device: {torch_device_name}")
 
 

invokeai/app/invocations/compel.py

@@ -5,26 +5,20 @@ from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
-from invokeai.app.invocations.fields import (
-    ConditioningField,
-    FieldDescriptions,
-    Input,
-    InputField,
-    OutputField,
-    TensorField,
-    UIComponent,
-)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIComponent
 from invokeai.app.invocations.primitives import ConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.ti_utils import generate_ti_list
-from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.lora_model_patcher import LoraModelPatcher
+from invokeai.backend.lora_model_raw import LoRAModelRaw
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
     ConditioningFieldData,
+    ExtraConditioningInfo,
     SDXLConditioningInfo,
 )
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import torch_dtype
 
 from .baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
 from .model import CLIPField
@@ -43,7 +37,7 @@ from .model import CLIPField
     title="Prompt",
     tags=["prompt", "compel"],
     category="conditioning",
-    version="1.2.0",
+    version="1.1.1",
 )
 class CompelInvocation(BaseInvocation):
     """Parse prompt using compel package to conditioning."""
@@ -58,9 +52,6 @@ class CompelInvocation(BaseInvocation):
         description=FieldDescriptions.clip,
         input=Input.Connection,
     )
-    mask: Optional[TensorField] = InputField(
-        default=None, description="A mask defining the region that this conditioning prompt applies to."
-    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
@@ -90,7 +81,8 @@ class CompelInvocation(BaseInvocation):
             ),
             text_encoder_info as text_encoder,
             # Apply the LoRA after text_encoder has been moved to its target device for faster patching.
-            ModelPatcher.apply_lora_text_encoder(text_encoder, _lora_loader()),
+            # ModelPatcher.apply_lora_text_encoder(text_encoder, _lora_loader()),
+            LoraModelPatcher.apply_lora_to_text_encoder(text_encoder, _lora_loader(), "text_encoder"),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
             ModelPatcher.apply_clip_skip(text_encoder_model, self.clip.skipped_layers),
         ):
@@ -99,7 +91,7 @@ class CompelInvocation(BaseInvocation):
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
                 textual_inversion_manager=ti_manager,
-                dtype_for_device_getter=TorchDevice.choose_torch_dtype,
+                dtype_for_device_getter=torch_dtype,
                 truncate_long_prompts=False,
             )
 
@@ -108,19 +100,27 @@ class CompelInvocation(BaseInvocation):
             if context.config.get().log_tokenization:
                 log_tokenization_for_conjunction(conjunction, tokenizer)
 
-            c, _options = compel.build_conditioning_tensor_for_conjunction(conjunction)
+            c, options = compel.build_conditioning_tensor_for_conjunction(conjunction)
+
+            ec = ExtraConditioningInfo(
+                tokens_count_including_eos_bos=get_max_token_count(tokenizer, conjunction),
+                cross_attention_control_args=options.get("cross_attention_control", None),
+            )
 
         c = c.detach().to("cpu")
 
-        conditioning_data = ConditioningFieldData(conditionings=[BasicConditioningInfo(embeds=c)])
+        conditioning_data = ConditioningFieldData(
+            conditionings=[
+                BasicConditioningInfo(
+                    embeds=c,
+                    extra_conditioning=ec,
+                )
+            ]
+        )
 
         conditioning_name = context.conditioning.save(conditioning_data)
-        return ConditioningOutput(
-            conditioning=ConditioningField(
-                conditioning_name=conditioning_name,
-                mask=self.mask,
-            )
-        )
+
+        return ConditioningOutput.build(conditioning_name)
 
 
 class SDXLPromptInvocationBase:
@@ -134,7 +134,7 @@ class SDXLPromptInvocationBase:
         get_pooled: bool,
         lora_prefix: str,
         zero_on_empty: bool,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[ExtraConditioningInfo]]:
         tokenizer_info = context.models.load(clip_field.tokenizer)
         tokenizer_model = tokenizer_info.model
         assert isinstance(tokenizer_model, CLIPTokenizer)
@@ -161,7 +161,7 @@ class SDXLPromptInvocationBase:
                 )
             else:
                 c_pooled = None
-            return c, c_pooled
+            return c, c_pooled, None
 
         def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in clip_field.loras:
@@ -183,7 +183,8 @@ class SDXLPromptInvocationBase:
             ),
             text_encoder_info as text_encoder,
             # Apply the LoRA after text_encoder has been moved to its target device for faster patching.
-            ModelPatcher.apply_lora(text_encoder, _lora_loader(), lora_prefix),
+            # ModelPatcher.apply_lora(text_encoder, _lora_loader(), lora_prefix),
+            LoraModelPatcher.apply_lora_to_text_encoder(text_encoder, _lora_loader(), lora_prefix),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
             ModelPatcher.apply_clip_skip(text_encoder_model, clip_field.skipped_layers),
         ):
@@ -193,7 +194,7 @@ class SDXLPromptInvocationBase:
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
                 textual_inversion_manager=ti_manager,
-                dtype_for_device_getter=TorchDevice.choose_torch_dtype,
+                dtype_for_device_getter=torch_dtype,
                 truncate_long_prompts=False,  # TODO:
                 returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,  # TODO: clip skip
                 requires_pooled=get_pooled,
@@ -206,12 +207,17 @@ class SDXLPromptInvocationBase:
                 log_tokenization_for_conjunction(conjunction, tokenizer)
 
             # TODO: ask for optimizations? to not run text_encoder twice
-            c, _options = compel.build_conditioning_tensor_for_conjunction(conjunction)
+            c, options = compel.build_conditioning_tensor_for_conjunction(conjunction)
             if get_pooled:
                 c_pooled = compel.conditioning_provider.get_pooled_embeddings([prompt])
             else:
                 c_pooled = None
 
+            ec = ExtraConditioningInfo(
+                tokens_count_including_eos_bos=get_max_token_count(tokenizer, conjunction),
+                cross_attention_control_args=options.get("cross_attention_control", None),
+            )
+
         del tokenizer
         del text_encoder
         del tokenizer_info
@@ -221,7 +227,7 @@ class SDXLPromptInvocationBase:
         if c_pooled is not None:
             c_pooled = c_pooled.detach().to("cpu")
 
-        return c, c_pooled
+        return c, c_pooled, ec
 
 
 @invocation(
@@ -229,7 +235,7 @@ class SDXLPromptInvocationBase:
     title="SDXL Prompt",
     tags=["sdxl", "compel", "prompt"],
     category="conditioning",
-    version="1.2.0",
+    version="1.1.1",
 )
 class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
@@ -252,19 +258,20 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     target_height: int = InputField(default=1024, description="")
     clip: CLIPField = InputField(description=FieldDescriptions.clip, input=Input.Connection, title="CLIP 1")
     clip2: CLIPField = InputField(description=FieldDescriptions.clip, input=Input.Connection, title="CLIP 2")
-    mask: Optional[TensorField] = InputField(
-        default=None, description="A mask defining the region that this conditioning prompt applies to."
-    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
-        c1, c1_pooled = self.run_clip_compel(context, self.clip, self.prompt, False, "lora_te1_", zero_on_empty=True)
+        c1, c1_pooled, ec1 = self.run_clip_compel(
+            context, self.clip, self.prompt, False, "text_encoder", zero_on_empty=True
+        )
         if self.style.strip() == "":
-            c2, c2_pooled = self.run_clip_compel(
-                context, self.clip2, self.prompt, True, "lora_te2_", zero_on_empty=True
+            c2, c2_pooled, ec2 = self.run_clip_compel(
+                context, self.clip2, self.prompt, True, "text_encoder_2", zero_on_empty=True
             )
         else:
-            c2, c2_pooled = self.run_clip_compel(context, self.clip2, self.style, True, "lora_te2_", zero_on_empty=True)
+            c2, c2_pooled, ec2 = self.run_clip_compel(
+                context, self.clip2, self.style, True, "text_encoder_2", zero_on_empty=True
+            )
 
         original_size = (self.original_height, self.original_width)
         crop_coords = (self.crop_top, self.crop_left)
@@ -303,19 +310,17 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
         conditioning_data = ConditioningFieldData(
             conditionings=[
                 SDXLConditioningInfo(
-                    embeds=torch.cat([c1, c2], dim=-1), pooled_embeds=c2_pooled, add_time_ids=add_time_ids
+                    embeds=torch.cat([c1, c2], dim=-1),
+                    pooled_embeds=c2_pooled,
+                    add_time_ids=add_time_ids,
+                    extra_conditioning=ec1,
                 )
             ]
         )
 
         conditioning_name = context.conditioning.save(conditioning_data)
 
-        return ConditioningOutput(
-            conditioning=ConditioningField(
-                conditioning_name=conditioning_name,
-                mask=self.mask,
-            )
-        )
+        return ConditioningOutput.build(conditioning_name)
 
 
 @invocation(
@@ -343,7 +348,7 @@ class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
         # TODO: if there will appear lora for refiner - write proper prefix
-        c2, c2_pooled = self.run_clip_compel(context, self.clip2, self.style, True, "<NONE>", zero_on_empty=False)
+        c2, c2_pooled, ec2 = self.run_clip_compel(context, self.clip2, self.style, True, "<NONE>", zero_on_empty=False)
 
         original_size = (self.original_height, self.original_width)
         crop_coords = (self.crop_top, self.crop_left)
@@ -352,7 +357,14 @@ class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase
 
         assert c2_pooled is not None
         conditioning_data = ConditioningFieldData(
-            conditionings=[SDXLConditioningInfo(embeds=c2, pooled_embeds=c2_pooled, add_time_ids=add_time_ids)]
+            conditionings=[
+                SDXLConditioningInfo(
+                    embeds=c2,
+                    pooled_embeds=c2_pooled,
+                    add_time_ids=add_time_ids,
+                    extra_conditioning=ec2,  # or None
+                )
+            ]
         )
 
         conditioning_name = context.conditioning.save(conditioning_data)

invokeai/app/invocations/custom_nodes/init.py

@@ -3,7 +3,6 @@ Invoke-managed custom node loader. See README.md for more information.
 """
 
 import sys
-import traceback
 from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
 
@@ -42,15 +41,11 @@ for d in Path(__file__).parent.iterdir():
 
     logger.info(f"Loading node pack {module_name}")
 
-    try:
-        module = module_from_spec(spec)
-        sys.modules[spec.name] = module
-        spec.loader.exec_module(module)
+    module = module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
 
-        loaded_count += 1
-    except Exception:
-        full_error = traceback.format_exc()
-        logger.error(f"Failed to load node pack {module_name}:\n{full_error}")
+    loaded_count += 1
 
     del init, module_name
 

invokeai/app/invocations/fields.py

@@ -203,12 +203,6 @@ class DenoiseMaskField(BaseModel):
     gradient: bool = Field(default=False, description="Used for gradient inpainting")
 
 
-class TensorField(BaseModel):
-    """A tensor primitive field."""
-
-    tensor_name: str = Field(description="The name of a tensor.")
-
-
 class LatentsField(BaseModel):
     """A latents tensor primitive field"""
 
@@ -232,11 +226,7 @@ class ConditioningField(BaseModel):
     """A conditioning tensor primitive value"""
 
     conditioning_name: str = Field(description="The name of conditioning tensor")
-    mask: Optional[TensorField] = Field(
-        default=None,
-        description="The mask associated with this conditioning tensor. Excluded regions should be set to False, "
-        "included regions should be set to True.",
-    )
+    # endregion
 
 
 class MetadataField(RootModel[dict[str, Any]]):

invokeai/app/invocations/infill.py

@@ -1,91 +1,154 @@
-from abc import abstractmethod
-from typing import Literal, get_args
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team
 
-from PIL import Image
+import math
+from typing import Literal, Optional, get_args
+
+import numpy as np
+from PIL import Image, ImageOps
 
 from invokeai.app.invocations.fields import ColorField, ImageField
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.util.download_with_progress import download_with_progress_bar
 from invokeai.app.util.misc import SEED_MAX
-from invokeai.backend.image_util.infill_methods.cv2_inpaint import cv2_inpaint
-from invokeai.backend.image_util.infill_methods.lama import LaMA
-from invokeai.backend.image_util.infill_methods.mosaic import infill_mosaic
-from invokeai.backend.image_util.infill_methods.patchmatch import PatchMatch, infill_patchmatch
-from invokeai.backend.image_util.infill_methods.tile import infill_tile
-from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.image_util.cv2_inpaint import cv2_inpaint
+from invokeai.backend.image_util.lama import LaMA
+from invokeai.backend.image_util.patchmatch import PatchMatch
 
 from .baseinvocation import BaseInvocation, invocation
 from .fields import InputField, WithBoard, WithMetadata
 from .image import PIL_RESAMPLING_MAP, PIL_RESAMPLING_MODES
 
-logger = InvokeAILogger.get_logger()
 
-
-def get_infill_methods():
-    methods = Literal["tile", "color", "lama", "cv2"]  # TODO: add mosaic back
+def infill_methods() -> list[str]:
+    methods = ["tile", "solid", "lama", "cv2"]
     if PatchMatch.patchmatch_available():
-        methods = Literal["patchmatch", "tile", "color", "lama", "cv2"]  # TODO: add mosaic back
+        methods.insert(0, "patchmatch")
     return methods
 
 
-INFILL_METHODS = get_infill_methods()
+INFILL_METHODS = Literal[tuple(infill_methods())]
 DEFAULT_INFILL_METHOD = "patchmatch" if "patchmatch" in get_args(INFILL_METHODS) else "tile"
 
 
-class InfillImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
-    """Base class for invocations that preprocess images for Infilling"""
+def infill_lama(im: Image.Image) -> Image.Image:
+    lama = LaMA()
+    return lama(im)
 
-    image: ImageField = InputField(description="The image to process")
 
-    @abstractmethod
-    def infill(self, image: Image.Image) -> Image.Image:
-        """Infill the image with the specified method"""
-        pass
+def infill_patchmatch(im: Image.Image) -> Image.Image:
+    if im.mode != "RGBA":
+        return im
 
-    def load_image(self, context: InvocationContext) -> tuple[Image.Image, bool]:
-        """Process the image to have an alpha channel before being infilled"""
-        image = context.images.get_pil(self.image.image_name)
-        has_alpha = True if image.mode == "RGBA" else False
-        return image, has_alpha
+    # Skip patchmatch if patchmatch isn't available
+    if not PatchMatch.patchmatch_available():
+        return im
 
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        # Retrieve and process image to be infilled
-        input_image, has_alpha = self.load_image(context)
+    # Patchmatch (note, we may want to expose patch_size? Increasing it significantly impacts performance though)
+    im_patched_np = PatchMatch.inpaint(im.convert("RGB"), ImageOps.invert(im.split()[-1]), patch_size=3)
+    im_patched = Image.fromarray(im_patched_np, mode="RGB")
+    return im_patched
 
-        # If the input image has no alpha channel, return it
-        if has_alpha is False:
-            return ImageOutput.build(context.images.get_dto(self.image.image_name))
 
-        # Perform Infill action
-        infilled_image = self.infill(input_image)
+def infill_cv2(im: Image.Image) -> Image.Image:
+    return cv2_inpaint(im)
 
-        # Create ImageDTO for Infilled Image
-        infilled_image_dto = context.images.save(image=infilled_image)
 
-        # Return Infilled Image
-        return ImageOutput.build(infilled_image_dto)
+def get_tile_images(image: np.ndarray, width=8, height=8):
+    _nrows, _ncols, depth = image.shape
+    _strides = image.strides
+
+    nrows, _m = divmod(_nrows, height)
+    ncols, _n = divmod(_ncols, width)
+    if _m != 0 or _n != 0:
+        return None
+
+    return np.lib.stride_tricks.as_strided(
+        np.ravel(image),
+        shape=(nrows, ncols, height, width, depth),
+        strides=(height * _strides[0], width * _strides[1], *_strides),
+        writeable=False,
+    )
+
+
+def tile_fill_missing(im: Image.Image, tile_size: int = 16, seed: Optional[int] = None) -> Image.Image:
+    # Only fill if there's an alpha layer
+    if im.mode != "RGBA":
+        return im
+
+    a = np.asarray(im, dtype=np.uint8)
+
+    tile_size_tuple = (tile_size, tile_size)
+
+    # Get the image as tiles of a specified size
+    tiles = get_tile_images(a, *tile_size_tuple).copy()
+
+    # Get the mask as tiles
+    tiles_mask = tiles[:, :, :, :, 3]
+
+    # Find any mask tiles with any fully transparent pixels (we will be replacing these later)
+    tmask_shape = tiles_mask.shape
+    tiles_mask = tiles_mask.reshape(math.prod(tiles_mask.shape))
+    n, ny = (math.prod(tmask_shape[0:2])), math.prod(tmask_shape[2:])
+    tiles_mask = tiles_mask > 0
+    tiles_mask = tiles_mask.reshape((n, ny)).all(axis=1)
+
+    # Get RGB tiles in single array and filter by the mask
+    tshape = tiles.shape
+    tiles_all = tiles.reshape((math.prod(tiles.shape[0:2]), *tiles.shape[2:]))
+    filtered_tiles = tiles_all[tiles_mask]
+
+    if len(filtered_tiles) == 0:
+        return im
+
+    # Find all invalid tiles and replace with a random valid tile
+    replace_count = (tiles_mask == False).sum()  # noqa: E712
+    rng = np.random.default_rng(seed=seed)
+    tiles_all[np.logical_not(tiles_mask)] = filtered_tiles[rng.choice(filtered_tiles.shape[0], replace_count), :, :, :]
+
+    # Convert back to an image
+    tiles_all = tiles_all.reshape(tshape)
+    tiles_all = tiles_all.swapaxes(1, 2)
+    st = tiles_all.reshape(
+        (
+            math.prod(tiles_all.shape[0:2]),
+            math.prod(tiles_all.shape[2:4]),
+            tiles_all.shape[4],
+        )
+    )
+    si = Image.fromarray(st, mode="RGBA")
+
+    return si
 
 
 @invocation("infill_rgba", title="Solid Color Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.2")
-class InfillColorInvocation(InfillImageProcessorInvocation):
+class InfillColorInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image with a solid color"""
 
+    image: ImageField = InputField(description="The image to infill")
     color: ColorField = InputField(
         default=ColorField(r=127, g=127, b=127, a=255),
         description="The color to use to infill",
     )
 
-    def infill(self, image: Image.Image):
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name)
+
         solid_bg = Image.new("RGBA", image.size, self.color.tuple())
         infilled = Image.alpha_composite(solid_bg, image.convert("RGBA"))
+
         infilled.paste(image, (0, 0), image.split()[-1])
-        return infilled
+
+        image_dto = context.images.save(image=infilled)
+
+        return ImageOutput.build(image_dto)
 
 
 @invocation("infill_tile", title="Tile Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.3")
-class InfillTileInvocation(InfillImageProcessorInvocation):
+class InfillTileInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image with tiles of the image"""
 
+    image: ImageField = InputField(description="The image to infill")
     tile_size: int = InputField(default=32, ge=1, description="The tile size (px)")
     seed: int = InputField(
         default=0,
@@ -94,74 +157,92 @@ class InfillTileInvocation(InfillImageProcessorInvocation):
         description="The seed to use for tile generation (omit for random)",
     )
 
-    def infill(self, image: Image.Image):
-        output = infill_tile(image, seed=self.seed, tile_size=self.tile_size)
-        return output.infilled
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name)
+
+        infilled = tile_fill_missing(image.copy(), seed=self.seed, tile_size=self.tile_size)
+        infilled.paste(image, (0, 0), image.split()[-1])
+
+        image_dto = context.images.save(image=infilled)
+
+        return ImageOutput.build(image_dto)
 
 
 @invocation(
     "infill_patchmatch", title="PatchMatch Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.2"
 )
-class InfillPatchMatchInvocation(InfillImageProcessorInvocation):
+class InfillPatchMatchInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image using the PatchMatch algorithm"""
 
+    image: ImageField = InputField(description="The image to infill")
     downscale: float = InputField(default=2.0, gt=0, description="Run patchmatch on downscaled image to speedup infill")
     resample_mode: PIL_RESAMPLING_MODES = InputField(default="bicubic", description="The resampling mode")
 
-    def infill(self, image: Image.Image):
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name).convert("RGBA")
+
         resample_mode = PIL_RESAMPLING_MAP[self.resample_mode]
 
+        infill_image = image.copy()
         width = int(image.width / self.downscale)
         height = int(image.height / self.downscale)
-
-        infilled = image.resize(
+        infill_image = infill_image.resize(
             (width, height),
             resample=resample_mode,
         )
-        infilled = infill_patchmatch(image)
+
+        if PatchMatch.patchmatch_available():
+            infilled = infill_patchmatch(infill_image)
+        else:
+            raise ValueError("PatchMatch is not available on this system")
+
         infilled = infilled.resize(
             (image.width, image.height),
             resample=resample_mode,
         )
-        infilled.paste(image, (0, 0), mask=image.split()[-1])
 
-        return infilled
+        infilled.paste(image, (0, 0), mask=image.split()[-1])
+        # image.paste(infilled, (0, 0), mask=image.split()[-1])
+
+        image_dto = context.images.save(image=infilled)
+
+        return ImageOutput.build(image_dto)
 
 
 @invocation("infill_lama", title="LaMa Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.2")
-class LaMaInfillInvocation(InfillImageProcessorInvocation):
+class LaMaInfillInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image using the LaMa model"""
 
-    def infill(self, image: Image.Image):
-        lama = LaMA()
-        return lama(image)
+    image: ImageField = InputField(description="The image to infill")
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name)
+
+        # Downloads the LaMa model if it doesn't already exist
+        download_with_progress_bar(
+            name="LaMa Inpainting Model",
+            url="https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+            dest_path=context.config.get().models_path / "core/misc/lama/lama.pt",
+        )
+
+        infilled = infill_lama(image.copy())
+
+        image_dto = context.images.save(image=infilled)
+
+        return ImageOutput.build(image_dto)
 
 
 @invocation("infill_cv2", title="CV2 Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.2")
-class CV2InfillInvocation(InfillImageProcessorInvocation):
+class CV2InfillInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image using OpenCV Inpainting"""
 
-    def infill(self, image: Image.Image):
-        return cv2_inpaint(image)
-
-
-# @invocation(
-#     "infill_mosaic", title="Mosaic Infill", tags=["image", "inpaint", "outpaint"], category="inpaint", version="1.0.0"
-# )
-class MosaicInfillInvocation(InfillImageProcessorInvocation):
-    """Infills transparent areas of an image with a mosaic pattern drawing colors from the rest of the image"""
-
     image: ImageField = InputField(description="The image to infill")
-    tile_width: int = InputField(default=64, description="Width of the tile")
-    tile_height: int = InputField(default=64, description="Height of the tile")
-    min_color: ColorField = InputField(
-        default=ColorField(r=0, g=0, b=0, a=255),
-        description="The min threshold for color",
-    )
-    max_color: ColorField = InputField(
-        default=ColorField(r=255, g=255, b=255, a=255),
-        description="The max threshold for color",
-    )
 
-    def infill(self, image: Image.Image):
-        return infill_mosaic(image, (self.tile_width, self.tile_height), self.min_color.tuple(), self.max_color.tuple())
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name)
+
+        infilled = infill_cv2(image.copy())
+
+        image_dto = context.images.save(image=infilled)
+
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/ip_adapter.py

@@ -1,41 +1,34 @@
 from builtins import float
-from typing import List, Literal, Optional, Union
+from typing import List, Union
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, TensorField, UIType
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.invocations.primitives import ImageField
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager.config import (
-    AnyModelConfig,
-    BaseModelType,
-    IPAdapterCheckpointConfig,
-    IPAdapterInvokeAIConfig,
-    ModelType,
-)
+from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, IPAdapterConfig, ModelType
 
 
 class IPAdapterField(BaseModel):
     image: Union[ImageField, List[ImageField]] = Field(description="The IP-Adapter image prompt(s).")
     ip_adapter_model: ModelIdentifierField = Field(description="The IP-Adapter model to use.")
     image_encoder_model: ModelIdentifierField = Field(description="The name of the CLIP image encoder model.")
-    weight: Union[float, List[float]] = Field(default=1, description="The weight given to the IP-Adapter.")
-    target_blocks: List[str] = Field(default=[], description="The IP Adapter blocks to apply")
+    weight: Union[float, List[float]] = Field(default=1, description="The weight given to the ControlNet")
     begin_step_percent: float = Field(
         default=0, ge=0, le=1, description="When the IP-Adapter is first applied (% of total steps)"
     )
     end_step_percent: float = Field(
         default=1, ge=0, le=1, description="When the IP-Adapter is last applied (% of total steps)"
     )
-    mask: Optional[TensorField] = Field(
-        default=None,
-        description="The bool mask associated with this IP-Adapter. Excluded regions should be set to False, included "
-        "regions should be set to True.",
-    )
 
     @field_validator("weight")
     @classmethod
@@ -55,15 +48,12 @@ class IPAdapterOutput(BaseInvocationOutput):
     ip_adapter: IPAdapterField = OutputField(description=FieldDescriptions.ip_adapter, title="IP-Adapter")
 
 
-CLIP_VISION_MODEL_MAP = {"ViT-H": "ip_adapter_sd_image_encoder", "ViT-G": "ip_adapter_sdxl_image_encoder"}
-
-
-@invocation("ip_adapter", title="IP-Adapter", tags=["ip_adapter", "control"], category="ip_adapter", version="1.4.0")
+@invocation("ip_adapter", title="IP-Adapter", tags=["ip_adapter", "control"], category="ip_adapter", version="1.2.2")
 class IPAdapterInvocation(BaseInvocation):
     """Collects IP-Adapter info to pass to other nodes."""
 
     # Inputs
-    image: Union[ImageField, List[ImageField]] = InputField(description="The IP-Adapter image prompt(s).", ui_order=1)
+    image: Union[ImageField, List[ImageField]] = InputField(description="The IP-Adapter image prompt(s).")
     ip_adapter_model: ModelIdentifierField = InputField(
         description="The IP-Adapter model.",
         title="IP-Adapter Model",
@@ -71,26 +61,16 @@ class IPAdapterInvocation(BaseInvocation):
         ui_order=-1,
         ui_type=UIType.IPAdapterModel,
     )
-    clip_vision_model: Literal["ViT-H", "ViT-G"] = InputField(
-        description="CLIP Vision model to use. Overrides model settings. Mandatory for checkpoint models.",
-        default="ViT-H",
-        ui_order=2,
-    )
+
     weight: Union[float, List[float]] = InputField(
         default=1, description="The weight given to the IP-Adapter", title="Weight"
     )
-    method: Literal["full", "style", "composition"] = InputField(
-        default="full", description="The method to apply the IP-Adapter"
-    )
     begin_step_percent: float = InputField(
         default=0, ge=0, le=1, description="When the IP-Adapter is first applied (% of total steps)"
     )
     end_step_percent: float = InputField(
         default=1, ge=0, le=1, description="When the IP-Adapter is last applied (% of total steps)"
     )
-    mask: Optional[TensorField] = InputField(
-        default=None, description="A mask defining the region that this IP-Adapter applies to."
-    )
 
     @field_validator("weight")
     @classmethod
@@ -106,68 +86,35 @@ class IPAdapterInvocation(BaseInvocation):
     def invoke(self, context: InvocationContext) -> IPAdapterOutput:
         # Lookup the CLIP Vision encoder that is intended to be used with the IP-Adapter model.
         ip_adapter_info = context.models.get_config(self.ip_adapter_model.key)
-        assert isinstance(ip_adapter_info, (IPAdapterInvokeAIConfig, IPAdapterCheckpointConfig))
-
-        if isinstance(ip_adapter_info, IPAdapterInvokeAIConfig):
-            image_encoder_model_id = ip_adapter_info.image_encoder_model_id
-            image_encoder_model_name = image_encoder_model_id.split("/")[-1].strip()
-        else:
-            image_encoder_model_name = CLIP_VISION_MODEL_MAP[self.clip_vision_model]
-
+        assert isinstance(ip_adapter_info, IPAdapterConfig)
+        image_encoder_model_id = ip_adapter_info.image_encoder_model_id
+        image_encoder_model_name = image_encoder_model_id.split("/")[-1].strip()
         image_encoder_model = self._get_image_encoder(context, image_encoder_model_name)
-
-        if self.method == "style":
-            if ip_adapter_info.base == "sd-1":
-                target_blocks = ["up_blocks.1"]
-            elif ip_adapter_info.base == "sdxl":
-                target_blocks = ["up_blocks.0.attentions.1"]
-            else:
-                raise ValueError(f"Unsupported IP-Adapter base type: '{ip_adapter_info.base}'.")
-        elif self.method == "composition":
-            if ip_adapter_info.base == "sd-1":
-                target_blocks = ["down_blocks.2", "mid_block"]
-            elif ip_adapter_info.base == "sdxl":
-                target_blocks = ["down_blocks.2.attentions.1"]
-            else:
-                raise ValueError(f"Unsupported IP-Adapter base type: '{ip_adapter_info.base}'.")
-        elif self.method == "full":
-            target_blocks = ["block"]
-        else:
-            raise ValueError(f"Unexpected IP-Adapter method: '{self.method}'.")
-
         return IPAdapterOutput(
             ip_adapter=IPAdapterField(
                 image=self.image,
                 ip_adapter_model=self.ip_adapter_model,
                 image_encoder_model=ModelIdentifierField.from_config(image_encoder_model),
                 weight=self.weight,
-                target_blocks=target_blocks,
                 begin_step_percent=self.begin_step_percent,
                 end_step_percent=self.end_step_percent,
-                mask=self.mask,
             ),
         )
 
     def _get_image_encoder(self, context: InvocationContext, image_encoder_model_name: str) -> AnyModelConfig:
-        image_encoder_models = context.models.search_by_attrs(
-            name=image_encoder_model_name, base=BaseModelType.Any, type=ModelType.CLIPVision
-        )
-
-        if not len(image_encoder_models) > 0:
-            context.logger.warning(
-                f"The image encoder required by this IP Adapter ({image_encoder_model_name}) is not installed. \
-                    Downloading and installing now. This may take a while."
-            )
-
-            installer = context._services.model_manager.install
-            job = installer.heuristic_import(f"InvokeAI/{image_encoder_model_name}")
-            installer.wait_for_job(job, timeout=600)  # Wait for up to 10 minutes
+        found = False
+        while not found:
             image_encoder_models = context.models.search_by_attrs(
                 name=image_encoder_model_name, base=BaseModelType.Any, type=ModelType.CLIPVision
             )
-
-            if len(image_encoder_models) == 0:
-                context.logger.error("Error while fetching CLIP Vision Image Encoder")
-                assert len(image_encoder_models) == 1
-
+            found = len(image_encoder_models) > 0
+            if not found:
+                context.logger.warning(
+                    f"The image encoder required by this IP Adapter ({image_encoder_model_name}) is not installed."
+                )
+                context.logger.warning("Downloading and installing now. This may take a while.")
+                installer = context._services.model_manager.install
+                job = installer.heuristic_import(f"InvokeAI/{image_encoder_model_name}")
+                installer.wait_for_job(job, timeout=600)  # wait up to 10 minutes - then raise a TimeoutException
+        assert len(image_encoder_models) == 1
         return image_encoder_models[0]

invokeai/app/invocations/latent.py

@@ -1,5 +1,5 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
-import inspect
+
 import math
 from contextlib import ExitStack
 from functools import singledispatchmethod
@@ -9,7 +9,6 @@ import einops
 import numpy as np
 import numpy.typing as npt
 import torch
-import torchvision
 import torchvision.transforms as T
 from diffusers import AutoencoderKL, AutoencoderTiny
 from diffusers.configuration_utils import ConfigMixin
@@ -44,40 +43,45 @@ from invokeai.app.invocations.fields import (
     WithMetadata,
 )
 from invokeai.app.invocations.ip_adapter import IPAdapterField
-from invokeai.app.invocations.primitives import DenoiseMaskOutput, ImageOutput, LatentsOutput
+from invokeai.app.invocations.primitives import (
+    DenoiseMaskOutput,
+    ImageOutput,
+    LatentsOutput,
+)
 from invokeai.app.invocations.t2i_adapter import T2IAdapterField
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.controlnet_utils import prepare_control_image
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter, IPAdapterPlus
-from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.lora_model_patcher import LoraModelPatcher
+from invokeai.backend.lora_model_raw import LoRAModelRaw
 from invokeai.backend.model_manager import BaseModelType, LoadedModel
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion import PipelineIntermediateState, set_seamless
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
-    BasicConditioningInfo,
-    IPAdapterConditioningInfo,
-    IPAdapterData,
-    Range,
-    SDXLConditioningInfo,
-    TextConditioningData,
-    TextConditioningRegions,
-)
-from invokeai.backend.util.mask import to_standard_float_mask
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningData, IPAdapterConditioningInfo
 from invokeai.backend.util.silence_warnings import SilenceWarnings
 
 from ...backend.stable_diffusion.diffusers_pipeline import (
     ControlNetData,
+    IPAdapterData,
     StableDiffusionGeneratorPipeline,
     T2IAdapterData,
     image_resized_to_grid_as_tensor,
 )
 from ...backend.stable_diffusion.schedulers import SCHEDULER_MAP
-from ...backend.util.devices import TorchDevice
-from .baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
+from ...backend.util.devices import choose_precision, choose_torch_device
+from .baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    invocation,
+    invocation_output,
+)
 from .controlnet_image_processors import ControlField
 from .model import ModelIdentifierField, UNetField, VAEField
 
-DEFAULT_PRECISION = TorchDevice.choose_torch_dtype()
+if choose_torch_device() == torch.device("mps"):
+    from torch import mps
+
+DEFAULT_PRECISION = choose_precision(choose_torch_device())
 
 
 @invocation_output("scheduler_output")
@@ -281,10 +285,10 @@ def get_scheduler(
 class DenoiseLatentsInvocation(BaseInvocation):
     """Denoises noisy latents to decodable images"""
 
-    positive_conditioning: Union[ConditioningField, list[ConditioningField]] = InputField(
+    positive_conditioning: ConditioningField = InputField(
         description=FieldDescriptions.positive_cond, input=Input.Connection, ui_order=0
     )
-    negative_conditioning: Union[ConditioningField, list[ConditioningField]] = InputField(
+    negative_conditioning: ConditioningField = InputField(
         description=FieldDescriptions.negative_cond, input=Input.Connection, ui_order=1
     )
     noise: Optional[LatentsField] = InputField(
@@ -362,169 +366,34 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 raise ValueError("cfg_scale must be greater than 1")
         return v
 
-    def _get_text_embeddings_and_masks(
-        self,
-        cond_list: list[ConditioningField],
-        context: InvocationContext,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> tuple[Union[list[BasicConditioningInfo], list[SDXLConditioningInfo]], list[Optional[torch.Tensor]]]:
-        """Get the text embeddings and masks from the input conditioning fields."""
-        text_embeddings: Union[list[BasicConditioningInfo], list[SDXLConditioningInfo]] = []
-        text_embeddings_masks: list[Optional[torch.Tensor]] = []
-        for cond in cond_list:
-            cond_data = context.conditioning.load(cond.conditioning_name)
-            text_embeddings.append(cond_data.conditionings[0].to(device=device, dtype=dtype))
-
-            mask = cond.mask
-            if mask is not None:
-                mask = context.tensors.load(mask.tensor_name)
-            text_embeddings_masks.append(mask)
-
-        return text_embeddings, text_embeddings_masks
-
-    def _preprocess_regional_prompt_mask(
-        self, mask: Optional[torch.Tensor], target_height: int, target_width: int, dtype: torch.dtype
-    ) -> torch.Tensor:
-        """Preprocess a regional prompt mask to match the target height and width.
-        If mask is None, returns a mask of all ones with the target height and width.
-        If mask is not None, resizes the mask to the target height and width using 'nearest' interpolation.
-
-        Returns:
-            torch.Tensor: The processed mask. shape: (1, 1, target_height, target_width).
-        """
-
-        if mask is None:
-            return torch.ones((1, 1, target_height, target_width), dtype=dtype)
-
-        mask = to_standard_float_mask(mask, out_dtype=dtype)
-
-        tf = torchvision.transforms.Resize(
-            (target_height, target_width), interpolation=torchvision.transforms.InterpolationMode.NEAREST
-        )
-
-        # Add a batch dimension to the mask, because torchvision expects shape (batch, channels, h, w).
-        mask = mask.unsqueeze(0)  # Shape: (1, h, w) -> (1, 1, h, w)
-        resized_mask = tf(mask)
-        return resized_mask
-
-    def _concat_regional_text_embeddings(
-        self,
-        text_conditionings: Union[list[BasicConditioningInfo], list[SDXLConditioningInfo]],
-        masks: Optional[list[Optional[torch.Tensor]]],
-        latent_height: int,
-        latent_width: int,
-        dtype: torch.dtype,
-    ) -> tuple[Union[BasicConditioningInfo, SDXLConditioningInfo], Optional[TextConditioningRegions]]:
-        """Concatenate regional text embeddings into a single embedding and track the region masks accordingly."""
-        if masks is None:
-            masks = [None] * len(text_conditionings)
-        assert len(text_conditionings) == len(masks)
-
-        is_sdxl = type(text_conditionings[0]) is SDXLConditioningInfo
-
-        all_masks_are_none = all(mask is None for mask in masks)
-
-        text_embedding = []
-        pooled_embedding = None
-        add_time_ids = None
-        cur_text_embedding_len = 0
-        processed_masks = []
-        embedding_ranges = []
-
-        for prompt_idx, text_embedding_info in enumerate(text_conditionings):
-            mask = masks[prompt_idx]
-
-            if is_sdxl:
-                # We choose a random SDXLConditioningInfo's pooled_embeds and add_time_ids here, with a preference for
-                # prompts without a mask. We prefer prompts without a mask, because they are more likely to contain
-                # global prompt information.  In an ideal case, there should be exactly one global prompt without a
-                # mask, but we don't enforce this.
-
-                # HACK(ryand): The fact that we have to choose a single pooled_embedding and add_time_ids here is a
-                # fundamental interface issue. The SDXL Compel nodes are not designed to be used in the way that we use
-                # them for regional prompting. Ideally, the DenoiseLatents invocation should accept a single
-                # pooled_embeds tensor and a list of standard text embeds with region masks. This change would be a
-                # pretty major breaking change to a popular node, so for now we use this hack.
-                if pooled_embedding is None or mask is None:
-                    pooled_embedding = text_embedding_info.pooled_embeds
-                if add_time_ids is None or mask is None:
-                    add_time_ids = text_embedding_info.add_time_ids
-
-            text_embedding.append(text_embedding_info.embeds)
-            if not all_masks_are_none:
-                embedding_ranges.append(
-                    Range(
-                        start=cur_text_embedding_len, end=cur_text_embedding_len + text_embedding_info.embeds.shape[1]
-                    )
-                )
-                processed_masks.append(
-                    self._preprocess_regional_prompt_mask(mask, latent_height, latent_width, dtype=dtype)
-                )
-
-            cur_text_embedding_len += text_embedding_info.embeds.shape[1]
-
-        text_embedding = torch.cat(text_embedding, dim=1)
-        assert len(text_embedding.shape) == 3  # batch_size, seq_len, token_len
-
-        regions = None
-        if not all_masks_are_none:
-            regions = TextConditioningRegions(
-                masks=torch.cat(processed_masks, dim=1),
-                ranges=embedding_ranges,
-            )
-
-        if is_sdxl:
-            return SDXLConditioningInfo(
-                embeds=text_embedding, pooled_embeds=pooled_embedding, add_time_ids=add_time_ids
-            ), regions
-        return BasicConditioningInfo(embeds=text_embedding), regions
-
     def get_conditioning_data(
         self,
         context: InvocationContext,
+        scheduler: Scheduler,
         unet: UNet2DConditionModel,
-        latent_height: int,
-        latent_width: int,
-    ) -> TextConditioningData:
-        # Normalize self.positive_conditioning and self.negative_conditioning to lists.
-        cond_list = self.positive_conditioning
-        if not isinstance(cond_list, list):
-            cond_list = [cond_list]
-        uncond_list = self.negative_conditioning
-        if not isinstance(uncond_list, list):
-            uncond_list = [uncond_list]
+        seed: int,
+    ) -> ConditioningData:
+        positive_cond_data = context.conditioning.load(self.positive_conditioning.conditioning_name)
+        c = positive_cond_data.conditionings[0].to(device=unet.device, dtype=unet.dtype)
 
-        cond_text_embeddings, cond_text_embedding_masks = self._get_text_embeddings_and_masks(
-            cond_list, context, unet.device, unet.dtype
-        )
-        uncond_text_embeddings, uncond_text_embedding_masks = self._get_text_embeddings_and_masks(
-            uncond_list, context, unet.device, unet.dtype
-        )
+        negative_cond_data = context.conditioning.load(self.negative_conditioning.conditioning_name)
+        uc = negative_cond_data.conditionings[0].to(device=unet.device, dtype=unet.dtype)
 
-        cond_text_embedding, cond_regions = self._concat_regional_text_embeddings(
-            text_conditionings=cond_text_embeddings,
-            masks=cond_text_embedding_masks,
-            latent_height=latent_height,
-            latent_width=latent_width,
-            dtype=unet.dtype,
-        )
-        uncond_text_embedding, uncond_regions = self._concat_regional_text_embeddings(
-            text_conditionings=uncond_text_embeddings,
-            masks=uncond_text_embedding_masks,
-            latent_height=latent_height,
-            latent_width=latent_width,
-            dtype=unet.dtype,
-        )
-
-        conditioning_data = TextConditioningData(
-            uncond_text=uncond_text_embedding,
-            cond_text=cond_text_embedding,
-            uncond_regions=uncond_regions,
-            cond_regions=cond_regions,
+        conditioning_data = ConditioningData(
+            unconditioned_embeddings=uc,
+            text_embeddings=c,
             guidance_scale=self.cfg_scale,
             guidance_rescale_multiplier=self.cfg_rescale_multiplier,
         )
+
+        conditioning_data = conditioning_data.add_scheduler_args_if_applicable(  # FIXME
+            scheduler,
+            # for ddim scheduler
+            eta=0.0,  # ddim_eta
+            # for ancestral and sde schedulers
+            # flip all bits to have noise different from initial
+            generator=torch.Generator(device=unet.device).manual_seed(seed ^ 0xFFFFFFFF),
+        )
         return conditioning_data
 
     def create_pipeline(
@@ -629,10 +498,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
         self,
         context: InvocationContext,
         ip_adapter: Optional[Union[IPAdapterField, list[IPAdapterField]]],
+        conditioning_data: ConditioningData,
         exit_stack: ExitStack,
-        latent_height: int,
-        latent_width: int,
-        dtype: torch.dtype,
     ) -> Optional[list[IPAdapterData]]:
         """If IP-Adapter is enabled, then this function loads the requisite models, and adds the image prompt embeddings
         to the `conditioning_data` (in-place).
@@ -648,6 +515,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
             return None
 
         ip_adapter_data_list = []
+        conditioning_data.ip_adapter_conditioning = []
         for single_ip_adapter in ip_adapter:
             ip_adapter_model: Union[IPAdapter, IPAdapterPlus] = exit_stack.enter_context(
                 context.models.load(single_ip_adapter.ip_adapter_model)
@@ -670,20 +538,16 @@ class DenoiseLatentsInvocation(BaseInvocation):
                     single_ipa_images, image_encoder_model
                 )
 
-            mask = single_ip_adapter.mask
-            if mask is not None:
-                mask = context.tensors.load(mask.tensor_name)
-            mask = self._preprocess_regional_prompt_mask(mask, latent_height, latent_width, dtype=dtype)
+                conditioning_data.ip_adapter_conditioning.append(
+                    IPAdapterConditioningInfo(image_prompt_embeds, uncond_image_prompt_embeds)
+                )
 
             ip_adapter_data_list.append(
                 IPAdapterData(
                     ip_adapter_model=ip_adapter_model,
                     weight=single_ip_adapter.weight,
-                    target_blocks=single_ip_adapter.target_blocks,
                     begin_step_percent=single_ip_adapter.begin_step_percent,
                     end_step_percent=single_ip_adapter.end_step_percent,
-                    ip_adapter_conditioning=IPAdapterConditioningInfo(image_prompt_embeds, uncond_image_prompt_embeds),
-                    mask=mask,
                 )
             )
 
@@ -773,7 +637,6 @@ class DenoiseLatentsInvocation(BaseInvocation):
         steps: int,
         denoising_start: float,
         denoising_end: float,
-        seed: int,
     ) -> Tuple[int, List[int], int]:
         assert isinstance(scheduler, ConfigMixin)
         if scheduler.config.get("cpu_only", False):
@@ -802,15 +665,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         timesteps = timesteps[t_start_idx : t_start_idx + t_end_idx]
         num_inference_steps = len(timesteps) // scheduler.order
 
-        scheduler_step_kwargs = {}
-        scheduler_step_signature = inspect.signature(scheduler.step)
-        if "generator" in scheduler_step_signature.parameters:
-            # At some point, someone decided that schedulers that accept a generator should use the original seed with
-            # all bits flipped. I don't know the original rationale for this, but now we must keep it like this for
-            # reproducibility.
-            scheduler_step_kwargs = {"generator": torch.Generator(device=device).manual_seed(seed ^ 0xFFFFFFFF)}
-
-        return num_inference_steps, timesteps, init_timestep, scheduler_step_kwargs
+        return num_inference_steps, timesteps, init_timestep
 
     def prep_inpaint_mask(
         self, context: InvocationContext, latents: torch.Tensor
@@ -885,7 +740,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 set_seamless(unet_info.model, self.unet.seamless_axes),  # FIXME
                 unet_info as unet,
                 # Apply the LoRA after unet has been moved to its target device for faster patching.
-                ModelPatcher.apply_lora_unet(unet, _lora_loader()),
+                # ModelPatcher.apply_lora_unet(unet, _lora_loader()),
+                LoraModelPatcher.apply_lora_to_unet(unet, _lora_loader()),
             ):
                 assert isinstance(unet, UNet2DConditionModel)
                 latents = latents.to(device=unet.device, dtype=unet.dtype)
@@ -904,11 +760,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 )
 
                 pipeline = self.create_pipeline(unet, scheduler)
-
-                _, _, latent_height, latent_width = latents.shape
-                conditioning_data = self.get_conditioning_data(
-                    context=context, unet=unet, latent_height=latent_height, latent_width=latent_width
-                )
+                conditioning_data = self.get_conditioning_data(context, scheduler, unet, seed)
 
                 controlnet_data = self.prep_control_data(
                     context=context,
@@ -922,19 +774,16 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 ip_adapter_data = self.prep_ip_adapter_data(
                     context=context,
                     ip_adapter=self.ip_adapter,
+                    conditioning_data=conditioning_data,
                     exit_stack=exit_stack,
-                    latent_height=latent_height,
-                    latent_width=latent_width,
-                    dtype=unet.dtype,
                 )
 
-                num_inference_steps, timesteps, init_timestep, scheduler_step_kwargs = self.init_scheduler(
+                num_inference_steps, timesteps, init_timestep = self.init_scheduler(
                     scheduler,
                     device=unet.device,
                     steps=self.steps,
                     denoising_start=self.denoising_start,
                     denoising_end=self.denoising_end,
-                    seed=seed,
                 )
 
                 result_latents = pipeline.latents_from_embeddings(
@@ -947,7 +796,6 @@ class DenoiseLatentsInvocation(BaseInvocation):
                     masked_latents=masked_latents,
                     gradient_mask=gradient_mask,
                     num_inference_steps=num_inference_steps,
-                    scheduler_step_kwargs=scheduler_step_kwargs,
                     conditioning_data=conditioning_data,
                     control_data=controlnet_data,
                     ip_adapter_data=ip_adapter_data,
@@ -957,10 +805,12 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
             # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
             result_latents = result_latents.to("cpu")
-            TorchDevice.empty_cache()
+            torch.cuda.empty_cache()
+            if choose_torch_device() == torch.device("mps"):
+                mps.empty_cache()
 
             name = context.tensors.save(tensor=result_latents)
-        return LatentsOutput.build(latents_name=name, latents=result_latents, seed=None)
+        return LatentsOutput.build(latents_name=name, latents=result_latents, seed=seed)
 
 
 @invocation(
@@ -1024,7 +874,9 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
                 vae.disable_tiling()
 
             # clear memory as vae decode can request a lot
-            TorchDevice.empty_cache()
+            torch.cuda.empty_cache()
+            if choose_torch_device() == torch.device("mps"):
+                mps.empty_cache()
 
             with torch.inference_mode():
                 # copied from diffusers pipeline
@@ -1036,7 +888,9 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
 
                 image = VaeImageProcessor.numpy_to_pil(np_image)[0]
 
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
+        if choose_torch_device() == torch.device("mps"):
+            mps.empty_cache()
 
         image_dto = context.images.save(image=image)
 
@@ -1075,7 +929,9 @@ class ResizeLatentsInvocation(BaseInvocation):
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         latents = context.tensors.load(self.latents.latents_name)
-        device = TorchDevice.choose_torch_device()
+
+        # TODO:
+        device = choose_torch_device()
 
         resized_latents = torch.nn.functional.interpolate(
             latents.to(device),
@@ -1086,8 +942,9 @@ class ResizeLatentsInvocation(BaseInvocation):
 
         # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
         resized_latents = resized_latents.to("cpu")
-
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
+        if device == torch.device("mps"):
+            mps.empty_cache()
 
         name = context.tensors.save(tensor=resized_latents)
         return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
@@ -1114,7 +971,8 @@ class ScaleLatentsInvocation(BaseInvocation):
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         latents = context.tensors.load(self.latents.latents_name)
 
-        device = TorchDevice.choose_torch_device()
+        # TODO:
+        device = choose_torch_device()
 
         # resizing
         resized_latents = torch.nn.functional.interpolate(
@@ -1126,7 +984,9 @@ class ScaleLatentsInvocation(BaseInvocation):
 
         # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
         resized_latents = resized_latents.to("cpu")
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
+        if device == torch.device("mps"):
+            mps.empty_cache()
 
         name = context.tensors.save(tensor=resized_latents)
         return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
@@ -1258,7 +1118,8 @@ class BlendLatentsInvocation(BaseInvocation):
         if latents_a.shape != latents_b.shape:
             raise Exception("Latents to blend must be the same size.")
 
-        device = TorchDevice.choose_torch_device()
+        # TODO:
+        device = choose_torch_device()
 
         def slerp(
             t: Union[float, npt.NDArray[Any]],  # FIXME: maybe use np.float32 here?
@@ -1311,8 +1172,9 @@ class BlendLatentsInvocation(BaseInvocation):
 
         # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
         blended_latents = blended_latents.to("cpu")
-
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
+        if device == torch.device("mps"):
+            mps.empty_cache()
 
         name = context.tensors.save(tensor=blended_latents)
         return LatentsOutput.build(latents_name=name, latents=blended_latents)
@@ -1403,7 +1265,7 @@ class IdealSizeInvocation(BaseInvocation):
         return tuple((x - x % multiple_of) for x in args)
 
     def invoke(self, context: InvocationContext) -> IdealSizeOutput:
-        unet_config = context.models.get_config(self.unet.unet.key)
+        unet_config = context.models.get_config(**self.unet.unet.model_dump())
         aspect = self.width / self.height
         dimension: float = 512
         if unet_config.base == BaseModelType.StableDiffusion2:

invokeai/app/invocations/mask.py

@@ -1,36 +0,0 @@
-import torch
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext, invocation
-from invokeai.app.invocations.fields import InputField, TensorField, WithMetadata
-from invokeai.app.invocations.primitives import MaskOutput
-
-
-@invocation(
-    "rectangle_mask",
-    title="Create Rectangle Mask",
-    tags=["conditioning"],
-    category="conditioning",
-    version="1.0.1",
-)
-class RectangleMaskInvocation(BaseInvocation, WithMetadata):
-    """Create a rectangular mask."""
-
-    width: int = InputField(description="The width of the entire mask.")
-    height: int = InputField(description="The height of the entire mask.")
-    x_left: int = InputField(description="The left x-coordinate of the rectangular masked region (inclusive).")
-    y_top: int = InputField(description="The top y-coordinate of the rectangular masked region (inclusive).")
-    rectangle_width: int = InputField(description="The width of the rectangular masked region.")
-    rectangle_height: int = InputField(description="The height of the rectangular masked region.")
-
-    def invoke(self, context: InvocationContext) -> MaskOutput:
-        mask = torch.zeros((1, self.height, self.width), dtype=torch.bool)
-        mask[:, self.y_top : self.y_top + self.rectangle_height, self.x_left : self.x_left + self.rectangle_width] = (
-            True
-        )
-
-        mask_tensor_name = context.tensors.save(mask)
-        return MaskOutput(
-            mask=TensorField(tensor_name=mask_tensor_name),
-            width=self.width,
-            height=self.height,
-        )

invokeai/app/invocations/metadata.py

@@ -2,8 +2,16 @@ from typing import Any, Literal, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.controlnet_image_processors import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.controlnet_image_processors import (
+    CONTROLNET_MODE_VALUES,
+    CONTROLNET_RESIZE_VALUES,
+)
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     ImageField,
@@ -35,8 +43,6 @@ class IPAdapterMetadataField(BaseModel):
 
     image: ImageField = Field(description="The IP-Adapter image prompt.")
     ip_adapter_model: ModelIdentifierField = Field(description="The IP-Adapter model.")
-    clip_vision_model: Literal["ViT-H", "ViT-G"] = Field(description="The CLIP Vision model")
-    method: Literal["full", "style", "composition"] = Field(description="Method to apply IP Weights with")
     weight: Union[float, list[float]] = Field(description="The weight given to the IP-Adapter")
     begin_step_percent: float = Field(description="When the IP-Adapter is first applied (% of total steps)")
     end_step_percent: float = Field(description="When the IP-Adapter is last applied (% of total steps)")

invokeai/app/invocations/noise.py

@@ -9,7 +9,7 @@ from invokeai.app.invocations.fields import FieldDescriptions, InputField, Laten
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.misc import SEED_MAX
 
-from ...backend.util.devices import TorchDevice
+from ...backend.util.devices import choose_torch_device, torch_dtype
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
@@ -46,7 +46,7 @@ def get_noise(
             height // downsampling_factor,
             width // downsampling_factor,
         ],
-        dtype=TorchDevice.choose_torch_dtype(device=device),
+        dtype=torch_dtype(device),
         device=noise_device_type,
         generator=generator,
     ).to("cpu")
@@ -111,14 +111,14 @@ class NoiseInvocation(BaseInvocation):
 
     @field_validator("seed", mode="before")
     def modulo_seed(cls, v):
-        """Return the seed modulo (SEED_MAX + 1) to ensure it is within the valid range."""
+        """Returns the seed modulo (SEED_MAX + 1) to ensure it is within the valid range."""
         return v % (SEED_MAX + 1)
 
     def invoke(self, context: InvocationContext) -> NoiseOutput:
         noise = get_noise(
             width=self.width,
             height=self.height,
-            device=TorchDevice.choose_torch_device(),
+            device=choose_torch_device(),
             seed=self.seed,
             use_cpu=self.use_cpu,
         )

invokeai/app/invocations/primitives.py

@@ -15,7 +15,6 @@ from invokeai.app.invocations.fields import (
     InputField,
     LatentsField,
     OutputField,
-    TensorField,
     UIComponent,
 )
 from invokeai.app.services.images.images_common import ImageDTO
@@ -406,19 +405,9 @@ class ColorInvocation(BaseInvocation):
 
 # endregion
 
-
 # region Conditioning
 
 
-@invocation_output("mask_output")
-class MaskOutput(BaseInvocationOutput):
-    """A torch mask tensor."""
-
-    mask: TensorField = OutputField(description="The mask.")
-    width: int = OutputField(description="The width of the mask in pixels.")
-    height: int = OutputField(description="The height of the mask in pixels.")
-
-
 @invocation_output("conditioning_output")
 class ConditioningOutput(BaseInvocationOutput):
     """Base class for nodes that output a single conditioning tensor"""

invokeai/app/invocations/upscale.py

@@ -4,6 +4,7 @@ from typing import Literal
 
 import cv2
 import numpy as np
+import torch
 from PIL import Image
 from pydantic import ConfigDict
 
@@ -13,7 +14,7 @@ from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.download_with_progress import download_with_progress_bar
 from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet
 from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGAN
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 
 from .baseinvocation import BaseInvocation, invocation
 from .fields import InputField, WithBoard, WithMetadata
@@ -34,6 +35,9 @@ ESRGAN_MODEL_URLS: dict[str, str] = {
     "RealESRGAN_x2plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
 }
 
+if choose_torch_device() == torch.device("mps"):
+    from torch import mps
+
 
 @invocation("esrgan", title="Upscale (RealESRGAN)", tags=["esrgan", "upscale"], category="esrgan", version="1.3.2")
 class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
@@ -116,7 +120,9 @@ class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
         upscaled_image = upscaler.upscale(cv2_image)
         pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA")
 
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
+        if choose_torch_device() == torch.device("mps"):
+            mps.empty_cache()
 
         image_dto = context.images.save(image=pil_image)
 

invokeai/app/services/config/config_default.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import locale
 import os
 import re
 import shutil
@@ -27,12 +26,12 @@ DEFAULT_RAM_CACHE = 10.0
 DEFAULT_VRAM_CACHE = 0.25
 DEFAULT_CONVERT_CACHE = 20.0
 DEVICE = Literal["auto", "cpu", "cuda", "cuda:1", "mps"]
-PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
+PRECISION = Literal["auto", "float16", "bfloat16", "float32", "autocast"]
 ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
 ATTENTION_SLICE_SIZE = Literal["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8]
 LOG_FORMAT = Literal["plain", "color", "syslog", "legacy"]
 LOG_LEVEL = Literal["debug", "info", "warning", "error", "critical"]
-CONFIG_SCHEMA_VERSION = "4.0.1"
+CONFIG_SCHEMA_VERSION = "4.0.0"
 
 
 def get_default_ram_cache_size() -> float:
@@ -105,7 +104,7 @@ class InvokeAIAppConfig(BaseSettings):
         lazy_offload: Keep models in VRAM until their space is needed.
         log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.
         device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `cuda:1`, `mps`
-        precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`
+        precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`, `autocast`
         sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
         attention_type: Attention type.<br>Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
         attention_slice_size: Slice size, valid when attention_type=="sliced".<br>Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`
@@ -318,10 +317,11 @@ class InvokeAIAppConfig(BaseSettings):
     @staticmethod
     def find_root() -> Path:
         """Choose the runtime root directory when not specified on command line or init file."""
+        venv = Path(os.environ.get("VIRTUAL_ENV") or ".")
         if os.environ.get("INVOKEAI_ROOT"):
             root = Path(os.environ["INVOKEAI_ROOT"])
-        elif venv := os.environ.get("VIRTUAL_ENV", None):
-            root = Path(venv).parent.resolve()
+        elif any((venv.parent / x).exists() for x in [INIT_FILE, LEGACY_INIT_FILE]):
+            root = (venv.parent).resolve()
         else:
             root = Path("~/invokeai").expanduser().resolve()
         return root
@@ -370,22 +370,16 @@ def migrate_v3_config_dict(config_dict: dict[str, Any]) -> InvokeAIAppConfig:
             # `max_vram_cache_size` was renamed to `vram` some time in v3, but both names were used
             if k == "max_vram_cache_size" and "vram" not in category_dict:
                 parsed_config_dict["vram"] = v
-            # autocast was removed in v4.0.1
-            if k == "precision" and v == "autocast":
-                parsed_config_dict["precision"] = "auto"
             if k == "conf_path":
                 parsed_config_dict["legacy_models_yaml_path"] = v
             if k == "legacy_conf_dir":
-                # The old default for this was "configs/stable-diffusion" ("configs\stable-diffusion" on Windows).
-                if v == "configs/stable-diffusion" or v == "configs\\stable-diffusion":
-                    # If if the incoming config has the default value, skip
-                    continue
-                elif Path(v).name == "stable-diffusion":
-                    # Else if the path ends in "stable-diffusion", we assume the parent is the new correct path.
-                    parsed_config_dict["legacy_conf_dir"] = str(Path(v).parent)
-                else:
-                    # Else we do not attempt to migrate this setting
+                # The old default for this was "configs/stable-diffusion". If if the incoming config has that as the value, we won't set it.
+                # Else if the path ends in "stable-diffusion", we assume the parent is the new correct path.
+                # Else we do not attempt to migrate this setting
+                if v != "configs/stable-diffusion":
                     parsed_config_dict["legacy_conf_dir"] = v
+                elif Path(v).name == "stable-diffusion":
+                    parsed_config_dict["legacy_conf_dir"] = str(Path(v).parent)
             elif k in InvokeAIAppConfig.model_fields:
                 # skip unknown fields
                 parsed_config_dict[k] = v
@@ -395,28 +389,6 @@ def migrate_v3_config_dict(config_dict: dict[str, Any]) -> InvokeAIAppConfig:
     return config
 
 
-def migrate_v4_0_0_config_dict(config_dict: dict[str, Any]) -> InvokeAIAppConfig:
-    """Migrate v4.0.0 config dictionary to a current config object.
-
-    Args:
-        config_dict: A dictionary of settings from a v4.0.0 config file.
-
-    Returns:
-        An instance of `InvokeAIAppConfig` with the migrated settings.
-    """
-    parsed_config_dict: dict[str, Any] = {}
-    for k, v in config_dict.items():
-        # autocast was removed from precision in v4.0.1
-        if k == "precision" and v == "autocast":
-            parsed_config_dict["precision"] = "auto"
-        else:
-            parsed_config_dict[k] = v
-        if k == "schema_version":
-            parsed_config_dict[k] = CONFIG_SCHEMA_VERSION
-    config = DefaultInvokeAIAppConfig.model_validate(parsed_config_dict)
-    return config
-
-
 def load_and_migrate_config(config_path: Path) -> InvokeAIAppConfig:
     """Load and migrate a config file to the latest version.
 
@@ -427,7 +399,7 @@ def load_and_migrate_config(config_path: Path) -> InvokeAIAppConfig:
         An instance of `InvokeAIAppConfig` with the loaded and migrated settings.
     """
     assert config_path.suffix == ".yaml"
-    with open(config_path, "rt", encoding=locale.getpreferredencoding()) as file:
+    with open(config_path) as file:
         loaded_config_dict = yaml.safe_load(file)
 
     assert isinstance(loaded_config_dict, dict)
@@ -443,21 +415,17 @@ def load_and_migrate_config(config_path: Path) -> InvokeAIAppConfig:
             raise RuntimeError(f"Failed to load and migrate v3 config file {config_path}: {e}") from e
         migrated_config.write_file(config_path)
         return migrated_config
-
-    if loaded_config_dict["schema_version"] == "4.0.0":
-        loaded_config_dict = migrate_v4_0_0_config_dict(loaded_config_dict)
-        loaded_config_dict.write_file(config_path)
-
-    # Attempt to load as a v4 config file
-    try:
-        # Meta is not included in the model fields, so we need to validate it separately
-        config = InvokeAIAppConfig.model_validate(loaded_config_dict)
-        assert (
-            config.schema_version == CONFIG_SCHEMA_VERSION
-        ), f"Invalid schema version, expected {CONFIG_SCHEMA_VERSION}: {config.schema_version}"
-        return config
-    except Exception as e:
-        raise RuntimeError(f"Failed to load config file {config_path}: {e}") from e
+    else:
+        # Attempt to load as a v4 config file
+        try:
+            # Meta is not included in the model fields, so we need to validate it separately
+            config = InvokeAIAppConfig.model_validate(loaded_config_dict)
+            assert (
+                config.schema_version == CONFIG_SCHEMA_VERSION
+            ), f"Invalid schema version, expected {CONFIG_SCHEMA_VERSION}: {config.schema_version}"
+            return config
+        except Exception as e:
+            raise RuntimeError(f"Failed to load config file {config_path}: {e}") from e
 
 
 @lru_cache(maxsize=1)

invokeai/app/services/model_install/model_install_default.py

@@ -1,6 +1,5 @@
 """Model installation class."""
 
-import locale
 import os
 import re
 import signal
@@ -13,7 +12,6 @@ from shutil import copyfile, copytree, move, rmtree
 from tempfile import mkdtemp
 from typing import Any, Dict, List, Optional, Union
 
-import torch
 import yaml
 from huggingface_hub import HfFolder
 from pydantic.networks import AnyHttpUrl
@@ -43,7 +41,7 @@ from invokeai.backend.model_manager.metadata.metadata_base import HuggingFaceMet
 from invokeai.backend.model_manager.probe import ModelProbe
 from invokeai.backend.model_manager.search import ModelSearch
 from invokeai.backend.util import InvokeAILogger
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import choose_precision, choose_torch_device
 
 from .model_install_base import (
     MODEL_SOURCE_TO_TYPE_MAP,
@@ -325,8 +323,7 @@ class ModelInstallService(ModelInstallServiceBase):
             legacy_models_yaml_path = Path(self._app_config.root_path, legacy_models_yaml_path)
 
         if legacy_models_yaml_path.exists():
-            with open(legacy_models_yaml_path, "rt", encoding=locale.getpreferredencoding()) as file:
-                legacy_models_yaml = yaml.safe_load(file)
+            legacy_models_yaml = yaml.safe_load(legacy_models_yaml_path.read_text())
 
             yaml_metadata = legacy_models_yaml.pop("__metadata__")
             yaml_version = yaml_metadata.get("version")
@@ -351,13 +348,8 @@ class ModelInstallService(ModelInstallServiceBase):
                     config: dict[str, Any] = {}
                     config["name"] = model_name
                     config["description"] = stanza.get("description")
-                    legacy_config_path = stanza.get("config")
-                    if legacy_config_path:
-                        # In v3, these paths were relative to the root. Migrate them to be relative to the legacy_conf_dir.
-                        legacy_config_path: Path = self._app_config.root_path / legacy_config_path
-                        if legacy_config_path.is_relative_to(self._app_config.legacy_conf_path):
-                            legacy_config_path = legacy_config_path.relative_to(self._app_config.legacy_conf_path)
-                        config["config_path"] = str(legacy_config_path)
+                    config["config_path"] = stanza.get("config")
+
                     try:
                         id = self.register_path(model_path=model_path, config=config)
                         self._logger.info(f"Migrated {model_name} with id {id}")
@@ -376,13 +368,11 @@ class ModelInstallService(ModelInstallServiceBase):
     def delete(self, key: str) -> None:  # noqa D102
         """Unregister the model. Delete its files only if they are within our models directory."""
         model = self.record_store.get_model(key)
-        model_path = self.app_config.models_path / model.path
-
-        if model_path.is_relative_to(self.app_config.models_path):
-            # If the models is in the Invoke-managed models dir, we delete it
+        models_dir = self.app_config.models_path
+        model_path = models_dir / Path(model.path)  # handle legacy relative model paths
+        if model_path.is_relative_to(models_dir):
             self.unconditionally_delete(key)
         else:
-            # Else we only unregister it, leaving the file in place
             self.unregister(key)
 
     def unconditionally_delete(self, key: str) -> None:  # noqa D102
@@ -510,9 +500,9 @@ class ModelInstallService(ModelInstallServiceBase):
     def _scan_for_missing_models(self) -> list[AnyModelConfig]:
         """Scan the models directory for missing models and return a list of them."""
         missing_models: list[AnyModelConfig] = []
-        for model_config in self.record_store.all_models():
-            if not (self.app_config.models_path / model_config.path).resolve().exists():
-                missing_models.append(model_config)
+        for x in self.record_store.all_models():
+            if not Path(x.path).resolve().exists():
+                missing_models.append(x)
         return missing_models
 
     def _register_orphaned_models(self) -> None:
@@ -522,9 +512,7 @@ class ModelInstallService(ModelInstallServiceBase):
         only situations in which we may have orphaned models in the models directory.
         """
 
-        installed_model_paths = {
-            (self._app_config.models_path / x.path).resolve() for x in self.record_store.all_models()
-        }
+        installed_model_paths = {Path(x.path).resolve() for x in self.record_store.all_models()}
 
         # The bool returned by this callback determines if the model is added to the list of models found by the search
         def on_model_found(model_path: Path) -> bool:
@@ -560,21 +548,20 @@ class ModelInstallService(ModelInstallServiceBase):
         May raise an UnknownModelException.
         """
         model = self.record_store.get_model(key)
-        models_dir = self.app_config.models_path
-        old_path = self.app_config.models_path / model.path
+        old_path = Path(model.path).resolve()
+        models_dir = self.app_config.models_path.resolve()
 
         if not old_path.is_relative_to(models_dir):
-            # The model is not in the models directory - we don't need to move it.
             return model
 
-        new_path = models_dir / model.base.value / model.type.value / old_path.name
+        new_path = (models_dir / model.base.value / model.type.value / model.name).with_suffix(old_path.suffix)
 
         if old_path == new_path or new_path.exists() and old_path == new_path.resolve():
             return model
 
         self._logger.info(f"Moving {model.name} to {new_path}.")
         new_path = self._move_model(old_path, new_path)
-        model.path = new_path.relative_to(models_dir).as_posix()
+        model.path = new_path.as_posix()
         self.record_store.update_model(key, ModelRecordChanges(path=model.path))
         return model
 
@@ -613,19 +600,12 @@ class ModelInstallService(ModelInstallServiceBase):
 
         model_path = model_path.resolve()
 
-        # Models in the Invoke-managed models dir should use relative paths.
-        if model_path.is_relative_to(self.app_config.models_path):
-            model_path = model_path.relative_to(self.app_config.models_path)
-
         info.path = model_path.as_posix()
 
+        # Checkpoints have a config file needed for conversion - resolve this to an absolute path
         if isinstance(info, CheckpointConfigBase):
-            # Checkpoints have a config file needed for conversion. Same handling as the model weights - if it's in the
-            # invoke-managed legacy config dir, we use a relative path.
-            legacy_config_path = self.app_config.legacy_conf_path / info.config_path
-            if legacy_config_path.is_relative_to(self.app_config.legacy_conf_path):
-                legacy_config_path = legacy_config_path.relative_to(self.app_config.legacy_conf_path)
-            info.config_path = legacy_config_path.as_posix()
+            legacy_conf = (self.app_config.legacy_conf_path / info.config_path).resolve()
+            info.config_path = legacy_conf.as_posix()
         self.record_store.add_model(info)
         return info.key
 
@@ -635,10 +615,11 @@ class ModelInstallService(ModelInstallServiceBase):
             self._next_job_id += 1
         return id
 
-    def _guess_variant(self) -> Optional[ModelRepoVariant]:
+    @staticmethod
+    def _guess_variant() -> Optional[ModelRepoVariant]:
         """Guess the best HuggingFace variant type to download."""
-        precision = TorchDevice.choose_torch_dtype()
-        return ModelRepoVariant.FP16 if precision == torch.float16 else None
+        precision = choose_precision(choose_torch_device())
+        return ModelRepoVariant.FP16 if precision == "float16" else None
 
     def _import_local_model(self, source: LocalModelSource, config: Optional[Dict[str, Any]]) -> ModelInstallJob:
         return ModelInstallJob(
@@ -754,8 +735,6 @@ class ModelInstallService(ModelInstallServiceBase):
             self._download_cache[download_job.source] = install_job  # matches a download job to an install job
             install_job.download_parts.add(download_job)
 
-        # only start the jobs once install_job.download_parts is fully populated
-        for download_job in install_job.download_parts:
             self._download_queue.submit_download_job(
                 download_job,
                 on_start=self._download_started_callback,
@@ -764,7 +743,6 @@ class ModelInstallService(ModelInstallServiceBase):
                 on_error=self._download_error_callback,
                 on_cancelled=self._download_cancelled_callback,
             )
-
         return install_job
 
     def _stat_size(self, path: Path) -> int:

invokeai/app/services/model_manager/model_manager_default.py

@@ -1,14 +1,12 @@
 # Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
 """Implementation of ModelManagerServiceBase."""
 
-from typing import Optional
-
 import torch
 from typing_extensions import Self
 
 from invokeai.app.services.invoker import Invoker
 from invokeai.backend.model_manager.load import ModelCache, ModelConvertCache, ModelLoaderRegistry
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 from invokeai.backend.util.logging import InvokeAILogger
 
 from ..config import InvokeAIAppConfig
@@ -69,7 +67,7 @@ class ModelManagerService(ModelManagerServiceBase):
         model_record_service: ModelRecordServiceBase,
         download_queue: DownloadQueueServiceBase,
         events: EventServiceBase,
-        execution_device: Optional[torch.device] = None,
+        execution_device: torch.device = choose_torch_device(),
     ) -> Self:
         """
         Construct the model manager service instance.
@@ -82,9 +80,8 @@ class ModelManagerService(ModelManagerServiceBase):
         ram_cache = ModelCache(
             max_cache_size=app_config.ram,
             max_vram_cache_size=app_config.vram,
-            lazy_offloading=app_config.lazy_offload,
             logger=logger,
-            execution_device=execution_device or TorchDevice.choose_torch_device(),
+            execution_device=execution_device,
         )
         convert_cache = ModelConvertCache(cache_path=app_config.convert_cache_path, max_size=app_config.convert_cache)
         loader = ModelLoadService(

invokeai/app/services/session_processor/session_processor_default.py

@@ -70,28 +70,12 @@ class DefaultSessionProcessor(SessionProcessorBase):
     async def _on_queue_event(self, event: FastAPIEvent) -> None:
         event_name = event[1]["event"]
 
-        if (
-            event_name == "session_canceled"
-            and self._queue_item
-            and self._queue_item.item_id == event[1]["data"]["queue_item_id"]
-        ):
-            self._cancel_event.set()
-            self._poll_now()
-        elif (
-            event_name == "queue_cleared"
-            and self._queue_item
-            and self._queue_item.queue_id == event[1]["data"]["queue_id"]
-        ):
+        if event_name == "session_canceled" or event_name == "queue_cleared":
+            # These both mean we should cancel the current session.
             self._cancel_event.set()
             self._poll_now()
         elif event_name == "batch_enqueued":
             self._poll_now()
-        elif event_name == "queue_item_status_changed" and event[1]["data"]["queue_item"]["status"] in [
-            "completed",
-            "failed",
-            "canceled",
-        ]:
-            self._poll_now()
 
     def resume(self) -> SessionProcessorStatus:
         if not self._resume_event.is_set():
@@ -127,146 +111,141 @@ class DefaultSessionProcessor(SessionProcessorBase):
                 poll_now_event.clear()
                 # Middle processor try block; any unhandled exception is a non-fatal processor error
                 try:
-                    # If we are paused, wait for resume event
-                    resume_event.wait()
-
                     # Get the next session to process
                     self._queue_item = self._invoker.services.session_queue.dequeue()
+                    if self._queue_item is not None and resume_event.is_set():
+                        self._invoker.services.logger.debug(f"Executing queue item {self._queue_item.item_id}")
+                        cancel_event.clear()
 
-                    if self._queue_item is None:
-                        # The queue was empty, wait for next polling interval or event to try again
-                        self._invoker.services.logger.debug("Waiting for next polling interval or event")
-                        poll_now_event.wait(self._polling_interval)
-                        continue
+                        # If profiling is enabled, start the profiler
+                        if self._profiler is not None:
+                            self._profiler.start(profile_id=self._queue_item.session_id)
 
-                    self._invoker.services.logger.debug(f"Executing queue item {self._queue_item.item_id}")
-                    cancel_event.clear()
+                        # Prepare invocations and take the first
+                        self._invocation = self._queue_item.session.next()
 
-                    # If profiling is enabled, start the profiler
-                    if self._profiler is not None:
-                        self._profiler.start(profile_id=self._queue_item.session_id)
+                        # Loop over invocations until the session is complete or canceled
+                        while self._invocation is not None and not cancel_event.is_set():
+                            # get the source node id to provide to clients (the prepared node id is not as useful)
+                            source_invocation_id = self._queue_item.session.prepared_source_mapping[self._invocation.id]
 
-                    # Prepare invocations and take the first
-                    self._invocation = self._queue_item.session.next()
+                            # Send starting event
+                            self._invoker.services.events.emit_invocation_started(
+                                queue_batch_id=self._queue_item.batch_id,
+                                queue_item_id=self._queue_item.item_id,
+                                queue_id=self._queue_item.queue_id,
+                                graph_execution_state_id=self._queue_item.session_id,
+                                node=self._invocation.model_dump(),
+                                source_node_id=source_invocation_id,
+                            )
 
-                    # Loop over invocations until the session is complete or canceled
-                    while self._invocation is not None and not cancel_event.is_set():
-                        # get the source node id to provide to clients (the prepared node id is not as useful)
-                        source_invocation_id = self._queue_item.session.prepared_source_mapping[self._invocation.id]
+                            # Innermost processor try block; any unhandled exception is an invocation error & will fail the graph
+                            try:
+                                with self._invoker.services.performance_statistics.collect_stats(
+                                    self._invocation, self._queue_item.session.id
+                                ):
+                                    # Build invocation context (the node-facing API)
+                                    data = InvocationContextData(
+                                        invocation=self._invocation,
+                                        source_invocation_id=source_invocation_id,
+                                        queue_item=self._queue_item,
+                                    )
+                                    context = build_invocation_context(
+                                        data=data,
+                                        services=self._invoker.services,
+                                        cancel_event=self._cancel_event,
+                                    )
 
-                        # Send starting event
-                        self._invoker.services.events.emit_invocation_started(
-                            queue_batch_id=self._queue_item.batch_id,
-                            queue_item_id=self._queue_item.item_id,
-                            queue_id=self._queue_item.queue_id,
-                            graph_execution_state_id=self._queue_item.session_id,
-                            node=self._invocation.model_dump(),
-                            source_node_id=source_invocation_id,
-                        )
+                                    # Invoke the node
+                                    outputs = self._invocation.invoke_internal(
+                                        context=context, services=self._invoker.services
+                                    )
 
-                        # Innermost processor try block; any unhandled exception is an invocation error & will fail the graph
-                        try:
-                            with self._invoker.services.performance_statistics.collect_stats(
-                                self._invocation, self._queue_item.session.id
-                            ):
-                                # Build invocation context (the node-facing API)
-                                data = InvocationContextData(
-                                    invocation=self._invocation,
-                                    source_invocation_id=source_invocation_id,
-                                    queue_item=self._queue_item,
-                                )
-                                context = build_invocation_context(
-                                    data=data,
-                                    services=self._invoker.services,
-                                    cancel_event=self._cancel_event,
+                                    # Save outputs and history
+                                    self._queue_item.session.complete(self._invocation.id, outputs)
+
+                                    # Send complete event
+                                    self._invoker.services.events.emit_invocation_complete(
+                                        queue_batch_id=self._queue_item.batch_id,
+                                        queue_item_id=self._queue_item.item_id,
+                                        queue_id=self._queue_item.queue_id,
+                                        graph_execution_state_id=self._queue_item.session.id,
+                                        node=self._invocation.model_dump(),
+                                        source_node_id=source_invocation_id,
+                                        result=outputs.model_dump(),
+                                    )
+
+                            except KeyboardInterrupt:
+                                # TODO(MM2): Create an event for this
+                                pass
+
+                            except CanceledException:
+                                # When the user cancels the graph, we first set the cancel event. The event is checked
+                                # between invocations, in this loop. Some invocations are long-running, and we need to
+                                # be able to cancel them mid-execution.
+                                #
+                                # For example, denoising is a long-running invocation with many steps. A step callback
+                                # is executed after each step. This step callback checks if the canceled event is set,
+                                # then raises a CanceledException to stop execution immediately.
+                                #
+                                # When we get a CanceledException, we don't need to do anything - just pass and let the
+                                # loop go to its next iteration, and the cancel event will be handled correctly.
+                                pass
+
+                            except Exception as e:
+                                error = traceback.format_exc()
+
+                                # Save error
+                                self._queue_item.session.set_node_error(self._invocation.id, error)
+                                self._invoker.services.logger.error(
+                                    f"Error while invoking session {self._queue_item.session_id}, invocation {self._invocation.id} ({self._invocation.get_type()}):\n{e}"
                                 )
+                                self._invoker.services.logger.error(error)
 
-                                # Invoke the node
-                                outputs = self._invocation.invoke_internal(
-                                    context=context, services=self._invoker.services
-                                )
-
-                                # Save outputs and history
-                                self._queue_item.session.complete(self._invocation.id, outputs)
-
-                                # Send complete event
-                                self._invoker.services.events.emit_invocation_complete(
-                                    queue_batch_id=self._queue_item.batch_id,
+                                # Send error event
+                                self._invoker.services.events.emit_invocation_error(
+                                    queue_batch_id=self._queue_item.session_id,
                                     queue_item_id=self._queue_item.item_id,
                                     queue_id=self._queue_item.queue_id,
                                     graph_execution_state_id=self._queue_item.session.id,
                                     node=self._invocation.model_dump(),
                                     source_node_id=source_invocation_id,
-                                    result=outputs.model_dump(),
+                                    error_type=e.__class__.__name__,
+                                    error=error,
                                 )
+                                pass
 
-                        except KeyboardInterrupt:
-                            # TODO(MM2): Create an event for this
-                            pass
-
-                        except CanceledException:
-                            # When the user cancels the graph, we first set the cancel event. The event is checked
-                            # between invocations, in this loop. Some invocations are long-running, and we need to
-                            # be able to cancel them mid-execution.
-                            #
-                            # For example, denoising is a long-running invocation with many steps. A step callback
-                            # is executed after each step. This step callback checks if the canceled event is set,
-                            # then raises a CanceledException to stop execution immediately.
-                            #
-                            # When we get a CanceledException, we don't need to do anything - just pass and let the
-                            # loop go to its next iteration, and the cancel event will be handled correctly.
-                            pass
-
-                        except Exception as e:
-                            error = traceback.format_exc()
-
-                            # Save error
-                            self._queue_item.session.set_node_error(self._invocation.id, error)
-                            self._invoker.services.logger.error(
-                                f"Error while invoking session {self._queue_item.session_id}, invocation {self._invocation.id} ({self._invocation.get_type()}):\n{e}"
-                            )
-                            self._invoker.services.logger.error(error)
-
-                            # Send error event
-                            self._invoker.services.events.emit_invocation_error(
-                                queue_batch_id=self._queue_item.session_id,
-                                queue_item_id=self._queue_item.item_id,
-                                queue_id=self._queue_item.queue_id,
-                                graph_execution_state_id=self._queue_item.session.id,
-                                node=self._invocation.model_dump(),
-                                source_node_id=source_invocation_id,
-                                error_type=e.__class__.__name__,
-                                error=error,
-                            )
-                            pass
-
-                        # The session is complete if the all invocations are complete or there was an error
-                        if self._queue_item.session.is_complete() or cancel_event.is_set():
-                            # Send complete event
-                            self._invoker.services.events.emit_graph_execution_complete(
-                                queue_batch_id=self._queue_item.batch_id,
-                                queue_item_id=self._queue_item.item_id,
-                                queue_id=self._queue_item.queue_id,
-                                graph_execution_state_id=self._queue_item.session.id,
-                            )
-                            # If we are profiling, stop the profiler and dump the profile & stats
-                            if self._profiler:
-                                profile_path = self._profiler.stop()
-                                stats_path = profile_path.with_suffix(".json")
-                                self._invoker.services.performance_statistics.dump_stats(
-                                    graph_execution_state_id=self._queue_item.session.id, output_path=stats_path
+                            # The session is complete if the all invocations are complete or there was an error
+                            if self._queue_item.session.is_complete() or cancel_event.is_set():
+                                # Send complete event
+                                self._invoker.services.events.emit_graph_execution_complete(
+                                    queue_batch_id=self._queue_item.batch_id,
+                                    queue_item_id=self._queue_item.item_id,
+                                    queue_id=self._queue_item.queue_id,
+                                    graph_execution_state_id=self._queue_item.session.id,
                                 )
-                            # We'll get a GESStatsNotFoundError if we try to log stats for an untracked graph, but in the processor
-                            # we don't care about that - suppress the error.
-                            with suppress(GESStatsNotFoundError):
-                                self._invoker.services.performance_statistics.log_stats(self._queue_item.session.id)
-                                self._invoker.services.performance_statistics.reset_stats()
+                                # If we are profiling, stop the profiler and dump the profile & stats
+                                if self._profiler:
+                                    profile_path = self._profiler.stop()
+                                    stats_path = profile_path.with_suffix(".json")
+                                    self._invoker.services.performance_statistics.dump_stats(
+                                        graph_execution_state_id=self._queue_item.session.id, output_path=stats_path
+                                    )
+                                # We'll get a GESStatsNotFoundError if we try to log stats for an untracked graph, but in the processor
+                                # we don't care about that - suppress the error.
+                                with suppress(GESStatsNotFoundError):
+                                    self._invoker.services.performance_statistics.log_stats(self._queue_item.session.id)
+                                    self._invoker.services.performance_statistics.reset_stats()
 
-                            # Set the invocation to None to prepare for the next session
-                            self._invocation = None
-                        else:
-                            # Prepare the next invocation
-                            self._invocation = self._queue_item.session.next()
+                                # Set the invocation to None to prepare for the next session
+                                self._invocation = None
+                            else:
+                                # Prepare the next invocation
+                                self._invocation = self._queue_item.session.next()
+
+                        # The session is complete, immediately poll for next session
+                        self._queue_item = None
+                        poll_now_event.set()
                     else:
                         # The queue was empty, wait for next polling interval or event to try again
                         self._invoker.services.logger.debug("Waiting for next polling interval or event")

invokeai/app/services/shared/invocation_context.py

@@ -245,18 +245,6 @@ class ImagesInterface(InvocationContextInterface):
         """
         return self._services.images.get_dto(image_name)
 
-    def get_path(self, image_name: str, thumbnail: bool = False) -> Path:
-        """Gets the internal path to an image or thumbnail.
-
-        Args:
-            image_name: The name of the image to get the path of.
-            thumbnail: Get the path of the thumbnail instead of the full image
-
-        Returns:
-            The local path of the image or thumbnail.
-        """
-        return self._services.images.get_path(image_name, thumbnail)
-
 
 class TensorsInterface(InvocationContextInterface):
     def save(self, tensor: Tensor) -> str:

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -10,8 +10,6 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_4 import
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_5 import build_migration_5
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_6 import build_migration_6
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_7 import build_migration_7
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_8 import build_migration_8
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_9 import build_migration_9
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -39,8 +37,6 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_5())
     migrator.register_migration(build_migration_6())
     migrator.register_migration(build_migration_7())
-    migrator.register_migration(build_migration_8(app_config=config))
-    migrator.register_migration(build_migration_9())
     migrator.run_migrations()
 
     return db

invokeai/app/services/shared/sqlite_migrator/migrations/migration_7.py

@@ -11,7 +11,7 @@ class Migration7Callback:
     def _drop_old_models_tables(self, cursor: sqlite3.Cursor) -> None:
         """Drops the old model_records, model_metadata, model_tags and tags tables."""
 
-        tables = ["model_config", "model_metadata", "model_tags", "tags"]
+        tables = ["model_records", "model_metadata", "model_tags", "tags"]
 
         for table in tables:
             cursor.execute(f"DROP TABLE IF EXISTS {table};")

invokeai/app/services/shared/sqlite_migrator/migrations/migration_8.py

@@ -1,91 +0,0 @@
-import sqlite3
-from pathlib import Path
-
-from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-
-class Migration8Callback:
-    def __init__(self, app_config: InvokeAIAppConfig) -> None:
-        self._app_config = app_config
-
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._drop_model_config_table(cursor)
-        self._migrate_abs_models_to_rel(cursor)
-
-    def _drop_model_config_table(self, cursor: sqlite3.Cursor) -> None:
-        """Drops the old model_config table. This was missed in a previous migration."""
-
-        cursor.execute("DROP TABLE IF EXISTS model_config;")
-
-    def _migrate_abs_models_to_rel(self, cursor: sqlite3.Cursor) -> None:
-        """Check all model paths & legacy config paths to determine if they are inside Invoke-managed directories. If
-        they are, update the paths to be relative to the managed directories.
-
-        This migration is a no-op for normal users (their paths will already be relative), but is necessary for users
-        who have been testing the RCs with their live databases. The paths were made absolute in the initial RC, but this
-        change was reverted. To smooth over the revert for our tests, we can migrate the paths back to relative.
-        """
-
-        models_path = self._app_config.models_path
-        legacy_conf_path = self._app_config.legacy_conf_path
-        legacy_conf_dir = self._app_config.legacy_conf_dir
-
-        stmt = """---sql
-        SELECT
-            id,
-            path,
-            json_extract(config, '$.config_path') as config_path
-        FROM models;
-        """
-
-        all_models = cursor.execute(stmt).fetchall()
-
-        for model_id, model_path, model_config_path in all_models:
-            # If the model path is inside the models directory, update it to be relative to the models directory.
-            if Path(model_path).is_relative_to(models_path):
-                new_path = Path(model_path).relative_to(models_path)
-                cursor.execute(
-                    """--sql
-                    UPDATE models
-                    SET config = json_set(config, '$.path', ?)
-                    WHERE id = ?;
-                    """,
-                    (str(new_path), model_id),
-                )
-            # If the model has a legacy config path and it is inside the legacy conf directory, update it to be
-            # relative to the legacy conf directory. This also fixes up cases in which the config path was
-            # incorrectly relativized to the root directory. It will now be relativized to the legacy conf directory.
-            if model_config_path:
-                if Path(model_config_path).is_relative_to(legacy_conf_path):
-                    new_config_path = Path(model_config_path).relative_to(legacy_conf_path)
-                elif Path(model_config_path).is_relative_to(legacy_conf_dir):
-                    new_config_path = Path(*Path(model_config_path).parts[1:])
-                else:
-                    new_config_path = None
-                if new_config_path:
-                    cursor.execute(
-                        """--sql
-                        UPDATE models
-                        SET config = json_set(config, '$.config_path', ?)
-                        WHERE id = ?;
-                        """,
-                        (str(new_config_path), model_id),
-                    )
-
-
-def build_migration_8(app_config: InvokeAIAppConfig) -> Migration:
-    """
-    Build the migration from database version 7 to 8.
-
-    This migration does the following:
-    - Removes the `model_config` table.
-    - Migrates absolute model & legacy config paths to be relative to the models directory.
-    """
-    migration_8 = Migration(
-        from_version=7,
-        to_version=8,
-        callback=Migration8Callback(app_config),
-    )
-
-    return migration_8

invokeai/app/services/shared/sqlite_migrator/migrations/migration_9.py

@@ -1,29 +0,0 @@
-import sqlite3
-
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-
-class Migration9Callback:
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._empty_session_queue(cursor)
-
-    def _empty_session_queue(self, cursor: sqlite3.Cursor) -> None:
-        """Empties the session queue. This is done to prevent any lingering session queue items from causing pydantic errors due to changed schemas."""
-
-        cursor.execute("DELETE FROM session_queue;")
-
-
-def build_migration_9() -> Migration:
-    """
-    Build the migration from database version 8 to 9.
-
-    This migration does the following:
-    - Empties the session queue. This is done to prevent any lingering session queue items from causing pydantic errors due to changed schemas.
-    """
-    migration_9 = Migration(
-        from_version=8,
-        to_version=9,
-        callback=Migration9Callback(),
-    )
-
-    return migration_9

invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_impl.py

@@ -1,6 +1,4 @@
 import sqlite3
-from contextlib import closing
-from datetime import datetime
 from pathlib import Path
 from typing import Optional
 
@@ -34,7 +32,6 @@ class SqliteMigrator:
         self._db = db
         self._logger = db.logger
         self._migration_set = MigrationSet()
-        self._backup_path: Optional[Path] = None
 
     def register_migration(self, migration: Migration) -> None:
         """Registers a migration."""
@@ -58,18 +55,6 @@ class SqliteMigrator:
                 return False
 
             self._logger.info("Database update needed")
-
-            # Make a backup of the db if it needs to be updated and is a file db
-            if self._db.db_path is not None:
-                timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
-                self._backup_path = self._db.db_path.parent / f"{self._db.db_path.stem}_backup_{timestamp}.db"
-                self._logger.info(f"Backing up database to {str(self._backup_path)}")
-                # Use SQLite to do the backup
-                with closing(sqlite3.connect(self._backup_path)) as backup_conn:
-                    self._db.conn.backup(backup_conn)
-            else:
-                self._logger.info("Using in-memory database, no backup needed")
-
             next_migration = self._migration_set.get(from_version=self._get_current_version(cursor))
             while next_migration is not None:
                 self._run_migration(next_migration)

invokeai/backend/image_util/__init__.py

@@ -2,7 +2,7 @@
 Initialization file for invokeai.backend.image_util methods.
 """
 
-from .infill_methods.patchmatch import PatchMatch  # noqa: F401
+from .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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 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 = TorchDevice.choose_torch_device()
+        self.device = 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(self.device)
+        self.model.to(choose_torch_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(self.device)
+        tensor_image = torch.from_numpy(np_image).unsqueeze(0).to(choose_torch_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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 
 from .onnxdet import inference_detector
 from .onnxpose import inference_pose
@@ -28,9 +28,9 @@ config = get_config()
 
 class Wholebody:
     def __init__(self):
-        device = TorchDevice.choose_torch_device()
+        device = choose_torch_device()
 
-        providers = ["CUDAExecutionProvider"] if device.type == "cuda" else ["CPUExecutionProvider"]
+        providers = ["CUDAExecutionProvider"] if device == "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/infill_methods/mosaic.py

@@ -1,60 +0,0 @@
-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

@@ -1,67 +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
-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

@@ -1,95 +0,0 @@
-{
- "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

@@ -1,122 +0,0 @@
-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/lama.py

@@ -7,8 +7,7 @@ from PIL import Image
 
 import invokeai.backend.util.logging as logger
 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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 
 
 def norm_img(np_img):
@@ -29,16 +28,8 @@ def load_jit_model(url_or_path, device):
 
 class LaMA:
     def __call__(self, input_image: Image.Image, *args: Any, **kwds: Any) -> Any:
-        device = TorchDevice.choose_torch_device()
+        device = 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/patchmatch.py

@@ -0,0 +1,49 @@
+"""
+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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 
 """
 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 = TorchDevice.choose_torch_device()
+        self.device = 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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 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 = TorchDevice.choose_torch_device()
+        device = choose_torch_device()
         features = cls.feature_extractor([image], return_tensors="pt")
         features.to(device)
         cls.safety_checker.to(device)

invokeai/backend/ip_adapter/attention_processor.py

@@ -0,0 +1,182 @@
+# 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,32 +1,21 @@
 # copied from https://github.com/tencent-ailab/IP-Adapter (Apache License 2.0)
 #   and modified as needed
 
-import pathlib
-from typing import List, Optional, TypedDict, Union
+from typing import Optional, Union
 
-import safetensors
-import safetensors.torch
 import torch
 from PIL import Image
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 
 from invokeai.backend.ip_adapter.ip_attention_weights import IPAttentionWeights
 
-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: int = 1024, clip_embeddings_dim: int = 1024, clip_extra_context_tokens: int = 4
-    ):
+    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
         super().__init__()
 
         self.cross_attention_dim = cross_attention_dim
@@ -35,7 +24,7 @@ class ImageProjModel(torch.nn.Module):
         self.norm = torch.nn.LayerNorm(cross_attention_dim)
 
     @classmethod
-    def from_state_dict(cls, state_dict: dict[str, torch.Tensor], clip_extra_context_tokens: int = 4):
+    def from_state_dict(cls, state_dict: dict[torch.Tensor], clip_extra_context_tokens=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.
@@ -55,7 +44,7 @@ class ImageProjModel(torch.nn.Module):
         model.load_state_dict(state_dict)
         return model
 
-    def forward(self, image_embeds: torch.Tensor):
+    def forward(self, image_embeds):
         embeds = image_embeds
         clip_extra_context_tokens = self.proj(embeds).reshape(
             -1, self.clip_extra_context_tokens, self.cross_attention_dim
@@ -67,7 +56,7 @@ class ImageProjModel(torch.nn.Module):
 class MLPProjModel(torch.nn.Module):
     """SD model with image prompt"""
 
-    def __init__(self, cross_attention_dim: int = 1024, clip_embeddings_dim: int = 1024):
+    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024):
         super().__init__()
 
         self.proj = torch.nn.Sequential(
@@ -78,7 +67,7 @@ class MLPProjModel(torch.nn.Module):
         )
 
     @classmethod
-    def from_state_dict(cls, state_dict: dict[str, torch.Tensor]):
+    def from_state_dict(cls, state_dict: dict[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.
@@ -97,17 +86,17 @@ class MLPProjModel(torch.nn.Module):
         model.load_state_dict(state_dict)
         return model
 
-    def forward(self, image_embeds: torch.Tensor):
+    def forward(self, image_embeds):
         clip_extra_context_tokens = self.proj(image_embeds)
         return clip_extra_context_tokens
 
 
-class IPAdapter(RawModel):
+class IPAdapter(torch.nn.Module):
     """IP-Adapter: https://arxiv.org/pdf/2308.06721.pdf"""
 
     def __init__(
         self,
-        state_dict: IPAdapterStateDict,
+        state_dict: dict[str, torch.Tensor],
         device: torch.device,
         dtype: torch.dtype = torch.float16,
         num_tokens: int = 4,
@@ -139,27 +128,24 @@ 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: dict[str, torch.Tensor]
-    ) -> Union[ImageProjModel, Resampler, MLPProjModel]:
+    def _init_image_proj_model(self, state_dict):
         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: List[Image.Image], image_encoder: CLIPVisionModelWithProjection):
+    def get_image_embeds(self, pil_image, image_encoder: CLIPVisionModelWithProjection):
+        if isinstance(pil_image, Image.Image):
+            pil_image = [pil_image]
         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
-        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
+        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
 
 
 class IPAdapterPlus(IPAdapter):
     """IP-Adapter with fine-grained features"""
 
-    def _init_image_proj_model(self, state_dict: dict[str, torch.Tensor]) -> Union[Resampler, MLPProjModel]:
+    def _init_image_proj_model(self, state_dict):
         return Resampler.from_state_dict(
             state_dict=state_dict,
             depth=4,
@@ -170,32 +156,31 @@ class IPAdapterPlus(IPAdapter):
         ).to(self.device, dtype=self.dtype)
 
     @torch.inference_mode()
-    def get_image_embeds(self, pil_image: List[Image.Image], image_encoder: CLIPVisionModelWithProjection):
+    def get_image_embeds(self, pil_image, image_encoder: CLIPVisionModelWithProjection):
+        if isinstance(pil_image, Image.Image):
+            pil_image = [pil_image]
         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
         ]
-        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
+        uncond_image_prompt_embeds = self._image_proj_model(uncond_clip_image_embeds)
+        return image_prompt_embeds, uncond_image_prompt_embeds
 
 
 class IPAdapterFull(IPAdapterPlus):
     """IP-Adapter Plus with full features."""
 
-    def _init_image_proj_model(self, state_dict: dict[str, torch.Tensor]):
+    def _init_image_proj_model(self, state_dict: dict[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: dict[str, torch.Tensor]):
+    def _init_image_proj_model(self, state_dict):
         return Resampler.from_state_dict(
             state_dict=state_dict,
             depth=4,
@@ -206,48 +191,24 @@ class IPAdapterPlusXL(IPAdapterPlus):
         ).to(self.device, dtype=self.dtype)
 
 
-def load_ip_adapter_tensors(ip_adapter_ckpt_path: pathlib.Path, device: str) -> IPAdapterStateDict:
-    state_dict: IPAdapterStateDict = {"ip_adapter": {}, "image_proj": {}}
-
-    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)
+    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")
 
-    # IPAdapter (with ImageProjModel)
-    if "proj.weight" in state_dict["image_proj"]:
+    if "proj.weight" in state_dict["image_proj"]:  # IPAdapter (with ImageProjModel).
         return IPAdapter(state_dict, device=device, dtype=dtype)
-
-    # IPAdaterPlus or IPAdapterPlusXL (with Resampler)
-    elif "proj_in.weight" in state_dict["image_proj"]:
+    elif "proj_in.weight" in state_dict["image_proj"]:  # IPAdaterPlus or IPAdapterPlusXL (with Resampler).
         cross_attention_dim = state_dict["ip_adapter"]["1.to_k_ip.weight"].shape[-1]
         if cross_attention_dim == 768:
-            return IPAdapterPlus(state_dict, device=device, dtype=dtype)  # SD1 IP-Adapter Plus
+            # SD1 IP-Adapter Plus
+            return IPAdapterPlus(state_dict, device=device, dtype=dtype)
         elif cross_attention_dim == 2048:
-            return IPAdapterPlusXL(state_dict, device=device, dtype=dtype)  # SDXL IP-Adapter Plus
+            # SDXL IP-Adapter Plus
+            return IPAdapterPlusXL(state_dict, device=device, dtype=dtype)
         else:
             raise Exception(f"Unsupported IP-Adapter Plus cross-attention dimension: {cross_attention_dim}.")
-
-    # IPAdapterFull (with MLPProjModel)
-    elif "proj.0.weight" in state_dict["image_proj"]:
+    elif "proj.0.weight" in state_dict["image_proj"]:  # IPAdapterFull (with MLPProjModel).
         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: int, mult: int = 4):
-    inner_dim = dim * mult
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
     return nn.Sequential(
         nn.LayerNorm(dim),
         nn.Linear(dim, inner_dim, bias=False),
@@ -19,8 +19,8 @@ def FeedForward(dim: int, mult: int = 4):
     )
 
 
-def reshape_tensor(x: torch.Tensor, heads: int):
-    bs, length, _ = x.shape
+def reshape_tensor(x, heads):
+    bs, length, width = 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: torch.Tensor, heads: int):
 
 
 class PerceiverAttention(nn.Module):
-    def __init__(self, *, dim: int, dim_head: int = 64, heads: int = 8):
+    def __init__(self, *, dim, dim_head=64, heads=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: torch.Tensor, latents: torch.Tensor):
+    def forward(self, x, latents):
         """
         Args:
             x (torch.Tensor): image features
@@ -80,14 +80,14 @@ class PerceiverAttention(nn.Module):
 class Resampler(nn.Module):
     def __init__(
         self,
-        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,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
     ):
         super().__init__()
 
@@ -110,15 +110,7 @@ class Resampler(nn.Module):
             )
 
     @classmethod
-    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,
-    ):
+    def from_state_dict(cls, state_dict: dict[torch.Tensor], depth=8, dim_head=64, heads=16, num_queries=8, ff_mult=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
@@ -153,7 +145,7 @@ class Resampler(nn.Module):
         model.load_state_dict(state_dict)
         return model
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x):
         latents = self.latents.repeat(x.size(0), 1, 1)
 
         x = self.proj_in(x)

invokeai/backend/ip_adapter/unet_patcher.py

@@ -0,0 +1,53 @@
+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/lora_model_patcher.py

@@ -0,0 +1,65 @@
+from contextlib import contextmanager
+from typing import Iterator, Tuple, Union
+
+from diffusers.loaders.lora import LoraLoaderMixin
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+from diffusers.utils.peft_utils import recurse_remove_peft_layers
+from transformers import CLIPTextModel
+
+from invokeai.backend.lora_model_raw import LoRAModelRaw
+
+
+class LoraModelPatcher:
+    @classmethod
+    def unload_lora_from_model(cls, m: Union[UNet2DConditionModel, CLIPTextModel]):
+        """Unload all LoRA models from a UNet or Text Encoder.
+        This implementation is base on LoraLoaderMixin.unload_lora_weights().
+        """
+        recurse_remove_peft_layers(m)
+        if hasattr(m, "peft_config"):
+            del m.peft_config  # type: ignore
+        if hasattr(m, "_hf_peft_config_loaded"):
+            m._hf_peft_config_loaded = None  # type: ignore
+
+    @classmethod
+    @contextmanager
+    def apply_lora_to_unet(cls, unet: UNet2DConditionModel, loras: Iterator[Tuple[LoRAModelRaw, float]]):
+        try:
+            # TODO(ryand): Test speed of low_cpu_mem_usage=True.
+            for lora, lora_weight in loras:
+                LoraLoaderMixin.load_lora_into_unet(
+                    state_dict=lora.state_dict,
+                    network_alphas=lora.network_alphas,
+                    unet=unet,
+                    low_cpu_mem_usage=True,
+                    adapter_name=lora.name,
+                    _pipeline=None,
+                )
+            yield
+        finally:
+            cls.unload_lora_from_model(unet)
+
+    @classmethod
+    @contextmanager
+    def apply_lora_to_text_encoder(
+        cls, text_encoder: CLIPTextModel, loras: Iterator[Tuple[LoRAModelRaw, float]], prefix: str
+    ):
+        assert prefix in ["text_encoder", "text_encoder_2"]
+        try:
+            for lora, lora_weight in loras:
+                # Filter the state_dict to only include the keys that start with the prefix.
+                text_encoder_state_dict = {
+                    key: value for key, value in lora.state_dict.items() if key.startswith(prefix + ".")
+                }
+                if len(text_encoder_state_dict) > 0:
+                    LoraLoaderMixin.load_lora_into_text_encoder(
+                        state_dict=text_encoder_state_dict,
+                        network_alphas=lora.network_alphas,
+                        text_encoder=text_encoder,
+                        low_cpu_mem_usage=True,
+                        adapter_name=lora.name,
+                        _pipeline=None,
+                    )
+            yield
+        finally:
+            cls.unload_lora_from_model(text_encoder)

invokeai/backend/lora_model_raw.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+from diffusers.loaders.lora import LoraLoaderMixin
+from typing_extensions import Self
+
+
+class LoRAModelRaw:
+    def __init__(
+        self,
+        name: str,
+        state_dict: dict[str, torch.Tensor],
+        network_alphas: Optional[dict[str, float]],
+    ):
+        self._name = name
+        self.state_dict = state_dict
+        self.network_alphas = network_alphas
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        for key, layer in self.state_dict.items():
+            self.state_dict[key] = layer.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        """Calculate the size of the model in bytes."""
+        model_size = 0
+        for layer in self.state_dict.values():
+            model_size += layer.numel() * layer.element_size()
+        return model_size
+
+    @classmethod
+    def from_checkpoint(
+        cls, file_path: Union[str, Path], device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None
+    ) -> Self:
+        """This function is based on diffusers LoraLoaderMixin.load_lora_weights()."""
+
+        file_path = Path(file_path)
+        if file_path.is_dir():
+            raise NotImplementedError("LoRA models from directories are not yet supported.")
+
+        dir_path = file_path.parent
+        file_name = file_path.name
+
+        state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(
+            pretrained_model_name_or_path_or_dict=str(file_path), local_files_only=True, weight_name=str(file_name)
+        )
+
+        is_correct_format = all("lora" in key for key in state_dict.keys())
+        if not is_correct_format:
+            raise ValueError("Invalid LoRA checkpoint.")
+
+        model = cls(
+            # TODO(ryand): Handle both files and directories here?
+            name=Path(file_path).stem,
+            state_dict=state_dict,
+            network_alphas=network_alphas,
+        )
+
+        device = device or torch.device("cpu")
+        dtype = dtype or torch.float32
+        model.to(device=device, dtype=dtype)
+        return model

invokeai/backend/lora_old.py

@@ -11,8 +11,6 @@ from typing_extensions import Self
 
 from invokeai.backend.model_manager import BaseModelType
 
-from .raw_model import RawModel
-
 
 class LoRALayerBase:
     # rank: Optional[int]
@@ -368,7 +366,7 @@ class IA3Layer(LoRALayerBase):
 AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer]
 
 
-class LoRAModelRaw(RawModel):  # (torch.nn.Module):
+class LoRAModelRaw(torch.nn.Module):
     _name: str
     layers: Dict[str, AnyLoRALayer]
 

invokeai/backend/model_manager/__init__.py

@@ -33,3 +33,42 @@ __all__ = [
     "SchedulerPredictionType",
     "SubModelType",
 ]
+
+########## to help populate the openapi_schema with format enums for each config ###########
+# This code is no longer necessary?
+# leave it here just in case
+#
+# import inspect
+# from enum import Enum
+# from typing import Any, Iterable, Dict, get_args, Set
+# def _expand(something: Any) -> Iterable[type]:
+#     if isinstance(something, type):
+#         yield something
+#     else:
+#         for x in get_args(something):
+#             for y in _expand(x):
+#                 yield y
+
+# def _find_format(cls: type) -> Iterable[Enum]:
+#     if hasattr(inspect, "get_annotations"):
+#         fields = inspect.get_annotations(cls)
+#     else:
+#         fields = cls.__annotations__
+#     if "format" in fields:
+#         for x in get_args(fields["format"]):
+#             yield x
+#     for parent_class in cls.__bases__:
+#         for x in _find_format(parent_class):
+#             yield x
+#     return None
+
+# def get_model_config_formats() -> Dict[str, Set[Enum]]:
+#     result: Dict[str, Set[Enum]] = {}
+#     for model_config in _expand(AnyModelConfig):
+#         for field in _find_format(model_config):
+#             if field is None:
+#                 continue
+#             if not result.get(model_config.__qualname__):
+#                 result[model_config.__qualname__] = set()
+#             result[model_config.__qualname__].add(field)
+#     return result

invokeai/backend/model_manager/config.py

@@ -31,12 +31,13 @@ from typing_extensions import Annotated, Any, Dict
 
 from invokeai.app.invocations.constants import SCHEDULER_NAME_VALUES
 from invokeai.app.util.misc import uuid_string
-
-from ..raw_model import RawModel
+from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
+from invokeai.backend.lora_model_raw import LoRAModelRaw
+from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
+from invokeai.backend.textual_inversion import TextualInversionModelRaw
 
 # ModelMixin is the base class for all diffusers and transformers models
-# RawModel is the InvokeAI wrapper class for ip_adapters, loras, textual_inversion and onnx runtime
-AnyModel = Union[ModelMixin, RawModel, torch.nn.Module]
+AnyModel = Union[ModelMixin, torch.nn.Module, IPAdapter, LoRAModelRaw, TextualInversionModelRaw, IAIOnnxRuntimeModel]
 
 
 class InvalidModelConfigException(Exception):
@@ -323,13 +324,10 @@ class MainDiffusersConfig(DiffusersConfigBase, MainConfigBase):
         return Tag(f"{ModelType.Main.value}.{ModelFormat.Diffusers.value}")
 
 
-class IPAdapterBaseConfig(ModelConfigBase):
+class IPAdapterConfig(ModelConfigBase):
+    """Model config for IP Adaptor format models."""
+
     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]
 
@@ -338,16 +336,6 @@ class IPAdapterInvokeAIConfig(IPAdapterBaseConfig):
         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."""
 
@@ -403,8 +391,7 @@ AnyModelConfig = Annotated[
         Annotated[LoRADiffusersConfig, LoRADiffusersConfig.get_tag()],
         Annotated[TextualInversionFileConfig, TextualInversionFileConfig.get_tag()],
         Annotated[TextualInversionFolderConfig, TextualInversionFolderConfig.get_tag()],
-        Annotated[IPAdapterInvokeAIConfig, IPAdapterInvokeAIConfig.get_tag()],
-        Annotated[IPAdapterCheckpointConfig, IPAdapterCheckpointConfig.get_tag()],
+        Annotated[IPAdapterConfig, IPAdapterConfig.get_tag()],
         Annotated[T2IAdapterConfig, T2IAdapterConfig.get_tag()],
         Annotated[CLIPVisionDiffusersConfig, CLIPVisionDiffusersConfig.get_tag()],
     ],

invokeai/backend/model_manager/convert_ckpt_to_diffusers.py

@@ -3,10 +3,10 @@
 """Conversion script for the Stable Diffusion checkpoints."""
 
 from pathlib import Path
-from typing import Optional
+from typing import Dict
 
 import torch
-from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers import AutoencoderKL
 from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
     convert_ldm_vae_checkpoint,
     create_vae_diffusers_config,
@@ -15,14 +15,11 @@ from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
 )
 from omegaconf import DictConfig
 
-from . import AnyModel
-
 
 def convert_ldm_vae_to_diffusers(
-    checkpoint: torch.Tensor | dict[str, torch.Tensor],
+    checkpoint: Dict[str, torch.Tensor],
     vae_config: DictConfig,
     image_size: int,
-    dump_path: Optional[Path] = None,
     precision: torch.dtype = torch.float16,
 ) -> AutoencoderKL:
     """Convert a checkpoint-style VAE into a Diffusers VAE"""
@@ -31,21 +28,16 @@ def convert_ldm_vae_to_diffusers(
 
     vae = AutoencoderKL(**vae_config)
     vae.load_state_dict(converted_vae_checkpoint)
-    vae.to(precision)
-
-    if dump_path:
-        vae.save_pretrained(dump_path, safe_serialization=True)
-
-    return vae
+    return vae.to(precision)
 
 
 def convert_ckpt_to_diffusers(
     checkpoint_path: str | Path,
-    dump_path: Optional[str | Path] = None,
+    dump_path: str | Path,
     precision: torch.dtype = torch.float16,
     use_safetensors: bool = True,
     **kwargs,
-) -> AnyModel:
+):
     """
     Takes all the arguments of download_from_original_stable_diffusion_ckpt(),
     and in addition a path-like object indicating the location of the desired diffusers
@@ -55,20 +47,18 @@ def convert_ckpt_to_diffusers(
     pipe = pipe.to(precision)
 
     # TO DO: save correct repo variant
-    if dump_path:
-        pipe.save_pretrained(
-            dump_path,
-            safe_serialization=use_safetensors,
-        )
-    return pipe
+    pipe.save_pretrained(
+        dump_path,
+        safe_serialization=use_safetensors,
+    )
 
 
 def convert_controlnet_to_diffusers(
     checkpoint_path: Path,
-    dump_path: Optional[Path] = None,
+    dump_path: Path,
     precision: torch.dtype = torch.float16,
     **kwargs,
-) -> AnyModel:
+):
     """
     Takes all the arguments of download_controlnet_from_original_ckpt(),
     and in addition a path-like object indicating the location of the desired diffusers
@@ -78,6 +68,4 @@ def convert_controlnet_to_diffusers(
     pipe = pipe.to(precision)
 
     # TO DO: save correct repo variant
-    if dump_path:
-        pipe.save_pretrained(dump_path, safe_serialization=True)
-    return pipe
+    pipe.save_pretrained(dump_path, safe_serialization=True)

invokeai/backend/model_manager/load/convert_cache/convert_cache_default.py

@@ -19,20 +19,11 @@ class ModelConvertCache(ModelConvertCacheBase):
         self._cache_path = cache_path
         self._max_size = max_size
 
-        # adjust cache size at startup in case it has been changed
-        if self._cache_path.exists():
-            self.make_room(0.0)
-
     @property
     def max_size(self) -> float:
         """Return the maximum size of this cache directory (GB)."""
         return self._max_size
 
-    @max_size.setter
-    def max_size(self, value: float) -> None:
-        """Set the maximum size of this cache directory (GB)."""
-        self._max_size = value
-
     def cache_path(self, key: str) -> Path:
         """Return the path for a model with the indicated key."""
         return self._cache_path / key

invokeai/backend/model_manager/load/load_base.py

@@ -83,15 +83,3 @@ class ModelLoaderBase(ABC):
     ) -> int:
         """Return size in bytes of the model, calculated before loading."""
         pass
-
-    @property
-    @abstractmethod
-    def convert_cache(self) -> ModelConvertCacheBase:
-        """Return the convert cache associated with this loader."""
-        pass
-
-    @property
-    @abstractmethod
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
-        """Return the ram cache associated with this loader."""
-        pass

invokeai/backend/model_manager/load/load_default.py

@@ -3,13 +3,14 @@
 
 from logging import Logger
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Tuple
 
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.backend.model_manager import (
     AnyModel,
     AnyModelConfig,
     InvalidModelConfigException,
+    ModelRepoVariant,
     SubModelType,
 )
 from invokeai.backend.model_manager.config import DiffusersConfigBase, ModelType
@@ -18,7 +19,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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device, torch_dtype
 
 
 # TO DO: The loader is not thread safe!
@@ -37,7 +38,7 @@ class ModelLoader(ModelLoaderBase):
         self._logger = logger
         self._ram_cache = ram_cache
         self._convert_cache = convert_cache
-        self._torch_dtype = TorchDevice.choose_torch_dtype()
+        self._torch_dtype = torch_dtype(choose_torch_device(), app_config)
 
     def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
         """
@@ -53,43 +54,51 @@ class ModelLoader(ModelLoaderBase):
         if model_config.type is ModelType.Main and not submodel_type:
             raise InvalidModelConfigException("submodel_type is required when loading a main model")
 
-        model_path = self._get_model_path(model_config)
+        model_path, model_config, submodel_type = self._get_model_path(model_config, submodel_type)
 
         if not model_path.exists():
             raise InvalidModelConfigException(f"Files for model '{model_config.name}' not found at {model_path}")
 
-        with skip_torch_weight_init():
-            locker = self._convert_and_load(model_config, model_path, submodel_type)
+        model_path = self._convert_if_needed(model_config, model_path, submodel_type)
+        locker = self._load_if_needed(model_config, model_path, submodel_type)
         return LoadedModel(config=model_config, _locker=locker)
 
-    @property
-    def convert_cache(self) -> ModelConvertCacheBase:
-        """Return the convert cache associated with this loader."""
-        return self._convert_cache
-
-    @property
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
-        """Return the ram cache associated with this loader."""
-        return self._ram_cache
-
-    def _get_model_path(self, config: AnyModelConfig) -> Path:
+    def _get_model_path(
+        self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+    ) -> Tuple[Path, AnyModelConfig, Optional[SubModelType]]:
         model_base = self._app_config.models_path
-        return (model_base / config.path).resolve()
+        result = (model_base / config.path).resolve(), config, submodel_type
+        return result
 
-    def _convert_and_load(
+    def _convert_if_needed(
+        self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None
+    ) -> Path:
+        cache_path: Path = self._convert_cache.cache_path(config.key)
+
+        if not self._needs_conversion(config, model_path, cache_path):
+            return cache_path if cache_path.exists() else model_path
+
+        self._convert_cache.make_room(self.get_size_fs(config, model_path, submodel_type))
+        return self._convert_model(config, model_path, cache_path)
+
+    def _needs_conversion(self, config: AnyModelConfig, model_path: Path, dest_path: Path) -> bool:
+        return False
+
+    def _load_if_needed(
         self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None
     ) -> ModelLockerBase:
+        # TO DO: This is not thread safe!
         try:
             return self._ram_cache.get(config.key, submodel_type)
         except IndexError:
             pass
 
-        cache_path: Path = self._convert_cache.cache_path(config.key)
-        if self._needs_conversion(config, model_path, cache_path):
-            loaded_model = self._do_convert(config, model_path, cache_path, submodel_type)
-        else:
-            config.path = str(cache_path) if cache_path.exists() else str(self._get_model_path(config))
-            loaded_model = self._load_model(config, submodel_type)
+        model_variant = getattr(config, "repo_variant", None)
+        self._ram_cache.make_room(self.get_size_fs(config, model_path, submodel_type))
+
+        # This is where the model is actually loaded!
+        with skip_torch_weight_init():
+            loaded_model = self._load_model(model_path, model_variant=model_variant, submodel_type=submodel_type)
 
         self._ram_cache.put(
             config.key,
@@ -114,34 +123,15 @@ class ModelLoader(ModelLoaderBase):
             variant=config.repo_variant if isinstance(config, DiffusersConfigBase) else None,
         )
 
-    def _do_convert(
-        self, config: AnyModelConfig, model_path: Path, cache_path: Path, submodel_type: Optional[SubModelType] = None
-    ) -> AnyModel:
-        self.convert_cache.make_room(calc_model_size_by_fs(model_path))
-        pipeline = self._convert_model(config, model_path, cache_path if self.convert_cache.max_size > 0 else None)
-        if submodel_type:
-            # Proactively load the various submodels into the RAM cache so that we don't have to re-convert
-            # the entire pipeline every time a new submodel is needed.
-            for subtype in SubModelType:
-                if subtype == submodel_type:
-                    continue
-                if submodel := getattr(pipeline, subtype.value, None):
-                    self._ram_cache.put(
-                        config.key, submodel_type=subtype, model=submodel, size=calc_model_size_by_data(submodel)
-                    )
-        return getattr(pipeline, submodel_type.value) if submodel_type else pipeline
-
-    def _needs_conversion(self, config: AnyModelConfig, model_path: Path, dest_path: Path) -> bool:
-        return False
-
     # This needs to be implemented in subclasses that handle checkpoints
-    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel:
+    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Path) -> Path:
         raise NotImplementedError
 
     # This needs to be implemented in the subclass
     def _load_model(
         self,
-        config: AnyModelConfig,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
         raise NotImplementedError

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) -> Optional[CacheStats]:
+    def stats(self) -> CacheStats:
         """Return collected CacheStats object."""
         pass
 

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

@@ -30,12 +30,15 @@ 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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device
 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
@@ -119,11 +122,6 @@ class ModelCache(ModelCacheBase[AnyModel]):
         """Return the cap on cache size."""
         return self._max_cache_size
 
-    @max_cache_size.setter
-    def max_cache_size(self, value: float) -> None:
-        """Set the cap on cache size."""
-        self._max_cache_size = value
-
     @property
     def stats(self) -> Optional[CacheStats]:
         """Return collected CacheStats object."""
@@ -159,9 +157,8 @@ class ModelCache(ModelCacheBase[AnyModel]):
     ) -> None:
         """Store model under key and optional submodel_type."""
         key = self._make_cache_key(key, submodel_type)
-        if key in self._cached_models:
-            return
-        self.make_room(size)
+        assert key not in self._cached_models
+
         cache_record = CacheRecord(key, model, size)
         self._cached_models[key] = cache_record
         self._cache_stack.append(key)
@@ -241,7 +238,9 @@ 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"
                 )
 
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
+        if choose_torch_device() == torch.device("mps"):
+            mps.empty_cache()
 
     def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device: torch.device) -> None:
         """Move model into the indicated device.
@@ -264,14 +263,12 @@ 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()
-        try:
-            cache_entry.model.to(target_device)
-        except Exception as e:  # blow away cache entry
-            self._delete_cache_entry(cache_entry)
-            raise e
-
+        cache_entry.model.to(target_device)
         snapshot_after = self._capture_memory_snapshot()
         end_model_to_time = time.time()
         self.logger.debug(
@@ -326,11 +323,11 @@ class ModelCache(ModelCacheBase[AnyModel]):
                     f" {in_ram_models}/{in_vram_models}({locked_in_vram_models})"
                 )
 
-    def make_room(self, size: int) -> None:
+    def make_room(self, model_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 = size
+        bytes_needed = model_size
         maximum_size = self.max_cache_size * GIG  # stored in GB, convert to bytes
         current_size = self.cache_size()
 
@@ -385,11 +382,12 @@ 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 {(size/GIG):.2f} GB (-{(cache_entry.size/GIG):.2f} GB)"
+                    f"Removing {model_key} from RAM cache to free at least {(model_size/GIG):.2f} GB (-{(cache_entry.size/GIG):.2f} GB)"
                 )
                 current_size -= cache_entry.size
                 models_cleared += 1
-                self._delete_cache_entry(cache_entry)
+                del self._cache_stack[pos]
+                del self._cached_models[model_key]
                 del cache_entry
 
             else:
@@ -407,13 +405,20 @@ class ModelCache(ModelCacheBase[AnyModel]):
             #
             # Keep in mind that gc is only responsible for handling reference cycles. Most objects should be cleaned up
             # immediately when their reference count hits 0.
-            if self.stats:
-                self.stats.cleared = models_cleared
             gc.collect()
 
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
+        if choose_torch_device() == torch.device("mps"):
+            mps.empty_cache()
+
         self.logger.debug(f"After making room: cached_models={len(self._cached_models)}")
 
-    def _delete_cache_entry(self, cache_entry: CacheRecord[AnyModel]) -> None:
-        self._cache_stack.remove(cache_entry.key)
-        del self._cached_models[cache_entry.key]
+    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

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

@@ -34,6 +34,7 @@ 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)
@@ -50,7 +51,6 @@ class ModelLocker(ModelLockerBase):
         except Exception:
             self._cache_entry.unlock()
             raise
-
         return self.model
 
     def unlock(self) -> None:

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

@@ -2,10 +2,8 @@
 """Class for ControlNet model loading in InvokeAI."""
 
 from pathlib import Path
-from typing import Optional
 
 from invokeai.backend.model_manager import (
-    AnyModel,
     AnyModelConfig,
     BaseModelType,
     ModelFormat,
@@ -35,7 +33,7 @@ class ControlNetLoader(GenericDiffusersLoader):
         else:
             return True
 
-    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel:
+    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Path) -> Path:
         assert isinstance(config, CheckpointConfigBase)
         image_size = (
             512
@@ -46,8 +44,8 @@ class ControlNetLoader(GenericDiffusersLoader):
         )
 
         self._logger.info(f"Converting {model_path} to diffusers format")
-        with open(self._app_config.legacy_conf_path / config.config_path, "r") as config_stream:
-            result = convert_controlnet_to_diffusers(
+        with open(self._app_config.root_path / config.config_path, "r") as config_stream:
+            convert_controlnet_to_diffusers(
                 model_path,
                 output_path,
                 original_config_file=config_stream,
@@ -55,4 +53,4 @@ class ControlNetLoader(GenericDiffusersLoader):
                 precision=self._torch_dtype,
                 from_safetensors=model_path.suffix == ".safetensors",
             )
-        return result
+        return output_path

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

@@ -10,14 +10,13 @@ from diffusers.models.modeling_utils import ModelMixin
 
 from invokeai.backend.model_manager import (
     AnyModel,
-    AnyModelConfig,
     BaseModelType,
     InvalidModelConfigException,
     ModelFormat,
+    ModelRepoVariant,
     ModelType,
     SubModelType,
 )
-from invokeai.backend.model_manager.config import DiffusersConfigBase
 
 from .. import ModelLoader, ModelLoaderRegistry
 
@@ -29,15 +28,14 @@ class GenericDiffusersLoader(ModelLoader):
 
     def _load_model(
         self,
-        config: AnyModelConfig,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
-        model_path = Path(config.path)
         model_class = self.get_hf_load_class(model_path)
         if submodel_type is not None:
             raise Exception(f"There are no submodels in models of type {model_class}")
-        repo_variant = config.repo_variant if isinstance(config, DiffusersConfigBase) else None
-        variant = repo_variant.value if repo_variant else None
+        variant = model_variant.value if model_variant else None
         try:
             result: AnyModel = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, variant=variant)
         except OSError as e:

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

@@ -7,26 +7,31 @@ 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,
+    BaseModelType,
+    ModelFormat,
+    ModelRepoVariant,
+    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."""
 
     def _load_model(
         self,
-        config: AnyModelConfig,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
         if submodel_type is not None:
             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=model_path,
+        model = build_ip_adapter(
+            ip_adapter_ckpt_path=str(model_path / "ip_adapter.bin"),
             device=torch.device("cpu"),
             dtype=self._torch_dtype,
         )

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

@@ -3,15 +3,16 @@
 
 from logging import Logger
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Tuple
 
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.lora_model_raw import LoRAModelRaw
 from invokeai.backend.model_manager import (
     AnyModel,
     AnyModelConfig,
     BaseModelType,
     ModelFormat,
+    ModelRepoVariant,
     ModelType,
     SubModelType,
 )
@@ -40,24 +41,26 @@ class LoRALoader(ModelLoader):
 
     def _load_model(
         self,
-        config: AnyModelConfig,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
         if submodel_type is not None:
             raise ValueError("There are no submodels in a LoRA model.")
-        model_path = Path(config.path)
         assert self._model_base is not None
         model = LoRAModelRaw.from_checkpoint(
             file_path=model_path,
             dtype=self._torch_dtype,
-            base_model=self._model_base,
         )
         return model
 
     # override
-    def _get_model_path(self, config: AnyModelConfig) -> Path:
-        # cheating a little - we remember this variable for using in the subsequent call to _load_model()
-        self._model_base = config.base
+    def _get_model_path(
+        self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+    ) -> Tuple[Path, AnyModelConfig, Optional[SubModelType]]:
+        self._model_base = (
+            config.base
+        )  # cheating a little - we remember this variable for using in the subsequent call to _load_model()
 
         model_base_path = self._app_config.models_path
         model_path = model_base_path / config.path
@@ -69,4 +72,5 @@ class LoRALoader(ModelLoader):
                     model_path = path
                     break
 
-        return model_path.resolve()
+        result = model_path.resolve(), config, submodel_type
+        return result

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

@@ -7,9 +7,9 @@ from typing import Optional
 
 from invokeai.backend.model_manager import (
     AnyModel,
-    AnyModelConfig,
     BaseModelType,
     ModelFormat,
+    ModelRepoVariant,
     ModelType,
     SubModelType,
 )
@@ -25,19 +25,18 @@ class OnnyxDiffusersModel(GenericDiffusersLoader):
 
     def _load_model(
         self,
-        config: AnyModelConfig,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
         if not submodel_type is not None:
             raise Exception("A submodel type must be provided when loading onnx pipelines.")
-        model_path = Path(config.path)
         load_class = self.get_hf_load_class(model_path, submodel_type)
-        repo_variant = getattr(config, "repo_variant", None)
-        variant = repo_variant.value if repo_variant else None
+        variant = model_variant.value if model_variant else None
         model_path = model_path / submodel_type.value
         result: AnyModel = load_class.from_pretrained(
             model_path,
             torch_dtype=self._torch_dtype,
             variant=variant,
-        )
+        )  # type: ignore
         return result

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

@@ -9,16 +9,12 @@ from invokeai.backend.model_manager import (
     AnyModelConfig,
     BaseModelType,
     ModelFormat,
+    ModelRepoVariant,
     ModelType,
     SchedulerPredictionType,
     SubModelType,
 )
-from invokeai.backend.model_manager.config import (
-    CheckpointConfigBase,
-    DiffusersConfigBase,
-    MainCheckpointConfig,
-    ModelVariantType,
-)
+from invokeai.backend.model_manager.config import CheckpointConfigBase, MainCheckpointConfig, ModelVariantType
 from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_ckpt_to_diffusers
 
 from .. import ModelLoaderRegistry
@@ -45,15 +41,14 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
 
     def _load_model(
         self,
-        config: AnyModelConfig,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
         if not submodel_type is not None:
             raise Exception("A submodel type must be provided when loading main pipelines.")
-        model_path = Path(config.path)
         load_class = self.get_hf_load_class(model_path, submodel_type)
-        repo_variant = config.repo_variant if isinstance(config, DiffusersConfigBase) else None
-        variant = repo_variant.value if repo_variant else None
+        variant = model_variant.value if model_variant else None
         model_path = model_path / submodel_type.value
         try:
             result: AnyModel = load_class.from_pretrained(
@@ -83,7 +78,7 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
         else:
             return True
 
-    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel:
+    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Path) -> Path:
         assert isinstance(config, MainCheckpointConfig)
         base = config.base
 
@@ -99,11 +94,11 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
 
         self._logger.info(f"Converting {model_path} to diffusers format")
 
-        loaded_model = convert_ckpt_to_diffusers(
+        convert_ckpt_to_diffusers(
             model_path,
             output_path,
             model_type=self.model_base_to_model_type[base],
-            original_config_file=self._app_config.legacy_conf_path / config.config_path,
+            original_config_file=self._app_config.root_path / config.config_path,
             extract_ema=True,
             from_safetensors=model_path.suffix == ".safetensors",
             precision=self._torch_dtype,
@@ -113,4 +108,4 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
             load_safety_checker=False,
             num_in_channels=VARIANT_TO_IN_CHANNEL_MAP[config.variant],
         )
-        return loaded_model
+        return output_path

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

@@ -2,13 +2,14 @@
 """Class for TI model loading in InvokeAI."""
 
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Tuple
 
 from invokeai.backend.model_manager import (
     AnyModel,
     AnyModelConfig,
     BaseModelType,
     ModelFormat,
+    ModelRepoVariant,
     ModelType,
     SubModelType,
 )
@@ -26,19 +27,22 @@ class TextualInversionLoader(ModelLoader):
 
     def _load_model(
         self,
-        config: AnyModelConfig,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
         if submodel_type is not None:
             raise ValueError("There are no submodels in a TI model.")
         model = TextualInversionModelRaw.from_checkpoint(
-            file_path=config.path,
+            file_path=model_path,
             dtype=self._torch_dtype,
         )
         return model
 
     # override
-    def _get_model_path(self, config: AnyModelConfig) -> Path:
+    def _get_model_path(
+        self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+    ) -> Tuple[Path, AnyModelConfig, Optional[SubModelType]]:
         model_path = self._app_config.models_path / config.path
 
         if config.format == ModelFormat.EmbeddingFolder:
@@ -49,4 +53,4 @@ class TextualInversionLoader(ModelLoader):
         if not path.exists():
             raise OSError(f"The embedding file at {path} was not found")
 
-        return path
+        return path, config, submodel_type

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

@@ -2,7 +2,6 @@
 """Class for VAE model loading in InvokeAI."""
 
 from pathlib import Path
-from typing import Optional
 
 import torch
 from omegaconf import DictConfig, OmegaConf
@@ -14,7 +13,7 @@ from invokeai.backend.model_manager import (
     ModelFormat,
     ModelType,
 )
-from invokeai.backend.model_manager.config import AnyModel, CheckpointConfigBase
+from invokeai.backend.model_manager.config import CheckpointConfigBase
 from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_ldm_vae_to_diffusers
 
 from .. import ModelLoaderRegistry
@@ -39,13 +38,13 @@ class VAELoader(GenericDiffusersLoader):
         else:
             return True
 
-    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel:
+    def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Path) -> Path:
         # TODO(MM2): check whether sdxl VAE models convert.
         if config.base not in {BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2}:
             raise Exception(f"VAE conversion not supported for model type: {config.base}")
         else:
             assert isinstance(config, CheckpointConfigBase)
-            config_file = self._app_config.legacy_conf_path / config.config_path
+            config_file = self._app_config.root_path / config.config_path
 
         if model_path.suffix == ".safetensors":
             checkpoint = safetensors_load_file(model_path, device="cpu")
@@ -64,6 +63,6 @@ class VAELoader(GenericDiffusersLoader):
             vae_config=ckpt_config,
             image_size=512,
             precision=self._torch_dtype,
-            dump_path=output_path,
         )
-        return vae_model
+        vae_model.save_pretrained(output_path, safe_serialization=True)
+        return output_path

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 TorchDevice
+from invokeai.backend.util.devices import choose_torch_device, torch_dtype
 
 from . import (
     AnyModelConfig,
@@ -43,7 +43,6 @@ 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,
@@ -69,7 +68,7 @@ class ModelMerger(object):
             warnings.simplefilter("ignore")
             verbosity = dlogging.get_verbosity()
             dlogging.set_verbosity_error()
-            dtype = torch.float16 if variant == "fp16" else self._dtype
+            dtype = torch.float16 if variant == "fp16" else torch_dtype(choose_torch_device())
 
             # 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.
@@ -152,7 +151,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 self._dtype
+        dtype = torch.float16 if variant == "fp16" else torch_dtype(choose_torch_device())
         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

@@ -230,10 +230,9 @@ 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()):
@@ -324,7 +323,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, map_location="cpu")
+                model = torch.load(model_path)
                 assert isinstance(model, dict)
                 return model
             else:
@@ -528,25 +527,8 @@ class ControlNetCheckpointProbe(CheckpointProbeBase):
 
 
 class IPAdapterCheckpointProbe(CheckpointProbeBase):
-    """Class for probing IP Adapters"""
-
     def get_base_type(self) -> BaseModelType:
-        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")
+        raise NotImplementedError()
 
 
 class CLIPVisionCheckpointProbe(CheckpointProbeBase):
@@ -786,7 +768,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_patcher.py

@@ -17,7 +17,7 @@ from invokeai.backend.model_manager import AnyModel
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
 
-from .lora import LoRAModelRaw
+from .lora_model_raw import LoRAModelRaw
 from .textual_inversion import TextualInversionManager, TextualInversionModelRaw
 
 """

invokeai/backend/onnx/onnx_runtime.py

@@ -6,17 +6,16 @@ from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
 import onnx
+import torch
 from onnx import numpy_helper
 from onnxruntime import InferenceSession, SessionOptions, get_available_providers
 
-from ..raw_model import RawModel
-
 ONNX_WEIGHTS_NAME = "model.onnx"
 
 
 # NOTE FROM LS: This was copied from Stalker's original implementation.
 # I have not yet gone through and fixed all the type hints
-class IAIOnnxRuntimeModel(RawModel):
+class IAIOnnxRuntimeModel(torch.nn.Module):
     class _tensor_access:
         def __init__(self, model):  # type: ignore
             self.model = model

invokeai/backend/raw_model.py

@@ -1,15 +0,0 @@
-"""Base class for 'Raw' models.
-
-The RawModel class is the base class of LoRAModelRaw and TextualInversionModelRaw,
-and is used for type checking of calls to the model patcher. Its main purpose
-is to avoid a circular import issues when lora.py tries to import BaseModelType
-from invokeai.backend.model_manager.config, and the latter tries to import LoRAModelRaw
-from lora.py.
-
-The term 'raw' was introduced to describe a wrapper around a torch.nn.Module
-that adds additional methods and attributes.
-"""
-
-
-class RawModel:
-    """Base class for 'Raw' model wrappers."""

invokeai/backend/stable_diffusion/diffusers_pipeline.py

@@ -21,11 +21,12 @@ from pydantic import Field
 from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 
 from invokeai.app.services.config.config_default import get_config
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import IPAdapterData, TextConditioningData
+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.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 TorchDevice
+from invokeai.backend.util.devices import normalize_device
 
 
 @dataclass
@@ -148,6 +149,16 @@ 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."""
@@ -255,7 +266,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(TorchDevice.normalize(self.unet.device))
+            mem_free, _ = torch.cuda.mem_get_info(normalize_device(self.unet.device))
         else:
             raise ValueError(f"unrecognized device {self.unet.device}")
         # input tensor of [1, 4, h/8, w/8]
@@ -284,8 +295,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         self,
         latents: torch.Tensor,
         num_inference_steps: int,
-        scheduler_step_kwargs: dict[str, Any],
-        conditioning_data: TextConditioningData,
+        conditioning_data: ConditioningData,
         *,
         noise: Optional[torch.Tensor],
         timesteps: torch.Tensor,
@@ -298,7 +308,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         mask: Optional[torch.Tensor] = None,
         masked_latents: Optional[torch.Tensor] = None,
         gradient_mask: Optional[bool] = False,
-        seed: int,
+        seed: Optional[int] = None,
     ) -> torch.Tensor:
         if init_timestep.shape[0] == 0:
             return latents
@@ -316,6 +326,20 @@ 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!")
@@ -324,15 +348,6 @@ 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:
@@ -340,7 +355,6 @@ 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,
@@ -366,8 +380,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         self,
         latents: torch.Tensor,
         timesteps,
-        conditioning_data: TextConditioningData,
-        scheduler_step_kwargs: dict[str, Any],
+        conditioning_data: ConditioningData,
         *,
         additional_guidance: List[Callable] = None,
         control_data: List[ControlNetData] = None,
@@ -384,22 +397,22 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         if timesteps.shape[0] == 0:
             return latents
 
-        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
+        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,
             )
-            unet_attention_patcher = UNetAttentionPatcher(ip_adapters)
-            attn_ctx = unet_attention_patcher.apply_ip_adapter_attention(self.invokeai_diffuser.model)
+            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()
 
         with attn_ctx:
             if callback is not None:
@@ -422,11 +435,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)
@@ -450,14 +463,14 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         self,
         t: torch.Tensor,
         latents: torch.Tensor,
-        conditioning_data: TextConditioningData,
+        conditioning_data: ConditioningData,
         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]
@@ -472,6 +485,23 @@ 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
@@ -520,7 +550,6 @@ 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
@@ -540,7 +569,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
             )
 
         # compute the previous noisy sample x_t -> x_t-1
-        step_output = self.scheduler.step(noise_pred, timestep, latents, **scheduler_step_kwargs)
+        step_output = self.scheduler.step(noise_pred, timestep, latents, **conditioning_data.scheduler_args)
 
         # 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,17 +1,27 @@
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Union
+import dataclasses
+import inspect
+from dataclasses import dataclass, field
+from typing import Any, List, Optional, Union
 
 import torch
 
-from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
+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
 
 
 @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)
@@ -25,8 +35,6 @@ class ConditioningFieldData:
 
 @dataclass
 class SDXLConditioningInfo(BasicConditioningInfo):
-    """SDXL text conditioning information produced by Compel."""
-
     pooled_embeds: torch.Tensor
     add_time_ids: torch.Tensor
 
@@ -49,75 +57,37 @@ class IPAdapterConditioningInfo:
 
 
 @dataclass
-class IPAdapterData:
-    ip_adapter_model: IPAdapter
-    ip_adapter_conditioning: IPAdapterConditioningInfo
-    mask: torch.Tensor
-    target_blocks: List[str]
+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)
 
-    # 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
+    ip_adapter_conditioning: Optional[list[IPAdapterConditioningInfo]] = None
 
-    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
+    @property
+    def dtype(self):
+        return self.text_embeddings.dtype
 
-
-@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)
+    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)

invokeai/backend/stable_diffusion/diffusion/cross_attention_control.py

@@ -0,0 +1,218 @@
+# 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

@@ -1,214 +0,0 @@
-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

@@ -1,72 +0,0 @@
-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

@@ -1,105 +0,0 @@
-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,20 +1,26 @@
 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 (
-    IPAdapterData,
-    Range,
-    TextConditioningData,
-    TextConditioningRegions,
+    ConditioningData,
+    ExtraConditioningInfo,
+    SDXLConditioningInfo,
+)
+
+from .cross_attention_control import (
+    CrossAttentionType,
+    CrossAttnControlContext,
+    SwapCrossAttnContext,
+    setup_cross_attention_control_attention_processors,
 )
-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]
@@ -52,8 +58,31 @@ 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,
@@ -61,7 +90,7 @@ class InvokeAIDiffuserComponent:
         timestep: torch.Tensor,
         step_index: int,
         total_step_count: int,
-        conditioning_data: TextConditioningData,
+        conditioning_data,
     ):
         down_block_res_samples, mid_block_res_sample = None, None
 
@@ -94,28 +123,28 @@ class InvokeAIDiffuserComponent:
                 added_cond_kwargs = None
 
                 if cfg_injection:  # only applying ControlNet to conditional instead of in unconditioned
-                    if conditioning_data.is_sdxl():
+                    if type(conditioning_data.text_embeddings) is SDXLConditioningInfo:
                         added_cond_kwargs = {
-                            "text_embeds": conditioning_data.cond_text.pooled_embeds,
-                            "time_ids": conditioning_data.cond_text.add_time_ids,
+                            "text_embeds": conditioning_data.text_embeddings.pooled_embeds,
+                            "time_ids": conditioning_data.text_embeddings.add_time_ids,
                         }
-                    encoder_hidden_states = conditioning_data.cond_text.embeds
+                    encoder_hidden_states = conditioning_data.text_embeddings.embeds
                     encoder_attention_mask = None
                 else:
-                    if conditioning_data.is_sdxl():
+                    if type(conditioning_data.text_embeddings) is SDXLConditioningInfo:
                         added_cond_kwargs = {
                             "text_embeds": torch.cat(
                                 [
                                     # TODO: how to pad? just by zeros? or even truncate?
-                                    conditioning_data.uncond_text.pooled_embeds,
-                                    conditioning_data.cond_text.pooled_embeds,
+                                    conditioning_data.unconditioned_embeddings.pooled_embeds,
+                                    conditioning_data.text_embeddings.pooled_embeds,
                                 ],
                                 dim=0,
                             ),
                             "time_ids": torch.cat(
                                 [
-                                    conditioning_data.uncond_text.add_time_ids,
-                                    conditioning_data.cond_text.add_time_ids,
+                                    conditioning_data.unconditioned_embeddings.add_time_ids,
+                                    conditioning_data.text_embeddings.add_time_ids,
                                 ],
                                 dim=0,
                             ),
@@ -124,8 +153,8 @@ class InvokeAIDiffuserComponent:
                         encoder_hidden_states,
                         encoder_attention_mask,
                     ) = self._concat_conditionings_for_batch(
-                        conditioning_data.uncond_text.embeds,
-                        conditioning_data.cond_text.embeds,
+                        conditioning_data.unconditioned_embeddings.embeds,
+                        conditioning_data.text_embeddings.embeds,
                     )
                 if isinstance(control_datum.weight, list):
                     # if controlnet has multiple weights, use the weight for the current step
@@ -169,15 +198,24 @@ class InvokeAIDiffuserComponent:
         self,
         sample: torch.Tensor,
         timestep: torch.Tensor,
-        conditioning_data: TextConditioningData,
-        ip_adapter_data: Optional[list[IPAdapterData]],
+        conditioning_data: ConditioningData,
         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
     ):
-        if self.sequential_guidance:
+        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.
             (
                 unconditioned_next_x,
                 conditioned_next_x,
@@ -185,9 +223,7 @@ 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,
+                cross_attention_control_types_to_do=cross_attention_control_types_to_do,
                 down_block_additional_residuals=down_block_additional_residuals,
                 mid_block_additional_residual=mid_block_additional_residual,
                 down_intrablock_additional_residuals=down_intrablock_additional_residuals,
@@ -200,9 +236,6 @@ 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,
@@ -261,84 +294,53 @@ class InvokeAIDiffuserComponent:
 
     def _apply_standard_conditioning(
         self,
-        x: torch.Tensor,
-        sigma: torch.Tensor,
-        conditioning_data: TextConditioningData,
-        ip_adapter_data: Optional[list[IPAdapterData]],
-        step_index: int,
-        total_step_count: int,
+        x,
+        sigma,
+        conditioning_data: ConditioningData,
         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 = {}
-        if ip_adapter_data is not None:
-            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
+        cross_attention_kwargs = None
+        if conditioning_data.ip_adapter_conditioning is not None:
             # Note that we 'stack' to produce tensors of shape (batch_size, num_ip_images, seq_len, token_len).
-            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
+            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
+                ]
+            }
 
         added_cond_kwargs = None
-        if conditioning_data.is_sdxl():
+        if type(conditioning_data.text_embeddings) is SDXLConditioningInfo:
             added_cond_kwargs = {
                 "text_embeds": torch.cat(
                     [
                         # TODO: how to pad? just by zeros? or even truncate?
-                        conditioning_data.uncond_text.pooled_embeds,
-                        conditioning_data.cond_text.pooled_embeds,
+                        conditioning_data.unconditioned_embeddings.pooled_embeds,
+                        conditioning_data.text_embeddings.pooled_embeds,
                     ],
                     dim=0,
                 ),
                 "time_ids": torch.cat(
                     [
-                        conditioning_data.uncond_text.add_time_ids,
-                        conditioning_data.cond_text.add_time_ids,
+                        conditioning_data.unconditioned_embeddings.add_time_ids,
+                        conditioning_data.text_embeddings.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.uncond_text.embeds, conditioning_data.cond_text.embeds
+            conditioning_data.unconditioned_embeddings.embeds, conditioning_data.text_embeddings.embeds
         )
         both_results = self.model_forward_callback(
             x_twice,
@@ -358,10 +360,8 @@ class InvokeAIDiffuserComponent:
         self,
         x: torch.Tensor,
         sigma,
-        conditioning_data: TextConditioningData,
-        ip_adapter_data: Optional[list[IPAdapterData]],
-        step_index: int,
-        total_step_count: int,
+        conditioning_data: ConditioningData,
+        cross_attention_control_types_to_do: list[CrossAttentionType],
         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,48 +391,53 @@ 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 = {}
+        cross_attention_kwargs = None
 
         # Prepare IP-Adapter cross-attention kwargs for the unconditioned pass.
-        if ip_adapter_data is not None:
-            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
+        if conditioning_data.ip_adapter_conditioning is not None:
             # Note that we 'unsqueeze' to produce tensors of shape (batch_size=1, num_ip_images, seq_len, token_len).
-            image_prompt_embeds = [
-                torch.unsqueeze(ipa_conditioning.uncond_image_prompt_embeds, dim=0)
-                for ipa_conditioning in ip_adapter_conditioning
-            ]
+            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
+                ]
+            }
 
-            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 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}
 
         # Prepare SDXL conditioning kwargs for the unconditioned pass.
         added_cond_kwargs = None
-        if conditioning_data.is_sdxl():
+        is_sdxl = type(conditioning_data.text_embeddings) is SDXLConditioningInfo
+        if is_sdxl:
             added_cond_kwargs = {
-                "text_embeds": conditioning_data.uncond_text.pooled_embeds,
-                "time_ids": conditioning_data.uncond_text.add_time_ids,
+                "text_embeds": conditioning_data.unconditioned_embeddings.pooled_embeds,
+                "time_ids": conditioning_data.unconditioned_embeddings.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.uncond_text.embeds,
+            conditioning_data.unconditioned_embeddings.embeds,
             cross_attention_kwargs=cross_attention_kwargs,
             down_block_additional_residuals=uncond_down_block,
             mid_block_additional_residual=uncond_mid_block,
@@ -444,43 +449,36 @@ class InvokeAIDiffuserComponent:
         # Conditioned pass
         ###################
 
-        cross_attention_kwargs = {}
+        cross_attention_kwargs = None
 
-        if ip_adapter_data is not None:
-            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
+        # Prepare IP-Adapter cross-attention kwargs for the conditioned pass.
+        if conditioning_data.ip_adapter_conditioning is not None:
             # Note that we 'unsqueeze' to produce tensors of shape (batch_size=1, num_ip_images, seq_len, token_len).
-            image_prompt_embeds = [
-                torch.unsqueeze(ipa_conditioning.cond_image_prompt_embeds, dim=0)
-                for ipa_conditioning in ip_adapter_conditioning
-            ]
+            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
+                ]
+            }
 
-            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 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}
 
         # Prepare SDXL conditioning kwargs for the conditioned pass.
         added_cond_kwargs = None
-        if conditioning_data.is_sdxl():
+        if is_sdxl:
             added_cond_kwargs = {
-                "text_embeds": conditioning_data.cond_text.pooled_embeds,
-                "time_ids": conditioning_data.cond_text.add_time_ids,
+                "text_embeds": conditioning_data.text_embeddings.pooled_embeds,
+                "time_ids": conditioning_data.text_embeddings.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.cond_text.embeds,
+            conditioning_data.text_embeddings.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

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