Enable LoRAPatcher.apply_smart_lora_patches(...) throughout the stack.

.github/workflows/typegen-checks.yml

@@ -1,85 +0,0 @@
-# Runs typegen schema quality checks.
-# Frontend types should match the server.
-#
-# Checks for changes to files before running the checks.
-# If always_run is true, always runs the checks.
-
-name: 'typegen checks'
-
-on:
-  push:
-    branches:
-      - 'main'
-  pull_request:
-    types:
-      - 'ready_for_review'
-      - 'opened'
-      - 'synchronize'
-  merge_group:
-  workflow_dispatch:
-    inputs:
-      always_run:
-        description: 'Always run the checks'
-        required: true
-        type: boolean
-        default: true
-  workflow_call:
-    inputs:
-      always_run:
-        description: 'Always run the checks'
-        required: true
-        type: boolean
-        default: true
-
-jobs:
-  typegen-checks:
-    runs-on: ubuntu-22.04
-    timeout-minutes: 15 # expected run time: <5 min
-    steps:
-      - name: checkout
-        uses: actions/checkout@v4
-
-      - name: check for changed files
-        if: ${{ inputs.always_run != true }}
-        id: changed-files
-        uses: tj-actions/changed-files@v42
-        with:
-          files_yaml: |
-            src:
-              - 'pyproject.toml'
-              - 'invokeai/**'
-
-      - name: setup python
-        if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-          cache: pip
-          cache-dependency-path: pyproject.toml
-
-      - name: install python dependencies
-        if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        run: pip3 install --use-pep517 --editable="."
-
-      - name: install frontend dependencies
-        if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        uses: ./.github/actions/install-frontend-deps
-
-      - name: copy schema
-        if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        run: cp invokeai/frontend/web/src/services/api/schema.ts invokeai/frontend/web/src/services/api/schema_orig.ts
-        shell: bash
-
-      - name: generate schema
-        if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        run: make frontend-typegen
-        shell: bash
-
-      - name: compare files
-        if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        run: |
-          if ! diff invokeai/frontend/web/src/services/api/schema.ts invokeai/frontend/web/src/services/api/schema_orig.ts; then
-            echo "Files are different!";
-            exit 1;
-          fi
-        shell: bash

README.md

@@ -30,12 +30,51 @@ Invoke is available in two editions:
 |----------------------------------------------------------------------------------------------------------------------------|
 |  [Installation and Updates][installation docs] - [Documentation and Tutorials][docs home] - [Bug Reports][github issues] - [Contributing][contributing docs]  | 
 
-# Installation
+</div>
 
-To get started with Invoke, [Download the Installer](https://www.invoke.com/downloads).
+## Quick Start
 
-For detailed step by step instructions, or for instructions on manual/docker installations, visit our documentation on [Installation and Updates][installation docs]
+1. Download and unzip the installer from the bottom of the [latest release][latest release link].
+2. Run the installer script.
 
+   - **Windows**: Double-click on the `install.bat` script.
+   - **macOS**: Open a Terminal window, drag the file `install.sh` from Finder into the Terminal, and press enter.
+   - **Linux**: Run `install.sh`.
+
+3. When prompted, enter a location for the install and select your GPU type.
+4. Once the install finishes, find the directory you selected during install. The default location is `C:\Users\Username\invokeai` for Windows or `~/invokeai` for Linux/macOS.
+5. Run the launcher script (`invoke.bat` for Windows, `invoke.sh` for macOS and Linux) the same way you ran the installer script in step 2.
+6. Select option 1 to start the application. Once it starts up, open your browser and go to <http://localhost:9090>.
+7. Open the model manager tab to install a starter model and then you'll be ready to generate.
+
+More detail, including hardware requirements and manual install instructions, are available in the [installation documentation][installation docs].
+
+## Docker Container
+
+We publish official container images in Github Container Registry: https://github.com/invoke-ai/InvokeAI/pkgs/container/invokeai. Both CUDA and ROCm images are available. Check the above link for relevant tags.
+
+> [!IMPORTANT]
+> Ensure that Docker is set up to use the GPU. Refer to [NVIDIA][nvidia docker docs] or [AMD][amd docker docs] documentation.
+
+### Generate!
+
+Run the container, modifying the command as necessary:
+
+```bash
+docker run --runtime=nvidia --gpus=all --publish 9090:9090 ghcr.io/invoke-ai/invokeai
+```
+
+Then open `http://localhost:9090` and install some models using the Model Manager tab to begin generating.
+
+For ROCm, add `--device /dev/kfd --device /dev/dri` to the `docker run` command.
+
+### Persist your data
+
+You will likely want to persist your workspace outside of the container. Use the `--volume /home/myuser/invokeai:/invokeai` flag to mount some local directory (using its **absolute** path) to the `/invokeai` path inside the container. Your generated images and models will reside there. You can use this directory with other InvokeAI installations, or switch between runtime directories as needed.
+
+### DIY
+
+Build your own image and customize the environment to match your needs using our `docker-compose` stack. See [README.md](./docker/README.md) in the [docker](./docker) directory.
 
 ## Troubleshooting, FAQ and Support
 

docs/configuration.md

@@ -114,10 +114,6 @@ remote_api_tokens:
 
 The provided token will be added as a `Bearer` token to the network requests to download the model files. As far as we know, this works for all model marketplaces that require authorization.
 
-!!! tip "HuggingFace Models"
-
-    If you get an error when installing a HF model using a URL instead of repo id, you may need to [set up a HF API token](https://huggingface.co/settings/tokens) and add an entry for it under `remote_api_tokens`. Use `huggingface.co` for `url_regex`.
-
 #### Model Hashing
 
 Models are hashed during installation, providing a stable identifier for models across all platforms. Hashing is a one-time operation.

docs/contributing/MODEL_MANAGER.md

@@ -1364,6 +1364,7 @@ the in-memory loaded model:
 |----------------|-----------------|------------------|
 | `config`       | AnyModelConfig         | A copy of the model's configuration record for retrieving base type, etc. |
 | `model`        | AnyModel               | The instantiated model (details below) |
+| `locker`       | ModelLockerBase        | A context manager that mediates the movement of the model into VRAM |
 
 ### get_model_by_key(key, [submodel]) -> LoadedModel
 

docs/contributing/dev-environment.md

@@ -1,10 +1,12 @@
 # Dev Environment
 
-To make changes to Invoke's backend, frontend or documentation, you'll need to set up a dev environment.
+To make changes to Invoke's backend, frontend, or documentation, you'll need to set up a dev environment.
 
-If you only want to make changes to the docs site, you can skip the frontend dev environment setup as described in the below guide.
+If you just want to use Invoke, you should use the [installer][installer link].
 
-If you just want to use Invoke, you should use the [launcher][launcher link].
+!!! info "Why do I need the frontend toolchain?"
+
+    The repo doesn't contain a build of the frontend. You'll be responsible for rebuilding it every time you pull in new changes, or run it in dev mode (which incurs a substantial performance penalty).
 
 !!! warning
 
@@ -15,66 +17,84 @@ If you just want to use Invoke, you should use the [launcher][launcher link].
 ## Setup
 
 1. Run through the [requirements][requirements link].
-
 2. [Fork and clone][forking link] the [InvokeAI repo][repo link].
-
 3. Create an directory for user data (images, models, db, etc). This is typically at `~/invokeai`, but if you already have a non-dev install, you may want to create a separate directory for the dev install.
-
-4. Follow the [manual install][manual install link] guide, with some modifications to the install command:
-
-      - Use `.` instead of `invokeai` to install from the current directory.
-
-      - Add `-e` after the `install` operation to make this an [editable install][editable install link]. That means your changes to the python code will be reflected when you restart the Invoke server.
-
-      - When installing the `invokeai` package, add the `dev`, `test` and `docs` package options to the package specifier. You may or may not need the `xformers` option - follow the manual install guide to figure that out. So, your package specifier will be either `".[dev,test,docs]"` or `".[dev,test,docs,xformers]"`. Note the quotes!
-
-     With the modifications made, the install command should look something like this:
+4. Create a python virtual environment inside the directory you just created:
 
       ```sh
-      uv pip install -e ".[dev,test,docs,xformers]" --python 3.11 --python-preference only-managed --index=https://download.pytorch.org/whl/cu124 --reinstall
+      python3 -m venv .venv --prompt InvokeAI-Dev
       ```
 
-5. At this point, you should have Invoke installed, a venv set up and activated, and the server running. But you will see a warning in the terminal that no UI was found. If you go to the URL for the server, you won't get a UI.
-
-      This is because the UI build is not distributed with the source code. You need to build it manually. End the running server instance.
-
-      If you only want to edit the docs, you can stop here and skip to the **Documentation** section below.
-
-6. Install the frontend dev toolchain:
-
-      - [`nodejs`](https://nodejs.org/) (v20+)
-
-      - [`pnpm`](https://pnpm.io/8.x/installation) (must be v8 - not v9!)
-
-7. Do a production build of the frontend:
+5. Activate the venv (you'll need to do this every time you want to run the app):
 
       ```sh
-      cd <PATH_TO_INVOKEAI_REPO>/invokeai/frontend/web
+      source .venv/bin/activate
+      ```
+
+6. Install the repo as an [editable install][editable install link]:
+
+      ```sh
+      pip install -e ".[dev,test,xformers]" --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu121
+      ```
+
+      Refer to the [manual installation][manual install link] instructions for more determining the correct install options. `xformers` is optional, but `dev` and `test` are not.
+
+7. Install the frontend dev toolchain:
+
+   - [`nodejs`](https://nodejs.org/) (recommend v20 LTS)
+   - [`pnpm`](https://pnpm.io/8.x/installation) (must be v8 - not v9!)
+
+8. Do a production build of the frontend:
+
+      ```sh
+      cd PATH_TO_INVOKEAI_REPO/invokeai/frontend/web
       pnpm i
       pnpm build
       ```
 
-8. Restart the server and navigate to the URL. You should get a UI. After making changes to the python code, restart the server to see those changes.
+9. Start the application:
+
+      ```sh
+      cd PATH_TO_INVOKEAI_REPO
+      python scripts/invokeai-web.py
+      ```
+
+10. Access the UI at `localhost:9090`.
 
 ## Updating the UI
 
-You'll need to run `pnpm build` every time you pull in new changes.
-
-Another option is to skip the build and instead run the UI in dev mode:
+You'll need to run `pnpm build` every time you pull in new changes. Another option is to skip the build and instead run the app in dev mode:
 
 ```sh
 pnpm dev
 ```
 
-This starts a vite dev server for the UI at `127.0.0.1:5173`, which you will use instead of `127.0.0.1:9090`.
+This starts a dev server at `localhost:5173`, which you will use instead of `localhost:9090`.
 
-The dev mode is substantially slower than the production build but may be more convenient if you just need to test things out. It will hot-reload the UI as you make changes to the frontend code. Sometimes the hot-reload doesn't work, and you need to manually refresh the browser tab.
+The dev mode is substantially slower than the production build but may be more convenient if you just need to test things out.
 
 ## Documentation
 
-The documentation is built with `mkdocs`. It provides a hot-reload dev server for the docs. Start it with `mkdocs serve`.
+The documentation is built with `mkdocs`. To preview it locally, you need a additional set of packages installed.
 
-[launcher link]: ../installation/quick_start.md
+```sh
+# after activating the venv
+pip install -e ".[docs]"
+```
+
+Then, you can start a live docs dev server, which will auto-refresh when you edit the docs:
+
+```sh
+mkdocs serve
+```
+
+On macOS and Linux, there is a `make` target for this:
+
+```sh
+make docs
+```
+
+[installer link]: ../installation/installer.md
 [forking link]: https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/fork-a-repo
 [requirements link]: ../installation/requirements.md
 [repo link]: https://github.com/invoke-ai/InvokeAI

docs/index.md

@@ -50,9 +50,11 @@ title: Invoke
 
 ## Installation
 
-The [Invoke Launcher](installation/quick_start.md) is the easiest way to install, update and run Invoke on Windows, macOS and Linux.
+The [installer script](installation/installer.md) is the easiest way to install and update the application.
 
-You can also install Invoke as [python package](installation/manual.md) or with [docker](installation/docker.md).
+You can also install Invoke as python package [via PyPI](installation/manual.md) or [docker](installation/docker.md).
+
+See the [installation section](./installation/index.md) for more information.
 
 ## Help
 

docs/installation/docker.md

@@ -4,7 +4,7 @@ title: Docker
 
 !!! warning "macOS users"
 
-    Docker can not access the GPU on macOS, so your generation speeds will be slow. Use the [launcher](./quick_start.md) instead.
+    Docker can not access the GPU on macOS, so your generation speeds will be slow. Use the [installer](./installer.md) instead.
 
 !!! tip "Linux and Windows Users"
 

docs/installation/index.md

@@ -0,0 +1,36 @@
+# Installation and Updating Overview
+
+Before installing, review the [installation requirements](./requirements.md) to ensure your system is set up properly.
+
+See the [FAQ](../faq.md) for frequently-encountered installation issues.
+
+If you need more help, join our [discord](https://discord.gg/ZmtBAhwWhy) or [create a GitHub issue](https://github.com/invoke-ai/InvokeAI/issues).
+
+## Automated Installer & Updates
+
+✅ The automated [installer](./installer.md) is the best way to install Invoke.
+
+⬆️ The same installer is also the best way to update Invoke - simply rerun it for the same folder you installed to.
+
+The installation process simply manages installation for the core libraries & application dependencies that run Invoke.
+
+Models, images, or other assets in the Invoke root folder won't be affected by the installation process.
+
+## Manual Install
+
+If you are familiar with python and want more control over the packages that are installed, you can [install Invoke manually via PyPI](./manual.md).
+
+Updates are managed by reinstalling the latest version through PyPi.
+
+## Developer Install
+
+If you want to contribute to InvokeAI, you'll need to set up a [dev environment](../contributing/dev-environment.md).
+
+## Docker
+
+Invoke publishes docker images. See the [docker installation guide](./docker.md) for details.
+
+## Other Installation Guides
+
+- [PyPatchMatch](./patchmatch.md)
+- [Installing Models](./models.md)

docs/installation/installer.md

@@ -1,10 +1,4 @@
-# Legacy Scripts
-
-!!! warning "Legacy Scripts"
-
-    We recommend using the Invoke Launcher to install and update Invoke. It's a desktop application for Windows, macOS and Linux. It takes care of a lot of nitty gritty details for you.
-
-    Follow the [quick start guide](./quick_start.md) to get started.
+# Automatic Install & Updates
 
 !!! tip "Use the installer to update"
 

docs/installation/manual.md

@@ -4,11 +4,11 @@
 
     **Python experience is mandatory.**
 
-    If you want to use Invoke locally, you should probably use the [launcher](./quick_start.md).
+    If you want to use Invoke locally, you should probably use the [installer](./installer.md).
 
-    If you want to contribute to Invoke or run the app on the latest dev branch, instead follow the [dev environment](../contributing/dev-environment.md) guide.
+    If you want to contribute to Invoke, instead follow the [dev environment](../contributing/dev-environment.md) guide.
 
-InvokeAI is distributed as a python package on PyPI, installable with `pip`. There are a few things that are handled by the launcher that you'll need to manage manually, described in this guide.
+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.
 
 ## Requirements
 
@@ -16,39 +16,43 @@ Before you start, go through the [installation requirements](./requirements.md).
 
 ## Walkthrough
 
-We'll use [`uv`](https://github.com/astral-sh/uv) to install python and create a virtual environment, then install the `invokeai` package. `uv` is a modern, very fast alternative to `pip`.
-
-The following commands vary depending on the version of Invoke being installed and the system onto which it is being installed.
-
-1. Install `uv` as described in its [docs](https://docs.astral.sh/uv/getting-started/installation/#standalone-installer). We suggest using the standalone installer method.
-
-    Run `uv --version` to confirm that `uv` is installed and working. After installation, you may need to restart your terminal to get access to `uv`.
-
-2. Create a directory for your installation, typically in your home directory (e.g. `~/invokeai` or `$Home/invokeai`):
+1. Create a directory to contain your InvokeAI library, configuration files, and models. This is known as the "runtime" or "root" directory, and typically lives in your home directory under the name `invokeai`.
 
     === "Linux/macOS"
 
         ```bash
         mkdir ~/invokeai
-        cd ~/invokeai
         ```
 
     === "Windows (PowerShell)"
 
         ```bash
         mkdir $Home/invokeai
-        cd $Home/invokeai
         ```
 
-3. Create a virtual environment in that directory:
+1. Enter the root directory and create a virtual Python environment within it named `.venv`.
 
-    ```sh
-    uv venv --relocatable --prompt invoke --python 3.11 --python-preference only-managed .venv
-    ```
+    !!! warning "Virtual Environment Location"
 
-    This command creates a portable virtual environment at `.venv` complete with a portable python 3.11. It doesn't matter if your system has no python installed, or has a different version - `uv` will handle everything.
+        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.
 
-4. Activate the virtual environment:
+        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.
+
+    === "Linux/macOS"
+
+        ```bash
+        cd ~/invokeai
+        python3 -m venv .venv --prompt InvokeAI
+        ```
+
+    === "Windows (PowerShell)"
+
+        ```bash
+        cd $Home/invokeai
+        python3 -m venv .venv --prompt InvokeAI
+        ```
+
+1. Activate the new environment:
 
     === "Linux/macOS"
 
@@ -56,48 +60,41 @@ The following commands vary depending on the version of Invoke being installed a
         source .venv/bin/activate
         ```
 
-    === "Windows (PowerShell)"
+    === "Windows"
 
         ```ps
         .venv\Scripts\activate
         ```
 
-5. Choose a version to install. Review the [GitHub releases page](https://github.com/invoke-ai/InvokeAI/releases).
+    !!! info "Permissions Error (Windows)"
 
-6. Determine the package package specifier to use when installing. This is a performance optimization.
+        If you get a permissions error at this point, run this command and try again.
 
-    - If you have an Nvidia 20xx series GPU or older, use `invokeai[xformers]`.
-    - If you have an Nvidia 30xx series GPU or newer, or do not have an Nvidia GPU, use `invokeai`.
+        `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser`
 
-7. Determine the `PyPI` index URL to use for installation, if any. This is necessary to get the right version of torch installed.
+    The command-line prompt should change to to show `(InvokeAI)`, indicating the venv is active.
 
-    === "Invoke v5 or later"
+1. Make sure that pip is installed in your virtual environment and up to date:
 
-        - If you are on Windows with an Nvidia GPU, use `https://download.pytorch.org/whl/cu124`.
-        - If you are on Linux with no GPU, use `https://download.pytorch.org/whl/cpu`.
-        - If you are on Linux with an AMD GPU, use `https://download.pytorch.org/whl/rocm62`.
-        - **In all other cases, do not use an index.**
-
-    === "Invoke v4"
-
-        - If you are on Windows with an Nvidia GPU, use `https://download.pytorch.org/whl/cu124`.
-        - If you are on Linux with no GPU, use `https://download.pytorch.org/whl/cpu`.
-        - If you are on Linux with an AMD GPU, use `https://download.pytorch.org/whl/rocm52`.
-        - **In all other cases, do not use an index.**
-
-8. Install the `invokeai` package. Substitute the package specifier and version.
-
-    ```sh
-    uv pip install <PACKAGE_SPECIFIER>=<VERSION> --python 3.11 --python-preference only-managed --force-reinstall
+    ```bash
+    python3 -m pip install --upgrade pip
     ```
 
-    If you determined you needed to use a `PyPI` index URL in the previous step, you'll need to add `--index=<INDEX_URL>` like this:
+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.
 
-    ```sh
-    uv pip install <PACKAGE_SPECIFIER>=<VERSION> --python 3.11 --python-preference only-managed --index=<INDEX_URL> --force-reinstall
-    ```
+    - You may need to provide an [extra index URL](https://pip.pypa.io/en/stable/cli/pip_install/#cmdoption-extra-index-url). Select your platform configuration using [this tool on the PyTorch website](https://pytorch.org/get-started/locally/). Copy the `--extra-index-url` string from this and append it to your install command.
 
-9. Deactivate and reactivate your venv so that the invokeai-specific commands become available in the environment:
+        ```bash
+        pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu121
+        ```
+
+    - 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 strictly necessary. PyTorch includes an implementation of the SDP attention algorithm with similar performance for most GPUs.
+
+        ```bash
+        pip install "InvokeAI[xformers]" --use-pep517
+        ```
+
+1. Deactivate and reactivate your venv so that the invokeai-specific commands become available in the environment:
 
     === "Linux/macOS"
 
@@ -105,31 +102,17 @@ The following commands vary depending on the version of Invoke being installed a
         deactivate && source .venv/bin/activate
         ```
 
-    === "Windows (PowerShell)"
+    === "Windows"
 
         ```ps
         deactivate
         .venv\Scripts\activate
         ```
 
-10. Run the application, specifying the directory you created earlier as the root directory:
+1. Run the application:
 
-    === "Linux/macOS"
+    Run `invokeai-web` to start the UI. You must activate the virtual environment before running the app.
 
-        ```bash
-        invokeai-web --root ~/invokeai
-        ```
+    !!! warning
 
-    === "Windows (PowerShell)"
-
-        ```bash
-        invokeai-web --root $Home/invokeai
-        ```
-
-## Headless Install and Launch Scripts
-
-If you run Invoke on a headless server, you might want to install and run Invoke on the command line.
-
-We do not plan to maintain scripts to do this moving forward, instead focusing our dev resources on the GUI [launcher](../installation/quick_start.md).
-
-You can create your own scripts for this by copying the handful of commands in this guide. `uv`'s [`pip` interface docs](https://docs.astral.sh/uv/reference/cli/#uv-pip-install) may be useful.
+        If the virtual environment is _not_ inside the root directory, then you _must_ specify the path to the root directory with `--root \path\to\invokeai` or the `INVOKEAI_ROOT` environment variable.

docs/installation/quick_start.md

@@ -1,114 +0,0 @@
-# Invoke Community Edition Quick Start
-
-Welcome to Invoke! Follow these steps to install, update, and get started creating.
-
-## Step 1: System Requirements
-
-Invoke runs on Windows 10+, macOS 14+ and Linux (Ubuntu 20.04+ is well-tested).
-
-Hardware requirements vary significantly depending on model and image output size. The requirements below are rough guidelines.
-
-- All Apple Silicon (M1, M2, etc) Macs work, but 16GB+ memory is recommended.
-- AMD GPUs are supported on Linux only. The VRAM requirements are the same as Nvidia GPUs.
-
-!!! info "Hardware Requirements (Windows/Linux)"
-
-    === "SD1.5 - 512×512"
-
-        - GPU: Nvidia 10xx series or later, 4GB+ VRAM.
-        - Memory: At least 8GB RAM.
-        - Disk: 10GB for base installation plus 30GB for models.
-
-    === "SDXL - 1024×1024"
-
-        - GPU: Nvidia 20xx series or later, 8GB+ VRAM.
-        - Memory: At least 16GB RAM.
-        - Disk: 10GB for base installation plus 100GB for models.
-
-    === "FLUX - 1024×1024"
-
-        - GPU: Nvidia 20xx series or later, 10GB+ VRAM.
-        - Memory: At least 32GB RAM.
-        - Disk: 10GB for base installation plus 200GB for models.
-
-More detail on system requirements can be found [here](./requirements.md).
-
-## Step 2: Download
-
-Download the most launcher for your operating system:
-
-- [Download for Windows](https://download.invoke.ai/Invoke%20Community%20Edition.exe)
-- [Download for macOS](https://download.invoke.ai/Invoke%20Community%20Edition.dmg)
-- [Download for Linux](https://download.invoke.ai/Invoke%20Community%20Edition.AppImage)
-
-## Step 3: Install or Update
-
-Run the launcher you just downloaded, click **Install** and follow the instructions to get set up.
-
-If you have an existing Invoke installation, you can select it and let the launcher manage the install. You'll be able to update or launch the installation.
-
-!!! warning "Problem running the launcher on macOS"
-
-    macOS may not allow you to run the launcher. We are working to resolve this by signing the launcher executable. Until that is done, you can either use the [legacy scripts](./legacy_scripts.md) to install, or manually flag the launcher as safe:
-
-    - Open the **Invoke-Installer-mac-arm64.dmg** file.
-    - Drag the launcher to **Applications**.
-    - Open a terminal.
-    - Run `xattr -cr /Applications/Invoke-Installer.app`.
-
-    You should now be able to run the launcher.
-
-## Step 4: Launch
-
-Once installed, click **Finish**, then **Launch** to start Invoke.
-
-The very first run after an installation or update will take a few extra moments to get ready.
-
-!!! tip "Server Mode"
-
-    The launcher runs Invoke as a desktop application. You can enable **Server Mode** in the launcher's settings to disable this and instead access the UI through your web browser.
-
-## Step 5: Install Models
-
-With Invoke started up, you'll need to install some models.
-
-The quickest way to get started is to install a **Starter Model** bundle. If you already have a model collection, Invoke can use it.
-
-!!! info "Install Models"
-
-    === "Install a Starter Model bundle"
-
-        1. Go to the **Models** tab.
-        2. Click **Starter Models** on the right.
-        3. Click one of the bundles to install its models. Refer to the [system requirements](#step-1-confirm-system-requirements) if you're unsure which model architecture will work for your system.
-
-    === "Use my model collection"
-
-        4. Go to the **Models** tab.
-        5. Click **Scan Folder** on the right.
-        6. Paste the path to your models collection and click **Scan Folder**.
-        7. With **In-place install** enabled, Invoke will leave the model files where they are. If you disable this, **Invoke will move the models into its own folders**.
-
-You’re now ready to start creating!
-
-## Step 6: Learn the Basics
-
-We recommend watching our [Getting Started Playlist](https://www.youtube.com/playlist?list=PLvWK1Kc8iXGrQy8r9TYg6QdUuJ5MMx-ZO). It covers essential features and workflows, including:
-
-- Generating your first image.
-- Using control layers and reference guides.
-- Refining images with advanced workflows.
-
-## Other Installation Methods
-
-- You can install the Invoke application as a python package. See our [manual install](./manual.md) docs.
-- You can run Invoke with docker. See our [docker install](./docker.md) docs.
-- You can still use our legacy scripts to install and run Invoke. See the [legacy scripts](./legacy_scripts.md) docs.
-
-## Need Help?
-
-- Visit our [Support Portal](https://support.invoke.ai).
-- Watch the [Getting Started Playlist](https://www.youtube.com/playlist?list=PLvWK1Kc8iXGrQy8r9TYg6QdUuJ5MMx-ZO).
-- Join the conversation on [Discord][discord link].
-
-[discord link]: https://discord.gg/ZmtBAhwWhy

docs/installation/requirements.md

@@ -1,33 +1,90 @@
 # Requirements
 
-Invoke runs on Windows 10+, macOS 14+ and Linux (Ubuntu 20.04+ is well-tested).
+## GPU
 
-## Hardware
+!!! warning "Problematic Nvidia GPUs"
 
-Hardware requirements vary significantly depending on model and image output size. The requirements below are rough guidelines.
+    We do not recommend these GPUs. They cannot operate with half precision, but have insufficient VRAM to generate 512x512 images at full precision.
 
-- All Apple Silicon (M1, M2, etc) Macs work, but 16GB+ memory is recommended.
-- AMD GPUs are supported on Linux only. The VRAM requirements are the same as Nvidia GPUs.
+    - NVIDIA 10xx series cards such as the 1080 TI
+    - GTX 1650 series cards
+    - GTX 1660 series cards
 
-!!! info "Hardware Requirements (Windows/Linux)"
+Invoke runs best with a dedicated GPU, but will fall back to running on CPU, albeit much slower. You'll need a beefier GPU for SDXL.
 
-    === "SD1.5 - 512×512"
+!!! example "Stable Diffusion 1.5"
 
-        - GPU: Nvidia 10xx series or later, 4GB+ VRAM.
-        - Memory: At least 8GB RAM.
-        - Disk: 10GB for base installation plus 30GB for models.
+    === "Nvidia"
 
-    === "SDXL - 1024×1024"
+        ```
+        Any GPU with at least 4GB VRAM.
+        ```
 
-        - GPU: Nvidia 20xx series or later, 8GB+ VRAM.
-        - Memory: At least 16GB RAM.
-        - Disk: 10GB for base installation plus 100GB for models.
+    === "AMD"
 
-    === "FLUX - 1024×1024"
+        ```
+        Any GPU with at least 4GB VRAM. Linux only.
+        ```
 
-        - GPU: Nvidia 20xx series or later, 10GB+ VRAM.
-        - Memory: At least 32GB RAM.
-        - Disk: 10GB for base installation plus 200GB for models.
+    === "Mac"
+
+        ```
+        Any Apple Silicon Mac with at least 8GB memory.
+        ```
+
+!!! example "Stable Diffusion XL"
+
+    === "Nvidia"
+
+        ```
+        Any GPU with at least 8GB VRAM.
+        ```
+
+    === "AMD"
+
+        ```
+        Any GPU with at least 16GB VRAM. Linux only.
+        ```
+
+    === "Mac"
+
+        ```
+        Any Apple Silicon Mac with at least 16GB memory.
+        ```
+
+## RAM
+
+At least 12GB of RAM.
+
+## Disk
+
+SSDs will, of course, offer the best performance.
+
+The base application disk usage depends on the torch backend.
+
+!!! example "Disk"
+
+    === "Nvidia (CUDA)"
+
+        ```
+        ~6.5GB
+        ```
+
+    === "AMD (ROCm)"
+
+        ```
+        ~12GB
+        ```
+
+    === "Mac (MPS)"
+
+        ```
+        ~3.5GB
+        ```
+
+You'll need to set aside some space for images, depending on how much you generate. A couple GB is enough to get started.
+
+You'll need a good chunk of space for models. Even if you only install the most popular models and the usual support models (ControlNet, IP Adapter ,etc), you will quickly hit 50GB of models.
 
 !!! info "`tmpfs` on Linux"
 
@@ -35,32 +92,26 @@ Hardware requirements vary significantly depending on model and image output siz
 
 ## Python
 
-!!! tip "The launcher installs python for you"
-
-    You don't need to do this if you are installing with the [Invoke Launcher](./quick_start.md).
-
 Invoke requires python 3.10 or 3.11. If you don't already have one of these versions installed, we suggest installing 3.11, as it will be supported for longer.
 
-Check that your system has an up-to-date Python installed by running `python3 --version` in the terminal (Linux, macOS) or cmd/powershell (Windows).
+Check that your system has an up-to-date Python installed by running `python --version` in the terminal (Linux, macOS) or cmd/powershell (Windows).
 
-!!! info "Installing Python"
+<h3>Installing Python (Windows)</h3>
 
-    === "Windows"
+- Install python 3.11 with [an official installer].
+- The installer includes an option to add python to your PATH. Be sure to enable this. If you missed it, re-run the installer, choose to modify an existing installation, and tick that checkbox.
+- You may need to install [Microsoft Visual C++ Redistributable].
 
-        - Install python 3.11 with [an official installer].
-        - The installer includes an option to add python to your PATH. Be sure to enable this. If you missed it, re-run the installer, choose to modify an existing installation, and tick that checkbox.
-        - You may need to install [Microsoft Visual C++ Redistributable].
+<h3>Installing Python (macOS)</h3>
 
-    === "macOS"
+- Install python 3.11 with [an official installer].
+- If model installs fail with a certificate error, you may need to run this command (changing the python version to match what you have installed): `/Applications/Python\ 3.10/Install\ Certificates.command`
+- If you haven't already, you will need to install the XCode CLI Tools by running `xcode-select --install` in a terminal.
 
-        - Install python 3.11 with [an official installer].
-        - If model installs fail with a certificate error, you may need to run this command (changing the python version to match what you have installed): `/Applications/Python\ 3.10/Install\ Certificates.command`
-        - If you haven't already, you will need to install the XCode CLI Tools by running `xcode-select --install` in a terminal.
+<h3>Installing Python (Linux)</h3>
 
-    === "Linux"
-
-        - Installing python varies depending on your system. On Ubuntu, you can use the [deadsnakes PPA](https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa).
-        - You'll need to install `libglib2.0-0` and `libgl1-mesa-glx` for OpenCV to work. For example, on a Debian system: `sudo apt update && sudo apt install -y libglib2.0-0 libgl1-mesa-glx`
+- Follow the [linux install instructions], being sure to install python 3.11.
+- You'll need to install `libglib2.0-0` and `libgl1-mesa-glx` for OpenCV to work. For example, on a Debian system: `sudo apt update && sudo apt install -y libglib2.0-0 libgl1-mesa-glx`
 
 ## Drivers
 
@@ -124,4 +175,7 @@ An alternative to installing ROCm locally is to use a [ROCm docker container] to
 [ROCm Documentation]: https://rocm.docs.amd.com/projects/install-on-linux/en/latest/tutorial/quick-start.html
 [cuDNN support matrix]: https://docs.nvidia.com/deeplearning/cudnn/support-matrix/index.html
 [Nvidia Container Runtime]: https://developer.nvidia.com/container-runtime
+[linux install instructions]: https://docs.python-guide.org/starting/install3/linux/
+[Microsoft Visual C++ Redistributable]: https://learn.microsoft.com/en-US/cpp/windows/latest-supported-vc-redist?view=msvc-170
+[an official installer]: https://www.python.org/downloads/
 [CUDA Toolkit Downloads]: https://developer.nvidia.com/cuda-downloads

docs/nodes/communityNodes.md

@@ -49,7 +49,6 @@ To use a community workflow, download the `.json` node graph file and load it in
     + [BriaAI Background Remove](#briaai-remove-background)
     + [Remove Background](#remove-background)    
     + [Retroize](#retroize)
-    + [Stereogram](#stereogram-nodes)
     + [Size Stepper Nodes](#size-stepper-nodes)
     + [Simple Skin Detection](#simple-skin-detection)
     + [Text font to Image](#text-font-to-image)
@@ -527,16 +526,6 @@ View:
 
 <img src="https://github.com/Ar7ific1al/InvokeAI_nodes_retroize/assets/2306586/de8b4fa6-324c-4c2d-b36c-297600c73974" width="500" />
 
---------------------------------
-### Stereogram Nodes
-
-**Description:** A set of custom nodes for InvokeAI to create cross-view or parallel-view stereograms. Stereograms are 2D images that, when viewed properly, reveal a 3D scene. Check out [r/crossview](https://www.reddit.com/r/CrossView/) for tutorials.
-
-**Node Link:** https://github.com/simonfuhrmann/invokeai-stereo
-
-**Example Workflow and Output**
-</br><img src="https://github.com/simonfuhrmann/invokeai-stereo/blob/main/docs/example_promo_03.jpg" width="500" />
-
 --------------------------------
 ### Simple Skin Detection
 

invokeai/app/api/routers/boards.py

@@ -31,7 +31,7 @@ class DeleteBoardResult(BaseModel):
     response_model=BoardDTO,
 )
 async def create_board(
-    board_name: str = Query(description="The name of the board to create", max_length=300),
+    board_name: str = Query(description="The name of the board to create"),
     is_private: bool = Query(default=False, description="Whether the board is private"),
 ) -> BoardDTO:
     """Creates a board"""

invokeai/app/api/routers/model_manager.py

@@ -4,6 +4,7 @@
 import contextlib
 import io
 import pathlib
+import shutil
 import traceback
 from copy import deepcopy
 from enum import Enum
@@ -20,6 +21,7 @@ from starlette.exceptions import HTTPException
 from typing_extensions import Annotated
 
 from invokeai.app.api.dependencies import ApiDependencies
+from invokeai.app.services.config import get_config
 from invokeai.app.services.model_images.model_images_common import ModelImageFileNotFoundException
 from invokeai.app.services.model_install.model_install_common import ModelInstallJob
 from invokeai.app.services.model_records import (
@@ -35,7 +37,7 @@ from invokeai.backend.model_manager.config import (
     ModelFormat,
     ModelType,
 )
-from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
+from invokeai.backend.model_manager.load.model_cache.model_cache_base import CacheStats
 from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch
 from invokeai.backend.model_manager.metadata.metadata_base import ModelMetadataWithFiles, UnknownMetadataException
 from invokeai.backend.model_manager.search import ModelSearch
@@ -846,6 +848,74 @@ async def get_starter_models() -> StarterModelResponse:
     return StarterModelResponse(starter_models=starter_models, starter_bundles=starter_bundles)
 
 
+@model_manager_router.get(
+    "/model_cache",
+    operation_id="get_cache_size",
+    response_model=float,
+    summary="Get maximum size of model manager RAM or VRAM cache.",
+)
+async def get_cache_size(cache_type: CacheType = Query(description="The cache type", default=CacheType.RAM)) -> float:
+    """Return the current RAM or VRAM cache size setting (in GB)."""
+    cache = ApiDependencies.invoker.services.model_manager.load.ram_cache
+    value = 0.0
+    if cache_type == CacheType.RAM:
+        value = cache.max_cache_size
+    elif cache_type == CacheType.VRAM:
+        value = cache.max_vram_cache_size
+    return value
+
+
+@model_manager_router.put(
+    "/model_cache",
+    operation_id="set_cache_size",
+    response_model=float,
+    summary="Set maximum size of model manager RAM or VRAM cache, optionally writing new value out to invokeai.yaml config file.",
+)
+async def set_cache_size(
+    value: float = Query(description="The new value for the maximum cache size"),
+    cache_type: CacheType = Query(description="The cache type", default=CacheType.RAM),
+    persist: bool = Query(description="Write new value out to invokeai.yaml", default=False),
+) -> float:
+    """Set the current RAM or VRAM cache size setting (in GB). ."""
+    cache = ApiDependencies.invoker.services.model_manager.load.ram_cache
+    app_config = get_config()
+    # Record initial state.
+    vram_old = app_config.vram
+    ram_old = app_config.ram
+
+    # Prepare target state.
+    vram_new = vram_old
+    ram_new = ram_old
+    if cache_type == CacheType.RAM:
+        ram_new = value
+    elif cache_type == CacheType.VRAM:
+        vram_new = value
+    else:
+        raise ValueError(f"Unexpected {cache_type=}.")
+
+    config_path = app_config.config_file_path
+    new_config_path = config_path.with_suffix(".yaml.new")
+
+    try:
+        # Try to apply the target state.
+        cache.max_vram_cache_size = vram_new
+        cache.max_cache_size = ram_new
+        app_config.ram = ram_new
+        app_config.vram = vram_new
+        if persist:
+            app_config.write_file(new_config_path)
+            shutil.move(new_config_path, config_path)
+    except Exception as e:
+        # If there was a failure, restore the initial state.
+        cache.max_cache_size = ram_old
+        cache.max_vram_cache_size = vram_old
+        app_config.ram = ram_old
+        app_config.vram = vram_old
+
+        raise RuntimeError("Failed to update cache size") from e
+    return value
+
+
 @model_manager_router.get(
     "/stats",
     operation_id="get_stats",

invokeai/app/api_app.py

@@ -59,32 +59,11 @@ logger.info(f"Using torch device: {torch_device_name}")
 
 loop = asyncio.new_event_loop()
 
-# We may change the port if the default is in use, this global variable is used to store the port so that we can log
-# the correct port when the server starts in the lifespan handler.
-port = app_config.port
-
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Add startup event to load dependencies
     ApiDependencies.initialize(config=app_config, event_handler_id=event_handler_id, loop=loop, logger=logger)
-
-    # Log the server address when it starts - in case the network log level is not high enough to see the startup log
-    proto = "https" if app_config.ssl_certfile else "http"
-    msg = f"Invoke running on {proto}://{app_config.host}:{port} (Press CTRL+C to quit)"
-
-    # Logging this way ignores the logger's log level and _always_ logs the message
-    record = logger.makeRecord(
-        name=logger.name,
-        level=logging.INFO,
-        fn="",
-        lno=0,
-        msg=msg,
-        args=(),
-        exc_info=None,
-    )
-    logger.handle(record)
-
     yield
     # Shut down threads
     ApiDependencies.shutdown()
@@ -227,7 +206,6 @@ def invoke_api() -> None:
         else:
             jurigged.watch(logger=InvokeAILogger.get_logger(name="jurigged").info)
 
-    global port
     port = find_port(app_config.port)
     if port != app_config.port:
         logger.warn(f"Port {app_config.port} in use, using port {port}")
@@ -239,17 +217,18 @@ def invoke_api() -> None:
         host=app_config.host,
         port=port,
         loop="asyncio",
-        log_level=app_config.log_level_network,
+        log_level=app_config.log_level,
         ssl_certfile=app_config.ssl_certfile,
         ssl_keyfile=app_config.ssl_keyfile,
     )
     server = uvicorn.Server(config)
 
     # replace uvicorn's loggers with InvokeAI's for consistent appearance
-    uvicorn_logger = InvokeAILogger.get_logger("uvicorn")
-    uvicorn_logger.handlers.clear()
-    for hdlr in logger.handlers:
-        uvicorn_logger.addHandler(hdlr)
+    for logname in ["uvicorn.access", "uvicorn"]:
+        log = InvokeAILogger.get_logger(logname)
+        log.handlers.clear()
+        for ch in logger.handlers:
+            log.addHandler(ch)
 
     loop.run_until_complete(server.serve())
 

invokeai/app/invocations/compel.py

@@ -19,9 +19,9 @@ from invokeai.app.invocations.model import CLIPField
 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.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoRAPatcher
 from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.patches.layer_patcher import LayerPatcher
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
     ConditioningFieldData,
@@ -63,28 +63,30 @@ class CompelInvocation(BaseInvocation):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
-        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
+        tokenizer_info = context.models.load(self.clip.tokenizer)
+        text_encoder_info = context.models.load(self.clip.text_encoder)
+
+        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.clip.loras:
                 lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, ModelPatchRaw)
+                assert isinstance(lora_info.model, LoRAModelRaw)
                 yield (lora_info.model, lora.weight)
                 del lora_info
             return
 
         # loras = [(context.models.get(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras]
 
-        text_encoder_info = context.models.load(self.clip.text_encoder)
         ti_list = generate_ti_list(self.prompt, text_encoder_info.config.base, context)
 
         with (
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
-            context.models.load(self.clip.tokenizer) as tokenizer,
-            LayerPatcher.apply_smart_model_patches(
+            tokenizer_info as tokenizer,
+            LoRAPatcher.apply_smart_lora_patches(
                 model=text_encoder,
                 patches=_lora_loader(),
                 prefix="lora_te_",
-                dtype=text_encoder.dtype,
+                dtype=TorchDevice.choose_torch_dtype(),
                 cached_weights=cached_weights,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
@@ -103,7 +105,6 @@ class CompelInvocation(BaseInvocation):
                 textual_inversion_manager=ti_manager,
                 dtype_for_device_getter=TorchDevice.choose_torch_dtype,
                 truncate_long_prompts=False,
-                device=TorchDevice.choose_torch_device(),
             )
 
             conjunction = Compel.parse_prompt_string(self.prompt)
@@ -138,7 +139,9 @@ class SDXLPromptInvocationBase:
         lora_prefix: str,
         zero_on_empty: bool,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        tokenizer_info = context.models.load(clip_field.tokenizer)
         text_encoder_info = context.models.load(clip_field.text_encoder)
+
         # return zero on empty
         if prompt == "" and zero_on_empty:
             cpu_text_encoder = text_encoder_info.model
@@ -160,11 +163,11 @@ class SDXLPromptInvocationBase:
                 c_pooled = None
             return c, c_pooled
 
-        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
+        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in clip_field.loras:
                 lora_info = context.models.load(lora.lora)
                 lora_model = lora_info.model
-                assert isinstance(lora_model, ModelPatchRaw)
+                assert isinstance(lora_model, LoRAModelRaw)
                 yield (lora_model, lora.weight)
                 del lora_info
             return
@@ -176,12 +179,12 @@ class SDXLPromptInvocationBase:
         with (
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
-            context.models.load(clip_field.tokenizer) as tokenizer,
-            LayerPatcher.apply_smart_model_patches(
-                model=text_encoder,
+            tokenizer_info as tokenizer,
+            LoRAPatcher.apply_smart_lora_patches(
+                text_encoder,
                 patches=_lora_loader(),
                 prefix=lora_prefix,
-                dtype=text_encoder.dtype,
+                dtype=TorchDevice.choose_torch_dtype(),
                 cached_weights=cached_weights,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
@@ -204,7 +207,6 @@ class SDXLPromptInvocationBase:
                 truncate_long_prompts=False,  # TODO:
                 returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,  # TODO: clip skip
                 requires_pooled=get_pooled,
-                device=TorchDevice.choose_torch_device(),
             )
 
             conjunction = Compel.parse_prompt_string(prompt)
@@ -222,6 +224,7 @@ class SDXLPromptInvocationBase:
 
         del tokenizer
         del text_encoder
+        del tokenizer_info
         del text_encoder_info
 
         c = c.detach().to("cpu")

invokeai/app/invocations/constants.py

@@ -1,5 +1,7 @@
 from typing import Literal
 
+from invokeai.backend.util.devices import TorchDevice
+
 LATENT_SCALE_FACTOR = 8
 """
 HACK: Many nodes are currently hard-coded to use a fixed latent scale factor of 8. This is fragile, and will need to
@@ -10,3 +12,5 @@ The ratio of image:latent dimensions is LATENT_SCALE_FACTOR:1, or 8:1.
 
 IMAGE_MODES = Literal["L", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "I", "F"]
 """A literal type for PIL image modes supported by Invoke"""
+
+DEFAULT_PRECISION = TorchDevice.choose_torch_dtype()

invokeai/app/invocations/create_denoise_mask.py

@@ -6,6 +6,7 @@ from PIL import Image
 from torchvision.transforms.functional import resize as tv_resize
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
+from invokeai.app.invocations.constants import DEFAULT_PRECISION
 from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField
 from invokeai.app.invocations.image_to_latents import ImageToLatentsInvocation
 from invokeai.app.invocations.model import VAEField
@@ -28,7 +29,11 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
     image: Optional[ImageField] = InputField(default=None, description="Image which will be masked", ui_order=1)
     mask: ImageField = InputField(description="The mask to use when pasting", ui_order=2)
     tiled: bool = InputField(default=False, description=FieldDescriptions.tiled, ui_order=3)
-    fp32: bool = InputField(default=False, description=FieldDescriptions.fp32, ui_order=4)
+    fp32: bool = InputField(
+        default=DEFAULT_PRECISION == torch.float32,
+        description=FieldDescriptions.fp32,
+        ui_order=4,
+    )
 
     def prep_mask_tensor(self, mask_image: Image.Image) -> torch.Tensor:
         if mask_image.mode != "L":

invokeai/app/invocations/create_gradient_mask.py

@@ -7,6 +7,7 @@ from PIL import Image, ImageFilter
 from torchvision.transforms.functional import resize as tv_resize
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
+from invokeai.app.invocations.constants import DEFAULT_PRECISION
 from invokeai.app.invocations.fields import (
     DenoiseMaskField,
     FieldDescriptions,
@@ -75,7 +76,11 @@ class CreateGradientMaskInvocation(BaseInvocation):
         ui_order=7,
     )
     tiled: bool = InputField(default=False, description=FieldDescriptions.tiled, ui_order=8)
-    fp32: bool = InputField(default=False, description=FieldDescriptions.fp32, ui_order=9)
+    fp32: bool = InputField(
+        default=DEFAULT_PRECISION == torch.float32,
+        description=FieldDescriptions.fp32,
+        ui_order=9,
+    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> GradientMaskOutput:

invokeai/app/invocations/denoise_latents.py

@@ -37,10 +37,10 @@ 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
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoRAPatcher
 from invokeai.backend.model_manager import BaseModelType, ModelVariantType
 from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.patches.layer_patcher import LayerPatcher
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.stable_diffusion import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext, DenoiseInputs
 from invokeai.backend.stable_diffusion.diffusers_pipeline import (
@@ -547,6 +547,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         for single_ip_adapter in ip_adapters:
             with context.models.load(single_ip_adapter.ip_adapter_model) as ip_adapter_model:
                 assert isinstance(ip_adapter_model, IPAdapter)
+                image_encoder_model_info = context.models.load(single_ip_adapter.image_encoder_model)
                 # `single_ip_adapter.image` could be a list or a single ImageField. Normalize to a list here.
                 single_ipa_image_fields = single_ip_adapter.image
                 if not isinstance(single_ipa_image_fields, list):
@@ -555,7 +556,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 single_ipa_images = [
                     context.images.get_pil(image.image_name, mode="RGB") for image in single_ipa_image_fields
                 ]
-                with context.models.load(single_ip_adapter.image_encoder_model) as image_encoder_model:
+                with image_encoder_model_info as image_encoder_model:
                     assert isinstance(image_encoder_model, CLIPVisionModelWithProjection)
                     # Get image embeddings from CLIP and ImageProjModel.
                     image_prompt_embeds, uncond_image_prompt_embeds = ip_adapter_model.get_image_embeds(
@@ -620,6 +621,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         t2i_adapter_data = []
         for t2i_adapter_field in t2i_adapter:
             t2i_adapter_model_config = context.models.get_config(t2i_adapter_field.t2i_adapter_model.key)
+            t2i_adapter_loaded_model = context.models.load(t2i_adapter_field.t2i_adapter_model)
             image = context.images.get_pil(t2i_adapter_field.image.image_name, mode="RGB")
 
             # The max_unet_downscale is the maximum amount that the UNet model downscales the latent image internally.
@@ -635,7 +637,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 raise ValueError(f"Unexpected T2I-Adapter base model type: '{t2i_adapter_model_config.base}'.")
 
             t2i_adapter_model: T2IAdapter
-            with context.models.load(t2i_adapter_field.t2i_adapter_model) as t2i_adapter_model:
+            with t2i_adapter_loaded_model as t2i_adapter_model:
                 total_downscale_factor = t2i_adapter_model.total_downscale_factor
 
                 # Note: We have hard-coded `do_classifier_free_guidance=False`. This is because we only want to prepare
@@ -924,8 +926,10 @@ class DenoiseLatentsInvocation(BaseInvocation):
             # ext: t2i/ip adapter
             ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx)
 
+            unet_info = context.models.load(self.unet.unet)
+            assert isinstance(unet_info.model, UNet2DConditionModel)
             with (
-                context.models.load(self.unet.unet).model_on_device() as (cached_weights, unet),
+                unet_info.model_on_device() as (cached_weights, unet),
                 ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls),
                 # ext: controlnet
                 ext_manager.patch_extensions(denoise_ctx),
@@ -983,21 +987,23 @@ class DenoiseLatentsInvocation(BaseInvocation):
         def step_callback(state: PipelineIntermediateState) -> None:
             context.util.sd_step_callback(state, unet_config.base)
 
-        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
+        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.unet.loras:
                 lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, ModelPatchRaw)
+                assert isinstance(lora_info.model, LoRAModelRaw)
                 yield (lora_info.model, lora.weight)
                 del lora_info
             return
 
+        unet_info = context.models.load(self.unet.unet)
+        assert isinstance(unet_info.model, UNet2DConditionModel)
         with (
             ExitStack() as exit_stack,
-            context.models.load(self.unet.unet).model_on_device() as (cached_weights, unet),
+            unet_info.model_on_device() as (cached_weights, unet),
             ModelPatcher.apply_freeu(unet, self.unet.freeu_config),
             SeamlessExt.static_patch_model(unet, self.unet.seamless_axes),  # FIXME
             # Apply the LoRA after unet has been moved to its target device for faster patching.
-            LayerPatcher.apply_smart_model_patches(
+            LoRAPatcher.apply_smart_lora_patches(
                 model=unet,
                 patches=_lora_loader(),
                 prefix="lora_unet_",

invokeai/app/invocations/fields.py

@@ -56,7 +56,7 @@ class UIType(str, Enum, metaclass=MetaEnum):
     CLIPLEmbedModel = "CLIPLEmbedModelField"
     CLIPGEmbedModel = "CLIPGEmbedModelField"
     SpandrelImageToImageModel = "SpandrelImageToImageModelField"
-    ControlLoRAModel = "ControlLoRAModelField"
+    StructuralLoRAModel = "StructuralLoRAModelField"
     # endregion
 
     # region Misc Field Types
@@ -144,7 +144,7 @@ class FieldDescriptions:
     controlnet_model = "ControlNet model to load"
     vae_model = "VAE model to load"
     lora_model = "LoRA model to load"
-    control_lora_model = "Control LoRA model to load"
+    structural_lora_model = "Structural LoRA model to load"
     main_model = "Main model (UNet, VAE, CLIP) to load"
     flux_model = "Flux model (Transformer) to load"
     sd3_model = "SD3 model (MMDiTX) to load"

invokeai/app/invocations/flux_control_lora_loader.py

@@ -1,49 +0,0 @@
-from invokeai.app.invocations.baseinvocation import (
-    BaseInvocation,
-    BaseInvocationOutput,
-    Classification,
-    invocation,
-    invocation_output,
-)
-from invokeai.app.invocations.fields import FieldDescriptions, ImageField, InputField, OutputField, UIType
-from invokeai.app.invocations.model import ControlLoRAField, ModelIdentifierField
-from invokeai.app.services.shared.invocation_context import InvocationContext
-
-
-@invocation_output("flux_control_lora_loader_output")
-class FluxControlLoRALoaderOutput(BaseInvocationOutput):
-    """Flux Control LoRA Loader Output"""
-
-    control_lora: ControlLoRAField = OutputField(
-        title="Flux Control LoRA", description="Control LoRAs to apply on model loading", default=None
-    )
-
-
-@invocation(
-    "flux_control_lora_loader",
-    title="Flux Control LoRA",
-    tags=["lora", "model", "flux"],
-    category="model",
-    version="1.1.0",
-    classification=Classification.Prototype,
-)
-class FluxControlLoRALoaderInvocation(BaseInvocation):
-    """LoRA model and Image to use with FLUX transformer generation."""
-
-    lora: ModelIdentifierField = InputField(
-        description=FieldDescriptions.control_lora_model, title="Control LoRA", ui_type=UIType.ControlLoRAModel
-    )
-    image: ImageField = InputField(description="The image to encode.")
-    weight: float = InputField(description="The weight of the LoRA.", default=1.0)
-
-    def invoke(self, context: InvocationContext) -> FluxControlLoRALoaderOutput:
-        if not context.models.exists(self.lora.key):
-            raise ValueError(f"Unknown lora: {self.lora.key}!")
-
-        return FluxControlLoRALoaderOutput(
-            control_lora=ControlLoRAField(
-                lora=self.lora,
-                img=self.image,
-                weight=self.weight,
-            )
-        )

invokeai/app/invocations/flux_denoise.py

@@ -1,15 +1,15 @@
 from contextlib import ExitStack
 from typing import Callable, Iterator, Optional, Tuple, Union
 
-import einops
 import numpy as np
 import numpy.typing as npt
 import torch
 import torchvision.transforms as tv_transforms
-from PIL import Image
 from torchvision.transforms.functional import resize as tv_resize
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+
 from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
 from invokeai.app.invocations.fields import (
     DenoiseMaskField,
@@ -23,9 +23,8 @@ from invokeai.app.invocations.fields import (
     WithMetadata,
 )
 from invokeai.app.invocations.flux_controlnet import FluxControlNetField
-from invokeai.app.invocations.flux_vae_encode import FluxVaeEncodeInvocation
 from invokeai.app.invocations.ip_adapter import IPAdapterField
-from invokeai.app.invocations.model import ControlLoRAField, LoRAField, TransformerField, VAEField
+from invokeai.app.invocations.model import TransformerField, VAEField, StructuralLoRAField, LoRAField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.controlnet.instantx_controlnet_flux import InstantXControlNetFlux
@@ -46,11 +45,13 @@ from invokeai.backend.flux.sampling_utils import (
     pack,
     unpack,
 )
+from invokeai.backend.flux.flux_tools_sampling_utils import prepare_control
+from invokeai.backend.flux.modules.conditioner import HFEncoder
 from invokeai.backend.flux.text_conditioning import FluxTextConditioning
+from invokeai.backend.lora.conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoRAPatcher
 from invokeai.backend.model_manager.config import ModelFormat
-from invokeai.backend.patches.layer_patcher import LayerPatcher
-from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
@@ -92,9 +93,6 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         input=Input.Connection,
         title="Transformer",
     )
-    control_lora: Optional[ControlLoRAField] = InputField(
-        description=FieldDescriptions.control_lora_model, input=Input.Connection, title="Control LoRA", default=None
-    )
     positive_text_conditioning: FluxConditioningField | list[FluxConditioningField] = InputField(
         description=FieldDescriptions.positive_cond, input=Input.Connection
     )
@@ -199,8 +197,8 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
             else None
         )
 
-        transformer_config = context.models.get_config(self.transformer.transformer)
-        is_schnell = "schnell" in getattr(transformer_config, "config_path", "")
+        transformer_info = context.models.load(self.transformer.transformer)
+        is_schnell = "schnell" in transformer_info.config.config_path
 
         # Calculate the timestep schedule.
         timesteps = get_schedule(
@@ -240,12 +238,6 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         if len(timesteps) <= 1:
             return x
 
-        if is_schnell and self.control_lora:
-            raise ValueError("Control LoRAs cannot be used with FLUX Schnell")
-
-        # Prepare the extra image conditioning tensor if a FLUX structural control image is provided.
-        img_cond = self._prep_structural_control_img_cond(context)
-
         inpaint_mask = self._prep_inpaint_mask(context, x)
 
         img_ids = generate_img_ids(h=latent_h, w=latent_w, batch_size=b, device=x.device, dtype=x.dtype)
@@ -253,7 +245,6 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         # Pack all latent tensors.
         init_latents = pack(init_latents) if init_latents is not None else None
         inpaint_mask = pack(inpaint_mask) if inpaint_mask is not None else None
-        img_cond = pack(img_cond) if img_cond is not None else None
         noise = pack(noise)
         x = pack(x)
 
@@ -297,42 +288,54 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 dtype=inference_dtype,
                 device=x.device,
             )
+            img_cond = None
+            if struct_lora := self.transformer.structural_lora:
+                # What should we do when we have multiple of these?
+                if not self.controlnet_vae:
+                    raise ValueError("controlnet_vae must be set when using a strutural lora")
+                ae_info = context.models.load(self.controlnet_vae.vae)
+                img = context.images.get_pil(struct_lora.img.image_name)
+                with ae_info as ae:
+                    assert isinstance(ae, AutoEncoder)
+                    img_cond = prepare_control(self.height, self.width, self.seed, ae, img)
 
             # Load the transformer model.
-            (cached_weights, transformer) = exit_stack.enter_context(
-                context.models.load(self.transformer.transformer).model_on_device()
-            )
+            (cached_weights, transformer) = exit_stack.enter_context(transformer_info.model_on_device())
             assert isinstance(transformer, Flux)
-            config = transformer_config
+            config = transformer_info.config
             assert config is not None
 
-            # Determine if the model is quantized.
-            # If the model is quantized, then we need to apply the LoRA weights as sidecar layers. This results in
-            # slower inference than direct patching, but is agnostic to the quantization format.
+            # Apply LoRA models to the transformer.
+            # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching.
             if config.format in [ModelFormat.Checkpoint]:
-                model_is_quantized = False
+                # The model is non-quantized, so we can apply the LoRA weights directly into the model.
+                exit_stack.enter_context(
+                    LoRAPatcher.apply_smart_lora_patches(
+                        model=transformer,
+                        patches=self._lora_iterator(context),
+                        prefix=FLUX_LORA_TRANSFORMER_PREFIX,
+                        dtype=inference_dtype,
+                        cached_weights=cached_weights,
+                    )
+                )
             elif config.format in [
                 ModelFormat.BnbQuantizedLlmInt8b,
                 ModelFormat.BnbQuantizednf4b,
                 ModelFormat.GGUFQuantized,
             ]:
-                model_is_quantized = True
+                # The model is quantized, so apply the LoRA weights as sidecar layers. This results in slower inference,
+                # than directly patching the weights, but is agnostic to the quantization format.
+                exit_stack.enter_context(
+                    LoRAPatcher.apply_lora_wrapper_patches(
+                        model=transformer,
+                        patches=self._lora_iterator(context),
+                        prefix=FLUX_LORA_TRANSFORMER_PREFIX,
+                        dtype=inference_dtype,
+                    )
+                )
             else:
                 raise ValueError(f"Unsupported model format: {config.format}")
 
-            # Apply LoRA models to the transformer.
-            # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching.
-            exit_stack.enter_context(
-                LayerPatcher.apply_smart_model_patches(
-                    model=transformer,
-                    patches=self._lora_iterator(context),
-                    prefix=FLUX_LORA_TRANSFORMER_PREFIX,
-                    dtype=inference_dtype,
-                    cached_weights=cached_weights,
-                    force_sidecar_patching=model_is_quantized,
-                )
-            )
-
             # Prepare IP-Adapter extensions.
             pos_ip_adapter_extensions, neg_ip_adapter_extensions = self._prep_ip_adapter_extensions(
                 pos_image_prompt_clip_embeds=pos_image_prompt_clip_embeds,
@@ -357,7 +360,7 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 controlnet_extensions=controlnet_extensions,
                 pos_ip_adapter_extensions=pos_ip_adapter_extensions,
                 neg_ip_adapter_extensions=neg_ip_adapter_extensions,
-                img_cond=img_cond,
+                img_cond=img_cond
             )
 
         x = unpack(x.float(), self.height, self.width)
@@ -514,18 +517,15 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         # before loading the models. Then make sure that all VAE encoding is done before loading the ControlNets to
         # minimize peak memory.
 
+        # First, load the ControlNet models so that we can determine the ControlNet types.
+        controlnet_models = [context.models.load(controlnet.control_model) for controlnet in controlnets]
+
         # Calculate the controlnet conditioning tensors.
         # We do this before loading the ControlNet models because it may require running the VAE, and we are trying to
         # keep peak memory down.
         controlnet_conds: list[torch.Tensor] = []
-        for controlnet in controlnets:
+        for controlnet, controlnet_model in zip(controlnets, controlnet_models, strict=True):
             image = context.images.get_pil(controlnet.image.image_name)
-
-            # HACK(ryand): We have to load the ControlNet model to determine whether the VAE needs to be run. We really
-            # shouldn't have to load the model here. There's a risk that the model will be dropped from the model cache
-            # before we load it into VRAM and thus we'll have to load it again (context:
-            # https://github.com/invoke-ai/InvokeAI/issues/7513).
-            controlnet_model = context.models.load(controlnet.control_model)
             if isinstance(controlnet_model.model, InstantXControlNetFlux):
                 if self.controlnet_vae is None:
                     raise ValueError("A ControlNet VAE is required when using an InstantX FLUX ControlNet.")
@@ -555,8 +555,10 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         # Finally, load the ControlNet models and initialize the ControlNet extensions.
         controlnet_extensions: list[XLabsControlNetExtension | InstantXControlNetExtension] = []
-        for controlnet, controlnet_cond in zip(controlnets, controlnet_conds, strict=True):
-            model = exit_stack.enter_context(context.models.load(controlnet.control_model))
+        for controlnet, controlnet_cond, controlnet_model in zip(
+            controlnets, controlnet_conds, controlnet_models, strict=True
+        ):
+            model = exit_stack.enter_context(controlnet_model)
 
             if isinstance(model, XLabsControlNetFlux):
                 controlnet_extensions.append(
@@ -589,29 +591,6 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         return controlnet_extensions
 
-    def _prep_structural_control_img_cond(self, context: InvocationContext) -> torch.Tensor | None:
-        if self.control_lora is None:
-            return None
-
-        if not self.controlnet_vae:
-            raise ValueError("controlnet_vae must be set when using a FLUX Control LoRA.")
-
-        # Load the conditioning image and resize it to the target image size.
-        cond_img = context.images.get_pil(self.control_lora.img.image_name)
-        cond_img = cond_img.convert("RGB")
-        cond_img = cond_img.resize((self.width, self.height), Image.Resampling.BICUBIC)
-        cond_img = np.array(cond_img)
-
-        # Normalize the conditioning image to the range [-1, 1].
-        # This normalization is based on the original implementations here:
-        # https://github.com/black-forest-labs/flux/blob/805da8571a0b49b6d4043950bd266a65328c243b/src/flux/modules/image_embedders.py#L34
-        # https://github.com/black-forest-labs/flux/blob/805da8571a0b49b6d4043950bd266a65328c243b/src/flux/modules/image_embedders.py#L60
-        img_cond = torch.from_numpy(cond_img).float() / 127.5 - 1.0
-        img_cond = einops.rearrange(img_cond, "h w c -> 1 c h w")
-
-        vae_info = context.models.load(self.controlnet_vae.vae)
-        return FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=img_cond)
-
     def _normalize_ip_adapter_fields(self) -> list[IPAdapterField]:
         if self.ip_adapter is None:
             return []
@@ -718,15 +697,13 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         return pos_ip_adapter_extensions, neg_ip_adapter_extensions
 
-    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
-        loras: list[Union[LoRAField, ControlLoRAField]] = [*self.transformer.loras]
-        if self.control_lora:
-            # Note: Since FLUX structural control LoRAs modify the shape of some weights, it is important that they are
-            # applied last.
-            loras.append(self.control_lora)
+    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[LoRAModelRaw, float]]:
+        loras: list[Union[LoRAField, StructuralLoRAField]] = [*self.transformer.loras]
+        if self.transformer.structural_lora:
+            loras.append(self.transformer.structural_lora)
         for lora in loras:
             lora_info = context.models.load(lora.lora)
-            assert isinstance(lora_info.model, ModelPatchRaw)
+            assert isinstance(lora_info.model, LoRAModelRaw)
             yield (lora_info.model, lora.weight)
             del lora_info
 

invokeai/app/invocations/flux_model_loader.py

@@ -81,8 +81,8 @@ class FluxModelLoaderInvocation(BaseInvocation):
         assert isinstance(transformer_config, CheckpointConfigBase)
 
         return FluxModelLoaderOutput(
-            transformer=TransformerField(transformer=transformer, loras=[]),
-            clip=CLIPField(tokenizer=tokenizer, text_encoder=clip_encoder, loras=[], skipped_layers=0),
+            transformer=TransformerField(transformer=transformer, loras=[], structural_loras=[]),
+            clip=CLIPField(tokenizer=tokenizer, text_encoder=clip_encoder, loras=[], structural_loras=[], skipped_layers=0),
             t5_encoder=T5EncoderField(tokenizer=tokenizer2, text_encoder=t5_encoder),
             vae=VAEField(vae=vae),
             max_seq_len=max_seq_lengths[transformer_config.config_path],

invokeai/app/invocations/flux_structural_lora_loader.py

@@ -0,0 +1,70 @@
+from typing import Optional, Literal
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType, ImageField
+from invokeai.app.invocations.model import VAEField, StructuralLoRAField, ModelIdentifierField, TransformerField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+
+@invocation_output("flux_structural_lora_loader_output")
+class FluxStructuralLoRALoaderOutput(BaseInvocationOutput):
+    """Flux Structural LoRA Loader Output"""
+
+    transformer: Optional[TransformerField] = OutputField(
+        default=None, description=FieldDescriptions.transformer, title="FLUX Transformer"
+    )
+
+
+@invocation(
+    "flux_structural_lora_loader",
+    title="Flux Structural LoRA",
+    tags=["lora", "model", "flux"],
+    category="model",
+    version="1.1.0",
+    classification=Classification.Prototype,
+)
+class FluxStructuralLoRALoaderInvocation(BaseInvocation):
+    """Apply a LoRA model to a FLUX transformer and/or text encoder."""
+
+    lora: ModelIdentifierField = InputField(
+        description=FieldDescriptions.structural_lora_model, title="Structural LoRA", ui_type=UIType.StructuralLoRAModel
+    )
+    transformer: TransformerField | None = InputField(
+        default=None,
+        description=FieldDescriptions.transformer,
+        input=Input.Connection,
+        title="FLUX Transformer",
+    )
+    image: ImageField = InputField(
+        description="The image to encode.",
+    )
+    weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
+
+    def invoke(self, context: InvocationContext) -> FluxStructuralLoRALoaderOutput:
+        lora_key = self.lora.key
+
+        if not context.models.exists(lora_key):
+            raise ValueError(f"Unknown lora: {lora_key}!")
+
+        # Check for existing LoRAs with the same key.
+        if self.transformer and self.transformer.structural_lora and self.transformer.structural_lora.lora.key == lora_key:
+            raise ValueError(f'Structural LoRA "{lora_key}" already applied to transformer.')
+
+        output = FluxStructuralLoRALoaderOutput()
+
+        # Attach LoRA layers to the models.
+        if self.transformer is not None:
+            output.transformer = self.transformer.model_copy(deep=True)
+            output.transformer.structural_lora = StructuralLoRAField(
+                lora=self.lora,
+                img=self.image,
+                weight=self.weight,
+            )
+
+        return output

invokeai/app/invocations/flux_text_encoder.py

@@ -17,11 +17,12 @@ from invokeai.app.invocations.model import CLIPField, T5EncoderField
 from invokeai.app.invocations.primitives import FluxConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.modules.conditioner import HFEncoder
+from invokeai.backend.lora.conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoRAPatcher
 from invokeai.backend.model_manager.config import ModelFormat
-from invokeai.backend.patches.layer_patcher import LayerPatcher
-from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, FLUXConditioningInfo
+from invokeai.backend.util.devices import TorchDevice
 
 
 @invocation(
@@ -69,11 +70,14 @@ class FluxTextEncoderInvocation(BaseInvocation):
         )
 
     def _t5_encode(self, context: InvocationContext) -> torch.Tensor:
+        t5_tokenizer_info = context.models.load(self.t5_encoder.tokenizer)
+        t5_text_encoder_info = context.models.load(self.t5_encoder.text_encoder)
+
         prompt = [self.prompt]
 
         with (
-            context.models.load(self.t5_encoder.text_encoder) as t5_text_encoder,
-            context.models.load(self.t5_encoder.tokenizer) as t5_tokenizer,
+            t5_text_encoder_info as t5_text_encoder,
+            t5_tokenizer_info as t5_tokenizer,
         ):
             assert isinstance(t5_text_encoder, T5EncoderModel)
             assert isinstance(t5_tokenizer, T5Tokenizer)
@@ -87,30 +91,32 @@ class FluxTextEncoderInvocation(BaseInvocation):
         return prompt_embeds
 
     def _clip_encode(self, context: InvocationContext) -> torch.Tensor:
-        prompt = [self.prompt]
-
+        clip_tokenizer_info = context.models.load(self.clip.tokenizer)
         clip_text_encoder_info = context.models.load(self.clip.text_encoder)
-        clip_text_encoder_config = clip_text_encoder_info.config
-        assert clip_text_encoder_config is not None
+
+        prompt = [self.prompt]
 
         with (
             clip_text_encoder_info.model_on_device() as (cached_weights, clip_text_encoder),
-            context.models.load(self.clip.tokenizer) as clip_tokenizer,
+            clip_tokenizer_info as clip_tokenizer,
             ExitStack() as exit_stack,
         ):
             assert isinstance(clip_text_encoder, CLIPTextModel)
             assert isinstance(clip_tokenizer, CLIPTokenizer)
 
+            clip_text_encoder_config = clip_text_encoder_info.config
+            assert clip_text_encoder_config is not None
+
             # Apply LoRA models to the CLIP encoder.
             # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching.
             if clip_text_encoder_config.format in [ModelFormat.Diffusers]:
                 # The model is non-quantized, so we can apply the LoRA weights directly into the model.
                 exit_stack.enter_context(
-                    LayerPatcher.apply_smart_model_patches(
+                    LoRAPatcher.apply_smart_lora_patches(
                         model=clip_text_encoder,
                         patches=self._clip_lora_iterator(context),
                         prefix=FLUX_LORA_CLIP_PREFIX,
-                        dtype=clip_text_encoder.dtype,
+                        dtype=TorchDevice.choose_torch_dtype(),
                         cached_weights=cached_weights,
                     )
                 )
@@ -126,9 +132,9 @@ class FluxTextEncoderInvocation(BaseInvocation):
         assert isinstance(pooled_prompt_embeds, torch.Tensor)
         return pooled_prompt_embeds
 
-    def _clip_lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
+    def _clip_lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[LoRAModelRaw, float]]:
         for lora in self.clip.loras:
             lora_info = context.models.load(lora.lora)
-            assert isinstance(lora_info.model, ModelPatchRaw)
+            assert isinstance(lora_info.model, LoRAModelRaw)
             yield (lora_info.model, lora.weight)
             del lora_info

invokeai/app/invocations/flux_vae_decode.py

@@ -3,7 +3,6 @@ from einops import rearrange
 from PIL import Image
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     Input,
@@ -25,7 +24,7 @@ from invokeai.backend.util.devices import TorchDevice
     title="FLUX Latents to Image",
     tags=["latents", "image", "vae", "l2i", "flux"],
     category="latents",
-    version="1.0.1",
+    version="1.0.0",
 )
 class FluxVaeDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents."""
@@ -39,23 +38,8 @@ class FluxVaeDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
         input=Input.Connection,
     )
 
-    def _estimate_working_memory(self, latents: torch.Tensor, vae: AutoEncoder) -> int:
-        """Estimate the working memory required by the invocation in bytes."""
-        # It was found experimentally that the peak working memory scales linearly with the number of pixels and the
-        # element size (precision).
-        out_h = LATENT_SCALE_FACTOR * latents.shape[-2]
-        out_w = LATENT_SCALE_FACTOR * latents.shape[-1]
-        element_size = next(vae.parameters()).element_size()
-        scaling_constant = 1090  # Determined experimentally.
-        working_memory = out_h * out_w * element_size * scaling_constant
-
-        # We add a 20% buffer to the working memory estimate to be safe.
-        working_memory = working_memory * 1.2
-        return int(working_memory)
-
     def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Image:
-        estimated_working_memory = self._estimate_working_memory(latents, vae_info.model)
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info as vae:
             assert isinstance(vae, AutoEncoder)
             vae_dtype = next(iter(vae.parameters())).dtype
             latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)

invokeai/app/invocations/image_to_latents.py

@@ -13,7 +13,7 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
+from invokeai.app.invocations.constants import DEFAULT_PRECISION, LATENT_SCALE_FACTOR
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     ImageField,
@@ -49,7 +49,7 @@ class ImageToLatentsInvocation(BaseInvocation):
     # NOTE: tile_size = 0 is a special value. We use this rather than `int | None`, because the workflow UI does not
     # offer a way to directly set None values.
     tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size)
-    fp32: bool = InputField(default=False, description=FieldDescriptions.fp32)
+    fp32: bool = InputField(default=DEFAULT_PRECISION == torch.float32, description=FieldDescriptions.fp32)
 
     @staticmethod
     def vae_encode(

invokeai/app/invocations/latents_to_image.py

@@ -12,7 +12,7 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
+from invokeai.app.invocations.constants import DEFAULT_PRECISION, LATENT_SCALE_FACTOR
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     Input,
@@ -34,7 +34,7 @@ from invokeai.backend.util.devices import TorchDevice
     title="Latents to Image",
     tags=["latents", "image", "vae", "l2i"],
     category="latents",
-    version="1.3.1",
+    version="1.3.0",
 )
 class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents."""
@@ -51,60 +51,18 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     # NOTE: tile_size = 0 is a special value. We use this rather than `int | None`, because the workflow UI does not
     # offer a way to directly set None values.
     tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size)
-    fp32: bool = InputField(default=False, description=FieldDescriptions.fp32)
-
-    def _estimate_working_memory(
-        self, latents: torch.Tensor, use_tiling: bool, vae: AutoencoderKL | AutoencoderTiny
-    ) -> int:
-        """Estimate the working memory required by the invocation in bytes."""
-        # It was found experimentally that the peak working memory scales linearly with the number of pixels and the
-        # element size (precision). This estimate is accurate for both SD1 and SDXL.
-        element_size = 4 if self.fp32 else 2
-        scaling_constant = 960  # Determined experimentally.
-
-        if use_tiling:
-            tile_size = self.tile_size
-            if tile_size == 0:
-                tile_size = vae.tile_sample_min_size
-                assert isinstance(tile_size, int)
-            out_h = tile_size
-            out_w = tile_size
-            working_memory = out_h * out_w * element_size * scaling_constant
-
-            # We add 25% to the working memory estimate when tiling is enabled to account for factors like tile overlap
-            # and number of tiles. We could make this more precise in the future, but this should be good enough for
-            # most use cases.
-            working_memory = working_memory * 1.25
-        else:
-            out_h = LATENT_SCALE_FACTOR * latents.shape[-2]
-            out_w = LATENT_SCALE_FACTOR * latents.shape[-1]
-            working_memory = out_h * out_w * element_size * scaling_constant
-
-        if self.fp32:
-            # If we are running in FP32, then we should account for the likely increase in model size (~250MB).
-            working_memory += 250 * 2**20
-
-        # We add 20% to the working memory estimate to be safe.
-        working_memory = int(working_memory * 1.2)
-        return working_memory
+    fp32: bool = InputField(default=DEFAULT_PRECISION == torch.float32, description=FieldDescriptions.fp32)
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ImageOutput:
         latents = context.tensors.load(self.latents.latents_name)
 
-        use_tiling = self.tiled or context.config.get().force_tiled_decode
-
         vae_info = context.models.load(self.vae.vae)
         assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny))
-
-        estimated_working_memory = self._estimate_working_memory(latents, use_tiling, vae_info.model)
-        with (
-            SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes),
-            vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae),
-        ):
+        with SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), vae_info as vae:
             context.util.signal_progress("Running VAE decoder")
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
-            latents = latents.to(TorchDevice.choose_torch_device())
+            latents = latents.to(vae.device)
             if self.fp32:
                 vae.to(dtype=torch.float32)
 
@@ -130,7 +88,7 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
                 vae.to(dtype=torch.float16)
                 latents = latents.half()
 
-            if use_tiling:
+            if self.tiled or context.config.get().force_tiled_decode:
                 vae.enable_tiling()
             else:
                 vae.disable_tiling()

invokeai/app/invocations/model.py

@@ -1,5 +1,5 @@
 import copy
-from typing import List, Optional
+from typing import List, Optional, Literal
 
 from pydantic import BaseModel, Field
 
@@ -10,7 +10,7 @@ from invokeai.app.invocations.baseinvocation import (
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField, OutputField, UIType
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType, ImageField
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.shared.models import FreeUConfig
 from invokeai.backend.model_manager.config import (
@@ -74,15 +74,13 @@ class VAEField(BaseModel):
     vae: ModelIdentifierField = Field(description="Info to load vae submodel")
     seamless_axes: List[str] = Field(default_factory=list, description='Axes("x" and "y") to which apply seamless')
 
-
-class ControlLoRAField(LoRAField):
+class StructuralLoRAField(LoRAField):
     img: ImageField = Field(description="Image to use in structural conditioning")
 
-
 class TransformerField(BaseModel):
     transformer: ModelIdentifierField = Field(description="Info to load Transformer submodel")
     loras: List[LoRAField] = Field(description="LoRAs to apply on model loading")
-
+    structural_lora: Optional[StructuralLoRAField] = Field(description="Structural LoRAs to apply on model loading", default=None)
 
 @invocation_output("unet_output")
 class UNetOutput(BaseInvocationOutput):

invokeai/app/invocations/sd3_latents_to_image.py

@@ -6,7 +6,6 @@ from einops import rearrange
 from PIL import Image
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     Input,
@@ -27,7 +26,7 @@ from invokeai.backend.util.devices import TorchDevice
     title="SD3 Latents to Image",
     tags=["latents", "image", "vae", "l2i", "sd3"],
     category="latents",
-    version="1.3.1",
+    version="1.3.0",
 )
 class SD3LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents."""
@@ -41,34 +40,16 @@ class SD3LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
         input=Input.Connection,
     )
 
-    def _estimate_working_memory(self, latents: torch.Tensor, vae: AutoencoderKL) -> int:
-        """Estimate the working memory required by the invocation in bytes."""
-        # It was found experimentally that the peak working memory scales linearly with the number of pixels and the
-        # element size (precision).
-        out_h = LATENT_SCALE_FACTOR * latents.shape[-2]
-        out_w = LATENT_SCALE_FACTOR * latents.shape[-1]
-        element_size = next(vae.parameters()).element_size()
-        scaling_constant = 1230  # Determined experimentally.
-        working_memory = out_h * out_w * element_size * scaling_constant
-
-        # We add a 20% buffer to the working memory estimate to be safe.
-        working_memory = working_memory * 1.2
-        return int(working_memory)
-
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ImageOutput:
         latents = context.tensors.load(self.latents.latents_name)
 
         vae_info = context.models.load(self.vae.vae)
         assert isinstance(vae_info.model, (AutoencoderKL))
-        estimated_working_memory = self._estimate_working_memory(latents, vae_info.model)
-        with (
-            SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes),
-            vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae),
-        ):
+        with SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), vae_info as vae:
             context.util.signal_progress("Running VAE")
             assert isinstance(vae, (AutoencoderKL))
-            latents = latents.to(TorchDevice.choose_torch_device())
+            latents = latents.to(vae.device)
 
             vae.disable_tiling()
 

invokeai/app/invocations/sd3_text_encoder.py

@@ -16,10 +16,10 @@ from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField
 from invokeai.app.invocations.model import CLIPField, T5EncoderField
 from invokeai.app.invocations.primitives import SD3ConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.lora.conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoRAPatcher
 from invokeai.backend.model_manager.config import ModelFormat
-from invokeai.backend.patches.layer_patcher import LayerPatcher
-from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, SD3ConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
 
@@ -87,11 +87,14 @@ class Sd3TextEncoderInvocation(BaseInvocation):
 
     def _t5_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Tensor:
         assert self.t5_encoder is not None
+        t5_tokenizer_info = context.models.load(self.t5_encoder.tokenizer)
+        t5_text_encoder_info = context.models.load(self.t5_encoder.text_encoder)
+
         prompt = [self.prompt]
 
         with (
-            context.models.load(self.t5_encoder.text_encoder) as t5_text_encoder,
-            context.models.load(self.t5_encoder.tokenizer) as t5_tokenizer,
+            t5_text_encoder_info as t5_text_encoder,
+            t5_tokenizer_info as t5_tokenizer,
         ):
             context.util.signal_progress("Running T5 encoder")
             assert isinstance(t5_text_encoder, T5EncoderModel)
@@ -118,7 +121,7 @@ class Sd3TextEncoderInvocation(BaseInvocation):
                     f" {max_seq_len} tokens: {removed_text}"
                 )
 
-            prompt_embeds = t5_text_encoder(text_input_ids.to(TorchDevice.choose_torch_device()))[0]
+            prompt_embeds = t5_text_encoder(text_input_ids.to(t5_text_encoder.device))[0]
 
         assert isinstance(prompt_embeds, torch.Tensor)
         return prompt_embeds
@@ -126,12 +129,14 @@ class Sd3TextEncoderInvocation(BaseInvocation):
     def _clip_encode(
         self, context: InvocationContext, clip_model: CLIPField, tokenizer_max_length: int = 77
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+        clip_tokenizer_info = context.models.load(clip_model.tokenizer)
+        clip_text_encoder_info = context.models.load(clip_model.text_encoder)
+
         prompt = [self.prompt]
 
-        clip_text_encoder_info = context.models.load(clip_model.text_encoder)
         with (
             clip_text_encoder_info.model_on_device() as (cached_weights, clip_text_encoder),
-            context.models.load(clip_model.tokenizer) as clip_tokenizer,
+            clip_tokenizer_info as clip_tokenizer,
             ExitStack() as exit_stack,
         ):
             context.util.signal_progress("Running CLIP encoder")
@@ -146,11 +151,11 @@ class Sd3TextEncoderInvocation(BaseInvocation):
             if clip_text_encoder_config.format in [ModelFormat.Diffusers]:
                 # The model is non-quantized, so we can apply the LoRA weights directly into the model.
                 exit_stack.enter_context(
-                    LayerPatcher.apply_smart_model_patches(
+                    LoRAPatcher.apply_smart_lora_patches(
                         model=clip_text_encoder,
                         patches=self._clip_lora_iterator(context, clip_model),
                         prefix=FLUX_LORA_CLIP_PREFIX,
-                        dtype=clip_text_encoder.dtype,
+                        dtype=TorchDevice.choose_torch_dtype(),
                         cached_weights=cached_weights,
                     )
                 )
@@ -181,7 +186,7 @@ class Sd3TextEncoderInvocation(BaseInvocation):
                     f" {tokenizer_max_length} tokens: {removed_text}"
                 )
             prompt_embeds = clip_text_encoder(
-                input_ids=text_input_ids.to(TorchDevice.choose_torch_device()), output_hidden_states=True
+                input_ids=text_input_ids.to(clip_text_encoder.device), output_hidden_states=True
             )
             pooled_prompt_embeds = prompt_embeds[0]
             prompt_embeds = prompt_embeds.hidden_states[-2]
@@ -190,9 +195,9 @@ class Sd3TextEncoderInvocation(BaseInvocation):
 
     def _clip_lora_iterator(
         self, context: InvocationContext, clip_model: CLIPField
-    ) -> Iterator[Tuple[ModelPatchRaw, float]]:
+    ) -> Iterator[Tuple[LoRAModelRaw, float]]:
         for lora in clip_model.loras:
             lora_info = context.models.load(lora.lora)
-            assert isinstance(lora_info.model, ModelPatchRaw)
+            assert isinstance(lora_info.model, LoRAModelRaw)
             yield (lora_info.model, lora.weight)
             del lora_info

invokeai/app/invocations/spandrel_image_to_image.py

@@ -157,6 +157,9 @@ class SpandrelImageToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
         # revisit this.
         image = context.images.get_pil(self.image.image_name, mode="RGB")
 
+        # Load the model.
+        spandrel_model_info = context.models.load(self.image_to_image_model)
+
         def step_callback(step: int, total_steps: int) -> None:
             context.util.signal_progress(
                 message=f"Processing tile {step}/{total_steps}",
@@ -164,7 +167,7 @@ class SpandrelImageToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
             )
 
         # Do the upscaling.
-        with context.models.load(self.image_to_image_model) as spandrel_model:
+        with spandrel_model_info as spandrel_model:
             assert isinstance(spandrel_model, SpandrelImageToImageModel)
 
             # Upscale the image
@@ -203,6 +206,9 @@ class SpandrelImageToImageAutoscaleInvocation(SpandrelImageToImageInvocation):
         # revisit this.
         image = context.images.get_pil(self.image.image_name, mode="RGB")
 
+        # Load the model.
+        spandrel_model_info = context.models.load(self.image_to_image_model)
+
         # The target size of the image, determined by the provided scale. We'll run the upscaler until we hit this size.
         # Later, we may mutate this value if the model doesn't upscale the image or if the user requested a multiple of 8.
         target_width = int(image.width * self.scale)
@@ -215,7 +221,7 @@ class SpandrelImageToImageAutoscaleInvocation(SpandrelImageToImageInvocation):
             )
 
         # Do the upscaling.
-        with context.models.load(self.image_to_image_model) as spandrel_model:
+        with spandrel_model_info as spandrel_model:
             assert isinstance(spandrel_model, SpandrelImageToImageModel)
 
             iteration = 1

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

@@ -22,8 +22,8 @@ from invokeai.app.invocations.fields import (
 from invokeai.app.invocations.model import UNetField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.patches.layer_patcher import LayerPatcher
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoRAPatcher
 from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState
 from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
     MultiDiffusionPipeline,
@@ -194,17 +194,20 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
             context.util.sd_step_callback(state, unet_config.base)
 
         # Prepare an iterator that yields the UNet's LoRA models and their weights.
-        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
+        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.unet.loras:
                 lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, ModelPatchRaw)
+                assert isinstance(lora_info.model, LoRAModelRaw)
                 yield (lora_info.model, lora.weight)
                 del lora_info
 
+        # Load the UNet model.
+        unet_info = context.models.load(self.unet.unet)
+
         with (
             ExitStack() as exit_stack,
-            context.models.load(self.unet.unet) as unet,
-            LayerPatcher.apply_smart_model_patches(
+            unet_info as unet,
+            LoRAPatcher.apply_smart_lora_patches(
                 model=unet, patches=_lora_loader(), prefix="lora_unet_", dtype=unet.dtype
             ),
         ):

invokeai/app/services/board_records/board_records_common.py

@@ -57,7 +57,7 @@ def deserialize_board_record(board_dict: dict) -> BoardRecord:
 
 
 class BoardChanges(BaseModel, extra="forbid"):
-    board_name: Optional[str] = Field(default=None, description="The board's new name.", max_length=300)
+    board_name: Optional[str] = Field(default=None, description="The board's new name.")
     cover_image_name: Optional[str] = Field(default=None, description="The name of the board's new cover image.")
     archived: Optional[bool] = Field(default=None, description="Whether or not the board is archived")
 

invokeai/app/services/config/config_default.py

@@ -13,6 +13,7 @@ from functools import lru_cache
 from pathlib import Path
 from typing import Any, Literal, Optional
 
+import psutil
 import yaml
 from pydantic import BaseModel, Field, PrivateAttr, field_validator
 from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, SettingsConfigDict
@@ -24,6 +25,8 @@ from invokeai.frontend.cli.arg_parser import InvokeAIArgs
 INIT_FILE = Path("invokeai.yaml")
 DB_FILE = Path("invokeai.db")
 LEGACY_INIT_FILE = Path("invokeai.init")
+DEFAULT_RAM_CACHE = 10.0
+DEFAULT_VRAM_CACHE = 0.25
 DEVICE = Literal["auto", "cpu", "cuda", "cuda:1", "mps"]
 PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
 ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
@@ -33,6 +36,24 @@ LOG_LEVEL = Literal["debug", "info", "warning", "error", "critical"]
 CONFIG_SCHEMA_VERSION = "4.0.2"
 
 
+def get_default_ram_cache_size() -> float:
+    """Run a heuristic for the default RAM cache based on installed RAM."""
+
+    # On some machines, psutil.virtual_memory().total gives a value that is slightly less than the actual RAM, so the
+    # limits are set slightly lower than than what we expect the actual RAM to be.
+
+    GB = 1024**3
+    max_ram = psutil.virtual_memory().total / GB
+
+    if max_ram >= 60:
+        return 15.0
+    if max_ram >= 30:
+        return 7.5
+    if max_ram >= 14:
+        return 4.0
+    return 2.1  # 2.1 is just large enough for sd 1.5 ;-)
+
+
 class URLRegexTokenPair(BaseModel):
     url_regex: str = Field(description="Regular expression to match against the URL")
     token: str = Field(description="Token to use when the URL matches the regex")
@@ -76,18 +97,15 @@ class InvokeAIAppConfig(BaseSettings):
         log_format: Log format. Use "plain" for text-only, "color" for colorized output, "legacy" for 2.3-style logging and "syslog" for syslog-style.<br>Valid values: `plain`, `color`, `syslog`, `legacy`
         log_level: Emit logging messages at this level or higher.<br>Valid values: `debug`, `info`, `warning`, `error`, `critical`
         log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.
-        log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.<br>Valid values: `debug`, `info`, `warning`, `error`, `critical`
         use_memory_db: Use in-memory database. Useful for development.
         dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.
         profile_graphs: Enable graph profiling using `cProfile`.
         profile_prefix: An optional prefix for profile output files.
         profiles_dir: Path to profiles output directory.
-        ram: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.
-        vram: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.
-        lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behaviour is out of beta.
+        ram: Maximum memory amount used by memory model cache for rapid switching (GB).
+        vram: Amount of VRAM reserved for model storage (GB).
+        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_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.
-        enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.
         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`
         sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
@@ -145,7 +163,6 @@ class InvokeAIAppConfig(BaseSettings):
     log_format:              LOG_FORMAT = Field(default="color",            description='Log format. Use "plain" for text-only, "color" for colorized output, "legacy" for 2.3-style logging and "syslog" for syslog-style.')
     log_level:                LOG_LEVEL = Field(default="info",             description="Emit logging messages at this level or higher.")
     log_sql:                       bool = Field(default=False,              description="Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.")
-    log_level_network:        LOG_LEVEL = Field(default='warning',          description="Log level for network-related messages. 'info' and 'debug' are very verbose.")
 
     # Development
     use_memory_db:                 bool = Field(default=False,              description="Use in-memory database. Useful for development.")
@@ -155,12 +172,10 @@ class InvokeAIAppConfig(BaseSettings):
     profiles_dir:                  Path = Field(default=Path("profiles"),   description="Path to profiles output directory.")
 
     # CACHE
-    ram:                Optional[float] = Field(default=None, gt=0,         description="The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.")
-    vram:               Optional[float] = Field(default=None, ge=0,         description="The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.")
-    lazy_offload:                  bool = Field(default=True,               description="DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behaviour is out of beta.")
+    ram:                           float = Field(default_factory=get_default_ram_cache_size, gt=0, description="Maximum memory amount used by memory model cache for rapid switching (GB).")
+    vram:                          float = Field(default=DEFAULT_VRAM_CACHE, ge=0, description="Amount of VRAM reserved for model storage (GB).")
+    lazy_offload:                  bool = Field(default=True,               description="Keep models in VRAM until their space is needed.")
     log_memory_usage:              bool = Field(default=False,              description="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_working_mem_gb:        float = Field(default=3,                  description="The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.")
-    enable_partial_loading:        bool = Field(default=False,              description="Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.")
 
     # DEVICE
     device:                      DEVICE = Field(default="auto",             description="Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.")

invokeai/app/services/download/download_default.py

@@ -8,7 +8,7 @@ import time
 import traceback
 from pathlib import Path
 from queue import Empty, PriorityQueue
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set
+from typing import Any, Dict, List, Literal, Optional, Set
 
 import requests
 from pydantic.networks import AnyHttpUrl
@@ -28,13 +28,11 @@ from invokeai.app.services.download.download_base import (
     ServiceInactiveException,
     UnknownJobIDException,
 )
+from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.util.misc import get_iso_timestamp
 from invokeai.backend.model_manager.metadata import RemoteModelFile
 from invokeai.backend.util.logging import InvokeAILogger
 
-if TYPE_CHECKING:
-    from invokeai.app.services.events.events_base import EventServiceBase
-
 # Maximum number of bytes to download during each call to requests.iter_content()
 DOWNLOAD_CHUNK_SIZE = 100000
 

invokeai/app/services/events/__init__.py

@@ -0,0 +1 @@
+from .events_base import EventServiceBase  # noqa F401

invokeai/app/services/events/events_common.py

@@ -4,7 +4,6 @@ from fastapi_events.handlers.local import local_handler
 from fastapi_events.registry.payload_schema import registry as payload_schema
 from pydantic import BaseModel, ConfigDict, Field
 
-from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
 from invokeai.app.services.session_processor.session_processor_common import ProgressImage
 from invokeai.app.services.session_queue.session_queue_common import (
     QUEUE_ITEM_STATUS,
@@ -19,7 +18,7 @@ from invokeai.backend.model_manager.config import AnyModelConfig, SubModelType
 
 if TYPE_CHECKING:
     from invokeai.app.services.download.download_base import DownloadJob
-    from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
+    from invokeai.app.services.model_install.model_install_common import ModelInstallJob
 
 
 class EventBase(BaseModel):
@@ -423,7 +422,7 @@ class ModelInstallDownloadStartedEvent(ModelEventBase):
     __event_name__ = "model_install_download_started"
 
     id: int = Field(description="The ID of the install job")
-    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
+    source: str = Field(description="Source of the model; local path, repo_id or url")
     local_path: str = Field(description="Where model is downloading to")
     bytes: int = Field(description="Number of bytes downloaded so far")
     total_bytes: int = Field(description="Total size of download, including all files")
@@ -444,7 +443,7 @@ class ModelInstallDownloadStartedEvent(ModelEventBase):
         ]
         return cls(
             id=job.id,
-            source=job.source,
+            source=str(job.source),
             local_path=job.local_path.as_posix(),
             parts=parts,
             bytes=job.bytes,
@@ -459,7 +458,7 @@ class ModelInstallDownloadProgressEvent(ModelEventBase):
     __event_name__ = "model_install_download_progress"
 
     id: int = Field(description="The ID of the install job")
-    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
+    source: str = Field(description="Source of the model; local path, repo_id or url")
     local_path: str = Field(description="Where model is downloading to")
     bytes: int = Field(description="Number of bytes downloaded so far")
     total_bytes: int = Field(description="Total size of download, including all files")
@@ -480,7 +479,7 @@ class ModelInstallDownloadProgressEvent(ModelEventBase):
         ]
         return cls(
             id=job.id,
-            source=job.source,
+            source=str(job.source),
             local_path=job.local_path.as_posix(),
             parts=parts,
             bytes=job.bytes,
@@ -495,11 +494,11 @@ class ModelInstallDownloadsCompleteEvent(ModelEventBase):
     __event_name__ = "model_install_downloads_complete"
 
     id: int = Field(description="The ID of the install job")
-    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
+    source: str = Field(description="Source of the model; local path, repo_id or url")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadsCompleteEvent":
-        return cls(id=job.id, source=job.source)
+        return cls(id=job.id, source=str(job.source))
 
 
 @payload_schema.register
@@ -509,11 +508,11 @@ class ModelInstallStartedEvent(ModelEventBase):
     __event_name__ = "model_install_started"
 
     id: int = Field(description="The ID of the install job")
-    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
+    source: str = Field(description="Source of the model; local path, repo_id or url")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallStartedEvent":
-        return cls(id=job.id, source=job.source)
+        return cls(id=job.id, source=str(job.source))
 
 
 @payload_schema.register
@@ -523,14 +522,14 @@ class ModelInstallCompleteEvent(ModelEventBase):
     __event_name__ = "model_install_complete"
 
     id: int = Field(description="The ID of the install job")
-    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
+    source: str = Field(description="Source of the model; local path, repo_id or url")
     key: str = Field(description="Model config record key")
     total_bytes: Optional[int] = Field(description="Size of the model (may be None for installation of a local path)")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallCompleteEvent":
         assert job.config_out is not None
-        return cls(id=job.id, source=job.source, key=(job.config_out.key), total_bytes=job.total_bytes)
+        return cls(id=job.id, source=str(job.source), key=(job.config_out.key), total_bytes=job.total_bytes)
 
 
 @payload_schema.register
@@ -540,11 +539,11 @@ class ModelInstallCancelledEvent(ModelEventBase):
     __event_name__ = "model_install_cancelled"
 
     id: int = Field(description="The ID of the install job")
-    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
+    source: str = Field(description="Source of the model; local path, repo_id or url")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallCancelledEvent":
-        return cls(id=job.id, source=job.source)
+        return cls(id=job.id, source=str(job.source))
 
 
 @payload_schema.register
@@ -554,7 +553,7 @@ class ModelInstallErrorEvent(ModelEventBase):
     __event_name__ = "model_install_error"
 
     id: int = Field(description="The ID of the install job")
-    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
+    source: str = Field(description="Source of the model; local path, repo_id or url")
     error_type: str = Field(description="The name of the exception")
     error: str = Field(description="A text description of the exception")
 
@@ -562,7 +561,7 @@ class ModelInstallErrorEvent(ModelEventBase):
     def build(cls, job: "ModelInstallJob") -> "ModelInstallErrorEvent":
         assert job.error_type is not None
         assert job.error is not None
-        return cls(id=job.id, source=job.source, error_type=job.error_type, error=job.error)
+        return cls(id=job.id, source=str(job.source), error_type=job.error_type, error=job.error)
 
 
 class BulkDownloadEventBase(EventBase):

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -20,7 +20,7 @@ from invokeai.app.services.invocation_stats.invocation_stats_common import (
     NodeExecutionStatsSummary,
 )
 from invokeai.app.services.invoker import Invoker
-from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
+from invokeai.backend.model_manager.load.model_cache import CacheStats
 
 # Size of 1GB in bytes.
 GB = 2**30

invokeai/app/services/model_install/model_install_base.py

@@ -3,20 +3,18 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import List, Optional, Union
 
 from pydantic.networks import AnyHttpUrl
 
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.download import DownloadQueueServiceBase
+from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
 from invokeai.app.services.model_records import ModelRecordChanges, ModelRecordServiceBase
 from invokeai.backend.model_manager import AnyModelConfig
 
-if TYPE_CHECKING:
-    from invokeai.app.services.events.events_base import EventServiceBase
-
 
 class ModelInstallServiceBase(ABC):
     """Abstract base class for InvokeAI model installation."""

invokeai/app/services/model_install/model_install_default.py

@@ -9,7 +9,7 @@ from pathlib import Path
 from queue import Empty, Queue
 from shutil import copyfile, copytree, move, rmtree
 from tempfile import mkdtemp
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import torch
 import yaml
@@ -20,6 +20,7 @@ from requests import Session
 
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.download import DownloadQueueServiceBase, MultiFileDownloadJob
+from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_install.model_install_base import ModelInstallServiceBase
 from invokeai.app.services.model_install.model_install_common import (
@@ -56,10 +57,6 @@ from invokeai.backend.util.catch_sigint import catch_sigint
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.util import slugify
 
-if TYPE_CHECKING:
-    from invokeai.app.services.events.events_base import EventServiceBase
-
-
 TMPDIR_PREFIX = "tmpinstall_"
 
 
@@ -441,10 +438,9 @@ class ModelInstallService(ModelInstallServiceBase):
         variants = "|".join(ModelRepoVariant.__members__.values())
         hf_repoid_re = f"^([^/:]+/[^/:]+)(?::({variants})?(?::/?([^:]+))?)?$"
         source_obj: Optional[StringLikeSource] = None
-        source_stripped = source.strip('"')
 
-        if Path(source_stripped).exists():  # A local file or directory
-            source_obj = LocalModelSource(path=Path(source_stripped))
+        if Path(source).exists():  # A local file or directory
+            source_obj = LocalModelSource(path=Path(source))
         elif match := re.match(hf_repoid_re, source):
             source_obj = HFModelSource(
                 repo_id=match.group(1),

invokeai/app/services/model_load/model_load_base.py

@@ -7,7 +7,7 @@ from typing import Callable, Optional
 
 from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType
 from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig
-from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
 
 
 class ModelLoadServiceBase(ABC):
@@ -24,7 +24,7 @@ class ModelLoadServiceBase(ABC):
 
     @property
     @abstractmethod
-    def ram_cache(self) -> ModelCache:
+    def ram_cache(self) -> ModelCacheBase[AnyModel]:
         """Return the RAM cache used by this loader."""
 
     @abstractmethod

invokeai/app/services/model_load/model_load_default.py

@@ -18,7 +18,7 @@ from invokeai.backend.model_manager.load import (
     ModelLoaderRegistry,
     ModelLoaderRegistryBase,
 )
-from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
 from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
@@ -30,7 +30,7 @@ class ModelLoadService(ModelLoadServiceBase):
     def __init__(
         self,
         app_config: InvokeAIAppConfig,
-        ram_cache: ModelCache,
+        ram_cache: ModelCacheBase[AnyModel],
         registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry,
     ):
         """Initialize the model load service."""
@@ -45,7 +45,7 @@ class ModelLoadService(ModelLoadServiceBase):
         self._invoker = invoker
 
     @property
-    def ram_cache(self) -> ModelCache:
+    def ram_cache(self) -> ModelCacheBase[AnyModel]:
         """Return the RAM cache used by this loader."""
         return self._ram_cache
 
@@ -78,8 +78,9 @@ class ModelLoadService(ModelLoadServiceBase):
         self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
     ) -> LoadedModelWithoutConfig:
         cache_key = str(model_path)
+        ram_cache = self.ram_cache
         try:
-            return LoadedModelWithoutConfig(cache_record=self._ram_cache.get(key=cache_key), cache=self._ram_cache)
+            return LoadedModelWithoutConfig(_locker=ram_cache.get(key=cache_key))
         except IndexError:
             pass
 
@@ -108,5 +109,5 @@ class ModelLoadService(ModelLoadServiceBase):
         )
         assert loader is not None
         raw_model = loader(model_path)
-        self._ram_cache.put(key=cache_key, model=raw_model)
-        return LoadedModelWithoutConfig(cache_record=self._ram_cache.get(key=cache_key), cache=self._ram_cache)
+        ram_cache.put(key=cache_key, model=raw_model)
+        return LoadedModelWithoutConfig(_locker=ram_cache.get(key=cache_key))

invokeai/app/services/model_manager/model_manager_default.py

@@ -16,8 +16,7 @@ from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBas
 from invokeai.app.services.model_load.model_load_default import ModelLoadService
 from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
 from invokeai.app.services.model_records.model_records_base import ModelRecordServiceBase
-from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
-from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.load import ModelCache, ModelLoaderRegistry
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 
@@ -82,12 +81,11 @@ class ModelManagerService(ModelManagerServiceBase):
         logger.setLevel(app_config.log_level.upper())
 
         ram_cache = ModelCache(
-            execution_device_working_mem_gb=app_config.device_working_mem_gb,
-            enable_partial_loading=app_config.enable_partial_loading,
-            max_ram_cache_size_gb=app_config.ram,
-            max_vram_cache_size_gb=app_config.vram,
-            execution_device=execution_device or TorchDevice.choose_torch_device(),
+            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(),
         )
         loader = ModelLoadService(
             app_config=app_config,

invokeai/app/services/session_processor/session_processor_default.py

@@ -439,9 +439,7 @@ class DefaultSessionProcessor(SessionProcessorBase):
                         poll_now_event.wait(self._polling_interval)
                         continue
 
-                    self._invoker.services.logger.info(
-                        f"Executing queue item {self._queue_item.item_id}, session {self._queue_item.session_id}"
-                    )
+                    self._invoker.services.logger.debug(f"Executing queue item {self._queue_item.item_id}")
                     cancel_event.clear()
 
                     # Run the graph

invokeai/backend/flux/denoise.py

@@ -31,7 +31,7 @@ def denoise(
     pos_ip_adapter_extensions: list[XLabsIPAdapterExtension],
     neg_ip_adapter_extensions: list[XLabsIPAdapterExtension],
     # extra img tokens
-    img_cond: torch.Tensor | None,
+    img_cond: torch.Tensor | None = None,
 ):
     # step 0 is the initial state
     total_steps = len(timesteps) - 1

invokeai/backend/flux/flux_tools_sampling_utils.py

@@ -0,0 +1,27 @@
+import torch
+import numpy as np
+from PIL import Image
+from einops import rearrange
+
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+
+def prepare_control(
+    height: int,
+    width: int,
+    seed: int,
+    ae: AutoEncoder,
+    cond_image: Image.Image,
+) -> torch.Tensor:
+    # load and encode the conditioning image
+    img_cond = cond_image.convert("RGB")
+    img_cond = img_cond.resize((width, height), Image.Resampling.LANCZOS)
+    img_cond = np.array(img_cond)
+    img_cond = torch.from_numpy(img_cond).float()
+    img_cond = rearrange(img_cond, "h w c -> 1 c h w")
+    ae_dtype = next(iter(ae.parameters())).dtype
+    ae_device = next(iter(ae.parameters())).device
+    img_cond = img_cond.to(device=ae_device, dtype=ae_dtype)
+    generator = torch.Generator(device=ae_device).manual_seed(seed)
+    img_cond = ae.encode(img_cond, sample=True, generator=generator)
+    img_cond = rearrange(img_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    return img_cond

invokeai/backend/flux/model.py

@@ -1,10 +1,10 @@
 # Initially pulled from https://github.com/black-forest-labs/flux
 
 from dataclasses import dataclass
-from typing import Optional
 
 import torch
 from torch import Tensor, nn
+from typing import Optional
 
 from invokeai.backend.flux.custom_block_processor import (
     CustomDoubleStreamBlockProcessor,

invokeai/backend/flux/modules/conditioner.py

@@ -3,8 +3,6 @@
 from torch import Tensor, nn
 from transformers import PreTrainedModel, PreTrainedTokenizer
 
-from invokeai.backend.util.devices import TorchDevice
-
 
 class HFEncoder(nn.Module):
     def __init__(self, encoder: PreTrainedModel, tokenizer: PreTrainedTokenizer, is_clip: bool, max_length: int):
@@ -28,7 +26,7 @@ class HFEncoder(nn.Module):
         )
 
         outputs = self.hf_module(
-            input_ids=batch_encoding["input_ids"].to(TorchDevice.choose_torch_device()),
+            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
             attention_mask=None,
             output_hidden_states=False,
         )

invokeai/backend/flux/modules/image_embedders.py

@@ -0,0 +1,50 @@
+import os
+import cv2
+import numpy as np
+import torch
+
+from einops import rearrange, repeat
+from PIL import Image
+from safetensors.torch import load_file as load_sft
+from torch import nn
+from transformers import AutoModelForDepthEstimation, AutoProcessor, SiglipImageProcessor, SiglipVisionModel
+
+class DepthImageEncoder:
+    depth_model_name = "LiheYoung/depth-anything-large-hf"
+    def __init__(self, device):
+        self.device = device
+        self.depth_model = AutoModelForDepthEstimation.from_pretrained(self.depth_model_name).to(device)
+        self.processor = AutoProcessor.from_pretrained(self.depth_model_name)
+    def __call__(self, img: torch.Tensor) -> torch.Tensor:
+        hw = img.shape[-2:]
+        img = torch.clamp(img, -1.0, 1.0)
+        img_byte = ((img + 1.0) * 127.5).byte()
+        img = self.processor(img_byte, return_tensors="pt")["pixel_values"]
+        depth = self.depth_model(img.to(self.device)).predicted_depth
+        depth = repeat(depth, "b h w -> b 3 h w")
+        depth = torch.nn.functional.interpolate(depth, hw, mode="bicubic", antialias=True)
+        depth = depth / 127.5 - 1.0
+        return depth
+
+class CannyImageEncoder:
+    def __init__(
+        self,
+        device,
+        min_t: int = 50,
+        max_t: int = 200,
+    ):
+        self.device = device
+        self.min_t = min_t
+        self.max_t = max_t
+    def __call__(self, img: torch.Tensor) -> torch.Tensor:
+        assert img.shape[0] == 1, "Only batch size 1 is supported"
+        img = rearrange(img[0], "c h w -> h w c")
+        img = torch.clamp(img, -1.0, 1.0)
+        img_np = ((img + 1.0) * 127.5).numpy().astype(np.uint8)
+        # Apply Canny edge detection
+        canny = cv2.Canny(img_np, self.min_t, self.max_t)
+        # Convert back to torch tensor and reshape
+        canny = torch.from_numpy(canny).float() / 127.5 - 1.0
+        canny = rearrange(canny, "h w -> 1 1 h w")
+        canny = repeat(canny, "b 1 ... -> b 3 ...")
+        return canny.to(self.device)

invokeai/backend/image_util/hed.py

@@ -18,7 +18,6 @@ from invokeai.backend.image_util.util import (
     resize_image_to_resolution,
     safe_step,
 )
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 
 
 class DoubleConvBlock(torch.nn.Module):
@@ -110,7 +109,7 @@ class HEDProcessor:
         Returns:
             The detected edges.
         """
-        device = get_effective_device(self.network)
+        device = next(iter(self.network.parameters())).device
         np_image = pil_to_np(input_image)
         np_image = normalize_image_channel_count(np_image)
         np_image = resize_image_to_resolution(np_image, detect_resolution)
@@ -184,7 +183,7 @@ class HEDEdgeDetector:
             The detected edges.
         """
 
-        device = get_effective_device(self.model)
+        device = next(iter(self.model.parameters())).device
 
         np_image = pil_to_np(image)
 

invokeai/backend/image_util/infill_methods/lama.py

@@ -7,7 +7,6 @@ from PIL import Image
 
 import invokeai.backend.util.logging as logger
 from invokeai.backend.model_manager.config import AnyModel
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 
 
 def norm_img(np_img):
@@ -32,7 +31,7 @@ class LaMA:
         mask = norm_img(mask)
         mask = (mask > 0) * 1
 
-        device = get_effective_device(self._model)
+        device = next(self._model.buffers()).device
         image = torch.from_numpy(image).unsqueeze(0).to(device)
         mask = torch.from_numpy(mask).unsqueeze(0).to(device)
 

invokeai/backend/image_util/lineart.py

@@ -17,7 +17,6 @@ from invokeai.backend.image_util.util import (
     pil_to_np,
     resize_image_to_resolution,
 )
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 
 
 class ResidualBlock(nn.Module):
@@ -131,7 +130,7 @@ class LineartProcessor:
         Returns:
             The detected lineart.
         """
-        device = get_effective_device(self.model)
+        device = next(iter(self.model.parameters())).device
 
         np_image = pil_to_np(input_image)
         np_image = normalize_image_channel_count(np_image)
@@ -202,7 +201,7 @@ class LineartEdgeDetector:
         Returns:
             The detected edges.
         """
-        device = get_effective_device(self.model)
+        device = next(iter(self.model.parameters())).device
 
         np_image = pil_to_np(image)
 

invokeai/backend/image_util/lineart_anime.py

@@ -19,7 +19,6 @@ from invokeai.backend.image_util.util import (
     pil_to_np,
     resize_image_to_resolution,
 )
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 
 
 class UnetGenerator(nn.Module):
@@ -172,7 +171,7 @@ class LineartAnimeProcessor:
         Returns:
             The detected lineart.
         """
-        device = get_effective_device(self.model)
+        device = next(iter(self.model.parameters())).device
         np_image = pil_to_np(input_image)
 
         np_image = normalize_image_channel_count(np_image)
@@ -240,7 +239,7 @@ class LineartAnimeEdgeDetector:
 
     def run(self, image: Image.Image) -> Image.Image:
         """Processes an image and returns the detected edges."""
-        device = get_effective_device(self.model)
+        device = next(iter(self.model.parameters())).device
 
         np_image = pil_to_np(image)
 

invokeai/backend/image_util/mlsd/utils.py

@@ -14,8 +14,6 @@ import numpy as np
 import torch
 from torch.nn import functional as F
 
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
-
 
 def deccode_output_score_and_ptss(tpMap, topk_n = 200, ksize = 5):
     '''
@@ -51,7 +49,7 @@ def pred_lines(image, model,
                dist_thr=20.0):
     h, w, _ = image.shape
 
-    device = get_effective_device(model)
+    device = next(iter(model.parameters())).device
     h_ratio, w_ratio = [h / input_shape[0], w / input_shape[1]]
 
     resized_image = np.concatenate([cv2.resize(image, (input_shape[1], input_shape[0]), interpolation=cv2.INTER_AREA),
@@ -110,7 +108,7 @@ def pred_squares(image,
     '''
     h, w, _ = image.shape
     original_shape = [h, w]
-    device = get_effective_device(model)
+    device = next(iter(model.parameters())).device
 
     resized_image = np.concatenate([cv2.resize(image, (input_shape[0], input_shape[1]), interpolation=cv2.INTER_AREA),
                                     np.ones([input_shape[0], input_shape[1], 1])], axis=-1)

invokeai/backend/image_util/normal_bae/__init__.py

@@ -13,7 +13,6 @@ from PIL import Image
 
 from invokeai.backend.image_util.normal_bae.nets.NNET import NNET
 from invokeai.backend.image_util.util import np_to_pil, pil_to_np, resize_to_multiple
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 
 
 class NormalMapDetector:
@@ -65,7 +64,7 @@ class NormalMapDetector:
     def run(self, image: Image.Image):
         """Processes an image and returns the detected normal map."""
 
-        device = get_effective_device(self.model)
+        device = next(iter(self.model.parameters())).device
         np_image = pil_to_np(image)
 
         height, width, _channels = np_image.shape

invokeai/backend/image_util/pidi/__init__.py

@@ -11,7 +11,6 @@ from PIL import Image
 
 from invokeai.backend.image_util.pidi.model import PiDiNet, pidinet
 from invokeai.backend.image_util.util import nms, normalize_image_channel_count, np_to_pil, pil_to_np, safe_step
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 
 
 class PIDINetDetector:
@@ -46,7 +45,7 @@ class PIDINetDetector:
     ) -> Image.Image:
         """Processes an image and returns the detected edges."""
 
-        device = get_effective_device(self.model)
+        device = next(iter(self.model.parameters())).device
 
         np_img = pil_to_np(image)
         np_img = normalize_image_channel_count(np_img)

invokeai/backend/lora/conversions/flux_control_lora_utils.py

@@ -0,0 +1,65 @@
+import re
+import torch
+
+from typing import Any, Dict
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.layers.set_parameter_layer import SetParameterLayer
+
+
+# A regex pattern that matches all of the keys in the Flux Dev/Canny LoRA format.
+# Example keys:
+#   guidance_in.in_layer.lora_B.bias
+#   single_blocks.0.linear1.lora_A.weight
+#   double_blocks.0.img_attn.norm.key_norm.scale
+FLUX_STRUCTURAL_TRANSFORMER_KEY_REGEX = r"(final_layer|vector_in|txt_in|time_in|img_in|guidance_in|\w+_blocks)(\.(\d+))?\.(lora_(A|B)|(in|out)_layer|adaLN_modulation|img_attn|img_mlp|img_mod|txt_attn|txt_mlp|txt_mod|linear|linear1|linear2|modulation|norm)\.?(.*)"
+
+def is_state_dict_likely_flux_control(state_dict: Dict[str, Any]) -> bool:
+    """Checks if the provided state dict is likely in the FLUX Control LoRA format.
+
+    This is intended to be a high-precision detector, but it is not guaranteed to have perfect precision. (A
+    perfect-precision detector would require checking all keys against a whitelist and verifying tensor shapes.)
+    """
+    return all(
+        re.match(FLUX_STRUCTURAL_TRANSFORMER_KEY_REGEX, k) or re.match(FLUX_STRUCTURAL_TRANSFORMER_KEY_REGEX, k)
+        for k in state_dict.keys()
+    )
+
+def lora_model_from_flux_control_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRAModelRaw:
+    # converted_state_dict = _convert_lora_bfl_control(state_dict=state_dict)
+    # Group keys by layer.
+    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {}
+    for key, value in state_dict.items():
+        key_props = key.split(".")
+        # Got it loading using lora_down and lora_up but it didn't seem to match this lora's structure
+        # Leaving this in since it doesn't hurt anything and may be better
+        layer_prop_size = -2 if any(prop in key for prop in ["lora_B", "lora_A"]) else -1
+        layer_name = ".".join(key_props[:layer_prop_size])
+        param_name = ".".join(key_props[layer_prop_size:])
+        if layer_name not in grouped_state_dict:
+            grouped_state_dict[layer_name] = {}
+        grouped_state_dict[layer_name][param_name] = value
+
+    # Create LoRA layers.
+    layers: dict[str, AnyLoRALayer] = {}
+    for layer_key, layer_state_dict in grouped_state_dict.items():
+        # Convert to a full layer diff
+        prefixed_key = f"{FLUX_LORA_TRANSFORMER_PREFIX}{layer_key}"
+        if all(k in layer_state_dict for k in ["lora_A.weight", "lora_B.bias", "lora_B.weight"]):
+            layers[prefixed_key] = LoRALayer(
+                layer_state_dict["lora_B.weight"],
+                None,
+                layer_state_dict["lora_A.weight"],
+                None,
+                layer_state_dict["lora_B.bias"]
+            )
+        elif "scale" in layer_state_dict:
+            layers[prefixed_key] = SetParameterLayer("scale", layer_state_dict["scale"])
+        else:
+            raise AssertionError(f"{layer_key} not expected")
+    # Create and return the LoRAModelRaw.
+    return LoRAModelRaw(layers=layers)
+

invokeai/backend/lora/conversions/flux_diffusers_lora_conversion_utils.py

@@ -2,11 +2,11 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
-from invokeai.backend.patches.layers.concatenated_lora_layer import ConcatenatedLoRALayer
-from invokeai.backend.patches.layers.lora_layer import LoRALayer
-from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.lora.conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
 
 
 def is_state_dict_likely_in_flux_diffusers_format(state_dict: Dict[str, torch.Tensor]) -> bool:
@@ -30,9 +30,7 @@ def is_state_dict_likely_in_flux_diffusers_format(state_dict: Dict[str, torch.Te
     return all_keys_in_peft_format and all_expected_keys_present
 
 
-def lora_model_from_flux_diffusers_state_dict(
-    state_dict: Dict[str, torch.Tensor], alpha: float | None
-) -> ModelPatchRaw:
+def lora_model_from_flux_diffusers_state_dict(state_dict: Dict[str, torch.Tensor], alpha: float | None) -> LoRAModelRaw:
     """Loads a state dict in the Diffusers FLUX LoRA format into a LoRAModelRaw object.
 
     This function is based on:
@@ -51,7 +49,7 @@ def lora_model_from_flux_diffusers_state_dict(
     mlp_ratio = 4.0
     mlp_hidden_dim = int(hidden_size * mlp_ratio)
 
-    layers: dict[str, BaseLayerPatch] = {}
+    layers: dict[str, AnyLoRALayer] = {}
 
     def add_lora_layer_if_present(src_key: str, dst_key: str) -> None:
         if src_key in grouped_state_dict:
@@ -217,7 +215,7 @@ def lora_model_from_flux_diffusers_state_dict(
 
     layers_with_prefix = {f"{FLUX_LORA_TRANSFORMER_PREFIX}{k}": v for k, v in layers.items()}
 
-    return ModelPatchRaw(layers=layers_with_prefix)
+    return LoRAModelRaw(layers=layers_with_prefix)
 
 
 def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:

invokeai/backend/lora/conversions/flux_kohya_lora_conversion_utils.py

@@ -3,13 +3,10 @@ from typing import Any, Dict, TypeVar
 
 import torch
 
-from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
-from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict
-from invokeai.backend.patches.lora_conversions.flux_lora_constants import (
-    FLUX_LORA_CLIP_PREFIX,
-    FLUX_LORA_TRANSFORMER_PREFIX,
-)
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.lora.conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX, FLUX_LORA_TRANSFORMER_PREFIX
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
 
 # A regex pattern that matches all of the transformer keys in the Kohya FLUX LoRA format.
 # Example keys:
@@ -39,7 +36,7 @@ def is_state_dict_likely_in_flux_kohya_format(state_dict: Dict[str, Any]) -> boo
     )
 
 
-def lora_model_from_flux_kohya_state_dict(state_dict: Dict[str, torch.Tensor]) -> ModelPatchRaw:
+def lora_model_from_flux_kohya_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRAModelRaw:
     # Group keys by layer.
     grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {}
     for key, value in state_dict.items():
@@ -64,14 +61,14 @@ def lora_model_from_flux_kohya_state_dict(state_dict: Dict[str, torch.Tensor]) -
     clip_grouped_sd = _convert_flux_clip_kohya_state_dict_to_invoke_format(clip_grouped_sd)
 
     # Create LoRA layers.
-    layers: dict[str, BaseLayerPatch] = {}
+    layers: dict[str, AnyLoRALayer] = {}
     for layer_key, layer_state_dict in transformer_grouped_sd.items():
         layers[FLUX_LORA_TRANSFORMER_PREFIX + layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
     for layer_key, layer_state_dict in clip_grouped_sd.items():
         layers[FLUX_LORA_CLIP_PREFIX + layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
 
     # Create and return the LoRAModelRaw.
-    return ModelPatchRaw(layers=layers)
+    return LoRAModelRaw(layers=layers)
 
 
 T = TypeVar("T")

invokeai/backend/lora/conversions/sd_lora_conversion_utils.py

@@ -2,19 +2,19 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
-from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict
-from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
 
 
-def lora_model_from_sd_state_dict(state_dict: Dict[str, torch.Tensor]) -> ModelPatchRaw:
+def lora_model_from_sd_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRAModelRaw:
     grouped_state_dict: dict[str, dict[str, torch.Tensor]] = _group_state(state_dict)
 
-    layers: dict[str, BaseLayerPatch] = {}
+    layers: dict[str, AnyLoRALayer] = {}
     for layer_key, values in grouped_state_dict.items():
         layers[layer_key] = any_lora_layer_from_state_dict(values)
 
-    return ModelPatchRaw(layers=layers)
+    return LoRAModelRaw(layers=layers)
 
 
 def _group_state(state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]:

invokeai/backend/lora/layers/any_lora_layer.py

@@ -0,0 +1,12 @@
+from typing import Union
+
+from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
+from invokeai.backend.lora.layers.full_layer import FullLayer
+from invokeai.backend.lora.layers.ia3_layer import IA3Layer
+from invokeai.backend.lora.layers.loha_layer import LoHALayer
+from invokeai.backend.lora.layers.lokr_layer import LoKRLayer
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.layers.norm_layer import NormLayer
+from invokeai.backend.lora.layers.set_parameter_layer import SetParameterLayer
+
+AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer, NormLayer, ConcatenatedLoRALayer, SetParameterLayer]

invokeai/backend/lora/layers/concatenated_lora_layer.py

@@ -2,8 +2,8 @@ from typing import Optional, Sequence
 
 import torch
 
-from invokeai.backend.patches.layers.lora_layer import LoRALayer
-from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
 
 
 class ConcatenatedLoRALayer(LoRALayerBase):
@@ -20,7 +20,7 @@ class ConcatenatedLoRALayer(LoRALayerBase):
         self.lora_layers = lora_layers
         self.concat_axis = concat_axis
 
-    def _rank(self) -> int | None:
+    def rank(self) -> int | None:
         return None
 
     def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
@@ -30,7 +30,7 @@ class ConcatenatedLoRALayer(LoRALayerBase):
         layer_weights = [lora_layer.get_weight(None) * lora_layer.scale() for lora_layer in self.lora_layers]  # pyright: ignore[reportArgumentType]
         return torch.cat(layer_weights, dim=self.concat_axis)
 
-    def get_bias(self, orig_bias: torch.Tensor | None) -> Optional[torch.Tensor]:
+    def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
         # TODO(ryand): Currently, we pass orig_bias=None to the sub-layers. If we want to support sub-layers that
         # require this value, we will need to implement chunking of the original bias tensor here.
         # Note that we must apply the sub-layer scales here.

invokeai/backend/lora/layers/full_layer.py

@@ -2,7 +2,7 @@ from typing import Dict, Optional
 
 import torch
 
-from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
 from invokeai.backend.util.calc_tensor_size import calc_tensor_size
 
 
@@ -20,7 +20,7 @@ class FullLayer(LoRALayerBase):
         cls.warn_on_unhandled_keys(values=values, handled_keys={"diff", "diff_b"})
         return layer
 
-    def _rank(self) -> int | None:
+    def rank(self) -> int | None:
         return None
 
     def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:

invokeai/backend/lora/layers/ia3_layer.py

@@ -2,8 +2,7 @@ from typing import Dict, Optional
 
 import torch
 
-from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
-from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
 
 
 class IA3Layer(LoRALayerBase):
@@ -17,7 +16,7 @@ class IA3Layer(LoRALayerBase):
         self.weight = weight
         self.on_input = on_input
 
-    def _rank(self) -> int | None:
+    def rank(self) -> int | None:
         return None
 
     @classmethod
@@ -51,7 +50,7 @@ class IA3Layer(LoRALayerBase):
         weight = self.weight
         if not self.on_input:
             weight = weight.reshape(-1, 1)
-        return cast_to_device(orig_weight, weight.device) * weight
+        return orig_weight * weight
 
     def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
         super().to(device, dtype)

invokeai/backend/lora/layers/loha_layer.py

@@ -2,7 +2,7 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
 from invokeai.backend.util.calc_tensor_size import calc_tensors_size
 
 
@@ -32,7 +32,7 @@ class LoHALayer(LoRALayerBase):
         self.t2 = t2
         assert (self.t1 is None) == (self.t2 is None)
 
-    def _rank(self) -> int | None:
+    def rank(self) -> int | None:
         return self.w1_b.shape[0]
 
     @classmethod

invokeai/backend/lora/layers/lokr_layer.py

@@ -2,7 +2,7 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
 from invokeai.backend.util.calc_tensor_size import calc_tensors_size
 
 
@@ -39,7 +39,7 @@ class LoKRLayer(LoRALayerBase):
         assert (self.w2 is None) != (self.w2_a is None)
         assert (self.w2_a is None) == (self.w2_b is None)
 
-    def _rank(self) -> int | None:
+    def rank(self) -> int | None:
         if self.w1_b is not None:
             return self.w1_b.shape[0]
         elif self.w2_b is not None:

invokeai/backend/lora/layers/lora_layer.py

@@ -2,7 +2,7 @@ from typing import Dict, Optional
 
 import torch
 
-from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
 from invokeai.backend.util.calc_tensor_size import calc_tensors_size
 
 
@@ -55,7 +55,7 @@ class LoRALayer(LoRALayerBase):
 
         return layer
 
-    def _rank(self) -> int:
+    def rank(self) -> int:
         return self.down.shape[0]
 
     def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:

invokeai/backend/lora/layers/lora_layer_base.py

@@ -1,13 +1,12 @@
-from typing import Optional
+from typing import Dict, Optional, Set
 
 import torch
 
 import invokeai.backend.util.logging as logger
-from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
 from invokeai.backend.util.calc_tensor_size import calc_tensors_size
 
 
-class LoRALayerBase(BaseLayerPatch):
+class LoRALayerBase:
     """Base class for all LoRA-like patching layers."""
 
     # Note: It is tempting to make this a torch.nn.Module sub-class and make all tensors 'torch.nn.Parameter's. Then we
@@ -24,7 +23,6 @@ class LoRALayerBase(BaseLayerPatch):
     def _parse_bias(
         cls, bias_indices: torch.Tensor | None, bias_values: torch.Tensor | None, bias_size: torch.Tensor | None
     ) -> torch.Tensor | None:
-        """Helper function to parse a bias tensor from a state dict in LyCORIS format."""
         assert (bias_indices is None) == (bias_values is None) == (bias_size is None)
 
         bias = None
@@ -39,14 +37,11 @@ class LoRALayerBase(BaseLayerPatch):
     ) -> float | None:
         return alpha.item() if alpha is not None else None
 
-    def _rank(self) -> int | None:
-        """Return the rank of the LoRA-like layer. Or None if the layer does not have a rank. This value is used to
-        calculate the scale.
-        """
+    def rank(self) -> int | None:
         raise NotImplementedError()
 
     def scale(self) -> float:
-        rank = self._rank()
+        rank = self.rank()
         if self._alpha is None or rank is None:
             return 1.0
         return self._alpha / rank
@@ -54,26 +49,18 @@ class LoRALayerBase(BaseLayerPatch):
     def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError()
 
-    def get_bias(self, orig_bias: torch.Tensor | None) -> Optional[torch.Tensor]:
+    def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
         return self.bias
 
-    def get_parameters(self, orig_parameters: dict[str, torch.Tensor], weight: float) -> dict[str, torch.Tensor]:
-        scale = self.scale()
-        params = {"weight": self.get_weight(orig_parameters["weight"]) * (weight * scale)}
-        bias = self.get_bias(orig_parameters.get("bias", None))
+    def get_parameters(self, orig_module: torch.nn.Module) -> Dict[str, torch.Tensor]:
+        params = {"weight": self.get_weight(orig_module.weight)}
+        bias = self.get_bias(orig_module.bias)
         if bias is not None:
-            params["bias"] = bias * (weight * scale)
-
-        # Reshape all params to match the original module's shape.
-        for param_name, param_weight in params.items():
-            orig_param = orig_parameters[param_name]
-            if param_weight.shape != orig_param.shape:
-                params[param_name] = param_weight.reshape(orig_param.shape)
-
+            params["bias"] = bias
         return params
 
     @classmethod
-    def warn_on_unhandled_keys(cls, values: dict[str, torch.Tensor], handled_keys: set[str]):
+    def warn_on_unhandled_keys(cls, values: Dict[str, torch.Tensor], handled_keys: Set[str]):
         """Log a warning if values contains unhandled keys."""
         unknown_keys = set(values.keys()) - handled_keys
         if unknown_keys:

invokeai/backend/lora/layers/norm_layer.py

@@ -2,7 +2,7 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
 from invokeai.backend.util.calc_tensor_size import calc_tensor_size
 
 
@@ -20,7 +20,7 @@ class NormLayer(LoRALayerBase):
         cls.warn_on_unhandled_keys(values, {"w_norm", "b_norm"})
         return layer
 
-    def _rank(self) -> int | None:
+    def rank(self) -> int | None:
         return None
 
     def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:

invokeai/backend/lora/layers/reshape_weight_layer.py

@@ -0,0 +1,34 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+
+
+class ReshapeWeightLayer(LoRALayerBase):
+    # TODO: Just everything in this class 
+    def __init__(self, weight: Optional[torch.Tensor], bias: Optional[torch.Tensor], scale: Optional[torch.Tensor]):
+        super().__init__(alpha=None, bias=bias)
+        self.weight = torch.nn.Parameter(weight) if weight is not None else None
+        self.bias = torch.nn.Parameter(bias) if bias is not None else None
+        self.manual_scale = scale
+
+    def scale(self):
+        return self.manual_scale.float() if self.manual_scale is not None else super().scale()
+
+    def rank(self) -> int | None:
+        return None
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        return orig_weight
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        super().to(device=device, dtype=dtype)
+        if self.weight is not None:
+            self.weight = self.weight.to(device=device, dtype=dtype)
+        if self.manual_scale is not None:
+            self.manual_scale = self.manual_scale.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        return super().calc_size() + calc_tensor_size(self.manual_scale)

invokeai/backend/lora/layers/set_parameter_layer.py

@@ -0,0 +1,29 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+
+
+class SetParameterLayer(LoRALayerBase):
+    def __init__(self, param_name: str, weight: torch.Tensor):
+        super().__init__(None, None)
+        self.weight = weight
+        self.param_name = param_name
+
+    def rank(self) -> int | None:
+        return None
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        return self.weight - orig_weight
+    
+    def get_parameters(self, orig_module: torch.nn.Module) -> Dict[str, torch.Tensor]:
+        return {self.param_name: self.get_weight(orig_module.get_parameter(self.param_name))}
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        super().to(device=device, dtype=dtype)
+        self.weight = self.weight.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        return super().calc_size() + calc_tensor_size(self.weight)

invokeai/backend/lora/layers/utils.py

@@ -2,16 +2,17 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
-from invokeai.backend.patches.layers.full_layer import FullLayer
-from invokeai.backend.patches.layers.ia3_layer import IA3Layer
-from invokeai.backend.patches.layers.loha_layer import LoHALayer
-from invokeai.backend.patches.layers.lokr_layer import LoKRLayer
-from invokeai.backend.patches.layers.lora_layer import LoRALayer
-from invokeai.backend.patches.layers.norm_layer import NormLayer
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.full_layer import FullLayer
+from invokeai.backend.lora.layers.ia3_layer import IA3Layer
+from invokeai.backend.lora.layers.loha_layer import LoHALayer
+from invokeai.backend.lora.layers.lokr_layer import LoKRLayer
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.layers.norm_layer import NormLayer
+from invokeai.backend.lora.layers.set_parameter_layer import SetParameterLayer
 
 
-def any_lora_layer_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> BaseLayerPatch:
+def any_lora_layer_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> AnyLoRALayer:
     # Detect layers according to LyCORIS detection logic(`weight_list_det`)
     # https://github.com/KohakuBlueleaf/LyCORIS/tree/8ad8000efb79e2b879054da8c9356e6143591bad/lycoris/modules
 

invokeai/backend/lora/lora_layer_wrappers.py

@@ -0,0 +1,133 @@
+import torch
+
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+
+
+class LoRASidecarWrapper(torch.nn.Module):
+    def __init__(self, orig_module: torch.nn.Module, lora_layers: list[AnyLoRALayer], lora_weights: list[float]):
+        super().__init__()
+        self._orig_module = orig_module
+        self._lora_layers = lora_layers
+        self._lora_weights = lora_weights
+
+    @property
+    def orig_module(self) -> torch.nn.Module:
+        return self._orig_module
+
+    def add_lora_layer(self, lora_layer: AnyLoRALayer, lora_weight: float):
+        self._lora_layers.append(lora_layer)
+        self._lora_weights.append(lora_weight)
+
+    @torch.no_grad()
+    def _get_lora_patched_parameters(
+        self, orig_params: dict[str, torch.Tensor], lora_layers: list[AnyLoRALayer], lora_weights: list[float]
+    ) -> dict[str, torch.Tensor]:
+        params: dict[str, torch.Tensor] = {}
+        for lora_layer, lora_weight in zip(lora_layers, lora_weights, strict=True):
+            layer_params = lora_layer.get_parameters(self._orig_module)
+            for param_name, param_weight in layer_params.items():
+                if orig_params[param_name].shape != param_weight.shape:
+                    param_weight = param_weight.reshape(orig_params[param_name].shape)
+
+                if param_name not in params:
+                    params[param_name] = param_weight * (lora_layer.scale() * lora_weight)
+                else:
+                    params[param_name] += param_weight * (lora_layer.scale() * lora_weight)
+
+        return params
+
+
+class LoRALinearWrapper(LoRASidecarWrapper):
+    def _lora_linear_forward(self, input: torch.Tensor, lora_layer: LoRALayer, lora_weight: float) -> torch.Tensor:
+        """An optimized implementation of the residual calculation for a Linear LoRALayer."""
+        x = torch.nn.functional.linear(input, lora_layer.down)
+        if lora_layer.mid is not None:
+            x = torch.nn.functional.linear(x, lora_layer.mid)
+        x = torch.nn.functional.linear(x, lora_layer.up, bias=lora_layer.bias)
+        x *= lora_weight * lora_layer.scale()
+        return x
+
+    def _concatenated_lora_forward(
+        self, input: torch.Tensor, concatenated_lora_layer: ConcatenatedLoRALayer, lora_weight: float
+    ) -> torch.Tensor:
+        """An optimized implementation of the residual calculation for a Linear ConcatenatedLoRALayer."""
+        x_chunks: list[torch.Tensor] = []
+        for lora_layer in concatenated_lora_layer.lora_layers:
+            x_chunk = torch.nn.functional.linear(input, lora_layer.down)
+            if lora_layer.mid is not None:
+                x_chunk = torch.nn.functional.linear(x_chunk, lora_layer.mid)
+            x_chunk = torch.nn.functional.linear(x_chunk, lora_layer.up, bias=lora_layer.bias)
+            x_chunk *= lora_weight * lora_layer.scale()
+            x_chunks.append(x_chunk)
+
+        # TODO(ryand): Generalize to support concat_axis != 0.
+        assert concatenated_lora_layer.concat_axis == 0
+        x = torch.cat(x_chunks, dim=-1)
+        return x
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        # Split the LoRA layers into those that have optimized implementations and those that don't.
+        optimized_layer_types = (LoRALayer, ConcatenatedLoRALayer)
+        optimized_layers = [
+            (layer, weight)
+            for layer, weight in zip(self._lora_layers, self._lora_weights, strict=True)
+            if isinstance(layer, optimized_layer_types)
+        ]
+        non_optimized_layers = [
+            (layer, weight)
+            for layer, weight in zip(self._lora_layers, self._lora_weights, strict=True)
+            if not isinstance(layer, optimized_layer_types)
+        ]
+
+        # First, calculate the residual for LoRA layers for which there is an optimized implementation.
+        residual = None
+        for lora_layer, lora_weight in optimized_layers:
+            if isinstance(lora_layer, LoRALayer):
+                added_residual = self._lora_linear_forward(input, lora_layer, lora_weight)
+            elif isinstance(lora_layer, ConcatenatedLoRALayer):
+                added_residual = self._concatenated_lora_forward(input, lora_layer, lora_weight)
+            else:
+                raise ValueError(f"Unsupported LoRA layer type: {type(lora_layer)}")
+
+            if residual is None:
+                residual = added_residual
+            else:
+                residual += added_residual
+
+        # Next, calculate the residuals for the LoRA layers for which there is no optimized implementation.
+        if non_optimized_layers:
+            unoptimized_layers, unoptimized_weights = zip(*non_optimized_layers, strict=True)
+            params = self._get_lora_patched_parameters(
+                orig_params={"weight": self._orig_module.weight, "bias": self._orig_module.bias},
+                lora_layers=unoptimized_layers,
+                lora_weights=unoptimized_weights,
+            )
+            added_residual = torch.nn.functional.linear(input, params["weight"], params.get("bias", None))
+            if residual is None:
+                residual = added_residual
+            else:
+                residual += added_residual
+
+        return self.orig_module(input) + residual
+
+
+class LoRAConv1dWrapper(LoRASidecarWrapper):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        params = self._get_lora_patched_parameters(
+            orig_params={"weight": self._orig_module.weight, "bias": self._orig_module.bias},
+            lora_layers=self._lora_layers,
+            lora_weights=self._lora_weights,
+        )
+        return self.orig_module(input) + torch.nn.functional.conv1d(input, params["weight"], params.get("bias", None))
+
+
+class LoRAConv2dWrapper(LoRASidecarWrapper):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        params = self._get_lora_patched_parameters(
+            orig_params={"weight": self._orig_module.weight, "bias": self._orig_module.bias},
+            lora_layers=self._lora_layers,
+            lora_weights=self._lora_weights,
+        )
+        return self.orig_module(input) + torch.nn.functional.conv2d(input, params["weight"], params.get("bias", None))

invokeai/backend/lora/lora_model_raw.py

@@ -3,17 +3,20 @@ from typing import Mapping, Optional
 
 import torch
 
-from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
 from invokeai.backend.raw_model import RawModel
 
 
-class ModelPatchRaw(RawModel):
-    def __init__(self, layers: Mapping[str, BaseLayerPatch]):
+class LoRAModelRaw(RawModel):  # (torch.nn.Module):
+    def __init__(self, layers: Mapping[str, AnyLoRALayer]):
         self.layers = layers
 
     def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        for layer in self.layers.values():
+        for _key, layer in self.layers.items():
             layer.to(device=device, dtype=dtype)
 
     def calc_size(self) -> int:
-        return sum(layer.calc_size() for layer in self.layers.values())
+        model_size = 0
+        for _, layer in self.layers.items():
+            model_size += layer.calc_size()
+        return model_size

invokeai/backend/lora/lora_patcher.py

@@ -0,0 +1,458 @@
+from contextlib import contextmanager
+from typing import Dict, Iterable, Optional, Tuple
+
+import torch
+
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.lora_layer_wrappers import (
+    LoRAConv1dWrapper,
+    LoRAConv2dWrapper,
+    LoRALinearWrapper,
+    LoRASidecarWrapper,
+)
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+
+
+class LoRAPatcher:
+    @staticmethod
+    @torch.no_grad()
+    @contextmanager
+    def apply_smart_lora_patches(
+        model: torch.nn.Module,
+        patches: Iterable[Tuple[LoRAModelRaw, float]],
+        prefix: str,
+        dtype: torch.dtype,
+        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
+    ):
+        """Apply 'smart' LoRA patching that chooses whether to use direct patching or a sidecar wrapper for each module."""
+
+        # original_weights are stored for unpatching layers that are directly patched.
+        original_weights = OriginalWeightsStorage(cached_weights)
+        # original_modules are stored for unpatching layers that are wrapped in a LoRASidecarWrapper.
+        original_modules: dict[str, torch.nn.Module] = {}
+        try:
+            for patch, patch_weight in patches:
+                LoRAPatcher._apply_smart_lora_patch(
+                    model=model,
+                    prefix=prefix,
+                    patch=patch,
+                    patch_weight=patch_weight,
+                    original_weights=original_weights,
+                    original_modules=original_modules,
+                    dtype=dtype,
+                )
+
+            yield
+        finally:
+            # Restore directly patched layers.
+            for param_key, weight in original_weights.get_changed_weights():
+                model.get_parameter(param_key).copy_(weight)
+
+            # Restore LoRASidecarWrapper modules.
+            # Note: This logic assumes no nested modules in original_modules.
+            for module_key, orig_module in original_modules.items():
+                module_parent_key, module_name = LoRAPatcher._split_parent_key(module_key)
+                parent_module = model.get_submodule(module_parent_key)
+                LoRAPatcher._set_submodule(parent_module, module_name, orig_module)
+
+    @staticmethod
+    @torch.no_grad()
+    def _apply_smart_lora_patch(
+        model: torch.nn.Module,
+        prefix: str,
+        patch: LoRAModelRaw,
+        patch_weight: float,
+        original_weights: OriginalWeightsStorage,
+        original_modules: dict[str, torch.nn.Module],
+        dtype: torch.dtype,
+    ):
+        """Apply a single LoRA patch to a model using the 'smart' patching strategy that chooses whether to use direct
+        patching or a sidecar wrapper for each module.
+        """
+        if patch_weight == 0:
+            return
+
+        # If the layer keys contain a dot, then they are not flattened, and can be directly used to access model
+        # submodules. If the layer keys do not contain a dot, then they are flattened, meaning that all '.' have been
+        # replaced with '_'. Non-flattened keys are preferred, because they allow submodules to be accessed directly
+        # without searching, but some legacy code still uses flattened keys.
+        layer_keys_are_flattened = "." not in next(iter(patch.layers.keys()))
+
+        prefix_len = len(prefix)
+
+        for layer_key, layer in patch.layers.items():
+            if not layer_key.startswith(prefix):
+                continue
+
+            module_key, module = LoRAPatcher._get_submodule(
+                model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
+            )
+
+            # Decide whether to use direct patching or a sidecar wrapper.
+            # Direct patching is preferred, because it results in better runtime speed.
+            # Reasons to use sidecar patching:
+            # - The module is already wrapped in a LoRASidecarWrapper.
+            # - The module is quantized.
+            # - The module is on the CPU (and we don't want to store a second full copy of the original weights on the
+            #   CPU, since this would double the RAM usage)
+            # NOTE: For now, we don't check if the layer is quantized here. We assume that this is checked in the caller
+            # and that the caller will use the 'apply_lora_wrapper_patches' method if the layer is quantized.
+            # TODO(ryand): Handle the case where we are running without a GPU. Should we set a config flag that allows
+            # forcing full patching even on the CPU?
+            if isinstance(module, LoRASidecarWrapper) or LoRAPatcher._is_any_part_of_layer_on_cpu(module):
+                LoRAPatcher._apply_lora_layer_wrapper_patch(
+                    model=model,
+                    module_to_patch=module,
+                    module_to_patch_key=module_key,
+                    patch=layer,
+                    patch_weight=patch_weight,
+                    original_modules=original_modules,
+                    dtype=dtype,
+                )
+            else:
+                LoRAPatcher._apply_lora_layer_patch(
+                    module_to_patch=module,
+                    module_to_patch_key=module_key,
+                    patch=layer,
+                    patch_weight=patch_weight,
+                    original_weights=original_weights,
+                )
+
+    @staticmethod
+    def _is_any_part_of_layer_on_cpu(layer: torch.nn.Module) -> bool:
+        return any(p.device.type == "cpu" for p in layer.parameters())
+
+    @staticmethod
+    @torch.no_grad()
+    @contextmanager
+    def apply_lora_patches(
+        model: torch.nn.Module,
+        patches: Iterable[Tuple[LoRAModelRaw, float]],
+        prefix: str,
+        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
+    ):
+        """Apply one or more LoRA patches to a model within a context manager.
+
+        Args:
+            model (torch.nn.Module): The model to patch.
+            patches (Iterable[Tuple[LoRAModelRaw, float]]): An iterator that returns tuples of LoRA patches and
+                associated weights. An iterator is used so that the LoRA patches do not need to be loaded into memory
+                all at once.
+            prefix (str): The keys in the patches will be filtered to only include weights with this prefix.
+            cached_weights (Optional[Dict[str, torch.Tensor]], optional): Read-only copy of the model's state dict in
+                CPU RAM, for efficient unpatching purposes.
+        """
+        original_weights = OriginalWeightsStorage(cached_weights)
+        try:
+            for patch, patch_weight in patches:
+                LoRAPatcher._apply_lora_patch(
+                    model=model,
+                    prefix=prefix,
+                    patch=patch,
+                    patch_weight=patch_weight,
+                    original_weights=original_weights,
+                )
+                del patch
+
+            yield
+        finally:
+            for param_key, weight in original_weights.get_changed_weights():
+                cur_param = model.get_parameter(param_key)
+                cur_param.data = weight.to(dtype=cur_param.dtype, device=cur_param.device, copy=True)
+
+    @staticmethod
+    @torch.no_grad()
+    def _apply_lora_patch(
+        model: torch.nn.Module,
+        prefix: str,
+        patch: LoRAModelRaw,
+        patch_weight: float,
+        original_weights: OriginalWeightsStorage,
+    ):
+        """Apply a single LoRA patch to a model.
+
+        Args:
+            model (torch.nn.Module): The model to patch.
+            prefix (str): A string prefix that precedes keys used in the LoRAs weight layers.
+            patch (LoRAModelRaw): The LoRA model to patch in.
+            patch_weight (float): The weight of the LoRA patch.
+            original_weights (OriginalWeightsStorage): Storage for the original weights of the model, for unpatching.
+        """
+        if patch_weight == 0:
+            return
+
+        # If the layer keys contain a dot, then they are not flattened, and can be directly used to access model
+        # submodules. If the layer keys do not contain a dot, then they are flattened, meaning that all '.' have been
+        # replaced with '_'. Non-flattened keys are preferred, because they allow submodules to be accessed directly
+        # without searching, but some legacy code still uses flattened keys.
+        layer_keys_are_flattened = "." not in next(iter(patch.layers.keys()))
+
+        prefix_len = len(prefix)
+
+        for layer_key, layer in patch.layers.items():
+            if not layer_key.startswith(prefix):
+                continue
+
+            module_key, module = LoRAPatcher._get_submodule(
+                model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
+            )
+
+            LoRAPatcher._apply_lora_layer_patch(
+                module_to_patch=module,
+                module_to_patch_key=module_key,
+                patch=layer,
+                patch_weight=patch_weight,
+                original_weights=original_weights,
+            )
+
+    @staticmethod
+    @torch.no_grad()
+    def _apply_lora_layer_patch(
+        module_to_patch: torch.nn.Module,
+        module_to_patch_key: str,
+        patch: AnyLoRALayer,
+        patch_weight: float,
+        original_weights: OriginalWeightsStorage,
+    ):
+        # All of the LoRA weight calculations will be done on the same device as the module weight.
+        # (Performance will be best if this is a CUDA device.)
+        first_param = next(module_to_patch.parameters())
+        device = first_param.device
+        dtype = first_param.dtype
+
+        layer_scale = patch.scale()
+
+        # We intentionally move to the target device first, then cast. Experimentally, this was found to
+        # be significantly faster for 16-bit CPU tensors being moved to a CUDA device than doing the
+        # same thing in a single call to '.to(...)'.
+        patch.to(device=device)
+        patch.to(dtype=torch.float32)
+
+        # TODO(ryand): Using torch.autocast(...) over explicit casting may offer a speed benefit on CUDA
+        # devices here. Experimentally, it was found to be very slow on CPU. More investigation needed.
+        for param_name, lora_param_weight in patch.get_parameters(module_to_patch).items():
+            param_key = module_to_patch_key + "." + param_name
+            module_param = module_to_patch.get_parameter(param_name)
+
+            # Save original weight
+            original_weights.save(param_key, module_param)
+
+            if module_param.shape != lora_param_weight.shape:
+                if module_param.nelement() == lora_param_weight.nelement():
+                    lora_param_weight = lora_param_weight.reshape(module_param.shape)
+                else:
+                    # This condition was added to handle layers in FLUX control LoRAs.
+                    # TODO(ryand): Move the weight update into the LoRA layer so that the LoRAPatcher doesn't need
+                    # to worry about this?
+                    expanded_weight = torch.zeros_like(
+                        lora_param_weight, dtype=module_param.dtype, device=module_param.device
+                    )
+                    slices = tuple(slice(0, dim) for dim in module_param.shape)
+                    expanded_weight[slices] = module_param
+                    setattr(
+                        module,
+                        param_name,
+                        torch.nn.Parameter(expanded_weight, requires_grad=module_param.requires_grad),
+                    )
+                    module_param = expanded_weight
+
+            lora_param_weight *= patch_weight * layer_scale
+            module_param += lora_param_weight.to(dtype=dtype)
+
+        patch.to(device=TorchDevice.CPU_DEVICE)
+
+    @staticmethod
+    @torch.no_grad()
+    @contextmanager
+    def apply_lora_wrapper_patches(
+        model: torch.nn.Module,
+        patches: Iterable[Tuple[LoRAModelRaw, float]],
+        prefix: str,
+        dtype: torch.dtype,
+    ):
+        """Apply one or more LoRA wrapper patches to a model within a context manager. Wrapper patches incur some
+        runtime overhead compared to normal LoRA patching, but they enable:
+        - LoRA layers to be applied to quantized models
+        - LoRA layers to be applied to CPU layers without needing to store a full copy of the original weights (i.e.
+          avoid doubling the memory requirements).
+
+        Args:
+            model (torch.nn.Module): The model to patch.
+            patches (Iterable[Tuple[LoRAModelRaw, float]]): An iterator that returns tuples of LoRA patches and
+                associated weights. An iterator is used so that the LoRA patches do not need to be loaded into memory
+                all at once.
+            prefix (str): The keys in the patches will be filtered to only include weights with this prefix.
+        """
+        original_modules: dict[str, torch.nn.Module] = {}
+        try:
+            for patch, patch_weight in patches:
+                LoRAPatcher._apply_lora_wrapper_patch(
+                    model=model,
+                    prefix=prefix,
+                    patch=patch,
+                    patch_weight=patch_weight,
+                    original_modules=original_modules,
+                    dtype=dtype,
+                )
+            yield
+        finally:
+            # Restore original modules.
+            # Note: This logic assumes no nested modules in original_modules.
+            for module_key, orig_module in original_modules.items():
+                module_parent_key, module_name = LoRAPatcher._split_parent_key(module_key)
+                parent_module = model.get_submodule(module_parent_key)
+                LoRAPatcher._set_submodule(parent_module, module_name, orig_module)
+
+    @staticmethod
+    def _apply_lora_wrapper_patch(
+        model: torch.nn.Module,
+        patch: LoRAModelRaw,
+        patch_weight: float,
+        prefix: str,
+        original_modules: dict[str, torch.nn.Module],
+        dtype: torch.dtype,
+    ):
+        """Apply a single LoRA wrapper patch to a model."""
+
+        if patch_weight == 0:
+            return
+
+        # If the layer keys contain a dot, then they are not flattened, and can be directly used to access model
+        # submodules. If the layer keys do not contain a dot, then they are flattened, meaning that all '.' have been
+        # replaced with '_'. Non-flattened keys are preferred, because they allow submodules to be accessed directly
+        # without searching, but some legacy code still uses flattened keys.
+        layer_keys_are_flattened = "." not in next(iter(patch.layers.keys()))
+
+        prefix_len = len(prefix)
+
+        for layer_key, layer in patch.layers.items():
+            if not layer_key.startswith(prefix):
+                continue
+
+            module_key, module = LoRAPatcher._get_submodule(
+                model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
+            )
+
+            LoRAPatcher._apply_lora_layer_wrapper_patch(
+                model=model,
+                module_to_patch=module,
+                module_to_patch_key=module_key,
+                patch=layer,
+                patch_weight=patch_weight,
+                original_modules=original_modules,
+                dtype=dtype,
+            )
+
+    @staticmethod
+    @torch.no_grad()
+    def _apply_lora_layer_wrapper_patch(
+        model: torch.nn.Module,
+        module_to_patch: torch.nn.Module,
+        module_to_patch_key: str,
+        patch: AnyLoRALayer,
+        patch_weight: float,
+        original_modules: dict[str, torch.nn.Module],
+        dtype: torch.dtype,
+    ):
+        """Apply a single LoRA wrapper patch to a model."""
+
+        # Replace the original module with a LoRASidecarWrapper if it has not already been done.
+        if not isinstance(module_to_patch, LoRASidecarWrapper):
+            lora_wrapper_layer = LoRAPatcher._initialize_lora_wrapper_layer(module_to_patch)
+            original_modules[module_to_patch_key] = module_to_patch
+            module_parent_key, module_name = LoRAPatcher._split_parent_key(module_to_patch_key)
+            module_parent = model.get_submodule(module_parent_key)
+            LoRAPatcher._set_submodule(module_parent, module_name, lora_wrapper_layer)
+            orig_module = module_to_patch
+        else:
+            assert module_to_patch_key in original_modules
+            lora_wrapper_layer = module_to_patch
+            orig_module = module_to_patch.orig_module
+
+        # Move the LoRA layer to the same device/dtype as the orig module.
+        patch.to(device=orig_module.weight.device, dtype=dtype)
+
+        # Add the LoRA wrapper layer to the LoRASidecarWrapper.
+        lora_wrapper_layer.add_lora_layer(patch, patch_weight)
+
+    @staticmethod
+    def _split_parent_key(module_key: str) -> tuple[str, str]:
+        """Split a module key into its parent key and module name.
+
+        Args:
+            module_key (str): The module key to split.
+
+        Returns:
+            tuple[str, str]: A tuple containing the parent key and module name.
+        """
+        split_key = module_key.rsplit(".", 1)
+        if len(split_key) == 2:
+            return tuple(split_key)
+        elif len(split_key) == 1:
+            return "", split_key[0]
+        else:
+            raise ValueError(f"Invalid module key: {module_key}")
+
+    @staticmethod
+    def _initialize_lora_wrapper_layer(orig_layer: torch.nn.Module):
+        if isinstance(orig_layer, torch.nn.Linear):
+            return LoRALinearWrapper(orig_layer, [], [])
+        elif isinstance(orig_layer, torch.nn.Conv1d):
+            return LoRAConv1dWrapper(orig_layer, [], [])
+        elif isinstance(orig_layer, torch.nn.Conv2d):
+            return LoRAConv2dWrapper(orig_layer, [], [])
+        else:
+            raise ValueError(f"Unsupported layer type: {type(orig_layer)}")
+
+    @staticmethod
+    def _set_submodule(parent_module: torch.nn.Module, module_name: str, submodule: torch.nn.Module):
+        try:
+            submodule_index = int(module_name)
+            # If the module name is an integer, then we use the __setitem__ method to set the submodule.
+            parent_module[submodule_index] = submodule  # type: ignore
+        except ValueError:
+            # If the module name is not an integer, then we use the setattr method to set the submodule.
+            setattr(parent_module, module_name, submodule)
+
+    @staticmethod
+    def _get_submodule(
+        model: torch.nn.Module, layer_key: str, layer_key_is_flattened: bool
+    ) -> tuple[str, torch.nn.Module]:
+        """Get the submodule corresponding to the given layer key.
+
+        Args:
+            model (torch.nn.Module): The model to search.
+            layer_key (str): The layer key to search for.
+            layer_key_is_flattened (bool): Whether the layer key is flattened. If flattened, then all '.' have been
+                replaced with '_'. Non-flattened keys are preferred, because they allow submodules to be accessed
+                directly without searching, but some legacy code still uses flattened keys.
+
+        Returns:
+            tuple[str, torch.nn.Module]: A tuple containing the module key and the submodule.
+        """
+        if not layer_key_is_flattened:
+            return layer_key, model.get_submodule(layer_key)
+
+        # Handle flattened keys.
+        assert "." not in layer_key
+
+        module = model
+        module_key = ""
+        key_parts = layer_key.split("_")
+
+        submodule_name = key_parts.pop(0)
+
+        while len(key_parts) > 0:
+            try:
+                module = module.get_submodule(submodule_name)
+                module_key += "." + submodule_name
+                submodule_name = key_parts.pop(0)
+            except Exception:
+                submodule_name += "_" + key_parts.pop(0)
+
+        module = module.get_submodule(submodule_name)
+        module_key = (module_key + "." + submodule_name).lstrip(".")
+
+        return module_key, module

invokeai/backend/model_manager/config.py

@@ -67,7 +67,7 @@ class ModelType(str, Enum):
     Main = "main"
     VAE = "vae"
     LoRA = "lora"
-    ControlLoRa = "control_lora"
+    StructuralLoRa = "structural_lora"
     ControlNet = "controlnet"  # used by model_probe
     TextualInversion = "embedding"
     IPAdapter = "ip_adapter"
@@ -274,34 +274,16 @@ class LoRALyCORISConfig(LoRAConfigBase):
         return Tag(f"{ModelType.LoRA.value}.{ModelFormat.LyCORIS.value}")
 
 
-class ControlAdapterConfigBase(BaseModel):
-    default_settings: Optional[ControlAdapterDefaultSettings] = Field(
-        description="Default settings for this model", default=None
-    )
+class StructuralLoRALyCORISConfig(ModelConfigBase):
+    """Model config for Structural LoRA/Lycoris models."""
 
-
-class ControlLoRALyCORISConfig(ModelConfigBase, ControlAdapterConfigBase):
-    """Model config for Control LoRA models."""
-
-    type: Literal[ModelType.ControlLoRa] = ModelType.ControlLoRa
+    type: Literal[ModelType.StructuralLoRa] = ModelType.StructuralLoRa
     trigger_phrases: Optional[set[str]] = Field(description="Set of trigger phrases for this model", default=None)
     format: Literal[ModelFormat.LyCORIS] = ModelFormat.LyCORIS
 
     @staticmethod
     def get_tag() -> Tag:
-        return Tag(f"{ModelType.ControlLoRa.value}.{ModelFormat.LyCORIS.value}")
-
-
-class ControlLoRADiffusersConfig(ModelConfigBase, ControlAdapterConfigBase):
-    """Model config for Control LoRA models."""
-
-    type: Literal[ModelType.ControlLoRa] = ModelType.ControlLoRa
-    trigger_phrases: Optional[set[str]] = Field(description="Set of trigger phrases for this model", default=None)
-    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
-
-    @staticmethod
-    def get_tag() -> Tag:
-        return Tag(f"{ModelType.ControlLoRa.value}.{ModelFormat.Diffusers.value}")
+        return Tag(f"{ModelType.StructuralLoRa.value}.{ModelFormat.LyCORIS.value}")
 
 
 class LoRADiffusersConfig(LoRAConfigBase):
@@ -335,6 +317,12 @@ class VAEDiffusersConfig(ModelConfigBase):
         return Tag(f"{ModelType.VAE.value}.{ModelFormat.Diffusers.value}")
 
 
+class ControlAdapterConfigBase(BaseModel):
+    default_settings: Optional[ControlAdapterDefaultSettings] = Field(
+        description="Default settings for this model", default=None
+    )
+
+
 class ControlNetDiffusersConfig(DiffusersConfigBase, ControlAdapterConfigBase):
     """Model config for ControlNet models (diffusers version)."""
 
@@ -560,8 +548,7 @@ AnyModelConfig = Annotated[
         Annotated[ControlNetDiffusersConfig, ControlNetDiffusersConfig.get_tag()],
         Annotated[ControlNetCheckpointConfig, ControlNetCheckpointConfig.get_tag()],
         Annotated[LoRALyCORISConfig, LoRALyCORISConfig.get_tag()],
-        Annotated[ControlLoRALyCORISConfig, ControlLoRALyCORISConfig.get_tag()],
-        Annotated[ControlLoRADiffusersConfig, ControlLoRADiffusersConfig.get_tag()],
+        Annotated[StructuralLoRALyCORISConfig, StructuralLoRALyCORISConfig.get_tag()],
         Annotated[LoRADiffusersConfig, LoRADiffusersConfig.get_tag()],
         Annotated[T5EncoderConfig, T5EncoderConfig.get_tag()],
         Annotated[T5EncoderBnbQuantizedLlmInt8bConfig, T5EncoderBnbQuantizedLlmInt8bConfig.get_tag()],

invokeai/backend/model_manager/load/__init__.py

@@ -8,7 +8,7 @@ from pathlib import Path
 
 from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig, ModelLoaderBase
 from invokeai.backend.model_manager.load.load_default import ModelLoader
-from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+from invokeai.backend.model_manager.load.model_cache.model_cache_default import ModelCache
 from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry, ModelLoaderRegistryBase
 
 # This registers the subclasses that implement loaders of specific model types

invokeai/backend/model_manager/load/load_base.py

@@ -5,6 +5,7 @@ Base class for model loading in InvokeAI.
 
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
+from dataclasses import dataclass
 from logging import Logger
 from pathlib import Path
 from typing import Any, Dict, Generator, Optional, Tuple
@@ -17,17 +18,19 @@ from invokeai.backend.model_manager.config import (
     AnyModelConfig,
     SubModelType,
 )
-from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
-from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase
 
 
+@dataclass
 class LoadedModelWithoutConfig:
-    """Context manager object that mediates transfer from RAM<->VRAM.
+    """
+    Context manager object that mediates transfer from RAM<->VRAM.
 
     This is a context manager object that has two distinct APIs:
 
     1. Older API (deprecated):
-    Use the LoadedModel object directly as a context manager.  It will move the model into VRAM (on CUDA devices), and
+    Use the LoadedModel object directly as a context manager.
+    It will move the model into VRAM (on CUDA devices), and
     return the model in a form suitable for passing to torch.
     Example:
     ```
@@ -37,9 +40,13 @@ class LoadedModelWithoutConfig:
     ```
 
     2. Newer API (recommended):
-    Call the LoadedModel's `model_on_device()` method in a context. It returns a tuple consisting of a copy of the
-    model's state dict in CPU RAM followed by a copy of the model in VRAM. The state dict is provided to allow LoRAs and
-    other model patchers to return the model to its unpatched state without expensive copy and restore operations.
+    Call the LoadedModel's `model_on_device()` method in a
+    context. It returns a tuple consisting of a copy of
+    the model's state dict in CPU RAM followed by a copy
+    of the model in VRAM. The state dict is provided to allow
+    LoRAs and other model patchers to return the model to
+    its unpatched state without expensive copy and restore
+    operations.
 
     Example:
     ```
@@ -48,48 +55,43 @@ class LoadedModelWithoutConfig:
         image = vae.decode(latents)[0]
     ```
 
-    The state_dict should be treated as a read-only object and never modified. Also be aware that some loadable models
-    do not have a state_dict, in which case this value will be None.
+    The state_dict should be treated as a read-only object and
+    never modified. Also be aware that some loadable models do
+    not have a state_dict, in which case this value will be None.
     """
 
-    def __init__(self, cache_record: CacheRecord, cache: ModelCache):
-        self._cache_record = cache_record
-        self._cache = cache
+    _locker: ModelLockerBase
 
     def __enter__(self) -> AnyModel:
-        self._cache.lock(self._cache_record, None)
+        """Context entry."""
+        self._locker.lock()
         return self.model
 
     def __exit__(self, *args: Any, **kwargs: Any) -> None:
-        self._cache.unlock(self._cache_record)
+        """Context exit."""
+        self._locker.unlock()
 
     @contextmanager
-    def model_on_device(
-        self, working_mem_bytes: Optional[int] = None
-    ) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]:
-        """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device.
-
-        :param working_mem_bytes: The amount of working memory to keep available on the compute device when loading the
-            model.
-        """
-        self._cache.lock(self._cache_record, working_mem_bytes)
+    def model_on_device(self) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]:
+        """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device."""
+        locked_model = self._locker.lock()
         try:
-            yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model)
+            state_dict = self._locker.get_state_dict()
+            yield (state_dict, locked_model)
         finally:
-            self._cache.unlock(self._cache_record)
+            self._locker.unlock()
 
     @property
     def model(self) -> AnyModel:
         """Return the model without locking it."""
-        return self._cache_record.cached_model.model
+        return self._locker.model
 
 
+@dataclass
 class LoadedModel(LoadedModelWithoutConfig):
     """Context manager object that mediates transfer from RAM<->VRAM."""
 
-    def __init__(self, config: Optional[AnyModelConfig], cache_record: CacheRecord, cache: ModelCache):
-        super().__init__(cache_record=cache_record, cache=cache)
-        self.config = config
+    config: Optional[AnyModelConfig] = None
 
 
 # TODO(MM2):
@@ -108,7 +110,7 @@ class ModelLoaderBase(ABC):
         self,
         app_config: InvokeAIAppConfig,
         logger: Logger,
-        ram_cache: ModelCache,
+        ram_cache: ModelCacheBase[AnyModel],
     ):
         """Initialize the loader."""
         pass
@@ -136,6 +138,6 @@ class ModelLoaderBase(ABC):
 
     @property
     @abstractmethod
-    def ram_cache(self) -> ModelCache:
+    def ram_cache(self) -> ModelCacheBase[AnyModel]:
         """Return the ram cache associated with this loader."""
         pass

invokeai/backend/model_manager/load/load_default.py

@@ -14,8 +14,7 @@ from invokeai.backend.model_manager import (
 )
 from invokeai.backend.model_manager.config import DiffusersConfigBase
 from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase
-from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
-from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache, get_model_cache_key
+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_fs
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.util.devices import TorchDevice
@@ -29,7 +28,7 @@ class ModelLoader(ModelLoaderBase):
         self,
         app_config: InvokeAIAppConfig,
         logger: Logger,
-        ram_cache: ModelCache,
+        ram_cache: ModelCacheBase[AnyModel],
     ):
         """Initialize the loader."""
         self._app_config = app_config
@@ -55,11 +54,11 @@ class ModelLoader(ModelLoaderBase):
             raise InvalidModelConfigException(f"Files for model '{model_config.name}' not found at {model_path}")
 
         with skip_torch_weight_init():
-            cache_record = self._load_and_cache(model_config, submodel_type)
-        return LoadedModel(config=model_config, cache_record=cache_record, cache=self._ram_cache)
+            locker = self._load_and_cache(model_config, submodel_type)
+        return LoadedModel(config=model_config, _locker=locker)
 
     @property
-    def ram_cache(self) -> ModelCache:
+    def ram_cache(self) -> ModelCacheBase[AnyModel]:
         """Return the ram cache associated with this loader."""
         return self._ram_cache
 
@@ -67,10 +66,10 @@ class ModelLoader(ModelLoaderBase):
         model_base = self._app_config.models_path
         return (model_base / config.path).resolve()
 
-    def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> CacheRecord:
+    def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> ModelLockerBase:
         stats_name = ":".join([config.base, config.type, config.name, (submodel_type or "")])
         try:
-            return self._ram_cache.get(key=get_model_cache_key(config.key, submodel_type), stats_name=stats_name)
+            return self._ram_cache.get(config.key, submodel_type, stats_name=stats_name)
         except IndexError:
             pass
 
@@ -79,11 +78,16 @@ class ModelLoader(ModelLoaderBase):
         loaded_model = self._load_model(config, submodel_type)
 
         self._ram_cache.put(
-            get_model_cache_key(config.key, submodel_type),
+            config.key,
+            submodel_type=submodel_type,
             model=loaded_model,
         )
 
-        return self._ram_cache.get(key=get_model_cache_key(config.key, submodel_type), stats_name=stats_name)
+        return self._ram_cache.get(
+            key=config.key,
+            submodel_type=submodel_type,
+            stats_name=stats_name,
+        )
 
     def get_size_fs(
         self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None

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

@@ -0,0 +1,6 @@
+"""Init file for ModelCache."""
+
+from .model_cache_base import ModelCacheBase, CacheStats  # noqa F401
+from .model_cache_default import ModelCache  # noqa F401
+
+_all__ = ["ModelCacheBase", "ModelCache", "CacheStats"]

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

@@ -1,33 +0,0 @@
-from dataclasses import dataclass
-
-from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_only_full_load import (
-    CachedModelOnlyFullLoad,
-)
-from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import (
-    CachedModelWithPartialLoad,
-)
-
-
-@dataclass
-class CacheRecord:
-    """A class that represents a model in the model cache."""
-
-    # Cache key.
-    key: str
-    # Model in memory.
-    cached_model: CachedModelWithPartialLoad | CachedModelOnlyFullLoad
-    _locks: int = 0
-
-    def lock(self) -> None:
-        """Lock this record."""
-        self._locks += 1
-
-    def unlock(self) -> None:
-        """Unlock this record."""
-        self._locks -= 1
-        assert self._locks >= 0
-
-    @property
-    def is_locked(self) -> bool:
-        """Return true if record is locked."""
-        return self._locks > 0

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

@@ -1,15 +0,0 @@
-from dataclasses import dataclass, field
-from typing import Dict
-
-
-@dataclass
-class CacheStats(object):
-    """Collect statistics on cache performance."""
-
-    hits: int = 0  # cache hits
-    misses: int = 0  # cache misses
-    high_watermark: int = 0  # amount of cache used
-    in_cache: int = 0  # number of models in cache
-    cleared: int = 0  # number of models cleared to make space
-    cache_size: int = 0  # total size of cache
-    loaded_model_sizes: Dict[str, int] = field(default_factory=dict)

invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py

@@ -1,93 +0,0 @@
-from typing import Any
-
-import torch
-
-
-class CachedModelOnlyFullLoad:
-    """A wrapper around a PyTorch model to handle full loads and unloads between the CPU and the compute device.
-    Note: "VRAM" is used throughout this class to refer to the memory on the compute device. It could be CUDA memory,
-    MPS memory, etc.
-    """
-
-    def __init__(self, model: torch.nn.Module | Any, compute_device: torch.device, total_bytes: int):
-        """Initialize a CachedModelOnlyFullLoad.
-        Args:
-            model (torch.nn.Module | Any): The model to wrap. Should be on the CPU.
-            compute_device (torch.device): The compute device to move the model to.
-            total_bytes (int): The total size (in bytes) of all the weights in the model.
-        """
-        # model is often a torch.nn.Module, but could be any model type. Throughout this class, we handle both cases.
-        self._model = model
-        self._compute_device = compute_device
-        self._offload_device = torch.device("cpu")
-
-        # A CPU read-only copy of the model's state dict.
-        self._cpu_state_dict: dict[str, torch.Tensor] | None = None
-        if isinstance(model, torch.nn.Module):
-            self._cpu_state_dict = model.state_dict()
-
-        self._total_bytes = total_bytes
-        self._is_in_vram = False
-
-    @property
-    def model(self) -> torch.nn.Module:
-        return self._model
-
-    def get_cpu_state_dict(self) -> dict[str, torch.Tensor] | None:
-        """Get a read-only copy of the model's state dict in RAM."""
-        # TODO(ryand): Document this better.
-        return self._cpu_state_dict
-
-    def total_bytes(self) -> int:
-        """Get the total size (in bytes) of all the weights in the model."""
-        return self._total_bytes
-
-    def cur_vram_bytes(self) -> int:
-        """Get the size (in bytes) of the weights that are currently in VRAM."""
-        if self._is_in_vram:
-            return self._total_bytes
-        else:
-            return 0
-
-    def is_in_vram(self) -> bool:
-        """Return true if the model is currently in VRAM."""
-        return self._is_in_vram
-
-    def full_load_to_vram(self) -> int:
-        """Load all weights into VRAM (if supported by the model).
-        Returns:
-            The number of bytes loaded into VRAM.
-        """
-        if self._is_in_vram:
-            # Already in VRAM.
-            return 0
-
-        if not hasattr(self._model, "to"):
-            # Model doesn't support moving to a device.
-            return 0
-
-        if self._cpu_state_dict is not None:
-            new_state_dict: dict[str, torch.Tensor] = {}
-            for k, v in self._cpu_state_dict.items():
-                new_state_dict[k] = v.to(self._compute_device, copy=True)
-            self._model.load_state_dict(new_state_dict, assign=True)
-        self._model.to(self._compute_device)
-
-        self._is_in_vram = True
-        return self._total_bytes
-
-    def full_unload_from_vram(self) -> int:
-        """Unload all weights from VRAM.
-        Returns:
-            The number of bytes unloaded from VRAM.
-        """
-        if not self._is_in_vram:
-            # Already in RAM.
-            return 0
-
-        if self._cpu_state_dict is not None:
-            self._model.load_state_dict(self._cpu_state_dict, assign=True)
-        self._model.to(self._offload_device)
-
-        self._is_in_vram = False
-        return self._total_bytes

invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py

@@ -1,206 +0,0 @@
-import torch
-
-from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
-    CustomModuleMixin,
-)
-from invokeai.backend.util.calc_tensor_size import calc_tensor_size
-from invokeai.backend.util.logging import InvokeAILogger
-
-
-class CachedModelWithPartialLoad:
-    """A wrapper around a PyTorch model to handle partial loads and unloads between the CPU and the compute device.
-
-    Note: "VRAM" is used throughout this class to refer to the memory on the compute device. It could be CUDA memory,
-    MPS memory, etc.
-    """
-
-    def __init__(self, model: torch.nn.Module, compute_device: torch.device):
-        self._model = model
-        self._compute_device = compute_device
-
-        # A CPU read-only copy of the model's state dict.
-        self._cpu_state_dict: dict[str, torch.Tensor] = model.state_dict()
-
-        # A dictionary of the size of each tensor in the state dict.
-        # HACK(ryand): We use this dictionary any time we are doing byte tracking calculations. We do this for
-        # consistency in case the application code has modified the model's size (e.g. by casting to a different
-        # precision). Of course, this means that we are making model cache load/unload decisions based on model size
-        # data that may not be fully accurate.
-        self._state_dict_bytes = {k: calc_tensor_size(v) for k, v in self._cpu_state_dict.items()}
-
-        self._total_bytes = sum(self._state_dict_bytes.values())
-        self._cur_vram_bytes: int | None = None
-
-        self._modules_that_support_autocast = self._find_modules_that_support_autocast()
-        self._keys_in_modules_that_do_not_support_autocast = self._find_keys_in_modules_that_do_not_support_autocast()
-
-    def _find_modules_that_support_autocast(self) -> dict[str, torch.nn.Module]:
-        """Find all modules that support autocasting."""
-        return {n: m for n, m in self._model.named_modules() if isinstance(m, CustomModuleMixin)}  # type: ignore
-
-    def _find_keys_in_modules_that_do_not_support_autocast(self) -> set[str]:
-        keys_in_modules_that_do_not_support_autocast: set[str] = set()
-        for key in self._cpu_state_dict.keys():
-            for module_name in self._modules_that_support_autocast.keys():
-                if key.startswith(module_name):
-                    break
-            else:
-                keys_in_modules_that_do_not_support_autocast.add(key)
-        return keys_in_modules_that_do_not_support_autocast
-
-    def _move_non_persistent_buffers_to_device(self, device: torch.device):
-        """Move the non-persistent buffers to the target device. These buffers are not included in the state dict,
-        so we need to move them manually.
-        """
-        # HACK(ryand): Typically, non-persistent buffers are moved when calling module.to(device). We don't move entire
-        # modules, because we manage the devices of individual tensors using the state dict. Since non-persistent
-        # buffers are not included in the state dict, we need to handle them manually. The only way to do this is by
-        # using private torch.nn.Module attributes.
-        for module in self._model.modules():
-            for name, buffer in module.named_buffers():
-                if name in module._non_persistent_buffers_set:
-                    module._buffers[name] = buffer.to(device, copy=True)
-
-    def _set_autocast_enabled_in_all_modules(self, enabled: bool):
-        """Set autocast_enabled flag in all modules that support device autocasting."""
-        for module in self._modules_that_support_autocast.values():
-            module.set_device_autocasting_enabled(enabled)
-
-    @property
-    def model(self) -> torch.nn.Module:
-        return self._model
-
-    def get_cpu_state_dict(self) -> dict[str, torch.Tensor] | None:
-        """Get a read-only copy of the model's state dict in RAM."""
-        # TODO(ryand): Document this better.
-        return self._cpu_state_dict
-
-    def total_bytes(self) -> int:
-        """Get the total size (in bytes) of all the weights in the model."""
-        return self._total_bytes
-
-    def cur_vram_bytes(self) -> int:
-        """Get the size (in bytes) of the weights that are currently in VRAM."""
-        if self._cur_vram_bytes is None:
-            cur_state_dict = self._model.state_dict()
-            self._cur_vram_bytes = sum(
-                self._state_dict_bytes[k]
-                for k, v in cur_state_dict.items()
-                if v.device.type == self._compute_device.type
-            )
-        return self._cur_vram_bytes
-
-    def full_load_to_vram(self) -> int:
-        """Load all weights into VRAM."""
-        return self.partial_load_to_vram(self.total_bytes())
-
-    def full_unload_from_vram(self) -> int:
-        """Unload all weights from VRAM."""
-        return self.partial_unload_from_vram(self.total_bytes())
-
-    @torch.no_grad()
-    def partial_load_to_vram(self, vram_bytes_to_load: int) -> int:
-        """Load more weights into VRAM without exceeding vram_bytes_to_load.
-
-        Returns:
-            The number of bytes loaded into VRAM.
-        """
-        # TODO(ryand): Handle the case where an exception is thrown while loading or unloading weights. At the very
-        # least, we should reset self._cur_vram_bytes to None.
-
-        vram_bytes_loaded = 0
-
-        cur_state_dict = self._model.state_dict()
-
-        # First, process the keys that *must* be loaded into VRAM.
-        for key in self._keys_in_modules_that_do_not_support_autocast:
-            param = cur_state_dict[key]
-            if param.device.type == self._compute_device.type:
-                continue
-
-            param_size = self._state_dict_bytes[key]
-            cur_state_dict[key] = param.to(self._compute_device, copy=True)
-            vram_bytes_loaded += param_size
-
-        if vram_bytes_loaded > vram_bytes_to_load:
-            logger = InvokeAILogger.get_logger()
-            logger.warning(
-                f"Loaded {vram_bytes_loaded / 2**20} MB into VRAM, but only {vram_bytes_to_load / 2**20} MB were "
-                "requested. This is the minimum set of weights in VRAM required to run the model."
-            )
-
-        # Next, process the keys that can optionally be loaded into VRAM.
-        fully_loaded = True
-        for key, param in cur_state_dict.items():
-            if param.device.type == self._compute_device.type:
-                continue
-
-            param_size = self._state_dict_bytes[key]
-            if vram_bytes_loaded + param_size > vram_bytes_to_load:
-                # TODO(ryand): Should we just break here? If we couldn't fit this parameter into VRAM, is it really
-                # worth continuing to search for a smaller parameter that would fit?
-                fully_loaded = False
-                continue
-
-            cur_state_dict[key] = param.to(self._compute_device, copy=True)
-            vram_bytes_loaded += param_size
-
-        if vram_bytes_loaded > 0:
-            # We load the entire state dict, not just the parameters that changed, in case there are modules that
-            # override _load_from_state_dict() and do some funky stuff that requires the entire state dict.
-            # Alternatively, in the future, grouping parameters by module could probably solve this problem.
-            self._model.load_state_dict(cur_state_dict, assign=True)
-
-        if self._cur_vram_bytes is not None:
-            self._cur_vram_bytes += vram_bytes_loaded
-
-        if fully_loaded:
-            self._set_autocast_enabled_in_all_modules(False)
-        else:
-            self._set_autocast_enabled_in_all_modules(True)
-
-        # Move all non-persistent buffers to the compute device. These are a weird edge case and do not participate in
-        # the vram_bytes_loaded tracking.
-        self._move_non_persistent_buffers_to_device(self._compute_device)
-
-        return vram_bytes_loaded
-
-    @torch.no_grad()
-    def partial_unload_from_vram(self, vram_bytes_to_free: int, keep_required_weights_in_vram: bool = False) -> int:
-        """Unload weights from VRAM until vram_bytes_to_free bytes are freed. Or the entire model is unloaded.
-
-        :param keep_required_weights_in_vram: If True, any weights that must be kept in VRAM to run the model will be
-            kept in VRAM.
-
-        Returns:
-            The number of bytes unloaded from VRAM.
-        """
-        vram_bytes_freed = 0
-        required_weights_in_vram = 0
-
-        offload_device = "cpu"
-        cur_state_dict = self._model.state_dict()
-        for key, param in cur_state_dict.items():
-            if vram_bytes_freed >= vram_bytes_to_free:
-                break
-
-            if param.device.type == offload_device:
-                continue
-
-            if keep_required_weights_in_vram and key in self._keys_in_modules_that_do_not_support_autocast:
-                required_weights_in_vram += self._state_dict_bytes[key]
-                continue
-
-            cur_state_dict[key] = self._cpu_state_dict[key]
-            vram_bytes_freed += self._state_dict_bytes[key]
-
-        if vram_bytes_freed > 0:
-            self._model.load_state_dict(cur_state_dict, assign=True)
-
-        if self._cur_vram_bytes is not None:
-            self._cur_vram_bytes -= vram_bytes_freed
-
-        # We may have gone from a fully-loaded model to a partially-loaded model, so we need to reapply the custom
-        # layers.
-        self._set_autocast_enabled_in_all_modules(True)
-        return vram_bytes_freed

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

@@ -1,33 +0,0 @@
-from contextlib import contextmanager
-
-import torch
-
-from invokeai.backend.util.logging import InvokeAILogger
-
-
-@contextmanager
-def log_operation_vram_usage(operation_name: str):
-    """A helper function for tuning working memory requirements for memory-intensive ops.
-
-    Sample usage:
-
-    ```python
-    with log_operation_vram_usage("some_operation"):
-        some_operation()
-    ```
-    """
-    torch.cuda.synchronize()
-    torch.cuda.reset_peak_memory_stats()
-    max_allocated_before = torch.cuda.max_memory_allocated()
-    max_reserved_before = torch.cuda.max_memory_reserved()
-    try:
-        yield
-    finally:
-        torch.cuda.synchronize()
-        max_allocated_after = torch.cuda.max_memory_allocated()
-        max_reserved_after = torch.cuda.max_memory_reserved()
-        logger = InvokeAILogger.get_logger()
-        logger.info(
-            f">>>{operation_name} Peak VRAM allocated: {(max_allocated_after - max_allocated_before) / 2**20} MB, "
-            f"Peak VRAM reserved: {(max_reserved_after - max_reserved_before) / 2**20} MB"
-        )

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

@@ -1,589 +0,0 @@
-import gc
-import logging
-import time
-from logging import Logger
-from typing import Dict, List, Optional
-
-import psutil
-import torch
-
-from invokeai.backend.model_manager import AnyModel, SubModelType
-from invokeai.backend.model_manager.load.memory_snapshot import MemorySnapshot
-from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
-from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
-from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_only_full_load import (
-    CachedModelOnlyFullLoad,
-)
-from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import (
-    CachedModelWithPartialLoad,
-)
-from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
-    apply_custom_layers_to_model,
-)
-from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
-from invokeai.backend.util.devices import TorchDevice
-from invokeai.backend.util.logging import InvokeAILogger
-from invokeai.backend.util.prefix_logger_adapter import PrefixedLoggerAdapter
-
-# Size of a GB in bytes.
-GB = 2**30
-
-# Size of a MB in bytes.
-MB = 2**20
-
-
-# TODO(ryand): Where should this go? The ModelCache shouldn't be concerned with submodels.
-def get_model_cache_key(model_key: str, submodel_type: Optional[SubModelType] = None) -> str:
-    """Get the cache key for a model based on the optional submodel type."""
-    if submodel_type:
-        return f"{model_key}:{submodel_type.value}"
-    else:
-        return model_key
-
-
-class ModelCache:
-    """A cache for managing models in memory.
-
-    The cache is based on two levels of model storage:
-    - execution_device: The device where most models are executed (typically "cuda", "mps", or "cpu").
-    - storage_device: The device where models are offloaded when not in active use (typically "cpu").
-
-    The model cache is based on the following assumptions:
-    - storage_device_mem_size > execution_device_mem_size
-    - disk_to_storage_device_transfer_time >> storage_device_to_execution_device_transfer_time
-
-    A copy of all models in the cache is always kept on the storage_device. A subset of the models also have a copy on
-    the execution_device.
-
-    Models are moved between the storage_device and the execution_device as necessary. Cache size limits are enforced
-    on both the storage_device and the execution_device. The execution_device cache uses a smallest-first offload
-    policy. The storage_device cache uses a least-recently-used (LRU) offload policy.
-
-    Note: Neither of these offload policies has really been compared against alternatives. It's likely that different
-    policies would be better, although the optimal policies are likely heavily dependent on usage patterns and HW
-    configuration.
-
-    The cache returns context manager generators designed to load the model into the execution device (often GPU) within
-    the context, and unload outside the context.
-
-    Example usage:
-    ```
-    cache = ModelCache(max_cache_size=7.5, max_vram_cache_size=6.0)
-    with cache.get_model('runwayml/stable-diffusion-1-5') as SD1:
-        do_something_on_gpu(SD1)
-    ```
-    """
-
-    def __init__(
-        self,
-        execution_device_working_mem_gb: float,
-        enable_partial_loading: bool,
-        max_ram_cache_size_gb: float | None = None,
-        max_vram_cache_size_gb: float | None = None,
-        execution_device: torch.device | str = "cuda",
-        storage_device: torch.device | str = "cpu",
-        log_memory_usage: bool = False,
-        logger: Optional[Logger] = None,
-    ):
-        """Initialize the model RAM cache.
-
-        :param execution_device_working_mem_gb: The amount of working memory to keep on the GPU (in GB) i.e. non-model
-            VRAM.
-        :param enable_partial_loading: Whether to enable partial loading of models.
-        :param max_ram_cache_size_gb: The maximum amount of CPU RAM to use for model caching in GB. This parameter is
-            kept to maintain compatibility with previous versions of the model cache, but should be deprecated in the
-            future. If set, this parameter overrides the default cache size logic.
-        :param max_vram_cache_size_gb: The amount of VRAM to use for model caching in GB. This parameter is kept to
-            maintain compatibility with previous versions of the model cache, but should be deprecated in the future.
-            If set, this parameter overrides the default cache size logic.
-        :param execution_device: Torch device to load active model into [torch.device('cuda')]
-        :param storage_device: Torch device to save inactive model in [torch.device('cpu')]
-        :param 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 disable this feature unless you are actively inspecting the model cache's
-            behaviour.
-        :param logger: InvokeAILogger to use (otherwise creates one)
-        """
-        self._enable_partial_loading = enable_partial_loading
-        self._execution_device_working_mem_gb = execution_device_working_mem_gb
-        self._execution_device: torch.device = torch.device(execution_device)
-        self._storage_device: torch.device = torch.device(storage_device)
-
-        self._max_ram_cache_size_gb = max_ram_cache_size_gb
-        self._max_vram_cache_size_gb = max_vram_cache_size_gb
-
-        self._logger = PrefixedLoggerAdapter(
-            logger or InvokeAILogger.get_logger(self.__class__.__name__), "MODEL CACHE"
-        )
-        self._log_memory_usage = log_memory_usage
-        self._stats: Optional[CacheStats] = None
-
-        self._cached_models: Dict[str, CacheRecord] = {}
-        self._cache_stack: List[str] = []
-
-    @property
-    def stats(self) -> Optional[CacheStats]:
-        """Return collected CacheStats object."""
-        return self._stats
-
-    @stats.setter
-    def stats(self, stats: CacheStats) -> None:
-        """Set the CacheStats object for collecting cache statistics."""
-        self._stats = stats
-
-    def put(self, key: str, model: AnyModel) -> None:
-        """Add a model to the cache."""
-        if key in self._cached_models:
-            self._logger.debug(
-                f"Attempted to add model {key} ({model.__class__.__name__}), but it already exists in the cache. No action necessary."
-            )
-            return
-
-        size = calc_model_size_by_data(self._logger, model)
-        self.make_room(size)
-
-        # Inject custom modules into the model.
-        if isinstance(model, torch.nn.Module):
-            apply_custom_layers_to_model(model)
-
-        # Partial loading only makes sense on CUDA.
-        # - When running on CPU, there is no 'loading' to do.
-        # - When running on MPS, memory is shared with the CPU, so the default OS memory management already handles this
-        #   well.
-        running_with_cuda = self._execution_device.type == "cuda"
-
-        # Wrap model.
-        if isinstance(model, torch.nn.Module) and running_with_cuda and self._enable_partial_loading:
-            wrapped_model = CachedModelWithPartialLoad(model, self._execution_device)
-        else:
-            wrapped_model = CachedModelOnlyFullLoad(model, self._execution_device, size)
-
-        cache_record = CacheRecord(key=key, cached_model=wrapped_model)
-        self._cached_models[key] = cache_record
-        self._cache_stack.append(key)
-        self._logger.debug(
-            f"Added model {key} (Type: {model.__class__.__name__}, Wrap mode: {wrapped_model.__class__.__name__}, Model size: {size/MB:.2f}MB)"
-        )
-
-    def get(self, key: str, stats_name: Optional[str] = None) -> CacheRecord:
-        """Retrieve a model from the cache.
-
-        :param key: Model key
-        :param stats_name: A human-readable id for the model for the purposes of stats reporting.
-
-        Raises IndexError if the model is not in the cache.
-        """
-        if key in self._cached_models:
-            if self.stats:
-                self.stats.hits += 1
-        else:
-            if self.stats:
-                self.stats.misses += 1
-            self._logger.debug(f"Cache miss: {key}")
-            raise IndexError(f"The model with key {key} is not in the cache.")
-
-        cache_entry = self._cached_models[key]
-
-        # more stats
-        if self.stats:
-            stats_name = stats_name or key
-            self.stats.high_watermark = max(self.stats.high_watermark, self._get_ram_in_use())
-            self.stats.in_cache = len(self._cached_models)
-            self.stats.loaded_model_sizes[stats_name] = max(
-                self.stats.loaded_model_sizes.get(stats_name, 0), cache_entry.cached_model.total_bytes()
-            )
-
-        # This moves the entry to the top (right end) of the stack.
-        self._cache_stack = [k for k in self._cache_stack if k != key]
-        self._cache_stack.append(key)
-
-        self._logger.debug(f"Cache hit: {key} (Type: {cache_entry.cached_model.model.__class__.__name__})")
-        return cache_entry
-
-    def lock(self, cache_entry: CacheRecord, working_mem_bytes: Optional[int]) -> None:
-        """Lock a model for use and move it into VRAM."""
-        if cache_entry.key not in self._cached_models:
-            self._logger.info(
-                f"Locking model cache entry {cache_entry.key} "
-                f"(Type: {cache_entry.cached_model.model.__class__.__name__}), but it has already been dropped from "
-                "the RAM cache. This is a sign that the model loading order is non-optimal in the invocation code "
-                "(See https://github.com/invoke-ai/InvokeAI/issues/7513)."
-            )
-        # cache_entry = self._cached_models[key]
-        cache_entry.lock()
-
-        self._logger.debug(
-            f"Locking model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})"
-        )
-
-        if self._execution_device.type == "cpu":
-            # Models don't need to be loaded into VRAM if we're running on CPU.
-            return
-
-        try:
-            self._load_locked_model(cache_entry, working_mem_bytes)
-            self._logger.debug(
-                f"Finished locking model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})"
-            )
-        except torch.cuda.OutOfMemoryError:
-            self._logger.warning("Insufficient GPU memory to load model. Aborting")
-            cache_entry.unlock()
-            raise
-        except Exception:
-            cache_entry.unlock()
-            raise
-
-        self._log_cache_state()
-
-    def unlock(self, cache_entry: CacheRecord) -> None:
-        """Unlock a model."""
-        if cache_entry.key not in self._cached_models:
-            self._logger.info(
-                f"Unlocking model cache entry {cache_entry.key} "
-                f"(Type: {cache_entry.cached_model.model.__class__.__name__}), but it has already been dropped from "
-                "the RAM cache. This is a sign that the model loading order is non-optimal in the invocation code "
-                "(See https://github.com/invoke-ai/InvokeAI/issues/7513)."
-            )
-        # cache_entry = self._cached_models[key]
-        cache_entry.unlock()
-        self._logger.debug(
-            f"Unlocked model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})"
-        )
-
-    def _load_locked_model(self, cache_entry: CacheRecord, working_mem_bytes: Optional[int] = None) -> None:
-        """Helper function for self.lock(). Loads a locked model into VRAM."""
-        start_time = time.time()
-
-        # Calculate model_vram_needed, the amount of additional VRAM that will be used if we fully load the model into
-        # VRAM.
-        model_cur_vram_bytes = cache_entry.cached_model.cur_vram_bytes()
-        model_total_bytes = cache_entry.cached_model.total_bytes()
-        model_vram_needed = model_total_bytes - model_cur_vram_bytes
-
-        vram_available = self._get_vram_available(working_mem_bytes)
-        self._logger.debug(
-            f"Before unloading: {self._get_vram_state_str(model_cur_vram_bytes, model_total_bytes, vram_available)}"
-        )
-
-        # Make room for the model in VRAM.
-        # 1. If the model can fit entirely in VRAM, then make enough room for it to be loaded fully.
-        # 2. If the model can't fit fully into VRAM, then unload all other models and load as much of the model as
-        #    possible.
-        vram_bytes_freed = self._offload_unlocked_models(model_vram_needed, working_mem_bytes)
-        self._logger.debug(f"Unloaded models (if necessary): vram_bytes_freed={(vram_bytes_freed/MB):.2f}MB")
-
-        # Check the updated vram_available after offloading.
-        vram_available = self._get_vram_available(working_mem_bytes)
-        self._logger.debug(
-            f"After unloading: {self._get_vram_state_str(model_cur_vram_bytes, model_total_bytes, vram_available)}"
-        )
-
-        if vram_available < 0:
-            # There is insufficient VRAM available. As a last resort, try to unload the model being locked from VRAM,
-            # as it may still be loaded from a previous use.
-            vram_bytes_freed_from_own_model = self._move_model_to_ram(cache_entry, -vram_available)
-            vram_available = self._get_vram_available(working_mem_bytes)
-            self._logger.debug(
-                f"Unloaded {vram_bytes_freed_from_own_model/MB:.2f}MB from the model being locked ({cache_entry.key})."
-            )
-
-        # Move as much of the model as possible into VRAM.
-        # For testing, only allow 10% of the model to be loaded into VRAM.
-        # vram_available = int(model_vram_needed * 0.1)
-        # We add 1 MB to the available VRAM to account for small errors in memory tracking (e.g. off-by-one). A fully
-        # loaded model is much faster than a 95% loaded model.
-        model_bytes_loaded = self._move_model_to_vram(cache_entry, vram_available + MB)
-
-        model_cur_vram_bytes = cache_entry.cached_model.cur_vram_bytes()
-        vram_available = self._get_vram_available(working_mem_bytes)
-        loaded_percent = model_cur_vram_bytes / model_total_bytes if model_total_bytes > 0 else 0
-        self._logger.info(
-            f"Loaded model '{cache_entry.key}' ({cache_entry.cached_model.model.__class__.__name__}) onto "
-            f"{self._execution_device.type} device in {(time.time() - start_time):.2f}s. "
-            f"Total model size: {model_total_bytes/MB:.2f}MB, "
-            f"VRAM: {model_cur_vram_bytes/MB:.2f}MB ({loaded_percent:.1%})"
-        )
-        self._logger.debug(f"Loaded model onto execution device: model_bytes_loaded={(model_bytes_loaded/MB):.2f}MB, ")
-        self._logger.debug(
-            f"After loading: {self._get_vram_state_str(model_cur_vram_bytes, model_total_bytes, vram_available)}"
-        )
-
-    def _move_model_to_vram(self, cache_entry: CacheRecord, vram_available: int) -> int:
-        try:
-            if isinstance(cache_entry.cached_model, CachedModelWithPartialLoad):
-                return cache_entry.cached_model.partial_load_to_vram(vram_available)
-            elif isinstance(cache_entry.cached_model, CachedModelOnlyFullLoad):  # type: ignore
-                # Partial load is not supported, so we have not choice but to try and fit it all into VRAM.
-                return cache_entry.cached_model.full_load_to_vram()
-            else:
-                raise ValueError(f"Unsupported cached model type: {type(cache_entry.cached_model)}")
-        except Exception as e:
-            if isinstance(e, torch.cuda.OutOfMemoryError):
-                self._logger.warning("Insufficient GPU memory to load model. Aborting")
-            # If an exception occurs, the model could be left in a bad state, so we delete it from the cache entirely.
-            self._delete_cache_entry(cache_entry)
-            raise
-
-    def _move_model_to_ram(self, cache_entry: CacheRecord, vram_bytes_to_free: int) -> int:
-        try:
-            if isinstance(cache_entry.cached_model, CachedModelWithPartialLoad):
-                return cache_entry.cached_model.partial_unload_from_vram(
-                    vram_bytes_to_free, keep_required_weights_in_vram=cache_entry.is_locked
-                )
-            elif isinstance(cache_entry.cached_model, CachedModelOnlyFullLoad):  # type: ignore
-                return cache_entry.cached_model.full_unload_from_vram()
-            else:
-                raise ValueError(f"Unsupported cached model type: {type(cache_entry.cached_model)}")
-        except Exception:
-            # If an exception occurs, the model could be left in a bad state, so we delete it from the cache entirely.
-            self._delete_cache_entry(cache_entry)
-            raise
-
-    def _get_vram_available(self, working_mem_bytes: Optional[int]) -> int:
-        """Calculate the amount of additional VRAM available for the cache to use (takes into account the working
-        memory).
-        """
-        # If self._max_vram_cache_size_gb is set, then it overrides the default logic.
-        if self._max_vram_cache_size_gb is not None:
-            vram_total_available_to_cache = int(self._max_vram_cache_size_gb * GB)
-            return vram_total_available_to_cache - self._get_vram_in_use()
-
-        working_mem_bytes_default = int(self._execution_device_working_mem_gb * GB)
-        working_mem_bytes = max(working_mem_bytes or working_mem_bytes_default, working_mem_bytes_default)
-
-        if self._execution_device.type == "cuda":
-            # TODO(ryand): It is debatable whether we should use memory_reserved() or memory_allocated() here.
-            # memory_reserved() includes memory reserved by the torch CUDA memory allocator that may or may not be
-            # re-used for future allocations. For now, we use memory_allocated() to be conservative.
-            # vram_reserved = torch.cuda.memory_reserved(self._execution_device)
-            vram_allocated = torch.cuda.memory_allocated(self._execution_device)
-            vram_free, _vram_total = torch.cuda.mem_get_info(self._execution_device)
-            vram_available_to_process = vram_free + vram_allocated
-        elif self._execution_device.type == "mps":
-            vram_reserved = torch.mps.driver_allocated_memory()
-            # TODO(ryand): Is it accurate that MPS shares memory with the CPU?
-            vram_free = psutil.virtual_memory().available
-            vram_available_to_process = vram_free + vram_reserved
-        else:
-            raise ValueError(f"Unsupported execution device: {self._execution_device.type}")
-
-        vram_total_available_to_cache = vram_available_to_process - working_mem_bytes
-        vram_cur_available_to_cache = vram_total_available_to_cache - self._get_vram_in_use()
-        return vram_cur_available_to_cache
-
-    def _get_vram_in_use(self) -> int:
-        """Get the amount of VRAM currently in use by the cache."""
-        if self._execution_device.type == "cuda":
-            return torch.cuda.memory_allocated()
-        elif self._execution_device.type == "mps":
-            return torch.mps.current_allocated_memory()
-        else:
-            raise ValueError(f"Unsupported execution device type: {self._execution_device.type}")
-        # Alternative definition of VRAM in use:
-        # return sum(ce.cached_model.cur_vram_bytes() for ce in self._cached_models.values())
-
-    def _get_ram_available(self) -> int:
-        """Get the amount of RAM available for the cache to use, while keeping memory pressure under control."""
-        # If self._max_ram_cache_size_gb is set, then it overrides the default logic.
-        if self._max_ram_cache_size_gb is not None:
-            ram_total_available_to_cache = int(self._max_ram_cache_size_gb * GB)
-            return ram_total_available_to_cache - self._get_ram_in_use()
-
-        virtual_memory = psutil.virtual_memory()
-        ram_total = virtual_memory.total
-        ram_available = virtual_memory.available
-        ram_used = ram_total - ram_available
-
-        # The total size of all the models in the cache will often be larger than the amount of RAM reported by psutil
-        # (due to lazy-loading and OS RAM caching behaviour). We could just rely on the psutil values, but it feels
-        # like a bad idea to over-fill the model cache. So, for now, we'll try to keep the total size of models in the
-        # cache under the total amount of system RAM.
-        cache_ram_used = self._get_ram_in_use()
-        ram_used = max(cache_ram_used, ram_used)
-
-        # Aim to keep 10% of RAM free.
-        ram_available_based_on_memory_usage = int(ram_total * 0.9) - ram_used
-
-        # If we are running out of RAM, then there's an increased likelihood that we will run into this issue:
-        # https://github.com/invoke-ai/InvokeAI/issues/7513
-        # To keep things running smoothly, there's a minimum RAM cache size that we always allow (even if this means
-        # using swap).
-        min_ram_cache_size_bytes = 4 * GB
-        ram_available_based_on_min_cache_size = min_ram_cache_size_bytes - cache_ram_used
-
-        return max(ram_available_based_on_memory_usage, ram_available_based_on_min_cache_size)
-
-    def _get_ram_in_use(self) -> int:
-        """Get the amount of RAM currently in use."""
-        return sum(ce.cached_model.total_bytes() for ce in self._cached_models.values())
-
-    def _capture_memory_snapshot(self) -> Optional[MemorySnapshot]:
-        if self._log_memory_usage:
-            return MemorySnapshot.capture()
-        return None
-
-    def _get_vram_state_str(self, model_cur_vram_bytes: int, model_total_bytes: int, vram_available: int) -> str:
-        """Helper function for preparing a VRAM state log string."""
-        model_cur_vram_bytes_percent = model_cur_vram_bytes / model_total_bytes if model_total_bytes > 0 else 0
-        return (
-            f"model_total={model_total_bytes/MB:.0f} MB, "
-            + f"model_vram={model_cur_vram_bytes/MB:.0f} MB ({model_cur_vram_bytes_percent:.1%} %), "
-            # + f"vram_total={int(self._max_vram_cache_size * GB)/MB:.0f} MB, "
-            + f"vram_available={(vram_available/MB):.0f} MB, "
-        )
-
-    def _offload_unlocked_models(self, vram_bytes_required: int, working_mem_bytes: Optional[int] = None) -> int:
-        """Offload models from the execution_device until vram_bytes_required bytes are available, or all models are
-        offloaded. Of course, locked models are not offloaded.
-
-        Returns:
-            int: The number of bytes freed based on believed model sizes. The actual change in VRAM may be different.
-        """
-        self._logger.debug(
-            f"Offloading unlocked models with goal of making room for {vram_bytes_required/MB:.2f}MB of VRAM."
-        )
-        vram_bytes_freed = 0
-        # TODO(ryand): Give more thought to the offloading policy used here.
-        cache_entries_increasing_size = sorted(self._cached_models.values(), key=lambda x: x.cached_model.total_bytes())
-        for cache_entry in cache_entries_increasing_size:
-            # We do not fully trust the count of bytes freed, so we check again on each iteration.
-            vram_available = self._get_vram_available(working_mem_bytes)
-            vram_bytes_to_free = vram_bytes_required - vram_available
-            if vram_bytes_to_free <= 0:
-                break
-            if cache_entry.is_locked:
-                # TODO(ryand): In the future, we may want to partially unload locked models, but this requires careful
-                # handling of model patches (e.g. LoRA).
-                continue
-            cache_entry_bytes_freed = self._move_model_to_ram(cache_entry, vram_bytes_to_free)
-            if cache_entry_bytes_freed > 0:
-                self._logger.debug(
-                    f"Unloaded {cache_entry.key} from VRAM to free {(cache_entry_bytes_freed/MB):.0f} MB."
-                )
-            vram_bytes_freed += cache_entry_bytes_freed
-
-        TorchDevice.empty_cache()
-        return vram_bytes_freed
-
-    def _log_cache_state(self, title: str = "Model cache state:", include_entry_details: bool = True):
-        if self._logger.getEffectiveLevel() > logging.DEBUG:
-            # Short circuit if the logger is not set to debug. Some of the data lookups could take a non-negligible
-            # amount of time.
-            return
-
-        log = f"{title}\n"
-
-        log_format = "  {:<30} Limit: {:>7.1f} MB, Used: {:>7.1f} MB ({:>5.1%}), Available: {:>7.1f} MB ({:>5.1%})\n"
-
-        ram_in_use_bytes = self._get_ram_in_use()
-        ram_available_bytes = self._get_ram_available()
-        ram_size_bytes = ram_in_use_bytes + ram_available_bytes
-        ram_in_use_bytes_percent = ram_in_use_bytes / ram_size_bytes if ram_size_bytes > 0 else 0
-        ram_available_bytes_percent = ram_available_bytes / ram_size_bytes if ram_size_bytes > 0 else 0
-        log += log_format.format(
-            f"Storage Device ({self._storage_device.type})",
-            ram_size_bytes / MB,
-            ram_in_use_bytes / MB,
-            ram_in_use_bytes_percent,
-            ram_available_bytes / MB,
-            ram_available_bytes_percent,
-        )
-
-        if self._execution_device.type != "cpu":
-            vram_in_use_bytes = self._get_vram_in_use()
-            vram_available_bytes = self._get_vram_available(None)
-            vram_size_bytes = vram_in_use_bytes + vram_available_bytes
-            vram_in_use_bytes_percent = vram_in_use_bytes / vram_size_bytes if vram_size_bytes > 0 else 0
-            vram_available_bytes_percent = vram_available_bytes / vram_size_bytes if vram_size_bytes > 0 else 0
-            log += log_format.format(
-                f"Compute Device ({self._execution_device.type})",
-                vram_size_bytes / MB,
-                vram_in_use_bytes / MB,
-                vram_in_use_bytes_percent,
-                vram_available_bytes / MB,
-                vram_available_bytes_percent,
-            )
-
-        if torch.cuda.is_available():
-            log += "  {:<30} {:.1f} MB\n".format("CUDA Memory Allocated:", torch.cuda.memory_allocated() / MB)
-        log += "  {:<30} {}\n".format("Total models:", len(self._cached_models))
-
-        if include_entry_details and len(self._cached_models) > 0:
-            log += "  Models:\n"
-            log_format = (
-                "    {:<80} total={:>7.1f} MB, vram={:>7.1f} MB ({:>5.1%}), ram={:>7.1f} MB ({:>5.1%}), locked={}\n"
-            )
-            for cache_record in self._cached_models.values():
-                total_bytes = cache_record.cached_model.total_bytes()
-                cur_vram_bytes = cache_record.cached_model.cur_vram_bytes()
-                cur_vram_bytes_percent = cur_vram_bytes / total_bytes if total_bytes > 0 else 0
-                cur_ram_bytes = total_bytes - cur_vram_bytes
-                cur_ram_bytes_percent = cur_ram_bytes / total_bytes if total_bytes > 0 else 0
-
-                log += log_format.format(
-                    f"{cache_record.key} ({cache_record.cached_model.model.__class__.__name__}):",
-                    total_bytes / MB,
-                    cur_vram_bytes / MB,
-                    cur_vram_bytes_percent,
-                    cur_ram_bytes / MB,
-                    cur_ram_bytes_percent,
-                    cache_record.is_locked,
-                )
-
-        self._logger.debug(log)
-
-    def make_room(self, bytes_needed: int) -> None:
-        """Make enough room in the cache to accommodate a new model of indicated size.
-
-        Note: This function deletes all of the cache's internal references to a model in order to free it. If there are
-        external references to the model, there's nothing that the cache can do about it, and those models will not be
-        garbage-collected.
-        """
-        self._logger.debug(f"Making room for {bytes_needed/MB:.2f}MB of RAM.")
-        self._log_cache_state(title="Before dropping models:")
-
-        ram_bytes_available = self._get_ram_available()
-        ram_bytes_to_free = max(0, bytes_needed - ram_bytes_available)
-
-        ram_bytes_freed = 0
-        pos = 0
-        models_cleared = 0
-        while ram_bytes_freed < ram_bytes_to_free and pos < len(self._cache_stack):
-            model_key = self._cache_stack[pos]
-            cache_entry = self._cached_models[model_key]
-
-            if not cache_entry.is_locked:
-                ram_bytes_freed += cache_entry.cached_model.total_bytes()
-                self._logger.debug(
-                    f"Dropping {model_key} from RAM cache to free {(cache_entry.cached_model.total_bytes()/MB):.2f}MB."
-                )
-                self._delete_cache_entry(cache_entry)
-                del cache_entry
-                models_cleared += 1
-            else:
-                pos += 1
-
-        if models_cleared > 0:
-            # There would likely be some 'garbage' to be collected regardless of whether a model was cleared or not, but
-            # there is a significant time cost to calling `gc.collect()`, so we want to use it sparingly. (The time cost
-            # is high even if no garbage gets collected.)
-            #
-            # Calling gc.collect(...) when a model is cleared seems like a good middle-ground:
-            # - If models had to be cleared, it's a signal that we are close to our memory limit.
-            # - If models were cleared, there's a good chance that there's a significant amount of garbage to be
-            #   collected.
-            #
-            # 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()
-        self._logger.debug(f"Dropped {models_cleared} models to free {ram_bytes_freed/MB:.2f}MB of RAM.")
-        self._log_cache_state(title="After dropping models:")
-
-    def _delete_cache_entry(self, cache_entry: CacheRecord) -> None:
-        """Delete cache_entry from the cache if it exists. No exception is thrown if it doesn't exist."""
-        self._cache_stack = [key for key in self._cache_stack if key != cache_entry.key]
-        self._cached_models.pop(cache_entry.key, None)

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

@@ -0,0 +1,221 @@
+# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development team
+# TODO: Add Stalker's proper name to copyright
+"""
+Manage a RAM cache of diffusion/transformer models for fast switching.
+They are moved between GPU VRAM and CPU RAM as necessary. If the cache
+grows larger than a preset maximum, then the least recently used
+model will be cleared and (re)loaded from disk when next needed.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from logging import Logger
+from typing import Dict, Generic, Optional, TypeVar
+
+import torch
+
+from invokeai.backend.model_manager.config import AnyModel, SubModelType
+
+
+class ModelLockerBase(ABC):
+    """Base class for the model locker used by the loader."""
+
+    @abstractmethod
+    def lock(self) -> AnyModel:
+        """Lock the contained model and move it into VRAM."""
+        pass
+
+    @abstractmethod
+    def unlock(self) -> None:
+        """Unlock the contained model, and remove it from VRAM."""
+        pass
+
+    @abstractmethod
+    def get_state_dict(self) -> Optional[Dict[str, torch.Tensor]]:
+        """Return the state dict (if any) for the cached model."""
+        pass
+
+    @property
+    @abstractmethod
+    def model(self) -> AnyModel:
+        """Return the model."""
+        pass
+
+
+T = TypeVar("T")
+
+
+@dataclass
+class CacheRecord(Generic[T]):
+    """
+    Elements of the cache:
+
+    key: Unique key for each model, same as used in the models database.
+    model: Model in memory.
+    state_dict: A read-only copy of the model's state dict in RAM. It will be
+                used as a template for creating a copy in the VRAM.
+    size: Size of the model
+    loaded: True if the model's state dict is currently in VRAM
+
+    Before a model is executed, the state_dict template is copied into VRAM,
+    and then injected into the model. When the model is finished, the VRAM
+    copy of the state dict is deleted, and the RAM version is reinjected
+    into the model.
+
+    The state_dict should be treated as a read-only attribute. Do not attempt
+    to patch or otherwise modify it. Instead, patch the copy of the state_dict
+    after it is loaded into the execution device (e.g. CUDA) using the `LoadedModel`
+    context manager call `model_on_device()`.
+    """
+
+    key: str
+    model: T
+    device: torch.device
+    state_dict: Optional[Dict[str, torch.Tensor]]
+    size: int
+    loaded: bool = False
+    _locks: int = 0
+
+    def lock(self) -> None:
+        """Lock this record."""
+        self._locks += 1
+
+    def unlock(self) -> None:
+        """Unlock this record."""
+        self._locks -= 1
+        assert self._locks >= 0
+
+    @property
+    def locked(self) -> bool:
+        """Return true if record is locked."""
+        return self._locks > 0
+
+
+@dataclass
+class CacheStats(object):
+    """Collect statistics on cache performance."""
+
+    hits: int = 0  # cache hits
+    misses: int = 0  # cache misses
+    high_watermark: int = 0  # amount of cache used
+    in_cache: int = 0  # number of models in cache
+    cleared: int = 0  # number of models cleared to make space
+    cache_size: int = 0  # total size of cache
+    loaded_model_sizes: Dict[str, int] = field(default_factory=dict)
+
+
+class ModelCacheBase(ABC, Generic[T]):
+    """Virtual base class for RAM model cache."""
+
+    @property
+    @abstractmethod
+    def storage_device(self) -> torch.device:
+        """Return the storage device (e.g. "CPU" for RAM)."""
+        pass
+
+    @property
+    @abstractmethod
+    def execution_device(self) -> torch.device:
+        """Return the exection device (e.g. "cuda" for VRAM)."""
+        pass
+
+    @property
+    @abstractmethod
+    def lazy_offloading(self) -> bool:
+        """Return true if the cache is configured to lazily offload models in VRAM."""
+        pass
+
+    @property
+    @abstractmethod
+    def max_cache_size(self) -> float:
+        """Return the maximum size the RAM cache can grow to."""
+        pass
+
+    @max_cache_size.setter
+    @abstractmethod
+    def max_cache_size(self, value: float) -> None:
+        """Set the cap on vram cache size."""
+
+    @property
+    @abstractmethod
+    def max_vram_cache_size(self) -> float:
+        """Return the maximum size the VRAM cache can grow to."""
+        pass
+
+    @max_vram_cache_size.setter
+    @abstractmethod
+    def max_vram_cache_size(self, value: float) -> float:
+        """Set the maximum size the VRAM cache can grow to."""
+        pass
+
+    @abstractmethod
+    def offload_unlocked_models(self, size_required: int) -> None:
+        """Offload from VRAM any models not actively in use."""
+        pass
+
+    @abstractmethod
+    def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device: torch.device) -> None:
+        """Move model into the indicated device."""
+        pass
+
+    @property
+    @abstractmethod
+    def stats(self) -> Optional[CacheStats]:
+        """Return collected CacheStats object."""
+        pass
+
+    @stats.setter
+    @abstractmethod
+    def stats(self, stats: CacheStats) -> None:
+        """Set the CacheStats object for collectin cache statistics."""
+        pass
+
+    @property
+    @abstractmethod
+    def logger(self) -> Logger:
+        """Return the logger used by the cache."""
+        pass
+
+    @abstractmethod
+    def make_room(self, size: int) -> None:
+        """Make enough room in the cache to accommodate a new model of indicated size."""
+        pass
+
+    @abstractmethod
+    def put(
+        self,
+        key: str,
+        model: T,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> None:
+        """Store model under key and optional submodel_type."""
+        pass
+
+    @abstractmethod
+    def get(
+        self,
+        key: str,
+        submodel_type: Optional[SubModelType] = None,
+        stats_name: Optional[str] = None,
+    ) -> ModelLockerBase:
+        """
+        Retrieve model using key and optional submodel_type.
+
+        :param key: Opaque model key
+        :param submodel_type: Type of the submodel to fetch
+        :param stats_name: A human-readable id for the model for the purposes of
+        stats reporting.
+
+        This may raise an IndexError if the model is not in the cache.
+        """
+        pass
+
+    @abstractmethod
+    def cache_size(self) -> int:
+        """Get the total size of the models currently cached."""
+        pass
+
+    @abstractmethod
+    def print_cuda_stats(self) -> None:
+        """Log debugging information on CUDA usage."""
+        pass

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

@@ -0,0 +1,426 @@
+# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development team
+# TODO: Add Stalker's proper name to copyright
+""" """
+
+import gc
+import math
+import time
+from contextlib import suppress
+from logging import Logger
+from typing import Dict, List, Optional
+
+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.model_manager.load.model_cache.model_cache_base import (
+    CacheRecord,
+    CacheStats,
+    ModelCacheBase,
+    ModelLockerBase,
+)
+from invokeai.backend.model_manager.load.model_cache.model_locker import ModelLocker
+from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.logging import InvokeAILogger
+
+# Size of a GB in bytes.
+GB = 2**30
+
+# Size of a MB in bytes.
+MB = 2**20
+
+
+class ModelCache(ModelCacheBase[AnyModel]):
+    """A cache for managing models in memory.
+
+    The cache is based on two levels of model storage:
+    - execution_device: The device where most models are executed (typically "cuda", "mps", or "cpu").
+    - storage_device: The device where models are offloaded when not in active use (typically "cpu").
+
+    The model cache is based on the following assumptions:
+    - storage_device_mem_size > execution_device_mem_size
+    - disk_to_storage_device_transfer_time >> storage_device_to_execution_device_transfer_time
+
+    A copy of all models in the cache is always kept on the storage_device. A subset of the models also have a copy on
+    the execution_device.
+
+    Models are moved between the storage_device and the execution_device as necessary. Cache size limits are enforced
+    on both the storage_device and the execution_device. The execution_device cache uses a smallest-first offload
+    policy. The storage_device cache uses a least-recently-used (LRU) offload policy.
+
+    Note: Neither of these offload policies has really been compared against alternatives. It's likely that different
+    policies would be better, although the optimal policies are likely heavily dependent on usage patterns and HW
+    configuration.
+
+    The cache returns context manager generators designed to load the model into the execution device (often GPU) within
+    the context, and unload outside the context.
+
+    Example usage:
+    ```
+    cache = ModelCache(max_cache_size=7.5, max_vram_cache_size=6.0)
+    with cache.get_model('runwayml/stable-diffusion-1-5') as SD1:
+        do_something_on_gpu(SD1)
+    ```
+    """
+
+    def __init__(
+        self,
+        max_cache_size: float,
+        max_vram_cache_size: float,
+        execution_device: torch.device = torch.device("cuda"),
+        storage_device: torch.device = torch.device("cpu"),
+        precision: torch.dtype = torch.float16,
+        lazy_offloading: bool = True,
+        log_memory_usage: bool = False,
+        logger: Optional[Logger] = None,
+    ):
+        """
+        Initialize the model RAM cache.
+
+        :param max_cache_size: Maximum size of the storage_device cache in GBs.
+        :param max_vram_cache_size: Maximum size of the execution_device cache in GBs.
+        :param execution_device: Torch device to load active model into [torch.device('cuda')]
+        :param storage_device: Torch device to save inactive model in [torch.device('cpu')]
+        :param precision: Precision for loaded models [torch.float16]
+        :param lazy_offloading: Keep model in VRAM until another model needs to be loaded
+        :param 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 disable this feature unless you are actively inspecting the model cache's
+            behaviour.
+        :param logger: InvokeAILogger to use (otherwise creates one)
+        """
+        # allow lazy offloading only when vram cache enabled
+        self._lazy_offloading = lazy_offloading and max_vram_cache_size > 0
+        self._max_cache_size: float = max_cache_size
+        self._max_vram_cache_size: float = max_vram_cache_size
+        self._execution_device: torch.device = execution_device
+        self._storage_device: torch.device = storage_device
+        self._logger = logger or InvokeAILogger.get_logger(self.__class__.__name__)
+        self._log_memory_usage = log_memory_usage
+        self._stats: Optional[CacheStats] = None
+
+        self._cached_models: Dict[str, CacheRecord[AnyModel]] = {}
+        self._cache_stack: List[str] = []
+
+    @property
+    def logger(self) -> Logger:
+        """Return the logger used by the cache."""
+        return self._logger
+
+    @property
+    def lazy_offloading(self) -> bool:
+        """Return true if the cache is configured to lazily offload models in VRAM."""
+        return self._lazy_offloading
+
+    @property
+    def storage_device(self) -> torch.device:
+        """Return the storage device (e.g. "CPU" for RAM)."""
+        return self._storage_device
+
+    @property
+    def execution_device(self) -> torch.device:
+        """Return the exection device (e.g. "cuda" for VRAM)."""
+        return self._execution_device
+
+    @property
+    def max_cache_size(self) -> float:
+        """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 max_vram_cache_size(self) -> float:
+        """Return the cap on vram cache size."""
+        return self._max_vram_cache_size
+
+    @max_vram_cache_size.setter
+    def max_vram_cache_size(self, value: float) -> None:
+        """Set the cap on vram cache size."""
+        self._max_vram_cache_size = value
+
+    @property
+    def stats(self) -> Optional[CacheStats]:
+        """Return collected CacheStats object."""
+        return self._stats
+
+    @stats.setter
+    def stats(self, stats: CacheStats) -> None:
+        """Set the CacheStats object for collectin cache statistics."""
+        self._stats = stats
+
+    def cache_size(self) -> int:
+        """Get the total size of the models currently cached."""
+        total = 0
+        for cache_record in self._cached_models.values():
+            total += cache_record.size
+        return total
+
+    def put(
+        self,
+        key: str,
+        model: AnyModel,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> None:
+        """Store model under key and optional submodel_type."""
+        key = self._make_cache_key(key, submodel_type)
+        if key in self._cached_models:
+            return
+        size = calc_model_size_by_data(self.logger, model)
+        self.make_room(size)
+
+        running_on_cpu = self.execution_device == torch.device("cpu")
+        state_dict = model.state_dict() if isinstance(model, torch.nn.Module) and not running_on_cpu else None
+        cache_record = CacheRecord(key=key, model=model, device=self.storage_device, state_dict=state_dict, size=size)
+        self._cached_models[key] = cache_record
+        self._cache_stack.append(key)
+
+    def get(
+        self,
+        key: str,
+        submodel_type: Optional[SubModelType] = None,
+        stats_name: Optional[str] = None,
+    ) -> ModelLockerBase:
+        """
+        Retrieve model using key and optional submodel_type.
+
+        :param key: Opaque model key
+        :param submodel_type: Type of the submodel to fetch
+        :param stats_name: A human-readable id for the model for the purposes of
+        stats reporting.
+
+        This may raise an IndexError if the model is not in the cache.
+        """
+        key = self._make_cache_key(key, submodel_type)
+        if key in self._cached_models:
+            if self.stats:
+                self.stats.hits += 1
+        else:
+            if self.stats:
+                self.stats.misses += 1
+            raise IndexError(f"The model with key {key} is not in the cache.")
+
+        cache_entry = self._cached_models[key]
+
+        # more stats
+        if self.stats:
+            stats_name = stats_name or key
+            self.stats.cache_size = int(self._max_cache_size * GB)
+            self.stats.high_watermark = max(self.stats.high_watermark, self.cache_size())
+            self.stats.in_cache = len(self._cached_models)
+            self.stats.loaded_model_sizes[stats_name] = max(
+                self.stats.loaded_model_sizes.get(stats_name, 0), cache_entry.size
+            )
+
+        # this moves the entry to the top (right end) of the stack
+        with suppress(Exception):
+            self._cache_stack.remove(key)
+        self._cache_stack.append(key)
+        return ModelLocker(
+            cache=self,
+            cache_entry=cache_entry,
+        )
+
+    def _capture_memory_snapshot(self) -> Optional[MemorySnapshot]:
+        if self._log_memory_usage:
+            return MemorySnapshot.capture()
+        return None
+
+    def _make_cache_key(self, model_key: str, submodel_type: Optional[SubModelType] = None) -> str:
+        if submodel_type:
+            return f"{model_key}:{submodel_type.value}"
+        else:
+            return model_key
+
+    def offload_unlocked_models(self, size_required: int) -> None:
+        """Offload models from the execution_device to make room for size_required.
+
+        :param size_required: The amount of space to clear in the execution_device cache, in bytes.
+        """
+        reserved = self._max_vram_cache_size * GB
+        vram_in_use = torch.cuda.memory_allocated() + size_required
+        self.logger.debug(f"{(vram_in_use/GB):.2f}GB VRAM needed for models; max allowed={(reserved/GB):.2f}GB")
+        for _, cache_entry in sorted(self._cached_models.items(), key=lambda x: x[1].size):
+            if vram_in_use <= reserved:
+                break
+            if not cache_entry.loaded:
+                continue
+            if not cache_entry.locked:
+                self.move_model_to_device(cache_entry, self.storage_device)
+                cache_entry.loaded = False
+                vram_in_use = torch.cuda.memory_allocated() + size_required
+                self.logger.debug(
+                    f"Removing {cache_entry.key} from VRAM to free {(cache_entry.size/GB):.2f}GB; vram free = {(torch.cuda.memory_allocated()/GB):.2f}GB"
+                )
+
+        TorchDevice.empty_cache()
+
+    def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device: torch.device) -> None:
+        """Move model into the indicated device.
+
+        :param cache_entry: The CacheRecord for the model
+        :param target_device: The torch.device to move the model into
+
+        May raise a torch.cuda.OutOfMemoryError
+        """
+        self.logger.debug(f"Called to move {cache_entry.key} to {target_device}")
+        source_device = cache_entry.device
+
+        # Note: We compare device types only so that 'cuda' == 'cuda:0'.
+        # This would need to be revised to support multi-GPU.
+        if torch.device(source_device).type == torch.device(target_device).type:
+            return
+
+        # Some models don't have a `to` method, in which case they run in RAM/CPU.
+        if not hasattr(cache_entry.model, "to"):
+            return
+
+        # This roundabout method for moving the model around is done to avoid
+        # the cost of moving the model from RAM to VRAM and then back from VRAM to RAM.
+        # When moving to VRAM, we copy (not move) each element of the state dict from
+        # RAM to a new state dict in VRAM, and then inject it into the model.
+        # This operation is slightly faster than running `to()` on the whole model.
+        #
+        # When the model needs to be removed from VRAM we simply delete the copy
+        # of the state dict in VRAM, and reinject the state dict that is cached
+        # in RAM into the model. So this operation is very fast.
+        start_model_to_time = time.time()
+        snapshot_before = self._capture_memory_snapshot()
+
+        try:
+            if cache_entry.state_dict is not None:
+                assert hasattr(cache_entry.model, "load_state_dict")
+                if target_device == self.storage_device:
+                    cache_entry.model.load_state_dict(cache_entry.state_dict, assign=True)
+                else:
+                    new_dict: Dict[str, torch.Tensor] = {}
+                    for k, v in cache_entry.state_dict.items():
+                        new_dict[k] = v.to(target_device, copy=True)
+                    cache_entry.model.load_state_dict(new_dict, assign=True)
+            cache_entry.model.to(target_device)
+            cache_entry.device = target_device
+        except Exception as e:  # blow away cache entry
+            self._delete_cache_entry(cache_entry)
+            raise e
+
+        snapshot_after = self._capture_memory_snapshot()
+        end_model_to_time = time.time()
+        self.logger.debug(
+            f"Moved model '{cache_entry.key}' from {source_device} to"
+            f" {target_device} in {(end_model_to_time-start_model_to_time):.2f}s."
+            f"Estimated model size: {(cache_entry.size/GB):.3f} GB."
+            f"{get_pretty_snapshot_diff(snapshot_before, snapshot_after)}"
+        )
+
+        if (
+            snapshot_before is not None
+            and snapshot_after is not None
+            and snapshot_before.vram is not None
+            and snapshot_after.vram is not None
+        ):
+            vram_change = abs(snapshot_before.vram - snapshot_after.vram)
+
+            # If the estimated model size does not match the change in VRAM, log a warning.
+            if not math.isclose(
+                vram_change,
+                cache_entry.size,
+                rel_tol=0.1,
+                abs_tol=10 * MB,
+            ):
+                self.logger.debug(
+                    f"Moving model '{cache_entry.key}' from {source_device} to"
+                    f" {target_device} caused an unexpected change in VRAM usage. The model's"
+                    " estimated size may be incorrect. Estimated model size:"
+                    f" {(cache_entry.size/GB):.3f} GB.\n"
+                    f"{get_pretty_snapshot_diff(snapshot_before, snapshot_after)}"
+                )
+
+    def print_cuda_stats(self) -> None:
+        """Log CUDA diagnostics."""
+        vram = "%4.2fG" % (torch.cuda.memory_allocated() / GB)
+        ram = "%4.2fG" % (self.cache_size() / GB)
+
+        in_ram_models = 0
+        in_vram_models = 0
+        locked_in_vram_models = 0
+        for cache_record in self._cached_models.values():
+            if hasattr(cache_record.model, "device"):
+                if cache_record.model.device == self.storage_device:
+                    in_ram_models += 1
+                else:
+                    in_vram_models += 1
+                if cache_record.locked:
+                    locked_in_vram_models += 1
+
+                self.logger.debug(
+                    f"Current VRAM/RAM usage: {vram}/{ram}; models_in_ram/models_in_vram(locked) ="
+                    f" {in_ram_models}/{in_vram_models}({locked_in_vram_models})"
+                )
+
+    def make_room(self, size: int) -> None:
+        """Make enough room in the cache to accommodate a new model of indicated size.
+
+        Note: This function deletes all of the cache's internal references to a model in order to free it. If there are
+        external references to the model, there's nothing that the cache can do about it, and those models will not be
+        garbage-collected.
+        """
+        bytes_needed = size
+        maximum_size = self.max_cache_size * GB  # stored in GB, convert to bytes
+        current_size = self.cache_size()
+
+        if current_size + bytes_needed > maximum_size:
+            self.logger.debug(
+                f"Max cache size exceeded: {(current_size/GB):.2f}/{self.max_cache_size:.2f} GB, need an additional"
+                f" {(bytes_needed/GB):.2f} GB"
+            )
+
+        self.logger.debug(f"Before making_room: cached_models={len(self._cached_models)}")
+
+        pos = 0
+        models_cleared = 0
+        while current_size + bytes_needed > maximum_size and pos < len(self._cache_stack):
+            model_key = self._cache_stack[pos]
+            cache_entry = self._cached_models[model_key]
+            device = cache_entry.model.device if hasattr(cache_entry.model, "device") else None
+            self.logger.debug(
+                f"Model: {model_key}, locks: {cache_entry._locks}, device: {device}, loaded: {cache_entry.loaded}"
+            )
+
+            if not cache_entry.locked:
+                self.logger.debug(
+                    f"Removing {model_key} from RAM cache to free at least {(size/GB):.2f} GB (-{(cache_entry.size/GB):.2f} GB)"
+                )
+                current_size -= cache_entry.size
+                models_cleared += 1
+                self._delete_cache_entry(cache_entry)
+                del cache_entry
+
+            else:
+                pos += 1
+
+        if models_cleared > 0:
+            # There would likely be some 'garbage' to be collected regardless of whether a model was cleared or not, but
+            # there is a significant time cost to calling `gc.collect()`, so we want to use it sparingly. (The time cost
+            # is high even if no garbage gets collected.)
+            #
+            # Calling gc.collect(...) when a model is cleared seems like a good middle-ground:
+            # - If models had to be cleared, it's a signal that we are close to our memory limit.
+            # - If models were cleared, there's a good chance that there's a significant amount of garbage to be
+            #   collected.
+            #
+            # 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()
+        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]

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

@@ -0,0 +1,64 @@
+"""
+Base class and implementation of a class that moves models in and out of VRAM.
+"""
+
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.model_manager import AnyModel
+from invokeai.backend.model_manager.load.model_cache.model_cache_base import (
+    CacheRecord,
+    ModelCacheBase,
+    ModelLockerBase,
+)
+
+
+class ModelLocker(ModelLockerBase):
+    """Internal class that mediates movement in and out of GPU."""
+
+    def __init__(self, cache: ModelCacheBase[AnyModel], cache_entry: CacheRecord[AnyModel]):
+        """
+        Initialize the model locker.
+
+        :param cache: The ModelCache object
+        :param cache_entry: The entry in the model cache
+        """
+        self._cache = cache
+        self._cache_entry = cache_entry
+
+    @property
+    def model(self) -> AnyModel:
+        """Return the model without moving it around."""
+        return self._cache_entry.model
+
+    def get_state_dict(self) -> Optional[Dict[str, torch.Tensor]]:
+        """Return the state dict (if any) for the cached model."""
+        return self._cache_entry.state_dict
+
+    def lock(self) -> AnyModel:
+        """Move the model into the execution device (GPU) and lock it."""
+        self._cache_entry.lock()
+        try:
+            if self._cache.lazy_offloading:
+                self._cache.offload_unlocked_models(self._cache_entry.size)
+            self._cache.move_model_to_device(self._cache_entry, self._cache.execution_device)
+            self._cache_entry.loaded = True
+            self._cache.logger.debug(f"Locking {self._cache_entry.key} in {self._cache.execution_device}")
+            self._cache.print_cuda_stats()
+        except torch.cuda.OutOfMemoryError:
+            self._cache.logger.warning("Insufficient GPU memory to load model. Aborting")
+            self._cache_entry.unlock()
+            raise
+        except Exception:
+            self._cache_entry.unlock()
+            raise
+
+        return self.model
+
+    def unlock(self) -> None:
+        """Call upon exit from context."""
+        self._cache_entry.unlock()
+        if not self._cache.lazy_offloading:
+            self._cache.offload_unlocked_models(0)
+            self._cache.print_cuda_stats()

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/cast_to_device.py

@@ -1,15 +0,0 @@
-from typing import TypeVar
-
-import torch
-
-T = TypeVar("T", torch.Tensor, None, torch.Tensor | None)
-
-
-def cast_to_device(t: T, to_device: torch.device) -> T:
-    """Helper function to cast an optional tensor to a target device."""
-    if t is None:
-        return t
-
-    if t.device.type != to_device.type:
-        return t.to(to_device)
-    return t

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/README.md

@@ -1,8 +0,0 @@
-
-This directory contains custom implementations of common torch.nn.Module classes that add support for:
-- Streaming weights to the execution device
-- Applying sidecar patches at execution time (e.g. sidecar LoRA layers)
-
-Each custom class sub-classes the original module type that is is replacing, so the following properties are preserved:
-- `isinstance(m, torch.nn.OrginalModule)` should still work.
-- Patching the weights directly (e.g. for LoRA) should still work. (Of course, this is not possible for quantized layers, hence the sidecar support.)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_conv1d.py

@@ -1,43 +0,0 @@
-import torch
-
-from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
-from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
-    CustomModuleMixin,
-)
-from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.utils import (
-    add_nullable_tensors,
-)
-
-
-class CustomConv1d(torch.nn.Conv1d, CustomModuleMixin):
-    def _autocast_forward_with_patches(self, input: torch.Tensor) -> torch.Tensor:
-        weight = cast_to_device(self.weight, input.device)
-        bias = cast_to_device(self.bias, input.device)
-
-        # Prepare the original parameters for the patch aggregation.
-        orig_params = {"weight": weight, "bias": bias}
-        # Filter out None values.
-        orig_params = {k: v for k, v in orig_params.items() if v is not None}
-
-        aggregated_param_residuals = self._aggregate_patch_parameters(
-            patches_and_weights=self._patches_and_weights,
-            orig_params=orig_params,
-            device=input.device,
-        )
-
-        weight = add_nullable_tensors(weight, aggregated_param_residuals.get("weight", None))
-        bias = add_nullable_tensors(bias, aggregated_param_residuals.get("bias", None))
-        return self._conv_forward(input, weight, bias)
-
-    def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor:
-        weight = cast_to_device(self.weight, input.device)
-        bias = cast_to_device(self.bias, input.device)
-        return self._conv_forward(input, weight, bias)
-
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        if len(self._patches_and_weights) > 0:
-            return self._autocast_forward_with_patches(input)
-        elif self._device_autocasting_enabled:
-            return self._autocast_forward(input)
-        else:
-            return super().forward(input)