diff --git a/invokeai/backend/model_manager/download/base.py b/invokeai/backend/model_manager/download/base.py index 66a34636ca..07e90261bd 100644 --- a/invokeai/backend/model_manager/download/base.py +++ b/invokeai/backend/model_manager/download/base.py @@ -9,9 +9,12 @@ from functools import total_ordering from pathlib import Path from typing import Callable, List, Optional, Union +import requests from pydantic import BaseModel, Field from pydantic.networks import AnyHttpUrl +from invokeai.app.services.config import InvokeAIAppConfig + # Used to distinguish between repo_id sources and URL sources REPO_ID_RE = r"^[\w-]+/[.\w-]+$" HTTP_RE = r"^https?://" @@ -107,6 +110,26 @@ class DownloadJobBase(BaseModel): class DownloadQueueBase(ABC): """Abstract base class for managing model downloads.""" + @abstractmethod + def __init__( + self, + max_parallel_dl: int = 5, + event_handlers: List[DownloadEventHandler] = [], + requests_session: Optional[requests.sessions.Session] = None, + config: Optional[InvokeAIAppConfig] = None, + quiet: bool = False, + ): + """ + Initialize DownloadQueue. + + :param max_parallel_dl: Number of simultaneous downloads allowed [5]. + :param event_handler: Optional callable that will be called each time a job status changes. + :param requests_session: Optional requests.sessions.Session object, for unit tests. + :param config: InvokeAIAppConfig object, used to configure the logger and other options. + :param quiet: If true, don't log the start of download jobs. Useful for subrequests. + """ + pass + @abstractmethod def create_download_job( self, diff --git a/invokeai/backend/model_manager/download/queue.py b/invokeai/backend/model_manager/download/queue.py index ce1e0ea72d..b4d67c08fa 100644 --- a/invokeai/backend/model_manager/download/queue.py +++ b/invokeai/backend/model_manager/download/queue.py @@ -359,11 +359,17 @@ class DownloadQueue(DownloadQueueBase): metadata.author = metadata.author or resp["creator"]["username"] metadata.tags = metadata.tags or resp["tags"] - metadata.thumbnail_url = metadata.thumbnail_url or resp["modelVersions"][0]["images"][0]["url"] metadata.license = ( metadata.license or f"allowCommercialUse={resp['allowCommercialUse']}; allowDerivatives={resp['allowDerivatives']}; allowNoCredit={resp['allowNoCredit']}" ) + first_version = resp["modelVersions"][0] + metadata.thumbnail_url = metadata.thumbnail_url or first_version.get("url") + metadata.description = metadata.description or ( + f"Trigger terms: {(', ').join(first_version.get('trainedWords'))}" + if first_version.get("trainedWords") + else first_version.get("description") + ) except (HTTPError, KeyError, TypeError, JSONDecodeError) as excp: self._logger.warn(excp) diff --git a/invokeai/frontend/install/model_install.py b/invokeai/frontend/install/model_install.py index bf6d2545c9..90caa9d24e 100644 --- a/invokeai/frontend/install/model_install.py +++ b/invokeai/frontend/install/model_install.py @@ -79,7 +79,7 @@ class UnifiedModelInfo(BaseModel): @dataclass class InstallSelections: install_models: List[UnifiedModelInfo] = field(default_factory=list) - remove_models: List[UnifiedModelInfo] = field(default_factory=list) + remove_models: List[str] = field(default_factory=list) def make_printable(s: str) -> str: @@ -576,7 +576,11 @@ def add_or_delete(installer: ModelInstall, selections: InstallSelections): ) for model in selections.remove_models: - base_model, model_type, model_name = model.split("/") + parts = model.split("/") + if len(parts) == 1: + base_model, model_type, model_name = (None, None, model) + else: + base_model, model_type, model_name = parts matches = installer.store.search_by_name(base_model=base_model, model_type=model_type, model_name=model_name) if len(matches) > 1: print(f"{model} is ambiguous. Please use model_type:model_name (e.g. main:my_model) to disambiguate.") @@ -601,7 +605,9 @@ def select_and_download_models(opt: Namespace): list_models(installer, opt.list_models) elif opt.add or opt.delete: - selections = InstallSelections(install_models=opt.add, remove_models=opt.delete) + selections = InstallSelections( + install_models=[UnifiedModelInfo(source=x) for x in (opt.add or [])], remove_models=opt.delete or [] + ) add_or_delete(installer, selections) elif opt.default_only: