67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
import comfy.samplers
|
|
import os
|
|
import yaml
|
|
|
|
base_path = os.path.dirname(os.path.realpath(__file__))
|
|
models_subdir = "../../models"
|
|
configs_subdir = "configs/checkpoints"
|
|
|
|
# Get the directory paths
|
|
models_dir = os.path.join(base_path, models_subdir)
|
|
config_file_directory = os.path.join(models_dir, configs_subdir)
|
|
|
|
config_files = [f for f in os.listdir(config_file_directory) if os.path.isfile(os.path.join(config_file_directory, f))]
|
|
|
|
# This node allows loading configuration per checkpoint.
|
|
class CheckpointLoaderConfigCustomNode:
|
|
CATEGORY = "utils"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"config_file": (config_files,),
|
|
}}
|
|
RETURN_TYPES = (
|
|
"FLOAT",
|
|
"INT",
|
|
comfy.samplers.KSampler.SAMPLERS,
|
|
comfy.samplers.KSampler.SCHEDULERS,
|
|
"INT",
|
|
"FLOAT"
|
|
)
|
|
RETURN_NAMES = (
|
|
"cfg",
|
|
"steps",
|
|
"sampler",
|
|
"scheduler",
|
|
"seed",
|
|
"denoise"
|
|
)
|
|
FUNCTION = "load_checkpoint_configuration"
|
|
|
|
# Load a config from the models/configs directory based on the checkpoint name.
|
|
def load_checkpoint_configuration(self, config_file):
|
|
config_file = os.path.join(config_file_directory, config_file)
|
|
if not os.path.exists(config_file):
|
|
raise ValueError(f"Config not found. Expected path: {config_file}")
|
|
|
|
with open(config_file, 'r') as file:
|
|
config = yaml.safe_load(file)
|
|
cfg = config.get('cfg', 4)
|
|
steps = config.get('steps', 10)
|
|
sampler_name = config.get('sampler', "euler")
|
|
scheduler_name = config.get('scheduler', "normal")
|
|
seed = config.get('seed', -1)
|
|
denoise = config.get('denoise', 1.0)
|
|
|
|
loaded_config = (
|
|
cfg,
|
|
steps,
|
|
sampler_name,
|
|
scheduler_name,
|
|
seed,
|
|
denoise
|
|
)
|
|
print(f"Loaded config: {loaded_config}")
|
|
return loaded_config
|