Refactor Model Management

[psychedelicious]

Refactor Model Installation

Our current install procedure involves two main scripts:

• A simple script to install python packages for the app
• A more complex configuration script, which installs the models in a TUI (terminal user interface)

The TUI is fairly difficult to maintain, with numerous edge cases related to terminal behavior across different platforms and the amount of content to be displayed in the TUI.

Functionally, it sits in an awkward place between the two typically use-cases:
-Headless installation would rather provide a config file that says exactly which models to install

• "Normal" users would rather use a GUI to do all model installation

A TUI is mostly good enough for either of these, but it's not exactly the right tool for either job.

Revised Installation flow

If we were to redesign the app and model installation flow from the ground up, what might it look like?

App Installation

This is a very simple script. It installs the python packages only.

The UI is launchable after this, though of course it cannot generate at this point, because models are not installed or set up.

App Launch

When the application launches, and the ModelManager is instantiated, it checks if its folder structure is correct. If not, it creates the default folder structure.

It also checks if all the support models are installed, and at least one SD model (?). The UI hits a route (maybe the app info route? or something on models/) that says if we have all the core models set up.

If something is missing, the user is prompted to go to the Model Manager UI to get everything set up.

Model Manager UI

New tabs here replace the TUI. Core model installation status is displayed.

A single button kicks off the core model installation if it hasn't been completed or is missing something.

Here we present recommended models for install. This includes ControlNet models, LoRAs, embeddings, VAEs and main models.

Integrated with models.invoke.ai, and retains the existing functionality to install models from repo IDs, or import from file.

Headless

The app may optionally be provided a model install file:
invokeai --models_from_file models.yaml

If this is received, it is parsed and used to queue up and kick off the model installation. The web API is not fired up until this is complete (?).

Model Services

The current Model Manager service is split up and new services are introduced.

The services are all narrow in scope and fairly low-level. Each has an ABC and light implementations for the OSS application, per our service architecture.

Services Overview

• ModelQueueManager: manages the simple model install queue
  • wraps a python queue library eg Celery
  • methods:
    • add_to_queue
    • remove_from_queue
    • get_status
    • reorder_queue
• ModelTaskHandler: handles all side-effects/actions
  • ModelQueueManager calls this service to action queue items
  • methods:
    • download_model
    • configure_model
• ModelStatusTracker: tracks the status of each individual model install, in-memory store
  • methods:
    • get_status
    • update_status
    • send_update (using events service)
  • maybe this could be wrapped up in another service?
• ModelRecordStorage: manages the db for models
  • see [enhancement]: model config storage & data structure #4043
• ModelFileStorage: manages model files
• ModelLoader: everything related to actually loading and providing models to the application; i.e. most of what the ModelManager does today including the model context manager

Model Install Queue Items

Queue items would look something like this:

• model_id: unique identifier for the model
• model_type: main, controlnet, etc
• model_url: where the model is downloaded from
• status: status of model installation; queued, downloading, configuring, ready, or error
• download_progress: float 0 to 100 indicating download progress
• error: error if there is any
• created_at: when the item was added to queue
• updated_at: last status update of the queue item

Model Install Events

These would include the queue item, or the relevant parts of it, as payload.

• model_queuedmodel install added to queue
• model_download_started
• model_download_progress: emitted once per second (?)
• model_download_completed
• model_configuring: the model configuration has started
• model_ready: download & config are complete, model is ready
• model_error
• model_removed
  • model install removed from queue, either user canceled or error

Web API

Routes would be something like this:

• install_model
• get_install_status
• get_install_queue
• cancel_model_install
• reorder_install_queue

Summary

This is a pretty chonky refactor, but I think necessary if we want to provide a robust model system with a good user experience.

Feels like any improvements we want to make to the current system are actually steps towards something like this proposal.

We may as well just implement it properly as services, rather than stuff it all into the ever-growing model manager service. Makes it way easier to reason about, test and debug.

The individual services are all narrow and easy to tackle one at at time, and none of them is trying to do anything complex. I suspect splitting up the model manager will be the hardest part.

[lstein]

I like the proposed refactoring, but I'd organize it a little differently. How about something like this?

class ModelSearch

The ModelSearch class provides a search interface on models, taking a series of desired attributes to search on and returning a series of records for matching models. Search dimensions can include:

• Base model
• Model name
• Model type
• Model format / architecture
• Model hash
• Model location (full or partial path match)
• Text match in the description field
• Search by installed state
• Search by tag (new optional attribute list)
• Search by license (new optional attribute)
• Search by author (new optional attribute)

As an abstract base class, we'd be able to create a compatible module that can search remote sites as well, such as HuggingFace and models.invoke.ai.

The records returned from the ModelSearch.search() call will contain all the metadata relevant to the matching models, along with the model's unique key.

class ModelLoader

This class manages instantiating a model and loading it from disk. It returns an object which can be used as a context manager to prepare the model for action (i.e. moving it into GPU).

The get_model() provided by this model take the model's unique key as returned by ModelSearch. It uses ModelRecordStorage as its backend.

class ModelRecordStorage

As per the first posting on this thread. This manages local storage for models, whether on disk, a DBM, or a

class ModelInstall

This class initiates the installation, deletion, editing and conversion of models, using the ModelQueueManager as its backend.

class ModelQueueManager

As per the first posting on this thread, this maintains the queue of installation tasks, such as downloading and conversion.

[bakkot]

Will this mean the end of the scourge of .cache/huggingface? (#4095, #4928, https://github.com/invoke-ai/InvokeAI/discussions/1361#discussioncomment-7561208, #4585, https://github.com/invoke-ai/InvokeAI/discussions/1799, #5029, #4991, https://github.com/invoke-ai/InvokeAI/discussions/4989, etc) 🤞

[joshistoast]

Will this take into account custom upscale models?