InvokeAI/tests/model_identification

Model Probe (Identification) Testing

Invoke's model Identification system is tested against example model files. Test cases are lightweight representations of real models which have been "stripped" of their tensor data.

Setup

Test cases are stored with git lfs. You must install git lfs to pull down the test cases and add to them.

# Only need to do this once
git lfs install
# Pull the actual model files down - if you just do `git pull` you'll only get pointers
git lfs pull

Running the Tests

To run the tests use:

pytest -v tests/test_model_probe/test_identification.py

Stripped Model Files

Invoke abstracts the loading of a model's state dict and metadata in a class called ModelOnDisk. This class loads real model weights. We use it to inspect models and identify them.

For testing purposes, we create a stripped-down version of model weights that contain only the model structure and metadata for each key, without the actual tensor data. The state dict structure is typically all we need to identify models; the tensors themselves are not needed. This allows us to store test cases in the repo without adding many gigabytes of data.

To see how this works, check out StrippedModelOnDisk. This class includes logic to strip models and to load these stripped models for testing.

Some Models Cannot Be Stripped

Certain models cannot be stripped because identification relies on inspecting the actual tensor data. We have to store the full model files for these test cases.

Currently, the only models that cannot be stripped are spandrel image-to-image models. spandrel supports many model architectures but doesn't provide a way to identify or assert support for a model by its state dict structure alone.

To positively identify these models, we must attempt to load the model using spandrel. If it loads successfully, we assume it is a supported model. Therefore, we cannot strip these models and must store the full model files in the test cases. We only store one such model to keep the test suite size manageable.

StrippedModelOnDisk will simply pass-through the "live" tensor data for these models when loading them to test.

Adding New Test Cases

Run the strip_model.py script to create a new test case. For example:

python strip_model.py /path/to/your/model --output_dir ./stripped_models

It supports single-file models and multi-file models (e.g. diffusers-style models). The output will be a directory named with a UUID, containing the stripped model files and a dummy __test_metadata__.json file.

Example output structure for a single-file model:

stripped_models/
└── 19fd1a40-c5b7-4734-bd3a-6e0e948cce0b/
    ├── __test_metadata__.json
    └── Standard Reference (XLabs FLUX IP-Adapter v2).safetensors

This test metadata file should contain a single JSON dict and must be filled out manually with the expected identification results.

Structure of __test_metadata__.json

This file contains a single JSON dict. Here's an example for a FLUX IP Adapter checkpoint:

{
  "source": "https://huggingface.co/XLabs-AI/flux-ip-adapter-v2/resolve/main/ip_adapter.safetensors",
  "file_name": "Standard Reference (XLabs FLUX IP-Adapter v2).safetensors",
  "expected_config_attrs": {
    "type": "ip_adapter",
    "format": "checkpoint",
    "base": "flux"
  }
}

See the details below for each field.

"source"

A string indicating the source of the model (e.g. a Hugging Face repo ID or URL). This is not used for identification, but is useful for reference so we know where the model came from. Nothing will break if this field is missing or incorrect, but it is good practice to fill it out.

• Example HF Repo ID: "RunDiffusion/Juggernaut-XL-v9"
• Example URL: "https://huggingface.co/XpucT/Deliberate/resolve/main/Deliberate_v5.safetensors"

"file_name"

If the model is a single file (e.g. a .safetensors file), this is the name of that file. The test suite will look for this file in the test case directory.

If the model is multi-file (e.g. diffusers-style), omit this key or set it to a falsey value like null or an empty string.

• Example: "model.safetensors"

The strip_model.py script will automatically fill this field in for single-file models.

"expected_config_attrs"

This field is a dict of expected configuration attributes for the model. It is required for all test cases.

It is used to verify that the model's configuration matches expectations. The keys and values in this dict depend on the specific model and its configuration.

These attributes must be included, as they are the primary discriminators for models:

• "type": The type of the model. This is the value of the ModelType enum.
• "format": The format of the model files. This is the value of the ModelFormat enum.
• "base": The base model pipeline architecture associated with this model. Many models do not have an associated base. For these, use "any". This is the value of the BaseModelType enum.

Depending on the kind of model, these additional keys may be useful:

• "prediction_type": The prediction type used by the model. This is the value of the SchedulerPredictionType enum.
• "variant": The variant of the model, if applicable. This is the value of the ModelVariantType enum.

To see all possible values for these enums, check out their definitions in invokeai/backend/model_manager/taxonomy.py.

For example, for a SD1.5 main (pipeline) inpainting model in diffusers format, you might have:

{
  "expected_config_attrs": {
    "type": "main",
    "format": "diffusers",
    "base": "sd-1",
    "prediction_type": "epsilon",
    "variant": "inpaint"
  }
}

"notes"

This is an optional string field where you can add any notes or comments about the test case. It can be useful for providing context or explaining any special considerations.

"override_fields"

In some rare cases, we may need to provide additional hints to the identification system to help it identify the model correctly.

Currently, the only known case where we need extra information is to differentiate between single-file SD1.x, SD2.x and SDXL VAEs. These models have identical structures, so we need to provide a hint. Though it is far from ideal, we use simple string matching on the model's name to provide this hint.

For example, when users install the taesdxl VAE from the HF repo madebyollin/taesdxl, the identification system will get the model name taesdxl. It sees "xl" in the name and infers that this is a SDXL VAE. To reproduce this in a test case, we add the following to __test_metadata__.json:

{
  "override_fields": {
    "name": "taesdxl"
  }
}