268 lines
9.3 KiB
Python
268 lines
9.3 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
ComfyUI WebDAV Output Sync Service
|
||
|
|
|
||
|
|
Monitors ComfyUI's output directory and automatically uploads new files
|
||
|
|
to HiDrive WebDAV storage.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
python3 comfyui_webdav_sync.py
|
||
|
|
|
||
|
|
Environment Variables:
|
||
|
|
WEBDAV_URL: WebDAV server URL (default: https://webdav.hidrive.ionos.com/)
|
||
|
|
WEBDAV_USERNAME: WebDAV username
|
||
|
|
WEBDAV_PASSWORD: WebDAV password
|
||
|
|
WEBDAV_REMOTE_PATH: Remote directory path (default: /users/valknar/Pictures/AI/ComfyUI)
|
||
|
|
COMFYUI_OUTPUT_DIR: Local directory to watch (default: /workspace/ComfyUI/output)
|
||
|
|
SYNC_DELAY: Seconds to wait after file write before upload (default: 2)
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
import logging
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Set
|
||
|
|
from watchdog.observers import Observer
|
||
|
|
from watchdog.events import FileSystemEventHandler, FileCreatedEvent, FileModifiedEvent
|
||
|
|
from webdav3.client import Client
|
||
|
|
|
||
|
|
# Configure logging
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.INFO,
|
||
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||
|
|
handlers=[
|
||
|
|
logging.StreamHandler(sys.stdout)
|
||
|
|
]
|
||
|
|
)
|
||
|
|
logger = logging.getLogger('comfyui-webdav-sync')
|
||
|
|
|
||
|
|
# Configuration from environment variables
|
||
|
|
WEBDAV_URL = os.getenv('WEBDAV_URL')
|
||
|
|
WEBDAV_USERNAME = os.getenv('WEBDAV_USERNAME')
|
||
|
|
WEBDAV_PASSWORD = os.getenv('WEBDAV_PASSWORD')
|
||
|
|
WEBDAV_REMOTE_PATH = os.getenv('WEBDAV_REMOTE_PATH')
|
||
|
|
COMFYUI_OUTPUT_DIR = os.getenv('COMFYUI_OUTPUT_DIR', '/workspace/ComfyUI/output')
|
||
|
|
SYNC_DELAY = int(os.getenv('SYNC_DELAY', '2')) # Wait 2 seconds after file write
|
||
|
|
RETRY_ATTEMPTS = int(os.getenv('RETRY_ATTEMPTS', '3'))
|
||
|
|
RETRY_DELAY = int(os.getenv('RETRY_DELAY', '5'))
|
||
|
|
|
||
|
|
# Validate required credentials
|
||
|
|
if not all([WEBDAV_URL, WEBDAV_USERNAME, WEBDAV_PASSWORD, WEBDAV_REMOTE_PATH]):
|
||
|
|
logger.error("Missing required WebDAV environment variables!")
|
||
|
|
logger.error("Required: WEBDAV_URL, WEBDAV_USERNAME, WEBDAV_PASSWORD, WEBDAV_REMOTE_PATH")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
class WebDAVClient:
|
||
|
|
"""WebDAV client wrapper with retry logic"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
options = {
|
||
|
|
'webdav_hostname': WEBDAV_URL,
|
||
|
|
'webdav_login': WEBDAV_USERNAME,
|
||
|
|
'webdav_password': WEBDAV_PASSWORD,
|
||
|
|
'webdav_timeout': 300, # 5 minutes for large files
|
||
|
|
}
|
||
|
|
self.client = Client(options)
|
||
|
|
self._ensure_remote_directory()
|
||
|
|
|
||
|
|
def _ensure_remote_directory(self):
|
||
|
|
"""Ensure the remote directory exists"""
|
||
|
|
try:
|
||
|
|
if not self.client.check(WEBDAV_REMOTE_PATH):
|
||
|
|
logger.info(f"Creating remote directory: {WEBDAV_REMOTE_PATH}")
|
||
|
|
# Create parent directories recursively
|
||
|
|
parts = Path(WEBDAV_REMOTE_PATH).parts
|
||
|
|
current = ''
|
||
|
|
for part in parts:
|
||
|
|
if not part or part == '/':
|
||
|
|
continue
|
||
|
|
current = f"{current}/{part}"
|
||
|
|
if not self.client.check(current):
|
||
|
|
self.client.mkdir(current)
|
||
|
|
logger.info(f"✓ Remote directory ready: {WEBDAV_REMOTE_PATH}")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to create remote directory: {e}")
|
||
|
|
raise
|
||
|
|
|
||
|
|
def upload_file(self, local_path: str, remote_path: str) -> bool:
|
||
|
|
"""Upload a file with retry logic"""
|
||
|
|
for attempt in range(1, RETRY_ATTEMPTS + 1):
|
||
|
|
try:
|
||
|
|
# Ensure parent directory exists
|
||
|
|
remote_dir = str(Path(remote_path).parent)
|
||
|
|
if not self.client.check(remote_dir):
|
||
|
|
self.client.mkdir(remote_dir)
|
||
|
|
|
||
|
|
# Upload file
|
||
|
|
logger.info(f"[{attempt}/{RETRY_ATTEMPTS}] Uploading {Path(local_path).name} -> {remote_path}")
|
||
|
|
self.client.upload_sync(remote_path=remote_path, local_path=local_path)
|
||
|
|
|
||
|
|
# Verify upload
|
||
|
|
if self.client.check(remote_path):
|
||
|
|
file_size = os.path.getsize(local_path)
|
||
|
|
logger.info(f"✓ Upload successful: {Path(local_path).name} ({file_size:,} bytes)")
|
||
|
|
return True
|
||
|
|
else:
|
||
|
|
logger.warning(f"Upload verification failed for {remote_path}")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Upload attempt {attempt} failed: {e}")
|
||
|
|
if attempt < RETRY_ATTEMPTS:
|
||
|
|
logger.info(f"Retrying in {RETRY_DELAY} seconds...")
|
||
|
|
time.sleep(RETRY_DELAY)
|
||
|
|
|
||
|
|
logger.error(f"✗ Failed to upload {local_path} after {RETRY_ATTEMPTS} attempts")
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
class ComfyUIOutputHandler(FileSystemEventHandler):
|
||
|
|
"""Handles file system events in ComfyUI output directory"""
|
||
|
|
|
||
|
|
def __init__(self, webdav_client: WebDAVClient):
|
||
|
|
self.webdav_client = webdav_client
|
||
|
|
self.pending_files: Set[str] = set() # Files waiting for write completion
|
||
|
|
self.uploaded_files: Set[str] = set() # Track uploaded files
|
||
|
|
self.watch_dir = Path(COMFYUI_OUTPUT_DIR)
|
||
|
|
|
||
|
|
def on_created(self, event):
|
||
|
|
"""Handle file creation events"""
|
||
|
|
if event.is_directory:
|
||
|
|
return
|
||
|
|
|
||
|
|
file_path = event.src_path
|
||
|
|
|
||
|
|
# Ignore temp files and hidden files
|
||
|
|
if self._should_ignore(file_path):
|
||
|
|
return
|
||
|
|
|
||
|
|
logger.info(f"New file detected: {Path(file_path).name}")
|
||
|
|
self.pending_files.add(file_path)
|
||
|
|
|
||
|
|
# Schedule upload after delay (to ensure file write is complete)
|
||
|
|
self._schedule_upload(file_path)
|
||
|
|
|
||
|
|
def on_modified(self, event):
|
||
|
|
"""Handle file modification events"""
|
||
|
|
if event.is_directory:
|
||
|
|
return
|
||
|
|
|
||
|
|
file_path = event.src_path
|
||
|
|
|
||
|
|
# Ignore if already uploaded or should be ignored
|
||
|
|
if file_path in self.uploaded_files or self._should_ignore(file_path):
|
||
|
|
return
|
||
|
|
|
||
|
|
# Add to pending if not already there
|
||
|
|
if file_path not in self.pending_files:
|
||
|
|
logger.info(f"Modified file detected: {Path(file_path).name}")
|
||
|
|
self.pending_files.add(file_path)
|
||
|
|
self._schedule_upload(file_path)
|
||
|
|
|
||
|
|
def _should_ignore(self, file_path: str) -> bool:
|
||
|
|
"""Check if file should be ignored"""
|
||
|
|
name = Path(file_path).name
|
||
|
|
|
||
|
|
# Ignore hidden files, temp files, and partial downloads
|
||
|
|
if name.startswith('.') or name.endswith('.tmp') or name.endswith('.part'):
|
||
|
|
return True
|
||
|
|
|
||
|
|
return False
|
||
|
|
|
||
|
|
def _schedule_upload(self, file_path: str):
|
||
|
|
"""Schedule file upload after ensuring write is complete"""
|
||
|
|
def upload_when_ready():
|
||
|
|
time.sleep(SYNC_DELAY) # Wait for file write to complete
|
||
|
|
|
||
|
|
# Verify file still exists and size is stable
|
||
|
|
if not os.path.exists(file_path):
|
||
|
|
logger.warning(f"File disappeared: {file_path}")
|
||
|
|
self.pending_files.discard(file_path)
|
||
|
|
return
|
||
|
|
|
||
|
|
# Check if file size is stable (not being written)
|
||
|
|
size1 = os.path.getsize(file_path)
|
||
|
|
time.sleep(0.5)
|
||
|
|
size2 = os.path.getsize(file_path)
|
||
|
|
|
||
|
|
if size1 != size2:
|
||
|
|
logger.info(f"File still being written: {Path(file_path).name}")
|
||
|
|
# Reschedule
|
||
|
|
self._schedule_upload(file_path)
|
||
|
|
return
|
||
|
|
|
||
|
|
# Upload file
|
||
|
|
self._upload_file(file_path)
|
||
|
|
|
||
|
|
# Run in background thread
|
||
|
|
import threading
|
||
|
|
threading.Thread(target=upload_when_ready, daemon=True).start()
|
||
|
|
|
||
|
|
def _upload_file(self, local_path: str):
|
||
|
|
"""Upload file to WebDAV"""
|
||
|
|
try:
|
||
|
|
# Calculate relative path from watch directory
|
||
|
|
rel_path = Path(local_path).relative_to(self.watch_dir)
|
||
|
|
remote_path = f"{WEBDAV_REMOTE_PATH}/{rel_path}".replace('\\', '/')
|
||
|
|
|
||
|
|
# Upload
|
||
|
|
success = self.webdav_client.upload_file(local_path, remote_path)
|
||
|
|
|
||
|
|
# Track result
|
||
|
|
self.pending_files.discard(local_path)
|
||
|
|
if success:
|
||
|
|
self.uploaded_files.add(local_path)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to upload {local_path}: {e}")
|
||
|
|
self.pending_files.discard(local_path)
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
"""Main entry point"""
|
||
|
|
logger.info("=" * 80)
|
||
|
|
logger.info("ComfyUI WebDAV Output Sync Service")
|
||
|
|
logger.info("=" * 80)
|
||
|
|
logger.info(f"Watch directory: {COMFYUI_OUTPUT_DIR}")
|
||
|
|
logger.info(f"WebDAV URL: {WEBDAV_URL}")
|
||
|
|
logger.info(f"Remote path: {WEBDAV_REMOTE_PATH}")
|
||
|
|
logger.info(f"Sync delay: {SYNC_DELAY}s")
|
||
|
|
logger.info("=" * 80)
|
||
|
|
|
||
|
|
# Verify watch directory exists
|
||
|
|
if not os.path.exists(COMFYUI_OUTPUT_DIR):
|
||
|
|
logger.error(f"Watch directory does not exist: {COMFYUI_OUTPUT_DIR}")
|
||
|
|
logger.info(f"Creating watch directory...")
|
||
|
|
os.makedirs(COMFYUI_OUTPUT_DIR, exist_ok=True)
|
||
|
|
|
||
|
|
# Initialize WebDAV client
|
||
|
|
try:
|
||
|
|
webdav_client = WebDAVClient()
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to initialize WebDAV client: {e}")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
# Create event handler and observer
|
||
|
|
event_handler = ComfyUIOutputHandler(webdav_client)
|
||
|
|
observer = Observer()
|
||
|
|
observer.schedule(event_handler, COMFYUI_OUTPUT_DIR, recursive=True)
|
||
|
|
|
||
|
|
# Start watching
|
||
|
|
observer.start()
|
||
|
|
logger.info("✓ Service started - watching for new files...")
|
||
|
|
|
||
|
|
try:
|
||
|
|
while True:
|
||
|
|
time.sleep(1)
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
logger.info("Shutting down...")
|
||
|
|
observer.stop()
|
||
|
|
|
||
|
|
observer.join()
|
||
|
|
logger.info("Service stopped")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|