284 lines
10 KiB
Python
284 lines
10 KiB
Python
"""Bounded projection of declared sampler roles, independent of the host."""
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
from .execution_contexts import PromptSnapshot, TextEvidence, text_size
|
|
from .prompt_pair_resolver import _DEFAULT_ADAPTERS
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SourceRef:
|
|
node_id: str
|
|
node_type: str
|
|
output_slot: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GraphRole:
|
|
sources: tuple[SourceRef, ...]
|
|
inferred: PromptSnapshot | None
|
|
reason: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DeclaredRecipe:
|
|
anchor_id: str
|
|
positive: GraphRole
|
|
negative: GraphRole
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GraphProjection:
|
|
recipes: tuple[DeclaredRecipe, ...]
|
|
sources: tuple[SourceRef, ...]
|
|
reason: str
|
|
|
|
|
|
_SAMPLERS = {"KSampler", "KSamplerAdvanced", "SamplerCustom"}
|
|
_PM = {
|
|
"PromptManager",
|
|
"PromptManagerFocused",
|
|
"PromptManagerSplit",
|
|
"PromptManagerText",
|
|
"PromptManagerPair",
|
|
}
|
|
|
|
|
|
class _Unknown(ValueError):
|
|
pass
|
|
|
|
|
|
def _identifier(value: object) -> str:
|
|
if type(value) is int:
|
|
if abs(value) > 10**127:
|
|
raise _Unknown("Invalid node identifier")
|
|
value = str(value)
|
|
try:
|
|
text_size(value, 512)
|
|
if len(value) > 128:
|
|
raise ValueError("Oversized node identifier")
|
|
except ValueError:
|
|
raise _Unknown("Invalid node identifier") from None
|
|
if not value:
|
|
raise _Unknown("Invalid node identifier")
|
|
return value
|
|
|
|
|
|
def _link(value: object) -> tuple[str, int]:
|
|
if (
|
|
type(value) not in (tuple, list)
|
|
or len(value) != 2
|
|
or type(value[1]) is not int
|
|
or value[1] < 0
|
|
):
|
|
raise _Unknown("Expected typed graph link")
|
|
return _identifier(value[0]), value[1]
|
|
|
|
|
|
class _Projector:
|
|
def __init__(self, graph: dict) -> None:
|
|
self.graph = graph
|
|
self.visits = 0
|
|
self.transformed = False
|
|
self.selected_seen = set()
|
|
|
|
def node(self, node_id: str, depth: int) -> tuple[str, dict]:
|
|
self.visits += 1
|
|
if self.visits > 500 or depth > 64:
|
|
raise _Unknown("Graph resolution limit reached")
|
|
node = self.graph.get(node_id)
|
|
if type(node) is not dict or type(node.get("inputs")) is not dict:
|
|
raise _Unknown("Missing or malformed node")
|
|
if len(node["inputs"]) > 500:
|
|
raise _Unknown("Node input limit reached")
|
|
kind = node.get("class_type")
|
|
try:
|
|
text_size(kind, 128)
|
|
except ValueError:
|
|
raise _Unknown("Invalid node type") from None
|
|
return kind, node["inputs"]
|
|
|
|
def text(
|
|
self, value: object, kind: str, depth: int = 1, active: frozenset = frozenset()
|
|
) -> tuple[PromptSnapshot, tuple]:
|
|
if type(value) is str and kind == "string":
|
|
text_size(value)
|
|
return self.snapshot(value, value, "", ""), ()
|
|
node_id, slot = _link(value)
|
|
key = (node_id, slot)
|
|
if key in active:
|
|
raise _Unknown("Cyclic prompt graph")
|
|
name, inputs = self.node(node_id, depth)
|
|
active = active | {key}
|
|
if name == "Reroute" and slot == 0:
|
|
keys = [key for key in ("value", "input", "in") if key in inputs]
|
|
if len(keys) != 1:
|
|
raise _Unknown("Ambiguous reroute")
|
|
return self.text(inputs[keys[0]], kind, depth + 1, active)
|
|
if (
|
|
name in ("StringConcatenate", "Text Concatenate")
|
|
and slot == 0
|
|
and kind == "string"
|
|
):
|
|
keys = (
|
|
["string_a", "string_b"]
|
|
if name == "StringConcatenate"
|
|
else sorted(
|
|
key
|
|
for key in inputs
|
|
if type(key) is str and re.fullmatch(r"text_[a-z]", key)
|
|
)
|
|
)
|
|
delimiter = inputs.get(
|
|
"delimiter", "" if name == "StringConcatenate" else " "
|
|
)
|
|
text_size(delimiter)
|
|
clean_value = inputs.get("clean_whitespace", "true")
|
|
if name == "Text Concatenate" and type(clean_value) not in (str, bool):
|
|
raise _Unknown("Invalid whitespace mode")
|
|
clean = name == "Text Concatenate" and (
|
|
clean_value is True
|
|
or type(clean_value) is str
|
|
and clean_value.lower() == "true"
|
|
)
|
|
if not keys:
|
|
raise _Unknown("Missing concatenation inputs")
|
|
parts, sources, size = [], [], 0
|
|
for input_name in keys:
|
|
child, refs = self.text(
|
|
inputs.get(input_name, ""), "string", depth + 1, active
|
|
)
|
|
part = child.effective.value.strip() if clean else child.effective.value
|
|
if part or name == "StringConcatenate":
|
|
size += text_size(part) + (text_size(delimiter) if parts else 0)
|
|
if size > 65536:
|
|
raise _Unknown("Text field limit reached")
|
|
parts.append(part)
|
|
sources.extend(refs)
|
|
if len(sources) > 64:
|
|
raise _Unknown("Source reference limit reached")
|
|
self.transformed = self.transformed or bool(sources)
|
|
text = delimiter.join(parts)
|
|
return self.snapshot(text, text, "", ""), tuple(sources)
|
|
adapter = _DEFAULT_ADAPTERS.get((name, slot))
|
|
if adapter is None or adapter.kind != kind:
|
|
raise _Unknown("Unsupported node semantics or output slot")
|
|
child, sources = self.text(
|
|
inputs.get(adapter.input_name), "string", depth + 1, active
|
|
)
|
|
if name in _PM:
|
|
sources = (SourceRef(node_id, name, slot),)
|
|
if adapter.compose:
|
|
before, before_sources = self.text(
|
|
inputs.get("prepend_text", ""), "string", depth + 1, active
|
|
)
|
|
after, after_sources = self.text(
|
|
inputs.get("append_text", ""), "string", depth + 1, active
|
|
)
|
|
raw = child.effective.value
|
|
prepend, append = before.effective.value, after.effective.value
|
|
effective = " ".join(
|
|
part for part in (prepend.strip(), raw, append.strip()) if part
|
|
)
|
|
text_size(effective)
|
|
snapshot = self.snapshot(raw, effective, prepend, append)
|
|
if name not in _PM:
|
|
sources += before_sources + after_sources
|
|
elif name in _PM:
|
|
snapshot = self.snapshot(
|
|
child.effective.value, child.effective.value, "", ""
|
|
)
|
|
else:
|
|
snapshot = child
|
|
if len(sources) > 64:
|
|
raise _Unknown("Source reference limit reached")
|
|
return snapshot, sources
|
|
|
|
@staticmethod
|
|
def snapshot(raw: str, effective: str, prepend: str, append: str) -> PromptSnapshot:
|
|
return PromptSnapshot(
|
|
*(
|
|
TextEvidence(value, "graph-inferred", "graph", None)
|
|
for value in (raw, effective, prepend, append)
|
|
)
|
|
)
|
|
|
|
def role(self, value: object) -> GraphRole:
|
|
try:
|
|
self.transformed = False
|
|
snapshot, sources = self.text(value, "conditioning")
|
|
return GraphRole(
|
|
sources, snapshot, "source-transformed" if self.transformed else ""
|
|
)
|
|
except ValueError as error:
|
|
reason = str(error) if isinstance(error, _Unknown) else "Invalid text value"
|
|
return GraphRole((), None, reason)
|
|
|
|
def anchors(self, node_id: str, depth: int = 1) -> tuple[str, ...]:
|
|
if node_id in self.selected_seen:
|
|
return ()
|
|
self.selected_seen.add(node_id)
|
|
name, inputs = self.node(node_id, depth)
|
|
anchors = [node_id] if name in _SAMPLERS else []
|
|
for value in inputs.values():
|
|
if type(value) not in (tuple, list):
|
|
continue
|
|
parent, _ = _link(value)
|
|
anchors.extend(self.anchors(parent, depth + 1))
|
|
if len(anchors) > 16:
|
|
raise _Unknown("Recipe context limit reached")
|
|
return tuple(anchors)
|
|
|
|
|
|
def project_contexts(graph: object, outputs: object) -> GraphProjection:
|
|
"""Keep declared roles and graph inference, never evidence of consumption.
|
|
|
|
Sources are PM capture points. An inferred snapshot may compose text; its
|
|
presence never promotes a source observation to recorded effective text.
|
|
"""
|
|
try:
|
|
if type(graph) is not dict or len(graph) > 500:
|
|
raise _Unknown("Invalid or oversized graph")
|
|
normalized = {_identifier(key): value for key, value in graph.items()}
|
|
if len(normalized) != len(graph):
|
|
raise _Unknown("Duplicate graph identifiers")
|
|
resolver = _Projector(normalized)
|
|
if outputs is None:
|
|
anchors = [
|
|
key
|
|
for key, value in normalized.items()
|
|
if type(value) is dict
|
|
and type(value.get("class_type")) is str
|
|
and value["class_type"] in _SAMPLERS
|
|
]
|
|
elif type(outputs) in (list, tuple) and len(outputs) <= 500:
|
|
anchors = []
|
|
for output in outputs:
|
|
anchors.extend(resolver.anchors(_identifier(output)))
|
|
anchors = list(dict.fromkeys(anchors))
|
|
else:
|
|
raise _Unknown("Invalid output selection")
|
|
if len(anchors) > 16:
|
|
raise _Unknown("Recipe context limit reached")
|
|
recipes, sources = [], []
|
|
for anchor in anchors:
|
|
_, inputs = resolver.node(anchor, 1)
|
|
positive = resolver.role(inputs.get("positive"))
|
|
negative = resolver.role(inputs.get("negative"))
|
|
recipes.append(DeclaredRecipe(anchor, positive, negative))
|
|
sources.extend(positive.sources + negative.sources)
|
|
sources = tuple(dict.fromkeys(sources))
|
|
if len(sources) > 64:
|
|
raise _Unknown("Source reference limit reached")
|
|
reason = "" if recipes else "No selected sampler anchors"
|
|
if any(
|
|
role.reason
|
|
for recipe in recipes
|
|
for role in (recipe.positive, recipe.negative)
|
|
):
|
|
reason = "Some roles require inferred or unresolved text"
|
|
return GraphProjection(tuple(recipes), sources, reason)
|
|
except ValueError as error:
|
|
reason = str(error) if isinstance(error, _Unknown) else "Invalid text value"
|
|
return GraphProjection((), (), reason)
|