From 8d7ca9c1b70918606fc11e97f5e72291bbe0a8ed Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Tue, 2 Jul 2024 16:48:33 -0400 Subject: [PATCH] More refactoring to help with circular imports. --- .../app/invocations/create_gradient_mask.py | 2 +- invokeai/app/invocations/image_to_latents.py | 2 +- .../services/model_load/model_load_base.py | 2 +- .../services/model_load/model_load_default.py | 8 ++---- .../model_manager/model_manager_default.py | 3 ++- invokeai/backend/model_manager/__init__.py | 2 -- .../backend/model_manager/load/__init__.py | 26 ------------------- .../model_manager/load/load_default.py | 2 +- .../load/model_loader_registry.py | 3 ++- .../load/model_loaders/generic_diffusers.py | 4 +-- .../load/model_loaders/ip_adapter.py | 3 ++- invokeai/backend/util/test_utils.py | 3 ++- .../model_manager/model_manager_fixtures.py | 2 +- 13 files changed, 17 insertions(+), 45 deletions(-) diff --git a/invokeai/app/invocations/create_gradient_mask.py b/invokeai/app/invocations/create_gradient_mask.py index 089313463b..cafd4b1ca7 100644 --- a/invokeai/app/invocations/create_gradient_mask.py +++ b/invokeai/app/invocations/create_gradient_mask.py @@ -19,8 +19,8 @@ from invokeai.app.invocations.fields import ( from invokeai.app.invocations.image_to_latents import ImageToLatentsInvocation from invokeai.app.invocations.model import UNetField, VAEField from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager import LoadedModel from invokeai.backend.model_manager.config import MainConfigBase, ModelVariantType +from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index dadd8bb3a1..bdff573b14 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/invokeai/app/invocations/image_to_latents.py @@ -23,7 +23,7 @@ from invokeai.app.invocations.fields import ( from invokeai.app.invocations.model import VAEField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager import LoadedModel +from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params diff --git a/invokeai/app/services/model_load/model_load_base.py b/invokeai/app/services/model_load/model_load_base.py index 03c2a81e0c..cc6c8b0eed 100644 --- a/invokeai/app/services/model_load/model_load_base.py +++ b/invokeai/app/services/model_load/model_load_base.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import Callable, Optional from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType -from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig +from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase diff --git a/invokeai/app/services/model_load/model_load_default.py b/invokeai/app/services/model_load/model_load_default.py index 8eb94616be..04284ee973 100644 --- a/invokeai/app/services/model_load/model_load_default.py +++ b/invokeai/app/services/model_load/model_load_default.py @@ -11,13 +11,9 @@ from torch import load as torch_load from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.invoker import Invoker from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType -from invokeai.backend.model_manager.load import ( - LoadedModel, - LoadedModelWithoutConfig, - ModelLoaderRegistry, - ModelLoaderRegistryBase, -) +from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry, ModelLoaderRegistryBase from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger diff --git a/invokeai/app/services/model_manager/model_manager_default.py b/invokeai/app/services/model_manager/model_manager_default.py index f695c3c8c1..e439638f45 100644 --- a/invokeai/app/services/model_manager/model_manager_default.py +++ b/invokeai/app/services/model_manager/model_manager_default.py @@ -7,7 +7,8 @@ import torch from typing_extensions import Self from invokeai.app.services.invoker import Invoker -from invokeai.backend.model_manager.load import ModelCache, ModelLoaderRegistry +from invokeai.backend.model_manager.load.model_cache.model_cache_default import ModelCache +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger diff --git a/invokeai/backend/model_manager/__init__.py b/invokeai/backend/model_manager/__init__.py index 98cc5054c7..5b89639a0a 100644 --- a/invokeai/backend/model_manager/__init__.py +++ b/invokeai/backend/model_manager/__init__.py @@ -13,7 +13,6 @@ from .config import ( SchedulerPredictionType, SubModelType, ) -from .load import LoadedModel from .probe import ModelProbe from .search import ModelSearch @@ -23,7 +22,6 @@ __all__ = [ "BaseModelType", "ModelRepoVariant", "InvalidModelConfigException", - "LoadedModel", "ModelConfigFactory", "ModelFormat", "ModelProbe", diff --git a/invokeai/backend/model_manager/load/__init__.py b/invokeai/backend/model_manager/load/__init__.py index 3d34f02727..8b13789179 100644 --- a/invokeai/backend/model_manager/load/__init__.py +++ b/invokeai/backend/model_manager/load/__init__.py @@ -1,27 +1 @@ -# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development Team -""" -Init file for the model loader. -""" -from importlib import import_module -from pathlib import Path - -from .load_base import LoadedModel, LoadedModelWithoutConfig, ModelLoaderBase -from .load_default import ModelLoader -from .model_cache.model_cache_default import ModelCache -from .model_loader_registry import ModelLoaderRegistry, ModelLoaderRegistryBase - -# This registers the subclasses that implement loaders of specific model types -loaders = [x.stem for x in Path(Path(__file__).parent, "model_loaders").glob("*.py") if x.stem != "__init__"] -for module in loaders: - import_module(f"{__package__}.model_loaders.{module}") - -__all__ = [ - "LoadedModel", - "LoadedModelWithoutConfig", - "ModelCache", - "ModelLoaderBase", - "ModelLoader", - "ModelLoaderRegistryBase", - "ModelLoaderRegistry", -] diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index ce9811534e..76c260d962 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -15,7 +15,7 @@ from invokeai.backend.model_manager import ( 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_fs +from invokeai.backend.model_manager.load.model_size_utils import calc_model_size_by_fs from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init from invokeai.backend.util.devices import TorchDevice diff --git a/invokeai/backend/model_manager/load/model_loader_registry.py b/invokeai/backend/model_manager/load/model_loader_registry.py index bb6bd18d7f..8d7d3cb176 100644 --- a/invokeai/backend/model_manager/load/model_loader_registry.py +++ b/invokeai/backend/model_manager/load/model_loader_registry.py @@ -18,6 +18,8 @@ Use like this: from abc import ABC, abstractmethod from typing import Callable, Dict, Optional, Tuple, Type, TypeVar +from invokeai.backend.model_manager.load.load_base import ModelLoaderBase + from ..config import ( AnyModelConfig, BaseModelType, @@ -26,7 +28,6 @@ from ..config import ( ModelType, SubModelType, ) -from . import ModelLoaderBase class ModelLoaderRegistryBase(ABC): diff --git a/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py b/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py index 6320797b8a..dfe38aa79c 100644 --- a/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py +++ b/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py @@ -18,8 +18,8 @@ from invokeai.backend.model_manager import ( SubModelType, ) from invokeai.backend.model_manager.config import DiffusersConfigBase - -from .. import ModelLoader, ModelLoaderRegistry +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers) diff --git a/invokeai/backend/model_manager/load/model_loaders/ip_adapter.py b/invokeai/backend/model_manager/load/model_loaders/ip_adapter.py index 55eed81fcd..d89d843738 100644 --- a/invokeai/backend/model_manager/load/model_loaders/ip_adapter.py +++ b/invokeai/backend/model_manager/load/model_loaders/ip_adapter.py @@ -8,7 +8,8 @@ import torch from invokeai.backend.ip_adapter.ip_adapter import build_ip_adapter from invokeai.backend.model_manager import AnyModel, AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType -from invokeai.backend.model_manager.load import ModelLoader, ModelLoaderRegistry +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.raw_model import RawModel diff --git a/invokeai/backend/util/test_utils.py b/invokeai/backend/util/test_utils.py index 8d7c0c4d04..7e83165aae 100644 --- a/invokeai/backend/util/test_utils.py +++ b/invokeai/backend/util/test_utils.py @@ -7,7 +7,8 @@ import torch from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase from invokeai.app.services.model_records import UnknownModelException -from invokeai.backend.model_manager import BaseModelType, LoadedModel, ModelType, SubModelType +from invokeai.backend.model_manager import BaseModelType, ModelType, SubModelType +from invokeai.backend.model_manager.load.load_base import LoadedModel @pytest.fixture(scope="session") diff --git a/tests/backend/model_manager/model_manager_fixtures.py b/tests/backend/model_manager/model_manager_fixtures.py index b7ef17970f..1c3cb7af69 100644 --- a/tests/backend/model_manager/model_manager_fixtures.py +++ b/tests/backend/model_manager/model_manager_fixtures.py @@ -26,7 +26,7 @@ from invokeai.backend.model_manager.config import ( ModelVariantType, VAEDiffusersConfig, ) -from invokeai.backend.model_manager.load import ModelCache +from invokeai.backend.model_manager.load.model_cache.model_cache_default import ModelCache from invokeai.backend.util.logging import InvokeAILogger from tests.backend.model_manager.model_metadata.metadata_examples import ( HFTestLoraMetadata,