From 852840dcaed74d91c2f513eb43bac23cd289fb10 Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Mon, 19 Jan 2026 22:02:47 -0800 Subject: [PATCH] refactor(core): restructure project into shared/ and nodes/ modules --- .env.example | 48 +++++++++++++++++++ __init__.py | 6 +-- bot/__main__.py | 2 +- bot/cogs/generate.py | 2 +- config.yaml.example | 48 +++++++++++++++++++ discordsend_utils/__init__.py | 21 -------- nodes/__init__.py | 10 ++++ discord_image_node.py => nodes/image_node.py | 6 +-- discord_video_node.py => nodes/video_node.py | 10 ++-- shared/__init__.py | 40 ++++++++++++++++ shared/discord/__init__.py | 21 ++++++++ .../discord/webhook_client.py | 0 .../github_integration.py | 0 .../logging_config.py | 0 shared/media/__init__.py | 11 +++++ .../media}/image_processing.py | 0 shared/workflow/__init__.py | 15 ++++++ .../workflow}/prompt_extractor.py | 0 .../workflow}/sanitizer.py | 0 .../workflow}/workflow_builder.py | 0 tests/test_discord_mock.py | 2 +- tests/test_github_validation.py | 2 +- tests/test_utils.py | 14 +++--- tests/test_workflow.py | 3 +- 24 files changed, 217 insertions(+), 44 deletions(-) create mode 100644 .env.example create mode 100644 config.yaml.example delete mode 100644 discordsend_utils/__init__.py create mode 100644 nodes/__init__.py rename discord_image_node.py => nodes/image_node.py (99%) rename discord_video_node.py => nodes/video_node.py (99%) create mode 100644 shared/__init__.py create mode 100644 shared/discord/__init__.py rename discordsend_utils/discord_api.py => shared/discord/webhook_client.py (100%) rename {discordsend_utils => shared}/github_integration.py (100%) rename {discordsend_utils => shared}/logging_config.py (100%) create mode 100644 shared/media/__init__.py rename {discordsend_utils => shared/media}/image_processing.py (100%) create mode 100644 shared/workflow/__init__.py rename {discordsend_utils => shared/workflow}/prompt_extractor.py (100%) rename {discordsend_utils => shared/workflow}/sanitizer.py (100%) rename {discordsend_utils => shared/workflow}/workflow_builder.py (100%) diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..73bfb79 --- /dev/null +++ b/.env.example @@ -0,0 +1,48 @@ +# ComfyUI-DiscordSend Bot Environment Variables +# Copy this file to .env and fill in your values +# These override settings in config.yaml + +# ============================================================================= +# REQUIRED +# ============================================================================= + +# Discord bot token from Discord Developer Portal +DISCORDBOT_DISCORD_TOKEN=your_bot_token_here + +# ============================================================================= +# OPTIONAL - ComfyUI Connection +# ============================================================================= + +# ComfyUI server URL (default: http://127.0.0.1:8188) +# DISCORDBOT_COMFYUI_URL=http://127.0.0.1:8188 + +# ComfyUI WebSocket URL (default: auto-derived from COMFYUI_URL) +# DISCORDBOT_COMFYUI_WS_URL=ws://127.0.0.1:8188/ws + +# Request timeout in seconds (default: 30) +# DISCORDBOT_COMFYUI_TIMEOUT=30 + +# ============================================================================= +# OPTIONAL - Database +# ============================================================================= + +# Database URL (default: SQLite in bot/data/bot.db) +# For PostgreSQL: postgresql+asyncpg://user:pass@localhost/dbname +# DISCORDBOT_DATABASE_URL=sqlite+aiosqlite:///bot/data/bot.db + +# ============================================================================= +# OPTIONAL - Defaults +# ============================================================================= + +# Maximum pending jobs per user (default: 3) +# DISCORDBOT_MAX_QUEUE_PER_USER=3 + +# Path to default workflow JSON file +# DISCORDBOT_WORKFLOW_PATH=bot/data/default_workflow_api.json + +# ============================================================================= +# OPTIONAL - Discord Application +# ============================================================================= + +# Discord application ID (for slash command registration) +# DISCORDBOT_APPLICATION_ID=your_app_id_here diff --git a/__init__.py b/__init__.py index f5f32f4..0077ea1 100644 --- a/__init__.py +++ b/__init__.py @@ -10,9 +10,9 @@ current_dir = os.path.dirname(os.path.realpath(__file__)) if current_dir not in sys.path: sys.path.insert(0, current_dir) -# Import nodes -from discord_image_node import DiscordSendSaveImage -from discord_video_node import DiscordSendSaveVideo +# Import nodes from nodes package +from nodes.image_node import DiscordSendSaveImage +from nodes.video_node import DiscordSendSaveVideo # Node class mappings for ComfyUI NODE_CLASS_MAPPINGS = { diff --git a/bot/__main__.py b/bot/__main__.py index a976f7c..5211c62 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -10,7 +10,7 @@ sys.path.append(str(project_root)) from bot.config import BotConfig from bot.bot import ComfyUIBot -from discordsend_utils.logging_config import setup_logging +from shared.logging_config import setup_logging def main(): # Setup logging diff --git a/bot/cogs/generate.py b/bot/cogs/generate.py index cfc29e3..f33be89 100644 --- a/bot/cogs/generate.py +++ b/bot/cogs/generate.py @@ -7,7 +7,7 @@ from pathlib import Path from ..embeds.builders import EmbedBuilder from ..services.permissions import require_permission, Permissions -from ...discordsend_utils.workflow_builder import WorkflowBuilder +from ...shared.workflow import WorkflowBuilder logger = logging.getLogger(__name__) diff --git a/config.yaml.example b/config.yaml.example new file mode 100644 index 0000000..379d097 --- /dev/null +++ b/config.yaml.example @@ -0,0 +1,48 @@ +# ComfyUI-DiscordSend Bot Configuration +# Copy this file to bot/config.yaml and fill in your values +# Environment variables will override these settings (see .env.example) + +discord: + # Your Discord bot token (required) + # Get one at: https://discord.com/developers/applications + token: "YOUR_BOT_TOKEN_HERE" + + # Your Discord application ID (optional) + # Used for slash command registration + application_id: "YOUR_APP_ID_HERE" + +comfyui: + # ComfyUI server URL (default: http://127.0.0.1:8188) + url: "http://127.0.0.1:8188" + + # WebSocket URL for real-time updates (auto-derived from url if not set) + # ws_url: "ws://127.0.0.1:8188/ws" + + # Request timeout in seconds + timeout: 30 + +defaults: + # Maximum pending jobs per user per server + max_queue_per_user: 3 + + # How often to update progress embeds (seconds) + progress_update_interval: 2.0 + + # Default generation parameters + default_steps: 20 + default_cfg: 7.0 + default_width: 512 + default_height: 512 + + # Path to default workflow (relative to bot/data/) + # workflow_path: "default_workflow_api.json" + +database: + # Database URL (default: SQLite in bot/data/bot.db) + # For PostgreSQL: postgresql+asyncpg://user:pass@localhost/dbname + # url: "sqlite+aiosqlite:///bot/data/bot.db" + +security: + # Restrict bot to specific guild IDs (empty = all guilds allowed) + allowed_guilds: [] + # Example: allowed_guilds: [123456789, 987654321] diff --git a/discordsend_utils/__init__.py b/discordsend_utils/__init__.py deleted file mode 100644 index 5ac0150..0000000 --- a/discordsend_utils/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -""" -ComfyUI-DiscordSend Utility Package - -Shared utilities for Discord integration, sanitization, and GitHub CDN operations. -""" - -from .sanitizer import sanitize_json_for_export -from .github_integration import update_github_cdn_urls -from .prompt_extractor import extract_prompts_from_workflow -from .discord_api import DiscordWebhookClient, validate_webhook_url, send_to_discord_with_retry -from .image_processing import tensor_to_numpy_uint8 - -__all__ = [ - 'sanitize_json_for_export', - 'update_github_cdn_urls', - 'extract_prompts_from_workflow', - 'DiscordWebhookClient', - 'validate_webhook_url', - 'send_to_discord_with_retry', - 'tensor_to_numpy_uint8', -] diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..c7bf1e6 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,10 @@ +""" +ComfyUI-DiscordSend Node Implementations + +This package contains the ComfyUI custom nodes for sending media to Discord. +""" + +from .image_node import DiscordSendSaveImage +from .video_node import DiscordSendSaveVideo + +__all__ = ['DiscordSendSaveImage', 'DiscordSendSaveVideo'] diff --git a/discord_image_node.py b/nodes/image_node.py similarity index 99% rename from discord_image_node.py rename to nodes/image_node.py index faf0a3c..73f9021 100644 --- a/discord_image_node.py +++ b/nodes/image_node.py @@ -17,9 +17,9 @@ from uuid import uuid4 from typing import Any, Union, List, Optional # Import shared utilities -from discordsend_utils import ( - sanitize_json_for_export, - update_github_cdn_urls, +from shared import ( + sanitize_json_for_export, + update_github_cdn_urls, extract_prompts_from_workflow, send_to_discord_with_retry, tensor_to_numpy_uint8 diff --git a/discord_video_node.py b/nodes/video_node.py similarity index 99% rename from discord_video_node.py rename to nodes/video_node.py index bea62c5..de5563b 100644 --- a/discord_video_node.py +++ b/nodes/video_node.py @@ -24,7 +24,7 @@ import functools import server # Import shared utilities -from discordsend_utils 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 # Define cached decorator for local use def cached(max_size=None): """ @@ -47,16 +47,18 @@ def cached(max_size=None): return decorator(func) return decorator -# Constants and configurations +# Define constants and variables previously imported from discordsend_utils ffmpeg_path = None ENCODE_ARGS = ("utf-8", "ignore") +floatOrInt = ("FLOAT", "INT") +imageOrLatent = ("IMAGE", "LATENT") +BIGMAX = 1000000 has_vhs_formats = False -# Try to import ComfyUI dependencies +# Try to import ProgressBar from comfy.utils try: from comfy.utils import ProgressBar except ImportError: - # Define minimal version of ProgressBar if not running in ComfyUI class ProgressBar: def __init__(self, total): self.total = total diff --git a/shared/__init__.py b/shared/__init__.py new file mode 100644 index 0000000..e1e7f59 --- /dev/null +++ b/shared/__init__.py @@ -0,0 +1,40 @@ +""" +ComfyUI-DiscordSend Shared Utilities + +This package contains shared utilities used by both ComfyUI nodes and the Discord bot. +Organized into subpackages: +- discord: Discord webhook and message utilities +- media: Image and video processing utilities +- workflow: ComfyUI workflow manipulation utilities +""" + +# Re-export commonly used utilities for convenience +from .workflow.sanitizer import sanitize_json_for_export +from .workflow.prompt_extractor import extract_prompts_from_workflow +from .workflow.workflow_builder import WorkflowBuilder +from .discord.webhook_client import ( + DiscordWebhookClient, + validate_webhook_url, + send_to_discord_with_retry +) +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 + +__all__ = [ + # Workflow utilities + 'sanitize_json_for_export', + 'extract_prompts_from_workflow', + 'WorkflowBuilder', + # Discord utilities + 'DiscordWebhookClient', + 'validate_webhook_url', + 'send_to_discord_with_retry', + # Media utilities + 'tensor_to_numpy_uint8', + # GitHub integration + 'update_github_cdn_urls', + # Logging + 'setup_logging', + 'get_logger', +] diff --git a/shared/discord/__init__.py b/shared/discord/__init__.py new file mode 100644 index 0000000..2b3c397 --- /dev/null +++ b/shared/discord/__init__.py @@ -0,0 +1,21 @@ +""" +Discord Integration Utilities + +Provides webhook client, message building, and CDN URL handling. +""" + +from .webhook_client import ( + DiscordWebhookClient, + validate_webhook_url, + sanitize_webhook_for_logging, + send_to_discord_with_retry, + validate_file_for_discord +) + +__all__ = [ + 'DiscordWebhookClient', + 'validate_webhook_url', + 'sanitize_webhook_for_logging', + 'send_to_discord_with_retry', + 'validate_file_for_discord', +] diff --git a/discordsend_utils/discord_api.py b/shared/discord/webhook_client.py similarity index 100% rename from discordsend_utils/discord_api.py rename to shared/discord/webhook_client.py diff --git a/discordsend_utils/github_integration.py b/shared/github_integration.py similarity index 100% rename from discordsend_utils/github_integration.py rename to shared/github_integration.py diff --git a/discordsend_utils/logging_config.py b/shared/logging_config.py similarity index 100% rename from discordsend_utils/logging_config.py rename to shared/logging_config.py diff --git a/shared/media/__init__.py b/shared/media/__init__.py new file mode 100644 index 0000000..f211be9 --- /dev/null +++ b/shared/media/__init__.py @@ -0,0 +1,11 @@ +""" +Media Processing Utilities + +Provides image and video processing functions. +""" + +from .image_processing import tensor_to_numpy_uint8 + +__all__ = [ + 'tensor_to_numpy_uint8', +] diff --git a/discordsend_utils/image_processing.py b/shared/media/image_processing.py similarity index 100% rename from discordsend_utils/image_processing.py rename to shared/media/image_processing.py diff --git a/shared/workflow/__init__.py b/shared/workflow/__init__.py new file mode 100644 index 0000000..126af73 --- /dev/null +++ b/shared/workflow/__init__.py @@ -0,0 +1,15 @@ +""" +Workflow Manipulation Utilities + +Provides sanitization, prompt extraction, and workflow building tools. +""" + +from .sanitizer import sanitize_json_for_export +from .prompt_extractor import extract_prompts_from_workflow +from .workflow_builder import WorkflowBuilder + +__all__ = [ + 'sanitize_json_for_export', + 'extract_prompts_from_workflow', + 'WorkflowBuilder', +] diff --git a/discordsend_utils/prompt_extractor.py b/shared/workflow/prompt_extractor.py similarity index 100% rename from discordsend_utils/prompt_extractor.py rename to shared/workflow/prompt_extractor.py diff --git a/discordsend_utils/sanitizer.py b/shared/workflow/sanitizer.py similarity index 100% rename from discordsend_utils/sanitizer.py rename to shared/workflow/sanitizer.py diff --git a/discordsend_utils/workflow_builder.py b/shared/workflow/workflow_builder.py similarity index 100% rename from discordsend_utils/workflow_builder.py rename to shared/workflow/workflow_builder.py diff --git a/tests/test_discord_mock.py b/tests/test_discord_mock.py index 46feb24..82ecb2d 100644 --- a/tests/test_discord_mock.py +++ b/tests/test_discord_mock.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock, patch # Add parent directory to path for imports sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from discordsend_utils.discord_api import send_to_discord_with_retry +from shared.discord import send_to_discord_with_retry class TestDiscordAPI(unittest.TestCase): """Tests for the Discord API utility with mocked network responses.""" diff --git a/tests/test_github_validation.py b/tests/test_github_validation.py index dd5dcb7..917b87c 100644 --- a/tests/test_github_validation.py +++ b/tests/test_github_validation.py @@ -13,7 +13,7 @@ sys.modules["comfy"] = MagicMock() sys.modules["comfy.cli_args"] = MagicMock() # Now we can safely import -from discordsend_utils.github_integration import validate_github_repo, validate_file_path, update_github_cdn_urls +from shared.github_integration import validate_github_repo, validate_file_path, update_github_cdn_urls class TestGitHubValidation(unittest.TestCase): diff --git a/tests/test_utils.py b/tests/test_utils.py index 7a97c2f..87108d7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -20,9 +20,9 @@ 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__)))) -from discordsend_utils.sanitizer import sanitize_json_for_export -from discordsend_utils.discord_api import validate_webhook_url, sanitize_webhook_for_logging, send_to_discord_with_retry, DiscordWebhookClient -from discordsend_utils.github_integration import update_github_cdn_urls +from shared.workflow.sanitizer import sanitize_json_for_export +from shared.discord.webhook_client import validate_webhook_url, sanitize_webhook_for_logging, send_to_discord_with_retry, DiscordWebhookClient +from shared.github_integration import update_github_cdn_urls class TestSanitizer(unittest.TestCase): @@ -176,7 +176,7 @@ class TestSSRFPrevention(unittest.TestCase): self.assertIn("Invalid webhook URL", str(cm.exception)) - @patch('discordsend_utils.discord_api.requests.post') + @patch('shared.discord.webhook_client.requests.post') def test_send_to_discord_allows_valid_url(self, mock_post): """Should allow valid Discord URLs.""" valid_url = "https://discord.com/api/webhooks/123/abc" @@ -210,7 +210,7 @@ class TestWebhookSanitization(unittest.TestCase): class TestDiscordWebhookClient(unittest.TestCase): """Tests for DiscordWebhookClient security features.""" - @patch('discordsend_utils.discord_api.requests.post') + @patch('shared.discord.webhook_client.requests.post') def test_exception_token_leakage(self, mock_post): """Should redact tokens from exception messages in last_error.""" token = "SUPER_SECRET_TOKEN" @@ -232,8 +232,8 @@ class TestDiscordWebhookClient(unittest.TestCase): class TestGitHubIntegration(unittest.TestCase): """Tests for GitHub integration security features.""" - @patch('discordsend_utils.github_integration.requests.put') - @patch('discordsend_utils.github_integration.requests.get') + @patch('shared.github_integration.requests.put') + @patch('shared.github_integration.requests.get') def test_github_token_redaction_in_response(self, mock_get, mock_put): """Should redact GitHub token from error messages including response text.""" token = "ghp_SECRET_TOKEN" diff --git a/tests/test_workflow.py b/tests/test_workflow.py index bd84f32..4335056 100644 --- a/tests/test_workflow.py +++ b/tests/test_workflow.py @@ -6,8 +6,7 @@ import json # Add parent directory to path for imports sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from discordsend_utils.workflow_builder import WorkflowBuilder -from discordsend_utils.prompt_extractor import extract_prompts_from_workflow +from shared.workflow import WorkflowBuilder, extract_prompts_from_workflow class TestWorkflowBuilder(unittest.TestCase): """Tests for the WorkflowBuilder class."""