From b2ea2597e664032d377f56849cee5e719013831b Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 8 Oct 2026 16:59:00 -0700 Subject: [PATCH] feat: scan reads prompts embedded in videos; unreadable files no longer log errors The output scan handed every .mp4 to Pillow and logged one ERROR per file. ComfyUI's video savers store the queued graph as JSON in the container's comment tag, so those files carry prompts worth linking. - utils/video_metadata.py reads the comment/description tag through ffprobe with a timeout, normalised to the PNG text-chunk shape so the existing parser handles both; without ffprobe videos are skipped and a single INFO notice is logged - _extract_comfyui_metadata routes videos there and logs unreadable images as warnings, since the scan carries on past them Test hygiene: every aiohttp test base class overrode tearDownAsync without calling super(), so no test client was ever closed. That produced thousands of ResourceWarnings, "Unclosed client session" errors from asyncio that broke log assertions, and doubled the suite's run time. --- py/api/__init__.py | 21 ++-- py/config.py | 3 +- tests/test_api.py | 1 + tests/test_api_admin.py | 85 +++++++++++++++ tests/test_api_autotag.py | 1 + tests/test_api_backup_restore.py | 1 + tests/test_api_core.py | 1 + tests/test_api_images.py | 1 + tests/test_api_logging.py | 1 + tests/test_api_lora.py | 1 + tests/test_api_prompts.py | 1 + tests/test_config.py | 12 ++- tests/test_utils_video_metadata.py | 165 +++++++++++++++++++++++++++++ utils/video_metadata.py | 123 +++++++++++++++++++++ 14 files changed, 404 insertions(+), 13 deletions(-) create mode 100644 tests/test_utils_video_metadata.py create mode 100644 utils/video_metadata.py diff --git a/py/api/__init__.py b/py/api/__init__.py index 4023f75..a71d300 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -31,6 +31,7 @@ from .lora_integration import LoraIntegrationMixin try: from ...database.operations import PromptDatabase from ...utils.logging_config import get_logger + from ...utils.video_metadata import is_video_path, read_video_metadata except ImportError: import sys @@ -39,6 +40,7 @@ except ImportError: ) from database.operations import PromptDatabase from utils.logging_config import get_logger + from utils.video_metadata import is_video_path, read_video_metadata def _get_project_root(): @@ -665,16 +667,21 @@ class PromptManagerAPI( return [] def _extract_comfyui_metadata(self, image_path): - """Extract ComfyUI workflow metadata from PNG image files.""" + """Embedded ComfyUI metadata of an image (PNG text chunks) or a video. + + Videos go through ffprobe (see utils/video_metadata). Unreadable files + are a warning, not an error: a scan over a large output tree meets + truncated and foreign files, and it carries on past them. + """ + if is_video_path(image_path): + return read_video_metadata(image_path) try: with Image.open(image_path) as img: - metadata = {} - if hasattr(img, "text"): - for key, value in img.text.items(): - metadata[key] = value - return metadata + return dict(getattr(img, "text", None) or {}) except Exception as e: - self.logger.error(f"Error reading {image_path}: {e}") + self.logger.warning( + f"Unreadable image {os.path.basename(str(image_path))}: {e}" + ) return {} def _parse_comfyui_prompt(self, metadata): diff --git a/py/config.py b/py/config.py index de8f56f..bf43b5d 100644 --- a/py/config.py +++ b/py/config.py @@ -506,7 +506,8 @@ class PromptManagerConfig: AUTO_BACKUP_INTERVAL = 24 # Hours # Threads for the output scan and thumbnail generation; half the cores by # default so ComfyUI's own generation keeps CPU headroom. 1..max_worker_threads(). - WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2) + DEFAULT_WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2) + WORKER_THREADS = DEFAULT_WORKER_THREADS @classmethod def max_worker_threads(cls) -> int: diff --git a/tests/test_api.py b/tests/test_api.py index c29f98e..f093393 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -38,6 +38,7 @@ class APITestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client # Close every thread's connection first: Windows refuses to delete # an open database file, and executor threads hold connections too. self.api.db.close_all() diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 13b4de9..9e32f55 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -69,6 +69,7 @@ class AdminAPITestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client GalleryConfig.MONITORING_DIRECTORIES = self._orig_gallery_dirs PromptManagerConfig.update_config(self._orig_pm_config) for key, value in self._orig_env.items(): @@ -1703,5 +1704,89 @@ class TestWorkerThreadsSetting(AdminAPITestCase): self.assertIn('"added": 1', body) +def _write_mp4(path, comment): + """Tiny black video with ``comment`` in its container tags (needs ffmpeg).""" + import subprocess + + subprocess.run( + [ + "ffmpeg", + "-v", + "quiet", + "-f", + "lavfi", + "-i", + "color=c=black:s=16x16:d=0.2", + "-metadata", + f"comment={comment}", + "-y", + str(path), + ], + check=True, + timeout=60, + ) + + +class TestScanReadsVideos(AdminAPITestCase): + """Videos saved by ComfyUI carry the graph in their comment tag.""" + + def setUp(self): + super().setUp() + from utils import video_metadata + + video_metadata.ffprobe_path.cache_clear() + self.addCleanup(video_metadata.ffprobe_path.cache_clear) + + @unittest.skipUnless( + shutil.which("ffmpeg") and shutil.which("ffprobe"), "ffmpeg not installed" + ) + async def test_scan_adds_prompts_found_in_videos(self): + graph = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}} + _write_mp4( + self.output_dir / "clip.mp4", json.dumps({"prompt": json.dumps(graph)}) + ) + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + + with self.assertNoLogs(level="ERROR"): + body = await ( + await self.client.request("POST", "/prompt_manager/scan") + ).text() + + self.assertIn('"added": 1', body) + prompts = self.api.db.get_recent_prompts(limit=10)["prompts"] + self.assertEqual([p["text"] for p in prompts], ["a cat"]) + images = self.api.db.get_prompt_images(prompts[0]["id"]) + self.assertTrue(images[0]["image_path"].endswith("clip.mp4")) + + async def test_scan_without_ffprobe_skips_videos_quietly(self): + from unittest.mock import patch + from utils import video_metadata + + (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypisom") + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + + with patch.object(video_metadata.shutil, "which", return_value=None): + with self.assertNoLogs("prompt_manager", level="WARNING"): + body = await ( + await self.client.request("POST", "/prompt_manager/scan") + ).text() + + self.assertIn('"processed": 1', body) + self.assertIn('"found": 0', body) + + async def test_unreadable_image_is_a_warning_not_an_error(self): + (self.output_dir / "bad.png").write_bytes(b"\x89PNG not really") + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + + with self.assertLogs(level="WARNING") as cm: + body = await ( + await self.client.request("POST", "/prompt_manager/scan") + ).text() + + self.assertIn('"type": "complete"', body) + self.assertTrue(any("bad.png" in line for line in cm.output)) + self.assertEqual([line for line in cm.output if line.startswith("ERROR")], []) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 5c90f17..fc13719 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -122,6 +122,7 @@ class AutotagAPITestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client self._service_patch.stop() self.api.db.close_all() for path in ( diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index 02c904b..326537a 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -46,6 +46,7 @@ class BackupRestoreTestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client self.api.db.close_all() for name in os.listdir(self.tmpdir): os.unlink(os.path.join(self.tmpdir, name)) diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 2d534e7..6d49b8d 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -372,6 +372,7 @@ class CoreRoutesTestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client self.api.db.close_all() shutil.rmtree(self.tmpdir, ignore_errors=True) diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 0e83a69..99b6b28 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -72,6 +72,7 @@ class ImageAPITestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client self._anchor_patch.stop() # Windows refuses to unlink a database that still has open handles. self.api.db.close_all() diff --git a/tests/test_api_logging.py b/tests/test_api_logging.py index cda66ed..1614909 100644 --- a/tests/test_api_logging.py +++ b/tests/test_api_logging.py @@ -112,6 +112,7 @@ class LoggingAPITestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client if self._saved_folder_paths is None: sys.modules.pop("folder_paths", None) else: diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index be65e2a..fe7e226 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -83,6 +83,7 @@ class LoraAPITestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client self._save_patch.stop() self._detect_patch.stop() self._config_patch.stop() diff --git a/tests/test_api_prompts.py b/tests/test_api_prompts.py index d27beff..93307a7 100644 --- a/tests/test_api_prompts.py +++ b/tests/test_api_prompts.py @@ -45,6 +45,7 @@ class PromptAPITestCase(AioHTTPTestCase): return app async def tearDownAsync(self): + await super().tearDownAsync() # closes the aiohttp test client # Windows refuses to unlink a database that still has open handles. self.api.db.close_all() for path in ( diff --git a/tests/test_config.py b/tests/test_config.py index 83af136..7daa152 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -136,13 +136,15 @@ class TestPromptManagerConfig(unittest.TestCase): PromptManagerConfig.update_config({"performance": {"max_search_results": 50}}) self.assertEqual(PromptManagerConfig.MAX_SEARCH_RESULTS, 50) - def test_worker_threads_default_is_within_the_machine(self): + def test_worker_threads_default_is_half_the_cores_and_current_is_in_range(self): + # The current value may come from config.json (loaded at import), so the + # rule is checked on the constant and only the range on the live value. cores = PromptManagerConfig.max_worker_threads() self.assertGreaterEqual(cores, 1) - default = PromptManagerConfig.get_config()["performance"]["worker_threads"] - self.assertGreaterEqual(default, 1) - self.assertLessEqual(default, cores) - self.assertEqual(default, max(1, cores // 2)) + self.assertEqual(PromptManagerConfig.DEFAULT_WORKER_THREADS, max(1, cores // 2)) + current = PromptManagerConfig.get_config()["performance"]["worker_threads"] + self.assertGreaterEqual(current, 1) + self.assertLessEqual(current, cores) def test_update_worker_threads_is_clamped_to_the_machine(self): cores = PromptManagerConfig.max_worker_threads() diff --git a/tests/test_utils_video_metadata.py b/tests/test_utils_video_metadata.py new file mode 100644 index 0000000..1d56267 --- /dev/null +++ b/tests/test_utils_video_metadata.py @@ -0,0 +1,165 @@ +"""Tests for utils/video_metadata.py: prompt metadata from video containers. + +ComfyUI's video savers store ``{"prompt": ..., "workflow": ...}`` as JSON in +the container's ``comment`` tag. ffprobe reads it; when ffprobe is missing +videos are skipped without noise. +""" + +import json +import os +import shutil +import subprocess +import sys +import tempfile +import unittest +from unittest import mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils import video_metadata # noqa: E402 +from utils.video_metadata import ( # noqa: E402 + VIDEO_EXTENSIONS, + is_video_path, + metadata_from_tags, + read_video_metadata, +) + +HAVE_FFMPEG = bool(shutil.which("ffmpeg") and shutil.which("ffprobe")) +GRAPH = json.dumps({"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}) + + +def write_mp4(path, comment): + """Write a tiny black video carrying ``comment`` in its container tags.""" + subprocess.run( + [ + "ffmpeg", + "-v", + "quiet", + "-f", + "lavfi", + "-i", + "color=c=black:s=16x16:d=0.2", + "-metadata", + f"comment={comment}", + "-y", + str(path), + ], + check=True, + timeout=60, + ) + + +class TestMetadataFromTags(unittest.TestCase): + + def test_prompt_string_and_workflow_dict_become_text_chunks(self): + tags = {"comment": json.dumps({"prompt": GRAPH, "workflow": {"nodes": []}})} + meta = metadata_from_tags(tags) + self.assertEqual(meta["prompt"], GRAPH) + self.assertEqual(json.loads(meta["workflow"]), {"nodes": []}) + + def test_missing_or_non_json_comment_is_empty(self): + self.assertEqual(metadata_from_tags({}), {}) + self.assertEqual(metadata_from_tags({"comment": "just a note"}), {}) + self.assertEqual(metadata_from_tags({"comment": "[1, 2]"}), {}) + self.assertEqual(metadata_from_tags({"comment": json.dumps({"other": 1})}), {}) + self.assertEqual(metadata_from_tags({"comment": 7}), {}) + + def test_description_tag_is_a_fallback(self): + tags = {"description": json.dumps({"prompt": GRAPH})} + self.assertEqual(metadata_from_tags(tags), {"prompt": GRAPH}) + + +class TestIsVideoPath(unittest.TestCase): + + def test_extensions(self): + for ext in VIDEO_EXTENSIONS: + self.assertTrue(is_video_path(f"clip{ext.upper()}"), ext) + self.assertFalse(is_video_path("image.png")) + self.assertFalse(is_video_path("noext")) + + +class TestReadVideoMetadataWithoutFfprobe(unittest.TestCase): + + def setUp(self): + video_metadata.ffprobe_path.cache_clear() + self.addCleanup(video_metadata.ffprobe_path.cache_clear) + + def test_missing_ffprobe_returns_empty_without_spawning(self): + with mock.patch.object(video_metadata.shutil, "which", return_value=None): + with mock.patch.object(video_metadata.subprocess, "run") as run: + self.assertEqual(read_video_metadata("clip.mp4"), {}) + run.assert_not_called() + + def test_missing_ffprobe_is_mentioned_once(self): + video_metadata.reset_missing_notice() + with mock.patch.object(video_metadata.shutil, "which", return_value=None): + with self.assertLogs("prompt_manager.video", level="INFO") as cm: + read_video_metadata("a.mp4") + read_video_metadata("b.mp4") + self.assertEqual(len(cm.output), 1) + self.assertIn("ffprobe", cm.output[0]) + + def _run_returning(self, **attrs): + return mock.patch.object( + video_metadata.subprocess, + "run", + return_value=mock.Mock(**{"returncode": 0, "stdout": "{}", **attrs}), + ) + + def test_ffprobe_failure_timeout_and_garbage_are_empty(self): + with mock.patch.object( + video_metadata.shutil, "which", return_value="/bin/ffprobe" + ): + with self._run_returning(returncode=1): + self.assertEqual(read_video_metadata("clip.mp4"), {}) + with self._run_returning(stdout="not json"): + self.assertEqual(read_video_metadata("clip.mp4"), {}) + with mock.patch.object( + video_metadata.subprocess, + "run", + side_effect=subprocess.TimeoutExpired("ffprobe", 1), + ): + self.assertEqual(read_video_metadata("clip.mp4"), {}) + with mock.patch.object( + video_metadata.subprocess, "run", side_effect=OSError + ): + self.assertEqual(read_video_metadata("clip.mp4"), {}) + + def test_ffprobe_output_is_parsed(self): + payload = {"format": {"tags": {"comment": json.dumps({"prompt": GRAPH})}}} + with mock.patch.object( + video_metadata.shutil, "which", return_value="/bin/ffprobe" + ): + with self._run_returning(stdout=json.dumps(payload)) as run: + self.assertEqual(read_video_metadata("clip.mp4"), {"prompt": GRAPH}) + cmd = run.call_args.args[0] + self.assertEqual(cmd[0], "/bin/ffprobe") + self.assertEqual(cmd[-1], "clip.mp4") + self.assertIn("format_tags", cmd) + + +@unittest.skipUnless(HAVE_FFMPEG, "ffmpeg/ffprobe not installed") +class TestReadVideoMetadataWithRealFfprobe(unittest.TestCase): + + def setUp(self): + video_metadata.ffprobe_path.cache_clear() + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir, True) + + def test_round_trip_through_a_real_mp4(self): + path = os.path.join(self.tmpdir, "clip.mp4") + write_mp4(path, json.dumps({"prompt": GRAPH, "workflow": {"v": 1}})) + + meta = read_video_metadata(path) + + self.assertEqual(meta["prompt"], GRAPH) + self.assertEqual(json.loads(meta["workflow"]), {"v": 1}) + + def test_video_without_a_comment_is_empty(self): + path = os.path.join(self.tmpdir, "plain.mp4") + write_mp4(path, "") + self.assertEqual(read_video_metadata(path), {}) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/video_metadata.py b/utils/video_metadata.py new file mode 100644 index 0000000..a7884dc --- /dev/null +++ b/utils/video_metadata.py @@ -0,0 +1,123 @@ +"""Prompt metadata from video containers. + +ComfyUI's video savers (core SaveVideo, VideoHelperSuite, ...) write the +queued graph into the container's ``comment`` tag as JSON: +``{"prompt": "", "workflow": {...}}``. ffprobe reads it on every +platform; when ffprobe is not installed, videos are skipped with a single +notice instead of one error per file. + +The result uses the same shape as PNG text chunks (string values keyed +``prompt`` / ``workflow`` / ``parameters``) so the existing prompt parser +handles both. +""" + +import functools +import json +import logging +import os +import shutil + +# Fixed argv, absolute binary, no shell. +import subprocess # nosec B404 + +VIDEO_EXTENSIONS = (".mp4", ".mov", ".m4v", ".webm", ".mkv", ".avi", ".wmv") +FFPROBE_TIMEOUT = 20 # seconds; a stuck probe must not hang a worker thread +_COMMENT_TAGS = ("comment", "description") +_METADATA_KEYS = ("prompt", "workflow", "parameters") +_NO_WINDOW = getattr(subprocess, "CREATE_NO_WINDOW", 0) # no console flash on Windows + +logger = logging.getLogger("prompt_manager.video") +_missing_notice_shown = False + + +def is_video_path(path) -> bool: + return os.path.splitext(str(path))[1].lower() in VIDEO_EXTENSIONS + + +@functools.lru_cache(maxsize=1) +def ffprobe_path(): + """Absolute path of ffprobe, or None. Cached for the life of the process.""" + return shutil.which("ffprobe") + + +def reset_missing_notice(): + """Allow the 'ffprobe not found' notice to be logged again (tests).""" + global _missing_notice_shown + _missing_notice_shown = False + + +def metadata_from_tags(tags) -> dict: + """Normalise container tags into PNG-text-chunk shaped metadata. + + Returns ``{}`` when no tag carries a JSON object with a known key. + """ + if not isinstance(tags, dict): + return {} + for tag in _COMMENT_TAGS: + raw = tags.get(tag) + if not isinstance(raw, str): + continue + try: + data = json.loads(raw) + except ValueError: + continue + if not isinstance(data, dict): + continue + found = {} + for key in _METADATA_KEYS: + value = data.get(key) + if value is None: + continue + found[key] = value if isinstance(value, str) else json.dumps(value) + if found: + return found + return {} + + +def read_video_metadata(path, timeout=FFPROBE_TIMEOUT) -> dict: + """Metadata embedded in the video at ``path``; ``{}`` when none or no ffprobe.""" + global _missing_notice_shown + probe = ffprobe_path() + if not probe: + if not _missing_notice_shown: + _missing_notice_shown = True + logger.info( + "ffprobe not found; prompts embedded in videos are skipped " + "during scans. Install ffmpeg to include them." + ) + return {} + + cmd = [ + probe, + "-v", + "quiet", + "-show_entries", + "format_tags", + "-of", + "json", + str(path), + ] + try: + result = subprocess.run( # nosec B603 + cmd, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + timeout=timeout, + creationflags=_NO_WINDOW, + ) + except (OSError, subprocess.TimeoutExpired, ValueError) as e: + logger.debug(f"ffprobe failed for {os.path.basename(str(path))}: {e}") + return {} + if result.returncode != 0: + return {} + try: + probed = json.loads(result.stdout or "{}") + except ValueError: + return {} + if not isinstance(probed, dict): + return {} + fmt = probed.get("format") + tags = fmt.get("tags") if isinstance(fmt, dict) else None + return metadata_from_tags(tags)