Reimplement model config storage backend, download and install modules

invokeai/backend/model_management2/model_config.py

@@ -0,0 +1,276 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+"""
+Configuration definitions for image generation models.
+
+Typical usage:
+
+  from invokeai.backend.model_management2.model_config 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',
+             model_variant='normal',
+             model_format='checkpoint'
+            )
+  config = ModelConfigFactory.make_config(raw)
+  print(config.name)
+
+Validation errors will raise an InvalidModelConfigException error.
+
+"""
+import pydantic
+from enum import Enum
+from pydantic import BaseModel, Field, Extra
+from typing import Optional, Literal, List, Union, Type
+from pydantic.error_wrappers import ValidationError
+from omegaconf.listconfig import ListConfig  # to support the yaml backend
+
+
+class InvalidModelConfigException(Exception):
+    """Exception raised when the config parser doesn't recognize the passed
+    combination of model type and format."""
+
+    pass
+
+
+class BaseModelType(str, Enum):
+    """Base model type."""
+
+    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"
+
+
+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"
+
+
+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
+    model_type: ModelType
+    model_format: ModelFormat
+    id: Optional[str] = Field(None)  # this may get added by the store
+    description: Optional[str] = Field(None)
+    author: Optional[str] = Field(description="Model author")
+    thumbnail_url: Optional[str] = Field(description="URL of thumbnail image")
+    license_url: Optional[str] = Field(description="URL of license")
+    source_url: Optional[str] = Field(description="Model download source")
+    tags: Optional[List[str]] = Field(description="Descriptive tags")  # Set would be better, but not JSON serializable
+
+    class Config:
+        """Pydantic configuration hint."""
+
+        use_enum_values = True
+        extra = Extra.forbid
+        validate_assignment = True
+
+    @pydantic.validator("tags", pre=True)
+    @classmethod
+    def _fix_tags(cls, v):
+        if isinstance(v, ListConfig):  # to support yaml backend
+            v = list(v)
+        return v
+
+
+class CheckpointConfig(ModelConfigBase):
+    """Model config for checkpoint-style models."""
+
+    model_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."""
+
+    model_format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+
+
+class LoRAConfig(ModelConfigBase):
+    """Model config for LoRA/Lycoris models."""
+
+    model_format: Literal[ModelFormat.Lycoris, ModelFormat.Diffusers]
+
+
+class VaeCheckpointConfig(ModelConfigBase):
+    """Model config for standalone VAE models."""
+
+    model_format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
+
+
+class VaeDiffusersConfig(ModelConfigBase):
+    """Model config for standalone VAE models (diffusers version)."""
+
+    model_format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+
+
+class TextualInversionConfig(ModelConfigBase):
+    """Model config for textual inversion embeddings."""
+
+    model_format: Literal[ModelFormat.EmbeddingFile, ModelFormat.EmbeddingFolder]
+
+
+class MainConfig(ModelConfigBase):
+    """Model config for main models."""
+
+    vae: Optional[str] = Field(None)
+    model_variant: ModelVariantType = ModelVariantType.Normal
+
+
+class MainCheckpointConfig(CheckpointConfig, MainConfig):
+    """Model config for main checkpoint models."""
+
+    pass
+
+
+class MainDiffusersConfig(DiffusersConfig, MainConfig):
+    """Model config for main diffusers models."""
+
+    pass
+
+
+class ONNXSD1Config(MainConfig):
+    """Model config for ONNX format models based on sd-1."""
+
+    model_format: Literal[ModelFormat.Onnx, ModelFormat.Olive]
+
+
+class ONNXSD2Config(MainConfig):
+    """Model config for ONNX format models based on sd-2."""
+
+    model_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 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,
+        },
+        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,
+        },
+    }
+
+    @classmethod
+    def make_config(
+        cls,
+        model_data: Union[dict, ModelConfigBase],
+        dest_class: Optional[Type] = None,
+    ) -> Union[
+        MainCheckpointConfig,
+        MainDiffusersConfig,
+        LoRAConfig,
+        TextualInversionConfig,
+        ONNXSD1Config,
+        ONNXSD2Config,
+    ]:
+        """
+        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):
+            return model_data
+        try:
+            model_format = model_data.get("model_format")
+            model_type = model_data.get("model_type")
+            model_base = model_data.get("base_model")
+            class_to_return = dest_class or cls._class_map[model_format][model_type]
+            if isinstance(class_to_return, dict):  # additional level allowed
+                class_to_return = class_to_return[model_base]
+            return class_to_return.parse_obj(model_data)
+        except KeyError:
+            raise InvalidModelConfigException(
+                f"Unknown combination of model_format '{model_format}' and model_type '{model_type}'"
+            )
+        except ValidationError as e:
+            raise InvalidModelConfigException(f"Invalid model configuration passed: {str(e)}") from e

invokeai/backend/model_management2/storage/base.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+"""
+Abstract base class for storing and retrieving model configuration records.
+"""
+
+
+from abc import ABC, abstractmethod
+from typing import Union, Set, List, Optional
+
+from ..model_config import ModelConfigBase, BaseModelType, ModelType
+
+
+class DuplicateModelException(Exception):
+    """Raised on an attempt to add a model with the same key twice."""
+
+    pass
+
+
+class UnknownModelException(Exception):
+    """Raised on an attempt to delete a model with a nonexistent key."""
+
+    pass
+
+
+class ModelConfigStore(ABC):
+    """Abstract base class for storage and retrieval of model configs."""
+
+    @abstractmethod
+    def add_model(self, key: str, config: Union[dict, ModelConfigBase]) -> None:
+        """
+        Add a model to the database.
+
+        :param key: Unique key for the model
+        :param config: Model configuration record, either a dict with the
+         required fields or a ModelConfigBase instance.
+
+        Can raise DuplicateModelException and InvalidModelConfig exceptions.
+        """
+        pass
+
+    @abstractmethod
+    def del_model(self, key: str) -> None:
+        """
+        Delete a model.
+
+        :param key: Unique key for the model to be deleted
+
+        Can raise an UnknownModelException
+        """
+        pass
+
+    @abstractmethod
+    def update_model(self, key: str, config: Union[dict, ModelConfigBase]) -> ModelConfigBase:
+        """
+        Update the model, returning the updated version.
+
+        :param key: Unique key for the model to be updated
+        :param config: Model configuration record. Either a dict with the
+         required fields, or a ModelConfigBase instance.
+        """
+        pass
+
+    @abstractmethod
+    def get_model(self, key: str) -> ModelConfigBase:
+        """
+        Retrieve the ModelConfigBase instance for the indicated model.
+
+        :param key: Key of model config to be fetched.
+
+        Exceptions: UnknownModelException
+        """
+        pass
+
+    @abstractmethod
+    def exists(self, key: str) -> bool:
+        """
+        Return True if a model with the indicated key exists in the databse.
+
+        :param key: Unique key for the model to be deleted
+        """
+        pass
+
+    @abstractmethod
+    def search_by_tag(self, tags: Set[str]) -> List[ModelConfigBase]:
+        """
+        Return models containing all of the listed tags.
+
+        :param tags: Set of tags to search on.
+        """
+        pass
+
+    @abstractmethod
+    def search_by_type(
+        self,
+        model_name: Optional[str] = None,
+        base_model: Optional[BaseModelType] = None,
+        model_type: Optional[ModelType] = None,
+    ) -> List[ModelConfigBase]:
+        """
+        Return models matching name, base and/or type.
+
+        :param model_name: Filter by name of model (optional)
+        :param base_model: Filter by base model (optional)
+        :param model_type: Filter by type of model (optional)
+
+        If none of the optional filters are passed, will return all
+        models in the database.
+        """
+        pass
+
+    def all_models(self) -> List[ModelConfigBase]:
+        """
+        Return all the model configs in the database.
+        """
+        return self.search_by_type()