409 lines
16 KiB
Python
409 lines
16 KiB
Python
"""Curated material snapshots built from generated ComfyUI images.
|
|
|
|
The first material type intentionally stays small: one output image, its exact
|
|
embedded UI workflow, and node-sized parameter blocks derived from that
|
|
workflow. Personal material data lives below ComfyUI's user directory.
|
|
"""
|
|
|
|
import asyncio
|
|
import copy
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import tempfile
|
|
import time
|
|
import uuid
|
|
|
|
from aiohttp import web
|
|
import folder_paths
|
|
|
|
from .recipes import (
|
|
MAX_RECIPE_BYTES,
|
|
_build_model_references,
|
|
_embedded_workflow_payload,
|
|
_normalise_source_image,
|
|
_output_source_path,
|
|
_recipe_cover_webp_bytes,
|
|
_validate_workflow,
|
|
_volatile_widget_indexes,
|
|
)
|
|
from .utils import require_filename, resolve_within
|
|
|
|
|
|
MATERIAL_SCHEMA_VERSION = 1
|
|
MAX_MATERIAL_NAME_LENGTH = 120
|
|
MAX_SOURCE_IMAGE_BYTES = 64 * 1024 * 1024
|
|
MODEL_PRIORITY = {"checkpoint": 0, "unet": 1, "lora": 2}
|
|
|
|
|
|
def get_materials_dir():
|
|
user_dir = (
|
|
folder_paths.get_user_directory()
|
|
if hasattr(folder_paths, "get_user_directory")
|
|
else None
|
|
)
|
|
if not user_dir:
|
|
user_dir = os.path.join(folder_paths.base_path, "user", "default")
|
|
materials_dir = os.path.join(user_dir, "workflows", "anomalous_materials")
|
|
os.makedirs(materials_dir, exist_ok=True)
|
|
return materials_dir
|
|
|
|
|
|
def _material_assets_dir(materials_dir, filename, create=False):
|
|
stem = os.path.splitext(require_filename(filename))[0]
|
|
assets_dir = resolve_within(materials_dir, ".assets", stem)
|
|
if create:
|
|
os.makedirs(assets_dir, exist_ok=True)
|
|
return assets_dir
|
|
|
|
|
|
def _atomic_write_json(path, value):
|
|
encoded = json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
|
if len(encoded) > MAX_RECIPE_BYTES:
|
|
raise ValueError("Material snapshot is too large")
|
|
target_dir = os.path.dirname(path)
|
|
fd, temp_path = tempfile.mkstemp(prefix=".material-", suffix=".tmp", dir=target_dir)
|
|
try:
|
|
with os.fdopen(fd, "wb") as output:
|
|
output.write(encoded)
|
|
output.flush()
|
|
os.fsync(output.fileno())
|
|
os.replace(temp_path, path)
|
|
finally:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
|
|
|
|
def _read_material(path):
|
|
if os.path.getsize(path) > MAX_RECIPE_BYTES:
|
|
raise ValueError("Material snapshot is too large")
|
|
with open(path, "r", encoding="utf-8") as material_file:
|
|
value = json.load(material_file)
|
|
if not isinstance(value, dict) or not isinstance(value.get("workflow"), dict):
|
|
raise ValueError("Invalid material snapshot")
|
|
return value
|
|
|
|
|
|
def _node_blocks(workflow, include_values=True):
|
|
blocks = []
|
|
occurrences = {}
|
|
for node in workflow.get("nodes", []):
|
|
if not isinstance(node, dict):
|
|
continue
|
|
node_type = str(node.get("type") or "").strip()
|
|
if not node_type:
|
|
continue
|
|
occurrences[node_type] = occurrences.get(node_type, 0) + 1
|
|
block = {
|
|
"node_id": node.get("id"),
|
|
"type": node_type,
|
|
"title": str(node.get("title") or node_type),
|
|
"occurrence": occurrences[node_type],
|
|
"widget_count": len(node.get("widgets_values") or []),
|
|
"volatile_widget_indexes": list(_volatile_widget_indexes(node)),
|
|
}
|
|
if include_values:
|
|
block["widgets_values"] = copy.deepcopy(node.get("widgets_values") or [])
|
|
blocks.append(block)
|
|
return blocks
|
|
|
|
|
|
def _prompt_excerpt(workflow):
|
|
for node in workflow.get("nodes", []):
|
|
if not isinstance(node, dict) or "cliptextencode" not in str(node.get("type") or "").lower():
|
|
continue
|
|
for value in node.get("widgets_values") or []:
|
|
if isinstance(value, str) and value.strip():
|
|
compact = re.sub(r"\s+", " ", value).strip()
|
|
return compact[:20]
|
|
return ""
|
|
|
|
|
|
def _display_model_name(reference):
|
|
value = str(reference.get("saved_value") or "").replace("\\", "/").split("/")[-1]
|
|
return re.sub(r"\.(?:safetensors|ckpt|pt|bin|sft)$", "", value, flags=re.IGNORECASE)
|
|
|
|
|
|
def _suggested_name(workflow, source_path, references=None):
|
|
if references is None:
|
|
references = _build_model_references({"workflow": workflow, "params": {}}, verify_identities=False)
|
|
references.sort(key=lambda item: MODEL_PRIORITY.get(item.get("category"), 99))
|
|
model_name = _display_model_name(references[0]) if references else "工作流快照"
|
|
excerpt = _prompt_excerpt(workflow)
|
|
date = time.strftime("%Y-%m-%d", time.localtime(os.path.getmtime(source_path)))
|
|
return " · ".join(part for part in (model_name, excerpt, date) if part)[:MAX_MATERIAL_NAME_LENGTH]
|
|
|
|
|
|
def _inspect_source_image(source_image):
|
|
source = _normalise_source_image(source_image)
|
|
if not source["filename"].lower().endswith(".png"):
|
|
raise ValueError("Only PNG output images can contain reusable workflow metadata")
|
|
source_path = _output_source_path(source)
|
|
payload = _embedded_workflow_payload(source_path)
|
|
workflow = payload.get("workflow") if isinstance(payload, dict) else None
|
|
if not isinstance(workflow, dict):
|
|
raise ValueError("Image has no reusable UI workflow")
|
|
_validate_workflow(workflow)
|
|
blocks = _node_blocks(workflow, include_values=False)
|
|
references = _build_model_references({"workflow": workflow, "params": {}}, verify_identities=False)
|
|
suggested_name = _suggested_name(workflow, source_path, references)
|
|
return source, source_path, workflow, blocks, references, suggested_name
|
|
|
|
|
|
def _store_assets(materials_dir, filename, source_path):
|
|
source_size = os.path.getsize(source_path)
|
|
if source_size > MAX_SOURCE_IMAGE_BYTES:
|
|
raise ValueError("Source image is too large")
|
|
digest = hashlib.sha256()
|
|
with open(source_path, "rb") as source_file:
|
|
for block in iter(lambda: source_file.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
assets_dir = _material_assets_dir(materials_dir, filename, create=True)
|
|
source_asset = f"source-{digest.hexdigest()}.png"
|
|
source_target = resolve_within(assets_dir, source_asset)
|
|
if not os.path.exists(source_target):
|
|
fd, temp_path = tempfile.mkstemp(prefix=".source-", suffix=".tmp", dir=assets_dir)
|
|
os.close(fd)
|
|
try:
|
|
shutil.copyfile(source_path, temp_path)
|
|
os.replace(temp_path, source_target)
|
|
finally:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
|
|
preview = _recipe_cover_webp_bytes(source_path)
|
|
preview_asset = None
|
|
if preview:
|
|
data, width, height = preview
|
|
preview_asset = f"preview-{hashlib.sha256(data).hexdigest()}.webp"
|
|
preview_target = resolve_within(assets_dir, preview_asset)
|
|
if not os.path.exists(preview_target):
|
|
with open(preview_target, "wb") as preview_file:
|
|
preview_file.write(data)
|
|
else:
|
|
width = height = None
|
|
return {
|
|
"source_asset_id": source_asset,
|
|
"preview_asset_id": preview_asset,
|
|
"preview_width": width,
|
|
"preview_height": height,
|
|
"source_sha256": digest.hexdigest(),
|
|
"source_size": source_size,
|
|
}
|
|
|
|
|
|
def _material_summary(filename, material):
|
|
workflow = material.get("workflow") or {}
|
|
image = material.get("image") or {}
|
|
return {
|
|
"filename": filename,
|
|
"id": material.get("id"),
|
|
"name": material.get("name") or "未命名素材",
|
|
"kind": material.get("kind"),
|
|
"timestamp": material.get("timestamp", 0),
|
|
"node_count": len(workflow.get("nodes") or []),
|
|
"node_types": sorted({block["type"] for block in _node_blocks(workflow, include_values=False)}),
|
|
"image": {
|
|
"preview_asset_id": image.get("preview_asset_id"),
|
|
"source_asset_id": image.get("source_asset_id"),
|
|
"width": image.get("preview_width"),
|
|
"height": image.get("preview_height"),
|
|
},
|
|
"capabilities": list(material.get("capabilities") or []),
|
|
}
|
|
|
|
|
|
def _workflow_hashes_for_blocks(workflow, blocks):
|
|
hashes = (workflow.get("extra") or {}).get("anomalous_hashes") if isinstance(workflow, dict) else None
|
|
if not isinstance(hashes, dict):
|
|
return {}
|
|
prefixes = tuple(f"{block.get('node_id')}_" for block in blocks)
|
|
return {
|
|
key: copy.deepcopy(value)
|
|
for key, value in hashes.items()
|
|
if isinstance(key, str) and key.startswith(prefixes)
|
|
}
|
|
|
|
|
|
def _list_materials(materials_dir):
|
|
result = []
|
|
try:
|
|
entries = os.scandir(materials_dir)
|
|
except OSError:
|
|
return result
|
|
with entries:
|
|
for entry in entries:
|
|
if not entry.is_file() or not entry.name.endswith(".json"):
|
|
continue
|
|
try:
|
|
result.append(_material_summary(entry.name, _read_material(entry.path)))
|
|
except (OSError, ValueError, json.JSONDecodeError):
|
|
continue
|
|
result.sort(key=lambda item: item.get("timestamp", 0), reverse=True)
|
|
return result
|
|
|
|
|
|
async def api_inspect_image_material(request):
|
|
try:
|
|
payload = await request.json()
|
|
source, source_path, workflow, blocks, references, suggested_name = await asyncio.to_thread(
|
|
_inspect_source_image, payload.get("source_image")
|
|
)
|
|
except (AttributeError, ValueError, json.JSONDecodeError):
|
|
return web.json_response({"status": "error", "message": "Image has no reusable UI workflow"}, status=400)
|
|
except FileNotFoundError:
|
|
return web.json_response({"status": "error", "message": "Output image not found"}, status=404)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not inspect output image"}, status=500)
|
|
return web.json_response({
|
|
"status": "success",
|
|
"source_image": source,
|
|
"suggested_name": suggested_name,
|
|
"node_count": len(blocks),
|
|
"node_blocks": blocks,
|
|
"model_references": references,
|
|
"prompt_excerpt": _prompt_excerpt(workflow),
|
|
})
|
|
|
|
|
|
async def api_save_image_material(request):
|
|
try:
|
|
payload = await request.json()
|
|
source, source_path, workflow, blocks, references, suggested_name = await asyncio.to_thread(
|
|
_inspect_source_image, payload.get("source_image")
|
|
)
|
|
name = payload.get("name") or suggested_name
|
|
if not isinstance(name, str) or not (name := name.strip()) or len(name) > MAX_MATERIAL_NAME_LENGTH:
|
|
raise ValueError("Invalid material name")
|
|
material_id = uuid.uuid4().hex
|
|
filename = f"material_{int(time.time())}_{material_id[:8]}.json"
|
|
materials_dir = get_materials_dir()
|
|
image = await asyncio.to_thread(_store_assets, materials_dir, filename, source_path)
|
|
now = int(time.time() * 1000)
|
|
material = {
|
|
"schema_version": MATERIAL_SCHEMA_VERSION,
|
|
"id": material_id,
|
|
"kind": "image_workflow_snapshot",
|
|
"name": name,
|
|
"timestamp": now,
|
|
"source": {"type": "generated_image", "image": source},
|
|
"image": image,
|
|
"workflow": workflow,
|
|
"model_references": references,
|
|
"capabilities": ["open_workflow", "apply_node_parameters", "reference_image"],
|
|
}
|
|
await asyncio.to_thread(_atomic_write_json, resolve_within(materials_dir, filename), material)
|
|
except (AttributeError, ValueError, json.JSONDecodeError):
|
|
return web.json_response({"status": "error", "message": "Could not save image material"}, status=400)
|
|
except FileNotFoundError:
|
|
return web.json_response({"status": "error", "message": "Output image not found"}, status=404)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not save image material"}, status=500)
|
|
return web.json_response({
|
|
"status": "success",
|
|
"filename": filename,
|
|
"material": _material_summary(filename, material),
|
|
"node_blocks": blocks,
|
|
})
|
|
|
|
|
|
async def api_get_materials(request):
|
|
materials = await asyncio.to_thread(_list_materials, get_materials_dir())
|
|
return web.json_response({"status": "success", "materials": materials})
|
|
|
|
|
|
async def api_get_material_full(request):
|
|
try:
|
|
filename = require_filename(request.query.get("filename", ""))
|
|
if not filename.endswith(".json"):
|
|
raise ValueError("Invalid material filename")
|
|
material = await asyncio.to_thread(_read_material, resolve_within(get_materials_dir(), filename))
|
|
except (AttributeError, ValueError, json.JSONDecodeError):
|
|
return web.json_response({"status": "error", "message": "Invalid material"}, status=400)
|
|
except FileNotFoundError:
|
|
return web.json_response({"status": "error", "message": "Material not found"}, status=404)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not read material"}, status=500)
|
|
return web.json_response({
|
|
"status": "success",
|
|
"data": material,
|
|
"node_blocks": _node_blocks(material["workflow"], include_values=True),
|
|
})
|
|
|
|
|
|
async def api_get_materials_by_node_type(request):
|
|
node_type = request.query.get("type", "")
|
|
if not isinstance(node_type, str) or not node_type.strip() or len(node_type) > 200:
|
|
return web.json_response({"status": "error", "message": "Invalid node type"}, status=400)
|
|
materials_dir = get_materials_dir()
|
|
|
|
def collect():
|
|
matches = []
|
|
for summary in _list_materials(materials_dir):
|
|
try:
|
|
material = _read_material(resolve_within(materials_dir, summary["filename"]))
|
|
except (OSError, ValueError, json.JSONDecodeError):
|
|
continue
|
|
blocks = [block for block in _node_blocks(material["workflow"], include_values=True) if block["type"] == node_type]
|
|
if blocks:
|
|
matches.append({
|
|
"filename": summary["filename"],
|
|
"name": summary["name"],
|
|
"kind": summary["kind"],
|
|
"timestamp": summary["timestamp"],
|
|
"blocks": blocks,
|
|
"workflow_hashes": _workflow_hashes_for_blocks(material["workflow"], blocks),
|
|
})
|
|
return matches
|
|
|
|
matches = await asyncio.to_thread(collect)
|
|
return web.json_response({"status": "success", "materials": matches})
|
|
|
|
|
|
async def api_get_material_asset(request):
|
|
try:
|
|
filename = require_filename(request.query.get("filename", ""))
|
|
asset_id = require_filename(request.query.get("asset", ""))
|
|
if not filename.endswith(".json") or not asset_id.lower().endswith((".png", ".webp")):
|
|
raise ValueError("Invalid material asset")
|
|
materials_dir = get_materials_dir()
|
|
await asyncio.to_thread(_read_material, resolve_within(materials_dir, filename))
|
|
asset_path = resolve_within(_material_assets_dir(materials_dir, filename), asset_id)
|
|
if not os.path.isfile(asset_path):
|
|
raise FileNotFoundError
|
|
except (AttributeError, ValueError, json.JSONDecodeError):
|
|
return web.json_response({"status": "error", "message": "Invalid material asset"}, status=400)
|
|
except FileNotFoundError:
|
|
return web.json_response({"status": "error", "message": "Material asset not found"}, status=404)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not read material asset"}, status=500)
|
|
return web.FileResponse(asset_path)
|
|
|
|
|
|
async def api_delete_material(request):
|
|
try:
|
|
payload = await request.json()
|
|
filename = require_filename(payload.get("filename", ""))
|
|
if not filename.endswith(".json"):
|
|
raise ValueError("Invalid material filename")
|
|
materials_dir = get_materials_dir()
|
|
path = resolve_within(materials_dir, filename)
|
|
await asyncio.to_thread(_delete_material, materials_dir, filename, path)
|
|
except (AttributeError, ValueError, json.JSONDecodeError):
|
|
return web.json_response({"status": "error", "message": "Invalid material"}, status=400)
|
|
except FileNotFoundError:
|
|
return web.json_response({"status": "error", "message": "Material not found"}, status=404)
|
|
except OSError:
|
|
return web.json_response({"status": "error", "message": "Could not delete material"}, status=500)
|
|
return web.json_response({"status": "success"})
|
|
|
|
|
|
def _delete_material(materials_dir, filename, path):
|
|
_read_material(path)
|
|
os.remove(path)
|
|
assets_dir = _material_assets_dir(materials_dir, filename)
|
|
if os.path.isdir(assets_dir):
|
|
shutil.rmtree(assets_dir)
|