Introduction to the Model Manager

The Model Manager is responsible for organizing the various machine learning models used by InvokeAI. It consists of a series of interdependent services that together handle the full lifecycle of a model. These are the:

• ModelRecordServiceBase: Responsible for managing model metadata and configuration information. Among other things, the record service tracks the type of the model, its provenance, and where it can be found on disk.
• ModelInstallServiceBase: A service for installing models to disk. It uses DownloadQueueServiceBase to download models and their metadata, and ModelRecordServiceBase to store that information. It is also responsible for managing the InvokeAI models directory and its contents.
• DownloadQueueServiceBase: A multithreaded downloader responsible for downloading models from a remote source to disk. The download queue has special methods for downloading repo_id folders from Hugging Face, as well as discriminating among model versions in Civitai, but can be used for arbitrary content.
• ModelLoadServiceBase Responsible for loading a model from disk into RAM and VRAM and getting it ready for inference.

Location of the Code

The four main services can be found in invokeai/app/services in the following directories:

• Directorioinvokeai

  • Directorioapp

    • Directorioservices Model manager services

      • Directoriomodel_records/

        • …
      • Directoriomodel_install/

        • …
      • Directoriodownloads/

        • …
      • Directoriomodel_load/

        • …
  • Directorioapi

    • Directoriorouters

      • model_manager_v2.py FastAPI web API for model management

---

What’s in a Model? The ModelRecordService

The ModelRecordService manages the model’s metadata. It supports a hierarchy of pydantic metadata “config” objects, which become increasingly specialized to support particular model types.

ModelConfigBase

All model metadata classes inherit from this pydantic class. it provides the following fields:

Field Name	Type	Description
key	str	Unique identifier for the model
name	str	Name of the model (not unique)
model_type	ModelType	The type of the model
model_format	ModelFormat	The format of the model (e.g. “diffusers”); also used as a Union discriminator
base_model	BaseModelType	The base model that the model is compatible with
path	str	Location of model on disk
hash	str	Hash of the model
description	str	Human-readable description of the model (optional)
source	str	Model’s source URL or repo id (optional)

The key is a unique 32-character random ID which was generated at install time. The hash field stores a hash of the model’s contents at install time obtained by sampling several parts of the model’s files using the imohash library. Over the course of the 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.

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.

ModelType, ModelFormat and BaseModelType are string enums that are defined in invokeai.backend.model_manager.config. They are also imported by, and can be reexported from, invokeai.app.services.model_manager.model_records:

invokeai.backend.model_manager.config

from invokeai.app.services.model_records import ModelType, ModelFormat, BaseModelType

CheckpointConfig

This adds support for checkpoint configurations, and adds the following field:

Field Name	Type	Description
config	str	Path to the checkpoint’s config file

config is the path to the checkpoint’s config file. If relative, it is taken to be relative to the InvokeAI root directory (e.g. configs/stable-diffusion/v1-inference.yaml)

MainConfig

This adds support for “main” Stable Diffusion models, and adds these fields:

Field Name	Type	Description
vae	str	Path to a VAE to use instead of the burnt-in one
variant	ModelVariantType	Model variant type, such as “inpainting”

vae can be an absolute or relative path. If relative, its base is taken to be the models_dir directory.

variant is an enumerated string class with values normal, inpaint and depth. If needed, it can be imported if needed from either invokeai.app.services.model_records or invokeai.backend.model_manager.config.

ONNXSD2Config

Field Name	Type	Description
prediction_type	SchedulerPredictionType	Scheduler prediction type to use, e.g. “epsilon”
upcast_attention	bool	Model requires its attention module to be upcast

The SchedulerPredictionType enum can be imported from either invokeai.app.services.model_records or invokeai.backend.model_manager.config.

Other config classes

There are a series of such classes each discriminated by their ModelFormat, including LoRAConfig, IPAdapterConfig, and so forth. These are rarely needed outside the model manager’s internal code, but available in invokeai.backend.model_manager.config if needed. There is also a Union of all ModelConfig classes, called AnyModelConfig that can be imported from the same file.

Limitations of the Data Model

The config hierarchy has a major limitation in its handling of the base model type. Each model can only be compatible with one base model, which breaks down in the event of models that are compatible with two or more base models. For example, SD-1 VAEs also work with SD-2 models. A partial workaround is to use BaseModelType.Any, which indicates that the model is compatible with any of the base models. This works OK for some models, such as the IP Adapter image encoders, but is an all-or-nothing proposition.

Reading and Writing Model Configuration Records

The model record service keeps each installed model’s configuration record in the models table of the InvokeAI database, as the config’s JSON; the columns models are searched by (name, base, type, format, path, hash, …) are generated from that document. ModelRecordServiceSQL implements ModelRecordServiceBase on the database layer.

An application-wide record service is created during API initialization and can be retrieved within an invocation from the InvocationContext object:

store = context.services.model_manager.store

or from elsewhere in the code by accessing ApiDependencies.invoker.services.model_manager.store.

Creating a ModelRecordServiceSQL

The service takes the application’s database and a logger. To open the records of an existing, migrated database outside the server:

from invokeai.app.services.config import get_config
from invokeai.app.services.model_records import ModelRecordServiceSQL
from invokeai.app.services.shared.database.database import Database
from invokeai.backend.util.logging import InvokeAILogger

config = get_config()
logger = InvokeAILogger.get_logger(config=config)
database = Database.open_sqlite(config.db_path, logger)
store = ModelRecordServiceSQL(database, logger)

Fetching a Model’s Configuration

Configurations can be retrieved in several ways.

get_model(key) -> AnyModelConfig

The basic functionality is to call the record store object’s get_model() method with the desired model’s unique key. It returns the appropriate subclass of ModelConfigBase:

model_conf = store.get_model('f13dd932c0c35c22dcb8d6cda4203764')
print(model_conf.path)

>> '/tmp/models/ckpts/v1-5-pruned-emaonly.safetensors'

If the key is unrecognized, this call raises an UnknownModelException.

exists(key) -> bool

Returns True if a model with the given key exists in the database.

search_by_path(path) -> List[AnyModelConfig]

Returns the configuration of the model whose path is path, in a list of at most one since paths are unique. The path is matched using a simple string comparison and won’t correctly match models referred to by different paths (e.g. using symbolic links).

search_by_hash(hash) -> List[AnyModelConfig]

Returns the configurations of the models whose file has the given hash.

search_by_attr(model_name, base_model, model_type, model_format, order_by, direction) -> List[AnyModelConfig]

Returns the models that match every filter given among model_name, base_model, model_type and model_format, ordered by a ModelRecordOrderBy (type, base, name and format by default; names and bases ignore case) in direction. Models that tie are ordered by key, in the same direction. Calling it without filters returns all the models in the database.

A stored config that no longer validates, for instance one written by a newer version of the application, is skipped with a warning in the log.

all_models() -> List[AnyModelConfig]

Return all the model configs in the database. Exactly equivalent to calling search_by_attr() with no arguments.

model_info_by_name(model_name, base_model, model_type) -> AnyModelConfig

Looks for the one model identified by name, base model and model type and returns it. It raises a DuplicateModelException if more than one model shares the same type, base and name, which is possible (one at /foo/my_model, another at /bar/my_model), and an UnknownModelException if none does. Prefer search_by_attr(), which can return several models, and pass the key of the one you want to get_model().

get_model_paths() -> List[str]

Returns the path of every record, including records whose config no longer validates. The orphaned-model scan uses it, so that it never reports the files of a registered model.

Writing model configs to the database

Several methods allow you to create and update stored model config records.

add_model(config) -> AnyModelConfig

Adds the configuration record; the config’s key becomes the record’s key. It returns the stored config, and raises a DuplicateModelException if a model with the same path or key is already installed.

update_model(key, changes, allow_class_change=False) -> AnyModelConfig

Applies a ModelRecordChanges to the record and returns the updated config. The changes are validated against the model’s config class; with allow_class_change, a change may turn the record into another class (a LoRA into a Main model, say), which must validate as well. The record is locked while the update is applied, so concurrent updates of one model apply one after the other. Raises UnknownModelException for an unknown key, pydantic’s ValidationError for invalid changes, and ValueError for a change of the key or a path longer than a record holds. With allow_class_change, a non-null change to a field that neither the current nor the resulting config class has also raises ValueError, so no change is silently dropped: None (clearing a value) is not refused this way, and a value only the current class stores is dropped when the class changes. Nothing is written when an update raises.

replace_model(key, new_config) -> AnyModelConfig

Replaces the whole record, for instance after a model was identified anew. new_config.key must equal key. Raises UnknownModelException for an unknown key.

del_model(key)

Deletes the record, and the model’s relationships with other models. Raises UnknownModelException for an unknown key.

---

Model installation

The ModelInstallService class implements the ModelInstallServiceBase abstract base class, and provides a one-stop shop for all your model install needs. It provides the following functionality:

• Registering a model config record for a model already located on the local filesystem, without moving it or changing its path.

• Installing a model alreadiy located on the local filesystem, by moving it into the InvokeAI root directory under the models folder (or wherever config parameter models_dir specifies).

• Probing of models to determine their type, base type and other key information.

• Interface with the InvokeAI event bus to provide status updates on the download, installation and registration process.

• Downloading a model from an arbitrary URL and installing it in models_dir.

• Special handling for HuggingFace repo_ids to recursively download the contents of the repository, paying attention to alternative variants such as fp16.

• Saving tags and other metadata about the model into the invokeai database when fetching from a repo that provides that type of information, (currently only HuggingFace).

Initializing the installer

A default installer is created at InvokeAI api startup time and stored in ApiDependencies.invoker.services.model_install and can also be retrieved from an invocation’s context argument with context.services.model_install.

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.model_records import ModelRecordServiceSQL
from invokeai.app.services.model_install import ModelInstallService
from invokeai.app.services.download import DownloadQueueService
from invokeai.app.services.shared.database.database import Database
from invokeai.backend.util.logging import InvokeAILogger

config = get_config()

logger = InvokeAILogger.get_logger(config=config)
database = Database.open_sqlite(config.db_path, logger)
record_store = ModelRecordServiceSQL(database, logger)
queue = DownloadQueueService()
queue.start()

installer = ModelInstallService(app_config=config,
                                record_store=record_store,
                                download_queue=queue
                                )
installer.start()

The full form of ModelInstallService() takes the following required parameters:

Argument	Type	Description
app_config	InvokeAIAppConfig	InvokeAI app configuration object
record_store	ModelRecordServiceBase	Config record storage database
download_queue	DownloadQueueServiceBase	Download queue object
session	Optional[requests.Session]	Swap in a different Session object (usually for debugging)

Once initialized, the installer will provide the following methods:

install_job = installer.heuristic_import(source, [config], [access_token])

This is a simplified interface to the installer which takes a source string, an optional model configuration dictionary and an optional access token.

The source is a string that can be any of these forms

1. A path on the local filesystem (C:\\users\\fred\\model.safetensors)
2. A Url pointing to a single downloadable model file (https://civitai.com/models/58390/detail-tweaker-lora-lora)
3. A HuggingFace repo_id with any of the following formats:
   • model/name — entire model
   • model/name:fp32 — entire model, using the fp32 variant
   • model/name:fp16:vae — vae submodel, using the fp16 variant
   • model/name::vae — vae submodel, using default precision
   • model/name:fp16:path/to/model.safetensors — an individual model file, fp16 variant
   • model/name::path/to/model.safetensors — an individual model file, default variant

Note that by specifying a relative path to the top of the HuggingFace repo, you can download and install arbitrary models files.

The variant, if not provided, will be automatically filled in with fp32 if the user has requested full precision, and fp16 otherwise. If a variant that does not exist is requested, then the method will install whatever HuggingFace returns as its default revision.

config is an optional dict of values that will override the autoprobed values for model type, base, scheduler prediction type, and so forth. See Model configuration and probing for details.

access_token is an optional access token for accessing resources that need authentication.

The method will return a ModelInstallJob. This object is discussed at length in the following section.

install_job = installer.import_model()

The import_model() method is the core of the installer. The following illustrates basic usage:

from invokeai.app.services.model_install import (
 LocalModelSource,
 HFModelSource,
 URLModelSource,
)

source1 = LocalModelSource(path='/opt/models/sushi.safetensors')   # a local safetensors file
source2 = LocalModelSource(path='/opt/models/sushi_diffusers')     # a local diffusers folder

source3 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5')  # a repo_id
source4 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5', subfolder='vae')  # a subfolder within a repo_id
source5 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5', variant='fp16')   # a named variant of a HF model
source6 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5', subfolder='OrangeMix/OrangeMix1.ckpt')   # path to an individual model file

source7 = URLModelSource(url='https://civitai.com/api/download/models/63006')       # model located at a URL
source8 = URLModelSource(url='https://civitai.com/api/download/models/63006', access_token='letmein') # with an access token

for source in [source1, source2, source3, source4, source5, source6, source7, source8]:
    installer.import_model(source)

jobs = installer.wait_for_installs(timeout=120)
for job in jobs:
    if job.complete:
        model_config = job.config_out
        model_key = model_config.key
        print(f"{job.source} installed as {model_key}")
    elif job.errored:
        print(f"{job.source}: {job.error_type}: {job.error}\nStack trace:\n{job.error_traceback}")
    elif job.paused:
        print(f"{job.source} paused; resume or retry its download to finish installation")

As shown here, the import_model() method accepts a variety of sources, including local safetensors files, local diffusers folders, HuggingFace repo_ids with and without a subfolder designation, Civitai model URLs and arbitrary URLs that point to checkpoint files (but not to folders).

Each call to import_model() return a ModelInstallJob job, an object which tracks the progress of the install.

If a remote model is requested, the model’s files are downloaded in parallel across a multiple set of threads using the download queue. During the download process, the ModelInstallJob is updated to provide status and progress information. After the files (if any) are downloaded, the remainder of the installation runs in a single serialized background thread. These are the model probing, file copying, and config record database update steps.

Multiple install jobs can be queued up. Call wait_for_installs() to wait for pending download and install work to quiesce. Its returned list contains the tracked jobs and may include a recoverably paused, nonterminal job; resume or retry that job before treating installation as complete. If a job errored out, its exception class, message and traceback are available as job.error_type, job.error and job.error_traceback.

The full list of arguments to import_model() is as follows:

Argument	Type	Default	Description
source	ModelSource	None	The source of the model, Path, URL or repo_id
config	Dict[str, Any]	None	Override all or a portion of model’s probed attributes

The next few sections describe the various types of ModelSource that can be passed to import_model().

config can be used to override all or a portion of the configuration attributes returned by the model prober. See the section below for details.

LocalModelSource

This is used for a model that is located on a locally-accessible Posix filesystem, such as a local disk or networked fileshare.

Argument	Type	Default	Description
path	str	Path	None
inplace	bool	False	If set, the model file(s) will be left in their location; otherwise they will be copied into the InvokeAI root’s models directory

URLModelSource

This is used for a single-file model that is accessible via a URL. The fields are:

Argument	Type	Default	Description
url	AnyHttpUrl	None	The URL for the model file.
access_token	str	None	An access token needed to gain access to this file.

The AnyHttpUrl class can be imported from pydantic.networks.

Ordinarily, no metadata is retrieved from these sources. However, there is special-case code in the installer that looks for HuggingFace and fetches the corresponding model metadata from the corresponding repo.

HFModelSource

HuggingFace has the most complicated ModelSource structure:

Argument	Type	Default	Description
repo_id	str	None	The ID of the desired model.
variant	ModelRepoVariant	ModelRepoVariant(‘fp16’)	The desired variant.
subfolder	Path	None	Look for the model in a subfolder of the repo.
access_token	str	None	An access token needed to gain access to a subscriber’s-only model.

The repo_id is the repository ID, such as stabilityai/sdxl-turbo.

The variant is one of the various diffusers formats that HuggingFace supports and is used to pick out from the hodgepodge of files that in a typical HuggingFace repository the particular components needed for a complete diffusers model. ModelRepoVariant is an enum that can be imported from invokeai.backend.model_manager and has the following values:

Name	String Value
ModelRepoVariant.DEFAULT	“default”
ModelRepoVariant.FP16	“fp16”
ModelRepoVariant.FP32	“fp32”
ModelRepoVariant.ONNX	“onnx”
ModelRepoVariant.OPENVINO	“openvino”
ModelRepoVariant.FLAX	“flax”

You can also pass the string forms to variant directly. Note that InvokeAI may not be able to load and run all variants. At the current time, specifying ModelRepoVariant.DEFAULT will retrieve model files that are unqualified, e.g. pytorch_model.safetensors rather than pytorch_model.fp16.safetensors. These are usually the 32-bit safetensors forms of the model.

If subfolder is specified, then the requested model resides in a subfolder of the main model repository. This is typically used to fetch and install VAEs.

Some models require you to be registered with HuggingFace and logged in. To download these files, you must provide an access_token. Internally, if no access token is provided, then HfFolder.get_token() will be called to fill it in with the cached one.

A Hugging Face install is pinned to the commit the repository is at when it starts. That covers an HFModelSource, a URLModelSource naming a repository page on huggingface.co (repository pages on a configured HF_ENDPOINT mirror are not recognized), and a URLModelSource naming one file by branch, tag or pull request ({endpoint}/{org}/{repo}/resolve/{revision}/{file}, http or https, on huggingface.co or the configured HF_ENDPOINT mirror; /blob/ page URLs on those hosts are converted to /resolve/ first), which is what /install/huggingface passes for a single-checkpoint repository. A repository’s commit comes from its metadata. A single file’s comes from the X-Repo-Commit header of a HEAD request to the URL’s own host. The request goes through the download queue (DownloadQueueServiceBase.head()), so it takes the downloads’ download_proxy and private-address policy. Like huggingface_hub, it follows redirects on the same host until a response names the commit (a renamed repository answers a relative 307, a small file one to its cache, http:// a 301 to https://), and never follows one to another host, such as the CDN serving large files. It is sent with the token the download sends: the source’s own token, else the remote_api_tokens entry matching the URL as given. That token is then sent explicitly with the download, whose URL names the commit instead of the branch. A file URL that already names a commit is used as given. If the host names no commit or refuses the request, the file is downloaded as given, unpinned, and the download reports any failure; a 429, 5xx or connection failure raises ModelMetadataUnavailableError, as a repository lookup does. Runtime downloads through download_and_cache_model() are not installs and never look up a commit. The install’s internal download URLs name the commit, and the install marker records it as revision. Resuming, retrying failed files and restoring after a restart all download from that commit, so a repository whose branch moves while the install is paused cannot have files from two commits mixed into one model. The install’s source, and URLs that users and model records see, such as the checkpoint URLs a repository lookup returns, keep naming the branch. If the pinned commit is no longer available (a repository lookup reports it missing, or a single file’s commit URL answers 404 with X-Error-Code EntryNotFound or RevisionNotFound), its staged files can never be completed: the resume or retry raises ModelRevisionUnavailableError (HTTP 409 from the API), the install ends in ERROR with a message saying to start it again, and its staging directory is removed. Any other answer for a pinned single file (a refusal such as 401 GatedRepo, a 404 RepoNotFound, a bare 404 from a mirror or proxy, an unreachable host) is left to the download, which keeps its staged files. The errored job keeps no download_parts, and resuming or restarting it does nothing; a retry submits the source again, which starts a new install pinned to the repository’s current commit. An install restored at startup errors the same way. Only a full commit hash pins an install: a branch can move after files were staged from it. So when a restored Hugging Face install’s marker names no commit (written before installs were pinned, a null revision, a branch, or an install whose host named no commit), restore deletes its staged files and partials, keeping the staging directory and marker, and the download starts over, pinned to the current commit if the host names one; a new install never reuses such staging. A resume, retry or restart of a Hugging Face install that is not pinned to a commit deletes its staged files the same way and downloads again, pinned if the host now names a commit, so an unpinned install never keeps a partial file. The marker of a Hugging Face file URL records the hf_endpoint it was recognized on, so after HF_ENDPOINT changes the install is still treated as Hugging Face: its unpinned staging is discarded, and its pinned download continues on the endpoint it started on. Downloads that have finished (markers with DOWNLOADS_DONE or RUNNING status) and sources without repository revisions keep their staged files.

Monitoring the install job process

When you create an install job with import_model(), it launches the download and installation process in the background and returns a ModelInstallJob object for monitoring the process.

The ModelInstallJob class has the following structure:

Attribute	Type	Description
id	int	Integer ID for this job
status	InstallStatus	An enum of [waiting, downloading, running, paused, completed, error and cancelled]
config_in	dict	Overriding configuration values provided by the caller
config_out	AnyModelConfig	After successful completion, contains the configuration record written to the database
inplace	boolean	True if the caller asked to install the model in place using its local path
source	ModelSource	The local path, remote URL or repo_id of the model to be installed
local_path	Path	If a remote model, holds the path of the model after it is downloaded; if a local model, same as source
error_type	str	Name of the exception that led to an error status
error	str	Message of the exception that led to an error status
error_traceback	str	Traceback of that exception

If the event_bus argument was provided, the installer also dispatches the following events through it (EventServiceBase.dispatch()), and clients receive them under these names. Every payload also carries a timestamp.

model_install_download_started and model_install_download_progress

For remote models, model_install_download_started is emitted when downloading begins, and model_install_download_progress is emitted as it progresses. Both payloads contain these keys:

Key	Type	Description
id	int	ID of the install job
source	ModelSource	Requested model source
local_path	str	Path where the model is downloading to (usually a temporary directory)
bytes	int	How many bytes downloaded so far
total_bytes	int	Total size of all the files that make up the model
parts	List[Dict]	Information on the progress of the individual files that make up the model

Each item in parts contains url, local_path, bytes and total_bytes.

These download events are not issued for local models. They occur before model_install_started.

model_install_downloads_complete

model_install_downloads_complete is emitted after all required downloads finish. Its payload contains id and source.

model_install_paused

model_install_paused is emitted when a remote install is paused, including when a transient file-download error preserves staging for resume or retry. Its payload contains id and source.

model_install_started

model_install_started is emitted after any download, when model probing, copying and registration begin. Its payload contains id and source.

model_install_complete

model_install_complete is emitted once an installation succeeds. Its payload contains id, source, key, total_bytes and config; key is the ID under which the model has been registered, and total_bytes may be null for a local model.

model_install_error

model_install_error is emitted if the installation process fails. Its payload contains id, source, error_type and error.

model_install_cancelled

model_install_cancelled is emitted when the installation is cancelled. Its payload contains id and source.

Following the model status

You may poll the ModelInstallJob object returned by import_model() to ascertain the state of the install. The job status can be read from the job’s status attribute, an InstallStatus enum which has the enumerated values WAITING, DOWNLOADING, RUNNING, PAUSED, COMPLETED, ERROR and CANCELLED. A paused remote install can be resumed or its failed files retried while its staging data is retained.

For convenience, install jobs also provide the following boolean properties: waiting, downloading, running, paused, complete, errored and cancelled, as well as in_terminal_state. The last will return True if the job is in the complete, errored or cancelled states.

Model configuration and probing

The install service uses the invokeai.backend.model_manager.probe module during import to determine the model’s type, base type, and other configuration parameters. Among other things, it assigns a default name and description for the model based on probed fields.

When downloading remote models is implemented, additional configuration information, such as list of trigger terms, will be retrieved from the HuggingFace and Civitai model repositories.

The probed values can be overridden by providing a dictionary in the optional config argument passed to import_model(). You may provide overriding values for any of the model’s configuration attributes. Here is an example of setting the SchedulerPredictionType and name for an sd-2 model:

install_job = installer.import_model(
               source=HFModelSource(repo_id='stabilityai/stable-diffusion-2-1',variant='fp32'),
      config=dict(
            prediction_type=SchedulerPredictionType('v_prediction')
      name='stable diffusion 2 base model',
            )
       )

Other installer methods

This section describes additional methods provided by the installer class.

jobs = installer.wait_for_installs([timeout])

Wait until pending download and install work has quiesced, then return the tracked jobs. Transient download failures can leave a job in PAUSED; this nonterminal job is returned and does not block the method. Resume or retry it before treating installation as complete. The optional timeout argument raises if work does not quiesce in time; 0 (the default) waits indefinitely.

jobs = installer.wait_for_job(job, [timeout])

Block until a specific job reaches a terminal state or pauses, then return it. While a resume or retry of a paused job is being set up, or the callback of its last download is still finishing, this keeps waiting for that to settle. A paused job is nonterminal and stays paused until it is resumed or its failed files are retried, so check job.paused before treating the install as finished. The optional timeout argument raises if the job neither finishes nor pauses in time; 0 (the default) waits indefinitely.

job = installer.resume_and_wait_for_job(job, [timeout])

For callers that need the model rather than a progress report, such as an invocation that installs a missing dependency. Resume the job if it is paused, then wait for it like wait_for_job(). resume_job() itself never waits: it raises InstallDownloadConflictError at once when another resume is underway or the job’s last download callback has not finished. This method waits out the callback and resumes the job afterwards if it is still paused; when another caller is resuming the job, it waits for that attempt’s outcome instead of starting a second one. A resume refused for any other reason, such as a stopping service or a partial file that must restart but cannot be deleted, is not retried: the job is returned paused. If the job’s pinned revision is gone, the job is returned in ERROR rather than raising. The optional timeout bounds the whole call.

jobs = installer.list_jobs()

Return a list of all tracked ModelInstallJob objects, including queued, paused, and terminal jobs.

jobs = installer.get_job_by_source(source)

Return a list of ModelInstallJob corresponding to the indicated model source.

jobs = installer.get_job_by_id(id)

Return a list of ModelInstallJob corresponding to the indicated model id.

jobs = installer.cancel_job(job)

Cancel the indicated job.

installer.prune_jobs

Remove jobs that are in a terminal state (i.e. complete, errored or cancelled) from the job list returned by list_jobs() and get_job().

Install jobs are in-memory progress records, not durable history. The install API exposes the jobs still retained by the current service through GET /install; DELETE /install removes terminal jobs from that list. An error’s type, traceback, and source are available on the job and error event while they are retained, but pruning or restarting the service makes that diagnostic history unavailable. Recovery markers may contain some install inputs, but they are not a supported job-history interface and do not guarantee the original error or the source-to-destination association.

When an install requires manual recovery, recovery sentinels protect the affected directories and their files across pruning, cleanup, and service restart. The supported recovery contract is artifact preservation: inspect the protected directories on disk and use register_path() after resolving the files. The installer does not restore a terminal recovery job or provide a history/list endpoint for protected roots after restart. Operators should capture the reported error and paths before pruning or restarting if they need those diagnostics later.

installer.app_config, installer.record_store, installer.event_bus

Properties that provide access to the installer’s InvokeAIAppConfig, ModelRecordServiceBase and EventServiceBase objects.

key = installer.register_path(model_path, config), key = installer.install_path(model_path, config)

These methods bypass the download queue and directly register or install the model at the indicated path, returning the unique ID for the installed model.

Both methods accept a Path object corresponding to a checkpoint or diffusers folder, and an optional dict of config attributes to use to override the values derived from model probing.

The difference between register_path() and install_path() is that the former creates a model configuration record without changing the location of the model in the filesystem. The latter moves the model into the InvokeAI models directory before registering it.

When an install transfer cannot be rolled back without risking data loss, the installer raises a recovery-required error and retains both transfer roots for manual recovery. Sibling sentinels keep recovery roots across restart, including when JSON marker writes fail, and prevent cancellation, restart, or orphan cleanup from deleting them. Successfully registering a manually recovered root removes its sentinel. Conversion writes its intermediate model under models/.convert_tmp, which the orphan scan skips while the conversion is active. If installation needs recovery, the conversion endpoint reports and retains that source directory; successful conversions and ordinary failures remove their scratch directory.

Rollback requires an atomic no-replace rename. Linux uses renameat2(RENAME_NOREPLACE), macOS uses renamex_np(RENAME_EXCL), and Windows relies on os.rename() refusing an existing destination. If the primitive is unavailable or the filesystem refuses it, rollback stops and reports the source and managed destination; it never substitutes an existence check followed by an overwriting rename. Any artifact still in the managed destination stays protected by its recovery sentinel for manual recovery, and a concurrently recreated source is left untouched. Cross-filesystem restores stage a copy but still require the same no-replace operation to publish it.

When a restore rename crosses filesystems, its temporary copy is created in a hidden .<source-name>.restore-* directory beside the original source. That parent may be outside the managed models directory. The copy is published with a no-replace rename before the managed copy is removed; a hard termination can therefore leave a partial or complete staged copy, both published copies, or (for directories) a partially removed managed copy. These scratch paths have no durable ownership record, so startup does not reconcile or delete them. Do not remove a restore scratch directory based on its name or age: first stop install activity, inspect the original source parent and the managed destination, and establish which copies are complete. Preserve every possible sole or incomplete copy until recovery is resolved. For sources outside the managed models directory, InvokeAI cannot discover or protect the original source parent; operators must inspect it directly. After resolving the files, use register_path() where appropriate.

Remote staging directories are claimed before downloads begin, and startup cleanup skips directories with a live claim. User-paused downloads keep their claim so orphan cleanup cannot expose or delete resumable staging; shutdown releases the claim only after the download queue has quiesced, preserving completed downloads for restart. Local sources under the managed models directory and orphan deletion also take an exclusive active claim, so probing or transfer cannot race with cleanup; startup removes claims from processes that have exited. Install jobs can be cancelled before file transfer starts. Once files are being moved, cancellation returns a conflict so the transfer can finish or preserve its recovery data. During shutdown, the installer waits for its worker to finish before cleaning queued jobs.

Active-claim sidecars accept the legacy PID process-start-time payload and new Linux claims append a stable linux:<boot-id>:<start-ticks> process identity. A live PID with a legacy timestamp mismatch is ambiguous because wall-clock changes and PID reuse cannot be distinguished; startup preserves that claim for operator recovery. New Linux claims use the boot ID and kernel start ticks, so a reused PID can be reclaimed without comparing wall-clock time. On other platforms, new claims retain the two-field format and a reused live PID is conservatively preserved; a Linux-format sidecar encountered there is unsupported and also preserved. Claim payloads are written and flushed to a temporary file, then published under the claim name without ever replacing an existing claim (invokeai/app/util/no_replace.py): as a hard link; where the filesystem has no hard links (exFAT, FAT32, many network shares), by an atomic no-replace rename; where neither works (macOS exFAT, NFS, many FUSE mounts), by creating the claim exclusively and copying the payload into it. Claim creation fails only when none of these works. A process killed before publication may leave an ignored temporary file. Only the copy fallback can expose a partial claim or leave one behind: after a crash, or when a failed copy cannot confirm that the file under the claim name is still its own and so keeps it. A partial claim still counts as held, and the sweep preserves it as malformed.

Startup serializes stale-claim scans in-process and across processes with the persistent hidden .invokeai_install_claims.lock file in the models root; it remains in place so concurrent scanners always lock the same file. The lock covers enumeration through stale-claim removal, so a waiting scan sees claims created after the preceding scan completes. Windows uses repeated nonblocking lock attempts rather than msvcrt.LK_LOCK’s finite retry count, which could fail during a long scan. An empty or malformed sidecar, denied process inspection, and other inspection failures do not prove that its owner is gone, so startup preserves the claim and emits at most ten individual warnings plus one summary per scan. The next service startup retries the check. If the claim remains, an operator may remove its sidecar manually only after confirming that no install or orphan deletion is using that path; age alone is not evidence that a claim is stale.

Install staging ledger

invokeai/app/services/model_install/install_staging_ledger.py owns everything an install keeps on disk under the models directory. Paths are relative to the models directory:

Path	Meaning
tmpinstall_*/	A remote install’s staging directory. Downloads land here before the install moves them.
tmpinstall_*/.invokeai_install.json	The staging marker, format version 1: version, source, access_token (Hugging Face and URL sources only), config_in, status, updated_at, per-file resume metadata in files (see below), the repository revision its downloads are pinned to (null for a source without one; absent in markers written before installs were pinned), hf_endpoint for a Hugging Face install (the endpoint its URLs were recognized on), and recovery_required: true once a transfer failed in a way that needs manual recovery.
.<name>.invokeai_install_active	An active claim on the sibling <name>, held by a live install or orphan deletion.
.<name>.invokeai_install_recovery_required	A recovery sentinel: <name> holds data a failed transfer could not roll back.
.invokeai_install_claims.lock	Serializes the startup sweep of stale claims.

Sidecars sit beside the root they protect, never inside it, so they do not travel with the files an install moves. The installer gives each install job one InstallStaging handle when it creates the job. The handle covers the job’s staging directory and, for a local source under the models directory, the claim on that source’s top-level root. It claims, persists the marker, marks both roots as recovery data before files move, and on release drops its claims before it clears the recovery marks it no longer needs, so a crash between the two steps leaves data protected rather than claimed. A job whose source already held recovery data keeps every mark until the job completes. A staging tree is only ever removed under a claim, taken first if the install does not hold one; a tree marked as recovery data is never removed, and cleanup checks the tree again once it holds the claim.

The destination root <key>/ is not part of the handle: the transfer code claims it through the ledger before creating it, marks it as recovery data as soon as it exists, and keeps the claim until the model is registered or the transfer has rolled back. A failed rollback retains the mark; a clean rollback removes the empty destination and clears the mark; registering the model clears it.

The orphaned models service asks the ledger one question before it lists or deletes a path: is it protected? A path is protected when it, an ancestor below the models directory, or anything beneath it is claimed or marked as recovery data. Protected directories are neither listed nor deleted. Startup registration of unregistered models skips protected paths the same way. A deletion claims the path’s top-level root for its duration, so an install cannot start using it meanwhile.

Marker writes replace the marker atomically: the new content goes to a temporary file in the staging directory, which is flushed, renamed over the marker, and followed by a sync of the directory on non-Windows platforms. Readers and a crash see the previous marker or the new one. On Windows a rename that a concurrent reader blocks is retried briefly. A temporary file left by a process killed mid-write is ignored and removed with the marker.

At startup the installer sweeps stale claims, then removes staging directories that no install can resume: those without a marker and those whose marker has a terminal status. It keeps claimed directories, directories holding a registered model, and recovery data. A staging directory whose marker exists but is unusable (torn, not JSON, or a format version this release does not know, for example after a downgrade) may hold the only copy of a download, so startup keeps it, logs a warning naming it, and never restores or reuses it. Orphaned model cleanup lists such a directory, unless it is claimed, holds recovery data, or holds a registered model, so it can be deleted once it is not needed.

Each files entry records one download part: its url, canonical_url and final_url, the validators etag and last_modified, expected_total_bytes, the download_path it writes, and resume_required, resume_refused and resume_message. The installer rewrites the marker as each part starts, once its response headers arrived, so a crash leaves the running part’s validator and file name on disk. On a resume or a restore, the download queue’s DownloadJob names the partial file a part would continue (resumable_partial()) and whether it may (may_resume_partial()): only with a recorded validator, sent as If-Range, and only if the server has not already refused to continue it. A server whose file changed answers with the whole file, and the part pauses as restart-required instead of splicing two revisions. Any other partial is deleted and its part reported with resume_from_scratch, as is a part that had started but whose partial and finished file are both gone. A recorded download_path outside the staging directory is ignored. Entries written before resume_refused existed count as refused when resume_required is set and the message is the queue’s refusal message. When a partial that must restart cannot be deleted (on Windows, while another process holds it open), the resume or restart fails with a conflict rather than continue it without a validator; during startup restore the job stays paused.

A failed-file retry (restart_failed()) or explicit file restart (restart_file()) deletes the affected file’s partial before downloading it again. Valid completed sibling files in staging remain reusable; a pinned Hugging Face install keeps downloading from its recorded commit. An unpinned Hugging Face install instead discards all staged files and resume metadata before retrying, then pins the new download if its host supplies a commit.

Startup restores nonterminal install markers automatically, except markers with PAUSED status. If shutdown interrupts restoration of a marker already visited, that job is saved as PAUSED; markers later in the scan remain unchanged. On the next service startup, non-paused markers resume automatically, while both user-paused and shutdown-paused jobs stay paused until resume_job() is called. Partial files and resume metadata are retained across this cycle, except that restoring a Hugging Face install without a full commit hash discards its staged files and resume metadata unless its downloads have finished (marker status DOWNLOADS_DONE or RUNNING).

Before a transfer, the installer flushes the recovery-sentinel file and, on non-Windows platforms, its containing directory. A file or directory fsync error aborts the install before source files move; unsupported directory fsync is treated as a failure, not as best-effort durability. Sentinel cleanup unlinks the file before syncing the containing directory. If that final sync fails, the installer logs a warning and continues cleanup; the unlink may not survive a crash, so a stale sentinel can reappear and require manual inspection. These steps do not promise crash durability on every filesystem or storage device.

installer.unregister(key)

This will remove the model config record for the model at key, and is equivalent to installer.record_store.del_model(key)

installer.delete(key)

This is similar to unregister() but has the additional effect of conditionally deleting the underlying model file(s) if they reside within the InvokeAI models directory

installer.unconditionally_delete(key)

This method is similar to unregister(), but also unconditionally deletes the corresponding model weights file(s), regardless of whether they are inside or outside the InvokeAI models hierarchy.

path = installer.download_and_cache(remote_source, [access_token], [timeout])

This utility routine will download the model file located at source, cache it, and return the path to the cached file. It does not attempt to determine the model type, probe its configuration values, or register it with the models database.

You may provide an access token if the remote source requires authorization. The call will block indefinitely until the file is completely downloaded, cancelled or raises an error of some sort. If you provide a timeout (in seconds), the call will raise a TimeoutError exception if the download hasn’t completed in the specified period.

You may use this mechanism to request any type of file, not just a model. The file will be stored in a subdirectory of INVOKEAI_ROOT/models/.cache. If the requested file is found in the cache, its path will be returned without redownloading it.

Be aware that the models cache is cleared of infrequently-used files and directories at regular intervals when the size of the cache exceeds the value specified in Invoke’s convert_cache configuration variable.

installer.start(invoker)

The start method is called by the API initialization routines when the API starts up. Its effect is to call sync_to_config() to synchronize the model record store database with what’s currently on disk.

---

Get on line: The Download Queue

InvokeAI can download arbitrary files using a multithreaded background download queue. Internally, the download queue is used for installing models located at remote locations. The queue is implemented by the DownloadQueueService defined in invokeai.app.services.download_manager. However, most of the implementation is spread out among several files in invokeai/backend/model_manager/download/*

A default download queue is located in ApiDependencies.invoker.services.download_queue. However, you can create additional instances if you need to isolate your queue from the main one.

A job for every task

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. Jobs come in a variety of shapes and colors as they are progressively specialized for particular download task.

The basic job is the DownloadJobBase, a pydantic object with the following fields:

Field	Type	Default	Description
id	int		Job ID, an integer >= 0
priority	int	10	Job priority. Lower priorities run before higher priorities
source	str		Where to download from (specialized types used in subclasses)
destination	Path		Where to download to
status	DownloadJobStatus	Idle	Job’s status (see below)
event_handlers	List[DownloadEventHandler]		Event handlers (see below)
job_started	float		Timestamp for when the job started running
job_ended	float		Timestamp for when the job completed or errored out
job_sequence	int		A counter that is incremented each time a model is dequeued
error	Exception		A copy of the Exception that caused an error during download

When you create a job, you can assign it a priority. If multiple jobs are queued, the job with the lowest priority runs first. (Don’t blame us! The Unix developers came up with this convention.)

Every job has a source and a destination. source is a string in the base class, but subclassses redefine it more specifically.

The destination must be the Path to a file or directory on the local filesystem. If the Path points to a new or existing file, then the source will be stored under that filename. If the Path ponts to an existing directory, then the downloaded file will be stored inside the directory, usually using the name assigned to it at the remote site in the content-disposition http field.

When the job is submitted, it is assigned a numeric id. The id can then be used to control the job, such as starting, stopping and cancelling its download.

The status field is updated by the queue to indicate where the job is in its lifecycle. Values are defined in the string enum DownloadJobStatus, a symbol available from invokeai.app.services.download_manager. Possible values are:

Value	String Value	Description
IDLE	idle	Job created, but not submitted to the queue
ENQUEUED	enqueued	Job is patiently waiting on the queue
RUNNING	running	Job is running!
PAUSED	paused	Job was paused and can be restarted
COMPLETED	completed	Job has finished its work without an error
ERROR	error	Job encountered an error and will not run again
CANCELLED	cancelled	Job was cancelled and will not run (again)

job_started, job_ended and job_sequence indicate when the job was started (using a python timestamp), when it completed, and the order in which it was taken off the queue. These are mostly used for debugging and performance testing.

In case of an error, the Exception that caused the error will be placed in the error field, and the job’s status will be set to DownloadJobStatus.ERROR.

Download errors can leave a .downloading file behind; cleanup is owned by the caller. When a part of a model install fails with a transient error (a connection failure or timeout, an interrupted transfer including a dropped TLS connection, HTTP 408, 429 or 5xx, or the .downloading file vanishing before its final rename) and the install’s staging directory holds downloaded bytes a resume would keep (a non-empty completed file, or a partial file with a recorded validator its server has not refused, from this attempt or an earlier one; a resume deletes any other partial), the installer pauses the install and preserves its staging directory and marker for resume or restart. The failed part keeps its error_type and error in the job’s download_parts. Any other error (for example HTTP 401, 403 or 404, a TLS certificate that fails verification, insufficient disk space, or an unsafe filename), or a transient one with no recoverable data, errors the install and removes the temporary staging directory.

There are a series of subclasses of DownloadJobBase that provide support for specific types of downloads. These are:

DownloadJobPath

This subclass redefines source to be a filesystem Path. It is used to move a file or directory from the source to the destination paths in the background using a uniform event-based infrastructure.

DownloadJobRemoteSource

This subclass adds the following fields to the job:

Field	Type	Default	Description
bytes	int	0	bytes downloaded so far
total_bytes	int	0	total size to download
access_token	Any	None	an authorization token to present to the remote source

The job will start out with 0/0 in its bytes/total_bytes fields. Once it starts running, total_bytes will be populated from information provided in the HTTP download header (if available), and the number of bytes downloaded so far will be progressively incremented.

DownloadJobURL

This is a subclass of DownloadJobBase. It redefines source to be a Pydantic AnyHttpUrl object, which enforces URL validation checking on the field.

Note that the installer service defines an additional subclass of DownloadJobRemoteSource that accepts HuggingFace repo_ids in addition to URLs. This is discussed later in this document.

Event handlers

While a job is being downloaded, the queue will emit events at periodic intervals. A typical series of events during a successful download session will look like this:

1. enqueued
2. running
3. running
4. running
5. completed

There will be a single enqueued event, followed by one or more running events, and finally one completed, error or cancelled events.

It is possible for a caller to pause download temporarily, in which case the events may look something like this:

1. enqueued
2. running
3. running
4. paused user paused the download
5. running
6. completed

The download queue logs when downloads start and end (unless quiet is set to True at initialization time) but doesn’t log any progress events. You will probably want to be alerted to events during the download job and provide more user feedback. In order to intercept and respond to events you may install a series of one or more event handlers in the job. Whenever the job’s status changes, the chain of event handlers is traversed and executed in the same thread that the download job is running in.

Event handlers have the signature Callable[["DownloadJobBase"], None], i.e.

def handler(job: DownloadJobBase):
   pass

A typical handler will examine job.status and decide if there’s something to be done. This can include cancelling or erroring the job, but more typically is used to report on the job status to the user interface or to perform certain actions on successful completion of the job.

Event handlers can be attached to a job at creation time. In addition, you can create a series of default handlers that are attached to the queue object itself. These handlers will be executed for each job after the job’s own handlers (if any) have run.

During a download, running events are issued every time roughly 1% of the file is transferred. This is to provide just enough granularity to update a tqdm progress bar smoothly.

Handlers can be added to a job after the fact using the job’s add_event_handler method:

job.add_event_handler(my_handler)

All handlers can be cleared using the job’s clear_event_handlers() method. Note that it might be a good idea to pause the job before altering its handlers.

Creating a download queue object

The DownloadQueueService constructor takes the following arguments:

Argument	Type	Default	Description
event_handlers	List[DownloadEventHandler]	[]	Event handlers
max_parallel_dl	int	5	Maximum number of simultaneous downloads allowed
requests_session	requests.sessions.Session	None	An alternative requests Session object to use for the download
quiet	bool	False	Do work quietly without issuing log messages

A typical initialization sequence will look like:

from invokeai.app.services.download_manager import DownloadQueueService

def log_download_event(job: DownloadJobBase):
 logger.info(f'job={job.id}: status={job.status}')

queue = DownloadQueueService(
       event_handlers=[log_download_event]
        )

Event handlers can be provided to the queue at initialization time as shown in the example. These will be automatically appended to the handler list for any job that is submitted to this queue.

max_parallel_dl sets the number of simultaneous active downloads that are allowed. The default of five has not been benchmarked in any way, but seems to give acceptable performance.

requests_session can be used to provide a requests module Session object that will be used to stream remote URLs to disk. This facility is intended for unit tests that simulate a remote web server. A test Session without Invoke’s socket-level SSRF guard must also set requests_session_is_trusted=True.

quiet will prevent the queue from issuing any log messages at the INFO or higher levels.

Submitting a download job

You can submit a download job to the queue either by creating the job manually and passing it to the queue’s submit_download_job() method, or using the create_download_job() method, which will do the same thing on your behalf.

To use the former method, follow this example:

job = DownloadJobRemoteSource(
   source='http://www.civitai.com/models/13456',
   destination='/tmp/models/',
   event_handlers=[my_handler1, my_handler2], # if desired
)
queue.submit_download_job(job, start=True)

submit_download_job() takes just two arguments: the job to submit, and a flag indicating whether to immediately start the job (defaulting to True). If you choose not to start the job immediately, you can start it later by calling the queue’s start_job() or start_all_jobs() methods, which are described later.

To have the queue create the job for you, follow this example instead:

job = queue.create_download_job(
   source='http://www.civitai.com/models/13456',
   destdir='/tmp/models/',
   filename='my_model.safetensors',
   event_handlers=[my_handler1, my_handler2], # if desired
   start=True,
 )

The filename argument forces the downloader to use the specified name for the file rather than the name provided by the remote source, and is equivalent to manually specifying a destination of `/tmp/models/my_model.safetensors’ in the submitted job.

Here is the full list of arguments that can be provided to create_download_job():

Argument	Type	Default	Description
source	Union[str, Path, AnyHttpUrl]		Download remote or local source
destdir	Path		Destination directory for downloaded file
filename	Path	None	Filename for downloaded file
start	bool	True	Enqueue the job immediately
priority	int	10	Starting priority for this job
access_token	str	None	Authorization token for this resource
event_handlers	List[DownloadEventHandler]	[]	Event handlers for this job

Internally, create_download_job() has a little bit of internal logic that looks at the type of the source and selects the right subclass of DownloadJobBase to create and enqueue.

TODO: move this logic into its own method for overriding in subclasses.

Job control

Prior to completion, jobs can be controlled with a series of queue method calls. Do not attempt to modify jobs by directly writing to their fields, as this is likely to lead to unexpected results.

Any method that accepts a job argument may raise an UnknownJobIDException if the job has not yet been submitted to the queue or was not created by this queue.

queue.join()

This method will block until all the active jobs in the queue have reached a terminal state (completed, errored or cancelled).

queue.wait_for_job(job, [timeout])

This method will block until the indicated job has reached a terminal state (completed, errored or cancelled). If the optional timeout is provided, the call will block for at most timeout seconds, and raise a TimeoutError otherwise.

jobs = queue.list_jobs()

This will return a list of all jobs, including ones that have not yet been enqueued and those that have completed or errored out.

job = queue.id_to_job(int)

This method allows you to recover a submitted job using its ID.

queue.prune_jobs()

Remove completed and errored jobs from the job list.

queue.start_job(job)

If the job was submitted with start=False, then it can be started using this method.

queue.pause_job(job)

This will temporarily pause the job, if possible. It can later be restarted and pick up where it left off using queue.start_job().

queue.cancel_job(job)

This will cancel the job if possible and clean up temporary files and other resources that it might have been using.

queue.start_all_jobs(), queue.pause_all_jobs(), queue.cancel_all_jobs()

This will start/pause/cancel all jobs that have been submitted to the queue and have not yet reached a terminal state.

---

This Meta be Good: Model Metadata Storage

The modules found under invokeai.backend.model_manager.metadata provide a straightforward API for fetching model metadatda from online repositories. Currently only HuggingFace is supported. However, the modules are easily extended for additional repos, provided that they have defined APIs for metadata access.

Metadata comprises any descriptive information that is not essential for getting the model to run. For example “author” is metadata, while “type”, “base” and “format” are not. The latter fields are part of the model’s config, as defined in invokeai.backend.model_manager.config.

Example Usage

from invokeai.backend.model_manager.metadata import (
   AnyModelRepoMetadata,
)
# to access the initialized sql database
from invokeai.app.api.dependencies import ApiDependencies

hf = HuggingFaceMetadataFetch()

# fetch the metadata
model_metadata = hf.from_id("<repo_id>")

assert isinstance(model_metadata, HuggingFaceMetadata)

Structure of the Metadata objects

There is a short class hierarchy of Metadata objects, all of which descend from the Pydantic BaseModel.

ModelMetadataBase

This is the common base class for metadata:

Field Name	Type	Description
name	str	Repository’s name for the model
author	str	Model’s author
tags	Set[str]	Model tags

Note that the model config record also has a name field. It is intended that the config record version be locally customizable, while the metadata version is read-only. However, enforcing this is expected to be part of the business logic.

Descendents of the base add additional fields.

HuggingFaceMetadata

This descends from ModelMetadataBase and adds the following fields:

Field Name	Type	Description
type	Literal[“huggingface”]	Used for the discriminated union of metadata classes
id	str	HuggingFace repo_id
tag_dict	Dict[str, Any]	A dictionary of tag/value pairs provided in addition to tags
last_modified	datetime	Date of last commit of this model to the repo
files	List[Path]	List of the files in the model repo

AnyModelRepoMetadata

This is a discriminated Union of HuggingFaceMetadata.

Fetching Metadata from Online Repos

The HuggingFaceMetadataFetch class will retrieve metadata from its corresponding repository and return AnyModelRepoMetadata objects. Their base class ModelMetadataFetchBase is an abstract class that defines two methods: from_url() and from_id(). The former accepts the type of model URLs that the user will try to cut and paste into the model import form. The latter accepts a string ID in the format recognized by the repository of choice. Both methods return an AnyModelRepoMetadata, and both accept an optional revision and pin_to_commit. By default, HuggingFaceMetadata file URLs name the branch (the variant’s, or main). pin_to_commit=True makes them name the commit that branch points to now; the installer uses it for its downloads. With revision, the files and URLs come from that revision, and a missing revision raises ModelRevisionUnavailableError. The returned metadata’s revision property holds the revision its URLs name; it is not serialized, so metadata read back from JSON takes it from its file URLs.

The base class also has a class method from_json() which will take the JSON representation of a ModelMetadata object, validate it, and return the corresponding AnyModelRepoMetadata object.

When initializing one of the metadata fetching classes, you may provide a requests.Session argument. This allows you to customize the low-level HTTP fetch requests and is used, for instance, in the testing suite to avoid hitting the internet.

Metadata HTTP calls pass a 15-second timeout to each request. This is a per-request transport timeout, not a total deadline for metadata lookup or service shutdown. Looking up a requested variant may make a second, sequential model-info request at the default revision when that variant is missing; filtering a model that needs model_index.json can make another sequential request. InvokeAI’s current metadata path does not add an application-level retry loop. A custom session or HTTP client may have its own retries and backoff, which can extend elapsed time; the request timeout does not bound those attempts as one total duration.

The installer checks for shutdown between restore markers and before it submits a prepared download, but it cannot interrupt a metadata request already in progress and does not pass a cancellation signal between the sequential metadata calls. stop() waits for the restore worker to finish before clearing jobs or releasing its staging claims. Therefore shutdown latency has no strict maximum: it includes the remaining metadata calls and any retries or waits performed by the configured transport. If shutdown arrives during metadata lookup, that lookup may continue through its fallback or model-index request. The current restore is then saved as PAUSED, and its partial files and resume metadata are preserved; later unvisited markers remain unchanged.

The HuggingFace fetcher subclass add additional repo-specific fetching methods:

HuggingFaceMetadataFetch

This overrides its base class from_json() method to return a HuggingFaceMetadata object directly.

Metadata Storage

The ModelConfigBase stores this response in the source_api_response field as a JSON blob.

---

The Lowdown on the ModelLoadService

The ModelLoadService is responsible for loading a named model into memory so that it can be used for inference. Despite the fact that it does a lot under the covers, it is very straightforward to use.

An application-wide model loader is created at API initialization time and stored in ApiDependencies.invoker.services.model_loader. However, you can create alternative instances if you wish.

Creating a ModelLoadService object

The class is defined in invokeai.app.services.model_load. It is initialized with an InvokeAIAppConfig object, from which it gets configuration information such as the user’s desired GPU and precision, and with a previously-created ModelRecordServiceBase object, from which it loads the requested model’s configuration information.

Here is a typical initialization pattern:

from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.app.services.model_load import ModelLoadService, ModelLoaderRegistry

config = InvokeAIAppConfig.get_config()

ram_cache = ModelCache(
  max_cache_size=config.ram_cache_size, max_vram_cache_size=config.vram_cache_size, logger=logger
)

convert_cache = ModelConvertCache(
  cache_path=config.models_convert_cache_path, max_size=config.convert_cache_size
)

loader = ModelLoadService(
  app_config=config,
  ram_cache=ram_cache,
  convert_cache=convert_cache,
  registry=ModelLoaderRegistry
)

load_model(model_config, [submodel_type], [context]) -> LoadedModel

The load_model() method takes an AnyModelConfig returned by ModelRecordService.get_model() and returns the corresponding loaded model. It loads the model into memory, gets the model ready for use, and returns a LoadedModel object.

The optional second argument, subtype is a SubModelType string enum, such as “vae”. It is mandatory when used with a main model, and is used to select which part of the main model to load.

The optional third argument, context can be provided by an invocation to trigger model load event reporting. See below for details.

If the model has to be built (it is not in the calling device’s cache) and its record was edited in a way that changes how it loads (for example its cpu_only or fp8_storage setting) after the caller read model_config, load_model() does not cache a model built from the superseded record: it reads the record again and loads that instead. The same applies when the model’s files are missing because it was moved and its record updated to the new path. If the record changes again during that retry, it raises StaleModelConfigError. Code that writes a model record in a way that may change how it loads must do so through ModelLoadService.apply_record_change(key, write): it brackets the write so overlapping loads are re-checked, compares the record before and after the same way the load check does, and when they differ drops the model’s entries from each device cache under MODEL_LOAD_LOCK. A write that raises an Exception may have committed first, so the record as stored is read back and judged the same way before the exception propagates; if it cannot be read back, the change counts as load-affecting. An interrupt such as KeyboardInterrupt propagates without the read-back or any drop; loads that overlapped the write are still re-checked. ModelManagerService.update_model_record() and ModelManagerService.reidentify_model() write records this way.

The returned LoadedModel object contains a copy of the configuration record returned by the model recordget_model() method, as well as the in-memory loaded model:

Attribute Name	Type	Description
config	AnyModelConfig	A copy of the model’s configuration record for retrieving base type, etc.
model	AnyModel	The instantiated model (details below)

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

LoadedModel acts as a context manager. The context loads the model into the execution device (e.g. VRAM on CUDA systems), locks the model in the execution device for the duration of the context, and returns the model. Use it like this:

loaded_model_= loader.get_model_by_key('f13dd932c0c35c22dcb8d6cda4203764', SubModelType('vae'))
with loaded_model 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.

An operation that needs more device memory than device_working_mem_gb reserves it when it locks the model, so the cache offloads other models instead of running out of memory. Pass the operation’s working-memory estimate:

estimate = estimate_vae_working_memory_flux2(
    operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device
)
with vae_info.model_on_device(working_memory=estimate) as (_, vae):
    image = vae.decode(latents)[0]

An estimator that declares an EstimatorSpec returns a WorkingMemoryEstimate (invokeai/backend/working_memory/estimate.py): the reservation, the terms it adds up to, the features it was evaluated from, and the source of its coefficients, which model_on_device() logs at debug level. A plain byte count goes in as working_mem_bytes=<int> instead; never both. The cache never reserves less than device_working_mem_gb.

To write such an estimator (invokeai/backend/util/vae_working_memory.py has five for the VAEs, invokeai/backend/z_image/denoise_working_memory.py, invokeai/backend/flux2/denoise_working_memory.py, invokeai/backend/flux3/denoise_working_memory.py, invokeai/backend/krea2/denoise_working_memory.py and invokeai/backend/moss_image/working_memory.py one each for a transformer forward, the Z-Image, the FLUX.2, the FLUX.3, the Krea-2 and the MOSS-Image denoise’s):

• Register an EstimatorSpec at module import, and list the module in ESTIMATOR_MODULES (working_memory/spec.py); the data test fails on a module that declares a spec without being listed.
• Split the estimator in two. The impure part reads the model and the device and records every runtime decision (tile, the dtype the operation runs in, score-matrix bytes, a quantized model’s dequantization transient) as a feature, and refuses a dtype the spec has no line for with a ValueError naming the estimator, dtype and platform. The pure terms(features, values) turns features and coefficient values into terms, building fitted terms with WorkingMemoryTerm.fitted(...). The public function returns SPEC.estimate(features, device).
• Record the dtype the operation will run in, not the one the cached model holds. The SD1/SDXL nodes cast the shared VAE with sd_vae_run_dtype: fp16 for every GPU run that does not ask for fp32 (bf16 never runs), fp32 when asked and for a CPU decode. Their estimate prices the same function’s answer, plus what the in-place cast adds to the resident weights. Uncast VAEs compute in their first floating-point parameter wider than a byte, so fp8 or packed SDNQ storage does not count.
• Declare the line the fitted terms price on with line= (a function returning a fit.Line): the suffix that names the line’s coefficients, the operation’s size on its x-axis, whether a measurement of these features fits the line, and how point ids label them. Every estimator declares it and nothing defaults it, so a new estimator cannot inherit another’s line, and with it which of its measurements fit. The VAE estimators declare the AutoencoderKL convolution line in output pixels (fit.convolution_line), the Z-Image, FLUX.2, FLUX.3 and Krea-2 denoise ones in attended tokens. A line may sit on another (Line.sits_on, the line beneath and the operation’s size on it): a tiled VAE operation prices its tile, clamped to the image per axis, on the untiled line and sits there, and its estimator’s tiled line tiled_<operation>_<dtype>, in the image’s pixels, prices what tiling adds on top. A line may also have a second phase that peaks apart from it (Line.assembly): a tiled decode assembles the image from its decoded tiles after the last one, priced on the assembly line assembled_decode_<dtype> in the image’s pixels, and the estimate reserves the larger phase. The fit subtracts nothing else from a point’s demand, so a line admits only measurements whose demand is the line’s alone, or the line’s above the line it sits on: an operation with a term of its own on top of the line (an in-place upcast, a batch, a ControlNet’s hints, regional prompting’s bias, LoRAs, a dequantization or fp8 cast transient, a materialized score matrix the estimate adds rather than takes the larger of) bounds its estimate but fits nothing.
• Declare legacy_formula= where the terms are shaped unlike the hand-set formula they replace, so that a legacy entry’s constants alone do not reproduce it (the Krea-2 denoise’s CFG and style multipliers, its count-based LoRA term). Where the values are legacy-priced (ResolvedCoefficients.legacy_priced: a legacy entry, or an unmeasured platform raised to an entry’s unmeasured_floor), the estimate adds what the formula gives beyond its terms, as a term legacy_formula; fitted values on a measured platform never read it.
• Opt in with user_calibrated=True where a user’s calibration may move the estimator: the five VAE estimators do. An estimator that does not is shipped-fit only; a calibration run still records its windows. The calibration summary and plan are built on pixel lines (they key sizes by width and height), so an estimator on another line, such as the Z-Image denoise’s in tokens, must not opt in until they key sizes by Line.size.
• Bump version whenever terms changes the formula.
• Key the spec by the platform traits its fitted terms depend on (platform_keys; the VAE estimators use conv_backend and allocator), and ship one coefficient entry per measured platform in invokeai/backend/working_memory/data/, with its own values, with values_from another estimator’s entry, which may carry more values than the borrower reads, or with both: the borrower’s own lines next to the borrowed ones, never a value it also borrows. vae.sd, vae.flux and vae.flux2 borrow the family’s untiled lines from vae.sd3 and carry their own tiled and assembly lines, which differ per family; on MIOpen with expandable segments vae.flux2 carries its own untiled lines too, as its fp32 decode differs from the family’s by more than the margin.
  • A platform is the convolution backend (cudnn, miopen, mps, xpu, cpu), torch’s caching allocator and the OS. The allocator is torch’s backend first (cudaMallocAsync ignores expandable segments), then expandable if the allocator environment variables enable expandable segments on the native allocator, else native; anything but a CUDA/HIP GPU is native. platform_for(device) reads them from this process; platform_for_environment(backend, hard) reads the same facts from a measurement’s recorded environment, through the same function.
  • MEASURED in working_memory/platform.py lists the measured backend, allocator and OS combinations (an OS of None for every OS). priced_as prices every other pair as one of them, keeping the OS: a backend nothing is measured on as cuDNN with the native allocator, an allocator not measured with its backend as that backend’s native allocator. There the entry that carries the values raises them to its unmeasured_floor: the pre-dtype slopes for the untiled VAE lines, for the tiled and assembly lines each value’s largest over the measured platforms (at every size of at least zero a line of per-value maxima lies on or above every measured platform’s line, whatever the signs, a negative intercept included). That is conservative by design until those platforms are measured. Its cost, over the shipped tiled points’ features against the most demanding measured platform’s estimate for them: with the default device_working_mem_gb of 3 GB a reservation below 3 GB changes nothing; a tiled decode on CPU, XPU or MPS reserves up to 1.18x (a FLUX.1 decode at a 512 px tile, 12288 px), a tiled encode there up to 1.29x (an SD fp32 encode at a 1024 px tile, 12288 px); on Linux ROCm decodes up to 1.71x and encodes 2.2-2.5x from 1536 px, mostly the untiled stand-ins at the tile. The denoise lines floor at their legacy constants. A fit describes the card it was measured on, not the platforms it stands in for.
  • An entry applies to the OS only when its platform names one, and then takes precedence there over the one that does not. Only MIOpen names one. With the native allocator, Windows demand is concave and up to 1.6x Linux’s, so Windows has its own fit and every other OS keeps the Linux one. Expandable segments are measured on Windows only, where they are Invoke’s ROCm default, so MEASURED lists them for Windows and every other OS prices them as the native allocator there. cuDNN is priced the same on every OS.
  • A spec whose fitted terms no platform trait moves declares platform_keys=() and ships one entry with an empty platform, read alike on every platform, measured or not, never priced as another or floored; its estimate’s source names it any platform, and its points name no platform. MOSS-Image’s denoise estimator does this: its fitted term is GEMM and fused-attention activations, not convolution workspaces, and a backend without a fused SDPA kernel is priced by the analytic score-matrix term.
• Ship the measured points next to them. A point’s method says how it was read: one fresh process per point for reserved_peak (the allocator’s peak reserved bytes), driver_peak (the driver-level device usage), and demand_peak (the largest of driver-level, reserved and allocated peak, over two runs, which catches the memory HIP maps beyond the allocator on a process’s first operation); device_peak, the drop in the driver’s free memory (torch.cuda.mem_get_info, polled during the operation) below its value just before it, which also counts what is allocated around the allocator’s reserve; or probe_demand, the probe’s demand_bytes, the largest over the valid, fit-eligible windows of calibration runs in a running server. scripts/working_memory/build_shipped.py builds probe_demand points from calibration store exports, takes fresh-process points from the point files a calibration script writes (--points), and refits the values (see the Family E2E guide). A point names the OS only where it measured an OS-specific entry. working_memory/bounds.py holds the rule the data test (tests/backend/working_memory/) and the builder apply: each estimate is at least its point, and where a fitted term decides an in-fit point, at most 1.3 times the point, its ratio in bounds.UPPER_BOUND_EXCEPTIONS (points no line can hold within 1.3x), or the point’s tighter max_ratio. A probe_demand point is a lower bound only (MeasuredPoint.holds_upper_bound) and takes no max_ratio: a warm window in a running server reads below the cold first run after a server start, the worst case a reservation must cover. It may raise a line; only the fresh-process points cap one. The device_working_mem_gb floor counts only where an analytic term decides.
• Fit the values, never hand-tune them. Fitted terms price a line per operation and dtype: k_<suffix> bytes per unit of the line’s size plus a base of b_<suffix> bytes. The AutoencoderKL estimators price their convolutions on lines in output pixels (k_<operation>_<dtype>), a tiled operation’s tile on the same lines and what tiling adds on its estimator’s tiled lines in image pixels (k_tiled_<operation>_<dtype>), a tiled decode’s assembly on its assembly lines in image pixels (k_assembled_decode_<dtype>), the Z-Image and Krea-2 denoise their forwards on lines in attended tokens (k_denoise_<dtype>), the FLUX.2 and FLUX.3 denoise their forwards on one such line per dtype and transformer width (k_denoise_<dtype>_h<width>).
  • working_memory/fit.py fits each line on the in-fit points its estimators’ line admits, of every estimator whose entry carries the line on its platform (its own entry, or the one it borrows the line from): for the convolution lines the single-sample points, leaving out those whose VAE an in-place cast grew (upcast_weight_bytes, a term of its own). The fit of the points that hold an upper bound (the fresh-process ones) is the rising edge of their upper convex hull (with the origin, so the base is never negative), or the flat line through the hull’s peak where it falls past it, that over-reserves least, rated on the highest reading per size of the capping points, the steeper on a tie and never a falling one. A probe_demand point is covered by the line and can raise it, but takes no part in that fit. A fresh VAE point below 512 px a side (fit.SMALLEST_APP_SIDE), on a tile or untiled, shapes that fit but rates no edge, and it holds no upper bound (Line.caps): a tile that small is set by hand in the workflow editor only (the nodes’ default resolves to 512 px for FLUX and 1024 px for SD, and the upscale tab sets 512 to 1536 px), and every estimate for an untiled image up to 384 px a side lies below the 3 GiB default device_working_mem_gb the model cache reserves at least. It needs at least two sizes twice apart. The shipped values add a 5% margin and round up: the slope to a whole byte per unit of size, the base to a whole MiB.
  • A line that sits on another covers its points’ demand above the untiled demand their own estimator measured at their size on the line beneath, on the same platform, operation and dtype (the highest of its in-fit, fresh-process points there: the untiled references a tiling run measures at each tile’s side). Only where no such measurement exists does the fitted price of the line beneath stand in, so the fit runs in dependency order: the lines that sit on none first, then each line that sits on one (fit.base_values, fit.fitted_values; the data test and the builder both fit through fitted_values). The reason is the tiles nobody measured. The shipped untiled line lies above the measured demand by a margin that differs per tile size; covered above that line, a tiled line would absorb the margin at the measured tiles and miss it at others (fitted so on the 512 and 1024 px tiles alone, it priced a 768 px tile at 0.99x on the RX 9060 XT with expandable segments and a 256 px tile at 0.74-0.89x on the RTX 4090). Covered above the measured demand, it is what tiling adds to the untiled demand at the tile, which holds for any tile whose growth is no larger than the measured tiles’. The estimate still prices the tile on the shipped untiled line, which covers the measured demand, and a line that sits on another is rated on that whole estimate, so its points hold the 1.3x bound like any other; a line on which no point’s demand lies above its points’ untiled demand is zero. The tiled VAE lines are fitted so on fresh-process points of scripts/calibrate_vae_tiling_working_memory.py, each the operation as its node runs it inside its model_on_device window (on the RTX 4090 with both allocators and the RX 9060 XT with both on Windows: image sides 1024 to 12288 px with 512 and 1024 px tiles, 2048 to 12288 px with a 1536 px tile on the RTX 4090, and 2048 to 8192 px with 256 and 768 px tiles, the 256 px ones bounding their estimates from below only; and past 12288 px: FLUX.1 and FLUX.2 in bf16 at a 512 px tile at 16384 to 24576 px under cudaMallocAsync and with expandable segments on the RX 9060 XT, and FLUX.2 at 26624 and 28672 px with the RTX 4090’s native allocator; FLUX.1 and FLUX.2 at a 1024 px tile at 16384 and 20480 px with the RTX 4090’s native allocator and with expandable segments on the RX 9060 XT; under cudaMallocAsync SD at a 1536 px tile at 14336 px and FLUX.1 at 16384 px; and with MIOpen’s native allocator SD at a 768 px tile at 12288 px, its fp32 encode also at 16384 px). Past the largest measured side the lines are extrapolated: at the tiles the app sets, each estimate stays above the extension of its last measured segment to at least 15,800 px (FLUX.2’s encode at a 512 px tile with MIOpen’s native allocator, measured to 12288 px). Where the floor decides the largest measured points, as for the FLUX.1 and FLUX.2 decodes at a 512 px tile, the estimate stays above that extension to 65,536 px at least. The upscale tab scales by up to 16x (UPSCALE_SCALE_MAX), so outputs past the covered range are reachable from the UI. A tiled encode needs more than its tile does, by up to about 42 bytes per image pixel at 4096 px, 29 at 8192 px and 20 at 12288 px, for its input image and moments. A tiled decode demands about what the untiled decode of one tile does, flat in the image, but for a jump some allocators take once the image spans several tiles, up to about 4096 px; above that the assembly takes over (every decoded tile, the blended rows and the concatenated image, about 3.8 image copies, and the node’s post-processing), up to 3x the tile’s own demand at 8192 px and 7x at 12288 px with a 512 px tile. One line on top of the tile cannot price both shapes, so a tiled decode reserves max(untiled(tile) + T(P), A(P)), its tiled line T and its assembly line A in the image’s pixels P. fit.two_phase_envelope fits the two together on the tiled decodes, by a deterministic search: each candidate A is the zero line or a line through two points’ readings with a slope of at least zero (its intercept may be negative), T the envelope of the points above it, and the candidate whose estimates over-reserve the capping points least as both lines ship (the largest ratio, then the sum) wins, the zero line on a tie, so no assembly line ships that no point needs. Both take the 5% margin, which scales the whole line, a negative intercept included; the search rates them with it, because it scales a point’s whole demand where A prices it and only what tiling adds where T does. An entry whose decodes need no assembly phase ships it at zero. However far the lines extrapolate, a tiled operation reserves at least its floor, a phase of its own, for every sample (vae_working_memory._tiled_convolutions), plus the 5% the shipped lines carry: what is counted varies across runs and cards as a line’s points do, so the floor gets the same headroom, as a term of its own in the estimate (floor_margin). A decode’s floor is the buffers of its assembly (every decoded tile as diffusers slices it, the edge tiles clipped, the rows blended from them and the concatenated image, three channels of the run dtype each), the latents it decodes, the one image its node post-processes inside its window (l2i, z_image_l2i; flux2_vae_decode and flux_vae_decode post-process after it), and the untiled line’s base. An encode’s is its input image, every encoded tile of moments (twice the latent channels), the rows blended from them and the concatenated moments, the distribution its node makes of them (three latent-wide tensors, six where it samples), and the base. The extra image and the distribution are what the measured allocators took beside the buffers, not what is live at once: the assembly’s tiles and rows are freed by then, and l2i and z_image_l2i hold the decoded image and two more at their post-processing’s peak. Under cudaMallocAsync the FLUX.1 decodes at 16384 to 24576 px took the counted floor and base to within 0.2%, and with expandable segments on Windows ROCm less; on the native allocators the fitted lines decide every such decode measured, to 12288 px at a 512 px tile and to 20480 px at a 1024 px tile, and past them one that keeps the freed blocks could take up to two images more than counted. The floor is counted at the 0.25 overlap every shipped VAE tiles at, from the recorded features alone, so the estimators refuse a tiled operation at another overlap. The fit leaves a point to the floor where the counted floor, without the margin, covers it, so that it reserves the point with the margin as a line would: the point lifts no line and passes no assembly line through it, and it rates a candidate only where the candidate prices it above its floor, since below it the point reserves its floor whichever candidate wins (fit.FitPoint.floor); bounds.py holds it to its upper bound against what is reserved. On the RTX 4090 under cudaMallocAsync and the RX 9060 XT with expandable segments the floor decides the largest measured operations at a 512 px tile from about 5,000 px a side (SD fp16 with expandable segments) to 12,000 px (FLUX.2 under cudaMallocAsync), later at larger tiles, and the measured points it decides took 0.75 to 0.95 of what it reserves. Images that large are reachable: the upscale tab scales by up to 16x, and an image made in the graph reaches the encode nodes without PIL’s decompression-bomb check, which only Image.open applies. Every VAE estimator (SD1/SDXL, FLUX.1 and FLUX.2) prices the batch its tensor carries: the convolution terms per sample, the assembly line’s base once where it is negative (k * P + max(b, n * b)), and the mid-block score matrix over batch times heads. The decode nodes that return one image (l2i, flux_vae_decode, z_image_l2i) decode and price only the first sample of their latents, so a multi-sample latent costs what one sample does.
  • One tiled line serves every tile of an estimator, operation and dtype, so the largest jump any tile takes above its untiled demand sets the line’s base for all tiles, and the fastest growth per pixel its slope. On the default allocators, cuDNN’s native one and expandable segments on MIOpen on Windows, the tiled estimates stay within 1.3x of their points but for 24 of 353, listed in UPPER_BOUND_EXCEPTIONS (up to 1.56x). FLUX.2’s decode lines on cuDNN with the native allocator have no headroom left: its 768 px tile at 8192 px sits at 1.41x, so a FLUX.1 or FLUX.2 bf16 decode probe there that raises the untiled line, from about 4.63 GB at 1024 px, breaks that bound, and the builder refuses the raise until the exceptions are adjusted. Under cudaMallocAsync and MIOpen’s native allocator 30 of 194 and 53 of 149 points sit more than 1.3x below their estimates, up to 2.38x (FLUX.2’s decode at a 512 px tile, 1024 px) and 1.72x (SD fp16’s decode at a 512 px tile, 1024 px; on MIOpen the untiled line’s headroom at the tile adds to it). A tiled line per tile would remove that; over-reserving is preferred to the extra parameter.
  • A point whose demand is not its line’s alone sets in_fit: false: it fits nothing and gets only the lower bound. That is a point measured in a regime its platform does not run (CUDA with the math attention kernel forced), and a fit-grid forward whose features its estimator’s line does not admit, such as one whose build materializes the score matrix. The test lists every such point.
  • A line is fitted only once its capping points alone (in-fit, fresh-process and at a size and tile the app sets, Line.caps) span two sizes twice apart (fit.capping_spans); fitted_values refuses any other line. For a line that sits on none the shipped pair is the fresh-process fit where it covers every other point of the line by the 5% margin. Where it does not, fit.raised_pair raises it by the least the upper bounds allow: of the lines at or above it, per coefficient, that cover every point by 5% and keep the upper bound of every point priced on the line (bounds.line_bounds, in the line’s units: batched and upcast VAE operations and denoise forwards with an analytic term included, though their line admits them to no fit, read where the line decides their estimate), the one whose reservation averaged over the capping sizes is smallest, the steeper on a tie. The candidates are the lines that ship, a whole byte of slope per unit of size and a whole MiB of base, checked against the conditions exactly, with no allowance. Each condition is a half-plane in slope and base, so the lines that keep them all form a convex region: the search scans each whole MiB of base across that region, taking the lowest whole slope the conditions leave at it, which finds the lowest rounded line or proves there is none (some milliseconds per line on the shipped data). Where no rounded line satisfies them all, also where only a line between whole bytes or MiB would, it refuses with fit.LineConflict, naming the probe and the bounds that conflict (each bound by its point’s id, measured bytes and the most it allows, since two points can share an id), or saying that the floor itself already breaks a bound. The data then conflicts with the bounds: build_shipped.py prints the changes it would make, writes nothing and exits 1, as for a result the bounds reject; it exits 2 only for input it cannot use. A probe_demand point thus lowers neither the slope nor the base, and since probe sizes do not enter the average, a probe the shipped line already covers leaves it where it is. A tiled operation bounds neither its tiled line nor the untiled line beneath: the tiled line is fitted after the untiled one, on what that one prices the tile at. A tiled point with an untiled reference measured at its tile is covered above that reference (fit.point_reading), so a raise beneath adds to its estimate byte for byte; only points without one, covered above the line beneath, can hand the tile’s cost from the tiled line to it. The data test and the builder check the result against every tiled bound, and the builder refuses (exit 1) a raise that breaks one. That is a known limitation, not the intent: fit.raised_pair picks the least raise without the tiled bounds, so it can refuse a probe for which a steeper raise would keep them all. On cuDNN with the native allocator, SD fp16 decode probes at 1024 px of 4.9 and 5.0 GB raise the untiled line, which puts the decodes at a 512 px tile at 1536 and 3072 px at 1.30-1.32x, over their bound; the earlier rule, the larger per coefficient of the fit on all points and on the fresh-process ones, raised the line otherwise and accepted them. A joint fit of the untiled and tiled lines would remove the limitation. A line that sits on another reads its points above offsets the line beneath sets, which no half-plane in its own slope and base states, so it is not raised so: its pair is the larger, per coefficient, of the fit on all its points and the fit on its fresh-process points alone, both rated on the capping points (fit.sitting_pair), so a probe_demand point that drops a fresh hull vertex lowers neither the slope nor the base there either. A tiled decode line and its assembly line take that per line and coefficient of the two two-phase fits (fit.fitted_two_phase).
  • Until then the line is a declared stand-in, listed in the entry’s stand_ins; an untiled one is never below the pre-dtype slope. probe_demand points on it bound it from below only: build_shipped.py keeps a line that sits on none where it covers them by 5%, otherwise raises it as fit.raised_pair raises a fitted line, the stand-in as the floor and the average over the sizes of its capping points or, without any, of all its points; a tiled line it raises to their fit where that lies above, read above its points’ offsets, with its assembly line. It notes the raise in the provenance and never lowers the line. MIOpen with the native allocator off Windows prices fp16 and fp32 at the pre-dtype figure, CogView4’s lines there at its expandable-segments ones, and every tiled and assembly line at its unmeasured floor. A tiled or assembly line of a dtype no tiled point measures on a platform (FLUX.1 fp16 and fp32, FLUX.2 fp32) takes there, value by value, the larger of the SDXL VAE’s line of that operation and dtype and the family’s own bf16 line scaled by the element size (twice for fp32): the FLUX encoders emit 4 to 8 times the SDXL VAE’s moment channels, and a FLUX decode’s node and allocator show in the family’s own lines, not in the SDXL VAE’s.
  • A value on no line its entry carries is hand-set, and only a spec keyed by no platform trait may carry one (hand_set_values in fit.py): a keyed spec prices lines whose every value is half of one, so its value off one fails the data test. MOSS-Image’s k_token is one, bytes per attended token (the image’s tokens at 1/16 of each side plus the prompt’s), whose entry ships no base for its token line. Its points still decide it and bound it from both sides; only the fit check leaves it out.
  • Every value an entry carries is a stand-in, a legacy constant, or decided by an in-fit point. The data test fails when a fitted value is not the fit of its points, naming the value to ship, and when in-fit fresh-process points that span a fit measure a stand-in’s line.
  • An entry marked legacy carries hand-set constants rather than a fit: no point decides them and fit.py does not reproduce them, though its estimator’s points still bound its estimates. LEGACY_ENTRIES in the data test lists the entries, by estimator and platform, allowed to be legacy, so a fit has to remove its entry there too. The builder replaces an entry’s legacy values with the fit of its points once they fit every line it carries within the upper bounds and each line has a point that holds the upper bound, a fresh-process one: probe windows alone never replace hand-set values. An entry whose platform they do not measure stays legacy.
• The Z-Image denoise (denoise.z_image, data/denoise.json) prices its activations as k_denoise_<dtype> bytes per attended token plus a base of b_denoise_<dtype>. The lines are fitted on fresh-process points of scripts/calibrate_z_image_working_memory.py (a random-init transformer of the real geometry with 2 main blocks, bf16, fp16 and fp32 from 512 to 2048 px): an RTX 4090 with the native allocator and with cudaMallocAsync, and an RX 9060 XT under ROCm 10 with AOTriton’s fused kernels, with expandable segments and with the native allocator on Windows (an entry of its own). In a running server the real 30-block Z-Image Turbo measured below its estimate on the RTX 4090 under the native allocator and cudaMallocAsync (bf16) and on the RX 9060 XT with expandable segments (an int8 build). ROCm with the native allocator off Windows has no points and keeps the legacy constants, 0.25 MiB per token and a 1 GiB base for every dtype; they cover bf16 and fp16 but not fp32 from 1440 px up, whose measured lines are steeper. Every other unmeasured platform (XPU, MPS, CPU, cuDNN with expandable segments, MIOpen under hipMallocAsync) is priced as a measured one, and the entries it borrows floor every value at the same legacy constants there (unmeasured_floor), so no unmeasured platform reserves less than the legacy formula.
  • The line admits only the plain forward, measured with a fused attention kernel: every analytic term zero and no attention mask, read off the estimator’s terms, so a term added later takes the forwards it prices out of the fit. A point whose score matrix is materialized fits nothing, because the score term is an upper bound, and fitting the demand less the term would leave the line low by the term’s slack, short for a build of the same platform keys that runs a fused kernel. A ControlNet’s hints (n + 4 hint-sized tensors) and a forced math kernel’s score matrix are checked against measurements of their own in tests/app/invocations/test_z_image_denoise_working_memory.py.
  • An fp8-stored transformer reserves the per-forward cast of its weights as a term of its own (util/fp8.py, peak_fp8_cast_transient_bytes): the largest layer’s float8 weights at the compute dtype, twice for a scaled checkpoint weight. The legacy base covers that cast too, so wherever the legacy constants price, such a model reserves it twice.
  • Denoise lines are shipped-fit only: denoise.z_image is not user_calibrated.
• The FLUX.2 denoise (denoise.flux2, data/denoise.json) prices its activations as k_denoise_bf16_h<width> bytes per attended token and sample plus a base of b_denoise_bf16_h<width>, a line per transformer width (Klein 4B 3072, Klein 9B 4096, [dev] 6144; an unrecognized variant prices as the widest). The attended tokens are the caption’s, the image’s and every reference image’s, and the line’s size is those times the batch. The node runs bf16 only.
  • The lines are fitted on fresh-process points of scripts/calibrate_flux2_denoise_working_memory.py (per width 512 to 2048 px with a 512-token caption, 1024 and 1440 px with a 77-token one, at 1024 px one and three 1024 px references and a batch of two, and long sequences up to 2048 px with one 2024 px reference (33025 tokens), on a random-init transformer of the loader’s geometry with 2 blocks per stream): an RTX 4090 with the native allocator and with cudaMallocAsync, and an RX 9060 XT under ROCm 10 with AOTriton’s fused kernels, with expandable segments and with the native allocator on Windows. On the RX 9060 XT [dev] does not fit 16 GB above 1440 px or with three references, so its [dev] lines there come from 512 to 1440 px, one reference included; every other line is measured up to 33025 tokens where the forward fits the card. The node caps neither the image’s side nor the number of references, so a longer sequence prices on its line extrapolated linearly: fused attention keeps the activations linear in the sequence, and from 16896 to 33025 tokens the line reserves a flat 1.05 to 1.13 times what each measured forward demands under the native allocator and expandable segments; under cudaMallocAsync the 33025-token forwards decide the line (1.05). Several references are resized to 1024 px each, as the node resizes them. 4 blocks per stream measured the same as 2 on the RTX 4090 and 5 to 7% more on the RX 9060 XT. In a running server the real models demanded less than the 2-block point at the same tokens: an fp8 Klein 4B and 9B 0.96 to 0.99 times it on the RTX 4090 (native allocator), a bf16 Klein 4B 0.81 to 0.95 and an fp8 Klein 9B 0.71 to 0.77 (after its fp8 cast term) on the RX 9060 XT (expandable segments). The Klein 4B line holds its 1440 px forward with a 77-token caption at 1.42x and 1.40x under the native allocator, listed in bounds.UPPER_BOUND_EXCEPTIONS. ROCm with the native allocator off Windows has no points and keeps the legacy constants the node reserved by hand, 0.42 MiB per token at width 4096 scaled linearly by the width and truncated per width (330300, 440401 and 660601 bytes) and a 1 GiB base, which tests/app/invocations/test_flux2_denoise_working_memory.py holds byte for byte to the formula they replace; they lie above every measured line. The entries unmeasured platforms borrow (cuDNN with the native allocator, MIOpen with the native allocator on Windows) floor every value at those constants, as the Z-Image entries do. A forced math kernel measured 8.74 bytes per score element over the default kernel on the RTX 4090, against the 17 the score term charges.
  • The line admits the plain forward, reference images and batches included: a batch adds sequences, which the line prices, and its base is paid once per forward. Every forward with an analytic term (regional prompting’s bias, a materialized score matrix, LoRAs, a dequantization or fp8 cast transient) or a mask bounds its estimate but fits nothing. An fp8-stored transformer reserves its cast transient on top, which the legacy base covers too.
  • Two analytic terms price what no fit point measures. Every point is a fully resident model, but the cache may load the transformer partially, and then copies each RAM-held layer to the device for its forward, alive beside the activations: the estimate reserves the largest layer as stored, read off the transformer (the largest module’s own parameters and buffers, calc_tensor_size): in bf16 the single-stream blocks’ fused projection of 9 × width² elements, 679,477,248 bytes for [dev]; for an fp8-stored or GGUF build its float8 or quantized bytes, whose widened copy the fp8 cast and dequantization terms already price. A real [dev] in bf16 with 15.0 of 64.4 GB resident on the RTX 4090 (native allocator) demanded 5509 MB at 8704 tokens (1024 px with one 1024 px reference), 5459 MB at 8612 (1440 px) and 3305 MB at 4608 (1024 px), 535, 518 and 719 MB above the resident 2-block points; the line alone covers 0.96, 0.96 and 0.85 times that, with this term 1.08, 1.08 and 1.05 times. It is charged on every forward, because whether the lock loads the model whole is decided after, and from, the estimate; the calibration script passes weights_resident=True for its resident models. With a regional mask, the forward runs on CUDA’s memory-efficient kernel, which materializes the bias again in a layout of its own: on top of the S × S bias the estimate reserves one copy per sample with its last dimension padded to a multiple of 8. On the RTX 4090 a 2-block Klein 9B-width forward at 1440 px with two 512-token regional captions (S = 9124, a 166.5 MB bias) took 158 MB more of the driver’s figure and 168 MB more reserved than the unmasked one, about one bias-sized transient.
  • CUDA on Linux prices with the RTX 4090’s Windows-measured lines without the unmeasured floor, by judgement rather than measurement: the activations come from the same allocator and kernels on both systems, and the Windows driver model maps more per allocation if anything.
  • denoise.flux2 is not user_calibrated.
• The FLUX.3 denoise (denoise.flux3, data/denoise.json) prices its activations as k_denoise_<dtype>_h<width> bytes per attended token and sample plus a base of b_denoise_<dtype>_h<width>, a line per dtype (bf16, fp16, fp32: the node runs the dtype TorchDevice.choose_bfloat16_safe_dtype picks) and transformer width, read off the transformer. Every released checkpoint is 3072 wide: Klein and Klein Base differ only by Klein’s target-timestep embedder and price on one line. The attended tokens are the joint blocks’ sequence: the caption’s, padded to a multiple of 80, the image’s at 1/16 of each side and every reference image’s, and the line’s size is those times the batch; Klein Base’s CFG passes run one after the other, so only the longer caption counts.
  • The lines are fitted on fresh-process points of scripts/calibrate_flux3_denoise_working_memory.py (bf16, fp16 and fp32 from 512 to 2048 px with a 560-token caption, 1024 and 1440 px with an 80-token one, at 1024 px one, three and (bf16, fp16) ten 1024 px references, a batch of two and Klein Base, and in bf16 long sequences up to 2048 px with ten references, 57904 tokens, on a random-init transformer of the loader’s geometry with 2 per-mode blocks per stream and 2 joint blocks): an RTX 4090 with the native allocator and with cudaMallocAsync, an RX 9060 XT under ROCm 10 with AOTriton’s fused kernels, with expandable segments and with the native allocator on Windows, and a Radeon Pro W7900 under ROCm 10 on Linux with the native allocator. The RX 9060 XT does not fit the ten-reference forwards in 16 GB, so its lines are measured up to 33328 tokens. Plain forwards take 0.289 MiB per token in bf16 and fp16 on the RTX 4090 under the native allocator (0.184 under cudaMallocAsync), 0.291 on the RX 9060 XT and the W7900 under the native allocator (0.167 on the RX 9060 XT with expandable segments), and 0.33 to 0.48 in fp32. Reference images cost up to 8% more per token under the native allocator and up to 17% more with expandable segments, and the fit is the upper envelope of every point, so the reference forwards decide most lines: the RTX 4090’s native bf16 line is 0.328 MiB per token with no base, 1.05x its ten-reference forward. Five RTX 4090 forwards without references, one of them a batch of two, sit just above 1.3x under their lines and are listed in bounds.UPPER_BOUND_EXCEPTIONS. From 26160 to 57904 tokens the bf16 lines reserve 1.05 to 1.18x each long forward. Klein Base demanded what Klein does (32 MiB more under cudaMallocAsync), and 4 blocks per stream the same as 2 on the RTX 4090 and the W7900 and 4.2% more on the RX 9060 XT. A forced math kernel measured 13.69 bytes per score element over the default kernel on the RTX 4090, against the 17 the score term charges; on ROCm, where the SDPA guard computes it in chunks and the term charges one 1 GiB chunk, it took 129 MiB more on the RX 9060 XT and 122 MiB less on the W7900. The released Klein in bf16, fully resident (--mode real, at 1024 and 2048 px and at 1024 px with one and three references), demanded exactly what the 2-block points of the same tokens do on the RTX 4090, and the same or up to 21 MiB less on the W7900; on the RX 9060 XT (expandable segments) 6 to 8% more at the short sequences, shrinking to 0.6% with three references, a near-constant 20 to 180 MiB for its extra depth. Fully resident the estimate reserves 1.07 to 1.15x those forwards on the RTX 4090, 1.04 to 1.16x on the RX 9060 XT and 1.05 to 1.13x on the W7900, 1.08 to 1.24x with the streamed layer the node adds. They ship as out-of-fit points, which hold every refit to the released model from below.
  • Every platform nothing measures prices as a measured one, and the entries it borrows, each with the native allocator (cuDNN; MIOpen on Windows; MIOpen on every other OS, the W7900’s), floor every value at the legacy constants the node reserved by hand, 0.33 MiB per token and sample truncated to 346030 bytes and a 1 GiB base for every dtype, which tests/app/invocations/test_flux3_denoise_working_memory.py holds every unmeasured platform to at least. They lie above every measured bf16 and fp16 line at any size a card holds: the W7900’s, 0.1% steeper, meets them only at 2.9 million tokens, 932 GiB.
  • Two judgements, not measurements, extend an entry past the card it was measured on. CUDA on Linux reuses the RTX 4090’s lines without the floor, as for FLUX.2. ROCm with the native allocator on every OS but Windows prices every card with the W7900’s lines, because two generations on two systems measured close: at every shape both measured, the RX 9060 XT (gfx1200, Windows) demanded 35 to 45 MiB more than the W7900 (gfx1100, Linux) in bf16 and 74 to 84 MiB more in fp16, on the same 0.291 MiB per token, and 290 to 454 MiB less in fp32. A ROCm card without a fused kernel for these shapes prices the score matrix on top, capped at the 1 GiB chunk the SDPA guard computes math attention in.
  • The line admits the plain forward, reference images, batches and Klein Base included. Its analytic terms are the largest layer partial loading streams (in bf16 the 20480 x 3072 caption embedder, 125,829,120 bytes; zero for the calibration script’s resident models), a materialized score matrix, a GGUF build’s dequantization or an fp8-stored build’s cast transient, and the LoRAs’ (lora_sidecar, patches/sidecar_working_memory.py, shared with the Qwen-Image-2.1 denoise): the tensors of every layer that runs as a sidecar, which on a quantized or fp8-stored transformer is every layer, and the largest update one patched module forms per forward, a row per attended token and sample, one per sample in the timestep embedders; where the LoRAs merge into the weights, one module’s worth for a layer partial loading leaves on the CPU. A forward with any of them bounds its estimate but fits nothing. The points are forwards without LoRAs, recorded without the term’s feature, which reads as zero. The node has no regional prompting or ControlNet path, so the estimate has no term for them.
  • denoise.flux3 is not user_calibrated.
• The Krea-2 denoise (denoise.krea2, data/denoise.json) prices the transformer’s forward as k_denoise_<dtype> bytes per attended token plus a base of b_denoise_<dtype>, for the bf16, fp16 and fp32 the node may run. The estimate sizes the transformer’s lock, which covers the whole denoise, so the line prices its longest pass: a target pass (the image’s tokens, one per 16 × 16 pixels, every joint edit reference’s and the longer caption of a CFG pair; the passes run one after the other, so CFG adds nothing), a style reference’s pass (the image and its own caption), or in reference K/V-cache mode the references’ precompute, whose cache’s exact size goes on top.
  • The lines are fitted on fresh-process points of scripts/calibrate_krea2_working_memory.py (a random-init transformer of the default geometry with 2 blocks and the node’s attention processors; bf16, fp16 and fp32 from 512 to 2048 px with a 512-token caption, 1024 and 1536 px with a 77-token one, one to three 1024 px joint references at 1024 px, and in bf16 joint sequences up to 29184 tokens): an RTX 4090 with the native allocator (bf16 and fp16 0.30 MiB per token plus 219 MiB, fp32 0.50 MiB plus 319 MiB) and with cudaMallocAsync (0.23 MiB plus 105 MiB, 0.43 MiB plus 193 MiB). 4 blocks measured the same as 2. Under the native allocator joint references lie on the line of the plain forwards; under cudaMallocAsync a warm forward with one reference read 18% above its cold one and holds the line up. Points name the card as torch does (NVIDIA GeForce RTX 4090), as the script writes them. An RX 9060 XT under ROCm 10 with AOTriton’s fused kernels measured the same grid on Windows with expandable segments and with the native allocator, each an entry for Windows alone, joint sequences up to 29184 tokens in bf16 and fp16 included; fp32 runs up to 22016 tokens with expandable segments and 16896 under the native allocator, as the shipped line plus the weights prices its longer sequences beyond the card’s ~14.9 GiB free: the native allocator 0.30 MiB per token plus 494 MiB in bf16 and 522 MiB in fp16, 0.50 MiB plus 464 MiB in fp32; expandable segments 0.20 MiB plus 407 and 435 MiB, 0.40 MiB plus 240 MiB. With expandable segments 4 blocks measured the same as 2 there too; the native allocator ran no 4-block forward. The card did not drive the desktop: per-process demand is the same, but WDDM paging on a card that also carries the desktop is not measured.
  • A caption token reserves more than an image token, with allocated memory growing alike (allocator reserve, consistent with the text fusion running each caption token as 12 tapped encoder layers before the blocks): under the native allocator 0.73 to 0.95 MiB against 0.30 in bf16 and fp16, 1.2 to 2.2 against 0.50 in fp32. One slope prices both up to the fit grid’s 512 caption tokens, so the 77-token forwards sit up to 1.35x below the line, listed in bounds.UPPER_BOUND_EXCEPTIONS. Beyond 512, which regional prompting reaches by concatenating its captions, each caption token adds caption_excess on top of the line. The caption checks measure it (--mode captions: 1024- and 2048-token captions at 1024 and 1536 px, every dtype, on every measured platform, shipped out of the fit), as do the same checks on Linux on an RTX 6000 Ada (the RTX 4090’s AD102 chip) under both CUDA allocators, shipped without its fit-grid points: under the native allocator it measured byte for byte what the RTX 4090 did, under cudaMallocAsync it scattered both ways. The line alone falls short at 2048 tokens by up to 23% (down to 0.77 times the demand). The term is the most a check demanded per caption token beyond 512 above the line’s slope, over every shipped caption check, times the 1.3 that bounds a fitted term (bounds.MAX_OVER_RESERVATION), rounded up to the next 0.01 MiB: 0.34 MiB in bf16 and fp16 (from 0.26 MiB, CUDA’s native allocator) and 0.74 MiB in fp32 (from 0.56 MiB, the RTX 6000 Ada under cudaMallocAsync at 1024 px with 2048 tokens; CUDA’s native allocator demanded at most 0.39 MiB, the RX 9060 XT 0.37). The estimate covers every check, 1.05 to 1.54 times in bf16 and fp16 and 1.08 to 1.67 in fp32 under CUDA’s native allocator, 1.21 to 1.50 and 1.14 to 1.34 under cudaMallocAsync on the RTX 4090 and 1.17 to 1.50 and 1.07 to 1.40 on the RTX 6000 Ada, 1.05 to 1.46 and 1.08 to 1.63 under ROCm’s native allocator, 1.17 to 1.43 and 1.33 to 1.45 with expandable segments. On the RTX 4090 each group’s largest ratio falls on a forward the line alone already covers 1.09 to 1.47 times.
  • SageAttention (attention_backend: sage) is priced on the SDPA lines: it needs about what SDPA needs. The RTX 4090 ran the fit grid, the long joint sequences, the K/V-cache and the caption checks with it in bf16 and fp16 under both CUDA allocators on Windows, 40 forwards each (--attention sage, SageAttention 2.2.0 serving every block’s attention; a point where it fell back to SDPA is refused). The 34 with an SDPA point of the same forward demanded 1.000 to 1.003 times it under the native allocator and 0.98 to 1.18 times under cudaMallocAsync, the top at 1024 px with a 512-token caption, where the cold forward, which includes SageAttention’s check of its first result, read 1.20 times the warm one. The other 6, the long joint sequences and K/V-cache forwards in fp16 (SDPA ran those in bf16 only), read what their bf16 twins did under the native allocator and up to 2.1% less under cudaMallocAsync. The estimate covers all 80, 1.05 to 1.54 times under the native allocator and 1.05 to 1.52 under cudaMallocAsync. The 80 points carry the feature attention_backend: sage, which the line refuses (ATTENTION_BACKEND_FEATURE), so they ship as out-of-fit checks that bound the estimate from below, never fitted and never merged with an SDPA point. Not measured on Linux, on another GPU or with 4 blocks; masked attention (regional prompting) stays on SDPA under SageAttention.
  • Cross-checks on the RTX 4090 (native allocator, ~84-token prompt), not shipped as points: in a running server the estimate reserved 1.22 to 1.66 times what the real 28-block model grew the reserved peak by over the transformer lock, GGUF Q4_K_M fully resident (1024 to 2048 px, with and without LoRAs, CFG, joint references and the reference K/V cache) and fp8-stored (partly resident in that run, 1.27 to 1.40 times). The top end falls at 1536 px, at 1536 × 2048 and at 1152 × 1536 with one joint reference, where the warm server’s allocator reserved below the line. Loaded partially (the cache told the card has 11.76 GiB, GGUF Q4_K_M at 1152 × 1536 with two plain LoRAs, 6.3 to 6.9 GiB of its weights resident), the model grew the reserved peak by 2.37 GiB against an estimate of 3.30 GiB (1.39 times), with no out-of-memory; the device peaked at 10.85 to 10.86 of the 11.76 GiB, 0.3 GiB of simulated foreign use included.
  • ROCm with the native allocator off Windows and every platform nothing measures price with the formula the node reserved by hand before the lines were measured, declared as the spec’s legacy_formula: 0.5 MiB per attended token and a 1.5 GiB base for every dtype, times 1.1 with CFG and 1.35 with a style reference, plus the whole model’s largest dequantization transient and half a GiB per LoRA, and the terms it lacked on top (the streamed layer, the fp8 cast, the sidecar residual). Those constants reserved 2.3 to 2.9 times what the real model demanded on the RTX 4090. ROCm’s entry with the native allocator for every OS but Windows is legacy; the entries priced_as lends to unmeasured platforms, CUDA’s and ROCm’s with the native allocator (Windows’ and the other OSes’), floor every value at the constants, which prices the formula there too. The reservation is the larger of the formula and the terms. tests/app/invocations/test_krea2_denoise_working_memory.py holds the legacy platforms byte for byte to the formula, frozen, and the unmeasured ones to at least it.
  • A forward has two phases that never peak together. The block phase holds the line, a materialized score matrix (for the config’s query heads), style reference’s extra activations, and the largest block layer’s streamed copy (charged on every forward, as for FLUX.2), dequantization or fp8 cast transient and sidecar residual. The layers outside the blocks never coincide with the block activations; the largest, the timestep modulation time_mod_proj (6144 × 36864), dequantizes to 1.23 GiB in Q4_K, yet GGUF Q4_K_M reserved only a near constant 0.35 GiB more than the fp8 build at every measured size. So the block layers’ transients price the block phase, and outer_layer_peak adds what the largest single layer anywhere needs beyond the whole block phase; it is the one term that reads fitted values, and bounds.line_bounds reads it at 0.
  • Resident through both phases: the caption excess, regional prompting’s masks and bias, the reference K/V cache, a style reference’s captured K/V, and the LoRA patches at their files’ size: a 0.03 GiB and a 0.21 GiB bf16 file of plain LoRALayers on the sidecar-patched GGUF Q4_K_M build measured +0.02 and +0.22 GiB, at 1152 × 1536 and 2048 px alike. Patch types that are no LoRALayer (LoKr, LoHa, DoRA, full layers) on a sidecar-patched layer (quantized, fp8-stored, or one partial loading may leave in RAM) rebuild a full residual weight per forward beside the base weight dequantized again (_aggregate_patch_parameters): the largest such layer’s streamed copy, transient, and weight at float32 plus the compute dtype is priced once (sidecar_patch_residual), derived from the code path; a packed nvfp4 weight counts half its logical elements, so the residual reads low there, which its dequantization transient covers in practice. Known and not priced: on a dense bf16 or fp16 build a LoRA on a resident layer is patched directly, once, before the loop, at about 8 bytes per element of the largest patched layer (~1.7 GiB for time_mod_proj); it never coincides with the block activations, and the device_working_mem_gb floor of every reservation covers it. A style reference’s extra activations are 0.35 × (0.5 MiB per step token + 1.5 GiB), because at 1024 px the styled run’s peak rose ~1.6 GiB over the unstyled one, 0.5 GiB of it the captured K/V.
  • Not measured: regional prompting, CFG with captions of unequal length, the patch types that are no LoRALayer, SageAttention anywhere but on the RTX 4090 under Windows, flash attention, ROCm off Windows, the real model on ROCm, int8 and bf16 builds of the real model. Known gap: on a backend that materializes score matrices, a styled block’s queries also see the reference’s keys, a larger matrix than priced. A family E2E calibration run would record real builds, LoRAs and CFG pairs as probe_demand points, which bound an estimate from below; none ship for Krea-2. The reference K/V-cache windows, the long-caption forwards, the 4-block forward and the SageAttention forwards ship as out-of-fit checks. denoise.krea2 is not user_calibrated.

Measuring estimates: the probe

invokeai/backend/working_memory/probe.py measures what an estimated operation really takes. It is active only inside a calibration_scope(...), which the session runner opens per node of a calibration queue item; without one, model_on_device() behaves as above and reads no counter.

Inside a scope, every model_on_device(working_memory=...) on a GPU is a window. The lock reserves the estimate plus calibration headroom (half the estimate, at least 1 GiB), so a low estimate cannot spill. The window runs from the settled lock to the unlock and collects a MeasuredWindow in the scope’s windows; nothing is written while it is open.

• Counters (working_memory/counters.py): torch’s reserved and allocated peaks, and this process’s usage as the driver counts it. That is D3DKMT on Windows for CUDA and ROCm and NVML on Linux CUDA (the readers verified on hardware; a window’s verified says whether its reader is one of them), DRM fdinfo on Linux ROCm, and the MPS driver figure. NVML lists host process ids, so in a container that does not share the host’s PID namespace this process’s id in that list is unknown: its own id may even be an unrelated host process’s. The reader therefore uses NVML only when /proc/self/ns/pid is the initial PID namespace (inode 0xEFFFFFFC), and windows elsewhere read torch’s counters alone (torch-only). (NSpid in /proc/self/status gives no host id there: it starts at the namespace of the procfs, which in a container is the container’s own.) XPU’s Sysman figure covers the whole device, so it only detects another process’s load. MPS has no peaks; a sampler thread keeps them while a window is open.
• Demand: per counter, the largest segment peak less the baseline and the bytes co-locks moved onto the device; demand_bytes is the largest of those deltas. Each counter misses something: the driver misses HIP’s retained pool, and reserved misses cuBLAS and cuDNN workspaces and, under cudaMallocAsync, part of the allocated peak.
• Segments: the window opens after a release, so cached blocks and tensors from before the lock form the baseline. Another model locked on the same thread inside the window is a co-lock: its weights go into the offset, and its lock interval is kept out of the peak. An estimated window inside another measures alone while the outer one is frozen. A TorchDevice.empty_cache() ends a segment. A direct torch.cuda.empty_cache() call (from a library, or the one cuDNN and MIOpen run after an algorithm search) does not: the driver figure can miss what it released, while the reserved and allocated peaks still bound the demand. Another thread’s TorchDevice.empty_cache() is read as a sample within the segment, before the release and without resetting the peaks; weights that left the device before it fall under weights_left below.
• Invalid (valid=False, with reasons): another thread locked a model on the same device’s cache during the window, held one when it opened, or was still locking one then (foreign_lock); weights left the device other than through one of the window’s co-locks (weights_left): another thread moved a model’s weights to RAM, or this thread did outside a lock, or the weights of a model the cache dropped (a load making room in RAM, Clear Model Cache, a changed model setting, the keep-alive timeout) died while the window measured, which can be after the drop when a reference cycle or a handle still holds the model, and whenever it was dropped, also between nodes with no scope armed (what it does not cover or covers coarsely: device tensors something else still references when the cache moves a model, which go only with that reference; a nested window’s co-lock, offset by that window alone, so what it moves off the device voids its suspended ancestors and what it moves on leaves them reading high); the allocator failed or retried; Windows still holds more than 256 MiB of the process in shared memory after a short settle; the operation raised or was canceled, or its node failed or was canceled after the window closed; the Krea-2 SDPA benchmark is active; XPU’s device-wide usage grew beyond this process’s own (foreign_device_load); a probe step failed (probe_error: a reading, a release, a peak reset or a synchronization, which also voids the scope’s later windows when it hit the lock bookkeeping; a window that cannot open is recorded invalid, and its operation runs unmeasured); or no counter could be read.
• Not fit-eligible: a model that was not fully resident when the window opened. The measurement stays valid, but streamed weights are working memory the estimate does not price.

working_memory/fingerprint.py keys measurements. device_class_for() gives the same id to identical cards. environment_for() splits hard facts, which change an operation’s memory and make up its id (the operating system as windows, linux or other, torch, runtime and convolution-library versions, allocator backend and configuration, the ROCm SDPA guard, the memory settings passed in), from soft facts kept for diagnosis (full versions, driver, OS build).

Calibration runs

A calibration run measures one model’s operations on this machine’s GPU; estimates read the result as user fits (below). The service is invokeai/app/services/working_memory_calibration/, the routes are in invokeai/app/api/routers/working_memory.py, and every route is admin-only (403 otherwise).

1. POST /api/v2/models/i/{key}/working_memory/calibrate takes a client-built graph, the identities of the models it names (graph_models), batch data, an optional idempotency_key, and an optional attempt_id with its attempt_steps. The server turns off use_cache on every node, sets is_intermediate on every node and clears the board of every WithBoard node (a node saves its images with its own flag and that board), sets runs=1, sets the graph’s id to the run id, and sets the origin to invokeai:wm-calibration:<run_id> and the destination to invokeai:wm-calibration. It records the run, enqueues the batch as the requesting admin, then records the run’s items.

   • An attempt_id (1–64 of letters, digits, _ - . :) groups the runs one owner enqueues for one calibration attempt (two owners that send one id start two attempts), and attempt_steps (1–256) says how many runs the attempt queues. They come together or not at all (422); the server does not read a count out of the id. Both are stored on the run (and left out of the file when absent, so older readers keep reading such runs) and returned in the run summaries of the status and the model summary. Every run of an attempt sends the same attempt_steps: a run that sends another count than the attempt’s runs already stored, or a repeat of a keyed request that changes it, is refused with 409 (checked under the store file’s lock, so two such runs cannot both be recorded).
   • One attempt runs on a device class at a time: its steps measure the GPU back to back, each on what the last one left, and another attempt’s items between them would measure neither. A run of an attempt is refused with 409 and detail {"code": "calibration_running", "message": "Another calibration is still running on this GPU. Wait until it has finished, or stop it, then calibrate again."} while the queue holds an item that is pending, in progress or waiting of another attempt’s run in the device class’s file (get_unfinished_origins, read under the file’s lock, which the queue takes only outside its transactions), and, with the message Another calibration on this GPU is between its steps. Try again in a moment, or once it has finished. under the same code, while another owner’s attempt is open: between its steps, short of attempt_steps runs, not aborted (no other run since its first included), its model as its runs recorded it, and active within ATTEMPT_IDLE_GRACE (two minutes) of its newest run’s creation or the latest observation of one of its items. The client posts the next step as soon as the last one finished; a step that runs long is running, and its last observation starts the wait. An attempt idle for longer is taken for abandoned: it blocks nothing, and, short of steps, never feeds. The grace is measured on the stored times against the clock now, so a jump of the wall clock lengthens or shortens it once. The owner’s own open attempt blocks nothing while none of its runs is running: an open attempt is almost always one whose submission stopped partway (a full queue, a lost connection, a closed tab), and the owner’s new attempt replaces it; the replaced one stays short, so aborted, and a later step of it is refused as interleaved. The attempt’s own next steps pass, and so does a repeat of a keyed request that finds its run. A later step of an attempt after any other run of the device class was created since the attempt’s first run, another attempt’s or one without an attempt (a script’s, the family E2E tool’s: it measured the GPU in between too), is refused with 409 and detail {"code": "calibration_interleaved", "message": "Another calibration ran on this GPU after this calibration started, so its measurements can no longer be combined. Calibrate again."}: it would measure the GPU after that calibration, not after its own earlier steps. The attempt then stays short and is aborted; an attempt another run was created between the first and the last run of is aborted too, should a store hold one (_interleaved_attempts). Two clients that start attempts at once cannot both pass: every creation of an attempt’s run holds one OS file lock, <db_dir>/working_memory/creations/attempts.lock, from finding or recording the run to recording its items, in every process serving the root, so the later creation finds the earlier one’s items queued. The creation takes it before its run’s creation lock (below), and it is a file of its own: on a file they shared, the second lock would wait for the first. A run without an attempt_id takes neither the lock nor the check.
   • The graph names each model it loads as the client compiled it (ModelIdentifierField: key, hash, name, base, type), in a node’s field or inside one, such as the LoRAs of a collection loader, or as a mapping of that shape in an untyped field, such as a metadata node’s record of a model (graph_model_identifiers): in the graph’s JSON a client cannot tell the two apart. The request states the identity of each of those models as the record the graph was compiled over had it, in graph_models: per model key its hash, base, type, format and variant (CalibrationCompiledModel), every one of them required (format and variant null when the record has none; hash, base and type never null). It names exactly the graph’s models; a missing or extra key or an unstated fact is refused with 422. The identities come from the records the graph was compiled over, never from a read at posting time, which would carry the change the comparison has to catch. Under the run’s creation locks (below), the creation reads each model’s record once and compares it with the graph’s identifier and with its graph_models entry. A model whose record differs in any of those facts, or that was removed, refuses the request with 409 and detail {"code": "model_changed", "message": "<name> changed after this calibration was planned (<facts>). Calibrate again."}, the facts among hash, base, type, format and variant, or with the message <name> was removed after this calibration was planned. Calibrate again. So a run always measures the configuration its graph was compiled for. The run records the identity of its model as that comparison read it: model_hash, and model_identity (base, type, format and variant); and that of every model the graph names, with its hash (graph_models), which equals what the client compiled against. A change after that is caught at admission. The run id does not derive from graph_models, so a keyed repeat that finds its run without items and has to enqueue its batch also compares its graph_models with what the run recorded (an entry the run lacks, or a recorded entry without a hash, compares only what it has); a difference is refused with the same coded 409, since the item would fail at admission on what the run recorded. A later run of an attempt is refused the same way when its model differs from what an earlier run of the attempt recorded for it, or a model its graph names differs from what an earlier run recorded for that model, the hash included (a graph compiled after a reinstall names the new hash): the attempt would measure two versions of a model, and, short of that step, it never feeds a fit. The message names the model that changed.
   • The queue takes all of the batch’s items or none (enqueue_batch(..., all_or_none=True), checked under its admission lock): a run with only some of its cases would measure only some of them. Without room for all of them nothing is enqueued, the key stays unused, and the run is removed if this request recorded it.
   • With a key, the run id is derived from the owner, the key, the request’s fingerprint (the model key, the graph without its id, and the batch data) and the attempt_id if given. A client that sends no graph id would otherwise get a fresh one each time and could never repeat a request. A repeated request finds its run and gets the same response; nothing new is enqueued, even after the queue collected the key’s receipt. A key therefore names one run. Once a reset removed the run, a repeat is refused (404) whether or not the queue still keeps the receipt: the store keeps the ids of removed runs, so the run is never started again and its queued items keep failing at admission. Another request that reuses the key (another model, graph, batch data or attempt_id) gets another run id and so another origin. While the key’s receipt exists, that request is refused (409) and the first run stays; the receipt is checked before the request’s model, so an unknown model does not turn the refusal into a 404. Once the queue has collected the acknowledged receipt, the request starts a run of its own. Clients make each key unique per run of an attempt, for example from the attempt id and the step. Acknowledge the key through POST /api/v1/queue/{queue_id}/enqueue_batch/acknowledge.
   • Once its items are queued, the request answers with them, also when the enqueue raised after its commit (the queue announces a batch after committing it). When the queue cannot tell whether it holds them, the run is kept and the enqueue’s own error answers.
   • If the store fails for a passing reason (an unwritable root, a lock timeout) to record the queued items, a request without a key still succeeds: its items run, and the status finds them by their origin. With a key that is a 503, and a repeat records them: from the key’s receipt while the queue keeps it, otherwise from the queue items with the run’s origin and the items the runner saw. It never enqueues a second batch for the run: the run is marked before its batch is enqueued (enqueue_attempted, set as the run is recorded), and only the attempt that set the mark clears it again, when its own enqueue queued nothing; a replay of the key’s receipt never clears it, since the attempt that queued the batch may have queued it. A repeat that finds a marked run without items, without the receipt and with no item in the queue or seen by the runner answers with the run and no items, and the run reads as canceled. That is the case when the items were cleared or pruned before the runner saw them, when the server shut down between the mark and the enqueue, or when an enqueue raised, the queue could not be read to tell what it had queued, and it had queued nothing. Should clearing the mark fail (it is logged), the run stays marked although nothing was queued, and the key cannot start it; a request with another key can. Recording them from the receipt checks neither the model nor the device, so a model removed meanwhile does not fail the repeat; a request that has to enqueue needs the model. A run removed meanwhile (404) or an unusable store file (409) fails the request with or without a key; its queued items then fail at admission.
   • Per run id, one creation runs at a time, and it runs to its end even when its request is canceled: the enqueue runs in a worker thread and can commit after the cancellation. Only the server’s shutdown cancels the creation with it; the run then stays as far as it got, and with a key a repeat of the request completes it. For a run with a retry key, “one at a time” holds across every process serving the root: the creation holds an OS file lock (the store’s own lock primitive) from finding or recording the run to recording its items, so another process’s repeat waits and then finds what the first attempt did. A holder that dies releases the lock with its process. A repeat waits up to 30 s (LOCK_TIMEOUT_S) and then answers 503, which with a key can be retried. The run ids share 256 empty lock files, <db_dir>/working_memory/creations/<xx>.lock by the first two hex digits of the id’s SHA-256. No lock file is ever removed (removing one could let two processes each lock a file of that name), and the folder never holds more than those 256 and attempts.lock. Unrelated runs on one file wait for each other, which costs little for a short administrator’s action. Within one process the repeats of a run first wait on an in-process lock per run, on the event loop, so only one of them at a time takes a worker thread to wait for the file lock: a burst of repeats cannot fill the default executor. A run without a key has a fresh id that no other request names, so it takes no lock. A run is removed only when this request recorded it and nothing of it was queued.
   • All of it before the request’s fingerprint is taken, so a keyed repeat stays the same request; batch data that sets use_cache or is_intermediate is refused (422); a board cannot come as batch data, whose items are strings, numbers, images and videos. The images a calibration item saves are therefore intermediates on no board, never in the gallery. They stay until Clear Intermediates removes them, which keeps them while their item runs, as it keeps every active session’s. The destination stops clients from routing the results to a gallery or canvas view.

2. enqueue_batch refuses any origin that starts with that prefix, in any case, for every user (400). Only the calibration route creates calibration items. Retries and workflow-call children inherit origin and owner from their source item.

3. DefaultSessionRunner.run() checks each item’s origin once. For a calibration item it calls admit() before any node runs, also for one left pending before a restart: the invoker starts the session processor after every other service. (The queue’s start cancels the items that were in progress.) Retries and children are checked here too. Otherwise the item fails with the reason. The checks:

   • the owner is an active admin (single-user mode counts as admin);
   • the machine has exactly one GPU generation device;
   • the run exists for that GPU’s device class;
   • the item measures the models the run was planned with: its model’s record still has the model_hash and the model_identity the run recorded (a run stored without the identity compares the hash alone), and every model the graph names still has the hash, base and type it names, and the hash, base, type, format and variant graph_models recorded for it (an entry stored without a hash compares the rest). Otherwise the item fails with Model <key> was removed after this calibration was planned. Calibrate again. or <name> changed after this calibration was planned (<facts>). Calibrate again., and the run’s item records error_code: model_changed, which the run’s error_code in the status and the model summary carries beside the message in error.

4. Each admitted node runs in probe.calibration_scope(...), keyed by the node’s source id. Its windows are persisted after the node completes, fails or is canceled. Unless the node completed, every one of them is persisted invalid (exception when it failed, canceled otherwise), including a window whose own operation finished before the node raised: run_node swallows the failure or cancellation, so the runner takes the outcome from its node hooks. Windows from a device of another class, or of a run that was removed, are dropped with a warning. So are windows of a model reset after the node started, also by another process sharing the root: the runner reads the device class’s reset epochs before the node’s scope opens, and the write drops every point attributed to a model whose epoch has moved since, as the reset attributes them (measured_models: a submodel’s <key>:vae window counts for its model). Without that, a reset while the node runs (of a VAE inside a transformer’s run, say) would be undone by the node’s later write. If they cannot be saved, the item fails and the session stops. Every other item gets no scope and runs as before.

5. After the item, the runner records the item’s outcome in the run, even when the session raised. A queue row that still exists has the last word, failed or canceled included. When the row is gone, the runner’s view of how the queue left the item after the session decides:

   • The queue completed or failed the item, and a prune or a clear removed it afterwards: the item keeps that outcome, a failure with its error.
   • Clear Queue removed the item before the queue settled it: the item is recorded as canceled. (A prune cannot do that: it removes finished items only.)
   • The runner raised before the queue settled the item: the item is recorded as failed with the error.

   Nothing is recorded ahead of the queue. If the process stops between the queue settling the item and this record, and the queue then prunes or clears the item, the item’s outcome reads as canceled, never as a success the queue did not confirm. On its own, such an item leaves the run canceled, or partial when it measured valid points; the run as a whole still reads failed when another of its items failed. The points it measured are kept either way.

The store has one JSON file per device class at <db_dir>/working_memory/<device_class_id>.json (CalibrationStoreFile, schema_version 1). It is machine data, so it lives next to the database rather than inside it. It holds:

• the device class, with its id;
• the runs: model key, the model’s hash and identity at creation (model_hash, model_identity), owner, environment and its soft facts at creation, queue items, whether an enqueue of the run’s batch may have committed (enqueue_attempted), and for each item the last status the runner saw, why admission failed it (error_code), and every retry of an item, recorded pending as the queue commits it until the runner sees it; how many of the queue’s retry operations that retry items of the run are in flight, however many items each retries (retries_in_flight); the models the graph names as they were then (graph_models); the soft facts, the model_identity, the graph_models, the retries_in_flight, the error_code, the attempt_id and the attempt_steps are left out when absent (or 0), so runs stored without them still load and older readers still read runs that lack them;
• the points: MeasuredWindow plus env_id;
• the environments those reference;
• the ids of the runs with a retry key that a reset removed (removed_run_ids), so that a repeat of the request that started one cannot start it again. They are kept for good, about 40 bytes per removed run. A run without a key gets a fresh id and is never repeated, so its id is not kept;
• per attempt, an owner’s (user_id) attempt_id, what resets removed of it (removed_attempt_steps, a list of such records): how many of its runs (removed), the attempt_steps they sent, and whether the attempt was running or already stopped when a reset removed one (aborted), judged then from the queue. A removed run counts as queued, so the runs left are judged on every run the attempt queued; a later run of the attempt must send the same attempt_steps. A record is dropped at a reset once no stored run of its attempt is left, unless that reset found the attempt running, since a later step of it may still come. The key is left out of a file without any record, so readers that predate it still read such files. Development builds stored the records by attempt id alone; such a record reads as the record of each owner whose stored runs carry the id, and one no stored run carries is dropped;
• per model key, its reset epoch (model_reset_epochs): how often the model was reset. The reset raises it in the same file change, so it is only written when the reset is.

A file without either of these two fields reads them as empty.

Writes and reads:

• Every change takes a process lock and an OS lock on the file’s .lock sibling, re-reads the file, applies the change, and replaces the file atomically. Two processes sharing a root therefore never lose each other’s writes.
• On Windows a process that ignores the lock (an antivirus scanner, a search indexer, an editor) can hold the file open and refuse the replace. The write then retries a few times with backoff, about 1.5 s in all (REPLACE_ATTEMPTS, REPLACE_RETRY_S).
• Readers take the lock as well, because Windows cannot replace a file another process has open.
• Each change rewrites the whole file, which is fine at calibration volumes.

A file is unusable when it has another schema version, invalid content, an id that does not match its fields or its file name, or a field this version does not know. Every store model forbids unknown fields: ignoring one would drop it at the next rewrite. An unusable file is reported in the status, makes calibration unavailable for its device class, and is never written, moved or deleted.

A run’s status is derived, never stored:

• The live queue items with the run’s origin come first. For items the queue has pruned, the runner’s recorded statuses are used. An item the queue no longer has and that the runner never saw finish (the process died mid-item) counts as canceled.
• A retry replaces the item it retried. The queue’s own record of the retry links them while it has the retry; the store records it too, with the retry’s commit: the announcement after it (QueueItemsRetriedEvent) cannot promise a record, since the server may stop before it is handled. SessionQueue.retry_items_by_id calls the calibration service before it commits retries (retries_starting), which counts the retry operation in flight on each run of a calibration item among the sources (retries_in_flight, under the store file’s lock, one per operation however many of the run’s items it retries), and again once the commit is done (retries_committed), which records each retry on its run, pending and standing for the item its lineage started from (the queue’s retried_from_item_id), and counts down. A full queue, or a failure before the queue’s insertions ended, which rolls them back, counts down without recording anything. A failure once they ended can be the commit’s own, which may have committed all the same (a lost connection to MySQL or MariaDB): which retries exist is unknown, so the count stays, and is logged. The count-down is tried once more after a lock timeout. The queue holds no transaction while it waits for the file’s lock. A retry inherits its source’s origin, which names the run: a recreated database that gave a recorded id to an unrelated item cannot make that item’s retry the run’s, and a bulk retry of generations touches no store file. A store that cannot count the retries refuses them: the queue retries nothing, also the request’s items that are no calibration’s, and the retry route answers 503. A count left over, because the server stopped, the commit failed that way, or the store could not be written between the count and the record (which is logged), means a retry may exist that the store does not know: the run’s attempt never feeds a fit, also after a restart, until a reset removes the run. So a retry canceled and pruned before the runner started it reads as canceled, also after a restart, and a retry the runner admitted before the record keeps what the runner saw. The run’s items include each item a known retry came from, even when nothing else knows that item any more.
• The run is running while any item is pending or running, failed when any failed (with the first failed item’s message as error and, when admission failed it, its error_code), and completed when all completed. Otherwise it is partial if valid points were recorded, or canceled if not: an invalid point measured nothing usable.

Routes, besides the 403 for non-admins and a 503 when the store cannot be read or written (an unwritable root, a lock timeout):

Route	Does	Other errors
POST .../i/{key}/working_memory/calibrate	starts a run (above)	404 unknown model, the run was removed while it was being created, or a reset removed the run the retry key started; 409 no GPU or several generation devices, unusable store file, no room in the queue for all of the run’s items (nothing is enqueued; the run is removed if this request recorded it), a retry key used for another request, a run of another attempt still running on the device class (calibration_running), or a model the graph names that changed or was removed since the client compiled the graph, or that a repeat of a keyed request states otherwise than the run it names recorded (model_changed), these two with detail {code, message} (above); the 404 and 409 for the run’s store, and with a key the 503, can come after the items were queued (above); 422 a call_saved_workflow node, batch data setting use_cache or is_intermediate, an invalid batch, fields the server sets, or graph_models that does not name exactly the graph’s models, every fact of each (hash, base and type never null); 429 too many unacknowledged retry keys
GET .../working_memory/status	every device class with its runs, point counts per estimator, version, validity and fit eligibility, and environments with hard and soft facts compared to the current environment (soft_matches_current compares the soft facts first recorded with the environment, not a later run’s); available says whether a run can start now: one GPU generation device, and a usable store file that can be written. The check stages a temporary file in the store’s folder (or the nearest folder above it that exists) and removes it at once, checks that the lock file is writable, and on Windows that the store file is not read-only. Per device class with a generation device here, measurement_method (the driver reader its windows use, counters.driver_method) and measurement_verified (counters.VERIFIED_METHODS), null for other classes	503 when the store folder cannot be listed
DELETE .../working_memory/runs/{run_id}	removes a run and its points; its items that have not started yet fail at admission, also after a repeat of the request that started it. A run of an attempt is counted in its owner’s attempt’s removed_attempt_steps record, so the runs left are judged as part of the whole attempt; a reset while the attempt runs, or of a run that ended canceled, marks the attempt stopped. The reset judges that from the queue inside the change that removes the run, under the store file’s lock, so a retry queued just before it is seen	404 no usable file holds the run
DELETE .../i/{key}/working_memory	removes the points the summary attributes to the model (measured_models): its runs with all their points, and every point of the model or one of its submodels taken in other runs (a VAE measured inside a transformer’s run, say); other models’ points on the same lines stay. The removed runs of attempts are counted as for a single run. A submodel’s window carries the model cache’s key for it, <key>:<submodel> (<key>:vae), and counts for its model; the reset fence uses the same attribution	409 no current device class, or an unusable file
DELETE .../working_memory	removes every run, point and environment of a device class; it keeps only the removed_attempt_steps records of attempts that were running, which a later step of them still meets	409 as above
GET .../working_memory/export	the file byte for byte: a CalibrationStoreFile when this version can read it, otherwise whatever the file holds	404 no file; 409 no current device class
GET .../i/{key}/working_memory/plan	the cases that calibrate the model here (WorkingMemoryCalibrationPlan, below), or supported: false with the reason	404 unknown model
GET .../i/{key}/working_memory/summary	per line and environment: the lines the model’s windows measured, and in the current environment each line the model’s decode prices on (its plan’s estimator and dtypes) that other models’ windows measured. Per line, every window that feeds it (all models’ counted windows of that environment that join the same user line): sizes, largest demand, shipped and user estimate; a size whose valid windows all are excluded is listed too, with counted: false and excluded_reason (fragmented, partial_residency, or aborted_attempt when all its windows come from stopped attempts), its largest such window and the user estimate there; the user line’s measured range and state (in_use, stale, unusable, insufficient); own_windows, the model’s share (0 for a line only others measured); measured_by, each contributing model with its name (null once uninstalled) and window count. Also the model’s windows that fit no line, and its runs, each with attempt_steps and held_back (running, aborted_attempt, or null)	404 unknown model; 409 no current device class

The reset, export and summary routes act on the current device class unless device_class_id is given.

User fits

backend/working_memory/user_fit.py owns the rule; the calibration service owns the store and the cache.

• Only the estimators that opt in (EstimatorSpec.user_calibrated, the five VAE estimators) take part. The windows of any other, a Z-Image denoise’s say, stay in the store and count among the windows that fit no line, but no user line is fitted on them, none moves their estimates, the summary lists none of their lines, and they size no plan: denoise lines are shipped-fit only.
• A user line is fitted per line of fit.py: the shipped entry that carries the values (one hop of values_from) and <operation>_<dtype>. A window joins where its estimator’s line admits it as the line’s own demand (Line.own_demand; for the convolution lines: untiled, one sample, no in-place upcast, an AutoencoderKL: the SD estimator’s AutoencoderTiny windows join no line, and tiled windows none either, since a tiled line sits on the untiled one and is the shipped fit’s alone), it is valid and fit-eligible, it is not fragmented, its demand is positive, and its estimator version has shipped values on the platform. Every user point is a probe_demand window, so all of them rate the hull’s edges. No fresh-process point caps a user line and no floor holds it, so that edge always covers its points and the raise of shipped lines (fit.raised_pair) has nothing to do there. fit.user_pair adds USER_MARGIN, the shipped 5%, to the envelope and rounds up as the shipped values do. A line needs two sizes MIN_AREA_RATIO apart.
• Fragmented windows measured the allocator’s history, not the operation. One rule, fragmentation.fragmented, judges them for the user fit and for scripts/working_memory/build_shipped.py, which flags shipped points by it. It judges a measurement: the fit-eligible windows of one attempt (an owner’s attempt_id, so two owners’ attempts with one id are two measurements; else the run; for the builder also one server lifetime) at one environment, estimator, version and features, in the order they ran. A measurement drops out whole when reserved grew by 10% from its first window to its largest while allocated stayed within 2%, or when its largest window’s reserved-to-allocated ratio is 1.2 times the median of the same line’s smaller measurements that count, from any attempt. Each attempt is judged on its own, because its windows ran back to back on one allocator state: a recalibration can bring back a size an earlier attempt measured fragmented, and a later fragmented attempt drops no clean one. The ratio rule catches a size whose windows all ran fragmented, or that kept only one fit-eligible window: on an RX 9060 XT the fp32 SDXL decode at 1216 px measured 1.47 against 1.10 at 576 and 832 px, with one fully resident window beside two partly resident ones. Windows of a running attempt are judged but set no reference, as they join no user line yet. Partly resident windows (partial_residency) never count and enter neither rule. A size counts when any attempt’s windows at it count, and its windows and measured_bytes cover those windows only.
• The windows are those of the current device class’s file taken under an environment whose hard facts equal the current ones (same env_id). A soft-fact difference keeps the fit. The summary reports a line as stale, with the facts that changed, when the soft facts of the newest finished run that fed it differ from the current ones; a run stored without soft facts stands for the environment’s, those first recorded with it. So recalibrating after a driver update clears stale. Older windows of the same hard environment keep feeding the line: the hard facts are what decide whether a window still measures this machine. A hard mismatch leaves the shipped values.
• EstimatorSpec.estimate asks the registered provider (user_fit.set_user_fit_provider; the service registers it in start() and clears it in stop()) for the user fit of the spec, platform and device. Without one, the estimate is the shipped one, bit for bit. With one, user_fit.resolve reads the operation’s line and its size there (output pixels) as EstimatorSpec.line places it. At or below the largest measured size the user line replaces the shipped one; above it, the larger of the two estimates wins. A tiled operation prices its tile on the larger of the user and the shipped untiled line, within the measured sizes and above them, and its tiled and assembly lines stay shipped: they were fitted above the untiled demand measured where the shipped line was, which a lower user line beneath them could undercut, while a card whose untiled decode needs more than the shipped line needs it for each tile too. The source starts with user:<device class id>:<env id> and names the line where the user line decides. A line the user did not measure stays shipped. Because estimates feed reservation_bytes, a user line also moves should_pretile_vae_decode.
• Only finished runs and attempts feed a fit. The windows of a run the queue holds an item of that is pending, in progress or waiting join no user line, whoever refreshes the fits meanwhile. The queue tells by the run’s origin, which an item carries from the moment it is queued, before the run records it or the runner admits it; a retry carries it too, so a run that failed holds its windows back again while its retry waits, and releases them once the retry finished or was canceled or removed. Nor do the windows of any run of the same attempt, the same owner’s attempt_id: the client posts each step of an attempt as its own run, so an attempt’s windows feed no fit while any of its runs is in the queue (held_back: running), and its later, larger cases are priced without its own earlier windows. Once none is, an attempt feeds a fit only if it is complete and was not stopped (aborted_attempt); its windows join at the first refresh after that, whether its runs completed or failed. Complete means it queued attempt_steps runs, those a reset removed since included (removed_attempt_steps), so removing a run of a finished attempt discards only that run’s windows. A reset while the attempt runs, or after it was stopped, stops it for good: its earlier steps are gone, so it cannot measure every size it planned. The UI disables Reset while a calibration runs, so only API clients do that. Stopped also means an item of it interrupted it, judged per item: the item that stands for a lineage, its latest retry or itself, was canceled (Stop cancels what the queue still holds), is gone from the queue before the runner saw it finish (a queue clear), or failed at admission because a model changed after the calibration was planned (error_code: model_changed: its later steps would measure another model). So does another run created between its first and its last run. The items are those the store records: the run’s batch and every retry, recorded as the queue commits it or by the runner, so a retry canceled and pruned before anything refreshed, a restart included, still stops its attempt; and an attempt one of whose runs has retries in flight left over (retries_in_flight) is stopped too, since a retry of it may exist that the store does not record. Any other failed item, such as an out-of-memory error at the largest size, ends that step and not the attempt, also beside a canceled one in a run whose batch data made several items, whose displayed status is then failed; a canceled item whose retry completed or failed no longer aborts it, and a retry canceled while it waited does. A run with no item at all, whose items were never recorded nor found by its origin, stops it as well. An attempt with fewer runs and none running is stopped too: Stop, or a submission that failed, ended it before it queued them all (a step not posted yet looks the same until it is). A stopped attempt measured only some of its sizes and never feeds a fit, also when those span enough sizes for a line: its runs read held_back: aborted_attempt, its windows are excluded with excluded_reason: aborted_attempt, and a line only it measured is insufficient, with a problem that says a stopped or interrupted calibration measured it. A run without an attempt_id is judged on its own and is never stopped this way: canceled, it feeds what it measured; another attempt that finished counts as usual. Runs stored with an attempt_id but no attempt_steps come only from development builds of the calibration; their attempt is judged on the runs it recorded, as before, rather than making the store file unusable. The summary shows a running run’s windows in the sizes, window counts and measured_by, while its user line, range, user estimates and state are those of the fit estimates read; a line only a running run measured is insufficient and says so.
• The provider answers from memory. refresh_user_fits() rebuilds that memory per GPU generation device. DefaultSessionRunner.run() calls it once when it starts any queue item, calibration items included; get_plan calls it first, so the case sizes price with the store as it is, and every reset calls it after its write, so a reset applies at once. It stats each device class’s file and re-reads it only when the mtime, size or inode or the environment changed. Each refresh of a file with runs, and each model summary, reads the queue to judge the attempts. To tell which runs are running it reads the origins of the queue’s unfinished calibration items (get_unfinished_origins): the queue finds those items by their status and tests the origin on their rows, so the read costs what the queue has pending or running, not its history; on SQLite about 0.05 ms beside a handful of unfinished items, whatever the history, and 3 ms beside 100 (34 ms beside 1000) pending items of 64 KB sessions, whose overflow pages the origin sits behind. Grouping the runs into attempts costs about 0.5 ms at 180 runs and 11 ms at 3600. To tell whether a settled attempt was stopped it reads by id (get_queue_item_summaries_by_ids, the primary key) the items the store records whose end the runner did not see. A finished item never changes, so its outcome is kept, per run and item id (a recreated or restored database can give two runs’ items one id), until its run is removed, by a reset or by a creation that queued nothing, and each item is read once it finished; a run with an item the queue shows pending or in progress, a retry queued since the read for running items, holds its attempt back as running for that refresh. Every judgement of these outcomes (a refresh, the model summary, a new attempt’s check, a reset), with its read of the store file, and every removal holds one lock of the service, taken before the file’s lock, so a removal’s runs leave no outcome behind. A row counts only while it carries the run’s origin: a recreated or restored database gives the id to an unrelated item, and the run’s own item is gone, so it reads as canceled. An unchanged file with an unchanged set of runs held back is not fitted again. A calibration item, a retry included, is in progress when the runner refreshes before it, so it is priced without its own attempt’s windows. Nothing on the estimate path reads a file. Several generation devices each use their own class’s fit; calibrating them is not supported.
• Shared lines: on cuDNN the SD, SD3, FLUX.1 and FLUX.2 estimators borrow one entry, so an SDXL fp16 decode calibrates every estimator’s fp16 decode there. The plan lists the installed models this affects.

The plan (working_memory_calibration_plan.py) covers the decode of SD1, SD2, SDXL, SD3, FLUX.1, FLUX.2, Z-Image and CogView4 main models, and standalone VAEs of those bases through an installed main model that accepts them (VaeFacet): through Generate’s vae slot where the base has one (FLUX.1, FLUX.2, SD3, Z-Image), otherwise with the VAE the main model decodes with, its own or the default VAE its settings name, which prices on the same lines. Each case is a two-step txt2img graph-export case per size and decode dtype; the SD bases set vaePrecision to fp16 and to fp32, the others run the dtype their loader picks: SD3 the precision setting; a standalone FLUX.1 or FLUX.2 VAE, and the one Z-Image decodes with, the configured dtype except fp16, which becomes bf16 where the device has it and fp32 otherwise (TorchDevice.dtype_avoiding_float16, which the loaders use too), so precision: float32 plans the fp32 line at the smaller sizes an fp32 decode fits untiled; a pipeline’s own VAE what its loader yields: bf16 for a FLUX.2 Diffusers pipeline, the device’s bf16-safe dtype for a Z-Image one, and fp32 for a self-contained SDNQ pipeline, whose loaders call from_pretrained without a dtype (bf16 for a FLUX.1 SDNQ pipeline’s SDNQ-quantized VAE); test_working_memory_calibration_plan.py holds these to the loaders. A pipeline decodes with its own VAE while the vae slot is empty: a FLUX.2 or Z-Image Diffusers pipeline, and a self-contained SDNQ pipeline (is_self_contained_sdnq_pipeline, for FLUX.1 is_self_contained_sdnq_flux1_pipeline), except that the FLUX.2 [dev] model loader node takes components only from a Diffusers pipeline, so an SDNQ [dev] model is no pipeline of its own. The slot is not empty when the main model’s default settings name an installed VAE it accepts (default_vae): selecting the model applies it (getModelDefaultVae), the SD, SD3, FLUX.1, FLUX.2 and Z-Image graphs decode with it, and the FLUX.2 and Z-Image model loader nodes take it over a pipeline’s own VAE, so the plan sizes for that standalone VAE; the graph of a self-contained FLUX.1 SDNQ pipeline, and the CogView4 graph, send no VAE. Another FLUX.2 model decodes with the VAE of the component source Generate picks on its own among the installed FLUX.2 pipelines its model loader node takes components from (Diffusers ones only for [dev]): the first whose text encoder fits it ([dev] for [dev], else the Klein variant’s Qwen3 size), else the first one; that graph keeps a pipeline that does not fit only beside a standalone encoder of the model’s (a Mistral encoder for [dev], else a Qwen3 encoder of the Klein variant’s size), which then encodes while the pipeline gives its VAE. Without a pipeline, and for every other FLUX.1 model, a standalone VAE decodes. The sizes are about 0.25, 0.5 and 1.0 of the largest square, in steps of 64 and at most 4096 px, whose calibration reservation (the estimate plus max(1 GiB, 0.5 × estimate), and never below device_working_mem_gb, as the lock applies it) leaves the VAE fully resident (decode_budget): it must fit both 90% of the card (VAE_PRETILE_VRAM_FRACTION) and what the model cache can give on the device, each less 1 GiB for the VAE’s weights. Nothing else stays beside it: inside a calibration scope, model_on_device first moves every other unlocked model on the device to RAM whole (ModelCache.offload_unlocked_models_from_vram, under the global load lock, inside the probe’s lock interval; locked models, a window’s co-locks among them, stay), and only then locks. The reason: a lock offloads only the shortfall it predicts and splits the last model it reaches, and a split can free a little less than predicted. On an RX 9060 XT, 9 of the 12 calibration decodes that had to split the SDXL UNet came up 35 to 55 MB short and ran the VAE partly resident, so their windows fit no line. Outside a calibration scope, locks are unchanged. That capacity is read when the plan is built, from the device’s model cache: ModelCache.calibration_capacity, what the calibration’s lock then gets from the same cache. It counts as ModelCache._get_vram_available counts a lock, through the same TorchDevice.memory_available_to_process, within max_cache_vram_gb, less reserve_vram_gb. On a dedicated card that is the driver’s free memory plus what torch has allocated, so the models the pre-lock offload moves out count as available. On a device whose memory is system RAM (MPS, an integrated Intel GPU) it is the RAM the host has available plus what the backend holds for this process, less what the cache’s models hold there now (_get_vram_in_use): the offload frees nothing on such a device, and the lock subtracts them too. The 90% share of its total would overstate what a decode finds there whenever the host is short of free memory. Memory other processes hold is thereby left out. What this process holds outside torch’s allocator counts at the larger of what it holds now and the most a valid stored window of the current environment recorded at its baseline (the driver’s figure for the process less torch’s reserve, recorded_outside_torch_bytes), over the windows of the estimators a user calibration moves only, because it grows with the first generations: on a 16 GB RX 9060 XT under ROCm, HIP’s retained pool and MIOpen’s workspaces took about 0.3 GB in a fresh server and 2.6 GiB once it had decoded, and sized for the fresh reading its largest cases ran partly resident. A first calibration has no stored window and sizes for the reading alone, so on such a card its largest cases may run partly resident; those windows do not count, and the next plan sizes for what the first one recorded. The recorded figure comes from the store refresh_user_fits just read; estimate_untiled_vae_decode prices them with the decode nodes’ own features. DECODE_LINES restates which estimator each base’s decode node prices with; test_working_memory_calibration_plan.py holds it to the nodes in invocations/vae/. Each case runs once to warm up and twice more; estimated_seconds is a rough order of magnitude, not a measurement. Encodes are not planned: text-to-image runs none. Two settings of the machine leave nothing to measure, and the plan is then supported: false with the way out: under force_tiled_decode the SD, FLUX.1, Z-Image and FLUX.2 decode nodes tile every decode, so a case measures no untiled line and no estimate would read one (FORCE_TILED_BASES; the plan does not override the setting for calibration items); and a FLUX.1, FLUX.2 or Z-Image model that decodes with a standalone VAE needs a compatible one that runs on the GPU, since a decode on the CPU opens no window. A main model whose default VAE runs on the CPU gets a vae pick of the first accepted VAE on the GPU where its base has the slot, and is otherwise supported: false (the SD bases have no slot a case could pick through); a standalone SD VAE goes through a compatible main model whose default VAE does not run on the CPU. A FLUX.2 model that is no pipeline needs a text encoder Generate can use, a fitting pipeline’s or a standalone one as above: without one no case compiles, and the plan is supported: false saying which encoder to install. A standalone FLUX.2 VAE goes through the first compatible main model that has one.

Emitting model loading events

When the context argument is passed to load_model_*(), it will retrieve the invocation event bus from the passed InvocationContext object to emit events on the invocation bus. The two events are “model_load_started” and “model_load_completed”. Both carry the following payload:

payload=dict(
  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_key=model_key,
  submodel_type=submodel,
  hash=model_info.hash,
  location=str(model_info.location),
  precision=str(model_info.precision),
)

Adding Model Loaders

Model loaders are small classes that inherit from the ModelLoader base class. They typically implement one method _load_model() whose signature is:

def _load_model(
  self,
  model_path: Path,
  model_variant: Optional[ModelRepoVariant] = None,
  submodel_type: Optional[SubModelType] = None,
) -> AnyModel:

_load_model() will be passed the path to the model on disk, an optional repository variant (used by the diffusers loaders to select, e.g. the fp16 variant, and an optional submodel_type for main and onnx models.

To install a new loader, place it in invokeai/backend/model_manager/load/model_loaders. Inherit from ModelLoader, which checks a cold load against edits of its record (the registry refuses any other class), and use the @ModelLoaderRegistry.register() decorator to indicate what type of models the loader can handle.

Here is a complete example from generic_diffusers.py, which is able to load several different diffusers types:

from pathlib import Path
from typing import Optional

from invokeai.backend.model_manager import (
  AnyModel,
  BaseModelType,
  ModelFormat,
  ModelRepoVariant,
  ModelType,
  SubModelType,
)
from .. import ModelLoader, ModelLoaderRegistry


@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers)
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.T2IAdapter, format=ModelFormat.Diffusers)
class GenericDiffusersLoader(ModelLoader):
    """Class to load simple diffusers models."""

    def _load_model(
        self,
        model_path: Path,
        model_variant: Optional[ModelRepoVariant] = None,
        submodel_type: Optional[SubModelType] = None,
    ) -> AnyModel:
        model_class = self._get_hf_load_class(model_path)
        if submodel_type is not None:
            raise Exception(f"There are no submodels in models of type {model_class}")
        variant = model_variant.value if model_variant else None
        result: AnyModel = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, variant=variant)  # type: ignore
        return result

Note

A loader can register itself to handle several different model types. An exception will be raised if more than one loader tries to register the same model type.

Conversion

Some models require conversion to diffusers format before they can be loaded. These loaders should override two additional methods:

_needs_conversion(self, config: AnyModelConfig, model_path: Path, dest_path: Path) -> bool
_convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Path) -> Path:

The first method accepts the model configuration, the path to where the unmodified model is currently installed, and a proposed destination for the converted model. This method returns True if the model needs to be converted. It typically does this by comparing the last modification time of the original model file to the modification time of the converted model. In some cases you will also want to check the modification date of the configuration record, in the event that the user has changed something like the scheduler prediction type that will require the model to be re-converted. See controlnet.py for an example of this logic.

The second method accepts the model configuration, the path to the original model on disk, and the desired output path for the converted model. It does whatever it needs to do to get the model into diffusers format, and returns the Path of the resulting model. (The path should ordinarily be the same as output_path.)

The ModelManagerService object

For convenience, the API provides a ModelManagerService object which gives a single point of access to the major model manager services. This object is created at initialization time and can be found in the global ApiDependencies.invoker.services.model_manager object, or in context.services.model_manager from within an invocation.

In the examples below, we have retrieved the manager using:

mm = ApiDependencies.invoker.services.model_manager

The following properties and methods will be available:

mm.store

This retrieves the ModelRecordService associated with the manager. Example:

configs = mm.store.get_model_by_attr(name='stable-diffusion-v1-5')

mm.install

This retrieves the ModelInstallService associated with the manager. Example:

job = mm.install.heuristic_import(`https://civitai.com/models/58390/detail-tweaker-lora-lora`)

mm.load

This retrieves the ModelLoaderService associated with the manager. Example:

configs = mm.store.get_model_by_attr(name='stable-diffusion-v1-5')
assert len(configs) > 0

loaded_model = mm.load.load_model(configs[0])

mm.update_model_record(key, changes) -> AnyModelConfig

Applies a ModelRecordChanges to the model’s record, allowing its config class to change, and returns the stored result. When the change affects how the model loads (for example its path, format, cpu_only or fp8_storage), every cached instance of the model is dropped so its next load rebuilds. Raises UnknownModelException for an unknown key and ValueError for a change the record cannot take, including a non-null value for a field neither the current nor the resulting config class has; nothing is written then. Use this rather than mm.store.update_model() for changes that may affect loading.

mm.reidentify_model(key) -> AnyModelConfig

Identifies the model’s files again and replaces its record with the result, keeping the name, description, cover image, trigger phrases, install source and CPU/FP8 preferences. The install source and an IP-Adapter’s image_encoder_model_id are passed to identification, since the files alone may not name them. Cached instances are dropped as for update_model_record(). Raises UnknownModelException for an unknown key and InvalidModelException when the files are not identified as any model.

The model manager also offers a few convenience shortcuts for loading models:

mm.load_model_by_config(model_config, [submodel], [context]) -> LoadedModel

Same as mm.load.load_model().

mm.load_model_by_attr(model_name, base_model, model_type, [submodel], [context]) -> LoadedModel

This accepts the combination of the model’s name, type and base, which it passes to the model record config store for retrieval. If a unique model config is found, this method returns a LoadedModel. It can raise the following exceptions:

• UnknownModelException — model with these attributes not known
• NotImplementedException — the loader doesn’t know how to load this type of model
• ValueError — more than one model matches this combination of base/type/name

mm.load_model_by_key(key, [submodel], [context]) -> LoadedModel

This method takes a model key, looks it up using the ModelRecordServiceBase object in mm.store, and passes the returned model configuration to load_model_by_config(). It may raise a NotImplementedException.

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