mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
b777dba430
When a batch creates a session, we need to alert the client of this. Because the sessions are created by the batch manager (not directly in response to a client action), we need to emit an event with the session id. To accomodate this, a secondary set of sio sub/unsub/event handlers are created. These are specifically for batch events. The room is the `batch_id`. When creating a batch, the client subscribes to this batch room. When the batch manager creates a batch session, a `batch_session_created` event is emitted in the appropriate room. It includes the session id. The client then may subscribe to the session room, and all socket stuff proceeds as it did before.
54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi_events.handlers.local import local_handler
|
|
from fastapi_events.typing import Event
|
|
from fastapi_socketio import SocketManager
|
|
|
|
from ..services.events import EventServiceBase
|
|
|
|
|
|
class SocketIO:
|
|
__sio: SocketManager
|
|
|
|
def __init__(self, app: FastAPI):
|
|
self.__sio = SocketManager(app=app)
|
|
|
|
self.__sio.on("subscribe_session", handler=self._handle_sub_session)
|
|
self.__sio.on("unsubscribe_session", handler=self._handle_unsub_session)
|
|
local_handler.register(event_name=EventServiceBase.session_event, _func=self._handle_session_event)
|
|
|
|
self.__sio.on("subscribe_batch", handler=self._handle_sub_batch)
|
|
self.__sio.on("unsubscribe_batch", handler=self._handle_unsub_batch)
|
|
local_handler.register(event_name=EventServiceBase.batch_event, _func=self._handle_batch_event)
|
|
|
|
async def _handle_session_event(self, event: Event):
|
|
await self.__sio.emit(
|
|
event=event[1]["event"],
|
|
data=event[1]["data"],
|
|
room=event[1]["data"]["graph_execution_state_id"],
|
|
)
|
|
|
|
async def _handle_sub_session(self, sid, data, *args, **kwargs):
|
|
if "session" in data:
|
|
self.__sio.enter_room(sid, data["session"])
|
|
|
|
async def _handle_unsub_session(self, sid, data, *args, **kwargs):
|
|
if "session" in data:
|
|
self.__sio.leave_room(sid, data["session"])
|
|
|
|
async def _handle_batch_event(self, event: Event):
|
|
await self.__sio.emit(
|
|
event=event[1]["event"],
|
|
data=event[1]["data"],
|
|
room=event[1]["data"]["batch_id"],
|
|
)
|
|
|
|
async def _handle_sub_batch(self, sid, data, *args, **kwargs):
|
|
if "batch_id" in data:
|
|
self.__sio.enter_room(sid, data["batch_id"])
|
|
|
|
async def _handle_unsub_batch(self, sid, data, *args, **kwargs):
|
|
if "batch_id" in data:
|
|
self.__sio.enter_room(sid, data["batch_id"])
|