Vito Sansevero
113d663fce
- Fix XSS in showFullPrompt() across admin.js, gallery.js, metadata.html by replacing HTML string interpolation with DOM manipulation (createElement + textContent), eliminating unsafe user content injection - Fix stack trace exposure in py/api error responses by replacing str(e) with generic messages and adding logger.exception() for server-side traceability
1139 lines
47 KiB
Python
1139 lines
47 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.exception("Scan error")
|
|
yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\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
|