diff --git a/oneapi.py b/oneapi.py index 7b585a0..c4247f1 100644 --- a/oneapi.py +++ b/oneapi.py @@ -1,5 +1,6 @@ import os import json +import traceback import uuid import time import copy @@ -21,6 +22,9 @@ path_custom_nodes = os.path.dirname(os.path.dirname(__file__)) path_comfyui_root = os.path.dirname(path_custom_nodes) path_workflows = os.path.join(path_comfyui_root, 'user/default/workflows') +# Output node types for workflow result extraction +OUTPUT_NODE_TYPES = ["SaveImage", "SaveVideo", "VHS_VideoCombine"] + @routes.post('/oneapi/v1/execute') async def execute_workflow(request): """ @@ -36,16 +40,18 @@ async def execute_workflow(request): - Return values: - status: Status (queued, processing, completed, timeout) - prompt_id: Prompt ID - - images: List of image URLs [string, ...] - - images_by_var: Mapped image URLs by variable name {var_name: [string, ...], ...} + - images: List of image URLs [string, ...], only present if there are image results + - images_by_var: Mapped image URLs by variable name {var_name: [string, ...], ...}, only present if there are image results + - videos: List of video URLs [string, ...], only present if there are video results + - videos_by_var: Mapped video URLs by variable name {var_name: [string, ...], ...}, only present if there are video results Node title markers: - Input: Use "$param.field" in node title to map parameter values - - Output: Use "$output" or "$output.name" in SaveImage node title to specify outputs - - "$output.name" - Marks a SaveImage node with a custom output name (added to "images_by_var[name]") + - Output: Use "$output" or "$output.name" in output node title (see OUTPUT_NODE_TYPES) to specify outputs + - "$output.name" - Marks an output node with a custom output name (added to "images_by_var[name]" or "videos_by_var[name]") - If no explicit output marker is set, the node_id is used as the variable name - - Only SaveImage nodes are considered for output - - Both images and images_by_var fields are always included in the response + - Only nodes in OUTPUT_NODE_TYPES are considered for output + - images/images_by_var/videos/videos_by_var fields are only included if there are corresponding results """ try: # Get request data @@ -111,7 +117,7 @@ async def execute_workflow(request): return web.json_response(result) except Exception as e: - print(f"Error executing workflow: {str(e)}") + print(f"Error executing workflow: {str(e)}, {traceback.format_exc()}") return web.json_response({"error": str(e)}, status=500) async def _apply_params_to_workflow(workflow, params): @@ -156,13 +162,18 @@ async def _process_node_params(node_data, params): async def _extract_output_nodes(workflow): """ - Extract SaveImage nodes and their output variable names from workflow + Extract output nodes and their output variable names from workflow + (output node types are defined in OUTPUT_NODE_TYPES) Args: workflow: Workflow JSON object Returns: Dictionary mapping node_id to output variable name + + Note: + Only nodes in OUTPUT_NODE_TYPES are considered as output nodes. + Output fields (images, videos, *_by_var) are only included in the response if there are corresponding results. """ output_id_2_var = {} @@ -171,8 +182,8 @@ async def _extract_output_nodes(workflow): if not _is_valid_node(node_data): continue - # Only process SaveImage nodes - if node_data.get('class_type') != 'SaveImage': + # Only process output nodes (see OUTPUT_NODE_TYPES) + if node_data.get('class_type') not in OUTPUT_NODE_TYPES: continue # Get node title @@ -194,7 +205,7 @@ async def _extract_output_nodes(workflow): # Simple $output without variable name is not valid raise Exception(f"Invalid output marker format (missing name): {part}. Use $output.name format.") - # For SaveImage nodes, always register in output_id_2_var + # For output nodes, always register in output_id_2_var # If no explicit marker, use node_id as the variable name output_id_2_var[node_id] = output_var if output_var else str(node_id) @@ -407,6 +418,108 @@ async def _get_base_url(request): return base_url +# Helper: extract URLs for a specific media type from a node output +def _extract_node_media_urls(node_output, base_url, media_key): + """ + Extract URLs for a specific media type (e.g., 'images', 'gifs') from a node output. + Args: + node_output: Output dict for a node + base_url: Base URL for constructing file URLs + media_key: Key in node_output ('images' or 'gifs') + Returns: + List of URLs for the specified media type + """ + urls = [] + for media_data in node_output.get(media_key, []): + filename = media_data.get("filename") + subfolder = media_data.get("subfolder", "") + media_type = media_data.get("type", "output") + url = f"{base_url}/view?filename={filename}" + if subfolder: + url += f"&subfolder={subfolder}" + if media_type: + url += f"&type={media_type}" + urls.append(url) + return urls + +# Helper: collect all outputs of a given media type from all nodes +def _collect_outputs_by_type(outputs, base_url, media_key): + """ + Collect all outputs of a given media type from all nodes. + Args: + outputs: All outputs dict + base_url: Base URL for constructing file URLs + media_key: Key in node_output ('images' or 'gifs') + Returns: + Dict mapping node_id to list of URLs + """ + result = {} + for node_id, node_output in outputs.items(): + urls = _extract_node_media_urls(node_output, base_url, media_key) + if urls: + result[node_id] = urls + return result + +# Helper: map outputs by variable name +def _map_outputs_by_var(output_id_2_var, output_id_2_media): + """ + Map outputs by variable name using output_id_2_var mapping. + Args: + output_id_2_var: Dict mapping node_id to variable name + output_id_2_media: Dict mapping node_id to list of URLs + Returns: + Dict mapping variable name to list of URLs + """ + result = {} + for node_id, var_name in output_id_2_var.items(): + if node_id in output_id_2_media: + result[var_name] = output_id_2_media[node_id] + return result + +# Helper: flatten all lists in a dict into a single list +def _extend_flat_list_from_dict(media_dict): + """ + Flatten all lists in a dict into a single list. + Args: + media_dict: Dict of lists + Returns: + Flat list of all items + """ + flat = [] + for items in media_dict.values(): + flat.extend(items) + return flat + +def _split_media_by_suffix(node_output, base_url): + """ + Split all media entries in node_output into images/videos by file extension. + Args: + node_output: Output dict for a node + base_url: Base URL for constructing file URLs + Returns: + (images: list, videos: list) - lists of URLs + """ + image_exts = {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff'} + video_exts = {'.mp4', '.mov', '.avi', '.webm', '.gif'} + images = [] + videos = [] + for media_key in ("images", "gifs"): + for media_data in node_output.get(media_key, []): + filename = media_data.get("filename") + subfolder = media_data.get("subfolder", "") + media_type = media_data.get("type", "output") + url = f"{base_url}/view?filename={filename}" + if subfolder: + url += f"&subfolder={subfolder}" + if media_type: + url += f"&type={media_type}" + ext = os.path.splitext(filename)[1].lower() + if ext in image_exts: + images.append(url) + elif ext in video_exts: + videos.append(url) + return images, videos + async def _wait_for_results(prompt_id, timeout=None, request=None, output_id_2_var=None): """Wait for workflow execution results, get history using HTTP API""" start_time = time.time() @@ -414,94 +527,75 @@ async def _wait_for_results(prompt_id, timeout=None, request=None, output_id_2_v "status": "processing", "prompt_id": prompt_id, "images": [], - "images_by_var": {} + "images_by_var": {}, + "videos": [], + "videos_by_var": {} } - - # Get base URL for image URLs + + # Get base URL for image/video URLs base_url = await _get_base_url(request) - + while True: # Check timeout if timeout is not None and timeout > 0 and (time.time() - start_time) > timeout: result["status"] = "timeout" return result - + # Get history using HTTP API try: async with aiohttp.ClientSession() as session: async with session.get(f"http://127.0.0.1:8188/history") as response: if response.status != 200: - # API call failed, retry after waiting await asyncio.sleep(1.0) continue - - # Get entire history history_data = await response.json() - - # Check if specified prompt_id is in history if prompt_id not in history_data: - # Workflow might not be completed yet, retry after waiting await asyncio.sleep(1.0) continue - - # Get history for specific prompt_id prompt_history = history_data[prompt_id] - - # Check if completed if "outputs" in prompt_history: + result["outputs"] = prompt_history["outputs"] result["status"] = "completed" - - # Store all SaveImage node outputs + + # Collect all image and video outputs by file extension output_id_2_images = {} - - # Process outputs, focusing on nodes with images + output_id_2_videos = {} for node_id, node_output in prompt_history["outputs"].items(): - if "images" in node_output: - # Store image URLs for this node - node_images = [] - for img_data in node_output["images"]: - filename = img_data.get("filename") - subfolder = img_data.get("subfolder", "") - img_type = img_data.get("type", "output") - - # Build image URL - img_url = f"{base_url}/view?filename={filename}" - if subfolder: - img_url += f"&subfolder={subfolder}" - if img_type: - img_url += f"&type={img_type}" - - node_images.append(img_url) - - output_id_2_images[node_id] = node_images - - # Process according to output mapping + images, videos = _split_media_by_suffix(node_output, base_url) + if images: + output_id_2_images[node_id] = images + if videos: + output_id_2_videos[node_id] = videos + + # Map by variable name if mapping is available if output_id_2_var and output_id_2_images: - # First pass: handle nodes with explicit mapping - for node_id, output_var in output_id_2_var.items(): - if node_id in output_id_2_images: - # Add to images_by_var using the variable name - result["images_by_var"][output_var] = output_id_2_images[node_id] - - # For backward compatibility: add all output images to the images array as well - for images in output_id_2_images.values(): - result["images"].extend(images) - else: - # No mapping or no image outputs, fallback to including all images in the images array - for node_images in output_id_2_images.values(): - result["images"].extend(node_images) - - # Return complete results + result["images_by_var"] = _map_outputs_by_var(output_id_2_var, output_id_2_images) + result["images"] = _extend_flat_list_from_dict(result["images_by_var"]) + elif output_id_2_images: + result["images"] = _extend_flat_list_from_dict(output_id_2_images) + + if output_id_2_var and output_id_2_videos: + result["videos_by_var"] = _map_outputs_by_var(output_id_2_var, output_id_2_videos) + result["videos"] = _extend_flat_list_from_dict(result["videos_by_var"]) + elif output_id_2_videos: + result["videos"] = _extend_flat_list_from_dict(output_id_2_videos) + + # Remove empty fields for images/videos + if not result["images"]: + result.pop("images") + if not result["images_by_var"]: + result.pop("images_by_var") + if "videos" in result and not result["videos"]: + result.pop("videos") + if "videos_by_var" in result and not result["videos_by_var"]: + result.pop("videos_by_var") + return result except Exception as e: print(f"Error getting history: {str(e)}") - # Continue trying after error - - # Wait before checking again await asyncio.sleep(1.0) # New: Load workflow from local file - def _load_workflow_from_local(filename): """ Load workflow JSON from user/default/workflows directory @@ -513,7 +607,6 @@ def _load_workflow_from_local(filename): return json.load(f) # New: Load workflow from URL - async def _load_workflow_from_url(url): """ Download workflow JSON from URL