Next: Remove deprecated app.on_event usage in api runner

The only major version is `query-string`. The breaking change for it is dropping support for old versions of node. Not a problem for us.

.github/workflows/lint-frontend.yml

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

Makefile

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

docs/contributing/INVOCATIONS.md

@@ -9,11 +9,15 @@ 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. 
+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:
 
-Example `nodes`  subfolder structure:
 ```py
 ├── __init__.py # Invoke-managed custom node loader
 │
@@ -30,14 +34,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
@@ -131,7 +135,6 @@ 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")
@@ -167,7 +170,6 @@ 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")
@@ -197,7 +199,6 @@ 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")
@@ -229,30 +230,17 @@ 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 image using InvokeAI's predefined Image Service. Returns the PIL image.
-        image = context.services.images.get_pil_image(self.image.image_name)
+        # Load the input image as a PIL image
+        image = context.images.get_pil(self.image.image_name)
 
-        # Resizing the image
+        # Resize the image
         resized_image = image.resize((self.width, self.height))
 
-        # 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,
-        )
+        # Save the image
+        image_dto = context.images.save(image=resized_image)
 
-        # Returning the Image
-        return ImageOutput(
-            image=ImageField(
-                image_name=output_image.image_name,
-            ),
-            width=output_image.width,
-            height=output_image.height,
-        )
+        # Return an ImageOutput
+        return ImageOutput.build(image_dto)
 ```
 
 **Note:** Do not be overwhelmed by the `ImageOutput` process. InvokeAI has a
@@ -343,27 +331,25 @@ 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 input types in InvokeAI do not meet your requirements. In such cases,
-you can create your own input types.
+existing fields in InvokeAI do not meet your requirements. In such cases, you
+can create your own fields.
 
 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.
 
-**Good Practices**
+### Best 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 input field.
+  `docstring` to describe your field.
 - Keep your field in the same file as the Invocation that it is made for or in
   another file where it is relevant.
 
@@ -378,10 +364,13 @@ class ColorField(BaseModel):
     pass
 ```
 
-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.
+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.
 
 ```python
 class ColorField(BaseModel):
@@ -396,25 +385,11 @@ That's it. We now have a new input field type that we can use in our Invocations
 like this.
 
 ```python
-color: ColorField = Field(default=ColorField(r=0, g=0, b=0, a=0), description='Background color of an image')
+color: ColorField = InputField(default=ColorField(r=0, g=0, b=0, a=0), description='Background color of an image')
 ```
 
-### Custom Components For Frontend
+### Using the custom field
 
-Every backend input type should have a corresponding frontend component so the
-UI knows what to render when you use a particular field type.
+When you start the UI, your custom field will be automatically recognized.
 
-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.
--->
+Custom fields only support connection inputs in the Workflow Editor.

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

invokeai/app/api/dependencies.py

@@ -2,9 +2,12 @@
 
 from logging import Logger
 
-from invokeai.app.services.item_storage.item_storage_memory import ItemStorageMemory
+import torch
+
+from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
+from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
 from invokeai.app.services.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__
 
@@ -12,26 +15,22 @@ from ..services.board_image_records.board_image_records_sqlite import SqliteBoar
 from ..services.board_images.board_images_default import BoardImagesService
 from ..services.board_records.board_records_sqlite import SqliteBoardRecordStorage
 from ..services.boards.boards_default import BoardService
+from ..services.bulk_download.bulk_download_default import BulkDownloadService
 from ..services.config import InvokeAIAppConfig
 from ..services.download import DownloadQueueService
 from ..services.image_files.image_files_disk import DiskImageFileStorage
 from ..services.image_records.image_records_sqlite import SqliteImageRecordStorage
 from ..services.images.images_default import ImageService
 from ..services.invocation_cache.invocation_cache_memory import MemoryInvocationCache
-from ..services.invocation_processor.invocation_processor_default import DefaultInvocationProcessor
-from ..services.invocation_queue.invocation_queue_memory import MemoryInvocationQueue
 from ..services.invocation_services import InvocationServices
 from ..services.invocation_stats.invocation_stats_default import InvocationStatsService
 from ..services.invoker import Invoker
-from ..services.latents_storage.latents_storage_disk import DiskLatentsStorage
-from ..services.latents_storage.latents_storage_forward_cache import ForwardCacheLatentsStorage
-from ..services.model_install import ModelInstallService
 from ..services.model_manager.model_manager_default import ModelManagerService
+from ..services.model_metadata import ModelMetadataStoreSQL
 from ..services.model_records import ModelRecordServiceSQL
 from ..services.names.names_default import SimpleNameService
 from ..services.session_processor.session_processor_default import DefaultSessionProcessor
 from ..services.session_queue.session_queue_sqlite import SqliteSessionQueue
-from ..services.shared.graph import GraphExecutionState
 from ..services.urls.urls_default import LocalUrlService
 from ..services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage
 from .events import FastAPIEventService
@@ -68,6 +67,9 @@ 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)
@@ -80,26 +82,26 @@ class ApiDependencies:
         board_records = SqliteBoardRecordStorage(db=db)
         boards = BoardService()
         events = FastAPIEventService(event_handler_id)
-        graph_execution_manager = ItemStorageMemory[GraphExecutionState]()
+        bulk_download = BulkDownloadService()
         image_records = SqliteImageRecordStorage(db=db)
         images = ImageService()
         invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size)
-        latents = ForwardCacheLatentsStorage(DiskLatentsStorage(f"{output_folder}/latents"))
-        model_manager = ModelManagerService(config, logger)
-        model_record_service = ModelRecordServiceSQL(db=db)
+        tensors = ObjectSerializerForwardCache(
+            ObjectSerializerDisk[torch.Tensor](output_folder / "tensors", ephemeral=True)
+        )
+        conditioning = ObjectSerializerForwardCache(
+            ObjectSerializerDisk[ConditioningFieldData](output_folder / "conditioning", ephemeral=True)
+        )
         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()
-        processor = DefaultInvocationProcessor()
-        queue = MemoryInvocationQueue()
         session_processor = DefaultSessionProcessor()
         session_queue = SqliteSessionQueue(db=db)
         urls = LocalUrlService()
@@ -110,27 +112,24 @@ class ApiDependencies:
             board_images=board_images,
             board_records=board_records,
             boards=boards,
+            bulk_download=bulk_download,
             configuration=configuration,
             events=events,
-            graph_execution_manager=graph_execution_manager,
             image_files=image_files,
             image_records=image_records,
             images=images,
             invocation_cache=invocation_cache,
-            latents=latents,
             logger=logger,
             model_manager=model_manager,
-            model_records=model_record_service,
             download_queue=download_queue_service,
-            model_install=model_install_service,
             names=names,
             performance_statistics=performance_statistics,
-            processor=processor,
-            queue=queue,
             session_processor=session_processor,
             session_queue=session_queue,
             urls=urls,
             workflow_records=workflow_records,
+            tensors=tensors,
+            conditioning=conditioning,
         )
 
         ApiDependencies.invoker = Invoker(services)

invokeai/app/api/routers/download_queue.py

@@ -36,7 +36,7 @@ async def list_downloads() -> List[DownloadJob]:
         400: {"description": "Bad request"},
     },
 )
-async def prune_downloads():
+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

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

invokeai/app/api/routers/model_manager.py

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

invokeai/app/api/routers/model_records.py

@@ -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/routers/sessions.py

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

invokeai/app/api/sockets.py

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

invokeai/app/api_app.py

@@ -2,10 +2,12 @@
 # which are imported/used before parse_args() is called will get the default config values instead of the
 # values from the command line or config file.
 import sys
+from contextlib import asynccontextmanager
 
 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()
@@ -47,18 +49,14 @@ 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,
         workflows,
     )
     from .api.sockets import SocketIO
     from .invocations.baseinvocation import (
         BaseInvocation,
-        InputFieldJSONSchemaExtra,
-        OutputFieldJSONSchemaExtra,
         UIConfigBase,
     )
 
@@ -74,9 +72,25 @@ logger = InvokeAILogger.get_logger(config=app_config)
 mimetypes.add_type("application/javascript", ".js")
 mimetypes.add_type("text/css", ".css")
 
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Add startup event to load dependencies
+    ApiDependencies.initialize(config=app_config, event_handler_id=event_handler_id, logger=logger)
+    yield
+    # Shut down threads
+    ApiDependencies.shutdown()
+
+
 # Create the app
 # TODO: create this all in a method so configuration/etc. can be passed in?
-app = FastAPI(title="Invoke - Community Edition", docs_url=None, redoc_url=None, separate_input_output_schemas=False)
+app = FastAPI(
+    title="Invoke - Community Edition",
+    docs_url=None,
+    redoc_url=None,
+    separate_input_output_schemas=False,
+    lifespan=lifespan,
+)
 
 # Add event handler
 event_handler_id: int = id(app)
@@ -99,24 +113,9 @@ app.add_middleware(
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 
 
-# Add startup event to load dependencies
-@app.on_event("startup")
-async def startup_event() -> None:
-    ApiDependencies.initialize(config=app_config, event_handler_id=event_handler_id, logger=logger)
-
-
-# Shut down threads
-@app.on_event("shutdown")
-async def shutdown_event() -> None:
-    ApiDependencies.shutdown()
-
-
 # Include all routers
-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")
@@ -154,6 +153,8 @@ def custom_openapi() -> dict[str, Any]:
         # TODO: note that we assume the schema_key here is the TYPE.__name__
         # This could break in some cases, figure out a better way to do it
         output_type_titles[schema_key] = output_schema["title"]
+        openapi_schema["components"]["schemas"][schema_key] = output_schema
+        openapi_schema["components"]["schemas"][schema_key]["class"] = "output"
 
     # Add Node Editor UI helper schemas
     ui_config_schemas = models_json_schema(
@@ -176,23 +177,24 @@ def custom_openapi() -> dict[str, Any]:
         outputs_ref = {"$ref": f"#/components/schemas/{output_type_title}"}
         invoker_schema["output"] = outputs_ref
         invoker_schema["class"] = "invocation"
-        openapi_schema["components"]["schemas"][f"{output_type_title}"]["class"] = "output"
 
-    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

@@ -8,17 +8,33 @@ import warnings
 from abc import ABC, abstractmethod
 from enum import Enum
 from inspect import signature
-from types import UnionType
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Literal, Optional, Type, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Callable,
+    ClassVar,
+    Iterable,
+    Literal,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import semver
-from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, create_model
-from pydantic.fields import FieldInfo, _Unset
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, create_model
+from pydantic.fields import FieldInfo
 from pydantic_core import PydanticUndefined
+from typing_extensions import TypeAliasType
 
+from invokeai.app.invocations.fields import (
+    FieldKind,
+    Input,
+)
 from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
-from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.metaenum import MetaEnum
 from invokeai.app.util.misc import uuid_string
 from invokeai.backend.util.logging import InvokeAILogger
@@ -52,393 +68,6 @@ 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.
@@ -460,33 +89,6 @@ 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.
@@ -495,6 +97,7 @@ class BaseInvocationOutput(BaseModel):
     """
 
     _output_classes: ClassVar[set[BaseInvocationOutput]] = set()
+    _typeadapter: ClassVar[Optional[TypeAdapter[Any]]] = None
 
     @classmethod
     def register_output(cls, output: BaseInvocationOutput) -> None:
@@ -507,10 +110,14 @@ class BaseInvocationOutput(BaseModel):
         return cls._output_classes
 
     @classmethod
-    def get_outputs_union(cls) -> UnionType:
-        """Gets a union of all invocation outputs."""
-        outputs_union = Union[tuple(cls._output_classes)]  # type: ignore [valid-type]
-        return outputs_union  # type: ignore [return-value]
+    def get_typeadapter(cls) -> TypeAdapter[Any]:
+        """Gets a pydantc TypeAdapter for the union of all invocation output types."""
+        if not cls._typeadapter:
+            InvocationOutputsUnion = TypeAliasType(
+                "InvocationOutputsUnion", Annotated[Union[tuple(cls._output_classes)], Field(discriminator="type")]
+            )
+            cls._typeadapter = TypeAdapter(InvocationOutputsUnion)
+        return cls._typeadapter
 
     @classmethod
     def get_output_types(cls) -> Iterable[str]:
@@ -559,6 +166,7 @@ class BaseInvocation(ABC, BaseModel):
     """
 
     _invocation_classes: ClassVar[set[BaseInvocation]] = set()
+    _typeadapter: ClassVar[Optional[TypeAdapter[Any]]] = None
 
     @classmethod
     def get_type(cls) -> str:
@@ -571,10 +179,14 @@ class BaseInvocation(ABC, BaseModel):
         cls._invocation_classes.add(invocation)
 
     @classmethod
-    def get_invocations_union(cls) -> UnionType:
-        """Gets a union of all invocation types."""
-        invocations_union = Union[tuple(cls._invocation_classes)]  # type: ignore [valid-type]
-        return invocations_union  # type: ignore [return-value]
+    def get_typeadapter(cls) -> TypeAdapter[Any]:
+        """Gets a pydantc TypeAdapter for the union of all invocation types."""
+        if not cls._typeadapter:
+            InvocationsUnion = TypeAliasType(
+                "InvocationsUnion", Annotated[Union[tuple(cls._invocation_classes)], Field(discriminator="type")]
+            )
+            cls._typeadapter = TypeAdapter(InvocationsUnion)
+        return cls._typeadapter
 
     @classmethod
     def get_invocations(cls) -> Iterable[BaseInvocation]:
@@ -632,7 +244,7 @@ class BaseInvocation(ABC, BaseModel):
         """Invoke with provided context and return outputs."""
         pass
 
-    def invoke_internal(self, context: InvocationContext) -> BaseInvocationOutput:
+    def invoke_internal(self, context: InvocationContext, services: "InvocationServices") -> BaseInvocationOutput:
         """
         Internal invoke method, calls `invoke()` after some prep.
         Handles optional fields that are required to call `invoke()` and invocation cache.
@@ -657,23 +269,23 @@ class BaseInvocation(ABC, BaseModel):
                     raise MissingInputException(self.model_fields["type"].default, field_name)
 
         # skip node cache codepath if it's disabled
-        if context.services.configuration.node_cache_size == 0:
+        if services.configuration.node_cache_size == 0:
             return self.invoke(context)
 
         output: BaseInvocationOutput
         if self.use_cache:
-            key = context.services.invocation_cache.create_key(self)
-            cached_value = context.services.invocation_cache.get(key)
+            key = services.invocation_cache.create_key(self)
+            cached_value = services.invocation_cache.get(key)
             if cached_value is None:
-                context.services.logger.debug(f'Invocation cache miss for type "{self.get_type()}": {self.id}')
+                services.logger.debug(f'Invocation cache miss for type "{self.get_type()}": {self.id}')
                 output = self.invoke(context)
-                context.services.invocation_cache.save(key, output)
+                services.invocation_cache.save(key, output)
                 return output
             else:
-                context.services.logger.debug(f'Invocation cache hit for type "{self.get_type()}": {self.id}')
+                services.logger.debug(f'Invocation cache hit for type "{self.get_type()}": {self.id}')
                 return cached_value
         else:
-            context.services.logger.debug(f'Skipping invocation cache for "{self.get_type()}": {self.id}')
+            services.logger.debug(f'Skipping invocation cache for "{self.get_type()}": {self.id}')
             return self.invoke(context)
 
     id: str = Field(
@@ -714,9 +326,7 @@ RESERVED_NODE_ATTRIBUTE_FIELD_NAMES = {
     "workflow",
 }
 
-RESERVED_INPUT_FIELD_NAMES = {
-    "metadata",
-}
+RESERVED_INPUT_FIELD_NAMES = {"metadata", "board"}
 
 RESERVED_OUTPUT_FIELD_NAMES = {"type"}
 
@@ -926,37 +536,3 @@ 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,9 +5,11 @@ 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, InputField, InvocationContext, invocation
+from .baseinvocation import BaseInvocation, invocation
+from .fields import InputField
 
 
 @invocation(

invokeai/app/invocations/compel.py

@@ -1,40 +1,28 @@
-from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import Iterator, List, Optional, Tuple, Union, cast
 
 import torch
 from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
+from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
-from invokeai.app.invocations.primitives import ConditioningField, ConditioningOutput
-from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIComponent
+from invokeai.app.invocations.primitives import ConditioningOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.util.ti_utils import generate_ti_list
+from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
+    ConditioningFieldData,
     ExtraConditioningInfo,
     SDXLConditioningInfo,
 )
+from invokeai.backend.util.devices import torch_dtype
 
-from ...backend.model_management.lora import ModelPatcher
-from ...backend.model_management.models import ModelNotFoundException, ModelType
-from ...backend.util.devices import torch_dtype
-from ..util.ti_utils import extract_ti_triggers_from_prompt
-from .baseinvocation import (
-    BaseInvocation,
-    BaseInvocationOutput,
-    Input,
-    InputField,
-    InvocationContext,
-    OutputField,
-    UIComponent,
-    invocation,
-    invocation_output,
-)
+from .baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
 from .model import ClipField
 
-
-@dataclass
-class ConditioningFieldData:
-    conditionings: List[BasicConditioningInfo]
-    # unconditioned: Optional[torch.Tensor]
+# unconditioned: Optional[torch.Tensor]
 
 
 # class ConditioningAlgo(str, Enum):
@@ -48,7 +36,7 @@ class ConditioningFieldData:
     title="Prompt",
     tags=["prompt", "compel"],
     category="conditioning",
-    version="1.0.0",
+    version="1.0.1",
 )
 class CompelInvocation(BaseInvocation):
     """Parse prompt using compel package to conditioning."""
@@ -66,49 +54,27 @@ class CompelInvocation(BaseInvocation):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
-        tokenizer_info = context.services.model_manager.get_model(
-            **self.clip.tokenizer.model_dump(),
-            context=context,
-        )
-        text_encoder_info = context.services.model_manager.get_model(
-            **self.clip.text_encoder.model_dump(),
-            context=context,
-        )
+        tokenizer_info = context.models.load(**self.clip.tokenizer.model_dump())
+        tokenizer_model = tokenizer_info.model
+        assert isinstance(tokenizer_model, CLIPTokenizer)
+        text_encoder_info = context.models.load(**self.clip.text_encoder.model_dump())
+        text_encoder_model = text_encoder_info.model
+        assert isinstance(text_encoder_model, CLIPTextModel)
 
-        def _lora_loader():
+        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.clip.loras:
-                lora_info = context.services.model_manager.get_model(
-                    **lora.model_dump(exclude={"weight"}), context=context
-                )
-                yield (lora_info.context.model, lora.weight)
+                lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
+                assert isinstance(lora_info.model, LoRAModelRaw)
+                yield (lora_info.model, lora.weight)
                 del lora_info
             return
 
-        # loras = [(context.services.model_manager.get_model(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras]
+        # loras = [(context.models.get(**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.services.model_manager.get_model(
-                            model_name=name,
-                            base_model=self.clip.text_encoder.base_model,
-                            model_type=ModelType.TextualInversion,
-                            context=context,
-                        ).context.model,
-                    )
-                )
-            except ModelNotFoundException:
-                # print(e)
-                # import traceback
-                # print(traceback.format_exc())
-                print(f'Warn: trigger: "{trigger}" not found')
+        ti_list = generate_ti_list(self.prompt, text_encoder_info.config.base, context)
 
         with (
-            ModelPatcher.apply_ti(tokenizer_info.context.model, text_encoder_info.context.model, ti_list) as (
+            ModelPatcher.apply_ti(tokenizer_model, text_encoder_model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -116,8 +82,9 @@ 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_model, self.clip.skipped_layers),
         ):
+            assert isinstance(text_encoder, CLIPTextModel)
             compel = Compel(
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
@@ -128,7 +95,7 @@ class CompelInvocation(BaseInvocation):
 
             conjunction = Compel.parse_prompt_string(self.prompt)
 
-            if context.services.configuration.log_tokenization:
+            if context.config.get().log_tokenization:
                 log_tokenization_for_conjunction(conjunction, tokenizer)
 
             c, options = compel.build_conditioning_tensor_for_conjunction(conjunction)
@@ -149,17 +116,14 @@ class CompelInvocation(BaseInvocation):
             ]
         )
 
-        conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning"
-        context.services.latents.save(conditioning_name, conditioning_data)
+        conditioning_name = context.conditioning.save(conditioning_data)
 
-        return ConditioningOutput(
-            conditioning=ConditioningField(
-                conditioning_name=conditioning_name,
-            ),
-        )
+        return ConditioningOutput.build(conditioning_name)
 
 
 class SDXLPromptInvocationBase:
+    """Prompt processor for SDXL models."""
+
     def run_clip_compel(
         self,
         context: InvocationContext,
@@ -168,26 +132,25 @@ class SDXLPromptInvocationBase:
         get_pooled: bool,
         lora_prefix: str,
         zero_on_empty: bool,
-    ):
-        tokenizer_info = context.services.model_manager.get_model(
-            **clip_field.tokenizer.model_dump(),
-            context=context,
-        )
-        text_encoder_info = context.services.model_manager.get_model(
-            **clip_field.text_encoder.model_dump(),
-            context=context,
-        )
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[ExtraConditioningInfo]]:
+        tokenizer_info = context.models.load(**clip_field.tokenizer.model_dump())
+        tokenizer_model = tokenizer_info.model
+        assert isinstance(tokenizer_model, CLIPTokenizer)
+        text_encoder_info = context.models.load(**clip_field.text_encoder.model_dump())
+        text_encoder_model = text_encoder_info.model
+        assert isinstance(text_encoder_model, (CLIPTextModel, CLIPTextModelWithProjection))
 
         # 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(
@@ -198,40 +161,21 @@ 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.services.model_manager.get_model(
-                    **lora.model_dump(exclude={"weight"}), context=context
-                )
-                yield (lora_info.context.model, lora.weight)
+                lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
+                lora_model = lora_info.model
+                assert isinstance(lora_model, LoRAModelRaw)
+                yield (lora_model, lora.weight)
                 del lora_info
             return
 
-        # loras = [(context.services.model_manager.get_model(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras]
+        # loras = [(context.models.get(**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.services.model_manager.get_model(
-                            model_name=name,
-                            base_model=clip_field.text_encoder.base_model,
-                            model_type=ModelType.TextualInversion,
-                            context=context,
-                        ).context.model,
-                    )
-                )
-            except ModelNotFoundException:
-                # print(e)
-                # import traceback
-                # print(traceback.format_exc())
-                print(f'Warn: trigger: "{trigger}" not found')
+        ti_list = generate_ti_list(prompt, text_encoder_info.config.base, context)
 
         with (
-            ModelPatcher.apply_ti(tokenizer_info.context.model, text_encoder_info.context.model, ti_list) as (
+            ModelPatcher.apply_ti(tokenizer_model, text_encoder_model, ti_list) as (
                 tokenizer,
                 ti_manager,
             ),
@@ -239,8 +183,10 @@ 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_model, clip_field.skipped_layers),
         ):
+            assert isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection))
+            text_encoder = cast(CLIPTextModel, text_encoder)
             compel = Compel(
                 tokenizer=tokenizer,
                 text_encoder=text_encoder,
@@ -253,7 +199,7 @@ class SDXLPromptInvocationBase:
 
             conjunction = Compel.parse_prompt_string(prompt)
 
-            if context.services.configuration.log_tokenization:
+            if context.config.get().log_tokenization:
                 # TODO: better logging for and syntax
                 log_tokenization_for_conjunction(conjunction, tokenizer)
 
@@ -286,7 +232,7 @@ class SDXLPromptInvocationBase:
     title="SDXL Prompt",
     tags=["sdxl", "compel", "prompt"],
     category="conditioning",
-    version="1.0.0",
+    version="1.0.1",
 )
 class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
@@ -357,6 +303,7 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
                 dim=1,
             )
 
+        assert c2_pooled is not None
         conditioning_data = ConditioningFieldData(
             conditionings=[
                 SDXLConditioningInfo(
@@ -368,14 +315,9 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
             ]
         )
 
-        conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning"
-        context.services.latents.save(conditioning_name, conditioning_data)
+        conditioning_name = context.conditioning.save(conditioning_data)
 
-        return ConditioningOutput(
-            conditioning=ConditioningField(
-                conditioning_name=conditioning_name,
-            ),
-        )
+        return ConditioningOutput.build(conditioning_name)
 
 
 @invocation(
@@ -383,7 +325,7 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     title="SDXL Refiner Prompt",
     tags=["sdxl", "compel", "prompt"],
     category="conditioning",
-    version="1.0.0",
+    version="1.0.1",
 )
 class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
@@ -410,6 +352,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(
@@ -421,14 +364,9 @@ class SDXLRefinerCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase
             ]
         )
 
-        conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning"
-        context.services.latents.save(conditioning_name, conditioning_data)
+        conditioning_name = context.conditioning.save(conditioning_data)
 
-        return ConditioningOutput(
-            conditioning=ConditioningField(
-                conditioning_name=conditioning_name,
-            ),
-        )
+        return ConditioningOutput.build(conditioning_name)
 
 
 @invocation_output("clip_skip_output")
@@ -449,7 +387,7 @@ class ClipSkipInvocation(BaseInvocation):
     """Skip layers in clip text_encoder model."""
 
     clip: ClipField = InputField(description=FieldDescriptions.clip, input=Input.Connection, title="CLIP")
-    skipped_layers: int = InputField(default=0, description=FieldDescriptions.skipped_layers)
+    skipped_layers: int = InputField(default=0, ge=0, description=FieldDescriptions.skipped_layers)
 
     def invoke(self, context: InvocationContext) -> ClipSkipInvocationOutput:
         self.clip.skipped_layers += self.skipped_layers
@@ -459,9 +397,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
@@ -473,7 +411,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")
 
@@ -486,24 +426,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
@@ -543,7 +488,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

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

invokeai/app/invocations/controlnet_image_processors.py

@@ -23,27 +23,24 @@ from controlnet_aux import (
 )
 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.primitives import ImageField, ImageOutput
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    ImageField,
+    Input,
+    InputField,
+    OutputField,
+    WithBoard,
+    WithMetadata,
+)
+from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
-from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
-from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.image_util.depth_anything import DepthAnythingDetector
 from invokeai.backend.image_util.dw_openpose import DWOpenposeDetector
 
-from ...backend.model_management import BaseModelType
-from .baseinvocation import (
-    BaseInvocation,
-    BaseInvocationOutput,
-    Input,
-    InputField,
-    InvocationContext,
-    OutputField,
-    WithMetadata,
-    invocation,
-    invocation_output,
-)
+from .baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
 
 CONTROLNET_MODE_VALUES = Literal["balanced", "more_prompt", "more_control", "unbalanced"]
 CONTROLNET_RESIZE_VALUES = Literal[
@@ -57,10 +54,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):
@@ -140,7 +134,7 @@ class ControlNetInvocation(BaseInvocation):
 
 
 # This invocation exists for other invocations to subclass it - do not register with @invocation!
-class ImageProcessorInvocation(BaseInvocation, WithMetadata):
+class ImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Base class for invocations that preprocess images for ControlNet"""
 
     image: ImageField = InputField(description="The image to process")
@@ -149,23 +143,18 @@ class ImageProcessorInvocation(BaseInvocation, WithMetadata):
         # superclass just passes through image without processing
         return image
 
+    def load_image(self, context: InvocationContext) -> Image.Image:
+        # allows override for any special formatting specific to the preprocessor
+        return context.images.get_pil(self.image.image_name, "RGB")
+
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        raw_image = context.services.images.get_pil_image(self.image.image_name)
+        raw_image = self.load_image(context)
         # 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.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,
-        )
+        image_dto = context.images.save(image=processed_image)
 
         """Builds an ImageOutput and its ImageField"""
         processed_image_field = ImageField(image_name=image_dto.image_name)
@@ -184,7 +173,7 @@ class ImageProcessorInvocation(BaseInvocation, WithMetadata):
     title="Canny Processor",
     tags=["controlnet", "canny"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class CannyImageProcessorInvocation(ImageProcessorInvocation):
     """Canny edge detection for ControlNet"""
@@ -196,6 +185,10 @@ class CannyImageProcessorInvocation(ImageProcessorInvocation):
         default=200, ge=0, le=255, description="The high threshold of the Canny pixel gradient (0-255)"
     )
 
+    def load_image(self, context: InvocationContext) -> Image.Image:
+        # Keep alpha channel for Canny processing to detect edges of transparent areas
+        return context.images.get_pil(self.image.image_name, "RGBA")
+
     def run_processor(self, image):
         canny_processor = CannyDetector()
         processed_image = canny_processor(image, self.low_threshold, self.high_threshold)
@@ -207,7 +200,7 @@ class CannyImageProcessorInvocation(ImageProcessorInvocation):
     title="HED (softedge) Processor",
     tags=["controlnet", "hed", "softedge"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class HedImageProcessorInvocation(ImageProcessorInvocation):
     """Applies HED edge detection to image"""
@@ -236,7 +229,7 @@ class HedImageProcessorInvocation(ImageProcessorInvocation):
     title="Lineart Processor",
     tags=["controlnet", "lineart"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class LineartImageProcessorInvocation(ImageProcessorInvocation):
     """Applies line art processing to image"""
@@ -258,7 +251,7 @@ class LineartImageProcessorInvocation(ImageProcessorInvocation):
     title="Lineart Anime Processor",
     tags=["controlnet", "lineart", "anime"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation):
     """Applies line art anime processing to image"""
@@ -281,7 +274,7 @@ class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation):
     title="Midas Depth Processor",
     tags=["controlnet", "midas"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class MidasDepthImageProcessorInvocation(ImageProcessorInvocation):
     """Applies Midas depth processing to image"""
@@ -308,7 +301,7 @@ class MidasDepthImageProcessorInvocation(ImageProcessorInvocation):
     title="Normal BAE Processor",
     tags=["controlnet"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class NormalbaeImageProcessorInvocation(ImageProcessorInvocation):
     """Applies NormalBae processing to image"""
@@ -325,7 +318,7 @@ class NormalbaeImageProcessorInvocation(ImageProcessorInvocation):
 
 
 @invocation(
-    "mlsd_image_processor", title="MLSD Processor", tags=["controlnet", "mlsd"], category="controlnet", version="1.2.0"
+    "mlsd_image_processor", title="MLSD Processor", tags=["controlnet", "mlsd"], category="controlnet", version="1.2.1"
 )
 class MlsdImageProcessorInvocation(ImageProcessorInvocation):
     """Applies MLSD processing to image"""
@@ -348,7 +341,7 @@ class MlsdImageProcessorInvocation(ImageProcessorInvocation):
 
 
 @invocation(
-    "pidi_image_processor", title="PIDI Processor", tags=["controlnet", "pidi"], category="controlnet", version="1.2.0"
+    "pidi_image_processor", title="PIDI Processor", tags=["controlnet", "pidi"], category="controlnet", version="1.2.1"
 )
 class PidiImageProcessorInvocation(ImageProcessorInvocation):
     """Applies PIDI processing to image"""
@@ -375,7 +368,7 @@ class PidiImageProcessorInvocation(ImageProcessorInvocation):
     title="Content Shuffle Processor",
     tags=["controlnet", "contentshuffle"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation):
     """Applies content shuffle processing to image"""
@@ -405,7 +398,7 @@ class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation):
     title="Zoe (Depth) Processor",
     tags=["controlnet", "zoe", "depth"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation):
     """Applies Zoe depth processing to image"""
@@ -421,7 +414,7 @@ class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation):
     title="Mediapipe Face Processor",
     tags=["controlnet", "mediapipe", "face"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     """Applies mediapipe face processing to image"""
@@ -430,10 +423,6 @@ class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     min_confidence: float = InputField(default=0.5, ge=0, le=1, description="Minimum confidence for face detection")
 
     def run_processor(self, image):
-        # MediaPipeFaceDetector throws an error if image has alpha channel
-        #     so convert to RGB if needed
-        if image.mode == "RGBA":
-            image = image.convert("RGB")
         mediapipe_face_processor = MediapipeFaceDetector()
         processed_image = mediapipe_face_processor(image, max_faces=self.max_faces, min_confidence=self.min_confidence)
         return processed_image
@@ -444,7 +433,7 @@ class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     title="Leres (Depth) Processor",
     tags=["controlnet", "leres", "depth"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class LeresImageProcessorInvocation(ImageProcessorInvocation):
     """Applies leres processing to image"""
@@ -473,7 +462,7 @@ class LeresImageProcessorInvocation(ImageProcessorInvocation):
     title="Tile Resample Processor",
     tags=["controlnet", "tile"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class TileResamplerProcessorInvocation(ImageProcessorInvocation):
     """Tile resampler processor"""
@@ -513,7 +502,7 @@ class TileResamplerProcessorInvocation(ImageProcessorInvocation):
     title="Segment Anything Processor",
     tags=["controlnet", "segmentanything"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class SegmentAnythingProcessorInvocation(ImageProcessorInvocation):
     """Applies segment anything processing to image"""
@@ -555,7 +544,7 @@ class SamDetectorReproducibleColors(SamDetector):
     title="Color Map Processor",
     tags=["controlnet"],
     category="controlnet",
-    version="1.2.0",
+    version="1.2.1",
 )
 class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
     """Generates a color map from the provided image"""
@@ -563,7 +552,6 @@ class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
     color_map_tile_size: int = InputField(default=64, ge=0, description=FieldDescriptions.tile_size)
 
     def run_processor(self, image: Image.Image):
-        image = image.convert("RGB")
         np_image = np.array(image, dtype=np.uint8)
         height, width = np_image.shape[:2]
 
@@ -603,9 +591,6 @@ class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
         depth_anything_detector = DepthAnythingDetector()
         depth_anything_detector.load_model(model_size=self.model_size)
 
-        if image.mode == "RGBA":
-            image = image.convert("RGB")
-
         processed_image = depth_anything_detector(image=image, resolution=self.resolution, offload=self.offload)
         return processed_image
 
@@ -625,7 +610,7 @@ class DWOpenposeImageProcessorInvocation(ImageProcessorInvocation):
     draw_hands: bool = InputField(default=False)
     image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
 
-    def run_processor(self, image):
+    def run_processor(self, image: Image.Image):
         dw_openpose = DWOpenposeDetector()
         processed_image = dw_openpose(
             image,

invokeai/app/invocations/cv.py

@@ -5,22 +5,24 @@ import cv2 as cv
 import numpy
 from PIL import Image, ImageOps
 
-from invokeai.app.invocations.primitives import ImageField, ImageOutput
-from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
+from invokeai.app.invocations.fields import ImageField
+from invokeai.app.invocations.primitives import ImageOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
 
-from .baseinvocation import BaseInvocation, InputField, InvocationContext, WithMetadata, invocation
+from .baseinvocation import BaseInvocation, invocation
+from .fields import InputField, WithBoard, WithMetadata
 
 
-@invocation("cv_inpaint", title="OpenCV Inpaint", tags=["opencv", "inpaint"], category="inpaint", version="1.2.0")
-class CvInpaintInvocation(BaseInvocation, WithMetadata):
+@invocation("cv_inpaint", title="OpenCV Inpaint", tags=["opencv", "inpaint"], category="inpaint", version="1.2.1")
+class CvInpaintInvocation(BaseInvocation, WithMetadata, WithBoard):
     """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.services.images.get_pil_image(self.image.image_name)
-        mask = context.services.images.get_pil_image(self.mask.image_name)
+        image = context.images.get_pil(self.image.image_name)
+        mask = context.images.get_pil(self.mask.image_name)
 
         # Convert to cv image/mask
         # TODO: consider making these utility functions
@@ -34,18 +36,6 @@ class CvInpaintInvocation(BaseInvocation, WithMetadata):
         # TODO: consider making a utility function
         image_inpainted = Image.fromarray(cv.cvtColor(cv_inpainted, cv.COLOR_BGR2RGB))
 
-        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,
-        )
+        image_dto = context.images.save(image=image_inpainted)
 
-        return ImageOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/facetools.py

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

invokeai/app/invocations/fields.py

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

invokeai/app/invocations/infill.py

@@ -6,14 +6,16 @@ from typing import Literal, Optional, get_args
 import numpy as np
 from PIL import Image, ImageOps
 
-from invokeai.app.invocations.primitives import ColorField, ImageField, ImageOutput
-from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
+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.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, InputField, InvocationContext, WithMetadata, invocation
+from .baseinvocation import BaseInvocation, invocation
+from .fields import InputField, WithBoard, WithMetadata
 from .image import PIL_RESAMPLING_MAP, PIL_RESAMPLING_MODES
 
 
@@ -118,8 +120,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.0")
-class InfillColorInvocation(BaseInvocation, WithMetadata):
+@invocation("infill_rgba", title="Solid Color Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
+class InfillColorInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image with a solid color"""
 
     image: ImageField = InputField(description="The image to infill")
@@ -129,33 +131,20 @@ class InfillColorInvocation(BaseInvocation, WithMetadata):
     )
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get_pil_image(self.image.image_name)
+        image = context.images.get_pil(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.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,
-        )
+        image_dto = context.images.save(image=infilled)
 
-        return ImageOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        return ImageOutput.build(image_dto)
 
 
-@invocation("infill_tile", title="Tile Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
-class InfillTileInvocation(BaseInvocation, WithMetadata):
+@invocation("infill_tile", title="Tile Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.2")
+class InfillTileInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image with tiles of the image"""
 
     image: ImageField = InputField(description="The image to infill")
@@ -168,33 +157,20 @@ class InfillTileInvocation(BaseInvocation, WithMetadata):
     )
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get_pil_image(self.image.image_name)
+        image = context.images.get_pil(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.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,
-        )
+        image_dto = context.images.save(image=infilled)
 
-        return ImageOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        return ImageOutput.build(image_dto)
 
 
 @invocation(
-    "infill_patchmatch", title="PatchMatch Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.0"
+    "infill_patchmatch", title="PatchMatch Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1"
 )
-class InfillPatchMatchInvocation(BaseInvocation, WithMetadata):
+class InfillPatchMatchInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Infills transparent areas of an image using the PatchMatch algorithm"""
 
     image: ImageField = InputField(description="The image to infill")
@@ -202,7 +178,7 @@ class InfillPatchMatchInvocation(BaseInvocation, WithMetadata):
     resample_mode: PIL_RESAMPLING_MODES = InputField(default="bicubic", description="The resampling mode")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get_pil_image(self.image.image_name).convert("RGBA")
+        image = context.images.get_pil(self.image.image_name).convert("RGBA")
 
         resample_mode = PIL_RESAMPLING_MAP[self.resample_mode]
 
@@ -227,77 +203,38 @@ class InfillPatchMatchInvocation(BaseInvocation, WithMetadata):
         infilled.paste(image, (0, 0), mask=image.split()[-1])
         # image.paste(infilled, (0, 0), mask=image.split()[-1])
 
-        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,
-        )
+        image_dto = context.images.save(image=infilled)
 
-        return ImageOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        return ImageOutput.build(image_dto)
 
 
-@invocation("infill_lama", title="LaMa Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.0")
-class LaMaInfillInvocation(BaseInvocation, WithMetadata):
+@invocation("infill_lama", title="LaMa Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
+class LaMaInfillInvocation(BaseInvocation, WithMetadata, WithBoard):
     """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.services.images.get_pil_image(self.image.image_name)
+        image = context.images.get_pil(self.image.image_name)
 
         infilled = infill_lama(image.copy())
 
-        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,
-        )
+        image_dto = context.images.save(image=infilled)
 
-        return ImageOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        return ImageOutput.build(image_dto)
 
 
-@invocation("infill_cv2", title="CV2 Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.0")
-class CV2InfillInvocation(BaseInvocation, WithMetadata):
+@invocation("infill_cv2", title="CV2 Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.1")
+class CV2InfillInvocation(BaseInvocation, WithMetadata, WithBoard):
     """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.services.images.get_pil_image(self.image.image_name)
+        image = context.images.get_pil(self.image.image_name)
 
         infilled = infill_cv2(image.copy())
 
-        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,
-        )
+        image_dto = context.images.save(image=infilled)
 
-        return ImageOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/ip_adapter.py

@@ -1,38 +1,29 @@
-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.shared.fields import FieldDescriptions
-from invokeai.backend.model_management.models.base import BaseModelType, ModelType
-from invokeai.backend.model_management.models.ip_adapter import get_ip_adapter_image_encoder_model_id
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.config 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):
@@ -49,12 +40,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
 
@@ -65,7 +56,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.1")
+@invocation("ip_adapter", title="IP-Adapter", tags=["ip_adapter", "control"], category="ip_adapter", version="1.1.2")
 class IPAdapterInvocation(BaseInvocation):
     """Collects IP-Adapter info to pass to other nodes."""
 
@@ -87,33 +78,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.services.model_manager.model_info(
-            self.ip_adapter_model.model_name, self.ip_adapter_model.base_model, ModelType.IPAdapter
-        )
-        # HACK(ryand): This is bad for a couple of reasons: 1) we are bypassing the model manager to read the model
-        # directly, and 2) we are reading from disk every time this invocation is called without caching the result.
-        # A better solution would be to store the image encoder model reference in the IP-Adapter model info, but this
-        # is currently messy due to differences between how the model info is generated when installing a model from
-        # disk vs. downloading the model.
-        image_encoder_model_id = get_ip_adapter_image_encoder_model_id(
-            os.path.join(context.services.configuration.get_config().models_path, ip_adapter_info["path"])
-        )
+        ip_adapter_info = context.models.get_config(self.ip_adapter_model.key)
+        image_encoder_model_id = ip_adapter_info.image_encoder_model_id
         image_encoder_model_name = image_encoder_model_id.split("/")[-1].strip()
-        image_encoder_model = CLIPVisionModelField(
-            model_name=image_encoder_model_name,
-            base_model=BaseModelType.Any,
+        image_encoder_models = context.models.search_by_attrs(
+            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,34 +20,44 @@ from diffusers.models.attention_processor import (
     LoRAXFormersAttnProcessor,
     XFormersAttnProcessor,
 )
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.schedulers import DPMSolverSDEScheduler
 from diffusers.schedulers import SchedulerMixin as Scheduler
+from PIL import Image, ImageFilter
 from pydantic import field_validator
 from torchvision.transforms.functional import resize as tv_resize
 
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR, SCHEDULER_NAME_VALUES
+from invokeai.app.invocations.fields import (
+    ConditioningField,
+    DenoiseMaskField,
+    FieldDescriptions,
+    ImageField,
+    Input,
+    InputField,
+    LatentsField,
+    OutputField,
+    UIType,
+    WithBoard,
+    WithMetadata,
+)
 from invokeai.app.invocations.ip_adapter import IPAdapterField
 from invokeai.app.invocations.primitives import (
-    DenoiseMaskField,
     DenoiseMaskOutput,
-    ImageField,
     ImageOutput,
-    LatentsField,
     LatentsOutput,
-    build_latents_output,
 )
 from invokeai.app.invocations.t2i_adapter import T2IAdapterField
-from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
-from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.services.shared.invocation_context import InvocationContext
 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,
@@ -59,16 +71,9 @@ 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
 
@@ -77,18 +82,10 @@ if choose_torch_device() == torch.device("mps"):
 
 DEFAULT_PRECISION = choose_precision(choose_torch_device())
 
-SAMPLER_NAME_VALUES = Literal[tuple(SCHEDULER_MAP.keys())]
-
-# HACK: Many nodes are currently hard-coded to use a fixed latent scale factor of 8. This is fragile, and will need to
-# be addressed if future models use a different latent scale factor. Also, note that there may be places where the scale
-# factor is hard-coded to a literal '8' rather than using this constant.
-# The ratio of image:latent dimensions is LATENT_SCALE_FACTOR:1, or 8:1.
-LATENT_SCALE_FACTOR = 8
-
 
 @invocation_output("scheduler_output")
 class SchedulerOutput(BaseInvocationOutput):
-    scheduler: SAMPLER_NAME_VALUES = OutputField(description=FieldDescriptions.scheduler, ui_type=UIType.Scheduler)
+    scheduler: SCHEDULER_NAME_VALUES = OutputField(description=FieldDescriptions.scheduler, ui_type=UIType.Scheduler)
 
 
 @invocation(
@@ -101,7 +98,7 @@ class SchedulerOutput(BaseInvocationOutput):
 class SchedulerInvocation(BaseInvocation):
     """Selects a scheduler."""
 
-    scheduler: SAMPLER_NAME_VALUES = InputField(
+    scheduler: SCHEDULER_NAME_VALUES = InputField(
         default="euler",
         description=FieldDescriptions.scheduler,
         ui_type=UIType.Scheduler,
@@ -116,7 +113,7 @@ class SchedulerInvocation(BaseInvocation):
     title="Create Denoise Mask",
     tags=["mask", "denoise"],
     category="latents",
-    version="1.0.0",
+    version="1.0.1",
 )
 class CreateDenoiseMaskInvocation(BaseInvocation):
     """Creates mask for denoising model run."""
@@ -131,10 +128,10 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
         ui_order=4,
     )
 
-    def prep_mask_tensor(self, mask_image):
+    def prep_mask_tensor(self, mask_image: 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:
@@ -144,41 +141,90 @@ class CreateDenoiseMaskInvocation(BaseInvocation):
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> DenoiseMaskOutput:
         if self.image is not None:
-            image = context.services.images.get_pil_image(self.image.image_name)
-            image = image_resized_to_grid_as_tensor(image.convert("RGB"))
-            if image.dim() == 3:
-                image = image.unsqueeze(0)
+            image = context.images.get_pil(self.image.image_name)
+            image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
+            if image_tensor.dim() == 3:
+                image_tensor = image_tensor.unsqueeze(0)
         else:
-            image = None
+            image_tensor = None
 
         mask = self.prep_mask_tensor(
-            context.services.images.get_pil_image(self.mask.image_name),
+            context.images.get_pil(self.mask.image_name),
         )
 
-        if image is not None:
-            vae_info = context.services.model_manager.get_model(
-                **self.vae.vae.model_dump(),
-                context=context,
-            )
+        if image_tensor is not None:
+            vae_info = context.models.load(**self.vae.vae.model_dump())
 
-            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 = f"{context.graph_execution_state_id}__{self.id}_masked_latents"
-            context.services.latents.save(masked_latents_name, masked_latents)
+            masked_latents_name = context.tensors.save(tensor=masked_latents)
         else:
             masked_latents_name = None
 
-        mask_name = f"{context.graph_execution_state_id}__{self.id}_mask"
-        context.services.latents.save(mask_name, mask)
+        mask_name = context.tensors.save(tensor=mask)
 
-        return DenoiseMaskOutput(
-            denoise_mask=DenoiseMaskField(
-                mask_name=mask_name,
-                masked_latents_name=masked_latents_name,
-            ),
+        return DenoiseMaskOutput.build(
+            mask_name=mask_name,
+            masked_latents_name=masked_latents_name,
+            gradient=False,
+        )
+
+
+@invocation(
+    "create_gradient_mask",
+    title="Create Gradient Mask",
+    tags=["mask", "denoise"],
+    category="latents",
+    version="1.0.0",
+)
+class CreateGradientMaskInvocation(BaseInvocation):
+    """Creates mask for denoising model run."""
+
+    mask: ImageField = InputField(default=None, description="Image which will be masked", ui_order=1)
+    edge_radius: int = InputField(
+        default=16, ge=0, description="How far to blur/expand the edges of the mask", ui_order=2
+    )
+    coherence_mode: Literal["Gaussian Blur", "Box Blur", "Staged"] = InputField(default="Gaussian Blur", ui_order=3)
+    minimum_denoise: float = InputField(
+        default=0.0, ge=0, le=1, description="Minimum denoise level for the coherence region", ui_order=4
+    )
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> DenoiseMaskOutput:
+        mask_image = context.images.get_pil(self.mask.image_name, mode="L")
+        if self.coherence_mode == "Box Blur":
+            blur_mask = mask_image.filter(ImageFilter.BoxBlur(self.edge_radius))
+        else:  # Gaussian Blur OR Staged
+            # Gaussian Blur uses standard deviation. 1/2 radius is a good approximation
+            blur_mask = mask_image.filter(ImageFilter.GaussianBlur(self.edge_radius / 2))
+
+        mask_tensor: torch.Tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
+        blur_tensor: torch.Tensor = image_resized_to_grid_as_tensor(blur_mask, normalize=False)
+
+        # redistribute blur so that the edges are 0 and blur out to 1
+        blur_tensor = (blur_tensor - 0.5) * 2
+
+        threshold = 1 - self.minimum_denoise
+
+        if self.coherence_mode == "Staged":
+            # wherever the blur_tensor is masked to any degree, convert it to threshold
+            blur_tensor = torch.where((blur_tensor < 1), threshold, blur_tensor)
+        else:
+            # wherever the blur_tensor is above threshold but less than 1, drop it to threshold
+            blur_tensor = torch.where((blur_tensor > threshold) & (blur_tensor < 1), threshold, blur_tensor)
+
+        # multiply original mask to force actually masked regions to 0
+        blur_tensor = mask_tensor * blur_tensor
+
+        mask_name = context.tensors.save(tensor=blur_tensor.unsqueeze(1))
+
+        return DenoiseMaskOutput.build(
+            mask_name=mask_name,
+            masked_latents_name=None,
+            gradient=True,
         )
 
 
@@ -189,10 +235,7 @@ def get_scheduler(
     seed: int,
 ) -> Scheduler:
     scheduler_class, scheduler_extra_config = SCHEDULER_MAP.get(scheduler_name, SCHEDULER_MAP["ddim"])
-    orig_scheduler_info = context.services.model_manager.get_model(
-        **scheduler_info.model_dump(),
-        context=context,
-    )
+    orig_scheduler_info = context.models.load(**scheduler_info.model_dump())
     with orig_scheduler_info as orig_scheduler:
         scheduler_config = orig_scheduler.config
 
@@ -200,7 +243,7 @@ def get_scheduler(
         scheduler_config = scheduler_config["_backup"]
     scheduler_config = {
         **scheduler_config,
-        **scheduler_extra_config,
+        **scheduler_extra_config,  # FIXME
         "_backup": scheduler_config,
     }
 
@@ -213,6 +256,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
 
 
@@ -221,7 +265,7 @@ def get_scheduler(
     title="Denoise Latents",
     tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
     category="latents",
-    version="1.5.1",
+    version="1.5.2",
 )
 class DenoiseLatentsInvocation(BaseInvocation):
     """Denoises noisy latents to decodable images"""
@@ -249,7 +293,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         description=FieldDescriptions.denoising_start,
     )
     denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
-    scheduler: SAMPLER_NAME_VALUES = InputField(
+    scheduler: SCHEDULER_NAME_VALUES = InputField(
         default="euler",
         description=FieldDescriptions.scheduler,
         ui_type=UIType.Scheduler,
@@ -296,7 +340,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:
@@ -307,34 +351,18 @@ 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.services.latents.get(self.positive_conditioning.conditioning_name)
+        positive_cond_data = context.conditioning.load(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.services.latents.get(self.negative_conditioning.conditioning_name)
+        negative_cond_data = context.conditioning.load(self.negative_conditioning.conditioning_name)
         uc = negative_cond_data.conditionings[0].to(device=unet.device, dtype=unet.dtype)
 
         conditioning_data = ConditioningData(
@@ -351,7 +379,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
@@ -363,8 +391,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
     def create_pipeline(
         self,
-        unet,
-        scheduler,
+        unet: UNet2DConditionModel,
+        scheduler: Scheduler,
     ) -> StableDiffusionGeneratorPipeline:
         # TODO:
         # configure_model_padding(
@@ -375,10 +403,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(
@@ -395,11 +423,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
@@ -421,18 +449,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
         #        and if weight is None, populate with default 1.0?
         controlnet_data = []
         for control_info in control_list:
-            control_model = exit_stack.enter_context(
-                context.services.model_manager.get_model(
-                    model_name=control_info.control_model.model_name,
-                    model_type=ModelType.ControlNet,
-                    base_model=control_info.control_model.base_model,
-                    context=context,
-                )
-            )
+            control_model = exit_stack.enter_context(context.models.load(key=control_info.control_model.key))
 
             # control_models.append(control_model)
             control_image_field = control_info.image
-            input_image = context.services.images.get_pil_image(control_image_field.image_name)
+            input_image = context.images.get_pil(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?
@@ -490,27 +511,17 @@ 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.services.model_manager.get_model(
-                    model_name=single_ip_adapter.ip_adapter_model.model_name,
-                    model_type=ModelType.IPAdapter,
-                    base_model=single_ip_adapter.ip_adapter_model.base_model,
-                    context=context,
-                )
+                context.models.load(key=single_ip_adapter.ip_adapter_model.key)
             )
 
-            image_encoder_model_info = context.services.model_manager.get_model(
-                model_name=single_ip_adapter.image_encoder_model.model_name,
-                model_type=ModelType.CLIPVision,
-                base_model=single_ip_adapter.image_encoder_model.base_model,
-                context=context,
-            )
+            image_encoder_model_info = context.models.load(key=single_ip_adapter.image_encoder_model.key)
 
             # `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.services.images.get_pil_image(image.image_name) for image in single_ipa_images]
+            single_ipa_images = [context.images.get_pil(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.
@@ -554,26 +565,20 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         t2i_adapter_data = []
         for t2i_adapter_field in t2i_adapter:
-            t2i_adapter_model_info = context.services.model_manager.get_model(
-                model_name=t2i_adapter_field.t2i_adapter_model.model_name,
-                model_type=ModelType.T2IAdapter,
-                base_model=t2i_adapter_field.t2i_adapter_model.base_model,
-                context=context,
-            )
-            image = context.services.images.get_pil_image(t2i_adapter_field.image.image_name)
+            t2i_adapter_model_config = context.models.get_config(key=t2i_adapter_field.t2i_adapter_model.key)
+            t2i_adapter_loaded_model = context.models.load(key=t2i_adapter_field.t2i_adapter_model.key)
+            image = context.images.get_pil(t2i_adapter_field.image.image_name)
 
             # 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_config.base == BaseModelType.StableDiffusion1:
                 max_unet_downscale = 8
-            elif t2i_adapter_field.t2i_adapter_model.base_model == BaseModelType.StableDiffusionXL:
+            elif t2i_adapter_model_config.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_config.base}'.")
 
             t2i_adapter_model: T2IAdapter
-            with t2i_adapter_model_info as t2i_adapter_model:
+            with t2i_adapter_loaded_model as t2i_adapter_model:
                 total_downscale_factor = t2i_adapter_model.total_downscale_factor
 
                 # Resize the T2I-Adapter input image.
@@ -593,7 +598,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,
@@ -618,7 +623,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)
@@ -630,11 +643,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
@@ -647,18 +660,20 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         return num_inference_steps, timesteps, init_timestep
 
-    def prep_inpaint_mask(self, context, latents):
+    def prep_inpaint_mask(
+        self, context: InvocationContext, latents: torch.Tensor
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], bool]:
         if self.denoise_mask is None:
-            return None, None
+            return None, None, False
 
-        mask = context.services.latents.get(self.denoise_mask.mask_name)
+        mask = context.tensors.load(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.services.latents.get(self.denoise_mask.masked_latents_name)
+            masked_latents = context.tensors.load(self.denoise_mask.masked_latents_name)
         else:
             masked_latents = None
 
-        return 1 - mask, masked_latents
+        return 1 - mask, masked_latents, self.denoise_mask.gradient
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
@@ -666,11 +681,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
             seed = None
             noise = None
             if self.noise is not None:
-                noise = context.services.latents.get(self.noise.latents_name)
+                noise = context.tensors.load(self.noise.latents_name)
                 seed = self.noise.seed
 
             if self.latents is not None:
-                latents = context.services.latents.get(self.latents.latents_name)
+                latents = context.tensors.load(self.latents.latents_name)
                 if seed is None:
                     seed = self.latents.seed
 
@@ -685,7 +700,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
             if seed is None:
                 seed = 0
 
-            mask, masked_latents = self.prep_inpaint_mask(context, latents)
+            mask, masked_latents, gradient_mask = self.prep_inpaint_mask(context, latents)
 
             # TODO(ryand): I have hard-coded `do_classifier_free_guidance=True` to mirror the behaviour of ControlNets,
             # below. Investigate whether this is appropriate.
@@ -696,35 +711,30 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 do_classifier_free_guidance=True,
             )
 
-            # 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]
+            # get the unet's config so that we can pass the base to dispatch_progress()
+            unet_config = context.models.get_config(self.unet.unet.key)
 
-            def step_callback(state: PipelineIntermediateState):
-                self.dispatch_progress(context, source_node_id, state, self.unet.unet.base_model)
+            def step_callback(state: PipelineIntermediateState) -> None:
+                context.util.sd_step_callback(state, unet_config.base)
 
-            def _lora_loader():
+            def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
                 for lora in self.unet.loras:
-                    lora_info = context.services.model_manager.get_model(
-                        **lora.model_dump(exclude={"weight"}),
-                        context=context,
-                    )
-                    yield (lora_info.context.model, lora.weight)
+                    lora_info = context.models.load(**lora.model_dump(exclude={"weight"}))
+                    yield (lora_info.model, lora.weight)
                     del lora_info
                 return
 
-            unet_info = context.services.model_manager.get_model(
-                **self.unet.unet.model_dump(),
-                context=context,
-            )
+            unet_info = context.models.load(**self.unet.unet.model_dump())
+            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)
@@ -778,6 +788,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
                     seed=seed,
                     mask=mask,
                     masked_latents=masked_latents,
+                    gradient_mask=gradient_mask,
                     num_inference_steps=num_inference_steps,
                     conditioning_data=conditioning_data,
                     control_data=controlnet_data,
@@ -792,9 +803,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
             if choose_torch_device() == torch.device("mps"):
                 mps.empty_cache()
 
-            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)
+            name = context.tensors.save(tensor=result_latents)
+        return LatentsOutput.build(latents_name=name, latents=result_latents, seed=seed)
 
 
 @invocation(
@@ -802,9 +812,9 @@ class DenoiseLatentsInvocation(BaseInvocation):
     title="Latents to Image",
     tags=["latents", "image", "vae", "l2i"],
     category="latents",
-    version="1.2.0",
+    version="1.2.1",
 )
-class LatentsToImageInvocation(BaseInvocation, WithMetadata):
+class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents."""
 
     latents: LatentsField = InputField(
@@ -820,14 +830,12 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        latents = context.services.latents.get(self.latents.latents_name)
+        latents = context.tensors.load(self.latents.latents_name)
 
-        vae_info = context.services.model_manager.get_model(
-            **self.vae.vae.model_dump(),
-            context=context,
-        )
+        vae_info = context.models.load(**self.vae.vae.model_dump())
 
-        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)
@@ -854,7 +862,7 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata):
                 vae.to(dtype=torch.float16)
                 latents = latents.half()
 
-            if self.tiled or context.services.configuration.tiled_decode:
+            if self.tiled or context.config.get().tiled_decode:
                 vae.enable_tiling()
             else:
                 vae.disable_tiling()
@@ -878,22 +886,9 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata):
         if choose_torch_device() == torch.device("mps"):
             mps.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,
-        )
+        image_dto = context.images.save(image=image)
 
-        return ImageOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        return ImageOutput.build(image_dto)
 
 
 LATENTS_INTERPOLATION_MODE = Literal["nearest", "linear", "bilinear", "bicubic", "trilinear", "area", "nearest-exact"]
@@ -904,7 +899,7 @@ LATENTS_INTERPOLATION_MODE = Literal["nearest", "linear", "bilinear", "bicubic",
     title="Resize Latents",
     tags=["latents", "resize"],
     category="latents",
-    version="1.0.0",
+    version="1.0.1",
 )
 class ResizeLatentsInvocation(BaseInvocation):
     """Resizes latents to explicit width/height (in pixels). Provided dimensions are floor-divided by 8."""
@@ -927,7 +922,7 @@ class ResizeLatentsInvocation(BaseInvocation):
     antialias: bool = InputField(default=False, description=FieldDescriptions.torch_antialias)
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.services.latents.get(self.latents.latents_name)
+        latents = context.tensors.load(self.latents.latents_name)
 
         # TODO:
         device = choose_torch_device()
@@ -945,10 +940,8 @@ class ResizeLatentsInvocation(BaseInvocation):
         if device == torch.device("mps"):
             mps.empty_cache()
 
-        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)
+        name = context.tensors.save(tensor=resized_latents)
+        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
 
 
 @invocation(
@@ -956,7 +949,7 @@ class ResizeLatentsInvocation(BaseInvocation):
     title="Scale Latents",
     tags=["latents", "resize"],
     category="latents",
-    version="1.0.0",
+    version="1.0.1",
 )
 class ScaleLatentsInvocation(BaseInvocation):
     """Scales latents by a given factor."""
@@ -970,7 +963,7 @@ class ScaleLatentsInvocation(BaseInvocation):
     antialias: bool = InputField(default=False, description=FieldDescriptions.torch_antialias)
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.services.latents.get(self.latents.latents_name)
+        latents = context.tensors.load(self.latents.latents_name)
 
         # TODO:
         device = choose_torch_device()
@@ -989,10 +982,8 @@ class ScaleLatentsInvocation(BaseInvocation):
         if device == torch.device("mps"):
             mps.empty_cache()
 
-        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)
+        name = context.tensors.save(tensor=resized_latents)
+        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
 
 
 @invocation(
@@ -1000,7 +991,7 @@ class ScaleLatentsInvocation(BaseInvocation):
     title="Image to Latents",
     tags=["latents", "image", "vae", "i2l"],
     category="latents",
-    version="1.0.0",
+    version="1.0.1",
 )
 class ImageToLatentsInvocation(BaseInvocation):
     """Encodes an image into latents."""
@@ -1016,8 +1007,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)
@@ -1061,12 +1053,9 @@ class ImageToLatentsInvocation(BaseInvocation):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        image = context.services.images.get_pil_image(self.image.image_name)
+        image = context.images.get_pil(self.image.image_name)
 
-        vae_info = context.services.model_manager.get_model(
-            **self.vae.vae.model_dump(),
-            context=context,
-        )
+        vae_info = context.models.load(**self.vae.vae.model_dump())
 
         image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
         if image_tensor.dim() == 3:
@@ -1074,22 +1063,26 @@ 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")
-        context.services.latents.save(name, latents)
-        return build_latents_output(latents_name=name, latents=latents, seed=None)
+        name = context.tensors.save(tensor=latents)
+        return LatentsOutput.build(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(
@@ -1097,7 +1090,7 @@ class ImageToLatentsInvocation(BaseInvocation):
     title="Blend Latents",
     tags=["latents", "blend"],
     category="latents",
-    version="1.0.0",
+    version="1.0.1",
 )
 class BlendLatentsInvocation(BaseInvocation):
     """Blend two latents using a given alpha. Latents must have same size."""
@@ -1113,8 +1106,8 @@ class BlendLatentsInvocation(BaseInvocation):
     alpha: float = InputField(default=0.5, description=FieldDescriptions.blend_alpha)
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents_a = context.services.latents.get(self.latents_a.latents_name)
-        latents_b = context.services.latents.get(self.latents_b.latents_name)
+        latents_a = context.tensors.load(self.latents_a.latents_name)
+        latents_b = context.tensors.load(self.latents_b.latents_name)
 
         if latents_a.shape != latents_b.shape:
             raise Exception("Latents to blend must be the same size.")
@@ -1122,7 +1115,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:
@@ -1155,12 +1153,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")
@@ -1168,10 +1170,8 @@ class BlendLatentsInvocation(BaseInvocation):
         if device == torch.device("mps"):
             mps.empty_cache()
 
-        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)
+        name = context.tensors.save(tensor=blended_latents)
+        return LatentsOutput.build(latents_name=name, latents=blended_latents)
 
 
 # The Crop Latents node was copied from @skunkworxdark's implementation here:
@@ -1181,7 +1181,7 @@ class BlendLatentsInvocation(BaseInvocation):
     title="Crop Latents",
     tags=["latents", "crop"],
     category="latents",
-    version="1.0.0",
+    version="1.0.1",
 )
 # TODO(ryand): Named `CropLatentsCoreInvocation` to prevent a conflict with custom node `CropLatentsInvocation`.
 # Currently, if the class names conflict then 'GET /openapi.json' fails.
@@ -1216,7 +1216,7 @@ class CropLatentsCoreInvocation(BaseInvocation):
     )
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.services.latents.get(self.latents.latents_name)
+        latents = context.tensors.load(self.latents.latents_name)
 
         x1 = self.x // LATENT_SCALE_FACTOR
         y1 = self.y // LATENT_SCALE_FACTOR
@@ -1225,10 +1225,9 @@ class CropLatentsCoreInvocation(BaseInvocation):
 
         cropped_latents = latents[..., y1:y2, x1:x2]
 
-        name = f"{context.graph_execution_state_id}__{self.id}"
-        context.services.latents.save(name, cropped_latents)
+        name = context.tensors.save(tensor=cropped_latents)
 
-        return build_latents_output(latents_name=name, latents=cropped_latents)
+        return LatentsOutput.build(latents_name=name, latents=cropped_latents)
 
 
 @invocation_output("ideal_size_output")
@@ -1256,15 +1255,16 @@ 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.models.get_config(**self.unet.unet.model_dump())
         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,10 +5,11 @@ 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.shared.fields import FieldDescriptions
+from invokeai.app.services.shared.invocation_context import InvocationContext
 
-from .baseinvocation import BaseInvocation, InputField, InvocationContext, invocation
+from .baseinvocation import BaseInvocation, invocation
 
 
 @invocation("add", title="Add Integers", tags=["math", "add"], category="math", version="1.0.0")

invokeai/app/invocations/metadata.py

@@ -5,20 +5,22 @@ 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.shared.fields import FieldDescriptions
+from invokeai.app.services.shared.invocation_context import InvocationContext
 
 from ...version import __version__
 
@@ -31,7 +33,7 @@ class MetadataItemField(BaseModel):
 class LoRAMetadataField(BaseModel):
     """LoRA Metadata Field"""
 
-    lora: LoRAModelField = Field(description=FieldDescriptions.lora_model)
+    model: LoRAModelField = Field(description=FieldDescriptions.lora_model)
     weight: float = Field(description=FieldDescriptions.lora_weight)
 
 
@@ -112,7 +114,7 @@ GENERATION_MODES = Literal[
 ]
 
 
-@invocation("core_metadata", title="Core Metadata", tags=["metadata"], category="metadata", version="1.0.1")
+@invocation("core_metadata", title="Core Metadata", tags=["metadata"], category="metadata", version="1.1.1")
 class CoreMetadataInvocation(BaseInvocation):
     """Collects core generation metadata into a MetadataField"""
 

invokeai/app/invocations/model.py

@@ -1,31 +1,24 @@
 import copy
 from typing import List, Optional
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, Field
 
-from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField
+from invokeai.app.services.shared.invocation_context import InvocationContext
 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):
@@ -55,7 +48,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")
 
@@ -84,20 +77,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(
@@ -105,7 +91,7 @@ class LoRAModelField(BaseModel):
     title="Main Model",
     tags=["model"],
     category="model",
-    version="1.0.0",
+    version="1.0.1",
 )
 class MainModelLoaderInvocation(BaseInvocation):
     """Loads a main model, outputting its submodels."""
@@ -114,85 +100,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.services.model_manager.model_exists(
-            model_name=model_name,
-            base_model=base_model,
-            model_type=model_type,
-        ):
-            raise Exception(f"Unknown {base_model} {model_type} model: {model_name}")
-
-        """
-        if not context.services.model_manager.model_exists(
-            model_name=self.model_name,
-            model_type=SDModelType.Diffusers,
-            submodel=SDModelType.Tokenizer,
-        ):
-            raise Exception(
-                f"Failed to find tokenizer submodel in {self.model_name}! Check if model corrupted"
-            )
-
-        if not context.services.model_manager.model_exists(
-            model_name=self.model_name,
-            model_type=SDModelType.Diffusers,
-            submodel=SDModelType.TextEncoder,
-        ):
-            raise Exception(
-                f"Failed to find text_encoder submodel in {self.model_name}! Check if model corrupted"
-            )
-
-        if not context.services.model_manager.model_exists(
-            model_name=self.model_name,
-            model_type=SDModelType.Diffusers,
-            submodel=SDModelType.UNet,
-        ):
-            raise Exception(
-                f"Failed to find unet submodel from {self.model_name}! Check if model corrupted"
-            )
-        """
+        if not context.models.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,
                 ),
             ),
         )
@@ -206,7 +147,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.0")
+@invocation("lora_loader", title="LoRA", tags=["model"], category="model", version="1.0.1")
 class LoraLoaderInvocation(BaseInvocation):
     """Apply selected lora to unet and text_encoder."""
 
@@ -229,21 +170,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.services.model_manager.model_exists(
-            base_model=base_model,
-            model_name=lora_name,
-            model_type=ModelType.Lora,
-        ):
-            raise Exception(f"Unkown lora name: {lora_name}!")
+        if not context.models.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()
 
@@ -251,10 +187,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,
                 )
             )
@@ -263,10 +197,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,
                 )
             )
@@ -288,7 +220,7 @@ class SDXLLoraLoaderOutput(BaseInvocationOutput):
     title="SDXL LoRA",
     tags=["lora", "model"],
     category="model",
-    version="1.0.0",
+    version="1.0.1",
 )
 class SDXLLoraLoaderInvocation(BaseInvocation):
     """Apply selected lora to unet and text_encoder."""
@@ -318,24 +250,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.services.model_manager.model_exists(
-            base_model=base_model,
-            model_name=lora_name,
-            model_type=ModelType.Lora,
-        ):
-            raise Exception(f"Unknown lora name: {lora_name}!")
+        if not context.models.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()
 
@@ -343,10 +270,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,
                 )
             )
@@ -355,10 +280,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,
                 )
             )
@@ -367,10 +290,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,
                 )
             )
@@ -381,13 +302,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.0")
+@invocation("vae_loader", title="VAE", tags=["vae", "model"], category="model", version="1.0.1")
 class VaeLoaderInvocation(BaseInvocation):
     """Loads a VAE model, outputting a VaeLoaderOutput"""
 
@@ -398,25 +316,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.services.model_manager.model_exists(
-            base_model=base_model,
-            model_name=model_name,
-            model_type=model_type,
-        ):
-            raise Exception(f"Unkown vae name: {model_name}!")
-        return VAEOutput(
-            vae=VaeField(
-                vae=ModelInfo(
-                    model_name=model_name,
-                    base_model=base_model,
-                    model_type=model_type,
-                )
-            )
-        )
+        if not context.models.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,17 +4,15 @@
 import torch
 from pydantic import field_validator
 
-from invokeai.app.invocations.latent import LatentsField
-from invokeai.app.shared.fields import FieldDescriptions
+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.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,
 )
@@ -69,13 +67,13 @@ class NoiseOutput(BaseInvocationOutput):
     width: int = OutputField(description=FieldDescriptions.width)
     height: int = OutputField(description=FieldDescriptions.height)
 
-
-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,
-    )
+    @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,
+        )
 
 
 @invocation(
@@ -96,13 +94,13 @@ class NoiseInvocation(BaseInvocation):
     )
     width: int = InputField(
         default=512,
-        multiple_of=8,
+        multiple_of=LATENT_SCALE_FACTOR,
         gt=0,
         description=FieldDescriptions.width,
     )
     height: int = InputField(
         default=512,
-        multiple_of=8,
+        multiple_of=LATENT_SCALE_FACTOR,
         gt=0,
         description=FieldDescriptions.height,
     )
@@ -124,6 +122,5 @@ class NoiseInvocation(BaseInvocation):
             seed=self.seed,
             use_cpu=self.use_cpu,
         )
-        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)
+        name = context.tensors.save(tensor=noise)
+        return NoiseOutput.build(latents_name=name, latents=noise, seed=self.seed)

invokeai/app/invocations/onnx.py

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

invokeai/app/invocations/param_easing.py

@@ -40,8 +40,10 @@ 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, InputField, InvocationContext, invocation
+from .baseinvocation import BaseInvocation, invocation
+from .fields import InputField
 
 
 @invocation(
@@ -109,7 +111,7 @@ EASING_FUNCTION_KEYS = Literal[tuple(EASING_FUNCTIONS_MAP.keys())]
     title="Step Param Easing",
     tags=["step", "easing"],
     category="step",
-    version="1.0.0",
+    version="1.0.1",
 )
 class StepParamEasingInvocation(BaseInvocation):
     """Experimental per-step parameter easing for denoising steps"""
@@ -148,19 +150,19 @@ class StepParamEasingInvocation(BaseInvocation):
         postlist = list(num_poststeps * [self.post_end_value])
 
         if log_diagnostics:
-            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))
+            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))
 
         easing_class = EASING_FUNCTIONS_MAP[self.easing]
         if log_diagnostics:
-            context.services.logger.debug("easing class: " + str(easing_class))
+            context.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
@@ -171,7 +173,7 @@ class StepParamEasingInvocation(BaseInvocation):
 
             base_easing_duration = int(np.ceil(num_easing_steps / 2.0))
             if log_diagnostics:
-                context.services.logger.debug("base easing duration: " + str(base_easing_duration))
+                context.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,
@@ -183,14 +185,14 @@ class StepParamEasingInvocation(BaseInvocation):
                 easing_val = easing_function.ease(step_index)
                 base_easing_vals.append(easing_val)
                 if log_diagnostics:
-                    context.services.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(easing_val))
+                    context.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.services.logger.debug("base easing vals: " + str(base_easing_vals))
-                context.services.logger.debug("mirror easing vals: " + str(mirror_easing_vals))
+                context.logger.debug("base easing vals: " + str(base_easing_vals))
+                context.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
@@ -225,12 +227,12 @@ class StepParamEasingInvocation(BaseInvocation):
                 step_val = easing_function.ease(step_index)
                 easing_list.append(step_val)
                 if log_diagnostics:
-                    context.services.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(step_val))
+                    context.logger.debug("step_index: " + str(step_index) + ", easing_val: " + str(step_val))
 
         if log_diagnostics:
-            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)))
+            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)))
 
         param_list = prelist + easing_list + postlist
 

invokeai/app/invocations/primitives.py

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

invokeai/app/invocations/sdxl.py

@@ -1,14 +1,10 @@
-from invokeai.app.shared.fields import FieldDescriptions
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager import SubModelType
 
-from ...backend.model_management import ModelType, SubModelType
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
-    Input,
-    InputField,
-    InvocationContext,
-    OutputField,
-    UIType,
     invocation,
     invocation_output,
 )
@@ -34,7 +30,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.0")
+@invocation("sdxl_model_loader", title="SDXL Main Model", tags=["model", "sdxl"], category="model", version="1.0.1")
 class SDXLModelLoaderInvocation(BaseInvocation):
     """Loads an sdxl base model, outputting its submodels."""
 
@@ -44,72 +40,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.services.model_manager.model_exists(
-            model_name=model_name,
-            base_model=base_model,
-            model_type=model_type,
-        ):
-            raise Exception(f"Unknown {base_model} {model_type} model: {model_name}")
+        if not context.models.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,
                 ),
             ),
         )
@@ -120,7 +96,7 @@ class SDXLModelLoaderInvocation(BaseInvocation):
     title="SDXL Refiner Model",
     tags=["model", "sdxl", "refiner"],
     category="model",
-    version="1.0.0",
+    version="1.0.1",
 )
 class SDXLRefinerModelLoaderInvocation(BaseInvocation):
     """Loads an sdxl refiner model, outputting its submodels."""
@@ -133,56 +109,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.services.model_manager.model_exists(
-            model_name=model_name,
-            base_model=base_model,
-            model_type=model_type,
-        ):
-            raise Exception(f"Unknown {base_model} {model_type} model: {model_name}")
+        if not context.models.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,16 +2,15 @@
 
 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,29 +1,21 @@
 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.primitives import ImageField
+from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField, OutputField
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
-from invokeai.app.shared.fields import FieldDescriptions
-from invokeai.backend.model_management.models.base import BaseModelType
+from invokeai.app.services.shared.invocation_context import InvocationContext
 
 
 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,16 +8,12 @@ from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
     Classification,
-    Input,
-    InputField,
-    InvocationContext,
-    OutputField,
-    WithMetadata,
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.primitives import ImageField, ImageOutput
-from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
+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.backend.tiles.tiles import (
     calc_tiles_even_split,
     calc_tiles_min_overlap,
@@ -236,7 +232,7 @@ BLEND_MODES = Literal["Linear", "Seam"]
     version="1.1.0",
     classification=Classification.Beta,
 )
-class MergeTilesToImageInvocation(BaseInvocation, WithMetadata):
+class MergeTilesToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Merge multiple tile images into a single image."""
 
     # Inputs
@@ -268,7 +264,7 @@ class MergeTilesToImageInvocation(BaseInvocation, WithMetadata):
         # existed in memory at an earlier point in the graph.
         tile_np_images: list[np.ndarray] = []
         for image in images:
-            pil_image = context.services.images.get_pil_image(image.image_name)
+            pil_image = context.images.get_pil(image.image_name)
             pil_image = pil_image.convert("RGB")
             tile_np_images.append(np.array(pil_image))
 
@@ -291,18 +287,5 @@ class MergeTilesToImageInvocation(BaseInvocation, WithMetadata):
         # Convert into a PIL image and save
         pil_image = Image.fromarray(np_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(
-            image=ImageField(image_name=image_dto.image_name),
-            width=image_dto.width,
-            height=image_dto.height,
-        )
+        image_dto = context.images.save(image=pil_image)
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/upscale.py

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

invokeai/app/services/bulk_download/bulk_download_base.py

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

invokeai/app/services/bulk_download/bulk_download_common.py

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

invokeai/app/services/bulk_download/bulk_download_default.py

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

invokeai/app/services/config/config_base.py

@@ -27,11 +27,11 @@ class InvokeAISettings(BaseSettings):
     """Runtime configuration settings in which default values are read from an omegaconf .yaml file."""
 
     initconf: ClassVar[Optional[DictConfig]] = None
-    argparse_groups: ClassVar[Dict] = {}
+    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"
 
@@ -155,10 +156,11 @@ class InvokeAISettings(BaseSettings):
             "lora_dir",
             "embedding_dir",
             "controlnet_dir",
+            "conf_path",
         ]
 
     @classmethod
-    def add_field_argument(cls, command_parser, name: str, field, default_override=None):
+    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

@@ -30,7 +30,6 @@ InvokeAI:
     lora_dir: null
     embedding_dir: null
     controlnet_dir: null
-    conf_path: configs/models.yaml
     models_dir: models
     legacy_conf_dir: configs/stable-diffusion
     db_dir: databases
@@ -123,7 +122,6 @@ a Path object:
 
  root_path          - path to InvokeAI root
  output_path        - path to default outputs directory
- model_conf_path    - path to models.yaml
  conf               - alias for the above
  embedding_path     - path to the embeddings directory
  lora_path          - path to the LoRA directory
@@ -163,7 +161,6 @@ two configs are kept in separate sections of the config file:
   InvokeAI:
      Paths:
         root: /home/lstein/invokeai-main
-        conf_path: configs/models.yaml
         legacy_conf_dir: configs/stable-diffusion
         outdir: outputs
      ...
@@ -173,7 +170,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 +182,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):
@@ -235,8 +234,8 @@ class InvokeAIAppConfig(InvokeAISettings):
     # PATHS
     root                : Optional[Path] = Field(default=None, description='InvokeAI runtime root directory', json_schema_extra=Categories.Paths)
     autoimport_dir      : Path = Field(default=Path('autoimport'), description='Path to a directory of models files to be imported on startup.', json_schema_extra=Categories.Paths)
-    conf_path           : Path = Field(default=Path('configs/models.yaml'), description='Path to models definition file', json_schema_extra=Categories.Paths)
     models_dir          : Path = Field(default=Path('models'), description='Path to the models directory', json_schema_extra=Categories.Paths)
+    convert_cache_dir   : Path = Field(default=Path('models/.cache'), description='Path to the converted models cache directory', json_schema_extra=Categories.Paths)
     legacy_conf_dir     : Path = Field(default=Path('configs/stable-diffusion'), description='Path to directory of legacy checkpoint config files', json_schema_extra=Categories.Paths)
     db_dir              : Path = Field(default=Path('databases'), description='Path to InvokeAI databases directory', json_schema_extra=Categories.Paths)
     outdir              : Path = Field(default=Path('outputs'), description='Default folder for output images', json_schema_extra=Categories.Paths)
@@ -260,8 +259,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)
 
@@ -296,6 +297,7 @@ class InvokeAIAppConfig(InvokeAISettings):
     lora_dir            : Optional[Path] = Field(default=None, description='Path to a directory of LoRA/LyCORIS models to be imported on startup.', json_schema_extra=Categories.Paths)
     embedding_dir       : Optional[Path] = Field(default=None, description='Path to a directory of Textual Inversion embeddings to be imported on startup.', json_schema_extra=Categories.Paths)
     controlnet_dir      : Optional[Path] = Field(default=None, description='Path to a directory of ControlNet embeddings to be imported on startup.', json_schema_extra=Categories.Paths)
+    conf_path           : Path = Field(default=Path('configs/models.yaml'), description='Path to models definition file', json_schema_extra=Categories.Paths)
 
     # this is not referred to in the source code and can be removed entirely
     #free_gpu_mem        : Optional[bool] = Field(default=None, description="If true, purge model from GPU after each generation.", json_schema_extra=Categories.MemoryPerformance)
@@ -404,6 +406,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 +440,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

@@ -3,7 +3,7 @@
 
 from typing import Any, Dict, List, Optional, Union
 
-from invokeai.app.services.invocation_processor.invocation_processor_common import ProgressImage
+from invokeai.app.services.session_processor.session_processor_common import ProgressImage
 from invokeai.app.services.session_queue.session_queue_common import (
     BatchStatus,
     EnqueueBatchResult,
@@ -11,12 +11,12 @@ from invokeai.app.services.session_queue.session_queue_common import (
     SessionQueueStatus,
 )
 from invokeai.app.util.misc import get_timestamp
-from invokeai.backend.model_management.model_manager import ModelInfo
-from invokeai.backend.model_management.models.base import BaseModelType, ModelType, SubModelType
+from invokeai.backend.model_manager import AnyModelConfig
 
 
 class EventServiceBase:
     queue_event: str = "queue_event"
+    bulk_download_event: str = "bulk_download_event"
     download_event: str = "download_event"
     model_event: str = "model_event"
 
@@ -25,6 +25,14 @@ class EventServiceBase:
     def dispatch(self, event_name: str, payload: Any) -> None:
         pass
 
+    def _emit_bulk_download_event(self, event_name: str, payload: dict) -> None:
+        """Bulk download events are emitted to a room with queue_id as the room name"""
+        payload["timestamp"] = get_timestamp()
+        self.dispatch(
+            event_name=EventServiceBase.bulk_download_event,
+            payload={"event": event_name, "data": payload},
+        )
+
     def __emit_queue_event(self, event_name: str, payload: dict) -> None:
         """Queue events are emitted to a room with queue_id as the room name"""
         payload["timestamp"] = get_timestamp()
@@ -55,7 +63,7 @@ class EventServiceBase:
         queue_item_id: int,
         queue_batch_id: str,
         graph_execution_state_id: str,
-        node: dict,
+        node_id: str,
         source_node_id: str,
         progress_image: Optional[ProgressImage],
         step: int,
@@ -70,7 +78,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.get("id"),
+                "node_id": node_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 +179,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 +189,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 +199,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_info: ModelInfo,
+        model_config: AnyModelConfig,
     ) -> None:
         """Emitted when a model is correctly loaded (returns model info)"""
         self.__emit_queue_event(
@@ -211,59 +209,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": model_info.hash,
-                "location": str(model_info.location),
-                "precision": str(model_info.precision),
-            },
-        )
-
-    def emit_session_retrieval_error(
-        self,
-        queue_id: str,
-        queue_item_id: int,
-        queue_batch_id: str,
-        graph_execution_state_id: str,
-        error_type: str,
-        error: str,
-    ) -> None:
-        """Emitted when session retrieval fails"""
-        self.__emit_queue_event(
-            event_name="session_retrieval_error",
-            payload={
-                "queue_id": queue_id,
-                "queue_item_id": queue_item_id,
-                "queue_batch_id": queue_batch_id,
-                "graph_execution_state_id": graph_execution_state_id,
-                "error_type": error_type,
-                "error": error,
-            },
-        )
-
-    def emit_invocation_retrieval_error(
-        self,
-        queue_id: str,
-        queue_item_id: int,
-        queue_batch_id: str,
-        graph_execution_state_id: str,
-        node_id: str,
-        error_type: str,
-        error: str,
-    ) -> None:
-        """Emitted when invocation retrieval fails"""
-        self.__emit_queue_event(
-            event_name="invocation_retrieval_error",
-            payload={
-                "queue_id": queue_id,
-                "queue_item_id": queue_item_id,
-                "queue_batch_id": queue_batch_id,
-                "graph_execution_state_id": graph_execution_state_id,
-                "node_id": node_id,
-                "error_type": error_type,
-                "error": error,
+                "model_config": model_config.model_dump(),
             },
         )
 
@@ -411,6 +357,7 @@ class EventServiceBase:
         bytes: int,
         total_bytes: int,
         parts: List[Dict[str, Union[str, int]]],
+        id: int,
     ) -> None:
         """
         Emit at intervals while the install job is in progress (remote models only).
@@ -430,6 +377,7 @@ class EventServiceBase:
                 "bytes": bytes,
                 "total_bytes": total_bytes,
                 "parts": parts,
+                "id": id,
             },
         )
 
@@ -444,7 +392,7 @@ class EventServiceBase:
             payload={"source": source},
         )
 
-    def emit_model_install_completed(self, source: str, key: str, total_bytes: Optional[int] = None) -> None:
+    def emit_model_install_completed(self, source: str, key: str, id: int, total_bytes: Optional[int] = None) -> None:
         """
         Emit when an install job is completed successfully.
 
@@ -454,11 +402,7 @@ class EventServiceBase:
         """
         self.__emit_model_event(
             event_name="model_install_completed",
-            payload={
-                "source": source,
-                "total_bytes": total_bytes,
-                "key": key,
-            },
+            payload={"source": source, "total_bytes": total_bytes, "key": key, "id": id},
         )
 
     def emit_model_install_cancelled(self, source: str) -> None:
@@ -472,12 +416,7 @@ class EventServiceBase:
             payload={"source": source},
         )
 
-    def emit_model_install_error(
-        self,
-        source: str,
-        error_type: str,
-        error: str,
-    ) -> None:
+    def emit_model_install_error(self, source: str, error_type: str, error: str, id: int) -> None:
         """
         Emit when an install job encounters an exception.
 
@@ -487,9 +426,45 @@ class EventServiceBase:
         """
         self.__emit_model_event(
             event_name="model_install_error",
+            payload={"source": source, "error_type": error_type, "error": error, "id": id},
+        )
+
+    def emit_bulk_download_started(
+        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
+    ) -> None:
+        """Emitted when a bulk download starts"""
+        self._emit_bulk_download_event(
+            event_name="bulk_download_started",
             payload={
-                "source": source,
-                "error_type": error_type,
+                "bulk_download_id": bulk_download_id,
+                "bulk_download_item_id": bulk_download_item_id,
+                "bulk_download_item_name": bulk_download_item_name,
+            },
+        )
+
+    def emit_bulk_download_completed(
+        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
+    ) -> None:
+        """Emitted when a bulk download completes"""
+        self._emit_bulk_download_event(
+            event_name="bulk_download_completed",
+            payload={
+                "bulk_download_id": bulk_download_id,
+                "bulk_download_item_id": bulk_download_item_id,
+                "bulk_download_item_name": bulk_download_item_name,
+            },
+        )
+
+    def emit_bulk_download_failed(
+        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, error: str
+    ) -> None:
+        """Emitted when a bulk download fails"""
+        self._emit_bulk_download_event(
+            event_name="bulk_download_failed",
+            payload={
+                "bulk_download_id": bulk_download_id,
+                "bulk_download_item_id": bulk_download_item_id,
+                "bulk_download_item_name": bulk_download_item_name,
                 "error": error,
             },
         )

invokeai/app/services/image_files/image_files_base.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 from PIL.Image import Image as PILImageType
 
-from invokeai.app.invocations.baseinvocation import MetadataField
+from invokeai.app.invocations.fields 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.baseinvocation import MetadataField
+from invokeai.app.invocations.fields 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.metadata import MetadataField
+from invokeai.app.invocations.fields 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.baseinvocation import MetadataField, MetadataFieldValidator
+from invokeai.app.invocations.fields 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.baseinvocation import MetadataField
+from invokeai.app.invocations.fields 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.baseinvocation import MetadataField
+from invokeai.app.invocations.fields 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

invokeai/app/services/invocation_cache/invocation_cache_memory.py

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

invokeai/app/services/invocation_processor/invocation_processor_base.py

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

invokeai/app/services/invocation_processor/invocation_processor_common.py

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

invokeai/app/services/invocation_processor/invocation_processor_default.py

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

invokeai/app/services/invocation_queue/invocation_queue_base.py

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

invokeai/app/services/invocation_queue/invocation_queue_common.py

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

invokeai/app/services/invocation_queue/invocation_queue_memory.py

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

invokeai/app/services/invocation_services.py

@@ -3,13 +3,20 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+from invokeai.app.services.object_serializer.object_serializer_base import ObjectSerializerBase
+
 if TYPE_CHECKING:
     from logging import Logger
 
+    import torch
+
+    from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
+
     from .board_image_records.board_image_records_base import BoardImageRecordStorageBase
     from .board_images.board_images_base import BoardImagesServiceABC
     from .board_records.board_records_base import BoardRecordStorageBase
     from .boards.boards_base import BoardServiceABC
+    from .bulk_download.bulk_download_base import BulkDownloadBase
     from .config import InvokeAIAppConfig
     from .download import DownloadQueueServiceBase
     from .events.events_base import EventServiceBase
@@ -17,18 +24,11 @@ if TYPE_CHECKING:
     from .image_records.image_records_base import ImageRecordStorageBase
     from .images.images_base import ImageServiceABC
     from .invocation_cache.invocation_cache_base import InvocationCacheBase
-    from .invocation_processor.invocation_processor_base import InvocationProcessorABC
-    from .invocation_queue.invocation_queue_base import InvocationQueueABC
     from .invocation_stats.invocation_stats_base import InvocationStatsServiceBase
-    from .item_storage.item_storage_base import ItemStorageABC
-    from .latents_storage.latents_storage_base import LatentsStorageBase
-    from .model_install import ModelInstallServiceBase
     from .model_manager.model_manager_base import ModelManagerServiceBase
-    from .model_records import ModelRecordServiceBase
     from .names.names_base import NameServiceBase
     from .session_processor.session_processor_base import SessionProcessorBase
     from .session_queue.session_queue_base import SessionQueueBase
-    from .shared.graph import GraphExecutionState
     from .urls.urls_base import UrlServiceBase
     from .workflow_records.workflow_records_base import WorkflowRecordsStorageBase
 
@@ -36,83 +36,50 @@ if TYPE_CHECKING:
 class InvocationServices:
     """Services that can be used by invocations"""
 
-    # TODO: Just forward-declared everything due to circular dependencies. Fix structure.
-    board_images: "BoardImagesServiceABC"
-    board_image_record_storage: "BoardImageRecordStorageBase"
-    boards: "BoardServiceABC"
-    board_records: "BoardRecordStorageBase"
-    configuration: "InvokeAIAppConfig"
-    events: "EventServiceBase"
-    graph_execution_manager: "ItemStorageABC[GraphExecutionState]"
-    images: "ImageServiceABC"
-    image_records: "ImageRecordStorageBase"
-    image_files: "ImageFileStorageBase"
-    latents: "LatentsStorageBase"
-    logger: "Logger"
-    model_manager: "ModelManagerServiceBase"
-    model_records: "ModelRecordServiceBase"
-    download_queue: "DownloadQueueServiceBase"
-    model_install: "ModelInstallServiceBase"
-    processor: "InvocationProcessorABC"
-    performance_statistics: "InvocationStatsServiceBase"
-    queue: "InvocationQueueABC"
-    session_queue: "SessionQueueBase"
-    session_processor: "SessionProcessorBase"
-    invocation_cache: "InvocationCacheBase"
-    names: "NameServiceBase"
-    urls: "UrlServiceBase"
-    workflow_records: "WorkflowRecordsStorageBase"
-
     def __init__(
         self,
         board_images: "BoardImagesServiceABC",
         board_image_records: "BoardImageRecordStorageBase",
         boards: "BoardServiceABC",
         board_records: "BoardRecordStorageBase",
+        bulk_download: "BulkDownloadBase",
         configuration: "InvokeAIAppConfig",
         events: "EventServiceBase",
-        graph_execution_manager: "ItemStorageABC[GraphExecutionState]",
         images: "ImageServiceABC",
         image_files: "ImageFileStorageBase",
         image_records: "ImageRecordStorageBase",
-        latents: "LatentsStorageBase",
         logger: "Logger",
         model_manager: "ModelManagerServiceBase",
-        model_records: "ModelRecordServiceBase",
         download_queue: "DownloadQueueServiceBase",
-        model_install: "ModelInstallServiceBase",
-        processor: "InvocationProcessorABC",
         performance_statistics: "InvocationStatsServiceBase",
-        queue: "InvocationQueueABC",
         session_queue: "SessionQueueBase",
         session_processor: "SessionProcessorBase",
         invocation_cache: "InvocationCacheBase",
         names: "NameServiceBase",
         urls: "UrlServiceBase",
         workflow_records: "WorkflowRecordsStorageBase",
+        tensors: "ObjectSerializerBase[torch.Tensor]",
+        conditioning: "ObjectSerializerBase[ConditioningFieldData]",
     ):
         self.board_images = board_images
         self.board_image_records = board_image_records
         self.boards = boards
         self.board_records = board_records
+        self.bulk_download = bulk_download
         self.configuration = configuration
         self.events = events
-        self.graph_execution_manager = graph_execution_manager
         self.images = images
         self.image_files = image_files
         self.image_records = image_records
-        self.latents = latents
         self.logger = logger
         self.model_manager = model_manager
-        self.model_records = model_records
         self.download_queue = download_queue
-        self.model_install = model_install
-        self.processor = processor
         self.performance_statistics = performance_statistics
-        self.queue = queue
         self.session_queue = session_queue
         self.session_processor = session_processor
         self.invocation_cache = invocation_cache
         self.names = names
         self.urls = urls
         self.workflow_records = workflow_records
+        self.tensors = tensors
+        self.conditioning = conditioning

invokeai/app/services/invocation_stats/invocation_stats_base.py

@@ -3,7 +3,7 @@
 
 Usage:
 
-statistics = InvocationStatsService(graph_execution_manager)
+statistics = InvocationStatsService()
 with statistics.collect_stats(invocation, graph_execution_state.id):
       ... execute graphs...
 statistics.log_stats()
@@ -29,8 +29,8 @@ writes to the system log is stored in InvocationServices.performance_statistics.
 """
 
 from abc import ABC, abstractmethod
-from contextlib import AbstractContextManager
 from pathlib import Path
+from typing import ContextManager
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
@@ -40,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:
+    ) -> ContextManager[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,16 +60,12 @@ class InvocationStatsServiceBase(ABC):
         pass
 
     @abstractmethod
-    def reset_stats(self, graph_execution_state_id: str):
-        """
-        Reset all statistics for the indicated graph.
-        :param graph_execution_state_id: The id of the session whose stats to reset.
-        :raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
-        """
+    def reset_stats(self):
+        """Reset all stored statistics."""
         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 Generator
 
 import psutil
 import torch
@@ -9,8 +10,7 @@ import torch
 import invokeai.backend.util.logging as logger
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invoker import Invoker
-from invokeai.app.services.item_storage.item_storage_common import ItemNotFoundError
-from invokeai.backend.model_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,22 +41,23 @@ 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) -> Generator[None, None, None]:
+        # This is to handle case of the model manager not being initialized, which happens
+        # during some tests.
+        services = self._invoker.services
         if not self._stats.get(graph_execution_state_id):
             # First time we're seeing this graph_execution_state_id.
             self._stats[graph_execution_state_id] = GraphExecutionStats()
             self._cache_stats[graph_execution_state_id] = CacheStats()
 
-            # Prune stale stats. There should be none since we're starting a new graph, but just in case.
-            self._prune_stale_stats()
-
         # Record state before the invocation.
         start_time = time.time()
         start_ram = psutil.Process().memory_info().rss
         if torch.cuda.is_available():
             torch.cuda.reset_peak_memory_stats()
-        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,42 +74,9 @@ class InvocationStatsService(InvocationStatsServiceBase):
             )
             self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
 
-    def _prune_stale_stats(self):
-        """Check all graphs being tracked and prune any that have completed/errored.
-
-        This shouldn't be necessary, but we don't have totally robust upstream handling of graph completions/errors, so
-        for now we call this function periodically to prevent them from accumulating.
-        """
-        to_prune: list[str] = []
-        for graph_execution_state_id in self._stats:
-            try:
-                graph_execution_state = self._invoker.services.graph_execution_manager.get(graph_execution_state_id)
-            except ItemNotFoundError:
-                # TODO(ryand): What would cause this? Should this exception just be allowed to propagate?
-                logger.warning(f"Failed to get graph state for {graph_execution_state_id}.")
-                continue
-
-            if not graph_execution_state.is_complete():
-                # The graph is still running, don't prune it.
-                continue
-
-            to_prune.append(graph_execution_state_id)
-
-        for graph_execution_state_id in to_prune:
-            del self._stats[graph_execution_state_id]
-            del self._cache_stats[graph_execution_state_id]
-
-        if len(to_prune) > 0:
-            logger.info(f"Pruned stale graph stats for {to_prune}.")
-
-    def reset_stats(self, graph_execution_state_id: str):
-        try:
-            del self._stats[graph_execution_state_id]
-            del self._cache_stats[graph_execution_state_id]
-        except KeyError as e:
-            raise GESStatsNotFoundError(
-                f"Attempted to clear statistics for unknown graph {graph_execution_state_id}: {e}."
-            ) from e
+    def reset_stats(self):
+        self._stats = {}
+        self._cache_stats = {}
 
     def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
         graph_stats_summary = self._get_graph_summary(graph_execution_state_id)

invokeai/app/services/invoker.py

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

invokeai/app/services/item_storage/item_storage_base.py

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

invokeai/app/services/latents_storage/latents_storage_base.py

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

invokeai/app/services/latents_storage/latents_storage_disk.py

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

invokeai/app/services/latents_storage/latents_storage_forward_cache.py

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

invokeai/app/services/model_install/model_install_base.py

@@ -14,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
 
 
@@ -154,6 +156,7 @@ class ModelInstallJob(BaseModel):
 
     id: int = Field(description="Unique ID for this job")
     status: InstallStatus = Field(default=InstallStatus.WAITING, description="Current status of install process")
+    error_reason: Optional[str] = Field(default=None, description="Information about why the job failed")
     config_in: Dict[str, Any] = Field(
         default_factory=dict, description="Configuration information (e.g. 'description') to apply to model."
     )
@@ -175,6 +178,12 @@ class ModelInstallJob(BaseModel):
     download_parts: Set[DownloadJob] = Field(
         default_factory=set, description="Download jobs contributing to this install"
     )
+    error: Optional[str] = Field(
+        default=None, description="On an error condition, this field will contain the text of the exception"
+    )
+    error_traceback: Optional[str] = Field(
+        default=None, description="On an error condition, this field will contain the exception traceback"
+    )
     # internal flags and transitory settings
     _install_tmpdir: Optional[Path] = PrivateAttr(default=None)
     _exception: Optional[Exception] = PrivateAttr(default=None)
@@ -182,7 +191,10 @@ class ModelInstallJob(BaseModel):
     def set_error(self, e: Exception) -> None:
         """Record the error and traceback from an exception."""
         self._exception = e
+        self.error = str(e)
+        self.error_traceback = self._format_error(e)
         self.status = InstallStatus.ERROR
+        self.error_reason = self._exception.__class__.__name__ if self._exception else None
 
     def cancel(self) -> None:
         """Call to cancel the job."""
@@ -193,10 +205,9 @@ class ModelInstallJob(BaseModel):
         """Class name of the exception that led to status==ERROR."""
         return self._exception.__class__.__name__ if self._exception else None
 
-    @property
-    def error(self) -> Optional[str]:
+    def _format_error(self, exception: Exception) -> str:
         """Error traceback."""
-        return "".join(traceback.format_exception(self._exception)) if self._exception else None
+        return "".join(traceback.format_exception(exception))
 
     @property
     def cancelled(self) -> bool:
@@ -243,7 +254,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 +335,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 +433,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 +454,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 +471,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,18 @@ 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
+
+        if preferred_name := config.get("name"):
+            preferred_name = Path(preferred_name).with_suffix(model_path.suffix)
+
+        dest_path = (
+            self.app_config.models_path / info.base.value / info.type.value / (preferred_name or model_path.name)
+        )
         try:
             new_path = self._copy_model(model_path, dest_path)
         except FileExistsError as excp:
@@ -177,7 +176,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 +239,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 +311,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 +375,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 +406,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")
@@ -465,8 +542,10 @@ class ModelInstallService(ModelInstallServiceBase):
     def _register(
         self, model_path: Path, config: Optional[Dict[str, Any]] = None, info: Optional[AnyModelConfig] = None
     ) -> str:
-        info = info or ModelProbe.probe(model_path, config)
         key = self._create_key()
+        if config and not config.get("key", None):
+            config["key"] = key
+        info = info or ModelProbe.probe(model_path, config)
 
         model_path = model_path.absolute()
         if model_path.is_relative_to(self.app_config.models_path):
@@ -479,8 +558,8 @@ class ModelInstallService(ModelInstallServiceBase):
             # make config relative to our root
             legacy_conf = (self.app_config.root_dir / self.app_config.legacy_conf_dir / info.config).resolve()
             info.config = legacy_conf.relative_to(self.app_config.root_dir).as_posix()
-        self.record_store.add_model(key, info)
-        return key
+        self.record_store.add_model(info.key, info)
+        return info.key
 
     def _next_id(self) -> int:
         with self._lock:
@@ -489,10 +568,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 +596,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 +644,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 +661,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 +680,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")
@@ -652,6 +743,7 @@ class ModelInstallService(ModelInstallServiceBase):
                 self._signal_job_downloading(install_job)
 
     def _download_complete_callback(self, download_job: DownloadJob) -> None:
+        self._logger.info(f"{download_job.source}: model download complete")
         with self._lock:
             install_job = self._download_cache[download_job.source]
             self._download_cache.pop(download_job.source, None)
@@ -684,7 +776,7 @@ class ModelInstallService(ModelInstallServiceBase):
             if not install_job:
                 return
             self._downloads_changed_event.set()
-            self._logger.warning(f"Download {download_job.source} cancelled.")
+            self._logger.warning(f"{download_job.source}: model download cancelled")
             # if install job has already registered an error, then do not replace its status with cancelled
             if not install_job.errored:
                 install_job.cancel()
@@ -731,6 +823,7 @@ class ModelInstallService(ModelInstallServiceBase):
                 parts=parts,
                 bytes=job.bytes,
                 total_bytes=job.total_bytes,
+                id=job.id,
             )
 
     def _signal_job_completed(self, job: ModelInstallJob) -> None:
@@ -743,7 +836,7 @@ class ModelInstallService(ModelInstallServiceBase):
             assert job.local_path is not None
             assert job.config_out is not None
             key = job.config_out.key
-            self._event_bus.emit_model_install_completed(str(job.source), key)
+            self._event_bus.emit_model_install_completed(str(job.source), key, id=job.id)
 
     def _signal_job_errored(self, job: ModelInstallJob) -> None:
         self._logger.info(f"{job.source}: model installation encountered an exception: {job.error_type}\n{job.error}")
@@ -752,7 +845,7 @@ class ModelInstallService(ModelInstallServiceBase):
             error = job.error
             assert error_type is not None
             assert error is not None
-            self._event_bus.emit_model_install_error(str(job.source), error_type, error)
+            self._event_bus.emit_model_install_error(str(job.source), error_type, error, id=job.id)
 
     def _signal_job_cancelled(self, job: ModelInstallJob) -> None:
         self._logger.info(f"{job.source}: model installation was cancelled")

invokeai/app/services/model_load/__init__.py

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

invokeai/app/services/model_load/model_load_default.py

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

invokeai/app/services/model_manager/__init__.py

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

invokeai/app/services/model_manager/model_manager_default.py

@@ -1,413 +1,155 @@
 # Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
+"""Implementation of ModelManagerServiceBase."""
 
-from __future__ import annotations
-
-from logging import Logger
-from pathlib import Path
-from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
+from typing import Optional
 
 import torch
-from pydantic import Field
+from typing_extensions import Self
 
-from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException
-from invokeai.backend.model_management import (
-    AddModelResult,
-    BaseModelType,
-    MergeInterpolationMethod,
-    ModelInfo,
-    ModelManager,
-    ModelMerger,
-    ModelNotFoundException,
-    ModelType,
-    SchedulerPredictionType,
-    SubModelType,
-)
-from invokeai.backend.model_management.model_cache import CacheStats
-from invokeai.backend.model_management.model_search import FindModels
-from invokeai.backend.util import choose_precision, choose_torch_device
+from invokeai.app.services.invoker import Invoker
+from invokeai.app.services.shared.invocation_context import InvocationContextData
+from invokeai.backend.model_manager import AnyModelConfig, BaseModelType, LoadedModel, ModelType, SubModelType
+from invokeai.backend.model_manager.load import ModelCache, ModelConvertCache, ModelLoaderRegistry
+from invokeai.backend.util.devices import choose_torch_device
+from invokeai.backend.util.logging import InvokeAILogger
 
+from ..config import InvokeAIAppConfig
+from ..download import DownloadQueueServiceBase
+from ..events.events_base import EventServiceBase
+from ..model_install import ModelInstallService, ModelInstallServiceBase
+from ..model_load import ModelLoadService, ModelLoadServiceBase
+from ..model_records import ModelRecordServiceBase, UnknownModelException
 from .model_manager_base import ModelManagerServiceBase
 
-if TYPE_CHECKING:
-    from invokeai.app.invocations.baseinvocation import InvocationContext
 
-
-# simple implementation
 class ModelManagerService(ModelManagerServiceBase):
-    """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 get_model(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        submodel: Optional[SubModelType] = None,
-        context: Optional[InvocationContext] = None,
-    ) -> ModelInfo:
-        """
-        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:
-            self._emit_load_event(
-                context=context,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-            )
-
-        model_info = self.mgr.get_model(
-            model_name,
-            base_model,
-            model_type,
-            submodel,
-        )
-
-        if context:
-            self._emit_load_event(
-                context=context,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-                model_info=model_info,
-            )
-
-        return model_info
-
-    def model_exists(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-    ) -> bool:
-        """
-        Given a model name, returns True if it is a valid
-        identifier.
-        """
-        return self.mgr.model_exists(
-            model_name,
-            base_model,
-            model_type,
-        )
-
-    def model_info(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> Union[dict, None]:
-        """
-        Given a model name returns a dict-like (OmegaConf) object describing it.
-        """
-        return self.mgr.model_info(model_name, base_model, model_type)
-
-    def model_names(self) -> List[Tuple[str, BaseModelType, ModelType]]:
-        """
-        Returns a list of all the model names known.
-        """
-        return self.mgr.model_names()
-
-    def list_models(
-        self, base_model: Optional[BaseModelType] = None, model_type: Optional[ModelType] = None
-    ) -> list[dict]:
-        """
-        Return a list of models.
-        """
-        return self.mgr.list_models(base_model, model_type)
-
-    def list_model(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> Union[dict, None]:
-        """
-        Return information about the model using the same format as list_models()
-        """
-        return self.mgr.list_model(model_name=model_name, base_model=base_model, model_type=model_type)
-
-    def add_model(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        model_attributes: dict,
-        clobber: bool = False,
-    ) -> AddModelResult:
-        """
-        Update the named model with a dictionary of attributes. Will fail with an
-        assertion error if the name already exists. Pass clobber=True to overwrite.
-        On a successful update, the config will be changed in memory. Will fail
-        with an assertion error if provided attributes are incorrect or
-        the model name is missing. Call commit() to write changes to disk.
-        """
-        self.logger.debug(f"add/update model {model_name}")
-        return self.mgr.add_model(model_name, base_model, model_type, model_attributes, clobber)
-
-    def update_model(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        model_attributes: dict,
-    ) -> AddModelResult:
-        """
-        Update the named model with a dictionary of attributes. Will fail with a
-        ModelNotFoundException exception if the name does not already exist.
-        On a successful update, the config will be changed in memory. Will fail
-        with an assertion error if provided attributes are incorrect or
-        the model name is missing. Call commit() to write changes to disk.
-        """
-        self.logger.debug(f"update model {model_name}")
-        if not self.model_exists(model_name, base_model, model_type):
-            raise ModelNotFoundException(f"Unknown model {model_name}")
-        return self.add_model(model_name, base_model, model_type, model_attributes, clobber=True)
-
-    def del_model(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-    ):
-        """
-        Delete the named model from configuration. If delete_files is true,
-        then the underlying weight file or diffusers directory will be deleted
-        as well.
-        """
-        self.logger.debug(f"delete model {model_name}")
-        self.mgr.del_model(model_name, base_model, model_type)
-        self.mgr.commit()
-
-    def convert_model(
-        self,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: Literal[ModelType.Main, ModelType.Vae],
-        convert_dest_directory: Optional[Path] = Field(
-            default=None, description="Optional directory location for merged model"
-        ),
-    ) -> AddModelResult:
-        """
-        Convert a checkpoint file into a diffusers folder, deleting the cached
-        version and deleting the original checkpoint file if it is in the models
-        directory.
-        :param model_name: Name of the model to convert
-        :param base_model: Base model type
-        :param model_type: Type of model ['vae' or 'main']
-        :param convert_dest_directory: Save the converted model to the designated directory (`models/etc/etc` by default)
-
-        This will raise a ValueError unless the model is not a checkpoint. It will
-        also raise a ValueError in the event that there is a similarly-named diffusers
-        directory already in place.
-        """
-        self.logger.debug(f"convert model {model_name}")
-        return self.mgr.convert_model(model_name, base_model, model_type, convert_dest_directory)
-
-    def collect_cache_stats(self, cache_stats: CacheStats):
-        """
-        Reset model cache statistics for graph with graph_id.
-        """
-        self.mgr.cache.stats = cache_stats
-
-    def commit(self, conf_file: Optional[Path] = None):
-        """
-        Write current configuration out to the indicated file.
-        If no conf_file is provided, then replaces the
-        original file/database used to initialize the object.
-        """
-        return self.mgr.commit(conf_file)
-
-    def _emit_load_event(
-        self,
-        context: InvocationContext,
-        model_name: str,
-        base_model: BaseModelType,
-        model_type: ModelType,
-        submodel: Optional[SubModelType] = None,
-        model_info: Optional[ModelInfo] = None,
-    ):
-        if context.services.queue.is_canceled(context.graph_execution_state_id):
-            raise CanceledException()
-
-        if model_info:
-            context.services.events.emit_model_load_completed(
-                queue_id=context.queue_id,
-                queue_item_id=context.queue_item_id,
-                queue_batch_id=context.queue_batch_id,
-                graph_execution_state_id=context.graph_execution_state_id,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-                model_info=model_info,
-            )
-        else:
-            context.services.events.emit_model_load_started(
-                queue_id=context.queue_id,
-                queue_item_id=context.queue_item_id,
-                queue_batch_id=context.queue_batch_id,
-                graph_execution_state_id=context.graph_execution_state_id,
-                model_name=model_name,
-                base_model=base_model,
-                model_type=model_type,
-                submodel=submodel,
-            )
+        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_data: Optional[InvocationContextData] = None,
+    ) -> LoadedModel:
+        return self.load.load_model(model_config, submodel_type, context_data)
 
-        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_data: Optional[InvocationContextData] = None,
+    ) -> LoadedModel:
+        config = self.store.get_model(key)
+        return self.load.load_model(config, submodel_type, context_data)
 
-    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_data: Optional[InvocationContextData] = 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_data)
+
+    @classmethod
+    def build_model_manager(
+        cls,
+        app_config: InvokeAIAppConfig,
+        model_record_service: ModelRecordServiceBase,
+        download_queue: DownloadQueueServiceBase,
+        events: EventServiceBase,
+        execution_device: torch.device = choose_torch_device(),
+    ) -> Self:
+        """
+        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,
+            execution_device=execution_device,
         )
+        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

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

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

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

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

invokeai/app/services/session_processor/session_processor_common.py

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

invokeai/app/services/session_processor/session_processor_default.py

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

invokeai/app/services/session_queue/session_queue_sqlite.py

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

invokeai/app/services/shared/default_graphs.py

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

invokeai/app/services/shared/graph.py

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

invokeai/app/services/shared/invocation_context.py

@@ -0,0 +1,470 @@
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+from PIL.Image import Image
+from torch import Tensor
+
+from invokeai.app.invocations.constants import IMAGE_MODES
+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.util.step_callback import stable_diffusion_step_callback
+from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType
+from invokeai.backend.model_manager.load.load_base import LoadedModel
+from invokeai.backend.model_manager.metadata.metadata_base import AnyModelRepoMetadata
+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
+    from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
+
+"""
+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:
+    queue_item: "SessionQueueItem"
+    """The queue item that is being executed."""
+    invocation: "BaseInvocation"
+    """The invocation that is being executed."""
+    source_invocation_id: str
+    """The ID of the invocation from which the currently executing invocation was prepared."""
+
+
+class InvocationContextInterface:
+    def __init__(self, services: InvocationServices, data: InvocationContextData) -> None:
+        self._services = services
+        self._data = 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._data.invocation, WithMetadata):
+            metadata_ = self._data.invocation.metadata
+
+        # If `board_id` is provided directly, use that. Else, use the board provided by `WithBoard`, falling back to None.
+        board_id_ = None
+        if board_id:
+            board_id_ = board_id
+        elif isinstance(self._data.invocation, WithBoard) and self._data.invocation.board:
+            board_id_ = self._data.invocation.board.board_id
+
+        return self._services.images.create(
+            image=image,
+            is_intermediate=self._data.invocation.is_intermediate,
+            image_category=image_category,
+            board_id=board_id_,
+            metadata=metadata_,
+            image_origin=ResourceOrigin.INTERNAL,
+            workflow=self._data.queue_item.workflow,
+            session_id=self._data.queue_item.session_id,
+            node_id=self._data.invocation.id,
+        )
+
+    def get_pil(self, image_name: str, mode: IMAGE_MODES | None = None) -> Image:
+        """
+        Gets an image as a PIL Image object.
+
+        :param image_name: The name of the image to get.
+        :param mode: The color mode to convert the image to. If None, the original mode is used.
+        """
+        image = self._services.images.get_pil_image(image_name)
+        if mode and mode != image.mode:
+            try:
+                image = image.convert(mode)
+            except ValueError:
+                self._services.logger.warning(
+                    f"Could not convert image from {image.mode} to {mode}. Using original mode instead."
+                )
+        return image
+
+    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_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, key: str) -> bool:
+        """
+        Checks if a model exists.
+
+        :param key: The key of the model.
+        """
+        return self._services.model_manager.store.exists(key)
+
+    def load(self, key: str, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+        """
+        Loads a model.
+
+        :param key: The key of the model.
+        :param submodel_type: 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.load_model_by_key(
+            key=key, submodel_type=submodel_type, context_data=self._data
+        )
+
+    def load_by_attrs(
+        self, model_name: str, base_model: BaseModelType, model_type: ModelType, submodel: Optional[SubModelType] = None
+    ) -> LoadedModel:
+        """
+        Loads a model by its attributes.
+
+        :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
+        """
+        return self._services.model_manager.load_model_by_attr(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+            submodel=submodel,
+            context_data=self._data,
+        )
+
+    def get_config(self, key: str) -> AnyModelConfig:
+        """
+        Gets a model's info, an dict-like object.
+
+        :param key: The key of the model.
+        """
+        return self._services.model_manager.store.get_model(key=key)
+
+    def get_metadata(self, key: str) -> Optional[AnyModelRepoMetadata]:
+        """
+        Gets a model's metadata, if it has any.
+
+        :param key: The key of the model.
+        """
+        return self._services.model_manager.store.get_metadata(key=key)
+
+    def search_by_path(self, path: Path) -> list[AnyModelConfig]:
+        """
+        Searches for models by path.
+
+        :param path: The path to search for.
+        """
+        return self._services.model_manager.store.search_by_path(path)
+
+    def search_by_attrs(
+        self,
+        model_name: Optional[str] = None,
+        base_model: Optional[BaseModelType] = None,
+        model_type: Optional[ModelType] = None,
+        model_format: Optional[ModelFormat] = None,
+    ) -> list[AnyModelConfig]:
+        """
+        Searches for models by attributes.
+
+        :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
+        """
+
+        return self._services.model_manager.store.search_by_attr(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+            model_format=model_format,
+        )
+
+
+class ConfigInterface(InvocationContextInterface):
+    def get(self) -> InvokeAIAppConfig:
+        """Gets the app's config."""
+
+        return self._services.configuration.get_config()
+
+
+class UtilInterface(InvocationContextInterface):
+    def __init__(
+        self, services: InvocationServices, data: InvocationContextData, cancel_event: threading.Event
+    ) -> None:
+        super().__init__(services, data)
+        self._cancel_event = cancel_event
+
+    def is_canceled(self) -> bool:
+        """Checks if the current invocation has been canceled."""
+        return self._cancel_event.is_set()
+
+    def sd_step_callback(self, intermediate_state: PipelineIntermediateState, base_model: BaseModelType) -> None:
+        """
+        The step callback emits a progress event with the current step, the total number of
+        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.
+        """
+
+        stable_diffusion_step_callback(
+            context_data=self._data,
+            intermediate_state=intermediate_state,
+            base_model=base_model,
+            events=self._services.events,
+            is_canceled=self.is_canceled,
+        )
+
+
+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,
+        data: InvocationContextData,
+        services: InvocationServices,
+    ) -> None:
+        self.images = images
+        """Methods to save, get and update images and their metadata."""
+        self.tensors = tensors
+        """Methods to save and get tensors, including image, noise, masks, and masked images."""
+        self.conditioning = conditioning
+        """Methods to save and get conditioning data."""
+        self.models = models
+        """Methods to check if a model exists, get a model, and get a model's info."""
+        self.logger = logger
+        """The app logger."""
+        self.config = config
+        """The app config."""
+        self.util = util
+        """Utility methods, including a method to check if an invocation was canceled and step callbacks."""
+        self.boards = boards
+        """Methods to interact with boards."""
+        self._data = data
+        """An internal API providing access to data about the current queue item and invocation. You probably shouldn't use this. It may change without warning."""
+        self._services = services
+        """An internal API providing access to all application services. You probably shouldn't use this. It may change without warning."""
+
+
+def build_invocation_context(
+    services: InvocationServices,
+    data: InvocationContextData,
+    cancel_event: threading.Event,
+) -> InvocationContext:
+    """
+    Builds the invocation context for a specific invocation execution.
+
+    :param services: The invocation services to wrap.
+    :param data: The invocation context data.
+    """
+
+    logger = LoggerInterface(services=services, data=data)
+    images = ImagesInterface(services=services, data=data)
+    tensors = TensorsInterface(services=services, data=data)
+    models = ModelsInterface(services=services, data=data)
+    config = ConfigInterface(services=services, data=data)
+    util = UtilInterface(services=services, data=data, cancel_event=cancel_event)
+    conditioning = ConditioningInterface(services=services, data=data)
+    boards = BoardsInterface(services=services, data=data)
+
+    ctx = InvocationContext(
+        images=images,
+        logger=logger,
+        config=config,
+        tensors=tensors,
+        models=models,
+        data=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

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

invokeai/app/util/ti_utils.py

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