add missing import

Makefile

@@ -18,7 +18,6 @@ help:
 	@echo "frontend-typegen         Generate types for the frontend from the OpenAPI schema"
 	@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 "openapi                  Generate the OpenAPI schema for the app, outputting to stdout"
 
 # Runs ruff, fixing any safely-fixable errors and formatting
 ruff:
@@ -71,6 +70,3 @@ installer-zip:
 tag-release:
 	cd installer && ./tag_release.sh
 
-# Generate the OpenAPI Schema for the app
-openapi:
-	python scripts/generate_openapi_schema.py

README.md

@@ -2,114 +2,21 @@
 
 ![project hero](https://github.com/invoke-ai/InvokeAI/assets/31807370/6e3728c7-e90e-4711-905c-3b55844ff5be)
 
-# Invoke - Professional Creative AI Tools for Visual Media
-
-#### To learn more about Invoke, or implement our Business solutions, visit [invoke.com]
-
-[![discord badge]][discord link] [![latest release badge]][latest release link] [![github stars badge]][github stars link] [![github forks badge]][github forks link] [![CI checks on main badge]][CI checks on main link] [![latest commit to main badge]][latest commit to main link] [![github open issues badge]][github open issues link] [![github open prs badge]][github open prs link] [![translation status badge]][translation status link]
-
-</div>
-
-Invoke is a leading creative engine built to empower professionals and enthusiasts alike. Generate and create stunning visual media using the latest AI-driven technologies. Invoke offers an industry leading web-based UI, and serves as the foundation for multiple commercial products.
-
-Invoke is available in two editions:
-
-| **Community Edition**                                                                                                      | **Professional Edition**                                                                            |
-|----------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------|
-| **For users looking for a locally installed, self-hosted and self-managed service**                                         | **For users or teams looking for a cloud-hosted, fully managed service**                            |
-| - Free to use under a commercially-friendly license                                                                         | - Monthly subscription fee with three different plan levels                                         |
-| - Download and install on compatible hardware                                                                               | - Offers additional benefits, including multi-user support, improved model training, and more                          |
-| - Includes all core studio features: generate, refine, iterate on images, and build workflows                               | - Hosted in the cloud for easy, secure model access and scalability                                               |
-| Quick Start -> [Installation and Updates][installation docs]                                                                     | More Information -> [www.invoke.com/pricing](https://www.invoke.com/pricing)                        |
+# Invoke - Professional Creative AI Tools for Visual Media 
+##  To learn more about Invoke, or implement our Business solutions, visit [invoke.com](https://www.invoke.com/about)
+  
 
 
-![Highlighted Features - Canvas and Workflows](https://github.com/invoke-ai/InvokeAI/assets/31807370/708f7a82-084f-4860-bfbe-e2588c53548d)
+[![discord badge]][discord link]
 
-# Documentation
-| **Quick Links**                                                                                                      | 
-|----------------------------------------------------------------------------------------------------------------------------|
-|  [Installation and Updates][installation docs] - [Documentation and Tutorials][docs home] - [Bug Reports][github issues] - [Contributing][contributing docs]  | 
+[![latest release badge]][latest release link] [![github stars badge]][github stars link] [![github forks badge]][github forks link]
 
-</div>
+[![CI checks on main badge]][CI checks on main link] [![latest commit to main badge]][latest commit to main link]
 
-## Quick Start
+[![github open issues badge]][github open issues link] [![github open prs badge]][github open prs link] [![translation status badge]][translation status link]
 
-1. Download and unzip the installer from the bottom of the [latest release][latest release link].
-2. Run the installer script.
-
-   - **Windows**: Double-click on the `install.bat` script.
-   - **macOS**: Open a Terminal window, drag the file `install.sh` from Finder into the Terminal, and press enter.
-   - **Linux**: Run `install.sh`.
-
-3. When prompted, enter a location for the install and select your GPU type.
-4. Once the install finishes, find the directory you selected during install. The default location is `C:\Users\Username\invokeai` for Windows or `~/invokeai` for Linux/macOS.
-5. Run the launcher script (`invoke.bat` for Windows, `invoke.sh` for macOS and Linux) the same way you ran the installer script in step 2.
-6. Select option 1 to start the application. Once it starts up, open your browser and go to <http://localhost:9090>.
-7. Open the model manager tab to install a starter model and then you'll be ready to generate.
-
-More detail, including hardware requirements and manual install instructions, are available in the [installation documentation][installation docs].
-
-## Troubleshooting, FAQ and Support
-
-Please review our [FAQ][faq] for solutions to common installation problems and other issues.
-
-For more help, please join our [Discord][discord link].
-
-## Features
-
-Full details on features can be found in [our documentation][features docs].
-
-### Web Server & UI
-
-Invoke runs a locally hosted web server & React UI with an industry-leading user experience.
-
-### Unified Canvas
-
-The Unified Canvas is a fully integrated canvas implementation with support for all core generation capabilities, in/out-painting, brush tools, and more. This creative tool unlocks the capability for artists to create with AI as a creative collaborator, and can be used to augment AI-generated imagery, sketches, photography, renders, and more.
-
-### Workflows & Nodes
-
-Invoke offers a fully featured workflow management solution, enabling users to combine the power of node-based workflows with the easy of a UI. This allows for customizable generation pipelines to be developed and shared by users looking to create specific workflows to support their production use-cases.
-
-### Board & Gallery Management
-
-Invoke features an organized gallery system for easily storing, accessing, and remixing your content in the Invoke workspace. Images can be dragged/dropped onto any Image-base UI element in the application, and rich metadata within the Image allows for easy recall of key prompts or settings used in your workflow.
-
-### Other features
-
-- Support for both ckpt and diffusers models
-- SD1.5, SD2.0, and SDXL support
-- Upscaling Tools
-- Embedding Manager & Support
-- Model Manager & Support
-- Workflow creation & management
-- Node-Based Architecture
-
-## Contributing
-
-Anyone who wishes to contribute to this project - whether documentation, features, bug fixes, code cleanup, testing, or code reviews - is very much encouraged to do so.
-
-Get started with contributing by reading our [contribution documentation][contributing docs], joining the [#dev-chat] or the GitHub discussion board.
-
-We hope you enjoy using Invoke as much as we enjoy creating it, and we hope you will elect to become part of our community.
-
-## Thanks
-
-Invoke is a combined effort of [passionate and talented people from across the world][contributors]. We thank them for their time, hard work and effort.
-
-Original portions of the software are Copyright © 2024 by respective contributors.
-
-[features docs]: https://invoke-ai.github.io/InvokeAI/features/
-[faq]: https://invoke-ai.github.io/InvokeAI/help/FAQ/
-[contributors]: https://invoke-ai.github.io/InvokeAI/other/CONTRIBUTORS/
-[invoke.com]: https://www.invoke.com/about
-[github issues]: https://github.com/invoke-ai/InvokeAI/issues
-[docs home]: https://invoke-ai.github.io/InvokeAI
-[installation docs]: https://invoke-ai.github.io/InvokeAI/installation/INSTALLATION/
-[#dev-chat]: https://discord.com/channels/1020123559063990373/1049495067846524939
-[contributing docs]: https://invoke-ai.github.io/InvokeAI/contributing/CONTRIBUTING/
 [CI checks on main badge]: https://flat.badgen.net/github/checks/invoke-ai/InvokeAI/main?label=CI%20status%20on%20main&cache=900&icon=github
-[CI checks on main link]: https://github.com/invoke-ai/InvokeAI/actions?query=branch%3Amain
+[CI checks on main link]:https://github.com/invoke-ai/InvokeAI/actions?query=branch%3Amain
 [discord badge]: https://flat.badgen.net/discord/members/ZmtBAhwWhy?icon=discord
 [discord link]: https://discord.gg/ZmtBAhwWhy
 [github forks badge]: https://flat.badgen.net/github/forks/invoke-ai/InvokeAI?icon=github
@@ -123,6 +30,402 @@ Original portions of the software are Copyright © 2024 by respective contributo
 [latest commit to main badge]: https://flat.badgen.net/github/last-commit/invoke-ai/InvokeAI/main?icon=github&color=yellow&label=last%20dev%20commit&cache=900
 [latest commit to main link]: https://github.com/invoke-ai/InvokeAI/commits/main
 [latest release badge]: https://flat.badgen.net/github/release/invoke-ai/InvokeAI/development?icon=github
-[latest release link]: https://github.com/invoke-ai/InvokeAI/releases/latest
+[latest release link]: https://github.com/invoke-ai/InvokeAI/releases
 [translation status badge]: https://hosted.weblate.org/widgets/invokeai/-/svg-badge.svg
 [translation status link]: https://hosted.weblate.org/engage/invokeai/
+
+</div>
+
+InvokeAI is a leading creative engine built to empower professionals
+and enthusiasts alike. Generate and create stunning visual media using
+the latest AI-driven technologies. InvokeAI offers an industry leading
+Web Interface, interactive Command Line Interface, and also serves as
+the foundation for multiple commercial products.
+
+**Quick links**: [[How to
+  Install](https://invoke-ai.github.io/InvokeAI/installation/INSTALLATION/)] [<a
+  href="https://discord.gg/ZmtBAhwWhy">Discord Server</a>] [<a
+  href="https://invoke-ai.github.io/InvokeAI/">Documentation and
+  Tutorials</a>]
+  [<a href="https://github.com/invoke-ai/InvokeAI/issues">Bug Reports</a>]
+  [<a
+  href="https://github.com/invoke-ai/InvokeAI/discussions">Discussion,
+  Ideas & Q&A</a>] 
+   [<a
+  href="https://invoke-ai.github.io/InvokeAI/contributing/CONTRIBUTING/">Contributing</a>] 
+
+<div align="center">
+
+
+![Highlighted Features - Canvas and Workflows](https://github.com/invoke-ai/InvokeAI/assets/31807370/708f7a82-084f-4860-bfbe-e2588c53548d)
+
+
+</div>
+
+## Table of Contents
+
+Table of Contents 📝
+
+**Getting Started**
+1. 🏁 [Quick Start](#quick-start) 
+3. 🖥️ [Hardware Requirements](#hardware-requirements) 
+
+**More About Invoke**
+1. 🌟 [Features](#features) 
+2. 📣 [Latest Changes](#latest-changes) 
+3. 🛠️ [Troubleshooting](#troubleshooting) 
+
+**Supporting the Project**
+1. 🤝 [Contributing](#contributing) 
+2. 👥 [Contributors](#contributors) 
+3. 💕 [Support](#support) 
+
+## Quick Start
+
+For full installation and upgrade instructions, please see:
+[InvokeAI Installation Overview](https://invoke-ai.github.io/InvokeAI/installation/INSTALLATION/)
+
+If upgrading from version 2.3, please read [Migrating a 2.3 root
+directory to 3.0](#migrating-to-3) first.
+
+### Automatic Installer (suggested for 1st time users)
+
+1. Go to the bottom of the [Latest Release Page](https://github.com/invoke-ai/InvokeAI/releases/latest)
+
+2. Download the .zip file for your OS (Windows/macOS/Linux).
+
+3. Unzip the file.
+
+4. **Windows:** double-click on the `install.bat` script. **macOS:** Open a Terminal window, drag the file `install.sh` from Finder
+into the Terminal, and press return. **Linux:** run `install.sh`.
+
+5. You'll be asked to confirm the location of the folder in which
+to install InvokeAI and its image generation model files. Pick a
+location with at least 15 GB of free memory. More if you plan on
+installing lots of models.
+
+6. Wait while the installer does its thing. After installing the software,
+the installer will launch a script that lets you configure InvokeAI and
+select a set of starting image generation models.
+
+7. Find the folder that InvokeAI was installed into (it is not the
+same as the unpacked zip file directory!) The default location of this
+folder (if you didn't change it in step 5) is `~/invokeai` on
+Linux/Mac systems, and `C:\Users\YourName\invokeai` on Windows. This directory will contain launcher scripts named `invoke.sh` and `invoke.bat`.
+
+8. On Windows systems, double-click on the `invoke.bat` file. On
+macOS, open a Terminal window, drag `invoke.sh` from the folder into
+the Terminal, and press return. On Linux, run `invoke.sh`
+
+9. Press 2 to open the "browser-based UI", press enter/return, wait a
+minute or two for Stable Diffusion to start up, then open your browser
+and go to http://localhost:9090.
+
+10. Type `banana sushi` in the box on the top left and click `Invoke`
+
+### Command-Line Installation (for developers and users familiar with Terminals)
+
+You must have Python 3.10 through 3.11 installed on your machine. Earlier or
+later versions are not supported.
+Node.js also needs to be installed along with `pnpm` (can be installed with
+the command `npm install -g pnpm` if needed)
+
+1. Open a command-line window on your machine. The PowerShell is recommended for Windows.
+2. Create a directory to install InvokeAI into. You'll need at least 15 GB of free space:
+
+    ```terminal
+    mkdir invokeai
+    ````
+
+3. Create a virtual environment named `.venv` inside this directory and activate it:
+
+    ```terminal
+    cd invokeai
+    python -m venv .venv --prompt InvokeAI
+    ```
+
+4. Activate the virtual environment (do it every time you run InvokeAI)
+
+    _For Linux/Mac users:_
+
+    ```sh
+    source .venv/bin/activate
+    ```
+
+    _For Windows users:_
+
+    ```ps
+    .venv\Scripts\activate
+    ```
+
+5. Install the InvokeAI module and its dependencies. Choose the command suited for your platform & GPU.
+
+    _For Windows/Linux with an NVIDIA GPU:_
+
+    ```terminal
+    pip install "InvokeAI[xformers]" --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu121
+    ```
+
+    _For Linux with an AMD GPU:_
+
+    ```sh
+    pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.6
+    ```
+
+    _For non-GPU systems:_
+    ```terminal
+    pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/cpu
+    ``` 
+
+    _For Macintoshes, either Intel or M1/M2/M3:_
+
+    ```sh
+    pip install InvokeAI --use-pep517
+    ```
+
+6. Configure InvokeAI and install a starting set of image generation models (you only need to do this once):
+
+    ```terminal
+    invokeai-configure --root .
+    ```
+	Don't miss the dot at the end!
+
+7. Launch the web server (do it every time you run InvokeAI):
+
+    ```terminal
+    invokeai-web
+    ```
+
+8. Point your browser to http://localhost:9090 to bring up the web interface.
+
+9. Type `banana sushi` in the box on the top left and click `Invoke`.
+
+Be sure to activate the virtual environment each time before re-launching InvokeAI,
+using `source .venv/bin/activate` or `.venv\Scripts\activate`.
+
+## Detailed Installation Instructions
+
+This fork is supported across Linux, Windows and Macintosh. Linux
+users can use either an Nvidia-based card (with CUDA support) or an
+AMD card (using the ROCm driver). For full installation and upgrade
+instructions, please see:
+[InvokeAI Installation Overview](https://invoke-ai.github.io/InvokeAI/installation/INSTALL_SOURCE/)
+
+<a name="migrating-to-3"></a>
+### Migrating a v2.3 InvokeAI root directory
+
+The InvokeAI root directory is where the InvokeAI startup file,
+installed models, and generated images are stored. It is ordinarily
+named `invokeai` and located in your home directory. The contents and
+layout of this directory has changed between versions 2.3 and 3.0 and
+cannot be used directly.
+
+We currently recommend that you use the installer to create a new root
+directory named differently from the 2.3 one, e.g. `invokeai-3` and
+then use a migration script to copy your 2.3 models into the new
+location. However, if you choose, you can upgrade this directory in
+place.  This section gives both recipes.
+
+#### Creating a new root directory and migrating old models
+
+This is the safer recipe because it leaves your old root directory in
+place to fall back on.
+
+1. Follow the instructions above to create and install InvokeAI in a
+directory that has a different name from the 2.3 invokeai directory.
+In this example, we will use "invokeai-3"
+
+2. When you are prompted to select models to install, select a minimal
+set of models, such as stable-diffusion-v1.5 only.
+
+3. After installation is complete launch `invokeai.sh` (Linux/Mac) or
+`invokeai.bat` and select option 8 "Open the developers console". This
+will take you to the command line.
+
+4. Issue the command `invokeai-migrate3 --from /path/to/v2.3-root --to
+/path/to/invokeai-3-root`. Provide the correct `--from` and `--to`
+paths for your v2.3 and v3.0 root directories respectively.
+
+This will copy and convert your old models from 2.3 format to 3.0
+format and create a new `models` directory in the 3.0 directory. The
+old models directory (which contains the models selected at install
+time) will be renamed `models.orig` and can be deleted once you have
+confirmed that the migration was successful.
+
+ If you wish, you can pass the 2.3 root directory to both `--from` and
+`--to` in order to update in place. Warning: this directory will no
+longer be usable with InvokeAI 2.3.
+
+#### Migrating in place
+
+For the adventurous, you may do an in-place upgrade from 2.3 to 3.0
+without touching the command line. ***This recipe does not work on
+Windows platforms due to a bug in the Windows version of the 2.3
+upgrade script.** See the next section for a Windows recipe.
+
+##### For Mac and Linux Users:
+
+1. Launch the InvokeAI launcher script in your current v2.3 root directory.
+
+2. Select option [9] "Update InvokeAI" to bring up the updater dialog.
+
+3. Select option [1] to upgrade to the latest release.
+
+4. Once the upgrade is finished you will be returned to the launcher
+menu. Select option [6] "Re-run the configure script to fix a broken
+install or to complete a major upgrade".
+
+This will run the configure script against the v2.3 directory and
+update it to the 3.0 format. The following files will be replaced:
+
+  - The invokeai.init file, replaced by invokeai.yaml
+  - The models directory
+  - The configs/models.yaml model index
+  
+The original versions of these files will be saved with the suffix
+".orig" appended to the end. Once you have confirmed that the upgrade
+worked, you can safely remove these files. Alternatively you can
+restore a working v2.3 directory by removing the new files and
+restoring the ".orig" files' original names.
+
+##### For Windows Users:
+
+Windows Users can upgrade with the
+
+1. Enter the 2.3 root directory you wish to upgrade
+2. Launch `invoke.sh` or `invoke.bat`
+3. Select the "Developer's console" option [8]
+4. Type the following commands
+
+```
+pip install "invokeai @ https://github.com/invoke-ai/InvokeAI/archive/refs/tags/v3.0.0" --use-pep517 --upgrade
+invokeai-configure --root .
+```
+(Replace `v3.0.0` with the current release number if this document is out of date).
+
+The first command will install and upgrade new software to run
+InvokeAI. The second will prepare the 2.3 directory for use with 3.0.
+You may now launch the WebUI in the usual way, by selecting option [1]
+from the launcher script
+
+#### Migrating Images
+
+The migration script will migrate your invokeai settings and models,
+including textual inversion models, LoRAs and merges that you may have
+installed previously. However it does **not** migrate the generated
+images stored in your 2.3-format outputs directory. To do this, you 
+need to run an additional step:
+
+1. From a working InvokeAI 3.0 root directory, start the launcher and
+enter menu option [8] to open the "developer's console".
+
+2. At the developer's console command line, type the command:
+
+```bash
+invokeai-import-images
+```
+
+3. This will lead you through the process of confirming the desired
+   source and destination for the imported images. The images will
+   appear in the gallery board of your choice, and contain the
+   original prompt, model name, and other parameters used to generate
+   the image.
+   
+(Many kudos to **techjedi** for contributing this script.)
+
+## Hardware Requirements
+
+InvokeAI is supported across Linux, Windows and macOS. Linux
+users can use either an Nvidia-based card (with CUDA support) or an
+AMD card (using the ROCm driver).
+
+### System
+
+You will need one of the following:
+
+- An NVIDIA-based graphics card with 4 GB or more VRAM memory. 6-8 GB
+  of VRAM is highly recommended for rendering using the Stable
+  Diffusion XL models
+- An Apple computer with an M1 chip.
+- An AMD-based graphics card with 4GB or more VRAM memory (Linux
+  only), 6-8 GB for XL rendering.
+
+We do not recommend the GTX 1650 or 1660 series video cards. They are
+unable to run in half-precision mode and do not have sufficient VRAM
+to render 512x512 images.
+
+**Memory** - At least 12 GB Main Memory RAM.
+
+**Disk** - At least 12 GB of free disk space for the machine learning model, Python, and all its dependencies.
+
+## Features
+
+Feature documentation can be reviewed by navigating to [the InvokeAI Documentation page](https://invoke-ai.github.io/InvokeAI/features/)
+
+### *Web Server & UI*
+
+InvokeAI offers a locally hosted Web Server & React Frontend, with an industry leading user experience. The Web-based UI allows for simple and intuitive workflows, and is responsive for use on mobile devices and tablets accessing the web server.
+
+### *Unified Canvas*
+
+The Unified Canvas is a fully integrated canvas implementation with support for all core generation capabilities, in/outpainting, brush tools, and more. This creative tool unlocks the capability for artists to create with AI as a creative collaborator, and can be used to augment AI-generated imagery, sketches, photography, renders, and more.
+
+### *Workflows & Nodes*
+
+InvokeAI offers a fully featured workflow management solution, enabling users to combine the power of nodes based workflows with the easy of a UI. This allows for customizable generation pipelines to be developed and shared by users looking to create specific workflows to support their production use-cases.
+
+### *Board & Gallery Management*
+
+Invoke AI provides an organized gallery system for easily storing, accessing, and remixing your content in the Invoke workspace. Images can be dragged/dropped onto any Image-base UI element in the application, and rich metadata within the Image allows for easy recall of key prompts or settings used in your workflow. 
+
+### Other features
+
+- *Support for both ckpt and diffusers models*
+- *SD 2.0, 2.1, XL support*
+- *Upscaling Tools*
+- *Embedding Manager & Support*
+- *Model Manager & Support*
+- *Workflow creation & management*
+- *Node-Based Architecture*
+
+
+### Latest Changes
+
+For our latest changes, view our [Release
+Notes](https://github.com/invoke-ai/InvokeAI/releases) and the
+[CHANGELOG](docs/CHANGELOG.md).
+
+### Troubleshooting / FAQ
+
+Please check out our **[FAQ](https://invoke-ai.github.io/InvokeAI/help/FAQ/)** to get solutions for common installation
+problems and other issues. For more help, please join our [Discord][discord link]
+
+## Contributing
+
+Anyone who wishes to contribute to this project, whether documentation, features, bug fixes, code
+cleanup, testing, or code reviews, is very much encouraged to do so.
+
+Get started with contributing by reading our [Contribution documentation](https://invoke-ai.github.io/InvokeAI/contributing/CONTRIBUTING/), joining the [#dev-chat](https://discord.com/channels/1020123559063990373/1049495067846524939) or the GitHub discussion board.
+
+If you are unfamiliar with how
+to contribute to GitHub projects, we have a new contributor checklist you can follow to get started contributing: 
+[New Contributor Checklist](https://invoke-ai.github.io/InvokeAI/contributing/contribution_guides/newContributorChecklist/).
+
+We hope you enjoy using our software as much as we enjoy creating it,
+and we hope that some of those of you who are reading this will elect
+to become part of our community.
+
+Welcome to InvokeAI!
+
+### Contributors
+
+This fork is a combined effort of various people from across the world.
+[Check out the list of all these amazing people](https://invoke-ai.github.io/InvokeAI/other/CONTRIBUTORS/). We thank them for
+their time, hard work and effort.
+
+### Support
+
+For support, please use this repository's GitHub Issues tracking service, or join the [Discord][discord link].
+
+Original portions of the software are Copyright (c) 2023 by respective contributors.
+

docker/README.md

@@ -64,7 +64,7 @@ GPU_DRIVER=nvidia
 
 Any environment variables supported by InvokeAI can be set here - please see the [Configuration docs](https://invoke-ai.github.io/InvokeAI/features/CONFIGURATION/) for further detail.
 
-## Even More Customizing!
+## Even Moar Customizing!
 
 See the `docker-compose.yml` file. The `command` instruction can be uncommented and used to run arbitrary startup commands. Some examples below.
 

docs/contributing/DOWNLOAD_QUEUE.md

@@ -128,8 +128,7 @@ The queue operates on a series of download job objects. These objects
 specify the source and destination of the download, and keep track of
 the progress of the download.
 
-Two job types are defined. `DownloadJob` and
-`MultiFileDownloadJob`. The former is a pydantic object with the
+The only job type currently implemented is `DownloadJob`, a pydantic object with the
 following fields:
 
 | **Field**      | **Type**        |  **Default**  | **Description** |
@@ -139,7 +138,7 @@ following fields:
 | `dest`           | Path            |               | Where to download to              |
 | `access_token`   | str             |               | [optional] string containing authentication token for access |
 | `on_start`       | Callable        |               | [optional] callback when the download starts |
-| `on_progress` | Callable | | [optional] callback called at intervals during download progress |
+| `on_progress`    | Callable        |               | [optional] callback called at intervals during download progress |
 | `on_complete`    | Callable        |               | [optional] callback called after successful download completion |
 | `on_error`       | Callable        |               | [optional] callback called after an error occurs  |
 | `id`             | int             | auto assigned | Job ID, an integer >= 0           |
@@ -191,33 +190,6 @@ A cancelled job will have status `DownloadJobStatus.ERROR` and an
 `error_type` field of "DownloadJobCancelledException". In addition,
 the job's `cancelled` property will be set to True.
 
-The `MultiFileDownloadJob` is used for diffusers model downloads,
-which contain multiple files and directories under a common root:
-
-| **Field**      | **Type**        |  **Default**  | **Description** |
-|----------------|-----------------|---------------|-----------------|
-| _Fields passed in at job creation time_                               |
-| `download_parts` | Set[DownloadJob]|               | Component download jobs |
-| `dest`           | Path            |               | Where to download to              |
-| `on_start`       | Callable        |               | [optional] callback when the download starts |
-| `on_progress` | Callable | | [optional] callback called at intervals during download progress |
-| `on_complete`    | Callable        |               | [optional] callback called after successful download completion |
-| `on_error`       | Callable        |               | [optional] callback called after an error occurs  |
-| `id`             | int             | auto assigned | Job ID, an integer >= 0           |
-| _Fields updated over the course of the download task_
-| `status`         | DownloadJobStatus|              | Status code                                |
-| `download_path`  | Path |              | Path to the root of the downloaded files |
-| `bytes`          | int              | 0            | Bytes downloaded so far   |
-| `total_bytes`    | int              | 0            | Total size of the file at the remote site  |
-| `error_type`     | str              |              | String version of the exception that caused an error during download |
-| `error`          | str              |              | String version of the traceback associated with an error |
-| `cancelled`      | bool             | False        | Set to true if the job was cancelled by the caller|
-
-Note that the MultiFileDownloadJob does not support the `priority`,
-`job_started`, `job_ended` or `content_type` attributes. You can get
-these from the individual download jobs in `download_parts`.
-
-
 ### Callbacks
 
 Download jobs can be associated with a series of callbacks, each with
@@ -279,40 +251,11 @@ jobs using `list_jobs()`, fetch a single job by its with
 running jobs with `cancel_all_jobs()`, and wait for all jobs to finish
 with `join()`.
 
-#### job = queue.download(source, dest, priority, access_token, on_start, on_progress, on_complete, on_cancelled, on_error)
+#### job = queue.download(source, dest, priority, access_token)
 
 Create a new download job and put it on the queue, returning the
 DownloadJob object.
 
-#### multifile_job = queue.multifile_download(parts, dest, access_token, on_start, on_progress, on_complete, on_cancelled, on_error)
-
-This is similar to download(), but instead of taking a single source,
-it accepts a `parts` argument consisting of a list of
-`RemoteModelFile` objects. Each part corresponds to a URL/Path pair,
-where the URL is the location of the remote file, and the Path is the
-destination.
-
-`RemoteModelFile` can be imported from `invokeai.backend.model_manager.metadata`, and
-consists of a url/path pair. Note that the path *must* be relative.
-
-The method returns a `MultiFileDownloadJob`.
-
-
-```
-from invokeai.backend.model_manager.metadata import RemoteModelFile
-remote_file_1 = RemoteModelFile(url='http://www.foo.bar/my/pytorch_model.safetensors'',
-                                path='my_model/textencoder/pytorch_model.safetensors'
-			 			  )
-remote_file_2 = RemoteModelFile(url='http://www.bar.baz/vae.ckpt',
-                                path='my_model/vae/diffusers_model.safetensors'
-			 			  )
-job = queue.multifile_download(parts=[remote_file_1, remote_file_2],
-                               dest='/tmp/downloads',
-                               on_progress=TqdmProgress().update)
-queue.wait_for_job(job)
-print(f"The files were downloaded to {job.download_path}")
-```
-
 #### jobs = queue.list_jobs()
 
 Return a list of all active and inactive `DownloadJob`s.

docs/contributing/MODEL_MANAGER.md

@@ -73,6 +73,15 @@ model's lifetime it may be transformed in various ways, such as
 changing its precision or converting it from a .safetensors to a
 diffusers model.
 
+`ModelType`, `ModelFormat` and `BaseModelType` are string enums that
+are defined in `invokeai.backend.model_manager.config`. They are also
+imported by, and can be reexported from,
+`invokeai.app.services.model_manager.model_records`:
+
+```
+from invokeai.app.services.model_records import ModelType, ModelFormat, BaseModelType
+```
+
 The `path` field can be absolute or relative. If relative, it is taken
 to be relative to the `models_dir` setting in the user's
 `invokeai.yaml` file.
@@ -388,25 +397,26 @@ In the event you wish to create a new installer, you may use the
 following initialization pattern:
 
 ```
-from invokeai.app.services.config import get_config
+from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.model_records import ModelRecordServiceSQL
 from invokeai.app.services.model_install import ModelInstallService
 from invokeai.app.services.download import DownloadQueueService
-from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+from invokeai.app.services.shared.sqlite import SqliteDatabase
 from invokeai.backend.util.logging import InvokeAILogger
 
-config = get_config()
+config = InvokeAIAppConfig.get_config()
+config.parse_args()
 
 logger = InvokeAILogger.get_logger(config=config)
-db = SqliteDatabase(config.db_path, logger)
+db = SqliteDatabase(config, logger)
 record_store = ModelRecordServiceSQL(db)
 queue = DownloadQueueService()
 queue.start()
 
-installer = ModelInstallService(app_config=config,
+installer = ModelInstallService(app_config=config, 
                                 record_store=record_store,
-                                download_queue=queue
-                                )
+              download_queue=queue
+           )
 installer.start()
 ```
 
@@ -1357,20 +1367,12 @@ the in-memory loaded model:
 | `model`        | AnyModel               | The instantiated model (details below) |
 | `locker`       | ModelLockerBase        | A context manager that mediates the movement of the model into VRAM |
 
-### get_model_by_key(key, [submodel]) -> LoadedModel
-
-The `get_model_by_key()` method will retrieve the model using its
-unique database key. For example:
-
-loaded_model = loader.get_model_by_key('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
-
-`get_model_by_key()` may raise any of the following exceptions:
-
-* `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
-
-### Using the Loaded Model in Inference
+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.
 
 `LoadedModel` acts as a context manager. The context loads the model
 into the execution device (e.g. VRAM on CUDA systems), locks the model
@@ -1378,33 +1380,17 @@ in the execution device for the duration of the context, and returns
 the model. Use it like this:
 
 ```
-loaded_model_= loader.get_model_by_key('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
-with loaded_model as vae:
+model_info = loader.get_model_by_key('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
+with model_info as vae:
  image = vae.decode(latents)[0]
 ```
 
-The object returned by the LoadedModel context manager is an
-`AnyModel`, which is a Union of `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.
-
-In addition, you may call `LoadedModel.model_on_device()`, a context
-manager that returns a tuple of the model's state dict in CPU and the
-model itself in VRAM. It is used to optimize the LoRA patching and
-unpatching process:
-
-```
-loaded_model_= loader.get_model_by_key('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
-with loaded_model.model_on_device() as (state_dict, vae):
- image = vae.decode(latents)[0]
-```
-
-Since not all models have state dicts, the `state_dict` return value
-can be None.
-
+`get_model_by_key()` may raise any of the following exceptions:
 
+* `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
+  
 ### Emitting model loading events
 
 When the `context` argument is passed to `load_model_*()`, it will
@@ -1592,59 +1578,3 @@ 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`.
-
-## Invocation Context Model Manager API
-
-Within invocations, the following methods are available from the
-`InvocationContext` object:
-
-### context.download_and_cache_model(source) -> Path
-
-This method accepts a `source` of a remote model, downloads and caches
-it locally, and then returns a Path to the local model. The source can
-be a direct download URL or a HuggingFace repo_id.
-
-In the case of HuggingFace repo_id, the following variants are
-recognized:
-
-* stabilityai/stable-diffusion-v4           -- default model
-* stabilityai/stable-diffusion-v4:fp16      -- fp16 variant
-* stabilityai/stable-diffusion-v4:fp16:vae  -- the fp16 vae subfolder
-* stabilityai/stable-diffusion-v4:onnx:vae  -- the onnx variant vae subfolder
-
-You can also point at an arbitrary individual file within a repo_id
-directory using this syntax:
-
-* stabilityai/stable-diffusion-v4::/checkpoints/sd4.safetensors
-
-### context.load_local_model(model_path, [loader]) -> LoadedModel
-
-This method loads a local model from the indicated path, returning a
-`LoadedModel`. The optional loader is a Callable that accepts a Path
-to the object, and returns a `AnyModel` object. If no loader is
-provided, then the method will use `torch.load()` for a .ckpt or .bin
-checkpoint file, `safetensors.torch.load_file()` for a safetensors
-checkpoint file, or `cls.from_pretrained()` for a directory that looks
-like a diffusers directory.
-
-### context.load_remote_model(source, [loader]) -> LoadedModel
-
-This method accepts a `source` of a remote model, downloads and caches
-it locally, loads it, and returns a `LoadedModel`. The source can be a
-direct download URL or a HuggingFace repo_id.
-
-In the case of HuggingFace repo_id, the following variants are
-recognized:
-
-* stabilityai/stable-diffusion-v4           -- default model
-* stabilityai/stable-diffusion-v4:fp16      -- fp16 variant
-* stabilityai/stable-diffusion-v4:fp16:vae  -- the fp16 vae subfolder
-* stabilityai/stable-diffusion-v4:onnx:vae  -- the onnx variant vae subfolder
-
-You can also point at an arbitrary individual file within a repo_id
-directory using this syntax:
-
-* stabilityai/stable-diffusion-v4::/checkpoints/sd4.safetensors
-
-
-

docs/contributing/frontend/WORKFLOWS.md

@@ -117,13 +117,13 @@ Stateless fields do not store their value in the node, so their field instances
 
 "Custom" fields will always be treated as stateless fields.
 
-##### Single and Collection Fields
+##### Collection and Scalar Fields
 
-Field types have a name and cardinality property which may identify it as a **SINGLE**, **COLLECTION** or **SINGLE_OR_COLLECTION** field.
+Field types have a name and two flags which may identify it as a **collection** or **collection or scalar** field.
 
-- If a field is annotated in python as a singular value or class, its field type is parsed as a **SINGLE** type (e.g. `int`, `ImageField`, `str`).
-- If a field is annotated in python as a list, its field type is parsed as a **COLLECTION** type (e.g. `list[int]`).
-- If it is annotated as a union of a type and list, the type will be parsed as a **SINGLE_OR_COLLECTION** type (e.g. `Union[int, list[int]]`). Fields may not be unions of different types (e.g. `Union[int, list[str]]` and `Union[int, str]` are not allowed).
+If a field is annotated in python as a list, its field type is parsed and flagged as a **collection** type (e.g. `list[int]`).
+
+If it is annotated as a union of a type and list, the type will be flagged as a **collection or scalar** type (e.g. `Union[int, list[int]]`). Fields may not be unions of different types (e.g. `Union[int, list[str]]` and `Union[int, str]` are not allowed).
 
 ## Implementation
 
@@ -173,7 +173,8 @@ Field types are represented as structured objects:
 ```ts
 type FieldType = {
   name: string;
-  cardinality: 'SINGLE' | 'COLLECTION' | 'SINGLE_OR_COLLECTION';
+  isCollection: boolean;
+  isCollectionOrScalar: boolean;
 };
 ```
 
@@ -185,7 +186,7 @@ There are 4 general cases for field type parsing.
 
 When a field is annotated as a primitive values (e.g. `int`, `str`, `float`), the field type parsing is fairly straightforward. The field is represented by a simple OpenAPI **schema object**, which has a `type` property.
 
-We create a field type name from this `type` string (e.g. `string` -> `StringField`). The cardinality is `"SINGLE"`.
+We create a field type name from this `type` string (e.g. `string` -> `StringField`).
 
 ##### Complex Types
 
@@ -199,13 +200,13 @@ We need to **dereference** the schema to pull these out. Dereferencing may requi
 
 When a field is annotated as a list of a single type, the schema object has an `items` property. They may be a schema object or reference object and must be parsed to determine the item type.
 
-We use the item type for field type name. The cardinality is `"COLLECTION"`.
+We use the item type for field type name, adding `isCollection: true` to the field type.
 
-##### Single or Collection Types
+##### Collection or Scalar Types
 
 When a field is annotated as a union of a type and list of that type, the schema object has an `anyOf` property, which holds a list of valid types for the union.
 
-After verifying that the union has two members (a type and list of the same type), we use the type for field type name, with cardinality `"SINGLE_OR_COLLECTION"`.
+After verifying that the union has two members (a type and list of the same type), we use the type for field type name, adding `isCollectionOrScalar: true` to the field type.
 
 ##### Optional Fields
 

docs/features/CONFIGURATION.md

@@ -51,11 +51,13 @@ The settings in this file will override the defaults. You only need
 to change this file if the default for a particular setting doesn't
 work for you.
 
-You'll find an example file next to `invokeai.yaml` that shows the default values.
-
 Some settings, like [Model Marketplace API Keys], require the YAML
 to be formatted correctly. Here is a [basic guide to YAML files].
 
+You can fix a broken `invokeai.yaml` by deleting it and running the
+configuration script again -- option [6] in the launcher, "Re-run the
+configure script".
+
 #### Custom Config File Location
 
 You can use any config file with the `--config` CLI arg. Pass in the path to the `invokeai.yaml` file you want to use.

docs/features/CONTROLNET.md

@@ -165,7 +165,7 @@ Additionally, each  section can be expanded with the "Show Advanced" button in o
 There are several ways to install IP-Adapter models with an existing InvokeAI installation:
 
 1. Through the command line interface launched from the invoke.sh / invoke.bat scripts, option [4] to download models.
-2. Through the Model Manager UI with models from the *Tools* section of [models.invoke.ai](https://models.invoke.ai). To do this, copy the repo ID from the desired model page, and paste it in the Add Model field of the model manager. **Note** Both the IP-Adapter and the Image Encoder must be installed for IP-Adapter to work. For example, the [SD 1.5 IP-Adapter](https://models.invoke.ai/InvokeAI/ip_adapter_plus_sd15) and [SD1.5 Image Encoder](https://models.invoke.ai/InvokeAI/ip_adapter_sd_image_encoder) must be installed to use IP-Adapter with SD1.5 based models.  
+2. Through the Model Manager UI with models from the *Tools* section of [www.models.invoke.ai](https://www.models.invoke.ai). To do this, copy the repo ID from the desired model page, and paste it in the Add Model field of the model manager. **Note** Both the IP-Adapter and the Image Encoder must be installed for IP-Adapter to work. For example, the [SD 1.5 IP-Adapter](https://models.invoke.ai/InvokeAI/ip_adapter_plus_sd15) and [SD1.5 Image Encoder](https://models.invoke.ai/InvokeAI/ip_adapter_sd_image_encoder) must be installed to use IP-Adapter with SD1.5 based models.  
 3. **Advanced -- Not recommended ** Manually downloading the IP-Adapter and Image Encoder files - Image Encoder folders shouid be placed in the `models\any\clip_vision` folders. IP Adapter Model folders should be placed in the relevant `ip-adapter` folder of relevant base model folder of Invoke root directory. For example, for the SDXL IP-Adapter, files should be added to the `model/sdxl/ip_adapter/` folder. 
 
 #### Using IP-Adapter

docs/features/TRAINING.md

@@ -4,6 +4,278 @@ title: Training
 
 # :material-file-document: Training
 
-Invoke Training has moved to its own repository, with a dedicated UI for accessing common scripts like Textual Inversion and LoRA training.
+# Textual Inversion Training
+## **Personalizing Text-to-Image Generation**
 
-You can find more by visiting the repo at https://github.com/invoke-ai/invoke-training
+You may personalize the generated images to provide your own styles or objects
+by training a new LDM checkpoint and introducing a new vocabulary to the fixed
+model as a (.pt) embeddings file. Alternatively, you may use or train
+HuggingFace Concepts embeddings files (.bin) from
+<https://huggingface.co/sd-concepts-library> and its associated
+notebooks.
+
+## **Hardware and Software Requirements**
+
+You will need a GPU to perform training in a reasonable length of
+time, and at least 12 GB of VRAM. We recommend using the [`xformers`
+library](../installation/070_INSTALL_XFORMERS.md) to accelerate the
+training process further. During training, about ~8 GB is temporarily
+needed in order to store intermediate models, checkpoints and logs.
+
+## **Preparing for Training**
+
+To train, prepare a folder that contains 3-5 images that illustrate
+the object or concept. It is good to provide a variety of examples or
+poses to avoid overtraining the system. Format these images as PNG
+(preferred) or JPG. You do not need to resize or crop the images in
+advance, but for more control you may wish to do so.
+
+Place the training images in a directory on the machine InvokeAI runs
+on. We recommend placing them in a subdirectory of the
+`text-inversion-training-data` folder located in the InvokeAI root
+directory, ordinarily `~/invokeai` (Linux/Mac), or
+`C:\Users\your_name\invokeai` (Windows). For example, to create an
+embedding for the "psychedelic" style, you'd place the training images
+into the directory
+`~invokeai/text-inversion-training-data/psychedelic`.
+
+## **Launching Training Using the Console Front End**
+
+InvokeAI 2.3 and higher comes with a text console-based training front
+end. From within the `invoke.sh`/`invoke.bat` Invoke launcher script,
+start training tool selecting choice (3):
+
+```sh
+1 "Generate images with a browser-based interface"
+2 "Explore InvokeAI nodes using a command-line interface"
+3 "Textual inversion training"
+4 "Merge models (diffusers type only)"
+5 "Download and install models"
+6 "Change InvokeAI startup options"
+7 "Re-run the configure script to fix a broken install or to complete a major upgrade"
+8 "Open the developer console"
+9 "Update InvokeAI"
+```
+
+Alternatively, you can select option (8) or from the command line, with the InvokeAI virtual environment active,
+you can then launch the front end with the command `invokeai-ti --gui`.
+
+This will launch a text-based front end that will look like this:
+
+<figure markdown>
+![ti-frontend](../assets/textual-inversion/ti-frontend.png)
+</figure>
+
+The interface is keyboard-based. Move from field to field using
+control-N (^N) to move to the next field and control-P (^P) to the
+previous one. <Tab> and <shift-TAB> work as well. Once a field is
+active, use the cursor keys. In a checkbox group, use the up and down
+cursor keys to move from choice to choice, and <space> to select a
+choice. In a scrollbar, use the left and right cursor keys to increase
+and decrease the value of the scroll. In textfields, type the desired
+values.
+
+The number of parameters may look intimidating, but in most cases the
+predefined defaults work fine. The red circled fields in the above
+illustration are the ones you will adjust most frequently.
+
+### Model Name
+
+This will list all the diffusers models that are currently
+installed. Select the one you wish to use as the basis for your
+embedding. Be aware that if you use a SD-1.X-based model for your
+training, you will only be able to use this embedding with other
+SD-1.X-based models. Similarly, if you train on SD-2.X, you will only
+be able to use the embeddings with models based on SD-2.X.
+
+### Trigger Term
+
+This is the prompt term you will use to trigger the embedding. Type a
+single word or phrase you wish to use as the trigger, example
+"psychedelic" (without angle brackets). Within InvokeAI, you will then
+be able to activate the trigger using the syntax `<psychedelic>`.
+
+### Initializer
+
+This is a single character that is used internally during the training
+process as a placeholder for the trigger term. It defaults to "*" and
+can usually be left alone.
+
+### Resume from last saved checkpoint
+
+As training proceeds, textual inversion will write a series of
+intermediate files that can be used to resume training from where it
+was left off in the case of an interruption. This checkbox will be
+automatically selected if you provide a previously used trigger term
+and at least one checkpoint file is found on disk.
+
+Note that as of 20 January 2023, resume does not seem to be working
+properly due to an issue with the upstream code.
+
+### Data Training Directory
+
+This is the location of the images to be used for training. When you
+select a trigger term like "my-trigger", the frontend will prepopulate
+this field with `~/invokeai/text-inversion-training-data/my-trigger`,
+but you can change the path to wherever you want.
+
+### Output Destination Directory
+
+This is the location of the logs, checkpoint files, and embedding
+files created during training. When you select a trigger term like
+"my-trigger", the frontend will prepopulate this field with
+`~/invokeai/text-inversion-output/my-trigger`, but you can change the
+path to wherever you want.
+
+### Image resolution
+
+The images in the training directory will be automatically scaled to
+the value you use here. For best results, you will want to use the
+same default resolution of the underlying model (512 pixels for
+SD-1.5, 768 for the larger version of SD-2.1).
+
+### Center crop images
+
+If this is selected, your images will be center cropped to make them
+square before resizing them to the desired resolution. Center cropping
+can indiscriminately cut off the top of subjects' heads for portrait
+aspect images, so if you have images like this, you may wish to use a
+photoeditor to manually crop them to a square aspect ratio.
+
+### Mixed precision
+
+Select the floating point precision for the embedding. "no" will
+result in a full 32-bit precision, "fp16" will provide 16-bit
+precision, and "bf16" will provide mixed precision (only available
+when XFormers is used).
+
+### Max training steps
+
+How many steps the training will take before the model converges. Most
+training sets will converge with 2000-3000 steps.
+
+### Batch size
+
+This adjusts how many training images are processed simultaneously in
+each step. Higher values will cause the training process to run more
+quickly, but use more memory. The default size will run with GPUs with
+as little as 12 GB.
+
+### Learning rate
+
+The rate at which the system adjusts its internal weights during
+training. Higher values risk overtraining (getting the same image each
+time), and lower values will take more steps to train a good
+model. The default of 0.0005 is conservative; you may wish to increase
+it to 0.005 to speed up training.
+
+### Scale learning rate by number of GPUs, steps and batch size
+
+If this is selected (the default) the system will adjust the provided
+learning rate to improve performance.
+
+### Use xformers acceleration
+
+This will activate XFormers memory-efficient attention. You need to
+have XFormers installed for this to have an effect.
+
+### Learning rate scheduler
+
+This adjusts how the learning rate changes over the course of
+training. The default "constant" means to use a constant learning rate
+for the entire training session. The other values scale the learning
+rate according to various formulas.
+
+Only "constant" is supported by the XFormers library.
+
+### Gradient accumulation steps
+
+This is a parameter that allows you to use bigger batch sizes than
+your GPU's VRAM would ordinarily accommodate, at the cost of some
+performance.
+
+### Warmup steps
+
+If "constant_with_warmup" is selected in the learning rate scheduler,
+then this provides the number of warmup steps. Warmup steps have a
+very low learning rate, and are one way of preventing early
+overtraining.
+
+## The training run
+
+Start the training run by advancing to the OK button (bottom right)
+and pressing <enter>. A series of progress messages will be displayed
+as the training process proceeds. This may take an hour or two,
+depending on settings and the speed of your system. Various log and
+checkpoint files will be written into the output directory (ordinarily
+`~/invokeai/text-inversion-output/my-model/`)
+
+At the end of successful training, the system will copy the file
+`learned_embeds.bin` into the InvokeAI root directory's `embeddings`
+directory, using a subdirectory named after the trigger token. For
+example, if the trigger token was `psychedelic`, then look for the
+embeddings file in
+`~/invokeai/embeddings/psychedelic/learned_embeds.bin`
+
+You may now launch InvokeAI and try out a prompt that uses the trigger
+term. For example `a plate of banana sushi in <psychedelic> style`.
+
+## **Training with the Command-Line Script**
+
+Training can also be done using a traditional command-line script. It
+can be launched from within the "developer's console", or from the
+command line after activating InvokeAI's virtual environment.
+
+It accepts a large number of arguments, which can be summarized by
+passing the `--help` argument:
+
+```sh
+invokeai-ti --help
+```
+
+Typical usage is shown here:
+```sh
+invokeai-ti \
+       --model=stable-diffusion-1.5 \
+       --resolution=512 \
+       --learnable_property=style \
+       --initializer_token='*' \
+       --placeholder_token='<psychedelic>' \
+       --train_data_dir=/home/lstein/invokeai/training-data/psychedelic \
+       --output_dir=/home/lstein/invokeai/text-inversion-training/psychedelic \
+       --scale_lr \
+       --train_batch_size=8 \
+       --gradient_accumulation_steps=4 \
+       --max_train_steps=3000 \
+       --learning_rate=0.0005 \
+       --resume_from_checkpoint=latest \
+       --lr_scheduler=constant \
+       --mixed_precision=fp16 \
+       --only_save_embeds
+```
+
+## Troubleshooting
+
+### `Cannot load embedding for <trigger>. It was trained on a model with token dimension 1024, but the current model has token dimension 768`
+
+Messages like this indicate you trained the embedding on a different base model than the currently selected one.
+
+For example, in the error above, the training was done on SD2.1 (768x768) but it was used on SD1.5 (512x512).
+
+## Reading
+
+For more information on textual inversion, please see the following
+resources:
+
+* The [textual inversion repository](https://github.com/rinongal/textual_inversion) and
+  associated paper for details and limitations.
+* [HuggingFace's textual inversion training
+  page](https://huggingface.co/docs/diffusers/training/text_inversion)
+* [HuggingFace example script
+  documentation](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion)
+  (Note that this script is similar to, but not identical, to
+  `textual_inversion`, but produces embed files that are completely compatible.
+
+---
+
+copyright (c) 2023, Lincoln Stein and the InvokeAI Development Team

docs/help/FAQ.md

@@ -154,18 +154,6 @@ This is caused by an invalid setting in the `invokeai.yaml` configuration file.
 
 Check the [configuration docs] for more detail about the settings and how to specify them.
 
-## `ModuleNotFoundError: No module named 'controlnet_aux'`
-
-`controlnet_aux` is a dependency of Invoke and appears to have been packaged or distributed strangely. Sometimes, it doesn't install correctly. This is outside our control.
-
-If you encounter this error, the solution is to remove the package from the `pip` cache and re-run the Invoke installer so a fresh, working version of `controlnet_aux` can be downloaded and installed:
-
-- Run the Invoke launcher
-- Choose the developer console option
-- Run this command: `pip cache remove controlnet_aux`
-- Close the terminal window
-- Download and run the [installer](https://github.com/invoke-ai/InvokeAI/releases/latest), selecting your current install location
-
 ## Out of Memory Issues
 
 The models are large, VRAM is expensive, and you may find yourself

docs/help/diffusion.md

@@ -20,7 +20,7 @@ When you generate an image using text-to-image, multiple steps occur in latent s
 4. The VAE decodes the final latent image from latent space into image space.
 
 Image-to-image is a similar process, with only step 1 being different:
-1. The input image is encoded from image space into latent space by the VAE. Noise is then added to the input latent image. Denoising Strength dictates how many noise steps are added, and the amount of noise added at each step. A Denoising Strength of 0 means there are 0 steps and no noise added, resulting in an unchanged image, while a Denoising Strength of 1 results in the image being completely replaced with noise and a full set of denoising steps are performance. The process is then the same as steps 2-4 in the text-to-image process. 
+1. The input image is encoded from image space into latent space by the VAE. Noise is then added to the input latent image. Denoising Strength dictates how may noise steps are added, and the amount of noise added at each step. A Denoising Strength of 0 means there are 0 steps and no noise added, resulting in an unchanged image, while a Denoising Strength of 1 results in the image being completely replaced with noise and a full set of denoising steps are performance. The process is then the same as steps 2-4 in the text-to-image process. 
 
 Furthermore, a model provides the CLIP prompt tokenizer, the VAE, and a U-Net (where noise prediction occurs given a prompt and initial noise tensor).
 

docs/installation/010_INSTALL_AUTOMATED.md

@@ -1,10 +1,8 @@
-# Automatic Install & Updates
+# Automatic Install
 
-**The same packaged installer file can be used for both new installs and updates.**
-Using the installer for updates will leave everything you've added since installation, and just update the core libraries used to run Invoke.
-Simply use the same path you installed to originally.
+The installer is used for both new installs and updates.
 
-Both release and pre-release versions can be installed using the installer. It also supports install through a wheel if needed.
+Both release and pre-release versions can be installed using it. It also supports install a wheel if needed.
 
 Be sure to review the [installation requirements] and ensure your system has everything it needs to install Invoke.
 
@@ -98,7 +96,7 @@ Updating is exactly the same as installing - download the latest installer, choo
 
 If you have installation issues, please review the [FAQ]. You can also [create an issue] or ask for help on [discord].
 
-[installation requirements]: INSTALL_REQUIREMENTS.md
+[installation requirements]: INSTALLATION.md#installation-requirements
 [FAQ]: ../help/FAQ.md
 [install some models]: 050_INSTALLING_MODELS.md
 [configuration docs]: ../features/CONFIGURATION.md

docs/installation/020_INSTALL_MANUAL.md

@@ -10,7 +10,7 @@ InvokeAI is distributed as a python package on PyPI, installable with `pip`. The
 
 ### Requirements
 
-Before you start, go through the [installation requirements](./INSTALL_REQUIREMENTS.md).
+Before you start, go through the [installation requirements].
 
 ### Installation Walkthrough
 
@@ -79,7 +79,7 @@ Before you start, go through the [installation requirements](./INSTALL_REQUIREME
 
 1. Install the InvokeAI Package. The base command is `pip install InvokeAI --use-pep517`, but you may need to change this depending on your system and the desired features.
 
-    - You may need to provide an [extra index URL](https://pip.pypa.io/en/stable/cli/pip_install/#cmdoption-extra-index-url). Select your platform configuration using [this tool on the PyTorch website](https://pytorch.org/get-started/locally/). Copy the `--extra-index-url` string from this and append it to your install command.
+    - You may need to provide an [extra index URL]. Select your platform configuration using [this tool on the PyTorch website]. Copy the `--extra-index-url` string from this and append it to your install command.
 
         !!! example "Install with an extra index URL"
 
@@ -116,4 +116,4 @@ Before you start, go through the [installation requirements](./INSTALL_REQUIREME
 
     !!! warning
 
-        If the virtual environment is _not_ inside the root directory, then you _must_ specify the path to the root directory with `--root \path\to\invokeai` or the `INVOKEAI_ROOT` environment variable.
+        If the virtual environment is _not_ inside the root directory, then you _must_ specify the path to the root directory with `--root_dir \path\to\invokeai` or the `INVOKEAI_ROOT` environment variable.

docs/installation/INSTALLATION.md

@@ -1,4 +1,4 @@
-# Installation and Updating Overview
+# Installation Overview
 
 Before installing, review the [installation requirements] to ensure your system is set up properly.
 
@@ -6,21 +6,14 @@ See the [FAQ] for frequently-encountered installation issues.
 
 If you need more help, join our [discord] or [create an issue].
 
-<h2>Automatic Install & Updates </h2>
+<h2>Automatic Install</h2>
 
 ✅ The automatic install is the best way to run InvokeAI. Check out the [installation guide] to get started.
 
-⬆️ The same installer is also the best way to update InvokeAI - Simply rerun it for the same folder you installed to.
-
-The installation process simply manages installation for the core libraries & application dependencies that run Invoke.
-Any models, images, or other assets in the Invoke root folder won't be affected by the installation process.
-
 <h2>Manual Install</h2>
 
 If you are familiar with python and want more control over the packages that are installed, you can [install InvokeAI manually via PyPI].
 
-Updates are managed by reinstalling the latest version through PyPi.
-
 <h2>Developer Install</h2>
 
 If you want to contribute to InvokeAI, consult the [developer install guide].

docs/installation/INSTALL_REQUIREMENTS.md

@@ -37,13 +37,13 @@ Invoke runs best with a dedicated GPU, but will fall back to running on CPU, alb
     === "Nvidia"
 
         ```
-        Any GPU with at least 8GB VRAM.
+        Any GPU with at least 8GB VRAM. Linux only.
         ```
 
     === "AMD"
 
         ```
-        Any GPU with at least 16GB VRAM. Linux only.
+        Any GPU with at least 16GB VRAM.
         ```
 
     === "Mac"

installer/templates/invoke.bat.in

@@ -10,7 +10,8 @@ set INVOKEAI_ROOT=.
 echo Desired action:
 echo 1. Generate images with the browser-based interface
 echo 2. Open the developer console
-echo 3. Command-line help
+echo 3. Run the InvokeAI image database maintenance script
+echo 4. Command-line help
 echo Q - Quit
 echo.
 echo To update, download and run the installer from https://github.com/invoke-ai/InvokeAI/releases/latest.
@@ -33,6 +34,9 @@ IF /I "%choice%" == "1" (
     echo *** Type `exit` to quit this shell and deactivate the Python virtual environment ***
     call cmd /k
 ) ELSE IF /I "%choice%" == "3" (
+   echo Running the db maintenance script...
+   python .venv\Scripts\invokeai-db-maintenance.exe
+) ELSE IF /I "%choice%" == "4" (
     echo Displaying command line help...
     python .venv\Scripts\invokeai-web.exe --help %*
     pause

installer/templates/invoke.sh.in

@@ -47,6 +47,11 @@ do_choice() {
         bash --init-file "$file_name"
         ;;
     3)
+        clear
+        printf "Running the db maintenance script\n"
+        invokeai-db-maintenance --root ${INVOKEAI_ROOT}
+        ;;
+    4)
         clear
         printf "Command-line help\n"
         invokeai-web --help
@@ -66,7 +71,8 @@ do_line_input() {
     printf "What would you like to do?\n"
     printf "1: Generate images using the browser-based interface\n"
     printf "2: Open the developer console\n"
-    printf "3: Command-line help\n"
+    printf "3: Run the InvokeAI image database maintenance script\n"
+    printf "4: Command-line help\n"
     printf "Q: Quit\n\n"
     printf "To update, download and run the installer from https://github.com/invoke-ai/InvokeAI/releases/latest.\n\n"
     read -p "Please enter 1-4, Q: [1] " yn

invokeai/app/api/dependencies.py

@@ -18,7 +18,6 @@ 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.events.events_fastapievents import FastAPIEventService
 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
@@ -30,10 +29,11 @@ from ..services.model_images.model_images_default import ModelImageFileStorageDi
 from ..services.model_manager.model_manager_default import ModelManagerService
 from ..services.model_records import ModelRecordServiceSQL
 from ..services.names.names_default import SimpleNameService
-from ..services.session_processor.session_processor_default import DefaultSessionProcessor, DefaultSessionRunner
+from ..services.session_processor.session_processor_default import DefaultSessionProcessor
 from ..services.session_queue.session_queue_sqlite import SqliteSessionQueue
 from ..services.urls.urls_default import LocalUrlService
 from ..services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage
+from .events import FastAPIEventService
 
 
 # TODO: is there a better way to achieve this?
@@ -93,7 +93,7 @@ class ApiDependencies:
         conditioning = ObjectSerializerForwardCache(
             ObjectSerializerDisk[ConditioningFieldData](output_folder / "conditioning", ephemeral=True)
         )
-        download_queue_service = DownloadQueueService(app_config=configuration, event_bus=events)
+        download_queue_service = DownloadQueueService(event_bus=events)
         model_images_service = ModelImageFileStorageDisk(model_images_folder / "model_images")
         model_manager = ModelManagerService.build_model_manager(
             app_config=configuration,
@@ -103,7 +103,7 @@ class ApiDependencies:
         )
         names = SimpleNameService()
         performance_statistics = InvocationStatsService()
-        session_processor = DefaultSessionProcessor(session_runner=DefaultSessionRunner())
+        session_processor = DefaultSessionProcessor()
         session_queue = SqliteSessionQueue(db=db)
         urls = LocalUrlService()
         workflow_records = SqliteWorkflowRecordsStorage(db=db)

invokeai/app/api/events.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+import asyncio
+import threading
+from queue import Empty, Queue
+from typing import Any
+
+from fastapi_events.dispatcher import dispatch
+
+from ..services.events.events_base import EventServiceBase
+
+
+class FastAPIEventService(EventServiceBase):
+    event_handler_id: int
+    __queue: Queue
+    __stop_event: threading.Event
+
+    def __init__(self, event_handler_id: int) -> None:
+        self.event_handler_id = event_handler_id
+        self.__queue = Queue()
+        self.__stop_event = threading.Event()
+        asyncio.create_task(self.__dispatch_from_queue(stop_event=self.__stop_event))
+
+        super().__init__()
+
+    def stop(self, *args, **kwargs):
+        self.__stop_event.set()
+        self.__queue.put(None)
+
+    def dispatch(self, event_name: str, payload: Any) -> None:
+        self.__queue.put({"event_name": event_name, "payload": payload})
+
+    async def __dispatch_from_queue(self, stop_event: threading.Event):
+        """Get events on from the queue and dispatch them, from the correct thread"""
+        while not stop_event.is_set():
+            try:
+                event = self.__queue.get(block=False)
+                if not event:  # Probably stopping
+                    continue
+
+                dispatch(
+                    event.get("event_name"),
+                    payload=event.get("payload"),
+                    middleware_id=self.event_handler_id,
+                )
+
+            except Empty:
+                await asyncio.sleep(0.1)
+                pass
+
+            except asyncio.CancelledError as e:
+                raise e  # Raise a proper error

invokeai/app/api/routers/app_info.py

@@ -13,6 +13,7 @@ from pydantic import BaseModel, Field
 from invokeai.app.invocations.upscale import ESRGAN_MODELS
 from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
 from invokeai.backend.image_util.infill_methods.patchmatch import PatchMatch
+from invokeai.backend.image_util.safety_checker import SafetyChecker
 from invokeai.backend.util.logging import logging
 from invokeai.version import __version__
 
@@ -108,7 +109,9 @@ async def get_config() -> AppConfig:
         upscaling_models.append(str(Path(model).stem))
     upscaler = Upscaler(upscaling_method="esrgan", upscaling_models=upscaling_models)
 
-    nsfw_methods = ["nsfw_checker"]
+    nsfw_methods = []
+    if SafetyChecker.safety_checker_available():
+        nsfw_methods.append("nsfw_checker")
 
     watermarking_methods = ["invisible_watermark"]
 

invokeai/app/api/routers/boards.py

@@ -118,13 +118,15 @@ async def list_boards(
     all: Optional[bool] = Query(default=None, description="Whether to list all boards"),
     offset: Optional[int] = Query(default=None, description="The page offset"),
     limit: Optional[int] = Query(default=None, description="The number of boards per page"),
-    include_archived: bool = Query(default=False, description="Whether or not to include archived boards in list"),
 ) -> Union[OffsetPaginatedResults[BoardDTO], list[BoardDTO]]:
     """Gets a list of boards"""
     if all:
-        return ApiDependencies.invoker.services.boards.get_all(include_archived)
+        return ApiDependencies.invoker.services.boards.get_all()
     elif offset is not None and limit is not None:
-        return ApiDependencies.invoker.services.boards.get_many(offset, limit, include_archived)
+        return ApiDependencies.invoker.services.boards.get_many(
+            offset,
+            limit,
+        )
     else:
         raise HTTPException(
             status_code=400,

invokeai/app/api/routers/images.py

@@ -6,17 +6,13 @@ from fastapi import BackgroundTasks, Body, HTTPException, Path, Query, Request,
 from fastapi.responses import FileResponse
 from fastapi.routing import APIRouter
 from PIL import Image
-from pydantic import BaseModel, Field, JsonValue
+from pydantic import BaseModel, Field, ValidationError
 
-from invokeai.app.invocations.fields import MetadataField
-from invokeai.app.services.image_records.image_records_common import (
-    ImageCategory,
-    ImageRecordChanges,
-    ResourceOrigin,
-)
+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
-from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID, WorkflowWithoutIDValidator
 
 from ..dependencies import ApiDependencies
 
@@ -46,17 +42,13 @@ async def upload_image(
     board_id: Optional[str] = Query(default=None, description="The board to add this image to, if any"),
     session_id: Optional[str] = Query(default=None, description="The session ID associated with this upload, if any"),
     crop_visible: Optional[bool] = Query(default=False, description="Whether to crop the image"),
-    metadata: Optional[JsonValue] = Body(
-        default=None, description="The metadata to associate with the image", embed=True
-    ),
 ) -> ImageDTO:
     """Uploads an image"""
     if not file.content_type or not file.content_type.startswith("image"):
         raise HTTPException(status_code=415, detail="Not an image")
 
-    _metadata = None
-    _workflow = None
-    _graph = None
+    metadata = None
+    workflow = None
 
     contents = await file.read()
     try:
@@ -70,28 +62,22 @@ async def upload_image(
 
     # TODO: retain non-invokeai metadata on upload?
     # attempt to parse metadata from image
-    metadata_raw = metadata if isinstance(metadata, str) else pil_image.info.get("invokeai_metadata", None)
-    if isinstance(metadata_raw, str):
-        _metadata = metadata_raw
-    else:
-        ApiDependencies.invoker.services.logger.debug("Failed to parse metadata for uploaded image")
-        pass
+    metadata_raw = pil_image.info.get("invokeai_metadata", None)
+    if metadata_raw:
+        try:
+            metadata = MetadataFieldValidator.validate_json(metadata_raw)
+        except ValidationError:
+            ApiDependencies.invoker.services.logger.warn("Failed to parse metadata for uploaded image")
+            pass
 
     # attempt to parse workflow from image
     workflow_raw = pil_image.info.get("invokeai_workflow", None)
-    if isinstance(workflow_raw, str):
-        _workflow = workflow_raw
-    else:
-        ApiDependencies.invoker.services.logger.debug("Failed to parse workflow for uploaded image")
-        pass
-
-    # attempt to extract graph from image
-    graph_raw = pil_image.info.get("invokeai_graph", None)
-    if isinstance(graph_raw, str):
-        _graph = graph_raw
-    else:
-        ApiDependencies.invoker.services.logger.debug("Failed to parse graph for uploaded image")
-        pass
+    if workflow_raw is not None:
+        try:
+            workflow = WorkflowWithoutIDValidator.validate_json(workflow_raw)
+        except ValidationError:
+            ApiDependencies.invoker.services.logger.warn("Failed to parse metadata for uploaded image")
+            pass
 
     try:
         image_dto = ApiDependencies.invoker.services.images.create(
@@ -100,9 +86,8 @@ async def upload_image(
             image_category=image_category,
             session_id=session_id,
             board_id=board_id,
-            metadata=_metadata,
-            workflow=_workflow,
-            graph=_graph,
+            metadata=metadata,
+            workflow=workflow,
             is_intermediate=is_intermediate,
         )
 
@@ -200,21 +185,14 @@ async def get_image_metadata(
         raise HTTPException(status_code=404)
 
 
-class WorkflowAndGraphResponse(BaseModel):
-    workflow: Optional[str] = Field(description="The workflow used to generate the image, as stringified JSON")
-    graph: Optional[str] = Field(description="The graph used to generate the image, as stringified JSON")
-
-
 @images_router.get(
-    "/i/{image_name}/workflow", operation_id="get_image_workflow", response_model=WorkflowAndGraphResponse
+    "/i/{image_name}/workflow", operation_id="get_image_workflow", response_model=Optional[WorkflowWithoutID]
 )
 async def get_image_workflow(
     image_name: str = Path(description="The name of image whose workflow to get"),
-) -> WorkflowAndGraphResponse:
+) -> Optional[WorkflowWithoutID]:
     try:
-        workflow = ApiDependencies.invoker.services.images.get_workflow(image_name)
-        graph = ApiDependencies.invoker.services.images.get_graph(image_name)
-        return WorkflowAndGraphResponse(workflow=workflow, graph=graph)
+        return ApiDependencies.invoker.services.images.get_workflow(image_name)
     except Exception:
         raise HTTPException(status_code=404)
 
@@ -321,14 +299,16 @@ async def list_image_dtos(
     ),
     offset: int = Query(default=0, description="The page offset"),
     limit: int = Query(default=10, description="The number of images per page"),
-    order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The order of sort"),
-    starred_first: bool = Query(default=True, description="Whether to sort by starred images first"),
-    search_term: Optional[str] = Query(default=None, description="The term to search for"),
 ) -> OffsetPaginatedResults[ImageDTO]:
     """Gets a list of image DTOs"""
 
     image_dtos = ApiDependencies.invoker.services.images.get_many(
-        offset, limit, starred_first, order_dir, image_origin, categories, is_intermediate, board_id, search_term
+        offset,
+        limit,
+        image_origin,
+        categories,
+        is_intermediate,
+        board_id,
     )
 
     return image_dtos

invokeai/app/api/routers/model_manager.py

@@ -3,22 +3,22 @@
 
 import io
 import pathlib
+import shutil
 import traceback
 from copy import deepcopy
-from tempfile import TemporaryDirectory
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional
 
 from fastapi import Body, Path, Query, Response, UploadFile
-from fastapi.responses import FileResponse, HTMLResponse
+from fastapi.responses import FileResponse
 from fastapi.routing import APIRouter
 from PIL import Image
 from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field
 from starlette.exceptions import HTTPException
 from typing_extensions import Annotated
 
-from invokeai.app.services.model_images.model_images_common import ModelImageFileNotFoundException
-from invokeai.app.services.model_install.model_install_common import ModelInstallJob
+from invokeai.app.services.model_install import ModelInstallJob
 from invokeai.app.services.model_records import (
+    DuplicateModelException,
     InvalidModelException,
     ModelRecordChanges,
     UnknownModelException,
@@ -29,6 +29,7 @@ from invokeai.backend.model_manager.config import (
     MainCheckpointConfig,
     ModelFormat,
     ModelType,
+    SubModelType,
 )
 from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch
 from invokeai.backend.model_manager.metadata.metadata_base import ModelMetadataWithFiles, UnknownMetadataException
@@ -51,13 +52,6 @@ class ModelsList(BaseModel):
     model_config = ConfigDict(use_enum_values=True)
 
 
-def add_cover_image_to_model_config(config: AnyModelConfig, dependencies: Type[ApiDependencies]) -> AnyModelConfig:
-    """Add a cover image URL to a model configuration."""
-    cover_image = dependencies.invoker.services.model_images.get_url(config.key)
-    config.cover_image = cover_image
-    return config
-
-
 ##############################################################################
 # These are example inputs and outputs that are used in places where Swagger
 # is unable to generate a correct example.
@@ -124,7 +118,8 @@ async def list_model_records(
             record_store.search_by_attr(model_type=model_type, model_name=model_name, model_format=model_format)
         )
     for model in found_models:
-        model = add_cover_image_to_model_config(model, ApiDependencies)
+        cover_image = ApiDependencies.invoker.services.model_images.get_url(model.key)
+        model.cover_image = cover_image
     return ModelsList(models=found_models)
 
 
@@ -165,13 +160,28 @@ 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 = ApiDependencies.invoker.services.model_manager.store.get_model(key)
-        return add_cover_image_to_model_config(config, ApiDependencies)
+        config: AnyModelConfig = record_store.get_model(key)
+        cover_image = ApiDependencies.invoker.services.model_images.get_url(key)
+        config.cover_image = cover_image
+        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
+
+
 class FoundModel(BaseModel):
     path: str = Field(description="Path to the model")
     is_installed: bool = Field(description="Whether or not the model is already installed")
@@ -284,15 +294,14 @@ async def update_model_record(
     installer = ApiDependencies.invoker.services.model_manager.install
     try:
         record_store.update_model(key, changes=changes)
-        config = installer.sync_model_path(key)
-        config = add_cover_image_to_model_config(config, ApiDependencies)
+        model_response: AnyModelConfig = installer.sync_model_path(key)
         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 config
+    return model_response
 
 
 @model_manager_router.get(
@@ -488,133 +497,6 @@ async def install_model(
     return result
 
 
-@model_manager_router.get(
-    "/install/huggingface",
-    operation_id="install_hugging_face_model",
-    responses={
-        201: {"description": "The model is being installed"},
-        400: {"description": "Bad request"},
-        409: {"description": "There is already a model corresponding to this path or repo_id"},
-    },
-    status_code=201,
-    response_class=HTMLResponse,
-)
-async def install_hugging_face_model(
-    source: str = Query(description="HuggingFace repo_id to install"),
-) -> HTMLResponse:
-    """Install a Hugging Face model using a string identifier."""
-
-    def generate_html(title: str, heading: str, repo_id: str, is_error: bool, message: str | None = "") -> str:
-        if message:
-            message = f"<p>{message}</p>"
-        title_class = "error" if is_error else "success"
-        return f"""
-            <html>
-
-            <head>
-                <title>{title}</title>
-                <style>
-                    body {{
-                        text-align: center;
-                        background-color: hsl(220 12% 10% / 1);
-                        font-family: Helvetica, sans-serif;
-                        color: hsl(220 12% 86% / 1);
-                    }}
-
-                    .repo-id {{
-                        color: hsl(220 12% 68% / 1);
-                    }}
-
-                    .error {{
-                        color: hsl(0 42% 68% / 1)
-                    }}
-
-                    .message-box {{
-                        display: inline-block;
-                        border-radius: 5px;
-                        background-color: hsl(220 12% 20% / 1);
-                        padding-inline-end: 30px;
-                        padding: 20px;
-                        padding-inline-start: 30px;
-                        padding-inline-end: 30px;
-                    }}
-
-                    .container {{
-                        display: flex;
-                        width: 100%;
-                        height: 100%;
-                        align-items: center;
-                        justify-content: center;
-                    }}
-
-                    a {{
-                        color: inherit
-                    }}
-
-                    a:visited {{
-                        color: inherit
-                    }}
-
-                    a:active {{
-                        color: inherit
-                    }}
-                </style>
-            </head>
-
-            <body style="background-color: hsl(220 12% 10% / 1);">
-                <div class="container">
-                    <div class="message-box">
-                        <h2 class="{title_class}">{heading}</h2>
-                        {message}
-                        <p class="repo-id">Repo ID: {repo_id}</p>
-                    </div>
-                </div>
-            </body>
-
-            </html>
-        """
-
-    try:
-        metadata = HuggingFaceMetadataFetch().from_id(source)
-        assert isinstance(metadata, ModelMetadataWithFiles)
-    except UnknownMetadataException:
-        title = "Unable to Install Model"
-        heading = "No HuggingFace repository found with that repo ID."
-        message = "Ensure the repo ID is correct and try again."
-        return HTMLResponse(content=generate_html(title, heading, source, True, message), status_code=400)
-
-    logger = ApiDependencies.invoker.services.logger
-
-    try:
-        installer = ApiDependencies.invoker.services.model_manager.install
-        if metadata.is_diffusers:
-            installer.heuristic_import(
-                source=source,
-                inplace=False,
-            )
-        elif metadata.ckpt_urls is not None and len(metadata.ckpt_urls) == 1:
-            installer.heuristic_import(
-                source=str(metadata.ckpt_urls[0]),
-                inplace=False,
-            )
-        else:
-            title = "Unable to Install Model"
-            heading = "This HuggingFace repo has multiple models."
-            message = "Please use the Model Manager to install this model."
-            return HTMLResponse(content=generate_html(title, heading, source, True, message), status_code=200)
-
-        title = "Model Install Started"
-        heading = "Your HuggingFace model is installing now."
-        message = "You can close this tab and check the Model Manager for installation progress."
-        return HTMLResponse(content=generate_html(title, heading, source, False, message), status_code=201)
-    except Exception as e:
-        logger.error(str(e))
-        title = "Unable to Install Model"
-        heading = "There was an problem installing this model."
-        message = 'Please use the Model Manager directly to install this model. If the issue persists, ask for help on <a href="https://discord.gg/ZmtBAhwWhy">discord</a>.'
-        return HTMLResponse(content=generate_html(title, heading, source, True, message), status_code=500)
-
-
 @model_manager_router.get(
     "/install",
     operation_id="list_model_installs",
@@ -732,54 +614,48 @@ async def convert_model(
         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.")
 
-    with TemporaryDirectory(dir=ApiDependencies.invoker.services.configuration.models_path) as tmpdir:
-        convert_path = pathlib.Path(tmpdir) / pathlib.Path(model_config.path).stem
-        converted_model = loader.load_model(model_config)
-        # write the converted file to the convert path
-        raw_model = converted_model.model
-        assert hasattr(raw_model, "save_pretrained")
-        raw_model.save_pretrained(convert_path)
-        assert convert_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"
-        changes = ModelRecordChanges(name=model_config.name)
-        store.update_model(key, changes=changes)
-
-        # install the diffusers
-        try:
-            new_key = installer.install_path(
-                convert_path,
-                config={
-                    "name": original_name,
-                    "description": model_config.description,
-                    "hash": model_config.hash,
-                    "source": model_config.source,
-                },
-            )
-        except Exception as e:
-            logger.error(str(e))
-            store.update_model(key, changes=ModelRecordChanges(name=original_name))
-            raise HTTPException(status_code=409, detail=str(e))
-
-    # Update the model image if the model had one
+    # loading the model will convert it into a cached diffusers file
     try:
-        model_image = ApiDependencies.invoker.services.model_images.get(key)
-        ApiDependencies.invoker.services.model_images.save(model_image, new_key)
-        ApiDependencies.invoker.services.model_images.delete(key)
-    except ModelImageFileNotFoundException:
-        pass
+        cc_size = loader.convert_cache.max_size
+        if cc_size == 0:  # temporary set the convert cache to a positive number so that cached model is written
+            loader._convert_cache.max_size = 1.0
+        loader.load_model(model_config, submodel_type=SubModelType.Scheduler)
+    finally:
+        loader._convert_cache.max_size = cc_size
+
+    # 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"
+    changes = ModelRecordChanges(name=model_config.name)
+    store.update_model(key, changes=changes)
+
+    # install the diffusers
+    try:
+        new_key = installer.install_path(
+            cache_path,
+            config={
+                "name": original_name,
+                "description": model_config.description,
+                "hash": model_config.hash,
+                "source": model_config.source,
+            },
+        )
+    except DuplicateModelException as e:
+        logger.error(str(e))
+        raise HTTPException(status_code=409, detail=str(e))
 
     # delete the original safetensors file
     installer.delete(key)
 
-    # delete the temporary directory
-    # shutil.rmtree(cache_path)
+    # delete the cached version
+    shutil.rmtree(cache_path)
 
     # return the config record for the new diffusers directory
-    new_config = store.get_model(new_key)
-    new_config = add_cover_image_to_model_config(new_config, ApiDependencies)
+    new_config: AnyModelConfig = store.get_model(new_key)
     return new_config
 
 

invokeai/app/api/routers/session_queue.py

@@ -203,7 +203,6 @@ async def get_batch_status(
     responses={
         200: {"model": SessionQueueItem},
     },
-    response_model_exclude_none=True,
 )
 async def get_queue_item(
     queue_id: str = Path(description="The queue id to perform this operation on"),

invokeai/app/api/sockets.py

@@ -1,125 +1,66 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
-from typing import Any
-
 from fastapi import FastAPI
-from pydantic import BaseModel
+from fastapi_events.handlers.local import local_handler
+from fastapi_events.typing import Event
 from socketio import ASGIApp, AsyncServer
 
-from invokeai.app.services.events.events_common import (
-    BatchEnqueuedEvent,
-    BulkDownloadCompleteEvent,
-    BulkDownloadErrorEvent,
-    BulkDownloadEventBase,
-    BulkDownloadStartedEvent,
-    DownloadCancelledEvent,
-    DownloadCompleteEvent,
-    DownloadErrorEvent,
-    DownloadEventBase,
-    DownloadProgressEvent,
-    DownloadStartedEvent,
-    FastAPIEvent,
-    InvocationCompleteEvent,
-    InvocationDenoiseProgressEvent,
-    InvocationErrorEvent,
-    InvocationStartedEvent,
-    ModelEventBase,
-    ModelInstallCancelledEvent,
-    ModelInstallCompleteEvent,
-    ModelInstallDownloadProgressEvent,
-    ModelInstallDownloadsCompleteEvent,
-    ModelInstallErrorEvent,
-    ModelInstallStartedEvent,
-    ModelLoadCompleteEvent,
-    ModelLoadStartedEvent,
-    QueueClearedEvent,
-    QueueEventBase,
-    QueueItemStatusChangedEvent,
-    register_events,
-)
-
-
-class QueueSubscriptionEvent(BaseModel):
-    """Event data for subscribing to the socket.io queue room.
-    This is a pydantic model to ensure the data is in the correct format."""
-
-    queue_id: str
-
-
-class BulkDownloadSubscriptionEvent(BaseModel):
-    """Event data for subscribing to the socket.io bulk downloads room.
-    This is a pydantic model to ensure the data is in the correct format."""
-
-    bulk_download_id: str
-
-
-QUEUE_EVENTS = {
-    InvocationStartedEvent,
-    InvocationDenoiseProgressEvent,
-    InvocationCompleteEvent,
-    InvocationErrorEvent,
-    QueueItemStatusChangedEvent,
-    BatchEnqueuedEvent,
-    QueueClearedEvent,
-}
-
-MODEL_EVENTS = {
-    DownloadCancelledEvent,
-    DownloadCompleteEvent,
-    DownloadErrorEvent,
-    DownloadProgressEvent,
-    DownloadStartedEvent,
-    ModelLoadStartedEvent,
-    ModelLoadCompleteEvent,
-    ModelInstallDownloadProgressEvent,
-    ModelInstallDownloadsCompleteEvent,
-    ModelInstallStartedEvent,
-    ModelInstallCompleteEvent,
-    ModelInstallCancelledEvent,
-    ModelInstallErrorEvent,
-}
-
-BULK_DOWNLOAD_EVENTS = {BulkDownloadStartedEvent, BulkDownloadCompleteEvent, BulkDownloadErrorEvent}
+from ..services.events.events_base import EventServiceBase
 
 
 class SocketIO:
-    _sub_queue = "subscribe_queue"
-    _unsub_queue = "unsubscribe_queue"
+    __sio: AsyncServer
+    __app: ASGIApp
 
-    _sub_bulk_download = "subscribe_bulk_download"
-    _unsub_bulk_download = "unsubscribe_bulk_download"
+    __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 = AsyncServer(async_mode="asgi", cors_allowed_origins="*")
+        self.__app = ASGIApp(socketio_server=self.__sio, socketio_path="/ws/socket.io")
+        app.mount("/ws", self.__app)
 
-        self._sio.on(self._sub_queue, handler=self._handle_sub_queue)
-        self._sio.on(self._unsub_queue, handler=self._handle_unsub_queue)
-        self._sio.on(self._sub_bulk_download, handler=self._handle_sub_bulk_download)
-        self._sio.on(self._unsub_bulk_download, handler=self._handle_unsub_bulk_download)
+        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)
 
-        register_events(QUEUE_EVENTS, self._handle_queue_event)
-        register_events(MODEL_EVENTS, self._handle_model_event)
-        register_events(BULK_DOWNLOAD_EVENTS, self._handle_bulk_image_download_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_sub_queue(self, sid: str, data: Any) -> None:
-        await self._sio.enter_room(sid, QueueSubscriptionEvent(**data).queue_id)
+    async def _handle_queue_event(self, event: Event):
+        await self.__sio.emit(
+            event=event[1]["event"],
+            data=event[1]["data"],
+            room=event[1]["data"]["queue_id"],
+        )
 
-    async def _handle_unsub_queue(self, sid: str, data: Any) -> None:
-        await self._sio.leave_room(sid, QueueSubscriptionEvent(**data).queue_id)
+    async def _handle_sub_queue(self, sid, data, *args, **kwargs) -> None:
+        if "queue_id" in data:
+            await self.__sio.enter_room(sid, data["queue_id"])
 
-    async def _handle_sub_bulk_download(self, sid: str, data: Any) -> None:
-        await self._sio.enter_room(sid, BulkDownloadSubscriptionEvent(**data).bulk_download_id)
+    async def _handle_unsub_queue(self, sid, data, *args, **kwargs) -> None:
+        if "queue_id" in data:
+            await self.__sio.leave_room(sid, data["queue_id"])
 
-    async def _handle_unsub_bulk_download(self, sid: str, data: Any) -> None:
-        await self._sio.leave_room(sid, BulkDownloadSubscriptionEvent(**data).bulk_download_id)
+    async def _handle_model_event(self, event: Event) -> None:
+        await self.__sio.emit(event=event[1]["event"], data=event[1]["data"])
 
-    async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]):
-        await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json"), room=event[1].queue_id)
+    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_model_event(self, event: FastAPIEvent[ModelEventBase | DownloadEventBase]) -> None:
-        await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json"))
+    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_bulk_image_download_event(self, event: FastAPIEvent[BulkDownloadEventBase]) -> None:
-        await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json"), room=event[1].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

@@ -3,7 +3,9 @@ import logging
 import mimetypes
 import socket
 from contextlib import asynccontextmanager
+from inspect import signature
 from pathlib import Path
+from typing import Any
 
 import torch
 import uvicorn
@@ -11,9 +13,11 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
+from fastapi.openapi.utils import get_openapi
 from fastapi.responses import HTMLResponse
 from fastapi_events.handlers.local import local_handler
 from fastapi_events.middleware import EventHandlerASGIMiddleware
+from pydantic.json_schema import models_json_schema
 from torch.backends.mps import is_available as is_mps_available
 
 # for PyCharm:
@@ -21,9 +25,10 @@ from torch.backends.mps import is_available as is_mps_available
 import invokeai.backend.util.hotfixes  # noqa: F401 (monkeypatching on import)
 import invokeai.frontend.web as web_dir
 from invokeai.app.api.no_cache_staticfiles import NoCacheStaticFiles
+from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.services.config.config_default import get_config
-from invokeai.app.util.custom_openapi import get_openapi_func
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.app.services.session_processor.session_processor_common import ProgressImage
+from invokeai.backend.util.devices import get_torch_device_name
 
 from ..backend.util.logging import InvokeAILogger
 from .api.dependencies import ApiDependencies
@@ -39,6 +44,11 @@ from .api.routers import (
     workflows,
 )
 from .api.sockets import SocketIO
+from .invocations.baseinvocation import (
+    BaseInvocation,
+    UIConfigBase,
+)
+from .invocations.fields import InputFieldJSONSchemaExtra, OutputFieldJSONSchemaExtra
 
 app_config = get_config()
 
@@ -53,7 +63,7 @@ logger = InvokeAILogger.get_logger(config=app_config)
 mimetypes.add_type("application/javascript", ".js")
 mimetypes.add_type("text/css", ".css")
 
-torch_device_name = TorchDevice.get_torch_device_name()
+torch_device_name = get_torch_device_name()
 logger.info(f"Using torch device: {torch_device_name}")
 
 
@@ -108,7 +118,85 @@ app.include_router(app_info.app_router, prefix="/api")
 app.include_router(session_queue.session_queue_router, prefix="/api")
 app.include_router(workflows.workflows_router, prefix="/api")
 
-app.openapi = get_openapi_func(app)
+
+# Build a custom OpenAPI to include all outputs
+# TODO: can outputs be included on metadata of invocation schemas somehow?
+def custom_openapi() -> dict[str, Any]:
+    if app.openapi_schema:
+        return app.openapi_schema
+    openapi_schema = get_openapi(
+        title=app.title,
+        description="An API for invoking AI image operations",
+        version="1.0.0",
+        routes=app.routes,
+        separate_input_output_schemas=False,  # https://fastapi.tiangolo.com/how-to/separate-openapi-schemas/
+    )
+
+    # Add all outputs
+    all_invocations = BaseInvocation.get_invocations()
+    output_types = set()
+    output_type_titles = {}
+    for invoker in all_invocations:
+        output_type = signature(invoker.invoke).return_annotation
+        output_types.add(output_type)
+
+    output_schemas = models_json_schema(
+        models=[(o, "serialization") for o in output_types], ref_template="#/components/schemas/{model}"
+    )
+    for schema_key, output_schema in output_schemas[1]["$defs"].items():
+        # 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"
+
+    # Some models don't end up in the schemas as standalone definitions
+    additional_schemas = models_json_schema(
+        [
+            (UIConfigBase, "serialization"),
+            (InputFieldJSONSchemaExtra, "serialization"),
+            (OutputFieldJSONSchemaExtra, "serialization"),
+            (ModelIdentifierField, "serialization"),
+            (ProgressImage, "serialization"),
+        ],
+        ref_template="#/components/schemas/{model}",
+    )
+    for schema_key, schema_json in additional_schemas[1]["$defs"].items():
+        openapi_schema["components"]["schemas"][schema_key] = schema_json
+
+    # Add a reference to the output type to additionalProperties of the invoker schema
+    for invoker in all_invocations:
+        invoker_name = invoker.__name__  # type: ignore [attr-defined] # this is a valid attribute
+        output_type = signature(obj=invoker.invoke).return_annotation
+        output_type_title = output_type_titles[output_type.__name__]
+        invoker_schema = openapi_schema["components"]["schemas"][f"{invoker_name}"]
+        outputs_ref = {"$ref": f"#/components/schemas/{output_type_title}"}
+        invoker_schema["output"] = outputs_ref
+        invoker_schema["class"] = "invocation"
+
+    # 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():
+
+    #     if model_config_name in openapi_schema["components"]["schemas"]:
+    #         # print(f"Config with name {name} already defined")
+    #         continue
+
+    #     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
+
+
+app.openapi = custom_openapi  # type: ignore [method-assign] # this is a valid assignment
 
 
 @app.get("/docs", include_in_schema=False)

invokeai/app/invocations/baseinvocation.py

@@ -98,13 +98,11 @@ class BaseInvocationOutput(BaseModel):
 
     _output_classes: ClassVar[set[BaseInvocationOutput]] = set()
     _typeadapter: ClassVar[Optional[TypeAdapter[Any]]] = None
-    _typeadapter_needs_update: ClassVar[bool] = False
 
     @classmethod
     def register_output(cls, output: BaseInvocationOutput) -> None:
         """Registers an invocation output."""
         cls._output_classes.add(output)
-        cls._typeadapter_needs_update = True
 
     @classmethod
     def get_outputs(cls) -> Iterable[BaseInvocationOutput]:
@@ -114,12 +112,11 @@ class BaseInvocationOutput(BaseModel):
     @classmethod
     def get_typeadapter(cls) -> TypeAdapter[Any]:
         """Gets a pydantc TypeAdapter for the union of all invocation output types."""
-        if not cls._typeadapter or cls._typeadapter_needs_update:
-            AnyInvocationOutput = TypeAliasType(
-                "AnyInvocationOutput", Annotated[Union[tuple(cls._output_classes)], Field(discriminator="type")]
+        if not cls._typeadapter:
+            InvocationOutputsUnion = TypeAliasType(
+                "InvocationOutputsUnion", Annotated[Union[tuple(cls._output_classes)], Field(discriminator="type")]
             )
-            cls._typeadapter = TypeAdapter(AnyInvocationOutput)
-            cls._typeadapter_needs_update = False
+            cls._typeadapter = TypeAdapter(InvocationOutputsUnion)
         return cls._typeadapter
 
     @classmethod
@@ -128,13 +125,12 @@ class BaseInvocationOutput(BaseModel):
         return (i.get_type() for i in BaseInvocationOutput.get_outputs())
 
     @staticmethod
-    def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseInvocationOutput]) -> None:
+    def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseModel]) -> None:
         """Adds various UI-facing attributes to the invocation output's OpenAPI schema."""
         # Because we use a pydantic Literal field with default value for the invocation type,
         # it will be typed as optional in the OpenAPI schema. Make it required manually.
         if "required" not in schema or not isinstance(schema["required"], list):
             schema["required"] = []
-        schema["class"] = "output"
         schema["required"].extend(["type"])
 
     @classmethod
@@ -171,7 +167,6 @@ class BaseInvocation(ABC, BaseModel):
 
     _invocation_classes: ClassVar[set[BaseInvocation]] = set()
     _typeadapter: ClassVar[Optional[TypeAdapter[Any]]] = None
-    _typeadapter_needs_update: ClassVar[bool] = False
 
     @classmethod
     def get_type(cls) -> str:
@@ -182,17 +177,15 @@ class BaseInvocation(ABC, BaseModel):
     def register_invocation(cls, invocation: BaseInvocation) -> None:
         """Registers an invocation."""
         cls._invocation_classes.add(invocation)
-        cls._typeadapter_needs_update = True
 
     @classmethod
     def get_typeadapter(cls) -> TypeAdapter[Any]:
         """Gets a pydantc TypeAdapter for the union of all invocation types."""
-        if not cls._typeadapter or cls._typeadapter_needs_update:
-            AnyInvocation = TypeAliasType(
-                "AnyInvocation", Annotated[Union[tuple(cls._invocation_classes)], Field(discriminator="type")]
+        if not cls._typeadapter:
+            InvocationsUnion = TypeAliasType(
+                "InvocationsUnion", Annotated[Union[tuple(cls._invocation_classes)], Field(discriminator="type")]
             )
-            cls._typeadapter = TypeAdapter(AnyInvocation)
-            cls._typeadapter_needs_update = False
+            cls._typeadapter = TypeAdapter(InvocationsUnion)
         return cls._typeadapter
 
     @classmethod
@@ -228,7 +221,7 @@ class BaseInvocation(ABC, BaseModel):
         return signature(cls.invoke).return_annotation
 
     @staticmethod
-    def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseInvocation]) -> None:
+    def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseModel], *args, **kwargs) -> None:
         """Adds various UI-facing attributes to the invocation's OpenAPI schema."""
         uiconfig = cast(UIConfigBase | None, getattr(model_class, "UIConfig", None))
         if uiconfig is not None:
@@ -244,7 +237,6 @@ class BaseInvocation(ABC, BaseModel):
             schema["version"] = uiconfig.version
         if "required" not in schema or not isinstance(schema["required"], list):
             schema["required"] = []
-        schema["class"] = "invocation"
         schema["required"].extend(["type", "id"])
 
     @abstractmethod
@@ -318,7 +310,7 @@ class BaseInvocation(ABC, BaseModel):
         protected_namespaces=(),
         validate_assignment=True,
         json_schema_extra=json_schema_extra,
-        json_schema_serialization_defaults_required=False,
+        json_schema_serialization_defaults_required=True,
         coerce_numbers_to_str=True,
     )
 

invokeai/app/invocations/blend_latents.py

@@ -1,98 +0,0 @@
-from typing import Any, Union
-
-import numpy as np
-import numpy.typing as npt
-import torch
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, LatentsField
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.util.devices import TorchDevice
-
-
-@invocation(
-    "lblend",
-    title="Blend Latents",
-    tags=["latents", "blend"],
-    category="latents",
-    version="1.0.3",
-)
-class BlendLatentsInvocation(BaseInvocation):
-    """Blend two latents using a given alpha. Latents must have same size."""
-
-    latents_a: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    latents_b: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    alpha: float = InputField(default=0.5, description=FieldDescriptions.blend_alpha)
-
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents_a = context.tensors.load(self.latents_a.latents_name)
-        latents_b = context.tensors.load(self.latents_b.latents_name)
-
-        if latents_a.shape != latents_b.shape:
-            raise Exception("Latents to blend must be the same size.")
-
-        device = TorchDevice.choose_torch_device()
-
-        def slerp(
-            t: Union[float, npt.NDArray[Any]],  # FIXME: maybe use np.float32 here?
-            v0: Union[torch.Tensor, npt.NDArray[Any]],
-            v1: Union[torch.Tensor, npt.NDArray[Any]],
-            DOT_THRESHOLD: float = 0.9995,
-        ) -> Union[torch.Tensor, npt.NDArray[Any]]:
-            """
-            Spherical linear interpolation
-            Args:
-                t (float/np.ndarray): Float value between 0.0 and 1.0
-                v0 (np.ndarray): Starting vector
-                v1 (np.ndarray): Final vector
-                DOT_THRESHOLD (float): Threshold for considering the two vectors as
-                                    colineal. Not recommended to alter this.
-            Returns:
-                v2 (np.ndarray): Interpolation vector between v0 and v1
-            """
-            inputs_are_torch = False
-            if not isinstance(v0, np.ndarray):
-                inputs_are_torch = True
-                v0 = v0.detach().cpu().numpy()
-            if not isinstance(v1, np.ndarray):
-                inputs_are_torch = True
-                v1 = v1.detach().cpu().numpy()
-
-            dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
-            if np.abs(dot) > DOT_THRESHOLD:
-                v2 = (1 - t) * v0 + t * v1
-            else:
-                theta_0 = np.arccos(dot)
-                sin_theta_0 = np.sin(theta_0)
-                theta_t = theta_0 * t
-                sin_theta_t = np.sin(theta_t)
-                s0 = np.sin(theta_0 - theta_t) / sin_theta_0
-                s1 = sin_theta_t / sin_theta_0
-                v2 = s0 * v0 + s1 * v1
-
-            if inputs_are_torch:
-                v2_torch: torch.Tensor = torch.from_numpy(v2).to(device)
-                return v2_torch
-            else:
-                assert isinstance(v2, np.ndarray)
-                return v2
-
-        # blend
-        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")
-
-        TorchDevice.empty_cache()
-
-        name = context.tensors.save(tensor=blended_latents)
-        return LatentsOutput.build(latents_name=name, latents=blended_latents, seed=self.latents_a.seed)

invokeai/app/invocations/compel.py

@@ -24,7 +24,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     ConditioningFieldData,
     SDXLConditioningInfo,
 )
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import torch_dtype
 
 from .baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
 from .model import CLIPField
@@ -65,7 +65,11 @@ class CompelInvocation(BaseInvocation):
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
         tokenizer_info = context.models.load(self.clip.tokenizer)
+        tokenizer_model = tokenizer_info.model
+        assert isinstance(tokenizer_model, CLIPTokenizer)
         text_encoder_info = context.models.load(self.clip.text_encoder)
+        text_encoder_model = text_encoder_info.model
+        assert isinstance(text_encoder_model, CLIPTextModel)
 
         def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.clip.loras:
@@ -80,35 +84,29 @@ class CompelInvocation(BaseInvocation):
         ti_list = generate_ti_list(self.prompt, text_encoder_info.config.base, context)
 
         with (
-            # apply all patches while the model is on the target device
-            text_encoder_info.model_on_device() as (model_state_dict, text_encoder),
-            tokenizer_info as tokenizer,
-            ModelPatcher.apply_lora_text_encoder(
-                text_encoder,
-                loras=_lora_loader(),
-                model_state_dict=model_state_dict,
-            ),
-            # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
-            ModelPatcher.apply_clip_skip(text_encoder, self.clip.skipped_layers),
-            ModelPatcher.apply_ti(tokenizer, text_encoder, ti_list) as (
-                patched_tokenizer,
+            ModelPatcher.apply_ti(tokenizer_model, text_encoder_model, ti_list) as (
+                tokenizer,
                 ti_manager,
             ),
+            text_encoder_info as text_encoder,
+            # Apply the LoRA after text_encoder has been moved to its target device for faster patching.
+            ModelPatcher.apply_lora_text_encoder(text_encoder, _lora_loader()),
+            # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
+            ModelPatcher.apply_clip_skip(text_encoder_model, self.clip.skipped_layers),
         ):
             assert isinstance(text_encoder, CLIPTextModel)
-            assert isinstance(tokenizer, CLIPTokenizer)
             compel = Compel(
-                tokenizer=patched_tokenizer,
+                tokenizer=tokenizer,
                 text_encoder=text_encoder,
                 textual_inversion_manager=ti_manager,
-                dtype_for_device_getter=TorchDevice.choose_torch_dtype,
+                dtype_for_device_getter=torch_dtype,
                 truncate_long_prompts=False,
             )
 
             conjunction = Compel.parse_prompt_string(self.prompt)
 
             if context.config.get().log_tokenization:
-                log_tokenization_for_conjunction(conjunction, patched_tokenizer)
+                log_tokenization_for_conjunction(conjunction, tokenizer)
 
             c, _options = compel.build_conditioning_tensor_for_conjunction(conjunction)
 
@@ -138,7 +136,11 @@ class SDXLPromptInvocationBase:
         zero_on_empty: bool,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         tokenizer_info = context.models.load(clip_field.tokenizer)
+        tokenizer_model = tokenizer_info.model
+        assert isinstance(tokenizer_model, CLIPTokenizer)
         text_encoder_info = context.models.load(clip_field.text_encoder)
+        text_encoder_model = text_encoder_info.model
+        assert isinstance(text_encoder_model, (CLIPTextModel, CLIPTextModelWithProjection))
 
         # return zero on empty
         if prompt == "" and zero_on_empty:
@@ -175,31 +177,23 @@ class SDXLPromptInvocationBase:
         ti_list = generate_ti_list(prompt, text_encoder_info.config.base, context)
 
         with (
-            # apply all patches while the model is on the target device
-            text_encoder_info.model_on_device() as (state_dict, text_encoder),
-            tokenizer_info as tokenizer,
-            ModelPatcher.apply_lora(
-                text_encoder,
-                loras=_lora_loader(),
-                prefix=lora_prefix,
-                model_state_dict=state_dict,
-            ),
-            # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
-            ModelPatcher.apply_clip_skip(text_encoder, clip_field.skipped_layers),
-            ModelPatcher.apply_ti(tokenizer, text_encoder, ti_list) as (
-                patched_tokenizer,
+            ModelPatcher.apply_ti(tokenizer_model, text_encoder_model, ti_list) as (
+                tokenizer,
                 ti_manager,
             ),
+            text_encoder_info as text_encoder,
+            # Apply the LoRA after text_encoder has been moved to its target device for faster patching.
+            ModelPatcher.apply_lora(text_encoder, _lora_loader(), lora_prefix),
+            # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
+            ModelPatcher.apply_clip_skip(text_encoder_model, clip_field.skipped_layers),
         ):
             assert isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection))
-            assert isinstance(tokenizer, CLIPTokenizer)
-
             text_encoder = cast(CLIPTextModel, text_encoder)
             compel = Compel(
-                tokenizer=patched_tokenizer,
+                tokenizer=tokenizer,
                 text_encoder=text_encoder,
                 textual_inversion_manager=ti_manager,
-                dtype_for_device_getter=TorchDevice.choose_torch_dtype,
+                dtype_for_device_getter=torch_dtype,
                 truncate_long_prompts=False,  # TODO:
                 returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,  # TODO: clip skip
                 requires_pooled=get_pooled,
@@ -209,7 +203,7 @@ class SDXLPromptInvocationBase:
 
             if context.config.get().log_tokenization:
                 # TODO: better logging for and syntax
-                log_tokenization_for_conjunction(conjunction, patched_tokenizer)
+                log_tokenization_for_conjunction(conjunction, tokenizer)
 
             # TODO: ask for optimizations? to not run text_encoder twice
             c, _options = compel.build_conditioning_tensor_for_conjunction(conjunction)

invokeai/app/invocations/constants.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
 
 LATENT_SCALE_FACTOR = 8
 """
@@ -10,7 +10,8 @@ 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"""
-
-DEFAULT_PRECISION = TorchDevice.choose_torch_dtype()

invokeai/app/invocations/controlnet_image_processors.py

@@ -2,7 +2,6 @@
 # initial implementation by Gregg Helt, 2023
 # heavily leverages controlnet_aux package: https://github.com/patrickvonplaten/controlnet_aux
 from builtins import bool, float
-from pathlib import Path
 from typing import Dict, List, Literal, Union
 
 import cv2
@@ -25,6 +24,7 @@ from pydantic import BaseModel, Field, field_validator, model_validator
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     ImageField,
+    Input,
     InputField,
     OutputField,
     UIType,
@@ -35,17 +35,22 @@ from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES, heuristic_resize
 from invokeai.backend.image_util.canny import get_canny_edges
-from invokeai.backend.image_util.depth_anything import DEPTH_ANYTHING_MODELS, DepthAnythingDetector
-from invokeai.backend.image_util.dw_openpose import DWPOSE_MODELS, DWOpenposeDetector
+from invokeai.backend.image_util.depth_anything import DepthAnythingDetector
+from invokeai.backend.image_util.dw_openpose import DWOpenposeDetector
 from invokeai.backend.image_util.hed import HEDProcessor
 from invokeai.backend.image_util.lineart import LineartProcessor
 from invokeai.backend.image_util.lineart_anime import LineartAnimeProcessor
-from invokeai.backend.image_util.util import np_to_pil, pil_to_np
-from invokeai.backend.util.devices import TorchDevice
 
-from .baseinvocation import BaseInvocation, BaseInvocationOutput, Classification, 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[
+    "just_resize",
+    "crop_resize",
+    "fill_resize",
+    "just_resize_simple",
+]
 
 
 class ControlField(BaseModel):
@@ -81,13 +86,13 @@ class ControlOutput(BaseInvocationOutput):
     control: ControlField = OutputField(description=FieldDescriptions.control)
 
 
-@invocation("controlnet", title="ControlNet", tags=["controlnet"], category="controlnet", version="1.1.2")
+@invocation("controlnet", title="ControlNet", tags=["controlnet"], category="controlnet", version="1.1.1")
 class ControlNetInvocation(BaseInvocation):
     """Collects ControlNet info to pass to other nodes"""
 
     image: ImageField = InputField(description="The control image")
     control_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.controlnet_model, ui_type=UIType.ControlNetModel
+        description=FieldDescriptions.controlnet_model, input=Input.Direct, ui_type=UIType.ControlNetModel
     )
     control_weight: Union[float, List[float]] = InputField(
         default=1.0, ge=-1, le=2, description="The weight given to the ControlNet"
@@ -141,7 +146,6 @@ class ImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
         return context.images.get_pil(self.image.image_name, "RGB")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        self._context = context
         raw_image = self.load_image(context)
         # image type should be PIL.PngImagePlugin.PngImageFile ?
         processed_image = self.run_processor(raw_image)
@@ -167,13 +171,13 @@ class ImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
     title="Canny Processor",
     tags=["controlnet", "canny"],
     category="controlnet",
-    version="1.3.3",
+    version="1.3.2",
 )
 class CannyImageProcessorInvocation(ImageProcessorInvocation):
     """Canny edge detection for ControlNet"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
     low_threshold: int = InputField(
         default=100, ge=0, le=255, description="The low threshold of the Canny pixel gradient (0-255)"
     )
@@ -201,13 +205,13 @@ class CannyImageProcessorInvocation(ImageProcessorInvocation):
     title="HED (softedge) Processor",
     tags=["controlnet", "hed", "softedge"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class HedImageProcessorInvocation(ImageProcessorInvocation):
     """Applies HED edge detection to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
     # safe not supported in controlnet_aux v0.0.3
     # safe: bool = InputField(default=False, description=FieldDescriptions.safe_mode)
     scribble: bool = InputField(default=False, description=FieldDescriptions.scribble_mode)
@@ -230,13 +234,13 @@ class HedImageProcessorInvocation(ImageProcessorInvocation):
     title="Lineart Processor",
     tags=["controlnet", "lineart"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class LineartImageProcessorInvocation(ImageProcessorInvocation):
     """Applies line art processing to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
     coarse: bool = InputField(default=False, description="Whether to use coarse mode")
 
     def run_processor(self, image: Image.Image) -> Image.Image:
@@ -252,13 +256,13 @@ class LineartImageProcessorInvocation(ImageProcessorInvocation):
     title="Lineart Anime Processor",
     tags=["controlnet", "lineart", "anime"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation):
     """Applies line art anime processing to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
 
     def run_processor(self, image: Image.Image) -> Image.Image:
         processor = LineartAnimeProcessor()
@@ -275,20 +279,19 @@ class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation):
     title="Midas Depth Processor",
     tags=["controlnet", "midas"],
     category="controlnet",
-    version="1.2.4",
+    version="1.2.3",
 )
 class MidasDepthImageProcessorInvocation(ImageProcessorInvocation):
     """Applies Midas depth processing to image"""
 
     a_mult: float = InputField(default=2.0, ge=0, description="Midas parameter `a_mult` (a = a_mult * PI)")
     bg_th: float = InputField(default=0.1, ge=0, description="Midas parameter `bg_th`")
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
     # depth_and_normal not supported in controlnet_aux v0.0.3
     # depth_and_normal: bool = InputField(default=False, description="whether to use depth and normal mode")
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
-        # TODO: replace from_pretrained() calls with context.models.download_and_cache() (or similar)
+    def run_processor(self, image):
         midas_processor = MidasDetector.from_pretrained("lllyasviel/Annotators")
         processed_image = midas_processor(
             image,
@@ -307,15 +310,15 @@ class MidasDepthImageProcessorInvocation(ImageProcessorInvocation):
     title="Normal BAE Processor",
     tags=["controlnet"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class NormalbaeImageProcessorInvocation(ImageProcessorInvocation):
     """Applies NormalBae processing to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         normalbae_processor = NormalBaeDetector.from_pretrained("lllyasviel/Annotators")
         processed_image = normalbae_processor(
             image, detect_resolution=self.detect_resolution, image_resolution=self.image_resolution
@@ -324,17 +327,17 @@ class NormalbaeImageProcessorInvocation(ImageProcessorInvocation):
 
 
 @invocation(
-    "mlsd_image_processor", title="MLSD Processor", tags=["controlnet", "mlsd"], category="controlnet", version="1.2.3"
+    "mlsd_image_processor", title="MLSD Processor", tags=["controlnet", "mlsd"], category="controlnet", version="1.2.2"
 )
 class MlsdImageProcessorInvocation(ImageProcessorInvocation):
     """Applies MLSD processing to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
     thr_v: float = InputField(default=0.1, ge=0, description="MLSD parameter `thr_v`")
     thr_d: float = InputField(default=0.1, ge=0, description="MLSD parameter `thr_d`")
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         mlsd_processor = MLSDdetector.from_pretrained("lllyasviel/Annotators")
         processed_image = mlsd_processor(
             image,
@@ -347,17 +350,17 @@ class MlsdImageProcessorInvocation(ImageProcessorInvocation):
 
 
 @invocation(
-    "pidi_image_processor", title="PIDI Processor", tags=["controlnet", "pidi"], category="controlnet", version="1.2.3"
+    "pidi_image_processor", title="PIDI Processor", tags=["controlnet", "pidi"], category="controlnet", version="1.2.2"
 )
 class PidiImageProcessorInvocation(ImageProcessorInvocation):
     """Applies PIDI processing to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
     safe: bool = InputField(default=False, description=FieldDescriptions.safe_mode)
     scribble: bool = InputField(default=False, description=FieldDescriptions.scribble_mode)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         pidi_processor = PidiNetDetector.from_pretrained("lllyasviel/Annotators")
         processed_image = pidi_processor(
             image,
@@ -374,18 +377,18 @@ class PidiImageProcessorInvocation(ImageProcessorInvocation):
     title="Content Shuffle Processor",
     tags=["controlnet", "contentshuffle"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation):
     """Applies content shuffle processing to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
     h: int = InputField(default=512, ge=0, description="Content shuffle `h` parameter")
     w: int = InputField(default=512, ge=0, description="Content shuffle `w` parameter")
     f: int = InputField(default=256, ge=0, description="Content shuffle `f` parameter")
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         content_shuffle_processor = ContentShuffleDetector()
         processed_image = content_shuffle_processor(
             image,
@@ -404,12 +407,12 @@ class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation):
     title="Zoe (Depth) Processor",
     tags=["controlnet", "zoe", "depth"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation):
     """Applies Zoe depth processing to image"""
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         zoe_depth_processor = ZoeDetector.from_pretrained("lllyasviel/Annotators")
         processed_image = zoe_depth_processor(image)
         return processed_image
@@ -420,17 +423,17 @@ class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation):
     title="Mediapipe Face Processor",
     tags=["controlnet", "mediapipe", "face"],
     category="controlnet",
-    version="1.2.4",
+    version="1.2.3",
 )
 class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     """Applies mediapipe face processing to image"""
 
     max_faces: int = InputField(default=1, ge=1, description="Maximum number of faces to detect")
     min_confidence: float = InputField(default=0.5, ge=0, le=1, description="Minimum confidence for face detection")
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         mediapipe_face_processor = MediapipeFaceDetector()
         processed_image = mediapipe_face_processor(
             image,
@@ -447,7 +450,7 @@ class MediapipeFaceProcessorInvocation(ImageProcessorInvocation):
     title="Leres (Depth) Processor",
     tags=["controlnet", "leres", "depth"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class LeresImageProcessorInvocation(ImageProcessorInvocation):
     """Applies leres processing to image"""
@@ -455,10 +458,10 @@ class LeresImageProcessorInvocation(ImageProcessorInvocation):
     thr_a: float = InputField(default=0, description="Leres parameter `thr_a`")
     thr_b: float = InputField(default=0, description="Leres parameter `thr_b`")
     boost: bool = InputField(default=False, description="Whether to use boost mode")
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         leres_processor = LeresDetector.from_pretrained("lllyasviel/Annotators")
         processed_image = leres_processor(
             image,
@@ -476,7 +479,7 @@ class LeresImageProcessorInvocation(ImageProcessorInvocation):
     title="Tile Resample Processor",
     tags=["controlnet", "tile"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class TileResamplerProcessorInvocation(ImageProcessorInvocation):
     """Tile resampler processor"""
@@ -500,8 +503,8 @@ class TileResamplerProcessorInvocation(ImageProcessorInvocation):
         np_img = cv2.resize(np_img, (W, H), interpolation=cv2.INTER_AREA)
         return np_img
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
-        np_img = np.array(image, dtype=np.uint8)
+    def run_processor(self, img):
+        np_img = np.array(img, dtype=np.uint8)
         processed_np_image = self.tile_resample(
             np_img,
             # res=self.tile_size,
@@ -516,15 +519,15 @@ class TileResamplerProcessorInvocation(ImageProcessorInvocation):
     title="Segment Anything Processor",
     tags=["controlnet", "segmentanything"],
     category="controlnet",
-    version="1.2.4",
+    version="1.2.3",
 )
 class SegmentAnythingProcessorInvocation(ImageProcessorInvocation):
     """Applies segment anything processing to image"""
 
-    detect_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.detect_res)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    detect_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.detect_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image):
         # segment_anything_processor = SamDetector.from_pretrained("ybelkada/segment-anything", subfolder="checkpoints")
         segment_anything_processor = SamDetectorReproducibleColors.from_pretrained(
             "ybelkada/segment-anything", subfolder="checkpoints"
@@ -563,14 +566,14 @@ class SamDetectorReproducibleColors(SamDetector):
     title="Color Map Processor",
     tags=["controlnet"],
     category="controlnet",
-    version="1.2.3",
+    version="1.2.2",
 )
 class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
     """Generates a color map from the provided image"""
 
-    color_map_tile_size: int = InputField(default=64, ge=1, description=FieldDescriptions.tile_size)
+    color_map_tile_size: int = InputField(default=64, ge=0, description=FieldDescriptions.tile_size)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
+    def run_processor(self, image: Image.Image):
         np_image = np.array(image, dtype=np.uint8)
         height, width = np_image.shape[:2]
 
@@ -595,7 +598,7 @@ DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small"]
     title="Depth Anything Processor",
     tags=["controlnet", "depth", "depth anything"],
     category="controlnet",
-    version="1.1.2",
+    version="1.1.1",
 )
 class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
     """Generates a depth map based on the Depth Anything algorithm"""
@@ -603,20 +606,14 @@ class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
     model_size: DEPTH_ANYTHING_MODEL_SIZES = InputField(
         default="small", description="The size of the depth model to use"
     )
-    resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    resolution: int = InputField(default=512, ge=64, multiple_of=64, description=FieldDescriptions.image_res)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
-        def loader(model_path: Path):
-            return DepthAnythingDetector.load_model(
-                model_path, model_size=self.model_size, device=TorchDevice.choose_torch_device()
-            )
+    def run_processor(self, image: Image.Image):
+        depth_anything_detector = DepthAnythingDetector()
+        depth_anything_detector.load_model(model_size=self.model_size)
 
-        with self._context.models.load_remote_model(
-            source=DEPTH_ANYTHING_MODELS[self.model_size], loader=loader
-        ) as model:
-            depth_anything_detector = DepthAnythingDetector(model, TorchDevice.choose_torch_device())
-            processed_image = depth_anything_detector(image=image, resolution=self.resolution)
-            return processed_image
+        processed_image = depth_anything_detector(image=image, resolution=self.resolution)
+        return processed_image
 
 
 @invocation(
@@ -624,7 +621,7 @@ class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
     title="DW Openpose Image Processor",
     tags=["controlnet", "dwpose", "openpose"],
     category="controlnet",
-    version="1.1.1",
+    version="1.1.0",
 )
 class DWOpenposeImageProcessorInvocation(ImageProcessorInvocation):
     """Generates an openpose pose from an image using DWPose"""
@@ -632,13 +629,10 @@ class DWOpenposeImageProcessorInvocation(ImageProcessorInvocation):
     draw_body: bool = InputField(default=True)
     draw_face: bool = InputField(default=False)
     draw_hands: bool = InputField(default=False)
-    image_resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
+    image_resolution: int = InputField(default=512, ge=0, description=FieldDescriptions.image_res)
 
-    def run_processor(self, image: Image.Image) -> Image.Image:
-        onnx_det = self._context.models.download_and_cache_model(DWPOSE_MODELS["yolox_l.onnx"])
-        onnx_pose = self._context.models.download_and_cache_model(DWPOSE_MODELS["dw-ll_ucoco_384.onnx"])
-
-        dw_openpose = DWOpenposeDetector(onnx_det=onnx_det, onnx_pose=onnx_pose)
+    def run_processor(self, image: Image.Image):
+        dw_openpose = DWOpenposeDetector()
         processed_image = dw_openpose(
             image,
             draw_face=self.draw_face,
@@ -647,27 +641,3 @@ class DWOpenposeImageProcessorInvocation(ImageProcessorInvocation):
             resolution=self.image_resolution,
         )
         return processed_image
-
-
-@invocation(
-    "heuristic_resize",
-    title="Heuristic Resize",
-    tags=["image, controlnet"],
-    category="image",
-    version="1.0.1",
-    classification=Classification.Prototype,
-)
-class HeuristicResizeInvocation(BaseInvocation):
-    """Resize an image using a heuristic method. Preserves edge maps."""
-
-    image: ImageField = InputField(description="The image to resize")
-    width: int = InputField(default=512, ge=1, description="The width to resize to (px)")
-    height: int = InputField(default=512, ge=1, description="The height to resize to (px)")
-
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.images.get_pil(self.image.image_name, "RGB")
-        np_img = pil_to_np(image)
-        np_resized = heuristic_resize(np_img, (self.width, self.height))
-        resized = np_to_pil(np_resized)
-        image_dto = context.images.save(image=resized)
-        return ImageOutput.build(image_dto)

invokeai/app/invocations/create_denoise_mask.py

@@ -1,80 +0,0 @@
-from typing import Optional
-
-import torch
-import torchvision.transforms as T
-from PIL import Image
-from torchvision.transforms.functional import resize as tv_resize
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import DEFAULT_PRECISION
-from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField
-from invokeai.app.invocations.image_to_latents import ImageToLatentsInvocation
-from invokeai.app.invocations.model import VAEField
-from invokeai.app.invocations.primitives import DenoiseMaskOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-
-
-@invocation(
-    "create_denoise_mask",
-    title="Create Denoise Mask",
-    tags=["mask", "denoise"],
-    category="latents",
-    version="1.0.2",
-)
-class CreateDenoiseMaskInvocation(BaseInvocation):
-    """Creates mask for denoising model run."""
-
-    vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection, ui_order=0)
-    image: Optional[ImageField] = InputField(default=None, description="Image which will be masked", ui_order=1)
-    mask: ImageField = InputField(description="The mask to use when pasting", ui_order=2)
-    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled, ui_order=3)
-    fp32: bool = InputField(
-        default=DEFAULT_PRECISION == torch.float32,
-        description=FieldDescriptions.fp32,
-        ui_order=4,
-    )
-
-    def prep_mask_tensor(self, mask_image: Image.Image) -> torch.Tensor:
-        if mask_image.mode != "L":
-            mask_image = mask_image.convert("L")
-        mask_tensor: torch.Tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
-        if mask_tensor.dim() == 3:
-            mask_tensor = mask_tensor.unsqueeze(0)
-        # if shape is not None:
-        #    mask_tensor = tv_resize(mask_tensor, shape, T.InterpolationMode.BILINEAR)
-        return mask_tensor
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> DenoiseMaskOutput:
-        if self.image is not None:
-            image = context.images.get_pil(self.image.image_name)
-            image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
-            if image_tensor.dim() == 3:
-                image_tensor = image_tensor.unsqueeze(0)
-        else:
-            image_tensor = None
-
-        mask = self.prep_mask_tensor(
-            context.images.get_pil(self.mask.image_name),
-        )
-
-        if image_tensor is not None:
-            vae_info = context.models.load(self.vae.vae)
-
-            img_mask = tv_resize(mask, image_tensor.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
-            masked_image = image_tensor * torch.where(img_mask < 0.5, 0.0, 1.0)
-            # TODO:
-            masked_latents = ImageToLatentsInvocation.vae_encode(vae_info, self.fp32, self.tiled, masked_image.clone())
-
-            masked_latents_name = context.tensors.save(tensor=masked_latents)
-        else:
-            masked_latents_name = None
-
-        mask_name = context.tensors.save(tensor=mask)
-
-        return DenoiseMaskOutput.build(
-            mask_name=mask_name,
-            masked_latents_name=masked_latents_name,
-            gradient=False,
-        )

invokeai/app/invocations/create_gradient_mask.py

@@ -1,138 +0,0 @@
-from typing import Literal, Optional
-
-import numpy as np
-import torch
-import torchvision.transforms as T
-from PIL import Image, ImageFilter
-from torchvision.transforms.functional import resize as tv_resize
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.constants import DEFAULT_PRECISION
-from invokeai.app.invocations.fields import (
-    DenoiseMaskField,
-    FieldDescriptions,
-    ImageField,
-    Input,
-    InputField,
-    OutputField,
-)
-from invokeai.app.invocations.image_to_latents import ImageToLatentsInvocation
-from invokeai.app.invocations.model import UNetField, VAEField
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager.config import MainConfigBase, ModelVariantType
-from invokeai.backend.model_manager.load.load_base import LoadedModel
-from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-
-
-@invocation_output("gradient_mask_output")
-class GradientMaskOutput(BaseInvocationOutput):
-    """Outputs a denoise mask and an image representing the total gradient of the mask."""
-
-    denoise_mask: DenoiseMaskField = OutputField(description="Mask for denoise model run")
-    expanded_mask_area: ImageField = OutputField(
-        description="Image representing the total gradient area of the mask. For paste-back purposes."
-    )
-
-
-@invocation(
-    "create_gradient_mask",
-    title="Create Gradient Mask",
-    tags=["mask", "denoise"],
-    category="latents",
-    version="1.1.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
-    )
-    image: Optional[ImageField] = InputField(
-        default=None,
-        description="OPTIONAL: Only connect for specialized Inpainting models, masked_latents will be generated from the image with the VAE",
-        title="[OPTIONAL] Image",
-        ui_order=6,
-    )
-    unet: Optional[UNetField] = InputField(
-        description="OPTIONAL: If the Unet is a specialized Inpainting model, masked_latents will be generated from the image with the VAE",
-        default=None,
-        input=Input.Connection,
-        title="[OPTIONAL] UNet",
-        ui_order=5,
-    )
-    vae: Optional[VAEField] = InputField(
-        default=None,
-        description="OPTIONAL: Only connect for specialized Inpainting models, masked_latents will be generated from the image with the VAE",
-        title="[OPTIONAL] VAE",
-        input=Input.Connection,
-        ui_order=7,
-    )
-    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled, ui_order=8)
-    fp32: bool = InputField(
-        default=DEFAULT_PRECISION == torch.float32,
-        description=FieldDescriptions.fp32,
-        ui_order=9,
-    )
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> GradientMaskOutput:
-        mask_image = context.images.get_pil(self.mask.image_name, mode="L")
-        if self.edge_radius > 0:
-            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))
-
-            blur_tensor: torch.Tensor = image_resized_to_grid_as_tensor(blur_mask, normalize=False)
-
-            # redistribute blur so that the original edges are 0 and blur outwards to 1
-            blur_tensor = (blur_tensor - 0.5) * 2
-
-            threshold = 1 - self.minimum_denoise
-
-            if self.coherence_mode == "Staged":
-                # wherever the blur_tensor is less than fully masked, convert it to threshold
-                blur_tensor = torch.where((blur_tensor < 1) & (blur_tensor > 0), threshold, blur_tensor)
-            else:
-                # wherever the blur_tensor is above threshold but less than 1, drop it to threshold
-                blur_tensor = torch.where((blur_tensor > threshold) & (blur_tensor < 1), threshold, blur_tensor)
-
-        else:
-            blur_tensor: torch.Tensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
-
-        mask_name = context.tensors.save(tensor=blur_tensor.unsqueeze(1))
-
-        # compute a [0, 1] mask from the blur_tensor
-        expanded_mask = torch.where((blur_tensor < 1), 0, 1)
-        expanded_mask_image = Image.fromarray((expanded_mask.squeeze(0).numpy() * 255).astype(np.uint8), mode="L")
-        expanded_image_dto = context.images.save(expanded_mask_image)
-
-        masked_latents_name = None
-        if self.unet is not None and self.vae is not None and self.image is not None:
-            # all three fields must be present at the same time
-            main_model_config = context.models.get_config(self.unet.unet.key)
-            assert isinstance(main_model_config, MainConfigBase)
-            if main_model_config.variant is ModelVariantType.Inpaint:
-                mask = blur_tensor
-                vae_info: LoadedModel = context.models.load(self.vae.vae)
-                image = context.images.get_pil(self.image.image_name)
-                image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
-                if image_tensor.dim() == 3:
-                    image_tensor = image_tensor.unsqueeze(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)
-                masked_latents = ImageToLatentsInvocation.vae_encode(
-                    vae_info, self.fp32, self.tiled, masked_image.clone()
-                )
-                masked_latents_name = context.tensors.save(tensor=masked_latents)
-
-        return GradientMaskOutput(
-            denoise_mask=DenoiseMaskField(mask_name=mask_name, masked_latents_name=masked_latents_name, gradient=True),
-            expanded_mask_area=ImageField(image_name=expanded_image_dto.image_name),
-        )

invokeai/app/invocations/crop_latents.py

@@ -1,61 +0,0 @@
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, LatentsField
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-
-
-# The Crop Latents node was copied from @skunkworxdark's implementation here:
-# https://github.com/skunkworxdark/XYGrid_nodes/blob/74647fa9c1fa57d317a94bd43ca689af7f0aae5e/images_to_grids.py#L1117C1-L1167C80
-@invocation(
-    "crop_latents",
-    title="Crop Latents",
-    tags=["latents", "crop"],
-    category="latents",
-    version="1.0.2",
-)
-# TODO(ryand): Named `CropLatentsCoreInvocation` to prevent a conflict with custom node `CropLatentsInvocation`.
-# Currently, if the class names conflict then 'GET /openapi.json' fails.
-class CropLatentsCoreInvocation(BaseInvocation):
-    """Crops a latent-space tensor to a box specified in image-space. The box dimensions and coordinates must be
-    divisible by the latent scale factor of 8.
-    """
-
-    latents: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    x: int = InputField(
-        ge=0,
-        multiple_of=LATENT_SCALE_FACTOR,
-        description="The left x coordinate (in px) of the crop rectangle in image space. This value will be converted to a dimension in latent space.",
-    )
-    y: int = InputField(
-        ge=0,
-        multiple_of=LATENT_SCALE_FACTOR,
-        description="The top y coordinate (in px) of the crop rectangle in image space. This value will be converted to a dimension in latent space.",
-    )
-    width: int = InputField(
-        ge=1,
-        multiple_of=LATENT_SCALE_FACTOR,
-        description="The width (in px) of the crop rectangle in image space. This value will be converted to a dimension in latent space.",
-    )
-    height: int = InputField(
-        ge=1,
-        multiple_of=LATENT_SCALE_FACTOR,
-        description="The height (in px) of the crop rectangle in image space. This value will be converted to a dimension in latent space.",
-    )
-
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.tensors.load(self.latents.latents_name)
-
-        x1 = self.x // LATENT_SCALE_FACTOR
-        y1 = self.y // LATENT_SCALE_FACTOR
-        x2 = x1 + (self.width // LATENT_SCALE_FACTOR)
-        y2 = y1 + (self.height // LATENT_SCALE_FACTOR)
-
-        cropped_latents = latents[..., y1:y2, x1:x2]
-
-        name = context.tensors.save(tensor=cropped_latents)
-
-        return LatentsOutput.build(latents_name=name, latents=cropped_latents)

invokeai/app/invocations/denoise_latents.py

@@ -1,848 +0,0 @@
-# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
-import inspect
-from contextlib import ExitStack
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
-
-import torch
-import torchvision
-import torchvision.transforms as T
-from diffusers.configuration_utils import ConfigMixin
-from diffusers.models.adapter import T2IAdapter
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
-from diffusers.schedulers.scheduling_dpmsolver_sde import DPMSolverSDEScheduler
-from diffusers.schedulers.scheduling_tcd import TCDScheduler
-from diffusers.schedulers.scheduling_utils import SchedulerMixin as Scheduler
-from pydantic import field_validator
-from torchvision.transforms.functional import resize as tv_resize
-from transformers import CLIPVisionModelWithProjection
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
-from invokeai.app.invocations.controlnet_image_processors import ControlField
-from invokeai.app.invocations.fields import (
-    ConditioningField,
-    DenoiseMaskField,
-    FieldDescriptions,
-    Input,
-    InputField,
-    LatentsField,
-    UIType,
-)
-from invokeai.app.invocations.ip_adapter import IPAdapterField
-from invokeai.app.invocations.model import ModelIdentifierField, UNetField
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.invocations.t2i_adapter import T2IAdapterField
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.controlnet_utils import prepare_control_image
-from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
-from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_manager import BaseModelType
-from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.stable_diffusion import PipelineIntermediateState, set_seamless
-from invokeai.backend.stable_diffusion.diffusers_pipeline import (
-    ControlNetData,
-    StableDiffusionGeneratorPipeline,
-    T2IAdapterData,
-)
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
-    BasicConditioningInfo,
-    IPAdapterConditioningInfo,
-    IPAdapterData,
-    Range,
-    SDXLConditioningInfo,
-    TextConditioningData,
-    TextConditioningRegions,
-)
-from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_MAP, SCHEDULER_NAME_VALUES
-from invokeai.backend.util.devices import TorchDevice
-from invokeai.backend.util.hotfixes import ControlNetModel
-from invokeai.backend.util.mask import to_standard_float_mask
-from invokeai.backend.util.silence_warnings import SilenceWarnings
-
-
-def get_scheduler(
-    context: InvocationContext,
-    scheduler_info: ModelIdentifierField,
-    scheduler_name: str,
-    seed: int,
-) -> Scheduler:
-    """Load a scheduler and apply some scheduler-specific overrides."""
-    # TODO(ryand): Silently falling back to ddim seems like a bad idea. Look into why this was added and remove if
-    # possible.
-    scheduler_class, scheduler_extra_config = SCHEDULER_MAP.get(scheduler_name, SCHEDULER_MAP["ddim"])
-    orig_scheduler_info = context.models.load(scheduler_info)
-    with orig_scheduler_info as orig_scheduler:
-        scheduler_config = orig_scheduler.config
-
-    if "_backup" in scheduler_config:
-        scheduler_config = scheduler_config["_backup"]
-    scheduler_config = {
-        **scheduler_config,
-        **scheduler_extra_config,  # FIXME
-        "_backup": scheduler_config,
-    }
-
-    # make dpmpp_sde reproducable(seed can be passed only in initializer)
-    if scheduler_class is DPMSolverSDEScheduler:
-        scheduler_config["noise_sampler_seed"] = seed
-
-    scheduler = scheduler_class.from_config(scheduler_config)
-
-    # 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
-
-
-@invocation(
-    "denoise_latents",
-    title="Denoise Latents",
-    tags=["latents", "denoise", "txt2img", "t2i", "t2l", "img2img", "i2i", "l2l"],
-    category="latents",
-    version="1.5.3",
-)
-class DenoiseLatentsInvocation(BaseInvocation):
-    """Denoises noisy latents to decodable images"""
-
-    positive_conditioning: Union[ConditioningField, list[ConditioningField]] = InputField(
-        description=FieldDescriptions.positive_cond, input=Input.Connection, ui_order=0
-    )
-    negative_conditioning: Union[ConditioningField, list[ConditioningField]] = InputField(
-        description=FieldDescriptions.negative_cond, input=Input.Connection, ui_order=1
-    )
-    noise: Optional[LatentsField] = InputField(
-        default=None,
-        description=FieldDescriptions.noise,
-        input=Input.Connection,
-        ui_order=3,
-    )
-    steps: int = InputField(default=10, gt=0, description=FieldDescriptions.steps)
-    cfg_scale: Union[float, List[float]] = InputField(
-        default=7.5, description=FieldDescriptions.cfg_scale, title="CFG Scale"
-    )
-    denoising_start: float = InputField(
-        default=0.0,
-        ge=0,
-        le=1,
-        description=FieldDescriptions.denoising_start,
-    )
-    denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
-    scheduler: SCHEDULER_NAME_VALUES = InputField(
-        default="euler",
-        description=FieldDescriptions.scheduler,
-        ui_type=UIType.Scheduler,
-    )
-    unet: UNetField = InputField(
-        description=FieldDescriptions.unet,
-        input=Input.Connection,
-        title="UNet",
-        ui_order=2,
-    )
-    control: Optional[Union[ControlField, list[ControlField]]] = InputField(
-        default=None,
-        input=Input.Connection,
-        ui_order=5,
-    )
-    ip_adapter: Optional[Union[IPAdapterField, list[IPAdapterField]]] = InputField(
-        description=FieldDescriptions.ip_adapter,
-        title="IP-Adapter",
-        default=None,
-        input=Input.Connection,
-        ui_order=6,
-    )
-    t2i_adapter: Optional[Union[T2IAdapterField, list[T2IAdapterField]]] = InputField(
-        description=FieldDescriptions.t2i_adapter,
-        title="T2I-Adapter",
-        default=None,
-        input=Input.Connection,
-        ui_order=7,
-    )
-    cfg_rescale_multiplier: float = InputField(
-        title="CFG Rescale Multiplier", default=0, ge=0, lt=1, description=FieldDescriptions.cfg_rescale_multiplier
-    )
-    latents: Optional[LatentsField] = InputField(
-        default=None,
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-        ui_order=4,
-    )
-    denoise_mask: Optional[DenoiseMaskField] = InputField(
-        default=None,
-        description=FieldDescriptions.mask,
-        input=Input.Connection,
-        ui_order=8,
-    )
-
-    @field_validator("cfg_scale")
-    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:
-                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
-
-    @staticmethod
-    def _get_text_embeddings_and_masks(
-        cond_list: list[ConditioningField],
-        context: InvocationContext,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> tuple[Union[list[BasicConditioningInfo], list[SDXLConditioningInfo]], list[Optional[torch.Tensor]]]:
-        """Get the text embeddings and masks from the input conditioning fields."""
-        text_embeddings: Union[list[BasicConditioningInfo], list[SDXLConditioningInfo]] = []
-        text_embeddings_masks: list[Optional[torch.Tensor]] = []
-        for cond in cond_list:
-            cond_data = context.conditioning.load(cond.conditioning_name)
-            text_embeddings.append(cond_data.conditionings[0].to(device=device, dtype=dtype))
-
-            mask = cond.mask
-            if mask is not None:
-                mask = context.tensors.load(mask.tensor_name)
-            text_embeddings_masks.append(mask)
-
-        return text_embeddings, text_embeddings_masks
-
-    @staticmethod
-    def _preprocess_regional_prompt_mask(
-        mask: Optional[torch.Tensor], target_height: int, target_width: int, dtype: torch.dtype
-    ) -> torch.Tensor:
-        """Preprocess a regional prompt mask to match the target height and width.
-        If mask is None, returns a mask of all ones with the target height and width.
-        If mask is not None, resizes the mask to the target height and width using 'nearest' interpolation.
-
-        Returns:
-            torch.Tensor: The processed mask. shape: (1, 1, target_height, target_width).
-        """
-
-        if mask is None:
-            return torch.ones((1, 1, target_height, target_width), dtype=dtype)
-
-        mask = to_standard_float_mask(mask, out_dtype=dtype)
-
-        tf = torchvision.transforms.Resize(
-            (target_height, target_width), interpolation=torchvision.transforms.InterpolationMode.NEAREST
-        )
-
-        # Add a batch dimension to the mask, because torchvision expects shape (batch, channels, h, w).
-        mask = mask.unsqueeze(0)  # Shape: (1, h, w) -> (1, 1, h, w)
-        resized_mask = tf(mask)
-        return resized_mask
-
-    @staticmethod
-    def _concat_regional_text_embeddings(
-        text_conditionings: Union[list[BasicConditioningInfo], list[SDXLConditioningInfo]],
-        masks: Optional[list[Optional[torch.Tensor]]],
-        latent_height: int,
-        latent_width: int,
-        dtype: torch.dtype,
-    ) -> tuple[Union[BasicConditioningInfo, SDXLConditioningInfo], Optional[TextConditioningRegions]]:
-        """Concatenate regional text embeddings into a single embedding and track the region masks accordingly."""
-        if masks is None:
-            masks = [None] * len(text_conditionings)
-        assert len(text_conditionings) == len(masks)
-
-        is_sdxl = type(text_conditionings[0]) is SDXLConditioningInfo
-
-        all_masks_are_none = all(mask is None for mask in masks)
-
-        text_embedding = []
-        pooled_embedding = None
-        add_time_ids = None
-        cur_text_embedding_len = 0
-        processed_masks = []
-        embedding_ranges = []
-
-        for prompt_idx, text_embedding_info in enumerate(text_conditionings):
-            mask = masks[prompt_idx]
-
-            if is_sdxl:
-                # We choose a random SDXLConditioningInfo's pooled_embeds and add_time_ids here, with a preference for
-                # prompts without a mask. We prefer prompts without a mask, because they are more likely to contain
-                # global prompt information.  In an ideal case, there should be exactly one global prompt without a
-                # mask, but we don't enforce this.
-
-                # HACK(ryand): The fact that we have to choose a single pooled_embedding and add_time_ids here is a
-                # fundamental interface issue. The SDXL Compel nodes are not designed to be used in the way that we use
-                # them for regional prompting. Ideally, the DenoiseLatents invocation should accept a single
-                # pooled_embeds tensor and a list of standard text embeds with region masks. This change would be a
-                # pretty major breaking change to a popular node, so for now we use this hack.
-                if pooled_embedding is None or mask is None:
-                    pooled_embedding = text_embedding_info.pooled_embeds
-                if add_time_ids is None or mask is None:
-                    add_time_ids = text_embedding_info.add_time_ids
-
-            text_embedding.append(text_embedding_info.embeds)
-            if not all_masks_are_none:
-                embedding_ranges.append(
-                    Range(
-                        start=cur_text_embedding_len, end=cur_text_embedding_len + text_embedding_info.embeds.shape[1]
-                    )
-                )
-                processed_masks.append(
-                    DenoiseLatentsInvocation._preprocess_regional_prompt_mask(
-                        mask, latent_height, latent_width, dtype=dtype
-                    )
-                )
-
-            cur_text_embedding_len += text_embedding_info.embeds.shape[1]
-
-        text_embedding = torch.cat(text_embedding, dim=1)
-        assert len(text_embedding.shape) == 3  # batch_size, seq_len, token_len
-
-        regions = None
-        if not all_masks_are_none:
-            regions = TextConditioningRegions(
-                masks=torch.cat(processed_masks, dim=1),
-                ranges=embedding_ranges,
-            )
-
-        if is_sdxl:
-            return (
-                SDXLConditioningInfo(embeds=text_embedding, pooled_embeds=pooled_embedding, add_time_ids=add_time_ids),
-                regions,
-            )
-        return BasicConditioningInfo(embeds=text_embedding), regions
-
-    @staticmethod
-    def get_conditioning_data(
-        context: InvocationContext,
-        positive_conditioning_field: Union[ConditioningField, list[ConditioningField]],
-        negative_conditioning_field: Union[ConditioningField, list[ConditioningField]],
-        unet: UNet2DConditionModel,
-        latent_height: int,
-        latent_width: int,
-        cfg_scale: float | list[float],
-        steps: int,
-        cfg_rescale_multiplier: float,
-    ) -> TextConditioningData:
-        # Normalize positive_conditioning_field and negative_conditioning_field to lists.
-        cond_list = positive_conditioning_field
-        if not isinstance(cond_list, list):
-            cond_list = [cond_list]
-        uncond_list = negative_conditioning_field
-        if not isinstance(uncond_list, list):
-            uncond_list = [uncond_list]
-
-        cond_text_embeddings, cond_text_embedding_masks = DenoiseLatentsInvocation._get_text_embeddings_and_masks(
-            cond_list, context, unet.device, unet.dtype
-        )
-        uncond_text_embeddings, uncond_text_embedding_masks = DenoiseLatentsInvocation._get_text_embeddings_and_masks(
-            uncond_list, context, unet.device, unet.dtype
-        )
-
-        cond_text_embedding, cond_regions = DenoiseLatentsInvocation._concat_regional_text_embeddings(
-            text_conditionings=cond_text_embeddings,
-            masks=cond_text_embedding_masks,
-            latent_height=latent_height,
-            latent_width=latent_width,
-            dtype=unet.dtype,
-        )
-        uncond_text_embedding, uncond_regions = DenoiseLatentsInvocation._concat_regional_text_embeddings(
-            text_conditionings=uncond_text_embeddings,
-            masks=uncond_text_embedding_masks,
-            latent_height=latent_height,
-            latent_width=latent_width,
-            dtype=unet.dtype,
-        )
-
-        if isinstance(cfg_scale, list):
-            assert len(cfg_scale) == steps, "cfg_scale (list) must have the same length as the number of steps"
-
-        conditioning_data = TextConditioningData(
-            uncond_text=uncond_text_embedding,
-            cond_text=cond_text_embedding,
-            uncond_regions=uncond_regions,
-            cond_regions=cond_regions,
-            guidance_scale=cfg_scale,
-            guidance_rescale_multiplier=cfg_rescale_multiplier,
-        )
-        return conditioning_data
-
-    @staticmethod
-    def create_pipeline(
-        unet: UNet2DConditionModel,
-        scheduler: Scheduler,
-    ) -> StableDiffusionGeneratorPipeline:
-        class FakeVae:
-            class FakeVaeConfig:
-                def __init__(self) -> None:
-                    self.block_out_channels = [0]
-
-            def __init__(self) -> None:
-                self.config = FakeVae.FakeVaeConfig()
-
-        return StableDiffusionGeneratorPipeline(
-            vae=FakeVae(),  # TODO: oh...
-            text_encoder=None,
-            tokenizer=None,
-            unet=unet,
-            scheduler=scheduler,
-            safety_checker=None,
-            feature_extractor=None,
-            requires_safety_checker=False,
-        )
-
-    @staticmethod
-    def prep_control_data(
-        context: InvocationContext,
-        control_input: ControlField | list[ControlField] | None,
-        latents_shape: List[int],
-        exit_stack: ExitStack,
-        do_classifier_free_guidance: bool = True,
-    ) -> list[ControlNetData] | None:
-        # Normalize control_input to a list.
-        control_list: list[ControlField]
-        if isinstance(control_input, ControlField):
-            control_list = [control_input]
-        elif isinstance(control_input, list):
-            control_list = control_input
-        elif control_input is None:
-            control_list = []
-        else:
-            raise ValueError(f"Unexpected control_input type: {type(control_input)}")
-
-        if len(control_list) == 0:
-            return None
-
-        # Assuming fixed dimensional scaling of LATENT_SCALE_FACTOR.
-        _, _, latent_height, latent_width = latents_shape
-        control_height_resize = latent_height * LATENT_SCALE_FACTOR
-        control_width_resize = latent_width * LATENT_SCALE_FACTOR
-
-        controlnet_data: list[ControlNetData] = []
-        for control_info in control_list:
-            control_model = exit_stack.enter_context(context.models.load(control_info.control_model))
-            assert isinstance(control_model, ControlNetModel)
-
-            control_image_field = control_info.image
-            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?
-            #        and do real check for classifier_free_guidance?
-            # prepare_control_image should return torch.Tensor of shape(batch_size, 3, height, width)
-            control_image = prepare_control_image(
-                image=input_image,
-                do_classifier_free_guidance=do_classifier_free_guidance,
-                width=control_width_resize,
-                height=control_height_resize,
-                # batch_size=batch_size * num_images_per_prompt,
-                # num_images_per_prompt=num_images_per_prompt,
-                device=control_model.device,
-                dtype=control_model.dtype,
-                control_mode=control_info.control_mode,
-                resize_mode=control_info.resize_mode,
-            )
-            control_item = ControlNetData(
-                model=control_model,
-                image_tensor=control_image,
-                weight=control_info.control_weight,
-                begin_step_percent=control_info.begin_step_percent,
-                end_step_percent=control_info.end_step_percent,
-                control_mode=control_info.control_mode,
-                # any resizing needed should currently be happening in prepare_control_image(),
-                #    but adding resize_mode to ControlNetData in case needed in the future
-                resize_mode=control_info.resize_mode,
-            )
-            controlnet_data.append(control_item)
-            # MultiControlNetModel has been refactored out, just need list[ControlNetData]
-
-        return controlnet_data
-
-    def prep_ip_adapter_image_prompts(
-        self,
-        context: InvocationContext,
-        ip_adapters: List[IPAdapterField],
-    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
-        """Run the IPAdapter CLIPVisionModel, returning image prompt embeddings."""
-        image_prompts = []
-        for single_ip_adapter in ip_adapters:
-            with context.models.load(single_ip_adapter.ip_adapter_model) as ip_adapter_model:
-                assert isinstance(ip_adapter_model, IPAdapter)
-                image_encoder_model_info = context.models.load(single_ip_adapter.image_encoder_model)
-                # `single_ip_adapter.image` could be a list or a single ImageField. Normalize to a list here.
-                single_ipa_image_fields = single_ip_adapter.image
-                if not isinstance(single_ipa_image_fields, list):
-                    single_ipa_image_fields = [single_ipa_image_fields]
-
-                single_ipa_images = [context.images.get_pil(image.image_name) for image in single_ipa_image_fields]
-                with image_encoder_model_info as image_encoder_model:
-                    assert isinstance(image_encoder_model, CLIPVisionModelWithProjection)
-                    # Get image embeddings from CLIP and ImageProjModel.
-                    image_prompt_embeds, uncond_image_prompt_embeds = ip_adapter_model.get_image_embeds(
-                        single_ipa_images, image_encoder_model
-                    )
-                    image_prompts.append((image_prompt_embeds, uncond_image_prompt_embeds))
-
-        return image_prompts
-
-    def prep_ip_adapter_data(
-        self,
-        context: InvocationContext,
-        ip_adapters: List[IPAdapterField],
-        image_prompts: List[Tuple[torch.Tensor, torch.Tensor]],
-        exit_stack: ExitStack,
-        latent_height: int,
-        latent_width: int,
-        dtype: torch.dtype,
-    ) -> Optional[List[IPAdapterData]]:
-        """If IP-Adapter is enabled, then this function loads the requisite models and adds the image prompt conditioning data."""
-        ip_adapter_data_list = []
-        for single_ip_adapter, (image_prompt_embeds, uncond_image_prompt_embeds) in zip(
-            ip_adapters, image_prompts, strict=True
-        ):
-            ip_adapter_model = exit_stack.enter_context(context.models.load(single_ip_adapter.ip_adapter_model))
-
-            mask_field = single_ip_adapter.mask
-            mask = context.tensors.load(mask_field.tensor_name) if mask_field is not None else None
-            mask = self._preprocess_regional_prompt_mask(mask, latent_height, latent_width, dtype=dtype)
-
-            ip_adapter_data_list.append(
-                IPAdapterData(
-                    ip_adapter_model=ip_adapter_model,
-                    weight=single_ip_adapter.weight,
-                    target_blocks=single_ip_adapter.target_blocks,
-                    begin_step_percent=single_ip_adapter.begin_step_percent,
-                    end_step_percent=single_ip_adapter.end_step_percent,
-                    ip_adapter_conditioning=IPAdapterConditioningInfo(image_prompt_embeds, uncond_image_prompt_embeds),
-                    mask=mask,
-                )
-            )
-
-        return ip_adapter_data_list if len(ip_adapter_data_list) > 0 else None
-
-    def run_t2i_adapters(
-        self,
-        context: InvocationContext,
-        t2i_adapter: Optional[Union[T2IAdapterField, list[T2IAdapterField]]],
-        latents_shape: list[int],
-        do_classifier_free_guidance: bool,
-    ) -> Optional[list[T2IAdapterData]]:
-        if t2i_adapter is None:
-            return None
-
-        # Handle the possibility that t2i_adapter could be a list or a single T2IAdapterField.
-        if isinstance(t2i_adapter, T2IAdapterField):
-            t2i_adapter = [t2i_adapter]
-
-        if len(t2i_adapter) == 0:
-            return None
-
-        t2i_adapter_data = []
-        for t2i_adapter_field in t2i_adapter:
-            t2i_adapter_model_config = context.models.get_config(t2i_adapter_field.t2i_adapter_model.key)
-            t2i_adapter_loaded_model = context.models.load(t2i_adapter_field.t2i_adapter_model)
-            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_model_config.base == BaseModelType.StableDiffusion1:
-                max_unet_downscale = 8
-            elif t2i_adapter_model_config.base == BaseModelType.StableDiffusionXL:
-                max_unet_downscale = 4
-            else:
-                raise ValueError(f"Unexpected T2I-Adapter base model type: '{t2i_adapter_model_config.base}'.")
-
-            t2i_adapter_model: T2IAdapter
-            with t2i_adapter_loaded_model as t2i_adapter_model:
-                total_downscale_factor = t2i_adapter_model.total_downscale_factor
-
-                # Resize the T2I-Adapter input image.
-                # We select the resize dimensions so that after the T2I-Adapter's total_downscale_factor is applied, the
-                # result will match the latent image's dimensions after max_unet_downscale is applied.
-                t2i_input_height = latents_shape[2] // max_unet_downscale * total_downscale_factor
-                t2i_input_width = latents_shape[3] // max_unet_downscale * total_downscale_factor
-
-                # Note: We have hard-coded `do_classifier_free_guidance=False`. This is because we only want to prepare
-                # a single image. If CFG is enabled, we will duplicate the resultant tensor after applying the
-                # T2I-Adapter model.
-                #
-                # Note: We re-use the `prepare_control_image(...)` from ControlNet for T2I-Adapter, because it has many
-                # of the same requirements (e.g. preserving binary masks during resize).
-                t2i_image = prepare_control_image(
-                    image=image,
-                    do_classifier_free_guidance=False,
-                    width=t2i_input_width,
-                    height=t2i_input_height,
-                    num_channels=t2i_adapter_model.config["in_channels"],  # mypy treats this as a FrozenDict
-                    device=t2i_adapter_model.device,
-                    dtype=t2i_adapter_model.dtype,
-                    resize_mode=t2i_adapter_field.resize_mode,
-                )
-
-                adapter_state = t2i_adapter_model(t2i_image)
-
-            if do_classifier_free_guidance:
-                for idx, value in enumerate(adapter_state):
-                    adapter_state[idx] = torch.cat([value] * 2, dim=0)
-
-            t2i_adapter_data.append(
-                T2IAdapterData(
-                    adapter_state=adapter_state,
-                    weight=t2i_adapter_field.weight,
-                    begin_step_percent=t2i_adapter_field.begin_step_percent,
-                    end_step_percent=t2i_adapter_field.end_step_percent,
-                )
-            )
-
-        return t2i_adapter_data
-
-    # original idea by https://github.com/AmericanPresidentJimmyCarter
-    # TODO: research more for second order schedulers timesteps
-    @staticmethod
-    def init_scheduler(
-        scheduler: Union[Scheduler, ConfigMixin],
-        device: torch.device,
-        steps: int,
-        denoising_start: float,
-        denoising_end: float,
-        seed: int,
-    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]:
-        assert isinstance(scheduler, ConfigMixin)
-        if scheduler.config.get("cpu_only", False):
-            scheduler.set_timesteps(steps, device="cpu")
-            timesteps = scheduler.timesteps.to(device=device)
-        else:
-            scheduler.set_timesteps(steps, device=device)
-            timesteps = scheduler.timesteps
-
-        # skip greater order timesteps
-        _timesteps = timesteps[:: scheduler.order]
-
-        # get start timestep index
-        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_idx = len(list(filter(lambda ts: ts >= t_end_val, _timesteps[t_start_idx:])))
-
-        # apply order to indexes
-        t_start_idx *= scheduler.order
-        t_end_idx *= scheduler.order
-
-        init_timestep = timesteps[t_start_idx : t_start_idx + 1]
-        timesteps = timesteps[t_start_idx : t_start_idx + t_end_idx]
-
-        scheduler_step_kwargs: Dict[str, Any] = {}
-        scheduler_step_signature = inspect.signature(scheduler.step)
-        if "generator" in scheduler_step_signature.parameters:
-            # At some point, someone decided that schedulers that accept a generator should use the original seed with
-            # all bits flipped. I don't know the original rationale for this, but now we must keep it like this for
-            # reproducibility.
-            #
-            # These Invoke-supported schedulers accept a generator as of 2024-06-04:
-            #   - DDIMScheduler
-            #   - DDPMScheduler
-            #   - DPMSolverMultistepScheduler
-            #   - EulerAncestralDiscreteScheduler
-            #   - EulerDiscreteScheduler
-            #   - KDPM2AncestralDiscreteScheduler
-            #   - LCMScheduler
-            #   - TCDScheduler
-            scheduler_step_kwargs.update({"generator": torch.Generator(device=device).manual_seed(seed ^ 0xFFFFFFFF)})
-        if isinstance(scheduler, TCDScheduler):
-            scheduler_step_kwargs.update({"eta": 1.0})
-
-        return timesteps, init_timestep, scheduler_step_kwargs
-
-    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, False
-
-        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.tensors.load(self.denoise_mask.masked_latents_name)
-        else:
-            masked_latents = torch.where(mask < 0.5, 0.0, latents)
-
-        return 1 - mask, masked_latents, self.denoise_mask.gradient
-
-    @staticmethod
-    def prepare_noise_and_latents(
-        context: InvocationContext, noise_field: LatentsField | None, latents_field: LatentsField | None
-    ) -> Tuple[int, torch.Tensor | None, torch.Tensor]:
-        """Depending on the workflow, we expect different combinations of noise and latents to be provided. This
-        function handles preparing these values accordingly.
-
-        Expected workflows:
-        - Text-to-Image Denoising: `noise` is provided, `latents` is not. `latents` is initialized to zeros.
-        - Image-to-Image Denoising: `noise` and `latents` are both provided.
-        - Text-to-Image SDXL Refiner Denoising: `latents` is provided, `noise` is not.
-        - Image-to-Image SDXL Refiner Denoising: `latents` is provided, `noise` is not.
-
-        NOTE(ryand): I wrote this docstring, but I am not the original author of this code. There may be other workflows
-        I haven't considered.
-        """
-        noise = None
-        if noise_field is not None:
-            noise = context.tensors.load(noise_field.latents_name)
-
-        if latents_field is not None:
-            latents = context.tensors.load(latents_field.latents_name)
-        elif noise is not None:
-            latents = torch.zeros_like(noise)
-        else:
-            raise ValueError("'latents' or 'noise' must be provided!")
-
-        if noise is not None and noise.shape[1:] != latents.shape[1:]:
-            raise ValueError(f"Incompatible 'noise' and 'latents' shapes: {latents.shape=} {noise.shape=}")
-
-        # The seed comes from (in order of priority): the noise field, the latents field, or 0.
-        seed = 0
-        if noise_field is not None and noise_field.seed is not None:
-            seed = noise_field.seed
-        elif latents_field is not None and latents_field.seed is not None:
-            seed = latents_field.seed
-        else:
-            seed = 0
-
-        return seed, noise, latents
-
-    @torch.no_grad()
-    @SilenceWarnings()  # This quenches the NSFW nag from diffusers.
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        seed, noise, latents = self.prepare_noise_and_latents(context, self.noise, self.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.
-        t2i_adapter_data = self.run_t2i_adapters(
-            context,
-            self.t2i_adapter,
-            latents.shape,
-            do_classifier_free_guidance=True,
-        )
-
-        ip_adapters: List[IPAdapterField] = []
-        if self.ip_adapter is not None:
-            # ip_adapter could be a list or a single IPAdapterField. Normalize to a list here.
-            if isinstance(self.ip_adapter, list):
-                ip_adapters = self.ip_adapter
-            else:
-                ip_adapters = [self.ip_adapter]
-
-        # If there are IP adapters, the following line runs the adapters' CLIPVision image encoders to return
-        # a series of image conditioning embeddings. This is being done here rather than in the
-        # big model context below in order to use less VRAM on low-VRAM systems.
-        # The image prompts are then passed to prep_ip_adapter_data().
-        image_prompts = self.prep_ip_adapter_image_prompts(context=context, ip_adapters=ip_adapters)
-
-        # get the unet's config so that we can pass the base to sd_step_callback()
-        unet_config = context.models.get_config(self.unet.unet.key)
-
-        def step_callback(state: PipelineIntermediateState) -> None:
-            context.util.sd_step_callback(state, unet_config.base)
-
-        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
-            for lora in self.unet.loras:
-                lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, LoRAModelRaw)
-                yield (lora_info.model, lora.weight)
-                del lora_info
-            return
-
-        unet_info = context.models.load(self.unet.unet)
-        assert isinstance(unet_info.model, UNet2DConditionModel)
-        with (
-            ExitStack() as exit_stack,
-            unet_info.model_on_device() as (model_state_dict, unet),
-            ModelPatcher.apply_freeu(unet, self.unet.freeu_config),
-            set_seamless(unet, self.unet.seamless_axes),  # FIXME
-            # Apply the LoRA after unet has been moved to its target device for faster patching.
-            ModelPatcher.apply_lora_unet(
-                unet,
-                loras=_lora_loader(),
-                model_state_dict=model_state_dict,
-            ),
-        ):
-            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)
-            if mask is not None:
-                mask = mask.to(device=unet.device, dtype=unet.dtype)
-            if masked_latents is not None:
-                masked_latents = masked_latents.to(device=unet.device, dtype=unet.dtype)
-
-            scheduler = get_scheduler(
-                context=context,
-                scheduler_info=self.unet.scheduler,
-                scheduler_name=self.scheduler,
-                seed=seed,
-            )
-
-            pipeline = self.create_pipeline(unet, scheduler)
-
-            _, _, latent_height, latent_width = latents.shape
-            conditioning_data = self.get_conditioning_data(
-                context=context,
-                positive_conditioning_field=self.positive_conditioning,
-                negative_conditioning_field=self.negative_conditioning,
-                unet=unet,
-                latent_height=latent_height,
-                latent_width=latent_width,
-                cfg_scale=self.cfg_scale,
-                steps=self.steps,
-                cfg_rescale_multiplier=self.cfg_rescale_multiplier,
-            )
-
-            controlnet_data = self.prep_control_data(
-                context=context,
-                control_input=self.control,
-                latents_shape=latents.shape,
-                # do_classifier_free_guidance=(self.cfg_scale >= 1.0))
-                do_classifier_free_guidance=True,
-                exit_stack=exit_stack,
-            )
-
-            ip_adapter_data = self.prep_ip_adapter_data(
-                context=context,
-                ip_adapters=ip_adapters,
-                image_prompts=image_prompts,
-                exit_stack=exit_stack,
-                latent_height=latent_height,
-                latent_width=latent_width,
-                dtype=unet.dtype,
-            )
-
-            timesteps, init_timestep, scheduler_step_kwargs = self.init_scheduler(
-                scheduler,
-                device=unet.device,
-                steps=self.steps,
-                denoising_start=self.denoising_start,
-                denoising_end=self.denoising_end,
-                seed=seed,
-            )
-
-            result_latents = pipeline.latents_from_embeddings(
-                latents=latents,
-                timesteps=timesteps,
-                init_timestep=init_timestep,
-                noise=noise,
-                seed=seed,
-                mask=mask,
-                masked_latents=masked_latents,
-                is_gradient_mask=gradient_mask,
-                scheduler_step_kwargs=scheduler_step_kwargs,
-                conditioning_data=conditioning_data,
-                control_data=controlnet_data,
-                ip_adapter_data=ip_adapter_data,
-                t2i_adapter_data=t2i_adapter_data,
-                callback=step_callback,
-            )
-
-        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
-        result_latents = result_latents.to("cpu")
-        TorchDevice.empty_cache()
-
-        name = context.tensors.save(tensor=result_latents)
-        return LatentsOutput.build(latents_name=name, latents=result_latents, seed=None)

invokeai/app/invocations/fields.py

@@ -160,8 +160,6 @@ class FieldDescriptions:
     fp32 = "Whether or not to use full float32 precision"
     precision = "Precision to use"
     tiled = "Processing using overlapping tiles (reduce memory consumption)"
-    vae_tile_size = "The tile size for VAE tiling in pixels (image space). If set to 0, the default tile size for the "
-    "model will be used. Larger tile sizes generally produce better results at the cost of higher memory usage."
     detect_res = "Pixel resolution for detection"
     image_res = "Pixel resolution for output image"
     safe_mode = "Whether or not to use safe mode"

invokeai/app/invocations/ideal_size.py

@@ -1,65 +0,0 @@
-import math
-from typing import Tuple
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
-from invokeai.app.invocations.fields import FieldDescriptions, InputField, OutputField
-from invokeai.app.invocations.model import UNetField
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager.config import BaseModelType
-
-
-@invocation_output("ideal_size_output")
-class IdealSizeOutput(BaseInvocationOutput):
-    """Base class for invocations that output an image"""
-
-    width: int = OutputField(description="The ideal width of the image (in pixels)")
-    height: int = OutputField(description="The ideal height of the image (in pixels)")
-
-
-@invocation(
-    "ideal_size",
-    title="Ideal Size",
-    tags=["latents", "math", "ideal_size"],
-    version="1.0.3",
-)
-class IdealSizeInvocation(BaseInvocation):
-    """Calculates the ideal size for generation to avoid duplication"""
-
-    width: int = InputField(default=1024, description="Final image width")
-    height: int = InputField(default=576, description="Final image height")
-    unet: UNetField = InputField(default=None, description=FieldDescriptions.unet)
-    multiplier: float = InputField(
-        default=1.0,
-        description="Amount to multiply the model's dimensions by when calculating the ideal size (may result in "
-        "initial generation artifacts if too large)",
-    )
-
-    def trim_to_multiple_of(self, *args: int, multiple_of: int = LATENT_SCALE_FACTOR) -> Tuple[int, ...]:
-        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.key)
-        aspect = self.width / self.height
-        dimension: float = 512
-        if unet_config.base == BaseModelType.StableDiffusion2:
-            dimension = 768
-        elif unet_config.base == BaseModelType.StableDiffusionXL:
-            dimension = 1024
-        dimension = dimension * self.multiplier
-        min_dimension = math.floor(dimension * 0.5)
-        model_area = dimension * dimension  # hardcoded for now since all models are trained on square images
-
-        if aspect > 1.0:
-            init_height = max(min_dimension, math.sqrt(model_area / aspect))
-            init_width = init_height * aspect
-        else:
-            init_width = max(min_dimension, math.sqrt(model_area * aspect))
-            init_height = init_width / aspect
-
-        scaled_width, scaled_height = self.trim_to_multiple_of(
-            math.floor(init_width),
-            math.floor(init_height),
-        )
-
-        return IdealSizeOutput(width=scaled_width, height=scaled_height)

invokeai/app/invocations/image.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
+from pathlib import Path
 from typing import Literal, Optional
 
 import cv2
@@ -503,7 +504,7 @@ class ImageInverseLerpInvocation(BaseInvocation, WithMetadata, WithBoard):
     title="Blur NSFW Image",
     tags=["image", "nsfw"],
     category="image",
-    version="1.2.3",
+    version="1.2.2",
 )
 class ImageNSFWBlurInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Add blur to NSFW-flagged images"""
@@ -515,12 +516,23 @@ class ImageNSFWBlurInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         logger = context.logger
         logger.debug("Running NSFW checker")
-        image = SafetyChecker.blur_if_nsfw(image)
+        if SafetyChecker.has_nsfw_concept(image):
+            logger.info("A potentially NSFW image has been detected. Image will be blurred.")
+            blurry_image = image.filter(filter=ImageFilter.GaussianBlur(radius=32))
+            caution = self._get_caution_img()
+            blurry_image.paste(caution, (0, 0), caution)
+            image = blurry_image
 
         image_dto = context.images.save(image=image)
 
         return ImageOutput.build(image_dto)
 
+    def _get_caution_img(self) -> Image.Image:
+        import invokeai.app.assets.images as image_assets
+
+        caution = Image.open(Path(image_assets.__path__[0]) / "caution.png")
+        return caution.resize((caution.width // 2, caution.height // 2))
+
 
 @invocation(
     "img_watermark",

invokeai/app/invocations/image_to_latents.py

@@ -1,143 +0,0 @@
-from contextlib import nullcontext
-from functools import singledispatchmethod
-
-import einops
-import torch
-from diffusers.models.attention_processor import (
-    AttnProcessor2_0,
-    LoRAAttnProcessor2_0,
-    LoRAXFormersAttnProcessor,
-    XFormersAttnProcessor,
-)
-from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
-from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import DEFAULT_PRECISION, LATENT_SCALE_FACTOR
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    ImageField,
-    Input,
-    InputField,
-)
-from invokeai.app.invocations.model import VAEField
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager.load.load_base import LoadedModel
-from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params
-
-
-@invocation(
-    "i2l",
-    title="Image to Latents",
-    tags=["latents", "image", "vae", "i2l"],
-    category="latents",
-    version="1.1.0",
-)
-class ImageToLatentsInvocation(BaseInvocation):
-    """Encodes an image into latents."""
-
-    image: ImageField = InputField(
-        description="The image to encode",
-    )
-    vae: VAEField = InputField(
-        description=FieldDescriptions.vae,
-        input=Input.Connection,
-    )
-    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
-    # NOTE: tile_size = 0 is a special value. We use this rather than `int | None`, because the workflow UI does not
-    # offer a way to directly set None values.
-    tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size)
-    fp32: bool = InputField(default=DEFAULT_PRECISION == torch.float32, description=FieldDescriptions.fp32)
-
-    @staticmethod
-    def vae_encode(
-        vae_info: LoadedModel, upcast: bool, tiled: bool, image_tensor: torch.Tensor, tile_size: int = 0
-    ) -> torch.Tensor:
-        with vae_info as vae:
-            assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
-            orig_dtype = vae.dtype
-            if upcast:
-                vae.to(dtype=torch.float32)
-
-                use_torch_2_0_or_xformers = hasattr(vae.decoder, "mid_block") and isinstance(
-                    vae.decoder.mid_block.attentions[0].processor,
-                    (
-                        AttnProcessor2_0,
-                        XFormersAttnProcessor,
-                        LoRAXFormersAttnProcessor,
-                        LoRAAttnProcessor2_0,
-                    ),
-                )
-                # if xformers or torch_2_0 is used attention block does not need
-                # to be in float32 which can save lots of memory
-                if use_torch_2_0_or_xformers:
-                    vae.post_quant_conv.to(orig_dtype)
-                    vae.decoder.conv_in.to(orig_dtype)
-                    vae.decoder.mid_block.to(orig_dtype)
-                # else:
-                #    latents = latents.float()
-
-            else:
-                vae.to(dtype=torch.float16)
-                # latents = latents.half()
-
-            if tiled:
-                vae.enable_tiling()
-            else:
-                vae.disable_tiling()
-
-            tiling_context = nullcontext()
-            if tile_size > 0:
-                tiling_context = patch_vae_tiling_params(
-                    vae,
-                    tile_sample_min_size=tile_size,
-                    tile_latent_min_size=tile_size // LATENT_SCALE_FACTOR,
-                    tile_overlap_factor=0.25,
-                )
-
-            # non_noised_latents_from_image
-            image_tensor = image_tensor.to(device=vae.device, dtype=vae.dtype)
-            with torch.inference_mode(), tiling_context:
-                latents = ImageToLatentsInvocation._encode_to_tensor(vae, image_tensor)
-
-            latents = vae.config.scaling_factor * latents
-            latents = latents.to(dtype=orig_dtype)
-
-        return latents
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        image = context.images.get_pil(self.image.image_name)
-
-        vae_info = context.models.load(self.vae.vae)
-
-        image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
-        if image_tensor.dim() == 3:
-            image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
-
-        latents = self.vae_encode(
-            vae_info=vae_info, upcast=self.fp32, tiled=self.tiled, image_tensor=image_tensor, tile_size=self.tile_size
-        )
-
-        latents = latents.to("cpu")
-        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: 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:
-        assert isinstance(vae, torch.nn.Module)
-        latents: torch.FloatTensor = vae.encode(image_tensor).latents
-        return latents

invokeai/app/invocations/infill.py

@@ -42,16 +42,15 @@ class InfillImageProcessorInvocation(BaseInvocation, WithMetadata, WithBoard):
         """Infill the image with the specified method"""
         pass
 
-    def load_image(self) -> tuple[Image.Image, bool]:
+    def load_image(self, context: InvocationContext) -> tuple[Image.Image, bool]:
         """Process the image to have an alpha channel before being infilled"""
-        image = self._context.images.get_pil(self.image.image_name)
+        image = context.images.get_pil(self.image.image_name)
         has_alpha = True if image.mode == "RGBA" else False
         return image, has_alpha
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        self._context = context
         # Retrieve and process image to be infilled
-        input_image, has_alpha = self.load_image()
+        input_image, has_alpha = self.load_image(context)
 
         # If the input image has no alpha channel, return it
         if has_alpha is False:
@@ -134,12 +133,8 @@ class LaMaInfillInvocation(InfillImageProcessorInvocation):
     """Infills transparent areas of an image using the LaMa model"""
 
     def infill(self, image: Image.Image):
-        with self._context.models.load_remote_model(
-            source="https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
-            loader=LaMA.load_jit_model,
-        ) as model:
-            lama = LaMA(model)
-            return lama(image)
+        lama = LaMA()
+        return lama(image)
 
 
 @invocation("infill_cv2", title="CV2 Infill", tags=["image", "inpaint"], category="inpaint", version="1.2.2")

invokeai/app/invocations/ip_adapter.py

@@ -4,8 +4,20 @@ from typing import List, Literal, Optional, Union
 from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.fields import FieldDescriptions, InputField, OutputField, TensorField, UIType
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    Input,
+    InputField,
+    OutputField,
+    TensorField,
+    UIType,
+)
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.invocations.primitives import ImageField
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
@@ -24,7 +36,6 @@ class IPAdapterField(BaseModel):
     ip_adapter_model: ModelIdentifierField = Field(description="The IP-Adapter model to use.")
     image_encoder_model: ModelIdentifierField = Field(description="The name of the CLIP image encoder model.")
     weight: Union[float, List[float]] = Field(default=1, description="The weight given to the IP-Adapter.")
-    target_blocks: List[str] = Field(default=[], description="The IP Adapter blocks to apply")
     begin_step_percent: float = Field(
         default=0, ge=0, le=1, description="When the IP-Adapter is first applied (% of total steps)"
     )
@@ -58,7 +69,7 @@ class IPAdapterOutput(BaseInvocationOutput):
 CLIP_VISION_MODEL_MAP = {"ViT-H": "ip_adapter_sd_image_encoder", "ViT-G": "ip_adapter_sdxl_image_encoder"}
 
 
-@invocation("ip_adapter", title="IP-Adapter", tags=["ip_adapter", "control"], category="ip_adapter", version="1.4.1")
+@invocation("ip_adapter", title="IP-Adapter", tags=["ip_adapter", "control"], category="ip_adapter", version="1.3.0")
 class IPAdapterInvocation(BaseInvocation):
     """Collects IP-Adapter info to pass to other nodes."""
 
@@ -67,6 +78,7 @@ class IPAdapterInvocation(BaseInvocation):
     ip_adapter_model: ModelIdentifierField = InputField(
         description="The IP-Adapter model.",
         title="IP-Adapter Model",
+        input=Input.Direct,
         ui_order=-1,
         ui_type=UIType.IPAdapterModel,
     )
@@ -78,9 +90,6 @@ class IPAdapterInvocation(BaseInvocation):
     weight: Union[float, List[float]] = InputField(
         default=1, description="The weight given to the IP-Adapter", title="Weight"
     )
-    method: Literal["full", "style", "composition"] = InputField(
-        default="full", description="The method to apply the IP-Adapter"
-    )
     begin_step_percent: float = InputField(
         default=0, ge=0, le=1, description="When the IP-Adapter is first applied (% of total steps)"
     )
@@ -115,32 +124,12 @@ class IPAdapterInvocation(BaseInvocation):
 
         image_encoder_model = self._get_image_encoder(context, image_encoder_model_name)
 
-        if self.method == "style":
-            if ip_adapter_info.base == "sd-1":
-                target_blocks = ["up_blocks.1"]
-            elif ip_adapter_info.base == "sdxl":
-                target_blocks = ["up_blocks.0.attentions.1"]
-            else:
-                raise ValueError(f"Unsupported IP-Adapter base type: '{ip_adapter_info.base}'.")
-        elif self.method == "composition":
-            if ip_adapter_info.base == "sd-1":
-                target_blocks = ["down_blocks.2", "mid_block"]
-            elif ip_adapter_info.base == "sdxl":
-                target_blocks = ["down_blocks.2.attentions.1"]
-            else:
-                raise ValueError(f"Unsupported IP-Adapter base type: '{ip_adapter_info.base}'.")
-        elif self.method == "full":
-            target_blocks = ["block"]
-        else:
-            raise ValueError(f"Unexpected IP-Adapter method: '{self.method}'.")
-
         return IPAdapterOutput(
             ip_adapter=IPAdapterField(
                 image=self.image,
                 ip_adapter_model=self.ip_adapter_model,
                 image_encoder_model=ModelIdentifierField.from_config(image_encoder_model),
                 weight=self.weight,
-                target_blocks=target_blocks,
                 begin_step_percent=self.begin_step_percent,
                 end_step_percent=self.end_step_percent,
                 mask=self.mask,

invokeai/app/invocations/latents_to_image.py

@@ -1,121 +0,0 @@
-from contextlib import nullcontext
-
-import torch
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.models.attention_processor import (
-    AttnProcessor2_0,
-    LoRAAttnProcessor2_0,
-    LoRAXFormersAttnProcessor,
-    XFormersAttnProcessor,
-)
-from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
-from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import DEFAULT_PRECISION, LATENT_SCALE_FACTOR
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    Input,
-    InputField,
-    LatentsField,
-    WithBoard,
-    WithMetadata,
-)
-from invokeai.app.invocations.model import VAEField
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.stable_diffusion import set_seamless
-from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params
-from invokeai.backend.util.devices import TorchDevice
-
-
-@invocation(
-    "l2i",
-    title="Latents to Image",
-    tags=["latents", "image", "vae", "l2i"],
-    category="latents",
-    version="1.3.0",
-)
-class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
-    """Generates an image from latents."""
-
-    latents: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    vae: VAEField = InputField(
-        description=FieldDescriptions.vae,
-        input=Input.Connection,
-    )
-    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
-    # NOTE: tile_size = 0 is a special value. We use this rather than `int | None`, because the workflow UI does not
-    # offer a way to directly set None values.
-    tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size)
-    fp32: bool = InputField(default=DEFAULT_PRECISION == torch.float32, description=FieldDescriptions.fp32)
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        latents = context.tensors.load(self.latents.latents_name)
-
-        vae_info = context.models.load(self.vae.vae)
-        assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny))
-        with set_seamless(vae_info.model, self.vae.seamless_axes), vae_info as vae:
-            assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
-            latents = latents.to(vae.device)
-            if self.fp32:
-                vae.to(dtype=torch.float32)
-
-                use_torch_2_0_or_xformers = hasattr(vae.decoder, "mid_block") and isinstance(
-                    vae.decoder.mid_block.attentions[0].processor,
-                    (
-                        AttnProcessor2_0,
-                        XFormersAttnProcessor,
-                        LoRAXFormersAttnProcessor,
-                        LoRAAttnProcessor2_0,
-                    ),
-                )
-                # if xformers or torch_2_0 is used attention block does not need
-                # to be in float32 which can save lots of memory
-                if use_torch_2_0_or_xformers:
-                    vae.post_quant_conv.to(latents.dtype)
-                    vae.decoder.conv_in.to(latents.dtype)
-                    vae.decoder.mid_block.to(latents.dtype)
-                else:
-                    latents = latents.float()
-
-            else:
-                vae.to(dtype=torch.float16)
-                latents = latents.half()
-
-            if self.tiled or context.config.get().force_tiled_decode:
-                vae.enable_tiling()
-            else:
-                vae.disable_tiling()
-
-            tiling_context = nullcontext()
-            if self.tile_size > 0:
-                tiling_context = patch_vae_tiling_params(
-                    vae,
-                    tile_sample_min_size=self.tile_size,
-                    tile_latent_min_size=self.tile_size // LATENT_SCALE_FACTOR,
-                    tile_overlap_factor=0.25,
-                )
-
-            # clear memory as vae decode can request a lot
-            TorchDevice.empty_cache()
-
-            with torch.inference_mode(), tiling_context:
-                # copied from diffusers pipeline
-                latents = latents / vae.config.scaling_factor
-                image = vae.decode(latents, return_dict=False)[0]
-                image = (image / 2 + 0.5).clamp(0, 1)  # denormalize
-                # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-                np_image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-
-                image = VaeImageProcessor.numpy_to_pil(np_image)[0]
-
-        TorchDevice.empty_cache()
-
-        image_dto = context.images.save(image=image)
-
-        return ImageOutput.build(image_dto)

invokeai/app/invocations/mask.py

@@ -1,8 +1,7 @@
-import numpy as np
 import torch
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, InvocationContext, invocation
-from invokeai.app.invocations.fields import ImageField, InputField, TensorField, WithMetadata
+from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext, invocation
+from invokeai.app.invocations.fields import InputField, TensorField, WithMetadata
 from invokeai.app.invocations.primitives import MaskOutput
 
 
@@ -35,86 +34,3 @@ class RectangleMaskInvocation(BaseInvocation, WithMetadata):
             width=self.width,
             height=self.height,
         )
-
-
-@invocation(
-    "alpha_mask_to_tensor",
-    title="Alpha Mask to Tensor",
-    tags=["conditioning"],
-    category="conditioning",
-    version="1.0.0",
-    classification=Classification.Beta,
-)
-class AlphaMaskToTensorInvocation(BaseInvocation):
-    """Convert a mask image to a tensor. Opaque regions are 1 and transparent regions are 0."""
-
-    image: ImageField = InputField(description="The mask image to convert.")
-    invert: bool = InputField(default=False, description="Whether to invert the mask.")
-
-    def invoke(self, context: InvocationContext) -> MaskOutput:
-        image = context.images.get_pil(self.image.image_name)
-        mask = torch.zeros((1, image.height, image.width), dtype=torch.bool)
-        if self.invert:
-            mask[0] = torch.tensor(np.array(image)[:, :, 3] == 0, dtype=torch.bool)
-        else:
-            mask[0] = torch.tensor(np.array(image)[:, :, 3] > 0, dtype=torch.bool)
-
-        return MaskOutput(
-            mask=TensorField(tensor_name=context.tensors.save(mask)),
-            height=mask.shape[1],
-            width=mask.shape[2],
-        )
-
-
-@invocation(
-    "invert_tensor_mask",
-    title="Invert Tensor Mask",
-    tags=["conditioning"],
-    category="conditioning",
-    version="1.0.0",
-    classification=Classification.Beta,
-)
-class InvertTensorMaskInvocation(BaseInvocation):
-    """Inverts a tensor mask."""
-
-    mask: TensorField = InputField(description="The tensor mask to convert.")
-
-    def invoke(self, context: InvocationContext) -> MaskOutput:
-        mask = context.tensors.load(self.mask.tensor_name)
-        inverted = ~mask
-
-        return MaskOutput(
-            mask=TensorField(tensor_name=context.tensors.save(inverted)),
-            height=inverted.shape[1],
-            width=inverted.shape[2],
-        )
-
-
-@invocation(
-    "image_mask_to_tensor",
-    title="Image Mask to Tensor",
-    tags=["conditioning"],
-    category="conditioning",
-    version="1.0.0",
-)
-class ImageMaskToTensorInvocation(BaseInvocation, WithMetadata):
-    """Convert a mask image to a tensor. Converts the image to grayscale and uses thresholding at the specified value."""
-
-    image: ImageField = InputField(description="The mask image to convert.")
-    cutoff: int = InputField(ge=0, le=255, description="Cutoff (<)", default=128)
-    invert: bool = InputField(default=False, description="Whether to invert the mask.")
-
-    def invoke(self, context: InvocationContext) -> MaskOutput:
-        image = context.images.get_pil(self.image.image_name, mode="L")
-
-        mask = torch.zeros((1, image.height, image.width), dtype=torch.bool)
-        if self.invert:
-            mask[0] = torch.tensor(np.array(image)[:, :] >= self.cutoff, dtype=torch.bool)
-        else:
-            mask[0] = torch.tensor(np.array(image)[:, :] < self.cutoff, dtype=torch.bool)
-
-        return MaskOutput(
-            mask=TensorField(tensor_name=context.tensors.save(mask)),
-            height=mask.shape[1],
-            width=mask.shape[2],
-        )

invokeai/app/invocations/metadata.py

@@ -3,6 +3,7 @@ from typing import Any, Literal, Optional, Union
 from pydantic import BaseModel, ConfigDict, Field
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
+from invokeai.app.invocations.controlnet_image_processors import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     ImageField,
@@ -13,7 +14,6 @@ from invokeai.app.invocations.fields import (
 )
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES
 
 from ...version import __version__
 
@@ -36,7 +36,6 @@ class IPAdapterMetadataField(BaseModel):
     image: ImageField = Field(description="The IP-Adapter image prompt.")
     ip_adapter_model: ModelIdentifierField = Field(description="The IP-Adapter model.")
     clip_vision_model: Literal["ViT-H", "ViT-G"] = Field(description="The CLIP Vision model")
-    method: Literal["full", "style", "composition"] = Field(description="Method to apply IP Weights with")
     weight: Union[float, list[float]] = Field(description="The weight given to the IP-Adapter")
     begin_step_percent: float = Field(description="When the IP-Adapter is first applied (% of total steps)")
     end_step_percent: float = Field(description="When the IP-Adapter is last applied (% of total steps)")

invokeai/app/invocations/model.py

@@ -11,7 +11,6 @@ from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType,
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
-    Classification,
     invocation,
     invocation_output,
 )
@@ -94,46 +93,19 @@ class ModelLoaderOutput(UNetOutput, CLIPOutput, VAEOutput):
     pass
 
 
-@invocation_output("model_identifier_output")
-class ModelIdentifierOutput(BaseInvocationOutput):
-    """Model identifier output"""
-
-    model: ModelIdentifierField = OutputField(description="Model identifier", title="Model")
-
-
-@invocation(
-    "model_identifier",
-    title="Model identifier",
-    tags=["model"],
-    category="model",
-    version="1.0.0",
-    classification=Classification.Prototype,
-)
-class ModelIdentifierInvocation(BaseInvocation):
-    """Selects any model, outputting it its identifier. Be careful with this one! The identifier will be accepted as
-    input for any model, even if the model types don't match. If you connect this to a mismatched input, you'll get an
-    error."""
-
-    model: ModelIdentifierField = InputField(description="The model to select", title="Model")
-
-    def invoke(self, context: InvocationContext) -> ModelIdentifierOutput:
-        if not context.models.exists(self.model.key):
-            raise Exception(f"Unknown model {self.model.key}")
-
-        return ModelIdentifierOutput(model=self.model)
-
-
 @invocation(
     "main_model_loader",
     title="Main Model",
     tags=["model"],
     category="model",
-    version="1.0.3",
+    version="1.0.2",
 )
 class MainModelLoaderInvocation(BaseInvocation):
     """Loads a main model, outputting its submodels."""
 
-    model: ModelIdentifierField = InputField(description=FieldDescriptions.main_model, ui_type=UIType.MainModel)
+    model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.main_model, input=Input.Direct, ui_type=UIType.MainModel
+    )
     # TODO: precision?
 
     def invoke(self, context: InvocationContext) -> ModelLoaderOutput:
@@ -162,12 +134,12 @@ 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.3")
+@invocation("lora_loader", title="LoRA", tags=["model"], category="model", version="1.0.2")
 class LoRALoaderInvocation(BaseInvocation):
     """Apply selected lora to unet and text_encoder."""
 
     lora: ModelIdentifierField = InputField(
-        description=FieldDescriptions.lora_model, title="LoRA", ui_type=UIType.LoRAModel
+        description=FieldDescriptions.lora_model, input=Input.Direct, title="LoRA", ui_type=UIType.LoRAModel
     )
     weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
     unet: Optional[UNetField] = InputField(
@@ -218,75 +190,6 @@ class LoRALoaderInvocation(BaseInvocation):
         return output
 
 
-@invocation_output("lora_selector_output")
-class LoRASelectorOutput(BaseInvocationOutput):
-    """Model loader output"""
-
-    lora: LoRAField = OutputField(description="LoRA model and weight", title="LoRA")
-
-
-@invocation("lora_selector", title="LoRA Selector", tags=["model"], category="model", version="1.0.1")
-class LoRASelectorInvocation(BaseInvocation):
-    """Selects a LoRA model and weight."""
-
-    lora: ModelIdentifierField = InputField(
-        description=FieldDescriptions.lora_model, title="LoRA", ui_type=UIType.LoRAModel
-    )
-    weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
-
-    def invoke(self, context: InvocationContext) -> LoRASelectorOutput:
-        return LoRASelectorOutput(lora=LoRAField(lora=self.lora, weight=self.weight))
-
-
-@invocation("lora_collection_loader", title="LoRA Collection Loader", tags=["model"], category="model", version="1.0.0")
-class LoRACollectionLoader(BaseInvocation):
-    """Applies a collection of LoRAs to the provided UNet and CLIP models."""
-
-    loras: LoRAField | list[LoRAField] = InputField(
-        description="LoRA models and weights. May be a single LoRA or collection.", title="LoRAs"
-    )
-    unet: Optional[UNetField] = InputField(
-        default=None,
-        description=FieldDescriptions.unet,
-        input=Input.Connection,
-        title="UNet",
-    )
-    clip: Optional[CLIPField] = InputField(
-        default=None,
-        description=FieldDescriptions.clip,
-        input=Input.Connection,
-        title="CLIP",
-    )
-
-    def invoke(self, context: InvocationContext) -> LoRALoaderOutput:
-        output = LoRALoaderOutput()
-        loras = self.loras if isinstance(self.loras, list) else [self.loras]
-        added_loras: list[str] = []
-
-        for lora in loras:
-            if lora.lora.key in added_loras:
-                continue
-
-            if not context.models.exists(lora.lora.key):
-                raise Exception(f"Unknown lora: {lora.lora.key}!")
-
-            assert lora.lora.base in (BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2)
-
-            added_loras.append(lora.lora.key)
-
-            if self.unet is not None:
-                if output.unet is None:
-                    output.unet = self.unet.model_copy(deep=True)
-                output.unet.loras.append(lora)
-
-            if self.clip is not None:
-                if output.clip is None:
-                    output.clip = self.clip.model_copy(deep=True)
-                output.clip.loras.append(lora)
-
-        return output
-
-
 @invocation_output("sdxl_lora_loader_output")
 class SDXLLoRALoaderOutput(BaseInvocationOutput):
     """SDXL LoRA Loader Output"""
@@ -301,13 +204,13 @@ class SDXLLoRALoaderOutput(BaseInvocationOutput):
     title="SDXL LoRA",
     tags=["lora", "model"],
     category="model",
-    version="1.0.3",
+    version="1.0.2",
 )
 class SDXLLoRALoaderInvocation(BaseInvocation):
     """Apply selected lora to unet and text_encoder."""
 
     lora: ModelIdentifierField = InputField(
-        description=FieldDescriptions.lora_model, title="LoRA", ui_type=UIType.LoRAModel
+        description=FieldDescriptions.lora_model, input=Input.Direct, title="LoRA", ui_type=UIType.LoRAModel
     )
     weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
     unet: Optional[UNetField] = InputField(
@@ -376,78 +279,12 @@ class SDXLLoRALoaderInvocation(BaseInvocation):
         return output
 
 
-@invocation(
-    "sdxl_lora_collection_loader",
-    title="SDXL LoRA Collection Loader",
-    tags=["model"],
-    category="model",
-    version="1.0.0",
-)
-class SDXLLoRACollectionLoader(BaseInvocation):
-    """Applies a collection of SDXL LoRAs to the provided UNet and CLIP models."""
-
-    loras: LoRAField | list[LoRAField] = InputField(
-        description="LoRA models and weights. May be a single LoRA or collection.", title="LoRAs"
-    )
-    unet: Optional[UNetField] = InputField(
-        default=None,
-        description=FieldDescriptions.unet,
-        input=Input.Connection,
-        title="UNet",
-    )
-    clip: Optional[CLIPField] = InputField(
-        default=None,
-        description=FieldDescriptions.clip,
-        input=Input.Connection,
-        title="CLIP",
-    )
-    clip2: Optional[CLIPField] = InputField(
-        default=None,
-        description=FieldDescriptions.clip,
-        input=Input.Connection,
-        title="CLIP 2",
-    )
-
-    def invoke(self, context: InvocationContext) -> SDXLLoRALoaderOutput:
-        output = SDXLLoRALoaderOutput()
-        loras = self.loras if isinstance(self.loras, list) else [self.loras]
-        added_loras: list[str] = []
-
-        for lora in loras:
-            if lora.lora.key in added_loras:
-                continue
-
-            if not context.models.exists(lora.lora.key):
-                raise Exception(f"Unknown lora: {lora.lora.key}!")
-
-            assert lora.lora.base is BaseModelType.StableDiffusionXL
-
-            added_loras.append(lora.lora.key)
-
-            if self.unet is not None:
-                if output.unet is None:
-                    output.unet = self.unet.model_copy(deep=True)
-                output.unet.loras.append(lora)
-
-            if self.clip is not None:
-                if output.clip is None:
-                    output.clip = self.clip.model_copy(deep=True)
-                output.clip.loras.append(lora)
-
-            if self.clip2 is not None:
-                if output.clip2 is None:
-                    output.clip2 = self.clip2.model_copy(deep=True)
-                output.clip2.loras.append(lora)
-
-        return output
-
-
-@invocation("vae_loader", title="VAE", tags=["vae", "model"], category="model", version="1.0.3")
+@invocation("vae_loader", title="VAE", tags=["vae", "model"], category="model", version="1.0.2")
 class VAELoaderInvocation(BaseInvocation):
     """Loads a VAE model, outputting a VaeLoaderOutput"""
 
     vae_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.vae_model, title="VAE", ui_type=UIType.VAEModel
+        description=FieldDescriptions.vae_model, input=Input.Direct, title="VAE", ui_type=UIType.VAEModel
     )
 
     def invoke(self, context: InvocationContext) -> VAEOutput:

invokeai/app/invocations/noise.py

@@ -9,7 +9,7 @@ from invokeai.app.invocations.fields import FieldDescriptions, InputField, Laten
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.misc import SEED_MAX
 
-from ...backend.util.devices import TorchDevice
+from ...backend.util.devices import choose_torch_device, torch_dtype
 from .baseinvocation import (
     BaseInvocation,
     BaseInvocationOutput,
@@ -46,7 +46,7 @@ def get_noise(
             height // downsampling_factor,
             width // downsampling_factor,
         ],
-        dtype=TorchDevice.choose_torch_dtype(device=device),
+        dtype=torch_dtype(device),
         device=noise_device_type,
         generator=generator,
     ).to("cpu")
@@ -111,14 +111,14 @@ class NoiseInvocation(BaseInvocation):
 
     @field_validator("seed", mode="before")
     def modulo_seed(cls, v):
-        """Return the seed modulo (SEED_MAX + 1) to ensure it is within the valid range."""
+        """Returns the seed modulo (SEED_MAX + 1) to ensure it is within the valid range."""
         return v % (SEED_MAX + 1)
 
     def invoke(self, context: InvocationContext) -> NoiseOutput:
         noise = get_noise(
             width=self.width,
             height=self.height,
-            device=TorchDevice.choose_torch_device(),
+            device=choose_torch_device(),
             seed=self.seed,
             use_cpu=self.use_cpu,
         )

invokeai/app/invocations/resize_latents.py

@@ -1,103 +0,0 @@
-from typing import Literal
-
-import torch
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    Input,
-    InputField,
-    LatentsField,
-)
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.util.devices import TorchDevice
-
-LATENTS_INTERPOLATION_MODE = Literal["nearest", "linear", "bilinear", "bicubic", "trilinear", "area", "nearest-exact"]
-
-
-@invocation(
-    "lresize",
-    title="Resize Latents",
-    tags=["latents", "resize"],
-    category="latents",
-    version="1.0.2",
-)
-class ResizeLatentsInvocation(BaseInvocation):
-    """Resizes latents to explicit width/height (in pixels). Provided dimensions are floor-divided by 8."""
-
-    latents: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    width: int = InputField(
-        ge=64,
-        multiple_of=LATENT_SCALE_FACTOR,
-        description=FieldDescriptions.width,
-    )
-    height: int = InputField(
-        ge=64,
-        multiple_of=LATENT_SCALE_FACTOR,
-        description=FieldDescriptions.width,
-    )
-    mode: LATENTS_INTERPOLATION_MODE = InputField(default="bilinear", description=FieldDescriptions.interp_mode)
-    antialias: bool = InputField(default=False, description=FieldDescriptions.torch_antialias)
-
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.tensors.load(self.latents.latents_name)
-        device = TorchDevice.choose_torch_device()
-
-        resized_latents = torch.nn.functional.interpolate(
-            latents.to(device),
-            size=(self.height // LATENT_SCALE_FACTOR, self.width // LATENT_SCALE_FACTOR),
-            mode=self.mode,
-            antialias=self.antialias if self.mode in ["bilinear", "bicubic"] else False,
-        )
-
-        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
-        resized_latents = resized_latents.to("cpu")
-
-        TorchDevice.empty_cache()
-
-        name = context.tensors.save(tensor=resized_latents)
-        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)
-
-
-@invocation(
-    "lscale",
-    title="Scale Latents",
-    tags=["latents", "resize"],
-    category="latents",
-    version="1.0.2",
-)
-class ScaleLatentsInvocation(BaseInvocation):
-    """Scales latents by a given factor."""
-
-    latents: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    scale_factor: float = InputField(gt=0, description=FieldDescriptions.scale_factor)
-    mode: LATENTS_INTERPOLATION_MODE = InputField(default="bilinear", description=FieldDescriptions.interp_mode)
-    antialias: bool = InputField(default=False, description=FieldDescriptions.torch_antialias)
-
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = context.tensors.load(self.latents.latents_name)
-
-        device = TorchDevice.choose_torch_device()
-
-        # resizing
-        resized_latents = torch.nn.functional.interpolate(
-            latents.to(device),
-            scale_factor=self.scale_factor,
-            mode=self.mode,
-            antialias=self.antialias if self.mode in ["bilinear", "bicubic"] else False,
-        )
-
-        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
-        resized_latents = resized_latents.to("cpu")
-        TorchDevice.empty_cache()
-
-        name = context.tensors.save(tensor=resized_latents)
-        return LatentsOutput.build(latents_name=name, latents=resized_latents, seed=self.latents.seed)

invokeai/app/invocations/scheduler.py

@@ -1,34 +0,0 @@
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    InputField,
-    OutputField,
-    UIType,
-)
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
-
-
-@invocation_output("scheduler_output")
-class SchedulerOutput(BaseInvocationOutput):
-    scheduler: SCHEDULER_NAME_VALUES = OutputField(description=FieldDescriptions.scheduler, ui_type=UIType.Scheduler)
-
-
-@invocation(
-    "scheduler",
-    title="Scheduler",
-    tags=["scheduler"],
-    category="latents",
-    version="1.0.0",
-)
-class SchedulerInvocation(BaseInvocation):
-    """Selects a scheduler."""
-
-    scheduler: SCHEDULER_NAME_VALUES = InputField(
-        default="euler",
-        description=FieldDescriptions.scheduler,
-        ui_type=UIType.Scheduler,
-    )
-
-    def invoke(self, context: InvocationContext) -> SchedulerOutput:
-        return SchedulerOutput(scheduler=self.scheduler)

invokeai/app/invocations/sdxl.py

@@ -1,4 +1,4 @@
-from invokeai.app.invocations.fields import FieldDescriptions, InputField, OutputField, UIType
+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
 
@@ -30,12 +30,12 @@ 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.3")
+@invocation("sdxl_model_loader", title="SDXL Main Model", tags=["model", "sdxl"], category="model", version="1.0.2")
 class SDXLModelLoaderInvocation(BaseInvocation):
     """Loads an sdxl base model, outputting its submodels."""
 
     model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.sdxl_main_model, ui_type=UIType.SDXLMainModel
+        description=FieldDescriptions.sdxl_main_model, input=Input.Direct, ui_type=UIType.SDXLMainModel
     )
     # TODO: precision?
 
@@ -67,13 +67,13 @@ class SDXLModelLoaderInvocation(BaseInvocation):
     title="SDXL Refiner Model",
     tags=["model", "sdxl", "refiner"],
     category="model",
-    version="1.0.3",
+    version="1.0.2",
 )
 class SDXLRefinerModelLoaderInvocation(BaseInvocation):
     """Loads an sdxl refiner model, outputting its submodels."""
 
     model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.sdxl_refiner_model, ui_type=UIType.SDXLRefinerModel
+        description=FieldDescriptions.sdxl_refiner_model, input=Input.Direct, ui_type=UIType.SDXLRefinerModel
     )
     # TODO: precision?
 

invokeai/app/invocations/t2i_adapter.py

@@ -8,11 +8,11 @@ from invokeai.app.invocations.baseinvocation import (
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import FieldDescriptions, ImageField, InputField, OutputField, UIType
+from invokeai.app.invocations.controlnet_image_processors import CONTROLNET_RESIZE_VALUES
+from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField, OutputField, UIType
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.invocations.util import validate_begin_end_step, validate_weights
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.controlnet_utils import CONTROLNET_RESIZE_VALUES
 
 
 class T2IAdapterField(BaseModel):
@@ -45,7 +45,7 @@ class T2IAdapterOutput(BaseInvocationOutput):
 
 
 @invocation(
-    "t2i_adapter", title="T2I-Adapter", tags=["t2i_adapter", "control"], category="t2i_adapter", version="1.0.3"
+    "t2i_adapter", title="T2I-Adapter", tags=["t2i_adapter", "control"], category="t2i_adapter", version="1.0.2"
 )
 class T2IAdapterInvocation(BaseInvocation):
     """Collects T2I-Adapter info to pass to other nodes."""
@@ -55,6 +55,7 @@ class T2IAdapterInvocation(BaseInvocation):
     t2i_adapter_model: ModelIdentifierField = InputField(
         description="The T2I-Adapter model.",
         title="T2I-Adapter Model",
+        input=Input.Direct,
         ui_order=-1,
         ui_type=UIType.T2IAdapterModel,
     )

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

@@ -1,282 +0,0 @@
-import copy
-from contextlib import ExitStack
-from typing import Iterator, Tuple
-
-import torch
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from pydantic import field_validator
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
-from invokeai.app.invocations.controlnet_image_processors import ControlField
-from invokeai.app.invocations.denoise_latents import DenoiseLatentsInvocation, get_scheduler
-from invokeai.app.invocations.fields import (
-    ConditioningField,
-    FieldDescriptions,
-    Input,
-    InputField,
-    LatentsField,
-    UIType,
-)
-from invokeai.app.invocations.model import UNetField
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState
-from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
-    MultiDiffusionPipeline,
-    MultiDiffusionRegionConditioning,
-)
-from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
-from invokeai.backend.tiles.tiles import (
-    calc_tiles_min_overlap,
-)
-from invokeai.backend.tiles.utils import TBLR
-from invokeai.backend.util.devices import TorchDevice
-
-
-def crop_controlnet_data(control_data: ControlNetData, latent_region: TBLR) -> ControlNetData:
-    """Crop a ControlNetData object to a region."""
-    # Create a shallow copy of the control_data object.
-    control_data_copy = copy.copy(control_data)
-    # The ControlNet reference image is the only attribute that needs to be cropped.
-    control_data_copy.image_tensor = control_data.image_tensor[
-        :,
-        :,
-        latent_region.top * LATENT_SCALE_FACTOR : latent_region.bottom * LATENT_SCALE_FACTOR,
-        latent_region.left * LATENT_SCALE_FACTOR : latent_region.right * LATENT_SCALE_FACTOR,
-    ]
-    return control_data_copy
-
-
-@invocation(
-    "tiled_multi_diffusion_denoise_latents",
-    title="Tiled Multi-Diffusion Denoise Latents",
-    tags=["upscale", "denoise"],
-    category="latents",
-    classification=Classification.Beta,
-    version="1.0.0",
-)
-class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
-    """Tiled Multi-Diffusion denoising.
-
-    This node handles automatically tiling the input image, and is primarily intended for global refinement of images
-    in tiled upscaling workflows. Future Multi-Diffusion nodes should allow the user to specify custom regions with
-    different parameters for each region to harness the full power of Multi-Diffusion.
-
-    This node has a similar interface to the `DenoiseLatents` node, but it has a reduced feature set (no IP-Adapter,
-    T2I-Adapter, masking, etc.).
-    """
-
-    positive_conditioning: ConditioningField = InputField(
-        description=FieldDescriptions.positive_cond, input=Input.Connection
-    )
-    negative_conditioning: ConditioningField = InputField(
-        description=FieldDescriptions.negative_cond, input=Input.Connection
-    )
-    noise: LatentsField | None = InputField(
-        default=None,
-        description=FieldDescriptions.noise,
-        input=Input.Connection,
-    )
-    latents: LatentsField | None = InputField(
-        default=None,
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    tile_height: int = InputField(
-        default=1024, gt=0, multiple_of=LATENT_SCALE_FACTOR, description="Height of the tiles in image space."
-    )
-    tile_width: int = InputField(
-        default=1024, gt=0, multiple_of=LATENT_SCALE_FACTOR, description="Width of the tiles in image space."
-    )
-    tile_overlap: int = InputField(
-        default=32,
-        multiple_of=LATENT_SCALE_FACTOR,
-        gt=0,
-        description="The overlap between adjacent tiles in pixel space. (Of course, tile merging is applied in latent "
-        "space.) Tiles will be cropped during merging (if necessary) to ensure that they overlap by exactly this "
-        "amount.",
-    )
-    steps: int = InputField(default=18, gt=0, description=FieldDescriptions.steps)
-    cfg_scale: float | list[float] = InputField(default=6.0, description=FieldDescriptions.cfg_scale, title="CFG Scale")
-    denoising_start: float = InputField(
-        default=0.0,
-        ge=0,
-        le=1,
-        description=FieldDescriptions.denoising_start,
-    )
-    denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
-    scheduler: SCHEDULER_NAME_VALUES = InputField(
-        default="euler",
-        description=FieldDescriptions.scheduler,
-        ui_type=UIType.Scheduler,
-    )
-    unet: UNetField = InputField(
-        description=FieldDescriptions.unet,
-        input=Input.Connection,
-        title="UNet",
-    )
-    cfg_rescale_multiplier: float = InputField(
-        title="CFG Rescale Multiplier", default=0, ge=0, lt=1, description=FieldDescriptions.cfg_rescale_multiplier
-    )
-    control: ControlField | list[ControlField] | None = InputField(
-        default=None,
-        input=Input.Connection,
-    )
-
-    @field_validator("cfg_scale")
-    def ge_one(cls, v: list[float] | float) -> list[float] | float:
-        """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
-
-    @staticmethod
-    def create_pipeline(
-        unet: UNet2DConditionModel,
-        scheduler: SchedulerMixin,
-    ) -> MultiDiffusionPipeline:
-        # TODO(ryand): Get rid of this FakeVae hack.
-        class FakeVae:
-            class FakeVaeConfig:
-                def __init__(self) -> None:
-                    self.block_out_channels = [0]
-
-            def __init__(self) -> None:
-                self.config = FakeVae.FakeVaeConfig()
-
-        return MultiDiffusionPipeline(
-            vae=FakeVae(),
-            text_encoder=None,
-            tokenizer=None,
-            unet=unet,
-            scheduler=scheduler,
-            safety_checker=None,
-            feature_extractor=None,
-            requires_safety_checker=False,
-        )
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        # Convert tile image-space dimensions to latent-space dimensions.
-        latent_tile_height = self.tile_height // LATENT_SCALE_FACTOR
-        latent_tile_width = self.tile_width // LATENT_SCALE_FACTOR
-        latent_tile_overlap = self.tile_overlap // LATENT_SCALE_FACTOR
-
-        seed, noise, latents = DenoiseLatentsInvocation.prepare_noise_and_latents(context, self.noise, self.latents)
-        _, _, latent_height, latent_width = latents.shape
-
-        # Calculate the tile locations to cover the latent-space image.
-        tiles = calc_tiles_min_overlap(
-            image_height=latent_height,
-            image_width=latent_width,
-            tile_height=latent_tile_height,
-            tile_width=latent_tile_width,
-            min_overlap=latent_tile_overlap,
-        )
-
-        # Get the unet's config so that we can pass the base to sd_step_callback().
-        unet_config = context.models.get_config(self.unet.unet.key)
-
-        def step_callback(state: PipelineIntermediateState) -> None:
-            context.util.sd_step_callback(state, unet_config.base)
-
-        # Prepare an iterator that yields the UNet's LoRA models and their weights.
-        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
-            for lora in self.unet.loras:
-                lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, LoRAModelRaw)
-                yield (lora_info.model, lora.weight)
-                del lora_info
-
-        # Load the UNet model.
-        unet_info = context.models.load(self.unet.unet)
-
-        with ExitStack() as exit_stack, unet_info as unet, 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)
-            scheduler = get_scheduler(
-                context=context,
-                scheduler_info=self.unet.scheduler,
-                scheduler_name=self.scheduler,
-                seed=seed,
-            )
-            pipeline = self.create_pipeline(unet=unet, scheduler=scheduler)
-
-            # Prepare the prompt conditioning data. The same prompt conditioning is applied to all tiles.
-            conditioning_data = DenoiseLatentsInvocation.get_conditioning_data(
-                context=context,
-                positive_conditioning_field=self.positive_conditioning,
-                negative_conditioning_field=self.negative_conditioning,
-                unet=unet,
-                latent_height=latent_tile_height,
-                latent_width=latent_tile_width,
-                cfg_scale=self.cfg_scale,
-                steps=self.steps,
-                cfg_rescale_multiplier=self.cfg_rescale_multiplier,
-            )
-
-            controlnet_data = DenoiseLatentsInvocation.prep_control_data(
-                context=context,
-                control_input=self.control,
-                latents_shape=list(latents.shape),
-                # do_classifier_free_guidance=(self.cfg_scale >= 1.0))
-                do_classifier_free_guidance=True,
-                exit_stack=exit_stack,
-            )
-
-            # Split the controlnet_data into tiles.
-            # controlnet_data_tiles[t][c] is the c'th control data for the t'th tile.
-            controlnet_data_tiles: list[list[ControlNetData]] = []
-            for tile in tiles:
-                tile_controlnet_data = [crop_controlnet_data(cn, tile.coords) for cn in controlnet_data or []]
-                controlnet_data_tiles.append(tile_controlnet_data)
-
-            # Prepare the MultiDiffusionRegionConditioning list.
-            multi_diffusion_conditioning: list[MultiDiffusionRegionConditioning] = []
-            for tile, tile_controlnet_data in zip(tiles, controlnet_data_tiles, strict=True):
-                multi_diffusion_conditioning.append(
-                    MultiDiffusionRegionConditioning(
-                        region=tile,
-                        text_conditioning_data=conditioning_data,
-                        control_data=tile_controlnet_data,
-                    )
-                )
-
-            timesteps, init_timestep, scheduler_step_kwargs = DenoiseLatentsInvocation.init_scheduler(
-                scheduler,
-                device=unet.device,
-                steps=self.steps,
-                denoising_start=self.denoising_start,
-                denoising_end=self.denoising_end,
-                seed=seed,
-            )
-
-            # Run Multi-Diffusion denoising.
-            result_latents = pipeline.multi_diffusion_denoise(
-                multi_diffusion_conditioning=multi_diffusion_conditioning,
-                target_overlap=latent_tile_overlap,
-                latents=latents,
-                scheduler_step_kwargs=scheduler_step_kwargs,
-                noise=noise,
-                timesteps=timesteps,
-                init_timestep=init_timestep,
-                callback=step_callback,
-            )
-
-        result_latents = result_latents.to("cpu")
-        # TODO(ryand): I copied this from DenoiseLatentsInvocation. I'm not sure if it's actually important.
-        TorchDevice.empty_cache()
-
-        name = context.tensors.save(tensor=result_latents)
-        return LatentsOutput.build(latents_name=name, latents=result_latents, seed=None)

invokeai/app/invocations/upscale.py

@@ -1,16 +1,20 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) & the InvokeAI Team
+from pathlib import Path
 from typing import Literal
 
 import cv2
 import numpy as np
+import torch
 from PIL import Image
 from pydantic import ConfigDict
 
 from invokeai.app.invocations.fields import ImageField
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.util.download_with_progress import download_with_progress_bar
 from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet
 from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGAN
+from invokeai.backend.util.devices import choose_torch_device
 
 from .baseinvocation import BaseInvocation, invocation
 from .fields import InputField, WithBoard, WithMetadata
@@ -31,6 +35,9 @@ ESRGAN_MODEL_URLS: dict[str, str] = {
     "RealESRGAN_x2plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
 }
 
+if choose_torch_device() == torch.device("mps"):
+    from torch import mps
+
 
 @invocation("esrgan", title="Upscale (RealESRGAN)", tags=["esrgan", "upscale"], category="esrgan", version="1.3.2")
 class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
@@ -49,6 +56,7 @@ class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         rrdbnet_model = None
         netscale = None
+        esrgan_model_path = None
 
         if self.model_name in [
             "RealESRGAN_x4plus.pth",
@@ -91,25 +99,30 @@ class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
             context.logger.error(msg)
             raise ValueError(msg)
 
-        loadnet = context.models.load_remote_model(
-            source=ESRGAN_MODEL_URLS[self.model_name],
+        esrgan_model_path = Path(context.config.get().models_path, f"core/upscaling/realesrgan/{self.model_name}")
+
+        # Downloads the ESRGAN model if it doesn't already exist
+        download_with_progress_bar(
+            name=self.model_name, url=ESRGAN_MODEL_URLS[self.model_name], dest_path=esrgan_model_path
         )
 
-        with loadnet as loadnet_model:
-            upscaler = RealESRGAN(
-                scale=netscale,
-                loadnet=loadnet_model,
-                model=rrdbnet_model,
-                half=False,
-                tile=self.tile_size,
-            )
+        upscaler = RealESRGAN(
+            scale=netscale,
+            model_path=esrgan_model_path,
+            model=rrdbnet_model,
+            half=False,
+            tile=self.tile_size,
+        )
 
-            # prepare image - Real-ESRGAN uses cv2 internally, and cv2 uses BGR vs RGB for PIL
-            # TODO: This strips the alpha... is that okay?
-            cv2_image = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
-            upscaled_image = upscaler.upscale(cv2_image)
+        # prepare image - Real-ESRGAN uses cv2 internally, and cv2 uses BGR vs RGB for PIL
+        # TODO: This strips the alpha... is that okay?
+        cv2_image = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
+        upscaled_image = upscaler.upscale(cv2_image)
+        pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA")
 
-            pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA")
+        torch.cuda.empty_cache()
+        if choose_torch_device() == torch.device("mps"):
+            mps.empty_cache()
 
         image_dto = context.images.save(image=pil_image)
 

invokeai/app/services/board_records/board_records_base.py

@@ -40,12 +40,16 @@ class BoardRecordStorageBase(ABC):
 
     @abstractmethod
     def get_many(
-        self, offset: int = 0, limit: int = 10, include_archived: bool = False
+        self,
+        offset: int = 0,
+        limit: int = 10,
     ) -> OffsetPaginatedResults[BoardRecord]:
         """Gets many board records."""
         pass
 
     @abstractmethod
-    def get_all(self, include_archived: bool = False) -> list[BoardRecord]:
+    def get_all(
+        self,
+    ) -> list[BoardRecord]:
         """Gets all board records."""
         pass

invokeai/app/services/board_records/board_records_common.py

@@ -22,8 +22,6 @@ class BoardRecord(BaseModelExcludeNull):
     """The updated timestamp of the image."""
     cover_image_name: Optional[str] = Field(default=None, description="The name of the cover image of the board.")
     """The name of the cover image of the board."""
-    archived: bool = Field(description="Whether or not the board is archived.")
-    """Whether or not the board is archived."""
 
 
 def deserialize_board_record(board_dict: dict) -> BoardRecord:
@@ -37,7 +35,6 @@ def deserialize_board_record(board_dict: dict) -> BoardRecord:
     created_at = board_dict.get("created_at", get_iso_timestamp())
     updated_at = board_dict.get("updated_at", get_iso_timestamp())
     deleted_at = board_dict.get("deleted_at", get_iso_timestamp())
-    archived = board_dict.get("archived", False)
 
     return BoardRecord(
         board_id=board_id,
@@ -46,14 +43,12 @@ def deserialize_board_record(board_dict: dict) -> BoardRecord:
         created_at=created_at,
         updated_at=updated_at,
         deleted_at=deleted_at,
-        archived=archived,
     )
 
 
 class BoardChanges(BaseModel, extra="forbid"):
     board_name: Optional[str] = Field(default=None, description="The board's new name.")
     cover_image_name: Optional[str] = Field(default=None, description="The name of the board's new cover image.")
-    archived: Optional[bool] = Field(default=None, description="Whether or not the board is archived")
 
 
 class BoardRecordNotFoundException(Exception):

invokeai/app/services/board_records/board_records_sqlite.py

@@ -125,17 +125,6 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
                     (changes.cover_image_name, board_id),
                 )
 
-            # Change the archived status of a board
-            if changes.archived is not None:
-                self._cursor.execute(
-                    """--sql
-                    UPDATE boards
-                    SET archived = ?
-                    WHERE board_id = ?;
-                    """,
-                    (changes.archived, board_id),
-                )
-
             self._conn.commit()
         except sqlite3.Error as e:
             self._conn.rollback()
@@ -145,49 +134,35 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
         return self.get(board_id)
 
     def get_many(
-        self, offset: int = 0, limit: int = 10, include_archived: bool = False
+        self,
+        offset: int = 0,
+        limit: int = 10,
     ) -> OffsetPaginatedResults[BoardRecord]:
         try:
             self._lock.acquire()
 
-            # Build base query
-            base_query = """
+            # Get all the boards
+            self._cursor.execute(
+                """--sql
                 SELECT *
                 FROM boards
-                {archived_filter}
                 ORDER BY created_at DESC
                 LIMIT ? OFFSET ?;
-            """
-
-            # Determine archived filter condition
-            if include_archived:
-                archived_filter = ""
-            else:
-                archived_filter = "WHERE archived = 0"
-
-            final_query = base_query.format(archived_filter=archived_filter)
-
-            # Execute query to fetch boards
-            self._cursor.execute(final_query, (limit, offset))
+                """,
+                (limit, offset),
+            )
 
             result = cast(list[sqlite3.Row], self._cursor.fetchall())
             boards = [deserialize_board_record(dict(r)) for r in result]
 
-            # Determine count query
-            if include_archived:
-                count_query = """
-                    SELECT COUNT(*)
-                    FROM boards;
+            # Get the total number of boards
+            self._cursor.execute(
+                """--sql
+                SELECT COUNT(*)
+                FROM boards
+                WHERE 1=1;
                 """
-            else:
-                count_query = """
-                    SELECT COUNT(*)
-                    FROM boards
-                    WHERE archived = 0;
-                """
-
-            # Execute count query
-            self._cursor.execute(count_query)
+            )
 
             count = cast(int, self._cursor.fetchone()[0])
 
@@ -199,25 +174,20 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
         finally:
             self._lock.release()
 
-    def get_all(self, include_archived: bool = False) -> list[BoardRecord]:
+    def get_all(
+        self,
+    ) -> list[BoardRecord]:
         try:
             self._lock.acquire()
 
-            base_query = """
+            # Get all the boards
+            self._cursor.execute(
+                """--sql
                 SELECT *
                 FROM boards
-                {archived_filter}
                 ORDER BY created_at DESC
-            """
-
-            if include_archived:
-                archived_filter = ""
-            else:
-                archived_filter = "WHERE archived = 0"
-
-            final_query = base_query.format(archived_filter=archived_filter)
-
-            self._cursor.execute(final_query)
+                """
+            )
 
             result = cast(list[sqlite3.Row], self._cursor.fetchall())
             boards = [deserialize_board_record(dict(r)) for r in result]

invokeai/app/services/boards/boards_base.py

@@ -44,12 +44,16 @@ class BoardServiceABC(ABC):
 
     @abstractmethod
     def get_many(
-        self, offset: int = 0, limit: int = 10, include_archived: bool = False
+        self,
+        offset: int = 0,
+        limit: int = 10,
     ) -> OffsetPaginatedResults[BoardDTO]:
         """Gets many boards."""
         pass
 
     @abstractmethod
-    def get_all(self, include_archived: bool = False) -> list[BoardDTO]:
+    def get_all(
+        self,
+    ) -> list[BoardDTO]:
         """Gets all boards."""
         pass

invokeai/app/services/boards/boards_default.py

@@ -48,10 +48,8 @@ class BoardService(BoardServiceABC):
     def delete(self, board_id: str) -> None:
         self.__invoker.services.board_records.delete(board_id)
 
-    def get_many(
-        self, offset: int = 0, limit: int = 10, include_archived: bool = False
-    ) -> OffsetPaginatedResults[BoardDTO]:
-        board_records = self.__invoker.services.board_records.get_many(offset, limit, include_archived)
+    def get_many(self, offset: int = 0, limit: int = 10) -> OffsetPaginatedResults[BoardDTO]:
+        board_records = self.__invoker.services.board_records.get_many(offset, limit)
         board_dtos = []
         for r in board_records.items:
             cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(r.board_id)
@@ -65,8 +63,8 @@ class BoardService(BoardServiceABC):
 
         return OffsetPaginatedResults[BoardDTO](items=board_dtos, offset=offset, limit=limit, total=len(board_dtos))
 
-    def get_all(self, include_archived: bool = False) -> list[BoardDTO]:
-        board_records = self.__invoker.services.board_records.get_all(include_archived)
+    def get_all(self) -> list[BoardDTO]:
+        board_records = self.__invoker.services.board_records.get_all()
         board_dtos = []
         for r in board_records:
             cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(r.board_id)

invokeai/app/services/bulk_download/bulk_download_default.py

@@ -106,7 +106,9 @@ class BulkDownloadService(BulkDownloadBase):
         if self._invoker:
             assert bulk_download_id is not None
             self._invoker.services.events.emit_bulk_download_started(
-                bulk_download_id, bulk_download_item_id, bulk_download_item_name
+                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(
@@ -116,8 +118,10 @@ class BulkDownloadService(BulkDownloadBase):
         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_complete(
-                bulk_download_id, bulk_download_item_id, bulk_download_item_name
+            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(
@@ -127,8 +131,11 @@ class BulkDownloadService(BulkDownloadBase):
         if self._invoker:
             assert bulk_download_id is not None
             assert exception is not None
-            self._invoker.services.events.emit_bulk_download_error(
-                bulk_download_id, bulk_download_item_id, bulk_download_item_name, str(exception)
+            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):

invokeai/app/services/config/config_default.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import copy
 import locale
 import os
 import re
@@ -26,13 +25,14 @@ DB_FILE = Path("invokeai.db")
 LEGACY_INIT_FILE = Path("invokeai.init")
 DEFAULT_RAM_CACHE = 10.0
 DEFAULT_VRAM_CACHE = 0.25
+DEFAULT_CONVERT_CACHE = 20.0
 DEVICE = Literal["auto", "cpu", "cuda", "cuda:1", "mps"]
-PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
+PRECISION = Literal["auto", "float16", "bfloat16", "float32", "autocast"]
 ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
 ATTENTION_SLICE_SIZE = Literal["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8]
 LOG_FORMAT = Literal["plain", "color", "syslog", "legacy"]
 LOG_LEVEL = Literal["debug", "info", "warning", "error", "critical"]
-CONFIG_SCHEMA_VERSION = "4.0.2"
+CONFIG_SCHEMA_VERSION = "4.0.0"
 
 
 def get_default_ram_cache_size() -> float:
@@ -85,8 +85,7 @@ class InvokeAIAppConfig(BaseSettings):
         log_tokenization: Enable logging of parsed prompt tokens.
         patchmatch: Enable patchmatch inpaint code.
         models_dir: Path to the models directory.
-        convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).
-        download_cache_dir: Path to the directory that contains dynamically downloaded models.
+        convert_cache_dir: Path to the converted models cache directory. When loading a non-diffusers model, it will be converted and store on disk at this location.
         legacy_conf_dir: Path to directory of legacy checkpoint config files.
         db_dir: Path to InvokeAI databases directory.
         outputs_dir: Path to directory for outputs.
@@ -102,17 +101,17 @@ class InvokeAIAppConfig(BaseSettings):
         profiles_dir: Path to profiles output directory.
         ram: Maximum memory amount used by memory model cache for rapid switching (GB).
         vram: Amount of VRAM reserved for model storage (GB).
+        convert_cache: Maximum size of on-disk converted models cache (GB).
         lazy_offload: Keep models in VRAM until their space is needed.
         log_memory_usage: 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.
         device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `cuda:1`, `mps`
-        precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`
+        precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`, `autocast`
         sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
         attention_type: Attention type.<br>Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
         attention_slice_size: Slice size, valid when attention_type=="sliced".<br>Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`
         force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).
         pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.
         max_queue_size: Maximum number of items in the session queue.
-        clear_queue_on_startup: Empties session queue on startup.
         allow_nodes: List of nodes to allow. Omit to allow all.
         deny_nodes: List of nodes to deny. Omit to deny none.
         node_cache_size: How many cached nodes to keep in memory.
@@ -147,8 +146,7 @@ class InvokeAIAppConfig(BaseSettings):
 
     # PATHS
     models_dir:                    Path = Field(default=Path("models"),     description="Path to the models directory.")
-    convert_cache_dir:             Path = Field(default=Path("models/.convert_cache"), description="Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).")
-    download_cache_dir:            Path = Field(default=Path("models/.download_cache"), description="Path to the directory that contains dynamically downloaded models.")
+    convert_cache_dir:             Path = Field(default=Path("models/.cache"), description="Path to the converted models cache directory. When loading a non-diffusers model, it will be converted and store on disk at this location.")
     legacy_conf_dir:               Path = Field(default=Path("configs"), description="Path to directory of legacy checkpoint config files.")
     db_dir:                        Path = Field(default=Path("databases"),  description="Path to InvokeAI databases directory.")
     outputs_dir:                   Path = Field(default=Path("outputs"),    description="Path to directory for outputs.")
@@ -169,8 +167,9 @@ class InvokeAIAppConfig(BaseSettings):
     profiles_dir:                  Path = Field(default=Path("profiles"),   description="Path to profiles output directory.")
 
     # CACHE
-    ram:                           float = Field(default_factory=get_default_ram_cache_size, gt=0, description="Maximum memory amount used by memory model cache for rapid switching (GB).")
-    vram:                          float = Field(default=DEFAULT_VRAM_CACHE, ge=0, description="Amount of VRAM reserved for model storage (GB).")
+    ram:                          float = Field(default_factory=get_default_ram_cache_size, gt=0, description="Maximum memory amount used by memory model cache for rapid switching (GB).")
+    vram:                         float = Field(default=DEFAULT_VRAM_CACHE, ge=0, description="Amount of VRAM reserved for model storage (GB).")
+    convert_cache:                float = Field(default=DEFAULT_CONVERT_CACHE, ge=0, description="Maximum size of on-disk converted models cache (GB).")
     lazy_offload:                  bool = Field(default=True,               description="Keep models in VRAM until their space is needed.")
     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.")
 
@@ -185,7 +184,6 @@ class InvokeAIAppConfig(BaseSettings):
     force_tiled_decode:            bool = Field(default=False,              description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).")
     pil_compress_level:             int = Field(default=1,                  description="The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.")
     max_queue_size:                 int = Field(default=10000, gt=0,        description="Maximum number of items in the session queue.")
-    clear_queue_on_startup:        bool = Field(default=False,              description="Empties session queue on startup.")
 
     # NODES
     allow_nodes:    Optional[list[str]] = Field(default=None,               description="List of nodes to allow. Omit to allow all.")
@@ -305,11 +303,6 @@ class InvokeAIAppConfig(BaseSettings):
         """Path to the converted cache models directory, resolved to an absolute path.."""
         return self._resolve(self.convert_cache_dir)
 
-    @property
-    def download_cache_path(self) -> Path:
-        """Path to the downloaded models directory, resolved to an absolute path.."""
-        return self._resolve(self.download_cache_dir)
-
     @property
     def custom_nodes_path(self) -> Path:
         """Path to the custom nodes directory, resolved to an absolute path.."""
@@ -355,14 +348,14 @@ class DefaultInvokeAIAppConfig(InvokeAIAppConfig):
         return (init_settings,)
 
 
-def migrate_v3_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
-    """Migrate a v3 config dictionary to a v4.0.0.
+def migrate_v3_config_dict(config_dict: dict[str, Any]) -> InvokeAIAppConfig:
+    """Migrate a v3 config dictionary to a current config object.
 
     Args:
         config_dict: A dictionary of settings from a v3 config file.
 
     Returns:
-        An `InvokeAIAppConfig` config dict.
+        An instance of `InvokeAIAppConfig` with the migrated settings.
 
     """
     parsed_config_dict: dict[str, Any] = {}
@@ -377,9 +370,6 @@ def migrate_v3_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
             # `max_vram_cache_size` was renamed to `vram` some time in v3, but both names were used
             if k == "max_vram_cache_size" and "vram" not in category_dict:
                 parsed_config_dict["vram"] = v
-            # autocast was removed in v4.0.1
-            if k == "precision" and v == "autocast":
-                parsed_config_dict["precision"] = "auto"
             if k == "conf_path":
                 parsed_config_dict["legacy_models_yaml_path"] = v
             if k == "legacy_conf_dir":
@@ -396,41 +386,10 @@ def migrate_v3_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
             elif k in InvokeAIAppConfig.model_fields:
                 # skip unknown fields
                 parsed_config_dict[k] = v
-    parsed_config_dict["schema_version"] = "4.0.0"
-    return parsed_config_dict
+    # When migrating the config file, we should not include currently-set environment variables.
+    config = DefaultInvokeAIAppConfig.model_validate(parsed_config_dict)
 
-
-def migrate_v4_0_0_to_4_0_1_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
-    """Migrate v4.0.0 config dictionary to a v4.0.1 config dictionary
-
-    Args:
-        config_dict: A dictionary of settings from a v4.0.0 config file.
-
-    Returns:
-        A config dict with the settings migrated to v4.0.1.
-    """
-    parsed_config_dict: dict[str, Any] = copy.deepcopy(config_dict)
-    # precision "autocast" was replaced by "auto" in v4.0.1
-    if parsed_config_dict.get("precision") == "autocast":
-        parsed_config_dict["precision"] = "auto"
-    parsed_config_dict["schema_version"] = "4.0.1"
-    return parsed_config_dict
-
-
-def migrate_v4_0_1_to_4_0_2_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
-    """Migrate v4.0.1 config dictionary to a v4.0.2 config dictionary.
-
-    Args:
-        config_dict: A dictionary of settings from a v4.0.1 config file.
-
-    Returns:
-        An config dict with the settings migrated to v4.0.2.
-    """
-    parsed_config_dict: dict[str, Any] = copy.deepcopy(config_dict)
-    # convert_cache was removed in 4.0.2
-    parsed_config_dict.pop("convert_cache", None)
-    parsed_config_dict["schema_version"] = "4.0.2"
-    return parsed_config_dict
+    return config
 
 
 def load_and_migrate_config(config_path: Path) -> InvokeAIAppConfig:
@@ -444,40 +403,32 @@ def load_and_migrate_config(config_path: Path) -> InvokeAIAppConfig:
     """
     assert config_path.suffix == ".yaml"
     with open(config_path, "rt", encoding=locale.getpreferredencoding()) as file:
-        loaded_config_dict: dict[str, Any] = yaml.safe_load(file)
+        loaded_config_dict = yaml.safe_load(file)
 
     assert isinstance(loaded_config_dict, dict)
 
-    migrated = False
     if "InvokeAI" in loaded_config_dict:
-        migrated = True
-        loaded_config_dict = migrate_v3_config_dict(loaded_config_dict)  # pyright: ignore [reportUnknownArgumentType]
-    if loaded_config_dict["schema_version"] == "4.0.0":
-        migrated = True
-        loaded_config_dict = migrate_v4_0_0_to_4_0_1_config_dict(loaded_config_dict)
-    if loaded_config_dict["schema_version"] == "4.0.1":
-        migrated = True
-        loaded_config_dict = migrate_v4_0_1_to_4_0_2_config_dict(loaded_config_dict)
-
-    if migrated:
+        # This is a v3 config file, attempt to migrate it
         shutil.copy(config_path, config_path.with_suffix(".yaml.bak"))
         try:
-            # load and write without environment variables
-            migrated_config = DefaultInvokeAIAppConfig.model_validate(loaded_config_dict)
-            migrated_config.write_file(config_path)
+            # loaded_config_dict could be the wrong shape, but we will catch all exceptions below
+            migrated_config = migrate_v3_config_dict(loaded_config_dict)  # pyright: ignore [reportUnknownArgumentType]
         except Exception as e:
             shutil.copy(config_path.with_suffix(".yaml.bak"), config_path)
             raise RuntimeError(f"Failed to load and migrate v3 config file {config_path}: {e}") from e
-
-    try:
-        # Meta is not included in the model fields, so we need to validate it separately
-        config = InvokeAIAppConfig.model_validate(loaded_config_dict)
-        assert (
-            config.schema_version == CONFIG_SCHEMA_VERSION
-        ), f"Invalid schema version, expected {CONFIG_SCHEMA_VERSION}: {config.schema_version}"
-        return config
-    except Exception as e:
-        raise RuntimeError(f"Failed to load config file {config_path}: {e}") from e
+        migrated_config.write_file(config_path)
+        return migrated_config
+    else:
+        # Attempt to load as a v4 config file
+        try:
+            # Meta is not included in the model fields, so we need to validate it separately
+            config = InvokeAIAppConfig.model_validate(loaded_config_dict)
+            assert (
+                config.schema_version == CONFIG_SCHEMA_VERSION
+            ), f"Invalid schema version, expected {CONFIG_SCHEMA_VERSION}: {config.schema_version}"
+            return config
+        except Exception as e:
+            raise RuntimeError(f"Failed to load config file {config_path}: {e}") from e
 
 
 @lru_cache(maxsize=1)

invokeai/app/services/download/__init__.py

@@ -1,17 +1,10 @@
 """Init file for download queue."""
 
-from .download_base import (
-    DownloadJob,
-    DownloadJobStatus,
-    DownloadQueueServiceBase,
-    MultiFileDownloadJob,
-    UnknownJobIDException,
-)
+from .download_base import DownloadJob, DownloadJobStatus, DownloadQueueServiceBase, UnknownJobIDException
 from .download_default import DownloadQueueService, TqdmProgress
 
 __all__ = [
     "DownloadJob",
-    "MultiFileDownloadJob",
     "DownloadQueueServiceBase",
     "DownloadQueueService",
     "TqdmProgress",

invokeai/app/services/download/download_base.py

@@ -5,13 +5,11 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from functools import total_ordering
 from pathlib import Path
-from typing import Any, Callable, List, Optional, Set, Union
+from typing import Any, Callable, List, Optional
 
 from pydantic import BaseModel, Field, PrivateAttr
 from pydantic.networks import AnyHttpUrl
 
-from invokeai.backend.model_manager.metadata import RemoteModelFile
-
 
 class DownloadJobStatus(str, Enum):
     """State of a download job."""
@@ -35,23 +33,30 @@ class ServiceInactiveException(Exception):
     """This exception is raised when user attempts to initiate a download before the service is started."""
 
 
-SingleFileDownloadEventHandler = Callable[["DownloadJob"], None]
-SingleFileDownloadExceptionHandler = Callable[["DownloadJob", Optional[Exception]], None]
-MultiFileDownloadEventHandler = Callable[["MultiFileDownloadJob"], None]
-MultiFileDownloadExceptionHandler = Callable[["MultiFileDownloadJob", Optional[Exception]], None]
-DownloadEventHandler = Union[SingleFileDownloadEventHandler, MultiFileDownloadEventHandler]
-DownloadExceptionHandler = Union[SingleFileDownloadExceptionHandler, MultiFileDownloadExceptionHandler]
+DownloadEventHandler = Callable[["DownloadJob"], None]
+DownloadExceptionHandler = Callable[["DownloadJob", Optional[Exception]], None]
 
 
-class DownloadJobBase(BaseModel):
-    """Base of classes to monitor and control downloads."""
+@total_ordering
+class DownloadJob(BaseModel):
+    """Class to monitor and control a model download request."""
 
+    # required variables to be passed in on creation
+    source: AnyHttpUrl = Field(description="Where to download from. Specific types specified in child classes.")
+    dest: Path = Field(description="Destination of downloaded model on local disk; a directory or file path")
+    access_token: Optional[str] = Field(default=None, description="authorization token for protected resources")
     # automatically assigned on creation
     id: int = Field(description="Numeric ID of this job", default=-1)  # default id is a sentinel
+    priority: int = Field(default=10, description="Queue priority; lower values are higher priority")
 
-    dest: Path = Field(description="Initial destination of downloaded model on local disk; a directory or file path")
-    download_path: Optional[Path] = Field(default=None, description="Final location of downloaded file or directory")
+    # set internally during download process
     status: DownloadJobStatus = Field(default=DownloadJobStatus.WAITING, description="Status of the download")
+    download_path: Optional[Path] = Field(default=None, description="Final location of downloaded file")
+    job_started: Optional[str] = Field(default=None, description="Timestamp for when the download job started")
+    job_ended: Optional[str] = Field(
+        default=None, description="Timestamp for when the download job ende1d (completed or errored)"
+    )
+    content_type: Optional[str] = Field(default=None, description="Content type of downloaded file")
     bytes: int = Field(default=0, description="Bytes downloaded so far")
     total_bytes: int = Field(default=0, description="Total file size (bytes)")
 
@@ -69,6 +74,14 @@ class DownloadJobBase(BaseModel):
     _on_cancelled: Optional[DownloadEventHandler] = PrivateAttr(default=None)
     _on_error: Optional[DownloadExceptionHandler] = PrivateAttr(default=None)
 
+    def __hash__(self) -> int:
+        """Return hash of the string representation of this object, for indexing."""
+        return hash(str(self))
+
+    def __le__(self, other: "DownloadJob") -> bool:
+        """Return True if this job's priority is less than another's."""
+        return self.priority <= other.priority
+
     def cancel(self) -> None:
         """Call to cancel the job."""
         self._cancelled = True
@@ -85,11 +98,6 @@ class DownloadJobBase(BaseModel):
         """Return true if job completed without errors."""
         return self.status == DownloadJobStatus.COMPLETED
 
-    @property
-    def waiting(self) -> bool:
-        """Return true if the job is waiting to run."""
-        return self.status == DownloadJobStatus.WAITING
-
     @property
     def running(self) -> bool:
         """Return true if the job is running."""
@@ -146,37 +154,6 @@ class DownloadJobBase(BaseModel):
         self._on_cancelled = on_cancelled
 
 
-@total_ordering
-class DownloadJob(DownloadJobBase):
-    """Class to monitor and control a model download request."""
-
-    # required variables to be passed in on creation
-    source: AnyHttpUrl = Field(description="Where to download from. Specific types specified in child classes.")
-    access_token: Optional[str] = Field(default=None, description="authorization token for protected resources")
-    priority: int = Field(default=10, description="Queue priority; lower values are higher priority")
-
-    # set internally during download process
-    job_started: Optional[str] = Field(default=None, description="Timestamp for when the download job started")
-    job_ended: Optional[str] = Field(
-        default=None, description="Timestamp for when the download job ende1d (completed or errored)"
-    )
-    content_type: Optional[str] = Field(default=None, description="Content type of downloaded file")
-
-    def __hash__(self) -> int:
-        """Return hash of the string representation of this object, for indexing."""
-        return hash(str(self))
-
-    def __le__(self, other: "DownloadJob") -> bool:
-        """Return True if this job's priority is less than another's."""
-        return self.priority <= other.priority
-
-
-class MultiFileDownloadJob(DownloadJobBase):
-    """Class to monitor and control multifile downloads."""
-
-    download_parts: Set[DownloadJob] = Field(default_factory=set, description="List of download parts.")
-
-
 class DownloadQueueServiceBase(ABC):
     """Multithreaded queue for downloading models via URL."""
 
@@ -224,48 +201,6 @@ class DownloadQueueServiceBase(ABC):
         """
         pass
 
-    @abstractmethod
-    def multifile_download(
-        self,
-        parts: List[RemoteModelFile],
-        dest: Path,
-        access_token: Optional[str] = None,
-        submit_job: bool = True,
-        on_start: Optional[DownloadEventHandler] = None,
-        on_progress: Optional[DownloadEventHandler] = None,
-        on_complete: Optional[DownloadEventHandler] = None,
-        on_cancelled: Optional[DownloadEventHandler] = None,
-        on_error: Optional[DownloadExceptionHandler] = None,
-    ) -> MultiFileDownloadJob:
-        """
-        Create and enqueue a multifile download job.
-
-        :param parts: Set of URL / filename pairs
-        :param dest: Path to download to. See below.
-        :param access_token: Access token to download the indicated files. If not provided,
-         each file's URL may be matched to an access token using the config file matching
-         system.
-        :param submit_job: If true [default] then submit the job for execution. Otherwise,
-         you will need to pass the job to submit_multifile_download().
-        :param on_start, on_progress, on_complete, on_error: Callbacks for the indicated
-         events.
-        :returns: A MultiFileDownloadJob object for monitoring the state of the download.
-
-        The `dest` argument is a Path object pointing to a directory. All downloads
-        with be placed inside this directory. The callbacks will receive the
-        MultiFileDownloadJob.
-        """
-        pass
-
-    @abstractmethod
-    def submit_multifile_download(self, job: MultiFileDownloadJob) -> None:
-        """
-        Enqueue a previously-created multi-file download job.
-
-        :param job: A MultiFileDownloadJob created with multifile_download()
-        """
-        pass
-
     @abstractmethod
     def submit_download_job(
         self,
@@ -317,7 +252,7 @@ class DownloadQueueServiceBase(ABC):
         pass
 
     @abstractmethod
-    def cancel_job(self, job: DownloadJobBase) -> None:
+    def cancel_job(self, job: DownloadJob) -> None:
         """Cancel the job, clearing partial downloads and putting it into ERROR state."""
         pass
 
@@ -327,7 +262,7 @@ class DownloadQueueServiceBase(ABC):
         pass
 
     @abstractmethod
-    def wait_for_job(self, job: DownloadJobBase, timeout: int = 0) -> DownloadJobBase:
+    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,

invokeai/app/services/download/download_default.py

@@ -8,28 +8,24 @@ import time
 import traceback
 from pathlib import Path
 from queue import Empty, PriorityQueue
-from typing import Any, Dict, List, Literal, Optional, Set
+from typing import Any, Dict, List, Optional, Set
 
 import requests
 from pydantic.networks import AnyHttpUrl
 from requests import HTTPError
 from tqdm import tqdm
 
-from invokeai.app.services.config import InvokeAIAppConfig, get_config
 from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.util.misc import get_iso_timestamp
-from invokeai.backend.model_manager.metadata import RemoteModelFile
 from invokeai.backend.util.logging import InvokeAILogger
 
 from .download_base import (
     DownloadEventHandler,
     DownloadExceptionHandler,
     DownloadJob,
-    DownloadJobBase,
     DownloadJobCancelledException,
     DownloadJobStatus,
     DownloadQueueServiceBase,
-    MultiFileDownloadJob,
     ServiceInactiveException,
     UnknownJobIDException,
 )
@@ -44,24 +40,20 @@ class DownloadQueueService(DownloadQueueServiceBase):
     def __init__(
         self,
         max_parallel_dl: int = 5,
-        app_config: Optional[InvokeAIAppConfig] = None,
-        event_bus: Optional["EventServiceBase"] = None,
+        event_bus: Optional[EventServiceBase] = None,
         requests_session: Optional[requests.sessions.Session] = None,
     ):
         """
         Initialize DownloadQueue.
 
-        :param app_config: InvokeAIAppConfig object
         :param max_parallel_dl: Number of simultaneous downloads allowed [5].
         :param requests_session: Optional requests.sessions.Session object, for unit tests.
         """
-        self._app_config = app_config or get_config()
         self._jobs: Dict[int, DownloadJob] = {}
-        self._download_part2parent: Dict[AnyHttpUrl, MultiFileDownloadJob] = {}
         self._next_job_id = 0
         self._queue: PriorityQueue[DownloadJob] = PriorityQueue()
         self._stop_event = threading.Event()
-        self._job_terminated_event = threading.Event()
+        self._job_completed_event = threading.Event()
         self._worker_pool: Set[threading.Thread] = set()
         self._lock = threading.Lock()
         self._logger = InvokeAILogger.get_logger("DownloadQueueService")
@@ -113,16 +105,18 @@ class DownloadQueueService(DownloadQueueServiceBase):
             raise ServiceInactiveException(
                 "The download service is not currently accepting requests. Please call start() to initialize the service."
             )
-        job.id = self._next_id()
-        job.set_callbacks(
-            on_start=on_start,
-            on_progress=on_progress,
-            on_complete=on_complete,
-            on_cancelled=on_cancelled,
-            on_error=on_error,
-        )
-        self._jobs[job.id] = job
-        self._queue.put(job)
+        with self._lock:
+            job.id = self._next_job_id
+            self._next_job_id += 1
+            job.set_callbacks(
+                on_start=on_start,
+                on_progress=on_progress,
+                on_complete=on_complete,
+                on_cancelled=on_cancelled,
+                on_error=on_error,
+            )
+            self._jobs[job.id] = job
+            self._queue.put(job)
 
     def download(
         self,
@@ -145,7 +139,7 @@ class DownloadQueueService(DownloadQueueServiceBase):
             source=source,
             dest=dest,
             priority=priority,
-            access_token=access_token or self._lookup_access_token(source),
+            access_token=access_token,
         )
         self.submit_download_job(
             job,
@@ -157,63 +151,10 @@ class DownloadQueueService(DownloadQueueServiceBase):
         )
         return job
 
-    def multifile_download(
-        self,
-        parts: List[RemoteModelFile],
-        dest: Path,
-        access_token: Optional[str] = None,
-        submit_job: bool = True,
-        on_start: Optional[DownloadEventHandler] = None,
-        on_progress: Optional[DownloadEventHandler] = None,
-        on_complete: Optional[DownloadEventHandler] = None,
-        on_cancelled: Optional[DownloadEventHandler] = None,
-        on_error: Optional[DownloadExceptionHandler] = None,
-    ) -> MultiFileDownloadJob:
-        mfdj = MultiFileDownloadJob(dest=dest, id=self._next_id())
-        mfdj.set_callbacks(
-            on_start=on_start,
-            on_progress=on_progress,
-            on_complete=on_complete,
-            on_cancelled=on_cancelled,
-            on_error=on_error,
-        )
-
-        for part in parts:
-            url = part.url
-            path = dest / part.path
-            assert path.is_relative_to(dest), "only relative download paths accepted"
-            job = DownloadJob(
-                source=url,
-                dest=path,
-                access_token=access_token,
-            )
-            mfdj.download_parts.add(job)
-            self._download_part2parent[job.source] = mfdj
-        if submit_job:
-            self.submit_multifile_download(mfdj)
-        return mfdj
-
-    def submit_multifile_download(self, job: MultiFileDownloadJob) -> None:
-        for download_job in job.download_parts:
-            self.submit_download_job(
-                download_job,
-                on_start=self._mfd_started,
-                on_progress=self._mfd_progress,
-                on_complete=self._mfd_complete,
-                on_cancelled=self._mfd_cancelled,
-                on_error=self._mfd_error,
-            )
-
     def join(self) -> None:
         """Wait for all jobs to complete."""
         self._queue.join()
 
-    def _next_id(self) -> int:
-        with self._lock:
-            id = self._next_job_id
-            self._next_job_id += 1
-        return id
-
     def list_jobs(self) -> List[DownloadJob]:
         """List all the jobs."""
         return list(self._jobs.values())
@@ -235,14 +176,14 @@ class DownloadQueueService(DownloadQueueServiceBase):
         except KeyError as excp:
             raise UnknownJobIDException("Unrecognized job") from excp
 
-    def cancel_job(self, job: DownloadJobBase) -> None:
+    def cancel_job(self, job: DownloadJob) -> None:
         """
         Cancel the indicated job.
 
         If it is running it will be stopped.
         job.status will be set to DownloadJobStatus.CANCELLED
         """
-        if job.status in [DownloadJobStatus.WAITING, DownloadJobStatus.RUNNING]:
+        with self._lock:
             job.cancel()
 
     def cancel_all_jobs(self) -> None:
@@ -251,12 +192,12 @@ class DownloadQueueService(DownloadQueueServiceBase):
             if not job.in_terminal_state:
                 self.cancel_job(job)
 
-    def wait_for_job(self, job: DownloadJobBase, timeout: int = 0) -> DownloadJobBase:
+    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_terminated_event.wait(timeout=0.25):  # in case we miss an event
-                self._job_terminated_event.clear()
+            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
@@ -285,25 +226,22 @@ class DownloadQueueService(DownloadQueueServiceBase):
                 job.job_started = get_iso_timestamp()
                 self._do_download(job)
                 self._signal_job_complete(job)
-            except DownloadJobCancelledException:
-                self._signal_job_cancelled(job)
-                self._cleanup_cancelled_job(job)
-            except Exception as excp:
+            except (OSError, HTTPError) as excp:
                 job.error_type = excp.__class__.__name__ + f"({str(excp)})"
                 job.error = traceback.format_exc()
                 self._signal_job_error(job, excp)
+            except DownloadJobCancelledException:
+                self._signal_job_cancelled(job)
+                self._cleanup_cancelled_job(job)
+
             finally:
                 job.job_ended = get_iso_timestamp()
-                self._job_terminated_event.set()  # signal a change to terminal state
-                self._download_part2parent.pop(job.source, None)  # if this is a subpart of a multipart job, remove it
-                self._job_terminated_event.set()
+                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.")
 
     def _do_download(self, job: DownloadJob) -> None:
         """Do the actual download."""
-
         url = job.source
         header = {"Authorization": f"Bearer {job.access_token}"} if job.access_token else {}
         open_mode = "wb"
@@ -332,7 +270,7 @@ class DownloadQueueService(DownloadQueueServiceBase):
             job.dest.parent.mkdir(parents=True, exist_ok=True)
             job.download_path = job.dest
 
-        assert job.download_path
+        assert job.download_path is not None
 
         # Don't clobber an existing file. See commit 82c2c85202f88c6d24ff84710f297cfc6ae174af
         # for code that instead resumes an interrupted download.
@@ -342,6 +280,9 @@ class DownloadQueueService(DownloadQueueServiceBase):
         # append ".downloading" to the path
         in_progress_path = self._in_progress_path(job.download_path)
 
+        # catch rare race condition that is appearing in unit tests.
+        assert in_progress_path.parent.exists(), f"Directory doesn't exist! in_progress_path={in_progress_path}; parent={in_progress_path.parent}"
+
         # signal caller that the download is starting. At this point, key fields such as
         # download_path and total_bytes will be populated. We call it here because the might
         # discover that the local file is already complete and generate a COMPLETED status.
@@ -380,8 +321,10 @@ class DownloadQueueService(DownloadQueueServiceBase):
         in_progress_path.rename(job.download_path)
 
     def _validate_filename(self, directory: str, filename: str) -> bool:
-        pc_name_max = get_pc_name_max(directory)
-        pc_path_max = get_pc_path_max(directory)
+        pc_name_max = os.pathconf(directory, "PC_NAME_MAX") if hasattr(os, "pathconf") else 260  # hardcoded for windows
+        pc_path_max = (
+            os.pathconf(directory, "PC_PATH_MAX") if hasattr(os, "pathconf") else 32767
+        )  # hardcoded for windows with long names enabled
         if "/" in filename:
             return False
         if filename.startswith(".."):
@@ -395,53 +338,79 @@ class DownloadQueueService(DownloadQueueServiceBase):
     def _in_progress_path(self, path: Path) -> Path:
         return path.with_name(path.name + ".downloading")
 
-    def _lookup_access_token(self, source: AnyHttpUrl) -> Optional[str]:
-        # Pull the token from config if it exists and matches the URL
-        token = None
-        for pair in self._app_config.remote_api_tokens or []:
-            if re.search(pair.url_regex, str(source)):
-                token = pair.token
-                break
-        return token
-
     def _signal_job_started(self, job: DownloadJob) -> None:
         job.status = DownloadJobStatus.RUNNING
-        self._execute_cb(job, "on_start")
+        if job.on_start:
+            try:
+                job.on_start(job)
+            except Exception as e:
+                self._logger.error(
+                    f"An error occurred while processing the on_start callback: {traceback.format_exception(e)}"
+                )
         if self._event_bus:
-            self._event_bus.emit_download_started(job)
+            assert job.download_path
+            self._event_bus.emit_download_started(str(job.source), job.download_path.as_posix())
 
     def _signal_job_progress(self, job: DownloadJob) -> None:
-        self._execute_cb(job, "on_progress")
+        if job.on_progress:
+            try:
+                job.on_progress(job)
+            except Exception as e:
+                self._logger.error(
+                    f"An error occurred while processing the on_progress callback: {traceback.format_exception(e)}"
+                )
         if self._event_bus:
-            self._event_bus.emit_download_progress(job)
+            assert job.download_path
+            self._event_bus.emit_download_progress(
+                str(job.source),
+                download_path=job.download_path.as_posix(),
+                current_bytes=job.bytes,
+                total_bytes=job.total_bytes,
+            )
 
     def _signal_job_complete(self, job: DownloadJob) -> None:
         job.status = DownloadJobStatus.COMPLETED
-        self._execute_cb(job, "on_complete")
+        if job.on_complete:
+            try:
+                job.on_complete(job)
+            except Exception as e:
+                self._logger.error(
+                    f"An error occurred while processing the on_complete callback: {traceback.format_exception(e)}"
+                )
         if self._event_bus:
-            self._event_bus.emit_download_complete(job)
+            assert job.download_path
+            self._event_bus.emit_download_complete(
+                str(job.source), download_path=job.download_path.as_posix(), total_bytes=job.total_bytes
+            )
 
     def _signal_job_cancelled(self, job: DownloadJob) -> None:
         if job.status not in [DownloadJobStatus.RUNNING, DownloadJobStatus.WAITING]:
             return
         job.status = DownloadJobStatus.CANCELLED
-        self._execute_cb(job, "on_cancelled")
+        if job.on_cancelled:
+            try:
+                job.on_cancelled(job)
+            except Exception as e:
+                self._logger.error(
+                    f"An error occurred while processing the on_cancelled callback: {traceback.format_exception(e)}"
+                )
         if self._event_bus:
-            self._event_bus.emit_download_cancelled(job)
-
-        # if multifile download, then signal the parent
-        if parent_job := self._download_part2parent.get(job.source, None):
-            if not parent_job.in_terminal_state:
-                parent_job.status = DownloadJobStatus.CANCELLED
-                self._execute_cb(parent_job, "on_cancelled")
+            self._event_bus.emit_download_cancelled(str(job.source))
 
     def _signal_job_error(self, job: DownloadJob, excp: Optional[Exception] = None) -> None:
         job.status = DownloadJobStatus.ERROR
         self._logger.error(f"{str(job.source)}: {traceback.format_exception(excp)}")
-        self._execute_cb(job, "on_error", excp)
-
+        if job.on_error:
+            try:
+                job.on_error(job, excp)
+            except Exception as e:
+                self._logger.error(
+                    f"An error occurred while processing the on_error callback: {traceback.format_exception(e)}"
+                )
         if self._event_bus:
-            self._event_bus.emit_download_error(job)
+            assert job.error_type
+            assert job.error
+            self._event_bus.emit_download_error(str(job.source), error_type=job.error_type, error=job.error)
 
     def _cleanup_cancelled_job(self, job: DownloadJob) -> None:
         self._logger.debug(f"Cleaning up leftover files from cancelled download job {job.download_path}")
@@ -452,117 +421,6 @@ class DownloadQueueService(DownloadQueueServiceBase):
         except OSError as excp:
             self._logger.warning(excp)
 
-    ########################################
-    # callbacks used for multifile downloads
-    ########################################
-    def _mfd_started(self, download_job: DownloadJob) -> None:
-        self._logger.info(f"File download started: {download_job.source}")
-        with self._lock:
-            mf_job = self._download_part2parent[download_job.source]
-            if mf_job.waiting:
-                mf_job.total_bytes = sum(x.total_bytes for x in mf_job.download_parts)
-                mf_job.status = DownloadJobStatus.RUNNING
-                assert download_job.download_path is not None
-                path_relative_to_destdir = download_job.download_path.relative_to(mf_job.dest)
-                mf_job.download_path = (
-                    mf_job.dest / path_relative_to_destdir.parts[0]
-                )  # keep just the first component of the path
-                self._execute_cb(mf_job, "on_start")
-
-    def _mfd_progress(self, download_job: DownloadJob) -> None:
-        with self._lock:
-            mf_job = self._download_part2parent[download_job.source]
-            if mf_job.cancelled:
-                for part in mf_job.download_parts:
-                    self.cancel_job(part)
-            elif mf_job.running:
-                mf_job.total_bytes = sum(x.total_bytes for x in mf_job.download_parts)
-                mf_job.bytes = sum(x.total_bytes for x in mf_job.download_parts)
-                self._execute_cb(mf_job, "on_progress")
-
-    def _mfd_complete(self, download_job: DownloadJob) -> None:
-        self._logger.info(f"Download complete: {download_job.source}")
-        with self._lock:
-            mf_job = self._download_part2parent[download_job.source]
-
-            # are there any more active jobs left in this task?
-            if mf_job.running and all(x.complete for x in mf_job.download_parts):
-                mf_job.status = DownloadJobStatus.COMPLETED
-                self._execute_cb(mf_job, "on_complete")
-
-            # we're done with this sub-job
-            self._job_terminated_event.set()
-
-    def _mfd_cancelled(self, download_job: DownloadJob) -> None:
-        with self._lock:
-            mf_job = self._download_part2parent[download_job.source]
-            assert mf_job is not None
-
-            if not mf_job.in_terminal_state:
-                self._logger.warning(f"Download cancelled: {download_job.source}")
-                mf_job.cancel()
-
-            for s in mf_job.download_parts:
-                self.cancel_job(s)
-
-    def _mfd_error(self, download_job: DownloadJob, excp: Optional[Exception] = None) -> None:
-        with self._lock:
-            mf_job = self._download_part2parent[download_job.source]
-            assert mf_job is not None
-            if not mf_job.in_terminal_state:
-                mf_job.status = download_job.status
-                mf_job.error = download_job.error
-                mf_job.error_type = download_job.error_type
-                self._execute_cb(mf_job, "on_error", excp)
-                self._logger.error(
-                    f"Cancelling {mf_job.dest} due to an error while downloading {download_job.source}: {str(excp)}"
-                )
-                for s in [x for x in mf_job.download_parts if x.running]:
-                    self.cancel_job(s)
-                self._download_part2parent.pop(download_job.source)
-                self._job_terminated_event.set()
-
-    def _execute_cb(
-        self,
-        job: DownloadJob | MultiFileDownloadJob,
-        callback_name: Literal[
-            "on_start",
-            "on_progress",
-            "on_complete",
-            "on_cancelled",
-            "on_error",
-        ],
-        excp: Optional[Exception] = None,
-    ) -> None:
-        if callback := getattr(job, callback_name, None):
-            args = [job, excp] if excp else [job]
-            try:
-                callback(*args)
-            except Exception as e:
-                self._logger.error(
-                    f"An error occurred while processing the {callback_name} callback: {traceback.format_exception(e)}"
-                )
-
-
-def get_pc_name_max(directory: str) -> int:
-    if hasattr(os, "pathconf"):
-        try:
-            return os.pathconf(directory, "PC_NAME_MAX")
-        except OSError:
-            # macOS w/ external drives raise OSError
-            pass
-    return 260  # hardcoded for windows
-
-
-def get_pc_path_max(directory: str) -> int:
-    if hasattr(os, "pathconf"):
-        try:
-            return os.pathconf(directory, "PC_PATH_MAX")
-        except OSError:
-            # some platforms may not have this value
-            pass
-    return 32767  # hardcoded for windows with long names enabled
-
 
 # Example on_progress event handler to display a TQDM status bar
 # Activate with:

invokeai/app/services/events/events_base.py

@@ -1,199 +1,486 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
 
-from typing import TYPE_CHECKING, Optional
+from typing import Any, Dict, List, Optional, Union
 
-from invokeai.app.services.events.events_common import (
-    BatchEnqueuedEvent,
-    BulkDownloadCompleteEvent,
-    BulkDownloadErrorEvent,
-    BulkDownloadStartedEvent,
-    DownloadCancelledEvent,
-    DownloadCompleteEvent,
-    DownloadErrorEvent,
-    DownloadProgressEvent,
-    DownloadStartedEvent,
-    EventBase,
-    InvocationCompleteEvent,
-    InvocationDenoiseProgressEvent,
-    InvocationErrorEvent,
-    InvocationStartedEvent,
-    ModelInstallCancelledEvent,
-    ModelInstallCompleteEvent,
-    ModelInstallDownloadProgressEvent,
-    ModelInstallDownloadsCompleteEvent,
-    ModelInstallDownloadStartedEvent,
-    ModelInstallErrorEvent,
-    ModelInstallStartedEvent,
-    ModelLoadCompleteEvent,
-    ModelLoadStartedEvent,
-    QueueClearedEvent,
-    QueueItemStatusChangedEvent,
+from invokeai.app.services.session_processor.session_processor_common import ProgressImage
+from invokeai.app.services.session_queue.session_queue_common import (
+    BatchStatus,
+    EnqueueBatchResult,
+    SessionQueueItem,
+    SessionQueueStatus,
 )
-from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
-
-if TYPE_CHECKING:
-    from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
-    from invokeai.app.services.download.download_base import DownloadJob
-    from invokeai.app.services.model_install.model_install_common import ModelInstallJob
-    from invokeai.app.services.session_processor.session_processor_common import ProgressImage
-    from invokeai.app.services.session_queue.session_queue_common import (
-        BatchStatus,
-        EnqueueBatchResult,
-        SessionQueueItem,
-        SessionQueueStatus,
-    )
-    from invokeai.backend.model_manager.config import AnyModelConfig, SubModelType
+from invokeai.app.util.misc import get_timestamp
+from invokeai.backend.model_manager import AnyModelConfig
+from invokeai.backend.model_manager.config import SubModelType
 
 
 class EventServiceBase:
+    queue_event: str = "queue_event"
+    bulk_download_event: str = "bulk_download_event"
+    download_event: str = "download_event"
+    model_event: str = "model_event"
+
     """Basic event bus, to have an empty stand-in when not needed"""
 
-    def dispatch(self, event: "EventBase") -> None:
+    def dispatch(self, event_name: str, payload: Any) -> None:
         pass
 
-    # region: Invocation
+    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_invocation_started(self, queue_item: "SessionQueueItem", invocation: "BaseInvocation") -> None:
-        """Emitted when an invocation is started"""
-        self.dispatch(InvocationStartedEvent.build(queue_item, invocation))
+    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()
+        self.dispatch(
+            event_name=EventServiceBase.queue_event,
+            payload={"event": event_name, "data": payload},
+        )
 
-    def emit_invocation_denoise_progress(
+    def __emit_download_event(self, event_name: str, payload: dict) -> None:
+        payload["timestamp"] = get_timestamp()
+        self.dispatch(
+            event_name=EventServiceBase.download_event,
+            payload={"event": event_name, "data": payload},
+        )
+
+    def __emit_model_event(self, event_name: str, payload: dict) -> None:
+        payload["timestamp"] = get_timestamp()
+        self.dispatch(
+            event_name=EventServiceBase.model_event,
+            payload={"event": event_name, "data": payload},
+        )
+
+    # Define events here for every event in the system.
+    # This will make them easier to integrate until we find a schema generator.
+    def emit_generator_progress(
         self,
-        queue_item: "SessionQueueItem",
-        invocation: "BaseInvocation",
-        intermediate_state: PipelineIntermediateState,
-        progress_image: "ProgressImage",
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        node_id: str,
+        source_node_id: str,
+        progress_image: Optional[ProgressImage],
+        step: int,
+        order: int,
+        total_steps: int,
     ) -> None:
-        """Emitted at each step during denoising of an invocation."""
-        self.dispatch(InvocationDenoiseProgressEvent.build(queue_item, invocation, intermediate_state, progress_image))
+        """Emitted when there is generation progress"""
+        self.__emit_queue_event(
+            event_name="generator_progress",
+            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,
+                "source_node_id": source_node_id,
+                "progress_image": progress_image.model_dump(mode="json") if progress_image is not None else None,
+                "step": step,
+                "order": order,
+                "total_steps": total_steps,
+            },
+        )
 
     def emit_invocation_complete(
-        self, queue_item: "SessionQueueItem", invocation: "BaseInvocation", output: "BaseInvocationOutput"
+        self,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        result: dict,
+        node: dict,
+        source_node_id: str,
     ) -> None:
-        """Emitted when an invocation is complete"""
-        self.dispatch(InvocationCompleteEvent.build(queue_item, invocation, output))
+        """Emitted when an invocation has completed"""
+        self.__emit_queue_event(
+            event_name="invocation_complete",
+            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": node,
+                "source_node_id": source_node_id,
+                "result": result,
+            },
+        )
 
     def emit_invocation_error(
         self,
-        queue_item: "SessionQueueItem",
-        invocation: "BaseInvocation",
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        node: dict,
+        source_node_id: str,
         error_type: str,
-        error_message: str,
-        error_traceback: str,
+        error: str,
     ) -> None:
-        """Emitted when an invocation encounters an error"""
-        self.dispatch(InvocationErrorEvent.build(queue_item, invocation, error_type, error_message, error_traceback))
+        """Emitted when an invocation has completed"""
+        self.__emit_queue_event(
+            event_name="invocation_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": node,
+                "source_node_id": source_node_id,
+                "error_type": error_type,
+                "error": error,
+            },
+        )
 
-    # endregion
+    def emit_invocation_started(
+        self,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        node: dict,
+        source_node_id: str,
+    ) -> None:
+        """Emitted when an invocation has started"""
+        self.__emit_queue_event(
+            event_name="invocation_started",
+            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": node,
+                "source_node_id": source_node_id,
+            },
+        )
 
-    # region Queue
+    def emit_graph_execution_complete(
+        self, queue_id: str, queue_item_id: int, queue_batch_id: str, graph_execution_state_id: str
+    ) -> None:
+        """Emitted when a session has completed all invocations"""
+        self.__emit_queue_event(
+            event_name="graph_execution_state_complete",
+            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,
+            },
+        )
+
+    def emit_model_load_started(
+        self,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        model_config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> None:
+        """Emitted when a model is requested"""
+        self.__emit_queue_event(
+            event_name="model_load_started",
+            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,
+                "model_config": model_config.model_dump(mode="json"),
+                "submodel_type": submodel_type,
+            },
+        )
+
+    def emit_model_load_completed(
+        self,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+        model_config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> None:
+        """Emitted when a model is correctly loaded (returns model info)"""
+        self.__emit_queue_event(
+            event_name="model_load_completed",
+            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,
+                "model_config": model_config.model_dump(mode="json"),
+                "submodel_type": submodel_type,
+            },
+        )
+
+    def emit_session_canceled(
+        self,
+        queue_id: str,
+        queue_item_id: int,
+        queue_batch_id: str,
+        graph_execution_state_id: str,
+    ) -> None:
+        """Emitted when a session is canceled"""
+        self.__emit_queue_event(
+            event_name="session_canceled",
+            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,
+            },
+        )
 
     def emit_queue_item_status_changed(
-        self, queue_item: "SessionQueueItem", batch_status: "BatchStatus", queue_status: "SessionQueueStatus"
+        self,
+        session_queue_item: SessionQueueItem,
+        batch_status: BatchStatus,
+        queue_status: SessionQueueStatus,
     ) -> None:
         """Emitted when a queue item's status changes"""
-        self.dispatch(QueueItemStatusChangedEvent.build(queue_item, batch_status, queue_status))
+        self.__emit_queue_event(
+            event_name="queue_item_status_changed",
+            payload={
+                "queue_id": queue_status.queue_id,
+                "queue_item": {
+                    "queue_id": session_queue_item.queue_id,
+                    "item_id": session_queue_item.item_id,
+                    "status": session_queue_item.status,
+                    "batch_id": session_queue_item.batch_id,
+                    "session_id": session_queue_item.session_id,
+                    "error": session_queue_item.error,
+                    "created_at": str(session_queue_item.created_at) if session_queue_item.created_at else None,
+                    "updated_at": str(session_queue_item.updated_at) if session_queue_item.updated_at else None,
+                    "started_at": str(session_queue_item.started_at) if session_queue_item.started_at else None,
+                    "completed_at": str(session_queue_item.completed_at) if session_queue_item.completed_at else None,
+                },
+                "batch_status": batch_status.model_dump(mode="json"),
+                "queue_status": queue_status.model_dump(mode="json"),
+            },
+        )
 
-    def emit_batch_enqueued(self, enqueue_result: "EnqueueBatchResult") -> None:
+    def emit_batch_enqueued(self, enqueue_result: EnqueueBatchResult) -> None:
         """Emitted when a batch is enqueued"""
-        self.dispatch(BatchEnqueuedEvent.build(enqueue_result))
+        self.__emit_queue_event(
+            event_name="batch_enqueued",
+            payload={
+                "queue_id": enqueue_result.queue_id,
+                "batch_id": enqueue_result.batch.batch_id,
+                "enqueued": enqueue_result.enqueued,
+            },
+        )
 
     def emit_queue_cleared(self, queue_id: str) -> None:
-        """Emitted when a queue is cleared"""
-        self.dispatch(QueueClearedEvent.build(queue_id))
+        """Emitted when the queue is cleared"""
+        self.__emit_queue_event(
+            event_name="queue_cleared",
+            payload={"queue_id": queue_id},
+        )
 
-    # endregion
+    def emit_download_started(self, source: str, download_path: str) -> None:
+        """
+        Emit when a download job is started.
 
-    # region Download
+        :param url: The downloaded url
+        """
+        self.__emit_download_event(
+            event_name="download_started",
+            payload={"source": source, "download_path": download_path},
+        )
 
-    def emit_download_started(self, job: "DownloadJob") -> None:
-        """Emitted when a download is started"""
-        self.dispatch(DownloadStartedEvent.build(job))
+    def emit_download_progress(self, source: str, download_path: str, current_bytes: int, total_bytes: int) -> None:
+        """
+        Emit "download_progress" events at regular intervals during a download job.
 
-    def emit_download_progress(self, job: "DownloadJob") -> None:
-        """Emitted at intervals during a download"""
-        self.dispatch(DownloadProgressEvent.build(job))
+        :param source: The downloaded source
+        :param download_path: The local downloaded file
+        :param current_bytes: Number of bytes downloaded so far
+        :param total_bytes: The size of the file being downloaded (if known)
+        """
+        self.__emit_download_event(
+            event_name="download_progress",
+            payload={
+                "source": source,
+                "download_path": download_path,
+                "current_bytes": current_bytes,
+                "total_bytes": total_bytes,
+            },
+        )
 
-    def emit_download_complete(self, job: "DownloadJob") -> None:
-        """Emitted when a download is completed"""
-        self.dispatch(DownloadCompleteEvent.build(job))
+    def emit_download_complete(self, source: str, download_path: str, total_bytes: int) -> None:
+        """
+        Emit a "download_complete" event at the end of a successful download.
 
-    def emit_download_cancelled(self, job: "DownloadJob") -> None:
-        """Emitted when a download is cancelled"""
-        self.dispatch(DownloadCancelledEvent.build(job))
+        :param source: Source URL
+        :param download_path: Path to the locally downloaded file
+        :param total_bytes: The size of the downloaded file
+        """
+        self.__emit_download_event(
+            event_name="download_complete",
+            payload={
+                "source": source,
+                "download_path": download_path,
+                "total_bytes": total_bytes,
+            },
+        )
 
-    def emit_download_error(self, job: "DownloadJob") -> None:
-        """Emitted when a download encounters an error"""
-        self.dispatch(DownloadErrorEvent.build(job))
+    def emit_download_cancelled(self, source: str) -> None:
+        """Emit a "download_cancelled" event in the event that the download was cancelled by user."""
+        self.__emit_download_event(
+            event_name="download_cancelled",
+            payload={
+                "source": source,
+            },
+        )
 
-    # endregion
+    def emit_download_error(self, source: str, error_type: str, error: str) -> None:
+        """
+        Emit a "download_error" event when an download job encounters an exception.
 
-    # region Model loading
+        :param source: Source URL
+        :param error_type: The name of the exception that raised the error
+        :param error: The traceback from this error
+        """
+        self.__emit_download_event(
+            event_name="download_error",
+            payload={
+                "source": source,
+                "error_type": error_type,
+                "error": error,
+            },
+        )
 
-    def emit_model_load_started(self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None) -> None:
-        """Emitted when a model load is started."""
-        self.dispatch(ModelLoadStartedEvent.build(config, submodel_type))
-
-    def emit_model_load_complete(
-        self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None
+    def emit_model_install_downloading(
+        self,
+        source: str,
+        local_path: str,
+        bytes: int,
+        total_bytes: int,
+        parts: List[Dict[str, Union[str, int]]],
+        id: int,
     ) -> None:
-        """Emitted when a model load is complete."""
-        self.dispatch(ModelLoadCompleteEvent.build(config, submodel_type))
+        """
+        Emit at intervals while the install job is in progress (remote models only).
 
-    # endregion
+        :param source: Source of the model
+        :param local_path: Where model is downloading to
+        :param parts: Progress of downloading URLs that comprise the model, if any.
+        :param bytes: Number of bytes downloaded so far.
+        :param total_bytes: Total size of download, including all files.
+        This emits a Dict with keys "source", "local_path", "bytes" and "total_bytes".
+        """
+        self.__emit_model_event(
+            event_name="model_install_downloading",
+            payload={
+                "source": source,
+                "local_path": local_path,
+                "bytes": bytes,
+                "total_bytes": total_bytes,
+                "parts": parts,
+                "id": id,
+            },
+        )
 
-    # region Model install
+    def emit_model_install_downloads_done(self, source: str) -> None:
+        """
+        Emit once when all parts are downloaded, but before the probing and registration start.
 
-    def emit_model_install_download_started(self, job: "ModelInstallJob") -> None:
-        """Emitted at intervals while the install job is started (remote models only)."""
-        self.dispatch(ModelInstallDownloadStartedEvent.build(job))
+        :param source: Source of the model; local path, repo_id or url
+        """
+        self.__emit_model_event(
+            event_name="model_install_downloads_done",
+            payload={"source": source},
+        )
 
-    def emit_model_install_download_progress(self, job: "ModelInstallJob") -> None:
-        """Emitted at intervals while the install job is in progress (remote models only)."""
-        self.dispatch(ModelInstallDownloadProgressEvent.build(job))
+    def emit_model_install_running(self, source: str) -> None:
+        """
+        Emit once when an install job becomes active.
 
-    def emit_model_install_downloads_complete(self, job: "ModelInstallJob") -> None:
-        self.dispatch(ModelInstallDownloadsCompleteEvent.build(job))
+        :param source: Source of the model; local path, repo_id or url
+        """
+        self.__emit_model_event(
+            event_name="model_install_running",
+            payload={"source": source},
+        )
 
-    def emit_model_install_started(self, job: "ModelInstallJob") -> None:
-        """Emitted once when an install job is started (after any download)."""
-        self.dispatch(ModelInstallStartedEvent.build(job))
+    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.
 
-    def emit_model_install_complete(self, job: "ModelInstallJob") -> None:
-        """Emitted when an install job is completed successfully."""
-        self.dispatch(ModelInstallCompleteEvent.build(job))
+        :param source: Source of the model; local path, repo_id or url
+        :param key: Model config record key
+        :param total_bytes: Size of the model (may be None for installation of a local path)
+        """
+        self.__emit_model_event(
+            event_name="model_install_completed",
+            payload={"source": source, "total_bytes": total_bytes, "key": key, "id": id},
+        )
 
-    def emit_model_install_cancelled(self, job: "ModelInstallJob") -> None:
-        """Emitted when an install job is cancelled."""
-        self.dispatch(ModelInstallCancelledEvent.build(job))
+    def emit_model_install_cancelled(self, source: str, id: int) -> None:
+        """
+        Emit when an install job is cancelled.
 
-    def emit_model_install_error(self, job: "ModelInstallJob") -> None:
-        """Emitted when an install job encounters an exception."""
-        self.dispatch(ModelInstallErrorEvent.build(job))
+        :param source: Source of the model; local path, repo_id or url
+        """
+        self.__emit_model_event(
+            event_name="model_install_cancelled",
+            payload={"source": source, "id": id},
+        )
 
-    # endregion
+    def emit_model_install_error(self, source: str, error_type: str, error: str, id: int) -> None:
+        """
+        Emit when an install job encounters an exception.
 
-    # region Bulk image download
+        :param source: Source of the model
+        :param error_type: The name of the exception
+        :param error: A text description of the exception
+        """
+        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 image download is started"""
-        self.dispatch(BulkDownloadStartedEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name))
-
-    def emit_bulk_download_complete(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
-    ) -> None:
-        """Emitted when a bulk image download is complete"""
-        self.dispatch(BulkDownloadCompleteEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name))
-
-    def emit_bulk_download_error(
-        self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, error: str
-    ) -> None:
-        """Emitted when a bulk image download has an error"""
-        self.dispatch(
-            BulkDownloadErrorEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name, error)
+        """Emitted when a bulk download starts"""
+        self._emit_bulk_download_event(
+            event_name="bulk_download_started",
+            payload={
+                "bulk_download_id": bulk_download_id,
+                "bulk_download_item_id": bulk_download_item_id,
+                "bulk_download_item_name": bulk_download_item_name,
+            },
         )
 
-    # endregion
+    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/events/events_common.py

@@ -1,628 +0,0 @@
-from math import floor
-from typing import TYPE_CHECKING, Any, ClassVar, Coroutine, Generic, Optional, Protocol, TypeAlias, TypeVar
-
-from fastapi_events.handlers.local import local_handler
-from fastapi_events.registry.payload_schema import registry as payload_schema
-from pydantic import BaseModel, ConfigDict, Field
-
-from invokeai.app.services.session_processor.session_processor_common import ProgressImage
-from invokeai.app.services.session_queue.session_queue_common import (
-    QUEUE_ITEM_STATUS,
-    BatchStatus,
-    EnqueueBatchResult,
-    SessionQueueItem,
-    SessionQueueStatus,
-)
-from invokeai.app.services.shared.graph import AnyInvocation, AnyInvocationOutput
-from invokeai.app.util.misc import get_timestamp
-from invokeai.backend.model_manager.config import AnyModelConfig, SubModelType
-from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
-
-if TYPE_CHECKING:
-    from invokeai.app.services.download.download_base import DownloadJob
-    from invokeai.app.services.model_install.model_install_common import ModelInstallJob
-
-
-class EventBase(BaseModel):
-    """Base class for all events. All events must inherit from this class.
-
-    Events must define a class attribute `__event_name__` to identify the event.
-
-    All other attributes should be defined as normal for a pydantic model.
-
-    A timestamp is automatically added to the event when it is created.
-    """
-
-    __event_name__: ClassVar[str]
-    timestamp: int = Field(description="The timestamp of the event", default_factory=get_timestamp)
-
-    model_config = ConfigDict(json_schema_serialization_defaults_required=True)
-
-    @classmethod
-    def get_events(cls) -> set[type["EventBase"]]:
-        """Get a set of all event models."""
-
-        event_subclasses: set[type["EventBase"]] = set()
-        for subclass in cls.__subclasses__():
-            # We only want to include subclasses that are event models, not intermediary classes
-            if hasattr(subclass, "__event_name__"):
-                event_subclasses.add(subclass)
-            event_subclasses.update(subclass.get_events())
-
-        return event_subclasses
-
-
-TEvent = TypeVar("TEvent", bound=EventBase, contravariant=True)
-
-FastAPIEvent: TypeAlias = tuple[str, TEvent]
-"""
-A tuple representing a `fastapi-events` event, with the event name and payload.
-Provide a generic type to `TEvent` to specify the payload type.
-"""
-
-
-class FastAPIEventFunc(Protocol, Generic[TEvent]):
-    def __call__(self, event: FastAPIEvent[TEvent]) -> Optional[Coroutine[Any, Any, None]]: ...
-
-
-def register_events(events: set[type[TEvent]] | type[TEvent], func: FastAPIEventFunc[TEvent]) -> None:
-    """Register a function to handle specific events.
-
-    :param events: An event or set of events to handle
-    :param func: The function to handle the events
-    """
-    events = events if isinstance(events, set) else {events}
-    for event in events:
-        assert hasattr(event, "__event_name__")
-        local_handler.register(event_name=event.__event_name__, _func=func)  # pyright: ignore [reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
-
-
-class QueueEventBase(EventBase):
-    """Base class for queue events"""
-
-    queue_id: str = Field(description="The ID of the queue")
-
-
-class QueueItemEventBase(QueueEventBase):
-    """Base class for queue item events"""
-
-    item_id: int = Field(description="The ID of the queue item")
-    batch_id: str = Field(description="The ID of the queue batch")
-
-
-class InvocationEventBase(QueueItemEventBase):
-    """Base class for invocation events"""
-
-    session_id: str = Field(description="The ID of the session (aka graph execution state)")
-    queue_id: str = Field(description="The ID of the queue")
-    item_id: int = Field(description="The ID of the queue item")
-    batch_id: str = Field(description="The ID of the queue batch")
-    session_id: str = Field(description="The ID of the session (aka graph execution state)")
-    invocation: AnyInvocation = Field(description="The ID of the invocation")
-    invocation_source_id: str = Field(description="The ID of the prepared invocation's source node")
-
-
-@payload_schema.register
-class InvocationStartedEvent(InvocationEventBase):
-    """Event model for invocation_started"""
-
-    __event_name__ = "invocation_started"
-
-    @classmethod
-    def build(cls, queue_item: SessionQueueItem, invocation: AnyInvocation) -> "InvocationStartedEvent":
-        return cls(
-            queue_id=queue_item.queue_id,
-            item_id=queue_item.item_id,
-            batch_id=queue_item.batch_id,
-            session_id=queue_item.session_id,
-            invocation=invocation,
-            invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
-        )
-
-
-@payload_schema.register
-class InvocationDenoiseProgressEvent(InvocationEventBase):
-    """Event model for invocation_denoise_progress"""
-
-    __event_name__ = "invocation_denoise_progress"
-
-    progress_image: ProgressImage = Field(description="The progress image sent at each step during processing")
-    step: int = Field(description="The current step of the invocation")
-    total_steps: int = Field(description="The total number of steps in the invocation")
-    order: int = Field(description="The order of the invocation in the session")
-    percentage: float = Field(description="The percentage of completion of the invocation")
-
-    @classmethod
-    def build(
-        cls,
-        queue_item: SessionQueueItem,
-        invocation: AnyInvocation,
-        intermediate_state: PipelineIntermediateState,
-        progress_image: ProgressImage,
-    ) -> "InvocationDenoiseProgressEvent":
-        step = intermediate_state.step
-        total_steps = intermediate_state.total_steps
-        order = intermediate_state.order
-        return cls(
-            queue_id=queue_item.queue_id,
-            item_id=queue_item.item_id,
-            batch_id=queue_item.batch_id,
-            session_id=queue_item.session_id,
-            invocation=invocation,
-            invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
-            progress_image=progress_image,
-            step=step,
-            total_steps=total_steps,
-            order=order,
-            percentage=cls.calc_percentage(step, total_steps, order),
-        )
-
-    @staticmethod
-    def calc_percentage(step: int, total_steps: int, scheduler_order: float) -> float:
-        """Calculate the percentage of completion of denoising."""
-        if total_steps == 0:
-            return 0.0
-        if scheduler_order == 2:
-            return floor((step + 1 + 1) / 2) / floor((total_steps + 1) / 2)
-        # order == 1
-        return (step + 1 + 1) / (total_steps + 1)
-
-
-@payload_schema.register
-class InvocationCompleteEvent(InvocationEventBase):
-    """Event model for invocation_complete"""
-
-    __event_name__ = "invocation_complete"
-
-    result: AnyInvocationOutput = Field(description="The result of the invocation")
-
-    @classmethod
-    def build(
-        cls, queue_item: SessionQueueItem, invocation: AnyInvocation, result: AnyInvocationOutput
-    ) -> "InvocationCompleteEvent":
-        return cls(
-            queue_id=queue_item.queue_id,
-            item_id=queue_item.item_id,
-            batch_id=queue_item.batch_id,
-            session_id=queue_item.session_id,
-            invocation=invocation,
-            invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
-            result=result,
-        )
-
-
-@payload_schema.register
-class InvocationErrorEvent(InvocationEventBase):
-    """Event model for invocation_error"""
-
-    __event_name__ = "invocation_error"
-
-    error_type: str = Field(description="The error type")
-    error_message: str = Field(description="The error message")
-    error_traceback: str = Field(description="The error traceback")
-    user_id: Optional[str] = Field(default=None, description="The ID of the user who created the invocation")
-    project_id: Optional[str] = Field(default=None, description="The ID of the user who created the invocation")
-
-    @classmethod
-    def build(
-        cls,
-        queue_item: SessionQueueItem,
-        invocation: AnyInvocation,
-        error_type: str,
-        error_message: str,
-        error_traceback: str,
-    ) -> "InvocationErrorEvent":
-        return cls(
-            queue_id=queue_item.queue_id,
-            item_id=queue_item.item_id,
-            batch_id=queue_item.batch_id,
-            session_id=queue_item.session_id,
-            invocation=invocation,
-            invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
-            error_type=error_type,
-            error_message=error_message,
-            error_traceback=error_traceback,
-            user_id=getattr(queue_item, "user_id", None),
-            project_id=getattr(queue_item, "project_id", None),
-        )
-
-
-@payload_schema.register
-class QueueItemStatusChangedEvent(QueueItemEventBase):
-    """Event model for queue_item_status_changed"""
-
-    __event_name__ = "queue_item_status_changed"
-
-    status: QUEUE_ITEM_STATUS = Field(description="The new status of the queue item")
-    error_type: Optional[str] = Field(default=None, description="The error type, if any")
-    error_message: Optional[str] = Field(default=None, description="The error message, if any")
-    error_traceback: Optional[str] = Field(default=None, description="The error traceback, if any")
-    created_at: Optional[str] = Field(default=None, description="The timestamp when the queue item was created")
-    updated_at: Optional[str] = Field(default=None, description="The timestamp when the queue item was last updated")
-    started_at: Optional[str] = Field(default=None, description="The timestamp when the queue item was started")
-    completed_at: Optional[str] = Field(default=None, description="The timestamp when the queue item was completed")
-    batch_status: BatchStatus = Field(description="The status of the batch")
-    queue_status: SessionQueueStatus = Field(description="The status of the queue")
-    session_id: str = Field(description="The ID of the session (aka graph execution state)")
-
-    @classmethod
-    def build(
-        cls, queue_item: SessionQueueItem, batch_status: BatchStatus, queue_status: SessionQueueStatus
-    ) -> "QueueItemStatusChangedEvent":
-        return cls(
-            queue_id=queue_item.queue_id,
-            item_id=queue_item.item_id,
-            batch_id=queue_item.batch_id,
-            session_id=queue_item.session_id,
-            status=queue_item.status,
-            error_type=queue_item.error_type,
-            error_message=queue_item.error_message,
-            error_traceback=queue_item.error_traceback,
-            created_at=str(queue_item.created_at) if queue_item.created_at else None,
-            updated_at=str(queue_item.updated_at) if queue_item.updated_at else None,
-            started_at=str(queue_item.started_at) if queue_item.started_at else None,
-            completed_at=str(queue_item.completed_at) if queue_item.completed_at else None,
-            batch_status=batch_status,
-            queue_status=queue_status,
-        )
-
-
-@payload_schema.register
-class BatchEnqueuedEvent(QueueEventBase):
-    """Event model for batch_enqueued"""
-
-    __event_name__ = "batch_enqueued"
-
-    batch_id: str = Field(description="The ID of the batch")
-    enqueued: int = Field(description="The number of invocations enqueued")
-    requested: int = Field(
-        description="The number of invocations initially requested to be enqueued (may be less than enqueued if queue was full)"
-    )
-    priority: int = Field(description="The priority of the batch")
-
-    @classmethod
-    def build(cls, enqueue_result: EnqueueBatchResult) -> "BatchEnqueuedEvent":
-        return cls(
-            queue_id=enqueue_result.queue_id,
-            batch_id=enqueue_result.batch.batch_id,
-            enqueued=enqueue_result.enqueued,
-            requested=enqueue_result.requested,
-            priority=enqueue_result.priority,
-        )
-
-
-@payload_schema.register
-class QueueClearedEvent(QueueEventBase):
-    """Event model for queue_cleared"""
-
-    __event_name__ = "queue_cleared"
-
-    @classmethod
-    def build(cls, queue_id: str) -> "QueueClearedEvent":
-        return cls(queue_id=queue_id)
-
-
-class DownloadEventBase(EventBase):
-    """Base class for events associated with a download"""
-
-    source: str = Field(description="The source of the download")
-
-
-@payload_schema.register
-class DownloadStartedEvent(DownloadEventBase):
-    """Event model for download_started"""
-
-    __event_name__ = "download_started"
-
-    download_path: str = Field(description="The local path where the download is saved")
-
-    @classmethod
-    def build(cls, job: "DownloadJob") -> "DownloadStartedEvent":
-        assert job.download_path
-        return cls(source=str(job.source), download_path=job.download_path.as_posix())
-
-
-@payload_schema.register
-class DownloadProgressEvent(DownloadEventBase):
-    """Event model for download_progress"""
-
-    __event_name__ = "download_progress"
-
-    download_path: str = Field(description="The local path where the download is saved")
-    current_bytes: int = Field(description="The number of bytes downloaded so far")
-    total_bytes: int = Field(description="The total number of bytes to be downloaded")
-
-    @classmethod
-    def build(cls, job: "DownloadJob") -> "DownloadProgressEvent":
-        assert job.download_path
-        return cls(
-            source=str(job.source),
-            download_path=job.download_path.as_posix(),
-            current_bytes=job.bytes,
-            total_bytes=job.total_bytes,
-        )
-
-
-@payload_schema.register
-class DownloadCompleteEvent(DownloadEventBase):
-    """Event model for download_complete"""
-
-    __event_name__ = "download_complete"
-
-    download_path: str = Field(description="The local path where the download is saved")
-    total_bytes: int = Field(description="The total number of bytes downloaded")
-
-    @classmethod
-    def build(cls, job: "DownloadJob") -> "DownloadCompleteEvent":
-        assert job.download_path
-        return cls(source=str(job.source), download_path=job.download_path.as_posix(), total_bytes=job.total_bytes)
-
-
-@payload_schema.register
-class DownloadCancelledEvent(DownloadEventBase):
-    """Event model for download_cancelled"""
-
-    __event_name__ = "download_cancelled"
-
-    @classmethod
-    def build(cls, job: "DownloadJob") -> "DownloadCancelledEvent":
-        return cls(source=str(job.source))
-
-
-@payload_schema.register
-class DownloadErrorEvent(DownloadEventBase):
-    """Event model for download_error"""
-
-    __event_name__ = "download_error"
-
-    error_type: str = Field(description="The type of error")
-    error: str = Field(description="The error message")
-
-    @classmethod
-    def build(cls, job: "DownloadJob") -> "DownloadErrorEvent":
-        assert job.error_type
-        assert job.error
-        return cls(source=str(job.source), error_type=job.error_type, error=job.error)
-
-
-class ModelEventBase(EventBase):
-    """Base class for events associated with a model"""
-
-
-@payload_schema.register
-class ModelLoadStartedEvent(ModelEventBase):
-    """Event model for model_load_started"""
-
-    __event_name__ = "model_load_started"
-
-    config: AnyModelConfig = Field(description="The model's config")
-    submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any")
-
-    @classmethod
-    def build(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> "ModelLoadStartedEvent":
-        return cls(config=config, submodel_type=submodel_type)
-
-
-@payload_schema.register
-class ModelLoadCompleteEvent(ModelEventBase):
-    """Event model for model_load_complete"""
-
-    __event_name__ = "model_load_complete"
-
-    config: AnyModelConfig = Field(description="The model's config")
-    submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any")
-
-    @classmethod
-    def build(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> "ModelLoadCompleteEvent":
-        return cls(config=config, submodel_type=submodel_type)
-
-
-@payload_schema.register
-class ModelInstallDownloadStartedEvent(ModelEventBase):
-    """Event model for model_install_download_started"""
-
-    __event_name__ = "model_install_download_started"
-
-    id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
-    local_path: str = Field(description="Where model is downloading to")
-    bytes: int = Field(description="Number of bytes downloaded so far")
-    total_bytes: int = Field(description="Total size of download, including all files")
-    parts: list[dict[str, int | str]] = Field(
-        description="Progress of downloading URLs that comprise the model, if any"
-    )
-
-    @classmethod
-    def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadStartedEvent":
-        parts: list[dict[str, str | int]] = [
-            {
-                "url": str(x.source),
-                "local_path": str(x.download_path),
-                "bytes": x.bytes,
-                "total_bytes": x.total_bytes,
-            }
-            for x in job.download_parts
-        ]
-        return cls(
-            id=job.id,
-            source=str(job.source),
-            local_path=job.local_path.as_posix(),
-            parts=parts,
-            bytes=job.bytes,
-            total_bytes=job.total_bytes,
-        )
-
-
-@payload_schema.register
-class ModelInstallDownloadProgressEvent(ModelEventBase):
-    """Event model for model_install_download_progress"""
-
-    __event_name__ = "model_install_download_progress"
-
-    id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
-    local_path: str = Field(description="Where model is downloading to")
-    bytes: int = Field(description="Number of bytes downloaded so far")
-    total_bytes: int = Field(description="Total size of download, including all files")
-    parts: list[dict[str, int | str]] = Field(
-        description="Progress of downloading URLs that comprise the model, if any"
-    )
-
-    @classmethod
-    def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadProgressEvent":
-        parts: list[dict[str, str | int]] = [
-            {
-                "url": str(x.source),
-                "local_path": str(x.download_path),
-                "bytes": x.bytes,
-                "total_bytes": x.total_bytes,
-            }
-            for x in job.download_parts
-        ]
-        return cls(
-            id=job.id,
-            source=str(job.source),
-            local_path=job.local_path.as_posix(),
-            parts=parts,
-            bytes=job.bytes,
-            total_bytes=job.total_bytes,
-        )
-
-
-@payload_schema.register
-class ModelInstallDownloadsCompleteEvent(ModelEventBase):
-    """Emitted once when an install job becomes active."""
-
-    __event_name__ = "model_install_downloads_complete"
-
-    id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
-
-    @classmethod
-    def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadsCompleteEvent":
-        return cls(id=job.id, source=str(job.source))
-
-
-@payload_schema.register
-class ModelInstallStartedEvent(ModelEventBase):
-    """Event model for model_install_started"""
-
-    __event_name__ = "model_install_started"
-
-    id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
-
-    @classmethod
-    def build(cls, job: "ModelInstallJob") -> "ModelInstallStartedEvent":
-        return cls(id=job.id, source=str(job.source))
-
-
-@payload_schema.register
-class ModelInstallCompleteEvent(ModelEventBase):
-    """Event model for model_install_complete"""
-
-    __event_name__ = "model_install_complete"
-
-    id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
-    key: str = Field(description="Model config record key")
-    total_bytes: Optional[int] = Field(description="Size of the model (may be None for installation of a local path)")
-
-    @classmethod
-    def build(cls, job: "ModelInstallJob") -> "ModelInstallCompleteEvent":
-        assert job.config_out is not None
-        return cls(id=job.id, source=str(job.source), key=(job.config_out.key), total_bytes=job.total_bytes)
-
-
-@payload_schema.register
-class ModelInstallCancelledEvent(ModelEventBase):
-    """Event model for model_install_cancelled"""
-
-    __event_name__ = "model_install_cancelled"
-
-    id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
-
-    @classmethod
-    def build(cls, job: "ModelInstallJob") -> "ModelInstallCancelledEvent":
-        return cls(id=job.id, source=str(job.source))
-
-
-@payload_schema.register
-class ModelInstallErrorEvent(ModelEventBase):
-    """Event model for model_install_error"""
-
-    __event_name__ = "model_install_error"
-
-    id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
-    error_type: str = Field(description="The name of the exception")
-    error: str = Field(description="A text description of the exception")
-
-    @classmethod
-    def build(cls, job: "ModelInstallJob") -> "ModelInstallErrorEvent":
-        assert job.error_type is not None
-        assert job.error is not None
-        return cls(id=job.id, source=str(job.source), error_type=job.error_type, error=job.error)
-
-
-class BulkDownloadEventBase(EventBase):
-    """Base class for events associated with a bulk image download"""
-
-    bulk_download_id: str = Field(description="The ID of the bulk image download")
-    bulk_download_item_id: str = Field(description="The ID of the bulk image download item")
-    bulk_download_item_name: str = Field(description="The name of the bulk image download item")
-
-
-@payload_schema.register
-class BulkDownloadStartedEvent(BulkDownloadEventBase):
-    """Event model for bulk_download_started"""
-
-    __event_name__ = "bulk_download_started"
-
-    @classmethod
-    def build(
-        cls, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
-    ) -> "BulkDownloadStartedEvent":
-        return cls(
-            bulk_download_id=bulk_download_id,
-            bulk_download_item_id=bulk_download_item_id,
-            bulk_download_item_name=bulk_download_item_name,
-        )
-
-
-@payload_schema.register
-class BulkDownloadCompleteEvent(BulkDownloadEventBase):
-    """Event model for bulk_download_complete"""
-
-    __event_name__ = "bulk_download_complete"
-
-    @classmethod
-    def build(
-        cls, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str
-    ) -> "BulkDownloadCompleteEvent":
-        return cls(
-            bulk_download_id=bulk_download_id,
-            bulk_download_item_id=bulk_download_item_id,
-            bulk_download_item_name=bulk_download_item_name,
-        )
-
-
-@payload_schema.register
-class BulkDownloadErrorEvent(BulkDownloadEventBase):
-    """Event model for bulk_download_error"""
-
-    __event_name__ = "bulk_download_error"
-
-    error: str = Field(description="The error message")
-
-    @classmethod
-    def build(
-        cls, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, error: str
-    ) -> "BulkDownloadErrorEvent":
-        return cls(
-            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/events/events_fastapievents.py

@@ -1,47 +0,0 @@
-# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
-
-import asyncio
-import threading
-from queue import Empty, Queue
-
-from fastapi_events.dispatcher import dispatch
-
-from invokeai.app.services.events.events_common import (
-    EventBase,
-)
-
-from .events_base import EventServiceBase
-
-
-class FastAPIEventService(EventServiceBase):
-    def __init__(self, event_handler_id: int) -> None:
-        self.event_handler_id = event_handler_id
-        self._queue = Queue[EventBase | None]()
-        self._stop_event = threading.Event()
-        asyncio.create_task(self._dispatch_from_queue(stop_event=self._stop_event))
-
-        super().__init__()
-
-    def stop(self, *args, **kwargs):
-        self._stop_event.set()
-        self._queue.put(None)
-
-    def dispatch(self, event: EventBase) -> None:
-        self._queue.put(event)
-
-    async def _dispatch_from_queue(self, stop_event: threading.Event):
-        """Get events on from the queue and dispatch them, from the correct thread"""
-        while not stop_event.is_set():
-            try:
-                event = self._queue.get(block=False)
-                if not event:  # Probably stopping
-                    continue
-                # Leave the payloads as live pydantic models
-                dispatch(event, middleware_id=self.event_handler_id, payload_schema_dump=False)
-
-            except Empty:
-                await asyncio.sleep(0.1)
-                pass
-
-            except asyncio.CancelledError as e:
-                raise e  # Raise a proper error

invokeai/app/services/image_files/image_files_base.py

@@ -4,6 +4,9 @@ from typing import Optional
 
 from PIL.Image import Image as PILImageType
 
+from invokeai.app.invocations.fields import MetadataField
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
+
 
 class ImageFileStorageBase(ABC):
     """Low-level service responsible for storing and retrieving image files."""
@@ -30,9 +33,8 @@ class ImageFileStorageBase(ABC):
         self,
         image: PILImageType,
         image_name: str,
-        metadata: Optional[str] = None,
-        workflow: Optional[str] = None,
-        graph: Optional[str] = None,
+        metadata: Optional[MetadataField] = None,
+        workflow: Optional[WorkflowWithoutID] = None,
         thumbnail_size: int = 256,
     ) -> None:
         """Saves an image and a 256x256 WEBP thumbnail. Returns a tuple of the image name, thumbnail name, and created timestamp."""
@@ -44,11 +46,6 @@ class ImageFileStorageBase(ABC):
         pass
 
     @abstractmethod
-    def get_workflow(self, image_name: str) -> Optional[str]:
+    def get_workflow(self, image_name: str) -> Optional[WorkflowWithoutID]:
         """Gets the workflow of an image."""
         pass
-
-    @abstractmethod
-    def get_graph(self, image_name: str) -> Optional[str]:
-        """Gets the graph of an image."""
-        pass

invokeai/app/services/image_files/image_files_disk.py

@@ -7,7 +7,9 @@ from PIL import Image, PngImagePlugin
 from PIL.Image import Image as PILImageType
 from send2trash import send2trash
 
+from invokeai.app.invocations.fields import MetadataField
 from invokeai.app.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
 
 from .image_files_base import ImageFileStorageBase
@@ -54,9 +56,8 @@ class DiskImageFileStorage(ImageFileStorageBase):
         self,
         image: PILImageType,
         image_name: str,
-        metadata: Optional[str] = None,
-        workflow: Optional[str] = None,
-        graph: Optional[str] = None,
+        metadata: Optional[MetadataField] = None,
+        workflow: Optional[WorkflowWithoutID] = None,
         thumbnail_size: int = 256,
     ) -> None:
         try:
@@ -67,14 +68,13 @@ class DiskImageFileStorage(ImageFileStorageBase):
             info_dict = {}
 
             if metadata is not None:
-                info_dict["invokeai_metadata"] = metadata
-                pnginfo.add_text("invokeai_metadata", metadata)
+                metadata_json = metadata.model_dump_json()
+                info_dict["invokeai_metadata"] = metadata_json
+                pnginfo.add_text("invokeai_metadata", metadata_json)
             if workflow is not None:
-                info_dict["invokeai_workflow"] = workflow
-                pnginfo.add_text("invokeai_workflow", workflow)
-            if graph is not None:
-                info_dict["invokeai_graph"] = graph
-                pnginfo.add_text("invokeai_graph", graph)
+                workflow_json = workflow.model_dump_json()
+                info_dict["invokeai_workflow"] = workflow_json
+                pnginfo.add_text("invokeai_workflow", workflow_json)
 
             # When saving the image, the image object's info field is not populated. We need to set it
             image.info = info_dict
@@ -129,18 +129,11 @@ class DiskImageFileStorage(ImageFileStorageBase):
         path = path if isinstance(path, Path) else Path(path)
         return path.exists()
 
-    def get_workflow(self, image_name: str) -> str | None:
+    def get_workflow(self, image_name: str) -> WorkflowWithoutID | None:
         image = self.get(image_name)
         workflow = image.info.get("invokeai_workflow", None)
-        if isinstance(workflow, str):
-            return workflow
-        return None
-
-    def get_graph(self, image_name: str) -> str | None:
-        image = self.get(image_name)
-        graph = image.info.get("invokeai_graph", None)
-        if isinstance(graph, str):
-            return graph
+        if workflow is not None:
+            return WorkflowWithoutID.model_validate_json(workflow)
         return None
 
     def __validate_storage_folders(self) -> None:

invokeai/app/services/image_records/image_records_base.py

@@ -4,7 +4,6 @@ from typing import Optional
 
 from invokeai.app.invocations.fields import MetadataField
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
-from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
 
 from .image_records_common import ImageCategory, ImageRecord, ImageRecordChanges, ResourceOrigin
 
@@ -38,13 +37,10 @@ class ImageRecordStorageBase(ABC):
         self,
         offset: int = 0,
         limit: int = 10,
-        starred_first: bool = True,
-        order_dir: SQLiteDirection = SQLiteDirection.Descending,
         image_origin: Optional[ResourceOrigin] = None,
         categories: Optional[list[ImageCategory]] = None,
         is_intermediate: Optional[bool] = None,
         board_id: Optional[str] = None,
-        search_term: Optional[str] = None,
     ) -> OffsetPaginatedResults[ImageRecord]:
         """Gets a page of image records."""
         pass
@@ -84,7 +80,7 @@ class ImageRecordStorageBase(ABC):
         starred: Optional[bool] = False,
         session_id: Optional[str] = None,
         node_id: Optional[str] = None,
-        metadata: Optional[str] = None,
+        metadata: Optional[MetadataField] = None,
     ) -> datetime:
         """Saves an image record."""
         pass

invokeai/app/services/image_records/image_records_sqlite.py

@@ -5,7 +5,6 @@ from typing import Optional, Union, cast
 
 from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
-from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
 
 from .image_records_base import ImageRecordStorageBase
@@ -145,13 +144,10 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
         self,
         offset: int = 0,
         limit: int = 10,
-        starred_first: bool = True,
-        order_dir: SQLiteDirection = SQLiteDirection.Descending,
         image_origin: Optional[ResourceOrigin] = None,
         categories: Optional[list[ImageCategory]] = None,
         is_intermediate: Optional[bool] = None,
         board_id: Optional[str] = None,
-        search_term: Optional[str] = None,
     ) -> OffsetPaginatedResults[ImageRecord]:
         try:
             self._lock.acquire()
@@ -212,21 +208,9 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
                 """
                 query_params.append(board_id)
 
-            # Search term condition
-            if search_term:
-                query_conditions += """--sql
-                AND images.metadata LIKE ?
-                """
-                query_params.append(f"%{search_term.lower()}%")
-
-            if starred_first:
-                query_pagination = f"""--sql
-                ORDER BY images.starred DESC, images.created_at {order_dir.value} LIMIT ? OFFSET ?
-                """
-            else:
-                query_pagination = f"""--sql
-                ORDER BY images.created_at {order_dir.value} LIMIT ? OFFSET ?
-                """
+            query_pagination = """--sql
+            ORDER BY images.starred DESC, images.created_at DESC LIMIT ? OFFSET ?
+            """
 
             # Final images query with pagination
             images_query += query_conditions + query_pagination + ";"
@@ -344,9 +328,10 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
         starred: Optional[bool] = False,
         session_id: Optional[str] = None,
         node_id: Optional[str] = None,
-        metadata: Optional[str] = None,
+        metadata: Optional[MetadataField] = None,
     ) -> datetime:
         try:
+            metadata_json = metadata.model_dump_json() if metadata is not None else None
             self._lock.acquire()
             self._cursor.execute(
                 """--sql
@@ -373,7 +358,7 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
                     height,
                     node_id,
                     session_id,
-                    metadata,
+                    metadata_json,
                     is_intermediate,
                     starred,
                     has_workflow,

invokeai/app/services/images/images_base.py

@@ -12,7 +12,7 @@ from invokeai.app.services.image_records.image_records_common import (
 )
 from invokeai.app.services.images.images_common import ImageDTO
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
-from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
 
 
 class ImageServiceABC(ABC):
@@ -51,9 +51,8 @@ class ImageServiceABC(ABC):
         session_id: Optional[str] = None,
         board_id: Optional[str] = None,
         is_intermediate: Optional[bool] = False,
-        metadata: Optional[str] = None,
-        workflow: Optional[str] = None,
-        graph: Optional[str] = None,
+        metadata: Optional[MetadataField] = None,
+        workflow: Optional[WorkflowWithoutID] = None,
     ) -> ImageDTO:
         """Creates an image, storing the file and its metadata."""
         pass
@@ -88,12 +87,7 @@ class ImageServiceABC(ABC):
         pass
 
     @abstractmethod
-    def get_workflow(self, image_name: str) -> Optional[str]:
-        """Gets an image's workflow."""
-        pass
-
-    @abstractmethod
-    def get_graph(self, image_name: str) -> Optional[str]:
+    def get_workflow(self, image_name: str) -> Optional[WorkflowWithoutID]:
         """Gets an image's workflow."""
         pass
 
@@ -117,13 +111,10 @@ class ImageServiceABC(ABC):
         self,
         offset: int = 0,
         limit: int = 10,
-        starred_first: bool = True,
-        order_dir: SQLiteDirection = SQLiteDirection.Descending,
         image_origin: Optional[ResourceOrigin] = None,
         categories: Optional[list[ImageCategory]] = None,
         is_intermediate: Optional[bool] = None,
         board_id: Optional[str] = None,
-        search_term: Optional[str] = None,
     ) -> OffsetPaginatedResults[ImageDTO]:
         """Gets a paginated list of image DTOs."""
         pass

invokeai/app/services/images/images_default.py

@@ -5,7 +5,7 @@ from PIL.Image import Image as PILImageType
 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.shared.sqlite.sqlite_common import SQLiteDirection
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID
 
 from ..image_files.image_files_common import (
     ImageFileDeleteException,
@@ -42,9 +42,8 @@ class ImageService(ImageServiceABC):
         session_id: Optional[str] = None,
         board_id: Optional[str] = None,
         is_intermediate: Optional[bool] = False,
-        metadata: Optional[str] = None,
-        workflow: Optional[str] = None,
-        graph: Optional[str] = None,
+        metadata: Optional[MetadataField] = None,
+        workflow: Optional[WorkflowWithoutID] = None,
     ) -> ImageDTO:
         if image_origin not in ResourceOrigin:
             raise InvalidOriginException
@@ -65,7 +64,7 @@ class ImageService(ImageServiceABC):
                 image_category=image_category,
                 width=width,
                 height=height,
-                has_workflow=workflow is not None or graph is not None,
+                has_workflow=workflow is not None,
                 # Meta fields
                 is_intermediate=is_intermediate,
                 # Nullable fields
@@ -74,14 +73,9 @@ class ImageService(ImageServiceABC):
                 session_id=session_id,
             )
             if board_id is not None:
-                try:
-                    self.__invoker.services.board_image_records.add_image_to_board(
-                        board_id=board_id, image_name=image_name
-                    )
-                except Exception as e:
-                    self.__invoker.services.logger.warn(f"Failed to add image to board {board_id}: {str(e)}")
+                self.__invoker.services.board_image_records.add_image_to_board(board_id=board_id, image_name=image_name)
             self.__invoker.services.image_files.save(
-                image_name=image_name, image=image, metadata=metadata, workflow=workflow, graph=graph
+                image_name=image_name, image=image, metadata=metadata, workflow=workflow
             )
             image_dto = self.get_dto(image_name)
 
@@ -163,7 +157,7 @@ class ImageService(ImageServiceABC):
             self.__invoker.services.logger.error("Problem getting image metadata")
             raise e
 
-    def get_workflow(self, image_name: str) -> Optional[str]:
+    def get_workflow(self, image_name: str) -> Optional[WorkflowWithoutID]:
         try:
             return self.__invoker.services.image_files.get_workflow(image_name)
         except ImageFileNotFoundException:
@@ -173,16 +167,6 @@ class ImageService(ImageServiceABC):
             self.__invoker.services.logger.error("Problem getting image workflow")
             raise
 
-    def get_graph(self, image_name: str) -> Optional[str]:
-        try:
-            return self.__invoker.services.image_files.get_graph(image_name)
-        except ImageFileNotFoundException:
-            self.__invoker.services.logger.error("Image file not found")
-            raise
-        except Exception:
-            self.__invoker.services.logger.error("Problem getting image graph")
-            raise
-
     def get_path(self, image_name: str, thumbnail: bool = False) -> str:
         try:
             return str(self.__invoker.services.image_files.get_path(image_name, thumbnail))
@@ -208,25 +192,19 @@ class ImageService(ImageServiceABC):
         self,
         offset: int = 0,
         limit: int = 10,
-        starred_first: bool = True,
-        order_dir: SQLiteDirection = SQLiteDirection.Descending,
         image_origin: Optional[ResourceOrigin] = None,
         categories: Optional[list[ImageCategory]] = None,
         is_intermediate: Optional[bool] = None,
         board_id: Optional[str] = None,
-        search_term: Optional[str] = None,
     ) -> OffsetPaginatedResults[ImageDTO]:
         try:
             results = self.__invoker.services.image_records.get_many(
                 offset,
                 limit,
-                starred_first,
-                order_dir,
                 image_origin,
                 categories,
                 is_intermediate,
                 board_id,
-                search_term,
             )
 
             image_dtos = [

invokeai/app/services/model_install/__init__.py

@@ -1,13 +1,11 @@
 """Initialization file for model install service package."""
 
 from .model_install_base import (
-    ModelInstallServiceBase,
-)
-from .model_install_common import (
     HFModelSource,
     InstallStatus,
     LocalModelSource,
     ModelInstallJob,
+    ModelInstallServiceBase,
     ModelSource,
     UnknownInstallJobException,
     URLModelSource,

invokeai/app/services/model_install/model_install_base.py

@@ -1,19 +1,244 @@
 # Copyright 2023 Lincoln D. Stein and the InvokeAI development team
 """Baseclass definitions for the model installer."""
 
+import re
+import traceback
 from abc import ABC, abstractmethod
+from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Set, Union
 
+from pydantic import BaseModel, Field, PrivateAttr, field_validator
 from pydantic.networks import AnyHttpUrl
+from typing_extensions import Annotated
 
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.app.services.download import DownloadQueueServiceBase
+from invokeai.app.services.download import DownloadJob, DownloadQueueServiceBase
 from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
-from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
 from invokeai.app.services.model_records import ModelRecordServiceBase
-from invokeai.backend.model_manager import AnyModelConfig
+from invokeai.backend.model_manager import AnyModelConfig, ModelRepoVariant
+from invokeai.backend.model_manager.config import ModelSourceType
+from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
+
+
+class InstallStatus(str, Enum):
+    """State of an install job running in the background."""
+
+    WAITING = "waiting"  # waiting to be dequeued
+    DOWNLOADING = "downloading"  # downloading of model files in process
+    DOWNLOADS_DONE = "downloads_done"  # downloading done, waiting to run
+    RUNNING = "running"  # being processed
+    COMPLETED = "completed"  # finished running
+    ERROR = "error"  # terminated with an error message
+    CANCELLED = "cancelled"  # terminated with an error message
+
+
+class ModelInstallPart(BaseModel):
+    url: AnyHttpUrl
+    path: Path
+    bytes: int = 0
+    total_bytes: int = 0
+
+
+class UnknownInstallJobException(Exception):
+    """Raised when the status of an unknown job is requested."""
+
+
+class StringLikeSource(BaseModel):
+    """
+    Base class for model sources, implements functions that lets the source be sorted and indexed.
+
+    These shenanigans let this stuff work:
+
+      source1 = LocalModelSource(path='C:/users/mort/foo.safetensors')
+      mydict = {source1: 'model 1'}
+      assert mydict['C:/users/mort/foo.safetensors'] == 'model 1'
+      assert mydict[LocalModelSource(path='C:/users/mort/foo.safetensors')] == 'model 1'
+
+      source2 = LocalModelSource(path=Path('C:/users/mort/foo.safetensors'))
+      assert source1 == source2
+      assert source1 == 'C:/users/mort/foo.safetensors'
+    """
+
+    def __hash__(self) -> int:
+        """Return hash of the path field, for indexing."""
+        return hash(str(self))
+
+    def __lt__(self, other: object) -> int:
+        """Return comparison of the stringified version, for sorting."""
+        return str(self) < str(other)
+
+    def __eq__(self, other: object) -> bool:
+        """Return equality on the stringified version."""
+        if isinstance(other, Path):
+            return str(self) == other.as_posix()
+        else:
+            return str(self) == str(other)
+
+
+class LocalModelSource(StringLikeSource):
+    """A local file or directory path."""
+
+    path: str | Path
+    inplace: Optional[bool] = False
+    type: Literal["local"] = "local"
+
+    # these methods allow the source to be used in a string-like way,
+    # for example as an index into a dict
+    def __str__(self) -> str:
+        """Return string version of path when string rep needed."""
+        return Path(self.path).as_posix()
+
+
+class HFModelSource(StringLikeSource):
+    """
+    A HuggingFace repo_id with optional variant, sub-folder and access token.
+    Note that the variant option, if not provided to the constructor, will default to fp16, which is
+    what people (almost) always want.
+    """
+
+    repo_id: str
+    variant: Optional[ModelRepoVariant] = ModelRepoVariant.FP16
+    subfolder: Optional[Path] = None
+    access_token: Optional[str] = None
+    type: Literal["hf"] = "hf"
+
+    @field_validator("repo_id")
+    @classmethod
+    def proper_repo_id(cls, v: str) -> str:  # noqa D102
+        if not re.match(r"^([.\w-]+/[.\w-]+)$", v):
+            raise ValueError(f"{v}: invalid repo_id format")
+        return v
+
+    def __str__(self) -> str:
+        """Return string version of repoid when string rep needed."""
+        base: str = self.repo_id
+        if self.variant:
+            base += f":{self.variant or ''}"
+        if self.subfolder:
+            base += f":{self.subfolder}"
+        return base
+
+
+class URLModelSource(StringLikeSource):
+    """A generic URL point to a checkpoint file."""
+
+    url: AnyHttpUrl
+    access_token: Optional[str] = None
+    type: Literal["url"] = "url"
+
+    def __str__(self) -> str:
+        """Return string version of the url when string rep needed."""
+        return str(self.url)
+
+
+ModelSource = Annotated[Union[LocalModelSource, HFModelSource, URLModelSource], Field(discriminator="type")]
+
+MODEL_SOURCE_TO_TYPE_MAP = {
+    URLModelSource: ModelSourceType.Url,
+    HFModelSource: ModelSourceType.HFRepoID,
+    LocalModelSource: ModelSourceType.Path,
+}
+
+
+class ModelInstallJob(BaseModel):
+    """Object that tracks the current status of an install request."""
+
+    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."
+    )
+    config_out: Optional[AnyModelConfig] = Field(
+        default=None, description="After successful installation, this will hold the configuration object."
+    )
+    inplace: bool = Field(
+        default=False, description="Leave model in its current location; otherwise install under models directory"
+    )
+    source: ModelSource = Field(description="Source (URL, repo_id, or local path) of model")
+    local_path: Path = Field(description="Path to locally-downloaded model; may be the same as the source")
+    bytes: int = Field(
+        default=0, description="For a remote model, the number of bytes downloaded so far (may not be available)"
+    )
+    total_bytes: int = Field(default=0, description="Total size of the model to be installed")
+    source_metadata: Optional[AnyModelRepoMetadata] = Field(
+        default=None, description="Metadata provided by the model source"
+    )
+    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)
+
+    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."""
+        self.status = InstallStatus.CANCELLED
+
+    @property
+    def error_type(self) -> Optional[str]:
+        """Class name of the exception that led to status==ERROR."""
+        return self._exception.__class__.__name__ if self._exception else None
+
+    def _format_error(self, exception: Exception) -> str:
+        """Error traceback."""
+        return "".join(traceback.format_exception(exception))
+
+    @property
+    def cancelled(self) -> bool:
+        """Set status to CANCELLED."""
+        return self.status == InstallStatus.CANCELLED
+
+    @property
+    def errored(self) -> bool:
+        """Return true if job has errored."""
+        return self.status == InstallStatus.ERROR
+
+    @property
+    def waiting(self) -> bool:
+        """Return true if job is waiting to run."""
+        return self.status == InstallStatus.WAITING
+
+    @property
+    def downloading(self) -> bool:
+        """Return true if job is downloading."""
+        return self.status == InstallStatus.DOWNLOADING
+
+    @property
+    def downloads_done(self) -> bool:
+        """Return true if job's downloads ae done."""
+        return self.status == InstallStatus.DOWNLOADS_DONE
+
+    @property
+    def running(self) -> bool:
+        """Return true if job is running."""
+        return self.status == InstallStatus.RUNNING
+
+    @property
+    def complete(self) -> bool:
+        """Return true if job completed without errors."""
+        return self.status == InstallStatus.COMPLETED
+
+    @property
+    def in_terminal_state(self) -> bool:
+        """Return true if job is in a terminal state."""
+        return self.status in [InstallStatus.COMPLETED, InstallStatus.ERROR, InstallStatus.CANCELLED]
 
 
 class ModelInstallServiceBase(ABC):
@@ -57,7 +282,7 @@ class ModelInstallServiceBase(ABC):
 
     @property
     @abstractmethod
-    def event_bus(self) -> Optional["EventServiceBase"]:
+    def event_bus(self) -> Optional[EventServiceBase]:
         """Return the event service base object associated with the installer."""
 
     @abstractmethod
@@ -243,11 +468,12 @@ class ModelInstallServiceBase(ABC):
         """
 
     @abstractmethod
-    def download_and_cache_model(self, source: str | AnyHttpUrl) -> Path:
+    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 string representing a URL or repo_id.
+        :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

invokeai/app/services/model_install/model_install_common.py

@@ -1,227 +0,0 @@
-import re
-import traceback
-from enum import Enum
-from pathlib import Path
-from typing import Any, Dict, Literal, Optional, Set, Union
-
-from pydantic import BaseModel, Field, PrivateAttr, field_validator
-from pydantic.networks import AnyHttpUrl
-from typing_extensions import Annotated
-
-from invokeai.app.services.download import DownloadJob, MultiFileDownloadJob
-from invokeai.backend.model_manager import AnyModelConfig, ModelRepoVariant
-from invokeai.backend.model_manager.config import ModelSourceType
-from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
-
-
-class InstallStatus(str, Enum):
-    """State of an install job running in the background."""
-
-    WAITING = "waiting"  # waiting to be dequeued
-    DOWNLOADING = "downloading"  # downloading of model files in process
-    DOWNLOADS_DONE = "downloads_done"  # downloading done, waiting to run
-    RUNNING = "running"  # being processed
-    COMPLETED = "completed"  # finished running
-    ERROR = "error"  # terminated with an error message
-    CANCELLED = "cancelled"  # terminated with an error message
-
-
-class UnknownInstallJobException(Exception):
-    """Raised when the status of an unknown job is requested."""
-
-
-class StringLikeSource(BaseModel):
-    """
-    Base class for model sources, implements functions that lets the source be sorted and indexed.
-
-    These shenanigans let this stuff work:
-
-      source1 = LocalModelSource(path='C:/users/mort/foo.safetensors')
-      mydict = {source1: 'model 1'}
-      assert mydict['C:/users/mort/foo.safetensors'] == 'model 1'
-      assert mydict[LocalModelSource(path='C:/users/mort/foo.safetensors')] == 'model 1'
-
-      source2 = LocalModelSource(path=Path('C:/users/mort/foo.safetensors'))
-      assert source1 == source2
-      assert source1 == 'C:/users/mort/foo.safetensors'
-    """
-
-    def __hash__(self) -> int:
-        """Return hash of the path field, for indexing."""
-        return hash(str(self))
-
-    def __lt__(self, other: object) -> int:
-        """Return comparison of the stringified version, for sorting."""
-        return str(self) < str(other)
-
-    def __eq__(self, other: object) -> bool:
-        """Return equality on the stringified version."""
-        if isinstance(other, Path):
-            return str(self) == other.as_posix()
-        else:
-            return str(self) == str(other)
-
-
-class LocalModelSource(StringLikeSource):
-    """A local file or directory path."""
-
-    path: str | Path
-    inplace: Optional[bool] = False
-    type: Literal["local"] = "local"
-
-    # these methods allow the source to be used in a string-like way,
-    # for example as an index into a dict
-    def __str__(self) -> str:
-        """Return string version of path when string rep needed."""
-        return Path(self.path).as_posix()
-
-
-class HFModelSource(StringLikeSource):
-    """
-    A HuggingFace repo_id with optional variant, sub-folder and access token.
-    Note that the variant option, if not provided to the constructor, will default to fp16, which is
-    what people (almost) always want.
-    """
-
-    repo_id: str
-    variant: Optional[ModelRepoVariant] = ModelRepoVariant.FP16
-    subfolder: Optional[Path] = None
-    access_token: Optional[str] = None
-    type: Literal["hf"] = "hf"
-
-    @field_validator("repo_id")
-    @classmethod
-    def proper_repo_id(cls, v: str) -> str:  # noqa D102
-        if not re.match(r"^([.\w-]+/[.\w-]+)$", v):
-            raise ValueError(f"{v}: invalid repo_id format")
-        return v
-
-    def __str__(self) -> str:
-        """Return string version of repoid when string rep needed."""
-        base: str = self.repo_id
-        if self.variant:
-            base += f":{self.variant or ''}"
-        if self.subfolder:
-            base += f":{self.subfolder}"
-        return base
-
-
-class URLModelSource(StringLikeSource):
-    """A generic URL point to a checkpoint file."""
-
-    url: AnyHttpUrl
-    access_token: Optional[str] = None
-    type: Literal["url"] = "url"
-
-    def __str__(self) -> str:
-        """Return string version of the url when string rep needed."""
-        return str(self.url)
-
-
-ModelSource = Annotated[Union[LocalModelSource, HFModelSource, URLModelSource], Field(discriminator="type")]
-
-MODEL_SOURCE_TO_TYPE_MAP = {
-    URLModelSource: ModelSourceType.Url,
-    HFModelSource: ModelSourceType.HFRepoID,
-    LocalModelSource: ModelSourceType.Path,
-}
-
-
-class ModelInstallJob(BaseModel):
-    """Object that tracks the current status of an install request."""
-
-    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."
-    )
-    config_out: Optional[AnyModelConfig] = Field(
-        default=None, description="After successful installation, this will hold the configuration object."
-    )
-    inplace: bool = Field(
-        default=False, description="Leave model in its current location; otherwise install under models directory"
-    )
-    source: ModelSource = Field(description="Source (URL, repo_id, or local path) of model")
-    local_path: Path = Field(description="Path to locally-downloaded model; may be the same as the source")
-    bytes: int = Field(
-        default=0, description="For a remote model, the number of bytes downloaded so far (may not be available)"
-    )
-    total_bytes: int = Field(default=0, description="Total size of the model to be installed")
-    source_metadata: Optional[AnyModelRepoMetadata] = Field(
-        default=None, description="Metadata provided by the model source"
-    )
-    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)
-    _multifile_job: Optional[MultiFileDownloadJob] = PrivateAttr(default=None)
-    _exception: Optional[Exception] = PrivateAttr(default=None)
-
-    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."""
-        self.status = InstallStatus.CANCELLED
-
-    @property
-    def error_type(self) -> Optional[str]:
-        """Class name of the exception that led to status==ERROR."""
-        return self._exception.__class__.__name__ if self._exception else None
-
-    def _format_error(self, exception: Exception) -> str:
-        """Error traceback."""
-        return "".join(traceback.format_exception(exception))
-
-    @property
-    def cancelled(self) -> bool:
-        """Set status to CANCELLED."""
-        return self.status == InstallStatus.CANCELLED
-
-    @property
-    def errored(self) -> bool:
-        """Return true if job has errored."""
-        return self.status == InstallStatus.ERROR
-
-    @property
-    def waiting(self) -> bool:
-        """Return true if job is waiting to run."""
-        return self.status == InstallStatus.WAITING
-
-    @property
-    def downloading(self) -> bool:
-        """Return true if job is downloading."""
-        return self.status == InstallStatus.DOWNLOADING
-
-    @property
-    def downloads_done(self) -> bool:
-        """Return true if job's downloads ae done."""
-        return self.status == InstallStatus.DOWNLOADS_DONE
-
-    @property
-    def running(self) -> bool:
-        """Return true if job is running."""
-        return self.status == InstallStatus.RUNNING
-
-    @property
-    def complete(self) -> bool:
-        """Return true if job completed without errors."""
-        return self.status == InstallStatus.COMPLETED
-
-    @property
-    def in_terminal_state(self) -> bool:
-        """Return true if job is in a terminal state."""
-        return self.status in [InstallStatus.COMPLETED, InstallStatus.ERROR, InstallStatus.CANCELLED]

invokeai/app/services/model_install/model_install_default.py

@@ -3,26 +3,25 @@
 import locale
 import os
 import re
+import signal
 import threading
 import time
+from hashlib import sha256
 from pathlib import Path
 from queue import Empty, Queue
 from shutil import copyfile, copytree, move, rmtree
 from tempfile import mkdtemp
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Union
 
-import torch
 import yaml
 from huggingface_hub import HfFolder
 from pydantic.networks import AnyHttpUrl
-from pydantic_core import Url
 from requests import Session
 
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.app.services.download import DownloadQueueServiceBase, MultiFileDownloadJob
+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_install.model_install_base import ModelInstallServiceBase
 from invokeai.app.services.model_records import DuplicateModelException, ModelRecordServiceBase
 from invokeai.app.services.model_records.model_records_base import ModelRecordChanges
 from invokeai.backend.model_manager.config import (
@@ -43,16 +42,15 @@ from invokeai.backend.model_manager.metadata.metadata_base import HuggingFaceMet
 from invokeai.backend.model_manager.probe import ModelProbe
 from invokeai.backend.model_manager.search import ModelSearch
 from invokeai.backend.util import InvokeAILogger
-from invokeai.backend.util.catch_sigint import catch_sigint
-from invokeai.backend.util.devices import TorchDevice
-from invokeai.backend.util.util import slugify
+from invokeai.backend.util.devices import choose_precision, choose_torch_device
 
-from .model_install_common import (
+from .model_install_base import (
     MODEL_SOURCE_TO_TYPE_MAP,
     HFModelSource,
     InstallStatus,
     LocalModelSource,
     ModelInstallJob,
+    ModelInstallServiceBase,
     ModelSource,
     StringLikeSource,
     URLModelSource,
@@ -69,7 +67,7 @@ class ModelInstallService(ModelInstallServiceBase):
         app_config: InvokeAIAppConfig,
         record_store: ModelRecordServiceBase,
         download_queue: DownloadQueueServiceBase,
-        event_bus: Optional["EventServiceBase"] = None,
+        event_bus: Optional[EventServiceBase] = None,
         session: Optional[Session] = None,
     ):
         """
@@ -90,7 +88,7 @@ class ModelInstallService(ModelInstallServiceBase):
         self._downloads_changed_event = threading.Event()
         self._install_completed_event = threading.Event()
         self._download_queue = download_queue
-        self._download_cache: Dict[int, ModelInstallJob] = {}
+        self._download_cache: Dict[AnyHttpUrl, ModelInstallJob] = {}
         self._running = False
         self._session = session
         self._install_thread: Optional[threading.Thread] = None
@@ -105,7 +103,7 @@ class ModelInstallService(ModelInstallServiceBase):
         return self._record_store
 
     @property
-    def event_bus(self) -> Optional["EventServiceBase"]:  # noqa D102
+    def event_bus(self) -> Optional[EventServiceBase]:  # noqa D102
         return self._event_bus
 
     # make the invoker optional here because we don't need it and it
@@ -113,6 +111,17 @@ class ModelInstallService(ModelInstallServiceBase):
     def start(self, invoker: Optional[Invoker] = None) -> None:
         """Start the installer thread."""
 
+        # Yes, this is weird. When the installer thread is running, the
+        # thread masks the ^C signal. When we receive a
+        # sigINT, we stop the thread, reset sigINT, and send a new
+        # sigINT to the parent process.
+        def sigint_handler(signum, frame):
+            self.stop()
+            signal.signal(signal.SIGINT, signal.SIG_DFL)
+            signal.raise_signal(signal.SIGINT)
+
+        signal.signal(signal.SIGINT, sigint_handler)
+
         with self._lock:
             if self._running:
                 raise Exception("Attempt to start the installer service twice")
@@ -122,8 +131,7 @@ class ModelInstallService(ModelInstallServiceBase):
             # In normal use, we do not want to scan the models directory - it should never have orphaned models.
             # We should only do the scan when the flag is set (which should only be set when testing).
             if self.app_config.scan_models_on_startup:
-                with catch_sigint():
-                    self._register_orphaned_models()
+                self._register_orphaned_models()
 
             # Check all models' paths and confirm they exist. A model could be missing if it was installed on a volume
             # that isn't currently mounted. In this case, we don't want to delete the model from the database, but we do
@@ -209,12 +217,33 @@ class ModelInstallService(ModelInstallServiceBase):
         access_token: Optional[str] = None,
         inplace: Optional[bool] = False,
     ) -> ModelInstallJob:
-        """Install a model using pattern matching to infer the type of source."""
-        source_obj = self._guess_source(source)
-        if isinstance(source_obj, LocalModelSource):
-            source_obj.inplace = inplace
-        elif isinstance(source_obj, HFModelSource) or isinstance(source_obj, URLModelSource):
-            source_obj.access_token = access_token
+        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), inplace=inplace)
+        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):
+            # Pull the token from config if it exists and matches the URL
+            _token = access_token
+            if _token is None:
+                for pair in self.app_config.remote_api_tokens or []:
+                    if re.search(pair.url_regex, source):
+                        _token = pair.token
+                        break
+            source_obj = URLModelSource(
+                url=AnyHttpUrl(source),
+                access_token=_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
@@ -275,9 +304,8 @@ class ModelInstallService(ModelInstallServiceBase):
     def cancel_job(self, job: ModelInstallJob) -> None:
         """Cancel the indicated job."""
         job.cancel()
-        self._logger.warning(f"Cancelling {job.source}")
-        if dj := job._multifile_job:
-            self._download_queue.cancel_job(dj)
+        with self._lock:
+            self._cancel_download_parts(job)
 
     def prune_jobs(self) -> None:
         """Prune all completed and errored jobs."""
@@ -325,7 +353,7 @@ class ModelInstallService(ModelInstallServiceBase):
                     legacy_config_path = stanza.get("config")
                     if legacy_config_path:
                         # In v3, these paths were relative to the root. Migrate them to be relative to the legacy_conf_dir.
-                        legacy_config_path = self._app_config.root_path / legacy_config_path
+                        legacy_config_path: Path = self._app_config.root_path / legacy_config_path
                         if legacy_config_path.is_relative_to(self._app_config.legacy_conf_path):
                             legacy_config_path = legacy_config_path.relative_to(self._app_config.legacy_conf_path)
                         config["config_path"] = str(legacy_config_path)
@@ -365,95 +393,38 @@ class ModelInstallService(ModelInstallServiceBase):
             rmtree(model_path)
         self.unregister(key)
 
-    @classmethod
-    def _download_cache_path(cls, source: Union[str, AnyHttpUrl], app_config: InvokeAIAppConfig) -> Path:
-        escaped_source = slugify(str(source))
-        return app_config.download_cache_path / escaped_source
-
-    def download_and_cache_model(
+    def download_and_cache(
         self,
-        source: str | AnyHttpUrl,
+        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_path = self._download_cache_path(str(source), self._app_config)
+        model_hash = sha256(str(source).encode("utf-8")).hexdigest()[0:32]
+        model_path = self._app_config.convert_cache_path / model_hash
 
-        # We expect the cache directory to contain one and only one downloaded file or directory.
+        # 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: List[Path] = list(model_path.iterdir())
+            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)
-        model_source = self._guess_source(str(source))
-        remote_files, _ = self._remote_files_from_source(model_source)
-        job = self._multifile_download(
+        job = self._download_queue.download(
+            source=AnyHttpUrl(str(source)),
             dest=model_path,
-            remote_files=remote_files,
-            subfolder=model_source.subfolder if isinstance(model_source, HFModelSource) else None,
+            access_token=access_token,
+            on_progress=TqdmProgress().update,
         )
-        files_string = "file" if len(remote_files) == 1 else "files"
-        self._logger.info(f"Queuing model download: {source} ({len(remote_files)} {files_string})")
-        self._download_queue.wait_for_job(job)
+        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)
 
-    def _remote_files_from_source(
-        self, source: ModelSource
-    ) -> Tuple[List[RemoteModelFile], Optional[AnyModelRepoMetadata]]:
-        metadata = None
-        if isinstance(source, HFModelSource):
-            metadata = HuggingFaceMetadataFetch(self._session).from_id(source.repo_id, source.variant)
-            assert isinstance(metadata, ModelMetadataWithFiles)
-            return (
-                metadata.download_urls(
-                    variant=source.variant or self._guess_variant(),
-                    subfolder=source.subfolder,
-                    session=self._session,
-                ),
-                metadata,
-            )
-
-        if isinstance(source, URLModelSource):
-            try:
-                fetcher = self.get_fetcher_from_url(str(source.url))
-                kwargs: dict[str, Any] = {"session": self._session}
-                metadata = fetcher(**kwargs).from_url(source.url)
-                assert isinstance(metadata, ModelMetadataWithFiles)
-                return metadata.download_urls(session=self._session), metadata
-            except ValueError:
-                pass
-
-            return [RemoteModelFile(url=source.url, path=Path("."), size=0)], None
-
-        raise Exception(f"No files associated with {source}")
-
-    def _guess_source(self, source: str) -> ModelSource:
-        """Turn a source string into a ModelSource object."""
-        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=ModelRepoVariant(match.group(2)) if match.group(2) else None,  # pass None rather than ''
-                subfolder=Path(match.group(3)) if match.group(3) else None,
-            )
-        elif re.match(r"^https?://[^/]+", source):
-            source_obj = URLModelSource(
-                url=Url(source),
-            )
-        else:
-            raise ValueError(f"Unsupported model source: '{source}'")
-        return source_obj
-
     # --------------------------------------------------------------------------------------------
     # Internal functions that manage the installer threads
     # --------------------------------------------------------------------------------------------
@@ -514,19 +485,16 @@ class ModelInstallService(ModelInstallServiceBase):
         job.config_out = self.record_store.get_model(key)
         self._signal_job_completed(job)
 
-    def _set_error(self, install_job: ModelInstallJob, excp: Exception) -> None:
-        multifile_download_job = install_job._multifile_job
-        if multifile_download_job and any(
-            x.content_type is not None and "text/html" in x.content_type for x in multifile_download_job.download_parts
-        ):
-            install_job.set_error(
+    def _set_error(self, job: ModelInstallJob, excp: Exception) -> None:
+        if any(x.content_type is not None and "text/html" in x.content_type for x in job.download_parts):
+            job.set_error(
                 InvalidModelConfigException(
-                    f"At least one file in {install_job.local_path} is an HTML page, not a model. This can happen when an access token is required to download."
+                    f"At least one file in {job.local_path} is an HTML page, not a model. This can happen when an access token is required to download."
                 )
             )
         else:
-            install_job.set_error(excp)
-        self._signal_job_errored(install_job)
+            job.set_error(excp)
+        self._signal_job_errored(job)
 
     # --------------------------------------------------------------------------------------------
     # Internal functions that manage the models directory
@@ -552,6 +520,7 @@ class ModelInstallService(ModelInstallServiceBase):
         This is typically only used during testing with a new DB or when using the memory DB, because those are the
         only situations in which we may have orphaned models in the models directory.
         """
+
         installed_model_paths = {
             (self._app_config.models_path / x.path).resolve() for x in self.record_store.all_models()
         }
@@ -563,13 +532,8 @@ class ModelInstallService(ModelInstallServiceBase):
             if resolved_path in installed_model_paths:
                 return True
             # Skip core models entirely - these aren't registered with the model manager.
-            for special_directory in [
-                self.app_config.models_path / "core",
-                self.app_config.convert_cache_dir,
-                self.app_config.download_cache_dir,
-            ]:
-                if resolved_path.is_relative_to(special_directory):
-                    return False
+            if str(resolved_path).startswith(str(self.app_config.models_path / "core")):
+                return False
             try:
                 model_id = self.register_path(model_path)
                 self._logger.info(f"Registered {model_path.name} with id {model_id}")
@@ -670,10 +634,11 @@ class ModelInstallService(ModelInstallServiceBase):
             self._next_job_id += 1
         return id
 
-    def _guess_variant(self) -> Optional[ModelRepoVariant]:
+    @staticmethod
+    def _guess_variant() -> Optional[ModelRepoVariant]:
         """Guess the best HuggingFace variant type to download."""
-        precision = TorchDevice.choose_torch_dtype()
-        return ModelRepoVariant.FP16 if precision == torch.float16 else None
+        precision = choose_precision(choose_torch_device())
+        return ModelRepoVariant.FP16 if precision == "float16" else None
 
     def _import_local_model(self, source: LocalModelSource, config: Optional[Dict[str, Any]]) -> ModelInstallJob:
         return ModelInstallJob(
@@ -684,15 +649,20 @@ class ModelInstallService(ModelInstallServiceBase):
             inplace=source.inplace or False,
         )
 
-    def _import_from_hf(
-        self,
-        source: HFModelSource,
-        config: Optional[Dict[str, Any]] = None,
-    ) -> ModelInstallJob:
+    def _import_from_hf(self, source: HFModelSource, config: Optional[Dict[str, Any]]) -> ModelInstallJob:
         # Add user's cached access token to HuggingFace requests
-        if source.access_token is None:
-            source.access_token = HfFolder.get_token()
-        remote_files, metadata = self._remote_files_from_source(source)
+        source.access_token = source.access_token or HfFolder.get_token()
+        if not source.access_token:
+            self._logger.info("No HuggingFace access token present; some models may not be downloadable.")
+
+        metadata = HuggingFaceMetadataFetch(self._session).from_id(source.repo_id, source.variant)
+        assert isinstance(metadata, ModelMetadataWithFiles)
+        remote_files = metadata.download_urls(
+            variant=source.variant or self._guess_variant(),
+            subfolder=source.subfolder,
+            session=self._session,
+        )
+
         return self._import_remote_model(
             source=source,
             config=config,
@@ -700,12 +670,22 @@ class ModelInstallService(ModelInstallServiceBase):
             metadata=metadata,
         )
 
-    def _import_from_url(
-        self,
-        source: URLModelSource,
-        config: Optional[Dict[str, Any]],
-    ) -> ModelInstallJob:
-        remote_files, metadata = self._remote_files_from_source(source)
+    def _import_from_url(self, source: URLModelSource, config: Optional[Dict[str, Any]]) -> ModelInstallJob:
+        # URLs from HuggingFace will be handled specially
+        metadata = None
+        fetcher = None
+        try:
+            fetcher = self.get_fetcher_from_url(str(source.url))
+        except ValueError:
+            pass
+        kwargs: dict[str, Any] = {"session": self._session}
+        if fetcher is not None:
+            metadata = fetcher(**kwargs).from_url(source.url)
+        self._logger.debug(f"metadata={metadata}")
+        if metadata and isinstance(metadata, ModelMetadataWithFiles):
+            remote_files = metadata.download_urls(session=self._session)
+        else:
+            remote_files = [RemoteModelFile(url=source.url, path=Path("."), size=0)]
         return self._import_remote_model(
             source=source,
             config=config,
@@ -720,9 +700,12 @@ class ModelInstallService(ModelInstallServiceBase):
         metadata: Optional[AnyModelRepoMetadata],
         config: Optional[Dict[str, Any]],
     ) -> ModelInstallJob:
+        # 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")
-        destdir = Path(
+        tmpdir = Path(
             mkdtemp(
                 dir=self._app_config.models_path,
                 prefix=TMPDIR_PREFIX,
@@ -733,28 +716,52 @@ class ModelInstallService(ModelInstallServiceBase):
             source=source,
             config_in=config or {},
             source_metadata=metadata,
-            local_path=destdir,  # local path may change once the download has started due to content-disposition handling
+            local_path=tmpdir,  # local path may change once the download has started due to content-disposition handling
             bytes=0,
             total_bytes=0,
         )
-        # remember the temporary directory for later removal
-        install_job._install_tmpdir = destdir
-        install_job.total_bytes = sum((x.size or 0) for x in remote_files)
+        # 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 isinstance(source, HFModelSource) and source.subfolder:
+            root = Path(remote_files[0].path.parts[0])
+            subfolder = root / source.subfolder
+        else:
+            root = Path(".")
+            subfolder = Path(".")
 
-        multifile_job = self._multifile_download(
-            remote_files=remote_files,
-            dest=destdir,
-            subfolder=source.subfolder if isinstance(source, HFModelSource) else None,
-            access_token=source.access_token,
-            submit_job=False,  # Important! Don't submit the job until we have set our _download_cache dict
-        )
-        self._download_cache[multifile_job.id] = install_job
-        install_job._multifile_job = multifile_job
+        # 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
+        assert install_job.total_bytes is not None  # to avoid type checking complaints in the loop below
 
-        files_string = "file" if len(remote_files) == 1 else "files"
-        self._logger.info(f"Queueing model install: {source} ({len(remote_files)} {files_string})")
+        files_string = "file" if len(remote_files) == 1 else "file"
+        self._logger.info(f"Queuing model install: {source} ({len(remote_files)} {files_string})")
         self._logger.debug(f"remote_files={remote_files}")
-        self._download_queue.submit_multifile_download(multifile_job)
+        for model_file in remote_files:
+            url = model_file.url
+            path = root / model_file.path.relative_to(subfolder)
+            self._logger.debug(f"Downloading {url} => {path}")
+            install_job.total_bytes += model_file.size
+            assert hasattr(source, "access_token")
+            dest = tmpdir / path.parent
+            dest.mkdir(parents=True, exist_ok=True)
+            download_job = DownloadJob(
+                source=url,
+                dest=dest,
+                access_token=source.access_token,
+            )
+            self._download_cache[download_job.source] = install_job  # matches a download job to an install job
+            install_job.download_parts.add(download_job)
+
+            self._download_queue.submit_download_job(
+                download_job,
+                on_start=self._download_started_callback,
+                on_progress=self._download_progress_callback,
+                on_complete=self._download_complete_callback,
+                on_error=self._download_error_callback,
+                on_cancelled=self._download_cancelled_callback,
+            )
         return install_job
 
     def _stat_size(self, path: Path) -> int:
@@ -766,104 +773,87 @@ class ModelInstallService(ModelInstallServiceBase):
                 size += sum(self._stat_size(Path(root, x)) for x in files)
         return size
 
-    def _multifile_download(
-        self,
-        remote_files: List[RemoteModelFile],
-        dest: Path,
-        subfolder: Optional[Path] = None,
-        access_token: Optional[str] = None,
-        submit_job: bool = True,
-    ) -> MultiFileDownloadJob:
-        # HuggingFace repo subfolders are a little tricky. If the name of the model is "sdxl-turbo", and
-        # we are installing the "vae" subfolder, we do not want to create an additional folder level, such
-        # as "sdxl-turbo/vae", nor do we want to put the contents of the vae folder directly into "sdxl-turbo".
-        # So what we do is to synthesize a folder named "sdxl-turbo_vae" here.
-        if subfolder:
-            top = Path(remote_files[0].path.parts[0])  # e.g. "sdxl-turbo/"
-            path_to_remove = top / subfolder.parts[-1]  # sdxl-turbo/vae/
-            path_to_add = Path(f"{top}_{subfolder}")
-        else:
-            path_to_remove = Path(".")
-            path_to_add = Path(".")
-
-        parts: List[RemoteModelFile] = []
-        for model_file in remote_files:
-            assert model_file.size is not None
-            parts.append(
-                RemoteModelFile(
-                    url=model_file.url,  # if a subfolder, then sdxl-turbo_vae/config.json
-                    path=path_to_add / model_file.path.relative_to(path_to_remove),
-                )
-            )
-
-        return self._download_queue.multifile_download(
-            parts=parts,
-            dest=dest,
-            access_token=access_token,
-            submit_job=submit_job,
-            on_start=self._download_started_callback,
-            on_progress=self._download_progress_callback,
-            on_complete=self._download_complete_callback,
-            on_error=self._download_error_callback,
-            on_cancelled=self._download_cancelled_callback,
-        )
-
     # ------------------------------------------------------------------
     # Callbacks are executed by the download queue in a separate thread
     # ------------------------------------------------------------------
-    def _download_started_callback(self, download_job: MultiFileDownloadJob) -> None:
+    def _download_started_callback(self, download_job: DownloadJob) -> None:
+        self._logger.info(f"Model download started: {download_job.source}")
         with self._lock:
-            if install_job := self._download_cache.get(download_job.id, None):
-                install_job.status = InstallStatus.DOWNLOADING
+            install_job = self._download_cache[download_job.source]
+            install_job.status = InstallStatus.DOWNLOADING
 
-                if install_job.local_path == install_job._install_tmpdir:  # first time
-                    assert download_job.download_path
-                    install_job.local_path = download_job.download_path
-                install_job.download_parts = download_job.download_parts
-                install_job.bytes = sum(x.bytes for x in download_job.download_parts)
-                install_job.total_bytes = download_job.total_bytes
-                self._signal_job_download_started(install_job)
+            assert download_job.download_path
+            if install_job.local_path == install_job._install_tmpdir:
+                partial_path = download_job.download_path.relative_to(install_job._install_tmpdir)
+                dest_name = partial_path.parts[0]
+                install_job.local_path = install_job._install_tmpdir / dest_name
 
-    def _download_progress_callback(self, download_job: MultiFileDownloadJob) -> None:
+            # Update the total bytes count for remote sources.
+            if not install_job.total_bytes:
+                install_job.total_bytes = sum(x.total_bytes for x in install_job.download_parts)
+
+    def _download_progress_callback(self, download_job: DownloadJob) -> None:
         with self._lock:
-            if install_job := self._download_cache.get(download_job.id, None):
-                if install_job.cancelled:  # This catches the case in which the caller directly calls job.cancel()
-                    self._download_queue.cancel_job(download_job)
-                else:
-                    # update sizes
-                    install_job.bytes = sum(x.bytes for x in download_job.download_parts)
-                    install_job.total_bytes = sum(x.total_bytes for x in download_job.download_parts)
-                    self._signal_job_downloading(install_job)
+            install_job = self._download_cache[download_job.source]
+            if install_job.cancelled:  # This catches the case in which the caller directly calls job.cancel()
+                self._cancel_download_parts(install_job)
+            else:
+                # update sizes
+                install_job.bytes = sum(x.bytes for x in install_job.download_parts)
+                self._signal_job_downloading(install_job)
 
-    def _download_complete_callback(self, download_job: MultiFileDownloadJob) -> None:
+    def _download_complete_callback(self, download_job: DownloadJob) -> None:
+        self._logger.info(f"Model download complete: {download_job.source}")
         with self._lock:
-            if install_job := self._download_cache.pop(download_job.id, None):
+            install_job = self._download_cache[download_job.source]
+
+            # are there any more active jobs left in this task?
+            if install_job.downloading and all(x.complete for x in install_job.download_parts):
                 self._signal_job_downloads_done(install_job)
-                self._put_in_queue(install_job)  # this starts the installation and registration
+                self._put_in_queue(install_job)
 
-                # Let other threads know that the number of downloads has changed
-                self._downloads_changed_event.set()
+            # Let other threads know that the number of downloads has changed
+            self._download_cache.pop(download_job.source, None)
+            self._downloads_changed_event.set()
 
-    def _download_error_callback(self, download_job: MultiFileDownloadJob, excp: Optional[Exception] = None) -> None:
+    def _download_error_callback(self, download_job: DownloadJob, excp: Optional[Exception] = None) -> None:
         with self._lock:
-            if install_job := self._download_cache.pop(download_job.id, None):
-                assert excp is not None
-                install_job.set_error(excp)
-                self._download_queue.cancel_job(download_job)
+            install_job = self._download_cache.pop(download_job.source, None)
+            assert install_job is not None
+            assert excp is not None
+            install_job.set_error(excp)
+            self._logger.error(
+                f"Cancelling {install_job.source} due to an error while downloading {download_job.source}: {str(excp)}"
+            )
+            self._cancel_download_parts(install_job)
 
-                # Let other threads know that the number of downloads has changed
-                self._downloads_changed_event.set()
+            # Let other threads know that the number of downloads has changed
+            self._downloads_changed_event.set()
 
-    def _download_cancelled_callback(self, download_job: MultiFileDownloadJob) -> None:
+    def _download_cancelled_callback(self, download_job: DownloadJob) -> None:
         with self._lock:
-            if install_job := self._download_cache.pop(download_job.id, None):
-                self._downloads_changed_event.set()
-                # if install job has already registered an error, then do not replace its status with cancelled
-                if not install_job.errored:
-                    install_job.cancel()
+            install_job = self._download_cache.pop(download_job.source, None)
+            if not install_job:
+                return
+            self._downloads_changed_event.set()
+            self._logger.warning(f"Model download canceled: {download_job.source}")
+            # if install job has already registered an error, then do not replace its status with cancelled
+            if not install_job.errored:
+                install_job.cancel()
+            self._cancel_download_parts(install_job)
 
-                # Let other threads know that the number of downloads has changed
-                self._downloads_changed_event.set()
+            # Let other threads know that the number of downloads has changed
+            self._downloads_changed_event.set()
+
+    def _cancel_download_parts(self, install_job: ModelInstallJob) -> None:
+        # on multipart downloads, _cancel_components() will get called repeatedly from the download callbacks
+        # do not lock here because it gets called within a locked context
+        for s in install_job.download_parts:
+            self._download_queue.cancel_job(s)
+
+        if all(x.in_terminal_state for x in install_job.download_parts):
+            # When all parts have reached their terminal state, we finalize the job to clean up the temporary directory and other resources
+            self._put_in_queue(install_job)
 
     # ------------------------------------------------------------------------------------------------
     # Internal methods that put events on the event bus
@@ -872,27 +862,35 @@ class ModelInstallService(ModelInstallServiceBase):
         job.status = InstallStatus.RUNNING
         self._logger.info(f"Model install started: {job.source}")
         if self._event_bus:
-            self._event_bus.emit_model_install_started(job)
-
-    def _signal_job_download_started(self, job: ModelInstallJob) -> None:
-        if self._event_bus:
-            assert job._multifile_job is not None
-            assert job.bytes is not None
-            assert job.total_bytes is not None
-            self._event_bus.emit_model_install_download_started(job)
+            self._event_bus.emit_model_install_running(str(job.source))
 
     def _signal_job_downloading(self, job: ModelInstallJob) -> None:
         if self._event_bus:
-            assert job._multifile_job is not None
+            parts: List[Dict[str, str | int]] = [
+                {
+                    "url": str(x.source),
+                    "local_path": str(x.download_path),
+                    "bytes": x.bytes,
+                    "total_bytes": x.total_bytes,
+                }
+                for x in job.download_parts
+            ]
             assert job.bytes is not None
             assert job.total_bytes is not None
-            self._event_bus.emit_model_install_download_progress(job)
+            self._event_bus.emit_model_install_downloading(
+                str(job.source),
+                local_path=job.local_path.as_posix(),
+                parts=parts,
+                bytes=job.bytes,
+                total_bytes=job.total_bytes,
+                id=job.id,
+            )
 
     def _signal_job_downloads_done(self, job: ModelInstallJob) -> None:
         job.status = InstallStatus.DOWNLOADS_DONE
         self._logger.info(f"Model download complete: {job.source}")
         if self._event_bus:
-            self._event_bus.emit_model_install_downloads_complete(job)
+            self._event_bus.emit_model_install_downloads_done(str(job.source))
 
     def _signal_job_completed(self, job: ModelInstallJob) -> None:
         job.status = InstallStatus.COMPLETED
@@ -902,28 +900,25 @@ class ModelInstallService(ModelInstallServiceBase):
         if self._event_bus:
             assert job.local_path is not None
             assert job.config_out is not None
-            self._event_bus.emit_model_install_complete(job)
+            key = job.config_out.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.error(f"Model install error: {job.source}\n{job.error_type}: {job.error}")
         if self._event_bus:
-            assert job.error_type is not None
-            assert job.error is not None
-            self._event_bus.emit_model_install_error(job)
+            error_type = job.error_type
+            error = job.error
+            assert error_type is not None
+            assert error is not None
+            self._event_bus.emit_model_install_error(str(job.source), error_type, error, id=job.id)
 
     def _signal_job_cancelled(self, job: ModelInstallJob) -> None:
         self._logger.info(f"Model install canceled: {job.source}")
         if self._event_bus:
-            self._event_bus.emit_model_install_cancelled(job)
+            self._event_bus.emit_model_install_cancelled(str(job.source), id=job.id)
 
     @staticmethod
-    def get_fetcher_from_url(url: str) -> Type[ModelMetadataFetchBase]:
-        """
-        Return a metadata fetcher appropriate for provided url.
-
-        This used to be more useful, but the number of supported model
-        sources has been reduced to HuggingFace alone.
-        """
+    def get_fetcher_from_url(url: str) -> ModelMetadataFetchBase:
         if re.match(r"^https?://huggingface.co/[^/]+/[^/]+$", url.lower()):
             return HuggingFaceMetadataFetch
         raise ValueError(f"Unsupported model source: '{url}'")

invokeai/app/services/model_load/model_load_base.py

@@ -2,11 +2,12 @@
 """Base class for model loader."""
 
 from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Callable, Optional
+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.load_base import LoadedModel, LoadedModelWithoutConfig
+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
 
 
@@ -14,12 +15,18 @@ class ModelLoadServiceBase(ABC):
     """Wrapper around AnyModelLoader."""
 
     @abstractmethod
-    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+    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
@@ -27,25 +34,7 @@ class ModelLoadServiceBase(ABC):
     def ram_cache(self) -> ModelCacheBase[AnyModel]:
         """Return the RAM cache used by this loader."""
 
+    @property
     @abstractmethod
-    def load_model_from_path(
-        self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
-    ) -> LoadedModelWithoutConfig:
-        """
-        Load the model file or directory located at the indicated Path.
-
-        This will load an arbitrary model file into the RAM cache. If the optional loader
-        argument is provided, the loader will be invoked to load the model into
-        memory. Otherwise the method will call safetensors.torch.load_file() or
-        torch.load() as appropriate to the file suffix.
-
-        Be aware that this returns a LoadedModelWithoutConfig object, which is the same as
-        LoadedModel, but without the config attribute.
-
-        Args:
-          model_path: A pathlib.Path to a checkpoint-style models file
-          loader: A Callable that expects a Path and returns a Dict[str, Tensor]
-
-        Returns:
-          A LoadedModel object.
-        """
+    def convert_cache(self) -> ModelConvertCacheBase:
+        """Return the checkpoint convert cache used by this loader."""

invokeai/app/services/model_load/model_load_default.py

@@ -1,21 +1,19 @@
 # Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Team
 """Implementation of model loader service."""
 
-from pathlib import Path
-from typing import Callable, Optional
-
-from picklescan.scanner import scan_file_path
-from safetensors.torch import load_file as safetensors_load_file
-from torch import load as torch_load
+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.load_base import LoadedModel, LoadedModelWithoutConfig
+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.model_manager.load.model_loader_registry import ModelLoaderRegistry
-from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
-from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 
 from .model_load_base import ModelLoadServiceBase
@@ -28,7 +26,8 @@ class ModelLoadService(ModelLoadServiceBase):
         self,
         app_config: InvokeAIAppConfig,
         ram_cache: ModelCacheBase[AnyModel],
-        registry: ModelLoaderRegistry,
+        convert_cache: ModelConvertCacheBase,
+        registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry,
     ):
         """Initialize the model load service."""
         logger = InvokeAILogger.get_logger(self.__class__.__name__)
@@ -36,6 +35,7 @@ class ModelLoadService(ModelLoadServiceBase):
         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:
@@ -46,65 +46,73 @@ class ModelLoadService(ModelLoadServiceBase):
         """Return the RAM cache used by this loader."""
         return self._ram_cache
 
-    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+    @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
         """
-
-        # We don't have an invoker during testing
-        # TODO(psyche): Mock this method on the invoker in the tests
-        if hasattr(self, "_invoker"):
-            self._invoker.services.events.emit_model_load_started(model_config, submodel_type)
+        if context_data:
+            self._emit_load_event(
+                context_data=context_data,
+                model_config=model_config,
+                submodel_type=submodel_type,
+            )
 
         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 hasattr(self, "_invoker"):
-            self._invoker.services.events.emit_model_load_complete(model_config, submodel_type)
-
+        if context_data:
+            self._emit_load_event(
+                context_data=context_data,
+                model_config=model_config,
+                submodel_type=submodel_type,
+                loaded=True,
+            )
         return loaded_model
 
-    def load_model_from_path(
-        self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
-    ) -> LoadedModelWithoutConfig:
-        cache_key = str(model_path)
-        ram_cache = self.ram_cache
-        try:
-            return LoadedModelWithoutConfig(_locker=ram_cache.get(key=cache_key))
-        except IndexError:
-            pass
+    def _emit_load_event(
+        self,
+        context_data: InvocationContextData,
+        model_config: AnyModelConfig,
+        loaded: Optional[bool] = False,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> None:
+        if not self._invoker:
+            return
 
-        def torch_load_file(checkpoint: Path) -> AnyModel:
-            scan_result = scan_file_path(checkpoint)
-            if scan_result.infected_files != 0:
-                raise Exception("The model at {checkpoint} is potentially infected by malware. Aborting load.")
-            result = torch_load(checkpoint, map_location="cpu")
-            return result
-
-        def diffusers_load_directory(directory: Path) -> AnyModel:
-            load_class = GenericDiffusersLoader(
-                app_config=self._app_config,
-                logger=self._logger,
-                ram_cache=self._ram_cache,
-                convert_cache=self.convert_cache,
-            ).get_hf_load_class(directory)
-            return load_class.from_pretrained(model_path, torch_dtype=TorchDevice.choose_torch_dtype())
-
-        loader = loader or (
-            diffusers_load_directory
-            if model_path.is_dir()
-            else torch_load_file
-            if model_path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin"))
-            else lambda path: safetensors_load_file(path, device="cpu")
-        )
-        assert loader is not None
-        raw_model = loader(model_path)
-        ram_cache.put(key=cache_key, model=raw_model)
-        return LoadedModelWithoutConfig(_locker=ram_cache.get(key=cache_key))
+        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,
+                submodel_type=submodel_type,
+            )
+        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,
+                submodel_type=submodel_type,
+            )

invokeai/app/services/model_manager/__init__.py

@@ -0,0 +1,17 @@
+"""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_default.py

@@ -1,15 +1,12 @@
 # Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
 """Implementation of ModelManagerServiceBase."""
 
-from typing import Optional
-
 import torch
 from typing_extensions import Self
 
 from invokeai.app.services.invoker import Invoker
-from invokeai.backend.model_manager.load.model_cache.model_cache_default import ModelCache
-from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
-from invokeai.backend.util.devices import TorchDevice
+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
@@ -70,7 +67,7 @@ class ModelManagerService(ModelManagerServiceBase):
         model_record_service: ModelRecordServiceBase,
         download_queue: DownloadQueueServiceBase,
         events: EventServiceBase,
-        execution_device: Optional[torch.device] = None,
+        execution_device: torch.device = choose_torch_device(),
     ) -> Self:
         """
         Construct the model manager service instance.
@@ -85,11 +82,13 @@ class ModelManagerService(ModelManagerServiceBase):
             max_vram_cache_size=app_config.vram,
             lazy_offloading=app_config.lazy_offload,
             logger=logger,
-            execution_device=execution_device or TorchDevice.choose_torch_device(),
+            execution_device=execution_device,
         )
+        convert_cache = ModelConvertCache(cache_path=app_config.convert_cache_path, max_size=app_config.convert_cache)
         loader = ModelLoadService(
             app_config=app_config,
             ram_cache=ram_cache,
+            convert_cache=convert_cache,
             registry=ModelLoaderRegistry,
         )
         installer = ModelInstallService(

invokeai/app/services/model_records/model_records_base.py

@@ -12,13 +12,15 @@ from pydantic import BaseModel, Field
 
 from invokeai.app.services.shared.pagination import PaginatedResults
 from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
-from invokeai.backend.model_manager.config import (
+from invokeai.backend.model_manager import (
     AnyModelConfig,
     BaseModelType,
-    ControlAdapterDefaultSettings,
-    MainModelDefaultSettings,
     ModelFormat,
     ModelType,
+)
+from invokeai.backend.model_manager.config import (
+    ControlAdapterDefaultSettings,
+    MainModelDefaultSettings,
     ModelVariantType,
     SchedulerPredictionType,
 )

invokeai/app/services/object_serializer/object_serializer_disk.py

@@ -1,6 +1,6 @@
-import shutil
 import tempfile
 import typing
+from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional, TypeVar
 
@@ -17,6 +17,12 @@ if TYPE_CHECKING:
 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`.
 
@@ -29,12 +35,6 @@ class ObjectSerializerDisk(ObjectSerializerBase[T]):
         self._ephemeral = ephemeral
         self._base_output_dir = output_dir
         self._base_output_dir.mkdir(parents=True, exist_ok=True)
-
-        if self._ephemeral:
-            # Remove dangling tempdirs that might have been left over from an earlier unplanned shutdown.
-            for temp_dir in filter(Path.is_dir, self._base_output_dir.glob("tmp*")):
-                shutil.rmtree(temp_dir)
-
         # 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

invokeai/app/services/session_processor/session_processor_base.py

@@ -1,49 +1,6 @@
 from abc import ABC, abstractmethod
-from threading import Event
-from typing import Optional, Protocol
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
-from invokeai.app.services.invocation_services import InvocationServices
 from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus
-from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
-from invokeai.app.util.profiler import Profiler
-
-
-class SessionRunnerBase(ABC):
-    """
-    Base class for session runner.
-    """
-
-    @abstractmethod
-    def start(self, services: InvocationServices, cancel_event: Event, profiler: Optional[Profiler] = None) -> None:
-        """Starts the session runner.
-
-        Args:
-            services: The invocation services.
-            cancel_event: The cancel event.
-            profiler: The profiler to use for session profiling via cProfile. Omit to disable profiling. Basic session
-                stats will be still be recorded and logged when profiling is disabled.
-        """
-        pass
-
-    @abstractmethod
-    def run(self, queue_item: SessionQueueItem) -> None:
-        """Runs a session.
-
-        Args:
-            queue_item: The session to run.
-        """
-        pass
-
-    @abstractmethod
-    def run_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem) -> None:
-        """Run a single node in the graph.
-
-        Args:
-            invocation: The invocation to run.
-            queue_item: The session queue item.
-        """
-        pass
 
 
 class SessionProcessorBase(ABC):
@@ -69,85 +26,3 @@ class SessionProcessorBase(ABC):
     def get_status(self) -> SessionProcessorStatus:
         """Gets the status of the session processor"""
         pass
-
-
-class OnBeforeRunNode(Protocol):
-    def __call__(self, invocation: BaseInvocation, queue_item: SessionQueueItem) -> None:
-        """Callback to run before executing a node.
-
-        Args:
-            invocation: The invocation that will be executed.
-            queue_item: The session queue item.
-        """
-        ...
-
-
-class OnAfterRunNode(Protocol):
-    def __call__(self, invocation: BaseInvocation, queue_item: SessionQueueItem, output: BaseInvocationOutput) -> None:
-        """Callback to run before executing a node.
-
-        Args:
-            invocation: The invocation that was executed.
-            queue_item: The session queue item.
-        """
-        ...
-
-
-class OnNodeError(Protocol):
-    def __call__(
-        self,
-        invocation: BaseInvocation,
-        queue_item: SessionQueueItem,
-        error_type: str,
-        error_message: str,
-        error_traceback: str,
-    ) -> None:
-        """Callback to run when a node has an error.
-
-        Args:
-            invocation: The invocation that errored.
-            queue_item: The session queue item.
-            error_type: The type of error, e.g. "ValueError".
-            error_message: The error message, e.g. "Invalid value".
-            error_traceback: The stringified error traceback.
-        """
-        ...
-
-
-class OnBeforeRunSession(Protocol):
-    def __call__(self, queue_item: SessionQueueItem) -> None:
-        """Callback to run before executing a session.
-
-        Args:
-            queue_item: The session queue item.
-        """
-        ...
-
-
-class OnAfterRunSession(Protocol):
-    def __call__(self, queue_item: SessionQueueItem) -> None:
-        """Callback to run after executing a session.
-
-        Args:
-            queue_item: The session queue item.
-        """
-        ...
-
-
-class OnNonFatalProcessorError(Protocol):
-    def __call__(
-        self,
-        queue_item: Optional[SessionQueueItem],
-        error_type: str,
-        error_message: str,
-        error_traceback: str,
-    ) -> None:
-        """Callback to run when a non-fatal error occurs in the processor.
-
-        Args:
-            queue_item: The session queue item, if one was being executed when the error occurred.
-            error_type: The type of error, e.g. "ValueError".
-            error_message: The error message, e.g. "Invalid value".
-            error_traceback: The stringified error traceback.
-        """
-        ...

invokeai/app/services/session_processor/session_processor_default.py

@@ -4,325 +4,24 @@ from threading import BoundedSemaphore, Thread
 from threading import Event as ThreadEvent
 from typing import Optional
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
-from invokeai.app.services.events.events_common import (
-    BatchEnqueuedEvent,
-    FastAPIEvent,
-    QueueClearedEvent,
-    QueueItemStatusChangedEvent,
-    register_events,
-)
+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_base import (
-    OnAfterRunNode,
-    OnAfterRunSession,
-    OnBeforeRunNode,
-    OnBeforeRunSession,
-    OnNodeError,
-    OnNonFatalProcessorError,
-)
 from invokeai.app.services.session_processor.session_processor_common import CanceledException
-from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem, SessionQueueItemNotFoundError
-from invokeai.app.services.shared.graph import NodeInputError
+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 InvocationServices, SessionProcessorBase, SessionRunnerBase
+from .session_processor_base import SessionProcessorBase
 from .session_processor_common import SessionProcessorStatus
 
 
-class DefaultSessionRunner(SessionRunnerBase):
-    """Processes a single session's invocations."""
-
-    def __init__(
-        self,
-        on_before_run_session_callbacks: Optional[list[OnBeforeRunSession]] = None,
-        on_before_run_node_callbacks: Optional[list[OnBeforeRunNode]] = None,
-        on_after_run_node_callbacks: Optional[list[OnAfterRunNode]] = None,
-        on_node_error_callbacks: Optional[list[OnNodeError]] = None,
-        on_after_run_session_callbacks: Optional[list[OnAfterRunSession]] = None,
-    ):
-        """
-        Args:
-            on_before_run_session_callbacks: Callbacks to run before the session starts.
-            on_before_run_node_callbacks: Callbacks to run before each node starts.
-            on_after_run_node_callbacks: Callbacks to run after each node completes.
-            on_node_error_callbacks: Callbacks to run when a node errors.
-            on_after_run_session_callbacks: Callbacks to run after the session completes.
-        """
-
-        self._on_before_run_session_callbacks = on_before_run_session_callbacks or []
-        self._on_before_run_node_callbacks = on_before_run_node_callbacks or []
-        self._on_after_run_node_callbacks = on_after_run_node_callbacks or []
-        self._on_node_error_callbacks = on_node_error_callbacks or []
-        self._on_after_run_session_callbacks = on_after_run_session_callbacks or []
-
-    def start(self, services: InvocationServices, cancel_event: ThreadEvent, profiler: Optional[Profiler] = None):
-        self._services = services
-        self._cancel_event = cancel_event
-        self._profiler = profiler
-
-    def _is_canceled(self) -> bool:
-        """Check if the cancel event is set. This is also passed to the invocation context builder and called during
-        denoising to check if the session has been canceled."""
-        return self._cancel_event.is_set()
-
-    def run(self, queue_item: SessionQueueItem):
-        # Exceptions raised outside `run_node` are handled by the processor. There is no need to catch them here.
-
-        self._on_before_run_session(queue_item=queue_item)
-
-        # Loop over invocations until the session is complete or canceled
-        while True:
-            try:
-                invocation = queue_item.session.next()
-            # Anything other than a `NodeInputError` is handled as a processor error
-            except NodeInputError as e:
-                error_type = e.__class__.__name__
-                error_message = str(e)
-                error_traceback = traceback.format_exc()
-                self._on_node_error(
-                    invocation=e.node,
-                    queue_item=queue_item,
-                    error_type=error_type,
-                    error_message=error_message,
-                    error_traceback=error_traceback,
-                )
-                break
-
-            if invocation is None or self._is_canceled():
-                break
-
-            self.run_node(invocation, queue_item)
-
-            # The session is complete if all invocations have been run or there is an error on the session.
-            # At this time, the queue item may be canceled, but the object itself here won't be updated yet. We must
-            # use the cancel event to check if the session is canceled.
-            if (
-                queue_item.session.is_complete()
-                or self._is_canceled()
-                or queue_item.status in ["failed", "canceled", "completed"]
-            ):
-                break
-
-        self._on_after_run_session(queue_item=queue_item)
-
-    def run_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem):
-        try:
-            # Any unhandled exception in this scope is an invocation error & will fail the graph
-            with self._services.performance_statistics.collect_stats(invocation, queue_item.session_id):
-                self._on_before_run_node(invocation, queue_item)
-
-                data = InvocationContextData(
-                    invocation=invocation,
-                    source_invocation_id=queue_item.session.prepared_source_mapping[invocation.id],
-                    queue_item=queue_item,
-                )
-                context = build_invocation_context(
-                    data=data,
-                    services=self._services,
-                    is_canceled=self._is_canceled,
-                )
-
-                # Invoke the node
-                output = invocation.invoke_internal(context=context, services=self._services)
-                # Save output and history
-                queue_item.session.complete(invocation.id, output)
-
-                self._on_after_run_node(invocation, queue_item, output)
-
-        except KeyboardInterrupt:
-            # TODO(psyche): This is expected to be caught in the main thread. Do we need to catch this here?
-            pass
-        except CanceledException:
-            # A CanceledException is raised during the denoising step callback if the cancel event is set. We don't need
-            # to do any handling here, and no error should be set - just pass and the cancellation will be handled
-            # correctly in the next iteration of the session runner loop.
-            #
-            # See the comment in the processor's `_on_queue_item_status_changed()` method for more details on how we
-            # handle cancellation.
-            pass
-        except Exception as e:
-            error_type = e.__class__.__name__
-            error_message = str(e)
-            error_traceback = traceback.format_exc()
-            self._on_node_error(
-                invocation=invocation,
-                queue_item=queue_item,
-                error_type=error_type,
-                error_message=error_message,
-                error_traceback=error_traceback,
-            )
-
-    def _on_before_run_session(self, queue_item: SessionQueueItem) -> None:
-        """Called before a session is run.
-
-        - Start the profiler if profiling is enabled.
-        - Run any callbacks registered for this event.
-        """
-
-        self._services.logger.debug(
-            f"On before run session: queue item {queue_item.item_id}, session {queue_item.session_id}"
-        )
-
-        # If profiling is enabled, start the profiler
-        if self._profiler is not None:
-            self._profiler.start(profile_id=queue_item.session_id)
-
-        for callback in self._on_before_run_session_callbacks:
-            callback(queue_item=queue_item)
-
-    def _on_after_run_session(self, queue_item: SessionQueueItem) -> None:
-        """Called after a session is run.
-
-        - Stop the profiler if profiling is enabled.
-        - Update the queue item's session object in the database.
-        - If not already canceled or failed, complete the queue item.
-        - Log and reset performance statistics.
-        - Run any callbacks registered for this event.
-        """
-
-        self._services.logger.debug(
-            f"On after run session: queue item {queue_item.item_id}, session {queue_item.session_id}"
-        )
-
-        # If we are profiling, stop the profiler and dump the profile & stats
-        if self._profiler is not None:
-            profile_path = self._profiler.stop()
-            stats_path = profile_path.with_suffix(".json")
-            self._services.performance_statistics.dump_stats(
-                graph_execution_state_id=queue_item.session.id, output_path=stats_path
-            )
-
-        try:
-            # Update the queue item with the completed session. If the queue item has been removed from the queue,
-            # we'll get a SessionQueueItemNotFoundError and we can ignore it. This can happen if the queue is cleared
-            # while the session is running.
-            queue_item = self._services.session_queue.set_queue_item_session(queue_item.item_id, queue_item.session)
-
-            # The queue item may have been canceled or failed while the session was running. We should only complete it
-            # if it is not already canceled or failed.
-            if queue_item.status not in ["canceled", "failed"]:
-                queue_item = self._services.session_queue.complete_queue_item(queue_item.item_id)
-
-            # 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._services.performance_statistics.log_stats(queue_item.session.id)
-                self._services.performance_statistics.reset_stats()
-
-            for callback in self._on_after_run_session_callbacks:
-                callback(queue_item=queue_item)
-        except SessionQueueItemNotFoundError:
-            pass
-
-    def _on_before_run_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem):
-        """Called before a node is run.
-
-        - Emits an invocation started event.
-        - Run any callbacks registered for this event.
-        """
-
-        self._services.logger.debug(
-            f"On before run node: queue item {queue_item.item_id}, session {queue_item.session_id}, node {invocation.id} ({invocation.get_type()})"
-        )
-
-        # Send starting event
-        self._services.events.emit_invocation_started(queue_item=queue_item, invocation=invocation)
-
-        for callback in self._on_before_run_node_callbacks:
-            callback(invocation=invocation, queue_item=queue_item)
-
-    def _on_after_run_node(
-        self, invocation: BaseInvocation, queue_item: SessionQueueItem, output: BaseInvocationOutput
-    ):
-        """Called after a node is run.
-
-        - Emits an invocation complete event.
-        - Run any callbacks registered for this event.
-        """
-
-        self._services.logger.debug(
-            f"On after run node: queue item {queue_item.item_id}, session {queue_item.session_id}, node {invocation.id} ({invocation.get_type()})"
-        )
-
-        # Send complete event on successful runs
-        self._services.events.emit_invocation_complete(invocation=invocation, queue_item=queue_item, output=output)
-
-        for callback in self._on_after_run_node_callbacks:
-            callback(invocation=invocation, queue_item=queue_item, output=output)
-
-    def _on_node_error(
-        self,
-        invocation: BaseInvocation,
-        queue_item: SessionQueueItem,
-        error_type: str,
-        error_message: str,
-        error_traceback: str,
-    ):
-        """Called when a node errors. Node errors may occur when running or preparing the node..
-
-        - Set the node error on the session object.
-        - Log the error.
-        - Fail the queue item.
-        - Emits an invocation error event.
-        - Run any callbacks registered for this event.
-        """
-
-        self._services.logger.debug(
-            f"On node error: queue item {queue_item.item_id}, session {queue_item.session_id}, node {invocation.id} ({invocation.get_type()})"
-        )
-
-        # Node errors do not get the full traceback. Only the queue item gets the full traceback.
-        node_error = f"{error_type}: {error_message}"
-        queue_item.session.set_node_error(invocation.id, node_error)
-        self._services.logger.error(
-            f"Error while invoking session {queue_item.session_id}, invocation {invocation.id} ({invocation.get_type()}): {error_message}"
-        )
-        self._services.logger.error(error_traceback)
-
-        # Fail the queue item
-        queue_item = self._services.session_queue.set_queue_item_session(queue_item.item_id, queue_item.session)
-        queue_item = self._services.session_queue.fail_queue_item(
-            queue_item.item_id, error_type, error_message, error_traceback
-        )
-
-        # Send error event
-        self._services.events.emit_invocation_error(
-            queue_item=queue_item,
-            invocation=invocation,
-            error_type=error_type,
-            error_message=error_message,
-            error_traceback=error_traceback,
-        )
-
-        for callback in self._on_node_error_callbacks:
-            callback(
-                invocation=invocation,
-                queue_item=queue_item,
-                error_type=error_type,
-                error_message=error_message,
-                error_traceback=error_traceback,
-            )
-
-
 class DefaultSessionProcessor(SessionProcessorBase):
-    def __init__(
-        self,
-        session_runner: Optional[SessionRunnerBase] = None,
-        on_non_fatal_processor_error_callbacks: Optional[list[OnNonFatalProcessorError]] = None,
-        thread_limit: int = 1,
-        polling_interval: int = 1,
-    ) -> None:
-        super().__init__()
-
-        self.session_runner = session_runner if session_runner else DefaultSessionRunner()
-        self._on_non_fatal_processor_error_callbacks = on_non_fatal_processor_error_callbacks or []
-        self._thread_limit = thread_limit
-        self._polling_interval = polling_interval
-
-    def start(self, invoker: Invoker) -> 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
@@ -332,11 +31,11 @@ class DefaultSessionProcessor(SessionProcessorBase):
         self._poll_now_event = ThreadEvent()
         self._cancel_event = ThreadEvent()
 
-        register_events(QueueClearedEvent, self._on_queue_cleared)
-        register_events(BatchEnqueuedEvent, self._on_batch_enqueued)
-        register_events(QueueItemStatusChangedEvent, self._on_queue_item_status_changed)
+        local_handler.register(event_name=EventServiceBase.queue_event, _func=self._on_queue_event)
 
-        self._thread_semaphore = BoundedSemaphore(self._thread_limit)
+        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.
@@ -350,7 +49,6 @@ class DefaultSessionProcessor(SessionProcessorBase):
             else None
         )
 
-        self.session_runner.start(services=invoker.services, cancel_event=self._cancel_event, profiler=self._profiler)
         self._thread = Thread(
             name="session_processor",
             target=self._process,
@@ -369,25 +67,30 @@ class DefaultSessionProcessor(SessionProcessorBase):
     def _poll_now(self) -> None:
         self._poll_now_event.set()
 
-    async def _on_queue_cleared(self, event: FastAPIEvent[QueueClearedEvent]) -> None:
-        if self._queue_item and self._queue_item.queue_id == event[1].queue_id:
+    async def _on_queue_event(self, event: FastAPIEvent) -> None:
+        event_name = event[1]["event"]
+
+        if (
+            event_name == "session_canceled"
+            and self._queue_item
+            and self._queue_item.item_id == event[1]["data"]["queue_item_id"]
+        ):
             self._cancel_event.set()
             self._poll_now()
-
-    async def _on_batch_enqueued(self, event: FastAPIEvent[BatchEnqueuedEvent]) -> None:
-        self._poll_now()
-
-    async def _on_queue_item_status_changed(self, event: FastAPIEvent[QueueItemStatusChangedEvent]) -> None:
-        if self._queue_item and event[1].status in ["completed", "failed", "canceled"]:
-            # When the queue item is canceled via HTTP, the queue item status is set to `"canceled"` and this event is
-            # emitted. We need to respond to this event and stop graph execution. This is done by setting the cancel
-            # event, which the session runner checks between invocations. If set, the session runner loop is broken.
-            #
-            # Long-running nodes that cannot be interrupted easily present a challenge. `denoise_latents` is one such
-            # node, but it gets a step callback, called on each step of denoising. This callback checks if the queue item
-            # is canceled, and if it is, raises a `CanceledException` to stop execution immediately.
-            if event[1].status == "canceled":
-                self._cancel_event.set()
+        elif (
+            event_name == "queue_cleared"
+            and self._queue_item
+            and self._queue_item.queue_id == event[1]["data"]["queue_id"]
+        ):
+            self._cancel_event.set()
+            self._poll_now()
+        elif event_name == "batch_enqueued":
+            self._poll_now()
+        elif event_name == "queue_item_status_changed" and event[1]["data"]["queue_item"]["status"] in [
+            "completed",
+            "failed",
+            "canceled",
+        ]:
             self._poll_now()
 
     def resume(self) -> SessionProcessorStatus:
@@ -413,8 +116,8 @@ class DefaultSessionProcessor(SessionProcessorBase):
         resume_event: ThreadEvent,
         cancel_event: ThreadEvent,
     ):
+        # Outermost processor try block; any unhandled exception is a fatal processor error
         try:
-            # Any unhandled exception in this block is a fatal processor error and will stop the processor.
             self._thread_semaphore.acquire()
             stop_event.clear()
             resume_event.set()
@@ -422,8 +125,8 @@ class DefaultSessionProcessor(SessionProcessorBase):
 
             while not stop_event.is_set():
                 poll_now_event.clear()
+                # Middle processor try block; any unhandled exception is a non-fatal processor error
                 try:
-                    # Any unhandled exception in this block is a nonfatal processor error and will be handled.
                     # If we are paused, wait for resume event
                     resume_event.wait()
 
@@ -439,69 +142,157 @@ class DefaultSessionProcessor(SessionProcessorBase):
                     self._invoker.services.logger.debug(f"Executing queue item {self._queue_item.item_id}")
                     cancel_event.clear()
 
-                    # Run the graph
-                    self.session_runner.run(queue_item=self._queue_item)
+                    # If profiling is enabled, start the profiler
+                    if self._profiler is not None:
+                        self._profiler.start(profile_id=self._queue_item.session_id)
 
-                except Exception as e:
-                    error_type = e.__class__.__name__
-                    error_message = str(e)
-                    error_traceback = traceback.format_exc()
-                    self._on_non_fatal_processor_error(
-                        queue_item=self._queue_item,
-                        error_type=error_type,
-                        error_message=error_message,
-                        error_traceback=error_traceback,
+                    # 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,
+                        )
+
+                        # 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}"
+                            )
+                            self._invoker.services.logger.error(error)
+
+                            # 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()
+                    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:
+                    # Non-fatal error in processor
+                    self._invoker.services.logger.error(
+                        f"Non-fatal error in session processor:\n{traceback.format_exc()}"
                     )
-                    # Wait for next polling interval or event to try again
+                    # 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()
+                        )
+                    # 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:
+        except Exception:
             # Fatal error in processor, log and pass - we're done here
-            error_type = e.__class__.__name__
-            error_message = str(e)
-            error_traceback = traceback.format_exc()
-            self._invoker.services.logger.error(f"Fatal Error in session processor {error_type}: {error_message}")
-            self._invoker.services.logger.error(error_traceback)
+            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._thread_semaphore.release()
-
-    def _on_non_fatal_processor_error(
-        self,
-        queue_item: Optional[SessionQueueItem],
-        error_type: str,
-        error_message: str,
-        error_traceback: str,
-    ) -> None:
-        """Called when a non-fatal error occurs in the processor.
-
-        - Log the error.
-        - If a queue item is provided, update the queue item with the completed session & fail it.
-        - Run any callbacks registered for this event.
-        """
-
-        self._invoker.services.logger.error(f"Non-fatal error in session processor {error_type}: {error_message}")
-        self._invoker.services.logger.error(error_traceback)
-
-        if queue_item is not None:
-            # Update the queue item with the completed session & fail it
-            queue_item = self._invoker.services.session_queue.set_queue_item_session(
-                queue_item.item_id, queue_item.session
-            )
-            queue_item = self._invoker.services.session_queue.fail_queue_item(
-                item_id=queue_item.item_id,
-                error_type=error_type,
-                error_message=error_message,
-                error_traceback=error_traceback,
-            )
-
-        for callback in self._on_non_fatal_processor_error_callbacks:
-            callback(
-                queue_item=queue_item,
-                error_type=error_type,
-                error_message=error_message,
-                error_traceback=error_traceback,
-            )

invokeai/app/services/session_queue/session_queue_base.py

@@ -16,7 +16,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     SessionQueueItemDTO,
     SessionQueueStatus,
 )
-from invokeai.app.services.shared.graph import GraphExecutionState
 from invokeai.app.services.shared.pagination import CursorPaginatedResults
 
 
@@ -74,22 +73,10 @@ class SessionQueueBase(ABC):
         pass
 
     @abstractmethod
-    def complete_queue_item(self, item_id: int) -> SessionQueueItem:
-        """Completes a session queue item"""
-        pass
-
-    @abstractmethod
-    def cancel_queue_item(self, item_id: int) -> SessionQueueItem:
+    def cancel_queue_item(self, item_id: int, error: Optional[str] = None) -> SessionQueueItem:
         """Cancels a session queue item"""
         pass
 
-    @abstractmethod
-    def fail_queue_item(
-        self, item_id: int, error_type: str, error_message: str, error_traceback: str
-    ) -> SessionQueueItem:
-        """Fails a session queue item"""
-        pass
-
     @abstractmethod
     def cancel_by_batch_ids(self, queue_id: str, batch_ids: list[str]) -> CancelByBatchIDsResult:
         """Cancels all queue items with matching batch IDs"""
@@ -116,8 +103,3 @@ class SessionQueueBase(ABC):
     def get_queue_item(self, item_id: int) -> SessionQueueItem:
         """Gets a session queue item by ID"""
         pass
-
-    @abstractmethod
-    def set_queue_item_session(self, item_id: int, session: GraphExecutionState) -> SessionQueueItem:
-        """Sets the session for a session queue item. Use this to update the session state."""
-        pass

invokeai/app/services/session_queue/session_queue_common.py

@@ -3,16 +3,7 @@ import json
 from itertools import chain, product
 from typing import Generator, Iterable, Literal, NamedTuple, Optional, TypeAlias, Union, cast
 
-from pydantic import (
-    AliasChoices,
-    BaseModel,
-    ConfigDict,
-    Field,
-    StrictStr,
-    TypeAdapter,
-    field_validator,
-    model_validator,
-)
+from pydantic import BaseModel, ConfigDict, Field, StrictStr, TypeAdapter, field_validator, model_validator
 from pydantic_core import to_jsonable_python
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
@@ -198,13 +189,7 @@ class SessionQueueItemWithoutGraph(BaseModel):
     session_id: str = Field(
         description="The ID of the session associated with this queue item. The session doesn't exist in graph_executions until the queue item is executed."
     )
-    error_type: Optional[str] = Field(default=None, description="The error type if this queue item errored")
-    error_message: Optional[str] = Field(default=None, description="The error message if this queue item errored")
-    error_traceback: Optional[str] = Field(
-        default=None,
-        description="The error traceback if this queue item errored",
-        validation_alias=AliasChoices("error_traceback", "error"),
-    )
+    error: Optional[str] = Field(default=None, description="The error message if this queue item errored")
     created_at: Union[datetime.datetime, str] = Field(description="When this queue item was created")
     updated_at: Union[datetime.datetime, str] = Field(description="When this queue item was updated")
     started_at: Optional[Union[datetime.datetime, str]] = Field(description="When this queue item was started")

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -2,6 +2,10 @@ import sqlite3
 import threading
 from typing import Optional, Union, cast
 
+from fastapi_events.handlers.local import local_handler
+from fastapi_events.typing import Event as FastAPIEvent
+
+from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase
 from invokeai.app.services.session_queue.session_queue_common import (
@@ -23,7 +27,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     calc_session_count,
     prepare_values_to_insert,
 )
-from invokeai.app.services.shared.graph import GraphExecutionState
 from invokeai.app.services.shared.pagination import CursorPaginatedResults
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
 
@@ -37,14 +40,10 @@ class SqliteSessionQueue(SessionQueueBase):
     def start(self, invoker: Invoker) -> None:
         self.__invoker = invoker
         self._set_in_progress_to_canceled()
-        if self.__invoker.services.configuration.clear_queue_on_startup:
-            clear_result = self.clear(DEFAULT_QUEUE_ID)
-            if clear_result.deleted > 0:
-                self.__invoker.services.logger.info(f"Cleared all {clear_result.deleted} queue items")
-        else:
-            prune_result = self.prune(DEFAULT_QUEUE_ID)
-            if prune_result.deleted > 0:
-                self.__invoker.services.logger.info(f"Pruned {prune_result.deleted} finished queue items")
+        prune_result = self.prune(DEFAULT_QUEUE_ID)
+        local_handler.register(event_name=EventServiceBase.queue_event, _func=self._on_session_event)
+        if prune_result.deleted > 0:
+            self.__invoker.services.logger.info(f"Pruned {prune_result.deleted} finished queue items")
 
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
@@ -52,6 +51,52 @@ class SqliteSessionQueue(SessionQueueBase):
         self.__conn = db.conn
         self.__cursor = self.__conn.cursor()
 
+    def _match_event_name(self, event: FastAPIEvent, match_in: list[str]) -> bool:
+        return event[1]["event"] in match_in
+
+    async def _on_session_event(self, event: FastAPIEvent) -> FastAPIEvent:
+        event_name = event[1]["event"]
+
+        # This was a match statement, but match is not supported on python 3.9
+        if event_name == "graph_execution_state_complete":
+            await self._handle_complete_event(event)
+        elif event_name == "invocation_error":
+            await self._handle_error_event(event)
+        elif event_name == "session_canceled":
+            await self._handle_cancel_event(event)
+        return event
+
+    async def _handle_complete_event(self, event: FastAPIEvent) -> None:
+        try:
+            item_id = event[1]["data"]["queue_item_id"]
+            # When a queue item has an error, we get an error event, then a completed event.
+            # Mark the queue item completed only if it isn't already marked completed, e.g.
+            # by a previously-handled error event.
+            queue_item = self.get_queue_item(item_id)
+            if queue_item.status not in ["completed", "failed", "canceled"]:
+                queue_item = self._set_queue_item_status(item_id=queue_item.item_id, status="completed")
+        except SessionQueueItemNotFoundError:
+            return
+
+    async def _handle_error_event(self, event: FastAPIEvent) -> None:
+        try:
+            item_id = event[1]["data"]["queue_item_id"]
+            error = event[1]["data"]["error"]
+            queue_item = self.get_queue_item(item_id)
+            # always set to failed if have an error, even if previously the item was marked completed or canceled
+            queue_item = self._set_queue_item_status(item_id=queue_item.item_id, status="failed", error=error)
+        except SessionQueueItemNotFoundError:
+            return
+
+    async def _handle_cancel_event(self, event: FastAPIEvent) -> None:
+        try:
+            item_id = event[1]["data"]["queue_item_id"]
+            queue_item = self.get_queue_item(item_id)
+            if queue_item.status not in ["completed", "failed", "canceled"]:
+                queue_item = self._set_queue_item_status(item_id=queue_item.item_id, status="canceled")
+        except SessionQueueItemNotFoundError:
+            return
+
     def _set_in_progress_to_canceled(self) -> None:
         """
         Sets all in_progress queue items to canceled. Run on app startup, not associated with any queue.
@@ -226,22 +271,17 @@ class SqliteSessionQueue(SessionQueueBase):
         return SessionQueueItem.queue_item_from_dict(dict(result))
 
     def _set_queue_item_status(
-        self,
-        item_id: int,
-        status: QUEUE_ITEM_STATUS,
-        error_type: Optional[str] = None,
-        error_message: Optional[str] = None,
-        error_traceback: Optional[str] = None,
+        self, item_id: int, status: QUEUE_ITEM_STATUS, error: Optional[str] = None
     ) -> SessionQueueItem:
         try:
             self.__lock.acquire()
             self.__cursor.execute(
                 """--sql
                 UPDATE session_queue
-                SET status = ?, error_type = ?, error_message = ?, error_traceback = ?
+                SET status = ?, error = ?
                 WHERE item_id = ?
                 """,
-                (status, error_type, error_message, error_traceback, item_id),
+                (status, error, item_id),
             )
             self.__conn.commit()
         except Exception:
@@ -252,7 +292,11 @@ class SqliteSessionQueue(SessionQueueBase):
         queue_item = self.get_queue_item(item_id)
         batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id)
         queue_status = self.get_queue_status(queue_id=queue_item.queue_id)
-        self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status)
+        self.__invoker.services.events.emit_queue_item_status_changed(
+            session_queue_item=queue_item,
+            batch_status=batch_status,
+            queue_status=queue_status,
+        )
         return queue_item
 
     def is_empty(self, queue_id: str) -> IsEmptyResult:
@@ -294,6 +338,26 @@ class SqliteSessionQueue(SessionQueueBase):
             self.__lock.release()
         return IsFullResult(is_full=is_full)
 
+    def delete_queue_item(self, item_id: int) -> SessionQueueItem:
+        queue_item = self.get_queue_item(item_id=item_id)
+        try:
+            self.__lock.acquire()
+            self.__cursor.execute(
+                """--sql
+                DELETE FROM session_queue
+                WHERE
+                  item_id = ?
+                """,
+                (item_id,),
+            )
+            self.__conn.commit()
+        except Exception:
+            self.__conn.rollback()
+            raise
+        finally:
+            self.__lock.release()
+        return queue_item
+
     def clear(self, queue_id: str) -> ClearResult:
         try:
             self.__lock.acquire()
@@ -360,28 +424,17 @@ class SqliteSessionQueue(SessionQueueBase):
             self.__lock.release()
         return PruneResult(deleted=count)
 
-    def cancel_queue_item(self, item_id: int) -> SessionQueueItem:
-        queue_item = self._set_queue_item_status(item_id=item_id, status="canceled")
-        return queue_item
-
-    def complete_queue_item(self, item_id: int) -> SessionQueueItem:
-        queue_item = self._set_queue_item_status(item_id=item_id, status="completed")
-        return queue_item
-
-    def fail_queue_item(
-        self,
-        item_id: int,
-        error_type: str,
-        error_message: str,
-        error_traceback: str,
-    ) -> SessionQueueItem:
-        queue_item = self._set_queue_item_status(
-            item_id=item_id,
-            status="failed",
-            error_type=error_type,
-            error_message=error_message,
-            error_traceback=error_traceback,
-        )
+    def cancel_queue_item(self, item_id: int, error: Optional[str] = None) -> SessionQueueItem:
+        queue_item = self.get_queue_item(item_id)
+        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.events.emit_session_canceled(
+                queue_item_id=queue_item.item_id,
+                queue_id=queue_item.queue_id,
+                queue_batch_id=queue_item.batch_id,
+                graph_execution_state_id=queue_item.session_id,
+            )
         return queue_item
 
     def cancel_by_batch_ids(self, queue_id: str, batch_ids: list[str]) -> CancelByBatchIDsResult:
@@ -417,10 +470,18 @@ 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.events.emit_session_canceled(
+                    queue_item_id=current_queue_item.item_id,
+                    queue_id=current_queue_item.queue_id,
+                    queue_batch_id=current_queue_item.batch_id,
+                    graph_execution_state_id=current_queue_item.session_id,
+                )
                 batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id)
                 queue_status = self.get_queue_status(queue_id=queue_id)
                 self.__invoker.services.events.emit_queue_item_status_changed(
-                    current_queue_item, batch_status, queue_status
+                    session_queue_item=current_queue_item,
+                    batch_status=batch_status,
+                    queue_status=queue_status,
                 )
         except Exception:
             self.__conn.rollback()
@@ -460,10 +521,18 @@ class SqliteSessionQueue(SessionQueueBase):
             )
             self.__conn.commit()
             if current_queue_item is not None and current_queue_item.queue_id == queue_id:
+                self.__invoker.services.events.emit_session_canceled(
+                    queue_item_id=current_queue_item.item_id,
+                    queue_id=current_queue_item.queue_id,
+                    queue_batch_id=current_queue_item.batch_id,
+                    graph_execution_state_id=current_queue_item.session_id,
+                )
                 batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id)
                 queue_status = self.get_queue_status(queue_id=queue_id)
                 self.__invoker.services.events.emit_queue_item_status_changed(
-                    current_queue_item, batch_status, queue_status
+                    session_queue_item=current_queue_item,
+                    batch_status=batch_status,
+                    queue_status=queue_status,
                 )
         except Exception:
             self.__conn.rollback()
@@ -493,29 +562,6 @@ class SqliteSessionQueue(SessionQueueBase):
             raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}")
         return SessionQueueItem.queue_item_from_dict(dict(result))
 
-    def set_queue_item_session(self, item_id: int, session: GraphExecutionState) -> SessionQueueItem:
-        try:
-            # Use exclude_none so we don't end up with a bunch of nulls in the graph - this can cause validation errors
-            # when the graph is loaded. Graph execution occurs purely in memory - the session saved here is not referenced
-            # during execution.
-            session_json = session.model_dump_json(warnings=False, exclude_none=True)
-            self.__lock.acquire()
-            self.__cursor.execute(
-                """--sql
-                UPDATE session_queue
-                SET session = ?
-                WHERE item_id = ?
-                """,
-                (session_json, item_id),
-            )
-            self.__conn.commit()
-        except Exception:
-            self.__conn.rollback()
-            raise
-        finally:
-            self.__lock.release()
-        return self.get_queue_item(item_id)
-
     def list_queue_items(
         self,
         queue_id: str,
@@ -532,9 +578,7 @@ class SqliteSessionQueue(SessionQueueBase):
                     status,
                     priority,
                     field_values,
-                    error_type,
-                    error_message,
-                    error_traceback,
+                    error,
                     created_at,
                     updated_at,
                     completed_at,

invokeai/app/services/shared/graph.py

@@ -2,19 +2,17 @@
 
 import copy
 import itertools
-from typing import Any, Optional, TypeVar, Union, get_args, get_origin, get_type_hints
+from typing import Annotated, Any, Optional, TypeVar, Union, get_args, get_origin, get_type_hints
 
 import networkx as nx
 from pydantic import (
     BaseModel,
-    GetCoreSchemaHandler,
     GetJsonSchemaHandler,
-    ValidationError,
     field_validator,
 )
 from pydantic.fields import Field
 from pydantic.json_schema import JsonSchemaValue
-from pydantic_core import core_schema
+from pydantic_core import CoreSchema
 
 # Importing * is bad karma but needed here for node detection
 from invokeai.app.invocations import *  # noqa: F401 F403
@@ -192,39 +190,6 @@ class UnknownGraphValidationError(ValueError):
     pass
 
 
-class NodeInputError(ValueError):
-    """Raised when a node fails preparation. This occurs when a node's inputs are being set from its incomers, but an
-    input fails validation.
-
-    Attributes:
-        node: The node that failed preparation. Note: only successfully set fields will be accurate. Review the error to
-            determine which field caused the failure.
-    """
-
-    def __init__(self, node: BaseInvocation, e: ValidationError):
-        self.original_error = e
-        self.node = node
-        # When preparing a node, we set each input one-at-a-time. We may thus safely assume that the first error
-        # represents the first input that failed.
-        self.failed_input = loc_to_dot_sep(e.errors()[0]["loc"])
-        super().__init__(f"Node {node.id} has invalid incoming input for {self.failed_input}")
-
-
-def loc_to_dot_sep(loc: tuple[Union[str, int], ...]) -> str:
-    """Helper to pretty-print pydantic error locations as dot-separated strings.
-    Taken from https://docs.pydantic.dev/latest/errors/errors/#customize-error-messages
-    """
-    path = ""
-    for i, x in enumerate(loc):
-        if isinstance(x, str):
-            if i > 0:
-                path += "."
-            path += x
-        else:
-            path += f"[{x}]"
-    return path
-
-
 @invocation_output("iterate_output")
 class IterateInvocationOutput(BaseInvocationOutput):
     """Used to connect iteration outputs. Will be expanded to a specific output."""
@@ -278,58 +243,73 @@ class CollectInvocation(BaseInvocation):
         return CollectInvocationOutput(collection=copy.copy(self.collection))
 
 
-class AnyInvocation(BaseInvocation):
-    @classmethod
-    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
-        def validate_invocation(v: Any) -> "AnyInvocation":
-            return BaseInvocation.get_typeadapter().validate_python(v)
-
-        return core_schema.no_info_plain_validator_function(validate_invocation)
-
-    @classmethod
-    def __get_pydantic_json_schema__(
-        cls, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
-    ) -> JsonSchemaValue:
-        # Nodes are too powerful, we have to make our own OpenAPI schema manually
-        # No but really, because the schema is dynamic depending on loaded nodes, we need to generate it manually
-        oneOf: list[dict[str, str]] = []
-        names = [i.__name__ for i in BaseInvocation.get_invocations()]
-        for name in sorted(names):
-            oneOf.append({"$ref": f"#/components/schemas/{name}"})
-        return {"oneOf": oneOf}
-
-
-class AnyInvocationOutput(BaseInvocationOutput):
-    @classmethod
-    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler):
-        def validate_invocation_output(v: Any) -> "AnyInvocationOutput":
-            return BaseInvocationOutput.get_typeadapter().validate_python(v)
-
-        return core_schema.no_info_plain_validator_function(validate_invocation_output)
-
-    @classmethod
-    def __get_pydantic_json_schema__(
-        cls, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
-    ) -> JsonSchemaValue:
-        # Nodes are too powerful, we have to make our own OpenAPI schema manually
-        # No but really, because the schema is dynamic depending on loaded nodes, we need to generate it manually
-
-        oneOf: list[dict[str, str]] = []
-        names = [i.__name__ for i in BaseInvocationOutput.get_outputs()]
-        for name in sorted(names):
-            oneOf.append({"$ref": f"#/components/schemas/{name}"})
-        return {"oneOf": oneOf}
-
-
 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, AnyInvocation] = 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
 
@@ -652,7 +632,7 @@ class Graph(BaseModel):
         output_fields = [get_input_field(self.get_node(e.node_id), e.field) for e in outputs]
 
         # Input type must be a list
-        if get_origin(input_field) is not list:
+        if get_origin(input_field) != list:
             return False
 
         # Validate that all outputs match the input type
@@ -760,7 +740,7 @@ class GraphExecutionState(BaseModel):
     )
 
     # The results of executed nodes
-    results: dict[str, AnyInvocationOutput] = 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)
@@ -777,12 +757,52 @@ class GraphExecutionState(BaseModel):
         default_factory=dict,
     )
 
+    @field_validator("results", mode="plain")
+    @classmethod
+    def validate_results(cls, v: dict[str, BaseInvocationOutput]):
+        """Validates the results in the GES by retrieving a union of all output types and validating each result."""
+
+        # See the comment in `Graph.validate_nodes` for an explanation of this logic.
+        results: dict[str, BaseInvocationOutput] = {}
+        typeadapter = BaseInvocationOutput.get_typeadapter()
+        for result_id, result in v.items():
+            results[result_id] = typeadapter.validate_python(result)
+        return results
+
     @field_validator("graph")
     def graph_is_valid(cls, v: Graph):
         """Validates that the graph is valid"""
         v.validate_self()
         return v
 
+    @classmethod
+    def __get_pydantic_json_schema__(cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
+        # See the comment in `Graph.__get_pydantic_json_schema__` for an explanation of this logic.
+        class GraphExecutionState(BaseModel):
+            """Tracks the state of a graph execution"""
+
+            id: str = Field(description="The id of the execution state")
+            graph: Graph = Field(description="The graph being executed")
+            execution_graph: Graph = Field(description="The expanded graph of activated and executed nodes")
+            executed: set[str] = Field(description="The set of node ids that have been executed")
+            executed_history: list[str] = Field(
+                description="The list of node ids that have been executed, in order of execution"
+            )
+            results: dict[
+                str, Annotated[Union[tuple(BaseInvocationOutput._output_classes)], Field(discriminator="type")]
+            ] = Field(description="The results of node executions")
+            errors: dict[str, str] = Field(description="Errors raised when executing nodes")
+            prepared_source_mapping: dict[str, str] = Field(
+                description="The map of prepared nodes to original graph nodes"
+            )
+            source_prepared_mapping: dict[str, set[str]] = Field(
+                description="The map of original graph nodes to prepared nodes"
+            )
+
+        json_schema = handler(GraphExecutionState.__pydantic_core_schema__)
+        json_schema = handler.resolve_ref_schema(json_schema)
+        return json_schema
+
     def next(self) -> Optional[BaseInvocation]:
         """Gets the next node ready to execute."""
 
@@ -801,10 +821,7 @@ class GraphExecutionState(BaseModel):
 
         # Get values from edges
         if next_node is not None:
-            try:
-                self._prepare_inputs(next_node)
-            except ValidationError as e:
-                raise NodeInputError(next_node, e)
+            self._prepare_inputs(next_node)
 
         # If next is still none, there's no next node, return None
         return next_node

invokeai/app/services/shared/invocation_context.py

@@ -1,9 +1,9 @@
+import threading
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from PIL.Image import Image
-from pydantic.networks import AnyHttpUrl
 from torch import Tensor
 
 from invokeai.app.invocations.constants import IMAGE_MODES
@@ -15,15 +15,8 @@ from invokeai.app.services.images.images_common import ImageDTO
 from invokeai.app.services.invocation_services import InvocationServices
 from invokeai.app.services.model_records.model_records_base import UnknownModelException
 from invokeai.app.util.step_callback import stable_diffusion_step_callback
-from invokeai.backend.model_manager.config import (
-    AnyModel,
-    AnyModelConfig,
-    BaseModelType,
-    ModelFormat,
-    ModelType,
-    SubModelType,
-)
-from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig
+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.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
 
@@ -187,9 +180,9 @@ class ImagesInterface(InvocationContextInterface):
         # If `metadata` is provided directly, use that. Else, use the metadata provided by `WithMetadata`, falling back to None.
         metadata_ = None
         if metadata:
-            metadata_ = metadata.model_dump_json()
-        elif isinstance(self._data.invocation, WithMetadata) and self._data.invocation.metadata:
-            metadata_ = self._data.invocation.metadata.model_dump_json()
+            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
@@ -198,14 +191,6 @@ class ImagesInterface(InvocationContextInterface):
         elif isinstance(self._data.invocation, WithBoard) and self._data.invocation.board:
             board_id_ = self._data.invocation.board.board_id
 
-        workflow_ = None
-        if self._data.queue_item.workflow:
-            workflow_ = self._data.queue_item.workflow.model_dump_json()
-
-        graph_ = None
-        if self._data.queue_item.session.graph:
-            graph_ = self._data.queue_item.session.graph.model_dump_json()
-
         return self._services.images.create(
             image=image,
             is_intermediate=self._data.invocation.is_intermediate,
@@ -213,8 +198,7 @@ class ImagesInterface(InvocationContextInterface):
             board_id=board_id_,
             metadata=metadata_,
             image_origin=ResourceOrigin.INTERNAL,
-            workflow=workflow_,
-            graph=graph_,
+            workflow=self._data.queue_item.workflow,
             session_id=self._data.queue_item.session_id,
             node_id=self._data.invocation.id,
         )
@@ -328,10 +312,8 @@ class ConditioningInterface(InvocationContextInterface):
 
 
 class ModelsInterface(InvocationContextInterface):
-    """Common API for loading, downloading and managing models."""
-
     def exists(self, identifier: Union[str, "ModelIdentifierField"]) -> bool:
-        """Check if a model exists.
+        """Checks if a model exists.
 
         Args:
             identifier: The key or ModelField representing the model.
@@ -341,13 +323,13 @@ class ModelsInterface(InvocationContextInterface):
         """
         if isinstance(identifier, str):
             return self._services.model_manager.store.exists(identifier)
-        else:
-            return self._services.model_manager.store.exists(identifier.key)
+
+        return self._services.model_manager.store.exists(identifier.key)
 
     def load(
         self, identifier: Union[str, "ModelIdentifierField"], submodel_type: Optional[SubModelType] = None
     ) -> LoadedModel:
-        """Load a model.
+        """Loads a model.
 
         Args:
             identifier: The key or ModelField representing the model.
@@ -362,16 +344,16 @@ class ModelsInterface(InvocationContextInterface):
 
         if isinstance(identifier, str):
             model = self._services.model_manager.store.get_model(identifier)
-            return self._services.model_manager.load.load_model(model, submodel_type)
+            return self._services.model_manager.load.load_model(model, submodel_type, self._data)
         else:
             _submodel_type = submodel_type or identifier.submodel_type
             model = self._services.model_manager.store.get_model(identifier.key)
-            return self._services.model_manager.load.load_model(model, _submodel_type)
+            return self._services.model_manager.load.load_model(model, _submodel_type, self._data)
 
     def load_by_attrs(
         self, name: str, base: BaseModelType, type: ModelType, submodel_type: Optional[SubModelType] = None
     ) -> LoadedModel:
-        """Load a model by its attributes.
+        """Loads a model by its attributes.
 
         Args:
             name: Name of the model.
@@ -391,10 +373,10 @@ class ModelsInterface(InvocationContextInterface):
         if len(configs) > 1:
             raise ValueError(f"More than one model found with name {name}, base {base}, and type {type}")
 
-        return self._services.model_manager.load.load_model(configs[0], submodel_type)
+        return self._services.model_manager.load.load_model(configs[0], submodel_type, self._data)
 
     def get_config(self, identifier: Union[str, "ModelIdentifierField"]) -> AnyModelConfig:
-        """Get a model's config.
+        """Gets a model's config.
 
         Args:
             identifier: The key or ModelField representing the model.
@@ -404,11 +386,11 @@ class ModelsInterface(InvocationContextInterface):
         """
         if isinstance(identifier, str):
             return self._services.model_manager.store.get_model(identifier)
-        else:
-            return self._services.model_manager.store.get_model(identifier.key)
+
+        return self._services.model_manager.store.get_model(identifier.key)
 
     def search_by_path(self, path: Path) -> list[AnyModelConfig]:
-        """Search for models by path.
+        """Searches for models by path.
 
         Args:
             path: The path to search for.
@@ -425,7 +407,7 @@ class ModelsInterface(InvocationContextInterface):
         type: Optional[ModelType] = None,
         format: Optional[ModelFormat] = None,
     ) -> list[AnyModelConfig]:
-        """Search for models by attributes.
+        """Searches for models by attributes.
 
         Args:
             name: The name to search for (exact match).
@@ -444,72 +426,6 @@ class ModelsInterface(InvocationContextInterface):
             model_format=format,
         )
 
-    def download_and_cache_model(
-        self,
-        source: str | AnyHttpUrl,
-    ) -> Path:
-        """
-        Download the model file located at source to the models cache and return its Path.
-
-        This can be used to single-file install models and other resources of arbitrary types
-        which should not get registered with the database. If the model is already
-        installed, the cached path will be returned. Otherwise it will be downloaded.
-
-        Args:
-            source: A URL that points to the model, or a huggingface repo_id.
-
-        Returns:
-            Path to the downloaded model
-        """
-        return self._services.model_manager.install.download_and_cache_model(source=source)
-
-    def load_local_model(
-        self,
-        model_path: Path,
-        loader: Optional[Callable[[Path], AnyModel]] = None,
-    ) -> LoadedModelWithoutConfig:
-        """
-        Load the model file located at the indicated path
-
-        If a loader callable is provided, it will be invoked to load the model. Otherwise,
-        `safetensors.torch.load_file()` or `torch.load()` will be called to load the model.
-
-        Be aware that the LoadedModelWithoutConfig object has no `config` attribute
-
-        Args:
-            path: A model Path
-            loader: A Callable that expects a Path and returns a dict[str|int, Any]
-
-        Returns:
-            A LoadedModelWithoutConfig object.
-        """
-        return self._services.model_manager.load.load_model_from_path(model_path=model_path, loader=loader)
-
-    def load_remote_model(
-        self,
-        source: str | AnyHttpUrl,
-        loader: Optional[Callable[[Path], AnyModel]] = None,
-    ) -> LoadedModelWithoutConfig:
-        """
-        Download, cache, and load the model file located at the indicated URL or repo_id.
-
-        If the model is already downloaded, it will be loaded from the cache.
-
-        If the a loader callable is provided, it will be invoked to load the model. Otherwise,
-        `safetensors.torch.load_file()` or `torch.load()` will be called to load the model.
-
-        Be aware that the LoadedModelWithoutConfig object has no `config` attribute
-
-        Args:
-            source: A URL or huggingface repoid.
-            loader: A Callable that expects a Path and returns a dict[str|int, Any]
-
-        Returns:
-            A LoadedModelWithoutConfig object.
-        """
-        model_path = self._services.model_manager.install.download_and_cache_model(source=str(source))
-        return self._services.model_manager.load.load_model_from_path(model_path=model_path, loader=loader)
-
 
 class ConfigInterface(InvocationContextInterface):
     def get(self) -> InvokeAIAppConfig:
@@ -524,10 +440,10 @@ class ConfigInterface(InvocationContextInterface):
 
 class UtilInterface(InvocationContextInterface):
     def __init__(
-        self, services: InvocationServices, data: InvocationContextData, is_canceled: Callable[[], bool]
+        self, services: InvocationServices, data: InvocationContextData, cancel_event: threading.Event
     ) -> None:
         super().__init__(services, data)
-        self._is_canceled = is_canceled
+        self._cancel_event = cancel_event
 
     def is_canceled(self) -> bool:
         """Checks if the current session has been canceled.
@@ -535,7 +451,7 @@ class UtilInterface(InvocationContextInterface):
         Returns:
             True if the current session has been canceled, False if not.
         """
-        return self._is_canceled()
+        return self._cancel_event.is_set()
 
     def sd_step_callback(self, intermediate_state: PipelineIntermediateState, base_model: BaseModelType) -> None:
         """
@@ -610,7 +526,7 @@ class InvocationContext:
 def build_invocation_context(
     services: InvocationServices,
     data: InvocationContextData,
-    is_canceled: Callable[[], bool],
+    cancel_event: threading.Event,
 ) -> InvocationContext:
     """Builds the invocation context for a specific invocation execution.
 
@@ -627,7 +543,7 @@ def build_invocation_context(
     tensors = TensorsInterface(services=services, data=data)
     models = ModelsInterface(services=services, data=data)
     config = ConfigInterface(services=services, data=data)
-    util = UtilInterface(services=services, data=data, is_canceled=is_canceled)
+    util = UtilInterface(services=services, data=data, cancel_event=cancel_event)
     conditioning = ConditioningInterface(services=services, data=data)
     boards = BoardsInterface(services=services, data=data)
 

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

@@ -12,10 +12,6 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_6 import
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_7 import build_migration_7
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_8 import build_migration_8
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_9 import build_migration_9
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_10 import build_migration_10
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_11 import build_migration_11
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_12 import build_migration_12
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_13 import build_migration_13
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -45,10 +41,6 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_7())
     migrator.register_migration(build_migration_8(app_config=config))
     migrator.register_migration(build_migration_9())
-    migrator.register_migration(build_migration_10())
-    migrator.register_migration(build_migration_11(app_config=config, logger=logger))
-    migrator.register_migration(build_migration_12(app_config=config))
-    migrator.register_migration(build_migration_13())
     migrator.run_migrations()
 
     return db

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

@@ -1,35 +0,0 @@
-import sqlite3
-
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-
-class Migration10Callback:
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._update_error_cols(cursor)
-
-    def _update_error_cols(self, cursor: sqlite3.Cursor) -> None:
-        """
-        - Adds `error_type` and `error_message` columns to the session queue table.
-        - Renames the `error` column to `error_traceback`.
-        """
-
-        cursor.execute("ALTER TABLE session_queue ADD COLUMN error_type TEXT;")
-        cursor.execute("ALTER TABLE session_queue ADD COLUMN error_message TEXT;")
-        cursor.execute("ALTER TABLE session_queue RENAME COLUMN error TO error_traceback;")
-
-
-def build_migration_10() -> Migration:
-    """
-    Build the migration from database version 9 to 10.
-
-    This migration does the following:
-    - Adds `error_type` and `error_message` columns to the session queue table.
-    - Renames the `error` column to `error_traceback`.
-    """
-    migration_10 = Migration(
-        from_version=9,
-        to_version=10,
-        callback=Migration10Callback(),
-    )
-
-    return migration_10

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

@@ -1,75 +0,0 @@
-import shutil
-import sqlite3
-from logging import Logger
-
-from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-LEGACY_CORE_MODELS = [
-    # OpenPose
-    "any/annotators/dwpose/yolox_l.onnx",
-    "any/annotators/dwpose/dw-ll_ucoco_384.onnx",
-    # DepthAnything
-    "any/annotators/depth_anything/depth_anything_vitl14.pth",
-    "any/annotators/depth_anything/depth_anything_vitb14.pth",
-    "any/annotators/depth_anything/depth_anything_vits14.pth",
-    # Lama inpaint
-    "core/misc/lama/lama.pt",
-    # RealESRGAN upscale
-    "core/upscaling/realesrgan/RealESRGAN_x4plus.pth",
-    "core/upscaling/realesrgan/RealESRGAN_x4plus_anime_6B.pth",
-    "core/upscaling/realesrgan/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth",
-    "core/upscaling/realesrgan/RealESRGAN_x2plus.pth",
-]
-
-
-class Migration11Callback:
-    def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
-        self._app_config = app_config
-        self._logger = logger
-
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._remove_convert_cache()
-        self._remove_downloaded_models()
-        self._remove_unused_core_models()
-
-    def _remove_convert_cache(self) -> None:
-        """Rename models/.cache to models/.convert_cache."""
-        self._logger.info("Removing .cache directory. Converted models will now be cached in .convert_cache.")
-        legacy_convert_path = self._app_config.root_path / "models" / ".cache"
-        shutil.rmtree(legacy_convert_path, ignore_errors=True)
-
-    def _remove_downloaded_models(self) -> None:
-        """Remove models from their old locations; they will re-download when needed."""
-        self._logger.info(
-            "Removing legacy just-in-time models. Downloaded models will now be cached in .download_cache."
-        )
-        for model_path in LEGACY_CORE_MODELS:
-            legacy_dest_path = self._app_config.models_path / model_path
-            legacy_dest_path.unlink(missing_ok=True)
-
-    def _remove_unused_core_models(self) -> None:
-        """Remove unused core models and their directories."""
-        self._logger.info("Removing defunct core models.")
-        for dir in ["face_restoration", "misc", "upscaling"]:
-            path_to_remove = self._app_config.models_path / "core" / dir
-            shutil.rmtree(path_to_remove, ignore_errors=True)
-        shutil.rmtree(self._app_config.models_path / "any" / "annotators", ignore_errors=True)
-
-
-def build_migration_11(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
-    """
-    Build the migration from database version 10 to 11.
-
-    This migration does the following:
-    - Moves "core" models previously downloaded with download_with_progress_bar() into new
-      "models/.download_cache" directory.
-    - Renames "models/.cache" to "models/.convert_cache".
-    """
-    migration_11 = Migration(
-        from_version=10,
-        to_version=11,
-        callback=Migration11Callback(app_config=app_config, logger=logger),
-    )
-
-    return migration_11

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

@@ -1,35 +0,0 @@
-import shutil
-import sqlite3
-
-from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-
-class Migration12Callback:
-    def __init__(self, app_config: InvokeAIAppConfig) -> None:
-        self._app_config = app_config
-
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._remove_model_convert_cache_dir()
-
-    def _remove_model_convert_cache_dir(self) -> None:
-        """
-        Removes unused model convert cache directory
-        """
-        convert_cache = self._app_config.convert_cache_path
-        shutil.rmtree(convert_cache, ignore_errors=True)
-
-
-def build_migration_12(app_config: InvokeAIAppConfig) -> Migration:
-    """
-    Build the migration from database version 11 to 12.
-
-    This migration removes the now-unused model convert cache directory.
-    """
-    migration_12 = Migration(
-        from_version=11,
-        to_version=12,
-        callback=Migration12Callback(app_config),
-    )
-
-    return migration_12

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

@@ -1,31 +0,0 @@
-import sqlite3
-
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-
-class Migration13Callback:
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._add_archived_col(cursor)
-
-    def _add_archived_col(self, cursor: sqlite3.Cursor) -> None:
-        """
-        - Adds `archived` columns to the board table.
-        """
-
-        cursor.execute("ALTER TABLE boards ADD COLUMN archived BOOLEAN DEFAULT FALSE;")
-
-
-def build_migration_13() -> Migration:
-    """
-    Build the migration from database version 12 to 13..
-
-    This migration does the following:
-    - Adds `archived` columns to the board table.
-    """
-    migration_13 = Migration(
-        from_version=12,
-        to_version=13,
-        callback=Migration13Callback(),
-    )
-
-    return migration_13

invokeai/app/util/controlnet_utils.py

@@ -1,21 +1,13 @@
-from typing import Any, Literal, Union
+from typing import Union
 
 import cv2
 import numpy as np
 import torch
+from controlnet_aux.util import HWC3
+from diffusers.utils import PIL_INTERPOLATION
 from einops import rearrange
 from PIL import Image
 
-from invokeai.backend.image_util.util import nms, normalize_image_channel_count
-
-CONTROLNET_RESIZE_VALUES = Literal[
-    "just_resize",
-    "crop_resize",
-    "fill_resize",
-    "just_resize_simple",
-]
-CONTROLNET_MODE_VALUES = Literal["balanced", "more_prompt", "more_control", "unbalanced"]
-
 ###################################################################
 # Copy of scripts/lvminthin.py from Mikubill/sd-webui-controlnet
 ###################################################################
@@ -76,6 +68,17 @@ def lvmin_thin(x, prunings=True):
     return y
 
 
+def nake_nms(x):
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+    y = np.zeros_like(x)
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+    return y
+
+
 ################################################################################
 # copied from Mikubill/sd-webui-controlnet external_code.py and modified for InvokeAI
 ################################################################################
@@ -131,122 +134,98 @@ def pixel_perfect_resolution(
     return int(np.round(estimation))
 
 
-def clone_contiguous(x: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
-    """Get a memory-contiguous clone of the given numpy array, as a safety measure and to improve computation efficiency."""
-    return np.ascontiguousarray(x).copy()
-
-
-def np_img_to_torch(np_img: np.ndarray[Any, Any], device: torch.device) -> torch.Tensor:
-    """Convert a numpy image to a PyTorch tensor. The image is normalized to 0-1, rearranged to BCHW format and sent to
-    the specified device."""
-
-    torch_img = torch.from_numpy(np_img)
-    normalized = torch_img.float() / 255.0
-    bchw = rearrange(normalized, "h w c -> 1 c h w")
-    on_device = bchw.to(device)
-    return on_device.clone()
-
-
-def heuristic_resize(np_img: np.ndarray[Any, Any], size: tuple[int, int]) -> np.ndarray[Any, Any]:
-    """Resizes an image using a heuristic to choose the best resizing strategy.
-
-    - If the image appears to be an edge map, special handling will be applied to ensure the edges are not distorted.
-    - Single-pixel edge maps use NMS and thinning to keep the edges as single-pixel lines.
-    - Low-color-count images are resized with nearest-neighbor to preserve color information (for e.g. segmentation maps).
-    - The alpha channel is handled separately to ensure it is resized correctly.
-
-    Args:
-        np_img (np.ndarray): The input image.
-        size (tuple[int, int]): The target size for the image.
-
-    Returns:
-        np.ndarray: The resized image.
-
-    Adapted from https://github.com/Mikubill/sd-webui-controlnet.
-    """
-
-    # Return early if the image is already at the requested size
-    if np_img.shape[0] == size[1] and np_img.shape[1] == size[0]:
-        return np_img
-
-    # If the image has an alpha channel, separate it for special handling later.
-    inpaint_mask = None
-    if np_img.ndim == 3 and np_img.shape[2] == 4:
-        inpaint_mask = np_img[:, :, 3]
-        np_img = np_img[:, :, 0:3]
-
-    new_size_is_smaller = (size[0] * size[1]) < (np_img.shape[0] * np_img.shape[1])
-    new_size_is_bigger = (size[0] * size[1]) > (np_img.shape[0] * np_img.shape[1])
-    unique_color_count = np.unique(np_img.reshape(-1, np_img.shape[2]), axis=0).shape[0]
-    is_one_pixel_edge = False
-    is_binary = False
-
-    if unique_color_count == 2:
-        # If the image has only two colors, it is likely binary. Check if the image has one-pixel edges.
-        is_binary = np.min(np_img) < 16 and np.max(np_img) > 240
-        if is_binary:
-            eroded = cv2.erode(np_img, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
-            dilated = cv2.dilate(eroded, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
-            one_pixel_edge_count = np.where(dilated < np_img)[0].shape[0]
-            all_edge_count = np.where(np_img > 127)[0].shape[0]
-            is_one_pixel_edge = one_pixel_edge_count * 2 > all_edge_count
-
-    if 2 < unique_color_count < 200:
-        # With a low color count, we assume this is a map where exact colors are important. Near-neighbor preserves
-        # the colors as needed.
-        interpolation = cv2.INTER_NEAREST
-    elif new_size_is_smaller:
-        # This works best for downscaling
-        interpolation = cv2.INTER_AREA
-    else:
-        # Fall back for other cases
-        interpolation = cv2.INTER_CUBIC  # Must be CUBIC because we now use nms. NEVER CHANGE THIS
-
-    # This may be further transformed depending on the binary nature of the image.
-    resized = cv2.resize(np_img, size, interpolation=interpolation)
-
-    if inpaint_mask is not None:
-        # Resize the inpaint mask to match the resized image using the same interpolation method.
-        inpaint_mask = cv2.resize(inpaint_mask, size, interpolation=interpolation)
-
-    # If the image is binary, we will perform some additional processing to ensure the edges are preserved.
-    if is_binary:
-        resized = np.mean(resized.astype(np.float32), axis=2).clip(0, 255).astype(np.uint8)
-        if is_one_pixel_edge:
-            # Use NMS and thinning to keep the edges as single-pixel lines.
-            resized = nms(resized)
-            _, resized = cv2.threshold(resized, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-            resized = lvmin_thin(resized, prunings=new_size_is_bigger)
-        else:
-            _, resized = cv2.threshold(resized, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-        resized = np.stack([resized] * 3, axis=2)
-
-    # Restore the alpha channel if it was present.
-    if inpaint_mask is not None:
-        inpaint_mask = (inpaint_mask > 127).astype(np.float32) * 255.0
-        inpaint_mask = inpaint_mask[:, :, None].clip(0, 255).astype(np.uint8)
-        resized = np.concatenate([resized, inpaint_mask], axis=2)
-
-    return resized
-
-
 ###########################################################################
 # Copied from detectmap_proc method in scripts/detectmap_proc.py in Mikubill/sd-webui-controlnet
 #    modified for InvokeAI
 ###########################################################################
-def np_img_resize(
-    np_img: np.ndarray,
-    resize_mode: CONTROLNET_RESIZE_VALUES,
-    h: int,
-    w: int,
-    device: torch.device = torch.device("cpu"),
-) -> tuple[torch.Tensor, np.ndarray[Any, Any]]:
-    np_img = normalize_image_channel_count(np_img)
+# def detectmap_proc(detected_map, module, resize_mode, h, w):
+def np_img_resize(np_img: np.ndarray, resize_mode: str, h: int, w: int, device: torch.device = torch.device("cpu")):
+    # if 'inpaint' in module:
+    #     np_img = np_img.astype(np.float32)
+    # else:
+    #     np_img = HWC3(np_img)
+    np_img = HWC3(np_img)
 
+    def safe_numpy(x):
+        # A very safe method to make sure that Apple/Mac works
+        y = x
+
+        # below is very boring but do not change these. If you change these Apple or Mac may fail.
+        y = y.copy()
+        y = np.ascontiguousarray(y)
+        y = y.copy()
+        return y
+
+    def get_pytorch_control(x):
+        # A very safe method to make sure that Apple/Mac works
+        y = x
+
+        # below is very boring but do not change these. If you change these Apple or Mac may fail.
+        y = torch.from_numpy(y)
+        y = y.float() / 255.0
+        y = rearrange(y, "h w c -> 1 c h w")
+        y = y.clone()
+        # y = y.to(devices.get_device_for("controlnet"))
+        y = y.to(device)
+        y = y.clone()
+        return y
+
+    def high_quality_resize(x: np.ndarray, size):
+        # Written by lvmin
+        # Super high-quality control map up-scaling, considering binary, seg, and one-pixel edges
+        inpaint_mask = None
+        if x.ndim == 3 and x.shape[2] == 4:
+            inpaint_mask = x[:, :, 3]
+            x = x[:, :, 0:3]
+
+        new_size_is_smaller = (size[0] * size[1]) < (x.shape[0] * x.shape[1])
+        new_size_is_bigger = (size[0] * size[1]) > (x.shape[0] * x.shape[1])
+        unique_color_count = np.unique(x.reshape(-1, x.shape[2]), axis=0).shape[0]
+        is_one_pixel_edge = False
+        is_binary = False
+        if unique_color_count == 2:
+            is_binary = np.min(x) < 16 and np.max(x) > 240
+            if is_binary:
+                xc = x
+                xc = cv2.erode(xc, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
+                xc = cv2.dilate(xc, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
+                one_pixel_edge_count = np.where(xc < x)[0].shape[0]
+                all_edge_count = np.where(x > 127)[0].shape[0]
+                is_one_pixel_edge = one_pixel_edge_count * 2 > all_edge_count
+
+        if 2 < unique_color_count < 200:
+            interpolation = cv2.INTER_NEAREST
+        elif new_size_is_smaller:
+            interpolation = cv2.INTER_AREA
+        else:
+            interpolation = cv2.INTER_CUBIC  # Must be CUBIC because we now use nms. NEVER CHANGE THIS
+
+        y = cv2.resize(x, size, interpolation=interpolation)
+        if inpaint_mask is not None:
+            inpaint_mask = cv2.resize(inpaint_mask, size, interpolation=interpolation)
+
+        if is_binary:
+            y = np.mean(y.astype(np.float32), axis=2).clip(0, 255).astype(np.uint8)
+            if is_one_pixel_edge:
+                y = nake_nms(y)
+                _, y = cv2.threshold(y, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+                y = lvmin_thin(y, prunings=new_size_is_bigger)
+            else:
+                _, y = cv2.threshold(y, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+            y = np.stack([y] * 3, axis=2)
+
+        if inpaint_mask is not None:
+            inpaint_mask = (inpaint_mask > 127).astype(np.float32) * 255.0
+            inpaint_mask = inpaint_mask[:, :, None].clip(0, 255).astype(np.uint8)
+            y = np.concatenate([y, inpaint_mask], axis=2)
+
+        return y
+
+    # if resize_mode == external_code.ResizeMode.RESIZE:
     if resize_mode == "just_resize":  # RESIZE
-        np_img = heuristic_resize(np_img, (w, h))
-        np_img = clone_contiguous(np_img)
-        return np_img_to_torch(np_img, device), np_img
+        np_img = high_quality_resize(np_img, (w, h))
+        np_img = safe_numpy(np_img)
+        return get_pytorch_control(np_img), np_img
 
     old_h, old_w, _ = np_img.shape
     old_w = float(old_w)
@@ -257,6 +236,7 @@ def np_img_resize(
     def safeint(x: Union[int, float]) -> int:
         return int(np.round(x))
 
+    # if resize_mode == external_code.ResizeMode.OUTER_FIT:
     if resize_mode == "fill_resize":  # OUTER_FIT
         k = min(k0, k1)
         borders = np.concatenate([np_img[0, :, :], np_img[-1, :, :], np_img[:, 0, :], np_img[:, -1, :]], axis=0)
@@ -265,23 +245,23 @@ def np_img_resize(
             # Inpaint hijack
             high_quality_border_color[3] = 255
         high_quality_background = np.tile(high_quality_border_color[None, None], [h, w, 1])
-        np_img = heuristic_resize(np_img, (safeint(old_w * k), safeint(old_h * k)))
+        np_img = high_quality_resize(np_img, (safeint(old_w * k), safeint(old_h * k)))
         new_h, new_w, _ = np_img.shape
         pad_h = max(0, (h - new_h) // 2)
         pad_w = max(0, (w - new_w) // 2)
         high_quality_background[pad_h : pad_h + new_h, pad_w : pad_w + new_w] = np_img
         np_img = high_quality_background
-        np_img = clone_contiguous(np_img)
-        return np_img_to_torch(np_img, device), np_img
+        np_img = safe_numpy(np_img)
+        return get_pytorch_control(np_img), np_img
     else:  # resize_mode == "crop_resize"  (INNER_FIT)
         k = max(k0, k1)
-        np_img = heuristic_resize(np_img, (safeint(old_w * k), safeint(old_h * k)))
+        np_img = high_quality_resize(np_img, (safeint(old_w * k), safeint(old_h * k)))
         new_h, new_w, _ = np_img.shape
         pad_h = max(0, (new_h - h) // 2)
         pad_w = max(0, (new_w - w) // 2)
         np_img = np_img[pad_h : pad_h + h, pad_w : pad_w + w]
-        np_img = clone_contiguous(np_img)
-        return np_img_to_torch(np_img, device), np_img
+        np_img = safe_numpy(np_img)
+        return get_pytorch_control(np_img), np_img
 
 
 def prepare_control_image(
@@ -289,12 +269,12 @@ def prepare_control_image(
     width: int,
     height: int,
     num_channels: int = 3,
-    device: str | torch.device = "cuda",
-    dtype: torch.dtype = torch.float16,
-    control_mode: CONTROLNET_MODE_VALUES = "balanced",
-    resize_mode: CONTROLNET_RESIZE_VALUES = "just_resize_simple",
-    do_classifier_free_guidance: bool = True,
-) -> torch.Tensor:
+    device="cuda",
+    dtype=torch.float16,
+    do_classifier_free_guidance=True,
+    control_mode="balanced",
+    resize_mode="just_resize_simple",
+):
     """Pre-process images for ControlNets or T2I-Adapters.
 
     Args:
@@ -304,7 +284,7 @@ def prepare_control_image(
         num_channels (int, optional): The target number of image channels. This is achieved by converting the input
             image to RGB, then naively taking the first `num_channels` channels. The primary use case is converting a
             RGB image to a single-channel grayscale image. Raises if `num_channels` cannot be achieved. Defaults to 3.
-        device (str | torch.Device, optional): The target device for the output image. Defaults to "cuda".
+        device (str, optional): The target device for the output image. Defaults to "cuda".
         dtype (_type_, optional): The dtype for the output image. Defaults to torch.float16.
         do_classifier_free_guidance (bool, optional): If True, repeat the output image along the batch dimension.
             Defaults to True.
@@ -312,15 +292,26 @@ def prepare_control_image(
         resize_mode (str, optional): Defaults to "just_resize_simple".
 
     Raises:
+        NotImplementedError: If resize_mode == "crop_resize_simple".
+        NotImplementedError: If resize_mode == "fill_resize_simple".
         ValueError: If `resize_mode` is not recognized.
         ValueError: If `num_channels` is out of range.
 
     Returns:
         torch.Tensor: The pre-processed input tensor.
     """
-    if resize_mode == "just_resize_simple":
+    if (
+        resize_mode == "just_resize_simple"
+        or resize_mode == "crop_resize_simple"
+        or resize_mode == "fill_resize_simple"
+    ):
         image = image.convert("RGB")
-        image = image.resize((width, height), resample=Image.LANCZOS)
+        if resize_mode == "just_resize_simple":
+            image = image.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])
+        elif resize_mode == "crop_resize_simple":
+            raise NotImplementedError(f"prepare_control_image is not implemented for resize_mode='{resize_mode}'.")
+        elif resize_mode == "fill_resize_simple":
+            raise NotImplementedError(f"prepare_control_image is not implemented for resize_mode='{resize_mode}'.")
         nimage = np.array(image)
         nimage = nimage[None, :]
         nimage = np.concatenate([nimage], axis=0)
@@ -337,7 +328,8 @@ def prepare_control_image(
             resize_mode=resize_mode,
             h=height,
             w=width,
-            device=torch.device(device),
+            # device=torch.device('cpu')
+            device=device,
         )
     else:
         raise ValueError(f"Unsupported resize_mode: '{resize_mode}'.")

invokeai/app/util/custom_openapi.py

@@ -1,116 +0,0 @@
-from typing import Any, Callable, Optional
-
-from fastapi import FastAPI
-from fastapi.openapi.utils import get_openapi
-from pydantic.json_schema import models_json_schema
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, UIConfigBase
-from invokeai.app.invocations.fields import InputFieldJSONSchemaExtra, OutputFieldJSONSchemaExtra
-from invokeai.app.invocations.model import ModelIdentifierField
-from invokeai.app.services.events.events_common import EventBase
-from invokeai.app.services.session_processor.session_processor_common import ProgressImage
-
-
-def move_defs_to_top_level(openapi_schema: dict[str, Any], component_schema: dict[str, Any]) -> None:
-    """Moves a component schema's $defs to the top level of the openapi schema. Useful when generating a schema
-    for a single model that needs to be added back to the top level of the schema. Mutates openapi_schema and
-    component_schema."""
-
-    defs = component_schema.pop("$defs", {})
-    for schema_key, json_schema in defs.items():
-        if schema_key in openapi_schema["components"]["schemas"]:
-            continue
-        openapi_schema["components"]["schemas"][schema_key] = json_schema
-
-
-def get_openapi_func(
-    app: FastAPI, post_transform: Optional[Callable[[dict[str, Any]], dict[str, Any]]] = None
-) -> Callable[[], dict[str, Any]]:
-    """Gets the OpenAPI schema generator function.
-
-    Args:
-        app (FastAPI): The FastAPI app to generate the schema for.
-        post_transform (Optional[Callable[[dict[str, Any]], dict[str, Any]]], optional): A function to apply to the
-            generated schema before returning it. Defaults to None.
-
-    Returns:
-        Callable[[], dict[str, Any]]: The OpenAPI schema generator function. When first called, the generated schema is
-            cached in `app.openapi_schema`. On subsequent calls, the cached schema is returned. This caching behaviour
-            matches FastAPI's default schema generation caching.
-    """
-
-    def openapi() -> dict[str, Any]:
-        if app.openapi_schema:
-            return app.openapi_schema
-
-        openapi_schema = get_openapi(
-            title=app.title,
-            description="An API for invoking AI image operations",
-            version="1.0.0",
-            routes=app.routes,
-            separate_input_output_schemas=False,  # https://fastapi.tiangolo.com/how-to/separate-openapi-schemas/
-        )
-
-        # We'll create a map of invocation type to output schema to make some types simpler on the client.
-        invocation_output_map_properties: dict[str, Any] = {}
-        invocation_output_map_required: list[str] = []
-
-        # We need to manually add all outputs to the schema - pydantic doesn't add them because they aren't used directly.
-        for output in BaseInvocationOutput.get_outputs():
-            json_schema = output.model_json_schema(mode="serialization", ref_template="#/components/schemas/{model}")
-            move_defs_to_top_level(openapi_schema, json_schema)
-            openapi_schema["components"]["schemas"][output.__name__] = json_schema
-
-        # Technically, invocations are added to the schema by pydantic, but we still need to manually set their output
-        # property, so we'll just do it all manually.
-        for invocation in BaseInvocation.get_invocations():
-            json_schema = invocation.model_json_schema(
-                mode="serialization", ref_template="#/components/schemas/{model}"
-            )
-            move_defs_to_top_level(openapi_schema, json_schema)
-            output_title = invocation.get_output_annotation().__name__
-            outputs_ref = {"$ref": f"#/components/schemas/{output_title}"}
-            json_schema["output"] = outputs_ref
-            openapi_schema["components"]["schemas"][invocation.__name__] = json_schema
-
-            # Add this invocation and its output to the output map
-            invocation_type = invocation.get_type()
-            invocation_output_map_properties[invocation_type] = json_schema["output"]
-            invocation_output_map_required.append(invocation_type)
-
-        # Add the output map to the schema
-        openapi_schema["components"]["schemas"]["InvocationOutputMap"] = {
-            "type": "object",
-            "properties": invocation_output_map_properties,
-            "required": invocation_output_map_required,
-        }
-
-        # Some models don't end up in the schemas as standalone definitions because they aren't used directly in the API.
-        # We need to add them manually here. WARNING: Pydantic can choke if you call `model.model_json_schema()` to get
-        # a schema. This has something to do with schema refs - not totally clear. For whatever reason, using
-        # `models_json_schema` seems to work fine.
-        additional_models = [
-            *EventBase.get_events(),
-            UIConfigBase,
-            InputFieldJSONSchemaExtra,
-            OutputFieldJSONSchemaExtra,
-            ModelIdentifierField,
-            ProgressImage,
-        ]
-
-        additional_schemas = models_json_schema(
-            [(m, "serialization") for m in additional_models],
-            ref_template="#/components/schemas/{model}",
-        )
-        # additional_schemas[1] is a dict of $defs that we need to add to the top level of the schema
-        move_defs_to_top_level(openapi_schema, additional_schemas[1])
-
-        if post_transform is not None:
-            openapi_schema = post_transform(openapi_schema)
-
-        openapi_schema["components"]["schemas"] = dict(sorted(openapi_schema["components"]["schemas"].items()))
-
-        app.openapi_schema = openapi_schema
-        return app.openapi_schema
-
-    return openapi

invokeai/app/util/download_with_progress.py

@@ -0,0 +1,51 @@
+from pathlib import Path
+from urllib import request
+
+from tqdm import tqdm
+
+from invokeai.backend.util.logging import InvokeAILogger
+
+
+class ProgressBar:
+    """Simple progress bar for urllib.request.urlretrieve using tqdm."""
+
+    def __init__(self, model_name: str = "file"):
+        self.pbar = None
+        self.name = model_name
+
+    def __call__(self, block_num: int, block_size: int, total_size: int):
+        if not self.pbar:
+            self.pbar = tqdm(
+                desc=self.name,
+                initial=0,
+                unit="iB",
+                unit_scale=True,
+                unit_divisor=1000,
+                total=total_size,
+            )
+        self.pbar.update(block_size)
+
+
+def download_with_progress_bar(name: str, url: str, dest_path: Path) -> bool:
+    """Download a file from a URL to a destination path, with a progress bar.
+    If the file already exists, it will not be downloaded again.
+
+    Exceptions are not caught.
+
+    Args:
+        name (str): Name of the file being downloaded.
+        url (str): URL to download the file from.
+        dest_path (Path): Destination path to save the file to.
+
+    Returns:
+        bool: True if the file was downloaded, False if it already existed.
+    """
+    if dest_path.exists():
+        return False  # already downloaded
+
+    InvokeAILogger.get_logger().info(f"Downloading {name}...")
+
+    dest_path.parent.mkdir(parents=True, exist_ok=True)
+    request.urlretrieve(url, dest_path, ProgressBar(name))
+
+    return True

invokeai/app/util/step_callback.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import TYPE_CHECKING, Callable
 
 import torch
 from PIL import Image
@@ -13,36 +13,8 @@ if TYPE_CHECKING:
     from invokeai.app.services.events.events_base import EventServiceBase
     from invokeai.app.services.shared.invocation_context import InvocationContextData
 
-# fast latents preview matrix for sdxl
-# generated by @StAlKeR7779
-SDXL_LATENT_RGB_FACTORS = [
-    #   R        G        B
-    [0.3816, 0.4930, 0.5320],
-    [-0.3753, 0.1631, 0.1739],
-    [0.1770, 0.3588, -0.2048],
-    [-0.4350, -0.2644, -0.4289],
-]
-SDXL_SMOOTH_MATRIX = [
-    [0.0358, 0.0964, 0.0358],
-    [0.0964, 0.4711, 0.0964],
-    [0.0358, 0.0964, 0.0358],
-]
 
-# origingally adapted from code by @erucipe and @keturn here:
-# https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7
-# these updated numbers for v1.5 are from @torridgristle
-SD1_5_LATENT_RGB_FACTORS = [
-    #    R        G        B
-    [0.3444, 0.1385, 0.0670],  # L1
-    [0.1247, 0.4027, 0.1494],  # L2
-    [-0.3192, 0.2513, 0.2103],  # L3
-    [-0.1307, -0.1874, -0.7445],  # L4
-]
-
-
-def sample_to_lowres_estimated_image(
-    samples: torch.Tensor, latent_rgb_factors: torch.Tensor, smooth_matrix: Optional[torch.Tensor] = None
-):
+def sample_to_lowres_estimated_image(samples, latent_rgb_factors, smooth_matrix=None):
     latent_image = samples[0].permute(1, 2, 0) @ latent_rgb_factors
 
     if smooth_matrix is not None:
@@ -75,12 +47,64 @@ def stable_diffusion_step_callback(
     else:
         sample = intermediate_state.latents
 
+    # TODO: This does not seem to be needed any more?
+    # # txt2img provides a Tensor in the step_callback
+    # # img2img provides a PipelineIntermediateState
+    # if isinstance(sample, PipelineIntermediateState):
+    #     # this was an img2img
+    #     print('img2img')
+    #     latents = sample.latents
+    #     step = sample.step
+    # else:
+    #     print('txt2img')
+    #     latents = sample
+    #     step = intermediate_state.step
+
+    # TODO: only output a preview image when requested
+
     if base_model in [BaseModelType.StableDiffusionXL, BaseModelType.StableDiffusionXLRefiner]:
-        sdxl_latent_rgb_factors = torch.tensor(SDXL_LATENT_RGB_FACTORS, dtype=sample.dtype, device=sample.device)
-        sdxl_smooth_matrix = torch.tensor(SDXL_SMOOTH_MATRIX, dtype=sample.dtype, device=sample.device)
+        # fast latents preview matrix for sdxl
+        # generated by @StAlKeR7779
+        sdxl_latent_rgb_factors = torch.tensor(
+            [
+                #   R        G        B
+                [0.3816, 0.4930, 0.5320],
+                [-0.3753, 0.1631, 0.1739],
+                [0.1770, 0.3588, -0.2048],
+                [-0.4350, -0.2644, -0.4289],
+            ],
+            dtype=sample.dtype,
+            device=sample.device,
+        )
+
+        sdxl_smooth_matrix = torch.tensor(
+            [
+                [0.0358, 0.0964, 0.0358],
+                [0.0964, 0.4711, 0.0964],
+                [0.0358, 0.0964, 0.0358],
+            ],
+            dtype=sample.dtype,
+            device=sample.device,
+        )
+
         image = sample_to_lowres_estimated_image(sample, sdxl_latent_rgb_factors, sdxl_smooth_matrix)
     else:
-        v1_5_latent_rgb_factors = torch.tensor(SD1_5_LATENT_RGB_FACTORS, dtype=sample.dtype, device=sample.device)
+        # origingally adapted from code by @erucipe and @keturn here:
+        # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7
+
+        # these updated numbers for v1.5 are from @torridgristle
+        v1_5_latent_rgb_factors = torch.tensor(
+            [
+                #    R        G        B
+                [0.3444, 0.1385, 0.0670],  # L1
+                [0.1247, 0.4027, 0.1494],  # L2
+                [-0.3192, 0.2513, 0.2103],  # L3
+                [-0.1307, -0.1874, -0.7445],  # L4
+            ],
+            dtype=sample.dtype,
+            device=sample.device,
+        )
+
         image = sample_to_lowres_estimated_image(sample, v1_5_latent_rgb_factors)
 
     (width, height) = image.size
@@ -89,9 +113,15 @@ def stable_diffusion_step_callback(
 
     dataURL = image_to_dataURL(image, image_format="JPEG")
 
-    events.emit_invocation_denoise_progress(
-        context_data.queue_item,
-        context_data.invocation,
-        intermediate_state,
-        ProgressImage(dataURL=dataURL, width=width, height=height),
+    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,
+        total_steps=intermediate_state.total_steps,
     )

invokeai/backend/image_util/__init__.py

@@ -4,4 +4,5 @@ Initialization file for invokeai.backend.image_util methods.
 
 from .infill_methods.patchmatch import PatchMatch  # noqa: F401
 from .pngwriter import PngWriter, PromptFormatter, retrieve_metadata, write_metadata  # noqa: F401
+from .seamless import configure_model_padding  # noqa: F401
 from .util import InitImageResizer, make_grid  # noqa: F401

invokeai/backend/image_util/depth_anything/__init__.py

@@ -1,5 +1,5 @@
-from pathlib import Path
-from typing import Literal
+import pathlib
+from typing import Literal, Union
 
 import cv2
 import numpy as np
@@ -10,17 +10,28 @@ from PIL import Image
 from torchvision.transforms import Compose
 
 from invokeai.app.services.config.config_default import get_config
+from invokeai.app.util.download_with_progress import download_with_progress_bar
 from invokeai.backend.image_util.depth_anything.model.dpt import DPT_DINOv2
 from invokeai.backend.image_util.depth_anything.utilities.util import NormalizeImage, PrepareForNet, Resize
+from invokeai.backend.util.devices import choose_torch_device
 from invokeai.backend.util.logging import InvokeAILogger
 
 config = get_config()
 logger = InvokeAILogger.get_logger(config=config)
 
 DEPTH_ANYTHING_MODELS = {
-    "large": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitl14.pth?download=true",
-    "base": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitb14.pth?download=true",
-    "small": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vits14.pth?download=true",
+    "large": {
+        "url": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitl14.pth?download=true",
+        "local": "any/annotators/depth_anything/depth_anything_vitl14.pth",
+    },
+    "base": {
+        "url": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitb14.pth?download=true",
+        "local": "any/annotators/depth_anything/depth_anything_vitb14.pth",
+    },
+    "small": {
+        "url": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vits14.pth?download=true",
+        "local": "any/annotators/depth_anything/depth_anything_vits14.pth",
+    },
 }
 
 
@@ -42,27 +53,36 @@ transform = Compose(
 
 
 class DepthAnythingDetector:
-    def __init__(self, model: DPT_DINOv2, device: torch.device) -> None:
-        self.model = model
-        self.device = device
+    def __init__(self) -> None:
+        self.model = None
+        self.model_size: Union[Literal["large", "base", "small"], None] = None
+        self.device = choose_torch_device()
 
-    @staticmethod
-    def load_model(
-        model_path: Path, device: torch.device, model_size: Literal["large", "base", "small"] = "small"
-    ) -> DPT_DINOv2:
-        match model_size:
-            case "small":
-                model = DPT_DINOv2(encoder="vits", features=64, out_channels=[48, 96, 192, 384])
-            case "base":
-                model = DPT_DINOv2(encoder="vitb", features=128, out_channels=[96, 192, 384, 768])
-            case "large":
-                model = DPT_DINOv2(encoder="vitl", features=256, out_channels=[256, 512, 1024, 1024])
+    def load_model(self, model_size: Literal["large", "base", "small"] = "small"):
+        DEPTH_ANYTHING_MODEL_PATH = config.models_path / DEPTH_ANYTHING_MODELS[model_size]["local"]
+        download_with_progress_bar(
+            pathlib.Path(DEPTH_ANYTHING_MODELS[model_size]["url"]).name,
+            DEPTH_ANYTHING_MODELS[model_size]["url"],
+            DEPTH_ANYTHING_MODEL_PATH,
+        )
 
-        model.load_state_dict(torch.load(model_path.as_posix(), map_location="cpu"))
-        model.eval()
+        if not self.model or model_size != self.model_size:
+            del self.model
+            self.model_size = model_size
 
-        model.to(device)
-        return model
+            match self.model_size:
+                case "small":
+                    self.model = DPT_DINOv2(encoder="vits", features=64, out_channels=[48, 96, 192, 384])
+                case "base":
+                    self.model = DPT_DINOv2(encoder="vitb", features=128, out_channels=[96, 192, 384, 768])
+                case "large":
+                    self.model = DPT_DINOv2(encoder="vitl", features=256, out_channels=[256, 512, 1024, 1024])
+
+            self.model.load_state_dict(torch.load(DEPTH_ANYTHING_MODEL_PATH.as_posix(), map_location="cpu"))
+            self.model.eval()
+
+        self.model.to(choose_torch_device())
+        return self.model
 
     def __call__(self, image: Image.Image, resolution: int = 512) -> Image.Image:
         if not self.model:
@@ -74,7 +94,7 @@ class DepthAnythingDetector:
 
         image_height, image_width = np_image.shape[:2]
         np_image = transform({"image": np_image})["image"]
-        tensor_image = torch.from_numpy(np_image).unsqueeze(0).to(self.device)
+        tensor_image = torch.from_numpy(np_image).unsqueeze(0).to(choose_torch_device())
 
         with torch.no_grad():
             depth = self.model(tensor_image)