319 lines
12 KiB
Python
319 lines
12 KiB
Python
"""Contained Recipe image assets and bounded output-gallery inspection."""
|
|
|
|
from io import BytesIO
|
|
import heapq
|
|
import json
|
|
import os
|
|
|
|
import folder_paths
|
|
from PIL import Image
|
|
|
|
from .recipe_constants import *
|
|
from .utils import require_filename, resolve_within
|
|
from .workflow_schema import _node_type, _parameter_signature, _workflow_node_signature
|
|
|
|
def _recipe_assets_dir(recipes_dir, filename, create=False):
|
|
stem = os.path.splitext(require_filename(filename))[0]
|
|
assets_dir = resolve_within(recipes_dir, ".assets", stem)
|
|
if create:
|
|
os.makedirs(assets_dir, exist_ok=True)
|
|
return assets_dir
|
|
|
|
|
|
def _preview_source_path(saved_value):
|
|
resolved = _resolve_exact_model_reference(saved_value)
|
|
if not resolved:
|
|
return None
|
|
base_path = os.path.splitext(resolved["path"])[0]
|
|
for suffix in STATIC_PREVIEW_SUFFIXES + STATIC_PREVIEW_EXTENSIONS:
|
|
candidate = f"{base_path}{suffix}"
|
|
if os.path.isfile(candidate):
|
|
return candidate
|
|
return None
|
|
|
|
|
|
def _thumbnail_webp_bytes(source_path):
|
|
try:
|
|
if os.path.getsize(source_path) > MAX_PREVIEW_SOURCE_BYTES:
|
|
return None
|
|
with Image.open(source_path) as image:
|
|
image.thumbnail((320, 320))
|
|
if image.mode not in ("RGB", "RGBA"):
|
|
image = image.convert("RGBA" if "A" in image.getbands() else "RGB")
|
|
width, height = image.size
|
|
for quality in (76, 62, 48):
|
|
output = BytesIO()
|
|
image.save(output, format="WEBP", quality=quality, method=4)
|
|
data = output.getvalue()
|
|
if len(data) <= MAX_PREVIEW_SNAPSHOT_BYTES:
|
|
return data, width, height
|
|
except (OSError, ValueError):
|
|
return None
|
|
return None
|
|
|
|
|
|
def _recipe_cover_webp_bytes(source_path):
|
|
"""Create a portable recipe-cover asset without retaining the full output image."""
|
|
try:
|
|
if os.path.getsize(source_path) > MAX_RECIPE_COVER_SOURCE_BYTES:
|
|
return None
|
|
with Image.open(source_path) as image:
|
|
image.thumbnail((640, 640))
|
|
if image.mode not in ("RGB", "RGBA"):
|
|
image = image.convert("RGBA" if "A" in image.getbands() else "RGB")
|
|
width, height = image.size
|
|
for quality in (84, 72, 60, 48):
|
|
output = BytesIO()
|
|
image.save(output, format="WEBP", quality=quality, method=4)
|
|
data = output.getvalue()
|
|
if len(data) <= MAX_RECIPE_COVER_BYTES:
|
|
return data, width, height
|
|
except (OSError, ValueError):
|
|
return None
|
|
return None
|
|
|
|
|
|
def _output_source_path(source_image):
|
|
output_dir = folder_paths.get_output_directory()
|
|
target_dir = resolve_within(output_dir, source_image.get("subfolder", ""))
|
|
source_path = resolve_within(target_dir, source_image["filename"])
|
|
if not os.path.isfile(source_path):
|
|
raise FileNotFoundError
|
|
return source_path
|
|
|
|
|
|
def _store_recipe_gallery_cover(recipes_dir, filename, source_image):
|
|
source_path = _output_source_path(source_image)
|
|
cover = _recipe_cover_webp_bytes(source_path)
|
|
if not cover:
|
|
raise ValueError("Could not create recipe cover")
|
|
data, width, height = cover
|
|
asset_id = f"cover-{hashlib.sha256(data).hexdigest()}.webp"
|
|
asset_path = resolve_within(_recipe_assets_dir(recipes_dir, filename, create=True), asset_id)
|
|
if not os.path.exists(asset_path):
|
|
with open(asset_path, "wb") as asset_file:
|
|
asset_file.write(data)
|
|
return {
|
|
"asset_id": asset_id,
|
|
"media_type": "image/webp",
|
|
"width": width,
|
|
"height": height,
|
|
}
|
|
|
|
|
|
def _decode_embedded_json(value):
|
|
if isinstance(value, bytes):
|
|
if len(value) > MAX_EMBEDDED_WORKFLOW_BYTES:
|
|
return None
|
|
value = value.decode("utf-8")
|
|
if not isinstance(value, str) or len(value.encode("utf-8")) > MAX_EMBEDDED_WORKFLOW_BYTES:
|
|
return None
|
|
parsed = json.loads(value)
|
|
return parsed if isinstance(parsed, dict) else None
|
|
|
|
|
|
def _embedded_workflow_payload(image_path):
|
|
"""Read bounded workflow/prompt metadata from one PNG without loading pixels."""
|
|
try:
|
|
with Image.open(image_path) as image:
|
|
return {
|
|
"workflow": _decode_embedded_json(image.info.get("workflow")),
|
|
"prompt": _decode_embedded_json(image.info.get("prompt")),
|
|
}
|
|
except (OSError, UnicodeDecodeError, ValueError, json.JSONDecodeError):
|
|
return None
|
|
|
|
|
|
def _embedded_workflow_node_signature(image_path):
|
|
payload = _embedded_workflow_payload(image_path)
|
|
if not payload:
|
|
return None
|
|
workflow = payload.get("workflow") or payload.get("prompt")
|
|
if not workflow:
|
|
return None
|
|
return _workflow_node_signature(workflow)["value"]
|
|
|
|
|
|
def _embedded_parameter_signature(image_path):
|
|
payload = _embedded_workflow_payload(image_path)
|
|
if not payload:
|
|
return None
|
|
workflow = payload.get("workflow") or payload.get("prompt")
|
|
if not workflow:
|
|
return None
|
|
return _parameter_signature(workflow)["value"]
|
|
|
|
|
|
def _recent_output_pngs(output_dir, limit=MAX_RECIPE_GALLERY_SCAN):
|
|
"""Return a bounded newest-first PNG list without trusting request paths."""
|
|
newest = []
|
|
try:
|
|
for root, _, files in os.walk(output_dir):
|
|
for name in files:
|
|
if not name.lower().endswith(".png"):
|
|
continue
|
|
path = os.path.join(root, name)
|
|
try:
|
|
stat = os.stat(path)
|
|
except OSError:
|
|
continue
|
|
candidate = (stat.st_mtime_ns, path)
|
|
if len(newest) < limit:
|
|
heapq.heappush(newest, candidate)
|
|
elif candidate > newest[0]:
|
|
heapq.heapreplace(newest, candidate)
|
|
except OSError:
|
|
return []
|
|
return sorted(newest, reverse=True)
|
|
|
|
|
|
def _recipe_gallery_images(node_signature):
|
|
output_dir = folder_paths.get_output_directory()
|
|
if not os.path.isdir(output_dir):
|
|
return [], 0
|
|
matches = []
|
|
scanned = 0
|
|
for mtime_ns, image_path in _recent_output_pngs(output_dir):
|
|
scanned += 1
|
|
if _embedded_workflow_node_signature(image_path) != node_signature:
|
|
continue
|
|
relative_dir = os.path.relpath(os.path.dirname(image_path), output_dir)
|
|
subfolder = "" if relative_dir == "." else relative_dir.replace("\\", "/")
|
|
matches.append({
|
|
"filename": os.path.basename(image_path),
|
|
"subfolder": subfolder,
|
|
"type": "output",
|
|
"mtime": mtime_ns // 1_000_000,
|
|
})
|
|
if len(matches) >= MAX_RECIPE_GALLERY_RESULTS:
|
|
break
|
|
return matches, scanned
|
|
|
|
|
|
def _parameter_gallery_images(parameter_signature):
|
|
output_dir = folder_paths.get_output_directory()
|
|
if not os.path.isdir(output_dir):
|
|
return [], 0
|
|
matches = []
|
|
scanned = 0
|
|
for mtime_ns, image_path in _recent_output_pngs(output_dir):
|
|
scanned += 1
|
|
if _embedded_parameter_signature(image_path) != parameter_signature:
|
|
continue
|
|
relative_dir = os.path.relpath(os.path.dirname(image_path), output_dir)
|
|
subfolder = "" if relative_dir == "." else relative_dir.replace("\\", "/")
|
|
matches.append({
|
|
"filename": os.path.basename(image_path),
|
|
"subfolder": subfolder,
|
|
"type": "output",
|
|
"mtime": mtime_ns // 1_000_000,
|
|
})
|
|
if len(matches) >= MAX_RECIPE_GALLERY_RESULTS:
|
|
break
|
|
return matches, scanned
|
|
|
|
|
|
def _bounded_gallery_value(value, depth=0):
|
|
"""Keep image-embedded parameter details useful without echoing huge metadata."""
|
|
if depth > 5:
|
|
return "…"
|
|
if isinstance(value, str):
|
|
return value if len(value) <= 2000 else f"{value[:1997]}..."
|
|
if isinstance(value, (int, float, bool)) or value is None:
|
|
return value
|
|
if isinstance(value, list):
|
|
return [_bounded_gallery_value(item, depth + 1) for item in value[:64]]
|
|
if isinstance(value, dict):
|
|
return {
|
|
str(key): _bounded_gallery_value(nested, depth + 1)
|
|
for key, nested in list(value.items())[:96]
|
|
}
|
|
return str(value)[:2000]
|
|
|
|
|
|
def _gallery_parameter_records(workflow):
|
|
if not isinstance(workflow, dict):
|
|
return []
|
|
nodes = workflow.get("nodes")
|
|
is_ui_workflow = isinstance(nodes, list)
|
|
values = nodes if is_ui_workflow else workflow.values()
|
|
records = []
|
|
occurrence = {}
|
|
for node in values:
|
|
if not isinstance(node, dict):
|
|
continue
|
|
node_type = _node_type(node) or str(node.get("class_type") or "").strip()
|
|
if not node_type:
|
|
continue
|
|
occurrence[node_type] = occurrence.get(node_type, 0) + 1
|
|
if is_ui_workflow:
|
|
parameters = {"widgets": _bounded_gallery_value(node.get("widgets_values", []))}
|
|
node_id = node.get("id")
|
|
else:
|
|
parameters = {"inputs": _bounded_gallery_value(node.get("inputs", {}))}
|
|
node_id = None
|
|
records.append({
|
|
"type": node_type,
|
|
"index": occurrence[node_type],
|
|
"node_id": node_id,
|
|
"parameters": parameters,
|
|
})
|
|
return records
|
|
|
|
|
|
def _gallery_parameter_diff(recipe_workflow, image_payload):
|
|
recipe_records = _gallery_parameter_records(recipe_workflow)
|
|
embedded_workflow = image_payload.get("workflow") or image_payload.get("prompt") if image_payload else None
|
|
image_records = _gallery_parameter_records(embedded_workflow)
|
|
recipe_map = {(item["type"], item["index"]): item for item in recipe_records}
|
|
image_map = {(item["type"], item["index"]): item for item in image_records}
|
|
changes = []
|
|
for key in sorted(set(recipe_map) | set(image_map), key=lambda item: (item[0].casefold(), item[1])):
|
|
recipe_item = recipe_map.get(key)
|
|
image_item = image_map.get(key)
|
|
if recipe_item is None or image_item is None or recipe_item["parameters"] != image_item["parameters"]:
|
|
changes.append({
|
|
"type": key[0],
|
|
"index": key[1],
|
|
"recipe": recipe_item["parameters"] if recipe_item else None,
|
|
"image": image_item["parameters"] if image_item else None,
|
|
})
|
|
return {
|
|
"recipe_nodes": recipe_records,
|
|
"image_nodes": image_records,
|
|
"changes": changes,
|
|
"comparable": bool(image_records) and all(
|
|
isinstance(item.get("workflow"), dict) for item in [image_payload or {}]
|
|
),
|
|
}
|
|
|
|
|
|
def _attach_preview_snapshots(recipe, recipes_dir, filename, references):
|
|
if not recipe.get("presentation", {}).get("save_model_preview_snapshots"):
|
|
return
|
|
assets_dir = _recipe_assets_dir(recipes_dir, filename, create=True)
|
|
total_bytes = 0
|
|
for reference in references[:MAX_PREVIEW_SNAPSHOTS]:
|
|
source_path = _preview_source_path(reference.get("saved_value"))
|
|
if not source_path:
|
|
continue
|
|
thumbnail = _thumbnail_webp_bytes(source_path)
|
|
if not thumbnail:
|
|
continue
|
|
data, width, height = thumbnail
|
|
if total_bytes + len(data) > MAX_PREVIEW_SNAPSHOT_TOTAL_BYTES:
|
|
break
|
|
asset_id = f"{hashlib.sha256(data).hexdigest()}.webp"
|
|
asset_path = resolve_within(assets_dir, asset_id)
|
|
if not os.path.exists(asset_path):
|
|
with open(asset_path, "wb") as asset_file:
|
|
asset_file.write(data)
|
|
total_bytes += len(data)
|
|
reference["preview"] = {
|
|
"snapshot_asset_id": asset_id,
|
|
"media_type": "image/webp",
|
|
"width": width,
|
|
"height": height,
|
|
"captured_at": recipe["timestamp"],
|
|
}
|