244 lines
9.8 KiB
Python
244 lines
9.8 KiB
Python
# drivesend_uploader_node.py
|
|
# DriveSend AutoUploader Node - monitors output folder and uploads to Google Drive
|
|
|
|
import os
|
|
import threading
|
|
import logging
|
|
import time
|
|
from pathlib import Path
|
|
from watchdog.observers import Observer
|
|
from .monitor_output import start_monitoring, stop_monitoring, watcher_observer, stop_queue_processor
|
|
from .encrypt_file import FileEncryptHandler, ENCRYPT_EXTENSIONS, get_encryption_key
|
|
from .encrypt_file import stop_queue_processor as stop_encrypt_queue_processor
|
|
from .encrypt_file import reset_queue_processor as reset_encrypt_queue_processor
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)s - %(message)s',
|
|
handlers=[
|
|
logging.FileHandler('drivesend.log'),
|
|
logging.StreamHandler()
|
|
]
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Define encrypt_observer at module level
|
|
encrypt_observer = None
|
|
|
|
|
|
class DriveSendAutoUploaderNode:
|
|
"""ComfyUI node for automatic Google Drive uploads."""
|
|
|
|
CATEGORY = "DriveSend"
|
|
FUNCTION = "start"
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("status",)
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Try to get default watch folder
|
|
default_watch = os.path.join(os.getcwd(), "output")
|
|
|
|
return {
|
|
"required": {
|
|
"watch_folder": ("STRING", {"default": default_watch}),
|
|
"auth_method": (["oauth", "service_account"], {"default": "oauth"}),
|
|
"enable_encryption": ("BOOLEAN", {"default": True}),
|
|
"Post_Delete_Enc": ("BOOLEAN", {"default": False}),
|
|
"Subfolder_Monitor": ("BOOLEAN", {"default": True}),
|
|
"run_process": ("BOOLEAN", {"default": True, "label": "Run Process"}),
|
|
}
|
|
}
|
|
|
|
def start(self, watch_folder, auth_method, enable_encryption, Post_Delete_Enc,
|
|
Subfolder_Monitor, run_process):
|
|
"""
|
|
Start or stop the Google Drive upload monitor.
|
|
|
|
When encryption is enabled, two watchers run:
|
|
1. Encrypt watcher: Creates .enc copies of new images (preserves originals)
|
|
2. Upload watcher: Uploads .enc files to Google Drive
|
|
|
|
This preserves original files so ComfyUI maintains sequential naming.
|
|
"""
|
|
global encrypt_observer
|
|
|
|
logger.info(f"[DriveSend] Starting AutoUploader: watch_folder={watch_folder}, "
|
|
f"auth_method={auth_method}, encryption={enable_encryption}, "
|
|
f"Post_Delete_Enc={Post_Delete_Enc}, Subfolder_Monitor={Subfolder_Monitor}, "
|
|
f"run_process={run_process}")
|
|
|
|
# Handle stopping the process
|
|
if not run_process:
|
|
logger.info("[DriveSend] Stopping AutoUploader monitoring")
|
|
|
|
# Stop upload watcher
|
|
if watcher_observer and watcher_observer.is_alive():
|
|
stop_monitoring()
|
|
logger.info("[DriveSend] Upload watcher stopped")
|
|
|
|
# Stop encryption watcher
|
|
if encrypt_observer and encrypt_observer.is_alive():
|
|
encrypt_observer.stop()
|
|
encrypt_observer.join()
|
|
logger.info("[DriveSend] Encryption watcher stopped")
|
|
|
|
# Stop queue processors
|
|
stop_queue_processor()
|
|
stop_encrypt_queue_processor()
|
|
|
|
stop_message = f"""
|
|
=====================================================================
|
|
🚙🛑 DriveSend - AutoUploader - STOPPED
|
|
=====================================================================
|
|
All monitoring, uploading, and encryption processes for {watch_folder} have been stopped.
|
|
Set 'run_process' to True and run the node again to resume.
|
|
=====================================================================
|
|
"""
|
|
print(stop_message)
|
|
return (f"All monitoring stopped for {watch_folder}",)
|
|
|
|
# Validate watch_folder
|
|
watch_path = Path(watch_folder)
|
|
if not watch_path.exists():
|
|
logger.error(f"[DriveSend] Watch folder does not exist: {watch_folder}")
|
|
return (f"❌ Error: Watch folder does not exist: {watch_folder}",)
|
|
if not watch_path.is_dir():
|
|
logger.error(f"[DriveSend] Watch folder is not a directory: {watch_folder}")
|
|
return (f"❌ Error: Watch folder is not a directory: {watch_folder}",)
|
|
|
|
# Get folder ID from environment
|
|
folder_id = os.environ.get('GOOGLE_DRIVE_FOLDER_ID')
|
|
if not folder_id:
|
|
return (
|
|
"❌ Error: GOOGLE_DRIVE_FOLDER_ID not set.\n\n"
|
|
"Run the DriveSend Setup node first, then:\n"
|
|
"1. Copy credentials to RunPod Secrets\n"
|
|
"2. Restart pod with environment variables loaded",
|
|
)
|
|
|
|
# Validate auth method requirements
|
|
if auth_method == 'oauth':
|
|
client_id = os.environ.get('GOOGLE_CLIENT_ID')
|
|
client_secret = os.environ.get('GOOGLE_CLIENT_SECRET')
|
|
refresh_token = os.environ.get('GOOGLE_REFRESH_TOKEN')
|
|
|
|
if not all([client_id, client_secret, refresh_token]):
|
|
token_file = Path(__file__).parent / 'token.json'
|
|
if not token_file.exists():
|
|
return (
|
|
"❌ Error: OAuth credentials not found.\n\n"
|
|
"Set these environment variables:\n"
|
|
"- GOOGLE_CLIENT_ID\n"
|
|
"- GOOGLE_CLIENT_SECRET\n"
|
|
"- GOOGLE_REFRESH_TOKEN\n\n"
|
|
"Or run DriveSend Setup to authorize.",
|
|
)
|
|
|
|
elif auth_method == 'service_account':
|
|
sa_json = os.environ.get('GOOGLE_SERVICE_ACCOUNT_JSON')
|
|
sa_file = Path(__file__).parent / 'service_account.json'
|
|
|
|
if not sa_json and not sa_file.exists():
|
|
return (
|
|
"❌ Error: Service account credentials not found.\n\n"
|
|
"Either:\n"
|
|
"1. Set GOOGLE_SERVICE_ACCOUNT_JSON environment variable\n"
|
|
"2. Place service_account.json in the node folder\n\n"
|
|
"⚠️ NOTE: Service accounts only work with Google Workspace.\n"
|
|
"For personal Gmail, use OAuth instead.",
|
|
)
|
|
|
|
# Validate encryption key if enabled
|
|
if enable_encryption:
|
|
enc_key = get_encryption_key()
|
|
if not enc_key:
|
|
logger.error("[DriveSend] Encryption enabled but no encryption key found")
|
|
return (
|
|
"❌ Error: Encryption enabled but no key found.\n\n"
|
|
"Set COMFYUI_ENCRYPTION_KEY environment variable\n"
|
|
"or run DriveSend Setup with encryption enabled.",
|
|
)
|
|
|
|
# Reset queue processors before starting
|
|
from .monitor_output import reset_queue_processor as reset_upload_queue
|
|
reset_upload_queue()
|
|
reset_encrypt_queue_processor()
|
|
|
|
# Start upload monitor
|
|
try:
|
|
start_monitoring(
|
|
watch_folder=watch_folder,
|
|
folder_id=folder_id,
|
|
auth_method=auth_method,
|
|
enable_encryption=enable_encryption,
|
|
delete_enc=Post_Delete_Enc,
|
|
subfolder_monitor=Subfolder_Monitor
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"[DriveSend] Failed to start upload monitor: {e}")
|
|
return (f"❌ Error starting upload monitor: {e}",)
|
|
|
|
# Start encryption watcher if enabled (separate from upload watcher)
|
|
if enable_encryption:
|
|
# Stop existing encrypt observer if running
|
|
if encrypt_observer and encrypt_observer.is_alive():
|
|
encrypt_observer.stop()
|
|
encrypt_observer.join()
|
|
|
|
encrypt_handler = FileEncryptHandler(watch_folder, False, Subfolder_Monitor)
|
|
encrypt_observer = Observer()
|
|
encrypt_observer.schedule(encrypt_handler, watch_folder, recursive=Subfolder_Monitor)
|
|
encrypt_observer.start()
|
|
logger.info(f"[DriveSend] Starting encryption monitor for {watch_folder}")
|
|
|
|
def keep_encrypt_alive():
|
|
try:
|
|
while True:
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
if encrypt_observer:
|
|
encrypt_observer.stop()
|
|
if encrypt_observer:
|
|
encrypt_observer.join()
|
|
|
|
threading.Thread(target=keep_encrypt_alive, daemon=True).start()
|
|
|
|
# Build status message
|
|
status_lines = [
|
|
"✅ DriveSend AutoUploader started!",
|
|
"",
|
|
f"📁 Watching: {watch_folder}",
|
|
f"☁️ Uploading to: Google Drive folder {folder_id[:20]}...",
|
|
f"🔐 Auth method: {auth_method}",
|
|
f"🔒 Encryption: {'Enabled (creating .enc copies)' if enable_encryption else 'Disabled'}",
|
|
f"📂 Subfolder monitoring: {'Enabled' if Subfolder_Monitor else 'Disabled'}",
|
|
f"🗑️ Delete .enc after upload: {'Yes' if Post_Delete_Enc else 'No'}",
|
|
"",
|
|
]
|
|
|
|
if enable_encryption:
|
|
status_lines.extend([
|
|
"⚠️ Note: Original files are preserved for ComfyUI naming.",
|
|
" You must manually delete files from output folder.",
|
|
"",
|
|
])
|
|
|
|
status_lines.extend([
|
|
"New files will be uploaded automatically.",
|
|
"Set run_process to False to stop.",
|
|
])
|
|
|
|
return ("\n".join(status_lines),)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DriveSendAutoUploader": DriveSendAutoUploaderNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DriveSendAutoUploader": "🚙📤 DriveSend - AutoUploader"
|
|
}
|