Isala Piyarisi
796bb396c0
HTTP API for ComfyUI workflow submission with: - POST /api/webhook endpoint for workflow submission - Real-time progress events via callback URL - Automatic output collection (no re-encoding) - Input/Output nodes for webhook data access
572 lines
17 KiB
Python
572 lines
17 KiB
Python
"""Output interception system for webhook delivery and progress events."""
|
|
|
|
import logging
|
|
import traceback
|
|
import time
|
|
import asyncio
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import datetime
|
|
from typing import Any, Callable, Optional
|
|
import uuid
|
|
|
|
from .types import WebhookContext, WebhookEvent, OutputInfo
|
|
from .context import get_webhook_context, unregister_webhook_context
|
|
from .client import get_webhook_client
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_original_task_done = None
|
|
_original_progress_hook = None
|
|
_interceptor_installed = False
|
|
_progress_hook_installed = False
|
|
_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="webhook_")
|
|
|
|
# Progress throttling - track last send time per prompt
|
|
_progress_last_sent: dict = {}
|
|
_PROGRESS_THROTTLE_MS = 500 # Minimum ms between progress updates
|
|
|
|
|
|
def install_output_interceptor() -> bool:
|
|
"""Install the output interceptor into ComfyUI's execution system."""
|
|
global _original_task_done, _interceptor_installed
|
|
if _interceptor_installed:
|
|
return True
|
|
try:
|
|
import execution
|
|
_original_task_done = execution.PromptQueue.task_done
|
|
execution.PromptQueue.task_done = _intercepted_task_done
|
|
_interceptor_installed = True
|
|
logger.info("Webhook output interceptor installed")
|
|
|
|
# Also install progress hook
|
|
install_progress_hook()
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to install output interceptor: {e}")
|
|
return False
|
|
|
|
|
|
def install_progress_hook() -> bool:
|
|
"""Install progress hook by wrapping the ProgressBar class."""
|
|
global _progress_hook_installed
|
|
|
|
if _progress_hook_installed:
|
|
return True
|
|
|
|
try:
|
|
import comfy.utils
|
|
|
|
# Wrap the ProgressBar.update_absolute method to intercept progress
|
|
original_update = comfy.utils.ProgressBar.update_absolute
|
|
|
|
def wrapped_update(self, value, total=None, preview=None):
|
|
# Call original first
|
|
result = original_update(self, value, total, preview)
|
|
|
|
# Then forward to webhook if applicable
|
|
try:
|
|
_handle_progress_update(value, self.total, preview, self.node_id)
|
|
except Exception as e:
|
|
logger.debug(f"Progress webhook error: {e}")
|
|
|
|
return result
|
|
|
|
comfy.utils.ProgressBar.update_absolute = wrapped_update
|
|
_progress_hook_installed = True
|
|
|
|
logger.info("Webhook progress hook installed")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to install progress hook: {e}")
|
|
return False
|
|
|
|
|
|
def _handle_progress_update(value, total, preview_image, node_id):
|
|
"""Handle progress update and forward to webhooks."""
|
|
global _progress_last_sent
|
|
|
|
# Get the current prompt_id from ComfyUI's server
|
|
try:
|
|
from server import PromptServer
|
|
prompt_id = PromptServer.instance.last_prompt_id
|
|
except Exception:
|
|
return
|
|
|
|
if not prompt_id:
|
|
return
|
|
|
|
# Look up webhook context for this prompt
|
|
context = get_webhook_context(prompt_id)
|
|
if context is None or not context.callback_url or not context.send_progress:
|
|
return
|
|
|
|
# Throttle progress updates
|
|
now_ms = time.time() * 1000
|
|
last_sent = _progress_last_sent.get(prompt_id, 0)
|
|
|
|
# Always send first and last updates
|
|
is_first = value <= 1
|
|
is_last = value >= total
|
|
|
|
if not is_first and not is_last:
|
|
if now_ms - last_sent < _PROGRESS_THROTTLE_MS:
|
|
return
|
|
|
|
_progress_last_sent[prompt_id] = now_ms
|
|
|
|
# Send progress event asynchronously
|
|
_executor.submit(
|
|
_send_progress_sync,
|
|
context.callback_url,
|
|
context.request_id,
|
|
prompt_id,
|
|
node_id or "unknown",
|
|
value,
|
|
total,
|
|
context.get_auth_headers() if hasattr(context, 'get_auth_headers') else {}
|
|
)
|
|
|
|
|
|
def _webhook_progress_hook(value, total, preview_image=None, prompt_id=None, node_id=None):
|
|
"""Legacy progress hook wrapper - kept for compatibility."""
|
|
global _original_progress_hook, _progress_last_sent
|
|
|
|
# Always call the original hook first (ComfyUI's websocket progress)
|
|
if _original_progress_hook is not None:
|
|
try:
|
|
_original_progress_hook(value, total, preview_image, prompt_id, node_id)
|
|
except Exception as e:
|
|
logger.warning(f"Original progress hook error: {e}")
|
|
|
|
# Now handle webhook progress
|
|
if prompt_id is None:
|
|
return
|
|
|
|
# Look up webhook context for this prompt
|
|
context = get_webhook_context(prompt_id)
|
|
if context is None or not context.callback_url or not context.send_progress:
|
|
return
|
|
|
|
# Throttle progress updates
|
|
now_ms = time.time() * 1000
|
|
last_sent = _progress_last_sent.get(prompt_id, 0)
|
|
|
|
# Always send first (value=1) and last (value=total) updates
|
|
is_first = value == 1 or value == 0
|
|
is_last = value >= total
|
|
|
|
if not is_first and not is_last:
|
|
if now_ms - last_sent < _PROGRESS_THROTTLE_MS:
|
|
return
|
|
|
|
_progress_last_sent[prompt_id] = now_ms
|
|
|
|
# Send progress event asynchronously
|
|
_executor.submit(
|
|
_send_progress_sync,
|
|
context.callback_url,
|
|
context.request_id,
|
|
prompt_id,
|
|
node_id or "unknown",
|
|
value,
|
|
total,
|
|
context.get_auth_headers() if hasattr(context, 'get_auth_headers') else {}
|
|
)
|
|
|
|
|
|
def _send_progress_sync(callback_url, request_id, prompt_id, node_id, value, total, headers):
|
|
"""Send progress event synchronously (for thread pool)."""
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
try:
|
|
loop.run_until_complete(
|
|
_send_progress_async(callback_url, request_id, prompt_id, node_id, value, total, headers)
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"Progress send error: {e}")
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
async def _send_progress_async(callback_url, request_id, prompt_id, node_id, value, total, headers):
|
|
"""Send progress event to callback URL."""
|
|
payload = {
|
|
"event": WebhookEvent.PROGRESS.value,
|
|
"request_id": request_id,
|
|
"prompt_id": prompt_id,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"node_id": node_id,
|
|
"progress": {
|
|
"value": value,
|
|
"max": total
|
|
}
|
|
}
|
|
|
|
try:
|
|
client = get_webhook_client()
|
|
result = await client.send_json(
|
|
url=callback_url,
|
|
payload=payload,
|
|
headers=headers
|
|
)
|
|
|
|
if result.success:
|
|
logger.debug(f"Progress {value}/{total} sent for {prompt_id}")
|
|
else:
|
|
logger.debug(f"Progress send failed: {result.error}")
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Progress send exception: {e}")
|
|
|
|
|
|
def uninstall_output_interceptor() -> bool:
|
|
"""Uninstall the output interceptor."""
|
|
global _original_task_done, _interceptor_installed
|
|
if not _interceptor_installed:
|
|
return True
|
|
try:
|
|
import execution
|
|
if _original_task_done is not None:
|
|
execution.PromptQueue.task_done = _original_task_done
|
|
_original_task_done = None
|
|
_interceptor_installed = False
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to uninstall output interceptor: {e}")
|
|
return False
|
|
|
|
|
|
def _intercepted_task_done(self, item_id: Any, history_result: dict, status: Any, process_item: Optional[Callable] = None):
|
|
"""Intercepted task_done to send completion events."""
|
|
global _original_task_done
|
|
|
|
prompt_tuple = self.currently_running.get(item_id)
|
|
prompt_id = None
|
|
extra_data = {}
|
|
|
|
if prompt_tuple:
|
|
prompt_id = prompt_tuple[1] if len(prompt_tuple) > 1 else None
|
|
extra_data = prompt_tuple[3] if len(prompt_tuple) > 3 else {}
|
|
|
|
# Check if this is a webhook request (has webhook_config in extra_data)
|
|
webhook_config = extra_data.get("webhook_config") if isinstance(extra_data, dict) else None
|
|
|
|
if webhook_config and prompt_id:
|
|
try:
|
|
outputs = history_result.get("outputs", {})
|
|
meta = history_result.get("meta", {})
|
|
|
|
# Submit webhook delivery to thread pool
|
|
_executor.submit(
|
|
_send_webhook_completion_sync,
|
|
webhook_config,
|
|
prompt_id,
|
|
outputs,
|
|
meta,
|
|
status
|
|
)
|
|
logger.debug(f"Webhook delivery queued for prompt {prompt_id}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to queue webhook send: {e}")
|
|
|
|
# Call original task_done - MUST happen for ComfyUI to work correctly
|
|
if _original_task_done:
|
|
return _original_task_done(self, item_id, history_result, status, process_item)
|
|
|
|
|
|
def _send_webhook_completion_sync(
|
|
webhook_config: dict,
|
|
prompt_id: str,
|
|
outputs: dict,
|
|
meta: dict,
|
|
status: Any
|
|
):
|
|
"""Send webhook completion in sync context (for thread pool)."""
|
|
import asyncio
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
try:
|
|
loop.run_until_complete(
|
|
_send_webhook_completion(webhook_config, prompt_id, outputs, meta, status)
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Webhook delivery failed for prompt {prompt_id}: {e}\n{traceback.format_exc()}")
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
async def _send_webhook_completion(
|
|
webhook_config: dict,
|
|
prompt_id: str,
|
|
outputs: dict,
|
|
meta: dict,
|
|
status: Any
|
|
):
|
|
"""Send completion event with outputs to callback URL."""
|
|
from ..processors.outputs import (
|
|
process_outputs_for_webhook,
|
|
load_output_file,
|
|
enrich_output_info
|
|
)
|
|
|
|
callback_url = webhook_config.get("callback_url")
|
|
if not callback_url:
|
|
logger.debug(f"No callback_url in webhook_config for prompt {prompt_id}")
|
|
return
|
|
|
|
request_id = webhook_config.get("request_id", str(uuid.uuid4()))
|
|
start_time = datetime.now()
|
|
|
|
try:
|
|
# Process all outputs from all nodes
|
|
all_outputs = process_outputs_for_webhook(outputs, meta)
|
|
|
|
if not all_outputs:
|
|
all_outputs = []
|
|
|
|
# Enrich outputs with dimensions, file sizes
|
|
for output in all_outputs:
|
|
try:
|
|
enrich_output_info(output)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to enrich output {output.filename}: {e}")
|
|
|
|
# Build files list for multipart
|
|
files = []
|
|
output_metadata = []
|
|
|
|
for idx, output in enumerate(all_outputs):
|
|
output_metadata.append(output.to_dict())
|
|
|
|
if output.filename:
|
|
try:
|
|
file_data = load_output_file(output)
|
|
if file_data:
|
|
files.append((
|
|
f"file_{idx}",
|
|
output.filename,
|
|
file_data,
|
|
output.mime_type or "application/octet-stream"
|
|
))
|
|
except Exception as e:
|
|
logger.warning(f"Failed to load output file {output.filename}: {e}")
|
|
|
|
# Extract status information
|
|
status_str = "unknown"
|
|
status_messages = []
|
|
if status is not None:
|
|
if hasattr(status, "status_str"):
|
|
status_str = status.status_str
|
|
elif hasattr(status, "_asdict"):
|
|
status_dict = status._asdict()
|
|
status_str = status_dict.get("status_str", "unknown")
|
|
status_messages = status_dict.get("messages", [])
|
|
|
|
# Calculate execution time from context if available
|
|
execution_time_ms = None
|
|
context = get_webhook_context(prompt_id)
|
|
if context and context.start_time:
|
|
execution_time_ms = (datetime.now() - context.start_time).total_seconds() * 1000
|
|
|
|
# Build metadata payload
|
|
event_type = WebhookEvent.COMPLETED.value if status_str == "success" else WebhookEvent.ERROR.value
|
|
metadata = {
|
|
"event": event_type,
|
|
"request_id": request_id,
|
|
"prompt_id": prompt_id,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"status": status_str,
|
|
"outputs": output_metadata,
|
|
"output_count": len(all_outputs),
|
|
"file_count": len(files),
|
|
"execution_time_ms": execution_time_ms,
|
|
"nodes_executed": len(outputs),
|
|
}
|
|
|
|
if status_messages:
|
|
metadata["messages"] = status_messages
|
|
|
|
# Get auth headers
|
|
headers = _get_auth_headers(webhook_config)
|
|
|
|
# Send the webhook
|
|
client = get_webhook_client()
|
|
|
|
if files:
|
|
result = await client.send_multipart(
|
|
url=callback_url,
|
|
metadata=metadata,
|
|
files=files,
|
|
headers=headers
|
|
)
|
|
else:
|
|
result = await client.send_json(
|
|
url=callback_url,
|
|
payload=metadata,
|
|
headers=headers
|
|
)
|
|
|
|
if result.success:
|
|
logger.info(
|
|
f"Webhook delivered for prompt {prompt_id}: "
|
|
f"{len(files)} files, {len(output_metadata)} outputs, "
|
|
f"status={result.status}, {result.elapsed_ms:.0f}ms"
|
|
)
|
|
else:
|
|
logger.error(
|
|
f"Webhook failed for prompt {prompt_id}: "
|
|
f"status={result.status}, error={result.error}"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error sending webhook for prompt {prompt_id}: {e}\n{traceback.format_exc()}")
|
|
|
|
finally:
|
|
# Clean up context
|
|
try:
|
|
unregister_webhook_context(prompt_id)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _get_auth_headers(webhook_config: dict) -> dict:
|
|
"""Build authentication headers from webhook config."""
|
|
headers = {}
|
|
|
|
auth_header = webhook_config.get("auth_header")
|
|
auth_value = webhook_config.get("auth_value")
|
|
|
|
if auth_header and auth_value:
|
|
headers[auth_header] = auth_value
|
|
|
|
custom_headers = webhook_config.get("headers")
|
|
if isinstance(custom_headers, dict):
|
|
headers.update(custom_headers)
|
|
|
|
return headers
|
|
|
|
|
|
async def send_progress_event(
|
|
context: WebhookContext,
|
|
node_id: str,
|
|
node_type: str,
|
|
progress: float,
|
|
message: str = ""
|
|
):
|
|
"""Send a progress event to the callback URL."""
|
|
if not context.send_progress or not context.callback_url:
|
|
return
|
|
|
|
elapsed_ms = None
|
|
if context.start_time:
|
|
elapsed_ms = (datetime.now() - context.start_time).total_seconds() * 1000
|
|
|
|
payload = {
|
|
"event": WebhookEvent.PROGRESS.value,
|
|
"request_id": context.request_id,
|
|
"prompt_id": context.prompt_id,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"node_id": node_id,
|
|
"node_type": node_type,
|
|
"progress": progress,
|
|
"message": message,
|
|
"elapsed_ms": elapsed_ms,
|
|
}
|
|
|
|
try:
|
|
client = get_webhook_client()
|
|
result = await client.send_json(
|
|
url=context.callback_url,
|
|
payload=payload,
|
|
headers=context.get_auth_headers()
|
|
)
|
|
|
|
if not result.success:
|
|
logger.warning(f"Progress event failed: {result.error}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to send progress event: {e}")
|
|
|
|
|
|
async def send_started_event(context: WebhookContext):
|
|
"""Send workflow started event."""
|
|
if not context.callback_url:
|
|
return
|
|
|
|
payload = {
|
|
"event": WebhookEvent.STARTED.value,
|
|
"request_id": context.request_id,
|
|
"prompt_id": context.prompt_id,
|
|
"timestamp": datetime.now().isoformat(),
|
|
}
|
|
|
|
try:
|
|
client = get_webhook_client()
|
|
await client.send_json(
|
|
url=context.callback_url,
|
|
payload=payload,
|
|
headers=context.get_auth_headers()
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to send started event: {e}")
|
|
|
|
|
|
async def send_error_event(
|
|
context: WebhookContext,
|
|
error_type: str,
|
|
error_message: str,
|
|
node_id: Optional[str] = None,
|
|
node_type: Optional[str] = None,
|
|
traceback_str: Optional[str] = None
|
|
):
|
|
"""Send workflow error event."""
|
|
if not context.callback_url:
|
|
return
|
|
|
|
elapsed_ms = None
|
|
if context.start_time:
|
|
elapsed_ms = (datetime.now() - context.start_time).total_seconds() * 1000
|
|
|
|
payload = {
|
|
"event": WebhookEvent.ERROR.value,
|
|
"request_id": context.request_id,
|
|
"prompt_id": context.prompt_id,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"error": {
|
|
"type": error_type,
|
|
"message": error_message,
|
|
"node_id": node_id,
|
|
"node_type": node_type,
|
|
},
|
|
"elapsed_ms": elapsed_ms,
|
|
}
|
|
|
|
if traceback_str:
|
|
payload["error"]["traceback"] = traceback_str[:2000]
|
|
|
|
try:
|
|
client = get_webhook_client()
|
|
await client.send_json(
|
|
url=context.callback_url,
|
|
payload=payload,
|
|
headers=context.get_auth_headers()
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to send error event: {e}")
|
|
|
|
|
|
def shutdown():
|
|
"""Shutdown the webhook executor gracefully."""
|
|
global _executor
|
|
try:
|
|
_executor.shutdown(wait=True, cancel_futures=False)
|
|
logger.info("Webhook executor shut down")
|
|
except Exception as e:
|
|
logger.warning(f"Error shutting down webhook executor: {e}")
|