205 lines
7.6 KiB
Python
205 lines
7.6 KiB
Python
"""Extraction of the ComfyUI workflow document from a library artifact.
|
|
|
|
The workflow library keeps the original uploaded file byte-identical — its
|
|
SHA256 is the identity CivitAI hashed (by-hash lookup and version tracking
|
|
depend on it). Consumers (send-to-ComfyUI, graph analysis) work on the
|
|
extracted workflow document instead:
|
|
|
|
- ``.json`` artifacts ARE the document; they are validated and read in place.
|
|
- ``.zip`` / ``.png`` / ``.webp`` containers get the document materialized
|
|
once into an LM-owned derived ``.workflow.json`` sidecar inside the
|
|
artifact's sidecar directory. Scanners exclude the suffix from the library
|
|
walk; lifecycle operations (move/rename/delete) carry it like any other
|
|
sidecar.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import zipfile
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from .sidecar_paths import get_sidecar_base_name, get_sidecar_dir
|
|
from .workflow_analysis import (
|
|
MAX_WORKFLOW_JSON_BYTES,
|
|
load_workflow_file,
|
|
parse_workflow_text,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
WORKFLOW_DOCUMENT_SUFFIX = ".workflow.json"
|
|
|
|
# Image containers whose metadata chunks can embed a ComfyUI workflow.
|
|
IMAGE_CONTAINER_EXTENSIONS = frozenset({".png", ".webp"})
|
|
|
|
# Every file type the workflow library accepts.
|
|
WORKFLOW_ARTIFACT_EXTENSIONS = frozenset({".json", ".zip"} | IMAGE_CONTAINER_EXTENSIONS)
|
|
|
|
# Bounds for zip exploration: workflow archives are user uploads with no
|
|
# structural contract, so candidate discovery is capped rather than exhaustive.
|
|
_MAX_ZIP_JSON_CANDIDATES = 64
|
|
|
|
|
|
def get_workflow_document_path(artifact_path: str) -> str:
|
|
"""Return the path of the extracted workflow document for an artifact.
|
|
|
|
``.json`` artifacts are their own document; containers map to a derived
|
|
``.workflow.json`` sidecar in the artifact's sidecar directory (alongside
|
|
the artifact or in the centralized mirror). Workflow-library containers
|
|
include their original extension in the sidecar name to keep same-stem
|
|
artifacts independent.
|
|
"""
|
|
|
|
if os.path.splitext(artifact_path)[1].lower() == ".json":
|
|
return artifact_path
|
|
base_name = get_sidecar_base_name(artifact_path)
|
|
return os.path.join(
|
|
get_sidecar_dir(artifact_path), base_name + WORKFLOW_DOCUMENT_SUFFIX
|
|
)
|
|
|
|
|
|
def _zip_json_candidates(archive: zipfile.ZipFile) -> List[zipfile.ZipInfo]:
|
|
"""Return the archive's plausible workflow documents, best first."""
|
|
|
|
candidates: List[zipfile.ZipInfo] = []
|
|
for info in archive.infolist():
|
|
if info.is_dir():
|
|
continue
|
|
parts = info.filename.replace("\\", "/").split("/")
|
|
# Skip macOS resource forks and hidden files/directories.
|
|
if any(part == "__MACOSX" or part.startswith(".") for part in parts):
|
|
continue
|
|
if not parts[-1].lower().endswith(".json"):
|
|
continue
|
|
if info.file_size > MAX_WORKFLOW_JSON_BYTES:
|
|
continue
|
|
candidates.append(info)
|
|
if len(candidates) >= _MAX_ZIP_JSON_CANDIDATES:
|
|
break
|
|
# Prefer shallow, small documents: the workflow is typically the
|
|
# archive's headline content, not a bundled config deep in a subfolder.
|
|
candidates.sort(key=lambda info: (len(info.filename.split("/")), info.file_size))
|
|
return candidates
|
|
|
|
|
|
def _extract_from_zip(archive_path: str) -> Optional[Dict[str, Any]]:
|
|
try:
|
|
with zipfile.ZipFile(archive_path, "r") as archive:
|
|
for info in _zip_json_candidates(archive):
|
|
try:
|
|
with archive.open(info) as handle:
|
|
raw = handle.read(MAX_WORKFLOW_JSON_BYTES + 1)
|
|
text = raw.decode("utf-8")
|
|
except (OSError, UnicodeDecodeError, zipfile.BadZipFile):
|
|
continue
|
|
document = parse_workflow_text(text)
|
|
if document is not None:
|
|
return document
|
|
except (OSError, zipfile.BadZipFile) as exc:
|
|
logger.debug("Workflow archive unreadable: %s (%s)", archive_path, exc)
|
|
return None
|
|
|
|
|
|
def _extract_from_image(image_path: str) -> Optional[Dict[str, Any]]:
|
|
# Local import: PIL/piexif are heavy and only needed for image artifacts.
|
|
from .exif_utils import ExifUtils
|
|
|
|
try:
|
|
metadata = ExifUtils.load_structured_metadata(image_path)
|
|
except Exception as exc:
|
|
logger.debug("Workflow image metadata unreadable: %s (%s)", image_path, exc)
|
|
return None
|
|
# The UI-format document ("workflow") is what ComfyUI's loadGraphData
|
|
# consumes; fall back to the API-format "prompt" payload.
|
|
for key in ("workflow", "prompt"):
|
|
text = metadata.get(key)
|
|
if not text:
|
|
continue
|
|
document = parse_workflow_text(text)
|
|
if document is not None:
|
|
return document
|
|
return None
|
|
|
|
|
|
def extract_workflow_document(artifact_path: str) -> Optional[Dict[str, Any]]:
|
|
"""Extract the ComfyUI workflow document from an artifact.
|
|
|
|
Returns the parsed document when the artifact carries a recognizable
|
|
workflow, ``None`` otherwise (unreadable file, unknown extension, no
|
|
workflow inside). Never raises.
|
|
"""
|
|
|
|
extension = os.path.splitext(artifact_path)[1].lower()
|
|
try:
|
|
if extension == ".json":
|
|
return load_workflow_file(artifact_path)
|
|
if extension == ".zip":
|
|
return _extract_from_zip(artifact_path)
|
|
if extension in IMAGE_CONTAINER_EXTENSIONS:
|
|
return _extract_from_image(artifact_path)
|
|
except Exception as exc: # defensive: extraction must never break a scan
|
|
logger.debug("Workflow extraction failed for %s: %s", artifact_path, exc)
|
|
return None
|
|
|
|
|
|
def materialize_workflow_document(
|
|
artifact_path: str, *, read_path: Optional[str] = None
|
|
) -> Optional[str]:
|
|
"""Ensure the extracted workflow document exists on disk.
|
|
|
|
Returns the document path, or ``None`` when the artifact carries no
|
|
recognizable workflow. ``.json`` artifacts are their own document (the
|
|
path is returned after validation; nothing is written). Containers get
|
|
the document written to the derived sidecar path; an up-to-date derived
|
|
file is reused without re-extracting.
|
|
|
|
``read_path`` overrides the path the artifact bytes are read from (the
|
|
scanner hashes/extracts via the real path while sidecar locations derive
|
|
from the business path).
|
|
"""
|
|
|
|
read_path = read_path or artifact_path
|
|
|
|
if os.path.splitext(artifact_path)[1].lower() == ".json":
|
|
return artifact_path if load_workflow_file(read_path) is not None else None
|
|
|
|
document_path = get_workflow_document_path(artifact_path)
|
|
try:
|
|
if os.path.exists(document_path) and os.path.getmtime(
|
|
document_path
|
|
) >= os.path.getmtime(read_path):
|
|
return document_path
|
|
except OSError:
|
|
pass
|
|
|
|
document = extract_workflow_document(read_path)
|
|
if document is None:
|
|
return None
|
|
|
|
try:
|
|
os.makedirs(os.path.dirname(document_path), exist_ok=True)
|
|
with open(document_path, "w", encoding="utf-8") as handle:
|
|
json.dump(document, handle, ensure_ascii=False)
|
|
except OSError as exc:
|
|
logger.warning(
|
|
"Failed to write workflow document %s: %s", document_path, exc
|
|
)
|
|
return None
|
|
return document_path
|
|
|
|
|
|
def load_artifact_workflow_document(artifact_path: str) -> Optional[Dict[str, Any]]:
|
|
"""Load the workflow document for a library artifact.
|
|
|
|
Reads the derived sidecar for containers (materializing it on demand
|
|
when missing or stale) and the artifact itself for bare JSON.
|
|
"""
|
|
|
|
document_path = materialize_workflow_document(artifact_path)
|
|
if document_path is None:
|
|
return None
|
|
return load_workflow_file(document_path)
|