142 lines
4.0 KiB
Python
142 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
|
|
from .adapters.samplers import is_sampler
|
|
from .graph import GraphIndex
|
|
from .models import NodeRecord
|
|
|
|
|
|
OUTPUT_INPUT_FIELDS = ("images", "image")
|
|
UPSTREAM_FIELDS = (
|
|
"images",
|
|
"image",
|
|
"samples",
|
|
"sample",
|
|
"latent",
|
|
"latent_image",
|
|
"pixels",
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class SamplerSelection:
|
|
sampler: NodeRecord
|
|
confidence: float
|
|
debug_trace: str
|
|
selected_by: str
|
|
|
|
|
|
def select_final_sampler(graph: GraphIndex) -> SamplerSelection:
|
|
for output_node in _output_nodes(graph):
|
|
selection = _select_from_output_node(graph, output_node)
|
|
if selection is not None:
|
|
return selection
|
|
|
|
fallback = _fallback_sampler(graph)
|
|
if fallback is None:
|
|
raise ValueError("No sampler candidate found")
|
|
|
|
return SamplerSelection(
|
|
sampler=fallback,
|
|
confidence=0.5,
|
|
debug_trace=(
|
|
"SAMPLER_SELECTION:\n"
|
|
f"sampler selected by fallback heuristic: "
|
|
f"{fallback.class_type} {fallback.node_id}; "
|
|
"no final SaveImage/PreviewImage output chain was found"
|
|
),
|
|
selected_by="fallback",
|
|
)
|
|
|
|
|
|
def _output_nodes(graph: GraphIndex) -> list[NodeRecord]:
|
|
return [node for node in graph.nodes() if _is_output_image_node(node)]
|
|
|
|
|
|
def _is_output_image_node(node: NodeRecord) -> bool:
|
|
lowered = node.class_type.lower()
|
|
if lowered in {"saveimage", "previewimage", "saveanimatedwebp"}:
|
|
return True
|
|
has_output_image_field = any(field in node.inputs for field in OUTPUT_INPUT_FIELDS)
|
|
return has_output_image_field and (
|
|
("save" in lowered and ("image" in lowered or "webp" in lowered))
|
|
or ("preview" in lowered and ("image" in lowered or "webp" in lowered))
|
|
)
|
|
|
|
|
|
def _select_from_output_node(
|
|
graph: GraphIndex,
|
|
output_node: NodeRecord,
|
|
) -> SamplerSelection | None:
|
|
queue: deque[tuple[NodeRecord, list[str]]] = deque()
|
|
for field in OUTPUT_INPUT_FIELDS:
|
|
target = graph.link_target(output_node.inputs.get(field))
|
|
if target is None:
|
|
continue
|
|
target_node = graph.get_node(target[0])
|
|
if target_node is None:
|
|
continue
|
|
queue.append(
|
|
(
|
|
target_node,
|
|
[f"{output_node.class_type} {output_node.node_id}.{field}"],
|
|
)
|
|
)
|
|
|
|
visited: set[str] = set()
|
|
while queue:
|
|
node, chain = queue.popleft()
|
|
if node.node_id in visited:
|
|
continue
|
|
visited.add(node.node_id)
|
|
|
|
if is_sampler(node):
|
|
trace_chain = " <- ".join(chain + [f"{node.class_type} {node.node_id}"])
|
|
return SamplerSelection(
|
|
sampler=node,
|
|
confidence=1.0,
|
|
debug_trace=(
|
|
"SAMPLER_SELECTION:\n"
|
|
f"{trace_chain}\n"
|
|
f"Selected {node.class_type} {node.node_id} because it is "
|
|
f"upstream of final {output_node.class_type}."
|
|
),
|
|
selected_by="output_chain",
|
|
)
|
|
|
|
for field in UPSTREAM_FIELDS:
|
|
target = graph.link_target(node.inputs.get(field))
|
|
if target is None:
|
|
continue
|
|
target_node = graph.get_node(target[0])
|
|
if target_node is None:
|
|
continue
|
|
queue.append(
|
|
(
|
|
target_node,
|
|
chain + [f"{node.class_type} {node.node_id}.{field}"],
|
|
)
|
|
)
|
|
return None
|
|
|
|
|
|
def _fallback_sampler(graph: GraphIndex) -> NodeRecord | None:
|
|
candidates = [
|
|
node
|
|
for node in graph.nodes()
|
|
if is_sampler(node) and _has_conditioning_inputs(node)
|
|
]
|
|
if candidates:
|
|
return candidates[-1]
|
|
|
|
samplers = [node for node in graph.nodes() if is_sampler(node)]
|
|
if samplers:
|
|
return samplers[-1]
|
|
return None
|
|
|
|
|
|
def _has_conditioning_inputs(node: NodeRecord) -> bool:
|
|
return "positive" in node.inputs or "negative" in node.inputs
|