tidy(mm): remove extraneous dependencies in model search

- `config` is unused
- `stats` is created on instantiation
- `logger` uses the app logger
This commit is contained in:
psychedelicious 2024-03-10 11:58:09 +11:00
parent bd5b43c00d
commit 56e7c04475
2 changed files with 3 additions and 11 deletions

View File

@ -334,7 +334,7 @@ class ModelInstallService(ModelInstallServiceBase):
def scan_directory(self, scan_dir: Path, install: bool = False) -> List[str]: # noqa D102 def scan_directory(self, scan_dir: Path, install: bool = False) -> List[str]: # noqa D102
self._cached_model_paths = {Path(x.path).absolute() for x in self.record_store.all_models()} self._cached_model_paths = {Path(x.path).absolute() for x in self.record_store.all_models()}
callback = self._scan_install if install else self._scan_register callback = self._scan_install if install else self._scan_register
search = ModelSearch(on_model_found=callback, config=self._app_config) search = ModelSearch(on_model_found=callback)
self._models_installed.clear() self._models_installed.clear()
search.search(scan_dir) search.search(scan_dir)
return list(self._models_installed) return list(self._models_installed)

View File

@ -22,15 +22,11 @@ Example usage:
import os import os
from dataclasses import dataclass from dataclasses import dataclass
from logging import Logger
from pathlib import Path from pathlib import Path
from typing import Callable, Optional, Set, Union from typing import Callable, Optional, Set, Union
from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.logging import InvokeAILogger
default_logger: Logger = InvokeAILogger.get_logger()
@dataclass @dataclass
class SearchStats: class SearchStats:
@ -59,12 +55,9 @@ class ModelSearch:
def __init__( def __init__(
self, self,
stats: Optional[SearchStats] = None,
logger: Optional[Logger] = None,
on_search_started: Optional[Callable[[Path], None]] = None, on_search_started: Optional[Callable[[Path], None]] = None,
on_model_found: Optional[Callable[[Path], bool]] = None, on_model_found: Optional[Callable[[Path], bool]] = None,
on_search_completed: Optional[Callable[[Set[Path]], None]] = None, on_search_completed: Optional[Callable[[Set[Path]], None]] = None,
config: Optional[InvokeAIAppConfig] = None,
) -> None: ) -> None:
"""Create a new ModelSearch object. """Create a new ModelSearch object.
@ -77,13 +70,12 @@ class ModelSearch:
on_search_completed: callback to be invoked when the search is completed on_search_completed: callback to be invoked when the search is completed
config: configuration object config: configuration object
""" """
self.stats = stats or SearchStats() self.stats = SearchStats()
self.logger = logger or default_logger self.logger = InvokeAILogger.get_logger()
self.on_search_started = on_search_started self.on_search_started = on_search_started
self.on_model_found = on_model_found self.on_model_found = on_model_found
self.on_search_completed = on_search_completed self.on_search_completed = on_search_completed
self.models_found: set[Path] = set() self.models_found: set[Path] = set()
self.config = config or InvokeAIAppConfig.get_config()
def search_started(self) -> None: def search_started(self) -> None:
self.models_found = set() self.models_found = set()