Vito Sansevero
4548beb723
Phase 4 structural refactoring: - Split monolithic py/api.py (5.3k lines) into domain mixins: prompts.py, images.py, admin.py, logging_routes.py, autotag_routes.py - Extract inline JS from admin.html into web/js/admin.js - Extract inline JS from gallery.html into web/js/gallery.js - Add gzip compression middleware to API responses (Phase 5.4) - Standardize API error envelope with success: false (Phase 5.3) - Add filmstrip image enrichment to prompt list responses (Phase 5.2)
964 lines
42 KiB
Python
964 lines
42 KiB
Python
"""Admin and maintenance API routes for PromptManager."""
|
|
|
|
import asyncio
|
|
import datetime
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sqlite3
|
|
import tempfile
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
from aiohttp import web
|
|
|
|
|
|
class AdminRoutesMixin:
|
|
"""Mixin providing admin, diagnostics, and maintenance API endpoints."""
|
|
|
|
def _register_admin_routes(self, routes):
|
|
@routes.get("/prompt_manager/scan_duplicates")
|
|
async def scan_duplicates_route(request):
|
|
return await self.scan_duplicates_endpoint(request)
|
|
|
|
@routes.post("/prompt_manager/delete_duplicate_images")
|
|
async def delete_duplicate_images_route(request):
|
|
return await self.delete_duplicate_images_endpoint(request)
|
|
|
|
@routes.post("/prompt_manager/cleanup")
|
|
async def cleanup_duplicates_route(request):
|
|
return await self.cleanup_duplicates_endpoint(request)
|
|
|
|
@routes.post("/prompt_manager/maintenance")
|
|
async def maintenance_route(request):
|
|
return await self.run_maintenance(request)
|
|
|
|
@routes.get("/prompt_manager/stats")
|
|
async def get_stats_route(request):
|
|
return await self.get_statistics(request)
|
|
|
|
@routes.get("/prompt_manager/settings")
|
|
async def get_settings_route(request):
|
|
return await self.get_settings(request)
|
|
|
|
@routes.post("/prompt_manager/settings")
|
|
async def save_settings_route(request):
|
|
return await self.save_settings(request)
|
|
|
|
@routes.get("/prompt_manager/backup")
|
|
async def backup_database_route(request):
|
|
return await self.backup_database(request)
|
|
|
|
@routes.post("/prompt_manager/restore")
|
|
async def restore_database_route(request):
|
|
return await self.restore_database(request)
|
|
|
|
@routes.get("/prompt_manager/diagnostics")
|
|
async def run_diagnostics_route(request):
|
|
return await self.run_diagnostics(request)
|
|
|
|
@routes.post("/prompt_manager/diagnostics/test-link")
|
|
async def test_image_link_route(request):
|
|
return await self.test_image_link(request)
|
|
|
|
@routes.post("/prompt_manager/scan")
|
|
async def scan_images_route(request):
|
|
return await self.scan_images(request)
|
|
|
|
async def scan_duplicates_endpoint(self, request):
|
|
"""Scan for duplicate images without removing them."""
|
|
try:
|
|
duplicates = await self.find_duplicate_images()
|
|
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"duplicates": duplicates,
|
|
"message": f"Found {len(duplicates)} groups of duplicate images",
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Scan duplicates error: {e}")
|
|
return web.json_response(
|
|
{"success": False, "error": f"Failed to scan duplicate images: {str(e)}"},
|
|
status=500,
|
|
)
|
|
|
|
async def cleanup_duplicates_endpoint(self, request):
|
|
"""Cleanup duplicate prompts endpoint."""
|
|
try:
|
|
removed_count = await self._run_in_executor(self.db.cleanup_duplicates)
|
|
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"message": "Cleanup completed",
|
|
"duplicates_removed": removed_count,
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Cleanup error: {e}")
|
|
return web.json_response(
|
|
{"success": False, "error": f"Failed to cleanup duplicates: {str(e)}"},
|
|
status=500,
|
|
)
|
|
|
|
async def find_duplicate_images(self):
|
|
"""Find duplicate images in ComfyUI output directory using content hashing."""
|
|
self.logger.info("Scanning for duplicate images")
|
|
|
|
try:
|
|
output_dir = self._find_comfyui_output_dir()
|
|
if not output_dir:
|
|
self.logger.warning("ComfyUI output directory not found")
|
|
return []
|
|
|
|
output_path = Path(output_dir)
|
|
|
|
image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.webp', '.bmp', '.tiff']
|
|
video_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.webm', '.gif']
|
|
all_extensions = image_extensions + video_extensions
|
|
|
|
media_files = []
|
|
seen_paths = set()
|
|
for ext in all_extensions:
|
|
for pattern in [f"*{ext.lower()}", f"*{ext.upper()}"]:
|
|
for media_path in output_path.rglob(pattern):
|
|
if 'thumbnails' not in media_path.parts:
|
|
normalized_path = str(media_path).lower()
|
|
if normalized_path not in seen_paths:
|
|
seen_paths.add(normalized_path)
|
|
media_files.append(media_path)
|
|
|
|
self.logger.info(f"Found {len(media_files)} media files to analyze")
|
|
|
|
file_hashes = {}
|
|
processed = 0
|
|
|
|
for media_path in media_files:
|
|
try:
|
|
file_hash = self._calculate_file_hash(media_path)
|
|
|
|
if file_hash not in file_hashes:
|
|
file_hashes[file_hash] = []
|
|
|
|
stat = media_path.stat()
|
|
rel_path = media_path.relative_to(output_path)
|
|
extension = media_path.suffix.lower()
|
|
is_video = extension in [ext.lower() for ext in video_extensions]
|
|
media_type = 'video' if is_video else 'image'
|
|
|
|
# Check if thumbnail exists
|
|
thumbnail_url = None
|
|
thumbnails_dir = output_path / "thumbnails"
|
|
if thumbnails_dir.exists():
|
|
thumbnail_ext = '.jpg' if is_video else extension
|
|
rel_path_no_ext = rel_path.with_suffix('')
|
|
thumbnail_rel_path = f"thumbnails/{rel_path_no_ext.as_posix()}_thumb{thumbnail_ext}"
|
|
thumbnail_abs_path = output_path / thumbnail_rel_path
|
|
|
|
if thumbnail_abs_path.exists():
|
|
from urllib.parse import quote
|
|
thumbnail_url = f'/prompt_manager/images/serve/{quote(thumbnail_rel_path, safe="/")}'
|
|
|
|
image_info = {
|
|
'id': str(hash(str(media_path))),
|
|
'filename': media_path.name,
|
|
'path': str(media_path),
|
|
'relative_path': str(rel_path),
|
|
'url': f'/prompt_manager/images/serve/{rel_path.as_posix()}',
|
|
'thumbnail_url': thumbnail_url,
|
|
'size': stat.st_size,
|
|
'modified_time': stat.st_mtime,
|
|
'extension': extension,
|
|
'media_type': media_type,
|
|
'is_video': is_video,
|
|
'hash': file_hash
|
|
}
|
|
|
|
file_hashes[file_hash].append(image_info)
|
|
processed += 1
|
|
|
|
if processed % 100 == 0:
|
|
self.logger.info(f"Processed {processed}/{len(media_files)} files for duplicate detection")
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error processing file {media_path}: {e}")
|
|
continue
|
|
|
|
# Find duplicates (groups with more than one file)
|
|
duplicates = []
|
|
for file_hash, images in file_hashes.items():
|
|
if len(images) > 1:
|
|
images.sort(key=lambda x: x['modified_time'])
|
|
duplicates.append({
|
|
'hash': file_hash,
|
|
'images': images,
|
|
'count': len(images)
|
|
})
|
|
|
|
self.logger.info(f"Found {len(duplicates)} groups of duplicate images")
|
|
return duplicates
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error finding duplicate images: {e}")
|
|
return []
|
|
|
|
def _calculate_file_hash(self, file_path):
|
|
"""Calculate SHA-256 hash of a file's content."""
|
|
hash_sha256 = hashlib.sha256()
|
|
with open(file_path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(4096), b""):
|
|
hash_sha256.update(chunk)
|
|
return hash_sha256.hexdigest()
|
|
|
|
async def delete_duplicate_images_endpoint(self, request):
|
|
"""Delete duplicate image files from disk."""
|
|
try:
|
|
data = await request.json()
|
|
image_paths = data.get('image_paths', [])
|
|
|
|
if not image_paths:
|
|
return web.json_response(
|
|
{"success": False, "error": "No image paths provided"},
|
|
status=400,
|
|
)
|
|
|
|
deleted_count = 0
|
|
failed_count = 0
|
|
failed_files = []
|
|
|
|
for image_path in image_paths:
|
|
try:
|
|
# Ensure the path is within the output directory for security
|
|
output_dir = self._find_comfyui_output_dir()
|
|
if not output_dir:
|
|
failed_files.append(f"{image_path} (output directory not found)")
|
|
failed_count += 1
|
|
continue
|
|
|
|
output_path = Path(output_dir)
|
|
file_path = Path(image_path)
|
|
|
|
# Security check - ensure file is within output directory
|
|
try:
|
|
file_path.resolve().relative_to(output_path.resolve())
|
|
except ValueError:
|
|
self.logger.warning(f"Attempted to delete file outside output directory: {image_path}")
|
|
failed_files.append(f"{image_path} (outside output directory)")
|
|
failed_count += 1
|
|
continue
|
|
|
|
if file_path.exists() and file_path.is_file():
|
|
os.remove(file_path)
|
|
deleted_count += 1
|
|
self.logger.info(f"Deleted duplicate image: {image_path}")
|
|
|
|
# Also try to delete associated thumbnail if it exists
|
|
try:
|
|
rel_path = file_path.relative_to(output_path)
|
|
rel_path_no_ext = rel_path.with_suffix('')
|
|
thumbnail_path = output_path / "thumbnails" / f"{rel_path_no_ext.as_posix()}_thumb{file_path.suffix}"
|
|
if thumbnail_path.exists():
|
|
os.remove(thumbnail_path)
|
|
self.logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
|
|
except Exception as e:
|
|
self.logger.warning(f"Could not delete thumbnail for {image_path}: {e}")
|
|
else:
|
|
failed_files.append(f"{image_path} (file not found)")
|
|
failed_count += 1
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error deleting file {image_path}: {e}")
|
|
failed_files.append(f"{image_path} ({str(e)})")
|
|
failed_count += 1
|
|
|
|
response_data = {
|
|
"success": True,
|
|
"deleted_count": deleted_count,
|
|
"failed_count": failed_count,
|
|
"message": f"Deleted {deleted_count} files successfully"
|
|
}
|
|
|
|
if failed_count > 0:
|
|
response_data["failed_files"] = failed_files
|
|
response_data["message"] += f", {failed_count} failed"
|
|
|
|
return web.json_response(response_data)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Delete duplicate images error: {e}")
|
|
return web.json_response(
|
|
{"success": False, "error": f"Failed to delete duplicate images: {str(e)}"},
|
|
status=500,
|
|
)
|
|
|
|
async def get_statistics(self, request):
|
|
"""Get database statistics."""
|
|
try:
|
|
stats = await self._run_in_executor(self.db.get_statistics)
|
|
|
|
return web.json_response({"success": True, "stats": stats})
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Stats error: {e}")
|
|
return web.json_response(
|
|
{"success": False, "error": f"Failed to get statistics: {str(e)}"},
|
|
status=500,
|
|
)
|
|
|
|
async def get_settings(self, request):
|
|
"""Get current settings."""
|
|
try:
|
|
from ..config import PromptManagerConfig, GalleryConfig
|
|
|
|
# Get monitored directories from image monitor singleton if available
|
|
monitored_dirs = []
|
|
try:
|
|
import sys
|
|
monitor_module = None
|
|
for mod_name in list(sys.modules.keys()):
|
|
if 'image_monitor' in mod_name and hasattr(sys.modules[mod_name], '_monitor_instance'):
|
|
monitor_module = sys.modules[mod_name]
|
|
break
|
|
|
|
if monitor_module and monitor_module._monitor_instance is not None:
|
|
monitored_dirs = getattr(monitor_module._monitor_instance, 'monitored_directories', [])
|
|
elif GalleryConfig.MONITORING_DIRECTORIES:
|
|
monitored_dirs = GalleryConfig.MONITORING_DIRECTORIES
|
|
except Exception:
|
|
if GalleryConfig.MONITORING_DIRECTORIES:
|
|
monitored_dirs = GalleryConfig.MONITORING_DIRECTORIES
|
|
|
|
return web.json_response({
|
|
"success": True,
|
|
"settings": {
|
|
"result_timeout": PromptManagerConfig.RESULT_TIMEOUT,
|
|
"webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE,
|
|
"gallery_root_path": GalleryConfig.MONITORING_DIRECTORIES[0] if GalleryConfig.MONITORING_DIRECTORIES else "",
|
|
"monitored_directories": monitored_dirs
|
|
}
|
|
})
|
|
except Exception as e:
|
|
return web.json_response(
|
|
{"success": False, "error": f"Failed to get settings: {str(e)}"},
|
|
status=500,
|
|
)
|
|
|
|
async def save_settings(self, request):
|
|
"""Save settings."""
|
|
try:
|
|
from ..config import PromptManagerConfig, GalleryConfig
|
|
|
|
data = await request.json()
|
|
restart_required = False
|
|
|
|
# Update in-memory config
|
|
if 'result_timeout' in data:
|
|
PromptManagerConfig.RESULT_TIMEOUT = data['result_timeout']
|
|
if 'webui_display_mode' in data:
|
|
PromptManagerConfig.WEBUI_DISPLAY_MODE = data['webui_display_mode']
|
|
|
|
# Handle gallery root path
|
|
if 'gallery_root_path' in data:
|
|
new_path = data['gallery_root_path'].strip()
|
|
old_path = GalleryConfig.MONITORING_DIRECTORIES[0] if GalleryConfig.MONITORING_DIRECTORIES else ""
|
|
|
|
if new_path != old_path:
|
|
if new_path:
|
|
from pathlib import Path as _Path
|
|
resolved = _Path(new_path).resolve()
|
|
if not resolved.is_dir():
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': f'Gallery path does not exist or is not a directory: {new_path}'
|
|
}, status=400)
|
|
blocked = ['/etc', '/usr', '/bin', '/sbin', '/boot', '/proc', '/sys', '/dev',
|
|
'/var/log', '/root', 'C:\\Windows', 'C:\\Program Files']
|
|
for b in blocked:
|
|
if str(resolved).startswith(b):
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': 'Gallery path cannot point to a system directory'
|
|
}, status=400)
|
|
GalleryConfig.MONITORING_DIRECTORIES = [new_path]
|
|
else:
|
|
GalleryConfig.MONITORING_DIRECTORIES = []
|
|
restart_required = True
|
|
|
|
# Save to config file for persistence
|
|
config_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
config_file = os.path.join(config_dir, 'config.json')
|
|
|
|
config_data = {
|
|
'web_ui': {
|
|
'result_timeout': PromptManagerConfig.RESULT_TIMEOUT,
|
|
'webui_display_mode': PromptManagerConfig.WEBUI_DISPLAY_MODE
|
|
},
|
|
'gallery': {
|
|
'monitoring': {
|
|
'directories': GalleryConfig.MONITORING_DIRECTORIES
|
|
}
|
|
}
|
|
}
|
|
|
|
try:
|
|
with open(config_file, 'w') as f:
|
|
json.dump(config_data, f, indent=2)
|
|
self.logger.info(f"Settings saved to {config_file}")
|
|
except Exception as save_err:
|
|
self.logger.warning(f"Could not save config file: {save_err}")
|
|
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": "Settings saved successfully",
|
|
"restart_required": restart_required
|
|
})
|
|
except Exception as e:
|
|
return web.json_response(
|
|
{"success": False, "error": f"Failed to save settings: {str(e)}"},
|
|
status=500,
|
|
)
|
|
|
|
async def run_diagnostics(self, request):
|
|
"""Run comprehensive system diagnostics and health checks."""
|
|
try:
|
|
results = {}
|
|
|
|
# Check database
|
|
try:
|
|
db_path = "prompts.db"
|
|
if os.path.exists(db_path):
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
cursor = conn.execute("SELECT COUNT(*) as count FROM prompts")
|
|
prompt_count = cursor.fetchone()['count']
|
|
|
|
cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='generated_images'")
|
|
has_images_table = cursor.fetchone() is not None
|
|
|
|
if has_images_table:
|
|
cursor = conn.execute("SELECT COUNT(*) as count FROM generated_images")
|
|
image_count = cursor.fetchone()['count']
|
|
else:
|
|
image_count = 0
|
|
|
|
results['database'] = {
|
|
'status': 'ok',
|
|
'prompt_count': prompt_count,
|
|
'has_images_table': has_images_table,
|
|
'image_count': image_count
|
|
}
|
|
else:
|
|
results['database'] = {
|
|
'status': 'error',
|
|
'message': f'Database file not found: {db_path}'
|
|
}
|
|
except Exception as e:
|
|
results['database'] = {
|
|
'status': 'error',
|
|
'message': f'Database error: {str(e)}'
|
|
}
|
|
|
|
# Check dependencies
|
|
dependencies = {}
|
|
try:
|
|
import watchdog
|
|
dependencies['watchdog'] = True
|
|
except ImportError:
|
|
dependencies['watchdog'] = False
|
|
|
|
try:
|
|
from PIL import Image
|
|
dependencies['PIL'] = True
|
|
except ImportError:
|
|
dependencies['PIL'] = False
|
|
|
|
dependencies['sqlite3'] = True # Always available in Python
|
|
|
|
results['dependencies'] = {
|
|
'status': 'ok' if all(dependencies.values()) else 'error',
|
|
'dependencies': dependencies
|
|
}
|
|
|
|
# Check output directories
|
|
output_dirs = []
|
|
potential_dirs = ["output", "../output", "../../output"]
|
|
|
|
for dir_path in potential_dirs:
|
|
abs_path = os.path.abspath(dir_path)
|
|
if os.path.exists(abs_path):
|
|
output_dirs.append(abs_path)
|
|
|
|
results['comfyui_output'] = {
|
|
'status': 'ok' if output_dirs else 'warning',
|
|
'output_dirs': output_dirs
|
|
}
|
|
|
|
# Check image monitor status
|
|
try:
|
|
from ...utils.image_monitor import _monitor_instance
|
|
if _monitor_instance is not None:
|
|
monitor_status = _monitor_instance.get_status()
|
|
results['image_monitor'] = {
|
|
'status': 'ok' if monitor_status.get('observer_alive') else 'error',
|
|
**monitor_status
|
|
}
|
|
else:
|
|
results['image_monitor'] = {
|
|
'status': 'error',
|
|
'message': 'Image monitor not initialized'
|
|
}
|
|
except Exception as e:
|
|
results['image_monitor'] = {
|
|
'status': 'error',
|
|
'message': f'Failed to get monitor status: {str(e)}'
|
|
}
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'diagnostics': results
|
|
})
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Diagnostics error: {e}", exc_info=True)
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=500)
|
|
|
|
async def test_image_link(self, request):
|
|
"""Test creating an image link."""
|
|
try:
|
|
data = await request.json()
|
|
prompt_id = data.get('prompt_id')
|
|
test_image_path = data.get('image_path', '/test/fake/image.png')
|
|
|
|
if not prompt_id:
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': 'prompt_id is required'
|
|
}, status=400)
|
|
|
|
test_metadata = {
|
|
'file_info': {
|
|
'size': 1024000,
|
|
'dimensions': [512, 512],
|
|
'format': 'PNG'
|
|
},
|
|
'workflow': {'test': True},
|
|
'prompt': {'test_prompt': 'This is a test image'}
|
|
}
|
|
|
|
try:
|
|
image_id = await self._run_in_executor(
|
|
self.db.link_image_to_prompt,
|
|
prompt_id=str(prompt_id),
|
|
image_path=test_image_path,
|
|
metadata=test_metadata
|
|
)
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'result': {
|
|
'status': 'ok',
|
|
'image_id': image_id,
|
|
'message': f'Test image linked successfully with ID {image_id}'
|
|
}
|
|
})
|
|
except Exception as e:
|
|
return web.json_response({
|
|
'success': False,
|
|
'result': {
|
|
'status': 'error',
|
|
'message': f'Failed to create test link: {str(e)}'
|
|
}
|
|
})
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Test link error: {e}", exc_info=True)
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=500)
|
|
|
|
async def run_maintenance(self, request):
|
|
"""Perform comprehensive database maintenance and optimization."""
|
|
try:
|
|
data = await request.json() if request.content_type == 'application/json' else {}
|
|
operations = data.get('operations', ['cleanup_duplicates', 'vacuum', 'cleanup_orphaned_images'])
|
|
|
|
results = {}
|
|
|
|
def _run_maintenance():
|
|
if 'cleanup_duplicates' in operations:
|
|
try:
|
|
duplicates_removed = self.db.cleanup_duplicates()
|
|
results['cleanup_duplicates'] = {
|
|
'success': True, 'removed_count': duplicates_removed,
|
|
'message': f'Removed {duplicates_removed} duplicate prompts'
|
|
}
|
|
except Exception as e:
|
|
results['cleanup_duplicates'] = {'success': False, 'error': str(e), 'message': 'Failed to cleanup duplicates'}
|
|
|
|
if 'vacuum' in operations:
|
|
try:
|
|
self.db.model.vacuum_database()
|
|
results['vacuum'] = {'success': True, 'message': 'Database vacuum completed successfully'}
|
|
except Exception as e:
|
|
results['vacuum'] = {'success': False, 'error': str(e), 'message': 'Failed to vacuum database'}
|
|
|
|
if 'cleanup_orphaned_images' in operations:
|
|
try:
|
|
orphaned_removed = self.db.cleanup_missing_images()
|
|
results['cleanup_orphaned_images'] = {
|
|
'success': True, 'removed_count': orphaned_removed,
|
|
'message': f'Removed {orphaned_removed} orphaned image records'
|
|
}
|
|
except Exception as e:
|
|
results['cleanup_orphaned_images'] = {'success': False, 'error': str(e), 'message': 'Failed to cleanup orphaned images'}
|
|
|
|
if 'check_hash_duplicates' in operations:
|
|
try:
|
|
hash_duplicates = self.db.check_hash_duplicates()
|
|
results['check_hash_duplicates'] = {
|
|
'success': True, 'duplicate_hashes': len(hash_duplicates),
|
|
'message': f'Found {len(hash_duplicates)} duplicate hash groups'
|
|
}
|
|
except Exception as e:
|
|
results['check_hash_duplicates'] = {'success': False, 'error': str(e), 'message': 'Failed to check hash duplicates'}
|
|
|
|
if 'statistics' in operations:
|
|
try:
|
|
db_info = self.db.model.get_database_info()
|
|
results['statistics'] = {'success': True, 'info': db_info, 'message': 'Database statistics retrieved'}
|
|
except Exception as e:
|
|
results['statistics'] = {'success': False, 'error': str(e), 'message': 'Failed to get database statistics'}
|
|
|
|
if 'prune_orphaned_prompts' in operations:
|
|
try:
|
|
removed_count = self.db.prune_orphaned_prompts()
|
|
results['prune_orphaned_prompts'] = {
|
|
'success': True, 'removed_count': removed_count,
|
|
'message': f'Removed {removed_count} orphaned prompts (prompts with no linked images, excluding protected prompts)'
|
|
}
|
|
except Exception as e:
|
|
results['prune_orphaned_prompts'] = {'success': False, 'error': str(e), 'message': 'Failed to prune orphaned prompts'}
|
|
|
|
if 'check_consistency' in operations:
|
|
try:
|
|
consistency_issues = self.db.check_consistency()
|
|
results['check_consistency'] = {
|
|
'success': True, 'issues_found': len(consistency_issues),
|
|
'issues': consistency_issues[:10],
|
|
'message': f'Found {len(consistency_issues)} consistency issues'
|
|
}
|
|
except Exception as e:
|
|
results['check_consistency'] = {'success': False, 'error': str(e), 'message': 'Failed to check database consistency'}
|
|
|
|
await self._run_in_executor(_run_maintenance)
|
|
|
|
all_successful = all(result.get('success', False) for result in results.values())
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'operations_completed': len(results),
|
|
'all_successful': all_successful,
|
|
'results': results,
|
|
'message': f'Maintenance completed: {len(results)} operations processed'
|
|
})
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Maintenance error: {e}", exc_info=True)
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': f'Maintenance failed: {str(e)}'
|
|
}, status=500)
|
|
|
|
async def backup_database(self, request):
|
|
"""Backup the entire prompts.db database file."""
|
|
try:
|
|
db_path = "prompts.db"
|
|
|
|
if not os.path.exists(db_path):
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': 'Database file not found'
|
|
}, status=404)
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix='.db') as temp_file:
|
|
temp_path = temp_file.name
|
|
|
|
shutil.copy2(db_path, temp_path)
|
|
|
|
with open(temp_path, 'rb') as f:
|
|
file_data = f.read()
|
|
|
|
os.unlink(temp_path)
|
|
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
filename = f"prompts_backup_{timestamp}.db"
|
|
|
|
return web.Response(
|
|
body=file_data,
|
|
content_type='application/octet-stream',
|
|
headers={
|
|
'Content-Disposition': f'attachment; filename="{filename}"',
|
|
'Content-Length': str(len(file_data))
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Backup error: {e}", exc_info=True)
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': f'Failed to backup database: {str(e)}'
|
|
}, status=500)
|
|
|
|
async def restore_database(self, request):
|
|
"""Restore the prompts.db database from uploaded file."""
|
|
try:
|
|
reader = await request.multipart()
|
|
field = await reader.next()
|
|
|
|
if not field or field.name != 'database_file':
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': 'No database file uploaded. Expected field name: database_file'
|
|
}, status=400)
|
|
|
|
MAX_RESTORE_SIZE = 100 * 1024 * 1024 # 100MB
|
|
file_data = await field.read()
|
|
|
|
if not file_data:
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': 'Uploaded file is empty'
|
|
}, status=400)
|
|
|
|
if len(file_data) > MAX_RESTORE_SIZE:
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': f'File too large. Maximum size is {MAX_RESTORE_SIZE // (1024*1024)}MB'
|
|
}, status=400)
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix='.db') as temp_file:
|
|
temp_path = temp_file.name
|
|
temp_file.write(file_data)
|
|
|
|
try:
|
|
with sqlite3.connect(temp_path) as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
|
|
cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'")
|
|
if not cursor.fetchone():
|
|
raise ValueError("Database does not contain a 'prompts' table")
|
|
|
|
cursor = conn.execute("PRAGMA table_info(prompts)")
|
|
columns = [row['name'] for row in cursor.fetchall()]
|
|
required_columns = ['id', 'text', 'created_at']
|
|
|
|
for col in required_columns:
|
|
if col not in columns:
|
|
raise ValueError(f"Database missing required column: {col}")
|
|
|
|
cursor = conn.execute("SELECT COUNT(*) as count FROM prompts")
|
|
prompt_count = cursor.fetchone()['count']
|
|
|
|
db_path = "prompts.db"
|
|
backup_path = f"{db_path}.backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
|
|
|
if os.path.exists(db_path):
|
|
shutil.copy2(db_path, backup_path)
|
|
self.logger.info(f"Current database backed up to: {backup_path}")
|
|
|
|
shutil.copy2(temp_path, db_path)
|
|
|
|
# Reinitialize the database connection
|
|
try:
|
|
from ...database.operations import PromptDatabase
|
|
except ImportError:
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
|
from database.operations import PromptDatabase
|
|
self.db = PromptDatabase()
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'message': f'Database restored successfully. Found {prompt_count} prompts.',
|
|
'prompt_count': prompt_count,
|
|
'backup_created': backup_path if os.path.exists(db_path) else None
|
|
})
|
|
|
|
except sqlite3.Error as e:
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': f'Invalid SQLite database: {str(e)}'
|
|
}, status=400)
|
|
except ValueError as e:
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=400)
|
|
finally:
|
|
if os.path.exists(temp_path):
|
|
os.unlink(temp_path)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Restore error: {e}", exc_info=True)
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': f'Failed to restore database: {str(e)}'
|
|
}, status=500)
|
|
|
|
async def scan_images(self, request):
|
|
"""Scan ComfyUI output images for prompt metadata and add them to the database."""
|
|
|
|
async def stream_response():
|
|
try:
|
|
self.logger.info("Starting image scan operation")
|
|
self.logger.info("Starting scan (timer clearing not implemented yet)")
|
|
|
|
output_dir = self._find_comfyui_output_dir()
|
|
if not output_dir:
|
|
self.logger.error("ComfyUI output directory not found")
|
|
yield f"data: {json.dumps({'type': 'error', 'message': 'ComfyUI output directory not found'})}\n\n"
|
|
return
|
|
|
|
yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Scanning for PNG files...', 'processed': 0, 'found': 0})}\n\n"
|
|
|
|
png_files = await self._run_in_executor(
|
|
lambda: list(Path(output_dir).rglob("*.png"))
|
|
)
|
|
total_files = len(png_files)
|
|
|
|
if total_files == 0:
|
|
yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'found': 0, 'added': 0})}\n\n"
|
|
return
|
|
|
|
yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': f'Found {total_files} PNG files to process...', 'processed': 0, 'found': 0})}\n\n"
|
|
|
|
processed_count = 0
|
|
found_count = 0
|
|
added_count = 0
|
|
linked_count = 0
|
|
|
|
try:
|
|
from ...utils.hashing import generate_prompt_hash
|
|
except ImportError:
|
|
import sys
|
|
current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
sys.path.insert(0, current_dir)
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
for i, png_file in enumerate(png_files):
|
|
try:
|
|
metadata = await self._run_in_executor(
|
|
self._extract_comfyui_metadata, str(png_file)
|
|
)
|
|
processed_count += 1
|
|
|
|
if metadata:
|
|
self.logger.debug(f"Found metadata in {os.path.basename(png_file)}: {list(metadata.keys())}")
|
|
|
|
parsed_data = self._parse_comfyui_prompt(metadata)
|
|
self.logger.debug(f"Parsed data keys: {list(parsed_data.keys())}, has prompt: {bool(parsed_data.get('prompt'))}, has parameters: {bool(parsed_data.get('parameters'))}")
|
|
|
|
if parsed_data.get('prompt') or parsed_data.get('parameters'):
|
|
found_count += 1
|
|
prompt_text = self._extract_readable_prompt(parsed_data)
|
|
|
|
if prompt_text:
|
|
self.logger.debug(f"Found prompt in {os.path.basename(png_file)} (type: {type(prompt_text)}): {str(prompt_text)[:100]}...")
|
|
else:
|
|
self.logger.debug(f"No readable prompt found in {os.path.basename(png_file)}, parsed_data keys: {list(parsed_data.keys())}")
|
|
|
|
if prompt_text and not isinstance(prompt_text, str):
|
|
self.logger.debug(f"Converting prompt_text from {type(prompt_text)} to string")
|
|
prompt_text = str(prompt_text)
|
|
|
|
if prompt_text and prompt_text.strip():
|
|
try:
|
|
prompt_hash = generate_prompt_hash(prompt_text.strip())
|
|
self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...")
|
|
|
|
existing = await self._run_in_executor(
|
|
self.db.get_prompt_by_hash, prompt_hash
|
|
)
|
|
if existing:
|
|
self.logger.debug(f"Found existing prompt ID {existing['id']} for image {os.path.basename(png_file)}")
|
|
try:
|
|
await self._run_in_executor(
|
|
self.db.link_image_to_prompt, existing['id'], str(png_file)
|
|
)
|
|
linked_count += 1
|
|
self.logger.debug(f"Linked image {os.path.basename(png_file)} to existing prompt {existing['id']}")
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to link image {png_file} to existing prompt: {e}")
|
|
else:
|
|
self.logger.debug(f"Saving new prompt from {os.path.basename(png_file)}")
|
|
prompt_id = await self._run_in_executor(
|
|
self.db.save_prompt,
|
|
prompt_text.strip(), 'scanned', ['auto-scanned'],
|
|
f'Auto-scanned from {os.path.basename(png_file)}',
|
|
prompt_hash
|
|
)
|
|
|
|
if prompt_id:
|
|
added_count += 1
|
|
self.logger.info(f"Successfully saved new prompt with ID {prompt_id} from {os.path.basename(png_file)}")
|
|
try:
|
|
await self._run_in_executor(
|
|
self.db.link_image_to_prompt, prompt_id, str(png_file)
|
|
)
|
|
self.logger.debug(f"Linked image {os.path.basename(png_file)} to new prompt {prompt_id}")
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to link image {png_file} to new prompt: {e}")
|
|
else:
|
|
self.logger.error(f"Failed to save prompt from {os.path.basename(png_file)} - no ID returned")
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to save prompt from {png_file}: {e}")
|
|
|
|
# Update progress every 10 files
|
|
if i % 10 == 0 or i == total_files - 1:
|
|
progress = int((i + 1) / total_files * 100)
|
|
status = f"Processing file {i + 1}/{total_files}..."
|
|
|
|
yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': status, 'processed': processed_count, 'found': found_count})}\n\n"
|
|
|
|
await asyncio.sleep(0.01)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error processing {png_file}: {e}")
|
|
continue
|
|
|
|
self.logger.info(f"Scan completed: processed={processed_count}, found={found_count}, new_prompts_added={added_count}, images_linked_to_existing={linked_count}")
|
|
yield f"data: {json.dumps({'type': 'complete', 'processed': processed_count, 'found': found_count, 'added': added_count, 'linked': linked_count})}\n\n"
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Scan error: {e}")
|
|
self.logger.error(f"Scan error traceback: {traceback.format_exc()}")
|
|
yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
|
|
|
|
response = web.StreamResponse(
|
|
status=200,
|
|
reason='OK',
|
|
headers={
|
|
'Content-Type': 'text/event-stream',
|
|
'Cache-Control': 'no-cache',
|
|
'Connection': 'keep-alive'
|
|
}
|
|
)
|
|
|
|
await response.prepare(request)
|
|
|
|
async for chunk in stream_response():
|
|
await response.write(chunk.encode('utf-8'))
|
|
|
|
await response.write_eof()
|
|
return response
|