Vito Sansevero
3dfb828ae5
The scan stores realpath spellings and the monitor the watcher's raw path, so a file reached through a symlinked or junction output folder could get a second generated_images row. file_link_rows now queries the path and its realpath (each normalized and raw) and dedupes by row id. TESTING.md notes the reverse order is matched only when the raw spelling is the realpath.
309 lines
11 KiB
Python
309 lines
11 KiB
Python
"""Atomic metadata enrichment for a single known generation context.
|
|
|
|
Also the helpers both ingestion store paths share (#247): the file's link
|
|
rows and the one-prompt-per-file rule (``file_link_rows``,
|
|
``link_conflict``) and the marks a prompt created by ingestion carries
|
|
(``Marks``, ``mark_created``).
|
|
"""
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from .models import normalize_image_path
|
|
from .operations import _coerce_file_info
|
|
from .workflow_documents import (
|
|
TEXT_COLUMNS,
|
|
DocumentReader,
|
|
release_documents,
|
|
resolve_image_row,
|
|
stored_values,
|
|
)
|
|
|
|
try:
|
|
from ..utils.hashing import generate_prompt_hash
|
|
from ..utils.json_strict import dumps_strict
|
|
from ..utils.prompt_pair_resolver import RESOLVER_VERSION
|
|
except ImportError:
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.json_strict import dumps_strict
|
|
from utils.prompt_pair_resolver import RESOLVER_VERSION
|
|
|
|
|
|
def blank_positive(resolution: Dict[str, Any]) -> bool:
|
|
"""A resolved pair whose stored positive text is empty (#342).
|
|
|
|
A PromptManager node whose editor is empty encodes only its prepend and
|
|
append text, which is never stored as a prompt: such a result names no
|
|
prompt to link.
|
|
"""
|
|
if resolution.get("state") != "resolved":
|
|
return False
|
|
positive = resolution.get("positive") or {}
|
|
text = positive.get("storage_text", positive.get("text"))
|
|
return not (isinstance(text, str) and text.strip())
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Marks:
|
|
"""What a prompt ingestion creates carries, so it can be found and removed.
|
|
|
|
``note`` is formatted with ``name``, the file's name.
|
|
"""
|
|
|
|
category: str
|
|
tags: Tuple[str, ...]
|
|
note: str
|
|
|
|
|
|
SCAN_MARKS = Marks("scanned", ("auto-scanned",), "Auto-scanned from {name}")
|
|
RUN_MARKS = Marks("scanned", ("auto-scanned",), "Saved from a workflow run: {name}")
|
|
|
|
|
|
def mark_created(conn: Any, db: Any, prompt_id: int, marks: Marks, name: str) -> None:
|
|
"""Apply ``marks`` to a prompt row the caller just created on ``conn``."""
|
|
conn.execute(
|
|
"UPDATE prompts SET category = ?, notes = ? WHERE id = ?",
|
|
(marks.category, marks.note.format(name=name), prompt_id),
|
|
)
|
|
db._sync_prompt_tags(conn, prompt_id, list(marks.tags))
|
|
|
|
|
|
def file_link_rows(conn: Any, image_path: str) -> List[Any]:
|
|
"""The file's link rows, each with its prompt's hash.
|
|
|
|
Matched under the path and its resolved spelling (an output folder
|
|
reached through a symlink or junction), each normalized and raw;
|
|
deduplicated by row id.
|
|
"""
|
|
rows: Dict[Any, Any] = {}
|
|
for spelling in dict.fromkeys((image_path, _resolved(image_path))):
|
|
for row in conn.execute(
|
|
"SELECT gi.*,p.hash AS positive_hash FROM generated_images gi "
|
|
"JOIN prompts p ON p.id=gi.prompt_id "
|
|
"WHERE gi.file_path=? OR gi.image_path=?",
|
|
(normalize_image_path(spelling), spelling),
|
|
).fetchall():
|
|
rows.setdefault(row["id"], row)
|
|
return list(rows.values())
|
|
|
|
|
|
def _resolved(path: str) -> str:
|
|
try:
|
|
return os.path.realpath(path)
|
|
except (OSError, ValueError):
|
|
return path
|
|
|
|
|
|
def link_conflict(rows: List[Any], prompt_hash: str) -> bool:
|
|
"""One prompt per file: more than one row, or a row of another prompt."""
|
|
return len(rows) > 1 or bool(rows and rows[0]["positive_hash"] != prompt_hash)
|
|
|
|
|
|
def ingest_resolved_pair(
|
|
db: Any,
|
|
image_path: str,
|
|
metadata: Dict[str, Any],
|
|
resolution: Dict[str, Any],
|
|
*,
|
|
marks: Optional[Marks] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Store a resolved pair and link the file to it in one transaction.
|
|
|
|
With ``marks`` (ingestion from files), each pair component this call
|
|
creates carries them; without (the PromptManager node path) components
|
|
are created bare, as always.
|
|
"""
|
|
state = resolution.get("state", "unresolved")
|
|
if state == "resolved" and blank_positive(resolution):
|
|
state = "unresolved"
|
|
if state != "resolved":
|
|
return {"status": state, "pairs_added": 0, "links_added": 0}
|
|
with db._pair_transaction() as conn:
|
|
return _ingest_resolved_pair_in(
|
|
conn, db, image_path, metadata, resolution, marks=marks
|
|
)
|
|
|
|
|
|
def _conflict() -> Dict[str, Any]:
|
|
return {"status": "conflict", "pairs_added": 0, "links_added": 0}
|
|
|
|
|
|
def _ingest_resolved_pair_in(
|
|
conn: Any,
|
|
db: Any,
|
|
image_path: str,
|
|
metadata: Dict[str, Any],
|
|
resolution: Dict[str, Any],
|
|
*,
|
|
provenance: Dict[str, Any] | None = None,
|
|
marks: Optional[Marks] = None,
|
|
) -> Dict[str, Any]:
|
|
"""``ingest_resolved_pair`` inside the caller's open ``_pair_transaction``.
|
|
|
|
The compatible-only rules are the public ingester's: an image with more
|
|
than one exact row, another positive, or a prior pair with another
|
|
negative is a ``conflict`` and nothing is written. The caller owns the
|
|
transaction, so its own later failure rolls the link back too.
|
|
``provenance`` replaces the link's default resolver provenance (Needs
|
|
review records the user's selection there). ``marks`` as in
|
|
``ingest_resolved_pair``.
|
|
"""
|
|
positive, negative = resolution["positive"], resolution["negative"]
|
|
positive_text = positive.get("storage_text", positive["text"])
|
|
negative_text = negative.get("storage_text", negative["text"])
|
|
normalized = normalize_image_path(image_path)
|
|
if provenance is None:
|
|
provenance = {
|
|
"resolver_version": RESOLVER_VERSION,
|
|
"positive": positive,
|
|
"negative": negative,
|
|
}
|
|
encoded = dumps_strict(provenance, sort_keys=True, separators=(",", ":"))
|
|
images = file_link_rows(conn, image_path)
|
|
if link_conflict(images, generate_prompt_hash(positive_text)):
|
|
return _conflict()
|
|
existing = images[0] if images else None
|
|
if existing:
|
|
prior = conn.execute(
|
|
"SELECT n.hash AS negative_hash FROM image_prompt_pairs ip "
|
|
"JOIN prompt_pairs pp ON pp.id=ip.pair_id "
|
|
"LEFT JOIN prompts n ON n.id=pp.negative_prompt_id WHERE ip.image_id=?",
|
|
(existing["id"],),
|
|
).fetchall()
|
|
expected = (
|
|
generate_prompt_hash(negative_text) if negative_text.strip() else None
|
|
)
|
|
if any(row["negative_hash"] != expected for row in prior):
|
|
return _conflict()
|
|
prior_pair = conn.execute(
|
|
"SELECT pp.id FROM prompt_pairs pp "
|
|
"JOIN prompts p ON p.id=pp.positive_prompt_id "
|
|
"LEFT JOIN prompts n ON n.id=pp.negative_prompt_id "
|
|
"WHERE p.hash=? AND n.hash IS ?",
|
|
(
|
|
generate_prompt_hash(positive_text),
|
|
generate_prompt_hash(negative_text) if negative_text.strip() else None,
|
|
),
|
|
).fetchone()
|
|
if marks is None:
|
|
pair_id = db._save_pair(conn, positive_text, negative_text)
|
|
else:
|
|
pair_id = _save_marked_pair(
|
|
conn, db, positive_text, negative_text, marks, os.path.basename(image_path)
|
|
)
|
|
pair = conn.execute("SELECT * FROM prompt_pairs WHERE id=?", (pair_id,)).fetchone()
|
|
pairs_added = int(prior_pair is None)
|
|
file_size, width, height, image_format = _coerce_file_info(metadata)
|
|
# ComfyUI's NaN/Infinity are stored as null: the snapshot is strict JSON.
|
|
snapshots = {
|
|
column: dumps_strict(metadata[key], sort_keys=True)
|
|
for column, key in (
|
|
("workflow_data", "workflow"),
|
|
("prompt_metadata", "prompt"),
|
|
("parameters", "parameters"),
|
|
)
|
|
if key in metadata
|
|
}
|
|
info = {
|
|
"file_size": file_size,
|
|
"width": width,
|
|
"height": height,
|
|
"format": image_format,
|
|
}
|
|
snapshots.update({key: value for key, value in info.items() if value is not None})
|
|
changed = False
|
|
if existing:
|
|
image_id = existing["id"]
|
|
# Compared as stored text, whichever form the row is in.
|
|
current = resolve_image_row(DocumentReader(conn), existing)
|
|
updates = {
|
|
key: value for key, value in snapshots.items() if current[key] != value
|
|
}
|
|
if updates:
|
|
replaced = [
|
|
existing[reference]
|
|
for column, reference in TEXT_COLUMNS
|
|
if column in updates
|
|
]
|
|
updates = _document_form(conn, updates)
|
|
conn.execute(
|
|
"UPDATE generated_images SET "
|
|
+ ",".join(f"{key}=?" for key in updates)
|
|
+ " WHERE id=?",
|
|
[*updates.values(), image_id],
|
|
)
|
|
release_documents(conn, replaced)
|
|
changed = True
|
|
else:
|
|
stored = _document_form(conn, snapshots)
|
|
columns = ["prompt_id", "image_path", "filename", "file_path", *stored]
|
|
values = [
|
|
pair["positive_prompt_id"],
|
|
image_path,
|
|
os.path.basename(image_path),
|
|
normalized,
|
|
*stored.values(),
|
|
]
|
|
image_id = conn.execute(
|
|
"INSERT INTO generated_images("
|
|
+ ",".join(columns)
|
|
+ ") VALUES("
|
|
+ ",".join("?" for _ in columns)
|
|
+ ")",
|
|
values,
|
|
).lastrowid
|
|
linked = conn.execute(
|
|
"SELECT 1 FROM image_prompt_pairs WHERE pair_id=? "
|
|
"AND image_id=? AND provenance=?",
|
|
(pair_id, image_id, encoded),
|
|
).fetchone()
|
|
db._link_image_to_pair(conn, pair_id, image_id, provenance)
|
|
status = (
|
|
"added"
|
|
if existing is None
|
|
else "enriched" if changed or not linked else "unchanged"
|
|
)
|
|
return {
|
|
"status": status,
|
|
"pairs_added": pairs_added,
|
|
"links_added": int(not linked),
|
|
"pair_id": pair_id,
|
|
"image_id": image_id,
|
|
}
|
|
|
|
|
|
def _save_marked_pair(
|
|
conn: Any,
|
|
db: Any,
|
|
positive_text: str,
|
|
negative_text: str,
|
|
marks: Marks,
|
|
name: str,
|
|
) -> int:
|
|
"""``_save_pair`` that marks each component it creates."""
|
|
texts = [positive_text] + ([negative_text] if negative_text.strip() else [])
|
|
ids = []
|
|
for text in texts:
|
|
prompt_id, created = db._pair_component_outcome(conn, text)
|
|
if created:
|
|
mark_created(conn, db, prompt_id, marks, name)
|
|
ids.append(prompt_id)
|
|
negative_id = ids[1] if len(ids) > 1 else None
|
|
return db._save_pair_components(conn, ids[0], negative_id)[0]
|
|
|
|
|
|
def _document_form(conn: Any, values: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""``values`` with the workflow and prompt graph texts stored once.
|
|
|
|
Each of those columns becomes NULL beside its document id (new form);
|
|
every other column is unchanged.
|
|
"""
|
|
texts = {
|
|
column: values[column]
|
|
for column, _reference in TEXT_COLUMNS
|
|
if column in values
|
|
}
|
|
rest = {key: value for key, value in values.items() if key not in texts}
|
|
return {**rest, **stored_values(conn, texts)}
|