66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
import socketio
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import folder_paths
|
|
import time
|
|
|
|
from ..utils.helpers import connect_print
|
|
from ..utils.gpu_utils import get_gpu_info, log_gpu_info
|
|
from ..config import config
|
|
|
|
class WebSocketController:
|
|
"""
|
|
WebSocket controller for handling real-time communication.
|
|
Routes WebSocket events to appropriate services, similar to HTTP controllers.
|
|
"""
|
|
|
|
def __init__(self, workflow_service):
|
|
self.sio = socketio.AsyncClient()
|
|
self.workflow_service = workflow_service
|
|
self.setup_event_handlers()
|
|
|
|
def setup_event_handlers(self):
|
|
"""Configure WebSocket event handlers (like HTTP routes)"""
|
|
|
|
@self.sio.event
|
|
async def connect():
|
|
connect_print("Client SocketIO connecté")
|
|
|
|
@self.sio.event
|
|
async def disconnect():
|
|
connect_print("Client SocketIO déconnecté")
|
|
|
|
@self.sio.on("run")
|
|
async def on_run(data):
|
|
connect_print(f"Événement 'run' reçu avec les données: {data}")
|
|
taskId = data.get("taskId")
|
|
name = data.get("name")
|
|
params = data.get("params")
|
|
result = await self.workflow_service.execute_workflow(name, params)
|
|
await self.sio.emit("return", {"taskId": taskId, "name": name, "result": result})
|
|
|
|
async def send_gpu_info(self):
|
|
"""Background task to send periodic GPU information"""
|
|
while True:
|
|
if self.sio.connected:
|
|
gpu_info = get_gpu_info()
|
|
# log_gpu_info(gpu_info)
|
|
await self.sio.emit("gpu_info", gpu_info)
|
|
await asyncio.sleep(config.GPU_INFO_INTERVAL)
|
|
|
|
async def start_socket_connection(self):
|
|
"""Initialize WebSocket connection to gateway server"""
|
|
socket_server_url = config.user_settings.get("Connect.Gateway")
|
|
if not socket_server_url:
|
|
connect_print("Connect.Gateway non configuré dans comfy.settings.json. Désactivation de SocketIO.")
|
|
return
|
|
|
|
await self.sio.connect(socket_server_url)
|
|
connect_print(f"Connecté au serveur SocketIO à {socket_server_url}")
|
|
await self.sio.wait()
|
|
|
|
async def initialize(self, app):
|
|
"""Initialize WebSocket tasks when application starts"""
|
|
asyncio.create_task(self.start_socket_connection())
|
|
asyncio.create_task(self.send_gpu_info()) |