diff --git a/__init__.py b/__init__.py
index e2b1de0..230ed6e 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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='''
-
-
-
-
-
-
-
-
-
-
-
-
-''', 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}
- )
diff --git a/config.py b/config.py
index c765819..37ddca2 100644
--- a/config.py
+++ b/config.py
@@ -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()
diff --git a/controllers/__init__.py b/controllers/__init__.py
new file mode 100644
index 0000000..b00ec38
--- /dev/null
+++ b/controllers/__init__.py
@@ -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
\ No newline at end of file
diff --git a/controllers/app_controller.py b/controllers/app_controller.py
new file mode 100644
index 0000000..4aefe8f
--- /dev/null
+++ b/controllers/app_controller.py
@@ -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='''
+
+
+
+
+
+
+
+
+
+
+
+
+''', 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())
\ No newline at end of file
diff --git a/websocket_manager.py b/controllers/websocket_controller.py
similarity index 52%
rename from websocket_manager.py
rename to controllers/websocket_controller.py
index 909b486..1b795c5 100644
--- a/websocket_manager.py
+++ b/controllers/websocket_controller.py
@@ -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())
\ No newline at end of file
diff --git a/controllers/workflow_controller.py b/controllers/workflow_controller.py
new file mode 100644
index 0000000..51de120
--- /dev/null
+++ b/controllers/workflow_controller.py
@@ -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}
+ )
\ No newline at end of file
diff --git a/entities/__init__.py b/entities/__init__.py
new file mode 100644
index 0000000..a38750c
--- /dev/null
+++ b/entities/__init__.py
@@ -0,0 +1,4 @@
+# Entities package for ComfyUI-Connect
+# Contains domain models and business entities
+
+from .workflow import Workflow
\ No newline at end of file
diff --git a/workflow_wrapper.py b/entities/workflow.py
similarity index 99%
rename from workflow_wrapper.py
rename to entities/workflow.py
index 2f86b51..d46d5cc 100644
--- a/workflow_wrapper.py
+++ b/entities/workflow.py
@@ -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.
diff --git a/services/__init__.py b/services/__init__.py
new file mode 100644
index 0000000..f667122
--- /dev/null
+++ b/services/__init__.py
@@ -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
\ No newline at end of file
diff --git a/comfyui_client.py b/services/comfyui_service.py
similarity index 68%
rename from comfyui_client.py
rename to services/comfyui_service.py
index d341cbc..c2d2f1c 100644
--- a/comfyui_client.py
+++ b/services/comfyui_service.py
@@ -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()
\ No newline at end of file
diff --git a/workflow_manager.py b/services/workflow_service.py
similarity index 95%
rename from workflow_manager.py
rename to services/workflow_service.py
index f2306fd..4edb962 100644
--- a/workflow_manager.py
+++ b/services/workflow_service.py
@@ -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(),
diff --git a/utils/__init__.py b/utils/__init__.py
new file mode 100644
index 0000000..e5d0c4e
--- /dev/null
+++ b/utils/__init__.py
@@ -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
\ No newline at end of file
diff --git a/gpu_info.py b/utils/gpu_utils.py
similarity index 98%
rename from gpu_info.py
rename to utils/gpu_utils.py
index 14ead09..65f97fd 100644
--- a/gpu_info.py
+++ b/utils/gpu_utils.py
@@ -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"""
diff --git a/utils.py b/utils/helpers.py
similarity index 100%
rename from utils.py
rename to utils/helpers.py
diff --git a/openapi_spec_generator.py b/utils/openapi_utils.py
similarity index 99%
rename from openapi_spec_generator.py
rename to utils/openapi_utils.py
index 96e10ab..b9aabf3 100644
--- a/openapi_spec_generator.py
+++ b/utils/openapi_utils.py
@@ -1,4 +1,4 @@
-from .config import config
+from ..config import config
class OpenAPISpecGenerator:
def __init__(self, workflows: list):