ComfyUI scheduler for Nvidia's Cosmos models

cosmos_scheduler_node.py

# Usage: Put in custom_nodes and restart ComfyUI/refresh browser.
import torch

class CosmosSchedulerNode:
    RETURN_TYPES = ("SIGMAS",)
    FUNCTION = "go"
    CATEGORY = "sampling/custom_sampling/schedulers"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "steps": ("INT", {"default": 20, "min": 1}),
                "sigma_max": ("FLOAT", {"default": 80.0, "step": 0.0001}),
                "sigma_min": ("FLOAT", {"default": 0.002, "step": 0.0001}),
                "order": ("FLOAT", {"default": 7.0, "min": 0.01}),
            },
        }

    @classmethod
    def go(cls, *, steps, sigma_max, sigma_min, order):
        adj_indices = torch.arange(steps, dtype=torch.float64).div_(steps)
        result = torch.zeros(steps + 1, dtype=torch.float)
        adj_order = 1 / order
        result[:-1] = (
            (
                sigma_max**adj_order
                + adj_indices * (sigma_min**adj_order - sigma_max**adj_order)
            )
            ** order
        ).float()
        return (result,)


NODE_CLASS_MAPPINGS = {
    "CosmosScheduler": CosmosSchedulerNode,
}