Merge branch 'main' into tiled-upscaling-graph

tests/aa_nodes/test_graph_execution_state.py

@@ -69,6 +69,7 @@ def mock_services() -> InvocationServices:
         logger=logging,  # type: ignore
         model_manager=None,  # type: ignore
         model_records=None,  # type: ignore
+        model_install=None,  # type: ignore
         names=None,  # type: ignore
         performance_statistics=InvocationStatsService(),
         processor=DefaultInvocationProcessor(),

tests/aa_nodes/test_invoker.py

@@ -74,6 +74,7 @@ def mock_services() -> InvocationServices:
         logger=logging,  # type: ignore
         model_manager=None,  # type: ignore
         model_records=None,  # type: ignore
+        model_install=None,  # type: ignore
         names=None,  # type: ignore
         performance_statistics=InvocationStatsService(),
         processor=DefaultInvocationProcessor(),

tests/aa_nodes/test_session_queue.py

@@ -12,7 +12,7 @@ from invokeai.app.services.session_queue.session_queue_common import (
     prepare_values_to_insert,
 )
 from invokeai.app.services.shared.graph import Graph, GraphExecutionState, GraphInvocation
-from tests.nodes.test_nodes import PromptTestInvocation
+from tests.aa_nodes.test_nodes import PromptTestInvocation
 
 
 @pytest.fixture

tests/app/services/model_install/test_model_install.py

@@ -0,0 +1,195 @@
+"""
+Test the model installer
+"""
+
+from pathlib import Path
+from typing import Any, Dict, List
+
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.events.events_base import EventServiceBase
+from invokeai.app.services.model_install import (
+    InstallStatus,
+    LocalModelSource,
+    ModelInstallJob,
+    ModelInstallService,
+    ModelInstallServiceBase,
+)
+from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL, UnknownModelException
+from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+from invokeai.backend.model_manager.config import BaseModelType, ModelType
+from invokeai.backend.util.logging import InvokeAILogger
+
+
+@pytest.fixture
+def test_file(datadir: Path) -> Path:
+    return datadir / "test_embedding.safetensors"
+
+
+@pytest.fixture
+def app_config(datadir: Path) -> InvokeAIAppConfig:
+    return InvokeAIAppConfig(
+        root=datadir / "root",
+        models_dir=datadir / "root/models",
+    )
+
+
+@pytest.fixture
+def store(app_config: InvokeAIAppConfig) -> ModelRecordServiceBase:
+    database = SqliteDatabase(app_config, InvokeAILogger.get_logger(config=app_config))
+    store: ModelRecordServiceBase = ModelRecordServiceSQL(database)
+    return store
+
+
+@pytest.fixture
+def installer(app_config: InvokeAIAppConfig, store: ModelRecordServiceBase) -> ModelInstallServiceBase:
+    return ModelInstallService(
+        app_config=app_config,
+        record_store=store,
+        event_bus=DummyEventService(),
+    )
+
+
+class DummyEvent(BaseModel):
+    """Dummy Event to use with Dummy Event service."""
+
+    event_name: str
+    payload: Dict[str, Any]
+
+
+class DummyEventService(EventServiceBase):
+    """Dummy event service for testing."""
+
+    events: List[DummyEvent]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.events = []
+
+    def dispatch(self, event_name: str, payload: Any) -> None:
+        """Dispatch an event by appending it to self.events."""
+        self.events.append(DummyEvent(event_name=payload["event"], payload=payload["data"]))
+
+
+def test_registration(installer: ModelInstallServiceBase, test_file: Path) -> None:
+    store = installer.record_store
+    matches = store.search_by_attr(model_name="test_embedding")
+    assert len(matches) == 0
+    key = installer.register_path(test_file)
+    assert key is not None
+    assert len(key) == 32
+
+
+def test_registration_meta(installer: ModelInstallServiceBase, test_file: Path) -> None:
+    store = installer.record_store
+    key = installer.register_path(test_file)
+    model_record = store.get_model(key)
+    assert model_record is not None
+    assert model_record.name == "test_embedding"
+    assert model_record.type == ModelType.TextualInversion
+    assert Path(model_record.path) == test_file
+    assert model_record.base == BaseModelType("sd-1")
+    assert model_record.description is not None
+    assert model_record.source is not None
+    assert Path(model_record.source) == test_file
+
+
+def test_registration_meta_override_fail(installer: ModelInstallServiceBase, test_file: Path) -> None:
+    key = None
+    with pytest.raises(ValidationError):
+        key = installer.register_path(test_file, {"name": "banana_sushi", "type": ModelType("lora")})
+    assert key is None
+
+
+def test_registration_meta_override_succeed(installer: ModelInstallServiceBase, test_file: Path) -> None:
+    store = installer.record_store
+    key = installer.register_path(
+        test_file, {"name": "banana_sushi", "source": "fake/repo_id", "current_hash": "New Hash"}
+    )
+    model_record = store.get_model(key)
+    assert model_record.name == "banana_sushi"
+    assert model_record.source == "fake/repo_id"
+    assert model_record.current_hash == "New Hash"
+
+
+def test_install(installer: ModelInstallServiceBase, test_file: Path, app_config: InvokeAIAppConfig) -> None:
+    store = installer.record_store
+    key = installer.install_path(test_file)
+    model_record = store.get_model(key)
+    assert model_record.path == "sd-1/embedding/test_embedding.safetensors"
+    assert model_record.source == test_file.as_posix()
+
+
+def test_background_install(installer: ModelInstallServiceBase, test_file: Path, app_config: InvokeAIAppConfig) -> None:
+    """Note: may want to break this down into several smaller unit tests."""
+    path = test_file
+    description = "Test of metadata assignment"
+    source = LocalModelSource(path=path, inplace=False)
+    job = installer.import_model(source, config={"description": description})
+    assert job is not None
+    assert isinstance(job, ModelInstallJob)
+
+    # See if job is registered properly
+    assert job in installer.get_job(source)
+
+    # test that the job object tracked installation correctly
+    jobs = installer.wait_for_installs()
+    assert len(jobs) > 0
+    my_job = [x for x in jobs if x.source == source]
+    assert len(my_job) == 1
+    assert my_job[0].status == InstallStatus.COMPLETED
+
+    # test that the expected events were issued
+    bus = installer.event_bus
+    assert bus is not None  # sigh - ruff is a stickler for type checking
+    assert isinstance(bus, DummyEventService)
+    assert len(bus.events) == 2
+    event_names = [x.event_name for x in bus.events]
+    assert "model_install_started" in event_names
+    assert "model_install_completed" in event_names
+    assert Path(bus.events[0].payload["source"]) == source
+    assert Path(bus.events[1].payload["source"]) == source
+    key = bus.events[1].payload["key"]
+    assert key is not None
+
+    # see if the thing actually got installed at the expected location
+    model_record = installer.record_store.get_model(key)
+    assert model_record is not None
+    assert model_record.path == "sd-1/embedding/test_embedding.safetensors"
+    assert Path(app_config.models_dir / model_record.path).exists()
+
+    # see if metadata was properly passed through
+    assert model_record.description == description
+
+    # see if prune works properly
+    installer.prune_jobs()
+    assert not installer.get_job(source)
+
+
+def test_delete_install(installer: ModelInstallServiceBase, test_file: Path, app_config: InvokeAIAppConfig):
+    store = installer.record_store
+    key = installer.install_path(test_file)
+    model_record = store.get_model(key)
+    assert Path(app_config.models_dir / model_record.path).exists()
+    assert test_file.exists()  # original should still be there after installation
+    installer.delete(key)
+    assert not Path(
+        app_config.models_dir / model_record.path
+    ).exists()  # after deletion, installed copy should not exist
+    assert test_file.exists()  # but original should still be there
+    with pytest.raises(UnknownModelException):
+        store.get_model(key)
+
+
+def test_delete_register(installer: ModelInstallServiceBase, test_file: Path, app_config: InvokeAIAppConfig):
+    store = installer.record_store
+    key = installer.register_path(test_file)
+    model_record = store.get_model(key)
+    assert Path(app_config.models_dir / model_record.path).exists()
+    assert test_file.exists()  # original should still be there after installation
+    installer.delete(key)
+    assert Path(app_config.models_dir / model_record.path).exists()
+    with pytest.raises(UnknownModelException):
+        store.get_model(key)

tests/app/services/model_install/test_model_install/README

@@ -0,0 +1 @@
+This directory is used by pytest-datadir.

tests/app/services/model_install/test_model_install/root/configs/stable-diffusion/v1-inference.yaml

@@ -0,0 +1,79 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: invokeai.backend.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False
+
+    scheduler_config: # 10000 warmup steps
+      target: invokeai.backend.stable_diffusion.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 10000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    personalization_config:
+      target: invokeai.backend.stable_diffusion.embedding_manager.EmbeddingManager
+      params:
+        placeholder_strings: ["*"]
+        initializer_words: ['sculpture']
+        per_image_tokens: false
+        num_vectors_per_token: 1
+        progressive_words: False
+
+    unet_config:
+      target: invokeai.backend.stable_diffusion.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: invokeai.backend.stable_diffusion.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: invokeai.backend.stable_diffusion.encoders.modules.WeightedFrozenCLIPEmbedder

tests/app/services/model_install/test_model_install/root/models/placeholder

@@ -0,0 +1 @@
+Dummy file to establish git path.

tests/test_path.py

@@ -10,7 +10,6 @@ from PIL import Image
 
 import invokeai.app.assets.images as image_assets
 import invokeai.configs as configs
-import invokeai.frontend.web.dist as frontend
 
 
 class ConfigsTestCase(unittest.TestCase):
@@ -21,20 +20,11 @@ class ConfigsTestCase(unittest.TestCase):
         configs_path = pathlib.Path(configs.__path__[0])
         return configs_path
 
-    def get_frontend_path(self) -> pathlib.Path:
-        """Get the path of the frontend dist folder"""
-        return pathlib.Path(frontend.__path__[0])
-
     def test_configs_path(self):
         """Test that the configs path is correct"""
         TEST_PATH = str(self.get_configs_path())
         assert TEST_PATH.endswith(str(osp.join("invokeai", "configs")))
 
-    def test_frontend_path(self):
-        """Test that the frontend path is correct"""
-        FRONTEND_PATH = str(self.get_frontend_path())
-        assert FRONTEND_PATH.endswith(osp.join("invokeai", "frontend", "web", "dist"))
-
     def test_caution_img(self):
         """Verify the caution image"""
         caution_img = Image.open(osp.join(image_assets.__path__[0], "caution.png"))