{release} 3.6.3rc1

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

docker/Dockerfile

@@ -18,8 +18,8 @@ ENV INVOKEAI_SRC=/opt/invokeai
 ENV VIRTUAL_ENV=/opt/venv/invokeai
 
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
-ARG TORCH_VERSION=2.1.2
-ARG TORCHVISION_VERSION=0.16.2
+ARG TORCH_VERSION=2.1.0
+ARG TORCHVISION_VERSION=0.16
 ARG GPU_DRIVER=cuda
 ARG TARGETPLATFORM="linux/amd64"
 # unused but available
@@ -35,7 +35,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then \
         extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cpu"; \
     elif [ "$GPU_DRIVER" = "rocm" ]; then \
-        extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/rocm5.6"; \
+        extra_index_url_arg="--index-url https://download.pytorch.org/whl/rocm5.6"; \
     else \
         extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cu121"; \
     fi &&\
@@ -54,7 +54,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     if [ "$GPU_DRIVER" = "cuda" ] && [ "$TARGETPLATFORM" = "linux/amd64" ]; then \
         pip install -e ".[xformers]"; \
     else \
-        pip install $extra_index_url_arg -e "."; \
+        pip install -e "."; \
     fi
 
 # #### Build the Web UI ------------------------------------

docker/README.md

@@ -28,7 +28,7 @@ This is done via Docker Desktop preferences
 
 ### Configure Invoke environment
 
-1. Make a copy of `.env.sample` and name it `.env` (`cp .env.sample .env` (Mac/Linux) or `copy example.env .env` (Windows)). Make changes as necessary. Set `INVOKEAI_ROOT` to an absolute path to:
+1. Make a copy of `env.sample` and name it `.env` (`cp env.sample .env` (Mac/Linux) or `copy example.env .env` (Windows)). Make changes as necessary. Set `INVOKEAI_ROOT` to an absolute path to:
     a. the desired location of the InvokeAI runtime directory, or
     b. an existing, v3.0.0 compatible runtime directory.
 1. Execute `run.sh`

docker/run.sh

@@ -21,7 +21,7 @@ run() {
     printf "%s\n" "$build_args"
   fi
 
-  docker compose build $build_args $service_name
+  docker compose build $build_args
   unset build_args
 
   printf "%s\n" "starting service $service_name"

docs/contributing/INVOCATIONS.md

@@ -9,15 +9,11 @@ complex functionality.
 
 ## Invocations Directory
 
-InvokeAI Nodes can be found in the `invokeai/app/invocations` directory. These
-can be used as examples to create your own nodes.
+InvokeAI Nodes can be found in the `invokeai/app/invocations` directory. These can be used as examples to create your own nodes.
 
-New nodes should be added to a subfolder in `nodes` direction found at the root
-level of the InvokeAI installation location. Nodes added to this folder will be
-able to be used upon application startup.
-
-Example `nodes` subfolder structure:
+New nodes should be added to a subfolder in `nodes` direction found at the root level of the InvokeAI installation location. Nodes added to this folder will be able to be used upon application startup. 
 
+Example `nodes`  subfolder structure:
 ```py
 ├── __init__.py # Invoke-managed custom node loader
 │
@@ -34,14 +30,14 @@ Example `nodes` subfolder structure:
         └── fancy_node.py
 ```
 
-Each node folder must have an `__init__.py` file that imports its nodes. Only
-nodes imported in the `__init__.py` file are loaded. See the README in the nodes
-folder for more examples:
+Each node folder must have an `__init__.py` file that imports its nodes. Only nodes imported in the `__init__.py` file are loaded.
+ See the README in the nodes folder for more examples: 
 
 ```py
 from .cool_node import CoolInvocation
 ```
 
+
 ## Creating A New Invocation
 
 In order to understand the process of creating a new Invocation, let us actually
@@ -135,6 +131,7 @@ from invokeai.app.invocations.primitives import ImageField
 class ResizeInvocation(BaseInvocation):
     '''Resizes an image'''
 
+    # Inputs
     image: ImageField = InputField(description="The input image")
     width: int = InputField(default=512, ge=64, le=2048, description="Width of the new image")
     height: int = InputField(default=512, ge=64, le=2048, description="Height of the new image")
@@ -170,6 +167,7 @@ from invokeai.app.invocations.primitives import ImageField
 class ResizeInvocation(BaseInvocation):
     '''Resizes an image'''
 
+    # Inputs
     image: ImageField = InputField(description="The input image")
     width: int = InputField(default=512, ge=64, le=2048, description="Width of the new image")
     height: int = InputField(default=512, ge=64, le=2048, description="Height of the new image")
@@ -199,6 +197,7 @@ from invokeai.app.invocations.image import ImageOutput
 class ResizeInvocation(BaseInvocation):
     '''Resizes an image'''
 
+    # Inputs
     image: ImageField = InputField(description="The input image")
     width: int = InputField(default=512, ge=64, le=2048, description="Width of the new image")
     height: int = InputField(default=512, ge=64, le=2048, description="Height of the new image")
@@ -230,17 +229,30 @@ class ResizeInvocation(BaseInvocation):
     height: int = InputField(default=512, ge=64, le=2048, description="Height of the new image")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        # Load the input image as a PIL image
-        image = context.images.get_pil(self.image.image_name)
+        # Load the image using InvokeAI's predefined Image Service. Returns the PIL image.
+        image = context.services.images.get_pil_image(self.image.image_name)
 
-        # Resize the image
+        # Resizing the image
         resized_image = image.resize((self.width, self.height))
 
-        # Save the image
-        image_dto = context.images.save(image=resized_image)
+        # Save the image using InvokeAI's predefined Image Service. Returns the prepared PIL image.
+        output_image = context.services.images.create(
+            image=resized_image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+        )
 
-        # Return an ImageOutput
-        return ImageOutput.build(image_dto)
+        # Returning the Image
+        return ImageOutput(
+            image=ImageField(
+                image_name=output_image.image_name,
+            ),
+            width=output_image.width,
+            height=output_image.height,
+        )
 ```
 
 **Note:** Do not be overwhelmed by the `ImageOutput` process. InvokeAI has a
@@ -331,25 +343,27 @@ class ImageColorStringOutput(BaseInvocationOutput):
 
 That's all there is to it.
 
+<!-- TODO: DANGER - we probably do not want people to create their own field types, because this requires a lot of work on the frontend to accomodate.
+
 ### Custom Input Fields
 
 Now that you know how to create your own Invocations, let us dive into slightly
 more advanced topics.
 
 While creating your own Invocations, you might run into a scenario where the
-existing fields in InvokeAI do not meet your requirements. In such cases, you
-can create your own fields.
+existing input types in InvokeAI do not meet your requirements. In such cases,
+you can create your own input types.
 
 Let us create one as an example. Let us say we want to create a color input
 field that represents a color code. But before we start on that here are some
 general good practices to keep in mind.
 
-### Best Practices
+**Good Practices**
 
 - There is no naming convention for input fields but we highly recommend that
   you name it something appropriate like `ColorField`.
 - It is not mandatory but it is heavily recommended to add a relevant
-  `docstring` to describe your field.
+  `docstring` to describe your input field.
 - Keep your field in the same file as the Invocation that it is made for or in
   another file where it is relevant.
 
@@ -364,13 +378,10 @@ class ColorField(BaseModel):
     pass
 ```
 
-Perfect. Now let us create the properties for our field. This is similar to how
-you created input fields for your Invocation. All the same rules apply. Let us
-create four fields representing the _red(r)_, _blue(b)_, _green(g)_ and
-_alpha(a)_ channel of the color.
-
-> Technically, the properties are _also_ called fields - but in this case, it
-> refers to a `pydantic` field.
+Perfect. Now let us create our custom inputs for our field. This is exactly
+similar how you created input fields for your Invocation. All the same rules
+apply. Let us create four fields representing the _red(r)_, _blue(b)_,
+_green(g)_ and _alpha(a)_ channel of the color.
 
 ```python
 class ColorField(BaseModel):
@@ -385,11 +396,25 @@ That's it. We now have a new input field type that we can use in our Invocations
 like this.
 
 ```python
-color: ColorField = InputField(default=ColorField(r=0, g=0, b=0, a=0), description='Background color of an image')
+color: ColorField = Field(default=ColorField(r=0, g=0, b=0, a=0), description='Background color of an image')
 ```
 
-### Using the custom field
+### Custom Components For Frontend
 
-When you start the UI, your custom field will be automatically recognized.
+Every backend input type should have a corresponding frontend component so the
+UI knows what to render when you use a particular field type.
 
-Custom fields only support connection inputs in the Workflow Editor.
+If you are using existing field types, we already have components for those. So
+you don't have to worry about creating anything new. But this might not always
+be the case. Sometimes you might want to create new field types and have the
+frontend UI deal with it in a different way.
+
+This is where we venture into the world of React and Javascript and create our
+own new components for our Invocations. Do not fear the world of JS. It's
+actually pretty straightforward.
+
+Let us create a new component for our custom color field we created above. When
+we use a color field, let us say we want the UI to display a color picker for
+the user to pick from rather than entering values. That is what we will build
+now.
+-->

docs/contributing/MODEL_MANAGER.md

@@ -28,7 +28,7 @@ model. These are the:
   Hugging Face, as well as discriminating among model versions in
   Civitai, but can be used for arbitrary content.
   
-  * _ModelLoadServiceBase_
+  * _ModelLoadServiceBase_ (**CURRENTLY UNDER DEVELOPMENT - NOT IMPLEMENTED**)
   Responsible for loading a model from disk
   into RAM and VRAM and getting it ready for inference.
 
@@ -41,10 +41,10 @@ The four main services can be found in
 * `invokeai/app/services/model_records/`
 * `invokeai/app/services/model_install/`
 * `invokeai/app/services/downloads/`
-* `invokeai/app/services/model_load/`
+* `invokeai/app/services/model_loader/` (**under development**)
 
 Code related to the FastAPI web API can be found in
-`invokeai/app/api/routers/model_manager_v2.py`.
+`invokeai/app/api/routers/model_records.py`.
 
 ***
 
@@ -84,10 +84,10 @@ diffusers model. When this happens, `original_hash` is unchanged, but
 `ModelType`, `ModelFormat` and `BaseModelType` are string enums that
 are defined in `invokeai.backend.model_manager.config`. They are also
 imported by, and can be reexported from,
-`invokeai.app.services.model_manager.model_records`:
+`invokeai.app.services.model_record_service`:
 
 ```
-from invokeai.app.services.model_records import ModelType, ModelFormat, BaseModelType
+from invokeai.app.services.model_record_service import ModelType, ModelFormat, BaseModelType
 ```
 
 The `path` field can be absolute or relative. If relative, it is taken
@@ -123,7 +123,7 @@ taken to be the `models_dir` directory.
 
 `variant` is an enumerated string class with values `normal`,
 `inpaint` and `depth`. If needed, it can be imported if needed from
-either `invokeai.app.services.model_records` or
+either `invokeai.app.services.model_record_service` or
 `invokeai.backend.model_manager.config`.
 
 ### ONNXSD2Config
@@ -134,7 +134,7 @@ either `invokeai.app.services.model_records` or
 | `upcast_attention`  | bool | Model requires its attention module to be upcast |
 
 The `SchedulerPredictionType` enum can be imported from either
-`invokeai.app.services.model_records` or
+`invokeai.app.services.model_record_service` or
 `invokeai.backend.model_manager.config`.
 
 ### Other config classes
@@ -157,6 +157,15 @@ indicates that the model is compatible with any of the base
 models. This works OK for some models, such as the IP Adapter image
 encoders, but is an all-or-nothing proposition.
 
+Another issue is that the config class hierarchy is paralleled to some
+extent by a `ModelBase` class hierarchy defined in
+`invokeai.backend.model_manager.models.base` and its subclasses. These
+are classes representing the models after they are loaded into RAM and
+include runtime information such as load status and bytes used. Some
+of the fields, including `name`, `model_type` and `base_model`, are
+shared between `ModelConfigBase` and `ModelBase`, and this is a
+potential source of confusion.
+
 ## Reading and Writing Model Configuration Records
 
 The `ModelRecordService` provides the ability to retrieve model
@@ -168,11 +177,11 @@ initialization and can be retrieved within an invocation from the
 `InvocationContext` object:
 
 ```
-store = context.services.model_manager.store
+store = context.services.model_record_store
 ```
 
 or from elsewhere in the code by accessing
-`ApiDependencies.invoker.services.model_manager.store`.
+`ApiDependencies.invoker.services.model_record_store`.
 
 ### Creating a `ModelRecordService` 
 
@@ -181,7 +190,7 @@ you can directly create either a `ModelRecordServiceSQL` or a
 `ModelRecordServiceFile` object:
 
 ```
-from invokeai.app.services.model_records import ModelRecordServiceSQL, ModelRecordServiceFile
+from invokeai.app.services.model_record_service import ModelRecordServiceSQL, ModelRecordServiceFile
 
 store = ModelRecordServiceSQL.from_connection(connection, lock)
 store = ModelRecordServiceSQL.from_db_file('/path/to/sqlite_database.db')
@@ -243,7 +252,7 @@ So a typical startup pattern would be:
 ```
 import sqlite3
 from invokeai.app.services.thread import lock
-from invokeai.app.services.model_records import ModelRecordServiceBase
+from invokeai.app.services.model_record_service import ModelRecordServiceBase
 from invokeai.app.services.config import InvokeAIAppConfig
 
 config = InvokeAIAppConfig.get_config()
@@ -251,6 +260,19 @@ db_conn = sqlite3.connect(config.db_path.as_posix(), check_same_thread=False)
 store = ModelRecordServiceBase.open(config, db_conn, lock)
 ```
 
+_A note on simultaneous access to `invokeai.db`_: The current InvokeAI
+service architecture for the image and graph databases is careful to
+use a shared sqlite3 connection and a thread lock to ensure that two
+threads don't attempt to access the database simultaneously. However,
+the default `sqlite3` library used by Python reports using
+**Serialized** mode, which allows multiple threads to access the
+database simultaneously using multiple database connections (see
+https://www.sqlite.org/threadsafe.html and
+https://ricardoanderegg.com/posts/python-sqlite-thread-safety/). Therefore
+it should be safe to allow the record service to open its own SQLite
+database connection. Opening a model record service should then be as
+simple as `ModelRecordServiceBase.open(config)`.
+
 ### Fetching a Model's Configuration from `ModelRecordServiceBase`
 
 Configurations can be retrieved in several ways.
@@ -446,44 +468,6 @@ required parameters:
 
 Once initialized, the installer will provide the following methods:
 
-#### install_job = installer.heuristic_import(source, [config], [access_token])
-
-This is a simplified interface to the installer which takes a source
-string, an optional model configuration dictionary and an optional
-access token.
-
-The `source` is a string that can be any of these forms
-
-1. A path on the local filesystem (`C:\\users\\fred\\model.safetensors`)
-2. A Url pointing to a single downloadable model file (`https://civitai.com/models/58390/detail-tweaker-lora-lora`)
-3. A HuggingFace repo_id with any of the following formats:
-   - `model/name` -- entire model
-   - `model/name:fp32` -- entire model, using the fp32 variant
-   - `model/name:fp16:vae` -- vae submodel, using the fp16 variant
-   - `model/name::vae` -- vae submodel, using default precision
-   - `model/name:fp16:path/to/model.safetensors` -- an individual model file, fp16 variant
-   - `model/name::path/to/model.safetensors` -- an individual model file, default variant
-
-Note that by specifying a relative path to the top of the HuggingFace
-repo, you can download and install arbitrary models files. 
-
-The variant, if not provided, will be automatically filled in with
-`fp32` if the user has requested full precision, and `fp16`
-otherwise. If a variant that does not exist is requested, then the
-method will install whatever HuggingFace returns as its default
-revision.
-
-`config` is an optional dict of values that will override the
-autoprobed values for model type, base, scheduler prediction type, and
-so forth. See [Model configuration and
-probing](#Model-configuration-and-probing) for details.
-
-`access_token` is an optional access token for accessing resources
-that need authentication.
-
-The method will return a `ModelInstallJob`. This object is discussed
-at length in the following section.
-
 #### install_job = installer.import_model()
 
 The `import_model()` method is the core of the installer. The
@@ -502,10 +486,9 @@ source2 = LocalModelSource(path='/opt/models/sushi_diffusers')     # a local dif
 source3 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5')  # a repo_id
 source4 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5', subfolder='vae')  # a subfolder within a repo_id
 source5 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5', variant='fp16')   # a named variant of a HF model
-source6 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5', subfolder='OrangeMix/OrangeMix1.ckpt')   # path to an individual model file
 
-source7 = URLModelSource(url='https://civitai.com/api/download/models/63006')       # model located at a URL
-source8 = URLModelSource(url='https://civitai.com/api/download/models/63006', access_token='letmein') # with an access token
+source6 = URLModelSource(url='https://civitai.com/api/download/models/63006')       # model located at a URL
+source7 = URLModelSource(url='https://civitai.com/api/download/models/63006', access_token='letmein') # with an access token
 
 for source in [source1, source2, source3, source4, source5, source6, source7]:
    install_job = installer.install_model(source)
@@ -561,6 +544,7 @@ can be passed to `import_model()`.
 attributes returned by the model prober. See the section below for
 details.
 
+
 #### LocalModelSource
 
 This is used for a model that is located on a locally-accessible Posix
@@ -753,7 +737,7 @@ and `cancelled`, as well as `in_terminal_state`. The last will return
 True if the job is in the complete, errored or cancelled states.
 
 
-#### Model configuration and probing
+#### Model confguration and probing
 
 The install service uses the `invokeai.backend.model_manager.probe`
 module during import to determine the model's type, base type, and
@@ -792,14 +776,6 @@ returns a list of completed jobs. The optional `timeout` argument will
 return from the call if jobs aren't completed in the specified
 time. An argument of 0 (the default) will block indefinitely.
 
-#### jobs = installer.wait_for_job(job, [timeout])
-
-Like `wait_for_installs()`, but block until a specific job has
-completed or errored, and then return the job.  The optional `timeout`
-argument will return from the call if the job doesn't complete in the
-specified time. An argument of 0 (the default) will block
-indefinitely.
-
 #### jobs = installer.list_jobs()
 
 Return a list of all active and complete `ModelInstallJobs`.
@@ -862,31 +838,6 @@ This method is similar to `unregister()`, but also unconditionally
 deletes the corresponding model weights file(s), regardless of whether
 they are inside or outside the InvokeAI models hierarchy.
 
-
-#### path = installer.download_and_cache(remote_source, [access_token], [timeout])
-
-This utility routine will download the model file located at source,
-cache it, and return the path to the cached file. It does not attempt
-to determine the model type, probe its configuration values, or
-register it with the models database.
-
-You may provide an access token if the remote source requires
-authorization. The call will block indefinitely until the file is
-completely downloaded, cancelled or raises an error of some sort. If
-you provide a timeout (in seconds), the call will raise a
-`TimeoutError` exception if the download hasn't completed in the
-specified period.
-
-You may use this mechanism to request any type of file, not just a
-model. The file will be stored in a subdirectory of
-`INVOKEAI_ROOT/models/.cache`. If the requested file is found in the
-cache, its path will be returned without redownloading it.
-
-Be aware that the models cache is cleared of infrequently-used files
-and directories at regular intervals when the size of the cache
-exceeds the value specified in Invoke's `convert_cache` configuration
-variable.
-
 #### List[str]=installer.scan_directory(scan_dir: Path, install: bool)
 
 This method will recursively scan the directory indicated in
@@ -1177,7 +1128,7 @@ job = queue.create_download_job(
 		 event_handlers=[my_handler1, my_handler2], # if desired
 		 start=True,
 	)
-```
+ ```
  
 The `filename` argument forces the downloader to use the specified
 name for the file rather than the name provided by the remote source,
@@ -1220,13 +1171,6 @@ queue or was not created by this queue.
 This method will block until all the active jobs in the queue have
 reached a terminal state (completed, errored or cancelled).
 
-#### queue.wait_for_job(job, [timeout])
-
-This method will block until the indicated job has reached a terminal
-state (completed, errored or cancelled). If the optional timeout is
-provided, the call will block for at most timeout seconds, and raise a
-TimeoutError otherwise.
-
 #### jobs = queue.list_jobs()
 
 This will return a list of all jobs, including ones that have not yet
@@ -1505,9 +1449,9 @@ set of keys to the corresponding model config objects.
 Find all model metadata records that have the given author and return
 a set of keys to the corresponding model config objects.
 
-***
+# The remainder of this documentation is provisional, pending implementation of the Load service
 
-## The Lowdown on the ModelLoadService
+## Let's get loaded, the lowdown on ModelLoadService
 
 The `ModelLoadService` is responsible for loading a named model into
 memory so that it can be used for inference. Despite the fact that it
@@ -1521,7 +1465,7 @@ create alternative instances if you wish.
 ### Creating a ModelLoadService object
 
 The class is defined in
-`invokeai.app.services.model_load`. It is initialized with
+`invokeai.app.services.model_loader_service`. It is initialized with
 an InvokeAIAppConfig object, from which it gets configuration
 information such as the user's desired GPU and precision, and with a
 previously-created `ModelRecordServiceBase` object, from which it
@@ -1531,29 +1475,26 @@ Here is a typical initialization pattern:
 
 ```
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.app.services.model_load import ModelLoadService, ModelLoaderRegistry
+from invokeai.app.services.model_record_service import ModelRecordServiceBase
+from invokeai.app.services.model_loader_service import ModelLoadService
 
 config = InvokeAIAppConfig.get_config()
-ram_cache = ModelCache(
-	max_cache_size=config.ram_cache_size, max_vram_cache_size=config.vram_cache_size, logger=logger
-)
-convert_cache = ModelConvertCache(
-	cache_path=config.models_convert_cache_path, max_size=config.convert_cache_size
-)
-loader = ModelLoadService(
-	app_config=config,
-	ram_cache=ram_cache,
-	convert_cache=convert_cache,
-	registry=ModelLoaderRegistry
-)
+store = ModelRecordServiceBase.open(config)
+loader = ModelLoadService(config, store)
 ```
 
-### load_model(model_config, [submodel_type], [context]) -> LoadedModel
+Note that we are relying on the contents of the application
+configuration to choose the implementation of
+`ModelRecordServiceBase`.
 
-The `load_model()` method takes an `AnyModelConfig` returned by
-`ModelRecordService.get_model()` and returns the corresponding loaded
+### get_model(key, [submodel_type], [context]) -> ModelInfo:
+
+*** TO DO: change to get_model(key, context=None, **kwargs)
+
+The `get_model()` method, like its similarly-named cousin in
+`ModelRecordService`, receives the unique key that identifies the
 model.  It loads the model into memory, gets the model ready for use,
-and returns a `LoadedModel` object.
+and returns a `ModelInfo` object. 
 
 The optional second argument, `subtype` is a `SubModelType` string
 enum, such as "vae". It is mandatory when used with a main model, and
@@ -1563,45 +1504,46 @@ The optional third argument, `context` can be provided by
 an invocation to trigger model load event reporting. See below for
 details.
 
-The returned `LoadedModel` object contains a copy of the configuration
-record returned by the model record `get_model()` method, as well as
-the in-memory loaded model:
+The returned `ModelInfo` object shares some fields in common with
+`ModelConfigBase`, but is otherwise a completely different beast:
 
-
-| **Attribute Name** | **Type**        |  **Description** |
+| **Field Name** | **Type**        |  **Description** |
 |----------------|-----------------|------------------|
-| `config`       | AnyModelConfig         | A copy of the model's configuration record for retrieving base type, etc. |
-| `model`        | AnyModel               | The instantiated model (details below) |
-| `locker`       | ModelLockerBase        | A context manager that mediates the movement of the model into VRAM |
+| `key`          | str                    | The model key derived from the ModelRecordService database |
+| `name`         | str                    | Name of this model |
+| `base_model`   | BaseModelType          | Base model for this model |
+| `type`         | ModelType or SubModelType   | Either the model type (non-main) or the submodel type (main models)|
+| `location`     | Path or str            | Location of the model on the filesystem |
+| `precision`    | torch.dtype            | The torch.precision to use for inference |
+| `context`      | ModelCache.ModelLocker | A context class used to lock the model in VRAM while in use |
 
-Because the loader can return multiple model types, it is typed to
-return `AnyModel`, a Union `ModelMixin`, `torch.nn.Module`,
-`IAIOnnxRuntimeModel`, `IPAdapter`, `IPAdapterPlus`, and
-`EmbeddingModelRaw`. `ModelMixin` is the base class of all diffusers
-models, `EmbeddingModelRaw` is used for LoRA and TextualInversion
-models. The others are obvious.
+The types for `ModelInfo` and `SubModelType` can be imported from
+`invokeai.app.services.model_loader_service`.
 
-
-`LoadedModel` acts as a context manager. The context loads the model
-into the execution device (e.g. VRAM on CUDA systems), locks the model
-in the execution device for the duration of the context, and returns
-the model. Use it like this:
+To use the model, you use the `ModelInfo` as a context manager using
+the following pattern:
 
 ```
-model_info = loader.get_model_by_key('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
+model_info = loader.get_model('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
 with model_info as vae:
 	image = vae.decode(latents)[0]
 ```
 
-`get_model_by_key()` may raise any of the following exceptions:
+The `vae` model will stay locked in the GPU during the period of time
+it is in the context manager's scope.
 
-- `UnknownModelException`   -- key not in database
-- `ModelNotFoundException`  -- key in database but model not found at path
-- `NotImplementedException` -- the loader doesn't know how to load this type of model
+`get_model()` may raise any of the following exceptions:
+
+- `UnknownModelException`  -- key not in database
+- `ModelNotFoundException` -- key in database but model not found at path
+- `InvalidModelException`  -- the model is guilty of a variety of sins
   
+** TO DO: ** Resolve discrepancy between ModelInfo.location and
+ModelConfig.path.
+
 ### Emitting model loading events
 
-When the `context` argument is passed to `load_model_*()`, it will
+When the `context` argument is passed to `get_model()`, it will
 retrieve the invocation event bus from the passed `InvocationContext`
 object to emit events on the invocation bus. The two events are
 "model_load_started" and "model_load_completed". Both carry the
@@ -1614,174 +1556,10 @@ payload=dict(
 	queue_batch_id=queue_batch_id,
 	graph_execution_state_id=graph_execution_state_id,
 	model_key=model_key,
-	submodel_type=submodel,
+	submodel=submodel,
 	hash=model_info.hash,
 	location=str(model_info.location),
 	precision=str(model_info.precision),
 )
 ```
 
-### Adding Model Loaders
-
-Model loaders are small classes that inherit from the `ModelLoader`
-base class. They typically implement one method `_load_model()` whose
-signature is:
-
-```
-def _load_model(
-    self,
-    model_path: Path,
-    model_variant: Optional[ModelRepoVariant] = None,
-    submodel_type: Optional[SubModelType] = None,
-) -> AnyModel:
-```
-
-`_load_model()` will be passed the path to the model on disk, an
-optional repository variant (used by the diffusers loaders to select,
-e.g.  the `fp16` variant, and an optional submodel_type for main and
-onnx models.
-
-To install a new loader, place it in
-`invokeai/backend/model_manager/load/model_loaders`. Inherit from
-`ModelLoader` and use the `@ModelLoaderRegistry.register()` decorator to
-indicate what type of models the loader can handle.
-
-Here is a complete example from `generic_diffusers.py`, which is able
-to load several different diffusers types:
-
-```
-from pathlib import Path
-from typing import Optional
-
-from invokeai.backend.model_manager import (
-    AnyModel,
-    BaseModelType,
-    ModelFormat,
-    ModelRepoVariant,
-    ModelType,
-    SubModelType,
-)
-from .. import ModelLoader, ModelLoaderRegistry
-
-
-@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers)
-@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.T2IAdapter, format=ModelFormat.Diffusers)
-class GenericDiffusersLoader(ModelLoader):
-    """Class to load simple diffusers models."""
-
-    def _load_model(
-        self,
-        model_path: Path,
-        model_variant: Optional[ModelRepoVariant] = None,
-        submodel_type: Optional[SubModelType] = None,
-    ) -> AnyModel:
-        model_class = self._get_hf_load_class(model_path)
-        if submodel_type is not None:
-            raise Exception(f"There are no submodels in models of type {model_class}")
-        variant = model_variant.value if model_variant else None
-        result: AnyModel = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, variant=variant)  # type: ignore
-        return result
-```
-
-Note that a loader can register itself to handle several different
-model types. An exception will be raised if more than one loader tries
-to register the same model type.
-
-#### Conversion
-
-Some models require conversion to diffusers format before they can be
-loaded. These loaders should override two additional methods:
-
-```
-_needs_conversion(self, config: AnyModelConfig, model_path: Path, dest_path: Path) -> bool
-_convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Path) -> Path:
-```
-
-The first method accepts the model configuration, the path to where
-the unmodified model is currently installed, and a proposed
-destination for the converted model. This method returns True if the
-model needs to be converted. It typically does this by comparing the
-last modification time of the original model file to the modification
-time of the converted model. In some cases you will also want to check
-the modification date of the configuration record, in the event that
-the user has changed something like the scheduler prediction type that
-will require the model to be re-converted. See `controlnet.py` for an
-example of this logic.
-
-The second method accepts the model configuration, the path to the
-original model on disk, and the desired output path for the converted
-model. It does whatever it needs to do to get the model into diffusers
-format, and returns the Path of the resulting model. (The path should
-ordinarily be the same as `output_path`.)
-
-## The ModelManagerService object
-
-For convenience, the API provides a `ModelManagerService` object which
-gives a single point of access to the major model manager
-services. This object is created at initialization time and can be
-found in the global `ApiDependencies.invoker.services.model_manager`
-object, or in `context.services.model_manager` from within an
-invocation.
-
-In the examples below, we have retrieved the manager using:
-```
-mm = ApiDependencies.invoker.services.model_manager
-```
-
-The following properties and methods will be available:
-
-### mm.store
-
-This retrieves the `ModelRecordService` associated with the
-manager. Example:
-
-```
-configs = mm.store.get_model_by_attr(name='stable-diffusion-v1-5')
-```
-
-### mm.install
-
-This retrieves the `ModelInstallService` associated with the manager.
-Example:
-
-```
-job = mm.install.heuristic_import(`https://civitai.com/models/58390/detail-tweaker-lora-lora`)
-```
-
-### mm.load
-
-This retrieves the `ModelLoaderService` associated with the manager. Example:
-
-```
-configs = mm.store.get_model_by_attr(name='stable-diffusion-v1-5')
-assert len(configs) > 0
-
-loaded_model = mm.load.load_model(configs[0])
-```
-
-The model manager also offers a few convenience shortcuts for loading
-models:
-
-### mm.load_model_by_config(model_config, [submodel], [context]) -> LoadedModel
-
-Same as `mm.load.load_model()`.
-
-### mm.load_model_by_attr(model_name, base_model, model_type, [submodel], [context]) -> LoadedModel
-
-This accepts the combination of the model's name, type and base, which
-it passes to the model record config store for retrieval. If a unique
-model config is found, this method returns a `LoadedModel`. It can
-raise the following exceptions:
-
-```
-UnknownModelException -- model with these attributes not known
-NotImplementedException -- the loader doesn't know how to load this type of model
-ValueError -- more than one model matches this combination of base/type/name
-```
-
-### mm.load_model_by_key(key, [submodel], [context]) -> LoadedModel
-
-This method takes a model key, looks it up using the
-`ModelRecordServiceBase` object in `mm.store`, and passes the returned
-model configuration to `load_model_by_config()`.  It may raise a
-`NotImplementedException`.

docs/features/CONTROLNET.md

@@ -94,8 +94,6 @@ A model that helps generate creative QR codes that still scan. Can also be used
 **Openpose**: 
 The OpenPose control model allows for the identification of the general pose of a character by pre-processing an existing image with a clear human structure. With advanced options, Openpose can also detect the face or hands in the image. 
 
-*Note:* The DWPose Processor has replaced the OpenPose processor in Invoke. Workflows and generations that relied on the OpenPose Processor will need to be updated to use the DWPose Processor instead.
-
 **Mediapipe Face**:
 
 The MediaPipe Face identification processor is able to clearly identify facial features in order to capture vivid expressions of human faces.

docs/installation/020_INSTALL_MANUAL.md

@@ -230,13 +230,13 @@ manager, please follow these steps:
         === "local Webserver"
 
             ```bash
-            invokeai-web
+            invokeai --web
             ```
 
         === "Public Webserver"
 
             ```bash
-            invokeai-web --host 0.0.0.0
+            invokeai --web --host 0.0.0.0
             ```
 
         === "CLI"
@@ -402,4 +402,4 @@ environment variable INVOKEAI_ROOT to point to the installation directory.
 Note that if you run into problems with the Conda installation, the InvokeAI
 staff will **not** be able to help you out. Caveat Emptor!
 
-[dev-chat]: https://discord.com/channels/1020123559063990373/1049495067846524939
+[dev-chat]: https://discord.com/channels/1020123559063990373/1049495067846524939

docs/installation/040_INSTALL_DOCKER.md

@@ -69,7 +69,7 @@ a token and copy it, since you will need in for the next step.
 
 ### Setup
 
-Set up your environmnent variables. In the `docker` directory, make a copy of `.env.sample` and name it `.env`. Make changes as necessary.
+Set up your environmnent variables. In the `docker` directory, make a copy of `env.sample` and name it `.env`. Make changes as necessary.
 
 Any environment variables supported by InvokeAI can be set here - please see the [CONFIGURATION](../features/CONFIGURATION.md) for further detail.
 

docs/nodes/defaultNodes.md

@@ -81,7 +81,7 @@ their descriptions.
 | ONNX Text to Latents                                          | Generates latents from conditionings.                                                                                                                |
 | ONNX Model Loader                                             | Loads a main model, outputting its submodels.                                                                                                        |
 | OpenCV Inpaint                                                | Simple inpaint using opencv.                                                                                                                         |
-| DW Openpose Processor                                            | Applies Openpose processing to image                                                                                                                 |
+| Openpose Processor                                            | Applies Openpose processing to image                                                                                                                 |
 | PIDI Processor                                                | Applies PIDI processing to image                                                                                                                     |
 | Prompts from File                                             | Loads prompts from a text file                                                                                                                       |
 | Random Integer                                                | Outputs a single random integer.                                                                                                                     |

installer/lib/messages.py

@@ -91,8 +91,8 @@ def choose_version(available_releases: tuple | None = None) -> str:
         complete_while_typing=True,
         completer=FuzzyWordCompleter(choices),
     )
-    console.print(f"   Version {choices[0] if response == '' else response} will be installed.")
 
+    console.print(f"   Version {choices[0]} will be installed.")
     console.line()
 
     return "stable" if response == "" else response

invokeai/app/api/dependencies.py

@@ -2,12 +2,9 @@
 
 from logging import Logger
 
-import torch
-
-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.item_storage.item_storage_memory import ItemStorageMemory
 from invokeai.app.services.shared.sqlite.sqlite_util import init_db
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
+from invokeai.backend.model_manager.metadata import ModelMetadataStore
 from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.version.invokeai_version import __version__
 
@@ -15,22 +12,26 @@ 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
+from ..services.latents_storage.latents_storage_disk import DiskLatentsStorage
+from ..services.latents_storage.latents_storage_forward_cache import ForwardCacheLatentsStorage
+from ..services.model_install import ModelInstallService
 from ..services.model_manager.model_manager_default import ModelManagerService
-from ..services.model_metadata import ModelMetadataStoreSQL
 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
@@ -67,9 +68,6 @@ class ApiDependencies:
         logger.debug(f"Internet connectivity is {config.internet_available}")
 
         output_folder = config.output_path
-        if output_folder is None:
-            raise ValueError("Output folder is not set")
-
         image_files = DiskImageFileStorage(f"{output_folder}/images")
 
         db = init_db(config=config, logger=logger, image_files=image_files)
@@ -82,26 +80,26 @@ 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)
-        tensors = ObjectSerializerForwardCache(
-            ObjectSerializerDisk[torch.Tensor](output_folder / "tensors", ephemeral=True)
-        )
-        conditioning = ObjectSerializerForwardCache(
-            ObjectSerializerDisk[ConditioningFieldData](output_folder / "conditioning", ephemeral=True)
-        )
+        latents = ForwardCacheLatentsStorage(DiskLatentsStorage(f"{output_folder}/latents"))
+        model_manager = ModelManagerService(config, logger)
+        model_record_service = ModelRecordServiceSQL(db=db)
         download_queue_service = DownloadQueueService(event_bus=events)
-        model_metadata_service = ModelMetadataStoreSQL(db=db)
-        model_manager = ModelManagerService.build_model_manager(
-            app_config=configuration,
-            model_record_service=ModelRecordServiceSQL(db=db, metadata_store=model_metadata_service),
+        metadata_store = ModelMetadataStore(db=db)
+        model_install_service = ModelInstallService(
+            app_config=config,
+            record_store=model_record_service,
             download_queue=download_queue_service,
-            events=events,
+            metadata_store=metadata_store,
+            event_bus=events,
         )
         names = SimpleNameService()
         performance_statistics = InvocationStatsService()
+        processor = DefaultInvocationProcessor()
+        queue = MemoryInvocationQueue()
         session_processor = DefaultSessionProcessor()
         session_queue = SqliteSessionQueue(db=db)
         urls = LocalUrlService()
@@ -112,24 +110,27 @@ 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,
             invocation_cache=invocation_cache,
+            latents=latents,
             logger=logger,
             model_manager=model_manager,
+            model_records=model_record_service,
             download_queue=download_queue_service,
+            model_install=model_install_service,
             names=names,
             performance_statistics=performance_statistics,
+            processor=processor,
+            queue=queue,
             session_processor=session_processor,
             session_queue=session_queue,
             urls=urls,
             workflow_records=workflow_records,
-            tensors=tensors,
-            conditioning=conditioning,
         )
 
         ApiDependencies.invoker = Invoker(services)

invokeai/app/api/routers/download_queue.py

@@ -36,7 +36,7 @@ async def list_downloads() -> List[DownloadJob]:
         400: {"description": "Bad request"},
     },
 )
-async def prune_downloads() -> Response:
+async def prune_downloads():
     """Prune completed and errored jobs."""
     queue = ApiDependencies.invoker.services.download_queue
     queue.prune_jobs()
@@ -55,7 +55,7 @@ async def download(
 ) -> DownloadJob:
     """Download the source URL to the file or directory indicted in dest."""
     queue = ApiDependencies.invoker.services.download_queue
-    return queue.download(source, Path(dest), priority, access_token)
+    return queue.download(source, dest, priority, access_token)
 
 
 @download_queue_router.get(
@@ -87,7 +87,7 @@ async def get_download_job(
 )
 async def cancel_download_job(
     id: int = Path(description="ID of the download job to cancel."),
-) -> Response:
+):
     """Cancel a download job using its ID."""
     try:
         queue = ApiDependencies.invoker.services.download_queue
@@ -105,7 +105,7 @@ async def cancel_download_job(
         204: {"description": "Download jobs have been cancelled"},
     },
 )
-async def cancel_all_download_jobs() -> Response:
+async def cancel_all_download_jobs():
     """Cancel all download jobs."""
     ApiDependencies.invoker.services.download_queue.cancel_all_jobs()
     return Response(status_code=204)

invokeai/app/api/routers/images.py

@@ -2,13 +2,13 @@ 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
 from pydantic import BaseModel, Field, ValidationError
 
-from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator
+from invokeai.app.invocations.baseinvocation import MetadataField, MetadataFieldValidator
 from invokeai.app.services.image_records.image_records_common import ImageCategory, ImageRecordChanges, ResourceOrigin
 from invokeai.app.services.images.images_common import ImageDTO, ImageUrlsDTO
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
@@ -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

@@ -1,751 +0,0 @@
-# Copyright (c) 2023 Lincoln D. Stein
-"""FastAPI route for model configuration records."""
-
-import pathlib
-import shutil
-from hashlib import sha1
-from random import randbytes
-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 starlette.exceptions import HTTPException
-from typing_extensions import Annotated
-
-from invokeai.app.services.model_install import ModelInstallJob
-from invokeai.app.services.model_records import (
-    DuplicateModelException,
-    InvalidModelException,
-    ModelRecordOrderBy,
-    ModelSummary,
-    UnknownModelException,
-)
-from invokeai.app.services.shared.pagination import PaginatedResults
-from invokeai.backend.model_manager.config import (
-    AnyModelConfig,
-    BaseModelType,
-    MainCheckpointConfig,
-    ModelFormat,
-    ModelType,
-    SubModelType,
-)
-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
-
-model_manager_router = APIRouter(prefix="/v2/models", tags=["model_manager"])
-
-
-class ModelsList(BaseModel):
-    """Return list of configs."""
-
-    models: List[AnyModelConfig]
-
-    model_config = ConfigDict(use_enum_values=True)
-
-
-class ModelTagSet(BaseModel):
-    """Return tags for a set of models."""
-
-    key: str
-    name: str
-    author: str
-    tags: Set[str]
-
-
-##############################################################################
-# These are example inputs and outputs that are used in places where Swagger
-# is unable to generate a correct example.
-##############################################################################
-example_model_config = {
-    "path": "string",
-    "name": "string",
-    "base": "sd-1",
-    "type": "main",
-    "format": "checkpoint",
-    "config": "string",
-    "key": "string",
-    "original_hash": "string",
-    "current_hash": "string",
-    "description": "string",
-    "source": "string",
-    "last_modified": 0,
-    "vae": "string",
-    "variant": "normal",
-    "prediction_type": "epsilon",
-    "repo_variant": "fp16",
-    "upcast_attention": False,
-    "ztsnr_training": False,
-}
-
-example_model_input = {
-    "path": "/path/to/model",
-    "name": "model_name",
-    "base": "sd-1",
-    "type": "main",
-    "format": "checkpoint",
-    "config": "configs/stable-diffusion/v1-inference.yaml",
-    "description": "Model description",
-    "vae": None,
-    "variant": "normal",
-}
-
-example_model_metadata = {
-    "name": "ip_adapter_sd_image_encoder",
-    "author": "InvokeAI",
-    "tags": [
-        "transformers",
-        "safetensors",
-        "clip_vision_model",
-        "endpoints_compatible",
-        "region:us",
-        "has_space",
-        "license:apache-2.0",
-    ],
-    "files": [
-        {
-            "url": "https://huggingface.co/InvokeAI/ip_adapter_sd_image_encoder/resolve/main/README.md",
-            "path": "ip_adapter_sd_image_encoder/README.md",
-            "size": 628,
-            "sha256": None,
-        },
-        {
-            "url": "https://huggingface.co/InvokeAI/ip_adapter_sd_image_encoder/resolve/main/config.json",
-            "path": "ip_adapter_sd_image_encoder/config.json",
-            "size": 560,
-            "sha256": None,
-        },
-        {
-            "url": "https://huggingface.co/InvokeAI/ip_adapter_sd_image_encoder/resolve/main/model.safetensors",
-            "path": "ip_adapter_sd_image_encoder/model.safetensors",
-            "size": 2528373448,
-            "sha256": "6ca9667da1ca9e0b0f75e46bb030f7e011f44f86cbfb8d5a36590fcd7507b030",
-        },
-    ],
-    "type": "huggingface",
-    "id": "InvokeAI/ip_adapter_sd_image_encoder",
-    "tag_dict": {"license": "apache-2.0"},
-    "last_modified": "2023-09-23T17:33:25Z",
-}
-
-##############################################################################
-# ROUTES
-##############################################################################
-
-
-@model_manager_router.get(
-    "/",
-    operation_id="list_model_records",
-)
-async def list_model_records(
-    base_models: Optional[List[BaseModelType]] = Query(default=None, description="Base models to include"),
-    model_type: Optional[ModelType] = Query(default=None, description="The type of model to get"),
-    model_name: Optional[str] = Query(default=None, description="Exact match on the name of the model"),
-    model_format: Optional[ModelFormat] = Query(
-        default=None, description="Exact match on the format of the model (e.g. 'diffusers')"
-    ),
-) -> ModelsList:
-    """Get a list of models."""
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    found_models: list[AnyModelConfig] = []
-    if base_models:
-        for base_model in base_models:
-            found_models.extend(
-                record_store.search_by_attr(
-                    base_model=base_model, model_type=model_type, model_name=model_name, model_format=model_format
-                )
-            )
-    else:
-        found_models.extend(
-            record_store.search_by_attr(model_type=model_type, model_name=model_name, model_format=model_format)
-        )
-    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",
-    responses={
-        200: {
-            "description": "The model configuration was retrieved successfully",
-            "content": {"application/json": {"example": example_model_config}},
-        },
-        400: {"description": "Bad request"},
-        404: {"description": "The model could not be found"},
-    },
-)
-async def get_model_record(
-    key: str = Path(description="Key of the model record to fetch."),
-) -> AnyModelConfig:
-    """Get a model record"""
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    try:
-        config: AnyModelConfig = record_store.get_model(key)
-        return config
-    except UnknownModelException as e:
-        raise HTTPException(status_code=404, detail=str(e))
-
-
-@model_manager_router.get("/summary", operation_id="list_model_summary")
-async def list_model_summary(
-    page: int = Query(default=0, description="The page to get"),
-    per_page: int = Query(default=10, description="The number of models per page"),
-    order_by: ModelRecordOrderBy = Query(default=ModelRecordOrderBy.Default, description="The attribute to order by"),
-) -> PaginatedResults[ModelSummary]:
-    """Gets a page of model summary data."""
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    results: PaginatedResults[ModelSummary] = record_store.list_models(page=page, per_page=per_page, order_by=order_by)
-    return results
-
-
-@model_manager_router.get(
-    "/i/{key}/metadata",
-    operation_id="get_model_metadata",
-    responses={
-        200: {
-            "description": "The model metadata was retrieved successfully",
-            "content": {"application/json": {"example": example_model_metadata}},
-        },
-        400: {"description": "Bad request"},
-    },
-)
-async def get_model_metadata(
-    key: str = Path(description="Key of the model repo metadata to fetch."),
-) -> Optional[AnyModelRepoMetadata]:
-    """Get a model metadata object."""
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    result: Optional[AnyModelRepoMetadata] = record_store.get_metadata(key)
-
-    return result
-
-
-@model_manager_router.get(
-    "/tags",
-    operation_id="list_tags",
-)
-async def list_tags() -> Set[str]:
-    """Get a unique set of all the model tags."""
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    result: Set[str] = record_store.list_tags()
-    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",
-)
-async def search_by_metadata_tags(
-    tags: Set[str] = Query(default=None, description="Tags to search for"),
-) -> ModelsList:
-    """Get a list of models."""
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    results = record_store.search_by_metadata_tag(tags)
-    return ModelsList(models=results)
-
-
-@model_manager_router.patch(
-    "/i/{key}",
-    operation_id="update_model_record",
-    responses={
-        200: {
-            "description": "The model was updated successfully",
-            "content": {"application/json": {"example": example_model_config}},
-        },
-        400: {"description": "Bad request"},
-        404: {"description": "The model could not be found"},
-        409: {"description": "There is already a model corresponding to the new name"},
-    },
-    status_code=200,
-)
-async def update_model_record(
-    key: Annotated[str, Path(description="Unique key of model")],
-    info: Annotated[
-        AnyModelConfig, Body(description="Model config", discriminator="type", example=example_model_input)
-    ],
-) -> AnyModelConfig:
-    """Update model contents with a new config. If the model name or base fields are changed, then the model is renamed."""
-    logger = ApiDependencies.invoker.services.logger
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    try:
-        model_response: AnyModelConfig = record_store.update_model(key, config=info)
-        logger.info(f"Updated model: {key}")
-    except UnknownModelException as e:
-        raise HTTPException(status_code=404, detail=str(e))
-    except ValueError as e:
-        logger.error(str(e))
-        raise HTTPException(status_code=409, detail=str(e))
-    return model_response
-
-
-@model_manager_router.delete(
-    "/i/{key}",
-    operation_id="del_model_record",
-    responses={
-        204: {"description": "Model deleted successfully"},
-        404: {"description": "Model not found"},
-    },
-    status_code=204,
-)
-async def del_model_record(
-    key: str = Path(description="Unique key of model to remove from model registry."),
-) -> Response:
-    """
-    Delete model record from database.
-
-    The configuration record will be removed. The corresponding weights files will be
-    deleted as well if they reside within the InvokeAI "models" directory.
-    """
-    logger = ApiDependencies.invoker.services.logger
-
-    try:
-        installer = ApiDependencies.invoker.services.model_manager.install
-        installer.delete(key)
-        logger.info(f"Deleted model: {key}")
-        return Response(status_code=204)
-    except UnknownModelException as e:
-        logger.error(str(e))
-        raise HTTPException(status_code=404, detail=str(e))
-
-
-@model_manager_router.post(
-    "/i/",
-    operation_id="add_model_record",
-    responses={
-        201: {
-            "description": "The model added successfully",
-            "content": {"application/json": {"example": example_model_config}},
-        },
-        409: {"description": "There is already a model corresponding to this path or repo_id"},
-        415: {"description": "Unrecognized file/folder format"},
-    },
-    status_code=201,
-)
-async def add_model_record(
-    config: Annotated[
-        AnyModelConfig, Body(description="Model config", discriminator="type", example=example_model_input)
-    ],
-) -> AnyModelConfig:
-    """Add a model using the configuration information appropriate for its type."""
-    logger = ApiDependencies.invoker.services.logger
-    record_store = ApiDependencies.invoker.services.model_manager.store
-    if config.key == "<NOKEY>":
-        config.key = sha1(randbytes(100)).hexdigest()
-        logger.info(f"Created model {config.key} for {config.name}")
-    try:
-        record_store.add_model(config.key, config)
-    except DuplicateModelException as e:
-        logger.error(str(e))
-        raise HTTPException(status_code=409, detail=str(e))
-    except InvalidModelException as e:
-        logger.error(str(e))
-        raise HTTPException(status_code=415)
-
-    # now fetch it out
-    result: AnyModelConfig = record_store.get_model(config.key)
-    return result
-
-
-@model_manager_router.post(
-    "/install",
-    operation_id="install_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 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?
-    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"},
-    ),
-    access_token: Optional[str] = None,
-) -> ModelInstallJob:
-    """Install a model using a string identifier.
-
-    `source` can be any of the following.
-
-    1. A path on the local filesystem ('C:\\users\\fred\\model.safetensors')
-    2. A Url pointing to a single downloadable model file
-    3. A HuggingFace repo_id with any of the following formats:
-       - model/name
-       - model/name:fp16:vae
-       - model/name::vae          -- use default precision
-       - model/name:fp16:path/to/model.safetensors
-       - model/name::path/to/model.safetensors
-
-    `config` is an optional dict containing model configuration values that will override
-    the ones that are probed automatically.
-
-    `access_token` is an optional access token for use with Urls that require
-    authentication.
-
-    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.
-
-    See the documentation for `import_model_record` for more information on
-    interpreting the job information returned by this route.
-    """
-    logger = ApiDependencies.invoker.services.logger
-
-    try:
-        installer = ApiDependencies.invoker.services.model_manager.install
-        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.get(
-    "/import",
-    operation_id="list_model_install_jobs",
-)
-async def list_model_install_jobs() -> List[ModelInstallJob]:
-    """Return the list of model install jobs.
-
-    Install jobs have a numeric `id`, a `status`, and other fields that provide information on
-    the nature of the job and its progress. The `status` is one of:
-
-    * "waiting" -- Job is waiting in the queue to run
-    * "downloading" -- Model file(s) are downloading
-    * "running" -- Model has downloaded and the model probing and registration process is running
-    * "completed" -- Installation completed successfully
-    * "error" -- An error occurred. Details will be in the "error_type" and "error" fields.
-    * "cancelled" -- Job was cancelled before completion.
-
-    Once completed, information about the model such as its size, base
-    model, type, and metadata can be retrieved from the `config_out`
-    field. For multi-file models such as diffusers, information on individual files
-    can be retrieved from `download_parts`.
-
-    See the example and schema below for more information.
-    """
-    jobs: List[ModelInstallJob] = ApiDependencies.invoker.services.model_manager.install.list_jobs()
-    return jobs
-
-
-@model_manager_router.get(
-    "/import/{id}",
-    operation_id="get_model_install_job",
-    responses={
-        200: {"description": "Success"},
-        404: {"description": "No such job"},
-    },
-)
-async def get_model_install_job(id: int = Path(description="Model install id")) -> ModelInstallJob:
-    """
-    Return model install job corresponding to the given source. See the documentation for 'List Model Install Jobs'
-    for information on the format of the return value.
-    """
-    try:
-        result: ModelInstallJob = ApiDependencies.invoker.services.model_manager.install.get_job_by_id(id)
-        return result
-    except ValueError as e:
-        raise HTTPException(status_code=404, detail=str(e))
-
-
-@model_manager_router.delete(
-    "/import/{id}",
-    operation_id="cancel_model_install_job",
-    responses={
-        201: {"description": "The job was cancelled successfully"},
-        415: {"description": "No such job"},
-    },
-    status_code=201,
-)
-async def cancel_model_install_job(id: int = Path(description="Model install job ID")) -> None:
-    """Cancel the model install job(s) corresponding to the given job ID."""
-    installer = ApiDependencies.invoker.services.model_manager.install
-    try:
-        job = installer.get_job_by_id(id)
-    except ValueError as e:
-        raise HTTPException(status_code=415, detail=str(e))
-    installer.cancel_job(job)
-
-
-@model_manager_router.patch(
-    "/import",
-    operation_id="prune_model_install_jobs",
-    responses={
-        204: {"description": "All completed and errored jobs have been pruned"},
-        400: {"description": "Bad request"},
-    },
-)
-async def prune_model_install_jobs() -> Response:
-    """Prune all completed and errored jobs from the install job list."""
-    ApiDependencies.invoker.services.model_manager.install.prune_jobs()
-    return Response(status_code=204)
-
-
-@model_manager_router.patch(
-    "/sync",
-    operation_id="sync_models_to_config",
-    responses={
-        204: {"description": "Model config record database resynced with files on disk"},
-        400: {"description": "Bad request"},
-    },
-)
-async def sync_models_to_config() -> Response:
-    """
-    Traverse the models and autoimport directories.
-
-    Model files without a corresponding
-    record in the database are added. Orphan records without a models file are deleted.
-    """
-    ApiDependencies.invoker.services.model_manager.install.sync_to_config()
-    return Response(status_code=204)
-
-
-@model_manager_router.put(
-    "/convert/{key}",
-    operation_id="convert_model",
-    responses={
-        200: {
-            "description": "Model converted successfully",
-            "content": {"application/json": {"example": example_model_config}},
-        },
-        400: {"description": "Bad request"},
-        404: {"description": "Model not found"},
-        409: {"description": "There is already a model registered at this location"},
-    },
-)
-async def convert_model(
-    key: str = Path(description="Unique key of the safetensors main model to convert to diffusers format."),
-) -> AnyModelConfig:
-    """
-    Permanently convert a model into diffusers format, replacing the safetensors version.
-    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
-    installer = ApiDependencies.invoker.services.model_manager.install
-
-    try:
-        model_config = store.get_model(key)
-    except UnknownModelException as e:
-        logger.error(str(e))
-        raise HTTPException(status_code=424, detail=str(e))
-
-    if not isinstance(model_config, MainCheckpointConfig):
-        logger.error(f"The model with key {key} is not a main checkpoint 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)
-
-    # Get the path of the converted model from the loader
-    cache_path = loader.convert_cache.cache_path(key)
-    assert cache_path.exists()
-
-    # temporarily rename the original safetensors file so that there is no naming conflict
-    original_name = model_config.name
-    model_config.name = f"{original_name}.DELETE"
-    store.update_model(key, config=model_config)
-
-    # install the diffusers
-    try:
-        new_key = installer.install_path(
-            cache_path,
-            config={
-                "name": original_name,
-                "description": model_config.description,
-                "original_hash": model_config.original_hash,
-                "source": model_config.source,
-            },
-        )
-    except DuplicateModelException as e:
-        logger.error(str(e))
-        raise HTTPException(status_code=409, detail=str(e))
-
-    # get the original metadata
-    if orig_metadata := store.get_metadata(key):
-        store.metadata_store.add_metadata(new_key, orig_metadata)
-
-    # delete the original safetensors file
-    installer.delete(key)
-
-    # delete the cached version
-    shutil.rmtree(cache_path)
-
-    # return the config record for the new diffusers directory
-    new_config: AnyModelConfig = store.get_model(new_key)
-    return new_config
-
-
-@model_manager_router.put(
-    "/merge",
-    operation_id="merge",
-    responses={
-        200: {
-            "description": "Model converted successfully",
-            "content": {"application/json": {"example": example_model_config}},
-        },
-        400: {"description": "Bad request"},
-        404: {"description": "Model not found"},
-        409: {"description": "There is already a model registered at this location"},
-    },
-)
-async def merge(
-    keys: List[str] = Body(description="Keys for two to three models to merge", min_length=2, max_length=3),
-    merged_model_name: Optional[str] = Body(description="Name of destination model", default=None),
-    alpha: float = Body(description="Alpha weighting strength to apply to 2d and 3d models", default=0.5),
-    force: bool = Body(
-        description="Force merging of models created with different versions of diffusers",
-        default=False,
-    ),
-    interp: Optional[MergeInterpolationMethod] = Body(description="Interpolation method", default=None),
-    merge_dest_directory: Optional[str] = Body(
-        description="Save the merged model to the designated directory (with 'merged_model_name' appended)",
-        default=None,
-    ),
-) -> AnyModelConfig:
-    """
-    Merge diffusers models. The process is controlled by a set parameters provided in the body of the request.
-    ```
-    Argument                Description [default]
-    --------               ----------------------
-    keys                   List of 2-3 model keys to merge together. All models must use the same base type.
-    merged_model_name      Name for the merged model [Concat model names]
-    alpha                  Alpha value (0.0-1.0). Higher values give more weight to the second model [0.5]
-    force                  If true, force the merge even if the models were generated by different versions of the diffusers library [False]
-    interp                 Interpolation method. One of "weighted_sum", "sigmoid", "inv_sigmoid" or "add_difference" [weighted_sum]
-    merge_dest_directory   Specify a directory to store the merged model in [models directory]
-    ```
-    """
-    logger = ApiDependencies.invoker.services.logger
-    try:
-        logger.info(f"Merging models: {keys} into {merge_dest_directory or '<MODELS>'}/{merged_model_name}")
-        dest = pathlib.Path(merge_dest_directory) if merge_dest_directory else None
-        installer = ApiDependencies.invoker.services.model_manager.install
-        merger = ModelMerger(installer)
-        model_names = [installer.record_store.get_model(x).name for x in keys]
-        response = merger.merge_diffusion_models_and_save(
-            model_keys=keys,
-            merged_model_name=merged_model_name or "+".join(model_names),
-            alpha=alpha,
-            interp=interp,
-            force=force,
-            merge_dest_directory=dest,
-        )
-    except UnknownModelException:
-        raise HTTPException(
-            status_code=404,
-            detail=f"One or more of the models '{keys}' not found",
-        )
-    except ValueError as e:
-        raise HTTPException(status_code=400, detail=str(e))
-    return response

invokeai/app/api/routers/model_records.py

@@ -0,0 +1,472 @@
+# Copyright (c) 2023 Lincoln D. Stein
+"""FastAPI route for model configuration records."""
+
+import pathlib
+from hashlib import sha1
+from random import randbytes
+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
+from starlette.exceptions import HTTPException
+from typing_extensions import Annotated
+
+from invokeai.app.services.model_install import ModelInstallJob, ModelSource
+from invokeai.app.services.model_records import (
+    DuplicateModelException,
+    InvalidModelException,
+    ModelRecordOrderBy,
+    ModelSummary,
+    UnknownModelException,
+)
+from invokeai.app.services.shared.pagination import PaginatedResults
+from invokeai.backend.model_manager.config import (
+    AnyModelConfig,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+from invokeai.backend.model_manager.merge import MergeInterpolationMethod, ModelMerger
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
+
+from ..dependencies import ApiDependencies
+
+model_records_router = APIRouter(prefix="/v1/model/record", tags=["model_manager_v2_unstable"])
+
+
+class ModelsList(BaseModel):
+    """Return list of configs."""
+
+    models: List[AnyModelConfig]
+
+    model_config = ConfigDict(use_enum_values=True)
+
+
+class ModelTagSet(BaseModel):
+    """Return tags for a set of models."""
+
+    key: str
+    name: str
+    author: str
+    tags: Set[str]
+
+
+@model_records_router.get(
+    "/",
+    operation_id="list_model_records",
+)
+async def list_model_records(
+    base_models: Optional[List[BaseModelType]] = Query(default=None, description="Base models to include"),
+    model_type: Optional[ModelType] = Query(default=None, description="The type of model to get"),
+    model_name: Optional[str] = Query(default=None, description="Exact match on the name of the model"),
+    model_format: Optional[ModelFormat] = Query(
+        default=None, description="Exact match on the format of the model (e.g. 'diffusers')"
+    ),
+) -> ModelsList:
+    """Get a list of models."""
+    record_store = ApiDependencies.invoker.services.model_records
+    found_models: list[AnyModelConfig] = []
+    if base_models:
+        for base_model in base_models:
+            found_models.extend(
+                record_store.search_by_attr(
+                    base_model=base_model, model_type=model_type, model_name=model_name, model_format=model_format
+                )
+            )
+    else:
+        found_models.extend(
+            record_store.search_by_attr(model_type=model_type, model_name=model_name, model_format=model_format)
+        )
+    return ModelsList(models=found_models)
+
+
+@model_records_router.get(
+    "/i/{key}",
+    operation_id="get_model_record",
+    responses={
+        200: {"description": "Success"},
+        400: {"description": "Bad request"},
+        404: {"description": "The model could not be found"},
+    },
+)
+async def get_model_record(
+    key: str = Path(description="Key of the model record to fetch."),
+) -> AnyModelConfig:
+    """Get a model record"""
+    record_store = ApiDependencies.invoker.services.model_records
+    try:
+        return record_store.get_model(key)
+    except UnknownModelException as e:
+        raise HTTPException(status_code=404, detail=str(e))
+
+
+@model_records_router.get("/meta", operation_id="list_model_summary")
+async def list_model_summary(
+    page: int = Query(default=0, description="The page to get"),
+    per_page: int = Query(default=10, description="The number of models per page"),
+    order_by: ModelRecordOrderBy = Query(default=ModelRecordOrderBy.Default, description="The attribute to order by"),
+) -> PaginatedResults[ModelSummary]:
+    """Gets a page of model summary data."""
+    return ApiDependencies.invoker.services.model_records.list_models(page=page, per_page=per_page, order_by=order_by)
+
+
+@model_records_router.get(
+    "/meta/i/{key}",
+    operation_id="get_model_metadata",
+    responses={
+        200: {"description": "Success"},
+        400: {"description": "Bad request"},
+        404: {"description": "No metadata available"},
+    },
+)
+async def get_model_metadata(
+    key: str = Path(description="Key of the model repo metadata to fetch."),
+) -> Optional[AnyModelRepoMetadata]:
+    """Get a model metadata object."""
+    record_store = ApiDependencies.invoker.services.model_records
+    result = record_store.get_metadata(key)
+    if not result:
+        raise HTTPException(status_code=404, detail="No metadata for a model with this key")
+    return result
+
+
+@model_records_router.get(
+    "/tags",
+    operation_id="list_tags",
+)
+async def list_tags() -> Set[str]:
+    """Get a unique set of all the model tags."""
+    record_store = ApiDependencies.invoker.services.model_records
+    return record_store.list_tags()
+
+
+@model_records_router.get(
+    "/tags/search",
+    operation_id="search_by_metadata_tags",
+)
+async def search_by_metadata_tags(
+    tags: Set[str] = Query(default=None, description="Tags to search for"),
+) -> ModelsList:
+    """Get a list of models."""
+    record_store = ApiDependencies.invoker.services.model_records
+    results = record_store.search_by_metadata_tag(tags)
+    return ModelsList(models=results)
+
+
+@model_records_router.patch(
+    "/i/{key}",
+    operation_id="update_model_record",
+    responses={
+        200: {"description": "The model was updated successfully"},
+        400: {"description": "Bad request"},
+        404: {"description": "The model could not be found"},
+        409: {"description": "There is already a model corresponding to the new name"},
+    },
+    status_code=200,
+    response_model=AnyModelConfig,
+)
+async def update_model_record(
+    key: Annotated[str, Path(description="Unique key of model")],
+    info: Annotated[AnyModelConfig, Body(description="Model config", discriminator="type")],
+) -> AnyModelConfig:
+    """Update model contents with a new config. If the model name or base fields are changed, then the model is renamed."""
+    logger = ApiDependencies.invoker.services.logger
+    record_store = ApiDependencies.invoker.services.model_records
+    try:
+        model_response = record_store.update_model(key, config=info)
+        logger.info(f"Updated model: {key}")
+    except UnknownModelException as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    except ValueError as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=409, detail=str(e))
+    return model_response
+
+
+@model_records_router.delete(
+    "/i/{key}",
+    operation_id="del_model_record",
+    responses={
+        204: {"description": "Model deleted successfully"},
+        404: {"description": "Model not found"},
+    },
+    status_code=204,
+)
+async def del_model_record(
+    key: str = Path(description="Unique key of model to remove from model registry."),
+) -> Response:
+    """
+    Delete model record from database.
+
+    The configuration record will be removed. The corresponding weights files will be
+    deleted as well if they reside within the InvokeAI "models" directory.
+    """
+    logger = ApiDependencies.invoker.services.logger
+
+    try:
+        installer = ApiDependencies.invoker.services.model_install
+        installer.delete(key)
+        logger.info(f"Deleted model: {key}")
+        return Response(status_code=204)
+    except UnknownModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=404, detail=str(e))
+
+
+@model_records_router.post(
+    "/i/",
+    operation_id="add_model_record",
+    responses={
+        201: {"description": "The model added successfully"},
+        409: {"description": "There is already a model corresponding to this path or repo_id"},
+        415: {"description": "Unrecognized file/folder format"},
+    },
+    status_code=201,
+)
+async def add_model_record(
+    config: Annotated[AnyModelConfig, Body(description="Model config", discriminator="type")],
+) -> AnyModelConfig:
+    """Add a model using the configuration information appropriate for its type."""
+    logger = ApiDependencies.invoker.services.logger
+    record_store = ApiDependencies.invoker.services.model_records
+    if config.key == "<NOKEY>":
+        config.key = sha1(randbytes(100)).hexdigest()
+        logger.info(f"Created model {config.key} for {config.name}")
+    try:
+        record_store.add_model(config.key, config)
+    except DuplicateModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=409, detail=str(e))
+    except InvalidModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=415)
+
+    # now fetch it out
+    return record_store.get_model(config.key)
+
+
+@model_records_router.post(
+    "/import",
+    operation_id="import_model_record",
+    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:
+    """Add 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_install
+        result: ModelInstallJob = installer.import_model(
+            source=source,
+            config=config,
+        )
+        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_records_router.get(
+    "/import",
+    operation_id="list_model_install_jobs",
+)
+async def list_model_install_jobs() -> List[ModelInstallJob]:
+    """Return list of model install jobs."""
+    jobs: List[ModelInstallJob] = ApiDependencies.invoker.services.model_install.list_jobs()
+    return jobs
+
+
+@model_records_router.get(
+    "/import/{id}",
+    operation_id="get_model_install_job",
+    responses={
+        200: {"description": "Success"},
+        404: {"description": "No such job"},
+    },
+)
+async def get_model_install_job(id: int = Path(description="Model install id")) -> ModelInstallJob:
+    """Return model install job corresponding to the given source."""
+    try:
+        return ApiDependencies.invoker.services.model_install.get_job_by_id(id)
+    except ValueError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+
+
+@model_records_router.delete(
+    "/import/{id}",
+    operation_id="cancel_model_install_job",
+    responses={
+        201: {"description": "The job was cancelled successfully"},
+        415: {"description": "No such job"},
+    },
+    status_code=201,
+)
+async def cancel_model_install_job(id: int = Path(description="Model install job ID")) -> None:
+    """Cancel the model install job(s) corresponding to the given job ID."""
+    installer = ApiDependencies.invoker.services.model_install
+    try:
+        job = installer.get_job_by_id(id)
+    except ValueError as e:
+        raise HTTPException(status_code=415, detail=str(e))
+    installer.cancel_job(job)
+
+
+@model_records_router.patch(
+    "/import",
+    operation_id="prune_model_install_jobs",
+    responses={
+        204: {"description": "All completed and errored jobs have been pruned"},
+        400: {"description": "Bad request"},
+    },
+)
+async def prune_model_install_jobs() -> Response:
+    """Prune all completed and errored jobs from the install job list."""
+    ApiDependencies.invoker.services.model_install.prune_jobs()
+    return Response(status_code=204)
+
+
+@model_records_router.patch(
+    "/sync",
+    operation_id="sync_models_to_config",
+    responses={
+        204: {"description": "Model config record database resynced with files on disk"},
+        400: {"description": "Bad request"},
+    },
+)
+async def sync_models_to_config() -> Response:
+    """
+    Traverse the models and autoimport directories.
+
+    Model files without a corresponding
+    record in the database are added. Orphan records without a models file are deleted.
+    """
+    ApiDependencies.invoker.services.model_install.sync_to_config()
+    return Response(status_code=204)
+
+
+@model_records_router.put(
+    "/merge",
+    operation_id="merge",
+)
+async def merge(
+    keys: List[str] = Body(description="Keys for two to three models to merge", min_length=2, max_length=3),
+    merged_model_name: Optional[str] = Body(description="Name of destination model", default=None),
+    alpha: float = Body(description="Alpha weighting strength to apply to 2d and 3d models", default=0.5),
+    force: bool = Body(
+        description="Force merging of models created with different versions of diffusers",
+        default=False,
+    ),
+    interp: Optional[MergeInterpolationMethod] = Body(description="Interpolation method", default=None),
+    merge_dest_directory: Optional[str] = Body(
+        description="Save the merged model to the designated directory (with 'merged_model_name' appended)",
+        default=None,
+    ),
+) -> AnyModelConfig:
+    """
+    Merge diffusers models.
+
+        keys: List of 2-3 model keys to merge together. All models must use the same base type.
+        merged_model_name: Name for the merged model [Concat model names]
+        alpha: Alpha value (0.0-1.0). Higher values give more weight to the second model [0.5]
+        force: If true, force the merge even if the models were generated by different versions of the diffusers library [False]
+        interp: Interpolation method. One of "weighted_sum", "sigmoid", "inv_sigmoid" or "add_difference" [weighted_sum]
+        merge_dest_directory: Specify a directory to store the merged model in [models directory]
+    """
+    print(f"here i am, keys={keys}")
+    logger = ApiDependencies.invoker.services.logger
+    try:
+        logger.info(f"Merging models: {keys} into {merge_dest_directory or '<MODELS>'}/{merged_model_name}")
+        dest = pathlib.Path(merge_dest_directory) if merge_dest_directory else None
+        installer = ApiDependencies.invoker.services.model_install
+        merger = ModelMerger(installer)
+        model_names = [installer.record_store.get_model(x).name for x in keys]
+        response = merger.merge_diffusion_models_and_save(
+            model_keys=keys,
+            merged_model_name=merged_model_name or "+".join(model_names),
+            alpha=alpha,
+            interp=interp,
+            force=force,
+            merge_dest_directory=dest,
+        )
+    except UnknownModelException:
+        raise HTTPException(
+            status_code=404,
+            detail=f"One or more of the models '{keys}' not found",
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    return response

invokeai/app/api/routers/models.py

@@ -0,0 +1,427 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654), 2023 Kent Keirsey (https://github.com/hipsterusername), 2023 Lincoln D. Stein
+
+import pathlib
+from typing import Annotated, List, Literal, Optional, Union
+
+from fastapi import Body, Path, Query, Response
+from fastapi.routing import APIRouter
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
+from starlette.exceptions import HTTPException
+
+from invokeai.backend import BaseModelType, ModelType
+from invokeai.backend.model_management import MergeInterpolationMethod
+from invokeai.backend.model_management.models import (
+    OPENAPI_MODEL_CONFIGS,
+    InvalidModelException,
+    ModelNotFoundException,
+    SchedulerPredictionType,
+)
+
+from ..dependencies import ApiDependencies
+
+models_router = APIRouter(prefix="/v1/models", tags=["models"])
+
+UpdateModelResponse = Union[tuple(OPENAPI_MODEL_CONFIGS)]
+UpdateModelResponseValidator = TypeAdapter(UpdateModelResponse)
+
+ImportModelResponse = Union[tuple(OPENAPI_MODEL_CONFIGS)]
+ImportModelResponseValidator = TypeAdapter(ImportModelResponse)
+
+ConvertModelResponse = Union[tuple(OPENAPI_MODEL_CONFIGS)]
+ConvertModelResponseValidator = TypeAdapter(ConvertModelResponse)
+
+MergeModelResponse = Union[tuple(OPENAPI_MODEL_CONFIGS)]
+ImportModelAttributes = Union[tuple(OPENAPI_MODEL_CONFIGS)]
+
+
+class ModelsList(BaseModel):
+    models: list[Union[tuple(OPENAPI_MODEL_CONFIGS)]]
+
+    model_config = ConfigDict(use_enum_values=True)
+
+
+ModelsListValidator = TypeAdapter(ModelsList)
+
+
+@models_router.get(
+    "/",
+    operation_id="list_models",
+    responses={200: {"model": ModelsList}},
+)
+async def list_models(
+    base_models: Optional[List[BaseModelType]] = Query(default=None, description="Base models to include"),
+    model_type: Optional[ModelType] = Query(default=None, description="The type of model to get"),
+) -> ModelsList:
+    """Gets a list of models"""
+    if base_models and len(base_models) > 0:
+        models_raw = []
+        for base_model in base_models:
+            models_raw.extend(ApiDependencies.invoker.services.model_manager.list_models(base_model, model_type))
+    else:
+        models_raw = ApiDependencies.invoker.services.model_manager.list_models(None, model_type)
+    models = ModelsListValidator.validate_python({"models": models_raw})
+    return models
+
+
+@models_router.patch(
+    "/{base_model}/{model_type}/{model_name}",
+    operation_id="update_model",
+    responses={
+        200: {"description": "The model was updated successfully"},
+        400: {"description": "Bad request"},
+        404: {"description": "The model could not be found"},
+        409: {"description": "There is already a model corresponding to the new name"},
+    },
+    status_code=200,
+    response_model=UpdateModelResponse,
+)
+async def update_model(
+    base_model: BaseModelType = Path(description="Base model"),
+    model_type: ModelType = Path(description="The type of model"),
+    model_name: str = Path(description="model name"),
+    info: Union[tuple(OPENAPI_MODEL_CONFIGS)] = Body(description="Model configuration"),
+) -> UpdateModelResponse:
+    """Update model contents with a new config. If the model name or base fields are changed, then the model is renamed."""
+    logger = ApiDependencies.invoker.services.logger
+
+    try:
+        previous_info = ApiDependencies.invoker.services.model_manager.list_model(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+        )
+
+        # rename operation requested
+        if info.model_name != model_name or info.base_model != base_model:
+            ApiDependencies.invoker.services.model_manager.rename_model(
+                base_model=base_model,
+                model_type=model_type,
+                model_name=model_name,
+                new_name=info.model_name,
+                new_base=info.base_model,
+            )
+            logger.info(f"Successfully renamed {base_model.value}/{model_name}=>{info.base_model}/{info.model_name}")
+            # update information to support an update of attributes
+            model_name = info.model_name
+            base_model = info.base_model
+            new_info = ApiDependencies.invoker.services.model_manager.list_model(
+                model_name=model_name,
+                base_model=base_model,
+                model_type=model_type,
+            )
+            if new_info.get("path") != previous_info.get(
+                "path"
+            ):  # model manager moved model path during rename - don't overwrite it
+                info.path = new_info.get("path")
+
+        # replace empty string values with None/null to avoid phenomenon of vae: ''
+        info_dict = info.model_dump()
+        info_dict = {x: info_dict[x] if info_dict[x] else None for x in info_dict.keys()}
+
+        ApiDependencies.invoker.services.model_manager.update_model(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+            model_attributes=info_dict,
+        )
+
+        model_raw = ApiDependencies.invoker.services.model_manager.list_model(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+        )
+        model_response = UpdateModelResponseValidator.validate_python(model_raw)
+    except ModelNotFoundException as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    except ValueError as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=409, detail=str(e))
+    except Exception as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=400, detail=str(e))
+
+    return model_response
+
+
+@models_router.post(
+    "/import",
+    operation_id="import_model",
+    responses={
+        201: {"description": "The model imported successfully"},
+        404: {"description": "The model could not be found"},
+        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,
+    response_model=ImportModelResponse,
+)
+async def import_model(
+    location: str = Body(description="A model path, repo_id or URL to import"),
+    prediction_type: Optional[Literal["v_prediction", "epsilon", "sample"]] = Body(
+        description="Prediction type for SDv2 checkpoints and rare SDv1 checkpoints",
+        default=None,
+    ),
+) -> ImportModelResponse:
+    """Add a model using its local path, repo_id, or remote URL. Model characteristics will be probed and configured automatically"""
+
+    location = location.strip("\"' ")
+    items_to_import = {location}
+    prediction_types = {x.value: x for x in SchedulerPredictionType}
+    logger = ApiDependencies.invoker.services.logger
+
+    try:
+        installed_models = ApiDependencies.invoker.services.model_manager.heuristic_import(
+            items_to_import=items_to_import,
+            prediction_type_helper=lambda x: prediction_types.get(prediction_type),
+        )
+        info = installed_models.get(location)
+
+        if not info:
+            logger.error("Import failed")
+            raise HTTPException(status_code=415)
+
+        logger.info(f"Successfully imported {location}, got {info}")
+        model_raw = ApiDependencies.invoker.services.model_manager.list_model(
+            model_name=info.name, base_model=info.base_model, model_type=info.model_type
+        )
+        return ImportModelResponseValidator.validate_python(model_raw)
+
+    except ModelNotFoundException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=404, 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))
+
+
+@models_router.post(
+    "/add",
+    operation_id="add_model",
+    responses={
+        201: {"description": "The model added successfully"},
+        404: {"description": "The model could not be found"},
+        424: {"description": "The model appeared to add 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,
+    response_model=ImportModelResponse,
+)
+async def add_model(
+    info: Union[tuple(OPENAPI_MODEL_CONFIGS)] = Body(description="Model configuration"),
+) -> ImportModelResponse:
+    """Add a model using the configuration information appropriate for its type. Only local models can be added by path"""
+
+    logger = ApiDependencies.invoker.services.logger
+
+    try:
+        ApiDependencies.invoker.services.model_manager.add_model(
+            info.model_name,
+            info.base_model,
+            info.model_type,
+            model_attributes=info.model_dump(),
+        )
+        logger.info(f"Successfully added {info.model_name}")
+        model_raw = ApiDependencies.invoker.services.model_manager.list_model(
+            model_name=info.model_name,
+            base_model=info.base_model,
+            model_type=info.model_type,
+        )
+        return ImportModelResponseValidator.validate_python(model_raw)
+    except ModelNotFoundException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=404, detail=str(e))
+    except ValueError as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=409, detail=str(e))
+
+
+@models_router.delete(
+    "/{base_model}/{model_type}/{model_name}",
+    operation_id="del_model",
+    responses={
+        204: {"description": "Model deleted successfully"},
+        404: {"description": "Model not found"},
+    },
+    status_code=204,
+    response_model=None,
+)
+async def delete_model(
+    base_model: BaseModelType = Path(description="Base model"),
+    model_type: ModelType = Path(description="The type of model"),
+    model_name: str = Path(description="model name"),
+) -> Response:
+    """Delete Model"""
+    logger = ApiDependencies.invoker.services.logger
+
+    try:
+        ApiDependencies.invoker.services.model_manager.del_model(
+            model_name, base_model=base_model, model_type=model_type
+        )
+        logger.info(f"Deleted model: {model_name}")
+        return Response(status_code=204)
+    except ModelNotFoundException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=404, detail=str(e))
+
+
+@models_router.put(
+    "/convert/{base_model}/{model_type}/{model_name}",
+    operation_id="convert_model",
+    responses={
+        200: {"description": "Model converted successfully"},
+        400: {"description": "Bad request"},
+        404: {"description": "Model not found"},
+    },
+    status_code=200,
+    response_model=ConvertModelResponse,
+)
+async def convert_model(
+    base_model: BaseModelType = Path(description="Base model"),
+    model_type: ModelType = Path(description="The type of model"),
+    model_name: str = Path(description="model name"),
+    convert_dest_directory: Optional[str] = Query(
+        default=None, description="Save the converted model to the designated directory"
+    ),
+) -> ConvertModelResponse:
+    """Convert a checkpoint model into a diffusers model, optionally saving to the indicated destination directory, or `models` if none."""
+    logger = ApiDependencies.invoker.services.logger
+    try:
+        logger.info(f"Converting model: {model_name}")
+        dest = pathlib.Path(convert_dest_directory) if convert_dest_directory else None
+        ApiDependencies.invoker.services.model_manager.convert_model(
+            model_name,
+            base_model=base_model,
+            model_type=model_type,
+            convert_dest_directory=dest,
+        )
+        model_raw = ApiDependencies.invoker.services.model_manager.list_model(
+            model_name, base_model=base_model, model_type=model_type
+        )
+        response = ConvertModelResponseValidator.validate_python(model_raw)
+    except ModelNotFoundException as e:
+        raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found: {str(e)}")
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    return response
+
+
+@models_router.get(
+    "/search",
+    operation_id="search_for_models",
+    responses={
+        200: {"description": "Directory searched successfully"},
+        404: {"description": "Invalid directory path"},
+    },
+    status_code=200,
+    response_model=List[pathlib.Path],
+)
+async def search_for_models(
+    search_path: pathlib.Path = Query(description="Directory path to search for models"),
+) -> List[pathlib.Path]:
+    if not search_path.is_dir():
+        raise HTTPException(
+            status_code=404,
+            detail=f"The search path '{search_path}' does not exist or is not directory",
+        )
+    return ApiDependencies.invoker.services.model_manager.search_for_models(search_path)
+
+
+@models_router.get(
+    "/ckpt_confs",
+    operation_id="list_ckpt_configs",
+    responses={
+        200: {"description": "paths retrieved successfully"},
+    },
+    status_code=200,
+    response_model=List[pathlib.Path],
+)
+async def list_ckpt_configs() -> List[pathlib.Path]:
+    """Return a list of the legacy checkpoint configuration files stored in `ROOT/configs/stable-diffusion`, relative to ROOT."""
+    return ApiDependencies.invoker.services.model_manager.list_checkpoint_configs()
+
+
+@models_router.post(
+    "/sync",
+    operation_id="sync_to_config",
+    responses={
+        201: {"description": "synchronization successful"},
+    },
+    status_code=201,
+    response_model=bool,
+)
+async def sync_to_config() -> bool:
+    """Call after making changes to models.yaml, autoimport directories or models directory to synchronize
+    in-memory data structures with disk data structures."""
+    ApiDependencies.invoker.services.model_manager.sync_to_config()
+    return True
+
+
+# There's some weird pydantic-fastapi behaviour that requires this to be a separate class
+# TODO: After a few updates, see if it works inside the route operation handler?
+class MergeModelsBody(BaseModel):
+    model_names: List[str] = Field(description="model name", min_length=2, max_length=3)
+    merged_model_name: Optional[str] = Field(description="Name of destination model")
+    alpha: Optional[float] = Field(description="Alpha weighting strength to apply to 2d and 3d models", default=0.5)
+    interp: Optional[MergeInterpolationMethod] = Field(description="Interpolation method")
+    force: Optional[bool] = Field(
+        description="Force merging of models created with different versions of diffusers",
+        default=False,
+    )
+
+    merge_dest_directory: Optional[str] = Field(
+        description="Save the merged model to the designated directory (with 'merged_model_name' appended)",
+        default=None,
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+@models_router.put(
+    "/merge/{base_model}",
+    operation_id="merge_models",
+    responses={
+        200: {"description": "Model converted successfully"},
+        400: {"description": "Incompatible models"},
+        404: {"description": "One or more models not found"},
+    },
+    status_code=200,
+    response_model=MergeModelResponse,
+)
+async def merge_models(
+    body: Annotated[MergeModelsBody, Body(description="Model configuration", embed=True)],
+    base_model: BaseModelType = Path(description="Base model"),
+) -> MergeModelResponse:
+    """Convert a checkpoint model into a diffusers model"""
+    logger = ApiDependencies.invoker.services.logger
+    try:
+        logger.info(
+            f"Merging models: {body.model_names} into {body.merge_dest_directory or '<MODELS>'}/{body.merged_model_name}"
+        )
+        dest = pathlib.Path(body.merge_dest_directory) if body.merge_dest_directory else None
+        result = ApiDependencies.invoker.services.model_manager.merge_models(
+            model_names=body.model_names,
+            base_model=base_model,
+            merged_model_name=body.merged_model_name or "+".join(body.model_names),
+            alpha=body.alpha,
+            interp=body.interp,
+            force=body.force,
+            merge_dest_directory=dest,
+        )
+        model_raw = ApiDependencies.invoker.services.model_manager.list_model(
+            result.name,
+            base_model=base_model,
+            model_type=ModelType.Main,
+        )
+        response = ConvertModelResponseValidator.validate_python(model_raw)
+    except ModelNotFoundException:
+        raise HTTPException(
+            status_code=404,
+            detail=f"One or more of the models '{body.model_names}' not found",
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    return response

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")
+        self.__app = ASGIApp(socketio_server=self.__sio, socketio_path="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

@@ -6,7 +6,6 @@ import sys
 from invokeai.app.api.no_cache_staticfiles import NoCacheStaticFiles
 from invokeai.version.invokeai_version import __version__
 
-from .invocations.fields import InputFieldJSONSchemaExtra, OutputFieldJSONSchemaExtra
 from .services.config import InvokeAIAppConfig
 
 app_config = InvokeAIAppConfig.get_config()
@@ -48,14 +47,18 @@ if True:  # hack to make flake8 happy with imports coming after setting up the c
         boards,
         download_queue,
         images,
-        model_manager,
+        model_records,
+        models,
         session_queue,
+        sessions,
         utilities,
         workflows,
     )
     from .api.sockets import SocketIO
     from .invocations.baseinvocation import (
         BaseInvocation,
+        InputFieldJSONSchemaExtra,
+        OutputFieldJSONSchemaExtra,
         UIConfigBase,
     )
 
@@ -109,8 +112,11 @@ 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(models.models_router, prefix="/api")
+app.include_router(model_records.model_records_router, prefix="/api")
 app.include_router(download_queue.download_queue_router, prefix="/api")
 app.include_router(images.images_router, prefix="/api")
 app.include_router(boards.boards_router, prefix="/api")
@@ -148,8 +154,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,24 +176,23 @@ 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
-    #
-    # from invokeai.backend.model_manager import get_model_config_formats
-    # formats = get_model_config_formats()
-    # for model_config_name, enum_set in formats.items():
+    from invokeai.backend.model_management.models import get_model_config_enums
 
-    #     if model_config_name in openapi_schema["components"]["schemas"]:
-    #         # print(f"Config with name {name} already defined")
-    #         continue
+    for model_config_format_enum in set(get_model_config_enums()):
+        name = model_config_format_enum.__qualname__
 
-    #     openapi_schema["components"]["schemas"][model_config_name] = {
-    #         "title": model_config_name,
-    #         "description": "An enumeration.",
-    #         "type": "string",
-    #         "enum": [v.value for v in enum_set],
-    #     }
+        if name in openapi_schema["components"]["schemas"]:
+            # print(f"Config with name {name} already defined")
+            continue
+
+        openapi_schema["components"]["schemas"][name] = {
+            "title": name,
+            "description": "An enumeration.",
+            "type": "string",
+            "enum": [v.value for v in model_config_format_enum],
+        }
 
     app.openapi_schema = openapi_schema
     return app.openapi_schema

invokeai/app/invocations/baseinvocation.py

@@ -8,33 +8,17 @@ 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.fields import FieldInfo
+from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, create_model
+from pydantic.fields import FieldInfo, _Unset
 from pydantic_core import PydanticUndefined
-from typing_extensions import TypeAliasType
 
-from invokeai.app.invocations.fields import (
-    FieldKind,
-    Input,
-)
 from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
+from invokeai.app.shared.fields import FieldDescriptions
 from invokeai.app.util.metaenum import MetaEnum
 from invokeai.app.util.misc import uuid_string
 from invokeai.backend.util.logging import InvokeAILogger
@@ -68,6 +52,393 @@ class Classification(str, Enum, metaclass=MetaEnum):
     Prototype = "prototype"
 
 
+class Input(str, Enum, metaclass=MetaEnum):
+    """
+    The type of input a field accepts.
+    - `Input.Direct`: The field must have its value provided directly, when the invocation and field \
+      are instantiated.
+    - `Input.Connection`: The field must have its value provided by a connection.
+    - `Input.Any`: The field may have its value provided either directly or by a connection.
+    """
+
+    Connection = "connection"
+    Direct = "direct"
+    Any = "any"
+
+
+class FieldKind(str, Enum, metaclass=MetaEnum):
+    """
+    The kind of field.
+    - `Input`: An input field on a node.
+    - `Output`: An output field on a node.
+    - `Internal`: A field which is treated as an input, but cannot be used in node definitions. Metadata is
+    one example. It is provided to nodes via the WithMetadata class, and we want to reserve the field name
+    "metadata" for this on all nodes. `FieldKind` is used to short-circuit the field name validation logic,
+    allowing "metadata" for that field.
+    - `NodeAttribute`: The field is a node attribute. These are fields which are not inputs or outputs,
+    but which are used to store information about the node. For example, the `id` and `type` fields are node
+    attributes.
+
+    The presence of this in `json_schema_extra["field_kind"]` is used when initializing node schemas on app
+    startup, and when generating the OpenAPI schema for the workflow editor.
+    """
+
+    Input = "input"
+    Output = "output"
+    Internal = "internal"
+    NodeAttribute = "node_attribute"
+
+
+class UIType(str, Enum, metaclass=MetaEnum):
+    """
+    Type hints for the UI for situations in which the field type is not enough to infer the correct UI type.
+
+    - Model Fields
+    The most common node-author-facing use will be for model fields. Internally, there is no difference
+    between SD-1, SD-2 and SDXL model fields - they all use the class `MainModelField`. To ensure the
+    base-model-specific UI is rendered, use e.g. `ui_type=UIType.SDXLMainModelField` to indicate that
+    the field is an SDXL main model field.
+
+    - Any Field
+    We cannot infer the usage of `typing.Any` via schema parsing, so you *must* use `ui_type=UIType.Any` to
+    indicate that the field accepts any type. Use with caution. This cannot be used on outputs.
+
+    - Scheduler Field
+    Special handling in the UI is needed for this field, which otherwise would be parsed as a plain enum field.
+
+    - Internal Fields
+    Similar to the Any Field, the `collect` and `iterate` nodes make use of `typing.Any`. To facilitate
+    handling these types in the client, we use `UIType._Collection` and `UIType._CollectionItem`. These
+    should not be used by node authors.
+
+    - DEPRECATED Fields
+    These types are deprecated and should not be used by node authors. A warning will be logged if one is
+    used, and the type will be ignored. They are included here for backwards compatibility.
+    """
+
+    # region Model Field Types
+    SDXLMainModel = "SDXLMainModelField"
+    SDXLRefinerModel = "SDXLRefinerModelField"
+    ONNXModel = "ONNXModelField"
+    VaeModel = "VAEModelField"
+    LoRAModel = "LoRAModelField"
+    ControlNetModel = "ControlNetModelField"
+    IPAdapterModel = "IPAdapterModelField"
+    # endregion
+
+    # region Misc Field Types
+    Scheduler = "SchedulerField"
+    Any = "AnyField"
+    # endregion
+
+    # region Internal Field Types
+    _Collection = "CollectionField"
+    _CollectionItem = "CollectionItemField"
+    # endregion
+
+    # region DEPRECATED
+    Boolean = "DEPRECATED_Boolean"
+    Color = "DEPRECATED_Color"
+    Conditioning = "DEPRECATED_Conditioning"
+    Control = "DEPRECATED_Control"
+    Float = "DEPRECATED_Float"
+    Image = "DEPRECATED_Image"
+    Integer = "DEPRECATED_Integer"
+    Latents = "DEPRECATED_Latents"
+    String = "DEPRECATED_String"
+    BooleanCollection = "DEPRECATED_BooleanCollection"
+    ColorCollection = "DEPRECATED_ColorCollection"
+    ConditioningCollection = "DEPRECATED_ConditioningCollection"
+    ControlCollection = "DEPRECATED_ControlCollection"
+    FloatCollection = "DEPRECATED_FloatCollection"
+    ImageCollection = "DEPRECATED_ImageCollection"
+    IntegerCollection = "DEPRECATED_IntegerCollection"
+    LatentsCollection = "DEPRECATED_LatentsCollection"
+    StringCollection = "DEPRECATED_StringCollection"
+    BooleanPolymorphic = "DEPRECATED_BooleanPolymorphic"
+    ColorPolymorphic = "DEPRECATED_ColorPolymorphic"
+    ConditioningPolymorphic = "DEPRECATED_ConditioningPolymorphic"
+    ControlPolymorphic = "DEPRECATED_ControlPolymorphic"
+    FloatPolymorphic = "DEPRECATED_FloatPolymorphic"
+    ImagePolymorphic = "DEPRECATED_ImagePolymorphic"
+    IntegerPolymorphic = "DEPRECATED_IntegerPolymorphic"
+    LatentsPolymorphic = "DEPRECATED_LatentsPolymorphic"
+    StringPolymorphic = "DEPRECATED_StringPolymorphic"
+    MainModel = "DEPRECATED_MainModel"
+    UNet = "DEPRECATED_UNet"
+    Vae = "DEPRECATED_Vae"
+    CLIP = "DEPRECATED_CLIP"
+    Collection = "DEPRECATED_Collection"
+    CollectionItem = "DEPRECATED_CollectionItem"
+    Enum = "DEPRECATED_Enum"
+    WorkflowField = "DEPRECATED_WorkflowField"
+    IsIntermediate = "DEPRECATED_IsIntermediate"
+    BoardField = "DEPRECATED_BoardField"
+    MetadataItem = "DEPRECATED_MetadataItem"
+    MetadataItemCollection = "DEPRECATED_MetadataItemCollection"
+    MetadataItemPolymorphic = "DEPRECATED_MetadataItemPolymorphic"
+    MetadataDict = "DEPRECATED_MetadataDict"
+    # endregion
+
+
+class UIComponent(str, Enum, metaclass=MetaEnum):
+    """
+    The type of UI component to use for a field, used to override the default components, which are
+    inferred from the field type.
+    """
+
+    None_ = "none"
+    Textarea = "textarea"
+    Slider = "slider"
+
+
+class InputFieldJSONSchemaExtra(BaseModel):
+    """
+    Extra attributes to be added to input fields and their OpenAPI schema. Used during graph execution,
+    and by the workflow editor during schema parsing and UI rendering.
+    """
+
+    input: Input
+    orig_required: bool
+    field_kind: FieldKind
+    default: Optional[Any] = None
+    orig_default: Optional[Any] = None
+    ui_hidden: bool = False
+    ui_type: Optional[UIType] = None
+    ui_component: Optional[UIComponent] = None
+    ui_order: Optional[int] = None
+    ui_choice_labels: Optional[dict[str, str]] = None
+
+    model_config = ConfigDict(
+        validate_assignment=True,
+        json_schema_serialization_defaults_required=True,
+    )
+
+
+class OutputFieldJSONSchemaExtra(BaseModel):
+    """
+    Extra attributes to be added to input fields and their OpenAPI schema. Used by the workflow editor
+    during schema parsing and UI rendering.
+    """
+
+    field_kind: FieldKind
+    ui_hidden: bool
+    ui_type: Optional[UIType]
+    ui_order: Optional[int]
+
+    model_config = ConfigDict(
+        validate_assignment=True,
+        json_schema_serialization_defaults_required=True,
+    )
+
+
+def InputField(
+    # copied from pydantic's Field
+    # TODO: Can we support default_factory?
+    default: Any = _Unset,
+    default_factory: Callable[[], Any] | None = _Unset,
+    title: str | None = _Unset,
+    description: str | None = _Unset,
+    pattern: str | None = _Unset,
+    strict: bool | None = _Unset,
+    gt: float | None = _Unset,
+    ge: float | None = _Unset,
+    lt: float | None = _Unset,
+    le: float | None = _Unset,
+    multiple_of: float | None = _Unset,
+    allow_inf_nan: bool | None = _Unset,
+    max_digits: int | None = _Unset,
+    decimal_places: int | None = _Unset,
+    min_length: int | None = _Unset,
+    max_length: int | None = _Unset,
+    # custom
+    input: Input = Input.Any,
+    ui_type: Optional[UIType] = None,
+    ui_component: Optional[UIComponent] = None,
+    ui_hidden: bool = False,
+    ui_order: Optional[int] = None,
+    ui_choice_labels: Optional[dict[str, str]] = None,
+) -> Any:
+    """
+    Creates an input field for an invocation.
+
+    This is a wrapper for Pydantic's [Field](https://docs.pydantic.dev/latest/api/fields/#pydantic.fields.Field) \
+    that adds a few extra parameters to support graph execution and the node editor UI.
+
+    :param Input input: [Input.Any] The kind of input this field requires. \
+      `Input.Direct` means a value must be provided on instantiation. \
+      `Input.Connection` means the value must be provided by a connection. \
+      `Input.Any` means either will do.
+
+    :param UIType ui_type: [None] Optionally provides an extra type hint for the UI. \
+      In some situations, the field's type is not enough to infer the correct UI type. \
+      For example, model selection fields should render a dropdown UI component to select a model. \
+      Internally, there is no difference between SD-1, SD-2 and SDXL model fields, they all use \
+      `MainModelField`. So to ensure the base-model-specific UI is rendered, you can use \
+      `UIType.SDXLMainModelField` to indicate that the field is an SDXL main model field.
+
+    :param UIComponent ui_component: [None] Optionally specifies a specific component to use in the UI. \
+      The UI will always render a suitable component, but sometimes you want something different than the default. \
+      For example, a `string` field will default to a single-line input, but you may want a multi-line textarea instead. \
+      For this case, you could provide `UIComponent.Textarea`.
+
+    :param bool ui_hidden: [False] Specifies whether or not this field should be hidden in the UI.
+
+    :param int ui_order: [None] Specifies the order in which this field should be rendered in the UI.
+
+    :param dict[str, str] ui_choice_labels: [None] Specifies the labels to use for the choices in an enum field.
+    """
+
+    json_schema_extra_ = InputFieldJSONSchemaExtra(
+        input=input,
+        ui_type=ui_type,
+        ui_component=ui_component,
+        ui_hidden=ui_hidden,
+        ui_order=ui_order,
+        ui_choice_labels=ui_choice_labels,
+        field_kind=FieldKind.Input,
+        orig_required=True,
+    )
+
+    """
+    There is a conflict between the typing of invocation definitions and the typing of an invocation's
+    `invoke()` function.
+
+    On instantiation of a node, the invocation definition is used to create the python class. At this time,
+    any number of fields may be optional, because they may be provided by connections.
+
+    On calling of `invoke()`, however, those fields may be required.
+
+    For example, consider an ResizeImageInvocation with an `image: ImageField` field.
+
+    `image` is required during the call to `invoke()`, but when the python class is instantiated,
+    the field may not be present. This is fine, because that image field will be provided by a
+    connection from an ancestor node, which outputs an image.
+
+    This means we want to type the `image` field as optional for the node class definition, but required
+    for the `invoke()` function.
+
+    If we use `typing.Optional` in the node class definition, the field will be typed as optional in the
+    `invoke()` method, and we'll have to do a lot of runtime checks to ensure the field is present - or
+    any static type analysis tools will complain.
+
+    To get around this, in node class definitions, we type all fields correctly for the `invoke()` function,
+    but secretly make them optional in `InputField()`. We also store the original required bool and/or default
+    value. When we call `invoke()`, we use this stored information to do an additional check on the class.
+    """
+
+    if default_factory is not _Unset and default_factory is not None:
+        default = default_factory()
+        logger.warn('"default_factory" is not supported, calling it now to set "default"')
+
+    # These are the args we may wish pass to the pydantic `Field()` function
+    field_args = {
+        "default": default,
+        "title": title,
+        "description": description,
+        "pattern": pattern,
+        "strict": strict,
+        "gt": gt,
+        "ge": ge,
+        "lt": lt,
+        "le": le,
+        "multiple_of": multiple_of,
+        "allow_inf_nan": allow_inf_nan,
+        "max_digits": max_digits,
+        "decimal_places": decimal_places,
+        "min_length": min_length,
+        "max_length": max_length,
+    }
+
+    # We only want to pass the args that were provided, otherwise the `Field()`` function won't work as expected
+    provided_args = {k: v for (k, v) in field_args.items() if v is not PydanticUndefined}
+
+    # Because we are manually making fields optional, we need to store the original required bool for reference later
+    json_schema_extra_.orig_required = default is PydanticUndefined
+
+    # Make Input.Any and Input.Connection fields optional, providing None as a default if the field doesn't already have one
+    if input is Input.Any or input is Input.Connection:
+        default_ = None if default is PydanticUndefined else default
+        provided_args.update({"default": default_})
+        if default is not PydanticUndefined:
+            # Before invoking, we'll check for the original default value and set it on the field if the field has no value
+            json_schema_extra_.default = default
+            json_schema_extra_.orig_default = default
+    elif default is not PydanticUndefined:
+        default_ = default
+        provided_args.update({"default": default_})
+        json_schema_extra_.orig_default = default_
+
+    return Field(
+        **provided_args,
+        json_schema_extra=json_schema_extra_.model_dump(exclude_none=True),
+    )
+
+
+def OutputField(
+    # copied from pydantic's Field
+    default: Any = _Unset,
+    title: str | None = _Unset,
+    description: str | None = _Unset,
+    pattern: str | None = _Unset,
+    strict: bool | None = _Unset,
+    gt: float | None = _Unset,
+    ge: float | None = _Unset,
+    lt: float | None = _Unset,
+    le: float | None = _Unset,
+    multiple_of: float | None = _Unset,
+    allow_inf_nan: bool | None = _Unset,
+    max_digits: int | None = _Unset,
+    decimal_places: int | None = _Unset,
+    min_length: int | None = _Unset,
+    max_length: int | None = _Unset,
+    # custom
+    ui_type: Optional[UIType] = None,
+    ui_hidden: bool = False,
+    ui_order: Optional[int] = None,
+) -> Any:
+    """
+    Creates an output field for an invocation output.
+
+    This is a wrapper for Pydantic's [Field](https://docs.pydantic.dev/1.10/usage/schema/#field-customization) \
+    that adds a few extra parameters to support graph execution and the node editor UI.
+
+    :param UIType ui_type: [None] Optionally provides an extra type hint for the UI. \
+      In some situations, the field's type is not enough to infer the correct UI type. \
+      For example, model selection fields should render a dropdown UI component to select a model. \
+      Internally, there is no difference between SD-1, SD-2 and SDXL model fields, they all use \
+      `MainModelField`. So to ensure the base-model-specific UI is rendered, you can use \
+      `UIType.SDXLMainModelField` to indicate that the field is an SDXL main model field.
+
+    :param bool ui_hidden: [False] Specifies whether or not this field should be hidden in the UI. \
+
+    :param int ui_order: [None] Specifies the order in which this field should be rendered in the UI. \
+    """
+    return Field(
+        default=default,
+        title=title,
+        description=description,
+        pattern=pattern,
+        strict=strict,
+        gt=gt,
+        ge=ge,
+        lt=lt,
+        le=le,
+        multiple_of=multiple_of,
+        allow_inf_nan=allow_inf_nan,
+        max_digits=max_digits,
+        decimal_places=decimal_places,
+        min_length=min_length,
+        max_length=max_length,
+        json_schema_extra=OutputFieldJSONSchemaExtra(
+            ui_type=ui_type,
+            ui_hidden=ui_hidden,
+            ui_order=ui_order,
+            field_kind=FieldKind.Output,
+        ).model_dump(exclude_none=True),
+    )
+
+
 class UIConfigBase(BaseModel):
     """
     Provides additional node configuration to the UI.
@@ -89,6 +460,33 @@ class UIConfigBase(BaseModel):
     )
 
 
+class InvocationContext:
+    """Initialized and provided to on execution of invocations."""
+
+    services: InvocationServices
+    graph_execution_state_id: str
+    queue_id: str
+    queue_item_id: int
+    queue_batch_id: str
+    workflow: Optional[WorkflowWithoutID]
+
+    def __init__(
+        self,
+        services: InvocationServices,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        workflow: Optional[WorkflowWithoutID],
+    ):
+        self.services = services
+        self.graph_execution_state_id = graph_execution_state_id
+        self.queue_id = queue_id
+        self.queue_item_id = queue_item_id
+        self.queue_batch_id = queue_batch_id
+        self.workflow = workflow
+
+
 class BaseInvocationOutput(BaseModel):
     """
     Base class for all invocation outputs.
@@ -97,7 +495,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 +507,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 +559,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 +571,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]:
@@ -244,7 +632,7 @@ class BaseInvocation(ABC, BaseModel):
         """Invoke with provided context and return outputs."""
         pass
 
-    def invoke_internal(self, context: InvocationContext, services: "InvocationServices") -> BaseInvocationOutput:
+    def invoke_internal(self, context: InvocationContext) -> BaseInvocationOutput:
         """
         Internal invoke method, calls `invoke()` after some prep.
         Handles optional fields that are required to call `invoke()` and invocation cache.
@@ -269,23 +657,23 @@ class BaseInvocation(ABC, BaseModel):
                     raise MissingInputException(self.model_fields["type"].default, field_name)
 
         # skip node cache codepath if it's disabled
-        if services.configuration.node_cache_size == 0:
+        if context.services.configuration.node_cache_size == 0:
             return self.invoke(context)
 
         output: BaseInvocationOutput
         if self.use_cache:
-            key = services.invocation_cache.create_key(self)
-            cached_value = services.invocation_cache.get(key)
+            key = context.services.invocation_cache.create_key(self)
+            cached_value = context.services.invocation_cache.get(key)
             if cached_value is None:
-                services.logger.debug(f'Invocation cache miss for type "{self.get_type()}": {self.id}')
+                context.services.logger.debug(f'Invocation cache miss for type "{self.get_type()}": {self.id}')
                 output = self.invoke(context)
-                services.invocation_cache.save(key, output)
+                context.services.invocation_cache.save(key, output)
                 return output
             else:
-                services.logger.debug(f'Invocation cache hit for type "{self.get_type()}": {self.id}')
+                context.services.logger.debug(f'Invocation cache hit for type "{self.get_type()}": {self.id}')
                 return cached_value
         else:
-            services.logger.debug(f'Skipping invocation cache for "{self.get_type()}": {self.id}')
+            context.services.logger.debug(f'Skipping invocation cache for "{self.get_type()}": {self.id}')
             return self.invoke(context)
 
     id: str = Field(
@@ -326,7 +714,9 @@ RESERVED_NODE_ATTRIBUTE_FIELD_NAMES = {
     "workflow",
 }
 
-RESERVED_INPUT_FIELD_NAMES = {"metadata", "board"}
+RESERVED_INPUT_FIELD_NAMES = {
+    "metadata",
+}
 
 RESERVED_OUTPUT_FIELD_NAMES = {"type"}
 
@@ -536,3 +926,37 @@ def invocation_output(
         return cls
 
     return wrapper
+
+
+class MetadataField(RootModel):
+    """
+    Pydantic model for metadata with custom root of type dict[str, Any].
+    Metadata is stored without a strict schema.
+    """
+
+    root: dict[str, Any] = Field(description="The metadata")
+
+
+MetadataFieldValidator = TypeAdapter(MetadataField)
+
+
+class WithMetadata(BaseModel):
+    metadata: Optional[MetadataField] = Field(
+        default=None,
+        description=FieldDescriptions.metadata,
+        json_schema_extra=InputFieldJSONSchemaExtra(
+            field_kind=FieldKind.Internal,
+            input=Input.Connection,
+            orig_required=False,
+        ).model_dump(exclude_none=True),
+    )
+
+
+class WithWorkflow:
+    workflow = None
+
+    def __init_subclass__(cls) -> None:
+        logger.warn(
+            f"{cls.__module__.split('.')[0]}.{cls.__name__}: WithWorkflow is deprecated. Use `context.workflow` to access the workflow."
+        )
+        super().__init_subclass__()

invokeai/app/invocations/collections.py

@@ -5,11 +5,9 @@ import numpy as np
 from pydantic import ValidationInfo, field_validator
 
 from invokeai.app.invocations.primitives import IntegerCollectionOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.misc import SEED_MAX
 
-from .baseinvocation import BaseInvocation, invocation
-from .fields import InputField
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, invocation
 
 
 @invocation(

invokeai/app/invocations/compel.py

@@ -1,39 +1,40 @@
-from typing import Iterator, List, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import List, Optional, Union
 
 import torch
 from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
-from transformers import CLIPTextModel, CLIPTokenizer
 
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    Input,
-    InputField,
-    OutputField,
-    UIComponent,
-)
-from invokeai.app.invocations.primitives import ConditioningOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.ti_utils import generate_ti_list
-from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_patcher import ModelPatcher
+from invokeai.app.invocations.primitives import ConditioningField, ConditioningOutput
+from invokeai.app.shared.fields import FieldDescriptions
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
-    ConditioningFieldData,
     ExtraConditioningInfo,
     SDXLConditioningInfo,
 )
-from invokeai.backend.util.devices import torch_dtype
 
+from ...backend.model_management.lora import ModelPatcher
+from ...backend.model_management.models import ModelNotFoundException, ModelType
+from ...backend.util.devices import torch_dtype
+from ..util.ti_utils import extract_ti_triggers_from_prompt
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    UIComponent,
     invocation,
     invocation_output,
 )
 from .model import ClipField
 
-# unconditioned: Optional[torch.Tensor]
+
+@dataclass
+class ConditioningFieldData:
+    conditionings: List[BasicConditioningInfo]
+    # unconditioned: Optional[torch.Tensor]
 
 
 # class ConditioningAlgo(str, Enum):
@@ -47,7 +48,7 @@ from .model import ClipField
     title="Prompt",
     tags=["prompt", "compel"],
     category="conditioning",
-    version="1.0.1",
+    version="1.0.0",
 )
 class CompelInvocation(BaseInvocation):
     """Parse prompt using compel package to conditioning."""
@@ -65,27 +66,49 @@ 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)
+        tokenizer_info = context.services.model_manager.get_model(
+            **self.clip.tokenizer.model_dump(),
+            context=context,
+        )
+        text_encoder_info = context.services.model_manager.get_model(
+            **self.clip.text_encoder.model_dump(),
+            context=context,
+        )
 
-        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
+        def _lora_loader():
             for lora in self.clip.loras:
-                lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
-                assert isinstance(lora_info.model, LoRAModelRaw)
-                yield (lora_info.model, lora.weight)
+                lora_info = context.services.model_manager.get_model(
+                    **lora.model_dump(exclude={"weight"}), context=context
+                )
+                yield (lora_info.context.model, lora.weight)
                 del lora_info
             return
 
-        # loras = [(context.models.get(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras]
+        # loras = [(context.services.model_manager.get_model(**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:
+                ti_list.append(
+                    (
+                        name,
+                        context.services.model_manager.get_model(
+                            model_name=name,
+                            base_model=self.clip.text_encoder.base_model,
+                            model_type=ModelType.TextualInversion,
+                            context=context,
+                        ).context.model,
+                    )
+                )
+            except ModelNotFoundException:
+                # 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.context.model, text_encoder_info.context.model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -93,9 +116,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.context.model, self.clip.skipped_layers),
         ):
-            assert isinstance(text_encoder, CLIPTextModel)
             compel = Compel(
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
@@ -106,7 +128,7 @@ class CompelInvocation(BaseInvocation):
 
             conjunction = Compel.parse_prompt_string(self.prompt)
 
-            if context.config.get().log_tokenization:
+            if context.services.configuration.log_tokenization:
                 log_tokenization_for_conjunction(conjunction, tokenizer)
 
             c, options = compel.build_conditioning_tensor_for_conjunction(conjunction)
@@ -127,14 +149,17 @@ class CompelInvocation(BaseInvocation):
             ]
         )
 
-        conditioning_name = context.conditioning.save(conditioning_data)
+        conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning"
+        context.services.latents.save(conditioning_name, conditioning_data)
 
-        return ConditioningOutput.build(conditioning_name)
+        return ConditioningOutput(
+            conditioning=ConditioningField(
+                conditioning_name=conditioning_name,
+            ),
+        )
 
 
 class SDXLPromptInvocationBase:
-    """Prompt processor for SDXL models."""
-
     def run_clip_compel(
         self,
         context: InvocationContext,
@@ -143,25 +168,26 @@ class SDXLPromptInvocationBase:
         get_pooled: bool,
         lora_prefix: str,
         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)
+    ):
+        tokenizer_info = context.services.model_manager.get_model(
+            **clip_field.tokenizer.model_dump(),
+            context=context,
+        )
+        text_encoder_info = context.services.model_manager.get_model(
+            **clip_field.text_encoder.model_dump(),
+            context=context,
+        )
 
         # return zero on empty
         if prompt == "" and zero_on_empty:
-            cpu_text_encoder = text_encoder_info.model
-            assert isinstance(cpu_text_encoder, torch.nn.Module)
+            cpu_text_encoder = text_encoder_info.context.model
             c = torch.zeros(
                 (
                     1,
                     cpu_text_encoder.config.max_position_embeddings,
                     cpu_text_encoder.config.hidden_size,
                 ),
-                dtype=cpu_text_encoder.dtype,
+                dtype=text_encoder_info.context.cache.precision,
             )
             if get_pooled:
                 c_pooled = torch.zeros(
@@ -172,21 +198,40 @@ class SDXLPromptInvocationBase:
                 c_pooled = None
             return c, c_pooled, None
 
-        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
+        def _lora_loader():
             for lora in clip_field.loras:
-                lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
-                lora_model = lora_info.model
-                assert isinstance(lora_model, LoRAModelRaw)
-                yield (lora_model, lora.weight)
+                lora_info = context.services.model_manager.get_model(
+                    **lora.model_dump(exclude={"weight"}), context=context
+                )
+                yield (lora_info.context.model, lora.weight)
                 del lora_info
             return
 
-        # loras = [(context.models.get(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras]
+        # loras = [(context.services.model_manager.get_model(**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_list.append(
+                    (
+                        name,
+                        context.services.model_manager.get_model(
+                            model_name=name,
+                            base_model=clip_field.text_encoder.base_model,
+                            model_type=ModelType.TextualInversion,
+                            context=context,
+                        ).context.model,
+                    )
+                )
+            except ModelNotFoundException:
+                # 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.context.model, text_encoder_info.context.model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -194,9 +239,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.context.model, clip_field.skipped_layers),
         ):
-            assert isinstance(text_encoder, CLIPTextModel)
             compel = Compel(
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
@@ -209,7 +253,7 @@ class SDXLPromptInvocationBase:
 
             conjunction = Compel.parse_prompt_string(prompt)
 
-            if context.config.get().log_tokenization:
+            if context.services.configuration.log_tokenization:
                 # TODO: better logging for and syntax
                 log_tokenization_for_conjunction(conjunction, tokenizer)
 
@@ -242,7 +286,7 @@ class SDXLPromptInvocationBase:
     title="SDXL Prompt",
     tags=["sdxl", "compel", "prompt"],
     category="conditioning",
-    version="1.0.1",
+    version="1.0.0",
 )
 class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
@@ -313,7 +357,6 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
                 dim=1,
             )
 
-        assert c2_pooled is not None
         conditioning_data = ConditioningFieldData(
             conditionings=[
                 SDXLConditioningInfo(
@@ -325,9 +368,14 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
             ]
         )
 
-        conditioning_name = context.conditioning.save(conditioning_data)
+        conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning"
+        context.services.latents.save(conditioning_name, conditioning_data)
 
-        return ConditioningOutput.build(conditioning_name)
+        return ConditioningOutput(
+            conditioning=ConditioningField(
+                conditioning_name=conditioning_name,
+            ),
+        )
 
 
 @invocation(
@@ -335,7 +383,7 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     title="SDXL Refiner Prompt",
     tags=["sdxl", "compel", "prompt"],
     category="conditioning",
-    version="1.0.1",
+    version="1.0.0",
 )
 class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
@@ -362,7 +410,6 @@ class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase
 
         add_time_ids = torch.tensor([original_size + crop_coords + (self.aesthetic_score,)])
 
-        assert c2_pooled is not None
         conditioning_data = ConditioningFieldData(
             conditionings=[
                 SDXLConditioningInfo(
@@ -374,9 +421,14 @@ class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase
             ]
         )
 
-        conditioning_name = context.conditioning.save(conditioning_data)
+        conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning"
+        context.services.latents.save(conditioning_name, conditioning_data)
 
-        return ConditioningOutput.build(conditioning_name)
+        return ConditioningOutput(
+            conditioning=ConditioningField(
+                conditioning_name=conditioning_name,
+            ),
+        )
 
 
 @invocation_output("clip_skip_output")
@@ -397,7 +449,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
@@ -407,9 +459,9 @@ class ClipSkipInvocation(BaseInvocation):
 
 
 def get_max_token_count(
-    tokenizer: CLIPTokenizer,
+    tokenizer,
     prompt: Union[FlattenedPrompt, Blend, Conjunction],
-    truncate_if_too_long: bool = False,
+    truncate_if_too_long=False,
 ) -> int:
     if type(prompt) is Blend:
         blend: Blend = prompt
@@ -421,9 +473,7 @@ def get_max_token_count(
         return len(get_tokens_for_prompt_object(tokenizer, prompt, truncate_if_too_long))
 
 
-def get_tokens_for_prompt_object(
-    tokenizer: CLIPTokenizer, parsed_prompt: FlattenedPrompt, truncate_if_too_long: bool = True
-) -> List[str]:
+def get_tokens_for_prompt_object(tokenizer, parsed_prompt: FlattenedPrompt, truncate_if_too_long=True) -> List[str]:
     if type(parsed_prompt) is Blend:
         raise ValueError("Blend is not supported here - you need to get tokens for each of its .children")
 
@@ -436,29 +486,24 @@ def get_tokens_for_prompt_object(
         for x in parsed_prompt.children
     ]
     text = " ".join(text_fragments)
-    tokens: List[str] = tokenizer.tokenize(text)
+    tokens = tokenizer.tokenize(text)
     if truncate_if_too_long:
         max_tokens_length = tokenizer.model_max_length - 2  # typically 75
         tokens = tokens[0:max_tokens_length]
     return tokens
 
 
-def log_tokenization_for_conjunction(
-    c: Conjunction, tokenizer: CLIPTokenizer, display_label_prefix: Optional[str] = None
-) -> None:
+def log_tokenization_for_conjunction(c: Conjunction, tokenizer, display_label_prefix=None):
     display_label_prefix = display_label_prefix or ""
     for i, p in enumerate(c.prompts):
         if len(c.prompts) > 1:
             this_display_label_prefix = f"{display_label_prefix}(conjunction part {i + 1}, weight={c.weights[i]})"
         else:
-            assert display_label_prefix is not None
             this_display_label_prefix = display_label_prefix
         log_tokenization_for_prompt_object(p, tokenizer, display_label_prefix=this_display_label_prefix)
 
 
-def log_tokenization_for_prompt_object(
-    p: Union[Blend, FlattenedPrompt], tokenizer: CLIPTokenizer, display_label_prefix: Optional[str] = None
-) -> None:
+def log_tokenization_for_prompt_object(p: Union[Blend, FlattenedPrompt], tokenizer, display_label_prefix=None):
     display_label_prefix = display_label_prefix or ""
     if type(p) is Blend:
         blend: Blend = p
@@ -498,12 +543,7 @@ def log_tokenization_for_prompt_object(
             log_tokenization_for_text(text, tokenizer, display_label=display_label_prefix)
 
 
-def log_tokenization_for_text(
-    text: str,
-    tokenizer: CLIPTokenizer,
-    display_label: Optional[str] = None,
-    truncate_if_too_long: Optional[bool] = False,
-) -> None:
+def log_tokenization_for_text(text, tokenizer, display_label=None, truncate_if_too_long=False):
     """shows how the prompt is tokenized
     # usually tokens have '</w>' to indicate end-of-word,
     # but for readability it has been replaced with ' '

invokeai/app/invocations/constants.py

@@ -1,17 +0,0 @@
-from typing import Literal
-
-from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
-
-LATENT_SCALE_FACTOR = 8
-"""
-HACK: Many nodes are currently hard-coded to use a fixed latent scale factor of 8. This is fragile, and will need to
-be addressed if future models use a different latent scale factor. Also, note that there may be places where the scale
-factor is hard-coded to a literal '8' rather than using this constant.
-The ratio of image:latent dimensions is LATENT_SCALE_FACTOR:1, or 8:1.
-"""
-
-SCHEDULER_NAME_VALUES = Literal[tuple(SCHEDULER_MAP.keys())]
-"""A literal type representing the valid scheduler names."""
-
-IMAGE_MODES = Literal["L", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "I", "F"]
-"""A literal type for PIL image modes supported by Invoke"""

invokeai/app/invocations/controlnet_image_processors.py

@@ -17,30 +17,33 @@ from controlnet_aux import (
     MidasDetector,
     MLSDdetector,
     NormalBaeDetector,
+    OpenposeDetector,
     PidiNetDetector,
     SamDetector,
     ZoeDetector,
 )
 from controlnet_aux.util import HWC3, ade_palette
 from PIL import Image
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    ImageField,
+from invokeai.app.invocations.primitives import ImageField, ImageOutput
+from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
+from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.backend.image_util.depth_anything import DepthAnythingDetector
+
+from ...backend.model_management import BaseModelType
+from .baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
     Input,
     InputField,
+    InvocationContext,
     OutputField,
-    WithBoard,
     WithMetadata,
+    invocation,
+    invocation_output,
 )
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.image_util.depth_anything import DepthAnythingDetector
-from invokeai.backend.image_util.dw_openpose import DWOpenposeDetector
-
-from .baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
 
 CONTROLNET_MODE_VALUES = Literal["balanced", "more_prompt", "more_control", "unbalanced"]
 CONTROLNET_RESIZE_VALUES = Literal[
@@ -54,7 +57,10 @@ CONTROLNET_RESIZE_VALUES = Literal[
 class ControlNetModelField(BaseModel):
     """ControlNet model field"""
 
-    key: str = Field(description="Model config record key for the ControlNet model")
+    model_name: str = Field(description="Name of the ControlNet model")
+    base_model: BaseModelType = Field(description="Base model")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class ControlField(BaseModel):
@@ -134,7 +140,7 @@ class ControlNetInvocation(BaseInvocation):
 
 
 # This invocation exists for other invocations to subclass it - do not register with @invocation!
-class ImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
+class ImageProcessorInvocation(BaseInvocation, WithMetadata):
     """Base class for invocations that preprocess images for ControlNet"""
 
     image: ImageField = InputField(description="The image to process")
@@ -143,18 +149,23 @@ class ImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
         # superclass just passes through image without processing
         return image
 
-    def load_image(self, context: InvocationContext) -> Image.Image:
-        # allows override for any special formatting specific to the preprocessor
-        return context.images.get_pil(self.image.image_name, "RGB")
-
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        raw_image = self.load_image(context)
+        raw_image = context.services.images.get_pil_image(self.image.image_name)
         # image type should be PIL.PngImagePlugin.PngImageFile ?
         processed_image = self.run_processor(raw_image)
 
         # currently can't see processed image in node UI without a showImage node,
         #    so for now setting image_type to RESULT instead of INTERMEDIATE so will get saved in gallery
-        image_dto = context.images.save(image=processed_image)
+        image_dto = context.services.images.create(
+            image=processed_image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.CONTROL,
+            session_id=context.graph_execution_state_id,
+            node_id=self.id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
         """Builds an ImageOutput and its ImageField"""
         processed_image_field = ImageField(image_name=image_dto.image_name)
@@ -173,7 +184,7 @@ class ImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
     title="Canny Processor",
     tags=["controlnet", "canny"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class CannyImageProcessorInvocation(ImageProcessorInvocation):
     """Canny edge detection for ControlNet"""
@@ -185,10 +196,6 @@ class CannyImageProcessorInvocation(ImageProcessorInvocation):
         default=200, ge=0, le=255, description="The high threshold of the Canny pixel gradient (0-255)"
     )
 
-    def load_image(self, context: InvocationContext) -> Image.Image:
-        # Keep alpha channel for Canny processing to detect edges of transparent areas
-        return context.images.get_pil(self.image.image_name, "RGBA")
-
     def run_processor(self, image):
         canny_processor = CannyDetector()
         processed_image = canny_processor(image, self.low_threshold, self.high_threshold)
@@ -200,7 +207,7 @@ class CannyImageProcessorInvocation(ImageProcessorInvocation):
     title="HED (softedge) Processor",
     tags=["controlnet", "hed", "softedge"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class HedImageProcessorInvocation(ImageProcessorInvocation):
     """Applies HED edge detection to image"""
@@ -229,7 +236,7 @@ class HedImageProcessorInvocation(ImageProcessorInvocation):
     title="Lineart Processor",
     tags=["controlnet", "lineart"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class LineartImageProcessorInvocation(ImageProcessorInvocation):
     """Applies line art processing to image"""
@@ -251,7 +258,7 @@ class LineartImageProcessorInvocation(ImageProcessorInvocation):
     title="Lineart Anime Processor",
     tags=["controlnet", "lineart", "anime"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation):
     """Applies line art anime processing to image"""
@@ -269,12 +276,37 @@ class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation):
         return processed_image
 
 
+@invocation(
+    "openpose_image_processor",
+    title="Openpose Processor",
+    tags=["controlnet", "openpose", "pose"],
+    category="controlnet",
+    version="1.2.0",
+)
+class OpenposeImageProcessorInvocation(ImageProcessorInvocation):
+    """Applies Openpose processing to image"""
+
+    hand_and_face: bool = InputField(default=False, description="Whether to use hands and face mode")
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
+
+    def run_processor(self, image):
+        openpose_processor = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+        processed_image = openpose_processor(
+            image,
+            detect_resolution=self.detect_resolution,
+            image_resolution=self.image_resolution,
+            hand_and_face=self.hand_and_face,
+        )
+        return processed_image
+
+
 @invocation(
     "midas_depth_image_processor",
     title="Midas Depth Processor",
     tags=["controlnet", "midas"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class MidasDepthImageProcessorInvocation(ImageProcessorInvocation):
     """Applies Midas depth processing to image"""
@@ -301,7 +333,7 @@ class MidasDepthImageProcessorInvocation(ImageProcessorInvocation):
     title="Normal BAE Processor",
     tags=["controlnet"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class NormalbaeImageProcessorInvocation(ImageProcessorInvocation):
     """Applies NormalBae processing to image"""
@@ -318,7 +350,7 @@ class NormalbaeImageProcessorInvocation(ImageProcessorInvocation):
 
 
 @invocation(
-    "mlsd_image_processor", title="MLSD Processor", tags=["controlnet", "mlsd"], category="controlnet", version="1.2.1"
+    "mlsd_image_processor", title="MLSD Processor", tags=["controlnet", "mlsd"], category="controlnet", version="1.2.0"
 )
 class MlsdImageProcessorInvocation(ImageProcessorInvocation):
     """Applies MLSD processing to image"""
@@ -341,7 +373,7 @@ class MlsdImageProcessorInvocation(ImageProcessorInvocation):
 
 
 @invocation(
-    "pidi_image_processor", title="PIDI Processor", tags=["controlnet", "pidi"], category="controlnet", version="1.2.1"
+    "pidi_image_processor", title="PIDI Processor", tags=["controlnet", "pidi"], category="controlnet", version="1.2.0"
 )
 class PidiImageProcessorInvocation(ImageProcessorInvocation):
     """Applies PIDI processing to image"""
@@ -368,7 +400,7 @@ class PidiImageProcessorInvocation(ImageProcessorInvocation):
     title="Content Shuffle Processor",
     tags=["controlnet", "contentshuffle"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation):
     """Applies content shuffle processing to image"""
@@ -398,7 +430,7 @@ class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation):
     title="Zoe (Depth) Processor",
     tags=["controlnet", "zoe", "depth"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation):
     """Applies Zoe depth processing to image"""
@@ -414,7 +446,7 @@ class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation):
     title="Mediapipe Face Processor",
     tags=["controlnet", "mediapipe", "face"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     """Applies mediapipe face processing to image"""
@@ -423,6 +455,10 @@ class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     min_confidence: float = InputField(default=0.5, ge=0, le=1, description="Minimum confidence for face detection")
 
     def run_processor(self, image):
+        # MediaPipeFaceDetector throws an error if image has alpha channel
+        #     so convert to RGB if needed
+        if image.mode == "RGBA":
+            image = image.convert("RGB")
         mediapipe_face_processor = MediapipeFaceDetector()
         processed_image = mediapipe_face_processor(image, max_faces=self.max_faces, min_confidence=self.min_confidence)
         return processed_image
@@ -433,7 +469,7 @@ class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     title="Leres (Depth) Processor",
     tags=["controlnet", "leres", "depth"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class LeresImageProcessorInvocation(ImageProcessorInvocation):
     """Applies leres processing to image"""
@@ -462,7 +498,7 @@ class LeresImageProcessorInvocation(ImageProcessorInvocation):
     title="Tile Resample Processor",
     tags=["controlnet", "tile"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class TileResamplerProcessorInvocation(ImageProcessorInvocation):
     """Tile resampler processor"""
@@ -502,7 +538,7 @@ class TileResamplerProcessorInvocation(ImageProcessorInvocation):
     title="Segment Anything Processor",
     tags=["controlnet", "segmentanything"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class SegmentAnythingProcessorInvocation(ImageProcessorInvocation):
     """Applies segment anything processing to image"""
@@ -544,7 +580,7 @@ class SamDetectorReproducibleColors(SamDetector):
     title="Color Map Processor",
     tags=["controlnet"],
     category="controlnet",
-    version="1.2.1",
+    version="1.2.0",
 )
 class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
     """Generates a color map from the provided image"""
@@ -552,6 +588,7 @@ class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
     color_map_tile_size: int = InputField(default=64, ge=0, description=FieldDescriptions.tile_size)
 
     def run_processor(self, image: Image.Image):
+        image = image.convert("RGB")
         np_image = np.array(image, dtype=np.uint8)
         height, width = np_image.shape[:2]
 
@@ -587,36 +624,12 @@ class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
     resolution: int = InputField(default=512, ge=64, multiple_of=64, description=FieldDescriptions.image_res)
     offload: bool = InputField(default=False)
 
-    def run_processor(self, image: Image.Image):
+    def run_processor(self, image):
         depth_anything_detector = DepthAnythingDetector()
         depth_anything_detector.load_model(model_size=self.model_size)
 
+        if image.mode == "RGBA":
+            image = image.convert("RGB")
+
         processed_image = depth_anything_detector(image=image, resolution=self.resolution, offload=self.offload)
         return processed_image
-
-
-@invocation(
-    "dw_openpose_image_processor",
-    title="DW Openpose Image Processor",
-    tags=["controlnet", "dwpose", "openpose"],
-    category="controlnet",
-    version="1.0.0",
-)
-class DWOpenposeImageProcessorInvocation(ImageProcessorInvocation):
-    """Generates an openpose pose from an image using DWPose"""
-
-    draw_body: bool = InputField(default=True)
-    draw_face: bool = InputField(default=False)
-    draw_hands: bool = InputField(default=False)
-    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
-
-    def run_processor(self, image: Image.Image):
-        dw_openpose = DWOpenposeDetector()
-        processed_image = dw_openpose(
-            image,
-            draw_face=self.draw_face,
-            draw_hands=self.draw_hands,
-            draw_body=self.draw_body,
-            resolution=self.image_resolution,
-        )
-        return processed_image

invokeai/app/invocations/cv.py

@@ -5,24 +5,22 @@ import cv2 as cv
 import numpy
 from PIL import Image, ImageOps
 
-from invokeai.app.invocations.fields import ImageField
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.invocations.primitives import ImageField, ImageOutput
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
 
-from .baseinvocation import BaseInvocation, invocation
-from .fields import InputField, WithBoard, WithMetadata
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, WithMetadata, invocation
 
 
-@invocation("cv_inpaint", title="OpenCV Inpaint", tags=["opencv", "inpaint"], category="inpaint", version="1.2.1")
-class CvInpaintInvocation(BaseInvocation, WithMetadata, WithBoard):
+@invocation("cv_inpaint", title="OpenCV Inpaint", tags=["opencv", "inpaint"], category="inpaint", version="1.2.0")
+class CvInpaintInvocation(BaseInvocation, WithMetadata):
     """Simple inpaint using opencv."""
 
     image: ImageField = InputField(description="The image to inpaint")
     mask: ImageField = InputField(description="The mask to use when inpainting")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
-        mask = context.images.get_pil(self.mask.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
+        mask = context.services.images.get_pil_image(self.mask.image_name)
 
         # Convert to cv image/mask
         # TODO: consider making these utility functions
@@ -36,6 +34,18 @@ class CvInpaintInvocation(BaseInvocation, WithMetadata, WithBoard):
         # TODO: consider making a utility function
         image_inpainted = Image.fromarray(cv.cvtColor(cv_inpainted, cv.COLOR_BGR2RGB))
 
-        image_dto = context.images.save(image=image_inpainted)
+        image_dto = context.services.images.create(
+            image=image_inpainted,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )

invokeai/app/invocations/facetools.py

@@ -13,13 +13,15 @@ from pydantic import field_validator
 import invokeai.assets.fonts as font_assets
 from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
+    InputField,
+    InvocationContext,
+    OutputField,
+    WithMetadata,
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import ImageField, InputField, OutputField, WithBoard, WithMetadata
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.image_records.image_records_common import ImageCategory
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.invocations.primitives import ImageField, ImageOutput
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
 
 
 @invocation_output("face_mask_output")
@@ -304,37 +306,37 @@ def extract_face(
 
     # Adjust the crop boundaries to stay within the original image's dimensions
     if x_min < 0:
-        context.logger.warning("FaceTools --> -X-axis padding reached image edge.")
+        context.services.logger.warning("FaceTools --> -X-axis padding reached image edge.")
         x_max -= x_min
         x_min = 0
     elif x_max > mask.width:
-        context.logger.warning("FaceTools --> +X-axis padding reached image edge.")
+        context.services.logger.warning("FaceTools --> +X-axis padding reached image edge.")
         x_min -= x_max - mask.width
         x_max = mask.width
 
     if y_min < 0:
-        context.logger.warning("FaceTools --> +Y-axis padding reached image edge.")
+        context.services.logger.warning("FaceTools --> +Y-axis padding reached image edge.")
         y_max -= y_min
         y_min = 0
     elif y_max > mask.height:
-        context.logger.warning("FaceTools --> -Y-axis padding reached image edge.")
+        context.services.logger.warning("FaceTools --> -Y-axis padding reached image edge.")
         y_min -= y_max - mask.height
         y_max = mask.height
 
     # Ensure the crop is square and adjust the boundaries if needed
     if x_max - x_min != crop_size:
-        context.logger.warning("FaceTools --> Limiting x-axis padding to constrain bounding box to a square.")
+        context.services.logger.warning("FaceTools --> Limiting x-axis padding to constrain bounding box to a square.")
         diff = crop_size - (x_max - x_min)
         x_min -= diff // 2
         x_max += diff - diff // 2
 
     if y_max - y_min != crop_size:
-        context.logger.warning("FaceTools --> Limiting y-axis padding to constrain bounding box to a square.")
+        context.services.logger.warning("FaceTools --> Limiting y-axis padding to constrain bounding box to a square.")
         diff = crop_size - (y_max - y_min)
         y_min -= diff // 2
         y_max += diff - diff // 2
 
-    context.logger.info(f"FaceTools --> Calculated bounding box (8 multiple): {crop_size}")
+    context.services.logger.info(f"FaceTools --> Calculated bounding box (8 multiple): {crop_size}")
 
     # Crop the output image to the specified size with the center of the face mesh as the center.
     mask = mask.crop((x_min, y_min, x_max, y_max))
@@ -366,7 +368,7 @@ def get_faces_list(
 
     # Generate the face box mask and get the center of the face.
     if not should_chunk:
-        context.logger.info("FaceTools --> Attempting full image face detection.")
+        context.services.logger.info("FaceTools --> Attempting full image face detection.")
         result = generate_face_box_mask(
             context=context,
             minimum_confidence=minimum_confidence,
@@ -378,7 +380,7 @@ def get_faces_list(
             draw_mesh=draw_mesh,
         )
     if should_chunk or len(result) == 0:
-        context.logger.info("FaceTools --> Chunking image (chunk toggled on, or no face found in full image).")
+        context.services.logger.info("FaceTools --> Chunking image (chunk toggled on, or no face found in full image).")
         width, height = image.size
         image_chunks = []
         x_offsets = []
@@ -397,7 +399,7 @@ def get_faces_list(
                 x_offsets.append(x)
                 y_offsets.append(0)
                 fx += increment
-                context.logger.info(f"FaceTools --> Chunk starting at x = {x}")
+                context.services.logger.info(f"FaceTools --> Chunk starting at x = {x}")
         elif height > width:
             # Portrait - slice the image vertically
             fy = 0.0
@@ -409,10 +411,10 @@ def get_faces_list(
                 x_offsets.append(0)
                 y_offsets.append(y)
                 fy += increment
-                context.logger.info(f"FaceTools --> Chunk starting at y = {y}")
+                context.services.logger.info(f"FaceTools --> Chunk starting at y = {y}")
 
         for idx in range(len(image_chunks)):
-            context.logger.info(f"FaceTools --> Evaluating faces in chunk {idx}")
+            context.services.logger.info(f"FaceTools --> Evaluating faces in chunk {idx}")
             result = result + generate_face_box_mask(
                 context=context,
                 minimum_confidence=minimum_confidence,
@@ -426,7 +428,7 @@ def get_faces_list(
 
         if len(result) == 0:
             # Give up
-            context.logger.warning(
+            context.services.logger.warning(
                 "FaceTools --> No face detected in chunked input image. Passing through original image."
             )
 
@@ -435,7 +437,7 @@ def get_faces_list(
     return all_faces
 
 
-@invocation("face_off", title="FaceOff", tags=["image", "faceoff", "face", "mask"], category="image", version="1.2.1")
+@invocation("face_off", title="FaceOff", tags=["image", "faceoff", "face", "mask"], category="image", version="1.2.0")
 class FaceOffInvocation(BaseInvocation, WithMetadata):
     """Bound, extract, and mask a face from an image using MediaPipe detection"""
 
@@ -468,11 +470,11 @@ class FaceOffInvocation(BaseInvocation, WithMetadata):
         )
 
         if len(all_faces) == 0:
-            context.logger.warning("FaceOff --> No faces detected. Passing through original image.")
+            context.services.logger.warning("FaceOff --> No faces detected. Passing through original image.")
             return None
 
         if self.face_id > len(all_faces) - 1:
-            context.logger.warning(
+            context.services.logger.warning(
                 f"FaceOff --> Face ID {self.face_id} is outside of the number of faces detected ({len(all_faces)}). Passing through original image."
             )
             return None
@@ -484,7 +486,7 @@ class FaceOffInvocation(BaseInvocation, WithMetadata):
         return face_data
 
     def invoke(self, context: InvocationContext) -> FaceOffOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
         result = self.faceoff(context=context, image=image)
 
         if result is None:
@@ -498,9 +500,24 @@ class FaceOffInvocation(BaseInvocation, WithMetadata):
             x = result["x_min"]
             y = result["y_min"]
 
-        image_dto = context.images.save(image=result_image)
+        image_dto = context.services.images.create(
+            image=result_image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            workflow=context.workflow,
+        )
 
-        mask_dto = context.images.save(image=result_mask, image_category=ImageCategory.MASK)
+        mask_dto = context.services.images.create(
+            image=result_mask,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.MASK,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+        )
 
         output = FaceOffOutput(
             image=ImageField(image_name=image_dto.image_name),
@@ -514,7 +531,7 @@ class FaceOffInvocation(BaseInvocation, WithMetadata):
         return output
 
 
-@invocation("face_mask_detection", title="FaceMask", tags=["image", "face", "mask"], category="image", version="1.2.1")
+@invocation("face_mask_detection", title="FaceMask", tags=["image", "face", "mask"], category="image", version="1.2.0")
 class FaceMaskInvocation(BaseInvocation, WithMetadata):
     """Face mask creation using mediapipe face detection"""
 
@@ -563,7 +580,7 @@ class FaceMaskInvocation(BaseInvocation, WithMetadata):
 
             if len(intersected_face_ids) == 0:
                 id_range_str = ",".join([str(id) for id in id_range])
-                context.logger.warning(
+                context.services.logger.warning(
                     f"Face IDs must be in range of detected faces - requested {self.face_ids}, detected {id_range_str}. Passing through original image."
                 )
                 return FaceMaskResult(
@@ -599,12 +616,27 @@ class FaceMaskInvocation(BaseInvocation, WithMetadata):
         )
 
     def invoke(self, context: InvocationContext) -> FaceMaskOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
         result = self.facemask(context=context, image=image)
 
-        image_dto = context.images.save(image=result["image"])
+        image_dto = context.services.images.create(
+            image=result["image"],
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            workflow=context.workflow,
+        )
 
-        mask_dto = context.images.save(image=result["mask"], image_category=ImageCategory.MASK)
+        mask_dto = context.services.images.create(
+            image=result["mask"],
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.MASK,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+        )
 
         output = FaceMaskOutput(
             image=ImageField(image_name=image_dto.image_name),
@@ -617,9 +649,9 @@ class FaceMaskInvocation(BaseInvocation, WithMetadata):
 
 
 @invocation(
-    "face_identifier", title="FaceIdentifier", tags=["image", "face", "identifier"], category="image", version="1.2.1"
+    "face_identifier", title="FaceIdentifier", tags=["image", "face", "identifier"], category="image", version="1.2.0"
 )
-class FaceIdentifierInvocation(BaseInvocation, WithMetadata, WithBoard):
+class FaceIdentifierInvocation(BaseInvocation, WithMetadata):
     """Outputs an image with detected face IDs printed on each face. For use with other FaceTools."""
 
     image: ImageField = InputField(description="Image to face detect")
@@ -673,9 +705,21 @@ class FaceIdentifierInvocation(BaseInvocation, WithMetadata, WithBoard):
         return image
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
         result_image = self.faceidentifier(context=context, image=image)
 
-        image_dto = context.images.save(image=result_image)
+        image_dto = context.services.images.create(
+            image=result_image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )

invokeai/app/invocations/fields.py

@@ -1,566 +0,0 @@
-from enum import Enum
-from typing import Any, Callable, Optional, Tuple
-
-from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter
-from pydantic.fields import _Unset
-from pydantic_core import PydanticUndefined
-
-from invokeai.app.util.metaenum import MetaEnum
-from invokeai.backend.util.logging import InvokeAILogger
-
-logger = InvokeAILogger.get_logger()
-
-
-class UIType(str, Enum, metaclass=MetaEnum):
-    """
-    Type hints for the UI for situations in which the field type is not enough to infer the correct UI type.
-
-    - Model Fields
-    The most common node-author-facing use will be for model fields. Internally, there is no difference
-    between SD-1, SD-2 and SDXL model fields - they all use the class `MainModelField`. To ensure the
-    base-model-specific UI is rendered, use e.g. `ui_type=UIType.SDXLMainModelField` to indicate that
-    the field is an SDXL main model field.
-
-    - Any Field
-    We cannot infer the usage of `typing.Any` via schema parsing, so you *must* use `ui_type=UIType.Any` to
-    indicate that the field accepts any type. Use with caution. This cannot be used on outputs.
-
-    - Scheduler Field
-    Special handling in the UI is needed for this field, which otherwise would be parsed as a plain enum field.
-
-    - Internal Fields
-    Similar to the Any Field, the `collect` and `iterate` nodes make use of `typing.Any`. To facilitate
-    handling these types in the client, we use `UIType._Collection` and `UIType._CollectionItem`. These
-    should not be used by node authors.
-
-    - DEPRECATED Fields
-    These types are deprecated and should not be used by node authors. A warning will be logged if one is
-    used, and the type will be ignored. They are included here for backwards compatibility.
-    """
-
-    # region Model Field Types
-    SDXLMainModel = "SDXLMainModelField"
-    SDXLRefinerModel = "SDXLRefinerModelField"
-    ONNXModel = "ONNXModelField"
-    VaeModel = "VAEModelField"
-    LoRAModel = "LoRAModelField"
-    ControlNetModel = "ControlNetModelField"
-    IPAdapterModel = "IPAdapterModelField"
-    # endregion
-
-    # region Misc Field Types
-    Scheduler = "SchedulerField"
-    Any = "AnyField"
-    # endregion
-
-    # region Internal Field Types
-    _Collection = "CollectionField"
-    _CollectionItem = "CollectionItemField"
-    # endregion
-
-    # region DEPRECATED
-    Boolean = "DEPRECATED_Boolean"
-    Color = "DEPRECATED_Color"
-    Conditioning = "DEPRECATED_Conditioning"
-    Control = "DEPRECATED_Control"
-    Float = "DEPRECATED_Float"
-    Image = "DEPRECATED_Image"
-    Integer = "DEPRECATED_Integer"
-    Latents = "DEPRECATED_Latents"
-    String = "DEPRECATED_String"
-    BooleanCollection = "DEPRECATED_BooleanCollection"
-    ColorCollection = "DEPRECATED_ColorCollection"
-    ConditioningCollection = "DEPRECATED_ConditioningCollection"
-    ControlCollection = "DEPRECATED_ControlCollection"
-    FloatCollection = "DEPRECATED_FloatCollection"
-    ImageCollection = "DEPRECATED_ImageCollection"
-    IntegerCollection = "DEPRECATED_IntegerCollection"
-    LatentsCollection = "DEPRECATED_LatentsCollection"
-    StringCollection = "DEPRECATED_StringCollection"
-    BooleanPolymorphic = "DEPRECATED_BooleanPolymorphic"
-    ColorPolymorphic = "DEPRECATED_ColorPolymorphic"
-    ConditioningPolymorphic = "DEPRECATED_ConditioningPolymorphic"
-    ControlPolymorphic = "DEPRECATED_ControlPolymorphic"
-    FloatPolymorphic = "DEPRECATED_FloatPolymorphic"
-    ImagePolymorphic = "DEPRECATED_ImagePolymorphic"
-    IntegerPolymorphic = "DEPRECATED_IntegerPolymorphic"
-    LatentsPolymorphic = "DEPRECATED_LatentsPolymorphic"
-    StringPolymorphic = "DEPRECATED_StringPolymorphic"
-    MainModel = "DEPRECATED_MainModel"
-    UNet = "DEPRECATED_UNet"
-    Vae = "DEPRECATED_Vae"
-    CLIP = "DEPRECATED_CLIP"
-    Collection = "DEPRECATED_Collection"
-    CollectionItem = "DEPRECATED_CollectionItem"
-    Enum = "DEPRECATED_Enum"
-    WorkflowField = "DEPRECATED_WorkflowField"
-    IsIntermediate = "DEPRECATED_IsIntermediate"
-    BoardField = "DEPRECATED_BoardField"
-    MetadataItem = "DEPRECATED_MetadataItem"
-    MetadataItemCollection = "DEPRECATED_MetadataItemCollection"
-    MetadataItemPolymorphic = "DEPRECATED_MetadataItemPolymorphic"
-    MetadataDict = "DEPRECATED_MetadataDict"
-
-
-class UIComponent(str, Enum, metaclass=MetaEnum):
-    """
-    The type of UI component to use for a field, used to override the default components, which are
-    inferred from the field type.
-    """
-
-    None_ = "none"
-    Textarea = "textarea"
-    Slider = "slider"
-
-
-class FieldDescriptions:
-    denoising_start = "When to start denoising, expressed a percentage of total steps"
-    denoising_end = "When to stop denoising, expressed a percentage of total steps"
-    cfg_scale = "Classifier-Free Guidance scale"
-    cfg_rescale_multiplier = "Rescale multiplier for CFG guidance, used for models trained with zero-terminal SNR"
-    scheduler = "Scheduler to use during inference"
-    positive_cond = "Positive conditioning tensor"
-    negative_cond = "Negative conditioning tensor"
-    noise = "Noise tensor"
-    clip = "CLIP (tokenizer, text encoder, LoRAs) and skipped layer count"
-    unet = "UNet (scheduler, LoRAs)"
-    vae = "VAE"
-    cond = "Conditioning tensor"
-    controlnet_model = "ControlNet model to load"
-    vae_model = "VAE model to load"
-    lora_model = "LoRA model to load"
-    main_model = "Main model (UNet, VAE, CLIP) to load"
-    sdxl_main_model = "SDXL Main model (UNet, VAE, CLIP1, CLIP2) to load"
-    sdxl_refiner_model = "SDXL Refiner Main Modde (UNet, VAE, CLIP2) to load"
-    onnx_main_model = "ONNX Main model (UNet, VAE, CLIP) to load"
-    lora_weight = "The weight at which the LoRA is applied to each model"
-    compel_prompt = "Prompt to be parsed by Compel to create a conditioning tensor"
-    raw_prompt = "Raw prompt text (no parsing)"
-    sdxl_aesthetic = "The aesthetic score to apply to the conditioning tensor"
-    skipped_layers = "Number of layers to skip in text encoder"
-    seed = "Seed for random number generation"
-    steps = "Number of steps to run"
-    width = "Width of output (px)"
-    height = "Height of output (px)"
-    control = "ControlNet(s) to apply"
-    ip_adapter = "IP-Adapter to apply"
-    t2i_adapter = "T2I-Adapter(s) to apply"
-    denoised_latents = "Denoised latents tensor"
-    latents = "Latents tensor"
-    strength = "Strength of denoising (proportional to steps)"
-    metadata = "Optional metadata to be saved with the image"
-    metadata_collection = "Collection of Metadata"
-    metadata_item_polymorphic = "A single metadata item or collection of metadata items"
-    metadata_item_label = "Label for this metadata item"
-    metadata_item_value = "The value for this metadata item (may be any type)"
-    workflow = "Optional workflow to be saved with the image"
-    interp_mode = "Interpolation mode"
-    torch_antialias = "Whether or not to apply antialiasing (bilinear or bicubic only)"
-    fp32 = "Whether or not to use full float32 precision"
-    precision = "Precision to use"
-    tiled = "Processing using overlapping tiles (reduce memory consumption)"
-    detect_res = "Pixel resolution for detection"
-    image_res = "Pixel resolution for output image"
-    safe_mode = "Whether or not to use safe mode"
-    scribble_mode = "Whether or not to use scribble mode"
-    scale_factor = "The factor by which to scale"
-    blend_alpha = (
-        "Blending factor. 0.0 = use input A only, 1.0 = use input B only, 0.5 = 50% mix of input A and input B."
-    )
-    num_1 = "The first number"
-    num_2 = "The second number"
-    mask = "The mask to use for the operation"
-    board = "The board to save the image to"
-    image = "The image to process"
-    tile_size = "Tile size"
-    inclusive_low = "The inclusive low value"
-    exclusive_high = "The exclusive high value"
-    decimal_places = "The number of decimal places to round to"
-    freeu_s1 = 'Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to mitigate the "oversmoothing effect" in the enhanced denoising process.'
-    freeu_s2 = 'Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to mitigate the "oversmoothing effect" in the enhanced denoising process.'
-    freeu_b1 = "Scaling factor for stage 1 to amplify the contributions of backbone features."
-    freeu_b2 = "Scaling factor for stage 2 to amplify the contributions of backbone features."
-
-
-class ImageField(BaseModel):
-    """An image primitive field"""
-
-    image_name: str = Field(description="The name of the image")
-
-
-class BoardField(BaseModel):
-    """A board primitive field"""
-
-    board_id: str = Field(description="The id of the board")
-
-
-class DenoiseMaskField(BaseModel):
-    """An inpaint mask field"""
-
-    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):
-    """A latents tensor primitive field"""
-
-    latents_name: str = Field(description="The name of the latents")
-    seed: Optional[int] = Field(default=None, description="Seed used to generate this latents")
-
-
-class ColorField(BaseModel):
-    """A color primitive field"""
-
-    r: int = Field(ge=0, le=255, description="The red component")
-    g: int = Field(ge=0, le=255, description="The green component")
-    b: int = Field(ge=0, le=255, description="The blue component")
-    a: int = Field(ge=0, le=255, description="The alpha component")
-
-    def tuple(self) -> Tuple[int, int, int, int]:
-        return (self.r, self.g, self.b, self.a)
-
-
-class ConditioningField(BaseModel):
-    """A conditioning tensor primitive value"""
-
-    conditioning_name: str = Field(description="The name of conditioning tensor")
-    # endregion
-
-
-class MetadataField(RootModel):
-    """
-    Pydantic model for metadata with custom root of type dict[str, Any].
-    Metadata is stored without a strict schema.
-    """
-
-    root: dict[str, Any] = Field(description="The metadata")
-
-
-MetadataFieldValidator = TypeAdapter(MetadataField)
-
-
-class Input(str, Enum, metaclass=MetaEnum):
-    """
-    The type of input a field accepts.
-    - `Input.Direct`: The field must have its value provided directly, when the invocation and field \
-      are instantiated.
-    - `Input.Connection`: The field must have its value provided by a connection.
-    - `Input.Any`: The field may have its value provided either directly or by a connection.
-    """
-
-    Connection = "connection"
-    Direct = "direct"
-    Any = "any"
-
-
-class FieldKind(str, Enum, metaclass=MetaEnum):
-    """
-    The kind of field.
-    - `Input`: An input field on a node.
-    - `Output`: An output field on a node.
-    - `Internal`: A field which is treated as an input, but cannot be used in node definitions. Metadata is
-    one example. It is provided to nodes via the WithMetadata class, and we want to reserve the field name
-    "metadata" for this on all nodes. `FieldKind` is used to short-circuit the field name validation logic,
-    allowing "metadata" for that field.
-    - `NodeAttribute`: The field is a node attribute. These are fields which are not inputs or outputs,
-    but which are used to store information about the node. For example, the `id` and `type` fields are node
-    attributes.
-
-    The presence of this in `json_schema_extra["field_kind"]` is used when initializing node schemas on app
-    startup, and when generating the OpenAPI schema for the workflow editor.
-    """
-
-    Input = "input"
-    Output = "output"
-    Internal = "internal"
-    NodeAttribute = "node_attribute"
-
-
-class InputFieldJSONSchemaExtra(BaseModel):
-    """
-    Extra attributes to be added to input fields and their OpenAPI schema. Used during graph execution,
-    and by the workflow editor during schema parsing and UI rendering.
-    """
-
-    input: Input
-    orig_required: bool
-    field_kind: FieldKind
-    default: Optional[Any] = None
-    orig_default: Optional[Any] = None
-    ui_hidden: bool = False
-    ui_type: Optional[UIType] = None
-    ui_component: Optional[UIComponent] = None
-    ui_order: Optional[int] = None
-    ui_choice_labels: Optional[dict[str, str]] = None
-
-    model_config = ConfigDict(
-        validate_assignment=True,
-        json_schema_serialization_defaults_required=True,
-    )
-
-
-class WithMetadata(BaseModel):
-    """
-    Inherit from this class if your node needs a metadata input field.
-    """
-
-    metadata: Optional[MetadataField] = Field(
-        default=None,
-        description=FieldDescriptions.metadata,
-        json_schema_extra=InputFieldJSONSchemaExtra(
-            field_kind=FieldKind.Internal,
-            input=Input.Connection,
-            orig_required=False,
-        ).model_dump(exclude_none=True),
-    )
-
-
-class WithWorkflow:
-    workflow = None
-
-    def __init_subclass__(cls) -> None:
-        logger.warn(
-            f"{cls.__module__.split('.')[0]}.{cls.__name__}: WithWorkflow is deprecated. Use `context.workflow` to access the workflow."
-        )
-        super().__init_subclass__()
-
-
-class WithBoard(BaseModel):
-    """
-    Inherit from this class if your node needs a board input field.
-    """
-
-    board: Optional[BoardField] = Field(
-        default=None,
-        description=FieldDescriptions.board,
-        json_schema_extra=InputFieldJSONSchemaExtra(
-            field_kind=FieldKind.Internal,
-            input=Input.Direct,
-            orig_required=False,
-        ).model_dump(exclude_none=True),
-    )
-
-
-class OutputFieldJSONSchemaExtra(BaseModel):
-    """
-    Extra attributes to be added to input fields and their OpenAPI schema. Used by the workflow editor
-    during schema parsing and UI rendering.
-    """
-
-    field_kind: FieldKind
-    ui_hidden: bool
-    ui_type: Optional[UIType]
-    ui_order: Optional[int]
-
-    model_config = ConfigDict(
-        validate_assignment=True,
-        json_schema_serialization_defaults_required=True,
-    )
-
-
-def InputField(
-    # copied from pydantic's Field
-    # TODO: Can we support default_factory?
-    default: Any = _Unset,
-    default_factory: Callable[[], Any] | None = _Unset,
-    title: str | None = _Unset,
-    description: str | None = _Unset,
-    pattern: str | None = _Unset,
-    strict: bool | None = _Unset,
-    gt: float | None = _Unset,
-    ge: float | None = _Unset,
-    lt: float | None = _Unset,
-    le: float | None = _Unset,
-    multiple_of: float | None = _Unset,
-    allow_inf_nan: bool | None = _Unset,
-    max_digits: int | None = _Unset,
-    decimal_places: int | None = _Unset,
-    min_length: int | None = _Unset,
-    max_length: int | None = _Unset,
-    # custom
-    input: Input = Input.Any,
-    ui_type: Optional[UIType] = None,
-    ui_component: Optional[UIComponent] = None,
-    ui_hidden: bool = False,
-    ui_order: Optional[int] = None,
-    ui_choice_labels: Optional[dict[str, str]] = None,
-) -> Any:
-    """
-    Creates an input field for an invocation.
-
-    This is a wrapper for Pydantic's [Field](https://docs.pydantic.dev/latest/api/fields/#pydantic.fields.Field) \
-    that adds a few extra parameters to support graph execution and the node editor UI.
-
-    :param Input input: [Input.Any] The kind of input this field requires. \
-      `Input.Direct` means a value must be provided on instantiation. \
-      `Input.Connection` means the value must be provided by a connection. \
-      `Input.Any` means either will do.
-
-    :param UIType ui_type: [None] Optionally provides an extra type hint for the UI. \
-      In some situations, the field's type is not enough to infer the correct UI type. \
-      For example, model selection fields should render a dropdown UI component to select a model. \
-      Internally, there is no difference between SD-1, SD-2 and SDXL model fields, they all use \
-      `MainModelField`. So to ensure the base-model-specific UI is rendered, you can use \
-      `UIType.SDXLMainModelField` to indicate that the field is an SDXL main model field.
-
-    :param UIComponent ui_component: [None] Optionally specifies a specific component to use in the UI. \
-      The UI will always render a suitable component, but sometimes you want something different than the default. \
-      For example, a `string` field will default to a single-line input, but you may want a multi-line textarea instead. \
-      For this case, you could provide `UIComponent.Textarea`.
-
-    :param bool ui_hidden: [False] Specifies whether or not this field should be hidden in the UI.
-
-    :param int ui_order: [None] Specifies the order in which this field should be rendered in the UI.
-
-    :param dict[str, str] ui_choice_labels: [None] Specifies the labels to use for the choices in an enum field.
-    """
-
-    json_schema_extra_ = InputFieldJSONSchemaExtra(
-        input=input,
-        ui_type=ui_type,
-        ui_component=ui_component,
-        ui_hidden=ui_hidden,
-        ui_order=ui_order,
-        ui_choice_labels=ui_choice_labels,
-        field_kind=FieldKind.Input,
-        orig_required=True,
-    )
-
-    """
-    There is a conflict between the typing of invocation definitions and the typing of an invocation's
-    `invoke()` function.
-
-    On instantiation of a node, the invocation definition is used to create the python class. At this time,
-    any number of fields may be optional, because they may be provided by connections.
-
-    On calling of `invoke()`, however, those fields may be required.
-
-    For example, consider an ResizeImageInvocation with an `image: ImageField` field.
-
-    `image` is required during the call to `invoke()`, but when the python class is instantiated,
-    the field may not be present. This is fine, because that image field will be provided by a
-    connection from an ancestor node, which outputs an image.
-
-    This means we want to type the `image` field as optional for the node class definition, but required
-    for the `invoke()` function.
-
-    If we use `typing.Optional` in the node class definition, the field will be typed as optional in the
-    `invoke()` method, and we'll have to do a lot of runtime checks to ensure the field is present - or
-    any static type analysis tools will complain.
-
-    To get around this, in node class definitions, we type all fields correctly for the `invoke()` function,
-    but secretly make them optional in `InputField()`. We also store the original required bool and/or default
-    value. When we call `invoke()`, we use this stored information to do an additional check on the class.
-    """
-
-    if default_factory is not _Unset and default_factory is not None:
-        default = default_factory()
-        logger.warn('"default_factory" is not supported, calling it now to set "default"')
-
-    # These are the args we may wish pass to the pydantic `Field()` function
-    field_args = {
-        "default": default,
-        "title": title,
-        "description": description,
-        "pattern": pattern,
-        "strict": strict,
-        "gt": gt,
-        "ge": ge,
-        "lt": lt,
-        "le": le,
-        "multiple_of": multiple_of,
-        "allow_inf_nan": allow_inf_nan,
-        "max_digits": max_digits,
-        "decimal_places": decimal_places,
-        "min_length": min_length,
-        "max_length": max_length,
-    }
-
-    # We only want to pass the args that were provided, otherwise the `Field()`` function won't work as expected
-    provided_args = {k: v for (k, v) in field_args.items() if v is not PydanticUndefined}
-
-    # Because we are manually making fields optional, we need to store the original required bool for reference later
-    json_schema_extra_.orig_required = default is PydanticUndefined
-
-    # Make Input.Any and Input.Connection fields optional, providing None as a default if the field doesn't already have one
-    if input is Input.Any or input is Input.Connection:
-        default_ = None if default is PydanticUndefined else default
-        provided_args.update({"default": default_})
-        if default is not PydanticUndefined:
-            # Before invoking, we'll check for the original default value and set it on the field if the field has no value
-            json_schema_extra_.default = default
-            json_schema_extra_.orig_default = default
-    elif default is not PydanticUndefined:
-        default_ = default
-        provided_args.update({"default": default_})
-        json_schema_extra_.orig_default = default_
-
-    return Field(
-        **provided_args,
-        json_schema_extra=json_schema_extra_.model_dump(exclude_none=True),
-    )
-
-
-def OutputField(
-    # copied from pydantic's Field
-    default: Any = _Unset,
-    title: str | None = _Unset,
-    description: str | None = _Unset,
-    pattern: str | None = _Unset,
-    strict: bool | None = _Unset,
-    gt: float | None = _Unset,
-    ge: float | None = _Unset,
-    lt: float | None = _Unset,
-    le: float | None = _Unset,
-    multiple_of: float | None = _Unset,
-    allow_inf_nan: bool | None = _Unset,
-    max_digits: int | None = _Unset,
-    decimal_places: int | None = _Unset,
-    min_length: int | None = _Unset,
-    max_length: int | None = _Unset,
-    # custom
-    ui_type: Optional[UIType] = None,
-    ui_hidden: bool = False,
-    ui_order: Optional[int] = None,
-) -> Any:
-    """
-    Creates an output field for an invocation output.
-
-    This is a wrapper for Pydantic's [Field](https://docs.pydantic.dev/1.10/usage/schema/#field-customization) \
-    that adds a few extra parameters to support graph execution and the node editor UI.
-
-    :param UIType ui_type: [None] Optionally provides an extra type hint for the UI. \
-      In some situations, the field's type is not enough to infer the correct UI type. \
-      For example, model selection fields should render a dropdown UI component to select a model. \
-      Internally, there is no difference between SD-1, SD-2 and SDXL model fields, they all use \
-      `MainModelField`. So to ensure the base-model-specific UI is rendered, you can use \
-      `UIType.SDXLMainModelField` to indicate that the field is an SDXL main model field.
-
-    :param bool ui_hidden: [False] Specifies whether or not this field should be hidden in the UI. \
-
-    :param int ui_order: [None] Specifies the order in which this field should be rendered in the UI. \
-    """
-    return Field(
-        default=default,
-        title=title,
-        description=description,
-        pattern=pattern,
-        strict=strict,
-        gt=gt,
-        ge=ge,
-        lt=lt,
-        le=le,
-        multiple_of=multiple_of,
-        allow_inf_nan=allow_inf_nan,
-        max_digits=max_digits,
-        decimal_places=decimal_places,
-        min_length=min_length,
-        max_length=max_length,
-        json_schema_extra=OutputFieldJSONSchemaExtra(
-            ui_type=ui_type,
-            ui_hidden=ui_hidden,
-            ui_order=ui_order,
-            field_kind=FieldKind.Output,
-        ).model_dump(exclude_none=True),
-    )

invokeai/app/invocations/infill.py

@@ -6,16 +6,14 @@ from typing import Literal, Optional, get_args
 import numpy as np
 from PIL import Image, ImageOps
 
-from invokeai.app.invocations.fields import ColorField, ImageField
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.invocations.primitives import ColorField, ImageField, ImageOutput
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
 from invokeai.app.util.misc import SEED_MAX
 from invokeai.backend.image_util.cv2_inpaint import cv2_inpaint
 from invokeai.backend.image_util.lama import LaMA
 from invokeai.backend.image_util.patchmatch import PatchMatch
 
-from .baseinvocation import BaseInvocation, invocation
-from .fields import InputField, WithBoard, WithMetadata
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, WithMetadata, invocation
 from .image import PIL_RESAMPLING_MAP, PIL_RESAMPLING_MODES
 
 
@@ -120,8 +118,8 @@ def tile_fill_missing(im: Image.Image, tile_size: int = 16, seed: Optional[int]
     return si
 
 
-@invocation("infill_rgba", title="Solid Color Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
-class InfillColorInvocation(BaseInvocation, WithMetadata, WithBoard):
+@invocation("infill_rgba", title="Solid Color Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.0")
+class InfillColorInvocation(BaseInvocation, WithMetadata):
     """Infills transparent areas of an image with a solid color"""
 
     image: ImageField = InputField(description="The image to infill")
@@ -131,20 +129,33 @@ class InfillColorInvocation(BaseInvocation, WithMetadata, WithBoard):
     )
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
 
         solid_bg = Image.new("RGBA", image.size, self.color.tuple())
         infilled = Image.alpha_composite(solid_bg, image.convert("RGBA"))
 
         infilled.paste(image, (0, 0), image.split()[-1])
 
-        image_dto = context.images.save(image=infilled)
+        image_dto = context.services.images.create(
+            image=infilled,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )
 
 
-@invocation("infill_tile", title="Tile Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.2")
-class InfillTileInvocation(BaseInvocation, WithMetadata, WithBoard):
+@invocation("infill_tile", title="Tile Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
+class InfillTileInvocation(BaseInvocation, WithMetadata):
     """Infills transparent areas of an image with tiles of the image"""
 
     image: ImageField = InputField(description="The image to infill")
@@ -157,20 +168,33 @@ class InfillTileInvocation(BaseInvocation, WithMetadata, WithBoard):
     )
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
 
         infilled = tile_fill_missing(image.copy(), seed=self.seed, tile_size=self.tile_size)
         infilled.paste(image, (0, 0), image.split()[-1])
 
-        image_dto = context.images.save(image=infilled)
+        image_dto = context.services.images.create(
+            image=infilled,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )
 
 
 @invocation(
-    "infill_patchmatch", title="PatchMatch Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1"
+    "infill_patchmatch", title="PatchMatch Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.0"
 )
-class InfillPatchMatchInvocation(BaseInvocation, WithMetadata, WithBoard):
+class InfillPatchMatchInvocation(BaseInvocation, WithMetadata):
     """Infills transparent areas of an image using the PatchMatch algorithm"""
 
     image: ImageField = InputField(description="The image to infill")
@@ -178,7 +202,7 @@ class InfillPatchMatchInvocation(BaseInvocation, WithMetadata, WithBoard):
     resample_mode: PIL_RESAMPLING_MODES = InputField(default="bicubic", description="The resampling mode")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name).convert("RGBA")
+        image = context.services.images.get_pil_image(self.image.image_name).convert("RGBA")
 
         resample_mode = PIL_RESAMPLING_MAP[self.resample_mode]
 
@@ -203,38 +227,77 @@ class InfillPatchMatchInvocation(BaseInvocation, WithMetadata, WithBoard):
         infilled.paste(image, (0, 0), mask=image.split()[-1])
         # image.paste(infilled, (0, 0), mask=image.split()[-1])
 
-        image_dto = context.images.save(image=infilled)
+        image_dto = context.services.images.create(
+            image=infilled,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )
 
 
-@invocation("infill_lama", title="LaMa Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
-class LaMaInfillInvocation(BaseInvocation, WithMetadata, WithBoard):
+@invocation("infill_lama", title="LaMa Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.0")
+class LaMaInfillInvocation(BaseInvocation, WithMetadata):
     """Infills transparent areas of an image using the LaMa model"""
 
     image: ImageField = InputField(description="The image to infill")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
 
         infilled = infill_lama(image.copy())
 
-        image_dto = context.images.save(image=infilled)
+        image_dto = context.services.images.create(
+            image=infilled,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )
 
 
-@invocation("infill_cv2", title="CV2 Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
-class CV2InfillInvocation(BaseInvocation, WithMetadata, WithBoard):
+@invocation("infill_cv2", title="CV2 Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.0")
+class CV2InfillInvocation(BaseInvocation, WithMetadata):
     """Infills transparent areas of an image using OpenCV Inpainting"""
 
     image: ImageField = InputField(description="The image to infill")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
 
         infilled = infill_cv2(image.copy())
 
-        image_dto = context.images.save(image=infilled)
+        image_dto = context.services.images.create(
+            image=infilled,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )

invokeai/app/invocations/ip_adapter.py

@@ -1,29 +1,38 @@
+import os
 from builtins import float
 from typing import List, Union
 
-from pydantic import BaseModel, Field, field_validator, model_validator
-from typing_extensions import Self
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField
 from invokeai.app.invocations.primitives import ImageField
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager.config import BaseModelType, ModelType
+from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.backend.model_management.models.base import BaseModelType, ModelType
+from invokeai.backend.model_management.models.ip_adapter import get_ip_adapter_image_encoder_model_id
 
 
-# LS: Consider moving these two classes into model.py
 class IPAdapterModelField(BaseModel):
-    key: str = Field(description="Key to the IP-Adapter model")
+    model_name: str = Field(description="Name of the IP-Adapter model")
+    base_model: BaseModelType = Field(description="Base model")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class CLIPVisionModelField(BaseModel):
-    key: str = Field(description="Key to the CLIP Vision image encoder model")
+    model_name: str = Field(description="Name of the CLIP Vision image encoder model")
+    base_model: BaseModelType = Field(description="Base model (usually 'Any')")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class IPAdapterField(BaseModel):
@@ -40,12 +49,12 @@ class IPAdapterField(BaseModel):
 
     @field_validator("weight")
     @classmethod
-    def validate_ip_adapter_weight(cls, v: float) -> float:
+    def validate_ip_adapter_weight(cls, v):
         validate_weights(v)
         return v
 
     @model_validator(mode="after")
-    def validate_begin_end_step_percent(self) -> Self:
+    def validate_begin_end_step_percent(self):
         validate_begin_end_step(self.begin_step_percent, self.end_step_percent)
         return self
 
@@ -56,7 +65,7 @@ class IPAdapterOutput(BaseInvocationOutput):
     ip_adapter: IPAdapterField = OutputField(description=FieldDescriptions.ip_adapter, title="IP-Adapter")
 
 
-@invocation("ip_adapter", title="IP-Adapter", tags=["ip_adapter", "control"], category="ip_adapter", version="1.1.2")
+@invocation("ip_adapter", title="IP-Adapter", tags=["ip_adapter", "control"], category="ip_adapter", version="1.1.1")
 class IPAdapterInvocation(BaseInvocation):
     """Collects IP-Adapter info to pass to other nodes."""
 
@@ -78,25 +87,33 @@ class IPAdapterInvocation(BaseInvocation):
 
     @field_validator("weight")
     @classmethod
-    def validate_ip_adapter_weight(cls, v: float) -> float:
+    def validate_ip_adapter_weight(cls, v):
         validate_weights(v)
         return v
 
     @model_validator(mode="after")
-    def validate_begin_end_step_percent(self) -> Self:
+    def validate_begin_end_step_percent(self):
         validate_begin_end_step(self.begin_step_percent, self.end_step_percent)
         return self
 
     def invoke(self, context: InvocationContext) -> IPAdapterOutput:
         # Lookup the CLIP Vision encoder that is intended to be used with the IP-Adapter model.
-        ip_adapter_info = context.models.get_config(self.ip_adapter_model.key)
-        image_encoder_model_id = ip_adapter_info.image_encoder_model_id
-        image_encoder_model_name = image_encoder_model_id.split("/")[-1].strip()
-        image_encoder_models = context.models.search_by_attrs(
-            model_name=image_encoder_model_name, base_model=BaseModelType.Any, model_type=ModelType.CLIPVision
+        ip_adapter_info = context.services.model_manager.model_info(
+            self.ip_adapter_model.model_name, self.ip_adapter_model.base_model, ModelType.IPAdapter
+        )
+        # HACK(ryand): This is bad for a couple of reasons: 1) we are bypassing the model manager to read the model
+        # directly, and 2) we are reading from disk every time this invocation is called without caching the result.
+        # A better solution would be to store the image encoder model reference in the IP-Adapter model info, but this
+        # is currently messy due to differences between how the model info is generated when installing a model from
+        # disk vs. downloading the model.
+        image_encoder_model_id = get_ip_adapter_image_encoder_model_id(
+            os.path.join(context.services.configuration.get_config().models_path, ip_adapter_info["path"])
+        )
+        image_encoder_model_name = image_encoder_model_id.split("/")[-1].strip()
+        image_encoder_model = CLIPVisionModelField(
+            model_name=image_encoder_model_name,
+            base_model=BaseModelType.Any,
         )
-        assert len(image_encoder_models) == 1
-        image_encoder_model = CLIPVisionModelField(key=image_encoder_models[0].key)
         return IPAdapterOutput(
             ip_adapter=IPAdapterField(
                 image=self.image,

invokeai/app/invocations/latent.py

@@ -3,15 +3,13 @@
 import math
 from contextlib import ExitStack
 from functools import singledispatchmethod
-from typing import Any, Iterator, List, Literal, Optional, Tuple, Union
+from typing import List, Literal, Optional, Union
 
 import einops
 import numpy as np
-import numpy.typing as npt
 import torch
 import torchvision.transforms as T
 from diffusers import AutoencoderKL, AutoencoderTiny
-from diffusers.configuration_utils import ConfigMixin
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.models.adapter import T2IAdapter
 from diffusers.models.attention_processor import (
@@ -20,44 +18,34 @@ from diffusers.models.attention_processor import (
     LoRAXFormersAttnProcessor,
     XFormersAttnProcessor,
 )
-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 pydantic import field_validator
 from torchvision.transforms.functional import resize as tv_resize
 
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR, SCHEDULER_NAME_VALUES
-from invokeai.app.invocations.fields import (
-    ConditioningField,
-    DenoiseMaskField,
-    FieldDescriptions,
-    ImageField,
-    Input,
-    InputField,
-    LatentsField,
-    OutputField,
-    UIType,
-    WithBoard,
-    WithMetadata,
-)
 from invokeai.app.invocations.ip_adapter import IPAdapterField
 from invokeai.app.invocations.primitives import (
+    DenoiseMaskField,
     DenoiseMaskOutput,
+    ImageField,
     ImageOutput,
+    LatentsField,
     LatentsOutput,
+    build_latents_output,
 )
 from invokeai.app.invocations.t2i_adapter import T2IAdapterField
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
+from invokeai.app.shared.fields import FieldDescriptions
 from invokeai.app.util.controlnet_utils import prepare_control_image
+from invokeai.app.util.step_callback import stable_diffusion_step_callback
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter, IPAdapterPlus
-from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_manager import BaseModelType, LoadedModel
-from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.stable_diffusion import PipelineIntermediateState, set_seamless
+from invokeai.backend.model_management.models import ModelType, SilenceWarnings
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningData, IPAdapterConditioningInfo
-from invokeai.backend.util.silence_warnings import SilenceWarnings
 
+from ...backend.model_management.lora import ModelPatcher
+from ...backend.model_management.models import BaseModelType
+from ...backend.model_management.seamless import set_seamless
+from ...backend.stable_diffusion import PipelineIntermediateState
 from ...backend.stable_diffusion.diffusers_pipeline import (
     ControlNetData,
     IPAdapterData,
@@ -71,9 +59,16 @@ from ...backend.util.devices import choose_precision, choose_torch_device
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    UIType,
+    WithMetadata,
     invocation,
     invocation_output,
 )
+from .compel import ConditioningField
 from .controlnet_image_processors import ControlField
 from .model import ModelInfo, UNetField, VaeField
 
@@ -82,10 +77,18 @@ if choose_torch_device() == torch.device("mps"):
 
 DEFAULT_PRECISION = choose_precision(choose_torch_device())
 
+SAMPLER_NAME_VALUES = Literal[tuple(SCHEDULER_MAP.keys())]
+
+# HACK: Many nodes are currently hard-coded to use a fixed latent scale factor of 8. This is fragile, and will need to
+# be addressed if future models use a different latent scale factor. Also, note that there may be places where the scale
+# factor is hard-coded to a literal '8' rather than using this constant.
+# The ratio of image:latent dimensions is LATENT_SCALE_FACTOR:1, or 8:1.
+LATENT_SCALE_FACTOR = 8
+
 
 @invocation_output("scheduler_output")
 class SchedulerOutput(BaseInvocationOutput):
-    scheduler: SCHEDULER_NAME_VALUES = OutputField(description=FieldDescriptions.scheduler, ui_type=UIType.Scheduler)
+    scheduler: SAMPLER_NAME_VALUES = OutputField(description=FieldDescriptions.scheduler, ui_type=UIType.Scheduler)
 
 
 @invocation(
@@ -98,7 +101,7 @@ class SchedulerOutput(BaseInvocationOutput):
 class SchedulerInvocation(BaseInvocation):
     """Selects a scheduler."""
 
-    scheduler: SCHEDULER_NAME_VALUES = InputField(
+    scheduler: SAMPLER_NAME_VALUES = InputField(
         default="euler",
         description=FieldDescriptions.scheduler,
         ui_type=UIType.Scheduler,
@@ -113,7 +116,7 @@ class SchedulerInvocation(BaseInvocation):
     title="Create Denoise Mask",
     tags=["mask", "denoise"],
     category="latents",
-    version="1.0.1",
+    version="1.0.0",
 )
 class CreateDenoiseMaskInvocation(BaseInvocation):
     """Creates mask for denoising model run."""
@@ -128,10 +131,10 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
         ui_order=4,
     )
 
-    def prep_mask_tensor(self, mask_image: Image.Image) -> torch.Tensor:
+    def prep_mask_tensor(self, mask_image):
         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)
+        mask_tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
         if mask_tensor.dim() == 3:
             mask_tensor = mask_tensor.unsqueeze(0)
         # if shape is not None:
@@ -141,90 +144,41 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> DenoiseMaskOutput:
         if self.image is not None:
-            image = context.images.get_pil(self.image.image_name)
-            image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
-            if image_tensor.dim() == 3:
-                image_tensor = image_tensor.unsqueeze(0)
+            image = context.services.images.get_pil_image(self.image.image_name)
+            image = image_resized_to_grid_as_tensor(image.convert("RGB"))
+            if image.dim() == 3:
+                image = image.unsqueeze(0)
         else:
-            image_tensor = None
+            image = None
 
         mask = self.prep_mask_tensor(
-            context.images.get_pil(self.mask.image_name),
+            context.services.images.get_pil_image(self.mask.image_name),
         )
 
-        if image_tensor is not None:
-            vae_info = context.models.load(**self.vae.vae.model_dump())
+        if image is not None:
+            vae_info = context.services.model_manager.get_model(
+                **self.vae.vae.model_dump(),
+                context=context,
+            )
 
-            img_mask = tv_resize(mask, image_tensor.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
-            masked_image = image_tensor * torch.where(img_mask < 0.5, 0.0, 1.0)
+            img_mask = tv_resize(mask, image.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
+            masked_image = image * torch.where(img_mask < 0.5, 0.0, 1.0)
             # TODO:
             masked_latents = ImageToLatentsInvocation.vae_encode(vae_info, self.fp32, self.tiled, masked_image.clone())
 
-            masked_latents_name = context.tensors.save(tensor=masked_latents)
+            masked_latents_name = f"{context.graph_execution_state_id}__{self.id}_masked_latents"
+            context.services.latents.save(masked_latents_name, masked_latents)
         else:
             masked_latents_name = None
 
-        mask_name = context.tensors.save(tensor=mask)
+        mask_name = f"{context.graph_execution_state_id}__{self.id}_mask"
+        context.services.latents.save(mask_name, mask)
 
-        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,
+        return DenoiseMaskOutput(
+            denoise_mask=DenoiseMaskField(
+                mask_name=mask_name,
+                masked_latents_name=masked_latents_name,
+            ),
         )
 
 
@@ -235,7 +189,10 @@ def get_scheduler(
     seed: int,
 ) -> Scheduler:
     scheduler_class, scheduler_extra_config = SCHEDULER_MAP.get(scheduler_name, SCHEDULER_MAP["ddim"])
-    orig_scheduler_info = context.models.load(**scheduler_info.model_dump())
+    orig_scheduler_info = context.services.model_manager.get_model(
+        **scheduler_info.model_dump(),
+        context=context,
+    )
     with orig_scheduler_info as orig_scheduler:
         scheduler_config = orig_scheduler.config
 
@@ -243,7 +200,7 @@ def get_scheduler(
         scheduler_config = scheduler_config["_backup"]
     scheduler_config = {
         **scheduler_config,
-        **scheduler_extra_config,  # FIXME
+        **scheduler_extra_config,
         "_backup": scheduler_config,
     }
 
@@ -256,7 +213,6 @@ def get_scheduler(
     # hack copied over from generate.py
     if not hasattr(scheduler, "uses_inpainting_model"):
         scheduler.uses_inpainting_model = lambda: False
-    assert isinstance(scheduler, Scheduler)
     return scheduler
 
 
@@ -265,7 +221,7 @@ def get_scheduler(
     title="Denoise Latents",
     tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
     category="latents",
-    version="1.5.2",
+    version="1.5.1",
 )
 class DenoiseLatentsInvocation(BaseInvocation):
     """Denoises noisy latents to decodable images"""
@@ -293,7 +249,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         description=FieldDescriptions.denoising_start,
     )
     denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
-    scheduler: SCHEDULER_NAME_VALUES = InputField(
+    scheduler: SAMPLER_NAME_VALUES = InputField(
         default="euler",
         description=FieldDescriptions.scheduler,
         ui_type=UIType.Scheduler,
@@ -340,7 +296,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
     )
 
     @field_validator("cfg_scale")
-    def ge_one(cls, v: Union[List[float], float]) -> Union[List[float], float]:
+    def ge_one(cls, v):
         """validate that all cfg_scale values are >= 1"""
         if isinstance(v, list):
             for i in v:
@@ -351,18 +307,34 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 raise ValueError("cfg_scale must be greater than 1")
         return v
 
+    # TODO: pass this an emitter method or something? or a session for dispatching?
+    def dispatch_progress(
+        self,
+        context: InvocationContext,
+        source_node_id: str,
+        intermediate_state: PipelineIntermediateState,
+        base_model: BaseModelType,
+    ) -> None:
+        stable_diffusion_step_callback(
+            context=context,
+            intermediate_state=intermediate_state,
+            node=self.model_dump(),
+            source_node_id=source_node_id,
+            base_model=base_model,
+        )
+
     def get_conditioning_data(
         self,
         context: InvocationContext,
-        scheduler: Scheduler,
-        unet: UNet2DConditionModel,
-        seed: int,
+        scheduler,
+        unet,
+        seed,
     ) -> ConditioningData:
-        positive_cond_data = context.conditioning.load(self.positive_conditioning.conditioning_name)
+        positive_cond_data = context.services.latents.get(self.positive_conditioning.conditioning_name)
         c = positive_cond_data.conditionings[0].to(device=unet.device, dtype=unet.dtype)
         extra_conditioning_info = c.extra_conditioning
 
-        negative_cond_data = context.conditioning.load(self.negative_conditioning.conditioning_name)
+        negative_cond_data = context.services.latents.get(self.negative_conditioning.conditioning_name)
         uc = negative_cond_data.conditionings[0].to(device=unet.device, dtype=unet.dtype)
 
         conditioning_data = ConditioningData(
@@ -379,7 +351,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
             ),
         )
 
-        conditioning_data = conditioning_data.add_scheduler_args_if_applicable(  # FIXME
+        conditioning_data = conditioning_data.add_scheduler_args_if_applicable(
             scheduler,
             # for ddim scheduler
             eta=0.0,  # ddim_eta
@@ -391,8 +363,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
     def create_pipeline(
         self,
-        unet: UNet2DConditionModel,
-        scheduler: Scheduler,
+        unet,
+        scheduler,
     ) -> StableDiffusionGeneratorPipeline:
         # TODO:
         # configure_model_padding(
@@ -403,10 +375,10 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         class FakeVae:
             class FakeVaeConfig:
-                def __init__(self) -> None:
+                def __init__(self):
                     self.block_out_channels = [0]
 
-            def __init__(self) -> None:
+            def __init__(self):
                 self.config = FakeVae.FakeVaeConfig()
 
         return StableDiffusionGeneratorPipeline(
@@ -423,11 +395,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
     def prep_control_data(
         self,
         context: InvocationContext,
-        control_input: Optional[Union[ControlField, List[ControlField]]],
+        control_input: Union[ControlField, List[ControlField]],
         latents_shape: List[int],
         exit_stack: ExitStack,
         do_classifier_free_guidance: bool = True,
-    ) -> Optional[List[ControlNetData]]:
+    ) -> List[ControlNetData]:
         # Assuming fixed dimensional scaling of LATENT_SCALE_FACTOR.
         control_height_resize = latents_shape[2] * LATENT_SCALE_FACTOR
         control_width_resize = latents_shape[3] * LATENT_SCALE_FACTOR
@@ -449,11 +421,18 @@ class DenoiseLatentsInvocation(BaseInvocation):
         #        and if weight is None, populate with default 1.0?
         controlnet_data = []
         for control_info in control_list:
-            control_model = exit_stack.enter_context(context.models.load(key=control_info.control_model.key))
+            control_model = exit_stack.enter_context(
+                context.services.model_manager.get_model(
+                    model_name=control_info.control_model.model_name,
+                    model_type=ModelType.ControlNet,
+                    base_model=control_info.control_model.base_model,
+                    context=context,
+                )
+            )
 
             # control_models.append(control_model)
             control_image_field = control_info.image
-            input_image = context.images.get_pil(control_image_field.image_name)
+            input_image = context.services.images.get_pil_image(control_image_field.image_name)
             # self.image.image_type, self.image.image_name
             # FIXME: still need to test with different widths, heights, devices, dtypes
             #        and add in batch_size, num_images_per_prompt?
@@ -511,17 +490,27 @@ class DenoiseLatentsInvocation(BaseInvocation):
         conditioning_data.ip_adapter_conditioning = []
         for single_ip_adapter in ip_adapter:
             ip_adapter_model: Union[IPAdapter, IPAdapterPlus] = exit_stack.enter_context(
-                context.models.load(key=single_ip_adapter.ip_adapter_model.key)
+                context.services.model_manager.get_model(
+                    model_name=single_ip_adapter.ip_adapter_model.model_name,
+                    model_type=ModelType.IPAdapter,
+                    base_model=single_ip_adapter.ip_adapter_model.base_model,
+                    context=context,
+                )
             )
 
-            image_encoder_model_info = context.models.load(key=single_ip_adapter.image_encoder_model.key)
+            image_encoder_model_info = context.services.model_manager.get_model(
+                model_name=single_ip_adapter.image_encoder_model.model_name,
+                model_type=ModelType.CLIPVision,
+                base_model=single_ip_adapter.image_encoder_model.base_model,
+                context=context,
+            )
 
             # `single_ip_adapter.image` could be a list or a single ImageField. Normalize to a list here.
-            single_ipa_image_fields = single_ip_adapter.image
-            if not isinstance(single_ipa_image_fields, list):
-                single_ipa_image_fields = [single_ipa_image_fields]
+            single_ipa_images = single_ip_adapter.image
+            if not isinstance(single_ipa_images, list):
+                single_ipa_images = [single_ipa_images]
 
-            single_ipa_images = [context.images.get_pil(image.image_name) for image in single_ipa_image_fields]
+            single_ipa_images = [context.services.images.get_pil_image(image.image_name) for image in single_ipa_images]
 
             # TODO(ryand): With some effort, the step of running the CLIP Vision encoder could be done before any other
             # models are needed in memory. This would help to reduce peak memory utilization in low-memory environments.
@@ -565,20 +554,26 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         t2i_adapter_data = []
         for t2i_adapter_field in t2i_adapter:
-            t2i_adapter_model_config = context.models.get_config(key=t2i_adapter_field.t2i_adapter_model.key)
-            t2i_adapter_loaded_model = context.models.load(key=t2i_adapter_field.t2i_adapter_model.key)
-            image = context.images.get_pil(t2i_adapter_field.image.image_name)
+            t2i_adapter_model_info = context.services.model_manager.get_model(
+                model_name=t2i_adapter_field.t2i_adapter_model.model_name,
+                model_type=ModelType.T2IAdapter,
+                base_model=t2i_adapter_field.t2i_adapter_model.base_model,
+                context=context,
+            )
+            image = context.services.images.get_pil_image(t2i_adapter_field.image.image_name)
 
             # The max_unet_downscale is the maximum amount that the UNet model downscales the latent image internally.
-            if t2i_adapter_model_config.base == BaseModelType.StableDiffusion1:
+            if t2i_adapter_field.t2i_adapter_model.base_model == BaseModelType.StableDiffusion1:
                 max_unet_downscale = 8
-            elif t2i_adapter_model_config.base == BaseModelType.StableDiffusionXL:
+            elif t2i_adapter_field.t2i_adapter_model.base_model == BaseModelType.StableDiffusionXL:
                 max_unet_downscale = 4
             else:
-                raise ValueError(f"Unexpected T2I-Adapter base model type: '{t2i_adapter_model_config.base}'.")
+                raise ValueError(
+                    f"Unexpected T2I-Adapter base model type: '{t2i_adapter_field.t2i_adapter_model.base_model}'."
+                )
 
             t2i_adapter_model: T2IAdapter
-            with t2i_adapter_loaded_model as t2i_adapter_model:
+            with t2i_adapter_model_info as t2i_adapter_model:
                 total_downscale_factor = t2i_adapter_model.total_downscale_factor
 
                 # Resize the T2I-Adapter input image.
@@ -598,7 +593,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
                     do_classifier_free_guidance=False,
                     width=t2i_input_width,
                     height=t2i_input_height,
-                    num_channels=t2i_adapter_model.config["in_channels"],  # mypy treats this as a FrozenDict
+                    num_channels=t2i_adapter_model.config.in_channels,
                     device=t2i_adapter_model.device,
                     dtype=t2i_adapter_model.dtype,
                     resize_mode=t2i_adapter_field.resize_mode,
@@ -623,15 +618,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
     # original idea by https://github.com/AmericanPresidentJimmyCarter
     # TODO: research more for second order schedulers timesteps
-    def init_scheduler(
-        self,
-        scheduler: Union[Scheduler, ConfigMixin],
-        device: torch.device,
-        steps: int,
-        denoising_start: float,
-        denoising_end: float,
-    ) -> Tuple[int, List[int], int]:
-        assert isinstance(scheduler, ConfigMixin)
+    def init_scheduler(self, scheduler, device, steps, denoising_start, denoising_end):
         if scheduler.config.get("cpu_only", False):
             scheduler.set_timesteps(steps, device="cpu")
             timesteps = scheduler.timesteps.to(device=device)
@@ -643,11 +630,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
         _timesteps = timesteps[:: scheduler.order]
 
         # get start timestep index
-        t_start_val = int(round(scheduler.config["num_train_timesteps"] * (1 - denoising_start)))
+        t_start_val = int(round(scheduler.config.num_train_timesteps * (1 - denoising_start)))
         t_start_idx = len(list(filter(lambda ts: ts >= t_start_val, _timesteps)))
 
         # get end timestep index
-        t_end_val = int(round(scheduler.config["num_train_timesteps"] * (1 - denoising_end)))
+        t_end_val = int(round(scheduler.config.num_train_timesteps * (1 - denoising_end)))
         t_end_idx = len(list(filter(lambda ts: ts >= t_end_val, _timesteps[t_start_idx:])))
 
         # apply order to indexes
@@ -660,20 +647,18 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         return num_inference_steps, timesteps, init_timestep
 
-    def prep_inpaint_mask(
-        self, context: InvocationContext, latents: torch.Tensor
-    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], bool]:
+    def prep_inpaint_mask(self, context, latents):
         if self.denoise_mask is None:
-            return None, None, False
+            return None, None
 
-        mask = context.tensors.load(self.denoise_mask.mask_name)
+        mask = context.services.latents.get(self.denoise_mask.mask_name)
         mask = tv_resize(mask, latents.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
         if self.denoise_mask.masked_latents_name is not None:
-            masked_latents = context.tensors.load(self.denoise_mask.masked_latents_name)
+            masked_latents = context.services.latents.get(self.denoise_mask.masked_latents_name)
         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:
@@ -681,11 +666,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
             seed = None
             noise = None
             if self.noise is not None:
-                noise = context.tensors.load(self.noise.latents_name)
+                noise = context.services.latents.get(self.noise.latents_name)
                 seed = self.noise.seed
 
             if self.latents is not None:
-                latents = context.tensors.load(self.latents.latents_name)
+                latents = context.services.latents.get(self.latents.latents_name)
                 if seed is None:
                     seed = self.latents.seed
 
@@ -700,7 +685,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.
@@ -711,30 +696,35 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 do_classifier_free_guidance=True,
             )
 
-            # get the unet's config so that we can pass the base to dispatch_progress()
-            unet_config = context.models.get_config(self.unet.unet.key)
+            # Get the source node id (we are invoking the prepared node)
+            graph_execution_state = context.services.graph_execution_manager.get(context.graph_execution_state_id)
+            source_node_id = graph_execution_state.prepared_source_mapping[self.id]
 
-            def step_callback(state: PipelineIntermediateState) -> None:
-                context.util.sd_step_callback(state, unet_config.base)
+            def step_callback(state: PipelineIntermediateState):
+                self.dispatch_progress(context, source_node_id, state, self.unet.unet.base_model)
 
-            def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
+            def _lora_loader():
                 for lora in self.unet.loras:
-                    lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
-                    yield (lora_info.model, lora.weight)
+                    lora_info = context.services.model_manager.get_model(
+                        **lora.model_dump(exclude={"weight"}),
+                        context=context,
+                    )
+                    yield (lora_info.context.model, lora.weight)
                     del lora_info
                 return
 
-            unet_info = context.models.load(**self.unet.unet.model_dump())
-            assert isinstance(unet_info.model, UNet2DConditionModel)
+            unet_info = context.services.model_manager.get_model(
+                **self.unet.unet.model_dump(),
+                context=context,
+            )
             with (
                 ExitStack() as exit_stack,
-                ModelPatcher.apply_freeu(unet_info.model, self.unet.freeu_config),
-                set_seamless(unet_info.model, self.unet.seamless_axes),  # FIXME
+                ModelPatcher.apply_freeu(unet_info.context.model, self.unet.freeu_config),
+                set_seamless(unet_info.context.model, self.unet.seamless_axes),
                 unet_info as unet,
                 # Apply the LoRA after unet has been moved to its target device for faster patching.
                 ModelPatcher.apply_lora_unet(unet, _lora_loader()),
             ):
-                assert isinstance(unet, UNet2DConditionModel)
                 latents = latents.to(device=unet.device, dtype=unet.dtype)
                 if noise is not None:
                     noise = noise.to(device=unet.device, dtype=unet.dtype)
@@ -788,7 +778,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,
@@ -803,8 +792,9 @@ class DenoiseLatentsInvocation(BaseInvocation):
             if choose_torch_device() == torch.device("mps"):
                 mps.empty_cache()
 
-            name = context.tensors.save(tensor=result_latents)
-        return LatentsOutput.build(latents_name=name, latents=result_latents, seed=seed)
+            name = f"{context.graph_execution_state_id}__{self.id}"
+            context.services.latents.save(name, result_latents)
+        return build_latents_output(latents_name=name, latents=result_latents, seed=seed)
 
 
 @invocation(
@@ -812,9 +802,9 @@ class DenoiseLatentsInvocation(BaseInvocation):
     title="Latents to Image",
     tags=["latents", "image", "vae", "l2i"],
     category="latents",
-    version="1.2.1",
+    version="1.2.0",
 )
-class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
+class LatentsToImageInvocation(BaseInvocation, WithMetadata):
     """Generates an image from latents."""
 
     latents: LatentsField = InputField(
@@ -830,12 +820,14 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        latents = context.tensors.load(self.latents.latents_name)
+        latents = context.services.latents.get(self.latents.latents_name)
 
-        vae_info = context.models.load(**self.vae.vae.model_dump())
+        vae_info = context.services.model_manager.get_model(
+            **self.vae.vae.model_dump(),
+            context=context,
+        )
 
-        with set_seamless(vae_info.model, self.vae.seamless_axes), vae_info as vae:
-            assert isinstance(vae, torch.nn.Module)
+        with set_seamless(vae_info.context.model, self.vae.seamless_axes), vae_info as vae:
             latents = latents.to(vae.device)
             if self.fp32:
                 vae.to(dtype=torch.float32)
@@ -862,7 +854,7 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
                 vae.to(dtype=torch.float16)
                 latents = latents.half()
 
-            if self.tiled or context.config.get().tiled_decode:
+            if self.tiled or context.services.configuration.tiled_decode:
                 vae.enable_tiling()
             else:
                 vae.disable_tiling()
@@ -886,9 +878,22 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
         if choose_torch_device() == torch.device("mps"):
             mps.empty_cache()
 
-        image_dto = context.images.save(image=image)
+        image_dto = context.services.images.create(
+            image=image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )
 
 
 LATENTS_INTERPOLATION_MODE = Literal["nearest", "linear", "bilinear", "bicubic", "trilinear", "area", "nearest-exact"]
@@ -899,7 +904,7 @@ LATENTS_INTERPOLATION_MODE = Literal["nearest", "linear", "bilinear", "bicubic",
     title="Resize Latents",
     tags=["latents", "resize"],
     category="latents",
-    version="1.0.1",
+    version="1.0.0",
 )
 class ResizeLatentsInvocation(BaseInvocation):
     """Resizes latents to explicit width/height (in pixels). Provided dimensions are floor-divided by 8."""
@@ -922,7 +927,7 @@ class ResizeLatentsInvocation(BaseInvocation):
     antialias: bool = InputField(default=False, description=FieldDescriptions.torch_antialias)
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.tensors.load(self.latents.latents_name)
+        latents = context.services.latents.get(self.latents.latents_name)
 
         # TODO:
         device = choose_torch_device()
@@ -940,8 +945,10 @@ class ResizeLatentsInvocation(BaseInvocation):
         if device == torch.device("mps"):
             mps.empty_cache()
 
-        name = context.tensors.save(tensor=resized_latents)
-        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
+        name = f"{context.graph_execution_state_id}__{self.id}"
+        # context.services.latents.set(name, resized_latents)
+        context.services.latents.save(name, resized_latents)
+        return build_latents_output(latents_name=name, latents=resized_latents, seed=self.latents.seed)
 
 
 @invocation(
@@ -949,7 +956,7 @@ class ResizeLatentsInvocation(BaseInvocation):
     title="Scale Latents",
     tags=["latents", "resize"],
     category="latents",
-    version="1.0.1",
+    version="1.0.0",
 )
 class ScaleLatentsInvocation(BaseInvocation):
     """Scales latents by a given factor."""
@@ -963,7 +970,7 @@ class ScaleLatentsInvocation(BaseInvocation):
     antialias: bool = InputField(default=False, description=FieldDescriptions.torch_antialias)
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.tensors.load(self.latents.latents_name)
+        latents = context.services.latents.get(self.latents.latents_name)
 
         # TODO:
         device = choose_torch_device()
@@ -982,8 +989,10 @@ class ScaleLatentsInvocation(BaseInvocation):
         if device == torch.device("mps"):
             mps.empty_cache()
 
-        name = context.tensors.save(tensor=resized_latents)
-        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
+        name = f"{context.graph_execution_state_id}__{self.id}"
+        # context.services.latents.set(name, resized_latents)
+        context.services.latents.save(name, resized_latents)
+        return build_latents_output(latents_name=name, latents=resized_latents, seed=self.latents.seed)
 
 
 @invocation(
@@ -991,7 +1000,7 @@ class ScaleLatentsInvocation(BaseInvocation):
     title="Image to Latents",
     tags=["latents", "image", "vae", "i2l"],
     category="latents",
-    version="1.0.1",
+    version="1.0.0",
 )
 class ImageToLatentsInvocation(BaseInvocation):
     """Encodes an image into latents."""
@@ -1007,9 +1016,8 @@ class ImageToLatentsInvocation(BaseInvocation):
     fp32: bool = InputField(default=DEFAULT_PRECISION == "float32", description=FieldDescriptions.fp32)
 
     @staticmethod
-    def vae_encode(vae_info: LoadedModel, upcast: bool, tiled: bool, image_tensor: torch.Tensor) -> torch.Tensor:
+    def vae_encode(vae_info, upcast, tiled, image_tensor):
         with vae_info as vae:
-            assert isinstance(vae, torch.nn.Module)
             orig_dtype = vae.dtype
             if upcast:
                 vae.to(dtype=torch.float32)
@@ -1053,9 +1061,12 @@ class ImageToLatentsInvocation(BaseInvocation):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
 
-        vae_info = context.models.load(**self.vae.vae.model_dump())
+        vae_info = context.services.model_manager.get_model(
+            **self.vae.vae.model_dump(),
+            context=context,
+        )
 
         image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
         if image_tensor.dim() == 3:
@@ -1063,26 +1074,22 @@ class ImageToLatentsInvocation(BaseInvocation):
 
         latents = self.vae_encode(vae_info, self.fp32, self.tiled, image_tensor)
 
+        name = f"{context.graph_execution_state_id}__{self.id}"
         latents = latents.to("cpu")
-        name = context.tensors.save(tensor=latents)
-        return LatentsOutput.build(latents_name=name, latents=latents, seed=None)
+        context.services.latents.save(name, latents)
+        return build_latents_output(latents_name=name, latents=latents, seed=None)
 
     @singledispatchmethod
     @staticmethod
     def _encode_to_tensor(vae: AutoencoderKL, image_tensor: torch.FloatTensor) -> torch.FloatTensor:
-        assert isinstance(vae, torch.nn.Module)
         image_tensor_dist = vae.encode(image_tensor).latent_dist
-        latents: torch.Tensor = image_tensor_dist.sample().to(
-            dtype=vae.dtype
-        )  # FIXME: uses torch.randn. make reproducible!
+        latents = image_tensor_dist.sample().to(dtype=vae.dtype)  # FIXME: uses torch.randn. make reproducible!
         return latents
 
     @_encode_to_tensor.register
     @staticmethod
     def _(vae: AutoencoderTiny, image_tensor: torch.FloatTensor) -> torch.FloatTensor:
-        assert isinstance(vae, torch.nn.Module)
-        latents: torch.FloatTensor = vae.encode(image_tensor).latents
-        return latents
+        return vae.encode(image_tensor).latents
 
 
 @invocation(
@@ -1090,7 +1097,7 @@ class ImageToLatentsInvocation(BaseInvocation):
     title="Blend Latents",
     tags=["latents", "blend"],
     category="latents",
-    version="1.0.1",
+    version="1.0.0",
 )
 class BlendLatentsInvocation(BaseInvocation):
     """Blend two latents using a given alpha. Latents must have same size."""
@@ -1106,8 +1113,8 @@ class BlendLatentsInvocation(BaseInvocation):
     alpha: float = InputField(default=0.5, description=FieldDescriptions.blend_alpha)
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents_a = context.tensors.load(self.latents_a.latents_name)
-        latents_b = context.tensors.load(self.latents_b.latents_name)
+        latents_a = context.services.latents.get(self.latents_a.latents_name)
+        latents_b = context.services.latents.get(self.latents_b.latents_name)
 
         if latents_a.shape != latents_b.shape:
             raise Exception("Latents to blend must be the same size.")
@@ -1115,12 +1122,7 @@ class BlendLatentsInvocation(BaseInvocation):
         # TODO:
         device = choose_torch_device()
 
-        def slerp(
-            t: Union[float, npt.NDArray[Any]],  # FIXME: maybe use np.float32 here?
-            v0: Union[torch.Tensor, npt.NDArray[Any]],
-            v1: Union[torch.Tensor, npt.NDArray[Any]],
-            DOT_THRESHOLD: float = 0.9995,
-        ) -> Union[torch.Tensor, npt.NDArray[Any]]:
+        def slerp(t, v0, v1, DOT_THRESHOLD=0.9995):
             """
             Spherical linear interpolation
             Args:
@@ -1153,16 +1155,12 @@ class BlendLatentsInvocation(BaseInvocation):
                 v2 = s0 * v0 + s1 * v1
 
             if inputs_are_torch:
-                v2_torch: torch.Tensor = torch.from_numpy(v2).to(device)
-                return v2_torch
-            else:
-                assert isinstance(v2, np.ndarray)
-                return v2
+                v2 = torch.from_numpy(v2).to(device)
+
+            return v2
 
         # blend
-        bl = slerp(self.alpha, latents_a, latents_b)
-        assert isinstance(bl, torch.Tensor)
-        blended_latents: torch.Tensor = bl  # for type checking convenience
+        blended_latents = slerp(self.alpha, latents_a, latents_b)
 
         # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
         blended_latents = blended_latents.to("cpu")
@@ -1170,8 +1168,10 @@ class BlendLatentsInvocation(BaseInvocation):
         if device == torch.device("mps"):
             mps.empty_cache()
 
-        name = context.tensors.save(tensor=blended_latents)
-        return LatentsOutput.build(latents_name=name, latents=blended_latents)
+        name = f"{context.graph_execution_state_id}__{self.id}"
+        # context.services.latents.set(name, resized_latents)
+        context.services.latents.save(name, blended_latents)
+        return build_latents_output(latents_name=name, latents=blended_latents)
 
 
 # The Crop Latents node was copied from @skunkworxdark's implementation here:
@@ -1181,7 +1181,7 @@ class BlendLatentsInvocation(BaseInvocation):
     title="Crop Latents",
     tags=["latents", "crop"],
     category="latents",
-    version="1.0.1",
+    version="1.0.0",
 )
 # TODO(ryand): Named `CropLatentsCoreInvocation` to prevent a conflict with custom node `CropLatentsInvocation`.
 # Currently, if the class names conflict then 'GET /openapi.json' fails.
@@ -1216,7 +1216,7 @@ class CropLatentsCoreInvocation(BaseInvocation):
     )
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.tensors.load(self.latents.latents_name)
+        latents = context.services.latents.get(self.latents.latents_name)
 
         x1 = self.x // LATENT_SCALE_FACTOR
         y1 = self.y // LATENT_SCALE_FACTOR
@@ -1225,9 +1225,10 @@ class CropLatentsCoreInvocation(BaseInvocation):
 
         cropped_latents = latents[..., y1:y2, x1:x2]
 
-        name = context.tensors.save(tensor=cropped_latents)
+        name = f"{context.graph_execution_state_id}__{self.id}"
+        context.services.latents.save(name, cropped_latents)
 
-        return LatentsOutput.build(latents_name=name, latents=cropped_latents)
+        return build_latents_output(latents_name=name, latents=cropped_latents)
 
 
 @invocation_output("ideal_size_output")
@@ -1255,16 +1256,15 @@ class IdealSizeInvocation(BaseInvocation):
         description="Amount to multiply the model's dimensions by when calculating the ideal size (may result in initial generation artifacts if too large)",
     )
 
-    def trim_to_multiple_of(self, *args: int, multiple_of: int = LATENT_SCALE_FACTOR) -> Tuple[int, ...]:
+    def trim_to_multiple_of(self, *args, multiple_of=LATENT_SCALE_FACTOR):
         return tuple((x - x % multiple_of) for x in args)
 
     def invoke(self, context: InvocationContext) -> IdealSizeOutput:
-        unet_config = context.models.get_config(**self.unet.unet.model_dump())
         aspect = self.width / self.height
-        dimension: float = 512
-        if unet_config.base == BaseModelType.StableDiffusion2:
+        dimension = 512
+        if self.unet.unet.base_model == BaseModelType.StableDiffusion2:
             dimension = 768
-        elif unet_config.base == BaseModelType.StableDiffusionXL:
+        elif self.unet.unet.base_model == BaseModelType.StableDiffusionXL:
             dimension = 1024
         dimension = dimension * self.multiplier
         min_dimension = math.floor(dimension * 0.5)

invokeai/app/invocations/math.py

@@ -5,11 +5,10 @@ from typing import Literal
 import numpy as np
 from pydantic import ValidationInfo, field_validator
 
-from invokeai.app.invocations.fields import FieldDescriptions, InputField
 from invokeai.app.invocations.primitives import FloatOutput, IntegerOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.shared.fields import FieldDescriptions
 
-from .baseinvocation import BaseInvocation, invocation
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, invocation
 
 
 @invocation("add", title="Add Integers", tags=["math", "add"], category="math", version="1.0.0")

invokeai/app/invocations/metadata.py

@@ -5,22 +5,20 @@ from pydantic import BaseModel, ConfigDict, Field
 from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    InputField,
+    InvocationContext,
+    MetadataField,
+    OutputField,
+    UIType,
     invocation,
     invocation_output,
 )
 from invokeai.app.invocations.controlnet_image_processors import ControlField
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    ImageField,
-    InputField,
-    MetadataField,
-    OutputField,
-    UIType,
-)
 from invokeai.app.invocations.ip_adapter import IPAdapterModelField
 from invokeai.app.invocations.model import LoRAModelField, MainModelField, VAEModelField
+from invokeai.app.invocations.primitives import ImageField
 from invokeai.app.invocations.t2i_adapter import T2IAdapterField
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.shared.fields import FieldDescriptions
 
 from ...version import __version__
 
@@ -33,7 +31,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 +112,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/model.py

@@ -1,24 +1,31 @@
 import copy
 from typing import List, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.shared.fields import FieldDescriptions
 from invokeai.app.shared.models import FreeUConfig
 
-from ...backend.model_manager import SubModelType
+from ...backend.model_management import BaseModelType, ModelType, SubModelType
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
     invocation,
     invocation_output,
 )
 
 
 class ModelInfo(BaseModel):
-    key: str = Field(description="Key of model as returned by ModelRecordServiceBase.get_model()")
-    submodel_type: Optional[SubModelType] = Field(default=None, description="Info to load submodel")
+    model_name: str = Field(description="Info to load submodel")
+    base_model: BaseModelType = Field(description="Base model")
+    model_type: ModelType = Field(description="Info to load submodel")
+    submodel: Optional[SubModelType] = Field(default=None, description="Info to load submodel")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class LoraInfo(ModelInfo):
@@ -48,7 +55,7 @@ class VaeField(BaseModel):
 
 @invocation_output("unet_output")
 class UNetOutput(BaseInvocationOutput):
-    """Base class for invocations that output a UNet field."""
+    """Base class for invocations that output a UNet field"""
 
     unet: UNetField = OutputField(description=FieldDescriptions.unet, title="UNet")
 
@@ -77,13 +84,20 @@ class ModelLoaderOutput(UNetOutput, CLIPOutput, VAEOutput):
 class MainModelField(BaseModel):
     """Main model field"""
 
-    key: str = Field(description="Model key")
+    model_name: str = Field(description="Name of the model")
+    base_model: BaseModelType = Field(description="Base model")
+    model_type: ModelType = Field(description="Model Type")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class LoRAModelField(BaseModel):
     """LoRA model field"""
 
-    key: str = Field(description="LoRA model key")
+    model_name: str = Field(description="Name of the LoRA model")
+    base_model: BaseModelType = Field(description="Base model")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 @invocation(
@@ -91,7 +105,7 @@ class LoRAModelField(BaseModel):
     title="Main Model",
     tags=["model"],
     category="model",
-    version="1.0.1",
+    version="1.0.0",
 )
 class MainModelLoaderInvocation(BaseInvocation):
     """Loads a main model, outputting its submodels."""
@@ -100,40 +114,85 @@ class MainModelLoaderInvocation(BaseInvocation):
     # TODO: precision?
 
     def invoke(self, context: InvocationContext) -> ModelLoaderOutput:
-        key = self.model.key
+        base_model = self.model.base_model
+        model_name = self.model.model_name
+        model_type = ModelType.Main
 
         # TODO: not found exceptions
-        if not context.models.exists(key):
-            raise Exception(f"Unknown model {key}")
+        if not context.services.model_manager.model_exists(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+        ):
+            raise Exception(f"Unknown {base_model} {model_type} model: {model_name}")
+
+        """
+        if not context.services.model_manager.model_exists(
+            model_name=self.model_name,
+            model_type=SDModelType.Diffusers,
+            submodel=SDModelType.Tokenizer,
+        ):
+            raise Exception(
+                f"Failed to find tokenizer submodel in {self.model_name}! Check if model corrupted"
+            )
+
+        if not context.services.model_manager.model_exists(
+            model_name=self.model_name,
+            model_type=SDModelType.Diffusers,
+            submodel=SDModelType.TextEncoder,
+        ):
+            raise Exception(
+                f"Failed to find text_encoder submodel in {self.model_name}! Check if model corrupted"
+            )
+
+        if not context.services.model_manager.model_exists(
+            model_name=self.model_name,
+            model_type=SDModelType.Diffusers,
+            submodel=SDModelType.UNet,
+        ):
+            raise Exception(
+                f"Failed to find unet submodel from {self.model_name}! Check if model corrupted"
+            )
+        """
 
         return ModelLoaderOutput(
             unet=UNetField(
                 unet=ModelInfo(
-                    key=key,
-                    submodel_type=SubModelType.UNet,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.UNet,
                 ),
                 scheduler=ModelInfo(
-                    key=key,
-                    submodel_type=SubModelType.Scheduler,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Scheduler,
                 ),
                 loras=[],
             ),
             clip=ClipField(
                 tokenizer=ModelInfo(
-                    key=key,
-                    submodel_type=SubModelType.Tokenizer,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Tokenizer,
                 ),
                 text_encoder=ModelInfo(
-                    key=key,
-                    submodel_type=SubModelType.TextEncoder,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.TextEncoder,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             vae=VaeField(
                 vae=ModelInfo(
-                    key=key,
-                    submodel_type=SubModelType.Vae,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Vae,
                 ),
             ),
         )
@@ -147,7 +206,7 @@ class LoraLoaderOutput(BaseInvocationOutput):
     clip: Optional[ClipField] = OutputField(default=None, description=FieldDescriptions.clip, title="CLIP")
 
 
-@invocation("lora_loader", title="LoRA", tags=["model"], category="model", version="1.0.1")
+@invocation("lora_loader", title="LoRA", tags=["model"], category="model", version="1.0.0")
 class LoraLoaderInvocation(BaseInvocation):
     """Apply selected lora to unet and text_encoder."""
 
@@ -170,16 +229,21 @@ class LoraLoaderInvocation(BaseInvocation):
         if self.lora is None:
             raise Exception("No LoRA provided")
 
-        lora_key = self.lora.key
+        base_model = self.lora.base_model
+        lora_name = self.lora.model_name
 
-        if not context.models.exists(lora_key):
-            raise Exception(f"Unkown lora: {lora_key}!")
+        if not context.services.model_manager.model_exists(
+            base_model=base_model,
+            model_name=lora_name,
+            model_type=ModelType.Lora,
+        ):
+            raise Exception(f"Unkown lora name: {lora_name}!")
 
-        if self.unet is not None and any(lora.key == lora_key for lora in self.unet.loras):
-            raise Exception(f'Lora "{lora_key}" already applied to unet')
+        if self.unet is not None and any(lora.model_name == lora_name for lora in self.unet.loras):
+            raise Exception(f'Lora "{lora_name}" already applied to unet')
 
-        if self.clip is not None and any(lora.key == lora_key for lora in self.clip.loras):
-            raise Exception(f'Lora "{lora_key}" already applied to clip')
+        if self.clip is not None and any(lora.model_name == lora_name for lora in self.clip.loras):
+            raise Exception(f'Lora "{lora_name}" already applied to clip')
 
         output = LoraLoaderOutput()
 
@@ -187,8 +251,10 @@ class LoraLoaderInvocation(BaseInvocation):
             output.unet = copy.deepcopy(self.unet)
             output.unet.loras.append(
                 LoraInfo(
-                    key=lora_key,
-                    submodel_type=None,
+                    base_model=base_model,
+                    model_name=lora_name,
+                    model_type=ModelType.Lora,
+                    submodel=None,
                     weight=self.weight,
                 )
             )
@@ -197,8 +263,10 @@ class LoraLoaderInvocation(BaseInvocation):
             output.clip = copy.deepcopy(self.clip)
             output.clip.loras.append(
                 LoraInfo(
-                    key=lora_key,
-                    submodel_type=None,
+                    base_model=base_model,
+                    model_name=lora_name,
+                    model_type=ModelType.Lora,
+                    submodel=None,
                     weight=self.weight,
                 )
             )
@@ -220,7 +288,7 @@ class SDXLLoraLoaderOutput(BaseInvocationOutput):
     title="SDXL LoRA",
     tags=["lora", "model"],
     category="model",
-    version="1.0.1",
+    version="1.0.0",
 )
 class SDXLLoraLoaderInvocation(BaseInvocation):
     """Apply selected lora to unet and text_encoder."""
@@ -250,19 +318,24 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
         if self.lora is None:
             raise Exception("No LoRA provided")
 
-        lora_key = self.lora.key
+        base_model = self.lora.base_model
+        lora_name = self.lora.model_name
 
-        if not context.models.exists(lora_key):
-            raise Exception(f"Unknown lora: {lora_key}!")
+        if not context.services.model_manager.model_exists(
+            base_model=base_model,
+            model_name=lora_name,
+            model_type=ModelType.Lora,
+        ):
+            raise Exception(f"Unknown lora name: {lora_name}!")
 
-        if self.unet is not None and any(lora.key == lora_key for lora in self.unet.loras):
-            raise Exception(f'Lora "{lora_key}" already applied to unet')
+        if self.unet is not None and any(lora.model_name == lora_name for lora in self.unet.loras):
+            raise Exception(f'Lora "{lora_name}" already applied to unet')
 
-        if self.clip is not None and any(lora.key == lora_key for lora in self.clip.loras):
-            raise Exception(f'Lora "{lora_key}" already applied to clip')
+        if self.clip is not None and any(lora.model_name == lora_name for lora in self.clip.loras):
+            raise Exception(f'Lora "{lora_name}" already applied to clip')
 
-        if self.clip2 is not None and any(lora.key == lora_key for lora in self.clip2.loras):
-            raise Exception(f'Lora "{lora_key}" already applied to clip2')
+        if self.clip2 is not None and any(lora.model_name == lora_name for lora in self.clip2.loras):
+            raise Exception(f'Lora "{lora_name}" already applied to clip2')
 
         output = SDXLLoraLoaderOutput()
 
@@ -270,8 +343,10 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
             output.unet = copy.deepcopy(self.unet)
             output.unet.loras.append(
                 LoraInfo(
-                    key=lora_key,
-                    submodel_type=None,
+                    base_model=base_model,
+                    model_name=lora_name,
+                    model_type=ModelType.Lora,
+                    submodel=None,
                     weight=self.weight,
                 )
             )
@@ -280,8 +355,10 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
             output.clip = copy.deepcopy(self.clip)
             output.clip.loras.append(
                 LoraInfo(
-                    key=lora_key,
-                    submodel_type=None,
+                    base_model=base_model,
+                    model_name=lora_name,
+                    model_type=ModelType.Lora,
+                    submodel=None,
                     weight=self.weight,
                 )
             )
@@ -290,8 +367,10 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
             output.clip2 = copy.deepcopy(self.clip2)
             output.clip2.loras.append(
                 LoraInfo(
-                    key=lora_key,
-                    submodel_type=None,
+                    base_model=base_model,
+                    model_name=lora_name,
+                    model_type=ModelType.Lora,
+                    submodel=None,
                     weight=self.weight,
                 )
             )
@@ -302,10 +381,13 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
 class VAEModelField(BaseModel):
     """Vae model field"""
 
-    key: str = Field(description="Model's key")
+    model_name: str = Field(description="Name of the model")
+    base_model: BaseModelType = Field(description="Base model")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
-@invocation("vae_loader", title="VAE", tags=["vae", "model"], category="model", version="1.0.1")
+@invocation("vae_loader", title="VAE", tags=["vae", "model"], category="model", version="1.0.0")
 class VaeLoaderInvocation(BaseInvocation):
     """Loads a VAE model, outputting a VaeLoaderOutput"""
 
@@ -316,12 +398,25 @@ class VaeLoaderInvocation(BaseInvocation):
     )
 
     def invoke(self, context: InvocationContext) -> VAEOutput:
-        key = self.vae_model.key
+        base_model = self.vae_model.base_model
+        model_name = self.vae_model.model_name
+        model_type = ModelType.Vae
 
-        if not context.models.exists(key):
-            raise Exception(f"Unkown vae: {key}!")
-
-        return VAEOutput(vae=VaeField(vae=ModelInfo(key=key)))
+        if not context.services.model_manager.model_exists(
+            base_model=base_model,
+            model_name=model_name,
+            model_type=model_type,
+        ):
+            raise Exception(f"Unkown vae name: {model_name}!")
+        return VAEOutput(
+            vae=VaeField(
+                vae=ModelInfo(
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                )
+            )
+        )
 
 
 @invocation_output("seamless_output")

invokeai/app/invocations/noise.py

@@ -4,15 +4,17 @@
 import torch
 from pydantic import field_validator
 
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
-from invokeai.app.invocations.fields import FieldDescriptions, InputField, LatentsField, OutputField
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.invocations.latent import LatentsField
+from invokeai.app.shared.fields import FieldDescriptions
 from invokeai.app.util.misc import SEED_MAX
 
 from ...backend.util.devices import choose_torch_device, torch_dtype
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    InputField,
+    InvocationContext,
+    OutputField,
     invocation,
     invocation_output,
 )
@@ -67,13 +69,13 @@ class NoiseOutput(BaseInvocationOutput):
     width: int = OutputField(description=FieldDescriptions.width)
     height: int = OutputField(description=FieldDescriptions.height)
 
-    @classmethod
-    def build(cls, latents_name: str, latents: torch.Tensor, seed: int) -> "NoiseOutput":
-        return cls(
-            noise=LatentsField(latents_name=latents_name, seed=seed),
-            width=latents.size()[3] * LATENT_SCALE_FACTOR,
-            height=latents.size()[2] * LATENT_SCALE_FACTOR,
-        )
+
+def build_noise_output(latents_name: str, latents: torch.Tensor, seed: int):
+    return NoiseOutput(
+        noise=LatentsField(latents_name=latents_name, seed=seed),
+        width=latents.size()[3] * 8,
+        height=latents.size()[2] * 8,
+    )
 
 
 @invocation(
@@ -94,13 +96,13 @@ class NoiseInvocation(BaseInvocation):
     )
     width: int = InputField(
         default=512,
-        multiple_of=LATENT_SCALE_FACTOR,
+        multiple_of=8,
         gt=0,
         description=FieldDescriptions.width,
     )
     height: int = InputField(
         default=512,
-        multiple_of=LATENT_SCALE_FACTOR,
+        multiple_of=8,
         gt=0,
         description=FieldDescriptions.height,
     )
@@ -122,5 +124,6 @@ class NoiseInvocation(BaseInvocation):
             seed=self.seed,
             use_cpu=self.use_cpu,
         )
-        name = context.tensors.save(tensor=noise)
-        return NoiseOutput.build(latents_name=name, latents=noise, seed=self.seed)
+        name = f"{context.graph_execution_state_id}__{self.id}"
+        context.services.latents.save(name, noise)
+        return build_noise_output(latents_name=name, latents=noise, seed=self.seed)

invokeai/app/invocations/onnx.py

@@ -0,0 +1,508 @@
+# Copyright (c) 2023 Borisov Sergey (https://github.com/StAlKeR7779)
+
+import inspect
+
+# from contextlib import ExitStack
+from typing import List, Literal, Union
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+from tqdm import tqdm
+
+from invokeai.app.invocations.primitives import ConditioningField, ConditioningOutput, ImageField, ImageOutput
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
+from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.util.step_callback import stable_diffusion_step_callback
+from invokeai.backend import BaseModelType, ModelType, SubModelType
+
+from ...backend.model_management import ONNXModelPatcher
+from ...backend.stable_diffusion import PipelineIntermediateState
+from ...backend.util import choose_torch_device
+from ..util.ti_utils import extract_ti_triggers_from_prompt
+from .baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    UIComponent,
+    UIType,
+    WithMetadata,
+    invocation,
+    invocation_output,
+)
+from .controlnet_image_processors import ControlField
+from .latent import SAMPLER_NAME_VALUES, LatentsField, LatentsOutput, build_latents_output, get_scheduler
+from .model import ClipField, ModelInfo, UNetField, VaeField
+
+ORT_TO_NP_TYPE = {
+    "tensor(bool)": np.bool_,
+    "tensor(int8)": np.int8,
+    "tensor(uint8)": np.uint8,
+    "tensor(int16)": np.int16,
+    "tensor(uint16)": np.uint16,
+    "tensor(int32)": np.int32,
+    "tensor(uint32)": np.uint32,
+    "tensor(int64)": np.int64,
+    "tensor(uint64)": np.uint64,
+    "tensor(float16)": np.float16,
+    "tensor(float)": np.float32,
+    "tensor(double)": np.float64,
+}
+
+PRECISION_VALUES = Literal[tuple(ORT_TO_NP_TYPE.keys())]
+
+
+@invocation("prompt_onnx", title="ONNX Prompt (Raw)", tags=["prompt", "onnx"], category="conditioning", version="1.0.0")
+class ONNXPromptInvocation(BaseInvocation):
+    prompt: str = InputField(default="", description=FieldDescriptions.raw_prompt, ui_component=UIComponent.Textarea)
+    clip: ClipField = InputField(description=FieldDescriptions.clip, input=Input.Connection)
+
+    def invoke(self, context: InvocationContext) -> ConditioningOutput:
+        tokenizer_info = context.services.model_manager.get_model(
+            **self.clip.tokenizer.model_dump(),
+        )
+        text_encoder_info = context.services.model_manager.get_model(
+            **self.clip.text_encoder.model_dump(),
+        )
+        with tokenizer_info as orig_tokenizer, text_encoder_info as text_encoder:  # , ExitStack() as stack:
+            loras = [
+                (
+                    context.services.model_manager.get_model(**lora.model_dump(exclude={"weight"})).context.model,
+                    lora.weight,
+                )
+                for lora in self.clip.loras
+            ]
+
+            ti_list = []
+            for trigger in extract_ti_triggers_from_prompt(self.prompt):
+                name = trigger[1:-1]
+                try:
+                    ti_list.append(
+                        (
+                            name,
+                            context.services.model_manager.get_model(
+                                model_name=name,
+                                base_model=self.clip.text_encoder.base_model,
+                                model_type=ModelType.TextualInversion,
+                            ).context.model,
+                        )
+                    )
+                except Exception:
+                    # print(e)
+                    # import traceback
+                    # print(traceback.format_exc())
+                    print(f'Warn: trigger: "{trigger}" not found')
+            if loras or ti_list:
+                text_encoder.release_session()
+            with (
+                ONNXModelPatcher.apply_lora_text_encoder(text_encoder, loras),
+                ONNXModelPatcher.apply_ti(orig_tokenizer, text_encoder, ti_list) as (tokenizer, ti_manager),
+            ):
+                text_encoder.create_session()
+
+                # copy from
+                # https://github.com/huggingface/diffusers/blob/3ebbaf7c96801271f9e6c21400033b6aa5ffcf29/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L153
+                text_inputs = tokenizer(
+                    self.prompt,
+                    padding="max_length",
+                    max_length=tokenizer.model_max_length,
+                    truncation=True,
+                    return_tensors="np",
+                )
+                text_input_ids = text_inputs.input_ids
+                """
+                untruncated_ids = tokenizer(prompt, padding="max_length", return_tensors="np").input_ids
+
+                if not np.array_equal(text_input_ids, untruncated_ids):
+                    removed_text = self.tokenizer.batch_decode(
+                        untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+                    )
+                    logger.warning(
+                        "The following part of your input was truncated because CLIP can only handle sequences up to"
+                        f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                    )
+                """
+
+                prompt_embeds = text_encoder(input_ids=text_input_ids.astype(np.int32))[0]
+
+        conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning"
+
+        # TODO: hacky but works ;D maybe rename latents somehow?
+        context.services.latents.save(conditioning_name, (prompt_embeds, None))
+
+        return ConditioningOutput(
+            conditioning=ConditioningField(
+                conditioning_name=conditioning_name,
+            ),
+        )
+
+
+# Text to image
+@invocation(
+    "t2l_onnx",
+    title="ONNX Text to Latents",
+    tags=["latents", "inference", "txt2img", "onnx"],
+    category="latents",
+    version="1.0.0",
+)
+class ONNXTextToLatentsInvocation(BaseInvocation):
+    """Generates latents from conditionings."""
+
+    positive_conditioning: ConditioningField = InputField(
+        description=FieldDescriptions.positive_cond,
+        input=Input.Connection,
+    )
+    negative_conditioning: ConditioningField = InputField(
+        description=FieldDescriptions.negative_cond,
+        input=Input.Connection,
+    )
+    noise: LatentsField = InputField(
+        description=FieldDescriptions.noise,
+        input=Input.Connection,
+    )
+    steps: int = InputField(default=10, gt=0, description=FieldDescriptions.steps)
+    cfg_scale: Union[float, List[float]] = InputField(
+        default=7.5,
+        ge=1,
+        description=FieldDescriptions.cfg_scale,
+    )
+    scheduler: SAMPLER_NAME_VALUES = InputField(
+        default="euler", description=FieldDescriptions.scheduler, input=Input.Direct, ui_type=UIType.Scheduler
+    )
+    precision: PRECISION_VALUES = InputField(default="tensor(float16)", description=FieldDescriptions.precision)
+    unet: UNetField = InputField(
+        description=FieldDescriptions.unet,
+        input=Input.Connection,
+    )
+    control: Union[ControlField, list[ControlField]] = InputField(
+        default=None,
+        description=FieldDescriptions.control,
+    )
+    # seamless:   bool = InputField(default=False, description="Whether or not to generate an image that can tile without seams", )
+    # seamless_axes: str = InputField(default="", description="The axes to tile the image on, 'x' and/or 'y'")
+
+    @field_validator("cfg_scale")
+    def ge_one(cls, v):
+        """validate that all cfg_scale values are >= 1"""
+        if isinstance(v, list):
+            for i in v:
+                if i < 1:
+                    raise ValueError("cfg_scale must be greater than 1")
+        else:
+            if v < 1:
+                raise ValueError("cfg_scale must be greater than 1")
+        return v
+
+    # based on
+    # https://github.com/huggingface/diffusers/blob/3ebbaf7c96801271f9e6c21400033b6aa5ffcf29/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L375
+    def invoke(self, context: InvocationContext) -> LatentsOutput:
+        c, _ = context.services.latents.get(self.positive_conditioning.conditioning_name)
+        uc, _ = context.services.latents.get(self.negative_conditioning.conditioning_name)
+        graph_execution_state = context.services.graph_execution_manager.get(context.graph_execution_state_id)
+        source_node_id = graph_execution_state.prepared_source_mapping[self.id]
+        if isinstance(c, torch.Tensor):
+            c = c.cpu().numpy()
+        if isinstance(uc, torch.Tensor):
+            uc = uc.cpu().numpy()
+        device = torch.device(choose_torch_device())
+        prompt_embeds = np.concatenate([uc, c])
+
+        latents = context.services.latents.get(self.noise.latents_name)
+        if isinstance(latents, torch.Tensor):
+            latents = latents.cpu().numpy()
+
+        # TODO: better execution device handling
+        latents = latents.astype(ORT_TO_NP_TYPE[self.precision])
+
+        # get the initial random noise unless the user supplied it
+        do_classifier_free_guidance = True
+        # latents_dtype = prompt_embeds.dtype
+        # latents_shape = (batch_size * num_images_per_prompt, 4, height // 8, width // 8)
+        # if latents.shape != latents_shape:
+        #    raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")
+
+        scheduler = get_scheduler(
+            context=context,
+            scheduler_info=self.unet.scheduler,
+            scheduler_name=self.scheduler,
+            seed=0,  # TODO: refactor this node
+        )
+
+        def torch2numpy(latent: torch.Tensor):
+            return latent.cpu().numpy()
+
+        def numpy2torch(latent, device):
+            return torch.from_numpy(latent).to(device)
+
+        def dispatch_progress(
+            self, context: InvocationContext, source_node_id: str, intermediate_state: PipelineIntermediateState
+        ) -> None:
+            stable_diffusion_step_callback(
+                context=context,
+                intermediate_state=intermediate_state,
+                node=self.model_dump(),
+                source_node_id=source_node_id,
+            )
+
+        scheduler.set_timesteps(self.steps)
+        latents = latents * np.float64(scheduler.init_noise_sigma)
+
+        extra_step_kwargs = {}
+        if "eta" in set(inspect.signature(scheduler.step).parameters.keys()):
+            extra_step_kwargs.update(
+                eta=0.0,
+            )
+
+        unet_info = context.services.model_manager.get_model(**self.unet.unet.model_dump())
+
+        with unet_info as unet:  # , ExitStack() as stack:
+            # loras = [(stack.enter_context(context.services.model_manager.get_model(**lora.dict(exclude={"weight"}))), lora.weight) for lora in self.unet.loras]
+            loras = [
+                (
+                    context.services.model_manager.get_model(**lora.model_dump(exclude={"weight"})).context.model,
+                    lora.weight,
+                )
+                for lora in self.unet.loras
+            ]
+
+            if loras:
+                unet.release_session()
+            with ONNXModelPatcher.apply_lora_unet(unet, loras):
+                # TODO:
+                _, _, h, w = latents.shape
+                unet.create_session(h, w)
+
+                timestep_dtype = next(
+                    (input.type for input in unet.session.get_inputs() if input.name == "timestep"), "tensor(float16)"
+                )
+                timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype]
+                for i in tqdm(range(len(scheduler.timesteps))):
+                    t = scheduler.timesteps[i]
+                    # expand the latents if we are doing classifier free guidance
+                    latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents
+                    latent_model_input = scheduler.scale_model_input(numpy2torch(latent_model_input, device), t)
+                    latent_model_input = latent_model_input.cpu().numpy()
+
+                    # predict the noise residual
+                    timestep = np.array([t], dtype=timestep_dtype)
+                    noise_pred = unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds)
+                    noise_pred = noise_pred[0]
+
+                    # perform guidance
+                    if do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
+                        noise_pred = noise_pred_uncond + self.cfg_scale * (noise_pred_text - noise_pred_uncond)
+
+                    # compute the previous noisy sample x_t -> x_t-1
+                    scheduler_output = scheduler.step(
+                        numpy2torch(noise_pred, device), t, numpy2torch(latents, device), **extra_step_kwargs
+                    )
+                    latents = torch2numpy(scheduler_output.prev_sample)
+
+                    state = PipelineIntermediateState(
+                        run_id="test", step=i, timestep=timestep, latents=scheduler_output.prev_sample
+                    )
+                    dispatch_progress(self, context=context, source_node_id=source_node_id, intermediate_state=state)
+
+                    # call the callback, if provided
+                    # if callback is not None and i % callback_steps == 0:
+                    #    callback(i, t, latents)
+
+        torch.cuda.empty_cache()
+
+        name = f"{context.graph_execution_state_id}__{self.id}"
+        context.services.latents.save(name, latents)
+        return build_latents_output(latents_name=name, latents=torch.from_numpy(latents))
+
+
+# Latent to image
+@invocation(
+    "l2i_onnx",
+    title="ONNX Latents to Image",
+    tags=["latents", "image", "vae", "onnx"],
+    category="image",
+    version="1.2.0",
+)
+class ONNXLatentsToImageInvocation(BaseInvocation, WithMetadata):
+    """Generates an image from latents."""
+
+    latents: LatentsField = InputField(
+        description=FieldDescriptions.denoised_latents,
+        input=Input.Connection,
+    )
+    vae: VaeField = InputField(
+        description=FieldDescriptions.vae,
+        input=Input.Connection,
+    )
+    # tiled: bool = InputField(default=False, description="Decode latents by overlaping tiles(less memory consumption)")
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        latents = context.services.latents.get(self.latents.latents_name)
+
+        if self.vae.vae.submodel != SubModelType.VaeDecoder:
+            raise Exception(f"Expected vae_decoder, found: {self.vae.vae.model_type}")
+
+        vae_info = context.services.model_manager.get_model(
+            **self.vae.vae.model_dump(),
+        )
+
+        # clear memory as vae decode can request a lot
+        torch.cuda.empty_cache()
+
+        with vae_info as vae:
+            vae.create_session()
+
+            # copied from
+            # https://github.com/huggingface/diffusers/blob/3ebbaf7c96801271f9e6c21400033b6aa5ffcf29/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L427
+            latents = 1 / 0.18215 * latents
+            # image = self.vae_decoder(latent_sample=latents)[0]
+            # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
+            image = np.concatenate([vae(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])])
+
+            image = np.clip(image / 2 + 0.5, 0, 1)
+            image = image.transpose((0, 2, 3, 1))
+            image = VaeImageProcessor.numpy_to_pil(image)[0]
+
+        torch.cuda.empty_cache()
+
+        image_dto = context.services.images.create(
+            image=image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
+
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )
+
+
+@invocation_output("model_loader_output_onnx")
+class ONNXModelLoaderOutput(BaseInvocationOutput):
+    """Model loader output"""
+
+    unet: UNetField = OutputField(default=None, description=FieldDescriptions.unet, title="UNet")
+    clip: ClipField = OutputField(default=None, description=FieldDescriptions.clip, title="CLIP")
+    vae_decoder: VaeField = OutputField(default=None, description=FieldDescriptions.vae, title="VAE Decoder")
+    vae_encoder: VaeField = OutputField(default=None, description=FieldDescriptions.vae, title="VAE Encoder")
+
+
+class OnnxModelField(BaseModel):
+    """Onnx model field"""
+
+    model_name: str = Field(description="Name of the model")
+    base_model: BaseModelType = Field(description="Base model")
+    model_type: ModelType = Field(description="Model Type")
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+@invocation("onnx_model_loader", title="ONNX Main Model", tags=["onnx", "model"], category="model", version="1.0.0")
+class OnnxModelLoaderInvocation(BaseInvocation):
+    """Loads a main model, outputting its submodels."""
+
+    model: OnnxModelField = InputField(
+        description=FieldDescriptions.onnx_main_model, input=Input.Direct, ui_type=UIType.ONNXModel
+    )
+
+    def invoke(self, context: InvocationContext) -> ONNXModelLoaderOutput:
+        base_model = self.model.base_model
+        model_name = self.model.model_name
+        model_type = ModelType.ONNX
+
+        # TODO: not found exceptions
+        if not context.services.model_manager.model_exists(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+        ):
+            raise Exception(f"Unknown {base_model} {model_type} model: {model_name}")
+
+        """
+        if not context.services.model_manager.model_exists(
+            model_name=self.model_name,
+            model_type=SDModelType.Diffusers,
+            submodel=SDModelType.Tokenizer,
+        ):
+            raise Exception(
+                f"Failed to find tokenizer submodel in {self.model_name}! Check if model corrupted"
+            )
+
+        if not context.services.model_manager.model_exists(
+            model_name=self.model_name,
+            model_type=SDModelType.Diffusers,
+            submodel=SDModelType.TextEncoder,
+        ):
+            raise Exception(
+                f"Failed to find text_encoder submodel in {self.model_name}! Check if model corrupted"
+            )
+
+        if not context.services.model_manager.model_exists(
+            model_name=self.model_name,
+            model_type=SDModelType.Diffusers,
+            submodel=SDModelType.UNet,
+        ):
+            raise Exception(
+                f"Failed to find unet submodel from {self.model_name}! Check if model corrupted"
+            )
+        """
+
+        return ONNXModelLoaderOutput(
+            unet=UNetField(
+                unet=ModelInfo(
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.UNet,
+                ),
+                scheduler=ModelInfo(
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Scheduler,
+                ),
+                loras=[],
+            ),
+            clip=ClipField(
+                tokenizer=ModelInfo(
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Tokenizer,
+                ),
+                text_encoder=ModelInfo(
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.TextEncoder,
+                ),
+                loras=[],
+                skipped_layers=0,
+            ),
+            vae_decoder=VaeField(
+                vae=ModelInfo(
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.VaeDecoder,
+                ),
+            ),
+            vae_encoder=VaeField(
+                vae=ModelInfo(
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.VaeEncoder,
+                ),
+            ),
+        )

invokeai/app/invocations/param_easing.py

@@ -40,10 +40,8 @@ from easing_functions import (
 from matplotlib.ticker import MaxNLocator
 
 from invokeai.app.invocations.primitives import FloatCollectionOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
 
-from .baseinvocation import BaseInvocation, invocation
-from .fields import InputField
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, invocation
 
 
 @invocation(
@@ -111,7 +109,7 @@ EASING_FUNCTION_KEYS = Literal[tuple(EASING_FUNCTIONS_MAP.keys())]
     title="Step Param Easing",
     tags=["step", "easing"],
     category="step",
-    version="1.0.1",
+    version="1.0.0",
 )
 class StepParamEasingInvocation(BaseInvocation):
     """Experimental per-step parameter easing for denoising steps"""
@@ -150,19 +148,19 @@ class StepParamEasingInvocation(BaseInvocation):
         postlist = list(num_poststeps * [self.post_end_value])
 
         if log_diagnostics:
-            context.logger.debug("start_step: " + str(start_step))
-            context.logger.debug("end_step: " + str(end_step))
-            context.logger.debug("num_easing_steps: " + str(num_easing_steps))
-            context.logger.debug("num_presteps: " + str(num_presteps))
-            context.logger.debug("num_poststeps: " + str(num_poststeps))
-            context.logger.debug("prelist size: " + str(len(prelist)))
-            context.logger.debug("postlist size: " + str(len(postlist)))
-            context.logger.debug("prelist: " + str(prelist))
-            context.logger.debug("postlist: " + str(postlist))
+            context.services.logger.debug("start_step: " + str(start_step))
+            context.services.logger.debug("end_step: " + str(end_step))
+            context.services.logger.debug("num_easing_steps: " + str(num_easing_steps))
+            context.services.logger.debug("num_presteps: " + str(num_presteps))
+            context.services.logger.debug("num_poststeps: " + str(num_poststeps))
+            context.services.logger.debug("prelist size: " + str(len(prelist)))
+            context.services.logger.debug("postlist size: " + str(len(postlist)))
+            context.services.logger.debug("prelist: " + str(prelist))
+            context.services.logger.debug("postlist: " + str(postlist))
 
         easing_class = EASING_FUNCTIONS_MAP[self.easing]
         if log_diagnostics:
-            context.logger.debug("easing class: " + str(easing_class))
+            context.services.logger.debug("easing class: " + str(easing_class))
         easing_list = []
         if self.mirror:  # "expected" mirroring
             # if number of steps is even, squeeze duration down to (number_of_steps)/2
@@ -173,7 +171,7 @@ class StepParamEasingInvocation(BaseInvocation):
 
             base_easing_duration = int(np.ceil(num_easing_steps / 2.0))
             if log_diagnostics:
-                context.logger.debug("base easing duration: " + str(base_easing_duration))
+                context.services.logger.debug("base easing duration: " + str(base_easing_duration))
             even_num_steps = num_easing_steps % 2 == 0  # even number of steps
             easing_function = easing_class(
                 start=self.start_value,
@@ -185,14 +183,14 @@ class StepParamEasingInvocation(BaseInvocation):
                 easing_val = easing_function.ease(step_index)
                 base_easing_vals.append(easing_val)
                 if log_diagnostics:
-                    context.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(easing_val))
+                    context.services.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(easing_val))
             if even_num_steps:
                 mirror_easing_vals = list(reversed(base_easing_vals))
             else:
                 mirror_easing_vals = list(reversed(base_easing_vals[0:-1]))
             if log_diagnostics:
-                context.logger.debug("base easing vals: " + str(base_easing_vals))
-                context.logger.debug("mirror easing vals: " + str(mirror_easing_vals))
+                context.services.logger.debug("base easing vals: " + str(base_easing_vals))
+                context.services.logger.debug("mirror easing vals: " + str(mirror_easing_vals))
             easing_list = base_easing_vals + mirror_easing_vals
 
         # FIXME: add alt_mirror option (alternative to default or mirror), or remove entirely
@@ -227,12 +225,12 @@ class StepParamEasingInvocation(BaseInvocation):
                 step_val = easing_function.ease(step_index)
                 easing_list.append(step_val)
                 if log_diagnostics:
-                    context.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(step_val))
+                    context.services.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(step_val))
 
         if log_diagnostics:
-            context.logger.debug("prelist size: " + str(len(prelist)))
-            context.logger.debug("easing_list size: " + str(len(easing_list)))
-            context.logger.debug("postlist size: " + str(len(postlist)))
+            context.services.logger.debug("prelist size: " + str(len(prelist)))
+            context.services.logger.debug("easing_list size: " + str(len(easing_list)))
+            context.services.logger.debug("postlist size: " + str(len(postlist)))
 
         param_list = prelist + easing_list + postlist
 

invokeai/app/invocations/primitives.py

@@ -1,28 +1,20 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
+from pydantic import BaseModel, Field
 
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
-from invokeai.app.invocations.fields import (
-    ColorField,
-    ConditioningField,
-    DenoiseMaskField,
-    FieldDescriptions,
-    ImageField,
-    Input,
-    InputField,
-    LatentsField,
-    OutputField,
-    UIComponent,
-)
-from invokeai.app.services.images.images_common import ImageDTO
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.shared.fields import FieldDescriptions
 
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    UIComponent,
     invocation,
     invocation_output,
 )
@@ -229,6 +221,18 @@ class StringCollectionInvocation(BaseInvocation):
 # region Image
 
 
+class ImageField(BaseModel):
+    """An image primitive field"""
+
+    image_name: str = Field(description="The name of the image")
+
+
+class BoardField(BaseModel):
+    """A board primitive field"""
+
+    board_id: str = Field(description="The id of the board")
+
+
 @invocation_output("image_output")
 class ImageOutput(BaseInvocationOutput):
     """Base class for nodes that output a single image"""
@@ -237,14 +241,6 @@ class ImageOutput(BaseInvocationOutput):
     width: int = OutputField(description="The width of the image in pixels")
     height: int = OutputField(description="The height of the image in pixels")
 
-    @classmethod
-    def build(cls, image_dto: ImageDTO) -> "ImageOutput":
-        return cls(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
-
 
 @invocation_output("image_collection_output")
 class ImageCollectionOutput(BaseInvocationOutput):
@@ -255,14 +251,16 @@ class ImageCollectionOutput(BaseInvocationOutput):
     )
 
 
-@invocation("image", title="Image Primitive", tags=["primitives", "image"], category="primitives", version="1.0.1")
-class ImageInvocation(BaseInvocation):
+@invocation("image", title="Image Primitive", tags=["primitives", "image"], category="primitives", version="1.0.0")
+class ImageInvocation(
+    BaseInvocation,
+):
     """An image primitive value"""
 
     image: ImageField = InputField(description="The image to load")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
+        image = context.services.images.get_pil_image(self.image.image_name)
 
         return ImageOutput(
             image=ImageField(image_name=self.image.image_name),
@@ -292,44 +290,42 @@ class ImageCollectionInvocation(BaseInvocation):
 # region DenoiseMask
 
 
+class DenoiseMaskField(BaseModel):
+    """An inpaint mask field"""
+
+    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")
+
+
 @invocation_output("denoise_mask_output")
 class DenoiseMaskOutput(BaseInvocationOutput):
     """Base class for nodes that output a single image"""
 
     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":
-        return cls(
-            denoise_mask=DenoiseMaskField(
-                mask_name=mask_name, masked_latents_name=masked_latents_name, gradient=gradient
-            ),
-        )
-
 
 # endregion
 
 # region Latents
 
 
+class LatentsField(BaseModel):
+    """A latents tensor primitive field"""
+
+    latents_name: str = Field(description="The name of the latents")
+    seed: Optional[int] = Field(default=None, description="Seed used to generate this latents")
+
+
 @invocation_output("latents_output")
 class LatentsOutput(BaseInvocationOutput):
     """Base class for nodes that output a single latents tensor"""
 
-    latents: LatentsField = OutputField(description=FieldDescriptions.latents)
+    latents: LatentsField = OutputField(
+        description=FieldDescriptions.latents,
+    )
     width: int = OutputField(description=FieldDescriptions.width)
     height: int = OutputField(description=FieldDescriptions.height)
 
-    @classmethod
-    def build(cls, latents_name: str, latents: torch.Tensor, seed: Optional[int] = None) -> "LatentsOutput":
-        return cls(
-            latents=LatentsField(latents_name=latents_name, seed=seed),
-            width=latents.size()[3] * LATENT_SCALE_FACTOR,
-            height=latents.size()[2] * LATENT_SCALE_FACTOR,
-        )
-
 
 @invocation_output("latents_collection_output")
 class LatentsCollectionOutput(BaseInvocationOutput):
@@ -341,7 +337,7 @@ class LatentsCollectionOutput(BaseInvocationOutput):
 
 
 @invocation(
-    "latents", title="Latents Primitive", tags=["primitives", "latents"], category="primitives", version="1.0.1"
+    "latents", title="Latents Primitive", tags=["primitives", "latents"], category="primitives", version="1.0.0"
 )
 class LatentsInvocation(BaseInvocation):
     """A latents tensor primitive value"""
@@ -349,9 +345,9 @@ class LatentsInvocation(BaseInvocation):
     latents: LatentsField = InputField(description="The latents tensor", input=Input.Connection)
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.tensors.load(self.latents.latents_name)
+        latents = context.services.latents.get(self.latents.latents_name)
 
-        return LatentsOutput.build(self.latents.latents_name, latents)
+        return build_latents_output(self.latents.latents_name, latents)
 
 
 @invocation(
@@ -372,11 +368,31 @@ class LatentsCollectionInvocation(BaseInvocation):
         return LatentsCollectionOutput(collection=self.collection)
 
 
+def build_latents_output(latents_name: str, latents: torch.Tensor, seed: Optional[int] = None):
+    return LatentsOutput(
+        latents=LatentsField(latents_name=latents_name, seed=seed),
+        width=latents.size()[3] * 8,
+        height=latents.size()[2] * 8,
+    )
+
+
 # endregion
 
 # region Color
 
 
+class ColorField(BaseModel):
+    """A color primitive field"""
+
+    r: int = Field(ge=0, le=255, description="The red component")
+    g: int = Field(ge=0, le=255, description="The green component")
+    b: int = Field(ge=0, le=255, description="The blue component")
+    a: int = Field(ge=0, le=255, description="The alpha component")
+
+    def tuple(self) -> Tuple[int, int, int, int]:
+        return (self.r, self.g, self.b, self.a)
+
+
 @invocation_output("color_output")
 class ColorOutput(BaseInvocationOutput):
     """Base class for nodes that output a single color"""
@@ -408,16 +424,18 @@ class ColorInvocation(BaseInvocation):
 # region Conditioning
 
 
+class ConditioningField(BaseModel):
+    """A conditioning tensor primitive value"""
+
+    conditioning_name: str = Field(description="The name of conditioning tensor")
+
+
 @invocation_output("conditioning_output")
 class ConditioningOutput(BaseInvocationOutput):
     """Base class for nodes that output a single conditioning tensor"""
 
     conditioning: ConditioningField = OutputField(description=FieldDescriptions.cond)
 
-    @classmethod
-    def build(cls, conditioning_name: str) -> "ConditioningOutput":
-        return cls(conditioning=ConditioningField(conditioning_name=conditioning_name))
-
 
 @invocation_output("conditioning_collection_output")
 class ConditioningCollectionOutput(BaseInvocationOutput):

invokeai/app/invocations/prompt.py

@@ -6,10 +6,8 @@ from dynamicprompts.generators import CombinatorialPromptGenerator, RandomPrompt
 from pydantic import field_validator
 
 from invokeai.app.invocations.primitives import StringCollectionOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
 
-from .baseinvocation import BaseInvocation, invocation
-from .fields import InputField, UIComponent
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, UIComponent, invocation
 
 
 @invocation(

invokeai/app/invocations/sdxl.py

@@ -1,10 +1,14 @@
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager import SubModelType
+from invokeai.app.shared.fields import FieldDescriptions
 
+from ...backend.model_management import ModelType, SubModelType
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    UIType,
     invocation,
     invocation_output,
 )
@@ -30,7 +34,7 @@ class SDXLRefinerModelLoaderOutput(BaseInvocationOutput):
     vae: VaeField = OutputField(description=FieldDescriptions.vae, title="VAE")
 
 
-@invocation("sdxl_model_loader", title="SDXL Main Model", tags=["model", "sdxl"], category="model", version="1.0.1")
+@invocation("sdxl_model_loader", title="SDXL Main Model", tags=["model", "sdxl"], category="model", version="1.0.0")
 class SDXLModelLoaderInvocation(BaseInvocation):
     """Loads an sdxl base model, outputting its submodels."""
 
@@ -40,52 +44,72 @@ class SDXLModelLoaderInvocation(BaseInvocation):
     # TODO: precision?
 
     def invoke(self, context: InvocationContext) -> SDXLModelLoaderOutput:
-        model_key = self.model.key
+        base_model = self.model.base_model
+        model_name = self.model.model_name
+        model_type = ModelType.Main
 
         # TODO: not found exceptions
-        if not context.models.exists(model_key):
-            raise Exception(f"Unknown model: {model_key}")
+        if not context.services.model_manager.model_exists(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+        ):
+            raise Exception(f"Unknown {base_model} {model_type} model: {model_name}")
 
         return SDXLModelLoaderOutput(
             unet=UNetField(
                 unet=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.UNet,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.UNet,
                 ),
                 scheduler=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.Scheduler,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Scheduler,
                 ),
                 loras=[],
             ),
             clip=ClipField(
                 tokenizer=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.Tokenizer,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Tokenizer,
                 ),
                 text_encoder=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.TextEncoder,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.TextEncoder,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             clip2=ClipField(
                 tokenizer=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.Tokenizer2,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Tokenizer2,
                 ),
                 text_encoder=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.TextEncoder2,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.TextEncoder2,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             vae=VaeField(
                 vae=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.Vae,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Vae,
                 ),
             ),
         )
@@ -96,7 +120,7 @@ class SDXLModelLoaderInvocation(BaseInvocation):
     title="SDXL Refiner Model",
     tags=["model", "sdxl", "refiner"],
     category="model",
-    version="1.0.1",
+    version="1.0.0",
 )
 class SDXLRefinerModelLoaderInvocation(BaseInvocation):
     """Loads an sdxl refiner model, outputting its submodels."""
@@ -109,40 +133,56 @@ class SDXLRefinerModelLoaderInvocation(BaseInvocation):
     # TODO: precision?
 
     def invoke(self, context: InvocationContext) -> SDXLRefinerModelLoaderOutput:
-        model_key = self.model.key
+        base_model = self.model.base_model
+        model_name = self.model.model_name
+        model_type = ModelType.Main
 
         # TODO: not found exceptions
-        if not context.models.exists(model_key):
-            raise Exception(f"Unknown model: {model_key}")
+        if not context.services.model_manager.model_exists(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+        ):
+            raise Exception(f"Unknown {base_model} {model_type} model: {model_name}")
 
         return SDXLRefinerModelLoaderOutput(
             unet=UNetField(
                 unet=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.UNet,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.UNet,
                 ),
                 scheduler=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.Scheduler,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Scheduler,
                 ),
                 loras=[],
             ),
             clip2=ClipField(
                 tokenizer=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.Tokenizer2,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Tokenizer2,
                 ),
                 text_encoder=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.TextEncoder2,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.TextEncoder2,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             vae=VaeField(
                 vae=ModelInfo(
-                    key=model_key,
-                    submodel_type=SubModelType.Vae,
+                    model_name=model_name,
+                    base_model=base_model,
+                    model_type=model_type,
+                    submodel=SubModelType.Vae,
                 ),
             ),
         )

invokeai/app/invocations/strings.py

@@ -2,15 +2,16 @@
 
 import re
 
-from invokeai.app.services.shared.invocation_context import InvocationContext
-
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    InputField,
+    InvocationContext,
+    OutputField,
+    UIComponent,
     invocation,
     invocation_output,
 )
-from .fields import InputField, OutputField, UIComponent
 from .primitives import StringOutput
 
 

invokeai/app/invocations/t2i_adapter.py

@@ -1,21 +1,29 @@
 from typing import Union
 
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
     invocation,
     invocation_output,
 )
 from invokeai.app.invocations.controlnet_image_processors import CONTROLNET_RESIZE_VALUES
-from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField, OutputField
+from invokeai.app.invocations.primitives import ImageField
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.backend.model_management.models.base import BaseModelType
 
 
 class T2IAdapterModelField(BaseModel):
-    key: str = Field(description="Model record key for the T2I-Adapter model")
+    model_name: str = Field(description="Name of the T2I-Adapter model")
+    base_model: BaseModelType = Field(description="Base model")
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class T2IAdapterField(BaseModel):

invokeai/app/invocations/tiles.py

@@ -8,12 +8,16 @@ from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
     Classification,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    WithMetadata,
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import ImageField, Input, InputField, OutputField, WithBoard, WithMetadata
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.invocations.primitives import ImageField, ImageOutput
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
 from invokeai.backend.tiles.tiles import (
     calc_tiles_even_split,
     calc_tiles_min_overlap,
@@ -232,7 +236,7 @@ BLEND_MODES = Literal["Linear", "Seam"]
     version="1.1.0",
     classification=Classification.Beta,
 )
-class MergeTilesToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
+class MergeTilesToImageInvocation(BaseInvocation, WithMetadata):
     """Merge multiple tile images into a single image."""
 
     # Inputs
@@ -264,7 +268,7 @@ class MergeTilesToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
         # existed in memory at an earlier point in the graph.
         tile_np_images: list[np.ndarray] = []
         for image in images:
-            pil_image = context.images.get_pil(image.image_name)
+            pil_image = context.services.images.get_pil_image(image.image_name)
             pil_image = pil_image.convert("RGB")
             tile_np_images.append(np.array(pil_image))
 
@@ -287,5 +291,18 @@ class MergeTilesToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
         # Convert into a PIL image and save
         pil_image = Image.fromarray(np_image)
 
-        image_dto = context.images.save(image=pil_image)
-        return ImageOutput.build(image_dto)
+        image_dto = context.services.images.create(
+            image=pil_image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )

invokeai/app/invocations/upscale.py

@@ -5,18 +5,16 @@ from typing import Literal
 import cv2
 import numpy as np
 import torch
+from basicsr.archs.rrdbnet_arch import RRDBNet
 from PIL import Image
 from pydantic import ConfigDict
 
-from invokeai.app.invocations.fields import ImageField
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet
+from invokeai.app.invocations.primitives import ImageField, ImageOutput
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
 from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGAN
 from invokeai.backend.util.devices import choose_torch_device
 
-from .baseinvocation import BaseInvocation, invocation
-from .fields import InputField, WithBoard, WithMetadata
+from .baseinvocation import BaseInvocation, InputField, InvocationContext, WithMetadata, invocation
 
 # TODO: Populate this from disk?
 # TODO: Use model manager to load?
@@ -31,8 +29,8 @@ if choose_torch_device() == torch.device("mps"):
     from torch import mps
 
 
-@invocation("esrgan", title="Upscale (RealESRGAN)", tags=["esrgan", "upscale"], category="esrgan", version="1.3.1")
-class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
+@invocation("esrgan", title="Upscale (RealESRGAN)", tags=["esrgan", "upscale"], category="esrgan", version="1.3.0")
+class ESRGANInvocation(BaseInvocation, WithMetadata):
     """Upscales an image using RealESRGAN."""
 
     image: ImageField = InputField(description="The input image")
@@ -44,8 +42,8 @@ class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
     model_config = ConfigDict(protected_namespaces=())
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name)
-        models_path = context.config.get().models_path
+        image = context.services.images.get_pil_image(self.image.image_name)
+        models_path = context.services.configuration.models_path
 
         rrdbnet_model = None
         netscale = None
@@ -89,7 +87,7 @@ class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
             netscale = 2
         else:
             msg = f"Invalid RealESRGAN model: {self.model_name}"
-            context.logger.error(msg)
+            context.services.logger.error(msg)
             raise ValueError(msg)
 
         esrgan_model_path = Path(f"core/upscaling/realesrgan/{self.model_name}")
@@ -112,6 +110,19 @@ class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
         if choose_torch_device() == torch.device("mps"):
             mps.empty_cache()
 
-        image_dto = context.images.save(image=pil_image)
+        image_dto = context.services.images.create(
+            image=pil_image,
+            image_origin=ResourceOrigin.INTERNAL,
+            image_category=ImageCategory.GENERAL,
+            node_id=self.id,
+            session_id=context.graph_execution_state_id,
+            is_intermediate=self.is_intermediate,
+            metadata=self.metadata,
+            workflow=context.workflow,
+        )
 
-        return ImageOutput.build(image_dto)
+        return ImageOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            width=image_dto.width,
+            height=image_dto.height,
+        )

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

@@ -27,11 +27,11 @@ class InvokeAISettings(BaseSettings):
     """Runtime configuration settings in which default values are read from an omegaconf .yaml file."""
 
     initconf: ClassVar[Optional[DictConfig]] = None
-    argparse_groups: ClassVar[Dict[str, Any]] = {}
+    argparse_groups: ClassVar[Dict] = {}
 
     model_config = SettingsConfigDict(env_file_encoding="utf-8", arbitrary_types_allowed=True, case_sensitive=True)
 
-    def parse_args(self, argv: Optional[List[str]] = sys.argv[1:]) -> None:
+    def parse_args(self, argv: Optional[list] = sys.argv[1:]):
         """Call to parse command-line arguments."""
         parser = self.get_parser()
         opt, unknown_opts = parser.parse_known_args(argv)
@@ -68,7 +68,7 @@ class InvokeAISettings(BaseSettings):
         return OmegaConf.to_yaml(conf)
 
     @classmethod
-    def add_parser_arguments(cls, parser: ArgumentParser) -> None:
+    def add_parser_arguments(cls, parser):
         """Dynamically create arguments for a settings parser."""
         if "type" in get_type_hints(cls):
             settings_stanza = get_args(get_type_hints(cls)["type"])[0]
@@ -117,8 +117,7 @@ class InvokeAISettings(BaseSettings):
         """Return the category of a setting."""
         hints = get_type_hints(cls)
         if command_field in hints:
-            result: str = get_args(hints[command_field])[0]
-            return result
+            return get_args(hints[command_field])[0]
         else:
             return "Uncategorized"
 
@@ -156,11 +155,10 @@ class InvokeAISettings(BaseSettings):
             "lora_dir",
             "embedding_dir",
             "controlnet_dir",
-            "conf_path",
         ]
 
     @classmethod
-    def add_field_argument(cls, command_parser, name: str, field, default_override=None) -> None:
+    def add_field_argument(cls, command_parser, name: str, field, default_override=None):
         """Add the argparse arguments for a setting parser."""
         field_type = get_type_hints(cls).get(name)
         default = (

invokeai/app/services/config/config_common.py

@@ -21,7 +21,7 @@ class PagingArgumentParser(argparse.ArgumentParser):
     It also supports reading defaults from an init file.
     """
 
-    def print_help(self, file=None) -> None:
+    def print_help(self, file=None):
         text = self.format_help()
         pydoc.pager(text)
 

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
      ...
@@ -170,7 +173,7 @@ from __future__ import annotations
 
 import os
 from pathlib import Path
-from typing import Any, ClassVar, Dict, List, Literal, Optional
+from typing import Any, ClassVar, Dict, List, Literal, Optional, Union
 
 from omegaconf import DictConfig, OmegaConf
 from pydantic import Field
@@ -182,9 +185,7 @@ from .config_base import InvokeAISettings
 INIT_FILE = Path("invokeai.yaml")
 DB_FILE = Path("invokeai.db")
 LEGACY_INIT_FILE = Path("invokeai.init")
-DEFAULT_RAM_CACHE = 10.0
-DEFAULT_VRAM_CACHE = 0.25
-DEFAULT_CONVERT_CACHE = 20.0
+DEFAULT_MAX_VRAM = 0.5
 
 
 class Categories(object):
@@ -234,8 +235,8 @@ 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)
     db_dir              : Path = Field(default=Path('databases'), description='Path to InvokeAI databases directory', json_schema_extra=Categories.Paths)
     outdir              : Path = Field(default=Path('outputs'), description='Default folder for output images', json_schema_extra=Categories.Paths)
@@ -259,10 +260,8 @@ class InvokeAIAppConfig(InvokeAISettings):
     version             : bool = Field(default=False, description="Show InvokeAI version and exit", json_schema_extra=Categories.Other)
 
     # CACHE
-    ram                 : float = Field(default=DEFAULT_RAM_CACHE, gt=0, description="Maximum memory amount used by model cache for rapid switching (floating point number, GB)", json_schema_extra=Categories.ModelCache, )
-    vram                : float = Field(default=DEFAULT_VRAM_CACHE, ge=0, description="Amount of VRAM reserved for model storage (floating point number, GB)", json_schema_extra=Categories.ModelCache, )
-    convert_cache       : float = Field(default=DEFAULT_CONVERT_CACHE, ge=0, description="Maximum size of on-disk converted models cache (GB)", json_schema_extra=Categories.ModelCache)
-
+    ram                 : float = Field(default=7.5, gt=0, description="Maximum memory amount used by model cache for rapid switching (floating point number, GB)", json_schema_extra=Categories.ModelCache, )
+    vram                : float = Field(default=0.25, ge=0, description="Amount of VRAM reserved for model storage (floating point number, GB)", json_schema_extra=Categories.ModelCache, )
     lazy_offload        : bool = Field(default=True, description="Keep models in VRAM until their space is needed", json_schema_extra=Categories.ModelCache, )
     log_memory_usage    : bool = Field(default=False, description="If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.", json_schema_extra=Categories.ModelCache)
 
@@ -297,7 +296,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)
@@ -406,11 +404,6 @@ class InvokeAIAppConfig(InvokeAISettings):
         """Path to the models directory."""
         return self._resolve(self.models_dir)
 
-    @property
-    def models_convert_cache_path(self) -> Path:
-        """Path to the converted cache models directory."""
-        return self._resolve(self.convert_cache_dir)
-
     @property
     def custom_nodes_path(self) -> Path:
         """Path to the custom nodes directory."""
@@ -440,20 +433,15 @@ class InvokeAIAppConfig(InvokeAISettings):
         return True
 
     @property
-    def ram_cache_size(self) -> float:
-        """Return the ram cache size using the legacy or modern setting (GB)."""
+    def ram_cache_size(self) -> Union[Literal["auto"], float]:
+        """Return the ram cache size using the legacy or modern setting."""
         return self.max_cache_size or self.ram
 
     @property
-    def vram_cache_size(self) -> float:
-        """Return the vram cache size using the legacy or modern setting (GB)."""
+    def vram_cache_size(self) -> Union[Literal["auto"], float]:
+        """Return the vram cache size using the legacy or modern setting."""
         return self.max_vram_cache_size or self.vram
 
-    @property
-    def convert_cache_size(self) -> float:
-        """Return the convert cache size on disk (GB)."""
-        return self.convert_cache
-
     @property
     def use_cpu(self) -> bool:
         """Return true if the device is set to CPU or the always_use_cpu flag is set."""

invokeai/app/services/download/download_base.py

@@ -260,16 +260,3 @@ class DownloadQueueServiceBase(ABC):
     def join(self) -> None:
         """Wait until all jobs are off the queue."""
         pass
-
-    @abstractmethod
-    def wait_for_job(self, job: DownloadJob, timeout: int = 0) -> DownloadJob:
-        """Wait until the indicated download job has reached a terminal state.
-
-        This will block until the indicated install job has completed,
-        been cancelled, or errored out.
-
-        :param job: The job to wait on.
-        :param timeout: Wait up to indicated number of seconds. Raise a TimeoutError if
-        the job hasn't completed within the indicated time.
-        """
-        pass

invokeai/app/services/download/download_default.py

@@ -4,11 +4,10 @@
 import os
 import re
 import threading
-import time
 import traceback
 from pathlib import Path
 from queue import Empty, PriorityQueue
-from typing import Any, Dict, List, Optional, Set
+from typing import Any, Dict, List, Optional
 
 import requests
 from pydantic.networks import AnyHttpUrl
@@ -49,12 +48,11 @@ class DownloadQueueService(DownloadQueueServiceBase):
         :param max_parallel_dl: Number of simultaneous downloads allowed [5].
         :param requests_session: Optional requests.sessions.Session object, for unit tests.
         """
-        self._jobs: Dict[int, DownloadJob] = {}
+        self._jobs = {}
         self._next_job_id = 0
-        self._queue: PriorityQueue[DownloadJob] = PriorityQueue()
+        self._queue = PriorityQueue()
         self._stop_event = threading.Event()
-        self._job_completed_event = threading.Event()
-        self._worker_pool: Set[threading.Thread] = set()
+        self._worker_pool = set()
         self._lock = threading.Lock()
         self._logger = InvokeAILogger.get_logger("DownloadQueueService")
         self._event_bus = event_bus
@@ -190,16 +188,6 @@ class DownloadQueueService(DownloadQueueServiceBase):
             if not job.in_terminal_state:
                 self.cancel_job(job)
 
-    def wait_for_job(self, job: DownloadJob, timeout: int = 0) -> DownloadJob:
-        """Block until the indicated job has reached terminal state, or when timeout limit reached."""
-        start = time.time()
-        while not job.in_terminal_state:
-            if self._job_completed_event.wait(timeout=0.25):  # in case we miss an event
-                self._job_completed_event.clear()
-            if timeout > 0 and time.time() - start > timeout:
-                raise TimeoutError("Timeout exceeded")
-        return job
-
     def _start_workers(self, max_workers: int) -> None:
         """Start the requested number of worker threads."""
         self._stop_event.clear()
@@ -228,20 +216,13 @@ class DownloadQueueService(DownloadQueueServiceBase):
             except (OSError, HTTPError) as excp:
                 job.error_type = excp.__class__.__name__ + f"({str(excp)})"
                 job.error = traceback.format_exc()
-                try:
-                    self._signal_job_error(job, excp)
-                except:
-                    pass
+                self._signal_job_error(job, excp)
             except DownloadJobCancelledException:
-                try:
-                    self._signal_job_cancelled(job)
-                    self._cleanup_cancelled_job(job)
-                except:
-                    pass
+                self._signal_job_cancelled(job)
+                self._cleanup_cancelled_job(job)
 
             finally:
                 job.job_ended = get_iso_timestamp()
-                self._job_completed_event.set()  # signal a change to terminal state
                 self._queue.task_done()
         self._logger.debug(f"Download queue worker thread {threading.current_thread().name} exiting.")
 
@@ -426,11 +407,11 @@ class DownloadQueueService(DownloadQueueServiceBase):
 
 # Example on_progress event handler to display a TQDM status bar
 # Activate with:
-#   download_service.download(DownloadJob('http://foo.bar/baz', '/tmp', on_progress=TqdmProgress().update))
+#   download_service.download('http://foo.bar/baz', '/tmp', on_progress=TqdmProgress().job_update
 class TqdmProgress(object):
     """TQDM-based progress bar object to use in on_progress handlers."""
 
-    _bars: Dict[int, tqdm]  # type: ignore
+    _bars: Dict[int, tqdm]  # the tqdm object
     _last: Dict[int, int]  # last bytes downloaded
 
     def __init__(self) -> None:  # noqa D107

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,
@@ -11,12 +11,12 @@ from invokeai.app.services.session_queue.session_queue_common import (
     SessionQueueStatus,
 )
 from invokeai.app.util.misc import get_timestamp
-from invokeai.backend.model_manager import AnyModelConfig
+from invokeai.backend.model_management.model_manager import ModelInfo
+from invokeai.backend.model_management.models.base import BaseModelType, ModelType, SubModelType
 
 
 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 +25,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()
@@ -63,7 +55,7 @@ class EventServiceBase:
         queue_item_id: int,
         queue_batch_id: str,
         graph_execution_state_id: str,
-        node_id: str,
+        node: dict,
         source_node_id: str,
         progress_image: Optional[ProgressImage],
         step: int,
@@ -78,7 +70,7 @@ class EventServiceBase:
                 "queue_item_id": queue_item_id,
                 "queue_batch_id": queue_batch_id,
                 "graph_execution_state_id": graph_execution_state_id,
-                "node_id": node_id,
+                "node_id": node.get("id"),
                 "source_node_id": source_node_id,
                 "progress_image": progress_image.model_dump() if progress_image is not None else None,
                 "step": step,
@@ -179,7 +171,10 @@ class EventServiceBase:
         queue_item_id: int,
         queue_batch_id: str,
         graph_execution_state_id: str,
-        model_config: AnyModelConfig,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        submodel: SubModelType,
     ) -> None:
         """Emitted when a model is requested"""
         self.__emit_queue_event(
@@ -189,7 +184,10 @@ class EventServiceBase:
                 "queue_item_id": queue_item_id,
                 "queue_batch_id": queue_batch_id,
                 "graph_execution_state_id": graph_execution_state_id,
-                "model_config": model_config.model_dump(),
+                "model_name": model_name,
+                "base_model": base_model,
+                "model_type": model_type,
+                "submodel": submodel,
             },
         )
 
@@ -199,7 +197,11 @@ class EventServiceBase:
         queue_item_id: int,
         queue_batch_id: str,
         graph_execution_state_id: str,
-        model_config: AnyModelConfig,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        submodel: SubModelType,
+        model_info: ModelInfo,
     ) -> None:
         """Emitted when a model is correctly loaded (returns model info)"""
         self.__emit_queue_event(
@@ -209,7 +211,59 @@ class EventServiceBase:
                 "queue_item_id": queue_item_id,
                 "queue_batch_id": queue_batch_id,
                 "graph_execution_state_id": graph_execution_state_id,
-                "model_config": model_config.model_dump(),
+                "model_name": model_name,
+                "base_model": base_model,
+                "model_type": model_type,
+                "submodel": submodel,
+                "hash": model_info.hash,
+                "location": str(model_info.location),
+                "precision": str(model_info.precision),
+            },
+        )
+
+    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,
             },
         )
 
@@ -357,7 +411,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 +430,6 @@ class EventServiceBase:
                 "bytes": bytes,
                 "total_bytes": total_bytes,
                 "parts": parts,
-                "id": id,
             },
         )
 
@@ -392,7 +444,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 +454,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 +472,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 +487,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/image_files/image_files_base.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 from PIL.Image import Image as PILImageType
 
-from invokeai.app.invocations.fields import MetadataField
+from invokeai.app.invocations.baseinvocation import MetadataField
 from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
 
 

invokeai/app/services/image_files/image_files_disk.py

@@ -7,7 +7,7 @@ from PIL import Image, PngImagePlugin
 from PIL.Image import Image as PILImageType
 from send2trash import send2trash
 
-from invokeai.app.invocations.fields import MetadataField
+from invokeai.app.invocations.baseinvocation import MetadataField
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
 from invokeai.app.util.thumbnails import get_thumbnail_name, make_thumbnail

invokeai/app/services/image_records/image_records_base.py

@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Optional
 
-from invokeai.app.invocations.fields import MetadataField
+from invokeai.app.invocations.metadata import MetadataField
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 
 from .image_records_common import ImageCategory, ImageRecord, ImageRecordChanges, ResourceOrigin

invokeai/app/services/image_records/image_records_sqlite.py

@@ -3,7 +3,7 @@ import threading
 from datetime import datetime
 from typing import Optional, Union, cast
 
-from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator
+from invokeai.app.invocations.baseinvocation import MetadataField, MetadataFieldValidator
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
 

invokeai/app/services/images/images_base.py

@@ -3,7 +3,7 @@ from typing import Callable, Optional
 
 from PIL.Image import Image as PILImageType
 
-from invokeai.app.invocations.fields import MetadataField
+from invokeai.app.invocations.baseinvocation import MetadataField
 from invokeai.app.services.image_records.image_records_common import (
     ImageCategory,
     ImageRecord,

invokeai/app/services/images/images_default.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 from PIL.Image import Image as PILImageType
 
-from invokeai.app.invocations.fields import MetadataField
+from invokeai.app.invocations.baseinvocation import MetadataField
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
@@ -154,7 +154,7 @@ class ImageService(ImageServiceABC):
             self.__invoker.services.logger.error("Image record not found")
             raise
         except Exception as e:
-            self.__invoker.services.logger.error("Problem getting image metadata")
+            self.__invoker.services.logger.error("Problem getting image DTO")
             raise e
 
     def get_workflow(self, image_name: str) -> Optional[WorkflowWithoutID]:

invokeai/app/services/invocation_cache/invocation_cache_memory.py

@@ -37,8 +37,7 @@ class MemoryInvocationCache(InvocationCacheBase):
         if self._max_cache_size == 0:
             return
         self._invoker.services.images.on_deleted(self._delete_by_match)
-        self._invoker.services.tensors.on_deleted(self._delete_by_match)
-        self._invoker.services.conditioning.on_deleted(self._delete_by_match)
+        self._invoker.services.latents.on_deleted(self._delete_by_match)
 
     def get(self, key: Union[int, str]) -> Optional[BaseInvocationOutput]:
         with self._lock:

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,237 @@
+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.invocations.baseinvocation import InvocationContext
+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.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:
+        # if we do want multithreading at some point, we could make this configurable
+        self.__threadLimit = BoundedSemaphore(1)
+        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
+                        outputs = invocation.invoke_internal(
+                            InvocationContext(
+                                services=self.__invoker.services,
+                                graph_execution_state_id=graph_execution_state.id,
+                                queue_item_id=queue_item.session_queue_item_id,
+                                queue_id=queue_item.session_queue_id,
+                                queue_batch_id=queue_item.session_queue_batch_id,
+                                workflow=queue_item.workflow,
+                            )
+                        )
+
+                        # 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

@@ -3,20 +3,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from invokeai.app.services.object_serializer.object_serializer_base import ObjectSerializerBase
-
 if TYPE_CHECKING:
     from logging import Logger
 
-    import torch
-
-    from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
-
     from .board_image_records.board_image_records_base import BoardImageRecordStorageBase
     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 +17,18 @@ 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 .latents_storage.latents_storage_base import LatentsStorageBase
+    from .model_install import ModelInstallServiceBase
     from .model_manager.model_manager_base import ModelManagerServiceBase
+    from .model_records import ModelRecordServiceBase
     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
 
@@ -36,50 +36,83 @@ if TYPE_CHECKING:
 class InvocationServices:
     """Services that can be used by invocations"""
 
+    # TODO: Just forward-declared everything due to circular dependencies. Fix structure.
+    board_images: "BoardImagesServiceABC"
+    board_image_record_storage: "BoardImageRecordStorageBase"
+    boards: "BoardServiceABC"
+    board_records: "BoardRecordStorageBase"
+    configuration: "InvokeAIAppConfig"
+    events: "EventServiceBase"
+    graph_execution_manager: "ItemStorageABC[GraphExecutionState]"
+    images: "ImageServiceABC"
+    image_records: "ImageRecordStorageBase"
+    image_files: "ImageFileStorageBase"
+    latents: "LatentsStorageBase"
+    logger: "Logger"
+    model_manager: "ModelManagerServiceBase"
+    model_records: "ModelRecordServiceBase"
+    download_queue: "DownloadQueueServiceBase"
+    model_install: "ModelInstallServiceBase"
+    processor: "InvocationProcessorABC"
+    performance_statistics: "InvocationStatsServiceBase"
+    queue: "InvocationQueueABC"
+    session_queue: "SessionQueueBase"
+    session_processor: "SessionProcessorBase"
+    invocation_cache: "InvocationCacheBase"
+    names: "NameServiceBase"
+    urls: "UrlServiceBase"
+    workflow_records: "WorkflowRecordsStorageBase"
+
     def __init__(
         self,
         board_images: "BoardImagesServiceABC",
         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",
+        latents: "LatentsStorageBase",
         logger: "Logger",
         model_manager: "ModelManagerServiceBase",
+        model_records: "ModelRecordServiceBase",
         download_queue: "DownloadQueueServiceBase",
+        model_install: "ModelInstallServiceBase",
+        processor: "InvocationProcessorABC",
         performance_statistics: "InvocationStatsServiceBase",
+        queue: "InvocationQueueABC",
         session_queue: "SessionQueueBase",
         session_processor: "SessionProcessorBase",
         invocation_cache: "InvocationCacheBase",
         names: "NameServiceBase",
         urls: "UrlServiceBase",
         workflow_records: "WorkflowRecordsStorageBase",
-        tensors: "ObjectSerializerBase[torch.Tensor]",
-        conditioning: "ObjectSerializerBase[ConditioningFieldData]",
     ):
         self.board_images = board_images
         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.latents = latents
         self.logger = logger
         self.model_manager = model_manager
+        self.model_records = model_records
         self.download_queue = download_queue
+        self.model_install = model_install
+        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
         self.names = names
         self.urls = urls
         self.workflow_records = workflow_records
-        self.tensors = tensors
-        self.conditioning = conditioning

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()
@@ -29,8 +29,8 @@ writes to the system log is stored in InvocationServices.performance_statistics.
 """
 
 from abc import ABC, abstractmethod
+from contextlib import AbstractContextManager
 from pathlib import Path
-from typing import ContextManager
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
@@ -40,17 +40,18 @@ class InvocationStatsServiceBase(ABC):
     "Abstract base class for recording node memory/time performance statistics"
 
     @abstractmethod
-    def __init__(self) -> None:
+    def __init__(self):
         """
         Initialize the InvocationStatsService and reset counters to zero
         """
+        pass
 
     @abstractmethod
     def collect_stats(
         self,
         invocation: BaseInvocation,
         graph_execution_state_id: str,
-    ) -> ContextManager[None]:
+    ) -> AbstractContextManager:
         """
         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,12 +61,16 @@ class InvocationStatsServiceBase(ABC):
         pass
 
     @abstractmethod
-    def reset_stats(self):
-        """Reset all stored statistics."""
+    def reset_stats(self, graph_execution_state_id: str):
+        """
+        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
-    def log_stats(self, graph_execution_state_id: str) -> None:
+    def log_stats(self, graph_execution_state_id: str):
         """
         Write out the accumulated statistics to the log or somewhere else.
         :param graph_execution_state_id: The id of the session whose stats to log.

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -2,7 +2,6 @@ import json
 import time
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Generator
 
 import psutil
 import torch
@@ -10,7 +9,8 @@ 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.backend.model_manager.load.model_cache import CacheStats
+from invokeai.app.services.item_storage.item_storage_common import ItemNotFoundError
+from invokeai.backend.model_management.model_cache import CacheStats
 
 from .invocation_stats_base import InvocationStatsServiceBase
 from .invocation_stats_common import (
@@ -41,23 +41,22 @@ class InvocationStatsService(InvocationStatsServiceBase):
         self._invoker = invoker
 
     @contextmanager
-    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Generator[None, None, None]:
-        # This is to handle case of the model manager not being initialized, which happens
-        # during some tests.
-        services = self._invoker.services
+    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str):
         if not self._stats.get(graph_execution_state_id):
             # First time we're seeing this graph_execution_state_id.
             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
         if torch.cuda.is_available():
             torch.cuda.reset_peak_memory_stats()
-
-        assert services.model_manager.load is not None
-        services.model_manager.load.ram_cache.stats = self._cache_stats[graph_execution_state_id]
+        if self._invoker.services.model_manager:
+            self._invoker.services.model_manager.collect_cache_stats(self._cache_stats[graph_execution_state_id])
 
         try:
             # Let the invocation run.
@@ -74,9 +73,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):
+        """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/item_storage/item_storage_base.py

@@ -30,7 +30,7 @@ class ItemStorageABC(ABC, Generic[T]):
     @abstractmethod
     def set(self, item: T) -> None:
         """
-        Sets the item.
+        Sets the item. The id will be extracted based on id_field.
         :param item: the item to set
         """
         pass

invokeai/app/services/latents_storage/latents_storage_base.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from abc import ABC, abstractmethod
+from typing import Callable
+
+import torch
+
+
+class LatentsStorageBase(ABC):
+    """Responsible for storing and retrieving latents."""
+
+    _on_changed_callbacks: list[Callable[[torch.Tensor], None]]
+    _on_deleted_callbacks: list[Callable[[str], None]]
+
+    def __init__(self) -> None:
+        self._on_changed_callbacks = []
+        self._on_deleted_callbacks = []
+
+    @abstractmethod
+    def get(self, name: str) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def save(self, name: str, data: torch.Tensor) -> None:
+        pass
+
+    @abstractmethod
+    def delete(self, name: str) -> None:
+        pass
+
+    def on_changed(self, on_changed: Callable[[torch.Tensor], None]) -> None:
+        """Register a callback for when an item is changed"""
+        self._on_changed_callbacks.append(on_changed)
+
+    def on_deleted(self, on_deleted: Callable[[str], None]) -> None:
+        """Register a callback for when an item is deleted"""
+        self._on_deleted_callbacks.append(on_deleted)
+
+    def _on_changed(self, item: torch.Tensor) -> None:
+        for callback in self._on_changed_callbacks:
+            callback(item)
+
+    def _on_deleted(self, item_id: str) -> None:
+        for callback in self._on_deleted_callbacks:
+            callback(item_id)

invokeai/app/services/latents_storage/latents_storage_disk.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from pathlib import Path
+from typing import Union
+
+import torch
+
+from invokeai.app.services.invoker import Invoker
+
+from .latents_storage_base import LatentsStorageBase
+
+
+class DiskLatentsStorage(LatentsStorageBase):
+    """Stores latents in a folder on disk without caching"""
+
+    __output_folder: Path
+
+    def __init__(self, output_folder: Union[str, Path]):
+        self.__output_folder = output_folder if isinstance(output_folder, Path) else Path(output_folder)
+        self.__output_folder.mkdir(parents=True, exist_ok=True)
+
+    def start(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+        self._delete_all_latents()
+
+    def get(self, name: str) -> torch.Tensor:
+        latent_path = self.get_path(name)
+        return torch.load(latent_path)
+
+    def save(self, name: str, data: torch.Tensor) -> None:
+        self.__output_folder.mkdir(parents=True, exist_ok=True)
+        latent_path = self.get_path(name)
+        torch.save(data, latent_path)
+
+    def delete(self, name: str) -> None:
+        latent_path = self.get_path(name)
+        latent_path.unlink()
+
+    def get_path(self, name: str) -> Path:
+        return self.__output_folder / name
+
+    def _delete_all_latents(self) -> None:
+        """
+        Deletes all latents from disk.
+        Must be called after we have access to `self._invoker` (e.g. in `start()`).
+        """
+        deleted_latents_count = 0
+        freed_space = 0
+        for latents_file in Path(self.__output_folder).glob("*"):
+            if latents_file.is_file():
+                freed_space += latents_file.stat().st_size
+                deleted_latents_count += 1
+                latents_file.unlink()
+        if deleted_latents_count > 0:
+            freed_space_in_mb = round(freed_space / 1024 / 1024, 2)
+            self._invoker.services.logger.info(
+                f"Deleted {deleted_latents_count} latents files (freed {freed_space_in_mb}MB)"
+            )

invokeai/app/services/latents_storage/latents_storage_forward_cache.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from queue import Queue
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.app.services.invoker import Invoker
+
+from .latents_storage_base import LatentsStorageBase
+
+
+class ForwardCacheLatentsStorage(LatentsStorageBase):
+    """Caches the latest N latents in memory, writing-thorugh to and reading from underlying storage"""
+
+    __cache: Dict[str, torch.Tensor]
+    __cache_ids: Queue
+    __max_cache_size: int
+    __underlying_storage: LatentsStorageBase
+
+    def __init__(self, underlying_storage: LatentsStorageBase, max_cache_size: int = 20):
+        super().__init__()
+        self.__underlying_storage = underlying_storage
+        self.__cache = {}
+        self.__cache_ids = Queue()
+        self.__max_cache_size = max_cache_size
+
+    def start(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+        start_op = getattr(self.__underlying_storage, "start", None)
+        if callable(start_op):
+            start_op(invoker)
+
+    def stop(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+        stop_op = getattr(self.__underlying_storage, "stop", None)
+        if callable(stop_op):
+            stop_op(invoker)
+
+    def get(self, name: str) -> torch.Tensor:
+        cache_item = self.__get_cache(name)
+        if cache_item is not None:
+            return cache_item
+
+        latent = self.__underlying_storage.get(name)
+        self.__set_cache(name, latent)
+        return latent
+
+    def save(self, name: str, data: torch.Tensor) -> None:
+        self.__underlying_storage.save(name, data)
+        self.__set_cache(name, data)
+        self._on_changed(data)
+
+    def delete(self, name: str) -> None:
+        self.__underlying_storage.delete(name)
+        if name in self.__cache:
+            del self.__cache[name]
+        self._on_deleted(name)
+
+    def __get_cache(self, name: str) -> Optional[torch.Tensor]:
+        return None if name not in self.__cache else self.__cache[name]
+
+    def __set_cache(self, name: str, data: torch.Tensor):
+        if name not in self.__cache:
+            self.__cache[name] = data
+            self.__cache_ids.put(name)
+            if self.__cache_ids.qsize() > self.__max_cache_size:
+                self.__cache.pop(self.__cache_ids.get())

invokeai/app/services/model_install/model_install_base.py

@@ -14,13 +14,11 @@ from typing_extensions import Annotated
 
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.download import DownloadJob, DownloadQueueServiceBase
-from invokeai.app.services.events.events_base import EventServiceBase
+from invokeai.app.services.events import EventServiceBase
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_records import ModelRecordServiceBase
 from invokeai.backend.model_manager import AnyModelConfig, ModelRepoVariant
-from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
-
-from ..model_metadata import ModelMetadataStoreBase
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata, ModelMetadataStore
 
 
 class InstallStatus(str, Enum):
@@ -129,8 +127,8 @@ class HFModelSource(StringLikeSource):
     def __str__(self) -> str:
         """Return string version of repoid when string rep needed."""
         base: str = self.repo_id
-        base += f":{self.variant or ''}"
         base += f":{self.subfolder}" if self.subfolder else ""
+        base += f" ({self.variant})" if self.variant else ""
         return base
 
 
@@ -156,7 +154,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 +175,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 +182,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 +193,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:
@@ -254,7 +243,7 @@ class ModelInstallServiceBase(ABC):
         app_config: InvokeAIAppConfig,
         record_store: ModelRecordServiceBase,
         download_queue: DownloadQueueServiceBase,
-        metadata_store: ModelMetadataStoreBase,
+        metadata_store: ModelMetadataStore,
         event_bus: Optional["EventServiceBase"] = None,
     ):
         """
@@ -335,43 +324,6 @@ class ModelInstallServiceBase(ABC):
         :returns id: The string ID of the registered model.
         """
 
-    @abstractmethod
-    def heuristic_import(
-        self,
-        source: str,
-        config: Optional[Dict[str, Any]] = None,
-        access_token: Optional[str] = None,
-    ) -> ModelInstallJob:
-        r"""Install the indicated model using heuristics to interpret user intentions.
-
-        :param source: String source
-        :param config: Optional dict. Any fields in this dict
-         will override corresponding autoassigned probe fields in the
-         model's config record as described in `import_model()`.
-        :param access_token: Optional access token for remote sources.
-
-        The source can be:
-        1. A local file path in posix() format (`/foo/bar` or `C:\foo\bar`)
-        2. An http or https URL (`https://foo.bar/foo`)
-        3. A HuggingFace repo_id (`foo/bar`, `foo/bar:fp16`, `foo/bar:fp16:vae`)
-
-        We extend the HuggingFace repo_id syntax to include the variant and the
-        subfolder or path. The following are acceptable alternatives:
-            stabilityai/stable-diffusion-v4
-            stabilityai/stable-diffusion-v4:fp16
-            stabilityai/stable-diffusion-v4:fp16:vae
-            stabilityai/stable-diffusion-v4::/checkpoints/sd4.safetensors
-            stabilityai/stable-diffusion-v4:onnx:vae
-
-        Because a local file path can look like a huggingface repo_id, the logic
-        first checks whether the path exists on disk, and if not, it is treated as
-        a parseable huggingface repo.
-
-        The previous support for recursing into a local folder and loading all model-like files
-        has been removed.
-        """
-        pass
-
     @abstractmethod
     def import_model(
         self,
@@ -433,18 +385,6 @@ class ModelInstallServiceBase(ABC):
     def cancel_job(self, job: ModelInstallJob) -> None:
         """Cancel the indicated job."""
 
-    @abstractmethod
-    def wait_for_job(self, job: ModelInstallJob, timeout: int = 0) -> ModelInstallJob:
-        """Wait for the indicated job to reach a terminal state.
-
-        This will block until the indicated install job has completed,
-        been cancelled, or errored out.
-
-        :param job: The job to wait on.
-        :param timeout: Wait up to indicated number of seconds. Raise a TimeoutError if
-        the job hasn't completed within the indicated time.
-        """
-
     @abstractmethod
     def wait_for_installs(self, timeout: int = 0) -> List[ModelInstallJob]:
         """
@@ -454,8 +394,7 @@ class ModelInstallServiceBase(ABC):
         completed, been cancelled, or errored out.
 
         :param timeout: Wait up to indicated number of seconds. Raise an Exception('timeout') if
-        installs do not complete within the indicated time. A timeout of zero (the default)
-        will block indefinitely until the installs complete.
+        installs do not complete within the indicated time.
         """
 
     @abstractmethod
@@ -471,22 +410,3 @@ class ModelInstallServiceBase(ABC):
     @abstractmethod
     def sync_to_config(self) -> None:
         """Synchronize models on disk to those in the model record database."""
-
-    @abstractmethod
-    def download_and_cache(self, source: Union[str, AnyHttpUrl], access_token: Optional[str] = None) -> Path:
-        """
-        Download the model file located at source to the models cache and return its Path.
-
-        :param source: A Url or a string that can be converted into one.
-        :param access_token: Optional access token to access restricted resources.
-
-        The model file will be downloaded into the system-wide model cache
-        (`models/.cache`) if it isn't already there. Note that the model cache
-        is periodically cleared of infrequently-used entries when the model
-        converter runs.
-
-        Note that this doesn't automaticallly install or register the model, but is
-        intended for use by nodes that need access to models that aren't directly
-        supported by InvokeAI. The downloading process takes advantage of the download queue
-        to avoid interrupting other operations.
-        """

invokeai/app/services/model_install/model_install_default.py

@@ -17,10 +17,10 @@ from pydantic.networks import AnyHttpUrl
 from requests import Session
 
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.app.services.download import DownloadJob, DownloadQueueServiceBase, TqdmProgress
+from invokeai.app.services.download import DownloadJob, DownloadQueueServiceBase
 from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
-from invokeai.app.services.model_records import DuplicateModelException, ModelRecordServiceBase
+from invokeai.app.services.model_records import DuplicateModelException, ModelRecordServiceBase, ModelRecordServiceSQL
 from invokeai.backend.model_manager.config import (
     AnyModelConfig,
     BaseModelType,
@@ -33,6 +33,7 @@ from invokeai.backend.model_manager.metadata import (
     AnyModelRepoMetadata,
     CivitaiMetadataFetch,
     HuggingFaceMetadataFetch,
+    ModelMetadataStore,
     ModelMetadataWithFiles,
     RemoteModelFile,
 )
@@ -49,7 +50,6 @@ from .model_install_base import (
     ModelInstallJob,
     ModelInstallServiceBase,
     ModelSource,
-    StringLikeSource,
     URLModelSource,
 )
 
@@ -64,6 +64,7 @@ class ModelInstallService(ModelInstallServiceBase):
         app_config: InvokeAIAppConfig,
         record_store: ModelRecordServiceBase,
         download_queue: DownloadQueueServiceBase,
+        metadata_store: Optional[ModelMetadataStore] = None,
         event_bus: Optional[EventServiceBase] = None,
         session: Optional[Session] = None,
     ):
@@ -85,13 +86,19 @@ class ModelInstallService(ModelInstallServiceBase):
         self._lock = threading.Lock()
         self._stop_event = threading.Event()
         self._downloads_changed_event = threading.Event()
-        self._install_completed_event = threading.Event()
         self._download_queue = download_queue
         self._download_cache: Dict[AnyHttpUrl, ModelInstallJob] = {}
         self._running = False
         self._session = session
         self._next_job_id = 0
-        self._metadata_store = record_store.metadata_store  # for convenience
+        # There may not necessarily be a metadata store initialized
+        # so we create one and initialize it with the same sql database
+        # used by the record store service.
+        if metadata_store:
+            self._metadata_store = metadata_store
+        else:
+            assert isinstance(record_store, ModelRecordServiceSQL)
+            self._metadata_store = ModelMetadataStore(record_store.db)
 
     @property
     def app_config(self) -> InvokeAIAppConfig:  # noqa D102
@@ -138,7 +145,7 @@ class ModelInstallService(ModelInstallServiceBase):
     ) -> str:  # noqa D102
         model_path = Path(model_path)
         config = config or {}
-        if not config.get("source"):
+        if config.get("source") is None:
             config["source"] = model_path.resolve().as_posix()
         return self._register(model_path, config)
 
@@ -149,18 +156,12 @@ class ModelInstallService(ModelInstallServiceBase):
     ) -> str:  # noqa D102
         model_path = Path(model_path)
         config = config or {}
-        if not config.get("source"):
+        if config.get("source") is None:
             config["source"] = model_path.resolve().as_posix()
 
         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)
-        )
+        old_hash = info.original_hash
+        dest_path = self.app_config.models_path / info.base.value / info.type.value / model_path.name
         try:
             new_path = self._copy_model(model_path, dest_path)
         except FileExistsError as excp:
@@ -176,40 +177,7 @@ class ModelInstallService(ModelInstallServiceBase):
             info,
         )
 
-    def heuristic_import(
-        self,
-        source: str,
-        config: Optional[Dict[str, Any]] = None,
-        access_token: Optional[str] = None,
-    ) -> ModelInstallJob:
-        variants = "|".join(ModelRepoVariant.__members__.values())
-        hf_repoid_re = f"^([^/:]+/[^/:]+)(?::({variants})?(?::/?([^:]+))?)?$"
-        source_obj: Optional[StringLikeSource] = None
-
-        if Path(source).exists():  # A local file or directory
-            source_obj = LocalModelSource(path=Path(source))
-        elif match := re.match(hf_repoid_re, source):
-            source_obj = HFModelSource(
-                repo_id=match.group(1),
-                variant=match.group(2) if match.group(2) else None,  # pass None rather than ''
-                subfolder=Path(match.group(3)) if match.group(3) else None,
-                access_token=access_token,
-            )
-        elif re.match(r"^https?://[^/]+", source):
-            source_obj = URLModelSource(
-                url=AnyHttpUrl(source),
-                access_token=access_token,
-            )
-        else:
-            raise ValueError(f"Unsupported model source: '{source}'")
-        return self.import_model(source_obj, config)
-
     def import_model(self, source: ModelSource, config: Optional[Dict[str, Any]] = None) -> ModelInstallJob:  # noqa D102
-        similar_jobs = [x for x in self.list_jobs() if x.source == source and not x.in_terminal_state]
-        if similar_jobs:
-            self._logger.warning(f"There is already an active install job for {source}. Not enqueuing.")
-            return similar_jobs[0]
-
         if isinstance(source, LocalModelSource):
             install_job = self._import_local_model(source, config)
             self._install_queue.put(install_job)  # synchronously install
@@ -239,25 +207,14 @@ class ModelInstallService(ModelInstallServiceBase):
         assert isinstance(jobs[0], ModelInstallJob)
         return jobs[0]
 
-    def wait_for_job(self, job: ModelInstallJob, timeout: int = 0) -> ModelInstallJob:
-        """Block until the indicated job has reached terminal state, or when timeout limit reached."""
-        start = time.time()
-        while not job.in_terminal_state:
-            if self._install_completed_event.wait(timeout=5):  # in case we miss an event
-                self._install_completed_event.clear()
-            if timeout > 0 and time.time() - start > timeout:
-                raise TimeoutError("Timeout exceeded")
-        return job
-
-    # TODO: Better name? Maybe wait_for_jobs()? Maybe too easily confused with above
     def wait_for_installs(self, timeout: int = 0) -> List[ModelInstallJob]:  # noqa D102
         """Block until all installation jobs are done."""
         start = time.time()
         while len(self._download_cache) > 0:
-            if self._downloads_changed_event.wait(timeout=0.25):  # in case we miss an event
+            if self._downloads_changed_event.wait(timeout=5):  # in case we miss an event
                 self._downloads_changed_event.clear()
             if timeout > 0 and time.time() - start > timeout:
-                raise TimeoutError("Timeout exceeded")
+                raise Exception("Timeout exceeded")
         self._install_queue.join()
         return self._install_jobs
 
@@ -311,38 +268,6 @@ class ModelInstallService(ModelInstallServiceBase):
             path.unlink()
         self.unregister(key)
 
-    def download_and_cache(
-        self,
-        source: Union[str, AnyHttpUrl],
-        access_token: Optional[str] = None,
-        timeout: int = 0,
-    ) -> Path:
-        """Download the model file located at source to the models cache and return its Path."""
-        model_hash = sha256(str(source).encode("utf-8")).hexdigest()[0:32]
-        model_path = self._app_config.models_convert_cache_path / model_hash
-
-        # We expect the cache directory to contain one and only one downloaded file.
-        # We don't know the file's name in advance, as it is set by the download
-        # content-disposition header.
-        if model_path.exists():
-            contents = [x for x in model_path.iterdir() if x.is_file()]
-            if len(contents) > 0:
-                return contents[0]
-
-        model_path.mkdir(parents=True, exist_ok=True)
-        job = self._download_queue.download(
-            source=AnyHttpUrl(str(source)),
-            dest=model_path,
-            access_token=access_token,
-            on_progress=TqdmProgress().update,
-        )
-        self._download_queue.wait_for_job(job, timeout)
-        if job.complete:
-            assert job.download_path is not None
-            return job.download_path
-        else:
-            raise Exception(job.error)
-
     # --------------------------------------------------------------------------------------------
     # Internal functions that manage the installer threads
     # --------------------------------------------------------------------------------------------
@@ -375,7 +300,6 @@ class ModelInstallService(ModelInstallServiceBase):
                     job.total_bytes = self._stat_size(job.local_path)
                     job.bytes = job.total_bytes
                     self._signal_job_running(job)
-                    job.config_in["source"] = str(job.source)
                     if job.inplace:
                         key = self.register_path(job.local_path, job.config_in)
                     else:
@@ -406,7 +330,6 @@ class ModelInstallService(ModelInstallServiceBase):
                 # if this is an install of a remote file, then clean up the temporary directory
                 if job._install_tmpdir is not None:
                     rmtree(job._install_tmpdir)
-                self._install_completed_event.set()
                 self._install_queue.task_done()
 
         self._logger.info("Install thread exiting")
@@ -542,10 +465,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):
@@ -558,8 +479,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:
@@ -568,10 +489,10 @@ class ModelInstallService(ModelInstallServiceBase):
         return id
 
     @staticmethod
-    def _guess_variant() -> Optional[ModelRepoVariant]:
+    def _guess_variant() -> ModelRepoVariant:
         """Guess the best HuggingFace variant type to download."""
         precision = choose_precision(choose_torch_device())
-        return ModelRepoVariant.FP16 if precision == "float16" else None
+        return ModelRepoVariant.FP16 if precision == "float16" else ModelRepoVariant.DEFAULT
 
     def _import_local_model(self, source: LocalModelSource, config: Optional[Dict[str, Any]]) -> ModelInstallJob:
         return ModelInstallJob(
@@ -596,7 +517,7 @@ class ModelInstallService(ModelInstallServiceBase):
         if not source.access_token:
             self._logger.info("No HuggingFace access token present; some models may not be downloadable.")
 
-        metadata = HuggingFaceMetadataFetch(self._session).from_id(source.repo_id, source.variant)
+        metadata = HuggingFaceMetadataFetch(self._session).from_id(source.repo_id)
         assert isinstance(metadata, ModelMetadataWithFiles)
         remote_files = metadata.download_urls(
             variant=source.variant or self._guess_variant(),
@@ -644,8 +565,6 @@ class ModelInstallService(ModelInstallServiceBase):
         # TODO: Replace with tempfile.tmpdir() when multithreading is cleaned up.
         # Currently the tmpdir isn't automatically removed at exit because it is
         # being held in a daemon thread.
-        if len(remote_files) == 0:
-            raise ValueError(f"{source}: No downloadable files found")
         tmpdir = Path(
             mkdtemp(
                 dir=self._app_config.models_path,
@@ -661,16 +580,6 @@ class ModelInstallService(ModelInstallServiceBase):
             bytes=0,
             total_bytes=0,
         )
-        # In the event that there is a subfolder specified in the source,
-        # we need to remove it from the destination path in order to avoid
-        # creating unwanted subfolders
-        if hasattr(source, "subfolder") and source.subfolder:
-            root = Path(remote_files[0].path.parts[0])
-            subfolder = root / source.subfolder
-        else:
-            root = Path(".")
-            subfolder = Path(".")
-
         # we remember the path up to the top of the tmpdir so that it may be
         # removed safely at the end of the install process.
         install_job._install_tmpdir = tmpdir
@@ -680,7 +589,7 @@ class ModelInstallService(ModelInstallServiceBase):
         self._logger.debug(f"remote_files={remote_files}")
         for model_file in remote_files:
             url = model_file.url
-            path = root / model_file.path.relative_to(subfolder)
+            path = model_file.path
             self._logger.info(f"Downloading {url} => {path}")
             install_job.total_bytes += model_file.size
             assert hasattr(source, "access_token")
@@ -743,7 +652,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)
@@ -776,7 +684,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()
@@ -823,7 +731,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:
@@ -836,7 +743,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}")
@@ -845,7 +752,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/__init__.py

@@ -1,6 +0,0 @@
-"""Initialization file for model load service module."""
-
-from .model_load_base import ModelLoadServiceBase
-from .model_load_default import ModelLoadService
-
-__all__ = ["ModelLoadServiceBase", "ModelLoadService"]

invokeai/app/services/model_load/model_load_base.py

@@ -1,40 +0,0 @@
-# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Team
-"""Base class for model loader."""
-
-from abc import ABC, abstractmethod
-from typing import Optional
-
-from invokeai.app.services.shared.invocation_context import InvocationContextData
-from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType
-from invokeai.backend.model_manager.load import LoadedModel
-from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase
-from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
-
-
-class ModelLoadServiceBase(ABC):
-    """Wrapper around AnyModelLoader."""
-
-    @abstractmethod
-    def load_model(
-        self,
-        model_config: AnyModelConfig,
-        submodel_type: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
-        """
-        Given a model's configuration, load it and return the LoadedModel object.
-
-        :param model_config: Model configuration record (as returned by ModelRecordBase.get_model())
-        :param submodel: For main (pipeline models), the submodel to fetch.
-        :param context_data: Invocation context data used for event reporting
-        """
-
-    @property
-    @abstractmethod
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
-        """Return the RAM cache used by this loader."""
-
-    @property
-    @abstractmethod
-    def convert_cache(self) -> ModelConvertCacheBase:
-        """Return the checkpoint convert cache used by this loader."""

invokeai/app/services/model_load/model_load_default.py

@@ -1,113 +0,0 @@
-# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Team
-"""Implementation of model loader service."""
-
-from typing import Optional, Type
-
-from invokeai.app.services.config import InvokeAIAppConfig
-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
-from invokeai.backend.model_manager.load import (
-    LoadedModel,
-    ModelLoaderRegistry,
-    ModelLoaderRegistryBase,
-)
-from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase
-from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
-from invokeai.backend.util.logging import InvokeAILogger
-
-from .model_load_base import ModelLoadServiceBase
-
-
-class ModelLoadService(ModelLoadServiceBase):
-    """Wrapper around ModelLoaderRegistry."""
-
-    def __init__(
-        self,
-        app_config: InvokeAIAppConfig,
-        ram_cache: ModelCacheBase[AnyModel],
-        convert_cache: ModelConvertCacheBase,
-        registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry,
-    ):
-        """Initialize the model load service."""
-        logger = InvokeAILogger.get_logger(self.__class__.__name__)
-        logger.setLevel(app_config.log_level.upper())
-        self._logger = logger
-        self._app_config = app_config
-        self._ram_cache = ram_cache
-        self._convert_cache = convert_cache
-        self._registry = registry
-
-    def start(self, invoker: Invoker) -> None:
-        self._invoker = invoker
-
-    @property
-    def ram_cache(self) -> ModelCacheBase[AnyModel]:
-        """Return the RAM cache used by this loader."""
-        return self._ram_cache
-
-    @property
-    def convert_cache(self) -> ModelConvertCacheBase:
-        """Return the checkpoint convert cache used by this loader."""
-        return self._convert_cache
-
-    def load_model(
-        self,
-        model_config: AnyModelConfig,
-        submodel_type: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
-        """
-        Given a model's configuration, load it and return the LoadedModel object.
-
-        :param model_config: Model configuration record (as returned by ModelRecordBase.get_model())
-        :param submodel: For main (pipeline models), the submodel to fetch.
-        :param context: Invocation context used for event reporting
-        """
-        if context_data:
-            self._emit_load_event(
-                context_data=context_data,
-                model_config=model_config,
-            )
-
-        implementation, model_config, submodel_type = self._registry.get_implementation(model_config, submodel_type)  # type: ignore
-        loaded_model: LoadedModel = implementation(
-            app_config=self._app_config,
-            logger=self._logger,
-            ram_cache=self._ram_cache,
-            convert_cache=self._convert_cache,
-        ).load_model(model_config, submodel_type)
-
-        if context_data:
-            self._emit_load_event(
-                context_data=context_data,
-                model_config=model_config,
-                loaded=True,
-            )
-        return loaded_model
-
-    def _emit_load_event(
-        self,
-        context_data: InvocationContextData,
-        model_config: AnyModelConfig,
-        loaded: Optional[bool] = False,
-    ) -> None:
-        if not self._invoker:
-            return
-
-        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,
-                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,
-                model_config=model_config,
-            )

invokeai/app/services/model_manager/__init__.py

@@ -1,17 +1 @@
-"""Initialization file for model manager service."""
-
-from invokeai.backend.model_manager import AnyModel, AnyModelConfig, BaseModelType, ModelType, SubModelType
-from invokeai.backend.model_manager.load import LoadedModel
-
-from .model_manager_default import ModelManagerService, ModelManagerServiceBase
-
-__all__ = [
-    "ModelManagerServiceBase",
-    "ModelManagerService",
-    "AnyModel",
-    "AnyModelConfig",
-    "BaseModelType",
-    "ModelType",
-    "SubModelType",
-    "LoadedModel",
-]
+from .model_manager_default import ModelManagerService  # noqa F401

invokeai/app/services/model_manager/model_manager_base.py

@@ -1,101 +1,286 @@
 # Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Optional
+from logging import Logger
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
 
-import torch
-from typing_extensions import Self
+from pydantic import Field
 
-from invokeai.app.services.invoker import Invoker
-from invokeai.app.services.shared.invocation_context import InvocationContextData
-from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelType, SubModelType
-from invokeai.backend.model_manager.load.load_base import LoadedModel
+from invokeai.app.services.config.config_default import InvokeAIAppConfig
+from invokeai.backend.model_management import (
+    AddModelResult,
+    BaseModelType,
+    MergeInterpolationMethod,
+    ModelInfo,
+    ModelType,
+    SchedulerPredictionType,
+    SubModelType,
+)
+from invokeai.backend.model_management.model_cache import CacheStats
 
-from ..config import InvokeAIAppConfig
-from ..download import DownloadQueueServiceBase
-from ..events.events_base import EventServiceBase
-from ..model_install import ModelInstallServiceBase
-from ..model_load import ModelLoadServiceBase
-from ..model_records import ModelRecordServiceBase
+if TYPE_CHECKING:
+    from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext
 
 
 class ModelManagerServiceBase(ABC):
-    """Abstract base class for the model manager service."""
+    """Responsible for managing models on disk and in memory"""
 
-    # attributes:
-    # store: ModelRecordServiceBase = Field(description="An instance of the model record configuration service.")
-    # install: ModelInstallServiceBase = Field(description="An instance of the model install service.")
-    # load: ModelLoadServiceBase = Field(description="An instance of the model load service.")
-
-    @classmethod
     @abstractmethod
-    def build_model_manager(
-        cls,
-        app_config: InvokeAIAppConfig,
-        model_record_service: ModelRecordServiceBase,
-        download_queue: DownloadQueueServiceBase,
-        events: EventServiceBase,
-        execution_device: torch.device,
-    ) -> Self:
+    def __init__(
+        self,
+        config: InvokeAIAppConfig,
+        logger: Logger,
+    ):
         """
-        Construct the model manager service instance.
-
-        Use it rather than the __init__ constructor. This class
-        method simplifies the construction considerably.
+        Initialize with the path to the models.yaml config file.
+        Optional parameters are the torch device type, precision, max_models,
+        and sequential_offload boolean. Note that the default device
+        type and precision are set up for a CUDA system running at half precision.
         """
         pass
 
-    @property
     @abstractmethod
-    def store(self) -> ModelRecordServiceBase:
-        """Return the ModelRecordServiceBase used to store and retrieve configuration records."""
-        pass
-
-    @property
-    @abstractmethod
-    def load(self) -> ModelLoadServiceBase:
-        """Return the ModelLoadServiceBase used to load models from their configuration records."""
-        pass
-
-    @property
-    @abstractmethod
-    def install(self) -> ModelInstallServiceBase:
-        """Return the ModelInstallServiceBase used to download and manipulate model files."""
-        pass
-
-    @abstractmethod
-    def start(self, invoker: Invoker) -> None:
-        pass
-
-    @abstractmethod
-    def stop(self, invoker: Invoker) -> None:
-        pass
-
-    @abstractmethod
-    def load_model_by_config(
-        self,
-        model_config: AnyModelConfig,
-        submodel_type: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
-        pass
-
-    @abstractmethod
-    def load_model_by_key(
-        self,
-        key: str,
-        submodel_type: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
-        pass
-
-    @abstractmethod
-    def load_model_by_attr(
+    def get_model(
         self,
         model_name: str,
         base_model: BaseModelType,
         model_type: ModelType,
         submodel: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
+        node: Optional[BaseInvocation] = None,
+        context: Optional[InvocationContext] = None,
+    ) -> ModelInfo:
+        """Retrieve the indicated model with name and type.
+        submodel can be used to get a part (such as the vae)
+        of a diffusers pipeline."""
+        pass
+
+    @property
+    @abstractmethod
+    def logger(self):
+        pass
+
+    @abstractmethod
+    def model_exists(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+    ) -> bool:
+        pass
+
+    @abstractmethod
+    def model_info(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> dict:
+        """
+        Given a model name returns a dict-like (OmegaConf) object describing it.
+        Uses the exact format as the omegaconf stanza.
+        """
+        pass
+
+    @abstractmethod
+    def list_models(self, base_model: Optional[BaseModelType] = None, model_type: Optional[ModelType] = None) -> dict:
+        """
+        Return a dict of models in the format:
+        { model_type1:
+          { model_name1: {'status': 'active'|'cached'|'not loaded',
+                         'model_name' : name,
+                         'model_type' : SDModelType,
+                         'description': description,
+                         'format': 'folder'|'safetensors'|'ckpt'
+                         },
+            model_name2: { etc }
+          },
+          model_type2:
+            { model_name_n: etc
+        }
+        """
+        pass
+
+    @abstractmethod
+    def list_model(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> dict:
+        """
+        Return information about the model using the same format as list_models()
+        """
+        pass
+
+    @abstractmethod
+    def model_names(self) -> List[Tuple[str, BaseModelType, ModelType]]:
+        """
+        Returns a list of all the model names known.
+        """
+        pass
+
+    @abstractmethod
+    def add_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        model_attributes: dict,
+        clobber: bool = False,
+    ) -> AddModelResult:
+        """
+        Update the named model with a dictionary of attributes. Will fail with an
+        assertion error if the name already exists. Pass clobber=True to overwrite.
+        On a successful update, the config will be changed in memory. Will fail
+        with an assertion error if provided attributes are incorrect or
+        the model name is missing. Call commit() to write changes to disk.
+        """
+        pass
+
+    @abstractmethod
+    def update_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        model_attributes: dict,
+    ) -> AddModelResult:
+        """
+        Update the named model with a dictionary of attributes. Will fail with a
+        ModelNotFoundException if the name does not already exist.
+
+        On a successful update, the config will be changed in memory. Will fail
+        with an assertion error if provided attributes are incorrect or
+        the model name is missing. Call commit() to write changes to disk.
+        """
+        pass
+
+    @abstractmethod
+    def del_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+    ):
+        """
+        Delete the named model from configuration. If delete_files is true,
+        then the underlying weight file or diffusers directory will be deleted
+        as well. Call commit() to write to disk.
+        """
+        pass
+
+    @abstractmethod
+    def rename_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        new_name: str,
+    ):
+        """
+        Rename the indicated model.
+        """
+        pass
+
+    @abstractmethod
+    def list_checkpoint_configs(self) -> List[Path]:
+        """
+        List the checkpoint config paths from ROOT/configs/stable-diffusion.
+        """
+        pass
+
+    @abstractmethod
+    def convert_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: Literal[ModelType.Main, ModelType.Vae],
+    ) -> AddModelResult:
+        """
+        Convert a checkpoint file into a diffusers folder, deleting the cached
+        version and deleting the original checkpoint file if it is in the models
+        directory.
+        :param model_name: Name of the model to convert
+        :param base_model: Base model type
+        :param model_type: Type of model ['vae' or 'main']
+
+        This will raise a ValueError unless the model is not a checkpoint. It will
+        also raise a ValueError in the event that there is a similarly-named diffusers
+        directory already in place.
+        """
+        pass
+
+    @abstractmethod
+    def heuristic_import(
+        self,
+        items_to_import: set[str],
+        prediction_type_helper: Optional[Callable[[Path], SchedulerPredictionType]] = None,
+    ) -> dict[str, AddModelResult]:
+        """Import a list of paths, repo_ids or URLs. Returns the set of
+        successfully imported items.
+        :param items_to_import: Set of strings corresponding to models to be imported.
+        :param prediction_type_helper: A callback that receives the Path of a Stable Diffusion 2 checkpoint model and returns a SchedulerPredictionType.
+
+        The prediction type helper is necessary to distinguish between
+        models based on Stable Diffusion 2 Base (requiring
+        SchedulerPredictionType.Epsilson) and Stable Diffusion 768
+        (requiring SchedulerPredictionType.VPrediction). It is
+        generally impossible to do this programmatically, so the
+        prediction_type_helper usually asks the user to choose.
+
+        The result is a set of successfully installed models. Each element
+        of the set is a dict corresponding to the newly-created OmegaConf stanza for
+        that model.
+        """
+        pass
+
+    @abstractmethod
+    def merge_models(
+        self,
+        model_names: List[str] = Field(
+            default=None, min_length=2, max_length=3, description="List of model names to merge"
+        ),
+        base_model: Union[BaseModelType, str] = Field(
+            default=None, description="Base model shared by all models to be merged"
+        ),
+        merged_model_name: str = Field(default=None, description="Name of destination model after merging"),
+        alpha: Optional[float] = 0.5,
+        interp: Optional[MergeInterpolationMethod] = None,
+        force: Optional[bool] = False,
+        merge_dest_directory: Optional[Path] = None,
+    ) -> AddModelResult:
+        """
+        Merge two to three diffusrs pipeline models and save as a new model.
+        :param model_names: List of 2-3 models to merge
+        :param base_model: Base model to use for all models
+        :param merged_model_name: Name of destination merged model
+        :param alpha: Alpha strength to apply to 2d and 3d model
+        :param interp: Interpolation method. None (default)
+        :param merge_dest_directory: Save the merged model to the designated directory (with 'merged_model_name' appended)
+        """
+        pass
+
+    @abstractmethod
+    def search_for_models(self, directory: Path) -> List[Path]:
+        """
+        Return list of all models found in the designated directory.
+        """
+        pass
+
+    @abstractmethod
+    def sync_to_config(self):
+        """
+        Re-read models.yaml, rescan the models directory, and reimport models
+        in the autoimport directories. Call after making changes outside the
+        model manager API.
+        """
+        pass
+
+    @abstractmethod
+    def collect_cache_stats(self, cache_stats: CacheStats):
+        """
+        Reset model cache statistics for graph with graph_id.
+        """
+        pass
+
+    @abstractmethod
+    def commit(self, conf_file: Optional[Path] = None) -> None:
+        """
+        Write current configuration out to the indicated file.
+        If no conf_file is provided, then replaces the
+        original file/database used to initialize the object.
+        """
         pass

invokeai/app/services/model_manager/model_manager_default.py

@@ -1,155 +1,413 @@
 # Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
-"""Implementation of ModelManagerServiceBase."""
 
-from typing import Optional
+from __future__ import annotations
+
+from logging import Logger
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
 
 import torch
-from typing_extensions import Self
+from pydantic import Field
 
-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 invokeai.app.services.config.config_default import InvokeAIAppConfig
+from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException
+from invokeai.backend.model_management import (
+    AddModelResult,
+    BaseModelType,
+    MergeInterpolationMethod,
+    ModelInfo,
+    ModelManager,
+    ModelMerger,
+    ModelNotFoundException,
+    ModelType,
+    SchedulerPredictionType,
+    SubModelType,
+)
+from invokeai.backend.model_management.model_cache import CacheStats
+from invokeai.backend.model_management.model_search import FindModels
+from invokeai.backend.util import choose_precision, choose_torch_device
 
-from ..config import InvokeAIAppConfig
-from ..download import DownloadQueueServiceBase
-from ..events.events_base import EventServiceBase
-from ..model_install import ModelInstallService, ModelInstallServiceBase
-from ..model_load import ModelLoadService, ModelLoadServiceBase
-from ..model_records import ModelRecordServiceBase, UnknownModelException
 from .model_manager_base import ModelManagerServiceBase
 
+if TYPE_CHECKING:
+    from invokeai.app.invocations.baseinvocation import InvocationContext
 
+
+# simple implementation
 class ModelManagerService(ModelManagerServiceBase):
-    """
-    The ModelManagerService handles various aspects of model installation, maintenance and loading.
-
-    It bundles three distinct services:
-    model_manager.store   -- Routines to manage the database of model configuration records.
-    model_manager.install -- Routines to install, move and delete models.
-    model_manager.load    -- Routines to load models into memory.
-    """
+    """Responsible for managing models on disk and in memory"""
 
     def __init__(
         self,
-        store: ModelRecordServiceBase,
-        install: ModelInstallServiceBase,
-        load: ModelLoadServiceBase,
+        config: InvokeAIAppConfig,
+        logger: Logger,
     ):
-        self._store = store
-        self._install = install
-        self._load = load
+        """
+        Initialize with the path to the models.yaml config file.
+        Optional parameters are the torch device type, precision, max_models,
+        and sequential_offload boolean. Note that the default device
+        type and precision are set up for a CUDA system running at half precision.
+        """
+        if config.model_conf_path and config.model_conf_path.exists():
+            config_file = config.model_conf_path
+        else:
+            config_file = config.root_dir / "configs/models.yaml"
 
-    @property
-    def store(self) -> ModelRecordServiceBase:
-        return self._store
+        logger.debug(f"Config file={config_file}")
 
-    @property
-    def install(self) -> ModelInstallServiceBase:
-        return self._install
+        device = torch.device(choose_torch_device())
+        device_name = torch.cuda.get_device_name() if device == torch.device("cuda") else ""
+        logger.info(f"GPU device = {device} {device_name}")
 
-    @property
-    def load(self) -> ModelLoadServiceBase:
-        return self._load
+        precision = config.precision
+        if precision == "auto":
+            precision = choose_precision(device)
+        dtype = torch.float32 if precision == "float32" else torch.float16
 
-    def start(self, invoker: Invoker) -> None:
-        for service in [self._store, self._install, self._load]:
-            if hasattr(service, "start"):
-                service.start(invoker)
+        # this is transitional backward compatibility
+        # support for the deprecated `max_loaded_models`
+        # configuration value. If present, then the
+        # cache size is set to 2.5 GB times
+        # the number of max_loaded_models. Otherwise
+        # use new `ram_cache_size` config setting
+        max_cache_size = config.ram_cache_size
 
-    def stop(self, invoker: Invoker) -> None:
-        for service in [self._store, self._install, self._load]:
-            if hasattr(service, "stop"):
-                service.stop(invoker)
+        logger.debug(f"Maximum RAM cache size: {max_cache_size} GiB")
 
-    def load_model_by_config(
-        self,
-        model_config: AnyModelConfig,
-        submodel_type: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
-        return self.load.load_model(model_config, submodel_type, context_data)
+        sequential_offload = config.sequential_guidance
 
-    def load_model_by_key(
-        self,
-        key: str,
-        submodel_type: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
-        config = self.store.get_model(key)
-        return self.load.load_model(config, submodel_type, context_data)
+        self.mgr = ModelManager(
+            config=config_file,
+            device_type=device,
+            precision=dtype,
+            max_cache_size=max_cache_size,
+            sequential_offload=sequential_offload,
+            logger=logger,
+        )
+        logger.info("Model manager service initialized")
 
-    def load_model_by_attr(
+    def get_model(
         self,
         model_name: str,
         base_model: BaseModelType,
         model_type: ModelType,
         submodel: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModel:
+        context: Optional[InvocationContext] = None,
+    ) -> ModelInfo:
         """
-        Given a model's attributes, search the database for it, and if found, load and return the LoadedModel object.
-
-        This is provided for API compatability with the get_model() method
-        in the original model manager. However, note that LoadedModel is
-        not the same as the original ModelInfo that ws returned.
-
-        :param model_name: Name of to be fetched.
-        :param base_model: Base model
-        :param model_type: Type of the model
-        :param submodel: For main (pipeline models), the submodel to fetch
-        :param context: The invocation context.
-
-        Exceptions: UnknownModelException -- model with this key not known
-                    NotImplementedException -- a model loader was not provided at initialization time
-                    ValueError -- more than one model matches this combination
+        Retrieve the indicated model. submodel can be used to get a
+        part (such as the vae) of a diffusers mode.
         """
-        configs = self.store.search_by_attr(model_name, base_model, model_type)
-        if len(configs) == 0:
-            raise UnknownModelException(f"{base_model}/{model_type}/{model_name}: Unknown model")
-        elif len(configs) > 1:
-            raise ValueError(f"{base_model}/{model_type}/{model_name}: More than one model matches.")
+
+        # we can emit model loading events if we are executing with access to the invocation context
+        if context:
+            self._emit_load_event(
+                context=context,
+                model_name=model_name,
+                base_model=base_model,
+                model_type=model_type,
+                submodel=submodel,
+            )
+
+        model_info = self.mgr.get_model(
+            model_name,
+            base_model,
+            model_type,
+            submodel,
+        )
+
+        if context:
+            self._emit_load_event(
+                context=context,
+                model_name=model_name,
+                base_model=base_model,
+                model_type=model_type,
+                submodel=submodel,
+                model_info=model_info,
+            )
+
+        return model_info
+
+    def model_exists(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+    ) -> bool:
+        """
+        Given a model name, returns True if it is a valid
+        identifier.
+        """
+        return self.mgr.model_exists(
+            model_name,
+            base_model,
+            model_type,
+        )
+
+    def model_info(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> Union[dict, None]:
+        """
+        Given a model name returns a dict-like (OmegaConf) object describing it.
+        """
+        return self.mgr.model_info(model_name, base_model, model_type)
+
+    def model_names(self) -> List[Tuple[str, BaseModelType, ModelType]]:
+        """
+        Returns a list of all the model names known.
+        """
+        return self.mgr.model_names()
+
+    def list_models(
+        self, base_model: Optional[BaseModelType] = None, model_type: Optional[ModelType] = None
+    ) -> list[dict]:
+        """
+        Return a list of models.
+        """
+        return self.mgr.list_models(base_model, model_type)
+
+    def list_model(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> Union[dict, None]:
+        """
+        Return information about the model using the same format as list_models()
+        """
+        return self.mgr.list_model(model_name=model_name, base_model=base_model, model_type=model_type)
+
+    def add_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        model_attributes: dict,
+        clobber: bool = False,
+    ) -> AddModelResult:
+        """
+        Update the named model with a dictionary of attributes. Will fail with an
+        assertion error if the name already exists. Pass clobber=True to overwrite.
+        On a successful update, the config will be changed in memory. Will fail
+        with an assertion error if provided attributes are incorrect or
+        the model name is missing. Call commit() to write changes to disk.
+        """
+        self.logger.debug(f"add/update model {model_name}")
+        return self.mgr.add_model(model_name, base_model, model_type, model_attributes, clobber)
+
+    def update_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        model_attributes: dict,
+    ) -> AddModelResult:
+        """
+        Update the named model with a dictionary of attributes. Will fail with a
+        ModelNotFoundException exception if the name does not already exist.
+        On a successful update, the config will be changed in memory. Will fail
+        with an assertion error if provided attributes are incorrect or
+        the model name is missing. Call commit() to write changes to disk.
+        """
+        self.logger.debug(f"update model {model_name}")
+        if not self.model_exists(model_name, base_model, model_type):
+            raise ModelNotFoundException(f"Unknown model {model_name}")
+        return self.add_model(model_name, base_model, model_type, model_attributes, clobber=True)
+
+    def del_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+    ):
+        """
+        Delete the named model from configuration. If delete_files is true,
+        then the underlying weight file or diffusers directory will be deleted
+        as well.
+        """
+        self.logger.debug(f"delete model {model_name}")
+        self.mgr.del_model(model_name, base_model, model_type)
+        self.mgr.commit()
+
+    def convert_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: Literal[ModelType.Main, ModelType.Vae],
+        convert_dest_directory: Optional[Path] = Field(
+            default=None, description="Optional directory location for merged model"
+        ),
+    ) -> AddModelResult:
+        """
+        Convert a checkpoint file into a diffusers folder, deleting the cached
+        version and deleting the original checkpoint file if it is in the models
+        directory.
+        :param model_name: Name of the model to convert
+        :param base_model: Base model type
+        :param model_type: Type of model ['vae' or 'main']
+        :param convert_dest_directory: Save the converted model to the designated directory (`models/etc/etc` by default)
+
+        This will raise a ValueError unless the model is not a checkpoint. It will
+        also raise a ValueError in the event that there is a similarly-named diffusers
+        directory already in place.
+        """
+        self.logger.debug(f"convert model {model_name}")
+        return self.mgr.convert_model(model_name, base_model, model_type, convert_dest_directory)
+
+    def collect_cache_stats(self, cache_stats: CacheStats):
+        """
+        Reset model cache statistics for graph with graph_id.
+        """
+        self.mgr.cache.stats = cache_stats
+
+    def commit(self, conf_file: Optional[Path] = None):
+        """
+        Write current configuration out to the indicated file.
+        If no conf_file is provided, then replaces the
+        original file/database used to initialize the object.
+        """
+        return self.mgr.commit(conf_file)
+
+    def _emit_load_event(
+        self,
+        context: InvocationContext,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        submodel: Optional[SubModelType] = None,
+        model_info: Optional[ModelInfo] = None,
+    ):
+        if context.services.queue.is_canceled(context.graph_execution_state_id):
+            raise CanceledException()
+
+        if model_info:
+            context.services.events.emit_model_load_completed(
+                queue_id=context.queue_id,
+                queue_item_id=context.queue_item_id,
+                queue_batch_id=context.queue_batch_id,
+                graph_execution_state_id=context.graph_execution_state_id,
+                model_name=model_name,
+                base_model=base_model,
+                model_type=model_type,
+                submodel=submodel,
+                model_info=model_info,
+            )
         else:
-            return self.load.load_model(configs[0], submodel, context_data)
+            context.services.events.emit_model_load_started(
+                queue_id=context.queue_id,
+                queue_item_id=context.queue_item_id,
+                queue_batch_id=context.queue_batch_id,
+                graph_execution_state_id=context.graph_execution_state_id,
+                model_name=model_name,
+                base_model=base_model,
+                model_type=model_type,
+                submodel=submodel,
+            )
 
-    @classmethod
-    def build_model_manager(
-        cls,
-        app_config: InvokeAIAppConfig,
-        model_record_service: ModelRecordServiceBase,
-        download_queue: DownloadQueueServiceBase,
-        events: EventServiceBase,
-        execution_device: torch.device = choose_torch_device(),
-    ) -> Self:
+    @property
+    def logger(self):
+        return self.mgr.logger
+
+    def heuristic_import(
+        self,
+        items_to_import: set[str],
+        prediction_type_helper: Optional[Callable[[Path], SchedulerPredictionType]] = None,
+    ) -> dict[str, AddModelResult]:
+        """Import a list of paths, repo_ids or URLs. Returns the set of
+        successfully imported items.
+        :param items_to_import: Set of strings corresponding to models to be imported.
+        :param prediction_type_helper: A callback that receives the Path of a Stable Diffusion 2 checkpoint model and returns a SchedulerPredictionType.
+
+        The prediction type helper is necessary to distinguish between
+        models based on Stable Diffusion 2 Base (requiring
+        SchedulerPredictionType.Epsilson) and Stable Diffusion 768
+        (requiring SchedulerPredictionType.VPrediction). It is
+        generally impossible to do this programmatically, so the
+        prediction_type_helper usually asks the user to choose.
+
+        The result is a set of successfully installed models. Each element
+        of the set is a dict corresponding to the newly-created OmegaConf stanza for
+        that model.
         """
-        Construct the model manager service instance.
+        return self.mgr.heuristic_import(items_to_import, prediction_type_helper)
 
-        For simplicity, use this class method rather than the __init__ constructor.
+    def merge_models(
+        self,
+        model_names: List[str] = Field(
+            default=None, min_length=2, max_length=3, description="List of model names to merge"
+        ),
+        base_model: Union[BaseModelType, str] = Field(
+            default=None, description="Base model shared by all models to be merged"
+        ),
+        merged_model_name: str = Field(default=None, description="Name of destination model after merging"),
+        alpha: float = 0.5,
+        interp: Optional[MergeInterpolationMethod] = None,
+        force: bool = False,
+        merge_dest_directory: Optional[Path] = Field(
+            default=None, description="Optional directory location for merged model"
+        ),
+    ) -> AddModelResult:
         """
-        logger = InvokeAILogger.get_logger(cls.__name__)
-        logger.setLevel(app_config.log_level.upper())
+        Merge two to three diffusrs pipeline models and save as a new model.
+        :param model_names: List of 2-3 models to merge
+        :param base_model: Base model to use for all models
+        :param merged_model_name: Name of destination merged model
+        :param alpha: Alpha strength to apply to 2d and 3d model
+        :param interp: Interpolation method. None (default)
+        :param merge_dest_directory: Save the merged model to the designated directory (with 'merged_model_name' appended)
+        """
+        merger = ModelMerger(self.mgr)
+        try:
+            result = merger.merge_diffusion_models_and_save(
+                model_names=model_names,
+                base_model=base_model,
+                merged_model_name=merged_model_name,
+                alpha=alpha,
+                interp=interp,
+                force=force,
+                merge_dest_directory=merge_dest_directory,
+            )
+        except AssertionError as e:
+            raise ValueError(e)
+        return result
 
-        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,
+    def search_for_models(self, directory: Path) -> List[Path]:
+        """
+        Return list of all models found in the designated directory.
+        """
+        search = FindModels([directory], self.logger)
+        return search.list_models()
+
+    def sync_to_config(self):
+        """
+        Re-read models.yaml, rescan the models directory, and reimport models
+        in the autoimport directories. Call after making changes outside the
+        model manager API.
+        """
+        return self.mgr.sync_to_config()
+
+    def list_checkpoint_configs(self) -> List[Path]:
+        """
+        List the checkpoint config paths from ROOT/configs/stable-diffusion.
+        """
+        config = self.mgr.app_config
+        conf_path = config.legacy_conf_path
+        root_path = config.root_path
+        return [(conf_path / x).relative_to(root_path) for x in conf_path.glob("**/*.yaml")]
+
+    def rename_model(
+        self,
+        model_name: str,
+        base_model: BaseModelType,
+        model_type: ModelType,
+        new_name: Optional[str] = None,
+        new_base: Optional[BaseModelType] = None,
+    ):
+        """
+        Rename the indicated model. Can provide a new name and/or a new base.
+        :param model_name: Current name of the model
+        :param base_model: Current base of the model
+        :param model_type: Model type (can't be changed)
+        :param new_name: New name for the model
+        :param new_base: New base for the model
+        """
+        self.mgr.rename_model(
+            base_model=base_model,
+            model_type=model_type,
+            model_name=model_name,
+            new_name=new_name,
+            new_base=new_base,
         )
-        convert_cache = ModelConvertCache(
-            cache_path=app_config.models_convert_cache_path, max_size=app_config.convert_cache_size
-        )
-        loader = ModelLoadService(
-            app_config=app_config,
-            ram_cache=ram_cache,
-            convert_cache=convert_cache,
-            registry=ModelLoaderRegistry,
-        )
-        installer = ModelInstallService(
-            app_config=app_config,
-            record_store=model_record_service,
-            download_queue=download_queue,
-            event_bus=events,
-        )
-        return cls(store=model_record_service, install=installer, load=loader)

invokeai/app/services/model_metadata/__init__.py

@@ -1,9 +0,0 @@
-"""Init file for ModelMetadataStoreService module."""
-
-from .metadata_store_base import ModelMetadataStoreBase
-from .metadata_store_sql import ModelMetadataStoreSQL
-
-__all__ = [
-    "ModelMetadataStoreBase",
-    "ModelMetadataStoreSQL",
-]

invokeai/app/services/model_metadata/metadata_store_base.py

@@ -1,65 +0,0 @@
-# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
-"""
-Storage for Model Metadata
-"""
-
-from abc import ABC, abstractmethod
-from typing import List, Set, Tuple
-
-from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
-
-
-class ModelMetadataStoreBase(ABC):
-    """Store, search and fetch model metadata retrieved from remote repositories."""
-
-    @abstractmethod
-    def add_metadata(self, model_key: str, metadata: AnyModelRepoMetadata) -> None:
-        """
-        Add a block of repo metadata to a model record.
-
-        The model record config must already exist in the database with the
-        same key. Otherwise a FOREIGN KEY constraint exception will be raised.
-
-        :param model_key: Existing model key in the `model_config` table
-        :param metadata: ModelRepoMetadata object to store
-        """
-
-    @abstractmethod
-    def get_metadata(self, model_key: str) -> AnyModelRepoMetadata:
-        """Retrieve the ModelRepoMetadata corresponding to model key."""
-
-    @abstractmethod
-    def list_all_metadata(self) -> List[Tuple[str, AnyModelRepoMetadata]]:  # key, metadata
-        """Dump out all the metadata."""
-
-    @abstractmethod
-    def update_metadata(self, model_key: str, metadata: AnyModelRepoMetadata) -> AnyModelRepoMetadata:
-        """
-        Update metadata corresponding to the model with the indicated key.
-
-        :param model_key: Existing model key in the `model_config` table
-        :param metadata: ModelRepoMetadata object to update
-        """
-
-    @abstractmethod
-    def list_tags(self) -> Set[str]:
-        """Return all tags in the tags table."""
-
-    @abstractmethod
-    def search_by_tag(self, tags: Set[str]) -> Set[str]:
-        """Return the keys of models containing all of the listed tags."""
-
-    @abstractmethod
-    def search_by_author(self, author: str) -> Set[str]:
-        """Return the keys of models authored by the indicated author."""
-
-    @abstractmethod
-    def search_by_name(self, name: str) -> Set[str]:
-        """
-        Return the keys of models with the indicated name.
-
-        Note that this is the name of the model given to it by
-        the remote source. The user may have changed the local
-        name. The local name will be located in the model config
-        record object.
-        """

invokeai/app/services/model_metadata/metadata_store_sql.py

@@ -1,222 +0,0 @@
-# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
-"""
-SQL Storage for Model Metadata
-"""
-
-import sqlite3
-from typing import List, Optional, Set, Tuple
-
-from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
-from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata, UnknownMetadataException
-from invokeai.backend.model_manager.metadata.fetch import ModelMetadataFetchBase
-
-from .metadata_store_base import ModelMetadataStoreBase
-
-
-class ModelMetadataStoreSQL(ModelMetadataStoreBase):
-    """Store, search and fetch model metadata retrieved from remote repositories."""
-
-    def __init__(self, db: SqliteDatabase):
-        """
-        Initialize a new object from preexisting sqlite3 connection and threading lock objects.
-
-        :param conn: sqlite3 connection object
-        :param lock: threading Lock object
-        """
-        super().__init__()
-        self._db = db
-        self._cursor = self._db.conn.cursor()
-
-    def add_metadata(self, model_key: str, metadata: AnyModelRepoMetadata) -> None:
-        """
-        Add a block of repo metadata to a model record.
-
-        The model record config must already exist in the database with the
-        same key. Otherwise a FOREIGN KEY constraint exception will be raised.
-
-        :param model_key: Existing model key in the `model_config` table
-        :param metadata: ModelRepoMetadata object to store
-        """
-        json_serialized = metadata.model_dump_json()
-        with self._db.lock:
-            try:
-                self._cursor.execute(
-                    """--sql
-                    INSERT INTO model_metadata(
-                       id,
-                       metadata
-                    )
-                    VALUES (?,?);
-                    """,
-                    (
-                        model_key,
-                        json_serialized,
-                    ),
-                )
-                self._update_tags(model_key, metadata.tags)
-                self._db.conn.commit()
-            except sqlite3.IntegrityError as excp:  # FOREIGN KEY error: the key was not in model_config table
-                self._db.conn.rollback()
-                raise UnknownMetadataException from excp
-            except sqlite3.Error as excp:
-                self._db.conn.rollback()
-                raise excp
-
-    def get_metadata(self, model_key: str) -> AnyModelRepoMetadata:
-        """Retrieve the ModelRepoMetadata corresponding to model key."""
-        with self._db.lock:
-            self._cursor.execute(
-                """--sql
-                SELECT metadata FROM model_metadata
-                WHERE id=?;
-                """,
-                (model_key,),
-            )
-            rows = self._cursor.fetchone()
-            if not rows:
-                raise UnknownMetadataException("model metadata not found")
-            return ModelMetadataFetchBase.from_json(rows[0])
-
-    def list_all_metadata(self) -> List[Tuple[str, AnyModelRepoMetadata]]:  # key, metadata
-        """Dump out all the metadata."""
-        with self._db.lock:
-            self._cursor.execute(
-                """--sql
-                SELECT id,metadata FROM model_metadata;
-                """,
-                (),
-            )
-            rows = self._cursor.fetchall()
-        return [(x[0], ModelMetadataFetchBase.from_json(x[1])) for x in rows]
-
-    def update_metadata(self, model_key: str, metadata: AnyModelRepoMetadata) -> AnyModelRepoMetadata:
-        """
-        Update metadata corresponding to the model with the indicated key.
-
-        :param model_key: Existing model key in the `model_config` table
-        :param metadata: ModelRepoMetadata object to update
-        """
-        json_serialized = metadata.model_dump_json()  # turn it into a json string.
-        with self._db.lock:
-            try:
-                self._cursor.execute(
-                    """--sql
-                    UPDATE model_metadata
-                    SET
-                        metadata=?
-                    WHERE id=?;
-                    """,
-                    (json_serialized, model_key),
-                )
-                if self._cursor.rowcount == 0:
-                    raise UnknownMetadataException("model metadata not found")
-                self._update_tags(model_key, metadata.tags)
-                self._db.conn.commit()
-            except sqlite3.Error as e:
-                self._db.conn.rollback()
-                raise e
-
-        return self.get_metadata(model_key)
-
-    def list_tags(self) -> Set[str]:
-        """Return all tags in the tags table."""
-        self._cursor.execute(
-            """--sql
-            select tag_text from tags;
-            """
-        )
-        return {x[0] for x in self._cursor.fetchall()}
-
-    def search_by_tag(self, tags: Set[str]) -> Set[str]:
-        """Return the keys of models containing all of the listed tags."""
-        with self._db.lock:
-            try:
-                matches: Optional[Set[str]] = None
-                for tag in tags:
-                    self._cursor.execute(
-                        """--sql
-                        SELECT a.model_id FROM model_tags AS a,
-                                                     tags AS b
-                        WHERE a.tag_id=b.tag_id
-                          AND b.tag_text=?;
-                        """,
-                        (tag,),
-                    )
-                    model_keys = {x[0] for x in self._cursor.fetchall()}
-                    if matches is None:
-                        matches = model_keys
-                    matches = matches.intersection(model_keys)
-            except sqlite3.Error as e:
-                raise e
-        return matches if matches else set()
-
-    def search_by_author(self, author: str) -> Set[str]:
-        """Return the keys of models authored by the indicated author."""
-        self._cursor.execute(
-            """--sql
-            SELECT id FROM model_metadata
-            WHERE author=?;
-            """,
-            (author,),
-        )
-        return {x[0] for x in self._cursor.fetchall()}
-
-    def search_by_name(self, name: str) -> Set[str]:
-        """
-        Return the keys of models with the indicated name.
-
-        Note that this is the name of the model given to it by
-        the remote source. The user may have changed the local
-        name. The local name will be located in the model config
-        record object.
-        """
-        self._cursor.execute(
-            """--sql
-            SELECT id FROM model_metadata
-            WHERE name=?;
-            """,
-            (name,),
-        )
-        return {x[0] for x in self._cursor.fetchall()}
-
-    def _update_tags(self, model_key: str, tags: Set[str]) -> None:
-        """Update tags for the model referenced by model_key."""
-        # remove previous tags from this model
-        self._cursor.execute(
-            """--sql
-            DELETE FROM model_tags
-            WHERE model_id=?;
-            """,
-            (model_key,),
-        )
-
-        for tag in tags:
-            self._cursor.execute(
-                """--sql
-                INSERT OR IGNORE INTO tags (
-                  tag_text
-                  )
-                VALUES (?);
-                """,
-                (tag,),
-            )
-            self._cursor.execute(
-                """--sql
-                SELECT tag_id
-                FROM tags
-                WHERE tag_text = ?
-                LIMIT 1;
-                """,
-                (tag,),
-            )
-            tag_id = self._cursor.fetchone()[0]
-            self._cursor.execute(
-                """--sql
-                INSERT OR IGNORE INTO model_tags (
-                   model_id,
-                   tag_id
-                  )
-                VALUES (?,?);
-                """,
-                (model_key, tag_id),
-            )

invokeai/app/services/model_records/model_records_base.py

@@ -11,15 +11,8 @@ from typing import Any, Dict, List, Optional, Set, Tuple, Union
 from pydantic import BaseModel, Field
 
 from invokeai.app.services.shared.pagination import PaginatedResults
-from invokeai.backend.model_manager import (
-    AnyModelConfig,
-    BaseModelType,
-    ModelFormat,
-    ModelType,
-)
-from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
-
-from ..model_metadata import ModelMetadataStoreBase
+from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelFormat, ModelType
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata, ModelMetadataStore
 
 
 class DuplicateModelException(Exception):
@@ -111,7 +104,7 @@ class ModelRecordServiceBase(ABC):
 
     @property
     @abstractmethod
-    def metadata_store(self) -> ModelMetadataStoreBase:
+    def metadata_store(self) -> ModelMetadataStore:
         """Return a ModelMetadataStore initialized on the same database."""
         pass
 
@@ -153,7 +146,7 @@ class ModelRecordServiceBase(ABC):
     @abstractmethod
     def exists(self, key: str) -> bool:
         """
-        Return True if a model with the indicated key exists in the database.
+        Return True if a model with the indicated key exists in the databse.
 
         :param key: Unique key for the model to be deleted
         """

invokeai/app/services/model_records/model_records_sql.py

@@ -54,9 +54,8 @@ from invokeai.backend.model_manager.config import (
     ModelFormat,
     ModelType,
 )
-from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata, UnknownMetadataException
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata, ModelMetadataStore, UnknownMetadataException
 
-from ..model_metadata import ModelMetadataStoreBase, ModelMetadataStoreSQL
 from ..shared.sqlite.sqlite_database import SqliteDatabase
 from .model_records_base import (
     DuplicateModelException,
@@ -70,16 +69,16 @@ from .model_records_base import (
 class ModelRecordServiceSQL(ModelRecordServiceBase):
     """Implementation of the ModelConfigStore ABC using a SQL database."""
 
-    def __init__(self, db: SqliteDatabase, metadata_store: ModelMetadataStoreBase):
+    def __init__(self, db: SqliteDatabase):
         """
         Initialize a new object from preexisting sqlite3 connection and threading lock objects.
 
-        :param db: Sqlite connection object
+        :param conn: sqlite3 connection object
+        :param lock: threading Lock object
         """
         super().__init__()
         self._db = db
-        self._cursor = db.conn.cursor()
-        self._metadata_store = metadata_store
+        self._cursor = self._db.conn.cursor()
 
     @property
     def db(self) -> SqliteDatabase:
@@ -159,7 +158,7 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
                 self._db.conn.rollback()
                 raise e
 
-    def update_model(self, key: str, config: Union[Dict[str, Any], AnyModelConfig]) -> AnyModelConfig:
+    def update_model(self, key: str, config: Union[dict, AnyModelConfig]) -> AnyModelConfig:
         """
         Update the model, returning the updated version.
 
@@ -200,7 +199,7 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
         with self._db.lock:
             self._cursor.execute(
                 """--sql
-                SELECT config, strftime('%s',updated_at) FROM model_config
+                SELECT config FROM model_config
                 WHERE id=?;
                 """,
                 (key,),
@@ -208,7 +207,7 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
             rows = self._cursor.fetchone()
             if not rows:
                 raise UnknownModelException("model not found")
-            model = ModelConfigFactory.make_config(json.loads(rows[0]), timestamp=rows[1])
+            model = ModelConfigFactory.make_config(json.loads(rows[0]))
         return model
 
     def exists(self, key: str) -> bool:
@@ -266,14 +265,12 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
         with self._db.lock:
             self._cursor.execute(
                 f"""--sql
-                select config, strftime('%s',updated_at) FROM model_config
+                select config FROM model_config
                 {where};
                 """,
                 tuple(bindings),
             )
-            results = [
-                ModelConfigFactory.make_config(json.loads(x[0]), timestamp=x[1]) for x in self._cursor.fetchall()
-            ]
+            results = [ModelConfigFactory.make_config(json.loads(x[0])) for x in self._cursor.fetchall()]
         return results
 
     def search_by_path(self, path: Union[str, Path]) -> List[AnyModelConfig]:
@@ -282,14 +279,12 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
         with self._db.lock:
             self._cursor.execute(
                 """--sql
-                SELECT config, strftime('%s',updated_at) FROM model_config
+                SELECT config FROM model_config
                 WHERE path=?;
                 """,
                 (str(path),),
             )
-            results = [
-                ModelConfigFactory.make_config(json.loads(x[0]), timestamp=x[1]) for x in self._cursor.fetchall()
-            ]
+            results = [ModelConfigFactory.make_config(json.loads(x[0])) for x in self._cursor.fetchall()]
         return results
 
     def search_by_hash(self, hash: str) -> List[AnyModelConfig]:
@@ -298,20 +293,18 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
         with self._db.lock:
             self._cursor.execute(
                 """--sql
-                SELECT config, strftime('%s',updated_at) FROM model_config
+                SELECT config FROM model_config
                 WHERE original_hash=?;
                 """,
                 (hash,),
             )
-            results = [
-                ModelConfigFactory.make_config(json.loads(x[0]), timestamp=x[1]) for x in self._cursor.fetchall()
-            ]
+            results = [ModelConfigFactory.make_config(json.loads(x[0])) for x in self._cursor.fetchall()]
         return results
 
     @property
-    def metadata_store(self) -> ModelMetadataStoreBase:
+    def metadata_store(self) -> ModelMetadataStore:
         """Return a ModelMetadataStore initialized on the same database."""
-        return self._metadata_store
+        return ModelMetadataStore(self._db)
 
     def get_metadata(self, key: str) -> Optional[AnyModelRepoMetadata]:
         """
@@ -332,18 +325,18 @@ class ModelRecordServiceSQL(ModelRecordServiceBase):
 
         :param tags: Set of tags to search for. All tags must be present.
         """
-        store = ModelMetadataStoreSQL(self._db)
+        store = ModelMetadataStore(self._db)
         keys = store.search_by_tag(tags)
         return [self.get_model(x) for x in keys]
 
     def list_tags(self) -> Set[str]:
         """Return a unique set of all the model tags in the metadata database."""
-        store = ModelMetadataStoreSQL(self._db)
+        store = ModelMetadataStore(self._db)
         return store.list_tags()
 
     def list_all_metadata(self) -> List[Tuple[str, AnyModelRepoMetadata]]:
         """List metadata for all models that have it."""
-        store = ModelMetadataStoreSQL(self._db)
+        store = ModelMetadataStore(self._db)
         return store.list_all_metadata()
 
     def list_models(

invokeai/app/services/object_serializer/object_serializer_base.py

@@ -1,44 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Callable, Generic, TypeVar
-
-T = TypeVar("T")
-
-
-class ObjectSerializerBase(ABC, Generic[T]):
-    """Saves and loads arbitrary python objects."""
-
-    def __init__(self) -> None:
-        self._on_deleted_callbacks: list[Callable[[str], None]] = []
-
-    @abstractmethod
-    def load(self, name: str) -> T:
-        """
-        Loads the object.
-        :param name: The name of the object to load.
-        :raises ObjectNotFoundError: if the object is not found
-        """
-        pass
-
-    @abstractmethod
-    def save(self, obj: T) -> str:
-        """
-        Saves the object, returning its name.
-        :param obj: The object to save.
-        """
-        pass
-
-    @abstractmethod
-    def delete(self, name: str) -> None:
-        """
-        Deletes the object, if it exists.
-        :param name: The name of the object to delete.
-        """
-        pass
-
-    def on_deleted(self, on_deleted: Callable[[str], None]) -> None:
-        """Register a callback for when an object is deleted"""
-        self._on_deleted_callbacks.append(on_deleted)
-
-    def _on_deleted(self, name: str) -> None:
-        for callback in self._on_deleted_callbacks:
-            callback(name)

invokeai/app/services/object_serializer/object_serializer_common.py

@@ -1,5 +0,0 @@
-class ObjectNotFoundError(KeyError):
-    """Raised when an object is not found while loading"""
-
-    def __init__(self, name: str) -> None:
-        super().__init__(f"Object with name {name} not found")

invokeai/app/services/object_serializer/object_serializer_disk.py

@@ -1,85 +0,0 @@
-import tempfile
-import typing
-from dataclasses import dataclass
-from pathlib import Path
-from typing import TYPE_CHECKING, Optional, TypeVar
-
-import torch
-
-from invokeai.app.services.object_serializer.object_serializer_base import ObjectSerializerBase
-from invokeai.app.services.object_serializer.object_serializer_common import ObjectNotFoundError
-from invokeai.app.util.misc import uuid_string
-
-if TYPE_CHECKING:
-    from invokeai.app.services.invoker import Invoker
-
-
-T = TypeVar("T")
-
-
-@dataclass
-class DeleteAllResult:
-    deleted_count: int
-    freed_space_bytes: float
-
-
-class ObjectSerializerDisk(ObjectSerializerBase[T]):
-    """Disk-backed storage for arbitrary python objects. Serialization is handled by `torch.save` and `torch.load`.
-
-    :param output_dir: The folder where the serialized objects will be stored
-    :param ephemeral: If True, objects will be stored in a temporary directory inside the given output_dir and cleaned up on exit
-    """
-
-    def __init__(self, output_dir: Path, ephemeral: bool = False):
-        super().__init__()
-        self._ephemeral = ephemeral
-        self._base_output_dir = output_dir
-        self._base_output_dir.mkdir(parents=True, exist_ok=True)
-        # Must specify `ignore_cleanup_errors` to avoid fatal errors during cleanup on Windows
-        self._tempdir = (
-            tempfile.TemporaryDirectory(dir=self._base_output_dir, ignore_cleanup_errors=True) if ephemeral else None
-        )
-        self._output_dir = Path(self._tempdir.name) if self._tempdir else self._base_output_dir
-        self.__obj_class_name: Optional[str] = None
-
-    def load(self, name: str) -> T:
-        file_path = self._get_path(name)
-        try:
-            return torch.load(file_path)  # pyright: ignore [reportUnknownMemberType]
-        except FileNotFoundError as e:
-            raise ObjectNotFoundError(name) from e
-
-    def save(self, obj: T) -> str:
-        name = self._new_name()
-        file_path = self._get_path(name)
-        torch.save(obj, file_path)  # pyright: ignore [reportUnknownMemberType]
-        return name
-
-    def delete(self, name: str) -> None:
-        file_path = self._get_path(name)
-        file_path.unlink()
-
-    @property
-    def _obj_class_name(self) -> str:
-        if not self.__obj_class_name:
-            # `__orig_class__` is not available in the constructor for some technical, undoubtedly very pythonic reason
-            self.__obj_class_name = typing.get_args(self.__orig_class__)[0].__name__  # pyright: ignore [reportUnknownMemberType, reportAttributeAccessIssue]
-        return self.__obj_class_name
-
-    def _get_path(self, name: str) -> Path:
-        return self._output_dir / name
-
-    def _new_name(self) -> str:
-        return f"{self._obj_class_name}_{uuid_string()}"
-
-    def _tempdir_cleanup(self) -> None:
-        """Calls `cleanup` on the temporary directory, if it exists."""
-        if self._tempdir:
-            self._tempdir.cleanup()
-
-    def __del__(self) -> None:
-        # In case the service is not properly stopped, clean up the temporary directory when the class instance is GC'd.
-        self._tempdir_cleanup()
-
-    def stop(self, invoker: "Invoker") -> None:
-        self._tempdir_cleanup()

invokeai/app/services/object_serializer/object_serializer_forward_cache.py

@@ -1,65 +0,0 @@
-from queue import Queue
-from typing import TYPE_CHECKING, Optional, TypeVar
-
-from invokeai.app.services.object_serializer.object_serializer_base import ObjectSerializerBase
-
-T = TypeVar("T")
-
-if TYPE_CHECKING:
-    from invokeai.app.services.invoker import Invoker
-
-
-class ObjectSerializerForwardCache(ObjectSerializerBase[T]):
-    """
-    Provides a LRU cache for an instance of `ObjectSerializerBase`.
-    Saving an object to the cache always writes through to the underlying storage.
-    """
-
-    def __init__(self, underlying_storage: ObjectSerializerBase[T], max_cache_size: int = 20):
-        super().__init__()
-        self._underlying_storage = underlying_storage
-        self._cache: dict[str, T] = {}
-        self._cache_ids = Queue[str]()
-        self._max_cache_size = max_cache_size
-
-    def start(self, invoker: "Invoker") -> None:
-        self._invoker = invoker
-        start_op = getattr(self._underlying_storage, "start", None)
-        if callable(start_op):
-            start_op(invoker)
-
-    def stop(self, invoker: "Invoker") -> None:
-        self._invoker = invoker
-        stop_op = getattr(self._underlying_storage, "stop", None)
-        if callable(stop_op):
-            stop_op(invoker)
-
-    def load(self, name: str) -> T:
-        cache_item = self._get_cache(name)
-        if cache_item is not None:
-            return cache_item
-
-        obj = self._underlying_storage.load(name)
-        self._set_cache(name, obj)
-        return obj
-
-    def save(self, obj: T) -> str:
-        name = self._underlying_storage.save(obj)
-        self._set_cache(name, obj)
-        return name
-
-    def delete(self, name: str) -> None:
-        self._underlying_storage.delete(name)
-        if name in self._cache:
-            del self._cache[name]
-        self._on_deleted(name)
-
-    def _get_cache(self, name: str) -> Optional[T]:
-        return None if name not in self._cache else self._cache[name]
-
-    def _set_cache(self, name: str, data: T):
-        if name not in self._cache:
-            self._cache[name] = data
-            self._cache_ids.put(name)
-            if self._cache_ids.qsize() > self._max_cache_size:
-                self._cache.pop(self._cache_ids.get())

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