Major refactoring with new config system

This commit is contained in:
IfnotFr
2025-06-28 10:18:17 +02:00
parent 92064dc9f2
commit afde99c7c2
15 changed files with 258 additions and 157 deletions
+13 -97
View File
@@ -1,10 +1,10 @@
import server
from aiohttp import web
from .workflow_manager import WorkflowManager
from .config import config
from .openapi_spec_generator import OpenAPISpecGenerator
from .utils import connect_print
from .websocket_manager import WebSocketManager
from .services.workflow_service import WorkflowService
from .controllers.websocket_controller import WebSocketController
from .controllers.workflow_controller import WorkflowController
from .controllers.app_controller import AppController
from .utils.helpers import connect_print
WEB_DIRECTORY = "./js"
NODE_CLASS_MAPPINGS = {}
@@ -13,99 +13,15 @@ version = "V0.0.1"
connect_print(f"Loading: ComfyUI Connect ({version})")
manager = WorkflowManager()
websocket_manager = WebSocketManager(manager)
# Initialize services
manager = WorkflowService()
# Initialize controllers
websocket_controller = WebSocketController(manager)
workflow_controller = WorkflowController(manager)
app_controller = AppController(manager)
async def init_socketio(app):
await websocket_manager.initialize(app)
await websocket_controller.initialize(app)
server.PromptServer.instance.app.on_startup.append(init_socketio)
# ####################
# ROUTES
# ####################
@server.PromptServer.instance.routes.get("/connect")
async def index(request):
return web.Response(text='''<!doctype html>
<html>
<head>
<meta charset="utf-8">
<script type="module" src="https://unpkg.com/rapidoc/dist/rapidoc-min.js"></script>
</head>
<body>
<rapi-doc
spec-url="/api/connect/openapi.json"
theme="dark"
show-info = 'false'
allow-authentication ='false'
allow-server-selection = 'false'
theme = 'dark'
>
</rapi-doc>
</body>
</html>''', content_type="text/html")
@server.PromptServer.instance.routes.get("/connect/openapi.json")
async def index(request):
workflows = []
names = await manager.list_workflows()
for name in names:
workflow = await manager.get_workflow(name)
workflows.append(workflow)
generator = OpenAPISpecGenerator(workflows)
return web.json_response(generator.generate())
@server.PromptServer.instance.routes.put("/connect/workflows")
async def save_workflow(request):
data = await request.json()
workflow = data["workflow"]
name = data["name"]
connect_print(f"PUT /connect/workflows - Saving workflow {name}")
await manager.save_workflow(name, workflow)
return web.json_response(
{"status": "success", "message": f"Workflow '{name}' saved."}
)
@server.PromptServer.instance.routes.delete("/connect/workflows/{name}")
async def delete_workflow(request):
name = request.match_info["name"]
connect_print(f"DELETE /connect/workflows/{name} - Deleting the workflow ...")
await manager.delete_workflow(name)
return web.json_response(
{"status": "success", "message": f"Workflow '{name}' deleted."}
)
@server.PromptServer.instance.routes.post("/connect/workflows/{name}")
async def execute_workflow(request):
params = await request.json()
name = request.match_info["name"]
connect_print(f"POST /connect/workflows/{name} - Running workflow ...")
result = await manager.execute_workflow(name, params)
return web.json_response({"status": "success", "workflow": name, "result": result})
@server.PromptServer.instance.routes.get("/connect/workflow/cache_nodes")
async def get_workflow(request):
cached_nodes = manager.get_workflows_cached_nodes()
return web.json_response({"status": "success", "nodes": cached_nodes})
@server.PromptServer.instance.routes.get("/connect/workflows/{name}")
async def get_workflow(request):
name = request.match_info["name"]
result = await manager.get_workflow(name)
return web.json_response(
{"status": "success", "workflow": name, "workflow": result}
)
+30 -9
View File
@@ -1,6 +1,6 @@
import os
import json
import folder_paths
from .comfyui_client import ComfyUIClient
class Config:
@@ -31,14 +31,35 @@ class Config:
API_TITLE: str = "Workflow API Documentation"
API_VERSION: str = "1.0.0"
CLIENT = None
async def client(self):
if not self.CLIENT:
self.CLIENT = ComfyUIClient(self.COMFY_ENDPOINT)
await self.CLIENT.connect()
return self.CLIENT
def __init__(self):
self._user_settings = None
self._settings_loaded = False
@property
def user_settings(self):
"""Lazy loading of user settings from comfy.settings.json"""
if not self._settings_loaded:
self._load_user_settings()
return self._user_settings
def _load_user_settings(self):
"""Load user settings from comfy.settings.json file"""
self._settings_loaded = True
settings_path = os.path.join(
folder_paths.get_user_directory(),
"default",
self.SETTINGS_FILENAME
)
try:
if not os.path.exists(settings_path):
self._user_settings = {}
return
with open(settings_path, "r", encoding="utf-8") as f:
self._user_settings = json.load(f)
except (json.JSONDecodeError, Exception):
self._user_settings = {}
config = Config()
+6
View File
@@ -0,0 +1,6 @@
# Controllers package for ComfyUI-Connect
# Handles HTTP routes and WebSocket events
from .workflow_controller import *
from .app_controller import *
from .websocket_controller import WebSocketController
+48
View File
@@ -0,0 +1,48 @@
import server
from aiohttp import web
from ..services.workflow_service import WorkflowService
from ..utils.openapi_utils import OpenAPISpecGenerator
class AppController:
def __init__(self, manager: WorkflowService):
self.manager = manager
self.setup_routes()
def setup_routes(self):
"""Setup API documentation and general routes"""
@server.PromptServer.instance.routes.get("/connect")
async def index(request):
return web.Response(text='''<!doctype html>
<html>
<head>
<meta charset="utf-8">
<script type="module" src="https://unpkg.com/rapidoc/dist/rapidoc-min.js"></script>
</head>
<body>
<rapi-doc
spec-url="/api/connect/openapi.json"
theme="dark"
show-info = 'false'
allow-authentication ='false'
allow-server-selection = 'false'
theme = 'dark'
>
</rapi-doc>
</body>
</html>''', content_type="text/html")
@server.PromptServer.instance.routes.get("/connect/openapi.json")
async def openapi_spec(request):
workflows = []
names = await self.manager.list_workflows()
for name in names:
workflow = await self.manager.get_workflow(name)
workflows.append(workflow)
generator = OpenAPISpecGenerator(workflows)
return web.json_response(generator.generate())
@@ -5,18 +5,23 @@ import os
import folder_paths
import time
from .utils import connect_print
from .gpu_info import get_gpu_info, log_gpu_info
from .config import config
from ..utils.helpers import connect_print
from ..utils.gpu_utils import get_gpu_info, log_gpu_info
from ..config import config
class WebSocketManager:
def __init__(self, workflow_manager):
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_manager = workflow_manager
self.workflow_service = workflow_service
self.setup_event_handlers()
def setup_event_handlers(self):
"""Configure les gestionnaires d'événements SocketIO"""
"""Configure WebSocket event handlers (like HTTP routes)"""
@self.sio.event
async def connect():
@@ -32,11 +37,11 @@ class WebSocketManager:
taskId = data.get("taskId")
name = data.get("name")
params = data.get("params")
result = await self.workflow_manager.execute_workflow(name, 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):
"""Tâche pour envoyer périodiquement les infos GPU"""
"""Background task to send periodic GPU information"""
while True:
if self.sio.connected:
gpu_info = get_gpu_info()
@@ -45,26 +50,8 @@ class WebSocketManager:
await asyncio.sleep(config.GPU_INFO_INTERVAL)
async def start_socket_connection(self):
"""Démarre la connexion SocketIO au serveur de passerelle"""
settings_path = os.path.join(os.path.dirname(folder_paths.__file__), "user", "default", config.SETTINGS_FILENAME)
try:
if not os.path.exists(settings_path):
connect_print(f"Fichier de paramètres non trouvé à: {settings_path}")
connect_print(f"Passerelle désactivée")
return
with open(settings_path, "r", encoding="utf-8") as f:
settings = json.load(f)
except json.JSONDecodeError as e:
connect_print(f"Erreur lors de l'analyse du fichier de paramètres {settings_path}: {e}")
connect_print(f"Passerelle désactivée")
return
except Exception as e:
connect_print(f"Erreur lors du chargement de {settings_path}: {e}")
connect_print(f"Passerelle désactivée")
return
socket_server_url = settings.get("Connect.Gateway")
"""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
@@ -74,6 +61,6 @@ class WebSocketManager:
await self.sio.wait()
async def initialize(self, app):
"""Initialise les tâches SocketIO lorsque l'application démarre"""
"""Initialize WebSocket tasks when application starts"""
asyncio.create_task(self.start_socket_connection())
asyncio.create_task(self.send_gpu_info())
+57
View File
@@ -0,0 +1,57 @@
import server
from aiohttp import web
from ..services.workflow_service import WorkflowService
from ..utils.helpers import connect_print
class WorkflowController:
def __init__(self, service: WorkflowService):
self.service = service
self.setup_routes()
def setup_routes(self):
"""Setup workflow-related routes"""
@server.PromptServer.instance.routes.put("/connect/workflows")
async def save_workflow(request):
data = await request.json()
workflow = data["workflow"]
name = data["name"]
connect_print(f"PUT /connect/workflows - Saving workflow {name}")
await self.service.save_workflow(name, workflow)
return web.json_response(
{"status": "success", "message": f"Workflow '{name}' saved."}
)
@server.PromptServer.instance.routes.delete("/connect/workflows/{name}")
async def delete_workflow(request):
name = request.match_info["name"]
connect_print(f"DELETE /connect/workflows/{name} - Deleting the workflow ...")
await self.service.delete_workflow(name)
return web.json_response(
{"status": "success", "message": f"Workflow '{name}' deleted."}
)
@server.PromptServer.instance.routes.post("/connect/workflows/{name}")
async def execute_workflow(request):
params = await request.json()
name = request.match_info["name"]
connect_print(f"POST /connect/workflows/{name} - Running workflow ...")
result = await self.service.execute_workflow(name, params)
return web.json_response({"status": "success", "workflow": name, "result": result})
@server.PromptServer.instance.routes.get("/connect/workflow/cache_nodes")
async def get_cached_nodes(request):
cached_nodes = self.service.get_workflows_cached_nodes()
return web.json_response({"status": "success", "nodes": cached_nodes})
@server.PromptServer.instance.routes.get("/connect/workflows/{name}")
async def get_workflow(request):
name = request.match_info["name"]
result = await self.service.get_workflow(name)
return web.json_response(
{"status": "success", "workflow": name, "workflow": result}
)
+4
View File
@@ -0,0 +1,4 @@
# Entities package for ComfyUI-Connect
# Contains domain models and business entities
from .workflow import Workflow
+1 -1
View File
@@ -8,7 +8,7 @@ def lowerSingular(string):
return string
class WorkflowWrapper(dict):
class Workflow(dict):
"""
A workflow wrapper that extends a dictionary of nodes to provide
additional functionalities such as tagging, filtering, and bypassing nodes.
+5
View File
@@ -0,0 +1,5 @@
# Services package for ComfyUI-Connect
# Contains business logic and service layer
from .workflow_service import WorkflowService
from .comfyui_service import ComfyUIService, comfyui_service
@@ -6,28 +6,53 @@ import aiohttp
import base64
import asyncio
from typing import Dict, List
from ..config import config
class ComfyUIClient:
def __init__(self, COMFY_ENDPOINT):
class ComfyUIService:
"""
Service for managing ComfyUI connections and workflow execution.
Combines client functionality with service-level management.
Implements singleton pattern for connection management.
"""
_instance = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
if self._initialized:
return
self._initialized = True
self.CLIENT_ID = str(uuid.uuid4())
self.COMFY_ENDPOINT = COMFY_ENDPOINT
self.COMFY_ENDPOINT = config.COMFY_ENDPOINT
self.ws = None
self.session = None
self._message_queue = asyncio.Queue()
self._prompt_events: Dict[str, asyncio.Event] = {}
self._listener_task = None
self._connected = False
async def connect(self):
"""Establish connection to ComfyUI"""
if self._connected:
return
self.session = aiohttp.ClientSession()
self.ws = await self.session.ws_connect(
f"ws://{self.COMFY_ENDPOINT}/ws?clientId={self.CLIENT_ID}"
)
# Start the global websocket listener
self._listener_task = asyncio.create_task(self._listen_websocket())
self._connected = True
async def _listen_websocket(self):
"""Listen for WebSocket messages from ComfyUI"""
try:
while True:
message = await self.ws.receive()
@@ -50,14 +75,23 @@ class ComfyUIClient:
self._listener_task = asyncio.create_task(self._listen_websocket())
async def close(self):
"""Close the ComfyUI connection"""
if self._listener_task:
self._listener_task.cancel()
if self.ws:
await self.ws.close()
if self.session:
await self.session.close()
self._connected = False
async def _ensure_connected(self):
"""Ensure we have an active connection to ComfyUI"""
if not self._connected:
await self.connect()
async def queue_prompt(self, prompt):
"""Queue a prompt for execution in ComfyUI"""
await self._ensure_connected()
payload = {"prompt": prompt, "client_id": self.CLIENT_ID}
data = json.dumps(payload).encode("utf-8")
async with self.session.post(
@@ -66,6 +100,8 @@ class ComfyUIClient:
return await response.json()
async def get_image(self, filename, subfolder, folder_type):
"""Retrieve an image from ComfyUI"""
await self._ensure_connected()
params = {"filename": filename, "subfolder": subfolder, "type": folder_type}
url_values = urllib.parse.urlencode(params)
async with self.session.get(
@@ -76,14 +112,24 @@ class ComfyUIClient:
return image_base64
async def get_history(self, prompt_id):
"""Get execution history for a prompt"""
await self._ensure_connected()
async with self.session.get(
f"http://{self.COMFY_ENDPOINT}/history/{prompt_id}"
) as response:
return await response.json()
async def run(self, prompt):
async def run_workflow(self, workflow: dict) -> dict:
"""
Execute a workflow and return the generated images.
:param workflow: The workflow to execute
:return: Dictionary of generated images by node ID
"""
await self._ensure_connected()
# Create an event for this prompt
prompt_id = (await self.queue_prompt(prompt))["prompt_id"]
prompt_id = (await self.queue_prompt(workflow))["prompt_id"]
self._prompt_events[prompt_id] = asyncio.Event()
try:
@@ -106,3 +152,7 @@ class ComfyUIClient:
finally:
# Clean up the event
del self._prompt_events[prompt_id]
# Global service instance
comfyui_service = ComfyUIService()
@@ -4,11 +4,12 @@ import copy
import aiofiles
import base64
import requests
from .workflow_wrapper import WorkflowWrapper
from .config import config
from ..entities.workflow import Workflow
from ..config import config
from .comfyui_service import comfyui_service
class WorkflowManager:
class WorkflowService:
"""
Manages workflows by loading them from JSON files, saving, deleting, and executing them.
Also handles updating workflows and caching certain nodes.
@@ -16,7 +17,7 @@ class WorkflowManager:
def __init__(self):
"""
Initializes the WorkflowManager by:
Initializes the WorkflowService by:
- Creating necessary directories if they don't exist.
- Loading JSON workflow files from disk into memory.
- Refreshing the cached nodes for all loaded workflows.
@@ -57,7 +58,7 @@ class WorkflowManager:
# Go through each workflow and check for cached nodes
for workflow_name, workflow_data in self.workflows.items():
wrapper = WorkflowWrapper(workflow_data)
wrapper = Workflow(workflow_data)
cached_nodes = wrapper.get_tagged_nodes("!cache")
# Store each cached node with the workflow name for reference
@@ -138,8 +139,8 @@ class WorkflowManager:
if name not in self.workflows:
raise FileNotFoundError(f"Workflow '{name}' not found.")
# Wrap the workflow in a WorkflowWrapper object for convenience
workflow = WorkflowWrapper(copy.deepcopy(self.workflows[name]))
# Wrap the workflow in a Workflow object for convenience
workflow = Workflow(copy.deepcopy(self.workflows[name]))
# Bypass any nodes tagged with "!bypass" if present
workflow.bypass_nodes("!bypass")
@@ -237,8 +238,8 @@ class WorkflowManager:
# Update the workflow with the value
workflow.update_tagged_nodes_input(tag, input_name, value)
# Run the workflow asynchronously using the configured client
images = await (await config.client()).run(workflow)
# Run the workflow asynchronously using the ComfyUI service
images = await comfyui_service.run_workflow(workflow)
response = {}
# Collect and group the resulting images by each node's tags
@@ -268,7 +269,7 @@ class WorkflowManager:
:param name: The name of the workflow to retrieve information from.
:return: A dictionary containing the workflow's name, its tagged inputs, and outputs.
"""
wrapper = WorkflowWrapper(self.workflows[name])
wrapper = Workflow(self.workflows[name])
return {
"name": name,
"inputs": wrapper.get_tagged_inputs(),
+6
View File
@@ -0,0 +1,6 @@
# Utils package for ComfyUI-Connect
# Contains utility functions and helpers
from .helpers import connect_print
from .gpu_utils import get_gpu_info, log_gpu_info
from .openapi_utils import OpenAPISpecGenerator
+2 -2
View File
@@ -1,7 +1,7 @@
import pynvml
import time
from .utils import connect_print
from .config import config
from .helpers import connect_print
from ..config import config
def get_gpu_info():
"""Retrieves detailed information about installed NVIDIA GPUs"""
View File
@@ -1,4 +1,4 @@
from .config import config
from ..config import config
class OpenAPISpecGenerator:
def __init__(self, workflows: list):