chore: bump version to v6.0.0rc1

Co-authored-by: kent <kent@invoke.ai>

.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/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 .

.gitignore

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

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/frontend/workflows.md

@@ -297,7 +297,7 @@ Migration logic is in [migrations.ts].
 <!-- links -->
 
 [pydantic]: https://github.com/pydantic/pydantic 'pydantic'
-[zod]: https://github.com/colinhacks/zod 'zod'
+[zod]: https://github.com/colinhacks/zod 'zod/v4'
 [openapi-types]: https://github.com/kogosoftwarellc/open-api/tree/main/packages/openapi-types 'openapi-types'
 [reactflow]: https://github.com/xyflow/xyflow 'reactflow'
 [reactflow-concepts]: https://reactflow.dev/learn/concepts/terms-and-definitions

docs/installation/manual.md

@@ -71,7 +71,14 @@ The following commands vary depending on the version of Invoke being installed a
 
 7. Determine the `PyPI` index URL to use for installation, if any. This is necessary to get the right version of torch installed.
 
-    === "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/cu128`.
+        - 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`.
+        - **In all other cases, do not use an index.**
+
+    === "Invoke v5.10.0 to v5.11.0"
 
         - 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`.

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

@@ -23,6 +23,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 +140,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())
@@ -161,6 +167,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,

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/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}",
+        example='"[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,
 )
 
@@ -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

@@ -14,13 +14,14 @@ from invokeai.app.services.session_queue.session_queue_common import (
     CancelByBatchIDsResult,
     CancelByDestinationResult,
     ClearResult,
+    DeleteAllExceptCurrentResult,
+    DeleteByDestinationResult,
     EnqueueBatchResult,
     FieldIdentifier,
     PruneResult,
     RetryItemsResult,
     SessionQueueCountsByDestination,
     SessionQueueItem,
-    SessionQueueItemDTO,
     SessionQueueStatus,
 )
 from invokeai.app.services.shared.pagination import CursorPaginatedResults
@@ -68,7 +69,7 @@ async def enqueue_batch(
     "/{queue_id}/list",
     operation_id="list_queue_items",
     responses={
-        200: {"model": CursorPaginatedResults[SessionQueueItemDTO]},
+        200: {"model": CursorPaginatedResults[SessionQueueItem]},
     },
 )
 async def list_queue_items(
@@ -77,11 +78,36 @@ 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
+        queue_id=queue_id,
+        limit=limit,
+        status=status,
+        cursor=cursor,
+        priority=priority,
+        destination=destination,
+    )
+
+
+@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"""
+
+    return ApiDependencies.invoker.services.session_queue.list_all_queue_items(
+        queue_id=queue_id,
+        destination=destination,
     )
 
 
@@ -121,6 +147,18 @@ async def cancel_all_except_current(
     return ApiDependencies.invoker.services.session_queue.cancel_all_except_current(queue_id=queue_id)
 
 
+@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"""
+    return ApiDependencies.invoker.services.session_queue.delete_all_except_current(queue_id=queue_id)
+
+
 @session_queue_router.put(
     "/{queue_id}/cancel_by_batch_ids",
     operation_id="cancel_by_batch_ids",
@@ -269,6 +307,18 @@ async def get_queue_item(
     return ApiDependencies.invoker.services.session_queue.get_queue_item(item_id)
 
 
+@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"""
+    ApiDependencies.invoker.services.session_queue.delete_queue_item(item_id)
+
+
 @session_queue_router.put(
     "/{queue_id}/i/{item_id}/cancel",
     operation_id="cancel_queue_item",
@@ -298,3 +348,18 @@ async def counts_by_destination(
     return ApiDependencies.invoker.services.session_queue.get_counts_by_destination(
         queue_id=queue_id, destination=destination
     )
+
+
+@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"""
+    return ApiDependencies.invoker.services.session_queue.delete_by_destination(
+        queue_id=queue_id, destination=destination
+    )

invokeai/app/api_app.py

@@ -22,6 +22,7 @@ from invokeai.app.api.routers import (
     download_queue,
     images,
     model_manager,
+    model_relationships,
     session_queue,
     style_presets,
     utilities,
@@ -125,6 +126,7 @@ 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")
@@ -156,7 +158,7 @@ web_root_path = Path(list(web_dir.__path__)[0])
 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,12 +22,14 @@ from typing import (
     Literal,
     Optional,
     Type,
+    TypedDict,
     TypeVar,
     Union,
+    cast,
 )
 
 import semver
-from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, create_model
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, create_model
 from pydantic.fields import FieldInfo
 from pydantic_core import PydanticUndefined
 
@@ -72,13 +76,24 @@ 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.
     This is used internally by the @invocation decorator logic. Do not use this directly.
     """
 
-    tags: Optional[list[str]] = Field(default_factory=None, description="The node's tags")
+    tags: Optional[list[str]] = Field(default=None, description="The node's tags")
     title: Optional[str] = Field(default=None, description="The node's display name")
     category: Optional[str] = Field(default=None, description="The node's category")
     version: str = Field(
@@ -93,6 +108,11 @@ class UIConfigBase(BaseModel):
     )
 
 
+class OriginalModelField(TypedDict):
+    annotation: Any
+    field_info: FieldInfo
+
+
 class BaseInvocationOutput(BaseModel):
     """
     Base class for all invocation outputs.
@@ -100,6 +120,12 @@ class BaseInvocationOutput(BaseModel):
     All invocation outputs must use the `@invocation_output` decorator to provide their unique type.
     """
 
+    output_meta: Optional[dict[str, JsonValue]] = Field(
+        default=None,
+        description="Optional dictionary of metadata for the invocation output, unrelated to the invocation's actual output value. This is not exposed as an output field.",
+        json_schema_extra={"field_kind": FieldKind.NodeAttribute},
+    )
+
     @staticmethod
     def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseInvocationOutput]) -> None:
         """Adds various UI-facing attributes to the invocation output's OpenAPI schema."""
@@ -115,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,
@@ -148,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
 
@@ -180,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
@@ -195,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:
@@ -235,6 +264,8 @@ class BaseInvocation(ABC, BaseModel):
         json_schema_extra={"field_kind": FieldKind.NodeAttribute},
     )
 
+    bottleneck: ClassVar[Bottleneck]
+
     UIConfig: ClassVar[UIConfigBase]
 
     model_config = ConfigDict(
@@ -245,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)
 
@@ -256,6 +290,26 @@ class InvocationRegistry:
     @classmethod
     def register_invocation(cls, invocation: type[BaseInvocation]) -> None:
         """Registers an invocation."""
+
+        invocation_type = invocation.get_type()
+        node_pack = invocation.UIConfig.node_pack
+
+        # Log a warning when an existing invocation is being clobbered by the one we are registering
+        clobbered_invocation = InvocationRegistry.get_invocation_for_type(invocation_type)
+        if clobbered_invocation is not None:
+            # This should always be true - we just checked if the invocation type was in the set
+            clobbered_node_pack = clobbered_invocation.UIConfig.node_pack
+
+            if clobbered_node_pack == "invokeai":
+                # The invocation being clobbered is a core invocation
+                logger.warning(f'Overriding core node "{invocation_type}" with node from "{node_pack}"')
+            else:
+                # The invocation being clobbered is a custom invocation
+                logger.warning(
+                    f'Overriding node "{invocation_type}" from "{node_pack}" with node from "{clobbered_node_pack}"'
+                )
+            cls._invocation_classes.remove(clobbered_invocation)
+
         cls._invocation_classes.add(invocation)
         cls.invalidate_invocation_typeadapter()
 
@@ -314,6 +368,15 @@ class InvocationRegistry:
     @classmethod
     def register_output(cls, output: "type[TBaseInvocationOutput]") -> None:
         """Registers an invocation output."""
+        output_type = output.get_type()
+
+        # Log a warning when an existing invocation is being clobbered by the one we are registering
+        clobbered_output = InvocationRegistry.get_output_for_type(output_type)
+        if clobbered_output is not None:
+            # TODO(psyche): We do not record the node pack of the output, so we cannot log it here
+            logger.warning(f'Overriding invocation output "{output_type}"')
+            cls._output_classes.remove(clobbered_output)
+
         cls._output_classes.add(output)
         cls.invalidate_output_typeadapter()
 
@@ -322,6 +385,11 @@ class InvocationRegistry:
         """Gets all invocation outputs."""
         return cls._output_classes
 
+    @classmethod
+    def get_outputs_map(cls) -> dict[str, type[BaseInvocationOutput]]:
+        """Gets a map of all output types to their output classes."""
+        return {i.get_type(): i for i in cls.get_output_classes()}
+
     @classmethod
     @lru_cache(maxsize=1)
     def get_output_typeadapter(cls) -> TypeAdapter[Any]:
@@ -347,6 +415,11 @@ class InvocationRegistry:
         """Gets all invocation output types."""
         return (i.get_type() for i in cls.get_output_classes())
 
+    @classmethod
+    def get_output_for_type(cls, output_type: str) -> type[BaseInvocationOutput] | None:
+        """Gets the output class for a given output type."""
+        return cls.get_outputs_map().get(output_type)
+
 
 RESERVED_NODE_ATTRIBUTE_FIELD_NAMES = {
     "id",
@@ -354,11 +427,12 @@ RESERVED_NODE_ATTRIBUTE_FIELD_NAMES = {
     "use_cache",
     "type",
     "workflow",
+    "bottleneck",
 }
 
 RESERVED_INPUT_FIELD_NAMES = {"metadata", "board"}
 
-RESERVED_OUTPUT_FIELD_NAMES = {"type"}
+RESERVED_OUTPUT_FIELD_NAMES = {"type", "output_meta"}
 
 
 class _Model(BaseModel):
@@ -425,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,
@@ -438,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.
@@ -449,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]:
@@ -460,27 +578,28 @@ def invocation(
         # The node pack is the module name - will be "invokeai" for built-in nodes
         node_pack = cls.__module__.split(".")[0]
 
-        # Handle the case where an existing node is being clobbered by the one we are registering
-        if invocation_type in InvocationRegistry.get_invocation_types():
-            clobbered_invocation = InvocationRegistry.get_invocation_for_type(invocation_type)
-            # This should always be true - we just checked if the invocation type was in the set
-            assert clobbered_invocation is not None
-
-            clobbered_node_pack = clobbered_invocation.UIConfig.node_pack
-
-            if clobbered_node_pack == "invokeai":
-                # The node being clobbered is a core node
-                raise ValueError(
-                    f'Cannot load node "{invocation_type}" from node pack "{node_pack}" - a core node with the same type already exists'
-                )
-            else:
-                # The node being clobbered is a custom node
-                raise ValueError(
-                    f'Cannot load node "{invocation_type}" from node pack "{node_pack}" - a node with the same type already exists in node pack "{clobbered_node_pack}"'
-                )
-
         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
@@ -496,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)
@@ -504,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:
@@ -513,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')
@@ -539,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
 
@@ -572,29 +705,41 @@ def invocation_output(
         if re.compile(r"^\S+$").match(output_type) is None:
             raise ValueError(f'"output_type" must consist of non-whitespace characters, got "{output_type}"')
 
-        if output_type in InvocationRegistry.get_output_types():
-            raise ValueError(f'Invocation type "{output_type}" already exists')
-
         validate_fields(cls.model_fields, output_type)
 
-        # Add the output type to the model.
+        fields: dict[str, tuple[Any, FieldInfo]] = {}
 
-        output_type_annotation = Literal[output_type]  # type: ignore
-        output_type_field = Field(
-            title="type", default=output_type, json_schema_extra={"field_kind": FieldKind.NodeAttribute}
+        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]
+
+        # 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.0.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: Optional[FluxKontextConditioningField] = 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,14 +385,34 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 dtype=inference_dtype,
             )
 
+            kontext_extension = None
+            if self.kontext_conditioning is not None:
+                if not self.controlnet_vae:
+                    raise ValueError("A VAE (e.g., controlnet_vae) must be provided to use Kontext conditioning.")
+
+                kontext_extension = KontextExtension(
+                    kontext_field=self.kontext_conditioning,
+                    context=context,
+                    vae_field=self.controlnet_vae,
+                    device=TorchDevice.choose_torch_device(),
+                    dtype=inference_dtype,
+                )
+
+            final_img, final_img_ids = x, img_ids
+            original_seq_len = x.shape[1]
+            if kontext_extension is not None:
+                final_img, final_img_ids = kontext_extension.apply(final_img, final_img_ids)
+
             x = denoise(
                 model=transformer,
-                img=x,
-                img_ids=img_ids,
+                img=final_img,
+                img_ids=final_img_ids,
                 pos_regional_prompting_extension=pos_regional_prompting_extension,
                 neg_regional_prompting_extension=neg_regional_prompting_extension,
                 timesteps=timesteps,
-                step_callback=self._build_step_callback(context),
+                step_callback=self._build_step_callback(
+                    context, original_seq_len if kontext_extension is not None else None
+                ),
                 guidance=self.guidance,
                 cfg_scale=cfg_scale,
                 inpaint_extension=inpaint_extension,
@@ -393,6 +422,9 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 img_cond=img_cond,
             )
 
+            if kontext_extension is not None:
+                x = x[:, :original_seq_len, :]  # Keep only the first original_seq_len tokens
+
         x = unpack(x.float(), self.height, self.width)
         return x
 
@@ -863,9 +895,15 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
             yield (lora_info.model, lora.weight)
             del lora_info
 
-    def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
+    def _build_step_callback(
+        self, context: InvocationContext, original_seq_len: Optional[int] = None
+    ) -> Callable[[PipelineIntermediateState], None]:
         def step_callback(state: PipelineIntermediateState) -> None:
-            state.latents = unpack(state.latents.float(), self.height, self.width).squeeze()
+            # Extract only main image tokens if Kontext conditioning was applied
+            latents = state.latents.float()
+            if original_seq_len is not None:
+                latents = latents[:, :original_seq_len, :]
+            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)
 

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/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

@@ -98,9 +98,18 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
                 FROM images
                 LEFT JOIN board_images ON board_images.image_name = images.image_name
                 WHERE 1=1
+                """
+
+        # 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)
+            params.append(board_id)
 
         # Add the category filter
         if categories is not None:

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`
@@ -176,7 +175,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

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,
@@ -196,9 +197,13 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
         # Search term condition
         if search_term:
             query_conditions += """--sql
-            AND images.metadata LIKE ?
+            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
@@ -382,3 +387,96 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
             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:
+        cursor = self._conn.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

@@ -27,6 +27,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 +58,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",
@@ -81,6 +87,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

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

@@ -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 (
@@ -147,7 +148,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:
@@ -646,14 +647,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

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,66 @@
+import sqlite3
+
+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._conn = db.conn
+
+    def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        if model_key_1 == model_key_2:
+            raise ValueError("Cannot relate a model to itself.")
+        a, b = sorted([model_key_1, model_key_2])
+        try:
+            cursor = self._conn.cursor()
+            cursor.execute(
+                "INSERT OR IGNORE INTO model_relationships (model_key_1, model_key_2) VALUES (?, ?)",
+                (a, b),
+            )
+            self._conn.commit()
+        except sqlite3.Error as e:
+            self._conn.rollback()
+            raise e
+
+    def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
+        a, b = sorted([model_key_1, model_key_2])
+        try:
+            cursor = self._conn.cursor()
+            cursor.execute(
+                "DELETE FROM model_relationships WHERE model_key_1 = ? AND model_key_2 = ?",
+                (a, b),
+            )
+            self._conn.commit()
+        except sqlite3.Error as e:
+            self._conn.rollback()
+            raise e
+
+    def get_related_model_keys(self, model_key: str) -> list[str]:
+        cursor = self._conn.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),
+        )
+        return [row[0] for row in cursor.fetchall()]
+
+    def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
+        cursor = self._conn.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,
+        )
+        return [row[0] for row in cursor.fetchall()]

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,10 +47,6 @@ 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__()
@@ -104,11 +101,7 @@ class SqliteSessionQueue(SessionQueueBase):
         return cast(Union[int, None], cursor.fetchone()[0]) or 0
 
     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)
-
-    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
@@ -118,8 +111,12 @@ class SqliteSessionQueue(SessionQueueBase):
             if prepend:
                 priority = self._get_highest_priority(queue_id) + 1
 
-            requested_count = calc_session_count(batch)
-            values_to_insert = prepare_values_to_insert(
+            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,
@@ -127,19 +124,28 @@ class SqliteSessionQueue(SessionQueueBase):
             )
             enqueued_count = len(values_to_insert)
 
-            if requested_count > enqueued_count:
-                values_to_insert = values_to_insert[:max_new_queue_items]
-
-            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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                """,
-                values_to_insert,
-            )
-            self._conn.commit()
+            with self._conn:
+                cursor = self._conn.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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                    values_to_insert,
+                )
+            with self._conn:
+                cursor = self._conn.cursor()
+                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()]
         except Exception:
-            self._conn.rollback()
             raise
         enqueue_result = EnqueueBatchResult(
             queue_id=queue_id,
@@ -147,6 +153,7 @@ class SqliteSessionQueue(SessionQueueBase):
             enqueued=enqueued_count,
             batch=batch,
             priority=priority,
+            item_ids=item_ids,
         )
         self.__invoker.services.events.emit_batch_enqueued(enqueue_result)
         return enqueue_result
@@ -220,6 +227,19 @@ class SqliteSessionQueue(SessionQueueBase):
     ) -> SessionQueueItem:
         try:
             cursor = self._conn.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)
             cursor.execute(
                 """--sql
                 UPDATE session_queue
@@ -331,6 +351,27 @@ class SqliteSessionQueue(SessionQueueBase):
         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
+        try:
+            cursor = self._conn.cursor()
+            cursor.execute(
+                """--sql
+                DELETE
+                FROM session_queue
+                WHERE item_id = ?
+                """,
+                (item_id,),
+            )
+            self._conn.commit()
+        except Exception:
+            self._conn.rollback()
+            raise
+
     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
@@ -428,6 +469,71 @@ class SqliteSessionQueue(SessionQueueBase):
             raise
         return CancelByDestinationResult(canceled=count)
 
+    def delete_by_destination(self, queue_id: str, destination: str) -> DeleteByDestinationResult:
+        try:
+            cursor = self._conn.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,
+            )
+            self._conn.commit()
+        except Exception:
+            self._conn.rollback()
+            raise
+        return DeleteByDestinationResult(deleted=count)
+
+    def delete_all_except_current(self, queue_id: str) -> DeleteAllExceptCurrentResult:
+        try:
+            cursor = self._conn.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,),
+            )
+            self._conn.commit()
+        except Exception:
+            self._conn.rollback()
+            raise
+        return DeleteAllExceptCurrentResult(deleted=count)
+
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
         try:
             cursor = self._conn.cursor()
@@ -543,26 +649,12 @@ class SqliteSessionQueue(SessionQueueBase):
         priority: int,
         cursor: Optional[int] = None,
         status: Optional[QUEUE_ITEM_STATUS] = None,
-    ) -> CursorPaginatedResults[SessionQueueItemDTO]:
+        destination: Optional[str] = None,
+    ) -> CursorPaginatedResults[SessionQueueItem]:
         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
+            SELECT *
             FROM session_queue
             WHERE queue_id = ?
         """
@@ -574,6 +666,12 @@ class SqliteSessionQueue(SessionQueueBase):
                 """
             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 > ?)
@@ -589,7 +687,7 @@ class SqliteSessionQueue(SessionQueueBase):
         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]
+        items = [SessionQueueItem.queue_item_from_dict(dict(result)) for result in results]
         has_more = False
         if len(items) > limit:
             # remove the extra item
@@ -597,6 +695,37 @@ 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"""
+        cursor_ = self._conn.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(

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_util.py

@@ -22,6 +22,7 @@ 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.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -61,6 +62,7 @@ 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.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/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:
@@ -61,6 +64,10 @@ def get_openapi_func(
         # We need to manually add all outputs to the schema - pydantic doesn't add them because they aren't used directly.
         for output in InvocationRegistry.get_output_classes():
             json_schema = output.model_json_schema(mode="serialization", ref_template="#/components/schemas/{model}")
+            # Remove output_metadata that is only used on back-end from the schema
+            if "output_meta" in json_schema["properties"]:
+                json_schema["properties"].pop("output_meta")
+
             move_defs_to_top_level(openapi_schema, json_schema)
             openapi_schema["components"]["schemas"][output.__name__] = json_schema
 

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/extensions/kontext_extension.py

@@ -0,0 +1,139 @@
+import einops
+import torch
+from einops import repeat
+
+from invokeai.app.invocations.fields import FluxKontextConditioningField
+from invokeai.app.invocations.flux_vae_encode import FluxVaeEncodeInvocation
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.sampling_utils import pack
+from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
+
+
+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,
+) -> torch.Tensor:
+    """Generate tensor of image position ids with an optional offset.
+
+    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.
+
+    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
+
+    # 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)
+    y_indices = torch.arange(packed_height, device=device, dtype=dtype)
+    # Broadcast y_indices to match the spatial dimensions [packed_height, 1]
+    img_ids[..., 1] = y_indices[:, None]
+
+    # Create x-coordinate indices (horizontal positions)
+    x_indices = torch.arange(packed_width, device=device, dtype=dtype)
+    # 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_field: FluxKontextConditioningField,
+        context: InvocationContext,
+        vae_field: VAEField,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        """
+        Initializes the KontextExtension, pre-processing the reference image
+        into latents and positional IDs.
+        """
+        self._context = context
+        self._device = device
+        self._dtype = dtype
+        self._vae_field = vae_field
+        self.kontext_field = kontext_field
+
+        # 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 image and prepares its latents and IDs."""
+        image = self._context.images.get_pil(self.kontext_field.image.image_name)
+
+        # Reuse VAE encoding logic from FluxVaeEncodeInvocation
+        vae_info = self._context.models.load(self._vae_field.vae)
+        image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
+        if image_tensor.dim() == 3:
+            image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
+        image_tensor = image_tensor.to(self._device)
+
+        kontext_latents_unpacked = FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+
+        # Extract tensor dimensions with descriptive names
+        # Latent tensor shape: [batch_size, channels, latent_height, latent_width]
+        batch_size, _, latent_height, latent_width = kontext_latents_unpacked.shape
+
+        # Pack the latents and generate IDs. The idx_offset distinguishes these
+        # tokens from the main image's tokens, which have an index of 0.
+        kontext_latents_packed = pack(kontext_latents_unpacked).to(self._device, self._dtype)
+        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,  # Distinguishes reference tokens from main image tokens
+        )
+
+        return kontext_latents_packed, kontext_ids
+
+    def apply(
+        self,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Concatenates the pre-processed kontext data to the main image sequence."""
+        # Ensure batch sizes match, repeating kontext data if necessary for batch operations.
+        if img.shape[0] != self.kontext_latents.shape[0]:
+            self.kontext_latents = self.kontext_latents.repeat(img.shape[0], 1, 1)
+            self.kontext_ids = self.kontext_ids.repeat(img.shape[0], 1, 1)
+
+        # Concatenate along the sequence dimension (dim=1)
+        combined_img = torch.cat([img, self.kontext_latents], dim=1)
+        combined_img_ids = torch.cat([img_ids, self.kontext_ids], dim=1)
+
+        return combined_img, combined_img_ids

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:
@@ -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_")):
@@ -600,6 +634,21 @@ class LlavaOnevisionConfig(DiffusersConfigBase, ModelConfigBase):
         }
 
 
+class ApiModelConfig(MainConfigBase, ModelConfigBase):
+    """Model config for API-based models."""
+
+    format: Literal[ModelFormat.Api] = ModelFormat.Api
+
+    @classmethod
+    def matches(cls, mod: ModelOnDisk) -> bool:
+        # API models are not stored on disk, so we can't match them.
+        return False
+
+    @classmethod
+    def parse(cls, mod: ModelOnDisk) -> dict[str, Any]:
+        raise NotImplementedError("API models are not parsed from disk.")
+
+
 def get_model_discriminator_value(v: Any) -> str:
     """
     Computes the discriminator value for a model config.
@@ -650,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()],
@@ -667,6 +717,7 @@ AnyModelConfig = Annotated[
         Annotated[SigLIPConfig, SigLIPConfig.get_tag()],
         Annotated[FluxReduxConfig, FluxReduxConfig.get_tag()],
         Annotated[LlavaOnevisionConfig, LlavaOnevisionConfig.get_tag()],
+        Annotated[ApiModelConfig, ApiModelConfig.get_tag()],
     ],
     Discriminator(get_model_discriminator_value),
 ]

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/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,6 +4,7 @@ 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.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
 from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
@@ -35,12 +36,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,18 +74,7 @@ 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")):
@@ -94,3 +93,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"]

invokeai/backend/model_manager/omi/omi.py

@@ -0,0 +1,21 @@
+from invokeai.backend.model_manager.model_on_disk import StateDict
+from invokeai.backend.model_manager.omi.vendor.convert.lora import (
+    convert_flux_lora as omi_flux,
+)
+from invokeai.backend.model_manager.omi.vendor.convert.lora import (
+    convert_lora_util as lora_util,
+)
+from invokeai.backend.model_manager.omi.vendor.convert.lora import (
+    convert_sdxl_lora as omi_sdxl,
+)
+from invokeai.backend.model_manager.taxonomy import BaseModelType
+
+
+def convert_from_omi(weights_sd: StateDict, base: BaseModelType):
+    keyset = {
+        BaseModelType.Flux: omi_flux.convert_flux_lora_key_sets(),
+        BaseModelType.StableDiffusionXL: omi_sdxl.convert_sdxl_lora_key_sets(),
+    }[base]
+    source = "omi"
+    target = "legacy_diffusers"
+    return lora_util.__convert(weights_sd, keyset, source, target)  # type: ignore

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_clip.py

@@ -0,0 +1,20 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+
+
+def map_clip(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("text_projection", "text_projection", parent=key_prefix)]
+
+    for k in map_prefix_range("text_model.encoder.layers", "text_model.encoder.layers", parent=key_prefix):
+        keys += [LoraConversionKeySet("mlp.fc1", "mlp.fc1", parent=k)]
+        keys += [LoraConversionKeySet("mlp.fc2", "mlp.fc2", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.k_proj", "self_attn.k_proj", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.out_proj", "self_attn.out_proj", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.q_proj", "self_attn.q_proj", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.v_proj", "self_attn.v_proj", parent=k)]
+
+    return keys

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_flux_lora.py

@@ -0,0 +1,84 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_clip import map_clip
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_t5 import map_t5
+
+
+def __map_double_transformer_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("img_attn.qkv.0", "attn.to_q", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_attn.qkv.1", "attn.to_k", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_attn.qkv.2", "attn.to_v", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("txt_attn.qkv.0", "attn.add_q_proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_attn.qkv.1", "attn.add_k_proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_attn.qkv.2", "attn.add_v_proj", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("img_attn.proj", "attn.to_out.0", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_mlp.0", "ff.net.0.proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_mlp.2", "ff.net.2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_mod.lin", "norm1.linear", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("txt_attn.proj", "attn.to_add_out", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_mlp.0", "ff_context.net.0.proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_mlp.2", "ff_context.net.2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_mod.lin", "norm1_context.linear", parent=key_prefix)]
+
+    return keys
+
+
+def __map_single_transformer_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("linear1.0", "attn.to_q", parent=key_prefix)]
+    keys += [LoraConversionKeySet("linear1.1", "attn.to_k", parent=key_prefix)]
+    keys += [LoraConversionKeySet("linear1.2", "attn.to_v", parent=key_prefix)]
+    keys += [LoraConversionKeySet("linear1.3", "proj_mlp", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("linear2", "proj_out", parent=key_prefix)]
+    keys += [LoraConversionKeySet("modulation.lin", "norm.linear", parent=key_prefix)]
+
+    return keys
+
+
+def __map_transformer(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("txt_in", "context_embedder", parent=key_prefix)]
+    keys += [
+        LoraConversionKeySet("final_layer.adaLN_modulation.1", "norm_out.linear", parent=key_prefix, swap_chunks=True)
+    ]
+    keys += [LoraConversionKeySet("final_layer.linear", "proj_out", parent=key_prefix)]
+    keys += [
+        LoraConversionKeySet("guidance_in.in_layer", "time_text_embed.guidance_embedder.linear_1", parent=key_prefix)
+    ]
+    keys += [
+        LoraConversionKeySet("guidance_in.out_layer", "time_text_embed.guidance_embedder.linear_2", parent=key_prefix)
+    ]
+    keys += [LoraConversionKeySet("vector_in.in_layer", "time_text_embed.text_embedder.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("vector_in.out_layer", "time_text_embed.text_embedder.linear_2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("time_in.in_layer", "time_text_embed.timestep_embedder.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("time_in.out_layer", "time_text_embed.timestep_embedder.linear_2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_in.proj", "x_embedder", parent=key_prefix)]
+
+    for k in map_prefix_range("double_blocks", "transformer_blocks", parent=key_prefix):
+        keys += __map_double_transformer_block(k)
+
+    for k in map_prefix_range("single_blocks", "single_transformer_blocks", parent=key_prefix):
+        keys += __map_single_transformer_block(k)
+
+    return keys
+
+
+def convert_flux_lora_key_sets() -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("bundle_emb", "bundle_emb")]
+    keys += __map_transformer(LoraConversionKeySet("transformer", "lora_transformer"))
+    keys += map_clip(LoraConversionKeySet("clip_l", "lora_te1"))
+    keys += map_t5(LoraConversionKeySet("t5", "lora_te2"))
+
+    return keys

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_lora_util.py

@@ -0,0 +1,217 @@
+import torch
+from torch import Tensor
+from typing_extensions import Self
+
+
+class LoraConversionKeySet:
+    def __init__(
+        self,
+        omi_prefix: str,
+        diffusers_prefix: str,
+        legacy_diffusers_prefix: str | None = None,
+        parent: Self | None = None,
+        swap_chunks: bool = False,
+        filter_is_last: bool | None = None,
+        next_omi_prefix: str | None = None,
+        next_diffusers_prefix: str | None = None,
+    ):
+        if parent is not None:
+            self.omi_prefix = combine(parent.omi_prefix, omi_prefix)
+            self.diffusers_prefix = combine(parent.diffusers_prefix, diffusers_prefix)
+        else:
+            self.omi_prefix = omi_prefix
+            self.diffusers_prefix = diffusers_prefix
+
+        if legacy_diffusers_prefix is None:
+            self.legacy_diffusers_prefix = self.diffusers_prefix.replace(".", "_")
+        elif parent is not None:
+            self.legacy_diffusers_prefix = combine(parent.legacy_diffusers_prefix, legacy_diffusers_prefix).replace(
+                ".", "_"
+            )
+        else:
+            self.legacy_diffusers_prefix = legacy_diffusers_prefix
+
+        self.parent = parent
+        self.swap_chunks = swap_chunks
+        self.filter_is_last = filter_is_last
+        self.prefix = parent
+
+        if next_omi_prefix is None and parent is not None:
+            self.next_omi_prefix = parent.next_omi_prefix
+            self.next_diffusers_prefix = parent.next_diffusers_prefix
+            self.next_legacy_diffusers_prefix = parent.next_legacy_diffusers_prefix
+        elif next_omi_prefix is not None and parent is not None:
+            self.next_omi_prefix = combine(parent.omi_prefix, next_omi_prefix)
+            self.next_diffusers_prefix = combine(parent.diffusers_prefix, next_diffusers_prefix)
+            self.next_legacy_diffusers_prefix = combine(parent.legacy_diffusers_prefix, next_diffusers_prefix).replace(
+                ".", "_"
+            )
+        elif next_omi_prefix is not None and parent is None:
+            self.next_omi_prefix = next_omi_prefix
+            self.next_diffusers_prefix = next_diffusers_prefix
+            self.next_legacy_diffusers_prefix = next_diffusers_prefix.replace(".", "_")
+        else:
+            self.next_omi_prefix = None
+            self.next_diffusers_prefix = None
+            self.next_legacy_diffusers_prefix = None
+
+    def __get_omi(self, in_prefix: str, key: str) -> str:
+        return self.omi_prefix + key.removeprefix(in_prefix)
+
+    def __get_diffusers(self, in_prefix: str, key: str) -> str:
+        return self.diffusers_prefix + key.removeprefix(in_prefix)
+
+    def __get_legacy_diffusers(self, in_prefix: str, key: str) -> str:
+        key = self.legacy_diffusers_prefix + key.removeprefix(in_prefix)
+
+        suffix = key[key.rfind(".") :]
+        if suffix not in [".alpha", ".dora_scale"]:  # some keys only have a single . in the suffix
+            suffix = key[key.removesuffix(suffix).rfind(".") :]
+        key = key.removesuffix(suffix)
+
+        return key.replace(".", "_") + suffix
+
+    def get_key(self, in_prefix: str, key: str, target: str) -> str:
+        if target == "omi":
+            return self.__get_omi(in_prefix, key)
+        elif target == "diffusers":
+            return self.__get_diffusers(in_prefix, key)
+        elif target == "legacy_diffusers":
+            return self.__get_legacy_diffusers(in_prefix, key)
+        return key
+
+    def __str__(self) -> str:
+        return f"omi: {self.omi_prefix}, diffusers: {self.diffusers_prefix}, legacy: {self.legacy_diffusers_prefix}"
+
+
+def combine(left: str, right: str) -> str:
+    left = left.rstrip(".")
+    right = right.lstrip(".")
+    if left == "" or left is None:
+        return right
+    elif right == "" or right is None:
+        return left
+    else:
+        return left + "." + right
+
+
+def map_prefix_range(
+    omi_prefix: str,
+    diffusers_prefix: str,
+    parent: LoraConversionKeySet,
+) -> list[LoraConversionKeySet]:
+    # 100 should be a safe upper bound. increase if it's not enough in the future
+    return [
+        LoraConversionKeySet(
+            omi_prefix=f"{omi_prefix}.{i}",
+            diffusers_prefix=f"{diffusers_prefix}.{i}",
+            parent=parent,
+            next_omi_prefix=f"{omi_prefix}.{i + 1}",
+            next_diffusers_prefix=f"{diffusers_prefix}.{i + 1}",
+        )
+        for i in range(100)
+    ]
+
+
+def __convert(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+    source: str,
+    target: str,
+) -> dict[str, Tensor]:
+    out_states = {}
+
+    if source == target:
+        return dict(state_dict)
+
+    # TODO: maybe replace with a non O(n^2) algorithm
+    for key, tensor in state_dict.items():
+        for key_set in key_sets:
+            in_prefix = ""
+
+            if source == "omi":
+                in_prefix = key_set.omi_prefix
+            elif source == "diffusers":
+                in_prefix = key_set.diffusers_prefix
+            elif source == "legacy_diffusers":
+                in_prefix = key_set.legacy_diffusers_prefix
+
+            if not key.startswith(in_prefix):
+                continue
+
+            if key_set.filter_is_last is not None:
+                next_prefix = None
+                if source == "omi":
+                    next_prefix = key_set.next_omi_prefix
+                elif source == "diffusers":
+                    next_prefix = key_set.next_diffusers_prefix
+                elif source == "legacy_diffusers":
+                    next_prefix = key_set.next_legacy_diffusers_prefix
+
+                is_last = not any(k.startswith(next_prefix) for k in state_dict)
+                if key_set.filter_is_last != is_last:
+                    continue
+
+            name = key_set.get_key(in_prefix, key, target)
+
+            can_swap_chunks = target == "omi" or source == "omi"
+            if key_set.swap_chunks and name.endswith(".lora_up.weight") and can_swap_chunks:
+                chunk_0, chunk_1 = tensor.chunk(2, dim=0)
+                tensor = torch.cat([chunk_1, chunk_0], dim=0)
+
+            out_states[name] = tensor
+
+            break  # only map the first matching key set
+
+    return out_states
+
+
+def __detect_source(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> str:
+    omi_count = 0
+    diffusers_count = 0
+    legacy_diffusers_count = 0
+
+    for key in state_dict:
+        for key_set in key_sets:
+            if key.startswith(key_set.omi_prefix):
+                omi_count += 1
+            if key.startswith(key_set.diffusers_prefix):
+                diffusers_count += 1
+            if key.startswith(key_set.legacy_diffusers_prefix):
+                legacy_diffusers_count += 1
+
+    if omi_count > diffusers_count and omi_count > legacy_diffusers_count:
+        return "omi"
+    if diffusers_count > omi_count and diffusers_count > legacy_diffusers_count:
+        return "diffusers"
+    if legacy_diffusers_count > omi_count and legacy_diffusers_count > diffusers_count:
+        return "legacy_diffusers"
+
+    return ""
+
+
+def convert_to_omi(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> dict[str, Tensor]:
+    source = __detect_source(state_dict, key_sets)
+    return __convert(state_dict, key_sets, source, "omi")
+
+
+def convert_to_diffusers(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> dict[str, Tensor]:
+    source = __detect_source(state_dict, key_sets)
+    return __convert(state_dict, key_sets, source, "diffusers")
+
+
+def convert_to_legacy_diffusers(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> dict[str, Tensor]:
+    source = __detect_source(state_dict, key_sets)
+    return __convert(state_dict, key_sets, source, "legacy_diffusers")

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_sdxl_lora.py

@@ -0,0 +1,125 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_clip import map_clip
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+
+
+def __map_unet_resnet_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("emb_layers.1", "time_emb_proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("in_layers.2", "conv1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("out_layers.3", "conv2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("skip_connection", "conv_shortcut", parent=key_prefix)]
+
+    return keys
+
+
+def __map_unet_attention_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("proj_in", "proj_in", parent=key_prefix)]
+    keys += [LoraConversionKeySet("proj_out", "proj_out", parent=key_prefix)]
+    for k in map_prefix_range("transformer_blocks", "transformer_blocks", parent=key_prefix):
+        keys += [LoraConversionKeySet("attn1.to_q", "attn1.to_q", parent=k)]
+        keys += [LoraConversionKeySet("attn1.to_k", "attn1.to_k", parent=k)]
+        keys += [LoraConversionKeySet("attn1.to_v", "attn1.to_v", parent=k)]
+        keys += [LoraConversionKeySet("attn1.to_out.0", "attn1.to_out.0", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_q", "attn2.to_q", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_k", "attn2.to_k", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_v", "attn2.to_v", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_out.0", "attn2.to_out.0", parent=k)]
+        keys += [LoraConversionKeySet("ff.net.0.proj", "ff.net.0.proj", parent=k)]
+        keys += [LoraConversionKeySet("ff.net.2", "ff.net.2", parent=k)]
+
+    return keys
+
+
+def __map_unet_down_blocks(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("1.0", "0.resnets.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("2.0", "0.resnets.1", parent=key_prefix))
+    keys += [LoraConversionKeySet("3.0.op", "0.downsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("4.0", "1.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("4.1", "1.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("5.0", "1.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("5.1", "1.attentions.1", parent=key_prefix))
+    keys += [LoraConversionKeySet("6.0.op", "1.downsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("7.0", "2.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("7.1", "2.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("8.0", "2.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("8.1", "2.attentions.1", parent=key_prefix))
+
+    return keys
+
+
+def __map_unet_mid_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("0", "resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("1", "attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("2", "resnets.1", parent=key_prefix))
+
+    return keys
+
+
+def __map_unet_up_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("0.0", "0.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("0.1", "0.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("1.0", "0.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("1.1", "0.attentions.1", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("2.0", "0.resnets.2", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("2.1", "0.attentions.2", parent=key_prefix))
+    keys += [LoraConversionKeySet("2.2.conv", "0.upsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("3.0", "1.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("3.1", "1.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("4.0", "1.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("4.1", "1.attentions.1", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("5.0", "1.resnets.2", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("5.1", "1.attentions.2", parent=key_prefix))
+    keys += [LoraConversionKeySet("5.2.conv", "1.upsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("6.0", "2.resnets.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("7.0", "2.resnets.1", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("8.0", "2.resnets.2", parent=key_prefix))
+
+    return keys
+
+
+def __map_unet(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("input_blocks.0.0", "conv_in", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("time_embed.0", "time_embedding.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("time_embed.2", "time_embedding.linear_2", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("label_emb.0.0", "add_embedding.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("label_emb.0.2", "add_embedding.linear_2", parent=key_prefix)]
+
+    keys += __map_unet_down_blocks(LoraConversionKeySet("input_blocks", "down_blocks", parent=key_prefix))
+    keys += __map_unet_mid_block(LoraConversionKeySet("middle_block", "mid_block", parent=key_prefix))
+    keys += __map_unet_up_block(LoraConversionKeySet("output_blocks", "up_blocks", parent=key_prefix))
+
+    keys += [LoraConversionKeySet("out.0", "conv_norm_out", parent=key_prefix)]
+    keys += [LoraConversionKeySet("out.2", "conv_out", parent=key_prefix)]
+
+    return keys
+
+
+def convert_sdxl_lora_key_sets() -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("bundle_emb", "bundle_emb")]
+    keys += __map_unet(LoraConversionKeySet("unet", "lora_unet"))
+    keys += map_clip(LoraConversionKeySet("clip_l", "lora_te1"))
+    keys += map_clip(LoraConversionKeySet("clip_g", "lora_te2"))
+
+    return keys

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_t5.py

@@ -0,0 +1,19 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+
+
+def map_t5(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    for k in map_prefix_range("encoder.block", "encoder.block", parent=key_prefix):
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.k", "layer.0.SelfAttention.k", parent=k)]
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.o", "layer.0.SelfAttention.o", parent=k)]
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.q", "layer.0.SelfAttention.q", parent=k)]
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.v", "layer.0.SelfAttention.v", parent=k)]
+        keys += [LoraConversionKeySet("layer.1.DenseReluDense.wi_0", "layer.1.DenseReluDense.wi_0", parent=k)]
+        keys += [LoraConversionKeySet("layer.1.DenseReluDense.wi_1", "layer.1.DenseReluDense.wi_1", parent=k)]
+        keys += [LoraConversionKeySet("layer.1.DenseReluDense.wo", "layer.1.DenseReluDense.wo", parent=k)]
+
+    return keys

invokeai/backend/model_manager/omi/vendor/model_spec/architecture.py

@@ -0,0 +1,31 @@
+stable_diffusion_1_lora = "stable-diffusion-v1/lora"
+stable_diffusion_1_inpainting_lora = "stable-diffusion-v1-inpainting/lora"
+
+stable_diffusion_2_512_lora = "stable-diffusion-v2-512/lora"
+stable_diffusion_2_768_v_lora = "stable-diffusion-v2-768-v/lora"
+stable_diffusion_2_depth_lora = "stable-diffusion-v2-depth/lora"
+stable_diffusion_2_inpainting_lora = "stable-diffusion-v2-inpainting/lora"
+
+stable_diffusion_3_medium_lora = "stable-diffusion-v3-medium/lora"
+stable_diffusion_35_medium_lora = "stable-diffusion-v3.5-medium/lora"
+stable_diffusion_35_large_lora = "stable-diffusion-v3.5-large/lora"
+
+stable_diffusion_xl_1_lora = "stable-diffusion-xl-v1-base/lora"
+stable_diffusion_xl_1_inpainting_lora = "stable-diffusion-xl-v1-base-inpainting/lora"
+
+wuerstchen_2_lora = "wuerstchen-v2-prior/lora"
+stable_cascade_1_stage_a_lora = "stable-cascade-v1-stage-a/lora"
+stable_cascade_1_stage_b_lora = "stable-cascade-v1-stage-b/lora"
+stable_cascade_1_stage_c_lora = "stable-cascade-v1-stage-c/lora"
+
+pixart_alpha_lora = "pixart-alpha/lora"
+pixart_sigma_lora = "pixart-sigma/lora"
+
+flux_dev_1_lora = "Flux.1-dev/lora"
+flux_fill_dev_1_lora = "Flux.1-fill-dev/lora"
+
+sana_lora = "sana/lora"
+
+hunyuan_video_lora = "hunyuan-video/lora"
+
+hi_dream_i1_lora = "hidream-i1/lora"

invokeai/backend/model_manager/starter_models.py

@@ -23,7 +23,7 @@ class StarterModel(StarterModelWithoutDependencies):
     dependencies: Optional[list[StarterModelWithoutDependencies]] = None
 
 
-class StarterModelBundles(BaseModel):
+class StarterModelBundle(BaseModel):
     name: str
     models: list[StarterModel]
 
@@ -109,7 +109,7 @@ flux_vae = StarterModel(
 
 # region: Main
 flux_schnell_quantized = StarterModel(
-    name="FLUX Schnell (Quantized)",
+    name="FLUX.1 schnell (quantized)",
     base=BaseModelType.Flux,
     source="InvokeAI/flux_schnell::transformer/bnb_nf4/flux1-schnell-bnb_nf4.safetensors",
     description="FLUX schnell transformer quantized to bitsandbytes NF4 format. Total size with dependencies: ~12GB",
@@ -117,7 +117,7 @@ flux_schnell_quantized = StarterModel(
     dependencies=[t5_8b_quantized_encoder, flux_vae, clip_l_encoder],
 )
 flux_dev_quantized = StarterModel(
-    name="FLUX Dev (Quantized)",
+    name="FLUX.1 dev (quantized)",
     base=BaseModelType.Flux,
     source="InvokeAI/flux_dev::transformer/bnb_nf4/flux1-dev-bnb_nf4.safetensors",
     description="FLUX dev transformer quantized to bitsandbytes NF4 format. Total size with dependencies: ~12GB",
@@ -125,7 +125,7 @@ flux_dev_quantized = StarterModel(
     dependencies=[t5_8b_quantized_encoder, flux_vae, clip_l_encoder],
 )
 flux_schnell = StarterModel(
-    name="FLUX Schnell",
+    name="FLUX.1 schnell",
     base=BaseModelType.Flux,
     source="InvokeAI/flux_schnell::transformer/base/flux1-schnell.safetensors",
     description="FLUX schnell transformer in bfloat16. Total size with dependencies: ~33GB",
@@ -133,13 +133,21 @@ flux_schnell = StarterModel(
     dependencies=[t5_base_encoder, flux_vae, clip_l_encoder],
 )
 flux_dev = StarterModel(
-    name="FLUX Dev",
+    name="FLUX.1 dev",
     base=BaseModelType.Flux,
     source="InvokeAI/flux_dev::transformer/base/flux1-dev.safetensors",
     description="FLUX dev transformer in bfloat16. Total size with dependencies: ~33GB",
     type=ModelType.Main,
     dependencies=[t5_base_encoder, flux_vae, clip_l_encoder],
 )
+flux_kontext = StarterModel(
+    name="FLUX.1 Kontext dev",
+    base=BaseModelType.Flux,
+    source="black-forest-labs/FLUX.1-Kontext-dev::flux1-kontext-dev.safetensors",
+    description="FLUX.1 Kontext dev transformer in bfloat16. Total size with dependencies: ~33GB",
+    type=ModelType.Main,
+    dependencies=[t5_base_encoder, flux_vae, clip_l_encoder],
+)
 sd35_medium = StarterModel(
     name="SD3.5 Medium",
     base=BaseModelType.StableDiffusion3,
@@ -297,6 +305,15 @@ ip_adapter_sdxl = StarterModel(
     dependencies=[ip_adapter_sdxl_image_encoder],
     previous_names=["IP Adapter SDXL"],
 )
+ip_adapter_plus_sdxl = StarterModel(
+    name="Precise Reference (IP Adapter Plus ViT-H)",
+    base=BaseModelType.StableDiffusionXL,
+    source="https://huggingface.co/InvokeAI/ip-adapter-plus_sdxl_vit-h/resolve/main/ip-adapter-plus_sdxl_vit-h.safetensors",
+    description="References images with a higher degree of precision.",
+    type=ModelType.IPAdapter,
+    dependencies=[ip_adapter_sdxl_image_encoder],
+    previous_names=["IP Adapter Plus SDXL"],
+)
 ip_adapter_flux = StarterModel(
     name="Standard Reference (XLabs FLUX IP-Adapter v2)",
     base=BaseModelType.Flux,
@@ -647,6 +664,7 @@ flux_fill = StarterModel(
 # List of starter models, displayed on the frontend.
 # The order/sort of this list is not changed by the frontend - set it how you want it here.
 STARTER_MODELS: list[StarterModel] = [
+    flux_kontext,
     flux_schnell_quantized,
     flux_dev_quantized,
     flux_schnell,
@@ -672,6 +690,7 @@ STARTER_MODELS: list[StarterModel] = [
     ip_adapter_plus_sd1,
     ip_adapter_plus_face_sd1,
     ip_adapter_sdxl,
+    ip_adapter_plus_sdxl,
     ip_adapter_flux,
     qr_code_cnet_sd1,
     qr_code_cnet_sdxl,
@@ -744,6 +763,7 @@ sdxl_bundle: list[StarterModel] = [
     juggernaut_sdxl,
     sdxl_fp16_vae_fix,
     ip_adapter_sdxl,
+    ip_adapter_plus_sdxl,
     canny_sdxl,
     depth_sdxl,
     softedge_sdxl,
@@ -765,12 +785,13 @@ flux_bundle: list[StarterModel] = [
     flux_depth_control_lora,
     flux_redux,
     flux_fill,
+    flux_kontext,
 ]
 
-STARTER_BUNDLES: dict[str, list[StarterModel]] = {
-    BaseModelType.StableDiffusion1: sd1_bundle,
-    BaseModelType.StableDiffusionXL: sdxl_bundle,
-    BaseModelType.Flux: flux_bundle,
+STARTER_BUNDLES: dict[str, StarterModelBundle] = {
+    BaseModelType.StableDiffusion1: StarterModelBundle(name="Stable Diffusion 1.5", models=sd1_bundle),
+    BaseModelType.StableDiffusionXL: StarterModelBundle(name="SDXL", models=sdxl_bundle),
+    BaseModelType.Flux: StarterModelBundle(name="FLUX.1 dev", models=flux_bundle),
 }
 
 assert len(STARTER_MODELS) == len({m.source for m in STARTER_MODELS}), "Duplicate starter models"

invokeai/backend/model_manager/taxonomy.py

@@ -26,7 +26,10 @@ class BaseModelType(str, Enum):
     StableDiffusionXLRefiner = "sdxl-refiner"
     Flux = "flux"
     CogView4 = "cogview4"
-    # Kandinsky2_1 = "kandinsky-2.1"
+    Imagen3 = "imagen3"
+    Imagen4 = "imagen4"
+    ChatGPT4o = "chatgpt-4o"
+    FluxKontext = "flux-kontext"
 
 
 class ModelType(str, Enum):
@@ -86,6 +89,7 @@ class ModelVariantType(str, Enum):
 class ModelFormat(str, Enum):
     """Storage format of model."""
 
+    OMI = "omi"
     Diffusers = "diffusers"
     Checkpoint = "checkpoint"
     LyCORIS = "lycoris"
@@ -98,6 +102,7 @@ class ModelFormat(str, Enum):
     BnbQuantizedLlmInt8b = "bnb_quantized_int8b"
     BnbQuantizednf4b = "bnb_quantized_nf4b"
     GGUFQuantized = "gguf_quantized"
+    Api = "api"
 
 
 class SchedulerPredictionType(str, Enum):
@@ -134,6 +139,7 @@ class FluxLoRAFormat(str, Enum):
     Kohya = "flux.kohya"
     OneTrainer = "flux.onetrainer"
     Control = "flux.control"
+    AIToolkit = "flux.aitoolkit"
 
 
 AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, None]

invokeai/backend/model_patcher.py

@@ -46,6 +46,10 @@ class ModelPatcher:
         text_encoder: Union[CLIPTextModel, CLIPTextModelWithProjection],
         ti_list: List[Tuple[str, TextualInversionModelRaw]],
     ) -> Iterator[Tuple[CLIPTokenizer, TextualInversionManager]]:
+        if len(ti_list) == 0:
+            yield tokenizer, TextualInversionManager(tokenizer)
+            return
+
         init_tokens_count = None
         new_tokens_added = None
 

invokeai/backend/patches/layer_patcher.py

@@ -1,3 +1,4 @@
+import re
 from contextlib import contextmanager
 from typing import Dict, Iterable, Optional, Tuple
 
@@ -7,6 +8,7 @@ from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
 from invokeai.backend.patches.layers.flux_control_lora_layer import FluxControlLoRALayer
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.patches.pad_with_zeros import pad_with_zeros
+from invokeai.backend.util import InvokeAILogger
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
 
@@ -23,6 +25,7 @@ class LayerPatcher:
         cached_weights: Optional[Dict[str, torch.Tensor]] = None,
         force_direct_patching: bool = False,
         force_sidecar_patching: bool = False,
+        suppress_warning_layers: Optional[re.Pattern] = None,
     ):
         """Apply 'smart' model patching that chooses whether to use direct patching or a sidecar wrapper for each
         module.
@@ -44,6 +47,7 @@ class LayerPatcher:
                     dtype=dtype,
                     force_direct_patching=force_direct_patching,
                     force_sidecar_patching=force_sidecar_patching,
+                    suppress_warning_layers=suppress_warning_layers,
                 )
 
             yield
@@ -70,6 +74,7 @@ class LayerPatcher:
         dtype: torch.dtype,
         force_direct_patching: bool,
         force_sidecar_patching: bool,
+        suppress_warning_layers: Optional[re.Pattern] = None,
     ):
         """Apply a single LoRA patch to a model using the 'smart' patching strategy that chooses whether to use direct
         patching or a sidecar wrapper for each module.
@@ -89,9 +94,17 @@ class LayerPatcher:
             if not layer_key.startswith(prefix):
                 continue
 
-            module_key, module = LayerPatcher._get_submodule(
-                model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
-            )
+            try:
+                module_key, module = LayerPatcher._get_submodule(
+                    model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
+                )
+            except AttributeError:
+                if suppress_warning_layers and suppress_warning_layers.search(layer_key):
+                    pass
+                else:
+                    logger = InvokeAILogger.get_logger(LayerPatcher.__name__)
+                    logger.warning("Failed to find module for LoRA layer key: %s", layer_key)
+                continue
 
             # Decide whether to use direct patching or a sidecar patch.
             # Direct patching is preferred, because it results in better runtime speed.

invokeai/backend/patches/layers/lora_layer.py

@@ -19,6 +19,7 @@ class LoRALayer(LoRALayerBase):
         self.up = up
         self.mid = mid
         self.down = down
+        self.are_ranks_equal = up.shape[1] == down.shape[0]
 
     @classmethod
     def from_state_dict_values(
@@ -58,12 +59,42 @@ class LoRALayer(LoRALayerBase):
     def _rank(self) -> int:
         return self.down.shape[0]
 
+    def fuse_weights(self, up: torch.Tensor, down: torch.Tensor) -> torch.Tensor:
+        """
+        Fuse the weights of the up and down matrices of a LoRA layer with different ranks.
+
+        Since the Huggingface implementation of KQV projections are fused, when we convert to Kohya format
+        the LoRA weights have different ranks. This function handles the fusion of these differently sized
+        matrices.
+        """
+
+        fused_lora = torch.zeros((up.shape[0], down.shape[1]), device=down.device, dtype=down.dtype)
+        rank_diff = down.shape[0] / up.shape[1]
+
+        if rank_diff > 1:
+            rank_diff = down.shape[0] / up.shape[1]
+            w_down = down.chunk(int(rank_diff), dim=0)
+            for w_down_chunk in w_down:
+                fused_lora = fused_lora + (torch.mm(up, w_down_chunk))
+        else:
+            rank_diff = up.shape[1] / down.shape[0]
+            w_up = up.chunk(int(rank_diff), dim=0)
+            for w_up_chunk in w_up:
+                fused_lora = fused_lora + (torch.mm(w_up_chunk, down))
+
+        return fused_lora
+
     def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
         if self.mid is not None:
             up = self.up.reshape(self.up.shape[0], self.up.shape[1])
             down = self.down.reshape(self.down.shape[0], self.down.shape[1])
             weight = torch.einsum("m n w h, i m, n j -> i j w h", self.mid, up, down)
         else:
+            # up matrix and down matrix have different ranks so we can't simply multiply them
+            if not self.are_ranks_equal:
+                weight = self.fuse_weights(self.up, self.down)
+                return weight
+
             weight = self.up.reshape(self.up.shape[0], -1) @ self.down.reshape(self.down.shape[0], -1)
 
         return weight

invokeai/backend/patches/lora_conversions/flux_aitoolkit_lora_conversion_utils.py

@@ -0,0 +1,63 @@
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import _group_by_layer
+from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.util import InvokeAILogger
+
+
+def is_state_dict_likely_in_flux_aitoolkit_format(state_dict: dict[str, Any], metadata: dict[str, Any] = None) -> bool:
+    if metadata:
+        try:
+            software = json.loads(metadata.get("software", "{}"))
+        except json.JSONDecodeError:
+            return False
+        return software.get("name") == "ai-toolkit"
+    # metadata got lost somewhere
+    return any("diffusion_model" == k.split(".", 1)[0] for k in state_dict.keys())
+
+
+@dataclass
+class GroupedStateDict:
+    transformer: dict[str, Any] = field(default_factory=dict)
+    # might also grow CLIP and T5 submodels
+
+
+def _group_state_by_submodel(state_dict: dict[str, Any]) -> GroupedStateDict:
+    logger = InvokeAILogger.get_logger()
+    grouped = GroupedStateDict()
+    for key, value in state_dict.items():
+        submodel_name, param_name = key.split(".", 1)
+        match submodel_name:
+            case "diffusion_model":
+                grouped.transformer[param_name] = value
+            case _:
+                logger.warning(f"Unexpected submodel name: {submodel_name}")
+    return grouped
+
+
+def _rename_peft_lora_keys(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Renames keys from the PEFT LoRA format to the InvokeAI format."""
+    renamed_state_dict = {}
+    for key, value in state_dict.items():
+        renamed_key = key.replace(".lora_A.", ".lora_down.").replace(".lora_B.", ".lora_up.")
+        renamed_state_dict[renamed_key] = value
+    return renamed_state_dict
+
+
+def lora_model_from_flux_aitoolkit_state_dict(state_dict: dict[str, torch.Tensor]) -> ModelPatchRaw:
+    state_dict = _rename_peft_lora_keys(state_dict)
+    by_layer = _group_by_layer(state_dict)
+    by_model = _group_state_by_submodel(by_layer)
+
+    layers: dict[str, BaseLayerPatch] = {}
+    for layer_key, layer_state_dict in by_model.transformer.items():
+        layers[FLUX_LORA_TRANSFORMER_PREFIX + layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
+
+    return ModelPatchRaw(layers=layers)

invokeai/backend/patches/lora_conversions/flux_kohya_lora_conversion_utils.py

@@ -20,6 +20,14 @@ from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 FLUX_KOHYA_TRANSFORMER_KEY_REGEX = (
     r"lora_unet_(\w+_blocks)_(\d+)_(img_attn|img_mlp|img_mod|txt_attn|txt_mlp|txt_mod|linear1|linear2|modulation)_?(.*)"
 )
+
+# A regex pattern that matches all of the last layer keys in the Kohya FLUX LoRA format.
+# Example keys:
+#   lora_unet_final_layer_linear.alpha
+#   lora_unet_final_layer_linear.lora_down.weight
+#   lora_unet_final_layer_linear.lora_up.weight
+FLUX_KOHYA_LAST_LAYER_KEY_REGEX = r"lora_unet_final_layer_(linear|linear1|linear2)_?(.*)"
+
 # A regex pattern that matches all of the CLIP keys in the Kohya FLUX LoRA format.
 # Example keys:
 #   lora_te1_text_model_encoder_layers_0_mlp_fc1.alpha
@@ -44,6 +52,7 @@ def is_state_dict_likely_in_flux_kohya_format(state_dict: Dict[str, Any]) -> boo
     """
     return all(
         re.match(FLUX_KOHYA_TRANSFORMER_KEY_REGEX, k)
+        or re.match(FLUX_KOHYA_LAST_LAYER_KEY_REGEX, k)
         or re.match(FLUX_KOHYA_CLIP_KEY_REGEX, k)
         or re.match(FLUX_KOHYA_T5_KEY_REGEX, k)
         for k in state_dict.keys()
@@ -65,6 +74,9 @@ def lora_model_from_flux_kohya_state_dict(state_dict: Dict[str, torch.Tensor]) -
     t5_grouped_sd: dict[str, dict[str, torch.Tensor]] = {}
     for layer_name, layer_state_dict in grouped_state_dict.items():
         if layer_name.startswith("lora_unet"):
+            # Skip the final layer. This is incompatible with current model definition.
+            if layer_name.startswith("lora_unet_final_layer"):
+                continue
             transformer_grouped_sd[layer_name] = layer_state_dict
         elif layer_name.startswith("lora_te1"):
             clip_grouped_sd[layer_name] = layer_state_dict

invokeai/backend/patches/lora_conversions/formats.py

@@ -1,4 +1,7 @@
 from invokeai.backend.model_manager.taxonomy import FluxLoRAFormat
+from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_aitoolkit_format,
+)
 from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import is_state_dict_likely_flux_control
 from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import (
     is_state_dict_likely_in_flux_diffusers_format,
@@ -11,7 +14,7 @@ from invokeai.backend.patches.lora_conversions.flux_onetrainer_lora_conversion_u
 )
 
 
-def flux_format_from_state_dict(state_dict):
+def flux_format_from_state_dict(state_dict: dict, metadata: dict | None = None) -> FluxLoRAFormat | None:
     if is_state_dict_likely_in_flux_kohya_format(state_dict):
         return FluxLoRAFormat.Kohya
     elif is_state_dict_likely_in_flux_onetrainer_format(state_dict):
@@ -20,5 +23,7 @@ def flux_format_from_state_dict(state_dict):
         return FluxLoRAFormat.Diffusers
     elif is_state_dict_likely_flux_control(state_dict):
         return FluxLoRAFormat.Control
+    elif is_state_dict_likely_in_flux_aitoolkit_format(state_dict, metadata):
+        return FluxLoRAFormat.AIToolkit
     else:
         return None

invokeai/backend/quantization/gguf/utils.py

@@ -5,7 +5,8 @@ from typing import Callable, Optional, Union
 import gguf
 import torch
 
-TORCH_COMPATIBLE_QTYPES = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
+# should not be a Set until this is resolved: https://github.com/pytorch/pytorch/issues/145761
+TORCH_COMPATIBLE_QTYPES = [None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]
 
 # K Quants #
 QK_K = 256

invokeai/backend/rectified_flow/rectified_flow_inpaint_extension.py

@@ -30,18 +30,13 @@ class RectifiedFlowInpaintExtension:
     def _apply_mask_gradient_adjustment(self, t_prev: float) -> torch.Tensor:
         """Applies inpaint mask gradient adjustment and returns the inpaint mask to be used at the current timestep."""
         # As we progress through the denoising process, we promote gradient regions of the mask to have a full weight of
-        # 1.0. This helps to produce more coherent seams around the inpainted region. We experimented with a (small)
-        # number of promotion strategies (e.g. gradual promotion based on timestep), but found that a simple cutoff
-        # threshold worked well.
+        # 1.0. This helps to produce more coherent seams around the inpainted region.
+
         # We use a small epsilon to avoid any potential issues with floating point precision.
         eps = 1e-4
-        mask_gradient_t_cutoff = 0.5
-        if t_prev > mask_gradient_t_cutoff:
-            # Early in the denoising process, use the inpaint mask as-is.
-            return self._inpaint_mask
-        else:
-            # After the cut-off, promote all non-zero mask values to 1.0.
-            mask = self._inpaint_mask.where(self._inpaint_mask <= (0.0 + eps), 1.0)
+        mask = torch.where(self._inpaint_mask >= t_prev + eps, 1.0, 0.0).to(
+            dtype=self._inpaint_mask.dtype, device=self._inpaint_mask.device
+        )
 
         return mask
 

invokeai/backend/stable_diffusion/diffusers_pipeline.py

@@ -371,7 +371,10 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
 
         if use_ip_adapter or use_regional_prompting:
             ip_adapters: Optional[List[UNetIPAdapterData]] = (
-                [{"ip_adapter": ipa.ip_adapter_model, "target_blocks": ipa.target_blocks} for ipa in ip_adapter_data]
+                [
+                    {"ip_adapter": ipa.ip_adapter_model, "target_blocks": ipa.target_blocks, "method": ipa.method}
+                    for ipa in ip_adapter_data
+                ]
                 if use_ip_adapter
                 else None
             )

invokeai/backend/stable_diffusion/diffusion/conditioning_data.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import math
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
@@ -104,15 +104,29 @@ class IPAdapterConditioningInfo:
 
 @dataclass
 class IPAdapterData:
+    """Data class for IP-Adapter configuration.
+
+    Attributes:
+        ip_adapter_model: The IP-Adapter model to use.
+        ip_adapter_conditioning: The IP-Adapter conditioning data.
+        mask: The mask to apply to the IP-Adapter conditioning.
+        target_blocks: List of target attention block names to apply IP-Adapter to.
+        negative_blocks: List of target attention block names that should use negative attention.
+        weight: The weight to apply to the IP-Adapter conditioning.
+        begin_step_percent: The percentage of steps at which to start applying the IP-Adapter.
+        end_step_percent: The percentage of steps at which to stop applying the IP-Adapter.
+        method: The method to use for applying the IP-Adapter ('full', 'style', 'composition').
+    """
+
     ip_adapter_model: IPAdapter
     ip_adapter_conditioning: IPAdapterConditioningInfo
     mask: torch.Tensor
     target_blocks: List[str]
-
-    # Either a single weight applied to all steps, or a list of weights for each step.
+    negative_blocks: List[str] = field(default_factory=list)
     weight: Union[float, List[float]] = 1.0
     begin_step_percent: float = 0.0
     end_step_percent: float = 1.0
+    method: str = "full"
 
     def scale_for_step(self, step_index: int, total_steps: int) -> float:
         first_adapter_step = math.floor(self.begin_step_percent * total_steps)

invokeai/backend/stable_diffusion/diffusion/custom_atttention.py

@@ -14,6 +14,7 @@ from invokeai.backend.stable_diffusion.diffusion.regional_prompt_data import Reg
 class IPAdapterAttentionWeights:
     ip_adapter_weights: IPAttentionProcessorWeights
     skip: bool
+    negative: bool
 
 
 class CustomAttnProcessor2_0(AttnProcessor2_0):
@@ -162,6 +163,10 @@ class CustomAttnProcessor2_0(AttnProcessor2_0):
                     # Expected ip_hidden_state shape: (batch_size, num_ip_images, ip_seq_len, ip_image_embedding)
 
                     if not self._ip_adapter_attention_weights[ipa_index].skip:
+                        # apply the IP-Adapter weights to the negative embeds
+                        if self._ip_adapter_attention_weights[ipa_index].negative:
+                            ip_hidden_states = torch.cat([ip_hidden_states[1], ip_hidden_states[0] * 0], dim=0)
+
                         ip_key = ipa_weights.to_k_ip(ip_hidden_states)
                         ip_value = ipa_weights.to_v_ip(ip_hidden_states)
 

invokeai/backend/stable_diffusion/diffusion/unet_attention_patcher.py

@@ -12,7 +12,8 @@ from invokeai.backend.stable_diffusion.diffusion.custom_atttention import (
 
 class UNetIPAdapterData(TypedDict):
     ip_adapter: IPAdapter
-    target_blocks: List[str]
+    target_blocks: List[str]  # Blocks where IP-Adapter should be applied
+    method: str  # Style or other method type
 
 
 class UNetAttentionPatcher:
@@ -39,12 +40,18 @@ class UNetAttentionPatcher:
                 for ip_adapter in self._ip_adapters:
                     ip_adapter_weights = ip_adapter["ip_adapter"].attn_weights.get_attention_processor_weights(idx)
                     skip = True
+                    negative = False
                     for block in ip_adapter["target_blocks"]:
                         if block in name:
                             skip = False
+                            negative = ip_adapter["method"] == "style_precise" and (
+                                block == "down_blocks.2.attentions.1"
+                                or block == "down_blocks.2"
+                                or block == "mid_block"
+                            )
                             break
                     ip_adapter_attention_weights: IPAdapterAttentionWeights = IPAdapterAttentionWeights(
-                        ip_adapter_weights=ip_adapter_weights, skip=skip
+                        ip_adapter_weights=ip_adapter_weights, skip=skip, negative=negative
                     )
                     ip_adapter_attention_weights_collection.append(ip_adapter_attention_weights)
 

invokeai/frontend/web/.eslintrc.js

@@ -9,13 +9,16 @@ module.exports = {
     // https://github.com/qdanik/eslint-plugin-path
     'path/no-relative-imports': ['error', { maxDepth: 0 }],
     // https://github.com/edvardchen/eslint-plugin-i18next/blob/HEAD/docs/rules/no-literal-string.md
-    'i18next/no-literal-string': 'error',
+    // TODO: ENABLE THIS RULE BEFORE v6.0.0
+    // 'i18next/no-literal-string': 'error',
     // https://eslint.org/docs/latest/rules/no-console
-    'no-console': 'error',
+    'no-console': 'warn',
     // https://eslint.org/docs/latest/rules/no-promise-executor-return
     'no-promise-executor-return': 'error',
     // https://eslint.org/docs/latest/rules/require-await
     'require-await': 'error',
+    // TODO: ENABLE THIS RULE BEFORE v6.0.0
+    'react/display-name': 'off',
     'no-restricted-properties': [
       'error',
       {
@@ -30,6 +33,27 @@ module.exports = {
           'The Clipboard API is not available by default in Firefox. Use the `useClipboard` hook instead, which wraps clipboard access to prevent errors.',
       },
     ],
+    'no-restricted-imports': [
+      'error',
+      {
+        paths: [
+          {
+            name: 'lodash-es',
+            importNames: ['isEqual'],
+            message: 'Please use objectEquals from @observ33r/object-equals instead.',
+          },
+          {
+            name: 'lodash-es',
+            message: 'Please use es-toolkit instead.',
+          },
+          {
+            name: 'es-toolkit',
+            importNames: ['isEqual'],
+            message: 'Please use objectEquals from @observ33r/object-equals instead.',
+          },
+        ],
+      },
+    ],
   },
   overrides: [
     /**

invokeai/frontend/web/knip.ts

@@ -3,6 +3,8 @@ import type { KnipConfig } from 'knip';
 const config: KnipConfig = {
   project: ['src/**/*.{ts,tsx}!'],
   ignore: [
+    // TODO(psyche): temporarily ignored all files for test build purposes
+    'src/**',
     // This file is only used during debugging
     'src/app/store/middleware/debugLoggerMiddleware.ts',
     // Autogenerated types - shouldn't ever touch these
@@ -14,6 +16,8 @@ const config: KnipConfig = {
     'src/features/controlLayers/konva/util.ts',
     // TODO(psyche): restore HRF functionality?
     'src/features/hrf/**',
+    // This feature is (temprarily?) disabled
+    'src/features/controlLayers/components/InpaintMask/InpaintMaskAddButtons.tsx',
   ],
   ignoreBinaries: ['only-allow'],
   paths: {