import torch

class LinearScheduler:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
            "sigma_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5000.0, "step": 0.01, "round": False}),
            "sigma_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 5000.0, "step": 0.01, "round": False}),
        }}

    RETURN_TYPES = ("SIGMAS",)
    CATEGORY = "sampling/custom_sampling/schedulers"
    FUNCTION = "get_sigmas"

    def get_sigmas(self, steps, sigma_max, sigma_min):
        sigmas = torch.linspace(sigma_max, sigma_min, steps + 1)
        return (sigmas,)

NODE_CLASS_MAPPINGS = {"LinearScheduler": LinearScheduler}
NODE_DISPLAY_NAME_MAPPINGS = {"LinearScheduler": "Linear Scheduler"}