refactor(phase2): extract shared utilities and reduce node duplication
- Create shared/filename_utils.py for date/time/dimension filename building - Create shared/path_utils.py for output directory handling - Create shared/discord/message_builder.py for Discord message construction - Create shared/discord/cdn_extractor.py for CDN URL extraction - Refactor image_node.py to use shared utilities (-161 lines, 16.3%) - Refactor video_node.py to use shared utilities (-361 lines, 23.1%) - Fix setup_logging missing from logging_config.py - Fix test imports to use new module paths (nodes.* instead of discord_*_node) - Total reduction: 522 lines (20.5%), exceeding PRD target of ~200 lines Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
AEmotionStudio
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
852840dcae
commit
a208cd482b
+68
-220
@@ -22,7 +22,13 @@ from shared import (
|
||||
update_github_cdn_urls,
|
||||
extract_prompts_from_workflow,
|
||||
send_to_discord_with_retry,
|
||||
tensor_to_numpy_uint8
|
||||
tensor_to_numpy_uint8,
|
||||
build_filename_with_metadata,
|
||||
get_output_directory,
|
||||
build_metadata_section,
|
||||
build_prompt_section,
|
||||
extract_cdn_urls_from_response,
|
||||
send_cdn_urls_file
|
||||
)
|
||||
|
||||
|
||||
@@ -240,48 +246,24 @@ class DiscordSendSaveImage:
|
||||
if send_workflow_json and extra_pnginfo is not None and "workflow" in extra_pnginfo:
|
||||
extra_pnginfo["workflow"] = sanitize_json_for_export(extra_pnginfo["workflow"])
|
||||
|
||||
# Add date and/or time if enabled
|
||||
date_time_parts = []
|
||||
|
||||
# Prepare info for Discord message
|
||||
# Build filename with date/time metadata using shared utility
|
||||
image_info = {}
|
||||
|
||||
if add_date:
|
||||
# Get ONLY the date in YYYY-MM-DD format
|
||||
current_date = time.strftime("%Y-%m-%d")
|
||||
date_time_parts.append(current_date)
|
||||
print(f"Adding date to filename: {current_date}")
|
||||
image_info["date"] = current_date
|
||||
|
||||
if add_time:
|
||||
# Get ONLY the time in HH-MM-SS format
|
||||
current_time = time.strftime("%H-%M-%S")
|
||||
date_time_parts.append(current_time)
|
||||
print(f"Adding time to filename: {current_time}")
|
||||
image_info["time"] = current_time
|
||||
|
||||
# Add date/time components to filename prefix if any were enabled
|
||||
if date_time_parts:
|
||||
date_time_suffix = "_" + "_".join(date_time_parts)
|
||||
filename_prefix += date_time_suffix
|
||||
print(f"Final timestamp suffix: {date_time_suffix}")
|
||||
|
||||
filename_prefix, image_info = build_filename_with_metadata(
|
||||
prefix=filename_prefix,
|
||||
add_date=add_date,
|
||||
add_time=add_time,
|
||||
info_dict=image_info
|
||||
)
|
||||
|
||||
# Add prefix append
|
||||
filename_prefix += self.prefix_append
|
||||
|
||||
# Get ComfyUI output directory for safe path handling
|
||||
comfy_output_dir = folder_paths.get_output_directory()
|
||||
|
||||
# Choose destination directory based on save_output flag
|
||||
if save_output:
|
||||
# Create a output subfolder in the ComfyUI output directory
|
||||
dest_folder = os.path.join(comfy_output_dir, "discord_output")
|
||||
os.makedirs(dest_folder, exist_ok=True)
|
||||
else:
|
||||
# Use ComfyUI's temporary directory for preview-only files
|
||||
dest_folder = folder_paths.get_temp_directory()
|
||||
os.makedirs(dest_folder, exist_ok=True)
|
||||
print(f"Using temporary directory for preview: {dest_folder}")
|
||||
# Get output directory using shared utility
|
||||
dest_folder = get_output_directory(
|
||||
save_output=save_output,
|
||||
comfy_output_dir=folder_paths.get_output_directory(),
|
||||
temp_dir=folder_paths.get_temp_directory()
|
||||
)
|
||||
|
||||
# Setup paths using ComfyUI's path validation
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
||||
@@ -358,71 +340,39 @@ class DiscordSendSaveImage:
|
||||
elif send_to_discord and not webhook_url:
|
||||
print("Discord integration was enabled but no webhook URL was provided")
|
||||
|
||||
# Add image info to Discord message if relevant options are enabled
|
||||
# Build image info message using shared utility
|
||||
if send_to_discord and webhook_url and (add_date or add_time or add_dimensions or resize_to_power_of_2 or include_format_in_message):
|
||||
info_message = "\n\n**Image Information:**\n"
|
||||
|
||||
if "date" in image_info:
|
||||
info_message += f"**Date:** {image_info['date']}\n"
|
||||
|
||||
if "time" in image_info:
|
||||
info_message += f"**Time:** {image_info['time']}\n"
|
||||
|
||||
# Add format to the message if the option is enabled
|
||||
if include_format_in_message:
|
||||
info_message += f"**Format:** {file_format.upper()}\n"
|
||||
|
||||
# Update the message prefix with the information
|
||||
info_message = build_metadata_section(
|
||||
info_dict=image_info,
|
||||
include_date=add_date,
|
||||
include_time=add_time,
|
||||
include_dimensions=False, # Dimensions added later after processing
|
||||
include_format=include_format_in_message,
|
||||
file_format=file_format,
|
||||
section_title="Image Information"
|
||||
)
|
||||
image_info["message_prefix"] = info_message
|
||||
|
||||
# Note: We don't add to discord_message yet, as dimensions aren't known until processing
|
||||
print("Prepared image information section for Discord message")
|
||||
|
||||
# Extract prompts if requested
|
||||
|
||||
# Extract and build prompts section using shared utilities
|
||||
if send_to_discord and include_prompts_in_message:
|
||||
workflow_data = None
|
||||
|
||||
|
||||
# First try to get workflow from extra_pnginfo
|
||||
if original_extra_pnginfo is not None and isinstance(original_extra_pnginfo, dict) and "workflow" in original_extra_pnginfo:
|
||||
workflow_data = original_extra_pnginfo["workflow"]
|
||||
|
||||
|
||||
# If no workflow in extra_pnginfo, check if prompt is actually a workflow
|
||||
if workflow_data is None and original_prompt is not None:
|
||||
# Check if prompt is already a workflow
|
||||
if isinstance(original_prompt, dict) and "nodes" in original_prompt:
|
||||
workflow_data = original_prompt
|
||||
|
||||
# Extract prompts from workflow data
|
||||
|
||||
# Extract and build prompts section
|
||||
if workflow_data is not None:
|
||||
positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data)
|
||||
|
||||
# Ensure the prompts are strings or None
|
||||
if positive_prompt is not False and positive_prompt is not None and not isinstance(positive_prompt, str):
|
||||
positive_prompt = str(positive_prompt)
|
||||
print(f"Converted positive prompt to string: {positive_prompt[:50]}...")
|
||||
|
||||
if negative_prompt is not False and negative_prompt is not None and not isinstance(negative_prompt, str):
|
||||
negative_prompt = str(negative_prompt)
|
||||
print(f"Converted negative prompt to string: {negative_prompt[:50]}...")
|
||||
|
||||
# Check if we have valid prompt data
|
||||
has_valid_prompt = (
|
||||
(isinstance(positive_prompt, str) and positive_prompt) or
|
||||
(isinstance(negative_prompt, str) and negative_prompt)
|
||||
)
|
||||
|
||||
# Add prompts to Discord message if found
|
||||
if has_valid_prompt:
|
||||
prompt_message = "\n\n**Generation Prompts:**\n"
|
||||
|
||||
if isinstance(positive_prompt, str) and positive_prompt:
|
||||
prompt_message += f"**Positive:**\n```\n{positive_prompt}\n```\n"
|
||||
|
||||
if isinstance(negative_prompt, str) and negative_prompt:
|
||||
prompt_message += f"**Negative:**\n```\n{negative_prompt}\n```\n"
|
||||
|
||||
# Store prompt message for adding after image info
|
||||
image_info["prompt_message"] = prompt_message
|
||||
prompt_section = build_prompt_section(positive_prompt, negative_prompt)
|
||||
if prompt_section:
|
||||
image_info["prompt_message"] = prompt_section
|
||||
print("Prepared prompts for Discord message")
|
||||
|
||||
# Optimization: Create metadata once for the entire batch
|
||||
@@ -712,58 +662,17 @@ class DiscordSendSaveImage:
|
||||
if send_workflow_json and "workflow" in files:
|
||||
print(f"Successfully sent workflow JSON for image {batch_number+1}")
|
||||
|
||||
# Try to extract CDN URLs from batch response
|
||||
if save_cdn_urls and response.status_code == 200:
|
||||
try:
|
||||
response_data = response.json()
|
||||
print(f"Received JSON response from Discord with {len(response_data) if isinstance(response_data, dict) else 'invalid'} fields")
|
||||
|
||||
if "attachments" in response_data and isinstance(response_data["attachments"], list):
|
||||
print(f"Found {len(response_data['attachments'])} attachments in Discord response")
|
||||
|
||||
for idx, attachment in enumerate(response_data["attachments"]):
|
||||
if "url" in attachment and "filename" in attachment:
|
||||
# Filter out workflow JSON files from URLs list
|
||||
if not attachment["filename"].endswith(".json"):
|
||||
batch_cdn_urls.append((attachment["filename"], attachment["url"]))
|
||||
print(f"Extracted CDN URL for batch image {idx+1}: {attachment['url']}")
|
||||
else:
|
||||
print(f"Skipping JSON file: {attachment['filename']}")
|
||||
else:
|
||||
print(f"Attachment {idx+1} missing URL or filename: {attachment.keys()}")
|
||||
|
||||
print(f"Total batch CDN URLs collected: {len(batch_cdn_urls)}")
|
||||
|
||||
# Create and send a text file with the CDN URLs if we have any
|
||||
if batch_cdn_urls:
|
||||
try:
|
||||
# Create the text file content
|
||||
url_text_content = "# Discord CDN URLs\n\n"
|
||||
for idx, (filename, url) in enumerate(batch_cdn_urls):
|
||||
url_text_content += f"{idx+1}. {filename}: {url}\n"
|
||||
|
||||
# Create a unique filename for the text file
|
||||
urls_filename = f"cdn_urls-{uuid4()}.txt"
|
||||
|
||||
# Prepare the request with just the URL file
|
||||
url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))}
|
||||
url_data = {"content": "Discord CDN URLs for the uploaded images:"}
|
||||
|
||||
# Send a follow-up message with just the URLs text file
|
||||
url_response = send_to_discord_with_retry(
|
||||
webhook_url,
|
||||
files=url_files,
|
||||
data=url_data
|
||||
)
|
||||
|
||||
if url_response.status_code in [200, 204]:
|
||||
print(f"Successfully sent CDN URLs text file to Discord")
|
||||
else:
|
||||
print(f"Error sending CDN URLs text file: Status code {url_response.status_code}")
|
||||
except Exception as e:
|
||||
print(f"Error creating or sending CDN URLs text file: {e}")
|
||||
except Exception as e:
|
||||
print(f"Error extracting CDN URLs from batch response: {e}")
|
||||
# Extract CDN URLs and send file using shared utility
|
||||
if save_cdn_urls:
|
||||
new_urls = extract_cdn_urls_from_response(response)
|
||||
batch_cdn_urls.extend(new_urls)
|
||||
if batch_cdn_urls:
|
||||
send_cdn_urls_file(
|
||||
webhook_url=webhook_url,
|
||||
urls=batch_cdn_urls,
|
||||
send_func=send_to_discord_with_retry,
|
||||
message="Discord CDN URLs for the uploaded images:"
|
||||
)
|
||||
else:
|
||||
print(f"Error: Discord returned status code {response.status_code}")
|
||||
discord_send_success = False
|
||||
@@ -792,32 +701,12 @@ class DiscordSendSaveImage:
|
||||
|
||||
# If we have CDN URLs and we're not in batch mode, send them as a text file
|
||||
if save_cdn_urls and discord_cdn_urls and not (group_batched_images and len(images) > 1):
|
||||
try:
|
||||
# Create the text file content
|
||||
url_text_content = "# Discord CDN URLs\n\n"
|
||||
for idx, (filename, url) in enumerate(discord_cdn_urls):
|
||||
url_text_content += f"{idx+1}. {filename}: {url}\n"
|
||||
|
||||
# Create a unique filename for the text file
|
||||
urls_filename = f"cdn_urls-{uuid4()}.txt"
|
||||
|
||||
# Prepare the request with just the URL file
|
||||
url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))}
|
||||
url_data = {"content": "Discord CDN URLs for the uploaded images:"}
|
||||
|
||||
# Send a follow-up message with just the URLs text file
|
||||
url_response = send_to_discord_with_retry(
|
||||
webhook_url,
|
||||
files=url_files,
|
||||
data=url_data
|
||||
)
|
||||
|
||||
if url_response.status_code in [200, 204]:
|
||||
print(f"Successfully sent CDN URLs text file to Discord")
|
||||
else:
|
||||
print(f"Error sending CDN URLs text file: Status code {url_response.status_code}")
|
||||
except Exception as e:
|
||||
print(f"Error creating or sending CDN URLs text file: {e}")
|
||||
send_cdn_urls_file(
|
||||
webhook_url=webhook_url,
|
||||
urls=discord_cdn_urls,
|
||||
send_func=send_to_discord_with_retry,
|
||||
message="Discord CDN URLs for the uploaded images:"
|
||||
)
|
||||
elif send_to_discord and not discord_send_success:
|
||||
print("DiscordSendSaveImage: There were errors sending some images to Discord")
|
||||
else:
|
||||
@@ -862,58 +751,17 @@ class DiscordSendSaveImage:
|
||||
discord_send_success = True
|
||||
discord_sent_files = ["batch_gallery"] # Mark as successfully sent
|
||||
|
||||
# Try to extract CDN URLs from batch response
|
||||
if save_cdn_urls and response.status_code == 200:
|
||||
try:
|
||||
response_data = response.json()
|
||||
print(f"Received JSON response from Discord with {len(response_data) if isinstance(response_data, dict) else 'invalid'} fields")
|
||||
|
||||
if "attachments" in response_data and isinstance(response_data["attachments"], list):
|
||||
print(f"Found {len(response_data['attachments'])} attachments in Discord response")
|
||||
|
||||
for idx, attachment in enumerate(response_data["attachments"]):
|
||||
if "url" in attachment and "filename" in attachment:
|
||||
# Filter out workflow JSON files from URLs list
|
||||
if not attachment["filename"].endswith(".json"):
|
||||
batch_cdn_urls.append((attachment["filename"], attachment["url"]))
|
||||
print(f"Extracted CDN URL for batch image {idx+1}: {attachment['url']}")
|
||||
else:
|
||||
print(f"Skipping JSON file: {attachment['filename']}")
|
||||
else:
|
||||
print(f"Attachment {idx+1} missing URL or filename: {attachment.keys()}")
|
||||
|
||||
print(f"Total batch CDN URLs collected: {len(batch_cdn_urls)}")
|
||||
|
||||
# Create and send a text file with the CDN URLs if we have any
|
||||
if batch_cdn_urls:
|
||||
try:
|
||||
# Create the text file content
|
||||
url_text_content = "# Discord CDN URLs\n\n"
|
||||
for idx, (filename, url) in enumerate(batch_cdn_urls):
|
||||
url_text_content += f"{idx+1}. {filename}: {url}\n"
|
||||
|
||||
# Create a unique filename for the text file
|
||||
urls_filename = f"cdn_urls-{uuid4()}.txt"
|
||||
|
||||
# Prepare the request with just the URL file
|
||||
url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))}
|
||||
url_data = {"content": "Discord CDN URLs for the uploaded images:"}
|
||||
|
||||
# Send a follow-up message with just the URLs text file
|
||||
url_response = send_to_discord_with_retry(
|
||||
webhook_url,
|
||||
files=url_files,
|
||||
data=url_data
|
||||
)
|
||||
|
||||
if url_response.status_code in [200, 204]:
|
||||
print(f"Successfully sent CDN URLs text file to Discord")
|
||||
else:
|
||||
print(f"Error sending CDN URLs text file: Status code {url_response.status_code}")
|
||||
except Exception as e:
|
||||
print(f"Error creating or sending CDN URLs text file: {e}")
|
||||
except Exception as e:
|
||||
print(f"Error extracting CDN URLs from batch response: {e}")
|
||||
# Extract CDN URLs and send file using shared utility
|
||||
if save_cdn_urls:
|
||||
new_urls = extract_cdn_urls_from_response(response)
|
||||
batch_cdn_urls.extend(new_urls)
|
||||
if batch_cdn_urls:
|
||||
send_cdn_urls_file(
|
||||
webhook_url=webhook_url,
|
||||
urls=batch_cdn_urls,
|
||||
send_func=send_to_discord_with_retry,
|
||||
message="Discord CDN URLs for the uploaded images:"
|
||||
)
|
||||
else:
|
||||
print(f"Error sending batch to Discord: Status code {response.status_code} - {response.text}")
|
||||
discord_send_success = False
|
||||
|
||||
+68
-434
@@ -24,7 +24,19 @@ import functools
|
||||
import server
|
||||
|
||||
# Import shared utilities
|
||||
from shared import sanitize_json_for_export, update_github_cdn_urls, send_to_discord_with_retry, tensor_to_numpy_uint8
|
||||
from shared import (
|
||||
sanitize_json_for_export,
|
||||
update_github_cdn_urls,
|
||||
send_to_discord_with_retry,
|
||||
tensor_to_numpy_uint8,
|
||||
build_filename_with_metadata,
|
||||
get_output_directory,
|
||||
build_metadata_section,
|
||||
build_prompt_section,
|
||||
extract_cdn_urls_from_response,
|
||||
send_cdn_urls_file,
|
||||
extract_prompts_from_workflow
|
||||
)
|
||||
# Define cached decorator for local use
|
||||
def cached(max_size=None):
|
||||
"""
|
||||
@@ -361,57 +373,28 @@ class DiscordSendSaveVideo:
|
||||
# Get first image for metadata
|
||||
first_image = images[0]
|
||||
|
||||
# Add date and/or time if enabled
|
||||
date_time_parts = []
|
||||
|
||||
# Video info for Discord message
|
||||
# Build filename with date/time/dimensions metadata using shared utility
|
||||
height, width = images[0].shape[0], images[0].shape[1]
|
||||
video_info = {}
|
||||
|
||||
if add_date:
|
||||
# Get ONLY the date in YYYY-MM-DD format
|
||||
current_date = time.strftime("%Y-%m-%d")
|
||||
date_time_parts.append(current_date)
|
||||
print(f"Adding date to filename: {current_date}")
|
||||
video_info["date"] = current_date
|
||||
|
||||
if add_time:
|
||||
# Get ONLY the time in HH-MM-SS format
|
||||
current_time = time.strftime("%H-%M-%S")
|
||||
date_time_parts.append(current_time)
|
||||
print(f"Adding time to filename: {current_time}")
|
||||
video_info["time"] = current_time
|
||||
|
||||
# Add dimensions if enabled
|
||||
if add_dimensions:
|
||||
# Get dimensions from first frame
|
||||
height, width = images[0].shape[0], images[0].shape[1]
|
||||
dim_text = f"{width}x{height}"
|
||||
date_time_parts.append(dim_text)
|
||||
print(f"Adding dimensions to filename: {dim_text}")
|
||||
video_info["dimensions"] = dim_text
|
||||
|
||||
# Add date/time/dimensions components to filename prefix if any were enabled
|
||||
if date_time_parts:
|
||||
date_time_suffix = "_" + "_".join(date_time_parts)
|
||||
filename_prefix += date_time_suffix
|
||||
print(f"Final metadata suffix: {date_time_suffix}")
|
||||
|
||||
filename_prefix, video_info = build_filename_with_metadata(
|
||||
prefix=filename_prefix,
|
||||
add_date=add_date,
|
||||
add_time=add_time,
|
||||
add_dimensions=add_dimensions,
|
||||
width=width,
|
||||
height=height,
|
||||
info_dict=video_info
|
||||
)
|
||||
|
||||
# Add prefix append
|
||||
filename_prefix += self.prefix_append
|
||||
|
||||
# Get ComfyUI output directory for safe path handling
|
||||
comfy_output_dir = folder_paths.get_output_directory()
|
||||
|
||||
# Choose destination directory based on save_output flag
|
||||
if save_output:
|
||||
# Create a video output subfolder in the ComfyUI output directory
|
||||
dest_folder = os.path.join(comfy_output_dir, "discord_output")
|
||||
os.makedirs(dest_folder, exist_ok=True)
|
||||
else:
|
||||
# Use ComfyUI's temporary directory for preview-only files
|
||||
dest_folder = folder_paths.get_temp_directory()
|
||||
os.makedirs(dest_folder, exist_ok=True)
|
||||
print(f"Using temporary directory for preview: {dest_folder}")
|
||||
|
||||
# Get output directory using shared utility
|
||||
dest_folder = get_output_directory(
|
||||
save_output=save_output,
|
||||
comfy_output_dir=folder_paths.get_output_directory(),
|
||||
temp_dir=folder_paths.get_temp_directory()
|
||||
)
|
||||
|
||||
# Setup paths using ComfyUI's path validation
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
||||
@@ -982,364 +965,45 @@ class DiscordSendSaveVideo:
|
||||
|
||||
# Prepare message content
|
||||
message_content = discord_message
|
||||
|
||||
# Add video metadata information to Discord message when enabled
|
||||
video_metadata_text = ""
|
||||
|
||||
# Only build video metadata text if the option is enabled
|
||||
|
||||
# Add video metadata to Discord message using shared utility
|
||||
if include_video_info:
|
||||
# Add date if enabled
|
||||
if add_date and "date" in video_info:
|
||||
video_metadata_text += f"**Date**: {video_info['date']}\n"
|
||||
|
||||
# Add time if enabled
|
||||
if add_time and "time" in video_info:
|
||||
video_metadata_text += f"**Time**: {video_info['time']}\n"
|
||||
|
||||
# Add dimensions if enabled
|
||||
if add_dimensions and "dimensions" in video_info:
|
||||
video_metadata_text += f"**Dimensions**: {video_info['dimensions']}\n"
|
||||
|
||||
# Add frame rate information
|
||||
video_metadata_text += f"**Frame Rate**: {frame_rate} fps\n"
|
||||
|
||||
# Add format information
|
||||
video_metadata_text += f"**Format**: {format}\n"
|
||||
|
||||
# Add video metadata to message if any was collected
|
||||
if video_metadata_text and message_content:
|
||||
message_content += "\n\n**Video Info**:\n" + video_metadata_text
|
||||
print(f"Added video metadata to Discord message: {len(video_metadata_text)} chars")
|
||||
elif video_metadata_text:
|
||||
message_content = "**Video Info**:\n" + video_metadata_text
|
||||
print(f"Added video metadata to Discord message: {len(video_metadata_text)} chars")
|
||||
metadata_section = build_metadata_section(
|
||||
info_dict=video_info,
|
||||
include_date=add_date,
|
||||
include_time=add_time,
|
||||
include_dimensions=add_dimensions,
|
||||
include_format=True,
|
||||
file_format=format,
|
||||
frame_rate=frame_rate,
|
||||
section_title="Video Info"
|
||||
)
|
||||
if metadata_section:
|
||||
message_content += metadata_section
|
||||
print(f"Added video metadata to Discord message: {len(metadata_section)} chars")
|
||||
else:
|
||||
print("Video info was not included in Discord message (disabled by user)")
|
||||
|
||||
# Include the generation prompts if requested
|
||||
# Include the generation prompts if requested using shared utilities
|
||||
if include_prompts_in_message:
|
||||
print(f"Prompt inclusion requested! Starting prompt extraction process...")
|
||||
print(f"original_prompt type: {type(original_prompt) if original_prompt is not None else 'None'}")
|
||||
print(f"original_extra_pnginfo type: {type(original_extra_pnginfo) if original_extra_pnginfo is not None else 'None'}")
|
||||
|
||||
if original_prompt is not None:
|
||||
print(f"original_prompt keys: {original_prompt.keys() if isinstance(original_prompt, dict) else 'Not a dict'}")
|
||||
|
||||
if original_extra_pnginfo is not None:
|
||||
print(f"original_extra_pnginfo keys: {original_extra_pnginfo.keys() if isinstance(original_extra_pnginfo, dict) else 'Not a dict'}")
|
||||
if isinstance(original_extra_pnginfo, dict) and "workflow" in original_extra_pnginfo:
|
||||
print(f"workflow keys: {original_extra_pnginfo['workflow'].keys() if isinstance(original_extra_pnginfo['workflow'], dict) else 'Workflow not a dict'}")
|
||||
|
||||
# First identify the workflow data to extract prompts from
|
||||
workflow_data = None
|
||||
|
||||
|
||||
# First try to get workflow from extra_pnginfo
|
||||
if original_extra_pnginfo is not None and isinstance(original_extra_pnginfo, dict) and "workflow" in original_extra_pnginfo:
|
||||
workflow_data = original_extra_pnginfo["workflow"]
|
||||
print("Found workflow data in extra_pnginfo")
|
||||
|
||||
|
||||
# If no workflow in extra_pnginfo, check if prompt is actually a workflow
|
||||
if workflow_data is None and original_prompt is not None:
|
||||
# Check if prompt is already a workflow
|
||||
if isinstance(original_prompt, dict) and "nodes" in original_prompt:
|
||||
workflow_data = original_prompt
|
||||
print("Found workflow data in original_prompt")
|
||||
|
||||
# Extract prompts if we found workflow data
|
||||
|
||||
# Extract and build prompts section using shared utilities
|
||||
if workflow_data is not None:
|
||||
print("Workflow data found, attempting to extract prompts...")
|
||||
|
||||
# Extract prompts and add to message
|
||||
try:
|
||||
# Try to use discord_image_node's function first
|
||||
try:
|
||||
from discord_image_node import extract_prompts_from_workflow
|
||||
positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data)
|
||||
|
||||
print(f"extract_prompts_from_workflow returned: positive_type={type(positive_prompt)}, negative_type={type(negative_prompt)}")
|
||||
|
||||
# Ensure the prompts are strings or None
|
||||
if positive_prompt is not False and positive_prompt is not None and not isinstance(positive_prompt, str):
|
||||
positive_prompt = str(positive_prompt)
|
||||
print(f"Converted positive prompt to string: {positive_prompt[:50]}...")
|
||||
|
||||
if negative_prompt is not False and negative_prompt is not None and not isinstance(negative_prompt, str):
|
||||
negative_prompt = str(negative_prompt)
|
||||
print(f"Converted negative prompt to string: {negative_prompt[:50]}...")
|
||||
|
||||
# Check if we have valid prompt data
|
||||
has_valid_prompt = (
|
||||
(isinstance(positive_prompt, str) and positive_prompt) or
|
||||
(isinstance(negative_prompt, str) and negative_prompt)
|
||||
)
|
||||
|
||||
if has_valid_prompt:
|
||||
print(f"Successfully extracted prompts via discord_image_node")
|
||||
|
||||
# Add to message
|
||||
prompt_text = "\n\n**Generation Prompts:**\n"
|
||||
|
||||
if isinstance(positive_prompt, str) and positive_prompt:
|
||||
prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n"
|
||||
|
||||
if isinstance(negative_prompt, str) and negative_prompt:
|
||||
prompt_text += f"**Negative:**\n```\n{negative_prompt}\n```\n"
|
||||
|
||||
# Add to message content
|
||||
message_content += prompt_text
|
||||
print("Successfully added prompts to Discord message")
|
||||
else:
|
||||
print("No valid prompts extracted from discord_image_node function")
|
||||
# Raise an exception to trigger the fallback extraction method
|
||||
raise ValueError("No valid prompts extracted")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Could not use discord_image_node extraction: {str(e)}")
|
||||
print("Falling back to built-in prompt extraction")
|
||||
|
||||
# Local extraction logic
|
||||
positive_prompt = None
|
||||
negative_prompt = None
|
||||
|
||||
# Find CLIP text nodes
|
||||
clip_nodes = []
|
||||
|
||||
# Define node types that can contain prompts
|
||||
prompt_node_types = [
|
||||
"CLIPTextEncode", # Standard SD 1.5 prompt node
|
||||
"SDXLPromptEncoder", # SDXL prompt encoder
|
||||
"SDXLTextEncode", # Another SDXL text node
|
||||
]
|
||||
|
||||
if "nodes" in workflow_data:
|
||||
nodes = workflow_data["nodes"]
|
||||
node_list = []
|
||||
|
||||
# Debug print full node data for the first few nodes
|
||||
print("\nExamining raw node data:")
|
||||
if isinstance(nodes, dict):
|
||||
for i, (node_id, node) in enumerate(list(nodes.items())[:3]):
|
||||
print(f"Node {i} (ID: {node_id}): {node.get('type', 'unknown')} - Active status: bypassed={node.get('bypassed', 'Not set')}, muted={node.get('muted', 'Not set')}, disabled={node.get('disabled', 'Not set')}")
|
||||
elif isinstance(nodes, list) and len(nodes) > 0:
|
||||
for i, node in enumerate(nodes[:3]):
|
||||
print(f"Node {i}: {node.get('type', 'unknown')} - Active status: bypassed={node.get('bypassed', 'Not set')}, muted={node.get('muted', 'Not set')}, disabled={node.get('disabled', 'Not set')}")
|
||||
|
||||
# Convert to list for processing
|
||||
if isinstance(nodes, dict):
|
||||
print("Nodes is a dictionary, converting to list...")
|
||||
for node_id, node in nodes.items():
|
||||
if isinstance(node, dict):
|
||||
# Store the ID in the node dict for reference
|
||||
node_with_id = node.copy() # Copy to avoid modifying original
|
||||
node_with_id["id"] = node_id
|
||||
node_list.append(node_with_id)
|
||||
elif isinstance(nodes, list):
|
||||
print("Nodes is already a list")
|
||||
node_list = nodes
|
||||
else:
|
||||
print(f"Unexpected type for nodes: {type(nodes)}")
|
||||
|
||||
print(f"Processed {len(node_list)} nodes from workflow data")
|
||||
|
||||
# Let's specifically search for all CLIP text nodes to analyze
|
||||
print("\nSearching for all possible prompt nodes (active or disabled):")
|
||||
all_prompt_nodes = []
|
||||
for node in node_list:
|
||||
if isinstance(node, dict) and "type" in node:
|
||||
node_type = node.get("type", "")
|
||||
if node_type in prompt_node_types or ("Text" in node_type and ("Encode" in node_type or "Prompt" in node_type)):
|
||||
# Found a potential prompt node
|
||||
title = node.get("title", "Untitled")
|
||||
is_active = not (node.get("bypassed", False) or node.get("muted", False) or
|
||||
node.get("disabled", False) or node.get("active", True) == False or
|
||||
node.get("enabled", True) == False)
|
||||
|
||||
# Print detailed info
|
||||
print(f"Prompt node found: {title} (Type: {node_type})")
|
||||
print(f" - Status: {'ACTIVE' if is_active else 'DISABLED'}")
|
||||
print(f" - Properties: bypassed={node.get('bypassed', 'Not set')}, muted={node.get('muted', 'Not set')}, disabled={node.get('disabled', 'Not set')}")
|
||||
if "widgets_values" in node and node["widgets_values"]:
|
||||
text_preview = node["widgets_values"][0][:50] + "..." if len(node["widgets_values"][0]) > 50 else node["widgets_values"][0]
|
||||
print(f" - Content preview: {text_preview}")
|
||||
|
||||
all_prompt_nodes.append(node)
|
||||
|
||||
print(f"Found {len(all_prompt_nodes)} total potential prompt nodes ({len([n for n in all_prompt_nodes if not (n.get('bypassed', False) or n.get('muted', False) or n.get('disabled', False) or n.get('active', True) == False or n.get('enabled', True) == False)])} active, {len([n for n in all_prompt_nodes if n.get('bypassed', False) or n.get('muted', False) or n.get('disabled', False) or n.get('active', True) == False or n.get('enabled', True) == False])} disabled)")
|
||||
|
||||
# Check for active prompt nodes (not muted/bypassed)
|
||||
for node in node_list:
|
||||
if isinstance(node, dict) and "type" in node:
|
||||
# Check for various ways a node might be disabled
|
||||
is_disabled = (
|
||||
node.get("bypassed", False) or
|
||||
node.get("muted", False) or
|
||||
node.get("disabled", False) or
|
||||
node.get("active", True) == False or
|
||||
node.get("enabled", True) == False
|
||||
)
|
||||
|
||||
# If no clear status property, let's check outputs to see if they're connected
|
||||
if not is_disabled and ("outputs" in node or "Output" in node):
|
||||
# No outputs might mean disconnected
|
||||
has_outputs = False
|
||||
if "outputs" in node and node["outputs"]:
|
||||
has_outputs = True
|
||||
elif "Output" in node and node["Output"]:
|
||||
has_outputs = True
|
||||
|
||||
# Check if we should include nodes with no connections (sometimes these are valid sources)
|
||||
# For prompt nodes, we'll consider them even if disconnected, but log it
|
||||
if not has_outputs and node["type"] in prompt_node_types:
|
||||
print(f"Note: Including prompt node that appears disconnected: {node.get('title', node['type'])}")
|
||||
|
||||
if is_disabled:
|
||||
print(f"Skipping disabled node: {node.get('title', node.get('type', 'unknown'))}")
|
||||
continue
|
||||
|
||||
# If we reach here, the node is active
|
||||
# Check standard prompt node types
|
||||
if node["type"] in prompt_node_types:
|
||||
if "widgets_values" in node and node["widgets_values"]:
|
||||
clip_nodes.append(node)
|
||||
print(f"Found active prompt node: {node.get('title', node['type'])}")
|
||||
# Also look for other common prompt nodes
|
||||
elif "Text" in node["type"] and ("Encode" in node["type"] or "Prompt" in node["type"]):
|
||||
print(f"Found potential prompt node of type: {node['type']}")
|
||||
if "widgets_values" in node and node["widgets_values"]:
|
||||
print(f"Values: {node['widgets_values']}")
|
||||
clip_nodes.append(node)
|
||||
|
||||
print(f"Found {len(clip_nodes)} active prompt nodes")
|
||||
|
||||
# Dump the first 5 nodes to debug
|
||||
print("Sampling first few nodes for debugging:")
|
||||
if isinstance(workflow_data["nodes"], dict):
|
||||
sample_nodes = list(workflow_data["nodes"].values())[:5]
|
||||
for i, node in enumerate(sample_nodes):
|
||||
if isinstance(node, dict):
|
||||
print(f"Node {i}: type={node.get('type', 'unknown')}, title={node.get('title', 'untitled')}")
|
||||
elif isinstance(workflow_data["nodes"], list):
|
||||
for i, node in enumerate(workflow_data["nodes"][:5]):
|
||||
if isinstance(node, dict):
|
||||
print(f"Node {i}: type={node.get('type', 'unknown')}, title={node.get('title', 'untitled')}")
|
||||
|
||||
# Process CLIP nodes
|
||||
if len(clip_nodes) == 2:
|
||||
# Identify which is positive/negative
|
||||
indicators = ["worst quality", "low quality", "bad quality", "nude", "nsfw"]
|
||||
scores = [0, 0]
|
||||
|
||||
for i, node in enumerate(clip_nodes):
|
||||
text = node["widgets_values"][0].lower()
|
||||
for indicator in indicators:
|
||||
if indicator in text:
|
||||
scores[i] += 1
|
||||
|
||||
# The one with more negative indicators is the negative prompt
|
||||
if scores[0] > scores[1]:
|
||||
negative_prompt = clip_nodes[0]["widgets_values"][0]
|
||||
positive_prompt = clip_nodes[1]["widgets_values"][0]
|
||||
else:
|
||||
negative_prompt = clip_nodes[1]["widgets_values"][0]
|
||||
positive_prompt = clip_nodes[0]["widgets_values"][0]
|
||||
|
||||
print(f"Identified prompts with scores: {scores}")
|
||||
|
||||
# Add to message
|
||||
prompt_text = "\n\n**Generation Prompts:**\n"
|
||||
prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n"
|
||||
prompt_text += f"**Negative:**\n```\n{negative_prompt}\n```\n"
|
||||
|
||||
# Add to message content
|
||||
message_content += prompt_text
|
||||
print("Added prompts to Discord message via built-in extraction")
|
||||
elif len(clip_nodes) == 1:
|
||||
# Just one prompt, assume it's positive
|
||||
positive_prompt = clip_nodes[0]["widgets_values"][0]
|
||||
|
||||
# Add to message
|
||||
prompt_text = "\n\n**Generation Prompts:**\n"
|
||||
prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n"
|
||||
|
||||
# Add to message content
|
||||
message_content += prompt_text
|
||||
print("Added single prompt to Discord message")
|
||||
else:
|
||||
# Handle multiple CLIP nodes (more than 2)
|
||||
print(f"Found {len(clip_nodes)} CLIP nodes, will analyze each to determine positive/negative")
|
||||
|
||||
# Collect potential positive and negative prompts
|
||||
positive_candidates = []
|
||||
negative_candidates = []
|
||||
|
||||
# Common negative prompt indicators
|
||||
negative_indicators = ["worst quality", "low quality", "bad quality", "nude", "nsfw",
|
||||
"negative", "bad", "worse", "poor", "deformed"]
|
||||
|
||||
for node in clip_nodes:
|
||||
# Get the node's text content
|
||||
if "widgets_values" in node and node["widgets_values"]:
|
||||
text = node["widgets_values"][0]
|
||||
title = node.get("title", "").lower()
|
||||
node_type = node.get("type", "").lower()
|
||||
|
||||
# Check if title explicitly indicates a negative prompt
|
||||
is_negative = False
|
||||
if "negative" in title:
|
||||
is_negative = True
|
||||
print(f"Found explicit negative prompt node by title: {title}")
|
||||
else:
|
||||
# Check content for negative indicators
|
||||
text_lower = text.lower()
|
||||
negative_score = 0
|
||||
for indicator in negative_indicators:
|
||||
if indicator in text_lower:
|
||||
negative_score += 1
|
||||
|
||||
if negative_score >= 2: # Threshold for considering as negative
|
||||
is_negative = True
|
||||
print(f"Identified negative prompt by content with score {negative_score}")
|
||||
|
||||
# Categorize the prompt
|
||||
if is_negative:
|
||||
negative_candidates.append(text)
|
||||
else:
|
||||
positive_candidates.append(text)
|
||||
|
||||
# Get the final prompts
|
||||
positive_prompt = None
|
||||
negative_prompt = None
|
||||
|
||||
if positive_candidates:
|
||||
# Use the longest positive prompt as it likely has more information
|
||||
positive_prompt = max(positive_candidates, key=len)
|
||||
print(f"Selected positive prompt ({len(positive_prompt)} chars)")
|
||||
|
||||
if negative_candidates:
|
||||
# Use the longest negative prompt
|
||||
negative_prompt = max(negative_candidates, key=len)
|
||||
print(f"Selected negative prompt ({len(negative_prompt)} chars)")
|
||||
|
||||
# Add to message if any prompts were found
|
||||
if positive_prompt or negative_prompt:
|
||||
prompt_text = "\n\n**Generation Prompts:**\n"
|
||||
|
||||
if positive_prompt:
|
||||
prompt_text += f"**Positive:**\n```\n{positive_prompt}\n```\n"
|
||||
|
||||
if negative_prompt:
|
||||
prompt_text += f"**Negative:**\n```\n{negative_prompt}\n```\n"
|
||||
|
||||
# Add to message content
|
||||
message_content += prompt_text
|
||||
print(f"Added prompts to Discord message from {len(positive_candidates)} positive and {len(negative_candidates)} negative candidates")
|
||||
else:
|
||||
print("Could not identify any positive or negative prompts from the CLIP nodes")
|
||||
except Exception as e:
|
||||
print(f"Error extracting prompts: {str(e)}")
|
||||
print(f"Error type: {type(e).__name__}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data)
|
||||
prompt_section = build_prompt_section(positive_prompt, negative_prompt)
|
||||
if prompt_section:
|
||||
message_content += prompt_section
|
||||
print("Successfully added prompts to Discord message")
|
||||
else:
|
||||
print("No workflow data found for prompt extraction")
|
||||
|
||||
@@ -1426,20 +1090,10 @@ class DiscordSendSaveVideo:
|
||||
discord_sent_files.append(discord_filename) # Store the Discord UUID filename instead of local path
|
||||
print(f"Successfully sent video to Discord with filename: {discord_filename}")
|
||||
|
||||
# Try to extract CDN URL if available
|
||||
if save_cdn_urls and response.status_code == 200:
|
||||
try:
|
||||
response_data = response.json()
|
||||
# Discord webhook responses include attachments with URLs
|
||||
if "attachments" in response_data and isinstance(response_data["attachments"], list):
|
||||
for attachment in response_data["attachments"]:
|
||||
if "url" in attachment and "filename" in attachment:
|
||||
# Filter out workflow JSON files
|
||||
if not attachment["filename"].endswith(".json"):
|
||||
discord_cdn_urls.append((attachment["filename"], attachment["url"]))
|
||||
print(f"Extracted CDN URL for video: {attachment['url']}")
|
||||
except Exception as e:
|
||||
print(f"Error extracting CDN URL from response: {e}")
|
||||
# Extract CDN URLs using shared utility
|
||||
if save_cdn_urls:
|
||||
new_urls = extract_cdn_urls_from_response(response)
|
||||
discord_cdn_urls.extend(new_urls)
|
||||
else:
|
||||
print(f"Discord API error: {response.status_code} - {response.text}")
|
||||
discord_send_success = False
|
||||
@@ -1457,34 +1111,14 @@ class DiscordSendSaveVideo:
|
||||
except Exception as e:
|
||||
print(f"Error cleaning up temporary file: {e}")
|
||||
|
||||
# If we have CDN URLs and the option is enabled, send them as a text file
|
||||
# Send CDN URLs file using shared utility
|
||||
if save_cdn_urls and discord_cdn_urls:
|
||||
try:
|
||||
# Create the text file content
|
||||
url_text_content = "# Discord CDN URLs\n\n"
|
||||
for idx, (filename, url) in enumerate(discord_cdn_urls):
|
||||
url_text_content += f"{idx+1}. {filename}: {url}\n"
|
||||
|
||||
# Create a unique filename for the text file
|
||||
urls_filename = f"cdn_urls-{uuid4()}.txt"
|
||||
|
||||
# Prepare the request with just the URL file
|
||||
url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))}
|
||||
url_data = {"content": "Discord CDN URLs for the uploaded videos:"}
|
||||
|
||||
# Send a follow-up message with just the URLs text file
|
||||
url_response = send_to_discord_with_retry(
|
||||
webhook_url,
|
||||
files=url_files,
|
||||
data=url_data
|
||||
)
|
||||
|
||||
if url_response.status_code in [200, 204]:
|
||||
print(f"Successfully sent CDN URLs text file to Discord")
|
||||
else:
|
||||
print(f"Error sending CDN URLs text file: Status code {url_response.status_code}")
|
||||
except Exception as e:
|
||||
print(f"Error creating or sending CDN URLs text file: {e}")
|
||||
send_cdn_urls_file(
|
||||
webhook_url=webhook_url,
|
||||
urls=discord_cdn_urls,
|
||||
send_func=send_to_discord_with_retry,
|
||||
message="Discord CDN URLs for the uploaded videos:"
|
||||
)
|
||||
|
||||
# Update GitHub repository with CDN URLs if enabled
|
||||
if github_cdn_update and send_to_discord and discord_cdn_urls:
|
||||
|
||||
@@ -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',
|
||||
]
|
||||
|
||||
@@ -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',
|
||||
]
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
Discord message building utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides functions for constructing Discord messages with metadata,
|
||||
prompts, and other formatted content.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional, Tuple, List
|
||||
|
||||
|
||||
def build_metadata_section(
|
||||
info_dict: Dict[str, Any],
|
||||
include_date: bool = True,
|
||||
include_time: bool = True,
|
||||
include_dimensions: bool = True,
|
||||
include_format: bool = True,
|
||||
file_format: Optional[str] = None,
|
||||
frame_rate: Optional[float] = None,
|
||||
section_title: str = "Information"
|
||||
) -> str:
|
||||
"""
|
||||
Build a formatted metadata section for Discord messages.
|
||||
|
||||
Args:
|
||||
info_dict: Dictionary containing metadata (date, time, dimensions, etc.)
|
||||
include_date: Whether to include date if present
|
||||
include_time: Whether to include time if present
|
||||
include_dimensions: Whether to include dimensions if present
|
||||
include_format: Whether to include format information
|
||||
file_format: File format string (e.g., "png", "mp4")
|
||||
frame_rate: Frame rate for video (optional)
|
||||
section_title: Title for the section (e.g., "Image Information", "Video Info")
|
||||
|
||||
Returns:
|
||||
Formatted metadata string, or empty string if no metadata
|
||||
"""
|
||||
metadata_lines = []
|
||||
|
||||
if include_date and "date" in info_dict:
|
||||
metadata_lines.append(f"**Date:** {info_dict['date']}")
|
||||
|
||||
if include_time and "time" in info_dict:
|
||||
metadata_lines.append(f"**Time:** {info_dict['time']}")
|
||||
|
||||
if include_dimensions and "dimensions" in info_dict:
|
||||
metadata_lines.append(f"**Dimensions:** {info_dict['dimensions']}")
|
||||
|
||||
if frame_rate is not None:
|
||||
metadata_lines.append(f"**Frame Rate:** {frame_rate} fps")
|
||||
|
||||
if include_format and file_format:
|
||||
metadata_lines.append(f"**Format:** {file_format.upper()}")
|
||||
|
||||
if not metadata_lines:
|
||||
return ""
|
||||
|
||||
section = f"\n\n**{section_title}:**\n"
|
||||
section += "\n".join(metadata_lines)
|
||||
return section
|
||||
|
||||
|
||||
def build_prompt_section(
|
||||
positive_prompt: Optional[str],
|
||||
negative_prompt: Optional[str],
|
||||
section_title: str = "Generation Prompts"
|
||||
) -> str:
|
||||
"""
|
||||
Build a formatted prompts section for Discord messages.
|
||||
|
||||
Args:
|
||||
positive_prompt: The positive/main prompt text
|
||||
negative_prompt: The negative prompt text
|
||||
section_title: Title for the section
|
||||
|
||||
Returns:
|
||||
Formatted prompt string, or empty string if no prompts
|
||||
"""
|
||||
# Validate and normalize prompts
|
||||
if positive_prompt is not None and not isinstance(positive_prompt, str):
|
||||
positive_prompt = str(positive_prompt)
|
||||
if negative_prompt is not None and not isinstance(negative_prompt, str):
|
||||
negative_prompt = str(negative_prompt)
|
||||
|
||||
has_positive = isinstance(positive_prompt, str) and positive_prompt.strip()
|
||||
has_negative = isinstance(negative_prompt, str) and negative_prompt.strip()
|
||||
|
||||
if not has_positive and not has_negative:
|
||||
return ""
|
||||
|
||||
section = f"\n\n**{section_title}:**\n"
|
||||
|
||||
if has_positive:
|
||||
section += f"**Positive:**\n```\n{positive_prompt.strip()}\n```\n"
|
||||
|
||||
if has_negative:
|
||||
section += f"**Negative:**\n```\n{negative_prompt.strip()}\n```\n"
|
||||
|
||||
return section
|
||||
|
||||
|
||||
def build_discord_message(
|
||||
base_message: str = "",
|
||||
metadata_section: str = "",
|
||||
prompt_section: str = "",
|
||||
additional_sections: Optional[List[str]] = None,
|
||||
max_length: int = 2000
|
||||
) -> str:
|
||||
"""
|
||||
Build a complete Discord message from components.
|
||||
|
||||
Args:
|
||||
base_message: The main message content
|
||||
metadata_section: Pre-built metadata section
|
||||
prompt_section: Pre-built prompt section
|
||||
additional_sections: List of additional section strings
|
||||
max_length: Maximum message length (Discord limit is 2000)
|
||||
|
||||
Returns:
|
||||
Complete formatted message, truncated if necessary
|
||||
"""
|
||||
parts = [base_message] if base_message else []
|
||||
|
||||
if metadata_section:
|
||||
parts.append(metadata_section)
|
||||
|
||||
if prompt_section:
|
||||
parts.append(prompt_section)
|
||||
|
||||
if additional_sections:
|
||||
parts.extend(additional_sections)
|
||||
|
||||
message = "".join(parts)
|
||||
|
||||
# Truncate if necessary
|
||||
if len(message) > max_length:
|
||||
truncation_notice = "\n...[Message truncated]"
|
||||
message = message[:max_length - len(truncation_notice)] + truncation_notice
|
||||
|
||||
return message
|
||||
|
||||
|
||||
def validate_message_content(message: str) -> Tuple[bool, str]:
|
||||
"""
|
||||
Validate Discord message content.
|
||||
|
||||
Args:
|
||||
message: Message content to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, validation_message)
|
||||
"""
|
||||
if not message:
|
||||
return True, "Empty message (valid for file-only uploads)"
|
||||
|
||||
if len(message) > 2000:
|
||||
return False, f"Message exceeds 2000 character limit ({len(message)} chars)"
|
||||
|
||||
# Check for required sections (informational)
|
||||
has_prompts = "Generation Prompts" in message
|
||||
|
||||
info_parts = []
|
||||
info_parts.append(f"Message has {message.count(chr(10))} lines")
|
||||
|
||||
if not has_prompts:
|
||||
info_parts.append("WARNING: Message does NOT contain 'Generation Prompts' section")
|
||||
|
||||
return True, "\n".join(info_parts)
|
||||
|
||||
|
||||
def format_file_info(
|
||||
filename: str,
|
||||
file_size: int,
|
||||
mime_type: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Format file information for logging/display.
|
||||
|
||||
Args:
|
||||
filename: Name of the file
|
||||
file_size: Size in bytes
|
||||
mime_type: MIME type of the file
|
||||
|
||||
Returns:
|
||||
Formatted string with file information
|
||||
"""
|
||||
size_str = format_file_size(file_size)
|
||||
info = f"File: {filename} ({size_str})"
|
||||
if mime_type:
|
||||
info += f" [{mime_type}]"
|
||||
return info
|
||||
|
||||
|
||||
def format_file_size(size_bytes: int) -> str:
|
||||
"""
|
||||
Format file size in human-readable format.
|
||||
|
||||
Args:
|
||||
size_bytes: Size in bytes
|
||||
|
||||
Returns:
|
||||
Formatted string (e.g., "1.5 MB", "256 KB")
|
||||
"""
|
||||
if size_bytes < 1024:
|
||||
return f"{size_bytes} bytes"
|
||||
elif size_bytes < 1024 * 1024:
|
||||
return f"{size_bytes / 1024:.1f} KB"
|
||||
elif size_bytes < 1024 * 1024 * 1024:
|
||||
return f"{size_bytes / (1024 * 1024):.1f} MB"
|
||||
else:
|
||||
return f"{size_bytes / (1024 * 1024 * 1024):.2f} GB"
|
||||
@@ -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 ""
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
@@ -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__))))
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
+1
-1
@@ -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."""
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user