class KarrasVeSchedulerState:
    # setable values
    num_inference_steps: Optional[int] = None
    timesteps: Optional[jnp.ndarray] = None
    schedule: Optional[jnp.ndarray] = None  # sigma(t_i)

    @classmethod
    def create(cls):
        return cls()
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/schedulers/scheduling_karras_ve_flax.py