diff --git a/README.md b/README.md
index daa797e..447cddb 100644
--- a/README.md
+++ b/README.md
@@ -2,15 +2,15 @@
> Structured workflow memory and human-in-the-loop creative feedback for ComfyUI.
-ComfyUI-Session-Memory gives your generation workflows **persistent memory across runs**. It accumulates structured run summaries, injects them into your LLM system prompt for context-aware generation, and provides an interactive in-node feedback panel where you can rate and annotate every prompt — all without leaving ComfyUI.
+ComfyUI-Session-Memory gives your generation workflows **persistent memory across runs**. It stores both the raw prompts and structured summaries, injects summaries into your LLM system prompt for context-aware generation, and provides an interactive in-node feedback panel where you can rate and annotate every prompt — all without leaving ComfyUI.
---
## ✨ Features
-- **Session Memory Reader** — Loads the full accumulated history of a session and outputs it as a string ready to inject into any LLM system prompt.
-- **Session Memory Writer** — Receives a structured run summary and appends it to the session's persistent JSON file. Outputs the current `run_count` for downstream nodes.
-- **Session Feedback Editor** — Interactive in-node UI to rate (⭐ 1–5) and annotate each prompt after every run. Supports visual thumbnail previews of generated images, organized per prompt and per batch pass.
+- **Session Memory Reader** — Loads the accumulated session history and outputs the **summaries** for injection into any LLM system prompt.
+- **Session Memory Writer** — Receives both the **raw prompt** and a **structured summary**, storing them as paired entries in the session's persistent JSON file.
+- **Session Feedback Editor** — Interactive in-node UI to rate (⭐ 1–5) and annotate each prompt. Displays the **raw prompt text** (not the summary) for accurate creative review. Supports visual thumbnail previews of generated images.
- **Feedback injection** — Rated prompts are automatically converted into `REPLICATE` / `AVOID` directives injected into the next run's system prompt, closing the creative feedback loop.
- **Thumbnail system** — Accepts a full image batch, downscales to ≤1 MP, and saves them organized as `sessions/{id}/run_{N}/prompt_{P}/img_{K}.jpg`. Supports multi-pass batches (`batch_count > 1`).
@@ -28,7 +28,9 @@ ComfyUI-Session-Memory gives your generation workflows **persistent memory acros
└────────────── sessions/{session_id}.json ◄────────┘
```
-The **Reader** always emits the history *before* the current run begins. The **Writer** appends the new summary *after* the run completes. The **Feedback Editor** operates non-destructively in parallel — it never blocks the Reader's next read until you explicitly save feedback.
+The **Reader** emits the accumulated **summaries** before the current run begins.
+The **Writer** stores both the **prompt** and **summary** after the run completes.
+The **Feedback Editor** displays the **prompt** text for review — it never shows the summary.
---
@@ -54,16 +56,18 @@ Loads and outputs the full accumulated session log.
### 💾 Session Memory Writer
-Appends a new run summary to the session file.
+Stores both the raw prompt and a structured summary as a paired entry.
| Input | Type | Description |
|---|---|---|
| `session_id` | `STRING` | Session identifier |
-| `run_summary` | `STRING` | Structured text summary from a Gemini/LLM summarizer node |
+| `prompt` | `STRING` | The raw prompt sent to the model. **Displayed** in the Feedback Editor |
+| `summary` | `STRING` | Structured summary from a Gemini/LLM summarizer. **Returned** by the Reader |
+| `max_runs` | `INT` | Sliding window: keep only the last N runs (0 = unlimited) |
| Output | Type | Description |
|---|---|---|
-| `run_count` | `INT` | Current run number after writing |
+| `run_summary` | `STRING` | Pass-through of the summary |
| `session_id` | `STRING` | Pass-through for chaining to Feedback Editor |
---
@@ -152,8 +156,10 @@ Search for `Session Memory` in the ComfyUI Manager and install.
{
"run_count": 5,
"entries": [
- "RUN 1:\nPROMPT 1: concept: urban decay...",
- "..."
+ {
+ "prompt": "--- RUN 1 ---\nPROMPT 1: concept: urban decay...\n",
+ "summary": "--- RUN 1 ---\nSummary of urban decay concept...\n"
+ }
],
"feedback": {
"3": {
@@ -166,6 +172,8 @@ Search for `Session Memory` in the ComfyUI Manager and install.
}
```
+> **Backwards compatibility:** Legacy sessions with plain-string entries are automatically migrated on load. Each old string becomes both the `prompt` and `summary`.
+
---
## 📄 License
diff --git a/session_memory.py b/session_memory.py
index 58cf5c4..1ee218e 100644
--- a/session_memory.py
+++ b/session_memory.py
@@ -50,7 +50,18 @@ def _load(session_id: str) -> dict:
if os.path.exists(path):
try:
with open(path, "r", encoding="utf-8") as f:
- return json.load(f)
+ data = json.load(f)
+ # Migrate legacy string entries to {prompt, summary} dicts
+ entries = data.get("entries", [])
+ migrated = False
+ for i, entry in enumerate(entries):
+ if isinstance(entry, str):
+ entries[i] = {"prompt": entry, "summary": entry}
+ migrated = True
+ if migrated:
+ data["entries"] = entries
+ logger.info(f"[Session Memory] Migrated {sum(1 for e in entries if isinstance(e, dict))} legacy entries in '{session_id}'.")
+ return data
except (json.JSONDecodeError, OSError) as e:
logger.warning(f"[Session Memory] Could not read {path}: {e}")
return {"run_count": 0, "entries": [], "feedback": {}}
@@ -66,12 +77,31 @@ def _save(session_id: str, data: dict) -> None:
def _parse_prompts_from_entry(entry_text: str) -> list[dict]:
- """Extract prompt headers, concepts, and full text from a run entry string."""
+ """Extract prompt headers, concepts, and full text from a run entry string.
+
+ Supports three formats (tried in order):
+ 1. XML-style ``text`` blocks (Visual DNA Synthesizer)
+ 2. Legacy ``PROMPT N:`` markers
+ 3. Fallback: whole entry treated as a single Prompt 1
+ """
prompts = []
+
+ # ── Strategy 1: XML-style tags ──────────────────────────
+ xml_blocks = re.findall(r'\s*(.*?)\s*', entry_text, re.DOTALL)
+ if xml_blocks:
+ for i, block in enumerate(xml_blocks, start=1):
+ prompts.append({
+ "number": i,
+ "concept": f"Prompt {i}",
+ "fullText": block.strip(),
+ })
+ return prompts
+
+ # ── Strategy 2: PROMPT N: markers ────────────────────────────────
prompt_blocks = re.split(r'(?=PROMPT \d+:)', entry_text)
# Markers that indicate the start of a trailing summary section (not prompt content)
_SUMMARY_MARKERS = re.compile(
- r'^(SESSION\s+TALLY|TALLY|SUMMARY|SESSION_FEEDBACK|RUN\s+\d+)', re.MULTILINE
+ r'^(SESSION\s+TALLY|TALLY|SUMMARY|SESSION_FEEDBACK|RUN\s+\d+:)', re.MULTILINE
)
for block in prompt_blocks:
m = re.match(r'PROMPT (\d+):', block)
@@ -91,20 +121,23 @@ def _parse_prompts_from_entry(entry_text: str) -> list[dict]:
"concept": concept,
"fullText": full_text,
})
-
- # Fallback: if no PROMPT X markers, treat the whole entry as Prompt 1
- if not prompts and entry_text.strip():
+
+ if prompts:
+ return prompts
+
+ # ── Fallback: whole entry as Prompt 1 ────────────────────────────
+ if entry_text.strip():
full_text = entry_text.strip()
trail = _SUMMARY_MARKERS.search(full_text)
if trail:
full_text = full_text[:trail.start()]
-
+
prompts.append({
"number": 1,
"concept": "Result",
"fullText": full_text.strip()
})
-
+
return prompts
@@ -123,8 +156,18 @@ def _save_thumbnails(session_id, run_number, num_prompts, batch_count,
Layout: sessions/{sid}/run_{N}/prompt_{P}/img_{K}.jpg
where K is the batch/pass index (0-based).
- Image order in tensor: [pass0_p1, pass0_p2, ..., pass1_p1, pass1_p2, ...]
- Total images = num_prompts × batch_count.
+ Image order in tensor: [p1, p2, ..., pN, p1, p2, ..., pN, ...]
+ distributed round-robin across prompts.
+
+ With ``batch_count > 1`` ComfyUI **re-executes the entire graph** for
+ each batch pass, sending the *full* image tensor every time. We must
+ avoid saving duplicate images. Strategy:
+
+ * On the **first call** (prompt dirs are empty) → save all images.
+ * On **subsequent calls** (prompt dirs already populated) → compute
+ how many images *per prompt* this call would add. If the folder
+ already holds ``>= expected_per_prompt`` images, skip entirely
+ (the images are duplicates from a graph re-execution).
"""
import numpy as np
from PIL import Image as PILImage
@@ -132,18 +175,37 @@ def _save_thumbnails(session_id, run_number, num_prompts, batch_count,
base_dir = _session_images_dir(session_id)
run_dir = os.path.join(base_dir, f"run_{run_number}")
- # NOTE: Do NOT rmtree run_dir here. ComfyUI with batch_count>1
- # re-executes the whole graph for each batch. Batch 1 saves img_0,
- # then batch 2 arrives and must APPEND img_1, not destroy img_0.
-
B = images_tensor.shape[0]
- saved = 0
logger.info(
f"[Session Feedback] Saving thumbnails: B={B}, "
f"num_prompts={num_prompts}, batch_count_widget={batch_count}"
)
+ # The user-specified batch_count is the authoritative expected number
+ # of images per prompt. Use it as the dedup threshold:
+ # - Combined tensor (B = num_prompts × batch_count): first call saves
+ # everything correctly; subsequent re-executions are skipped.
+ # - Partial batches (B = num_prompts per call): each call appends
+ # until batch_count is reached, then subsequent calls are skipped.
+ expected_per_prompt = max(1, batch_count)
+
+ first_prompt_dir = os.path.join(run_dir, "prompt_1")
+ if os.path.isdir(first_prompt_dir):
+ existing_count = len([
+ f for f in os.listdir(first_prompt_dir)
+ if f.startswith("img_") and f.endswith(".jpg")
+ ])
+ if existing_count >= expected_per_prompt:
+ logger.info(
+ f"[Session Feedback] Prompt_1 already has {existing_count} "
+ f"images (>= {expected_per_prompt} expected per batch_count). "
+ f"Skipping duplicate re-execution."
+ )
+ return
+
+ saved = 0
+
# Distribute images round-robin: image[i] → prompt (i % num_prompts)
for i in range(B):
p = i % num_prompts # 0-based prompt index
@@ -258,7 +320,10 @@ def _format_feedback_block(feedback: dict) -> str:
@PromptServer.instance.routes.get("/session_feedback/load")
async def _api_load_feedback(request):
- """Load run data + existing feedback for the frontend editor."""
+ """Load run data + existing feedback for the frontend editor.
+
+ The feedback editor shows the PROMPT text (not the summary).
+ """
session_id = request.query.get("session_id", "default")
run_number = request.query.get("run_number", "")
@@ -284,7 +349,12 @@ async def _api_load_feedback(request):
sel_run = max(1, min(sel_run, run_count))
idx = sel_run - 1
- entry_text = entries[idx] if idx < len(entries) else ""
+ entry = entries[idx] if idx < len(entries) else {}
+ # Use the prompt field for display in the feedback editor
+ if isinstance(entry, dict):
+ entry_text = entry.get("prompt", "")
+ else:
+ entry_text = entry # legacy string fallback
prompts = _parse_prompts_from_entry(entry_text)
run_feedback = feedback.get(str(sel_run), {}).get("prompts", {})
@@ -393,7 +463,22 @@ async def _api_delete_run(request):
# 1. Update entries
if 0 <= idx < len(entries):
entries.pop(idx)
+
session["run_count"] = max(0, run_count - 1)
+
+ # Renumber the headers inside the entries to maintain 1..N contiguity
+ for i in range(len(entries)):
+ if isinstance(entries[i], dict):
+ for field in ("prompt", "summary"):
+ if field in entries[i]:
+ entries[i][field] = re.sub(
+ r'^---\s*RUN\s+\d+\s*---\n',
+ f'--- RUN {i+1} ---\n',
+ entries[i][field], count=1, flags=re.MULTILINE,
+ )
+ else:
+ entries[i] = re.sub(r'^---\s*RUN\s+\d+\s*---\n', f'--- RUN {i+1} ---\n', entries[i], count=1, flags=re.MULTILINE)
+
session["entries"] = entries
# 2. Update feedback keys
@@ -646,9 +731,14 @@ class SessionFeedbackEditorNode:
)
# Get latest run prompts + existing feedback
+ # Feedback editor shows the PROMPT text (not summary)
entries = session.get("entries", [])
- latest_entry = entries[-1] if entries else ""
- prompts = _parse_prompts_from_entry(latest_entry)
+ latest_entry = entries[-1] if entries else {}
+ if isinstance(latest_entry, dict):
+ entry_text = latest_entry.get("prompt", "")
+ else:
+ entry_text = latest_entry # legacy string fallback
+ prompts = _parse_prompts_from_entry(entry_text)
feedback = session.get("feedback", {})
run_feedback = feedback.get(str(run_count), {}).get("prompts", {})
@@ -743,7 +833,28 @@ class SessionMemoryReaderNode:
run_count = session.get("run_count", 0)
feedback = session.get("feedback", {})
- session_memory = "\n".join(entries) if entries else ""
+ # Reader returns the SUMMARY field (not the prompt).
+ # When summary is empty / placeholder, try extracting
+ # tags from the prompt field (Visual DNA Synthesizer format).
+ _PLACEHOLDER = "(No summary provided for this run)"
+ summaries = []
+ for entry in entries:
+ if isinstance(entry, dict):
+ summary_text = entry.get("summary", "")
+ # Detect placeholder / empty summary → extract tags
+ stripped = summary_text.strip().lstrip("- RUN 0123456789\n")
+ if not stripped or stripped == _PLACEHOLDER:
+ prompt_text = entry.get("prompt", "")
+ xml_sums = re.findall(
+ r'\s*(.*?)\s*',
+ prompt_text, re.DOTALL,
+ )
+ if xml_sums:
+ summary_text = "\n".join(s.strip() for s in xml_sums)
+ summaries.append(summary_text)
+ else:
+ summaries.append(entry) # legacy string fallback
+ session_memory = "\n".join(summaries) if summaries else ""
# Append feedback block if any
feedback_block = _format_feedback_block(feedback)
@@ -776,10 +887,15 @@ class SessionMemoryWriterNode:
"default": "default",
"tooltip": "Session identifier. Must match Reader/Feedback.",
}),
- "run_summary": ("STRING", {
+ "prompt": ("STRING", {
"multiline": True,
"default": "",
- "tooltip": "Structured summary from the Gemini Summarizer.",
+ "tooltip": "The raw prompt sent to the model. Displayed in the Feedback Editor.",
+ }),
+ "summary": ("STRING", {
+ "multiline": True,
+ "default": "",
+ "tooltip": "Structured summary from the Gemini Summarizer. Returned by the Reader.",
}),
},
"optional": {
@@ -800,26 +916,37 @@ class SessionMemoryWriterNode:
FUNCTION = "write"
CATEGORY = "\U0001f9e0 Memory"
+ @classmethod
+ def IS_CHANGED(cls, **kwargs):
+ return float("NaN")
+
def write(
self,
session_id: str = "default",
- run_summary: str = "",
+ prompt: str = "",
+ summary: str = "",
max_runs: int = 0,
):
- stripped = run_summary.strip()
- if not stripped:
+ prompt_stripped = prompt.strip()
+ summary_stripped = summary.strip()
+
+ if not prompt_stripped:
+ prompt_stripped = "(No prompt provided for this run)"
+ if not summary_stripped:
+ summary_stripped = "(No summary provided for this run)"
logger.info(
- f"[Session Memory Writer] Empty summary — nothing written "
- f"to session '{session_id}'."
+ f"[Session Memory Writer] Empty summary — creating a new run with placeholder in session '{session_id}'."
)
- return (run_summary, session_id)
session = _load(session_id)
run_count = session.get("run_count", 0) + 1
- entries: list[str] = session.get("entries", [])
+ entries: list = session.get("entries", [])
feedback = session.get("feedback", {})
- entry = f"--- RUN {run_count} ---\n{stripped}\n"
+ entry = {
+ "prompt": f"--- RUN {run_count} ---\n{prompt_stripped}\n",
+ "summary": f"--- RUN {run_count} ---\n{summary_stripped}\n",
+ }
entries.append(entry)
# Sliding window
@@ -838,9 +965,10 @@ class SessionMemoryWriterNode:
})
logger.info(
f"[Session Memory Writer] Written Run {run_count} to session "
- f"'{session_id}' ({len(stripped)} chars, {len(entries)} entries)."
+ f"'{session_id}' (prompt={len(prompt_stripped)} chars, "
+ f"summary={len(summary_stripped)} chars, {len(entries)} entries)."
)
- return (run_summary, session_id)
+ return (summary, session_id)
def _get_available_sessions():