nodes/AEmotionStudio-ComfyUI-DiscordSend
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #35 from AEmotionStudio/refactor/separation-of-concerns

Browse Source 
refactor(phase2): extract shared utilities and reduce node duplication
This commit is contained in:
Æmotion Studio
2026-01-20 03:20:29 -08:00
committed by GitHub
parent 9295546c17 0d6ed9d9c5
commit 569f76f484
19 changed files with 1878 additions and 840 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/__init__.py
+2 -1
View File
@@ -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']
nodes/base_node.py
+343
View File
@@ -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
nodes/image_node.py
+86 -227
View File
@@ -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
@@ -495,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:
@@ -712,58 +671,18 @@ 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
# Extract CDN URLs and send file using shared utility
# Only extract when status is 200 (has content), not 204 (no content)
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}")
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 +711,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:
@@ -861,61 +760,21 @@ 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
# Try to extract CDN URLs from batch response
# Extract CDN URLs and send file using shared utility
# Only extract when status is 200 (has content), not 204 (no content)
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}")
else:
print(f"Error sending batch to Discord: Status code {response.status_code} - {response.text}")
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
except Exception as e:
print(f"Error sending batch to Discord: {e}")
nodes/video_node.py
+84 -559
View File
@@ -24,7 +24,25 @@ 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
)
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):
"""
@@ -47,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
@@ -63,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.
@@ -361,57 +324,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(
@@ -874,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 = []
@@ -982,364 +856,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 +981,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 +1002,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:
shared/__init__.py
+26
View File
@@ -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',
]
shared/discord/__init__.py
+27
View File
@@ -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',
]
shared/discord/cdn_extractor.py
+168
View File
@@ -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
shared/discord/message_builder.py
+210
View File
@@ -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) + "\n"
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"
shared/filename_utils.py
+83
View File
@@ -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 ""
shared/logging_config.py
+16
View File
@@ -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.
shared/media/__init__.py
+29
View File
@@ -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',
]
shared/media/format_utils.py
+137
View File
@@ -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
shared/media/video_encoder.py
+510
View File
@@ -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
shared/path_utils.py
+83
View File
@@ -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)
shared/workflow/prompt_extractor.py
+57 -43
View File
@@ -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:
tests/test_discord_mock.py
+5
View File
@@ -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__))))
tests/test_image_node_sanitization.py
+3 -1
View File
@@ -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):
tests/test_media.py
+1 -1
View File
@@ -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."""
tests/test_security_fixes.py
+8 -8
View File
@@ -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()