138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
"""Owned, deterministic Library data for the native recipe editor workflows."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import shutil
|
|
from typing import TYPE_CHECKING
|
|
|
|
from aiohttp import web
|
|
from PIL import Image
|
|
|
|
if TYPE_CHECKING:
|
|
from py.api import PromptManagerAPI
|
|
|
|
|
|
EDITOR_POSITIVE = "Editor portrait, soft window light"
|
|
FIRST_NEGATIVE = "blurry, watermark"
|
|
SELECTED_NEGATIVE = "low contrast, muddy shadows"
|
|
ARCHIVED_NEGATIVE = "oversaturated, harsh flash"
|
|
STANDALONE = "Standalone sketch without saved negatives"
|
|
EDITOR_METADATA = {
|
|
"title": "Editor portrait",
|
|
"category": "Portraits",
|
|
"tags": ["studio", "window light"],
|
|
"rating": 4,
|
|
"notes": "Keep the soft light.",
|
|
}
|
|
_TABLES = (
|
|
"image_favorites",
|
|
"image_prompt_pairs",
|
|
"prompt_pair_tags",
|
|
"prompt_pairs",
|
|
"generated_images",
|
|
"prompt_tags",
|
|
"tags",
|
|
"prompt_library_state",
|
|
"prompts",
|
|
)
|
|
|
|
|
|
def seed_library_editor(api: PromptManagerAPI, root: Path) -> dict:
|
|
"""Replace this server's disposable rows with one fixed editor Library."""
|
|
db = api.db
|
|
with db.model.get_connection() as conn:
|
|
for table in _TABLES:
|
|
conn.execute(f"DELETE FROM {table}")
|
|
for name in ("output", "input"):
|
|
shutil.rmtree(root / name)
|
|
(root / name).mkdir()
|
|
epoch = db.get_execution_epoch()
|
|
receipt = db.create_library_recipe(
|
|
EDITOR_POSITIVE,
|
|
FIRST_NEGATIVE,
|
|
expected_epoch=epoch,
|
|
parent_metadata=dict(EDITOR_METADATA),
|
|
)
|
|
positive_id = receipt["positive"]["id"]
|
|
first_pair_id = receipt["pair"]["id"]
|
|
pairs = {}
|
|
for name, text in (
|
|
("pair_id", SELECTED_NEGATIVE),
|
|
("empty_pair_id", ""),
|
|
("archived_pair_id", ARCHIVED_NEGATIVE),
|
|
):
|
|
receipt = db.add_library_negative(
|
|
positive_id,
|
|
text,
|
|
expected_epoch=epoch,
|
|
if_match=receipt["positive"]["edit_token"],
|
|
)
|
|
pairs[name] = receipt["pair"]["id"]
|
|
db.archive_prompt_pair(pairs["archived_pair_id"], True)
|
|
image_path = root / "output" / "editor" / "portrait.png"
|
|
image_path.parent.mkdir()
|
|
with Image.new("RGB", (64, 64), (40, 90, 120)) as image:
|
|
image.save(image_path)
|
|
image_id = db.link_image_to_prompt(positive_id, str(image_path))
|
|
db.link_image_to_pair(pairs["pair_id"], image_id, {"source": "editor fixture"})
|
|
db.record_prompt_pair_use(pairs["pair_id"], times=2)
|
|
standalone_id = db.save_prompt(STANDALONE, category="Sketches", tags=["draft"])
|
|
with db.model.get_connection() as conn:
|
|
negative_only_id = conn.execute(
|
|
"SELECT negative_prompt_id FROM prompt_pairs WHERE id=?",
|
|
(first_pair_id,),
|
|
).fetchone()[0]
|
|
api.invalidate_gallery_cache()
|
|
return {
|
|
"positive_id": positive_id,
|
|
"first_pair_id": first_pair_id,
|
|
**pairs,
|
|
"standalone_id": standalone_id,
|
|
"negative_only_id": negative_only_id,
|
|
"positive_text": EDITOR_POSITIVE,
|
|
"first_negative": FIRST_NEGATIVE,
|
|
"selected_negative": SELECTED_NEGATIVE,
|
|
"archived_negative": ARCHIVED_NEGATIVE,
|
|
"metadata": dict(EDITOR_METADATA),
|
|
}
|
|
|
|
|
|
class _LibraryEditorCommands:
|
|
"""One disposable fixture lifetime; an ordinary reset expires it."""
|
|
|
|
def __init__(self) -> None:
|
|
self.manifest: dict | None = None
|
|
|
|
def reset(self) -> None:
|
|
self.manifest = None
|
|
|
|
|
|
def register_library_editor_commands(
|
|
routes: web.RouteTableDef, api: PromptManagerAPI, root: Path
|
|
) -> _LibraryEditorCommands:
|
|
"""Register only fixed seed and snapshot commands; no paths or SQL input."""
|
|
state = _LibraryEditorCommands()
|
|
|
|
@routes.post("/__test__/library-editor/reset")
|
|
async def reset(request: web.Request) -> web.Response:
|
|
data = await request.json() if request.can_read_body else {}
|
|
if data != {} or request.query:
|
|
raise web.HTTPBadRequest()
|
|
state.manifest = await api._run_in_executor(seed_library_editor, api, root)
|
|
return web.json_response({"success": True, **state.manifest})
|
|
|
|
@routes.get("/__test__/library-editor/snapshot")
|
|
async def snapshot(request: web.Request) -> web.Response:
|
|
if state.manifest is None or request.query:
|
|
raise web.HTTPConflict()
|
|
path = root / "owned-library-editor-snapshot.db"
|
|
created = await api._run_in_executor(api.db.model.backup_database, str(path))
|
|
if not created:
|
|
raise web.HTTPInternalServerError()
|
|
return web.Response(
|
|
body=path.read_bytes(), content_type="application/octet-stream"
|
|
)
|
|
|
|
return state
|