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)