feat(images): link summary on workspace image items
This commit is contained in:
+23
-4
@@ -13,10 +13,29 @@ _ELIGIBLE = """(EXISTS(SELECT 1 FROM prompt_pairs ep WHERE ep.positive_prompt_id
|
||||
_FROM = "prompts p LEFT JOIN prompt_library_state ls ON ls.prompt_id=p.id"
|
||||
# The characters JavaScript's String.prototype.trim() removes, so a title the
|
||||
# display treats as blank also falls back to the prompt text for name sorting.
|
||||
_BLANK = (
|
||||
"char(9,10,11,12,13,32,160,5760,8192,8193,8194,8195,8196,8197,8198,8199,"
|
||||
"8200,8201,8202,8232,8233,8239,8287,12288,65279)"
|
||||
)
|
||||
_TRIM_CODES = (9, 10, 11, 12, 13, 32, 160, 5760, *range(8192, 8203))
|
||||
_TRIM_CODES += (8232, 8233, 8239, 8287, 12288, 65279)
|
||||
_BLANK = f"char({','.join(map(str, _TRIM_CODES))})"
|
||||
_JS_TRIM = "".join(map(chr, _TRIM_CODES))
|
||||
# The Library's display title: a stored title, else the prompt's first line.
|
||||
DISPLAY_TITLE_LIMIT = 80
|
||||
|
||||
|
||||
def display_title(stored_titles: tuple, text: object) -> str | None:
|
||||
"""Return the first non-blank stored title, else the text's first line.
|
||||
|
||||
Blank follows JavaScript's trim, as in the Library view; the first line is
|
||||
trimmed and cut to DISPLAY_TITLE_LIMIT characters. None when both are empty.
|
||||
"""
|
||||
for title in stored_titles:
|
||||
if isinstance(title, str) and title.strip(_JS_TRIM):
|
||||
return title
|
||||
if not isinstance(text, str):
|
||||
return None
|
||||
first = text.split("\n")[0].strip(_JS_TRIM)[:DISPLAY_TITLE_LIMIT]
|
||||
return first or None
|
||||
|
||||
|
||||
_NAME = (
|
||||
f"LOWER(CASE WHEN TRIM(COALESCE(ls.title,''),{_BLANK})!='' "
|
||||
"THEN ls.title ELSE p.text END)"
|
||||
|
||||
+103
-1
@@ -10,7 +10,7 @@ from typing import Optional, List, Dict, Any
|
||||
|
||||
from .models import PromptModel, normalize_image_path, utc_now_iso
|
||||
from .pairs import PairDependencyError, PairOperations
|
||||
from .library import LibraryOperations
|
||||
from .library import LibraryOperations, display_title
|
||||
from .library_commands import LibraryCommandOperations
|
||||
from .vocabulary import VocabularyOperations
|
||||
from .execution_history import ExecutionHistoryOperations
|
||||
@@ -179,6 +179,59 @@ AND p.id NOT IN (SELECT prompt_id FROM live_hint_prompts)
|
||||
"""
|
||||
|
||||
|
||||
# One SELECT reads the link summaries of at most this many paths: three bound
|
||||
# parameters per path stays under SQLite's historical 999-variable limit.
|
||||
LINK_SUMMARY_BATCH = 200
|
||||
# A pair link counts only when an image has exactly one distinct pair, as in
|
||||
# get_image_prompt_info; its positive, title and metadata then describe the card.
|
||||
_LINK_SUMMARY_SQL = """SELECT gi.image_path, gi.file_path, gi.prompt_id,
|
||||
p.text, p.category, p.rating, ls.title,
|
||||
pp.positive_prompt_id, pp.title, pt.text, pt.category, pt.rating, pls.title
|
||||
FROM generated_images gi
|
||||
JOIN prompts p ON gi.prompt_id = p.id
|
||||
LEFT JOIN prompt_library_state ls ON ls.prompt_id = p.id
|
||||
LEFT JOIN prompt_pairs pp ON pp.id = (
|
||||
SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END
|
||||
FROM image_prompt_pairs ip WHERE ip.image_id = gi.id)
|
||||
LEFT JOIN prompts pt ON pt.id = pp.positive_prompt_id
|
||||
LEFT JOIN prompt_library_state pls ON pls.prompt_id = pp.positive_prompt_id
|
||||
WHERE gi.image_path IN ({marks}) OR gi.file_path IN ({marks}){fallback}"""
|
||||
_ASCII_FOLD = str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz")
|
||||
|
||||
|
||||
def _like_suffix(value: object, suffix: str) -> bool:
|
||||
"""Python form of SQLite's ``value LIKE '%' || suffix``: ASCII-only case."""
|
||||
return isinstance(value, str) and value.translate(_ASCII_FOLD).endswith(
|
||||
suffix.translate(_ASCII_FOLD)
|
||||
)
|
||||
|
||||
|
||||
def _empty_link_summary(status: str) -> Dict[str, Any]:
|
||||
return {
|
||||
"status": status,
|
||||
"positive_id": None,
|
||||
"title": None,
|
||||
"category": None,
|
||||
"rating": None,
|
||||
}
|
||||
|
||||
|
||||
def _link_summary(row: tuple) -> Dict[str, Any]:
|
||||
if row[7] is not None:
|
||||
positive, text, category, rating = row[7], row[9], row[10], row[11]
|
||||
titles = (row[8], row[12])
|
||||
else:
|
||||
positive, text, category, rating = row[2], row[3], row[4], row[5]
|
||||
titles = (row[6],)
|
||||
return {
|
||||
"status": "linked",
|
||||
"positive_id": positive,
|
||||
"title": display_title(titles, text),
|
||||
"category": category,
|
||||
"rating": rating,
|
||||
}
|
||||
|
||||
|
||||
class PromptDatabase(
|
||||
PairOperations,
|
||||
LibraryOperations,
|
||||
@@ -2160,6 +2213,55 @@ class PromptDatabase(
|
||||
"parameters": row[8] if len(row) > 8 else None,
|
||||
}
|
||||
|
||||
def get_image_link_summaries(
|
||||
self, paths: List[str], allow_filename_fallback: bool = True
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
"""Summarize each path's linked saved positive for an image card.
|
||||
|
||||
Matching equals get_image_prompt_info: exact image_path or normalized
|
||||
file_path rows win; otherwise, with the fallback, rows whose image_path
|
||||
ends with the basename. One match is linked, several are ambiguous.
|
||||
One SELECT reads up to LINK_SUMMARY_BATCH paths.
|
||||
"""
|
||||
unique = list(dict.fromkeys(paths))
|
||||
summaries: Dict[str, Dict[str, Any]] = {}
|
||||
for start in range(0, len(unique), LINK_SUMMARY_BATCH):
|
||||
batch = unique[start : start + LINK_SUMMARY_BATCH]
|
||||
summaries.update(self._link_summary_batch(batch, allow_filename_fallback))
|
||||
return summaries
|
||||
|
||||
def _link_summary_batch(
|
||||
self, paths: List[str], allow_filename_fallback: bool
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
normalized = {path: normalize_image_path(path) for path in paths}
|
||||
basenames = {path: os.path.basename(path) for path in paths}
|
||||
marks = ",".join("?" for _ in paths)
|
||||
params: List[Any] = [*paths, *normalized.values()]
|
||||
fallback = ""
|
||||
if allow_filename_fallback:
|
||||
fallback = "".join(" OR gi.image_path LIKE ? ESCAPE '\\'" for _ in paths)
|
||||
params.extend(f"%{escape_like(basenames[path])}" for path in paths)
|
||||
sql = _LINK_SUMMARY_SQL.format(marks=marks, fallback=fallback)
|
||||
with self.model.get_connection() as conn:
|
||||
rows = conn.execute(sql, params).fetchall()
|
||||
summaries = {}
|
||||
for path in paths:
|
||||
exact = [
|
||||
row for row in rows if row[0] == path or row[1] == normalized[path]
|
||||
]
|
||||
matches = exact or (
|
||||
[row for row in rows if _like_suffix(row[0], basenames[path])]
|
||||
if allow_filename_fallback
|
||||
else []
|
||||
)
|
||||
if len(matches) == 1:
|
||||
summaries[path] = _link_summary(matches[0])
|
||||
else:
|
||||
summaries[path] = _empty_link_summary(
|
||||
"ambiguous" if matches else "none"
|
||||
)
|
||||
return summaries
|
||||
|
||||
def get_prompt_id_for_image(self, image_path: str) -> Optional[int]:
|
||||
"""Return the prompt_id linked to an image, or None."""
|
||||
with self.model.get_connection() as conn:
|
||||
|
||||
+64
-5
@@ -29,9 +29,11 @@ from .inspection_service import (
|
||||
)
|
||||
|
||||
try:
|
||||
from ...database.library import display_title
|
||||
from ...utils.parallel import map_parallel
|
||||
from ...utils.inspection_response import encode_inspection_response
|
||||
except ImportError:
|
||||
from database.library import display_title
|
||||
from utils.parallel import map_parallel
|
||||
from utils.inspection_response import encode_inspection_response
|
||||
|
||||
@@ -39,6 +41,35 @@ except ImportError:
|
||||
IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS
|
||||
|
||||
|
||||
def _unavailable_link() -> dict:
|
||||
"""A card link summary when the database could not be read."""
|
||||
return {
|
||||
"status": "unavailable",
|
||||
"positive_id": None,
|
||||
"title": None,
|
||||
"category": None,
|
||||
"rating": None,
|
||||
}
|
||||
|
||||
|
||||
def _verified_title(saved: dict | None, title: str | None, observed: dict):
|
||||
"""The inspector's Library title, withholding unverified saved prompt text.
|
||||
|
||||
A stored title is always shown. A title taken from the saved positive's
|
||||
first line is shown only when the file's own resolved positive gives the
|
||||
same title, because a saved association never replaces file evidence.
|
||||
"""
|
||||
if title is None or not isinstance(saved, dict):
|
||||
return title
|
||||
pair = saved.get("pair")
|
||||
text = pair.get("positive_text") if isinstance(pair, dict) else saved.get("text")
|
||||
if title != display_title((), text):
|
||||
return title
|
||||
side = observed.get("prompts", {}).get("positive") or {}
|
||||
resolved = side.get("status") == "resolved"
|
||||
return title if resolved and display_title((), side.get("text")) == title else None
|
||||
|
||||
|
||||
def _public_path(path):
|
||||
"""Response-safe form of *path* (relative to the ComfyUI tree).
|
||||
|
||||
@@ -1712,11 +1743,14 @@ class ImageRoutesMixin:
|
||||
try:
|
||||
|
||||
def inspect() -> bytes | None:
|
||||
saved, unavailable = None, False
|
||||
saved, unavailable, title = None, False, None
|
||||
try:
|
||||
saved = self.db.get_image_prompt_info(
|
||||
str(path), False, report_ambiguity=True
|
||||
)
|
||||
if saved and saved.get("association_status") != "ambiguous":
|
||||
summary = self.db.get_image_link_summaries([str(path)], False)
|
||||
title = summary[str(path)]["title"]
|
||||
except Exception:
|
||||
unavailable = True
|
||||
self.logger.warning("Saved inspection display could not be read")
|
||||
@@ -1726,11 +1760,12 @@ class ImageRoutesMixin:
|
||||
source_label=_public_path(roots[index]),
|
||||
saved=saved,
|
||||
library=self._inspection_library(
|
||||
saved, unavailable=unavailable
|
||||
saved, unavailable=unavailable, title=title
|
||||
),
|
||||
)
|
||||
except CurrentFileChanged:
|
||||
return None
|
||||
observed["library"]["title"] = _verified_title(saved, title, observed)
|
||||
return encode_inspection_response(
|
||||
image={
|
||||
"root_index": index,
|
||||
@@ -1763,7 +1798,10 @@ class ImageRoutesMixin:
|
||||
{"success": False, "error": "Prompt inspection failed"}, status=500
|
||||
)
|
||||
|
||||
def _inspection_library(self, saved: dict | None, *, unavailable: bool) -> dict:
|
||||
def _inspection_library(
|
||||
self, saved: dict | None, *, unavailable: bool, title: str | None = None
|
||||
) -> dict:
|
||||
"""The inspector's Library facts; ``title`` is the card's display title."""
|
||||
result = {
|
||||
"status": "unavailable" if unavailable else "none",
|
||||
"scope": "saved-library",
|
||||
@@ -1785,7 +1823,7 @@ class ImageRoutesMixin:
|
||||
status="linked",
|
||||
positive_id=pair.get("positive_prompt_id"),
|
||||
pair_id=pair.get("id"),
|
||||
title=pair.get("title"),
|
||||
title=title,
|
||||
category=pair.get("category"),
|
||||
tags=pair.get("tags"),
|
||||
metadata_scope="pair",
|
||||
@@ -1794,6 +1832,7 @@ class ImageRoutesMixin:
|
||||
result.update(
|
||||
status="linked",
|
||||
positive_id=saved.get("prompt_id"),
|
||||
title=title,
|
||||
category=saved.get("category"),
|
||||
tags=saved.get("tags"),
|
||||
metadata_scope="positive",
|
||||
@@ -1897,13 +1936,33 @@ class ImageRoutesMixin:
|
||||
if _in_subfolder(str(entry[0].relative_to(entry[2]).parent), subfolder)
|
||||
]
|
||||
|
||||
def _format_workspace_page(self, page: list[tuple]) -> list[dict]:
|
||||
def _workspace_link_summaries(self, paths: list[str]) -> dict[str, dict]:
|
||||
"""Card link summaries, matched exactly as the inspection route does."""
|
||||
try:
|
||||
return self.db.get_image_link_summaries(paths, False)
|
||||
except Exception:
|
||||
# The listing stays usable; cards fall back to their filenames.
|
||||
self.logger.warning("Workspace link summaries could not be read")
|
||||
return {}
|
||||
|
||||
def _format_workspace_page(
|
||||
self, page: list[tuple], *, linked: bool = True
|
||||
) -> list[dict]:
|
||||
"""Workspace descriptors; ``linked`` adds each file's link summary."""
|
||||
images = self._format_output_page(page, scoped=True)
|
||||
paths = {
|
||||
(entry[3], entry[0].relative_to(entry[2]).as_posix()): entry[0]
|
||||
for entry in page
|
||||
}
|
||||
summaries = (
|
||||
self._workspace_link_summaries([str(path) for path in paths.values()])
|
||||
if linked
|
||||
else {}
|
||||
)
|
||||
for image in images:
|
||||
if linked:
|
||||
key = str(paths[(image["root_index"], image["path"])])
|
||||
image["linked"] = summaries.get(key) or _unavailable_link()
|
||||
image["width"] = None
|
||||
image["height"] = None
|
||||
if image["is_video"]:
|
||||
|
||||
+8
-3
@@ -233,7 +233,9 @@ class LibraryRoutesMixin:
|
||||
self.add_library_negative
|
||||
)
|
||||
|
||||
def _library_media(self, records: list[dict], roots: list[Path]) -> list[dict]:
|
||||
def _library_media(
|
||||
self, records: list[dict], roots: list[Path], *, linked: bool = True
|
||||
) -> list[dict]:
|
||||
page = []
|
||||
for record in records:
|
||||
try:
|
||||
@@ -251,7 +253,7 @@ class LibraryRoutesMixin:
|
||||
break
|
||||
except (OSError, ValueError, RuntimeError):
|
||||
continue
|
||||
return self._format_workspace_page(page)
|
||||
return self._format_workspace_page(page, linked=linked)
|
||||
|
||||
@pair_response
|
||||
async def get_library_options(self, request: web.Request) -> web.Response:
|
||||
@@ -296,7 +298,10 @@ class LibraryRoutesMixin:
|
||||
|
||||
def previews() -> None:
|
||||
for parent in result["items"]:
|
||||
parent["previews"] = self._library_media(records[parent["id"]], roots)
|
||||
# Previews belong to their parent; a per-parent link read is N+1.
|
||||
parent["previews"] = self._library_media(
|
||||
records[parent["id"]], roots, linked=False
|
||||
)
|
||||
|
||||
await self._run_in_executor(previews)
|
||||
return web.json_response({"success": True, **result})
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
"""Opt-in workspace gallery contracts against disposable SQLite and media."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
try:
|
||||
from test_api_images import ImageAPITestCase, make_png
|
||||
from open_handles import assert_closed
|
||||
@@ -13,6 +16,24 @@ except ImportError:
|
||||
from tests.open_handles import assert_closed
|
||||
|
||||
|
||||
def embedded_png(path, positive, negative=""):
|
||||
"""A PNG whose API prompt graph resolves to *positive* and *negative*."""
|
||||
graph = {
|
||||
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
||||
},
|
||||
}
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("prompt", json.dumps(graph))
|
||||
with Image.new("RGB", (4, 4)) as image:
|
||||
image.save(path, pnginfo=info)
|
||||
os.utime(path, (100, 100))
|
||||
return path
|
||||
|
||||
|
||||
class TestWorkspaceImages(ImageAPITestCase):
|
||||
async def tearDownAsync(self):
|
||||
self.api.db.close_all()
|
||||
@@ -300,6 +321,156 @@ class TestWorkspaceImages(ImageAPITestCase):
|
||||
},
|
||||
)
|
||||
|
||||
def _link_fixtures(self):
|
||||
"""Linked (titled, first-line, pair, untitled pair), ambiguous and none."""
|
||||
db = self.api.db
|
||||
long_line = "lantern " * 15
|
||||
titled = self._save_prompt("Titled positive", category="Night", rating=5)
|
||||
with db.model.get_connection() as connection:
|
||||
connection.execute(
|
||||
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
||||
(titled, "Stored library title"),
|
||||
)
|
||||
plain = self._save_prompt(long_line)
|
||||
pair = db.save_prompt_pair("Pair positive", "Noise", title="Pair title")
|
||||
bare = db.save_prompt_pair("Bare pair positive\nmore", "Blur")
|
||||
first, second = self._save_prompt("Owner one"), self._save_prompt("Owner two")
|
||||
files = {
|
||||
name: self.media(f"{name}.png")
|
||||
for name in ("titled", "pair", "ambiguous", "none")
|
||||
}
|
||||
# First-line titles come from files whose own prompt evidence agrees.
|
||||
files["plain"] = embedded_png(self.output_dir / "plain.png", long_line)
|
||||
files["bare"] = embedded_png(
|
||||
self.output_dir / "bare.png", "Bare pair positive\nmore", "Blur"
|
||||
)
|
||||
db.link_image_to_prompt(titled, str(files["titled"]))
|
||||
db.link_image_to_prompt(plain, str(files["plain"]))
|
||||
for name, pair_id in (("pair", pair), ("bare", bare)):
|
||||
positive = db.get_prompt_pair(pair_id)["positive_prompt_id"]
|
||||
image = db.link_image_to_prompt(positive, str(files[name]))
|
||||
db.link_image_to_pair(pair_id, image, {"source": "fixture"})
|
||||
db.link_image_to_prompt(first, str(files["ambiguous"]))
|
||||
db.link_image_to_prompt(second, str(files["ambiguous"]))
|
||||
return {
|
||||
"titled": ("linked", titled, "Stored library title"),
|
||||
"plain": ("linked", plain, long_line.strip()[:80]),
|
||||
"pair": (
|
||||
"linked",
|
||||
db.get_prompt_pair(pair)["positive_prompt_id"],
|
||||
"Pair title",
|
||||
),
|
||||
"bare": (
|
||||
"linked",
|
||||
db.get_prompt_pair(bare)["positive_prompt_id"],
|
||||
"Bare pair positive",
|
||||
),
|
||||
"ambiguous": ("ambiguous", None, None),
|
||||
"none": ("none", None, None),
|
||||
}
|
||||
|
||||
async def test_workspace_items_carry_link_summary_matching_inspection(self):
|
||||
expected = self._link_fixtures()
|
||||
_, data = await self.get(limit=50)
|
||||
items = {Path(item["filename"]).stem: item for item in data["images"]}
|
||||
self.assertEqual(set(items), set(expected))
|
||||
for name, (status, positive_id, title) in expected.items():
|
||||
with self.subTest(file=name):
|
||||
linked = items[name]["linked"]
|
||||
self.assertEqual(
|
||||
set(linked),
|
||||
{"status", "positive_id", "title", "category", "rating"},
|
||||
)
|
||||
self.assertEqual(
|
||||
(linked["status"], linked["positive_id"], linked["title"]),
|
||||
(status, positive_id, title),
|
||||
)
|
||||
response = await self.client.get(
|
||||
f"/prompt_manager/images/inspect/{name}.png",
|
||||
params={"root": 0, "root_scope": items[name]["root_scope"]},
|
||||
)
|
||||
self.assertEqual(response.status, 200)
|
||||
library = (await response.json())["library"]
|
||||
self.assertEqual(
|
||||
(library["status"], library["positive_id"], library["title"]),
|
||||
(linked["status"], linked["positive_id"], linked["title"]),
|
||||
)
|
||||
self.assertEqual(
|
||||
(
|
||||
items["titled"]["linked"]["category"],
|
||||
items["titled"]["linked"]["rating"],
|
||||
),
|
||||
("Night", 5),
|
||||
)
|
||||
|
||||
async def test_unverified_first_line_title_stays_out_of_inspection(self):
|
||||
prompt = self._save_prompt("Secret first line\nrest")
|
||||
path = self.media("plain.png")
|
||||
self.api.db.link_image_to_prompt(prompt, str(path))
|
||||
_, data = await self.get()
|
||||
item = data["images"][0]
|
||||
self.assertEqual(item["linked"]["title"], "Secret first line")
|
||||
response = await self.client.get(
|
||||
"/prompt_manager/images/inspect/plain.png",
|
||||
params={"root": 0, "root_scope": item["root_scope"]},
|
||||
)
|
||||
body = await response.json()
|
||||
self.assertEqual(body["library"]["status"], "linked")
|
||||
self.assertEqual(body["library"]["positive_id"], prompt)
|
||||
self.assertIsNone(body["library"]["title"])
|
||||
self.assertNotIn("Secret first line", json.dumps(body))
|
||||
|
||||
async def test_link_summary_uses_one_query_per_page(self):
|
||||
prompt = self._save_prompt("Paged prompt")
|
||||
for index in range(12):
|
||||
path = self.media(f"{index:02}.png", mtime=100 + index)
|
||||
if index % 2:
|
||||
self.api.db.link_image_to_prompt(prompt, str(path))
|
||||
statements = []
|
||||
original = self.api.db.model.get_connection
|
||||
|
||||
def trace_connection():
|
||||
connection = original()
|
||||
connection.set_trace_callback(statements.append)
|
||||
return connection
|
||||
|
||||
with patch.object(
|
||||
self.api.db.model, "get_connection", side_effect=trace_connection
|
||||
):
|
||||
response, data = await self.get(limit=12)
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(len(data["images"]), 12)
|
||||
self.assertEqual(
|
||||
sorted(item["linked"]["status"] for item in data["images"]),
|
||||
["linked"] * 6 + ["none"] * 6,
|
||||
)
|
||||
selects = [
|
||||
query
|
||||
for query in statements
|
||||
if query.lstrip().upper().startswith(("SELECT", "WITH"))
|
||||
]
|
||||
self.assertEqual(len(selects), 1, selects)
|
||||
|
||||
async def test_link_summary_runs_in_executor_and_survives_database_errors(self):
|
||||
import threading
|
||||
|
||||
self.media("a.png")
|
||||
event_thread = threading.get_ident()
|
||||
threads = []
|
||||
|
||||
def broken(paths, allow_filename_fallback=True):
|
||||
threads.append(threading.get_ident())
|
||||
raise RuntimeError("database is gone")
|
||||
|
||||
with patch.object(self.api.db, "get_image_link_summaries", side_effect=broken):
|
||||
response, data = await self.get()
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue(threads)
|
||||
self.assertNotEqual(threads[0], event_thread)
|
||||
self.assertEqual(data["images"][0]["linked"]["status"], "unavailable")
|
||||
self.assertIsNone(data["images"][0]["linked"]["title"])
|
||||
self.assertNotIn("database is gone", str(data))
|
||||
|
||||
|
||||
class TestWorkspaceCursorEncoding(unittest.TestCase):
|
||||
def test_ntfs_unicode_basename_cursor_is_bounded_and_round_trips(self):
|
||||
|
||||
@@ -1432,3 +1432,151 @@ class TestImagePaging(DatabaseTestCase):
|
||||
self.assertEqual(len(second), 2)
|
||||
self.assertNotEqual({img["id"] for img in first}, {img["id"] for img in second})
|
||||
self.assertEqual(len(self.db.search_images_by_prompt("paged")), 5)
|
||||
|
||||
|
||||
class TestImageLinkSummaries(DatabaseTestCase):
|
||||
"""Batched card summaries follow get_image_prompt_info's matching rules."""
|
||||
|
||||
def _statements(self, call):
|
||||
statements = []
|
||||
original = self.db.model.get_connection
|
||||
|
||||
def traced():
|
||||
connection = original()
|
||||
connection.set_trace_callback(statements.append)
|
||||
return connection
|
||||
|
||||
try:
|
||||
with patch.object(self.db.model, "get_connection", side_effect=traced):
|
||||
result = call()
|
||||
finally:
|
||||
original().set_trace_callback(None)
|
||||
selects = [s for s in statements if s.lstrip().upper().startswith("SELECT")]
|
||||
return result, selects
|
||||
|
||||
def test_summary_status_matches_image_prompt_info_for_every_rule(self):
|
||||
owner = self._save("Owner prompt\nsecond line", category="Cat", rating=4)
|
||||
rival = self._save("Rival prompt")
|
||||
exact = os.path.join("out", "one", "same.png")
|
||||
other = os.path.join("out", "two", "same.png")
|
||||
unique = os.path.join("out", "three", "unique.png")
|
||||
shared = os.path.join("out", "shared.png")
|
||||
self.db.link_image_to_prompt(owner, exact)
|
||||
self.db.link_image_to_prompt(rival, other)
|
||||
self.db.link_image_to_prompt(owner, unique)
|
||||
self.db.link_image_to_prompt(owner, shared)
|
||||
self.db.link_image_to_prompt(rival, shared)
|
||||
probes = [
|
||||
exact,
|
||||
os.path.join("out", "x", "..", "one", "same.png"),
|
||||
os.path.join("elsewhere", "same.png"),
|
||||
os.path.join("elsewhere", "unique.png"),
|
||||
shared,
|
||||
"missing.png",
|
||||
]
|
||||
for fallback in (True, False):
|
||||
summaries = self.db.get_image_link_summaries(probes, fallback)
|
||||
self.assertEqual(set(summaries), set(probes))
|
||||
for probe in probes:
|
||||
with self.subTest(probe=probe, fallback=fallback):
|
||||
info = self.db.get_image_prompt_info(
|
||||
probe, fallback, report_ambiguity=True
|
||||
)
|
||||
summary = summaries[probe]
|
||||
if info is None:
|
||||
expected = "none"
|
||||
elif info.get("association_status") == "ambiguous":
|
||||
expected = "ambiguous"
|
||||
else:
|
||||
expected = "linked"
|
||||
self.assertEqual(summary["status"], expected)
|
||||
if expected == "linked":
|
||||
self.assertEqual(summary["positive_id"], info["prompt_id"])
|
||||
else:
|
||||
self.assertEqual(
|
||||
summary,
|
||||
{
|
||||
"status": expected,
|
||||
"positive_id": None,
|
||||
"title": None,
|
||||
"category": None,
|
||||
"rating": None,
|
||||
},
|
||||
)
|
||||
linked = self.db.get_image_link_summaries([unique])[unique]
|
||||
self.assertEqual(
|
||||
linked,
|
||||
{
|
||||
"status": "linked",
|
||||
"positive_id": owner,
|
||||
"title": "Owner prompt",
|
||||
"category": "Cat",
|
||||
"rating": 4,
|
||||
},
|
||||
)
|
||||
|
||||
def test_summary_title_prefers_stored_title_then_first_line(self):
|
||||
long_line = "word " * 30
|
||||
titled = self._save("Untitled text")
|
||||
plain = self._save(f" {long_line}\nsecond")
|
||||
blank = self._save("Blank title text")
|
||||
with self.db.model.get_connection() as conn:
|
||||
conn.executemany(
|
||||
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
||||
[(titled, "Stored title"), (blank, " \u3000 ")],
|
||||
)
|
||||
paths = {}
|
||||
for name, prompt in (("titled", titled), ("plain", plain), ("blank", blank)):
|
||||
paths[name] = os.path.join("out", f"{name}.png")
|
||||
self.db.link_image_to_prompt(prompt, paths[name])
|
||||
summaries = self.db.get_image_link_summaries(list(paths.values()))
|
||||
self.assertEqual(summaries[paths["titled"]]["title"], "Stored title")
|
||||
self.assertEqual(summaries[paths["plain"]]["title"], long_line.strip()[:80])
|
||||
self.assertLessEqual(len(summaries[paths["plain"]]["title"]), 80)
|
||||
self.assertEqual(summaries[paths["blank"]]["title"], "Blank title text")
|
||||
|
||||
def test_summary_uses_one_select_per_200_paths(self):
|
||||
prompt = self._save("Batch")
|
||||
paths = [os.path.join("out", f"{index}.png") for index in range(201)]
|
||||
for path in paths[:5]:
|
||||
self.db.link_image_to_prompt(prompt, path)
|
||||
result, selects = self._statements(
|
||||
lambda: self.db.get_image_link_summaries(paths[:200])
|
||||
)
|
||||
self.assertEqual(len(selects), 1)
|
||||
self.assertEqual(len(result), 200)
|
||||
self.assertEqual(
|
||||
[result[path]["status"] for path in paths[:6]], ["linked"] * 5 + ["none"]
|
||||
)
|
||||
_, selects = self._statements(lambda: self.db.get_image_link_summaries(paths))
|
||||
self.assertEqual(len(selects), 2)
|
||||
_, selects = self._statements(lambda: self.db.get_image_link_summaries([]))
|
||||
self.assertEqual(selects, [])
|
||||
|
||||
def test_windows_style_paths_normalize(self):
|
||||
import ntpath
|
||||
|
||||
def windows_normalize(path):
|
||||
return ntpath.normcase(ntpath.normpath(path))
|
||||
|
||||
prompt = self._save("Windows prompt")
|
||||
stored = r"C:\ComfyUI\Output\Sub\Image.PNG"
|
||||
with patch("database.operations.normalize_image_path", windows_normalize):
|
||||
self.db.link_image_to_prompt(prompt, stored)
|
||||
probes = [
|
||||
stored,
|
||||
r"c:\comfyui\output\SUB\image.png",
|
||||
r"C:\ComfyUI\Output\x\..\Sub\Image.PNG",
|
||||
"C:/ComfyUI/Output/Sub/Image.PNG",
|
||||
r"D:\ComfyUI\Output\Sub\Other.png",
|
||||
]
|
||||
summaries = self.db.get_image_link_summaries(probes, False)
|
||||
for probe in probes:
|
||||
with self.subTest(probe=probe):
|
||||
info = self.db.get_image_prompt_info(probe, False)
|
||||
expected = "linked" if info else "none"
|
||||
self.assertEqual(summaries[probe]["status"], expected)
|
||||
self.assertEqual(
|
||||
[summaries[probe]["status"] for probe in probes],
|
||||
["linked"] * 4 + ["none"],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user