chore: bump version to v6.3.0

This is useful for FLUX Kontext, where you typically want the generation
size to at least roughly match the first ref image size.

.github/CODEOWNERS

@@ -1,32 +1,31 @@
 # continuous integration
-/.github/workflows/  @lstein @blessedcoolant @hipsterusername @ebr @jazzhaiku
+/.github/workflows/  @lstein @blessedcoolant @hipsterusername @ebr @jazzhaiku @psychedelicious
 
 # documentation
 /docs/ @lstein @blessedcoolant @hipsterusername @psychedelicious
 /mkdocs.yml @lstein @blessedcoolant @hipsterusername @psychedelicious
 
 # nodes
-/invokeai/app/ @blessedcoolant @psychedelicious @brandonrising @hipsterusername @jazzhaiku
+/invokeai/app/ @blessedcoolant @psychedelicious @hipsterusername @jazzhaiku
 
 # installation and configuration
-/pyproject.toml  @lstein @blessedcoolant @hipsterusername
-/docker/  @lstein @blessedcoolant @hipsterusername @ebr
-/scripts/ @ebr @lstein @hipsterusername
-/installer/ @lstein @ebr @hipsterusername
-/invokeai/assets @lstein @ebr @hipsterusername
-/invokeai/configs @lstein @hipsterusername
-/invokeai/version @lstein @blessedcoolant @hipsterusername
+/pyproject.toml  @lstein @blessedcoolant @psychedelicious @hipsterusername
+/docker/  @lstein @blessedcoolant @psychedelicious @hipsterusername @ebr
+/scripts/ @ebr @lstein @psychedelicious @hipsterusername
+/installer/ @lstein @ebr @psychedelicious @hipsterusername
+/invokeai/assets @lstein @ebr @psychedelicious @hipsterusername
+/invokeai/configs @lstein @psychedelicious @hipsterusername
+/invokeai/version @lstein @blessedcoolant @psychedelicious @hipsterusername
 
 # web ui
 /invokeai/frontend @blessedcoolant @psychedelicious @lstein @maryhipp @hipsterusername
-/invokeai/backend @blessedcoolant @psychedelicious @lstein @maryhipp @hipsterusername
 
 # generation, model management, postprocessing
-/invokeai/backend  @lstein @blessedcoolant @brandonrising @hipsterusername @jazzhaiku
+/invokeai/backend  @lstein @blessedcoolant @hipsterusername @jazzhaiku @psychedelicious @maryhipp 
 
 # front ends
-/invokeai/frontend/CLI @lstein @hipsterusername
-/invokeai/frontend/install @lstein @ebr @hipsterusername
-/invokeai/frontend/merge @lstein @blessedcoolant @hipsterusername
-/invokeai/frontend/training @lstein @blessedcoolant @hipsterusername
+/invokeai/frontend/CLI @lstein @psychedelicious @hipsterusername
+/invokeai/frontend/install @lstein @ebr @psychedelicious @hipsterusername
+/invokeai/frontend/merge @lstein @blessedcoolant @psychedelicious @hipsterusername
+/invokeai/frontend/training @lstein @blessedcoolant @psychedelicious @hipsterusername
 /invokeai/frontend/web @psychedelicious @blessedcoolant @maryhipp @hipsterusername

.github/ISSUE_TEMPLATE/BUG_REPORT.yml

@@ -21,6 +21,20 @@ body:
         - label: I have searched the existing issues
           required: true
 
+  - type: dropdown
+    id: install_method
+    attributes:
+      label: Install method
+      description: How did you install Invoke?
+      multiple: false
+      options:
+        - "Invoke's Launcher"
+        - 'Stability Matrix'
+        - 'Pinokio'
+        - 'Manual'
+    validations:
+      required: true
+
   - type: markdown
     attributes:
       value: __Describe your environment__
@@ -76,8 +90,8 @@ body:
     attributes:
       label: Version number
       description: |
-        The version of Invoke you have installed. If it is not the latest version, please update and try again to confirm the issue still exists. If you are testing main, please include the commit hash instead.
-      placeholder: ex. 3.6.1
+        The version of Invoke you have installed. If it is not the [latest version](https://github.com/invoke-ai/InvokeAI/releases/latest), please update and try again to confirm the issue still exists. If you are testing main, please include the commit hash instead.
+      placeholder: ex. v6.0.2
     validations:
       required: true
 
@@ -85,17 +99,17 @@ body:
     id: browser-version
     attributes:
       label: Browser
-      description: Your web browser and version.
+      description: Your web browser and version, if you do not use the Launcher's provided GUI.
       placeholder: ex. Firefox 123.0b3
     validations:
-      required: true
+      required: false
 
   - type: textarea
     id: python-deps
     attributes:
-      label: Python dependencies
+      label: System Information
       description: |
-        If the problem occurred during image generation, click the gear icon at the bottom left corner, click "About", click the copy button and then paste here.
+        Click the gear icon at the bottom left corner, then click "About". Click the copy button and then paste here.
     validations:
       required: false
 

.github/actions/install-frontend-deps/action.yml

@@ -3,15 +3,15 @@ description: Installs frontend dependencies with pnpm, with caching
 runs:
   using: 'composite'
   steps:
-    - name: setup node 18
+    - name: setup node 20
       uses: actions/setup-node@v4
       with:
-        node-version: '18'
+        node-version: '20'
 
     - name: setup pnpm
       uses: pnpm/action-setup@v4
       with:
-        version: 8.15.6
+        version: 10
         run_install: false
 
     - name: get pnpm store directory

.github/workflows/python-checks.yml

@@ -67,6 +67,10 @@ jobs:
           version: '0.6.10'
           enable-cache: true
 
+      - name: check pypi classifiers
+        if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
+        run: uv run --no-project scripts/check_classifiers.py ./pyproject.toml
+
       - name: ruff check
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
         run: uv tool run ruff@0.11.2 check --output-format=github .

.github/workflows/typegen-checks.yml

@@ -39,6 +39,18 @@ jobs:
       - name: checkout
         uses: actions/checkout@v4
 
+      - name: Free up more disk space on the runner
+        # https://github.com/actions/runner-images/issues/2840#issuecomment-1284059930
+        run: |
+          echo "----- Free space before cleanup"
+          df -h
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+          sudo swapoff /mnt/swapfile
+          sudo rm -rf /mnt/swapfile
+          echo "----- Free space after cleanup"
+          df -h
+
       - name: check for changed files
         if: ${{ inputs.always_run != true }}
         id: changed-files

.gitignore

@@ -180,6 +180,7 @@ cython_debug/
 # Scratch folder
 .scratch/
 .vscode/
+.zed/
 
 # source installer files
 installer/*zip
@@ -188,3 +189,6 @@ installer/install.sh
 installer/update.bat
 installer/update.sh
 installer/InvokeAI-Installer/
+.aider*
+
+.claude/

docker/.env.sample

@@ -22,6 +22,10 @@
 ## GPU_DRIVER can be set to either `cuda` or `rocm` to enable GPU support in the container accordingly.
 # GPU_DRIVER=cuda #| rocm
 
+## If you are using ROCM, you will need to ensure that the render group within the container and the host system use the same group ID.
+## To obtain the group ID of the render group on the host system, run `getent group render` and grab the number.
+# RENDER_GROUP_ID=
+
 ## CONTAINER_UID can be set to the UID of the user on the host system that should own the files in the container.
 ## It is usually not necessary to change this. Use `id -u` on the host system to find the UID.
 # CONTAINER_UID=1000

docker/Dockerfile

@@ -5,8 +5,7 @@
 FROM docker.io/node:22-slim AS web-builder
 ENV PNPM_HOME="/pnpm"
 ENV PATH="$PNPM_HOME:$PATH"
-RUN corepack use pnpm@8.x
-RUN corepack enable
+RUN corepack use pnpm@10.x && corepack enable
 
 WORKDIR /build
 COPY invokeai/frontend/web/ ./
@@ -44,7 +43,6 @@ ENV \
     UV_MANAGED_PYTHON=1 \
     UV_LINK_MODE=copy \
     UV_PROJECT_ENVIRONMENT=/opt/venv \
-    UV_INDEX="https://download.pytorch.org/whl/cu124" \
     INVOKEAI_ROOT=/invokeai \
     INVOKEAI_HOST=0.0.0.0 \
     INVOKEAI_PORT=9090 \
@@ -75,19 +73,17 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=uv.lock,target=uv.lock \
     # this is just to get the package manager to recognize that the project exists, without making changes to the docker layer
     --mount=type=bind,source=invokeai/version,target=invokeai/version \
-    if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then UV_INDEX="https://download.pytorch.org/whl/cpu"; \
-    elif [ "$GPU_DRIVER" = "rocm" ]; then UV_INDEX="https://download.pytorch.org/whl/rocm6.2"; \
-    fi && \
-    uv sync --frozen
-
-# build patchmatch
-RUN cd /usr/lib/$(uname -p)-linux-gnu/pkgconfig/ && ln -sf opencv4.pc opencv.pc
-RUN python -c "from patchmatch import patch_match"
+    ulimit -n 30000 && \
+    uv sync --extra $GPU_DRIVER --frozen
 
 # Link amdgpu.ids for ROCm builds
 # contributed by https://github.com/Rubonnek
 RUN mkdir -p "/opt/amdgpu/share/libdrm" &&\
-    ln -s "/usr/share/libdrm/amdgpu.ids" "/opt/amdgpu/share/libdrm/amdgpu.ids"
+    ln -s "/usr/share/libdrm/amdgpu.ids" "/opt/amdgpu/share/libdrm/amdgpu.ids" && groupadd render
+
+# build patchmatch
+RUN cd /usr/lib/$(uname -p)-linux-gnu/pkgconfig/ && ln -sf opencv4.pc opencv.pc
+RUN python -c "from patchmatch import patch_match"
 
 RUN mkdir -p ${INVOKEAI_ROOT} && chown -R ${CONTAINER_UID}:${CONTAINER_GID} ${INVOKEAI_ROOT}
 
@@ -106,8 +102,6 @@ COPY invokeai ${INVOKEAI_SRC}/invokeai
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
     --mount=type=bind,source=uv.lock,target=uv.lock \
-    if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then UV_INDEX="https://download.pytorch.org/whl/cpu"; \
-    elif [ "$GPU_DRIVER" = "rocm" ]; then UV_INDEX="https://download.pytorch.org/whl/rocm6.2"; \
-    fi && \
-    uv pip install -e .
+    ulimit -n 30000 && \
+    uv pip install -e .[$GPU_DRIVER]
 

docker/Dockerfile-rocm-full

@@ -0,0 +1,136 @@
+# syntax=docker/dockerfile:1.4
+
+#### Web UI ------------------------------------
+
+FROM docker.io/node:22-slim AS web-builder
+ENV PNPM_HOME="/pnpm"
+ENV PATH="$PNPM_HOME:$PATH"
+RUN corepack use pnpm@8.x
+RUN corepack enable
+
+WORKDIR /build
+COPY invokeai/frontend/web/ ./
+RUN --mount=type=cache,target=/pnpm/store \
+    pnpm install --frozen-lockfile
+RUN npx vite build
+
+## Backend ---------------------------------------
+
+FROM library/ubuntu:24.04
+
+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 \
+    --mount=type=cache,target=/var/lib/apt \
+    apt update && apt install -y --no-install-recommends \
+    ca-certificates \
+    git \
+    gosu \
+    libglib2.0-0 \
+    libgl1 \
+    libglx-mesa0 \
+    build-essential \
+    libopencv-dev \
+    libstdc++-10-dev \
+    wget
+
+ENV \
+    PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    VIRTUAL_ENV=/opt/venv \
+    INVOKEAI_SRC=/opt/invokeai \
+    PYTHON_VERSION=3.12 \
+    UV_PYTHON=3.12 \
+    UV_COMPILE_BYTECODE=1 \
+    UV_MANAGED_PYTHON=1 \
+    UV_LINK_MODE=copy \
+    UV_PROJECT_ENVIRONMENT=/opt/venv \
+    INVOKEAI_ROOT=/invokeai \
+    INVOKEAI_HOST=0.0.0.0 \
+    INVOKEAI_PORT=9090 \
+    PATH="/opt/venv/bin:$PATH" \
+    CONTAINER_UID=${CONTAINER_UID:-1000} \
+    CONTAINER_GID=${CONTAINER_GID:-1000}
+
+ARG GPU_DRIVER=cuda
+
+# Install `uv` for package management
+COPY --from=ghcr.io/astral-sh/uv:0.6.9 /uv /uvx /bin/
+
+# Install python & allow non-root user to use it by traversing the /root dir without read permissions
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv python install ${PYTHON_VERSION} && \
+    # chmod --recursive a+rX /root/.local/share/uv/python
+    chmod 711 /root
+
+WORKDIR ${INVOKEAI_SRC}
+
+# Install project's dependencies as a separate layer so they aren't rebuilt every commit.
+# bind-mount instead of copy to defer adding sources to the image until next layer.
+#
+# NOTE: there are no pytorch builds for arm64 + cuda, only cpu
+# x86_64/CUDA is the default
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    # this is just to get the package manager to recognize that the project exists, without making changes to the docker layer
+    --mount=type=bind,source=invokeai/version,target=invokeai/version \
+    ulimit -n 30000 && \
+    uv sync --extra $GPU_DRIVER --frozen
+
+RUN --mount=type=cache,target=/var/cache/apt \
+    --mount=type=cache,target=/var/lib/apt \
+    if [ "$GPU_DRIVER" = "rocm" ]; then \
+    wget -O /tmp/amdgpu-install.deb \
+    https://repo.radeon.com/amdgpu-install/6.3.4/ubuntu/noble/amdgpu-install_6.3.60304-1_all.deb && \
+    apt install -y /tmp/amdgpu-install.deb && \
+    apt update && \
+    amdgpu-install --usecase=rocm -y && \
+    apt-get autoclean && \
+    apt clean && \
+    rm -rf /tmp/* /var/tmp/* && \
+    usermod -a -G render ubuntu && \
+    usermod -a -G video ubuntu && \
+    echo "\\n/opt/rocm/lib\\n/opt/rocm/lib64" >> /etc/ld.so.conf.d/rocm.conf && \
+    ldconfig && \
+    update-alternatives --auto rocm; \
+    fi
+
+## Heathen711: Leaving this for review input, will remove before merge
+# RUN --mount=type=cache,target=/var/cache/apt \
+#     --mount=type=cache,target=/var/lib/apt \
+#     if [ "$GPU_DRIVER" = "rocm" ]; then \
+#     groupadd render && \
+#     usermod -a -G render ubuntu && \
+#     usermod -a -G video ubuntu; \
+#     fi
+
+## Link amdgpu.ids for ROCm builds
+## contributed by https://github.com/Rubonnek
+# RUN mkdir -p "/opt/amdgpu/share/libdrm" &&\
+#     ln -s "/usr/share/libdrm/amdgpu.ids" "/opt/amdgpu/share/libdrm/amdgpu.ids"
+
+# build patchmatch
+RUN cd /usr/lib/$(uname -p)-linux-gnu/pkgconfig/ && ln -sf opencv4.pc opencv.pc
+RUN python -c "from patchmatch import patch_match"
+
+RUN mkdir -p ${INVOKEAI_ROOT} && chown -R ${CONTAINER_UID}:${CONTAINER_GID} ${INVOKEAI_ROOT}
+
+COPY docker/docker-entrypoint.sh ./
+ENTRYPOINT ["/opt/invokeai/docker-entrypoint.sh"]
+CMD ["invokeai-web"]
+
+# --link requires buldkit w/ dockerfile syntax 1.4, does not work with podman
+COPY --link --from=web-builder /build/dist ${INVOKEAI_SRC}/invokeai/frontend/web/dist
+
+# add sources last to minimize image changes on code changes
+COPY invokeai ${INVOKEAI_SRC}/invokeai
+
+# this should not increase image size because we've already installed dependencies
+# in a previous layer
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    ulimit -n 30000 && \
+    uv pip install -e .[$GPU_DRIVER]
+

docker/docker-compose.yml

@@ -47,8 +47,9 @@ services:
 
   invokeai-rocm:
     <<: *invokeai
-    devices:
-      - /dev/kfd:/dev/kfd
-      - /dev/dri:/dev/dri
+    environment:
+      - AMD_VISIBLE_DEVICES=all
+      - RENDER_GROUP_ID=${RENDER_GROUP_ID}
+    runtime: amd
     profiles:
       - rocm

docker/docker-entrypoint.sh

@@ -21,6 +21,17 @@ _=$(id ${USER} 2>&1) || useradd -u ${USER_ID} ${USER}
 # ensure the UID is correct
 usermod -u ${USER_ID} ${USER} 1>/dev/null
 
+## ROCM specific configuration
+# render group within the container must match the host render group
+# otherwise the container will not be able to access the host GPU.
+if [[ -v "RENDER_GROUP_ID" ]] && [[ ! -z "${RENDER_GROUP_ID}" ]]; then
+  # ensure the render group exists
+  groupmod -g ${RENDER_GROUP_ID} render
+  usermod -a -G render ${USER}
+  usermod -a -G video ${USER}
+fi
+
+
 ### Set the $PUBLIC_KEY env var to enable SSH access.
 # We do not install openssh-server in the image by default to avoid bloat.
 # but it is useful to have the full SSH server e.g. on Runpod.

docker/run.sh

@@ -13,7 +13,7 @@ run() {
 
   # parse .env file for build args
   build_args=$(awk '$1 ~ /=[^$]/ && $0 !~ /^#/ {print "--build-arg " $0 " "}' .env) &&
-  profile="$(awk -F '=' '/GPU_DRIVER/ {print $2}' .env)"
+  profile="$(awk -F '=' '/GPU_DRIVER=/ {print $2}' .env)"
 
   # default to 'cuda' profile
   [[ -z "$profile" ]] && profile="cuda"
@@ -30,7 +30,7 @@ run() {
 
   printf "%s\n" "starting service $service_name"
   docker compose --profile "$profile" up -d "$service_name"
-  docker compose logs -f
+  docker compose --profile "$profile" logs -f
 }
 
 run

docs/contributing/INVOCATIONS.md

@@ -39,7 +39,7 @@ nodes imported in the `__init__.py` file are loaded. See the README in the nodes
 folder for more examples:
 
 ```py
-from .cool_node import CoolInvocation
+from .cool_node import ResizeInvocation
 ```
 
 ## Creating A New Invocation
@@ -69,7 +69,10 @@ The first set of things we need to do when creating a new Invocation are -
 So let us do that.
 
 ```python
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
+from invokeai.invocation_api import (
+    BaseInvocation,
+    invocation,
+)
 
 @invocation('resize')
 class ResizeInvocation(BaseInvocation):
@@ -103,8 +106,12 @@ create your own custom field types later in this guide. For now, let's go ahead
 and use it.
 
 ```python
-from invokeai.app.invocations.baseinvocation import BaseInvocation, InputField, invocation
-from invokeai.app.invocations.primitives import ImageField
+from invokeai.invocation_api import (
+    BaseInvocation,
+    ImageField,
+    InputField,
+    invocation,
+)
 
 @invocation('resize')
 class ResizeInvocation(BaseInvocation):
@@ -128,8 +135,12 @@ image: ImageField = InputField(description="The input image")
 Great. Now let us create our other inputs for `width` and `height`
 
 ```python
-from invokeai.app.invocations.baseinvocation import BaseInvocation, InputField, invocation
-from invokeai.app.invocations.primitives import ImageField
+from invokeai.invocation_api import (
+    BaseInvocation,
+    ImageField,
+    InputField,
+    invocation,
+)
 
 @invocation('resize')
 class ResizeInvocation(BaseInvocation):
@@ -163,8 +174,13 @@ that are provided by it by InvokeAI.
 Let us create this function first.
 
 ```python
-from invokeai.app.invocations.baseinvocation import BaseInvocation, InputField, invocation, InvocationContext
-from invokeai.app.invocations.primitives import ImageField
+from invokeai.invocation_api import (
+    BaseInvocation,
+    ImageField,
+    InputField,
+    InvocationContext,
+    invocation,
+)
 
 @invocation('resize')
 class ResizeInvocation(BaseInvocation):
@@ -191,8 +207,14 @@ all the necessary info related to image outputs. So let us use that.
 We will cover how to create your own output types later in this guide.
 
 ```python
-from invokeai.app.invocations.baseinvocation import BaseInvocation, InputField, invocation, InvocationContext
-from invokeai.app.invocations.primitives import ImageField
+from invokeai.invocation_api import (
+    BaseInvocation,
+    ImageField,
+    InputField,
+    InvocationContext,
+    invocation,
+)
+
 from invokeai.app.invocations.image import ImageOutput
 
 @invocation('resize')
@@ -217,9 +239,15 @@ Perfect. Now that we have our Invocation setup, let us do what we want to do.
 So let's do that.
 
 ```python
-from invokeai.app.invocations.baseinvocation import BaseInvocation, InputField, invocation, InvocationContext
-from invokeai.app.invocations.primitives import ImageField
-from invokeai.app.invocations.image import ImageOutput, ResourceOrigin, ImageCategory
+from invokeai.invocation_api import (
+    BaseInvocation,
+    ImageField,
+    InputField,
+    InvocationContext,
+    invocation,
+)
+
+from invokeai.app.invocations.image import ImageOutput
 
 @invocation("resize")
 class ResizeInvocation(BaseInvocation):

docs/contributing/dev-environment.md

@@ -41,7 +41,7 @@ If you just want to use Invoke, you should use the [launcher][launcher link].
      With the modifications made, the install command should look something like this:
 
       ```sh
-      uv pip install -e ".[dev,test,docs,xformers]" --python 3.12 --python-preference only-managed --index=https://download.pytorch.org/whl/cu126 --reinstall
+      uv pip install -e ".[dev,test,docs,xformers]" --python 3.12 --python-preference only-managed --index=https://download.pytorch.org/whl/cu128 --reinstall
       ```
 
 6. At this point, you should have Invoke installed, a venv set up and activated, and the server running. But you will see a warning in the terminal that no UI was found. If you go to the URL for the server, you won't get a UI.
@@ -50,11 +50,11 @@ If you just want to use Invoke, you should use the [launcher][launcher link].
 
       If you only want to edit the docs, you can stop here and skip to the **Documentation** section below.
 
-7. Install the frontend dev toolchain:
+7. Install the frontend dev toolchain, paying attention to versions:
 
-      - [`nodejs`](https://nodejs.org/) (v20+)
+      - [`nodejs`](https://nodejs.org/) (tested on LTS, v22)
 
-      - [`pnpm`](https://pnpm.io/8.x/installation) (must be v8 - not v9!)
+      - [`pnpm`](https://pnpm.io/installation) (tested on v10)
 
 8. Do a production build of the frontend:
 

docs/installation/manual.md

@@ -69,27 +69,34 @@ The following commands vary depending on the version of Invoke being installed a
     - If you have an Nvidia 20xx series GPU or older, use `invokeai[xformers]`.
     - If you have an Nvidia 30xx series GPU or newer, or do not have an Nvidia GPU, use `invokeai`.
 
-7. Determine the `PyPI` index URL to use for installation, if any. This is necessary to get the right version of torch installed.
+7. Determine the torch backend to use for installation, if any. This is necessary to get the right version of torch installed. This is acheived by using [UV's built in torch support.](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection)
 
-    === "Invoke v5.10.0 and later"
+    === "Invoke v5.12 and later"
 
-        - If you are on Windows or Linux with an Nvidia GPU, use `https://download.pytorch.org/whl/cu126`.
-        - If you are on Linux with no GPU, use `https://download.pytorch.org/whl/cpu`.
-        - If you are on Linux with an AMD GPU, use `https://download.pytorch.org/whl/rocm6.2.4`.
+        - If you are on Windows or Linux with an Nvidia GPU, use `--torch-backend=cu128`.
+        - If you are on Linux with no GPU, use `--torch-backend=cpu`.
+        - If you are on Linux with an AMD GPU, use `--torch-backend=rocm6.3`.
+        - **In all other cases, do not use a torch backend.**
+
+    === "Invoke v5.10.0 to v5.11.0"
+
+        - If you are on Windows or Linux with an Nvidia GPU, use `--torch-backend=cu126`.
+        - If you are on Linux with no GPU, use `--torch-backend=cpu`.
+        - If you are on Linux with an AMD GPU, use `--torch-backend=rocm6.2.4`.
         - **In all other cases, do not use an index.**
 
     === "Invoke v5.0.0 to v5.9.1"
 
-        - If you are on Windows with an Nvidia GPU, use `https://download.pytorch.org/whl/cu124`.
-        - If you are on Linux with no GPU, use `https://download.pytorch.org/whl/cpu`.
-        - If you are on Linux with an AMD GPU, use `https://download.pytorch.org/whl/rocm6.1`.
+        - If you are on Windows with an Nvidia GPU, use `--torch-backend=cu124`.
+        - If you are on Linux with no GPU, use `--torch-backend=cpu`.
+        - If you are on Linux with an AMD GPU, use `--torch-backend=rocm6.1`.
         - **In all other cases, do not use an index.**
 
     === "Invoke v4"
 
-        - If you are on Windows with an Nvidia GPU, use `https://download.pytorch.org/whl/cu124`.
-        - If you are on Linux with no GPU, use `https://download.pytorch.org/whl/cpu`.
-        - If you are on Linux with an AMD GPU, use `https://download.pytorch.org/whl/rocm5.2`.
+        - If you are on Windows with an Nvidia GPU, use `--torch-backend=cu124`.
+        - If you are on Linux with no GPU, use `--torch-backend=cpu`.
+        - If you are on Linux with an AMD GPU, use `--torch-backend=rocm5.2`.
         - **In all other cases, do not use an index.**
 
 8. Install the `invokeai` package. Substitute the package specifier and version.
@@ -98,10 +105,10 @@ The following commands vary depending on the version of Invoke being installed a
     uv pip install <PACKAGE_SPECIFIER>==<VERSION> --python 3.12 --python-preference only-managed --force-reinstall
     ```
 
-    If you determined you needed to use a `PyPI` index URL in the previous step, you'll need to add `--index=<INDEX_URL>` like this:
+    If you determined you needed to use a torch backend in the previous step, you'll need to set the backend like this:
 
     ```sh
-    uv pip install <PACKAGE_SPECIFIER>==<VERSION> --python 3.12 --python-preference only-managed --index=<INDEX_URL> --force-reinstall
+    uv pip install <PACKAGE_SPECIFIER>==<VERSION> --python 3.12 --python-preference only-managed --torch-backend=<VERSION> --force-reinstall
     ```
 
 9. Deactivate and reactivate your venv so that the invokeai-specific commands become available in the environment:

docs/installation/quick_start.md

@@ -35,7 +35,7 @@ More detail on system requirements can be found [here](./requirements.md).
 
 ## Step 2: Download
 
-Download the most launcher for your operating system:
+Download the most recent launcher for your operating system:
 
 - [Download for Windows](https://download.invoke.ai/Invoke%20Community%20Edition.exe)
 - [Download for macOS](https://download.invoke.ai/Invoke%20Community%20Edition.dmg)

docs/nodes/communityNodes.md

@@ -13,6 +13,7 @@ If you'd prefer, you can also just download the whole node folder from the linke
 To use a community workflow, download the `.json` node graph file and load it into Invoke AI via the **Load Workflow** button in the Workflow Editor. 
 
 - Community Nodes
+    + [Anamorphic Tools](#anamorphic-tools)
     + [Adapters-Linked](#adapters-linked-nodes)
     + [Autostereogram](#autostereogram-nodes)
     + [Average Images](#average-images)
@@ -20,9 +21,12 @@ To use a community workflow, download the `.json` node graph file and load it in
     + [Close Color Mask](#close-color-mask) 
     + [Clothing Mask](#clothing-mask)
     + [Contrast Limited Adaptive Histogram Equalization](#contrast-limited-adaptive-histogram-equalization)
+    + [Curves](#curves)
     + [Depth Map from Wavefront OBJ](#depth-map-from-wavefront-obj)
     + [Enhance Detail](#enhance-detail)
     + [Film Grain](#film-grain)
+    + [Flip Pose](#flip-pose)
+    + [Flux Ideal Size](#flux-ideal-size)
     + [Generative Grammar-Based Prompt Nodes](#generative-grammar-based-prompt-nodes)
     + [GPT2RandomPromptMaker](#gpt2randompromptmaker)
     + [Grid to Gif](#grid-to-gif)
@@ -61,6 +65,13 @@ To use a community workflow, download the `.json` node graph file and load it in
 - [Help](#help)
 
 
+--------------------------------
+### Anamorphic Tools
+
+**Description:** A set of nodes to perform anamorphic modifications to images, like lens blur, streaks, spherical distortion, and vignetting.
+
+**Node Link:** https://github.com/JPPhoto/anamorphic-tools
+
 --------------------------------
 ### Adapters Linked Nodes
 
@@ -132,6 +143,13 @@ Node Link: https://github.com/VeyDlin/clahe-node
 View:
 </br><img src="https://raw.githubusercontent.com/VeyDlin/clahe-node/master/.readme/node.png" width="500" />
 
+--------------------------------
+### Curves
+
+**Description:** Adjust an image's curve based on a user-defined string.
+
+**Node Link:** https://github.com/JPPhoto/curves-node
+
 --------------------------------
 ### Depth Map from Wavefront OBJ
 
@@ -162,6 +180,20 @@ To be imported, an .obj must use triangulated meshes, so make sure to enable tha
 
 **Node Link:** https://github.com/JPPhoto/film-grain-node
 
+--------------------------------
+### Flip Pose
+
+**Description:** This node will flip an openpose image horizontally, recoloring it to make sure that it isn't facing the wrong direction. Note that it does not work with openpose hands.
+
+**Node Link:** https://github.com/JPPhoto/flip-pose-node
+
+--------------------------------
+### Flux Ideal Size
+
+**Description:** This node returns an ideal size to use for the first stage of a Flux image generation pipeline. Generating at the right size helps limit duplication and odd subject placement.
+
+**Node Link:** https://github.com/JPPhoto/flux-ideal-size
+
 --------------------------------
 ### Generative Grammar-Based Prompt Nodes
 

invokeai/app/api/dependencies.py

@@ -10,6 +10,7 @@ from invokeai.app.services.board_images.board_images_default import BoardImagesS
 from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage
 from invokeai.app.services.boards.boards_default import BoardService
 from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService
+from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ClientStatePersistenceSqlite
 from invokeai.app.services.config.config_default import InvokeAIAppConfig
 from invokeai.app.services.download.download_default import DownloadQueueService
 from invokeai.app.services.events.events_fastapievents import FastAPIEventService
@@ -23,6 +24,10 @@ from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_images.model_images_default import ModelImageFileStorageDisk
 from invokeai.app.services.model_manager.model_manager_default import ModelManagerService
 from invokeai.app.services.model_records.model_records_sql import ModelRecordServiceSQL
+from invokeai.app.services.model_relationship_records.model_relationship_records_sqlite import (
+    SqliteModelRelationshipRecordStorage,
+)
+from invokeai.app.services.model_relationships.model_relationships_default import ModelRelationshipsService
 from invokeai.app.services.names.names_default import SimpleNameService
 from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
 from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
@@ -136,6 +141,8 @@ class ApiDependencies:
             download_queue=download_queue_service,
             events=events,
         )
+        model_relationships = ModelRelationshipsService()
+        model_relationship_records = SqliteModelRelationshipRecordStorage(db=db)
         names = SimpleNameService()
         performance_statistics = InvocationStatsService()
         session_processor = DefaultSessionProcessor(session_runner=DefaultSessionRunner())
@@ -145,6 +152,7 @@ class ApiDependencies:
         style_preset_records = SqliteStylePresetRecordsStorage(db=db)
         style_preset_image_files = StylePresetImageFileStorageDisk(style_presets_folder / "images")
         workflow_thumbnails = WorkflowThumbnailFileStorageDisk(workflow_thumbnails_folder)
+        client_state_persistence = ClientStatePersistenceSqlite(db=db)
 
         services = InvocationServices(
             board_image_records=board_image_records,
@@ -161,6 +169,8 @@ class ApiDependencies:
             logger=logger,
             model_images=model_images_service,
             model_manager=model_manager,
+            model_relationships=model_relationships,
+            model_relationship_records=model_relationship_records,
             download_queue=download_queue_service,
             names=names,
             performance_statistics=performance_statistics,
@@ -173,6 +183,7 @@ class ApiDependencies:
             style_preset_records=style_preset_records,
             style_preset_image_files=style_preset_image_files,
             workflow_thumbnails=workflow_thumbnails,
+            client_state_persistence=client_state_persistence,
         )
 
         ApiDependencies.invoker = Invoker(services)

invokeai/app/api/routers/app_info.py

@@ -1,8 +1,7 @@
 import typing
 from enum import Enum
-from importlib.metadata import PackageNotFoundError, version
+from importlib.metadata import distributions
 from pathlib import Path
-from platform import python_version
 from typing import Optional
 
 import torch
@@ -44,24 +43,6 @@ class AppVersion(BaseModel):
     highlights: Optional[list[str]] = Field(default=None, description="Highlights of release")
 
 
-class AppDependencyVersions(BaseModel):
-    """App depencency Versions Response"""
-
-    accelerate: str = Field(description="accelerate version")
-    compel: str = Field(description="compel version")
-    cuda: Optional[str] = Field(description="CUDA version")
-    diffusers: str = Field(description="diffusers version")
-    numpy: str = Field(description="Numpy version")
-    opencv: str = Field(description="OpenCV version")
-    onnx: str = Field(description="ONNX version")
-    pillow: str = Field(description="Pillow (PIL) version")
-    python: str = Field(description="Python version")
-    torch: str = Field(description="PyTorch version")
-    torchvision: str = Field(description="PyTorch Vision version")
-    transformers: str = Field(description="transformers version")
-    xformers: Optional[str] = Field(description="xformers version")
-
-
 class AppConfig(BaseModel):
     """App Config Response"""
 
@@ -76,27 +57,19 @@ async def get_version() -> AppVersion:
     return AppVersion(version=__version__)
 
 
-@app_router.get("/app_deps", operation_id="get_app_deps", status_code=200, response_model=AppDependencyVersions)
-async def get_app_deps() -> AppDependencyVersions:
+@app_router.get("/app_deps", operation_id="get_app_deps", status_code=200, response_model=dict[str, str])
+async def get_app_deps() -> dict[str, str]:
+    deps: dict[str, str] = {dist.metadata["Name"]: dist.version for dist in distributions()}
     try:
-        xformers = version("xformers")
-    except PackageNotFoundError:
-        xformers = None
-    return AppDependencyVersions(
-        accelerate=version("accelerate"),
-        compel=version("compel"),
-        cuda=torch.version.cuda,
-        diffusers=version("diffusers"),
-        numpy=version("numpy"),
-        opencv=version("opencv-python"),
-        onnx=version("onnx"),
-        pillow=version("pillow"),
-        python=python_version(),
-        torch=torch.version.__version__,
-        torchvision=version("torchvision"),
-        transformers=version("transformers"),
-        xformers=xformers,
-    )
+        cuda = torch.version.cuda or "N/A"
+    except Exception:
+        cuda = "N/A"
+
+    deps["CUDA"] = cuda
+
+    sorted_deps = dict(sorted(deps.items(), key=lambda item: item[0].lower()))
+
+    return sorted_deps
 
 
 @app_router.get("/config", operation_id="get_config", status_code=200, response_model=AppConfig)

invokeai/app/api/routers/board_images.py

@@ -1,21 +1,12 @@
 from fastapi import Body, HTTPException
 from fastapi.routing import APIRouter
-from pydantic import BaseModel, Field
 
 from invokeai.app.api.dependencies import ApiDependencies
+from invokeai.app.services.images.images_common import AddImagesToBoardResult, RemoveImagesFromBoardResult
 
 board_images_router = APIRouter(prefix="/v1/board_images", tags=["boards"])
 
 
-class AddImagesToBoardResult(BaseModel):
-    board_id: str = Field(description="The id of the board the images were added to")
-    added_image_names: list[str] = Field(description="The image names that were added to the board")
-
-
-class RemoveImagesFromBoardResult(BaseModel):
-    removed_image_names: list[str] = Field(description="The image names that were removed from their board")
-
-
 @board_images_router.post(
     "/",
     operation_id="add_image_to_board",
@@ -23,17 +14,26 @@ class RemoveImagesFromBoardResult(BaseModel):
         201: {"description": "The image was added to a board successfully"},
     },
     status_code=201,
+    response_model=AddImagesToBoardResult,
 )
 async def add_image_to_board(
     board_id: str = Body(description="The id of the board to add to"),
     image_name: str = Body(description="The name of the image to add"),
-):
+) -> AddImagesToBoardResult:
     """Creates a board_image"""
     try:
-        result = ApiDependencies.invoker.services.board_images.add_image_to_board(
-            board_id=board_id, image_name=image_name
+        added_images: set[str] = set()
+        affected_boards: set[str] = set()
+        old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none"
+        ApiDependencies.invoker.services.board_images.add_image_to_board(board_id=board_id, image_name=image_name)
+        added_images.add(image_name)
+        affected_boards.add(board_id)
+        affected_boards.add(old_board_id)
+
+        return AddImagesToBoardResult(
+            added_images=list(added_images),
+            affected_boards=list(affected_boards),
         )
-        return result
     except Exception:
         raise HTTPException(status_code=500, detail="Failed to add image to board")
 
@@ -45,14 +45,25 @@ async def add_image_to_board(
         201: {"description": "The image was removed from the board successfully"},
     },
     status_code=201,
+    response_model=RemoveImagesFromBoardResult,
 )
 async def remove_image_from_board(
     image_name: str = Body(description="The name of the image to remove", embed=True),
-):
+) -> RemoveImagesFromBoardResult:
     """Removes an image from its board, if it had one"""
     try:
-        result = ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name)
-        return result
+        removed_images: set[str] = set()
+        affected_boards: set[str] = set()
+        old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none"
+        ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name)
+        removed_images.add(image_name)
+        affected_boards.add("none")
+        affected_boards.add(old_board_id)
+        return RemoveImagesFromBoardResult(
+            removed_images=list(removed_images),
+            affected_boards=list(affected_boards),
+        )
+
     except Exception:
         raise HTTPException(status_code=500, detail="Failed to remove image from board")
 
@@ -72,16 +83,25 @@ async def add_images_to_board(
 ) -> AddImagesToBoardResult:
     """Adds a list of images to a board"""
     try:
-        added_image_names: list[str] = []
+        added_images: set[str] = set()
+        affected_boards: set[str] = set()
         for image_name in image_names:
             try:
+                old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none"
                 ApiDependencies.invoker.services.board_images.add_image_to_board(
-                    board_id=board_id, image_name=image_name
+                    board_id=board_id,
+                    image_name=image_name,
                 )
-                added_image_names.append(image_name)
+                added_images.add(image_name)
+                affected_boards.add(board_id)
+                affected_boards.add(old_board_id)
+
             except Exception:
                 pass
-        return AddImagesToBoardResult(board_id=board_id, added_image_names=added_image_names)
+        return AddImagesToBoardResult(
+            added_images=list(added_images),
+            affected_boards=list(affected_boards),
+        )
     except Exception:
         raise HTTPException(status_code=500, detail="Failed to add images to board")
 
@@ -100,13 +120,20 @@ async def remove_images_from_board(
 ) -> RemoveImagesFromBoardResult:
     """Removes a list of images from their board, if they had one"""
     try:
-        removed_image_names: list[str] = []
+        removed_images: set[str] = set()
+        affected_boards: set[str] = set()
         for image_name in image_names:
             try:
+                old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none"
                 ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name)
-                removed_image_names.append(image_name)
+                removed_images.add(image_name)
+                affected_boards.add("none")
+                affected_boards.add(old_board_id)
             except Exception:
                 pass
-        return RemoveImagesFromBoardResult(removed_image_names=removed_image_names)
+        return RemoveImagesFromBoardResult(
+            removed_images=list(removed_images),
+            affected_boards=list(affected_boards),
+        )
     except Exception:
         raise HTTPException(status_code=500, detail="Failed to remove images from board")

invokeai/app/api/routers/boards.py

@@ -146,7 +146,7 @@ async def list_boards(
     response_model=list[str],
 )
 async def list_all_board_image_names(
-    board_id: str = Path(description="The id of the board"),
+    board_id: str = Path(description="The id of the board or 'none' for uncategorized images"),
     categories: list[ImageCategory] | None = Query(default=None, description="The categories of image to include."),
     is_intermediate: bool | None = Query(default=None, description="Whether to list intermediate images."),
 ) -> list[str]:

invokeai/app/api/routers/client_state.py

@@ -0,0 +1,58 @@
+from fastapi import Body, HTTPException, Path, Query
+from fastapi.routing import APIRouter
+
+from invokeai.app.api.dependencies import ApiDependencies
+from invokeai.backend.util.logging import logging
+
+client_state_router = APIRouter(prefix="/v1/client_state", tags=["client_state"])
+
+
+@client_state_router.get(
+    "/{queue_id}/get_by_key",
+    operation_id="get_client_state_by_key",
+    response_model=str | None,
+)
+async def get_client_state_by_key(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+    key: str = Query(..., description="Key to get"),
+) -> str | None:
+    """Gets the client state"""
+    try:
+        return ApiDependencies.invoker.services.client_state_persistence.get_by_key(queue_id, key)
+    except Exception as e:
+        logging.error(f"Error getting client state: {e}")
+        raise HTTPException(status_code=500, detail="Error setting client state")
+
+
+@client_state_router.post(
+    "/{queue_id}/set_by_key",
+    operation_id="set_client_state",
+    response_model=str,
+)
+async def set_client_state(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+    key: str = Query(..., description="Key to set"),
+    value: str = Body(..., description="Stringified value to set"),
+) -> str:
+    """Sets the client state"""
+    try:
+        return ApiDependencies.invoker.services.client_state_persistence.set_by_key(queue_id, key, value)
+    except Exception as e:
+        logging.error(f"Error setting client state: {e}")
+        raise HTTPException(status_code=500, detail="Error setting client state")
+
+
+@client_state_router.post(
+    "/{queue_id}/delete",
+    operation_id="delete_client_state",
+    responses={204: {"description": "Client state deleted"}},
+)
+async def delete_client_state(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+) -> None:
+    """Deletes the client state"""
+    try:
+        ApiDependencies.invoker.services.client_state_persistence.delete(queue_id)
+    except Exception as e:
+        logging.error(f"Error deleting client state: {e}")
+        raise HTTPException(status_code=500, detail="Error deleting client state")

invokeai/app/api/routers/images.py

@@ -1,24 +1,34 @@
 import io
+import json
 import traceback
-from typing import Optional
+from typing import ClassVar, Optional
 
 from fastapi import BackgroundTasks, Body, HTTPException, Path, Query, Request, Response, UploadFile
 from fastapi.responses import FileResponse
 from fastapi.routing import APIRouter
 from PIL import Image
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 from invokeai.app.api.dependencies import ApiDependencies
 from invokeai.app.api.extract_metadata_from_image import extract_metadata_from_image
 from invokeai.app.invocations.fields import MetadataField
 from invokeai.app.services.image_records.image_records_common import (
     ImageCategory,
+    ImageNamesResult,
     ImageRecordChanges,
     ResourceOrigin,
 )
-from invokeai.app.services.images.images_common import ImageDTO, ImageUrlsDTO
+from invokeai.app.services.images.images_common import (
+    DeleteImagesResult,
+    ImageDTO,
+    ImageUrlsDTO,
+    StarredImagesResult,
+    UnstarredImagesResult,
+)
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
+from invokeai.app.util.controlnet_utils import heuristic_resize_fast
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np
 
 images_router = APIRouter(prefix="/v1/images", tags=["images"])
 
@@ -27,6 +37,19 @@ images_router = APIRouter(prefix="/v1/images", tags=["images"])
 IMAGE_MAX_AGE = 31536000
 
 
+class ResizeToDimensions(BaseModel):
+    width: int = Field(..., gt=0)
+    height: int = Field(..., gt=0)
+
+    MAX_SIZE: ClassVar[int] = 4096 * 4096
+
+    @model_validator(mode="after")
+    def validate_total_output_size(self):
+        if self.width * self.height > self.MAX_SIZE:
+            raise ValueError(f"Max total output size for resizing is {self.MAX_SIZE} pixels")
+        return self
+
+
 @images_router.post(
     "/upload",
     operation_id="upload_image",
@@ -46,6 +69,11 @@ async def upload_image(
     board_id: Optional[str] = Query(default=None, description="The board to add this image to, if any"),
     session_id: Optional[str] = Query(default=None, description="The session ID associated with this upload, if any"),
     crop_visible: Optional[bool] = Query(default=False, description="Whether to crop the image"),
+    resize_to: Optional[str] = Body(
+        default=None,
+        description=f"Dimensions to resize the image to, must be stringified tuple of 2 integers. Max total pixel count: {ResizeToDimensions.MAX_SIZE}",
+        examples=['"[1024,1024]"'],
+    ),
     metadata: Optional[str] = Body(
         default=None,
         description="The metadata to associate with the image, must be a stringified JSON dict",
@@ -59,13 +87,33 @@ async def upload_image(
     contents = await file.read()
     try:
         pil_image = Image.open(io.BytesIO(contents))
-        if crop_visible:
-            bbox = pil_image.getbbox()
-            pil_image = pil_image.crop(bbox)
     except Exception:
         ApiDependencies.invoker.services.logger.error(traceback.format_exc())
         raise HTTPException(status_code=415, detail="Failed to read image")
 
+    if crop_visible:
+        try:
+            bbox = pil_image.getbbox()
+            pil_image = pil_image.crop(bbox)
+        except Exception:
+            raise HTTPException(status_code=500, detail="Failed to crop image")
+
+    if resize_to:
+        try:
+            dims = json.loads(resize_to)
+            resize_dims = ResizeToDimensions(**dims)
+        except Exception:
+            raise HTTPException(status_code=400, detail="Invalid resize_to format or size")
+
+        try:
+            # heuristic_resize_fast expects an RGB or RGBA image
+            pil_rgba = pil_image.convert("RGBA")
+            np_image = pil_to_np(pil_rgba)
+            np_image = heuristic_resize_fast(np_image, (resize_dims.width, resize_dims.height))
+            pil_image = np_to_pil(np_image)
+        except Exception:
+            raise HTTPException(status_code=500, detail="Failed to resize image")
+
     extracted_metadata = extract_metadata_from_image(
         pil_image=pil_image,
         invokeai_metadata_override=metadata,
@@ -112,18 +160,30 @@ async def create_image_upload_entry(
     raise HTTPException(status_code=501, detail="Not implemented")
 
 
-@images_router.delete("/i/{image_name}", operation_id="delete_image")
+@images_router.delete("/i/{image_name}", operation_id="delete_image", response_model=DeleteImagesResult)
 async def delete_image(
     image_name: str = Path(description="The name of the image to delete"),
-) -> None:
+) -> DeleteImagesResult:
     """Deletes an image"""
 
+    deleted_images: set[str] = set()
+    affected_boards: set[str] = set()
+
     try:
+        image_dto = ApiDependencies.invoker.services.images.get_dto(image_name)
+        board_id = image_dto.board_id or "none"
         ApiDependencies.invoker.services.images.delete(image_name)
+        deleted_images.add(image_name)
+        affected_boards.add(board_id)
     except Exception:
         # TODO: Does this need any exception handling at all?
         pass
 
+    return DeleteImagesResult(
+        deleted_images=list(deleted_images),
+        affected_boards=list(affected_boards),
+    )
+
 
 @images_router.delete("/intermediates", operation_id="clear_intermediates")
 async def clear_intermediates() -> int:
@@ -335,23 +395,52 @@ async def list_image_dtos(
     return image_dtos
 
 
-class DeleteImagesFromListResult(BaseModel):
-    deleted_images: list[str]
-
-
-@images_router.post("/delete", operation_id="delete_images_from_list", response_model=DeleteImagesFromListResult)
+@images_router.post("/delete", operation_id="delete_images_from_list", response_model=DeleteImagesResult)
 async def delete_images_from_list(
     image_names: list[str] = Body(description="The list of names of images to delete", embed=True),
-) -> DeleteImagesFromListResult:
+) -> DeleteImagesResult:
     try:
-        deleted_images: list[str] = []
+        deleted_images: set[str] = set()
+        affected_boards: set[str] = set()
+        for image_name in image_names:
+            try:
+                image_dto = ApiDependencies.invoker.services.images.get_dto(image_name)
+                board_id = image_dto.board_id or "none"
+                ApiDependencies.invoker.services.images.delete(image_name)
+                deleted_images.add(image_name)
+                affected_boards.add(board_id)
+            except Exception:
+                pass
+        return DeleteImagesResult(
+            deleted_images=list(deleted_images),
+            affected_boards=list(affected_boards),
+        )
+    except Exception:
+        raise HTTPException(status_code=500, detail="Failed to delete images")
+
+
+@images_router.delete("/uncategorized", operation_id="delete_uncategorized_images", response_model=DeleteImagesResult)
+async def delete_uncategorized_images() -> DeleteImagesResult:
+    """Deletes all images that are uncategorized"""
+
+    image_names = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board(
+        board_id="none", categories=None, is_intermediate=None
+    )
+
+    try:
+        deleted_images: set[str] = set()
+        affected_boards: set[str] = set()
         for image_name in image_names:
             try:
                 ApiDependencies.invoker.services.images.delete(image_name)
-                deleted_images.append(image_name)
+                deleted_images.add(image_name)
+                affected_boards.add("none")
             except Exception:
                 pass
-        return DeleteImagesFromListResult(deleted_images=deleted_images)
+        return DeleteImagesResult(
+            deleted_images=list(deleted_images),
+            affected_boards=list(affected_boards),
+        )
     except Exception:
         raise HTTPException(status_code=500, detail="Failed to delete images")
 
@@ -360,36 +449,50 @@ class ImagesUpdatedFromListResult(BaseModel):
     updated_image_names: list[str] = Field(description="The image names that were updated")
 
 
-@images_router.post("/star", operation_id="star_images_in_list", response_model=ImagesUpdatedFromListResult)
+@images_router.post("/star", operation_id="star_images_in_list", response_model=StarredImagesResult)
 async def star_images_in_list(
     image_names: list[str] = Body(description="The list of names of images to star", embed=True),
-) -> ImagesUpdatedFromListResult:
+) -> StarredImagesResult:
     try:
-        updated_image_names: list[str] = []
+        starred_images: set[str] = set()
+        affected_boards: set[str] = set()
         for image_name in image_names:
             try:
-                ApiDependencies.invoker.services.images.update(image_name, changes=ImageRecordChanges(starred=True))
-                updated_image_names.append(image_name)
+                updated_image_dto = ApiDependencies.invoker.services.images.update(
+                    image_name, changes=ImageRecordChanges(starred=True)
+                )
+                starred_images.add(image_name)
+                affected_boards.add(updated_image_dto.board_id or "none")
             except Exception:
                 pass
-        return ImagesUpdatedFromListResult(updated_image_names=updated_image_names)
+        return StarredImagesResult(
+            starred_images=list(starred_images),
+            affected_boards=list(affected_boards),
+        )
     except Exception:
         raise HTTPException(status_code=500, detail="Failed to star images")
 
 
-@images_router.post("/unstar", operation_id="unstar_images_in_list", response_model=ImagesUpdatedFromListResult)
+@images_router.post("/unstar", operation_id="unstar_images_in_list", response_model=UnstarredImagesResult)
 async def unstar_images_in_list(
     image_names: list[str] = Body(description="The list of names of images to unstar", embed=True),
-) -> ImagesUpdatedFromListResult:
+) -> UnstarredImagesResult:
     try:
-        updated_image_names: list[str] = []
+        unstarred_images: set[str] = set()
+        affected_boards: set[str] = set()
         for image_name in image_names:
             try:
-                ApiDependencies.invoker.services.images.update(image_name, changes=ImageRecordChanges(starred=False))
-                updated_image_names.append(image_name)
+                updated_image_dto = ApiDependencies.invoker.services.images.update(
+                    image_name, changes=ImageRecordChanges(starred=False)
+                )
+                unstarred_images.add(image_name)
+                affected_boards.add(updated_image_dto.board_id or "none")
             except Exception:
                 pass
-        return ImagesUpdatedFromListResult(updated_image_names=updated_image_names)
+        return UnstarredImagesResult(
+            unstarred_images=list(unstarred_images),
+            affected_boards=list(affected_boards),
+        )
     except Exception:
         raise HTTPException(status_code=500, detail="Failed to unstar images")
 
@@ -460,3 +563,61 @@ async def get_bulk_download_item(
         return response
     except Exception:
         raise HTTPException(status_code=404)
+
+
+@images_router.get("/names", operation_id="get_image_names")
+async def get_image_names(
+    image_origin: Optional[ResourceOrigin] = Query(default=None, description="The origin of images to list."),
+    categories: Optional[list[ImageCategory]] = Query(default=None, description="The categories of image to include."),
+    is_intermediate: Optional[bool] = Query(default=None, description="Whether to list intermediate images."),
+    board_id: Optional[str] = Query(
+        default=None,
+        description="The board id to filter by. Use 'none' to find images without a board.",
+    ),
+    order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The order of sort"),
+    starred_first: bool = Query(default=True, description="Whether to sort by starred images first"),
+    search_term: Optional[str] = Query(default=None, description="The term to search for"),
+) -> ImageNamesResult:
+    """Gets ordered list of image names with metadata for optimistic updates"""
+
+    try:
+        result = ApiDependencies.invoker.services.images.get_image_names(
+            starred_first=starred_first,
+            order_dir=order_dir,
+            image_origin=image_origin,
+            categories=categories,
+            is_intermediate=is_intermediate,
+            board_id=board_id,
+            search_term=search_term,
+        )
+        return result
+    except Exception:
+        raise HTTPException(status_code=500, detail="Failed to get image names")
+
+
+@images_router.post(
+    "/images_by_names",
+    operation_id="get_images_by_names",
+    responses={200: {"model": list[ImageDTO]}},
+)
+async def get_images_by_names(
+    image_names: list[str] = Body(embed=True, description="Object containing list of image names to fetch DTOs for"),
+) -> list[ImageDTO]:
+    """Gets image DTOs for the specified image names. Maintains order of input names."""
+
+    try:
+        image_service = ApiDependencies.invoker.services.images
+
+        # Fetch DTOs preserving the order of requested names
+        image_dtos: list[ImageDTO] = []
+        for name in image_names:
+            try:
+                dto = image_service.get_dto(name)
+                image_dtos.append(dto)
+            except Exception:
+                # Skip missing images - they may have been deleted between name fetch and DTO fetch
+                continue
+
+        return image_dtos
+    except Exception:
+        raise HTTPException(status_code=500, detail="Failed to get image DTOs")

invokeai/app/api/routers/model_manager.py

@@ -41,6 +41,7 @@ from invokeai.backend.model_manager.starter_models import (
     STARTER_BUNDLES,
     STARTER_MODELS,
     StarterModel,
+    StarterModelBundle,
     StarterModelWithoutDependencies,
 )
 
@@ -291,7 +292,7 @@ async def get_hugging_face_models(
 )
 async def update_model_record(
     key: Annotated[str, Path(description="Unique key of model")],
-    changes: Annotated[ModelRecordChanges, Body(description="Model config", example=example_model_input)],
+    changes: Annotated[ModelRecordChanges, Body(description="Model config", examples=[example_model_input])],
 ) -> AnyModelConfig:
     """Update a model's config."""
     logger = ApiDependencies.invoker.services.logger
@@ -449,7 +450,7 @@ async def install_model(
     access_token: Optional[str] = Query(description="access token for the remote resource", default=None),
     config: ModelRecordChanges = Body(
         description="Object containing fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
-        example={"name": "string", "description": "string"},
+        examples=[{"name": "string", "description": "string"}],
     ),
 ) -> ModelInstallJob:
     """Install a model using a string identifier.
@@ -799,7 +800,7 @@ async def convert_model(
 
 class StarterModelResponse(BaseModel):
     starter_models: list[StarterModel]
-    starter_bundles: dict[str, list[StarterModel]]
+    starter_bundles: dict[str, StarterModelBundle]
 
 
 def get_is_installed(
@@ -833,7 +834,7 @@ async def get_starter_models() -> StarterModelResponse:
         model.dependencies = missing_deps
 
     for bundle in starter_bundles.values():
-        for model in bundle:
+        for model in bundle.models:
             model.is_installed = get_is_installed(model, installed_models)
             # Remove already-installed dependencies
             missing_deps: list[StarterModelWithoutDependencies] = []
@@ -893,6 +894,12 @@ class HFTokenHelper:
             huggingface_hub.login(token=token, add_to_git_credential=False)
         return cls.get_status()
 
+    @classmethod
+    def reset_token(cls) -> HFTokenStatus:
+        with SuppressOutput(), contextlib.suppress(Exception):
+            huggingface_hub.logout()
+        return cls.get_status()
+
 
 @model_manager_router.get("/hf_login", operation_id="get_hf_login_status", response_model=HFTokenStatus)
 async def get_hf_login_status() -> HFTokenStatus:
@@ -915,3 +922,8 @@ async def do_hf_login(
         ApiDependencies.invoker.services.logger.warning("Unable to verify HF token")
 
     return token_status
+
+
+@model_manager_router.delete("/hf_login", operation_id="reset_hf_token", response_model=HFTokenStatus)
+async def reset_hf_token() -> HFTokenStatus:
+    return HFTokenHelper.reset_token()

invokeai/app/api/routers/model_relationships.py

@@ -0,0 +1,215 @@
+"""FastAPI route for model relationship records."""
+
+from typing import List
+
+from fastapi import APIRouter, Body, HTTPException, Path, status
+from pydantic import BaseModel, Field
+
+from invokeai.app.api.dependencies import ApiDependencies
+
+model_relationships_router = APIRouter(prefix="/v1/model_relationships", tags=["model_relationships"])
+
+# === Schemas ===
+
+
+class ModelRelationshipCreateRequest(BaseModel):
+    model_key_1: str = Field(
+        ...,
+        description="The key of the first model in the relationship",
+        examples=[
+            "aa3b247f-90c9-4416-bfcd-aeaa57a5339e",
+            "ac32b914-10ab-496e-a24a-3068724b9c35",
+            "d944abfd-c7c3-42e2-a4ff-da640b29b8b4",
+            "b1c2d3e4-f5a6-7890-abcd-ef1234567890",
+            "12345678-90ab-cdef-1234-567890abcdef",
+            "fedcba98-7654-3210-fedc-ba9876543210",
+        ],
+    )
+    model_key_2: str = Field(
+        ...,
+        description="The key of the second model in the relationship",
+        examples=[
+            "3bb7c0eb-b6c8-469c-ad8c-4d69c06075e4",
+            "f0c3da4e-d9ff-42b5-a45c-23be75c887c9",
+            "38170dd8-f1e5-431e-866c-2c81f1277fcc",
+            "c57fea2d-7646-424c-b9ad-c0ba60fc68be",
+            "10f7807b-ab54-46a9-ab03-600e88c630a1",
+            "f6c1d267-cf87-4ee0-bee0-37e791eacab7",
+        ],
+    )
+
+
+class ModelRelationshipBatchRequest(BaseModel):
+    model_keys: List[str] = Field(
+        ...,
+        description="List of model keys to fetch related models for",
+        examples=[
+            [
+                "aa3b247f-90c9-4416-bfcd-aeaa57a5339e",
+                "ac32b914-10ab-496e-a24a-3068724b9c35",
+            ],
+            [
+                "b1c2d3e4-f5a6-7890-abcd-ef1234567890",
+                "12345678-90ab-cdef-1234-567890abcdef",
+                "fedcba98-7654-3210-fedc-ba9876543210",
+            ],
+            [
+                "3bb7c0eb-b6c8-469c-ad8c-4d69c06075e4",
+            ],
+        ],
+    )
+
+
+# === Routes ===
+
+
+@model_relationships_router.get(
+    "/i/{model_key}",
+    operation_id="get_related_models",
+    response_model=list[str],
+    responses={
+        200: {
+            "description": "A list of related model keys was retrieved successfully",
+            "content": {
+                "application/json": {
+                    "example": [
+                        "15e9eb28-8cfe-47c9-b610-37907a79fc3c",
+                        "71272e82-0e5f-46d5-bca9-9a61f4bd8a82",
+                        "a5d7cd49-1b98-4534-a475-aeee4ccf5fa2",
+                    ]
+                }
+            },
+        },
+        404: {"description": "The specified model could not be found"},
+        422: {"description": "Validation error"},
+    },
+)
+async def get_related_models(
+    model_key: str = Path(..., description="The key of the model to get relationships for"),
+) -> list[str]:
+    """
+    Get a list of model keys related to a given model.
+    """
+    try:
+        return ApiDependencies.invoker.services.model_relationships.get_related_model_keys(model_key)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@model_relationships_router.post(
+    "/",
+    status_code=status.HTTP_204_NO_CONTENT,
+    responses={
+        204: {"description": "The relationship was successfully created"},
+        400: {"description": "Invalid model keys or self-referential relationship"},
+        409: {"description": "The relationship already exists"},
+        422: {"description": "Validation error"},
+        500: {"description": "Internal server error"},
+    },
+    summary="Add Model Relationship",
+    description="Creates a **bidirectional** relationship between two models, allowing each to reference the other as related.",
+)
+async def add_model_relationship(
+    req: ModelRelationshipCreateRequest = Body(..., description="The model keys to relate"),
+) -> None:
+    """
+    Add a relationship between two models.
+
+    Relationships are bidirectional and will be accessible from both models.
+
+    - Raises 400 if keys are invalid or identical.
+    - Raises 409 if the relationship already exists.
+    """
+    try:
+        if req.model_key_1 == req.model_key_2:
+            raise HTTPException(status_code=400, detail="Cannot relate a model to itself.")
+
+        ApiDependencies.invoker.services.model_relationships.add_model_relationship(
+            req.model_key_1,
+            req.model_key_2,
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=409, detail=str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@model_relationships_router.delete(
+    "/",
+    status_code=status.HTTP_204_NO_CONTENT,
+    responses={
+        204: {"description": "The relationship was successfully removed"},
+        400: {"description": "Invalid model keys or self-referential relationship"},
+        404: {"description": "The relationship does not exist"},
+        422: {"description": "Validation error"},
+        500: {"description": "Internal server error"},
+    },
+    summary="Remove Model Relationship",
+    description="Removes a **bidirectional** relationship between two models. The relationship must already exist.",
+)
+async def remove_model_relationship(
+    req: ModelRelationshipCreateRequest = Body(..., description="The model keys to disconnect"),
+) -> None:
+    """
+    Removes a bidirectional relationship between two model keys.
+
+    - Raises 400 if attempting to unlink a model from itself.
+    - Raises 404 if the relationship was not found.
+    """
+    try:
+        if req.model_key_1 == req.model_key_2:
+            raise HTTPException(status_code=400, detail="Cannot unlink a model from itself.")
+
+        ApiDependencies.invoker.services.model_relationships.remove_model_relationship(
+            req.model_key_1,
+            req.model_key_2,
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@model_relationships_router.post(
+    "/batch",
+    operation_id="get_related_models_batch",
+    response_model=List[str],
+    responses={
+        200: {
+            "description": "Related model keys retrieved successfully",
+            "content": {
+                "application/json": {
+                    "example": [
+                        "ca562b14-995e-4a42-90c1-9528f1a5921d",
+                        "cc0c2b8a-c62e-41d6-878e-cc74dde5ca8f",
+                        "18ca7649-6a9e-47d5-bc17-41ab1e8cec81",
+                        "7c12d1b2-0ef9-4bec-ba55-797b2d8f2ee1",
+                        "c382eaa3-0e28-4ab0-9446-408667699aeb",
+                        "71272e82-0e5f-46d5-bca9-9a61f4bd8a82",
+                        "a5d7cd49-1b98-4534-a475-aeee4ccf5fa2",
+                    ]
+                }
+            },
+        },
+        422: {"description": "Validation error"},
+        500: {"description": "Internal server error"},
+    },
+    summary="Get Related Model Keys (Batch)",
+    description="Retrieves all **unique related model keys** for a list of given models. This is useful for contextual suggestions or filtering.",
+)
+async def get_related_models_batch(
+    req: ModelRelationshipBatchRequest = Body(..., description="Model keys to check for related connections"),
+) -> list[str]:
+    """
+    Accepts multiple model keys and returns a flat list of all unique related keys.
+
+    Useful when working with multiple selections in the UI or cross-model comparisons.
+    """
+    try:
+        all_related: set[str] = set()
+        for key in req.model_keys:
+            related = ApiDependencies.invoker.services.model_relationships.get_related_model_keys(key)
+            all_related.update(related)
+        return list(all_related)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

invokeai/app/api/routers/session_queue.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from fastapi import Body, Path, Query
+from fastapi import Body, HTTPException, Path, Query
 from fastapi.routing import APIRouter
 from pydantic import BaseModel, Field
 
@@ -14,13 +14,15 @@ from invokeai.app.services.session_queue.session_queue_common import (
     CancelByBatchIDsResult,
     CancelByDestinationResult,
     ClearResult,
+    DeleteAllExceptCurrentResult,
+    DeleteByDestinationResult,
     EnqueueBatchResult,
     FieldIdentifier,
     PruneResult,
     RetryItemsResult,
     SessionQueueCountsByDestination,
     SessionQueueItem,
-    SessionQueueItemDTO,
+    SessionQueueItemNotFoundError,
     SessionQueueStatus,
 )
 from invokeai.app.services.shared.pagination import CursorPaginatedResults
@@ -58,17 +60,19 @@ async def enqueue_batch(
     ),
 ) -> EnqueueBatchResult:
     """Processes a batch and enqueues the output graphs for execution."""
-
-    return await ApiDependencies.invoker.services.session_queue.enqueue_batch(
-        queue_id=queue_id, batch=batch, prepend=prepend
-    )
+    try:
+        return await ApiDependencies.invoker.services.session_queue.enqueue_batch(
+            queue_id=queue_id, batch=batch, prepend=prepend
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while enqueuing batch: {e}")
 
 
 @session_queue_router.get(
     "/{queue_id}/list",
     operation_id="list_queue_items",
     responses={
-        200: {"model": CursorPaginatedResults[SessionQueueItemDTO]},
+        200: {"model": CursorPaginatedResults[SessionQueueItem]},
     },
 )
 async def list_queue_items(
@@ -77,12 +81,42 @@ async def list_queue_items(
     status: Optional[QUEUE_ITEM_STATUS] = Query(default=None, description="The status of items to fetch"),
     cursor: Optional[int] = Query(default=None, description="The pagination cursor"),
     priority: int = Query(default=0, description="The pagination cursor priority"),
-) -> CursorPaginatedResults[SessionQueueItemDTO]:
-    """Gets all queue items (without graphs)"""
+    destination: Optional[str] = Query(default=None, description="The destination of queue items to fetch"),
+) -> CursorPaginatedResults[SessionQueueItem]:
+    """Gets cursor-paginated queue items"""
 
-    return ApiDependencies.invoker.services.session_queue.list_queue_items(
-        queue_id=queue_id, limit=limit, status=status, cursor=cursor, priority=priority
-    )
+    try:
+        return ApiDependencies.invoker.services.session_queue.list_queue_items(
+            queue_id=queue_id,
+            limit=limit,
+            status=status,
+            cursor=cursor,
+            priority=priority,
+            destination=destination,
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while listing all items: {e}")
+
+
+@session_queue_router.get(
+    "/{queue_id}/list_all",
+    operation_id="list_all_queue_items",
+    responses={
+        200: {"model": list[SessionQueueItem]},
+    },
+)
+async def list_all_queue_items(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+    destination: Optional[str] = Query(default=None, description="The destination of queue items to fetch"),
+) -> list[SessionQueueItem]:
+    """Gets all queue items"""
+    try:
+        return ApiDependencies.invoker.services.session_queue.list_all_queue_items(
+            queue_id=queue_id,
+            destination=destination,
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while listing all queue items: {e}")
 
 
 @session_queue_router.put(
@@ -94,7 +128,10 @@ async def resume(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> SessionProcessorStatus:
     """Resumes session processor"""
-    return ApiDependencies.invoker.services.session_processor.resume()
+    try:
+        return ApiDependencies.invoker.services.session_processor.resume()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while resuming queue: {e}")
 
 
 @session_queue_router.put(
@@ -106,7 +143,10 @@ async def Pause(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> SessionProcessorStatus:
     """Pauses session processor"""
-    return ApiDependencies.invoker.services.session_processor.pause()
+    try:
+        return ApiDependencies.invoker.services.session_processor.pause()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while pausing queue: {e}")
 
 
 @session_queue_router.put(
@@ -118,7 +158,25 @@ async def cancel_all_except_current(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> CancelAllExceptCurrentResult:
     """Immediately cancels all queue items except in-processing items"""
-    return ApiDependencies.invoker.services.session_queue.cancel_all_except_current(queue_id=queue_id)
+    try:
+        return ApiDependencies.invoker.services.session_queue.cancel_all_except_current(queue_id=queue_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while canceling all except current: {e}")
+
+
+@session_queue_router.put(
+    "/{queue_id}/delete_all_except_current",
+    operation_id="delete_all_except_current",
+    responses={200: {"model": DeleteAllExceptCurrentResult}},
+)
+async def delete_all_except_current(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+) -> DeleteAllExceptCurrentResult:
+    """Immediately deletes all queue items except in-processing items"""
+    try:
+        return ApiDependencies.invoker.services.session_queue.delete_all_except_current(queue_id=queue_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while deleting all except current: {e}")
 
 
 @session_queue_router.put(
@@ -131,7 +189,12 @@ async def cancel_by_batch_ids(
     batch_ids: list[str] = Body(description="The list of batch_ids to cancel all queue items for", embed=True),
 ) -> CancelByBatchIDsResult:
     """Immediately cancels all queue items from the given batch ids"""
-    return ApiDependencies.invoker.services.session_queue.cancel_by_batch_ids(queue_id=queue_id, batch_ids=batch_ids)
+    try:
+        return ApiDependencies.invoker.services.session_queue.cancel_by_batch_ids(
+            queue_id=queue_id, batch_ids=batch_ids
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while canceling by batch id: {e}")
 
 
 @session_queue_router.put(
@@ -144,9 +207,12 @@ async def cancel_by_destination(
     destination: str = Query(description="The destination to cancel all queue items for"),
 ) -> CancelByDestinationResult:
     """Immediately cancels all queue items with the given origin"""
-    return ApiDependencies.invoker.services.session_queue.cancel_by_destination(
-        queue_id=queue_id, destination=destination
-    )
+    try:
+        return ApiDependencies.invoker.services.session_queue.cancel_by_destination(
+            queue_id=queue_id, destination=destination
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while canceling by destination: {e}")
 
 
 @session_queue_router.put(
@@ -159,7 +225,10 @@ async def retry_items_by_id(
     item_ids: list[int] = Body(description="The queue item ids to retry"),
 ) -> RetryItemsResult:
     """Immediately cancels all queue items with the given origin"""
-    return ApiDependencies.invoker.services.session_queue.retry_items_by_id(queue_id=queue_id, item_ids=item_ids)
+    try:
+        return ApiDependencies.invoker.services.session_queue.retry_items_by_id(queue_id=queue_id, item_ids=item_ids)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while retrying queue items: {e}")
 
 
 @session_queue_router.put(
@@ -173,11 +242,14 @@ async def clear(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> ClearResult:
     """Clears the queue entirely, immediately canceling the currently-executing session"""
-    queue_item = ApiDependencies.invoker.services.session_queue.get_current(queue_id)
-    if queue_item is not None:
-        ApiDependencies.invoker.services.session_queue.cancel_queue_item(queue_item.item_id)
-    clear_result = ApiDependencies.invoker.services.session_queue.clear(queue_id)
-    return clear_result
+    try:
+        queue_item = ApiDependencies.invoker.services.session_queue.get_current(queue_id)
+        if queue_item is not None:
+            ApiDependencies.invoker.services.session_queue.cancel_queue_item(queue_item.item_id)
+        clear_result = ApiDependencies.invoker.services.session_queue.clear(queue_id)
+        return clear_result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while clearing queue: {e}")
 
 
 @session_queue_router.put(
@@ -191,7 +263,10 @@ async def prune(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> PruneResult:
     """Prunes all completed or errored queue items"""
-    return ApiDependencies.invoker.services.session_queue.prune(queue_id)
+    try:
+        return ApiDependencies.invoker.services.session_queue.prune(queue_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while pruning queue: {e}")
 
 
 @session_queue_router.get(
@@ -205,7 +280,10 @@ async def get_current_queue_item(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> Optional[SessionQueueItem]:
     """Gets the currently execution queue item"""
-    return ApiDependencies.invoker.services.session_queue.get_current(queue_id)
+    try:
+        return ApiDependencies.invoker.services.session_queue.get_current(queue_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while getting current queue item: {e}")
 
 
 @session_queue_router.get(
@@ -219,7 +297,10 @@ async def get_next_queue_item(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> Optional[SessionQueueItem]:
     """Gets the next queue item, without executing it"""
-    return ApiDependencies.invoker.services.session_queue.get_next(queue_id)
+    try:
+        return ApiDependencies.invoker.services.session_queue.get_next(queue_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while getting next queue item: {e}")
 
 
 @session_queue_router.get(
@@ -233,9 +314,12 @@ async def get_queue_status(
     queue_id: str = Path(description="The queue id to perform this operation on"),
 ) -> SessionQueueAndProcessorStatus:
     """Gets the status of the session queue"""
-    queue = ApiDependencies.invoker.services.session_queue.get_queue_status(queue_id)
-    processor = ApiDependencies.invoker.services.session_processor.get_status()
-    return SessionQueueAndProcessorStatus(queue=queue, processor=processor)
+    try:
+        queue = ApiDependencies.invoker.services.session_queue.get_queue_status(queue_id)
+        processor = ApiDependencies.invoker.services.session_processor.get_status()
+        return SessionQueueAndProcessorStatus(queue=queue, processor=processor)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while getting queue status: {e}")
 
 
 @session_queue_router.get(
@@ -250,7 +334,10 @@ async def get_batch_status(
     batch_id: str = Path(description="The batch to get the status of"),
 ) -> BatchStatus:
     """Gets the status of the session queue"""
-    return ApiDependencies.invoker.services.session_queue.get_batch_status(queue_id=queue_id, batch_id=batch_id)
+    try:
+        return ApiDependencies.invoker.services.session_queue.get_batch_status(queue_id=queue_id, batch_id=batch_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while getting batch status: {e}")
 
 
 @session_queue_router.get(
@@ -266,7 +353,27 @@ async def get_queue_item(
     item_id: int = Path(description="The queue item to get"),
 ) -> SessionQueueItem:
     """Gets a queue item"""
-    return ApiDependencies.invoker.services.session_queue.get_queue_item(item_id)
+    try:
+        return ApiDependencies.invoker.services.session_queue.get_queue_item(item_id)
+    except SessionQueueItemNotFoundError:
+        raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while fetching queue item: {e}")
+
+
+@session_queue_router.delete(
+    "/{queue_id}/i/{item_id}",
+    operation_id="delete_queue_item",
+)
+async def delete_queue_item(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+    item_id: int = Path(description="The queue item to delete"),
+) -> None:
+    """Deletes a queue item"""
+    try:
+        ApiDependencies.invoker.services.session_queue.delete_queue_item(item_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while deleting queue item: {e}")
 
 
 @session_queue_router.put(
@@ -281,8 +388,12 @@ async def cancel_queue_item(
     item_id: int = Path(description="The queue item to cancel"),
 ) -> SessionQueueItem:
     """Deletes a queue item"""
-
-    return ApiDependencies.invoker.services.session_queue.cancel_queue_item(item_id)
+    try:
+        return ApiDependencies.invoker.services.session_queue.cancel_queue_item(item_id)
+    except SessionQueueItemNotFoundError:
+        raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while canceling queue item: {e}")
 
 
 @session_queue_router.get(
@@ -295,6 +406,27 @@ async def counts_by_destination(
     destination: str = Query(description="The destination to query"),
 ) -> SessionQueueCountsByDestination:
     """Gets the counts of queue items by destination"""
-    return ApiDependencies.invoker.services.session_queue.get_counts_by_destination(
-        queue_id=queue_id, destination=destination
-    )
+    try:
+        return ApiDependencies.invoker.services.session_queue.get_counts_by_destination(
+            queue_id=queue_id, destination=destination
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while fetching counts by destination: {e}")
+
+
+@session_queue_router.delete(
+    "/{queue_id}/d/{destination}",
+    operation_id="delete_by_destination",
+    responses={200: {"model": DeleteByDestinationResult}},
+)
+async def delete_by_destination(
+    queue_id: str = Path(description="The queue id to query"),
+    destination: str = Path(description="The destination to query"),
+) -> DeleteByDestinationResult:
+    """Deletes all items with the given destination"""
+    try:
+        return ApiDependencies.invoker.services.session_queue.delete_by_destination(
+            queue_id=queue_id, destination=destination
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while deleting by destination: {e}")

invokeai/app/api_app.py

@@ -19,9 +19,11 @@ from invokeai.app.api.routers import (
     app_info,
     board_images,
     boards,
+    client_state,
     download_queue,
     images,
     model_manager,
+    model_relationships,
     session_queue,
     style_presets,
     utilities,
@@ -125,10 +127,12 @@ app.include_router(download_queue.download_queue_router, prefix="/api")
 app.include_router(images.images_router, prefix="/api")
 app.include_router(boards.boards_router, prefix="/api")
 app.include_router(board_images.board_images_router, prefix="/api")
+app.include_router(model_relationships.model_relationships_router, prefix="/api")
 app.include_router(app_info.app_router, prefix="/api")
 app.include_router(session_queue.session_queue_router, prefix="/api")
 app.include_router(workflows.workflows_router, prefix="/api")
 app.include_router(style_presets.style_presets_router, prefix="/api")
+app.include_router(client_state.client_state_router, prefix="/api")
 
 app.openapi = get_openapi_func(app)
 
@@ -153,10 +157,16 @@ def overridden_redoc() -> HTMLResponse:
 
 web_root_path = Path(list(web_dir.__path__)[0])
 
+if app_config.unsafe_disable_picklescan:
+    logger.warning(
+        "The unsafe_disable_picklescan option is enabled. This disables malware scanning while installing and"
+        "loading models, which may allow malicious code to be executed. Use at your own risk."
+    )
+
 try:
     app.mount("/", NoCacheStaticFiles(directory=Path(web_root_path, "dist"), html=True), name="ui")
 except RuntimeError:
-    logger.warn(f"No UI found at {web_root_path}/dist, skipping UI mount")
+    logger.warning(f"No UI found at {web_root_path}/dist, skipping UI mount")
 app.mount(
     "/static", NoCacheStaticFiles(directory=Path(web_root_path, "static/")), name="static"
 )  # docs favicon is in here

invokeai/app/invocations/baseinvocation.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 import inspect
 import re
 import sys
+import types
+import typing
 import warnings
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -20,8 +22,10 @@ from typing import (
     Literal,
     Optional,
     Type,
+    TypedDict,
     TypeVar,
     Union,
+    cast,
 )
 
 import semver
@@ -72,6 +76,17 @@ class Classification(str, Enum, metaclass=MetaEnum):
     Special = "special"
 
 
+class Bottleneck(str, Enum, metaclass=MetaEnum):
+    """
+    The bottleneck of an invocation.
+    - `Network`: The invocation's execution is network-bound.
+    - `GPU`: The invocation's execution is GPU-bound.
+    """
+
+    Network = "network"
+    GPU = "gpu"
+
+
 class UIConfigBase(BaseModel):
     """
     Provides additional node configuration to the UI.
@@ -93,6 +108,11 @@ class UIConfigBase(BaseModel):
     )
 
 
+class OriginalModelField(TypedDict):
+    annotation: Any
+    field_info: FieldInfo
+
+
 class BaseInvocationOutput(BaseModel):
     """
     Base class for all invocation outputs.
@@ -121,6 +141,9 @@ class BaseInvocationOutput(BaseModel):
         """Gets the invocation output's type, as provided by the `@invocation_output` decorator."""
         return cls.model_fields["type"].default
 
+    _original_model_fields: ClassVar[dict[str, OriginalModelField]] = {}
+    """The original model fields, before any modifications were made by the @invocation_output decorator."""
+
     model_config = ConfigDict(
         protected_namespaces=(),
         validate_assignment=True,
@@ -154,7 +177,7 @@ class BaseInvocation(ABC, BaseModel):
         return cls.model_fields["type"].default
 
     @classmethod
-    def get_output_annotation(cls) -> BaseInvocationOutput:
+    def get_output_annotation(cls) -> Type[BaseInvocationOutput]:
         """Gets the invocation's output annotation (i.e. the return annotation of its `invoke()` method)."""
         return signature(cls.invoke).return_annotation
 
@@ -186,7 +209,7 @@ class BaseInvocation(ABC, BaseModel):
         Internal invoke method, calls `invoke()` after some prep.
         Handles optional fields that are required to call `invoke()` and invocation cache.
         """
-        for field_name, field in self.model_fields.items():
+        for field_name, field in type(self).model_fields.items():
             if not field.json_schema_extra or callable(field.json_schema_extra):
                 # something has gone terribly awry, we should always have this and it should be a dict
                 continue
@@ -201,9 +224,9 @@ class BaseInvocation(ABC, BaseModel):
                 setattr(self, field_name, orig_default)
             if orig_required and orig_default is PydanticUndefined and getattr(self, field_name) is None:
                 if input_ == Input.Connection:
-                    raise RequiredConnectionException(self.model_fields["type"].default, field_name)
+                    raise RequiredConnectionException(type(self).model_fields["type"].default, field_name)
                 elif input_ == Input.Any:
-                    raise MissingInputException(self.model_fields["type"].default, field_name)
+                    raise MissingInputException(type(self).model_fields["type"].default, field_name)
 
         # skip node cache codepath if it's disabled
         if services.configuration.node_cache_size == 0:
@@ -241,6 +264,8 @@ class BaseInvocation(ABC, BaseModel):
         json_schema_extra={"field_kind": FieldKind.NodeAttribute},
     )
 
+    bottleneck: ClassVar[Bottleneck]
+
     UIConfig: ClassVar[UIConfigBase]
 
     model_config = ConfigDict(
@@ -251,6 +276,9 @@ class BaseInvocation(ABC, BaseModel):
         coerce_numbers_to_str=True,
     )
 
+    _original_model_fields: ClassVar[dict[str, OriginalModelField]] = {}
+    """The original model fields, before any modifications were made by the @invocation decorator."""
+
 
 TBaseInvocation = TypeVar("TBaseInvocation", bound=BaseInvocation)
 
@@ -399,6 +427,7 @@ RESERVED_NODE_ATTRIBUTE_FIELD_NAMES = {
     "use_cache",
     "type",
     "workflow",
+    "bottleneck",
 }
 
 RESERVED_INPUT_FIELD_NAMES = {"metadata", "board"}
@@ -470,11 +499,53 @@ def validate_fields(model_fields: dict[str, FieldInfo], model_type: str) -> None
 
         ui_type = field.json_schema_extra.get("ui_type", None)
         if isinstance(ui_type, str) and ui_type.startswith("DEPRECATED_"):
-            logger.warn(f'"UIType.{ui_type.split("_")[-1]}" is deprecated, ignoring')
+            logger.warning(f'"UIType.{ui_type.split("_")[-1]}" is deprecated, ignoring')
             field.json_schema_extra.pop("ui_type")
     return None
 
 
+class NoDefaultSentinel:
+    pass
+
+
+def validate_field_default(
+    cls_name: str, field_name: str, invocation_type: str, annotation: Any, field_info: FieldInfo
+) -> None:
+    """Validates the default value of a field against its pydantic field definition."""
+
+    assert isinstance(field_info.json_schema_extra, dict), "json_schema_extra is not a dict"
+
+    # By the time we are doing this, we've already done some pydantic magic by overriding the original default value.
+    # We store the original default value in the json_schema_extra dict, so we can validate it here.
+    orig_default = field_info.json_schema_extra.get("orig_default", NoDefaultSentinel)
+
+    if orig_default is NoDefaultSentinel:
+        return
+
+    # To validate the default value, we can create a temporary pydantic model with the field we are validating as its
+    # only field. Then validate the default value against this temporary model.
+    TempDefaultValidator = cast(BaseModel, create_model(cls_name, **{field_name: (annotation, field_info)}))
+
+    try:
+        TempDefaultValidator.model_validate({field_name: orig_default})
+    except Exception as e:
+        raise InvalidFieldError(
+            f'Default value for field "{field_name}" on invocation "{invocation_type}" is invalid, {e}'
+        ) from e
+
+
+def is_optional(annotation: Any) -> bool:
+    """
+    Checks if the given annotation is optional (i.e. Optional[X], Union[X, None] or X | None).
+    """
+    origin = typing.get_origin(annotation)
+    # PEP 604 unions (int|None) have origin types.UnionType
+    is_union = origin is typing.Union or origin is types.UnionType
+    if not is_union:
+        return False
+    return any(arg is type(None) for arg in typing.get_args(annotation))
+
+
 def invocation(
     invocation_type: str,
     title: Optional[str] = None,
@@ -483,6 +554,7 @@ def invocation(
     version: Optional[str] = None,
     use_cache: Optional[bool] = True,
     classification: Classification = Classification.Stable,
+    bottleneck: Bottleneck = Bottleneck.GPU,
 ) -> Callable[[Type[TBaseInvocation]], Type[TBaseInvocation]]:
     """
     Registers an invocation.
@@ -494,6 +566,7 @@ def invocation(
     :param Optional[str] version: Adds a version to the invocation. Must be a valid semver string. Defaults to None.
     :param Optional[bool] use_cache: Whether or not to use the invocation cache. Defaults to True. The user may override this in the workflow editor.
     :param Classification classification: The classification of the invocation. Defaults to FeatureClassification.Stable. Use Beta or Prototype if the invocation is unstable.
+    :param Bottleneck bottleneck: The bottleneck of the invocation. Defaults to Bottleneck.GPU. Use Network if the invocation is network-bound.
     """
 
     def wrapper(cls: Type[TBaseInvocation]) -> Type[TBaseInvocation]:
@@ -507,6 +580,26 @@ def invocation(
 
         validate_fields(cls.model_fields, invocation_type)
 
+        fields: dict[str, tuple[Any, FieldInfo]] = {}
+
+        original_model_fields: dict[str, OriginalModelField] = {}
+
+        for field_name, field_info in cls.model_fields.items():
+            annotation = field_info.annotation
+            assert annotation is not None, f"{field_name} on invocation {invocation_type} has no type annotation."
+            assert isinstance(field_info.json_schema_extra, dict), (
+                f"{field_name} on invocation {invocation_type} has a non-dict json_schema_extra, did you forget to use InputField?"
+            )
+
+            original_model_fields[field_name] = OriginalModelField(annotation=annotation, field_info=field_info)
+
+            validate_field_default(cls.__name__, field_name, invocation_type, annotation, field_info)
+
+            if field_info.default is None and not is_optional(annotation):
+                annotation = annotation | None
+
+            fields[field_name] = (annotation, field_info)
+
         # Add OpenAPI schema extras
         uiconfig: dict[str, Any] = {}
         uiconfig["title"] = title
@@ -522,7 +615,7 @@ def invocation(
                 raise InvalidVersionError(f'Invalid version string for node "{invocation_type}": "{version}"') from e
             uiconfig["version"] = version
         else:
-            logger.warn(f'No version specified for node "{invocation_type}", using "1.0.0"')
+            logger.warning(f'No version specified for node "{invocation_type}", using "1.0.0"')
             uiconfig["version"] = "1.0.0"
 
         cls.UIConfig = UIConfigBase(**uiconfig)
@@ -530,6 +623,8 @@ def invocation(
         if use_cache is not None:
             cls.model_fields["use_cache"].default = use_cache
 
+        cls.bottleneck = bottleneck
+
         # Add the invocation type to the model.
 
         # You'd be tempted to just add the type field and rebuild the model, like this:
@@ -539,11 +634,27 @@ def invocation(
         # Unfortunately, because the `GraphInvocation` uses a forward ref in its `graph` field's annotation, this does
         # not work. Instead, we have to create a new class with the type field and patch the original class with it.
 
-        invocation_type_annotation = Literal[invocation_type]  # type: ignore
-        invocation_type_field = Field(
-            title="type", default=invocation_type, json_schema_extra={"field_kind": FieldKind.NodeAttribute}
+        invocation_type_annotation = Literal[invocation_type]
+
+        # Field() returns an instance of FieldInfo, but thanks to a pydantic implementation detail, it is _typed_ as Any.
+        # This cast makes the type annotation match the class's true type.
+        invocation_type_field_info = cast(
+            FieldInfo,
+            Field(title="type", default=invocation_type, json_schema_extra={"field_kind": FieldKind.NodeAttribute}),
         )
 
+        fields["type"] = (invocation_type_annotation, invocation_type_field_info)
+
+        # Invocation outputs must be registered using the @invocation_output decorator, but it is possible that the
+        # output is registered _after_ this invocation is registered. It depends on module import ordering.
+        #
+        # We can only confirm the output for an invocation is registered after all modules are imported. There's
+        # only really one good time to do that - during application startup, in `run_app.py`, after loading all
+        # custom nodes.
+        #
+        # We can still do some basic validation here - ensure the invoke method is defined and returns an instance
+        # of BaseInvocationOutput.
+
         # Validate the `invoke()` method is implemented
         if "invoke" in cls.__abstractmethods__:
             raise ValueError(f'Invocation "{invocation_type}" must implement the "invoke" method')
@@ -565,17 +676,13 @@ def invocation(
             )
 
         docstring = cls.__doc__
-        cls = create_model(
-            cls.__qualname__,
-            __base__=cls,
-            __module__=cls.__module__,
-            type=(invocation_type_annotation, invocation_type_field),
-        )
-        cls.__doc__ = docstring
+        new_class = create_model(cls.__qualname__, __base__=cls, __module__=cls.__module__, **fields)  # type: ignore
+        new_class.__doc__ = docstring
+        new_class._original_model_fields = original_model_fields
 
-        InvocationRegistry.register_invocation(cls)
+        InvocationRegistry.register_invocation(new_class)
 
-        return cls
+        return new_class
 
     return wrapper
 
@@ -600,23 +707,39 @@ def invocation_output(
 
         validate_fields(cls.model_fields, output_type)
 
+        fields: dict[str, tuple[Any, FieldInfo]] = {}
+
+        for field_name, field_info in cls.model_fields.items():
+            annotation = field_info.annotation
+            assert annotation is not None, f"{field_name} on invocation output {output_type} has no type annotation."
+            assert isinstance(field_info.json_schema_extra, dict), (
+                f"{field_name} on invocation output {output_type} has a non-dict json_schema_extra, did you forget to use InputField?"
+            )
+
+            cls._original_model_fields[field_name] = OriginalModelField(annotation=annotation, field_info=field_info)
+
+            if field_info.default is not PydanticUndefined and is_optional(annotation):
+                annotation = annotation | None
+            fields[field_name] = (annotation, field_info)
+
         # Add the output type to the model.
-        output_type_annotation = Literal[output_type]  # type: ignore
-        output_type_field = Field(
-            title="type", default=output_type, json_schema_extra={"field_kind": FieldKind.NodeAttribute}
+        output_type_annotation = Literal[output_type]
+
+        # Field() returns an instance of FieldInfo, but thanks to a pydantic implementation detail, it is _typed_ as Any.
+        # This cast makes the type annotation match the class's true type.
+        output_type_field_info = cast(
+            FieldInfo,
+            Field(title="type", default=output_type, json_schema_extra={"field_kind": FieldKind.NodeAttribute}),
         )
 
+        fields["type"] = (output_type_annotation, output_type_field_info)
+
         docstring = cls.__doc__
-        cls = create_model(
-            cls.__qualname__,
-            __base__=cls,
-            __module__=cls.__module__,
-            type=(output_type_annotation, output_type_field),
-        )
-        cls.__doc__ = docstring
+        new_class = create_model(cls.__qualname__, __base__=cls, __module__=cls.__module__, **fields)
+        new_class.__doc__ = docstring
 
-        InvocationRegistry.register_output(cls)
+        InvocationRegistry.register_output(new_class)
 
-        return cls
+        return new_class
 
     return wrapper

invokeai/app/invocations/batch.py

@@ -64,7 +64,6 @@ class ImageBatchInvocation(BaseBatchInvocation):
     """Create a batched generation, where the workflow is executed once for each image in the batch."""
 
     images: list[ImageField] = InputField(
-        default=[],
         min_length=1,
         description="The images to batch over",
     )
@@ -120,7 +119,6 @@ class StringBatchInvocation(BaseBatchInvocation):
     """Create a batched generation, where the workflow is executed once for each string in the batch."""
 
     strings: list[str] = InputField(
-        default=[],
         min_length=1,
         description="The strings to batch over",
     )
@@ -176,7 +174,6 @@ class IntegerBatchInvocation(BaseBatchInvocation):
     """Create a batched generation, where the workflow is executed once for each integer in the batch."""
 
     integers: list[int] = InputField(
-        default=[],
         min_length=1,
         description="The integers to batch over",
     )
@@ -230,7 +227,6 @@ class FloatBatchInvocation(BaseBatchInvocation):
     """Create a batched generation, where the workflow is executed once for each float in the batch."""
 
     floats: list[float] = InputField(
-        default=[],
         min_length=1,
         description="The floats to batch over",
     )

invokeai/app/invocations/compel.py

@@ -1,7 +1,7 @@
 from typing import Iterator, List, Optional, Tuple, Union, cast
 
 import torch
-from compel import Compel, ReturnedEmbeddingsType
+from compel import Compel, ReturnedEmbeddingsType, SplitLongTextMode
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
@@ -104,6 +104,7 @@ class CompelInvocation(BaseInvocation):
                 dtype_for_device_getter=TorchDevice.choose_torch_dtype,
                 truncate_long_prompts=False,
                 device=TorchDevice.choose_torch_device(),
+                split_long_text_mode=SplitLongTextMode.SENTENCES,
             )
 
             conjunction = Compel.parse_prompt_string(self.prompt)
@@ -113,6 +114,13 @@ class CompelInvocation(BaseInvocation):
 
             c, _options = compel.build_conditioning_tensor_for_conjunction(conjunction)
 
+        del compel
+        del patched_tokenizer
+        del tokenizer
+        del ti_manager
+        del text_encoder
+        del text_encoder_info
+
         c = c.detach().to("cpu")
 
         conditioning_data = ConditioningFieldData(conditionings=[BasicConditioningInfo(embeds=c)])
@@ -205,6 +213,7 @@ class SDXLPromptInvocationBase:
                 returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,  # TODO: clip skip
                 requires_pooled=get_pooled,
                 device=TorchDevice.choose_torch_device(),
+                split_long_text_mode=SplitLongTextMode.SENTENCES,
             )
 
             conjunction = Compel.parse_prompt_string(prompt)
@@ -220,7 +229,10 @@ class SDXLPromptInvocationBase:
             else:
                 c_pooled = None
 
+        del compel
+        del patched_tokenizer
         del tokenizer
+        del ti_manager
         del text_encoder
         del text_encoder_info
 

invokeai/app/invocations/composition-nodes.py

@@ -274,12 +274,12 @@ class InvokeAdjustImageHuePlusInvocation(BaseInvocation, WithMetadata, WithBoard
     title="Enhance Image",
     tags=["enhance", "image"],
     category="image",
-    version="1.2.0",
+    version="1.2.1",
 )
 class InvokeImageEnhanceInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Applies processing from PIL's ImageEnhance module. Originally created by @dwringer"""
 
-    image: ImageField = InputField(default=None, description="The image for which to apply processing")
+    image: ImageField = InputField(description="The image for which to apply processing")
     invert: bool = InputField(default=False, description="Whether to invert the image colors")
     color: float = InputField(ge=0, default=1.0, description="Color enhancement factor")
     contrast: float = InputField(ge=0, default=1.0, description="Contrast enhancement factor")

invokeai/app/invocations/controlnet.py

@@ -22,7 +22,11 @@ from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES, heuristic_resize
+from invokeai.app.util.controlnet_utils import (
+    CONTROLNET_MODE_VALUES,
+    CONTROLNET_RESIZE_VALUES,
+    heuristic_resize_fast,
+)
 from invokeai.backend.image_util.util import np_to_pil, pil_to_np
 
 
@@ -109,7 +113,7 @@ class ControlNetInvocation(BaseInvocation):
     title="Heuristic Resize",
     tags=["image, controlnet"],
     category="image",
-    version="1.0.1",
+    version="1.1.1",
     classification=Classification.Prototype,
 )
 class HeuristicResizeInvocation(BaseInvocation):
@@ -122,7 +126,7 @@ class HeuristicResizeInvocation(BaseInvocation):
     def invoke(self, context: InvocationContext) -> ImageOutput:
         image = context.images.get_pil(self.image.image_name, "RGB")
         np_img = pil_to_np(image)
-        np_resized = heuristic_resize(np_img, (self.width, self.height))
+        np_resized = heuristic_resize_fast(np_img, (self.width, self.height))
         resized = np_to_pil(np_resized)
         image_dto = context.images.save(image=resized)
         return ImageOutput.build(image_dto)

invokeai/app/invocations/create_gradient_mask.py

@@ -1,12 +1,14 @@
 from typing import Literal, Optional
 
+import cv2
 import numpy as np
 import torch
 import torchvision.transforms as T
-from PIL import Image, ImageFilter
+from PIL import Image
 from torchvision.transforms.functional import resize as tv_resize
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
 from invokeai.app.invocations.fields import (
     DenoiseMaskField,
     FieldDescriptions,
@@ -42,15 +44,13 @@ class GradientMaskOutput(BaseInvocationOutput):
     title="Create Gradient Mask",
     tags=["mask", "denoise"],
     category="latents",
-    version="1.2.0",
+    version="1.3.0",
 )
 class CreateGradientMaskInvocation(BaseInvocation):
-    """Creates mask for denoising model run."""
+    """Creates mask for denoising."""
 
-    mask: ImageField = InputField(default=None, description="Image which will be masked", ui_order=1)
-    edge_radius: int = InputField(
-        default=16, ge=0, description="How far to blur/expand the edges of the mask", ui_order=2
-    )
+    mask: ImageField = InputField(description="Image which will be masked", ui_order=1)
+    edge_radius: int = InputField(default=16, ge=0, description="How far to expand the edges of the mask", ui_order=2)
     coherence_mode: Literal["Gaussian Blur", "Box Blur", "Staged"] = InputField(default="Gaussian Blur", ui_order=3)
     minimum_denoise: float = InputField(
         default=0.0, ge=0, le=1, description="Minimum denoise level for the coherence region", ui_order=4
@@ -81,45 +81,110 @@ class CreateGradientMaskInvocation(BaseInvocation):
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> GradientMaskOutput:
         mask_image = context.images.get_pil(self.mask.image_name, mode="L")
+
+        # Resize the mask_image. Makes the filter 64x faster and doesn't hurt quality in latent scale anyway
+        mask_image = mask_image.resize(
+            (
+                mask_image.width // LATENT_SCALE_FACTOR,
+                mask_image.height // LATENT_SCALE_FACTOR,
+            ),
+            resample=Image.Resampling.BILINEAR,
+        )
+
+        mask_np_orig = np.array(mask_image, dtype=np.float32)
+
+        self.edge_radius = self.edge_radius // LATENT_SCALE_FACTOR  # scale the edge radius to match the mask size
+
         if self.edge_radius > 0:
+            mask_np = 255 - mask_np_orig  # invert so 0 is unmasked (higher values = higher denoise strength)
+            dilated_mask = mask_np.copy()
+
+            # Create kernel based on coherence mode
             if self.coherence_mode == "Box Blur":
-                blur_mask = mask_image.filter(ImageFilter.BoxBlur(self.edge_radius))
-            else:  # Gaussian Blur OR Staged
-                # Gaussian Blur uses standard deviation. 1/2 radius is a good approximation
-                blur_mask = mask_image.filter(ImageFilter.GaussianBlur(self.edge_radius / 2))
+                # Create a circular distance kernel that fades from center outward
+                kernel_size = self.edge_radius * 2 + 1
+                center = self.edge_radius
+                kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)
+                for i in range(kernel_size):
+                    for j in range(kernel_size):
+                        dist = np.sqrt((i - center) ** 2 + (j - center) ** 2)
+                        if dist <= self.edge_radius:
+                            kernel[i, j] = 1.0 - (dist / self.edge_radius)
+            else:  # Gaussian Blur or Staged
+                # Create a Gaussian kernel
+                kernel_size = self.edge_radius * 2 + 1
+                kernel = cv2.getGaussianKernel(
+                    kernel_size, self.edge_radius / 2.5
+                )  # 2.5 is a magic number (standard deviation capturing)
+                kernel = kernel * kernel.T  # Make 2D gaussian kernel
+                kernel = kernel / np.max(kernel)  # Normalize center to 1.0
 
-            blur_tensor: torch.Tensor = image_resized_to_grid_as_tensor(blur_mask, normalize=False)
+                # Ensure values outside radius are 0
+                center = self.edge_radius
+                for i in range(kernel_size):
+                    for j in range(kernel_size):
+                        dist = np.sqrt((i - center) ** 2 + (j - center) ** 2)
+                        if dist > self.edge_radius:
+                            kernel[i, j] = 0
 
-            # redistribute blur so that the original edges are 0 and blur outwards to 1
-            blur_tensor = (blur_tensor - 0.5) * 2
-            blur_tensor[blur_tensor < 0] = 0.0
+            # 2D max filter
+            mask_tensor = torch.tensor(mask_np)
+            kernel_tensor = torch.tensor(kernel)
+            dilated_mask = 255 - self.max_filter2D_torch(mask_tensor, kernel_tensor).cpu()
+            dilated_mask = dilated_mask.numpy()
 
-            threshold = 1 - self.minimum_denoise
+            threshold = (1 - self.minimum_denoise) * 255
 
             if self.coherence_mode == "Staged":
-                # wherever the blur_tensor is less than fully masked, convert it to threshold
-                blur_tensor = torch.where((blur_tensor < 1) & (blur_tensor > 0), threshold, blur_tensor)
-            else:
-                # wherever the blur_tensor is above threshold but less than 1, drop it to threshold
-                blur_tensor = torch.where((blur_tensor > threshold) & (blur_tensor < 1), threshold, blur_tensor)
+                # wherever expanded mask is darker than the original mask but original was above threshhold, set it to the threshold
+                # makes any expansion areas drop to threshhold. Raising minimum across the image happen outside of this if
+                threshold_mask = (dilated_mask < mask_np_orig) & (mask_np_orig > threshold)
+                dilated_mask = np.where(threshold_mask, threshold, mask_np_orig)
+
+            # wherever expanded mask is less than 255 but greater than threshold, drop it to threshold (minimum denoise)
+            threshold_mask = (dilated_mask > threshold) & (dilated_mask < 255)
+            dilated_mask = np.where(threshold_mask, threshold, dilated_mask)
 
         else:
-            blur_tensor: torch.Tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
+            dilated_mask = mask_np_orig.copy()
 
-        mask_name = context.tensors.save(tensor=blur_tensor.unsqueeze(1))
+        # convert to tensor
+        dilated_mask = np.clip(dilated_mask, 0, 255).astype(np.uint8)
+        mask_tensor = torch.tensor(dilated_mask, device=torch.device("cpu"))
 
-        # compute a [0, 1] mask from the blur_tensor
-        expanded_mask = torch.where((blur_tensor < 1), 0, 1)
-        expanded_mask_image = Image.fromarray((expanded_mask.squeeze(0).numpy() * 255).astype(np.uint8), mode="L")
+        # binary mask for compositing
+        expanded_mask = np.where((dilated_mask < 255), 0, 255)
+        expanded_mask_image = Image.fromarray(expanded_mask.astype(np.uint8), mode="L")
+        expanded_mask_image = expanded_mask_image.resize(
+            (
+                mask_image.width * LATENT_SCALE_FACTOR,
+                mask_image.height * LATENT_SCALE_FACTOR,
+            ),
+            resample=Image.Resampling.NEAREST,
+        )
         expanded_image_dto = context.images.save(expanded_mask_image)
 
+        # restore the original mask size
+        dilated_mask = Image.fromarray(dilated_mask.astype(np.uint8))
+        dilated_mask = dilated_mask.resize(
+            (
+                mask_image.width * LATENT_SCALE_FACTOR,
+                mask_image.height * LATENT_SCALE_FACTOR,
+            ),
+            resample=Image.Resampling.NEAREST,
+        )
+
+        # stack the mask as a tensor, repeating 4 times on dimmension 1
+        dilated_mask_tensor = image_resized_to_grid_as_tensor(dilated_mask, normalize=False)
+        mask_name = context.tensors.save(tensor=dilated_mask_tensor.unsqueeze(0))
+
         masked_latents_name = None
         if self.unet is not None and self.vae is not None and self.image is not None:
             # all three fields must be present at the same time
             main_model_config = context.models.get_config(self.unet.unet.key)
             assert isinstance(main_model_config, MainConfigBase)
             if main_model_config.variant is ModelVariantType.Inpaint:
-                mask = blur_tensor
+                mask = dilated_mask_tensor
                 vae_info: LoadedModel = context.models.load(self.vae.vae)
                 image = context.images.get_pil(self.image.image_name)
                 image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
@@ -137,3 +202,29 @@ class CreateGradientMaskInvocation(BaseInvocation):
             denoise_mask=DenoiseMaskField(mask_name=mask_name, masked_latents_name=masked_latents_name, gradient=True),
             expanded_mask_area=ImageField(image_name=expanded_image_dto.image_name),
         )
+
+    def max_filter2D_torch(self, image: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
+        """
+        This morphological operation is much faster in torch than numpy or opencv
+        For reasonable kernel sizes, the overhead of copying the data to the GPU is not worth it.
+        """
+        h, w = kernel.shape
+        pad_h, pad_w = h // 2, w // 2
+
+        padded = torch.nn.functional.pad(image, (pad_w, pad_w, pad_h, pad_h), mode="constant", value=0)
+        result = torch.zeros_like(image)
+
+        # This looks like it's inside out, but it does the same thing and is more efficient
+        for i in range(h):
+            for j in range(w):
+                weight = kernel[i, j]
+                if weight <= 0:
+                    continue
+
+                # Extract the region from padded tensor
+                region = padded[i : i + image.shape[0], j : j + image.shape[1]]
+
+                # Apply weight and update max
+                result = torch.maximum(result, region * weight)
+
+        return result

invokeai/app/invocations/denoise_latents.py

@@ -608,6 +608,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
                     end_step_percent=single_ip_adapter.end_step_percent,
                     ip_adapter_conditioning=IPAdapterConditioningInfo(image_prompt_embeds, uncond_image_prompt_embeds),
                     mask=mask,
+                    method=single_ip_adapter.method,
                 )
             )
 

invokeai/app/invocations/fields.py

@@ -61,6 +61,10 @@ class UIType(str, Enum, metaclass=MetaEnum):
     SigLipModel = "SigLipModelField"
     FluxReduxModel = "FluxReduxModelField"
     LlavaOnevisionModel = "LLaVAModelField"
+    Imagen3Model = "Imagen3ModelField"
+    Imagen4Model = "Imagen4ModelField"
+    ChatGPT4oModel = "ChatGPT4oModelField"
+    FluxKontextModel = "FluxKontextModelField"
     # endregion
 
     # region Misc Field Types
@@ -211,6 +215,7 @@ class FieldDescriptions:
     flux_redux_conditioning = "FLUX Redux conditioning tensor"
     vllm_model = "The VLLM model to use"
     flux_fill_conditioning = "FLUX Fill conditioning tensor"
+    flux_kontext_conditioning = "FLUX Kontext conditioning (reference image)"
 
 
 class ImageField(BaseModel):
@@ -287,6 +292,12 @@ class FluxFillConditioningField(BaseModel):
     mask: TensorField = Field(description="The FLUX Fill inpaint mask.")
 
 
+class FluxKontextConditioningField(BaseModel):
+    """A conditioning field for FLUX Kontext (reference image)."""
+
+    image: ImageField = Field(description="The Kontext reference image.")
+
+
 class SD3ConditioningField(BaseModel):
     """A conditioning tensor primitive value"""
 
@@ -398,8 +409,8 @@ class InputFieldJSONSchemaExtra(BaseModel):
     """
 
     input: Input
-    orig_required: bool
     field_kind: FieldKind
+    orig_required: bool = True
     default: Optional[Any] = None
     orig_default: Optional[Any] = None
     ui_hidden: bool = False
@@ -434,7 +445,7 @@ class WithWorkflow:
     workflow = None
 
     def __init_subclass__(cls) -> None:
-        logger.warn(
+        logger.warning(
             f"{cls.__module__.split('.')[0]}.{cls.__name__}: WithWorkflow is deprecated. Use `context.workflow` to access the workflow."
         )
         super().__init_subclass__()
@@ -496,7 +507,7 @@ def InputField(
     input: Input = Input.Any,
     ui_type: Optional[UIType] = None,
     ui_component: Optional[UIComponent] = None,
-    ui_hidden: bool = False,
+    ui_hidden: Optional[bool] = None,
     ui_order: Optional[int] = None,
     ui_choice_labels: Optional[dict[str, str]] = None,
 ) -> Any:
@@ -532,15 +543,20 @@ def InputField(
 
     json_schema_extra_ = InputFieldJSONSchemaExtra(
         input=input,
-        ui_type=ui_type,
-        ui_component=ui_component,
-        ui_hidden=ui_hidden,
-        ui_order=ui_order,
-        ui_choice_labels=ui_choice_labels,
         field_kind=FieldKind.Input,
-        orig_required=True,
     )
 
+    if ui_type is not None:
+        json_schema_extra_.ui_type = ui_type
+    if ui_component is not None:
+        json_schema_extra_.ui_component = ui_component
+    if ui_hidden is not None:
+        json_schema_extra_.ui_hidden = ui_hidden
+    if ui_order is not None:
+        json_schema_extra_.ui_order = ui_order
+    if ui_choice_labels is not None:
+        json_schema_extra_.ui_choice_labels = ui_choice_labels
+
     """
     There is a conflict between the typing of invocation definitions and the typing of an invocation's
     `invoke()` function.
@@ -570,7 +586,7 @@ def InputField(
 
     if default_factory is not _Unset and default_factory is not None:
         default = default_factory()
-        logger.warn('"default_factory" is not supported, calling it now to set "default"')
+        logger.warning('"default_factory" is not supported, calling it now to set "default"')
 
     # These are the args we may wish pass to the pydantic `Field()` function
     field_args = {
@@ -612,7 +628,7 @@ def InputField(
 
     return Field(
         **provided_args,
-        json_schema_extra=json_schema_extra_.model_dump(exclude_none=True),
+        json_schema_extra=json_schema_extra_.model_dump(exclude_unset=True),
     )
 
 

invokeai/app/invocations/flux_denoise.py

@@ -16,13 +16,12 @@ from invokeai.app.invocations.fields import (
     FieldDescriptions,
     FluxConditioningField,
     FluxFillConditioningField,
+    FluxKontextConditioningField,
     FluxReduxConditioningField,
     ImageField,
     Input,
     InputField,
     LatentsField,
-    WithBoard,
-    WithMetadata,
 )
 from invokeai.app.invocations.flux_controlnet import FluxControlNetField
 from invokeai.app.invocations.flux_vae_encode import FluxVaeEncodeInvocation
@@ -34,6 +33,7 @@ from invokeai.backend.flux.controlnet.instantx_controlnet_flux import InstantXCo
 from invokeai.backend.flux.controlnet.xlabs_controlnet_flux import XLabsControlNetFlux
 from invokeai.backend.flux.denoise import denoise
 from invokeai.backend.flux.extensions.instantx_controlnet_extension import InstantXControlNetExtension
+from invokeai.backend.flux.extensions.kontext_extension import KontextExtension
 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
@@ -63,9 +63,9 @@ from invokeai.backend.util.devices import TorchDevice
     title="FLUX Denoise",
     tags=["image", "flux"],
     category="image",
-    version="3.3.0",
+    version="4.1.0",
 )
-class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
+class FluxDenoiseInvocation(BaseInvocation):
     """Run denoising process with a FLUX transformer model."""
 
     # If latents is provided, this means we are doing image-to-image.
@@ -145,11 +145,20 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         description=FieldDescriptions.vae,
         input=Input.Connection,
     )
+    # This node accepts a images for features like FLUX Fill, ControlNet, and Kontext, but needs to operate on them in
+    # latent space. We'll run the VAE to encode them in this node instead of requiring the user to run the VAE in
+    # upstream nodes.
 
     ip_adapter: IPAdapterField | list[IPAdapterField] | None = InputField(
         description=FieldDescriptions.ip_adapter, title="IP-Adapter", default=None, input=Input.Connection
     )
 
+    kontext_conditioning: FluxKontextConditioningField | list[FluxKontextConditioningField] | None = InputField(
+        default=None,
+        description="FLUX Kontext conditioning (reference image).",
+        input=Input.Connection,
+    )
+
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         latents = self._run_diffusion(context)
@@ -376,6 +385,29 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 dtype=inference_dtype,
             )
 
+            kontext_extension = None
+            if self.kontext_conditioning:
+                if not self.controlnet_vae:
+                    raise ValueError("A VAE (e.g., controlnet_vae) must be provided to use Kontext conditioning.")
+
+                kontext_extension = KontextExtension(
+                    context=context,
+                    kontext_conditioning=self.kontext_conditioning
+                    if isinstance(self.kontext_conditioning, list)
+                    else [self.kontext_conditioning],
+                    vae_field=self.controlnet_vae,
+                    device=TorchDevice.choose_torch_device(),
+                    dtype=inference_dtype,
+                )
+
+            # Prepare Kontext conditioning if provided
+            img_cond_seq = None
+            img_cond_seq_ids = None
+            if kontext_extension is not None:
+                # Ensure batch sizes match
+                kontext_extension.ensure_batch_size(x.shape[0])
+                img_cond_seq, img_cond_seq_ids = kontext_extension.kontext_latents, kontext_extension.kontext_ids
+
             x = denoise(
                 model=transformer,
                 img=x,
@@ -391,6 +423,8 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 pos_ip_adapter_extensions=pos_ip_adapter_extensions,
                 neg_ip_adapter_extensions=neg_ip_adapter_extensions,
                 img_cond=img_cond,
+                img_cond_seq=img_cond_seq,
+                img_cond_seq_ids=img_cond_seq_ids,
             )
 
         x = unpack(x.float(), self.height, self.width)
@@ -865,7 +899,10 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
 
     def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
         def step_callback(state: PipelineIntermediateState) -> None:
-            state.latents = unpack(state.latents.float(), self.height, self.width).squeeze()
+            # The denoise function now handles Kontext conditioning correctly,
+            # so we don't need to slice the latents here
+            latents = state.latents.float()
+            state.latents = unpack(latents, self.height, self.width).squeeze()
             context.util.flux_step_callback(state)
 
         return step_callback

invokeai/app/invocations/flux_kontext.py

@@ -0,0 +1,40 @@
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    FluxKontextConditioningField,
+    InputField,
+    OutputField,
+)
+from invokeai.app.invocations.primitives import ImageField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+
+@invocation_output("flux_kontext_output")
+class FluxKontextOutput(BaseInvocationOutput):
+    """The conditioning output of a FLUX Kontext invocation."""
+
+    kontext_cond: FluxKontextConditioningField = OutputField(
+        description=FieldDescriptions.flux_kontext_conditioning, title="Kontext Conditioning"
+    )
+
+
+@invocation(
+    "flux_kontext",
+    title="Kontext Conditioning - FLUX",
+    tags=["conditioning", "kontext", "flux"],
+    category="conditioning",
+    version="1.0.0",
+)
+class FluxKontextInvocation(BaseInvocation):
+    """Prepares a reference image for FLUX Kontext conditioning."""
+
+    image: ImageField = InputField(description="The Kontext reference image.")
+
+    def invoke(self, context: InvocationContext) -> FluxKontextOutput:
+        """Packages the provided image into a Kontext conditioning field."""
+        return FluxKontextOutput(kontext_cond=FluxKontextConditioningField(image=self.image))

invokeai/app/invocations/flux_text_encoder.py

@@ -1,5 +1,5 @@
 from contextlib import ExitStack
-from typing import Iterator, Literal, Optional, Tuple
+from typing import Iterator, Literal, Optional, Tuple, Union
 
 import torch
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer, T5TokenizerFast
@@ -111,6 +111,9 @@ class FluxTextEncoderInvocation(BaseInvocation):
 
             t5_encoder = HFEncoder(t5_text_encoder, t5_tokenizer, False, self.t5_max_seq_len)
 
+            if context.config.get().log_tokenization:
+                self._log_t5_tokenization(context, t5_tokenizer)
+
             context.util.signal_progress("Running T5 encoder")
             prompt_embeds = t5_encoder(prompt)
 
@@ -151,6 +154,9 @@ class FluxTextEncoderInvocation(BaseInvocation):
 
             clip_encoder = HFEncoder(clip_text_encoder, clip_tokenizer, True, 77)
 
+            if context.config.get().log_tokenization:
+                self._log_clip_tokenization(context, clip_tokenizer)
+
             context.util.signal_progress("Running CLIP encoder")
             pooled_prompt_embeds = clip_encoder(prompt)
 
@@ -170,3 +176,88 @@ class FluxTextEncoderInvocation(BaseInvocation):
             assert isinstance(lora_info.model, ModelPatchRaw)
             yield (lora_info.model, lora.weight)
             del lora_info
+
+    def _log_t5_tokenization(
+        self,
+        context: InvocationContext,
+        tokenizer: Union[T5Tokenizer, T5TokenizerFast],
+    ) -> None:
+        """Logs the tokenization of a prompt for a T5-based model like FLUX."""
+
+        # Tokenize the prompt using the same parameters as the model's text encoder.
+        # T5 tokenizers add an EOS token (</s>) and then pad to max_length.
+        tokenized_output = tokenizer(
+            self.prompt,
+            padding="max_length",
+            max_length=self.t5_max_seq_len,
+            truncation=True,
+            add_special_tokens=True,  # This is important for T5 to add the EOS token.
+            return_tensors="pt",
+        )
+
+        input_ids = tokenized_output.input_ids[0]
+        tokens = tokenizer.convert_ids_to_tokens(input_ids)
+
+        # The T5 tokenizer uses a space-like character ' ' (U+2581) to denote spaces.
+        # We'll replace it with a regular space for readability.
+        tokens = [t.replace("\u2581", " ") for t in tokens]
+
+        tokenized_str = ""
+        used_tokens = 0
+        for token in tokens:
+            if token == tokenizer.eos_token:
+                tokenized_str += f"\x1b[0;31m{token}\x1b[0m"  # Red for EOS
+                used_tokens += 1
+            elif token == tokenizer.pad_token:
+                # tokenized_str += f"\x1b[0;34m{token}\x1b[0m"  # Blue for PAD
+                continue
+            else:
+                color = (used_tokens % 6) + 1  # Cycle through 6 colors
+                tokenized_str += f"\x1b[0;3{color}m{token}\x1b[0m"
+                used_tokens += 1
+
+        context.logger.info(f">> [T5 TOKENLOG] Tokens ({used_tokens}/{self.t5_max_seq_len}):")
+        context.logger.info(f"{tokenized_str}\x1b[0m")
+
+    def _log_clip_tokenization(
+        self,
+        context: InvocationContext,
+        tokenizer: CLIPTokenizer,
+    ) -> None:
+        """Logs the tokenization of a prompt for a CLIP-based model."""
+        max_length = tokenizer.model_max_length
+
+        tokenized_output = tokenizer(
+            self.prompt,
+            padding="max_length",
+            max_length=max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        input_ids = tokenized_output.input_ids[0]
+        attention_mask = tokenized_output.attention_mask[0]
+        tokens = tokenizer.convert_ids_to_tokens(input_ids)
+
+        # The CLIP tokenizer uses '</w>' to denote spaces.
+        # We'll replace it with a regular space for readability.
+        tokens = [t.replace("</w>", " ") for t in tokens]
+
+        tokenized_str = ""
+        used_tokens = 0
+        for i, token in enumerate(tokens):
+            if attention_mask[i] == 0:
+                # Do not log padding tokens.
+                continue
+
+            if token == tokenizer.bos_token:
+                tokenized_str += f"\x1b[0;32m{token}\x1b[0m"  # Green for BOS
+            elif token == tokenizer.eos_token:
+                tokenized_str += f"\x1b[0;31m{token}\x1b[0m"  # Red for EOS
+            else:
+                color = (used_tokens % 6) + 1  # Cycle through 6 colors
+                tokenized_str += f"\x1b[0;3{color}m{token}\x1b[0m"
+            used_tokens += 1
+
+        context.logger.info(f">> [CLIP TOKENLOG] Tokens ({used_tokens}/{max_length}):")
+        context.logger.info(f"{tokenized_str}\x1b[0m")

invokeai/app/invocations/ideal_size.py

@@ -21,14 +21,14 @@ class IdealSizeOutput(BaseInvocationOutput):
     "ideal_size",
     title="Ideal Size - SD1.5, SDXL",
     tags=["latents", "math", "ideal_size"],
-    version="1.0.5",
+    version="1.0.6",
 )
 class IdealSizeInvocation(BaseInvocation):
     """Calculates the ideal size for generation to avoid duplication"""
 
     width: int = InputField(default=1024, description="Final image width")
     height: int = InputField(default=576, description="Final image height")
-    unet: UNetField = InputField(default=None, description=FieldDescriptions.unet)
+    unet: UNetField = InputField(description=FieldDescriptions.unet)
     multiplier: float = InputField(
         default=1.0,
         description="Amount to multiply the model's dimensions by when calculating the ideal size (may result in "

invokeai/app/invocations/image.py

@@ -975,13 +975,13 @@ class SaveImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     title="Canvas Paste Back",
     tags=["image", "combine"],
     category="image",
-    version="1.0.0",
+    version="1.0.1",
 )
 class CanvasPasteBackInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Combines two images by using the mask provided. Intended for use on the Unified Canvas."""
 
     source_image: ImageField = InputField(description="The source image")
-    target_image: ImageField = InputField(default=None, description="The target image")
+    target_image: ImageField = InputField(description="The target image")
     mask: ImageField = InputField(
         description="The mask to use when pasting",
     )
@@ -1218,12 +1218,15 @@ class ApplyMaskToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     title="Add Image Noise",
     tags=["image", "noise"],
     category="image",
-    version="1.0.1",
+    version="1.1.0",
 )
 class ImageNoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Add noise to an image"""
 
     image: ImageField = InputField(description="The image to add noise to")
+    mask: Optional[ImageField] = InputField(
+        default=None, description="Optional mask determining where to apply noise (black=noise, white=no noise)"
+    )
     seed: int = InputField(
         default=0,
         ge=0,
@@ -1267,12 +1270,27 @@ class ImageNoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         noise = Image.fromarray(noise.astype(numpy.uint8), mode="RGB").resize(
             (image.width, image.height), Image.Resampling.NEAREST
         )
+
+        # Create a noisy version of the input image
         noisy_image = Image.blend(image.convert("RGB"), noise, self.amount).convert("RGBA")
 
-        # Paste back the alpha channel
-        noisy_image.putalpha(alpha)
+        # Apply mask if provided
+        if self.mask is not None:
+            mask_image = context.images.get_pil(self.mask.image_name, mode="L")
 
-        image_dto = context.images.save(image=noisy_image)
+            if mask_image.size != image.size:
+                mask_image = mask_image.resize(image.size, Image.Resampling.LANCZOS)
+
+            result_image = image.copy()
+            mask_image = ImageOps.invert(mask_image)
+            result_image.paste(noisy_image, (0, 0), mask=mask_image)
+        else:
+            result_image = noisy_image
+
+        # Paste back the alpha channel from the original image
+        result_image.putalpha(alpha)
+
+        image_dto = context.images.save(image=result_image)
 
         return ImageOutput.build(image_dto)
 
@@ -1329,3 +1347,96 @@ class PasteImageIntoBoundingBoxInvocation(BaseInvocation, WithMetadata, WithBoar
 
         image_dto = context.images.save(image=target_image)
         return ImageOutput.build(image_dto)
+
+
+@invocation(
+    "flux_kontext_image_prep",
+    title="FLUX Kontext Image Prep",
+    tags=["image", "concatenate", "flux", "kontext"],
+    category="image",
+    version="1.0.0",
+)
+class FluxKontextConcatenateImagesInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Prepares an image or images for use with FLUX Kontext. The first/single image is resized to the nearest
+    preferred Kontext resolution. All other images are concatenated horizontally, maintaining their aspect ratio."""
+
+    images: list[ImageField] = InputField(
+        description="The images to concatenate",
+        min_length=1,
+        max_length=10,
+    )
+
+    use_preferred_resolution: bool = InputField(
+        default=True, description="Use FLUX preferred resolutions for the first image"
+    )
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        from invokeai.backend.flux.util import PREFERED_KONTEXT_RESOLUTIONS
+
+        # Step 1: Load all images
+        pil_images = []
+        for image_field in self.images:
+            image = context.images.get_pil(image_field.image_name, mode="RGBA")
+            pil_images.append(image)
+
+        # Step 2: Determine target resolution for the first image
+        first_image = pil_images[0]
+        width, height = first_image.size
+
+        if self.use_preferred_resolution:
+            aspect_ratio = width / height
+
+            # Find the closest preferred resolution for the first image
+            _, target_width, target_height = min(
+                ((abs(aspect_ratio - w / h), w, h) for w, h in PREFERED_KONTEXT_RESOLUTIONS), key=lambda x: x[0]
+            )
+
+            # Apply BFL's scaling formula
+            scaled_height = 2 * int(target_height / 16)
+            final_height = 8 * scaled_height  # This will be consistent for all images
+            scaled_width = 2 * int(target_width / 16)
+            first_width = 8 * scaled_width
+        else:
+            # Use original dimensions of first image, ensuring divisibility by 16
+            final_height = 16 * (height // 16)
+            first_width = 16 * (width // 16)
+            # Ensure minimum dimensions
+            if final_height < 16:
+                final_height = 16
+            if first_width < 16:
+                first_width = 16
+
+        # Step 3: Process and resize all images with consistent height
+        processed_images = []
+        total_width = 0
+
+        for i, image in enumerate(pil_images):
+            if i == 0:
+                # First image uses the calculated dimensions
+                final_width = first_width
+            else:
+                # Subsequent images maintain aspect ratio with the same height
+                img_aspect_ratio = image.width / image.height
+                # Calculate width that maintains aspect ratio at the target height
+                calculated_width = int(final_height * img_aspect_ratio)
+                # Ensure width is divisible by 16 for proper VAE encoding
+                final_width = 16 * (calculated_width // 16)
+                # Ensure minimum width
+                if final_width < 16:
+                    final_width = 16
+
+            # Resize image to calculated dimensions
+            resized_image = image.resize((final_width, final_height), Image.Resampling.LANCZOS)
+            processed_images.append(resized_image)
+            total_width += final_width
+
+        # Step 4: Concatenate images horizontally
+        concatenated_image = Image.new("RGB", (total_width, final_height))
+        x_offset = 0
+        for img in processed_images:
+            concatenated_image.paste(img, (x_offset, 0))
+            x_offset += img.width
+
+        # Save the concatenated image
+        image_dto = context.images.save(image=concatenated_image)
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/ip_adapter.py

@@ -31,6 +31,7 @@ class IPAdapterField(BaseModel):
     image_encoder_model: ModelIdentifierField = Field(description="The name of the CLIP image encoder model.")
     weight: Union[float, List[float]] = Field(default=1, description="The weight given to the IP-Adapter.")
     target_blocks: List[str] = Field(default=[], description="The IP Adapter blocks to apply")
+    method: str = Field(default="full", description="Weight apply method")
     begin_step_percent: float = Field(
         default=0, ge=0, le=1, description="When the IP-Adapter is first applied (% of total steps)"
     )
@@ -94,7 +95,7 @@ class IPAdapterInvocation(BaseInvocation):
     weight: Union[float, List[float]] = InputField(
         default=1, description="The weight given to the IP-Adapter", title="Weight"
     )
-    method: Literal["full", "style", "composition"] = InputField(
+    method: Literal["full", "style", "composition", "style_strong", "style_precise"] = InputField(
         default="full", description="The method to apply the IP-Adapter"
     )
     begin_step_percent: float = InputField(
@@ -147,6 +148,38 @@ class IPAdapterInvocation(BaseInvocation):
                 target_blocks = ["down_blocks.2.attentions.1"]
             else:
                 raise ValueError(f"Unsupported IP-Adapter base type: '{ip_adapter_info.base}'.")
+        elif self.method == "style_precise":
+            if ip_adapter_info.base == "sd-1":
+                target_blocks = ["up_blocks.1", "down_blocks.2", "mid_block"]
+            elif ip_adapter_info.base == "sdxl":
+                target_blocks = ["up_blocks.0.attentions.1", "down_blocks.2.attentions.1"]
+            else:
+                raise ValueError(f"Unsupported IP-Adapter base type: '{ip_adapter_info.base}'.")
+        elif self.method == "style_strong":
+            if ip_adapter_info.base == "sd-1":
+                target_blocks = ["up_blocks.0", "up_blocks.1", "up_blocks.2", "down_blocks.0", "down_blocks.1"]
+            elif ip_adapter_info.base == "sdxl":
+                target_blocks = [
+                    "up_blocks.0.attentions.1",
+                    "up_blocks.1.attentions.1",
+                    "up_blocks.2.attentions.1",
+                    "up_blocks.0.attentions.2",
+                    "up_blocks.1.attentions.2",
+                    "up_blocks.2.attentions.2",
+                    "up_blocks.0.attentions.0",
+                    "up_blocks.1.attentions.0",
+                    "up_blocks.2.attentions.0",
+                    "down_blocks.0.attentions.0",
+                    "down_blocks.0.attentions.1",
+                    "down_blocks.0.attentions.2",
+                    "down_blocks.1.attentions.0",
+                    "down_blocks.1.attentions.1",
+                    "down_blocks.1.attentions.2",
+                    "down_blocks.2.attentions.0",
+                    "down_blocks.2.attentions.2",
+                ]
+            else:
+                raise ValueError(f"Unsupported IP-Adapter base type: '{ip_adapter_info.base}'.")
         elif self.method == "full":
             target_blocks = ["block"]
         else:
@@ -162,6 +195,7 @@ class IPAdapterInvocation(BaseInvocation):
                 begin_step_percent=self.begin_step_percent,
                 end_step_percent=self.end_step_percent,
                 mask=self.mask,
+                method=self.method,
             ),
         )
 

invokeai/app/invocations/metadata.py

@@ -42,7 +42,9 @@ class IPAdapterMetadataField(BaseModel):
     image: ImageField = Field(description="The IP-Adapter image prompt.")
     ip_adapter_model: ModelIdentifierField = Field(description="The IP-Adapter model.")
     clip_vision_model: Literal["ViT-L", "ViT-H", "ViT-G"] = Field(description="The CLIP Vision model")
-    method: Literal["full", "style", "composition"] = Field(description="Method to apply IP Weights with")
+    method: Literal["full", "style", "composition", "style_strong", "style_precise"] = Field(
+        description="Method to apply IP Weights with"
+    )
     weight: Union[float, list[float]] = Field(description="The weight given to the IP-Adapter")
     begin_step_percent: float = Field(description="When the IP-Adapter is first applied (% of total steps)")
     end_step_percent: float = Field(description="When the IP-Adapter is last applied (% of total steps)")

invokeai/app/invocations/primitives.py

@@ -430,6 +430,15 @@ class FluxConditioningOutput(BaseInvocationOutput):
         return cls(conditioning=FluxConditioningField(conditioning_name=conditioning_name))
 
 
+@invocation_output("flux_conditioning_collection_output")
+class FluxConditioningCollectionOutput(BaseInvocationOutput):
+    """Base class for nodes that output a collection of conditioning tensors"""
+
+    collection: list[FluxConditioningField] = OutputField(
+        description="The output conditioning tensors",
+    )
+
+
 @invocation_output("sd3_conditioning_output")
 class SD3ConditioningOutput(BaseInvocationOutput):
     """Base class for nodes that output a single SD3 conditioning tensor"""

invokeai/app/invocations/segment_anything.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 from PIL import Image
 from pydantic import BaseModel, Field
-from transformers import AutoModelForMaskGeneration, AutoProcessor
+from transformers import AutoProcessor
 from transformers.models.sam import SamModel
 from transformers.models.sam.processing_sam import SamProcessor
 
@@ -104,14 +104,13 @@ class SegmentAnythingInvocation(BaseInvocation):
 
     @staticmethod
     def _load_sam_model(model_path: Path):
-        sam_model = AutoModelForMaskGeneration.from_pretrained(
+        sam_model = SamModel.from_pretrained(
             model_path,
             local_files_only=True,
             # TODO(ryand): Setting the torch_dtype here doesn't work. Investigate whether fp16 is supported by the
             # model, and figure out how to make it work in the pipeline.
             # torch_dtype=TorchDevice.choose_torch_dtype(),
         )
-        assert isinstance(sam_model, SamModel)
 
         sam_processor = AutoProcessor.from_pretrained(model_path, local_files_only=True)
         assert isinstance(sam_processor, SamProcessor)

invokeai/app/run_app.py

@@ -1,12 +1,3 @@
-import uvicorn
-
-from invokeai.app.invocations.load_custom_nodes import load_custom_nodes
-from invokeai.app.services.config.config_default import get_config
-from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
-from invokeai.backend.util.logging import InvokeAILogger
-from invokeai.frontend.cli.arg_parser import InvokeAIArgs
-
-
 def get_app():
     """Import the app and event loop. We wrap this in a function to more explicitly control when it happens, because
     importing from api_app does a bunch of stuff - it's more like calling a function than importing a module.
@@ -18,9 +9,18 @@ def get_app():
 
 def run_app() -> None:
     """The main entrypoint for the app."""
-    # Parse the CLI arguments.
+    from invokeai.frontend.cli.arg_parser import InvokeAIArgs
+
+    # Parse the CLI arguments before doing anything else, which ensures CLI args correctly override settings from other
+    # sources like `invokeai.yaml` or env vars.
     InvokeAIArgs.parse_args()
 
+    import uvicorn
+
+    from invokeai.app.services.config.config_default import get_config
+    from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator
+    from invokeai.backend.util.logging import InvokeAILogger
+
     # Load config.
     app_config = get_config()
 
@@ -32,6 +32,8 @@ def run_app() -> None:
         configure_torch_cuda_allocator(app_config.pytorch_cuda_alloc_conf, logger)
 
     # This import must happen after configure_torch_cuda_allocator() is called, because the module imports torch.
+    from invokeai.app.invocations.baseinvocation import InvocationRegistry
+    from invokeai.app.invocations.load_custom_nodes import load_custom_nodes
     from invokeai.backend.util.devices import TorchDevice
 
     torch_device_name = TorchDevice.get_torch_device_name()
@@ -66,6 +68,15 @@ def run_app() -> None:
     # core nodes have been imported so that we can catch when a custom node clobbers a core node.
     load_custom_nodes(custom_nodes_path=app_config.custom_nodes_path, logger=logger)
 
+    # Check all invocations and ensure their outputs are registered.
+    for invocation in InvocationRegistry.get_invocation_classes():
+        invocation_type = invocation.get_type()
+        output_annotation = invocation.get_output_annotation()
+        if output_annotation not in InvocationRegistry.get_output_classes():
+            logger.warning(
+                f'Invocation "{invocation_type}" has unregistered output class "{output_annotation.__name__}"'
+            )
+
     if app_config.dev_reload:
         # load_custom_nodes seems to bypass jurrigged's import sniffer, so be sure to call it *after* they're already
         # imported.

invokeai/app/services/board_image_records/board_image_records_sqlite.py

@@ -14,15 +14,14 @@ from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
 class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
-        self._conn = db.conn
+        self._db = db
 
     def add_image_to_board(
         self,
         board_id: str,
         image_name: str,
     ) -> None:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 INSERT INTO board_images (board_id, image_name)
@@ -31,17 +30,12 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
                 """,
                 (board_id, image_name, board_id),
             )
-            self._conn.commit()
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise e
 
     def remove_image_from_board(
         self,
         image_name: str,
     ) -> None:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 DELETE FROM board_images
@@ -49,10 +43,6 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
                 """,
                 (image_name,),
             )
-            self._conn.commit()
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise e
 
     def get_images_for_board(
         self,
@@ -60,27 +50,26 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
         offset: int = 0,
         limit: int = 10,
     ) -> OffsetPaginatedResults[ImageRecord]:
-        # TODO: this isn't paginated yet?
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT images.*
-            FROM board_images
-            INNER JOIN images ON board_images.image_name = images.image_name
-            WHERE board_images.board_id = ?
-            ORDER BY board_images.updated_at DESC;
-            """,
-            (board_id,),
-        )
-        result = cast(list[sqlite3.Row], cursor.fetchall())
-        images = [deserialize_image_record(dict(r)) for r in result]
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT images.*
+                FROM board_images
+                INNER JOIN images ON board_images.image_name = images.image_name
+                WHERE board_images.board_id = ?
+                ORDER BY board_images.updated_at DESC;
+                """,
+                (board_id,),
+            )
+            result = cast(list[sqlite3.Row], cursor.fetchall())
+            images = [deserialize_image_record(dict(r)) for r in result]
 
-        cursor.execute(
-            """--sql
-            SELECT COUNT(*) FROM images WHERE 1=1;
-            """
-        )
-        count = cast(int, cursor.fetchone()[0])
+            cursor.execute(
+                """--sql
+                SELECT COUNT(*) FROM images WHERE 1=1;
+                """
+            )
+            count = cast(int, cursor.fetchone()[0])
 
         return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=count)
 
@@ -90,47 +79,55 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
         categories: list[ImageCategory] | None,
         is_intermediate: bool | None,
     ) -> list[str]:
-        params: list[str | bool] = []
+        with self._db.transaction() as cursor:
+            params: list[str | bool] = []
 
-        # Base query is a join between images and board_images
-        stmt = """
-                SELECT images.image_name
-                FROM images
-                LEFT JOIN board_images ON board_images.image_name = images.image_name
-                WHERE 1=1
-                AND board_images.board_id = ?
-                """
-        params.append(board_id)
+            # Base query is a join between images and board_images
+            stmt = """
+                    SELECT images.image_name
+                    FROM images
+                    LEFT JOIN board_images ON board_images.image_name = images.image_name
+                    WHERE 1=1
+                    """
 
-        # Add the category filter
-        if categories is not None:
-            # Convert the enum values to unique list of strings
-            category_strings = [c.value for c in set(categories)]
-            # Create the correct length of placeholders
-            placeholders = ",".join("?" * len(category_strings))
-            stmt += f"""--sql
-                AND images.image_category IN ( {placeholders} )
-                """
+            # Handle board_id filter
+            if board_id == "none":
+                stmt += """--sql
+                    AND board_images.board_id IS NULL
+                    """
+            else:
+                stmt += """--sql
+                    AND board_images.board_id = ?
+                    """
+                params.append(board_id)
 
-            # Unpack the included categories into the query params
-            for c in category_strings:
-                params.append(c)
+            # Add the category filter
+            if categories is not None:
+                # Convert the enum values to unique list of strings
+                category_strings = [c.value for c in set(categories)]
+                # Create the correct length of placeholders
+                placeholders = ",".join("?" * len(category_strings))
+                stmt += f"""--sql
+                    AND images.image_category IN ( {placeholders} )
+                    """
 
-        # Add the is_intermediate filter
-        if is_intermediate is not None:
-            stmt += """--sql
-                AND images.is_intermediate = ?
-                """
-            params.append(is_intermediate)
+                # Unpack the included categories into the query params
+                for c in category_strings:
+                    params.append(c)
 
-        # Put a ring on it
-        stmt += ";"
+            # Add the is_intermediate filter
+            if is_intermediate is not None:
+                stmt += """--sql
+                    AND images.is_intermediate = ?
+                    """
+                params.append(is_intermediate)
 
-        # Execute the query
-        cursor = self._conn.cursor()
-        cursor.execute(stmt, params)
+            # Put a ring on it
+            stmt += ";"
 
-        result = cast(list[sqlite3.Row], cursor.fetchall())
+            cursor.execute(stmt, params)
+
+            result = cast(list[sqlite3.Row], cursor.fetchall())
         image_names = [r[0] for r in result]
         return image_names
 
@@ -138,31 +135,31 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
         self,
         image_name: str,
     ) -> Optional[str]:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-                SELECT board_id
-                FROM board_images
-                WHERE image_name = ?;
-                """,
-            (image_name,),
-        )
-        result = cursor.fetchone()
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                    SELECT board_id
+                    FROM board_images
+                    WHERE image_name = ?;
+                    """,
+                (image_name,),
+            )
+            result = cursor.fetchone()
         if result is None:
             return None
         return cast(str, result[0])
 
     def get_image_count_for_board(self, board_id: str) -> int:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-                SELECT COUNT(*)
-                FROM board_images
-                INNER JOIN images ON board_images.image_name = images.image_name
-                WHERE images.is_intermediate = FALSE
-                AND board_images.board_id = ?;
-                """,
-            (board_id,),
-        )
-        count = cast(int, cursor.fetchone()[0])
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                    SELECT COUNT(*)
+                    FROM board_images
+                    INNER JOIN images ON board_images.image_name = images.image_name
+                    WHERE images.is_intermediate = FALSE
+                    AND board_images.board_id = ?;
+                    """,
+                (board_id,),
+            )
+            count = cast(int, cursor.fetchone()[0])
         return count

invokeai/app/services/board_records/board_records_sqlite.py

@@ -20,61 +20,57 @@ from invokeai.app.util.misc import uuid_string
 class SqliteBoardRecordStorage(BoardRecordStorageBase):
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
-        self._conn = db.conn
+        self._db = db
 
     def delete(self, board_id: str) -> None:
-        try:
-            cursor = self._conn.cursor()
-            cursor.execute(
-                """--sql
-                DELETE FROM boards
-                WHERE board_id = ?;
-                """,
-                (board_id,),
-            )
-            self._conn.commit()
-        except Exception as e:
-            self._conn.rollback()
-            raise BoardRecordDeleteException from e
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    """--sql
+                    DELETE FROM boards
+                    WHERE board_id = ?;
+                    """,
+                    (board_id,),
+                )
+            except Exception as e:
+                raise BoardRecordDeleteException from e
 
     def save(
         self,
         board_name: str,
     ) -> BoardRecord:
-        try:
-            board_id = uuid_string()
-            cursor = self._conn.cursor()
-            cursor.execute(
-                """--sql
-                INSERT OR IGNORE INTO boards (board_id, board_name)
-                VALUES (?, ?);
-                """,
-                (board_id, board_name),
-            )
-            self._conn.commit()
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise BoardRecordSaveException from e
+        with self._db.transaction() as cursor:
+            try:
+                board_id = uuid_string()
+                cursor.execute(
+                    """--sql
+                    INSERT OR IGNORE INTO boards (board_id, board_name)
+                    VALUES (?, ?);
+                    """,
+                    (board_id, board_name),
+                )
+            except sqlite3.Error as e:
+                raise BoardRecordSaveException from e
         return self.get(board_id)
 
     def get(
         self,
         board_id: str,
     ) -> BoardRecord:
-        try:
-            cursor = self._conn.cursor()
-            cursor.execute(
-                """--sql
-                SELECT *
-                FROM boards
-                WHERE board_id = ?;
-                """,
-                (board_id,),
-            )
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    """--sql
+                    SELECT *
+                    FROM boards
+                    WHERE board_id = ?;
+                    """,
+                    (board_id,),
+                )
 
-            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
-        except sqlite3.Error as e:
-            raise BoardRecordNotFoundException from e
+                result = cast(Union[sqlite3.Row, None], cursor.fetchone())
+            except sqlite3.Error as e:
+                raise BoardRecordNotFoundException from e
         if result is None:
             raise BoardRecordNotFoundException
         return BoardRecord(**dict(result))
@@ -84,45 +80,43 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
         board_id: str,
         changes: BoardChanges,
     ) -> BoardRecord:
-        try:
-            cursor = self._conn.cursor()
-            # Change the name of a board
-            if changes.board_name is not None:
-                cursor.execute(
-                    """--sql
-                    UPDATE boards
-                    SET board_name = ?
-                    WHERE board_id = ?;
-                    """,
-                    (changes.board_name, board_id),
-                )
+        with self._db.transaction() as cursor:
+            try:
+                # Change the name of a board
+                if changes.board_name is not None:
+                    cursor.execute(
+                        """--sql
+                        UPDATE boards
+                        SET board_name = ?
+                        WHERE board_id = ?;
+                        """,
+                        (changes.board_name, board_id),
+                    )
 
-            # Change the cover image of a board
-            if changes.cover_image_name is not None:
-                cursor.execute(
-                    """--sql
-                    UPDATE boards
-                    SET cover_image_name = ?
-                    WHERE board_id = ?;
-                    """,
-                    (changes.cover_image_name, board_id),
-                )
+                # Change the cover image of a board
+                if changes.cover_image_name is not None:
+                    cursor.execute(
+                        """--sql
+                        UPDATE boards
+                        SET cover_image_name = ?
+                        WHERE board_id = ?;
+                        """,
+                        (changes.cover_image_name, board_id),
+                    )
 
-            # Change the archived status of a board
-            if changes.archived is not None:
-                cursor.execute(
-                    """--sql
-                    UPDATE boards
-                    SET archived = ?
-                    WHERE board_id = ?;
-                    """,
-                    (changes.archived, board_id),
-                )
+                # Change the archived status of a board
+                if changes.archived is not None:
+                    cursor.execute(
+                        """--sql
+                        UPDATE boards
+                        SET archived = ?
+                        WHERE board_id = ?;
+                        """,
+                        (changes.archived, board_id),
+                    )
 
-            self._conn.commit()
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise BoardRecordSaveException from e
+            except sqlite3.Error as e:
+                raise BoardRecordSaveException from e
         return self.get(board_id)
 
     def get_many(
@@ -133,78 +127,77 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
         limit: int = 10,
         include_archived: bool = False,
     ) -> OffsetPaginatedResults[BoardRecord]:
-        cursor = self._conn.cursor()
-
-        # Build base query
-        base_query = """
-                SELECT *
-                FROM boards
-                {archived_filter}
-                ORDER BY {order_by} {direction}
-                LIMIT ? OFFSET ?;
-            """
-
-        # Determine archived filter condition
-        archived_filter = "" if include_archived else "WHERE archived = 0"
-
-        final_query = base_query.format(
-            archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
-        )
-
-        # Execute query to fetch boards
-        cursor.execute(final_query, (limit, offset))
-
-        result = cast(list[sqlite3.Row], cursor.fetchall())
-        boards = [deserialize_board_record(dict(r)) for r in result]
-
-        # Determine count query
-        if include_archived:
-            count_query = """
-                    SELECT COUNT(*)
-                    FROM boards;
-                """
-        else:
-            count_query = """
-                    SELECT COUNT(*)
+        with self._db.transaction() as cursor:
+            # Build base query
+            base_query = """
+                    SELECT *
                     FROM boards
-                    WHERE archived = 0;
+                    {archived_filter}
+                    ORDER BY {order_by} {direction}
+                    LIMIT ? OFFSET ?;
                 """
 
-        # Execute count query
-        cursor.execute(count_query)
+            # Determine archived filter condition
+            archived_filter = "" if include_archived else "WHERE archived = 0"
 
-        count = cast(int, cursor.fetchone()[0])
+            final_query = base_query.format(
+                archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
+            )
+
+            # Execute query to fetch boards
+            cursor.execute(final_query, (limit, offset))
+
+            result = cast(list[sqlite3.Row], cursor.fetchall())
+            boards = [deserialize_board_record(dict(r)) for r in result]
+
+            # Determine count query
+            if include_archived:
+                count_query = """
+                        SELECT COUNT(*)
+                        FROM boards;
+                    """
+            else:
+                count_query = """
+                        SELECT COUNT(*)
+                        FROM boards
+                        WHERE archived = 0;
+                    """
+
+            # Execute count query
+            cursor.execute(count_query)
+
+            count = cast(int, cursor.fetchone()[0])
 
         return OffsetPaginatedResults[BoardRecord](items=boards, offset=offset, limit=limit, total=count)
 
     def get_all(
         self, order_by: BoardRecordOrderBy, direction: SQLiteDirection, include_archived: bool = False
     ) -> list[BoardRecord]:
-        cursor = self._conn.cursor()
-        if order_by == BoardRecordOrderBy.Name:
-            base_query = """
-                    SELECT *
-                    FROM boards
-                    {archived_filter}
-                    ORDER BY LOWER(board_name) {direction}
-                """
-        else:
-            base_query = """
-                    SELECT *
-                    FROM boards
-                    {archived_filter}
-                    ORDER BY {order_by} {direction}
-                """
+        with self._db.transaction() as cursor:
+            if order_by == BoardRecordOrderBy.Name:
+                base_query = """
+                        SELECT *
+                        FROM boards
+                        {archived_filter}
+                        ORDER BY LOWER(board_name) {direction}
+                    """
+            else:
+                base_query = """
+                        SELECT *
+                        FROM boards
+                        {archived_filter}
+                        ORDER BY {order_by} {direction}
+                    """
 
-        archived_filter = "" if include_archived else "WHERE archived = 0"
+            archived_filter = "" if include_archived else "WHERE archived = 0"
 
-        final_query = base_query.format(
-            archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
-        )
+            final_query = base_query.format(
+                archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
+            )
 
-        cursor.execute(final_query)
+            cursor.execute(final_query)
 
-        result = cast(list[sqlite3.Row], cursor.fetchall())
+            result = cast(list[sqlite3.Row], cursor.fetchall())
         boards = [deserialize_board_record(dict(r)) for r in result]
 
         return boards

invokeai/app/services/client_state_persistence/client_state_persistence_base.py

@@ -0,0 +1,42 @@
+from abc import ABC, abstractmethod
+
+
+class ClientStatePersistenceABC(ABC):
+    """
+    Base class for client persistence implementations.
+    This class defines the interface for persisting client data.
+    """
+
+    @abstractmethod
+    def set_by_key(self, queue_id: str, key: str, value: str) -> str:
+        """
+        Set a key-value pair for the client.
+
+        Args:
+            key (str): The key to set.
+            value (str): The value to set for the key.
+
+        Returns:
+            str: The value that was set.
+        """
+        pass
+
+    @abstractmethod
+    def get_by_key(self, queue_id: str, key: str) -> str | None:
+        """
+        Get the value for a specific key of the client.
+
+        Args:
+            key (str): The key to retrieve the value for.
+
+        Returns:
+            str | None: The value associated with the key, or None if the key does not exist.
+        """
+        pass
+
+    @abstractmethod
+    def delete(self, queue_id: str) -> None:
+        """
+        Delete all client state.
+        """
+        pass

invokeai/app/services/client_state_persistence/client_state_persistence_sqlite.py

@@ -0,0 +1,65 @@
+import json
+
+from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC
+from invokeai.app.services.invoker import Invoker
+from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+
+
+class ClientStatePersistenceSqlite(ClientStatePersistenceABC):
+    """
+    Base class for client persistence implementations.
+    This class defines the interface for persisting client data.
+    """
+
+    def __init__(self, db: SqliteDatabase) -> None:
+        super().__init__()
+        self._db = db
+        self._default_row_id = 1
+
+    def start(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+
+    def _get(self) -> dict[str, str] | None:
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                f"""
+                SELECT data FROM client_state
+                WHERE id = {self._default_row_id}
+                """
+            )
+            row = cursor.fetchone()
+            if row is None:
+                return None
+            return json.loads(row[0])
+
+    def set_by_key(self, queue_id: str, key: str, value: str) -> str:
+        state = self._get() or {}
+        state.update({key: value})
+
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                f"""
+                INSERT INTO client_state (id, data)
+                VALUES ({self._default_row_id}, ?)
+                ON CONFLICT(id) DO UPDATE
+                  SET data = excluded.data;
+                """,
+                (json.dumps(state),),
+            )
+
+        return value
+
+    def get_by_key(self, queue_id: str, key: str) -> str | None:
+        state = self._get()
+        if state is None:
+            return None
+        return state.get(key, None)
+
+    def delete(self, queue_id: str) -> None:
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                f"""
+                DELETE FROM client_state
+                WHERE id = {self._default_row_id}
+                """
+            )

invokeai/app/services/config/config_default.py

@@ -24,7 +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")
-DEVICE = Literal["auto", "cpu", "cuda", "cuda:1", "mps"]
 PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
 ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
 ATTENTION_SLICE_SIZE = Literal["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8]
@@ -93,7 +92,7 @@ class InvokeAIAppConfig(BaseSettings):
         vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.
         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 behavior is stable.
         pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.
-        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`
+        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`, `mps`, `cuda:N` (where N is a device number)
         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.
         attention_type: Attention type.<br>Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
@@ -108,6 +107,7 @@ class InvokeAIAppConfig(BaseSettings):
         hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.<br>Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`
         remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.
         scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.
+        unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.
     """
 
     _root: Optional[Path] = PrivateAttr(default=None)
@@ -176,7 +176,7 @@ class InvokeAIAppConfig(BaseSettings):
     pytorch_cuda_alloc_conf: Optional[str] = Field(default=None,            description="Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.")
 
     # 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.")
+    device:                      str = Field(default="auto",                description="Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)", pattern=r"^(auto|cpu|mps|cuda(:\d+)?)$")
     precision:                PRECISION = Field(default="auto",             description="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.")
 
     # GENERATION
@@ -197,6 +197,7 @@ class InvokeAIAppConfig(BaseSettings):
     hashing_algorithm: HASHING_ALGORITHMS = Field(default="blake3_single",  description="Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.")
     remote_api_tokens: Optional[list[URLRegexTokenPair]] = Field(default=None, description="List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.")
     scan_models_on_startup:        bool = Field(default=False,              description="Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.")
+    unsafe_disable_picklescan:     bool = Field(default=False,              description="UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.")
 
     # fmt: on
 

invokeai/app/services/download/download_default.py

@@ -8,6 +8,7 @@ import time
 import traceback
 from pathlib import Path
 from queue import Empty, PriorityQueue
+from shutil import disk_usage
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set
 
 import requests
@@ -335,6 +336,14 @@ class DownloadQueueService(DownloadQueueServiceBase):
 
         assert job.download_path
 
+        free_space = disk_usage(job.download_path.parent).free
+        GB = 2**30
+        self._logger.debug(f"Download is {job.total_bytes / GB:.2f} GB of {free_space / GB:.2f} GB free.")
+        if free_space < job.total_bytes:
+            raise RuntimeError(
+                f"Free disk space {free_space / GB:.2f} GB is not enough for download of {job.total_bytes / GB:.2f} GB."
+            )
+
         # Don't clobber an existing file. See commit 82c2c85202f88c6d24ff84710f297cfc6ae174af
         # for code that instead resumes an interrupted download.
         if job.download_path.exists():

invokeai/app/services/events/events_common.py

@@ -241,6 +241,7 @@ class QueueItemStatusChangedEvent(QueueItemEventBase):
     batch_status: BatchStatus = Field(description="The status of the batch")
     queue_status: SessionQueueStatus = Field(description="The status of the queue")
     session_id: str = Field(description="The ID of the session (aka graph execution state)")
+    credits: Optional[float] = Field(default=None, description="The total credits used for this queue item")
 
     @classmethod
     def build(
@@ -263,6 +264,7 @@ class QueueItemStatusChangedEvent(QueueItemEventBase):
             completed_at=str(queue_item.completed_at) if queue_item.completed_at else None,
             batch_status=batch_status,
             queue_status=queue_status,
+            credits=queue_item.credits,
         )
 
 

invokeai/app/services/image_records/image_records_base.py

@@ -5,6 +5,7 @@ from typing import Optional
 from invokeai.app.invocations.fields import MetadataField
 from invokeai.app.services.image_records.image_records_common import (
     ImageCategory,
+    ImageNamesResult,
     ImageRecord,
     ImageRecordChanges,
     ResourceOrigin,
@@ -97,3 +98,17 @@ class ImageRecordStorageBase(ABC):
     def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]:
         """Gets the most recent image for a board."""
         pass
+
+    @abstractmethod
+    def get_image_names(
+        self,
+        starred_first: bool = True,
+        order_dir: SQLiteDirection = SQLiteDirection.Descending,
+        image_origin: Optional[ResourceOrigin] = None,
+        categories: Optional[list[ImageCategory]] = None,
+        is_intermediate: Optional[bool] = None,
+        board_id: Optional[str] = None,
+        search_term: Optional[str] = None,
+    ) -> ImageNamesResult:
+        """Gets ordered list of image names with metadata for optimistic updates."""
+        pass

invokeai/app/services/image_records/image_records_common.py

@@ -3,7 +3,7 @@ import datetime
 from enum import Enum
 from typing import Optional, Union
 
-from pydantic import Field, StrictBool, StrictStr
+from pydantic import BaseModel, Field, StrictBool, StrictStr
 
 from invokeai.app.util.metaenum import MetaEnum
 from invokeai.app.util.misc import get_iso_timestamp
@@ -207,3 +207,16 @@ def deserialize_image_record(image_dict: dict) -> ImageRecord:
         starred=starred,
         has_workflow=has_workflow,
     )
+
+
+class ImageCollectionCounts(BaseModel):
+    starred_count: int = Field(description="The number of starred images in the collection.")
+    unstarred_count: int = Field(description="The number of unstarred images in the collection.")
+
+
+class ImageNamesResult(BaseModel):
+    """Response containing ordered image names with metadata for optimistic updates."""
+
+    image_names: list[str] = Field(description="Ordered list of image names")
+    starred_count: int = Field(description="Number of starred images (when starred_first=True)")
+    total_count: int = Field(description="Total number of images matching the query")

invokeai/app/services/image_records/image_records_sqlite.py

@@ -7,6 +7,7 @@ from invokeai.app.services.image_records.image_records_base import ImageRecordSt
 from invokeai.app.services.image_records.image_records_common import (
     IMAGE_DTO_COLS,
     ImageCategory,
+    ImageNamesResult,
     ImageRecord,
     ImageRecordChanges,
     ImageRecordDeleteException,
@@ -23,22 +24,22 @@ from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
 class SqliteImageRecordStorage(ImageRecordStorageBase):
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
-        self._conn = db.conn
+        self._db = db
 
     def get(self, image_name: str) -> ImageRecord:
-        try:
-            cursor = self._conn.cursor()
-            cursor.execute(
-                f"""--sql
-                SELECT {IMAGE_DTO_COLS} FROM images
-                WHERE image_name = ?;
-                """,
-                (image_name,),
-            )
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    f"""--sql
+                    SELECT {IMAGE_DTO_COLS} FROM images
+                    WHERE image_name = ?;
+                    """,
+                    (image_name,),
+                )
 
-            result = cast(Optional[sqlite3.Row], cursor.fetchone())
-        except sqlite3.Error as e:
-            raise ImageRecordNotFoundException from e
+                result = cast(Optional[sqlite3.Row], cursor.fetchone())
+            except sqlite3.Error as e:
+                raise ImageRecordNotFoundException from e
 
         if not result:
             raise ImageRecordNotFoundException
@@ -46,17 +47,20 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
         return deserialize_image_record(dict(result))
 
     def get_metadata(self, image_name: str) -> Optional[MetadataField]:
-        try:
-            cursor = self._conn.cursor()
-            cursor.execute(
-                """--sql
-                SELECT metadata FROM images
-                WHERE image_name = ?;
-                """,
-                (image_name,),
-            )
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    """--sql
+                    SELECT metadata FROM images
+                    WHERE image_name = ?;
+                    """,
+                    (image_name,),
+                )
 
-            result = cast(Optional[sqlite3.Row], cursor.fetchone())
+                result = cast(Optional[sqlite3.Row], cursor.fetchone())
+
+            except sqlite3.Error as e:
+                raise ImageRecordNotFoundException from e
 
             if not result:
                 raise ImageRecordNotFoundException
@@ -64,64 +68,60 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
             as_dict = dict(result)
             metadata_raw = cast(Optional[str], as_dict.get("metadata", None))
             return MetadataFieldValidator.validate_json(metadata_raw) if metadata_raw is not None else None
-        except sqlite3.Error as e:
-            raise ImageRecordNotFoundException from e
 
     def update(
         self,
         image_name: str,
         changes: ImageRecordChanges,
     ) -> None:
-        try:
-            cursor = self._conn.cursor()
-            # Change the category of the image
-            if changes.image_category is not None:
-                cursor.execute(
-                    """--sql
-                    UPDATE images
-                    SET image_category = ?
-                    WHERE image_name = ?;
-                    """,
-                    (changes.image_category, image_name),
-                )
+        with self._db.transaction() as cursor:
+            try:
+                # Change the category of the image
+                if changes.image_category is not None:
+                    cursor.execute(
+                        """--sql
+                        UPDATE images
+                        SET image_category = ?
+                        WHERE image_name = ?;
+                        """,
+                        (changes.image_category, image_name),
+                    )
 
-            # Change the session associated with the image
-            if changes.session_id is not None:
-                cursor.execute(
-                    """--sql
-                    UPDATE images
-                    SET session_id = ?
-                    WHERE image_name = ?;
-                    """,
-                    (changes.session_id, image_name),
-                )
+                # Change the session associated with the image
+                if changes.session_id is not None:
+                    cursor.execute(
+                        """--sql
+                        UPDATE images
+                        SET session_id = ?
+                        WHERE image_name = ?;
+                        """,
+                        (changes.session_id, image_name),
+                    )
 
-            # Change the image's `is_intermediate`` flag
-            if changes.is_intermediate is not None:
-                cursor.execute(
-                    """--sql
-                    UPDATE images
-                    SET is_intermediate = ?
-                    WHERE image_name = ?;
-                    """,
-                    (changes.is_intermediate, image_name),
-                )
+                # Change the image's `is_intermediate`` flag
+                if changes.is_intermediate is not None:
+                    cursor.execute(
+                        """--sql
+                        UPDATE images
+                        SET is_intermediate = ?
+                        WHERE image_name = ?;
+                        """,
+                        (changes.is_intermediate, image_name),
+                    )
 
-            # Change the image's `starred`` state
-            if changes.starred is not None:
-                cursor.execute(
-                    """--sql
-                    UPDATE images
-                    SET starred = ?
-                    WHERE image_name = ?;
-                    """,
-                    (changes.starred, image_name),
-                )
+                # Change the image's `starred`` state
+                if changes.starred is not None:
+                    cursor.execute(
+                        """--sql
+                        UPDATE images
+                        SET starred = ?
+                        WHERE image_name = ?;
+                        """,
+                        (changes.starred, image_name),
+                    )
 
-            self._conn.commit()
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise ImageRecordSaveException from e
+            except sqlite3.Error as e:
+                raise ImageRecordSaveException from e
 
     def get_many(
         self,
@@ -135,166 +135,162 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
         board_id: Optional[str] = None,
         search_term: Optional[str] = None,
     ) -> OffsetPaginatedResults[ImageRecord]:
-        cursor = self._conn.cursor()
-
-        # Manually build two queries - one for the count, one for the records
-        count_query = """--sql
-        SELECT COUNT(*)
-        FROM images
-        LEFT JOIN board_images ON board_images.image_name = images.image_name
-        WHERE 1=1
-        """
-
-        images_query = f"""--sql
-        SELECT {IMAGE_DTO_COLS}
-        FROM images
-        LEFT JOIN board_images ON board_images.image_name = images.image_name
-        WHERE 1=1
-        """
-
-        query_conditions = ""
-        query_params: list[Union[int, str, bool]] = []
-
-        if image_origin is not None:
-            query_conditions += """--sql
-            AND images.image_origin = ?
-            """
-            query_params.append(image_origin.value)
-
-        if categories is not None:
-            # Convert the enum values to unique list of strings
-            category_strings = [c.value for c in set(categories)]
-            # Create the correct length of placeholders
-            placeholders = ",".join("?" * len(category_strings))
-
-            query_conditions += f"""--sql
-            AND images.image_category IN ( {placeholders} )
+        with self._db.transaction() as cursor:
+            # Manually build two queries - one for the count, one for the records
+            count_query = """--sql
+            SELECT COUNT(*)
+            FROM images
+            LEFT JOIN board_images ON board_images.image_name = images.image_name
+            WHERE 1=1
             """
 
-            # Unpack the included categories into the query params
-            for c in category_strings:
-                query_params.append(c)
-
-        if is_intermediate is not None:
-            query_conditions += """--sql
-            AND images.is_intermediate = ?
+            images_query = f"""--sql
+            SELECT {IMAGE_DTO_COLS}
+            FROM images
+            LEFT JOIN board_images ON board_images.image_name = images.image_name
+            WHERE 1=1
             """
 
-            query_params.append(is_intermediate)
+            query_conditions = ""
+            query_params: list[Union[int, str, bool]] = []
 
-        # board_id of "none" is reserved for images without a board
-        if board_id == "none":
-            query_conditions += """--sql
-            AND board_images.board_id IS NULL
-            """
-        elif board_id is not None:
-            query_conditions += """--sql
-            AND board_images.board_id = ?
-            """
-            query_params.append(board_id)
+            if image_origin is not None:
+                query_conditions += """--sql
+                AND images.image_origin = ?
+                """
+                query_params.append(image_origin.value)
 
-        # Search term condition
-        if search_term:
-            query_conditions += """--sql
-            AND images.metadata LIKE ?
-            """
-            query_params.append(f"%{search_term.lower()}%")
+            if categories is not None:
+                # Convert the enum values to unique list of strings
+                category_strings = [c.value for c in set(categories)]
+                # Create the correct length of placeholders
+                placeholders = ",".join("?" * len(category_strings))
 
-        if starred_first:
-            query_pagination = f"""--sql
-            ORDER BY images.starred DESC, images.created_at {order_dir.value} LIMIT ? OFFSET ?
-            """
-        else:
-            query_pagination = f"""--sql
-            ORDER BY images.created_at {order_dir.value} LIMIT ? OFFSET ?
-            """
+                query_conditions += f"""--sql
+                AND images.image_category IN ( {placeholders} )
+                """
 
-        # Final images query with pagination
-        images_query += query_conditions + query_pagination + ";"
-        # Add all the parameters
-        images_params = query_params.copy()
-        # Add the pagination parameters
-        images_params.extend([limit, offset])
+                # Unpack the included categories into the query params
+                for c in category_strings:
+                    query_params.append(c)
 
-        # Build the list of images, deserializing each row
-        cursor.execute(images_query, images_params)
-        result = cast(list[sqlite3.Row], cursor.fetchall())
-        images = [deserialize_image_record(dict(r)) for r in result]
+            if is_intermediate is not None:
+                query_conditions += """--sql
+                AND images.is_intermediate = ?
+                """
 
-        # Set up and execute the count query, without pagination
-        count_query += query_conditions + ";"
-        count_params = query_params.copy()
-        cursor.execute(count_query, count_params)
-        count = cast(int, cursor.fetchone()[0])
+                query_params.append(is_intermediate)
+
+            # board_id of "none" is reserved for images without a board
+            if board_id == "none":
+                query_conditions += """--sql
+                AND board_images.board_id IS NULL
+                """
+            elif board_id is not None:
+                query_conditions += """--sql
+                AND board_images.board_id = ?
+                """
+                query_params.append(board_id)
+
+            # Search term condition
+            if search_term:
+                query_conditions += """--sql
+                AND (
+                    images.metadata LIKE ?
+                    OR images.created_at LIKE ?
+                )
+                """
+                query_params.append(f"%{search_term.lower()}%")
+                query_params.append(f"%{search_term.lower()}%")
+
+            if starred_first:
+                query_pagination = f"""--sql
+                ORDER BY images.starred DESC, images.created_at {order_dir.value} LIMIT ? OFFSET ?
+                """
+            else:
+                query_pagination = f"""--sql
+                ORDER BY images.created_at {order_dir.value} LIMIT ? OFFSET ?
+                """
+
+            # Final images query with pagination
+            images_query += query_conditions + query_pagination + ";"
+            # Add all the parameters
+            images_params = query_params.copy()
+            # Add the pagination parameters
+            images_params.extend([limit, offset])
+
+            # Build the list of images, deserializing each row
+            cursor.execute(images_query, images_params)
+            result = cast(list[sqlite3.Row], cursor.fetchall())
+
+            images = [deserialize_image_record(dict(r)) for r in result]
+
+            # Set up and execute the count query, without pagination
+            count_query += query_conditions + ";"
+            count_params = query_params.copy()
+            cursor.execute(count_query, count_params)
+            count = cast(int, cursor.fetchone()[0])
 
         return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=count)
 
     def delete(self, image_name: str) -> None:
-        try:
-            cursor = self._conn.cursor()
-            cursor.execute(
-                """--sql
-                DELETE FROM images
-                WHERE image_name = ?;
-                """,
-                (image_name,),
-            )
-            self._conn.commit()
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise ImageRecordDeleteException from e
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    """--sql
+                    DELETE FROM images
+                    WHERE image_name = ?;
+                    """,
+                    (image_name,),
+                )
+            except sqlite3.Error as e:
+                raise ImageRecordDeleteException from e
 
     def delete_many(self, image_names: list[str]) -> None:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
+            try:
+                placeholders = ",".join("?" for _ in image_names)
 
-            placeholders = ",".join("?" for _ in image_names)
+                # Construct the SQLite query with the placeholders
+                query = f"DELETE FROM images WHERE image_name IN ({placeholders})"
 
-            # Construct the SQLite query with the placeholders
-            query = f"DELETE FROM images WHERE image_name IN ({placeholders})"
+                # Execute the query with the list of IDs as parameters
+                cursor.execute(query, image_names)
 
-            # Execute the query with the list of IDs as parameters
-            cursor.execute(query, image_names)
-
-            self._conn.commit()
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise ImageRecordDeleteException from e
+            except sqlite3.Error as e:
+                raise ImageRecordDeleteException from e
 
     def get_intermediates_count(self) -> int:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT COUNT(*) FROM images
-            WHERE is_intermediate = TRUE;
-            """
-        )
-        count = cast(int, cursor.fetchone()[0])
-        self._conn.commit()
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT COUNT(*) FROM images
+                WHERE is_intermediate = TRUE;
+                """
+            )
+            count = cast(int, cursor.fetchone()[0])
         return count
 
     def delete_intermediates(self) -> list[str]:
-        try:
-            cursor = self._conn.cursor()
-            cursor.execute(
-                """--sql
-                SELECT image_name FROM images
-                WHERE is_intermediate = TRUE;
-                """
-            )
-            result = cast(list[sqlite3.Row], cursor.fetchall())
-            image_names = [r[0] for r in result]
-            cursor.execute(
-                """--sql
-                DELETE FROM images
-                WHERE is_intermediate = TRUE;
-                """
-            )
-            self._conn.commit()
-            return image_names
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise ImageRecordDeleteException from e
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    """--sql
+                    SELECT image_name FROM images
+                    WHERE is_intermediate = TRUE;
+                    """
+                )
+                result = cast(list[sqlite3.Row], cursor.fetchall())
+                image_names = [r[0] for r in result]
+                cursor.execute(
+                    """--sql
+                    DELETE FROM images
+                    WHERE is_intermediate = TRUE;
+                    """
+                )
+            except sqlite3.Error as e:
+                raise ImageRecordDeleteException from e
+        return image_names
 
     def save(
         self,
@@ -310,75 +306,165 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
         node_id: Optional[str] = None,
         metadata: Optional[str] = None,
     ) -> datetime:
-        try:
-            cursor = self._conn.cursor()
-            cursor.execute(
-                """--sql
-                INSERT OR IGNORE INTO images (
-                    image_name,
-                    image_origin,
-                    image_category,
-                    width,
-                    height,
-                    node_id,
-                    session_id,
-                    metadata,
-                    is_intermediate,
-                    starred,
-                    has_workflow
-                    )
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
-                """,
-                (
-                    image_name,
-                    image_origin.value,
-                    image_category.value,
-                    width,
-                    height,
-                    node_id,
-                    session_id,
-                    metadata,
-                    is_intermediate,
-                    starred,
-                    has_workflow,
-                ),
-            )
-            self._conn.commit()
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    """--sql
+                    INSERT OR IGNORE INTO images (
+                        image_name,
+                        image_origin,
+                        image_category,
+                        width,
+                        height,
+                        node_id,
+                        session_id,
+                        metadata,
+                        is_intermediate,
+                        starred,
+                        has_workflow
+                        )
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
+                    """,
+                    (
+                        image_name,
+                        image_origin.value,
+                        image_category.value,
+                        width,
+                        height,
+                        node_id,
+                        session_id,
+                        metadata,
+                        is_intermediate,
+                        starred,
+                        has_workflow,
+                    ),
+                )
 
-            cursor.execute(
-                """--sql
-                SELECT created_at
-                FROM images
-                WHERE image_name = ?;
-                """,
-                (image_name,),
-            )
+                cursor.execute(
+                    """--sql
+                    SELECT created_at
+                    FROM images
+                    WHERE image_name = ?;
+                    """,
+                    (image_name,),
+                )
 
-            created_at = datetime.fromisoformat(cursor.fetchone()[0])
+                created_at = datetime.fromisoformat(cursor.fetchone()[0])
 
-            return created_at
-        except sqlite3.Error as e:
-            self._conn.rollback()
-            raise ImageRecordSaveException from e
+            except sqlite3.Error as e:
+                raise ImageRecordSaveException from e
+        return created_at
 
     def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT images.*
-            FROM images
-            JOIN board_images ON images.image_name = board_images.image_name
-            WHERE board_images.board_id = ?
-            AND images.is_intermediate = FALSE
-            ORDER BY images.starred DESC, images.created_at DESC
-            LIMIT 1;
-            """,
-            (board_id,),
-        )
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT images.*
+                FROM images
+                JOIN board_images ON images.image_name = board_images.image_name
+                WHERE board_images.board_id = ?
+                AND images.is_intermediate = FALSE
+                ORDER BY images.starred DESC, images.created_at DESC
+                LIMIT 1;
+                """,
+                (board_id,),
+            )
 
-        result = cast(Optional[sqlite3.Row], cursor.fetchone())
+            result = cast(Optional[sqlite3.Row], cursor.fetchone())
 
         if result is None:
             return None
 
         return deserialize_image_record(dict(result))
+
+    def get_image_names(
+        self,
+        starred_first: bool = True,
+        order_dir: SQLiteDirection = SQLiteDirection.Descending,
+        image_origin: Optional[ResourceOrigin] = None,
+        categories: Optional[list[ImageCategory]] = None,
+        is_intermediate: Optional[bool] = None,
+        board_id: Optional[str] = None,
+        search_term: Optional[str] = None,
+    ) -> ImageNamesResult:
+        with self._db.transaction() as cursor:
+            # Build query conditions (reused for both starred count and image names queries)
+            query_conditions = ""
+            query_params: list[Union[int, str, bool]] = []
+
+            if image_origin is not None:
+                query_conditions += """--sql
+                AND images.image_origin = ?
+                """
+                query_params.append(image_origin.value)
+
+            if categories is not None:
+                category_strings = [c.value for c in set(categories)]
+                placeholders = ",".join("?" * len(category_strings))
+                query_conditions += f"""--sql
+                AND images.image_category IN ( {placeholders} )
+                """
+                for c in category_strings:
+                    query_params.append(c)
+
+            if is_intermediate is not None:
+                query_conditions += """--sql
+                AND images.is_intermediate = ?
+                """
+                query_params.append(is_intermediate)
+
+            if board_id == "none":
+                query_conditions += """--sql
+                AND board_images.board_id IS NULL
+                """
+            elif board_id is not None:
+                query_conditions += """--sql
+                AND board_images.board_id = ?
+                """
+                query_params.append(board_id)
+
+            if search_term:
+                query_conditions += """--sql
+                AND (
+                    images.metadata LIKE ?
+                    OR images.created_at LIKE ?
+                )
+                """
+                query_params.append(f"%{search_term.lower()}%")
+                query_params.append(f"%{search_term.lower()}%")
+
+            # Get starred count if starred_first is enabled
+            starred_count = 0
+            if starred_first:
+                starred_count_query = f"""--sql
+                SELECT COUNT(*)
+                FROM images
+                LEFT JOIN board_images ON board_images.image_name = images.image_name
+                WHERE images.starred = TRUE AND (1=1{query_conditions})
+                """
+                cursor.execute(starred_count_query, query_params)
+                starred_count = cast(int, cursor.fetchone()[0])
+
+            # Get all image names with proper ordering
+            if starred_first:
+                names_query = f"""--sql
+                SELECT images.image_name
+                FROM images
+                LEFT JOIN board_images ON board_images.image_name = images.image_name
+                WHERE 1=1{query_conditions}
+                ORDER BY images.starred DESC, images.created_at {order_dir.value}
+                """
+            else:
+                names_query = f"""--sql
+                SELECT images.image_name
+                FROM images
+                LEFT JOIN board_images ON board_images.image_name = images.image_name
+                WHERE 1=1{query_conditions}
+                ORDER BY images.created_at {order_dir.value}
+                """
+
+            cursor.execute(names_query, query_params)
+            result = cast(list[sqlite3.Row], cursor.fetchall())
+        image_names = [row[0] for row in result]
+
+        return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names))

invokeai/app/services/images/images_base.py

@@ -6,6 +6,7 @@ from PIL.Image import Image as PILImageType
 from invokeai.app.invocations.fields import MetadataField
 from invokeai.app.services.image_records.image_records_common import (
     ImageCategory,
+    ImageNamesResult,
     ImageRecord,
     ImageRecordChanges,
     ResourceOrigin,
@@ -125,7 +126,7 @@ class ImageServiceABC(ABC):
         board_id: Optional[str] = None,
         search_term: Optional[str] = None,
     ) -> OffsetPaginatedResults[ImageDTO]:
-        """Gets a paginated list of image DTOs."""
+        """Gets a paginated list of image DTOs with starred images first when starred_first=True."""
         pass
 
     @abstractmethod
@@ -147,3 +148,17 @@ class ImageServiceABC(ABC):
     def delete_images_on_board(self, board_id: str):
         """Deletes all images on a board."""
         pass
+
+    @abstractmethod
+    def get_image_names(
+        self,
+        starred_first: bool = True,
+        order_dir: SQLiteDirection = SQLiteDirection.Descending,
+        image_origin: Optional[ResourceOrigin] = None,
+        categories: Optional[list[ImageCategory]] = None,
+        is_intermediate: Optional[bool] = None,
+        board_id: Optional[str] = None,
+        search_term: Optional[str] = None,
+    ) -> ImageNamesResult:
+        """Gets ordered list of image names with metadata for optimistic updates."""
+        pass

invokeai/app/services/images/images_common.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from pydantic import Field
+from pydantic import BaseModel, Field
 
 from invokeai.app.services.image_records.image_records_common import ImageRecord
 from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
@@ -39,3 +39,27 @@ def image_record_to_dto(
         thumbnail_url=thumbnail_url,
         board_id=board_id,
     )
+
+
+class ResultWithAffectedBoards(BaseModel):
+    affected_boards: list[str] = Field(description="The ids of boards affected by the delete operation")
+
+
+class DeleteImagesResult(ResultWithAffectedBoards):
+    deleted_images: list[str] = Field(description="The names of the images that were deleted")
+
+
+class StarredImagesResult(ResultWithAffectedBoards):
+    starred_images: list[str] = Field(description="The names of the images that were starred")
+
+
+class UnstarredImagesResult(ResultWithAffectedBoards):
+    unstarred_images: list[str] = Field(description="The names of the images that were unstarred")
+
+
+class AddImagesToBoardResult(ResultWithAffectedBoards):
+    added_images: list[str] = Field(description="The image names that were added to the board")
+
+
+class RemoveImagesFromBoardResult(ResultWithAffectedBoards):
+    removed_images: list[str] = Field(description="The image names that were removed from their board")

invokeai/app/services/images/images_default.py

@@ -10,6 +10,7 @@ from invokeai.app.services.image_files.image_files_common import (
 )
 from invokeai.app.services.image_records.image_records_common import (
     ImageCategory,
+    ImageNamesResult,
     ImageRecord,
     ImageRecordChanges,
     ImageRecordDeleteException,
@@ -78,7 +79,7 @@ class ImageService(ImageServiceABC):
                         board_id=board_id, image_name=image_name
                     )
                 except Exception as e:
-                    self.__invoker.services.logger.warn(f"Failed to add image to board {board_id}: {str(e)}")
+                    self.__invoker.services.logger.warning(f"Failed to add image to board {board_id}: {str(e)}")
             self.__invoker.services.image_files.save(
                 image_name=image_name, image=image, metadata=metadata, workflow=workflow, graph=graph
             )
@@ -309,3 +310,27 @@ class ImageService(ImageServiceABC):
         except Exception as e:
             self.__invoker.services.logger.error("Problem getting intermediates count")
             raise e
+
+    def get_image_names(
+        self,
+        starred_first: bool = True,
+        order_dir: SQLiteDirection = SQLiteDirection.Descending,
+        image_origin: Optional[ResourceOrigin] = None,
+        categories: Optional[list[ImageCategory]] = None,
+        is_intermediate: Optional[bool] = None,
+        board_id: Optional[str] = None,
+        search_term: Optional[str] = None,
+    ) -> ImageNamesResult:
+        try:
+            return self.__invoker.services.image_records.get_image_names(
+                starred_first=starred_first,
+                order_dir=order_dir,
+                image_origin=image_origin,
+                categories=categories,
+                is_intermediate=is_intermediate,
+                board_id=board_id,
+                search_term=search_term,
+            )
+        except Exception as e:
+            self.__invoker.services.logger.error("Problem getting image names")
+            raise e

invokeai/app/services/invocation_services.py

@@ -17,6 +17,7 @@ if TYPE_CHECKING:
     from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase
     from invokeai.app.services.boards.boards_base import BoardServiceABC
     from invokeai.app.services.bulk_download.bulk_download_base import BulkDownloadBase
+    from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC
     from invokeai.app.services.config import InvokeAIAppConfig
     from invokeai.app.services.download import DownloadQueueServiceBase
     from invokeai.app.services.events.events_base import EventServiceBase
@@ -27,6 +28,10 @@ if TYPE_CHECKING:
     from invokeai.app.services.invocation_stats.invocation_stats_base import InvocationStatsServiceBase
     from invokeai.app.services.model_images.model_images_base import ModelImageFileStorageBase
     from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
+    from invokeai.app.services.model_relationship_records.model_relationship_records_base import (
+        ModelRelationshipRecordStorageBase,
+    )
+    from invokeai.app.services.model_relationships.model_relationships_base import ModelRelationshipsServiceABC
     from invokeai.app.services.names.names_base import NameServiceBase
     from invokeai.app.services.session_processor.session_processor_base import SessionProcessorBase
     from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase
@@ -54,6 +59,8 @@ class InvocationServices:
         logger: "Logger",
         model_images: "ModelImageFileStorageBase",
         model_manager: "ModelManagerServiceBase",
+        model_relationships: "ModelRelationshipsServiceABC",
+        model_relationship_records: "ModelRelationshipRecordStorageBase",
         download_queue: "DownloadQueueServiceBase",
         performance_statistics: "InvocationStatsServiceBase",
         session_queue: "SessionQueueBase",
@@ -67,6 +74,7 @@ class InvocationServices:
         style_preset_records: "StylePresetRecordsStorageBase",
         style_preset_image_files: "StylePresetImageFileStorageBase",
         workflow_thumbnails: "WorkflowThumbnailServiceBase",
+        client_state_persistence: "ClientStatePersistenceABC",
     ):
         self.board_images = board_images
         self.board_image_records = board_image_records
@@ -81,6 +89,8 @@ class InvocationServices:
         self.logger = logger
         self.model_images = model_images
         self.model_manager = model_manager
+        self.model_relationships = model_relationships
+        self.model_relationship_records = model_relationship_records
         self.download_queue = download_queue
         self.performance_statistics = performance_statistics
         self.session_queue = session_queue
@@ -94,3 +104,4 @@ class InvocationServices:
         self.style_preset_records = style_preset_records
         self.style_preset_image_files = style_preset_image_files
         self.workflow_thumbnails = workflow_thumbnails
+        self.client_state_persistence = client_state_persistence

invokeai/app/services/invocation_stats/invocation_stats_base.py

@@ -60,7 +60,7 @@ class InvocationStatsServiceBase(ABC):
         pass
 
     @abstractmethod
-    def reset_stats(self):
+    def reset_stats(self, graph_execution_state_id: str) -> None:
         """Reset all stored statistics."""
         pass
 

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -73,9 +73,9 @@ class InvocationStatsService(InvocationStatsServiceBase):
             )
             self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
 
-    def reset_stats(self):
-        self._stats = {}
-        self._cache_stats = {}
+    def reset_stats(self, graph_execution_state_id: str) -> None:
+        self._stats.pop(graph_execution_state_id, None)
+        self._cache_stats.pop(graph_execution_state_id, None)
 
     def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
         graph_stats_summary = self._get_graph_summary(graph_execution_state_id)

invokeai/app/services/model_install/model_install_default.py

@@ -7,7 +7,7 @@ import threading
 import time
 from pathlib import Path
 from queue import Empty, Queue
-from shutil import copyfile, copytree, move, rmtree
+from shutil import move, rmtree
 from tempfile import mkdtemp
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 
@@ -38,6 +38,7 @@ from invokeai.backend.model_manager.config import (
     AnyModelConfig,
     CheckpointConfigBase,
     InvalidModelConfigException,
+    ModelConfigBase,
 )
 from invokeai.backend.model_manager.legacy_probe import ModelProbe
 from invokeai.backend.model_manager.metadata import (
@@ -50,6 +51,7 @@ from invokeai.backend.model_manager.metadata import (
 from invokeai.backend.model_manager.metadata.metadata_base import HuggingFaceMetadata
 from invokeai.backend.model_manager.search import ModelSearch
 from invokeai.backend.model_manager.taxonomy import ModelRepoVariant, ModelSourceType
+from invokeai.backend.model_manager.util.lora_metadata_extractor import apply_lora_metadata
 from invokeai.backend.util import InvokeAILogger
 from invokeai.backend.util.catch_sigint import catch_sigint
 from invokeai.backend.util.devices import TorchDevice
@@ -147,7 +149,7 @@ class ModelInstallService(ModelInstallServiceBase):
     def _clear_pending_jobs(self) -> None:
         for job in self.list_jobs():
             if not job.in_terminal_state:
-                self._logger.warning("Cancelling job {job.id}")
+                self._logger.warning(f"Cancelling job {job.id}")
                 self.cancel_job(job)
         while True:
             try:
@@ -184,13 +186,14 @@ class ModelInstallService(ModelInstallServiceBase):
         info: AnyModelConfig = self._probe(Path(model_path), config)  # type: ignore
 
         if preferred_name := config.name:
-            preferred_name = Path(preferred_name).with_suffix(model_path.suffix)
+            # Careful! Don't use pathlib.Path(...).with_suffix - it can will strip everything after the first dot.
+            preferred_name = f"{preferred_name}{model_path.suffix}"
 
         dest_path = (
             self.app_config.models_path / info.base.value / info.type.value / (preferred_name or model_path.name)
         )
         try:
-            new_path = self._copy_model(model_path, dest_path)
+            new_path = self._move_model(model_path, dest_path)
         except FileExistsError as excp:
             raise DuplicateModelException(
                 f"A model named {model_path.name} is already installed at {dest_path.as_posix()}"
@@ -615,16 +618,6 @@ class ModelInstallService(ModelInstallServiceBase):
         self.record_store.update_model(key, ModelRecordChanges(path=model.path))
         return model
 
-    def _copy_model(self, old_path: Path, new_path: Path) -> Path:
-        if old_path == new_path:
-            return old_path
-        new_path.parent.mkdir(parents=True, exist_ok=True)
-        if old_path.is_dir():
-            copytree(old_path, new_path)
-        else:
-            copyfile(old_path, new_path)
-        return new_path
-
     def _move_model(self, old_path: Path, new_path: Path) -> Path:
         if old_path == new_path:
             return old_path
@@ -646,14 +639,18 @@ class ModelInstallService(ModelInstallServiceBase):
         hash_algo = self._app_config.hashing_algorithm
         fields = config.model_dump()
 
-        return ModelProbe.probe(model_path=model_path, fields=fields, hash_algo=hash_algo)
-
-        # New model probe API is disabled pending resolution of issue caused by a change of the ordering of checks.
-        # See commit message for details.
-        # try:
-        #     return ModelConfigBase.classify(model_path=model_path, hash_algo=hash_algo, **fields)
-        # except InvalidModelConfigException:
-        #     return ModelProbe.probe(model_path=model_path, fields=fields, hash_algo=hash_algo)  # type: ignore
+        # WARNING!
+        # The legacy probe relies on the implicit order of tests to determine model classification.
+        # This can lead to regressions between the legacy and new probes.
+        # Do NOT change the order of `probe` and `classify` without implementing one of the following fixes:
+        # Short-term fix: `classify` tests `matches` in the same order as the legacy probe.
+        # Long-term fix: Improve `matches` to be more specific so that only one config matches
+        #   any given model - eliminating ambiguity and removing reliance on order.
+        # After implementing either of these fixes, remove @pytest.mark.xfail from `test_regression_against_model_probe`
+        try:
+            return ModelProbe.probe(model_path=model_path, fields=fields, hash_algo=hash_algo)  # type: ignore
+        except InvalidModelConfigException:
+            return ModelConfigBase.classify(model_path, hash_algo, **fields)
 
     def _register(
         self, model_path: Path, config: Optional[ModelRecordChanges] = None, info: Optional[AnyModelConfig] = None
@@ -662,6 +659,10 @@ class ModelInstallService(ModelInstallServiceBase):
 
         info = info or self._probe(model_path, config)
 
+        # Apply LoRA metadata if applicable
+        model_images_path = self.app_config.models_path / "model_images"
+        apply_lora_metadata(info, model_path.resolve(), model_images_path)
+
         model_path = model_path.resolve()
 
         # Models in the Invoke-managed models dir should use relative paths.

invokeai/app/services/model_load/model_load_default.py

@@ -87,9 +87,21 @@ class ModelLoadService(ModelLoadServiceBase):
         def torch_load_file(checkpoint: Path) -> AnyModel:
             scan_result = scan_file_path(checkpoint)
             if scan_result.infected_files != 0:
-                raise Exception(f"The model at {checkpoint} is potentially infected by malware. Aborting load.")
+                if self._app_config.unsafe_disable_picklescan:
+                    self._logger.warning(
+                        f"Model at {checkpoint} is potentially infected by malware, but picklescan is disabled. "
+                        "Proceeding with caution."
+                    )
+                else:
+                    raise Exception(f"The model at {checkpoint} is potentially infected by malware. Aborting load.")
             if scan_result.scan_err:
-                raise Exception(f"Error scanning model at {checkpoint} for malware. Aborting load.")
+                if self._app_config.unsafe_disable_picklescan:
+                    self._logger.warning(
+                        f"Error scanning model at {checkpoint} for malware, but picklescan is disabled. "
+                        "Proceeding with caution."
+                    )
+                else:
+                    raise Exception(f"Error scanning model at {checkpoint} for malware. Aborting load.")
 
             result = torch_load(checkpoint, map_location="cpu")
             return result

invokeai/app/services/model_records/model_records_sql.py

@@ -78,11 +78,6 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
         self._db = db
         self._logger = logger
 
-    @property
-    def db(self) -> SqliteDatabase:
-        """Return the underlying database."""
-        return self._db
-
     def add_model(self, config: AnyModelConfig) -> AnyModelConfig:
         """
         Add a model to the database.
@@ -93,38 +88,33 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
 
         Can raise DuplicateModelException and InvalidModelConfigException exceptions.
         """
-        try:
-            cursor = self._db.conn.cursor()
-            cursor.execute(
-                """--sql
-                INSERT INTO models (
-                    id,
-                    config
-                    )
-                VALUES (?,?);
-                """,
-                (
-                    config.key,
-                    config.model_dump_json(),
-                ),
-            )
-            self._db.conn.commit()
+        with self._db.transaction() as cursor:
+            try:
+                cursor.execute(
+                    """--sql
+                    INSERT INTO models (
+                        id,
+                        config
+                        )
+                    VALUES (?,?);
+                    """,
+                    (
+                        config.key,
+                        config.model_dump_json(),
+                    ),
+                )
 
-        except sqlite3.IntegrityError as e:
-            self._db.conn.rollback()
-            if "UNIQUE constraint failed" in str(e):
-                if "models.path" in str(e):
-                    msg = f"A model with path '{config.path}' is already installed"
-                elif "models.name" in str(e):
-                    msg = f"A model with name='{config.name}', type='{config.type}', base='{config.base}' is already installed"
+            except sqlite3.IntegrityError as e:
+                if "UNIQUE constraint failed" in str(e):
+                    if "models.path" in str(e):
+                        msg = f"A model with path '{config.path}' is already installed"
+                    elif "models.name" in str(e):
+                        msg = f"A model with name='{config.name}', type='{config.type}', base='{config.base}' is already installed"
+                    else:
+                        msg = f"A model with key '{config.key}' is already installed"
+                    raise DuplicateModelException(msg) from e
                 else:
-                    msg = f"A model with key '{config.key}' is already installed"
-                raise DuplicateModelException(msg) from e
-            else:
-                raise e
-        except sqlite3.Error as e:
-            self._db.conn.rollback()
-            raise e
+                    raise e
 
         return self.get_model(config.key)
 
@@ -136,8 +126,7 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
 
         Can raise an UnknownModelException
         """
-        try:
-            cursor = self._db.conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 DELETE FROM models
@@ -147,22 +136,17 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
             )
             if cursor.rowcount == 0:
                 raise UnknownModelException("model not found")
-            self._db.conn.commit()
-        except sqlite3.Error as e:
-            self._db.conn.rollback()
-            raise e
 
     def update_model(self, key: str, changes: ModelRecordChanges) -> AnyModelConfig:
-        record = self.get_model(key)
+        with self._db.transaction() as cursor:
+            record = self.get_model(key)
 
-        # Model configs use pydantic's `validate_assignment`, so each change is validated by pydantic.
-        for field_name in changes.model_fields_set:
-            setattr(record, field_name, getattr(changes, field_name))
+            # Model configs use pydantic's `validate_assignment`, so each change is validated by pydantic.
+            for field_name in changes.model_fields_set:
+                setattr(record, field_name, getattr(changes, field_name))
 
-        json_serialized = record.model_dump_json()
+            json_serialized = record.model_dump_json()
 
-        try:
-            cursor = self._db.conn.cursor()
             cursor.execute(
                 """--sql
                 UPDATE models
@@ -174,10 +158,6 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
             )
             if cursor.rowcount == 0:
                 raise UnknownModelException("model not found")
-            self._db.conn.commit()
-        except sqlite3.Error as e:
-            self._db.conn.rollback()
-            raise e
 
         return self.get_model(key)
 
@@ -189,30 +169,30 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
 
         Exceptions: UnknownModelException
         """
-        cursor = self._db.conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT config, strftime('%s',updated_at) FROM models
-            WHERE id=?;
-            """,
-            (key,),
-        )
-        rows = cursor.fetchone()
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT config, strftime('%s',updated_at) FROM models
+                WHERE id=?;
+                """,
+                (key,),
+            )
+            rows = cursor.fetchone()
         if not rows:
             raise UnknownModelException("model not found")
         model = ModelConfigFactory.make_config(json.loads(rows[0]), timestamp=rows[1])
         return model
 
     def get_model_by_hash(self, hash: str) -> AnyModelConfig:
-        cursor = self._db.conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT config, strftime('%s',updated_at) FROM models
-            WHERE hash=?;
-            """,
-            (hash,),
-        )
-        rows = cursor.fetchone()
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT config, strftime('%s',updated_at) FROM models
+                WHERE hash=?;
+                """,
+                (hash,),
+            )
+            rows = cursor.fetchone()
         if not rows:
             raise UnknownModelException("model not found")
         model = ModelConfigFactory.make_config(json.loads(rows[0]), timestamp=rows[1])
@@ -224,15 +204,15 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
 
         :param key: Unique key for the model to be deleted
         """
-        cursor = self._db.conn.cursor()
-        cursor.execute(
-            """--sql
-            select count(*) FROM models
-            WHERE id=?;
-            """,
-            (key,),
-        )
-        count = cursor.fetchone()[0]
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                select count(*) FROM models
+                WHERE id=?;
+                """,
+                (key,),
+            )
+            count = cursor.fetchone()[0]
         return count > 0
 
     def search_by_attr(
@@ -255,43 +235,42 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
         If none of the optional filters are passed, will return all
         models in the database.
         """
+        with self._db.transaction() as cursor:
+            assert isinstance(order_by, ModelRecordOrderBy)
+            ordering = {
+                ModelRecordOrderBy.Default: "type, base, name, format",
+                ModelRecordOrderBy.Type: "type",
+                ModelRecordOrderBy.Base: "base",
+                ModelRecordOrderBy.Name: "name",
+                ModelRecordOrderBy.Format: "format",
+            }
 
-        assert isinstance(order_by, ModelRecordOrderBy)
-        ordering = {
-            ModelRecordOrderBy.Default: "type, base, name, format",
-            ModelRecordOrderBy.Type: "type",
-            ModelRecordOrderBy.Base: "base",
-            ModelRecordOrderBy.Name: "name",
-            ModelRecordOrderBy.Format: "format",
-        }
+            where_clause: list[str] = []
+            bindings: list[str] = []
+            if model_name:
+                where_clause.append("name=?")
+                bindings.append(model_name)
+            if base_model:
+                where_clause.append("base=?")
+                bindings.append(base_model)
+            if model_type:
+                where_clause.append("type=?")
+                bindings.append(model_type)
+            if model_format:
+                where_clause.append("format=?")
+                bindings.append(model_format)
+            where = f"WHERE {' AND '.join(where_clause)}" if where_clause else ""
 
-        where_clause: list[str] = []
-        bindings: list[str] = []
-        if model_name:
-            where_clause.append("name=?")
-            bindings.append(model_name)
-        if base_model:
-            where_clause.append("base=?")
-            bindings.append(base_model)
-        if model_type:
-            where_clause.append("type=?")
-            bindings.append(model_type)
-        if model_format:
-            where_clause.append("format=?")
-            bindings.append(model_format)
-        where = f"WHERE {' AND '.join(where_clause)}" if where_clause else ""
-
-        cursor = self._db.conn.cursor()
-        cursor.execute(
-            f"""--sql
-            SELECT config, strftime('%s',updated_at)
-            FROM models
-            {where}
-            ORDER BY {ordering[order_by]} -- using ? to bind doesn't work here for some reason;
-            """,
-            tuple(bindings),
-        )
-        result = cursor.fetchall()
+            cursor.execute(
+                f"""--sql
+                SELECT config, strftime('%s',updated_at)
+                FROM models
+                {where}
+                ORDER BY {ordering[order_by]} -- using ? to bind doesn't work here for some reason;
+                """,
+                tuple(bindings),
+            )
+            result = cursor.fetchall()
 
         # Parse the model configs.
         results: list[AnyModelConfig] = []
@@ -313,69 +292,68 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
 
     def search_by_path(self, path: Union[str, Path]) -> List[AnyModelConfig]:
         """Return models with the indicated path."""
-        cursor = self._db.conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT config, strftime('%s',updated_at) FROM models
-            WHERE path=?;
-            """,
-            (str(path),),
-        )
-        results = [ModelConfigFactory.make_config(json.loads(x[0]), timestamp=x[1]) for x in cursor.fetchall()]
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT config, strftime('%s',updated_at) FROM models
+                WHERE path=?;
+                """,
+                (str(path),),
+            )
+            results = [ModelConfigFactory.make_config(json.loads(x[0]), timestamp=x[1]) for x in cursor.fetchall()]
         return results
 
     def search_by_hash(self, hash: str) -> List[AnyModelConfig]:
         """Return models with the indicated hash."""
-        cursor = self._db.conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT config, strftime('%s',updated_at) FROM models
-            WHERE hash=?;
-            """,
-            (hash,),
-        )
-        results = [ModelConfigFactory.make_config(json.loads(x[0]), timestamp=x[1]) for x in cursor.fetchall()]
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT config, strftime('%s',updated_at) FROM models
+                WHERE hash=?;
+                """,
+                (hash,),
+            )
+            results = [ModelConfigFactory.make_config(json.loads(x[0]), timestamp=x[1]) for x in cursor.fetchall()]
         return results
 
     def list_models(
         self, page: int = 0, per_page: int = 10, order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default
     ) -> PaginatedResults[ModelSummary]:
         """Return a paginated summary listing of each model in the database."""
-        assert isinstance(order_by, ModelRecordOrderBy)
-        ordering = {
-            ModelRecordOrderBy.Default: "type, base, name, format",
-            ModelRecordOrderBy.Type: "type",
-            ModelRecordOrderBy.Base: "base",
-            ModelRecordOrderBy.Name: "name",
-            ModelRecordOrderBy.Format: "format",
-        }
+        with self._db.transaction() as cursor:
+            assert isinstance(order_by, ModelRecordOrderBy)
+            ordering = {
+                ModelRecordOrderBy.Default: "type, base, name, format",
+                ModelRecordOrderBy.Type: "type",
+                ModelRecordOrderBy.Base: "base",
+                ModelRecordOrderBy.Name: "name",
+                ModelRecordOrderBy.Format: "format",
+            }
 
-        cursor = self._db.conn.cursor()
+            # Lock so that the database isn't updated while we're doing the two queries.
+            # query1: get the total number of model configs
+            cursor.execute(
+                """--sql
+                select count(*) from models;
+                """,
+                (),
+            )
+            total = int(cursor.fetchone()[0])
 
-        # Lock so that the database isn't updated while we're doing the two queries.
-        # query1: get the total number of model configs
-        cursor.execute(
-            """--sql
-            select count(*) from models;
-            """,
-            (),
-        )
-        total = int(cursor.fetchone()[0])
-
-        # query2: fetch key fields
-        cursor.execute(
-            f"""--sql
-            SELECT config
-            FROM models
-            ORDER BY {ordering[order_by]} -- using ? to bind doesn't work here for some reason
-            LIMIT ?
-            OFFSET ?;
-            """,
-            (
-                per_page,
-                page * per_page,
-            ),
-        )
-        rows = cursor.fetchall()
+            # query2: fetch key fields
+            cursor.execute(
+                f"""--sql
+                SELECT config
+                FROM models
+                ORDER BY {ordering[order_by]} -- using ? to bind doesn't work here for some reason
+                LIMIT ?
+                OFFSET ?;
+                """,
+                (
+                    per_page,
+                    page * per_page,
+                ),
+            )
+            rows = cursor.fetchall()
         items = [ModelSummary.model_validate(dict(x)) for x in rows]
         return PaginatedResults(page=page, pages=ceil(total / per_page), per_page=per_page, total=total, items=items)

invokeai/app/services/model_relationship_records/model_relationship_records_base.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+
+
+class ModelRelationshipRecordStorageBase(ABC):
+    """Abstract base class for model-to-model relationship record storage."""
+
+    @abstractmethod
+    def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        """Creates a relationship between two models by keys."""
+        pass
+
+    @abstractmethod
+    def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        """Removes a relationship between two models by keys."""
+        pass
+
+    @abstractmethod
+    def get_related_model_keys(self, model_key: str) -> list[str]:
+        """Gets all models keys related to a given model key."""
+        pass
+
+    @abstractmethod
+    def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
+        """Get related model keys for multiple models given a list of keys."""
+        pass

invokeai/app/services/model_relationship_records/model_relationship_records_sqlite.py

@@ -0,0 +1,55 @@
+from invokeai.app.services.model_relationship_records.model_relationship_records_base import (
+    ModelRelationshipRecordStorageBase,
+)
+from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+
+
+class SqliteModelRelationshipRecordStorage(ModelRelationshipRecordStorageBase):
+    def __init__(self, db: SqliteDatabase) -> None:
+        super().__init__()
+        self._db = db
+
+    def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        with self._db.transaction() as cursor:
+            if model_key_1 == model_key_2:
+                raise ValueError("Cannot relate a model to itself.")
+            a, b = sorted([model_key_1, model_key_2])
+            cursor.execute(
+                "INSERT OR IGNORE INTO model_relationships (model_key_1, model_key_2) VALUES (?, ?)",
+                (a, b),
+            )
+
+    def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        with self._db.transaction() as cursor:
+            a, b = sorted([model_key_1, model_key_2])
+            cursor.execute(
+                "DELETE FROM model_relationships WHERE model_key_1 = ? AND model_key_2 = ?",
+                (a, b),
+            )
+
+    def get_related_model_keys(self, model_key: str) -> list[str]:
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """
+                SELECT model_key_2 FROM model_relationships WHERE model_key_1 = ?
+                UNION
+                SELECT model_key_1 FROM model_relationships WHERE model_key_2 = ?
+                """,
+                (model_key, model_key),
+            )
+            result = [row[0] for row in cursor.fetchall()]
+        return result
+
+    def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
+        with self._db.transaction() as cursor:
+            key_list = ",".join("?" for _ in model_keys)
+            cursor.execute(
+                f"""
+                SELECT model_key_2 FROM model_relationships WHERE model_key_1 IN ({key_list})
+                UNION
+                SELECT model_key_1 FROM model_relationships WHERE model_key_2 IN ({key_list})
+                """,
+                model_keys + model_keys,
+            )
+            result = [row[0] for row in cursor.fetchall()]
+        return result

invokeai/app/services/model_relationships/model_relationships_base.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+
+
+class ModelRelationshipsServiceABC(ABC):
+    """High-level service for managing model-to-model relationships."""
+
+    @abstractmethod
+    def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        """Creates a relationship between two models keys."""
+        pass
+
+    @abstractmethod
+    def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        """Removes a relationship between two models keys."""
+        pass
+
+    @abstractmethod
+    def get_related_model_keys(self, model_key: str) -> list[str]:
+        """Gets all models keys related to a given model key."""
+        pass
+
+    @abstractmethod
+    def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
+        """Get related model keys for multiple models."""
+        pass

invokeai/app/services/model_relationships/model_relationships_common.py

@@ -0,0 +1,9 @@
+from datetime import datetime
+
+from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
+
+
+class ModelRelationship(BaseModelExcludeNull):
+    model_key_1: str
+    model_key_2: str
+    created_at: datetime

invokeai/app/services/model_relationships/model_relationships_default.py

@@ -0,0 +1,31 @@
+from invokeai.app.services.invoker import Invoker
+from invokeai.app.services.model_relationships.model_relationships_base import ModelRelationshipsServiceABC
+from invokeai.backend.model_manager.config import AnyModelConfig
+
+
+class ModelRelationshipsService(ModelRelationshipsServiceABC):
+    __invoker: Invoker
+
+    def start(self, invoker: Invoker) -> None:
+        self.__invoker = invoker
+
+    def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        self.__invoker.services.model_relationship_records.add_model_relationship(model_key_1, model_key_2)
+
+    def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        self.__invoker.services.model_relationship_records.remove_model_relationship(model_key_1, model_key_2)
+
+    def get_related_model_keys(self, model_key: str) -> list[str]:
+        return self.__invoker.services.model_relationship_records.get_related_model_keys(model_key)
+
+    def add_relationship_from_models(self, model_1: AnyModelConfig, model_2: AnyModelConfig) -> None:
+        self.add_model_relationship(model_1.key, model_2.key)
+
+    def remove_relationship_from_models(self, model_1: AnyModelConfig, model_2: AnyModelConfig) -> None:
+        self.remove_model_relationship(model_1.key, model_2.key)
+
+    def get_related_keys_from_model(self, model: AnyModelConfig) -> list[str]:
+        return self.get_related_model_keys(model.key)
+
+    def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
+        return self.__invoker.services.model_relationship_records.get_related_model_keys_batch(model_keys)

invokeai/app/services/session_processor/session_processor_default.py

@@ -1,3 +1,4 @@
+import gc
 import traceback
 from contextlib import suppress
 from threading import BoundedSemaphore, Thread
@@ -210,7 +211,7 @@ class DefaultSessionRunner(SessionRunnerBase):
             # we don't care about that - suppress the error.
             with suppress(GESStatsNotFoundError):
                 self._services.performance_statistics.log_stats(queue_item.session.id)
-                self._services.performance_statistics.reset_stats()
+                self._services.performance_statistics.reset_stats(queue_item.session.id)
 
             for callback in self._on_after_run_session_callbacks:
                 callback(queue_item=queue_item)
@@ -439,6 +440,12 @@ class DefaultSessionProcessor(SessionProcessorBase):
                         poll_now_event.wait(self._polling_interval)
                         continue
 
+                    # GC-ing here can reduce peak memory usage of the invoke process by freeing allocated memory blocks.
+                    # Most queue items take seconds to execute, so the relative cost of a GC is very small.
+                    # Python will never cede allocated memory back to the OS, so anything we can do to reduce the peak
+                    # allocation is well worth it.
+                    gc.collect()
+
                     self._invoker.services.logger.info(
                         f"Executing queue item {self._queue_item.item_id}, session {self._queue_item.session_id}"
                     )

invokeai/app/services/session_queue/session_queue_base.py

@@ -10,6 +10,8 @@ from invokeai.app.services.session_queue.session_queue_common import (
     CancelByDestinationResult,
     CancelByQueueIDResult,
     ClearResult,
+    DeleteAllExceptCurrentResult,
+    DeleteByDestinationResult,
     EnqueueBatchResult,
     IsEmptyResult,
     IsFullResult,
@@ -17,7 +19,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     RetryItemsResult,
     SessionQueueCountsByDestination,
     SessionQueueItem,
-    SessionQueueItemDTO,
     SessionQueueStatus,
 )
 from invokeai.app.services.shared.graph import GraphExecutionState
@@ -92,6 +93,11 @@ class SessionQueueBase(ABC):
         """Cancels a session queue item"""
         pass
 
+    @abstractmethod
+    def delete_queue_item(self, item_id: int) -> None:
+        """Deletes a session queue item"""
+        pass
+
     @abstractmethod
     def fail_queue_item(
         self, item_id: int, error_type: str, error_message: str, error_traceback: str
@@ -109,6 +115,11 @@ class SessionQueueBase(ABC):
         """Cancels all queue items with the given batch destination"""
         pass
 
+    @abstractmethod
+    def delete_by_destination(self, queue_id: str, destination: str) -> DeleteByDestinationResult:
+        """Deletes all queue items with the given batch destination"""
+        pass
+
     @abstractmethod
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
         """Cancels all queue items with matching queue ID"""
@@ -119,6 +130,11 @@ class SessionQueueBase(ABC):
         """Cancels all queue items except in-progress items"""
         pass
 
+    @abstractmethod
+    def delete_all_except_current(self, queue_id: str) -> DeleteAllExceptCurrentResult:
+        """Deletes all queue items except in-progress items"""
+        pass
+
     @abstractmethod
     def list_queue_items(
         self,
@@ -127,10 +143,20 @@ class SessionQueueBase(ABC):
         priority: int,
         cursor: Optional[int] = None,
         status: Optional[QUEUE_ITEM_STATUS] = None,
-    ) -> CursorPaginatedResults[SessionQueueItemDTO]:
+        destination: Optional[str] = None,
+    ) -> CursorPaginatedResults[SessionQueueItem]:
         """Gets a page of session queue items"""
         pass
 
+    @abstractmethod
+    def list_all_queue_items(
+        self,
+        queue_id: str,
+        destination: Optional[str] = None,
+    ) -> list[SessionQueueItem]:
+        """Gets all queue items that match the given parameters"""
+        pass
+
     @abstractmethod
     def get_queue_item(self, item_id: int) -> SessionQueueItem:
         """Gets a session queue item by ID"""

invokeai/app/services/session_queue/session_queue_common.py

@@ -148,7 +148,7 @@ class Batch(BaseModel):
                     node = cast(BaseInvocation, graph.get_node(batch_data.node_path))
                 except NodeNotFoundError:
                     raise NodeNotFoundError(f"Node {batch_data.node_path} not found in graph")
-                if batch_data.field_name not in node.model_fields:
+                if batch_data.field_name not in type(node).model_fields:
                     raise NodeNotFoundError(f"Field {batch_data.field_name} not found in node {batch_data.node_path}")
         return values
 
@@ -205,9 +205,10 @@ class FieldIdentifier(BaseModel):
     kind: Literal["input", "output"] = Field(description="The kind of field")
     node_id: str = Field(description="The ID of the node")
     field_name: str = Field(description="The name of the field")
+    user_label: str | None = Field(description="The user label of the field, if any")
 
 
-class SessionQueueItemWithoutGraph(BaseModel):
+class SessionQueueItem(BaseModel):
     """Session queue item without the full graph. Used for serialization."""
 
     item_id: int = Field(description="The identifier of the session queue item")
@@ -251,41 +252,7 @@ class SessionQueueItemWithoutGraph(BaseModel):
         default=None,
         description="The ID of the published workflow associated with this queue item",
     )
-    api_input_fields: Optional[list[FieldIdentifier]] = Field(
-        default=None, description="The fields that were used as input to the API"
-    )
-    api_output_fields: Optional[list[FieldIdentifier]] = Field(
-        default=None, description="The nodes that were used as output from the API"
-    )
-
-    @classmethod
-    def queue_item_dto_from_dict(cls, queue_item_dict: dict) -> "SessionQueueItemDTO":
-        # must parse these manually
-        queue_item_dict["field_values"] = get_field_values(queue_item_dict)
-        return SessionQueueItemDTO(**queue_item_dict)
-
-    model_config = ConfigDict(
-        json_schema_extra={
-            "required": [
-                "item_id",
-                "status",
-                "batch_id",
-                "queue_id",
-                "session_id",
-                "priority",
-                "session_id",
-                "created_at",
-                "updated_at",
-            ]
-        }
-    )
-
-
-class SessionQueueItemDTO(SessionQueueItemWithoutGraph):
-    pass
-
-
-class SessionQueueItem(SessionQueueItemWithoutGraph):
+    credits: Optional[float] = Field(default=None, description="The total credits used for this queue item")
     session: GraphExecutionState = Field(description="The fully-populated session to be executed")
     workflow: Optional[WorkflowWithoutID] = Field(
         default=None, description="The workflow associated with this queue item"
@@ -365,6 +332,7 @@ class EnqueueBatchResult(BaseModel):
     requested: int = Field(description="The total number of queue items requested to be enqueued")
     batch: Batch = Field(description="The batch that was enqueued")
     priority: int = Field(description="The priority of the enqueued batch")
+    item_ids: list[int] = Field(description="The IDs of the queue items that were enqueued")
 
 
 class RetryItemsResult(BaseModel):
@@ -396,6 +364,18 @@ class CancelByDestinationResult(CancelByBatchIDsResult):
     pass
 
 
+class DeleteByDestinationResult(BaseModel):
+    """Result of deleting by a destination"""
+
+    deleted: int = Field(..., description="Number of queue items deleted")
+
+
+class DeleteAllExceptCurrentResult(DeleteByDestinationResult):
+    """Result of deleting all except current"""
+
+    pass
+
+
 class CancelByQueueIDResult(CancelByBatchIDsResult):
     """Result of canceling by queue id"""
 

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -17,6 +17,8 @@ from invokeai.app.services.session_queue.session_queue_common import (
     CancelByDestinationResult,
     CancelByQueueIDResult,
     ClearResult,
+    DeleteAllExceptCurrentResult,
+    DeleteByDestinationResult,
     EnqueueBatchResult,
     IsEmptyResult,
     IsFullResult,
@@ -24,7 +26,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     RetryItemsResult,
     SessionQueueCountsByDestination,
     SessionQueueItem,
-    SessionQueueItemDTO,
     SessionQueueItemNotFoundError,
     SessionQueueStatus,
     ValueToInsertTuple,
@@ -46,22 +47,17 @@ class SqliteSessionQueue(SessionQueueBase):
             clear_result = self.clear(DEFAULT_QUEUE_ID)
             if clear_result.deleted > 0:
                 self.__invoker.services.logger.info(f"Cleared all {clear_result.deleted} queue items")
-        else:
-            prune_result = self.prune(DEFAULT_QUEUE_ID)
-            if prune_result.deleted > 0:
-                self.__invoker.services.logger.info(f"Pruned {prune_result.deleted} finished queue items")
 
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
-        self._conn = db.conn
+        self._db = db
 
     def _set_in_progress_to_canceled(self) -> None:
         """
         Sets all in_progress queue items to canceled. Run on app startup, not associated with any queue.
         This is necessary because the invoker may have been killed while processing a queue item.
         """
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 UPDATE session_queue
@@ -69,102 +65,104 @@ class SqliteSessionQueue(SessionQueueBase):
                 WHERE status = 'in_progress';
                 """
             )
-        except Exception:
-            self._conn.rollback()
-            raise
 
     def _get_current_queue_size(self, queue_id: str) -> int:
         """Gets the current number of pending queue items"""
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT count(*)
-            FROM session_queue
-            WHERE
-              queue_id = ?
-              AND status = 'pending'
-            """,
-            (queue_id,),
-        )
-        return cast(int, cursor.fetchone()[0])
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT count(*)
+                FROM session_queue
+                WHERE
+                queue_id = ?
+                AND status = 'pending'
+                """,
+                (queue_id,),
+            )
+            count = cast(int, cursor.fetchone()[0])
+        return count
 
     def _get_highest_priority(self, queue_id: str) -> int:
         """Gets the highest priority value in the queue"""
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT MAX(priority)
-            FROM session_queue
-            WHERE
-              queue_id = ?
-              AND status = 'pending'
-            """,
-            (queue_id,),
-        )
-        return cast(Union[int, None], cursor.fetchone()[0]) or 0
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT MAX(priority)
+                FROM session_queue
+                WHERE
+                queue_id = ?
+                AND status = 'pending'
+                """,
+                (queue_id,),
+            )
+            priority = cast(Union[int, None], cursor.fetchone()[0]) or 0
+        return priority
 
     async def enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool) -> EnqueueBatchResult:
-        return await asyncio.to_thread(self._enqueue_batch, queue_id, batch, prepend)
+        current_queue_size = self._get_current_queue_size(queue_id)
+        max_queue_size = self.__invoker.services.configuration.max_queue_size
+        max_new_queue_items = max_queue_size - current_queue_size
 
-    def _enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool) -> EnqueueBatchResult:
-        try:
-            cursor = self._conn.cursor()
-            # TODO: how does this work in a multi-user scenario?
-            current_queue_size = self._get_current_queue_size(queue_id)
-            max_queue_size = self.__invoker.services.configuration.max_queue_size
-            max_new_queue_items = max_queue_size - current_queue_size
+        priority = 0
+        if prepend:
+            priority = self._get_highest_priority(queue_id) + 1
 
-            priority = 0
-            if prepend:
-                priority = self._get_highest_priority(queue_id) + 1
-
-            requested_count = calc_session_count(batch)
-            values_to_insert = prepare_values_to_insert(
-                queue_id=queue_id,
-                batch=batch,
-                priority=priority,
-                max_new_queue_items=max_new_queue_items,
-            )
-            enqueued_count = len(values_to_insert)
-
-            if requested_count > enqueued_count:
-                values_to_insert = values_to_insert[:max_new_queue_items]
+        requested_count = await asyncio.to_thread(
+            calc_session_count,
+            batch=batch,
+        )
+        values_to_insert = await asyncio.to_thread(
+            prepare_values_to_insert,
+            queue_id=queue_id,
+            batch=batch,
+            priority=priority,
+            max_new_queue_items=max_new_queue_items,
+        )
+        enqueued_count = len(values_to_insert)
 
+        with self._db.transaction() as cursor:
             cursor.executemany(
                 """--sql
-                INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                """,
+                    INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
                 values_to_insert,
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
+            cursor.execute(
+                """--sql
+                    SELECT item_id
+                    FROM session_queue
+                    WHERE batch_id = ?
+                    ORDER BY item_id DESC;
+                    """,
+                (batch.batch_id,),
+            )
+            item_ids = [row[0] for row in cursor.fetchall()]
         enqueue_result = EnqueueBatchResult(
             queue_id=queue_id,
             requested=requested_count,
             enqueued=enqueued_count,
             batch=batch,
             priority=priority,
+            item_ids=item_ids,
         )
         self.__invoker.services.events.emit_batch_enqueued(enqueue_result)
         return enqueue_result
 
     def dequeue(self) -> Optional[SessionQueueItem]:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT *
-            FROM session_queue
-            WHERE status = 'pending'
-            ORDER BY
-                priority DESC,
-                item_id ASC
-            LIMIT 1
-            """
-        )
-        result = cast(Union[sqlite3.Row, None], cursor.fetchone())
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT *
+                FROM session_queue
+                WHERE status = 'pending'
+                ORDER BY
+                    priority DESC,
+                    item_id ASC
+                LIMIT 1
+                """
+            )
+            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
         if result is None:
             return None
         queue_item = SessionQueueItem.queue_item_from_dict(dict(result))
@@ -172,40 +170,40 @@ class SqliteSessionQueue(SessionQueueBase):
         return queue_item
 
     def get_next(self, queue_id: str) -> Optional[SessionQueueItem]:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT *
-            FROM session_queue
-            WHERE
-                queue_id = ?
-                AND status = 'pending'
-            ORDER BY
-                priority DESC,
-                created_at ASC
-            LIMIT 1
-            """,
-            (queue_id,),
-        )
-        result = cast(Union[sqlite3.Row, None], cursor.fetchone())
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT *
+                FROM session_queue
+                WHERE
+                    queue_id = ?
+                    AND status = 'pending'
+                ORDER BY
+                    priority DESC,
+                    created_at ASC
+                LIMIT 1
+                """,
+                (queue_id,),
+            )
+            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
         if result is None:
             return None
         return SessionQueueItem.queue_item_from_dict(dict(result))
 
     def get_current(self, queue_id: str) -> Optional[SessionQueueItem]:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT *
-            FROM session_queue
-            WHERE
-                queue_id = ?
-                AND status = 'in_progress'
-            LIMIT 1
-            """,
-            (queue_id,),
-        )
-        result = cast(Union[sqlite3.Row, None], cursor.fetchone())
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT *
+                FROM session_queue
+                WHERE
+                    queue_id = ?
+                    AND status = 'in_progress'
+                LIMIT 1
+                """,
+                (queue_id,),
+            )
+            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
         if result is None:
             return None
         return SessionQueueItem.queue_item_from_dict(dict(result))
@@ -218,8 +216,23 @@ class SqliteSessionQueue(SessionQueueBase):
         error_message: Optional[str] = None,
         error_traceback: Optional[str] = None,
     ) -> SessionQueueItem:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT status FROM session_queue WHERE item_id = ?
+                """,
+                (item_id,),
+            )
+            row = cursor.fetchone()
+        if row is None:
+            raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}")
+        current_status = row[0]
+
+        # Only update if not already finished (completed, failed or canceled)
+        if current_status in ("completed", "failed", "canceled"):
+            return self.get_queue_item(item_id)
+
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 UPDATE session_queue
@@ -228,10 +241,7 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 (status, error_type, error_message, error_traceback, item_id),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
+
         queue_item = self.get_queue_item(item_id)
         batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id)
         queue_status = self.get_queue_status(queue_id=queue_item.queue_id)
@@ -239,35 +249,34 @@ class SqliteSessionQueue(SessionQueueBase):
         return queue_item
 
     def is_empty(self, queue_id: str) -> IsEmptyResult:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT count(*)
-            FROM session_queue
-            WHERE queue_id = ?
-            """,
-            (queue_id,),
-        )
-        is_empty = cast(int, cursor.fetchone()[0]) == 0
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT count(*)
+                FROM session_queue
+                WHERE queue_id = ?
+                """,
+                (queue_id,),
+            )
+            is_empty = cast(int, cursor.fetchone()[0]) == 0
         return IsEmptyResult(is_empty=is_empty)
 
     def is_full(self, queue_id: str) -> IsFullResult:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT count(*)
-            FROM session_queue
-            WHERE queue_id = ?
-            """,
-            (queue_id,),
-        )
-        max_queue_size = self.__invoker.services.configuration.max_queue_size
-        is_full = cast(int, cursor.fetchone()[0]) >= max_queue_size
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT count(*)
+                FROM session_queue
+                WHERE queue_id = ?
+                """,
+                (queue_id,),
+            )
+            max_queue_size = self.__invoker.services.configuration.max_queue_size
+            is_full = cast(int, cursor.fetchone()[0]) >= max_queue_size
         return IsFullResult(is_full=is_full)
 
     def clear(self, queue_id: str) -> ClearResult:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT COUNT(*)
@@ -285,24 +294,19 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 (queue_id,),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         self.__invoker.services.events.emit_queue_cleared(queue_id)
         return ClearResult(deleted=count)
 
     def prune(self, queue_id: str) -> PruneResult:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             where = """--sql
                 WHERE
-                  queue_id = ?
-                  AND (
+                queue_id = ?
+                AND (
                     status = 'completed'
                     OR status = 'failed'
                     OR status = 'canceled'
-                  )
+                )
                 """
             cursor.execute(
                 f"""--sql
@@ -321,16 +325,28 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 (queue_id,),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return PruneResult(deleted=count)
 
     def cancel_queue_item(self, item_id: int) -> SessionQueueItem:
         queue_item = self._set_queue_item_status(item_id=item_id, status="canceled")
         return queue_item
 
+    def delete_queue_item(self, item_id: int) -> None:
+        """Deletes a session queue item"""
+        try:
+            self.cancel_queue_item(item_id)
+        except SessionQueueItemNotFoundError:
+            pass
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                DELETE
+                FROM session_queue
+                WHERE item_id = ?
+                """,
+                (item_id,),
+            )
+
     def complete_queue_item(self, item_id: int) -> SessionQueueItem:
         queue_item = self._set_queue_item_status(item_id=item_id, status="completed")
         return queue_item
@@ -352,8 +368,7 @@ class SqliteSessionQueue(SessionQueueBase):
         return queue_item
 
     def cancel_by_batch_ids(self, queue_id: str, batch_ids: list[str]) -> CancelByBatchIDsResult:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             current_queue_item = self.get_current(queue_id)
             placeholders = ", ".join(["?" for _ in batch_ids])
             where = f"""--sql
@@ -363,6 +378,8 @@ class SqliteSessionQueue(SessionQueueBase):
                   AND status != 'canceled'
                   AND status != 'completed'
                   AND status != 'failed'
+                  -- We will cancel the current item separately below - skip it here
+                  AND status != 'in_progress'
                 """
             params = [queue_id] + batch_ids
             cursor.execute(
@@ -382,17 +399,14 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 tuple(params),
             )
-            self._conn.commit()
-            if current_queue_item is not None and current_queue_item.batch_id in batch_ids:
-                self._set_queue_item_status(current_queue_item.item_id, "canceled")
-        except Exception:
-            self._conn.rollback()
-            raise
+
+        if current_queue_item is not None and current_queue_item.batch_id in batch_ids:
+            self._set_queue_item_status(current_queue_item.item_id, "canceled")
+
         return CancelByBatchIDsResult(canceled=count)
 
     def cancel_by_destination(self, queue_id: str, destination: str) -> CancelByDestinationResult:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             current_queue_item = self.get_current(queue_id)
             where = """--sql
                 WHERE
@@ -401,6 +415,8 @@ class SqliteSessionQueue(SessionQueueBase):
                   AND status != 'canceled'
                   AND status != 'completed'
                   AND status != 'failed'
+                  -- We will cancel the current item separately below - skip it here
+                  AND status != 'in_progress'
                 """
             params = (queue_id, destination)
             cursor.execute(
@@ -420,17 +436,67 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 params,
             )
-            self._conn.commit()
-            if current_queue_item is not None and current_queue_item.destination == destination:
-                self._set_queue_item_status(current_queue_item.item_id, "canceled")
-        except Exception:
-            self._conn.rollback()
-            raise
+        if current_queue_item is not None and current_queue_item.destination == destination:
+            self._set_queue_item_status(current_queue_item.item_id, "canceled")
         return CancelByDestinationResult(canceled=count)
 
+    def delete_by_destination(self, queue_id: str, destination: str) -> DeleteByDestinationResult:
+        with self._db.transaction() as cursor:
+            current_queue_item = self.get_current(queue_id)
+            if current_queue_item is not None and current_queue_item.destination == destination:
+                self.cancel_queue_item(current_queue_item.item_id)
+            params = (queue_id, destination)
+            cursor.execute(
+                """--sql
+                SELECT COUNT(*)
+                FROM session_queue
+                WHERE
+                  queue_id = ?
+                  AND destination = ?;
+                """,
+                params,
+            )
+            count = cursor.fetchone()[0]
+            cursor.execute(
+                """--sql
+                DELETE
+                FROM session_queue
+                WHERE
+                  queue_id = ?
+                  AND destination = ?;
+                """,
+                params,
+            )
+        return DeleteByDestinationResult(deleted=count)
+
+    def delete_all_except_current(self, queue_id: str) -> DeleteAllExceptCurrentResult:
+        with self._db.transaction() as cursor:
+            where = """--sql
+                WHERE
+                  queue_id == ?
+                  AND status == 'pending'
+                """
+            cursor.execute(
+                f"""--sql
+                SELECT COUNT(*)
+                FROM session_queue
+                {where};
+                """,
+                (queue_id,),
+            )
+            count = cursor.fetchone()[0]
+            cursor.execute(
+                f"""--sql
+                DELETE
+                FROM session_queue
+                {where};
+                """,
+                (queue_id,),
+            )
+        return DeleteAllExceptCurrentResult(deleted=count)
+
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             current_queue_item = self.get_current(queue_id)
             where = """--sql
                 WHERE
@@ -438,6 +504,8 @@ class SqliteSessionQueue(SessionQueueBase):
                   AND status != 'canceled'
                   AND status != 'completed'
                   AND status != 'failed'
+                  -- We will cancel the current item separately below - skip it here
+                  AND status != 'in_progress'
                 """
             params = [queue_id]
             cursor.execute(
@@ -457,21 +525,13 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 tuple(params),
             )
-            self._conn.commit()
-            if current_queue_item is not None and current_queue_item.queue_id == queue_id:
-                batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id)
-                queue_status = self.get_queue_status(queue_id=queue_id)
-                self.__invoker.services.events.emit_queue_item_status_changed(
-                    current_queue_item, batch_status, queue_status
-                )
-        except Exception:
-            self._conn.rollback()
-            raise
+
+        if current_queue_item is not None and current_queue_item.queue_id == queue_id:
+            self._set_queue_item_status(current_queue_item.item_id, "canceled")
         return CancelByQueueIDResult(canceled=count)
 
     def cancel_all_except_current(self, queue_id: str) -> CancelAllExceptCurrentResult:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             where = """--sql
                 WHERE
                   queue_id == ?
@@ -494,30 +554,25 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 (queue_id,),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return CancelAllExceptCurrentResult(canceled=count)
 
     def get_queue_item(self, item_id: int) -> SessionQueueItem:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT * FROM session_queue
-            WHERE
-                item_id = ?
-            """,
-            (item_id,),
-        )
-        result = cast(Union[sqlite3.Row, None], cursor.fetchone())
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT * FROM session_queue
+                WHERE
+                    item_id = ?
+                """,
+                (item_id,),
+            )
+            result = cast(Union[sqlite3.Row, None], cursor.fetchone())
         if result is None:
             raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}")
         return SessionQueueItem.queue_item_from_dict(dict(result))
 
     def set_queue_item_session(self, item_id: int, session: GraphExecutionState) -> SessionQueueItem:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             # Use exclude_none so we don't end up with a bunch of nulls in the graph - this can cause validation errors
             # when the graph is loaded. Graph execution occurs purely in memory - the session saved here is not referenced
             # during execution.
@@ -530,10 +585,6 @@ class SqliteSessionQueue(SessionQueueBase):
                 """,
                 (session_json, item_id),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return self.get_queue_item(item_id)
 
     def list_queue_items(
@@ -543,53 +594,45 @@ class SqliteSessionQueue(SessionQueueBase):
         priority: int,
         cursor: Optional[int] = None,
         status: Optional[QUEUE_ITEM_STATUS] = None,
-    ) -> CursorPaginatedResults[SessionQueueItemDTO]:
-        cursor_ = self._conn.cursor()
-        item_id = cursor
-        query = """--sql
-            SELECT item_id,
-                status,
-                priority,
-                field_values,
-                error_type,
-                error_message,
-                error_traceback,
-                created_at,
-                updated_at,
-                completed_at,
-                started_at,
-                session_id,
-                batch_id,
-                queue_id,
-                origin,
-                destination
-            FROM session_queue
-            WHERE queue_id = ?
-        """
-        params: list[Union[str, int]] = [queue_id]
-
-        if status is not None:
-            query += """--sql
-                AND status = ?
-                """
-            params.append(status)
-
-        if item_id is not None:
-            query += """--sql
-                AND (priority < ?) OR (priority = ? AND item_id > ?)
-                """
-            params.extend([priority, priority, item_id])
-
-        query += """--sql
-            ORDER BY
-                priority DESC,
-                item_id ASC
-            LIMIT ?
+        destination: Optional[str] = None,
+    ) -> CursorPaginatedResults[SessionQueueItem]:
+        with self._db.transaction() as cursor_:
+            item_id = cursor
+            query = """--sql
+                SELECT *
+                FROM session_queue
+                WHERE queue_id = ?
             """
-        params.append(limit + 1)
-        cursor_.execute(query, params)
-        results = cast(list[sqlite3.Row], cursor_.fetchall())
-        items = [SessionQueueItemDTO.queue_item_dto_from_dict(dict(result)) for result in results]
+            params: list[Union[str, int]] = [queue_id]
+
+            if status is not None:
+                query += """--sql
+                    AND status = ?
+                    """
+                params.append(status)
+
+            if destination is not None:
+                query += """---sql
+                    AND destination = ?
+                """
+                params.append(destination)
+
+            if item_id is not None:
+                query += """--sql
+                    AND (priority < ?) OR (priority = ? AND item_id > ?)
+                    """
+                params.extend([priority, priority, item_id])
+
+            query += """--sql
+                ORDER BY
+                    priority DESC,
+                    item_id ASC
+                LIMIT ?
+                """
+            params.append(limit + 1)
+            cursor_.execute(query, params)
+            results = cast(list[sqlite3.Row], cursor_.fetchall())
+        items = [SessionQueueItem.queue_item_from_dict(dict(result)) for result in results]
         has_more = False
         if len(items) > limit:
             # remove the extra item
@@ -597,21 +640,52 @@ class SqliteSessionQueue(SessionQueueBase):
             has_more = True
         return CursorPaginatedResults(items=items, limit=limit, has_more=has_more)
 
+    def list_all_queue_items(
+        self,
+        queue_id: str,
+        destination: Optional[str] = None,
+    ) -> list[SessionQueueItem]:
+        """Gets all queue items that match the given parameters"""
+        with self._db.transaction() as cursor:
+            query = """--sql
+                SELECT *
+                FROM session_queue
+                WHERE queue_id = ?
+            """
+            params: list[Union[str, int]] = [queue_id]
+
+            if destination is not None:
+                query += """---sql
+                    AND destination = ?
+                """
+                params.append(destination)
+
+            query += """--sql
+                ORDER BY
+                    priority DESC,
+                    item_id ASC
+                ;
+                """
+            cursor.execute(query, params)
+            results = cast(list[sqlite3.Row], cursor.fetchall())
+        items = [SessionQueueItem.queue_item_from_dict(dict(result)) for result in results]
+        return items
+
     def get_queue_status(self, queue_id: str) -> SessionQueueStatus:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT status, count(*)
-            FROM session_queue
-            WHERE queue_id = ?
-            GROUP BY status
-            """,
-            (queue_id,),
-        )
-        counts_result = cast(list[sqlite3.Row], cursor.fetchall())
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT status, count(*)
+                FROM session_queue
+                WHERE queue_id = ?
+                GROUP BY status
+                """,
+                (queue_id,),
+            )
+            counts_result = cast(list[sqlite3.Row], cursor.fetchall())
 
         current_item = self.get_current(queue_id=queue_id)
-        total = sum(row[1] for row in counts_result)
+        total = sum(row[1] or 0 for row in counts_result)
         counts: dict[str, int] = {row[0]: row[1] for row in counts_result}
         return SessionQueueStatus(
             queue_id=queue_id,
@@ -627,20 +701,20 @@ class SqliteSessionQueue(SessionQueueBase):
         )
 
     def get_batch_status(self, queue_id: str, batch_id: str) -> BatchStatus:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT status, count(*), origin, destination
-            FROM session_queue
-            WHERE
-                queue_id = ?
-                AND batch_id = ?
-            GROUP BY status
-            """,
-            (queue_id, batch_id),
-        )
-        result = cast(list[sqlite3.Row], cursor.fetchall())
-        total = sum(row[1] for row in result)
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT status, count(*), origin, destination
+                FROM session_queue
+                WHERE
+                    queue_id = ?
+                    AND batch_id = ?
+                GROUP BY status
+                """,
+                (queue_id, batch_id),
+            )
+            result = cast(list[sqlite3.Row], cursor.fetchall())
+        total = sum(row[1] or 0 for row in result)
         counts: dict[str, int] = {row[0]: row[1] for row in result}
         origin = result[0]["origin"] if result else None
         destination = result[0]["destination"] if result else None
@@ -659,20 +733,20 @@ class SqliteSessionQueue(SessionQueueBase):
         )
 
     def get_counts_by_destination(self, queue_id: str, destination: str) -> SessionQueueCountsByDestination:
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT status, count(*)
-            FROM session_queue
-            WHERE queue_id = ?
-            AND destination = ?
-            GROUP BY status
-            """,
-            (queue_id, destination),
-        )
-        counts_result = cast(list[sqlite3.Row], cursor.fetchall())
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT status, count(*)
+                FROM session_queue
+                WHERE queue_id = ?
+                AND destination = ?
+                GROUP BY status
+                """,
+                (queue_id, destination),
+            )
+            counts_result = cast(list[sqlite3.Row], cursor.fetchall())
 
-        total = sum(row[1] for row in counts_result)
+        total = sum(row[1] or 0 for row in counts_result)
         counts: dict[str, int] = {row[0]: row[1] for row in counts_result}
 
         return SessionQueueCountsByDestination(
@@ -688,8 +762,7 @@ class SqliteSessionQueue(SessionQueueBase):
 
     def retry_items_by_id(self, queue_id: str, item_ids: list[int]) -> RetryItemsResult:
         """Retries the given queue items"""
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             values_to_insert: list[ValueToInsertTuple] = []
             retried_item_ids: list[int] = []
 
@@ -740,10 +813,6 @@ class SqliteSessionQueue(SessionQueueBase):
                 values_to_insert,
             )
 
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         retry_result = RetryItemsResult(
             queue_id=queue_id,
             retried_item_ids=retried_item_ids,

invokeai/app/services/shared/graph.py

@@ -2,11 +2,12 @@
 
 import copy
 import itertools
-from typing import Any, Optional, TypeVar, Union, get_args, get_origin, get_type_hints
+from typing import Any, Optional, TypeVar, Union, get_args, get_origin
 
 import networkx as nx
 from pydantic import (
     BaseModel,
+    ConfigDict,
     GetCoreSchemaHandler,
     GetJsonSchemaHandler,
     ValidationError,
@@ -57,17 +58,32 @@ class Edge(BaseModel):
 
 
 def get_output_field_type(node: BaseInvocation, field: str) -> Any:
-    node_type = type(node)
-    node_outputs = get_type_hints(node_type.get_output_annotation())
-    node_output_field = node_outputs.get(field) or None
-    return node_output_field
+    # TODO(psyche): This is awkward - if field_info is None, it means the field is not defined in the output, which
+    # really should raise. The consumers of this utility expect it to never raise, and return None instead. Fixing this
+    # would require some fairly significant changes and I don't want risk breaking anything.
+    try:
+        invocation_class = type(node)
+        invocation_output_class = invocation_class.get_output_annotation()
+        field_info = invocation_output_class.model_fields.get(field)
+        assert field_info is not None, f"Output field '{field}' not found in {invocation_output_class.get_type()}"
+        output_field_type = field_info.annotation
+        return output_field_type
+    except Exception:
+        return None
 
 
 def get_input_field_type(node: BaseInvocation, field: str) -> Any:
-    node_type = type(node)
-    node_inputs = get_type_hints(node_type)
-    node_input_field = node_inputs.get(field) or None
-    return node_input_field
+    # TODO(psyche): This is awkward - if field_info is None, it means the field is not defined in the output, which
+    # really should raise. The consumers of this utility expect it to never raise, and return None instead. Fixing this
+    # would require some fairly significant changes and I don't want risk breaking anything.
+    try:
+        invocation_class = type(node)
+        field_info = invocation_class.model_fields.get(field)
+        assert field_info is not None, f"Input field '{field}' not found in {invocation_class.get_type()}"
+        input_field_type = field_info.annotation
+        return input_field_type
+    except Exception:
+        return None
 
 
 def is_union_subtype(t1, t2):
@@ -424,7 +440,7 @@ class Graph(BaseModel):
                 )
 
             # input fields are on the node
-            if edge.destination.field not in destination_node.model_fields:
+            if edge.destination.field not in type(destination_node).model_fields:
                 raise NodeFieldNotFoundError(
                     f"Edge destination field {edge.destination.field} does not exist in node {edge.destination.node_id}"
                 )
@@ -787,6 +803,22 @@ class GraphExecutionState(BaseModel):
         default_factory=dict,
     )
 
+    model_config = ConfigDict(
+        json_schema_extra={
+            "required": [
+                "id",
+                "graph",
+                "execution_graph",
+                "executed",
+                "executed_history",
+                "results",
+                "errors",
+                "prepared_source_mapping",
+                "source_prepared_mapping",
+            ]
+        }
+    )
+
     @field_validator("graph")
     def graph_is_valid(cls, v: Graph):
         """Validates that the graph is valid"""
@@ -975,10 +1007,11 @@ class GraphExecutionState(BaseModel):
         new_node_ids = []
         if isinstance(next_node, CollectInvocation):
             # Collapse all iterator input mappings and create a single execution node for the collect invocation
-            all_iteration_mappings = list(
-                itertools.chain(*(((s, p) for p in self.source_prepared_mapping[s]) for s in next_node_parents))
-            )
-            # all_iteration_mappings = list(set(itertools.chain(*prepared_parent_mappings)))
+            all_iteration_mappings = []
+            for source_node_id in next_node_parents:
+                prepared_nodes = self.source_prepared_mapping[source_node_id]
+                all_iteration_mappings.extend([(source_node_id, p) for p in prepared_nodes])
+
             create_results = self._create_execution_node(next_node_id, all_iteration_mappings)
             if create_results is not None:
                 new_node_ids.extend(create_results)

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

@@ -1,4 +1,7 @@
 import sqlite3
+import threading
+from collections.abc import Generator
+from contextlib import contextmanager
 from logging import Logger
 from pathlib import Path
 
@@ -26,46 +29,65 @@ class SqliteDatabase:
 
     def __init__(self, db_path: Path | None, logger: Logger, verbose: bool = False) -> None:
         """Initializes the database. This is used internally by the class constructor."""
-        self.logger = logger
-        self.db_path = db_path
-        self.verbose = verbose
+        self._logger = logger
+        self._db_path = db_path
+        self._verbose = verbose
+        self._lock = threading.RLock()
 
-        if not self.db_path:
+        if not self._db_path:
             logger.info("Initializing in-memory database")
         else:
-            self.db_path.parent.mkdir(parents=True, exist_ok=True)
-            self.logger.info(f"Initializing database at {self.db_path}")
+            self._db_path.parent.mkdir(parents=True, exist_ok=True)
+            self._logger.info(f"Initializing database at {self._db_path}")
 
-        self.conn = sqlite3.connect(database=self.db_path or sqlite_memory, check_same_thread=False)
-        self.conn.row_factory = sqlite3.Row
+        self._conn = sqlite3.connect(database=self._db_path or sqlite_memory, check_same_thread=False)
+        self._conn.row_factory = sqlite3.Row
 
-        if self.verbose:
-            self.conn.set_trace_callback(self.logger.debug)
+        if self._verbose:
+            self._conn.set_trace_callback(self._logger.debug)
 
         # Enable foreign key constraints
-        self.conn.execute("PRAGMA foreign_keys = ON;")
+        self._conn.execute("PRAGMA foreign_keys = ON;")
 
         # Enable Write-Ahead Logging (WAL) mode for better concurrency
-        self.conn.execute("PRAGMA journal_mode = WAL;")
+        self._conn.execute("PRAGMA journal_mode = WAL;")
 
         # Set a busy timeout to prevent database lockups during writes
-        self.conn.execute("PRAGMA busy_timeout = 5000;")  # 5 seconds
+        self._conn.execute("PRAGMA busy_timeout = 5000;")  # 5 seconds
 
     def clean(self) -> None:
         """
         Cleans the database by running the VACUUM command, reporting on the freed space.
         """
         # No need to clean in-memory database
-        if not self.db_path:
+        if not self._db_path:
             return
         try:
-            initial_db_size = Path(self.db_path).stat().st_size
-            self.conn.execute("VACUUM;")
-            self.conn.commit()
-            final_db_size = Path(self.db_path).stat().st_size
-            freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2)
-            if freed_space_in_mb > 0:
-                self.logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)")
+            with self._conn as conn:
+                initial_db_size = Path(self._db_path).stat().st_size
+                conn.execute("VACUUM;")
+                conn.commit()
+                final_db_size = Path(self._db_path).stat().st_size
+                freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2)
+                if freed_space_in_mb > 0:
+                    self._logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)")
         except Exception as e:
-            self.logger.error(f"Error cleaning database: {e}")
+            self._logger.error(f"Error cleaning database: {e}")
             raise
+
+    @contextmanager
+    def transaction(self) -> Generator[sqlite3.Cursor, None, None]:
+        """
+        Thread-safe context manager for DB work.
+        Acquires the RLock, yields a Cursor, then commits or rolls back.
+        """
+        with self._lock:
+            cursor = self._conn.cursor()
+            try:
+                yield cursor
+                self._conn.commit()
+            except:
+                self._conn.rollback()
+                raise
+            finally:
+                cursor.close()

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

@@ -22,6 +22,8 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_16 import
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_17 import build_migration_17
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_18 import build_migration_18
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_19 import build_migration_19
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_20 import build_migration_20
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_21 import build_migration_21
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -61,6 +63,8 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_17())
     migrator.register_migration(build_migration_18())
     migrator.register_migration(build_migration_19(app_config=config))
+    migrator.register_migration(build_migration_20())
+    migrator.register_migration(build_migration_21())
     migrator.run_migrations()
 
     return db

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

@@ -0,0 +1,37 @@
+import sqlite3
+
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+
+
+class Migration20Callback:
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        cursor.execute(
+            """
+            -- many-to-many relationship table for models
+            CREATE TABLE IF NOT EXISTS model_relationships (
+                -- model_key_1 and model_key_2 are the same as the key(primary key) in the models table
+                model_key_1 TEXT NOT NULL,
+                model_key_2 TEXT NOT NULL,
+                created_at TEXT DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
+                PRIMARY KEY (model_key_1, model_key_2),
+                -- model_key_1 < model_key_2, to ensure uniqueness and prevent duplicates
+                FOREIGN KEY (model_key_1) REFERENCES models(id) ON DELETE CASCADE,
+                FOREIGN KEY (model_key_2) REFERENCES models(id) ON DELETE CASCADE
+            );
+            """
+        )
+        cursor.execute(
+            """
+            -- Creates an index to keep performance equal when searching for model_key_1 or model_key_2
+            CREATE INDEX IF NOT EXISTS keyx_model_relationships_model_key_2
+            ON model_relationships(model_key_2)
+            """
+        )
+
+
+def build_migration_20() -> Migration:
+    return Migration(
+        from_version=19,
+        to_version=20,
+        callback=Migration20Callback(),
+    )

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

@@ -0,0 +1,40 @@
+import sqlite3
+
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+
+
+class Migration21Callback:
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        cursor.execute(
+            """
+            CREATE TABLE client_state (
+              id          INTEGER PRIMARY KEY CHECK(id = 1),
+              data        TEXT    NOT NULL, -- Frontend will handle the shape of this data
+              updated_at  DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP)
+            );
+            """
+        )
+        cursor.execute(
+            """
+            CREATE TRIGGER tg_client_state_updated_at
+            AFTER UPDATE ON client_state
+            FOR EACH ROW
+            BEGIN
+              UPDATE client_state
+                SET updated_at = CURRENT_TIMESTAMP
+              WHERE id = OLD.id;
+            END;
+            """
+        )
+
+
+def build_migration_21() -> Migration:
+    """Builds the migration object for migrating from version 20 to version 21. This includes:
+    - Creating the `client_state` table.
+    - Adding a trigger to update the `updated_at` field on updates.
+    """
+    return Migration(
+        from_version=20,
+        to_version=21,
+        callback=Migration21Callback(),
+    )

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

@@ -32,7 +32,7 @@ class SqliteMigrator:
 
     def __init__(self, db: SqliteDatabase) -> None:
         self._db = db
-        self._logger = db.logger
+        self._logger = db._logger
         self._migration_set = MigrationSet()
         self._backup_path: Optional[Path] = None
 
@@ -45,7 +45,7 @@ class SqliteMigrator:
         """Migrates the database to the latest version."""
         # This throws if there is a problem.
         self._migration_set.validate_migration_chain()
-        cursor = self._db.conn.cursor()
+        cursor = self._db._conn.cursor()
         self._create_migrations_table(cursor=cursor)
 
         if self._migration_set.count == 0:
@@ -59,13 +59,13 @@ class SqliteMigrator:
         self._logger.info("Database update needed")
 
         # Make a backup of the db if it needs to be updated and is a file db
-        if self._db.db_path is not None:
+        if self._db._db_path is not None:
             timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
-            self._backup_path = self._db.db_path.parent / f"{self._db.db_path.stem}_backup_{timestamp}.db"
+            self._backup_path = self._db._db_path.parent / f"{self._db._db_path.stem}_backup_{timestamp}.db"
             self._logger.info(f"Backing up database to {str(self._backup_path)}")
             # Use SQLite to do the backup
             with closing(sqlite3.connect(self._backup_path)) as backup_conn:
-                self._db.conn.backup(backup_conn)
+                self._db._conn.backup(backup_conn)
         else:
             self._logger.info("Using in-memory database, no backup needed")
 
@@ -81,7 +81,7 @@ class SqliteMigrator:
         try:
             # Using sqlite3.Connection as a context manager commits a the transaction on exit, or rolls it back if an
             # exception is raised.
-            with self._db.conn as conn:
+            with self._db._conn as conn:
                 cursor = conn.cursor()
                 if self._get_current_version(cursor) != migration.from_version:
                     raise MigrationError(

invokeai/app/services/style_preset_records/style_preset_records_sqlite.py

@@ -17,7 +17,7 @@ from invokeai.app.util.misc import uuid_string
 class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase):
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
-        self._conn = db.conn
+        self._db = db
 
     def start(self, invoker: Invoker) -> None:
         self._invoker = invoker
@@ -25,24 +25,23 @@ class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase):
 
     def get(self, style_preset_id: str) -> StylePresetRecordDTO:
         """Gets a style preset by ID."""
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT *
-            FROM style_presets
-            WHERE id = ?;
-            """,
-            (style_preset_id,),
-        )
-        row = cursor.fetchone()
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT *
+                FROM style_presets
+                WHERE id = ?;
+                """,
+                (style_preset_id,),
+            )
+            row = cursor.fetchone()
         if row is None:
             raise StylePresetNotFoundError(f"Style preset with id {style_preset_id} not found")
         return StylePresetRecordDTO.from_dict(dict(row))
 
     def create(self, style_preset: StylePresetWithoutId) -> StylePresetRecordDTO:
         style_preset_id = uuid_string()
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 INSERT OR IGNORE INTO style_presets (
@@ -60,16 +59,11 @@ class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase):
                     style_preset.type,
                 ),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return self.get(style_preset_id)
 
     def create_many(self, style_presets: list[StylePresetWithoutId]) -> None:
         style_preset_ids = []
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             for style_preset in style_presets:
                 style_preset_id = uuid_string()
                 style_preset_ids.append(style_preset_id)
@@ -90,16 +84,11 @@ class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase):
                         style_preset.type,
                     ),
                 )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
 
         return None
 
     def update(self, style_preset_id: str, changes: StylePresetChanges) -> StylePresetRecordDTO:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             # Change the name of a style preset
             if changes.name is not None:
                 cursor.execute(
@@ -122,15 +111,10 @@ class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase):
                     (changes.preset_data.model_dump_json(), style_preset_id),
                 )
 
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return self.get(style_preset_id)
 
     def delete(self, style_preset_id: str) -> None:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 DELETE from style_presets
@@ -138,51 +122,41 @@ class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase):
                 """,
                 (style_preset_id,),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return None
 
     def get_many(self, type: PresetType | None = None) -> list[StylePresetRecordDTO]:
-        main_query = """
-            SELECT
-                *
-            FROM style_presets
-            """
+        with self._db.transaction() as cursor:
+            main_query = """
+                SELECT
+                    *
+                FROM style_presets
+                """
 
-        if type is not None:
-            main_query += "WHERE type = ? "
+            if type is not None:
+                main_query += "WHERE type = ? "
 
-        main_query += "ORDER BY LOWER(name) ASC"
+            main_query += "ORDER BY LOWER(name) ASC"
 
-        cursor = self._conn.cursor()
-        if type is not None:
-            cursor.execute(main_query, (type,))
-        else:
-            cursor.execute(main_query)
+            if type is not None:
+                cursor.execute(main_query, (type,))
+            else:
+                cursor.execute(main_query)
 
-        rows = cursor.fetchall()
+            rows = cursor.fetchall()
         style_presets = [StylePresetRecordDTO.from_dict(dict(row)) for row in rows]
 
         return style_presets
 
     def _sync_default_style_presets(self) -> None:
         """Syncs default style presets to the database. Internal use only."""
-
-        # First delete all existing default style presets
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
+            # First delete all existing default style presets
             cursor.execute(
                 """--sql
                 DELETE FROM style_presets
                 WHERE type = "default";
                 """
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         # Next, parse and create the default style presets
         with open(Path(__file__).parent / Path("default_style_presets.json"), "r") as file:
             presets = json.load(file)

invokeai/app/services/workflow_records/workflow_records_sqlite.py

@@ -25,7 +25,7 @@ SQL_TIME_FORMAT = "%Y-%m-%d %H:%M:%f"
 class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
-        self._conn = db.conn
+        self._db = db
 
     def start(self, invoker: Invoker) -> None:
         self._invoker = invoker
@@ -33,16 +33,16 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
 
     def get(self, workflow_id: str) -> WorkflowRecordDTO:
         """Gets a workflow by ID. Updates the opened_at column."""
-        cursor = self._conn.cursor()
-        cursor.execute(
-            """--sql
-            SELECT workflow_id, workflow, name, created_at, updated_at, opened_at
-            FROM workflow_library
-            WHERE workflow_id = ?;
-            """,
-            (workflow_id,),
-        )
-        row = cursor.fetchone()
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                SELECT workflow_id, workflow, name, created_at, updated_at, opened_at
+                FROM workflow_library
+                WHERE workflow_id = ?;
+                """,
+                (workflow_id,),
+            )
+            row = cursor.fetchone()
         if row is None:
             raise WorkflowNotFoundError(f"Workflow with id {workflow_id} not found")
         return WorkflowRecordDTO.from_dict(dict(row))
@@ -51,9 +51,8 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
         if workflow.meta.category is WorkflowCategory.Default:
             raise ValueError("Default workflows cannot be created via this method")
 
-        try:
+        with self._db.transaction() as cursor:
             workflow_with_id = Workflow(**workflow.model_dump(), id=uuid_string())
-            cursor = self._conn.cursor()
             cursor.execute(
                 """--sql
                 INSERT OR IGNORE INTO workflow_library (
@@ -64,18 +63,13 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
                 """,
                 (workflow_with_id.id, workflow_with_id.model_dump_json()),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return self.get(workflow_with_id.id)
 
     def update(self, workflow: Workflow) -> WorkflowRecordDTO:
         if workflow.meta.category is WorkflowCategory.Default:
             raise ValueError("Default workflows cannot be updated")
 
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 UPDATE workflow_library
@@ -84,18 +78,13 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
                 """,
                 (workflow.model_dump_json(), workflow.id),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return self.get(workflow.id)
 
     def delete(self, workflow_id: str) -> None:
         if self.get(workflow_id).workflow.meta.category is WorkflowCategory.Default:
             raise ValueError("Default workflows cannot be deleted")
 
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 DELETE from workflow_library
@@ -103,10 +92,6 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
                 """,
                 (workflow_id,),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
         return None
 
     def get_many(
@@ -121,108 +106,108 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
         has_been_opened: Optional[bool] = None,
         is_published: Optional[bool] = None,
     ) -> PaginatedResults[WorkflowRecordListItemDTO]:
-        # sanitize!
-        assert order_by in WorkflowRecordOrderBy
-        assert direction in SQLiteDirection
+        with self._db.transaction() as cursor:
+            # sanitize!
+            assert order_by in WorkflowRecordOrderBy
+            assert direction in SQLiteDirection
 
-        # We will construct the query dynamically based on the query params
+            # We will construct the query dynamically based on the query params
 
-        # The main query to get the workflows / counts
-        main_query = """
-                SELECT
-                    workflow_id,
-                    category,
-                    name,
-                    description,
-                    created_at,
-                    updated_at,
-                    opened_at,
-                    tags
-                FROM workflow_library
-                """
-        count_query = "SELECT COUNT(*) FROM workflow_library"
+            # The main query to get the workflows / counts
+            main_query = """
+                    SELECT
+                        workflow_id,
+                        category,
+                        name,
+                        description,
+                        created_at,
+                        updated_at,
+                        opened_at,
+                        tags
+                    FROM workflow_library
+                    """
+            count_query = "SELECT COUNT(*) FROM workflow_library"
 
-        # Start with an empty list of conditions and params
-        conditions: list[str] = []
-        params: list[str | int] = []
+            # Start with an empty list of conditions and params
+            conditions: list[str] = []
+            params: list[str | int] = []
 
-        if categories:
-            # Categories is a list of WorkflowCategory enum values, and a single string in the DB
+            if categories:
+                # Categories is a list of WorkflowCategory enum values, and a single string in the DB
 
-            # Ensure all categories are valid (is this necessary?)
-            assert all(c in WorkflowCategory for c in categories)
+                # Ensure all categories are valid (is this necessary?)
+                assert all(c in WorkflowCategory for c in categories)
 
-            # Construct a placeholder string for the number of categories
-            placeholders = ", ".join("?" for _ in categories)
+                # Construct a placeholder string for the number of categories
+                placeholders = ", ".join("?" for _ in categories)
 
-            # Construct the condition string & params
-            category_condition = f"category IN ({placeholders})"
-            category_params = [category.value for category in categories]
+                # Construct the condition string & params
+                category_condition = f"category IN ({placeholders})"
+                category_params = [category.value for category in categories]
 
-            conditions.append(category_condition)
-            params.extend(category_params)
+                conditions.append(category_condition)
+                params.extend(category_params)
 
-        if tags:
-            # Tags is a list of strings, and a single string in the DB
-            # The string in the DB has no guaranteed format
+            if tags:
+                # Tags is a list of strings, and a single string in the DB
+                # The string in the DB has no guaranteed format
 
-            # Construct a list of conditions for each tag
-            tags_conditions = ["tags LIKE ?" for _ in tags]
-            tags_conditions_joined = " OR ".join(tags_conditions)
-            tags_condition = f"({tags_conditions_joined})"
+                # Construct a list of conditions for each tag
+                tags_conditions = ["tags LIKE ?" for _ in tags]
+                tags_conditions_joined = " OR ".join(tags_conditions)
+                tags_condition = f"({tags_conditions_joined})"
 
-            # And the params for the tags, case-insensitive
-            tags_params = [f"%{t.strip()}%" for t in tags]
+                # And the params for the tags, case-insensitive
+                tags_params = [f"%{t.strip()}%" for t in tags]
 
-            conditions.append(tags_condition)
-            params.extend(tags_params)
+                conditions.append(tags_condition)
+                params.extend(tags_params)
 
-        if has_been_opened:
-            conditions.append("opened_at IS NOT NULL")
-        elif has_been_opened is False:
-            conditions.append("opened_at IS NULL")
+            if has_been_opened:
+                conditions.append("opened_at IS NOT NULL")
+            elif has_been_opened is False:
+                conditions.append("opened_at IS NULL")
 
-        # Ignore whitespace in the query
-        stripped_query = query.strip() if query else None
-        if stripped_query:
-            # Construct a wildcard query for the name, description, and tags
-            wildcard_query = "%" + stripped_query + "%"
-            query_condition = "(name LIKE ? OR description LIKE ? OR tags LIKE ?)"
+            # Ignore whitespace in the query
+            stripped_query = query.strip() if query else None
+            if stripped_query:
+                # Construct a wildcard query for the name, description, and tags
+                wildcard_query = "%" + stripped_query + "%"
+                query_condition = "(name LIKE ? OR description LIKE ? OR tags LIKE ?)"
 
-            conditions.append(query_condition)
-            params.extend([wildcard_query, wildcard_query, wildcard_query])
+                conditions.append(query_condition)
+                params.extend([wildcard_query, wildcard_query, wildcard_query])
 
-        if conditions:
-            # If there are conditions, add a WHERE clause and then join the conditions
-            main_query += " WHERE "
-            count_query += " WHERE "
+            if conditions:
+                # If there are conditions, add a WHERE clause and then join the conditions
+                main_query += " WHERE "
+                count_query += " WHERE "
 
-            all_conditions = " AND ".join(conditions)
-            main_query += all_conditions
-            count_query += all_conditions
+                all_conditions = " AND ".join(conditions)
+                main_query += all_conditions
+                count_query += all_conditions
 
-        # After this point, the query and params differ for the main query and the count query
-        main_params = params.copy()
-        count_params = params.copy()
+            # After this point, the query and params differ for the main query and the count query
+            main_params = params.copy()
+            count_params = params.copy()
 
-        # Main query also gets ORDER BY and LIMIT/OFFSET
-        main_query += f" ORDER BY {order_by.value} {direction.value}"
+            # Main query also gets ORDER BY and LIMIT/OFFSET
+            main_query += f" ORDER BY {order_by.value} {direction.value}"
 
-        if per_page:
-            main_query += " LIMIT ? OFFSET ?"
-            main_params.extend([per_page, page * per_page])
+            if per_page:
+                main_query += " LIMIT ? OFFSET ?"
+                main_params.extend([per_page, page * per_page])
 
-        # Put a ring on it
-        main_query += ";"
-        count_query += ";"
+            # Put a ring on it
+            main_query += ";"
+            count_query += ";"
 
-        cursor = self._conn.cursor()
-        cursor.execute(main_query, main_params)
-        rows = cursor.fetchall()
-        workflows = [WorkflowRecordListItemDTOValidator.validate_python(dict(row)) for row in rows]
+            cursor.execute(main_query, main_params)
+            rows = cursor.fetchall()
+            workflows = [WorkflowRecordListItemDTOValidator.validate_python(dict(row)) for row in rows]
 
-        cursor.execute(count_query, count_params)
-        total = cursor.fetchone()[0]
+            cursor.execute(count_query, count_params)
+            total = cursor.fetchone()[0]
 
         if per_page:
             pages = total // per_page + (total % per_page > 0)
@@ -247,46 +232,46 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
         if not tags:
             return {}
 
-        cursor = self._conn.cursor()
-        result: dict[str, int] = {}
-        # Base conditions for categories and selected tags
-        base_conditions: list[str] = []
-        base_params: list[str | int] = []
+        with self._db.transaction() as cursor:
+            result: dict[str, int] = {}
+            # Base conditions for categories and selected tags
+            base_conditions: list[str] = []
+            base_params: list[str | int] = []
 
-        # Add category conditions
-        if categories:
-            assert all(c in WorkflowCategory for c in categories)
-            placeholders = ", ".join("?" for _ in categories)
-            base_conditions.append(f"category IN ({placeholders})")
-            base_params.extend([category.value for category in categories])
+            # Add category conditions
+            if categories:
+                assert all(c in WorkflowCategory for c in categories)
+                placeholders = ", ".join("?" for _ in categories)
+                base_conditions.append(f"category IN ({placeholders})")
+                base_params.extend([category.value for category in categories])
 
-        if has_been_opened:
-            base_conditions.append("opened_at IS NOT NULL")
-        elif has_been_opened is False:
-            base_conditions.append("opened_at IS NULL")
+            if has_been_opened:
+                base_conditions.append("opened_at IS NOT NULL")
+            elif has_been_opened is False:
+                base_conditions.append("opened_at IS NULL")
 
-        # For each tag to count, run a separate query
-        for tag in tags:
-            # Start with the base conditions
-            conditions = base_conditions.copy()
-            params = base_params.copy()
+            # For each tag to count, run a separate query
+            for tag in tags:
+                # Start with the base conditions
+                conditions = base_conditions.copy()
+                params = base_params.copy()
 
-            # Add this specific tag condition
-            conditions.append("tags LIKE ?")
-            params.append(f"%{tag.strip()}%")
+                # Add this specific tag condition
+                conditions.append("tags LIKE ?")
+                params.append(f"%{tag.strip()}%")
 
-            # Construct the full query
-            stmt = """--sql
-                SELECT COUNT(*)
-                FROM workflow_library
-                """
+                # Construct the full query
+                stmt = """--sql
+                    SELECT COUNT(*)
+                    FROM workflow_library
+                    """
 
-            if conditions:
-                stmt += " WHERE " + " AND ".join(conditions)
+                if conditions:
+                    stmt += " WHERE " + " AND ".join(conditions)
 
-            cursor.execute(stmt, params)
-            count = cursor.fetchone()[0]
-            result[tag] = count
+                cursor.execute(stmt, params)
+                count = cursor.fetchone()[0]
+                result[tag] = count
 
         return result
 
@@ -296,52 +281,51 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
         has_been_opened: Optional[bool] = None,
         is_published: Optional[bool] = None,
     ) -> dict[str, int]:
-        cursor = self._conn.cursor()
-        result: dict[str, int] = {}
-        # Base conditions for categories
-        base_conditions: list[str] = []
-        base_params: list[str | int] = []
+        with self._db.transaction() as cursor:
+            result: dict[str, int] = {}
+            # Base conditions for categories
+            base_conditions: list[str] = []
+            base_params: list[str | int] = []
 
-        # Add category conditions
-        if categories:
-            assert all(c in WorkflowCategory for c in categories)
-            placeholders = ", ".join("?" for _ in categories)
-            base_conditions.append(f"category IN ({placeholders})")
-            base_params.extend([category.value for category in categories])
+            # Add category conditions
+            if categories:
+                assert all(c in WorkflowCategory for c in categories)
+                placeholders = ", ".join("?" for _ in categories)
+                base_conditions.append(f"category IN ({placeholders})")
+                base_params.extend([category.value for category in categories])
 
-        if has_been_opened:
-            base_conditions.append("opened_at IS NOT NULL")
-        elif has_been_opened is False:
-            base_conditions.append("opened_at IS NULL")
+            if has_been_opened:
+                base_conditions.append("opened_at IS NOT NULL")
+            elif has_been_opened is False:
+                base_conditions.append("opened_at IS NULL")
 
-        # For each category to count, run a separate query
-        for category in categories:
-            # Start with the base conditions
-            conditions = base_conditions.copy()
-            params = base_params.copy()
+            # For each category to count, run a separate query
+            for category in categories:
+                # Start with the base conditions
+                conditions = base_conditions.copy()
+                params = base_params.copy()
 
-            # Add this specific category condition
-            conditions.append("category = ?")
-            params.append(category.value)
+                # Add this specific category condition
+                conditions.append("category = ?")
+                params.append(category.value)
 
-            # Construct the full query
-            stmt = """--sql
-                SELECT COUNT(*)
-                FROM workflow_library
-                """
+                # Construct the full query
+                stmt = """--sql
+                    SELECT COUNT(*)
+                    FROM workflow_library
+                    """
 
-            if conditions:
-                stmt += " WHERE " + " AND ".join(conditions)
+                if conditions:
+                    stmt += " WHERE " + " AND ".join(conditions)
 
-            cursor.execute(stmt, params)
-            count = cursor.fetchone()[0]
-            result[category.value] = count
+                cursor.execute(stmt, params)
+                count = cursor.fetchone()[0]
+                result[category.value] = count
 
         return result
 
     def update_opened_at(self, workflow_id: str) -> None:
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 f"""--sql
                 UPDATE workflow_library
@@ -350,10 +334,6 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
                 """,
                 (workflow_id,),
             )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
 
     def _sync_default_workflows(self) -> None:
         """Syncs default workflows to the database. Internal use only."""
@@ -368,8 +348,7 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
         meaningless, as they are overwritten every time the server starts.
         """
 
-        try:
-            cursor = self._conn.cursor()
+        with self._db.transaction() as cursor:
             workflows_from_file: list[Workflow] = []
             workflows_to_update: list[Workflow] = []
             workflows_to_add: list[Workflow] = []
@@ -449,8 +428,3 @@ class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase):
                     """,
                     (w.model_dump_json(), w.id),
                 )
-
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise

invokeai/app/util/controlnet_utils.py

@@ -230,6 +230,86 @@ def heuristic_resize(np_img: np.ndarray[Any, Any], size: tuple[int, int]) -> np.
     return resized
 
 
+# precompute common kernels
+_KERNEL3 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+# directional masks for NMS
+_DIRS = [
+    np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], np.uint8),
+    np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], np.uint8),
+    np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], np.uint8),
+    np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], np.uint8),
+]
+
+
+def heuristic_resize_fast(np_img: np.ndarray, size: tuple[int, int]) -> np.ndarray:
+    h, w = np_img.shape[:2]
+    # early exit
+    if (w, h) == size:
+        return np_img
+
+    # separate alpha channel
+    img = np_img
+    alpha = None
+    if img.ndim == 3 and img.shape[2] == 4:
+        alpha, img = img[:, :, 3], img[:, :, :3]
+
+    # build small sample for unique‐color & binary detection
+    flat = img.reshape(-1, img.shape[-1])
+    N = flat.shape[0]
+    # include four corners to avoid missing extreme values
+    corners = np.vstack([img[0, 0], img[0, w - 1], img[h - 1, 0], img[h - 1, w - 1]])
+    cnt = min(N, 100_000)
+    samp = np.vstack([corners, flat[np.random.choice(N, cnt, replace=False)]])
+    uc = np.unique(samp, axis=0).shape[0]
+    vmin, vmax = samp.min(), samp.max()
+
+    # detect binary edge map & one‐pixel‐edge case
+    is_binary = uc == 2 and vmin < 16 and vmax > 240
+    one_pixel_edge = False
+    if is_binary:
+        # single gray conversion
+        gray0 = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        grad = cv2.morphologyEx(gray0, cv2.MORPH_GRADIENT, _KERNEL3)
+        cnt_edge = cv2.countNonZero(grad)
+        cnt_all = cv2.countNonZero((gray0 > 127).astype(np.uint8))
+        one_pixel_edge = (2 * cnt_edge) > cnt_all
+
+    # choose interp for color/seg/grayscale
+    area_new, area_old = size[0] * size[1], w * h
+    if 2 < uc < 200:  # segmentation map
+        interp = cv2.INTER_NEAREST
+    elif area_new < area_old:
+        interp = cv2.INTER_AREA
+    else:
+        interp = cv2.INTER_CUBIC
+
+    # single resize pass on RGB
+    resized = cv2.resize(img, size, interpolation=interp)
+
+    if is_binary:
+        # convert to gray & apply NMS via C++ dilate
+        gray_r = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)
+        nms = np.zeros_like(gray_r)
+        for K in _DIRS:
+            d = cv2.dilate(gray_r, K)
+            mask = d == gray_r
+            nms[mask] = gray_r[mask]
+
+        # threshold + thinning if needed
+        _, bw = cv2.threshold(nms, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        out_bin = cv2.ximgproc.thinning(bw) if one_pixel_edge else bw
+        # restore 3 channels
+        resized = np.stack([out_bin] * 3, axis=2)
+
+    # restore alpha with same interp as RGB for consistency
+    if alpha is not None:
+        am = cv2.resize(alpha, size, interpolation=interp)
+        am = (am > 127).astype(np.uint8) * 255
+        resized = np.dstack((resized, am))
+
+    return resized
+
+
 ###########################################################################
 # Copied from detectmap_proc method in scripts/detectmap_proc.py in Mikubill/sd-webui-controlnet
 #    modified for InvokeAI
@@ -244,7 +324,7 @@ def np_img_resize(
     np_img = normalize_image_channel_count(np_img)
 
     if resize_mode == "just_resize":  # RESIZE
-        np_img = heuristic_resize(np_img, (w, h))
+        np_img = heuristic_resize_fast(np_img, (w, h))
         np_img = clone_contiguous(np_img)
         return np_img_to_torch(np_img, device), np_img
 
@@ -265,7 +345,7 @@ def np_img_resize(
             # Inpaint hijack
             high_quality_border_color[3] = 255
         high_quality_background = np.tile(high_quality_border_color[None, None], [h, w, 1])
-        np_img = heuristic_resize(np_img, (safeint(old_w * k), safeint(old_h * k)))
+        np_img = heuristic_resize_fast(np_img, (safeint(old_w * k), safeint(old_h * k)))
         new_h, new_w, _ = np_img.shape
         pad_h = max(0, (h - new_h) // 2)
         pad_w = max(0, (w - new_w) // 2)
@@ -275,7 +355,7 @@ def np_img_resize(
         return np_img_to_torch(np_img, device), np_img
     else:  # resize_mode == "crop_resize"  (INNER_FIT)
         k = max(k0, k1)
-        np_img = heuristic_resize(np_img, (safeint(old_w * k), safeint(old_h * k)))
+        np_img = heuristic_resize_fast(np_img, (safeint(old_w * k), safeint(old_h * k)))
         new_h, new_w, _ = np_img.shape
         pad_h = max(0, (new_h - h) // 2)
         pad_w = max(0, (new_w - w) // 2)

invokeai/app/util/custom_openapi.py

@@ -12,6 +12,9 @@ from invokeai.app.invocations.fields import InputFieldJSONSchemaExtra, OutputFie
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.services.events.events_common import EventBase
 from invokeai.app.services.session_processor.session_processor_common import ProgressImage
+from invokeai.backend.util.logging import InvokeAILogger
+
+logger = InvokeAILogger.get_logger()
 
 
 def move_defs_to_top_level(openapi_schema: dict[str, Any], component_schema: dict[str, Any]) -> None:

invokeai/app/util/misc.py

@@ -10,7 +10,7 @@ def get_timestamp() -> int:
 
 
 def get_iso_timestamp() -> str:
-    return datetime.datetime.utcnow().isoformat()
+    return datetime.datetime.now(datetime.timezone.utc).isoformat()
 
 
 def get_datetime_from_iso_timestamp(iso_timestamp: str) -> datetime.datetime:

invokeai/app/util/step_callback.py

@@ -123,7 +123,11 @@ def calc_percentage(intermediate_state: PipelineIntermediateState) -> float:
     if total_steps == 0:
         return 0.0
     if order == 2:
-        return floor(step / 2) / floor(total_steps / 2)
+        # Prevent division by zero when total_steps is 1 or 2
+        denominator = floor(total_steps / 2)
+        if denominator == 0:
+            return 0.0
+        return floor(step / 2) / denominator
     # order == 1
     return step / total_steps
 

invokeai/backend/flux/denoise.py

@@ -30,8 +30,11 @@ def denoise(
     controlnet_extensions: list[XLabsControlNetExtension | InstantXControlNetExtension],
     pos_ip_adapter_extensions: list[XLabsIPAdapterExtension],
     neg_ip_adapter_extensions: list[XLabsIPAdapterExtension],
-    # extra img tokens
+    # extra img tokens (channel-wise)
     img_cond: torch.Tensor | None,
+    # extra img tokens (sequence-wise) - for Kontext conditioning
+    img_cond_seq: torch.Tensor | None = None,
+    img_cond_seq_ids: torch.Tensor | None = None,
 ):
     # step 0 is the initial state
     total_steps = len(timesteps) - 1
@@ -46,6 +49,10 @@ def denoise(
     )
     # guidance_vec is ignored for schnell.
     guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+
+    # Store original sequence length for slicing predictions
+    original_seq_len = img.shape[1]
+
     for step_index, (t_curr, t_prev) in tqdm(list(enumerate(zip(timesteps[:-1], timesteps[1:], strict=True)))):
         t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
 
@@ -71,10 +78,26 @@ 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
+
+        # Prepare input for model - concatenate fresh each step
+        img_input = img
+        img_input_ids = img_ids
+
+        # Add channel-wise conditioning (for ControlNet, FLUX Fill, etc.)
+        if img_cond is not None:
+            img_input = torch.cat((img_input, img_cond), dim=-1)
+
+        # Add sequence-wise conditioning (for Kontext)
+        if img_cond_seq is not None:
+            assert img_cond_seq_ids is not None, (
+                "You need to provide either both or neither of the sequence conditioning"
+            )
+            img_input = torch.cat((img_input, img_cond_seq), dim=1)
+            img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids), dim=1)
+
         pred = model(
-            img=pred_img,
-            img_ids=img_ids,
+            img=img_input,
+            img_ids=img_input_ids,
             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,
@@ -88,6 +111,10 @@ def denoise(
             regional_prompting_extension=pos_regional_prompting_extension,
         )
 
+        # Slice prediction to only include the main image tokens
+        if img_cond_seq is not None:
+            pred = pred[:, :original_seq_len]
+
         step_cfg_scale = cfg_scale[step_index]
 
         # If step_cfg_scale, is 1.0, then we don't need to run the negative prediction.
@@ -98,9 +125,26 @@ def denoise(
             if neg_regional_prompting_extension is None:
                 raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
 
+            # For negative prediction with Kontext, we need to include the reference images
+            # to maintain consistency between positive and negative passes. Without this,
+            # CFG would create artifacts as the attention mechanism would see different
+            # spatial structures in each pass
+            neg_img_input = img
+            neg_img_input_ids = img_ids
+
+            # Add channel-wise conditioning for negative pass if present
+            if img_cond is not None:
+                neg_img_input = torch.cat((neg_img_input, img_cond), dim=-1)
+
+            # Add sequence-wise conditioning (Kontext) for negative pass
+            # This ensures reference images are processed consistently
+            if img_cond_seq is not None:
+                neg_img_input = torch.cat((neg_img_input, img_cond_seq), dim=1)
+                neg_img_input_ids = torch.cat((neg_img_input_ids, img_cond_seq_ids), dim=1)
+
             neg_pred = model(
-                img=img,
-                img_ids=img_ids,
+                img=neg_img_input,
+                img_ids=neg_img_input_ids,
                 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,
@@ -113,6 +157,10 @@ def denoise(
                 ip_adapter_extensions=neg_ip_adapter_extensions,
                 regional_prompting_extension=neg_regional_prompting_extension,
             )
+
+            # Slice negative prediction to match main image tokens
+            if img_cond_seq is not None:
+                neg_pred = neg_pred[:, :original_seq_len]
             pred = neg_pred + step_cfg_scale * (pred - neg_pred)
 
         preview_img = img - t_curr * pred

invokeai/backend/flux/extensions/kontext_extension.py

@@ -0,0 +1,203 @@
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+from einops import repeat
+
+from invokeai.app.invocations.fields import FluxKontextConditioningField
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+from invokeai.backend.flux.sampling_utils import pack
+from invokeai.backend.util.devices import TorchDevice
+
+
+def generate_img_ids_with_offset(
+    latent_height: int,
+    latent_width: int,
+    batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    idx_offset: int = 0,
+    h_offset: int = 0,
+    w_offset: int = 0,
+) -> torch.Tensor:
+    """Generate tensor of image position ids with optional index and spatial offsets.
+
+    Args:
+        latent_height (int): Height of image in latent space (after packing, this becomes h//2).
+        latent_width (int): Width of image in latent space (after packing, this becomes w//2).
+        batch_size (int): Number of images in the batch.
+        device (torch.device): Device to create tensors on.
+        dtype (torch.dtype): Data type for the tensors.
+        idx_offset (int): Offset to add to the first dimension of the image ids (default: 0).
+        h_offset (int): Spatial offset for height/y-coordinates in latent space (default: 0).
+        w_offset (int): Spatial offset for width/x-coordinates in latent space (default: 0).
+
+    Returns:
+        torch.Tensor: Image position ids with shape [batch_size, (latent_height//2 * latent_width//2), 3].
+    """
+
+    if device.type == "mps":
+        orig_dtype = dtype
+        dtype = torch.float16
+
+    # After packing, the spatial dimensions are halved due to the 2x2 patch structure
+    packed_height = latent_height // 2
+    packed_width = latent_width // 2
+
+    # Convert spatial offsets from latent space to packed space
+    packed_h_offset = h_offset // 2
+    packed_w_offset = w_offset // 2
+
+    # Create base tensor for position IDs with shape [packed_height, packed_width, 3]
+    # The 3 channels represent: [batch_offset, y_position, x_position]
+    img_ids = torch.zeros(packed_height, packed_width, 3, device=device, dtype=dtype)
+
+    # Set the batch offset for all positions
+    img_ids[..., 0] = idx_offset
+
+    # Create y-coordinate indices (vertical positions) with spatial offset
+    y_indices = torch.arange(packed_height, device=device, dtype=dtype) + packed_h_offset
+    # Broadcast y_indices to match the spatial dimensions [packed_height, 1]
+    img_ids[..., 1] = y_indices[:, None]
+
+    # Create x-coordinate indices (horizontal positions) with spatial offset
+    x_indices = torch.arange(packed_width, device=device, dtype=dtype) + packed_w_offset
+    # Broadcast x_indices to match the spatial dimensions [1, packed_width]
+    img_ids[..., 2] = x_indices[None, :]
+
+    # Expand to include batch dimension: [batch_size, (packed_height * packed_width), 3]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=batch_size)
+
+    if device.type == "mps":
+        img_ids = img_ids.to(orig_dtype)
+
+    return img_ids
+
+
+class KontextExtension:
+    """Applies FLUX Kontext (reference image) conditioning."""
+
+    def __init__(
+        self,
+        kontext_conditioning: list[FluxKontextConditioningField],
+        context: InvocationContext,
+        vae_field: VAEField,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        """
+        Initializes the KontextExtension, pre-processing the reference images
+        into latents and positional IDs.
+        """
+        self._context = context
+        self._device = device
+        self._dtype = dtype
+        self._vae_field = vae_field
+        self.kontext_conditioning = kontext_conditioning
+
+        # Pre-process and cache the kontext latents and ids upon initialization.
+        self.kontext_latents, self.kontext_ids = self._prepare_kontext()
+
+    def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encodes the reference images and prepares their concatenated latents and IDs with spatial tiling."""
+        all_latents = []
+        all_ids = []
+
+        # Track cumulative dimensions for spatial tiling
+        # These track the running extent of the virtual canvas in latent space
+        h = 0  # Running height extent
+        w = 0  # Running width extent
+
+        vae_info = self._context.models.load(self._vae_field.vae)
+
+        for idx, kontext_field in enumerate(self.kontext_conditioning):
+            image = self._context.images.get_pil(kontext_field.image.image_name)
+
+            # Convert to RGB
+            image = image.convert("RGB")
+
+            # Convert to tensor using torchvision transforms for consistency
+            transformation = T.Compose(
+                [
+                    T.ToTensor(),  # Converts PIL image to tensor and scales to [0, 1]
+                ]
+            )
+            image_tensor = transformation(image)
+            # Convert from [0, 1] to [-1, 1] range expected by VAE
+            image_tensor = image_tensor * 2.0 - 1.0
+            image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension
+            image_tensor = image_tensor.to(self._device)
+
+            # Continue with VAE encoding
+            # Don't sample from the distribution for reference images - use the mean (matching ComfyUI)
+            with vae_info as vae:
+                assert isinstance(vae, AutoEncoder)
+                vae_dtype = next(iter(vae.parameters())).dtype
+                image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
+                # Use sample=False to get the distribution mean without noise
+                kontext_latents_unpacked = vae.encode(image_tensor, sample=False)
+
+            # Extract tensor dimensions
+            batch_size, _, latent_height, latent_width = kontext_latents_unpacked.shape
+
+            # Pad latents to be compatible with patch_size=2
+            # This ensures dimensions are even for the pack() function
+            pad_h = (2 - latent_height % 2) % 2
+            pad_w = (2 - latent_width % 2) % 2
+            if pad_h > 0 or pad_w > 0:
+                kontext_latents_unpacked = F.pad(kontext_latents_unpacked, (0, pad_w, 0, pad_h), mode="circular")
+                # Update dimensions after padding
+                _, _, latent_height, latent_width = kontext_latents_unpacked.shape
+
+            # Pack the latents
+            kontext_latents_packed = pack(kontext_latents_unpacked).to(self._device, self._dtype)
+
+            # Determine spatial offsets for this reference image
+            # - Compare the potential new canvas dimensions if we add the image vertically vs horizontally
+            # - Choose the placement that results in a more square-like canvas
+            h_offset = 0
+            w_offset = 0
+
+            if idx > 0:  # First image starts at (0, 0)
+                # Check which placement would result in better canvas dimensions
+                # If adding to height would make the canvas taller than wide, tile horizontally
+                # Otherwise, tile vertically
+                if latent_height + h > latent_width + w:
+                    # Tile horizontally (to the right of existing images)
+                    w_offset = w
+                else:
+                    # Tile vertically (below existing images)
+                    h_offset = h
+
+            # Generate IDs with both index offset and spatial offsets
+            kontext_ids = generate_img_ids_with_offset(
+                latent_height=latent_height,
+                latent_width=latent_width,
+                batch_size=batch_size,
+                device=self._device,
+                dtype=self._dtype,
+                idx_offset=1,  # All reference images use index=1 (matching ComfyUI implementation)
+                h_offset=h_offset,
+                w_offset=w_offset,
+            )
+
+            # Update cumulative dimensions
+            # Track the maximum extent of the virtual canvas after placing this image
+            h = max(h, latent_height + h_offset)
+            w = max(w, latent_width + w_offset)
+
+            all_latents.append(kontext_latents_packed)
+            all_ids.append(kontext_ids)
+
+        # Concatenate all latents and IDs along the sequence dimension
+        concatenated_latents = torch.cat(all_latents, dim=1)  # Concatenate along sequence dimension
+        concatenated_ids = torch.cat(all_ids, dim=1)  # Concatenate along sequence dimension
+
+        return concatenated_latents, concatenated_ids
+
+    def ensure_batch_size(self, target_batch_size: int) -> None:
+        """Ensures the kontext latents and IDs match the target batch size by repeating if necessary."""
+        if self.kontext_latents.shape[0] != target_batch_size:
+            self.kontext_latents = self.kontext_latents.repeat(target_batch_size, 1, 1)
+            self.kontext_ids = self.kontext_ids.repeat(target_batch_size, 1, 1)

invokeai/backend/flux/sampling_utils.py

@@ -174,11 +174,13 @@ def generate_img_ids(h: int, w: int, batch_size: int, device: torch.device, dtyp
         dtype = torch.float16
 
     img_ids = torch.zeros(h // 2, w // 2, 3, device=device, dtype=dtype)
+    # Set batch offset to 0 for main image tokens
+    img_ids[..., 0] = 0
     img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2, device=device, dtype=dtype)[:, None]
     img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2, device=device, dtype=dtype)[None, :]
     img_ids = repeat(img_ids, "h w c -> b (h w) c", b=batch_size)
 
     if device.type == "mps":
-        img_ids.to(orig_dtype)
+        img_ids = img_ids.to(orig_dtype)
 
     return img_ids

invokeai/backend/flux/util.py

@@ -18,6 +18,29 @@ class ModelSpec:
     repo_ae: str | None
 
 
+# Preferred resolutions for Kontext models to avoid tiling artifacts
+# These are the specific resolutions the model was trained on
+PREFERED_KONTEXT_RESOLUTIONS = [
+    (672, 1568),
+    (688, 1504),
+    (720, 1456),
+    (752, 1392),
+    (800, 1328),
+    (832, 1248),
+    (880, 1184),
+    (944, 1104),
+    (1024, 1024),
+    (1104, 944),
+    (1184, 880),
+    (1248, 832),
+    (1328, 800),
+    (1392, 752),
+    (1456, 720),
+    (1504, 688),
+    (1568, 672),
+]
+
+
 max_seq_lengths: Dict[str, Literal[256, 512]] = {
     "flux-dev": 512,
     "flux-dev-fill": 512,

invokeai/backend/ip_adapter/README.md

@@ -42,4 +42,5 @@ IP-Adapters:
 - [InvokeAI/ip_adapter_plus_sd15](https://huggingface.co/InvokeAI/ip_adapter_plus_sd15)
 - [InvokeAI/ip_adapter_plus_face_sd15](https://huggingface.co/InvokeAI/ip_adapter_plus_face_sd15)
 - [InvokeAI/ip_adapter_sdxl](https://huggingface.co/InvokeAI/ip_adapter_sdxl)
-- [InvokeAI/ip_adapter_sdxl_vit_h](https://huggingface.co/InvokeAI/ip_adapter_sdxl_vit_h)
+- [InvokeAI/ip_adapter_sdxl_vit_h](https://huggingface.co/InvokeAI/ip_adapter_sdxl_vit_h)
+- [InvokeAI/ip-adapter-plus_sdxl_vit-h](https://huggingface.co/InvokeAI/ip-adapter-plus_sdxl_vit-h)

invokeai/backend/model_manager/config.py

@@ -37,6 +37,7 @@ from invokeai.app.util.misc import uuid_string
 from invokeai.backend.model_hash.hash_validator import validate_hash
 from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS
 from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.omi import flux_dev_1_lora, stable_diffusion_xl_1_lora
 from invokeai.backend.model_manager.taxonomy import (
     AnyVariant,
     BaseModelType,
@@ -144,34 +145,37 @@ class ModelConfigBase(ABC, BaseModel):
     submodels: Optional[Dict[SubModelType, SubmodelDefinition]] = Field(
         description="Loadable submodels in this model", default=None
     )
+    usage_info: Optional[str] = Field(default=None, description="Usage information for this model")
 
-    _USING_LEGACY_PROBE: ClassVar[set] = set()
-    _USING_CLASSIFY_API: ClassVar[set] = set()
+    USING_LEGACY_PROBE: ClassVar[set] = set()
+    USING_CLASSIFY_API: ClassVar[set] = set()
     _MATCH_SPEED: ClassVar[MatchSpeed] = MatchSpeed.MED
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         if issubclass(cls, LegacyProbeMixin):
-            ModelConfigBase._USING_LEGACY_PROBE.add(cls)
+            ModelConfigBase.USING_LEGACY_PROBE.add(cls)
         else:
-            ModelConfigBase._USING_CLASSIFY_API.add(cls)
+            ModelConfigBase.USING_CLASSIFY_API.add(cls)
 
     @staticmethod
     def all_config_classes():
-        subclasses = ModelConfigBase._USING_LEGACY_PROBE | ModelConfigBase._USING_CLASSIFY_API
+        subclasses = ModelConfigBase.USING_LEGACY_PROBE | ModelConfigBase.USING_CLASSIFY_API
         concrete = {cls for cls in subclasses if not isabstract(cls)}
         return concrete
 
     @staticmethod
-    def classify(model_path: Path, hash_algo: HASHING_ALGORITHMS = "blake3_single", **overrides):
+    def classify(mod: str | Path | ModelOnDisk, hash_algo: HASHING_ALGORITHMS = "blake3_single", **overrides):
         """
         Returns the best matching ModelConfig instance from a model's file/folder path.
         Raises InvalidModelConfigException if no valid configuration is found.
         Created to deprecate ModelProbe.probe
         """
-        candidates = ModelConfigBase._USING_CLASSIFY_API
+        if isinstance(mod, Path | str):
+            mod = ModelOnDisk(mod, hash_algo)
+
+        candidates = ModelConfigBase.USING_CLASSIFY_API
         sorted_by_match_speed = sorted(candidates, key=lambda cls: (cls._MATCH_SPEED, cls.__name__))
-        mod = ModelOnDisk(model_path, hash_algo)
 
         for config_cls in sorted_by_match_speed:
             try:
@@ -183,7 +187,7 @@ class ModelConfigBase(ABC, BaseModel):
             else:
                 return config_cls.from_model_on_disk(mod, **overrides)
 
-        raise InvalidModelConfigException("No valid config found")
+        raise InvalidModelConfigException("Unable to determine model type")
 
     @classmethod
     def get_tag(cls) -> Tag:
@@ -293,7 +297,7 @@ class LoRAConfigBase(ABC, BaseModel):
         from invokeai.backend.patches.lora_conversions.formats import flux_format_from_state_dict
 
         sd = mod.load_state_dict(mod.path)
-        value = flux_format_from_state_dict(sd)
+        value = flux_format_from_state_dict(sd, mod.metadata())
         mod.cache[key] = value
         return value
 
@@ -331,6 +335,36 @@ class T5EncoderBnbQuantizedLlmInt8bConfig(T5EncoderConfigBase, LegacyProbeMixin,
     format: Literal[ModelFormat.BnbQuantizedLlmInt8b] = ModelFormat.BnbQuantizedLlmInt8b
 
 
+class LoRAOmiConfig(LoRAConfigBase, ModelConfigBase):
+    format: Literal[ModelFormat.OMI] = ModelFormat.OMI
+
+    @classmethod
+    def matches(cls, mod: ModelOnDisk) -> bool:
+        if mod.path.is_dir():
+            return False
+
+        metadata = mod.metadata()
+        return (
+            metadata.get("modelspec.sai_model_spec")
+            and metadata.get("ot_branch") == "omi_format"
+            and metadata["modelspec.architecture"].split("/")[1].lower() == "lora"
+        )
+
+    @classmethod
+    def parse(cls, mod: ModelOnDisk) -> dict[str, Any]:
+        metadata = mod.metadata()
+        architecture = metadata["modelspec.architecture"]
+
+        if architecture == stable_diffusion_xl_1_lora:
+            base = BaseModelType.StableDiffusionXL
+        elif architecture == flux_dev_1_lora:
+            base = BaseModelType.Flux
+        else:
+            raise InvalidModelConfigException(f"Unrecognised/unsupported architecture for OMI LoRA: {architecture}")
+
+        return {"base": base}
+
+
 class LoRALyCORISConfig(LoRAConfigBase, ModelConfigBase):
     """Model config for LoRA/Lycoris models."""
 
@@ -347,7 +381,7 @@ class LoRALyCORISConfig(LoRAConfigBase, ModelConfigBase):
 
         state_dict = mod.load_state_dict()
         for key in state_dict.keys():
-            if type(key) is int:
+            if isinstance(key, int):
                 continue
 
             if key.startswith(("lora_te_", "lora_unet_", "lora_te1_", "lora_te2_", "lora_transformer_")):
@@ -665,6 +699,7 @@ AnyModelConfig = Annotated[
         Annotated[ControlNetDiffusersConfig, ControlNetDiffusersConfig.get_tag()],
         Annotated[ControlNetCheckpointConfig, ControlNetCheckpointConfig.get_tag()],
         Annotated[LoRALyCORISConfig, LoRALyCORISConfig.get_tag()],
+        Annotated[LoRAOmiConfig, LoRAOmiConfig.get_tag()],
         Annotated[ControlLoRALyCORISConfig, ControlLoRALyCORISConfig.get_tag()],
         Annotated[ControlLoRADiffusersConfig, ControlLoRADiffusersConfig.get_tag()],
         Annotated[LoRADiffusersConfig, LoRADiffusersConfig.get_tag()],

invokeai/backend/model_manager/legacy_probe.py

@@ -9,6 +9,7 @@ import spandrel
 import torch
 
 import invokeai.backend.util.logging as logger
+from invokeai.app.services.config.config_default import get_config
 from invokeai.app.util.misc import uuid_string
 from invokeai.backend.flux.controlnet.state_dict_utils import (
     is_state_dict_instantx_controlnet,
@@ -493,9 +494,21 @@ class ModelProbe(object):
         # scan model
         scan_result = pscan.scan_file_path(checkpoint)
         if scan_result.infected_files != 0:
-            raise Exception(f"The model {model_name} is potentially infected by malware. Aborting import.")
+            if get_config().unsafe_disable_picklescan:
+                logger.warning(
+                    f"The model {model_name} is potentially infected by malware, but picklescan is disabled. "
+                    "Proceeding with caution."
+                )
+            else:
+                raise RuntimeError(f"The model {model_name} is potentially infected by malware. Aborting import.")
         if scan_result.scan_err:
-            raise Exception(f"Error scanning model {model_name} for malware. Aborting import.")
+            if get_config().unsafe_disable_picklescan:
+                logger.warning(
+                    f"Error scanning the model at {model_name} for malware, but picklescan is disabled. "
+                    "Proceeding with caution."
+                )
+            else:
+                raise RuntimeError(f"Error scanning the model at {model_name} for malware. Aborting import.")
 
 
 # Probing utilities

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

@@ -2,6 +2,8 @@ from typing import Any
 
 import torch
 
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
 
 class CachedModelOnlyFullLoad:
     """A wrapper around a PyTorch model to handle full loads and unloads between the CPU and the compute device.
@@ -76,7 +78,15 @@ class CachedModelOnlyFullLoad:
             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)
+
+        check_for_gguf = hasattr(self._model, "state_dict") and self._model.state_dict().get("img_in.weight")
+        if isinstance(check_for_gguf, GGMLTensor):
+            old_value = torch.__future__.get_overwrite_module_params_on_conversion()
+            torch.__future__.set_overwrite_module_params_on_conversion(True)
+            self._model.to(self._compute_device)
+            torch.__future__.set_overwrite_module_params_on_conversion(old_value)
+        else:
+            self._model.to(self._compute_device)
 
         self._is_in_vram = True
         return self._total_bytes
@@ -92,7 +102,15 @@ class CachedModelOnlyFullLoad:
 
         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)
+
+        check_for_gguf = hasattr(self._model, "state_dict") and self._model.state_dict().get("img_in.weight")
+        if isinstance(check_for_gguf, GGMLTensor):
+            old_value = torch.__future__.get_overwrite_module_params_on_conversion()
+            torch.__future__.set_overwrite_module_params_on_conversion(True)
+            self._model.to(self._offload_device)
+            torch.__future__.set_overwrite_module_params_on_conversion(old_value)
+        else:
+            self._model.to(self._offload_device)
 
         self._is_in_vram = False
         return self._total_bytes

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

@@ -2,9 +2,10 @@ import gc
 import logging
 import threading
 import time
+from dataclasses import dataclass
 from functools import wraps
 from logging import Logger
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Protocol
 
 import psutil
 import torch
@@ -54,6 +55,39 @@ def synchronized(method: Callable[..., Any]) -> Callable[..., Any]:
     return wrapper
 
 
+@dataclass
+class CacheEntrySnapshot:
+    cache_key: str
+    total_bytes: int
+    current_vram_bytes: int
+
+
+class CacheMissCallback(Protocol):
+    def __call__(
+        self,
+        model_key: str,
+        cache_snapshot: dict[str, CacheEntrySnapshot],
+    ) -> None: ...
+
+
+class CacheHitCallback(Protocol):
+    def __call__(
+        self,
+        model_key: str,
+        cache_snapshot: dict[str, CacheEntrySnapshot],
+    ) -> None: ...
+
+
+class CacheModelsClearedCallback(Protocol):
+    def __call__(
+        self,
+        models_cleared: int,
+        bytes_requested: int,
+        bytes_freed: int,
+        cache_snapshot: dict[str, CacheEntrySnapshot],
+    ) -> None: ...
+
+
 class ModelCache:
     """A cache for managing models in memory.
 
@@ -144,6 +178,34 @@ class ModelCache:
         # - Requests to empty the cache from a separate thread
         self._lock = threading.RLock()
 
+        self._on_cache_hit_callbacks: set[CacheHitCallback] = set()
+        self._on_cache_miss_callbacks: set[CacheMissCallback] = set()
+        self._on_cache_models_cleared_callbacks: set[CacheModelsClearedCallback] = set()
+
+    def on_cache_hit(self, cb: CacheHitCallback) -> Callable[[], None]:
+        self._on_cache_hit_callbacks.add(cb)
+
+        def unsubscribe() -> None:
+            self._on_cache_hit_callbacks.discard(cb)
+
+        return unsubscribe
+
+    def on_cache_miss(self, cb: CacheHitCallback) -> Callable[[], None]:
+        self._on_cache_miss_callbacks.add(cb)
+
+        def unsubscribe() -> None:
+            self._on_cache_miss_callbacks.discard(cb)
+
+        return unsubscribe
+
+    def on_cache_models_cleared(self, cb: CacheModelsClearedCallback) -> Callable[[], None]:
+        self._on_cache_models_cleared_callbacks.add(cb)
+
+        def unsubscribe() -> None:
+            self._on_cache_models_cleared_callbacks.discard(cb)
+
+        return unsubscribe
+
     @property
     @synchronized
     def stats(self) -> Optional[CacheStats]:
@@ -195,6 +257,20 @@ class ModelCache:
             f"Added model {key} (Type: {model.__class__.__name__}, Wrap mode: {wrapped_model.__class__.__name__}, Model size: {size / MB:.2f}MB)"
         )
 
+    @synchronized
+    def _get_cache_snapshot(self) -> dict[str, CacheEntrySnapshot]:
+        overview: dict[str, CacheEntrySnapshot] = {}
+        for cache_key, cache_entry in self._cached_models.items():
+            total_bytes = cache_entry.cached_model.total_bytes()
+            current_vram_bytes = cache_entry.cached_model.cur_vram_bytes()
+            overview[cache_key] = CacheEntrySnapshot(
+                cache_key=cache_key,
+                total_bytes=total_bytes,
+                current_vram_bytes=current_vram_bytes,
+            )
+
+        return overview
+
     @synchronized
     def get(self, key: str, stats_name: Optional[str] = None) -> CacheRecord:
         """Retrieve a model from the cache.
@@ -208,6 +284,8 @@ class ModelCache:
             if self.stats:
                 self.stats.hits += 1
         else:
+            for cb in self._on_cache_miss_callbacks:
+                cb(model_key=key, cache_snapshot=self._get_cache_snapshot())
             if self.stats:
                 self.stats.misses += 1
             self._logger.debug(f"Cache miss: {key}")
@@ -229,6 +307,8 @@ class ModelCache:
         self._cache_stack.append(key)
 
         self._logger.debug(f"Cache hit: {key} (Type: {cache_entry.cached_model.model.__class__.__name__})")
+        for cb in self._on_cache_hit_callbacks:
+            cb(model_key=key, cache_snapshot=self._get_cache_snapshot())
         return cache_entry
 
     @synchronized
@@ -649,6 +729,13 @@ class ModelCache:
             # immediately when their reference count hits 0.
             if self.stats:
                 self.stats.cleared = models_cleared
+            for cb in self._on_cache_models_cleared_callbacks:
+                cb(
+                    models_cleared=models_cleared,
+                    bytes_requested=bytes_needed,
+                    bytes_freed=ram_bytes_freed,
+                    cache_snapshot=self._get_cache_snapshot(),
+                )
             gc.collect()
 
         TorchDevice.empty_cache()

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

@@ -13,6 +13,12 @@ from invokeai.backend.patches.layers.lora_layer import LoRALayer
 
 def linear_lora_forward(input: torch.Tensor, lora_layer: LoRALayer, lora_weight: float) -> torch.Tensor:
     """An optimized implementation of the residual calculation for a sidecar linear LoRALayer."""
+    # up matrix and down matrix have different ranks so we can't simply multiply them
+    if lora_layer.up.shape[1] != lora_layer.down.shape[0]:
+        x = torch.nn.functional.linear(input, lora_layer.get_weight(lora_weight), bias=lora_layer.bias)
+        x *= lora_weight * lora_layer.scale()
+        return x
+
     x = torch.nn.functional.linear(input, lora_layer.down)
     if lora_layer.mid is not None:
         x = torch.nn.functional.linear(x, lora_layer.mid)

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

@@ -7,7 +7,14 @@ from typing import Optional
 import accelerate
 import torch
 from safetensors.torch import load_file
-from transformers import AutoConfig, AutoModelForTextEncoding, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
+from transformers import (
+    AutoConfig,
+    AutoModelForTextEncoding,
+    CLIPTextModel,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
 
 from invokeai.app.services.config.config_default import get_config
 from invokeai.backend.flux.controlnet.instantx_controlnet_flux import InstantXControlNetFlux
@@ -139,7 +146,7 @@ class BnbQuantizedLlmInt8bCheckpointModel(ModelLoader):
             )
         match submodel_type:
             case SubModelType.Tokenizer2 | SubModelType.Tokenizer3:
-                return T5Tokenizer.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
+                return T5TokenizerFast.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
             case SubModelType.TextEncoder2 | SubModelType.TextEncoder3:
                 te2_model_path = Path(config.path) / "text_encoder_2"
                 model_config = AutoConfig.from_pretrained(te2_model_path)
@@ -183,7 +190,7 @@ class T5EncoderCheckpointModel(ModelLoader):
 
         match submodel_type:
             case SubModelType.Tokenizer2 | SubModelType.Tokenizer3:
-                return T5Tokenizer.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
+                return T5TokenizerFast.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
             case SubModelType.TextEncoder2 | SubModelType.TextEncoder3:
                 return T5EncoderModel.from_pretrained(
                     Path(config.path) / "text_encoder_2", torch_dtype="auto", low_cpu_mem_usage=True

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

@@ -13,6 +13,7 @@ from invokeai.backend.model_manager.config import AnyModelConfig
 from invokeai.backend.model_manager.load.load_default import ModelLoader
 from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
 from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.omi.omi import convert_from_omi
 from invokeai.backend.model_manager.taxonomy import (
     AnyModel,
     BaseModelType,
@@ -20,6 +21,10 @@ from invokeai.backend.model_manager.taxonomy import (
     ModelType,
     SubModelType,
 )
+from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_aitoolkit_format,
+    lora_model_from_flux_aitoolkit_state_dict,
+)
 from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import (
     is_state_dict_likely_flux_control,
     lora_model_from_flux_control_state_dict,
@@ -39,6 +44,8 @@ from invokeai.backend.patches.lora_conversions.sd_lora_conversion_utils import l
 from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import convert_sdxl_keys_to_diffusers_format
 
 
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.LoRA, format=ModelFormat.OMI)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.LoRA, format=ModelFormat.OMI)
 @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)
@@ -73,12 +80,23 @@ class LoRALoader(ModelLoader):
         else:
             state_dict = torch.load(model_path, map_location="cpu")
 
+        # Strip 'bundle_emb' keys - these are unused and currently cause downstream errors.
+        # To revisit later to determine if they're needed/useful.
+        state_dict = {k: v for k, v in state_dict.items() if not k.startswith("bundle_emb")}
+
+        # At the time of writing, we support the OMI standard for base models Flux and SDXL
+        if config.format == ModelFormat.OMI and self._model_base in [
+            BaseModelType.StableDiffusionXL,
+            BaseModelType.Flux,
+        ]:
+            state_dict = convert_from_omi(state_dict, config.base)  # type: ignore
+
         # Apply state_dict key conversions, if necessary.
         if self._model_base == BaseModelType.StableDiffusionXL:
             state_dict = convert_sdxl_keys_to_diffusers_format(state_dict)
             model = lora_model_from_sd_state_dict(state_dict=state_dict)
         elif self._model_base == BaseModelType.Flux:
-            if config.format == ModelFormat.Diffusers:
+            if config.format in [ModelFormat.Diffusers, ModelFormat.OMI]:
                 # HACK(ryand): We set alpha=None for diffusers PEFT format models. These models are typically
                 # distributed as a single file without the associated metadata containing the alpha value. We chose
                 # alpha=None, because this is treated as alpha=rank internally in `LoRALayerBase.scale()`. alpha=rank
@@ -92,8 +110,10 @@ class LoRALoader(ModelLoader):
                     model = lora_model_from_flux_onetrainer_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)
+                elif is_state_dict_likely_in_flux_aitoolkit_format(state_dict=state_dict):
+                    model = lora_model_from_flux_aitoolkit_state_dict(state_dict=state_dict)
                 else:
-                    raise ValueError(f"LoRA model is in unsupported FLUX format: {config.format}")
+                    raise ValueError("LoRA model is in unsupported FLUX format")
             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/metadata/fetch/huggingface.py

@@ -62,11 +62,14 @@ class HuggingFaceMetadataFetch(ModelMetadataFetchBase):
         # If this too fails, raise exception.
 
         model_info = None
+
+        # Handling for our special syntax - we only want the base HF `org/repo` here.
+        repo_id = id.split("::")[0] or id
         while not model_info:
             try:
-                model_info = HfApi().model_info(repo_id=id, files_metadata=True, revision=variant)
+                model_info = HfApi().model_info(repo_id=repo_id, files_metadata=True, revision=variant)
             except RepositoryNotFoundError as excp:
-                raise UnknownMetadataException(f"'{id}' not found. See trace for details.") from excp
+                raise UnknownMetadataException(f"'{repo_id}' not found. See trace for details.") from excp
             except RevisionNotFoundError:
                 if variant is None:
                     raise
@@ -75,14 +78,14 @@ class HuggingFaceMetadataFetch(ModelMetadataFetchBase):
 
         files: list[RemoteModelFile] = []
 
-        _, name = id.split("/")
+        _, name = repo_id.split("/")
 
         for s in model_info.siblings or []:
             assert s.rfilename is not None
             assert s.size is not None
             files.append(
                 RemoteModelFile(
-                    url=hf_hub_url(id, s.rfilename, revision=variant or "main"),
+                    url=hf_hub_url(repo_id, s.rfilename, revision=variant or "main"),
                     path=Path(name, s.rfilename),
                     size=s.size,
                     sha256=s.lfs.get("sha256") if s.lfs else None,

invokeai/backend/model_manager/model_on_disk.py

@@ -4,14 +4,19 @@ from typing import Any, Optional, TypeAlias
 import safetensors.torch
 import torch
 from picklescan.scanner import scan_file_path
+from safetensors import safe_open
 
+from invokeai.app.services.config.config_default import get_config
 from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
 from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
 from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
+from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.backend.util.silence_warnings import SilenceWarnings
 
 StateDict: TypeAlias = dict[str | int, Any]  # When are the keys int?
 
+logger = InvokeAILogger.get_logger()
+
 
 class ModelOnDisk:
     """A utility class representing a model stored on disk."""
@@ -35,12 +40,21 @@ class ModelOnDisk:
             return self.path.stat().st_size
         return sum(file.stat().st_size for file in self.path.rglob("*"))
 
-    def component_paths(self) -> set[Path]:
+    def weight_files(self) -> set[Path]:
         if self.path.is_file():
             return {self.path}
         extensions = {".safetensors", ".pt", ".pth", ".ckpt", ".bin", ".gguf"}
         return {f for f in self.path.rglob("*") if f.suffix in extensions}
 
+    def metadata(self, path: Optional[Path] = None) -> dict[str, str]:
+        try:
+            with safe_open(self.path, framework="pt", device="cpu") as f:
+                metadata = f.metadata()
+                assert isinstance(metadata, dict)
+                return metadata
+        except Exception:
+            return {}
+
     def repo_variant(self) -> Optional[ModelRepoVariant]:
         if self.path.is_file():
             return None
@@ -64,24 +78,29 @@ class ModelOnDisk:
         if path in sd_cache:
             return sd_cache[path]
 
-        if not path:
-            components = list(self.component_paths())
-            match components:
-                case []:
-                    raise ValueError("No weight files found for this model")
-                case [p]:
-                    path = p
-                case ps if len(ps) >= 2:
-                    raise ValueError(
-                        f"Multiple weight files found for this model: {ps}. "
-                        f"Please specify the intended file using the 'path' argument"
-                    )
+        path = self.resolve_weight_file(path)
 
         with SilenceWarnings():
             if path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin")):
                 scan_result = scan_file_path(path)
-                if scan_result.infected_files != 0 or scan_result.scan_err:
-                    raise RuntimeError(f"The model {path.stem} is potentially infected by malware. Aborting import.")
+                if scan_result.infected_files != 0:
+                    if get_config().unsafe_disable_picklescan:
+                        logger.warning(
+                            f"The model {path.stem} is potentially infected by malware, but picklescan is disabled. "
+                            "Proceeding with caution."
+                        )
+                    else:
+                        raise RuntimeError(
+                            f"The model {path.stem} is potentially infected by malware. Aborting import."
+                        )
+                if scan_result.scan_err:
+                    if get_config().unsafe_disable_picklescan:
+                        logger.warning(
+                            f"Error scanning the model at {path.stem} for malware, but picklescan is disabled. "
+                            "Proceeding with caution."
+                        )
+                    else:
+                        raise RuntimeError(f"Error scanning the model at {path.stem} for malware. Aborting import.")
                 checkpoint = torch.load(path, map_location="cpu")
                 assert isinstance(checkpoint, dict)
             elif path.suffix.endswith(".gguf"):
@@ -94,3 +113,18 @@ class ModelOnDisk:
         state_dict = checkpoint.get("state_dict", checkpoint)
         sd_cache[path] = state_dict
         return state_dict
+
+    def resolve_weight_file(self, path: Optional[Path] = None) -> Path:
+        if not path:
+            weight_files = list(self.weight_files())
+            match weight_files:
+                case []:
+                    raise ValueError("No weight files found for this model")
+                case [p]:
+                    return p
+                case ps if len(ps) >= 2:
+                    raise ValueError(
+                        f"Multiple weight files found for this model: {ps}. "
+                        f"Please specify the intended file using the 'path' argument"
+                    )
+        return path

invokeai/backend/model_manager/omi/__init__.py

@@ -0,0 +1,7 @@
+from invokeai.backend.model_manager.omi.omi import convert_from_omi
+from invokeai.backend.model_manager.omi.vendor.model_spec.architecture import (
+    flux_dev_1_lora,
+    stable_diffusion_xl_1_lora,
+)
+
+__all__ = ["flux_dev_1_lora", "stable_diffusion_xl_1_lora", "convert_from_omi"]