92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
import os
|
|
|
|
import folder_paths
|
|
from comfy.samplers import KSampler
|
|
|
|
|
|
def _with_none_option(items):
|
|
return ["(none)"] + list(items)
|
|
|
|
|
|
def _strip_extension(filename):
|
|
return os.path.splitext(filename)[0]
|
|
|
|
|
|
def _pathless_stem(value):
|
|
# Keep only the filename part (drop folders), then remove extension.
|
|
leaf = str(value).replace("\\", "/").split("/")[-1]
|
|
return _strip_extension(leaf)
|
|
|
|
|
|
class AIOCoreSettingsNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
checkpoints = _with_none_option(folder_paths.get_filename_list("checkpoints"))
|
|
loras = _with_none_option(folder_paths.get_filename_list("loras"))
|
|
samplers = _with_none_option(KSampler.SAMPLERS)
|
|
schedulers = _with_none_option(KSampler.SCHEDULERS)
|
|
|
|
return {
|
|
"required": {
|
|
"checkpoint": (checkpoints,),
|
|
"lora": (loras,),
|
|
"positive_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"negative_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"sampler": (samplers,),
|
|
"scheduler": (schedulers,),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000, "step": 1}),
|
|
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
|
"filename": ("STRING", {"default": "", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("*", "*", "STRING", "STRING", "*", "*", "INT", "FLOAT", "STRING")
|
|
RETURN_NAMES = ("checkpoint", "lora", "positive_prompt", "negative_prompt", "sampler", "scheduler", "steps", "cfg", "filename")
|
|
FUNCTION = "collect"
|
|
CATEGORY = "AIO Core"
|
|
|
|
def collect(self, checkpoint, lora, positive_prompt, negative_prompt, sampler, scheduler, steps, cfg, filename):
|
|
checkpoint_value = "" if checkpoint == "(none)" else checkpoint
|
|
lora_value = "" if lora == "(none)" else lora
|
|
sampler_value = "" if sampler == "(none)" else sampler
|
|
scheduler_value = "" if scheduler == "(none)" else scheduler
|
|
|
|
filename_text = filename if isinstance(filename, str) else ""
|
|
final_filename = filename_text.strip()
|
|
if not final_filename:
|
|
parts = []
|
|
if checkpoint_value:
|
|
parts.append(_pathless_stem(checkpoint_value))
|
|
if lora_value:
|
|
parts.append(_pathless_stem(lora_value))
|
|
if sampler_value:
|
|
parts.append(sampler_value)
|
|
|
|
if parts:
|
|
final_filename = "_".join(parts) + "_"
|
|
else:
|
|
final_filename = "ComfyUI_"
|
|
if not final_filename:
|
|
final_filename = "ComfyUI_"
|
|
|
|
return (
|
|
checkpoint_value,
|
|
lora_value,
|
|
positive_prompt,
|
|
negative_prompt,
|
|
sampler_value,
|
|
scheduler_value,
|
|
steps,
|
|
cfg,
|
|
final_filename,
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AIOCoreSettingsNode": AIOCoreSettingsNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AIOCoreSettingsNode": "AIO Core Settings",
|
|
}
|