68 lines
1.7 KiB
Python
68 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from ..models import NodeRecord
|
|
|
|
|
|
SAMPLER_CLASS_TYPES = {
|
|
"KSampler",
|
|
"KSamplerAdvanced",
|
|
"SamplerCustom",
|
|
"SamplerCustomAdvanced",
|
|
}
|
|
|
|
SAMPLER_LIKE_INPUT_FIELDS = {
|
|
"positive",
|
|
"negative",
|
|
"model",
|
|
"vae",
|
|
"seed",
|
|
"noise_seed",
|
|
"steps",
|
|
"cfg",
|
|
"cfg_scale",
|
|
"sampler_name",
|
|
"sampler",
|
|
"scheduler",
|
|
"scheduler_name",
|
|
"denoise",
|
|
}
|
|
|
|
|
|
def is_sampler(node: NodeRecord) -> bool:
|
|
return node.class_type in SAMPLER_CLASS_TYPES or _is_usdu_sampler_like(node)
|
|
|
|
|
|
def extract_sampler_settings(sampler: NodeRecord) -> dict[str, Any]:
|
|
settings: dict[str, Any] = {}
|
|
field_map = {
|
|
"seed": ("seed", "noise_seed"),
|
|
"steps": ("steps",),
|
|
"cfg": ("cfg", "cfg_scale"),
|
|
"sampler_name": ("sampler_name", "sampler"),
|
|
"scheduler": ("scheduler", "scheduler_name", "schedule"),
|
|
"denoise": ("denoise",),
|
|
}
|
|
|
|
for output_name, field_names in field_map.items():
|
|
value = _first_existing_input(sampler, field_names)
|
|
if value is not None:
|
|
settings[output_name] = value
|
|
return settings
|
|
|
|
|
|
def _first_existing_input(node: NodeRecord, field_names: tuple[str, ...]) -> Any | None:
|
|
for field_name in field_names:
|
|
if field_name in node.inputs:
|
|
return node.inputs[field_name]
|
|
return None
|
|
|
|
|
|
def _is_usdu_sampler_like(node: NodeRecord) -> bool:
|
|
lowered = node.class_type.lower().replace("_", " ")
|
|
looks_like_usdu = ("ultimate" in lowered and "upscale" in lowered) or "usdu" in lowered
|
|
if not looks_like_usdu:
|
|
return False
|
|
return any(field in node.inputs for field in SAMPLER_LIKE_INPUT_FIELDS)
|