add sql-based model config store and api

invokeai/backend/model_manager/config.py

@@ -0,0 +1,319 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+"""
+Configuration definitions for image generation models.
+
+Typical usage:
+
+  from invokeai.backend.model_manager import ModelConfigFactory
+  raw = dict(path='models/sd-1/main/foo.ckpt',
+             name='foo',
+             base_model='sd-1',
+             model_type='main',
+             config='configs/stable-diffusion/v1-inference.yaml',
+             variant='normal',
+             format='checkpoint'
+            )
+  config = ModelConfigFactory.make_config(raw)
+  print(config.name)
+
+Validation errors will raise an InvalidModelConfigException error.
+
+"""
+from enum import Enum
+from typing import Literal, Optional, Type, Union
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+
+class InvalidModelConfigException(Exception):
+    """Exception for when config parser doesn't recognized this combination of model type and format."""
+
+
+class BaseModelType(str, Enum):
+    """Base model type."""
+
+    Any = "any"
+    StableDiffusion1 = "sd-1"
+    StableDiffusion2 = "sd-2"
+    StableDiffusionXL = "sdxl"
+    StableDiffusionXLRefiner = "sdxl-refiner"
+    # Kandinsky2_1 = "kandinsky-2.1"
+
+
+class ModelType(str, Enum):
+    """Model type."""
+
+    ONNX = "onnx"
+    Main = "main"
+    Vae = "vae"
+    Lora = "lora"
+    ControlNet = "controlnet"  # used by model_probe
+    TextualInversion = "embedding"
+    IPAdapter = "ip_adapter"
+    CLIPVision = "clip_vision"
+    T2IAdapter = "t2i_adapter"
+
+
+class SubModelType(str, Enum):
+    """Submodel type."""
+
+    UNet = "unet"
+    TextEncoder = "text_encoder"
+    TextEncoder2 = "text_encoder_2"
+    Tokenizer = "tokenizer"
+    Tokenizer2 = "tokenizer_2"
+    Vae = "vae"
+    VaeDecoder = "vae_decoder"
+    VaeEncoder = "vae_encoder"
+    Scheduler = "scheduler"
+    SafetyChecker = "safety_checker"
+
+
+class ModelVariantType(str, Enum):
+    """Variant type."""
+
+    Normal = "normal"
+    Inpaint = "inpaint"
+    Depth = "depth"
+
+
+class ModelFormat(str, Enum):
+    """Storage format of model."""
+
+    Diffusers = "diffusers"
+    Checkpoint = "checkpoint"
+    Lycoris = "lycoris"
+    Onnx = "onnx"
+    Olive = "olive"
+    EmbeddingFile = "embedding_file"
+    EmbeddingFolder = "embedding_folder"
+    InvokeAI = "invokeai"
+
+
+class SchedulerPredictionType(str, Enum):
+    """Scheduler prediction type."""
+
+    Epsilon = "epsilon"
+    VPrediction = "v_prediction"
+    Sample = "sample"
+
+
+class ModelConfigBase(BaseModel):
+    """Base class for model configuration information."""
+
+    path: str
+    name: str
+    base_model: BaseModelType
+    type: ModelType
+    format: ModelFormat
+    key: str = Field(description="unique key for model", default="<NOKEY>")
+    original_hash: Optional[str] = Field(
+        description="original fasthash of model contents", default=None
+    )  # this is assigned at install time and will not change
+    current_hash: Optional[str] = Field(
+        description="current fasthash of model contents", default=None
+    )  # if model is converted or otherwise modified, this will hold updated hash
+    description: Optional[str] = Field(None)
+    source: Optional[str] = Field(description="Model download source (URL or repo_id)", default=None)
+
+    model_config = ConfigDict(
+        use_enum_values=False,
+        validate_assignment=True,
+    )
+
+    def update(self, attributes: dict):
+        """Update the object with fields in dict."""
+        for key, value in attributes.items():
+            setattr(self, key, value)  # may raise a validation error
+
+
+class CheckpointConfig(ModelConfigBase):
+    """Model config for checkpoint-style models."""
+
+    format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
+    config: str = Field(description="path to the checkpoint model config file")
+
+
+class DiffusersConfig(ModelConfigBase):
+    """Model config for diffusers-style models."""
+
+    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+
+
+class LoRAConfig(ModelConfigBase):
+    """Model config for LoRA/Lycoris models."""
+
+    format: Literal[ModelFormat.Lycoris, ModelFormat.Diffusers]
+
+
+class VaeCheckpointConfig(ModelConfigBase):
+    """Model config for standalone VAE models."""
+
+    format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
+
+
+class VaeDiffusersConfig(ModelConfigBase):
+    """Model config for standalone VAE models (diffusers version)."""
+
+    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+
+
+class ControlNetDiffusersConfig(DiffusersConfig):
+    """Model config for ControlNet models (diffusers version)."""
+
+    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+
+
+class ControlNetCheckpointConfig(CheckpointConfig):
+    """Model config for ControlNet models (diffusers version)."""
+
+    format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
+
+
+class TextualInversionConfig(ModelConfigBase):
+    """Model config for textual inversion embeddings."""
+
+    format: Literal[ModelFormat.EmbeddingFile, ModelFormat.EmbeddingFolder]
+
+
+class MainConfig(ModelConfigBase):
+    """Model config for main models."""
+
+    vae: Optional[str] = Field(None)
+    variant: ModelVariantType = ModelVariantType.Normal
+
+
+class MainCheckpointConfig(CheckpointConfig, MainConfig):
+    """Model config for main checkpoint models."""
+
+
+class MainDiffusersConfig(DiffusersConfig, MainConfig):
+    """Model config for main diffusers models."""
+
+
+class ONNXSD1Config(MainConfig):
+    """Model config for ONNX format models based on sd-1."""
+
+    format: Literal[ModelFormat.Onnx, ModelFormat.Olive]
+
+
+class ONNXSD2Config(MainConfig):
+    """Model config for ONNX format models based on sd-2."""
+
+    format: Literal[ModelFormat.Onnx, ModelFormat.Olive]
+    # No yaml config file for ONNX, so these are part of config
+    prediction_type: SchedulerPredictionType
+    upcast_attention: bool
+
+
+class IPAdapterConfig(ModelConfigBase):
+    """Model config for IP Adaptor format models."""
+
+    format: Literal[ModelFormat.InvokeAI]
+
+
+class CLIPVisionDiffusersConfig(ModelConfigBase):
+    """Model config for ClipVision."""
+
+    format: Literal[ModelFormat.Diffusers]
+
+
+class T2IConfig(ModelConfigBase):
+    """Model config for T2I."""
+
+    format: Literal[ModelFormat.Diffusers]
+
+
+AnyModelConfig = Union[
+    MainCheckpointConfig,
+    MainDiffusersConfig,
+    LoRAConfig,
+    TextualInversionConfig,
+    ONNXSD1Config,
+    ONNXSD2Config,
+    VaeCheckpointConfig,
+    VaeDiffusersConfig,
+    ControlNetDiffusersConfig,
+    ControlNetCheckpointConfig,
+    IPAdapterConfig,
+    CLIPVisionDiffusersConfig,
+    T2IConfig,
+]
+
+
+class ModelConfigFactory(object):
+    """Class for parsing config dicts into StableDiffusion Config obects."""
+
+    _class_map: dict = {
+        ModelFormat.Checkpoint: {
+            ModelType.Main: MainCheckpointConfig,
+            ModelType.Vae: VaeCheckpointConfig,
+        },
+        ModelFormat.Diffusers: {
+            ModelType.Main: MainDiffusersConfig,
+            ModelType.Lora: LoRAConfig,
+            ModelType.Vae: VaeDiffusersConfig,
+            ModelType.ControlNet: ControlNetDiffusersConfig,
+            ModelType.CLIPVision: CLIPVisionDiffusersConfig,
+        },
+        ModelFormat.Lycoris: {
+            ModelType.Lora: LoRAConfig,
+        },
+        ModelFormat.Onnx: {
+            ModelType.ONNX: {
+                BaseModelType.StableDiffusion1: ONNXSD1Config,
+                BaseModelType.StableDiffusion2: ONNXSD2Config,
+            },
+        },
+        ModelFormat.Olive: {
+            ModelType.ONNX: {
+                BaseModelType.StableDiffusion1: ONNXSD1Config,
+                BaseModelType.StableDiffusion2: ONNXSD2Config,
+            },
+        },
+        ModelFormat.EmbeddingFile: {
+            ModelType.TextualInversion: TextualInversionConfig,
+        },
+        ModelFormat.EmbeddingFolder: {
+            ModelType.TextualInversion: TextualInversionConfig,
+        },
+        ModelFormat.InvokeAI: {
+            ModelType.IPAdapter: IPAdapterConfig,
+        },
+    }
+
+    @classmethod
+    def make_config(
+        cls,
+        model_data: Union[dict, ModelConfigBase],
+        key: Optional[str] = None,
+        dest_class: Optional[Type] = None,
+    ) -> AnyModelConfig:
+        """
+        Return the appropriate config object from raw dict values.
+
+        :param model_data: A raw dict corresponding the obect fields to be
+        parsed into a ModelConfigBase obect (or descendent), or a ModelConfigBase
+        object, which will be passed through unchanged.
+        :param dest_class: The config class to be returned. If not provided, will
+        be selected automatically.
+        """
+        if isinstance(model_data, ModelConfigBase):
+            if key:
+                model_data.key = key
+            return model_data
+        try:
+            format = model_data.get("format")
+            type = model_data.get("type")
+            model_base = model_data.get("base_model")
+            class_to_return = dest_class or cls._class_map[format][type]
+            if isinstance(class_to_return, dict):  # additional level allowed
+                class_to_return = class_to_return[model_base]
+            model = class_to_return.model_validate(model_data)
+            if key:
+                model.key = key  # ensure consistency
+            return model
+        except KeyError as exc:
+            raise InvalidModelConfigException(f"Unknown combination of format '{format}' and type '{type}'") from exc
+        except ValidationError as exc:
+            raise InvalidModelConfigException(f"Invalid model configuration passed: {str(exc)}") from exc

invokeai/backend/model_manager/hash.py

@@ -0,0 +1,66 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+"""
+Fast hashing of diffusers and checkpoint-style models.
+
+Usage:
+from invokeai.backend.model_managre.model_hash import FastModelHash
+>>> FastModelHash.hash('/home/models/stable-diffusion-v1.5')
+'a8e693a126ea5b831c96064dc569956f'
+"""
+
+import hashlib
+import os
+from pathlib import Path
+from typing import Dict, Union
+
+from imohash import hashfile
+
+
+class FastModelHash(object):
+    """FastModelHash obect provides one public class method, hash()."""
+
+    @classmethod
+    def hash(cls, model_location: Union[str, Path]) -> str:
+        """
+        Return hexdigest string for model located at model_location.
+
+        :param model_location: Path to the model
+        """
+        model_location = Path(model_location)
+        if model_location.is_file():
+            return cls._hash_file(model_location)
+        elif model_location.is_dir():
+            return cls._hash_dir(model_location)
+        else:
+            raise OSError(f"Not a valid file or directory: {model_location}")
+
+    @classmethod
+    def _hash_file(cls, model_location: Union[str, Path]) -> str:
+        """
+        Fasthash a single file and return its hexdigest.
+
+        :param model_location: Path to the model file
+        """
+        # we return md5 hash of the filehash to make it shorter
+        # cryptographic security not needed here
+        return hashlib.md5(hashfile(model_location)).hexdigest()
+
+    @classmethod
+    def _hash_dir(cls, model_location: Union[str, Path]) -> str:
+        components: Dict[str, str] = {}
+
+        for root, dirs, files in os.walk(model_location):
+            for file in files:
+                # only tally tensor files because diffusers config files change slightly
+                # depending on how the model was downloaded/converted.
+                if not file.endswith((".ckpt", ".safetensors", ".bin", ".pt", ".pth")):
+                    continue
+                path = (Path(root) / file).as_posix()
+                fast_hash = cls._hash_file(path)
+                components.update({path: fast_hash})
+
+        # hash all the model hashes together, using alphabetic file order
+        md5 = hashlib.md5()
+        for path, fast_hash in sorted(components.items()):
+            md5.update(fast_hash.encode("utf-8"))
+        return md5.hexdigest()

invokeai/backend/model_manager/migrate_to_db.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2023 Lincoln D. Stein
+"""Migrate from the InvokeAI v2 models.yaml format to the v3 sqlite format."""
+
+from hashlib import sha1
+
+from omegaconf import DictConfig, OmegaConf
+from pydantic import TypeAdapter
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.model_records import DuplicateModelException, ModelRecordServiceSQL
+from invokeai.app.services.shared.sqlite import SqliteDatabase
+from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelType
+from invokeai.backend.model_manager.hash import FastModelHash
+from invokeai.backend.util.logging import InvokeAILogger
+
+ModelsValidator = TypeAdapter(AnyModelConfig)
+
+
+class Migrate:
+    """Migration class."""
+
+    config: InvokeAIAppConfig
+    logger: InvokeAILogger
+
+    def __init__(self):
+        self.config = InvokeAIAppConfig.get_config()
+        self.config.parse_args()
+        self.logger = InvokeAILogger.get_logger()
+
+    def get_db(self) -> ModelRecordServiceSQL:
+        db = SqliteDatabase(self.config, self.logger)
+        return ModelRecordServiceSQL(db)
+
+    def get_yaml(self) -> DictConfig:
+        yaml_path = self.config.model_conf_path
+        return OmegaConf.load(yaml_path)
+
+    def migrate(self):
+        db = self.get_db()
+        yaml = self.get_yaml()
+
+        for model_key, stanza in yaml.items():
+            if model_key == "__metadata__":
+                assert (
+                    stanza["version"] == "3.0.0"
+                ), f"This script works on version 3.0.0 yaml files, but your configuration points to a {stanza['version']} version"
+                continue
+
+            base_type, model_type, model_name = str(model_key).split("/")
+            hash = FastModelHash.hash(self.config.models_path / stanza.path)
+            new_key = sha1(model_key.encode("utf-8")).hexdigest()
+
+            stanza["base_model"] = BaseModelType(base_type)
+            stanza["type"] = ModelType(model_type)
+            stanza["name"] = model_name
+            stanza["original_hash"] = hash
+            stanza["current_hash"] = hash
+
+            new_config = ModelsValidator.validate_python(stanza)
+            self.logger.info(f"Adding model {model_name} with key {model_key}")
+            try:
+                db.add_model(new_key, new_config)
+            except DuplicateModelException:
+                self.logger.warning(f"Model {model_name} is already in the database")
+
+
+def main():
+    Migrate().migrate()
+
+
+if __name__ == "__main__":
+    main()