Initial Commit

.gitignore

@@ -0,0 +1,156 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*.pyc
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintainted in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+*.pyc
+
+.idea
+/node_modules

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 daxcay
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

__init__.py

@@ -0,0 +1,419 @@
+import logging
+import json
+import os
+import shutil
+import datetime
+import uuid
+from aiohttp import web
+from server import PromptServer
+
+from .lib.WebSocketManager import WebSocketManager
+
+DATA_FOLDER = './nexus'
+
+PERMISSIONS_FILE = os.path.join(DATA_FOLDER, 'permissions.json')
+ADMINS_FILE = os.path.join(DATA_FOLDER, 'admins.json')
+TOKENS_FILE = os.path.join(DATA_FOLDER, 'tokens.json')
+NAMES_FILE = os.path.join(DATA_FOLDER, 'names.json')
+
+SERVER_FILE = os.path.join(DATA_FOLDER, 'server.json')
+WORKFLOWS_DIR = os.path.join(DATA_FOLDER, 'workflows')
+
+if not os.path.exists(DATA_FOLDER):
+    os.makedirs(DATA_FOLDER)
+
+if not os.path.exists(ADMINS_FILE):
+    with open(ADMINS_FILE, 'w') as f:
+        admins_data = {'epic': 'comfynexus'}
+        json.dump(admins_data, f)
+
+for file_path in [PERMISSIONS_FILE, ADMINS_FILE, TOKENS_FILE, NAMES_FILE]:
+    if not os.path.exists(file_path):
+        with open(file_path, 'w') as f:
+            json.dump({}, f)
+
+if not os.path.exists(WORKFLOWS_DIR):
+    os.makedirs(WORKFLOWS_DIR)
+
+if not os.path.exists(SERVER_FILE):
+    with open(SERVER_FILE, 'w') as f:
+        default_config = {
+            "name": "Default Nexus Server",
+            "max_user": 15,
+            "config": {
+                "last_saved_user_id": "id"
+            }
+        }
+        json.dump(default_config, f, indent=4)
+
+
+if hasattr(PromptServer, 'instance'):
+    server = PromptServer.instance
+    routes = web.RouteTableDef()
+    socket_manager = WebSocketManager()
+
+    async def broadcast_message(message, sender):
+        for id, socket in socket_manager.sockets.items():
+            if id != sender:
+                try:
+                    await socket.send_str(json.dumps(message))
+                except Exception as e:
+                    logging.error(f"[ComfyUI-Nexus]: Error sending message to {id}: {e}")
+
+    async def broadcast_message_except(message, excluder):
+        for id, socket in socket_manager.sockets.items():
+            if id != excluder:
+                try:
+                    await socket.send_str(json.dumps(message))
+                except Exception as e:
+                    logging.error(f"[ComfyUI-Nexus]: Error sending message to {id}: {e}")
+
+    async def broadcast_message_all(message):
+        for id, socket in socket_manager.sockets.items():
+            try:
+                await socket.send_str(json.dumps(message))
+            except Exception as e:
+                logging.error(f"[ComfyUI-Nexus]: Error sending message to {id}: {e}")
+
+    async def send_message(message, sender, receiver):
+        for id, socket in socket_manager.sockets.items():
+            if id == receiver and sender != receiver:
+                try:
+                    await socket.send_str(json.dumps(message))
+                except Exception as e:
+                    logging.error(f"[ComfyUI-Nexus]: Error sending message to {id}: {e}")
+
+    @routes.get("/nexus")
+    async def websocket_handler(request):
+        id = request.query.get("id")
+        socket = web.WebSocketResponse()
+        await socket.prepare(request)
+
+        socket_manager.set(id, socket)
+        logging.info(f"[ComfyUI-Nexus]: Socket {id} connected")
+
+        try:
+            async for msg in socket:
+                if msg.type == web.WSMsgType.TEXT:
+                    
+                    data = json.loads(msg.data)
+                    event_name = data.get('name')
+                    event_data = data.get('data')
+
+                    sender = data.get("from")
+                    all = data.get('all')
+                    receiver = data.get('receiver')
+                    exclude = data.get('exclude')
+
+                    # logging.info(f"[ComfyUI-Nexus]: Received {sender} {all} {receiver} {exclude}")
+                    
+                    if event_name == "chat-join":
+                        user_id = data.get("from")
+                        user_name = event_data.get("name")
+
+                        with open(NAMES_FILE, 'r+') as f:
+                            names = json.load(f)
+                            names[user_id] = user_name
+                            f.seek(0)
+                            json.dump(names, f, indent=4)
+                            f.truncate()
+
+                    if event_name == "nick":
+                        user_id = data.get("from")
+                        user_name = event_data.get("new_name")
+
+                        with open(NAMES_FILE, 'r+') as f:
+                            names = json.load(f)
+                            names[user_id] = user_name
+                            f.seek(0)
+                            json.dump(names, f, indent=4)
+                            f.truncate()
+
+                    if all:
+                        await broadcast_message_all(data)
+                    elif receiver:
+                        await send_message(data, sender, receiver)
+                    elif exclude:
+                        await broadcast_message_except(data, exclude)
+                    else:
+                        await broadcast_message(data, sender)
+
+                elif msg.type == web.WSMsgType.ERROR:
+                    logging.error(f"[ComfyUI-Nexus]: WebSocket error: {msg.exception()}")
+        except Exception as e:
+            logging.error(f"[ComfyUI-Nexus]: WebSocket error: {e}")
+        finally:
+            logging.info(f"[ComfyUI-Nexus]: Socket {id} disconnected")
+            socket_manager.delete(id)
+
+        return socket
+
+    @routes.post("/nexus/workflows")
+    async def get_user_specific_workflows(request):
+
+        user_id = (await request.json()).get('user_id')
+        if not user_id:
+            return web.json_response({'error': 'User ID required'}, status=400)
+
+        with open(PERMISSIONS_FILE, 'r') as f:
+            permissions = json.load(f)
+
+        user_permissions = permissions.get(user_id, {})
+        if not user_permissions.get('editor', False):
+            return web.json_response({'error': 'Forbidden'}, status=403)
+
+        user_workflow_dir = os.path.join(WORKFLOWS_DIR, user_id)
+        if not os.path.exists(user_workflow_dir):
+            return web.json_response({'error': 'User not found'}, status=404)
+
+        workflows = {}
+        for file_name in sorted([f for f in os.listdir(user_workflow_dir) if f.isdigit()], key=int):
+            file_path = os.path.join(user_workflow_dir, file_name)
+            with open(file_path, 'r') as f:
+                workflows[file_name] = json.load(f)
+
+        return web.json_response(workflows)
+
+    @routes.post("/nexus/workflows/meta")
+    async def get_user_specific_workflows_meta(request):
+ 
+        user_id = (await request.json()).get('user_id')
+        if not user_id:
+            return web.json_response({'error': 'User ID required'}, status=400)
+        with open(PERMISSIONS_FILE, 'r') as f:
+            permissions = json.load(f)
+
+        user_permissions = permissions.get(user_id, {})
+        if not user_permissions.get('editor', False):
+            return web.json_response({'error': 'Forbidden'}, status=403)
+
+        user_workflow_dir = os.path.join(WORKFLOWS_DIR, user_id)
+        if not os.path.exists(user_workflow_dir):
+            return web.json_response({'error': 'User not found'}, status=404)
+
+        meta_data = []
+        for file_name in sorted([f for f in os.listdir(user_workflow_dir) if f.isdigit()], key=int):
+            file_path = os.path.join(user_workflow_dir, file_name)
+            last_modified_time = os.path.getmtime(file_path)
+            last_modified_readable = datetime.datetime.fromtimestamp(last_modified_time).strftime('%Y-%m-%d %H:%M:%S')
+
+            meta_data.append({
+                'file_name': file_name,
+                'last_saved': last_modified_readable
+            })
+
+        return web.json_response(meta_data)    
+    
+    @routes.post("/nexus/workflows/{id}")
+    async def get_user_specific_workflow(request):
+        user_id = (await request.json()).get('user_id')
+        workflow_id = request.match_info.get('id')
+
+        if not user_id:
+            return web.json_response({'error': 'User ID required'}, status=400)
+        if not workflow_id:
+            return web.json_response({'error': 'Workflow ID required'}, status=400)
+
+        with open(PERMISSIONS_FILE, 'r') as f:
+            permissions = json.load(f)
+
+        user_permissions = permissions.get(user_id, {})
+        if not user_permissions.get('editor', False):
+            return web.json_response({'error': 'Forbidden'}, status=403)
+
+        user_workflow_dir = os.path.join(WORKFLOWS_DIR, user_id)
+        if not os.path.exists(user_workflow_dir):
+            return web.json_response({'error': 'User not found'}, status=404)
+
+        workflow_file_path = os.path.join(user_workflow_dir, workflow_id)
+        if not os.path.exists(workflow_file_path):
+            return web.json_response({'error': 'Workflow not found'}, status=404)
+
+        with open(workflow_file_path, 'r') as f:
+            workflow = json.load(f)
+
+        return web.json_response(workflow)
+
+    @routes.post("/nexus/workflow")
+    async def post_workflow(request):
+        user_id = (await request.json()).get('user_id')
+        if not user_id:
+            return web.json_response({'error': 'User ID required'}, status=400)
+
+        with open(PERMISSIONS_FILE, 'r') as f:
+            permissions = json.load(f)
+
+        user_permissions = permissions.get(user_id, {})
+        if not user_permissions.get('editor', False):
+            return web.json_response({'error': 'Forbidden'}, status=403)
+
+        new_data = await request.json()
+
+        user_workflow_dir = os.path.join(WORKFLOWS_DIR, user_id)
+        if not os.path.exists(user_workflow_dir):
+            os.makedirs(user_workflow_dir)
+
+        existing_files = sorted([f for f in os.listdir(user_workflow_dir) if f.isdigit()], key=int)
+
+        if len(existing_files) == 5:
+            os.remove(os.path.join(user_workflow_dir, existing_files[0]))
+            existing_files.pop(0)
+
+        for i, file_name in enumerate(existing_files):
+            new_name = str(i + 1)
+            shutil.move(os.path.join(user_workflow_dir, file_name), os.path.join(user_workflow_dir, new_name))
+
+        new_file_path = os.path.join(user_workflow_dir, '5')
+        with open(new_file_path, 'w') as f:
+            json.dump(new_data, f, indent=4)
+
+        # Update the server.json file with the last saved user_id
+        with open(SERVER_FILE, 'r+') as f:
+            server_config = json.load(f)
+            server_config['config']['last_saved_user_id'] = user_id
+            f.seek(0)
+            json.dump(server_config, f, indent=4)
+            f.truncate()
+
+        return web.json_response({'status': 'success'}, status=200)
+    
+    @routes.get("/nexus/workflow")
+    async def get_latest_workflow(request):
+        with open(SERVER_FILE, 'r') as f:
+            server_config = json.load(f)
+        
+        last_saved_user_id = server_config['config'].get('last_saved_user_id')
+        if not last_saved_user_id:
+            return web.json_response({'error': 'No workflow saved'}, status=404)
+
+        user_workflow_dir = os.path.join(WORKFLOWS_DIR, last_saved_user_id)
+        if not os.path.exists(user_workflow_dir):
+            return web.json_response({'error': 'User not found'}, status=404)
+
+        last_file_path = os.path.join(user_workflow_dir, '5')
+        if not os.path.exists(last_file_path):
+            return web.json_response({'error': 'Workflow not found'}, status=404)
+
+        with open(last_file_path, 'r') as f:
+            workflow_data = json.load(f)
+
+        return web.json_response(workflow_data)
+    
+    @routes.get("/nexus/permission/{id}")
+    async def get_permission_by_id(request):
+        user_id = request.match_info['id']
+        with open(PERMISSIONS_FILE, 'r') as f:
+            data = json.load(f)
+        user_permission = data.get(user_id)
+        if user_permission:
+            return web.json_response(user_permission)
+        else:
+            return web.json_response({'error': 'User not found'}, status=404)
+
+    @routes.post("/nexus/permission/{id}")
+    async def post_permission_by_id(request):
+        user_id = request.match_info['id']
+        data = await request.json()
+        admin_id = data.get('admin_id')
+        if not admin_id:
+            return web.json_response({'error': 'Admin ID required'}, status=400)
+
+        token = request.headers.get('Authorization')
+        if not token:
+            return web.json_response({'error': 'Authorization token required'}, status=400)
+
+        with open(TOKENS_FILE, 'r') as f:
+            tokens = json.load(f)
+
+        if tokens.get(token) != admin_id:
+            return web.json_response({'error': 'Invalid or expired token'}, status=403)
+
+        new_permission = data.get('data')
+        if not new_permission:
+            return web.json_response({'error': 'Permission data required'}, status=400)
+
+        with open(PERMISSIONS_FILE, 'r+') as f:
+            permissions = json.load(f)
+
+            if not permissions.get(admin_id, {}).get('admin', False):
+                return web.json_response({'error': 'Forbidden'}, status=403)
+
+            # Merge the new permissions with the existing ones
+            existing_permissions = permissions.get(user_id, {})
+            existing_permissions.update(new_permission)
+
+            # Save the updated permissions
+            permissions[user_id] = existing_permissions
+            f.seek(0)
+            json.dump(permissions, f, indent=4)
+            f.truncate()
+
+        return web.json_response({'status': 'success'}, status=200)
+
+    @routes.post("/nexus/login")
+    async def post_login(request):
+
+        credentials = await request.json()
+        user_id = credentials.get('uuid')
+        account = credentials.get('account')
+        password = credentials.get('password')
+
+        with open(ADMINS_FILE, 'r') as f:
+            admins = json.load(f)
+
+        if admins.get(account) == password:
+            token = str(uuid.uuid4())
+
+            with open(TOKENS_FILE, 'r+') as f:
+                tokens = json.load(f)
+                tokens[token] = user_id
+                f.seek(0)
+                json.dump(tokens, f, indent=4)
+                f.truncate()
+
+            with open(PERMISSIONS_FILE, 'r+') as f:
+                permissions = json.load(f)
+                permissions[user_id] = {'editor': True, 'admin': True}
+                f.seek(0)
+                json.dump(permissions, f, indent=4)
+                f.truncate()
+
+            return web.json_response({'token': token}, status=200)
+
+        else:
+            return web.json_response({'error': 'Invalid credentials'}, status=401)
+
+    @routes.post("/nexus/verify")
+    async def post_verify(request):
+        credentials = await request.json()
+        token = credentials.get('token')
+
+        with open(TOKENS_FILE, 'r') as f:
+            tokens = json.load(f)
+
+        user_id = tokens.get(token)
+        if user_id:
+            with open(PERMISSIONS_FILE, 'r') as f:
+                permissions = json.load(f)
+
+            user_permissions = permissions.get(user_id, {})
+            return web.json_response({'user_id': user_id, 'permissions': user_permissions}, status=200)
+        else:
+            return web.json_response({'error': 'Invalid token'}, status=401)
+
+    @routes.get("/nexus/name/{id}")
+    async def get_name_by_id(request):
+        user_id = request.match_info['id']
+        with open(NAMES_FILE, 'r') as f:
+            names = json.load(f)
+        user_name = names.get(user_id)
+        if user_name:
+            return web.json_response({'name': user_name})
+        else:
+            return web.json_response({'error': 'User not found'}, status=404)
+
+    server.app.add_routes(routes)
+
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+WEB_DIRECTORY = "./web"

pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "comfyui-nexus"
+description = "Seamless Multiuser Workflow Collaboration"
+version = "1.0"
+license = "LICENSE"
+dependencies = []
+
+[project.urls]
+Repository = "https://github.com/daxcay/ComfyUI-Nexus"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "daxcay"
+DisplayName = "ComfyUI-Nexus"
+Icon = "https://raw.githubusercontent.com/daxcay/ComfyUI-Nexus/master/web/assets/img/logo.jpg"