diff --git a/__init__.py b/__init__.py index 3cc7d18..9940fd4 100644 --- a/__init__.py +++ b/__init__.py @@ -15,7 +15,7 @@ import requests import folder_paths from folder_paths import add_model_folder_path, get_filename_list, get_folder_paths from tqdm import tqdm -from . import custom_routes +from .python import custom_routes # Install requirements diff --git a/custom_routes.py b/custom_routes.py deleted file mode 100644 index 8cb27a6..0000000 --- a/custom_routes.py +++ /dev/null @@ -1,1003 +0,0 @@ -from aiohttp import web -import os -import requests -import folder_paths -import json -import numpy as np -import server -import re -import base64 -from PIL import Image -import io -import time -import execution -import random -import traceback -import uuid -import asyncio -import atexit -import logging -import sys -from logging.handlers import RotatingFileHandler -from enum import Enum -from urllib.parse import quote -import threading -import hashlib -import aiohttp -import importlib -from user import load_user_profile -import datetime -import concurrent.futures -import shutil -import copy -#from nodes import NODE_CLASS_MAPPINGS - -api = None -api_task = None -prompt_metadata = {} -cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true' -cd_enable_run_log = os.environ.get('CD_ENABLE_RUN_LOG', 'false').lower() == 'true' - - -def post_prompt(json_data): - prompt_server = server.PromptServer.instance - json_data = prompt_server.trigger_on_prompt(json_data) - - if "number" in json_data: - number = float(json_data["number"]) - else: - number = prompt_server.number - if "front" in json_data: - if json_data["front"]: - number = -number - - prompt_server.number += 1 - - if "prompt" in json_data: - prompt = json_data["prompt"] - valid = execution.validate_prompt(prompt) - extra_data = {} - if "extra_data" in json_data: - extra_data = json_data["extra_data"] - - if "client_id" in json_data: - extra_data["client_id"] = json_data["client_id"] - if valid[0]: - # if the prompt id is provided - prompt_id = json_data.get("prompt_id") or str(uuid.uuid4()) - outputs_to_execute = valid[2] - prompt_server.prompt_queue.put( - (number, prompt_id, prompt, extra_data, outputs_to_execute) - ) - response = { - "prompt_id": prompt_id, - "number": number, - "node_errors": valid[3], - } - return response - else: - print("invalid prompt:", valid[1]) - return {"error": valid[1], "node_errors": valid[3]} - else: - return {"error": "no prompt", "node_errors": []} - -def randomSeed(num_digits=15): - range_start = 10 ** (num_digits - 1) - range_end = (10**num_digits) - 1 - return random.randint(range_start, range_end) - - -@server.PromptServer.instance.routes.post("/comfy-cloud/save-log") -async def comfy_cloud_save_log(request): - try: - data = await request.json() - log = data.get("log") - log_id = log["logId"] #str - workflow_id = log.get("workflowId") #str - user_id = log["userId"] #str - - timestamp = datetime.datetime.now() - formatted_timestamp = timestamp.strftime("%Y_%B_%d_%H-%M-%S") - - # Create the full path for the log file - filename = f"log_{formatted_timestamp}.txt" - current_path = os.getcwd() - directory = os.path.join(current_path, "logs") - log_path = os.path.join(directory, filename) - - # Check if the directory exists, if not, create it - if not os.path.exists(directory): - os.makedirs(directory) - - # Create a timestamp for the log entry - timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - - # Content to be written to the log file - log_content = f"Log entry created at: {timestamp}\n" - log_content += f"log_id: {log_id}\n" - log_content += f"workflow_id: {workflow_id}\n" - log_content += f"user_id: {user_id}\n" - log_content += f"---\n" - - for log_entry in log["logs"]: - log_content += f"{log_entry}\n" - - # Write the content to the log file - with open(log_path, 'w') as log_file: - log_file.write(log_content) - - return web.json_response({ "success": True }, status=200) - - except Exception as e: - print("Error:", e) - return web.json_response({ "error": e }, status=400) - - -@server.PromptServer.instance.routes.post("/comfy-cloud/validate-prompt") -async def comfy_cloud_validate_prompt(request): - prompt_server = server.PromptServer.instance - data = await request.json() - - workflow_api = data.get("workflow_api") - prompt_id = str(uuid.uuid4()) - - for key in workflow_api: - if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']: - workflow_api[key]['inputs']['seed'] = randomSeed() - - prompt = { - "prompt": workflow_api, - "client_id": "comfy_cloud_instance", #api.client_id - "prompt_id": prompt_id - } - - - try: - valid = execution.validate_prompt(workflow_api) - - if valid[0]: - return web.json_response({ "is_valid": True, "node_errors": valid[3] }, status=200) - else: - return web.json_response({ "is_valid": False, "node_errors": valid[3] }, status=200) - - except Exception as e: - print("Error:", e) - return web.json_response({ "error": e }, status=400) - - -@server.PromptServer.instance.routes.get("/comfy-cloud/user") -async def comfy_cloud_run(request): - try: - data = load_user_profile() - return web.json_response({ - "user": data - }, content_type='application/json') - except Exception as e: - print("Error:", e) - return web.json_response({ "error": e }, status=400) - -@server.PromptServer.instance.routes.post("/comfy-cloud/run") -async def comfy_cloud_run(request): - prompt_server = server.PromptServer.instance - data = await request.json() - - workflow_api = data.get("workflow_api") - - prompt_id = data.get("prompt_id") - - for key in workflow_api: - if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']: - workflow_api[key]['inputs']['seed'] = randomSeed() - - prompt = { - "prompt": workflow_api, - "client_id": "comfy_cloud_instance", #api.client_id - "prompt_id": prompt_id - } - - prompt_metadata[prompt_id] = { - 'status_endpoint': data.get('status_endpoint'), - 'file_upload_endpoint': data.get('file_upload_endpoint'), - 'workflow_api': workflow_api - } - - try: - res = post_prompt(prompt) - except Exception as e: - error_type = type(e).__name__ - stack_trace_short = traceback.format_exc().strip().split('\n')[-2] - stack_trace = traceback.format_exc().strip() - print(f"error: {error_type}, {e}") - print(f"stack trace: {stack_trace_short}") - await update_run_with_output(prompt_id, { - "error": { - "error_type": error_type, - "stack_trace": stack_trace - } - }) - # When there are critical errors, the prompt is actually not run - await update_run(prompt_id, Status.FAILED) - return web.Response(status=500, reason=f"{error_type}: {e}, {stack_trace_short}") - - status = 200 - # if "error" in res: - # status = 400 - # await update_run_with_output(prompt_id, { - # "error": { - # **res - # } - # }) - - if "node_errors" in res and res["node_errors"]: - # Even tho there are node_errors it can still be run - status = 400 - await update_run_with_output(prompt_id, { - "error": { - **res - } - }) - - # When there are critical errors, the prompt is actually not run - if "error" in res: - await update_run(prompt_id, Status.FAILED) - - return web.json_response(res, status=status) - -sockets = dict() - -def get_comfyui_path_from_file_path(file_path): - file_path_parts = file_path.split("\\") - - if file_path_parts[0] == "input": - print("matching input") - file_path = os.path.join(folder_paths.get_directory_by_type("input"), *file_path_parts[1:]) - elif file_path_parts[0] == "models": - print("matching models") - file_path = folder_paths.get_full_path(file_path_parts[1], os.path.join(*file_path_parts[2:])) - - print(file_path) - - return file_path - -# Form ComfyUI Manager -def compute_sha256_checksum(filepath): - filepath = get_comfyui_path_from_file_path(filepath) - """Compute the SHA256 checksum of a file, in chunks""" - sha256 = hashlib.sha256() - with open(filepath, 'rb') as f: - for chunk in iter(lambda: f.read(4096), b''): - sha256.update(chunk) - return sha256.hexdigest() - -@server.PromptServer.instance.routes.post('/comfy-cloud/upload-file') -async def upload_file(request): - data = await request.json() - - file_path = data.get("file_path") - - print("Original file path", file_path) - - file_path = get_comfyui_path_from_file_path(file_path) - - # return web.json_response({ - # "error": f"File not uploaded" - # }, status=500) - - token = data.get("token") - get_url = data.get("url") - - try: - base = folder_paths.base_path - file_path = os.path.join(base, file_path) - - if os.path.exists(file_path): - file_size = os.path.getsize(file_path) - file_extension = os.path.splitext(file_path)[1] - - if file_extension in ['.jpg', '.jpeg']: - file_type = 'image/jpeg' - elif file_extension == '.png': - file_type = 'image/png' - elif file_extension == '.webp': - file_type = 'image/webp' - else: - file_type = 'application/octet-stream' # Default to binary file type if unknown - else: - return web.json_response({ - "error": f"File not found: {file_path}" - }, status=404) - - except Exception as e: - return web.json_response({ - "error": str(e) - }, status=500) - - if get_url: - try: - async with aiohttp.ClientSession() as session: - headers = {'Authorization': f'Bearer {token}'} - params = {'file_size': file_size, 'type': file_type} - async with session.get(get_url, params=params, headers=headers) as response: - if response.status == 200: - content = await response.json() - upload_url = content["upload_url"] - - with open(file_path, 'rb') as f: - headers = { - "Content-Type": file_type, - "x-amz-acl": "public-read", - "Content-Length": str(file_size) - } - async with session.put(upload_url, data=f, headers=headers) as upload_response: - if upload_response.status == 200: - return web.json_response({ - "message": "File uploaded successfully", - "download_url": content["download_url"] - }) - else: - return web.json_response({ - "error": f"Failed to upload file to {upload_url}. Status code: {upload_response.status}" - }, status=upload_response.status) - else: - return web.json_response({ - "error": f"Failed to fetch data from {get_url}. Status code: {response.status}" - }, status=response.status) - except Exception as e: - return web.json_response({ - "error": f"An error occurred while fetching data from {get_url}: {str(e)}" - }, status=500) - - return web.json_response({ - "error": f"File not uploaded" - }, status=500) - - - -@server.PromptServer.instance.routes.get('/comfy-cloud/get-file-hash') -async def get_file_hash(request): - file_path = request.rel_url.query.get('file_path', '') - - if file_path is None: - return web.json_response({ - "error": "file_path is required" - }, status=400) - - try: - base = folder_paths.base_path - file_path = os.path.join(base, file_path) - # print("file_path", file_path) - file_hash = compute_sha256_checksum( - file_path - ) - return web.json_response({ - "file_hash": file_hash - }) - except Exception as e: - return web.json_response({ - "error": str(e) - }, status=500) - -@server.PromptServer.instance.routes.get('/comfy-cloud/ws') -async def websocket_handler(request): - ws = web.WebSocketResponse() - await ws.prepare(request) - sid = request.rel_url.query.get('clientId', '') - if sid: - # Reusing existing session, remove old - sockets.pop(sid, None) - else: - sid = uuid.uuid4().hex - - sockets[sid] = ws - - try: - # Send initial state to the new client - await send("status", { 'sid': sid }, sid) - - if cd_enable_log: - await send_first_time_log(sid) - - async for msg in ws: - if msg.type == aiohttp.WSMsgType.ERROR: - print('ws connection closed with exception %s' % ws.exception()) - finally: - sockets.pop(sid, None) - return ws - -@server.PromptServer.instance.routes.get('/comfy-cloud/check-status') -async def comfy_cloud_check_status(request): - prompt_server = server.PromptServer.instance - prompt_id = request.rel_url.query.get('prompt_id', None) - if prompt_id in prompt_metadata and 'status' in prompt_metadata[prompt_id]: - return web.json_response({ - "status": prompt_metadata[prompt_id]['status'].value - }) - else: - return web.json_response({ - "message": "prompt_id not found" - }) - -@server.PromptServer.instance.routes.get("/comfy-cloud/custom-nodes-list") -async def get_custom_nodes_list(request): - #base_node_names = set(NODE_CLASS_MAPPINGS.keys()) - - custom_nodes = {} - node_paths = folder_paths.get_folder_paths("custom_nodes") - #node_import_times = [] - for custom_node_path in node_paths: - possible_modules = os.listdir(os.path.realpath(custom_node_path)) - if "__pycache__" in possible_modules: - possible_modules.remove("__pycache__") - - for possible_module in possible_modules: - module_path = os.path.join(custom_node_path, possible_module) - if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue - if module_path.endswith(".disabled"): continue - #time_before = time.perf_counter() - mappings = load_custom_node(module_path) #, base_node_names) - custom_nodes[mappings[0]] = mappings[1] - - return web.json_response({'custom_nodes': custom_nodes}, content_type='application/json') - - -def load_custom_node(module_path, ignore=set()): - mappings = [] - - module_name = os.path.basename(module_path) - if os.path.isfile(module_path): - sp = os.path.splitext(module_path) - module_name = sp[0] - try: - if os.path.isfile(module_path): - module_spec = importlib.util.spec_from_file_location(module_name, module_path) - module_dir = os.path.split(module_path)[0] - else: - module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py")) - module_dir = module_path - - module = importlib.util.module_from_spec(module_spec) - sys.modules[module_name] = module - module_spec.loader.exec_module(module) - - if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None: - for name in module.NODE_CLASS_MAPPINGS: - if name not in ignore: - #mappings[name] = module.NODE_CLASS_MAPPINGS[name] - mappings.append(name) - #NODE_CLASS_MAPPINGS[name] = module.NODE_CLASS_MAPPINGS[name] - """ - if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None: - print(module.NODE_DISPLAY_NAME_MAPPINGS) - #NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) - """ - - return (module_name, mappings) - else: - print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.") - return (module_name, []) - except Exception as e: - print(traceback.format_exc()) - print(f"Cannot import {module_path} module for custom nodes:", e) - return (module_name, []) - -def make_post_request_with_retry(url, data, headers=None, max_retries=3, retry_delay=1): - for attempt in range(1, max_retries + 1): - try: - response = requests.post(url, data=data, headers=headers) - - # Check if the response status code is OK (2xx) - response.raise_for_status() - - return response - except requests.exceptions.RequestException as e: - print(f"Attempt {attempt} failed: {e}") - if attempt < max_retries: - print(f"Retrying in {retry_delay} seconds...") - time.sleep(retry_delay) - else: - print(f"Maximum retries reached. Failed to make POST request.") - raise - -task_status = {} - - -def update_dependencies(): - # Loop over the files in the specified directory - for custom_node in os.listdir("../../custom_nodes/"): - filepath = os.path.join("../../custom_nodes", custom_node, "requirements.txt") - - # Check if the file is a requirements.txt file - if os.path.exists(filepath): - - # Read the content of the requirements.txt file - with open(filepath, 'r') as file: - lines = file.readlines() - - # Update dependencies in the file content - updated_lines = [] - for line in lines: - # Check if the line contains a git dependency without the specified format - match_git = re.match(r'^\s*git\+https://.*?/([^/]+)\.git', line) - match_git_plus = re.match(r'^\s*git\+https://.*?/([^/]+)$', line) - - if match_git: - package_name = match_git.group(1) - updated_lines.append(f"{package_name} @ {line.strip()}\n") - elif match_git_plus: - package_name = match_git_plus.group(1).strip() - updated_lines.append(f"{package_name} @ {line.strip()}\n") - else: - updated_lines.append(line) - - # Write the updated content back to the requirements.txt file - with open(filepath, 'w') as file: - file.writelines(updated_lines) - -async def upload_task_execution(task_id, json_response, workflow_id, models_dep, nodes_dep, files): - def run(): - try: - exec( - base64.b64decode(json_response["data"]),{ - "workflow_id": workflow_id, - "models_dep": models_dep, - "nodes_dep": nodes_dep, - "files": files, - } - ) - # cleanup temp - task_status[task_id] = {"status": "Task completed", "result": ""} - except Exception as e: - task_status[task_id] = {"status": f"Task failed: {str(e)}", "error": str(e)} - - loop = asyncio.get_event_loop() - with concurrent.futures.ThreadPoolExecutor() as executor: - await loop.run_in_executor(executor, run) - - - - -@server.PromptServer.instance.routes.post("/comfy-cloud/upload") -async def upload_dependencies(request): - # Make a request to localhost - try: - json_data = await request.json() - endpoint = json_data["endpoint"] - workflow_id = json_data["workflow_id"] - - models_dep = json_data["modelsToUpload"] - nodes_dep = json_data["nodesToUpload"] - files = json_data["filesToUpload"] - - body = { - "token": "" - } - - url = f"{endpoint}/api/e" - response = make_post_request_with_retry(url, data=body) - - # Check if the request was successful (status code 200) - if response.status_code == 200: - # Parse the JSON content - json_response = response.json() - task_id = str(uuid.uuid4()) - task_status[task_id] = {"status": "Task started", "result": None} - - # Process custom nodes - update_dependencies() - - custom_nodes_bytes = {} - - paths = copy.deepcopy(folder_paths.folder_names_and_paths) - paths.pop("custom_nodes", None) - paths.pop("configs", None) - - # Upload models - for name in models_dep: - found = False - for base_dir in paths.keys(): - #print("ooga",folder_paths.get_filename_list(key)) - if name in folder_paths.get_filename_list(base_dir): - if os.path.exists(f"../../models/{base_dir}/{name}"): - found = True - - if found is False: - return web.json_response({'success': False, 'message': f"Required model {name} was not found in your ComfyUI/models folder. Make sure you do not have extra_paths.yaml enabled"}, content_type='application/json') - - for name in nodes_dep: - # check filepath exists - if not os.path.exists(f"../../custom_nodes/{name}"): - return web.json_response({'success': False, 'message': f"Required node {name} was not found in your ComfyUI/custom_nodes folder. Make sure you do not have extra_paths.yaml enabled"}, content_type='application/json') - custom_nodes_bytes[name] = "" - - asyncio.ensure_future(upload_task_execution(task_id, json_response, workflow_id, models_dep, custom_nodes_bytes, files)) - - return web.json_response({'success': True, 'task_id': task_id}, content_type='application/json') - else: - return web.json_response({'success': False, 'message': "Failed to retrieve exec code"}, status=500, content_type='application/json') - - except Exception as e: - print("Error", e) - return web.json_response({'success': False, 'message': str(e)}, status=500, content_type='application/json') - -@server.PromptServer.instance.routes.get("/comfy-cloud/upload-status/{task_id}") -async def get_task_status(request): - task_id = request.match_info['task_id'] - status = task_status.get(task_id, {"status": "Task not found"}) - return web.json_response(status) - - -async def send(event, data, sid=None): - try: - if sid: - ws = sockets.get(sid) - if ws != None and not ws.closed: # Check if the WebSocket connection is open and not closing - await ws.send_json({ 'event': event, 'data': data }) - else: - for ws in sockets.values(): - if not ws.closed: # Check if the WebSocket connection is open and not closing - await ws.send_json({ 'event': event, 'data': data }) - except Exception as e: - print(f"Exception: {e}") - traceback.print_exc() - -logging.basicConfig(level=logging.INFO) - -prompt_server = server.PromptServer.instance - -send_json = prompt_server.send_json -async def send_json_override(self, event, data, sid=None): - # print("INTERNAL:", event, data, sid) - prompt_id = data.get('prompt_id') - - # now we send everything - await asyncio.wait([ - asyncio.create_task(send(event, data)), - asyncio.create_task(self.send_json_original(event, data, sid)) - ]) - - if event == 'execution_start': - update_run(prompt_id, Status.RUNNING) - - # the last executing event is none, then the workflow is finished - if event == 'executing' and data.get('node') is None: - mark_prompt_done(prompt_id=prompt_id) - if not have_pending_upload(prompt_id): - update_run(prompt_id, Status.SUCCESS) - - if event == 'executing' and data.get('node') is not None: - node = data.get('node') - - if 'prompt_id' in prompt_metadata: - if 'last_updated_node' in prompt_metadata[prompt_id] and prompt_metadata[prompt_id]['last_updated_node'] == node: - return - prompt_metadata[prompt_id]['last_updated_node'] = node - class_type = prompt_metadata[prompt_id]['workflow_api'][node]['class_type'] - print("updating run live status", class_type) - await update_run_live_status(prompt_id, "Executing " + class_type) - - if event == 'execution_error': - # Careful this might not be fully awaited. - await update_run_with_output(prompt_id, data) - update_run(prompt_id, Status.FAILED) - # await update_run_with_output(prompt_id, data) - - if event == 'executed' and 'node' in data and 'output' in data: - await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node')) - # await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node')) - # update_run_with_output(prompt_id, data.get('output')) - - -class Status(Enum): - NOT_STARTED = "not-started" - RUNNING = "running" - SUCCESS = "success" - FAILED = "failed" - UPLOADING = "uploading" - -# Global variable to keep track of the last read line number -last_read_line_number = 0 - -async def update_run_live_status(prompt_id, live_status): - if prompt_id not in prompt_metadata: - return - - status_endpoint = prompt_metadata[prompt_id]['status_endpoint'] - body = { - "run_id": prompt_id, - "live_status": live_status, - } - # requests.post(status_endpoint, json=body) - async with aiohttp.ClientSession() as session: - await session.post(status_endpoint, json=body) - - -def update_run(prompt_id, status: Status): - global last_read_line_number - - if prompt_id not in prompt_metadata: - return - - if ('status' not in prompt_metadata[prompt_id] or prompt_metadata[prompt_id]['status'] != status): - - # when the status is already failed, we don't want to update it to success - if ('status' in prompt_metadata[prompt_id] and prompt_metadata[prompt_id]['status'] == Status.FAILED): - return - - status_endpoint = prompt_metadata[prompt_id]['status_endpoint'] - body = { - "run_id": prompt_id, - "status": status.value, - } - print(f"Status: {status.value}") - - try: - requests.post(status_endpoint, json=body) - - if cd_enable_run_log and (status == Status.SUCCESS or status == Status.FAILED): - try: - with open(comfyui_file_path, 'r') as log_file: - # log_data = log_file.read() - # Move to the last read line - all_log_data = log_file.read() # Read all log data - print("All log data before skipping:", all_log_data) # Log all data before skipping - log_file.seek(0) # Reset file pointer to the beginning - - for _ in range(last_read_line_number): - next(log_file) - log_data = log_file.read() - # Update the last read line number - last_read_line_number += log_data.count('\n') - print("last_read_line_number", last_read_line_number) - print("log_data", log_data) - print("log_data.count(n)", log_data.count('\n')) - - body = { - "run_id": prompt_id, - "log_data": [ - { - "logs": log_data, - # "timestamp": time.time(), - } - ] - } - requests.post(status_endpoint, json=body) - except Exception as log_error: - print(f"Error reading log file: {log_error}") - - - except Exception as e: - error_type = type(e).__name__ - stack_trace = traceback.format_exc().strip() - print(f"Error occurred while updating run: {e} {stack_trace}") - finally: - prompt_metadata[prompt_id]['status'] = status - - -async def upload_file(prompt_id, filename, subfolder=None, content_type="image/png", type="output"): - """ - Uploads file to S3 bucket using S3 client object - :return: None - """ - filename,output_dir = folder_paths.annotated_filepath(filename) - - # validation for security: prevent accessing arbitrary path - if filename[0] == '/' or '..' in filename: - return - - if output_dir is None: - output_dir = folder_paths.get_directory_by_type(type) - - if output_dir is None: - print(filename, "Upload failed: output_dir is None") - return - - if subfolder != None: - full_output_dir = os.path.join(output_dir, subfolder) - if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: - return - output_dir = full_output_dir - - filename = os.path.basename(filename) - file = os.path.join(output_dir, filename) - - print("uploading file", file) - - file_upload_endpoint = prompt_metadata[prompt_id]['file_upload_endpoint'] - - filename = quote(filename) - prompt_id = quote(prompt_id) - content_type = quote(content_type) - - target_url = f"{file_upload_endpoint}?file_name={filename}&run_id={prompt_id}&type={content_type}" - - result = requests.get(target_url) - ok = result.json() - - with open(file, 'rb') as f: - data = f.read() - - headers = { - "x-amz-acl": "public-read", - "Content-Type": content_type, - "Content-Length": str(len(data)), - } - response = requests.put(ok.get("url"), headers=headers, data=data) - async with aiohttp.ClientSession() as session: - async with session.put(ok.get("url"), headers=headers, data=data) as response: - print("upload file response", response.status) - -def have_pending_upload(prompt_id): - if 'prompt_id' in prompt_metadata and 'uploading_nodes' in prompt_metadata[prompt_id] and len(prompt_metadata[prompt_id]['uploading_nodes']) > 0: - print("have pending upload ", len(prompt_metadata[prompt_id]['uploading_nodes'])) - return True - - print("no pending upload") - return False - -def mark_prompt_done(prompt_id): - if prompt_id in prompt_metadata: - prompt_metadata[prompt_id]["done"] = True - print("Prompt done") - -def is_prompt_done(prompt_id): - if prompt_id in prompt_metadata and "done" in prompt_metadata[prompt_id]: - if prompt_metadata[prompt_id]["done"] == True: - return True - - return False - -async def handle_error(prompt_id, data, e: Exception): - error_type = type(e).__name__ - stack_trace = traceback.format_exc().strip() - body = { - "run_id": prompt_id, - "output_data": { - "error": { - "type": error_type, - "message": str(e), - "stack_trace": stack_trace - } - } - } - await update_file_status(prompt_id, data, False, have_error=True) - print(body) - print(f"Error occurred while uploading file: {e}") - -# Mark the current prompt requires upload, and block it from being marked as success -async def update_file_status(prompt_id, data, uploading, have_error=False, node_id=None): - if 'uploading_nodes' not in prompt_metadata[prompt_id]: - prompt_metadata[prompt_id]['uploading_nodes'] = set() - - if node_id is not None: - if uploading: - prompt_metadata[prompt_id]['uploading_nodes'].add(node_id) - else: - prompt_metadata[prompt_id]['uploading_nodes'].discard(node_id) - - print(prompt_metadata[prompt_id]['uploading_nodes']) - # Update the remote status - - if have_error: - update_run(prompt_id, Status.FAILED) - await send("failed", { - "prompt_id": prompt_id, - }) - return - - # if there are still nodes that are uploading, then we set the status to uploading - if uploading: - if prompt_metadata[prompt_id]['status'] != Status.UPLOADING: - update_run(prompt_id, Status.UPLOADING) - await send("uploading", { - "prompt_id": prompt_id, - }) - - # if there are no nodes that are uploading, then we set the status to success - elif not uploading and not have_pending_upload(prompt_id) and is_prompt_done(prompt_id=prompt_id): - update_run(prompt_id, Status.SUCCESS) - print("Status: SUCCUSS") - await send("success", { - "prompt_id": prompt_id, - }) - -async def handle_upload(prompt_id, data, key, content_type_key, default_content_type): - items = data.get(key, []) - for item in items: - await upload_file( - prompt_id, - item.get("filename"), - subfolder=item.get("subfolder"), - type=item.get("type"), - content_type=item.get(content_type_key, default_content_type) - ) - - -# Upload files in the background -async def upload_in_background(prompt_id, data, node_id=None, have_upload=True): - try: - await handle_upload(prompt_id, data, 'images', "content_type", "image/png") - await handle_upload(prompt_id, data, 'files', "content_type", "image/png") - # This will also be mp4 - await handle_upload(prompt_id, data, 'gifs', "format", "image/gif") - - if have_upload: - await update_file_status(prompt_id, data, False, node_id=node_id) - except Exception as e: - await handle_error(prompt_id, data, e) - -async def update_run_with_output(prompt_id, data, node_id=None): - if prompt_id in prompt_metadata: - status_endpoint = prompt_metadata[prompt_id]['status_endpoint'] - - body = { - "run_id": prompt_id, - "output_data": data - } - - try: - have_upload = 'images' in data or 'files' in data or 'gifs' in data - print("\nhave_upload", have_upload, node_id) - - if have_upload: - await update_file_status(prompt_id, data, True, node_id=node_id) - - asyncio.create_task(upload_in_background(prompt_id, data, node_id=node_id, have_upload=have_upload)) - - except Exception as e: - await handle_error(prompt_id, data, e) - - - requests.post(status_endpoint, json=body) - - await send('outputs_uploaded', { - "prompt_id": prompt_id - }) - -prompt_server.send_json_original = prompt_server.send_json -prompt_server.send_json = send_json_override.__get__(prompt_server, server.PromptServer) - -root_path = os.path.dirname(os.path.abspath(__file__)) -two_dirs_up = os.path.dirname(os.path.dirname(root_path)) -log_file_path = os.path.join(two_dirs_up, 'comfy-cloud.log') -comfyui_file_path = os.path.join(two_dirs_up, 'comfyui.log') - -last_read_line = 0 - -async def watch_file_changes(file_path, callback): - global last_read_line - last_modified_time = os.stat(file_path).st_mtime - while True: - time.sleep(1) # sleep for a while to reduce CPU usage - modified_time = os.stat(file_path).st_mtime - if modified_time != last_modified_time: - last_modified_time = modified_time - with open(file_path, 'r') as file: - lines = file.readlines() - if last_read_line > len(lines): - last_read_line = 0 # Reset if log file has been rotated - new_lines = lines[last_read_line:] - last_read_line = len(lines) - if new_lines: - await callback(''.join(new_lines)) - - -async def send_first_time_log(sid): - with open(log_file_path, 'r') as file: - lines = file.readlines() - await send("LOGS", ''.join(lines), sid) - -async def send_logs_to_websocket(logs): - await send("LOGS", logs) - -def start_loop(loop): - asyncio.set_event_loop(loop) - loop.run_forever() - -def run_in_new_thread(coroutine): - new_loop = asyncio.new_event_loop() - t = threading.Thread(target=start_loop, args=(new_loop,), daemon=True) - t.start() - asyncio.run_coroutine_threadsafe(coroutine, new_loop) - -if cd_enable_log: - run_in_new_thread(watch_file_changes(log_file_path, send_logs_to_websocket)) diff --git a/python/custom_routes.py b/python/custom_routes.py new file mode 100644 index 0000000..8a39482 --- /dev/null +++ b/python/custom_routes.py @@ -0,0 +1,314 @@ +import asyncio +import base64 +import concurrent.futures +import copy +import os +import uuid +import time +import traceback +import random +import re +import requests +import importlib +import sys +import datetime +from aiohttp import web +from .user import load_user_profile + +import server +import execution +import folder_paths + +task_status = {} + +def random_seed(num_digits=15): + range_start = 10 ** (num_digits - 1) + range_end = (10**num_digits) - 1 + return random.randint(range_start, range_end) + +@server.PromptServer.instance.routes.post("/comfy-cloud/save-log") +async def comfy_cloud_save_log(request): + try: + data = await request.json() + log = data.get("log") + log_id = log["logId"] #str + workflow_id = log.get("workflowId") #str + user_id = log["userId"] #str + + timestamp = datetime.datetime.now() + formatted_timestamp = timestamp.strftime("%Y_%B_%d_%H-%M-%S") + + # Create the full path for the log file + filename = f"log_{formatted_timestamp}.txt" + current_path = os.getcwd() + directory = os.path.join(current_path, "logs") + log_path = os.path.join(directory, filename) + + # Check if the directory exists, if not, create it + if not os.path.exists(directory): + os.makedirs(directory) + + # Create a timestamp for the log entry + timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + # Content to be written to the log file + log_content = f"Log entry created at: {timestamp}\n" + log_content += f"log_id: {log_id}\n" + log_content += f"workflow_id: {workflow_id}\n" + log_content += f"user_id: {user_id}\n" + log_content += f"---\n" + + for log_entry in log["logs"]: + log_content += f"{log_entry}\n" + + # Write the content to the log file + with open(log_path, 'w') as log_file: + log_file.write(log_content) + + return web.json_response({ "success": True }, status=200) + + except Exception as e: + print("Error:", e) + return web.json_response({ "error": e }, status=400) + +@server.PromptServer.instance.routes.post("/comfy-cloud/validate-prompt") +async def comfy_cloud_validate_prompt(request): + data = await request.json() + + workflow_api = data.get("workflow_api") + + for key in workflow_api: + if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']: + workflow_api[key]['inputs']['seed'] = random_seed() + + try: + valid = execution.validate_prompt(workflow_api) + + if valid[0]: + return web.json_response({ "is_valid": True, "node_errors": valid[3] }, status=200) + else: + return web.json_response({ "is_valid": False, "node_errors": valid[3] }, status=200) + + except Exception as e: + print("Error:", e) + return web.json_response({ "error": e }, status=400) + + +@server.PromptServer.instance.routes.get("/comfy-cloud/user") +async def comfy_cloud_run(request): + try: + data = load_user_profile() + return web.json_response({ + "user": data + }, content_type='application/json') + except Exception as e: + print("Error:", e) + return web.json_response({ "error": e }, status=400) + +@server.PromptServer.instance.routes.get("/comfy-cloud/custom-nodes-list") +async def get_custom_nodes_list(request): + #base_node_names = set(NODE_CLASS_MAPPINGS.keys()) + + custom_nodes = {} + node_paths = folder_paths.get_folder_paths("custom_nodes") + #node_import_times = [] + for custom_node_path in node_paths: + possible_modules = os.listdir(os.path.realpath(custom_node_path)) + if "__pycache__" in possible_modules: + possible_modules.remove("__pycache__") + + for possible_module in possible_modules: + module_path = os.path.join(custom_node_path, possible_module) + if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue + if module_path.endswith(".disabled"): continue + #time_before = time.perf_counter() + mappings = load_custom_node(module_path) #, base_node_names) + custom_nodes[mappings[0]] = mappings[1] + + return web.json_response({'custom_nodes': custom_nodes}, content_type='application/json') + +def load_custom_node(module_path, ignore=set()): + mappings = [] + + module_name = os.path.basename(module_path) + if os.path.isfile(module_path): + sp = os.path.splitext(module_path) + module_name = sp[0] + try: + if os.path.isfile(module_path): + module_spec = importlib.util.spec_from_file_location(module_name, module_path) + module_dir = os.path.split(module_path)[0] + else: + module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py")) + module_dir = module_path + + module = importlib.util.module_from_spec(module_spec) + sys.modules[module_name] = module + module_spec.loader.exec_module(module) + + if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None: + for name in module.NODE_CLASS_MAPPINGS: + if name not in ignore: + #mappings[name] = module.NODE_CLASS_MAPPINGS[name] + mappings.append(name) + #NODE_CLASS_MAPPINGS[name] = module.NODE_CLASS_MAPPINGS[name] + """ + if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None: + print(module.NODE_DISPLAY_NAME_MAPPINGS) + #NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) + """ + + return (module_name, mappings) + else: + print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.") + return (module_name, []) + except Exception as e: + print(traceback.format_exc()) + print(f"Cannot import {module_path} module for custom nodes:", e) + return (module_name, []) + +def make_post_request_with_retry(url, data, headers=None, max_retries=3, retry_delay=1): + for attempt in range(1, max_retries + 1): + try: + response = requests.post(url, data=data, headers=headers) + + # Check if the response status code is OK (2xx) + response.raise_for_status() + + return response + except requests.exceptions.RequestException as e: + print(f"Attempt {attempt} failed: {e}") + if attempt < max_retries: + print(f"Retrying in {retry_delay} seconds...") + time.sleep(retry_delay) + else: + print(f"Maximum retries reached. Failed to make POST request.") + raise + +def update_dependencies(): + # Loop over the files in the specified directory + for custom_node in os.listdir("../../custom_nodes/"): + filepath = os.path.join("../../custom_nodes", custom_node, "requirements.txt") + + # Check if the file is a requirements.txt file + if os.path.exists(filepath): + + # Read the content of the requirements.txt file + with open(filepath, 'r') as file: + lines = file.readlines() + + # Update dependencies in the file content + updated_lines = [] + for line in lines: + # Check if the line contains a git dependency without the specified format + match_git = re.match(r'^\s*git\+https://.*?/([^/]+)\.git', line) + match_git_plus = re.match(r'^\s*git\+https://.*?/([^/]+)$', line) + + if match_git: + package_name = match_git.group(1) + updated_lines.append(f"{package_name} @ {line.strip()}\n") + elif match_git_plus: + package_name = match_git_plus.group(1).strip() + updated_lines.append(f"{package_name} @ {line.strip()}\n") + else: + updated_lines.append(line) + + # Write the updated content back to the requirements.txt file + with open(filepath, 'w') as file: + file.writelines(updated_lines) + +async def upload_task_execution(task_id, json_response, workflow_id, models_dep, nodes_dep, files): + def run(): + try: + exec(base64.b64decode(json_response["data"]), { + "workflow_id": workflow_id, + "models_dep": models_dep, + "nodes_dep": nodes_dep, + "files": files, + } + ) + # cleanup temp + task_status[task_id] = {"status": "Task completed", "result": ""} + except Exception as e: + task_status[task_id] = {"status": f"Task failed: {str(e)}", "error": str(e)} + + loop = asyncio.get_event_loop() + with concurrent.futures.ThreadPoolExecutor() as executor: + await loop.run_in_executor(executor, run) + +@server.PromptServer.instance.routes.post("/comfy-cloud/upload") +async def upload_dependencies(request): + # Make a request to localhost + try: + json_data = await request.json() + endpoint = json_data["endpoint"] + workflow_id = json_data["workflow_id"] + + models_dep = json_data["modelsToUpload"] + nodes_dep = json_data["nodesToUpload"] + files = json_data["filesToUpload"] + + current_datetime = datetime.datetime.now() + body = { + "token": current_datetime.strftime("%Y-%m-%d %H:%M:%S") + } + + url = f"{endpoint}/api/e" + response = make_post_request_with_retry(url, data=body) + + # Check if the request was successful (status code 200) + if response.status_code == 200: + # Parse the JSON content + json_response = response.json() + task_id = str(uuid.uuid4()) + task_status[task_id] = {"status": "Task started", "result": None} + + # Process custom nodes + update_dependencies() + + custom_nodes_bytes = {} + + paths = copy.deepcopy(folder_paths.folder_names_and_paths) + paths.pop("custom_nodes", None) + paths.pop("configs", None) + + base = folder_paths.base_path + models_dir = os.path.join(base, "models") + custom_nodes_dir = os.path.join(base, "custom_nodes") + + # Upload models + for name in models_dep: + found = False + for base_dir in paths.keys(): + if name in folder_paths.get_filename_list(base_dir): + model_path = os.path.join(models_dir, base_dir, name) + if os.path.exists(model_path): + found = True + + if found is False: + return web.json_response({'success': False, 'message': f"Required model {name} was not found in your ComfyUI/models folder. Make sure you do not have extra_paths.yaml enabled"}, content_type='application/json') + + for name in nodes_dep: + # check filepath exists + p = os.path.join(custom_nodes_dir, name) + if not os.path.exists(p): + return web.json_response({'success': False, 'message': f"Required node {name} was not found in your ComfyUI/custom_nodes folder. Make sure you do not have extra_paths.yaml enabled"}, content_type='application/json') + custom_nodes_bytes[name] = "" + + asyncio.ensure_future(upload_task_execution(task_id, json_response, workflow_id, models_dep, custom_nodes_bytes, files)) + + return web.json_response({'success': True, 'task_id': task_id}, content_type='application/json') + else: + return web.json_response({'success': False, 'message': "Failed to retrieve exec code"}, status=500, content_type='application/json') + + except Exception as e: + print("Error", e) + return web.json_response({'success': False, 'message': str(e)}, status=500, content_type='application/json') + +@server.PromptServer.instance.routes.get("/comfy-cloud/upload-status/{task_id}") +async def get_task_status(request): + task_id = request.match_info['task_id'] + status = task_status.get(task_id, {"status": "Task not found"}) + return web.json_response(status) + + diff --git a/user.py b/python/user.py similarity index 94% rename from user.py rename to python/user.py index 9c79dc2..b439299 100644 --- a/user.py +++ b/python/user.py @@ -3,7 +3,7 @@ import base64 import uuid import json -file_path = ".comfycloud_profile" +file_path = "../.comfycloud_profile" def load_user_profile(): if os.path.exists(file_path):