From c025c9c4ed3207d037c4064837c63dab5f281052 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 30 Sep 2023 13:57:13 -0400 Subject: [PATCH] speed up model scanning at startup --- invokeai/backend/model_manager/install.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/invokeai/backend/model_manager/install.py b/invokeai/backend/model_manager/install.py index 040eb2ceac..2ea7e4e54a 100644 --- a/invokeai/backend/model_manager/install.py +++ b/invokeai/backend/model_manager/install.py @@ -328,6 +328,7 @@ class ModelInstall(ModelInstallBase): _async_installs: Dict[Union[str, Path, AnyHttpUrl], Optional[str]] _installed: Set[str] = Field(default=set) _tmpdir: Optional[tempfile.TemporaryDirectory] # used for downloads + _cached_model_paths: Set[Path] = Field(default=set) # used to speed up directory scanning _legacy_configs: Dict[BaseModelType, Dict[ModelVariantType, Union[str, dict]]] = { BaseModelType.StableDiffusion1: { @@ -451,7 +452,9 @@ class ModelInstall(ModelInstallBase): # if path already exists then we jigger the name to make it unique counter: int = 1 while new_path.exists(): - new_path = new_path.with_stem(new_path.stem + f"_{counter:02d}") + path = new_path.with_stem(new_path.stem + f"_{counter:02d}") + if not path.exists(): + new_path = path counter += 1 return old_path.replace(new_path) @@ -630,6 +633,7 @@ class ModelInstall(ModelInstallBase): return id_map def scan_directory(self, scan_dir: Path, install: bool = False) -> List[str]: # noqa D102 + self._cached_model_paths = set([Path(x.path) for x in self._store.all_models()]) callback = self._scan_install if install else self._scan_register search = ModelSearch(on_model_found=callback) self._installed = set() @@ -704,6 +708,8 @@ class ModelInstall(ModelInstallBase): # the following two methods are callbacks to the ModelSearch object def _scan_register(self, model: Path) -> bool: + if model in self._cached_model_paths: + return True try: id = self.register_path(model) self.sync_model_path(id) # possibly move it to right place in `models` @@ -714,6 +720,8 @@ class ModelInstall(ModelInstallBase): return True def _scan_install(self, model: Path) -> bool: + if model in self._cached_model_paths: + return True try: id = self.install_path(model) self._logger.info(f"Installed {model} with id {id}")