diff --git a/src/custom_schedulers/get_sigmas_sigmoid.py b/src/custom_schedulers/get_sigmas_sigmoid.py index 6fdd64d..9d64454 100644 --- a/src/custom_schedulers/get_sigmas_sigmoid.py +++ b/src/custom_schedulers/get_sigmas_sigmoid.py @@ -1,9 +1,13 @@ import torch -# Necessary map for it to work with simple sampler, do not include steps in the settings as it gets added automatically. +# Define the settings for the sigmoid function scheduler +# Note: Do not include 'steps' in the settings as it gets added automatically. +# The settings dictionary should have the same names as the parameters in the get_sigmas function. settings = { "name": "sigmoid", - "settings": { # name: [type, default value, min val, max val, snap to, round] Note: this currently only supports floats. + "settings": { + # Format: "parameter_name": [type, default_value, min_value, max_value, step_value, round_flag] + # Types supported: FLOAT, INT, STRING, BOOLEAN "sigma_max_sig": ["FLOAT", 25.0, 0.0, 5000.0, 0.01, False], "sigma_min_sig": ["FLOAT", 0.0, 0.0, 5000.0, 0.01, False], "steepness": ["FLOAT", 3.5, 0.0, 10.0, 0.01, False], @@ -11,7 +15,8 @@ settings = { } } -# A scheduler must have a function called "get_sigmas" to work properly, its first input must be the steps for it to work properly. +# A scheduler must have a function called "get_sigmas" to work properly. +# The first input must be the steps for it to work properly as the steps setting is automatically added. def get_sigmas(n, sigma_max_sig, sigma_min_sig, steepness=10, midpoint_ratio=0.5, device='cpu'): """ Generates 'n' sigmas on an S-curve from sigma_max to sigma_min, with a controllable midpoint. diff --git a/src/dynamic_scheduler_loader.py b/src/dynamic_scheduler_loader.py index a2bd605..ed17ce9 100644 --- a/src/dynamic_scheduler_loader.py +++ b/src/dynamic_scheduler_loader.py @@ -2,13 +2,27 @@ from .custom_schedulers.custom_schedulers import CustomSchedulers def generate_scheduler_node_class(settings, get_sigmas_function): class DynamicSchedulerNode: + + @classmethod + def create_setting_entry(setting_type, setting_value): + if setting_type == "INT": + return ("INT", {"default": setting_value[1], "min": setting_value[2], "max": setting_value[3]}) + elif setting_type == "FLOAT": + return ("FLOAT", {"default": setting_value[1], "min": setting_value[2], "max": setting_value[3], "step": setting_value[4]}) + elif setting_type == "STRING": + return ("STRING", {"default": setting_value[1]}) + elif setting_type == "BOOLEAN": + return ("BOOLEAN", {"default": setting_value[1]}) + else: + raise ValueError(f"Unsupported setting type: {setting_type}") + @classmethod def INPUT_TYPES(cls): inputs = { "required": { "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), **{ - key: ("FLOAT", {"default": value[1], "min": value[2], "max": value[3], "step": value[4]}) + key: cls.create_setting_entry(value[0], value) for key, value in settings['settings'].items() } }