add support for multi-gpu rendering

This commit adds speculative support for parallel rendering across
multiple GPUs. The parallelism is at the level of a session. Each
session is given access to a different GPU. When all GPUs are busy,
execution of the session will block until a GPU becomes available.

The code is untested at the current time, and is being posted for
comment.

.github/workflows/lint-frontend.yml

@@ -36,10 +36,8 @@ jobs:
       - name: Typescript
         run: 'pnpm run lint:tsc'
       - name: Madge
-        run: 'pnpm run lint:dpdm'
+        run: 'pnpm run lint:madge'
       - name: ESLint
         run: 'pnpm run lint:eslint'
       - name: Prettier
         run: 'pnpm run lint:prettier'
-      - name: Knip
-        run: 'pnpm run lint:knip'

Makefile

@@ -6,44 +6,33 @@ default: help
 help:
 	@echo Developer commands:
 	@echo
-	@echo "ruff              Run ruff, fixing any safely-fixable errors and formatting"
-	@echo "ruff-unsafe       Run ruff, fixing all fixable errors and formatting"
-	@echo "mypy              Run mypy using the config in pyproject.toml to identify type mismatches and other coding errors"
-	@echo "mypy-all          Run mypy ignoring the config in pyproject.tom but still ignoring missing imports"
-	@echo "test"             Run the unit tests.
-	@echo "frontend-install" Install the pnpm modules needed for the front end
-	@echo "frontend-build    Build the frontend in order to run on localhost:9090"
-	@echo "frontend-dev      Run the frontend in developer mode on localhost:5173"
-	@echo "installer-zip     Build the installer .zip file for the current version"
-	@echo "tag-release       Tag the GitHub repository with the current version (use at release time only!)"
+	@echo "ruff           Run ruff, fixing any safely-fixable errors and formatting"
+	@echo "ruff-unsafe    Run ruff, fixing all fixable errors and formatting"
+	@echo "mypy           Run mypy using the config in pyproject.toml to identify type mismatches and other coding errors"
+	@echo "mypy-all       Run mypy ignoring the config in pyproject.tom but still ignoring missing imports"
+	@echo "frontend-build Build the frontend in order to run on localhost:9090"
+	@echo "frontend-dev   Run the frontend in developer mode on localhost:5173"
+	@echo "installer-zip  Build the installer .zip file for the current version"
+	@echo "tag-release    Tag the GitHub repository with the current version (use at release time only!)"
 
 # Runs ruff, fixing any safely-fixable errors and formatting
 ruff:
-	ruff check . --fix
-	ruff format .
+		ruff check . --fix
+		ruff format .
 
 # Runs ruff, fixing all errors it can fix and formatting
 ruff-unsafe:
-	ruff check . --fix --unsafe-fixes
-	ruff format .
+		ruff check . --fix --unsafe-fixes
+		ruff format .
 
 # Runs mypy, using the config in pyproject.toml
 mypy:
-	mypy scripts/invokeai-web.py
+		mypy scripts/invokeai-web.py
 
 # Runs mypy, ignoring the config in pyproject.toml but still ignoring missing (untyped) imports
 # (many files are ignored by the config, so this is useful for checking all files)
 mypy-all:
-	mypy scripts/invokeai-web.py --config-file= --ignore-missing-imports
-
-# Run the unit tests
-test:
-	pytest ./tests
-
-# Install the pnpm modules needed for the front end
-frontend-install:
-	rm -rf invokeai/frontend/web/node_modules
-	cd invokeai/frontend/web && pnpm install
+		mypy scripts/invokeai-web.py --config-file= --ignore-missing-imports
 
 # Build the frontend
 frontend-build:

invokeai/app/api/dependencies.py

@@ -4,6 +4,7 @@ from logging import Logger
 
 import torch
 
+from invokeai.app.services.item_storage.item_storage_memory import ItemStorageMemory
 from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
 from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
 from invokeai.app.services.shared.sqlite.sqlite_util import init_db
@@ -15,13 +16,14 @@ from ..services.board_image_records.board_image_records_sqlite import SqliteBoar
 from ..services.board_images.board_images_default import BoardImagesService
 from ..services.board_records.board_records_sqlite import SqliteBoardRecordStorage
 from ..services.boards.boards_default import BoardService
-from ..services.bulk_download.bulk_download_default import BulkDownloadService
 from ..services.config import InvokeAIAppConfig
 from ..services.download import DownloadQueueService
 from ..services.image_files.image_files_disk import DiskImageFileStorage
 from ..services.image_records.image_records_sqlite import SqliteImageRecordStorage
 from ..services.images.images_default import ImageService
 from ..services.invocation_cache.invocation_cache_memory import MemoryInvocationCache
+from ..services.invocation_processor.invocation_processor_default import DefaultInvocationProcessor
+from ..services.invocation_queue.invocation_queue_memory import MemoryInvocationQueue
 from ..services.invocation_services import InvocationServices
 from ..services.invocation_stats.invocation_stats_default import InvocationStatsService
 from ..services.invoker import Invoker
@@ -31,6 +33,7 @@ from ..services.model_records import ModelRecordServiceSQL
 from ..services.names.names_default import SimpleNameService
 from ..services.session_processor.session_processor_default import DefaultSessionProcessor
 from ..services.session_queue.session_queue_sqlite import SqliteSessionQueue
+from ..services.shared.graph import GraphExecutionState
 from ..services.urls.urls_default import LocalUrlService
 from ..services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage
 from .events import FastAPIEventService
@@ -82,7 +85,7 @@ class ApiDependencies:
         board_records = SqliteBoardRecordStorage(db=db)
         boards = BoardService()
         events = FastAPIEventService(event_handler_id)
-        bulk_download = BulkDownloadService()
+        graph_execution_manager = ItemStorageMemory[GraphExecutionState]()
         image_records = SqliteImageRecordStorage(db=db)
         images = ImageService()
         invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size)
@@ -102,6 +105,8 @@ class ApiDependencies:
         )
         names = SimpleNameService()
         performance_statistics = InvocationStatsService()
+        processor = DefaultInvocationProcessor()
+        queue = MemoryInvocationQueue()
         session_processor = DefaultSessionProcessor()
         session_queue = SqliteSessionQueue(db=db)
         urls = LocalUrlService()
@@ -112,9 +117,9 @@ class ApiDependencies:
             board_images=board_images,
             board_records=board_records,
             boards=boards,
-            bulk_download=bulk_download,
             configuration=configuration,
             events=events,
+            graph_execution_manager=graph_execution_manager,
             image_files=image_files,
             image_records=image_records,
             images=images,
@@ -124,6 +129,8 @@ class ApiDependencies:
             download_queue=download_queue_service,
             names=names,
             performance_statistics=performance_statistics,
+            processor=processor,
+            queue=queue,
             session_processor=session_processor,
             session_queue=session_queue,
             urls=urls,

invokeai/app/api/routers/images.py

@@ -2,7 +2,7 @@ import io
 import traceback
 from typing import Optional
 
-from fastapi import BackgroundTasks, Body, HTTPException, Path, Query, Request, Response, UploadFile
+from fastapi import Body, HTTPException, Path, Query, Request, Response, UploadFile
 from fastapi.responses import FileResponse
 from fastapi.routing import APIRouter
 from PIL import Image
@@ -375,67 +375,16 @@ async def unstar_images_in_list(
 
 class ImagesDownloaded(BaseModel):
     response: Optional[str] = Field(
-        default=None, description="The message to display to the user when images begin downloading"
-    )
-    bulk_download_item_name: Optional[str] = Field(
-        default=None, description="The name of the bulk download item for which events will be emitted"
+        description="If defined, the message to display to the user when images begin downloading"
     )
 
 
-@images_router.post(
-    "/download", operation_id="download_images_from_list", response_model=ImagesDownloaded, status_code=202
-)
+@images_router.post("/download", operation_id="download_images_from_list", response_model=ImagesDownloaded)
 async def download_images_from_list(
-    background_tasks: BackgroundTasks,
-    image_names: Optional[list[str]] = Body(
-        default=None, description="The list of names of images to download", embed=True
-    ),
+    image_names: list[str] = Body(description="The list of names of images to download", embed=True),
     board_id: Optional[str] = Body(
-        default=None, description="The board from which image should be downloaded", embed=True
+        default=None, description="The board from which image should be downloaded from", embed=True
     ),
 ) -> ImagesDownloaded:
-    if (image_names is None or len(image_names) == 0) and board_id is None:
-        raise HTTPException(status_code=400, detail="No images or board id specified.")
-    bulk_download_item_id: str = ApiDependencies.invoker.services.bulk_download.generate_item_id(board_id)
-
-    background_tasks.add_task(
-        ApiDependencies.invoker.services.bulk_download.handler,
-        image_names,
-        board_id,
-        bulk_download_item_id,
-    )
-    return ImagesDownloaded(bulk_download_item_name=bulk_download_item_id + ".zip")
-
-
-@images_router.api_route(
-    "/download/{bulk_download_item_name}",
-    methods=["GET"],
-    operation_id="get_bulk_download_item",
-    response_class=Response,
-    responses={
-        200: {
-            "description": "Return the complete bulk download item",
-            "content": {"application/zip": {}},
-        },
-        404: {"description": "Image not found"},
-    },
-)
-async def get_bulk_download_item(
-    background_tasks: BackgroundTasks,
-    bulk_download_item_name: str = Path(description="The bulk_download_item_name of the bulk download item to get"),
-) -> FileResponse:
-    """Gets a bulk download zip file"""
-    try:
-        path = ApiDependencies.invoker.services.bulk_download.get_path(bulk_download_item_name)
-
-        response = FileResponse(
-            path,
-            media_type="application/zip",
-            filename=bulk_download_item_name,
-            content_disposition_type="inline",
-        )
-        response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
-        background_tasks.add_task(ApiDependencies.invoker.services.bulk_download.delete, bulk_download_item_name)
-        return response
-    except Exception:
-        raise HTTPException(status_code=404)
+    # return ImagesDownloaded(response="Your images are downloading")
+    raise HTTPException(status_code=501, detail="Endpoint is not yet implemented")

invokeai/app/api/routers/model_manager.py

@@ -9,11 +9,11 @@ from typing import Any, Dict, List, Optional, Set
 
 from fastapi import Body, Path, Query, Response
 from fastapi.routing import APIRouter
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict
 from starlette.exceptions import HTTPException
 from typing_extensions import Annotated
 
-from invokeai.app.services.model_install import ModelInstallJob
+from invokeai.app.services.model_install import ModelInstallJob, ModelSource
 from invokeai.app.services.model_records import (
     DuplicateModelException,
     InvalidModelException,
@@ -32,7 +32,6 @@ from invokeai.backend.model_manager.config import (
 )
 from invokeai.backend.model_manager.merge import MergeInterpolationMethod, ModelMerger
 from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
-from invokeai.backend.model_manager.search import ModelSearch
 
 from ..dependencies import ApiDependencies
 
@@ -165,27 +164,6 @@ async def list_model_records(
     return ModelsList(models=found_models)
 
 
-@model_manager_router.get(
-    "/get_by_attrs",
-    operation_id="get_model_records_by_attrs",
-    response_model=AnyModelConfig,
-)
-async def get_model_records_by_attrs(
-    name: str = Query(description="The name of the model"),
-    type: ModelType = Query(description="The type of the model"),
-    base: BaseModelType = Query(description="The base model of the model"),
-) -> AnyModelConfig:
-    """Gets a model by its attributes. The main use of this route is to provide backwards compatibility with the old
-    model manager, which identified models by a combination of name, base and type."""
-    configs = ApiDependencies.invoker.services.model_manager.store.search_by_attr(
-        base_model=base, model_type=type, model_name=name
-    )
-    if not configs:
-        raise HTTPException(status_code=404, detail="No model found with these attributes")
-
-    return configs[0]
-
-
 @model_manager_router.get(
     "/i/{key}",
     operation_id="get_model_record",
@@ -223,7 +201,7 @@ async def list_model_summary(
 
 
 @model_manager_router.get(
-    "/i/{key}/metadata",
+    "/meta/i/{key}",
     operation_id="get_model_metadata",
     responses={
         200: {
@@ -231,6 +209,7 @@ async def list_model_summary(
             "content": {"application/json": {"example": example_model_metadata}},
         },
         400: {"description": "Bad request"},
+        404: {"description": "No metadata available"},
     },
 )
 async def get_model_metadata(
@@ -239,7 +218,8 @@ async def get_model_metadata(
     """Get a model metadata object."""
     record_store = ApiDependencies.invoker.services.model_manager.store
     result: Optional[AnyModelRepoMetadata] = record_store.get_metadata(key)
-
+    if not result:
+        raise HTTPException(status_code=404, detail="No metadata for a model with this key")
     return result
 
 
@@ -254,75 +234,6 @@ async def list_tags() -> Set[str]:
     return result
 
 
-class FoundModel(BaseModel):
-    path: str = Field(description="Path to the model")
-    is_installed: bool = Field(description="Whether or not the model is already installed")
-
-
-@model_manager_router.get(
-    "/scan_folder",
-    operation_id="scan_for_models",
-    responses={
-        200: {"description": "Directory scanned successfully"},
-        400: {"description": "Invalid directory path"},
-    },
-    status_code=200,
-    response_model=List[FoundModel],
-)
-async def scan_for_models(
-    scan_path: str = Query(description="Directory path to search for models", default=None),
-) -> List[FoundModel]:
-    path = pathlib.Path(scan_path)
-    if not scan_path or not path.is_dir():
-        raise HTTPException(
-            status_code=400,
-            detail=f"The search path '{scan_path}' does not exist or is not directory",
-        )
-
-    search = ModelSearch()
-    try:
-        found_model_paths = search.search(path)
-        models_path = ApiDependencies.invoker.services.configuration.models_path
-
-        # If the search path includes the main models directory, we need to exclude core models from the list.
-        # TODO(MM2): Core models should be handled by the model manager so we can determine if they are installed
-        # without needing to crawl the filesystem.
-        core_models_path = pathlib.Path(models_path, "core").resolve()
-        non_core_model_paths = [p for p in found_model_paths if not p.is_relative_to(core_models_path)]
-
-        installed_models = ApiDependencies.invoker.services.model_manager.store.search_by_attr()
-        resolved_installed_model_paths: list[str] = []
-        installed_model_sources: list[str] = []
-
-        # This call lists all installed models.
-        for model in installed_models:
-            path = pathlib.Path(model.path)
-            # If the model has a source, we need to add it to the list of installed sources.
-            if model.source:
-                installed_model_sources.append(model.source)
-            # If the path is not absolute, that means it is in the app models directory, and we need to join it with
-            # the models path before resolving.
-            if not path.is_absolute():
-                resolved_installed_model_paths.append(str(pathlib.Path(models_path, path).resolve()))
-                continue
-            resolved_installed_model_paths.append(str(path.resolve()))
-
-        scan_results: list[FoundModel] = []
-
-        # Check if the model is installed by comparing the resolved paths, appending to the scan result.
-        for p in non_core_model_paths:
-            path = str(p)
-            is_installed = path in resolved_installed_model_paths or path in installed_model_sources
-            found_model = FoundModel(path=path, is_installed=is_installed)
-            scan_results.append(found_model)
-    except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail=f"An error occurred while searching the directory: {e}",
-        )
-    return scan_results
-
-
 @model_manager_router.get(
     "/tags/search",
     operation_id="search_by_metadata_tags",
@@ -439,8 +350,8 @@ async def add_model_record(
 
 
 @model_manager_router.post(
-    "/install",
-    operation_id="install_model",
+    "/heuristic_import",
+    operation_id="heuristic_import_model",
     responses={
         201: {"description": "The model imported successfully"},
         415: {"description": "Unrecognized file/folder format"},
@@ -449,13 +360,12 @@ async def add_model_record(
     },
     status_code=201,
 )
-async def install_model(
-    source: str = Query(description="Model source to install, can be a local path, repo_id, or remote URL"),
-    # TODO(MM2): Can we type this?
+async def heuristic_import(
+    source: str,
     config: Optional[Dict[str, Any]] = Body(
         description="Dict of fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
         default=None,
-        example={"name": "string", "description": "string"},
+        example={"name": "modelT", "description": "antique cars"},
     ),
     access_token: Optional[str] = None,
 ) -> ModelInstallJob:
@@ -492,7 +402,106 @@ async def install_model(
         result: ModelInstallJob = installer.heuristic_import(
             source=source,
             config=config,
-            access_token=access_token,
+        )
+        logger.info(f"Started installation of {source}")
+    except UnknownModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=424, detail=str(e))
+    except InvalidModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=415)
+    except ValueError as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=409, detail=str(e))
+    return result
+
+
+@model_manager_router.post(
+    "/install",
+    operation_id="import_model",
+    responses={
+        201: {"description": "The model imported successfully"},
+        415: {"description": "Unrecognized file/folder format"},
+        424: {"description": "The model appeared to import successfully, but could not be found in the model manager"},
+        409: {"description": "There is already a model corresponding to this path or repo_id"},
+    },
+    status_code=201,
+)
+async def import_model(
+    source: ModelSource,
+    config: Optional[Dict[str, Any]] = Body(
+        description="Dict of fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
+        default=None,
+    ),
+) -> ModelInstallJob:
+    """Install a model using its local path, repo_id, or remote URL.
+
+    Models will be downloaded, probed, configured and installed in a
+    series of background threads. The return object has `status` attribute
+    that can be used to monitor progress.
+
+    The source object is a discriminated Union of LocalModelSource,
+    HFModelSource and URLModelSource. Set the "type" field to the
+    appropriate value:
+
+    * To install a local path using LocalModelSource, pass a source of form:
+      ```
+      {
+        "type": "local",
+        "path": "/path/to/model",
+        "inplace": false
+      }
+      ```
+      The "inplace" flag, if true, will register the model in place in its
+      current filesystem location. Otherwise, the model will be copied
+      into the InvokeAI models directory.
+
+    * To install a HuggingFace repo_id using HFModelSource, pass a source of form:
+      ```
+      {
+        "type": "hf",
+        "repo_id": "stabilityai/stable-diffusion-2.0",
+        "variant": "fp16",
+        "subfolder": "vae",
+        "access_token": "f5820a918aaf01"
+      }
+      ```
+      The `variant`, `subfolder` and `access_token` fields are optional.
+
+    * To install a remote model using an arbitrary URL, pass:
+      ```
+      {
+        "type": "url",
+        "url": "http://www.civitai.com/models/123456",
+        "access_token": "f5820a918aaf01"
+      }
+      ```
+      The `access_token` field is optonal
+
+    The model's configuration record will be probed and filled in
+    automatically.  To override the default guesses, pass "metadata"
+    with a Dict containing the attributes you wish to override.
+
+    Installation occurs in the background. Either use list_model_install_jobs()
+    to poll for completion, or listen on the event bus for the following events:
+
+      * "model_install_running"
+      * "model_install_completed"
+      * "model_install_error"
+
+    On successful completion, the event's payload will contain the field "key"
+    containing the installed ID of the model. On an error, the event's payload
+    will contain the fields "error_type" and "error" describing the nature of the
+    error and its traceback, respectively.
+
+    """
+    logger = ApiDependencies.invoker.services.logger
+
+    try:
+        installer = ApiDependencies.invoker.services.model_manager.install
+        result: ModelInstallJob = installer.import_model(
+            source=source,
+            config=config,
         )
         logger.info(f"Started installation of {source}")
     except UnknownModelException as e:
@@ -628,7 +637,6 @@ async def convert_model(
     Note that during the conversion process the key and model hash will change.
     The return value is the model configuration for the converted model.
     """
-    model_manager = ApiDependencies.invoker.services.model_manager
     logger = ApiDependencies.invoker.services.logger
     loader = ApiDependencies.invoker.services.model_manager.load
     store = ApiDependencies.invoker.services.model_manager.store
@@ -645,7 +653,7 @@ async def convert_model(
         raise HTTPException(400, f"The model with key {key} is not a main checkpoint model.")
 
     # loading the model will convert it into a cached diffusers file
-    model_manager.load_model_by_config(model_config, submodel_type=SubModelType.Scheduler)
+    loader.load_model_by_config(model_config, submodel_type=SubModelType.Scheduler)
 
     # Get the path of the converted model from the loader
     cache_path = loader.convert_cache.cache_path(key)

invokeai/app/api/routers/sessions.py

@@ -0,0 +1,276 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+
+from fastapi import HTTPException, Path
+from fastapi.routing import APIRouter
+
+from ...services.shared.graph import GraphExecutionState
+from ..dependencies import ApiDependencies
+
+session_router = APIRouter(prefix="/v1/sessions", tags=["sessions"])
+
+
+# @session_router.post(
+#     "/",
+#     operation_id="create_session",
+#     responses={
+#         200: {"model": GraphExecutionState},
+#         400: {"description": "Invalid json"},
+#     },
+#     deprecated=True,
+# )
+# async def create_session(
+#     queue_id: str = Query(default="", description="The id of the queue to associate the session with"),
+#     graph: Optional[Graph] = Body(default=None, description="The graph to initialize the session with"),
+# ) -> GraphExecutionState:
+#     """Creates a new session, optionally initializing it with an invocation graph"""
+#     session = ApiDependencies.invoker.create_execution_state(queue_id=queue_id, graph=graph)
+#     return session
+
+
+# @session_router.get(
+#     "/",
+#     operation_id="list_sessions",
+#     responses={200: {"model": PaginatedResults[GraphExecutionState]}},
+#     deprecated=True,
+# )
+# async def list_sessions(
+#     page: int = Query(default=0, description="The page of results to get"),
+#     per_page: int = Query(default=10, description="The number of results per page"),
+#     query: str = Query(default="", description="The query string to search for"),
+# ) -> PaginatedResults[GraphExecutionState]:
+#     """Gets a list of sessions, optionally searching"""
+#     if query == "":
+#         result = ApiDependencies.invoker.services.graph_execution_manager.list(page, per_page)
+#     else:
+#         result = ApiDependencies.invoker.services.graph_execution_manager.search(query, page, per_page)
+#     return result
+
+
+@session_router.get(
+    "/{session_id}",
+    operation_id="get_session",
+    responses={
+        200: {"model": GraphExecutionState},
+        404: {"description": "Session not found"},
+    },
+)
+async def get_session(
+    session_id: str = Path(description="The id of the session to get"),
+) -> GraphExecutionState:
+    """Gets a session"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        raise HTTPException(status_code=404)
+    else:
+        return session
+
+
+# @session_router.post(
+#     "/{session_id}/nodes",
+#     operation_id="add_node",
+#     responses={
+#         200: {"model": str},
+#         400: {"description": "Invalid node or link"},
+#         404: {"description": "Session not found"},
+#     },
+#     deprecated=True,
+# )
+# async def add_node(
+#     session_id: str = Path(description="The id of the session"),
+#     node: Annotated[Union[BaseInvocation.get_invocations()], Field(discriminator="type")] = Body(  # type: ignore
+#         description="The node to add"
+#     ),
+# ) -> str:
+#     """Adds a node to the graph"""
+#     session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+#     if session is None:
+#         raise HTTPException(status_code=404)
+
+#     try:
+#         session.add_node(node)
+#         ApiDependencies.invoker.services.graph_execution_manager.set(
+#             session
+#         )  # TODO: can this be done automatically, or add node through an API?
+#         return session.id
+#     except NodeAlreadyExecutedError:
+#         raise HTTPException(status_code=400)
+#     except IndexError:
+#         raise HTTPException(status_code=400)
+
+
+# @session_router.put(
+#     "/{session_id}/nodes/{node_path}",
+#     operation_id="update_node",
+#     responses={
+#         200: {"model": GraphExecutionState},
+#         400: {"description": "Invalid node or link"},
+#         404: {"description": "Session not found"},
+#     },
+#     deprecated=True,
+# )
+# async def update_node(
+#     session_id: str = Path(description="The id of the session"),
+#     node_path: str = Path(description="The path to the node in the graph"),
+#     node: Annotated[Union[BaseInvocation.get_invocations()], Field(discriminator="type")] = Body(  # type: ignore
+#         description="The new node"
+#     ),
+# ) -> GraphExecutionState:
+#     """Updates a node in the graph and removes all linked edges"""
+#     session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+#     if session is None:
+#         raise HTTPException(status_code=404)
+
+#     try:
+#         session.update_node(node_path, node)
+#         ApiDependencies.invoker.services.graph_execution_manager.set(
+#             session
+#         )  # TODO: can this be done automatically, or add node through an API?
+#         return session
+#     except NodeAlreadyExecutedError:
+#         raise HTTPException(status_code=400)
+#     except IndexError:
+#         raise HTTPException(status_code=400)
+
+
+# @session_router.delete(
+#     "/{session_id}/nodes/{node_path}",
+#     operation_id="delete_node",
+#     responses={
+#         200: {"model": GraphExecutionState},
+#         400: {"description": "Invalid node or link"},
+#         404: {"description": "Session not found"},
+#     },
+#     deprecated=True,
+# )
+# async def delete_node(
+#     session_id: str = Path(description="The id of the session"),
+#     node_path: str = Path(description="The path to the node to delete"),
+# ) -> GraphExecutionState:
+#     """Deletes a node in the graph and removes all linked edges"""
+#     session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+#     if session is None:
+#         raise HTTPException(status_code=404)
+
+#     try:
+#         session.delete_node(node_path)
+#         ApiDependencies.invoker.services.graph_execution_manager.set(
+#             session
+#         )  # TODO: can this be done automatically, or add node through an API?
+#         return session
+#     except NodeAlreadyExecutedError:
+#         raise HTTPException(status_code=400)
+#     except IndexError:
+#         raise HTTPException(status_code=400)
+
+
+# @session_router.post(
+#     "/{session_id}/edges",
+#     operation_id="add_edge",
+#     responses={
+#         200: {"model": GraphExecutionState},
+#         400: {"description": "Invalid node or link"},
+#         404: {"description": "Session not found"},
+#     },
+#     deprecated=True,
+# )
+# async def add_edge(
+#     session_id: str = Path(description="The id of the session"),
+#     edge: Edge = Body(description="The edge to add"),
+# ) -> GraphExecutionState:
+#     """Adds an edge to the graph"""
+#     session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+#     if session is None:
+#         raise HTTPException(status_code=404)
+
+#     try:
+#         session.add_edge(edge)
+#         ApiDependencies.invoker.services.graph_execution_manager.set(
+#             session
+#         )  # TODO: can this be done automatically, or add node through an API?
+#         return session
+#     except NodeAlreadyExecutedError:
+#         raise HTTPException(status_code=400)
+#     except IndexError:
+#         raise HTTPException(status_code=400)
+
+
+# # TODO: the edge being in the path here is really ugly, find a better solution
+# @session_router.delete(
+#     "/{session_id}/edges/{from_node_id}/{from_field}/{to_node_id}/{to_field}",
+#     operation_id="delete_edge",
+#     responses={
+#         200: {"model": GraphExecutionState},
+#         400: {"description": "Invalid node or link"},
+#         404: {"description": "Session not found"},
+#     },
+#     deprecated=True,
+# )
+# async def delete_edge(
+#     session_id: str = Path(description="The id of the session"),
+#     from_node_id: str = Path(description="The id of the node the edge is coming from"),
+#     from_field: str = Path(description="The field of the node the edge is coming from"),
+#     to_node_id: str = Path(description="The id of the node the edge is going to"),
+#     to_field: str = Path(description="The field of the node the edge is going to"),
+# ) -> GraphExecutionState:
+#     """Deletes an edge from the graph"""
+#     session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+#     if session is None:
+#         raise HTTPException(status_code=404)
+
+#     try:
+#         edge = Edge(
+#             source=EdgeConnection(node_id=from_node_id, field=from_field),
+#             destination=EdgeConnection(node_id=to_node_id, field=to_field),
+#         )
+#         session.delete_edge(edge)
+#         ApiDependencies.invoker.services.graph_execution_manager.set(
+#             session
+#         )  # TODO: can this be done automatically, or add node through an API?
+#         return session
+#     except NodeAlreadyExecutedError:
+#         raise HTTPException(status_code=400)
+#     except IndexError:
+#         raise HTTPException(status_code=400)
+
+
+# @session_router.put(
+#     "/{session_id}/invoke",
+#     operation_id="invoke_session",
+#     responses={
+#         200: {"model": None},
+#         202: {"description": "The invocation is queued"},
+#         400: {"description": "The session has no invocations ready to invoke"},
+#         404: {"description": "Session not found"},
+#     },
+#     deprecated=True,
+# )
+# async def invoke_session(
+#     queue_id: str = Query(description="The id of the queue to associate the session with"),
+#     session_id: str = Path(description="The id of the session to invoke"),
+#     all: bool = Query(default=False, description="Whether or not to invoke all remaining invocations"),
+# ) -> Response:
+#     """Invokes a session"""
+#     session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+#     if session is None:
+#         raise HTTPException(status_code=404)
+
+#     if session.is_complete():
+#         raise HTTPException(status_code=400)
+
+#     ApiDependencies.invoker.invoke(queue_id, session, invoke_all=all)
+#     return Response(status_code=202)
+
+
+# @session_router.delete(
+#     "/{session_id}/invoke",
+#     operation_id="cancel_session_invoke",
+#     responses={202: {"description": "The invocation is canceled"}},
+#     deprecated=True,
+# )
+# async def cancel_session_invoke(
+#     session_id: str = Path(description="The id of the session to cancel"),
+# ) -> Response:
+#     """Invokes a session"""
+#     ApiDependencies.invoker.cancel(session_id)
+#     return Response(status_code=202)

invokeai/app/api/sockets.py

@@ -12,26 +12,16 @@ class SocketIO:
     __sio: AsyncServer
     __app: ASGIApp
 
-    __sub_queue: str = "subscribe_queue"
-    __unsub_queue: str = "unsubscribe_queue"
-
-    __sub_bulk_download: str = "subscribe_bulk_download"
-    __unsub_bulk_download: str = "unsubscribe_bulk_download"
-
     def __init__(self, app: FastAPI):
         self.__sio = AsyncServer(async_mode="asgi", cors_allowed_origins="*")
         self.__app = ASGIApp(socketio_server=self.__sio, socketio_path="/ws/socket.io")
         app.mount("/ws", self.__app)
 
-        self.__sio.on(self.__sub_queue, handler=self._handle_sub_queue)
-        self.__sio.on(self.__unsub_queue, handler=self._handle_unsub_queue)
+        self.__sio.on("subscribe_queue", handler=self._handle_sub_queue)
+        self.__sio.on("unsubscribe_queue", handler=self._handle_unsub_queue)
         local_handler.register(event_name=EventServiceBase.queue_event, _func=self._handle_queue_event)
         local_handler.register(event_name=EventServiceBase.model_event, _func=self._handle_model_event)
 
-        self.__sio.on(self.__sub_bulk_download, handler=self._handle_sub_bulk_download)
-        self.__sio.on(self.__unsub_bulk_download, handler=self._handle_unsub_bulk_download)
-        local_handler.register(event_name=EventServiceBase.bulk_download_event, _func=self._handle_bulk_download_event)
-
     async def _handle_queue_event(self, event: Event):
         await self.__sio.emit(
             event=event[1]["event"],
@@ -49,18 +39,3 @@ class SocketIO:
 
     async def _handle_model_event(self, event: Event) -> None:
         await self.__sio.emit(event=event[1]["event"], data=event[1]["data"])
-
-    async def _handle_bulk_download_event(self, event: Event):
-        await self.__sio.emit(
-            event=event[1]["event"],
-            data=event[1]["data"],
-            room=event[1]["data"]["bulk_download_id"],
-        )
-
-    async def _handle_sub_bulk_download(self, sid, data, *args, **kwargs):
-        if "bulk_download_id" in data:
-            await self.__sio.enter_room(sid, data["bulk_download_id"])
-
-    async def _handle_unsub_bulk_download(self, sid, data, *args, **kwargs):
-        if "bulk_download_id" in data:
-            await self.__sio.leave_room(sid, data["bulk_download_id"])

invokeai/app/api_app.py

@@ -50,6 +50,7 @@ if True:  # hack to make flake8 happy with imports coming after setting up the c
         images,
         model_manager,
         session_queue,
+        sessions,
         utilities,
         workflows,
     )
@@ -109,6 +110,8 @@ async def shutdown_event() -> None:
 
 
 # Include all routers
+app.include_router(sessions.session_router, prefix="/api")
+
 app.include_router(utilities.utilities_router, prefix="/api")
 app.include_router(model_manager.model_manager_router, prefix="/api")
 app.include_router(download_queue.download_queue_router, prefix="/api")
@@ -148,8 +151,6 @@ def custom_openapi() -> dict[str, Any]:
         # TODO: note that we assume the schema_key here is the TYPE.__name__
         # This could break in some cases, figure out a better way to do it
         output_type_titles[schema_key] = output_schema["title"]
-        openapi_schema["components"]["schemas"][schema_key] = output_schema
-        openapi_schema["components"]["schemas"][schema_key]["class"] = "output"
 
     # Add Node Editor UI helper schemas
     ui_config_schemas = models_json_schema(
@@ -172,6 +173,7 @@ def custom_openapi() -> dict[str, Any]:
         outputs_ref = {"$ref": f"#/components/schemas/{output_type_title}"}
         invoker_schema["output"] = outputs_ref
         invoker_schema["class"] = "invocation"
+        openapi_schema["components"]["schemas"][f"{output_type_title}"]["class"] = "output"
 
     # This code no longer seems to be necessary?
     # Leave it here just in case

invokeai/app/invocations/baseinvocation.py

@@ -8,26 +8,13 @@ import warnings
 from abc import ABC, abstractmethod
 from enum import Enum
 from inspect import signature
-from typing import (
-    TYPE_CHECKING,
-    Annotated,
-    Any,
-    Callable,
-    ClassVar,
-    Iterable,
-    Literal,
-    Optional,
-    Type,
-    TypeVar,
-    Union,
-    cast,
-)
+from types import UnionType
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Literal, Optional, Type, TypeVar, Union, cast
 
 import semver
-from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, create_model
+from pydantic import BaseModel, ConfigDict, Field, create_model
 from pydantic.fields import FieldInfo
 from pydantic_core import PydanticUndefined
-from typing_extensions import TypeAliasType
 
 from invokeai.app.invocations.fields import (
     FieldKind,
@@ -97,7 +84,6 @@ class BaseInvocationOutput(BaseModel):
     """
 
     _output_classes: ClassVar[set[BaseInvocationOutput]] = set()
-    _typeadapter: ClassVar[Optional[TypeAdapter[Any]]] = None
 
     @classmethod
     def register_output(cls, output: BaseInvocationOutput) -> None:
@@ -110,14 +96,10 @@ class BaseInvocationOutput(BaseModel):
         return cls._output_classes
 
     @classmethod
-    def get_typeadapter(cls) -> TypeAdapter[Any]:
-        """Gets a pydantc TypeAdapter for the union of all invocation output types."""
-        if not cls._typeadapter:
-            InvocationOutputsUnion = TypeAliasType(
-                "InvocationOutputsUnion", Annotated[Union[tuple(cls._output_classes)], Field(discriminator="type")]
-            )
-            cls._typeadapter = TypeAdapter(InvocationOutputsUnion)
-        return cls._typeadapter
+    def get_outputs_union(cls) -> UnionType:
+        """Gets a union of all invocation outputs."""
+        outputs_union = Union[tuple(cls._output_classes)]  # type: ignore [valid-type]
+        return outputs_union  # type: ignore [return-value]
 
     @classmethod
     def get_output_types(cls) -> Iterable[str]:
@@ -166,7 +148,6 @@ class BaseInvocation(ABC, BaseModel):
     """
 
     _invocation_classes: ClassVar[set[BaseInvocation]] = set()
-    _typeadapter: ClassVar[Optional[TypeAdapter[Any]]] = None
 
     @classmethod
     def get_type(cls) -> str:
@@ -179,14 +160,10 @@ class BaseInvocation(ABC, BaseModel):
         cls._invocation_classes.add(invocation)
 
     @classmethod
-    def get_typeadapter(cls) -> TypeAdapter[Any]:
-        """Gets a pydantc TypeAdapter for the union of all invocation types."""
-        if not cls._typeadapter:
-            InvocationsUnion = TypeAliasType(
-                "InvocationsUnion", Annotated[Union[tuple(cls._invocation_classes)], Field(discriminator="type")]
-            )
-            cls._typeadapter = TypeAdapter(InvocationsUnion)
-        return cls._typeadapter
+    def get_invocations_union(cls) -> UnionType:
+        """Gets a union of all invocation types."""
+        invocations_union = Union[tuple(cls._invocation_classes)]  # type: ignore [valid-type]
+        return invocations_union  # type: ignore [return-value]
 
     @classmethod
     def get_invocations(cls) -> Iterable[BaseInvocation]:

invokeai/app/invocations/compel.py

@@ -3,8 +3,9 @@ from typing import Iterator, List, Optional, Tuple, Union
 import torch
 from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
-from transformers import CLIPTextModel, CLIPTokenizer
+from transformers import CLIPTokenizer
 
+import invokeai.backend.util.logging as logger
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     Input,
@@ -13,9 +14,11 @@ from invokeai.app.invocations.fields import (
     UIComponent,
 )
 from invokeai.app.invocations.primitives import ConditioningOutput
+from invokeai.app.services.model_records import UnknownModelException
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.ti_utils import generate_ti_list
+from invokeai.app.util.ti_utils import extract_ti_triggers_from_prompt
 from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.model_manager import ModelType
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
@@ -23,6 +26,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     ExtraConditioningInfo,
     SDXLConditioningInfo,
 )
+from invokeai.backend.textual_inversion import TextualInversionModelRaw
 from invokeai.backend.util.devices import torch_dtype
 
 from .baseinvocation import (
@@ -66,11 +70,7 @@ class CompelInvocation(BaseInvocation):
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
         tokenizer_info = context.models.load(**self.clip.tokenizer.model_dump())
-        tokenizer_model = tokenizer_info.model
-        assert isinstance(tokenizer_model, CLIPTokenizer)
         text_encoder_info = context.models.load(**self.clip.text_encoder.model_dump())
-        text_encoder_model = text_encoder_info.model
-        assert isinstance(text_encoder_model, CLIPTextModel)
 
         def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.clip.loras:
@@ -82,10 +82,21 @@ class CompelInvocation(BaseInvocation):
 
         # loras = [(context.models.get(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras]
 
-        ti_list = generate_ti_list(self.prompt, text_encoder_info.config.base, context)
+        ti_list = []
+        for trigger in extract_ti_triggers_from_prompt(self.prompt):
+            name = trigger[1:-1]
+            try:
+                loaded_model = context.models.load(**self.clip.text_encoder.model_dump()).model
+                assert isinstance(loaded_model, TextualInversionModelRaw)
+                ti_list.append((name, loaded_model))
+            except UnknownModelException:
+                # print(e)
+                # import traceback
+                # print(traceback.format_exc())
+                print(f'Warn: trigger: "{trigger}" not found')
 
         with (
-            ModelPatcher.apply_ti(tokenizer_model, text_encoder_model, ti_list) as (
+            ModelPatcher.apply_ti(tokenizer_info.model, text_encoder_info.model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -93,9 +104,8 @@ class CompelInvocation(BaseInvocation):
             # Apply the LoRA after text_encoder has been moved to its target device for faster patching.
             ModelPatcher.apply_lora_text_encoder(text_encoder, _lora_loader()),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
-            ModelPatcher.apply_clip_skip(text_encoder_model, self.clip.skipped_layers),
+            ModelPatcher.apply_clip_skip(text_encoder_info.model, self.clip.skipped_layers),
         ):
-            assert isinstance(text_encoder, CLIPTextModel)
             compel = Compel(
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
@@ -145,11 +155,7 @@ class SDXLPromptInvocationBase:
         zero_on_empty: bool,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[ExtraConditioningInfo]]:
         tokenizer_info = context.models.load(**clip_field.tokenizer.model_dump())
-        tokenizer_model = tokenizer_info.model
-        assert isinstance(tokenizer_model, CLIPTokenizer)
         text_encoder_info = context.models.load(**clip_field.text_encoder.model_dump())
-        text_encoder_model = text_encoder_info.model
-        assert isinstance(text_encoder_model, CLIPTextModel)
 
         # return zero on empty
         if prompt == "" and zero_on_empty:
@@ -183,10 +189,25 @@ class SDXLPromptInvocationBase:
 
         # loras = [(context.models.get(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras]
 
-        ti_list = generate_ti_list(prompt, text_encoder_info.config.base, context)
+        ti_list = []
+        for trigger in extract_ti_triggers_from_prompt(prompt):
+            name = trigger[1:-1]
+            try:
+                ti_model = context.models.load_by_attrs(
+                    model_name=name, base_model=text_encoder_info.config.base, model_type=ModelType.TextualInversion
+                ).model
+                assert isinstance(ti_model, TextualInversionModelRaw)
+                ti_list.append((name, ti_model))
+            except UnknownModelException:
+                # print(e)
+                # import traceback
+                # print(traceback.format_exc())
+                logger.warning(f'trigger: "{trigger}" not found')
+            except ValueError:
+                logger.warning(f'trigger: "{trigger}" more than one similarly-named textual inversion models')
 
         with (
-            ModelPatcher.apply_ti(tokenizer_model, text_encoder_model, ti_list) as (
+            ModelPatcher.apply_ti(tokenizer_info.model, text_encoder_info.model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -194,9 +215,8 @@ class SDXLPromptInvocationBase:
             # Apply the LoRA after text_encoder has been moved to its target device for faster patching.
             ModelPatcher.apply_lora(text_encoder, _lora_loader(), lora_prefix),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
-            ModelPatcher.apply_clip_skip(text_encoder_model, clip_field.skipped_layers),
+            ModelPatcher.apply_clip_skip(text_encoder_info.model, clip_field.skipped_layers),
         ):
-            assert isinstance(text_encoder, CLIPTextModel)
             compel = Compel(
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
@@ -397,7 +417,7 @@ class ClipSkipInvocation(BaseInvocation):
     """Skip layers in clip text_encoder model."""
 
     clip: ClipField = InputField(description=FieldDescriptions.clip, input=Input.Connection, title="CLIP")
-    skipped_layers: int = InputField(default=0, ge=0, description=FieldDescriptions.skipped_layers)
+    skipped_layers: int = InputField(default=0, description=FieldDescriptions.skipped_layers)
 
     def invoke(self, context: InvocationContext) -> ClipSkipInvocationOutput:
         self.clip.skipped_layers += self.skipped_layers

invokeai/app/invocations/fields.py

@@ -199,7 +199,6 @@ class DenoiseMaskField(BaseModel):
 
     mask_name: str = Field(description="The name of the mask image")
     masked_latents_name: Optional[str] = Field(default=None, description="The name of the masked image latents")
-    gradient: bool = Field(default=False, description="Used for gradient inpainting")
 
 
 class LatentsField(BaseModel):

invokeai/app/invocations/image.py

@@ -22,7 +22,11 @@ from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark
 from invokeai.backend.image_util.safety_checker import SafetyChecker
 
-from .baseinvocation import BaseInvocation, Classification, invocation
+from .baseinvocation import (
+    BaseInvocation,
+    Classification,
+    invocation,
+)
 
 
 @invocation("show_image", title="Show Image", tags=["image"], category="image", version="1.0.1")
@@ -930,40 +934,3 @@ class SaveImageInvocation(BaseInvocation, WithMetadata, WithBoard):
         image_dto = context.images.save(image=image)
 
         return ImageOutput.build(image_dto)
-
-
-@invocation(
-    "canvas_paste_back",
-    title="Canvas Paste Back",
-    tags=["image", "combine"],
-    category="image",
-    version="1.0.0",
-)
-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")
-    mask: ImageField = InputField(
-        description="The mask to use when pasting",
-    )
-    mask_blur: int = InputField(default=0, ge=0, description="The amount to blur the mask by")
-
-    def _prepare_mask(self, mask: Image.Image) -> Image.Image:
-        mask_array = numpy.array(mask)
-        kernel = numpy.ones((self.mask_blur, self.mask_blur), numpy.uint8)
-        dilated_mask_array = cv2.erode(mask_array, kernel, iterations=3)
-        dilated_mask = Image.fromarray(dilated_mask_array)
-        if self.mask_blur > 0:
-            mask = dilated_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))
-        return ImageOps.invert(mask.convert("L"))
-
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        source_image = context.images.get_pil(self.source_image.image_name)
-        target_image = context.images.get_pil(self.target_image.image_name)
-        mask = self._prepare_mask(context.images.get_pil(self.mask.image_name))
-
-        source_image.paste(target_image, (0, 0), mask)
-
-        image_dto = context.images.save(image=source_image)
-        return ImageOutput.build(image_dto)

invokeai/app/invocations/latent.py

@@ -23,7 +23,7 @@ from diffusers.models.attention_processor import (
 from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.schedulers import DPMSolverSDEScheduler
 from diffusers.schedulers import SchedulerMixin as Scheduler
-from PIL import Image, ImageFilter
+from PIL import Image
 from pydantic import field_validator
 from torchvision.transforms.functional import resize as tv_resize
 
@@ -128,7 +128,7 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
         ui_order=4,
     )
 
-    def prep_mask_tensor(self, mask_image: Image.Image) -> torch.Tensor:
+    def prep_mask_tensor(self, mask_image: Image) -> torch.Tensor:
         if mask_image.mode != "L":
             mask_image = mask_image.convert("L")
         mask_tensor: torch.Tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
@@ -169,62 +169,6 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
         return DenoiseMaskOutput.build(
             mask_name=mask_name,
             masked_latents_name=masked_latents_name,
-            gradient=False,
-        )
-
-
-@invocation(
-    "create_gradient_mask",
-    title="Create Gradient Mask",
-    tags=["mask", "denoise"],
-    category="latents",
-    version="1.0.0",
-)
-class CreateGradientMaskInvocation(BaseInvocation):
-    """Creates mask for denoising model run."""
-
-    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
-    )
-    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
-    )
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> DenoiseMaskOutput:
-        mask_image = context.images.get_pil(self.mask.image_name, mode="L")
-        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))
-
-        mask_tensor: torch.Tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
-        blur_tensor: torch.Tensor = image_resized_to_grid_as_tensor(blur_mask, normalize=False)
-
-        # redistribute blur so that the edges are 0 and blur out to 1
-        blur_tensor = (blur_tensor - 0.5) * 2
-
-        threshold = 1 - self.minimum_denoise
-
-        if self.coherence_mode == "Staged":
-            # wherever the blur_tensor is masked to any degree, convert it to threshold
-            blur_tensor = torch.where((blur_tensor < 1), 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)
-
-        # multiply original mask to force actually masked regions to 0
-        blur_tensor = mask_tensor * blur_tensor
-
-        mask_name = context.tensors.save(tensor=blur_tensor.unsqueeze(1))
-
-        return DenoiseMaskOutput.build(
-            mask_name=mask_name,
-            masked_latents_name=None,
-            gradient=True,
         )
 
 
@@ -662,9 +606,9 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
     def prep_inpaint_mask(
         self, context: InvocationContext, latents: torch.Tensor
-    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], bool]:
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
         if self.denoise_mask is None:
-            return None, None, False
+            return None, None
 
         mask = context.tensors.load(self.denoise_mask.mask_name)
         mask = tv_resize(mask, latents.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
@@ -673,7 +617,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         else:
             masked_latents = None
 
-        return 1 - mask, masked_latents, self.denoise_mask.gradient
+        return 1 - mask, masked_latents
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
@@ -700,7 +644,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
             if seed is None:
                 seed = 0
 
-            mask, masked_latents, gradient_mask = self.prep_inpaint_mask(context, latents)
+            mask, masked_latents = self.prep_inpaint_mask(context, latents)
 
             # TODO(ryand): I have hard-coded `do_classifier_free_guidance=True` to mirror the behaviour of ControlNets,
             # below. Investigate whether this is appropriate.
@@ -788,7 +732,6 @@ class DenoiseLatentsInvocation(BaseInvocation):
                     seed=seed,
                     mask=mask,
                     masked_latents=masked_latents,
-                    gradient_mask=gradient_mask,
                     num_inference_steps=num_inference_steps,
                     conditioning_data=conditioning_data,
                     control_data=controlnet_data,

invokeai/app/invocations/metadata.py

@@ -33,7 +33,7 @@ class MetadataItemField(BaseModel):
 class LoRAMetadataField(BaseModel):
     """LoRA Metadata Field"""
 
-    model: LoRAModelField = Field(description=FieldDescriptions.lora_model)
+    lora: LoRAModelField = Field(description=FieldDescriptions.lora_model)
     weight: float = Field(description=FieldDescriptions.lora_weight)
 
 
@@ -114,7 +114,7 @@ GENERATION_MODES = Literal[
 ]
 
 
-@invocation("core_metadata", title="Core Metadata", tags=["metadata"], category="metadata", version="1.1.1")
+@invocation("core_metadata", title="Core Metadata", tags=["metadata"], category="metadata", version="1.0.1")
 class CoreMetadataInvocation(BaseInvocation):
     """Collects core generation metadata into a MetadataField"""
 

invokeai/app/invocations/primitives.py

@@ -299,13 +299,9 @@ class DenoiseMaskOutput(BaseInvocationOutput):
     denoise_mask: DenoiseMaskField = OutputField(description="Mask for denoise model run")
 
     @classmethod
-    def build(
-        cls, mask_name: str, masked_latents_name: Optional[str] = None, gradient: bool = False
-    ) -> "DenoiseMaskOutput":
+    def build(cls, mask_name: str, masked_latents_name: Optional[str] = None) -> "DenoiseMaskOutput":
         return cls(
-            denoise_mask=DenoiseMaskField(
-                mask_name=mask_name, masked_latents_name=masked_latents_name, gradient=gradient
-            ),
+            denoise_mask=DenoiseMaskField(mask_name=mask_name, masked_latents_name=masked_latents_name),
         )
 
 

invokeai/app/services/bulk_download/bulk_download_base.py

@@ -1,44 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Optional
-
-
-class BulkDownloadBase(ABC):
-    """Responsible for creating a zip file containing the images specified by the given image names or board id."""
-
-    @abstractmethod
-    def handler(
-        self, image_names: Optional[list[str]], board_id: Optional[str], bulk_download_item_id: Optional[str]
-    ) -> None:
-        """
-        Create a zip file containing the images specified by the given image names or board id.
-
-        :param image_names: A list of image names to include in the zip file.
-        :param board_id: The ID of the board. If provided, all images associated with the board will be included in the zip file.
-        :param bulk_download_item_id: The bulk_download_item_id that will be used to retrieve the bulk download item when it is prepared, if none is provided a uuid will be generated.
-        """
-
-    @abstractmethod
-    def get_path(self, bulk_download_item_name: str) -> str:
-        """
-        Get the path to the bulk download file.
-
-        :param bulk_download_item_name: The name of the bulk download item.
-        :return: The path to the bulk download file.
-        """
-
-    @abstractmethod
-    def generate_item_id(self, board_id: Optional[str]) -> str:
-        """
-        Generate an item ID for a bulk download item.
-
-        :param board_id: The ID of the board whose name is to be included in the item id.
-        :return: The generated item ID.
-        """
-
-    @abstractmethod
-    def delete(self, bulk_download_item_name: str) -> None:
-        """
-        Delete the bulk download file.
-
-        :param bulk_download_item_name: The name of the bulk download item.
-        """

invokeai/app/services/bulk_download/bulk_download_common.py

@@ -1,25 +0,0 @@
-DEFAULT_BULK_DOWNLOAD_ID = "default"
-
-
-class BulkDownloadException(Exception):
-    """Exception raised when a bulk download fails."""
-
-    def __init__(self, message="Bulk download failed"):
-        super().__init__(message)
-        self.message = message
-
-
-class BulkDownloadTargetException(BulkDownloadException):
-    """Exception raised when a bulk download target is not found."""
-
-    def __init__(self, message="The bulk download target was not found"):
-        super().__init__(message)
-        self.message = message
-
-
-class BulkDownloadParametersException(BulkDownloadException):
-    """Exception raised when a bulk download parameter is invalid."""
-
-    def __init__(self, message="No image names or board ID provided"):
-        super().__init__(message)
-        self.message = message

invokeai/app/services/bulk_download/bulk_download_default.py

@@ -1,157 +0,0 @@
-from pathlib import Path
-from tempfile import TemporaryDirectory
-from typing import Optional, Union
-from zipfile import ZipFile
-
-from invokeai.app.services.board_records.board_records_common import BoardRecordNotFoundException
-from invokeai.app.services.bulk_download.bulk_download_common import (
-    DEFAULT_BULK_DOWNLOAD_ID,
-    BulkDownloadException,
-    BulkDownloadParametersException,
-    BulkDownloadTargetException,
-)
-from invokeai.app.services.image_records.image_records_common import ImageRecordNotFoundException
-from invokeai.app.services.images.images_common import ImageDTO
-from invokeai.app.services.invoker import Invoker
-from invokeai.app.util.misc import uuid_string
-
-from .bulk_download_base import BulkDownloadBase
-
-
-class BulkDownloadService(BulkDownloadBase):
-    def start(self, invoker: Invoker) -> None:
-        self._invoker = invoker
-
-    def __init__(self):
-        self._temp_directory = TemporaryDirectory()
-        self._bulk_downloads_folder = Path(self._temp_directory.name) / "bulk_downloads"
-        self._bulk_downloads_folder.mkdir(parents=True, exist_ok=True)
-
-    def handler(
-        self, image_names: Optional[list[str]], board_id: Optional[str], bulk_download_item_id: Optional[str]
-    ) -> None:
-        bulk_download_id: str = DEFAULT_BULK_DOWNLOAD_ID
-        bulk_download_item_id = bulk_download_item_id or uuid_string()
-        bulk_download_item_name = bulk_download_item_id + ".zip"
-
-        self._signal_job_started(bulk_download_id, bulk_download_item_id, bulk_download_item_name)
-
-        try:
-            image_dtos: list[ImageDTO] = []
-
-            if board_id:
-                image_dtos = self._board_handler(board_id)
-            elif image_names:
-                image_dtos = self._image_handler(image_names)
-            else:
-                raise BulkDownloadParametersException()
-
-            bulk_download_item_name: str = self._create_zip_file(image_dtos, bulk_download_item_id)
-            self._signal_job_completed(bulk_download_id, bulk_download_item_id, bulk_download_item_name)
-        except (
-            ImageRecordNotFoundException,
-            BoardRecordNotFoundException,
-            BulkDownloadException,
-            BulkDownloadParametersException,
-        ) as e:
-            self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e)
-        except Exception as e:
-            self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e)
-            self._invoker.services.logger.error("Problem bulk downloading images.")
-            raise e
-
-    def _image_handler(self, image_names: list[str]) -> list[ImageDTO]:
-        return [self._invoker.services.images.get_dto(image_name) for image_name in image_names]
-
-    def _board_handler(self, board_id: str) -> list[ImageDTO]:
-        image_names = self._invoker.services.board_image_records.get_all_board_image_names_for_board(board_id)
-        return self._image_handler(image_names)
-
-    def generate_item_id(self, board_id: Optional[str]) -> str:
-        return uuid_string() if board_id is None else self._get_clean_board_name(board_id) + "_" + uuid_string()
-
-    def _get_clean_board_name(self, board_id: str) -> str:
-        if board_id == "none":
-            return "Uncategorized"
-
-        return self._clean_string_to_path_safe(self._invoker.services.board_records.get(board_id).board_name)
-
-    def _create_zip_file(self, image_dtos: list[ImageDTO], bulk_download_item_id: str) -> str:
-        """
-        Create a zip file containing the images specified by the given image names or board id.
-        If download with the same bulk_download_id already exists, it will be overwritten.
-
-        :return: The name of the zip file.
-        """
-        zip_file_name = bulk_download_item_id + ".zip"
-        zip_file_path = self._bulk_downloads_folder / (zip_file_name)
-
-        with ZipFile(zip_file_path, "w") as zip_file:
-            for image_dto in image_dtos:
-                image_zip_path = Path(image_dto.image_category.value) / image_dto.image_name
-                image_disk_path = self._invoker.services.images.get_path(image_dto.image_name)
-                zip_file.write(image_disk_path, arcname=image_zip_path)
-
-        return str(zip_file_name)
-
-    # from https://stackoverflow.com/questions/7406102/create-sane-safe-filename-from-any-unsafe-string
-    def _clean_string_to_path_safe(self, s: str) -> str:
-        """Clean a string to be path safe."""
-        return "".join([c for c in s if c.isalpha() or c.isdigit() or c == " " or c == "_" or c == "-"]).rstrip()
-
-    def _signal_job_started(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
-    ) -> None:
-        """Signal that a bulk download job has started."""
-        if self._invoker:
-            assert bulk_download_id is not None
-            self._invoker.services.events.emit_bulk_download_started(
-                bulk_download_id=bulk_download_id,
-                bulk_download_item_id=bulk_download_item_id,
-                bulk_download_item_name=bulk_download_item_name,
-            )
-
-    def _signal_job_completed(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
-    ) -> None:
-        """Signal that a bulk download job has completed."""
-        if self._invoker:
-            assert bulk_download_id is not None
-            assert bulk_download_item_name is not None
-            self._invoker.services.events.emit_bulk_download_completed(
-                bulk_download_id=bulk_download_id,
-                bulk_download_item_id=bulk_download_item_id,
-                bulk_download_item_name=bulk_download_item_name,
-            )
-
-    def _signal_job_failed(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, exception: Exception
-    ) -> None:
-        """Signal that a bulk download job has failed."""
-        if self._invoker:
-            assert bulk_download_id is not None
-            assert exception is not None
-            self._invoker.services.events.emit_bulk_download_failed(
-                bulk_download_id=bulk_download_id,
-                bulk_download_item_id=bulk_download_item_id,
-                bulk_download_item_name=bulk_download_item_name,
-                error=str(exception),
-            )
-
-    def stop(self, *args, **kwargs):
-        self._temp_directory.cleanup()
-
-    def delete(self, bulk_download_item_name: str) -> None:
-        path = self.get_path(bulk_download_item_name)
-        Path(path).unlink()
-
-    def get_path(self, bulk_download_item_name: str) -> str:
-        path = str(self._bulk_downloads_folder / bulk_download_item_name)
-        if not self._is_valid_path(path):
-            raise BulkDownloadTargetException()
-        return path
-
-    def _is_valid_path(self, path: Union[str, Path]) -> bool:
-        """Validates the path given for a bulk download."""
-        path = path if isinstance(path, Path) else Path(path)
-        return path.exists()

invokeai/app/services/config/config_base.py

@@ -156,7 +156,6 @@ class InvokeAISettings(BaseSettings):
             "lora_dir",
             "embedding_dir",
             "controlnet_dir",
-            "conf_path",
         ]
 
     @classmethod

invokeai/app/services/config/config_default.py

@@ -30,6 +30,7 @@ InvokeAI:
     lora_dir: null
     embedding_dir: null
     controlnet_dir: null
+    conf_path: configs/models.yaml
     models_dir: models
     legacy_conf_dir: configs/stable-diffusion
     db_dir: databases
@@ -122,6 +123,7 @@ a Path object:
 
  root_path          - path to InvokeAI root
  output_path        - path to default outputs directory
+ model_conf_path    - path to models.yaml
  conf               - alias for the above
  embedding_path     - path to the embeddings directory
  lora_path          - path to the LoRA directory
@@ -161,6 +163,7 @@ two configs are kept in separate sections of the config file:
   InvokeAI:
      Paths:
         root: /home/lstein/invokeai-main
+        conf_path: configs/models.yaml
         legacy_conf_dir: configs/stable-diffusion
         outdir: outputs
      ...
@@ -197,7 +200,6 @@ class Categories(object):
     Development: JsonDict = {"category": "Development"}
     Other: JsonDict = {"category": "Other"}
     ModelCache: JsonDict = {"category": "Model Cache"}
-    ModelImport: JsonDict = {"category": "Model Import"}
     Device: JsonDict = {"category": "Device"}
     Generation: JsonDict = {"category": "Generation"}
     Queue: JsonDict = {"category": "Queue"}
@@ -235,6 +237,7 @@ class InvokeAIAppConfig(InvokeAISettings):
     # PATHS
     root                : Optional[Path] = Field(default=None, description='InvokeAI runtime root directory', json_schema_extra=Categories.Paths)
     autoimport_dir      : Path = Field(default=Path('autoimport'), description='Path to a directory of models files to be imported on startup.', json_schema_extra=Categories.Paths)
+    conf_path           : Path = Field(default=Path('configs/models.yaml'), description='Path to models definition file', json_schema_extra=Categories.Paths)
     models_dir          : Path = Field(default=Path('models'), description='Path to the models directory', json_schema_extra=Categories.Paths)
     convert_cache_dir   : Path = Field(default=Path('models/.cache'), description='Path to the converted models cache directory', json_schema_extra=Categories.Paths)
     legacy_conf_dir     : Path = Field(default=Path('configs/stable-diffusion'), description='Path to directory of legacy checkpoint config files', json_schema_extra=Categories.Paths)
@@ -287,8 +290,7 @@ class InvokeAIAppConfig(InvokeAISettings):
     node_cache_size     : int = Field(default=512, description="How many cached nodes to keep in memory", json_schema_extra=Categories.Nodes)
 
     # MODEL IMPORT
-    civitai_api_key     : Optional[str] = Field(default=os.environ.get("CIVITAI_API_KEY"), description="API key for CivitAI", json_schema_extra=Categories.ModelImport)
-    model_sym_links     : bool = Field(default=False, description="If true, create symbolic links to models instead of copying them. [REQUIRES ADMIN PERMISSIONS OR DEVELOPER MODE IN WINDOWS]", json_schema_extra=Categories.ModelImport)
+    civitai_api_key       : Optional[str] = Field(default=os.environ.get("CIVITAI_API_KEY"), description="API key for CivitAI", json_schema_extra=Categories.Other)
 
     # DEPRECATED FIELDS - STILL HERE IN ORDER TO OBTAN VALUES FROM PRE-3.1 CONFIG FILES
     always_use_cpu      : bool = Field(default=False, description="If true, use the CPU for rendering even if a GPU is available.", json_schema_extra=Categories.MemoryPerformance)
@@ -299,7 +301,6 @@ class InvokeAIAppConfig(InvokeAISettings):
     lora_dir            : Optional[Path] = Field(default=None, description='Path to a directory of LoRA/LyCORIS models to be imported on startup.', json_schema_extra=Categories.Paths)
     embedding_dir       : Optional[Path] = Field(default=None, description='Path to a directory of Textual Inversion embeddings to be imported on startup.', json_schema_extra=Categories.Paths)
     controlnet_dir      : Optional[Path] = Field(default=None, description='Path to a directory of ControlNet embeddings to be imported on startup.', json_schema_extra=Categories.Paths)
-    conf_path           : Path = Field(default=Path('configs/models.yaml'), description='Path to models definition file', json_schema_extra=Categories.Paths)
 
     # this is not referred to in the source code and can be removed entirely
     #free_gpu_mem        : Optional[bool] = Field(default=None, description="If true, purge model from GPU after each generation.", json_schema_extra=Categories.MemoryPerformance)

invokeai/app/services/events/events_base.py

@@ -3,7 +3,7 @@
 
 from typing import Any, Dict, List, Optional, Union
 
-from invokeai.app.services.session_processor.session_processor_common import ProgressImage
+from invokeai.app.services.invocation_processor.invocation_processor_common import ProgressImage
 from invokeai.app.services.session_queue.session_queue_common import (
     BatchStatus,
     EnqueueBatchResult,
@@ -16,7 +16,6 @@ from invokeai.backend.model_manager import AnyModelConfig
 
 class EventServiceBase:
     queue_event: str = "queue_event"
-    bulk_download_event: str = "bulk_download_event"
     download_event: str = "download_event"
     model_event: str = "model_event"
 
@@ -25,14 +24,6 @@ class EventServiceBase:
     def dispatch(self, event_name: str, payload: Any) -> None:
         pass
 
-    def _emit_bulk_download_event(self, event_name: str, payload: dict) -> None:
-        """Bulk download events are emitted to a room with queue_id as the room name"""
-        payload["timestamp"] = get_timestamp()
-        self.dispatch(
-            event_name=EventServiceBase.bulk_download_event,
-            payload={"event": event_name, "data": payload},
-        )
-
     def __emit_queue_event(self, event_name: str, payload: dict) -> None:
         """Queue events are emitted to a room with queue_id as the room name"""
         payload["timestamp"] = get_timestamp()
@@ -213,6 +204,52 @@ class EventServiceBase:
             },
         )
 
+    def emit_session_retrieval_error(
+        self,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        error_type: str,
+        error: str,
+    ) -> None:
+        """Emitted when session retrieval fails"""
+        self.__emit_queue_event(
+            event_name="session_retrieval_error",
+            payload={
+                "queue_id": queue_id,
+                "queue_item_id": queue_item_id,
+                "queue_batch_id": queue_batch_id,
+                "graph_execution_state_id": graph_execution_state_id,
+                "error_type": error_type,
+                "error": error,
+            },
+        )
+
+    def emit_invocation_retrieval_error(
+        self,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        node_id: str,
+        error_type: str,
+        error: str,
+    ) -> None:
+        """Emitted when invocation retrieval fails"""
+        self.__emit_queue_event(
+            event_name="invocation_retrieval_error",
+            payload={
+                "queue_id": queue_id,
+                "queue_item_id": queue_item_id,
+                "queue_batch_id": queue_batch_id,
+                "graph_execution_state_id": graph_execution_state_id,
+                "node_id": node_id,
+                "error_type": error_type,
+                "error": error,
+            },
+        )
+
     def emit_session_canceled(
         self,
         queue_id: str,
@@ -357,7 +394,6 @@ class EventServiceBase:
         bytes: int,
         total_bytes: int,
         parts: List[Dict[str, Union[str, int]]],
-        id: int,
     ) -> None:
         """
         Emit at intervals while the install job is in progress (remote models only).
@@ -377,7 +413,6 @@ class EventServiceBase:
                 "bytes": bytes,
                 "total_bytes": total_bytes,
                 "parts": parts,
-                "id": id,
             },
         )
 
@@ -392,7 +427,7 @@ class EventServiceBase:
             payload={"source": source},
         )
 
-    def emit_model_install_completed(self, source: str, key: str, id: int, total_bytes: Optional[int] = None) -> None:
+    def emit_model_install_completed(self, source: str, key: str, total_bytes: Optional[int] = None) -> None:
         """
         Emit when an install job is completed successfully.
 
@@ -402,7 +437,11 @@ class EventServiceBase:
         """
         self.__emit_model_event(
             event_name="model_install_completed",
-            payload={"source": source, "total_bytes": total_bytes, "key": key, "id": id},
+            payload={
+                "source": source,
+                "total_bytes": total_bytes,
+                "key": key,
+            },
         )
 
     def emit_model_install_cancelled(self, source: str) -> None:
@@ -416,7 +455,12 @@ class EventServiceBase:
             payload={"source": source},
         )
 
-    def emit_model_install_error(self, source: str, error_type: str, error: str, id: int) -> None:
+    def emit_model_install_error(
+        self,
+        source: str,
+        error_type: str,
+        error: str,
+    ) -> None:
         """
         Emit when an install job encounters an exception.
 
@@ -426,45 +470,9 @@ class EventServiceBase:
         """
         self.__emit_model_event(
             event_name="model_install_error",
-            payload={"source": source, "error_type": error_type, "error": error, "id": id},
-        )
-
-    def emit_bulk_download_started(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
-    ) -> None:
-        """Emitted when a bulk download starts"""
-        self._emit_bulk_download_event(
-            event_name="bulk_download_started",
             payload={
-                "bulk_download_id": bulk_download_id,
-                "bulk_download_item_id": bulk_download_item_id,
-                "bulk_download_item_name": bulk_download_item_name,
-            },
-        )
-
-    def emit_bulk_download_completed(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
-    ) -> None:
-        """Emitted when a bulk download completes"""
-        self._emit_bulk_download_event(
-            event_name="bulk_download_completed",
-            payload={
-                "bulk_download_id": bulk_download_id,
-                "bulk_download_item_id": bulk_download_item_id,
-                "bulk_download_item_name": bulk_download_item_name,
-            },
-        )
-
-    def emit_bulk_download_failed(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, error: str
-    ) -> None:
-        """Emitted when a bulk download fails"""
-        self._emit_bulk_download_event(
-            event_name="bulk_download_failed",
-            payload={
-                "bulk_download_id": bulk_download_id,
-                "bulk_download_item_id": bulk_download_item_id,
-                "bulk_download_item_name": bulk_download_item_name,
+                "source": source,
+                "error_type": error_type,
                 "error": error,
             },
         )

invokeai/app/services/invocation_processor/invocation_processor_base.py

@@ -0,0 +1,5 @@
+from abc import ABC
+
+
+class InvocationProcessorABC(ABC):  # noqa: B024
+    pass

invokeai/app/services/invocation_processor/invocation_processor_common.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, Field
+
+
+class ProgressImage(BaseModel):
+    """The progress image sent intermittently during processing"""
+
+    width: int = Field(description="The effective width of the image in pixels")
+    height: int = Field(description="The effective height of the image in pixels")
+    dataURL: str = Field(description="The image data as a b64 data URL")
+
+
+class CanceledException(Exception):
+    """Execution canceled by user."""
+
+    pass

invokeai/app/services/invocation_processor/invocation_processor_default.py

@@ -0,0 +1,243 @@
+import time
+import traceback
+from contextlib import suppress
+from threading import BoundedSemaphore, Event, Thread
+from typing import Optional
+
+import invokeai.backend.util.logging as logger
+from invokeai.app.services.invocation_queue.invocation_queue_common import InvocationQueueItem
+from invokeai.app.services.invocation_stats.invocation_stats_common import (
+    GESStatsNotFoundError,
+)
+from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context
+from invokeai.app.util.profiler import Profiler
+
+from ..invoker import Invoker
+from .invocation_processor_base import InvocationProcessorABC
+from .invocation_processor_common import CanceledException
+
+
+class DefaultInvocationProcessor(InvocationProcessorABC):
+    __invoker_thread: Thread
+    __stop_event: Event
+    __invoker: Invoker
+    __threadLimit: BoundedSemaphore
+
+    def start(self, invoker: Invoker) -> None:
+        # LS - this will probably break
+        # but the idea is to enable multithreading up to the number of available
+        # GPUs. Nodes will block on model loading if no GPU is free.
+        self.__threadLimit = BoundedSemaphore(invoker.services.model_manager.gpu_count)
+        self.__invoker = invoker
+        self.__stop_event = Event()
+        self.__invoker_thread = Thread(
+            name="invoker_processor",
+            target=self.__process,
+            kwargs={"stop_event": self.__stop_event},
+        )
+        self.__invoker_thread.daemon = True  # TODO: make async and do not use threads
+        self.__invoker_thread.start()
+
+    def stop(self, *args, **kwargs) -> None:
+        self.__stop_event.set()
+
+    def __process(self, stop_event: Event):
+        try:
+            self.__threadLimit.acquire()
+            queue_item: Optional[InvocationQueueItem] = None
+
+            profiler = (
+                Profiler(
+                    logger=self.__invoker.services.logger,
+                    output_dir=self.__invoker.services.configuration.profiles_path,
+                    prefix=self.__invoker.services.configuration.profile_prefix,
+                )
+                if self.__invoker.services.configuration.profile_graphs
+                else None
+            )
+
+            def stats_cleanup(graph_execution_state_id: str) -> None:
+                if profiler:
+                    profile_path = profiler.stop()
+                    stats_path = profile_path.with_suffix(".json")
+                    self.__invoker.services.performance_statistics.dump_stats(
+                        graph_execution_state_id=graph_execution_state_id, output_path=stats_path
+                    )
+                with suppress(GESStatsNotFoundError):
+                    self.__invoker.services.performance_statistics.log_stats(graph_execution_state_id)
+                    self.__invoker.services.performance_statistics.reset_stats(graph_execution_state_id)
+
+            while not stop_event.is_set():
+                try:
+                    queue_item = self.__invoker.services.queue.get()
+                except Exception as e:
+                    self.__invoker.services.logger.error("Exception while getting from queue:\n%s" % e)
+
+                if not queue_item:  # Probably stopping
+                    # do not hammer the queue
+                    time.sleep(0.5)
+                    continue
+
+                if profiler and profiler.profile_id != queue_item.graph_execution_state_id:
+                    profiler.start(profile_id=queue_item.graph_execution_state_id)
+
+                try:
+                    graph_execution_state = self.__invoker.services.graph_execution_manager.get(
+                        queue_item.graph_execution_state_id
+                    )
+                except Exception as e:
+                    self.__invoker.services.logger.error("Exception while retrieving session:\n%s" % e)
+                    self.__invoker.services.events.emit_session_retrieval_error(
+                        queue_batch_id=queue_item.session_queue_batch_id,
+                        queue_item_id=queue_item.session_queue_item_id,
+                        queue_id=queue_item.session_queue_id,
+                        graph_execution_state_id=queue_item.graph_execution_state_id,
+                        error_type=e.__class__.__name__,
+                        error=traceback.format_exc(),
+                    )
+                    continue
+
+                try:
+                    invocation = graph_execution_state.execution_graph.get_node(queue_item.invocation_id)
+                except Exception as e:
+                    self.__invoker.services.logger.error("Exception while retrieving invocation:\n%s" % e)
+                    self.__invoker.services.events.emit_invocation_retrieval_error(
+                        queue_batch_id=queue_item.session_queue_batch_id,
+                        queue_item_id=queue_item.session_queue_item_id,
+                        queue_id=queue_item.session_queue_id,
+                        graph_execution_state_id=queue_item.graph_execution_state_id,
+                        node_id=queue_item.invocation_id,
+                        error_type=e.__class__.__name__,
+                        error=traceback.format_exc(),
+                    )
+                    continue
+
+                # get the source node id to provide to clients (the prepared node id is not as useful)
+                source_node_id = graph_execution_state.prepared_source_mapping[invocation.id]
+
+                # Send starting event
+                self.__invoker.services.events.emit_invocation_started(
+                    queue_batch_id=queue_item.session_queue_batch_id,
+                    queue_item_id=queue_item.session_queue_item_id,
+                    queue_id=queue_item.session_queue_id,
+                    graph_execution_state_id=graph_execution_state.id,
+                    node=invocation.model_dump(),
+                    source_node_id=source_node_id,
+                )
+
+                # Invoke
+                try:
+                    graph_id = graph_execution_state.id
+                    with self.__invoker.services.performance_statistics.collect_stats(invocation, graph_id):
+                        # use the internal invoke_internal(), which wraps the node's invoke() method,
+                        # which handles a few things:
+                        # - nodes that require a value, but get it only from a connection
+                        # - referencing the invocation cache instead of executing the node
+                        context_data = InvocationContextData(
+                            invocation=invocation,
+                            session_id=graph_id,
+                            workflow=queue_item.workflow,
+                            source_node_id=source_node_id,
+                            queue_id=queue_item.session_queue_id,
+                            queue_item_id=queue_item.session_queue_item_id,
+                            batch_id=queue_item.session_queue_batch_id,
+                        )
+                        context = build_invocation_context(
+                            services=self.__invoker.services,
+                            context_data=context_data,
+                        )
+                        outputs = invocation.invoke_internal(context=context, services=self.__invoker.services)
+
+                        # Check queue to see if this is canceled, and skip if so
+                        if self.__invoker.services.queue.is_canceled(graph_execution_state.id):
+                            continue
+
+                        # Save outputs and history
+                        graph_execution_state.complete(invocation.id, outputs)
+
+                        # Save the state changes
+                        self.__invoker.services.graph_execution_manager.set(graph_execution_state)
+
+                        # Send complete event
+                        self.__invoker.services.events.emit_invocation_complete(
+                            queue_batch_id=queue_item.session_queue_batch_id,
+                            queue_item_id=queue_item.session_queue_item_id,
+                            queue_id=queue_item.session_queue_id,
+                            graph_execution_state_id=graph_execution_state.id,
+                            node=invocation.model_dump(),
+                            source_node_id=source_node_id,
+                            result=outputs.model_dump(),
+                        )
+
+                except KeyboardInterrupt:
+                    pass
+
+                except CanceledException:
+                    stats_cleanup(graph_execution_state.id)
+                    pass
+
+                except Exception as e:
+                    error = traceback.format_exc()
+                    logger.error(error)
+
+                    # Save error
+                    graph_execution_state.set_node_error(invocation.id, error)
+
+                    # Save the state changes
+                    self.__invoker.services.graph_execution_manager.set(graph_execution_state)
+
+                    self.__invoker.services.logger.error("Error while invoking:\n%s" % e)
+                    # Send error event
+                    self.__invoker.services.events.emit_invocation_error(
+                        queue_batch_id=queue_item.session_queue_batch_id,
+                        queue_item_id=queue_item.session_queue_item_id,
+                        queue_id=queue_item.session_queue_id,
+                        graph_execution_state_id=graph_execution_state.id,
+                        node=invocation.model_dump(),
+                        source_node_id=source_node_id,
+                        error_type=e.__class__.__name__,
+                        error=error,
+                    )
+                    pass
+
+                # Check queue to see if this is canceled, and skip if so
+                if self.__invoker.services.queue.is_canceled(graph_execution_state.id):
+                    continue
+
+                # Queue any further commands if invoking all
+                is_complete = graph_execution_state.is_complete()
+                if queue_item.invoke_all and not is_complete:
+                    try:
+                        self.__invoker.invoke(
+                            session_queue_batch_id=queue_item.session_queue_batch_id,
+                            session_queue_item_id=queue_item.session_queue_item_id,
+                            session_queue_id=queue_item.session_queue_id,
+                            graph_execution_state=graph_execution_state,
+                            workflow=queue_item.workflow,
+                            invoke_all=True,
+                        )
+                    except Exception as e:
+                        self.__invoker.services.logger.error("Error while invoking:\n%s" % e)
+                        self.__invoker.services.events.emit_invocation_error(
+                            queue_batch_id=queue_item.session_queue_batch_id,
+                            queue_item_id=queue_item.session_queue_item_id,
+                            queue_id=queue_item.session_queue_id,
+                            graph_execution_state_id=graph_execution_state.id,
+                            node=invocation.model_dump(),
+                            source_node_id=source_node_id,
+                            error_type=e.__class__.__name__,
+                            error=traceback.format_exc(),
+                        )
+                elif is_complete:
+                    self.__invoker.services.events.emit_graph_execution_complete(
+                        queue_batch_id=queue_item.session_queue_batch_id,
+                        queue_item_id=queue_item.session_queue_item_id,
+                        queue_id=queue_item.session_queue_id,
+                        graph_execution_state_id=graph_execution_state.id,
+                    )
+                    stats_cleanup(graph_execution_state.id)
+
+        except KeyboardInterrupt:
+            pass  # Log something? KeyboardInterrupt is probably not going to be seen by the processor
+        finally:
+            self.__threadLimit.release()

invokeai/app/services/invocation_queue/invocation_queue_base.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from .invocation_queue_common import InvocationQueueItem
+
+
+class InvocationQueueABC(ABC):
+    """Abstract base class for all invocation queues"""
+
+    @abstractmethod
+    def get(self) -> InvocationQueueItem:
+        pass
+
+    @abstractmethod
+    def put(self, item: Optional[InvocationQueueItem]) -> None:
+        pass
+
+    @abstractmethod
+    def cancel(self, graph_execution_state_id: str) -> None:
+        pass
+
+    @abstractmethod
+    def is_canceled(self, graph_execution_state_id: str) -> bool:
+        pass

invokeai/app/services/invocation_queue/invocation_queue_common.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+import time
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
+
+
+class InvocationQueueItem(BaseModel):
+    graph_execution_state_id: str = Field(description="The ID of the graph execution state")
+    invocation_id: str = Field(description="The ID of the node being invoked")
+    session_queue_id: str = Field(description="The ID of the session queue from which this invocation queue item came")
+    session_queue_item_id: int = Field(
+        description="The ID of session queue item from which this invocation queue item came"
+    )
+    session_queue_batch_id: str = Field(
+        description="The ID of the session batch from which this invocation queue item came"
+    )
+    workflow: Optional[WorkflowWithoutID] = Field(description="The workflow associated with this queue item")
+    invoke_all: bool = Field(default=False)
+    timestamp: float = Field(default_factory=time.time)

invokeai/app/services/invocation_queue/invocation_queue_memory.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+import time
+from queue import Queue
+from typing import Optional
+
+from .invocation_queue_base import InvocationQueueABC
+from .invocation_queue_common import InvocationQueueItem
+
+
+class MemoryInvocationQueue(InvocationQueueABC):
+    __queue: Queue
+    __cancellations: dict[str, float]
+
+    def __init__(self):
+        self.__queue = Queue()
+        self.__cancellations = {}
+
+    def get(self) -> InvocationQueueItem:
+        item = self.__queue.get()
+
+        while (
+            isinstance(item, InvocationQueueItem)
+            and item.graph_execution_state_id in self.__cancellations
+            and self.__cancellations[item.graph_execution_state_id] > item.timestamp
+        ):
+            item = self.__queue.get()
+
+        # Clear old items
+        for graph_execution_state_id in list(self.__cancellations.keys()):
+            if self.__cancellations[graph_execution_state_id] < item.timestamp:
+                del self.__cancellations[graph_execution_state_id]
+
+        return item
+
+    def put(self, item: Optional[InvocationQueueItem]) -> None:
+        self.__queue.put(item)
+
+    def cancel(self, graph_execution_state_id: str) -> None:
+        if graph_execution_state_id not in self.__cancellations:
+            self.__cancellations[graph_execution_state_id] = time.time()
+
+    def is_canceled(self, graph_execution_state_id: str) -> bool:
+        return graph_execution_state_id in self.__cancellations

invokeai/app/services/invocation_services.py

@@ -16,7 +16,6 @@ if TYPE_CHECKING:
     from .board_images.board_images_base import BoardImagesServiceABC
     from .board_records.board_records_base import BoardRecordStorageBase
     from .boards.boards_base import BoardServiceABC
-    from .bulk_download.bulk_download_base import BulkDownloadBase
     from .config import InvokeAIAppConfig
     from .download import DownloadQueueServiceBase
     from .events.events_base import EventServiceBase
@@ -24,11 +23,15 @@ if TYPE_CHECKING:
     from .image_records.image_records_base import ImageRecordStorageBase
     from .images.images_base import ImageServiceABC
     from .invocation_cache.invocation_cache_base import InvocationCacheBase
+    from .invocation_processor.invocation_processor_base import InvocationProcessorABC
+    from .invocation_queue.invocation_queue_base import InvocationQueueABC
     from .invocation_stats.invocation_stats_base import InvocationStatsServiceBase
+    from .item_storage.item_storage_base import ItemStorageABC
     from .model_manager.model_manager_base import ModelManagerServiceBase
     from .names.names_base import NameServiceBase
     from .session_processor.session_processor_base import SessionProcessorBase
     from .session_queue.session_queue_base import SessionQueueBase
+    from .shared.graph import GraphExecutionState
     from .urls.urls_base import UrlServiceBase
     from .workflow_records.workflow_records_base import WorkflowRecordsStorageBase
 
@@ -42,16 +45,18 @@ class InvocationServices:
         board_image_records: "BoardImageRecordStorageBase",
         boards: "BoardServiceABC",
         board_records: "BoardRecordStorageBase",
-        bulk_download: "BulkDownloadBase",
         configuration: "InvokeAIAppConfig",
         events: "EventServiceBase",
+        graph_execution_manager: "ItemStorageABC[GraphExecutionState]",
         images: "ImageServiceABC",
         image_files: "ImageFileStorageBase",
         image_records: "ImageRecordStorageBase",
         logger: "Logger",
         model_manager: "ModelManagerServiceBase",
         download_queue: "DownloadQueueServiceBase",
+        processor: "InvocationProcessorABC",
         performance_statistics: "InvocationStatsServiceBase",
+        queue: "InvocationQueueABC",
         session_queue: "SessionQueueBase",
         session_processor: "SessionProcessorBase",
         invocation_cache: "InvocationCacheBase",
@@ -65,16 +70,18 @@ class InvocationServices:
         self.board_image_records = board_image_records
         self.boards = boards
         self.board_records = board_records
-        self.bulk_download = bulk_download
         self.configuration = configuration
         self.events = events
+        self.graph_execution_manager = graph_execution_manager
         self.images = images
         self.image_files = image_files
         self.image_records = image_records
         self.logger = logger
         self.model_manager = model_manager
         self.download_queue = download_queue
+        self.processor = processor
         self.performance_statistics = performance_statistics
+        self.queue = queue
         self.session_queue = session_queue
         self.session_processor = session_processor
         self.invocation_cache = invocation_cache

invokeai/app/services/invocation_stats/invocation_stats_base.py

@@ -3,7 +3,7 @@
 
 Usage:
 
-statistics = InvocationStatsService()
+statistics = InvocationStatsService(graph_execution_manager)
 with statistics.collect_stats(invocation, graph_execution_state.id):
       ... execute graphs...
 statistics.log_stats()
@@ -30,7 +30,7 @@ writes to the system log is stored in InvocationServices.performance_statistics.
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import ContextManager
+from typing import Iterator
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
@@ -50,7 +50,7 @@ class InvocationStatsServiceBase(ABC):
         self,
         invocation: BaseInvocation,
         graph_execution_state_id: str,
-    ) -> ContextManager[None]:
+    ) -> Iterator[None]:
         """
         Return a context object that will capture the statistics on the execution
         of invocaation. Use with: to place around the part of the code that executes the invocation.
@@ -60,8 +60,12 @@ class InvocationStatsServiceBase(ABC):
         pass
 
     @abstractmethod
-    def reset_stats(self):
-        """Reset all stored statistics."""
+    def reset_stats(self, graph_execution_state_id: str) -> None:
+        """
+        Reset all statistics for the indicated graph.
+        :param graph_execution_state_id: The id of the session whose stats to reset.
+        :raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
+        """
         pass
 
     @abstractmethod

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -2,7 +2,7 @@ import json
 import time
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Generator
+from typing import Iterator
 
 import psutil
 import torch
@@ -10,6 +10,7 @@ import torch
 import invokeai.backend.util.logging as logger
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invoker import Invoker
+from invokeai.app.services.item_storage.item_storage_common import ItemNotFoundError
 from invokeai.backend.model_manager.load.model_cache import CacheStats
 
 from .invocation_stats_base import InvocationStatsServiceBase
@@ -41,7 +42,7 @@ class InvocationStatsService(InvocationStatsServiceBase):
         self._invoker = invoker
 
     @contextmanager
-    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Generator[None, None, None]:
+    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Iterator[None]:
         # This is to handle case of the model manager not being initialized, which happens
         # during some tests.
         services = self._invoker.services
@@ -50,6 +51,9 @@ class InvocationStatsService(InvocationStatsServiceBase):
             self._stats[graph_execution_state_id] = GraphExecutionStats()
             self._cache_stats[graph_execution_state_id] = CacheStats()
 
+            # Prune stale stats. There should be none since we're starting a new graph, but just in case.
+            self._prune_stale_stats()
+
         # Record state before the invocation.
         start_time = time.time()
         start_ram = psutil.Process().memory_info().rss
@@ -74,9 +78,42 @@ 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 _prune_stale_stats(self) -> None:
+        """Check all graphs being tracked and prune any that have completed/errored.
+
+        This shouldn't be necessary, but we don't have totally robust upstream handling of graph completions/errors, so
+        for now we call this function periodically to prevent them from accumulating.
+        """
+        to_prune: list[str] = []
+        for graph_execution_state_id in self._stats:
+            try:
+                graph_execution_state = self._invoker.services.graph_execution_manager.get(graph_execution_state_id)
+            except ItemNotFoundError:
+                # TODO(ryand): What would cause this? Should this exception just be allowed to propagate?
+                logger.warning(f"Failed to get graph state for {graph_execution_state_id}.")
+                continue
+
+            if not graph_execution_state.is_complete():
+                # The graph is still running, don't prune it.
+                continue
+
+            to_prune.append(graph_execution_state_id)
+
+        for graph_execution_state_id in to_prune:
+            del self._stats[graph_execution_state_id]
+            del self._cache_stats[graph_execution_state_id]
+
+        if len(to_prune) > 0:
+            logger.info(f"Pruned stale graph stats for {to_prune}.")
+
+    def reset_stats(self, graph_execution_state_id: str):
+        try:
+            del self._stats[graph_execution_state_id]
+            del self._cache_stats[graph_execution_state_id]
+        except KeyError as e:
+            raise GESStatsNotFoundError(
+                f"Attempted to clear statistics for unknown graph {graph_execution_state_id}: {e}."
+            ) from e
 
     def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
         graph_stats_summary = self._get_graph_summary(graph_execution_state_id)

invokeai/app/services/invoker.py

@@ -1,7 +1,12 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
+from typing import Optional
 
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
+
+from .invocation_queue.invocation_queue_common import InvocationQueueItem
 from .invocation_services import InvocationServices
+from .shared.graph import Graph, GraphExecutionState
 
 
 class Invoker:
@@ -13,6 +18,51 @@ class Invoker:
         self.services = services
         self._start()
 
+    def invoke(
+        self,
+        session_queue_id: str,
+        session_queue_item_id: int,
+        session_queue_batch_id: str,
+        graph_execution_state: GraphExecutionState,
+        workflow: Optional[WorkflowWithoutID] = None,
+        invoke_all: bool = False,
+    ) -> Optional[str]:
+        """Determines the next node to invoke and enqueues it, preparing if needed.
+        Returns the id of the queued node, or `None` if there are no nodes left to enqueue."""
+
+        # Get the next invocation
+        invocation = graph_execution_state.next()
+        if not invocation:
+            return None
+
+        # Save the execution state
+        self.services.graph_execution_manager.set(graph_execution_state)
+
+        # Queue the invocation
+        self.services.queue.put(
+            InvocationQueueItem(
+                session_queue_id=session_queue_id,
+                session_queue_item_id=session_queue_item_id,
+                session_queue_batch_id=session_queue_batch_id,
+                graph_execution_state_id=graph_execution_state.id,
+                invocation_id=invocation.id,
+                workflow=workflow,
+                invoke_all=invoke_all,
+            )
+        )
+
+        return invocation.id
+
+    def create_execution_state(self, graph: Optional[Graph] = None) -> GraphExecutionState:
+        """Creates a new execution state for the given graph"""
+        new_state = GraphExecutionState(graph=Graph() if graph is None else graph)
+        self.services.graph_execution_manager.set(new_state)
+        return new_state
+
+    def cancel(self, graph_execution_state_id: str) -> None:
+        """Cancels the given execution state"""
+        self.services.queue.cancel(graph_execution_state_id)
+
     def __start_service(self, service) -> None:
         # Call start() method on any services that have it
         start_op = getattr(service, "start", None)
@@ -35,3 +85,5 @@ class Invoker:
         # First stop all services
         for service in vars(self.services):
             self.__stop_service(getattr(self.services, service))
+
+        self.services.queue.put(None)

invokeai/app/services/model_install/model_install_base.py

@@ -156,7 +156,6 @@ class ModelInstallJob(BaseModel):
 
     id: int = Field(description="Unique ID for this job")
     status: InstallStatus = Field(default=InstallStatus.WAITING, description="Current status of install process")
-    error_reason: Optional[str] = Field(default=None, description="Information about why the job failed")
     config_in: Dict[str, Any] = Field(
         default_factory=dict, description="Configuration information (e.g. 'description') to apply to model."
     )
@@ -178,12 +177,6 @@ class ModelInstallJob(BaseModel):
     download_parts: Set[DownloadJob] = Field(
         default_factory=set, description="Download jobs contributing to this install"
     )
-    error: Optional[str] = Field(
-        default=None, description="On an error condition, this field will contain the text of the exception"
-    )
-    error_traceback: Optional[str] = Field(
-        default=None, description="On an error condition, this field will contain the exception traceback"
-    )
     # internal flags and transitory settings
     _install_tmpdir: Optional[Path] = PrivateAttr(default=None)
     _exception: Optional[Exception] = PrivateAttr(default=None)
@@ -191,10 +184,7 @@ class ModelInstallJob(BaseModel):
     def set_error(self, e: Exception) -> None:
         """Record the error and traceback from an exception."""
         self._exception = e
-        self.error = str(e)
-        self.error_traceback = self._format_error(e)
         self.status = InstallStatus.ERROR
-        self.error_reason = self._exception.__class__.__name__ if self._exception else None
 
     def cancel(self) -> None:
         """Call to cancel the job."""
@@ -205,9 +195,10 @@ class ModelInstallJob(BaseModel):
         """Class name of the exception that led to status==ERROR."""
         return self._exception.__class__.__name__ if self._exception else None
 
-    def _format_error(self, exception: Exception) -> str:
+    @property
+    def error(self) -> Optional[str]:
         """Error traceback."""
-        return "".join(traceback.format_exception(exception))
+        return "".join(traceback.format_exception(self._exception)) if self._exception else None
 
     @property
     def cancelled(self) -> bool:

invokeai/app/services/model_install/model_install_default.py

@@ -154,12 +154,8 @@ class ModelInstallService(ModelInstallServiceBase):
 
         info: AnyModelConfig = self._probe_model(Path(model_path), config)
         old_hash = info.current_hash
-
-        if preferred_name := config.get("name"):
-            preferred_name = Path(preferred_name).with_suffix(model_path.suffix)
-
         dest_path = (
-            self.app_config.models_path / info.base.value / info.type.value / (preferred_name or model_path.name)
+            self.app_config.models_path / info.base.value / info.type.value / (config.get("name") or model_path.name)
         )
         try:
             new_path = self._copy_model(model_path, dest_path)
@@ -507,13 +503,10 @@ class ModelInstallService(ModelInstallServiceBase):
         if old_path == new_path:
             return old_path
         new_path.parent.mkdir(parents=True, exist_ok=True)
-        if self.app_config.model_sym_links:
-            new_path.symlink_to(old_path, target_is_directory=old_path.is_dir())
+        if old_path.is_dir():
+            copytree(old_path, new_path)
         else:
-            if old_path.is_dir():
-                copytree(old_path, new_path)
-            else:
-                copyfile(old_path, new_path)
+            copyfile(old_path, new_path)
         return new_path
 
     def _move_model(self, old_path: Path, new_path: Path) -> Path:
@@ -545,10 +538,8 @@ class ModelInstallService(ModelInstallServiceBase):
     def _register(
         self, model_path: Path, config: Optional[Dict[str, Any]] = None, info: Optional[AnyModelConfig] = None
     ) -> str:
-        key = self._create_key()
-        if config and not config.get("key", None):
-            config["key"] = key
         info = info or ModelProbe.probe(model_path, config)
+        key = self._create_key()
 
         model_path = model_path.absolute()
         if model_path.is_relative_to(self.app_config.models_path):
@@ -561,8 +552,8 @@ class ModelInstallService(ModelInstallServiceBase):
             # make config relative to our root
             legacy_conf = (self.app_config.root_dir / self.app_config.legacy_conf_dir / info.config).resolve()
             info.config = legacy_conf.relative_to(self.app_config.root_dir).as_posix()
-        self.record_store.add_model(info.key, info)
-        return info.key
+        self.record_store.add_model(key, info)
+        return key
 
     def _next_id(self) -> int:
         with self._lock:
@@ -746,7 +737,6 @@ class ModelInstallService(ModelInstallServiceBase):
                 self._signal_job_downloading(install_job)
 
     def _download_complete_callback(self, download_job: DownloadJob) -> None:
-        self._logger.info(f"{download_job.source}: model download complete")
         with self._lock:
             install_job = self._download_cache[download_job.source]
             self._download_cache.pop(download_job.source, None)
@@ -779,7 +769,7 @@ class ModelInstallService(ModelInstallServiceBase):
             if not install_job:
                 return
             self._downloads_changed_event.set()
-            self._logger.warning(f"{download_job.source}: model download cancelled")
+            self._logger.warning(f"Download {download_job.source} cancelled.")
             # if install job has already registered an error, then do not replace its status with cancelled
             if not install_job.errored:
                 install_job.cancel()
@@ -826,7 +816,6 @@ class ModelInstallService(ModelInstallServiceBase):
                 parts=parts,
                 bytes=job.bytes,
                 total_bytes=job.total_bytes,
-                id=job.id,
             )
 
     def _signal_job_completed(self, job: ModelInstallJob) -> None:
@@ -839,7 +828,7 @@ class ModelInstallService(ModelInstallServiceBase):
             assert job.local_path is not None
             assert job.config_out is not None
             key = job.config_out.key
-            self._event_bus.emit_model_install_completed(str(job.source), key, id=job.id)
+            self._event_bus.emit_model_install_completed(str(job.source), key)
 
     def _signal_job_errored(self, job: ModelInstallJob) -> None:
         self._logger.info(f"{job.source}: model installation encountered an exception: {job.error_type}\n{job.error}")
@@ -848,7 +837,7 @@ class ModelInstallService(ModelInstallServiceBase):
             error = job.error
             assert error_type is not None
             assert error is not None
-            self._event_bus.emit_model_install_error(str(job.source), error_type, error, id=job.id)
+            self._event_bus.emit_model_install_error(str(job.source), error_type, error)
 
     def _signal_job_cancelled(self, job: ModelInstallJob) -> None:
         self._logger.info(f"{job.source}: model installation was cancelled")

invokeai/app/services/model_load/model_load_base.py

@@ -38,3 +38,8 @@ class ModelLoadServiceBase(ABC):
     @abstractmethod
     def convert_cache(self) -> ModelConvertCacheBase:
         """Return the checkpoint convert cache used by this loader."""
+
+    @property
+    @abstractmethod
+    def gpu_count(self) -> int:
+        """Return the number of GPUs we are configured to use."""

invokeai/app/services/model_load/model_load_default.py

@@ -4,6 +4,7 @@
 from typing import Optional, Type
 
 from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.shared.invocation_context import InvocationContextData
 from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType
@@ -39,6 +40,7 @@ class ModelLoadService(ModelLoadServiceBase):
         self._registry = registry
 
     def start(self, invoker: Invoker) -> None:
+        """Start the service."""
         self._invoker = invoker
 
     @property
@@ -46,6 +48,11 @@ class ModelLoadService(ModelLoadServiceBase):
         """Return the RAM cache used by this loader."""
         return self._ram_cache
 
+    @property
+    def gpu_count(self) -> int:
+        """Return the number of GPUs available for our uses."""
+        return len(self._ram_cache.execution_devices)
+
     @property
     def convert_cache(self) -> ModelConvertCacheBase:
         """Return the checkpoint convert cache used by this loader."""
@@ -94,20 +101,22 @@ class ModelLoadService(ModelLoadServiceBase):
     ) -> None:
         if not self._invoker:
             return
+        if self._invoker.services.queue.is_canceled(context_data.session_id):
+            raise CanceledException()
 
         if not loaded:
             self._invoker.services.events.emit_model_load_started(
-                queue_id=context_data.queue_item.queue_id,
-                queue_item_id=context_data.queue_item.item_id,
-                queue_batch_id=context_data.queue_item.batch_id,
-                graph_execution_state_id=context_data.queue_item.session_id,
+                queue_id=context_data.queue_id,
+                queue_item_id=context_data.queue_item_id,
+                queue_batch_id=context_data.batch_id,
+                graph_execution_state_id=context_data.session_id,
                 model_config=model_config,
             )
         else:
             self._invoker.services.events.emit_model_load_completed(
-                queue_id=context_data.queue_item.queue_id,
-                queue_item_id=context_data.queue_item.item_id,
-                queue_batch_id=context_data.queue_item.batch_id,
-                graph_execution_state_id=context_data.queue_item.session_id,
+                queue_id=context_data.queue_id,
+                queue_item_id=context_data.queue_item_id,
+                queue_batch_id=context_data.batch_id,
+                graph_execution_state_id=context_data.session_id,
                 model_config=model_config,
             )

invokeai/app/services/model_manager/model_manager_base.py

@@ -3,7 +3,6 @@
 from abc import ABC, abstractmethod
 from typing import Optional
 
-import torch
 from typing_extensions import Self
 
 from invokeai.app.services.invoker import Invoker
@@ -17,6 +16,7 @@ from ..events.events_base import EventServiceBase
 from ..model_install import ModelInstallServiceBase
 from ..model_load import ModelLoadServiceBase
 from ..model_records import ModelRecordServiceBase
+from ..shared.sqlite.sqlite_database import SqliteDatabase
 
 
 class ModelManagerServiceBase(ABC):
@@ -32,10 +32,9 @@ class ModelManagerServiceBase(ABC):
     def build_model_manager(
         cls,
         app_config: InvokeAIAppConfig,
-        model_record_service: ModelRecordServiceBase,
+        db: SqliteDatabase,
         download_queue: DownloadQueueServiceBase,
         events: EventServiceBase,
-        execution_device: torch.device,
     ) -> Self:
         """
         Construct the model manager service instance.
@@ -99,3 +98,8 @@ class ModelManagerServiceBase(ABC):
         context_data: Optional[InvocationContextData] = None,
     ) -> LoadedModel:
         pass
+
+    @property
+    @abstractmethod
+    def gpu_count(self) -> int:
+        """Return the number of GPUs we are configured to use."""

invokeai/app/services/model_manager/model_manager_default.py

@@ -3,14 +3,12 @@
 
 from typing import Optional
 
-import torch
 from typing_extensions import Self
 
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.shared.invocation_context import InvocationContextData
 from invokeai.backend.model_manager import AnyModelConfig, BaseModelType, LoadedModel, ModelType, SubModelType
 from invokeai.backend.model_manager.load import ModelCache, ModelConvertCache, ModelLoaderRegistry
-from invokeai.backend.util.devices import choose_torch_device
 from invokeai.backend.util.logging import InvokeAILogger
 
 from ..config import InvokeAIAppConfig
@@ -114,6 +112,11 @@ class ModelManagerService(ModelManagerServiceBase):
         else:
             return self.load.load_model(configs[0], submodel, context_data)
 
+    @property
+    def gpu_count(self) -> int:
+        """Return the number of GPUs we are using."""
+        return self.load.gpu_count
+
     @classmethod
     def build_model_manager(
         cls,
@@ -121,7 +124,6 @@ class ModelManagerService(ModelManagerServiceBase):
         model_record_service: ModelRecordServiceBase,
         download_queue: DownloadQueueServiceBase,
         events: EventServiceBase,
-        execution_device: torch.device = choose_torch_device(),
     ) -> Self:
         """
         Construct the model manager service instance.
@@ -132,10 +134,7 @@ class ModelManagerService(ModelManagerServiceBase):
         logger.setLevel(app_config.log_level.upper())
 
         ram_cache = ModelCache(
-            max_cache_size=app_config.ram_cache_size,
-            max_vram_cache_size=app_config.vram_cache_size,
-            logger=logger,
-            execution_device=execution_device,
+            max_cache_size=app_config.ram_cache_size, max_vram_cache_size=app_config.vram_cache_size, logger=logger
         )
         convert_cache = ModelConvertCache(
             cache_path=app_config.models_convert_cache_path, max_size=app_config.convert_cache_size

invokeai/app/services/session_processor/session_processor_common.py

@@ -4,17 +4,3 @@ from pydantic import BaseModel, Field
 class SessionProcessorStatus(BaseModel):
     is_started: bool = Field(description="Whether the session processor is started")
     is_processing: bool = Field(description="Whether a session is being processed")
-
-
-class CanceledException(Exception):
-    """Execution canceled by user."""
-
-    pass
-
-
-class ProgressImage(BaseModel):
-    """The progress image sent intermittently during processing"""
-
-    width: int = Field(description="The effective width of the image in pixels")
-    height: int = Field(description="The effective height of the image in pixels")
-    dataURL: str = Field(description="The image data as a b64 data URL")

invokeai/app/services/session_processor/session_processor_default.py

@@ -1,5 +1,4 @@
 import traceback
-from contextlib import suppress
 from threading import BoundedSemaphore, Thread
 from threading import Event as ThreadEvent
 from typing import Optional
@@ -7,270 +6,136 @@ from typing import Optional
 from fastapi_events.handlers.local import local_handler
 from fastapi_events.typing import Event as FastAPIEvent
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.events.events_base import EventServiceBase
-from invokeai.app.services.invocation_stats.invocation_stats_common import GESStatsNotFoundError
-from invokeai.app.services.session_processor.session_processor_common import CanceledException
 from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
-from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context
-from invokeai.app.util.profiler import Profiler
 
 from ..invoker import Invoker
 from .session_processor_base import SessionProcessorBase
 from .session_processor_common import SessionProcessorStatus
 
+POLLING_INTERVAL = 1
+THREAD_LIMIT = 1
+
 
 class DefaultSessionProcessor(SessionProcessorBase):
-    def start(self, invoker: Invoker, thread_limit: int = 1, polling_interval: int = 1) -> None:
-        self._invoker: Invoker = invoker
-        self._queue_item: Optional[SessionQueueItem] = None
-        self._invocation: Optional[BaseInvocation] = None
+    def start(self, invoker: Invoker) -> None:
+        self.__invoker: Invoker = invoker
+        self.__queue_item: Optional[SessionQueueItem] = None
 
-        self._resume_event = ThreadEvent()
-        self._stop_event = ThreadEvent()
-        self._poll_now_event = ThreadEvent()
-        self._cancel_event = ThreadEvent()
+        self.__resume_event = ThreadEvent()
+        self.__stop_event = ThreadEvent()
+        self.__poll_now_event = ThreadEvent()
 
         local_handler.register(event_name=EventServiceBase.queue_event, _func=self._on_queue_event)
 
-        self._thread_limit = thread_limit
-        self._thread_semaphore = BoundedSemaphore(thread_limit)
-        self._polling_interval = polling_interval
-
-        # If profiling is enabled, create a profiler. The same profiler will be used for all sessions. Internally,
-        # the profiler will create a new profile for each session.
-        self._profiler = (
-            Profiler(
-                logger=self._invoker.services.logger,
-                output_dir=self._invoker.services.configuration.profiles_path,
-                prefix=self._invoker.services.configuration.profile_prefix,
-            )
-            if self._invoker.services.configuration.profile_graphs
-            else None
-        )
-
-        self._thread = Thread(
+        self.__threadLimit = BoundedSemaphore(THREAD_LIMIT)
+        self.__thread = Thread(
             name="session_processor",
-            target=self._process,
+            target=self.__process,
             kwargs={
-                "stop_event": self._stop_event,
-                "poll_now_event": self._poll_now_event,
-                "resume_event": self._resume_event,
-                "cancel_event": self._cancel_event,
+                "stop_event": self.__stop_event,
+                "poll_now_event": self.__poll_now_event,
+                "resume_event": self.__resume_event,
             },
         )
-        self._thread.start()
+        self.__thread.start()
 
     def stop(self, *args, **kwargs) -> None:
-        self._stop_event.set()
+        self.__stop_event.set()
 
     def _poll_now(self) -> None:
-        self._poll_now_event.set()
+        self.__poll_now_event.set()
 
     async def _on_queue_event(self, event: FastAPIEvent) -> None:
         event_name = event[1]["event"]
 
-        if event_name == "session_canceled" or event_name == "queue_cleared":
-            # These both mean we should cancel the current session.
-            self._cancel_event.set()
+        # This was a match statement, but match is not supported on python 3.9
+        if event_name in [
+            "graph_execution_state_complete",
+            "invocation_error",
+            "session_retrieval_error",
+            "invocation_retrieval_error",
+        ]:
+            self.__queue_item = None
+            self._poll_now()
+        elif (
+            event_name == "session_canceled"
+            and self.__queue_item is not None
+            and self.__queue_item.session_id == event[1]["data"]["graph_execution_state_id"]
+        ):
+            self.__queue_item = None
             self._poll_now()
         elif event_name == "batch_enqueued":
             self._poll_now()
+        elif event_name == "queue_cleared":
+            self.__queue_item = None
+            self._poll_now()
 
     def resume(self) -> SessionProcessorStatus:
-        if not self._resume_event.is_set():
-            self._resume_event.set()
+        if not self.__resume_event.is_set():
+            self.__resume_event.set()
         return self.get_status()
 
     def pause(self) -> SessionProcessorStatus:
-        if self._resume_event.is_set():
-            self._resume_event.clear()
+        if self.__resume_event.is_set():
+            self.__resume_event.clear()
         return self.get_status()
 
     def get_status(self) -> SessionProcessorStatus:
         return SessionProcessorStatus(
-            is_started=self._resume_event.is_set(),
-            is_processing=self._queue_item is not None,
+            is_started=self.__resume_event.is_set(),
+            is_processing=self.__queue_item is not None,
         )
 
-    def _process(
+    def __process(
         self,
         stop_event: ThreadEvent,
         poll_now_event: ThreadEvent,
         resume_event: ThreadEvent,
-        cancel_event: ThreadEvent,
     ):
-        # Outermost processor try block; any unhandled exception is a fatal processor error
         try:
-            self._thread_semaphore.acquire()
             stop_event.clear()
             resume_event.set()
-            cancel_event.clear()
-
+            self.__threadLimit.acquire()
+            queue_item: Optional[SessionQueueItem] = None
             while not stop_event.is_set():
                 poll_now_event.clear()
-                # Middle processor try block; any unhandled exception is a non-fatal processor error
                 try:
-                    # Get the next session to process
-                    self._queue_item = self._invoker.services.session_queue.dequeue()
-                    if self._queue_item is not None and resume_event.is_set():
-                        self._invoker.services.logger.debug(f"Executing queue item {self._queue_item.item_id}")
-                        cancel_event.clear()
+                    # do not dequeue if there is already a session running
+                    if self.__queue_item is None and resume_event.is_set():
+                        queue_item = self.__invoker.services.session_queue.dequeue()
 
-                        # If profiling is enabled, start the profiler
-                        if self._profiler is not None:
-                            self._profiler.start(profile_id=self._queue_item.session_id)
-
-                        # Prepare invocations and take the first
-                        self._invocation = self._queue_item.session.next()
-
-                        # Loop over invocations until the session is complete or canceled
-                        while self._invocation is not None and not cancel_event.is_set():
-                            # get the source node id to provide to clients (the prepared node id is not as useful)
-                            source_invocation_id = self._queue_item.session.prepared_source_mapping[self._invocation.id]
-
-                            # Send starting event
-                            self._invoker.services.events.emit_invocation_started(
-                                queue_batch_id=self._queue_item.batch_id,
-                                queue_item_id=self._queue_item.item_id,
-                                queue_id=self._queue_item.queue_id,
-                                graph_execution_state_id=self._queue_item.session_id,
-                                node=self._invocation.model_dump(),
-                                source_node_id=source_invocation_id,
+                        if queue_item is not None:
+                            self.__invoker.services.logger.debug(f"Executing queue item {queue_item.item_id}")
+                            self.__queue_item = queue_item
+                            self.__invoker.services.graph_execution_manager.set(queue_item.session)
+                            self.__invoker.invoke(
+                                session_queue_batch_id=queue_item.batch_id,
+                                session_queue_id=queue_item.queue_id,
+                                session_queue_item_id=queue_item.item_id,
+                                graph_execution_state=queue_item.session,
+                                workflow=queue_item.workflow,
+                                invoke_all=True,
                             )
+                            queue_item = None
 
-                            # Innermost processor try block; any unhandled exception is an invocation error & will fail the graph
-                            try:
-                                with self._invoker.services.performance_statistics.collect_stats(
-                                    self._invocation, self._queue_item.session.id
-                                ):
-                                    # Build invocation context (the node-facing API)
-                                    data = InvocationContextData(
-                                        invocation=self._invocation,
-                                        source_invocation_id=source_invocation_id,
-                                        queue_item=self._queue_item,
-                                    )
-                                    context = build_invocation_context(
-                                        data=data,
-                                        services=self._invoker.services,
-                                        cancel_event=self._cancel_event,
-                                    )
-
-                                    # Invoke the node
-                                    outputs = self._invocation.invoke_internal(
-                                        context=context, services=self._invoker.services
-                                    )
-
-                                    # Save outputs and history
-                                    self._queue_item.session.complete(self._invocation.id, outputs)
-
-                                    # Send complete event
-                                    self._invoker.services.events.emit_invocation_complete(
-                                        queue_batch_id=self._queue_item.batch_id,
-                                        queue_item_id=self._queue_item.item_id,
-                                        queue_id=self._queue_item.queue_id,
-                                        graph_execution_state_id=self._queue_item.session.id,
-                                        node=self._invocation.model_dump(),
-                                        source_node_id=source_invocation_id,
-                                        result=outputs.model_dump(),
-                                    )
-
-                            except KeyboardInterrupt:
-                                # TODO(MM2): Create an event for this
-                                pass
-
-                            except CanceledException:
-                                # When the user cancels the graph, we first set the cancel event. The event is checked
-                                # between invocations, in this loop. Some invocations are long-running, and we need to
-                                # be able to cancel them mid-execution.
-                                #
-                                # For example, denoising is a long-running invocation with many steps. A step callback
-                                # is executed after each step. This step callback checks if the canceled event is set,
-                                # then raises a CanceledException to stop execution immediately.
-                                #
-                                # When we get a CanceledException, we don't need to do anything - just pass and let the
-                                # loop go to its next iteration, and the cancel event will be handled correctly.
-                                pass
-
-                            except Exception as e:
-                                error = traceback.format_exc()
-
-                                # Save error
-                                self._queue_item.session.set_node_error(self._invocation.id, error)
-                                self._invoker.services.logger.error(
-                                    f"Error while invoking session {self._queue_item.session_id}, invocation {self._invocation.id} ({self._invocation.get_type()}):\n{e}"
-                                )
-
-                                # Send error event
-                                self._invoker.services.events.emit_invocation_error(
-                                    queue_batch_id=self._queue_item.session_id,
-                                    queue_item_id=self._queue_item.item_id,
-                                    queue_id=self._queue_item.queue_id,
-                                    graph_execution_state_id=self._queue_item.session.id,
-                                    node=self._invocation.model_dump(),
-                                    source_node_id=source_invocation_id,
-                                    error_type=e.__class__.__name__,
-                                    error=error,
-                                )
-                                pass
-
-                            # The session is complete if the all invocations are complete or there was an error
-                            if self._queue_item.session.is_complete() or cancel_event.is_set():
-                                # Send complete event
-                                self._invoker.services.events.emit_graph_execution_complete(
-                                    queue_batch_id=self._queue_item.batch_id,
-                                    queue_item_id=self._queue_item.item_id,
-                                    queue_id=self._queue_item.queue_id,
-                                    graph_execution_state_id=self._queue_item.session.id,
-                                )
-                                # If we are profiling, stop the profiler and dump the profile & stats
-                                if self._profiler:
-                                    profile_path = self._profiler.stop()
-                                    stats_path = profile_path.with_suffix(".json")
-                                    self._invoker.services.performance_statistics.dump_stats(
-                                        graph_execution_state_id=self._queue_item.session.id, output_path=stats_path
-                                    )
-                                # We'll get a GESStatsNotFoundError if we try to log stats for an untracked graph, but in the processor
-                                # we don't care about that - suppress the error.
-                                with suppress(GESStatsNotFoundError):
-                                    self._invoker.services.performance_statistics.log_stats(self._queue_item.session.id)
-                                    self._invoker.services.performance_statistics.reset_stats()
-
-                                # Set the invocation to None to prepare for the next session
-                                self._invocation = None
-                            else:
-                                # Prepare the next invocation
-                                self._invocation = self._queue_item.session.next()
-
-                        # The session is complete, immediately poll for next session
-                        self._queue_item = None
-                        poll_now_event.set()
-                    else:
-                        # The queue was empty, wait for next polling interval or event to try again
-                        self._invoker.services.logger.debug("Waiting for next polling interval or event")
-                        poll_now_event.wait(self._polling_interval)
+                    if queue_item is None:
+                        self.__invoker.services.logger.debug("Waiting for next polling interval or event")
+                        poll_now_event.wait(POLLING_INTERVAL)
                         continue
-                except Exception:
-                    # Non-fatal error in processor
-                    self._invoker.services.logger.error(
-                        f"Non-fatal error in session processor:\n{traceback.format_exc()}"
-                    )
-                    # Cancel the queue item
-                    if self._queue_item is not None:
-                        self._invoker.services.session_queue.cancel_queue_item(
-                            self._queue_item.item_id, error=traceback.format_exc()
+                except Exception as e:
+                    self.__invoker.services.logger.error(f"Error in session processor: {e}")
+                    if queue_item is not None:
+                        self.__invoker.services.session_queue.cancel_queue_item(
+                            queue_item.item_id, error=traceback.format_exc()
                         )
-                    # Reset the invocation to None to prepare for the next session
-                    self._invocation = None
-                    # Immediately poll for next queue item
-                    poll_now_event.wait(self._polling_interval)
+                    poll_now_event.wait(POLLING_INTERVAL)
                     continue
-        except Exception:
-            # Fatal error in processor, log and pass - we're done here
-            self._invoker.services.logger.error(f"Fatal Error in session processor:\n{traceback.format_exc()}")
+        except Exception as e:
+            self.__invoker.services.logger.error(f"Fatal Error in session processor: {e}")
             pass
         finally:
             stop_event.clear()
             poll_now_event.clear()
-            self._queue_item = None
-            self._thread_semaphore.release()
+            self.__queue_item = None
+            self.__threadLimit.release()

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -60,7 +60,7 @@ class SqliteSessionQueue(SessionQueueBase):
         # This was a match statement, but match is not supported on python 3.9
         if event_name == "graph_execution_state_complete":
             await self._handle_complete_event(event)
-        elif event_name == "invocation_error":
+        elif event_name in ["invocation_error", "session_retrieval_error", "invocation_retrieval_error"]:
             await self._handle_error_event(event)
         elif event_name == "session_canceled":
             await self._handle_cancel_event(event)
@@ -429,6 +429,7 @@ class SqliteSessionQueue(SessionQueueBase):
         if queue_item.status not in ["canceled", "failed", "completed"]:
             status = "failed" if error is not None else "canceled"
             queue_item = self._set_queue_item_status(item_id=item_id, status=status, error=error)  # type: ignore [arg-type] # mypy seems to not narrow the Literals here
+            self.__invoker.services.queue.cancel(queue_item.session_id)
             self.__invoker.services.events.emit_session_canceled(
                 queue_item_id=queue_item.item_id,
                 queue_id=queue_item.queue_id,
@@ -470,6 +471,7 @@ class SqliteSessionQueue(SessionQueueBase):
             )
             self.__conn.commit()
             if current_queue_item is not None and current_queue_item.batch_id in batch_ids:
+                self.__invoker.services.queue.cancel(current_queue_item.session_id)
                 self.__invoker.services.events.emit_session_canceled(
                     queue_item_id=current_queue_item.item_id,
                     queue_id=current_queue_item.queue_id,
@@ -521,6 +523,7 @@ class SqliteSessionQueue(SessionQueueBase):
             )
             self.__conn.commit()
             if current_queue_item is not None and current_queue_item.queue_id == queue_id:
+                self.__invoker.services.queue.cancel(current_queue_item.session_id)
                 self.__invoker.services.events.emit_session_canceled(
                     queue_item_id=current_queue_item.item_id,
                     queue_id=current_queue_item.queue_id,

invokeai/app/services/shared/default_graphs.py

@@ -0,0 +1,92 @@
+from invokeai.app.services.item_storage.item_storage_base import ItemStorageABC
+
+from ...invocations.compel import CompelInvocation
+from ...invocations.image import ImageNSFWBlurInvocation
+from ...invocations.latent import DenoiseLatentsInvocation, LatentsToImageInvocation
+from ...invocations.noise import NoiseInvocation
+from ...invocations.primitives import IntegerInvocation
+from .graph import Edge, EdgeConnection, ExposedNodeInput, ExposedNodeOutput, Graph, LibraryGraph
+
+default_text_to_image_graph_id = "539b2af5-2b4d-4d8c-8071-e54a3255fc74"
+
+
+def create_text_to_image() -> LibraryGraph:
+    graph = Graph(
+        nodes={
+            "width": IntegerInvocation(id="width", value=512),
+            "height": IntegerInvocation(id="height", value=512),
+            "seed": IntegerInvocation(id="seed", value=-1),
+            "3": NoiseInvocation(id="3"),
+            "4": CompelInvocation(id="4"),
+            "5": CompelInvocation(id="5"),
+            "6": DenoiseLatentsInvocation(id="6"),
+            "7": LatentsToImageInvocation(id="7"),
+            "8": ImageNSFWBlurInvocation(id="8"),
+        },
+        edges=[
+            Edge(
+                source=EdgeConnection(node_id="width", field="value"),
+                destination=EdgeConnection(node_id="3", field="width"),
+            ),
+            Edge(
+                source=EdgeConnection(node_id="height", field="value"),
+                destination=EdgeConnection(node_id="3", field="height"),
+            ),
+            Edge(
+                source=EdgeConnection(node_id="seed", field="value"),
+                destination=EdgeConnection(node_id="3", field="seed"),
+            ),
+            Edge(
+                source=EdgeConnection(node_id="3", field="noise"),
+                destination=EdgeConnection(node_id="6", field="noise"),
+            ),
+            Edge(
+                source=EdgeConnection(node_id="6", field="latents"),
+                destination=EdgeConnection(node_id="7", field="latents"),
+            ),
+            Edge(
+                source=EdgeConnection(node_id="4", field="conditioning"),
+                destination=EdgeConnection(node_id="6", field="positive_conditioning"),
+            ),
+            Edge(
+                source=EdgeConnection(node_id="5", field="conditioning"),
+                destination=EdgeConnection(node_id="6", field="negative_conditioning"),
+            ),
+            Edge(
+                source=EdgeConnection(node_id="7", field="image"),
+                destination=EdgeConnection(node_id="8", field="image"),
+            ),
+        ],
+    )
+    return LibraryGraph(
+        id=default_text_to_image_graph_id,
+        name="t2i",
+        description="Converts text to an image",
+        graph=graph,
+        exposed_inputs=[
+            ExposedNodeInput(node_path="4", field="prompt", alias="positive_prompt"),
+            ExposedNodeInput(node_path="5", field="prompt", alias="negative_prompt"),
+            ExposedNodeInput(node_path="width", field="value", alias="width"),
+            ExposedNodeInput(node_path="height", field="value", alias="height"),
+            ExposedNodeInput(node_path="seed", field="value", alias="seed"),
+        ],
+        exposed_outputs=[ExposedNodeOutput(node_path="8", field="image", alias="image")],
+    )
+
+
+def create_system_graphs(graph_library: ItemStorageABC[LibraryGraph]) -> list[LibraryGraph]:
+    """Creates the default system graphs, or adds new versions if the old ones don't match"""
+
+    # TODO: Uncomment this when we are ready to fix this up to prevent breaking changes
+    graphs: list[LibraryGraph] = []
+
+    text_to_image = graph_library.get(default_text_to_image_graph_id)
+
+    # TODO: Check if the graph is the same as the default one, and if not, update it
+    # if text_to_image is None:
+    text_to_image = create_text_to_image()
+    graph_library.set(text_to_image)
+
+    graphs.append(text_to_image)
+
+    return graphs

invokeai/app/services/shared/graph.py

@@ -5,14 +5,8 @@ import itertools
 from typing import Annotated, Any, Optional, TypeVar, Union, get_args, get_origin, get_type_hints
 
 import networkx as nx
-from pydantic import (
-    BaseModel,
-    GetJsonSchemaHandler,
-    field_validator,
-)
+from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 from pydantic.fields import Field
-from pydantic.json_schema import JsonSchemaValue
-from pydantic_core import CoreSchema
 
 # Importing * is bad karma but needed here for node detection
 from invokeai.app.invocations import *  # noqa: F401 F403
@@ -182,6 +176,10 @@ class NodeIdMismatchError(ValueError):
     pass
 
 
+class InvalidSubGraphError(ValueError):
+    pass
+
+
 class CyclicalGraphError(ValueError):
     pass
 
@@ -190,6 +188,25 @@ class UnknownGraphValidationError(ValueError):
     pass
 
 
+# TODO: Create and use an Empty output?
+@invocation_output("graph_output")
+class GraphInvocationOutput(BaseInvocationOutput):
+    pass
+
+
+# TODO: Fill this out and move to invocations
+@invocation("graph", version="1.0.0")
+class GraphInvocation(BaseInvocation):
+    """Execute a graph"""
+
+    # TODO: figure out how to create a default here
+    graph: "Graph" = InputField(description="The graph to run", default=None)
+
+    def invoke(self, context: InvocationContext) -> GraphInvocationOutput:
+        """Invoke with provided services and return outputs."""
+        return GraphInvocationOutput()
+
+
 @invocation_output("iterate_output")
 class IterateInvocationOutput(BaseInvocationOutput):
     """Used to connect iteration outputs. Will be expanded to a specific output."""
@@ -243,73 +260,21 @@ class CollectInvocation(BaseInvocation):
         return CollectInvocationOutput(collection=copy.copy(self.collection))
 
 
+InvocationsUnion: Any = BaseInvocation.get_invocations_union()
+InvocationOutputsUnion: Any = BaseInvocationOutput.get_outputs_union()
+
+
 class Graph(BaseModel):
     id: str = Field(description="The id of this graph", default_factory=uuid_string)
     # TODO: use a list (and never use dict in a BaseModel) because pydantic/fastapi hates me
-    nodes: dict[str, BaseInvocation] = Field(description="The nodes in this graph", default_factory=dict)
+    nodes: dict[str, Annotated[InvocationsUnion, Field(discriminator="type")]] = Field(
+        description="The nodes in this graph", default_factory=dict
+    )
     edges: list[Edge] = Field(
         description="The connections between nodes and their fields in this graph",
         default_factory=list,
     )
 
-    @field_validator("nodes", mode="plain")
-    @classmethod
-    def validate_nodes(cls, v: dict[str, Any]):
-        """Validates the nodes in the graph by retrieving a union of all node types and validating each node."""
-
-        # Invocations register themselves as their python modules are executed. The union of all invocations is
-        # constructed at runtime. We use pydantic to validate `Graph.nodes` using that union.
-        #
-        # It's possible that when `graph.py` is executed, not all invocation-containing modules will have executed. If
-        # we construct the invocation union as `graph.py` is executed, we may miss some invocations. Those missing
-        # invocations will cause a graph to fail if they are used.
-        #
-        # We can get around this by validating the nodes in the graph using a "plain" validator, which overrides the
-        # pydantic validation entirely. This allows us to validate the nodes using the union of invocations at runtime.
-        #
-        # This same pattern is used in `GraphExecutionState`.
-
-        nodes: dict[str, BaseInvocation] = {}
-        typeadapter = BaseInvocation.get_typeadapter()
-        for node_id, node in v.items():
-            nodes[node_id] = typeadapter.validate_python(node)
-        return nodes
-
-    @classmethod
-    def __get_pydantic_json_schema__(cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
-        # We use a "plain" validator to validate the nodes in the graph. Pydantic is unable to create a JSON Schema for
-        # fields that use "plain" validators, so we have to hack around this. Also, we need to add all invocations to
-        # the generated schema as options for the `nodes` field.
-        #
-        # The workaround is to create a new BaseModel that has the same fields as `Graph` but without the validator and
-        # with the invocation union as the type for the `nodes` field. Pydantic then generates the JSON Schema as
-        # expected.
-        #
-        # You might be tempted to do something like this:
-        #
-        # ```py
-        # cloned_model = create_model(cls.__name__, __base__=cls, nodes=...)
-        # delattr(cloned_model, "validate_nodes")
-        # cloned_model.model_rebuild(force=True)
-        # json_schema = handler(cloned_model.__pydantic_core_schema__)
-        # ```
-        #
-        # Unfortunately, this does not work. Calling `handler` here results in infinite recursion as pydantic attempts
-        # to build the JSON Schema for the cloned model. Instead, we have to manually clone the model.
-        #
-        # This same pattern is used in `GraphExecutionState`.
-
-        class Graph(BaseModel):
-            id: Optional[str] = Field(default=None, description="The id of this graph")
-            nodes: dict[
-                str, Annotated[Union[tuple(BaseInvocation._invocation_classes)], Field(discriminator="type")]
-            ] = Field(description="The nodes in this graph")
-            edges: list[Edge] = Field(description="The connections between nodes and their fields in this graph")
-
-        json_schema = handler(Graph.__pydantic_core_schema__)
-        json_schema = handler.resolve_ref_schema(json_schema)
-        return json_schema
-
     def add_node(self, node: BaseInvocation) -> None:
         """Adds a node to a graph
 
@@ -321,21 +286,41 @@ class Graph(BaseModel):
 
         self.nodes[node.id] = node
 
-    def delete_node(self, node_id: str) -> None:
+    def _get_graph_and_node(self, node_path: str) -> tuple["Graph", str]:
+        """Returns the graph and node id for a node path."""
+        # Materialized graphs may have nodes at the top level
+        if node_path in self.nodes:
+            return (self, node_path)
+
+        node_id = node_path if "." not in node_path else node_path[: node_path.index(".")]
+        if node_id not in self.nodes:
+            raise NodeNotFoundError(f"Node {node_path} not found in graph")
+
+        node = self.nodes[node_id]
+
+        if not isinstance(node, GraphInvocation):
+            # There's more node path left but this isn't a graph - failure
+            raise NodeNotFoundError("Node path terminated early at a non-graph node")
+
+        return node.graph._get_graph_and_node(node_path[node_path.index(".") + 1 :])
+
+    def delete_node(self, node_path: str) -> None:
         """Deletes a node from a graph"""
 
         try:
+            graph, node_id = self._get_graph_and_node(node_path)
+
             # Delete edges for this node
-            input_edges = self._get_input_edges(node_id)
-            output_edges = self._get_output_edges(node_id)
+            input_edges = self._get_input_edges_and_graphs(node_path)
+            output_edges = self._get_output_edges_and_graphs(node_path)
 
-            for edge in input_edges:
-                self.delete_edge(edge)
+            for edge_graph, _, edge in input_edges:
+                edge_graph.delete_edge(edge)
 
-            for edge in output_edges:
-                self.delete_edge(edge)
+            for edge_graph, _, edge in output_edges:
+                edge_graph.delete_edge(edge)
 
-            del self.nodes[node_id]
+            del graph.nodes[node_id]
 
         except NodeNotFoundError:
             pass  # Ignore, not doesn't exist (should this throw?)
@@ -385,6 +370,13 @@ class Graph(BaseModel):
             if k != v.id:
                 raise NodeIdMismatchError(f"Node ids must match, got {k} and {v.id}")
 
+        # Validate all subgraphs
+        for gn in (n for n in self.nodes.values() if isinstance(n, GraphInvocation)):
+            try:
+                gn.graph.validate_self()
+            except Exception as e:
+                raise InvalidSubGraphError(f"Subgraph {gn.id} is invalid") from e
+
         # Validate that all edges match nodes and fields in the graph
         for edge in self.edges:
             source_node = self.nodes.get(edge.source.node_id, None)
@@ -446,6 +438,7 @@ class Graph(BaseModel):
         except (
             DuplicateNodeIdError,
             NodeIdMismatchError,
+            InvalidSubGraphError,
             NodeNotFoundError,
             NodeFieldNotFoundError,
             CyclicalGraphError,
@@ -466,7 +459,7 @@ class Graph(BaseModel):
     def _validate_edge(self, edge: Edge):
         """Validates that a new edge doesn't create a cycle in the graph"""
 
-        # Validate that the nodes exist
+        # Validate that the nodes exist (edges may contain node paths, so we can't just check for nodes directly)
         try:
             from_node = self.get_node(edge.source.node_id)
             to_node = self.get_node(edge.destination.node_id)
@@ -533,90 +526,171 @@ class Graph(BaseModel):
                     f"Collector input type does not match collector output type: {edge.source.node_id}.{edge.source.field} to {edge.destination.node_id}.{edge.destination.field}"
                 )
 
-    def has_node(self, node_id: str) -> bool:
+    def has_node(self, node_path: str) -> bool:
         """Determines whether or not a node exists in the graph."""
         try:
-            _ = self.get_node(node_id)
-            return True
+            n = self.get_node(node_path)
+            if n is not None:
+                return True
+            else:
+                return False
         except NodeNotFoundError:
             return False
 
-    def get_node(self, node_id: str) -> BaseInvocation:
-        """Gets a node from the graph."""
-        try:
-            return self.nodes[node_id]
-        except KeyError as e:
-            raise NodeNotFoundError(f"Node {node_id} not found in graph") from e
+    def get_node(self, node_path: str) -> BaseInvocation:
+        """Gets a node from the graph using a node path."""
+        # Materialized graphs may have nodes at the top level
+        graph, node_id = self._get_graph_and_node(node_path)
+        return graph.nodes[node_id]
 
-    def update_node(self, node_id: str, new_node: BaseInvocation) -> None:
+    def _get_node_path(self, node_id: str, prefix: Optional[str] = None) -> str:
+        return node_id if prefix is None or prefix == "" else f"{prefix}.{node_id}"
+
+    def update_node(self, node_path: str, new_node: BaseInvocation) -> None:
         """Updates a node in the graph."""
-        node = self.nodes[node_id]
+        graph, node_id = self._get_graph_and_node(node_path)
+        node = graph.nodes[node_id]
 
         # Ensure the node type matches the new node
         if type(node) is not type(new_node):
-            raise TypeError(f"Node {node_id} is type {type(node)} but new node is type {type(new_node)}")
+            raise TypeError(f"Node {node_path} is type {type(node)} but new node is type {type(new_node)}")
 
         # Ensure the new id is either the same or is not in the graph
-        if new_node.id != node.id and self.has_node(new_node.id):
-            raise NodeAlreadyInGraphError(f"Node with id {new_node.id} already exists in graph")
+        prefix = None if "." not in node_path else node_path[: node_path.rindex(".")]
+        new_path = self._get_node_path(new_node.id, prefix=prefix)
+        if new_node.id != node.id and self.has_node(new_path):
+            raise NodeAlreadyInGraphError("Node with id {new_node.id} already exists in graph")
 
         # Set the new node in the graph
-        self.nodes[new_node.id] = new_node
+        graph.nodes[new_node.id] = new_node
         if new_node.id != node.id:
-            input_edges = self._get_input_edges(node_id)
-            output_edges = self._get_output_edges(node_id)
+            input_edges = self._get_input_edges_and_graphs(node_path)
+            output_edges = self._get_output_edges_and_graphs(node_path)
 
             # Delete node and all edges
-            self.delete_node(node_id)
+            graph.delete_node(node_path)
 
             # Create new edges for each input and output
-            for edge in input_edges:
-                self.add_edge(
+            for graph, _, edge in input_edges:
+                # Remove the graph prefix from the node path
+                new_graph_node_path = (
+                    new_node.id
+                    if "." not in edge.destination.node_id
+                    else f'{edge.destination.node_id[edge.destination.node_id.rindex("."):]}.{new_node.id}'
+                )
+                graph.add_edge(
                     Edge(
                         source=edge.source,
-                        destination=EdgeConnection(node_id=new_node.id, field=edge.destination.field),
+                        destination=EdgeConnection(node_id=new_graph_node_path, field=edge.destination.field),
                     )
                 )
 
-            for edge in output_edges:
-                self.add_edge(
+            for graph, _, edge in output_edges:
+                # Remove the graph prefix from the node path
+                new_graph_node_path = (
+                    new_node.id
+                    if "." not in edge.source.node_id
+                    else f'{edge.source.node_id[edge.source.node_id.rindex("."):]}.{new_node.id}'
+                )
+                graph.add_edge(
                     Edge(
-                        source=EdgeConnection(node_id=new_node.id, field=edge.source.field),
+                        source=EdgeConnection(node_id=new_graph_node_path, field=edge.source.field),
                         destination=edge.destination,
                     )
                 )
 
-    def _get_input_edges(self, node_id: str, field: Optional[str] = None) -> list[Edge]:
-        """Gets all input edges for a node. If field is provided, only edges to that field are returned."""
+    def _get_input_edges(self, node_path: str, field: Optional[str] = None) -> list[Edge]:
+        """Gets all input edges for a node"""
+        edges = self._get_input_edges_and_graphs(node_path)
 
-        edges = [e for e in self.edges if e.destination.node_id == node_id]
+        # Filter to edges that match the field
+        filtered_edges = (e for e in edges if field is None or e[2].destination.field == field)
 
-        if field is None:
-            return edges
+        # Create full node paths for each edge
+        return [
+            Edge(
+                source=EdgeConnection(
+                    node_id=self._get_node_path(e.source.node_id, prefix=prefix),
+                    field=e.source.field,
+                ),
+                destination=EdgeConnection(
+                    node_id=self._get_node_path(e.destination.node_id, prefix=prefix),
+                    field=e.destination.field,
+                ),
+            )
+            for _, prefix, e in filtered_edges
+        ]
 
-        filtered_edges = [e for e in edges if e.destination.field == field]
+    def _get_input_edges_and_graphs(
+        self, node_path: str, prefix: Optional[str] = None
+    ) -> list[tuple["Graph", Union[str, None], Edge]]:
+        """Gets all input edges for a node along with the graph they are in and the graph's path"""
+        edges = []
 
-        return filtered_edges
+        # Return any input edges that appear in this graph
+        edges.extend([(self, prefix, e) for e in self.edges if e.destination.node_id == node_path])
 
-    def _get_output_edges(self, node_id: str, field: Optional[str] = None) -> list[Edge]:
-        """Gets all output edges for a node. If field is provided, only edges from that field are returned."""
-        edges = [e for e in self.edges if e.source.node_id == node_id]
+        node_id = node_path if "." not in node_path else node_path[: node_path.index(".")]
+        node = self.nodes[node_id]
 
-        if field is None:
-            return edges
+        if isinstance(node, GraphInvocation):
+            graph = node.graph
+            graph_path = node.id if prefix is None or prefix == "" else self._get_node_path(node.id, prefix=prefix)
+            graph_edges = graph._get_input_edges_and_graphs(node_path[(len(node_id) + 1) :], prefix=graph_path)
+            edges.extend(graph_edges)
 
-        filtered_edges = [e for e in edges if e.source.field == field]
+        return edges
 
-        return filtered_edges
+    def _get_output_edges(self, node_path: str, field: str) -> list[Edge]:
+        """Gets all output edges for a node"""
+        edges = self._get_output_edges_and_graphs(node_path)
+
+        # Filter to edges that match the field
+        filtered_edges = (e for e in edges if e[2].source.field == field)
+
+        # Create full node paths for each edge
+        return [
+            Edge(
+                source=EdgeConnection(
+                    node_id=self._get_node_path(e.source.node_id, prefix=prefix),
+                    field=e.source.field,
+                ),
+                destination=EdgeConnection(
+                    node_id=self._get_node_path(e.destination.node_id, prefix=prefix),
+                    field=e.destination.field,
+                ),
+            )
+            for _, prefix, e in filtered_edges
+        ]
+
+    def _get_output_edges_and_graphs(
+        self, node_path: str, prefix: Optional[str] = None
+    ) -> list[tuple["Graph", Union[str, None], Edge]]:
+        """Gets all output edges for a node along with the graph they are in and the graph's path"""
+        edges = []
+
+        # Return any input edges that appear in this graph
+        edges.extend([(self, prefix, e) for e in self.edges if e.source.node_id == node_path])
+
+        node_id = node_path if "." not in node_path else node_path[: node_path.index(".")]
+        node = self.nodes[node_id]
+
+        if isinstance(node, GraphInvocation):
+            graph = node.graph
+            graph_path = node.id if prefix is None or prefix == "" else self._get_node_path(node.id, prefix=prefix)
+            graph_edges = graph._get_output_edges_and_graphs(node_path[(len(node_id) + 1) :], prefix=graph_path)
+            edges.extend(graph_edges)
+
+        return edges
 
     def _is_iterator_connection_valid(
         self,
-        node_id: str,
+        node_path: str,
         new_input: Optional[EdgeConnection] = None,
         new_output: Optional[EdgeConnection] = None,
     ) -> bool:
-        inputs = [e.source for e in self._get_input_edges(node_id, "collection")]
-        outputs = [e.destination for e in self._get_output_edges(node_id, "item")]
+        inputs = [e.source for e in self._get_input_edges(node_path, "collection")]
+        outputs = [e.destination for e in self._get_output_edges(node_path, "item")]
 
         if new_input is not None:
             inputs.append(new_input)
@@ -644,12 +718,12 @@ class Graph(BaseModel):
 
     def _is_collector_connection_valid(
         self,
-        node_id: str,
+        node_path: str,
         new_input: Optional[EdgeConnection] = None,
         new_output: Optional[EdgeConnection] = None,
     ) -> bool:
-        inputs = [e.source for e in self._get_input_edges(node_id, "item")]
-        outputs = [e.destination for e in self._get_output_edges(node_id, "collection")]
+        inputs = [e.source for e in self._get_input_edges(node_path, "item")]
+        outputs = [e.destination for e in self._get_output_edges(node_path, "collection")]
 
         if new_input is not None:
             inputs.append(new_input)
@@ -705,17 +779,27 @@ class Graph(BaseModel):
         g.add_edges_from({(e.source.node_id, e.destination.node_id) for e in self.edges})
         return g
 
-    def nx_graph_flat(self, nx_graph: Optional[nx.DiGraph] = None) -> nx.DiGraph:
+    def nx_graph_flat(self, nx_graph: Optional[nx.DiGraph] = None, prefix: Optional[str] = None) -> nx.DiGraph:
         """Returns a flattened NetworkX DiGraph, including all subgraphs (but not with iterations expanded)"""
         g = nx_graph or nx.DiGraph()
 
         # Add all nodes from this graph except graph/iteration nodes
-        g.add_nodes_from([n.id for n in self.nodes.values() if not isinstance(n, IterateInvocation)])
+        g.add_nodes_from(
+            [
+                self._get_node_path(n.id, prefix)
+                for n in self.nodes.values()
+                if not isinstance(n, GraphInvocation) and not isinstance(n, IterateInvocation)
+            ]
+        )
+
+        # Expand graph nodes
+        for sgn in (gn for gn in self.nodes.values() if isinstance(gn, GraphInvocation)):
+            g = sgn.graph.nx_graph_flat(g, self._get_node_path(sgn.id, prefix))
 
         # TODO: figure out if iteration nodes need to be expanded
 
         unique_edges = {(e.source.node_id, e.destination.node_id) for e in self.edges}
-        g.add_edges_from([(e[0], e[1]) for e in unique_edges])
+        g.add_edges_from([(self._get_node_path(e[0], prefix), self._get_node_path(e[1], prefix)) for e in unique_edges])
         return g
 
 
@@ -740,7 +824,9 @@ class GraphExecutionState(BaseModel):
     )
 
     # The results of executed nodes
-    results: dict[str, BaseInvocationOutput] = Field(description="The results of node executions", default_factory=dict)
+    results: dict[str, Annotated[InvocationOutputsUnion, Field(discriminator="type")]] = Field(
+        description="The results of node executions", default_factory=dict
+    )
 
     # Errors raised when executing nodes
     errors: dict[str, str] = Field(description="Errors raised when executing nodes", default_factory=dict)
@@ -757,51 +843,27 @@ class GraphExecutionState(BaseModel):
         default_factory=dict,
     )
 
-    @field_validator("results", mode="plain")
-    @classmethod
-    def validate_results(cls, v: dict[str, BaseInvocationOutput]):
-        """Validates the results in the GES by retrieving a union of all output types and validating each result."""
-
-        # See the comment in `Graph.validate_nodes` for an explanation of this logic.
-        results: dict[str, BaseInvocationOutput] = {}
-        typeadapter = BaseInvocationOutput.get_typeadapter()
-        for result_id, result in v.items():
-            results[result_id] = typeadapter.validate_python(result)
-        return results
-
     @field_validator("graph")
     def graph_is_valid(cls, v: Graph):
         """Validates that the graph is valid"""
         v.validate_self()
         return v
 
-    @classmethod
-    def __get_pydantic_json_schema__(cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
-        # See the comment in `Graph.__get_pydantic_json_schema__` for an explanation of this logic.
-        class GraphExecutionState(BaseModel):
-            """Tracks the state of a graph execution"""
-
-            id: str = Field(description="The id of the execution state")
-            graph: Graph = Field(description="The graph being executed")
-            execution_graph: Graph = Field(description="The expanded graph of activated and executed nodes")
-            executed: set[str] = Field(description="The set of node ids that have been executed")
-            executed_history: list[str] = Field(
-                description="The list of node ids that have been executed, in order of execution"
-            )
-            results: dict[
-                str, Annotated[Union[tuple(BaseInvocationOutput._output_classes)], Field(discriminator="type")]
-            ] = Field(description="The results of node executions")
-            errors: dict[str, str] = Field(description="Errors raised when executing nodes")
-            prepared_source_mapping: dict[str, str] = Field(
-                description="The map of prepared nodes to original graph nodes"
-            )
-            source_prepared_mapping: dict[str, set[str]] = Field(
-                description="The map of original graph nodes to prepared nodes"
-            )
-
-        json_schema = handler(GraphExecutionState.__pydantic_core_schema__)
-        json_schema = handler.resolve_ref_schema(json_schema)
-        return json_schema
+    model_config = ConfigDict(
+        json_schema_extra={
+            "required": [
+                "id",
+                "graph",
+                "execution_graph",
+                "executed",
+                "executed_history",
+                "results",
+                "errors",
+                "prepared_source_mapping",
+                "source_prepared_mapping",
+            ]
+        }
+    )
 
     def next(self) -> Optional[BaseInvocation]:
         """Gets the next node ready to execute."""
@@ -857,17 +919,17 @@ class GraphExecutionState(BaseModel):
         """Returns true if the graph has any errors"""
         return len(self.errors) > 0
 
-    def _create_execution_node(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> list[str]:
+    def _create_execution_node(self, node_path: str, iteration_node_map: list[tuple[str, str]]) -> list[str]:
         """Prepares an iteration node and connects all edges, returning the new node id"""
 
-        node = self.graph.get_node(node_id)
+        node = self.graph.get_node(node_path)
 
         self_iteration_count = -1
 
         # If this is an iterator node, we must create a copy for each iteration
         if isinstance(node, IterateInvocation):
             # Get input collection edge (should error if there are no inputs)
-            input_collection_edge = next(iter(self.graph._get_input_edges(node_id, "collection")))
+            input_collection_edge = next(iter(self.graph._get_input_edges(node_path, "collection")))
             input_collection_prepared_node_id = next(
                 n[1] for n in iteration_node_map if n[0] == input_collection_edge.source.node_id
             )
@@ -881,7 +943,7 @@ class GraphExecutionState(BaseModel):
             return new_nodes
 
         # Get all input edges
-        input_edges = self.graph._get_input_edges(node_id)
+        input_edges = self.graph._get_input_edges(node_path)
 
         # Create new edges for this iteration
         # For collect nodes, this may contain multiple inputs to the same field
@@ -908,10 +970,10 @@ class GraphExecutionState(BaseModel):
 
             # Add to execution graph
             self.execution_graph.add_node(new_node)
-            self.prepared_source_mapping[new_node.id] = node_id
-            if node_id not in self.source_prepared_mapping:
-                self.source_prepared_mapping[node_id] = set()
-            self.source_prepared_mapping[node_id].add(new_node.id)
+            self.prepared_source_mapping[new_node.id] = node_path
+            if node_path not in self.source_prepared_mapping:
+                self.source_prepared_mapping[node_path] = set()
+            self.source_prepared_mapping[node_path].add(new_node.id)
 
             # Add new edges to execution graph
             for edge in new_edges:
@@ -1015,13 +1077,13 @@ class GraphExecutionState(BaseModel):
 
     def _get_iteration_node(
         self,
-        source_node_id: str,
+        source_node_path: str,
         graph: nx.DiGraph,
         execution_graph: nx.DiGraph,
         prepared_iterator_nodes: list[str],
     ) -> Optional[str]:
         """Gets the prepared version of the specified source node that matches every iteration specified"""
-        prepared_nodes = self.source_prepared_mapping[source_node_id]
+        prepared_nodes = self.source_prepared_mapping[source_node_path]
         if len(prepared_nodes) == 1:
             return next(iter(prepared_nodes))
 
@@ -1032,7 +1094,7 @@ class GraphExecutionState(BaseModel):
 
         # Filter to only iterator nodes that are a parent of the specified node, in tuple format (prepared, source)
         iterator_source_node_mapping = [(n, self.prepared_source_mapping[n]) for n in prepared_iterator_nodes]
-        parent_iterators = [itn for itn in iterator_source_node_mapping if nx.has_path(graph, itn[1], source_node_id)]
+        parent_iterators = [itn for itn in iterator_source_node_mapping if nx.has_path(graph, itn[1], source_node_path)]
 
         return next(
             (n for n in prepared_nodes if all(nx.has_path(execution_graph, pit[0], n) for pit in parent_iterators)),
@@ -1101,19 +1163,19 @@ class GraphExecutionState(BaseModel):
     def add_node(self, node: BaseInvocation) -> None:
         self.graph.add_node(node)
 
-    def update_node(self, node_id: str, new_node: BaseInvocation) -> None:
-        if not self._is_node_updatable(node_id):
+    def update_node(self, node_path: str, new_node: BaseInvocation) -> None:
+        if not self._is_node_updatable(node_path):
             raise NodeAlreadyExecutedError(
-                f"Node {node_id} has already been prepared or executed and cannot be updated"
+                f"Node {node_path} has already been prepared or executed and cannot be updated"
             )
-        self.graph.update_node(node_id, new_node)
+        self.graph.update_node(node_path, new_node)
 
-    def delete_node(self, node_id: str) -> None:
-        if not self._is_node_updatable(node_id):
+    def delete_node(self, node_path: str) -> None:
+        if not self._is_node_updatable(node_path):
             raise NodeAlreadyExecutedError(
-                f"Node {node_id} has already been prepared or executed and cannot be deleted"
+                f"Node {node_path} has already been prepared or executed and cannot be deleted"
             )
-        self.graph.delete_node(node_id)
+        self.graph.delete_node(node_path)
 
     def add_edge(self, edge: Edge) -> None:
         if not self._is_node_updatable(edge.destination.node_id):
@@ -1128,3 +1190,63 @@ class GraphExecutionState(BaseModel):
                 f"Destination node {edge.destination.node_id} has already been prepared or executed and cannot have a source edge deleted"
             )
         self.graph.delete_edge(edge)
+
+
+class ExposedNodeInput(BaseModel):
+    node_path: str = Field(description="The node path to the node with the input")
+    field: str = Field(description="The field name of the input")
+    alias: str = Field(description="The alias of the input")
+
+
+class ExposedNodeOutput(BaseModel):
+    node_path: str = Field(description="The node path to the node with the output")
+    field: str = Field(description="The field name of the output")
+    alias: str = Field(description="The alias of the output")
+
+
+class LibraryGraph(BaseModel):
+    id: str = Field(description="The unique identifier for this library graph", default_factory=uuid_string)
+    graph: Graph = Field(description="The graph")
+    name: str = Field(description="The name of the graph")
+    description: str = Field(description="The description of the graph")
+    exposed_inputs: list[ExposedNodeInput] = Field(description="The inputs exposed by this graph", default_factory=list)
+    exposed_outputs: list[ExposedNodeOutput] = Field(
+        description="The outputs exposed by this graph", default_factory=list
+    )
+
+    @field_validator("exposed_inputs", "exposed_outputs")
+    def validate_exposed_aliases(cls, v: list[Union[ExposedNodeInput, ExposedNodeOutput]]):
+        if len(v) != len({i.alias for i in v}):
+            raise ValueError("Duplicate exposed alias")
+        return v
+
+    @model_validator(mode="after")
+    def validate_exposed_nodes(cls, values):
+        graph = values.graph
+
+        # Validate exposed inputs
+        for exposed_input in values.exposed_inputs:
+            if not graph.has_node(exposed_input.node_path):
+                raise ValueError(f"Exposed input node {exposed_input.node_path} does not exist")
+            node = graph.get_node(exposed_input.node_path)
+            if get_input_field(node, exposed_input.field) is None:
+                raise ValueError(
+                    f"Exposed input field {exposed_input.field} does not exist on node {exposed_input.node_path}"
+                )
+
+        # Validate exposed outputs
+        for exposed_output in values.exposed_outputs:
+            if not graph.has_node(exposed_output.node_path):
+                raise ValueError(f"Exposed output node {exposed_output.node_path} does not exist")
+            node = graph.get_node(exposed_output.node_path)
+            if get_output_field(node, exposed_output.field) is None:
+                raise ValueError(
+                    f"Exposed output field {exposed_output.field} does not exist on node {exposed_output.node_path}"
+                )
+
+        return values
+
+
+GraphInvocation.model_rebuild(force=True)
+Graph.model_rebuild(force=True)
+GraphExecutionState.model_rebuild(force=True)

invokeai/app/services/shared/invocation_context.py

@@ -1,4 +1,3 @@
-import threading
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
@@ -13,6 +12,7 @@ from invokeai.app.services.config.config_default import InvokeAIAppConfig
 from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
 from invokeai.app.services.images.images_common import ImageDTO
 from invokeai.app.services.invocation_services import InvocationServices
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
 from invokeai.app.util.step_callback import stable_diffusion_step_callback
 from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType
 from invokeai.backend.model_manager.load.load_base import LoadedModel
@@ -22,7 +22,6 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Condit
 
 if TYPE_CHECKING:
     from invokeai.app.invocations.baseinvocation import BaseInvocation
-    from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
 
 """
 The InvocationContext provides access to various services and data about the current invocation.
@@ -49,18 +48,26 @@ Note: The docstrings are in weird places, but that's where they must be to get I
 
 @dataclass
 class InvocationContextData:
-    queue_item: "SessionQueueItem"
-    """The queue item that is being executed."""
     invocation: "BaseInvocation"
     """The invocation that is being executed."""
-    source_invocation_id: str
-    """The ID of the invocation from which the currently executing invocation was prepared."""
+    session_id: str
+    """The session that is being executed."""
+    queue_id: str
+    """The queue in which the session is being executed."""
+    source_node_id: str
+    """The ID of the node from which the currently executing invocation was prepared."""
+    queue_item_id: int
+    """The ID of the queue item that is being executed."""
+    batch_id: str
+    """The ID of the batch that is being executed."""
+    workflow: Optional[WorkflowWithoutID] = None
+    """The workflow associated with this queue item, if any."""
 
 
 class InvocationContextInterface:
-    def __init__(self, services: InvocationServices, data: InvocationContextData) -> None:
+    def __init__(self, services: InvocationServices, context_data: InvocationContextData) -> None:
         self._services = services
-        self._data = data
+        self._context_data = context_data
 
 
 class BoardsInterface(InvocationContextInterface):
@@ -166,26 +173,26 @@ class ImagesInterface(InvocationContextInterface):
         metadata_ = None
         if metadata:
             metadata_ = metadata
-        elif isinstance(self._data.invocation, WithMetadata):
-            metadata_ = self._data.invocation.metadata
+        elif isinstance(self._context_data.invocation, WithMetadata):
+            metadata_ = self._context_data.invocation.metadata
 
         # If `board_id` is provided directly, use that. Else, use the board provided by `WithBoard`, falling back to None.
         board_id_ = None
         if board_id:
             board_id_ = board_id
-        elif isinstance(self._data.invocation, WithBoard) and self._data.invocation.board:
-            board_id_ = self._data.invocation.board.board_id
+        elif isinstance(self._context_data.invocation, WithBoard) and self._context_data.invocation.board:
+            board_id_ = self._context_data.invocation.board.board_id
 
         return self._services.images.create(
             image=image,
-            is_intermediate=self._data.invocation.is_intermediate,
+            is_intermediate=self._context_data.invocation.is_intermediate,
             image_category=image_category,
             board_id=board_id_,
             metadata=metadata_,
             image_origin=ResourceOrigin.INTERNAL,
-            workflow=self._data.queue_item.workflow,
-            session_id=self._data.queue_item.session_id,
-            node_id=self._data.invocation.id,
+            workflow=self._context_data.workflow,
+            session_id=self._context_data.session_id,
+            node_id=self._context_data.invocation.id,
         )
 
     def get_pil(self, image_name: str, mode: IMAGE_MODES | None = None) -> Image:
@@ -247,7 +254,7 @@ class ConditioningInterface(InvocationContextInterface):
         """
         Saves a conditioning data object, returning its name.
 
-        :param conditioning_data: The conditioning data to save.
+        :param conditioning_context_data: The conditioning data to save.
         """
 
         name = self._services.conditioning.save(obj=conditioning_data)
@@ -285,7 +292,7 @@ class ModelsInterface(InvocationContextInterface):
         # the event payloads.
 
         return self._services.model_manager.load_model_by_key(
-            key=key, submodel_type=submodel_type, context_data=self._data
+            key=key, submodel_type=submodel_type, context_data=self._context_data
         )
 
     def load_by_attrs(
@@ -304,7 +311,7 @@ class ModelsInterface(InvocationContextInterface):
             base_model=base_model,
             model_type=model_type,
             submodel=submodel,
-            context_data=self._data,
+            context_data=self._context_data,
         )
 
     def get_config(self, key: str) -> AnyModelConfig:
@@ -363,16 +370,6 @@ class ConfigInterface(InvocationContextInterface):
 
 
 class UtilInterface(InvocationContextInterface):
-    def __init__(
-        self, services: InvocationServices, data: InvocationContextData, cancel_event: threading.Event
-    ) -> None:
-        super().__init__(services, data)
-        self._cancel_event = cancel_event
-
-    def is_canceled(self) -> bool:
-        """Checks if the current invocation has been canceled."""
-        return self._cancel_event.is_set()
-
     def sd_step_callback(self, intermediate_state: PipelineIntermediateState, base_model: BaseModelType) -> None:
         """
         The step callback emits a progress event with the current step, the total number of
@@ -384,12 +381,17 @@ class UtilInterface(InvocationContextInterface):
         :param base_model: The base model for the current denoising step.
         """
 
+        # The step callback needs access to the events and the invocation queue services, but this
+        # represents a dangerous level of access.
+        #
+        # We wrap the step callback so that nodes do not have direct access to these services.
+
         stable_diffusion_step_callback(
-            context_data=self._data,
+            context_data=self._context_data,
             intermediate_state=intermediate_state,
             base_model=base_model,
+            invocation_queue=self._services.queue,
             events=self._services.events,
-            is_canceled=self.is_canceled,
         )
 
 
@@ -408,51 +410,50 @@ class InvocationContext:
         config: ConfigInterface,
         util: UtilInterface,
         boards: BoardsInterface,
-        data: InvocationContextData,
+        context_data: InvocationContextData,
         services: InvocationServices,
     ) -> None:
         self.images = images
-        """Methods to save, get and update images and their metadata."""
+        """Provides methods to save, get and update images and their metadata."""
         self.tensors = tensors
-        """Methods to save and get tensors, including image, noise, masks, and masked images."""
+        """Provides methods to save and get tensors, including image, noise, masks, and masked images."""
         self.conditioning = conditioning
-        """Methods to save and get conditioning data."""
+        """Provides methods to save and get conditioning data."""
         self.models = models
-        """Methods to check if a model exists, get a model, and get a model's info."""
+        """Provides methods to check if a model exists, get a model, and get a model's info."""
         self.logger = logger
-        """The app logger."""
+        """Provides access to the app logger."""
         self.config = config
-        """The app config."""
+        """Provides access to the app's config."""
         self.util = util
-        """Utility methods, including a method to check if an invocation was canceled and step callbacks."""
+        """Provides utility methods."""
         self.boards = boards
-        """Methods to interact with boards."""
-        self._data = data
-        """An internal API providing access to data about the current queue item and invocation. You probably shouldn't use this. It may change without warning."""
+        """Provides methods to interact with boards."""
+        self._data = context_data
+        """Provides data about the current queue item and invocation. This is an internal API and may change without warning."""
         self._services = services
-        """An internal API providing access to all application services. You probably shouldn't use this. It may change without warning."""
+        """Provides access to the full application services. This is an internal API and may change without warning."""
 
 
 def build_invocation_context(
     services: InvocationServices,
-    data: InvocationContextData,
-    cancel_event: threading.Event,
+    context_data: InvocationContextData,
 ) -> InvocationContext:
     """
     Builds the invocation context for a specific invocation execution.
 
-    :param services: The invocation services to wrap.
-    :param data: The invocation context data.
+    :param invocation_services: The invocation services to wrap.
+    :param invocation_context_data: The invocation context data.
     """
 
-    logger = LoggerInterface(services=services, data=data)
-    images = ImagesInterface(services=services, data=data)
-    tensors = TensorsInterface(services=services, data=data)
-    models = ModelsInterface(services=services, data=data)
-    config = ConfigInterface(services=services, data=data)
-    util = UtilInterface(services=services, data=data, cancel_event=cancel_event)
-    conditioning = ConditioningInterface(services=services, data=data)
-    boards = BoardsInterface(services=services, data=data)
+    logger = LoggerInterface(services=services, context_data=context_data)
+    images = ImagesInterface(services=services, context_data=context_data)
+    tensors = TensorsInterface(services=services, context_data=context_data)
+    models = ModelsInterface(services=services, context_data=context_data)
+    config = ConfigInterface(services=services, context_data=context_data)
+    util = UtilInterface(services=services, context_data=context_data)
+    conditioning = ConditioningInterface(services=services, context_data=context_data)
+    boards = BoardsInterface(services=services, context_data=context_data)
 
     ctx = InvocationContext(
         images=images,
@@ -460,7 +461,7 @@ def build_invocation_context(
         config=config,
         tensors=tensors,
         models=models,
-        data=data,
+        context_data=context_data,
         util=util,
         conditioning=conditioning,
         services=services,

invokeai/app/util/step_callback.py

@@ -1,9 +1,9 @@
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 import torch
 from PIL import Image
 
-from invokeai.app.services.session_processor.session_processor_common import CanceledException, ProgressImage
+from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException, ProgressImage
 from invokeai.backend.model_manager.config import BaseModelType
 
 from ...backend.stable_diffusion import PipelineIntermediateState
@@ -11,6 +11,7 @@ from ...backend.util.util import image_to_dataURL
 
 if TYPE_CHECKING:
     from invokeai.app.services.events.events_base import EventServiceBase
+    from invokeai.app.services.invocation_queue.invocation_queue_base import InvocationQueueABC
     from invokeai.app.services.shared.invocation_context import InvocationContextData
 
 
@@ -33,10 +34,10 @@ def stable_diffusion_step_callback(
     context_data: "InvocationContextData",
     intermediate_state: PipelineIntermediateState,
     base_model: BaseModelType,
+    invocation_queue: "InvocationQueueABC",
     events: "EventServiceBase",
-    is_canceled: Callable[[], bool],
 ) -> None:
-    if is_canceled():
+    if invocation_queue.is_canceled(context_data.session_id):
         raise CanceledException
 
     # Some schedulers report not only the noisy latents at the current timestep,
@@ -114,12 +115,12 @@ def stable_diffusion_step_callback(
     dataURL = image_to_dataURL(image, image_format="JPEG")
 
     events.emit_generator_progress(
-        queue_id=context_data.queue_item.queue_id,
-        queue_item_id=context_data.queue_item.item_id,
-        queue_batch_id=context_data.queue_item.batch_id,
-        graph_execution_state_id=context_data.queue_item.session_id,
+        queue_id=context_data.queue_id,
+        queue_item_id=context_data.queue_item_id,
+        queue_batch_id=context_data.batch_id,
+        graph_execution_state_id=context_data.session_id,
         node_id=context_data.invocation.id,
-        source_node_id=context_data.source_invocation_id,
+        source_node_id=context_data.source_node_id,
         progress_image=ProgressImage(width=width, height=height, dataURL=dataURL),
         step=intermediate_state.step,
         order=intermediate_state.order,

invokeai/app/util/ti_utils.py

@@ -1,47 +1,8 @@
 import re
-from typing import List, Tuple
-
-import invokeai.backend.util.logging as logger
-from invokeai.app.services.model_records import UnknownModelException
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager.config import BaseModelType, ModelType
-from invokeai.backend.textual_inversion import TextualInversionModelRaw
 
 
-def extract_ti_triggers_from_prompt(prompt: str) -> List[str]:
-    ti_triggers: List[str] = []
+def extract_ti_triggers_from_prompt(prompt: str) -> list[str]:
+    ti_triggers = []
     for trigger in re.findall(r"<[a-zA-Z0-9., _-]+>", prompt):
-        ti_triggers.append(str(trigger))
+        ti_triggers.append(trigger)
     return ti_triggers
-
-
-def generate_ti_list(
-    prompt: str, base: BaseModelType, context: InvocationContext
-) -> List[Tuple[str, TextualInversionModelRaw]]:
-    ti_list: List[Tuple[str, TextualInversionModelRaw]] = []
-    for trigger in extract_ti_triggers_from_prompt(prompt):
-        name_or_key = trigger[1:-1]
-        try:
-            loaded_model = context.models.load(key=name_or_key)
-            model = loaded_model.model
-            assert isinstance(model, TextualInversionModelRaw)
-            assert loaded_model.config.base == base
-            ti_list.append((name_or_key, model))
-        except UnknownModelException:
-            try:
-                loaded_model = context.models.load_by_attrs(
-                    model_name=name_or_key, base_model=base, model_type=ModelType.TextualInversion
-                )
-                model = loaded_model.model
-                assert isinstance(model, TextualInversionModelRaw)
-                assert loaded_model.config.base == base
-                ti_list.append((name_or_key, model))
-            except UnknownModelException:
-                pass
-        except ValueError:
-            logger.warning(f'trigger: "{trigger}" more than one similarly-named textual inversion models')
-        except AssertionError:
-            logger.warning(f'trigger: "{trigger}" not a valid textual inversion model for this graph')
-        except Exception:
-            logger.warning(f'Failed to load TI model for trigger: "{trigger}"')
-    return ti_list

invokeai/backend/install/check_root.py

@@ -8,6 +8,7 @@ from invokeai.app.services.config import InvokeAIAppConfig
 
 def check_invokeai_root(config: InvokeAIAppConfig):
     try:
+        assert config.model_conf_path.exists(), f"{config.model_conf_path} not found"
         assert config.db_path.parent.exists(), f"{config.db_path.parent} not found"
         assert config.models_path.exists(), f"{config.models_path} not found"
         if not config.ignore_missing_core_models:

invokeai/backend/install/install_helper.py

@@ -1,11 +1,14 @@
 """Utility (backend) functions used by model_install.py"""
+import re
 from logging import Logger
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
 import omegaconf
+from huggingface_hub import HfFolder
 from pydantic import BaseModel, Field
 from pydantic.dataclasses import dataclass
+from pydantic.networks import AnyHttpUrl
 from requests import HTTPError
 from tqdm import tqdm
 
@@ -15,8 +18,12 @@ from invokeai.app.services.download import DownloadQueueService
 from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage
 from invokeai.app.services.model_install import (
+    HFModelSource,
+    LocalModelSource,
     ModelInstallService,
     ModelInstallServiceBase,
+    ModelSource,
+    URLModelSource,
 )
 from invokeai.app.services.model_metadata import ModelMetadataStoreSQL
 from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL
@@ -24,6 +31,7 @@ from invokeai.app.services.shared.sqlite.sqlite_util import init_db
 from invokeai.backend.model_manager import (
     BaseModelType,
     InvalidModelConfigException,
+    ModelRepoVariant,
     ModelType,
 )
 from invokeai.backend.model_manager.metadata import UnknownMetadataException
@@ -218,13 +226,37 @@ class InstallHelper(object):
                     additional_models.append(reverse_source[requirement])
         model_list.extend(additional_models)
 
+    def _make_install_source(self, model_info: UnifiedModelInfo) -> ModelSource:
+        assert model_info.source
+        model_path_id_or_url = model_info.source.strip("\"' ")
+        model_path = Path(model_path_id_or_url)
+
+        if model_path.exists():  # local file on disk
+            return LocalModelSource(path=model_path.absolute(), inplace=True)
+
+        # parsing huggingface repo ids
+        # we're going to do a little trick that allows for extended repo_ids of form "foo/bar:fp16"
+        variants = "|".join([x.lower() for x in ModelRepoVariant.__members__])
+        if match := re.match(f"^([^/]+/[^/]+?)(?::({variants}))?$", model_path_id_or_url):
+            repo_id = match.group(1)
+            repo_variant = ModelRepoVariant(match.group(2)) if match.group(2) else None
+            subfolder = Path(model_info.subfolder) if model_info.subfolder else None
+            return HFModelSource(
+                repo_id=repo_id,
+                access_token=HfFolder.get_token(),
+                subfolder=subfolder,
+                variant=repo_variant,
+            )
+        if re.match(r"^(http|https):", model_path_id_or_url):
+            return URLModelSource(url=AnyHttpUrl(model_path_id_or_url))
+        raise ValueError(f"Unsupported model source: {model_path_id_or_url}")
+
     def add_or_delete(self, selections: InstallSelections) -> None:
         """Add or delete selected models."""
         installer = self._installer
         self._add_required_models(selections.install_models)
         for model in selections.install_models:
-            assert model.source
-            model_path_id_or_url = model.source.strip("\"' ")
+            source = self._make_install_source(model)
             config = (
                 {
                     "description": model.description,
@@ -235,12 +267,12 @@ class InstallHelper(object):
             )
 
             try:
-                installer.heuristic_import(
-                    source=model_path_id_or_url,
+                installer.import_model(
+                    source=source,
                     config=config,
                 )
             except (UnknownMetadataException, InvalidModelConfigException, HTTPError, OSError) as e:
-                self._logger.warning(f"{model.source}: {e}")
+                self._logger.warning(f"{source}: {e}")
 
         for model_to_remove in selections.remove_models:
             parts = model_to_remove.split("/")

invokeai/backend/install/invokeai_configure.py

@@ -939,7 +939,7 @@ def main() -> None:
         # run this unconditionally in case new directories need to be added
         initialize_rootdir(config.root_path, opt.yes_to_all)
 
-        # this will initialize and populate the models tables if not present
+        # this will initialize the models.yaml file if not present
         install_helper = InstallHelper(config, logger)
 
         models_to_download = default_user_selections(opt, install_helper)

invokeai/backend/model_manager/config.py

@@ -138,16 +138,9 @@ class ModelConfigBase(BaseModel):
     source: Optional[str] = Field(description="model original source (path, URL or repo_id)", default=None)
     last_modified: Optional[float] = Field(description="timestamp for modification time", default_factory=time.time)
 
-    @staticmethod
-    def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseModel]) -> None:
-        schema["required"].extend(
-            ["key", "base", "type", "format", "original_hash", "current_hash", "source", "last_modified"]
-        )
-
     model_config = ConfigDict(
         use_enum_values=False,
         validate_assignment=True,
-        json_schema_extra=json_schema_extra,
     )
 
     def update(self, attributes: Dict[str, Any]) -> None:
@@ -234,6 +227,37 @@ class MainDiffusersConfig(_DiffusersConfig, _MainConfig):
     type: Literal[ModelType.Main] = ModelType.Main
 
 
+class ONNXSD1Config(_MainConfig):
+    """Model config for ONNX format models based on sd-1."""
+
+    type: Literal[ModelType.ONNX] = ModelType.ONNX
+    format: Literal[ModelFormat.Onnx, ModelFormat.Olive]
+    base: Literal[BaseModelType.StableDiffusion1] = BaseModelType.StableDiffusion1
+    prediction_type: SchedulerPredictionType = SchedulerPredictionType.Epsilon
+    upcast_attention: bool = False
+
+
+class ONNXSD2Config(_MainConfig):
+    """Model config for ONNX format models based on sd-2."""
+
+    type: Literal[ModelType.ONNX] = ModelType.ONNX
+    format: Literal[ModelFormat.Onnx, ModelFormat.Olive]
+    # No yaml config file for ONNX, so these are part of config
+    base: Literal[BaseModelType.StableDiffusion2] = BaseModelType.StableDiffusion2
+    prediction_type: SchedulerPredictionType = SchedulerPredictionType.VPrediction
+    upcast_attention: bool = True
+
+
+class ONNXSDXLConfig(_MainConfig):
+    """Model config for ONNX format models based on sdxl."""
+
+    type: Literal[ModelType.ONNX] = ModelType.ONNX
+    format: Literal[ModelFormat.Onnx, ModelFormat.Olive]
+    # No yaml config file for ONNX, so these are part of config
+    base: Literal[BaseModelType.StableDiffusionXL] = BaseModelType.StableDiffusionXL
+    prediction_type: SchedulerPredictionType = SchedulerPredictionType.VPrediction
+
+
 class IPAdapterConfig(ModelConfigBase):
     """Model config for IP Adaptor format models."""
 
@@ -256,6 +280,7 @@ class T2IConfig(ModelConfigBase):
     format: Literal[ModelFormat.Diffusers]
 
 
+_ONNXConfig = Annotated[Union[ONNXSD1Config, ONNXSD2Config, ONNXSDXLConfig], Field(discriminator="base")]
 _ControlNetConfig = Annotated[
     Union[ControlNetDiffusersConfig, ControlNetCheckpointConfig],
     Field(discriminator="format"),
@@ -265,6 +290,7 @@ _MainModelConfig = Annotated[Union[MainDiffusersConfig, MainCheckpointConfig], F
 
 AnyModelConfig = Union[
     _MainModelConfig,
+    _ONNXConfig,
     _VaeConfig,
     _ControlNetConfig,
     # ModelConfigBase,

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

@@ -10,7 +10,7 @@ model will be cleared and (re)loaded from disk when next needed.
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from logging import Logger
-from typing import Dict, Generic, Optional, TypeVar
+from typing import Dict, Generic, Optional, Set, TypeVar
 
 import torch
 
@@ -89,8 +89,24 @@ class ModelCacheBase(ABC, Generic[T]):
 
     @property
     @abstractmethod
-    def execution_device(self) -> torch.device:
-        """Return the exection device (e.g. "cuda" for VRAM)."""
+    def execution_devices(self) -> Set[torch.device]:
+        """Return the set of available execution devices."""
+        pass
+
+    @abstractmethod
+    def acquire_execution_device(self, timeout: int = 0) -> torch.device:
+        """
+        Pick the next available execution device.
+
+        If all devices are currently engaged (locked), then
+        block until timeout seconds have passed and raise a
+        TimeoutError if no devices are available.
+        """
+        pass
+
+    @abstractmethod
+    def release_execution_device(self, device: torch.device) -> None:
+        """Release a previously-acquired execution device."""
         pass
 
     @property
@@ -111,7 +127,7 @@ class ModelCacheBase(ABC, Generic[T]):
         pass
 
     @abstractmethod
-    def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device: torch.device) -> None:
+    def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], device: torch.device) -> None:
         """Move model into the indicated device."""
         pass
 

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

@@ -25,7 +25,8 @@ import sys
 import time
 from contextlib import suppress
 from logging import Logger
-from typing import Dict, List, Optional
+from threading import BoundedSemaphore, Lock
+from typing import Dict, List, Optional, Set
 
 import torch
 
@@ -61,8 +62,8 @@ class ModelCache(ModelCacheBase[AnyModel]):
         self,
         max_cache_size: float = DEFAULT_MAX_CACHE_SIZE,
         max_vram_cache_size: float = DEFAULT_MAX_VRAM_CACHE_SIZE,
-        execution_device: torch.device = torch.device("cuda"),
         storage_device: torch.device = torch.device("cpu"),
+        execution_devices: Optional[Set[torch.device]] = None,
         precision: torch.dtype = torch.float16,
         sequential_offload: bool = False,
         lazy_offloading: bool = True,
@@ -74,7 +75,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
         Initialize the model RAM cache.
 
         :param max_cache_size: Maximum size of the RAM cache [6.0 GB]
-        :param execution_device: Torch device to load active model into [torch.device('cuda')]
+        :param execution_devices: Set of torch device to load active model into [calculated]
         :param storage_device: Torch device to save inactive model in [torch.device('cpu')]
         :param precision: Precision for loaded models [torch.float16]
         :param lazy_offloading: Keep model in VRAM until another model needs to be loaded
@@ -89,7 +90,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
         self._precision: torch.dtype = precision
         self._max_cache_size: float = max_cache_size
         self._max_vram_cache_size: float = max_vram_cache_size
-        self._execution_device: torch.device = execution_device
+        self._execution_devices: Set[torch.device] = execution_devices or self._get_execution_devices()
         self._storage_device: torch.device = storage_device
         self._logger = logger or InvokeAILogger.get_logger(self.__class__.__name__)
         self._log_memory_usage = log_memory_usage or self._logger.level == logging.DEBUG
@@ -99,6 +100,10 @@ class ModelCache(ModelCacheBase[AnyModel]):
         self._cached_models: Dict[str, CacheRecord[AnyModel]] = {}
         self._cache_stack: List[str] = []
 
+        self._lock = Lock()
+        self._free_execution_device = BoundedSemaphore(len(self._execution_devices))
+        self._busy_execution_devices: Set[torch.device] = set()
+
     @property
     def logger(self) -> Logger:
         """Return the logger used by the cache."""
@@ -115,9 +120,24 @@ class ModelCache(ModelCacheBase[AnyModel]):
         return self._storage_device
 
     @property
-    def execution_device(self) -> torch.device:
-        """Return the exection device (e.g. "cuda" for VRAM)."""
-        return self._execution_device
+    def execution_devices(self) -> Set[torch.device]:
+        """Return the set of available execution devices."""
+        return self._execution_devices
+
+    def acquire_execution_device(self, timeout: int = 0) -> torch.device:
+        """Acquire and return an execution device (e.g. "cuda" for VRAM)."""
+        with self._lock:
+            self._free_execution_device.acquire(timeout=timeout)
+            free_devices = self.execution_devices - self._busy_execution_devices
+            chosen_device = list(free_devices)[0]
+            self._busy_execution_devices.add(chosen_device)
+        return chosen_device
+
+    def release_execution_device(self, device: torch.device) -> None:
+        """Mark this execution device as unused."""
+        with self._lock:
+            self._free_execution_device.release()
+            self._busy_execution_devices.remove(device)
 
     @property
     def max_cache_size(self) -> float:
@@ -245,13 +265,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
             mps.empty_cache()
 
     def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device: torch.device) -> None:
-        """Move model into the indicated device.
-
-        :param cache_entry: The CacheRecord for the model
-        :param target_device: The torch.device to move the model into
-
-        May raise a torch.cuda.OutOfMemoryError
-        """
+        """Move model into the indicated device."""
         # These attributes are not in the base ModelMixin class but in various derived classes.
         # Some models don't have these attributes, in which case they run in RAM/CPU.
         self.logger.debug(f"Called to move {cache_entry.key} to {target_device}")
@@ -265,9 +279,6 @@ class ModelCache(ModelCacheBase[AnyModel]):
         if torch.device(source_device).type == torch.device(target_device).type:
             return
 
-        # may raise an exception here if insufficient GPU VRAM
-        self._check_free_vram(target_device, cache_entry.size)
-
         start_model_to_time = time.time()
         snapshot_before = self._capture_memory_snapshot()
         cache_entry.model.to(target_device)
@@ -415,12 +426,12 @@ class ModelCache(ModelCacheBase[AnyModel]):
 
         self.logger.debug(f"After making room: cached_models={len(self._cached_models)}")
 
-    def _check_free_vram(self, target_device: torch.device, needed_size: int) -> None:
-        if target_device.type != "cuda":
-            return
-        vram_device = (  # mem_get_info() needs an indexed device
-            target_device if target_device.index is not None else torch.device(str(target_device), index=0)
-        )
-        free_mem, _ = torch.cuda.mem_get_info(torch.device(vram_device))
-        if needed_size > free_mem:
-            raise torch.cuda.OutOfMemoryError
+    @staticmethod
+    def _get_execution_devices() -> Set[torch.device]:
+        default_device = choose_torch_device()
+        if default_device != torch.device("cuda"):
+            return {default_device}
+
+        # we get here if the default device is cuda, and return each of the
+        # cuda devices.
+        return {torch.device(f"cuda:{x}") for x in range(0, torch.cuda.device_count())}

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

@@ -2,12 +2,16 @@
 Base class and implementation of a class that moves models in and out of VRAM.
 """
 
+from typing import Optional
+
 import torch
 
 from invokeai.backend.model_manager import AnyModel
 
 from .model_cache_base import CacheRecord, ModelCacheBase, ModelLockerBase
 
+MAX_GPU_WAIT = 600  # wait up to 10 minutes for a GPU to become free
+
 
 class ModelLocker(ModelLockerBase):
     """Internal class that mediates movement in and out of GPU."""
@@ -21,6 +25,7 @@ class ModelLocker(ModelLockerBase):
         """
         self._cache = cache
         self._cache_entry = cache_entry
+        self._execution_device: Optional[torch.device] = None
 
     @property
     def model(self) -> AnyModel:
@@ -39,15 +44,14 @@ class ModelLocker(ModelLockerBase):
             if self._cache.lazy_offloading:
                 self._cache.offload_unlocked_models(self._cache_entry.size)
 
-            self._cache.move_model_to_device(self._cache_entry, self._cache.execution_device)
+            # We wait for a gpu to be free - may raise a TimeoutError
+            self._execution_device = self._cache.acquire_execution_device(MAX_GPU_WAIT)
+            self._cache.move_model_to_device(self._cache_entry, self._execution_device)
             self._cache_entry.loaded = True
 
-            self._cache.logger.debug(f"Locking {self._cache_entry.key} in {self._cache.execution_device}")
+            self._cache.logger.debug(f"Locking {self._cache_entry.key} in {self._execution_device}")
             self._cache.print_cuda_stats()
-        except torch.cuda.OutOfMemoryError:
-            self._cache.logger.warning("Insufficient GPU memory to load model. Aborting")
-            self._cache_entry.unlock()
-            raise
+
         except Exception:
             self._cache_entry.unlock()
             raise
@@ -59,6 +63,8 @@ class ModelLocker(ModelLockerBase):
             return
 
         self._cache_entry.unlock()
+        if self._execution_device:
+            self._cache.release_execution_device(self._execution_device)
         if not self._cache.lazy_offloading:
             self._cache.offload_unlocked_models(self._cache_entry.size)
             self._cache.print_cuda_stats()

invokeai/backend/model_manager/merge.py

@@ -1,7 +1,7 @@
 """
 invokeai.backend.model_manager.merge exports:
 merge_diffusion_models() -- combine multiple models by location and return a pipeline object
-merge_diffusion_models_and_commit() -- combine multiple models by ModelManager ID and write to the models tables
+merge_diffusion_models_and_commit() -- combine multiple models by ModelManager ID and write to models.yaml
 
 Copyright (c) 2023 Lincoln Stein and the InvokeAI Development Team
 """
@@ -101,7 +101,7 @@ class ModelMerger(object):
         **kwargs: Any,
     ) -> AnyModelConfig:
         """
-        :param models: up to three models, designated by their registered InvokeAI model name
+        :param models: up to three models, designated by their InvokeAI models.yaml model name
         :param merged_model_name: name for new model
         :param alpha: The interpolation parameter. Ranges from 0 to 1.  It affects the ratio in which the checkpoints are merged. A 0.8 alpha
                    would mean that the first model checkpoints would affect the final result far less than an alpha of 0.2

invokeai/backend/model_manager/metadata/fetch/civitai.py

@@ -160,7 +160,7 @@ class CivitaiMetadataFetch(ModelMetadataFetchBase):
             nsfw=model_json["nsfw"],
             restrictions=LicenseRestrictions(
                 AllowNoCredit=model_json["allowNoCredit"],
-                AllowCommercialUse={CommercialUsage(x) for x in model_json["allowCommercialUse"]},
+                AllowCommercialUse=CommercialUsage(model_json["allowCommercialUse"]),
                 AllowDerivatives=model_json["allowDerivatives"],
                 AllowDifferentLicense=model_json["allowDifferentLicense"],
             ),

invokeai/backend/model_manager/metadata/metadata_base.py

@@ -54,8 +54,8 @@ class LicenseRestrictions(BaseModel):
     AllowDifferentLicense: bool = Field(
         description="if true, derivatives of this model be redistributed under a different license", default=False
     )
-    AllowCommercialUse: Optional[Set[CommercialUsage] | CommercialUsage] = Field(
-        description="Type of commercial use allowed if no commercial use is allowed.", default=None
+    AllowCommercialUse: Optional[CommercialUsage] = Field(
+        description="Type of commercial use allowed or 'No' if no commercial use is allowed.", default=None
     )
 
 
@@ -142,10 +142,7 @@ class CivitaiMetadata(ModelMetadataWithFiles):
         if self.restrictions.AllowCommercialUse is None:
             return False
         else:
-            # accommodate schema change
-            acu = self.restrictions.AllowCommercialUse
-            commercial_usage = acu if isinstance(acu, set) else {acu}
-            return CommercialUsage.No not in commercial_usage
+            return self.restrictions.AllowCommercialUse != CommercialUsage("None")
 
     @property
     def allow_derivatives(self) -> bool:

invokeai/backend/model_manager/probe.py

@@ -188,7 +188,7 @@ class ModelProbe(object):
                 and fields["prediction_type"] == SchedulerPredictionType.VPrediction
             )
 
-        model_info = ModelConfigFactory.make_config(fields, key=fields.get("key", None))
+        model_info = ModelConfigFactory.make_config(fields)
         return model_info
 
     @classmethod

invokeai/backend/stable_diffusion/diffusers_pipeline.py

@@ -86,7 +86,6 @@ class AddsMaskGuidance:
     mask_latents: torch.FloatTensor
     scheduler: SchedulerMixin
     noise: torch.Tensor
-    gradient_mask: bool
 
     def __call__(self, step_output: Union[BaseOutput, SchedulerOutput], t: torch.Tensor, conditioning) -> BaseOutput:
         output_class = step_output.__class__  # We'll create a new one with masked data.
@@ -122,12 +121,7 @@ class AddsMaskGuidance:
         # TODO: Do we need to also apply scheduler.scale_model_input? Or is add_noise appropriately scaled already?
         # mask_latents = self.scheduler.scale_model_input(mask_latents, t)
         mask_latents = einops.repeat(mask_latents, "b c h w -> (repeat b) c h w", repeat=batch_size)
-        if self.gradient_mask:
-            threshhold = (t.item()) / self.scheduler.config.num_train_timesteps
-            mask_bool = mask > threshhold  # I don't know when mask got inverted, but it did
-            masked_input = torch.where(mask_bool, latents, mask_latents)
-        else:
-            masked_input = torch.lerp(mask_latents.to(dtype=latents.dtype), latents, mask.to(dtype=latents.dtype))
+        masked_input = torch.lerp(mask_latents.to(dtype=latents.dtype), latents, mask.to(dtype=latents.dtype))
         return masked_input
 
 
@@ -341,7 +335,6 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         t2i_adapter_data: Optional[list[T2IAdapterData]] = None,
         mask: Optional[torch.Tensor] = None,
         masked_latents: Optional[torch.Tensor] = None,
-        gradient_mask: Optional[bool] = False,
         seed: Optional[int] = None,
     ) -> tuple[torch.Tensor, Optional[AttentionMapSaver]]:
         if init_timestep.shape[0] == 0:
@@ -382,7 +375,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
                     self._unet_forward, mask, masked_latents
                 )
             else:
-                additional_guidance.append(AddsMaskGuidance(mask, orig_latents, self.scheduler, noise, gradient_mask))
+                additional_guidance.append(AddsMaskGuidance(mask, orig_latents, self.scheduler, noise))
 
         try:
             latents, attention_map_saver = self.generate_latents_from_embeddings(
@@ -399,7 +392,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
             self.invokeai_diffuser.model_forward_callback = self._unet_forward
 
         # restore unmasked part
-        if mask is not None and not gradient_mask:
+        if mask is not None:
             latents = torch.lerp(orig_latents, latents.to(dtype=orig_latents.dtype), mask.to(dtype=orig_latents.dtype))
 
         return latents, attention_map_saver

invokeai/backend/training/textual_inversion_training.py

@@ -120,7 +120,7 @@ def parse_args() -> Namespace:
         "--model",
         type=str,
         default="sd-1/main/stable-diffusion-v1-5",
-        help="Name of the diffusers model to train against.",
+        help="Name of the diffusers model to train against, as defined in configs/models.yaml.",
     )
     model_group.add_argument(
         "--revision",

invokeai/configs/INITIAL_MODELS.yaml

@@ -34,7 +34,7 @@ sd-1/main/Analog-Diffusion:
    recommended: False
 sd-1/main/Deliberate:
    description: Versatile model that produces detailed images up to 768px (4.27 GB)
-   source: https://huggingface.co/XpucT/Deliberate/resolve/main/Deliberate_v5.safetensors?download=true
+   source: XpucT/Deliberate
    recommended: False
 sd-1/main/Dungeons-and-Diffusion:
    description: Dungeons & Dragons characters (2.13 GB)

invokeai/frontend/install/model_install.py

@@ -455,7 +455,7 @@ class addModelsForm(CyclingForm, npyscreen.FormMultiPage):
         selections = self.parentApp.install_selections
         all_models = self.all_models
 
-        # Defined models (in INITIAL_CONFIG.yaml or invokeai.db) to add/remove
+        # Defined models (in INITIAL_CONFIG.yaml or models.yaml) to add/remove
         ui_sections = [
             self.starter_pipelines,
             self.pipeline_models,

invokeai/frontend/merge/merge_diffusers.py

@@ -435,7 +435,7 @@ def main():
             run_cli(args)
     except widget.NotEnoughSpaceForWidget as e:
         if str(e).startswith("Height of 1 allocated"):
-            logger.error("You need to have at least two diffusers models in order to merge")
+            logger.error("You need to have at least two diffusers models defined in models.yaml in order to merge")
         else:
             logger.error("Not enough room for the user interface. Try making this window larger.")
         sys.exit(-1)

invokeai/frontend/training/textual_inversion.py

@@ -261,7 +261,7 @@ class textualInversionForm(npyscreen.FormMultiPageAction):
     def validate_field_values(self) -> bool:
         bad_fields = []
         if self.model.value is None:
-            bad_fields.append("Model Name must correspond to a known model in invokeai.db")
+            bad_fields.append("Model Name must correspond to a known model in models.yaml")
         if not re.match("^[a-zA-Z0-9.-]+$", self.placeholder_token.value):
             bad_fields.append("Trigger term must only contain alphanumeric characters, the dot and hyphen")
         if self.train_data_dir.value is None:
@@ -442,7 +442,7 @@ def main() -> None:
         pass
     except (widget.NotEnoughSpaceForWidget, Exception) as e:
         if str(e).startswith("Height of 1 allocated"):
-            logger.error("You need to have at least one diffusers models defined in invokeai.db in order to train")
+            logger.error("You need to have at least one diffusers models defined in models.yaml in order to train")
         elif str(e).startswith("addwstr"):
             logger.error("Not enough window space for the interface. Please make your window larger and try again.")
         else:

invokeai/frontend/training/textual_inversion2.py

@@ -0,0 +1,454 @@
+#!/usr/bin/env python
+
+"""
+This is the frontend to "textual_inversion_training.py".
+
+Copyright (c) 2023-24 Lincoln Stein and the InvokeAI Development Team
+"""
+
+
+import os
+import re
+import shutil
+import sys
+import traceback
+from argparse import Namespace
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import npyscreen
+from npyscreen import widget
+from omegaconf import OmegaConf
+
+import invokeai.backend.util.logging as logger
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.install.install_helper import initialize_installer
+from invokeai.backend.model_manager import ModelType
+from invokeai.backend.training import do_textual_inversion_training, parse_args
+
+TRAINING_DATA = "text-inversion-training-data"
+TRAINING_DIR = "text-inversion-output"
+CONF_FILE = "preferences.conf"
+config = None
+
+
+class textualInversionForm(npyscreen.FormMultiPageAction):
+    resolutions = [512, 768, 1024]
+    lr_schedulers = [
+        "linear",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "constant",
+        "constant_with_warmup",
+    ]
+    precisions = ["no", "fp16", "bf16"]
+    learnable_properties = ["object", "style"]
+
+    def __init__(self, parentApp: npyscreen.NPSAppManaged, name: str, saved_args: Optional[Dict[str, str]] = None):
+        self.saved_args = saved_args or {}
+        super().__init__(parentApp, name)
+
+    def afterEditing(self) -> None:
+        self.parentApp.setNextForm(None)
+
+    def create(self) -> None:
+        self.model_names, default = self.get_model_names()
+        default_initializer_token = "★"
+        default_placeholder_token = ""
+        saved_args = self.saved_args
+
+        assert config is not None
+
+        try:
+            default = self.model_names.index(saved_args["model"])
+        except Exception:
+            pass
+
+        self.add_widget_intelligent(
+            npyscreen.FixedText,
+            value="Use ctrl-N and ctrl-P to move to the <N>ext and <P>revious fields, cursor arrows to make a selection, and space to toggle checkboxes.",
+            editable=False,
+        )
+
+        self.model = self.add_widget_intelligent(
+            npyscreen.TitleSelectOne,
+            name="Model Name:",
+            values=sorted(self.model_names),
+            value=default,
+            max_height=len(self.model_names) + 1,
+            scroll_exit=True,
+        )
+        self.placeholder_token = self.add_widget_intelligent(
+            npyscreen.TitleText,
+            name="Trigger Term:",
+            value="",  # saved_args.get('placeholder_token',''), # to restore previous term
+            scroll_exit=True,
+        )
+        self.placeholder_token.when_value_edited = self.initializer_changed
+        self.nextrely -= 1
+        self.nextrelx += 30
+        self.prompt_token = self.add_widget_intelligent(
+            npyscreen.FixedText,
+            name="Trigger term for use in prompt",
+            value="",
+            editable=False,
+            scroll_exit=True,
+        )
+        self.nextrelx -= 30
+        self.initializer_token = self.add_widget_intelligent(
+            npyscreen.TitleText,
+            name="Initializer:",
+            value=saved_args.get("initializer_token", default_initializer_token),
+            scroll_exit=True,
+        )
+        self.resume_from_checkpoint = self.add_widget_intelligent(
+            npyscreen.Checkbox,
+            name="Resume from last saved checkpoint",
+            value=False,
+            scroll_exit=True,
+        )
+        self.learnable_property = self.add_widget_intelligent(
+            npyscreen.TitleSelectOne,
+            name="Learnable property:",
+            values=self.learnable_properties,
+            value=self.learnable_properties.index(saved_args.get("learnable_property", "object")),
+            max_height=4,
+            scroll_exit=True,
+        )
+        self.train_data_dir = self.add_widget_intelligent(
+            npyscreen.TitleFilename,
+            name="Data Training Directory:",
+            select_dir=True,
+            must_exist=False,
+            value=str(
+                saved_args.get(
+                    "train_data_dir",
+                    config.root_dir / TRAINING_DATA / default_placeholder_token,
+                )
+            ),
+            scroll_exit=True,
+        )
+        self.output_dir = self.add_widget_intelligent(
+            npyscreen.TitleFilename,
+            name="Output Destination Directory:",
+            select_dir=True,
+            must_exist=False,
+            value=str(
+                saved_args.get(
+                    "output_dir",
+                    config.root_dir / TRAINING_DIR / default_placeholder_token,
+                )
+            ),
+            scroll_exit=True,
+        )
+        self.resolution = self.add_widget_intelligent(
+            npyscreen.TitleSelectOne,
+            name="Image resolution (pixels):",
+            values=self.resolutions,
+            value=self.resolutions.index(saved_args.get("resolution", 512)),
+            max_height=4,
+            scroll_exit=True,
+        )
+        self.center_crop = self.add_widget_intelligent(
+            npyscreen.Checkbox,
+            name="Center crop images before resizing to resolution",
+            value=saved_args.get("center_crop", False),
+            scroll_exit=True,
+        )
+        self.mixed_precision = self.add_widget_intelligent(
+            npyscreen.TitleSelectOne,
+            name="Mixed Precision:",
+            values=self.precisions,
+            value=self.precisions.index(saved_args.get("mixed_precision", "fp16")),
+            max_height=4,
+            scroll_exit=True,
+        )
+        self.num_train_epochs = self.add_widget_intelligent(
+            npyscreen.TitleSlider,
+            name="Number of training epochs:",
+            out_of=1000,
+            step=50,
+            lowest=1,
+            value=saved_args.get("num_train_epochs", 100),
+            scroll_exit=True,
+        )
+        self.max_train_steps = self.add_widget_intelligent(
+            npyscreen.TitleSlider,
+            name="Max Training Steps:",
+            out_of=10000,
+            step=500,
+            lowest=1,
+            value=saved_args.get("max_train_steps", 3000),
+            scroll_exit=True,
+        )
+        self.train_batch_size = self.add_widget_intelligent(
+            npyscreen.TitleSlider,
+            name="Batch Size (reduce if you run out of memory):",
+            out_of=50,
+            step=1,
+            lowest=1,
+            value=saved_args.get("train_batch_size", 8),
+            scroll_exit=True,
+        )
+        self.gradient_accumulation_steps = self.add_widget_intelligent(
+            npyscreen.TitleSlider,
+            name="Gradient Accumulation Steps (may need to decrease this to resume from a checkpoint):",
+            out_of=10,
+            step=1,
+            lowest=1,
+            value=saved_args.get("gradient_accumulation_steps", 4),
+            scroll_exit=True,
+        )
+        self.lr_warmup_steps = self.add_widget_intelligent(
+            npyscreen.TitleSlider,
+            name="Warmup Steps:",
+            out_of=100,
+            step=1,
+            lowest=0,
+            value=saved_args.get("lr_warmup_steps", 0),
+            scroll_exit=True,
+        )
+        self.learning_rate = self.add_widget_intelligent(
+            npyscreen.TitleText,
+            name="Learning Rate:",
+            value=str(
+                saved_args.get("learning_rate", "5.0e-04"),
+            ),
+            scroll_exit=True,
+        )
+        self.scale_lr = self.add_widget_intelligent(
+            npyscreen.Checkbox,
+            name="Scale learning rate by number GPUs, steps and batch size",
+            value=saved_args.get("scale_lr", True),
+            scroll_exit=True,
+        )
+        self.enable_xformers_memory_efficient_attention = self.add_widget_intelligent(
+            npyscreen.Checkbox,
+            name="Use xformers acceleration",
+            value=saved_args.get("enable_xformers_memory_efficient_attention", False),
+            scroll_exit=True,
+        )
+        self.lr_scheduler = self.add_widget_intelligent(
+            npyscreen.TitleSelectOne,
+            name="Learning rate scheduler:",
+            values=self.lr_schedulers,
+            max_height=7,
+            value=self.lr_schedulers.index(saved_args.get("lr_scheduler", "constant")),
+            scroll_exit=True,
+        )
+        self.model.editing = True
+
+    def initializer_changed(self) -> None:
+        placeholder = self.placeholder_token.value
+        self.prompt_token.value = f"(Trigger by using <{placeholder}> in your prompts)"
+        self.train_data_dir.value = str(config.root_dir / TRAINING_DATA / placeholder)
+        self.output_dir.value = str(config.root_dir / TRAINING_DIR / placeholder)
+        self.resume_from_checkpoint.value = Path(self.output_dir.value).exists()
+
+    def on_ok(self):
+        if self.validate_field_values():
+            self.parentApp.setNextForm(None)
+            self.editing = False
+            self.parentApp.ti_arguments = self.marshall_arguments()
+            npyscreen.notify("Launching textual inversion training. This will take a while...")
+        else:
+            self.editing = True
+
+    def ok_cancel(self):
+        sys.exit(0)
+
+    def validate_field_values(self) -> bool:
+        bad_fields = []
+        if self.model.value is None:
+            bad_fields.append("Model Name must correspond to a known model in models.yaml")
+        if not re.match("^[a-zA-Z0-9.-]+$", self.placeholder_token.value):
+            bad_fields.append("Trigger term must only contain alphanumeric characters, the dot and hyphen")
+        if self.train_data_dir.value is None:
+            bad_fields.append("Data Training Directory cannot be empty")
+        if self.output_dir.value is None:
+            bad_fields.append("The Output Destination Directory cannot be empty")
+        if len(bad_fields) > 0:
+            message = "The following problems were detected and must be corrected:"
+            for problem in bad_fields:
+                message += f"\n* {problem}"
+            npyscreen.notify_confirm(message)
+            return False
+        else:
+            return True
+
+    def get_model_names(self) -> Tuple[List[str], int]:
+        global config
+        assert config is not None
+        installer = initialize_installer(config)
+        store = installer.record_store
+        main_models = store.search_by_attr(model_type=ModelType.Main)
+        model_names = [f"{x.base.value}/{x.type.value}/{x.name}" for x in main_models if x.format == "diffusers"]
+        default = 0
+        return (model_names, default)
+
+    def marshall_arguments(self) -> dict:
+        args = {}
+
+        # the choices
+        args.update(
+            model=self.model_names[self.model.value[0]],
+            resolution=self.resolutions[self.resolution.value[0]],
+            lr_scheduler=self.lr_schedulers[self.lr_scheduler.value[0]],
+            mixed_precision=self.precisions[self.mixed_precision.value[0]],
+            learnable_property=self.learnable_properties[self.learnable_property.value[0]],
+        )
+
+        # all the strings and booleans
+        for attr in (
+            "initializer_token",
+            "placeholder_token",
+            "train_data_dir",
+            "output_dir",
+            "scale_lr",
+            "center_crop",
+            "enable_xformers_memory_efficient_attention",
+        ):
+            args[attr] = getattr(self, attr).value
+
+        # all the integers
+        for attr in (
+            "train_batch_size",
+            "gradient_accumulation_steps",
+            "num_train_epochs",
+            "max_train_steps",
+            "lr_warmup_steps",
+        ):
+            args[attr] = int(getattr(self, attr).value)
+
+        # the floats (just one)
+        args.update(learning_rate=float(self.learning_rate.value))
+
+        # a special case
+        if self.resume_from_checkpoint.value and Path(self.output_dir.value).exists():
+            args["resume_from_checkpoint"] = "latest"
+
+        return args
+
+
+class MyApplication(npyscreen.NPSAppManaged):
+    def __init__(self, saved_args: Optional[Dict[str, str]] = None):
+        super().__init__()
+        self.ti_arguments = None
+        self.saved_args = saved_args
+
+    def onStart(self):
+        npyscreen.setTheme(npyscreen.Themes.DefaultTheme)
+        self.main = self.addForm(
+            "MAIN",
+            textualInversionForm,
+            name="Textual Inversion Settings",
+            saved_args=self.saved_args,
+        )
+
+
+def copy_to_embeddings_folder(args: Dict[str, str]) -> None:
+    """
+    Copy learned_embeds.bin into the embeddings folder, and offer to
+    delete the full model and checkpoints.
+    """
+    assert config is not None
+    source = Path(args["output_dir"], "learned_embeds.bin")
+    dest_dir_name = args["placeholder_token"].strip("<>")
+    destination = config.root_dir / "embeddings" / dest_dir_name
+    os.makedirs(destination, exist_ok=True)
+    logger.info(f"Training completed. Copying learned_embeds.bin into {str(destination)}")
+    shutil.copy(source, destination)
+    if (input("Delete training logs and intermediate checkpoints? [y] ") or "y").startswith(("y", "Y")):
+        shutil.rmtree(Path(args["output_dir"]))
+    else:
+        logger.info(f'Keeping {args["output_dir"]}')
+
+
+def save_args(args: dict) -> None:
+    """
+    Save the current argument values to an omegaconf file
+    """
+    assert config is not None
+    dest_dir = config.root_dir / TRAINING_DIR
+    os.makedirs(dest_dir, exist_ok=True)
+    conf_file = dest_dir / CONF_FILE
+    conf = OmegaConf.create(args)
+    OmegaConf.save(config=conf, f=conf_file)
+
+
+def previous_args() -> dict:
+    """
+    Get the previous arguments used.
+    """
+    assert config is not None
+    conf_file = config.root_dir / TRAINING_DIR / CONF_FILE
+    try:
+        conf = OmegaConf.load(conf_file)
+        conf["placeholder_token"] = conf["placeholder_token"].strip("<>")
+    except Exception:
+        conf = None
+
+    return conf
+
+
+def do_front_end() -> None:
+    global config
+    saved_args = previous_args()
+    myapplication = MyApplication(saved_args=saved_args)
+    myapplication.run()
+
+    if my_args := myapplication.ti_arguments:
+        os.makedirs(my_args["output_dir"], exist_ok=True)
+
+        # Automatically add angle brackets around the trigger
+        if not re.match("^<.+>$", my_args["placeholder_token"]):
+            my_args["placeholder_token"] = f"<{my_args['placeholder_token']}>"
+
+        my_args["only_save_embeds"] = True
+        save_args(my_args)
+
+        try:
+            print(my_args)
+            do_textual_inversion_training(config, **my_args)
+            copy_to_embeddings_folder(my_args)
+        except Exception as e:
+            logger.error("An exception occurred during training. The exception was:")
+            logger.error(str(e))
+            logger.error("DETAILS:")
+            logger.error(traceback.format_exc())
+
+
+def main() -> None:
+    global config
+
+    args: Namespace = parse_args()
+    config = InvokeAIAppConfig.get_config()
+    config.parse_args([])
+
+    # change root if needed
+    if args.root_dir:
+        config.root = args.root_dir
+
+    try:
+        if args.front_end:
+            do_front_end()
+        else:
+            do_textual_inversion_training(config, **vars(args))
+    except AssertionError as e:
+        logger.error(e)
+        sys.exit(-1)
+    except KeyboardInterrupt:
+        pass
+    except (widget.NotEnoughSpaceForWidget, Exception) as e:
+        if str(e).startswith("Height of 1 allocated"):
+            logger.error("You need to have at least one diffusers models defined in models.yaml in order to train")
+        elif str(e).startswith("addwstr"):
+            logger.error("Not enough window space for the interface. Please make your window larger and try again.")
+        else:
+            logger.error(e)
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    main()

invokeai/frontend/web/.unimportedrc.json

@@ -0,0 +1,9 @@
+{
+  "entry": ["src/main.tsx"],
+  "extensions": [".ts", ".tsx"],
+  "ignorePatterns": ["**/node_modules/**", "dist/**", "public/**", "**/*.stories.tsx", "config/**"],
+  "ignoreUnresolved": [],
+  "ignoreUnimported": ["src/i18.d.ts", "vite.config.ts", "src/vite-env.d.ts"],
+  "respectGitignore": true,
+  "ignoreUnused": []
+}

invokeai/frontend/web/knip.ts

@@ -1,27 +0,0 @@
-import type { KnipConfig } from 'knip';
-
-const config: KnipConfig = {
-  ignore: [
-    // This file is only used during debugging
-    'src/app/store/middleware/debugLoggerMiddleware.ts',
-    // Autogenerated types - shouldn't ever touch these
-    'src/services/api/schema.ts',
-  ],
-  ignoreBinaries: ['only-allow'],
-  rules: {
-    files: 'warn',
-    dependencies: 'warn',
-    unlisted: 'warn',
-    binaries: 'warn',
-    unresolved: 'warn',
-    exports: 'warn',
-    types: 'warn',
-    nsExports: 'warn',
-    nsTypes: 'warn',
-    enumMembers: 'warn',
-    classMembers: 'warn',
-    duplicates: 'warn',
-  },
-};
-
-export default config;

invokeai/frontend/web/package.json

@@ -24,16 +24,16 @@
     "build": "pnpm run lint && vite build",
     "typegen": "node scripts/typegen.js",
     "preview": "vite preview",
-    "lint:knip": "knip",
-    "lint:dpdm": "dpdm --no-warning --no-tree --transform --exit-code circular:1 src/main.tsx",
+    "lint:madge": "madge --circular src/main.tsx",
     "lint:eslint": "eslint --max-warnings=0 .",
     "lint:prettier": "prettier --check .",
     "lint:tsc": "tsc --noEmit",
-    "lint": "concurrently -g -c red,green,yellow,blue,magenta pnpm:lint:*",
-    "fix": "knip --fix && eslint --fix . && prettier --log-level warn --write .",
+    "lint": "concurrently -g -n eslint,prettier,tsc,madge -c cyan,green,magenta,yellow \"pnpm run lint:eslint\" \"pnpm run lint:prettier\" \"pnpm run lint:tsc\" \"pnpm run lint:madge\"",
+    "fix": "eslint --fix . && prettier --log-level warn --write .",
     "preinstall": "npx only-allow pnpm",
     "storybook": "storybook dev -p 6006",
     "build-storybook": "storybook build",
+    "unimported": "npx unimported",
     "test": "vitest",
     "test:no-watch": "vitest --no-watch"
   },
@@ -57,51 +57,54 @@
     "@dnd-kit/sortable": "^8.0.0",
     "@dnd-kit/utilities": "^3.2.2",
     "@fontsource-variable/inter": "^5.0.16",
-    "@invoke-ai/ui-library": "^0.0.21",
-    "@nanostores/react": "^0.7.2",
-    "@reduxjs/toolkit": "2.2.1",
+    "@invoke-ai/ui-library": "^0.0.18",
+    "@mantine/form": "6.0.21",
+    "@nanostores/react": "^0.7.1",
+    "@reduxjs/toolkit": "2.0.1",
     "@roarr/browser-log-writer": "^1.3.0",
     "chakra-react-select": "^4.7.6",
     "compare-versions": "^6.1.0",
     "dateformat": "^5.0.3",
-    "framer-motion": "^11.0.6",
-    "i18next": "^23.10.0",
-    "i18next-http-backend": "^2.5.0",
+    "framer-motion": "^10.18.0",
+    "i18next": "^23.7.16",
+    "i18next-http-backend": "^2.4.2",
     "idb-keyval": "^6.2.1",
     "jsondiffpatch": "^0.6.0",
-    "konva": "^9.3.3",
+    "konva": "^9.3.1",
     "lodash-es": "^4.17.21",
-    "nanostores": "^0.10.0",
+    "nanostores": "^0.9.5",
     "new-github-issue-url": "^1.0.0",
-    "overlayscrollbars": "^2.5.0",
-    "overlayscrollbars-react": "^0.5.4",
-    "query-string": "^9.0.0",
+    "overlayscrollbars": "^2.4.6",
+    "overlayscrollbars-react": "^0.5.3",
+    "query-string": "^8.1.0",
     "react": "^18.2.0",
     "react-colorful": "^5.6.1",
     "react-dom": "^18.2.0",
     "react-dropzone": "^14.2.3",
     "react-error-boundary": "^4.0.12",
-    "react-hook-form": "^7.50.1",
-    "react-hotkeys-hook": "4.5.0",
-    "react-i18next": "^14.0.5",
+    "react-hook-form": "^7.49.3",
+    "react-hotkeys-hook": "4.4.4",
+    "react-i18next": "^14.0.0",
     "react-icons": "^5.0.1",
     "react-konva": "^18.2.10",
     "react-redux": "9.1.0",
-    "react-resizable-panels": "^2.0.11",
+    "react-resizable-panels": "^1.0.9",
     "react-select": "5.8.0",
-    "react-use": "^17.5.0",
-    "react-virtuoso": "^4.7.1",
-    "reactflow": "^11.10.4",
+    "react-textarea-autosize": "^8.5.3",
+    "react-use": "^17.4.3",
+    "react-virtuoso": "^4.6.2",
+    "reactflow": "^11.10.2",
     "redux-dynamic-middlewares": "^2.2.0",
     "redux-remember": "^5.1.0",
     "roarr": "^7.21.0",
     "serialize-error": "^11.0.3",
     "socket.io-client": "^4.7.4",
+    "type-fest": "^4.9.0",
     "use-debounce": "^10.0.0",
     "use-image": "^1.1.1",
     "uuid": "^9.0.1",
     "zod": "^3.22.4",
-    "zod-validation-error": "^3.0.2"
+    "zod-validation-error": "^3.0.0"
   },
   "peerDependencies": {
     "@chakra-ui/react": "^2.8.2",
@@ -110,42 +113,59 @@
     "ts-toolbelt": "^9.6.0"
   },
   "devDependencies": {
-    "@invoke-ai/eslint-config-react": "^0.0.14",
-    "@invoke-ai/prettier-config-react": "^0.0.7",
-    "@storybook/addon-essentials": "^7.6.17",
-    "@storybook/addon-interactions": "^7.6.17",
-    "@storybook/addon-links": "^7.6.17",
-    "@storybook/addon-storysource": "^7.6.17",
-    "@storybook/manager-api": "^7.6.17",
-    "@storybook/react": "^7.6.17",
-    "@storybook/react-vite": "^7.6.17",
-    "@storybook/theming": "^7.6.17",
+    "@arthurgeron/eslint-plugin-react-usememo": "^2.2.3",
+    "@invoke-ai/eslint-config-react": "^0.0.13",
+    "@invoke-ai/prettier-config-react": "^0.0.6",
+    "@storybook/addon-docs": "^7.6.10",
+    "@storybook/addon-essentials": "^7.6.10",
+    "@storybook/addon-interactions": "^7.6.10",
+    "@storybook/addon-links": "^7.6.10",
+    "@storybook/addon-storysource": "^7.6.10",
+    "@storybook/blocks": "^7.6.10",
+    "@storybook/manager-api": "^7.6.10",
+    "@storybook/react": "^7.6.10",
+    "@storybook/react-vite": "^7.6.10",
+    "@storybook/test": "^7.6.10",
+    "@storybook/theming": "^7.6.10",
     "@types/dateformat": "^5.0.2",
     "@types/lodash-es": "^4.17.12",
-    "@types/node": "^20.11.20",
-    "@types/react": "^18.2.59",
-    "@types/react-dom": "^18.2.19",
-    "@types/uuid": "^9.0.8",
-    "@vitejs/plugin-react-swc": "^3.6.0",
+    "@types/node": "^20.11.5",
+    "@types/react": "^18.2.48",
+    "@types/react-dom": "^18.2.18",
+    "@types/uuid": "^9.0.7",
+    "@typescript-eslint/eslint-plugin": "^6.19.0",
+    "@typescript-eslint/parser": "^6.19.0",
+    "@vitejs/plugin-react-swc": "^3.5.0",
     "concurrently": "^8.2.2",
-    "dpdm": "^3.14.0",
-    "eslint": "^8.57.0",
+    "eslint": "^8.56.0",
+    "eslint-config-prettier": "^9.1.0",
     "eslint-plugin-i18next": "^6.0.3",
+    "eslint-plugin-import": "^2.29.1",
     "eslint-plugin-path": "^1.2.4",
-    "knip": "^5.0.2",
+    "eslint-plugin-react": "^7.33.2",
+    "eslint-plugin-react-hooks": "^4.6.0",
+    "eslint-plugin-simple-import-sort": "^10.0.0",
+    "eslint-plugin-storybook": "^0.6.15",
+    "eslint-plugin-unused-imports": "^3.0.0",
+    "madge": "^6.1.0",
     "openapi-types": "^12.1.3",
-    "openapi-typescript": "^6.7.4",
-    "prettier": "^3.2.5",
+    "openapi-typescript": "^6.7.3",
+    "prettier": "^3.2.4",
     "rollup-plugin-visualizer": "^5.12.0",
-    "storybook": "^7.6.17",
+    "storybook": "^7.6.10",
     "ts-toolbelt": "^9.6.0",
     "tsafe": "^1.6.6",
     "typescript": "^5.3.3",
-    "vite": "^5.1.4",
-    "vite-plugin-css-injected-by-js": "^3.4.0",
-    "vite-plugin-dts": "^3.7.3",
+    "vite": "^5.0.12",
+    "vite-plugin-css-injected-by-js": "^3.3.1",
+    "vite-plugin-dts": "^3.7.1",
     "vite-plugin-eslint": "^1.8.1",
     "vite-tsconfig-paths": "^4.3.1",
-    "vitest": "^1.3.1"
+    "vitest": "^1.2.2"
+  },
+  "pnpm": {
+    "patchedDependencies": {
+      "reselect@5.0.1": "patches/reselect@5.0.1.patch"
+    }
   }
 }

invokeai/frontend/web/public/locales/en.json

@@ -424,11 +424,8 @@
         "uploads": "Uploads",
         "deleteSelection": "Delete Selection",
         "downloadSelection": "Download Selection",
-        "bulkDownloadRequested": "Preparing Download",
-        "bulkDownloadRequestedDesc": "Your download request is being prepared. This may take a few moments.",
-        "bulkDownloadRequestFailed": "Problem Preparing Download",
-        "bulkDownloadStarting": "Download Starting",
-        "bulkDownloadFailed": "Download Failed",
+        "preparingDownload": "Preparing Download",
+        "preparingDownloadFailed": "Problem Preparing Download",
         "problemDeletingImages": "Problem Deleting Images",
         "problemDeletingImagesDesc": "One or more images could not be deleted"
     },
@@ -656,7 +653,6 @@
         }
     },
     "metadata": {
-        "allPrompts": "All Prompts",
         "cfgScale": "CFG scale",
         "cfgRescaleMultiplier": "$t(parameters.cfgRescaleMultiplier)",
         "createdBy": "Created By",
@@ -665,7 +661,6 @@
         "height": "Height",
         "hiresFix": "High Resolution Optimization",
         "imageDetails": "Image Details",
-        "imageDimensions": "Image Dimensions",
         "initImage": "Initial image",
         "metadata": "Metadata",
         "model": "Model",
@@ -673,12 +668,9 @@
         "noImageDetails": "No image details found",
         "noMetaData": "No metadata found",
         "noRecallParameters": "No parameters to recall found",
-        "parameterSet": "Parameter {{parameter}} set",
-        "parsingFailed": "Parsing Failed",
         "perlin": "Perlin Noise",
         "positivePrompt": "Positive Prompt",
         "recallParameters": "Recall Parameters",
-        "recallParameter": "Recall {{label}}",
         "scheduler": "Scheduler",
         "seamless": "Seamless",
         "seed": "Seed",
@@ -692,24 +684,20 @@
     },
     "modelManager": {
         "active": "active",
-        "addAll": "Add All",
         "addCheckpointModel": "Add Checkpoint / Safetensor Model",
         "addDifference": "Add Difference",
         "addDiffuserModel": "Add Diffusers",
         "addManually": "Add Manually",
         "addModel": "Add Model",
-        "addModels": "Add Models",
         "addNew": "Add New",
         "addNewModel": "Add New Model",
         "addSelected": "Add Selected",
         "advanced": "Advanced",
-        "advancedImportInfo": "The advanced tab allows for manual configuration of core model settings. Only use this tab if you are confident that you know the correct model type and configuration for the selected model.",
         "allModels": "All Models",
         "alpha": "Alpha",
         "availableModels": "Available Models",
         "baseModel": "Base Model",
         "cached": "cached",
-        "cancel": "Cancel",
         "cannotUseSpaces": "Cannot Use Spaces",
         "checkpointFolder": "Checkpoint Folder",
         "checkpointModels": "Checkpoints",
@@ -743,7 +731,6 @@
         "descriptionValidationMsg": "Add a description for your model",
         "deselectAll": "Deselect All",
         "diffusersModels": "Diffusers",
-        "edit": "Edit",
         "findModels": "Find Models",
         "formMessageDiffusersModelLocation": "Diffusers Model Location",
         "formMessageDiffusersModelLocationDesc": "Please enter at least one.",
@@ -752,9 +739,7 @@
         "height": "Height",
         "heightValidationMsg": "Default height of your model.",
         "ignoreMismatch": "Ignore Mismatches Between Selected Models",
-        "imageEncoderModelId": "Image Encoder Model ID",
         "importModels": "Import Models",
-        "importQueue": "Import Queue",
         "inpainting": "v1 Inpainting",
         "interpolationType": "Interpolation Type",
         "inverseSigmoid": "Inverse Sigmoid",
@@ -783,11 +768,8 @@
         "modelMergeHeaderHelp1": "You can merge up to three different models to create a blend that suits your needs.",
         "modelMergeHeaderHelp2": "Only Diffusers are available for merging. If you want to merge a checkpoint model, please convert it to Diffusers first.",
         "modelMergeInterpAddDifferenceHelp": "In this mode, Model 3 is first subtracted from Model 2. The resulting version is blended with Model 1 with the alpha rate set above.",
-        "modelMetadata": "Model Metadata",
-        "modelName": "Model Name",
         "modelOne": "Model 1",
         "modelsFound": "Models Found",
-        "modelSettings": "Model Settings",
         "modelsMerged": "Models Merged",
         "modelsMergeFailed": "Model Merge Failed",
         "modelsSynced": "Models Synced",
@@ -807,24 +789,16 @@
         "notLoaded": "not loaded",
         "oliveModels": "Olives",
         "onnxModels": "Onnx",
-        "path": "Path",
         "pathToCustomConfig": "Path To Custom Config",
         "pickModelType": "Pick Model Type",
-        "predictionType": "Prediction Type",
-        "prune": "Prune",
-        "pruneTooltip": "Prune finished imports from queue",
+        "predictionType": "Prediction Type (for Stable Diffusion 2.x Models and occasional Stable Diffusion 1.x Models)",
         "quickAdd": "Quick Add",
-        "removeFromQueue": "Remove From Queue",
         "repo_id": "Repo ID",
         "repoIDValidationMsg": "Online repository of your model",
-        "repoVariant": "Repo Variant",
         "safetensorModels": "SafeTensors",
         "sameFolder": "Same folder",
-        "scan": "Scan",
-        "scanFolder": "Scan folder",
         "scanAgain": "Scan Again",
         "scanForModels": "Scan For Models",
-        "scanResults": "Scan Results",
         "search": "Search",
         "selectAll": "Select All",
         "selectAndAdd": "Select and Add Models Listed Below",
@@ -835,11 +809,9 @@
         "showExisting": "Show Existing",
         "sigmoid": "Sigmoid",
         "simpleModelDesc": "Provide a path to a local Diffusers model, local checkpoint / safetensors model a HuggingFace Repo ID, or a checkpoint/diffusers model URL.",
-        "source": "Source",
         "statusConverting": "Converting",
         "syncModels": "Sync Models",
-        "syncModelsDesc": "If your models are out of sync with the backend, you can refresh them up using this option. This is generally handy in cases where you add models to the InvokeAI root folder or autoimport directory after the application has booted.",
-        "upcastAttention": "Upcast Attention",
+        "syncModelsDesc": "If your models are out of sync with the backend, you can refresh them up using this option. This is generally handy in cases where you manually update your models.yaml file or add models to the InvokeAI root folder after the application has booted.",
         "updateModel": "Update Model",
         "useCustomConfig": "Use Custom Config",
         "v1": "v1",
@@ -854,8 +826,7 @@
         "variant": "Variant",
         "weightedSum": "Weighted Sum",
         "width": "Width",
-        "widthValidationMsg": "Default width of your model.",
-        "ztsnrTraining": "ZTSNR Training"
+        "widthValidationMsg": "Default width of your model."
     },
     "models": {
         "addLora": "Add LoRA",
@@ -1149,8 +1120,8 @@
         "codeformerFidelity": "Fidelity",
         "coherenceMode": "Mode",
         "coherencePassHeader": "Coherence Pass",
-        "coherenceEdgeSize": "Edge Size",
-        "coherenceMinDenoise": "Min Denoise",
+        "coherenceSteps": "Steps",
+        "coherenceStrength": "Strength",
         "compositingSettingsHeader": "Compositing Settings",
         "controlNetControlMode": "Control Mode",
         "copyImage": "Copy Image",
@@ -1194,8 +1165,8 @@
             "unableToInvoke": "Unable to Invoke"
         },
         "maskAdjustmentsHeader": "Mask Adjustments",
-        "maskBlur": "Mask Blur",
-        "maskBlurMethod": "Mask Blur Method",
+        "maskBlur": "Blur",
+        "maskBlurMethod": "Blur Method",
         "maskEdge": "Mask Edge",
         "negativePromptPlaceholder": "Negative Prompt",
         "noiseSettings": "Noise",
@@ -1377,8 +1348,6 @@
         "modelAdded": "Model Added: {{modelName}}",
         "modelAddedSimple": "Model Added",
         "modelAddFailed": "Model Add Failed",
-        "modelImportCanceled": "Model Import Canceled",
-        "modelImportRemoved": "Model Import Removed",
         "nodesBrokenConnections": "Cannot load. Some connections are broken.",
         "nodesCorruptedGraph": "Cannot load. Graph seems to be corrupted.",
         "nodesLoaded": "Nodes Loaded",
@@ -1387,8 +1356,8 @@
         "nodesNotValidJSON": "Not a valid JSON",
         "nodesSaved": "Nodes Saved",
         "nodesUnrecognizedTypes": "Cannot load. Graph has unrecognized types",
-        "parameterNotSet": "{{parameter}} not set",
-        "parameterSet": "{{parameter}} set",
+        "parameterNotSet": "Parameter not set",
+        "parameterSet": "Parameter set",
         "parametersFailed": "Problem loading parameters",
         "parametersFailedDesc": "Unable to load init image.",
         "parametersNotSet": "Parameters Not Set",
@@ -1412,7 +1381,6 @@
         "promptNotSet": "Prompt Not Set",
         "promptNotSetDesc": "Could not find prompt for this image.",
         "promptSet": "Prompt Set",
-        "prunedQueue": "Pruned Queue",
         "resetInitialImage": "Reset Initial Image",
         "seedNotSet": "Seed Not Set",
         "seedNotSetDesc": "Could not find seed for this image.",
@@ -1481,8 +1449,8 @@
                 "Scheduler defines how to iteratively add noise to an image or how to update a sample based on a model's output."
             ]
         },
-        "compositingMaskBlur": {
-            "heading": "Mask Blur",
+        "compositingBlur": {
+            "heading": "Blur",
             "paragraphs": ["The blur radius of the mask."]
         },
         "compositingBlurMethod": {
@@ -1497,15 +1465,15 @@
             "heading": "Mode",
             "paragraphs": ["The mode of the Coherence Pass."]
         },
-        "compositingCoherenceEdgeSize": {
-            "heading": "Edge Size",
-            "paragraphs": ["The edge size of the coherence pass."]
+        "compositingCoherenceSteps": {
+            "heading": "Steps",
+            "paragraphs": ["Number of denoising steps used in the Coherence Pass.", "Same as the main Steps parameter."]
         },
-        "compositingCoherenceMinDenoise": {
-            "heading": "Minimum Denoise",
+        "compositingStrength": {
+            "heading": "Strength",
             "paragraphs": [
-                "Minimum denoise strength for the Coherence mode",
-                "The minimum denoise strength for the coherence region when inpainting or outpainting"
+                "Denoising strength for the Coherence Pass.",
+                "Same as the Image to Image Denoising Strength parameter."
             ]
         },
         "compositingMaskAdjustments": {
@@ -1674,9 +1642,6 @@
         "clearCanvasHistoryMessage": "Clearing the canvas history leaves your current canvas intact, but irreversibly clears the undo and redo history.",
         "clearHistory": "Clear History",
         "clearMask": "Clear Mask (Shift+C)",
-        "coherenceModeGaussianBlur": "Gaussian Blur",
-        "coherenceModeBoxBlur": "Box Blur",
-        "coherenceModeStaged": "Staged",
         "colorPicker": "Color Picker",
         "copyToClipboard": "Copy to Clipboard",
         "cursorPosition": "Cursor Position",

invokeai/frontend/web/scripts/colors.js

@@ -0,0 +1,34 @@
+export const COLORS = {
+  reset: '\x1b[0m',
+  bright: '\x1b[1m',
+  dim: '\x1b[2m',
+  underscore: '\x1b[4m',
+  blink: '\x1b[5m',
+  reverse: '\x1b[7m',
+  hidden: '\x1b[8m',
+
+  fg: {
+    black: '\x1b[30m',
+    red: '\x1b[31m',
+    green: '\x1b[32m',
+    yellow: '\x1b[33m',
+    blue: '\x1b[34m',
+    magenta: '\x1b[35m',
+    cyan: '\x1b[36m',
+    white: '\x1b[37m',
+    gray: '\x1b[90m',
+    crimson: '\x1b[38m',
+  },
+  bg: {
+    black: '\x1b[40m',
+    red: '\x1b[41m',
+    green: '\x1b[42m',
+    yellow: '\x1b[43m',
+    blue: '\x1b[44m',
+    magenta: '\x1b[45m',
+    cyan: '\x1b[46m',
+    white: '\x1b[47m',
+    gray: '\x1b[100m',
+    crimson: '\x1b[48m',
+  },
+};

invokeai/frontend/web/src/app/hooks/useSocketIO.ts

@@ -19,7 +19,7 @@ declare global {
 }
 
 export const $socketOptions = map<Partial<ManagerOptions & SocketOptions>>({});
-const $isSocketInitialized = atom<boolean>(false);
+export const $isSocketInitialized = atom<boolean>(false);
 
 /**
  * Initializes the socket.io connection and sets up event listeners.

invokeai/frontend/web/src/app/logging/logger.ts

@@ -12,6 +12,7 @@ ROARR.serializeMessage = serializeMessage;
 ROARR.write = createLogWriter();
 
 export const BASE_CONTEXT = {};
+export const log = Roarr.child(BASE_CONTEXT);
 
 export const $logger = atom<Logger>(Roarr.child(BASE_CONTEXT));
 

invokeai/frontend/web/src/app/store/actions.ts

@@ -1,7 +1,10 @@
 import { createAction } from '@reduxjs/toolkit';
 import type { InvokeTabName } from 'features/ui/store/tabMap';
+import type { BatchConfig } from 'services/api/types';
 
 export const enqueueRequested = createAction<{
   tabName: InvokeTabName;
   prepend: boolean;
 }>('app/enqueueRequested');
+
+export const batchEnqueued = createAction<BatchConfig>('app/batchEnqueued');

invokeai/frontend/web/src/app/store/constants.ts

@@ -1,2 +1 @@
 export const STORAGE_PREFIX = '@@invokeai-';
-export const EMPTY_ARRAY = [];

invokeai/frontend/web/src/app/store/createMemoizedSelector.ts

@@ -13,9 +13,19 @@ export const createMemoizedSelector = createSelectorCreator({
   argsMemoize: lruMemoize,
 });
 
+/**
+ * A memoized selector creator that uses LRU cache default shallow equality check.
+ */
+export const createLruSelector = createSelectorCreator({
+  memoize: lruMemoize,
+  argsMemoize: lruMemoize,
+});
+
+export const createLruDraftSafeSelector = createDraftSafeSelectorCreator({
+  memoize: lruMemoize,
+  argsMemoize: lruMemoize,
+});
+
 export const getSelectorsOptions: GetSelectorsOptions = {
-  createSelector: createDraftSafeSelectorCreator({
-    memoize: lruMemoize,
-    argsMemoize: lruMemoize,
-  }),
+  createSelector: createLruDraftSafeSelector,
 };

invokeai/frontend/web/src/app/store/enhancers/reduxRemember/driver.ts

@@ -2,15 +2,15 @@ import { StorageError } from 'app/store/enhancers/reduxRemember/errors';
 import { $projectId } from 'app/store/nanostores/projectId';
 import type { UseStore } from 'idb-keyval';
 import { clear, createStore as createIDBKeyValStore, get, set } from 'idb-keyval';
-import { atom } from 'nanostores';
+import { action, atom } from 'nanostores';
 import type { Driver } from 'redux-remember';
 
 // Create a custom idb-keyval store (just needed to customize the name)
-const $idbKeyValStore = atom<UseStore>(createIDBKeyValStore('invoke', 'invoke-store'));
+export const $idbKeyValStore = atom<UseStore>(createIDBKeyValStore('invoke', 'invoke-store'));
 
-export const clearIdbKeyValStore = () => {
-  clear($idbKeyValStore.get());
-};
+export const clearIdbKeyValStore = action($idbKeyValStore, 'clear', (store) => {
+  clear(store.get());
+});
 
 // Create redux-remember driver, wrapping idb-keyval
 export const idbKeyValDriver: Driver = {

invokeai/frontend/web/src/app/store/enhancers/reduxRemember/errors.ts

@@ -3,7 +3,7 @@ import { parseify } from 'common/util/serialize';
 import { PersistError, RehydrateError } from 'redux-remember';
 import { serializeError } from 'serialize-error';
 
-type StorageErrorArgs = {
+export type StorageErrorArgs = {
   key: string;
   /* eslint-disable-next-line @typescript-eslint/no-explicit-any */ // any is correct
   value?: any;

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -1,65 +1,80 @@
-import type { TypedStartListening } from '@reduxjs/toolkit';
-import { createListenerMiddleware } from '@reduxjs/toolkit';
-import { addCommitStagingAreaImageListener } from 'app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener';
-import { addFirstListImagesListener } from 'app/store/middleware/listenerMiddleware/listeners/addFirstListImagesListener.ts';
-import { addAnyEnqueuedListener } from 'app/store/middleware/listenerMiddleware/listeners/anyEnqueued';
-import { addAppConfigReceivedListener } from 'app/store/middleware/listenerMiddleware/listeners/appConfigReceived';
-import { addAppStartedListener } from 'app/store/middleware/listenerMiddleware/listeners/appStarted';
-import { addBatchEnqueuedListener } from 'app/store/middleware/listenerMiddleware/listeners/batchEnqueued';
-import { addDeleteBoardAndImagesFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/boardAndImagesDeleted';
-import { addBoardIdSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/boardIdSelected';
-import { addBulkDownloadListeners } from 'app/store/middleware/listenerMiddleware/listeners/bulkDownload';
-import { addCanvasCopiedToClipboardListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasCopiedToClipboard';
-import { addCanvasDownloadedAsImageListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasDownloadedAsImage';
-import { addCanvasImageToControlNetListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet';
-import { addCanvasMaskSavedToGalleryListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasMaskSavedToGallery';
-import { addCanvasMaskToControlNetListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet';
-import { addCanvasMergedListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasMerged';
-import { addCanvasSavedToGalleryListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasSavedToGallery';
-import { addControlNetAutoProcessListener } from 'app/store/middleware/listenerMiddleware/listeners/controlNetAutoProcess';
-import { addControlNetImageProcessedListener } from 'app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed';
-import { addEnqueueRequestedCanvasListener } from 'app/store/middleware/listenerMiddleware/listeners/enqueueRequestedCanvas';
-import { addEnqueueRequestedLinear } from 'app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear';
-import { addEnqueueRequestedNodes } from 'app/store/middleware/listenerMiddleware/listeners/enqueueRequestedNodes';
+import type { ListenerEffect, TypedAddListener, TypedStartListening, UnknownAction } from '@reduxjs/toolkit';
+import { addListener, createListenerMiddleware } from '@reduxjs/toolkit';
 import { addGalleryImageClickedListener } from 'app/store/middleware/listenerMiddleware/listeners/galleryImageClicked';
-import { addGetOpenAPISchemaListener } from 'app/store/middleware/listenerMiddleware/listeners/getOpenAPISchema';
-import { addImageAddedToBoardFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/imageAddedToBoard';
-import { addRequestedSingleImageDeletionListener } from 'app/store/middleware/listenerMiddleware/listeners/imageDeleted';
-import { addImageDroppedListener } from 'app/store/middleware/listenerMiddleware/listeners/imageDropped';
-import { addImageRemovedFromBoardFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/imageRemovedFromBoard';
-import { addImagesStarredListener } from 'app/store/middleware/listenerMiddleware/listeners/imagesStarred';
-import { addImagesUnstarredListener } from 'app/store/middleware/listenerMiddleware/listeners/imagesUnstarred';
-import { addImageToDeleteSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/imageToDeleteSelected';
-import { addImageUploadedFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/imageUploaded';
-import { addInitialImageSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/initialImageSelected';
-import { addModelSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/modelSelected';
-import { addModelsLoadedListener } from 'app/store/middleware/listenerMiddleware/listeners/modelsLoaded';
-import { addDynamicPromptsListener } from 'app/store/middleware/listenerMiddleware/listeners/promptChanged';
-import { addSocketConnectedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected';
-import { addSocketDisconnectedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketDisconnected';
-import { addGeneratorProgressEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketGeneratorProgress';
-import { addGraphExecutionStateCompleteEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketGraphExecutionStateComplete';
-import { addInvocationCompleteEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete';
-import { addInvocationErrorEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationError';
-import { addInvocationRetrievalErrorEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationRetrievalError';
-import { addInvocationStartedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationStarted';
-import { addModelInstallEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall';
-import { addModelLoadEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketModelLoad';
-import { addSocketQueueItemStatusChangedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged';
-import { addSessionRetrievalErrorEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketSessionRetrievalError';
-import { addSocketSubscribedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketSubscribed';
-import { addSocketUnsubscribedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketUnsubscribed';
-import { addStagingAreaImageSavedListener } from 'app/store/middleware/listenerMiddleware/listeners/stagingAreaImageSaved';
-import { addUpdateAllNodesRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/updateAllNodesRequested';
-import { addUpscaleRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/upscaleRequested';
-import { addWorkflowLoadRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/workflowLoadRequested';
 import type { AppDispatch, RootState } from 'app/store/store';
 
+import { addCommitStagingAreaImageListener } from './listeners/addCommitStagingAreaImageListener';
+import { addFirstListImagesListener } from './listeners/addFirstListImagesListener.ts';
+import { addAnyEnqueuedListener } from './listeners/anyEnqueued';
+import { addAppConfigReceivedListener } from './listeners/appConfigReceived';
+import { addAppStartedListener } from './listeners/appStarted';
+import { addBatchEnqueuedListener } from './listeners/batchEnqueued';
+import { addDeleteBoardAndImagesFulfilledListener } from './listeners/boardAndImagesDeleted';
+import { addBoardIdSelectedListener } from './listeners/boardIdSelected';
+import { addCanvasCopiedToClipboardListener } from './listeners/canvasCopiedToClipboard';
+import { addCanvasDownloadedAsImageListener } from './listeners/canvasDownloadedAsImage';
+import { addCanvasImageToControlNetListener } from './listeners/canvasImageToControlNet';
+import { addCanvasMaskSavedToGalleryListener } from './listeners/canvasMaskSavedToGallery';
+import { addCanvasMaskToControlNetListener } from './listeners/canvasMaskToControlNet';
+import { addCanvasMergedListener } from './listeners/canvasMerged';
+import { addCanvasSavedToGalleryListener } from './listeners/canvasSavedToGallery';
+import { addControlNetAutoProcessListener } from './listeners/controlNetAutoProcess';
+import { addControlNetImageProcessedListener } from './listeners/controlNetImageProcessed';
+import { addEnqueueRequestedCanvasListener } from './listeners/enqueueRequestedCanvas';
+import { addEnqueueRequestedLinear } from './listeners/enqueueRequestedLinear';
+import { addEnqueueRequestedNodes } from './listeners/enqueueRequestedNodes';
+import { addGetOpenAPISchemaListener } from './listeners/getOpenAPISchema';
+import {
+  addImageAddedToBoardFulfilledListener,
+  addImageAddedToBoardRejectedListener,
+} from './listeners/imageAddedToBoard';
+import {
+  addImageDeletedFulfilledListener,
+  addImageDeletedPendingListener,
+  addImageDeletedRejectedListener,
+  addRequestedMultipleImageDeletionListener,
+  addRequestedSingleImageDeletionListener,
+} from './listeners/imageDeleted';
+import { addImageDroppedListener } from './listeners/imageDropped';
+import {
+  addImageRemovedFromBoardFulfilledListener,
+  addImageRemovedFromBoardRejectedListener,
+} from './listeners/imageRemovedFromBoard';
+import { addImagesStarredListener } from './listeners/imagesStarred';
+import { addImagesUnstarredListener } from './listeners/imagesUnstarred';
+import { addImageToDeleteSelectedListener } from './listeners/imageToDeleteSelected';
+import { addImageUploadedFulfilledListener, addImageUploadedRejectedListener } from './listeners/imageUploaded';
+import { addInitialImageSelectedListener } from './listeners/initialImageSelected';
+import { addModelSelectedListener } from './listeners/modelSelected';
+import { addModelsLoadedListener } from './listeners/modelsLoaded';
+import { addDynamicPromptsListener } from './listeners/promptChanged';
+import { addSocketConnectedEventListener as addSocketConnectedListener } from './listeners/socketio/socketConnected';
+import { addSocketDisconnectedEventListener as addSocketDisconnectedListener } from './listeners/socketio/socketDisconnected';
+import { addGeneratorProgressEventListener as addGeneratorProgressListener } from './listeners/socketio/socketGeneratorProgress';
+import { addGraphExecutionStateCompleteEventListener as addGraphExecutionStateCompleteListener } from './listeners/socketio/socketGraphExecutionStateComplete';
+import { addInvocationCompleteEventListener as addInvocationCompleteListener } from './listeners/socketio/socketInvocationComplete';
+import { addInvocationErrorEventListener as addInvocationErrorListener } from './listeners/socketio/socketInvocationError';
+import { addInvocationRetrievalErrorEventListener } from './listeners/socketio/socketInvocationRetrievalError';
+import { addInvocationStartedEventListener as addInvocationStartedListener } from './listeners/socketio/socketInvocationStarted';
+import { addModelLoadEventListener } from './listeners/socketio/socketModelLoad';
+import { addSocketQueueItemStatusChangedEventListener } from './listeners/socketio/socketQueueItemStatusChanged';
+import { addSessionRetrievalErrorEventListener } from './listeners/socketio/socketSessionRetrievalError';
+import { addSocketSubscribedEventListener as addSocketSubscribedListener } from './listeners/socketio/socketSubscribed';
+import { addSocketUnsubscribedEventListener as addSocketUnsubscribedListener } from './listeners/socketio/socketUnsubscribed';
+import { addStagingAreaImageSavedListener } from './listeners/stagingAreaImageSaved';
+import { addUpdateAllNodesRequestedListener } from './listeners/updateAllNodesRequested';
+import { addUpscaleRequestedListener } from './listeners/upscaleRequested';
+import { addWorkflowLoadRequestedListener } from './listeners/workflowLoadRequested';
+
 export const listenerMiddleware = createListenerMiddleware();
 
 export type AppStartListening = TypedStartListening<RootState, AppDispatch>;
 
-const startAppListening = listenerMiddleware.startListening as AppStartListening;
+export const startAppListening = listenerMiddleware.startListening as AppStartListening;
+
+export const addAppListener = addListener as TypedAddListener<RootState, AppDispatch>;
+
+export type AppListenerEffect = ListenerEffect<UnknownAction, RootState, AppDispatch>;
 
 /**
  * The RTK listener middleware is a lightweight alternative sagas/observables.
@@ -68,88 +83,93 @@ const startAppListening = listenerMiddleware.startListening as AppStartListening
  */
 
 // Image uploaded
-addImageUploadedFulfilledListener(startAppListening);
+addImageUploadedFulfilledListener();
+addImageUploadedRejectedListener();
 
 // Image selected
-addInitialImageSelectedListener(startAppListening);
+addInitialImageSelectedListener();
 
 // Image deleted
-addRequestedSingleImageDeletionListener(startAppListening);
-addDeleteBoardAndImagesFulfilledListener(startAppListening);
-addImageToDeleteSelectedListener(startAppListening);
+addRequestedSingleImageDeletionListener();
+addRequestedMultipleImageDeletionListener();
+addImageDeletedPendingListener();
+addImageDeletedFulfilledListener();
+addImageDeletedRejectedListener();
+addDeleteBoardAndImagesFulfilledListener();
+addImageToDeleteSelectedListener();
 
 // Image starred
-addImagesStarredListener(startAppListening);
-addImagesUnstarredListener(startAppListening);
+addImagesStarredListener();
+addImagesUnstarredListener();
 
 // Gallery
-addGalleryImageClickedListener(startAppListening);
+addGalleryImageClickedListener();
 
 // User Invoked
-addEnqueueRequestedCanvasListener(startAppListening);
-addEnqueueRequestedNodes(startAppListening);
-addEnqueueRequestedLinear(startAppListening);
-addAnyEnqueuedListener(startAppListening);
-addBatchEnqueuedListener(startAppListening);
+addEnqueueRequestedCanvasListener();
+addEnqueueRequestedNodes();
+addEnqueueRequestedLinear();
+addAnyEnqueuedListener();
+addBatchEnqueuedListener();
 
 // Canvas actions
-addCanvasSavedToGalleryListener(startAppListening);
-addCanvasMaskSavedToGalleryListener(startAppListening);
-addCanvasImageToControlNetListener(startAppListening);
-addCanvasMaskToControlNetListener(startAppListening);
-addCanvasDownloadedAsImageListener(startAppListening);
-addCanvasCopiedToClipboardListener(startAppListening);
-addCanvasMergedListener(startAppListening);
-addStagingAreaImageSavedListener(startAppListening);
-addCommitStagingAreaImageListener(startAppListening);
+addCanvasSavedToGalleryListener();
+addCanvasMaskSavedToGalleryListener();
+addCanvasImageToControlNetListener();
+addCanvasMaskToControlNetListener();
+addCanvasDownloadedAsImageListener();
+addCanvasCopiedToClipboardListener();
+addCanvasMergedListener();
+addStagingAreaImageSavedListener();
+addCommitStagingAreaImageListener();
 
 // Socket.IO
-addGeneratorProgressEventListener(startAppListening);
-addGraphExecutionStateCompleteEventListener(startAppListening);
-addInvocationCompleteEventListener(startAppListening);
-addInvocationErrorEventListener(startAppListening);
-addInvocationStartedEventListener(startAppListening);
-addSocketConnectedEventListener(startAppListening);
-addSocketDisconnectedEventListener(startAppListening);
-addSocketSubscribedEventListener(startAppListening);
-addSocketUnsubscribedEventListener(startAppListening);
-addModelLoadEventListener(startAppListening);
-addModelInstallEventListener(startAppListening);
-addSessionRetrievalErrorEventListener(startAppListening);
-addInvocationRetrievalErrorEventListener(startAppListening);
-addSocketQueueItemStatusChangedEventListener(startAppListening);
-addBulkDownloadListeners(startAppListening);
+addGeneratorProgressListener();
+addGraphExecutionStateCompleteListener();
+addInvocationCompleteListener();
+addInvocationErrorListener();
+addInvocationStartedListener();
+addSocketConnectedListener();
+addSocketDisconnectedListener();
+addSocketSubscribedListener();
+addSocketUnsubscribedListener();
+addModelLoadEventListener();
+addSessionRetrievalErrorEventListener();
+addInvocationRetrievalErrorEventListener();
+addSocketQueueItemStatusChangedEventListener();
 
 // ControlNet
-addControlNetImageProcessedListener(startAppListening);
-addControlNetAutoProcessListener(startAppListening);
+addControlNetImageProcessedListener();
+addControlNetAutoProcessListener();
 
 // Boards
-addImageAddedToBoardFulfilledListener(startAppListening);
-addImageRemovedFromBoardFulfilledListener(startAppListening);
-addBoardIdSelectedListener(startAppListening);
+addImageAddedToBoardFulfilledListener();
+addImageAddedToBoardRejectedListener();
+addImageRemovedFromBoardFulfilledListener();
+addImageRemovedFromBoardRejectedListener();
+addBoardIdSelectedListener();
 
 // Node schemas
-addGetOpenAPISchemaListener(startAppListening);
+addGetOpenAPISchemaListener();
 
 // Workflows
-addWorkflowLoadRequestedListener(startAppListening);
-addUpdateAllNodesRequestedListener(startAppListening);
+addWorkflowLoadRequestedListener();
+addUpdateAllNodesRequestedListener();
 
 // DND
-addImageDroppedListener(startAppListening);
+addImageDroppedListener();
 
 // Models
-addModelSelectedListener(startAppListening);
+addModelSelectedListener();
 
 // app startup
-addAppStartedListener(startAppListening);
-addModelsLoadedListener(startAppListening);
-addAppConfigReceivedListener(startAppListening);
-addFirstListImagesListener(startAppListening);
+addAppStartedListener();
+addModelsLoadedListener();
+addAppConfigReceivedListener();
+addFirstListImagesListener();
 
 // Ad-hoc upscale workflwo
-addUpscaleRequestedListener(startAppListening);
+addUpscaleRequestedListener();
 
 // Dynamic prompts
-addDynamicPromptsListener(startAppListening);
+addDynamicPromptsListener();

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener.ts

@@ -1,14 +1,15 @@
 import { isAnyOf } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasBatchIdsReset, commitStagingAreaImage, discardStagedImages } from 'features/canvas/store/canvasSlice';
 import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { queueApi } from 'services/api/endpoints/queue';
 
+import { startAppListening } from '..';
+
 const matcher = isAnyOf(commitStagingAreaImage, discardStagedImages);
 
-export const addCommitStagingAreaImageListener = (startAppListening: AppStartListening) => {
+export const addCommitStagingAreaImageListener = () => {
   startAppListening({
     matcher,
     effect: async (_, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addFirstListImagesListener.ts.ts

@@ -1,11 +1,15 @@
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
+import { createAction } from '@reduxjs/toolkit';
 import { imageSelected } from 'features/gallery/store/gallerySlice';
 import { IMAGE_CATEGORIES } from 'features/gallery/store/types';
 import { imagesApi } from 'services/api/endpoints/images';
 import type { ImageCache } from 'services/api/types';
 import { getListImagesUrl, imagesSelectors } from 'services/api/util';
 
-export const addFirstListImagesListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const appStarted = createAction('app/appStarted');
+
+export const addFirstListImagesListener = () => {
   startAppListening({
     matcher: imagesApi.endpoints.listImages.matchFulfilled,
     effect: async (action, { dispatch, unsubscribe, cancelActiveListeners }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/anyEnqueued.ts

@@ -1,7 +1,8 @@
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { queueApi, selectQueueStatus } from 'services/api/endpoints/queue';
 
-export const addAnyEnqueuedListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addAnyEnqueuedListener = () => {
   startAppListening({
     matcher: queueApi.endpoints.enqueueBatch.matchFulfilled,
     effect: async (_, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/appConfigReceived.ts

@@ -1,9 +1,10 @@
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { setInfillMethod } from 'features/parameters/store/generationSlice';
 import { shouldUseNSFWCheckerChanged, shouldUseWatermarkerChanged } from 'features/system/store/systemSlice';
 import { appInfoApi } from 'services/api/endpoints/appInfo';
 
-export const addAppConfigReceivedListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addAppConfigReceivedListener = () => {
   startAppListening({
     matcher: appInfoApi.endpoints.getAppConfig.matchFulfilled,
     effect: async (action, { getState, dispatch }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/appStarted.ts

@@ -1,9 +1,10 @@
 import { createAction } from '@reduxjs/toolkit';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
+
+import { startAppListening } from '..';
 
 export const appStarted = createAction('app/appStarted');
 
-export const addAppStartedListener = (startAppListening: AppStartListening) => {
+export const addAppStartedListener = () => {
   startAppListening({
     actionCreator: appStarted,
     effect: async (action, { unsubscribe, cancelActiveListeners }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/batchEnqueued.ts

@@ -1,13 +1,19 @@
+import { createStandaloneToast, theme, TOAST_OPTIONS } from '@invoke-ai/ui-library';
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { parseify } from 'common/util/serialize';
-import { toast } from 'common/util/toast';
 import { zPydanticValidationError } from 'features/system/store/zodSchemas';
 import { t } from 'i18next';
 import { truncate, upperFirst } from 'lodash-es';
 import { queueApi } from 'services/api/endpoints/queue';
 
-export const addBatchEnqueuedListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+const { toast } = createStandaloneToast({
+  theme: theme,
+  defaultOptions: TOAST_OPTIONS.defaultOptions,
+});
+
+export const addBatchEnqueuedListener = () => {
   // success
   startAppListening({
     matcher: queueApi.endpoints.enqueueBatch.matchFulfilled,

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardAndImagesDeleted.ts

@@ -1,4 +1,3 @@
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { resetCanvas } from 'features/canvas/store/canvasSlice';
 import { controlAdaptersReset } from 'features/controlAdapters/store/controlAdaptersSlice';
 import { getImageUsage } from 'features/deleteImageModal/store/selectors';
@@ -6,7 +5,9 @@ import { nodeEditorReset } from 'features/nodes/store/nodesSlice';
 import { clearInitialImage } from 'features/parameters/store/generationSlice';
 import { imagesApi } from 'services/api/endpoints/images';
 
-export const addDeleteBoardAndImagesFulfilledListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addDeleteBoardAndImagesFulfilledListener = () => {
   startAppListening({
     matcher: imagesApi.endpoints.deleteBoardAndImages.matchFulfilled,
     effect: async (action, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts

@@ -1,11 +1,12 @@
 import { isAnyOf } from '@reduxjs/toolkit';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { boardIdSelected, galleryViewChanged, imageSelected } from 'features/gallery/store/gallerySlice';
 import { ASSETS_CATEGORIES, IMAGE_CATEGORIES } from 'features/gallery/store/types';
 import { imagesApi } from 'services/api/endpoints/images';
 import { imagesSelectors } from 'services/api/util';
 
-export const addBoardIdSelectedListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addBoardIdSelectedListener = () => {
   startAppListening({
     matcher: isAnyOf(boardIdSelected, galleryViewChanged),
     effect: async (action, { getState, dispatch, condition, cancelActiveListeners }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/bulkDownload.tsx

@@ -1,114 +0,0 @@
-import type { UseToastOptions } from '@invoke-ai/ui-library';
-import { ExternalLink } from '@invoke-ai/ui-library';
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { toast } from 'common/util/toast';
-import { t } from 'i18next';
-import { imagesApi } from 'services/api/endpoints/images';
-import {
-  socketBulkDownloadCompleted,
-  socketBulkDownloadFailed,
-  socketBulkDownloadStarted,
-} from 'services/events/actions';
-
-const log = logger('images');
-
-export const addBulkDownloadListeners = (startAppListening: AppStartListening) => {
-  startAppListening({
-    matcher: imagesApi.endpoints.bulkDownloadImages.matchFulfilled,
-    effect: async (action) => {
-      log.debug(action.payload, 'Bulk download requested');
-
-      // If we have an item name, we are processing the bulk download locally and should use it as the toast id to
-      // prevent multiple toasts for the same item.
-      toast({
-        id: action.payload.bulk_download_item_name ?? undefined,
-        title: t('gallery.bulkDownloadRequested'),
-        status: 'success',
-        // Show the response message if it exists, otherwise show the default message
-        description: action.payload.response || t('gallery.bulkDownloadRequestedDesc'),
-        duration: null,
-        isClosable: true,
-      });
-    },
-  });
-
-  startAppListening({
-    matcher: imagesApi.endpoints.bulkDownloadImages.matchRejected,
-    effect: async () => {
-      log.debug('Bulk download request failed');
-
-      // There isn't any toast to update if we get this event.
-      toast({
-        title: t('gallery.bulkDownloadRequestFailed'),
-        status: 'success',
-        isClosable: true,
-      });
-    },
-  });
-
-  startAppListening({
-    actionCreator: socketBulkDownloadStarted,
-    effect: async (action) => {
-      // This should always happen immediately after the bulk download request, so we don't need to show a toast here.
-      log.debug(action.payload.data, 'Bulk download preparation started');
-    },
-  });
-
-  startAppListening({
-    actionCreator: socketBulkDownloadCompleted,
-    effect: async (action) => {
-      log.debug(action.payload.data, 'Bulk download preparation completed');
-
-      const { bulk_download_item_name } = action.payload.data;
-
-      // TODO(psyche): This URL may break in in some environments (e.g. Nvidia workbench) but we need to test it first
-      const url = `/api/v1/images/download/${bulk_download_item_name}`;
-
-      const toastOptions: UseToastOptions = {
-        id: bulk_download_item_name,
-        title: t('gallery.bulkDownloadReady', 'Download ready'),
-        status: 'success',
-        description: (
-          <ExternalLink
-            label={t('gallery.clickToDownload', 'Click here to download')}
-            href={url}
-            download={bulk_download_item_name}
-          />
-        ),
-        duration: null,
-        isClosable: true,
-      };
-
-      if (toast.isActive(bulk_download_item_name)) {
-        toast.update(bulk_download_item_name, toastOptions);
-      } else {
-        toast(toastOptions);
-      }
-    },
-  });
-
-  startAppListening({
-    actionCreator: socketBulkDownloadFailed,
-    effect: async (action) => {
-      log.debug(action.payload.data, 'Bulk download preparation failed');
-
-      const { bulk_download_item_name } = action.payload.data;
-
-      const toastOptions: UseToastOptions = {
-        id: bulk_download_item_name,
-        title: t('gallery.bulkDownloadFailed'),
-        status: 'error',
-        description: action.payload.data.error,
-        duration: null,
-        isClosable: true,
-      };
-
-      if (toast.isActive(bulk_download_item_name)) {
-        toast.update(bulk_download_item_name, toastOptions);
-      } else {
-        toast(toastOptions);
-      }
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasCopiedToClipboard.ts

@@ -1,12 +1,13 @@
 import { $logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasCopiedToClipboard } from 'features/canvas/store/actions';
 import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
 import { addToast } from 'features/system/store/systemSlice';
 import { copyBlobToClipboard } from 'features/system/util/copyBlobToClipboard';
 import { t } from 'i18next';
 
-export const addCanvasCopiedToClipboardListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addCanvasCopiedToClipboardListener = () => {
   startAppListening({
     actionCreator: canvasCopiedToClipboard,
     effect: async (action, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasDownloadedAsImage.ts

@@ -1,12 +1,13 @@
 import { $logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasDownloadedAsImage } from 'features/canvas/store/actions';
 import { downloadBlob } from 'features/canvas/util/downloadBlob';
 import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
 import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 
-export const addCanvasDownloadedAsImageListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addCanvasDownloadedAsImageListener = () => {
   startAppListening({
     actionCreator: canvasDownloadedAsImage,
     effect: async (action, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet.ts

@@ -1,5 +1,4 @@
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasImageToControlAdapter } from 'features/canvas/store/actions';
 import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
 import { controlAdapterImageChanged } from 'features/controlAdapters/store/controlAdaptersSlice';
@@ -7,7 +6,9 @@ import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { imagesApi } from 'services/api/endpoints/images';
 
-export const addCanvasImageToControlNetListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addCanvasImageToControlNetListener = () => {
   startAppListening({
     actionCreator: canvasImageToControlAdapter,
     effect: async (action, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskSavedToGallery.ts

@@ -1,12 +1,13 @@
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasMaskSavedToGallery } from 'features/canvas/store/actions';
 import { getCanvasData } from 'features/canvas/util/getCanvasData';
 import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { imagesApi } from 'services/api/endpoints/images';
 
-export const addCanvasMaskSavedToGalleryListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addCanvasMaskSavedToGalleryListener = () => {
   startAppListening({
     actionCreator: canvasMaskSavedToGallery,
     effect: async (action, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet.ts

@@ -1,5 +1,4 @@
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasMaskToControlAdapter } from 'features/canvas/store/actions';
 import { getCanvasData } from 'features/canvas/util/getCanvasData';
 import { controlAdapterImageChanged } from 'features/controlAdapters/store/controlAdaptersSlice';
@@ -7,7 +6,9 @@ import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { imagesApi } from 'services/api/endpoints/images';
 
-export const addCanvasMaskToControlNetListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addCanvasMaskToControlNetListener = () => {
   startAppListening({
     actionCreator: canvasMaskToControlAdapter,
     effect: async (action, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMerged.ts

@@ -1,5 +1,4 @@
 import { $logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasMerged } from 'features/canvas/store/actions';
 import { $canvasBaseLayer } from 'features/canvas/store/canvasNanostore';
 import { setMergedCanvas } from 'features/canvas/store/canvasSlice';
@@ -8,7 +7,9 @@ import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { imagesApi } from 'services/api/endpoints/images';
 
-export const addCanvasMergedListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addCanvasMergedListener = () => {
   startAppListening({
     actionCreator: canvasMerged,
     effect: async (action, { dispatch }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasSavedToGallery.ts

@@ -1,12 +1,13 @@
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { canvasSavedToGallery } from 'features/canvas/store/actions';
 import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
 import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { imagesApi } from 'services/api/endpoints/images';
 
-export const addCanvasSavedToGalleryListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addCanvasSavedToGalleryListener = () => {
   startAppListening({
     actionCreator: canvasSavedToGallery,
     effect: async (action, { dispatch, getState }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetAutoProcess.ts

@@ -1,6 +1,5 @@
 import type { AnyListenerPredicate } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import type { RootState } from 'app/store/store';
 import { controlAdapterImageProcessed } from 'features/controlAdapters/store/actions';
 import {
@@ -13,6 +12,8 @@ import {
 } from 'features/controlAdapters/store/controlAdaptersSlice';
 import { isControlNetOrT2IAdapter } from 'features/controlAdapters/store/types';
 
+import { startAppListening } from '..';
+
 type AnyControlAdapterParamChangeAction =
   | ReturnType<typeof controlAdapterProcessorParamsChanged>
   | ReturnType<typeof controlAdapterModelChanged>
@@ -66,7 +67,7 @@ const DEBOUNCE_MS = 300;
  *
  * The network request is debounced.
  */
-export const addControlNetAutoProcessListener = (startAppListening: AppStartListening) => {
+export const addControlNetAutoProcessListener = () => {
   startAppListening({
     predicate,
     effect: async (action, { dispatch, cancelActiveListeners, delay }) => {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed.ts

@@ -1,5 +1,4 @@
 import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { parseify } from 'common/util/serialize';
 import { controlAdapterImageProcessed } from 'features/controlAdapters/store/actions';
 import {
@@ -17,7 +16,9 @@ import { queueApi } from 'services/api/endpoints/queue';
 import type { BatchConfig, ImageDTO } from 'services/api/types';
 import { socketInvocationComplete } from 'services/events/actions';
 
-export const addControlNetImageProcessedListener = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addControlNetImageProcessedListener = () => {
   startAppListening({
     actionCreator: controlAdapterImageProcessed,
     effect: async (action, { dispatch, getState, take }) => {
@@ -50,7 +51,6 @@ export const addControlNetImageProcessedListener = (startAppListening: AppStartL
                 image: { image_name: ca.controlImage },
               },
             },
-            edges: [],
           },
           runs: 1,
         },

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedCanvas.ts

@@ -1,6 +1,5 @@
 import { logger } from 'app/logging/logger';
 import { enqueueRequested } from 'app/store/actions';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import openBase64ImageInTab from 'common/util/openBase64ImageInTab';
 import { parseify } from 'common/util/serialize';
 import { canvasBatchIdAdded, stagingAreaInitialized } from 'features/canvas/store/canvasSlice';
@@ -14,6 +13,8 @@ import { imagesApi } from 'services/api/endpoints/images';
 import { queueApi } from 'services/api/endpoints/queue';
 import type { ImageDTO } from 'services/api/types';
 
+import { startAppListening } from '..';
+
 /**
  * This listener is responsible invoking the canvas. This involves a number of steps:
  *
@@ -27,7 +28,7 @@ import type { ImageDTO } from 'services/api/types';
  * 8. Initialize the staging area if not yet initialized
  * 9. Dispatch the sessionReadyToInvoke action to invoke the session
  */
-export const addEnqueueRequestedCanvasListener = (startAppListening: AppStartListening) => {
+export const addEnqueueRequestedCanvasListener = () => {
   startAppListening({
     predicate: (action): action is ReturnType<typeof enqueueRequested> =>
       enqueueRequested.match(action) && action.payload.tabName === 'unifiedCanvas',

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts

@@ -1,5 +1,4 @@
 import { enqueueRequested } from 'app/store/actions';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
 import { buildLinearImageToImageGraph } from 'features/nodes/util/graph/buildLinearImageToImageGraph';
 import { buildLinearSDXLImageToImageGraph } from 'features/nodes/util/graph/buildLinearSDXLImageToImageGraph';
@@ -7,7 +6,9 @@ import { buildLinearSDXLTextToImageGraph } from 'features/nodes/util/graph/build
 import { buildLinearTextToImageGraph } from 'features/nodes/util/graph/buildLinearTextToImageGraph';
 import { queueApi } from 'services/api/endpoints/queue';
 
-export const addEnqueueRequestedLinear = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addEnqueueRequestedLinear = () => {
   startAppListening({
     predicate: (action): action is ReturnType<typeof enqueueRequested> =>
       enqueueRequested.match(action) && (action.payload.tabName === 'txt2img' || action.payload.tabName === 'img2img'),

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedNodes.ts

@@ -1,11 +1,12 @@
 import { enqueueRequested } from 'app/store/actions';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { buildNodesGraph } from 'features/nodes/util/graph/buildNodesGraph';
 import { buildWorkflowWithValidation } from 'features/nodes/util/workflow/buildWorkflow';
 import { queueApi } from 'services/api/endpoints/queue';
 import type { BatchConfig } from 'services/api/types';
 
-export const addEnqueueRequestedNodes = (startAppListening: AppStartListening) => {
+import { startAppListening } from '..';
+
+export const addEnqueueRequestedNodes = () => {
   startAppListening({
     predicate: (action): action is ReturnType<typeof enqueueRequested> =>
       enqueueRequested.match(action) && action.payload.tabName === 'nodes',