79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
import comfy.samplers
|
|
|
|
from .sampler_combo_common import (
|
|
DEFAULT_SAMPLER,
|
|
DEFAULT_SCHEDULER,
|
|
DEFAULT_STEPS,
|
|
MAX_PAIRS,
|
|
PRESETS_TYPE,
|
|
selected_index,
|
|
)
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
|
|
ANY_TYPE = AnyType("*")
|
|
|
|
|
|
class SamplerComboSelectorNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
required = {
|
|
"presets": (PRESETS_TYPE,),
|
|
"pair_count": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": MAX_PAIRS,
|
|
"step": 1,
|
|
}),
|
|
}
|
|
for i in range(1, MAX_PAIRS + 1):
|
|
required[f"label_{i}"] = ("STRING", {
|
|
"default": f"Preset {i}",
|
|
"multiline": False,
|
|
})
|
|
required[f"enabled_{i}"] = ("BOOLEAN", {"default": i == 1})
|
|
return {"required": required}
|
|
|
|
RETURN_TYPES = (
|
|
comfy.samplers.KSampler.SAMPLERS,
|
|
ANY_TYPE,
|
|
"INT",
|
|
)
|
|
RETURN_NAMES = ("sampler_name", "scheduler", "steps")
|
|
FUNCTION = "select_combo"
|
|
CATEGORY = "sampling"
|
|
|
|
def select_combo(self, presets, pair_count, **kwargs):
|
|
pairs = []
|
|
if isinstance(presets, dict):
|
|
pairs = presets.get("pairs", [])
|
|
if not isinstance(pairs, list):
|
|
pairs = []
|
|
|
|
usable_count = min(
|
|
max(pair_count, 1),
|
|
len(pairs) if len(pairs) > 0 else MAX_PAIRS,
|
|
)
|
|
idx = selected_index(usable_count, kwargs)
|
|
|
|
if 0 <= idx < len(pairs):
|
|
selected = pairs[idx]
|
|
return (
|
|
selected.get("sampler", DEFAULT_SAMPLER),
|
|
selected.get("scheduler", DEFAULT_SCHEDULER),
|
|
int(selected.get("steps", DEFAULT_STEPS)),
|
|
)
|
|
|
|
if len(pairs) > 0:
|
|
selected = pairs[0]
|
|
return (
|
|
selected.get("sampler", DEFAULT_SAMPLER),
|
|
selected.get("scheduler", DEFAULT_SCHEDULER),
|
|
int(selected.get("steps", DEFAULT_STEPS)),
|
|
)
|
|
|
|
return (DEFAULT_SAMPLER, DEFAULT_SCHEDULER, DEFAULT_STEPS)
|