"""Utility (backend) functions used by model_install.py""" from logging import Logger from pathlib import Path from typing import Any, Dict, List, Optional import omegaconf from pydantic import BaseModel, Field from pydantic.dataclasses import dataclass from requests import HTTPError from tqdm import tqdm import invokeai.configs as configs from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.download import DownloadQueueService from invokeai.app.services.events.events_base import EventServiceBase from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage from invokeai.app.services.model_install import ( ModelInstallService, ModelInstallServiceBase, ) from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL from invokeai.app.services.shared.sqlite.sqlite_util import init_db from invokeai.backend.model_manager import ( BaseModelType, InvalidModelConfigException, ModelType, ) from invokeai.backend.model_manager.metadata import UnknownMetadataException from invokeai.backend.util.logging import InvokeAILogger # name of the starter models file INITIAL_MODELS = "INITIAL_MODELS.yaml" def initialize_record_store(app_config: InvokeAIAppConfig) -> ModelRecordServiceBase: """Return an initialized ModelConfigRecordServiceBase object.""" logger = InvokeAILogger.get_logger(config=app_config) image_files = DiskImageFileStorage(f"{app_config.output_path}/images") db = init_db(config=app_config, logger=logger, image_files=image_files) obj: ModelRecordServiceBase = ModelRecordServiceSQL(db) return obj def initialize_installer( app_config: InvokeAIAppConfig, event_bus: Optional[EventServiceBase] = None ) -> ModelInstallServiceBase: """Return an initialized ModelInstallService object.""" record_store = initialize_record_store(app_config) download_queue = DownloadQueueService() installer = ModelInstallService( app_config=app_config, record_store=record_store, download_queue=download_queue, event_bus=event_bus, ) download_queue.start() installer.start() return installer class UnifiedModelInfo(BaseModel): """Catchall class for information in INITIAL_MODELS2.yaml.""" name: Optional[str] = None base: Optional[BaseModelType] = None type: Optional[ModelType] = None source: Optional[str] = None subfolder: Optional[str] = None description: Optional[str] = None recommended: bool = False installed: bool = False default: bool = False requires: List[str] = Field(default_factory=list) @dataclass class InstallSelections: """Lists of models to install and remove.""" install_models: List[UnifiedModelInfo] = Field(default_factory=list) remove_models: List[str] = Field(default_factory=list) class TqdmEventService(EventServiceBase): """An event service to track downloads.""" def __init__(self) -> None: """Create a new TqdmEventService object.""" super().__init__() self._bars: Dict[str, tqdm] = {} self._last: Dict[str, int] = {} self._logger = InvokeAILogger.get_logger(__name__) def dispatch(self, event_name: str, payload: Any) -> None: """Dispatch an event by appending it to self.events.""" data = payload["data"] source = data["source"] if payload["event"] == "model_install_downloading": dest = data["local_path"] total_bytes = data["total_bytes"] bytes = data["bytes"] if dest not in self._bars: self._bars[dest] = tqdm(desc=Path(dest).name, initial=0, total=total_bytes, unit="iB", unit_scale=True) self._last[dest] = 0 self._bars[dest].update(bytes - self._last[dest]) self._last[dest] = bytes elif payload["event"] == "model_install_completed": self._logger.info(f"{source}: installed successfully.") elif payload["event"] == "model_install_error": self._logger.warning(f"{source}: installation failed with error {data['error']}") elif payload["event"] == "model_install_cancelled": self._logger.warning(f"{source}: installation cancelled") class InstallHelper(object): """Capture information stored jointly in INITIAL_MODELS.yaml and the installed models db.""" def __init__(self, app_config: InvokeAIAppConfig, logger: Logger): """Create new InstallHelper object.""" self._app_config = app_config self.all_models: Dict[str, UnifiedModelInfo] = {} omega = omegaconf.OmegaConf.load(Path(configs.__path__[0]) / INITIAL_MODELS) assert isinstance(omega, omegaconf.dictconfig.DictConfig) self._installer = initialize_installer(app_config, TqdmEventService()) self._initial_models = omega self._installed_models: List[str] = [] self._starter_models: List[str] = [] self._default_model: Optional[str] = None self._logger = logger self._initialize_model_lists() @property def installer(self) -> ModelInstallServiceBase: """Return the installer object used internally.""" return self._installer def _initialize_model_lists(self) -> None: """ Initialize our model slots. Set up the following: installed_models -- list of installed model keys starter_models -- list of starter model keys from INITIAL_MODELS all_models -- dict of key => UnifiedModelInfo default_model -- key to default model """ # previously-installed models for model in self._installer.record_store.all_models(): info = UnifiedModelInfo.parse_obj(model.dict()) info.installed = True model_key = f"{model.base.value}/{model.type.value}/{model.name}" self.all_models[model_key] = info self._installed_models.append(model_key) for key in self._initial_models.keys(): assert isinstance(key, str) if key in self.all_models: # we want to preserve the description description = self.all_models[key].description or self._initial_models[key].get("description") self.all_models[key].description = description else: base_model, model_type, model_name = key.split("/") info = UnifiedModelInfo( name=model_name, type=ModelType(model_type), base=BaseModelType(base_model), source=self._initial_models[key].source, description=self._initial_models[key].get("description"), recommended=self._initial_models[key].get("recommended", False), default=self._initial_models[key].get("default", False), subfolder=self._initial_models[key].get("subfolder"), requires=list(self._initial_models[key].get("requires", [])), ) self.all_models[key] = info if not self.default_model(): self._default_model = key elif self._initial_models[key].get("default", False): self._default_model = key self._starter_models.append(key) # previously-installed models for model in self._installer.record_store.all_models(): info = UnifiedModelInfo.parse_obj(model.dict()) info.installed = True model_key = f"{model.base.value}/{model.type.value}/{model.name}" self.all_models[model_key] = info self._installed_models.append(model_key) def recommended_models(self) -> List[UnifiedModelInfo]: """List of the models recommended in INITIAL_MODELS.yaml.""" return [self._to_model(x) for x in self._starter_models if self._to_model(x).recommended] def installed_models(self) -> List[UnifiedModelInfo]: """List of models already installed.""" return [self._to_model(x) for x in self._installed_models] def starter_models(self) -> List[UnifiedModelInfo]: """List of starter models.""" return [self._to_model(x) for x in self._starter_models] def default_model(self) -> Optional[UnifiedModelInfo]: """Return the default model.""" return self._to_model(self._default_model) if self._default_model else None def _to_model(self, key: str) -> UnifiedModelInfo: return self.all_models[key] def _add_required_models(self, model_list: List[UnifiedModelInfo]) -> None: installed = {x.source for x in self.installed_models()} reverse_source = {x.source: x for x in self.all_models.values()} additional_models: List[UnifiedModelInfo] = [] for model_info in model_list: for requirement in model_info.requires: if requirement not in installed and reverse_source.get(requirement): additional_models.append(reverse_source[requirement]) model_list.extend(additional_models) def add_or_delete(self, selections: InstallSelections) -> None: """Add or delete selected models.""" installer = self._installer self._add_required_models(selections.install_models) for model in selections.install_models: assert model.source model_path_id_or_url = model.source.strip("\"' ") config = ( { "description": model.description, "name": model.name, } if model.name else None ) try: installer.heuristic_import( source=model_path_id_or_url, config=config, ) except (UnknownMetadataException, InvalidModelConfigException, HTTPError, OSError) as e: self._logger.warning(f"{model.source}: {e}") for model_to_remove in selections.remove_models: parts = model_to_remove.split("/") if len(parts) == 1: base_model, model_type, model_name = (None, None, model_to_remove) else: base_model, model_type, model_name = parts matches = installer.record_store.search_by_attr( base_model=BaseModelType(base_model) if base_model else None, model_type=ModelType(model_type) if model_type else None, model_name=model_name, ) if len(matches) > 1: self._logger.error( "{model_to_remove} is ambiguous. Please use model_base/model_type/model_name (e.g. sd-1/main/my_model) to disambiguate" ) elif not matches: self._logger.error(f"{model_to_remove}: unknown model") else: for m in matches: self._logger.info(f"Deleting {m.type}:{m.name}") installer.delete(m.key) installer.wait_for_installs()