scheduler.py

import torch


class BaseScheduler():
    def __init__(self, args):
        self.args = args
        self.step_index = 0
        self.latents = None
    
    def step_pre(self, step_index):
        self.step_index = step_index
        self.latents = self.latents.to(dtype=torch.bfloat16)