mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-02-11 01:45:10 -05:00
5a3195f757
- Replace AnyModelLoader with ModelLoaderRegistry - Fix type check errors in multiple files - Remove apparently unneeded `get_model_config_enum()` method from model manager - Remove last vestiges of old model manager - Updated tests and documentation resolve conflict with seamless.py
91 lines
4.0 KiB
Python
91 lines
4.0 KiB
Python
# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
|
|
"""Class for simple diffusers model loading in InvokeAI."""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
from diffusers import ConfigMixin, ModelMixin
|
|
|
|
from invokeai.backend.model_manager import (
|
|
AnyModel,
|
|
BaseModelType,
|
|
InvalidModelConfigException,
|
|
ModelFormat,
|
|
ModelRepoVariant,
|
|
ModelType,
|
|
SubModelType,
|
|
)
|
|
|
|
from .. import ModelLoader, ModelLoaderRegistry
|
|
|
|
|
|
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers)
|
|
@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.T2IAdapter, format=ModelFormat.Diffusers)
|
|
class GenericDiffusersLoader(ModelLoader):
|
|
"""Class to load simple diffusers models."""
|
|
|
|
def _load_model(
|
|
self,
|
|
model_path: Path,
|
|
model_variant: Optional[ModelRepoVariant] = None,
|
|
submodel_type: Optional[SubModelType] = None,
|
|
) -> AnyModel:
|
|
model_class = self.get_hf_load_class(model_path)
|
|
if submodel_type is not None:
|
|
raise Exception(f"There are no submodels in models of type {model_class}")
|
|
variant = model_variant.value if model_variant else None
|
|
result: AnyModel = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, variant=variant) # type: ignore
|
|
return result
|
|
|
|
# TO DO: Add exception handling
|
|
def get_hf_load_class(self, model_path: Path, submodel_type: Optional[SubModelType] = None) -> ModelMixin:
|
|
"""Given the model path and submodel, returns the diffusers ModelMixin subclass needed to load."""
|
|
if submodel_type:
|
|
try:
|
|
config = self._load_diffusers_config(model_path, config_name="model_index.json")
|
|
module, class_name = config[submodel_type.value]
|
|
result = self._hf_definition_to_type(module=module, class_name=class_name)
|
|
except KeyError as e:
|
|
raise InvalidModelConfigException(
|
|
f'The "{submodel_type}" submodel is not available for this model.'
|
|
) from e
|
|
else:
|
|
try:
|
|
config = self._load_diffusers_config(model_path, config_name="config.json")
|
|
class_name = config.get("_class_name", None)
|
|
if class_name:
|
|
result = self._hf_definition_to_type(module="diffusers", class_name=class_name)
|
|
if config.get("model_type", None) == "clip_vision_model":
|
|
class_name = config.get("architectures")
|
|
assert class_name is not None
|
|
result = self._hf_definition_to_type(module="transformers", class_name=class_name[0])
|
|
if not class_name:
|
|
raise InvalidModelConfigException("Unable to decifer Load Class based on given config.json")
|
|
except KeyError as e:
|
|
raise InvalidModelConfigException("An expected config.json file is missing from this model.") from e
|
|
return result
|
|
|
|
# TO DO: Add exception handling
|
|
def _hf_definition_to_type(self, module: str, class_name: str) -> ModelMixin: # fix with correct type
|
|
if module in ["diffusers", "transformers"]:
|
|
res_type = sys.modules[module]
|
|
else:
|
|
res_type = sys.modules["diffusers"].pipelines
|
|
result: ModelMixin = getattr(res_type, class_name)
|
|
return result
|
|
|
|
def _load_diffusers_config(self, model_path: Path, config_name: str = "config.json") -> Dict[str, Any]:
|
|
return ConfigLoader.load_config(model_path, config_name=config_name)
|
|
|
|
|
|
class ConfigLoader(ConfigMixin):
|
|
"""Subclass of ConfigMixin for loading diffusers configuration files."""
|
|
|
|
@classmethod
|
|
def load_config(cls, *args: Any, **kwargs: Any) -> Dict[str, Any]:
|
|
"""Load a diffusrs ConfigMixin configuration."""
|
|
cls.config_name = kwargs.pop("config_name")
|
|
# Diffusers doesn't provide typing info
|
|
return super().load_config(*args, **kwargs) # type: ignore
|