perf: run duplicate and output scans off the event loop
The duplicate scan walked the output tree and hashed every file on the aiohttp event loop, and the duplicate delete removed files inline; the output scan ran its per-file database work one executor hop at a time. - _find_duplicate_images_sync / _describe_media_file do the walk and hashing; find_duplicate_images() runs them via _run_in_executor. - _delete_duplicate_images_sync / _delete_one_duplicate do the removal (absolute or output-relative paths) off the loop. - The output scan is split into _collect_output_media_sync, _extract_batch_metadata_sync and _ingest_scan_batch_sync (one executor call per batch for all DB work) with a small SSE generator. - Media collection is shared and de-duplicates by normcase'd path.
This commit is contained in:
Vito Sansevero
+335
-347
@@ -8,7 +8,6 @@ import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
@@ -24,6 +23,47 @@ except ImportError:
|
||||
from database.operations import PromptDatabase
|
||||
|
||||
|
||||
try:
|
||||
from ...utils.hashing import generate_prompt_hash
|
||||
except ImportError:
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff")
|
||||
VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv")
|
||||
SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan
|
||||
|
||||
|
||||
def _collect_media_files(output_dirs, extensions):
|
||||
"""All media files under ``output_dirs`` (blocking), skipping thumbnails.
|
||||
|
||||
Matches extensions case-insensitively and de-duplicates by
|
||||
case-folded path so case-insensitive filesystems do not list a file twice.
|
||||
"""
|
||||
found = []
|
||||
seen = set()
|
||||
for output_dir in output_dirs:
|
||||
for ext in extensions:
|
||||
for pattern in (f"*{ext}", f"*{ext.upper()}"):
|
||||
for media_path in Path(output_dir).rglob(pattern):
|
||||
if "thumbnails" in media_path.parts:
|
||||
continue
|
||||
key = os.path.normcase(str(media_path))
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
found.append(media_path)
|
||||
return found
|
||||
|
||||
|
||||
def _thumbnail_rel_path(rel_path, thumbnail_ext):
|
||||
"""Relative path of the thumbnail the gallery would generate for ``rel_path``."""
|
||||
return f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb{thumbnail_ext}"
|
||||
|
||||
|
||||
def _sse(payload):
|
||||
"""Encode one server-sent event."""
|
||||
return f"data: {json.dumps(payload)}\n\n"
|
||||
|
||||
|
||||
def _image_monitor_module():
|
||||
"""The utils.image_monitor module under either package identity."""
|
||||
try:
|
||||
@@ -131,115 +171,74 @@ class AdminRoutesMixin:
|
||||
)
|
||||
|
||||
async def find_duplicate_images(self):
|
||||
"""Find duplicate images in ComfyUI output directory using content hashing."""
|
||||
self.logger.info("Scanning for duplicate images")
|
||||
"""Find duplicate media in the output directory, off the event loop."""
|
||||
output_dir = await self._run_in_executor(self._find_comfyui_output_dir)
|
||||
return await self._run_in_executor(self._find_duplicate_images_sync, output_dir)
|
||||
|
||||
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(rel_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}")
|
||||
def _find_duplicate_images_sync(self, output_dir):
|
||||
"""Group media files under ``output_dir`` by content hash (blocking)."""
|
||||
if not output_dir or not os.path.isdir(output_dir):
|
||||
self.logger.warning("ComfyUI output directory not found")
|
||||
return []
|
||||
|
||||
output_path = Path(output_dir)
|
||||
extensions = IMAGE_EXTENSIONS + VIDEO_EXTENSIONS
|
||||
media_files = _collect_media_files([output_path], extensions)
|
||||
self.logger.info(f"Found {len(media_files)} media files to analyze")
|
||||
|
||||
file_hashes = {}
|
||||
for index, media_path in enumerate(media_files, start=1):
|
||||
try:
|
||||
info = self._describe_media_file(media_path, output_path)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing file {media_path}: {e}")
|
||||
continue
|
||||
file_hashes.setdefault(info["hash"], []).append(info)
|
||||
if index % 100 == 0:
|
||||
self.logger.info(
|
||||
f"Processed {index}/{len(media_files)} files for duplicate detection"
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
def _describe_media_file(self, media_path, output_path):
|
||||
"""Hash one media file and describe it for the duplicates response."""
|
||||
from urllib.parse import quote
|
||||
|
||||
stat = media_path.stat()
|
||||
rel_path = media_path.relative_to(output_path)
|
||||
extension = media_path.suffix.lower()
|
||||
is_video = extension in VIDEO_EXTENSIONS
|
||||
|
||||
thumbnail_url = None
|
||||
thumb_rel = _thumbnail_rel_path(rel_path, ".jpg" if is_video else extension)
|
||||
if (output_path / thumb_rel).exists():
|
||||
thumbnail_url = f"/prompt_manager/images/serve/{quote(thumb_rel, safe='/')}"
|
||||
|
||||
return {
|
||||
"id": str(hash(str(media_path))),
|
||||
"filename": media_path.name,
|
||||
"path": str(rel_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": "video" if is_video else "image",
|
||||
"is_video": is_video,
|
||||
"hash": self._calculate_file_hash(media_path),
|
||||
}
|
||||
|
||||
def _calculate_file_hash(self, file_path):
|
||||
"""Calculate SHA-256 hash of a file's content."""
|
||||
hash_sha256 = hashlib.sha256()
|
||||
@@ -254,85 +253,26 @@ class AdminRoutesMixin:
|
||||
data = await request.json()
|
||||
image_paths = data.get("image_paths", [])
|
||||
|
||||
if not image_paths:
|
||||
if not image_paths or not isinstance(image_paths, list):
|
||||
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)
|
||||
if not file_path.is_absolute():
|
||||
file_path = output_path / file_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} ({self._public_error(e)})")
|
||||
failed_count += 1
|
||||
output_dir = await self._run_in_executor(self._find_comfyui_output_dir)
|
||||
result = await self._run_in_executor(
|
||||
self._delete_duplicate_images_sync, image_paths, output_dir
|
||||
)
|
||||
|
||||
response_data = {
|
||||
"success": True,
|
||||
"deleted_count": deleted_count,
|
||||
"failed_count": failed_count,
|
||||
"message": f"Deleted {deleted_count} files successfully",
|
||||
"deleted_count": result["deleted_count"],
|
||||
"failed_count": result["failed_count"],
|
||||
"message": f"Deleted {result['deleted_count']} files successfully",
|
||||
}
|
||||
|
||||
if failed_count > 0:
|
||||
response_data["failed_files"] = failed_files
|
||||
response_data["message"] += f", {failed_count} failed"
|
||||
if result["failed_count"] > 0:
|
||||
response_data["failed_files"] = result["failed_files"]
|
||||
response_data["message"] += f", {result['failed_count']} failed"
|
||||
|
||||
return web.json_response(response_data)
|
||||
|
||||
@@ -343,6 +283,58 @@ class AdminRoutesMixin:
|
||||
status=500,
|
||||
)
|
||||
|
||||
def _delete_duplicate_images_sync(self, image_paths, output_dir):
|
||||
"""Delete the given files (absolute or output-relative) within ``output_dir``."""
|
||||
result = {"deleted_count": 0, "failed_count": 0, "failed_files": []}
|
||||
output_path = Path(output_dir) if output_dir else None
|
||||
|
||||
for image_path in image_paths:
|
||||
if output_path is None:
|
||||
failure = "output directory not found"
|
||||
else:
|
||||
try:
|
||||
failure = self._delete_one_duplicate(str(image_path), output_path)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error deleting file {image_path}: {e}")
|
||||
failure = self._public_error(e)
|
||||
|
||||
if failure is None:
|
||||
result["deleted_count"] += 1
|
||||
else:
|
||||
result["failed_count"] += 1
|
||||
result["failed_files"].append(f"{image_path} ({failure})")
|
||||
|
||||
return result
|
||||
|
||||
def _delete_one_duplicate(self, image_path, output_path):
|
||||
"""Delete one file and its thumbnail; return a failure reason or None."""
|
||||
file_path = Path(image_path)
|
||||
if not file_path.is_absolute():
|
||||
file_path = output_path / file_path
|
||||
|
||||
try:
|
||||
rel_path = file_path.resolve().relative_to(output_path.resolve())
|
||||
except ValueError:
|
||||
self.logger.warning(
|
||||
f"Attempted to delete file outside output directory: {image_path}"
|
||||
)
|
||||
return "outside output directory"
|
||||
|
||||
if not file_path.is_file():
|
||||
return "file not found"
|
||||
|
||||
os.remove(file_path)
|
||||
self.logger.info(f"Deleted duplicate image: {image_path}")
|
||||
|
||||
thumbnail_path = output_path / _thumbnail_rel_path(rel_path, file_path.suffix)
|
||||
try:
|
||||
if thumbnail_path.exists():
|
||||
os.remove(thumbnail_path)
|
||||
self.logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
|
||||
except OSError as e:
|
||||
self.logger.warning(f"Could not delete thumbnail for {image_path}: {e}")
|
||||
return None
|
||||
|
||||
async def get_statistics(self, request):
|
||||
"""Get database statistics."""
|
||||
try:
|
||||
@@ -997,175 +989,11 @@ class AdminRoutesMixin:
|
||||
os.unlink(temp_path + suffix)
|
||||
|
||||
async def scan_images(self, request):
|
||||
"""Scan ComfyUI output images for prompt metadata and add them to the database."""
|
||||
|
||||
BATCH_SIZE = 50 # Files per executor call for metadata extraction
|
||||
|
||||
def _collect_media_files(output_dirs):
|
||||
"""Collect all media files from output directories (blocking I/O)."""
|
||||
image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"]
|
||||
video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"]
|
||||
media_extensions = image_extensions + video_extensions
|
||||
all_files = []
|
||||
seen = set()
|
||||
for output_dir in output_dirs:
|
||||
for ext in media_extensions:
|
||||
for pattern in [f"*{ext}", f"*{ext.upper()}"]:
|
||||
for f in output_dir.rglob(pattern):
|
||||
if "thumbnails" not in f.parts:
|
||||
norm = str(f).lower()
|
||||
if norm not in seen:
|
||||
seen.add(norm)
|
||||
all_files.append(f)
|
||||
return all_files
|
||||
|
||||
def _extract_batch_metadata(file_batch):
|
||||
"""Extract metadata from a batch of files in one executor call."""
|
||||
results = []
|
||||
for f in file_batch:
|
||||
try:
|
||||
meta = self._extract_comfyui_metadata(str(f))
|
||||
results.append((f, meta))
|
||||
except Exception:
|
||||
results.append((f, {}))
|
||||
return results
|
||||
|
||||
async def stream_response():
|
||||
try:
|
||||
self.logger.info("Starting image scan operation")
|
||||
|
||||
output_dirs = self._get_all_output_dirs()
|
||||
if not output_dirs:
|
||||
self.logger.error("No output directories found")
|
||||
yield f"data: {json.dumps({'type': 'error', 'message': 'No output directories found. Configure scan directories in Settings.'})}\n\n"
|
||||
return
|
||||
|
||||
dir_names = [self._public_path(d) for d in output_dirs]
|
||||
yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': f'Scanning {len(output_dirs)} directory(ies) for media files...', 'processed': 0, 'found': 0})}\n\n"
|
||||
|
||||
media_files = await self._run_in_executor(
|
||||
_collect_media_files, output_dirs
|
||||
)
|
||||
total_files = len(media_files)
|
||||
|
||||
if total_files == 0:
|
||||
yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'found': 0, 'added': 0, 'linked': 0, 'directories': dir_names})}\n\n"
|
||||
return
|
||||
|
||||
yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': f'Found {total_files} media 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
|
||||
|
||||
# Process files in batches for performance
|
||||
for batch_start in range(0, total_files, BATCH_SIZE):
|
||||
batch = media_files[batch_start : batch_start + BATCH_SIZE]
|
||||
|
||||
# Extract metadata for entire batch in one executor call
|
||||
batch_results = await self._run_in_executor(
|
||||
_extract_batch_metadata, batch
|
||||
)
|
||||
|
||||
for media_file, metadata in batch_results:
|
||||
try:
|
||||
processed_count += 1
|
||||
|
||||
if not metadata:
|
||||
continue
|
||||
|
||||
parsed_data = self._parse_comfyui_prompt(metadata)
|
||||
|
||||
if not (
|
||||
parsed_data.get("prompt")
|
||||
or parsed_data.get("parameters")
|
||||
):
|
||||
continue
|
||||
|
||||
found_count += 1
|
||||
prompt_text = self._extract_readable_prompt(parsed_data)
|
||||
|
||||
if prompt_text and not isinstance(prompt_text, str):
|
||||
prompt_text = str(prompt_text)
|
||||
|
||||
if not (prompt_text and prompt_text.strip()):
|
||||
continue
|
||||
|
||||
prompt_hash = generate_prompt_hash(prompt_text.strip())
|
||||
existing = await self._run_in_executor(
|
||||
self.db.get_prompt_by_hash, prompt_hash
|
||||
)
|
||||
|
||||
if existing:
|
||||
try:
|
||||
await self._run_in_executor(
|
||||
self.db.link_image_to_prompt,
|
||||
existing["id"],
|
||||
str(media_file),
|
||||
)
|
||||
linked_count += 1
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to link {media_file.name} to existing prompt: {e}"
|
||||
)
|
||||
else:
|
||||
prompt_id = await self._run_in_executor(
|
||||
self.db.save_prompt,
|
||||
prompt_text.strip(),
|
||||
"scanned",
|
||||
["auto-scanned"],
|
||||
None,
|
||||
f"Auto-scanned from {media_file.name}",
|
||||
prompt_hash,
|
||||
)
|
||||
if prompt_id:
|
||||
added_count += 1
|
||||
try:
|
||||
await self._run_in_executor(
|
||||
self.db.link_image_to_prompt,
|
||||
prompt_id,
|
||||
str(media_file),
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to link {media_file.name} to new prompt: {e}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Error processing {media_file.name}: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
# Progress update after each batch
|
||||
progress = int(
|
||||
min(batch_start + len(batch), total_files) / total_files * 100
|
||||
)
|
||||
yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Processing file {min(batch_start + len(batch), total_files)}/{total_files}...', 'processed': processed_count, 'found': found_count})}\n\n"
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self.logger.info(
|
||||
f"Scan completed: processed={processed_count}, found={found_count}, "
|
||||
f"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, 'directories': dir_names})}\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"
|
||||
"""Scan ComfyUI output images for prompt metadata and add them to the database.
|
||||
|
||||
Streams progress as server-sent events; all filesystem and database
|
||||
work runs through the executor.
|
||||
"""
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
reason="OK",
|
||||
@@ -1175,11 +1003,171 @@ class AdminRoutesMixin:
|
||||
"Connection": "keep-alive",
|
||||
},
|
||||
)
|
||||
|
||||
await response.prepare(request)
|
||||
|
||||
async for chunk in stream_response():
|
||||
async for chunk in self._scan_images_events():
|
||||
await response.write(chunk.encode("utf-8"))
|
||||
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
async def _scan_images_events(self):
|
||||
"""Yield SSE strings while scanning every configured output directory."""
|
||||
try:
|
||||
self.logger.info("Starting image scan operation")
|
||||
output_dirs = self._get_all_output_dirs()
|
||||
if not output_dirs:
|
||||
self.logger.error("No output directories found")
|
||||
yield _sse(
|
||||
{
|
||||
"type": "error",
|
||||
"message": "No output directories found. "
|
||||
"Configure scan directories in Settings.",
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
dir_names = [self._public_path(d) for d in output_dirs]
|
||||
yield self._scan_progress(
|
||||
0, f"Scanning {len(output_dirs)} directory(ies) for media files..."
|
||||
)
|
||||
|
||||
media_files = await self._run_in_executor(
|
||||
self._collect_output_media_sync, output_dirs
|
||||
)
|
||||
counts = {"processed": 0, "found": 0, "added": 0, "linked": 0}
|
||||
if media_files:
|
||||
yield self._scan_progress(
|
||||
5, f"Found {len(media_files)} media files to process...", counts
|
||||
)
|
||||
async for event in self._scan_batches(media_files, counts):
|
||||
yield event
|
||||
|
||||
self.logger.info(
|
||||
f"Scan completed: processed={counts['processed']}, "
|
||||
f"found={counts['found']}, new_prompts_added={counts['added']}, "
|
||||
f"images_linked_to_existing={counts['linked']}"
|
||||
)
|
||||
yield _sse({"type": "complete", **counts, "directories": dir_names})
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("Scan error")
|
||||
yield _sse(
|
||||
{
|
||||
"type": "error",
|
||||
"message": "An internal error occurred. Check server logs for details.",
|
||||
}
|
||||
)
|
||||
|
||||
async def _scan_batches(self, media_files, counts):
|
||||
"""Process ``media_files`` in executor batches, updating ``counts`` in place.
|
||||
|
||||
Yields one progress event per batch.
|
||||
"""
|
||||
total = len(media_files)
|
||||
for batch_start in range(0, total, SCAN_BATCH_SIZE):
|
||||
batch = media_files[batch_start : batch_start + SCAN_BATCH_SIZE]
|
||||
batch_results = await self._run_in_executor(
|
||||
self._extract_batch_metadata_sync, batch
|
||||
)
|
||||
batch_counts = await self._run_in_executor(
|
||||
self._ingest_scan_batch_sync, batch_results
|
||||
)
|
||||
for key in counts:
|
||||
counts[key] += batch_counts[key]
|
||||
|
||||
done = min(batch_start + len(batch), total)
|
||||
yield self._scan_progress(
|
||||
int(done / total * 100), f"Processing file {done}/{total}...", counts
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
@staticmethod
|
||||
def _scan_progress(progress, status, counts=None):
|
||||
"""SSE progress event for the output scan."""
|
||||
counts = counts or {}
|
||||
return _sse(
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": progress,
|
||||
"status": status,
|
||||
"processed": counts.get("processed", 0),
|
||||
"found": counts.get("found", 0),
|
||||
}
|
||||
)
|
||||
|
||||
def _collect_output_media_sync(self, output_dirs):
|
||||
"""Collect all media files from the output directories (blocking I/O)."""
|
||||
return _collect_media_files(output_dirs, IMAGE_EXTENSIONS + VIDEO_EXTENSIONS)
|
||||
|
||||
def _extract_batch_metadata_sync(self, file_batch):
|
||||
"""Extract metadata from a batch of files in one executor call."""
|
||||
results = []
|
||||
for media_file in file_batch:
|
||||
try:
|
||||
results.append(
|
||||
(media_file, self._extract_comfyui_metadata(str(media_file)))
|
||||
)
|
||||
except Exception:
|
||||
results.append((media_file, {}))
|
||||
return results
|
||||
|
||||
def _ingest_scan_batch_sync(self, batch_results):
|
||||
"""Store the prompts found in a metadata batch (blocking DB work)."""
|
||||
counts = {"processed": 0, "found": 0, "added": 0, "linked": 0}
|
||||
for media_file, metadata in batch_results:
|
||||
counts["processed"] += 1
|
||||
try:
|
||||
outcome = self._ingest_scanned_file(media_file, metadata)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing {media_file.name}: {e}")
|
||||
continue
|
||||
if outcome is None:
|
||||
continue
|
||||
counts["found"] += 1
|
||||
if outcome in ("added", "linked"):
|
||||
counts[outcome] += 1
|
||||
return counts
|
||||
|
||||
def _ingest_scanned_file(self, media_file, metadata):
|
||||
"""Save or link one scanned file; returns 'added', 'linked', 'found' or None."""
|
||||
if not metadata:
|
||||
return None
|
||||
parsed = self._parse_comfyui_prompt(metadata)
|
||||
if not (parsed.get("prompt") or parsed.get("parameters")):
|
||||
return None
|
||||
|
||||
prompt_text = self._extract_readable_prompt(parsed)
|
||||
if prompt_text is not None and not isinstance(prompt_text, str):
|
||||
prompt_text = str(prompt_text)
|
||||
if not (prompt_text and prompt_text.strip()):
|
||||
return "found"
|
||||
prompt_text = prompt_text.strip()
|
||||
|
||||
prompt_hash = generate_prompt_hash(prompt_text)
|
||||
existing = self.db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
try:
|
||||
self.db.link_image_to_prompt(existing["id"], str(media_file))
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to link {media_file.name} to existing prompt: {e}"
|
||||
)
|
||||
return "found"
|
||||
return "linked"
|
||||
|
||||
prompt_id = self.db.save_prompt(
|
||||
prompt_text,
|
||||
"scanned",
|
||||
["auto-scanned"],
|
||||
None,
|
||||
f"Auto-scanned from {media_file.name}",
|
||||
prompt_hash,
|
||||
)
|
||||
if not prompt_id:
|
||||
return "found"
|
||||
try:
|
||||
self.db.link_image_to_prompt(prompt_id, str(media_file))
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to link {media_file.name} to new prompt: {e}")
|
||||
return "added"
|
||||
|
||||
@@ -348,5 +348,219 @@ class TestNoAbsolutePathsInResponses(AdminAPITestCase):
|
||||
self.assertIn('"directories": ["output"]', body)
|
||||
|
||||
|
||||
def _write_png(path, prompt_text=None):
|
||||
"""Write a tiny PNG, optionally carrying a ComfyUI-style 'prompt' chunk."""
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
img = Image.new("RGB", (2, 2), (255, 0, 0))
|
||||
info = PngImagePlugin.PngInfo()
|
||||
if prompt_text is not None:
|
||||
graph = {
|
||||
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt_text}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "bad"}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
||||
},
|
||||
}
|
||||
info.add_text("prompt", json.dumps(graph))
|
||||
img.save(str(path), pnginfo=info)
|
||||
|
||||
|
||||
class TestScansRunOffTheEventLoop(AdminAPITestCase):
|
||||
"""Blocking filesystem work must go through _run_in_executor."""
|
||||
|
||||
def _spy_executor(self):
|
||||
calls = []
|
||||
original = self.api._run_in_executor
|
||||
|
||||
async def spy(func, *args, **kwargs):
|
||||
calls.append(getattr(func, "__name__", repr(func)))
|
||||
return await original(func, *args, **kwargs)
|
||||
|
||||
self.api._run_in_executor = spy
|
||||
return calls
|
||||
|
||||
async def test_scan_duplicates_runs_in_executor(self):
|
||||
(self.output_dir / "a.png").write_bytes(b"dup")
|
||||
(self.output_dir / "b.png").write_bytes(b"dup")
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
calls = self._spy_executor()
|
||||
|
||||
resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual((await resp.json())["duplicates"][0]["count"], 2)
|
||||
self.assertIn("_find_duplicate_images_sync", calls)
|
||||
|
||||
async def test_delete_duplicates_runs_in_executor(self):
|
||||
(self.output_dir / "a.png").write_bytes(b"dup")
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
calls = self._spy_executor()
|
||||
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/delete_duplicate_images",
|
||||
json={"image_paths": ["a.png"]},
|
||||
)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual((await resp.json())["deleted_count"], 1)
|
||||
self.assertIn("_delete_duplicate_images_sync", calls)
|
||||
|
||||
async def test_output_scan_runs_in_executor(self):
|
||||
_write_png(self.output_dir / "gen.png", prompt_text="a red square")
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
calls = self._spy_executor()
|
||||
|
||||
resp = await self.client.request("POST", "/prompt_manager/scan")
|
||||
|
||||
body = await resp.text()
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertIn('"type": "complete"', body)
|
||||
self.assertIn('"added": 1', body)
|
||||
for name in (
|
||||
"_collect_output_media_sync",
|
||||
"_extract_batch_metadata_sync",
|
||||
"_ingest_scan_batch_sync",
|
||||
):
|
||||
self.assertIn(name, calls)
|
||||
|
||||
async def test_output_scan_links_existing_prompt_on_second_pass(self):
|
||||
_write_png(self.output_dir / "gen.png", prompt_text="a red square")
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
|
||||
first = await (await self.client.request("POST", "/prompt_manager/scan")).text()
|
||||
second = await (
|
||||
await self.client.request("POST", "/prompt_manager/scan")
|
||||
).text()
|
||||
|
||||
self.assertIn('"added": 1', first)
|
||||
self.assertIn('"linked": 1', second)
|
||||
self.assertIn('"added": 0', second)
|
||||
prompts = self.api.db.get_recent_prompts(limit=10)["prompts"]
|
||||
self.assertEqual([p["text"] for p in prompts], ["a red square"])
|
||||
|
||||
async def test_output_scan_without_directories_reports_error(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.outside_dir)]
|
||||
sys.modules.pop("folder_paths", None)
|
||||
self.api._find_comfyui_output_dir = lambda: None
|
||||
|
||||
body = await (await self.client.request("POST", "/prompt_manager/scan")).text()
|
||||
|
||||
self.assertIn('"type": "error"', body)
|
||||
|
||||
|
||||
class TestDuplicateScanSyncBody(AdminAPITestCase):
|
||||
"""Behaviour of the synchronous duplicate scan against a real tree."""
|
||||
|
||||
def _scan(self):
|
||||
return self.api._find_duplicate_images_sync(str(self.output_dir))
|
||||
|
||||
def test_groups_identical_files_and_sorts_by_mtime(self):
|
||||
older = self.output_dir / "older.png"
|
||||
newer = self.output_dir / "sub" / "newer.PNG"
|
||||
newer.parent.mkdir()
|
||||
older.write_bytes(b"same")
|
||||
newer.write_bytes(b"same")
|
||||
os.utime(older, (1_000_000, 1_000_000))
|
||||
os.utime(newer, (2_000_000, 2_000_000))
|
||||
(self.output_dir / "unique.png").write_bytes(b"different")
|
||||
|
||||
groups = self._scan()
|
||||
|
||||
self.assertEqual(len(groups), 1)
|
||||
self.assertEqual(groups[0]["count"], 2)
|
||||
self.assertEqual(
|
||||
[img["filename"] for img in groups[0]["images"]], ["older.png", "newer.PNG"]
|
||||
)
|
||||
self.assertEqual(
|
||||
groups[0]["images"][1]["path"], os.path.join("sub", "newer.PNG")
|
||||
)
|
||||
self.assertEqual(
|
||||
groups[0]["images"][1]["url"], "/prompt_manager/images/serve/sub/newer.PNG"
|
||||
)
|
||||
self.assertNotIn(self.tmpdir, json.dumps(groups))
|
||||
|
||||
def test_thumbnails_dir_is_skipped_and_thumbnail_url_reported(self):
|
||||
(self.output_dir / "a.png").write_bytes(b"same")
|
||||
(self.output_dir / "b.png").write_bytes(b"same")
|
||||
thumbs = self.output_dir / "thumbnails"
|
||||
thumbs.mkdir()
|
||||
(thumbs / "a_thumb.png").write_bytes(b"same")
|
||||
(thumbs / "b_thumb.png").write_bytes(b"other")
|
||||
|
||||
groups = self._scan()
|
||||
|
||||
self.assertEqual(len(groups), 1)
|
||||
by_name = {img["filename"]: img for img in groups[0]["images"]}
|
||||
self.assertEqual(set(by_name), {"a.png", "b.png"})
|
||||
self.assertEqual(
|
||||
by_name["a.png"]["thumbnail_url"],
|
||||
"/prompt_manager/images/serve/thumbnails/a_thumb.png",
|
||||
)
|
||||
self.assertEqual(
|
||||
by_name["b.png"]["thumbnail_url"],
|
||||
"/prompt_manager/images/serve/thumbnails/b_thumb.png",
|
||||
)
|
||||
|
||||
def test_videos_are_flagged_and_unreadable_entries_skipped(self):
|
||||
(self.output_dir / "clip.mp4").write_bytes(b"vid")
|
||||
(self.output_dir / "copy.mp4").write_bytes(b"vid")
|
||||
(self.output_dir / "not_a_file.png").mkdir() # rglob matches, hashing fails
|
||||
|
||||
groups = self._scan()
|
||||
|
||||
self.assertEqual(len(groups), 1)
|
||||
self.assertTrue(all(img["is_video"] for img in groups[0]["images"]))
|
||||
self.assertTrue(
|
||||
all(img["media_type"] == "video" for img in groups[0]["images"])
|
||||
)
|
||||
|
||||
def test_missing_output_dir_gives_empty_list(self):
|
||||
self.assertEqual(
|
||||
self.api._find_duplicate_images_sync(os.path.join(self.tmpdir, "nope")), []
|
||||
)
|
||||
|
||||
|
||||
class TestDeleteDuplicatesSyncBody(AdminAPITestCase):
|
||||
|
||||
def test_deletes_file_and_thumbnail_inside_output(self):
|
||||
target = self.output_dir / "sub" / "x.png"
|
||||
target.parent.mkdir()
|
||||
target.write_bytes(b"x")
|
||||
thumb = self.output_dir / "thumbnails" / "sub" / "x_thumb.png"
|
||||
thumb.parent.mkdir(parents=True)
|
||||
thumb.write_bytes(b"t")
|
||||
|
||||
result = self.api._delete_duplicate_images_sync(
|
||||
[str(target)], str(self.output_dir)
|
||||
)
|
||||
|
||||
self.assertEqual(result["deleted_count"], 1)
|
||||
self.assertEqual(result["failed_count"], 0)
|
||||
self.assertFalse(target.exists())
|
||||
self.assertFalse(thumb.exists())
|
||||
|
||||
def test_refuses_paths_outside_output_and_reports_missing(self):
|
||||
secret = self.outside_dir / "secret.txt"
|
||||
secret.write_text("s")
|
||||
|
||||
result = self.api._delete_duplicate_images_sync(
|
||||
[str(secret), "ghost.png"], str(self.output_dir)
|
||||
)
|
||||
|
||||
self.assertEqual(result["deleted_count"], 0)
|
||||
self.assertEqual(result["failed_count"], 2)
|
||||
self.assertTrue(secret.exists())
|
||||
self.assertTrue(any("outside" in f for f in result["failed_files"]))
|
||||
self.assertTrue(any("not found" in f for f in result["failed_files"]))
|
||||
|
||||
def test_without_output_dir_everything_fails(self):
|
||||
result = self.api._delete_duplicate_images_sync(["a.png"], None)
|
||||
self.assertEqual(result["failed_count"], 1)
|
||||
self.assertEqual(result["deleted_count"], 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user