refactor(events): include full model source in model install events

This is required to fix an issue with the MM UI's error handling.

Previously, we only included the model source as a string. That could be an arbitrary URL, file path or HF repo id, but the frontend has no parsing logic to differentiate between these different model sources.

Without access to the type of model source, it is difficult to determine how the user should proceed. For example, if it's HF URL with an HTTP unauthorized error, we should direct the user to log in to HF. But if it's a civitai URL with the same error, we should not direct the user to HF.

There are a variety of related edge cases.

With this change, the full `ModelSource` object is included in each model install event, including error events.

I had to fix some circular import issues, hence the import changes to files other than `events_common.py`.

Author: psychedelicious

invokeai/app/services/download/download_default.py

@@ -8,7 +8,7 @@
 import traceback
 from pathlib import Path
 from queue import Empty, PriorityQueue
-from typing import Any, Dict, List, Literal, Optional, Set
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set
 
 import requests
 from pydantic.networks import AnyHttpUrl
@@ -28,11 +28,13 @@
     ServiceInactiveException,
     UnknownJobIDException,
 )
-from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.util.misc import get_iso_timestamp
 from invokeai.backend.model_manager.metadata import RemoteModelFile
 from invokeai.backend.util.logging import InvokeAILogger
 
+if TYPE_CHECKING:
+    from invokeai.app.services.events.events_base import EventServiceBase
+
 # Maximum number of bytes to download during each call to requests.iter_content()
 DOWNLOAD_CHUNK_SIZE = 100000
 

invokeai/app/services/events/__init__.py

@@ -1 +0,0 @@
-from .events_base import EventServiceBase  # noqa F401

invokeai/app/services/events/events_common.py

@@ -4,6 +4,7 @@
 from fastapi_events.registry.payload_schema import registry as payload_schema
 from pydantic import BaseModel, ConfigDict, Field
 
+from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
 from invokeai.app.services.session_processor.session_processor_common import ProgressImage
 from invokeai.app.services.session_queue.session_queue_common import (
     QUEUE_ITEM_STATUS,
@@ -18,7 +19,7 @@
 
 if TYPE_CHECKING:
     from invokeai.app.services.download.download_base import DownloadJob
-    from invokeai.app.services.model_install.model_install_common import ModelInstallJob
+    from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
 
 
 class EventBase(BaseModel):
@@ -422,7 +423,7 @@ class ModelInstallDownloadStartedEvent(ModelEventBase):
     __event_name__ = "model_install_download_started"
 
     id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
+    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
     local_path: str = Field(description="Where model is downloading to")
     bytes: int = Field(description="Number of bytes downloaded so far")
     total_bytes: int = Field(description="Total size of download, including all files")
@@ -443,7 +444,7 @@ def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadStartedEvent":
         ]
         return cls(
             id=job.id,
-            source=str(job.source),
+            source=job.source,
             local_path=job.local_path.as_posix(),
             parts=parts,
             bytes=job.bytes,
@@ -458,7 +459,7 @@ class ModelInstallDownloadProgressEvent(ModelEventBase):
     __event_name__ = "model_install_download_progress"
 
     id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
+    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
     local_path: str = Field(description="Where model is downloading to")
     bytes: int = Field(description="Number of bytes downloaded so far")
     total_bytes: int = Field(description="Total size of download, including all files")
@@ -479,7 +480,7 @@ def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadProgressEvent":
         ]
         return cls(
             id=job.id,
-            source=str(job.source),
+            source=job.source,
             local_path=job.local_path.as_posix(),
             parts=parts,
             bytes=job.bytes,
@@ -494,11 +495,11 @@ class ModelInstallDownloadsCompleteEvent(ModelEventBase):
     __event_name__ = "model_install_downloads_complete"
 
     id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
+    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadsCompleteEvent":
-        return cls(id=job.id, source=str(job.source))
+        return cls(id=job.id, source=job.source)
 
 
 @payload_schema.register
@@ -508,11 +509,11 @@ class ModelInstallStartedEvent(ModelEventBase):
     __event_name__ = "model_install_started"
 
     id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
+    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallStartedEvent":
-        return cls(id=job.id, source=str(job.source))
+        return cls(id=job.id, source=job.source)
 
 
 @payload_schema.register
@@ -522,14 +523,14 @@ class ModelInstallCompleteEvent(ModelEventBase):
     __event_name__ = "model_install_complete"
 
     id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
+    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
     key: str = Field(description="Model config record key")
     total_bytes: Optional[int] = Field(description="Size of the model (may be None for installation of a local path)")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallCompleteEvent":
         assert job.config_out is not None
-        return cls(id=job.id, source=str(job.source), key=(job.config_out.key), total_bytes=job.total_bytes)
+        return cls(id=job.id, source=job.source, key=(job.config_out.key), total_bytes=job.total_bytes)
 
 
 @payload_schema.register
@@ -539,11 +540,11 @@ class ModelInstallCancelledEvent(ModelEventBase):
     __event_name__ = "model_install_cancelled"
 
     id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
+    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallCancelledEvent":
-        return cls(id=job.id, source=str(job.source))
+        return cls(id=job.id, source=job.source)
 
 
 @payload_schema.register
@@ -553,15 +554,15 @@ class ModelInstallErrorEvent(ModelEventBase):
     __event_name__ = "model_install_error"
 
     id: int = Field(description="The ID of the install job")
-    source: str = Field(description="Source of the model; local path, repo_id or url")
+    source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
     error_type: str = Field(description="The name of the exception")
     error: str = Field(description="A text description of the exception")
 
     @classmethod
     def build(cls, job: "ModelInstallJob") -> "ModelInstallErrorEvent":
         assert job.error_type is not None
         assert job.error is not None
-        return cls(id=job.id, source=str(job.source), error_type=job.error_type, error=job.error)
+        return cls(id=job.id, source=job.source, error_type=job.error_type, error=job.error)
 
 
 class BulkDownloadEventBase(EventBase):

invokeai/app/services/model_install/model_install_base.py

@@ -3,18 +3,20 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 
 from pydantic.networks import AnyHttpUrl
 
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.download import DownloadQueueServiceBase
-from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
 from invokeai.app.services.model_records import ModelRecordChanges, ModelRecordServiceBase
 from invokeai.backend.model_manager import AnyModelConfig
 
+if TYPE_CHECKING:
+    from invokeai.app.services.events.events_base import EventServiceBase
+
 
 class ModelInstallServiceBase(ABC):
     """Abstract base class for InvokeAI model installation."""

invokeai/app/services/model_install/model_install_default.py

@@ -9,7 +9,7 @@
 from queue import Empty, Queue
 from shutil import copyfile, copytree, move, rmtree
 from tempfile import mkdtemp
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 
 import torch
 import yaml
@@ -20,7 +20,6 @@
 
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.download import DownloadQueueServiceBase, MultiFileDownloadJob
-from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_install.model_install_base import ModelInstallServiceBase
 from invokeai.app.services.model_install.model_install_common import (
@@ -57,6 +56,10 @@
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.util import slugify
 
+if TYPE_CHECKING:
+    from invokeai.app.services.events.events_base import EventServiceBase
+
+
 TMPDIR_PREFIX = "tmpinstall_"