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
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
search = ModelSearch(on_model_found=callback, config=self._app_config)
search = ModelSearch(on_model_found=callback)
self._models_installed.clear()
search.search(scan_dir)
return list(self._models_installed)

View File

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