346 lines
11 KiB
Python
346 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
from .a1111_parser import parse_a1111_parameters
|
|
from .adapters.samplers import extract_sampler_settings
|
|
from .graph import GraphIndex
|
|
from .lora_extractor import extract_loras, extract_model_name, extract_vae_name
|
|
from .models import MetadataBundle, ParseResult, PromptSegment
|
|
from .prompt_merge import merge_prompt_segments
|
|
from .resolver import ConditioningResolver
|
|
from .sampler_selector import select_final_sampler
|
|
from .settings_formatter import format_setting
|
|
|
|
|
|
def _metadata_value(info: dict[str, Any], key: str) -> str | None:
|
|
value = info.get(key)
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, bytes):
|
|
return value.decode("utf-8", errors="replace")
|
|
return str(value)
|
|
|
|
|
|
def _source_format(prompt_raw: str | None, workflow_raw: str | None, parameters_raw: str | None) -> str:
|
|
if prompt_raw or workflow_raw:
|
|
return "ComfyUI prompt/workflow"
|
|
if parameters_raw:
|
|
return "A1111 parameters"
|
|
return "none"
|
|
|
|
|
|
def read_metadata(image_path: str | Path) -> MetadataBundle:
|
|
path = Path(image_path)
|
|
with Image.open(path) as image:
|
|
prompt_raw = _metadata_value(image.info, "prompt")
|
|
workflow_raw = _metadata_value(image.info, "workflow")
|
|
parameters_raw = _metadata_value(image.info, "parameters")
|
|
width, height = image.size
|
|
|
|
return MetadataBundle(
|
|
filename=path.name,
|
|
width=width,
|
|
height=height,
|
|
prompt_raw=prompt_raw,
|
|
workflow_raw=workflow_raw,
|
|
parameters_raw=parameters_raw,
|
|
source_format=_source_format(prompt_raw, workflow_raw, parameters_raw),
|
|
)
|
|
|
|
|
|
def parse_metadata_bundle(
|
|
bundle: MetadataBundle,
|
|
parameter_index: int = 0,
|
|
prefer_cached_text: bool = True,
|
|
include_raw_json: bool = True,
|
|
max_depth: int = 40,
|
|
) -> ParseResult:
|
|
if bundle.prompt_raw:
|
|
try:
|
|
return _parse_comfyui_bundle(
|
|
bundle=bundle,
|
|
prefer_cached_text=prefer_cached_text,
|
|
include_raw_json=include_raw_json,
|
|
max_depth=max_depth,
|
|
)
|
|
except _ComfyGraphStartError as exc:
|
|
if bundle.parameters_raw:
|
|
return _parse_a1111_bundle(
|
|
bundle=bundle,
|
|
parameter_index=parameter_index,
|
|
fallback_reason=str(exc),
|
|
)
|
|
return _failed_result(bundle, "FAILED", str(exc), include_raw_json)
|
|
|
|
if bundle.parameters_raw:
|
|
return _parse_a1111_bundle(
|
|
bundle=bundle,
|
|
parameter_index=parameter_index,
|
|
fallback_reason="missing ComfyUI prompt JSON",
|
|
)
|
|
|
|
return _failed_result(
|
|
bundle=bundle,
|
|
status_message="NO_METADATA",
|
|
reason="no ComfyUI prompt/workflow or A1111 parameters metadata found",
|
|
include_raw_json=include_raw_json,
|
|
)
|
|
|
|
|
|
class _ComfyGraphStartError(Exception):
|
|
pass
|
|
|
|
|
|
def _parse_comfyui_bundle(
|
|
bundle: MetadataBundle,
|
|
prefer_cached_text: bool,
|
|
include_raw_json: bool,
|
|
max_depth: int,
|
|
) -> ParseResult:
|
|
prompt = _load_required_json(bundle.prompt_raw, "prompt")
|
|
workflow, workflow_error = _load_optional_json(bundle.workflow_raw, "workflow")
|
|
graph = GraphIndex(prompt=prompt, workflow=workflow)
|
|
|
|
try:
|
|
sampler_selection = select_final_sampler(graph)
|
|
except ValueError as exc:
|
|
raise _ComfyGraphStartError(str(exc)) from exc
|
|
|
|
sampler = sampler_selection.sampler
|
|
resolver = ConditioningResolver(
|
|
graph=graph,
|
|
max_depth=max_depth,
|
|
prefer_cached_text=prefer_cached_text,
|
|
)
|
|
positive_segments = resolver.resolve_positive(sampler)
|
|
negative_segments = resolver.resolve_negative(sampler)
|
|
|
|
settings = extract_sampler_settings(sampler)
|
|
loras = extract_loras(graph, sampler)
|
|
partial_result = _partial_result(
|
|
sampler_selection=sampler_selection,
|
|
resolver=resolver,
|
|
workflow_error=workflow_error,
|
|
)
|
|
debug_trace = _debug_trace(
|
|
sampler_trace=sampler_selection.debug_trace,
|
|
resolver_trace=resolver.debug_trace,
|
|
loras=loras,
|
|
partial_result=partial_result,
|
|
workflow_error=workflow_error,
|
|
)
|
|
confidence = _confidence(
|
|
selection_confidence=sampler_selection.confidence,
|
|
segments=positive_segments + negative_segments,
|
|
has_unresolved=bool(resolver.unresolved),
|
|
used_fallback_sampler=sampler_selection.selected_by == "fallback",
|
|
has_workflow_error=workflow_error is not None,
|
|
has_llm_template_skip=_has_llm_template_skip(debug_trace),
|
|
)
|
|
status_message = _status_message(
|
|
has_unresolved=bool(resolver.unresolved),
|
|
used_fallback_sampler=sampler_selection.selected_by == "fallback",
|
|
has_workflow_error=workflow_error is not None,
|
|
has_llm_template_skip=_has_llm_template_skip(debug_trace),
|
|
)
|
|
|
|
result = ParseResult(
|
|
positive=merge_prompt_segments(positive_segments),
|
|
negative=merge_prompt_segments(negative_segments),
|
|
seed=_int_or_default(settings.get("seed"), -1),
|
|
steps=_int_or_default(settings.get("steps"), 0),
|
|
cfg=_float_or_default(settings.get("cfg"), 0.0),
|
|
width=bundle.width,
|
|
height=bundle.height,
|
|
model_name=extract_model_name(graph, sampler),
|
|
filename=bundle.filename,
|
|
vae_name=extract_vae_name(graph),
|
|
sampler_name=_str_or_empty(settings.get("sampler_name")),
|
|
scheduler=_str_or_empty(settings.get("scheduler")),
|
|
loras=loras,
|
|
raw_prompt_json=bundle.prompt_raw if include_raw_json and bundle.prompt_raw else "",
|
|
raw_workflow_json=bundle.workflow_raw if include_raw_json and bundle.workflow_raw else "",
|
|
partial_result=partial_result,
|
|
debug_trace=debug_trace,
|
|
confidence=confidence,
|
|
status_message=status_message,
|
|
)
|
|
result.setting = format_setting(result)
|
|
return result
|
|
|
|
|
|
def _load_required_json(raw: str | None, label: str) -> dict[str, Any]:
|
|
if raw is None or not raw.strip():
|
|
raise _ComfyGraphStartError(f"missing {label} JSON")
|
|
try:
|
|
value = json.loads(raw)
|
|
except json.JSONDecodeError as exc:
|
|
raise _ComfyGraphStartError(f"invalid {label} JSON: {exc}") from exc
|
|
if not isinstance(value, dict):
|
|
raise _ComfyGraphStartError(f"{label} JSON is not an object")
|
|
return value
|
|
|
|
|
|
def _load_optional_json(
|
|
raw: str | None,
|
|
label: str,
|
|
) -> tuple[dict[str, Any] | None, str | None]:
|
|
if raw is None or not raw.strip():
|
|
return None, None
|
|
try:
|
|
value = json.loads(raw)
|
|
except json.JSONDecodeError as exc:
|
|
return None, f"invalid {label} JSON: {exc}"
|
|
if not isinstance(value, dict):
|
|
return None, f"{label} JSON is not an object"
|
|
return value, None
|
|
|
|
|
|
def _partial_result(
|
|
sampler_selection: Any,
|
|
resolver: ConditioningResolver,
|
|
workflow_error: str | None,
|
|
) -> dict[str, Any]:
|
|
partial: dict[str, Any] = {
|
|
"source_format": "ComfyUI prompt/workflow",
|
|
"sampler_selection": {
|
|
"node_id": sampler_selection.sampler.node_id,
|
|
"class_type": sampler_selection.sampler.class_type,
|
|
"selected_by": sampler_selection.selected_by,
|
|
"confidence": sampler_selection.confidence,
|
|
},
|
|
}
|
|
if resolver.unresolved:
|
|
partial["unresolved"] = list(resolver.unresolved)
|
|
if workflow_error:
|
|
partial["workflow_error"] = workflow_error
|
|
return partial
|
|
|
|
|
|
def _debug_trace(
|
|
sampler_trace: str,
|
|
resolver_trace: str,
|
|
loras: list[Any],
|
|
partial_result: dict[str, Any],
|
|
workflow_error: str | None,
|
|
) -> str:
|
|
parts = [sampler_trace]
|
|
if workflow_error:
|
|
parts.append(f"WORKFLOW_FALLBACK:\n{workflow_error}")
|
|
if resolver_trace:
|
|
parts.append(resolver_trace)
|
|
if loras:
|
|
parts.append(_lora_trace(loras))
|
|
if partial_result.get("unresolved"):
|
|
parts.append("PARTIAL:\n" + json.dumps(partial_result["unresolved"], indent=2))
|
|
return "\n\n".join(part for part in parts if part)
|
|
|
|
|
|
def _lora_trace(loras: list[Any]) -> str:
|
|
lines = ["LORA_TRACE:"]
|
|
for lora in loras:
|
|
path = " <- ".join(lora.path)
|
|
lines.append(f"{lora.lora_name}: {path}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _confidence(
|
|
selection_confidence: float,
|
|
segments: list[PromptSegment],
|
|
has_unresolved: bool,
|
|
used_fallback_sampler: bool,
|
|
has_workflow_error: bool,
|
|
has_llm_template_skip: bool,
|
|
) -> float:
|
|
confidence = min(selection_confidence, 0.95)
|
|
if used_fallback_sampler:
|
|
confidence = min(confidence, 0.7)
|
|
if segments:
|
|
confidence = min(confidence, min(segment.confidence for segment in segments))
|
|
if has_workflow_error:
|
|
confidence = min(confidence, 0.85)
|
|
if has_unresolved or has_llm_template_skip:
|
|
confidence = min(confidence, 0.6)
|
|
return max(0.0, round(confidence, 2))
|
|
|
|
|
|
def _status_message(
|
|
has_unresolved: bool,
|
|
used_fallback_sampler: bool,
|
|
has_workflow_error: bool,
|
|
has_llm_template_skip: bool,
|
|
) -> str:
|
|
if has_unresolved or used_fallback_sampler or has_workflow_error or has_llm_template_skip:
|
|
return "PARTIAL"
|
|
return "OK"
|
|
|
|
|
|
def _has_llm_template_skip(debug_trace: str) -> bool:
|
|
return "LLM template input skipped" in debug_trace
|
|
|
|
|
|
def _parse_a1111_bundle(
|
|
bundle: MetadataBundle,
|
|
parameter_index: int,
|
|
fallback_reason: str,
|
|
) -> ParseResult:
|
|
result = parse_a1111_parameters(bundle.parameters_raw or "", parameter_index)
|
|
result.filename = bundle.filename
|
|
result.width = bundle.width or result.width
|
|
result.height = bundle.height or result.height
|
|
result.partial_result.setdefault("source_format", "A1111 parameters")
|
|
result.partial_result["fallback_reason"] = fallback_reason
|
|
result.debug_trace = (
|
|
"A1111_FALLBACK:\n"
|
|
f"{fallback_reason}\n"
|
|
"Parsed metadata from parameters field."
|
|
)
|
|
result.setting = format_setting(result)
|
|
return result
|
|
|
|
|
|
def _failed_result(
|
|
bundle: MetadataBundle,
|
|
status_message: str,
|
|
reason: str,
|
|
include_raw_json: bool,
|
|
) -> ParseResult:
|
|
result = ParseResult(
|
|
width=bundle.width,
|
|
height=bundle.height,
|
|
filename=bundle.filename,
|
|
raw_prompt_json=bundle.prompt_raw if include_raw_json and bundle.prompt_raw else "",
|
|
raw_workflow_json=bundle.workflow_raw if include_raw_json and bundle.workflow_raw else "",
|
|
partial_result={"source_format": bundle.source_format, "reason": reason},
|
|
debug_trace=f"{status_message}:\n{reason}",
|
|
confidence=0.0,
|
|
status_message=status_message,
|
|
)
|
|
result.setting = format_setting(result)
|
|
return result
|
|
|
|
|
|
def _int_or_default(value: Any, default: int) -> int:
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def _float_or_default(value: Any, default: float) -> float:
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def _str_or_empty(value: Any) -> str:
|
|
if value is None:
|
|
return ""
|
|
return str(value)
|