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()