diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs index c6b833cf58..5c04dc964e 100644 --- a/.git-blame-ignore-revs +++ b/.git-blame-ignore-revs @@ -1,2 +1,5 @@ b3dccfaeb636599c02effc377cdd8a87d658256c 218b6d0546b990fc449c876fb99f44b50c4daa35 +182580ff6970caed400be178c5b888514b75d7f2 +8e9d5c1187b0d36da80571ce4c8ba9b3a37b6c46 +99aac5870e1092b182e6c5f21abcaab6936a4ad1 \ No newline at end of file diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 8642085bd9..94c0e88d4f 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -28,12 +28,10 @@ from invokeai.app.services.model_records import ( UnknownModelException, ) from invokeai.app.util.suppress_output import SuppressOutput +from invokeai.backend.model_manager import BaseModelType, ModelFormat, ModelType from invokeai.backend.model_manager.config import ( AnyModelConfig, - BaseModelType, MainCheckpointConfig, - ModelFormat, - ModelType, ) from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch diff --git a/invokeai/app/invocations/create_gradient_mask.py b/invokeai/app/invocations/create_gradient_mask.py index 0c5b4c4418..97ccdb5133 100644 --- a/invokeai/app/invocations/create_gradient_mask.py +++ b/invokeai/app/invocations/create_gradient_mask.py @@ -19,7 +19,8 @@ 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.config import MainConfigBase +from invokeai.backend.model_manager.taxonomy import ModelVariantType from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 4a98cf0f3d..ab65298e0b 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -39,8 +39,8 @@ from invokeai.app.invocations.t2i_adapter import T2IAdapterField from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.util.controlnet_utils import prepare_control_image from invokeai.backend.ip_adapter.ip_adapter import IPAdapter -from invokeai.backend.model_manager import BaseModelType, ModelVariantType from invokeai.backend.model_manager.config import AnyModelConfig +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelVariantType from invokeai.backend.model_patcher import ModelPatcher from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.model_patch_raw import ModelPatchRaw diff --git a/invokeai/app/invocations/flux_denoise.py b/invokeai/app/invocations/flux_denoise.py index c79e63360c..fb063acbc7 100644 --- a/invokeai/app/invocations/flux_denoise.py +++ b/invokeai/app/invocations/flux_denoise.py @@ -49,7 +49,7 @@ from invokeai.backend.flux.sampling_utils import ( unpack, ) from invokeai.backend.flux.text_conditioning import FluxReduxConditioning, FluxTextConditioning -from invokeai.backend.model_manager.config import ModelFormat, ModelVariantType +from invokeai.backend.model_manager.taxonomy import ModelFormat, ModelVariantType from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw diff --git a/invokeai/app/invocations/flux_lora_loader.py b/invokeai/app/invocations/flux_lora_loader.py index fe1e399676..7c155ddbd4 100644 --- a/invokeai/app/invocations/flux_lora_loader.py +++ b/invokeai/app/invocations/flux_lora_loader.py @@ -10,7 +10,7 @@ from invokeai.app.invocations.baseinvocation import ( from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType from invokeai.app.invocations.model import CLIPField, LoRAField, ModelIdentifierField, T5EncoderField, TransformerField from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager.config import BaseModelType +from invokeai.backend.model_manager.taxonomy import BaseModelType @invocation_output("flux_lora_loader_output") diff --git a/invokeai/app/invocations/flux_model_loader.py b/invokeai/app/invocations/flux_model_loader.py index 88483f71ae..6803f0edc3 100644 --- a/invokeai/app/invocations/flux_model_loader.py +++ b/invokeai/app/invocations/flux_model_loader.py @@ -17,8 +17,8 @@ from invokeai.app.util.t5_model_identifier import ( from invokeai.backend.flux.util import max_seq_lengths from invokeai.backend.model_manager.config import ( CheckpointConfigBase, - SubModelType, ) +from invokeai.backend.model_manager.taxonomy import SubModelType @invocation_output("flux_model_loader_output") diff --git a/invokeai/app/invocations/flux_redux.py b/invokeai/app/invocations/flux_redux.py index f6581bd73d..e055d94f8d 100644 --- a/invokeai/app/invocations/flux_redux.py +++ b/invokeai/app/invocations/flux_redux.py @@ -23,7 +23,8 @@ from invokeai.app.invocations.primitives import ImageField from invokeai.app.services.model_records.model_records_base import ModelRecordChanges from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.redux.flux_redux_model import FluxReduxModel -from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelType +from invokeai.backend.model_manager import BaseModelType, ModelType +from invokeai.backend.model_manager.config import AnyModelConfig from invokeai.backend.model_manager.starter_models import siglip from invokeai.backend.sig_lip.sig_lip_pipeline import SigLipPipeline from invokeai.backend.util.devices import TorchDevice diff --git a/invokeai/app/invocations/flux_text_encoder.py b/invokeai/app/invocations/flux_text_encoder.py index de49ee2485..aab40ed85b 100644 --- a/invokeai/app/invocations/flux_text_encoder.py +++ b/invokeai/app/invocations/flux_text_encoder.py @@ -17,7 +17,7 @@ from invokeai.app.invocations.model import CLIPField, T5EncoderField from invokeai.app.invocations.primitives import FluxConditioningOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.modules.conditioner import HFEncoder -from invokeai.backend.model_manager.config import ModelFormat +from invokeai.backend.model_manager import ModelFormat from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX, FLUX_LORA_T5_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw diff --git a/invokeai/app/invocations/ideal_size.py b/invokeai/app/invocations/ideal_size.py index 954fc95d94..5b087e31d6 100644 --- a/invokeai/app/invocations/ideal_size.py +++ b/invokeai/app/invocations/ideal_size.py @@ -6,7 +6,7 @@ from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.app.invocations.fields import FieldDescriptions, InputField, OutputField from invokeai.app.invocations.model import UNetField from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager.config import BaseModelType +from invokeai.backend.model_manager.taxonomy import BaseModelType @invocation_output("ideal_size_output") diff --git a/invokeai/app/invocations/ip_adapter.py b/invokeai/app/invocations/ip_adapter.py index e3fc26c9db..f6b4cb63ec 100644 --- a/invokeai/app/invocations/ip_adapter.py +++ b/invokeai/app/invocations/ip_adapter.py @@ -13,10 +13,8 @@ from invokeai.app.services.model_records.model_records_base import ModelRecordCh from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.config import ( AnyModelConfig, - BaseModelType, IPAdapterCheckpointConfig, IPAdapterInvokeAIConfig, - ModelType, ) from invokeai.backend.model_manager.starter_models import ( StarterModel, @@ -24,6 +22,7 @@ from invokeai.backend.model_manager.starter_models import ( ip_adapter_sd_image_encoder, ip_adapter_sdxl_image_encoder, ) +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType class IPAdapterField(BaseModel): diff --git a/invokeai/app/invocations/metadata_linked.py b/invokeai/app/invocations/metadata_linked.py index 0803e10074..c212d34f7a 100644 --- a/invokeai/app/invocations/metadata_linked.py +++ b/invokeai/app/invocations/metadata_linked.py @@ -43,7 +43,7 @@ from invokeai.app.invocations.primitives import BooleanOutput, FloatOutput, Inte from invokeai.app.invocations.scheduler import SchedulerOutput from invokeai.app.invocations.t2i_adapter import T2IAdapterField, T2IAdapterInvocation from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager.config import ModelType, SubModelType +from invokeai.backend.model_manager.taxonomy import ModelType, SubModelType from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES from invokeai.version import __version__ diff --git a/invokeai/app/invocations/model.py b/invokeai/app/invocations/model.py index 8889074897..5a58bc18a8 100644 --- a/invokeai/app/invocations/model.py +++ b/invokeai/app/invocations/model.py @@ -15,10 +15,8 @@ from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.shared.models import FreeUConfig from invokeai.backend.model_manager.config import ( AnyModelConfig, - BaseModelType, - ModelType, - SubModelType, ) +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, SubModelType class ModelIdentifierField(BaseModel): diff --git a/invokeai/app/invocations/sd3_denoise.py b/invokeai/app/invocations/sd3_denoise.py index 99dbebb0b6..ea7c5f39d1 100644 --- a/invokeai/app/invocations/sd3_denoise.py +++ b/invokeai/app/invocations/sd3_denoise.py @@ -23,7 +23,7 @@ from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.invocations.sd3_text_encoder import SD3_T5_MAX_SEQ_LEN from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional -from invokeai.backend.model_manager.config import BaseModelType +from invokeai.backend.model_manager import BaseModelType from invokeai.backend.sd3.extensions.inpaint_extension import InpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import SD3ConditioningInfo diff --git a/invokeai/app/invocations/sd3_model_loader.py b/invokeai/app/invocations/sd3_model_loader.py index 70a95dbe3a..7bc31e7aec 100644 --- a/invokeai/app/invocations/sd3_model_loader.py +++ b/invokeai/app/invocations/sd3_model_loader.py @@ -14,7 +14,7 @@ from invokeai.app.util.t5_model_identifier import ( preprocess_t5_encoder_model_identifier, preprocess_t5_tokenizer_model_identifier, ) -from invokeai.backend.model_manager.config import SubModelType +from invokeai.backend.model_manager.taxonomy import SubModelType @invocation_output("sd3_model_loader_output") diff --git a/invokeai/app/invocations/sd3_text_encoder.py b/invokeai/app/invocations/sd3_text_encoder.py index 76b46c73aa..2cdd16b2b4 100644 --- a/invokeai/app/invocations/sd3_text_encoder.py +++ b/invokeai/app/invocations/sd3_text_encoder.py @@ -16,7 +16,7 @@ from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField from invokeai.app.invocations.model import CLIPField, T5EncoderField from invokeai.app.invocations.primitives import SD3ConditioningOutput from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager.config import ModelFormat +from invokeai.backend.model_manager.taxonomy import ModelFormat from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw diff --git a/invokeai/app/invocations/sdxl.py b/invokeai/app/invocations/sdxl.py index 8d7843bc50..9bc8450d9b 100644 --- a/invokeai/app/invocations/sdxl.py +++ b/invokeai/app/invocations/sdxl.py @@ -2,7 +2,7 @@ from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocati from invokeai.app.invocations.fields import FieldDescriptions, InputField, OutputField, UIType from invokeai.app.invocations.model import CLIPField, ModelIdentifierField, UNetField, VAEField from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager import SubModelType +from invokeai.backend.model_manager.taxonomy import SubModelType @invocation_output("sdxl_model_loader_output") diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py index 7fa4ea8fe6..fc0f0bb2c6 100644 --- a/invokeai/app/services/events/events_base.py +++ b/invokeai/app/services/events/events_base.py @@ -44,7 +44,8 @@ if TYPE_CHECKING: SessionQueueItem, SessionQueueStatus, ) - from invokeai.backend.model_manager.config import AnyModelConfig, SubModelType + from invokeai.backend.model_manager import SubModelType + from invokeai.backend.model_manager.config import AnyModelConfig class EventServiceBase: diff --git a/invokeai/app/services/events/events_common.py b/invokeai/app/services/events/events_common.py index 640788f48e..be114b70d0 100644 --- a/invokeai/app/services/events/events_common.py +++ b/invokeai/app/services/events/events_common.py @@ -16,7 +16,8 @@ from invokeai.app.services.session_queue.session_queue_common import ( ) from invokeai.app.services.shared.graph import AnyInvocation, AnyInvocationOutput from invokeai.app.util.misc import get_timestamp -from invokeai.backend.model_manager.config import AnyModelConfig, SubModelType +from invokeai.backend.model_manager import SubModelType +from invokeai.backend.model_manager.config import AnyModelConfig if TYPE_CHECKING: from invokeai.app.services.download.download_base import DownloadJob diff --git a/invokeai/app/services/model_install/model_install_common.py b/invokeai/app/services/model_install/model_install_common.py index 4dbd1ef1c1..fea75d7375 100644 --- a/invokeai/app/services/model_install/model_install_common.py +++ b/invokeai/app/services/model_install/model_install_common.py @@ -10,9 +10,9 @@ from typing_extensions import Annotated from invokeai.app.services.download import DownloadJob, MultiFileDownloadJob from invokeai.app.services.model_records import ModelRecordChanges -from invokeai.backend.model_manager import AnyModelConfig, ModelRepoVariant -from invokeai.backend.model_manager.config import ModelSourceType +from invokeai.backend.model_manager.config import AnyModelConfig from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata +from invokeai.backend.model_manager.taxonomy import ModelRepoVariant, ModelSourceType class InstallStatus(str, Enum): diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 38ab2676c9..6a780256ed 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -39,8 +39,6 @@ from invokeai.backend.model_manager.config import ( CheckpointConfigBase, InvalidModelConfigException, ModelConfigBase, - ModelRepoVariant, - ModelSourceType, ) from invokeai.backend.model_manager.legacy_probe import ModelProbe from invokeai.backend.model_manager.metadata import ( @@ -52,6 +50,7 @@ from invokeai.backend.model_manager.metadata import ( ) from invokeai.backend.model_manager.metadata.metadata_base import HuggingFaceMetadata from invokeai.backend.model_manager.search import ModelSearch +from invokeai.backend.model_manager.taxonomy import ModelRepoVariant, ModelSourceType from invokeai.backend.util import InvokeAILogger from invokeai.backend.util.catch_sigint import catch_sigint from invokeai.backend.util.devices import TorchDevice diff --git a/invokeai/app/services/model_load/model_load_base.py b/invokeai/app/services/model_load/model_load_base.py index cbda6114bc..8aae80e29d 100644 --- a/invokeai/app/services/model_load/model_load_base.py +++ b/invokeai/app/services/model_load/model_load_base.py @@ -5,9 +5,10 @@ from abc import ABC, abstractmethod from pathlib import Path from typing import Callable, Optional -from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType +from invokeai.backend.model_manager.config import AnyModelConfig from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache +from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType class ModelLoadServiceBase(ABC): diff --git a/invokeai/app/services/model_load/model_load_default.py b/invokeai/app/services/model_load/model_load_default.py index 9acf213c43..dac45c7025 100644 --- a/invokeai/app/services/model_load/model_load_default.py +++ b/invokeai/app/services/model_load/model_load_default.py @@ -11,7 +11,7 @@ from torch import load as torch_load from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.invoker import Invoker from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBase -from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType +from invokeai.backend.model_manager.config import AnyModelConfig from invokeai.backend.model_manager.load import ( LoadedModel, LoadedModelWithoutConfig, @@ -20,6 +20,7 @@ from invokeai.backend.model_manager.load import ( ) from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader +from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger diff --git a/invokeai/app/services/model_manager/__init__.py b/invokeai/app/services/model_manager/__init__.py index 07c27cee31..aad67ff352 100644 --- a/invokeai/app/services/model_manager/__init__.py +++ b/invokeai/app/services/model_manager/__init__.py @@ -1,16 +1,12 @@ """Initialization file for model manager service.""" from invokeai.app.services.model_manager.model_manager_default import ModelManagerService, ModelManagerServiceBase -from invokeai.backend.model_manager import AnyModel, AnyModelConfig, BaseModelType, ModelType, SubModelType +from invokeai.backend.model_manager import AnyModelConfig from invokeai.backend.model_manager.load import LoadedModel __all__ = [ "ModelManagerServiceBase", "ModelManagerService", - "AnyModel", "AnyModelConfig", - "BaseModelType", - "ModelType", - "SubModelType", "LoadedModel", ] diff --git a/invokeai/app/services/model_records/model_records_base.py b/invokeai/app/services/model_records/model_records_base.py index 4f40da37d4..4403a1eeb6 100644 --- a/invokeai/app/services/model_records/model_records_base.py +++ b/invokeai/app/services/model_records/model_records_base.py @@ -14,10 +14,12 @@ from invokeai.app.services.shared.pagination import PaginatedResults from invokeai.app.util.model_exclude_null import BaseModelExcludeNull from invokeai.backend.model_manager.config import ( AnyModelConfig, - BaseModelType, - ClipVariantType, ControlAdapterDefaultSettings, MainModelDefaultSettings, +) +from invokeai.backend.model_manager.taxonomy import ( + BaseModelType, + ClipVariantType, ModelFormat, ModelSourceType, ModelType, diff --git a/invokeai/app/services/model_records/model_records_sql.py b/invokeai/app/services/model_records/model_records_sql.py index 5b982e5cdf..3b35124b1e 100644 --- a/invokeai/app/services/model_records/model_records_sql.py +++ b/invokeai/app/services/model_records/model_records_sql.py @@ -60,11 +60,9 @@ from invokeai.app.services.shared.pagination import PaginatedResults from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase from invokeai.backend.model_manager.config import ( AnyModelConfig, - BaseModelType, ModelConfigFactory, - ModelFormat, - ModelType, ) +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType class ModelRecordServiceSQL(ModelRecordServiceBase): diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index 443587ad37..ce49ad0304 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -20,14 +20,10 @@ from invokeai.app.services.session_processor.session_processor_common import Pro from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection from invokeai.app.util.step_callback import flux_step_callback, stable_diffusion_step_callback from invokeai.backend.model_manager.config import ( - AnyModel, AnyModelConfig, - BaseModelType, - ModelFormat, - ModelType, - SubModelType, ) from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig +from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData diff --git a/invokeai/app/util/step_callback.py b/invokeai/app/util/step_callback.py index fd7d6f67ea..21d45280db 100644 --- a/invokeai/app/util/step_callback.py +++ b/invokeai/app/util/step_callback.py @@ -5,7 +5,7 @@ import torch from PIL import Image from invokeai.app.services.session_processor.session_processor_common import CanceledException -from invokeai.backend.model_manager.config import BaseModelType +from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState # fast latents preview matrix for sdxl diff --git a/invokeai/app/util/t5_model_identifier.py b/invokeai/app/util/t5_model_identifier.py index eb3a84aee5..a0d999920c 100644 --- a/invokeai/app/util/t5_model_identifier.py +++ b/invokeai/app/util/t5_model_identifier.py @@ -1,5 +1,5 @@ from invokeai.app.invocations.model import ModelIdentifierField -from invokeai.backend.model_manager.config import BaseModelType, SubModelType +from invokeai.backend.model_manager.taxonomy import BaseModelType, SubModelType def preprocess_t5_encoder_model_identifier(model_identifier: ModelIdentifierField) -> ModelIdentifierField: diff --git a/invokeai/app/util/ti_utils.py b/invokeai/app/util/ti_utils.py index 34669fe64e..8f18b14d66 100644 --- a/invokeai/app/util/ti_utils.py +++ b/invokeai/app/util/ti_utils.py @@ -4,7 +4,7 @@ from typing import List, Tuple import invokeai.backend.util.logging as logger from invokeai.app.services.model_records import UnknownModelException from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.model_manager.config import BaseModelType, ModelType +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType from invokeai.backend.textual_inversion import TextualInversionModelRaw diff --git a/invokeai/backend/image_util/infill_methods/lama.py b/invokeai/backend/image_util/infill_methods/lama.py index faf25e44a4..5b3b685709 100644 --- a/invokeai/backend/image_util/infill_methods/lama.py +++ b/invokeai/backend/image_util/infill_methods/lama.py @@ -6,8 +6,8 @@ import torch from PIL import Image import invokeai.backend.util.logging as logger -from invokeai.backend.model_manager.config import AnyModel from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device +from invokeai.backend.model_manager.taxonomy import AnyModel def norm_img(np_img): diff --git a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/conv2d_layers.py b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/conv2d_layers.py index d8467460c4..2369f7de2c 100644 --- a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/conv2d_layers.py +++ b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/conv2d_layers.py @@ -16,7 +16,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from .config import * +from .config import is_exportable, is_scriptable # From PyTorch internals diff --git a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/efficientnet_builder.py b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/efficientnet_builder.py index 95dd63d400..4da0585389 100644 --- a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/efficientnet_builder.py +++ b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/efficientnet_builder.py @@ -5,8 +5,8 @@ Copyright 2020 Ross Wightman import re from copy import deepcopy -from .conv2d_layers import * -from geffnet.activations import * +from .conv2d_layers import CondConv2d, get_condconv_initializer, math, partial, select_conv2d +from geffnet.activations import F, get_act_layer, nn, sigmoid, torch __all__ = ['get_bn_args_tf', 'resolve_bn_args', 'resolve_se_args', 'resolve_act_layer', 'make_divisible', 'round_channels', 'drop_connect', 'SqueezeExcite', 'ConvBnAct', 'DepthwiseSeparableConv', diff --git a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/gen_efficientnet.py b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/gen_efficientnet.py index cd170d4cc5..029799305a 100644 --- a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/gen_efficientnet.py +++ b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/gen_efficientnet.py @@ -32,7 +32,9 @@ import torch.nn.functional as F from .config import layer_config_kwargs, is_scriptable from .conv2d_layers import select_conv2d from .helpers import load_pretrained -from .efficientnet_builder import * +from .efficientnet_builder import (BN_EPS_TF_DEFAULT, EfficientNetBuilder, decode_arch_def, + initialize_weight_default, initialize_weight_goog, + resolve_act_layer, resolve_bn_args, round_channels) __all__ = ['GenEfficientNet', 'mnasnet_050', 'mnasnet_075', 'mnasnet_100', 'mnasnet_b1', 'mnasnet_140', 'semnasnet_050', 'semnasnet_075', 'semnasnet_100', 'mnasnet_a1', 'semnasnet_140', 'mnasnet_small', diff --git a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/mobilenetv3.py b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/mobilenetv3.py index b5966c28f7..7c09bb3a16 100644 --- a/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/mobilenetv3.py +++ b/invokeai/backend/image_util/normal_bae/nets/submodules/efficientnet_repo/geffnet/mobilenetv3.py @@ -13,7 +13,9 @@ from .activations import get_act_fn, get_act_layer, HardSwish from .config import layer_config_kwargs from .conv2d_layers import select_conv2d from .helpers import load_pretrained -from .efficientnet_builder import * +from .efficientnet_builder import (BN_EPS_TF_DEFAULT, EfficientNetBuilder, decode_arch_def, + initialize_weight_default, initialize_weight_goog, + resolve_act_layer, resolve_bn_args, round_channels) __all__ = ['mobilenetv3_rw', 'mobilenetv3_large_075', 'mobilenetv3_large_100', 'mobilenetv3_large_minimal_100', 'mobilenetv3_small_075', 'mobilenetv3_small_100', 'mobilenetv3_small_minimal_100', diff --git a/invokeai/backend/image_util/realesrgan/realesrgan.py b/invokeai/backend/image_util/realesrgan/realesrgan.py index c5fe3fa598..37853401c2 100644 --- a/invokeai/backend/image_util/realesrgan/realesrgan.py +++ b/invokeai/backend/image_util/realesrgan/realesrgan.py @@ -10,7 +10,7 @@ from cv2.typing import MatLike from tqdm import tqdm from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet -from invokeai.backend.model_manager.config import AnyModel +from invokeai.backend.model_manager.taxonomy import AnyModel from invokeai.backend.util.devices import TorchDevice """ diff --git a/invokeai/backend/model_manager/__init__.py b/invokeai/backend/model_manager/__init__.py index ed29b6589a..dca72f170e 100644 --- a/invokeai/backend/model_manager/__init__.py +++ b/invokeai/backend/model_manager/__init__.py @@ -1,37 +1,45 @@ """Re-export frequently-used symbols from the Model Manager backend.""" from invokeai.backend.model_manager.config import ( - AnyModel, AnyModelConfig, - BaseModelType, InvalidModelConfigException, ModelConfigBase, ModelConfigFactory, +) +from invokeai.backend.model_manager.legacy_probe import ModelProbe +from invokeai.backend.model_manager.load import LoadedModel +from invokeai.backend.model_manager.search import ModelSearch +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, + AnyVariant, + BaseModelType, + ClipVariantType, ModelFormat, ModelRepoVariant, + ModelSourceType, ModelType, ModelVariantType, SchedulerPredictionType, SubModelType, ) -from invokeai.backend.model_manager.legacy_probe import ModelProbe -from invokeai.backend.model_manager.load import LoadedModel -from invokeai.backend.model_manager.search import ModelSearch __all__ = [ - "AnyModel", "AnyModelConfig", - "BaseModelType", - "ModelRepoVariant", "InvalidModelConfigException", "LoadedModel", "ModelConfigFactory", - "ModelFormat", "ModelProbe", "ModelSearch", + "ModelConfigBase", + "AnyModel", + "AnyVariant", + "BaseModelType", + "ClipVariantType", + "ModelFormat", + "ModelRepoVariant", + "ModelSourceType", "ModelType", "ModelVariantType", "SchedulerPredictionType", "SubModelType", - "ModelConfigBase", ] diff --git a/invokeai/backend/model_manager/config.py b/invokeai/backend/model_manager/config.py index 3a9db9a2b5..b150572106 100644 --- a/invokeai/backend/model_manager/config.py +++ b/invokeai/backend/model_manager/config.py @@ -30,11 +30,8 @@ from inspect import isabstract from pathlib import Path from typing import ClassVar, Literal, Optional, TypeAlias, Union -import diffusers -import onnxruntime as ort import safetensors.torch import torch -from diffusers.models.modeling_utils import ModelMixin from picklescan.scanner import scan_file_path from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, TypeAdapter from typing_extensions import Annotated, Any, Dict @@ -42,141 +39,34 @@ from typing_extensions import Annotated, Any, Dict from invokeai.app.util.misc import uuid_string from invokeai.backend.model_hash.hash_validator import validate_hash from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash +from invokeai.backend.model_manager.taxonomy import ( + AnyVariant, + BaseModelType, + ClipVariantType, + ModelFormat, + ModelRepoVariant, + ModelSourceType, + ModelType, + ModelVariantType, + SchedulerPredictionType, + SubModelType, +) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader -from invokeai.backend.raw_model import RawModel from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES from invokeai.backend.util.silence_warnings import SilenceWarnings logger = logging.getLogger(__name__) -# 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, Dict[str, torch.Tensor], diffusers.DiffusionPipeline, ort.InferenceSession -] - class InvalidModelConfigException(Exception): """Exception for when config parser doesn't recognize this combination of model type and format.""" - -class BaseModelType(str, Enum): - """Base model type.""" - - Any = "any" - StableDiffusion1 = "sd-1" - StableDiffusion2 = "sd-2" - StableDiffusion3 = "sd-3" - StableDiffusionXL = "sdxl" - StableDiffusionXLRefiner = "sdxl-refiner" - Flux = "flux" - # Kandinsky2_1 = "kandinsky-2.1" - - -class ModelType(str, Enum): - """Model type.""" - - ONNX = "onnx" - Main = "main" - VAE = "vae" - LoRA = "lora" - ControlLoRa = "control_lora" - ControlNet = "controlnet" # used by model_probe - TextualInversion = "embedding" - IPAdapter = "ip_adapter" - CLIPVision = "clip_vision" - CLIPEmbed = "clip_embed" - T2IAdapter = "t2i_adapter" - T5Encoder = "t5_encoder" - SpandrelImageToImage = "spandrel_image_to_image" - SigLIP = "siglip" - FluxRedux = "flux_redux" - LlavaOnevision = "llava_onevision" - - -class SubModelType(str, Enum): - """Submodel type.""" - - UNet = "unet" - Transformer = "transformer" - TextEncoder = "text_encoder" - TextEncoder2 = "text_encoder_2" - TextEncoder3 = "text_encoder_3" - Tokenizer = "tokenizer" - Tokenizer2 = "tokenizer_2" - Tokenizer3 = "tokenizer_3" - VAE = "vae" - VAEDecoder = "vae_decoder" - VAEEncoder = "vae_encoder" - Scheduler = "scheduler" - SafetyChecker = "safety_checker" - - -class ClipVariantType(str, Enum): - """Variant type.""" - - L = "large" - G = "gigantic" - - -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" - T5Encoder = "t5_encoder" - BnbQuantizedLlmInt8b = "bnb_quantized_int8b" - BnbQuantizednf4b = "bnb_quantized_nf4b" - GGUFQuantized = "gguf_quantized" - - -class SchedulerPredictionType(str, Enum): - """Scheduler prediction type.""" - - Epsilon = "epsilon" - VPrediction = "v_prediction" - Sample = "sample" - - -class ModelRepoVariant(str, Enum): - """Various hugging face variants on the diffusers format.""" - - Default = "" # model files without "fp16" or other qualifier - FP16 = "fp16" - FP32 = "fp32" - ONNX = "onnx" - OpenVINO = "openvino" - Flax = "flax" - - -class ModelSourceType(str, Enum): - """Model source type.""" - - Path = "path" - Url = "url" - HFRepoID = "hf_repo_id" + pass DEFAULTS_PRECISION = Literal["fp16", "fp32"] -AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, None] - - class SubmodelDefinition(BaseModel): path_or_prefix: str model_type: ModelType diff --git a/invokeai/backend/model_manager/legacy_probe.py b/invokeai/backend/model_manager/legacy_probe.py index b37136dfd7..304fbce346 100644 --- a/invokeai/backend/model_manager/legacy_probe.py +++ b/invokeai/backend/model_manager/legacy_probe.py @@ -19,22 +19,24 @@ from invokeai.backend.flux.redux.flux_redux_state_dict_utils import is_state_dic from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash from invokeai.backend.model_manager.config import ( AnyModelConfig, - AnyVariant, - BaseModelType, ControlAdapterDefaultSettings, InvalidModelConfigException, MainModelDefaultSettings, ModelConfigFactory, + SubmodelDefinition, +) +from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import ConfigLoader +from invokeai.backend.model_manager.taxonomy import ( + AnyVariant, + BaseModelType, ModelFormat, ModelRepoVariant, ModelSourceType, ModelType, ModelVariantType, SchedulerPredictionType, - SubmodelDefinition, SubModelType, ) -from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import ConfigLoader from invokeai.backend.model_manager.util.model_util import ( get_clip_variant_type, lora_token_vector_length, diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 94206c8f91..458fc0cfc0 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -13,12 +13,11 @@ import torch from invokeai.app.services.config import InvokeAIAppConfig from invokeai.backend.model_manager.config import ( - AnyModel, AnyModelConfig, - SubModelType, ) from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache +from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType class LoadedModelWithoutConfig: diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index f6cd7972d6..3c26a956b7 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -6,18 +6,16 @@ from pathlib import Path from typing import Optional from invokeai.app.services.config import InvokeAIAppConfig -from invokeai.backend.model_manager import ( - AnyModel, - AnyModelConfig, - InvalidModelConfigException, - SubModelType, -) -from invokeai.backend.model_manager.config import DiffusersConfigBase +from invokeai.backend.model_manager.config import AnyModelConfig, DiffusersConfigBase, InvalidModelConfigException from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache, get_model_cache_key from invokeai.backend.model_manager.load.model_util import calc_model_size_by_fs from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, + SubModelType, +) from invokeai.backend.util.devices import TorchDevice diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 466f3c399c..b6b4d3c7d9 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -9,7 +9,6 @@ from typing import Any, Callable, Dict, List, Optional import psutil import torch -from invokeai.backend.model_manager import AnyModel, SubModelType from invokeai.backend.model_manager.load.memory_snapshot import MemorySnapshot from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats @@ -23,6 +22,7 @@ from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch apply_custom_layers_to_model, ) from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data +from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.prefix_logger_adapter import PrefixedLoggerAdapter diff --git a/invokeai/backend/model_manager/load/model_loader_registry.py b/invokeai/backend/model_manager/load/model_loader_registry.py index 0ce8f8a6b4..ecc4d1fe93 100644 --- a/invokeai/backend/model_manager/load/model_loader_registry.py +++ b/invokeai/backend/model_manager/load/model_loader_registry.py @@ -20,13 +20,10 @@ from typing import Callable, Dict, Optional, Tuple, Type, TypeVar from invokeai.backend.model_manager.config import ( AnyModelConfig, - BaseModelType, ModelConfigBase, - ModelFormat, - ModelType, - SubModelType, ) from invokeai.backend.model_manager.load import ModelLoaderBase +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType class ModelLoaderRegistryBase(ABC): diff --git a/invokeai/backend/model_manager/load/model_loaders/clip_vision.py b/invokeai/backend/model_manager/load/model_loaders/clip_vision.py index cef1c962f9..29d7bc691c 100644 --- a/invokeai/backend/model_manager/load/model_loaders/clip_vision.py +++ b/invokeai/backend/model_manager/load/model_loaders/clip_vision.py @@ -4,16 +4,12 @@ from typing import Optional from transformers import CLIPVisionModelWithProjection from invokeai.backend.model_manager.config import ( - AnyModel, AnyModelConfig, - BaseModelType, DiffusersConfigBase, - ModelFormat, - ModelType, - SubModelType, ) from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers) diff --git a/invokeai/backend/model_manager/load/model_loaders/controlnet.py b/invokeai/backend/model_manager/load/model_loaders/controlnet.py index 5db9520892..5bf93db381 100644 --- a/invokeai/backend/model_manager/load/model_loaders/controlnet.py +++ b/invokeai/backend/model_manager/load/model_loaders/controlnet.py @@ -5,19 +5,19 @@ from typing import Optional from diffusers import ControlNetModel -from invokeai.backend.model_manager import ( - AnyModel, - AnyModelConfig, -) from invokeai.backend.model_manager.config import ( - BaseModelType, + AnyModelConfig, ControlNetCheckpointConfig, +) +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, + BaseModelType, ModelFormat, ModelType, SubModelType, ) -from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry -from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader @ModelLoaderRegistry.register( diff --git a/invokeai/backend/model_manager/load/model_loaders/flux.py b/invokeai/backend/model_manager/load/model_loaders/flux.py index ac3bafc2ff..7e9062aac5 100644 --- a/invokeai/backend/model_manager/load/model_loaders/flux.py +++ b/invokeai/backend/model_manager/load/model_loaders/flux.py @@ -27,15 +27,8 @@ from invokeai.backend.flux.model import Flux from invokeai.backend.flux.modules.autoencoder import AutoEncoder from invokeai.backend.flux.redux.flux_redux_model import FluxReduxModel from invokeai.backend.flux.util import ae_params, params -from invokeai.backend.model_manager import ( - AnyModel, - AnyModelConfig, - BaseModelType, - ModelFormat, - ModelType, - SubModelType, -) from invokeai.backend.model_manager.config import ( + AnyModelConfig, CheckpointConfigBase, CLIPEmbedDiffusersConfig, ControlNetCheckpointConfig, @@ -51,6 +44,13 @@ from invokeai.backend.model_manager.config import ( ) from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, + BaseModelType, + ModelFormat, + ModelType, + SubModelType, +) from invokeai.backend.model_manager.util.model_util import ( convert_bundle_to_flux_transformer_checkpoint, ) 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 4ce51a56d0..8a690583d5 100644 --- a/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py +++ b/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py @@ -8,18 +8,16 @@ from typing import Any, Optional from diffusers.configuration_utils import ConfigMixin from diffusers.models.modeling_utils import ModelMixin -from invokeai.backend.model_manager import ( +from invokeai.backend.model_manager.config import AnyModelConfig, DiffusersConfigBase, InvalidModelConfigException +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import ( AnyModel, - AnyModelConfig, BaseModelType, - InvalidModelConfigException, ModelFormat, ModelType, SubModelType, ) -from invokeai.backend.model_manager.config import DiffusersConfigBase -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.T2IAdapter, 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..d103bc5dbc 100644 --- a/invokeai/backend/model_manager/load/model_loaders/ip_adapter.py +++ b/invokeai/backend/model_manager/load/model_loaders/ip_adapter.py @@ -7,8 +7,9 @@ from typing import Optional 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.config import AnyModelConfig from invokeai.backend.model_manager.load import ModelLoader, ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.raw_model import RawModel diff --git a/invokeai/backend/model_manager/load/model_loaders/llava_onevision.py b/invokeai/backend/model_manager/load/model_loaders/llava_onevision.py index 32c98c6a35..3321382a22 100644 --- a/invokeai/backend/model_manager/load/model_loaders/llava_onevision.py +++ b/invokeai/backend/model_manager/load/model_loaders/llava_onevision.py @@ -3,15 +3,11 @@ from typing import Optional from invokeai.backend.llava_onevision_model import LlavaOnevisionModel from invokeai.backend.model_manager.config import ( - AnyModel, AnyModelConfig, - BaseModelType, - ModelFormat, - ModelType, - SubModelType, ) from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LlavaOnevision, format=ModelFormat.Diffusers) diff --git a/invokeai/backend/model_manager/load/model_loaders/lora.py b/invokeai/backend/model_manager/load/model_loaders/lora.py index 116dee4913..071713316a 100644 --- a/invokeai/backend/model_manager/load/model_loaders/lora.py +++ b/invokeai/backend/model_manager/load/model_loaders/lora.py @@ -9,17 +9,17 @@ import torch from safetensors.torch import load_file from invokeai.app.services.config import InvokeAIAppConfig -from invokeai.backend.model_manager import ( +from invokeai.backend.model_manager.config import AnyModelConfig +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import ( AnyModel, - AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType, ) -from invokeai.backend.model_manager.load.load_default import ModelLoader -from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache -from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import ( is_state_dict_likely_flux_control, lora_model_from_flux_control_state_dict, diff --git a/invokeai/backend/model_manager/load/model_loaders/onnx.py b/invokeai/backend/model_manager/load/model_loaders/onnx.py index 0a5d8477c4..3078d622b4 100644 --- a/invokeai/backend/model_manager/load/model_loaders/onnx.py +++ b/invokeai/backend/model_manager/load/model_loaders/onnx.py @@ -5,16 +5,16 @@ from pathlib import Path from typing import Optional -from invokeai.backend.model_manager import ( +from invokeai.backend.model_manager.config import AnyModelConfig +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader +from invokeai.backend.model_manager.taxonomy import ( AnyModel, - AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType, ) -from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry -from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.ONNX, format=ModelFormat.ONNX) diff --git a/invokeai/backend/model_manager/load/model_loaders/sig_lip_pipeline.py b/invokeai/backend/model_manager/load/model_loaders/sig_lip_pipeline.py index 235a852dd7..7678dda3a7 100644 --- a/invokeai/backend/model_manager/load/model_loaders/sig_lip_pipeline.py +++ b/invokeai/backend/model_manager/load/model_loaders/sig_lip_pipeline.py @@ -2,15 +2,11 @@ from pathlib import Path from typing import Optional from invokeai.backend.model_manager.config import ( - AnyModel, AnyModelConfig, - BaseModelType, - ModelFormat, - ModelType, - SubModelType, ) from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.sig_lip.sig_lip_pipeline import SigLipPipeline diff --git a/invokeai/backend/model_manager/load/model_loaders/spandrel_image_to_image.py b/invokeai/backend/model_manager/load/model_loaders/spandrel_image_to_image.py index 7a57c5cf59..44cb0277fc 100644 --- a/invokeai/backend/model_manager/load/model_loaders/spandrel_image_to_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/spandrel_image_to_image.py @@ -4,15 +4,11 @@ from typing import Optional import torch from invokeai.backend.model_manager.config import ( - AnyModel, AnyModelConfig, - BaseModelType, - ModelFormat, - ModelType, - SubModelType, ) from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel 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 92b80d2c7a..aa692478ca 100644 --- a/invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py +++ b/invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py @@ -11,16 +11,8 @@ from diffusers import ( StableDiffusionXLPipeline, ) -from invokeai.backend.model_manager import ( - AnyModel, - AnyModelConfig, - BaseModelType, - ModelFormat, - ModelType, - ModelVariantType, - SubModelType, -) from invokeai.backend.model_manager.config import ( + AnyModelConfig, CheckpointConfigBase, DiffusersConfigBase, MainCheckpointConfig, @@ -28,6 +20,14 @@ from invokeai.backend.model_manager.config import ( from invokeai.backend.model_manager.load.model_cache.model_cache import get_model_cache_key from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, + BaseModelType, + ModelFormat, + ModelType, + ModelVariantType, + SubModelType, +) from invokeai.backend.util.silence_warnings import SilenceWarnings VARIANT_TO_IN_CHANNEL_MAP = { diff --git a/invokeai/backend/model_manager/load/model_loaders/textual_inversion.py b/invokeai/backend/model_manager/load/model_loaders/textual_inversion.py index 8d0f08f91a..60ae4ea08b 100644 --- a/invokeai/backend/model_manager/load/model_loaders/textual_inversion.py +++ b/invokeai/backend/model_manager/load/model_loaders/textual_inversion.py @@ -4,16 +4,16 @@ from pathlib import Path from typing import Optional -from invokeai.backend.model_manager import ( +from invokeai.backend.model_manager.config import AnyModelConfig +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import ( AnyModel, - AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType, ) -from invokeai.backend.model_manager.load.load_default import ModelLoader -from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.textual_inversion import TextualInversionModelRaw diff --git a/invokeai/backend/model_manager/load/model_loaders/vae.py b/invokeai/backend/model_manager/load/model_loaders/vae.py index bae29ea773..365fa0a547 100644 --- a/invokeai/backend/model_manager/load/model_loaders/vae.py +++ b/invokeai/backend/model_manager/load/model_loaders/vae.py @@ -5,15 +5,16 @@ from typing import Optional from diffusers import AutoencoderKL -from invokeai.backend.model_manager import ( - AnyModelConfig, +from invokeai.backend.model_manager.config import AnyModelConfig, VAECheckpointConfig +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, BaseModelType, ModelFormat, ModelType, + SubModelType, ) -from invokeai.backend.model_manager.config import AnyModel, SubModelType, VAECheckpointConfig -from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry -from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.VAE, format=ModelFormat.Diffusers) diff --git a/invokeai/backend/model_manager/load/model_util.py b/invokeai/backend/model_manager/load/model_util.py index 6d1e463db6..0fa0b99626 100644 --- a/invokeai/backend/model_manager/load/model_util.py +++ b/invokeai/backend/model_manager/load/model_util.py @@ -15,7 +15,7 @@ from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import D from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline from invokeai.backend.ip_adapter.ip_adapter import IPAdapter -from invokeai.backend.model_manager.config import AnyModel +from invokeai.backend.model_manager.taxonomy import AnyModel from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel from invokeai.backend.patches.model_patch_raw import ModelPatchRaw from invokeai.backend.sig_lip.sig_lip_pipeline import SigLipPipeline diff --git a/invokeai/backend/model_manager/metadata/fetch/fetch_base.py b/invokeai/backend/model_manager/metadata/fetch/fetch_base.py index b86a029b3e..1dee3b1a75 100644 --- a/invokeai/backend/model_manager/metadata/fetch/fetch_base.py +++ b/invokeai/backend/model_manager/metadata/fetch/fetch_base.py @@ -17,12 +17,12 @@ from typing import Optional from pydantic.networks import AnyHttpUrl from requests.sessions import Session -from invokeai.backend.model_manager import ModelRepoVariant from invokeai.backend.model_manager.metadata.metadata_base import ( AnyModelRepoMetadata, AnyModelRepoMetadataValidator, BaseMetadata, ) +from invokeai.backend.model_manager.taxonomy import ModelRepoVariant class ModelMetadataFetchBase(ABC): diff --git a/invokeai/backend/model_manager/metadata/fetch/huggingface.py b/invokeai/backend/model_manager/metadata/fetch/huggingface.py index 8c1d6e74ae..a623ab4143 100644 --- a/invokeai/backend/model_manager/metadata/fetch/huggingface.py +++ b/invokeai/backend/model_manager/metadata/fetch/huggingface.py @@ -24,7 +24,6 @@ from huggingface_hub.errors import RepositoryNotFoundError, RevisionNotFoundErro from pydantic.networks import AnyHttpUrl from requests.sessions import Session -from invokeai.backend.model_manager.config import ModelRepoVariant from invokeai.backend.model_manager.metadata.fetch.fetch_base import ModelMetadataFetchBase from invokeai.backend.model_manager.metadata.metadata_base import ( AnyModelRepoMetadata, @@ -32,6 +31,7 @@ from invokeai.backend.model_manager.metadata.metadata_base import ( RemoteModelFile, UnknownMetadataException, ) +from invokeai.backend.model_manager.taxonomy import ModelRepoVariant HF_MODEL_RE = r"https?://huggingface.co/([\w\-.]+/[\w\-.]+)" diff --git a/invokeai/backend/model_manager/metadata/metadata_base.py b/invokeai/backend/model_manager/metadata/metadata_base.py index 97fc598380..d33aa7d91b 100644 --- a/invokeai/backend/model_manager/metadata/metadata_base.py +++ b/invokeai/backend/model_manager/metadata/metadata_base.py @@ -23,7 +23,7 @@ from pydantic.networks import AnyHttpUrl from requests.sessions import Session from typing_extensions import Annotated -from invokeai.backend.model_manager import ModelRepoVariant +from invokeai.backend.model_manager.taxonomy import ModelRepoVariant from invokeai.backend.model_manager.util.select_hf_files import filter_files diff --git a/invokeai/backend/model_manager/starter_models.py b/invokeai/backend/model_manager/starter_models.py index 4b4318a612..bad6bc7aad 100644 --- a/invokeai/backend/model_manager/starter_models.py +++ b/invokeai/backend/model_manager/starter_models.py @@ -2,7 +2,7 @@ from typing import Optional from pydantic import BaseModel -from invokeai.backend.model_manager.config import BaseModelType, ModelFormat, ModelType +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType class StarterModelWithoutDependencies(BaseModel): diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py new file mode 100644 index 0000000000..79134cef32 --- /dev/null +++ b/invokeai/backend/model_manager/taxonomy.py @@ -0,0 +1,129 @@ +from enum import Enum +from typing import Dict, TypeAlias, Union + +import diffusers +import onnxruntime as ort +import torch +from diffusers import ModelMixin + +from invokeai.backend.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, Dict[str, torch.Tensor], diffusers.DiffusionPipeline, ort.InferenceSession +] + + +class BaseModelType(str, Enum): + """Base model type.""" + + Any = "any" + StableDiffusion1 = "sd-1" + StableDiffusion2 = "sd-2" + StableDiffusion3 = "sd-3" + StableDiffusionXL = "sdxl" + StableDiffusionXLRefiner = "sdxl-refiner" + Flux = "flux" + # Kandinsky2_1 = "kandinsky-2.1" + + +class ModelType(str, Enum): + """Model type.""" + + ONNX = "onnx" + Main = "main" + VAE = "vae" + LoRA = "lora" + ControlLoRa = "control_lora" + ControlNet = "controlnet" # used by model_probe + TextualInversion = "embedding" + IPAdapter = "ip_adapter" + CLIPVision = "clip_vision" + CLIPEmbed = "clip_embed" + T2IAdapter = "t2i_adapter" + T5Encoder = "t5_encoder" + SpandrelImageToImage = "spandrel_image_to_image" + SigLIP = "siglip" + FluxRedux = "flux_redux" + LlavaOnevision = "llava_onevision" + + +class SubModelType(str, Enum): + """Submodel type.""" + + UNet = "unet" + Transformer = "transformer" + TextEncoder = "text_encoder" + TextEncoder2 = "text_encoder_2" + TextEncoder3 = "text_encoder_3" + Tokenizer = "tokenizer" + Tokenizer2 = "tokenizer_2" + Tokenizer3 = "tokenizer_3" + VAE = "vae" + VAEDecoder = "vae_decoder" + VAEEncoder = "vae_encoder" + Scheduler = "scheduler" + SafetyChecker = "safety_checker" + + +class ClipVariantType(str, Enum): + """Variant type.""" + + L = "large" + G = "gigantic" + + +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" + T5Encoder = "t5_encoder" + BnbQuantizedLlmInt8b = "bnb_quantized_int8b" + BnbQuantizednf4b = "bnb_quantized_nf4b" + GGUFQuantized = "gguf_quantized" + + +class SchedulerPredictionType(str, Enum): + """Scheduler prediction type.""" + + Epsilon = "epsilon" + VPrediction = "v_prediction" + Sample = "sample" + + +class ModelRepoVariant(str, Enum): + """Various hugging face variants on the diffusers format.""" + + Default = "" # model files without "fp16" or other qualifier + FP16 = "fp16" + FP32 = "fp32" + ONNX = "onnx" + OpenVINO = "openvino" + Flax = "flax" + + +class ModelSourceType(str, Enum): + """Model source type.""" + + Path = "path" + Url = "url" + HFRepoID = "hf_repo_id" + + +AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, None] diff --git a/invokeai/backend/model_manager/util/model_util.py b/invokeai/backend/model_manager/util/model_util.py index 4fc54c34f1..13b36c009d 100644 --- a/invokeai/backend/model_manager/util/model_util.py +++ b/invokeai/backend/model_manager/util/model_util.py @@ -8,7 +8,7 @@ import picklescan.scanner as pscan import safetensors import torch -from invokeai.backend.model_manager.config import ClipVariantType +from invokeai.backend.model_manager.taxonomy import ClipVariantType from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader diff --git a/invokeai/backend/model_manager/util/select_hf_files.py b/invokeai/backend/model_manager/util/select_hf_files.py index 0f5cc763b2..ec42d1550b 100644 --- a/invokeai/backend/model_manager/util/select_hf_files.py +++ b/invokeai/backend/model_manager/util/select_hf_files.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Optional, Set -from invokeai.backend.model_manager.config import ModelRepoVariant +from invokeai.backend.model_manager.taxonomy import ModelRepoVariant def filter_files( diff --git a/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py b/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py index 3835f53d45..67fede9366 100644 --- a/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py +++ b/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py @@ -8,7 +8,7 @@ from diffusers import T2IAdapter from PIL.Image import Image from invokeai.app.util.controlnet_utils import prepare_control_image -from invokeai.backend.model_manager import BaseModelType +from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningMode from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback diff --git a/invokeai/invocation_api/__init__.py b/invokeai/invocation_api/__init__.py index 4d8958c79b..be8fefbae3 100644 --- a/invokeai/invocation_api/__init__.py +++ b/invokeai/invocation_api/__init__.py @@ -71,7 +71,6 @@ from invokeai.app.services.image_records.image_records_common import ImageCatego from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID from invokeai.app.util.misc import SEED_MAX, get_random_seed -from invokeai.backend.model_manager.config import BaseModelType, ModelType, SubModelType from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( @@ -164,9 +163,6 @@ __all__ = [ # invokeai.backend.model_management.model_manager "LoadedModel", # invokeai.backend.model_management.models.base - "BaseModelType", - "ModelType", - "SubModelType", # invokeai.backend.stable_diffusion.schedulers.schedulers "SCHEDULER_NAME_VALUES", # invokeai.version diff --git a/scripts/strip_models.py b/scripts/strip_models.py index 7756cce736..7f6c9da3c7 100644 --- a/scripts/strip_models.py +++ b/scripts/strip_models.py @@ -22,8 +22,9 @@ from pathlib import Path import humanize import torch -from invokeai.backend.model_manager.config import ModelFormat, ModelOnDisk +from invokeai.backend.model_manager.config import ModelOnDisk from invokeai.backend.model_manager.search import ModelSearch +from invokeai.backend.model_manager.taxonomy import ModelFormat def strip(v): diff --git a/tests/app/services/download/test_download_queue.py b/tests/app/services/download/test_download_queue.py index 1c5c8ec6af..edf3c115ac 100644 --- a/tests/app/services/download/test_download_queue.py +++ b/tests/app/services/download/test_download_queue.py @@ -22,7 +22,6 @@ from invokeai.app.services.events.events_common import ( DownloadStartedEvent, ) from invokeai.backend.model_manager.metadata import HuggingFaceMetadataFetch, ModelMetadataWithFiles, RemoteModelFile -from tests.backend.model_manager.model_manager_fixtures import * # noqa F403 from tests.test_nodes import TestEventService # Prevent pytest deprecation warnings diff --git a/tests/app/services/model_records/test_model_records_sql.py b/tests/app/services/model_records/test_model_records_sql.py index e6a89dff06..58131da0eb 100644 --- a/tests/app/services/model_records/test_model_records_sql.py +++ b/tests/app/services/model_records/test_model_records_sql.py @@ -16,19 +16,16 @@ from invokeai.app.services.model_records import ( UnknownModelException, ) from invokeai.app.services.model_records.model_records_base import ModelRecordChanges +from invokeai.backend.model_manager import BaseModelType, ModelFormat, ModelType from invokeai.backend.model_manager.config import ( - BaseModelType, ControlAdapterDefaultSettings, MainDiffusersConfig, MainModelDefaultSettings, - ModelFormat, - ModelSourceType, - ModelType, TextualInversionFileConfig, VAEDiffusersConfig, ) +from invokeai.backend.model_manager.taxonomy import ModelSourceType from invokeai.backend.util.logging import InvokeAILogger -from tests.backend.model_manager.model_manager_fixtures import * # noqa F403 from tests.fixtures.sqlite_database import create_mock_sqlite_database diff --git a/tests/backend/model_manager/model_loading/test_model_load.py b/tests/backend/model_manager/model_loading/test_model_load.py index 3f12f7f8ee..8a331033f6 100644 --- a/tests/backend/model_manager/model_loading/test_model_load.py +++ b/tests/backend/model_manager/model_loading/test_model_load.py @@ -6,7 +6,6 @@ from pathlib import Path from invokeai.app.services.model_manager import ModelManagerServiceBase from invokeai.backend.textual_inversion import TextualInversionModelRaw -from tests.backend.model_manager.model_manager_fixtures import * # noqa F403 def test_loading(mm2_model_manager: ModelManagerServiceBase, embedding_file: Path): diff --git a/tests/backend/model_manager/model_manager_fixtures.py b/tests/backend/model_manager/model_manager_fixtures.py index 87d617662e..59dc71fefe 100644 --- a/tests/backend/model_manager/model_manager_fixtures.py +++ b/tests/backend/model_manager/model_manager_fixtures.py @@ -14,18 +14,15 @@ from invokeai.app.services.model_install import ModelInstallService, ModelInstal from invokeai.app.services.model_load import ModelLoadService, ModelLoadServiceBase from invokeai.app.services.model_manager import ModelManagerService, ModelManagerServiceBase from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL +from invokeai.backend.model_manager import BaseModelType, ModelFormat, ModelType, ModelVariantType from invokeai.backend.model_manager.config import ( - BaseModelType, LoRADiffusersConfig, MainCheckpointConfig, MainDiffusersConfig, - ModelFormat, - ModelSourceType, - ModelType, - ModelVariantType, VAEDiffusersConfig, ) from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache +from invokeai.backend.model_manager.taxonomy import ModelSourceType from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger from tests.backend.model_manager.model_metadata.metadata_examples import ( diff --git a/tests/backend/model_manager/util/test_hf_model_select.py b/tests/backend/model_manager/util/test_hf_model_select.py index 8b5a395fdb..9ef61bdcda 100644 --- a/tests/backend/model_manager/util/test_hf_model_select.py +++ b/tests/backend/model_manager/util/test_hf_model_select.py @@ -3,7 +3,7 @@ from typing import List import pytest -from invokeai.backend.model_manager.config import ModelRepoVariant +from invokeai.backend.model_manager import ModelRepoVariant from invokeai.backend.model_manager.util.select_hf_files import filter_files diff --git a/tests/test_model_probe.py b/tests/test_model_probe.py index 8d7bb8b227..c4a0b34062 100644 --- a/tests/test_model_probe.py +++ b/tests/test_model_probe.py @@ -10,18 +10,14 @@ from polyfactory.factories.pydantic_factory import ModelFactory from sympy.testing.pytest import slow from torch import tensor +from invokeai.backend.model_manager import BaseModelType, ModelFormat, ModelRepoVariant, ModelType, ModelVariantType from invokeai.backend.model_manager.config import ( AnyModelConfig, - BaseModelType, InvalidModelConfigException, MainDiffusersConfig, ModelConfigBase, ModelConfigFactory, - ModelFormat, ModelOnDisk, - ModelRepoVariant, - ModelType, - ModelVariantType, get_model_discriminator_value, ) from invokeai.backend.model_manager.legacy_probe import (