From a208cd482bf4a0547603a42e5a4d37f2996d2649 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Mon, 19 Jan 2026 23:33:10 -0800 Subject: [PATCH 1/8] refactor(phase2): extract shared utilities and reduce node duplication - Create shared/filename_utils.py for date/time/dimension filename building - Create shared/path_utils.py for output directory handling - Create shared/discord/message_builder.py for Discord message construction - Create shared/discord/cdn_extractor.py for CDN URL extraction - Refactor image_node.py to use shared utilities (-161 lines, 16.3%) - Refactor video_node.py to use shared utilities (-361 lines, 23.1%) - Fix setup_logging missing from logging_config.py - Fix test imports to use new module paths (nodes.* instead of discord_*_node) - Total reduction: 522 lines (20.5%), exceeding PRD target of ~200 lines Co-Authored-By: Claude Opus 4.5 --- nodes/image_node.py | 288 ++++----------- nodes/video_node.py | 502 ++++---------------------- shared/__init__.py | 26 ++ shared/discord/__init__.py | 27 ++ shared/discord/cdn_extractor.py | 168 +++++++++ shared/discord/message_builder.py | 210 +++++++++++ shared/filename_utils.py | 83 +++++ shared/logging_config.py | 16 + shared/path_utils.py | 83 +++++ tests/test_discord_mock.py | 5 + tests/test_image_node_sanitization.py | 4 +- tests/test_media.py | 2 +- tests/test_security_fixes.py | 16 +- 13 files changed, 766 insertions(+), 664 deletions(-) create mode 100644 shared/discord/cdn_extractor.py create mode 100644 shared/discord/message_builder.py create mode 100644 shared/filename_utils.py create mode 100644 shared/path_utils.py diff --git a/nodes/image_node.py b/nodes/image_node.py index 73f9021..657a19a 100644 --- a/nodes/image_node.py +++ b/nodes/image_node.py @@ -22,7 +22,13 @@ from shared import ( update_github_cdn_urls, extract_prompts_from_workflow, send_to_discord_with_retry, - tensor_to_numpy_uint8 + tensor_to_numpy_uint8, + build_filename_with_metadata, + get_output_directory, + build_metadata_section, + build_prompt_section, + extract_cdn_urls_from_response, + send_cdn_urls_file ) @@ -240,48 +246,24 @@ class DiscordSendSaveImage: if send_workflow_json and extra_pnginfo is not None and "workflow" in extra_pnginfo: extra_pnginfo["workflow"] = sanitize_json_for_export(extra_pnginfo["workflow"]) - # Add date and/or time if enabled - date_time_parts = [] - - # Prepare info for Discord message + # Build filename with date/time metadata using shared utility image_info = {} - - if add_date: - # Get ONLY the date in YYYY-MM-DD format - current_date = time.strftime("%Y-%m-%d") - date_time_parts.append(current_date) - print(f"Adding date to filename: {current_date}") - image_info["date"] = current_date - - if add_time: - # Get ONLY the time in HH-MM-SS format - current_time = time.strftime("%H-%M-%S") - date_time_parts.append(current_time) - print(f"Adding time to filename: {current_time}") - image_info["time"] = current_time - - # Add date/time components to filename prefix if any were enabled - if date_time_parts: - date_time_suffix = "_" + "_".join(date_time_parts) - filename_prefix += date_time_suffix - print(f"Final timestamp suffix: {date_time_suffix}") - + filename_prefix, image_info = build_filename_with_metadata( + prefix=filename_prefix, + add_date=add_date, + add_time=add_time, + info_dict=image_info + ) + # Add prefix append filename_prefix += self.prefix_append - # Get ComfyUI output directory for safe path handling - comfy_output_dir = folder_paths.get_output_directory() - - # Choose destination directory based on save_output flag - if save_output: - # Create a output subfolder in the ComfyUI output directory - dest_folder = os.path.join(comfy_output_dir, "discord_output") - os.makedirs(dest_folder, exist_ok=True) - else: - # Use ComfyUI's temporary directory for preview-only files - dest_folder = folder_paths.get_temp_directory() - os.makedirs(dest_folder, exist_ok=True) - print(f"Using temporary directory for preview: {dest_folder}") + # Get output directory using shared utility + dest_folder = get_output_directory( + save_output=save_output, + comfy_output_dir=folder_paths.get_output_directory(), + temp_dir=folder_paths.get_temp_directory() + ) # Setup paths using ComfyUI's path validation full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( @@ -358,71 +340,39 @@ class DiscordSendSaveImage: elif send_to_discord and not webhook_url: print("Discord integration was enabled but no webhook URL was provided") - # Add image info to Discord message if relevant options are enabled + # Build image info message using shared utility if send_to_discord and webhook_url and (add_date or add_time or add_dimensions or resize_to_power_of_2 or include_format_in_message): - info_message = "\n\n**Image Information:**\n" - - if "date" in image_info: - info_message += f"**Date:** {image_info['date']}\n" - - if "time" in image_info: - info_message += f"**Time:** {image_info['time']}\n" - - # Add format to the message if the option is enabled - if include_format_in_message: - info_message += f"**Format:** {file_format.upper()}\n" - - # Update the message prefix with the information + info_message = build_metadata_section( + info_dict=image_info, + include_date=add_date, + include_time=add_time, + include_dimensions=False, # Dimensions added later after processing + include_format=include_format_in_message, + file_format=file_format, + section_title="Image Information" + ) image_info["message_prefix"] = info_message - - # Note: We don't add to discord_message yet, as dimensions aren't known until processing print("Prepared image information section for Discord message") - - # Extract prompts if requested + + # Extract and build prompts section using shared utilities if send_to_discord and include_prompts_in_message: workflow_data = None - + # First try to get workflow from extra_pnginfo if original_extra_pnginfo is not None and isinstance(original_extra_pnginfo, dict) and "workflow" in original_extra_pnginfo: workflow_data = original_extra_pnginfo["workflow"] - + # If no workflow in extra_pnginfo, check if prompt is actually a workflow if workflow_data is None and original_prompt is not None: - # Check if prompt is already a workflow if isinstance(original_prompt, dict) and "nodes" in original_prompt: workflow_data = original_prompt - - # Extract prompts from workflow data + + # Extract and build prompts section if workflow_data is not None: positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data) - - # Ensure the prompts are strings or None - if positive_prompt is not False and positive_prompt is not None and not isinstance(positive_prompt, str): - positive_prompt = str(positive_prompt) - print(f"Converted positive prompt to string: {positive_prompt[:50]}...") - - if negative_prompt is not False and negative_prompt is not None and not isinstance(negative_prompt, str): - negative_prompt = str(negative_prompt) - print(f"Converted negative prompt to string: {negative_prompt[:50]}...") - - # Check if we have valid prompt data - has_valid_prompt = ( - (isinstance(positive_prompt, str) and positive_prompt) or - (isinstance(negative_prompt, str) and negative_prompt) - ) - - # Add prompts to Discord message if found - if has_valid_prompt: - prompt_message = "\n\n**Generation Prompts:**\n" - - if isinstance(positive_prompt, str) and positive_prompt: - prompt_message += f"**Positive:**\n```\n{positive_prompt}\n```\n" - - if isinstance(negative_prompt, str) and negative_prompt: - prompt_message += f"**Negative:**\n```\n{negative_prompt}\n```\n" - - # Store prompt message for adding after image info - image_info["prompt_message"] = prompt_message + prompt_section = build_prompt_section(positive_prompt, negative_prompt) + if prompt_section: + image_info["prompt_message"] = prompt_section print("Prepared prompts for Discord message") # Optimization: Create metadata once for the entire batch @@ -712,58 +662,17 @@ class DiscordSendSaveImage: if send_workflow_json and "workflow" in files: print(f"Successfully sent workflow JSON for image {batch_number+1}") - # Try to extract CDN URLs from batch response - if save_cdn_urls and response.status_code == 200: - try: - response_data = response.json() - print(f"Received JSON response from Discord with {len(response_data) if isinstance(response_data, dict) else 'invalid'} fields") - - if "attachments" in response_data and isinstance(response_data["attachments"], list): - print(f"Found {len(response_data['attachments'])} attachments in Discord response") - - for idx, attachment in enumerate(response_data["attachments"]): - if "url" in attachment and "filename" in attachment: - # Filter out workflow JSON files from URLs list - if not attachment["filename"].endswith(".json"): - batch_cdn_urls.append((attachment["filename"], attachment["url"])) - print(f"Extracted CDN URL for batch image {idx+1}: {attachment['url']}") - else: - print(f"Skipping JSON file: {attachment['filename']}") - else: - print(f"Attachment {idx+1} missing URL or filename: {attachment.keys()}") - - print(f"Total batch CDN URLs collected: {len(batch_cdn_urls)}") - - # Create and send a text file with the CDN URLs if we have any - if batch_cdn_urls: - try: - # Create the text file content - url_text_content = "# Discord CDN URLs\n\n" - for idx, (filename, url) in enumerate(batch_cdn_urls): - url_text_content += f"{idx+1}. {filename}: {url}\n" - - # Create a unique filename for the text file - urls_filename = f"cdn_urls-{uuid4()}.txt" - - # Prepare the request with just the URL file - url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))} - url_data = {"content": "Discord CDN URLs for the uploaded images:"} - - # Send a follow-up message with just the URLs text file - url_response = send_to_discord_with_retry( - webhook_url, - files=url_files, - data=url_data - ) - - if url_response.status_code in [200, 204]: - print(f"Successfully sent CDN URLs text file to Discord") - else: - print(f"Error sending CDN URLs text file: Status code {url_response.status_code}") - except Exception as e: - print(f"Error creating or sending CDN URLs text file: {e}") - except Exception as e: - print(f"Error extracting CDN URLs from batch response: {e}") + # Extract CDN URLs and send file using shared utility + if save_cdn_urls: + new_urls = extract_cdn_urls_from_response(response) + batch_cdn_urls.extend(new_urls) + if batch_cdn_urls: + send_cdn_urls_file( + webhook_url=webhook_url, + urls=batch_cdn_urls, + send_func=send_to_discord_with_retry, + message="Discord CDN URLs for the uploaded images:" + ) else: print(f"Error: Discord returned status code {response.status_code}") discord_send_success = False @@ -792,32 +701,12 @@ class DiscordSendSaveImage: # If we have CDN URLs and we're not in batch mode, send them as a text file if save_cdn_urls and discord_cdn_urls and not (group_batched_images and len(images) > 1): - try: - # Create the text file content - url_text_content = "# Discord CDN URLs\n\n" - for idx, (filename, url) in enumerate(discord_cdn_urls): - url_text_content += f"{idx+1}. {filename}: {url}\n" - - # Create a unique filename for the text file - urls_filename = f"cdn_urls-{uuid4()}.txt" - - # Prepare the request with just the URL file - url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))} - url_data = {"content": "Discord CDN URLs for the uploaded images:"} - - # Send a follow-up message with just the URLs text file - url_response = send_to_discord_with_retry( - webhook_url, - files=url_files, - data=url_data - ) - - if url_response.status_code in [200, 204]: - print(f"Successfully sent CDN URLs text file to Discord") - else: - print(f"Error sending CDN URLs text file: Status code {url_response.status_code}") - except Exception as e: - print(f"Error creating or sending CDN URLs text file: {e}") + send_cdn_urls_file( + webhook_url=webhook_url, + urls=discord_cdn_urls, + send_func=send_to_discord_with_retry, + message="Discord CDN URLs for the uploaded images:" + ) elif send_to_discord and not discord_send_success: print("DiscordSendSaveImage: There were errors sending some images to Discord") else: @@ -862,58 +751,17 @@ class DiscordSendSaveImage: discord_send_success = True discord_sent_files = ["batch_gallery"] # Mark as successfully sent - # Try to extract CDN URLs from batch response - if save_cdn_urls and response.status_code == 200: - try: - response_data = response.json() - print(f"Received JSON response from Discord with {len(response_data) if isinstance(response_data, dict) else 'invalid'} fields") - - if "attachments" in response_data and isinstance(response_data["attachments"], list): - print(f"Found {len(response_data['attachments'])} attachments in Discord response") - - for idx, attachment in enumerate(response_data["attachments"]): - if "url" in attachment and "filename" in attachment: - # Filter out workflow JSON files from URLs list - if not attachment["filename"].endswith(".json"): - batch_cdn_urls.append((attachment["filename"], attachment["url"])) - print(f"Extracted CDN URL for batch image {idx+1}: {attachment['url']}") - else: - print(f"Skipping JSON file: {attachment['filename']}") - else: - print(f"Attachment {idx+1} missing URL or filename: {attachment.keys()}") - - print(f"Total batch CDN URLs collected: {len(batch_cdn_urls)}") - - # Create and send a text file with the CDN URLs if we have any - if batch_cdn_urls: - try: - # Create the text file content - url_text_content = "# Discord CDN URLs\n\n" - for idx, (filename, url) in enumerate(batch_cdn_urls): - url_text_content += f"{idx+1}. {filename}: {url}\n" - - # Create a unique filename for the text file - urls_filename = f"cdn_urls-{uuid4()}.txt" - - # Prepare the request with just the URL file - url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))} - url_data = {"content": "Discord CDN URLs for the uploaded images:"} - - # Send a follow-up message with just the URLs text file - url_response = send_to_discord_with_retry( - webhook_url, - files=url_files, - data=url_data - ) - - if url_response.status_code in [200, 204]: - print(f"Successfully sent CDN URLs text file to Discord") - else: - print(f"Error sending CDN URLs text file: Status code {url_response.status_code}") - except Exception as e: - print(f"Error creating or sending CDN URLs text file: {e}") - except Exception as e: - print(f"Error extracting CDN URLs from batch response: {e}") + # Extract CDN URLs and send file using shared utility + if save_cdn_urls: + new_urls = extract_cdn_urls_from_response(response) + batch_cdn_urls.extend(new_urls) + if batch_cdn_urls: + send_cdn_urls_file( + webhook_url=webhook_url, + urls=batch_cdn_urls, + send_func=send_to_discord_with_retry, + message="Discord CDN URLs for the uploaded images:" + ) else: print(f"Error sending batch to Discord: Status code {response.status_code} - {response.text}") discord_send_success = False diff --git a/nodes/video_node.py b/nodes/video_node.py index de5563b..b885095 100644 --- a/nodes/video_node.py +++ b/nodes/video_node.py @@ -24,7 +24,19 @@ import functools import server # Import shared utilities -from shared import sanitize_json_for_export, update_github_cdn_urls, send_to_discord_with_retry, tensor_to_numpy_uint8 +from shared import ( + sanitize_json_for_export, + update_github_cdn_urls, + send_to_discord_with_retry, + tensor_to_numpy_uint8, + build_filename_with_metadata, + get_output_directory, + build_metadata_section, + build_prompt_section, + extract_cdn_urls_from_response, + send_cdn_urls_file, + extract_prompts_from_workflow +) # Define cached decorator for local use def cached(max_size=None): """ @@ -361,57 +373,28 @@ class DiscordSendSaveVideo: # Get first image for metadata first_image = images[0] - # Add date and/or time if enabled - date_time_parts = [] - - # Video info for Discord message + # Build filename with date/time/dimensions metadata using shared utility + height, width = images[0].shape[0], images[0].shape[1] video_info = {} - - if add_date: - # Get ONLY the date in YYYY-MM-DD format - current_date = time.strftime("%Y-%m-%d") - date_time_parts.append(current_date) - print(f"Adding date to filename: {current_date}") - video_info["date"] = current_date - - if add_time: - # Get ONLY the time in HH-MM-SS format - current_time = time.strftime("%H-%M-%S") - date_time_parts.append(current_time) - print(f"Adding time to filename: {current_time}") - video_info["time"] = current_time - - # Add dimensions if enabled - if add_dimensions: - # Get dimensions from first frame - height, width = images[0].shape[0], images[0].shape[1] - dim_text = f"{width}x{height}" - date_time_parts.append(dim_text) - print(f"Adding dimensions to filename: {dim_text}") - video_info["dimensions"] = dim_text - - # Add date/time/dimensions components to filename prefix if any were enabled - if date_time_parts: - date_time_suffix = "_" + "_".join(date_time_parts) - filename_prefix += date_time_suffix - print(f"Final metadata suffix: {date_time_suffix}") - + filename_prefix, video_info = build_filename_with_metadata( + prefix=filename_prefix, + add_date=add_date, + add_time=add_time, + add_dimensions=add_dimensions, + width=width, + height=height, + info_dict=video_info + ) + # Add prefix append filename_prefix += self.prefix_append - - # Get ComfyUI output directory for safe path handling - comfy_output_dir = folder_paths.get_output_directory() - - # Choose destination directory based on save_output flag - if save_output: - # Create a video output subfolder in the ComfyUI output directory - dest_folder = os.path.join(comfy_output_dir, "discord_output") - os.makedirs(dest_folder, exist_ok=True) - else: - # Use ComfyUI's temporary directory for preview-only files - dest_folder = folder_paths.get_temp_directory() - os.makedirs(dest_folder, exist_ok=True) - print(f"Using temporary directory for preview: {dest_folder}") + + # Get output directory using shared utility + dest_folder = get_output_directory( + save_output=save_output, + comfy_output_dir=folder_paths.get_output_directory(), + temp_dir=folder_paths.get_temp_directory() + ) # Setup paths using ComfyUI's path validation full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( @@ -982,364 +965,45 @@ class DiscordSendSaveVideo: # Prepare message content message_content = discord_message - - # Add video metadata information to Discord message when enabled - video_metadata_text = "" - - # Only build video metadata text if the option is enabled + + # Add video metadata to Discord message using shared utility if include_video_info: - # Add date if enabled - if add_date and "date" in video_info: - video_metadata_text += f"**Date**: {video_info['date']}\n" - - # Add time if enabled - if add_time and "time" in video_info: - video_metadata_text += f"**Time**: {video_info['time']}\n" - - # Add dimensions if enabled - if add_dimensions and "dimensions" in video_info: - video_metadata_text += f"**Dimensions**: {video_info['dimensions']}\n" - - # Add frame rate information - video_metadata_text += f"**Frame Rate**: {frame_rate} fps\n" - - # Add format information - video_metadata_text += f"**Format**: {format}\n" - - # Add video metadata to message if any was collected - if video_metadata_text and message_content: - message_content += "\n\n**Video Info**:\n" + video_metadata_text - print(f"Added video metadata to Discord message: {len(video_metadata_text)} chars") - elif video_metadata_text: - message_content = "**Video Info**:\n" + video_metadata_text - print(f"Added video metadata to Discord message: {len(video_metadata_text)} chars") + metadata_section = build_metadata_section( + info_dict=video_info, + include_date=add_date, + include_time=add_time, + include_dimensions=add_dimensions, + include_format=True, + file_format=format, + frame_rate=frame_rate, + section_title="Video Info" + ) + if metadata_section: + message_content += metadata_section + print(f"Added video metadata to Discord message: {len(metadata_section)} chars") else: print("Video info was not included in Discord message (disabled by user)") - # Include the generation prompts if requested + # Include the generation prompts if requested using shared utilities if include_prompts_in_message: - print(f"Prompt inclusion requested! Starting prompt extraction process...") - print(f"original_prompt type: {type(original_prompt) if original_prompt is not None else 'None'}") - print(f"original_extra_pnginfo type: {type(original_extra_pnginfo) if original_extra_pnginfo is not None else 'None'}") - - if original_prompt is not None: - print(f"original_prompt keys: {original_prompt.keys() if isinstance(original_prompt, dict) else 'Not a dict'}") - - if original_extra_pnginfo is not None: - print(f"original_extra_pnginfo keys: {original_extra_pnginfo.keys() if isinstance(original_extra_pnginfo, dict) else 'Not a dict'}") - if isinstance(original_extra_pnginfo, dict) and "workflow" in original_extra_pnginfo: - print(f"workflow keys: {original_extra_pnginfo['workflow'].keys() if isinstance(original_extra_pnginfo['workflow'], dict) else 'Workflow not a dict'}") - - # First identify the workflow data to extract prompts from workflow_data = None - + # First try to get workflow from extra_pnginfo if original_extra_pnginfo is not None and isinstance(original_extra_pnginfo, dict) and "workflow" in original_extra_pnginfo: workflow_data = original_extra_pnginfo["workflow"] - print("Found workflow data in extra_pnginfo") - + # If no workflow in extra_pnginfo, check if prompt is actually a workflow if workflow_data is None and original_prompt is not None: - # Check if prompt is already a workflow if isinstance(original_prompt, dict) and "nodes" in original_prompt: workflow_data = original_prompt - print("Found workflow data in original_prompt") - - # Extract prompts if we found workflow data + + # Extract and build prompts section using shared utilities if workflow_data is not None: - print("Workflow data found, attempting to extract prompts...") - - # Extract prompts and add to message - try: - # Try to use discord_image_node's function first - try: - from discord_image_node import extract_prompts_from_workflow - positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data) - - print(f"extract_prompts_from_workflow returned: positive_type={type(positive_prompt)}, negative_type={type(negative_prompt)}") - - # Ensure the prompts are strings or None - if positive_prompt is not False and positive_prompt is not None and not isinstance(positive_prompt, str): - positive_prompt = str(positive_prompt) - print(f"Converted positive prompt to string: {positive_prompt[:50]}...") - - if negative_prompt is not False and negative_prompt is not None and not isinstance(negative_prompt, str): - negative_prompt = str(negative_prompt) - print(f"Converted negative prompt to string: {negative_prompt[:50]}...") - - # Check if we have valid prompt data - has_valid_prompt = ( - (isinstance(positive_prompt, str) and positive_prompt) or - (isinstance(negative_prompt, str) and negative_prompt) - ) - - if has_valid_prompt: - print(f"Successfully extracted prompts via discord_image_node") - - # Add to message - prompt_text = "\n\n**Generation Prompts:**\n" - - if isinstance(positive_prompt, str) and positive_prompt: - prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n" - - if isinstance(negative_prompt, str) and negative_prompt: - prompt_text += f"**Negative:**\n```\n{negative_prompt}\n```\n" - - # Add to message content - message_content += prompt_text - print("Successfully added prompts to Discord message") - else: - print("No valid prompts extracted from discord_image_node function") - # Raise an exception to trigger the fallback extraction method - raise ValueError("No valid prompts extracted") - - except Exception as e: - print(f"Could not use discord_image_node extraction: {str(e)}") - print("Falling back to built-in prompt extraction") - - # Local extraction logic - positive_prompt = None - negative_prompt = None - - # Find CLIP text nodes - clip_nodes = [] - - # Define node types that can contain prompts - prompt_node_types = [ - "CLIPTextEncode", # Standard SD 1.5 prompt node - "SDXLPromptEncoder", # SDXL prompt encoder - "SDXLTextEncode", # Another SDXL text node - ] - - if "nodes" in workflow_data: - nodes = workflow_data["nodes"] - node_list = [] - - # Debug print full node data for the first few nodes - print("\nExamining raw node data:") - if isinstance(nodes, dict): - for i, (node_id, node) in enumerate(list(nodes.items())[:3]): - print(f"Node {i} (ID: {node_id}): {node.get('type', 'unknown')} - Active status: bypassed={node.get('bypassed', 'Not set')}, muted={node.get('muted', 'Not set')}, disabled={node.get('disabled', 'Not set')}") - elif isinstance(nodes, list) and len(nodes) > 0: - for i, node in enumerate(nodes[:3]): - print(f"Node {i}: {node.get('type', 'unknown')} - Active status: bypassed={node.get('bypassed', 'Not set')}, muted={node.get('muted', 'Not set')}, disabled={node.get('disabled', 'Not set')}") - - # Convert to list for processing - if isinstance(nodes, dict): - print("Nodes is a dictionary, converting to list...") - for node_id, node in nodes.items(): - if isinstance(node, dict): - # Store the ID in the node dict for reference - node_with_id = node.copy() # Copy to avoid modifying original - node_with_id["id"] = node_id - node_list.append(node_with_id) - elif isinstance(nodes, list): - print("Nodes is already a list") - node_list = nodes - else: - print(f"Unexpected type for nodes: {type(nodes)}") - - print(f"Processed {len(node_list)} nodes from workflow data") - - # Let's specifically search for all CLIP text nodes to analyze - print("\nSearching for all possible prompt nodes (active or disabled):") - all_prompt_nodes = [] - for node in node_list: - if isinstance(node, dict) and "type" in node: - node_type = node.get("type", "") - if node_type in prompt_node_types or ("Text" in node_type and ("Encode" in node_type or "Prompt" in node_type)): - # Found a potential prompt node - title = node.get("title", "Untitled") - is_active = not (node.get("bypassed", False) or node.get("muted", False) or - node.get("disabled", False) or node.get("active", True) == False or - node.get("enabled", True) == False) - - # Print detailed info - print(f"Prompt node found: {title} (Type: {node_type})") - print(f" - Status: {'ACTIVE' if is_active else 'DISABLED'}") - print(f" - Properties: bypassed={node.get('bypassed', 'Not set')}, muted={node.get('muted', 'Not set')}, disabled={node.get('disabled', 'Not set')}") - if "widgets_values" in node and node["widgets_values"]: - text_preview = node["widgets_values"][0][:50] + "..." if len(node["widgets_values"][0]) > 50 else node["widgets_values"][0] - print(f" - Content preview: {text_preview}") - - all_prompt_nodes.append(node) - - print(f"Found {len(all_prompt_nodes)} total potential prompt nodes ({len([n for n in all_prompt_nodes if not (n.get('bypassed', False) or n.get('muted', False) or n.get('disabled', False) or n.get('active', True) == False or n.get('enabled', True) == False)])} active, {len([n for n in all_prompt_nodes if n.get('bypassed', False) or n.get('muted', False) or n.get('disabled', False) or n.get('active', True) == False or n.get('enabled', True) == False])} disabled)") - - # Check for active prompt nodes (not muted/bypassed) - for node in node_list: - if isinstance(node, dict) and "type" in node: - # Check for various ways a node might be disabled - is_disabled = ( - node.get("bypassed", False) or - node.get("muted", False) or - node.get("disabled", False) or - node.get("active", True) == False or - node.get("enabled", True) == False - ) - - # If no clear status property, let's check outputs to see if they're connected - if not is_disabled and ("outputs" in node or "Output" in node): - # No outputs might mean disconnected - has_outputs = False - if "outputs" in node and node["outputs"]: - has_outputs = True - elif "Output" in node and node["Output"]: - has_outputs = True - - # Check if we should include nodes with no connections (sometimes these are valid sources) - # For prompt nodes, we'll consider them even if disconnected, but log it - if not has_outputs and node["type"] in prompt_node_types: - print(f"Note: Including prompt node that appears disconnected: {node.get('title', node['type'])}") - - if is_disabled: - print(f"Skipping disabled node: {node.get('title', node.get('type', 'unknown'))}") - continue - - # If we reach here, the node is active - # Check standard prompt node types - if node["type"] in prompt_node_types: - if "widgets_values" in node and node["widgets_values"]: - clip_nodes.append(node) - print(f"Found active prompt node: {node.get('title', node['type'])}") - # Also look for other common prompt nodes - elif "Text" in node["type"] and ("Encode" in node["type"] or "Prompt" in node["type"]): - print(f"Found potential prompt node of type: {node['type']}") - if "widgets_values" in node and node["widgets_values"]: - print(f"Values: {node['widgets_values']}") - clip_nodes.append(node) - - print(f"Found {len(clip_nodes)} active prompt nodes") - - # Dump the first 5 nodes to debug - print("Sampling first few nodes for debugging:") - if isinstance(workflow_data["nodes"], dict): - sample_nodes = list(workflow_data["nodes"].values())[:5] - for i, node in enumerate(sample_nodes): - if isinstance(node, dict): - print(f"Node {i}: type={node.get('type', 'unknown')}, title={node.get('title', 'untitled')}") - elif isinstance(workflow_data["nodes"], list): - for i, node in enumerate(workflow_data["nodes"][:5]): - if isinstance(node, dict): - print(f"Node {i}: type={node.get('type', 'unknown')}, title={node.get('title', 'untitled')}") - - # Process CLIP nodes - if len(clip_nodes) == 2: - # Identify which is positive/negative - indicators = ["worst quality", "low quality", "bad quality", "nude", "nsfw"] - scores = [0, 0] - - for i, node in enumerate(clip_nodes): - text = node["widgets_values"][0].lower() - for indicator in indicators: - if indicator in text: - scores[i] += 1 - - # The one with more negative indicators is the negative prompt - if scores[0] > scores[1]: - negative_prompt = clip_nodes[0]["widgets_values"][0] - positive_prompt = clip_nodes[1]["widgets_values"][0] - else: - negative_prompt = clip_nodes[1]["widgets_values"][0] - positive_prompt = clip_nodes[0]["widgets_values"][0] - - print(f"Identified prompts with scores: {scores}") - - # Add to message - prompt_text = "\n\n**Generation Prompts:**\n" - prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n" - prompt_text += f"**Negative:**\n```\n{negative_prompt}\n```\n" - - # Add to message content - message_content += prompt_text - print("Added prompts to Discord message via built-in extraction") - elif len(clip_nodes) == 1: - # Just one prompt, assume it's positive - positive_prompt = clip_nodes[0]["widgets_values"][0] - - # Add to message - prompt_text = "\n\n**Generation Prompts:**\n" - prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n" - - # Add to message content - message_content += prompt_text - print("Added single prompt to Discord message") - else: - # Handle multiple CLIP nodes (more than 2) - print(f"Found {len(clip_nodes)} CLIP nodes, will analyze each to determine positive/negative") - - # Collect potential positive and negative prompts - positive_candidates = [] - negative_candidates = [] - - # Common negative prompt indicators - negative_indicators = ["worst quality", "low quality", "bad quality", "nude", "nsfw", - "negative", "bad", "worse", "poor", "deformed"] - - for node in clip_nodes: - # Get the node's text content - if "widgets_values" in node and node["widgets_values"]: - text = node["widgets_values"][0] - title = node.get("title", "").lower() - node_type = node.get("type", "").lower() - - # Check if title explicitly indicates a negative prompt - is_negative = False - if "negative" in title: - is_negative = True - print(f"Found explicit negative prompt node by title: {title}") - else: - # Check content for negative indicators - text_lower = text.lower() - negative_score = 0 - for indicator in negative_indicators: - if indicator in text_lower: - negative_score += 1 - - if negative_score >= 2: # Threshold for considering as negative - is_negative = True - print(f"Identified negative prompt by content with score {negative_score}") - - # Categorize the prompt - if is_negative: - negative_candidates.append(text) - else: - positive_candidates.append(text) - - # Get the final prompts - positive_prompt = None - negative_prompt = None - - if positive_candidates: - # Use the longest positive prompt as it likely has more information - positive_prompt = max(positive_candidates, key=len) - print(f"Selected positive prompt ({len(positive_prompt)} chars)") - - if negative_candidates: - # Use the longest negative prompt - negative_prompt = max(negative_candidates, key=len) - print(f"Selected negative prompt ({len(negative_prompt)} chars)") - - # Add to message if any prompts were found - if positive_prompt or negative_prompt: - prompt_text = "\n\n**Generation Prompts:**\n" - - if positive_prompt: - prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n" - - if negative_prompt: - prompt_text += f"**Negative:**\n```\n{negative_prompt}\n```\n" - - # Add to message content - message_content += prompt_text - print(f"Added prompts to Discord message from {len(positive_candidates)} positive and {len(negative_candidates)} negative candidates") - else: - print("Could not identify any positive or negative prompts from the CLIP nodes") - except Exception as e: - print(f"Error extracting prompts: {str(e)}") - print(f"Error type: {type(e).__name__}") - import traceback - traceback.print_exc() + positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data) + prompt_section = build_prompt_section(positive_prompt, negative_prompt) + if prompt_section: + message_content += prompt_section + print("Successfully added prompts to Discord message") else: print("No workflow data found for prompt extraction") @@ -1426,20 +1090,10 @@ class DiscordSendSaveVideo: discord_sent_files.append(discord_filename) # Store the Discord UUID filename instead of local path print(f"Successfully sent video to Discord with filename: {discord_filename}") - # Try to extract CDN URL if available - if save_cdn_urls and response.status_code == 200: - try: - response_data = response.json() - # Discord webhook responses include attachments with URLs - if "attachments" in response_data and isinstance(response_data["attachments"], list): - for attachment in response_data["attachments"]: - if "url" in attachment and "filename" in attachment: - # Filter out workflow JSON files - if not attachment["filename"].endswith(".json"): - discord_cdn_urls.append((attachment["filename"], attachment["url"])) - print(f"Extracted CDN URL for video: {attachment['url']}") - except Exception as e: - print(f"Error extracting CDN URL from response: {e}") + # Extract CDN URLs using shared utility + if save_cdn_urls: + new_urls = extract_cdn_urls_from_response(response) + discord_cdn_urls.extend(new_urls) else: print(f"Discord API error: {response.status_code} - {response.text}") discord_send_success = False @@ -1457,34 +1111,14 @@ class DiscordSendSaveVideo: except Exception as e: print(f"Error cleaning up temporary file: {e}") - # If we have CDN URLs and the option is enabled, send them as a text file + # Send CDN URLs file using shared utility if save_cdn_urls and discord_cdn_urls: - try: - # Create the text file content - url_text_content = "# Discord CDN URLs\n\n" - for idx, (filename, url) in enumerate(discord_cdn_urls): - url_text_content += f"{idx+1}. {filename}: {url}\n" - - # Create a unique filename for the text file - urls_filename = f"cdn_urls-{uuid4()}.txt" - - # Prepare the request with just the URL file - url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))} - url_data = {"content": "Discord CDN URLs for the uploaded videos:"} - - # Send a follow-up message with just the URLs text file - url_response = send_to_discord_with_retry( - webhook_url, - files=url_files, - data=url_data - ) - - if url_response.status_code in [200, 204]: - print(f"Successfully sent CDN URLs text file to Discord") - else: - print(f"Error sending CDN URLs text file: Status code {url_response.status_code}") - except Exception as e: - print(f"Error creating or sending CDN URLs text file: {e}") + send_cdn_urls_file( + webhook_url=webhook_url, + urls=discord_cdn_urls, + send_func=send_to_discord_with_retry, + message="Discord CDN URLs for the uploaded videos:" + ) # Update GitHub repository with CDN URLs if enabled if github_cdn_update and send_to_discord and discord_cdn_urls: diff --git a/shared/__init__.py b/shared/__init__.py index e1e7f59..4b89ccc 100644 --- a/shared/__init__.py +++ b/shared/__init__.py @@ -17,9 +17,21 @@ from .discord.webhook_client import ( validate_webhook_url, send_to_discord_with_retry ) +from .discord.message_builder import ( + build_metadata_section, + build_prompt_section, + build_discord_message, + format_file_size +) +from .discord.cdn_extractor import ( + extract_cdn_urls_from_response, + send_cdn_urls_file +) from .media.image_processing import tensor_to_numpy_uint8 from .github_integration import update_github_cdn_urls from .logging_config import setup_logging, get_logger +from .filename_utils import build_filename_with_metadata, get_timestamp_string +from .path_utils import get_output_directory, ensure_directory_exists __all__ = [ # Workflow utilities @@ -30,6 +42,14 @@ __all__ = [ 'DiscordWebhookClient', 'validate_webhook_url', 'send_to_discord_with_retry', + # Discord message building + 'build_metadata_section', + 'build_prompt_section', + 'build_discord_message', + 'format_file_size', + # CDN extraction + 'extract_cdn_urls_from_response', + 'send_cdn_urls_file', # Media utilities 'tensor_to_numpy_uint8', # GitHub integration @@ -37,4 +57,10 @@ __all__ = [ # Logging 'setup_logging', 'get_logger', + # Filename utilities + 'build_filename_with_metadata', + 'get_timestamp_string', + # Path utilities + 'get_output_directory', + 'ensure_directory_exists', ] diff --git a/shared/discord/__init__.py b/shared/discord/__init__.py index 2b3c397..f321e32 100644 --- a/shared/discord/__init__.py +++ b/shared/discord/__init__.py @@ -11,11 +11,38 @@ from .webhook_client import ( send_to_discord_with_retry, validate_file_for_discord ) +from .message_builder import ( + build_metadata_section, + build_prompt_section, + build_discord_message, + validate_message_content, + format_file_info, + format_file_size +) +from .cdn_extractor import ( + extract_cdn_urls_from_response, + create_cdn_urls_content, + send_cdn_urls_file, + collect_and_send_cdn_urls +) __all__ = [ + # Webhook client 'DiscordWebhookClient', 'validate_webhook_url', 'sanitize_webhook_for_logging', 'send_to_discord_with_retry', 'validate_file_for_discord', + # Message building + 'build_metadata_section', + 'build_prompt_section', + 'build_discord_message', + 'validate_message_content', + 'format_file_info', + 'format_file_size', + # CDN extraction + 'extract_cdn_urls_from_response', + 'create_cdn_urls_content', + 'send_cdn_urls_file', + 'collect_and_send_cdn_urls', ] diff --git a/shared/discord/cdn_extractor.py b/shared/discord/cdn_extractor.py new file mode 100644 index 0000000..f3f1033 --- /dev/null +++ b/shared/discord/cdn_extractor.py @@ -0,0 +1,168 @@ +""" +Discord CDN URL extraction utilities for ComfyUI-DiscordSend + +Provides functions for extracting CDN URLs from Discord responses +and creating/sending URL text files. +""" + +from typing import List, Tuple, Optional, Any +from uuid import uuid4 + + +def extract_cdn_urls_from_response( + response: Any, + exclude_json: bool = True +) -> List[Tuple[str, str]]: + """ + Extract CDN URLs from a Discord webhook response. + + Args: + response: Response object from Discord API (must have .status_code and .json()) + exclude_json: Whether to exclude .json files from results + + Returns: + List of (filename, url) tuples + """ + cdn_urls = [] + + if response.status_code != 200: + return cdn_urls + + try: + response_data = response.json() + print(f"Received JSON response from Discord with " + f"{len(response_data) if isinstance(response_data, dict) else 'invalid'} fields") + + if "attachments" in response_data and isinstance(response_data["attachments"], list): + print(f"Found {len(response_data['attachments'])} attachments in Discord response") + + for idx, attachment in enumerate(response_data["attachments"]): + if "url" in attachment and "filename" in attachment: + filename = attachment["filename"] + url = attachment["url"] + + # Filter out workflow JSON files if requested + if exclude_json and filename.endswith(".json"): + print(f"Skipping JSON file: {filename}") + continue + + cdn_urls.append((filename, url)) + print(f"Extracted CDN URL for attachment {idx + 1}: {url}") + else: + print(f"Attachment {idx + 1} missing URL or filename: {attachment.keys()}") + + print(f"Total CDN URLs collected: {len(cdn_urls)}") + + except Exception as e: + print(f"Error extracting CDN URLs from response: {e}") + + return cdn_urls + + +def create_cdn_urls_content( + urls: List[Tuple[str, str]], + header: str = "# Discord CDN URLs\n\n" +) -> str: + """ + Create text content from a list of CDN URLs. + + Args: + urls: List of (filename, url) tuples + header: Header text for the content + + Returns: + Formatted text content + """ + content = header + for idx, (filename, url) in enumerate(urls): + content += f"{idx + 1}. {filename}: {url}\n" + return content + + +def send_cdn_urls_file( + webhook_url: str, + urls: List[Tuple[str, str]], + send_func: Any, + message: str = "Discord CDN URLs for the uploaded files:", + filename_prefix: str = "cdn_urls" +) -> bool: + """ + Create and send a text file containing CDN URLs to Discord. + + Args: + webhook_url: Discord webhook URL + urls: List of (filename, url) tuples + send_func: Function to send to Discord (send_to_discord_with_retry) + message: Message to accompany the file + filename_prefix: Prefix for the generated filename + + Returns: + True if successful, False otherwise + """ + if not urls: + print("No CDN URLs to send") + return False + + try: + # Create the text file content + url_text_content = create_cdn_urls_content(urls) + + # Create a unique filename for the text file + urls_filename = f"{filename_prefix}-{uuid4()}.txt" + + # Prepare the request with just the URL file + url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))} + url_data = {"content": message} + + # Send a follow-up message with just the URLs text file + url_response = send_func( + webhook_url, + files=url_files, + data=url_data + ) + + if url_response.status_code in [200, 204]: + print(f"Successfully sent CDN URLs text file to Discord") + return True + else: + print(f"Error sending CDN URLs text file: Status code {url_response.status_code}") + return False + + except Exception as e: + print(f"Error creating or sending CDN URLs text file: {e}") + return False + + +def collect_and_send_cdn_urls( + response: Any, + webhook_url: str, + send_func: Any, + save_cdn_urls: bool, + existing_urls: Optional[List[Tuple[str, str]]] = None, + message: str = "Discord CDN URLs for the uploaded files:" +) -> List[Tuple[str, str]]: + """ + Convenience function to extract CDN URLs from a response and optionally send them. + + Args: + response: Discord webhook response + webhook_url: Webhook URL for sending the URLs file + send_func: Function to send to Discord + save_cdn_urls: Whether to extract and save CDN URLs + existing_urls: Existing URLs to append to (for batch operations) + message: Message to accompany the URLs file + + Returns: + List of all collected CDN URLs + """ + if existing_urls is None: + existing_urls = [] + + if not save_cdn_urls: + return existing_urls + + # Extract URLs from this response + new_urls = extract_cdn_urls_from_response(response) + all_urls = existing_urls + new_urls + + return all_urls diff --git a/shared/discord/message_builder.py b/shared/discord/message_builder.py new file mode 100644 index 0000000..b7cc106 --- /dev/null +++ b/shared/discord/message_builder.py @@ -0,0 +1,210 @@ +""" +Discord message building utilities for ComfyUI-DiscordSend + +Provides functions for constructing Discord messages with metadata, +prompts, and other formatted content. +""" + +from typing import Dict, Any, Optional, Tuple, List + + +def build_metadata_section( + info_dict: Dict[str, Any], + include_date: bool = True, + include_time: bool = True, + include_dimensions: bool = True, + include_format: bool = True, + file_format: Optional[str] = None, + frame_rate: Optional[float] = None, + section_title: str = "Information" +) -> str: + """ + Build a formatted metadata section for Discord messages. + + Args: + info_dict: Dictionary containing metadata (date, time, dimensions, etc.) + include_date: Whether to include date if present + include_time: Whether to include time if present + include_dimensions: Whether to include dimensions if present + include_format: Whether to include format information + file_format: File format string (e.g., "png", "mp4") + frame_rate: Frame rate for video (optional) + section_title: Title for the section (e.g., "Image Information", "Video Info") + + Returns: + Formatted metadata string, or empty string if no metadata + """ + metadata_lines = [] + + if include_date and "date" in info_dict: + metadata_lines.append(f"**Date:** {info_dict['date']}") + + if include_time and "time" in info_dict: + metadata_lines.append(f"**Time:** {info_dict['time']}") + + if include_dimensions and "dimensions" in info_dict: + metadata_lines.append(f"**Dimensions:** {info_dict['dimensions']}") + + if frame_rate is not None: + metadata_lines.append(f"**Frame Rate:** {frame_rate} fps") + + if include_format and file_format: + metadata_lines.append(f"**Format:** {file_format.upper()}") + + if not metadata_lines: + return "" + + section = f"\n\n**{section_title}:**\n" + section += "\n".join(metadata_lines) + return section + + +def build_prompt_section( + positive_prompt: Optional[str], + negative_prompt: Optional[str], + section_title: str = "Generation Prompts" +) -> str: + """ + Build a formatted prompts section for Discord messages. + + Args: + positive_prompt: The positive/main prompt text + negative_prompt: The negative prompt text + section_title: Title for the section + + Returns: + Formatted prompt string, or empty string if no prompts + """ + # Validate and normalize prompts + if positive_prompt is not None and not isinstance(positive_prompt, str): + positive_prompt = str(positive_prompt) + if negative_prompt is not None and not isinstance(negative_prompt, str): + negative_prompt = str(negative_prompt) + + has_positive = isinstance(positive_prompt, str) and positive_prompt.strip() + has_negative = isinstance(negative_prompt, str) and negative_prompt.strip() + + if not has_positive and not has_negative: + return "" + + section = f"\n\n**{section_title}:**\n" + + if has_positive: + section += f"**Positive:**\n```\n{positive_prompt.strip()}\n```\n" + + if has_negative: + section += f"**Negative:**\n```\n{negative_prompt.strip()}\n```\n" + + return section + + +def build_discord_message( + base_message: str = "", + metadata_section: str = "", + prompt_section: str = "", + additional_sections: Optional[List[str]] = None, + max_length: int = 2000 +) -> str: + """ + Build a complete Discord message from components. + + Args: + base_message: The main message content + metadata_section: Pre-built metadata section + prompt_section: Pre-built prompt section + additional_sections: List of additional section strings + max_length: Maximum message length (Discord limit is 2000) + + Returns: + Complete formatted message, truncated if necessary + """ + parts = [base_message] if base_message else [] + + if metadata_section: + parts.append(metadata_section) + + if prompt_section: + parts.append(prompt_section) + + if additional_sections: + parts.extend(additional_sections) + + message = "".join(parts) + + # Truncate if necessary + if len(message) > max_length: + truncation_notice = "\n...[Message truncated]" + message = message[:max_length - len(truncation_notice)] + truncation_notice + + return message + + +def validate_message_content(message: str) -> Tuple[bool, str]: + """ + Validate Discord message content. + + Args: + message: Message content to validate + + Returns: + Tuple of (is_valid, validation_message) + """ + if not message: + return True, "Empty message (valid for file-only uploads)" + + if len(message) > 2000: + return False, f"Message exceeds 2000 character limit ({len(message)} chars)" + + # Check for required sections (informational) + has_prompts = "Generation Prompts" in message + + info_parts = [] + info_parts.append(f"Message has {message.count(chr(10))} lines") + + if not has_prompts: + info_parts.append("WARNING: Message does NOT contain 'Generation Prompts' section") + + return True, "\n".join(info_parts) + + +def format_file_info( + filename: str, + file_size: int, + mime_type: Optional[str] = None +) -> str: + """ + Format file information for logging/display. + + Args: + filename: Name of the file + file_size: Size in bytes + mime_type: MIME type of the file + + Returns: + Formatted string with file information + """ + size_str = format_file_size(file_size) + info = f"File: {filename} ({size_str})" + if mime_type: + info += f" [{mime_type}]" + return info + + +def format_file_size(size_bytes: int) -> str: + """ + Format file size in human-readable format. + + Args: + size_bytes: Size in bytes + + Returns: + Formatted string (e.g., "1.5 MB", "256 KB") + """ + if size_bytes < 1024: + return f"{size_bytes} bytes" + elif size_bytes < 1024 * 1024: + return f"{size_bytes / 1024:.1f} KB" + elif size_bytes < 1024 * 1024 * 1024: + return f"{size_bytes / (1024 * 1024):.1f} MB" + else: + return f"{size_bytes / (1024 * 1024 * 1024):.2f} GB" diff --git a/shared/filename_utils.py b/shared/filename_utils.py new file mode 100644 index 0000000..75d7883 --- /dev/null +++ b/shared/filename_utils.py @@ -0,0 +1,83 @@ +""" +Filename utilities for ComfyUI-DiscordSend + +Provides functions for building filenames with date, time, and dimension metadata. +""" + +import time +from typing import Dict, Optional, Tuple, Any + + +def build_filename_with_metadata( + prefix: str, + add_date: bool = False, + add_time: bool = False, + add_dimensions: bool = False, + width: Optional[int] = None, + height: Optional[int] = None, + info_dict: Optional[Dict[str, Any]] = None +) -> Tuple[str, Dict[str, Any]]: + """ + Build a filename with optional date, time, and dimension suffixes. + + Args: + prefix: The base filename prefix + add_date: Whether to add the current date (YYYY-MM-DD) + add_time: Whether to add the current time (HH-MM-SS) + add_dimensions: Whether to add dimensions (WxH) + width: Image/video width (required if add_dimensions is True) + height: Image/video height (required if add_dimensions is True) + info_dict: Optional dict to update with metadata (creates new if None) + + Returns: + Tuple of (modified_prefix, info_dict with metadata) + """ + if info_dict is None: + info_dict = {} + + metadata_parts = [] + + if add_date: + current_date = time.strftime("%Y-%m-%d") + metadata_parts.append(current_date) + info_dict["date"] = current_date + print(f"Adding date to filename: {current_date}") + + if add_time: + current_time = time.strftime("%H-%M-%S") + metadata_parts.append(current_time) + info_dict["time"] = current_time + print(f"Adding time to filename: {current_time}") + + if add_dimensions and width is not None and height is not None: + dim_text = f"{width}x{height}" + metadata_parts.append(dim_text) + info_dict["dimensions"] = dim_text + print(f"Adding dimensions to filename: {dim_text}") + + modified_prefix = prefix + if metadata_parts: + metadata_suffix = "_" + "_".join(metadata_parts) + modified_prefix += metadata_suffix + print(f"Final metadata suffix: {metadata_suffix}") + + return modified_prefix, info_dict + + +def get_timestamp_string(include_date: bool = True, include_time: bool = True) -> str: + """ + Get a formatted timestamp string. + + Args: + include_date: Include date in format YYYY-MM-DD + include_time: Include time in format HH-MM-SS + + Returns: + Formatted timestamp string + """ + parts = [] + if include_date: + parts.append(time.strftime("%Y-%m-%d")) + if include_time: + parts.append(time.strftime("%H-%M-%S")) + return "_".join(parts) if parts else "" diff --git a/shared/logging_config.py b/shared/logging_config.py index 4bff9c0..323ee96 100644 --- a/shared/logging_config.py +++ b/shared/logging_config.py @@ -8,6 +8,22 @@ import logging import sys +def setup_logging(level: int = logging.INFO) -> None: + """ + Set up logging configuration for the application. + + Args: + level: The logging level to use (default: INFO) + """ + logging.basicConfig( + level=level, + format='[%(name)s] %(levelname)s: %(message)s', + handlers=[ + logging.StreamHandler(sys.stdout) + ] + ) + + def get_logger(name: str = "comfyui_discordsend") -> logging.Logger: """ Get a configured logger for the extension. diff --git a/shared/path_utils.py b/shared/path_utils.py new file mode 100644 index 0000000..d55a0c4 --- /dev/null +++ b/shared/path_utils.py @@ -0,0 +1,83 @@ +""" +Path utilities for ComfyUI-DiscordSend + +Provides functions for handling output directories and file paths. +""" + +import os +from typing import Optional + + +def get_output_directory( + save_output: bool, + comfy_output_dir: str, + temp_dir: str, + subfolder: str = "discord_output" +) -> str: + """ + Determine the appropriate output directory based on save settings. + + Args: + save_output: Whether files should be saved permanently + comfy_output_dir: ComfyUI's output directory path + temp_dir: ComfyUI's temporary directory path + subfolder: Subfolder name within output directory (default: "discord_output") + + Returns: + Path to the destination directory + """ + if save_output: + # Create output subfolder in the ComfyUI output directory + dest_folder = os.path.join(comfy_output_dir, subfolder) + os.makedirs(dest_folder, exist_ok=True) + else: + # Use ComfyUI's temporary directory for preview-only files + dest_folder = temp_dir + os.makedirs(dest_folder, exist_ok=True) + print(f"Using temporary directory for preview: {dest_folder}") + + return dest_folder + + +def ensure_directory_exists(path: str) -> str: + """ + Ensure a directory exists, creating it if necessary. + + Args: + path: Directory path to ensure exists + + Returns: + The same path (for chaining) + """ + os.makedirs(path, exist_ok=True) + return path + + +def get_unique_filepath( + directory: str, + filename: str, + extension: str, + counter: Optional[int] = None +) -> str: + """ + Generate a unique filepath, optionally with a counter. + + Args: + directory: Base directory + filename: Base filename (without extension) + extension: File extension (with or without leading dot) + counter: Optional counter to append to filename + + Returns: + Full filepath + """ + # Ensure extension has leading dot + if not extension.startswith("."): + extension = "." + extension + + if counter is not None: + full_filename = f"{filename}_{counter:05d}{extension}" + else: + full_filename = f"{filename}{extension}" + + return os.path.join(directory, full_filename) diff --git a/tests/test_discord_mock.py b/tests/test_discord_mock.py index 82ecb2d..24ccbd6 100644 --- a/tests/test_discord_mock.py +++ b/tests/test_discord_mock.py @@ -3,6 +3,11 @@ import sys import os from unittest.mock import MagicMock, patch +# Mock dependencies before importing project modules +sys.modules["torch"] = MagicMock() +sys.modules["numpy"] = MagicMock() +sys.modules["cv2"] = MagicMock() + # Add parent directory to path for imports sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) diff --git a/tests/test_image_node_sanitization.py b/tests/test_image_node_sanitization.py index aed9b0e..16212c0 100644 --- a/tests/test_image_node_sanitization.py +++ b/tests/test_image_node_sanitization.py @@ -27,10 +27,12 @@ sys.modules['comfy'] = MagicMock() sys.modules['comfy.cli_args'] = MagicMock() sys.modules['comfy.cli_args'].args = MagicMock() sys.modules['comfy.cli_args'].args.disable_metadata = False +sys.modules['comfy.utils'] = MagicMock() sys.modules['folder_paths'] = MagicMock() sys.modules['folder_paths'].get_output_directory = MagicMock(return_value="/tmp") sys.modules['folder_paths'].get_temp_directory = MagicMock(return_value="/tmp") sys.modules['folder_paths'].get_save_image_path = MagicMock(return_value=("/tmp", "test", 0, "", "test")) +sys.modules['server'] = MagicMock() # We need real PIL for this test to verify PngInfo try: @@ -39,7 +41,7 @@ except ImportError: # If it failed because it was mocked out and we deleted it, reload pass -from discord_image_node import DiscordSendSaveImage +from nodes.image_node import DiscordSendSaveImage class TestDiscordImageNodeOptimization(unittest.TestCase): def setUp(self): diff --git a/tests/test_media.py b/tests/test_media.py index d78fa9c..ee2cab7 100644 --- a/tests/test_media.py +++ b/tests/test_media.py @@ -20,7 +20,7 @@ sys.modules["server"] = MagicMock() import os sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from discord_video_node import validate_video_for_discord +from nodes.video_node import validate_video_for_discord class TestPathLogic(unittest.TestCase): """Tests for path and file validation logic.""" diff --git a/tests/test_security_fixes.py b/tests/test_security_fixes.py index 62942b9..5a6e132 100644 --- a/tests/test_security_fixes.py +++ b/tests/test_security_fixes.py @@ -26,7 +26,7 @@ sys.modules["folder_paths"] = mock_folder_paths sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # Import the node -from discord_video_node import DiscordSendSaveVideo +from nodes.video_node import DiscordSendSaveVideo class TestTempFileLeak(unittest.TestCase): def setUp(self): @@ -55,7 +55,7 @@ class TestTempFileLeak(unittest.TestCase): self.dummy_image.shape = (512, 512, 3) # height, width, channels # Mock tensor_to_numpy_uint8 in discordsend_utils - self.patcher_numpy = patch("discord_video_node.tensor_to_numpy_uint8") + self.patcher_numpy = patch("nodes.video_node.tensor_to_numpy_uint8") self.mock_numpy_conv = self.patcher_numpy.start() # Return a dummy numpy array import numpy as np @@ -65,12 +65,12 @@ class TestTempFileLeak(unittest.TestCase): self.patcher_numpy.stop() shutil.rmtree(self.test_dir) - @patch("discord_video_node.subprocess.Popen") - @patch("discord_video_node.subprocess.run") - @patch("discord_video_node.send_to_discord_with_retry") - @patch("discord_video_node.Image") - @patch("discord_video_node.os.path.getsize") - @patch("discord_video_node.validate_video_for_discord") + @patch("nodes.video_node.subprocess.Popen") + @patch("nodes.video_node.subprocess.run") + @patch("nodes.video_node.send_to_discord_with_retry") + @patch("nodes.video_node.Image") + @patch("nodes.video_node.os.path.getsize") + @patch("nodes.video_node.validate_video_for_discord") def test_temp_file_leak(self, mock_validate, mock_getsize, mock_image, mock_send, mock_run, mock_popen): # Setup mocks mock_process = MagicMock() From 76da850b3d9e2b72891874fe1a2dcc1b8bb42c91 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Mon, 19 Jan 2026 23:38:58 -0800 Subject: [PATCH 2/8] refactor(phase3): extract video encoding utilities to shared modules - Create shared/media/format_utils.py with format detection and validation - Create shared/media/video_encoder.py with FFmpegEncoder and PILEncoder classes - Extract validate_video_for_discord to shared utility - Extract ffmpeg detection to shared detect_ffmpeg() function - Extract Discord video optimization to shared optimize_video_for_discord() - Remove duplicate code from video_node.py (-109 lines) Phase 3 summary: - video_node.py: 1201 -> 1092 lines (9% reduction this phase) - Total reduction from original: 1562 -> 1092 lines (30% reduction) Co-Authored-By: Claude Opus 4.5 --- nodes/video_node.py | 141 ++-------- shared/media/__init__.py | 29 ++ shared/media/format_utils.py | 137 +++++++++ shared/media/video_encoder.py | 510 ++++++++++++++++++++++++++++++++++ 4 files changed, 692 insertions(+), 125 deletions(-) create mode 100644 shared/media/format_utils.py create mode 100644 shared/media/video_encoder.py diff --git a/nodes/video_node.py b/nodes/video_node.py index b885095..cc9ab70 100644 --- a/nodes/video_node.py +++ b/nodes/video_node.py @@ -37,6 +37,12 @@ from shared import ( send_cdn_urls_file, extract_prompts_from_workflow ) +from shared.media import ( + validate_video_for_discord, + normalize_video_extension, + optimize_video_for_discord as shared_optimize_video, + detect_ffmpeg +) # Define cached decorator for local use def cached(max_size=None): """ @@ -59,14 +65,16 @@ def cached(max_size=None): return decorator(func) return decorator -# Define constants and variables previously imported from discordsend_utils -ffmpeg_path = None +# Define constants ENCODE_ARGS = ("utf-8", "ignore") floatOrInt = ("FLOAT", "INT") imageOrLatent = ("IMAGE", "LATENT") BIGMAX = 1000000 has_vhs_formats = False +# Detect ffmpeg using shared utility +ffmpeg_path = detect_ffmpeg() + # Try to import ProgressBar from comfy.utils try: from comfy.utils import ProgressBar @@ -75,67 +83,10 @@ except ImportError: def __init__(self, total): self.total = total self.current = 0 - + def update(self, advance=1): self.current += advance print(f"Progress: {self.current}/{self.total}", end="\r") - -# Fallback if ffmpeg_path is None, try to detect it -if ffmpeg_path is None: - # Try direct detection methods for ffmpeg - try: - # Try imageio-ffmpeg first (common in Python environments) - try: - import imageio_ffmpeg - ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe() - print(f"Found ffmpeg via imageio_ffmpeg: {ffmpeg_path}") - except (ImportError, Exception): - # Fall back to checking the system path - from shutil import which - ffmpeg_path = which("ffmpeg") - if ffmpeg_path: - print(f"Found ffmpeg in system path: {ffmpeg_path}") - except Exception as e: - print(f"Error during ffmpeg detection: {str(e)}") - -# Function to validate video files for Discord compatibility -def validate_video_for_discord(file_path): - """ - Validate that a video file is compatible with Discord. - Returns a tuple of (is_valid, message) - """ - if not os.path.exists(file_path): - return False, f"File does not exist: {file_path}" - - # Check if file is empty or too small - file_size = os.path.getsize(file_path) - if file_size == 0: - return False, "File is empty" - if file_size < 1024: # Less than 1KB - return False, f"File is suspiciously small: {file_size} bytes" - - # Check if file is too large for Discord (Discord limit is 25MB for regular users, 50MB for Nitro) - max_size = 25 * 1024 * 1024 # 25MB in bytes - if file_size > max_size: - return False, f"File exceeds Discord's size limit: {file_size} bytes (max {max_size} bytes)" - - # Get file extension - ext = os.path.splitext(file_path)[1].lower().lstrip('.') - - # Return validation result based on file type - if ext in ['mp4', 'webm', 'gif']: - # These formats are well supported by Discord - return True, "Valid video format for Discord" - elif ext in ['mov']: - # MOV files (ProRes) may need conversion for Discord - return False, "MOV files may need conversion for Discord compatibility" - elif ext in ['png']: - # PNG sequences are not directly supported by Discord - return False, "PNG sequences are not directly supported by Discord and require compilation into a video format" - else: - # For any other extension, warn but allow sending - return False, f"Unknown format {ext}, may not be compatible with Discord" - class DiscordSendSaveVideo: """ A ComfyUI node that can send videos to Discord and save them with advanced options. @@ -857,73 +808,13 @@ class DiscordSendSaveVideo: discord_optimized_file = None try: # For Discord compatibility, create a special Discord-optimized copy of the video - # This is particularly important when add_time is disabled input_file = output_files[-1] # Get input file (the last output file) + temp_dir = folder_paths.get_temp_directory() - try: - # Create optimized output file in a temp location - temp_dir = folder_paths.get_temp_directory() - discord_optimized_file = os.path.join(temp_dir, f"discord_optimized_{uuid4()}{os.path.splitext(input_file)[1]}") - - # Set up ffmpeg arguments for optimized Discord conversion - # This creates a new file specifically optimized for Discord playback - format_ext = os.path.splitext(input_file)[1].lstrip('.').lower() - optimize_args = [] - - if format_ext == "mp4": - # MP4 optimization for Discord - optimize_args = [ - ffmpeg_path, "-i", input_file, - "-c:v", "libx264", "-pix_fmt", "yuv420p", - "-movflags", "faststart", "-preset", "fast", - "-profile:v", "baseline", "-level", "3.0", - "-crf", "23" - ] - - # Add audio if present in original file - optimize_args.extend(["-c:a", "aac", "-b:a", "128k"]) - - # Add output file - optimize_args.append(discord_optimized_file) - elif format_ext == "webm": - # WebM optimization for Discord - optimize_args = [ - ffmpeg_path, "-i", input_file, - "-c:v", "libvpx-vp9", - "-pix_fmt", "yuv420p", - "-crf", "30", "-b:v", "0", - "-deadline", "good" - ] - - # Add audio if present in original file - optimize_args.extend(["-c:a", "libopus", "-b:a", "96k"]) - - # Add output file - optimize_args.append(discord_optimized_file) - elif format_ext == "gif": - # GIF optimization for Discord - optimize_args = [ - ffmpeg_path, "-i", input_file, - "-vf", "fps=15,scale=trunc(iw/2)*2:trunc(ih/2)*2", - ] - - # Add output file - optimize_args.append(discord_optimized_file) - - if optimize_args: - print(f"Creating Discord-optimized version of {format_ext.upper()} file...") - subprocess.run(optimize_args, check=True, capture_output=True) - print(f"Discord-optimized file created: {discord_optimized_file}") - else: - # If no optimization needed, just use the original file - # Set to None to indicate we're using the original file so we don't clean it up - discord_optimized_file = None - print(f"Using original file for Discord: {input_file}") - - except Exception as e: - print(f"Warning: Failed to create Discord-optimized file: {str(e)}") - print("Falling back to original file") - discord_optimized_file = None # Using original file + # Use shared utility for Discord optimization + discord_optimized_file = shared_optimize_video(input_file, ffmpeg_path, temp_dir) + if discord_optimized_file is None: + print(f"Using original file for Discord: {input_file}") # Prepare Discord files and message discord_files = [] diff --git a/shared/media/__init__.py b/shared/media/__init__.py index f211be9..d051449 100644 --- a/shared/media/__init__.py +++ b/shared/media/__init__.py @@ -5,7 +5,36 @@ Provides image and video processing functions. """ from .image_processing import tensor_to_numpy_uint8 +from .format_utils import ( + parse_format_string, + normalize_video_extension, + get_mime_type, + validate_video_for_discord, + is_animated_format, + supports_alpha +) +from .video_encoder import ( + detect_ffmpeg, + FFmpegEncoder, + PILEncoder, + optimize_video_for_discord, + mux_audio_to_video +) __all__ = [ + # Image processing 'tensor_to_numpy_uint8', + # Format utilities + 'parse_format_string', + 'normalize_video_extension', + 'get_mime_type', + 'validate_video_for_discord', + 'is_animated_format', + 'supports_alpha', + # Video encoding + 'detect_ffmpeg', + 'FFmpegEncoder', + 'PILEncoder', + 'optimize_video_for_discord', + 'mux_audio_to_video', ] diff --git a/shared/media/format_utils.py b/shared/media/format_utils.py new file mode 100644 index 0000000..70110c1 --- /dev/null +++ b/shared/media/format_utils.py @@ -0,0 +1,137 @@ +""" +Video format utilities for ComfyUI-DiscordSend + +Provides format detection, extension mapping, and validation. +""" + +import os +from typing import Tuple, Optional + + +def parse_format_string(format_str: str) -> Tuple[str, str]: + """ + Parse a format string into type and extension. + + Args: + format_str: Format string like "video/h264-mp4" or "image/gif" + + Returns: + Tuple of (format_type, format_extension) + """ + if "/" in format_str: + format_type, format_ext = format_str.split("/", 1) + else: + format_type = "video" + format_ext = format_str + + return format_type, format_ext + + +def normalize_video_extension(format_str: str) -> str: + """ + Normalize a format string to a file extension. + + Args: + format_str: Format string like "video/h264-mp4" + + Returns: + Normalized extension (e.g., "mp4", "webm", "gif") + """ + _, format_ext = parse_format_string(format_str) + + # Map format strings to extensions + extension_map = { + "h264-mp4": "mp4", + "h265-mp4": "mp4", + "vp9-webm": "webm", + "prores": "mov", + } + + return extension_map.get(format_ext, format_ext) + + +def get_mime_type(extension: str) -> str: + """ + Get MIME type for a video extension. + + Args: + extension: File extension (without dot) + + Returns: + MIME type string + """ + mime_types = { + "mp4": "video/mp4", + "webm": "video/webm", + "gif": "image/gif", + "mov": "video/quicktime", + "avi": "video/x-msvideo", + "mkv": "video/x-matroska", + } + return mime_types.get(extension.lower(), "application/octet-stream") + + +def validate_video_for_discord(file_path: str, max_size_mb: int = 25) -> Tuple[bool, str]: + """ + Validate that a video file is compatible with Discord. + + Args: + file_path: Path to the video file + max_size_mb: Maximum file size in megabytes (default 25MB for Discord) + + Returns: + Tuple of (is_valid, message) + """ + if not os.path.exists(file_path): + return False, f"File does not exist: {file_path}" + + file_size = os.path.getsize(file_path) + + if file_size == 0: + return False, "File is empty" + + if file_size < 1024: + return False, f"File is suspiciously small: {file_size} bytes" + + max_size = max_size_mb * 1024 * 1024 + if file_size > max_size: + return False, f"File exceeds Discord's size limit of {max_size_mb}MB ({file_size / (1024*1024):.2f}MB)" + + ext = os.path.splitext(file_path)[1].lower().lstrip('.') + + if ext in ['mp4', 'webm', 'gif']: + return True, "Valid" + elif ext in ['mov']: + return False, "MOV files may need conversion for Discord compatibility" + elif ext in ['png', 'apng']: + return False, "PNG/APNG sequence may need compilation for Discord" + else: + return False, f"Unknown format '{ext}' - may not be compatible with Discord" + + +def is_animated_format(extension: str) -> bool: + """ + Check if a format supports animation. + + Args: + extension: File extension (without dot) + + Returns: + True if the format supports animation + """ + animated_formats = {'gif', 'webp', 'mp4', 'webm', 'mov', 'avi', 'mkv', 'apng'} + return extension.lower() in animated_formats + + +def supports_alpha(extension: str) -> bool: + """ + Check if a format supports alpha channel (transparency). + + Args: + extension: File extension (without dot) + + Returns: + True if the format supports alpha + """ + alpha_formats = {'webm', 'gif', 'webp', 'png', 'apng', 'mov'} + return extension.lower() in alpha_formats diff --git a/shared/media/video_encoder.py b/shared/media/video_encoder.py new file mode 100644 index 0000000..83e663c --- /dev/null +++ b/shared/media/video_encoder.py @@ -0,0 +1,510 @@ +""" +Video encoding utilities for ComfyUI-DiscordSend + +Provides FFmpeg-based video encoding with fallback to PIL for GIF/WebP. +""" + +import os +import subprocess +from typing import List, Tuple, Optional, Iterator, Any, Callable +from uuid import uuid4 +import numpy as np + + +def detect_ffmpeg() -> Optional[str]: + """ + Detect FFmpeg binary location. + + Returns: + Path to FFmpeg executable, or None if not found + """ + ffmpeg_path = None + + # Try imageio-ffmpeg first (common in Python environments) + try: + import imageio_ffmpeg + ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe() + print(f"Found ffmpeg via imageio_ffmpeg: {ffmpeg_path}") + return ffmpeg_path + except (ImportError, Exception): + pass + + # Try system PATH + try: + from shutil import which + ffmpeg_path = which("ffmpeg") + if ffmpeg_path: + print(f"Found ffmpeg in system path: {ffmpeg_path}") + return ffmpeg_path + except Exception: + pass + + return None + + +class FFmpegEncoder: + """ + FFmpeg-based video encoder supporting multiple formats. + """ + + def __init__(self, ffmpeg_path: Optional[str] = None): + """ + Initialize the encoder. + + Args: + ffmpeg_path: Path to FFmpeg executable (auto-detected if None) + """ + self.ffmpeg_path = ffmpeg_path or detect_ffmpeg() + if not self.ffmpeg_path: + raise RuntimeError("FFmpeg not found. Install ffmpeg or imageio-ffmpeg.") + + def encode( + self, + images: List[np.ndarray], + output_path: str, + format_ext: str, + frame_rate: float = 24.0, + quality: int = 85, + lossless: bool = False, + loop_count: int = 0, + codec: Optional[str] = None, + progress_callback: Optional[Callable[[int, int], None]] = None + ) -> str: + """ + Encode images to video using FFmpeg. + + Args: + images: List of numpy arrays (H, W, C) in uint8 format + output_path: Output file path + format_ext: Output format extension (mp4, webm, gif) + frame_rate: Frame rate in FPS + quality: Quality level 1-100 + lossless: Use lossless encoding if supported + loop_count: Number of loops (0 = infinite for GIF) + codec: Specific codec to use (h264, h265, vp9, etc.) + progress_callback: Optional callback(current, total) for progress + + Returns: + Path to the encoded file + """ + if not images: + raise ValueError("No images provided for encoding") + + # Get dimensions from first image + height, width = images[0].shape[:2] + has_alpha = images[0].shape[2] == 4 if len(images[0].shape) > 2 else False + + # Determine input pixel format + i_pix_fmt = "rgba" if has_alpha else "rgb24" + dimensions = f"{width}x{height}" + + # Build FFmpeg arguments + args = self._build_ffmpeg_args( + format_ext=format_ext, + dimensions=dimensions, + frame_rate=frame_rate, + quality=quality, + lossless=lossless, + loop_count=loop_count, + i_pix_fmt=i_pix_fmt, + has_alpha=has_alpha, + codec=codec, + output_path=output_path + ) + + # Execute encoding + self._execute_encoding(args, images, i_pix_fmt, progress_callback) + + return output_path + + def _build_ffmpeg_args( + self, + format_ext: str, + dimensions: str, + frame_rate: float, + quality: int, + lossless: bool, + loop_count: int, + i_pix_fmt: str, + has_alpha: bool, + codec: Optional[str], + output_path: str + ) -> List[str]: + """Build FFmpeg command arguments.""" + + # Loop arguments + loop_args = [] + if format_ext == "gif": + loop_args = ["-loop", "0" if loop_count == 0 else str(loop_count)] + + # Base input arguments + args = [ + self.ffmpeg_path, "-v", "error", + "-f", "rawvideo", + "-pix_fmt", i_pix_fmt, + "-s", dimensions, + "-r", str(frame_rate), + "-i", "-" + ] + loop_args + + # Format-specific encoding arguments + if format_ext == "gif": + args.extend(self._get_gif_args(quality)) + elif format_ext == "mp4": + args.extend(self._get_mp4_args(quality, lossless, codec)) + elif format_ext == "webm": + args.extend(self._get_webm_args(quality, lossless, has_alpha)) + else: + # Default to MP4-like encoding + args.extend(self._get_mp4_args(quality, lossless, codec)) + + # Add output path + args.extend(["-y", output_path]) + + return args + + def _get_gif_args(self, quality: int) -> List[str]: + """Get FFmpeg arguments for GIF encoding.""" + # Use palettegen for better quality + if quality >= 80: + return [ + "-vf", "split[s0][s1];[s0]palettegen=max_colors=256:stats_mode=diff[p];[s1][p]paletteuse=dither=sierra2", + "-f", "gif" + ] + else: + return [ + "-vf", "split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse", + "-f", "gif" + ] + + def _get_mp4_args(self, quality: int, lossless: bool, codec: Optional[str]) -> List[str]: + """Get FFmpeg arguments for MP4 encoding.""" + args = [] + + # Determine codec + use_h265 = codec == "h265" or codec == "hevc" + + if lossless: + if use_h265: + args.extend(["-c:v", "libx265", "-x265-params", "lossless=1"]) + else: + args.extend(["-c:v", "libx264", "-crf", "0"]) + else: + # Map quality (1-100) to CRF (51-0 for h264, lower is better) + crf = int(51 - (quality / 100 * 33)) # Maps 1->51, 100->18 + + if use_h265: + args.extend(["-c:v", "libx265", "-crf", str(crf + 5)]) # H.265 uses different CRF scale + else: + args.extend(["-c:v", "libx264", "-crf", str(crf)]) + + # Always use yuv420p for Discord compatibility + args.extend(["-pix_fmt", "yuv420p", "-movflags", "faststart"]) + + return args + + def _get_webm_args(self, quality: int, lossless: bool, has_alpha: bool) -> List[str]: + """Get FFmpeg arguments for WebM encoding.""" + args = ["-c:v", "libvpx-vp9"] + + if lossless: + args.extend(["-lossless", "1"]) + else: + # Map quality to CRF (63-0 for VP9) + crf = int(63 - (quality / 100 * 33)) # Maps 1->63, 100->30 + args.extend(["-crf", str(crf), "-b:v", "0"]) + + # Pixel format - support alpha if present + pix_fmt = "yuva420p" if has_alpha else "yuv420p" + args.extend(["-pix_fmt", pix_fmt]) + + # VP9 threading + args.extend(["-row-mt", "1"]) + + return args + + def _execute_encoding( + self, + args: List[str], + images: List[np.ndarray], + i_pix_fmt: str, + progress_callback: Optional[Callable[[int, int], None]] = None + ) -> None: + """Execute FFmpeg process and feed frames.""" + total_frames = len(images) + + # Create image chunk iterator for memory efficiency + def image_chunks() -> Iterator[bytes]: + for i, img in enumerate(images): + # Ensure contiguous array for subprocess + chunk = np.ascontiguousarray(img) + if progress_callback: + progress_callback(i + 1, total_frames) + yield chunk.tobytes() + + # Start FFmpeg process + process = subprocess.Popen( + args, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE + ) + + # Feed frames + try: + for chunk in image_chunks(): + process.stdin.write(chunk) + process.stdin.close() + process.wait() + + if process.returncode != 0: + stderr = process.stderr.read().decode('utf-8', errors='ignore') + raise RuntimeError(f"FFmpeg encoding failed: {stderr}") + finally: + if process.stdin: + process.stdin.close() + if process.stdout: + process.stdout.close() + if process.stderr: + process.stderr.close() + + +class PILEncoder: + """ + PIL-based encoder for GIF and WebP formats. + Fallback when FFmpeg is not available. + """ + + def encode( + self, + images: List[Any], # PIL Images or numpy arrays + output_path: str, + format_ext: str, + frame_rate: float = 24.0, + quality: int = 85, + lossless: bool = False, + loop_count: int = 0, + tensor_to_numpy_func: Optional[Callable] = None + ) -> str: + """ + Encode images using PIL. + + Args: + images: List of PIL Images or numpy arrays + output_path: Output file path + format_ext: Output format (gif, webp) + frame_rate: Frame rate in FPS + quality: Quality level 1-100 + lossless: Use lossless encoding for WebP + loop_count: Number of loops (0 = infinite) + tensor_to_numpy_func: Optional function to convert tensors to numpy + + Returns: + Path to the encoded file + """ + from PIL import Image + + # Convert to PIL images if needed + pil_images = [] + for img in images: + if hasattr(img, 'shape'): # numpy array or tensor + if tensor_to_numpy_func and hasattr(img, 'cpu'): + img = tensor_to_numpy_func(img) + elif hasattr(img, 'numpy'): + img = img.numpy() + pil_images.append(Image.fromarray(img.astype(np.uint8))) + else: + pil_images.append(img) + + if not pil_images: + raise ValueError("No images provided for encoding") + + # Calculate frame duration in milliseconds + duration = int(1000 / frame_rate) + + if format_ext.lower() == "gif": + self._encode_gif(pil_images, output_path, duration, loop_count) + elif format_ext.lower() == "webp": + self._encode_webp(pil_images, output_path, duration, loop_count, quality, lossless) + else: + # Single frame fallback + pil_images[0].save(output_path, format=format_ext.upper()) + + return output_path + + def _encode_gif( + self, + images: List[Any], + output_path: str, + duration: int, + loop_count: int + ) -> None: + """Encode images as GIF.""" + durations = [duration] * len(images) + images[0].save( + output_path, + format="GIF", + append_images=images[1:] if len(images) > 1 else [], + save_all=True, + duration=durations, + loop=0 if loop_count == 0 else loop_count, + optimize=False + ) + + def _encode_webp( + self, + images: List[Any], + output_path: str, + duration: int, + loop_count: int, + quality: int, + lossless: bool + ) -> None: + """Encode images as WebP.""" + save_kwargs = { + "format": "WEBP", + "append_images": images[1:] if len(images) > 1 else [], + "save_all": True, + "duration": duration, + "loop": 0 if loop_count == 0 else loop_count, + } + + if lossless: + save_kwargs["lossless"] = True + else: + save_kwargs["quality"] = quality + + images[0].save(output_path, **save_kwargs) + + +def optimize_video_for_discord( + input_file: str, + ffmpeg_path: str, + temp_dir: str +) -> Optional[str]: + """ + Create a Discord-optimized version of a video file. + + Args: + input_file: Path to the input video file + ffmpeg_path: Path to FFmpeg executable + temp_dir: Directory for temporary files + + Returns: + Path to the optimized file, or None if optimization failed + """ + format_ext = os.path.splitext(input_file)[1].lstrip('.').lower() + discord_optimized_file = os.path.join(temp_dir, f"discord_optimized_{uuid4()}.{format_ext}") + + try: + if format_ext == "mp4": + optimize_args = [ + ffmpeg_path, "-i", input_file, + "-c:v", "libx264", "-pix_fmt", "yuv420p", + "-movflags", "faststart", "-preset", "fast", + "-profile:v", "baseline", "-level", "3.0", + "-crf", "23", + "-c:a", "aac", "-b:a", "128k", + "-y", discord_optimized_file + ] + elif format_ext == "webm": + optimize_args = [ + ffmpeg_path, "-i", input_file, + "-c:v", "libvpx-vp9", + "-pix_fmt", "yuv420p", + "-crf", "30", "-b:v", "0", + "-deadline", "good", + "-c:a", "libopus", "-b:a", "96k", + "-y", discord_optimized_file + ] + elif format_ext == "gif": + optimize_args = [ + ffmpeg_path, "-i", input_file, + "-vf", "fps=15,scale=trunc(iw/2)*2:trunc(ih/2)*2", + "-y", discord_optimized_file + ] + else: + print(f"No optimization rules for format: {format_ext}") + return None + + print(f"Creating Discord-optimized version of {format_ext.upper()} file...") + result = subprocess.run( + optimize_args, + capture_output=True, + text=True + ) + + if result.returncode == 0 and os.path.exists(discord_optimized_file): + print(f"Discord-optimized file created: {discord_optimized_file}") + return discord_optimized_file + else: + print(f"Optimization failed: {result.stderr}") + return None + + except Exception as e: + print(f"Error during Discord optimization: {e}") + return None + + +def mux_audio_to_video( + video_path: str, + audio_waveform: np.ndarray, + sample_rate: int, + format_ext: str, + ffmpeg_path: str, + output_path: str, + channels: int = 2 +) -> bool: + """ + Mux audio into a video file. + + Args: + video_path: Path to the video file + audio_waveform: Audio data as numpy array + sample_rate: Audio sample rate + format_ext: Video format extension + ffmpeg_path: Path to FFmpeg executable + output_path: Output path for the muxed file + channels: Number of audio channels + + Returns: + True if successful, False otherwise + """ + try: + # Determine audio codec based on format + if format_ext == "mp4": + audio_pass = ["-c:a", "aac", "-b:a", "192k"] + elif format_ext == "webm": + audio_pass = ["-c:a", "libopus", "-b:a", "128k"] + else: + audio_pass = ["-c:a", "libopus", "-b:a", "128k"] + + mux_args = [ + ffmpeg_path, "-v", "error", "-y", + "-i", video_path, + "-ar", str(sample_rate), + "-ac", str(channels), + "-f", "f32le", + "-i", "-", + "-c:v", "copy" + ] + audio_pass + ["-shortest", output_path] + + # Ensure contiguous array for subprocess + audio_data = np.ascontiguousarray(audio_waveform) + + result = subprocess.run( + mux_args, + input=memoryview(audio_data), + capture_output=True + ) + + if result.returncode == 0: + print(f"Successfully muxed audio to video: {output_path}") + return True + else: + print(f"Audio muxing failed: {result.stderr.decode('utf-8', errors='ignore')}") + return False + + except Exception as e: + print(f"Error muxing audio: {e}") + return False From 87c97e04a3237d061671c880c430ab2e86ab0866 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Mon, 19 Jan 2026 23:41:05 -0800 Subject: [PATCH 3/8] refactor(phase4): create BaseDiscordNode with shared utilities - Create nodes/base_node.py with BaseDiscordNode class - Provide common INPUT_TYPES generators (discord, cdn, filename) - Extract shared methods: - sanitize_workflow_data() - build_filename_prefix() - get_dest_folder() - extract_workflow_from_metadata() - build_prompt_message() - send_discord_files() - send_cdn_urls_to_discord() - update_github_cdn() The base class is ready for nodes to inherit from, enabling gradual adoption of shared functionality without disrupting existing node implementations. Co-Authored-By: Claude Opus 4.5 --- nodes/__init__.py | 3 +- nodes/base_node.py | 343 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 345 insertions(+), 1 deletion(-) create mode 100644 nodes/base_node.py diff --git a/nodes/__init__.py b/nodes/__init__.py index c7bf1e6..c76de77 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -4,7 +4,8 @@ ComfyUI-DiscordSend Node Implementations This package contains the ComfyUI custom nodes for sending media to Discord. """ +from .base_node import BaseDiscordNode from .image_node import DiscordSendSaveImage from .video_node import DiscordSendSaveVideo -__all__ = ['DiscordSendSaveImage', 'DiscordSendSaveVideo'] +__all__ = ['BaseDiscordNode', 'DiscordSendSaveImage', 'DiscordSendSaveVideo'] diff --git a/nodes/base_node.py b/nodes/base_node.py new file mode 100644 index 0000000..fdf27d8 --- /dev/null +++ b/nodes/base_node.py @@ -0,0 +1,343 @@ +""" +Base class for Discord-enabled ComfyUI nodes. + +Provides common functionality for sending media to Discord, +including INPUT_TYPES definitions, sanitization, and Discord integration. +""" + +import os +import folder_paths + +from shared import ( + sanitize_json_for_export, + update_github_cdn_urls, + send_to_discord_with_retry, + build_filename_with_metadata, + get_output_directory, + build_metadata_section, + build_prompt_section, + extract_cdn_urls_from_response, + send_cdn_urls_file, + extract_prompts_from_workflow +) + + +class BaseDiscordNode: + """ + Base class for Discord-enabled ComfyUI nodes. + + Provides common functionality for: + - Filename generation with metadata + - Output directory management + - Discord webhook integration + - GitHub CDN URL updates + - Workflow data sanitization + """ + + def __init__(self): + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + self.output_dir = None + + @staticmethod + def get_discord_input_types(): + """ + Returns Discord-related INPUT_TYPES fields. + + These can be merged into a node's INPUT_TYPES definition. + """ + return { + "send_to_discord": ("BOOLEAN", { + "default": False, + "tooltip": "If enabled, will send the media to Discord via webhook." + }), + "webhook_url": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "Discord webhook URL. Get this from Discord server settings > Integrations > Webhooks." + }), + "discord_message": ("STRING", { + "default": "", + "multiline": True, + "tooltip": "Message to include with the media when sending to Discord." + }), + "include_prompts_in_message": ("BOOLEAN", { + "default": False, + "tooltip": "If enabled, will include the generation prompts in the Discord message." + }), + "send_workflow_json": ("BOOLEAN", { + "default": False, + "tooltip": "If enabled, will send the workflow JSON alongside the media." + }), + } + + @staticmethod + def get_cdn_input_types(): + """ + Returns CDN and GitHub-related INPUT_TYPES fields. + """ + return { + "save_cdn_urls": ("BOOLEAN", { + "default": False, + "tooltip": "If enabled, will extract and save Discord CDN URLs." + }), + "github_cdn_update": ("BOOLEAN", { + "default": False, + "tooltip": "If enabled, will update a GitHub repository with the CDN URLs." + }), + "github_repo": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "GitHub repository in format 'username/repo'." + }), + "github_token": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "GitHub personal access token with repo write permissions." + }), + "github_file_path": ("STRING", { + "default": "cdn_urls.md", + "multiline": False, + "tooltip": "Path to the file in the GitHub repository to update." + }), + } + + @staticmethod + def get_filename_input_types( + add_date_default: bool = False, + add_time_default: bool = True, + add_dimensions_default: bool = False + ): + """ + Returns filename metadata INPUT_TYPES fields. + + Args: + add_date_default: Default value for add_date + add_time_default: Default value for add_time + add_dimensions_default: Default value for add_dimensions + """ + return { + "add_date": ("BOOLEAN", { + "default": add_date_default, + "tooltip": "Add date (YYYY-MM-DD) to the filename." + }), + "add_time": ("BOOLEAN", { + "default": add_time_default, + "tooltip": "Add time (HH-MM-SS) to the filename." + }), + "add_dimensions": ("BOOLEAN", { + "default": add_dimensions_default, + "tooltip": "Add dimensions (WxH) to the filename." + }), + } + + def sanitize_workflow_data(self, prompt, extra_pnginfo): + """ + Sanitize workflow data by removing sensitive information. + + Args: + prompt: The prompt data + extra_pnginfo: Extra PNG info including workflow + + Returns: + Tuple of (sanitized_prompt, sanitized_extra_pnginfo, + original_prompt, original_extra_pnginfo) + """ + # Store original references for prompt extraction + original_prompt = prompt + original_extra_pnginfo = extra_pnginfo + + # Sanitize workflow data + if prompt is not None: + prompt = sanitize_json_for_export(prompt) + + if extra_pnginfo is not None: + extra_pnginfo = sanitize_json_for_export(extra_pnginfo) + + return prompt, extra_pnginfo, original_prompt, original_extra_pnginfo + + def build_filename_prefix( + self, + filename_prefix: str, + add_date: bool, + add_time: bool, + add_dimensions: bool = False, + width: int = None, + height: int = None + ): + """ + Build filename prefix with metadata. + + Returns: + Tuple of (modified_prefix, info_dict) + """ + info_dict = {} + filename_prefix, info_dict = build_filename_with_metadata( + prefix=filename_prefix, + add_date=add_date, + add_time=add_time, + add_dimensions=add_dimensions, + width=width, + height=height, + info_dict=info_dict + ) + filename_prefix += self.prefix_append + return filename_prefix, info_dict + + def get_dest_folder(self, save_output: bool): + """ + Get the destination folder for output files. + + Args: + save_output: Whether to save to output directory (True) or temp (False) + + Returns: + Path to the destination folder + """ + return get_output_directory( + save_output=save_output, + comfy_output_dir=folder_paths.get_output_directory(), + temp_dir=folder_paths.get_temp_directory() + ) + + def extract_workflow_from_metadata(self, original_prompt, original_extra_pnginfo): + """ + Extract workflow data from metadata. + + Args: + original_prompt: Original prompt data + original_extra_pnginfo: Original extra PNG info + + Returns: + Workflow data dict, or None if not found + """ + workflow_data = None + + # First try to get workflow from extra_pnginfo + if (original_extra_pnginfo is not None and + isinstance(original_extra_pnginfo, dict) and + "workflow" in original_extra_pnginfo): + workflow_data = original_extra_pnginfo["workflow"] + + # If no workflow in extra_pnginfo, check if prompt is actually a workflow + if workflow_data is None and original_prompt is not None: + if isinstance(original_prompt, dict) and "nodes" in original_prompt: + workflow_data = original_prompt + + return workflow_data + + def build_prompt_message(self, workflow_data): + """ + Extract and build prompt message from workflow data. + + Args: + workflow_data: Workflow data dict + + Returns: + Formatted prompt section string, or empty string + """ + if workflow_data is None: + return "" + + positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data) + return build_prompt_section(positive_prompt, negative_prompt) + + def send_discord_files( + self, + webhook_url: str, + files: dict, + data: dict, + save_cdn_urls: bool = False + ): + """ + Send files to Discord via webhook. + + Args: + webhook_url: Discord webhook URL + files: Files dict for the request + data: Data dict for the request + save_cdn_urls: Whether to extract CDN URLs from response + + Returns: + Tuple of (success, response, cdn_urls) + """ + cdn_urls = [] + + response = send_to_discord_with_retry( + webhook_url, + files=files, + data=data + ) + + success = response.status_code in [200, 204] + + if success and save_cdn_urls: + cdn_urls = extract_cdn_urls_from_response(response) + + return success, response, cdn_urls + + def send_cdn_urls_to_discord( + self, + webhook_url: str, + cdn_urls: list, + message: str = "Discord CDN URLs:" + ): + """ + Send CDN URLs as a text file to Discord. + + Args: + webhook_url: Discord webhook URL + cdn_urls: List of (filename, url) tuples + message: Message to accompany the file + + Returns: + True if successful + """ + if not cdn_urls: + return False + + return send_cdn_urls_file( + webhook_url=webhook_url, + urls=cdn_urls, + send_func=send_to_discord_with_retry, + message=message + ) + + def update_github_cdn( + self, + cdn_urls: list, + github_repo: str, + github_token: str, + github_file_path: str + ): + """ + Update GitHub repository with CDN URLs. + + Args: + cdn_urls: List of (filename, url) tuples + github_repo: GitHub repository in 'owner/repo' format + github_token: GitHub personal access token + github_file_path: Path to file in repository + + Returns: + Tuple of (success, message) + """ + if not cdn_urls: + return False, "No CDN URLs to update" + + print(f"Updating GitHub repository {github_repo} with {len(cdn_urls)} CDN URLs...") + + success, message = update_github_cdn_urls( + github_repo=github_repo, + github_token=github_token, + file_path=github_file_path, + cdn_urls=cdn_urls + ) + + if success: + print(f"GitHub update successful: {message}") + else: + print(f"GitHub update failed: {message}") + + return success, message From 1e94a69e5f5f948bd8920226d3d7ab0873431ba3 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Tue, 20 Jan 2026 02:13:16 -0800 Subject: [PATCH 4/8] fix(image_node): correct else block indentation in batch Discord send The else block for error handling was incorrectly associated with the `if save_cdn_urls:` condition instead of the `if response.status_code in [200, 204]:` check. This caused discord_send_success to always be set to False after a successful batch send, regardless of the actual outcome. Fixes control flow bug identified in PR #35 review. Co-Authored-By: Claude Opus 4.5 --- nodes/image_node.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nodes/image_node.py b/nodes/image_node.py index 657a19a..3441347 100644 --- a/nodes/image_node.py +++ b/nodes/image_node.py @@ -750,7 +750,7 @@ class DiscordSendSaveImage: print(f"Successfully sent batch of {len(batch_discord_files)} images to Discord as a gallery") discord_send_success = True discord_sent_files = ["batch_gallery"] # Mark as successfully sent - + # Extract CDN URLs and send file using shared utility if save_cdn_urls: new_urls = extract_cdn_urls_from_response(response) @@ -762,8 +762,8 @@ class DiscordSendSaveImage: send_func=send_to_discord_with_retry, message="Discord CDN URLs for the uploaded images:" ) - else: - print(f"Error sending batch to Discord: Status code {response.status_code} - {response.text}") + else: + print(f"Error sending batch to Discord: Status code {response.status_code} - {response.text}") discord_send_success = False except Exception as e: print(f"Error sending batch to Discord: {e}") From b8aadd7faa5fc3e1bf07bfc83ec2e04133679db3 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Tue, 20 Jan 2026 02:27:42 -0800 Subject: [PATCH 5/8] fix(image_node): add section header when only dimensions are displayed When add_dimensions or resize_to_power_of_2 was enabled but add_date, add_time, and include_format_in_message were all disabled, dimensions would be appended to the Discord message without the "Image Information:" section header. Now checks if dimensions will be added but no other metadata exists, and adds the section header in that case. Fixes issue identified in PR #35 review. Co-Authored-By: Claude Opus 4.5 --- nodes/image_node.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/nodes/image_node.py b/nodes/image_node.py index 3441347..9298cf6 100644 --- a/nodes/image_node.py +++ b/nodes/image_node.py @@ -445,17 +445,26 @@ class DiscordSendSaveImage: # Add image info if available if "message_prefix" in image_info: info_message = image_info["message_prefix"] - + + # Check if we need to add dimensions + has_resize_dimensions = "original_dimensions" in image_info and "resized_dimensions" in image_info + has_dimensions = "dimensions" in image_info + + # Add section header if dimensions will be added but no other metadata exists + if (has_resize_dimensions or has_dimensions) and not info_message: + info_message = "\n\n**Image Information:**\n" + # Add dimensions info if available - if "original_dimensions" in image_info and "resized_dimensions" in image_info: + if has_resize_dimensions: info_message += f"**Original Dimensions:** {image_info['original_dimensions']}\n" info_message += f"**Resized Dimensions:** {image_info['resized_dimensions']} (Power of 2)\n" - elif "dimensions" in image_info: + elif has_dimensions: info_message += f"**Dimensions:** {image_info['dimensions']}\n" - + # Add the complete info message to the Discord message - discord_message += info_message - print("Added image information to Discord message") + if info_message: + discord_message += info_message + print("Added image information to Discord message") # Add prompts after image info if available (decoupled from image info presence) if "prompt_message" in image_info: From 36c5d06cf8eda9c46dc136ef4b7477dd8a6f9d10 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Tue, 20 Jan 2026 02:42:25 -0800 Subject: [PATCH 6/8] fix(prompt_extractor): restore SDXL workflow prompt extraction support The refactored prompt extractor only detected CLIPTextEncode nodes, losing support for SDXL-specific nodes that was present in the original code. This caused SDXL workflows using SDXLPromptEncoder or SDXLTextEncode nodes to not have their prompts included in Discord messages. Changes: - Add PROMPT_NODE_TYPES list with CLIPTextEncode, SDXLPromptEncoder, and SDXLTextEncode - Rename functions from clip-specific to generic prompt node naming - Add fallback detection for custom text/encode/prompt nodes - Update docstrings to reflect broader node type support Fixes regression identified in PR #35 review. Co-Authored-By: Claude Opus 4.5 --- shared/workflow/prompt_extractor.py | 100 ++++++++++++++++------------ 1 file changed, 57 insertions(+), 43 deletions(-) diff --git a/shared/workflow/prompt_extractor.py b/shared/workflow/prompt_extractor.py index 6938b2d..4141005 100644 --- a/shared/workflow/prompt_extractor.py +++ b/shared/workflow/prompt_extractor.py @@ -15,17 +15,25 @@ NEGATIVE_INDICATORS = [ "extra limbs", "bad anatomy", "watermark", "text", "signature" ] +# Node types that can contain prompts +PROMPT_NODE_TYPES = [ + "CLIPTextEncode", # Standard SD 1.5 prompt node + "SDXLPromptEncoder", # SDXL prompt encoder + "SDXLTextEncode", # Another SDXL text node +] + def extract_prompts_from_workflow(workflow_data: Any) -> Tuple[Optional[str], Optional[str]]: """ Extract positive and negative prompts from workflow data. - - Analyzes ComfyUI workflow structure to find CLIPTextEncode nodes and + + Analyzes ComfyUI workflow structure to find prompt nodes (CLIPTextEncode, + SDXLPromptEncoder, SDXLTextEncode, and other text encoding nodes) and determine which contains the positive vs negative prompt. - + Args: workflow_data: The workflow data dictionary or JSON string - + Returns: A tuple of (positive_prompt, negative_prompt) or (None, None) if not found """ @@ -47,25 +55,25 @@ def extract_prompts_from_workflow(workflow_data: Any) -> Tuple[Optional[str], Op positive_prompt = None negative_prompt = None - # Find CLIPTextEncode nodes + # Find prompt nodes (CLIPTextEncode, SDXL nodes, etc.) if "nodes" in data: nodes = data["nodes"] else: # Check if it's API format (dict of nodes) nodes = data - - clip_nodes = _find_clip_text_encode_nodes(nodes) - - if not clip_nodes: + + prompt_nodes = _find_prompt_nodes(nodes) + + if not prompt_nodes: return None, None # Determine positive/negative based on content and structure - if len(clip_nodes) == 1: - # Single CLIP node - assume it's the positive prompt - positive_prompt = _get_prompt_text(clip_nodes[0]) - elif len(clip_nodes) >= 2: - # Multiple CLIP nodes - need to determine which is which - positive_prompt, negative_prompt = _classify_prompts(clip_nodes, data) + if len(prompt_nodes) == 1: + # Single prompt node - assume it's the positive prompt + positive_prompt = _get_prompt_text(prompt_nodes[0]) + elif len(prompt_nodes) >= 2: + # Multiple prompt nodes - need to determine which is which + positive_prompt, negative_prompt = _classify_prompts(prompt_nodes, data) # Ensure we return empty string for negative if we have positive but not negative if positive_prompt is not None and negative_prompt is None: @@ -74,41 +82,47 @@ def extract_prompts_from_workflow(workflow_data: Any) -> Tuple[Optional[str], Op return positive_prompt, negative_prompt -def _find_clip_text_encode_nodes(nodes: Union[List, Dict]) -> List[Dict]: - """Find all CLIPTextEncode nodes in the workflow.""" - clip_nodes = [] - +def _find_prompt_nodes(nodes: Union[List, Dict]) -> List[Dict]: + """Find all prompt nodes (CLIPTextEncode, SDXL nodes, etc.) in the workflow.""" + prompt_nodes = [] + if isinstance(nodes, list): for node in nodes: - if _is_clip_text_encode(node): - clip_nodes.append(node) + if _is_prompt_node(node): + prompt_nodes.append(node) elif isinstance(nodes, dict): for node_id, node in nodes.items(): - if _is_clip_text_encode(node): + if _is_prompt_node(node): node_copy = dict(node) node_copy["id"] = node_id - clip_nodes.append(node_copy) - - return clip_nodes + prompt_nodes.append(node_copy) + + return prompt_nodes -def _is_clip_text_encode(node: Any) -> bool: - """Check if a node is a CLIPTextEncode node with valid text.""" +def _is_prompt_node(node: Any) -> bool: + """Check if a node is a prompt node (CLIPTextEncode, SDXL, etc.) with valid text.""" if not isinstance(node, dict): return False - + # Handle both Workflow format (type) and API format (class_type) node_type = node.get("type") or node.get("class_type") - if node_type != "CLIPTextEncode": - return False - + + # Check against known prompt node types + if node_type not in PROMPT_NODE_TYPES: + # Also check for dynamic text/prompt nodes (e.g., custom nodes) + if node_type and ("Text" in node_type and ("Encode" in node_type or "Prompt" in node_type)): + pass # Allow these through + else: + return False + # Check for text in either widgets_values (Workflow) or inputs (API) text = _get_prompt_text(node) return text is not None def _get_prompt_text(node: Dict) -> Optional[str]: - """Extract the prompt text from a CLIP node.""" + """Extract the prompt text from a prompt node.""" # Workflow format (widgets_values) widgets = node.get("widgets_values", []) if isinstance(widgets, list) and len(widgets) > 0 and isinstance(widgets[0], str): @@ -122,20 +136,20 @@ def _get_prompt_text(node: Dict) -> Optional[str]: return None -def _classify_prompts(clip_nodes: List[Dict], workflow_data: Dict) -> Tuple[Optional[str], Optional[str]]: +def _classify_prompts(prompt_nodes: List[Dict], workflow_data: Dict) -> Tuple[Optional[str], Optional[str]]: """ - Classify which CLIP nodes contain positive vs negative prompts. - + Classify which prompt nodes contain positive vs negative prompts. + Uses multiple heuristics: 1. Content analysis (negative prompts often contain quality-related terms) 2. Connection analysis (traces connections to sampler nodes) """ - if not clip_nodes: + if not prompt_nodes: return None, None # First pass: Score all nodes based on content node_scores = [] - for node in clip_nodes: + for node in prompt_nodes: prompt_text = _get_prompt_text(node) # Skip empty or None text if not prompt_text or not prompt_text.strip(): @@ -169,14 +183,14 @@ def _classify_prompts(clip_nodes: List[Dict], workflow_data: Dict) -> Tuple[Opti else: # All scores are 0, use connection analysis positive_prompt, negative_prompt = _classify_by_connections( - clip_nodes, workflow_data, None, None + prompt_nodes, workflow_data, None, None ) # Fallback: if we still can't determine, use first two nodes - if positive_prompt is None and negative_prompt is None and len(clip_nodes) >= 2: + if positive_prompt is None and negative_prompt is None and len(prompt_nodes) >= 2: # Convention: assume first is positive, second is negative - positive_prompt = _get_prompt_text(clip_nodes[0]) - negative_prompt = _get_prompt_text(clip_nodes[1]) + positive_prompt = _get_prompt_text(prompt_nodes[0]) + negative_prompt = _get_prompt_text(prompt_nodes[1]) elif positive_prompt is None and negative_prompt is not None: # Find the other prompt for _, _, text in node_scores: @@ -194,7 +208,7 @@ def _classify_prompts(clip_nodes: List[Dict], workflow_data: Dict) -> Tuple[Opti def _classify_by_connections( - clip_nodes: List[Dict], + prompt_nodes: List[Dict], workflow_data: Dict, existing_positive: Optional[str], existing_negative: Optional[str] @@ -235,7 +249,7 @@ def _classify_by_connections( to_slot = link[3] # Find matching CLIP node and sampler - for clip_node in clip_nodes: + for clip_node in prompt_nodes: clip_id = clip_node.get("id") if clip_id == from_node_id: for sampler in samplers: From cc07d20f86aa449a9155b3141c145a974eec2428 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Tue, 20 Jan 2026 02:54:58 -0800 Subject: [PATCH 7/8] fix(image_node): prevent redundant CDN URL sends on 204 responses The refactored code checked only `if save_cdn_urls:` instead of the original `if save_cdn_urls and response.status_code == 200:`. While extract_cdn_urls_from_response correctly returns empty for 204 responses, accumulated URLs from prior iterations could still trigger the `if batch_cdn_urls:` condition, causing duplicate sends. Fixed in both individual image send and batch send sections. Fixes issue identified in PR #35 review. Co-Authored-By: Claude Opus 4.5 --- nodes/image_node.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nodes/image_node.py b/nodes/image_node.py index 9298cf6..01580bb 100644 --- a/nodes/image_node.py +++ b/nodes/image_node.py @@ -672,7 +672,8 @@ class DiscordSendSaveImage: print(f"Successfully sent workflow JSON for image {batch_number+1}") # Extract CDN URLs and send file using shared utility - if save_cdn_urls: + # Only extract when status is 200 (has content), not 204 (no content) + if save_cdn_urls and response.status_code == 200: new_urls = extract_cdn_urls_from_response(response) batch_cdn_urls.extend(new_urls) if batch_cdn_urls: @@ -761,7 +762,8 @@ class DiscordSendSaveImage: discord_sent_files = ["batch_gallery"] # Mark as successfully sent # Extract CDN URLs and send file using shared utility - if save_cdn_urls: + # Only extract when status is 200 (has content), not 204 (no content) + if save_cdn_urls and response.status_code == 200: new_urls = extract_cdn_urls_from_response(response) batch_cdn_urls.extend(new_urls) if batch_cdn_urls: From 0d6ed9d9c56aae82f78d92a8194f2a538709d151 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Tue, 20 Jan 2026 03:07:13 -0800 Subject: [PATCH 8/8] fix(message_builder): add trailing newline to metadata section The build_metadata_section function used "\n".join(metadata_lines) which lacks a trailing newline. When dimensions were appended later, they concatenated directly without proper spacing, producing malformed output like "**Format:** PNG**Original Dimensions:** 1024x1024". Added trailing newline so subsequent content appears on a new line. Fixes issue identified in PR #35 review. Co-Authored-By: Claude Opus 4.5 --- shared/discord/message_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shared/discord/message_builder.py b/shared/discord/message_builder.py index b7cc106..8397f32 100644 --- a/shared/discord/message_builder.py +++ b/shared/discord/message_builder.py @@ -55,7 +55,7 @@ def build_metadata_section( return "" section = f"\n\n**{section_title}:**\n" - section += "\n".join(metadata_lines) + section += "\n".join(metadata_lines) + "\n" return section