import requests import time import json from PIL import Image from io import BytesIO import numpy as np import torch import os import websocket # 需要安装 websocket-client 包 import threading import comfy.utils # Import comfy utils for ProgressBar # Try importing folder_paths safely try: import folder_paths comfyui_env_available = True # Use a more generic name except ImportError: comfyui_env_available = False print("ComfyUI folder_paths not found. Some features like specific output paths might use fallbacks.") class ExecuteNode: ESTIMATED_TOTAL_NODES = 10 # Default estimate def __init__(self): self.ws = None self.task_completed = False self.ws_error = None self.executed_nodes = set() self.prompt_tips = "{}" self.pbar = None self.node_lock = threading.Lock() self.total_nodes = None self.current_steps = 0 # Track current steps for logging def update_progress(self): """Increments the progress bar by one step and logs, stopping at total_nodes.""" # --- Guard Condition --- # Use lock to ensure thread safety when checking/updating steps and flag with self.node_lock: if self.task_completed or (self.pbar and self.current_steps >= self.total_nodes): # Print only if trying to update *after* completion for debugging if self.task_completed: print(f"Skipping progress update because task is already completed.") return if self.pbar: self.current_steps += 1 # Increment the ComfyUI progress bar by 1 self.pbar.update(1) # Log the current state display_steps = min(self.current_steps, self.total_nodes) # Ensure log doesn't exceed total print(f"Progress Update: Step {display_steps}/{self.total_nodes} ({(display_steps/self.total_nodes)*100:.1f}%)") def complete_progress(self): """Sets the progress bar to 100% and marks task as completed.""" # --- Use lock for thread safety --- with self.node_lock: # Check if already completed to prevent redundant calls/logs if self.task_completed: return print(f"Finalizing progress: Setting task_completed = True") # --- Set completion flag FIRST --- self.task_completed = True if self.pbar: # Ensure internal step count reflects completion self.current_steps = self.total_nodes # Use update_absolute for the final state self.pbar.update_absolute(1.0) print(f"Progress Finalized: {self.total_nodes}/{self.total_nodes} (100.0%)") else: print("Progress bar not available during finalization.") @classmethod def INPUT_TYPES(cls): return { "required": { "apiConfig": ("STRUCT",), }, "optional": { "nodeInfoList": ("ARRAY", {"default": []}), "run_timeout": ("INT", {"default": 600}), # query_interval is no longer used for the main loop timing # "query_interval": ("INT", {"default": 10}), "concurrency_limit": ("INT", {"default": 1, "min": 1}), "estimated_total_nodes": ("INT", {"default": cls.ESTIMATED_TOTAL_NODES, "min": 1}), }, } RETURN_TYPES = ("IMAGE", "VIDEO") RETURN_NAMES = ("images", "videos") CATEGORY = "RunningHub" FUNCTION = "process" OUTPUT_NODE = True # Indicate support for progress display # --- WebSocket Handlers --- def on_ws_message(self, ws, message): """处理 WebSocket 消息，更新内部状态和进度条""" try: # --- Check completion status AT THE START --- # This check is implicitly thread-safe due to complete_progress lock if self.task_completed: # print("WS Message received after task completion, ignoring.") # Optional: reduce log spam return data = json.loads(message) message_type = data.get("type") if message_type == "executing": node_data = data.get("data", {}) node_id = node_data.get("node") if node_id is not None: # Lock is handled within update_progress now # Check if it's a new node before calling update if node_id not in self.executed_nodes: self.executed_nodes.add(node_id) # Add before update call self.update_progress() # This method is now guarded internally print(f"WS: Node {node_id} executed.") else: # Null node signal check remains guarded by the top-level check print("WS: Received null node signal, waiting for final success signal...") elif message_type == "execution_success": # The internal check in complete_progress handles redundancy print("WS: Task execution success signal received.") self.complete_progress() # No need for sleep or setting task_completed here except Exception as e: print(f"Error processing WebSocket message: {e}") self.ws_error = e # Call complete_progress which handles the task_completed flag and lock self.complete_progress() def on_ws_error(self, ws, error): """处理 WebSocket 错误""" print(f"WebSocket error: {error}") self.ws_error = error # Mark task as complete via the centralized method self.complete_progress() def on_ws_close(self, ws, close_status_code, close_msg): """处理 WebSocket 关闭""" print(f"WebSocket closed: {close_status_code} - {close_msg}") # If closed unexpectedly, mark as complete to end loop # Use lock temporarily just to read task_completed safely with self.node_lock: should_complete = not self.task_completed if should_complete: print("Warning: WebSocket closed unexpectedly. Forcing task completion.") self.ws_error = self.ws_error or IOError(f"WebSocket closed unexpectedly ({close_status_code})") # Mark task as complete via the centralized method self.complete_progress() def on_ws_open(self, ws): """处理 WebSocket 连接打开""" print("WebSocket connection established") # Note: executed_nodes should be cleared at the start of 'process' def connect_websocket(self, wss_url): """建立 WebSocket 连接""" print(f"Connecting to WebSocket: {wss_url}") websocket.enableTrace(False) # Keep this false unless debugging WS protocol self.ws = websocket.WebSocketApp( wss_url, on_message=self.on_ws_message, on_error=self.on_ws_error, on_close=self.on_ws_close, on_open=self.on_ws_open ) ws_thread = threading.Thread(target=self.ws.run_forever, name="RH_ExecuteNode_WSThread") ws_thread.daemon = True ws_thread.start() print("WebSocket thread started.") def check_and_complete_task(self): """If task times out after null node, force completion.""" # complete_progress now checks the flag internally and uses lock print("Task completion timeout after null node signal - attempting forced completion.") self.complete_progress() # --- Main Process Method --- def process(self, apiConfig, nodeInfoList=None, run_timeout=600, concurrency_limit=1, estimated_total_nodes=ESTIMATED_TOTAL_NODES): # Reset state with self.node_lock: # Use lock for resetting shared state self.executed_nodes.clear() self.task_completed = False self.ws_error = None self.prompt_tips = "{}" self.current_steps = 0 # Reset step counter # 设置总节点数并初始化进度条 self.total_nodes = max(1, estimated_total_nodes) print(f"Using total nodes for progress: {self.total_nodes}") # Initialize ComfyUI progress bar self.pbar = comfy.utils.ProgressBar(self.total_nodes) print("Progress bar initialized at 0") # --- Concurrency Check --- api_key = None base_url = None try: api_key = apiConfig.get("apiKey") base_url = apiConfig.get("base_url") if not api_key or not base_url: raise ValueError("apiKey and base_url missing from apiConfig") account_status = self.check_account_status(api_key, base_url) current_tasks = int(account_status["currentTaskCounts"]) print(f"There are {current_tasks} tasks running") if current_tasks >= concurrency_limit: print(f"Concurrency limit ({concurrency_limit}) reached, waiting...") start_wait_time = time.time() # Use a shorter sleep interval while waiting for concurrency wait_interval = 2 # seconds while current_tasks >= concurrency_limit: if time.time() - start_wait_time > run_timeout: if self.pbar: self.pbar.update_absolute(1.0) # Use absolute directly for setup failure raise Exception(f"Timeout waiting for concurrent tasks ({current_tasks}/{concurrency_limit}) to finish.") print(f"Waiting for concurrent tasks... ({current_tasks}/{concurrency_limit})") time.sleep(wait_interval) account_status = self.check_account_status(api_key, base_url) current_tasks = int(account_status["currentTaskCounts"]) print("Concurrency slot available.") except Exception as e: print(f"Error checking account status or waiting: {e}") if self.pbar: self.pbar.update_absolute(1.0) # Use absolute directly for setup failure raise # --- Task Creation & WebSocket --- task_id = None try: print(f"ExecuteNode NodeInfoList: {nodeInfoList}") # Pass base_url explicitly from the validated config task_creation_result = self.create_task(apiConfig, nodeInfoList or [], base_url) print(f"Task Creation Result: {json.dumps(task_creation_result, indent=2, ensure_ascii=False)}") # Validate task creation response structure before accessing data if not isinstance(task_creation_result.get("data"), dict): raise ValueError("Invalid task creation response data structure.") self.prompt_tips = task_creation_result["data"].get("promptTips", "{}") task_id = task_creation_result["data"].get("taskId") wss_url = task_creation_result["data"].get("netWssUrl") if not task_id or not wss_url: raise ValueError("Missing taskId or netWssUrl in task creation response.") print(f"Task created successfully, taskId: {task_id}") self.connect_websocket(wss_url) except Exception as e: print(f"Error creating task or connecting WS: {e}") if self.pbar: self.pbar.update_absolute(1.0) # Use absolute directly for setup failure raise # --- Task Monitoring Loop --- task_start_time = time.time() loop_sleep_interval = 0.1 print("Starting task monitoring loop...") timeout_timer = None try: # Setup global timeout timer def force_timeout(): # Use lock to safely check task_completed with self.node_lock: is_completed = self.task_completed if not is_completed: print("Global timeout reached - forcing task completion.") self.ws_error = Exception("Global timeout reached") # Let the main loop call complete_progress via the finally block or error handling # Just set the flags here to break loop self.task_completed = True # Set flag directly here to break loop timeout_timer = threading.Timer(run_timeout, force_timeout) timeout_timer.daemon = True timeout_timer.start() # Main wait loop while True: # Check completion flags (read safely with lock) with self.node_lock: is_completed = self.task_completed current_error = self.ws_error if is_completed or current_error: break # Exit loop if completed or error occurred # Check for timeout explicitly in loop as backup/alternative to timer if time.time() - task_start_time > run_timeout: print("Task monitoring loop timeout check triggered.") # Set flags to exit loop; rely on finally block for completion with self.node_lock: if not self.task_completed: # Avoid overwriting WS error self.ws_error = self.ws_error or Exception(f"Timeout: Task {task_id} did not complete within {run_timeout} seconds.") self.task_completed = True # Ensure loop exit break # Exit loop time.sleep(loop_sleep_interval) # Yield CPU # Handle exit conditions after loop with self.node_lock: # Read error flag safely final_error = self.ws_error if final_error: print(f"Task ended with error: {final_error}") # Only complete progress if not already completed by WS handler with self.node_lock: if not self.task_completed: self.complete_progress() raise final_error # Re-raise the error else: # Task completed normally print("Task monitoring completed successfully.") # complete_progress should have been called by WS handler finally: # Cleanup if timeout_timer: timeout_timer.cancel() if self.ws: try: self.ws.close() except Exception as e: print(f"Error closing WebSocket: {e}") self.ws = None # Final safety net: Only complete progress if not already completed with self.node_lock: is_finally_completed = self.task_completed if not is_finally_completed: print("Warning: Monitoring loop ended unexpectedly. Finalizing progress via finally block.") self.complete_progress() # Call the safe completion method # --- Process Output --- print("Processing task output...") # Pass the validated api_key and base_url again return self.process_task_output(task_id, api_key, base_url) def process_task_output(self, task_id, api_key, base_url): """处理任务输出，包含轮询等待机制""" max_retries = 30 # 最多等待30次 retry_interval = 1 # 初始等待1秒 max_retry_interval = 5 # 最大等待间隔5秒 for attempt in range(max_retries): try: task_status_result = self.check_task_status(task_id, api_key, base_url) print(f"Check output attempt {attempt + 1}/{max_retries}") print("Task Status Result:", json.dumps(task_status_result, indent=2, ensure_ascii=False)) # 如果任务仍在运行，等待后重试 if isinstance(task_status_result, dict) and task_status_result.get("taskStatus") in ["RUNNING", "QUEUED"]: wait_time = min(retry_interval * (1.5 ** attempt), max_retry_interval) print(f"Task still running, waiting {wait_time:.1f} seconds before next check...") time.sleep(wait_time) continue # 如果获取到了实际的输出结果（文件列表） if isinstance(task_status_result, list) and len(task_status_result) > 0: print("Got valid output result, processing files...") image_urls = [] video_urls = [] for output in task_status_result: if isinstance(output, dict): file_url = output.get("fileUrl") file_type = output.get("fileType") if file_url and file_type: file_type_lower = file_type.lower() if file_type_lower in ["png", "jpg", "jpeg", "webp", "bmp", "gif"]: image_urls.append(file_url) elif file_type_lower in ["mp4", "avi", "mov", "webm"]: video_urls.append(file_url) # 处理图片和视频 image_data_list = [] if image_urls: print(f"Downloading {len(image_urls)} images...") for url in image_urls: try: print(f"Downloading image: {url}") image_tensor = self.download_image(url) if image_tensor is not None: image_data_list.append(image_tensor) except Exception as img_e: print(f"Error downloading image {url}: {img_e}") video_data_list = [] if video_urls: print(f"Downloading {len(video_urls)} videos...") for url in video_urls: try: print(f"Downloading video: {url}") video_path = self.download_video(url) if video_path is not None: video_data_list.append(video_path) except Exception as vid_e: print(f"Error downloading video {url}: {vid_e}") if image_data_list or video_data_list: print(f"Successfully got results: {len(image_data_list)} images, {len(video_data_list)} videos") return (image_data_list, video_data_list) # 如果是错误状态 if isinstance(task_status_result, dict) and task_status_result.get("taskStatus") == "error": print(f"Task error: {task_status_result.get('error', 'Unknown error')}") return ([], []) except Exception as e: print(f"Error checking task status (attempt {attempt + 1}): {e}") time.sleep(retry_interval) print(f"Failed to get valid output after {max_retries} attempts") return ([], []) def download_image(self, image_url): """ 从 URL 下载图像并转换为适合预览或保存的 torch.Tensor 格式。 包含重试机制，最多重试5次。 Returns tensor [H, W, C] or None on failure. """ max_retries = 5 retry_delay = 1 last_exception = None for attempt in range(max_retries): try: response = requests.get(image_url, timeout=30) print(f"Download image attempt {attempt + 1} ({image_url}): Status code: {response.status_code}") response.raise_for_status() content_type = response.headers.get('Content-Type', '').lower() img = Image.open(BytesIO(response.content)).convert("RGB") img_array = np.array(img).astype(np.float32) / 255.0 img_tensor = torch.from_numpy(img_array) # Shape: [H, W, C] return img_tensor except (requests.exceptions.RequestException, IOError, Image.UnidentifiedImageError) as e: print(f"Download image attempt {attempt + 1} failed: {e}") last_exception = e if attempt < max_retries - 1: print(f"Retrying in {retry_delay} seconds...") time.sleep(retry_delay) retry_delay *= 2 else: print(f"Failed to download image {image_url} after {max_retries} attempts.") return None return None def download_video(self, video_url): """ 从 URL 下载视频并保存到本地。 包含重试机制，最多重试5次。 Returns the local file path string or None on failure. """ max_retries = 5 retry_delay = 1 last_exception = None for attempt in range(max_retries): video_path = None try: output_dir = "output" if comfyui_env_available and hasattr(folder_paths, 'get_output_directory'): try: output_dir = folder_paths.get_output_directory() except Exception as e_dir: print(f"Warning: Could not get output directory from folder_paths: {e_dir}. Using default 'output'.") if not os.path.exists(output_dir): os.makedirs(output_dir) print(f"Created output directory: {output_dir}") try: parsed_url = requests.utils.urlparse(video_url) filename_from_url = os.path.basename(parsed_url.path) if parsed_url.path and '.' in os.path.basename(parsed_url.path) else None safe_filename_from_url = "".join(c if c.isalnum() or c in ['.', '_', '-'] else '_' for c in filename_from_url) if filename_from_url else None base_filename = safe_filename_from_url if safe_filename_from_url else f"RH_output_video_{str(int(time.time()*1000))}.mp4" counter = 0 video_path_base, video_ext = os.path.splitext(base_filename) max_base_len = 100 video_path_base = video_path_base[:max_base_len] video_path = os.path.join(output_dir, f"{video_path_base}{video_ext}") while os.path.exists(video_path): counter += 1 video_path = os.path.join(output_dir, f"{video_path_base}_{counter}{video_ext}") if counter > 100: print("Warning: Could not find unique filename after 100 attempts.") video_path = os.path.join(output_dir, f"{video_path_base}_{str(int(time.time()*1000))}{video_ext}") break except Exception as path_e: print(f"Error determining video path: {path_e}") fallback_filename = f"RH_output_video_fallback_{str(int(time.time()*1000))}.mp4" video_path = os.path.join(output_dir, fallback_filename) print(f"Attempting to download video to: {video_path}") response = requests.get(video_url, stream=True, timeout=60) print(f"Download video attempt {attempt + 1} ({video_url}): Status code: {response.status_code}") response.raise_for_status() content_type = response.headers.get('Content-Type', '').lower() if not any(vid_type in content_type for vid_type in ['video/', 'octet-stream']): print(f"Warning: Content-Type '{content_type}' may not be a video for URL {video_url}") downloaded_size = 0 with open(video_path, "wb") as f: for chunk in response.iter_content(chunk_size=65536): if chunk: f.write(chunk) downloaded_size += len(chunk) if downloaded_size > 0: print(f"Video saved successfully to {video_path} ({downloaded_size / (1024*1024):.2f} MB)") return video_path else: print(f"Warning: Downloaded video file is empty: {video_path}") if os.path.exists(video_path): os.remove(video_path) last_exception = IOError("Downloaded video file is empty.") except (requests.exceptions.RequestException, IOError) as e: print(f"Download video attempt {attempt + 1} failed: {e}") last_exception = e if video_path and os.path.exists(video_path): try: os.remove(video_path) print(f"Removed partial/failed download file: {video_path}") except OSError as rm_e: print(f"Error removing partial file {video_path}: {rm_e}") # Retry logic if attempt < max_retries - 1: print(f"Retrying in {retry_delay} seconds...") time.sleep(retry_delay) retry_delay *= 2 else: print(f"Failed to download video {video_url} after {max_retries} attempts.") return None return None def check_account_status(self, api_key, base_url): """ 查询账户状态，检查是否可以提交新任务。包含重试机制。 """ if not api_key or not base_url: raise ValueError("API Key and Base URL are required for checking account status.") url = f"{base_url}/uc/openapi/accountStatus" headers = { "User-Agent": "ComfyUI-RH-APICall-Node/1.0", "Content-Type": "application/json", } data = {"apikey": api_key} max_retries = 5 retry_delay = 1 last_exception = None for attempt in range(max_retries): response = None try: print(f"Attempt {attempt + 1}/{max_retries} to check account status...") response = requests.post(url, json=data, headers=headers, timeout=15) response.raise_for_status() result = response.json() if result.get("code") != 0: api_msg = result.get('msg', 'Unknown API error') print(f"API error on attempt {attempt + 1}: {api_msg}") raise Exception(f"API error getting account status: {api_msg}") account_data = result.get("data") if not account_data or "currentTaskCounts" not in account_data: raise ValueError("Invalid response structure for account status.") try: current_task_counts = int(account_data["currentTaskCounts"]) account_data["currentTaskCounts"] = current_task_counts print("Account status check successful.") return account_data # Success except (ValueError, TypeError) as e: raise ValueError(f"Invalid value for currentTaskCounts: {account_data.get('currentTaskCounts')}. Error: {e}") except (requests.exceptions.RequestException, json.JSONDecodeError, ValueError, Exception) as e: print(f"Error on attempt {attempt + 1}/{max_retries}: {e}") last_exception = e if isinstance(e, json.JSONDecodeError) and response is not None: print(f"Raw response text on JSON decode error: {response.text}") if attempt < max_retries - 1: print(f"Retrying in {retry_delay} seconds...") time.sleep(retry_delay) retry_delay *= 2 else: print("Max retries reached for checking account status.") raise Exception(f"Failed to check account status after {max_retries} attempts. Last error: {last_exception}") from last_exception raise Exception(f"Failed to check account status after {max_retries} attempts (unexpected loop end). Last error: {last_exception}") def create_task(self, apiConfig, nodeInfoList, base_url): """ 创建任务，包含重试机制，最多重试5次 """ safe_base_url = apiConfig.get('base_url') safe_workflow_id = apiConfig.get("workflowId") safe_api_key = apiConfig.get("apiKey") if not safe_base_url or not safe_workflow_id or not safe_api_key: raise ValueError("Missing required apiConfig fields: 'base_url', 'workflowId', 'apiKey'") url = f"{safe_base_url}/task/openapi/create" headers = { "Content-Type": "application/json", "User-Agent": "ComfyUI-RH-APICall-Node/1.0", } data = { "workflowId": safe_workflow_id, "apiKey": safe_api_key, "nodeInfoList": nodeInfoList, } max_retries = 5 retry_delay = 1 last_exception = None for attempt in range(max_retries): response = None current_last_exception = None try: print(f"Create task attempt {attempt + 1}/{max_retries}...") response = requests.post(url, json=data, headers=headers, timeout=30) print(f"Create task attempt {attempt + 1}: Status code {response.status_code}") response.raise_for_status() result = response.json() if result.get("code") == 0: if "data" in result and "taskId" in result["data"] and "netWssUrl" in result["data"]: print("Task created successfully.") return result else: print(f"API success code 0, but response structure invalid: {result}") current_last_exception = ValueError(f"API success code 0, but response structure invalid.") else: api_msg = result.get('msg', 'Unknown API error') print(f"API error creating task (code {result.get('code')}): {api_msg}") current_last_exception = Exception(f"API error (code {result.get('code')}): {api_msg}") except requests.exceptions.Timeout as e: print(f"Create task attempt {attempt + 1} timed out.") current_last_exception = e except requests.exceptions.RequestException as e: print(f"Create task attempt {attempt + 1} network error: {e}") current_last_exception = e except json.JSONDecodeError as e: print(f"Create task attempt {attempt + 1} failed to decode JSON response.") if response is not None: print(f"Raw response text: {response.text}") current_last_exception = e except Exception as e: print(f"Create task attempt {attempt + 1} unexpected error: {e}") current_last_exception = e if current_last_exception is not None: last_exception = current_last_exception if attempt < max_retries - 1: print(f"Retrying task creation in {retry_delay} seconds...") time.sleep(retry_delay) retry_delay *= 2 else: error_message = f"Failed to create task after {max_retries} attempts." if last_exception: error_message += f" Last error: {last_exception}" print(error_message) raise Exception(error_message) from last_exception else: if attempt == max_retries - 1: error_message = f"Failed to create task after {max_retries} attempts (unknown reason)." if last_exception: error_message += f" Last error: {last_exception}" print(error_message) raise Exception(error_message) from last_exception raise Exception("Task creation failed unexpectedly after retry loop.") def check_task_status(self, task_id, api_key, base_url): """ 查询任务状态。 Returns a dictionary representing status or list of outputs on success. """ if not task_id or not api_key or not base_url: raise ValueError("Task ID, API Key, and Base URL are required for checking task status.") url = f"{base_url}/task/openapi/outputs" headers = { "User-Agent": "ComfyUI-RH-APICall-Node/1.0", "Content-Type": "application/json", } data = { "taskId": task_id, "apiKey": api_key } try: response = requests.post(url, json=data, headers=headers, timeout=20) print(f"Check status ({task_id}): Response Status Code: {response.status_code}") try: result = response.json() print(f"Check status ({task_id}): Response JSON: {json.dumps(result, indent=2, ensure_ascii=False)}") except json.JSONDecodeError: print(f"Check status ({task_id}): Failed to decode JSON. Response Text: {response.text}") error_msg = f"HTTP Error {response.status_code} and Invalid JSON" if response.status_code != 200 else "Invalid JSON response" return {"taskStatus": "error", "error": error_msg} api_code = result.get("code") api_msg = result.get("msg", "") api_data = result.get("data") if response.status_code != 200: error_detail = api_msg if api_msg else f"HTTP Error {response.status_code}" print(f"Warning: Non-200 status code ({response.status_code}). API Message: {api_msg}") return {"taskStatus": "error", "error": error_detail} if api_code == 0 and isinstance(api_data, list) and api_data: return api_data elif api_msg == "APIKEY_TASK_IS_RUNNING": return {"taskStatus": "RUNNING"} elif api_msg == "APIKEY_TASK_IS_QUEUED": return {"taskStatus": "QUEUED"} elif api_code != 0: print(f"API Error checking status (code {api_code}): {api_msg}") return {"taskStatus": "error", "error": api_msg} elif api_code == 0 and (api_data is None or (isinstance(api_data, list) and not api_data)): print("Task status check returned code 0 but no data - assuming still running.") return {"taskStatus": "RUNNING"} else: print(f"Unknown task status response: {result}") return {"taskStatus": "unknown", "details": result} except requests.exceptions.Timeout: print(f"Network timeout checking task status for {task_id}") return {"taskStatus": "error", "error": "Network Timeout"} except requests.exceptions.RequestException as e: print(f"Network error checking task status: {e}") return {"taskStatus": "error", "error": f"Network Error: {e}"} except Exception as e: print(f"Unexpected error checking task status: {e}") return {"taskStatus": "error", "error": f"Unexpected Error: {e}"}