refactor(events): use pydantic schemas for events

Our events handling and implementation is not very friendly:
- We have to define functions for every type of event we want to emit and call those specific functions.
- Adding or removing data from payloads requires changes wherever the events are dispatched from, and in the events service.
- We have no type safety for events and need to rely on string matching and dict access when interacting with events.

`fastapi_events` has a neat feature where you can create a pydantic model as an event payload, give it an `__event_name__` attr, and then dispatch the model directly.

This allows us to eliminate a layer of indirection and some unpleasant complexity:
- We do not need functions for every event type. Define the event in a single model, and dispatch the model directly. The events service only has a single `dispatch` method.
- Event handler callbacks get type hints for their event payloads, and can use `isinstance` on them if needed. *see note below*
- Event payload construction is now the responsibility of the event itself, not the service. Every event model has a `build` class method, encapsulating this logic.
- We can generate OpenAPI schemas for the event payloads and get type safety on the frontend. Previously, the types were manually created (and bugs _have_ occurred when a payload changed on the backend).
- When registering event callbacks, we can now register the event model itself instead of its event name. Impossible to make typos.

This commit moves the backend over to this improved event handling setup.

*Note*
Actually, `fastapi_events` has baked in conversion of pydantic events to dicts, so event callbacks get an untyped dict instead of a pydantic model. I've raised a PR to make this behaviour configurable: melvinkcx/fastapi-events#57

I used my `fastapi_events` PR branch while working on this PR, so you'd need to use it to test: https://github.com/psychedelicious/fastapi-events/tree/psyche/feat/payload-schema-dump

chore(ui): typegen

refactor(ui): update frontend to use new events setup

TODO: Support session_started and session_canceled events.

feat(events): improved types

build: pin fastapi-events to unreleased version

It's merged, but not released yet.

fix(events): fix merge issue

chore(ui): lint

feat(events): migrate MM install events to pydantic

feat(ui): update events for new mm install

chore: bump fastapi_events

feat(events): revise how events are dispatched

Dispatching the fully-formed event payloads directly ended up causing circular import issues. Revised so that each event has its own `emit_...` method on the events service, which is how it was originally. This also makes it clear what events are valid in the system.

invokeai/app/api/dependencies.py

@@ -18,6 +18,7 @@
 from ..services.bulk_download.bulk_download_default import BulkDownloadService
 from ..services.config import InvokeAIAppConfig
 from ..services.download import DownloadQueueService
+from ..services.events.events_fastapievents import FastAPIEventService
 from ..services.image_files.image_files_disk import DiskImageFileStorage
 from ..services.image_records.image_records_sqlite import SqliteImageRecordStorage
 from ..services.images.images_default import ImageService
@@ -33,7 +34,6 @@
 from ..services.session_queue.session_queue_sqlite import SqliteSessionQueue
 from ..services.urls.urls_default import LocalUrlService
 from ..services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage
-from .events import FastAPIEventService
 
 
 # TODO: is there a better way to achieve this?

invokeai/app/api/routers/model_manager.py

@@ -15,7 +15,7 @@
 from starlette.exceptions import HTTPException
 from typing_extensions import Annotated
 
-from invokeai.app.services.model_install import ModelInstallJob
+from invokeai.app.services.model_install.model_install_common import ModelInstallJob
 from invokeai.app.services.model_records import (
     InvalidModelException,
     UnknownModelException,

invokeai/app/api/sockets.py

@@ -1,66 +1,119 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
+from typing import Any
+
 from fastapi import FastAPI
-from fastapi_events.handlers.local import local_handler
-from fastapi_events.typing import Event
+from pydantic import BaseModel
 from socketio import ASGIApp, AsyncServer
 
-from ..services.events.events_base import EventServiceBase
+from invokeai.app.services.events.events_common import (
+    BatchEnqueuedEvent,
+    BulkDownloadCompleteEvent,
+    BulkDownloadErrorEvent,
+    BulkDownloadEvent,
+    BulkDownloadStartedEvent,
+    FastAPIEvent,
+    InvocationCompleteEvent,
+    InvocationDenoiseProgressEvent,
+    InvocationErrorEvent,
+    InvocationStartedEvent,
+    ModelEvent,
+    ModelInstallCancelledEvent,
+    ModelInstallCompleteEvent,
+    ModelInstallDownloadProgressEvent,
+    ModelInstallErrorEvent,
+    ModelInstallStartedEvent,
+    ModelLoadCompleteEvent,
+    ModelLoadStartedEvent,
+    QueueClearedEvent,
+    QueueEvent,
+    QueueItemStatusChangedEvent,
+    SessionCanceledEvent,
+    SessionCompleteEvent,
+    SessionStartedEvent,
+    register_events,
+)
+
+
+class QueueSubscriptionEvent(BaseModel):
+    queue_id: str
+
+
+class BulkDownloadSubscriptionEvent(BaseModel):
+    bulk_download_id: str
 
 
 class SocketIO:
-    __sio: AsyncServer
-    __app: ASGIApp
-
-    __sub_queue: str = "subscribe_queue"
-    __unsub_queue: str = "unsubscribe_queue"
+    _sub_queue = "subscribe_queue"
+    _unsub_queue = "unsubscribe_queue"
 
-    __sub_bulk_download: str = "subscribe_bulk_download"
-    __unsub_bulk_download: str = "unsubscribe_bulk_download"
+    _sub_bulk_download = "subscribe_bulk_download"
+    _unsub_bulk_download = "unsubscribe_bulk_download"
 
     def __init__(self, app: FastAPI):
-        self.__sio = AsyncServer(async_mode="asgi", cors_allowed_origins="*")
-        self.__app = ASGIApp(socketio_server=self.__sio, socketio_path="/ws/socket.io")
-        app.mount("/ws", self.__app)
-
-        self.__sio.on(self.__sub_queue, handler=self._handle_sub_queue)
-        self.__sio.on(self.__unsub_queue, handler=self._handle_unsub_queue)
-        local_handler.register(event_name=EventServiceBase.queue_event, _func=self._handle_queue_event)
-        local_handler.register(event_name=EventServiceBase.model_event, _func=self._handle_model_event)
-
-        self.__sio.on(self.__sub_bulk_download, handler=self._handle_sub_bulk_download)
-        self.__sio.on(self.__unsub_bulk_download, handler=self._handle_unsub_bulk_download)
-        local_handler.register(event_name=EventServiceBase.bulk_download_event, _func=self._handle_bulk_download_event)
-
-    async def _handle_queue_event(self, event: Event):
-        await self.__sio.emit(
-            event=event[1]["event"],
-            data=event[1]["data"],
-            room=event[1]["data"]["queue_id"],
+        self._sio = AsyncServer(async_mode="asgi", cors_allowed_origins="*")
+        self._app = ASGIApp(socketio_server=self._sio, socketio_path="/ws/socket.io")
+        app.mount("/ws", self._app)
+
+        self._sio.on(self._sub_queue, handler=self._handle_sub_queue)
+        self._sio.on(self._unsub_queue, handler=self._handle_unsub_queue)
+        self._sio.on(self._sub_bulk_download, handler=self._handle_sub_bulk_download)
+        self._sio.on(self._unsub_bulk_download, handler=self._handle_unsub_bulk_download)
+
+        register_events(
+            {
+                InvocationStartedEvent,
+                InvocationDenoiseProgressEvent,
+                InvocationCompleteEvent,
+                InvocationErrorEvent,
+                SessionStartedEvent,
+                SessionCompleteEvent,
+                SessionCanceledEvent,
+                QueueItemStatusChangedEvent,
+                BatchEnqueuedEvent,
+                QueueClearedEvent,
+            },
+            self._handle_queue_event,
+        )
+
+        register_events(
+            {
+                ModelLoadStartedEvent,
+                ModelLoadCompleteEvent,
+                ModelInstallDownloadProgressEvent,
+                ModelInstallStartedEvent,
+                ModelInstallCompleteEvent,
+                ModelInstallCancelledEvent,
+                ModelInstallErrorEvent,
+            },
+            self._handle_model_event,
         )
 
-    async def _handle_sub_queue(self, sid, data, *args, **kwargs) -> None:
-        if "queue_id" in data:
-            await self.__sio.enter_room(sid, data["queue_id"])
+        register_events(
+            {BulkDownloadStartedEvent, BulkDownloadCompleteEvent, BulkDownloadErrorEvent},
+            self._handle_bulk_image_download_event,
+        )
 
-    async def _handle_unsub_queue(self, sid, data, *args, **kwargs) -> None:
-        if "queue_id" in data:
-            await self.__sio.leave_room(sid, data["queue_id"])
+    async def _handle_sub_queue(self, sid: str, data: Any) -> None:
+        await self._sio.enter_room(sid, QueueSubscriptionEvent(**data).queue_id)
 
-    async def _handle_model_event(self, event: Event) -> None:
-        await self.__sio.emit(event=event[1]["event"], data=event[1]["data"])
+    async def _handle_unsub_queue(self, sid: str, data: Any) -> None:
+        await self._sio.leave_room(sid, QueueSubscriptionEvent(**data).queue_id)
 
-    async def _handle_bulk_download_event(self, event: Event):
-        await self.__sio.emit(
-            event=event[1]["event"],
-            data=event[1]["data"],
-            room=event[1]["data"]["bulk_download_id"],
-        )
+    async def _handle_sub_bulk_download(self, sid: str, data: Any) -> None:
+        await self._sio.enter_room(sid, BulkDownloadSubscriptionEvent(**data).bulk_download_id)
+
+    async def _handle_unsub_bulk_download(self, sid: str, data: Any) -> None:
+        await self._sio.leave_room(sid, BulkDownloadSubscriptionEvent(**data).bulk_download_id)
+
+    async def _handle_queue_event(self, event: FastAPIEvent[QueueEvent]):
+        event_name, payload = event
+        await self._sio.emit(event=event_name, data=payload.model_dump(), room=payload.queue_id)
 
-    async def _handle_sub_bulk_download(self, sid, data, *args, **kwargs):
-        if "bulk_download_id" in data:
-            await self.__sio.enter_room(sid, data["bulk_download_id"])
+    async def _handle_model_event(self, event: FastAPIEvent[ModelEvent]) -> None:
+        event_name, payload = event
+        await self._sio.emit(event=event_name, data=payload.model_dump())
 
-    async def _handle_unsub_bulk_download(self, sid, data, *args, **kwargs):
-        if "bulk_download_id" in data:
-            await self.__sio.leave_room(sid, data["bulk_download_id"])
+    async def _handle_bulk_image_download_event(self, event: FastAPIEvent[BulkDownloadEvent]) -> None:
+        event_name, payload = event
+        await self._sio.emit(event=event_name, data=payload.model_dump())

invokeai/app/api_app.py

@@ -2,6 +2,9 @@
 # which are imported/used before parse_args() is called will get the default config values instead of the
 # values from the command line or config file.
 import sys
+from typing import cast
+
+from pydantic import BaseModel
 
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.services.session_processor.session_processor_common import ProgressImage
@@ -33,6 +36,7 @@
     from fastapi.responses import HTMLResponse
     from fastapi_events.handlers.local import local_handler
     from fastapi_events.middleware import EventHandlerASGIMiddleware
+    from fastapi_events.registry.payload_schema import registry as fastapi_events_registry
     from pydantic.json_schema import models_json_schema
     from torch.backends.mps import is_available as is_mps_available
 
@@ -182,23 +186,16 @@ def custom_openapi() -> dict[str, Any]:
         invoker_schema["output"] = outputs_ref
         invoker_schema["class"] = "invocation"
 
-    # This code no longer seems to be necessary?
-    # Leave it here just in case
-    #
-    # from invokeai.backend.model_manager import get_model_config_formats
-    # formats = get_model_config_formats()
-    # for model_config_name, enum_set in formats.items():
-
-    #     if model_config_name in openapi_schema["components"]["schemas"]:
-    #         # print(f"Config with name {name} already defined")
-    #         continue
-
-    #     openapi_schema["components"]["schemas"][model_config_name] = {
-    #         "title": model_config_name,
-    #         "description": "An enumeration.",
-    #         "type": "string",
-    #         "enum": [v.value for v in enum_set],
-    #     }
+    # Add all pydantic event schemas registered with fastapi-events
+    for payload in fastapi_events_registry.data.values():
+        json_schema = cast(BaseModel, payload).model_json_schema(
+            mode="serialization", ref_template="#/components/schemas/{model}"
+        )
+        if "$defs" in json_schema:
+            for schema_key, schema in json_schema["$defs"].items():
+                openapi_schema["components"]["schemas"][schema_key] = schema
+            del json_schema["$defs"]
+        openapi_schema["components"]["schemas"][payload.__name__] = json_schema
 
     app.openapi_schema = openapi_schema
     return app.openapi_schema

invokeai/app/services/bulk_download/bulk_download_default.py

@@ -106,9 +106,7 @@ def _signal_job_started(
         if self._invoker:
             assert bulk_download_id is not None
             self._invoker.services.events.emit_bulk_download_started(
-                bulk_download_id=bulk_download_id,
-                bulk_download_item_id=bulk_download_item_id,
-                bulk_download_item_name=bulk_download_item_name,
+                bulk_download_id, bulk_download_item_id, bulk_download_item_name
             )
 
     def _signal_job_completed(
@@ -118,10 +116,8 @@ def _signal_job_completed(
         if self._invoker:
             assert bulk_download_id is not None
             assert bulk_download_item_name is not None
-            self._invoker.services.events.emit_bulk_download_completed(
-                bulk_download_id=bulk_download_id,
-                bulk_download_item_id=bulk_download_item_id,
-                bulk_download_item_name=bulk_download_item_name,
+            self._invoker.services.events.emit_bulk_download_complete(
+                bulk_download_id, bulk_download_item_id, bulk_download_item_name
             )
 
     def _signal_job_failed(
@@ -131,11 +127,8 @@ def _signal_job_failed(
         if self._invoker:
             assert bulk_download_id is not None
             assert exception is not None
-            self._invoker.services.events.emit_bulk_download_failed(
-                bulk_download_id=bulk_download_id,
-                bulk_download_item_id=bulk_download_item_id,
-                bulk_download_item_name=bulk_download_item_name,
-                error=str(exception),
+            self._invoker.services.events.emit_bulk_download_error(
+                bulk_download_id, bulk_download_item_id, bulk_download_item_name, str(exception)
             )
 
     def stop(self, *args, **kwargs):

invokeai/app/services/download/download_default.py

@@ -8,14 +8,13 @@
 import traceback
 from pathlib import Path
 from queue import Empty, PriorityQueue
-from typing import Any, Dict, List, Optional, Set
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set
 
 import requests
 from pydantic.networks import AnyHttpUrl
 from requests import HTTPError
 from tqdm import tqdm
 
-from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.util.misc import get_iso_timestamp
 from invokeai.backend.util.logging import InvokeAILogger
 
@@ -30,6 +29,9 @@
     UnknownJobIDException,
 )
 
+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
 
@@ -40,7 +42,7 @@ class DownloadQueueService(DownloadQueueServiceBase):
     def __init__(
         self,
         max_parallel_dl: int = 5,
-        event_bus: Optional[EventServiceBase] = None,
+        event_bus: Optional["EventServiceBase"] = None,
         requests_session: Optional[requests.sessions.Session] = None,
     ):
         """