feat: PromptSearchList batch workflow support & image linking v3.2.2 (#139)
* fix(search): PromptSearchList crash, filters, and batch output (#52) - Fix empty list crash (OUTPUT_IS_LIST requires at least one element) - Add partial tag matching (LIKE instead of exact) - Add preview and count outputs for batch visibility - Add skip_multipart filter for Clip_N video prompts - Add LoRA-only prompt filter - Collapse newlines for StringOutputList compatibility * feat(db): add tag_partial parameter for LIKE tag matching (#52) Allows search_prompts to use LIKE matching for tags instead of exact. Used by PromptSearchList for discovery-style search. Existing callers default to exact matching (unchanged behavior). * feat(tracker): add FIFO prompt queue for batch image linking (#52) Batch workflows encode all prompts first, then save all images later. The queue preserves encode order so each image links to the correct prompt by position (pop on save matches push on encode). * fix(monitor): multi-strategy image-to-prompt linking (#52) Previous approach relied on live prompt tracker state, which was always stale for batch workflows (all images linked to last prompt). New 5-strategy priority chain: 1. FIFO queue pop (batch-correct by position) 2. Metadata-based DB lookup (extracts prompt from PNG) 3. Prompt snapshot at file-creation time 4. Live prompt tracker 5. Most recent DB prompt fallback * test: add tests for prompt queue and search filters (#52) - 6 queue tests: FIFO ordering, batch simulation, exhaustion - 15 filter tests: newline collapse, Clip_ filter, LoRA-only filter - Fix _make_tracker helper to include queue fields * chore: bump version to 3.2.2 * fix: address code review — regex backtracking and queue expiry (#52) - Replace backtracking-prone LoRA-only regex with subtraction approach (re.sub removes tags, check if anything remains). Fixes CodeQL alert. - Add expiry check to pop_next_prompt — skips stale entries to prevent incorrect linking from accumulated queue entries across sessions. - Cleanup thread now also prunes expired queue entries.
This commit is contained in:
Vito
+16
-6
@@ -190,6 +190,7 @@ class PromptDatabase:
|
||||
folder: Optional[str] = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
tag_partial: bool = False,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search prompts with various filters.
|
||||
@@ -205,6 +206,7 @@ class PromptDatabase:
|
||||
folder: Filter by subfolder name in generated image paths
|
||||
limit: Maximum number of results
|
||||
offset: Number of results to skip
|
||||
tag_partial: Use LIKE matching for tags instead of exact match
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing prompt data
|
||||
@@ -222,12 +224,20 @@ class PromptDatabase:
|
||||
|
||||
if tags:
|
||||
for tag in tags:
|
||||
query_parts.append(
|
||||
"AND prompts.id IN ("
|
||||
" SELECT pt.prompt_id FROM prompt_tags pt"
|
||||
" JOIN tags t ON pt.tag_id = t.id WHERE t.name = ?)"
|
||||
)
|
||||
params.append(tag)
|
||||
if tag_partial:
|
||||
query_parts.append(
|
||||
"AND prompts.id IN ("
|
||||
" SELECT pt.prompt_id FROM prompt_tags pt"
|
||||
" JOIN tags t ON pt.tag_id = t.id WHERE t.name LIKE ?)"
|
||||
)
|
||||
params.append(f"%{tag}%")
|
||||
else:
|
||||
query_parts.append(
|
||||
"AND prompts.id IN ("
|
||||
" SELECT pt.prompt_id FROM prompt_tags pt"
|
||||
" JOIN tags t ON pt.tag_id = t.id WHERE t.name = ?)"
|
||||
)
|
||||
params.append(tag)
|
||||
|
||||
if folder:
|
||||
# Escape LIKE wildcards in the folder name
|
||||
|
||||
+68
-22
@@ -5,6 +5,7 @@ This node enables batch processing workflows by outputting prompt texts with
|
||||
OUTPUT_IS_LIST=True, allowing direct connection to nodes that accept list inputs.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
try:
|
||||
@@ -79,7 +80,7 @@ class PromptSearchList(ComfyNodeABC):
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Comma-separated list of tags to filter by",
|
||||
"tooltip": "Comma-separated list of tags to filter by (partial match)",
|
||||
},
|
||||
),
|
||||
"min_rating": (
|
||||
@@ -100,20 +101,33 @@ class PromptSearchList(ComfyNodeABC):
|
||||
"tooltip": "Maximum number of results to return",
|
||||
},
|
||||
),
|
||||
"skip_multipart": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Skip prompts containing Clip_1/Clip_2/etc. multi-part markers",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.STRING,)
|
||||
RETURN_NAMES = ("prompts",)
|
||||
OUTPUT_IS_LIST = (True,)
|
||||
OUTPUT_TOOLTIPS = ("List of prompt texts matching the search criteria.",)
|
||||
RETURN_TYPES = (IO.STRING, IO.STRING, IO.INT)
|
||||
RETURN_NAMES = ("prompts", "preview", "count")
|
||||
OUTPUT_IS_LIST = (True, False, False)
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"List of prompt texts matching the search criteria.",
|
||||
"Preview of all found prompts (for display).",
|
||||
"Number of prompts found.",
|
||||
)
|
||||
FUNCTION = "search"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
|
||||
DESCRIPTION = (
|
||||
"Searches the prompt database and outputs matching prompts as a list. "
|
||||
"Use this node to retrieve stored prompts for batch processing workflows. "
|
||||
"Connect to nodes that accept list inputs like String OutputList or batch processors."
|
||||
"'prompts' output sends each prompt individually to downstream nodes (batch). "
|
||||
"'preview' output shows all found prompts as one text block. "
|
||||
"'count' output shows how many prompts were found."
|
||||
)
|
||||
|
||||
def search(
|
||||
@@ -123,7 +137,8 @@ class PromptSearchList(ComfyNodeABC):
|
||||
tags: str = "",
|
||||
min_rating: int = 0,
|
||||
limit: int = 50,
|
||||
) -> Tuple[List[str]]:
|
||||
skip_multipart: bool = True,
|
||||
):
|
||||
"""
|
||||
Search for prompts matching the given criteria.
|
||||
|
||||
@@ -135,7 +150,7 @@ class PromptSearchList(ComfyNodeABC):
|
||||
limit: Maximum number of results to return
|
||||
|
||||
Returns:
|
||||
Tuple containing a list of prompt text strings
|
||||
Dict with 'ui' display info and 'result' tuple
|
||||
"""
|
||||
try:
|
||||
# Parse tags from comma-separated string
|
||||
@@ -143,38 +158,69 @@ class PromptSearchList(ComfyNodeABC):
|
||||
if tags and tags.strip():
|
||||
tags_list = [tag.strip() for tag in tags.split(",") if tag.strip()]
|
||||
|
||||
# Perform search
|
||||
# Perform search with partial tag matching for discovery
|
||||
results = self.db.search_prompts(
|
||||
text=text.strip() if text and text.strip() else None,
|
||||
category=category.strip() if category and category.strip() else None,
|
||||
tags=tags_list,
|
||||
tag_partial=True,
|
||||
rating_min=min_rating if min_rating > 0 else None,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
# Extract just the prompt text from results
|
||||
prompt_texts = [r["text"] for r in results if r.get("text")]
|
||||
# Extract prompt text, collapsing newlines to spaces so downstream
|
||||
# nodes that split on \n (e.g. StringOutputList) treat each DB
|
||||
# entry as a single prompt.
|
||||
prompt_texts = [
|
||||
" ".join(r["text"].split()) for r in results if r.get("text")
|
||||
]
|
||||
|
||||
self.logger.debug(
|
||||
f"Search returned {len(prompt_texts)} prompts "
|
||||
f"(text='{text[:20]}...' if text else '', category='{category}', "
|
||||
f"tags={tags_list}, min_rating={min_rating}, limit={limit})"
|
||||
)
|
||||
# Filter out multi-part prompts (Clip_1/Clip_2 markers)
|
||||
if skip_multipart:
|
||||
prompt_texts = [
|
||||
p for p in prompt_texts if not re.search(r"Clip_\d+", p)
|
||||
]
|
||||
|
||||
# Return empty list if no results (not an error)
|
||||
# Filter out prompts that are only LoRA tags with no actual content.
|
||||
# Uses subtraction instead of repeated groups to avoid backtracking.
|
||||
prompt_texts = [
|
||||
p for p in prompt_texts if re.sub(r"<lora:[^>]+>", "", p).strip()
|
||||
]
|
||||
|
||||
count = len(prompt_texts)
|
||||
|
||||
# Build a preview: numbered list of truncated prompts
|
||||
preview_lines = []
|
||||
for i, p in enumerate(prompt_texts, 1):
|
||||
truncated = p[:120] + "..." if len(p) > 120 else p
|
||||
preview_lines.append(f"[{i}] {truncated}")
|
||||
preview = "\n".join(preview_lines) if preview_lines else "No results found"
|
||||
|
||||
# Must return at least one element — ComfyUI's slice_dict
|
||||
# indexes into OUTPUT_IS_LIST outputs and crashes on empty lists.
|
||||
if not prompt_texts:
|
||||
self.logger.info("Search returned no results")
|
||||
return ([],)
|
||||
return {
|
||||
"ui": {"text": ["No results found"]},
|
||||
"result": ([""], preview, count),
|
||||
}
|
||||
|
||||
return (prompt_texts,)
|
||||
return {
|
||||
"ui": {"text": [f"Found {count} prompts"]},
|
||||
"result": (prompt_texts, preview, count),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Search error: {e}", exc_info=True)
|
||||
# Return empty list on error to avoid breaking workflows
|
||||
return ([],)
|
||||
return {
|
||||
"ui": {"text": [f"Search error: {e}"]},
|
||||
"result": ([""], f"Error: {e}", 0),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, text="", category="", tags="", min_rating=0, limit=50):
|
||||
def IS_CHANGED(
|
||||
cls, text="", category="", tags="", min_rating=0, limit=50, skip_multipart=True
|
||||
):
|
||||
"""
|
||||
Always re-execute to get fresh results from the database.
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "promptmanager"
|
||||
description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite."
|
||||
version = "3.2.1"
|
||||
version = "3.2.2"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["# Core dependencies for PromptManager", "# Note: Most dependencies are already included with ComfyUI", "# Already included with Python standard library:", "# - sqlite3", "# - hashlib", "# - json", "# - datetime", "# - os", "# - typing", "# - threading", "# - uuid", "# Required for gallery functionality:", "watchdog>=2.1.0 # For file system monitoring", "Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)", "# Optional dependencies for enhanced search functionality:", "# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)", "# sqlalchemy>=1.4.0 # For advanced ORM features (optional)", "# Development dependencies (optional):", "# pytest>=6.0.0 # For running tests", "# black>=22.0.0 # For code formatting", "# flake8>=4.0.0 # For linting", "# mypy>=0.910 # For type checking"]
|
||||
|
||||
|
||||
@@ -34,6 +34,8 @@ def _make_tracker(prompt_timeout=600, cleanup_interval=300):
|
||||
tracker._local = threading.local()
|
||||
tracker.active_prompts = {}
|
||||
tracker.lock = threading.Lock()
|
||||
tracker._prompt_queue = []
|
||||
tracker._queue_lock = threading.Lock()
|
||||
tracker.cleanup_interval = cleanup_interval
|
||||
tracker.prompt_timeout = prompt_timeout
|
||||
# Don't start cleanup thread in tests
|
||||
@@ -256,5 +258,59 @@ class TestCleanupExpiredPrompts(unittest.TestCase):
|
||||
self.assertEqual(len(expired_ids), 1)
|
||||
|
||||
|
||||
class TestPromptQueue(unittest.TestCase):
|
||||
"""Test the FIFO prompt queue for batch workflow image linking."""
|
||||
|
||||
def setUp(self):
|
||||
self.tracker = _make_tracker()
|
||||
|
||||
def test_queue_empty_initially(self):
|
||||
result = self.tracker.pop_next_prompt()
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_set_prompt_pushes_to_queue(self):
|
||||
self.tracker.set_current_prompt("p1", {"prompt_id": 10})
|
||||
self.assertEqual(len(self.tracker._prompt_queue), 1)
|
||||
|
||||
def test_pop_returns_fifo_order(self):
|
||||
self.tracker.set_current_prompt("first", {"prompt_id": 1})
|
||||
self.tracker.set_current_prompt("second", {"prompt_id": 2})
|
||||
self.tracker.set_current_prompt("third", {"prompt_id": 3})
|
||||
|
||||
p1 = self.tracker.pop_next_prompt()
|
||||
p2 = self.tracker.pop_next_prompt()
|
||||
p3 = self.tracker.pop_next_prompt()
|
||||
|
||||
self.assertEqual(p1["id"], 1)
|
||||
self.assertEqual(p2["id"], 2)
|
||||
self.assertEqual(p3["id"], 3)
|
||||
|
||||
def test_pop_returns_none_when_exhausted(self):
|
||||
self.tracker.set_current_prompt("only", {"prompt_id": 1})
|
||||
self.tracker.pop_next_prompt()
|
||||
result = self.tracker.pop_next_prompt()
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_queue_preserves_prompt_text(self):
|
||||
self.tracker.set_current_prompt("hello world", {"prompt_id": 5})
|
||||
ctx = self.tracker.pop_next_prompt()
|
||||
self.assertEqual(ctx["text"], "hello world")
|
||||
|
||||
def test_batch_simulation(self):
|
||||
"""Simulate a batch workflow: encode 10 prompts, then pop 10 in order."""
|
||||
ids = list(range(100, 110))
|
||||
for pid in ids:
|
||||
self.tracker.set_current_prompt(f"prompt_{pid}", {"prompt_id": pid})
|
||||
|
||||
popped_ids = []
|
||||
while True:
|
||||
ctx = self.tracker.pop_next_prompt()
|
||||
if ctx is None:
|
||||
break
|
||||
popped_ids.append(ctx["id"])
|
||||
|
||||
self.assertEqual(popped_ids, ids)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
"""
|
||||
Tests for PromptSearchList output filtering.
|
||||
|
||||
Tests the newline collapsing, multi-part (Clip_) filtering,
|
||||
and LoRA-only filtering applied to search results.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
|
||||
# --- Helpers: replicate the filter logic from prompt_search_list.py ---
|
||||
|
||||
|
||||
def apply_filters(texts, skip_multipart=True):
|
||||
"""Apply the same filtering chain as PromptSearchList.search()."""
|
||||
# Collapse newlines
|
||||
result = [" ".join(t.split()) for t in texts if t]
|
||||
|
||||
# Filter multi-part
|
||||
if skip_multipart:
|
||||
result = [p for p in result if not re.search(r"Clip_\d+", p)]
|
||||
|
||||
# Filter LoRA-only (subtraction approach — no backtracking risk)
|
||||
result = [p for p in result if re.sub(r"<lora:[^>]+>", "", p).strip()]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class TestNewlineCollapse(unittest.TestCase):
|
||||
"""Newlines must be collapsed so StringOutputList treats each DB entry as one prompt."""
|
||||
|
||||
def test_single_line_unchanged(self):
|
||||
result = apply_filters(["a beautiful landscape"])
|
||||
self.assertEqual(result, ["a beautiful landscape"])
|
||||
|
||||
def test_multiline_collapsed_to_spaces(self):
|
||||
result = apply_filters(["line one\nline two\nline three"])
|
||||
self.assertEqual(result, ["line one line two line three"])
|
||||
|
||||
def test_tabs_and_extra_spaces_collapsed(self):
|
||||
result = apply_filters(["word1 \t word2\n\nword3"])
|
||||
self.assertEqual(result, ["word1 word2 word3"])
|
||||
|
||||
def test_leading_trailing_whitespace_stripped(self):
|
||||
result = apply_filters([" padded prompt \n"])
|
||||
self.assertEqual(result, ["padded prompt"])
|
||||
|
||||
|
||||
class TestMultipartFilter(unittest.TestCase):
|
||||
"""Prompts containing Clip_N markers should be filtered when skip_multipart=True."""
|
||||
|
||||
def test_clip_markers_filtered(self):
|
||||
texts = [
|
||||
"Clip_1 a video prompt Clip_2 continuation",
|
||||
"a normal prompt",
|
||||
]
|
||||
result = apply_filters(texts, skip_multipart=True)
|
||||
self.assertEqual(result, ["a normal prompt"])
|
||||
|
||||
def test_clip_markers_kept_when_disabled(self):
|
||||
texts = ["Clip_1 a video prompt Clip_2 continuation"]
|
||||
result = apply_filters(texts, skip_multipart=False)
|
||||
self.assertEqual(len(result), 1)
|
||||
|
||||
def test_clip_in_word_filtered(self):
|
||||
# "Clip_3" anywhere in text should match
|
||||
texts = ["some text with Clip_3 marker"]
|
||||
result = apply_filters(texts, skip_multipart=True)
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_clip_without_number_not_filtered(self):
|
||||
texts = ["clip art style painting"]
|
||||
result = apply_filters(texts, skip_multipart=True)
|
||||
self.assertEqual(result, ["clip art style painting"])
|
||||
|
||||
|
||||
class TestLoraOnlyFilter(unittest.TestCase):
|
||||
"""Prompts that are only LoRA tags with no actual content should be filtered."""
|
||||
|
||||
def test_single_lora_tag_filtered(self):
|
||||
result = apply_filters(["<lora:some_model:0.7>"])
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_multiple_lora_tags_filtered(self):
|
||||
result = apply_filters(["<lora:model_a:1> <lora:model_b:0.5>"])
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_lora_with_content_kept(self):
|
||||
result = apply_filters(["<lora:style:0.8> a beautiful painting"])
|
||||
self.assertEqual(result, ["<lora:style:0.8> a beautiful painting"])
|
||||
|
||||
def test_content_with_lora_in_middle_kept(self):
|
||||
result = apply_filters(["masterpiece, <lora:fix:1>, best quality"])
|
||||
self.assertEqual(result, ["masterpiece, <lora:fix:1>, best quality"])
|
||||
|
||||
def test_lora_with_newlines_collapsed_then_filtered(self):
|
||||
# After collapsing newlines, this becomes a single line of LoRA tags
|
||||
result = apply_filters(["<lora:a:1>\n<lora:b:1>"])
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_empty_string_filtered(self):
|
||||
result = apply_filters([""])
|
||||
self.assertEqual(result, [])
|
||||
|
||||
|
||||
class TestFilterCombination(unittest.TestCase):
|
||||
"""Test all filters working together on a mixed batch."""
|
||||
|
||||
def test_mixed_batch(self):
|
||||
texts = [
|
||||
"a good prompt",
|
||||
"<lora:style:1>",
|
||||
"Clip_1\nfirst segment\nClip_2\nsecond segment",
|
||||
"<lora:model:0.5> nice painting, flowers",
|
||||
"another\ngood\nprompt",
|
||||
"<lora:a:1> <lora:b:1>",
|
||||
]
|
||||
result = apply_filters(texts)
|
||||
self.assertEqual(
|
||||
result,
|
||||
[
|
||||
"a good prompt",
|
||||
"<lora:model:0.5> nice painting, flowers",
|
||||
"another good prompt",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+151
-32
@@ -80,14 +80,26 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
directory. It filters for image files and schedules them for processing after
|
||||
a small delay to ensure the file is fully written.
|
||||
|
||||
The prompt context is captured immediately (before the delay) to handle
|
||||
batch workflows where the prompt tracker advances before images are processed.
|
||||
|
||||
Args:
|
||||
event: FileSystemEvent object containing event details
|
||||
"""
|
||||
if not event.is_directory and self.is_image_file(event.src_path):
|
||||
self.logger.info(f"New image detected: {event.src_path}")
|
||||
# Small delay to ensure file is fully written
|
||||
# Snapshot prompt context NOW before the delay — in batch workflows
|
||||
# the tracker advances to the next prompt before images are processed.
|
||||
prompt_snapshot = self.prompt_tracker.get_current_prompt()
|
||||
if prompt_snapshot:
|
||||
self.logger.debug(
|
||||
f"Snapshot prompt {prompt_snapshot.get('id', '?')} for {os.path.basename(event.src_path)}"
|
||||
)
|
||||
threading.Timer(
|
||||
self.processing_delay, self.process_new_image, args=[event.src_path]
|
||||
self.processing_delay,
|
||||
self.process_new_image,
|
||||
args=[event.src_path],
|
||||
kwargs={"prompt_snapshot": prompt_snapshot},
|
||||
).start()
|
||||
|
||||
def is_image_file(self, filepath: str) -> bool:
|
||||
@@ -105,18 +117,19 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
return False
|
||||
return filepath.lower().endswith(self.supported_extensions)
|
||||
|
||||
def process_new_image(self, image_path: str):
|
||||
def process_new_image(self, image_path: str, prompt_snapshot=None):
|
||||
"""Process a newly created image file for gallery integration.
|
||||
|
||||
This method handles the complete processing pipeline for a new image:
|
||||
1. Verifies the file still exists
|
||||
2. Gets the current prompt context from the tracker
|
||||
3. Extracts ComfyUI metadata from the image
|
||||
4. Links the image to the appropriate prompt in the database
|
||||
5. Handles fallback scenarios when no active prompt is available
|
||||
2. Extracts ComfyUI metadata from the image
|
||||
3. Tries to identify the correct prompt via metadata-based lookup
|
||||
4. Falls back to the prompt snapshot captured at event time
|
||||
5. Links the image to the appropriate prompt in the database
|
||||
|
||||
Args:
|
||||
image_path: Full path to the newly created image file
|
||||
prompt_snapshot: Prompt context captured at file-creation time (optional)
|
||||
"""
|
||||
try:
|
||||
self.logger.info(f"Processing image: {image_path}")
|
||||
@@ -125,35 +138,61 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
self.logger.warning(f"Image file no longer exists: {image_path}")
|
||||
return
|
||||
|
||||
# Get current prompt context first
|
||||
current_prompt = self.prompt_tracker.get_current_prompt()
|
||||
self.logger.info(
|
||||
f"Current prompt context: {current_prompt['id'] if current_prompt else 'None'}"
|
||||
)
|
||||
|
||||
if not current_prompt:
|
||||
self.logger.debug(f"No active prompt context for image: {image_path}")
|
||||
# Fallback: try to link to the most recent prompt in database
|
||||
current_prompt = self._get_fallback_prompt()
|
||||
if current_prompt:
|
||||
self.logger.debug(f"Using fallback prompt: {current_prompt['id']}")
|
||||
else:
|
||||
self.logger.warning(f"No fallback prompt available, skipping image")
|
||||
return
|
||||
else:
|
||||
# Extend the timeout for this prompt since we're still getting images
|
||||
if "execution_id" in current_prompt:
|
||||
self.prompt_tracker.extend_prompt_timeout(
|
||||
current_prompt["execution_id"], 300
|
||||
) # Add 5 more minutes
|
||||
|
||||
# Extract ComfyUI metadata
|
||||
# Extract ComfyUI metadata first — needed both for linking and prompt lookup
|
||||
metadata = None
|
||||
try:
|
||||
metadata = self.metadata_extractor.extract_metadata(image_path)
|
||||
self.logger.debug(f"Extracted metadata: {bool(metadata)}")
|
||||
except Exception as meta_error:
|
||||
self.logger.warning(f"Metadata extraction failed: {meta_error}")
|
||||
metadata = None
|
||||
|
||||
# Strategy 1: Pop from batch queue (most reliable for batch workflows).
|
||||
# Prompts are queued during CLIP encoding in order; images save in
|
||||
# the same order, so FIFO pop gives the correct prompt per image.
|
||||
current_prompt = self.prompt_tracker.pop_next_prompt()
|
||||
if current_prompt:
|
||||
self.logger.info(
|
||||
f"Queue match: prompt {current_prompt['id']} for "
|
||||
f"{os.path.basename(image_path)}"
|
||||
)
|
||||
|
||||
# Strategy 2: Find prompt from image metadata
|
||||
if not current_prompt:
|
||||
current_prompt = self._find_prompt_from_metadata(metadata)
|
||||
if current_prompt:
|
||||
self.logger.info(f"Metadata match: prompt {current_prompt['id']}")
|
||||
|
||||
# Strategy 3: Use snapshot captured at file-creation time
|
||||
if not current_prompt and prompt_snapshot:
|
||||
current_prompt = prompt_snapshot
|
||||
self.logger.info(
|
||||
f"Snapshot match: prompt {current_prompt.get('id', 'unknown')}"
|
||||
)
|
||||
|
||||
# Strategy 4: Check live prompt tracker
|
||||
if not current_prompt:
|
||||
current_prompt = self.prompt_tracker.get_current_prompt()
|
||||
if current_prompt:
|
||||
self.logger.info(
|
||||
f"Live tracker match: prompt {current_prompt['id']}"
|
||||
)
|
||||
|
||||
# Strategy 5: Fallback to most recent prompt in DB
|
||||
if not current_prompt:
|
||||
current_prompt = self._get_fallback_prompt()
|
||||
if current_prompt:
|
||||
self.logger.debug(f"Fallback match: prompt {current_prompt['id']}")
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"No prompt context available, skipping image: {image_path}"
|
||||
)
|
||||
return
|
||||
|
||||
# Extend timeout if we have an active execution
|
||||
if current_prompt.get("execution_id"):
|
||||
self.prompt_tracker.extend_prompt_timeout(
|
||||
current_prompt["execution_id"], 300
|
||||
)
|
||||
|
||||
if metadata:
|
||||
self.logger.debug(
|
||||
@@ -164,7 +203,6 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
self.logger.debug(
|
||||
f"Linking image with basic info to prompt {current_prompt['id']}"
|
||||
)
|
||||
# Link with basic file info even without metadata
|
||||
basic_metadata = self.get_basic_file_info(image_path)
|
||||
self.link_image_to_prompt(
|
||||
image_path, current_prompt, {"file_info": basic_metadata}
|
||||
@@ -176,6 +214,87 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
|
||||
self.logger.error(traceback.format_exc())
|
||||
|
||||
def _find_prompt_from_metadata(self, metadata):
|
||||
"""Extract prompt text from image metadata and look up the matching DB prompt.
|
||||
|
||||
Parses the ComfyUI workflow/prompt data embedded in the image to find
|
||||
PromptManager node inputs, then matches against the database by hash.
|
||||
|
||||
Args:
|
||||
metadata: Extracted metadata dict from the image, or None
|
||||
|
||||
Returns:
|
||||
Prompt context dict with 'id' and 'text', or None if not found
|
||||
"""
|
||||
if not metadata:
|
||||
return None
|
||||
|
||||
prompt_text = None
|
||||
|
||||
# Try to find PromptManager node in the prompt execution data
|
||||
prompt_data = metadata.get("prompt")
|
||||
if isinstance(prompt_data, dict):
|
||||
for node_id, node_info in prompt_data.items():
|
||||
if not isinstance(node_info, dict):
|
||||
continue
|
||||
class_type = node_info.get("class_type", "")
|
||||
if class_type in ("PromptManager", "PromptManagerText"):
|
||||
inputs = node_info.get("inputs", {})
|
||||
text = inputs.get("text", "")
|
||||
if text and isinstance(text, str) and text.strip():
|
||||
prompt_text = text.strip()
|
||||
break
|
||||
|
||||
# Fallback: check text_encoder_nodes from workflow, but only if
|
||||
# the text input is NOT connected (connected inputs override widget values,
|
||||
# so the widget value would be stale in batch workflows).
|
||||
if not prompt_text:
|
||||
text_nodes = metadata.get("text_encoder_nodes", [])
|
||||
for node in text_nodes:
|
||||
node_type = node.get("type") or node.get("class_type") or ""
|
||||
if "PromptManager" in node_type:
|
||||
# Check if text input is connected — if so, widget value is stale
|
||||
text_connected = False
|
||||
for inp in node.get("inputs", []):
|
||||
if isinstance(inp, dict) and inp.get("name") == "text":
|
||||
if inp.get("link") is not None:
|
||||
text_connected = True
|
||||
break
|
||||
if text_connected:
|
||||
self.logger.debug(
|
||||
"PromptManager text input is connected — "
|
||||
"skipping stale widget value, deferring to snapshot"
|
||||
)
|
||||
continue
|
||||
widgets = node.get("widgets_values", [])
|
||||
if widgets and isinstance(widgets[0], str) and widgets[0].strip():
|
||||
prompt_text = widgets[0].strip()
|
||||
break
|
||||
|
||||
if not prompt_text:
|
||||
return None
|
||||
|
||||
# Look up by hash in database
|
||||
try:
|
||||
import hashlib
|
||||
|
||||
normalized = prompt_text.strip().lower()
|
||||
prompt_hash = hashlib.sha256(normalized.encode("utf-8")).hexdigest()
|
||||
existing = self.db_manager.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
self.logger.debug(
|
||||
f"Found DB prompt {existing['id']} from metadata text"
|
||||
)
|
||||
return {
|
||||
"id": existing["id"],
|
||||
"text": existing["text"],
|
||||
"from_metadata": True,
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Metadata-based prompt lookup failed: {e}")
|
||||
|
||||
return None
|
||||
|
||||
def get_basic_file_info(self, image_path: str) -> Dict[str, Any]:
|
||||
"""Get basic file information when metadata extraction fails.
|
||||
|
||||
|
||||
@@ -76,6 +76,10 @@ class PromptTracker:
|
||||
self._local = threading.local()
|
||||
self.active_prompts = {} # Global tracking for multiple threads
|
||||
self.lock = threading.Lock()
|
||||
# FIFO queue for batch workflows: prompt_ids are pushed during encoding
|
||||
# and popped during image linking, preserving correct ordering.
|
||||
self._prompt_queue = []
|
||||
self._queue_lock = threading.Lock()
|
||||
|
||||
# Read from GalleryConfig if available, otherwise use defaults
|
||||
try:
|
||||
@@ -159,6 +163,13 @@ class PromptTracker:
|
||||
with self.lock:
|
||||
self.active_prompts[execution_id] = execution_context
|
||||
|
||||
# Push to batch queue for ordered image linking
|
||||
with self._queue_lock:
|
||||
self._prompt_queue.append(execution_context)
|
||||
self.logger.debug(
|
||||
f"Queue size: {len(self._prompt_queue)} after push for prompt {prompt_id}"
|
||||
)
|
||||
|
||||
self.logger.debug(
|
||||
f"Set current prompt: {execution_id} -> {prompt_text[:50]}... (thread: {threading.current_thread().ident})"
|
||||
)
|
||||
@@ -207,6 +218,32 @@ class PromptTracker:
|
||||
)
|
||||
return None
|
||||
|
||||
def pop_next_prompt(self) -> Optional[Dict[str, Any]]:
|
||||
"""Pop the next prompt from the batch queue (FIFO).
|
||||
|
||||
In batch workflows, prompts are encoded in order and images are saved
|
||||
in the same order. This method pops the oldest queued prompt, ensuring
|
||||
each image links to the correct prompt by position.
|
||||
|
||||
Skips and discards expired entries to prevent stale prompts from
|
||||
accumulating across sessions.
|
||||
|
||||
Returns:
|
||||
Prompt context dict, or None if queue is empty or all entries expired
|
||||
"""
|
||||
now = time.time()
|
||||
with self._queue_lock:
|
||||
while self._prompt_queue:
|
||||
ctx = self._prompt_queue.pop(0)
|
||||
if now - ctx.get("timestamp", 0) < self.prompt_timeout:
|
||||
self.logger.debug(
|
||||
f"Queue pop: prompt {ctx.get('id')} "
|
||||
f"({len(self._prompt_queue)} remaining)"
|
||||
)
|
||||
return ctx
|
||||
self.logger.debug(f"Queue discard expired: prompt {ctx.get('id')}")
|
||||
return None
|
||||
|
||||
def _find_recent_prompt(self) -> Optional[Dict[str, Any]]:
|
||||
"""Find the most recent prompt that's still valid.
|
||||
|
||||
@@ -335,6 +372,18 @@ class PromptTracker:
|
||||
if expired_ids:
|
||||
self.logger.debug(f"Cleaned up {len(expired_ids)} expired prompts")
|
||||
|
||||
# Also prune expired entries from the batch queue
|
||||
with self._queue_lock:
|
||||
before = len(self._prompt_queue)
|
||||
self._prompt_queue = [
|
||||
ctx
|
||||
for ctx in self._prompt_queue
|
||||
if current_time - ctx.get("timestamp", 0) < self.prompt_timeout
|
||||
]
|
||||
pruned = before - len(self._prompt_queue)
|
||||
if pruned:
|
||||
self.logger.debug(f"Pruned {pruned} expired queue entries")
|
||||
|
||||
time.sleep(self.cleanup_interval)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
Reference in New Issue
Block a user