File size: 496 Bytes
82ea528
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import torch


class MochiPrepareSigmasNode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sigmas": ("SIGMAS", {"tooltip": "Override sigma schedule and steps"}),
            }
        }

    RETURN_TYPES = ("SIGMAS",)
    FUNCTION = "process"
    CATEGORY = "MochiEdit"

    def process(self, sigmas):
        sigmas = sigmas.tolist()
        if sigmas[-1] != 0.0:
            sigmas = [*sigmas, 0.0]

        return (torch.Tensor(sigmas),)