InvokeAI/invokeai/app/services/events.py

# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)

from typing import Any, Dict, TypedDict

ProgressImage = TypedDict(
    "ProgressImage", {"dataURL": str, "width": int, "height": int}
)

class EventServiceBase:
    session_event: str = "session_event"

    """Basic event bus, to have an empty stand-in when not needed"""

    def dispatch(self, event_name: str, payload: Any) -> None:
        pass

    def __emit_session_event(self, event_name: str, payload: Dict) -> None:
        self.dispatch(
            event_name=EventServiceBase.session_event,
            payload=dict(event=event_name, data=payload),
        )

    # Define events here for every event in the system.
    # This will make them easier to integrate until we find a schema generator.
    def emit_generator_progress(
        self,
        graph_execution_state_id: str,
        invocation_id: str,
        progress_image: ProgressImage | None,
        step: int,
        total_steps: int,
    ) -> None:
        """Emitted when there is generation progress"""
        self.__emit_session_event(
            event_name="generator_progress",
            payload=dict(
                graph_execution_state_id=graph_execution_state_id,
                invocation_id=invocation_id,
                progress_image=progress_image,
                step=step,
                total_steps=total_steps,
            ),
        )

    def emit_invocation_complete(
        self, graph_execution_state_id: str, invocation_id: str, result: Dict
    ) -> None:
        """Emitted when an invocation has completed"""
        self.__emit_session_event(
            event_name="invocation_complete",
            payload=dict(
                graph_execution_state_id=graph_execution_state_id,
                invocation_id=invocation_id,
                result=result,
            ),
        )

    def emit_invocation_error(
        self, graph_execution_state_id: str, invocation_id: str, error: str
    ) -> None:
        """Emitted when an invocation has completed"""
        self.__emit_session_event(
            event_name="invocation_error",
            payload=dict(
                graph_execution_state_id=graph_execution_state_id,
                invocation_id=invocation_id,
                error=error,
            ),
        )

    def emit_invocation_started(
        self, graph_execution_state_id: str, invocation_id: str
    ) -> None:
        """Emitted when an invocation has started"""
        self.__emit_session_event(
            event_name="invocation_started",
            payload=dict(
                graph_execution_state_id=graph_execution_state_id,
                invocation_id=invocation_id,
            ),
        )

    def emit_graph_execution_complete(self, graph_execution_state_id: str) -> None:
        """Emitted when a session has completed all invocations"""
        self.__emit_session_event(
            event_name="graph_execution_state_complete",
            payload=dict(graph_execution_state_id=graph_execution_state_id),
        )