diff --git a/invokeai/backend/model_manager/probe.py b/invokeai/backend/model_manager/probe.py index f2f70f3038..2e433049ff 100644 --- a/invokeai/backend/model_manager/probe.py +++ b/invokeai/backend/model_manager/probe.py @@ -132,7 +132,7 @@ class ModelProbe(object): format_type = ModelFormat.Diffusers if model_path.is_dir() else ModelFormat.Checkpoint model_info = None - model_type = fields['type'] if 'type' in fields else None + model_type = fields["type"] if "type" in fields else None model_type = ModelType(model_type) if isinstance(model_type, str) else model_type if not model_type: if format_type is ModelFormat.Diffusers: diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index b4265da9f7..bb507fb12e 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -5,11 +5,12 @@ Test the model installer import platform import uuid from pathlib import Path +from time import sleep +from typing import Any, Dict import pytest from pydantic import ValidationError from pydantic.networks import Url -from time import sleep from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.events.events_base import EventServiceBase @@ -21,7 +22,7 @@ from invokeai.app.services.model_install import ( URLModelSource, ) from invokeai.app.services.model_records import UnknownModelException -from invokeai.backend.model_manager.config import BaseModelType, ModelFormat, ModelType, InvalidModelConfigException +from invokeai.backend.model_manager.config import BaseModelType, InvalidModelConfigException, ModelFormat, ModelType from tests.backend.model_manager.model_manager_fixtures import * # noqa F403 OS = platform.uname().system @@ -273,13 +274,13 @@ def test_404_download(mm2_installer: ModelInstallServiceBase, mm2_app_config: In { "repo_id": "InvokeAI-test/textual_inversion_tests::learned_embeds-steps-1000.safetensors", "name": "test_lora", - "type": 'embedding', + "type": "embedding", }, # SDXL, Lora - incorrect type { "repo_id": "InvokeAI-test/textual_inversion_tests::learned_embeds-steps-1000.safetensors", "name": "test_lora", - "type": 'lora', + "type": "lora", }, ], ) @@ -289,11 +290,11 @@ def test_heuristic_import_with_type(mm2_installer: ModelInstallServiceBase, mode "type": model_params["type"], } try: - assert("repo_id" in model_params) + assert "repo_id" in model_params install_job = mm2_installer.heuristic_import(source=model_params["repo_id"], config=config) while not install_job.in_terminal_state: - sleep(.01) - assert(install_job.config_out if model_params["type"] == "embedding" else not install_job.config_out) + sleep(0.01) + assert install_job.config_out if model_params["type"] == "embedding" else not install_job.config_out except InvalidModelConfigException: assert model_params["type"] != "embedding" diff --git a/tests/backend/model_manager/model_manager_fixtures.py b/tests/backend/model_manager/model_manager_fixtures.py index afc11c1ac1..16194841b0 100644 --- a/tests/backend/model_manager/model_manager_fixtures.py +++ b/tests/backend/model_manager/model_manager_fixtures.py @@ -33,12 +33,12 @@ from invokeai.backend.model_manager.config import ( from invokeai.backend.model_manager.load import ModelCache, ModelConvertCache from invokeai.backend.util.logging import InvokeAILogger from tests.backend.model_manager.model_metadata.metadata_examples import ( + HFTestLoraMetadata, RepoCivitaiModelMetadata1, RepoCivitaiVersionMetadata1, RepoHFMetadata1, RepoHFMetadata1_nofp16, RepoHFModelJson1, - HFTestLoraMetadata, ) from tests.fixtures.sqlite_database import create_mock_sqlite_database @@ -301,7 +301,7 @@ def mm2_session(embedding_file: Path, diffusers_dir: Path) -> Session: headers={"Content-Type": "application/json; charset=utf-8", "Content-Length": len(RepoHFMetadata1)}, ), ) - + with open(embedding_file, "rb") as f: data = f.read() # file is small - just 15K sess.mount(