final tidying before marking PR as ready for review

- Replace AnyModelLoader with ModelLoaderRegistry
- Fix type check errors in multiple files
- Remove apparently unneeded `get_model_config_enum()` method from model manager
- Remove last vestiges of old model manager
- Updated tests and documentation

resolve conflict with seamless.py

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_ (**CURRENTLY UNDER DEVELOPMENT - NOT IMPLEMENTED**)
+  * _ModelLoadServiceBase_
   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_loader/` (**under development**)
+* `invokeai/app/services/model_load/`
 
 Code related to the FastAPI web API can be found in
-`invokeai/app/api/routers/model_records.py`.
+`invokeai/app/api/routers/model_manager_v2.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_record_service`:
+`invokeai.app.services.model_manager.model_records`:
 
 ```
-from invokeai.app.services.model_record_service import ModelType, ModelFormat, BaseModelType
+from invokeai.app.services.model_records 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_record_service` or
+either `invokeai.app.services.model_records` or
 `invokeai.backend.model_manager.config`.
 
 ### ONNXSD2Config
@@ -134,7 +134,7 @@ either `invokeai.app.services.model_record_service` or
 | `upcast_attention`  | bool | Model requires its attention module to be upcast |
 
 The `SchedulerPredictionType` enum can be imported from either
-`invokeai.app.services.model_record_service` or
+`invokeai.app.services.model_records` or
 `invokeai.backend.model_manager.config`.
 
 ### Other config classes
@@ -157,15 +157,6 @@ 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
@@ -177,11 +168,11 @@ initialization and can be retrieved within an invocation from the
 `InvocationContext` object:
 
 ```
-store = context.services.model_record_store
+store = context.services.model_manager.store
 ```
 
 or from elsewhere in the code by accessing
-`ApiDependencies.invoker.services.model_record_store`.
+`ApiDependencies.invoker.services.model_manager.store`.
 
 ### Creating a `ModelRecordService` 
 
@@ -190,7 +181,7 @@ you can directly create either a `ModelRecordServiceSQL` or a
 `ModelRecordServiceFile` object:
 
 ```
-from invokeai.app.services.model_record_service import ModelRecordServiceSQL, ModelRecordServiceFile
+from invokeai.app.services.model_records import ModelRecordServiceSQL, ModelRecordServiceFile
 
 store = ModelRecordServiceSQL.from_connection(connection, lock)
 store = ModelRecordServiceSQL.from_db_file('/path/to/sqlite_database.db')
@@ -252,7 +243,7 @@ So a typical startup pattern would be:
 ```
 import sqlite3
 from invokeai.app.services.thread import lock
-from invokeai.app.services.model_record_service import ModelRecordServiceBase
+from invokeai.app.services.model_records import ModelRecordServiceBase
 from invokeai.app.services.config import InvokeAIAppConfig
 
 config = InvokeAIAppConfig.get_config()
@@ -260,19 +251,6 @@ 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.
@@ -468,6 +446,44 @@ 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
@@ -486,9 +502,10 @@ 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
 
-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
+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
 
 for source in [source1, source2, source3, source4, source5, source6, source7]:
    install_job = installer.install_model(source)
@@ -544,7 +561,6 @@ 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
@@ -737,7 +753,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 confguration and probing
+#### Model configuration and probing
 
 The install service uses the `invokeai.backend.model_manager.probe`
 module during import to determine the model's type, base type, and
@@ -776,6 +792,14 @@ 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`.
@@ -838,6 +862,31 @@ 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
@@ -1128,7 +1177,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,
@@ -1171,6 +1220,13 @@ 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
@@ -1449,9 +1505,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
+***
 
-## Let's get loaded, the lowdown on ModelLoadService
+## The Lowdown on the 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
@@ -1465,7 +1521,7 @@ create alternative instances if you wish.
 ### Creating a ModelLoadService object
 
 The class is defined in
-`invokeai.app.services.model_loader_service`. It is initialized with
+`invokeai.app.services.model_load`. 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
@@ -1475,26 +1531,29 @@ Here is a typical initialization pattern:
 
 ```
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.app.services.model_record_service import ModelRecordServiceBase
-from invokeai.app.services.model_loader_service import ModelLoadService
+from invokeai.app.services.model_load import ModelLoadService, ModelLoaderRegistry
 
 config = InvokeAIAppConfig.get_config()
-store = ModelRecordServiceBase.open(config)
-loader = ModelLoadService(config, store)
+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
+)
 ```
 
-Note that we are relying on the contents of the application
-configuration to choose the implementation of
-`ModelRecordServiceBase`.
+### load_model(model_config, [submodel_type], [context]) -> LoadedModel
 
-### 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
+The `load_model()` method takes an `AnyModelConfig` returned by
+`ModelRecordService.get_model()` and returns the corresponding loaded
 model.  It loads the model into memory, gets the model ready for use,
-and returns a `ModelInfo` object. 
+and returns a `LoadedModel` object.
 
 The optional second argument, `subtype` is a `SubModelType` string
 enum, such as "vae". It is mandatory when used with a main model, and
@@ -1504,46 +1563,45 @@ The optional third argument, `context` can be provided by
 an invocation to trigger model load event reporting. See below for
 details.
 
-The returned `ModelInfo` object shares some fields in common with
-`ModelConfigBase`, but is otherwise a completely different beast:
+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:
 
-| **Field Name** | **Type**        |  **Description** |
+
+| **Attribute Name** | **Type**        |  **Description** |
 |----------------|-----------------|------------------|
-| `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 |
+| `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 |
 
-The types for `ModelInfo` and `SubModelType` can be imported from
-`invokeai.app.services.model_loader_service`.
+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.
 
-To use the model, you use the `ModelInfo` as a context manager using
-the following pattern:
+
+`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:
 
 ```
-model_info = loader.get_model('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
+model_info = loader.get_model_by_key('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
 with model_info as vae:
 	image = vae.decode(latents)[0]
 ```
 
-The `vae` model will stay locked in the GPU during the period of time
-it is in the context manager's scope.
+`get_model_by_key()` may raise any of the following exceptions:
 
-`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
+- `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
   
-** TO DO: ** Resolve discrepancy between ModelInfo.location and
-ModelConfig.path.
-
 ### Emitting model loading events
 
-When the `context` argument is passed to `get_model()`, it will
+When the `context` argument is passed to `load_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
@@ -1556,10 +1614,174 @@ payload=dict(
 	queue_batch_id=queue_batch_id,
 	graph_execution_state_id=graph_execution_state_id,
 	model_key=model_key,
-	submodel=submodel,
+	submodel_type=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,7 +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] if response == "" else response} will be installed.")
 
     console.line()
 

invokeai/app/api/dependencies.py

@@ -2,14 +2,8 @@
 
 from logging import Logger
 
-import torch
-
 from invokeai.app.services.item_storage.item_storage_memory import ItemStorageMemory
-from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
-from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
 from invokeai.app.services.shared.sqlite.sqlite_util import init_db
-from invokeai.backend.model_manager.metadata import ModelMetadataStore
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
 from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.version.invokeai_version import __version__
 
@@ -28,8 +22,10 @@ from ..services.invocation_queue.invocation_queue_memory import MemoryInvocation
 from ..services.invocation_services import InvocationServices
 from ..services.invocation_stats.invocation_stats_default import InvocationStatsService
 from ..services.invoker import Invoker
-from ..services.model_install import ModelInstallService
+from ..services.latents_storage.latents_storage_disk import DiskLatentsStorage
+from ..services.latents_storage.latents_storage_forward_cache import ForwardCacheLatentsStorage
 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
@@ -71,9 +67,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)
@@ -90,22 +83,14 @@ class ApiDependencies:
         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)
-        )
-        model_manager = ModelManagerService(config, logger)
-        model_record_service = ModelRecordServiceSQL(db=db)
+        latents = ForwardCacheLatentsStorage(DiskLatentsStorage(f"{output_folder}/latents"))
         download_queue_service = DownloadQueueService(event_bus=events)
-        metadata_store = ModelMetadataStore(db=db)
-        model_install_service = ModelInstallService(
-            app_config=config,
-            record_store=model_record_service,
+        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),
             download_queue=download_queue_service,
-            metadata_store=metadata_store,
-            event_bus=events,
+            events=events,
         )
         names = SimpleNameService()
         performance_statistics = InvocationStatsService()
@@ -128,11 +113,10 @@ class ApiDependencies:
             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,
@@ -141,8 +125,6 @@ class ApiDependencies:
             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():
+async def prune_downloads() -> Response:
     """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, dest, priority, access_token)
+    return queue.download(source, Path(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():
+async def cancel_all_download_jobs() -> Response:
     """Cancel all download jobs."""
     ApiDependencies.invoker.services.download_queue.cancel_all_jobs()
     return Response(status_code=204)

invokeai/app/api/routers/images.py

@@ -8,7 +8,7 @@ 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

invokeai/app/api/routers/model_manager.py

@@ -0,0 +1,759 @@
+# 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
+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,
+    MainCheckpointConfig,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+from invokeai.backend.model_manager.merge import MergeInterpolationMethod, ModelMerger
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
+
+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(
+    "/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(
+    "/meta/i/{key}",
+    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"},
+        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_manager.store
+    result: Optional[AnyModelRepoMetadata] = record_store.get_metadata(key)
+    if not result:
+        raise HTTPException(status_code=404, detail="No metadata for a model with this key")
+    return result
+
+
+@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
+
+
+@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(
+    "/heuristic_import",
+    operation_id="heuristic_import_model",
+    responses={
+        201: {"description": "The model imported successfully"},
+        415: {"description": "Unrecognized file/folder format"},
+        424: {"description": "The model appeared to import successfully, but could not be found in the model manager"},
+        409: {"description": "There is already a model corresponding to this path or repo_id"},
+    },
+    status_code=201,
+)
+async def heuristic_import(
+    source: str,
+    config: Optional[Dict[str, Any]] = Body(
+        description="Dict of fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
+        default=None,
+        example={"name": "modelT", "description": "antique cars"},
+    ),
+    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,
+        )
+        logger.info(f"Started installation of {source}")
+    except UnknownModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=424, detail=str(e))
+    except InvalidModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=415)
+    except ValueError as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=409, detail=str(e))
+    return result
+
+
+@model_manager_router.post(
+    "/install",
+    operation_id="import_model",
+    responses={
+        201: {"description": "The model imported successfully"},
+        415: {"description": "Unrecognized file/folder format"},
+        424: {"description": "The model appeared to import successfully, but could not be found in the model manager"},
+        409: {"description": "There is already a model corresponding to this path or repo_id"},
+    },
+    status_code=201,
+)
+async def import_model(
+    source: ModelSource,
+    config: Optional[Dict[str, Any]] = Body(
+        description="Dict of fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
+        default=None,
+    ),
+) -> ModelInstallJob:
+    """Install a model using its local path, repo_id, or remote URL.
+
+    Models will be downloaded, probed, configured and installed in a
+    series of background threads. The return object has `status` attribute
+    that can be used to monitor progress.
+
+    The source object is a discriminated Union of LocalModelSource,
+    HFModelSource and URLModelSource. Set the "type" field to the
+    appropriate value:
+
+    * To install a local path using LocalModelSource, pass a source of form:
+      ```
+      {
+        "type": "local",
+        "path": "/path/to/model",
+        "inplace": false
+      }
+      ```
+      The "inplace" flag, if true, will register the model in place in its
+      current filesystem location. Otherwise, the model will be copied
+      into the InvokeAI models directory.
+
+    * To install a HuggingFace repo_id using HFModelSource, pass a source of form:
+      ```
+      {
+        "type": "hf",
+        "repo_id": "stabilityai/stable-diffusion-2.0",
+        "variant": "fp16",
+        "subfolder": "vae",
+        "access_token": "f5820a918aaf01"
+      }
+      ```
+      The `variant`, `subfolder` and `access_token` fields are optional.
+
+    * To install a remote model using an arbitrary URL, pass:
+      ```
+      {
+        "type": "url",
+        "url": "http://www.civitai.com/models/123456",
+        "access_token": "f5820a918aaf01"
+      }
+      ```
+      The `access_token` field is optonal
+
+    The model's configuration record will be probed and filled in
+    automatically.  To override the default guesses, pass "metadata"
+    with a Dict containing the attributes you wish to override.
+
+    Installation occurs in the background. Either use list_model_install_jobs()
+    to poll for completion, or listen on the event bus for the following events:
+
+      * "model_install_running"
+      * "model_install_completed"
+      * "model_install_error"
+
+    On successful completion, the event's payload will contain the field "key"
+    containing the installed ID of the model. On an error, the event's payload
+    will contain the fields "error_type" and "error" describing the nature of the
+    error and its traceback, respectively.
+
+    """
+    logger = ApiDependencies.invoker.services.logger
+
+    try:
+        installer = ApiDependencies.invoker.services.model_manager.install
+        result: ModelInstallJob = installer.import_model(
+            source=source,
+            config=config,
+        )
+        logger.info(f"Started installation of {source}")
+    except UnknownModelException as e:
+        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.
+    """
+    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
+    loader.load_model_by_config(model_config, submodel_type=SubModelType.Scheduler)
+
+    # Get the path of the converted model from the loader
+    cache_path = loader.convert_cache.cache_path(key)
+    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

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

@@ -1,427 +0,0 @@
-# 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_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,8 +47,7 @@ if True:  # hack to make flake8 happy with imports coming after setting up the c
         boards,
         download_queue,
         images,
-        model_records,
-        models,
+        model_manager,
         session_queue,
         sessions,
         utilities,
@@ -58,6 +56,8 @@ if True:  # hack to make flake8 happy with imports coming after setting up the c
     from .api.sockets import SocketIO
     from .invocations.baseinvocation import (
         BaseInvocation,
+        InputFieldJSONSchemaExtra,
+        OutputFieldJSONSchemaExtra,
         UIConfigBase,
     )
 
@@ -114,8 +114,7 @@ async def shutdown_event() -> None:
 app.include_router(sessions.session_router, prefix="/api")
 
 app.include_router(utilities.utilities_router, prefix="/api")
-app.include_router(models.models_router, prefix="/api")
-app.include_router(model_records.model_records_router, prefix="/api")
+app.include_router(model_manager.model_manager_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")
@@ -177,21 +176,23 @@ def custom_openapi() -> dict[str, Any]:
         invoker_schema["class"] = "invocation"
         openapi_schema["components"]["schemas"][f"{output_type_title}"]["class"] = "output"
 
-    from invokeai.backend.model_management.models import get_model_config_enums
+    # 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():
 
-    for model_config_format_enum in set(get_model_config_enums()):
-        name = model_config_format_enum.__qualname__
+    #     if model_config_name in openapi_schema["components"]["schemas"]:
+    #         # print(f"Config with name {name} already defined")
+    #         continue
 
-        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],
-        }
+    #     openapi_schema["components"]["schemas"][model_config_name] = {
+    #         "title": model_config_name,
+    #         "description": "An enumeration.",
+    #         "type": "string",
+    #         "enum": [v.value for v in enum_set],
+    #     }
 
     app.openapi_schema = openapi_schema
     return app.openapi_schema

invokeai/app/invocations/baseinvocation.py

@@ -12,16 +12,13 @@ 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, 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 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
@@ -55,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.
@@ -76,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.
@@ -221,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.
@@ -246,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(
@@ -303,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"}
 
@@ -513,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,38 +1,45 @@
-from typing import List, Optional, Union
+from dataclasses import dataclass
+from typing import Iterator, List, Optional, Tuple, Union
 
 import torch
 from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
+from transformers import 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
+import invokeai.backend.util.logging as logger
+from invokeai.app.invocations.primitives import ConditioningField, ConditioningOutput
+from invokeai.app.services.model_records import UnknownModelException
+from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.util.ti_utils import extract_ti_triggers_from_prompt
+from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.model_manager import ModelType
+from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
-    ConditioningFieldData,
     ExtraConditioningInfo,
     SDXLConditioningInfo,
 )
+from invokeai.backend.textual_inversion import TextualInversionModelRaw
+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):
@@ -46,7 +53,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."""
@@ -64,40 +71,45 @@ class CompelInvocation(BaseInvocation):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
-        tokenizer_info = context.models.load(**self.clip.tokenizer.model_dump())
-        text_encoder_info = context.models.load(**self.clip.text_encoder.model_dump())
+        tokenizer_info = context.services.model_manager.load_model_by_key(
+            **self.clip.tokenizer.model_dump(),
+            context=context,
+        )
+        text_encoder_info = context.services.model_manager.load_model_by_key(
+            **self.clip.text_encoder.model_dump(),
+            context=context,
+        )
 
-        def _lora_loader():
+        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.clip.loras:
-                lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
-                yield (lora_info.context.model, lora.weight)
+                lora_info = context.services.model_manager.load_model_by_key(
+                    **lora.model_dump(exclude={"weight"}), context=context
+                )
+                assert isinstance(lora_info.model, LoRAModelRaw)
+                yield (lora_info.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 = []
         for trigger in extract_ti_triggers_from_prompt(self.prompt):
             name = trigger[1:-1]
             try:
-                ti_list.append(
-                    (
-                        name,
-                        context.models.load(
-                            model_name=name,
-                            base_model=self.clip.text_encoder.base_model,
-                            model_type=ModelType.TextualInversion,
-                        ).context.model,
-                    )
-                )
-            except ModelNotFoundException:
+                loaded_model = context.services.model_manager.load_model_by_key(
+                    **self.clip.text_encoder.model_dump(),
+                    context=context,
+                ).model
+                assert isinstance(loaded_model, TextualInversionModelRaw)
+                ti_list.append((name, loaded_model))
+            except UnknownModelException:
                 # print(e)
                 # import traceback
                 # print(traceback.format_exc())
                 print(f'Warn: trigger: "{trigger}" not found')
 
         with (
-            ModelPatcher.apply_ti(tokenizer_info.context.model, text_encoder_info.context.model, ti_list) as (
+            ModelPatcher.apply_ti(tokenizer_info.model, text_encoder_info.model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -105,7 +117,7 @@ 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_info.context.model, self.clip.skipped_layers),
+            ModelPatcher.apply_clip_skip(text_encoder_info.model, self.clip.skipped_layers),
         ):
             compel = Compel(
                 tokenizer=tokenizer,
@@ -117,7 +129,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)
@@ -138,12 +150,19 @@ 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)  # TODO: fix type mismatch here
 
-        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,
@@ -152,20 +171,27 @@ class SDXLPromptInvocationBase:
         get_pooled: bool,
         lora_prefix: str,
         zero_on_empty: bool,
-    ):
-        tokenizer_info = context.models.load(**clip_field.tokenizer.model_dump())
-        text_encoder_info = context.models.load(**clip_field.text_encoder.model_dump())
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[ExtraConditioningInfo]]:
+        tokenizer_info = context.services.model_manager.load_model_by_key(
+            **clip_field.tokenizer.model_dump(),
+            context=context,
+        )
+        text_encoder_info = context.services.model_manager.load_model_by_key(
+            **clip_field.text_encoder.model_dump(),
+            context=context,
+        )
 
         # return zero on empty
         if prompt == "" and zero_on_empty:
-            cpu_text_encoder = text_encoder_info.context.model
+            cpu_text_encoder = text_encoder_info.model
+            assert isinstance(cpu_text_encoder, torch.nn.Module)
             c = torch.zeros(
                 (
                     1,
                     cpu_text_encoder.config.max_position_embeddings,
                     cpu_text_encoder.config.hidden_size,
                 ),
-                dtype=text_encoder_info.context.cache.precision,
+                dtype=cpu_text_encoder.dtype,
             )
             if get_pooled:
                 c_pooled = torch.zeros(
@@ -176,37 +202,41 @@ class SDXLPromptInvocationBase:
                 c_pooled = None
             return c, c_pooled, None
 
-        def _lora_loader():
+        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in clip_field.loras:
-                lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
-                yield (lora_info.context.model, lora.weight)
+                lora_info = context.services.model_manager.load_model_by_key(
+                    **lora.model_dump(exclude={"weight"}), context=context
+                )
+                lora_model = lora_info.model
+                assert isinstance(lora_model, LoRAModelRaw)
+                yield (lora_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 = []
         for trigger in extract_ti_triggers_from_prompt(prompt):
             name = trigger[1:-1]
             try:
-                ti_list.append(
-                    (
-                        name,
-                        context.models.load(
-                            model_name=name,
-                            base_model=clip_field.text_encoder.base_model,
-                            model_type=ModelType.TextualInversion,
-                        ).context.model,
-                    )
-                )
-            except ModelNotFoundException:
+                ti_model = context.services.model_manager.load_model_by_attr(
+                    model_name=name,
+                    base_model=text_encoder_info.config.base,
+                    model_type=ModelType.TextualInversion,
+                    context=context,
+                ).model
+                assert isinstance(ti_model, TextualInversionModelRaw)
+                ti_list.append((name, ti_model))
+            except UnknownModelException:
                 # print(e)
                 # import traceback
                 # print(traceback.format_exc())
-                print(f'Warn: trigger: "{trigger}" not found')
+                logger.warning(f'trigger: "{trigger}" not found')
+            except ValueError:
+                logger.warning(f'trigger: "{trigger}" more than one similarly-named textual inversion models')
 
         with (
-            ModelPatcher.apply_ti(tokenizer_info.context.model, text_encoder_info.context.model, ti_list) as (
+            ModelPatcher.apply_ti(tokenizer_info.model, text_encoder_info.model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -214,7 +244,7 @@ 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_info.context.model, clip_field.skipped_layers),
+            ModelPatcher.apply_clip_skip(text_encoder_info.model, clip_field.skipped_layers),
         ):
             compel = Compel(
                 tokenizer=tokenizer,
@@ -228,7 +258,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)
 
@@ -261,7 +291,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."""
@@ -332,6 +362,7 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
                 dim=1,
             )
 
+        assert c2_pooled is not None
         conditioning_data = ConditioningFieldData(
             conditionings=[
                 SDXLConditioningInfo(
@@ -343,9 +374,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(
@@ -353,7 +389,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."""
@@ -380,6 +416,7 @@ 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(
@@ -391,9 +428,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")
@@ -424,9 +466,9 @@ class ClipSkipInvocation(BaseInvocation):
 
 
 def get_max_token_count(
-    tokenizer,
+    tokenizer: CLIPTokenizer,
     prompt: Union[FlattenedPrompt, Blend, Conjunction],
-    truncate_if_too_long=False,
+    truncate_if_too_long: bool = False,
 ) -> int:
     if type(prompt) is Blend:
         blend: Blend = prompt
@@ -438,7 +480,9 @@ 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, parsed_prompt: FlattenedPrompt, truncate_if_too_long=True) -> List[str]:
+def get_tokens_for_prompt_object(
+    tokenizer: CLIPTokenizer, parsed_prompt: FlattenedPrompt, truncate_if_too_long: bool = 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")
 
@@ -451,24 +495,29 @@ def get_tokens_for_prompt_object(tokenizer, parsed_prompt: FlattenedPrompt, trun
         for x in parsed_prompt.children
     ]
     text = " ".join(text_fragments)
-    tokens = tokenizer.tokenize(text)
+    tokens: List[str] = 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, display_label_prefix=None):
+def log_tokenization_for_conjunction(
+    c: Conjunction, tokenizer: CLIPTokenizer, display_label_prefix: Optional[str] = None
+) -> 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, display_label_prefix=None):
+def log_tokenization_for_prompt_object(
+    p: Union[Blend, FlattenedPrompt], tokenizer: CLIPTokenizer, display_label_prefix: Optional[str] = None
+) -> None:
     display_label_prefix = display_label_prefix or ""
     if type(p) is Blend:
         blend: Blend = p
@@ -508,7 +557,12 @@ def log_tokenization_for_prompt_object(p: Union[Blend, FlattenedPrompt], tokeniz
             log_tokenization_for_text(text, tokenizer, display_label=display_label_prefix)
 
 
-def log_tokenization_for_text(text, tokenizer, display_label=None, truncate_if_too_long=False):
+def log_tokenization_for_text(
+    text: str,
+    tokenizer: CLIPTokenizer,
+    display_label: Optional[str] = None,
+    truncate_if_too_long: Optional[bool] = False,
+) -> None:
     """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,14 +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."""

invokeai/app/invocations/controlnet_image_processors.py

@@ -17,33 +17,29 @@ 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, ConfigDict, Field, field_validator, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    ImageField,
-    Input,
-    InputField,
-    OutputField,
-    WithBoard,
-    WithMetadata,
-)
-from invokeai.app.invocations.primitives import ImageOutput
+from invokeai.app.invocations.primitives import ImageField, ImageOutput
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
-from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.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 invokeai.backend.image_util.dw_openpose import DWOpenposeDetector
-from invokeai.backend.model_management.models.base import BaseModelType
 
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    WithMetadata,
     invocation,
     invocation_output,
 )
@@ -60,10 +56,7 @@ CONTROLNET_RESIZE_VALUES = Literal[
 class ControlNetModelField(BaseModel):
     """ControlNet model field"""
 
-    model_name: str = Field(description="Name of the ControlNet model")
-    base_model: BaseModelType = Field(description="Base model")
-
-    model_config = ConfigDict(protected_namespaces=())
+    key: str = Field(description="Model config record key for the ControlNet model")
 
 
 class ControlField(BaseModel):
@@ -143,7 +136,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")
@@ -153,13 +146,22 @@ class ImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
         return image
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        raw_image = context.images.get_pil(self.image.image_name)
+        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)
@@ -178,7 +180,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"""
@@ -201,7 +203,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"""
@@ -230,7 +232,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"""
@@ -252,7 +254,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"""
@@ -270,12 +272,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"""
@@ -302,7 +329,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"""
@@ -319,7 +346,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"""
@@ -342,7 +369,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"""
@@ -369,7 +396,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"""
@@ -399,7 +426,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"""
@@ -415,7 +442,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"""
@@ -438,7 +465,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"""
@@ -467,7 +494,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"""
@@ -507,7 +534,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"""
@@ -549,7 +576,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"""
@@ -593,7 +620,7 @@ 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)
 
@@ -602,30 +629,3 @@ class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
 
         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):
-        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,565 +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")
-
-
-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,35 +1,32 @@
-import os
 from builtins import float
 from typing import List, Union
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
+from typing_extensions import Self
 
 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_management.models.base import BaseModelType, ModelType
-from invokeai.backend.model_management.models.ip_adapter import get_ip_adapter_image_encoder_model_id
+from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.backend.model_manager import BaseModelType, ModelType
 
 
+# LS: Consider moving these two classes into model.py
 class IPAdapterModelField(BaseModel):
-    model_name: str = Field(description="Name of the IP-Adapter model")
-    base_model: BaseModelType = Field(description="Base model")
-
-    model_config = ConfigDict(protected_namespaces=())
+    key: str = Field(description="Key to the IP-Adapter model")
 
 
 class CLIPVisionModelField(BaseModel):
-    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=())
+    key: str = Field(description="Key to the CLIP Vision image encoder model")
 
 
 class IPAdapterField(BaseModel):
@@ -46,12 +43,12 @@ class IPAdapterField(BaseModel):
 
     @field_validator("weight")
     @classmethod
-    def validate_ip_adapter_weight(cls, v):
+    def validate_ip_adapter_weight(cls, v: float) -> float:
         validate_weights(v)
         return v
 
     @model_validator(mode="after")
-    def validate_begin_end_step_percent(self):
+    def validate_begin_end_step_percent(self) -> Self:
         validate_begin_end_step(self.begin_step_percent, self.end_step_percent)
         return self
 
@@ -62,7 +59,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."""
 
@@ -84,33 +81,25 @@ class IPAdapterInvocation(BaseInvocation):
 
     @field_validator("weight")
     @classmethod
-    def validate_ip_adapter_weight(cls, v):
+    def validate_ip_adapter_weight(cls, v: float) -> float:
         validate_weights(v)
         return v
 
     @model_validator(mode="after")
-    def validate_begin_end_step_percent(self):
+    def validate_begin_end_step_percent(self) -> 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_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.config.get().models_path, ip_adapter_info["path"])
-        )
+        ip_adapter_info = context.services.model_manager.store.get_model(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_model = CLIPVisionModelField(
-            model_name=image_encoder_model_name,
-            base_model=BaseModelType.Any,
+        image_encoder_models = context.services.model_manager.store.search_by_attr(
+            model_name=image_encoder_model_name, base_model=BaseModelType.Any, model_type=ModelType.CLIPVision
         )
+        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,13 +3,15 @@
 import math
 from contextlib import ExitStack
 from functools import singledispatchmethod
-from typing import List, Literal, Optional, Union
+from typing import Any, Iterator, List, Literal, Optional, Tuple, 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 (
@@ -18,42 +20,36 @@ 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
 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.model_management.models import ModelType, SilenceWarnings
+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.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,
@@ -67,9 +63,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
 
@@ -78,10 +81,20 @@ if choose_torch_device() == torch.device("mps"):
 
 DEFAULT_PRECISION = choose_precision(choose_torch_device())
 
+SAMPLER_NAME_VALUES = Literal[
+    tuple(SCHEDULER_MAP.keys())
+]  # FIXME: "Invalid type alias". This defeats static type checking.
+
+# 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(
@@ -94,7 +107,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,
@@ -109,7 +122,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."""
@@ -124,10 +137,10 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
         ui_order=4,
     )
 
-    def prep_mask_tensor(self, mask_image):
+    def prep_mask_tensor(self, mask_image: Image) -> torch.Tensor:
         if mask_image.mode != "L":
             mask_image = mask_image.convert("L")
-        mask_tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
+        mask_tensor: torch.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:
@@ -137,34 +150,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 = image_resized_to_grid_as_tensor(image.convert("RGB"))
-            if image.dim() == 3:
-                image = image.unsqueeze(0)
+            image = context.services.images.get_pil_image(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)
         else:
-            image = None
+            image_tensor = 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 is not None:
-            vae_info = context.models.load(**self.vae.vae.model_dump())
+        if image_tensor is not None:
+            vae_info = context.services.model_manager.load_model_by_key(
+                **self.vae.vae.model_dump(),
+                context=context,
+            )
 
-            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)
+            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)
             # 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,
+        return DenoiseMaskOutput(
+            denoise_mask=DenoiseMaskField(
+                mask_name=mask_name,
+                masked_latents_name=masked_latents_name,
+            ),
         )
 
 
@@ -175,7 +195,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.load_model_by_key(
+        **scheduler_info.model_dump(),
+        context=context,
+    )
     with orig_scheduler_info as orig_scheduler:
         scheduler_config = orig_scheduler.config
 
@@ -183,7 +206,7 @@ def get_scheduler(
         scheduler_config = scheduler_config["_backup"]
     scheduler_config = {
         **scheduler_config,
-        **scheduler_extra_config,
+        **scheduler_extra_config,  # FIXME
         "_backup": scheduler_config,
     }
 
@@ -196,6 +219,7 @@ 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
 
 
@@ -204,7 +228,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"""
@@ -232,7 +256,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,
@@ -279,7 +303,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
     )
 
     @field_validator("cfg_scale")
-    def ge_one(cls, v):
+    def ge_one(cls, v: Union[List[float], float]) -> Union[List[float], float]:
         """validate that all cfg_scale values are >= 1"""
         if isinstance(v, list):
             for i in v:
@@ -290,18 +314,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,
-        unet,
-        seed,
+        scheduler: Scheduler,
+        unet: UNet2DConditionModel,
+        seed: int,
     ) -> 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(
@@ -318,7 +358,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
             ),
         )
 
-        conditioning_data = conditioning_data.add_scheduler_args_if_applicable(
+        conditioning_data = conditioning_data.add_scheduler_args_if_applicable(  # FIXME
             scheduler,
             # for ddim scheduler
             eta=0.0,  # ddim_eta
@@ -330,8 +370,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
     def create_pipeline(
         self,
-        unet,
-        scheduler,
+        unet: UNet2DConditionModel,
+        scheduler: Scheduler,
     ) -> StableDiffusionGeneratorPipeline:
         # TODO:
         # configure_model_padding(
@@ -342,10 +382,10 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         class FakeVae:
             class FakeVaeConfig:
-                def __init__(self):
+                def __init__(self) -> None:
                     self.block_out_channels = [0]
 
-            def __init__(self):
+            def __init__(self) -> None:
                 self.config = FakeVae.FakeVaeConfig()
 
         return StableDiffusionGeneratorPipeline(
@@ -362,11 +402,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
     def prep_control_data(
         self,
         context: InvocationContext,
-        control_input: Union[ControlField, List[ControlField]],
+        control_input: Optional[Union[ControlField, List[ControlField]]],
         latents_shape: List[int],
         exit_stack: ExitStack,
         do_classifier_free_guidance: bool = True,
-    ) -> List[ControlNetData]:
+    ) -> Optional[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
@@ -389,16 +429,15 @@ class DenoiseLatentsInvocation(BaseInvocation):
         controlnet_data = []
         for control_info in control_list:
             control_model = exit_stack.enter_context(
-                context.models.load(
-                    model_name=control_info.control_model.model_name,
-                    model_type=ModelType.ControlNet,
-                    base_model=control_info.control_model.base_model,
+                context.services.model_manager.load_model_by_key(
+                    key=control_info.control_model.key,
+                    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?
@@ -456,25 +495,25 @@ 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(
-                    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.services.model_manager.load_model_by_key(
+                    key=single_ip_adapter.ip_adapter_model.key,
+                    context=context,
                 )
             )
 
-            image_encoder_model_info = context.models.load(
-                model_name=single_ip_adapter.image_encoder_model.model_name,
-                model_type=ModelType.CLIPVision,
-                base_model=single_ip_adapter.image_encoder_model.base_model,
+            image_encoder_model_info = context.services.model_manager.load_model_by_key(
+                key=single_ip_adapter.image_encoder_model.key,
+                context=context,
             )
 
             # `single_ip_adapter.image` could be a list or a single ImageField. Normalize to a list here.
-            single_ipa_images = single_ip_adapter.image
-            if not isinstance(single_ipa_images, list):
-                single_ipa_images = [single_ipa_images]
+            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 = [context.images.get_pil(image.image_name) for image in single_ipa_images]
+            single_ipa_images = [
+                context.services.images.get_pil_image(image.image_name) for image in single_ipa_image_fields
+            ]
 
             # 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.
@@ -518,22 +557,19 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         t2i_adapter_data = []
         for t2i_adapter_field in t2i_adapter:
-            t2i_adapter_model_info = context.models.load(
-                model_name=t2i_adapter_field.t2i_adapter_model.model_name,
-                model_type=ModelType.T2IAdapter,
-                base_model=t2i_adapter_field.t2i_adapter_model.base_model,
+            t2i_adapter_model_info = context.services.model_manager.load_model_by_key(
+                key=t2i_adapter_field.t2i_adapter_model.key,
+                context=context,
             )
-            image = context.images.get_pil(t2i_adapter_field.image.image_name)
+            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_field.t2i_adapter_model.base_model == BaseModelType.StableDiffusion1:
+            if t2i_adapter_model_info.base == BaseModelType.StableDiffusion1:
                 max_unet_downscale = 8
-            elif t2i_adapter_field.t2i_adapter_model.base_model == BaseModelType.StableDiffusionXL:
+            elif t2i_adapter_model_info.base == BaseModelType.StableDiffusionXL:
                 max_unet_downscale = 4
             else:
-                raise ValueError(
-                    f"Unexpected T2I-Adapter base model type: '{t2i_adapter_field.t2i_adapter_model.base_model}'."
-                )
+                raise ValueError(f"Unexpected T2I-Adapter base model type: '{t2i_adapter_model_info.base}'.")
 
             t2i_adapter_model: T2IAdapter
             with t2i_adapter_model_info as t2i_adapter_model:
@@ -556,7 +592,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,
+                    num_channels=t2i_adapter_model.config["in_channels"],  # mypy treats this as a FrozenDict
                     device=t2i_adapter_model.device,
                     dtype=t2i_adapter_model.dtype,
                     resize_mode=t2i_adapter_field.resize_mode,
@@ -581,7 +617,15 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
     # original idea by https://github.com/AmericanPresidentJimmyCarter
     # TODO: research more for second order schedulers timesteps
-    def init_scheduler(self, scheduler, device, steps, denoising_start, denoising_end):
+    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)
         if scheduler.config.get("cpu_only", False):
             scheduler.set_timesteps(steps, device="cpu")
             timesteps = scheduler.timesteps.to(device=device)
@@ -593,11 +637,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
@@ -610,14 +654,16 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         return num_inference_steps, timesteps, init_timestep
 
-    def prep_inpaint_mask(self, context: InvocationContext, latents):
+    def prep_inpaint_mask(
+        self, context: InvocationContext, latents: torch.Tensor
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
         if self.denoise_mask is None:
             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
 
@@ -629,11 +675,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
 
@@ -659,25 +705,40 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 do_classifier_free_guidance=True,
             )
 
-            def step_callback(state: PipelineIntermediateState):
-                context.util.sd_step_callback(state, self.unet.unet.base_model)
+            # 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 _lora_loader():
+            # get the unet's config so that we can pass the base to dispatch_progress()
+            unet_config = context.services.model_manager.store.get_model(self.unet.unet.key)
+
+            def step_callback(state: PipelineIntermediateState) -> None:
+                self.dispatch_progress(context, source_node_id, state, unet_config.base)
+
+            def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
                 for lora in self.unet.loras:
-                    lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
-                    yield (lora_info.context.model, lora.weight)
+                    lora_info = context.services.model_manager.load_model_by_key(
+                        **lora.model_dump(exclude={"weight"}),
+                        context=context,
+                    )
+                    yield (lora_info.model, lora.weight)
                     del lora_info
                 return
 
-            unet_info = context.models.load(**self.unet.unet.model_dump())
+            unet_info = context.services.model_manager.load_model_by_key(
+                **self.unet.unet.model_dump(),
+                context=context,
+            )
+            assert isinstance(unet_info.model, UNet2DConditionModel)
             with (
                 ExitStack() as exit_stack,
-                ModelPatcher.apply_freeu(unet_info.context.model, self.unet.freeu_config),
-                set_seamless(unet_info.context.model, self.unet.seamless_axes),
+                ModelPatcher.apply_freeu(unet_info.model, self.unet.freeu_config),
+                set_seamless(unet_info.model, self.unet.seamless_axes),  # FIXME
                 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)
@@ -745,8 +806,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(
@@ -754,9 +816,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(
@@ -772,11 +834,15 @@ 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.load_model_by_key(
+            **self.vae.vae.model_dump(),
+            context=context,
+        )
 
-        with set_seamless(vae_info.context.model, self.vae.seamless_axes), vae_info as vae:
+        with set_seamless(vae_info.model, self.vae.seamless_axes), vae_info as vae:
+            assert isinstance(vae, torch.nn.Module)
             latents = latents.to(vae.device)
             if self.fp32:
                 vae.to(dtype=torch.float32)
@@ -803,7 +869,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()
@@ -827,9 +893,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"]
@@ -840,7 +919,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."""
@@ -863,7 +942,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()
@@ -881,8 +960,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(
@@ -890,7 +971,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."""
@@ -904,7 +985,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()
@@ -923,8 +1004,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(
@@ -932,7 +1015,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."""
@@ -948,8 +1031,9 @@ class ImageToLatentsInvocation(BaseInvocation):
     fp32: bool = InputField(default=DEFAULT_PRECISION == "float32", description=FieldDescriptions.fp32)
 
     @staticmethod
-    def vae_encode(vae_info, upcast, tiled, image_tensor):
+    def vae_encode(vae_info: LoadedModel, upcast: bool, tiled: bool, image_tensor: torch.Tensor) -> torch.Tensor:
         with vae_info as vae:
+            assert isinstance(vae, torch.nn.Module)
             orig_dtype = vae.dtype
             if upcast:
                 vae.to(dtype=torch.float32)
@@ -993,9 +1077,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.load_model_by_key(
+            **self.vae.vae.model_dump(),
+            context=context,
+        )
 
         image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
         if image_tensor.dim() == 3:
@@ -1003,21 +1090,27 @@ 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 = image_tensor_dist.sample().to(dtype=vae.dtype)  # FIXME: uses torch.randn. make reproducible!
+        latents: torch.Tensor = 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:
-        return vae.encode(image_tensor).latents
+        assert isinstance(vae, torch.nn.Module)
+        latents: torch.FloatTensor = vae.encode(image_tensor).latents
+        return latents
 
 
 @invocation(
@@ -1025,7 +1118,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."""
@@ -1041,8 +1134,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.")
@@ -1050,7 +1143,12 @@ class BlendLatentsInvocation(BaseInvocation):
         # TODO:
         device = choose_torch_device()
 
-        def slerp(t, v0, v1, DOT_THRESHOLD=0.9995):
+        def slerp(
+            t: Union[float, npt.NDArray[Any]],  # FIXME: maybe use np.float32 here?
+            v0: Union[torch.Tensor, npt.NDArray[Any]],
+            v1: Union[torch.Tensor, npt.NDArray[Any]],
+            DOT_THRESHOLD: float = 0.9995,
+        ) -> Union[torch.Tensor, npt.NDArray[Any]]:
             """
             Spherical linear interpolation
             Args:
@@ -1083,12 +1181,16 @@ class BlendLatentsInvocation(BaseInvocation):
                 v2 = s0 * v0 + s1 * v1
 
             if inputs_are_torch:
-                v2 = torch.from_numpy(v2).to(device)
-
-            return v2
+                v2_torch: torch.Tensor = torch.from_numpy(v2).to(device)
+                return v2_torch
+            else:
+                assert isinstance(v2, np.ndarray)
+                return v2
 
         # blend
-        blended_latents = slerp(self.alpha, latents_a, latents_b)
+        bl = slerp(self.alpha, latents_a, latents_b)
+        assert isinstance(bl, torch.Tensor)
+        blended_latents: torch.Tensor = bl  # for type checking convenience
 
         # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
         blended_latents = blended_latents.to("cpu")
@@ -1096,8 +1198,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:
@@ -1107,7 +1211,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.
@@ -1142,7 +1246,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
@@ -1151,9 +1255,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")
@@ -1181,15 +1286,19 @@ 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, multiple_of=LATENT_SCALE_FACTOR):
+    def trim_to_multiple_of(self, *args: int, multiple_of: int = LATENT_SCALE_FACTOR) -> Tuple[int, ...]:
         return tuple((x - x % multiple_of) for x in args)
 
     def invoke(self, context: InvocationContext) -> IdealSizeOutput:
+        unet_config = context.services.model_manager.load_model_by_key(
+            **self.unet.unet.model_dump(),
+            context=context,
+        )
         aspect = self.width / self.height
-        dimension = 512
-        if self.unet.unet.base_model == BaseModelType.StableDiffusion2:
+        dimension: float = 512
+        if unet_config.base == BaseModelType.StableDiffusion2:
             dimension = 768
-        elif self.unet.unet.base_model == BaseModelType.StableDiffusionXL:
+        elif unet_config.base == 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__
 

invokeai/app/invocations/model.py

@@ -1,28 +1,27 @@
 import copy
 from typing import List, Optional
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, 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_management import BaseModelType, ModelType, SubModelType
+from ...backend.model_manager import SubModelType
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
     invocation,
     invocation_output,
 )
 
 
 class ModelInfo(BaseModel):
-    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=())
+    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")
 
 
 class LoraInfo(ModelInfo):
@@ -52,7 +51,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")
 
@@ -81,20 +80,13 @@ class ModelLoaderOutput(UNetOutput, CLIPOutput, VAEOutput):
 class MainModelField(BaseModel):
     """Main 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=())
+    key: str = Field(description="Model key")
 
 
 class LoRAModelField(BaseModel):
     """LoRA model field"""
 
-    model_name: str = Field(description="Name of the LoRA model")
-    base_model: BaseModelType = Field(description="Base model")
-
-    model_config = ConfigDict(protected_namespaces=())
+    key: str = Field(description="LoRA model key")
 
 
 @invocation(
@@ -102,7 +94,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."""
@@ -111,85 +103,40 @@ class MainModelLoaderInvocation(BaseInvocation):
     # TODO: precision?
 
     def invoke(self, context: InvocationContext) -> ModelLoaderOutput:
-        base_model = self.model.base_model
-        model_name = self.model.model_name
-        model_type = ModelType.Main
+        key = self.model.key
 
         # TODO: not found exceptions
-        if not context.models.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"
-            )
-        """
+        if not context.services.model_manager.store.exists(key):
+            raise Exception(f"Unknown model {key}")
 
         return ModelLoaderOutput(
             unet=UNetField(
                 unet=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.UNet,
+                    key=key,
+                    submodel_type=SubModelType.UNet,
                 ),
                 scheduler=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Scheduler,
+                    key=key,
+                    submodel_type=SubModelType.Scheduler,
                 ),
                 loras=[],
             ),
             clip=ClipField(
                 tokenizer=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Tokenizer,
+                    key=key,
+                    submodel_type=SubModelType.Tokenizer,
                 ),
                 text_encoder=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.TextEncoder,
+                    key=key,
+                    submodel_type=SubModelType.TextEncoder,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             vae=VaeField(
                 vae=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Vae,
+                    key=key,
+                    submodel_type=SubModelType.Vae,
                 ),
             ),
         )
@@ -203,7 +150,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."""
 
@@ -226,21 +173,16 @@ class LoraLoaderInvocation(BaseInvocation):
         if self.lora is None:
             raise Exception("No LoRA provided")
 
-        base_model = self.lora.base_model
-        lora_name = self.lora.model_name
+        lora_key = self.lora.key
 
-        if not context.models.exists(
-            base_model=base_model,
-            model_name=lora_name,
-            model_type=ModelType.Lora,
-        ):
-            raise Exception(f"Unkown lora name: {lora_name}!")
+        if not context.services.model_manager.store.exists(lora_key):
+            raise Exception(f"Unkown lora: {lora_key}!")
 
-        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.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.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.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')
 
         output = LoraLoaderOutput()
 
@@ -248,10 +190,8 @@ class LoraLoaderInvocation(BaseInvocation):
             output.unet = copy.deepcopy(self.unet)
             output.unet.loras.append(
                 LoraInfo(
-                    base_model=base_model,
-                    model_name=lora_name,
-                    model_type=ModelType.Lora,
-                    submodel=None,
+                    key=lora_key,
+                    submodel_type=None,
                     weight=self.weight,
                 )
             )
@@ -260,10 +200,8 @@ class LoraLoaderInvocation(BaseInvocation):
             output.clip = copy.deepcopy(self.clip)
             output.clip.loras.append(
                 LoraInfo(
-                    base_model=base_model,
-                    model_name=lora_name,
-                    model_type=ModelType.Lora,
-                    submodel=None,
+                    key=lora_key,
+                    submodel_type=None,
                     weight=self.weight,
                 )
             )
@@ -285,7 +223,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."""
@@ -315,24 +253,19 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
         if self.lora is None:
             raise Exception("No LoRA provided")
 
-        base_model = self.lora.base_model
-        lora_name = self.lora.model_name
+        lora_key = self.lora.key
 
-        if not context.models.exists(
-            base_model=base_model,
-            model_name=lora_name,
-            model_type=ModelType.Lora,
-        ):
-            raise Exception(f"Unknown lora name: {lora_name}!")
+        if not context.services.model_manager.store.exists(lora_key):
+            raise Exception(f"Unknown lora: {lora_key}!")
 
-        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.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.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.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.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')
+        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')
 
         output = SDXLLoraLoaderOutput()
 
@@ -340,10 +273,8 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
             output.unet = copy.deepcopy(self.unet)
             output.unet.loras.append(
                 LoraInfo(
-                    base_model=base_model,
-                    model_name=lora_name,
-                    model_type=ModelType.Lora,
-                    submodel=None,
+                    key=lora_key,
+                    submodel_type=None,
                     weight=self.weight,
                 )
             )
@@ -352,10 +283,8 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
             output.clip = copy.deepcopy(self.clip)
             output.clip.loras.append(
                 LoraInfo(
-                    base_model=base_model,
-                    model_name=lora_name,
-                    model_type=ModelType.Lora,
-                    submodel=None,
+                    key=lora_key,
+                    submodel_type=None,
                     weight=self.weight,
                 )
             )
@@ -364,10 +293,8 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
             output.clip2 = copy.deepcopy(self.clip2)
             output.clip2.loras.append(
                 LoraInfo(
-                    base_model=base_model,
-                    model_name=lora_name,
-                    model_type=ModelType.Lora,
-                    submodel=None,
+                    key=lora_key,
+                    submodel_type=None,
                     weight=self.weight,
                 )
             )
@@ -378,13 +305,10 @@ class SDXLLoraLoaderInvocation(BaseInvocation):
 class VAEModelField(BaseModel):
     """Vae model field"""
 
-    model_name: str = Field(description="Name of the model")
-    base_model: BaseModelType = Field(description="Base model")
-
-    model_config = ConfigDict(protected_namespaces=())
+    key: str = Field(description="Model's key")
 
 
-@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"""
 
@@ -395,25 +319,12 @@ class VaeLoaderInvocation(BaseInvocation):
     )
 
     def invoke(self, context: InvocationContext) -> VAEOutput:
-        base_model = self.vae_model.base_model
-        model_name = self.vae_model.model_name
-        model_type = ModelType.Vae
+        key = self.vae_model.key
 
-        if not context.models.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,
-                )
-            )
-        )
+        if not context.services.model_manager.store.exists(key):
+            raise Exception(f"Unkown vae: {key}!")
+
+        return VAEOutput(vae=VaeField(vae=ModelInfo(key=key)))
 
 
 @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,457 @@
+# 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, 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.model_manager import ModelType, SubModelType
+from invokeai.backend.model_patcher 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.load_model_by_key(
+            **self.clip.tokenizer.model_dump(),
+        )
+        text_encoder_info = context.services.model_manager.load_model_by_key(
+            **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.load_model_by_key(**lora.model_dump(exclude={"weight"})).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.load_model_by_attr(
+                                model_name=name,
+                                base_model=text_encoder_info.config.base,
+                                model_type=ModelType.TextualInversion,
+                            ).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.load_model_by_key(**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.load_model_by_key(**lora.model_dump(exclude={"weight"})).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.submodel}")
+
+        vae_info = context.services.model_manager.load_model_by_key(
+            **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"""
+
+    key: str = Field(description="Model ID")
+
+
+@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:
+        model_key = self.model.key
+
+        # TODO: not found exceptions
+        if not context.services.model_manager.store.exists(model_key):
+            raise Exception(f"Unknown model: {model_key}")
+
+        return ONNXModelLoaderOutput(
+            unet=UNetField(
+                unet=ModelInfo(
+                    key=model_key,
+                    submodel_type=SubModelType.UNet,
+                ),
+                scheduler=ModelInfo(
+                    key=model_key,
+                    submodel_type=SubModelType.Scheduler,
+                ),
+                loras=[],
+            ),
+            clip=ClipField(
+                tokenizer=ModelInfo(
+                    key=model_key,
+                    submodel_type=SubModelType.Tokenizer,
+                ),
+                text_encoder=ModelInfo(
+                    key=model_key,
+                    submodel_type=SubModelType.TextEncoder,
+                ),
+                loras=[],
+                skipped_layers=0,
+            ),
+            vae_decoder=VaeField(
+                vae=ModelInfo(
+                    key=model_key,
+                    submodel_type=SubModelType.VaeDecoder,
+                ),
+            ),
+            vae_encoder=VaeField(
+                vae=ModelInfo(
+                    key=model_key,
+                    submodel_type=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,40 +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) -> "DenoiseMaskOutput":
-        return cls(
-            denoise_mask=DenoiseMaskField(mask_name=mask_name, masked_latents_name=masked_latents_name),
-        )
-
 
 # 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):
@@ -337,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"""
@@ -345,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(
@@ -368,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) -> LatentsOutput:
+    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"""
@@ -404,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.app.shared.fields import FieldDescriptions
+from invokeai.backend.model_manager import SubModelType
 
-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,72 +44,52 @@ class SDXLModelLoaderInvocation(BaseInvocation):
     # TODO: precision?
 
     def invoke(self, context: InvocationContext) -> SDXLModelLoaderOutput:
-        base_model = self.model.base_model
-        model_name = self.model.model_name
-        model_type = ModelType.Main
+        model_key = self.model.key
 
         # TODO: not found exceptions
-        if not context.models.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.store.exists(model_key):
+            raise Exception(f"Unknown model: {model_key}")
 
         return SDXLModelLoaderOutput(
             unet=UNetField(
                 unet=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.UNet,
+                    key=model_key,
+                    submodel_type=SubModelType.UNet,
                 ),
                 scheduler=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Scheduler,
+                    key=model_key,
+                    submodel_type=SubModelType.Scheduler,
                 ),
                 loras=[],
             ),
             clip=ClipField(
                 tokenizer=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Tokenizer,
+                    key=model_key,
+                    submodel_type=SubModelType.Tokenizer,
                 ),
                 text_encoder=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.TextEncoder,
+                    key=model_key,
+                    submodel_type=SubModelType.TextEncoder,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             clip2=ClipField(
                 tokenizer=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Tokenizer2,
+                    key=model_key,
+                    submodel_type=SubModelType.Tokenizer2,
                 ),
                 text_encoder=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.TextEncoder2,
+                    key=model_key,
+                    submodel_type=SubModelType.TextEncoder2,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             vae=VaeField(
                 vae=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Vae,
+                    key=model_key,
+                    submodel_type=SubModelType.Vae,
                 ),
             ),
         )
@@ -116,7 +100,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."""
@@ -129,56 +113,40 @@ class SDXLRefinerModelLoaderInvocation(BaseInvocation):
     # TODO: precision?
 
     def invoke(self, context: InvocationContext) -> SDXLRefinerModelLoaderOutput:
-        base_model = self.model.base_model
-        model_name = self.model.model_name
-        model_type = ModelType.Main
+        model_key = self.model.key
 
         # TODO: not found exceptions
-        if not context.models.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.store.exists(model_key):
+            raise Exception(f"Unknown model: {model_key}")
 
         return SDXLRefinerModelLoaderOutput(
             unet=UNetField(
                 unet=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.UNet,
+                    key=model_key,
+                    submodel_type=SubModelType.UNet,
                 ),
                 scheduler=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Scheduler,
+                    key=model_key,
+                    submodel_type=SubModelType.Scheduler,
                 ),
                 loras=[],
             ),
             clip2=ClipField(
                 tokenizer=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Tokenizer2,
+                    key=model_key,
+                    submodel_type=SubModelType.Tokenizer2,
                 ),
                 text_encoder=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.TextEncoder2,
+                    key=model_key,
+                    submodel_type=SubModelType.TextEncoder2,
                 ),
                 loras=[],
                 skipped_layers=0,
             ),
             vae=VaeField(
                 vae=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                    submodel=SubModelType.Vae,
+                    key=model_key,
+                    submodel_type=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,25 +1,25 @@
 from typing import Union
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import BaseModel, 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.backend.model_management.models.base import BaseModelType
+from invokeai.app.shared.fields import FieldDescriptions
 
 
 class T2IAdapterModelField(BaseModel):
-    model_name: str = Field(description="Name of the T2I-Adapter model")
-    base_model: BaseModelType = Field(description="Base model")
-
-    model_config = ConfigDict(protected_namespaces=())
+    key: str = Field(description="Model record key for the T2I-Adapter model")
 
 
 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

@@ -8,15 +8,13 @@ import torch
 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.app.invocations.primitives import ImageField, ImageOutput
+from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
 from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet
 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/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] = {}
+    argparse_groups: ClassVar[Dict[str, Any]] = {}
 
     model_config = SettingsConfigDict(env_file_encoding="utf-8", arbitrary_types_allowed=True, case_sensitive=True)
 
-    def parse_args(self, argv: Optional[list] = sys.argv[1:]):
+    def parse_args(self, argv: Optional[List[str]] = sys.argv[1:]) -> None:
         """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):
+    def add_parser_arguments(cls, parser: ArgumentParser) -> None:
         """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,7 +117,8 @@ class InvokeAISettings(BaseSettings):
         """Return the category of a setting."""
         hints = get_type_hints(cls)
         if command_field in hints:
-            return get_args(hints[command_field])[0]
+            result: str = get_args(hints[command_field])[0]
+            return result
         else:
             return "Uncategorized"
 
@@ -158,7 +159,7 @@ class InvokeAISettings(BaseSettings):
         ]
 
     @classmethod
-    def add_field_argument(cls, command_parser, name: str, field, default_override=None):
+    def add_field_argument(cls, command_parser, name: str, field, default_override=None) -> 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):
+    def print_help(self, file=None) -> None:
         text = self.format_help()
         pydoc.pager(text)
 

invokeai/app/services/config/config_default.py

@@ -173,7 +173,7 @@ from __future__ import annotations
 
 import os
 from pathlib import Path
-from typing import Any, ClassVar, Dict, List, Literal, Optional, Union
+from typing import Any, ClassVar, Dict, List, Literal, Optional
 
 from omegaconf import DictConfig, OmegaConf
 from pydantic import Field
@@ -185,7 +185,9 @@ from .config_base import InvokeAISettings
 INIT_FILE = Path("invokeai.yaml")
 DB_FILE = Path("invokeai.db")
 LEGACY_INIT_FILE = Path("invokeai.init")
-DEFAULT_MAX_VRAM = 0.5
+DEFAULT_RAM_CACHE = 10.0
+DEFAULT_VRAM_CACHE = 0.25
+DEFAULT_CONVERT_CACHE = 20.0
 
 
 class Categories(object):
@@ -237,6 +239,7 @@ class InvokeAIAppConfig(InvokeAISettings):
     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)
@@ -260,8 +263,10 @@ class InvokeAIAppConfig(InvokeAISettings):
     version             : bool = Field(default=False, description="Show InvokeAI version and exit", json_schema_extra=Categories.Other)
 
     # CACHE
-    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, )
+    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)
+
     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)
 
@@ -404,6 +409,11 @@ 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."""
@@ -433,15 +443,20 @@ class InvokeAIAppConfig(InvokeAISettings):
         return True
 
     @property
-    def ram_cache_size(self) -> Union[Literal["auto"], float]:
-        """Return the ram cache size using the legacy or modern setting."""
+    def ram_cache_size(self) -> float:
+        """Return the ram cache size using the legacy or modern setting (GB)."""
         return self.max_cache_size or self.ram
 
     @property
-    def vram_cache_size(self) -> Union[Literal["auto"], float]:
-        """Return the vram cache size using the legacy or modern setting."""
+    def vram_cache_size(self) -> float:
+        """Return the vram cache size using the legacy or modern setting (GB)."""
         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,3 +260,16 @@ 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,10 +4,11 @@
 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
+from typing import Any, Dict, List, Optional, Set
 
 import requests
 from pydantic.networks import AnyHttpUrl
@@ -48,11 +49,12 @@ 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 = {}
+        self._jobs: Dict[int, DownloadJob] = {}
         self._next_job_id = 0
-        self._queue = PriorityQueue()
+        self._queue: PriorityQueue[DownloadJob] = PriorityQueue()
         self._stop_event = threading.Event()
-        self._worker_pool = set()
+        self._job_completed_event = threading.Event()
+        self._worker_pool: Set[threading.Thread] = set()
         self._lock = threading.Lock()
         self._logger = InvokeAILogger.get_logger("DownloadQueueService")
         self._event_bus = event_bus
@@ -188,6 +190,16 @@ 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()
@@ -223,6 +235,7 @@ class DownloadQueueService(DownloadQueueServiceBase):
 
             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.")
 
@@ -407,11 +420,11 @@ class DownloadQueueService(DownloadQueueServiceBase):
 
 # Example on_progress event handler to display a TQDM status bar
 # Activate with:
-#   download_service.download('http://foo.bar/baz', '/tmp', on_progress=TqdmProgress().job_update
+#   download_service.download(DownloadJob('http://foo.bar/baz', '/tmp', on_progress=TqdmProgress().update))
 class TqdmProgress(object):
     """TQDM-based progress bar object to use in on_progress handlers."""
 
-    _bars: Dict[int, tqdm]  # the tqdm object
+    _bars: Dict[int, tqdm]  # type: ignore
     _last: Dict[int, int]  # last bytes downloaded
 
     def __init__(self) -> None:  # noqa D107

invokeai/app/services/events/events_base.py

@@ -11,8 +11,7 @@ from invokeai.app.services.session_queue.session_queue_common import (
     SessionQueueStatus,
 )
 from invokeai.app.util.misc import get_timestamp
-from invokeai.backend.model_management.model_manager import LoadedModelInfo
-from invokeai.backend.model_management.models.base import BaseModelType, ModelType, SubModelType
+from invokeai.backend.model_manager import AnyModelConfig
 
 
 class EventServiceBase:
@@ -55,7 +54,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,
@@ -70,7 +69,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,
@@ -171,10 +170,7 @@ class EventServiceBase:
         queue_item_id: int,
         queue_batch_id: str,
         graph_execution_state_id: str,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        submodel: SubModelType,
+        model_config: AnyModelConfig,
     ) -> None:
         """Emitted when a model is requested"""
         self.__emit_queue_event(
@@ -184,10 +180,7 @@ class EventServiceBase:
                 "queue_item_id": queue_item_id,
                 "queue_batch_id": queue_batch_id,
                 "graph_execution_state_id": graph_execution_state_id,
-                "model_name": model_name,
-                "base_model": base_model,
-                "model_type": model_type,
-                "submodel": submodel,
+                "model_config": model_config.model_dump(),
             },
         )
 
@@ -197,11 +190,7 @@ class EventServiceBase:
         queue_item_id: int,
         queue_batch_id: str,
         graph_execution_state_id: str,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        submodel: SubModelType,
-        loaded_model_info: LoadedModelInfo,
+        model_config: AnyModelConfig,
     ) -> None:
         """Emitted when a model is correctly loaded (returns model info)"""
         self.__emit_queue_event(
@@ -211,13 +200,7 @@ class EventServiceBase:
                 "queue_item_id": queue_item_id,
                 "queue_batch_id": queue_batch_id,
                 "graph_execution_state_id": graph_execution_state_id,
-                "model_name": model_name,
-                "base_model": base_model,
-                "model_type": model_type,
-                "submodel": submodel,
-                "hash": loaded_model_info.hash,
-                "location": str(loaded_model_info.location),
-                "precision": str(loaded_model_info.precision),
+                "model_config": model_config.model_dump(),
             },
         )
 

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

@@ -5,11 +5,11 @@ 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.services.shared.invocation_context import InvocationContextData, build_invocation_context
 from invokeai.app.util.profiler import Profiler
 
 from ..invoker import Invoker
@@ -131,20 +131,16 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
                         # which handles a few things:
                         # - nodes that require a value, but get it only from a connection
                         # - referencing the invocation cache instead of executing the node
-                        context_data = InvocationContextData(
-                            invocation=invocation,
-                            session_id=graph_id,
-                            workflow=queue_item.workflow,
-                            source_node_id=source_node_id,
-                            queue_id=queue_item.session_queue_id,
-                            queue_item_id=queue_item.session_queue_item_id,
-                            batch_id=queue_item.session_queue_batch_id,
+                        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,
+                            )
                         )
-                        context = build_invocation_context(
-                            services=self.__invoker.services,
-                            context_data=context_data,
-                        )
-                        outputs = invocation.invoke_internal(context=context, services=self.__invoker.services)
 
                         # Check queue to see if this is canceled, and skip if so
                         if self.__invoker.services.queue.is_canceled(graph_execution_state.id):

invokeai/app/services/invocation_services.py

@@ -3,15 +3,9 @@ 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
@@ -27,9 +21,8 @@ if TYPE_CHECKING:
     from .invocation_queue.invocation_queue_base import InvocationQueueABC
     from .invocation_stats.invocation_stats_base import InvocationStatsServiceBase
     from .item_storage.item_storage_base import ItemStorageABC
-    from .model_install import ModelInstallServiceBase
+    from .latents_storage.latents_storage_base import LatentsStorageBase
     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
@@ -41,6 +34,31 @@ 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"
+    download_queue: "DownloadQueueServiceBase"
+    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",
@@ -53,11 +71,10 @@ class InvocationServices:
         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",
@@ -67,8 +84,6 @@ class InvocationServices:
         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
@@ -80,11 +95,10 @@ class InvocationServices:
         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
@@ -94,5 +108,3 @@ class InvocationServices:
         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

@@ -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 Iterator
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
@@ -40,18 +40,17 @@ class InvocationStatsServiceBase(ABC):
     "Abstract base class for recording node memory/time performance statistics"
 
     @abstractmethod
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initialize the InvocationStatsService and reset counters to zero
         """
-        pass
 
     @abstractmethod
     def collect_stats(
         self,
         invocation: BaseInvocation,
         graph_execution_state_id: str,
-    ) -> AbstractContextManager:
+    ) -> Iterator[None]:
         """
         Return a context object that will capture the statistics on the execution
         of invocaation. Use with: to place around the part of the code that executes the invocation.
@@ -61,7 +60,7 @@ class InvocationStatsServiceBase(ABC):
         pass
 
     @abstractmethod
-    def reset_stats(self, graph_execution_state_id: str):
+    def reset_stats(self, graph_execution_state_id: str) -> None:
         """
         Reset all statistics for the indicated graph.
         :param graph_execution_state_id: The id of the session whose stats to reset.
@@ -70,7 +69,7 @@ class InvocationStatsServiceBase(ABC):
         pass
 
     @abstractmethod
-    def log_stats(self, graph_execution_state_id: str):
+    def log_stats(self, graph_execution_state_id: str) -> None:
         """
         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,6 +2,7 @@ import json
 import time
 from contextlib import contextmanager
 from pathlib import Path
+from typing import Iterator
 
 import psutil
 import torch
@@ -10,7 +11,7 @@ import invokeai.backend.util.logging as logger
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.item_storage.item_storage_common import ItemNotFoundError
-from invokeai.backend.model_management.model_cache import CacheStats
+from invokeai.backend.model_manager.load.model_cache import CacheStats
 
 from .invocation_stats_base import InvocationStatsServiceBase
 from .invocation_stats_common import (
@@ -41,7 +42,10 @@ class InvocationStatsService(InvocationStatsServiceBase):
         self._invoker = invoker
 
     @contextmanager
-    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str):
+    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Iterator[None]:
+        # This is to handle case of the model manager not being initialized, which happens
+        # during some tests.
+        services = self._invoker.services
         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()
@@ -55,8 +59,9 @@ class InvocationStatsService(InvocationStatsServiceBase):
         start_ram = psutil.Process().memory_info().rss
         if torch.cuda.is_available():
             torch.cuda.reset_peak_memory_stats()
-        if self._invoker.services.model_manager:
-            self._invoker.services.model_manager.collect_cache_stats(self._cache_stats[graph_execution_state_id])
+
+        assert services.model_manager.load is not None
+        services.model_manager.load.ram_cache.stats = self._cache_stats[graph_execution_state_id]
 
         try:
             # Let the invocation run.
@@ -73,7 +78,7 @@ class InvocationStatsService(InvocationStatsServiceBase):
             )
             self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
 
-    def _prune_stale_stats(self):
+    def _prune_stale_stats(self) -> None:
         """Check all graphs being tracked and prune any that have completed/errored.
 
         This shouldn't be necessary, but we don't have totally robust upstream handling of graph completions/errors, so

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,49 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from abc import ABC, abstractmethod
+from typing import Callable, Union
+
+import torch
+
+from invokeai.app.invocations.compel import ConditioningFieldData
+
+
+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
+
+    # (LS) Added a Union with ConditioningFieldData to fix type mismatch errors in compel.py
+    # Not 100% sure this isn't an existing bug.
+    @abstractmethod
+    def save(self, name: str, data: Union[torch.Tensor, ConditioningFieldData]) -> 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,59 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from pathlib import Path
+from typing import Union
+
+import torch
+
+from invokeai.app.invocations.compel import ConditioningFieldData
+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: Union[torch.Tensor, ConditioningFieldData]) -> 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,71 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from queue import Queue
+from typing import Dict, Optional, Union
+
+import torch
+
+from invokeai.app.invocations.compel import ConditioningFieldData
+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
+
+    # TODO: (LS) ConditioningFieldData added as Union because of type-checking errors
+    # in compel.py. Unclear whether this is a long-standing bug, but seems to run.
+    def save(self, name: str, data: Union[torch.Tensor, ConditioningFieldData]) -> 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,11 +14,13 @@ 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 import EventServiceBase
+from invokeai.app.services.events.events_base 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, ModelMetadataStore
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
+
+from ..model_metadata import ModelMetadataStoreBase
 
 
 class InstallStatus(str, Enum):
@@ -127,8 +129,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
 
 
@@ -243,7 +245,7 @@ class ModelInstallServiceBase(ABC):
         app_config: InvokeAIAppConfig,
         record_store: ModelRecordServiceBase,
         download_queue: DownloadQueueServiceBase,
-        metadata_store: ModelMetadataStore,
+        metadata_store: ModelMetadataStoreBase,
         event_bus: Optional["EventServiceBase"] = None,
     ):
         """
@@ -324,6 +326,43 @@ 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,
@@ -385,6 +424,18 @@ 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]:
         """
@@ -394,7 +445,8 @@ 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.
+        installs do not complete within the indicated time. A timeout of zero (the default)
+        will block indefinitely until the installs complete.
         """
 
     @abstractmethod
@@ -410,3 +462,22 @@ 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
+from invokeai.app.services.download import DownloadJob, DownloadQueueServiceBase, TqdmProgress
 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, ModelRecordServiceSQL
+from invokeai.app.services.model_records import DuplicateModelException, ModelRecordServiceBase
 from invokeai.backend.model_manager.config import (
     AnyModelConfig,
     BaseModelType,
@@ -33,7 +33,6 @@ from invokeai.backend.model_manager.metadata import (
     AnyModelRepoMetadata,
     CivitaiMetadataFetch,
     HuggingFaceMetadataFetch,
-    ModelMetadataStore,
     ModelMetadataWithFiles,
     RemoteModelFile,
 )
@@ -50,6 +49,7 @@ from .model_install_base import (
     ModelInstallJob,
     ModelInstallServiceBase,
     ModelSource,
+    StringLikeSource,
     URLModelSource,
 )
 
@@ -64,7 +64,6 @@ 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,
     ):
@@ -86,19 +85,13 @@ 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
-        # 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)
+        self._metadata_store = record_store.metadata_store  # for convenience
 
     @property
     def app_config(self) -> InvokeAIAppConfig:  # noqa D102
@@ -145,7 +138,7 @@ class ModelInstallService(ModelInstallServiceBase):
     ) -> str:  # noqa D102
         model_path = Path(model_path)
         config = config or {}
-        if config.get("source") is None:
+        if not config.get("source"):
             config["source"] = model_path.resolve().as_posix()
         return self._register(model_path, config)
 
@@ -156,12 +149,14 @@ class ModelInstallService(ModelInstallServiceBase):
     ) -> str:  # noqa D102
         model_path = Path(model_path)
         config = config or {}
-        if config.get("source") is None:
+        if not config.get("source"):
             config["source"] = model_path.resolve().as_posix()
 
         info: AnyModelConfig = self._probe_model(Path(model_path), config)
-        old_hash = info.original_hash
-        dest_path = self.app_config.models_path / info.base.value / info.type.value / model_path.name
+        old_hash = info.current_hash
+        dest_path = (
+            self.app_config.models_path / info.base.value / info.type.value / (config.get("name") or model_path.name)
+        )
         try:
             new_path = self._copy_model(model_path, dest_path)
         except FileExistsError as excp:
@@ -177,7 +172,40 @@ 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
@@ -207,14 +235,25 @@ 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=5):  # in case we miss an event
+            if self._downloads_changed_event.wait(timeout=0.25):  # in case we miss an event
                 self._downloads_changed_event.clear()
             if timeout > 0 and time.time() - start > timeout:
-                raise Exception("Timeout exceeded")
+                raise TimeoutError("Timeout exceeded")
         self._install_queue.join()
         return self._install_jobs
 
@@ -268,6 +307,38 @@ 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
     # --------------------------------------------------------------------------------------------
@@ -300,6 +371,7 @@ 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:
@@ -330,6 +402,7 @@ 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")
@@ -489,10 +562,10 @@ class ModelInstallService(ModelInstallServiceBase):
         return id
 
     @staticmethod
-    def _guess_variant() -> ModelRepoVariant:
+    def _guess_variant() -> Optional[ModelRepoVariant]:
         """Guess the best HuggingFace variant type to download."""
         precision = choose_precision(choose_torch_device())
-        return ModelRepoVariant.FP16 if precision == "float16" else ModelRepoVariant.DEFAULT
+        return ModelRepoVariant.FP16 if precision == "float16" else None
 
     def _import_local_model(self, source: LocalModelSource, config: Optional[Dict[str, Any]]) -> ModelInstallJob:
         return ModelInstallJob(
@@ -517,7 +590,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)
+        metadata = HuggingFaceMetadataFetch(self._session).from_id(source.repo_id, source.variant)
         assert isinstance(metadata, ModelMetadataWithFiles)
         remote_files = metadata.download_urls(
             variant=source.variant or self._guess_variant(),
@@ -565,6 +638,8 @@ 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,
@@ -580,6 +655,16 @@ 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
@@ -589,7 +674,7 @@ class ModelInstallService(ModelInstallServiceBase):
         self._logger.debug(f"remote_files={remote_files}")
         for model_file in remote_files:
             url = model_file.url
-            path = model_file.path
+            path = root / model_file.path.relative_to(subfolder)
             self._logger.info(f"Downloading {url} => {path}")
             install_job.total_bytes += model_file.size
             assert hasattr(source, "access_token")

invokeai/app/services/model_load/__init__.py

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

@@ -0,0 +1,40 @@
+# 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.invocations.baseinvocation import InvocationContext
+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: Optional[InvocationContext] = 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
+        """
+
+    @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

@@ -0,0 +1,106 @@
+# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Team
+"""Implementation of model loader service."""
+
+from typing import Optional, Type
+
+from invokeai.app.invocations.baseinvocation import InvocationContext
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException
+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
+
+    @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: Optional[InvocationContext] = 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:
+            self._emit_load_event(
+                context=context,
+                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:
+            self._emit_load_event(
+                context=context,
+                model_config=model_config,
+                loaded=True,
+            )
+        return loaded_model
+
+    def _emit_load_event(
+        self,
+        context: InvocationContext,
+        model_config: AnyModelConfig,
+        loaded: Optional[bool] = False,
+    ) -> None:
+        if context.services.queue.is_canceled(context.graph_execution_state_id):
+            raise CanceledException()
+
+        if not loaded:
+            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_config=model_config,
+            )
+        else:
+            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_config=model_config,
+            )

invokeai/app/services/model_manager/__init__.py

@@ -1 +1,17 @@
-from .model_manager_default import ModelManagerService  # noqa F401
+"""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",
+]

invokeai/app/services/model_manager/model_manager_base.py

@@ -1,283 +1,67 @@
 # Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
 
-from __future__ import annotations
-
 from abc import ABC, abstractmethod
-from logging import Logger
-from pathlib import Path
-from typing import Callable, List, Literal, Optional, Tuple, Union
 
-from pydantic import Field
+from typing_extensions import Self
 
-from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.shared.invocation_context import InvocationContextData
-from invokeai.backend.model_management import (
-    AddModelResult,
-    BaseModelType,
-    LoadedModelInfo,
-    MergeInterpolationMethod,
-    ModelType,
-    SchedulerPredictionType,
-    SubModelType,
-)
-from invokeai.backend.model_management.model_cache import CacheStats
+from invokeai.app.services.invoker import Invoker
+
+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
+from ..shared.sqlite.sqlite_database import SqliteDatabase
 
 
 class ModelManagerServiceBase(ABC):
-    """Responsible for managing models on disk and in memory"""
+    """Abstract base class for the model manager service."""
 
-    @abstractmethod
-    def __init__(
-        self,
-        config: InvokeAIAppConfig,
-        logger: Logger,
-    ):
-        """
-        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
+    # 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 get_model(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        submodel: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModelInfo:
-        """Retrieve the indicated model with name and type.
-        submodel can be used to get a part (such as the vae)
-        of a diffusers pipeline."""
+    def build_model_manager(
+        cls,
+        app_config: InvokeAIAppConfig,
+        db: SqliteDatabase,
+        download_queue: DownloadQueueServiceBase,
+        events: EventServiceBase,
+    ) -> Self:
+        """
+        Construct the model manager service instance.
+
+        Use it rather than the __init__ constructor. This class
+        method simplifies the construction considerably.
+        """
         pass
 
     @property
     @abstractmethod
-    def logger(self):
+    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 model_exists(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-    ) -> bool:
+    def start(self, invoker: Invoker) -> None:
         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.
-        """
+    def stop(self, invoker: Invoker) -> None:
         pass

invokeai/app/services/model_manager/model_manager_default.py

@@ -1,421 +1,149 @@
 # Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
+"""Implementation of ModelManagerServiceBase."""
 
-from __future__ import annotations
+from typing import Optional
 
-from logging import Logger
-from pathlib import Path
-from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
+from typing_extensions import Self
 
-import torch
-from pydantic import Field
-
-from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException
+from invokeai.app.invocations.baseinvocation import InvocationContext
 from invokeai.app.services.invoker import Invoker
-from invokeai.app.services.shared.invocation_context import InvocationContextData
-from invokeai.backend.model_management import (
-    AddModelResult,
-    BaseModelType,
-    LoadedModelInfo,
-    MergeInterpolationMethod,
-    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 invokeai.backend.model_manager import AnyModelConfig, BaseModelType, LoadedModel, ModelType, SubModelType
+from invokeai.backend.model_manager.load import ModelCache, ModelConvertCache, ModelLoaderRegistry
+from invokeai.backend.util.logging import InvokeAILogger
 
+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:
-    pass
 
-
-# simple implementation
 class ModelManagerService(ModelManagerServiceBase):
-    """Responsible for managing models on disk and in memory"""
+    """
+    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.
+    """
 
     def __init__(
         self,
-        config: InvokeAIAppConfig,
-        logger: Logger,
+        store: ModelRecordServiceBase,
+        install: ModelInstallServiceBase,
+        load: ModelLoadServiceBase,
     ):
-        """
-        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"
-
-        logger.debug(f"Config file={config_file}")
-
-        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}")
-
-        precision = config.precision
-        if precision == "auto":
-            precision = choose_precision(device)
-        dtype = torch.float32 if precision == "float32" else torch.float16
-
-        # 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
-
-        logger.debug(f"Maximum RAM cache size: {max_cache_size} GiB")
-
-        sequential_offload = config.sequential_guidance
-
-        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 start(self, invoker: Invoker) -> None:
-        self._invoker: Optional[Invoker] = invoker
-
-    def get_model(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        submodel: Optional[SubModelType] = None,
-        context_data: Optional[InvocationContextData] = None,
-    ) -> LoadedModelInfo:
-        """
-        Retrieve the indicated model. submodel can be used to get a
-        part (such as the vae) of a diffusers mode.
-        """
-
-        # we can emit model loading events if we are executing with access to the invocation context
-        if context_data is not None:
-            self._emit_load_event(
-                context_data=context_data,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-            )
-
-        loaded_model_info = self.mgr.get_model(
-            model_name,
-            base_model,
-            model_type,
-            submodel,
-        )
-
-        if context_data is not None:
-            self._emit_load_event(
-                context_data=context_data,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-                loaded_model_info=loaded_model_info,
-            )
-
-        return loaded_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_data: InvocationContextData,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        submodel: Optional[SubModelType] = None,
-        loaded_model_info: Optional[LoadedModelInfo] = None,
-    ):
-        if self._invoker is None:
-            return
-
-        if self._invoker.services.queue.is_canceled(context_data.session_id):
-            raise CanceledException()
-
-        if loaded_model_info:
-            self._invoker.services.events.emit_model_load_completed(
-                queue_id=context_data.queue_id,
-                queue_item_id=context_data.queue_item_id,
-                queue_batch_id=context_data.batch_id,
-                graph_execution_state_id=context_data.session_id,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-                loaded_model_info=loaded_model_info,
-            )
-        else:
-            self._invoker.services.events.emit_model_load_started(
-                queue_id=context_data.queue_id,
-                queue_item_id=context_data.queue_item_id,
-                queue_batch_id=context_data.batch_id,
-                graph_execution_state_id=context_data.session_id,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-            )
+        self._store = store
+        self._install = install
+        self._load = load
 
     @property
-    def logger(self):
-        return self.mgr.logger
+    def store(self) -> ModelRecordServiceBase:
+        return self._store
 
-    def heuristic_import(
+    @property
+    def install(self) -> ModelInstallServiceBase:
+        return self._install
+
+    @property
+    def load(self) -> ModelLoadServiceBase:
+        return self._load
+
+    def start(self, invoker: Invoker) -> None:
+        for service in [self._store, self._install, self._load]:
+            if hasattr(service, "start"):
+                service.start(invoker)
+
+    def stop(self, invoker: Invoker) -> None:
+        for service in [self._store, self._install, self._load]:
+            if hasattr(service, "stop"):
+                service.stop(invoker)
+
+    def load_model_by_config(
         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.
+        model_config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+        context: Optional[InvocationContext] = None,
+    ) -> LoadedModel:
+        return self.load.load_model(model_config, submodel_type, context)
 
-        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.
-        """
-        return self.mgr.heuristic_import(items_to_import, prediction_type_helper)
-
-    def merge_models(
+    def load_model_by_key(
         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:
-        """
-        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
+        key: str,
+        submodel_type: Optional[SubModelType] = None,
+        context: Optional[InvocationContext] = None,
+    ) -> LoadedModel:
+        config = self.store.get_model(key)
+        return self.load.load_model(config, submodel_type, context)
 
-    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(
+    def load_model_by_attr(
         self,
         model_name: str,
         base_model: BaseModelType,
         model_type: ModelType,
-        new_name: Optional[str] = None,
-        new_base: Optional[BaseModelType] = None,
-    ):
+        submodel: Optional[SubModelType] = None,
+        context: Optional[InvocationContext] = None,
+    ) -> LoadedModel:
         """
-        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
+        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
         """
-        self.mgr.rename_model(
-            base_model=base_model,
-            model_type=model_type,
-            model_name=model_name,
-            new_name=new_name,
-            new_base=new_base,
+        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.")
+        else:
+            return self.load.load_model(configs[0], submodel, context)
+
+    @classmethod
+    def build_model_manager(
+        cls,
+        app_config: InvokeAIAppConfig,
+        model_record_service: ModelRecordServiceBase,
+        download_queue: DownloadQueueServiceBase,
+        events: EventServiceBase,
+    ) -> Self:
+        """
+        Construct the model manager service instance.
+
+        For simplicity, use this class method rather than the __init__ constructor.
+        """
+        logger = InvokeAILogger.get_logger(cls.__name__)
+        logger.setLevel(app_config.log_level.upper())
+
+        ram_cache = ModelCache(
+            max_cache_size=app_config.ram_cache_size, max_vram_cache_size=app_config.vram_cache_size, logger=logger
         )
+        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

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

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

@@ -0,0 +1,222 @@
+# 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,8 +11,15 @@ 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.config import AnyModelConfig, BaseModelType, ModelFormat, ModelType
-from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata, ModelMetadataStore
+from invokeai.backend.model_manager import (
+    AnyModelConfig,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
+
+from ..model_metadata import ModelMetadataStoreBase
 
 
 class DuplicateModelException(Exception):
@@ -104,7 +111,7 @@ class ModelRecordServiceBase(ABC):
 
     @property
     @abstractmethod
-    def metadata_store(self) -> ModelMetadataStore:
+    def metadata_store(self) -> ModelMetadataStoreBase:
         """Return a ModelMetadataStore initialized on the same database."""
         pass
 
@@ -146,7 +153,7 @@ class ModelRecordServiceBase(ABC):
     @abstractmethod
     def exists(self, key: str) -> bool:
         """
-        Return True if a model with the indicated key exists in the databse.
+        Return True if a model with the indicated key exists in the database.
 
         :param key: Unique key for the model to be deleted
         """

invokeai/app/services/model_records/model_records_sql.py

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

@@ -13,11 +13,14 @@ from invokeai.app.invocations import *  # noqa: F401 F403
 from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
+    Input,
+    InputField,
+    InvocationContext,
+    OutputField,
+    UIType,
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import Input, InputField, OutputField, UIType
-from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.misc import uuid_string
 
 # in 3.10 this would be "from types import NoneType"
@@ -537,7 +540,7 @@ class Graph(BaseModel):
         except NodeNotFoundError:
             return False
 
-    def get_node(self, node_path: str) -> BaseInvocation:
+    def get_node(self, node_path: str) -> InvocationsUnion:
         """Gets a node from the graph using a node path."""
         # Materialized graphs may have nodes at the top level
         graph, node_id = self._get_graph_and_node(node_path)
@@ -888,7 +891,7 @@ class GraphExecutionState(BaseModel):
         # If next is still none, there's no next node, return None
         return next_node
 
-    def complete(self, node_id: str, output: BaseInvocationOutput) -> None:
+    def complete(self, node_id: str, output: InvocationOutputsUnion):
         """Marks a node as complete"""
 
         if node_id not in self.execution_graph.nodes:

invokeai/app/services/shared/invocation_context.py

@@ -1,409 +0,0 @@
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
-
-from PIL.Image import Image
-from torch import Tensor
-
-from invokeai.app.invocations.fields import MetadataField, WithBoard, WithMetadata
-from invokeai.app.services.boards.boards_common import BoardDTO
-from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
-from invokeai.app.services.images.images_common import ImageDTO
-from invokeai.app.services.invocation_services import InvocationServices
-from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
-from invokeai.app.util.step_callback import stable_diffusion_step_callback
-from invokeai.backend.model_management.model_manager import LoadedModelInfo
-from invokeai.backend.model_management.models.base import BaseModelType, ModelType, SubModelType
-from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
-
-if TYPE_CHECKING:
-    from invokeai.app.invocations.baseinvocation import BaseInvocation
-
-"""
-The InvocationContext provides access to various services and data about the current invocation.
-
-We do not provide the invocation services directly, as their methods are both dangerous and
-inconvenient to use.
-
-For example:
-- The `images` service allows nodes to delete or unsafely modify existing images.
-- The `configuration` service allows nodes to change the app's config at runtime.
-- The `events` service allows nodes to emit arbitrary events.
-
-Wrapping these services provides a simpler and safer interface for nodes to use.
-
-When a node executes, a fresh `InvocationContext` is built for it, ensuring nodes cannot interfere
-with each other.
-
-Many of the wrappers have the same signature as the methods they wrap. This allows us to write
-user-facing docstrings and not need to go and update the internal services to match.
-
-Note: The docstrings are in weird places, but that's where they must be to get IDEs to see them.
-"""
-
-
-@dataclass
-class InvocationContextData:
-    invocation: "BaseInvocation"
-    """The invocation that is being executed."""
-    session_id: str
-    """The session that is being executed."""
-    queue_id: str
-    """The queue in which the session is being executed."""
-    source_node_id: str
-    """The ID of the node from which the currently executing invocation was prepared."""
-    queue_item_id: int
-    """The ID of the queue item that is being executed."""
-    batch_id: str
-    """The ID of the batch that is being executed."""
-    workflow: Optional[WorkflowWithoutID] = None
-    """The workflow associated with this queue item, if any."""
-
-
-class InvocationContextInterface:
-    def __init__(self, services: InvocationServices, context_data: InvocationContextData) -> None:
-        self._services = services
-        self._context_data = context_data
-
-
-class BoardsInterface(InvocationContextInterface):
-    def create(self, board_name: str) -> BoardDTO:
-        """
-        Creates a board.
-
-        :param board_name: The name of the board to create.
-        """
-        return self._services.boards.create(board_name)
-
-    def get_dto(self, board_id: str) -> BoardDTO:
-        """
-        Gets a board DTO.
-
-        :param board_id: The ID of the board to get.
-        """
-        return self._services.boards.get_dto(board_id)
-
-    def get_all(self) -> list[BoardDTO]:
-        """
-        Gets all boards.
-        """
-        return self._services.boards.get_all()
-
-    def add_image_to_board(self, board_id: str, image_name: str) -> None:
-        """
-        Adds an image to a board.
-
-        :param board_id: The ID of the board to add the image to.
-        :param image_name: The name of the image to add to the board.
-        """
-        return self._services.board_images.add_image_to_board(board_id, image_name)
-
-    def get_all_image_names_for_board(self, board_id: str) -> list[str]:
-        """
-        Gets all image names for a board.
-
-        :param board_id: The ID of the board to get the image names for.
-        """
-        return self._services.board_images.get_all_board_image_names_for_board(board_id)
-
-
-class LoggerInterface(InvocationContextInterface):
-    def debug(self, message: str) -> None:
-        """
-        Logs a debug message.
-
-        :param message: The message to log.
-        """
-        self._services.logger.debug(message)
-
-    def info(self, message: str) -> None:
-        """
-        Logs an info message.
-
-        :param message: The message to log.
-        """
-        self._services.logger.info(message)
-
-    def warning(self, message: str) -> None:
-        """
-        Logs a warning message.
-
-        :param message: The message to log.
-        """
-        self._services.logger.warning(message)
-
-    def error(self, message: str) -> None:
-        """
-        Logs an error message.
-
-        :param message: The message to log.
-        """
-        self._services.logger.error(message)
-
-
-class ImagesInterface(InvocationContextInterface):
-    def save(
-        self,
-        image: Image,
-        board_id: Optional[str] = None,
-        image_category: ImageCategory = ImageCategory.GENERAL,
-        metadata: Optional[MetadataField] = None,
-    ) -> ImageDTO:
-        """
-        Saves an image, returning its DTO.
-
-        If the current queue item has a workflow or metadata, it is automatically saved with the image.
-
-        :param image: The image to save, as a PIL image.
-        :param board_id: The board ID to add the image to, if it should be added. It the invocation \
-            inherits from `WithBoard`, that board will be used automatically. **Use this only if \
-            you want to override or provide a board manually!**
-        :param image_category: The category of the image. Only the GENERAL category is added \
-            to the gallery.
-        :param metadata: The metadata to save with the image, if it should have any. If the \
-            invocation inherits from `WithMetadata`, that metadata will be used automatically. \
-            **Use this only if you want to override or provide metadata manually!**
-        """
-
-        # If `metadata` is provided directly, use that. Else, use the metadata provided by `WithMetadata`, falling back to None.
-        metadata_ = None
-        if metadata:
-            metadata_ = metadata
-        elif isinstance(self._context_data.invocation, WithMetadata):
-            metadata_ = self._context_data.invocation.metadata
-
-        # If `board_id` is provided directly, use that. Else, use the board provided by `WithBoard`, falling back to None.
-        board_id_ = None
-        if board_id:
-            board_id_ = board_id
-        elif isinstance(self._context_data.invocation, WithBoard) and self._context_data.invocation.board:
-            board_id_ = self._context_data.invocation.board.board_id
-
-        return self._services.images.create(
-            image=image,
-            is_intermediate=self._context_data.invocation.is_intermediate,
-            image_category=image_category,
-            board_id=board_id_,
-            metadata=metadata_,
-            image_origin=ResourceOrigin.INTERNAL,
-            workflow=self._context_data.workflow,
-            session_id=self._context_data.session_id,
-            node_id=self._context_data.invocation.id,
-        )
-
-    def get_pil(self, image_name: str) -> Image:
-        """
-        Gets an image as a PIL Image object.
-
-        :param image_name: The name of the image to get.
-        """
-        return self._services.images.get_pil_image(image_name)
-
-    def get_metadata(self, image_name: str) -> Optional[MetadataField]:
-        """
-        Gets an image's metadata, if it has any.
-
-        :param image_name: The name of the image to get the metadata for.
-        """
-        return self._services.images.get_metadata(image_name)
-
-    def get_dto(self, image_name: str) -> ImageDTO:
-        """
-        Gets an image as an ImageDTO object.
-
-        :param image_name: The name of the image to get.
-        """
-        return self._services.images.get_dto(image_name)
-
-
-class TensorsInterface(InvocationContextInterface):
-    def save(self, tensor: Tensor) -> str:
-        """
-        Saves a tensor, returning its name.
-
-        :param tensor: The tensor to save.
-        """
-
-        name = self._services.tensors.save(obj=tensor)
-        return name
-
-    def load(self, name: str) -> Tensor:
-        """
-        Loads a tensor by name.
-
-        :param name: The name of the tensor to load.
-        """
-        return self._services.tensors.load(name)
-
-
-class ConditioningInterface(InvocationContextInterface):
-    def save(self, conditioning_data: ConditioningFieldData) -> str:
-        """
-        Saves a conditioning data object, returning its name.
-
-        :param conditioning_context_data: The conditioning data to save.
-        """
-
-        name = self._services.conditioning.save(obj=conditioning_data)
-        return name
-
-    def load(self, name: str) -> ConditioningFieldData:
-        """
-        Loads conditioning data by name.
-
-        :param name: The name of the conditioning data to load.
-        """
-
-        return self._services.conditioning.load(name)
-
-
-class ModelsInterface(InvocationContextInterface):
-    def exists(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> bool:
-        """
-        Checks if a model exists.
-
-        :param model_name: The name of the model to check.
-        :param base_model: The base model of the model to check.
-        :param model_type: The type of the model to check.
-        """
-        return self._services.model_manager.model_exists(model_name, base_model, model_type)
-
-    def load(
-        self, model_name: str, base_model: BaseModelType, model_type: ModelType, submodel: Optional[SubModelType] = None
-    ) -> LoadedModelInfo:
-        """
-        Loads a model.
-
-        :param model_name: The name of the model to get.
-        :param base_model: The base model of the model to get.
-        :param model_type: The type of the model to get.
-        :param submodel: The submodel of the model to get.
-        :returns: An object representing the loaded model.
-        """
-
-        # The model manager emits events as it loads the model. It needs the context data to build
-        # the event payloads.
-
-        return self._services.model_manager.get_model(
-            model_name, base_model, model_type, submodel, context_data=self._context_data
-        )
-
-    def get_info(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> dict:
-        """
-        Gets a model's info, an dict-like object.
-
-        :param model_name: The name of the model to get.
-        :param base_model: The base model of the model to get.
-        :param model_type: The type of the model to get.
-        """
-        return self._services.model_manager.model_info(model_name, base_model, model_type)
-
-
-class ConfigInterface(InvocationContextInterface):
-    def get(self) -> InvokeAIAppConfig:
-        """Gets the app's config."""
-
-        return self._services.configuration.get_config()
-
-
-class UtilInterface(InvocationContextInterface):
-    def sd_step_callback(self, intermediate_state: PipelineIntermediateState, base_model: BaseModelType) -> None:
-        """
-        The step callback emits a progress event with the current step, the total number of
-        steps, a preview image, and some other internal metadata.
-
-        This should be called after each denoising step.
-
-        :param intermediate_state: The intermediate state of the diffusion pipeline.
-        :param base_model: The base model for the current denoising step.
-        """
-
-        # The step callback needs access to the events and the invocation queue services, but this
-        # represents a dangerous level of access.
-        #
-        # We wrap the step callback so that nodes do not have direct access to these services.
-
-        stable_diffusion_step_callback(
-            context_data=self._context_data,
-            intermediate_state=intermediate_state,
-            base_model=base_model,
-            invocation_queue=self._services.queue,
-            events=self._services.events,
-        )
-
-
-class InvocationContext:
-    """
-    The `InvocationContext` provides access to various services and data for the current invocation.
-    """
-
-    def __init__(
-        self,
-        images: ImagesInterface,
-        tensors: TensorsInterface,
-        conditioning: ConditioningInterface,
-        models: ModelsInterface,
-        logger: LoggerInterface,
-        config: ConfigInterface,
-        util: UtilInterface,
-        boards: BoardsInterface,
-        context_data: InvocationContextData,
-        services: InvocationServices,
-    ) -> None:
-        self.images = images
-        """Provides methods to save, get and update images and their metadata."""
-        self.tensors = tensors
-        """Provides methods to save and get tensors, including image, noise, masks, and masked images."""
-        self.conditioning = conditioning
-        """Provides methods to save and get conditioning data."""
-        self.models = models
-        """Provides methods to check if a model exists, get a model, and get a model's info."""
-        self.logger = logger
-        """Provides access to the app logger."""
-        self.config = config
-        """Provides access to the app's config."""
-        self.util = util
-        """Provides utility methods."""
-        self.boards = boards
-        """Provides methods to interact with boards."""
-        self._data = context_data
-        """Provides data about the current queue item and invocation. This is an internal API and may change without warning."""
-        self._services = services
-        """Provides access to the full application services. This is an internal API and may change without warning."""
-
-
-def build_invocation_context(
-    services: InvocationServices,
-    context_data: InvocationContextData,
-) -> InvocationContext:
-    """
-    Builds the invocation context for a specific invocation execution.
-
-    :param invocation_services: The invocation services to wrap.
-    :param invocation_context_data: The invocation context data.
-    """
-
-    logger = LoggerInterface(services=services, context_data=context_data)
-    images = ImagesInterface(services=services, context_data=context_data)
-    tensors = TensorsInterface(services=services, context_data=context_data)
-    models = ModelsInterface(services=services, context_data=context_data)
-    config = ConfigInterface(services=services, context_data=context_data)
-    util = UtilInterface(services=services, context_data=context_data)
-    conditioning = ConditioningInterface(services=services, context_data=context_data)
-    boards = BoardsInterface(services=services, context_data=context_data)
-
-    ctx = InvocationContext(
-        images=images,
-        logger=logger,
-        config=config,
-        tensors=tensors,
-        models=models,
-        context_data=context_data,
-        util=util,
-        conditioning=conditioning,
-        services=services,
-        boards=boards,
-    )
-
-    return ctx

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

@@ -8,6 +8,7 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2 import
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_3 import build_migration_3
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_4 import build_migration_4
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_5 import build_migration_5
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_6 import build_migration_6
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -33,6 +34,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_3(app_config=config, logger=logger))
     migrator.register_migration(build_migration_4())
     migrator.register_migration(build_migration_5())
+    migrator.register_migration(build_migration_6())
     migrator.run_migrations()
 
     return db

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

@@ -0,0 +1,62 @@
+import sqlite3
+
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+
+
+class Migration6Callback:
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        self._recreate_model_triggers(cursor)
+        self._delete_ip_adapters(cursor)
+
+    def _recreate_model_triggers(self, cursor: sqlite3.Cursor) -> None:
+        """
+        Adds the timestamp trigger to the model_config table.
+
+        This trigger was inadvertently dropped in earlier migration scripts.
+        """
+
+        cursor.execute(
+            """--sql
+            CREATE TRIGGER IF NOT EXISTS model_config_updated_at
+            AFTER UPDATE
+            ON model_config FOR EACH ROW
+            BEGIN
+                UPDATE model_config SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
+                    WHERE id = old.id;
+            END;
+            """
+        )
+
+    def _delete_ip_adapters(self, cursor: sqlite3.Cursor) -> None:
+        """
+        Delete all the IP adapters.
+
+        The model manager will automatically find and re-add them after the migration
+        is done. This allows the manager to add the correct image encoder to their
+        configuration records.
+        """
+
+        cursor.execute(
+            """--sql
+            DELETE FROM model_config
+              WHERE type='ip_adapter';
+            """
+        )
+
+
+def build_migration_6() -> Migration:
+    """
+    Build the migration from database version 5 to 6.
+
+    This migration does the following:
+    - Adds the model_config_updated_at trigger if it does not exist
+    - Delete all ip_adapter models so that the model prober can find and
+      update with the correct image processor model.
+    """
+    migration_6 = Migration(
+        from_version=5,
+        to_version=6,
+        callback=Migration6Callback(),
+    )
+
+    return migration_6

invokeai/app/shared/fields.py

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

invokeai/app/shared/models.py

@@ -1,6 +1,6 @@
 from pydantic import BaseModel, Field
 
-from invokeai.app.invocations.fields import FieldDescriptions
+from invokeai.app.shared.fields import FieldDescriptions
 
 
 class FreeUConfig(BaseModel):

invokeai/app/util/misc.py

@@ -5,7 +5,7 @@ import uuid
 import numpy as np
 
 
-def get_timestamp():
+def get_timestamp() -> int:
     return int(datetime.datetime.now(datetime.timezone.utc).timestamp())
 
 
@@ -20,16 +20,16 @@ def get_datetime_from_iso_timestamp(iso_timestamp: str) -> datetime.datetime:
 SEED_MAX = np.iinfo(np.uint32).max
 
 
-def get_random_seed():
+def get_random_seed() -> int:
     rng = np.random.default_rng(seed=None)
     return int(rng.integers(0, SEED_MAX))
 
 
-def uuid_string():
+def uuid_string() -> str:
     res = uuid.uuid4()
     return str(res)
 
 
-def is_optional(value: typing.Any):
+def is_optional(value: typing.Any) -> bool:
     """Checks if a value is typed as Optional. Note that Optional is sugar for Union[x, None]."""
     return typing.get_origin(value) is typing.Union and type(None) in typing.get_args(value)

invokeai/app/util/step_callback.py

@@ -1,18 +1,12 @@
-from typing import TYPE_CHECKING
-
 import torch
 from PIL import Image
 
 from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException, ProgressImage
 
-from ...backend.model_management.models import BaseModelType
+from ...backend.model_manager import BaseModelType
 from ...backend.stable_diffusion import PipelineIntermediateState
 from ...backend.util.util import image_to_dataURL
-
-if TYPE_CHECKING:
-    from invokeai.app.services.events.events_base import EventServiceBase
-    from invokeai.app.services.invocation_queue.invocation_queue_base import InvocationQueueABC
-    from invokeai.app.services.shared.invocation_context import InvocationContextData
+from ..invocations.baseinvocation import InvocationContext
 
 
 def sample_to_lowres_estimated_image(samples, latent_rgb_factors, smooth_matrix=None):
@@ -31,13 +25,13 @@ def sample_to_lowres_estimated_image(samples, latent_rgb_factors, smooth_matrix=
 
 
 def stable_diffusion_step_callback(
-    context_data: "InvocationContextData",
+    context: InvocationContext,
     intermediate_state: PipelineIntermediateState,
+    node: dict,
+    source_node_id: str,
     base_model: BaseModelType,
-    invocation_queue: "InvocationQueueABC",
-    events: "EventServiceBase",
-) -> None:
-    if invocation_queue.is_canceled(context_data.session_id):
+):
+    if context.services.queue.is_canceled(context.graph_execution_state_id):
         raise CanceledException
 
     # Some schedulers report not only the noisy latents at the current timestep,
@@ -114,13 +108,13 @@ def stable_diffusion_step_callback(
 
     dataURL = image_to_dataURL(image, image_format="JPEG")
 
-    events.emit_generator_progress(
-        queue_id=context_data.queue_id,
-        queue_item_id=context_data.queue_item_id,
-        queue_batch_id=context_data.batch_id,
-        graph_execution_state_id=context_data.session_id,
-        node_id=context_data.invocation.id,
-        source_node_id=context_data.source_node_id,
+    context.services.events.emit_generator_progress(
+        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,
+        node=node,
+        source_node_id=source_node_id,
         progress_image=ProgressImage(width=width, height=height, dataURL=dataURL),
         step=intermediate_state.step,
         order=intermediate_state.order,

invokeai/backend/__init__.py

@@ -1,12 +1,3 @@
 """
 Initialization file for invokeai.backend
 """
-from .model_management import (  # noqa: F401
-    BaseModelType,
-    LoadedModelInfo,
-    ModelCache,
-    ModelManager,
-    ModelType,
-    SubModelType,
-)
-from .model_management.models import SilenceWarnings  # noqa: F401

invokeai/backend/embeddings/__init__.py

@@ -0,0 +1,4 @@
+"""Initialization file for invokeai.backend.embeddings modules."""
+
+# from .model_patcher import ModelPatcher
+# __all__ = ["ModelPatcher"]

invokeai/backend/embeddings/embedding_base.py

@@ -0,0 +1,12 @@
+"""Base class for LoRA and Textual Inversion models.
+
+The EmbeddingRaw class is the base class of LoRAModelRaw and TextualInversionModelRaw,
+and is used for type checking of calls to the model patcher.
+
+The use of "Raw" here is a historical artifact, and carried forward in
+order to avoid confusion.
+"""
+
+
+class EmbeddingModelRaw:
+    """Base class for LoRA and Textual Inversion models."""

invokeai/backend/image_util/dw_openpose/__init__.py

@@ -1,81 +0,0 @@
-import numpy as np
-import torch
-from controlnet_aux.util import resize_image
-from PIL import Image
-
-from invokeai.backend.image_util.dw_openpose.utils import draw_bodypose, draw_facepose, draw_handpose
-from invokeai.backend.image_util.dw_openpose.wholebody import Wholebody
-
-
-def draw_pose(pose, H, W, draw_face=True, draw_body=True, draw_hands=True, resolution=512):
-    bodies = pose["bodies"]
-    faces = pose["faces"]
-    hands = pose["hands"]
-    candidate = bodies["candidate"]
-    subset = bodies["subset"]
-    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
-
-    if draw_body:
-        canvas = draw_bodypose(canvas, candidate, subset)
-
-    if draw_hands:
-        canvas = draw_handpose(canvas, hands)
-
-    if draw_face:
-        canvas = draw_facepose(canvas, faces)
-
-    dwpose_image = resize_image(
-        canvas,
-        resolution,
-    )
-    dwpose_image = Image.fromarray(dwpose_image)
-
-    return dwpose_image
-
-
-class DWOpenposeDetector:
-    """
-    Code from the original implementation of the DW Openpose Detector.
-    Credits: https://github.com/IDEA-Research/DWPose
-    """
-
-    def __init__(self) -> None:
-        self.pose_estimation = Wholebody()
-
-    def __call__(
-        self, image: Image.Image, draw_face=False, draw_body=True, draw_hands=False, resolution=512
-    ) -> Image.Image:
-        np_image = np.array(image)
-        H, W, C = np_image.shape
-
-        with torch.no_grad():
-            candidate, subset = self.pose_estimation(np_image)
-            nums, keys, locs = candidate.shape
-            candidate[..., 0] /= float(W)
-            candidate[..., 1] /= float(H)
-            body = candidate[:, :18].copy()
-            body = body.reshape(nums * 18, locs)
-            score = subset[:, :18]
-            for i in range(len(score)):
-                for j in range(len(score[i])):
-                    if score[i][j] > 0.3:
-                        score[i][j] = int(18 * i + j)
-                    else:
-                        score[i][j] = -1
-
-            un_visible = subset < 0.3
-            candidate[un_visible] = -1
-
-            # foot = candidate[:, 18:24]
-
-            faces = candidate[:, 24:92]
-
-            hands = candidate[:, 92:113]
-            hands = np.vstack([hands, candidate[:, 113:]])
-
-            bodies = {"candidate": body, "subset": score}
-            pose = {"bodies": bodies, "hands": hands, "faces": faces}
-
-            return draw_pose(
-                pose, H, W, draw_face=draw_face, draw_hands=draw_hands, draw_body=draw_body, resolution=resolution
-            )

invokeai/backend/image_util/dw_openpose/onnxdet.py

@@ -1,128 +0,0 @@
-# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
-
-import cv2
-import numpy as np
-
-
-def nms(boxes, scores, nms_thr):
-    """Single class NMS implemented in Numpy."""
-    x1 = boxes[:, 0]
-    y1 = boxes[:, 1]
-    x2 = boxes[:, 2]
-    y2 = boxes[:, 3]
-
-    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
-    order = scores.argsort()[::-1]
-
-    keep = []
-    while order.size > 0:
-        i = order[0]
-        keep.append(i)
-        xx1 = np.maximum(x1[i], x1[order[1:]])
-        yy1 = np.maximum(y1[i], y1[order[1:]])
-        xx2 = np.minimum(x2[i], x2[order[1:]])
-        yy2 = np.minimum(y2[i], y2[order[1:]])
-
-        w = np.maximum(0.0, xx2 - xx1 + 1)
-        h = np.maximum(0.0, yy2 - yy1 + 1)
-        inter = w * h
-        ovr = inter / (areas[i] + areas[order[1:]] - inter)
-
-        inds = np.where(ovr <= nms_thr)[0]
-        order = order[inds + 1]
-
-    return keep
-
-
-def multiclass_nms(boxes, scores, nms_thr, score_thr):
-    """Multiclass NMS implemented in Numpy. Class-aware version."""
-    final_dets = []
-    num_classes = scores.shape[1]
-    for cls_ind in range(num_classes):
-        cls_scores = scores[:, cls_ind]
-        valid_score_mask = cls_scores > score_thr
-        if valid_score_mask.sum() == 0:
-            continue
-        else:
-            valid_scores = cls_scores[valid_score_mask]
-            valid_boxes = boxes[valid_score_mask]
-            keep = nms(valid_boxes, valid_scores, nms_thr)
-            if len(keep) > 0:
-                cls_inds = np.ones((len(keep), 1)) * cls_ind
-                dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
-                final_dets.append(dets)
-    if len(final_dets) == 0:
-        return None
-    return np.concatenate(final_dets, 0)
-
-
-def demo_postprocess(outputs, img_size, p6=False):
-    grids = []
-    expanded_strides = []
-    strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
-
-    hsizes = [img_size[0] // stride for stride in strides]
-    wsizes = [img_size[1] // stride for stride in strides]
-
-    for hsize, wsize, stride in zip(hsizes, wsizes, strides, strict=False):
-        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
-        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
-        grids.append(grid)
-        shape = grid.shape[:2]
-        expanded_strides.append(np.full((*shape, 1), stride))
-
-    grids = np.concatenate(grids, 1)
-    expanded_strides = np.concatenate(expanded_strides, 1)
-    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
-    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
-
-    return outputs
-
-
-def preprocess(img, input_size, swap=(2, 0, 1)):
-    if len(img.shape) == 3:
-        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
-    else:
-        padded_img = np.ones(input_size, dtype=np.uint8) * 114
-
-    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
-    resized_img = cv2.resize(
-        img,
-        (int(img.shape[1] * r), int(img.shape[0] * r)),
-        interpolation=cv2.INTER_LINEAR,
-    ).astype(np.uint8)
-    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
-
-    padded_img = padded_img.transpose(swap)
-    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
-    return padded_img, r
-
-
-def inference_detector(session, oriImg):
-    input_shape = (640, 640)
-    img, ratio = preprocess(oriImg, input_shape)
-
-    ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
-    output = session.run(None, ort_inputs)
-    predictions = demo_postprocess(output[0], input_shape)[0]
-
-    boxes = predictions[:, :4]
-    scores = predictions[:, 4:5] * predictions[:, 5:]
-
-    boxes_xyxy = np.ones_like(boxes)
-    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
-    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
-    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
-    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
-    boxes_xyxy /= ratio
-    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
-    if dets is not None:
-        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
-        isscore = final_scores > 0.3
-        iscat = final_cls_inds == 0
-        isbbox = [i and j for (i, j) in zip(isscore, iscat, strict=False)]
-        final_boxes = final_boxes[isbbox]
-    else:
-        final_boxes = np.array([])
-
-    return final_boxes

invokeai/backend/image_util/dw_openpose/onnxpose.py

@@ -1,361 +0,0 @@
-# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
-
-from typing import List, Tuple
-
-import cv2
-import numpy as np
-import onnxruntime as ort
-
-
-def preprocess(
-    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Do preprocessing for RTMPose model inference.
-
-    Args:
-        img (np.ndarray): Input image in shape.
-        input_size (tuple): Input image size in shape (w, h).
-
-    Returns:
-        tuple:
-        - resized_img (np.ndarray): Preprocessed image.
-        - center (np.ndarray): Center of image.
-        - scale (np.ndarray): Scale of image.
-    """
-    # get shape of image
-    img_shape = img.shape[:2]
-    out_img, out_center, out_scale = [], [], []
-    if len(out_bbox) == 0:
-        out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
-    for i in range(len(out_bbox)):
-        x0 = out_bbox[i][0]
-        y0 = out_bbox[i][1]
-        x1 = out_bbox[i][2]
-        y1 = out_bbox[i][3]
-        bbox = np.array([x0, y0, x1, y1])
-
-        # get center and scale
-        center, scale = bbox_xyxy2cs(bbox, padding=1.25)
-
-        # do affine transformation
-        resized_img, scale = top_down_affine(input_size, scale, center, img)
-
-        # normalize image
-        mean = np.array([123.675, 116.28, 103.53])
-        std = np.array([58.395, 57.12, 57.375])
-        resized_img = (resized_img - mean) / std
-
-        out_img.append(resized_img)
-        out_center.append(center)
-        out_scale.append(scale)
-
-    return out_img, out_center, out_scale
-
-
-def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray:
-    """Inference RTMPose model.
-
-    Args:
-        sess (ort.InferenceSession): ONNXRuntime session.
-        img (np.ndarray): Input image in shape.
-
-    Returns:
-        outputs (np.ndarray): Output of RTMPose model.
-    """
-    all_out = []
-    # build input
-    for i in range(len(img)):
-        input = [img[i].transpose(2, 0, 1)]
-
-        # build output
-        sess_input = {sess.get_inputs()[0].name: input}
-        sess_output = []
-        for out in sess.get_outputs():
-            sess_output.append(out.name)
-
-        # run model
-        outputs = sess.run(sess_output, sess_input)
-        all_out.append(outputs)
-
-    return all_out
-
-
-def postprocess(
-    outputs: List[np.ndarray],
-    model_input_size: Tuple[int, int],
-    center: Tuple[int, int],
-    scale: Tuple[int, int],
-    simcc_split_ratio: float = 2.0,
-) -> Tuple[np.ndarray, np.ndarray]:
-    """Postprocess for RTMPose model output.
-
-    Args:
-        outputs (np.ndarray): Output of RTMPose model.
-        model_input_size (tuple): RTMPose model Input image size.
-        center (tuple): Center of bbox in shape (x, y).
-        scale (tuple): Scale of bbox in shape (w, h).
-        simcc_split_ratio (float): Split ratio of simcc.
-
-    Returns:
-        tuple:
-        - keypoints (np.ndarray): Rescaled keypoints.
-        - scores (np.ndarray): Model predict scores.
-    """
-    all_key = []
-    all_score = []
-    for i in range(len(outputs)):
-        # use simcc to decode
-        simcc_x, simcc_y = outputs[i]
-        keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
-
-        # rescale keypoints
-        keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
-        all_key.append(keypoints[0])
-        all_score.append(scores[0])
-
-    return np.array(all_key), np.array(all_score)
-
-
-def bbox_xyxy2cs(bbox: np.ndarray, padding: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
-    """Transform the bbox format from (x,y,w,h) into (center, scale)
-
-    Args:
-        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
-            as (left, top, right, bottom)
-        padding (float): BBox padding factor that will be multilied to scale.
-            Default: 1.0
-
-    Returns:
-        tuple: A tuple containing center and scale.
-        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
-            (n, 2)
-        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
-            (n, 2)
-    """
-    # convert single bbox from (4, ) to (1, 4)
-    dim = bbox.ndim
-    if dim == 1:
-        bbox = bbox[None, :]
-
-    # get bbox center and scale
-    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
-    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
-    scale = np.hstack([x2 - x1, y2 - y1]) * padding
-
-    if dim == 1:
-        center = center[0]
-        scale = scale[0]
-
-    return center, scale
-
-
-def _fix_aspect_ratio(bbox_scale: np.ndarray, aspect_ratio: float) -> np.ndarray:
-    """Extend the scale to match the given aspect ratio.
-
-    Args:
-        scale (np.ndarray): The image scale (w, h) in shape (2, )
-        aspect_ratio (float): The ratio of ``w/h``
-
-    Returns:
-        np.ndarray: The reshaped image scale in (2, )
-    """
-    w, h = np.hsplit(bbox_scale, [1])
-    bbox_scale = np.where(w > h * aspect_ratio, np.hstack([w, w / aspect_ratio]), np.hstack([h * aspect_ratio, h]))
-    return bbox_scale
-
-
-def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
-    """Rotate a point by an angle.
-
-    Args:
-        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
-        angle_rad (float): rotation angle in radian
-
-    Returns:
-        np.ndarray: Rotated point in shape (2, )
-    """
-    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
-    rot_mat = np.array([[cs, -sn], [sn, cs]])
-    return rot_mat @ pt
-
-
-def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
-    """To calculate the affine matrix, three pairs of points are required. This
-    function is used to get the 3rd point, given 2D points a & b.
-
-    The 3rd point is defined by rotating vector `a - b` by 90 degrees
-    anticlockwise, using b as the rotation center.
-
-    Args:
-        a (np.ndarray): The 1st point (x,y) in shape (2, )
-        b (np.ndarray): The 2nd point (x,y) in shape (2, )
-
-    Returns:
-        np.ndarray: The 3rd point.
-    """
-    direction = a - b
-    c = b + np.r_[-direction[1], direction[0]]
-    return c
-
-
-def get_warp_matrix(
-    center: np.ndarray,
-    scale: np.ndarray,
-    rot: float,
-    output_size: Tuple[int, int],
-    shift: Tuple[float, float] = (0.0, 0.0),
-    inv: bool = False,
-) -> np.ndarray:
-    """Calculate the affine transformation matrix that can warp the bbox area
-    in the input image to the output size.
-
-    Args:
-        center (np.ndarray[2, ]): Center of the bounding box (x, y).
-        scale (np.ndarray[2, ]): Scale of the bounding box
-            wrt [width, height].
-        rot (float): Rotation angle (degree).
-        output_size (np.ndarray[2, ] | list(2,)): Size of the
-            destination heatmaps.
-        shift (0-100%): Shift translation ratio wrt the width/height.
-            Default (0., 0.).
-        inv (bool): Option to inverse the affine transform direction.
-            (inv=False: src->dst or inv=True: dst->src)
-
-    Returns:
-        np.ndarray: A 2x3 transformation matrix
-    """
-    shift = np.array(shift)
-    src_w = scale[0]
-    dst_w = output_size[0]
-    dst_h = output_size[1]
-
-    # compute transformation matrix
-    rot_rad = np.deg2rad(rot)
-    src_dir = _rotate_point(np.array([0.0, src_w * -0.5]), rot_rad)
-    dst_dir = np.array([0.0, dst_w * -0.5])
-
-    # get four corners of the src rectangle in the original image
-    src = np.zeros((3, 2), dtype=np.float32)
-    src[0, :] = center + scale * shift
-    src[1, :] = center + src_dir + scale * shift
-    src[2, :] = _get_3rd_point(src[0, :], src[1, :])
-
-    # get four corners of the dst rectangle in the input image
-    dst = np.zeros((3, 2), dtype=np.float32)
-    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
-    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
-    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
-
-    if inv:
-        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
-    else:
-        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
-
-    return warp_mat
-
-
-def top_down_affine(
-    input_size: dict, bbox_scale: dict, bbox_center: dict, img: np.ndarray
-) -> Tuple[np.ndarray, np.ndarray]:
-    """Get the bbox image as the model input by affine transform.
-
-    Args:
-        input_size (dict): The input size of the model.
-        bbox_scale (dict): The bbox scale of the img.
-        bbox_center (dict): The bbox center of the img.
-        img (np.ndarray): The original image.
-
-    Returns:
-        tuple: A tuple containing center and scale.
-        - np.ndarray[float32]: img after affine transform.
-        - np.ndarray[float32]: bbox scale after affine transform.
-    """
-    w, h = input_size
-    warp_size = (int(w), int(h))
-
-    # reshape bbox to fixed aspect ratio
-    bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
-
-    # get the affine matrix
-    center = bbox_center
-    scale = bbox_scale
-    rot = 0
-    warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
-
-    # do affine transform
-    img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
-
-    return img, bbox_scale
-
-
-def get_simcc_maximum(simcc_x: np.ndarray, simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
-    """Get maximum response location and value from simcc representations.
-
-    Note:
-        instance number: N
-        num_keypoints: K
-        heatmap height: H
-        heatmap width: W
-
-    Args:
-        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
-        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
-
-    Returns:
-        tuple:
-        - locs (np.ndarray): locations of maximum heatmap responses in shape
-            (K, 2) or (N, K, 2)
-        - vals (np.ndarray): values of maximum heatmap responses in shape
-            (K,) or (N, K)
-    """
-    N, K, Wx = simcc_x.shape
-    simcc_x = simcc_x.reshape(N * K, -1)
-    simcc_y = simcc_y.reshape(N * K, -1)
-
-    # get maximum value locations
-    x_locs = np.argmax(simcc_x, axis=1)
-    y_locs = np.argmax(simcc_y, axis=1)
-    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
-    max_val_x = np.amax(simcc_x, axis=1)
-    max_val_y = np.amax(simcc_y, axis=1)
-
-    # get maximum value across x and y axis
-    mask = max_val_x > max_val_y
-    max_val_x[mask] = max_val_y[mask]
-    vals = max_val_x
-    locs[vals <= 0.0] = -1
-
-    # reshape
-    locs = locs.reshape(N, K, 2)
-    vals = vals.reshape(N, K)
-
-    return locs, vals
-
-
-def decode(simcc_x: np.ndarray, simcc_y: np.ndarray, simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]:
-    """Modulate simcc distribution with Gaussian.
-
-    Args:
-        simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
-        simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
-        simcc_split_ratio (int): The split ratio of simcc.
-
-    Returns:
-        tuple: A tuple containing center and scale.
-        - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
-        - np.ndarray[float32]: scores in shape (K,) or (n, K)
-    """
-    keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
-    keypoints /= simcc_split_ratio
-
-    return keypoints, scores
-
-
-def inference_pose(session, out_bbox, oriImg):
-    h, w = session.get_inputs()[0].shape[2:]
-    model_input_size = (w, h)
-    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
-    outputs = inference(session, resized_img)
-    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
-
-    return keypoints, scores

invokeai/backend/image_util/dw_openpose/utils.py

@@ -1,155 +0,0 @@
-# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
-
-import math
-
-import cv2
-import matplotlib
-import numpy as np
-
-eps = 0.01
-
-
-def draw_bodypose(canvas, candidate, subset):
-    H, W, C = canvas.shape
-    candidate = np.array(candidate)
-    subset = np.array(subset)
-
-    stickwidth = 4
-
-    limbSeq = [
-        [2, 3],
-        [2, 6],
-        [3, 4],
-        [4, 5],
-        [6, 7],
-        [7, 8],
-        [2, 9],
-        [9, 10],
-        [10, 11],
-        [2, 12],
-        [12, 13],
-        [13, 14],
-        [2, 1],
-        [1, 15],
-        [15, 17],
-        [1, 16],
-        [16, 18],
-        [3, 17],
-        [6, 18],
-    ]
-
-    colors = [
-        [255, 0, 0],
-        [255, 85, 0],
-        [255, 170, 0],
-        [255, 255, 0],
-        [170, 255, 0],
-        [85, 255, 0],
-        [0, 255, 0],
-        [0, 255, 85],
-        [0, 255, 170],
-        [0, 255, 255],
-        [0, 170, 255],
-        [0, 85, 255],
-        [0, 0, 255],
-        [85, 0, 255],
-        [170, 0, 255],
-        [255, 0, 255],
-        [255, 0, 170],
-        [255, 0, 85],
-    ]
-
-    for i in range(17):
-        for n in range(len(subset)):
-            index = subset[n][np.array(limbSeq[i]) - 1]
-            if -1 in index:
-                continue
-            Y = candidate[index.astype(int), 0] * float(W)
-            X = candidate[index.astype(int), 1] * float(H)
-            mX = np.mean(X)
-            mY = np.mean(Y)
-            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
-            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
-            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
-            cv2.fillConvexPoly(canvas, polygon, colors[i])
-
-    canvas = (canvas * 0.6).astype(np.uint8)
-
-    for i in range(18):
-        for n in range(len(subset)):
-            index = int(subset[n][i])
-            if index == -1:
-                continue
-            x, y = candidate[index][0:2]
-            x = int(x * W)
-            y = int(y * H)
-            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
-
-    return canvas
-
-
-def draw_handpose(canvas, all_hand_peaks):
-    H, W, C = canvas.shape
-
-    edges = [
-        [0, 1],
-        [1, 2],
-        [2, 3],
-        [3, 4],
-        [0, 5],
-        [5, 6],
-        [6, 7],
-        [7, 8],
-        [0, 9],
-        [9, 10],
-        [10, 11],
-        [11, 12],
-        [0, 13],
-        [13, 14],
-        [14, 15],
-        [15, 16],
-        [0, 17],
-        [17, 18],
-        [18, 19],
-        [19, 20],
-    ]
-
-    for peaks in all_hand_peaks:
-        peaks = np.array(peaks)
-
-        for ie, e in enumerate(edges):
-            x1, y1 = peaks[e[0]]
-            x2, y2 = peaks[e[1]]
-            x1 = int(x1 * W)
-            y1 = int(y1 * H)
-            x2 = int(x2 * W)
-            y2 = int(y2 * H)
-            if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
-                cv2.line(
-                    canvas,
-                    (x1, y1),
-                    (x2, y2),
-                    matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255,
-                    thickness=2,
-                )
-
-        for _, keyponit in enumerate(peaks):
-            x, y = keyponit
-            x = int(x * W)
-            y = int(y * H)
-            if x > eps and y > eps:
-                cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
-    return canvas
-
-
-def draw_facepose(canvas, all_lmks):
-    H, W, C = canvas.shape
-    for lmks in all_lmks:
-        lmks = np.array(lmks)
-        for lmk in lmks:
-            x, y = lmk
-            x = int(x * W)
-            y = int(y * H)
-            if x > eps and y > eps:
-                cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
-    return canvas

invokeai/backend/image_util/dw_openpose/wholebody.py

@@ -1,67 +0,0 @@
-# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
-# Modified pathing to suit Invoke
-
-import pathlib
-
-import numpy as np
-import onnxruntime as ort
-
-from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.backend.util.devices import choose_torch_device
-from invokeai.backend.util.util import download_with_progress_bar
-
-from .onnxdet import inference_detector
-from .onnxpose import inference_pose
-
-DWPOSE_MODELS = {
-    "yolox_l.onnx": {
-        "local": "any/annotators/dwpose/yolox_l.onnx",
-        "url": "https://huggingface.co/yzd-v/DWPose/resolve/main/yolox_l.onnx?download=true",
-    },
-    "dw-ll_ucoco_384.onnx": {
-        "local": "any/annotators/dwpose/dw-ll_ucoco_384.onnx",
-        "url": "https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.onnx?download=true",
-    },
-}
-
-config = InvokeAIAppConfig.get_config()
-
-
-class Wholebody:
-    def __init__(self):
-        device = choose_torch_device()
-
-        providers = ["CUDAExecutionProvider"] if device == "cuda" else ["CPUExecutionProvider"]
-
-        DET_MODEL_PATH = pathlib.Path(config.models_path / DWPOSE_MODELS["yolox_l.onnx"]["local"])
-        if not DET_MODEL_PATH.exists():
-            download_with_progress_bar(DWPOSE_MODELS["yolox_l.onnx"]["url"], DET_MODEL_PATH)
-
-        POSE_MODEL_PATH = pathlib.Path(config.models_path / DWPOSE_MODELS["dw-ll_ucoco_384.onnx"]["local"])
-        if not POSE_MODEL_PATH.exists():
-            download_with_progress_bar(DWPOSE_MODELS["dw-ll_ucoco_384.onnx"]["url"], POSE_MODEL_PATH)
-
-        onnx_det = DET_MODEL_PATH
-        onnx_pose = POSE_MODEL_PATH
-
-        self.session_det = ort.InferenceSession(path_or_bytes=onnx_det, providers=providers)
-        self.session_pose = ort.InferenceSession(path_or_bytes=onnx_pose, providers=providers)
-
-    def __call__(self, oriImg):
-        det_result = inference_detector(self.session_det, oriImg)
-        keypoints, scores = inference_pose(self.session_pose, det_result, oriImg)
-
-        keypoints_info = np.concatenate((keypoints, scores[..., None]), axis=-1)
-        # compute neck joint
-        neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
-        # neck score when visualizing pred
-        neck[:, 2:4] = np.logical_and(keypoints_info[:, 5, 2:4] > 0.3, keypoints_info[:, 6, 2:4] > 0.3).astype(int)
-        new_keypoints_info = np.insert(keypoints_info, 17, neck, axis=1)
-        mmpose_idx = [17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3]
-        openpose_idx = [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17]
-        new_keypoints_info[:, openpose_idx] = new_keypoints_info[:, mmpose_idx]
-        keypoints_info = new_keypoints_info
-
-        keypoints, scores = keypoints_info[..., :2], keypoints_info[..., 2]
-
-        return keypoints, scores

invokeai/backend/image_util/safety_checker.py

@@ -8,8 +8,8 @@ from PIL import Image
 
 import invokeai.backend.util.logging as logger
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.backend import SilenceWarnings
 from invokeai.backend.util.devices import choose_torch_device
+from invokeai.backend.util.silence_warnings import SilenceWarnings
 
 config = InvokeAIAppConfig.get_config()
 

invokeai/backend/install/install_helper.py

@@ -25,18 +25,20 @@ from invokeai.app.services.model_install import (
     ModelSource,
     URLModelSource,
 )
+from invokeai.app.services.model_metadata import ModelMetadataStoreSQL
 from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL
 from invokeai.app.services.shared.sqlite.sqlite_util import init_db
 from invokeai.backend.model_manager import (
     BaseModelType,
     InvalidModelConfigException,
+    ModelRepoVariant,
     ModelType,
 )
 from invokeai.backend.model_manager.metadata import UnknownMetadataException
 from invokeai.backend.util.logging import InvokeAILogger
 
 # name of the starter models file
-INITIAL_MODELS = "INITIAL_MODELS2.yaml"
+INITIAL_MODELS = "INITIAL_MODELS.yaml"
 
 
 def initialize_record_store(app_config: InvokeAIAppConfig) -> ModelRecordServiceBase:
@@ -44,7 +46,7 @@ def initialize_record_store(app_config: InvokeAIAppConfig) -> ModelRecordService
     logger = InvokeAILogger.get_logger(config=app_config)
     image_files = DiskImageFileStorage(f"{app_config.output_path}/images")
     db = init_db(config=app_config, logger=logger, image_files=image_files)
-    obj: ModelRecordServiceBase = ModelRecordServiceSQL(db)
+    obj: ModelRecordServiceBase = ModelRecordServiceSQL(db, ModelMetadataStoreSQL(db))
     return obj
 
 
@@ -53,12 +55,10 @@ def initialize_installer(
 ) -> ModelInstallServiceBase:
     """Return an initialized ModelInstallService object."""
     record_store = initialize_record_store(app_config)
-    metadata_store = record_store.metadata_store
     download_queue = DownloadQueueService()
     installer = ModelInstallService(
         app_config=app_config,
         record_store=record_store,
-        metadata_store=metadata_store,
         download_queue=download_queue,
         event_bus=event_bus,
     )
@@ -98,11 +98,13 @@ class TqdmEventService(EventServiceBase):
         super().__init__()
         self._bars: Dict[str, tqdm] = {}
         self._last: Dict[str, int] = {}
+        self._logger = InvokeAILogger.get_logger(__name__)
 
     def dispatch(self, event_name: str, payload: Any) -> None:
         """Dispatch an event by appending it to self.events."""
+        data = payload["data"]
+        source = data["source"]
         if payload["event"] == "model_install_downloading":
-            data = payload["data"]
             dest = data["local_path"]
             total_bytes = data["total_bytes"]
             bytes = data["bytes"]
@@ -111,6 +113,12 @@ class TqdmEventService(EventServiceBase):
                 self._last[dest] = 0
             self._bars[dest].update(bytes - self._last[dest])
             self._last[dest] = bytes
+        elif payload["event"] == "model_install_completed":
+            self._logger.info(f"{source}: installed successfully.")
+        elif payload["event"] == "model_install_error":
+            self._logger.warning(f"{source}: installation failed with error {data['error']}")
+        elif payload["event"] == "model_install_cancelled":
+            self._logger.warning(f"{source}: installation cancelled")
 
 
 class InstallHelper(object):
@@ -225,11 +233,19 @@ class InstallHelper(object):
 
         if model_path.exists():  # local file on disk
             return LocalModelSource(path=model_path.absolute(), inplace=True)
-        if re.match(r"^[^/]+/[^/]+$", model_path_id_or_url):  # hugging face repo_id
+
+        # parsing huggingface repo ids
+        # we're going to do a little trick that allows for extended repo_ids of form "foo/bar:fp16"
+        variants = "|".join([x.lower() for x in ModelRepoVariant.__members__])
+        if match := re.match(f"^([^/]+/[^/]+?)(?::({variants}))?$", model_path_id_or_url):
+            repo_id = match.group(1)
+            repo_variant = ModelRepoVariant(match.group(2)) if match.group(2) else None
+            subfolder = Path(model_info.subfolder) if model_info.subfolder else None
             return HFModelSource(
-                repo_id=model_path_id_or_url,
+                repo_id=repo_id,
                 access_token=HfFolder.get_token(),
-                subfolder=model_info.subfolder,
+                subfolder=subfolder,
+                variant=repo_variant,
             )
         if re.match(r"^(http|https):", model_path_id_or_url):
             return URLModelSource(url=AnyHttpUrl(model_path_id_or_url))
@@ -270,12 +286,14 @@ class InstallHelper(object):
                 model_name=model_name,
             )
             if len(matches) > 1:
-                print(f"{model} is ambiguous. Please use model_type:model_name (e.g. main:my_model) to disambiguate.")
+                self._logger.error(
+                    "{model_to_remove} is ambiguous. Please use model_base/model_type/model_name (e.g. sd-1/main/my_model) to disambiguate"
+                )
             elif not matches:
-                print(f"{model}: unknown model")
+                self._logger.error(f"{model_to_remove}: unknown model")
             else:
                 for m in matches:
-                    print(f"Deleting {m.type}:{m.name}")
+                    self._logger.info(f"Deleting {m.type}:{m.name}")
                     installer.delete(m.key)
 
         installer.wait_for_installs()

invokeai/backend/install/invokeai_configure.py

@@ -18,31 +18,30 @@ from argparse import Namespace
 from enum import Enum
 from pathlib import Path
 from shutil import get_terminal_size
-from typing import Any, get_args, get_type_hints
+from typing import Any, Optional, Set, Tuple, Type, get_args, get_type_hints
 from urllib import request
 
 import npyscreen
-import omegaconf
 import psutil
 import torch
 import transformers
-import yaml
-from diffusers import AutoencoderKL
+from diffusers import AutoencoderKL, ModelMixin
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 from huggingface_hub import HfFolder
 from huggingface_hub import login as hf_hub_login
-from omegaconf import OmegaConf
-from pydantic import ValidationError
+from omegaconf import DictConfig, OmegaConf
+from pydantic.error_wrappers import ValidationError
 from tqdm import tqdm
 from transformers import AutoFeatureExtractor, BertTokenizerFast, CLIPTextConfig, CLIPTextModel, CLIPTokenizer
 
 import invokeai.configs as configs
 from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.install.install_helper import InstallHelper, InstallSelections
 from invokeai.backend.install.legacy_arg_parsing import legacy_parser
-from invokeai.backend.install.model_install_backend import InstallSelections, ModelInstall, hf_download_from_pretrained
-from invokeai.backend.model_management.model_probe import BaseModelType, ModelType
+from invokeai.backend.model_manager import BaseModelType, ModelType
+from invokeai.backend.util import choose_precision, choose_torch_device
 from invokeai.backend.util.logging import InvokeAILogger
-from invokeai.frontend.install.model_install import addModelsForm, process_and_execute
+from invokeai.frontend.install.model_install import addModelsForm
 
 # TO DO - Move all the frontend code into invokeai.frontend.install
 from invokeai.frontend.install.widgets import (
@@ -61,7 +60,7 @@ warnings.filterwarnings("ignore")
 transformers.logging.set_verbosity_error()
 
 
-def get_literal_fields(field) -> list[Any]:
+def get_literal_fields(field: str) -> Tuple[Any]:
     return get_args(get_type_hints(InvokeAIAppConfig).get(field))
 
 
@@ -80,8 +79,7 @@ ATTENTION_SLICE_CHOICES = get_literal_fields("attention_slice_size")
 GENERATION_OPT_CHOICES = ["sequential_guidance", "force_tiled_decode", "lazy_offload"]
 GB = 1073741824  # GB in bytes
 HAS_CUDA = torch.cuda.is_available()
-_, MAX_VRAM = torch.cuda.mem_get_info() if HAS_CUDA else (0, 0)
-
+_, MAX_VRAM = torch.cuda.mem_get_info() if HAS_CUDA else (0.0, 0.0)
 
 MAX_VRAM /= GB
 MAX_RAM = psutil.virtual_memory().total / GB
@@ -96,13 +94,15 @@ logger = InvokeAILogger.get_logger()
 
 
 class DummyWidgetValue(Enum):
+    """Dummy widget values."""
+
     zero = 0
     true = True
     false = False
 
 
 # --------------------------------------------
-def postscript(errors: None):
+def postscript(errors: Set[str]) -> None:
     if not any(errors):
         message = f"""
 ** INVOKEAI INSTALLATION SUCCESSFUL **
@@ -143,7 +143,7 @@ def yes_or_no(prompt: str, default_yes=True):
 
 
 # ---------------------------------------------
-def HfLogin(access_token) -> str:
+def HfLogin(access_token) -> None:
     """
     Helper for logging in to Huggingface
     The stdout capture is needed to hide the irrelevant "git credential helper" warning
@@ -162,7 +162,7 @@ def HfLogin(access_token) -> str:
 
 # -------------------------------------
 class ProgressBar:
-    def __init__(self, model_name="file"):
+    def __init__(self, model_name: str = "file"):
         self.pbar = None
         self.name = model_name
 
@@ -179,6 +179,22 @@ class ProgressBar:
         self.pbar.update(block_size)
 
 
+# ---------------------------------------------
+def hf_download_from_pretrained(model_class: Type[ModelMixin], model_name: str, destination: Path, **kwargs: Any):
+    filter = lambda x: "fp16 is not a valid" not in x.getMessage()  # noqa E731
+    logger.addFilter(filter)
+    try:
+        model = model_class.from_pretrained(
+            model_name,
+            resume_download=True,
+            **kwargs,
+        )
+        model.save_pretrained(destination, safe_serialization=True)
+    finally:
+        logger.removeFilter(filter)
+    return destination
+
+
 # ---------------------------------------------
 def download_with_progress_bar(model_url: str, model_dest: str, label: str = "the"):
     try:
@@ -249,6 +265,7 @@ def download_conversion_models():
 
 
 # ---------------------------------------------
+# TO DO: use the download queue here.
 def download_realesrgan():
     logger.info("Installing ESRGAN Upscaling models...")
     URLs = [
@@ -288,18 +305,19 @@ def download_lama():
 
 
 # ---------------------------------------------
-def download_support_models():
+def download_support_models() -> None:
     download_realesrgan()
     download_lama()
     download_conversion_models()
 
 
 # -------------------------------------
-def get_root(root: str = None) -> str:
+def get_root(root: Optional[str] = None) -> str:
     if root:
         return root
-    elif os.environ.get("INVOKEAI_ROOT"):
-        return os.environ.get("INVOKEAI_ROOT")
+    elif root := os.environ.get("INVOKEAI_ROOT"):
+        assert root is not None
+        return root
     else:
         return str(config.root_path)
 
@@ -455,6 +473,25 @@ Use cursor arrows to make a checkbox selection, and space to toggle.
             max_width=110,
             scroll_exit=True,
         )
+        self.add_widget_intelligent(
+            npyscreen.TitleFixedText,
+            name="Model disk conversion cache size (GB). This is used to cache safetensors files that need to be converted to diffusers..",
+            begin_entry_at=0,
+            editable=False,
+            color="CONTROL",
+            scroll_exit=True,
+        )
+        self.nextrely -= 1
+        self.disk = self.add_widget_intelligent(
+            npyscreen.Slider,
+            value=clip(old_opts.convert_cache, range=(0, 100), step=0.5),
+            out_of=100,
+            lowest=0.0,
+            step=0.5,
+            relx=8,
+            scroll_exit=True,
+        )
+        self.nextrely += 1
         self.add_widget_intelligent(
             npyscreen.TitleFixedText,
             name="Model RAM cache size (GB). Make this at least large enough to hold a single full model (2GB for SD-1, 6GB for SDXL).",
@@ -495,6 +532,14 @@ Use cursor arrows to make a checkbox selection, and space to toggle.
             )
         else:
             self.vram = DummyWidgetValue.zero
+
+        self.nextrely += 1
+        self.add_widget_intelligent(
+            npyscreen.FixedText,
+            value="Location of the database used to store model path and configuration information:",
+            editable=False,
+            color="CONTROL",
+        )
         self.nextrely += 1
         self.outdir = self.add_widget_intelligent(
             FileBox,
@@ -506,19 +551,21 @@ Use cursor arrows to make a checkbox selection, and space to toggle.
             labelColor="GOOD",
             begin_entry_at=40,
             max_height=3,
+            max_width=127,
             scroll_exit=True,
         )
         self.autoimport_dirs = {}
         self.autoimport_dirs["autoimport_dir"] = self.add_widget_intelligent(
             FileBox,
-            name="Folder to recursively scan for new checkpoints, ControlNets, LoRAs and TI models",
-            value=str(config.root_path / config.autoimport_dir),
+            name="Optional folder to scan for new checkpoints, ControlNets, LoRAs and TI models",
+            value=str(config.root_path / config.autoimport_dir) if config.autoimport_dir else "",
             select_dir=True,
             must_exist=False,
             use_two_lines=False,
             labelColor="GOOD",
             begin_entry_at=32,
             max_height=3,
+            max_width=127,
             scroll_exit=True,
         )
         self.nextrely += 1
@@ -555,6 +602,10 @@ https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENS
         self.attention_slice_label.hidden = not show
         self.attention_slice_size.hidden = not show
 
+    def show_hide_model_conf_override(self, value):
+        self.model_conf_override.hidden = value
+        self.model_conf_override.display()
+
     def on_ok(self):
         options = self.marshall_arguments()
         if self.validate_field_values(options):
@@ -584,18 +635,21 @@ https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENS
         else:
             return True
 
-    def marshall_arguments(self):
+    def marshall_arguments(self) -> Namespace:
         new_opts = Namespace()
 
         for attr in [
             "ram",
             "vram",
+            "convert_cache",
             "outdir",
         ]:
             if hasattr(self, attr):
                 setattr(new_opts, attr, getattr(self, attr).value)
 
         for attr in self.autoimport_dirs:
+            if not self.autoimport_dirs[attr].value:
+                continue
             directory = Path(self.autoimport_dirs[attr].value)
             if directory.is_relative_to(config.root_path):
                 directory = directory.relative_to(config.root_path)
@@ -615,13 +669,14 @@ https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENS
 
 
 class EditOptApplication(npyscreen.NPSAppManaged):
-    def __init__(self, program_opts: Namespace, invokeai_opts: Namespace):
+    def __init__(self, program_opts: Namespace, invokeai_opts: InvokeAIAppConfig, install_helper: InstallHelper):
         super().__init__()
         self.program_opts = program_opts
         self.invokeai_opts = invokeai_opts
         self.user_cancelled = False
         self.autoload_pending = True
-        self.install_selections = default_user_selections(program_opts)
+        self.install_helper = install_helper
+        self.install_selections = default_user_selections(program_opts, install_helper)
 
     def onStart(self):
         npyscreen.setTheme(npyscreen.Themes.DefaultTheme)
@@ -640,16 +695,10 @@ class EditOptApplication(npyscreen.NPSAppManaged):
                 cycle_widgets=False,
             )
 
-    def new_opts(self):
+    def new_opts(self) -> Namespace:
         return self.options.marshall_arguments()
 
 
-def edit_opts(program_opts: Namespace, invokeai_opts: Namespace) -> argparse.Namespace:
-    editApp = EditOptApplication(program_opts, invokeai_opts)
-    editApp.run()
-    return editApp.new_opts()
-
-
 def default_ramcache() -> float:
     """Run a heuristic for the default RAM cache based on installed RAM."""
 
@@ -660,27 +709,18 @@ def default_ramcache() -> float:
     )  # 2.1 is just large enough for sd 1.5 ;-)
 
 
-def default_startup_options(init_file: Path) -> Namespace:
+def default_startup_options(init_file: Path) -> InvokeAIAppConfig:
     opts = InvokeAIAppConfig.get_config()
-    opts.ram = opts.ram or default_ramcache()
+    opts.ram = default_ramcache()
     return opts
 
 
-def default_user_selections(program_opts: Namespace) -> InstallSelections:
-    try:
-        installer = ModelInstall(config)
-    except omegaconf.errors.ConfigKeyError:
-        logger.warning("Your models.yaml file is corrupt or out of date. Reinitializing")
-        initialize_rootdir(config.root_path, True)
-        installer = ModelInstall(config)
-
-    models = installer.all_models()
+def default_user_selections(program_opts: Namespace, install_helper: InstallHelper) -> InstallSelections:
+    default_model = install_helper.default_model()
+    assert default_model is not None
+    default_models = [default_model] if program_opts.default_only else install_helper.recommended_models()
     return InstallSelections(
-        install_models=[models[installer.default_model()].path or models[installer.default_model()].repo_id]
-        if program_opts.default_only
-        else [models[x].path or models[x].repo_id for x in installer.recommended_models()]
-        if program_opts.yes_to_all
-        else [],
+        install_models=default_models if program_opts.yes_to_all else [],
     )
 
 
@@ -716,21 +756,10 @@ def initialize_rootdir(root: Path, yes_to_all: bool = False):
     path.mkdir(parents=True, exist_ok=True)
 
 
-def maybe_create_models_yaml(root: Path):
-    models_yaml = root / "configs" / "models.yaml"
-    if models_yaml.exists():
-        if OmegaConf.load(models_yaml).get("__metadata__"):  # up to date
-            return
-        else:
-            logger.info("Creating new models.yaml, original saved as models.yaml.orig")
-            models_yaml.rename(models_yaml.parent / "models.yaml.orig")
-
-    with open(models_yaml, "w") as yaml_file:
-        yaml_file.write(yaml.dump({"__metadata__": {"version": "3.0.0"}}))
-
-
 # -------------------------------------
-def run_console_ui(program_opts: Namespace, initfile: Path = None) -> (Namespace, Namespace):
+def run_console_ui(
+    program_opts: Namespace, initfile: Path, install_helper: InstallHelper
+) -> Tuple[Optional[Namespace], Optional[InstallSelections]]:
     invokeai_opts = default_startup_options(initfile)
     invokeai_opts.root = program_opts.root
 
@@ -739,22 +768,16 @@ def run_console_ui(program_opts: Namespace, initfile: Path = None) -> (Namespace
             "Could not increase terminal size. Try running again with a larger window or smaller font size."
         )
 
-    # the install-models application spawns a subprocess to install
-    # models, and will crash unless this is set before running.
-    import torch
-
-    torch.multiprocessing.set_start_method("spawn")
-
-    editApp = EditOptApplication(program_opts, invokeai_opts)
+    editApp = EditOptApplication(program_opts, invokeai_opts, install_helper)
     editApp.run()
     if editApp.user_cancelled:
         return (None, None)
     else:
-        return (editApp.new_opts, editApp.install_selections)
+        return (editApp.new_opts(), editApp.install_selections)
 
 
 # -------------------------------------
-def write_opts(opts: Namespace, init_file: Path):
+def write_opts(opts: InvokeAIAppConfig, init_file: Path) -> None:
     """
     Update the invokeai.yaml file with values from current settings.
     """
@@ -762,7 +785,7 @@ def write_opts(opts: Namespace, init_file: Path):
     new_config = InvokeAIAppConfig.get_config()
     new_config.root = config.root
 
-    for key, value in opts.__dict__.items():
+    for key, value in opts.model_dump().items():
         if hasattr(new_config, key):
             setattr(new_config, key, value)
 
@@ -779,7 +802,7 @@ def default_output_dir() -> Path:
 
 
 # -------------------------------------
-def write_default_options(program_opts: Namespace, initfile: Path):
+def write_default_options(program_opts: Namespace, initfile: Path) -> None:
     opt = default_startup_options(initfile)
     write_opts(opt, initfile)
 
@@ -789,16 +812,11 @@ def write_default_options(program_opts: Namespace, initfile: Path):
 # the legacy Args object in order to parse
 # the old init file and write out the new
 # yaml format.
-def migrate_init_file(legacy_format: Path):
+def migrate_init_file(legacy_format: Path) -> None:
     old = legacy_parser.parse_args([f"@{str(legacy_format)}"])
     new = InvokeAIAppConfig.get_config()
 
-    fields = [
-        x
-        for x, y in InvokeAIAppConfig.model_fields.items()
-        if (y.json_schema_extra.get("category", None) if y.json_schema_extra else None) != "DEPRECATED"
-    ]
-    for attr in fields:
+    for attr in InvokeAIAppConfig.model_fields.keys():
         if hasattr(old, attr):
             try:
                 setattr(new, attr, getattr(old, attr))
@@ -819,7 +837,7 @@ def migrate_init_file(legacy_format: Path):
 
 
 # -------------------------------------
-def migrate_models(root: Path):
+def migrate_models(root: Path) -> None:
     from invokeai.backend.install.migrate_to_3 import do_migrate
 
     do_migrate(root, root)
@@ -838,7 +856,9 @@ def migrate_if_needed(opt: Namespace, root: Path) -> bool:
         ):
             logger.info("** Migrating invokeai.init to invokeai.yaml")
             migrate_init_file(old_init_file)
-            config.parse_args(argv=[], conf=OmegaConf.load(new_init_file))
+            omegaconf = OmegaConf.load(new_init_file)
+            assert isinstance(omegaconf, DictConfig)
+            config.parse_args(argv=[], conf=omegaconf)
 
             if old_hub.exists():
                 migrate_models(config.root_path)
@@ -849,7 +869,7 @@ def migrate_if_needed(opt: Namespace, root: Path) -> bool:
 
 
 # -------------------------------------
-def main() -> None:
+def main():
     parser = argparse.ArgumentParser(description="InvokeAI model downloader")
     parser.add_argument(
         "--skip-sd-weights",
@@ -908,6 +928,7 @@ def main() -> None:
     if opt.full_precision:
         invoke_args.extend(["--precision", "float32"])
     config.parse_args(invoke_args)
+    config.precision = "float32" if opt.full_precision else choose_precision(torch.device(choose_torch_device()))
     logger = InvokeAILogger().get_logger(config=config)
 
     errors = set()
@@ -921,14 +942,18 @@ def main() -> None:
         # run this unconditionally in case new directories need to be added
         initialize_rootdir(config.root_path, opt.yes_to_all)
 
-        models_to_download = default_user_selections(opt)
+        # this will initialize the models.yaml file if not present
+        install_helper = InstallHelper(config, logger)
+
+        models_to_download = default_user_selections(opt, install_helper)
         new_init_file = config.root_path / "invokeai.yaml"
 
         if opt.yes_to_all:
             write_default_options(opt, new_init_file)
             init_options = Namespace(precision="float32" if opt.full_precision else "float16")
+
         else:
-            init_options, models_to_download = run_console_ui(opt, new_init_file)
+            init_options, models_to_download = run_console_ui(opt, new_init_file, install_helper)
             if init_options:
                 write_opts(init_options, new_init_file)
             else:
@@ -943,10 +968,12 @@ def main() -> None:
 
         if opt.skip_sd_weights:
             logger.warning("Skipping diffusion weights download per user request")
+
         elif models_to_download:
-            process_and_execute(opt, models_to_download)
+            install_helper.add_or_delete(models_to_download)
 
         postscript(errors=errors)
+
         if not opt.yes_to_all:
             input("Press any key to continue...")
     except WindowTooSmallException as e: