diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index ec260c7684..c3e5e50c4b 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -3,9 +3,9 @@ import io import pathlib -import shutil import traceback from copy import deepcopy +from tempfile import TemporaryDirectory from typing import Any, Dict, List, Optional, Type from fastapi import Body, Path, Query, Response, UploadFile @@ -19,7 +19,6 @@ from typing_extensions import Annotated from invokeai.app.services.model_images.model_images_common import ModelImageFileNotFoundException from invokeai.app.services.model_install.model_install_common import ModelInstallJob from invokeai.app.services.model_records import ( - DuplicateModelException, InvalidModelException, ModelRecordChanges, UnknownModelException, @@ -30,7 +29,6 @@ from invokeai.backend.model_manager.config import ( MainCheckpointConfig, ModelFormat, ModelType, - SubModelType, ) from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch from invokeai.backend.model_manager.metadata.metadata_base import ModelMetadataWithFiles, UnknownMetadataException @@ -174,18 +172,6 @@ async def get_model_record( raise HTTPException(status_code=404, detail=str(e)) -# @model_manager_router.get("/summary", operation_id="list_model_summary") -# async def list_model_summary( -# page: int = Query(default=0, description="The page to get"), -# per_page: int = Query(default=10, description="The number of models per page"), -# order_by: ModelRecordOrderBy = Query(default=ModelRecordOrderBy.Default, description="The attribute to order by"), -# ) -> PaginatedResults[ModelSummary]: -# """Gets a page of model summary data.""" -# record_store = ApiDependencies.invoker.services.model_manager.store -# results: PaginatedResults[ModelSummary] = record_store.list_models(page=page, per_page=per_page, order_by=order_by) -# return results - - class FoundModel(BaseModel): path: str = Field(description="Path to the model") is_installed: bool = Field(description="Whether or not the model is already installed") @@ -619,34 +605,38 @@ async def convert_model( logger.error(f"The model with key {key} is not a main checkpoint model.") raise HTTPException(400, f"The model with key {key} is not a main checkpoint model.") - cache_path = loader.convert_cache.cache_path(key) - converted_model = loader.load_model(model_config, submodel_type=SubModelType.Scheduler) - # write the converted file to the model cache directory - raw_model = converted_model.model - assert hasattr(raw_model, 'save_pretrained') - raw_model.save_pretrained(cache_path) - assert cache_path.exists() + with TemporaryDirectory(dir=ApiDependencies.invoker.services.configuration.models_path) as tmpdir: + convert_path = pathlib.Path(tmpdir) / pathlib.Path(model_config.path).stem + print(f"DEBUG: convert_path={convert_path}") + converted_model = loader.load_model(model_config) + # write the converted file to the convert path + raw_model = converted_model.model + print(f"DEBUG: raw_model = {raw_model}") + assert hasattr(raw_model, "save_pretrained") + raw_model.save_pretrained(convert_path) + assert convert_path.exists() - # temporarily rename the original safetensors file so that there is no naming conflict - original_name = model_config.name - model_config.name = f"{original_name}.DELETE" - changes = ModelRecordChanges(name=model_config.name) - store.update_model(key, changes=changes) + # temporarily rename the original safetensors file so that there is no naming conflict + original_name = model_config.name + model_config.name = f"{original_name}.DELETE" + changes = ModelRecordChanges(name=model_config.name) + store.update_model(key, changes=changes) - # install the diffusers - try: - new_key = installer.install_path( - cache_path, - config={ - "name": original_name, - "description": model_config.description, - "hash": model_config.hash, - "source": model_config.source, - }, - ) - except DuplicateModelException as e: - logger.error(str(e)) - raise HTTPException(status_code=409, detail=str(e)) + # install the diffusers + try: + new_key = installer.install_path( + convert_path, + config={ + "name": original_name, + "description": model_config.description, + "hash": model_config.hash, + "source": model_config.source, + }, + ) + except Exception as e: + logger.error(str(e)) + store.update_model(key, changes=ModelRecordChanges(name=original_name)) + raise HTTPException(status_code=409, detail=str(e)) # Update the model image if the model had one try: @@ -659,8 +649,8 @@ async def convert_model( # delete the original safetensors file installer.delete(key) - # delete the cached version - shutil.rmtree(cache_path) + # delete the temporary directory + # shutil.rmtree(cache_path) # return the config record for the new diffusers directory new_config = store.get_model(new_key) diff --git a/invokeai/app/services/model_load/model_load_base.py b/invokeai/app/services/model_load/model_load_base.py index 9d75aafde1..a2328a98a5 100644 --- a/invokeai/app/services/model_load/model_load_base.py +++ b/invokeai/app/services/model_load/model_load_base.py @@ -6,7 +6,6 @@ from typing import Optional from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType from invokeai.backend.model_manager.load import LoadedModel -from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase @@ -26,8 +25,3 @@ class ModelLoadServiceBase(ABC): @abstractmethod def ram_cache(self) -> ModelCacheBase[AnyModel]: """Return the RAM cache used by this loader.""" - - @property - @abstractmethod - def convert_cache(self) -> ModelConvertCacheBase: - """Return the checkpoint convert cache used by this loader.""" diff --git a/invokeai/app/services/model_load/model_load_default.py b/invokeai/app/services/model_load/model_load_default.py index e9f527bc86..85440fe3ca 100644 --- a/invokeai/app/services/model_load/model_load_default.py +++ b/invokeai/app/services/model_load/model_load_default.py @@ -11,7 +11,6 @@ from invokeai.backend.model_manager.load import ( ModelLoaderRegistry, ModelLoaderRegistryBase, ) -from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase from invokeai.backend.util.logging import InvokeAILogger @@ -25,7 +24,6 @@ class ModelLoadService(ModelLoadServiceBase): self, app_config: InvokeAIAppConfig, ram_cache: ModelCacheBase[AnyModel], - convert_cache: ModelConvertCacheBase, registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry, ): """Initialize the model load service.""" @@ -34,7 +32,6 @@ class ModelLoadService(ModelLoadServiceBase): self._logger = logger self._app_config = app_config self._ram_cache = ram_cache - self._convert_cache = convert_cache self._registry = registry def start(self, invoker: Invoker) -> None: @@ -45,11 +42,6 @@ class ModelLoadService(ModelLoadServiceBase): """Return the RAM cache used by this loader.""" return self._ram_cache - @property - def convert_cache(self) -> ModelConvertCacheBase: - """Return the checkpoint convert cache used by this loader.""" - return self._convert_cache - def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel: """ Given a model's configuration, load it and return the LoadedModel object. @@ -68,7 +60,6 @@ class ModelLoadService(ModelLoadServiceBase): app_config=self._app_config, logger=self._logger, ram_cache=self._ram_cache, - convert_cache=self._convert_cache, ).load_model(model_config, submodel_type) if hasattr(self, "_invoker"): diff --git a/invokeai/app/services/model_manager/model_manager_default.py b/invokeai/app/services/model_manager/model_manager_default.py index 1a2b9a3402..f695c3c8c1 100644 --- a/invokeai/app/services/model_manager/model_manager_default.py +++ b/invokeai/app/services/model_manager/model_manager_default.py @@ -7,7 +7,7 @@ import torch from typing_extensions import Self from invokeai.app.services.invoker import Invoker -from invokeai.backend.model_manager.load import ModelCache, ModelConvertCache, ModelLoaderRegistry +from invokeai.backend.model_manager.load import ModelCache, ModelLoaderRegistry from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger @@ -86,11 +86,9 @@ class ModelManagerService(ModelManagerServiceBase): logger=logger, execution_device=execution_device or TorchDevice.choose_torch_device(), ) - convert_cache = ModelConvertCache(cache_path=app_config.convert_cache_path, max_size=app_config.convert_cache) loader = ModelLoadService( app_config=app_config, ram_cache=ram_cache, - convert_cache=convert_cache, registry=ModelLoaderRegistry, ) installer = ModelInstallService( diff --git a/invokeai/app/services/shared/sqlite/sqlite_util.py b/invokeai/app/services/shared/sqlite/sqlite_util.py index cadf09f457..05e430c4db 100644 --- a/invokeai/app/services/shared/sqlite/sqlite_util.py +++ b/invokeai/app/services/shared/sqlite/sqlite_util.py @@ -13,6 +13,7 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_7 import from invokeai.app.services.shared.sqlite_migrator.migrations.migration_8 import build_migration_8 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_9 import build_migration_9 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_10 import build_migration_10 +from invokeai.app.services.shared.sqlite_migrator.migrations.migration_11 import build_migration_11 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator @@ -43,6 +44,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto migrator.register_migration(build_migration_8(app_config=config)) migrator.register_migration(build_migration_9()) migrator.register_migration(build_migration_10()) + migrator.register_migration(build_migration_11(app_config=config)) migrator.run_migrations() return db diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_11.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_11.py new file mode 100644 index 0000000000..d2c281110a --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_11.py @@ -0,0 +1,36 @@ +import sqlite3 +import shutil + +from invokeai.app.services.config import InvokeAIAppConfig +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration + + +class Migration11Callback: + def __init__(self, app_config: InvokeAIAppConfig) -> None: + self._app_config = app_config + + def __call__(self, cursor: sqlite3.Cursor) -> None: + self._remove_model_convert_cache_dir() + + def _remove_model_convert_cache_dir(self) -> None: + """ + Removes unused model convert cache directory + """ + convert_cache = self._app_config.convert_cache_path + print(f'DEBUG: convert_cache = {convert_cache}') + # shutil.rmtree(convert_cache) + + +def build_migration_11(app_config: InvokeAIAppConfig) -> Migration: + """ + Build the migration from database version 10 to 11. + + This migration removes the now-unused model convert cache directory. + """ + migration_11 = Migration( + from_version=10, + to_version=11, + callback=Migration11Callback(app_config), + ) + + return migration_11 diff --git a/invokeai/backend/model_manager/config.py b/invokeai/backend/model_manager/config.py index b19501843c..1e92499d30 100644 --- a/invokeai/backend/model_manager/config.py +++ b/invokeai/backend/model_manager/config.py @@ -24,6 +24,7 @@ import time from enum import Enum from typing import Literal, Optional, Type, TypeAlias, Union +import diffusers import torch from diffusers.models.modeling_utils import ModelMixin from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, TypeAdapter @@ -36,7 +37,7 @@ from ..raw_model import RawModel # ModelMixin is the base class for all diffusers and transformers models # RawModel is the InvokeAI wrapper class for ip_adapters, loras, textual_inversion and onnx runtime -AnyModel = Union[ModelMixin, RawModel, torch.nn.Module] +AnyModel = Union[ModelMixin, RawModel, torch.nn.Module, diffusers.DiffusionPipeline] class InvalidModelConfigException(Exception): diff --git a/invokeai/backend/model_manager/convert_ckpt_to_diffusers.py b/invokeai/backend/model_manager/convert_ckpt_to_diffusers.py deleted file mode 100644 index 450e69cf38..0000000000 --- a/invokeai/backend/model_manager/convert_ckpt_to_diffusers.py +++ /dev/null @@ -1,83 +0,0 @@ -# Adapted for use in InvokeAI by Lincoln Stein, July 2023 -# -"""Conversion script for the Stable Diffusion checkpoints.""" - -from pathlib import Path -from typing import Optional - -import torch -from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL -from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( - convert_ldm_vae_checkpoint, - create_vae_diffusers_config, - download_controlnet_from_original_ckpt, - download_from_original_stable_diffusion_ckpt, -) -from omegaconf import DictConfig - -from . import AnyModel - - -def convert_ldm_vae_to_diffusers( - checkpoint: torch.Tensor | dict[str, torch.Tensor], - vae_config: DictConfig, - image_size: int, - dump_path: Optional[Path] = None, - precision: torch.dtype = torch.float16, -) -> AutoencoderKL: - """Convert a checkpoint-style VAE into a Diffusers VAE""" - vae_config = create_vae_diffusers_config(vae_config, image_size=image_size) - converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) - - vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - vae.to(precision) - - if dump_path: - vae.save_pretrained(dump_path, safe_serialization=True) - - return vae - - -def convert_ckpt_to_diffusers( - checkpoint_path: str | Path, - dump_path: Optional[str | Path] = None, - precision: torch.dtype = torch.float16, - use_safetensors: bool = True, - **kwargs, -) -> AnyModel: - """ - Takes all the arguments of download_from_original_stable_diffusion_ckpt(), - and in addition a path-like object indicating the location of the desired diffusers - model to be written. - """ - pipe = download_from_original_stable_diffusion_ckpt(Path(checkpoint_path).as_posix(), **kwargs) - pipe = pipe.to(precision) - - # TO DO: save correct repo variant - if dump_path: - pipe.save_pretrained( - dump_path, - safe_serialization=use_safetensors, - ) - return pipe - - -def convert_controlnet_to_diffusers( - checkpoint_path: Path, - dump_path: Optional[Path] = None, - precision: torch.dtype = torch.float16, - **kwargs, -) -> AnyModel: - """ - Takes all the arguments of download_controlnet_from_original_ckpt(), - and in addition a path-like object indicating the location of the desired diffusers - model to be written. - """ - pipe = download_controlnet_from_original_ckpt(checkpoint_path.as_posix(), **kwargs) - pipe = pipe.to(precision) - - # TO DO: save correct repo variant - if dump_path: - pipe.save_pretrained(dump_path, safe_serialization=True) - return pipe diff --git a/invokeai/backend/model_manager/load/__init__.py b/invokeai/backend/model_manager/load/__init__.py index f47a2c4368..f862e071d4 100644 --- a/invokeai/backend/model_manager/load/__init__.py +++ b/invokeai/backend/model_manager/load/__init__.py @@ -6,7 +6,6 @@ Init file for the model loader. from importlib import import_module from pathlib import Path -from .convert_cache.convert_cache_default import ModelConvertCache from .load_base import LoadedModel, ModelLoaderBase from .load_default import ModelLoader from .model_cache.model_cache_default import ModelCache @@ -20,7 +19,6 @@ for module in loaders: __all__ = [ "LoadedModel", "ModelCache", - "ModelConvertCache", "ModelLoaderBase", "ModelLoader", "ModelLoaderRegistryBase", diff --git a/invokeai/backend/model_manager/load/convert_cache/__init__.py b/invokeai/backend/model_manager/load/convert_cache/__init__.py deleted file mode 100644 index 5be56d2d58..0000000000 --- a/invokeai/backend/model_manager/load/convert_cache/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .convert_cache_base import ModelConvertCacheBase -from .convert_cache_default import ModelConvertCache - -__all__ = ["ModelConvertCacheBase", "ModelConvertCache"] diff --git a/invokeai/backend/model_manager/load/convert_cache/convert_cache_base.py b/invokeai/backend/model_manager/load/convert_cache/convert_cache_base.py deleted file mode 100644 index ef363cc7f4..0000000000 --- a/invokeai/backend/model_manager/load/convert_cache/convert_cache_base.py +++ /dev/null @@ -1,28 +0,0 @@ -""" -Disk-based converted model cache. -""" - -from abc import ABC, abstractmethod -from pathlib import Path - - -class ModelConvertCacheBase(ABC): - @property - @abstractmethod - def max_size(self) -> float: - """Return the maximum size of this cache directory.""" - pass - - @abstractmethod - def make_room(self, size: float) -> None: - """ - Make sufficient room in the cache directory for a model of max_size. - - :param size: Size required (GB) - """ - pass - - @abstractmethod - def cache_path(self, key: str) -> Path: - """Return the path for a model with the indicated key.""" - pass diff --git a/invokeai/backend/model_manager/load/convert_cache/convert_cache_default.py b/invokeai/backend/model_manager/load/convert_cache/convert_cache_default.py deleted file mode 100644 index 8dc2aff74b..0000000000 --- a/invokeai/backend/model_manager/load/convert_cache/convert_cache_default.py +++ /dev/null @@ -1,81 +0,0 @@ -""" -Placeholder for convert cache implementation. -""" - -import shutil -from pathlib import Path - -from invokeai.backend.util import GIG, directory_size -from invokeai.backend.util.logging import InvokeAILogger - -from .convert_cache_base import ModelConvertCacheBase - - -class ModelConvertCache(ModelConvertCacheBase): - def __init__(self, cache_path: Path, max_size: float = 10.0): - """Initialize the convert cache with the base directory and a limit on its maximum size (in GBs).""" - if not cache_path.exists(): - cache_path.mkdir(parents=True) - self._cache_path = cache_path - self._max_size = max_size - - # adjust cache size at startup in case it has been changed - if self._cache_path.exists(): - self.make_room(0.0) - - @property - def max_size(self) -> float: - """Return the maximum size of this cache directory (GB).""" - return self._max_size - - @max_size.setter - def max_size(self, value: float) -> None: - """Set the maximum size of this cache directory (GB).""" - self._max_size = value - - def cache_path(self, key: str) -> Path: - """Return the path for a model with the indicated key.""" - return self._cache_path / key - - def make_room(self, size: float) -> None: - """ - Make sufficient room in the cache directory for a model of max_size. - - :param size: Size required (GB) - """ - size_needed = directory_size(self._cache_path) + size - max_size = int(self.max_size) * GIG - logger = InvokeAILogger.get_logger() - - if size_needed <= max_size: - return - - logger.debug( - f"Convert cache has gotten too large {(size_needed / GIG):4.2f} > {(max_size / GIG):4.2f}G.. Trimming." - ) - - # For this to work, we make the assumption that the directory contains - # a 'model_index.json', 'unet/config.json' file, or a 'config.json' file at top level. - # This should be true for any diffusers model. - def by_atime(path: Path) -> float: - for config in ["model_index.json", "unet/config.json", "config.json"]: - sentinel = path / config - if sentinel.exists(): - return sentinel.stat().st_atime - - # no sentinel file found! - pick the most recent file in the directory - try: - atimes = sorted([x.stat().st_atime for x in path.iterdir() if x.is_file()], reverse=True) - return atimes[0] - except IndexError: - return 0.0 - - # sort by last access time - least accessed files will be at the end - lru_models = sorted(self._cache_path.iterdir(), key=by_atime, reverse=True) - logger.debug(f"cached models in descending atime order: {lru_models}") - while size_needed > max_size and len(lru_models) > 0: - next_victim = lru_models.pop() - victim_size = directory_size(next_victim) - logger.debug(f"Removing cached converted model {next_victim} to free {victim_size / GIG} GB") - shutil.rmtree(next_victim) - size_needed -= victim_size diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 1bb093a990..4a6aba87ca 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -18,7 +18,6 @@ from invokeai.backend.model_manager.config import ( AnyModelConfig, SubModelType, ) -from invokeai.backend.model_manager.load.convert_cache.convert_cache_base import ModelConvertCacheBase from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase @@ -106,7 +105,6 @@ class ModelLoaderBase(ABC): app_config: InvokeAIAppConfig, logger: Logger, ram_cache: ModelCacheBase[AnyModel], - convert_cache: ModelConvertCacheBase, ): """Initialize the loader.""" pass @@ -132,12 +130,6 @@ class ModelLoaderBase(ABC): """Return size in bytes of the model, calculated before loading.""" pass - @property - @abstractmethod - def convert_cache(self) -> ModelConvertCacheBase: - """Return the convert cache associated with this loader.""" - pass - @property @abstractmethod def ram_cache(self) -> ModelCacheBase[AnyModel]: diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index f406b1a391..21a119f6bc 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -12,8 +12,7 @@ from invokeai.backend.model_manager import ( InvalidModelConfigException, SubModelType, ) -from invokeai.backend.model_manager.config import DiffusersConfigBase, ModelType -from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase +from invokeai.backend.model_manager.config import DiffusersConfigBase from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data, calc_model_size_by_fs @@ -30,13 +29,11 @@ class ModelLoader(ModelLoaderBase): app_config: InvokeAIAppConfig, logger: Logger, ram_cache: ModelCacheBase[AnyModel], - convert_cache: ModelConvertCacheBase, ): """Initialize the loader.""" self._app_config = app_config self._logger = logger self._ram_cache = ram_cache - self._convert_cache = convert_cache self._torch_dtype = TorchDevice.choose_torch_dtype() def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel: @@ -50,23 +47,15 @@ class ModelLoader(ModelLoaderBase): :param submodel_type: an ModelType enum indicating the portion of the model to retrieve (e.g. ModelType.Vae) """ - if model_config.type is ModelType.Main and not submodel_type: - raise InvalidModelConfigException("submodel_type is required when loading a main model") - model_path = self._get_model_path(model_config) if not model_path.exists(): raise InvalidModelConfigException(f"Files for model '{model_config.name}' not found at {model_path}") with skip_torch_weight_init(): - locker = self._convert_and_load(model_config, model_path, submodel_type) + locker = self._load_and_cache(model_config, submodel_type) return LoadedModel(config=model_config, _locker=locker) - @property - def convert_cache(self) -> ModelConvertCacheBase: - """Return the convert cache associated with this loader.""" - return self._convert_cache - @property def ram_cache(self) -> ModelCacheBase[AnyModel]: """Return the ram cache associated with this loader.""" @@ -76,9 +65,7 @@ class ModelLoader(ModelLoaderBase): model_base = self._app_config.models_path return (model_base / config.path).resolve() - def _convert_and_load( - self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None - ) -> ModelLockerBase: + def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> ModelLockerBase: try: return self._ram_cache.get(config.key, submodel_type) except IndexError: @@ -86,13 +73,6 @@ class ModelLoader(ModelLoaderBase): config.path = str(self._get_model_path(config)) loaded_model = self._load_model(config, submodel_type) - - # cache_path: Path = self._convert_cache.cache_path(config.key) - # if self._needs_conversion(config, model_path, cache_path): - # loaded_model = self._do_convert(config, model_path, cache_path, submodel_type) - # else: - # config.path = str(cache_path) if cache_path.exists() else str(self._get_model_path(config)) - # loaded_model = self._load_model(config, submodel_type) self._ram_cache.put( config.key, @@ -117,30 +97,6 @@ class ModelLoader(ModelLoaderBase): variant=config.repo_variant if isinstance(config, DiffusersConfigBase) else None, ) - def _do_convert( - self, config: AnyModelConfig, model_path: Path, cache_path: Path, submodel_type: Optional[SubModelType] = None - ) -> AnyModel: - self.convert_cache.make_room(calc_model_size_by_fs(model_path)) - pipeline = self._convert_model(config, model_path, cache_path if self.convert_cache.max_size > 0 else None) - if submodel_type: - # Proactively load the various submodels into the RAM cache so that we don't have to re-convert - # the entire pipeline every time a new submodel is needed. - for subtype in SubModelType: - if subtype == submodel_type: - continue - if submodel := getattr(pipeline, subtype.value, None): - self._ram_cache.put( - config.key, submodel_type=subtype, model=submodel, size=calc_model_size_by_data(submodel) - ) - return getattr(pipeline, submodel_type.value) if submodel_type else pipeline - - def _needs_conversion(self, config: AnyModelConfig, model_path: Path, dest_path: Path) -> bool: - return False - - # This needs to be implemented in subclasses that handle checkpoints - def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel: - raise NotImplementedError - # This needs to be implemented in the subclass def _load_model( self, diff --git a/invokeai/backend/model_manager/load/model_loaders/controlnet.py b/invokeai/backend/model_manager/load/model_loaders/controlnet.py index ffd851e31f..6b88e279ba 100644 --- a/invokeai/backend/model_manager/load/model_loaders/controlnet.py +++ b/invokeai/backend/model_manager/load/model_loaders/controlnet.py @@ -1,9 +1,10 @@ # Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team """Class for ControlNet model loading in InvokeAI.""" -from pathlib import Path from typing import Optional + from diffusers import ControlNetModel + from invokeai.backend.model_manager import ( AnyModel, AnyModelConfig, @@ -11,7 +12,7 @@ from invokeai.backend.model_manager import ( ModelFormat, ModelType, ) -from invokeai.backend.model_manager.config import SubModelType, ControlNetCheckpointConfig +from invokeai.backend.model_manager.config import ControlNetCheckpointConfig, SubModelType from .. import ModelLoaderRegistry from .generic_diffusers import GenericDiffusersLoader @@ -28,32 +29,11 @@ class ControlNetLoader(GenericDiffusersLoader): submodel_type: Optional[SubModelType] = None, ) -> AnyModel: if isinstance(config, ControlNetCheckpointConfig): - return ControlNetModel.from_single_file(config.path, - config=self._app_config.legacy_conf_path / config.config_path, - torch_dtype=self._torch_dtype, - local_files_only=True, - ) + return ControlNetModel.from_single_file( + config.path, + config=self._app_config.legacy_conf_path / config.config_path, + torch_dtype=self._torch_dtype, + local_files_only=True, + ) else: return super()._load_model(config, submodel_type) - - # def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel: - # assert isinstance(config, CheckpointConfigBase) - # image_size = ( - # 512 - # if config.base == BaseModelType.StableDiffusion1 - # else 768 - # if config.base == BaseModelType.StableDiffusion2 - # else 1024 - # ) - - # self._logger.info(f"Converting {model_path} to diffusers format") - # with open(self._app_config.legacy_conf_path / config.config_path, "r") as config_stream: - # result = convert_controlnet_to_diffusers( - # model_path, - # output_path, - # original_config_file=config_stream, - # image_size=image_size, - # precision=self._torch_dtype, - # from_safetensors=model_path.suffix == ".safetensors", - # ) - # return result diff --git a/invokeai/backend/model_manager/load/model_loaders/lora.py b/invokeai/backend/model_manager/load/model_loaders/lora.py index 53814279ec..aa0acab6bc 100644 --- a/invokeai/backend/model_manager/load/model_loaders/lora.py +++ b/invokeai/backend/model_manager/load/model_loaders/lora.py @@ -15,7 +15,6 @@ from invokeai.backend.model_manager import ( ModelType, SubModelType, ) -from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase from .. import ModelLoader, ModelLoaderRegistry @@ -32,10 +31,9 @@ class LoRALoader(ModelLoader): app_config: InvokeAIAppConfig, logger: Logger, ram_cache: ModelCacheBase[AnyModel], - convert_cache: ModelConvertCacheBase, ): """Initialize the loader.""" - super().__init__(app_config, logger, ram_cache, convert_cache) + super().__init__(app_config, logger, ram_cache) self._model_base: Optional[BaseModelType] = None def _load_model( diff --git a/invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py b/invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py index dce3990f37..15abfe593a 100644 --- a/invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py +++ b/invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py @@ -3,14 +3,14 @@ from pathlib import Path from typing import Optional + from diffusers import ( - StableDiffusionPipeline, StableDiffusionInpaintPipeline, - StableDiffusionXLPipeline, + StableDiffusionPipeline, StableDiffusionXLInpaintPipeline, + StableDiffusionXLPipeline, ) -from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data from invokeai.backend.model_manager import ( AnyModel, AnyModelConfig, @@ -18,16 +18,14 @@ from invokeai.backend.model_manager import ( ModelFormat, ModelType, ModelVariantType, - SchedulerPredictionType, SubModelType, ) from invokeai.backend.model_manager.config import ( CheckpointConfigBase, DiffusersConfigBase, MainCheckpointConfig, - ModelVariantType, ) -from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_ckpt_to_diffusers +from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data from .. import ModelLoaderRegistry from .generic_diffusers import GenericDiffusersLoader @@ -56,12 +54,12 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader): config: AnyModelConfig, submodel_type: Optional[SubModelType] = None, ) -> AnyModel: - if not submodel_type is not None: - raise Exception("A submodel type must be provided when loading main pipelines.") - if isinstance(config, CheckpointConfigBase): return self._load_from_singlefile(config, submodel_type) + if not submodel_type is not None: + raise Exception("A submodel type must be provided when loading main pipelines.") + model_path = Path(config.path) load_class = self.get_hf_load_class(model_path, submodel_type) repo_variant = config.repo_variant if isinstance(config, DiffusersConfigBase) else None @@ -84,9 +82,9 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader): return result def _load_from_singlefile( - self, - config: AnyModelConfig, - submodel_type: SubModelType, + self, + config: AnyModelConfig, + submodel_type: Optional[SubModelType] = None, ) -> AnyModel: load_classes = { BaseModelType.StableDiffusion1: { @@ -100,22 +98,32 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader): BaseModelType.StableDiffusionXL: { ModelVariantType.Normal: StableDiffusionXLPipeline, ModelVariantType.Inpaint: StableDiffusionXLInpaintPipeline, - } + }, } assert isinstance(config, MainCheckpointConfig) try: load_class = load_classes[config.base][config.variant] except KeyError as e: - raise Exception(f'No diffusers pipeline known for base={config.base}, variant={config.variant}') from e - original_config_file=self._app_config.legacy_conf_path / config.config_path # should try without using this... - pipeline = load_class.from_single_file(config.path, - config=original_config_file, - torch_dtype=self._torch_dtype, - local_files_only=True, - ) + raise Exception(f"No diffusers pipeline known for base={config.base}, variant={config.variant}") from e + original_config_file = self._app_config.legacy_conf_path / config.config_path + prediction_type = config.prediction_type.value + upcast_attention = config.upcast_attention - # Proactively load the various submodels into the RAM cache so that we don't have to re-convert + pipeline = load_class.from_single_file( + config.path, + config=original_config_file, + torch_dtype=self._torch_dtype, + local_files_only=True, + prediction_type=prediction_type, + upcast_attention=upcast_attention, + load_safety_checker=False, + ) + + # Proactively load the various submodels into the RAM cache so that we don't have to re-load # the entire pipeline every time a new submodel is needed. + if not submodel_type: + return pipeline + for subtype in SubModelType: if subtype == submodel_type: continue @@ -124,48 +132,3 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader): config.key, submodel_type=subtype, model=submodel, size=calc_model_size_by_data(submodel) ) return getattr(pipeline, submodel_type.value) - - - # def _needs_conversion(self, config: AnyModelConfig, model_path: Path, dest_path: Path) -> bool: - # if not isinstance(config, CheckpointConfigBase): - # return False - # elif ( - # dest_path.exists() - # and (dest_path / "model_index.json").stat().st_mtime >= (config.converted_at or 0.0) - # and (dest_path / "model_index.json").stat().st_mtime >= model_path.stat().st_mtime - # ): - # return False - # else: - # return True - - # def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel: - # assert isinstance(config, MainCheckpointConfig) - # base = config.base - - # prediction_type = config.prediction_type.value - # upcast_attention = config.upcast_attention - # image_size = ( - # 1024 - # if base == BaseModelType.StableDiffusionXL - # else 768 - # if config.prediction_type == SchedulerPredictionType.VPrediction and base == BaseModelType.StableDiffusion2 - # else 512 - # ) - - # self._logger.info(f"Converting {model_path} to diffusers format") - - # loaded_model = convert_ckpt_to_diffusers( - # model_path, - # output_path, - # model_type=self.model_base_to_model_type[base], - # original_config_file=self._app_config.legacy_conf_path / config.config_path, - # extract_ema=True, - # from_safetensors=model_path.suffix == ".safetensors", - # precision=self._torch_dtype, - # prediction_type=prediction_type, - # image_size=image_size, - # upcast_attention=upcast_attention, - # load_safety_checker=False, - # num_in_channels=VARIANT_TO_IN_CHANNEL_MAP[config.variant], - # ) - # return loaded_model diff --git a/invokeai/backend/model_manager/load/model_loaders/vae.py b/invokeai/backend/model_manager/load/model_loaders/vae.py index 81cf36130d..f278a2069e 100644 --- a/invokeai/backend/model_manager/load/model_loaders/vae.py +++ b/invokeai/backend/model_manager/load/model_loaders/vae.py @@ -1,13 +1,9 @@ # Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team """Class for VAE model loading in InvokeAI.""" -from diffusers import AutoencoderKL -from pathlib import Path from typing import Optional -import torch -from omegaconf import DictConfig, OmegaConf -from safetensors.torch import load_file as safetensors_load_file +from diffusers import AutoencoderKL from invokeai.backend.model_manager import ( AnyModelConfig, @@ -15,8 +11,7 @@ from invokeai.backend.model_manager import ( ModelFormat, ModelType, ) -from invokeai.backend.model_manager.config import AnyModel, CheckpointConfigBase, SubModelType, VAECheckpointConfig -from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_ldm_vae_to_diffusers +from invokeai.backend.model_manager.config import AnyModel, SubModelType, VAECheckpointConfig from .. import ModelLoaderRegistry from .generic_diffusers import GenericDiffusersLoader @@ -34,39 +29,11 @@ class VAELoader(GenericDiffusersLoader): submodel_type: Optional[SubModelType] = None, ) -> AnyModel: if isinstance(config, VAECheckpointConfig): - return AutoencoderKL.from_single_file(config.path, - config=self._app_config.legacy_conf_path / config.config_path, - torch_dtype=self._torch_dtype, - local_files_only=True, - ) + return AutoencoderKL.from_single_file( + config.path, + config=self._app_config.legacy_conf_path / config.config_path, + torch_dtype=self._torch_dtype, + local_files_only=True, + ) else: return super()._load_model(config, submodel_type) - - # def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel: - # # TODO(MM2): check whether sdxl VAE models convert. - # if config.base not in {BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2}: - # raise Exception(f"VAE conversion not supported for model type: {config.base}") - # else: - # assert isinstance(config, CheckpointConfigBase) - # config_file = self._app_config.legacy_conf_path / config.config_path - - # if model_path.suffix == ".safetensors": - # checkpoint = safetensors_load_file(model_path, device="cpu") - # else: - # checkpoint = torch.load(model_path, map_location="cpu") - - # # sometimes weights are hidden under "state_dict", and sometimes not - # if "state_dict" in checkpoint: - # checkpoint = checkpoint["state_dict"] - - # ckpt_config = OmegaConf.load(config_file) - # assert isinstance(ckpt_config, DictConfig) - # self._logger.info(f"Converting {model_path} to diffusers format") - # vae_model = convert_ldm_vae_to_diffusers( - # checkpoint=checkpoint, - # vae_config=ckpt_config, - # image_size=512, - # precision=self._torch_dtype, - # dump_path=output_path, - # ) - # return vae_model diff --git a/invokeai/backend/model_manager/probe.py b/invokeai/backend/model_manager/probe.py index 8f33e4b49f..2e34c2fb1d 100644 --- a/invokeai/backend/model_manager/probe.py +++ b/invokeai/backend/model_manager/probe.py @@ -257,6 +257,8 @@ class ModelProbe(object): if (folder_path / "image_encoder.txt").exists(): return ModelType.IPAdapter + print(f'DEBUG: {folder_path} contents = {list(folder_path.glob("**"))}') + i = folder_path / "model_index.json" c = folder_path / "config.json" config_path = i if i.exists() else c if c.exists() else None diff --git a/tests/backend/model_manager/model_manager_fixtures.py b/tests/backend/model_manager/model_manager_fixtures.py index 5ddccd05bb..60da14db6d 100644 --- a/tests/backend/model_manager/model_manager_fixtures.py +++ b/tests/backend/model_manager/model_manager_fixtures.py @@ -25,7 +25,7 @@ from invokeai.backend.model_manager.config import ( ModelVariantType, VAEDiffusersConfig, ) -from invokeai.backend.model_manager.load import ModelCache, ModelConvertCache +from invokeai.backend.model_manager.load import ModelCache from invokeai.backend.util.logging import InvokeAILogger from tests.backend.model_manager.model_metadata.metadata_examples import ( HFTestLoraMetadata, @@ -82,17 +82,15 @@ def mm2_download_queue(mm2_session: Session) -> DownloadQueueServiceBase: @pytest.fixture -def mm2_loader(mm2_app_config: InvokeAIAppConfig, mm2_record_store: ModelRecordServiceBase) -> ModelLoadServiceBase: +def mm2_loader(mm2_app_config: InvokeAIAppConfig) -> ModelLoadServiceBase: ram_cache = ModelCache( logger=InvokeAILogger.get_logger(), max_cache_size=mm2_app_config.ram, max_vram_cache_size=mm2_app_config.vram, ) - convert_cache = ModelConvertCache(mm2_app_config.convert_cache_path) return ModelLoadService( app_config=mm2_app_config, ram_cache=ram_cache, - convert_cache=convert_cache, )