Files
doubletwisted-ComfyUI-Deadl…/plugins/ComfyUI/ComfyUI.py
T

1104 lines
47 KiB
Python

from __future__ import absolute_import
from Deadline.Plugins import DeadlinePlugin, PluginType
from System.Diagnostics import ProcessPriorityClass
from Deadline.Scripting import RepositoryUtils, SystemUtils, FileUtils
import os
import re
import sys
import json
import time
import socket
import threading
import urllib.request
import urllib.error
import urllib.parse
import traceback
import random
"""
ComfyUI Deadline Plugin
by Dominik Bargiel dominikbargiel97@gmail.com
A Deadline plugin for rendering ComfyUI workflows. Handles workflow submission,
progress monitoring, seed manipulation, batch processing, and multi-GPU support.
Supports both existing ComfyUI instances and launching new ones.
"""
# Constants
DEFAULT_PORT = 8188
PORT_OFFSET_PER_GPU = 100
MAX_PORT_SEARCH_RANGE = 100
DEFAULT_TIMEOUT = 6000 # 100 minutes
DEFAULT_POLLING_INTERVAL = 10 # seconds
MAX_SEED_VALUE = 2147483647
PROGRESS_LOG_INTERVAL = 10 # Log every 10 polls
FILE_WRITE_DELAY = 2 # seconds to wait for files to be written
# Seed parameter names to search for in workflows
SEED_PARAMETER_NAMES = ["seed", "noise_seed", "value"]
# Output node types that indicate the workflow will produce output
OUTPUT_NODE_TYPES = ["SaveImage", "PreviewImage", "SaveVideo"]
def GetDeadlinePlugin():
return ComfyUI()
def CleanupDeadlinePlugin(deadlinePlugin):
deadlinePlugin.Cleanup()
class ComfyUIError(Exception):
"""Custom exception for ComfyUI plugin errors"""
pass
class ComfyUI(DeadlinePlugin):
def __init__(self):
if sys.version_info.major == 3:
super().__init__()
self._setup_callbacks()
self._setup_stdout_handlers()
self._initialize_member_variables()
def _setup_callbacks(self):
"""Setup all plugin callbacks"""
self.InitializeProcessCallback += self.InitializeProcess
self.RenderExecutableCallback += self.RenderExecutable
self.RenderArgumentCallback += self.RenderArgument
self.PreRenderTasksCallback += self.PreRenderTasks
self.PostRenderTasksCallback += self.PostRenderTasks
def _setup_stdout_handlers(self):
"""Setup stdout handlers for ComfyUI output parsing"""
# Server startup handlers
self.AddStdoutHandlerCallback(".*Starting server.*").HandleCallback += self.HandleServerStarted
self.AddStdoutHandlerCallback(".*Error:.*").HandleCallback += self.HandleStdoutError
self.AddStdoutHandlerCallback(".*Exception:.*").HandleCallback += self.HandleStdoutError
# Progress handlers
self.AddStdoutHandlerCallback(r"\s*([0-9]+)%\|.*\|\s*([0-9]+)/([0-9]+).*").HandleCallback += self.HandleStdoutProgressBar
self.AddStdoutHandlerCallback(r"Progress: ([0-9.]+)%.*").HandleCallback += self.HandleStdoutProgressPercent
self.AddStdoutHandlerCallback(r"Prompt executed in ([0-9.]+) seconds").HandleCallback += self.HandleStdoutPromptExecuted
def _initialize_member_variables(self):
"""Initialize all member variables"""
# Core state variables
self.comfyui_output_dir = ""
self.server_started = False
self.task_completed = False
self.comfyui_process = None
self.comfyui_api_url = None
self.temp_dir = None
self.workflow_submitted = False
self.client_id = None
self.prompt_id = None
self.progress_value = 0
self.thread_running = True
self.custom_output_dir_specified = False
# Batch processing variables
self.chunk_size = 1
self.prompts_executed = 0
self.batch_mode = False
# Prompt tracking variables
self.prompt_ids = []
self.completed_prompts = set()
self.current_tracking_index = 0
def Cleanup(self):
"""Clean up plugin resources"""
self.thread_running = False
# Clean up callbacks
del self.InitializeProcessCallback
del self.RenderExecutableCallback
del self.RenderArgumentCallback
del self.PreRenderTasksCallback
del self.PostRenderTasksCallback
# Clean up stdout handlers
for stdoutHandler in self.StdoutHandlers:
del stdoutHandler.HandleCallback
def InitializeProcess(self):
"""Initialize process settings"""
self.SingleFramesOnly = True
self.PluginType = PluginType.Simple
self.ProcessPriority = ProcessPriorityClass.BelowNormal
self.UseProcessTree = True
self.StdoutHandling = True
self.PopupHandling = False
def _get_cuda_device_arg(self) -> str:
"""
Get CUDA device argument based on plugin configuration and worker settings.
Returns:
str: CUDA device argument string or empty string
"""
assigned_gpu = None
# 1. Check for specific CudaDeviceID from job plugin info
cuda_device_id_plugin_info = self.GetPluginInfoEntryWithDefault("CudaDeviceID", "").strip()
if cuda_device_id_plugin_info:
try:
int(cuda_device_id_plugin_info) # Validate it's an integer string
assigned_gpu = cuda_device_id_plugin_info
self.LogInfo(f"Using specific CUDA device ID from plugin info: {assigned_gpu}")
except ValueError:
self.LogWarning(f"Invalid CudaDeviceID value '{cuda_device_id_plugin_info}' in plugin info. Ignoring.")
# 2. Check Deadline worker GPU affinity if not set by plugin info
if assigned_gpu is None:
assigned_gpu = self._get_gpu_from_worker_affinity()
# 3. Use default device if still no assignment
if assigned_gpu is None:
assigned_gpu = self._get_default_cuda_device()
return f"--cuda-device {assigned_gpu}" if assigned_gpu is not None else ""
def _get_gpu_from_worker_affinity(self) -> str:
"""Get GPU assignment from Deadline worker affinity settings"""
if not self.OverrideGpuAffinity():
self.LogInfo("Deadline worker GPU affinity is not overridden for this worker.")
return None
available_gpus = self.GpuAffinity()
if not available_gpus:
self.LogInfo("Worker GPU affinity is overridden but no specific GPUs are assigned.")
return None
self.LogInfo(f"Worker has GPU affinity set by Deadline: {available_gpus}")
selected_gpu_device_id = available_gpus[self.GetThreadNumber() % len(available_gpus)]
assigned_gpu = str(selected_gpu_device_id)
self.LogInfo(f"Assigning CUDA device based on worker affinity: {assigned_gpu}")
return assigned_gpu
def _get_default_cuda_device(self) -> str:
"""Get default CUDA device if configured"""
use_default_device_zero = self.GetBooleanPluginInfoEntryWithDefault("DefaultCudaDeviceZero", True)
if use_default_device_zero:
self.LogInfo("No CUDA device assigned. Defaulting to CUDA device 0.")
return "0"
else:
self.LogInfo("No CUDA device assigned and DefaultCudaDeviceZero is false. ComfyUI will use default GPU behavior.")
return None
def _setup_batch_processing(self):
"""Setup batch processing configuration"""
self.batch_mode = self.GetBooleanPluginInfoEntryWithDefault("BatchMode", False)
if self.batch_mode:
self.chunk_size = int(self.GetJob().ChunkSize)
self.LogInfo(f"Batch mode enabled. Chunk size: {self.chunk_size}")
else:
self.chunk_size = 1
self.LogInfo("Batch mode disabled. Processing single task.")
self.prompts_executed = 0
def _setup_output_directory(self):
"""Setup output directory configuration"""
job_output_directory_plugin = self.GetPluginInfoEntryWithDefault("JobOutputDirectory", "")
if job_output_directory_plugin:
self._setup_custom_output_directory(job_output_directory_plugin)
else:
self._setup_default_output_directory()
def _setup_custom_output_directory(self, output_dir: str):
"""Setup custom output directory"""
self.comfyui_output_dir = os.path.abspath(output_dir)
self.custom_output_dir_specified = True
self.LogInfo(f"ComfyUI will output directly to user-specified directory: {self.comfyui_output_dir}")
if not os.path.exists(self.comfyui_output_dir):
self._create_directory(self.comfyui_output_dir, "user-specified output")
def _setup_default_output_directory(self):
"""Setup default ComfyUI output directory"""
# Check for configured default output directory first
default_output_dir = self.GetConfigEntryWithDefault("DefaultOutputDirectory", "")
if default_output_dir:
self.comfyui_output_dir = os.path.abspath(default_output_dir)
self.custom_output_dir_specified = False
self.LogInfo(f"Using configured default output directory: {self.comfyui_output_dir}")
else:
# Fall back to ComfyUI's standard output directory
comfyui_path = self.GetConfigEntry("ComfyUIPath")
self.comfyui_output_dir = os.path.join(comfyui_path, "ComfyUI", "output")
self.custom_output_dir_specified = False
self.LogInfo(f"Using ComfyUI's default output directory: {self.comfyui_output_dir}")
if not os.path.exists(self.comfyui_output_dir):
self._create_directory(self.comfyui_output_dir, "default output")
def _create_directory(self, directory_path: str, description: str):
"""Create a directory with error handling"""
try:
os.makedirs(directory_path)
self.LogInfo(f"Created {description} directory: {directory_path}")
except Exception as e:
self.LogWarning(f"Could not create {description} directory {directory_path}: {e}")
def _calculate_comfyui_port(self) -> str:
"""Calculate the port for ComfyUI based on CUDA device"""
cuda_arg = self._get_cuda_device_arg()
cuda_device_id = None
if cuda_arg:
cuda_match = re.search(r'--cuda-device\s+(\d+)', cuda_arg)
if cuda_match:
cuda_device_id = int(cuda_match.group(1))
self.LogInfo(f"Using CUDA device ID {cuda_device_id} for port calculation")
default_port = int(self.GetPluginInfoEntryWithDefault("ComfyUIPort", str(DEFAULT_PORT)))
if cuda_device_id is not None:
base_port = default_port + (cuda_device_id * PORT_OFFSET_PER_GPU)
self.LogInfo(f"Calculated base port {base_port} for CUDA device {cuda_device_id}")
else:
base_port = default_port
self.LogInfo(f"No CUDA device ID available, using default port {base_port}")
return self._determine_final_port(base_port)
def _determine_final_port(self, base_port: int) -> str:
"""Determine final port to use, checking for existing instances"""
self.LogInfo(f"Checking if ComfyUI is already running on port {base_port}")
if self._is_port_in_use(base_port):
self.LogInfo(f"ComfyUI is already running on port {base_port}, will use existing instance")
self.use_existing_comfyui = True
self.comfyui_port = str(base_port)
self.comfyui_api_url = f"http://127.0.0.1:{self.comfyui_port}"
self.server_started = True
else:
self.LogInfo(f"No ComfyUI instance detected on port {base_port}")
self.use_existing_comfyui = False
self.comfyui_port = self._find_available_port(base_port)
self.LogInfo(f"Will use port {self.comfyui_port} for ComfyUI")
self.comfyui_api_url = f"http://127.0.0.1:{self.comfyui_port}"
return self.comfyui_port
def PreRenderTasks(self):
"""Setup tasks before rendering"""
self.LogInfo("ComfyUI PreRenderTasks started.")
try:
self._setup_batch_processing()
self._setup_output_directory()
self._setup_temp_directory()
self._calculate_comfyui_port()
self.task_completed = False
self.LogInfo("PreRenderTasks completed successfully.")
except Exception as e:
self.LogWarning(f"Error in PreRenderTasks: {e}")
raise ComfyUIError(f"PreRenderTasks failed: {str(e)}")
def _setup_temp_directory(self):
"""Create temporary directory for job files"""
self.temp_dir = self.CreateTempDirectory("comfyui_job")
def _is_port_in_use(self, port: int) -> bool:
"""Check if the given port is already in use"""
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.settimeout(1)
result = s.connect_ex(('127.0.0.1', port))
return result == 0
except:
return False
def _find_available_port(self, start_port: int) -> str:
"""Find an available port starting from start_port"""
port = start_port
max_port = start_port + MAX_PORT_SEARCH_RANGE
while port < max_port:
if not self._is_port_in_use(port):
return str(port)
port += 1
return str(start_port)
def PostRenderTasks(self):
"""Cleanup tasks after rendering"""
self.LogInfo("ComfyUI PostRenderTasks started.")
# Wait for files to be written
self.LogInfo("Waiting for files to be written to disk...")
time.sleep(FILE_WRITE_DELAY)
self._log_output_directory_status()
self.LogInfo("PostRenderTasks finished.")
def _log_output_directory_status(self):
"""Log the status of output directories"""
if self.custom_output_dir_specified:
self._log_custom_directory_status()
else:
self._log_default_directory_status()
def _log_custom_directory_status(self):
"""Log status of custom output directory"""
self.LogInfo(f"ComfyUI was instructed to output to: {self.comfyui_output_dir}")
if os.path.exists(self.comfyui_output_dir):
try:
final_outputs = os.listdir(self.comfyui_output_dir)
self.LogInfo(f"Output directory contains {len(final_outputs)} item(s). Examples: {final_outputs[:5]}")
except Exception as e:
self.LogWarning(f"Could not list contents of output directory: {e}")
else:
self.LogWarning(f"Output directory was not found: {self.comfyui_output_dir}")
def _log_default_directory_status(self):
"""Log status of default output directory"""
self.LogInfo(f"ComfyUI used default output directory: {self.comfyui_output_dir}")
if os.path.exists(self.comfyui_output_dir):
try:
default_outputs = os.listdir(self.comfyui_output_dir)
self.LogInfo(f"Default output directory contains {len(default_outputs)} item(s). Examples: {default_outputs[:5]}")
except Exception as e:
self.LogWarning(f"Could not list contents of default output directory: {e}")
else:
self.LogWarning(f"Default output directory was not found: {self.comfyui_output_dir}")
def RenderExecutable(self):
"""Get the Python executable for ComfyUI"""
comfyui_path = self.GetConfigEntry("ComfyUIPath")
python_exe = os.path.join(comfyui_path, "python_embeded", "python.exe")
if os.path.exists(python_exe):
self.LogInfo(f"Using ComfyUI embedded Python: {python_exe}")
return python_exe
else:
error_msg = f"ComfyUI embedded Python not found at: {python_exe}. Please check your ComfyUIPath configuration."
self.LogWarning(error_msg)
self.FailRender(error_msg)
return ""
def RenderArgument(self):
"""Build command line arguments for ComfyUI"""
comfyui_path = self.GetConfigEntry("ComfyUIPath")
comfyui_main_py = os.path.join(comfyui_path, "ComfyUI", "main.py")
# Validate that main.py exists
if not os.path.exists(comfyui_main_py):
error_msg = f"ComfyUI main.py not found at: {comfyui_main_py}. Please check your ComfyUIPath configuration."
self.LogWarning(error_msg)
self.FailRender(error_msg)
return ""
# If using existing ComfyUI instance, return dummy command
if self.use_existing_comfyui:
return self._create_dummy_command()
# Build arguments for new ComfyUI instance
return self._build_comfyui_arguments(comfyui_main_py)
def _create_dummy_command(self) -> str:
"""Create dummy command for existing ComfyUI instances"""
self.LogInfo("Using existing ComfyUI instance - returning dummy command.")
# Start workflow submission thread
workflow_thread = threading.Thread(target=self.submit_workflow)
workflow_thread.daemon = True
workflow_thread.start()
dummy_script = self._get_dummy_script()
return f'-c "{dummy_script}"'
def _get_dummy_script(self) -> str:
"""Get dummy Python script for workflow waiting"""
return """
import time, os, sys
print('Waiting for ComfyUI workflow to complete via API...')
max_wait = 600 # 10 minute timeout
for i in range(max_wait):
if i % 10 == 0:
progress_val = float(os.environ.get('DEADLINE_TASK_PROGRESS', i * (100.0 / max_wait)))
print(f'Progress: {progress_val:.1f}%')
time.sleep(1)
print('Dummy command timeout reached or task completed.')
"""
def _build_comfyui_arguments(self, comfyui_main_py: str) -> str:
"""Build command line arguments for new ComfyUI instance"""
port = getattr(self, "comfyui_port", str(DEFAULT_PORT))
args_list = []
# Add Python flags
if self.GetBooleanPluginInfoEntryWithDefault("PythonNoUserSite", True):
args_list.append("-s")
# Add main script and port
args_list.append(f'"{comfyui_main_py}"')
args_list.append(f"--port {port}")
# Add CUDA device argument
cuda_arg = self._get_cuda_device_arg()
if cuda_arg:
args_list.append(cuda_arg)
# Add output directory if custom one was specified
if self.custom_output_dir_specified and self.comfyui_output_dir:
args_list.append(f'--output-directory "{self.comfyui_output_dir}"')
self.LogInfo(f"Passing --output-directory \"{self.comfyui_output_dir}\" to ComfyUI.")
else:
self.LogInfo("Not passing --output-directory to ComfyUI, it will use its default.")
args = " ".join(args_list)
self.LogInfo(f"Render Arguments: {args}")
return args
def HandleServerStarted(self):
"""Called when the ComfyUI server has started"""
self.LogInfo("ComfyUI server has started")
self.server_started = True
# Start workflow submission if not using existing instance
if not self.use_existing_comfyui and not self.workflow_submitted:
workflow_thread = threading.Thread(target=self.submit_workflow)
workflow_thread.daemon = True
workflow_thread.start()
def http_request(self, url: str, method: str = "GET", data=None, headers=None, verbose: bool = True) -> dict:
"""Make an HTTP request to the ComfyUI API"""
if not self.thread_running:
self.LogInfo("Thread stopping due to task completion")
return {'status_code': 0, 'text': '', 'json': lambda: {}}
if verbose:
self.LogInfo(f"Making {method} request to {url}")
if headers is None:
headers = {}
if data is not None and not isinstance(data, bytes):
data = json.dumps(data).encode('utf-8')
headers['Content-Type'] = 'application/json'
req = urllib.request.Request(url, data=data, headers=headers, method=method)
try:
with urllib.request.urlopen(req) as response:
response_data = response.read().decode('utf-8')
return {
'status_code': response.status,
'text': response_data,
'json': lambda: json.loads(response_data) if response_data else {}
}
except urllib.error.HTTPError as e:
self.LogWarning(f"HTTP Error: {e.code} {e.reason}")
return {
'status_code': e.code,
'text': e.read().decode('utf-8'),
'json': lambda: {}
}
except Exception as e:
self.LogWarning(f"Error in HTTP request: {str(e)}")
raise
def modify_workflow_seeds(self, workflow_data: dict, task_id: int) -> bool:
"""
Modify seeds in the workflow based on task ID and seed mode.
Args:
workflow_data: The workflow data
task_id: Current task ID (or frame number)
Returns:
bool: True if seeds were modified, False otherwise
"""
seed_mode = self.GetPluginInfoEntryWithDefault("SeedMode", "fixed")
if seed_mode == "fixed":
self.LogInfo(f"No seed manipulation: SeedMode is 'fixed' (keep)")
return False
self.LogInfo(f"Seed Control: Using '{seed_mode}' mode for task ID {task_id}")
seeds_modified = False
for node_id, node in workflow_data.items():
if not isinstance(node, dict) or "inputs" not in node:
continue
seeds_modified |= self._modify_node_seeds(node, node_id, task_id, seed_mode)
if not seeds_modified:
self.LogInfo(f"Seed Control: No seed parameters found in workflow that could be modified")
return seeds_modified
def _modify_node_seeds(self, node: dict, node_id: str, task_id: int, seed_mode: str) -> bool:
"""Modify seeds in a single node"""
inputs = node.get("inputs", {})
if not inputs:
return False
node_type = node.get("class_type", "unknown")
seeds_modified = False
for param_name in SEED_PARAMETER_NAMES:
if param_name in inputs:
try:
original_seed = int(inputs[param_name])
new_seed = self._calculate_new_seed(original_seed, task_id, seed_mode, inputs)
if new_seed != original_seed:
node["inputs"][param_name] = new_seed
self.LogInfo(f"Seed Control: Modified node {node_id} ({node_type}) seed from {original_seed} to {new_seed}")
seeds_modified = True
else:
self.LogInfo(f"Seed Control: Node {node_id} ({node_type}) seed kept at {original_seed}")
except (ValueError, TypeError):
self.LogInfo(f"Seed Control: Node {node_id} ({node_type}) has non-numeric {param_name}: {inputs[param_name]}")
return seeds_modified
def _calculate_new_seed(self, original_seed: int, task_id: int, seed_mode: str, inputs: dict) -> int:
"""Calculate new seed based on mode and parameters"""
if seed_mode == "auto":
control_mode = inputs.get("control_after_generate", "increment")
if control_mode == "fixed":
return original_seed
elif control_mode == "increment":
return original_seed + task_id
elif control_mode == "decrement":
return max(0, original_seed - task_id)
elif control_mode == "randomize":
return random.randint(0, MAX_SEED_VALUE)
else:
return original_seed + task_id
elif seed_mode == "change":
return random.randint(0, MAX_SEED_VALUE)
else:
return random.randint(0, MAX_SEED_VALUE)
def load_and_validate_workflow(self) -> dict:
"""Load workflow file and validate its structure"""
workflow_file = self._get_workflow_file_path()
if not workflow_file or not os.path.exists(workflow_file):
self.LogWarning(f"Workflow file does not exist at: {workflow_file}")
self.FailRender(f"Workflow file not found: {workflow_file}")
return None
try:
workflow_data = self._load_workflow_from_file(workflow_file)
workflow_data = self.validate_workflow(workflow_data)
# Apply seed manipulation
task_id = self.GetCurrentTaskId()
seeds_modified = self.modify_workflow_seeds(workflow_data, task_id)
if seeds_modified:
self.LogInfo(f"Applied seed manipulation for task ID {task_id}")
else:
self.LogInfo(f"No seed manipulation applied for task ID {task_id}")
return workflow_data
except Exception as e:
self.LogWarning(f"Error loading or validating workflow file '{workflow_file}': {e}")
self.FailRender(f"Error loading or validating workflow file: {str(e)}")
return None
def _get_workflow_file_path(self) -> str:
"""Get the workflow file path from plugin settings"""
workflow_file = self.GetPluginInfoEntryWithDefault("ComfyWorkflowFile", self.GetDataFilename())
workflow_file = RepositoryUtils.CheckPathMapping(workflow_file)
self.LogInfo(f"Workflow file setting from plugin info: '{workflow_file}'")
return workflow_file
def _load_workflow_from_file(self, workflow_file: str) -> dict:
"""Load workflow data from JSON file"""
with open(workflow_file, 'r') as f:
workflow_content = f.read()
workflow_data = json.loads(workflow_content)
self.LogInfo(f"Successfully loaded workflow file: {workflow_file}")
return workflow_data
def validate_workflow(self, workflow_data: dict) -> dict:
"""Check workflow for output nodes and convert to UI format if needed"""
has_save_image, has_output_node = self._check_workflow_output_nodes(workflow_data)
if not has_output_node:
self.LogWarning("No output nodes found in workflow. You need at least one SaveImage, PreviewImage, or SaveVideo node.")
if not has_save_image:
self.LogWarning("No SaveImage node found in workflow. Images may not be saved to disk.")
# Convert from API format to UI format if needed
if "nodes" in workflow_data:
self.LogInfo("Converting from API format to UI format...")
workflow_data = self._convert_api_to_ui_format(workflow_data)
return workflow_data
def _check_workflow_output_nodes(self, workflow_data: dict) -> tuple:
"""Check for output nodes in workflow"""
has_save_image = False
has_output_node = False
# Check UI format (numbered keys)
for key, node in workflow_data.items():
if isinstance(node, dict):
class_type = node.get("class_type", "")
if class_type == "SaveImage":
has_save_image = True
has_output_node = True
self.LogInfo(f"Found SaveImage node in workflow")
elif class_type in OUTPUT_NODE_TYPES:
has_output_node = True
self.LogInfo(f"Found output node {class_type} in workflow")
# Check API format (nodes array)
if not has_output_node and "nodes" in workflow_data:
for node in workflow_data["nodes"]:
if isinstance(node, dict):
class_type = node.get("class_type", "")
if class_type == "SaveImage":
has_save_image = True
has_output_node = True
self.LogInfo(f"Found SaveImage node in workflow")
elif class_type in OUTPUT_NODE_TYPES:
has_output_node = True
self.LogInfo(f"Found output node {class_type} in workflow")
return has_save_image, has_output_node
def _convert_api_to_ui_format(self, workflow_data: dict) -> dict:
"""Convert workflow from API format to UI format"""
converted_workflow = {}
for node in workflow_data["nodes"]:
node_id = str(node.get("id", 0))
converted_workflow[node_id] = node
return converted_workflow
def HandleStdoutProgressBar(self):
"""Handle progress in the format ' 4%|4 | 1/25 [00:02<00:59, 2.50s/it]'"""
try:
percent = float(self.GetRegexMatch(1))
current_step = int(self.GetRegexMatch(2))
total_steps = int(self.GetRegexMatch(3))
# Calculate overall chunk progress if needed
if self.chunk_size > 1:
completed_progress = (self.prompts_executed / self.chunk_size) * 100
current_contribution = (percent / self.chunk_size)
overall_progress = min(99, completed_progress + current_contribution) if self.prompts_executed < self.chunk_size else 100
self.SetProgress(overall_progress)
self.progress_value = overall_progress
self.SetStatusMessage(f"Chunk {self.prompts_executed + 1}/{self.chunk_size} - Step {current_step}/{total_steps} ({overall_progress:.2f}%)")
self.LogInfo(f"Chunk Progress: {overall_progress:.2f}% (Prompt {self.prompts_executed + 1}/{self.chunk_size}, Step {current_step}/{total_steps})")
else:
self.SetProgress(percent)
self.progress_value = percent
self.SetStatusMessage(f"Step {current_step}/{total_steps} ({percent:.2f}%)")
self.LogInfo(f"Progress: {percent}% ({current_step}/{total_steps})")
except ValueError:
self.LogWarning(f"Could not parse progress from: {self.GetRegexMatch(0)}")
def HandleStdoutProgressPercent(self):
"""Handle progress in the format 'Progress: 45.5%'"""
try:
percent = float(self.GetRegexMatch(1))
# Calculate overall chunk progress if needed
if self.chunk_size > 1:
completed_progress = (self.prompts_executed / self.chunk_size) * 100
current_contribution = (percent / self.chunk_size)
overall_progress = min(99, completed_progress + current_contribution) if self.prompts_executed < self.chunk_size else 100
self.SetProgress(overall_progress)
self.progress_value = overall_progress
self.SetStatusMessage(f"Chunk {self.prompts_executed + 1}/{self.chunk_size} - Rendering: {overall_progress:.2f}%")
self.LogInfo(f"Chunk Progress: {overall_progress:.2f}% (Prompt {self.prompts_executed + 1}/{self.chunk_size} at {percent:.1f}%)")
else:
self.SetProgress(percent)
self.progress_value = percent
self.SetStatusMessage(f"Rendering: {percent:.2f}%")
self.LogInfo(f"Progress: {percent}%")
except ValueError:
self.LogWarning(f"Could not parse progress from: {self.GetRegexMatch(0)}")
def HandleStdoutPromptExecuted(self):
"""Handle completion message 'Prompt executed in X seconds'"""
execution_time = self.GetRegexMatch(1)
self.LogInfo(f"Workflow completed in {execution_time} seconds")
# Check if we already counted this prompt
if self.prompt_id and self.prompt_id in self.completed_prompts:
self.LogInfo(f"Prompt {self.prompt_id} already counted")
else:
self._handle_stdout_prompt_completion()
def _handle_stdout_prompt_completion(self):
"""Handle prompt completion detected from stdout"""
if self.prompt_id:
self.completed_prompts.add(self.prompt_id)
self.prompts_executed += 1
self.LogInfo(f"Prompt execution {self.prompts_executed} of {self.chunk_size} completed")
# Move to next prompt
if self.prompt_id:
self._move_to_next_prompt()
# Check if all prompts completed
if self.prompts_executed >= self.chunk_size:
self._complete_task()
time.sleep(FILE_WRITE_DELAY) # Wait for files to be written
else:
self._update_progress()
self.LogInfo(f"Waiting for remaining prompts. {self.prompts_executed} of {self.chunk_size} completed")
def HandleStdoutError(self):
"""Handle errors from ComfyUI"""
error_msg = self.GetRegexMatch(0)
self.LogWarning(f"ComfyUI error: {error_msg}")
if not self.task_completed:
self.FailRender(f"ComfyUI error: {error_msg}")
def initialize_api_connection(self) -> bool:
"""Initialize connection to ComfyUI API and get client ID"""
try:
time.sleep(2) # Wait for server to be fully initialized
response = self.http_request(f"{self.comfyui_api_url}/prompt")
if response['status_code'] != 200:
self.LogWarning(f"Error connecting to ComfyUI API: {response['status_code']}")
self.FailRender(f"Error connecting to ComfyUI API: {response['status_code']}")
return False
self.client_id = response['json']().get('client_id', '')
self.LogInfo(f"Got client ID: {self.client_id}")
return True
except Exception as e:
self.LogWarning(f"Error initializing API connection: {e}")
self.FailRender(f"Error initializing API connection: {str(e)}")
return False
def queue_workflow(self, workflow_data: dict) -> bool:
"""Submit workflow to ComfyUI queue"""
try:
self._reset_prompt_tracking()
# Queue initial prompt
if not self._queue_single_prompt(workflow_data):
return False
# Queue additional prompts for batch mode
if self.batch_mode and self.chunk_size > 1:
self._queue_batch_prompts(workflow_data)
self.LogInfo(f"Queued total of {len(self.prompt_ids)} prompts: {self.prompt_ids}")
return True
except Exception as e:
self.LogWarning(f"Error queuing workflow: {e}")
self.FailRender(f"Error queuing workflow: {str(e)}")
return False
def _reset_prompt_tracking(self):
"""Reset prompt tracking variables"""
self.prompt_ids = []
self.completed_prompts = set()
self.current_tracking_index = 0
def _queue_single_prompt(self, workflow_data: dict) -> bool:
"""Queue a single prompt to ComfyUI"""
data = {"prompt": workflow_data, "client_id": self.client_id}
response = self.http_request(f"{self.comfyui_api_url}/prompt", method="POST", data=data)
if response['status_code'] != 200:
self.LogWarning(f"Error queuing prompt: {response['text']}")
self.FailRender(f"Error queuing prompt: {response['text']}")
return False
self.prompt_id = response['json']()['prompt_id']
self.prompt_ids.append(self.prompt_id)
self.LogInfo(f"Queued prompt with ID: {self.prompt_id}")
self.workflow_submitted = True
return True
def _queue_batch_prompts(self, workflow_data: dict):
"""Queue additional prompts for batch processing"""
self.LogInfo(f"Batch mode with chunk size {self.chunk_size}. Queueing additional prompts...")
import copy
for i in range(1, self.chunk_size):
prompt_workflow = copy.deepcopy(workflow_data)
# Modify seeds if needed
if self.GetPluginInfoEntryWithDefault("SeedMode", "fixed") != "fixed":
self.modify_workflow_seeds(prompt_workflow, i)
self.LogInfo(f"Modified seeds for additional prompt {i}")
# Queue the workflow
data = {"prompt": prompt_workflow, "client_id": self.client_id}
response = self.http_request(f"{self.comfyui_api_url}/prompt", method="POST", data=data)
if response['status_code'] != 200:
self.LogWarning(f"Error queuing additional prompt {i}: {response['text']}")
break
prompt_id = response['json']()['prompt_id']
self.prompt_ids.append(prompt_id)
self.LogInfo(f"Queued additional prompt {i} with ID: {prompt_id}")
time.sleep(0.5) # Small delay between submissions
def process_history_data(self, history_data: dict) -> bool:
"""Process history data and update task status"""
if self.prompt_id not in history_data:
return False
# Check for outputs indicating completion
if 'outputs' in history_data[self.prompt_id]:
return self._handle_prompt_completion(history_data[self.prompt_id]['outputs'])
# Check for error status
elif 'status' in history_data[self.prompt_id]:
return self._handle_prompt_status(history_data[self.prompt_id]['status'])
return False
def _handle_prompt_completion(self, outputs: dict) -> bool:
"""Handle completed prompt outputs"""
if not outputs:
return False
self.LogInfo(f"Workflow complete: Found outputs in history for prompt {self.prompt_id}")
# Mark prompt as completed
if self.prompt_id not in self.completed_prompts:
self.completed_prompts.add(self.prompt_id)
self.prompts_executed += 1
self.LogInfo(f"Prompt {self.prompt_id} execution {self.prompts_executed} of {self.chunk_size} completed")
# Log output information
self._log_output_information(outputs)
# Check if all prompts completed
if self.prompts_executed >= self.chunk_size:
self._complete_task()
return True
else:
self._move_to_next_prompt()
self._update_progress()
return False
def _log_output_information(self, outputs: dict):
"""Log information about generated outputs"""
output_nodes = []
for node_id, node_outputs in outputs.items():
if 'images' in node_outputs:
output_nodes.append(node_id)
for img in node_outputs['images']:
self.LogInfo(f"Generated image: {img['filename']}")
if output_nodes:
self.LogInfo(f"Output producing nodes: {output_nodes}")
def _complete_task(self):
"""Mark task as complete"""
self.SetProgress(100)
self.SetStatusMessage("Finished Render")
self.task_completed = True
self.signal_task_completion()
self.LogInfo(f"All {self.chunk_size} prompts in chunk completed, task marked as complete")
def _move_to_next_prompt(self):
"""Move to tracking the next prompt"""
self.current_tracking_index += 1
if self.current_tracking_index < len(self.prompt_ids):
self.prompt_id = self.prompt_ids[self.current_tracking_index]
self.LogInfo(f"Moving to track next prompt: {self.prompt_id}")
else:
self.prompt_id = None
self.LogInfo(f"No more prompts to track. Waiting for {self.chunk_size - self.prompts_executed} more executions.")
def _update_progress(self):
"""Update progress based on completed prompts"""
progress_percent = (self.prompts_executed / self.chunk_size) * 100
self.SetProgress(progress_percent)
self.SetStatusMessage(f"Completed {self.prompts_executed} of {self.chunk_size} prompts ({progress_percent:.1f}%)")
def _handle_prompt_status(self, status: dict) -> bool:
"""Handle prompt status information"""
if status.get('status') == 'error':
return self._handle_prompt_error(status)
# Update progress from execution status
if 'exec_info' in status and 'progress' in status['exec_info']:
self._update_execution_progress(status['exec_info']['progress'])
return False
def _handle_prompt_error(self, status: dict) -> bool:
"""Handle prompt execution errors"""
error_msg = status.get('error', 'Unknown error')
self.LogWarning(f"ComfyUI reported error for prompt {self.prompt_id}: {error_msg}")
if self.chunk_size > 1:
# Continue with remaining prompts in batch
self.LogWarning(f"Continuing with remaining prompts in chunk")
self.completed_prompts.add(self.prompt_id)
self._move_to_next_prompt()
return False
else:
# Single prompt mode, fail the task
self.FailRender(f"ComfyUI workflow failed: {error_msg}")
return True
def _update_execution_progress(self, progress: float):
"""Update progress from execution information"""
current_prompt_progress = float(progress) * 100
# Calculate overall chunk progress if needed
if self.chunk_size > 1:
completed_progress = (self.prompts_executed / self.chunk_size) * 100
current_contribution = (current_prompt_progress / self.chunk_size)
overall_progress = min(99, completed_progress + current_contribution) if self.prompts_executed < self.chunk_size else 100
self.SetProgress(overall_progress)
self.progress_value = overall_progress
else:
self.SetProgress(current_prompt_progress)
self.progress_value = current_prompt_progress
def signal_task_completion(self):
"""Signal to Deadline that the task is complete"""
try:
job = self.GetJob()
task_id = self.GetCurrentTaskId()
slave_name = self.GetSlaveName()
self.LogInfo(f"Signaling Deadline that task {task_id} for job {job.JobId} is complete")
tasks = RepositoryUtils.GetJobTasks(job, True)
current_task = self._find_current_task(tasks, task_id)
if current_task:
self.LogInfo(f"Completing task {current_task.TaskID}")
RepositoryUtils.CompleteTasks(job, [current_task], slave_name)
else:
self.LogWarning(f"Could not find task with ID {task_id}")
except Exception as e:
self.LogWarning(f"Error signaling task completion: {e}")
self.LogWarning(traceback.format_exc())
def _find_current_task(self, tasks, task_id):
"""Find the current task in the task list"""
for task in tasks:
if str(task.TaskID) == str(task_id):
return task
return None
def monitor_workflow_execution(self) -> bool:
"""Poll history endpoint and wait for workflow completion"""
start_time = time.time()
self.LogInfo(f"Beginning to monitor workflow execution for chunk size {self.chunk_size}")
self.LogInfo(f"Monitoring prompts in this order: {self.prompt_ids}")
if self.prompt_ids:
self.prompt_id = self.prompt_ids[0]
poll_count = 0
while time.time() - start_time < DEFAULT_TIMEOUT and self.thread_running:
if self.task_completed:
self.LogInfo("Task already marked as complete by stdout handler")
self.signal_task_completion()
return True
if self.prompt_id:
if self._poll_prompt_status(poll_count):
break
else:
self._check_for_missed_prompts()
poll_count += 1
time.sleep(DEFAULT_POLLING_INTERVAL)
# Check for timeout
if not self.task_completed and time.time() - start_time >= DEFAULT_TIMEOUT:
self.LogWarning(f"Timeout waiting for workflow to complete after {DEFAULT_TIMEOUT} seconds")
self.FailRender(f"Timeout waiting for workflow to complete")
return False
if self.task_completed:
self.signal_task_completion()
return self.task_completed
def _poll_prompt_status(self, poll_count: int) -> bool:
"""Poll the status of the current prompt"""
try:
verbose_log = (poll_count == 0) or (poll_count % PROGRESS_LOG_INTERVAL == 0)
history_response = self.http_request(f"{self.comfyui_api_url}/history/{self.prompt_id}", verbose=verbose_log)
if history_response['status_code'] == 200:
history_data = history_response['json']()
if self.process_history_data(history_data):
return True # Task completed
if verbose_log and self.progress_value > 0:
self.SetStatusMessage(f"Executing: {self.progress_value:.1f}%")
elif history_response['status_code'] == 404:
if verbose_log:
self.LogInfo(f"History entry not found yet for prompt {self.prompt_id}")
else:
if verbose_log:
self.LogWarning(f"Unexpected response from history endpoint: {history_response['status_code']}")
except Exception as e:
self.LogWarning(f"Error checking history endpoint: {e}")
return False
def _check_for_missed_prompts(self):
"""Check for any completed prompts that weren't tracked"""
if self.prompts_executed >= self.chunk_size:
return
try:
history_response = self.http_request(f"{self.comfyui_api_url}/history", verbose=False)
if history_response['status_code'] == 200:
all_history = history_response['json']()
for i, prompt_id in enumerate(self.prompt_ids):
if prompt_id in self.completed_prompts:
continue
if prompt_id in all_history and 'outputs' in all_history[prompt_id]:
self.LogInfo(f"Found completed prompt {prompt_id} that wasn't tracked")
self.prompt_id = prompt_id
self.current_tracking_index = i
break
except Exception as e:
self.LogWarning(f"Error looking for completed prompts: {e}")
def submit_workflow(self):
"""Submit the workflow to ComfyUI's API and wait for completion"""
try:
workflow_data = self.load_and_validate_workflow()
if not workflow_data:
return
if not self.initialize_api_connection():
return
if not self.queue_workflow(workflow_data):
return
self.monitor_workflow_execution()
except Exception as e:
self.LogWarning(f"Error during workflow submission: {e}")
traceback.print_exc()
self.FailRender(f"Error during workflow submission: {str(e)}")