Vito Sansevero
e3191853c8
database/scanned_prompts.py links a file to the prompt of one text, creating it when new, inside one BEGIN IMMEDIATE transaction: a file linked to another prompt is a conflict, a failed link leaves no prompt behind and two writers store one prompt. pair_ingestion gains file_link_rows, link_conflict and Marks/mark_created; ingest_resolved_pair takes optional marks for the components it creates. recorded_runs and named_prompt in run_attribution become public (#247).
293 lines
10 KiB
Python
293 lines
10 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 (either spelling), each with its prompt's hash."""
|
|
return 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(image_path), image_path),
|
|
).fetchall()
|
|
|
|
|
|
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)}
|