First pass at dynamically calculating the working memory requirements for the VAE decoding operation. Still need to tune SD3 and FLUX.

## Summary

Add the ability to control the weight of a FLUX Control LoRA.

## Example

Original image:
<div style="display: flex; gap: 10px;">
<img
src="https://github.com/user-attachments/assets/4a2d9f4a-b58b-4df6-af90-67b018763a38"
alt="Image 1" width="300"/>
</div>

Prompt: `a scarecrow playing tennis`
Weights: 0.4, 0.6, 0.8, 1.0
<div style="display: flex; gap: 10px;">
<img
src="https://github.com/user-attachments/assets/62b83fd6-46ce-460a-8d51-9c2cda9b05c9"
alt="Image 1" width="300"/>
<img
src="https://github.com/user-attachments/assets/75442207-1538-46bc-9d6b-08ac5c235c93"
alt="Image 2" width="300"/>
</div>
<div style="display: flex; gap: 10px;">
<img
src="https://github.com/user-attachments/assets/4a9dc9ea-9757-4965-837e-197fc9243007"
alt="Image 1" width="300"/>
<img
src="https://github.com/user-attachments/assets/846f6918-ca82-4482-8c19-19172752fa8c"
alt="Image 2" width="300"/>
</div>

## QA Instructions

- [x] weight control changes strength of control image
- [x] Test that results match across both quantized and non-quantized.

## Merge Plan

**_Do not merge this PR yet._**

1. Merge #7450 
2. Merge #7446 
3. Change target branch to main
4. Merge this branch.

## Checklist

- [ ] _The PR has a short but descriptive title, suitable for a
changelog_
- [ ] _Tests added / updated (if applicable)_
- [ ] _Documentation added / updated (if applicable)_
- [ ] _Updated `What's New` copy (if doing a release after this PR)_

SECURITY.md

@@ -0,0 +1,14 @@
+# Security Policy
+
+## Supported Versions
+
+Only the latest version of Invoke will receive security updates. 
+We do not currently maintain multiple versions of the application with updates.
+
+## Reporting a Vulnerability
+
+To report a vulnerability, contact the Invoke team directly at security@invoke.ai
+
+At this time, we do not maintain a formal bug bounty program. 
+
+You can also share identified security issues with our team on huntr.com

docker/Dockerfile

@@ -2,29 +2,42 @@
 
 ## Builder stage
 
-FROM library/ubuntu:23.04 AS builder
+FROM library/ubuntu:24.04 AS builder
 
 ARG DEBIAN_FRONTEND=noninteractive
 RUN rm -f /etc/apt/apt.conf.d/docker-clean; echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' > /etc/apt/apt.conf.d/keep-cache
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt update && apt-get install -y \
-        git \
-        python3-venv \
-        python3-pip \
-        build-essential
+        build-essential \
+        git
 
-ENV INVOKEAI_SRC=/opt/invokeai
-ENV VIRTUAL_ENV=/opt/venv/invokeai
+# Install `uv` for package management
+COPY --from=ghcr.io/astral-sh/uv:0.5.5 /uv /uvx /bin/
 
+ENV VIRTUAL_ENV=/opt/venv
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+ENV INVOKEAI_SRC=/opt/invokeai
+ENV PYTHON_VERSION=3.11
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_LINK_MODE=copy
+
 ARG GPU_DRIVER=cuda
 ARG TARGETPLATFORM="linux/amd64"
 # unused but available
 ARG BUILDPLATFORM
 
-WORKDIR ${INVOKEAI_SRC}
+# Switch to the `ubuntu` user to work around dependency issues with uv-installed python
+RUN mkdir -p ${VIRTUAL_ENV} && \
+    mkdir -p ${INVOKEAI_SRC} && \
+    chmod -R a+w /opt
+USER ubuntu
 
+# Install python and create the venv
+RUN uv python install ${PYTHON_VERSION} && \
+    uv venv --relocatable --prompt "invoke" --python ${PYTHON_VERSION} ${VIRTUAL_ENV}
+
+WORKDIR ${INVOKEAI_SRC}
 COPY invokeai ./invokeai
 COPY pyproject.toml ./
 
@@ -32,25 +45,18 @@ COPY pyproject.toml ./
 # the local working copy can be bind-mounted into the image
 # at path defined by ${INVOKEAI_SRC}
 # NOTE: there are no pytorch builds for arm64 + cuda, only cpu
-# x86_64/CUDA is default
-RUN --mount=type=cache,target=/root/.cache/pip \
-    python3 -m venv ${VIRTUAL_ENV} &&\
+# x86_64/CUDA is the default
+RUN --mount=type=cache,target=/home/ubuntu/.cache/uv,uid=1000,gid=1000 \
     if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then \
         extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cpu"; \
     elif [ "$GPU_DRIVER" = "rocm" ]; then \
         extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/rocm6.1"; \
     else \
         extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cu124"; \
-    fi &&\
+    fi && \
+    uv pip install --python ${PYTHON_VERSION} $extra_index_url_arg -e "."
 
-    # xformers + triton fails to install on arm64
-    if [ "$GPU_DRIVER" = "cuda" ] && [ "$TARGETPLATFORM" = "linux/amd64" ]; then \
-        pip install $extra_index_url_arg -e ".[xformers]"; \
-    else \
-        pip install $extra_index_url_arg -e "."; \
-    fi
-
-# #### Build the Web UI ------------------------------------
+#### Build the Web UI ------------------------------------
 
 FROM node:20-slim AS web-builder
 ENV PNPM_HOME="/pnpm"
@@ -66,7 +72,7 @@ RUN npx vite build
 
 #### Runtime stage ---------------------------------------
 
-FROM library/ubuntu:23.04 AS runtime
+FROM library/ubuntu:24.04 AS runtime
 
 ARG DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
@@ -83,17 +89,16 @@ RUN apt update && apt install -y --no-install-recommends \
         gosu \
         magic-wormhole \
         libglib2.0-0 \
-        libgl1-mesa-glx \
-        python3-venv \
-        python3-pip \
+        libgl1 \
+        libglx-mesa0 \
         build-essential \
         libopencv-dev \
         libstdc++-10-dev &&\
     apt-get clean && apt-get autoclean
 
-
 ENV INVOKEAI_SRC=/opt/invokeai
-ENV VIRTUAL_ENV=/opt/venv/invokeai
+ENV VIRTUAL_ENV=/opt/venv
+ENV PYTHON_VERSION=3.11
 ENV INVOKEAI_ROOT=/invokeai
 ENV INVOKEAI_HOST=0.0.0.0
 ENV INVOKEAI_PORT=9090
@@ -101,6 +106,14 @@ ENV PATH="$VIRTUAL_ENV/bin:$INVOKEAI_SRC:$PATH"
 ENV CONTAINER_UID=${CONTAINER_UID:-1000}
 ENV CONTAINER_GID=${CONTAINER_GID:-1000}
 
+# Install `uv` for package management
+# and install python for the ubuntu user (expected to exist on ubuntu >=24.x)
+# this is too tiny to optimize with multi-stage builds, but maybe we'll come back to it
+COPY --from=ghcr.io/astral-sh/uv:0.5.5 /uv /uvx /bin/
+USER ubuntu
+RUN uv python install ${PYTHON_VERSION}
+USER root
+
 # --link requires buldkit w/ dockerfile syntax 1.4
 COPY --link --from=builder ${INVOKEAI_SRC} ${INVOKEAI_SRC}
 COPY --link --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
@@ -115,7 +128,7 @@ WORKDIR ${INVOKEAI_SRC}
 
 # build patchmatch
 RUN cd /usr/lib/$(uname -p)-linux-gnu/pkgconfig/ && ln -sf opencv4.pc opencv.pc
-RUN python3 -c "from patchmatch import patch_match"
+RUN python -c "from patchmatch import patch_match"
 
 RUN mkdir -p ${INVOKEAI_ROOT} && chown -R ${CONTAINER_UID}:${CONTAINER_GID} ${INVOKEAI_ROOT}
 

docker/docker-entrypoint.sh

@@ -16,6 +16,9 @@ set -e -o pipefail
 
 USER_ID=${CONTAINER_UID:-1000}
 USER=ubuntu
+# if the user does not exist, create it. It is expected to be present on ubuntu >=24.x
+_=$(id ${USER} 2>&1) || useradd -u ${USER_ID} ${USER}
+# ensure the UID is correct
 usermod -u ${USER_ID} ${USER} 1>/dev/null
 
 ### Set the $PUBLIC_KEY env var to enable SSH access.
@@ -36,6 +39,8 @@ fi
 mkdir -p "${INVOKEAI_ROOT}"
 chown --recursive ${USER} "${INVOKEAI_ROOT}" || true
 cd "${INVOKEAI_ROOT}"
+export HF_HOME=${HF_HOME:-$INVOKEAI_ROOT/.cache/huggingface}
+export MPLCONFIGDIR=${MPLCONFIGDIR:-$INVOKEAI_ROOT/.matplotlib}
 
 # Run the CMD as the Container User (not root).
 exec gosu ${USER} "$@"

docs/contributing/ARCHITECTURE.md

@@ -50,7 +50,7 @@ Applications are built on top of the invoke framework. They should construct `in
 
 ### Web UI
 
-The Web UI is built on top of an HTTP API built with [FastAPI](https://fastapi.tiangolo.com/) and [Socket.IO](https://socket.io/). The frontend code is found in `/frontend` and the backend code is found in `/ldm/invoke/app/api_app.py` and `/ldm/invoke/app/api/`. The code is further organized as such:
+The Web UI is built on top of an HTTP API built with [FastAPI](https://fastapi.tiangolo.com/) and [Socket.IO](https://socket.io/). The frontend code is found in `/invokeai/frontend` and the backend code is found in `/invokeai/app/api_app.py` and `/invokeai/app/api/`. The code is further organized as such:
 
 | Component | Description |
 | --- | --- |
@@ -62,7 +62,7 @@ The Web UI is built on top of an HTTP API built with [FastAPI](https://fastapi.t
 
 ### CLI
 
-The CLI is built automatically from invocation metadata, and also supports invocation piping and auto-linking. Code is available in `/ldm/invoke/app/cli_app.py`.
+The CLI is built automatically from invocation metadata, and also supports invocation piping and auto-linking. Code is available in `/invokeai/frontend/cli`.
 
 ## Invoke
 
@@ -70,7 +70,7 @@ The Invoke framework provides the interface to the underlying AI systems and is
 
 ### Invoker
 
-The invoker (`/ldm/invoke/app/services/invoker.py`) is the primary interface through which applications interact with the framework. Its primary purpose is to create, manage, and invoke sessions. It also maintains two sets of services:
+The invoker (`/invokeai/app/services/invoker.py`) is the primary interface through which applications interact with the framework. Its primary purpose is to create, manage, and invoke sessions. It also maintains two sets of services:
 - **invocation services**, which are used by invocations to interact with core functionality.
 - **invoker services**, which are used by the invoker to manage sessions and manage the invocation queue.
 
@@ -82,12 +82,12 @@ The session graph does not support looping. This is left as an application probl
 
 ### Invocations
 
-Invocations represent individual units of execution, with inputs and outputs. All invocations are located in `/ldm/invoke/app/invocations`, and are all automatically discovered and made available in the applications. These are the primary way to expose new functionality in Invoke.AI, and the [implementation guide](INVOCATIONS.md) explains how to add new invocations.
+Invocations represent individual units of execution, with inputs and outputs. All invocations are located in `/invokeai/app/invocations`, and are all automatically discovered and made available in the applications. These are the primary way to expose new functionality in Invoke.AI, and the [implementation guide](INVOCATIONS.md) explains how to add new invocations.
 
 ### Services
 
-Services provide invocations access AI Core functionality and other necessary functionality (e.g. image storage). These are available in `/ldm/invoke/app/services`. As a general rule, new services should provide an interface as an abstract base class, and may provide a lightweight local implementation by default in their module. The goal for all services should be to enable the usage of different implementations (e.g. using cloud storage for image storage), but should not load any module dependencies unless that implementation has been used (i.e. don't import anything that won't be used, especially if it's expensive to import).
+Services provide invocations access AI Core functionality and other necessary functionality (e.g. image storage). These are available in `/invokeai/app/services`. As a general rule, new services should provide an interface as an abstract base class, and may provide a lightweight local implementation by default in their module. The goal for all services should be to enable the usage of different implementations (e.g. using cloud storage for image storage), but should not load any module dependencies unless that implementation has been used (i.e. don't import anything that won't be used, especially if it's expensive to import).
 
 ## AI Core
 
-The AI Core is represented by the rest of the code base (i.e. the code outside of `/ldm/invoke/app/`).
+The AI Core is represented by the rest of the code base (i.e. the code outside of `/invokeai/app/`).

docs/contributing/INVOCATIONS.md

@@ -287,8 +287,8 @@ new Invocation ready to be used.
 
 Once you've created a Node, the next step is to share it with the community! The
 best way to do this is to submit a Pull Request to add the Node to the
-[Community Nodes](nodes/communityNodes) list. If you're not sure how to do that,
-take a look a at our [contributing nodes overview](contributingNodes).
+[Community Nodes](../nodes/communityNodes.md) list. If you're not sure how to do that,
+take a look a at our [contributing nodes overview](../nodes/contributingNodes.md).
 
 ## Advanced
 

docs/contributing/MODEL_MANAGER.md

@@ -9,20 +9,20 @@ model. These are the:
   configuration information. Among other things, the record service
   tracks the type of the model, its provenance, and where it can be
   found on disk.
-  
+
 * _ModelInstallServiceBase_ A service for installing models to
   disk. It uses `DownloadQueueServiceBase` to download models and
   their metadata, and `ModelRecordServiceBase` to store that
   information. It is also responsible for managing the InvokeAI
   `models` directory and its contents.
-  
+
 * _DownloadQueueServiceBase_
   A multithreaded downloader responsible
   for downloading models from a remote source to disk. The download
   queue has special methods for downloading repo_id folders from
   Hugging Face, as well as discriminating among model versions in
   Civitai, but can be used for arbitrary content.
-  
+
   * _ModelLoadServiceBase_
   Responsible for loading a model from disk
   into RAM and VRAM and getting it ready for inference.
@@ -207,9 +207,9 @@ for use in the InvokeAI web server. Its signature is:
 
 ```
 def open(
-       cls, 
-    config: InvokeAIAppConfig, 
-    conn: Optional[sqlite3.Connection] = None, 
+       cls,
+    config: InvokeAIAppConfig,
+    conn: Optional[sqlite3.Connection] = None,
     lock: Optional[threading.Lock] = None
     ) -> Union[ModelRecordServiceSQL, ModelRecordServiceFile]:
 ```
@@ -363,7 +363,7 @@ functionality:
 
 * Registering a model config record for a model already located on the
   local filesystem, without moving it or changing its path.
-  
+
 * Installing a model alreadiy located on the local filesystem, by
   moving it into the InvokeAI root directory under the
   `models` folder (or wherever config parameter `models_dir`
@@ -371,21 +371,21 @@ functionality:
 
 * Probing of models to determine their type, base type and other key
   information.
-  
+
 * Interface with the InvokeAI event bus to provide status updates on
   the download, installation and registration process.
-  
+
 * Downloading a model from an arbitrary URL and installing it in
   `models_dir`.
 
 * Special handling for HuggingFace repo_ids to recursively download
   the contents of the repository, paying attention to alternative
   variants such as fp16.
-  
+
 * Saving tags and other metadata about the model into the invokeai database
   when fetching from a repo that provides that type of information,
   (currently only HuggingFace).
-  
+
 ### Initializing the installer
 
 A default installer is created at InvokeAI api startup time and stored
@@ -461,7 +461,7 @@ revision.
 `config` is an optional dict of values that will override the
 autoprobed values for model type, base, scheduler prediction type, and
 so forth. See [Model configuration and
-probing](#Model-configuration-and-probing) for details.
+probing](#model-configuration-and-probing) for details.
 
 `access_token` is an optional access token for accessing resources
 that need authentication.
@@ -494,7 +494,7 @@ source8 = URLModelSource(url='https://civitai.com/api/download/models/63006', ac
 
 for source in [source1, source2, source3, source4, source5, source6, source7]:
    install_job = installer.install_model(source)
-   
+
 source2job = installer.wait_for_installs(timeout=120)
 for source in sources:
     job = source2job[source]
@@ -504,7 +504,7 @@ for source in sources:
   print(f"{source} installed as {model_key}")
  elif job.errored:
      print(f"{source}: {job.error_type}.\nStack trace:\n{job.error}")
- 
+
 ```
 
 As shown here, the `import_model()` method accepts a variety of
@@ -1364,7 +1364,6 @@ 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/TESTS.md

@@ -1,6 +1,6 @@
 # InvokeAI Backend Tests
 
-We use `pytest` to run the backend python tests. (See [pyproject.toml](/pyproject.toml) for the default `pytest` options.)
+We use `pytest` to run the backend python tests. (See [pyproject.toml](https://github.com/invoke-ai/InvokeAI/blob/main/pyproject.toml) for the default `pytest` options.)
 
 ## Fast vs. Slow
 All tests are categorized as either 'fast' (no test annotation) or 'slow' (annotated with the `@pytest.mark.slow` decorator).
@@ -33,7 +33,7 @@ pytest tests -m ""
 
 ## Test Organization
 
-All backend tests are in the [`tests/`](/tests/) directory. This directory mirrors the organization of the `invokeai/` directory. For example, tests for `invokeai/model_management/model_manager.py` would be found in `tests/model_management/test_model_manager.py`.
+All backend tests are in the [`tests/`](https://github.com/invoke-ai/InvokeAI/tree/main/tests) directory. This directory mirrors the organization of the `invokeai/` directory. For example, tests for `invokeai/model_management/model_manager.py` would be found in `tests/model_management/test_model_manager.py`.
 
 TODO: The above statement is aspirational. A re-organization of legacy tests is required to make it true.
 

docs/contributing/contribution_guides/development.md

@@ -2,7 +2,7 @@
 
 ## **What do I need to know to help?**
 
-If you are looking to help with a code contribution, InvokeAI uses several different technologies under the hood: Python (Pydantic, FastAPI, diffusers) and Typescript (React, Redux Toolkit, ChakraUI, Mantine, Konva). Familiarity with StableDiffusion and image generation concepts is helpful, but not essential. 
+If you are looking to help with a code contribution, InvokeAI uses several different technologies under the hood: Python (Pydantic, FastAPI, diffusers) and Typescript (React, Redux Toolkit, ChakraUI, Mantine, Konva). Familiarity with StableDiffusion and image generation concepts is helpful, but not essential.
 
 
 ## **Get Started**
@@ -12,7 +12,7 @@ To get started, take a look at our [new contributors checklist](newContributorCh
 Once you're setup, for more information, you can review the documentation specific to your area of interest:
 
 * #### [InvokeAI Architecure](../ARCHITECTURE.md)
-* #### [Frontend Documentation](https://github.com/invoke-ai/InvokeAI/tree/main/invokeai/frontend/web)
+* #### [Frontend Documentation](../frontend/index.md)
 * #### [Node Documentation](../INVOCATIONS.md)
 * #### [Local Development](../LOCAL_DEVELOPMENT.md)
 
@@ -20,15 +20,15 @@ Once you're setup, for more information, you can review the documentation specif
 
 If you don't feel ready to make a code contribution yet, no problem! You can also help out in other ways, such as [documentation](documentation.md), [translation](translation.md) or helping support other users and triage issues as they're reported in GitHub.
 
-There are two paths to making a development contribution: 
+There are two paths to making a development contribution:
 
 1. Choosing an open issue to address. Open issues can be found in the [Issues](https://github.com/invoke-ai/InvokeAI/issues?q=is%3Aissue+is%3Aopen) section of the InvokeAI repository. These are tagged by the issue type (bug, enhancement, etc.) along with the “good first issues” tag denoting if they are suitable for first time contributors.
-    1. Additional items can be found on our [roadmap](https://github.com/orgs/invoke-ai/projects/7). The roadmap is organized in terms of priority, and contains features of varying size and complexity. If there is an inflight item you’d like to help with, reach out to the contributor assigned to the item to see how you can help. 
+    1. Additional items can be found on our [roadmap](https://github.com/orgs/invoke-ai/projects/7). The roadmap is organized in terms of priority, and contains features of varying size and complexity. If there is an inflight item you’d like to help with, reach out to the contributor assigned to the item to see how you can help.
 2. Opening a new issue or feature to add. **Please make sure you have searched through existing issues before creating new ones.**
 
 *Regardless of what you choose, please post in the  [#dev-chat](https://discord.com/channels/1020123559063990373/1049495067846524939) channel of the Discord before you start development in order to confirm that the issue or feature is aligned with the current direction of the project. We value our contributors time and effort and want to ensure that no one’s time is being misspent.*
 
-## Best Practices: 
+## Best Practices:
 * Keep your pull requests small. Smaller pull requests are more likely to be accepted and merged
 * Comments! Commenting your code helps reviewers easily understand your contribution
 * Use Python and Typescript’s typing systems, and consider using an editor with [LSP](https://microsoft.github.io/language-server-protocol/) support to streamline development
@@ -38,7 +38,7 @@ There are two paths to making a development contribution:
 
 If you need help, you can ask questions in the [#dev-chat](https://discord.com/channels/1020123559063990373/1049495067846524939) channel of the Discord.
 
-For frontend related work, **@psychedelicious** is the best person to reach out to. 
+For frontend related work, **@psychedelicious** is the best person to reach out to.
 
 For backend related work, please reach out to **@blessedcoolant**, **@lstein**, **@StAlKeR7779** or **@psychedelicious**.
 

docs/contributing/contribution_guides/newContributorChecklist.md

@@ -22,15 +22,15 @@ Before starting these steps, ensure you have your local environment [configured
 2. Fork the [InvokeAI](https://github.com/invoke-ai/InvokeAI) repository to your GitHub profile. This means that you will have a copy of the repository under **your-GitHub-username/InvokeAI**.
 3. Clone the repository to your local machine using:
 
-   ```bash
-   git clone https://github.com/your-GitHub-username/InvokeAI.git
-   ```
+    ```bash
+    git clone https://github.com/your-GitHub-username/InvokeAI.git
+    ```
 
 If you're unfamiliar with using Git through the commandline, [GitHub Desktop](https://desktop.github.com) is a easy-to-use alternative with a UI. You can do all the same steps listed here, but through the interface. 4. Create a new branch for your fix using:
 
-    ```bash
-    git checkout -b branch-name-here
-    ```
+  ```bash
+  git checkout -b branch-name-here
+  ```
 
 5. Make the appropriate changes for the issue you are trying to address or the feature that you want to add.
 6. Add the file contents of the changed files to the "snapshot" git uses to manage the state of the project, also known as the index:

docs/contributing/dev-environment.md

@@ -27,9 +27,9 @@ If you just want to use Invoke, you should use the [installer][installer link].
 
 5. Activate the venv (you'll need to do this every time you want to run the app):
 
-        ```sh
-        source .venv/bin/activate
-        ```
+      ```sh
+      source .venv/bin/activate
+      ```
 
 6. Install the repo as an [editable install][editable install link]:
 
@@ -37,7 +37,7 @@ If you just want to use Invoke, you should use the [installer][installer link].
       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.
+      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:
 

docs/contributing/index.md

@@ -34,11 +34,11 @@ Please reach out to @hipsterusername on [Discord](https://discord.gg/ZmtBAhwWhy)
 
 ## Contributors
 
-This project is a combined effort of dedicated people from across the world. [Check out the list of all these amazing people](https://invoke-ai.github.io/InvokeAI/other/CONTRIBUTORS/). We thank them for their time, hard work and effort.
+This project is a combined effort of dedicated people from across the world. [Check out the list of all these amazing people](contributors.md). We thank them for their time, hard work and effort.
 
 ## Code of Conduct
 
-The InvokeAI community is a welcoming place, and we want your help in maintaining that. Please review our [Code of Conduct](https://github.com/invoke-ai/InvokeAI/blob/main/CODE_OF_CONDUCT.md) to learn more - it's essential to maintaining a respectful and inclusive environment.
+The InvokeAI community is a welcoming place, and we want your help in maintaining that. Please review our [Code of Conduct](../CODE_OF_CONDUCT.md) to learn more - it's essential to maintaining a respectful and inclusive environment.
 
 By making a contribution to this project, you certify that:
 

docs/nodes/defaultNodes.md

@@ -99,7 +99,6 @@ their descriptions.
 | Scale Latents                                                 | Scales latents by a given factor.                                                                                                                    |
 | Segment Anything Processor                                    | Applies segment anything processing to image                                                                                                         |
 | Show Image                                                    | Displays a provided image, and passes it forward in the pipeline.                                                                                    |
-| Step Param Easing                                             | Experimental per-step parameter easing for denoising steps                                                                                           |
 | String Primitive Collection                                   | A collection of string primitive values                                                                                                              |
 | String Primitive                                              | A string primitive value                                                                                                                             |
 | Subtract Integers                                             | Subtracts two numbers                                                                                                                                |

invokeai/app/api/routers/model_manager.py

@@ -4,7 +4,6 @@
 import contextlib
 import io
 import pathlib
-import shutil
 import traceback
 from copy import deepcopy
 from enum import Enum
@@ -21,7 +20,6 @@ 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 (
@@ -37,7 +35,7 @@ from invokeai.backend.model_manager.config import (
     ModelFormat,
     ModelType,
 )
-from invokeai.backend.model_manager.load.model_cache.model_cache_base import CacheStats
+from invokeai.backend.model_manager.load.model_cache.cache_stats 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
@@ -848,74 +846,6 @@ 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/routers/session_queue.py

@@ -110,7 +110,7 @@ async def cancel_by_batch_ids(
 @session_queue_router.put(
     "/{queue_id}/cancel_by_destination",
     operation_id="cancel_by_destination",
-    responses={200: {"model": CancelByBatchIDsResult}},
+    responses={200: {"model": CancelByDestinationResult}},
 )
 async def cancel_by_destination(
     queue_id: str = Path(description="The queue id to perform this operation on"),

invokeai/app/invocations/__init__.py

@@ -15,6 +15,11 @@ custom_nodes_readme_path = str(custom_nodes_path / "README.md")
 shutil.copy(Path(__file__).parent / "custom_nodes/init.py", custom_nodes_init_path)
 shutil.copy(Path(__file__).parent / "custom_nodes/README.md", custom_nodes_readme_path)
 
+# set the same permissions as the destination directory, in case our source is read-only,
+# so that the files are user-writable
+for p in custom_nodes_path.glob("**/*"):
+    p.chmod(custom_nodes_path.stat().st_mode)
+
 # Import custom nodes, see https://docs.python.org/3/library/importlib.html#importing-programmatically
 spec = spec_from_file_location("custom_nodes", custom_nodes_init_path)
 if spec is None or spec.loader is None:

invokeai/app/invocations/baseinvocation.py

@@ -63,6 +63,7 @@ class Classification(str, Enum, metaclass=MetaEnum):
     - `Prototype`: The invocation is not yet stable and may be removed from the application at any time. Workflows built around this invocation may break, and we are *not* committed to supporting this invocation.
     - `Deprecated`: The invocation is deprecated and may be removed in a future version.
     - `Internal`: The invocation is not intended for use by end-users. It may be changed or removed at any time, but is exposed for users to play with.
+    - `Special`: The invocation is a special case and does not fit into any of the other classifications.
     """
 
     Stable = "stable"
@@ -70,6 +71,7 @@ class Classification(str, Enum, metaclass=MetaEnum):
     Prototype = "prototype"
     Deprecated = "deprecated"
     Internal = "internal"
+    Special = "special"
 
 
 class UIConfigBase(BaseModel):

invokeai/app/invocations/blend_latents.py

@@ -1,98 +1,120 @@
-from typing import Any, Union
+from typing import Optional, Union
 
 import numpy as np
-import numpy.typing as npt
 import torch
+import torchvision.transforms as T
+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.fields import FieldDescriptions, Input, InputField, LatentsField
+from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField, LatentsField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 from invokeai.backend.util.devices import TorchDevice
 
 
+def slerp(
+    t: Union[float, np.ndarray],
+    v0: Union[torch.Tensor, np.ndarray],
+    v1: Union[torch.Tensor, np.ndarray],
+    device: torch.device,
+    DOT_THRESHOLD: float = 0.9995,
+):
+    """
+    Spherical linear interpolation
+    Args:
+        t (float/np.ndarray): Float value between 0.0 and 1.0
+        v0 (np.ndarray): Starting vector
+        v1 (np.ndarray): Final vector
+        DOT_THRESHOLD (float): Threshold for considering the two vectors as
+                            colineal. Not recommended to alter this.
+    Returns:
+        v2 (np.ndarray): Interpolation vector between v0 and v1
+    """
+    inputs_are_torch = False
+    if not isinstance(v0, np.ndarray):
+        inputs_are_torch = True
+        v0 = v0.detach().cpu().numpy()
+    if not isinstance(v1, np.ndarray):
+        inputs_are_torch = True
+        v1 = v1.detach().cpu().numpy()
+
+    dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
+    if np.abs(dot) > DOT_THRESHOLD:
+        v2 = (1 - t) * v0 + t * v1
+    else:
+        theta_0 = np.arccos(dot)
+        sin_theta_0 = np.sin(theta_0)
+        theta_t = theta_0 * t
+        sin_theta_t = np.sin(theta_t)
+        s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+        s1 = sin_theta_t / sin_theta_0
+        v2 = s0 * v0 + s1 * v1
+
+    if inputs_are_torch:
+        v2 = torch.from_numpy(v2).to(device)
+
+    return v2
+
+
 @invocation(
     "lblend",
     title="Blend Latents",
-    tags=["latents", "blend"],
+    tags=["latents", "blend", "mask"],
     category="latents",
-    version="1.0.3",
+    version="1.1.0",
 )
 class BlendLatentsInvocation(BaseInvocation):
-    """Blend two latents using a given alpha. Latents must have same size."""
+    """Blend two latents using a given alpha. If a mask is provided, the second latents will be masked before blending.
+    Latents must have same size. Masking functionality added by @dwringer."""
 
-    latents_a: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    latents_b: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    alpha: float = InputField(default=0.5, description=FieldDescriptions.blend_alpha)
+    latents_a: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection)
+    latents_b: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection)
+    mask: Optional[ImageField] = InputField(default=None, description="Mask for blending in latents B")
+    alpha: float = InputField(ge=0, default=0.5, description=FieldDescriptions.blend_alpha)
+
+    def prep_mask_tensor(self, mask_image: Image.Image) -> torch.Tensor:
+        if mask_image.mode != "L":
+            mask_image = mask_image.convert("L")
+        mask_tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
+        if mask_tensor.dim() == 3:
+            mask_tensor = mask_tensor.unsqueeze(0)
+        return mask_tensor
+
+    def replace_tensor_from_masked_tensor(
+        self, tensor: torch.Tensor, other_tensor: torch.Tensor, mask_tensor: torch.Tensor
+    ):
+        output = tensor.clone()
+        mask_tensor = mask_tensor.expand(output.shape)
+        if output.dtype != torch.float16:
+            output = torch.add(output, mask_tensor * torch.sub(other_tensor, tensor))
+        else:
+            output = torch.add(output, mask_tensor.half() * torch.sub(other_tensor, tensor))
+        return output
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         latents_a = context.tensors.load(self.latents_a.latents_name)
         latents_b = context.tensors.load(self.latents_b.latents_name)
+        if self.mask is None:
+            mask_tensor = torch.zeros(latents_a.shape[-2:])
+        else:
+            mask_tensor = self.prep_mask_tensor(context.images.get_pil(self.mask.image_name))
+            mask_tensor = tv_resize(mask_tensor, latents_a.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
+
+        latents_b = self.replace_tensor_from_masked_tensor(latents_b, latents_a, mask_tensor)
 
         if latents_a.shape != latents_b.shape:
-            raise Exception("Latents to blend must be the same size.")
+            raise ValueError("Latents to blend must be the same size.")
 
         device = TorchDevice.choose_torch_device()
 
-        def slerp(
-            t: Union[float, npt.NDArray[Any]],  # FIXME: maybe use np.float32 here?
-            v0: Union[torch.Tensor, npt.NDArray[Any]],
-            v1: Union[torch.Tensor, npt.NDArray[Any]],
-            DOT_THRESHOLD: float = 0.9995,
-        ) -> Union[torch.Tensor, npt.NDArray[Any]]:
-            """
-            Spherical linear interpolation
-            Args:
-                t (float/np.ndarray): Float value between 0.0 and 1.0
-                v0 (np.ndarray): Starting vector
-                v1 (np.ndarray): Final vector
-                DOT_THRESHOLD (float): Threshold for considering the two vectors as
-                                    colineal. Not recommended to alter this.
-            Returns:
-                v2 (np.ndarray): Interpolation vector between v0 and v1
-            """
-            inputs_are_torch = False
-            if not isinstance(v0, np.ndarray):
-                inputs_are_torch = True
-                v0 = v0.detach().cpu().numpy()
-            if not isinstance(v1, np.ndarray):
-                inputs_are_torch = True
-                v1 = v1.detach().cpu().numpy()
-
-            dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
-            if np.abs(dot) > DOT_THRESHOLD:
-                v2 = (1 - t) * v0 + t * v1
-            else:
-                theta_0 = np.arccos(dot)
-                sin_theta_0 = np.sin(theta_0)
-                theta_t = theta_0 * t
-                sin_theta_t = np.sin(theta_t)
-                s0 = np.sin(theta_0 - theta_t) / sin_theta_0
-                s1 = sin_theta_t / sin_theta_0
-                v2 = s0 * v0 + s1 * v1
-
-            if inputs_are_torch:
-                v2_torch: torch.Tensor = torch.from_numpy(v2).to(device)
-                return v2_torch
-            else:
-                assert isinstance(v2, np.ndarray)
-                return v2
-
         # blend
-        bl = slerp(self.alpha, latents_a, latents_b)
-        assert isinstance(bl, torch.Tensor)
-        blended_latents: torch.Tensor = bl  # for type checking convenience
+        blended_latents = slerp(self.alpha, latents_a, latents_b, device)
 
         # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
         blended_latents = blended_latents.to("cpu")
-
-        TorchDevice.empty_cache()
+        torch.cuda.empty_cache()
 
         name = context.tensors.save(tensor=blended_latents)
-        return LatentsOutput.build(latents_name=name, latents=blended_latents, seed=self.latents_a.seed)
+        return LatentsOutput.build(latents_name=name, latents=blended_latents)

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,
@@ -66,10 +66,10 @@ class CompelInvocation(BaseInvocation):
         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]]:
+        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
             for lora in self.clip.loras:
                 lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, LoRAModelRaw)
+                assert isinstance(lora_info.model, ModelPatchRaw)
                 yield (lora_info.model, lora.weight)
                 del lora_info
             return
@@ -82,10 +82,11 @@ class CompelInvocation(BaseInvocation):
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
             tokenizer_info as tokenizer,
-            LoRAPatcher.apply_lora_patches(
+            LayerPatcher.apply_smart_model_patches(
                 model=text_encoder,
                 patches=_lora_loader(),
                 prefix="lora_te_",
+                dtype=TorchDevice.choose_torch_dtype(),
                 cached_weights=cached_weights,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
@@ -162,11 +163,11 @@ class SDXLPromptInvocationBase:
                 c_pooled = None
             return c, c_pooled
 
-        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
+        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
             for lora in clip_field.loras:
                 lora_info = context.models.load(lora.lora)
                 lora_model = lora_info.model
-                assert isinstance(lora_model, LoRAModelRaw)
+                assert isinstance(lora_model, ModelPatchRaw)
                 yield (lora_model, lora.weight)
                 del lora_info
             return
@@ -179,10 +180,11 @@ class SDXLPromptInvocationBase:
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
             tokenizer_info as tokenizer,
-            LoRAPatcher.apply_lora_patches(
-                text_encoder,
+            LayerPatcher.apply_smart_model_patches(
+                model=text_encoder,
                 patches=_lora_loader(),
                 prefix=lora_prefix,
+                dtype=TorchDevice.choose_torch_dtype(),
                 cached_weights=cached_weights,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.

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 (
@@ -987,10 +987,10 @@ class DenoiseLatentsInvocation(BaseInvocation):
         def step_callback(state: PipelineIntermediateState) -> None:
             context.util.sd_step_callback(state, unet_config.base)
 
-        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
+        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
             for lora in self.unet.loras:
                 lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, LoRAModelRaw)
+                assert isinstance(lora_info.model, ModelPatchRaw)
                 yield (lora_info.model, lora.weight)
                 del lora_info
             return
@@ -1003,10 +1003,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
             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.
-            LoRAPatcher.apply_lora_patches(
+            LayerPatcher.apply_smart_model_patches(
                 model=unet,
                 patches=_lora_loader(),
                 prefix="lora_unet_",
+                dtype=unet.dtype,
                 cached_weights=cached_weights,
             ),
         ):

invokeai/app/invocations/fields.py

@@ -56,6 +56,7 @@ class UIType(str, Enum, metaclass=MetaEnum):
     CLIPLEmbedModel = "CLIPLEmbedModelField"
     CLIPGEmbedModel = "CLIPGEmbedModelField"
     SpandrelImageToImageModel = "SpandrelImageToImageModelField"
+    ControlLoRAModel = "ControlLoRAModelField"
     # endregion
 
     # region Misc Field Types
@@ -143,6 +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"
     main_model = "Main model (UNet, VAE, CLIP) to load"
     flux_model = "Flux model (Transformer) to load"
     sd3_model = "SD3 model (MMDiTX) to load"
@@ -250,6 +252,11 @@ class FluxConditioningField(BaseModel):
     """A conditioning tensor primitive value"""
 
     conditioning_name: str = Field(description="The name of conditioning tensor")
+    mask: Optional[TensorField] = Field(
+        default=None,
+        description="The mask associated with this conditioning tensor. Excluded regions should be set to False, "
+        "included regions should be set to True.",
+    )
 
 
 class SD3ConditioningField(BaseModel):

invokeai/app/invocations/flux_control_lora_loader.py

@@ -0,0 +1,49 @@
+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,10 +1,12 @@
 from contextlib import ExitStack
-from typing import Callable, Iterator, Optional, Tuple
+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
 
@@ -21,8 +23,9 @@ 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 TransformerField, VAEField
+from invokeai.app.invocations.model import ControlLoRAField, LoRAField, TransformerField, VAEField
 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
@@ -30,6 +33,7 @@ from invokeai.backend.flux.controlnet.xlabs_controlnet_flux import XLabsControlN
 from invokeai.backend.flux.denoise import denoise
 from invokeai.backend.flux.extensions.inpaint_extension import InpaintExtension
 from invokeai.backend.flux.extensions.instantx_controlnet_extension import InstantXControlNetExtension
+from invokeai.backend.flux.extensions.regional_prompting_extension import RegionalPromptingExtension
 from invokeai.backend.flux.extensions.xlabs_controlnet_extension import XLabsControlNetExtension
 from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
 from invokeai.backend.flux.ip_adapter.xlabs_ip_adapter_flux import XlabsIpAdapterFlux
@@ -42,10 +46,11 @@ from invokeai.backend.flux.sampling_utils import (
     pack,
     unpack,
 )
-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.flux.text_conditioning import FluxTextConditioning
 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
@@ -56,7 +61,7 @@ from invokeai.backend.util.devices import TorchDevice
     title="FLUX Denoise",
     tags=["image", "flux"],
     category="image",
-    version="3.2.1",
+    version="3.2.2",
     classification=Classification.Prototype,
 )
 class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
@@ -87,10 +92,13 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         input=Input.Connection,
         title="Transformer",
     )
-    positive_text_conditioning: FluxConditioningField = InputField(
+    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
     )
-    negative_text_conditioning: FluxConditioningField | None = InputField(
+    negative_text_conditioning: FluxConditioningField | list[FluxConditioningField] | None = InputField(
         default=None,
         description="Negative conditioning tensor. Can be None if cfg_scale is 1.0.",
         input=Input.Connection,
@@ -139,36 +147,12 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         name = context.tensors.save(tensor=latents)
         return LatentsOutput.build(latents_name=name, latents=latents, seed=None)
 
-    def _load_text_conditioning(
-        self, context: InvocationContext, conditioning_name: str, dtype: torch.dtype
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # Load the conditioning data.
-        cond_data = context.conditioning.load(conditioning_name)
-        assert len(cond_data.conditionings) == 1
-        flux_conditioning = cond_data.conditionings[0]
-        assert isinstance(flux_conditioning, FLUXConditioningInfo)
-        flux_conditioning = flux_conditioning.to(dtype=dtype)
-        t5_embeddings = flux_conditioning.t5_embeds
-        clip_embeddings = flux_conditioning.clip_embeds
-        return t5_embeddings, clip_embeddings
-
     def _run_diffusion(
         self,
         context: InvocationContext,
     ):
         inference_dtype = torch.bfloat16
 
-        # Load the conditioning data.
-        pos_t5_embeddings, pos_clip_embeddings = self._load_text_conditioning(
-            context, self.positive_text_conditioning.conditioning_name, inference_dtype
-        )
-        neg_t5_embeddings: torch.Tensor | None = None
-        neg_clip_embeddings: torch.Tensor | None = None
-        if self.negative_text_conditioning is not None:
-            neg_t5_embeddings, neg_clip_embeddings = self._load_text_conditioning(
-                context, self.negative_text_conditioning.conditioning_name, inference_dtype
-            )
-
         # Load the input latents, if provided.
         init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None
         if init_latents is not None:
@@ -183,15 +167,45 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
             dtype=inference_dtype,
             seed=self.seed,
         )
+        b, _c, latent_h, latent_w = noise.shape
+        packed_h = latent_h // 2
+        packed_w = latent_w // 2
+
+        # Load the conditioning data.
+        pos_text_conditionings = self._load_text_conditioning(
+            context=context,
+            cond_field=self.positive_text_conditioning,
+            packed_height=packed_h,
+            packed_width=packed_w,
+            dtype=inference_dtype,
+            device=TorchDevice.choose_torch_device(),
+        )
+        neg_text_conditionings: list[FluxTextConditioning] | None = None
+        if self.negative_text_conditioning is not None:
+            neg_text_conditionings = self._load_text_conditioning(
+                context=context,
+                cond_field=self.negative_text_conditioning,
+                packed_height=packed_h,
+                packed_width=packed_w,
+                dtype=inference_dtype,
+                device=TorchDevice.choose_torch_device(),
+            )
+        pos_regional_prompting_extension = RegionalPromptingExtension.from_text_conditioning(
+            pos_text_conditionings, img_seq_len=packed_h * packed_w
+        )
+        neg_regional_prompting_extension = (
+            RegionalPromptingExtension.from_text_conditioning(neg_text_conditionings, img_seq_len=packed_h * packed_w)
+            if neg_text_conditionings
+            else None
+        )
 
         transformer_info = context.models.load(self.transformer.transformer)
-        is_schnell = "schnell" in transformer_info.config.config_path
+        is_schnell = "schnell" in getattr(transformer_info.config, "config_path", "")
 
         # Calculate the timestep schedule.
-        image_seq_len = noise.shape[-1] * noise.shape[-2] // 4
         timesteps = get_schedule(
             num_steps=self.num_steps,
-            image_seq_len=image_seq_len,
+            image_seq_len=packed_h * packed_w,
             shift=not is_schnell,
         )
 
@@ -226,30 +240,23 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         if len(timesteps) <= 1:
             return x
 
+        # 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)
 
-        b, _c, latent_h, latent_w = x.shape
         img_ids = generate_img_ids(h=latent_h, w=latent_w, batch_size=b, device=x.device, dtype=x.dtype)
 
-        pos_bs, pos_t5_seq_len, _ = pos_t5_embeddings.shape
-        pos_txt_ids = torch.zeros(
-            pos_bs, pos_t5_seq_len, 3, dtype=inference_dtype, device=TorchDevice.choose_torch_device()
-        )
-        neg_txt_ids: torch.Tensor | None = None
-        if neg_t5_embeddings is not None:
-            neg_bs, neg_t5_seq_len, _ = neg_t5_embeddings.shape
-            neg_txt_ids = torch.zeros(
-                neg_bs, neg_t5_seq_len, 3, dtype=inference_dtype, device=TorchDevice.choose_torch_device()
-            )
-
         # 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)
 
-        # Now that we have 'packed' the latent tensors, verify that we calculated the image_seq_len correctly.
-        assert image_seq_len == x.shape[1]
+        # Now that we have 'packed' the latent tensors, verify that we calculated the image_seq_len, packed_h, and
+        # packed_w correctly.
+        assert packed_h * packed_w == x.shape[1]
 
         # Prepare inpaint extension.
         inpaint_extension: InpaintExtension | None = None
@@ -294,36 +301,33 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
             config = transformer_info.config
             assert config is not None
 
-            # Apply LoRA models to the transformer.
-            # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching.
+            # 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.
             if config.format in [ModelFormat.Checkpoint]:
-                # The model is non-quantized, so we can apply the LoRA weights directly into the model.
-                exit_stack.enter_context(
-                    LoRAPatcher.apply_lora_patches(
-                        model=transformer,
-                        patches=self._lora_iterator(context),
-                        prefix=FLUX_LORA_TRANSFORMER_PREFIX,
-                        cached_weights=cached_weights,
-                    )
-                )
+                model_is_quantized = False
             elif config.format in [
                 ModelFormat.BnbQuantizedLlmInt8b,
                 ModelFormat.BnbQuantizednf4b,
                 ModelFormat.GGUFQuantized,
             ]:
-                # 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_sidecar_patches(
-                        model=transformer,
-                        patches=self._lora_iterator(context),
-                        prefix=FLUX_LORA_TRANSFORMER_PREFIX,
-                        dtype=inference_dtype,
-                    )
-                )
+                model_is_quantized = True
             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,
@@ -338,12 +342,8 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 model=transformer,
                 img=x,
                 img_ids=img_ids,
-                txt=pos_t5_embeddings,
-                txt_ids=pos_txt_ids,
-                vec=pos_clip_embeddings,
-                neg_txt=neg_t5_embeddings,
-                neg_txt_ids=neg_txt_ids,
-                neg_vec=neg_clip_embeddings,
+                pos_regional_prompting_extension=pos_regional_prompting_extension,
+                neg_regional_prompting_extension=neg_regional_prompting_extension,
                 timesteps=timesteps,
                 step_callback=self._build_step_callback(context),
                 guidance=self.guidance,
@@ -352,11 +352,49 @@ 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,
             )
 
         x = unpack(x.float(), self.height, self.width)
         return x
 
+    def _load_text_conditioning(
+        self,
+        context: InvocationContext,
+        cond_field: FluxConditioningField | list[FluxConditioningField],
+        packed_height: int,
+        packed_width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> list[FluxTextConditioning]:
+        """Load text conditioning data from a FluxConditioningField or a list of FluxConditioningFields."""
+        # Normalize to a list of FluxConditioningFields.
+        cond_list = [cond_field] if isinstance(cond_field, FluxConditioningField) else cond_field
+
+        text_conditionings: list[FluxTextConditioning] = []
+        for cond_field in cond_list:
+            # Load the text embeddings.
+            cond_data = context.conditioning.load(cond_field.conditioning_name)
+            assert len(cond_data.conditionings) == 1
+            flux_conditioning = cond_data.conditionings[0]
+            assert isinstance(flux_conditioning, FLUXConditioningInfo)
+            flux_conditioning = flux_conditioning.to(dtype=dtype, device=device)
+            t5_embeddings = flux_conditioning.t5_embeds
+            clip_embeddings = flux_conditioning.clip_embeds
+
+            # Load the mask, if provided.
+            mask: Optional[torch.Tensor] = None
+            if cond_field.mask is not None:
+                mask = context.tensors.load(cond_field.mask.tensor_name)
+                mask = mask.to(device=device)
+                mask = RegionalPromptingExtension.preprocess_regional_prompt_mask(
+                    mask, packed_height, packed_width, dtype, device
+                )
+
+            text_conditionings.append(FluxTextConditioning(t5_embeddings, clip_embeddings, mask))
+
+        return text_conditionings
+
     @classmethod
     def prep_cfg_scale(
         cls, cfg_scale: float | list[float], timesteps: list[float], cfg_scale_start_step: int, cfg_scale_end_step: int
@@ -545,6 +583,29 @@ 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 []
@@ -651,10 +712,15 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         return pos_ip_adapter_extensions, neg_ip_adapter_extensions
 
-    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[LoRAModelRaw, float]]:
-        for lora in self.transformer.loras:
+    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)
+        for lora in loras:
             lora_info = context.models.load(lora.lora)
-            assert isinstance(lora_info.model, LoRAModelRaw)
+            assert isinstance(lora_info.model, ModelPatchRaw)
             yield (lora_info.model, lora.weight)
             del lora_info
 

invokeai/app/invocations/flux_text_encoder.py

@@ -1,20 +1,28 @@
 from contextlib import ExitStack
-from typing import Iterator, Literal, Tuple
+from typing import Iterator, Literal, Optional, Tuple
 
 import torch
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    FluxConditioningField,
+    Input,
+    InputField,
+    TensorField,
+    UIComponent,
+)
 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(
@@ -22,7 +30,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Condit
     title="FLUX Text Encoding",
     tags=["prompt", "conditioning", "flux"],
     category="conditioning",
-    version="1.1.0",
+    version="1.1.1",
     classification=Classification.Prototype,
 )
 class FluxTextEncoderInvocation(BaseInvocation):
@@ -41,7 +49,10 @@ class FluxTextEncoderInvocation(BaseInvocation):
     t5_max_seq_len: Literal[256, 512] = InputField(
         description="Max sequence length for the T5 encoder. Expected to be 256 for FLUX schnell models and 512 for FLUX dev models."
     )
-    prompt: str = InputField(description="Text prompt to encode.")
+    prompt: str = InputField(description="Text prompt to encode.", ui_component=UIComponent.Textarea)
+    mask: Optional[TensorField] = InputField(
+        default=None, description="A mask defining the region that this conditioning prompt applies to."
+    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> FluxConditioningOutput:
@@ -54,17 +65,16 @@ class FluxTextEncoderInvocation(BaseInvocation):
         )
 
         conditioning_name = context.conditioning.save(conditioning_data)
-        return FluxConditioningOutput.build(conditioning_name)
+        return FluxConditioningOutput(
+            conditioning=FluxConditioningField(conditioning_name=conditioning_name, mask=self.mask)
+        )
 
     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 (
-            t5_text_encoder_info as t5_text_encoder,
-            t5_tokenizer_info as t5_tokenizer,
+            context.models.load(self.t5_encoder.text_encoder) as t5_text_encoder,
+            context.models.load(self.t5_encoder.tokenizer) as t5_tokenizer,
         ):
             assert isinstance(t5_text_encoder, T5EncoderModel)
             assert isinstance(t5_tokenizer, T5Tokenizer)
@@ -78,14 +88,12 @@ class FluxTextEncoderInvocation(BaseInvocation):
         return prompt_embeds
 
     def _clip_encode(self, context: InvocationContext) -> torch.Tensor:
-        clip_tokenizer_info = context.models.load(self.clip.tokenizer)
-        clip_text_encoder_info = context.models.load(self.clip.text_encoder)
-
         prompt = [self.prompt]
+        clip_text_encoder_info = context.models.load(self.clip.text_encoder)
 
         with (
             clip_text_encoder_info.model_on_device() as (cached_weights, clip_text_encoder),
-            clip_tokenizer_info as clip_tokenizer,
+            context.models.load(self.clip.tokenizer) as clip_tokenizer,
             ExitStack() as exit_stack,
         ):
             assert isinstance(clip_text_encoder, CLIPTextModel)
@@ -99,10 +107,11 @@ class FluxTextEncoderInvocation(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(
-                    LoRAPatcher.apply_lora_patches(
+                    LayerPatcher.apply_smart_model_patches(
                         model=clip_text_encoder,
                         patches=self._clip_lora_iterator(context),
                         prefix=FLUX_LORA_CLIP_PREFIX,
+                        dtype=TorchDevice.choose_torch_dtype(),
                         cached_weights=cached_weights,
                     )
                 )
@@ -118,9 +127,9 @@ class FluxTextEncoderInvocation(BaseInvocation):
         assert isinstance(pooled_prompt_embeds, torch.Tensor)
         return pooled_prompt_embeds
 
-    def _clip_lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[LoRAModelRaw, float]]:
+    def _clip_lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
         for lora in self.clip.loras:
             lora_info = context.models.load(lora.lora)
-            assert isinstance(lora_info.model, LoRAModelRaw)
+            assert isinstance(lora_info.model, ModelPatchRaw)
             yield (lora_info.model, lora.weight)
             del lora_info

invokeai/app/invocations/flux_vae_decode.py

@@ -3,6 +3,7 @@ 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,
@@ -38,8 +39,22 @@ 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). This estimate is accurate for both SD1 and SDXL.
+        out_h = LATENT_SCALE_FACTOR * latents.shape[-2]
+        out_w = LATENT_SCALE_FACTOR * latents.shape[-1]
+        element_size = next(vae.parameters()).element_size()
+        # TODO(ryand): Need to tune this value, it was copied from the SD1 implementation.
+        scaling_constant = 960  # Determined experimentally.
+        working_memory = out_h * out_w * element_size * scaling_constant
+
+        return working_memory
+
     def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Image:
-        with vae_info as vae:
+        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):
             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_panels.py

@@ -0,0 +1,59 @@
+from pydantic import ValidationInfo, field_validator
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import InputField, OutputField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+
+@invocation_output("image_panel_coordinate_output")
+class ImagePanelCoordinateOutput(BaseInvocationOutput):
+    x_left: int = OutputField(description="The left x-coordinate of the panel.")
+    y_top: int = OutputField(description="The top y-coordinate of the panel.")
+    width: int = OutputField(description="The width of the panel.")
+    height: int = OutputField(description="The height of the panel.")
+
+
+@invocation(
+    "image_panel_layout",
+    title="Image Panel Layout",
+    tags=["image", "panel", "layout"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class ImagePanelLayoutInvocation(BaseInvocation):
+    """Get the coordinates of a single panel in a grid. (If the full image shape cannot be divided evenly into panels,
+    then the grid may not cover the entire image.)
+    """
+
+    width: int = InputField(description="The width of the entire grid.")
+    height: int = InputField(description="The height of the entire grid.")
+    num_cols: int = InputField(ge=1, default=1, description="The number of columns in the grid.")
+    num_rows: int = InputField(ge=1, default=1, description="The number of rows in the grid.")
+    panel_col_idx: int = InputField(ge=0, default=0, description="The column index of the panel to be processed.")
+    panel_row_idx: int = InputField(ge=0, default=0, description="The row index of the panel to be processed.")
+
+    @field_validator("panel_col_idx")
+    def validate_panel_col_idx(cls, v: int, info: ValidationInfo) -> int:
+        if v < 0 or v >= info.data["num_cols"]:
+            raise ValueError(f"panel_col_idx must be between 0 and {info.data['num_cols'] - 1}")
+        return v
+
+    @field_validator("panel_row_idx")
+    def validate_panel_row_idx(cls, v: int, info: ValidationInfo) -> int:
+        if v < 0 or v >= info.data["num_rows"]:
+            raise ValueError(f"panel_row_idx must be between 0 and {info.data['num_rows'] - 1}")
+        return v
+
+    def invoke(self, context: InvocationContext) -> ImagePanelCoordinateOutput:
+        x_left = self.panel_col_idx * (self.width // self.num_cols)
+        y_top = self.panel_row_idx * (self.height // self.num_rows)
+        width = self.width // self.num_cols
+        height = self.height // self.num_rows
+        return ImagePanelCoordinateOutput(x_left=x_left, y_top=y_top, width=width, height=height)

invokeai/app/invocations/latents_to_image.py

@@ -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.0",
+    version="1.3.1",
 )
 class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents."""
@@ -53,13 +53,32 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size)
     fp32: bool = InputField(default=DEFAULT_PRECISION == torch.float32, description=FieldDescriptions.fp32)
 
+    def _estimate_working_memory(self, latents: torch.Tensor) -> 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.
+        out_h = LATENT_SCALE_FACTOR * latents.shape[-2]
+        out_w = LATENT_SCALE_FACTOR * latents.shape[-1]
+        element_size = 4 if self.fp32 else 2
+        scaling_constant = 960  # Determined experimentally.
+        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
+
+        return 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, AutoencoderTiny))
-        with SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), vae_info as vae:
+        with (
+            SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes),
+            vae_info.model_on_device(working_mem_bytes=self._estimate_working_memory(latents)) as (_, vae),
+        ):
             context.util.signal_progress("Running VAE decoder")
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
             latents = latents.to(vae.device)

invokeai/app/invocations/model.py

@@ -10,7 +10,7 @@ from invokeai.app.invocations.baseinvocation import (
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
+from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField, OutputField, UIType
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.shared.models import FreeUConfig
 from invokeai.backend.model_manager.config import (
@@ -65,11 +65,6 @@ class CLIPField(BaseModel):
     loras: List[LoRAField] = Field(description="LoRAs to apply on model loading")
 
 
-class TransformerField(BaseModel):
-    transformer: ModelIdentifierField = Field(description="Info to load Transformer submodel")
-    loras: List[LoRAField] = Field(description="LoRAs to apply on model loading")
-
-
 class T5EncoderField(BaseModel):
     tokenizer: ModelIdentifierField = Field(description="Info to load tokenizer submodel")
     text_encoder: ModelIdentifierField = Field(description="Info to load text_encoder submodel")
@@ -80,6 +75,15 @@ class VAEField(BaseModel):
     seamless_axes: List[str] = Field(default_factory=list, description='Axes("x" and "y") to which apply seamless')
 
 
+class ControlLoRAField(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")
+
+
 @invocation_output("unet_output")
 class UNetOutput(BaseInvocationOutput):
     """Base class for invocations that output a UNet field."""

invokeai/app/invocations/param_easing.py

@@ -1,43 +1,4 @@
-import io
-from typing import Literal, Optional
-
-import matplotlib.pyplot as plt
 import numpy as np
-import PIL.Image
-from easing_functions import (
-    BackEaseIn,
-    BackEaseInOut,
-    BackEaseOut,
-    BounceEaseIn,
-    BounceEaseInOut,
-    BounceEaseOut,
-    CircularEaseIn,
-    CircularEaseInOut,
-    CircularEaseOut,
-    CubicEaseIn,
-    CubicEaseInOut,
-    CubicEaseOut,
-    ElasticEaseIn,
-    ElasticEaseInOut,
-    ElasticEaseOut,
-    ExponentialEaseIn,
-    ExponentialEaseInOut,
-    ExponentialEaseOut,
-    LinearInOut,
-    QuadEaseIn,
-    QuadEaseInOut,
-    QuadEaseOut,
-    QuarticEaseIn,
-    QuarticEaseInOut,
-    QuarticEaseOut,
-    QuinticEaseIn,
-    QuinticEaseInOut,
-    QuinticEaseOut,
-    SineEaseIn,
-    SineEaseInOut,
-    SineEaseOut,
-)
-from matplotlib.ticker import MaxNLocator
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
 from invokeai.app.invocations.fields import InputField
@@ -65,191 +26,3 @@ class FloatLinearRangeInvocation(BaseInvocation):
     def invoke(self, context: InvocationContext) -> FloatCollectionOutput:
         param_list = list(np.linspace(self.start, self.stop, self.steps))
         return FloatCollectionOutput(collection=param_list)
-
-
-EASING_FUNCTIONS_MAP = {
-    "Linear": LinearInOut,
-    "QuadIn": QuadEaseIn,
-    "QuadOut": QuadEaseOut,
-    "QuadInOut": QuadEaseInOut,
-    "CubicIn": CubicEaseIn,
-    "CubicOut": CubicEaseOut,
-    "CubicInOut": CubicEaseInOut,
-    "QuarticIn": QuarticEaseIn,
-    "QuarticOut": QuarticEaseOut,
-    "QuarticInOut": QuarticEaseInOut,
-    "QuinticIn": QuinticEaseIn,
-    "QuinticOut": QuinticEaseOut,
-    "QuinticInOut": QuinticEaseInOut,
-    "SineIn": SineEaseIn,
-    "SineOut": SineEaseOut,
-    "SineInOut": SineEaseInOut,
-    "CircularIn": CircularEaseIn,
-    "CircularOut": CircularEaseOut,
-    "CircularInOut": CircularEaseInOut,
-    "ExponentialIn": ExponentialEaseIn,
-    "ExponentialOut": ExponentialEaseOut,
-    "ExponentialInOut": ExponentialEaseInOut,
-    "ElasticIn": ElasticEaseIn,
-    "ElasticOut": ElasticEaseOut,
-    "ElasticInOut": ElasticEaseInOut,
-    "BackIn": BackEaseIn,
-    "BackOut": BackEaseOut,
-    "BackInOut": BackEaseInOut,
-    "BounceIn": BounceEaseIn,
-    "BounceOut": BounceEaseOut,
-    "BounceInOut": BounceEaseInOut,
-}
-
-EASING_FUNCTION_KEYS = Literal[tuple(EASING_FUNCTIONS_MAP.keys())]
-
-
-# actually I think for now could just use CollectionOutput (which is list[Any]
-@invocation(
-    "step_param_easing",
-    title="Step Param Easing",
-    tags=["step", "easing"],
-    category="step",
-    version="1.0.2",
-)
-class StepParamEasingInvocation(BaseInvocation):
-    """Experimental per-step parameter easing for denoising steps"""
-
-    easing: EASING_FUNCTION_KEYS = InputField(default="Linear", description="The easing function to use")
-    num_steps: int = InputField(default=20, description="number of denoising steps")
-    start_value: float = InputField(default=0.0, description="easing starting value")
-    end_value: float = InputField(default=1.0, description="easing ending value")
-    start_step_percent: float = InputField(default=0.0, description="fraction of steps at which to start easing")
-    end_step_percent: float = InputField(default=1.0, description="fraction of steps after which to end easing")
-    # if None, then start_value is used prior to easing start
-    pre_start_value: Optional[float] = InputField(default=None, description="value before easing start")
-    # if None, then end value is used prior to easing end
-    post_end_value: Optional[float] = InputField(default=None, description="value after easing end")
-    mirror: bool = InputField(default=False, description="include mirror of easing function")
-    # FIXME: add alt_mirror option (alternative to default or mirror), or remove entirely
-    # alt_mirror: bool = InputField(default=False, description="alternative mirroring by dual easing")
-    show_easing_plot: bool = InputField(default=False, description="show easing plot")
-
-    def invoke(self, context: InvocationContext) -> FloatCollectionOutput:
-        log_diagnostics = False
-        # convert from start_step_percent to nearest step <= (steps * start_step_percent)
-        # start_step = int(np.floor(self.num_steps * self.start_step_percent))
-        start_step = int(np.round(self.num_steps * self.start_step_percent))
-        # convert from end_step_percent to nearest step >= (steps * end_step_percent)
-        # end_step = int(np.ceil((self.num_steps - 1) * self.end_step_percent))
-        end_step = int(np.round((self.num_steps - 1) * self.end_step_percent))
-
-        # end_step = int(np.ceil(self.num_steps * self.end_step_percent))
-        num_easing_steps = end_step - start_step + 1
-
-        # num_presteps = max(start_step - 1, 0)
-        num_presteps = start_step
-        num_poststeps = self.num_steps - (num_presteps + num_easing_steps)
-        prelist = list(num_presteps * [self.pre_start_value])
-        postlist = list(num_poststeps * [self.post_end_value])
-
-        if log_diagnostics:
-            context.logger.debug("start_step: " + str(start_step))
-            context.logger.debug("end_step: " + str(end_step))
-            context.logger.debug("num_easing_steps: " + str(num_easing_steps))
-            context.logger.debug("num_presteps: " + str(num_presteps))
-            context.logger.debug("num_poststeps: " + str(num_poststeps))
-            context.logger.debug("prelist size: " + str(len(prelist)))
-            context.logger.debug("postlist size: " + str(len(postlist)))
-            context.logger.debug("prelist: " + str(prelist))
-            context.logger.debug("postlist: " + str(postlist))
-
-        easing_class = EASING_FUNCTIONS_MAP[self.easing]
-        if log_diagnostics:
-            context.logger.debug("easing class: " + str(easing_class))
-        easing_list = []
-        if self.mirror:  # "expected" mirroring
-            # if number of steps is even, squeeze duration down to (number_of_steps)/2
-            # and create reverse copy of list to append
-            # if number of steps is odd, squeeze duration down to ceil(number_of_steps/2)
-            # and create reverse copy of list[1:end-1]
-            # but if even then number_of_steps/2 === ceil(number_of_steps/2), so can just use ceil always
-
-            base_easing_duration = int(np.ceil(num_easing_steps / 2.0))
-            if log_diagnostics:
-                context.logger.debug("base easing duration: " + str(base_easing_duration))
-            even_num_steps = num_easing_steps % 2 == 0  # even number of steps
-            easing_function = easing_class(
-                start=self.start_value,
-                end=self.end_value,
-                duration=base_easing_duration - 1,
-            )
-            base_easing_vals = []
-            for step_index in range(base_easing_duration):
-                easing_val = easing_function.ease(step_index)
-                base_easing_vals.append(easing_val)
-                if log_diagnostics:
-                    context.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(easing_val))
-            if even_num_steps:
-                mirror_easing_vals = list(reversed(base_easing_vals))
-            else:
-                mirror_easing_vals = list(reversed(base_easing_vals[0:-1]))
-            if log_diagnostics:
-                context.logger.debug("base easing vals: " + str(base_easing_vals))
-                context.logger.debug("mirror easing vals: " + str(mirror_easing_vals))
-            easing_list = base_easing_vals + mirror_easing_vals
-
-        # FIXME: add alt_mirror option (alternative to default or mirror), or remove entirely
-        # elif self.alt_mirror:  # function mirroring (unintuitive behavior (at least to me))
-        #     # half_ease_duration = round(num_easing_steps - 1 / 2)
-        #     half_ease_duration = round((num_easing_steps - 1) / 2)
-        #     easing_function = easing_class(start=self.start_value,
-        #                                    end=self.end_value,
-        #                                    duration=half_ease_duration,
-        #                                    )
-        #
-        #     mirror_function = easing_class(start=self.end_value,
-        #                                    end=self.start_value,
-        #                                    duration=half_ease_duration,
-        #                                    )
-        #     for step_index in range(num_easing_steps):
-        #         if step_index <= half_ease_duration:
-        #             step_val = easing_function.ease(step_index)
-        #         else:
-        #             step_val = mirror_function.ease(step_index - half_ease_duration)
-        #         easing_list.append(step_val)
-        #         if log_diagnostics: logger.debug(step_index, step_val)
-        #
-
-        else:  # no mirroring (default)
-            easing_function = easing_class(
-                start=self.start_value,
-                end=self.end_value,
-                duration=num_easing_steps - 1,
-            )
-            for step_index in range(num_easing_steps):
-                step_val = easing_function.ease(step_index)
-                easing_list.append(step_val)
-                if log_diagnostics:
-                    context.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(step_val))
-
-        if log_diagnostics:
-            context.logger.debug("prelist size: " + str(len(prelist)))
-            context.logger.debug("easing_list size: " + str(len(easing_list)))
-            context.logger.debug("postlist size: " + str(len(postlist)))
-
-        param_list = prelist + easing_list + postlist
-
-        if self.show_easing_plot:
-            plt.figure()
-            plt.xlabel("Step")
-            plt.ylabel("Param Value")
-            plt.title("Per-Step Values Based On Easing: " + self.easing)
-            plt.bar(range(len(param_list)), param_list)
-            # plt.plot(param_list)
-            ax = plt.gca()
-            ax.xaxis.set_major_locator(MaxNLocator(integer=True))
-            buf = io.BytesIO()
-            plt.savefig(buf, format="png")
-            buf.seek(0)
-            im = PIL.Image.open(buf)
-            im.show()
-            buf.close()
-
-        # output array of size steps, each entry list[i] is param value for step i
-        return FloatCollectionOutput(collection=param_list)

invokeai/app/invocations/primitives.py

@@ -4,7 +4,13 @@ from typing import Optional
 
 import torch
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
 from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
 from invokeai.app.invocations.fields import (
     BoundingBoxField,
@@ -533,3 +539,23 @@ class BoundingBoxInvocation(BaseInvocation):
 
 
 # endregion
+
+
+@invocation(
+    "image_batch",
+    title="Image Batch",
+    tags=["primitives", "image", "batch", "internal"],
+    category="primitives",
+    version="1.0.0",
+    classification=Classification.Special,
+)
+class ImageBatchInvocation(BaseInvocation):
+    """Create a batched generation, where the workflow is executed once for each image in the batch."""
+
+    images: list[ImageField] = InputField(min_length=1, description="The images to batch over", input=Input.Direct)
+
+    def __init__(self):
+        raise NotImplementedError("This class should never be executed or instantiated directly.")
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        raise NotImplementedError("This class should never be executed or instantiated directly.")

invokeai/app/invocations/sd3_latents_to_image.py

@@ -6,6 +6,7 @@ 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,
@@ -26,7 +27,7 @@ from invokeai.backend.util.devices import TorchDevice
     title="SD3 Latents to Image",
     tags=["latents", "image", "vae", "l2i", "sd3"],
     category="latents",
-    version="1.3.0",
+    version="1.3.1",
 )
 class SD3LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents."""
@@ -40,13 +41,30 @@ 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). This estimate is accurate for both SD1 and SDXL.
+        out_h = LATENT_SCALE_FACTOR * latents.shape[-2]
+        out_w = LATENT_SCALE_FACTOR * latents.shape[-1]
+        element_size = next(vae.parameters()).element_size()
+        # TODO(ryand): Need to tune this value, it was copied from the SD1 implementation.
+        scaling_constant = 960  # Determined experimentally.
+        working_memory = out_h * out_w * element_size * scaling_constant
+
+        return 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))
-        with SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), vae_info as vae:
+        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),
+        ):
             context.util.signal_progress("Running VAE")
             assert isinstance(vae, (AutoencoderKL))
             latents = latents.to(vae.device)

invokeai/app/invocations/sd3_text_encoder.py

@@ -16,11 +16,12 @@ 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
 
 # The SD3 T5 Max Sequence Length set based on the default in diffusers.
 SD3_T5_MAX_SEQ_LEN = 256
@@ -150,10 +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(
-                    LoRAPatcher.apply_lora_patches(
+                    LayerPatcher.apply_smart_model_patches(
                         model=clip_text_encoder,
                         patches=self._clip_lora_iterator(context, clip_model),
                         prefix=FLUX_LORA_CLIP_PREFIX,
+                        dtype=TorchDevice.choose_torch_dtype(),
                         cached_weights=cached_weights,
                     )
                 )
@@ -193,9 +195,9 @@ class Sd3TextEncoderInvocation(BaseInvocation):
 
     def _clip_lora_iterator(
         self, context: InvocationContext, clip_model: CLIPField
-    ) -> Iterator[Tuple[LoRAModelRaw, float]]:
+    ) -> Iterator[Tuple[ModelPatchRaw, float]]:
         for lora in clip_model.loras:
             lora_info = context.models.load(lora.lora)
-            assert isinstance(lora_info.model, LoRAModelRaw)
+            assert isinstance(lora_info.model, ModelPatchRaw)
             yield (lora_info.model, lora.weight)
             del lora_info

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.lora.lora_model_raw import LoRAModelRaw
-from invokeai.backend.lora.lora_patcher import LoRAPatcher
+from invokeai.backend.patches.layer_patcher import LayerPatcher
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState
 from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
     MultiDiffusionPipeline,
@@ -194,10 +194,10 @@ 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[LoRAModelRaw, float]]:
+        def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
             for lora in self.unet.loras:
                 lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, LoRAModelRaw)
+                assert isinstance(lora_info.model, ModelPatchRaw)
                 yield (lora_info.model, lora.weight)
                 del lora_info
 
@@ -207,7 +207,9 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
         with (
             ExitStack() as exit_stack,
             unet_info as unet,
-            LoRAPatcher.apply_lora_patches(model=unet, patches=_lora_loader(), prefix="lora_unet_"),
+            LayerPatcher.apply_smart_model_patches(
+                model=unet, patches=_lora_loader(), prefix="lora_unet_", dtype=unet.dtype
+            ),
         ):
             assert isinstance(unet, UNet2DConditionModel)
             latents = latents.to(device=unet.device, dtype=unet.dtype)

invokeai/app/services/config/config_default.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import copy
+import filecmp
 import locale
 import os
 import re
@@ -12,7 +13,6 @@ 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,8 +24,6 @@ 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"]
@@ -35,24 +33,6 @@ 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")
@@ -101,10 +81,12 @@ class InvokeAIAppConfig(BaseSettings):
         profile_graphs: Enable graph profiling using `cProfile`.
         profile_prefix: An optional prefix for profile output files.
         profiles_dir: Path to profiles output directory.
-        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.
+        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.
         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.
@@ -171,10 +153,12 @@ class InvokeAIAppConfig(BaseSettings):
     profiles_dir:                  Path = Field(default=Path("profiles"),   description="Path to profiles output directory.")
 
     # CACHE
-    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.")
+    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.")
     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=2,                  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.")
@@ -525,9 +509,35 @@ def get_config() -> InvokeAIAppConfig:
     ]
     example_config.write_file(config.config_file_path.with_suffix(".example.yaml"), as_example=True)
 
-    # Copy all legacy configs - We know `__path__[0]` is correct here
+    # Copy all legacy configs only if needed
+    # We know `__path__[0]` is correct here
     configs_src = Path(model_configs.__path__[0])  # pyright: ignore [reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
-    shutil.copytree(configs_src, config.legacy_conf_path, dirs_exist_ok=True)
+    dest_path = config.legacy_conf_path
+
+    # Create destination (we don't need to check for existence)
+    dest_path.mkdir(parents=True, exist_ok=True)
+
+    # Compare directories recursively
+    comparison = filecmp.dircmp(configs_src, dest_path)
+    need_copy = any(
+        [
+            comparison.left_only,  # Files exist only in source
+            comparison.diff_files,  # Files that differ
+            comparison.common_funny,  # Files that couldn't be compared
+        ]
+    )
+
+    if need_copy:
+        # Get permissions from destination directory
+        dest_mode = dest_path.stat().st_mode
+
+        # Copy directory tree
+        shutil.copytree(configs_src, dest_path, dirs_exist_ok=True)
+
+        # Set permissions on copied files to match destination directory
+        dest_path.chmod(dest_mode)
+        for p in dest_path.glob("**/*"):
+            p.chmod(dest_mode)
 
     if config.config_file_path.exists():
         config_from_file = load_and_migrate_config(config.config_file_path)

invokeai/app/services/invocation_stats/invocation_stats_common.py

@@ -22,7 +22,6 @@ class ModelCacheStatsSummary:
     """The stats for the model cache."""
 
     high_water_mark_gb: float
-    cache_size_gb: float
     total_usage_gb: float
     cache_hits: int
     cache_misses: int
@@ -79,7 +78,7 @@ class InvocationStatsSummary:
         _str += f"   Model cache misses: {self.model_cache_stats.cache_misses}\n"
         _str += f"   Models cached: {self.model_cache_stats.models_cached}\n"
         _str += f"   Models cleared from cache: {self.model_cache_stats.models_cleared}\n"
-        _str += f"   Cache high water mark: {self.model_cache_stats.high_water_mark_gb:4.2f}/{self.model_cache_stats.cache_size_gb:4.2f}G\n"
+        _str += f"   Cache high water mark: {self.model_cache_stats.high_water_mark_gb:4.2f}G\n"
 
         return _str
 

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 import CacheStats
+from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
 
 # Size of 1GB in bytes.
 GB = 2**30
@@ -111,7 +111,6 @@ class InvocationStatsService(InvocationStatsServiceBase):
             cache_hits=cache_stats.hits,
             cache_misses=cache_stats.misses,
             high_water_mark_gb=cache_stats.high_watermark / GB,
-            cache_size_gb=cache_stats.cache_size / GB,
             total_usage_gb=sum(list(cache_stats.loaded_model_sizes.values())) / GB,
             models_cached=cache_stats.in_cache,
             models_cleared=cache_stats.cleared,

invokeai/app/services/model_load/model_load_base.py

@@ -7,14 +7,18 @@ 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_base import ModelCacheBase
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
 
 
 class ModelLoadServiceBase(ABC):
     """Wrapper around AnyModelLoader."""
 
     @abstractmethod
-    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+    def load_model(
+        self,
+        model_config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> LoadedModel:
         """
         Given a model's configuration, load it and return the LoadedModel object.
 
@@ -24,7 +28,7 @@ class ModelLoadServiceBase(ABC):
 
     @property
     @abstractmethod
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
+    def ram_cache(self) -> ModelCache:
         """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_base import ModelCacheBase
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
 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: ModelCacheBase[AnyModel],
+        ram_cache: ModelCache,
         registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry,
     ):
         """Initialize the model load service."""
@@ -45,11 +45,15 @@ class ModelLoadService(ModelLoadServiceBase):
         self._invoker = invoker
 
     @property
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
+    def ram_cache(self) -> ModelCache:
         """Return the RAM cache used by this loader."""
         return self._ram_cache
 
-    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+    def load_model(
+        self,
+        model_config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> LoadedModel:
         """
         Given a model's configuration, load it and return the LoadedModel object.
 
@@ -78,15 +82,14 @@ 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(_locker=ram_cache.get(key=cache_key))
+            return LoadedModelWithoutConfig(cache_record=self._ram_cache.get(key=cache_key), cache=self._ram_cache)
         except IndexError:
             pass
 
         def torch_load_file(checkpoint: Path) -> AnyModel:
             scan_result = scan_file_path(checkpoint)
-            if scan_result.infected_files != 0:
+            if scan_result.infected_files != 0 or scan_result.scan_err:
                 raise Exception("The model at {checkpoint} is potentially infected by malware. Aborting load.")
             result = torch_load(checkpoint, map_location="cpu")
             return result
@@ -109,5 +112,5 @@ class ModelLoadService(ModelLoadServiceBase):
         )
         assert loader is not None
         raw_model = loader(model_path)
-        ram_cache.put(key=cache_key, model=raw_model)
-        return LoadedModelWithoutConfig(_locker=ram_cache.get(key=cache_key))
+        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)

invokeai/app/services/model_manager/model_manager_default.py

@@ -16,7 +16,8 @@ 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 import ModelCache, ModelLoaderRegistry
+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.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 
@@ -81,9 +82,10 @@ class ModelManagerService(ModelManagerServiceBase):
         logger.setLevel(app_config.log_level.upper())
 
         ram_cache = ModelCache(
-            max_cache_size=app_config.ram,
-            max_vram_cache_size=app_config.vram,
-            lazy_offloading=app_config.lazy_offload,
+            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,
             logger=logger,
             execution_device=execution_device or TorchDevice.choose_torch_device(),
         )

invokeai/app/services/session_processor/session_processor_default.py

@@ -378,6 +378,9 @@ class DefaultSessionProcessor(SessionProcessorBase):
         self._poll_now()
 
     async def _on_queue_item_status_changed(self, event: FastAPIEvent[QueueItemStatusChangedEvent]) -> None:
+        # Make sure the cancel event is for the currently processing queue item
+        if self._queue_item and self._queue_item.item_id != event[1].item_id:
+            return
         if self._queue_item and event[1].status in ["completed", "failed", "canceled"]:
             # When the queue item is canceled via HTTP, the queue item status is set to `"canceled"` and this event is
             # emitted. We need to respond to this event and stop graph execution. This is done by setting the cancel

invokeai/app/services/session_queue/session_queue_common.py

@@ -16,6 +16,7 @@ from pydantic import (
 from pydantic_core import to_jsonable_python
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
+from invokeai.app.invocations.fields import ImageField
 from invokeai.app.services.shared.graph import Graph, GraphExecutionState, NodeNotFoundError
 from invokeai.app.services.workflow_records.workflow_records_common import (
     WorkflowWithoutID,
@@ -51,11 +52,7 @@ class SessionQueueItemNotFoundError(ValueError):
 
 # region Batch
 
-BatchDataType = Union[
-    StrictStr,
-    float,
-    int,
-]
+BatchDataType = Union[StrictStr, float, int, ImageField]
 
 
 class NodeFieldValue(BaseModel):

invokeai/app/services/shared/invocation_context.py

@@ -361,7 +361,9 @@ class ModelsInterface(InvocationContextInterface):
             return self._services.model_manager.store.exists(identifier.key)
 
     def load(
-        self, identifier: Union[str, "ModelIdentifierField"], submodel_type: Optional[SubModelType] = None
+        self,
+        identifier: Union[str, "ModelIdentifierField"],
+        submodel_type: Optional[SubModelType] = None,
     ) -> LoadedModel:
         """Load a model.
 

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

@@ -35,7 +35,7 @@ class Migration11Callback:
 
     def _remove_convert_cache(self) -> None:
         """Rename models/.cache to models/.convert_cache."""
-        self._logger.info("Removing .cache directory. Converted models will now be cached in .convert_cache.")
+        self._logger.info("Removing models/.cache directory. Converted models will now be cached in .convert_cache.")
         legacy_convert_path = self._app_config.root_path / "models" / ".cache"
         shutil.rmtree(legacy_convert_path, ignore_errors=True)
 

invokeai/backend/flux/custom_block_processor.py

@@ -1,9 +1,10 @@
 import einops
 import torch
 
+from invokeai.backend.flux.extensions.regional_prompting_extension import RegionalPromptingExtension
 from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
 from invokeai.backend.flux.math import attention
-from invokeai.backend.flux.modules.layers import DoubleStreamBlock
+from invokeai.backend.flux.modules.layers import DoubleStreamBlock, SingleStreamBlock
 
 
 class CustomDoubleStreamBlockProcessor:
@@ -13,7 +14,12 @@ class CustomDoubleStreamBlockProcessor:
 
     @staticmethod
     def _double_stream_block_forward(
-        block: DoubleStreamBlock, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor
+        block: DoubleStreamBlock,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """This function is a direct copy of DoubleStreamBlock.forward(), but it returns some of the intermediate
         values.
@@ -40,7 +46,7 @@ class CustomDoubleStreamBlockProcessor:
         k = torch.cat((txt_k, img_k), dim=2)
         v = torch.cat((txt_v, img_v), dim=2)
 
-        attn = attention(q, k, v, pe=pe)
+        attn = attention(q, k, v, pe=pe, attn_mask=attn_mask)
         txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
 
         # calculate the img bloks
@@ -63,11 +69,15 @@ class CustomDoubleStreamBlockProcessor:
         vec: torch.Tensor,
         pe: torch.Tensor,
         ip_adapter_extensions: list[XLabsIPAdapterExtension],
+        regional_prompting_extension: RegionalPromptingExtension,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """A custom implementation of DoubleStreamBlock.forward() with additional features:
         - IP-Adapter support
         """
-        img, txt, img_q = CustomDoubleStreamBlockProcessor._double_stream_block_forward(block, img, txt, vec, pe)
+        attn_mask = regional_prompting_extension.get_double_stream_attn_mask(block_index)
+        img, txt, img_q = CustomDoubleStreamBlockProcessor._double_stream_block_forward(
+            block, img, txt, vec, pe, attn_mask=attn_mask
+        )
 
         # Apply IP-Adapter conditioning.
         for ip_adapter_extension in ip_adapter_extensions:
@@ -81,3 +91,48 @@ class CustomDoubleStreamBlockProcessor:
             )
 
         return img, txt
+
+
+class CustomSingleStreamBlockProcessor:
+    """A class containing a custom implementation of SingleStreamBlock.forward() with additional features (masking,
+    etc.)
+    """
+
+    @staticmethod
+    def _single_stream_block_forward(
+        block: SingleStreamBlock,
+        x: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """This function is a direct copy of SingleStreamBlock.forward()."""
+        mod, _ = block.modulation(vec)
+        x_mod = (1 + mod.scale) * block.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(block.linear1(x_mod), [3 * block.hidden_size, block.mlp_hidden_dim], dim=-1)
+
+        q, k, v = einops.rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=block.num_heads)
+        q, k = block.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe, attn_mask=attn_mask)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = block.linear2(torch.cat((attn, block.mlp_act(mlp)), 2))
+        return x + mod.gate * output
+
+    @staticmethod
+    def custom_single_block_forward(
+        timestep_index: int,
+        total_num_timesteps: int,
+        block_index: int,
+        block: SingleStreamBlock,
+        img: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        regional_prompting_extension: RegionalPromptingExtension,
+    ) -> torch.Tensor:
+        """A custom implementation of SingleStreamBlock.forward() with additional features:
+        - Masking
+        """
+        attn_mask = regional_prompting_extension.get_single_stream_attn_mask(block_index)
+        return CustomSingleStreamBlockProcessor._single_stream_block_forward(block, img, vec, pe, attn_mask=attn_mask)

invokeai/backend/flux/denoise.py

@@ -7,6 +7,7 @@ from tqdm import tqdm
 from invokeai.backend.flux.controlnet.controlnet_flux_output import ControlNetFluxOutput, sum_controlnet_flux_outputs
 from invokeai.backend.flux.extensions.inpaint_extension import InpaintExtension
 from invokeai.backend.flux.extensions.instantx_controlnet_extension import InstantXControlNetExtension
+from invokeai.backend.flux.extensions.regional_prompting_extension import RegionalPromptingExtension
 from invokeai.backend.flux.extensions.xlabs_controlnet_extension import XLabsControlNetExtension
 from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
 from invokeai.backend.flux.model import Flux
@@ -18,14 +19,8 @@ def denoise(
     # model input
     img: torch.Tensor,
     img_ids: torch.Tensor,
-    # positive text conditioning
-    txt: torch.Tensor,
-    txt_ids: torch.Tensor,
-    vec: torch.Tensor,
-    # negative text conditioning
-    neg_txt: torch.Tensor | None,
-    neg_txt_ids: torch.Tensor | None,
-    neg_vec: torch.Tensor | None,
+    pos_regional_prompting_extension: RegionalPromptingExtension,
+    neg_regional_prompting_extension: RegionalPromptingExtension | None,
     # sampling parameters
     timesteps: list[float],
     step_callback: Callable[[PipelineIntermediateState], None],
@@ -35,6 +30,8 @@ def denoise(
     controlnet_extensions: list[XLabsControlNetExtension | InstantXControlNetExtension],
     pos_ip_adapter_extensions: list[XLabsIPAdapterExtension],
     neg_ip_adapter_extensions: list[XLabsIPAdapterExtension],
+    # extra img tokens
+    img_cond: torch.Tensor | None,
 ):
     # step 0 is the initial state
     total_steps = len(timesteps) - 1
@@ -61,9 +58,9 @@ def denoise(
                     total_num_timesteps=total_steps,
                     img=img,
                     img_ids=img_ids,
-                    txt=txt,
-                    txt_ids=txt_ids,
-                    y=vec,
+                    txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                    txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                    y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
                     timesteps=t_vec,
                     guidance=guidance_vec,
                 )
@@ -74,13 +71,13 @@ def denoise(
         # controlnet_residuals datastructure is efficient in that it likely contains multiple references to the same
         # tensors. Calculating the sum materializes each tensor into its own instance.
         merged_controlnet_residuals = sum_controlnet_flux_outputs(controlnet_residuals)
-
+        pred_img = torch.cat((img, img_cond), dim=-1) if img_cond is not None else img
         pred = model(
-            img=img,
+            img=pred_img,
             img_ids=img_ids,
-            txt=txt,
-            txt_ids=txt_ids,
-            y=vec,
+            txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+            txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+            y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
             timesteps=t_vec,
             guidance=guidance_vec,
             timestep_index=step_index,
@@ -88,6 +85,7 @@ def denoise(
             controlnet_double_block_residuals=merged_controlnet_residuals.double_block_residuals,
             controlnet_single_block_residuals=merged_controlnet_residuals.single_block_residuals,
             ip_adapter_extensions=pos_ip_adapter_extensions,
+            regional_prompting_extension=pos_regional_prompting_extension,
         )
 
         step_cfg_scale = cfg_scale[step_index]
@@ -97,15 +95,15 @@ def denoise(
             # TODO(ryand): Add option to run positive and negative predictions in a single batch for better performance
             # on systems with sufficient VRAM.
 
-            if neg_txt is None or neg_txt_ids is None or neg_vec is None:
+            if neg_regional_prompting_extension is None:
                 raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
 
             neg_pred = model(
                 img=img,
                 img_ids=img_ids,
-                txt=neg_txt,
-                txt_ids=neg_txt_ids,
-                y=neg_vec,
+                txt=neg_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                txt_ids=neg_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                y=neg_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
                 timesteps=t_vec,
                 guidance=guidance_vec,
                 timestep_index=step_index,
@@ -113,6 +111,7 @@ def denoise(
                 controlnet_double_block_residuals=None,
                 controlnet_single_block_residuals=None,
                 ip_adapter_extensions=neg_ip_adapter_extensions,
+                regional_prompting_extension=neg_regional_prompting_extension,
             )
             pred = neg_pred + step_cfg_scale * (pred - neg_pred)
 

invokeai/backend/flux/extensions/regional_prompting_extension.py

@@ -0,0 +1,276 @@
+from typing import Optional
+
+import torch
+import torchvision
+
+from invokeai.backend.flux.text_conditioning import FluxRegionalTextConditioning, FluxTextConditioning
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.mask import to_standard_float_mask
+
+
+class RegionalPromptingExtension:
+    """A class for managing regional prompting with FLUX.
+
+    This implementation is inspired by https://arxiv.org/pdf/2411.02395 (though there are significant differences).
+    """
+
+    def __init__(
+        self,
+        regional_text_conditioning: FluxRegionalTextConditioning,
+        restricted_attn_mask: torch.Tensor | None = None,
+    ):
+        self.regional_text_conditioning = regional_text_conditioning
+        self.restricted_attn_mask = restricted_attn_mask
+
+    def get_double_stream_attn_mask(self, block_index: int) -> torch.Tensor | None:
+        order = [self.restricted_attn_mask, None]
+        return order[block_index % len(order)]
+
+    def get_single_stream_attn_mask(self, block_index: int) -> torch.Tensor | None:
+        order = [self.restricted_attn_mask, None]
+        return order[block_index % len(order)]
+
+    @classmethod
+    def from_text_conditioning(cls, text_conditioning: list[FluxTextConditioning], img_seq_len: int):
+        """Create a RegionalPromptingExtension from a list of text conditionings.
+
+        Args:
+            text_conditioning (list[FluxTextConditioning]): The text conditionings to use for regional prompting.
+            img_seq_len (int): The image sequence length (i.e. packed_height * packed_width).
+        """
+        regional_text_conditioning = cls._concat_regional_text_conditioning(text_conditioning)
+        attn_mask_with_restricted_img_self_attn = cls._prepare_restricted_attn_mask(
+            regional_text_conditioning, img_seq_len
+        )
+        return cls(
+            regional_text_conditioning=regional_text_conditioning,
+            restricted_attn_mask=attn_mask_with_restricted_img_self_attn,
+        )
+
+    # Keeping _prepare_unrestricted_attn_mask for reference as an alternative masking strategy:
+    #
+    # @classmethod
+    # def _prepare_unrestricted_attn_mask(
+    #     cls,
+    #     regional_text_conditioning: FluxRegionalTextConditioning,
+    #     img_seq_len: int,
+    # ) -> torch.Tensor:
+    #     """Prepare an 'unrestricted' attention mask. In this context, 'unrestricted' means that:
+    #     - img self-attention is not masked.
+    #     - img regions attend to both txt within their own region and to global prompts.
+    #     """
+    #     device = TorchDevice.choose_torch_device()
+
+    #     # Infer txt_seq_len from the t5_embeddings tensor.
+    #     txt_seq_len = regional_text_conditioning.t5_embeddings.shape[1]
+
+    #     # In the attention blocks, the txt seq and img seq are concatenated and then attention is applied.
+    #     # Concatenation happens in the following order: [txt_seq, img_seq].
+    #     # There are 4 portions of the attention mask to consider as we prepare it:
+    #     # 1. txt attends to itself
+    #     # 2. txt attends to corresponding regional img
+    #     # 3. regional img attends to corresponding txt
+    #     # 4. regional img attends to itself
+
+    #     # Initialize empty attention mask.
+    #     regional_attention_mask = torch.zeros(
+    #         (txt_seq_len + img_seq_len, txt_seq_len + img_seq_len), device=device, dtype=torch.float16
+    #     )
+
+    #     for image_mask, t5_embedding_range in zip(
+    #         regional_text_conditioning.image_masks, regional_text_conditioning.t5_embedding_ranges, strict=True
+    #     ):
+    #         # 1. txt attends to itself
+    #         regional_attention_mask[
+    #             t5_embedding_range.start : t5_embedding_range.end, t5_embedding_range.start : t5_embedding_range.end
+    #         ] = 1.0
+
+    #         # 2. txt attends to corresponding regional img
+    #         # Note that we reshape to (1, img_seq_len) to ensure broadcasting works as desired.
+    #         fill_value = image_mask.view(1, img_seq_len) if image_mask is not None else 1.0
+    #         regional_attention_mask[t5_embedding_range.start : t5_embedding_range.end, txt_seq_len:] = fill_value
+
+    #         # 3. regional img attends to corresponding txt
+    #         # Note that we reshape to (img_seq_len, 1) to ensure broadcasting works as desired.
+    #         fill_value = image_mask.view(img_seq_len, 1) if image_mask is not None else 1.0
+    #         regional_attention_mask[txt_seq_len:, t5_embedding_range.start : t5_embedding_range.end] = fill_value
+
+    #     # 4. regional img attends to itself
+    #     # Allow unrestricted img self attention.
+    #     regional_attention_mask[txt_seq_len:, txt_seq_len:] = 1.0
+
+    #     # Convert attention mask to boolean.
+    #     regional_attention_mask = regional_attention_mask > 0.5
+
+    #     return regional_attention_mask
+
+    @classmethod
+    def _prepare_restricted_attn_mask(
+        cls,
+        regional_text_conditioning: FluxRegionalTextConditioning,
+        img_seq_len: int,
+    ) -> torch.Tensor | None:
+        """Prepare a 'restricted' attention mask. In this context, 'restricted' means that:
+        - img self-attention is only allowed within regions.
+        - img regions only attend to txt within their own region, not to global prompts.
+        """
+        # Identify background region. I.e. the region that is not covered by any region masks.
+        background_region_mask: None | torch.Tensor = None
+        for image_mask in regional_text_conditioning.image_masks:
+            if image_mask is not None:
+                if background_region_mask is None:
+                    background_region_mask = torch.ones_like(image_mask)
+                background_region_mask *= 1 - image_mask
+
+        if background_region_mask is None:
+            # There are no region masks, short-circuit and return None.
+            # TODO(ryand): We could restrict txt-txt attention across multiple global prompts, but this would
+            # is a rare use case and would make the logic here significantly more complicated.
+            return None
+
+        device = TorchDevice.choose_torch_device()
+
+        # Infer txt_seq_len from the t5_embeddings tensor.
+        txt_seq_len = regional_text_conditioning.t5_embeddings.shape[1]
+
+        # In the attention blocks, the txt seq and img seq are concatenated and then attention is applied.
+        # Concatenation happens in the following order: [txt_seq, img_seq].
+        # There are 4 portions of the attention mask to consider as we prepare it:
+        # 1. txt attends to itself
+        # 2. txt attends to corresponding regional img
+        # 3. regional img attends to corresponding txt
+        # 4. regional img attends to itself
+
+        # Initialize empty attention mask.
+        regional_attention_mask = torch.zeros(
+            (txt_seq_len + img_seq_len, txt_seq_len + img_seq_len), device=device, dtype=torch.float16
+        )
+
+        for image_mask, t5_embedding_range in zip(
+            regional_text_conditioning.image_masks, regional_text_conditioning.t5_embedding_ranges, strict=True
+        ):
+            # 1. txt attends to itself
+            regional_attention_mask[
+                t5_embedding_range.start : t5_embedding_range.end, t5_embedding_range.start : t5_embedding_range.end
+            ] = 1.0
+
+            if image_mask is not None:
+                # 2. txt attends to corresponding regional img
+                # Note that we reshape to (1, img_seq_len) to ensure broadcasting works as desired.
+                regional_attention_mask[t5_embedding_range.start : t5_embedding_range.end, txt_seq_len:] = (
+                    image_mask.view(1, img_seq_len)
+                )
+
+                # 3. regional img attends to corresponding txt
+                # Note that we reshape to (img_seq_len, 1) to ensure broadcasting works as desired.
+                regional_attention_mask[txt_seq_len:, t5_embedding_range.start : t5_embedding_range.end] = (
+                    image_mask.view(img_seq_len, 1)
+                )
+
+                # 4. regional img attends to itself
+                image_mask = image_mask.view(img_seq_len, 1)
+                regional_attention_mask[txt_seq_len:, txt_seq_len:] += image_mask @ image_mask.T
+            else:
+                # We don't allow attention between non-background image regions and global prompts. This helps to ensure
+                # that regions focus on their local prompts. We do, however, allow attention between background regions
+                # and global prompts. If we didn't do this, then the background regions would not attend to any txt
+                # embeddings, which we found experimentally to cause artifacts.
+
+                # 2. global txt attends to background region
+                # Note that we reshape to (1, img_seq_len) to ensure broadcasting works as desired.
+                regional_attention_mask[t5_embedding_range.start : t5_embedding_range.end, txt_seq_len:] = (
+                    background_region_mask.view(1, img_seq_len)
+                )
+
+                # 3. background region attends to global txt
+                # Note that we reshape to (img_seq_len, 1) to ensure broadcasting works as desired.
+                regional_attention_mask[txt_seq_len:, t5_embedding_range.start : t5_embedding_range.end] = (
+                    background_region_mask.view(img_seq_len, 1)
+                )
+
+        # Allow background regions to attend to themselves.
+        regional_attention_mask[txt_seq_len:, txt_seq_len:] += background_region_mask.view(img_seq_len, 1)
+        regional_attention_mask[txt_seq_len:, txt_seq_len:] += background_region_mask.view(1, img_seq_len)
+
+        # Convert attention mask to boolean.
+        regional_attention_mask = regional_attention_mask > 0.5
+
+        return regional_attention_mask
+
+    @classmethod
+    def _concat_regional_text_conditioning(
+        cls,
+        text_conditionings: list[FluxTextConditioning],
+    ) -> FluxRegionalTextConditioning:
+        """Concatenate regional text conditioning data into a single conditioning tensor (with associated masks)."""
+        concat_t5_embeddings: list[torch.Tensor] = []
+        concat_t5_embedding_ranges: list[Range] = []
+        image_masks: list[torch.Tensor | None] = []
+
+        # Choose global CLIP embedding.
+        # Use the first global prompt's CLIP embedding as the global CLIP embedding. If there is no global prompt, use
+        # the first prompt's CLIP embedding.
+        global_clip_embedding: torch.Tensor = text_conditionings[0].clip_embeddings
+        for text_conditioning in text_conditionings:
+            if text_conditioning.mask is None:
+                global_clip_embedding = text_conditioning.clip_embeddings
+                break
+
+        cur_t5_embedding_len = 0
+        for text_conditioning in text_conditionings:
+            concat_t5_embeddings.append(text_conditioning.t5_embeddings)
+
+            concat_t5_embedding_ranges.append(
+                Range(start=cur_t5_embedding_len, end=cur_t5_embedding_len + text_conditioning.t5_embeddings.shape[1])
+            )
+
+            image_masks.append(text_conditioning.mask)
+
+            cur_t5_embedding_len += text_conditioning.t5_embeddings.shape[1]
+
+        t5_embeddings = torch.cat(concat_t5_embeddings, dim=1)
+
+        # Initialize the txt_ids tensor.
+        pos_bs, pos_t5_seq_len, _ = t5_embeddings.shape
+        t5_txt_ids = torch.zeros(
+            pos_bs, pos_t5_seq_len, 3, dtype=t5_embeddings.dtype, device=TorchDevice.choose_torch_device()
+        )
+
+        return FluxRegionalTextConditioning(
+            t5_embeddings=t5_embeddings,
+            clip_embeddings=global_clip_embedding,
+            t5_txt_ids=t5_txt_ids,
+            image_masks=image_masks,
+            t5_embedding_ranges=concat_t5_embedding_ranges,
+        )
+
+    @staticmethod
+    def preprocess_regional_prompt_mask(
+        mask: Optional[torch.Tensor], packed_height: int, packed_width: int, dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        """Preprocess a regional prompt mask to match the target height and width.
+        If mask is None, returns a mask of all ones with the target height and width.
+        If mask is not None, resizes the mask to the target height and width using 'nearest' interpolation.
+
+        packed_height and packed_width are the target height and width of the mask in the 'packed' latent space.
+
+        Returns:
+            torch.Tensor: The processed mask. shape: (1, 1, packed_height * packed_width).
+        """
+
+        if mask is None:
+            return torch.ones((1, 1, packed_height * packed_width), dtype=dtype, device=device)
+
+        mask = to_standard_float_mask(mask, out_dtype=dtype)
+
+        tf = torchvision.transforms.Resize(
+            (packed_height, packed_width), interpolation=torchvision.transforms.InterpolationMode.NEAREST
+        )
+
+        # Add a batch dimension to the mask, because torchvision expects shape (batch, channels, h, w).
+        mask = mask.unsqueeze(0)  # Shape: (1, h, w) -> (1, 1, h, w)
+        resized_mask = tf(mask)
+
+        # Flatten the height and width dimensions into a single image_seq_len dimension.
+        return resized_mask.flatten(start_dim=2)

invokeai/backend/flux/ip_adapter/state_dict_utils.py

@@ -41,10 +41,12 @@ def infer_xlabs_ip_adapter_params_from_state_dict(state_dict: dict[str, torch.Te
     hidden_dim = state_dict["double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight"].shape[0]
     context_dim = state_dict["double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight"].shape[1]
     clip_embeddings_dim = state_dict["ip_adapter_proj_model.proj.weight"].shape[1]
+    clip_extra_context_tokens = state_dict["ip_adapter_proj_model.proj.weight"].shape[0] // context_dim
 
     return XlabsIpAdapterParams(
         num_double_blocks=num_double_blocks,
         context_dim=context_dim,
         hidden_dim=hidden_dim,
         clip_embeddings_dim=clip_embeddings_dim,
+        clip_extra_context_tokens=clip_extra_context_tokens,
     )

invokeai/backend/flux/ip_adapter/xlabs_ip_adapter_flux.py

@@ -31,13 +31,16 @@ class XlabsIpAdapterParams:
     hidden_dim: int
 
     clip_embeddings_dim: int
+    clip_extra_context_tokens: int
 
 
 class XlabsIpAdapterFlux(torch.nn.Module):
     def __init__(self, params: XlabsIpAdapterParams):
         super().__init__()
         self.image_proj = ImageProjModel(
-            cross_attention_dim=params.context_dim, clip_embeddings_dim=params.clip_embeddings_dim
+            cross_attention_dim=params.context_dim,
+            clip_embeddings_dim=params.clip_embeddings_dim,
+            clip_extra_context_tokens=params.clip_extra_context_tokens,
         )
         self.ip_adapter_double_blocks = IPAdapterDoubleBlocks(
             num_double_blocks=params.num_double_blocks, context_dim=params.context_dim, hidden_dim=params.hidden_dim

invokeai/backend/flux/math.py

@@ -5,10 +5,10 @@ from einops import rearrange
 from torch import Tensor
 
 
-def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, attn_mask: Tensor | None = None) -> Tensor:
     q, k = apply_rope(q, k, pe)
 
-    x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+    x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
     x = rearrange(x, "B H L D -> B L (H D)")
 
     return x
@@ -24,12 +24,12 @@ def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
     out = torch.einsum("...n,d->...nd", pos, omega)
     out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
     out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
-    return out.float()
+    return out.to(dtype=pos.dtype, device=pos.device)
 
 
 def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
-    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
-    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_ = xq.view(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.view(*xk.shape[:-1], -1, 1, 2)
     xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
     xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
-    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+    return xq_out.view(*xq.shape).type_as(xq), xk_out.view(*xk.shape).type_as(xk)

invokeai/backend/flux/model.py

@@ -1,11 +1,16 @@
 # 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 invokeai.backend.flux.custom_block_processor import CustomDoubleStreamBlockProcessor
+from invokeai.backend.flux.custom_block_processor import (
+    CustomDoubleStreamBlockProcessor,
+    CustomSingleStreamBlockProcessor,
+)
+from invokeai.backend.flux.extensions.regional_prompting_extension import RegionalPromptingExtension
 from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
 from invokeai.backend.flux.modules.layers import (
     DoubleStreamBlock,
@@ -31,6 +36,7 @@ class FluxParams:
     theta: int
     qkv_bias: bool
     guidance_embed: bool
+    out_channels: Optional[int] = None
 
 
 class Flux(nn.Module):
@@ -43,7 +49,7 @@ class Flux(nn.Module):
 
         self.params = params
         self.in_channels = params.in_channels
-        self.out_channels = self.in_channels
+        self.out_channels = params.out_channels or self.in_channels
         if params.hidden_size % params.num_heads != 0:
             raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}")
         pe_dim = params.hidden_size // params.num_heads
@@ -95,6 +101,7 @@ class Flux(nn.Module):
         controlnet_double_block_residuals: list[Tensor] | None,
         controlnet_single_block_residuals: list[Tensor] | None,
         ip_adapter_extensions: list[XLabsIPAdapterExtension],
+        regional_prompting_extension: RegionalPromptingExtension,
     ) -> Tensor:
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -117,7 +124,6 @@ class Flux(nn.Module):
             assert len(controlnet_double_block_residuals) == len(self.double_blocks)
         for block_index, block in enumerate(self.double_blocks):
             assert isinstance(block, DoubleStreamBlock)
-
             img, txt = CustomDoubleStreamBlockProcessor.custom_double_block_forward(
                 timestep_index=timestep_index,
                 total_num_timesteps=total_num_timesteps,
@@ -128,6 +134,7 @@ class Flux(nn.Module):
                 vec=vec,
                 pe=pe,
                 ip_adapter_extensions=ip_adapter_extensions,
+                regional_prompting_extension=regional_prompting_extension,
             )
 
             if controlnet_double_block_residuals is not None:
@@ -140,7 +147,17 @@ class Flux(nn.Module):
             assert len(controlnet_single_block_residuals) == len(self.single_blocks)
 
         for block_index, block in enumerate(self.single_blocks):
-            img = block(img, vec=vec, pe=pe)
+            assert isinstance(block, SingleStreamBlock)
+            img = CustomSingleStreamBlockProcessor.custom_single_block_forward(
+                timestep_index=timestep_index,
+                total_num_timesteps=total_num_timesteps,
+                block_index=block_index,
+                block=block,
+                img=img,
+                vec=vec,
+                pe=pe,
+                regional_prompting_extension=regional_prompting_extension,
+            )
 
             if controlnet_single_block_residuals is not None:
                 img[:, txt.shape[1] :, ...] += controlnet_single_block_residuals[block_index]

invokeai/backend/flux/modules/layers.py

@@ -66,10 +66,7 @@ class RMSNorm(torch.nn.Module):
         self.scale = nn.Parameter(torch.ones(dim))
 
     def forward(self, x: Tensor):
-        x_dtype = x.dtype
-        x = x.float()
-        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
-        return (x * rrms).to(dtype=x_dtype) * self.scale
+        return torch.nn.functional.rms_norm(x, self.scale.shape, self.scale, eps=1e-6)
 
 
 class QKNorm(torch.nn.Module):

invokeai/backend/flux/text_conditioning.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+
+import torch
+
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range
+
+
+@dataclass
+class FluxTextConditioning:
+    t5_embeddings: torch.Tensor
+    clip_embeddings: torch.Tensor
+    # If mask is None, the prompt is a global prompt.
+    mask: torch.Tensor | None
+
+
+@dataclass
+class FluxRegionalTextConditioning:
+    # Concatenated text embeddings.
+    # Shape: (1, concatenated_txt_seq_len, 4096)
+    t5_embeddings: torch.Tensor
+    # Shape: (1, concatenated_txt_seq_len, 3)
+    t5_txt_ids: torch.Tensor
+
+    # Global CLIP embeddings.
+    # Shape: (1, 768)
+    clip_embeddings: torch.Tensor
+
+    # A binary mask indicating the regions of the image that the prompt should be applied to. If None, the prompt is a
+    # global prompt.
+    # image_masks[i] is the mask for the ith prompt.
+    # image_masks[i] has shape (1, image_seq_len) and dtype torch.bool.
+    image_masks: list[torch.Tensor | None]
+
+    # List of ranges that represent the embedding ranges for each mask.
+    # t5_embedding_ranges[i] contains the range of the t5 embeddings that correspond to image_masks[i].
+    t5_embedding_ranges: list[Range]

invokeai/backend/lora/layers/any_lora_layer.py

@@ -1,11 +0,0 @@
-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
-
-AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer, NormLayer, ConcatenatedLoRALayer]

invokeai/backend/lora/lora_patcher.py

@@ -1,302 +0,0 @@
-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.layers.concatenated_lora_layer import ConcatenatedLoRALayer
-from invokeai.backend.lora.layers.lora_layer import LoRALayer
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
-from invokeai.backend.lora.sidecar_layers.concatenated_lora.concatenated_lora_linear_sidecar_layer import (
-    ConcatenatedLoRALinearSidecarLayer,
-)
-from invokeai.backend.lora.sidecar_layers.lora.lora_linear_sidecar_layer import LoRALinearSidecarLayer
-from invokeai.backend.lora.sidecar_layers.lora_sidecar_module import LoRASidecarModule
-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_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():
-                model.get_parameter(param_key).copy_(weight)
-
-    @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
-            )
-
-            # 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.)
-            device = module.weight.device
-            dtype = module.weight.dtype
-
-            layer_scale = layer.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(...)'.
-            layer.to(device=device)
-            layer.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 layer.get_parameters(module).items():
-                param_key = module_key + "." + param_name
-                module_param = module.get_parameter(param_name)
-
-                # Save original weight
-                original_weights.save(param_key, module_param)
-
-                if module_param.shape != lora_param_weight.shape:
-                    lora_param_weight = lora_param_weight.reshape(module_param.shape)
-
-                lora_param_weight *= patch_weight * layer_scale
-                module_param += lora_param_weight.to(dtype=dtype)
-
-            layer.to(device=TorchDevice.CPU_DEVICE)
-
-    @staticmethod
-    @torch.no_grad()
-    @contextmanager
-    def apply_lora_sidecar_patches(
-        model: torch.nn.Module,
-        patches: Iterable[Tuple[LoRAModelRaw, float]],
-        prefix: str,
-        dtype: torch.dtype,
-    ):
-        """Apply one or more LoRA sidecar patches to a model within a context manager. Sidecar patches incur some
-        overhead compared to normal LoRA patching, but they allow for LoRA layers to applied to base layers in any
-        quantization format.
-
-        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.
-            dtype (torch.dtype): The compute dtype of the sidecar layers. This cannot easily be inferred from the model,
-                since the sidecar layers are typically applied on top of quantized layers whose weight dtype is
-                different from their compute dtype.
-        """
-        original_modules: dict[str, torch.nn.Module] = {}
-        try:
-            for patch, patch_weight in patches:
-                LoRAPatcher._apply_lora_sidecar_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_sidecar_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 sidecar 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
-            )
-
-            # Initialize the LoRA sidecar layer.
-            lora_sidecar_layer = LoRAPatcher._initialize_lora_sidecar_layer(module, layer, patch_weight)
-
-            # Replace the original module with a LoRASidecarModule if it has not already been done.
-            if module_key in original_modules:
-                # The module has already been patched with a LoRASidecarModule. Append to it.
-                assert isinstance(module, LoRASidecarModule)
-                lora_sidecar_module = module
-            else:
-                # The module has not yet been patched with a LoRASidecarModule. Create one.
-                lora_sidecar_module = LoRASidecarModule(module, [])
-                original_modules[module_key] = module
-                module_parent_key, module_name = LoRAPatcher._split_parent_key(module_key)
-                module_parent = model.get_submodule(module_parent_key)
-                LoRAPatcher._set_submodule(module_parent, module_name, lora_sidecar_module)
-
-            # Move the LoRA sidecar layer to the same device/dtype as the orig module.
-            # TODO(ryand): Experiment with moving to the device first, then casting. This could be faster.
-            lora_sidecar_layer.to(device=lora_sidecar_module.orig_module.weight.device, dtype=dtype)
-
-            # Add the LoRA sidecar layer to the LoRASidecarModule.
-            lora_sidecar_module.add_lora_layer(lora_sidecar_layer)
-
-    @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_sidecar_layer(orig_layer: torch.nn.Module, lora_layer: AnyLoRALayer, patch_weight: float):
-        # TODO(ryand): Add support for more original layer types and LoRA layer types.
-        if isinstance(orig_layer, torch.nn.Linear) or (
-            isinstance(orig_layer, LoRASidecarModule) and isinstance(orig_layer.orig_module, torch.nn.Linear)
-        ):
-            if isinstance(lora_layer, LoRALayer):
-                return LoRALinearSidecarLayer(lora_layer=lora_layer, weight=patch_weight)
-            elif isinstance(lora_layer, ConcatenatedLoRALayer):
-                return ConcatenatedLoRALinearSidecarLayer(concatenated_lora_layer=lora_layer, weight=patch_weight)
-            else:
-                raise ValueError(f"Unsupported Linear LoRA layer type: {type(lora_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/lora/sidecar_layers/concatenated_lora/concatenated_lora_linear_sidecar_layer.py

@@ -1,34 +0,0 @@
-import torch
-
-from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
-
-
-class ConcatenatedLoRALinearSidecarLayer(torch.nn.Module):
-    def __init__(
-        self,
-        concatenated_lora_layer: ConcatenatedLoRALayer,
-        weight: float,
-    ):
-        super().__init__()
-
-        self._concatenated_lora_layer = concatenated_lora_layer
-        self._weight = weight
-
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        x_chunks: list[torch.Tensor] = []
-        for lora_layer in self._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 *= self._weight * lora_layer.scale()
-            x_chunks.append(x_chunk)
-
-        # TODO(ryand): Generalize to support concat_axis != 0.
-        assert self._concatenated_lora_layer.concat_axis == 0
-        x = torch.cat(x_chunks, dim=-1)
-        return x
-
-    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
-        self._concatenated_lora_layer.to(device=device, dtype=dtype)
-        return self

invokeai/backend/lora/sidecar_layers/lora/lora_linear_sidecar_layer.py

@@ -1,27 +0,0 @@
-import torch
-
-from invokeai.backend.lora.layers.lora_layer import LoRALayer
-
-
-class LoRALinearSidecarLayer(torch.nn.Module):
-    def __init__(
-        self,
-        lora_layer: LoRALayer,
-        weight: float,
-    ):
-        super().__init__()
-
-        self._lora_layer = lora_layer
-        self._weight = weight
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = torch.nn.functional.linear(x, self._lora_layer.down)
-        if self._lora_layer.mid is not None:
-            x = torch.nn.functional.linear(x, self._lora_layer.mid)
-        x = torch.nn.functional.linear(x, self._lora_layer.up, bias=self._lora_layer.bias)
-        x *= self._weight * self._lora_layer.scale()
-        return x
-
-    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
-        self._lora_layer.to(device=device, dtype=dtype)
-        return self

invokeai/backend/lora/sidecar_layers/lora_sidecar_module.py

@@ -1,24 +0,0 @@
-import torch
-
-
-class LoRASidecarModule(torch.nn.Module):
-    """A LoRA sidecar module that wraps an original module and adds LoRA layers to it."""
-
-    def __init__(self, orig_module: torch.nn.Module, lora_layers: list[torch.nn.Module]):
-        super().__init__()
-        self.orig_module = orig_module
-        self._lora_layers = lora_layers
-
-    def add_lora_layer(self, lora_layer: torch.nn.Module):
-        self._lora_layers.append(lora_layer)
-
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        x = self.orig_module(input)
-        for lora_layer in self._lora_layers:
-            x += lora_layer(input)
-        return x
-
-    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
-        self._orig_module.to(device=device, dtype=dtype)
-        for lora_layer in self._lora_layers:
-            lora_layer.to(device=device, dtype=dtype)

invokeai/backend/model_manager/config.py

@@ -67,6 +67,7 @@ class ModelType(str, Enum):
     Main = "main"
     VAE = "vae"
     LoRA = "lora"
+    ControlLoRa = "control_lora"
     ControlNet = "controlnet"  # used by model_probe
     TextualInversion = "embedding"
     IPAdapter = "ip_adapter"
@@ -273,6 +274,24 @@ 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 ControlLoRALyCORISConfig(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.LyCORIS] = ModelFormat.LyCORIS
+
+    @staticmethod
+    def get_tag() -> Tag:
+        return Tag(f"{ModelType.ControlLoRa.value}.{ModelFormat.LyCORIS.value}")
+
+
 class LoRADiffusersConfig(LoRAConfigBase):
     """Model config for LoRA/Diffusers models."""
 
@@ -304,12 +323,6 @@ 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)."""
 
@@ -535,6 +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[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_default import ModelCache
+from invokeai.backend.model_manager.load.model_cache.model_cache 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,7 +5,6 @@ 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
@@ -18,19 +17,17 @@ from invokeai.backend.model_manager.config import (
     AnyModelConfig,
     SubModelType,
 )
-from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase
+from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
 
 
-@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:
     ```
@@ -40,13 +37,9 @@ 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:
     ```
@@ -55,43 +48,53 @@ 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.
     """
 
-    _locker: ModelLockerBase
+    def __init__(self, cache_record: CacheRecord, cache: ModelCache):
+        self._cache_record = cache_record
+        self._cache = cache
 
     def __enter__(self) -> AnyModel:
-        """Context entry."""
-        self._locker.lock()
+        self._cache.lock(self._cache_record.key, None)
         return self.model
 
     def __exit__(self, *args: Any, **kwargs: Any) -> None:
-        """Context exit."""
-        self._locker.unlock()
+        self._cache.unlock(self._cache_record.key)
 
     @contextmanager
-    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()
+    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.key, working_mem_bytes)
         try:
-            state_dict = self._locker.get_state_dict()
-            yield (state_dict, locked_model)
+            yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model)
         finally:
-            self._locker.unlock()
+            self._cache.unlock(self._cache_record.key)
 
     @property
     def model(self) -> AnyModel:
         """Return the model without locking it."""
-        return self._locker.model
+        return self._cache_record.cached_model.model
 
 
-@dataclass
 class LoadedModel(LoadedModelWithoutConfig):
     """Context manager object that mediates transfer from RAM<->VRAM."""
 
-    config: Optional[AnyModelConfig] = None
+    def __init__(
+        self,
+        config: Optional[AnyModelConfig],
+        cache_record: CacheRecord,
+        cache: ModelCache,
+    ):
+        super().__init__(cache_record=cache_record, cache=cache)
+        self.config = config
 
 
 # TODO(MM2):
@@ -110,13 +113,17 @@ class ModelLoaderBase(ABC):
         self,
         app_config: InvokeAIAppConfig,
         logger: Logger,
-        ram_cache: ModelCacheBase[AnyModel],
+        ram_cache: ModelCache,
     ):
         """Initialize the loader."""
         pass
 
     @abstractmethod
-    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+    def load_model(
+        self,
+        model_config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> LoadedModel:
         """
         Return a model given its confguration.
 
@@ -138,6 +145,6 @@ class ModelLoaderBase(ABC):
 
     @property
     @abstractmethod
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
+    def ram_cache(self) -> ModelCache:
         """Return the ram cache associated with this loader."""
         pass

invokeai/backend/model_manager/load/load_default.py

@@ -14,7 +14,8 @@ 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.model_cache_base import ModelCacheBase, ModelLockerBase
+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_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
@@ -28,7 +29,7 @@ class ModelLoader(ModelLoaderBase):
         self,
         app_config: InvokeAIAppConfig,
         logger: Logger,
-        ram_cache: ModelCacheBase[AnyModel],
+        ram_cache: ModelCache,
     ):
         """Initialize the loader."""
         self._app_config = app_config
@@ -37,7 +38,11 @@ class ModelLoader(ModelLoaderBase):
         self._torch_dtype = TorchDevice.choose_torch_dtype()
         self._torch_device = TorchDevice.choose_torch_device()
 
-    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+    def load_model(
+        self,
+        model_config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> LoadedModel:
         """
         Return a model given its configuration.
 
@@ -54,11 +59,11 @@ class ModelLoader(ModelLoaderBase):
             raise InvalidModelConfigException(f"Files for model '{model_config.name}' not found at {model_path}")
 
         with skip_torch_weight_init():
-            locker = self._load_and_cache(model_config, submodel_type)
-        return LoadedModel(config=model_config, _locker=locker)
+            cache_record = self._load_and_cache(model_config, submodel_type)
+        return LoadedModel(config=model_config, cache_record=cache_record, cache=self._ram_cache)
 
     @property
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
+    def ram_cache(self) -> ModelCache:
         """Return the ram cache associated with this loader."""
         return self._ram_cache
 
@@ -66,10 +71,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) -> ModelLockerBase:
+    def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> CacheRecord:
         stats_name = ":".join([config.base, config.type, config.name, (submodel_type or "")])
         try:
-            return self._ram_cache.get(config.key, submodel_type, stats_name=stats_name)
+            return self._ram_cache.get(key=get_model_cache_key(config.key, submodel_type), stats_name=stats_name)
         except IndexError:
             pass
 
@@ -78,16 +83,11 @@ class ModelLoader(ModelLoaderBase):
         loaded_model = self._load_model(config, submodel_type)
 
         self._ram_cache.put(
-            config.key,
-            submodel_type=submodel_type,
+            get_model_cache_key(config.key, submodel_type),
             model=loaded_model,
         )
 
-        return self._ram_cache.get(
-            key=config.key,
-            submodel_type=submodel_type,
-            stats_name=stats_name,
-        )
+        return self._ram_cache.get(key=get_model_cache_key(config.key, 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

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

@@ -0,0 +1,31 @@
+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
+    # If locks > 0, the model is actively being used, so we should do our best to keep it on the compute device.
+    _locks: int = 0
+
+    def lock(self) -> None:
+        self._locks += 1
+
+    def unlock(self) -> None:
+        self._locks -= 1
+        assert self._locks >= 0
+
+    @property
+    def is_locked(self) -> bool:
+        return self._locks > 0

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

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

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

@@ -0,0 +1,157 @@
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_function_autocast_context import (
+    add_autocast_to_module_forward,
+    remove_autocast_from_module_forward,
+)
+from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+
+
+def set_nested_attr(obj: object, attr: str, value: object):
+    """A helper function that extends setattr() to support nested attributes.
+
+    Example:
+        set_nested_attr(model, "module.encoder.conv1.weight", new_conv1_weight)
+    """
+    attrs = attr.split(".")
+    for attr in attrs[:-1]:
+        obj = getattr(obj, attr)
+    setattr(obj, attrs[-1], value)
+
+
+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()
+
+        # TODO(ryand): Handle the case where the model sizes changes after initial load (e.g. due to dtype casting).
+        # Consider how we should handle this for both self._total_bytes and self._cur_vram_bytes.
+        self._total_bytes = sum(calc_tensor_size(p) for p in self._cpu_state_dict.values())
+        self._cur_vram_bytes: int | None = None
+
+        self._update_model_autocast_context()
+
+    @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(
+                calc_tensor_size(p) for p in cur_state_dict.values() if p.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()
+
+        for key, param in cur_state_dict.items():
+            if param.device.type == self._compute_device.type:
+                continue
+
+            param_size = calc_tensor_size(param)
+            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?
+                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 self._cur_vram_bytes == self.total_bytes():
+            # HACK(ryand): The model should already be on the compute device, but we have to call this to ensure that
+            # all non-persistent buffers are moved (i.e. buffers that are not registered in the state dict).
+            self._model.to(self._compute_device)
+
+        self._update_model_autocast_context()
+        return vram_bytes_loaded
+
+    @torch.no_grad()
+    def partial_unload_from_vram(self, vram_bytes_to_free: int) -> int:
+        """Unload weights from VRAM until vram_bytes_to_free bytes are freed. Or the entire model is unloaded.
+
+        Returns:
+            The number of bytes unloaded from VRAM.
+        """
+        vram_bytes_freed = 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
+
+            cur_state_dict[key] = self._cpu_state_dict[key]
+            vram_bytes_freed += calc_tensor_size(param)
+
+        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
+
+        self._update_model_autocast_context()
+        return vram_bytes_freed
+
+    def _update_model_autocast_context(self):
+        """A helper function that should be called whenever the model's VRAM usage changes to add/remove the autocast
+        context.
+        """
+        if self.cur_vram_bytes() == self.total_bytes():
+            # We remove the autocast context when the model is fully loaded into VRAM, because the context causes some
+            # runtime overhead.
+            remove_autocast_from_module_forward(self._model)
+        else:
+            # Monkey-patch the model to add autocasting to the model's forward method.
+            add_autocast_to_module_forward(self._model, self._compute_device)

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

@@ -0,0 +1,554 @@
+import gc
+import logging
+import time
+from logging import Logger
+from typing import Dict, List, Optional
+
+import psutil
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+from invokeai.backend.flux.ip_adapter.xlabs_ip_adapter_flux import XlabsIpAdapterFlux
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+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_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)
+
+        # 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"
+
+        # Specific models that opt-out of partial loading.
+        partial_loading_opt_out_models = (
+            # The following models have multiple entrypoints. Our auto-casting context management is only applied to the
+            # forward method, so a partially loaded AutoEncoder could fail if another entrypoint is used. These models
+            # can be supported in the future by improving the autocast context management.
+            # AutoEncoder has three entrypoints: encode, decode, and forward.
+            AutoEncoder,
+            # XLabsIPAdapterFlux is a wrapper around two models that are called directly.
+            XlabsIpAdapterFlux,
+            AutoencoderKL,
+        )
+
+        # Wrap model.
+        if (
+            isinstance(model, torch.nn.Module)
+            and running_with_cuda
+            and self._enable_partial_loading
+            and not isinstance(model, partial_loading_opt_out_models)
+        ):
+            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, key: str, working_mem_bytes: Optional[int]) -> None:
+        """Lock a model for use and move it into VRAM.
+
+        :param working_mem_bytes: The number of bytes of working memory to keep on the GPU while this model is loaded on
+            the GPU. If None, self._execution_device_working_mem_gb is used.
+        """
+        cache_entry = self._cached_models[key]
+        cache_entry.lock()
+
+        self._logger.debug(f"Locking model {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 {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, key: str) -> None:
+        """Unlock a model."""
+        cache_entry = self._cached_models[key]
+        cache_entry.unlock()
+        self._logger.debug(f"Unlocked model {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()
+        vram_available = self._get_vram_available(working_mem_bytes)
+
+        # 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
+
+        # The amount of VRAM that must be freed to make room for model_vram_needed.
+        vram_bytes_to_free = max(0, model_vram_needed - vram_available)
+
+        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(vram_bytes_to_free)
+        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)}"
+        )
+
+        # 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)
+        model_bytes_loaded = self._move_model_to_vram(cache_entry, vram_available)
+
+        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)
+            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] = None) -> 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":
+            vram_reserved = torch.cuda.memory_reserved(self._execution_device)
+            vram_free, _vram_total = torch.cuda.mem_get_info(self._execution_device)
+            vram_available_to_process = vram_free + vram_reserved
+        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."""
+        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.
+        return int(ram_total * 0.9) - ram_used
+
+    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_to_free: int) -> int:
+        """Offload models from the execution_device until vram_bytes_to_free bytes are freed, or all models are
+        offloaded. Of course, locked models are not offloaded.
+
+        Returns:
+            int: The number of bytes freed.
+        """
+        self._logger.debug(f"Offloading unlocked models with goal of freeing {vram_bytes_to_free/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:
+            if vram_bytes_freed >= vram_bytes_to_free:
+                break
+            if cache_entry.is_locked:
+                continue
+
+            cache_entry_bytes_freed = self._move_model_to_ram(cache_entry, vram_bytes_to_free - vram_bytes_freed)
+            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()
+            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

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

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

@@ -1,64 +0,0 @@
-"""
-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_function_autocast_context.py

@@ -0,0 +1,42 @@
+from typing import Any, Callable
+
+import torch
+from torch.overrides import TorchFunctionMode
+
+
+def add_autocast_to_module_forward(m: torch.nn.Module, to_device: torch.device):
+    """Monkey-patch m.forward(...) with a new forward(...) method that activates device autocasting for its duration."""
+    old_forward = m.forward
+
+    def new_forward(*args: Any, **kwargs: Any):
+        with TorchFunctionAutocastDeviceContext(to_device):
+            return old_forward(*args, **kwargs)
+
+    m.old_forward = old_forward  # type: ignore
+    m.forward = new_forward
+
+
+def remove_autocast_from_module_forward(m: torch.nn.Module):
+    """Remove the autocast context from m.forward(...) and restore the old forward method."""
+    if not hasattr(m, "old_forward"):
+        return
+    m.forward = m.old_forward
+    del m.old_forward
+
+
+def _cast_to_device_and_run(
+    func: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any], to_device: torch.device
+):
+    args_on_device = [a.to(to_device) if isinstance(a, torch.Tensor) else a for a in args]
+    kwargs_on_device = {k: v.to(to_device) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
+    return func(*args_on_device, **kwargs_on_device)
+
+
+class TorchFunctionAutocastDeviceContext(TorchFunctionMode):
+    def __init__(self, to_device: torch.device):
+        self._to_device = to_device
+
+    def __torch_function__(
+        self, func: Callable[..., Any], types, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None
+    ):
+        return _cast_to_device_and_run(func, args, kwargs or {}, self._to_device)

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

@@ -9,14 +9,6 @@ import torch
 from safetensors.torch import load_file
 
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.backend.lora.conversions.flux_diffusers_lora_conversion_utils import (
-    lora_model_from_flux_diffusers_state_dict,
-)
-from invokeai.backend.lora.conversions.flux_kohya_lora_conversion_utils import (
-    lora_model_from_flux_kohya_state_dict,
-)
-from invokeai.backend.lora.conversions.sd_lora_conversion_utils import lora_model_from_sd_state_dict
-from invokeai.backend.lora.conversions.sdxl_lora_conversion_utils import convert_sdxl_keys_to_diffusers_format
 from invokeai.backend.model_manager import (
     AnyModel,
     AnyModelConfig,
@@ -26,12 +18,26 @@ from invokeai.backend.model_manager import (
     SubModelType,
 )
 from invokeai.backend.model_manager.load.load_default import ModelLoader
-from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
+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.patches.lora_conversions.flux_control_lora_utils import (
+    is_state_dict_likely_flux_control,
+    lora_model_from_flux_control_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import (
+    lora_model_from_flux_diffusers_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_kohya_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_kohya_format,
+    lora_model_from_flux_kohya_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.sd_lora_conversion_utils import lora_model_from_sd_state_dict
+from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import convert_sdxl_keys_to_diffusers_format
 
 
 @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LoRA, format=ModelFormat.Diffusers)
 @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LoRA, format=ModelFormat.LyCORIS)
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.ControlLoRa, format=ModelFormat.LyCORIS)
 class LoRALoader(ModelLoader):
     """Class to load LoRA models."""
 
@@ -40,7 +46,7 @@ class LoRALoader(ModelLoader):
         self,
         app_config: InvokeAIAppConfig,
         logger: Logger,
-        ram_cache: ModelCacheBase[AnyModel],
+        ram_cache: ModelCache,
     ):
         """Initialize the loader."""
         super().__init__(app_config, logger, ram_cache)
@@ -75,7 +81,10 @@ class LoRALoader(ModelLoader):
                 # https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_flux.py#L1194
                 model = lora_model_from_flux_diffusers_state_dict(state_dict=state_dict, alpha=None)
             elif config.format == ModelFormat.LyCORIS:
-                model = lora_model_from_flux_kohya_state_dict(state_dict=state_dict)
+                if is_state_dict_likely_in_flux_kohya_format(state_dict=state_dict):
+                    model = lora_model_from_flux_kohya_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_flux_control(state_dict=state_dict):
+                    model = lora_model_from_flux_control_state_dict(state_dict=state_dict)
             else:
                 raise ValueError(f"LoRA model is in unsupported FLUX format: {config.format}")
         elif self._model_base in [BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]:

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

@@ -25,6 +25,7 @@ from invokeai.backend.model_manager.config import (
     DiffusersConfigBase,
     MainCheckpointConfig,
 )
+from invokeai.backend.model_manager.load.model_cache.model_cache import get_model_cache_key
 from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
 from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
 from invokeai.backend.util.silence_warnings import SilenceWarnings
@@ -132,5 +133,5 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
             if subtype == submodel_type:
                 continue
             if submodel := getattr(pipeline, subtype.value, None):
-                self._ram_cache.put(config.key, submodel_type=subtype, model=submodel)
+                self._ram_cache.put(get_model_cache_key(config.key, subtype), model=submodel)
         return getattr(pipeline, submodel_type.value)

invokeai/backend/model_manager/load/model_util.py

@@ -15,9 +15,9 @@ from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import D
 from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
 from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
 from invokeai.backend.model_manager.config import AnyModel
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
 from invokeai.backend.textual_inversion import TextualInversionModelRaw
 from invokeai.backend.util.calc_tensor_size import calc_tensor_size
@@ -43,7 +43,7 @@ def calc_model_size_by_data(logger: logging.Logger, model: AnyModel) -> int:
         (
             TextualInversionModelRaw,
             IPAdapter,
-            LoRAModelRaw,
+            ModelPatchRaw,
             SpandrelImageToImageModel,
             GroundingDinoPipeline,
             SegmentAnythingPipeline,

invokeai/backend/model_manager/probe.py

@@ -15,10 +15,6 @@ from invokeai.backend.flux.controlnet.state_dict_utils import (
     is_state_dict_xlabs_controlnet,
 )
 from invokeai.backend.flux.ip_adapter.state_dict_utils import is_state_dict_xlabs_ip_adapter
-from invokeai.backend.lora.conversions.flux_diffusers_lora_conversion_utils import (
-    is_state_dict_likely_in_flux_diffusers_format,
-)
-from invokeai.backend.lora.conversions.flux_kohya_lora_conversion_utils import is_state_dict_likely_in_flux_kohya_format
 from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
 from invokeai.backend.model_manager.config import (
     AnyModelConfig,
@@ -43,6 +39,13 @@ from invokeai.backend.model_manager.util.model_util import (
     lora_token_vector_length,
     read_checkpoint_meta,
 )
+from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import is_state_dict_likely_flux_control
+from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_diffusers_format,
+)
+from invokeai.backend.patches.lora_conversions.flux_kohya_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_kohya_format,
+)
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
 from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
 from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
@@ -199,8 +202,8 @@ class ModelProbe(object):
         fields["default_settings"] = fields.get("default_settings")
 
         if not fields["default_settings"]:
-            if fields["type"] in {ModelType.ControlNet, ModelType.T2IAdapter}:
-                fields["default_settings"] = get_default_settings_controlnet_t2i_adapter(fields["name"])
+            if fields["type"] in {ModelType.ControlNet, ModelType.T2IAdapter, ModelType.ControlLoRa}:
+                fields["default_settings"] = get_default_settings_control_adapters(fields["name"])
             elif fields["type"] is ModelType.Main:
                 fields["default_settings"] = get_default_settings_main(fields["base"])
 
@@ -258,6 +261,9 @@ class ModelProbe(object):
         ckpt = checkpoint if checkpoint else read_checkpoint_meta(model_path, scan=True)
         ckpt = ckpt.get("state_dict", ckpt)
 
+        if isinstance(ckpt, dict) and is_state_dict_likely_flux_control(ckpt):
+            return ModelType.ControlLoRa
+
         for key in [str(k) for k in ckpt.keys()]:
             if key.startswith(
                 (
@@ -469,7 +475,7 @@ class ModelProbe(object):
         """
         # scan model
         scan_result = scan_file_path(checkpoint)
-        if scan_result.infected_files != 0:
+        if scan_result.infected_files != 0 or scan_result.scan_err:
             raise Exception("The model {model_name} is potentially infected by malware. Aborting import.")
 
 
@@ -485,6 +491,7 @@ MODEL_NAME_TO_PREPROCESSOR = {
     "lineart anime": "lineart_anime_image_processor",
     "lineart_anime": "lineart_anime_image_processor",
     "lineart": "lineart_image_processor",
+    "soft": "hed_image_processor",
     "softedge": "hed_image_processor",
     "hed": "hed_image_processor",
     "shuffle": "content_shuffle_image_processor",
@@ -496,7 +503,7 @@ MODEL_NAME_TO_PREPROCESSOR = {
 }
 
 
-def get_default_settings_controlnet_t2i_adapter(model_name: str) -> Optional[ControlAdapterDefaultSettings]:
+def get_default_settings_control_adapters(model_name: str) -> Optional[ControlAdapterDefaultSettings]:
     for k, v in MODEL_NAME_TO_PREPROCESSOR.items():
         model_name_lower = model_name.lower()
         if k in model_name_lower:
@@ -623,8 +630,10 @@ class LoRACheckpointProbe(CheckpointProbeBase):
         return ModelFormat.LyCORIS
 
     def get_base_type(self) -> BaseModelType:
-        if is_state_dict_likely_in_flux_kohya_format(self.checkpoint) or is_state_dict_likely_in_flux_diffusers_format(
-            self.checkpoint
+        if (
+            is_state_dict_likely_in_flux_kohya_format(self.checkpoint)
+            or is_state_dict_likely_in_flux_diffusers_format(self.checkpoint)
+            or is_state_dict_likely_flux_control(self.checkpoint)
         ):
             return BaseModelType.Flux
 
@@ -1045,6 +1054,7 @@ ModelProbe.register_probe("diffusers", ModelType.SpandrelImageToImage, SpandrelI
 ModelProbe.register_probe("checkpoint", ModelType.Main, PipelineCheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.VAE, VaeCheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.LoRA, LoRACheckpointProbe)
+ModelProbe.register_probe("checkpoint", ModelType.ControlLoRa, LoRACheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.TextualInversion, TextualInversionCheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.ControlNet, ControlNetCheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.IPAdapter, IPAdapterCheckpointProbe)

invokeai/backend/model_manager/starter_models.py

@@ -298,13 +298,12 @@ ip_adapter_sdxl = StarterModel(
     previous_names=["IP Adapter SDXL"],
 )
 ip_adapter_flux = StarterModel(
-    name="Standard Reference (XLabs FLUX IP-Adapter)",
+    name="Standard Reference (XLabs FLUX IP-Adapter v2)",
     base=BaseModelType.Flux,
-    source="https://huggingface.co/XLabs-AI/flux-ip-adapter/resolve/main/ip_adapter.safetensors",
+    source="https://huggingface.co/XLabs-AI/flux-ip-adapter-v2/resolve/main/ip_adapter.safetensors",
     description="References images with a more generalized/looser degree of precision.",
     type=ModelType.IPAdapter,
     dependencies=[clip_vit_l_image_encoder],
-    previous_names=["XLabs FLUX IP-Adapter"],
 )
 # endregion
 # region ControlNet
@@ -489,6 +488,22 @@ union_cnet_flux = StarterModel(
     type=ModelType.ControlNet,
 )
 # endregion
+# region Control LoRA
+flux_canny_control_lora = StarterModel(
+    name="Hard Edge Detection (Canny)",
+    base=BaseModelType.Flux,
+    source="black-forest-labs/FLUX.1-Canny-dev-lora::flux1-canny-dev-lora.safetensors",
+    description="Uses detected edges in the image to control composition.",
+    type=ModelType.ControlLoRa,
+)
+flux_depth_control_lora = StarterModel(
+    name="Depth Map",
+    base=BaseModelType.Flux,
+    source="black-forest-labs/FLUX.1-Depth-dev-lora::flux1-depth-dev-lora.safetensors",
+    description="Uses depth information in the image to control the depth in the generation.",
+    type=ModelType.ControlLoRa,
+)
+# endregion
 # region T2I Adapter
 t2i_canny_sd1 = StarterModel(
     name="Hard Edge Detection (canny)",
@@ -631,6 +646,8 @@ STARTER_MODELS: list[StarterModel] = [
     tile_sdxl,
     union_cnet_sdxl,
     union_cnet_flux,
+    flux_canny_control_lora,
+    flux_depth_control_lora,
     t2i_canny_sd1,
     t2i_sketch_sd1,
     t2i_depth_sd1,
@@ -689,6 +706,8 @@ flux_bundle: list[StarterModel] = [
     clip_l_encoder,
     union_cnet_flux,
     ip_adapter_flux,
+    flux_canny_control_lora,
+    flux_depth_control_lora,
 ]
 
 STARTER_BUNDLES: dict[str, list[StarterModel]] = {

invokeai/backend/model_manager/util/model_util.py

@@ -44,7 +44,7 @@ def _fast_safetensors_reader(path: str) -> Dict[str, torch.Tensor]:
     return checkpoint
 
 
-def read_checkpoint_meta(path: Union[str, Path], scan: bool = False) -> Dict[str, torch.Tensor]:
+def read_checkpoint_meta(path: Union[str, Path], scan: bool = True) -> Dict[str, torch.Tensor]:
     if str(path).endswith(".safetensors"):
         try:
             path_str = path.as_posix() if isinstance(path, Path) else path
@@ -52,16 +52,15 @@ def read_checkpoint_meta(path: Union[str, Path], scan: bool = False) -> Dict[str
         except Exception:
             # TODO: create issue for support "meta"?
             checkpoint = safetensors.torch.load_file(path, device="cpu")
+    elif str(path).endswith(".gguf"):
+        # The GGUF reader used here uses numpy memmap, so these tensors are not loaded into memory during this function
+        checkpoint = gguf_sd_loader(Path(path), compute_dtype=torch.float32)
     else:
         if scan:
             scan_result = scan_file_path(path)
-            if scan_result.infected_files != 0:
+            if scan_result.infected_files != 0 or scan_result.scan_err:
                 raise Exception(f'The model file "{path}" is potentially infected by malware. Aborting import.')
-        if str(path).endswith(".gguf"):
-            # The GGUF reader used here uses numpy memmap, so these tensors are not loaded into memory during this function
-            checkpoint = gguf_sd_loader(Path(path), compute_dtype=torch.float32)
-        else:
-            checkpoint = torch.load(path, map_location=torch.device("meta"))
+        checkpoint = torch.load(path, map_location=torch.device("meta"))
     return checkpoint
 
 

invokeai/backend/model_patcher.py

@@ -5,17 +5,14 @@ from __future__ import annotations
 
 import pickle
 from contextlib import contextmanager
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union
+from typing import Any, Iterator, List, Optional, Tuple, Type, Union
 
-import numpy as np
 import torch
 from diffusers import UNet2DConditionModel
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
 from invokeai.app.shared.models import FreeUConfig
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
-from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
 from invokeai.backend.textual_inversion import TextualInversionManager, TextualInversionModelRaw
 
 
@@ -176,180 +173,3 @@ class ModelPatcher:
             assert hasattr(unet, "disable_freeu")  # mypy doesn't pick up this attribute?
             if did_apply_freeu:
                 unet.disable_freeu()
-
-
-class ONNXModelPatcher:
-    # based on
-    # https://github.com/ssube/onnx-web/blob/ca2e436f0623e18b4cfe8a0363fcfcf10508acf7/api/onnx_web/convert/diffusion/lora.py#L323
-    @classmethod
-    @contextmanager
-    def apply_lora(
-        cls,
-        model: IAIOnnxRuntimeModel,
-        loras: List[Tuple[LoRAModelRaw, float]],
-        prefix: str,
-    ) -> None:
-        from invokeai.backend.models.base import IAIOnnxRuntimeModel
-
-        if not isinstance(model, IAIOnnxRuntimeModel):
-            raise Exception("Only IAIOnnxRuntimeModel models supported")
-
-        orig_weights = {}
-
-        try:
-            blended_loras: Dict[str, torch.Tensor] = {}
-
-            for lora, lora_weight in loras:
-                for layer_key, layer in lora.layers.items():
-                    if not layer_key.startswith(prefix):
-                        continue
-
-                    layer.to(dtype=torch.float32)
-                    layer_key = layer_key.replace(prefix, "")
-                    # TODO: rewrite to pass original tensor weight(required by ia3)
-                    layer_weight = layer.get_weight(None).detach().cpu().numpy() * lora_weight
-                    if layer_key in blended_loras:
-                        blended_loras[layer_key] += layer_weight
-                    else:
-                        blended_loras[layer_key] = layer_weight
-
-            node_names = {}
-            for node in model.nodes.values():
-                node_names[node.name.replace("/", "_").replace(".", "_").lstrip("_")] = node.name
-
-            for layer_key, lora_weight in blended_loras.items():
-                conv_key = layer_key + "_Conv"
-                gemm_key = layer_key + "_Gemm"
-                matmul_key = layer_key + "_MatMul"
-
-                if conv_key in node_names or gemm_key in node_names:
-                    if conv_key in node_names:
-                        conv_node = model.nodes[node_names[conv_key]]
-                    else:
-                        conv_node = model.nodes[node_names[gemm_key]]
-
-                    weight_name = [n for n in conv_node.input if ".weight" in n][0]
-                    orig_weight = model.tensors[weight_name]
-
-                    if orig_weight.shape[-2:] == (1, 1):
-                        if lora_weight.shape[-2:] == (1, 1):
-                            new_weight = orig_weight.squeeze((3, 2)) + lora_weight.squeeze((3, 2))
-                        else:
-                            new_weight = orig_weight.squeeze((3, 2)) + lora_weight
-
-                        new_weight = np.expand_dims(new_weight, (2, 3))
-                    else:
-                        if orig_weight.shape != lora_weight.shape:
-                            new_weight = orig_weight + lora_weight.reshape(orig_weight.shape)
-                        else:
-                            new_weight = orig_weight + lora_weight
-
-                    orig_weights[weight_name] = orig_weight
-                    model.tensors[weight_name] = new_weight.astype(orig_weight.dtype)
-
-                elif matmul_key in node_names:
-                    weight_node = model.nodes[node_names[matmul_key]]
-                    matmul_name = [n for n in weight_node.input if "MatMul" in n][0]
-
-                    orig_weight = model.tensors[matmul_name]
-                    new_weight = orig_weight + lora_weight.transpose()
-
-                    orig_weights[matmul_name] = orig_weight
-                    model.tensors[matmul_name] = new_weight.astype(orig_weight.dtype)
-
-                else:
-                    # warn? err?
-                    pass
-
-            yield
-
-        finally:
-            # restore original weights
-            for name, orig_weight in orig_weights.items():
-                model.tensors[name] = orig_weight
-
-    @classmethod
-    @contextmanager
-    def apply_ti(
-        cls,
-        tokenizer: CLIPTokenizer,
-        text_encoder: IAIOnnxRuntimeModel,
-        ti_list: List[Tuple[str, Any]],
-    ) -> Iterator[Tuple[CLIPTokenizer, TextualInversionManager]]:
-        from invokeai.backend.models.base import IAIOnnxRuntimeModel
-
-        if not isinstance(text_encoder, IAIOnnxRuntimeModel):
-            raise Exception("Only IAIOnnxRuntimeModel models supported")
-
-        orig_embeddings = None
-
-        try:
-            # HACK: The CLIPTokenizer API does not include a way to remove tokens after calling add_tokens(...). As a
-            # workaround, we create a full copy of `tokenizer` so that its original behavior can be restored after
-            # exiting this `apply_ti(...)` context manager.
-            #
-            # In a previous implementation, the deep copy was obtained with `ti_tokenizer = copy.deepcopy(tokenizer)`,
-            # but a pickle roundtrip was found to be much faster (1 sec vs. 0.05 secs).
-            ti_tokenizer = pickle.loads(pickle.dumps(tokenizer))
-            ti_manager = TextualInversionManager(ti_tokenizer)
-
-            def _get_trigger(ti_name: str, index: int) -> str:
-                trigger = ti_name
-                if index > 0:
-                    trigger += f"-!pad-{i}"
-                return f"<{trigger}>"
-
-            # modify text_encoder
-            orig_embeddings = text_encoder.tensors["text_model.embeddings.token_embedding.weight"]
-
-            # modify tokenizer
-            new_tokens_added = 0
-            for ti_name, ti in ti_list:
-                if ti.embedding_2 is not None:
-                    ti_embedding = (
-                        ti.embedding_2 if ti.embedding_2.shape[1] == orig_embeddings.shape[0] else ti.embedding
-                    )
-                else:
-                    ti_embedding = ti.embedding
-
-                for i in range(ti_embedding.shape[0]):
-                    new_tokens_added += ti_tokenizer.add_tokens(_get_trigger(ti_name, i))
-
-            embeddings = np.concatenate(
-                (np.copy(orig_embeddings), np.zeros((new_tokens_added, orig_embeddings.shape[1]))),
-                axis=0,
-            )
-
-            for ti_name, _ in ti_list:
-                ti_tokens = []
-                for i in range(ti_embedding.shape[0]):
-                    embedding = ti_embedding[i].detach().numpy()
-                    trigger = _get_trigger(ti_name, i)
-
-                    token_id = ti_tokenizer.convert_tokens_to_ids(trigger)
-                    if token_id == ti_tokenizer.unk_token_id:
-                        raise RuntimeError(f"Unable to find token id for token '{trigger}'")
-
-                    if embeddings[token_id].shape != embedding.shape:
-                        raise ValueError(
-                            f"Cannot load embedding for {trigger}. It was trained on a model with token dimension"
-                            f" {embedding.shape[0]}, but the current model has token dimension"
-                            f" {embeddings[token_id].shape[0]}."
-                        )
-
-                    embeddings[token_id] = embedding
-                    ti_tokens.append(token_id)
-
-                if len(ti_tokens) > 1:
-                    ti_manager.pad_tokens[ti_tokens[0]] = ti_tokens[1:]
-
-            text_encoder.tensors["text_model.embeddings.token_embedding.weight"] = embeddings.astype(
-                orig_embeddings.dtype
-            )
-
-            yield ti_tokenizer, ti_manager
-
-        finally:
-            # restore
-            if orig_embeddings is not None:
-                text_encoder.tensors["text_model.embeddings.token_embedding.weight"] = orig_embeddings

invokeai/backend/patches/layer_patcher.py

@@ -0,0 +1,290 @@
+from contextlib import contextmanager
+from typing import Dict, Iterable, Optional, Tuple
+
+import torch
+
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.flux_control_lora_layer import FluxControlLoRALayer
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.patches.pad_with_zeros import pad_with_zeros
+from invokeai.backend.patches.sidecar_wrappers.base_sidecar_wrapper import BaseSidecarWrapper
+from invokeai.backend.patches.sidecar_wrappers.utils import wrap_module_with_sidecar_wrapper
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+
+
+class LayerPatcher:
+    @staticmethod
+    @torch.no_grad()
+    @contextmanager
+    def apply_smart_model_patches(
+        model: torch.nn.Module,
+        patches: Iterable[Tuple[ModelPatchRaw, float]],
+        prefix: str,
+        dtype: torch.dtype,
+        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
+        force_direct_patching: bool = False,
+        force_sidecar_patching: bool = False,
+    ):
+        """Apply 'smart' model 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:
+                LayerPatcher.apply_smart_model_patch(
+                    model=model,
+                    prefix=prefix,
+                    patch=patch,
+                    patch_weight=patch_weight,
+                    original_weights=original_weights,
+                    original_modules=original_modules,
+                    dtype=dtype,
+                    force_direct_patching=force_direct_patching,
+                    force_sidecar_patching=force_sidecar_patching,
+                )
+
+            yield
+        finally:
+            # Restore directly patched layers.
+            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)
+
+            # 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 = LayerPatcher._split_parent_key(module_key)
+                parent_module = model.get_submodule(module_parent_key)
+                LayerPatcher._set_submodule(parent_module, module_name, orig_module)
+
+    @staticmethod
+    @torch.no_grad()
+    def apply_smart_model_patch(
+        model: torch.nn.Module,
+        prefix: str,
+        patch: ModelPatchRaw,
+        patch_weight: float,
+        original_weights: OriginalWeightsStorage,
+        original_modules: dict[str, torch.nn.Module],
+        dtype: torch.dtype,
+        force_direct_patching: bool,
+        force_sidecar_patching: bool,
+    ):
+        """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 = LayerPatcher._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 quantized, so the caller passed force_sidecar_patching=True.
+            # - The module is already wrapped in a BaseSidecarWrapper.
+            # - 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 set force_sidecar_patching=True 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?
+            use_sidecar_patching = False
+            if force_direct_patching and force_sidecar_patching:
+                raise ValueError("Cannot force both direct and sidecar patching.")
+            elif force_direct_patching:
+                use_sidecar_patching = False
+            elif force_sidecar_patching:
+                use_sidecar_patching = True
+            elif isinstance(module, BaseSidecarWrapper):
+                use_sidecar_patching = True
+            elif LayerPatcher._is_any_part_of_layer_on_cpu(module):
+                use_sidecar_patching = True
+
+            if use_sidecar_patching:
+                LayerPatcher._apply_model_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:
+                LayerPatcher._apply_model_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()
+    def _apply_model_layer_patch(
+        module_to_patch: torch.nn.Module,
+        module_to_patch_key: str,
+        patch: BaseLayerPatch,
+        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
+
+        # 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, param_weight in patch.get_parameters(module_to_patch, weight=patch_weight).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)
+
+            # HACK(ryand): This condition is only necessary to handle layers in FLUX control LoRAs that change the
+            # shape of the original layer.
+            if module_param.nelement() != param_weight.nelement():
+                assert isinstance(patch, FluxControlLoRALayer)
+                expanded_weight = pad_with_zeros(module_param, param_weight.shape)
+                setattr(
+                    module_to_patch,
+                    param_name,
+                    torch.nn.Parameter(expanded_weight, requires_grad=module_param.requires_grad),
+                )
+                module_param = expanded_weight
+
+            module_param += param_weight.to(dtype=dtype)
+
+        patch.to(device=TorchDevice.CPU_DEVICE)
+
+    @staticmethod
+    @torch.no_grad()
+    def _apply_model_layer_wrapper_patch(
+        model: torch.nn.Module,
+        module_to_patch: torch.nn.Module,
+        module_to_patch_key: str,
+        patch: BaseLayerPatch,
+        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 BaseSidecarWrapper if it has not already been done.
+        if not isinstance(module_to_patch, BaseSidecarWrapper):
+            wrapped_module = wrap_module_with_sidecar_wrapper(orig_module=module_to_patch)
+            original_modules[module_to_patch_key] = module_to_patch
+            module_parent_key, module_name = LayerPatcher._split_parent_key(module_to_patch_key)
+            module_parent = model.get_submodule(module_parent_key)
+            LayerPatcher._set_submodule(module_parent, module_name, wrapped_module)
+        else:
+            assert module_to_patch_key in original_modules
+            wrapped_module = module_to_patch
+
+        # Move the LoRA layer to the same device/dtype as the orig module.
+        first_param = next(module_to_patch.parameters())
+        device = first_param.device
+        patch.to(device=device, dtype=dtype)
+
+        # Add the patch to the sidecar wrapper.
+        wrapped_module.add_patch(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 _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/patches/layers/base_layer_patch.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+
+import torch
+
+
+class BaseLayerPatch(ABC):
+    @abstractmethod
+    def get_parameters(self, orig_module: torch.nn.Module, weight: float) -> dict[str, torch.Tensor]:
+        """Get the parameter residual updates that should be applied to the original parameters. Parameters omitted
+        from the returned dict are not updated.
+        """
+        ...
+
+    @abstractmethod
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        """Move all internal tensors to the specified device and dtype."""
+        ...
+
+    @abstractmethod
+    def calc_size(self) -> int:
+        """Calculate the total size of all internal tensors in bytes."""
+        ...

invokeai/backend/patches/layers/concatenated_lora_layer.py

@@ -2,8 +2,8 @@ from typing import Optional, Sequence
 
 import torch
 
-from invokeai.backend.lora.layers.lora_layer import LoRALayer
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.patches.layers.lora_layer import LoRALayer
+from invokeai.backend.patches.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:

invokeai/backend/patches/layers/flux_control_lora_layer.py

@@ -0,0 +1,19 @@
+import torch
+
+from invokeai.backend.patches.layers.lora_layer import LoRALayer
+
+
+class FluxControlLoRALayer(LoRALayer):
+    """A special case of LoRALayer for use with FLUX Control LoRAs that pads the target parameter with zeros if the
+    shapes don't match.
+    """
+
+    def get_parameters(self, orig_module: torch.nn.Module, weight: float) -> dict[str, torch.Tensor]:
+        """This overrides the base class behavior to skip the reshaping step."""
+        scale = self.scale()
+        params = {"weight": self.get_weight(orig_module.weight) * (weight * scale)}
+        bias = self.get_bias(orig_module.bias)
+        if bias is not None:
+            params["bias"] = bias * (weight * scale)
+
+        return params

invokeai/backend/patches/layers/full_layer.py

@@ -2,7 +2,7 @@ from typing import Dict, Optional
 
 import torch
 
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.patches.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/patches/layers/ia3_layer.py

@@ -2,7 +2,7 @@ from typing import Dict, Optional
 
 import torch
 
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.patches.layers.lora_layer_base import LoRALayerBase
 
 
 class IA3Layer(LoRALayerBase):
@@ -16,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

invokeai/backend/patches/layers/loha_layer.py

@@ -2,7 +2,7 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.patches.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/patches/layers/lokr_layer.py

@@ -2,7 +2,7 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.patches.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/patches/layers/lora_layer.py

@@ -2,7 +2,7 @@ from typing import Dict, Optional
 
 import torch
 
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.patches.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/patches/layers/lora_layer_base.py

@@ -1,12 +1,13 @@
-from typing import Dict, Optional, Set
+from typing import Optional
 
 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:
+class LoRALayerBase(BaseLayerPatch):
     """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
@@ -23,6 +24,7 @@ class LoRALayerBase:
     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
@@ -37,11 +39,14 @@ class LoRALayerBase:
     ) -> float | None:
         return alpha.item() if alpha is not None else None
 
-    def rank(self) -> int | 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.
+        """
         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
@@ -52,15 +57,23 @@ class LoRALayerBase:
     def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
         return self.bias
 
-    def get_parameters(self, orig_module: torch.nn.Module) -> Dict[str, torch.Tensor]:
-        params = {"weight": self.get_weight(orig_module.weight)}
+    def get_parameters(self, orig_module: torch.nn.Module, weight: float) -> dict[str, torch.Tensor]:
+        scale = self.scale()
+        params = {"weight": self.get_weight(orig_module.weight) * (weight * scale)}
         bias = self.get_bias(orig_module.bias)
         if bias is not None:
-            params["bias"] = bias
+            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_module.get_parameter(param_name)
+            if param_weight.shape != orig_param.shape:
+                params[param_name] = param_weight.reshape(orig_param.shape)
+
         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/patches/layers/norm_layer.py

@@ -2,7 +2,7 @@ from typing import Dict
 
 import torch
 
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.patches.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/patches/layers/set_parameter_layer.py

@@ -0,0 +1,27 @@
+import torch
+
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+
+
+class SetParameterLayer(BaseLayerPatch):
+    """A layer that sets a single parameter to a new target value.
+    (The diff between the target value and current value is calculated internally.)
+    """
+
+    def __init__(self, param_name: str, weight: torch.Tensor):
+        super().__init__()
+        self.weight = weight
+        self.param_name = param_name
+
+    def get_parameters(self, orig_module: torch.nn.Module, weight: float) -> dict[str, torch.Tensor]:
+        # Note: We intentionally ignore the weight parameter here. This matches the behavior in the official FLUX
+        # Control LoRA implementation.
+        diff = self.weight - orig_module.get_parameter(self.param_name)
+        return {self.param_name: diff}
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.weight = self.weight.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        return calc_tensor_size(self.weight)

invokeai/backend/patches/layers/utils.py

@@ -2,16 +2,16 @@ from typing import Dict
 
 import torch
 
-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.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
 
 
-def any_lora_layer_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> AnyLoRALayer:
+def any_lora_layer_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> BaseLayerPatch:
     # Detect layers according to LyCORIS detection logic(`weight_list_det`)
     # https://github.com/KohakuBlueleaf/LyCORIS/tree/8ad8000efb79e2b879054da8c9356e6143591bad/lycoris/modules
 

invokeai/backend/patches/lora_conversions/flux_control_lora_utils.py

@@ -0,0 +1,84 @@
+import re
+from typing import Any, Dict
+
+import torch
+
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.flux_control_lora_layer import FluxControlLoRALayer
+from invokeai.backend.patches.layers.lora_layer import LoRALayer
+from invokeai.backend.patches.layers.set_parameter_layer import SetParameterLayer
+from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+
+# 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_CONTROL_TRANSFORMER_KEY_REGEX = r"(\w+\.)+(lora_A\.weight|lora_B\.weight|lora_B\.bias|scale)"
+
+
+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.)
+    """
+
+    all_keys_match = all(re.match(FLUX_CONTROL_TRANSFORMER_KEY_REGEX, str(k)) for k in state_dict.keys())
+
+    # Check the shape of the img_in weight, because this layer shape is modified by FLUX control LoRAs.
+    lora_a_weight = state_dict.get("img_in.lora_A.weight", None)
+    lora_b_bias = state_dict.get("img_in.lora_B.bias", None)
+    lora_b_weight = state_dict.get("img_in.lora_B.weight", None)
+
+    return (
+        all_keys_match
+        and lora_a_weight is not None
+        and lora_b_bias is not None
+        and lora_b_weight is not None
+        and lora_a_weight.shape[1] == 128
+        and lora_b_weight.shape[0] == 3072
+        and lora_b_bias.shape[0] == 3072
+    )
+
+
+def lora_model_from_flux_control_state_dict(state_dict: Dict[str, torch.Tensor]) -> ModelPatchRaw:
+    # Group keys by layer.
+    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {}
+    for key, value in state_dict.items():
+        key_props = key.split(".")
+        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, BaseLayerPatch] = {}
+    for layer_key, layer_state_dict in grouped_state_dict.items():
+        prefixed_key = f"{FLUX_LORA_TRANSFORMER_PREFIX}{layer_key}"
+        if layer_key == "img_in":
+            # img_in is a special case because it changes the shape of the original weight.
+            layers[prefixed_key] = FluxControlLoRALayer(
+                layer_state_dict["lora_B.weight"],
+                None,
+                layer_state_dict["lora_A.weight"],
+                None,
+                layer_state_dict["lora_B.bias"],
+            )
+        elif 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 ValueError(f"{layer_key} not expected")
+
+    return ModelPatchRaw(layers=layers)