commit 796bb396c041b4cb4b6b4fa40d6c232c90ab279b Author: Isala Piyarisi Date: Wed Jan 14 23:27:00 2026 +0530 Initial commit: ComfyUI Workflow API 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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3f73620 --- /dev/null +++ b/.gitignore @@ -0,0 +1,50 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# Virtual environments +.env +.venv +env/ +venv/ +ENV/ + +# IDE +.idea/ +.vscode/ +*.swp +*.swo +*~ + +# Testing +.pytest_cache/ +.coverage +htmlcov/ +.tox/ +.nox/ + +# OS +.DS_Store +Thumbs.db + +# Project specific +outputs/ +*.log diff --git a/README.md b/README.md new file mode 100644 index 0000000..58a62ac --- /dev/null +++ b/README.md @@ -0,0 +1,202 @@ +# ComfyUI Workflow API + +HTTP API for ComfyUI with webhook-based workflow execution. Submit workflows via HTTP, receive real-time progress updates and outputs at your callback URL. + +## Installation + +```bash +cd /path/to/ComfyUI/custom_nodes +git clone https://github.com/isala404/comfy-workflow-api.git +pip install -r comfy-workflow-api/requirements.txt +``` + +## Use Cases + +- **Backend Integration**: Run ComfyUI as a headless service behind your application +- **Batch Processing**: Submit multiple workflows and receive results asynchronously +- **Distributed Systems**: Decouple workflow submission from result handling +- **Mobile/Web Apps**: Trigger image generation from any HTTP client + +## How It Works + +``` +Client ComfyUI Your Server + │ │ │ + ├─── POST /api/webhook ─────────►│ │ + │ (workflow + callback_url) │ │ + │ │ │ + │◄── {request_id, prompt_id} ────┤ │ + │ │ │ + │ ├── workflow.progress ──────────►│ + │ ├── workflow.progress ──────────►│ + │ ├── ... │ + │ │ │ + │ ├── multipart (files) ──────────►│ + │ │ │ +``` + +## Quick Start + +### 1. Create a workflow using webhook nodes + +Use **WebhookReceiver** → **WebhookTransformer** → **WebhookSend** to read inputs from HTTP and send outputs back: + +``` + ┌─► WebhookTransformer ──► CLIPTextEncode ──┐ +WebhookReceiver ────┤ (extracts "prompt") ├──► KSampler ──► WebhookSend + └─► WebhookTransformer ──► VAEEncode ───────┘ + (extracts "image") +``` + +### 2. Submit via HTTP + +```bash +# Image-to-image example +curl -X POST http://localhost:8188/api/webhook \ + -F "workflow=@workflow.json" \ + -F "callback_url=http://your-server.com/webhook" \ + -F "prompt=turn it into a watercolor painting" \ + -F "image=@input.png" +``` + +### 3. Receive at your callback + +Your server receives: +- **Progress events** (JSON): `{"event": "workflow.progress", "progress": {"value": 5, "max": 20}}` +- **Output files** (gzip multipart): Images, audio, video, etc. + +## Nodes + +### WebhookReceiver + +Entry point for webhook workflows. Reads request configuration from the HTTP call. + +| Output | Type | Description | +|--------|------|-------------| +| webhook_context | WEBHOOK_CONTEXT | Context for other webhook nodes | + +### WebhookTransformer + +Extracts a field from the HTTP request. Connects to any input type. + +| Input | Description | +|-------|-------------| +| webhook_context | From WebhookReceiver | +| field | Field name to extract (e.g., "prompt") | +| default_value | Optional fallback if field missing | + +| Output | Type | Description | +|--------|------|-------------| +| value | * (wildcard) | Extracted value, auto-typed | + +**Type inference**: Strings are converted to INT, FLOAT, BOOLEAN when applicable. File uploads become IMAGE, AUDIO, VIDEO, MESH, or LATENT based on content type. + +### WebhookSend + +Sends outputs to the callback URL. Supports 5 output fields. + +| Input | Description | +|-------|-------------| +| webhook_context | From WebhookReceiver | +| field_1..5 | Any ComfyUI type (IMAGE, AUDIO, STRING, etc.) | +| field_1..5_name | Names for each output | + +**Encoding**: IMAGE→PNG, AUDIO→FLAC, VIDEO→MP4, MESH→GLB, LATENT→safetensors + +## API Reference + +### POST /api/webhook + +| Field | Required | Description | +|-------|----------|-------------| +| workflow | Yes | Workflow JSON | +| callback_url | Yes | URL to receive events/outputs | +| auth_header | No | Auth header name | +| auth_value | No | Auth header value | +| timeout | No | HTTP timeout (default: 60s) | +| max_retries | No | Retry attempts (default: 3) | +| *custom fields* | No | Any field accessible via WebhookTransformer | + +**Response:** +```json +{"request_id": "...", "prompt_id": "...", "status": "queued"} +``` + +### GET /api/webhook/status/{request_id} + +Check request status. + +### DELETE /api/webhook/{request_id} + +Cancel a queued/running workflow. + +## Webhook Events + +| Event | Content-Type | Description | +|-------|--------------|-------------| +| workflow.started | application/json | Execution began | +| workflow.progress | application/json | Step progress (throttled 500ms) | +| workflow.completed | multipart/form-data (gzip) | Output files + metadata | +| workflow.error | application/json | Execution failed | + +**Note**: Completion payload is gzip compressed. Decompress before parsing. + +## Example + +**workflow.json** (see `example_workflow.json`): +```json +{ + "10": { + "class_type": "WebhookReceiver", + "inputs": {"debug": true} + }, + "11": { + "class_type": "WebhookTransformer", + "inputs": {"field": "prompt", "webhook_context": ["10", 0]} + }, + "6": { + "class_type": "CLIPTextEncode", + "inputs": {"text": ["11", 0], "clip": ["4", 1]} + }, + "...": "... rest of workflow ...", + "12": { + "class_type": "WebhookSend", + "inputs": {"webhook_context": ["10", 0], "field_1": ["8", 0]} + } +} +``` + +**Python client** (see `example_client.py`): +```python +import requests, json + +# Load image for img2img +image_data = open("input.png", "rb").read() + +response = requests.post( + "http://localhost:8188/api/webhook", + files={ + "workflow": (None, json.dumps(workflow), "application/json"), + "callback_url": (None, "http://localhost:5001/webhook"), + "prompt": (None, "make it look like a painting"), + "image": ("input.png", image_data, "image/png"), + } +) +``` + +## Testing + +```bash +# Terminal 1: Start ComfyUI +python main.py + +# Terminal 2: Run example client (image-to-image) +cd custom_nodes/comfy-workflow-api +python example_client.py sample_input.png "turn it into a watercolor painting" +``` + +A sample image (`sample_input.png`) is included in the repo for testing. + +## License + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..5270cc6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,77 @@ +""" +ComfyUI Workflow API - Server extension for workflow integration. + +Provides: +1. HTTP endpoint: POST /api/webhook for submitting workflows +2. WebhookReceiver node - Entry point for API-triggered workflows +3. WebhookTransformer node - Extract individual fields by name +4. WebhookSend node - Send outputs back to callback URL +5. Real-time progress updates via callback URL +""" + +__version__ = "1.0.0" + +import logging + +logger = logging.getLogger(__name__) + +# Import nodes +from .nodes.receiver import WebhookReceiver +from .nodes.transformer import WebhookTransformer +from .nodes.send import WebhookSend + +# Node mappings for ComfyUI +NODE_CLASS_MAPPINGS = { + "WebhookReceiver": WebhookReceiver, + "WebhookTransformer": WebhookTransformer, + "WebhookSend": WebhookSend, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "WebhookReceiver": "Workflow API Receiver", + "WebhookTransformer": "Workflow API Transformer", + "WebhookSend": "Workflow API Send", +} + +# Initialize server extension +_initialized = False + + +def _initialize(): + """Initialize server extension and interceptor hooks.""" + global _initialized + + if _initialized: + return + + try: + # Register HTTP routes + from .server.routes import register_routes + if register_routes(): + logger.info(f"[comfy-workflow-api] v{__version__} - Routes registered") + + # Install output interceptor for automatic webhook delivery + from .core.interceptor import install_output_interceptor + if install_output_interceptor(): + logger.info(f"[comfy-workflow-api] v{__version__} - Output interceptor installed") + + _initialized = True + logger.info(f"[comfy-workflow-api] v{__version__} loaded successfully") + + except Exception as e: + logger.error(f"[comfy-workflow-api] Failed to initialize: {e}", exc_info=True) + + +# Initialize on import +_initialize() + + +# Cleanup function +def _cleanup(): + """Cleanup on unload.""" + try: + from .core.interceptor import shutdown, uninstall_output_interceptor + uninstall_output_interceptor() + shutdown() + except Exception as e: + logger.warning(f"[comfy-workflow-api] Cleanup error: {e}") diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..39fe0d0 --- /dev/null +++ b/core/__init__.py @@ -0,0 +1,45 @@ +"""Core components for webhook functionality.""" + +from .types import WebhookContext, WebhookEvent, WebhookResult, OutputInfo, WebhookField +from .context import ( + register_webhook_context, + get_webhook_context, + unregister_webhook_context, + get_all_contexts, + clear_all_contexts, + cleanup_old_contexts, + # Legacy aliases + register_context, + get_context, + get_context_by_prompt, + update_context, + unregister_context, +) +from .client import get_webhook_client, WebhookClient +from .collector import collect_outputs + +__all__ = [ + # Types + "WebhookContext", + "WebhookEvent", + "WebhookResult", + "OutputInfo", + "WebhookField", + # Context functions + "register_webhook_context", + "get_webhook_context", + "unregister_webhook_context", + "get_all_contexts", + "clear_all_contexts", + "cleanup_old_contexts", + # Legacy aliases + "register_context", + "get_context", + "get_context_by_prompt", + "update_context", + "unregister_context", + # Client + "get_webhook_client", + "WebhookClient", + "collect_outputs", +] diff --git a/core/client.py b/core/client.py new file mode 100644 index 0000000..1adfb53 --- /dev/null +++ b/core/client.py @@ -0,0 +1,472 @@ +""" +Production-grade async HTTP client for webhook delivery. +Supports multipart/form-data for file uploads with retry logic. +""" + +import asyncio +import gzip +import io +import json +import logging +import random +import time +from typing import Optional +from pathlib import Path + +import aiohttp + +from .types import WebhookResult + +logger = logging.getLogger(__name__) + +# Singleton client instance +_webhook_client: Optional["WebhookClient"] = None + + +def get_webhook_client() -> "WebhookClient": + """Get or create the singleton webhook client instance.""" + global _webhook_client + if _webhook_client is None: + _webhook_client = WebhookClient() + return _webhook_client + + +class WebhookClient: + """ + Production-grade async HTTP client for multipart webhook delivery. + + Features: + - Async HTTP with connection pooling + - Multipart/form-data support for file uploads + - Exponential backoff with jitter for retries + - Configurable timeout and retry settings + """ + + # Status codes that should trigger retry + RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504} + + def __init__( + self, + timeout: float = 60.0, + max_retries: int = 3, + base_delay: float = 1.0, + max_delay: float = 60.0, + jitter_factor: float = 0.2, + max_connections: int = 10, + max_connections_per_host: int = 5, + ): + """ + Initialize the webhook client. + + Args: + timeout: Request timeout in seconds + max_retries: Maximum number of retry attempts + base_delay: Base delay for exponential backoff + max_delay: Maximum delay between retries + jitter_factor: Random jitter factor (0-1) + max_connections: Maximum total connections + max_connections_per_host: Maximum connections per host + """ + self._session: Optional[aiohttp.ClientSession] = None + self._session_loop: Optional[asyncio.AbstractEventLoop] = None + self.timeout = aiohttp.ClientTimeout(total=timeout) + self.max_retries = max_retries + self.base_delay = base_delay + self.max_delay = max_delay + self.jitter_factor = jitter_factor + self.max_connections = max_connections + self.max_connections_per_host = max_connections_per_host + + async def get_session(self) -> aiohttp.ClientSession: + """Get or create aiohttp session with connection pooling.""" + current_loop = asyncio.get_running_loop() + + # Create new session if none exists, if closed, or if event loop changed + if (self._session is None or + self._session.closed or + self._session_loop is not current_loop): + # Close old session if it exists and loop is still valid + if self._session and not self._session.closed: + try: + await self._session.close() + except Exception: + pass + + connector = aiohttp.TCPConnector( + limit=self.max_connections, + limit_per_host=self.max_connections_per_host, + keepalive_timeout=30, + enable_cleanup_closed=True, + ) + self._session = aiohttp.ClientSession( + connector=connector, + timeout=self.timeout, + ) + self._session_loop = current_loop + return self._session + + async def send_json( + self, + url: str, + payload: dict, + headers: Optional[dict] = None, + ) -> WebhookResult: + """ + Send JSON payload to webhook URL. + + Used for progress events, errors, and status updates. + + Args: + url: Webhook URL + payload: JSON-serializable dictionary + headers: Optional additional headers (e.g., auth) + + Returns: + WebhookResult with success status and details + """ + session = await self.get_session() + all_headers = {"Content-Type": "application/json"} + if headers: + all_headers.update(headers) + + return await self._execute_with_retry( + url=url, + session=session, + request_kwargs={"json": payload, "headers": all_headers} + ) + + async def send_multipart( + self, + url: str, + metadata: dict, + files: list[tuple[str, str, bytes, str]], + headers: Optional[dict] = None, + timeout: Optional[float] = None, + max_retries: Optional[int] = None, + compress: bool = False, + ) -> WebhookResult: + """ + Send multipart/form-data request with files. + + Args: + url: Webhook URL + metadata: JSON metadata dictionary + files: List of (field_name, filename, binary_data, mime_type) + headers: Optional additional headers (e.g., auth) + timeout: Optional request timeout (overrides default) + max_retries: Optional max retries (overrides default) + compress: If True, compress payload with gzip + + Returns: + WebhookResult with success status and details + """ + session = await self.get_session() + + if compress: + # Build compressed payload + return await self._send_multipart_compressed( + url=url, + session=session, + metadata=metadata, + files=files, + headers=headers, + timeout=timeout, + max_retries=max_retries, + ) + + def build_form(): + # Build FormData for each request (cannot reuse) + form = aiohttp.FormData() + + # Add metadata as JSON part + form.add_field( + "metadata", + json.dumps(metadata), + content_type="application/json", + ) + + # Add files + for name, filename, data, content_type in files: + form.add_field( + name, + data, + filename=filename, + content_type=content_type, + ) + return form + + return await self._execute_with_retry( + url=url, + session=session, + request_kwargs={"headers": headers}, + form_builder=build_form, + timeout=timeout, + max_retries=max_retries, + ) + + async def _send_multipart_compressed( + self, + url: str, + session: aiohttp.ClientSession, + metadata: dict, + files: list[tuple[str, str, bytes, str]], + headers: Optional[dict] = None, + timeout: Optional[float] = None, + max_retries: Optional[int] = None, + ) -> WebhookResult: + """ + Send gzip-compressed multipart/form-data request. + + Builds the multipart body, compresses it with gzip, and sends + with Content-Encoding: gzip header. + """ + import secrets + + def build_compressed_payload(): + # Generate a unique boundary + boundary = f"----WebhookBoundary{secrets.token_hex(16)}" + boundary_bytes = boundary.encode("utf-8") + + # Build multipart body manually + parts = [] + + # Add metadata part + parts.append(b"--" + boundary_bytes) + parts.append(b'Content-Disposition: form-data; name="metadata"') + parts.append(b"Content-Type: application/json") + parts.append(b"") + parts.append(json.dumps(metadata).encode("utf-8")) + + # Add file parts + for name, filename, data, content_type in files: + parts.append(b"--" + boundary_bytes) + disposition = f'Content-Disposition: form-data; name="{name}"; filename="{filename}"' + parts.append(disposition.encode("utf-8")) + parts.append(f"Content-Type: {content_type}".encode("utf-8")) + parts.append(b"") + parts.append(data) + + # Final boundary + parts.append(b"--" + boundary_bytes + b"--") + parts.append(b"") + + # Join with CRLF + body = b"\r\n".join(parts) + + # Compress with gzip + compressed = gzip.compress(body, compresslevel=6) + + return compressed, boundary + + compressed_data, boundary = build_compressed_payload() + + # Build headers + all_headers = { + "Content-Type": f"multipart/form-data; boundary={boundary}", + "Content-Encoding": "gzip", + "Content-Length": str(len(compressed_data)), + } + if headers: + all_headers.update(headers) + + return await self._execute_with_retry( + url=url, + session=session, + request_kwargs={"headers": all_headers, "data": compressed_data}, + timeout=timeout, + max_retries=max_retries, + ) + + async def send_multipart_streaming( + self, + url: str, + metadata: dict, + file_paths: list[tuple[str, str, str, str]], + headers: Optional[dict] = None, + ) -> WebhookResult: + """ + Send multipart request with files streamed from disk. + + Use this for large files (>50MB) to avoid memory issues. + + Args: + url: Webhook URL + metadata: JSON metadata dictionary + file_paths: List of (field_name, filename, file_path, mime_type) + headers: Optional additional headers (e.g., auth) + + Returns: + WebhookResult with success status and details + """ + session = await self.get_session() + + def build_writer(): + # Build multipart writer for streaming + writer = aiohttp.MultipartWriter("form-data") + + # Add metadata + metadata_payload = aiohttp.payload.JsonPayload(metadata) + metadata_payload.set_content_disposition("form-data", name="metadata") + writer.append_payload(metadata_payload) + + # Add files (streamed from disk) + for name, filename, path, content_type in file_paths: + file_path = Path(path) + if file_path.exists(): + with open(file_path, "rb") as f: + file_data = f.read() + payload = aiohttp.payload.BytesPayload( + file_data, + content_type=content_type, + ) + payload.set_content_disposition( + "form-data", name=name, filename=filename + ) + writer.append_payload(payload) + return writer + + return await self._execute_with_retry( + url=url, + session=session, + request_kwargs={"headers": headers}, + data_builder=build_writer + ) + + async def _execute_with_retry( + self, + url: str, + session: aiohttp.ClientSession, + request_kwargs: Optional[dict] = None, + form_builder=None, + data_builder=None, + timeout: Optional[float] = None, + max_retries: Optional[int] = None, + ) -> WebhookResult: + """ + Execute request with exponential backoff and jitter. + + Args: + url: Target URL + session: aiohttp session + request_kwargs: Additional kwargs for session.post() + form_builder: Function that returns FormData (for multipart) + data_builder: Function that returns MultipartWriter (for streaming) + timeout: Optional request timeout (overrides default) + max_retries: Optional max retries (overrides default) + + Returns: + WebhookResult with success status and details + """ + last_error = None + start_time = time.time() + retries = 0 + kwargs = request_kwargs or {} + + # Use instance defaults if not specified + effective_max_retries = max_retries if max_retries is not None else self.max_retries + + # Add timeout to kwargs if specified + if timeout is not None: + kwargs["timeout"] = aiohttp.ClientTimeout(total=timeout) + + for attempt in range(effective_max_retries + 1): + try: + if attempt > 0: + delay = self._calculate_delay(attempt) + logger.info( + f"Webhook retry {attempt}/{effective_max_retries} after {delay:.2f}s" + ) + await asyncio.sleep(delay) + retries = attempt + + # Build request data + post_kwargs = dict(kwargs) + if form_builder: + post_kwargs["data"] = form_builder() + elif data_builder: + post_kwargs["data"] = data_builder() + + # Make the request + response = await session.post(url, **post_kwargs) + try: + elapsed_ms = (time.time() - start_time) * 1000 + + if response.status < 400: + return WebhookResult( + success=True, + status=response.status, + elapsed_ms=elapsed_ms, + retries=retries, + ) + + # Check if retryable + if response.status in self.RETRYABLE_STATUS_CODES: + last_error = f"HTTP {response.status}" + logger.warning( + f"Webhook request failed with retryable status {response.status}" + ) + continue + + # Non-retryable error + body = await response.text() + return WebhookResult( + success=False, + status=response.status, + error=f"HTTP {response.status}: {body[:500]}", + elapsed_ms=elapsed_ms, + retries=retries, + ) + finally: + response.close() + + except aiohttp.ClientError as e: + last_error = f"Client error: {str(e)}" + logger.warning(f"Webhook client error: {e}") + continue + except asyncio.TimeoutError: + last_error = "Request timeout" + logger.warning("Webhook request timed out") + continue + except Exception as e: + last_error = f"Unexpected error: {str(e)}" + logger.error(f"Webhook unexpected error: {e}", exc_info=True) + continue + + elapsed_ms = (time.time() - start_time) * 1000 + return WebhookResult( + success=False, + error=last_error, + elapsed_ms=elapsed_ms, + retries=retries, + ) + + def _calculate_delay(self, attempt: int) -> float: + """ + Calculate retry delay with exponential backoff and jitter. + + Args: + attempt: Current attempt number (1-based) + + Returns: + Delay in seconds + """ + # Exponential backoff + delay = min(self.base_delay * (2 ** (attempt - 1)), self.max_delay) + # Add jitter + jitter = delay * self.jitter_factor * (random.random() * 2 - 1) + return max(0.1, delay + jitter) + + async def close(self): + """Close the HTTP session and release resources.""" + if self._session and not self._session.closed: + await self._session.close() + self._session = None + logger.debug("Webhook client session closed") + + +async def close_webhook_client(): + """Close the global webhook client instance.""" + global _webhook_client + if _webhook_client: + await _webhook_client.close() + _webhook_client = None diff --git a/core/collector.py b/core/collector.py new file mode 100644 index 0000000..7ca44a4 --- /dev/null +++ b/core/collector.py @@ -0,0 +1,237 @@ +""" +Output collector for webhook delivery. + +IMPORTANT: This collector does NOT re-encode outputs. +It reads the actual files saved by ComfyUI output nodes and sends them as-is. + +- User saves as WebP → WebP is sent +- User saves as FLAC → FLAC is sent +- User saves as GIF → GIF is sent + +Zero information loss, zero re-encoding overhead. +""" + +import logging +import os +from typing import Optional + +from .types import OutputInfo +from ..utils.mime import guess_mime_type +from ..utils.files import get_output_path, get_image_dimensions + +logger = logging.getLogger(__name__) + + +async def collect_outputs( + outputs: dict, + meta: Optional[dict] = None +) -> tuple[list[tuple[str, str, bytes, str]], list[dict]]: + """ + Collect all outputs from a workflow execution. + + This function reads the ACTUAL files saved by output nodes. + No re-encoding - files are sent exactly as saved. + + Args: + outputs: Dictionary mapping node_id -> UI output dict + e.g., {"9": {"images": [{"filename": "out.webp", ...}]}} + meta: Optional metadata dict mapping node_id -> node metadata + + Returns: + Tuple of: + - files: List of (field_name, filename, data, mime_type) tuples + - metadata: List of output info dicts for the metadata payload + """ + files = [] + metadata = [] + file_index = 0 + + for node_id, ui_output in outputs.items(): + if not ui_output or not isinstance(ui_output, dict): + continue + + # Get node type from metadata if available + node_type = None + if meta and node_id in meta: + node_meta = meta[node_id] + if isinstance(node_meta, dict): + node_type = node_meta.get("node_type") + + # Process all output keys generically + # ComfyUI uses: images, audio, video, gifs, 3d, latents, text, etc. + for output_key, items in ui_output.items(): + if not isinstance(items, list): + # Handle non-list outputs (e.g., text as tuple) + if output_key == "text": + text_content = _extract_text(items) + if text_content: + info = OutputInfo( + filename=f"text_{node_id}.txt", + mime_type="text/plain", + node_id=node_id, + node_type=node_type, + output_key=output_key, + file_size_bytes=len(text_content.encode("utf-8")), + ) + files.append(( + f"file_{file_index}", + info.filename, + text_content.encode("utf-8"), + info.mime_type, + )) + metadata.append(info.to_dict()) + file_index += 1 + continue + + # Process list of file references + for idx, item in enumerate(items): + if not isinstance(item, dict): + continue + + if "filename" not in item: + continue + + # This is a file reference + filename = item["filename"] + subfolder = item.get("subfolder", "") + folder_type = item.get("type", "output") + + # Get full path and read file + file_path = get_output_path(filename, subfolder, folder_type) + + if not os.path.exists(file_path): + logger.warning(f"Output file not found: {file_path}") + continue + + try: + with open(file_path, "rb") as f: + file_data = f.read() + except Exception as e: + logger.error(f"Failed to read output file {file_path}: {e}") + continue + + # Determine MIME type from actual filename + mime_type = guess_mime_type(filename) + + # Build output info + info = OutputInfo( + filename=filename, + mime_type=mime_type, + node_id=node_id, + node_type=node_type, + output_key=output_key, + file_size_bytes=len(file_data), + subfolder=subfolder if subfolder else None, + folder_type=folder_type, + ) + + # Add type-specific metadata + _enrich_output_info(info, file_path, file_data) + + # Add batch info if multiple items + if len(items) > 1: + info.batch_index = idx + info.batch_size = len(items) + + files.append(( + f"file_{file_index}", + filename, + file_data, + mime_type, + )) + metadata.append(info.to_dict()) + file_index += 1 + + logger.debug(f"Collected {len(files)} output files") + return files, metadata + + +def _extract_text(value) -> Optional[str]: + """Extract text content from various formats.""" + if isinstance(value, str): + return value + if isinstance(value, tuple) and len(value) > 0: + return str(value[0]) + if isinstance(value, list) and len(value) > 0: + return str(value[0]) + return None + + +def _enrich_output_info(info: OutputInfo, file_path: str, file_data: bytes): + """ + Add type-specific metadata to OutputInfo. + + This reads metadata from the file without modifying it. + """ + # Extract format from filename + if "." in info.filename: + info.format = info.filename.rsplit(".", 1)[1].lower() + + # Image dimensions + if info.mime_type.startswith("image/"): + try: + width, height = get_image_dimensions(file_path) + if width and height: + info.width = width + info.height = height + except Exception as e: + logger.debug(f"Could not get image dimensions: {e}") + + # Audio metadata + elif info.mime_type.startswith("audio/"): + try: + _extract_audio_metadata(info, file_data) + except Exception as e: + logger.debug(f"Could not get audio metadata: {e}") + + # Video metadata + elif info.mime_type.startswith("video/"): + try: + _extract_video_metadata(info, file_path) + except Exception as e: + logger.debug(f"Could not get video metadata: {e}") + + +def _extract_audio_metadata(info: OutputInfo, file_data: bytes): + """Extract audio metadata without decoding the full file.""" + # Simple detection based on file headers + if file_data[:4] == b"fLaC": + # FLAC - could parse header for sample rate + pass + elif file_data[:4] == b"RIFF": + # WAV - parse header + try: + import struct + if len(file_data) >= 44: + # Channels at offset 22 (2 bytes) + # Sample rate at offset 24 (4 bytes) + sample_rate = struct.unpack(" None: + """ + Register a webhook context for tracking. + + Args: + context: The WebhookContext to register + """ + with _lock: + if context.prompt_id: + _contexts[context.prompt_id] = context + context.start_time = datetime.now() + logger.debug( + f"Registered webhook context: prompt={context.prompt_id}, " + f"request={context.request_id}" + ) + else: + logger.warning("Cannot register context without prompt_id") + + +def get_webhook_context(prompt_id: str) -> Optional[WebhookContext]: + """ + Get a context by prompt ID. + + Args: + prompt_id: The ComfyUI prompt ID + + Returns: + The WebhookContext if found, None otherwise + """ + with _lock: + return _contexts.get(prompt_id) + + +def unregister_webhook_context(prompt_id: str) -> Optional[WebhookContext]: + """ + Remove and return a context from the registry. + + Args: + prompt_id: The prompt ID to unregister + + Returns: + The removed WebhookContext if found, None otherwise + """ + with _lock: + context = _contexts.pop(prompt_id, None) + if context: + logger.debug(f"Unregistered webhook context: {prompt_id}") + return context + + +def get_all_contexts() -> Dict[str, WebhookContext]: + """ + Get a copy of all registered contexts. + Useful for debugging and monitoring. + + Returns: + Dictionary mapping prompt_id to WebhookContext + """ + with _lock: + return dict(_contexts) + + +def clear_all_contexts() -> int: + """ + Clear all registered contexts. + Useful for testing and cleanup. + + Returns: + Number of contexts that were cleared + """ + with _lock: + count = len(_contexts) + _contexts.clear() + logger.debug(f"Cleared {count} webhook contexts") + return count + + +def cleanup_old_contexts(max_age_hours: int = 24) -> int: + """ + Remove old completed contexts. + + Args: + max_age_hours: Remove contexts older than this + + Returns: + Number of contexts removed + """ + cutoff = datetime.now() - timedelta(hours=max_age_hours) + count = 0 + + with _lock: + to_remove = [] + + for prompt_id, context in _contexts.items(): + # Clean up old contexts based on completed_at or start_time + if context.completed_at and context.completed_at < cutoff: + to_remove.append(prompt_id) + elif context.start_time and context.start_time < cutoff: + to_remove.append(prompt_id) + + for prompt_id in to_remove: + _contexts.pop(prompt_id, None) + count += 1 + + if count > 0: + logger.info(f"Cleaned up {count} old webhook contexts") + + return count + + +# Legacy function aliases for compatibility +def register_context(context: WebhookContext) -> None: + """Legacy alias for register_webhook_context.""" + register_webhook_context(context) + + +def get_context(request_id: str) -> Optional[WebhookContext]: + """ + Get a context by request ID (searches all contexts). + + Args: + request_id: The request ID to look up + + Returns: + The WebhookContext if found, None otherwise + """ + with _lock: + for context in _contexts.values(): + if context.request_id == request_id: + return context + return None + + +def get_context_by_prompt(prompt_id: str) -> Optional[WebhookContext]: + """Legacy alias for get_webhook_context.""" + return get_webhook_context(prompt_id) + + +def update_context(request_id: str, **kwargs) -> bool: + """ + Update fields on a context by request ID. + + Args: + request_id: The request ID to update + **kwargs: Fields to update + + Returns: + True if context was found and updated + """ + with _lock: + for context in _contexts.values(): + if context.request_id == request_id: + for key, value in kwargs.items(): + if hasattr(context, key): + setattr(context, key, value) + return True + return False + + +def unregister_context(request_id: str) -> Optional[WebhookContext]: + """ + Remove a context by request ID. + + Args: + request_id: The request ID to unregister + + Returns: + The removed WebhookContext if found, None otherwise + """ + with _lock: + to_remove = None + for prompt_id, context in _contexts.items(): + if context.request_id == request_id: + to_remove = prompt_id + break + + if to_remove: + return _contexts.pop(to_remove, None) + return None diff --git a/core/interceptor.py b/core/interceptor.py new file mode 100644 index 0000000..a5dada8 --- /dev/null +++ b/core/interceptor.py @@ -0,0 +1,571 @@ +"""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}") diff --git a/core/types.py b/core/types.py new file mode 100644 index 0000000..ac09792 --- /dev/null +++ b/core/types.py @@ -0,0 +1,189 @@ +"""Core data types for webhook functionality.""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, Optional +from datetime import datetime +import uuid + + +class WebhookEvent(Enum): + """Types of webhook events.""" + STARTED = "workflow.started" + PROGRESS = "workflow.progress" + COMPLETED = "workflow.completed" + ERROR = "workflow.error" + INTERRUPTED = "workflow.interrupted" + + +@dataclass +class WebhookField: + """Individual field from webhook request.""" + name: str + value: Any + is_file: bool = False + filename: Optional[str] = None + content_type: Optional[str] = None + size_bytes: Optional[int] = None + + +@dataclass +class WebhookContext: + """ + Context for a webhook request. + + Tracks the entire lifecycle of a webhook request from submission + through execution and completion. + """ + # Callback configuration + callback_url: str = "" + auth_header: Optional[str] = None + auth_value: Optional[str] = None + timeout: int = 60 + max_retries: int = 3 + + # Identity + request_id: str = field(default_factory=lambda: str(uuid.uuid4())) + prompt_id: Optional[str] = None + + # Parsed webhook fields + fields: Dict[str, WebhookField] = field(default_factory=dict) + + # Raw inputs for fallback + inputs: dict = field(default_factory=dict) + + # Progress settings + send_progress: bool = True + progress_interval_ms: int = 500 + + # Request metadata + remote_ip: Optional[str] = None + content_type: Optional[str] = None + content_length: Optional[int] = None + user_agent: Optional[str] = None + + # Workflow settings + include_workflow_in_response: bool = False + + # Timestamps + timestamp: datetime = field(default_factory=datetime.now) + start_time: Optional[datetime] = None + completed_at: Optional[datetime] = None + + def get_auth_headers(self) -> dict: + """Get authentication headers for HTTP requests.""" + if self.auth_header and self.auth_value: + return {self.auth_header: self.auth_value} + return {} + + def censor_auth(self) -> str: + """Get censored auth header for logging.""" + if self.auth_header and self.auth_value: + value = self.auth_value + if len(value) > 12: + censored = value[:4] + "*" * (len(value) - 8) + value[-4:] + else: + censored = "*" * len(value) + return f"{self.auth_header}: {censored}" + return "(not set)" + + def has_field(self, name: str) -> bool: + """Check if a field exists.""" + return name in self.fields or name in self.inputs + + def get_field(self, name: str) -> Optional[WebhookField]: + """Get a WebhookField by name.""" + return self.fields.get(name) + + def get_field_value(self, name: str, default: Any = None) -> Any: + """Get a field's value.""" + if name in self.fields: + return self.fields[name].value + if name in self.inputs: + return self.inputs[name] + return default + + def get_field_names(self) -> list: + """Get all field names.""" + names = set(self.fields.keys()) + names.update(self.inputs.keys()) + return sorted(names) + + def to_dict(self) -> dict: + """Convert to dictionary for serialization.""" + return { + "request_id": self.request_id, + "prompt_id": self.prompt_id, + "callback_url": self.callback_url, + "fields": list(self.fields.keys()), + "timestamp": self.timestamp.isoformat() if self.timestamp else None, + } + + +@dataclass +class WebhookResult: + """Result of a webhook HTTP request.""" + success: bool + status: Optional[int] = None + error: Optional[str] = None + response_body: Optional[str] = None + elapsed_ms: float = 0.0 + retries: int = 0 + + +@dataclass +class OutputInfo: + """Information about a collected output file.""" + # Required fields + type: str = "UNKNOWN" + filename: str = "" + mime_type: str = "application/octet-stream" + + # File info + file_size_bytes: Optional[int] = None + + # Source info (for automatic collection) + node_id: Optional[str] = None + node_type: Optional[str] = None + output_key: Optional[str] = None + subfolder: Optional[str] = None + folder_type: str = "output" + + # Type-specific metadata + width: Optional[int] = None + height: Optional[int] = None + channels: Optional[int] = None + dtype: Optional[str] = None + format: Optional[str] = None + + # Audio/Video metadata + duration_seconds: Optional[float] = None + sample_rate: Optional[int] = None + frame_count: Optional[int] = None + fps: Optional[float] = None + + # Batch info + batch_index: Optional[int] = None + batch_size: Optional[int] = None + + # Inline content for small text + inline_content: Optional[str] = None + + def to_dict(self) -> dict: + """Convert to dictionary for metadata payload.""" + result = { + "type": self.type, + "filename": self.filename, + "mime_type": self.mime_type, + } + + # Add optional fields if present + for key in ["file_size_bytes", "node_id", "node_type", "output_key", + "width", "height", "channels", "dtype", "format", + "duration_seconds", "sample_rate", "frame_count", "fps", + "batch_index", "batch_size", "inline_content"]: + val = getattr(self, key) + if val is not None: + result[key] = val + + return result diff --git a/example_client.py b/example_client.py new file mode 100644 index 0000000..63ddf76 --- /dev/null +++ b/example_client.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +""" +Example client for comfy-workflow-api (Image-to-Image). + +Submits an image + prompt to ComfyUI and receives the transformed image via webhook. + +Usage: + python example_client.py [prompt] + +Examples: + python example_client.py input.png "make it look like a painting" + python example_client.py photo.jpg "add snow to the scene" + +Requirements: + pip install requests + +Make sure ComfyUI is running on http://localhost:8188 +""" + +import gzip +import json +import sys +import threading +from http.server import HTTPServer, BaseHTTPRequestHandler +from pathlib import Path + +import requests + +# Configuration +COMFYUI_URL = "http://localhost:8188" +CALLBACK_PORT = 5001 +OUTPUT_DIR = Path("./outputs") +WORKFLOW_FILE = Path(__file__).parent / "example_workflow.json" + + +class WebhookHandler(BaseHTTPRequestHandler): + """Handles webhook callbacks from ComfyUI.""" + + files_received = [] + done = threading.Event() + + def log_message(self, *args): + pass # Suppress logging + + def do_POST(self): + content_type = self.headers.get("Content-Type", "") + content_encoding = self.headers.get("Content-Encoding", "") + length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(length) + + # Decompress gzip if needed (WebhookSend uses gzip compression) + if "gzip" in content_encoding: + body = gzip.decompress(body) + + if "application/json" in content_type: + data = json.loads(body) + event = data.get("event") + + if event == "workflow.progress": + progress = data.get("progress", {}) + print(f" Progress: {progress.get('value')}/{progress.get('max')}") + elif event == "workflow.completed": + print(f" Completed in {data.get('execution_time_ms', 'N/A')}ms") + elif event == "workflow.error": + print(f" Error: {data.get('error')}") + WebhookHandler.done.set() + + elif "multipart" in content_type: + # Parse multipart and save files + boundary = content_type.split("boundary=")[1] + parts = body.split(f"--{boundary}".encode()) + + OUTPUT_DIR.mkdir(exist_ok=True) + + for part in parts: + if b"filename=" in part and b"\r\n\r\n" in part: + header, content = part.split(b"\r\n\r\n", 1) + filename = header.decode().split('filename="')[1].split('"')[0] + content = content.rstrip(b"\r\n--") + + path = OUTPUT_DIR / filename + path.write_bytes(content) + WebhookHandler.files_received.append(str(path)) + print(f" Saved: {path}") + + WebhookHandler.done.set() + + self.send_response(200) + self.end_headers() + self.wfile.write(b"OK") + + +def main(): + if len(sys.argv) < 2: + print("Usage: python example_client.py [prompt]") + print("Example: python example_client.py input.png \"make it look like a painting\"") + sys.exit(1) + + image_path = Path(sys.argv[1]) + if not image_path.exists(): + print(f"Error: Image not found: {image_path}") + sys.exit(1) + + prompt = " ".join(sys.argv[2:]) if len(sys.argv) > 2 else "enhance the image" + + print(f"Image: {image_path}") + print(f"Prompt: {prompt}") + print(f"Output: {OUTPUT_DIR.absolute()}\n") + + # Load workflow + workflow = json.loads(WORKFLOW_FILE.read_text()) + + # Load image + image_data = image_path.read_bytes() + image_mime = "image/png" if image_path.suffix.lower() == ".png" else "image/jpeg" + + # Start webhook server + server = HTTPServer(("localhost", CALLBACK_PORT), WebhookHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + + callback_url = f"http://localhost:{CALLBACK_PORT}/webhook" + + # Submit workflow with image and prompt + print("Submitting workflow...") + response = requests.post( + f"{COMFYUI_URL}/api/webhook", + files={ + "workflow": (None, json.dumps(workflow), "application/json"), + "callback_url": (None, callback_url), + "prompt": (None, prompt), + "image": (image_path.name, image_data, image_mime), + } + ) + + result = response.json() + print(f" Request ID: {result.get('request_id')}") + print(f" Prompt ID: {result.get('prompt_id')}\n") + + # Wait for completion + print("Processing...") + WebhookHandler.done.wait(timeout=300) + + server.shutdown() + + if WebhookHandler.files_received: + print(f"\nDone! Files saved to {OUTPUT_DIR.absolute()}") + else: + print("\nNo files received.") + + +if __name__ == "__main__": + main() diff --git a/example_workflow.json b/example_workflow.json new file mode 100644 index 0000000..627ac76 --- /dev/null +++ b/example_workflow.json @@ -0,0 +1,95 @@ +{ + "3": { + "inputs": { + "seed": 798343486140665, + "steps": 20, + "cfg": 7, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 0.7, + "model": ["4", 0], + "positive": ["6", 0], + "negative": ["7", 0], + "latent_image": ["5", 0] + }, + "class_type": "KSampler", + "_meta": {"title": "KSampler"} + }, + "4": { + "inputs": { + "ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": {"title": "Load Checkpoint"} + }, + "5": { + "inputs": { + "pixels": ["13", 0], + "vae": ["4", 2] + }, + "class_type": "VAEEncode", + "_meta": {"title": "VAE Encode"} + }, + "6": { + "inputs": { + "text": ["11", 0], + "clip": ["4", 1] + }, + "class_type": "CLIPTextEncode", + "_meta": {"title": "CLIP Text Encode (Prompt)"} + }, + "7": { + "inputs": { + "text": "text, watermark, blurry, low quality", + "clip": ["4", 1] + }, + "class_type": "CLIPTextEncode", + "_meta": {"title": "CLIP Text Encode (Negative)"} + }, + "8": { + "inputs": { + "samples": ["3", 0], + "vae": ["4", 2] + }, + "class_type": "VAEDecode", + "_meta": {"title": "VAE Decode"} + }, + "10": { + "inputs": { + "default_callback_url": "", + "default_timeout": 60, + "default_max_retries": 3, + "debug": true + }, + "class_type": "WebhookReceiver", + "_meta": {"title": "Webhook Receiver"} + }, + "11": { + "inputs": { + "field": "prompt", + "debug": true, + "webhook_context": ["10", 0] + }, + "class_type": "WebhookTransformer", + "_meta": {"title": "Extract Prompt"} + }, + "12": { + "inputs": { + "field_1_name": "output", + "debug": false, + "webhook_context": ["10", 0], + "field_1": ["8", 0] + }, + "class_type": "WebhookSend", + "_meta": {"title": "Webhook Send"} + }, + "13": { + "inputs": { + "field": "image", + "debug": true, + "webhook_context": ["10", 0] + }, + "class_type": "WebhookTransformer", + "_meta": {"title": "Extract Image"} + } +} diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..ce73547 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,11 @@ +"""ComfyUI nodes for webhook integration.""" + +from .receiver import WebhookReceiver +from .transformer import WebhookTransformer +from .send import WebhookSend + +__all__ = [ + "WebhookReceiver", + "WebhookTransformer", + "WebhookSend", +] diff --git a/nodes/receiver.py b/nodes/receiver.py new file mode 100644 index 0000000..e70887b --- /dev/null +++ b/nodes/receiver.py @@ -0,0 +1,258 @@ +"""WebhookReceiver node - Entry point for webhook-triggered workflows.""" + +import asyncio +import logging +import uuid +from datetime import datetime +from typing import Any, Dict + +from ..core.types import WebhookContext, WebhookField +from ..core.context import register_webhook_context +from ..core.interceptor import send_started_event +from ..utils.helpers import format_size + +logger = logging.getLogger(__name__) + + +class WebhookReceiver: + """ + Entry point for webhook-triggered workflows. + + Extracts webhook configuration from extra_pnginfo (passed from /api/webhook) + and creates a WebhookContext for use by downstream nodes. + + Configuration comes from webhook_config in extra_pnginfo: + - callback_url: Where to send results + - auth_header: Auth header name (e.g., "Authorization") + - auth_value: Auth header value (e.g., "Bearer token123") + - request_id: Unique request identifier + - timeout: HTTP timeout + - max_retries: Retry attempts + + Fields/uploads come from webhook_inputs in extra_pnginfo. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": { + "default_callback_url": ("STRING", { + "default": "", + "tooltip": "Fallback callback URL if not provided in request" + }), + "default_timeout": ("INT", { + "default": 60, + "min": 1, + "max": 300, + "tooltip": "HTTP request timeout in seconds" + }), + "default_max_retries": ("INT", { + "default": 3, + "min": 0, + "max": 10, + "tooltip": "Number of retry attempts on failure" + }), + "debug": ("BOOLEAN", { + "default": False, + "tooltip": "Enable debug logging" + }), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = ("WEBHOOK_CONTEXT",) + RETURN_NAMES = ("webhook_context",) + FUNCTION = "receive" + CATEGORY = "webhook" + DESCRIPTION = "Entry point for webhook-triggered workflows. Extracts webhook config and creates context for downstream nodes." + + @classmethod + def IS_CHANGED(cls, **kwargs): + """Always return a unique value to prevent caching.""" + # Webhook requests should never be cached as each request has unique data + import uuid + return str(uuid.uuid4()) + + def receive( + self, + prompt=None, + extra_pnginfo=None, + unique_id=None, + default_callback_url="", + default_timeout=60, + default_max_retries=3, + debug=False, + ): + """ + Extract webhook configuration and create context. + + The context is registered for output interception and passed + to WebhookTransformer nodes for field extraction. + """ + # Extract webhook_config and webhook_inputs from extra_data + extra_data = extra_pnginfo or {} + webhook_config = extra_data.get("webhook_config", {}) + webhook_inputs = extra_data.get("webhook_inputs", {}) + + # Extract configuration from webhook_config + callback_url = webhook_config.get("callback_url") or default_callback_url + auth_header = webhook_config.get("auth_header") + auth_value = webhook_config.get("auth_value") + request_id = webhook_config.get("request_id") or str(uuid.uuid4()) + timeout = webhook_config.get("timeout") or default_timeout + max_retries = webhook_config.get("max_retries") or default_max_retries + + # Extract request metadata + remote_ip = webhook_config.get("remote_ip", "unknown") + content_type = webhook_config.get("content_type", "multipart/form-data") + content_length = webhook_config.get("content_length") + user_agent = webhook_config.get("user_agent", "unknown") + + # Build fields dictionary from webhook_inputs + fields: Dict[str, WebhookField] = {} + for name, value in webhook_inputs.items(): + # Skip reserved fields that start with _ + if name.startswith("_"): + continue + + # Detect if this is a file upload (has special structure) + is_file = False + filename = None + content_type_field = None + size_bytes = None + + if isinstance(value, dict) and value.get("_is_file"): + # This is a file upload - load from saved path + is_file = True + filename = value.get("_filename") + content_type_field = value.get("_content_type") + size_bytes = value.get("_size_bytes", 0) + saved_path = value.get("_saved_path") + + # Load file data from disk + if saved_path: + try: + from pathlib import Path + value = Path(saved_path).read_bytes() + except Exception as e: + logger.warning(f"Failed to load file from {saved_path}: {e}") + value = b"" + else: + value = b"" + elif isinstance(value, bytes): + # Raw bytes = file + is_file = True + size_bytes = len(value) + elif isinstance(value, str): + size_bytes = len(value.encode("utf-8")) + + fields[name] = WebhookField( + name=name, + value=value, + is_file=is_file, + filename=filename, + content_type=content_type_field, + size_bytes=size_bytes, + ) + + # Create context + context = WebhookContext( + callback_url=callback_url, + request_id=request_id, + auth_header=auth_header, + auth_value=auth_value, + timeout=timeout, + max_retries=max_retries, + fields=fields, + remote_ip=remote_ip, + content_type=content_type, + content_length=content_length, + user_agent=user_agent, + prompt_id=unique_id, + inputs=webhook_inputs, # Raw inputs for fallback + send_progress=webhook_config.get("send_progress", True), + progress_interval_ms=webhook_config.get("progress_interval_ms", 500), + include_workflow_in_response=webhook_config.get("include_workflow_in_response", False), + ) + + # Debug output + if debug: + self._print_debug(context, fields) + + # Register context for output interception + if context.prompt_id: + register_webhook_context(context) + if not debug: + logger.info( + f"WebhookReceiver: Registered context for prompt {context.prompt_id}, " + f"callback: {callback_url or '(none)'}, " + f"fields: {list(fields.keys())}" + ) + + # Send started event if callback URL is configured + if context.callback_url: + try: + # Try to get running loop (Python 3.10+ safe) + loop = asyncio.get_running_loop() + asyncio.create_task(send_started_event(context)) + except RuntimeError: + # No running loop - run synchronously + try: + asyncio.run(send_started_event(context)) + except RuntimeError: + # Fallback for nested event loops + from ..utils.helpers import run_async + run_async(send_started_event(context)) + except Exception as e: + logger.warning(f"Failed to send started event: {e}") + + return (context,) + + def _print_debug(self, context: WebhookContext, fields: Dict[str, WebhookField]): + """Print debug information about the incoming request.""" + timestamp = context.timestamp.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" + + lines = [ + "", + "[Webhook Receiver] Incoming Request", + "=" * 50, + f" Remote IP: {context.remote_ip or 'unknown'}", + f" Timestamp: {timestamp}", + f" Request ID: {context.request_id}", + f" Content-Type: {context.content_type or 'unknown'}", + f" Content-Length: {format_size(context.content_length)}", + f" User-Agent: {context.user_agent or 'unknown'}", + "", + " Configuration:", + f" Callback URL: {context.callback_url or '(none)'}", + f" Auth Header: {context.censor_auth()}", + f" Timeout: {context.timeout}s", + f" Max Retries: {context.max_retries}", + "", + f" Fields Received ({len(fields)}):", + ] + + if fields: + field_names = sorted(fields.keys()) + for i, name in enumerate(field_names): + field = fields[name] + prefix = " +-" if i == len(field_names) - 1 else " |-" + + if field.is_file: + size_str = format_size(field.size_bytes) + ct = field.content_type or "application/octet-stream" + lines.append(f"{prefix} {name:20} (file) {ct}, {size_str}") + else: + size_str = format_size(field.size_bytes) + lines.append(f"{prefix} {name:20} (text) {size_str}") + else: + lines.append(" (no fields)") + + lines.append("=" * 50) + + print("\n".join(lines)) diff --git a/nodes/send.py b/nodes/send.py new file mode 100644 index 0000000..be85b77 --- /dev/null +++ b/nodes/send.py @@ -0,0 +1,774 @@ +"""WebhookSend node - Send workflow outputs to callback URL.""" + +import io +import json +import logging +import struct +from datetime import datetime +from typing import Optional, Any, Dict, List, Tuple, Union + +import numpy as np +import torch +from PIL import Image + +from ..core.types import WebhookContext, WebhookResult, OutputInfo +from ..core.client import get_webhook_client +from ..utils.helpers import format_size, run_async + +logger = logging.getLogger(__name__) + + +class WebhookSend: + """ + Send workflow outputs to the callback URL as multipart/form-data. + + Supports 5 customizable wildcard fields. Each field accepts any + ComfyUI type (IMAGE, AUDIO, VIDEO, MESH, LATENT, MASK, STRING, etc.) + and automatically encodes it appropriately. + + Field Encoding: + - IMAGE: PNG file + - AUDIO: FLAC file + - VIDEO: MP4 file + - MESH: GLB file + - LATENT: Safetensors file + - MASK: PNG file (grayscale) + - STRING: Text part + - INT/FLOAT/BOOLEAN: Text part + + Always includes a metadata JSON field with request info. + Payload is always gzip compressed (Content-Encoding: gzip). + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "webhook_context": ("WEBHOOK_CONTEXT",), + }, + "optional": { + # Field 1 + "field_1": ("*", {"tooltip": "First output (any type)"}), + "field_1_name": ("STRING", { + "default": "output_1", + "tooltip": "Field name for first output" + }), + # Field 2 + "field_2": ("*", {"tooltip": "Second output (any type)"}), + "field_2_name": ("STRING", { + "default": "output_2", + "tooltip": "Field name for second output" + }), + # Field 3 + "field_3": ("*", {"tooltip": "Third output (any type)"}), + "field_3_name": ("STRING", { + "default": "output_3", + "tooltip": "Field name for third output" + }), + # Field 4 + "field_4": ("*", {"tooltip": "Fourth output (any type)"}), + "field_4_name": ("STRING", { + "default": "output_4", + "tooltip": "Field name for fourth output" + }), + # Field 5 + "field_5": ("*", {"tooltip": "Fifth output (any type)"}), + "field_5_name": ("STRING", { + "default": "output_5", + "tooltip": "Field name for fifth output" + }), + # Debug + "debug": ("BOOLEAN", { + "default": False, + "tooltip": "Enable debug logging" + }), + # Batch outputs + "batch_outputs": ("BOOLEAN", { + "default": False, + "tooltip": "Send all items in batched outputs (images/masks). Field names become prefixes with _0, _1, etc. suffixes." + }), + }, + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "send" + CATEGORY = "webhook" + DESCRIPTION = "Send workflow outputs to callback URL as gzip-compressed multipart/form-data. Supports 5 customizable wildcard fields." + + def send( + self, + webhook_context, + field_1=None, + field_1_name="output_1", + field_2=None, + field_2_name="output_2", + field_3=None, + field_3_name="output_3", + field_4=None, + field_4_name="output_4", + field_5=None, + field_5_name="output_5", + debug=False, + batch_outputs=False, + ): + if not isinstance(webhook_context, WebhookContext): + logger.warning("WebhookSend: Invalid webhook context") + return {} + + if not webhook_context.callback_url: + logger.warning("WebhookSend: No callback URL configured") + return {} + + # Collect non-None fields + fields_to_send = [] + for field, name in [ + (field_1, field_1_name), + (field_2, field_2_name), + (field_3, field_3_name), + (field_4, field_4_name), + (field_5, field_5_name), + ]: + if field is not None: + fields_to_send.append((name, field)) + + try: + result = run_async( + self._send_async(webhook_context, fields_to_send, debug, batch_outputs) + ) + return {} + except Exception as e: + logger.error(f"WebhookSend failed: {e}", exc_info=True) + if debug: + self._print_error(webhook_context, str(e)) + return {} + + async def _send_async( + self, + context: WebhookContext, + fields_to_send: List[Tuple[str, Any]], + debug: bool, + batch_outputs: bool = False, + ): + """Async implementation of send.""" + files = [] + outputs = [] + start_time = datetime.now() + + # Process each field + for name, value in fields_to_send: + encoded = self._encode_value(name, value, batch_outputs) + if encoded: + # Check if we got a list of encoded items (batched) or single item + if isinstance(encoded, list): + # Batched output - expand with _N suffixes + for i, (filename, data, mime_type, output_info) in enumerate(encoded): + batch_name = f"{name}_{i}" + files.append((batch_name, filename, data, mime_type)) + outputs.append(output_info.to_dict()) + else: + # Single item + filename, data, mime_type, output_info = encoded + files.append((name, filename, data, mime_type)) + outputs.append(output_info.to_dict()) + + # Build metadata + metadata = { + "request_id": context.request_id, + "prompt_id": context.prompt_id, + "timestamp": datetime.now().isoformat(), + "fields": outputs, + } + + # Send request + client = get_webhook_client() + + result = await client.send_multipart( + url=context.callback_url, + metadata=metadata, + files=files, + headers=context.get_auth_headers(), + timeout=context.timeout, + max_retries=context.max_retries, + compress=True, # Always use gzip compression + ) + + elapsed_ms = (datetime.now() - start_time).total_seconds() * 1000 + + if debug: + self._print_debug(context, files, outputs, result, elapsed_ms) + + if not result.success: + logger.error(f"WebhookSend failed: {result.error}") + else: + logger.info( + f"WebhookSend: Sent {len(files)} fields to {context.callback_url}" + ) + + def _encode_value( + self, + name: str, + value: Any, + batch_outputs: bool = False, + ) -> Optional[Union[Tuple[str, bytes, str, OutputInfo], List[Tuple[str, bytes, str, OutputInfo]]]]: + """ + Encode value to appropriate format for multipart upload. + + Returns: + tuple: (filename, data_bytes, mime_type, output_info) or None + list: List of tuples when batch_outputs=True and value is batched + """ + # IMAGE tensor + if torch.is_tensor(value) and len(value.shape) == 4: + # Batched images [B, H, W, C] + batch_size = value.shape[0] + + if batch_outputs and batch_size > 1: + # Return all images as separate items + results = [] + for i in range(batch_size): + img_bytes = self._tensor_to_png(value[i]) + filename = f"{name}_{i}.png" + output_info = OutputInfo( + type="IMAGE", + filename=filename, + mime_type="image/png", + file_size_bytes=len(img_bytes), + width=value.shape[2], + height=value.shape[1], + channels=value.shape[3] if len(value.shape) > 3 else 3, + dtype="float32", + format="png", + batch_index=i, + batch_size=batch_size, + ) + results.append((filename, img_bytes, "image/png", output_info)) + return results + else: + # Send only first image (original behavior) + data = self._tensor_to_png(value[0]) + filename = f"{name}.png" + + output_info = OutputInfo( + type="IMAGE", + filename=filename, + mime_type="image/png", + file_size_bytes=len(data), + width=value.shape[2], + height=value.shape[1], + channels=value.shape[3] if len(value.shape) > 3 else 3, + dtype="float32", + format="png", + batch_index=0, + batch_size=batch_size, + ) + return (filename, data, "image/png", output_info) + + # MASK tensor [B, H, W] or [H, W] - check before IMAGE to avoid confusion + # Masks are 2D or 3D with small batch dimension (no color channel) + if torch.is_tensor(value) and len(value.shape) in (2, 3): + # Check if this looks like a mask: + # - 2D tensor [H, W] is definitely a mask + # - 3D tensor [B, H, W] where B < 10 is a batched mask + # - 3D tensor [H, W, C] where C is 3 or 4 is an image + is_mask = False + is_batched_mask = False + batch_size = 1 + + if len(value.shape) == 2: + is_mask = True + elif len(value.shape) == 3: + # [B, H, W] mask vs [H, W, C] image + # If last dim is 1, 3, or 4, it's likely an image + # If first dim is small and last dim is large, it's likely a mask batch + if value.shape[-1] not in (1, 3, 4) and value.shape[0] < 10: + is_mask = True + is_batched_mask = value.shape[0] > 1 + batch_size = value.shape[0] + elif value.shape[0] == 1 and value.shape[-1] > 10: + # [1, H, W] is a batched mask + is_mask = True + is_batched_mask = False + batch_size = 1 + + if is_mask: + if batch_outputs and is_batched_mask and batch_size > 1: + # Return all masks as separate items + results = [] + for i in range(batch_size): + mask = value[i] + mask_bytes = self._mask_to_png(mask) + filename = f"{name}_{i}.png" + output_info = OutputInfo( + type="MASK", + filename=filename, + mime_type="image/png", + file_size_bytes=len(mask_bytes), + width=mask.shape[1], + height=mask.shape[0], + format="png", + batch_index=i, + batch_size=batch_size, + ) + results.append((filename, mask_bytes, "image/png", output_info)) + return results + else: + # Send only first mask (original behavior) + mask = value[0] if len(value.shape) == 3 else value + data = self._mask_to_png(mask) + filename = f"{name}.png" + + output_info = OutputInfo( + type="MASK", + filename=filename, + mime_type="image/png", + file_size_bytes=len(data), + width=mask.shape[1], + height=mask.shape[0], + format="png", + batch_index=0, + batch_size=batch_size, + ) + return (filename, data, "image/png", output_info) + + # IMAGE tensor (single) [H, W, C] + if torch.is_tensor(value) and len(value.shape) == 3: + data = self._tensor_to_png(value) + filename = f"{name}.png" + + output_info = OutputInfo( + type="IMAGE", + filename=filename, + mime_type="image/png", + file_size_bytes=len(data), + width=value.shape[1], + height=value.shape[0], + channels=value.shape[2] if len(value.shape) > 2 else 3, + dtype="float32", + format="png", + ) + return (filename, data, "image/png", output_info) + + # AUDIO dict + if isinstance(value, dict) and "waveform" in value: + data = self._audio_to_flac(value) + if data: + filename = f"{name}.flac" + output_info = OutputInfo( + type="AUDIO", + filename=filename, + mime_type="audio/flac", + file_size_bytes=len(data), + sample_rate=value.get("sample_rate", 44100), + format="flac", + ) + return (filename, data, "audio/flac", output_info) + + # VIDEO dict or tensor + if isinstance(value, dict) and "images" in value: + data = self._video_to_mp4(value) + if data: + filename = f"{name}.mp4" + fps = value.get("fps", 24) + images = value.get("images") + output_info = OutputInfo( + type="VIDEO", + filename=filename, + mime_type="video/mp4", + file_size_bytes=len(data), + width=images.shape[2] if torch.is_tensor(images) else None, + height=images.shape[1] if torch.is_tensor(images) else None, + format="mp4", + ) + return (filename, data, "video/mp4", output_info) + + # MESH dict + if isinstance(value, dict) and ("vertices" in value or "faces" in value): + data = self._mesh_to_glb(value) + if data: + filename = f"{name}.glb" + vertices = value.get("vertices") + faces = value.get("faces") + output_info = OutputInfo( + type="MESH", + filename=filename, + mime_type="model/gltf-binary", + file_size_bytes=len(data), + format="glb", + ) + return (filename, data, "model/gltf-binary", output_info) + + # LATENT dict + if isinstance(value, dict) and "samples" in value: + data = self._latent_to_safetensors(value) + if data: + filename = f"{name}.safetensors" + samples = value.get("samples") + output_info = OutputInfo( + type="LATENT", + filename=filename, + mime_type="application/octet-stream", + file_size_bytes=len(data), + format="safetensors", + ) + return (filename, data, "application/octet-stream", output_info) + + # STRING + if isinstance(value, str): + data = value.encode("utf-8") + filename = f"{name}.txt" + output_info = OutputInfo( + type="STRING", + filename=filename, + mime_type="text/plain", + file_size_bytes=len(data), + inline_content=value if len(value) < 1000 else None, + ) + return (filename, data, "text/plain", output_info) + + # INT, FLOAT, BOOLEAN + if isinstance(value, (int, float, bool)): + str_value = str(value).lower() if isinstance(value, bool) else str(value) + data = str_value.encode("utf-8") + filename = f"{name}.txt" + type_name = type(value).__name__.upper() + output_info = OutputInfo( + type=type_name, + filename=filename, + mime_type="text/plain", + file_size_bytes=len(data), + inline_content=str_value, + ) + return (filename, data, "text/plain", output_info) + + # BYTES + if isinstance(value, bytes): + filename = f"{name}.bin" + output_info = OutputInfo( + type="BYTES", + filename=filename, + mime_type="application/octet-stream", + file_size_bytes=len(value), + ) + return (filename, value, "application/octet-stream", output_info) + + # Try JSON serialization for other types + try: + json_str = json.dumps(value, default=str) + data = json_str.encode("utf-8") + filename = f"{name}.json" + output_info = OutputInfo( + type="JSON", + filename=filename, + mime_type="application/json", + file_size_bytes=len(data), + ) + return (filename, data, "application/json", output_info) + except Exception: + pass + + logger.warning(f"WebhookSend: Unable to encode field '{name}' of type {type(value)}") + return None + + def _tensor_to_png(self, tensor: torch.Tensor) -> bytes: + """Convert image tensor [H, W, C] to PNG bytes.""" + np_image = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8) + pil_image = Image.fromarray(np_image) + + buffer = io.BytesIO() + pil_image.save(buffer, format="PNG", compress_level=4) + return buffer.getvalue() + + def _mask_to_png(self, tensor: torch.Tensor) -> bytes: + """Convert mask tensor [H, W] to grayscale PNG bytes.""" + np_mask = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8) + pil_image = Image.fromarray(np_mask, mode='L') + + buffer = io.BytesIO() + pil_image.save(buffer, format="PNG") + return buffer.getvalue() + + def _audio_to_flac(self, audio: dict) -> Optional[bytes]: + """Convert audio dict to FLAC bytes.""" + try: + import av + except ImportError: + logger.warning("av library not available for audio encoding") + return None + + waveform = audio.get("waveform") + sample_rate = audio.get("sample_rate", 44100) + + if waveform is None: + return None + + if len(waveform.shape) == 3: + waveform = waveform[0] + + try: + buffer = io.BytesIO() + container = av.open(buffer, mode="w", format="flac") + + layout = "mono" if waveform.shape[0] == 1 else "stereo" + stream = container.add_stream("flac", rate=sample_rate, layout=layout) + + audio_np = waveform.cpu().numpy() + + frame = av.AudioFrame.from_ndarray( + audio_np.reshape(1, -1).astype(np.float32), + format="flt", + layout=layout, + ) + frame.sample_rate = sample_rate + + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(None): + container.mux(packet) + + container.close() + buffer.seek(0) + return buffer.getvalue() + + except Exception as e: + logger.error(f"Failed to encode audio: {e}") + return None + + def _video_to_mp4(self, video: dict) -> Optional[bytes]: + """Convert video dict to MP4 bytes.""" + try: + import av + except ImportError: + logger.warning("av library not available for video encoding") + return None + + images = video.get("images") + fps = video.get("fps", 24) + + if images is None: + return None + + try: + buffer = io.BytesIO() + container = av.open(buffer, mode="w", format="mp4") + stream = container.add_stream("libx264", rate=fps) + + if len(images.shape) == 4: + height, width = images.shape[1], images.shape[2] + else: + height, width = images.shape[0], images.shape[1] + + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + + if len(images.shape) == 4: + for i in range(images.shape[0]): + np_frame = (images[i].cpu().numpy() * 255).clip(0, 255).astype(np.uint8) + frame = av.VideoFrame.from_ndarray(np_frame, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + + for packet in stream.encode(None): + container.mux(packet) + + container.close() + buffer.seek(0) + return buffer.getvalue() + + except Exception as e: + logger.error(f"Failed to encode video: {e}") + return None + + def _mesh_to_glb(self, mesh: dict) -> Optional[bytes]: + """Convert mesh dict to GLB bytes.""" + try: + vertices = mesh.get("vertices") + faces = mesh.get("faces") + + if vertices is None or faces is None: + return None + + if torch.is_tensor(vertices): + vertices = vertices.cpu().numpy() + if torch.is_tensor(faces): + faces = faces.cpu().numpy() + + # Handle batch dimension + if len(vertices.shape) == 3: + vertices = vertices[0] + if len(faces.shape) == 3: + faces = faces[0] + + vertices = vertices.astype(np.float32) + faces = faces.astype(np.uint32) + + # Calculate bounding box + v_min = vertices.min(axis=0).tolist() + v_max = vertices.max(axis=0).tolist() + + # Create buffer data + vertex_data = vertices.tobytes() + index_data = faces.flatten().tobytes() + + def pad4(data): + remainder = len(data) % 4 + if remainder: + data += b'\x00' * (4 - remainder) + return data + + vertex_data = pad4(vertex_data) + index_data = pad4(index_data) + buffer_data = vertex_data + index_data + + # Create glTF JSON + gltf = { + "asset": {"version": "2.0", "generator": "ComfyUI-Webhook"}, + "scene": 0, + "scenes": [{"nodes": [0]}], + "nodes": [{"mesh": 0}], + "meshes": [{ + "primitives": [{ + "attributes": {"POSITION": 0}, + "indices": 1 + }] + }], + "accessors": [ + { + "bufferView": 0, + "componentType": 5126, + "count": len(vertices), + "type": "VEC3", + "min": v_min, + "max": v_max + }, + { + "bufferView": 1, + "componentType": 5125, + "count": faces.size, + "type": "SCALAR" + } + ], + "bufferViews": [ + { + "buffer": 0, + "byteOffset": 0, + "byteLength": len(vertex_data) + }, + { + "buffer": 0, + "byteOffset": len(vertex_data), + "byteLength": len(index_data) + } + ], + "buffers": [{"byteLength": len(buffer_data)}] + } + + json_data = json.dumps(gltf, separators=(',', ':')).encode('utf-8') + json_data = pad4(json_data) + + # Build GLB + glb = io.BytesIO() + glb.write(struct.pack('<4sII', b'glTF', 2, 12 + 8 + len(json_data) + 8 + len(buffer_data))) + glb.write(struct.pack(' Optional[bytes]: + try: + from safetensors.torch import save + samples = latent.get("samples") + if samples is None: + return None + return save({"samples": samples}) + except Exception as e: + logger.error(f"Failed to encode latent: {e}") + return None + + def _print_debug( + self, + context: WebhookContext, + files: List[Tuple[str, str, bytes, str]], + outputs: List[Dict], + result: WebhookResult, + elapsed_ms: float + ): + """Print debug information about the outgoing request.""" + lines = [ + "", + "[Webhook Send] Outgoing Request", + "=" * 50, + f" Target URL: {context.callback_url}", + f" Method: POST", + f" Content-Type: multipart/form-data", + f" Auth Header: {context.censor_auth()}", + f" Request ID: {context.request_id}", + "", + f" Fields ({len(files)}):", + ] + + for i, (name, filename, data, mime_type) in enumerate(files): + prefix = " +-" if i == len(files) - 1 else " |-" + + # Find output info for additional details + output = outputs[i] if i < len(outputs) else {} + type_name = output.get("type", "UNKNOWN") + + if type_name == "IMAGE": + w = output.get("width", "?") + h = output.get("height", "?") + lines.append(f"{prefix} {name:20} ({type_name}) {w}x{h}, PNG, {format_size(len(data))}") + elif type_name == "MASK": + w = output.get("width", "?") + h = output.get("height", "?") + lines.append(f"{prefix} {name:20} ({type_name}) {w}x{h}, PNG, {format_size(len(data))}") + elif type_name == "STRING": + content = output.get("inline_content") + if content and len(content) < 30: + lines.append(f"{prefix} {name:20} ({type_name}) \"{content}\"") + else: + lines.append(f"{prefix} {name:20} ({type_name}) {format_size(len(data))}") + else: + lines.append(f"{prefix} {name:20} ({type_name}) {format_size(len(data))}") + + lines.append("") + lines.append(" Response:") + if result.success: + lines.append(f" Status: {result.status} OK") + else: + lines.append(f" Status: {result.status or 'N/A'} {result.error or 'Failed'}") + + lines.append(f" Time: {elapsed_ms:.0f} ms") + + if result.response_body: + body = result.response_body + if len(body) > 100: + body = body[:97] + "..." + lines.append(f" Body: {body}") + + if result.retries > 0: + lines.append(f" Retries: {result.retries}") + + lines.append("=" * 50) + print("\n".join(lines)) + + def _print_error(self, context: WebhookContext, error: str): + """Print error information.""" + lines = [ + "", + "[Webhook Send] Request Failed", + "=" * 50, + f" Target URL: {context.callback_url}", + f" Request ID: {context.request_id}", + "", + f" Error: {error}", + "=" * 50, + ] + print("\n".join(lines)) diff --git a/nodes/transformer.py b/nodes/transformer.py new file mode 100644 index 0000000..aa77008 --- /dev/null +++ b/nodes/transformer.py @@ -0,0 +1,581 @@ +"""WebhookTransformer node - Extract a single atomic field from webhook context.""" + +import io +import logging +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import torch +from PIL import Image + +from ..core.types import WebhookContext, WebhookField +from ..processors.inputs import WebhookInputProcessor +from ..utils.helpers import format_size, run_async + +logger = logging.getLogger(__name__) + +_processor = None + + +def get_processor() -> WebhookInputProcessor: + global _processor + if _processor is None: + _processor = WebhookInputProcessor() + return _processor + + +class WebhookTransformer: + """ + Extract a single atomic field from the webhook request. + + Field names are atomic (flat) - no JSON paths, no nesting. + Outputs wildcard type (*) that connects to any ComfyUI input. + + Type Inference: + - Text fields with numbers -> INT or FLOAT + - Text fields -> STRING + - File uploads -> IMAGE, AUDIO, VIDEO, MESH, LATENT based on content type + - Boolean strings ("true"/"false") -> BOOLEAN + + Missing Field Handling: + - If field missing AND default_value provided -> use default + - If field missing AND NO default -> STOP workflow with error + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "webhook_context": ("WEBHOOK_CONTEXT",), + "field": ("STRING", { + "default": "prompt", + "tooltip": "Field name to extract (atomic, no nesting)" + }), + }, + "optional": { + "default_value": ("*", { + "tooltip": "Fallback value if field not found (leave empty to require field)" + }), + "debug": ("BOOLEAN", { + "default": False, + "tooltip": "Enable debug logging" + }), + }, + } + + RETURN_TYPES = ("*",) + RETURN_NAMES = ("value",) + FUNCTION = "transform" + CATEGORY = "webhook" + DESCRIPTION = "Extract a single atomic field from webhook request. Outputs wildcard type that connects to any input." + + @classmethod + def VALIDATE_INPUTS(cls, webhook_context=None, field="", default_value=None, debug=False): + """Validate inputs before execution.""" + # Basic validation - actual field checking happens in transform() + if not field or not field.strip(): + return "Field name cannot be empty" + return True + + def transform( + self, + webhook_context, + field: str, + default_value=None, + debug=False, + ): + """ + Extract field from webhook context. + + Raises: + ValueError: If field is missing and no default provided + """ + # Validate context (use duck typing for flexibility with module loading) + if not hasattr(webhook_context, 'get_field_value') or not hasattr(webhook_context, 'has_field'): + if default_value is not None: + if debug: + self._print_debug_default(field, default_value, "Invalid webhook context") + return (default_value,) + raise ValueError( + f"[Webhook Transformer] ERROR: Invalid webhook context for field '{field}'" + ) + + # Check if field exists + has_field = webhook_context.has_field(field) + + if not has_field: + # Field not found + if default_value is not None: + if debug: + self._print_debug_default(field, default_value, "Field not found") + return (default_value,) + + # No default - STOP workflow + available_fields = webhook_context.get_field_names() + self._print_error(field, available_fields) + raise ValueError( + f"[Webhook Transformer] Required field '{field}' not found. " + f"Available fields: {', '.join(available_fields) if available_fields else '(none)'}. " + f"Provide field or set default_value." + ) + + # Get field value + webhook_field = webhook_context.get_field(field) + raw_value = webhook_context.get_field_value(field) + + # Process and convert value based on type inference + processed_value, value_type, debug_info = self._process_value( + raw_value, webhook_field + ) + + if debug: + self._print_debug_found(field, processed_value, value_type, debug_info) + + return (processed_value,) + + def _process_value( + self, + raw_value: Any, + webhook_field: Optional[WebhookField] + ) -> Tuple[Any, str, Dict]: + """ + Process raw value and infer type. + + Returns: + tuple: (processed_value, type_name, debug_info_dict) + """ + debug_info = {} + + # Handle file uploads + if webhook_field and webhook_field.is_file: + return self._process_file(raw_value, webhook_field) + + # Handle string values + if isinstance(raw_value, str): + # Try to infer type from string content + + # Boolean + if raw_value.lower() in ("true", "false", "yes", "no", "1", "0"): + bool_value = raw_value.lower() in ("true", "yes", "1") + debug_info["original"] = raw_value + return (bool_value, "BOOLEAN", debug_info) + + # Integer + try: + if raw_value.isdigit() or (raw_value.startswith("-") and raw_value[1:].isdigit()): + int_value = int(raw_value) + debug_info["original"] = raw_value + return (int_value, "INT", debug_info) + except (ValueError, AttributeError): + pass + + # Float + try: + if "." in raw_value or "e" in raw_value.lower(): + float_value = float(raw_value) + debug_info["original"] = raw_value + return (float_value, "FLOAT", debug_info) + except (ValueError, AttributeError): + pass + + # Check if it's base64 encoded data + if raw_value.startswith("data:"): + return self._process_data_uri(raw_value) + + # Plain string + debug_info["length"] = len(raw_value) + return (raw_value, "STRING", debug_info) + + # Handle numeric types directly + if isinstance(raw_value, bool): + return (raw_value, "BOOLEAN", debug_info) + if isinstance(raw_value, int): + return (raw_value, "INT", debug_info) + if isinstance(raw_value, float): + return (raw_value, "FLOAT", debug_info) + + # Handle bytes (file data) + if isinstance(raw_value, bytes): + return self._process_bytes(raw_value) + + # Handle torch tensors + if torch.is_tensor(raw_value): + debug_info["shape"] = list(raw_value.shape) + debug_info["dtype"] = str(raw_value.dtype) + return (raw_value, "TENSOR", debug_info) + + # Handle dicts (could be LATENT, AUDIO, etc.) + if isinstance(raw_value, dict): + # Check for known dict types + if "samples" in raw_value: + return (raw_value, "LATENT", {"type": "latent_dict"}) + if "waveform" in raw_value: + return (raw_value, "AUDIO", {"type": "audio_dict"}) + + # Generic dict - return as is + debug_info["keys"] = list(raw_value.keys()) + return (raw_value, "DICT", debug_info) + + # Handle lists + if isinstance(raw_value, list): + debug_info["length"] = len(raw_value) + return (raw_value, "LIST", debug_info) + + # Default: return as-is + debug_info["python_type"] = type(raw_value).__name__ + return (raw_value, "ANY", debug_info) + + def _process_file( + self, + raw_value: Any, + webhook_field: WebhookField + ) -> Tuple[Any, str, Dict]: + """Process file upload based on content type.""" + content_type = webhook_field.content_type or "" + debug_info = { + "filename": webhook_field.filename, + "content_type": content_type, + "size_bytes": webhook_field.size_bytes, + } + + data = raw_value if isinstance(raw_value, bytes) else None + + if data is None: + return (raw_value, "UNKNOWN", debug_info) + + # Image types + if content_type.startswith("image/"): + try: + tensor, info = self._decode_image(data) + debug_info.update(info) + return (tensor, "IMAGE", debug_info) + except Exception as e: + logger.warning(f"Failed to decode image: {e}") + return (data, "BYTES", debug_info) + + # Audio types + if content_type.startswith("audio/"): + try: + audio_dict, info = self._decode_audio(data) + debug_info.update(info) + return (audio_dict, "AUDIO", debug_info) + except Exception as e: + logger.warning(f"Failed to decode audio: {e}") + return (data, "BYTES", debug_info) + + # Video types + if content_type.startswith("video/"): + try: + video, info = self._decode_video(data) + debug_info.update(info) + return (video, "VIDEO", debug_info) + except Exception as e: + logger.warning(f"Failed to decode video: {e}") + return (data, "BYTES", debug_info) + + # 3D mesh types + if content_type in ("model/gltf-binary", "model/obj", "model/stl", "application/octet-stream"): + # Try to detect from filename + filename = webhook_field.filename or "" + if filename.endswith((".glb", ".gltf", ".obj", ".stl")): + try: + mesh, info = self._decode_mesh(data, filename) + debug_info.update(info) + return (mesh, "MESH", debug_info) + except Exception as e: + logger.warning(f"Failed to decode mesh: {e}") + return (data, "BYTES", debug_info) + + # Safetensors (latent) + if webhook_field.filename and webhook_field.filename.endswith(".safetensors"): + try: + latent, info = self._decode_latent(data) + debug_info.update(info) + return (latent, "LATENT", debug_info) + except Exception as e: + logger.warning(f"Failed to decode latent: {e}") + return (data, "BYTES", debug_info) + + # Default: return raw bytes + return (data, "BYTES", debug_info) + + def _process_data_uri(self, data_uri: str) -> Tuple[Any, str, Dict]: + """Process data URI (base64 encoded data).""" + processor = get_processor() + debug_info = {} + + try: + # Detect type from data URI + if data_uri.startswith("data:image/"): + tensor = run_async(processor.process_image(data_uri)) + debug_info["source"] = "base64" + debug_info["dimensions"] = f"{tensor.shape[2]}x{tensor.shape[1]}" + return (tensor, "IMAGE", debug_info) + + elif data_uri.startswith("data:audio/"): + audio = run_async(processor.process_audio(data_uri)) + debug_info["source"] = "base64" + return (audio, "AUDIO", debug_info) + + elif data_uri.startswith("data:video/"): + video = run_async(processor.process_video(data_uri)) + debug_info["source"] = "base64" + return (video, "VIDEO", debug_info) + + except Exception as e: + logger.warning(f"Failed to process data URI: {e}") + + # Return as string if can't process + debug_info["length"] = len(data_uri) + return (data_uri, "STRING", debug_info) + + def _process_bytes(self, data: bytes) -> Tuple[Any, str, Dict]: + """Process raw bytes by detecting format.""" + debug_info = {"size_bytes": len(data)} + + # Try to detect image from magic bytes + if data[:8] == b'\x89PNG\r\n\x1a\n' or data[:2] == b'\xff\xd8': + try: + tensor, info = self._decode_image(data) + debug_info.update(info) + return (tensor, "IMAGE", debug_info) + except Exception: + pass + + # Try to detect audio + if data[:4] == b'fLaC' or data[:4] == b'RIFF': + try: + audio, info = self._decode_audio(data) + debug_info.update(info) + return (audio, "AUDIO", debug_info) + except Exception: + pass + + # Default: return raw bytes + return (data, "BYTES", debug_info) + + def _decode_image(self, data: bytes) -> Tuple[torch.Tensor, Dict]: + """Decode image bytes to tensor.""" + img = Image.open(io.BytesIO(data)) + img = img.convert("RGB") + + np_image = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(np_image).unsqueeze(0) # [1, H, W, C] + + info = { + "dimensions": f"{img.width}x{img.height}", + "channels": 3, + "dtype": "float32", + } + return tensor, info + + def _decode_audio(self, data: bytes) -> Tuple[Dict, Dict]: + """Decode audio bytes to audio dict.""" + try: + import av + except ImportError: + raise ImportError("av library required for audio decoding") + + buffer = io.BytesIO(data) + container = av.open(buffer) + + audio_stream = container.streams.audio[0] + sample_rate = audio_stream.rate + + frames = [] + for frame in container.decode(audio=0): + frames.append(frame.to_ndarray()) + + container.close() + + audio_np = np.concatenate(frames, axis=1) + waveform = torch.from_numpy(audio_np).unsqueeze(0) + + info = { + "sample_rate": sample_rate, + "channels": audio_np.shape[0], + "duration_seconds": audio_np.shape[1] / sample_rate, + } + + return {"waveform": waveform, "sample_rate": sample_rate}, info + + def _decode_video(self, data: bytes) -> Tuple[Dict, Dict]: + """Decode video bytes to video dict.""" + try: + import av + except ImportError: + raise ImportError("av library required for video decoding") + + buffer = io.BytesIO(data) + container = av.open(buffer) + + video_stream = container.streams.video[0] + fps = float(video_stream.average_rate) + + frames = [] + for frame in container.decode(video=0): + np_frame = frame.to_ndarray(format="rgb24") + frames.append(np_frame) + + container.close() + + if not frames: + raise ValueError("No video frames found") + + frames_np = np.stack(frames) + frames_tensor = torch.from_numpy(frames_np).float() / 255.0 + + info = { + "fps": fps, + "frame_count": len(frames), + "dimensions": f"{frames_np.shape[2]}x{frames_np.shape[1]}", + } + + return {"images": frames_tensor, "fps": fps}, info + + def _decode_mesh(self, data: bytes, filename: str) -> Tuple[Dict, Dict]: + """Decode mesh bytes.""" + processor = get_processor() + + # Use processor's mesh parsing + if filename.endswith(".glb") or filename.endswith(".gltf"): + mesh = processor._parse_glb(data) + elif filename.endswith(".obj"): + mesh = processor._parse_obj(data.decode("utf-8")) + elif filename.endswith(".stl"): + mesh = processor._parse_stl_binary(data) + else: + raise ValueError(f"Unsupported mesh format: {filename}") + + vertices = mesh.get("vertices") + faces = mesh.get("faces") + + info = { + "vertex_count": vertices.shape[1] if vertices is not None else 0, + "face_count": faces.shape[1] if faces is not None else 0, + } + + return mesh, info + + def _decode_latent(self, data: bytes) -> Tuple[Dict, Dict]: + """Decode latent from safetensors.""" + from safetensors.torch import load + + buffer = io.BytesIO(data) + tensors = load(buffer.read()) + + samples = tensors.get("samples") + if samples is None: + raise ValueError("No 'samples' key in safetensors") + + info = { + "shape": list(samples.shape), + "dtype": str(samples.dtype), + } + + return {"samples": samples}, info + + def _print_debug_found( + self, + field: str, + value: Any, + value_type: str, + debug_info: dict + ): + """Print debug info for found field.""" + lines = [ + "", + f"[Webhook Transformer] Field: \"{field}\"", + "=" * 50, + f" Status: Found", + f" Type: {value_type}", + ] + + # Add type-specific info + if value_type == "STRING": + preview = str(value) + if len(preview) > 50: + preview = preview[:47] + "..." + lines.append(f" Value: \"{preview}\"") + lines.append(f" Length: {debug_info.get('length', len(str(value)))} characters") + + elif value_type == "IMAGE": + lines.append(f" Dimensions: {debug_info.get('dimensions', 'unknown')}") + lines.append(f" Channels: {debug_info.get('channels', 'unknown')}") + lines.append(f" Dtype: {debug_info.get('dtype', 'unknown')}") + if debug_info.get('size_bytes'): + lines.append(f" Size: {format_size(debug_info['size_bytes'])}") + lines.append(f" Source: {debug_info.get('source', 'file upload')}") + + elif value_type == "AUDIO": + lines.append(f" Sample Rate: {debug_info.get('sample_rate', 'unknown')} Hz") + lines.append(f" Channels: {debug_info.get('channels', 'unknown')}") + if debug_info.get('duration_seconds'): + lines.append(f" Duration: {debug_info['duration_seconds']:.2f}s") + + elif value_type in ("INT", "FLOAT", "BOOLEAN"): + lines.append(f" Value: {value}") + if debug_info.get('original'): + lines.append(f" Original: \"{debug_info['original']}\"") + + elif value_type == "LATENT": + if debug_info.get('shape'): + lines.append(f" Shape: {debug_info['shape']}") + + elif value_type == "MESH": + lines.append(f" Vertices: {debug_info.get('vertex_count', 'unknown')}") + lines.append(f" Faces: {debug_info.get('face_count', 'unknown')}") + + else: + # Generic info + preview = str(value) + if len(preview) > 50: + preview = preview[:47] + "..." + lines.append(f" Value: {preview}") + + lines.append("=" * 50) + print("\n".join(lines)) + + def _print_debug_default(self, field: str, default_value: Any, reason: str): + """Print debug info when using default value.""" + lines = [ + "", + f"[Webhook Transformer] Field: \"{field}\"", + "=" * 50, + f" Status: Not Found -> Using Default", + f" Reason: {reason}", + f" Type: {type(default_value).__name__}", + ] + + preview = str(default_value) + if len(preview) > 50: + preview = preview[:47] + "..." + lines.append(f" Value: {preview}") + + lines.append("=" * 50) + print("\n".join(lines)) + + def _print_error(self, field: str, available_fields: List[str]): + """Print error when field is missing and no default.""" + lines = [ + "", + "[Webhook Transformer] ERROR", + "=" * 50, + f" Field: \"{field}\"", + f" Status: NOT FOUND", + "", + " Available fields:", + ] + + if available_fields: + for i, name in enumerate(available_fields): + prefix = " +-" if i == len(available_fields) - 1 else " |-" + lines.append(f"{prefix} {name}") + else: + lines.append(" (none)") + + lines.append("") + lines.append(" -> Workflow stopped. Provide field or set default_value.") + lines.append("=" * 50) + + print("\n".join(lines)) diff --git a/processors/__init__.py b/processors/__init__.py new file mode 100644 index 0000000..48668ee --- /dev/null +++ b/processors/__init__.py @@ -0,0 +1 @@ +"""Processors for input and output handling.""" diff --git a/processors/inputs.py b/processors/inputs.py new file mode 100644 index 0000000..8fd8a71 --- /dev/null +++ b/processors/inputs.py @@ -0,0 +1,746 @@ +"""Input processing for webhook data.""" + +import asyncio +import base64 +import io +import json +import logging +import re +from typing import Any, List, Optional, Union + +import numpy as np +import torch +from PIL import Image, ImageOps + +try: + import aiohttp +except ImportError: + aiohttp = None + +try: + import torchaudio +except ImportError: + torchaudio = None + +logger = logging.getLogger(__name__) + +DATA_URL_PATTERN = re.compile(r'^data:([^;,]+)?(;base64)?,(.*)$', re.DOTALL) +HTTP_URL_PATTERN = re.compile(r'^https?://', re.IGNORECASE) + + +class WebhookInputProcessor: + """Processes webhook input data into ComfyUI-compatible formats.""" + + def __init__(self, download_timeout: float = 30.0): + self.download_timeout = download_timeout + + async def process_inputs(self, webhook_inputs: dict) -> dict: + """ + Process all inputs and return ComfyUI-compatible values. + + Args: + webhook_inputs: Dictionary of input name -> value + + Returns: + Dictionary of input name -> processed value + """ + processed = {} + + for key, value in webhook_inputs.items(): + try: + processed[key] = await self._process_value(value) + except Exception as e: + logger.error(f"Failed to process input '{key}': {e}") + # Keep original value on error + processed[key] = value + + return processed + + async def _process_value(self, value: Any) -> Any: + """Process a single value based on its type.""" + if isinstance(value, str): + return await self._process_string(value) + elif isinstance(value, list): + return await self._process_array(value) + elif isinstance(value, dict): + # JSON data - serialize to string + return json.dumps(value) + else: + # Numbers, booleans - pass through + return value + + async def _process_string(self, value: str) -> Any: + """Process a string value, detecting if it's base64/URL data.""" + if self._is_base64_image(value): + return await self._decode_base64_image(value) + elif self._is_base64_audio(value): + return await self._decode_base64_audio(value) + elif self._is_url(value): + return await self._download_and_decode(value) + else: + # Plain string + return value + + async def _process_array(self, items: list) -> Union[torch.Tensor, list]: + """ + Process array of items. + + If all items are images, returns a batched tensor [N, H, W, C]. + Otherwise returns a list of processed items. + """ + if not items: + return [] + + # Check if all items are image data + if all(self._is_image_data(item) for item in items): + tensors = [] + for item in items: + if isinstance(item, str): + if self._is_base64_image(item): + tensor = await self._decode_base64_image(item) + elif self._is_url(item): + tensor = await self._download_and_decode(item) + else: + continue + tensors.append(tensor) + + if tensors: + # Stack into batch tensor [N, H, W, C] + return self._batch_images(tensors) + return [] + + # Otherwise process each item individually + return [await self._process_value(item) for item in items] + + def _batch_images(self, tensors: List[torch.Tensor]) -> torch.Tensor: + """ + Batch multiple image tensors, handling different sizes. + + Args: + tensors: List of [1, H, W, C] tensors + + Returns: + Batched tensor [N, H, W, C] + """ + if not tensors: + return torch.zeros(0, 64, 64, 3) + + # Check if all images have the same size + shapes = [t.shape[1:3] for t in tensors] + if all(s == shapes[0] for s in shapes): + # All same size, simple cat + return torch.cat(tensors, dim=0) + + # Different sizes - pad to largest + max_h = max(t.shape[1] for t in tensors) + max_w = max(t.shape[2] for t in tensors) + + padded = [] + for t in tensors: + h, w = t.shape[1], t.shape[2] + if h < max_h or w < max_w: + # Center pad + pad_h_before = (max_h - h) // 2 + pad_h_after = max_h - h - pad_h_before + pad_w_before = (max_w - w) // 2 + pad_w_after = max_w - w - pad_w_before + # F.pad expects (left, right, top, bottom) for 4D tensor + t = torch.nn.functional.pad( + t.permute(0, 3, 1, 2), # [B, C, H, W] + (pad_w_before, pad_w_after, pad_h_before, pad_h_after), + mode='constant', + value=0 + ).permute(0, 2, 3, 1) # Back to [B, H, W, C] + padded.append(t) + + return torch.cat(padded, dim=0) + + # Detection methods + + def _is_base64_image(self, value: str) -> bool: + """Check if string is a base64-encoded image.""" + if value.startswith("data:image/"): + return True + # Check for raw base64 that might be an image (heuristic) + if len(value) > 100 and re.match(r'^[A-Za-z0-9+/]+=*$', value): + # Try to detect PNG/JPEG magic bytes + try: + decoded = base64.b64decode(value[:32]) + # PNG magic: 89 50 4E 47 + if decoded.startswith(b'\x89PNG'): + return True + # JPEG magic: FF D8 FF + if decoded.startswith(b'\xff\xd8\xff'): + return True + # WebP magic: RIFF....WEBP + if decoded.startswith(b'RIFF') and b'WEBP' in decoded[:16]: + return True + except: + pass + return False + + def _is_base64_audio(self, value: str) -> bool: + """Check if string is a base64-encoded audio.""" + return value.startswith("data:audio/") + + def _is_url(self, value: str) -> bool: + """Check if string is an HTTP(S) URL.""" + return bool(HTTP_URL_PATTERN.match(value)) + + def _is_image_data(self, value: Any) -> bool: + """Check if value is image data (base64 or URL).""" + if not isinstance(value, str): + return False + return self._is_base64_image(value) or ( + self._is_url(value) and + any(ext in value.lower() for ext in ['.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp']) + ) + + # Decoding methods + + async def _decode_base64_image(self, data: str) -> torch.Tensor: + """ + Decode base64 image to tensor. + + Supports: + - data:image/png;base64,... + - data:image/jpeg;base64,... + - data:image/webp;base64,... + - Raw base64 (no prefix) + + Returns: + Tensor [1, H, W, C] with values 0-1 + """ + # Strip data URL prefix if present + if data.startswith("data:"): + match = DATA_URL_PATTERN.match(data) + if match: + data = match.group(3) + else: + _, data = data.split(",", 1) + + # Decode base64 + image_bytes = base64.b64decode(data) + + # Load with PIL + pil_image = Image.open(io.BytesIO(image_bytes)) + pil_image = ImageOps.exif_transpose(pil_image) + + # Convert to RGB if needed + if pil_image.mode == "I": + # 16-bit images + pil_image = pil_image.point(lambda i: i * (1 / 255)) + if pil_image.mode != "RGB": + pil_image = pil_image.convert("RGB") + + # Convert to tensor [1, H, W, C] with values 0-1 + np_image = np.array(pil_image).astype(np.float32) / 255.0 + tensor = torch.from_numpy(np_image).unsqueeze(0) + + return tensor + + async def _decode_base64_audio(self, data: str) -> dict: + """ + Decode base64 audio to ComfyUI audio format. + + Returns: + {"waveform": tensor [1, channels, samples], "sample_rate": int} + """ + if torchaudio is None: + raise ImportError("torchaudio is required for audio processing") + + # Strip data URL prefix if present + if data.startswith("data:"): + _, data = data.split(",", 1) + + audio_bytes = base64.b64decode(data) + + # Use torchaudio to load + buffer = io.BytesIO(audio_bytes) + waveform, sample_rate = torchaudio.load(buffer) + + # Resample to 44100 if needed (common ComfyUI sample rate) + target_rate = 44100 + if sample_rate != target_rate: + resampler = torchaudio.transforms.Resample(sample_rate, target_rate) + waveform = resampler(waveform) + sample_rate = target_rate + + return { + "waveform": waveform.unsqueeze(0), # [1, channels, samples] + "sample_rate": sample_rate + } + + async def _download_and_decode(self, url: str) -> Any: + """ + Download content from URL and decode based on content type. + + Returns: + Decoded content (tensor for images, dict for audio, etc.) + """ + if aiohttp is None: + raise ImportError("aiohttp is required for URL downloads") + + timeout = aiohttp.ClientTimeout(total=self.download_timeout) + + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.get(url) as response: + response.raise_for_status() + content_type = response.headers.get("Content-Type", "").lower() + data = await response.read() + + # Detect type from content-type header or URL + if "image" in content_type or self._looks_like_image_url(url): + return await self._decode_image_bytes(data) + elif "audio" in content_type or self._looks_like_audio_url(url): + return await self._decode_audio_bytes(data) + else: + # Return as string + return data.decode("utf-8", errors="replace") + + async def _decode_image_bytes(self, data: bytes) -> torch.Tensor: + """Decode image bytes to tensor.""" + pil_image = Image.open(io.BytesIO(data)) + pil_image = ImageOps.exif_transpose(pil_image) + + if pil_image.mode == "I": + pil_image = pil_image.point(lambda i: i * (1 / 255)) + if pil_image.mode != "RGB": + pil_image = pil_image.convert("RGB") + + np_image = np.array(pil_image).astype(np.float32) / 255.0 + tensor = torch.from_numpy(np_image).unsqueeze(0) + return tensor + + async def _decode_audio_bytes(self, data: bytes) -> dict: + """Decode audio bytes to ComfyUI format.""" + if torchaudio is None: + raise ImportError("torchaudio is required for audio processing") + + buffer = io.BytesIO(data) + waveform, sample_rate = torchaudio.load(buffer) + + target_rate = 44100 + if sample_rate != target_rate: + resampler = torchaudio.transforms.Resample(sample_rate, target_rate) + waveform = resampler(waveform) + sample_rate = target_rate + + return { + "waveform": waveform.unsqueeze(0), + "sample_rate": sample_rate + } + + def _looks_like_image_url(self, url: str) -> bool: + """Check if URL looks like an image based on extension.""" + lower = url.lower() + return any(ext in lower for ext in ['.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp']) + + def _looks_like_audio_url(self, url: str) -> bool: + """Check if URL looks like audio based on extension.""" + lower = url.lower() + return any(ext in lower for ext in ['.mp3', '.wav', '.flac', '.ogg', '.opus', '.m4a']) + + # Specialized extraction methods for nodes + + async def process_image(self, value: Any) -> torch.Tensor: + """ + Process a single image input. + + Args: + value: Base64 string, URL, or already a tensor + + Returns: + Tensor [1, H, W, C] + """ + if isinstance(value, torch.Tensor): + return value + if isinstance(value, str): + if self._is_base64_image(value): + return await self._decode_base64_image(value) + elif self._is_url(value): + return await self._download_and_decode(value) + # Return empty tensor for invalid input + return torch.zeros(1, 64, 64, 3) + + async def process_audio(self, value: Any) -> dict: + """ + Process a single audio input. + + Args: + value: Base64 string, URL, or already an audio dict + + Returns: + {"waveform": tensor, "sample_rate": int} + """ + if isinstance(value, dict) and "waveform" in value: + return value + if isinstance(value, str): + if self._is_base64_audio(value): + return await self._decode_base64_audio(value) + elif self._is_url(value): + return await self._download_and_decode(value) + # Return empty audio for invalid input + return {"waveform": torch.zeros(1, 1, 44100), "sample_rate": 44100} + + async def process_mask(self, value: Any) -> torch.Tensor: + """ + Process a mask input (grayscale image). + + Args: + value: Base64 string or URL of grayscale image + + Returns: + Tensor [1, H, W] for MASK type + """ + if isinstance(value, torch.Tensor): + if len(value.shape) == 3: + return value + elif len(value.shape) == 4: + # Convert [B, H, W, C] to [B, H, W] using first channel + return value[:, :, :, 0] + + if isinstance(value, str): + # Decode image + if self._is_base64_image(value): + tensor = await self._decode_base64_image(value) + elif self._is_url(value): + tensor = await self._download_and_decode(value) + else: + return torch.zeros(1, 64, 64) + + # Convert to grayscale mask + if len(tensor.shape) == 4: + # Average RGB channels + mask = tensor.mean(dim=-1) + else: + mask = tensor + return mask + + return torch.zeros(1, 64, 64) + + async def process_video(self, value: Any) -> Any: + """ + Process video input. + + Note: Video processing is more complex and may require + additional dependencies like decord or av. + + Args: + value: Base64 string or URL + + Returns: + Video data (format depends on ComfyUI video nodes) + """ + # Basic implementation - return None for now + # Full video support would require decord or av + logger.warning("Video input processing not fully implemented") + return None + + async def process_latent(self, value: Any) -> dict: + """ + Process latent input (safetensors or torch file). + + Args: + value: Base64 string (safetensors/pt) or URL + + Returns: + {"samples": tensor, ...} latent dict + """ + if isinstance(value, dict) and "samples" in value: + return value + + default_latent = {"samples": torch.zeros(1, 4, 64, 64)} + + if isinstance(value, str): + try: + # Get data bytes + if value.startswith("data:"): + _, data_str = value.split(",", 1) + data = base64.b64decode(data_str) + elif self._is_url(value): + if aiohttp is None: + return default_latent + timeout = aiohttp.ClientTimeout(total=self.download_timeout) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.get(value) as response: + response.raise_for_status() + data = await response.read() + else: + # Raw base64 + data = base64.b64decode(value) + + # Try safetensors first + try: + from safetensors.torch import load + buffer = io.BytesIO(data) + tensors = load(buffer.getvalue()) + if "samples" in tensors: + return {"samples": tensors["samples"]} + # Return first tensor as samples + for key, tensor in tensors.items(): + return {"samples": tensor} + except ImportError: + pass + except Exception: + pass + + # Try torch load + try: + buffer = io.BytesIO(data) + loaded = torch.load(buffer, weights_only=True) + if isinstance(loaded, dict) and "samples" in loaded: + return loaded + if isinstance(loaded, torch.Tensor): + return {"samples": loaded} + except Exception: + pass + + except Exception as e: + logger.error(f"Failed to process latent: {e}") + + return default_latent + + async def process_mesh(self, value: Any) -> dict: + """ + Process 3D mesh input (GLB, OBJ, STL). + + Args: + value: Base64 string or URL + + Returns: + {"vertices": tensor, "faces": tensor} mesh dict + """ + default_mesh = { + "vertices": torch.zeros(1, 3, 3), + "faces": torch.zeros(1, 1, 3, dtype=torch.long) + } + + if isinstance(value, dict) and "vertices" in value: + return value + + if isinstance(value, str): + try: + # Get data bytes + if value.startswith("data:"): + _, data_str = value.split(",", 1) + data = base64.b64decode(data_str) + elif self._is_url(value): + if aiohttp is None: + return default_mesh + timeout = aiohttp.ClientTimeout(total=self.download_timeout) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.get(value) as response: + response.raise_for_status() + data = await response.read() + else: + # Raw base64 + data = base64.b64decode(value) + + # Detect format and parse + return self._parse_mesh_bytes(data, value) + + except Exception as e: + logger.error(f"Failed to process mesh: {e}") + + return default_mesh + + def _parse_mesh_bytes(self, data: bytes, hint: str = "") -> dict: + """Parse mesh from bytes based on format detection.""" + default_mesh = { + "vertices": torch.zeros(1, 3, 3), + "faces": torch.zeros(1, 1, 3, dtype=torch.long) + } + + # Detect format + if data.startswith(b'glTF'): + return self._parse_glb(data) + elif data.startswith(b'solid ') or data[:80].count(b'\x00') < 10: + # Text STL or ASCII format + if b'v ' in data[:1000]: + return self._parse_obj(data) + elif len(data) > 84: + # Binary STL check + try: + import struct + num_triangles = struct.unpack(' dict: + """Parse GLB (binary glTF) file.""" + import struct + import json as json_mod + + default_mesh = { + "vertices": torch.zeros(1, 3, 3), + "faces": torch.zeros(1, 1, 3, dtype=torch.long) + } + + try: + # GLB header: magic (4), version (4), length (4) + if len(data) < 12: + return default_mesh + + magic, version, length = struct.unpack('<4sII', data[:12]) + if magic != b'glTF': + return default_mesh + + # Read JSON chunk + chunk_length, chunk_type = struct.unpack(' len(data): + return default_mesh + + bin_length, bin_type = struct.unpack(' dict: + """Parse OBJ file.""" + default_mesh = { + "vertices": torch.zeros(1, 3, 3), + "faces": torch.zeros(1, 1, 3, dtype=torch.long) + } + + try: + text = data.decode('utf-8', errors='replace') + vertices = [] + faces = [] + + for line in text.split('\n'): + line = line.strip() + if line.startswith('v '): + parts = line.split()[1:4] + vertices.append([float(p) for p in parts]) + elif line.startswith('f '): + parts = line.split()[1:] + # Handle v/vt/vn format + face_indices = [] + for p in parts[:3]: # Only first 3 vertices for triangles + idx = int(p.split('/')[0]) - 1 # OBJ is 1-indexed + face_indices.append(idx) + if len(face_indices) == 3: + faces.append(face_indices) + + if vertices and faces: + return { + "vertices": torch.tensor(vertices).unsqueeze(0).float(), + "faces": torch.tensor(faces, dtype=torch.long).unsqueeze(0) + } + + except Exception as e: + logger.error(f"Failed to parse OBJ: {e}") + + return default_mesh + + def _parse_stl_binary(self, data: bytes) -> dict: + """Parse binary STL file.""" + import struct + + default_mesh = { + "vertices": torch.zeros(1, 3, 3), + "faces": torch.zeros(1, 1, 3, dtype=torch.long) + } + + try: + num_triangles = struct.unpack(' str: + """ + Resolve the full path to an output file. + + Args: + filename: Name of the file + subfolder: Subfolder within the directory + folder_type: Type of folder (output, temp, input) + + Returns: + Full filesystem path to the file + """ + if folder_paths is None: + # Fallback if folder_paths not available + base_dir = os.getcwd() + return os.path.join(base_dir, folder_type, subfolder, filename) + + if folder_type == "output": + base_dir = folder_paths.get_output_directory() + elif folder_type == "temp": + base_dir = folder_paths.get_temp_directory() + elif folder_type == "input": + base_dir = folder_paths.get_input_directory() + else: + base_dir = folder_paths.get_output_directory() + + if subfolder: + return os.path.join(base_dir, subfolder, filename) + return os.path.join(base_dir, filename) + + +def detect_output_type( + ui_output: dict, + node_id: Optional[str] = None, + node_type: Optional[str] = None +) -> list[OutputInfo]: + """ + Detect output type from ComfyUI UI return structure. + + Known UI structures: + - {"images": [{filename, subfolder, type}, ...]} + - {"audio": [{filename, subfolder, type}, ...]} + - {"video": [{filename, subfolder, type}, ...]} + - {"3d": [{filename, subfolder, type}, ...]} + - {"latents": [{filename, subfolder, type}, ...]} + - {"text": (value,)} or {"text": ["value"]} + - {"result": [model_file, camera_info, bg_image]} + + Args: + ui_output: The UI output dictionary from a node + node_id: Optional node ID for metadata + node_type: Optional node type for metadata + + Returns: + List of OutputInfo objects describing the outputs + """ + outputs = [] + + if not ui_output or not isinstance(ui_output, dict): + return outputs + + # Handle images + if "images" in ui_output: + images = ui_output["images"] + if isinstance(images, list): + for idx, item in enumerate(images): + if isinstance(item, dict) and "filename" in item: + info = _create_output_info_from_file( + item, "image", node_id, node_type + ) + outputs.append(info) + + # Handle audio + if "audio" in ui_output: + audio_items = ui_output["audio"] + if isinstance(audio_items, list): + for item in audio_items: + if isinstance(item, dict) and "filename" in item: + info = _create_output_info_from_file( + item, "audio", node_id, node_type + ) + outputs.append(info) + + # Handle video + if "video" in ui_output: + video_items = ui_output["video"] + if isinstance(video_items, list): + for item in video_items: + if isinstance(item, dict) and "filename" in item: + info = _create_output_info_from_file( + item, "video", node_id, node_type + ) + outputs.append(info) + + # Handle 3D meshes + if "3d" in ui_output: + mesh_items = ui_output["3d"] + if isinstance(mesh_items, list): + for item in mesh_items: + if isinstance(item, dict) and "filename" in item: + info = _create_output_info_from_file( + item, "mesh", node_id, node_type + ) + outputs.append(info) + + # Handle latents + if "latents" in ui_output: + latent_items = ui_output["latents"] + if isinstance(latent_items, list): + for item in latent_items: + if isinstance(item, dict) and "filename" in item: + info = _create_output_info_from_file( + item, "latent", node_id, node_type + ) + outputs.append(info) + + # Handle text outputs + if "text" in ui_output: + text_data = ui_output["text"] + if isinstance(text_data, tuple) and len(text_data) > 0: + # Text is often returned as a tuple + text_value = str(text_data[0]) + elif isinstance(text_data, list) and len(text_data) > 0: + text_value = str(text_data[0]) + elif isinstance(text_data, str): + text_value = text_data + else: + text_value = str(text_data) + + outputs.append(OutputInfo( + type="text", + inline_content=text_value, + mime_type="text/plain", + node_id=node_id, + node_type=node_type, + )) + + # Handle 3D preview (result structure) + if "result" in ui_output and "images" not in ui_output: + result_data = ui_output["result"] + outputs.append(OutputInfo( + type="3d_preview", + inline_content=json.dumps(result_data), + mime_type="application/json", + node_id=node_id, + node_type=node_type, + )) + + # Handle animated flag (for videos/GIFs saved as images) + if "animated" in ui_output: + # Mark the last image output as animated if applicable + for output in reversed(outputs): + if output.type == "image": + output.format = "animated" + break + + # Fallback: try to extract any file references from unknown structures + if not outputs: + outputs = _extract_unknown_outputs(ui_output, node_id, node_type) + + return outputs + + +def _create_output_info_from_file( + item: dict, + output_type: str, + node_id: Optional[str], + node_type: Optional[str] +) -> OutputInfo: + """Create OutputInfo from a file reference dictionary.""" + filename = item.get("filename", "") + subfolder = item.get("subfolder", "") + folder_type = item.get("type", "output") + + # Get file path and size + file_path = get_output_path(filename, subfolder, folder_type) + file_size = None + if os.path.exists(file_path): + file_size = os.path.getsize(file_path) + + return OutputInfo( + type=output_type, + filename=filename, + subfolder=subfolder, + folder_type=folder_type, + mime_type=guess_mime_type(filename), + format=get_format_from_filename(filename), + file_size_bytes=file_size, + node_id=node_id, + node_type=node_type, + ) + + +def _extract_unknown_outputs( + ui_output: dict, + node_id: Optional[str], + node_type: Optional[str] +) -> list[OutputInfo]: + """ + Extract outputs from unknown UI structures. + Looks for any dictionary with 'filename' key. + """ + outputs = [] + + for key, value in ui_output.items(): + if isinstance(value, list): + for item in value: + if isinstance(item, dict) and "filename" in item: + info = _create_output_info_from_file( + item, "unknown", node_id, node_type + ) + outputs.append(info) + logger.info( + f"Extracted unknown output type '{key}': {item.get('filename')}" + ) + + # If still no outputs, serialize the entire structure as JSON + if not outputs and ui_output: + logger.warning( + f"Unknown output format, serializing as JSON: {list(ui_output.keys())}" + ) + outputs.append(OutputInfo( + type="unknown", + inline_content=json.dumps(ui_output), + mime_type="application/json", + node_id=node_id, + node_type=node_type, + )) + + return outputs + + +def process_outputs_for_webhook( + outputs: dict, + meta: Optional[dict] = None +) -> list[OutputInfo]: + """ + Process all outputs from a workflow execution for webhook delivery. + + Args: + outputs: Dictionary mapping node_id -> UI output dict + meta: Optional metadata dict mapping node_id -> node metadata + + Returns: + List of all OutputInfo objects from all nodes + """ + all_outputs = [] + + for node_id, ui_output in outputs.items(): + # Get node type from metadata if available + node_type = None + if meta and node_id in meta: + node_meta = meta[node_id] + if isinstance(node_meta, dict): + node_type = node_meta.get("node_type") + + # Detect outputs for this node + node_outputs = detect_output_type(ui_output, node_id, node_type) + all_outputs.extend(node_outputs) + + return all_outputs + + +def load_output_file(output_info: OutputInfo) -> Optional[bytes]: + """ + Load the binary content of an output file. + + Args: + output_info: OutputInfo with file location details + + Returns: + File bytes or None if file doesn't exist + """ + if not output_info.filename: + return None + + file_path = get_output_path( + output_info.filename, + output_info.subfolder or "", + output_info.folder_type or "output" + ) + + if not os.path.exists(file_path): + logger.warning(f"Output file not found: {file_path}") + return None + + try: + with open(file_path, "rb") as f: + return f.read() + except Exception as e: + logger.error(f"Failed to read output file {file_path}: {e}") + return None + + +def get_image_dimensions(output_info: OutputInfo) -> tuple[Optional[int], Optional[int]]: + """ + Get dimensions of an image output. + + Args: + output_info: OutputInfo for an image + + Returns: + Tuple of (width, height) or (None, None) if unable to determine + """ + if output_info.type != "image" or not output_info.filename: + return None, None + + try: + from PIL import Image + + file_path = get_output_path( + output_info.filename, + output_info.subfolder or "", + output_info.folder_type or "output" + ) + + if os.path.exists(file_path): + with Image.open(file_path) as img: + return img.size # (width, height) + except Exception as e: + logger.warning(f"Failed to get image dimensions: {e}") + + return None, None + + +def enrich_output_info(output_info: OutputInfo) -> OutputInfo: + """ + Enrich OutputInfo with additional metadata. + + Adds dimensions for images, file size, etc. + + Args: + output_info: The OutputInfo to enrich + + Returns: + The same OutputInfo with additional fields populated + """ + # Get file size if not already set + if output_info.file_size_bytes is None and output_info.filename: + file_path = get_output_path( + output_info.filename, + output_info.subfolder or "", + output_info.folder_type or "output" + ) + if os.path.exists(file_path): + output_info.file_size_bytes = os.path.getsize(file_path) + + # Get image dimensions + if output_info.type == "image" and (output_info.width is None or output_info.height is None): + width, height = get_image_dimensions(output_info) + if width and height: + output_info.width = width + output_info.height = height + + return output_info diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..2fb1caf --- /dev/null +++ b/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +testpaths = tests +python_files = test_*.py +asyncio_mode = auto diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d16ee03 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,18 @@ +# ComfyUI Workflow API - Dependencies + +# Required +aiohttp>=3.8.0 # Async HTTP client for webhook delivery + +# Optional - for enhanced validation (uncomment if desired) +# pydantic>=2.0.0 # Data validation and serialization + +# Development dependencies (install with: pip install -r requirements.txt[dev]) +# pytest>=7.0.0 # Testing framework +# pytest-asyncio>=0.21.0 # Async test support +# pytest-cov>=4.0.0 # Coverage reporting + +# Note: The following are typically already available in ComfyUI: +# - torch # Tensor operations +# - numpy # Array operations +# - Pillow # Image processing +# - av # Video encoding (for video output) diff --git a/run_tests.py b/run_tests.py new file mode 100644 index 0000000..9c071c8 --- /dev/null +++ b/run_tests.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""Run all tests.""" + +import subprocess +import sys +from pathlib import Path + +def main(): + test_dir = Path(__file__).parent / "tests" + test_files = list(test_dir.glob("test_*.py")) + + if not test_files: + print("No test files found") + return 1 + + failed = 0 + for test_file in test_files: + print(f"\n{'='*50}") + print(f"Running {test_file.name}") + print('='*50) + result = subprocess.run([sys.executable, str(test_file)], cwd=Path(__file__).parent) + if result.returncode != 0: + failed += 1 + + print(f"\n{'='*50}") + if failed: + print(f"FAILED: {failed} test file(s) had failures") + else: + print("All tests passed!") + return failed + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/sample_input.png b/sample_input.png new file mode 100644 index 0000000..19c6017 Binary files /dev/null and b/sample_input.png differ diff --git a/server/__init__.py b/server/__init__.py new file mode 100644 index 0000000..cfe7477 --- /dev/null +++ b/server/__init__.py @@ -0,0 +1,6 @@ +"""Server extension for ComfyUI webhook integration.""" + +from .routes import register_routes +from .uploads import save_upload, cleanup_uploads + +__all__ = ["register_routes", "save_upload", "cleanup_uploads"] diff --git a/server/routes.py b/server/routes.py new file mode 100644 index 0000000..fe26f55 --- /dev/null +++ b/server/routes.py @@ -0,0 +1,308 @@ +"""HTTP routes for webhook integration.""" + +import json +import logging +import uuid +from datetime import datetime +from typing import Optional + +from aiohttp import web + +logger = logging.getLogger(__name__) + + +def register_routes(): + """Register webhook routes with ComfyUI's PromptServer.""" + try: + from server import PromptServer + + routes = PromptServer.instance.routes + + @routes.post("/api/webhook") + async def webhook_submit(request): + """ + Submit a workflow via multipart/form-data. + + Fields: + - workflow: JSON workflow definition (required) + - callback_url: URL to receive events and outputs (required) + - auth_header: Auth header name (optional) + - auth_value: Auth header value (optional) + - timeout: HTTP timeout in seconds (optional, default: 60) + - max_retries: Max retry attempts (optional, default: 3) + - : Uploaded files, accessible by field name + - : Text fields, accessible by field name + + Returns: + - request_id: Unique identifier for this request + - prompt_id: ComfyUI prompt ID + - status: "queued" + """ + from .uploads import save_upload + + request_id = str(uuid.uuid4()) + + try: + # Parse multipart request + reader = await request.multipart() + + workflow = None + callback_url = None + auth_header = None + auth_value = None + timeout = 60 + max_retries = 3 + send_progress = True + progress_interval_ms = 500 + webhook_inputs = {} + + # Get request metadata + remote_ip = request.remote + content_type = request.content_type + content_length = request.content_length + user_agent = request.headers.get("User-Agent", "unknown") + + async for part in reader: + name = part.name + + if name == "workflow": + content = await part.text() + try: + workflow = json.loads(content) + except json.JSONDecodeError as e: + return web.json_response( + {"error": f"Invalid workflow JSON: {e}"}, + status=400 + ) + + elif name == "callback_url": + callback_url = await part.text() + + elif name == "auth_header": + auth_header = await part.text() + + elif name == "auth_value": + auth_value = await part.text() + + elif name == "timeout": + try: + timeout = int(await part.text()) + except ValueError: + pass + + elif name == "max_retries": + try: + max_retries = int(await part.text()) + except ValueError: + pass + + elif name == "send_progress": + value = await part.text() + send_progress = value.lower() in ("true", "1", "yes") + + elif name == "progress_interval_ms": + try: + progress_interval_ms = int(await part.text()) + except ValueError: + pass + + elif part.filename: + # File upload - save to disk and store path (not raw data to avoid serialization issues) + data = await part.read() + filename = part.filename + file_content_type = part.headers.get("Content-Type", "application/octet-stream") + + # Save to ComfyUI input directory + saved_path = save_upload(data, filename, request_id) + + # Store file metadata (path, not raw bytes) in webhook_inputs + webhook_inputs[name] = { + "_is_file": True, + "_filename": filename, + "_content_type": file_content_type, + "_saved_path": saved_path, + "_size_bytes": len(data), + } + logger.debug(f"Saved upload '{name}' ({len(data)} bytes) to {saved_path}") + + else: + # Text field + webhook_inputs[name] = await part.text() + + # Validate required fields + if workflow is None: + return web.json_response( + {"error": "Missing required field: workflow"}, + status=400 + ) + + if callback_url is None: + return web.json_response( + {"error": "Missing required field: callback_url"}, + status=400 + ) + + # Generate prompt_id + prompt_id = str(uuid.uuid4()) + + # Build webhook_config for extra_data + webhook_config = { + "callback_url": callback_url, + "auth_header": auth_header, + "auth_value": auth_value, + "request_id": request_id, + "timeout": timeout, + "max_retries": max_retries, + "send_progress": send_progress, + "progress_interval_ms": progress_interval_ms, + "remote_ip": remote_ip, + "content_type": content_type, + "content_length": content_length, + "user_agent": user_agent, + } + + # Build extra_data for workflow + # ComfyUI passes extra_data['extra_pnginfo'] to nodes with EXTRA_PNGINFO hidden input + extra_data = { + "webhook_config": webhook_config, + "webhook_inputs": webhook_inputs, + "extra_pnginfo": { + "webhook_config": webhook_config, + "webhook_inputs": webhook_inputs, + }, + } + + # Validate prompt and get outputs_to_execute + import execution + valid = await execution.validate_prompt(prompt_id, workflow, None) + + if not valid[0]: + # Validation failed + logger.warning(f"Invalid workflow: {valid[1]}") + return web.json_response( + {"error": valid[1], "node_errors": valid[3]}, + status=400 + ) + + outputs_to_execute = valid[2] + + # Register webhook context for progress tracking + from ..core.context import register_webhook_context + from ..core.types import WebhookContext + + context = WebhookContext( + request_id=request_id, + prompt_id=prompt_id, + callback_url=callback_url, + auth_header=auth_header, + auth_value=auth_value, + timeout=timeout, + max_retries=max_retries, + send_progress=send_progress, + progress_interval_ms=progress_interval_ms, + ) + register_webhook_context(context) + + # Queue the workflow + queue = PromptServer.instance.prompt_queue + + # ComfyUI expects: (priority, prompt_id, prompt, extra_data, outputs_to_execute, sensitive) + queue.put((0, prompt_id, workflow, extra_data, outputs_to_execute, {})) + + logger.info(f"Webhook request {request_id} queued as prompt {prompt_id}") + + return web.json_response({ + "request_id": request_id, + "prompt_id": prompt_id, + "status": "queued", + }) + + except Exception as e: + logger.error(f"Webhook submit error: {e}", exc_info=True) + return web.json_response( + {"error": str(e)}, + status=500 + ) + + @routes.get("/api/webhook/status/{request_id}") + async def webhook_status(request): + """ + Get status of a webhook request. + + Returns: + - request_id: The request ID + - prompt_id: ComfyUI prompt ID + - status: queued|running|completed|error|interrupted + - progress: Current progress info (if running) + - error: Error details (if error) + """ + from ..core.context import get_context + + request_id = request.match_info["request_id"] + context = get_context(request_id) + + if context is None: + return web.json_response( + {"error": "Request not found"}, + status=404 + ) + + response = { + "request_id": context.request_id, + "prompt_id": context.prompt_id, + "callback_url": context.callback_url, + "fields": context.get_field_names(), + } + + if context.start_time: + response["started_at"] = context.start_time.isoformat() + + if context.completed_at: + response["completed_at"] = context.completed_at.isoformat() + if context.start_time: + elapsed = (context.completed_at - context.start_time).total_seconds() * 1000 + response["execution_time_ms"] = elapsed + + return web.json_response(response) + + @routes.delete("/api/webhook/{request_id}") + async def webhook_cancel(request): + """Cancel a queued or running webhook request.""" + from ..core.context import get_context, update_context + + request_id = request.match_info["request_id"] + context = get_context(request_id) + + if context is None: + return web.json_response( + {"error": "Request not found"}, + status=404 + ) + + # Try to interrupt the execution + try: + if context.prompt_id: + PromptServer.instance.prompt_queue.delete_queue_item(context.prompt_id) + + update_context(request_id, completed_at=datetime.now()) + + return web.json_response({ + "request_id": request_id, + "status": "cancelled", + }) + except Exception as e: + logger.error(f"Failed to cancel request {request_id}: {e}") + return web.json_response( + {"error": str(e)}, + status=500 + ) + + logger.info("Webhook routes registered: POST /api/webhook, GET /api/webhook/status/{id}") + return True + + except ImportError as e: + logger.warning(f"Could not register webhook routes (ComfyUI server not available): {e}") + return False + except Exception as e: + logger.error(f"Failed to register webhook routes: {e}", exc_info=True) + return False diff --git a/server/uploads.py b/server/uploads.py new file mode 100644 index 0000000..b1c6993 --- /dev/null +++ b/server/uploads.py @@ -0,0 +1,187 @@ +"""Handle file uploads for webhook requests.""" + +import logging +import os +import shutil +import uuid +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + +# Track uploads per request for cleanup +_upload_registry: dict[str, list[str]] = {} + + +def get_upload_directory() -> str: + """Get the directory for webhook uploads.""" + try: + import folder_paths + # Use ComfyUI's input directory for uploads + input_dir = folder_paths.get_input_directory() + webhook_dir = os.path.join(input_dir, "webhook_uploads") + except ImportError: + # Fallback for testing + webhook_dir = os.path.join(os.getcwd(), "input", "webhook_uploads") + + os.makedirs(webhook_dir, exist_ok=True) + return webhook_dir + + +def save_upload(data: bytes, filename: str, request_id: str) -> str: + """ + Save an uploaded file to the webhook uploads directory. + + Args: + data: File content as bytes + filename: Original filename + request_id: Request ID for organizing uploads + + Returns: + Full path to the saved file + """ + upload_dir = get_upload_directory() + + # Create request-specific subdirectory + request_dir = os.path.join(upload_dir, request_id) + os.makedirs(request_dir, exist_ok=True) + + # Sanitize filename and make unique + safe_filename = sanitize_filename(filename) + base, ext = os.path.splitext(safe_filename) + + # Add short UUID to avoid collisions + unique_filename = f"{base}_{uuid.uuid4().hex[:8]}{ext}" + file_path = os.path.join(request_dir, unique_filename) + + # Write file + with open(file_path, "wb") as f: + f.write(data) + + # Track for cleanup + if request_id not in _upload_registry: + _upload_registry[request_id] = [] + _upload_registry[request_id].append(file_path) + + logger.debug(f"Saved upload: {file_path} ({len(data)} bytes)") + return file_path + + +def sanitize_filename(filename: str) -> str: + """Sanitize a filename for safe storage.""" + # Remove path separators and null bytes + filename = filename.replace("/", "_").replace("\\", "_").replace("\x00", "") + + # Remove leading dots (hidden files) + while filename.startswith("."): + filename = filename[1:] + + # Default name if empty + if not filename: + filename = "upload" + + # Limit length + if len(filename) > 200: + base, ext = os.path.splitext(filename) + filename = base[:200 - len(ext)] + ext + + return filename + + +def cleanup_uploads(request_id: str) -> int: + """ + Clean up uploaded files for a request. + + Args: + request_id: Request ID to clean up + + Returns: + Number of files removed + """ + count = 0 + + # Remove tracked files + if request_id in _upload_registry: + for file_path in _upload_registry[request_id]: + try: + if os.path.exists(file_path): + os.remove(file_path) + count += 1 + except Exception as e: + logger.warning(f"Failed to remove upload {file_path}: {e}") + + del _upload_registry[request_id] + + # Also try to remove the request directory + upload_dir = get_upload_directory() + request_dir = os.path.join(upload_dir, request_id) + + try: + if os.path.exists(request_dir): + shutil.rmtree(request_dir) + logger.debug(f"Removed upload directory: {request_dir}") + except Exception as e: + logger.warning(f"Failed to remove upload directory {request_dir}: {e}") + + return count + + +def get_upload_path(request_id: str, field_name: str) -> Optional[str]: + """ + Get the path to an uploaded file by field name. + + This is used by WebhookInput node to load files. + + Args: + request_id: Request ID + field_name: Field name from the upload + + Returns: + File path if found, None otherwise + """ + # This is primarily resolved via extra_data["webhook_uploads"] + # But we can search the registry as fallback + if request_id in _upload_registry: + for path in _upload_registry[request_id]: + if field_name in path: + return path + return None + + +def cleanup_old_uploads(max_age_hours: int = 24) -> int: + """ + Clean up old upload directories. + + Args: + max_age_hours: Remove directories older than this + + Returns: + Number of directories removed + """ + import time + + upload_dir = get_upload_directory() + if not os.path.exists(upload_dir): + return 0 + + count = 0 + cutoff = time.time() - (max_age_hours * 3600) + + try: + for entry in os.scandir(upload_dir): + if entry.is_dir(): + try: + mtime = entry.stat().st_mtime + if mtime < cutoff: + shutil.rmtree(entry.path) + count += 1 + logger.debug(f"Cleaned up old upload directory: {entry.path}") + except Exception as e: + logger.warning(f"Failed to clean up {entry.path}: {e}") + except Exception as e: + logger.warning(f"Failed to scan upload directory: {e}") + + if count > 0: + logger.info(f"Cleaned up {count} old upload directories") + + return count diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..74864ec --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,7 @@ +"""Pytest configuration.""" + +import sys +from pathlib import Path + +# Add parent directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent)) diff --git a/tests/test_core.py b/tests/test_core.py new file mode 100644 index 0000000..97da314 --- /dev/null +++ b/tests/test_core.py @@ -0,0 +1,121 @@ +"""Core functionality tests.""" + +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from core.types import WebhookContext, WebhookField, WebhookEvent, WebhookResult, OutputInfo +from core.context import register_webhook_context, get_webhook_context, unregister_webhook_context + + +def test_context_defaults(): + ctx = WebhookContext() + assert ctx.callback_url == "" + assert ctx.timeout == 60 + assert ctx.max_retries == 3 + assert ctx.request_id + # start_time is None until register_webhook_context() is called + assert ctx.start_time is None + assert ctx.timestamp is not None + + +def test_context_with_values(): + ctx = WebhookContext( + callback_url="https://example.com/webhook", + auth_header="Authorization", + auth_value="Bearer token123", + timeout=30, + ) + assert ctx.callback_url == "https://example.com/webhook" + assert ctx.auth_header == "Authorization" + assert ctx.timeout == 30 + + +def test_auth_headers(): + ctx = WebhookContext(auth_header="X-API-Key", auth_value="secret") + assert ctx.get_auth_headers() == {"X-API-Key": "secret"} + + ctx_no_auth = WebhookContext() + assert ctx_no_auth.get_auth_headers() == {} + + +def test_field_access(): + field = WebhookField(name="prompt", value="hello world", size_bytes=11) + ctx = WebhookContext(fields={"prompt": field}) + + assert ctx.has_field("prompt") + assert not ctx.has_field("missing") + assert ctx.get_field("prompt") == field + assert ctx.get_field_value("prompt") == "hello world" + assert ctx.get_field_value("missing", "default") == "default" + + +def test_legacy_inputs(): + ctx = WebhookContext(inputs={"old_field": "old_value"}) + assert ctx.get_field_value("old_field") == "old_value" + assert ctx.has_field("old_field") + + +def test_censor_auth(): + ctx = WebhookContext(auth_header="Authorization", auth_value="Bearer secrettoken123") + censored = ctx.censor_auth() + assert "Bear" in censored + assert "123" in censored + assert "secrettoken" not in censored + + +def test_context_registry(): + ctx = WebhookContext(prompt_id="test-123", callback_url="https://example.com") + register_webhook_context(ctx) + assert get_webhook_context("test-123") is ctx + unregister_webhook_context("test-123") + assert get_webhook_context("test-123") is None + + +def test_output_info(): + info = OutputInfo(type="image", filename="test.png", width=1024, height=768) + d = info.to_dict() + assert d["type"] == "image" + assert d["filename"] == "test.png" + assert d["width"] == 1024 + + +def test_webhook_events(): + assert WebhookEvent.STARTED.value == "workflow.started" + assert WebhookEvent.COMPLETED.value == "workflow.completed" + assert WebhookEvent.ERROR.value == "workflow.error" + + +def run_all(): + tests = [ + test_context_defaults, + test_context_with_values, + test_auth_headers, + test_field_access, + test_legacy_inputs, + test_censor_auth, + test_context_registry, + test_output_info, + test_webhook_events, + ] + + passed = 0 + failed = 0 + + for test in tests: + try: + test() + print(f" PASS: {test.__name__}") + passed += 1 + except Exception as e: + print(f" FAIL: {test.__name__} - {e}") + failed += 1 + + print(f"\n{passed} passed, {failed} failed") + return failed == 0 + + +if __name__ == "__main__": + print("Running core tests...") + success = run_all() + sys.exit(0 if success else 1) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..56c1bf0 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,16 @@ +"""Utility functions for webhook functionality.""" + +from .mime import guess_mime_type, get_extension_for_mime, MIME_TYPES +from .files import get_output_path, get_image_dimensions, get_temp_directory +from .helpers import format_size, run_async + +__all__ = [ + "guess_mime_type", + "get_extension_for_mime", + "MIME_TYPES", + "get_output_path", + "get_image_dimensions", + "get_temp_directory", + "format_size", + "run_async", +] diff --git a/utils/files.py b/utils/files.py new file mode 100644 index 0000000..02a2265 --- /dev/null +++ b/utils/files.py @@ -0,0 +1,192 @@ +"""File path resolution and utilities.""" + +import logging +import os +from typing import Optional, Tuple + +logger = logging.getLogger(__name__) + + +def get_output_path( + filename: str, + subfolder: str = "", + folder_type: str = "output" +) -> str: + """ + Resolve the full path to an output file. + + Args: + filename: Name of the file + subfolder: Subfolder within the directory + folder_type: Type of folder (output, temp, input) + + Returns: + Full filesystem path to the file + """ + try: + import folder_paths + + if folder_type == "output": + base_dir = folder_paths.get_output_directory() + elif folder_type == "temp": + base_dir = folder_paths.get_temp_directory() + elif folder_type == "input": + base_dir = folder_paths.get_input_directory() + else: + base_dir = folder_paths.get_output_directory() + + except ImportError: + # Fallback if folder_paths not available + base_dir = os.path.join(os.getcwd(), folder_type) + + if subfolder: + return os.path.join(base_dir, subfolder, filename) + return os.path.join(base_dir, filename) + + +def get_image_dimensions(file_path: str) -> Tuple[Optional[int], Optional[int]]: + """ + Get dimensions of an image file. + + Args: + file_path: Path to the image file + + Returns: + Tuple of (width, height) or (None, None) if unable to determine + """ + if not os.path.exists(file_path): + return None, None + + try: + from PIL import Image + with Image.open(file_path) as img: + return img.size # (width, height) + except ImportError: + logger.debug("PIL not available for image dimension detection") + except Exception as e: + logger.debug(f"Failed to get image dimensions: {e}") + + return None, None + + +def get_temp_directory() -> str: + """Get the webhook temp directory.""" + try: + import folder_paths + temp_dir = folder_paths.get_temp_directory() + except ImportError: + temp_dir = os.path.join(os.getcwd(), "temp") + + webhook_temp = os.path.join(temp_dir, "webhook") + os.makedirs(webhook_temp, exist_ok=True) + return webhook_temp + + +def save_temp_file(data: bytes, filename: str, request_id: str = "") -> str: + """ + Save data to a temp file. + + Args: + data: File content + filename: Filename to use + request_id: Optional request ID for organization + + Returns: + Path to saved file + """ + temp_dir = get_temp_directory() + + if request_id: + subdir = os.path.join(temp_dir, request_id) + os.makedirs(subdir, exist_ok=True) + file_path = os.path.join(subdir, filename) + else: + file_path = os.path.join(temp_dir, filename) + + with open(file_path, "wb") as f: + f.write(data) + + return file_path + + +def cleanup_temp_files(request_id: str) -> int: + """ + Clean up temp files for a request. + + Args: + request_id: Request ID + + Returns: + Number of files removed + """ + import shutil + + temp_dir = get_temp_directory() + subdir = os.path.join(temp_dir, request_id) + + if not os.path.exists(subdir): + return 0 + + try: + count = len(os.listdir(subdir)) + shutil.rmtree(subdir) + return count + except Exception as e: + logger.warning(f"Failed to clean up temp files: {e}") + return 0 + + +def detect_file_type(file_path: str) -> Optional[str]: + """ + Detect file type from magic bytes. + + Args: + file_path: Path to file + + Returns: + Detected type: "image", "audio", "video", "3d", or None + """ + if not os.path.exists(file_path): + return None + + try: + with open(file_path, "rb") as f: + header = f.read(32) + except Exception: + return None + + # Image formats + if header.startswith(b"\x89PNG"): + return "image" + if header.startswith(b"\xff\xd8\xff"): + return "image" + if header.startswith(b"RIFF") and b"WEBP" in header[:16]: + return "image" + if header.startswith(b"GIF8"): + return "image" + if header.startswith(b"BM"): + return "image" + + # Audio formats + if header.startswith(b"fLaC"): + return "audio" + if header.startswith(b"RIFF") and b"WAVE" in header[:16]: + return "audio" + if header.startswith(b"ID3") or header.startswith(b"\xff\xfb"): + return "audio" + if header.startswith(b"OggS"): + return "audio" + + # Video formats + if header[4:8] == b"ftyp": # MP4/MOV + return "video" + if header.startswith(b"\x1a\x45\xdf\xa3"): # WebM/MKV + return "video" + + # 3D formats + if header.startswith(b"glTF"): + return "3d" + if b"solid " in header.lower() or header.startswith(b"\x00"): # STL + return "3d" + + return None diff --git a/utils/helpers.py b/utils/helpers.py new file mode 100644 index 0000000..e0bba63 --- /dev/null +++ b/utils/helpers.py @@ -0,0 +1,30 @@ +"""Shared helper functions.""" + +import asyncio +import concurrent.futures + + +def format_size(size_bytes: int) -> str: + """Format bytes as human readable.""" + if size_bytes is None: + return "unknown" + if size_bytes < 1024: + return f"{size_bytes} bytes" + elif size_bytes < 1024 * 1024: + return f"{size_bytes / 1024:.1f} KB" + else: + return f"{size_bytes / (1024 * 1024):.1f} MB" + + +def run_async(coro): + """Run an async coroutine in sync context.""" + try: + loop = asyncio.get_event_loop() + if loop.is_running(): + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(asyncio.run, coro) + return future.result() + else: + return loop.run_until_complete(coro) + except RuntimeError: + return asyncio.run(coro) diff --git a/utils/mime.py b/utils/mime.py new file mode 100644 index 0000000..32153f8 --- /dev/null +++ b/utils/mime.py @@ -0,0 +1,141 @@ +"""MIME type utilities for webhook file handling.""" + +import os +from typing import Optional + +# Comprehensive MIME type mappings for ComfyUI output types +MIME_TYPES = { + # Images + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".webp": "image/webp", + ".gif": "image/gif", + ".bmp": "image/bmp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + + # Audio + ".flac": "audio/flac", + ".mp3": "audio/mpeg", + ".wav": "audio/wav", + ".ogg": "audio/ogg", + ".opus": "audio/opus", + ".m4a": "audio/mp4", + ".aac": "audio/aac", + + # Video + ".mp4": "video/mp4", + ".webm": "video/webm", + ".mov": "video/quicktime", + ".avi": "video/x-msvideo", + ".mkv": "video/x-matroska", + + # 3D Models + ".glb": "model/gltf-binary", + ".gltf": "model/gltf+json", + ".obj": "text/plain", + ".stl": "model/stl", + ".fbx": "application/octet-stream", + ".ply": "application/octet-stream", + + # Data/Other + ".json": "application/json", + ".safetensors": "application/octet-stream", + ".pt": "application/octet-stream", + ".pth": "application/octet-stream", + ".ckpt": "application/octet-stream", + ".latent": "application/octet-stream", + ".txt": "text/plain", + ".csv": "text/csv", +} + +# Reverse mapping for getting extension from MIME type +EXTENSION_FOR_MIME = { + "image/png": ".png", + "image/jpeg": ".jpg", + "image/webp": ".webp", + "image/gif": ".gif", + "audio/flac": ".flac", + "audio/mpeg": ".mp3", + "audio/wav": ".wav", + "audio/ogg": ".ogg", + "audio/opus": ".opus", + "video/mp4": ".mp4", + "video/webm": ".webm", + "model/gltf-binary": ".glb", + "model/gltf+json": ".gltf", + "application/json": ".json", + "text/plain": ".txt", +} + + +def guess_mime_type(filename: str) -> str: + """ + Guess MIME type from filename extension. + + Args: + filename: Filename with extension + + Returns: + MIME type string, defaults to "application/octet-stream" + """ + if not filename: + return "application/octet-stream" + + ext = os.path.splitext(filename.lower())[1] + return MIME_TYPES.get(ext, "application/octet-stream") + + +def get_extension_for_mime(mime_type: str) -> Optional[str]: + """ + Get file extension for a MIME type. + + Args: + mime_type: MIME type string + + Returns: + File extension including dot, or None if unknown + """ + return EXTENSION_FOR_MIME.get(mime_type) + + +def get_output_type_from_mime(mime_type: str) -> str: + """ + Determine output type category from MIME type. + + Args: + mime_type: MIME type string + + Returns: + Output type: image, audio, video, mesh, text, or unknown + """ + if mime_type.startswith("image/"): + return "image" + elif mime_type.startswith("audio/"): + return "audio" + elif mime_type.startswith("video/"): + return "video" + elif mime_type.startswith("model/") or mime_type in ("application/octet-stream",): + # Check if it's a known 3D format + return "mesh" + elif mime_type.startswith("text/") or mime_type == "application/json": + return "text" + else: + return "unknown" + + +def get_format_from_filename(filename: str) -> Optional[str]: + """ + Extract format (extension without dot) from filename. + + Args: + filename: Filename with extension + + Returns: + Format string (e.g., "png", "mp3") or None + """ + if not filename: + return None + ext = os.path.splitext(filename.lower())[1] + return ext[1:] if ext else None diff --git a/utils/validation.py b/utils/validation.py new file mode 100644 index 0000000..789b24a --- /dev/null +++ b/utils/validation.py @@ -0,0 +1,191 @@ +"""Validation utilities for webhook functionality.""" + +import re +import ipaddress +import logging +from urllib.parse import urlparse +from typing import Optional + +logger = logging.getLogger(__name__) + +# Regex for data URL detection +DATA_URL_PATTERN = re.compile(r'^data:([^;,]+)?(;base64)?,(.*)$', re.DOTALL) +URL_PATTERN = re.compile(r'^https?://') + + +def validate_url( + url: str, + require_https: bool = False, + block_private_ips: bool = False, +) -> tuple[bool, Optional[str]]: + """ + Validate a webhook URL for safety and correctness. + + Args: + url: URL to validate + require_https: If True, only allow HTTPS URLs + block_private_ips: If True, block private/local IP addresses + + Returns: + Tuple of (is_valid, error_message) + """ + if not url: + return False, "URL is empty" + + try: + parsed = urlparse(url) + + # Check scheme + if parsed.scheme not in ("http", "https"): + return False, f"Invalid URL scheme: {parsed.scheme}" + + if require_https and parsed.scheme != "https": + return False, "HTTPS is required" + + # Check host + if not parsed.netloc: + return False, "URL has no host" + + # Check for private IPs if requested + if block_private_ips: + hostname = parsed.hostname + if hostname: + # Check for localhost + if hostname in ("localhost", "127.0.0.1", "::1"): + return False, "Localhost URLs are not allowed" + + # Check for private IP ranges + try: + ip = ipaddress.ip_address(hostname) + if ip.is_private or ip.is_loopback or ip.is_link_local: + return False, "Private IP addresses are not allowed" + except ValueError: + # Not an IP address, hostname is fine + pass + + return True, None + + except Exception as e: + return False, f"URL parsing error: {str(e)}" + + +def sanitize_for_logging(data: dict, sensitive_keys: Optional[set] = None) -> dict: + """ + Sanitize a dictionary for safe logging by redacting sensitive values. + + Args: + data: Dictionary to sanitize + sensitive_keys: Set of key names to redact (case-insensitive) + + Returns: + Sanitized copy of the dictionary + """ + if sensitive_keys is None: + sensitive_keys = { + "auth_value", + "authorization", + "api_key", + "apikey", + "token", + "password", + "secret", + "credential", + "bearer", + } + + def redact_value(key: str, value): + if isinstance(value, dict): + return {k: redact_value(k, v) for k, v in value.items()} + elif isinstance(value, list): + return [redact_value(key, item) for item in value] + elif isinstance(value, str): + # Check if key is sensitive + if key.lower() in sensitive_keys: + return "[REDACTED]" + # Check if value looks like base64 data (long strings) + if len(value) > 100 and is_base64_data(value): + return f"[BASE64_DATA: {len(value)} chars]" + return value + + return {k: redact_value(k, v) for k, v in data.items()} + + +def is_base64_data(value: str) -> bool: + """ + Check if a string appears to be base64 encoded data. + + Args: + value: String to check + + Returns: + True if the string looks like base64 data + """ + # Check for data URL prefix + if value.startswith("data:"): + return True + + # Check for base64 character pattern (rough check) + # Base64 uses A-Z, a-z, 0-9, +, /, and = for padding + if len(value) > 50 and re.match(r'^[A-Za-z0-9+/]+=*$', value): + return True + + return False + + +def is_url(value: str) -> bool: + """ + Check if a string is an HTTP(S) URL. + + Args: + value: String to check + + Returns: + True if the string is a valid URL + """ + return bool(URL_PATTERN.match(value)) + + +def is_data_url(value: str) -> bool: + """ + Check if a string is a data URL. + + Args: + value: String to check + + Returns: + True if the string is a data URL + """ + return bool(DATA_URL_PATTERN.match(value)) + + +def parse_data_url(value: str) -> tuple[Optional[str], Optional[bytes]]: + """ + Parse a data URL and extract the content. + + Args: + value: Data URL string + + Returns: + Tuple of (mime_type, decoded_bytes) or (None, None) if invalid + """ + import base64 + + match = DATA_URL_PATTERN.match(value) + if not match: + return None, None + + mime_type = match.group(1) or "application/octet-stream" + is_base64 = match.group(2) is not None + data = match.group(3) + + try: + if is_base64: + decoded = base64.b64decode(data) + else: + # URL-encoded data + from urllib.parse import unquote + decoded = unquote(data).encode() + return mime_type, decoded + except Exception as e: + logger.warning(f"Failed to parse data URL: {e}") + return None, None