mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
|
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
||
|
|
|
||
|
|
from fastapi import FastAPI
|
||
|
|
from fastapi_socketio import SocketManager
|
||
|
|
from fastapi_events.handlers.local import local_handler
|
||
|
|
from fastapi_events.typing import Event
|
||
|
|
from ..services.events import EventServiceBase
|
||
|
|
|
||
|
|
class SocketIO:
|
||
|
|
__sio: SocketManager
|
||
|
|
|
||
|
|
def __init__(self, app: FastAPI):
|
||
|
|
self.__sio = SocketManager(app = app)
|
||
|
|
self.__sio.on('subscribe', handler=self._handle_sub)
|
||
|
|
self.__sio.on('unsubscribe', handler=self._handle_unsub)
|
||
|
|
|
||
|
|
local_handler.register(
|
||
|
|
event_name = EventServiceBase.session_event,
|
||
|
|
_func=self._handle_session_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(self, sid, data, *args, **kwargs):
|
||
|
|
if 'session' in data:
|
||
|
|
self.__sio.enter_room(sid, data['session'])
|
||
|
|
|
||
|
|
# @app.sio.on('unsubscribe')
|
||
|
|
async def _handle_unsub(self, sid, data, *args, **kwargs):
|
||
|
|
if 'session' in data:
|
||
|
|
self.__sio.leave_room(sid, data['session'])
|