diff --git a/py/api/admin.py b/py/api/admin.py index 122dde2..914698c 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -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" diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 9534063..e3a8949 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -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()