786 lines
31 KiB
Python
786 lines
31 KiB
Python
"""
|
|
ComfyUI Deadline submission nodes.
|
|
|
|
This module owns the ComfyUI-side submission flow. It packages the current API
|
|
prompt, stages referenced input assets, and submits a render job to Deadline.
|
|
"""
|
|
|
|
import copy
|
|
import filecmp
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import uuid
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
|
|
DEADLINE_COMMAND_PATHS = {
|
|
"windows": "C:\\Program Files\\Thinkbox\\Deadline10\\bin\\deadlinecommand.exe",
|
|
"linux": "/opt/Thinkbox/Deadline10/bin/deadlinecommand",
|
|
}
|
|
|
|
DEADLINE_SUBMIT_NODE_TYPES = {"DeadlineSubmit", "SaveAndSubmitNode"}
|
|
OUTPUT_NODE_TYPES = {"SaveImage", "PreviewImage", "SaveVideo", "VHS_VideoCombine"}
|
|
INPUT_LOADER_FIELDS = {
|
|
"LoadImage": ("image",),
|
|
"LoadImageMask": ("image",),
|
|
"LoadAudio": ("audio",),
|
|
"LoadVideo": ("file",),
|
|
# VideoHelperSuite file-based loaders. Directory loaders such as
|
|
# VHS_LoadImagesPath intentionally remain unsupported here.
|
|
"VHS_LoadVideo": ("video",),
|
|
"VHS_LoadVideoFFmpeg": ("video",),
|
|
"VHS_LoadAudioUpload": ("audio",),
|
|
}
|
|
ABSOLUTE_FILE_LOADER_FIELDS = {
|
|
"VHS_LoadVideoPath": ("video",),
|
|
"VHS_LoadVideoFFmpegPath": ("video",),
|
|
"VHS_LoadImagePath": ("image",),
|
|
"VHS_LoadAudio": ("audio_file",),
|
|
}
|
|
MEDIA_EXTENSIONS = {
|
|
".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff", ".exr",
|
|
".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v",
|
|
".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aac",
|
|
}
|
|
|
|
# A worker that is allowed to reuse a long-running GUI endpoint must prove that
|
|
# the endpoint belongs to this installation. A random value created at module
|
|
# import also distinguishes a restarted ComfyUI session from the old process.
|
|
DEADLINE_ENDPOINT_SESSION_ID = uuid.uuid4().hex
|
|
|
|
|
|
def _coerce_bool(value: Any) -> bool:
|
|
if isinstance(value, str):
|
|
return value.strip().lower() in {"1", "true", "yes", "on"}
|
|
return bool(value)
|
|
|
|
|
|
class NodeDefaults:
|
|
JOB_NAME = "ComfyUI via Deadline"
|
|
PRIORITY = 50
|
|
POOL = "none"
|
|
GROUP = "none"
|
|
BATCH_COUNT = 1
|
|
CHUNK_SIZE = 1
|
|
MAX_BATCH_COUNT = 10000
|
|
MAX_CHUNK_SIZE = 256
|
|
MAX_PRIORITY = 100
|
|
|
|
|
|
class DeadlineCommandHelper:
|
|
@staticmethod
|
|
def get_deadline_command() -> str:
|
|
deadline_bin = os.environ.get("DEADLINE_PATH", "")
|
|
|
|
if not deadline_bin and os.path.exists("/Users/Shared/Thinkbox/DEADLINE_PATH"):
|
|
try:
|
|
with open("/Users/Shared/Thinkbox/DEADLINE_PATH", "r", encoding="utf-8") as handle:
|
|
deadline_bin = handle.read().strip()
|
|
except Exception:
|
|
deadline_bin = ""
|
|
|
|
candidates = []
|
|
if deadline_bin:
|
|
candidates.append(os.path.join(deadline_bin, "deadlinecommand.exe" if os.name == "nt" else "deadlinecommand"))
|
|
candidates.append(DEADLINE_COMMAND_PATHS["windows"] if sys.platform.startswith("win") else DEADLINE_COMMAND_PATHS["linux"])
|
|
|
|
for candidate in candidates:
|
|
if os.path.exists(candidate):
|
|
return candidate
|
|
return ""
|
|
|
|
@staticmethod
|
|
def call_deadline_command(arguments: List[str], hide_window: bool = True) -> str:
|
|
deadline_command = DeadlineCommandHelper.get_deadline_command()
|
|
if not deadline_command:
|
|
raise RuntimeError("Deadline command not found. Set DEADLINE_PATH or install Deadline Client.")
|
|
|
|
startupinfo = None
|
|
creationflags = 0
|
|
if os.name == "nt" and hide_window:
|
|
try:
|
|
startupinfo = subprocess.STARTUPINFO()
|
|
startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
|
|
except Exception:
|
|
startupinfo = None
|
|
elif os.name == "nt":
|
|
creationflags = 0x08000000
|
|
|
|
process = subprocess.Popen(
|
|
[deadline_command] + arguments,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
startupinfo=startupinfo,
|
|
creationflags=creationflags,
|
|
)
|
|
stdout, stderr = process.communicate()
|
|
output = stdout.decode(errors="replace") if isinstance(stdout, bytes) else str(stdout)
|
|
errors = stderr.decode(errors="replace") if isinstance(stderr, bytes) else str(stderr)
|
|
|
|
if process.returncode != 0:
|
|
raise RuntimeError(f"deadlinecommand failed with code {process.returncode}: {errors or output}")
|
|
return output
|
|
|
|
@staticmethod
|
|
def get_job_id_from_submission(submission_results: str) -> str:
|
|
for token in submission_results.replace("\r", "\n").split():
|
|
if token.startswith("JobID="):
|
|
return token.split("=", 1)[1].strip()
|
|
return ""
|
|
|
|
|
|
class WorkflowProcessor:
|
|
@staticmethod
|
|
def normalize_prompt(prompt: Dict[str, Any]) -> Dict[str, Any]:
|
|
if not isinstance(prompt, dict) or not prompt:
|
|
raise ValueError("ComfyUI did not provide a valid API prompt.")
|
|
return copy.deepcopy(prompt)
|
|
|
|
@staticmethod
|
|
def prepare_for_worker(prompt: Dict[str, Any]) -> Dict[str, Any]:
|
|
prepared = WorkflowProcessor.normalize_prompt(prompt)
|
|
removed = []
|
|
for node_id, node in list(prepared.items()):
|
|
if isinstance(node, dict) and node.get("class_type") in DEADLINE_SUBMIT_NODE_TYPES:
|
|
removed.append(node_id)
|
|
del prepared[node_id]
|
|
|
|
if removed:
|
|
print(f"Deadline Submission: Removed submit node(s) from worker prompt: {', '.join(map(str, removed))}")
|
|
|
|
WorkflowProcessor.validate_worker_prompt(prepared)
|
|
return prepared
|
|
|
|
@staticmethod
|
|
def validate_worker_prompt(prompt: Dict[str, Any]) -> None:
|
|
if not prompt:
|
|
raise ValueError("Worker prompt is empty after removing Deadline submit nodes.")
|
|
|
|
has_output = any(
|
|
isinstance(node, dict) and node.get("class_type") in OUTPUT_NODE_TYPES
|
|
for node in prompt.values()
|
|
)
|
|
if not has_output:
|
|
print("Deadline Submission: Warning - worker prompt has no known output node.")
|
|
|
|
|
|
class InputAssetStager:
|
|
def __init__(self, output_directory: str, job_name: str, submission_id: str):
|
|
self.output_directory = os.path.abspath(output_directory)
|
|
self.job_name = job_name
|
|
self.submission_id = submission_id
|
|
|
|
def stage_referenced_assets(self, prompt: Dict[str, Any]) -> Tuple[str, str, List[Dict[str, Any]]]:
|
|
input_dir = self._get_local_input_directory()
|
|
references = self._collect_references(prompt, input_dir)
|
|
staging_dir = self._staging_directory()
|
|
manifest_path = os.path.join(staging_dir, f"deadline_input_manifest_{self.submission_id}.json")
|
|
|
|
if not references:
|
|
os.makedirs(staging_dir, exist_ok=True)
|
|
manifest = {
|
|
"submission_id": self.submission_id,
|
|
"input_directory": staging_dir,
|
|
"assets": [],
|
|
}
|
|
self._write_manifest(manifest_path, manifest)
|
|
return staging_dir, manifest_path, []
|
|
|
|
os.makedirs(staging_dir, exist_ok=True)
|
|
assets = []
|
|
for original_reference, source_path in sorted(references.items()):
|
|
# Absolute-path loader nodes are converted to a farm-visible staged
|
|
# absolute path later. Use only the basename inside the staging
|
|
# directory; never reproduce the submitter's local directory tree.
|
|
destination_rel_path = (
|
|
os.path.basename(source_path)
|
|
if os.path.isabs(original_reference)
|
|
else original_reference
|
|
)
|
|
staged_rel_path, destination_path = self._resolve_destination(staging_dir, destination_rel_path, source_path)
|
|
os.makedirs(os.path.dirname(destination_path), exist_ok=True)
|
|
if not os.path.exists(destination_path):
|
|
shutil.copy2(source_path, destination_path)
|
|
assets.append({
|
|
"original_relative_path": original_reference.replace("\\", "/"),
|
|
"staged_relative_path": staged_rel_path.replace("\\", "/"),
|
|
"relative_path": staged_rel_path.replace("\\", "/"),
|
|
"source": source_path,
|
|
"destination": destination_path,
|
|
"size": os.path.getsize(destination_path),
|
|
})
|
|
|
|
manifest = {
|
|
"submission_id": self.submission_id,
|
|
"input_directory": staging_dir,
|
|
"assets": assets,
|
|
}
|
|
self._write_manifest(manifest_path, manifest)
|
|
print(f"Deadline Submission: Staged {len(assets)} input asset(s) to {staging_dir}")
|
|
return staging_dir, manifest_path, assets
|
|
|
|
def _get_local_input_directory(self) -> str:
|
|
try:
|
|
import folder_paths
|
|
return os.path.abspath(folder_paths.get_input_directory())
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Could not resolve ComfyUI input directory: {exc}")
|
|
|
|
def _collect_references(self, prompt: Dict[str, Any], input_dir: str) -> Dict[str, str]:
|
|
references: Dict[str, str] = {}
|
|
for node in prompt.values():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
|
|
class_type = node.get("class_type", "")
|
|
inputs = node.get("inputs", {})
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
|
|
candidate_values: List[Tuple[Any, bool, bool]] = []
|
|
for field_name in INPUT_LOADER_FIELDS.get(class_type, ()):
|
|
if field_name in inputs:
|
|
candidate_values.append((inputs[field_name], True, False))
|
|
|
|
for field_name in ABSOLUTE_FILE_LOADER_FIELDS.get(class_type, ()):
|
|
if field_name in inputs:
|
|
candidate_values.append((inputs[field_name], True, True))
|
|
|
|
for value in inputs.values():
|
|
if isinstance(value, str):
|
|
candidate_values.append((value, False, False))
|
|
|
|
for value, strict, allow_absolute in candidate_values:
|
|
if not isinstance(value, str):
|
|
continue
|
|
resolved = self._resolve_input_file(value, input_dir, strict, allow_absolute)
|
|
if not resolved:
|
|
continue
|
|
rel_path, source_path = resolved
|
|
references[rel_path] = source_path
|
|
return references
|
|
|
|
def _resolve_destination(self, staging_dir: str, rel_path: str, source_path: str) -> Tuple[str, str]:
|
|
destination_path = os.path.abspath(os.path.join(staging_dir, rel_path))
|
|
if not os.path.exists(destination_path):
|
|
return rel_path, destination_path
|
|
|
|
if os.path.isfile(destination_path) and filecmp.cmp(source_path, destination_path, shallow=False):
|
|
return rel_path, destination_path
|
|
|
|
stem, extension = os.path.splitext(rel_path)
|
|
staged_rel_path = f"{stem}_{self.submission_id}{extension}"
|
|
return staged_rel_path, os.path.abspath(os.path.join(staging_dir, staged_rel_path))
|
|
|
|
def _resolve_input_file(
|
|
self,
|
|
value: str,
|
|
input_dir: str,
|
|
strict: bool,
|
|
allow_absolute: bool = False,
|
|
) -> Optional[Tuple[str, str]]:
|
|
clean_value, annotation = self._strip_annotation(value)
|
|
if annotation in {"output", "temp"}:
|
|
return None
|
|
if not clean_value:
|
|
return None
|
|
if os.path.splitext(clean_value)[1].lower() not in MEDIA_EXTENSIONS:
|
|
return None
|
|
|
|
if os.path.isabs(clean_value):
|
|
if not allow_absolute:
|
|
return None
|
|
candidate = os.path.abspath(clean_value)
|
|
if not os.path.isfile(candidate):
|
|
if not strict:
|
|
return None
|
|
raise FileNotFoundError(f"Referenced input asset was not found: {value} ({candidate})")
|
|
return clean_value, candidate
|
|
|
|
candidate = os.path.abspath(os.path.join(input_dir, clean_value))
|
|
try:
|
|
common = os.path.commonpath([input_dir, candidate])
|
|
except ValueError:
|
|
common = ""
|
|
if common != input_dir:
|
|
raise ValueError(f"Input asset escapes ComfyUI input directory: {value}")
|
|
if not os.path.isfile(candidate):
|
|
if not strict:
|
|
return None
|
|
raise FileNotFoundError(f"Referenced input asset was not found: {value} ({candidate})")
|
|
|
|
rel_path = os.path.relpath(candidate, input_dir)
|
|
return rel_path, candidate
|
|
|
|
def _strip_annotation(self, value: str) -> Tuple[str, Optional[str]]:
|
|
match = re.match(r"^(.*)\s+\[(input|output|temp)\]\s*$", value)
|
|
if not match:
|
|
return value.strip().replace("/", os.sep), None
|
|
return match.group(1).strip().replace("/", os.sep), match.group(2)
|
|
|
|
def _staging_directory(self) -> str:
|
|
parent = os.path.dirname(self.output_directory.rstrip("\\/"))
|
|
return os.path.join(parent, "input")
|
|
|
|
def _write_manifest(self, manifest_path: str, manifest: Dict[str, Any]) -> None:
|
|
with open(manifest_path, "w", encoding="utf-8") as handle:
|
|
json.dump(manifest, handle, indent=2)
|
|
|
|
|
|
class DeadlineJobSubmitter:
|
|
def __init__(self, workflow_data: Dict[str, Any], job_config: Dict[str, Any]):
|
|
self.workflow_data = workflow_data
|
|
self.job_config = job_config
|
|
|
|
def submit_job(self) -> Tuple[bool, str]:
|
|
try:
|
|
submission_dir = tempfile.mkdtemp(prefix="comfy_deadline_job_")
|
|
job_info_file, plugin_info_file, auxiliary_files = self._create_submission_files(submission_dir)
|
|
result = DeadlineCommandHelper.call_deadline_command([job_info_file, plugin_info_file] + auxiliary_files)
|
|
job_id = DeadlineCommandHelper.get_job_id_from_submission(result)
|
|
if not job_id:
|
|
return False, f"Deadline submission did not return a JobID. Output: {result}"
|
|
return True, job_id
|
|
except Exception as exc:
|
|
return False, str(exc)
|
|
|
|
def _create_submission_files(self, submission_dir: str) -> Tuple[str, str, List[str]]:
|
|
job_info_file = os.path.join(submission_dir, "job_info.txt")
|
|
plugin_info_file = os.path.join(submission_dir, "plugin_info.txt")
|
|
prompt_file = os.path.join(submission_dir, "prompt_to_execute.json")
|
|
standard_workflow_file = os.path.join(submission_dir, "workflow.json")
|
|
|
|
with open(prompt_file, "w", encoding="utf-8") as handle:
|
|
json.dump(self.workflow_data, handle, indent=2)
|
|
|
|
auxiliary_files = [prompt_file]
|
|
standard_workflow = self.job_config.get("standard_workflow")
|
|
if standard_workflow:
|
|
with open(standard_workflow_file, "w", encoding="utf-8") as handle:
|
|
json.dump(standard_workflow, handle, indent=2)
|
|
auxiliary_files.append(standard_workflow_file)
|
|
|
|
self._write_job_info(job_info_file)
|
|
self._write_plugin_info(plugin_info_file)
|
|
return job_info_file, plugin_info_file, auxiliary_files
|
|
|
|
def _write_job_info(self, path: str) -> None:
|
|
config = self.job_config
|
|
batch_count = int(config["batch_count"])
|
|
chunk_size = max(1, int(config["chunk_size"]))
|
|
output_dir = os.path.abspath(config["output_directory"])
|
|
|
|
with open(path, "w", encoding="utf-8") as handle:
|
|
handle.write("Plugin=ComfyUI\n")
|
|
handle.write(f"Name={config['job_name']}\n")
|
|
handle.write(f"Comment={config.get('comment', '')}\n")
|
|
handle.write(f"Department={config.get('department', '')}\n")
|
|
handle.write(f"Pool={'' if config['pool'] == 'none' else config['pool']}\n")
|
|
handle.write(f"Group={'' if config['group'] == 'none' else config['group']}\n")
|
|
handle.write(f"Priority={int(config['priority'])}\n")
|
|
handle.write(f"Frames=0-{batch_count - 1}\n")
|
|
handle.write(f"ChunkSize={chunk_size}\n")
|
|
handle.write(f"OutputDirectory0={output_dir}\n")
|
|
machine_list = config.get("machine_list", "")
|
|
if machine_list:
|
|
if not re.fullmatch(r"[A-Za-z0-9_,.\-]+", machine_list):
|
|
raise ValueError("Machine list must contain comma-separated Worker names.")
|
|
handle.write(f"Whitelist={machine_list}\n")
|
|
|
|
def _write_plugin_info(self, path: str) -> None:
|
|
config = self.job_config
|
|
entries = {
|
|
"StandardWorkflowFile": "workflow.json" if config.get("standard_workflow") else "",
|
|
"JobOutputDirectory": os.path.abspath(config["output_directory"]),
|
|
"JobInputDirectory": os.path.abspath(config["input_directory"]),
|
|
"InputManifestFile": os.path.abspath(config["input_manifest"]),
|
|
"SubmissionId": config["submission_id"],
|
|
"BatchCount": str(int(config["batch_count"])),
|
|
"BatchMode": "True",
|
|
"DefaultCudaDeviceZero": "True",
|
|
"SeedMode": "fixed",
|
|
"WorkerMode": "False",
|
|
"DistributedMode": "False",
|
|
"ForceNewInstance": "True",
|
|
"DisableDynamicVRAM": str(bool(config.get("disable_dynamic_vram", False))),
|
|
"ReserveVRAM": str(max(0.0, min(64.0, float(config.get("reserve_vram", 0.0))))),
|
|
}
|
|
|
|
with open(path, "w", encoding="utf-8") as handle:
|
|
for key, value in entries.items():
|
|
if value != "":
|
|
handle.write(f"{key}={value}\n")
|
|
|
|
|
|
class DeadlineSeed:
|
|
"""
|
|
Deadline-compatible seed node.
|
|
Batch mode varies by Deadline task ID; distributed mode varies by worker ID.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {
|
|
"default": 1125899906842,
|
|
"min": 0,
|
|
"max": 1125899906842624,
|
|
"forceInput": False,
|
|
}),
|
|
},
|
|
"hidden": {
|
|
"task_id": ("INT", {"default": 0}),
|
|
"batch_mode": ("BOOLEAN", {"default": False}),
|
|
"is_worker": ("BOOLEAN", {"default": False}),
|
|
"worker_id": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("seed",)
|
|
FUNCTION = "distribute"
|
|
CATEGORY = "deadline"
|
|
|
|
def distribute(self, seed, task_id=0, batch_mode=False, is_worker=False, worker_id=""):
|
|
seed = int(seed)
|
|
|
|
if _coerce_bool(batch_mode):
|
|
try:
|
|
task_id = int(task_id)
|
|
except (TypeError, ValueError):
|
|
task_id = 0
|
|
return (seed + task_id,)
|
|
|
|
if _coerce_bool(is_worker):
|
|
try:
|
|
worker_id = str(worker_id)
|
|
if worker_id.startswith("worker_"):
|
|
worker_index = int(worker_id.split("_")[1])
|
|
else:
|
|
worker_index = int(worker_id)
|
|
return (seed + worker_index + 1,)
|
|
except (TypeError, ValueError, IndexError):
|
|
return (seed,)
|
|
|
|
return (seed,)
|
|
|
|
|
|
class LegacyDeadlineSeedAlias(DeadlineSeed):
|
|
DEPRECATED = True
|
|
|
|
|
|
class DeadlineSubmitNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
pools = cls._get_deadline_pools()
|
|
groups = cls._get_deadline_groups()
|
|
return {
|
|
"required": {
|
|
"output_directory": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"placeholder": "Farm-visible output directory",
|
|
}),
|
|
"batch_count": ("INT", {
|
|
"default": NodeDefaults.BATCH_COUNT,
|
|
"min": 1,
|
|
"max": NodeDefaults.MAX_BATCH_COUNT,
|
|
"step": 1,
|
|
}),
|
|
"chunk_size": ("INT", {
|
|
"default": NodeDefaults.CHUNK_SIZE,
|
|
"min": 1,
|
|
"max": NodeDefaults.MAX_CHUNK_SIZE,
|
|
"step": 1,
|
|
}),
|
|
"priority": ("INT", {
|
|
"default": NodeDefaults.PRIORITY,
|
|
"min": 0,
|
|
"max": NodeDefaults.MAX_PRIORITY,
|
|
}),
|
|
"pool": (pools, {"default": NodeDefaults.POOL}),
|
|
"group": (groups, {"default": NodeDefaults.GROUP}),
|
|
"job_name": ("STRING", {"default": NodeDefaults.JOB_NAME}),
|
|
},
|
|
"optional": {
|
|
"comment": ("STRING", {"default": ""}),
|
|
"department": ("STRING", {"default": ""}),
|
|
"machine_list": ("STRING", {"default": "", "tooltip": "Optional comma-separated Deadline Worker allowlist."}),
|
|
"disable_dynamic_vram": ("BOOLEAN", {"default": False}),
|
|
"reserve_vram": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 64.0, "step": 0.5, "tooltip": "GB reserved for activations and other GPU users."}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("job_id",)
|
|
FUNCTION = "submit_to_deadline"
|
|
CATEGORY = "deadline"
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return f"deadline_submit_{time.time()}_{uuid.uuid4()}"
|
|
|
|
@classmethod
|
|
def _get_deadline_pools(cls) -> List[str]:
|
|
try:
|
|
output = DeadlineCommandHelper.call_deadline_command(["-pools"])
|
|
pools = [line.strip() for line in output.splitlines() if line.strip()]
|
|
return pools or [NodeDefaults.POOL]
|
|
except Exception as exc:
|
|
print(f"Deadline Submission: Could not query Deadline pools: {exc}")
|
|
return [NodeDefaults.POOL]
|
|
|
|
@classmethod
|
|
def _get_deadline_groups(cls) -> List[str]:
|
|
try:
|
|
output = DeadlineCommandHelper.call_deadline_command(["-groups"])
|
|
groups = [line.strip() for line in output.splitlines() if line.strip()]
|
|
return groups or [NodeDefaults.GROUP]
|
|
except Exception as exc:
|
|
print(f"Deadline Submission: Could not query Deadline groups: {exc}")
|
|
return [NodeDefaults.GROUP]
|
|
|
|
def submit_to_deadline(
|
|
self,
|
|
output_directory,
|
|
batch_count,
|
|
chunk_size,
|
|
priority,
|
|
pool,
|
|
group,
|
|
job_name,
|
|
comment="",
|
|
department="",
|
|
machine_list="",
|
|
disable_dynamic_vram=False,
|
|
reserve_vram=0.0,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
**_legacy_inputs,
|
|
):
|
|
try:
|
|
output_directory = self._prepare_output_directory(output_directory)
|
|
batch_count = max(1, int(batch_count))
|
|
chunk_size = max(1, min(int(chunk_size), batch_count))
|
|
submission_id = uuid.uuid4().hex[:12]
|
|
|
|
if prompt is None:
|
|
raise ValueError("ComfyUI did not inject the current API prompt.")
|
|
|
|
worker_prompt = WorkflowProcessor.prepare_for_worker(prompt)
|
|
stager = InputAssetStager(output_directory, job_name, submission_id)
|
|
input_dir, manifest_file, assets = stager.stage_referenced_assets(worker_prompt)
|
|
self._rewrite_prompt_asset_references(worker_prompt, assets)
|
|
standard_workflow = self._extract_standard_workflow(extra_pnginfo)
|
|
self._rewrite_standard_workflow_assets(standard_workflow, assets)
|
|
|
|
job_config = {
|
|
"submission_id": submission_id,
|
|
"job_name": job_name.strip() or NodeDefaults.JOB_NAME,
|
|
"priority": int(priority),
|
|
"pool": pool,
|
|
"group": group,
|
|
"batch_count": batch_count,
|
|
"chunk_size": chunk_size,
|
|
"output_directory": output_directory,
|
|
"input_directory": input_dir,
|
|
"input_manifest": manifest_file,
|
|
"standard_workflow": standard_workflow,
|
|
"comment": comment,
|
|
"department": department,
|
|
"machine_list": ",".join(part.strip() for part in machine_list.split(",") if part.strip()),
|
|
"disable_dynamic_vram": disable_dynamic_vram,
|
|
"reserve_vram": reserve_vram,
|
|
}
|
|
|
|
submitter = DeadlineJobSubmitter(worker_prompt, job_config)
|
|
success, result = submitter.submit_job()
|
|
if not success:
|
|
raise RuntimeError(result)
|
|
|
|
print(f"Deadline Submission: Submitted job {result} with {batch_count} variation(s), chunk size {chunk_size}, {len(assets)} staged asset(s).")
|
|
return (result,)
|
|
except Exception as exc:
|
|
raise
|
|
|
|
def _rewrite_prompt_asset_references(self, prompt: Dict[str, Any], assets: List[Dict[str, Any]]) -> None:
|
|
staged_by_original = self._asset_map(assets, "staged_relative_path", allow_absolute=True)
|
|
staged_absolute_by_original = self._asset_map(assets, "destination", allow_absolute=True)
|
|
if not staged_by_original:
|
|
return
|
|
|
|
for node in prompt.values():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
|
|
class_type = node.get("class_type", "")
|
|
for field_name in ABSOLUTE_FILE_LOADER_FIELDS.get(class_type, ()):
|
|
value = inputs.get(field_name)
|
|
normalized = self._normalize_asset_reference(value, allow_absolute=True)
|
|
if normalized in staged_absolute_by_original:
|
|
inputs[field_name] = staged_absolute_by_original[normalized]
|
|
|
|
for field_name in INPUT_LOADER_FIELDS.get(class_type, ()):
|
|
value = inputs.get(field_name)
|
|
normalized = self._normalize_asset_reference(value)
|
|
if normalized in staged_by_original:
|
|
inputs[field_name] = staged_by_original[normalized]
|
|
|
|
for field_name, value in list(inputs.items()):
|
|
normalized = self._normalize_asset_reference(value)
|
|
if normalized in staged_by_original:
|
|
inputs[field_name] = staged_by_original[normalized]
|
|
|
|
def _rewrite_standard_workflow_assets(self, workflow: Optional[Dict[str, Any]], assets: List[Dict[str, Any]]) -> None:
|
|
staged_absolute_by_original = self._asset_map(assets, "destination", allow_absolute=True)
|
|
if not workflow or not staged_absolute_by_original:
|
|
return
|
|
|
|
nodes = workflow.get("nodes", [])
|
|
if not isinstance(nodes, list):
|
|
return
|
|
|
|
for node in nodes:
|
|
if not isinstance(node, dict):
|
|
continue
|
|
widgets = node.get("widgets_values")
|
|
if not isinstance(widgets, list):
|
|
continue
|
|
|
|
for index, value in enumerate(widgets):
|
|
normalized = self._normalize_asset_reference(value)
|
|
if normalized in staged_absolute_by_original:
|
|
widgets[index] = staged_absolute_by_original[normalized]
|
|
|
|
def _asset_map(
|
|
self,
|
|
assets: List[Dict[str, Any]],
|
|
target_key: str,
|
|
allow_absolute: bool = False,
|
|
) -> Dict[str, str]:
|
|
mapping = {}
|
|
for asset in assets:
|
|
original = self._normalize_asset_reference(
|
|
asset.get("original_relative_path"), allow_absolute=allow_absolute
|
|
)
|
|
target = asset.get(target_key)
|
|
if original and target:
|
|
mapping[original] = target
|
|
return mapping
|
|
|
|
def _normalize_asset_reference(self, value: Any, allow_absolute: bool = False) -> Optional[str]:
|
|
if not isinstance(value, str) or not value:
|
|
return None
|
|
cleaned = re.sub(r"\s+\[(input|output|temp)\]\s*$", "", value.strip())
|
|
if os.path.isabs(cleaned) and not allow_absolute:
|
|
return None
|
|
return os.path.normpath(cleaned.replace("/", os.sep)).replace("\\", "/")
|
|
|
|
def _extract_standard_workflow(self, extra_pnginfo: Any) -> Optional[Dict[str, Any]]:
|
|
if not isinstance(extra_pnginfo, dict):
|
|
return None
|
|
|
|
workflow = extra_pnginfo.get("workflow")
|
|
if isinstance(workflow, dict):
|
|
return copy.deepcopy(workflow)
|
|
if isinstance(workflow, str):
|
|
try:
|
|
parsed = json.loads(workflow)
|
|
return parsed if isinstance(parsed, dict) else None
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return None
|
|
|
|
def _prepare_output_directory(self, output_directory: str) -> str:
|
|
output_directory = (output_directory or "").strip().strip("\"")
|
|
if not output_directory:
|
|
raise ValueError(
|
|
"Set DeadlineSubmit output_directory to a shared farm path, "
|
|
"e.g. \\\\server\\share\\output."
|
|
)
|
|
|
|
output_directory = os.path.abspath(os.path.expandvars(output_directory))
|
|
os.makedirs(output_directory, exist_ok=True)
|
|
if not os.path.isdir(output_directory):
|
|
raise ValueError(f"Output path is not a directory: {output_directory}")
|
|
return output_directory
|
|
|
|
|
|
def on_prompt(json_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
prompt = json_data.get("prompt")
|
|
if not isinstance(prompt, dict):
|
|
return json_data
|
|
|
|
submit_ids = [
|
|
str(node_id)
|
|
for node_id, node in prompt.items()
|
|
if isinstance(node, dict) and node.get("class_type") in DEADLINE_SUBMIT_NODE_TYPES
|
|
]
|
|
|
|
if submit_ids:
|
|
json_data["partial_execution_targets"] = submit_ids[:1]
|
|
if len(submit_ids) > 1:
|
|
print(f"Deadline Submission: Multiple submit nodes detected; only node {submit_ids[0]} will execute locally.")
|
|
return json_data
|
|
|
|
|
|
def register_on_prompt_handler() -> None:
|
|
try:
|
|
import server
|
|
instance = getattr(getattr(server, "PromptServer", None), "instance", None)
|
|
if instance and hasattr(instance, "add_on_prompt_handler"):
|
|
instance.add_on_prompt_handler(on_prompt)
|
|
print("Deadline Submission: Registered submit-only prompt handler.")
|
|
except Exception as exc:
|
|
print(f"Deadline Submission: Could not register prompt handler: {exc}")
|
|
|
|
try:
|
|
from aiohttp import web
|
|
import folder_paths
|
|
import server
|
|
|
|
instance = getattr(getattr(server, "PromptServer", None), "instance", None)
|
|
if instance and hasattr(instance, "routes"):
|
|
@instance.routes.get("/deadline/session")
|
|
async def deadline_session(_request):
|
|
return web.json_response({
|
|
"product": "ComfyUI-Deadline-Plugin",
|
|
"protocol": 1,
|
|
"session_id": DEADLINE_ENDPOINT_SESSION_ID,
|
|
"pid": os.getpid(),
|
|
"comfyui_root": os.path.abspath(folder_paths.base_path),
|
|
})
|
|
print(f"Deadline Submission: Registered endpoint identity {DEADLINE_ENDPOINT_SESSION_ID}.")
|
|
except Exception as exc:
|
|
print(f"Deadline Submission: Could not register endpoint identity route: {exc}")
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DeadlineSubmit": DeadlineSubmitNode,
|
|
"DeadlineSeed": DeadlineSeed,
|
|
"DeadlineDistributedSeed": LegacyDeadlineSeedAlias,
|
|
"DistributedSeed": LegacyDeadlineSeedAlias,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DeadlineSubmit": "Submit to Deadline",
|
|
"DeadlineSeed": "Deadline Seed",
|
|
"DeadlineDistributedSeed": "Deadline Seed",
|
|
"DistributedSeed": "Deadline Seed",
|
|
}
|