# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/scheduler/scheduler.ipynb.

# %% auto 0
__all__ = ['Scheduler']

# %% ../../src/scheduler/scheduler.ipynb 3
from ..imports import *
from ..util import virtual
from ..config_loader import *

# %% ../../src/scheduler/scheduler.ipynb 4
class Scheduler:
    """Base class for all diffusion schedulers"""
    def __init__(self): 
        pass
    
    @virtual
    def set_timesteps(self): pass
     
    @virtual
    def step(self): pass

    @virtual
    def add_noise(self): pass

    #---------------------------------------

    def get_config(self, without_metadata=False):
        if not without_metadata:       
            config = {}
            config["target"] = class_to_str(type(self))
            config["params"] = self.params_config  
        else:
            config = self.params_config  

        self.config = config        
        return config
    
    @property
    @virtual
    def params_config(self): return None

    #---------------------------------------

    def unsqueeze_vector_to_shape(self, vec, shape):
        vec = vec.flatten()
        while len(vec.shape) < len(shape):
            vec = vec.unsqueeze(-1)
        return vec 
