diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 78ef965d5a..58950d6692 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -75,7 +75,6 @@ example_model_config = { "description": "string", "source": "string", "last_modified": 0, - "vae": "string", "variant": "normal", "prediction_type": "epsilon", "repo_variant": "fp16", diff --git a/invokeai/backend/model_manager/config.py b/invokeai/backend/model_manager/config.py index 87c2feab5b..5aa07f44ee 100644 --- a/invokeai/backend/model_manager/config.py +++ b/invokeai/backend/model_manager/config.py @@ -260,7 +260,6 @@ class TextualInversionFolderConfig(ModelConfigBase): class _MainConfig(ModelConfigBase): """Model config for main models.""" - vae: Optional[str] = Field(default=None) variant: ModelVariantType = ModelVariantType.Normal prediction_type: SchedulerPredictionType = SchedulerPredictionType.Epsilon upcast_attention: bool = False diff --git a/invokeai/backend/model_manager/load/model_loader_registry.py b/invokeai/backend/model_manager/load/model_loader_registry.py index e7747650d4..4b6a6ff57f 100644 --- a/invokeai/backend/model_manager/load/model_loader_registry.py +++ b/invokeai/backend/model_manager/load/model_loader_registry.py @@ -15,9 +15,7 @@ Use like this: """ -import hashlib from abc import ABC, abstractmethod -from pathlib import Path from typing import Callable, Dict, Optional, Tuple, Type from ..config import ( @@ -27,8 +25,6 @@ from ..config import ( ModelFormat, ModelType, SubModelType, - VaeCheckpointConfig, - VaeDiffusersConfig, ) from . import ModelLoaderBase @@ -90,33 +86,15 @@ class ModelLoaderRegistry: cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] ) -> Tuple[Type[ModelLoaderBase], ModelConfigBase, Optional[SubModelType]]: """Get subclass of ModelLoaderBase registered to handle base and type.""" - # We have to handle VAE overrides here because this will change the model type and the corresponding implementation returned - conf2, submodel_type = cls._handle_subtype_overrides(config, submodel_type) - key1 = cls._to_registry_key(conf2.base, conf2.type, conf2.format) # for a specific base type - key2 = cls._to_registry_key(BaseModelType.Any, conf2.type, conf2.format) # with wildcard Any + key1 = cls._to_registry_key(config.base, config.type, config.format) # for a specific base type + key2 = cls._to_registry_key(BaseModelType.Any, config.type, config.format) # with wildcard Any implementation = cls._registry.get(key1) or cls._registry.get(key2) if not implementation: raise NotImplementedError( f"No subclass of LoadedModel is registered for base={config.base}, type={config.type}, format={config.format}" ) - return implementation, conf2, submodel_type - - @classmethod - def _handle_subtype_overrides( - cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] - ) -> Tuple[ModelConfigBase, Optional[SubModelType]]: - if submodel_type == SubModelType.Vae and hasattr(config, "vae") and config.vae is not None: - model_path = Path(config.vae) - config_class = ( - VaeCheckpointConfig if model_path.suffix in [".pt", ".safetensors", ".ckpt"] else VaeDiffusersConfig - ) - hash = hashlib.md5(model_path.as_posix().encode("utf-8")).hexdigest() - new_conf = config_class(path=model_path.as_posix(), name=model_path.stem, base=config.base, key=hash) - submodel_type = None - else: - new_conf = config - return new_conf, submodel_type + return implementation, config, submodel_type @staticmethod def _to_registry_key(base: BaseModelType, type: ModelType, format: ModelFormat) -> str: