hcpdiff 2.2__py3-none-any.whl → 2.3__py3-none-any.whl

hcpdiff/ckpt_manager/__init__.py

@@ -1,4 +1,4 @@
 from .format import EmbFormat, DiffusersSD15Format, DiffusersModelFormat, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSDXLFormat, \
     OfficialSD15Format, LoraWebuiFormat
-from .ckpt import EmbSaver, easy_emb_saver
+from .ckpt import EmbSaver
 from .loader import HCPLoraLoader

hcpdiff/ckpt_manager/ckpt.py

@@ -1,24 +1,28 @@
-from rainbowneko.ckpt_manager import NekoSaver, CkptFormat, LocalCkptSource, PKLFormat
-from torch import nn
+from rainbowneko.ckpt_manager import NekoSaver, CkptFormat, LocalCkptSource, PKLFormat, LAYERS_ALL, LAYERS_TRAINABLE
+from torch import Tensor
 from typing import Dict, Any
 
 class EmbSaver(NekoSaver):
-    def __init__(self, format: CkptFormat, source: LocalCkptSource, target_key='embs', prefix=None):
-        super().__init__(format, source)
-        self.target_key = target_key
+    def __init__(self, format: CkptFormat=None, source: LocalCkptSource=None, layers='all', key_map=None, prefix=None):
+        if format is None:
+            format = PKLFormat()
+        if source is None:
+            source = LocalCkptSource()
+        key_map = key_map or ('name -> name', 'embs -> embs', 'name_template -> name_template')
+        super().__init__(format, source, layers=layers, key_map=key_map)
         self.prefix = prefix
 
-    def save_to(self, name, model: nn.Module, plugin_groups: Dict[str, Any], model_ema=None, exclude_key=None,
-                name_template=None):
-        train_pts = plugin_groups[self.target_key]
-        for pt_name, pt in train_pts.items():
-            self.save(pt_name, (pt_name, pt), prefix=self.prefix)
+    def _save_to(self, name, embs: Dict[str, Tensor], name_template=None):
+        for pt_name, pt in embs.items():
+            if self.layers == LAYERS_ALL:
+                pass
+            elif self.layers == LAYERS_TRAINABLE:
+                if not pt.requires_grad:
+                    continue
+            elif pt_name not in self.layers:
+                continue
+
+            self.save((pt_name, pt), pt_name, prefix=self.prefix)
             if name_template is not None:
                 pt_name = name_template.format(pt_name)
-                self.save(pt_name, (pt_name, pt), prefix=self.prefix)
-
-def easy_emb_saver():
-    return EmbSaver(
-        format=PKLFormat(),
-        source=LocalCkptSource(),
-    )
+                self.save((pt_name, pt), pt_name, prefix=self.prefix)

hcpdiff/ckpt_manager/format/diffusers.py

@@ -3,7 +3,7 @@ from diffusers import ModelMixin, AutoencoderKL, UNet2DConditionModel, PixArtTra
 from rainbowneko.ckpt_manager.format import CkptFormat
 from transformers import CLIPTextModel, AutoTokenizer, T5EncoderModel
 
-from hcpdiff.diffusion.sampler import DDPMSampler, DDPMDiscreteSigmaScheduler
+from hcpdiff.diffusion.sampler import VPSampler, DDPMDiscreteSigmaScheduler
 from hcpdiff.models.compose import SDXLTokenizer, SDXLTextEncoder
 
 class DiffusersModelFormat(CkptFormat):
@@ -23,7 +23,7 @@ class DiffusersSD15Format(CkptFormat):
             pretrained_model, subfolder="unet", revision=revision, torch_dtype=dtype
         )
         vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
-        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+        noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler())
 
         TE = TE or CLIPTextModel.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
         tokenizer = tokenizer or AutoTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)
@@ -37,7 +37,7 @@ class DiffusersSDXLFormat(CkptFormat):
             pretrained_model, subfolder="unet", revision=revision, torch_dtype=dtype
         )
         vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
-        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+        noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler())
 
         TE = TE or SDXLTextEncoder.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
         tokenizer = tokenizer or SDXLTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)
@@ -51,7 +51,7 @@ class DiffusersPixArtFormat(CkptFormat):
             pretrained_model, subfolder="transformer", revision=revision, torch_dtype=dtype
         )
         vae = vae or AutoencoderKL.from_pretrained(pretrained_model, subfolder="vae", revision=revision, torch_dtype=dtype)
-        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+        noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler())
 
         TE = TE or T5EncoderModel.from_pretrained(pretrained_model, subfolder="text_encoder", revision=revision, torch_dtype=dtype)
         tokenizer = tokenizer or AutoTokenizer.from_pretrained(pretrained_model, subfolder="tokenizer", revision=revision, use_fast=False)

hcpdiff/ckpt_manager/format/sd_single.py

@@ -2,7 +2,7 @@ import torch
 from diffusers import AutoencoderKL, StableDiffusionPipeline, StableDiffusionXLPipeline
 from rainbowneko.ckpt_manager.format import CkptFormat
 
-from hcpdiff.diffusion.sampler import DDPMSampler, DDPMDiscreteSigmaScheduler
+from hcpdiff.diffusion.sampler import VPSampler, DDPMDiscreteSigmaScheduler
 from hcpdiff.models.compose import SDXLTextEncoder, SDXLTokenizer
 
 class OfficialSD15Format(CkptFormat):
@@ -14,7 +14,7 @@ class OfficialSD15Format(CkptFormat):
         pipe = StableDiffusionPipeline.from_single_file(
             pretrained_model, revision=revision, torch_dtype=dtype, **pipe_args
         )
-        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+        noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler())
         return dict(denoiser=pipe.unet, TE=pipe.text_encoder, vae=pipe.vae, noise_sampler=noise_sampler, tokenizer=pipe.tokenizer)
 
 class OfficialSDXLFormat(CkptFormat):
@@ -34,7 +34,7 @@ class OfficialSDXLFormat(CkptFormat):
             pretrained_model, revision=revision, torch_dtype=dtype, **pipe_args
         )
 
-        noise_sampler = noise_sampler or DDPMSampler(DDPMDiscreteSigmaScheduler())
+        noise_sampler = noise_sampler or VPSampler(DDPMDiscreteSigmaScheduler())
         TE = SDXLTextEncoder([('clip_L', pipe.text_encoder), ('clip_bigG', pipe.text_encoder_2)])
         tokenizer = SDXLTokenizer([('clip_L', pipe.tokenizer), ('clip_bigG', pipe.tokenizer_2)])
 

hcpdiff/ckpt_manager/loader.py

@@ -14,17 +14,24 @@ def get_lora_rank_and_cls(lora_state):
 
 class HCPLoraLoader(NekoPluginLoader):
     def __init__(self, format: CkptFormat=None, source: LocalCkptSource=None, path: str = None, layers='all', target_plugin=None,
-                 state_prefix=None, base_model_alpha=0.0, load_ema=False, module_to_load='', **plugin_kwargs):
+                 state_prefix=None, base_model_alpha=0.0, load_ema=False, module_to_load='', key_map=None, **plugin_kwargs):
+        key_map = key_map or ('name -> name', 'model -> model')
         super().__init__(format, source, path=path, layers=layers, target_plugin=target_plugin, state_prefix=state_prefix,
-                         base_model_alpha=base_model_alpha, load_ema=load_ema, **plugin_kwargs)
+                         base_model_alpha=base_model_alpha, load_ema=load_ema, key_map=key_map, **plugin_kwargs)
         self.module_to_load = module_to_load
 
-    def load_to(self, name, model):
+    def _load_to(self, name, model):
         # get model to load plugin and its named_modules
         model = model if self.module_to_load == '' else eval(f"model.{self.module_to_load}")
 
         named_modules = {k:v for k, v in model.named_modules()}
-        plugin_state = self.load(self.path, map_location='cpu')['base_ema' if self.load_ema else 'base']
+        state_dict = self.load(self.path, map_location='cpu')
+        if 'base' in state_dict or 'base_ema' in state_dict:
+            plugin_state = state_dict['base_ema' if self.load_ema else 'base']
+        elif 'plugin' in state_dict or 'plugin_ema' in state_dict:
+            plugin_state = state_dict['plugin_ema' if self.load_ema else 'plugin']
+        else:
+            plugin_state = state_dict
 
         # filter layers to load
         if self.layers != 'all':

hcpdiff/diffusion/noise/__init__.py

@@ -1,2 +1 @@
 from .pyramid_noise import PyramidNoiseSampler
-from .zero_terminal import ZeroTerminalSampler

hcpdiff/diffusion/sampler/VP.py

@@ -0,0 +1,27 @@
+from .base import Sampler
+
+class VPSampler(Sampler):
+    # closed-form: \alpha(t)^2 + \sigma(t)^2 = 1
+    def velocity_to_eps(self, v_pred, x_t, t):
+        alpha = self.sigma_scheduler.alpha(t)
+        sigma = self.sigma_scheduler.sigma(t)
+        return alpha*v_pred+sigma*x_t
+
+    def eps_to_velocity(self, eps, x_t, t, x_0=None):
+        alpha = self.sigma_scheduler.alpha(t)
+        sigma = self.sigma_scheduler.sigma(t)
+        if x_0 is None:
+            x_0 = self.eps_to_x0(eps, x_t, t)
+        return alpha*eps-sigma*x_0
+
+    def velocity_to_x0(self, v_pred, x_t, t):
+        alpha = self.sigma_scheduler.alpha(t)
+        sigma = self.sigma_scheduler.sigma(t)
+        return alpha*x_t-sigma*v_pred
+
+    def x0_to_velocity(self, x_0, x_t, t, eps=None):
+        alpha = self.sigma_scheduler.alpha(t)
+        sigma = self.sigma_scheduler.sigma(t)
+        if eps is None:
+            eps = self.x0_to_eps(x_0, x_t, t)
+        return alpha*eps-sigma*x_0

hcpdiff/diffusion/sampler/__init__.py

@@ -1,5 +1,4 @@
 from .sigma_scheduler import *
-from .base import BaseSampler
-from .ddpm import DDPMSampler
-from .edm import EDMSampler
+from .base import BaseSampler, Sampler
+from .VP import VPSampler
 from .diffusers import DiffusersSampler

hcpdiff/diffusion/sampler/base.py

@@ -1,72 +1,134 @@
 from typing import Tuple
+
 import torch
-from .sigma_scheduler import SigmaScheduler
-from diffusers import DDPMScheduler
 
-class BaseSampler:
-    def __init__(self, sigma_scheduler: SigmaScheduler, generator: torch.Generator = None):
-        self.sigma_scheduler = sigma_scheduler
-        self.generator = generator
+from .sigma_scheduler import SigmaScheduler
 
-    def c_in(self, sigma):
-        return 1
+try:
+    from diffusers.utils import randn_tensor
+except:
+    # new version of diffusers
+    from diffusers.utils.torch_utils import randn_tensor
 
-    def c_out(self, sigma):
-        return 1
+class BaseSampler:
+    def __init__(self, sigma_scheduler: SigmaScheduler, pred_type='eps', target_type='eps', generator: torch.Generator = None):
+        '''
+        Some losses can only be calculated in a specific space. Such as SSIM in x0 space.
+        The model pred need convert to target space.
 
-    def c_skip(self, sigma):
-        return 1
+        :param pred_type: ['x0', 'eps', 'velocity', ..., None]  The output space of the model
+        :param target_type: ['x0', 'eps', 'velocity', ..., None]  The space to calculate the loss
+        '''
 
-    @property
-    def num_timesteps(self):
-        return getattr(self.sigma_scheduler, 'num_timesteps', 1000.)
+        self.sigma_scheduler = sigma_scheduler
+        self.generator = generator
+        self.pred_type = pred_type
+        self.target_type = target_type
 
     def get_timesteps(self, N_steps, device='cuda'):
-        return torch.linspace(0, self.num_timesteps, N_steps, device=device)
+        times = torch.linspace(0., 1., N_steps, device=device)
+        return self.sigma_scheduler.scale_t(times)
 
     def make_nosie(self, shape, device='cuda', dtype=torch.float32):
-        return torch.randn(shape, generator=self.generator, device=device, dtype=dtype)
+        return randn_tensor(shape, generator=self.generator, device=device, dtype=dtype)
 
     def init_noise(self, shape, device='cuda', dtype=torch.float32):
-        sigma = self.sigma_scheduler.sigma_max
+        sigma = self.sigma_scheduler.sigma_end
         return self.make_nosie(shape, device, dtype)*sigma
 
-    def add_noise(self, x, sigma) -> Tuple[torch.Tensor, torch.Tensor]:
+    def add_noise(self, x, t) -> Tuple[torch.Tensor, torch.Tensor]:
         noise = self.make_nosie(x.shape, device=x.device)
-        noisy_x = (x.to(dtype=torch.float32)-self.c_out(sigma)*noise)/self.c_skip(sigma)
+        alpha = self.sigma_scheduler.alpha(t).view(-1, 1, 1, 1).to(x.device)
+        sigma = self.sigma_scheduler.sigma(t).view(-1, 1, 1, 1).to(x.device)
+        noisy_x = alpha*x+sigma*noise
         return noisy_x.to(dtype=x.dtype), noise.to(dtype=x.dtype)
 
     def add_noise_rand_t(self, x):
         bs = x.shape[0]
         # timesteps: [0, 1]
-        sigma, timesteps = self.sigma_scheduler.sample_sigma(shape=(bs,))
-        sigma = sigma.view(-1, 1, 1, 1).to(x.device)
+        timesteps = self.sigma_scheduler.sample(shape=(bs,))
         timesteps = timesteps.to(x.device)
-        noisy_x, noise = self.add_noise(x, sigma)
+        noisy_x, noise = self.add_noise(x, timesteps)
 
         # Sample a random timestep for each image
-        timesteps = timesteps*(self.num_timesteps-1)
-        return noisy_x, noise, sigma, timesteps
+        return noisy_x, noise, timesteps
 
     def denoise(self, x, sigma, eps=None, generator=None):
         raise NotImplementedError
 
-    def eps_to_x0(self, eps, x_t, sigma):
-        return self.c_skip(sigma)*x_t+self.c_out(sigma)*eps
-
-    def velocity_to_eps(self, v_pred, x_t, sigma):
-        alpha = 1/(sigma**2+1)
-        sqrt_alpha = alpha.sqrt()
-        one_sqrt_alpha = (1-alpha).sqrt()
-        return sqrt_alpha*v_pred + one_sqrt_alpha*(x_t*sqrt_alpha)
-
-    def eps_to_velocity(self, eps, x_t, sigma):
-        alpha = 1/(sigma**2+1)
-        sqrt_alpha = alpha.sqrt()
-        one_sqrt_alpha = (1-alpha).sqrt()
-        return eps/sqrt_alpha - one_sqrt_alpha*x_t
-
-    def velocity_to_x0(self, v_pred, x_t, sigma):
-        alpha = 1/(sigma**2+1)
-        one_sqrt_alpha = (1-alpha).sqrt()
-        return alpha*x_t - one_sqrt_alpha*v_pred
+    def get_target(self, x0, x_t, t, eps=None, target_type=None):
+        raise x0
+
+    def pred_for_target(self, pred, x_t, t, eps=None, target_type=None):
+        return self.sigma_scheduler.c_skip(t)*x_t+self.sigma_scheduler.c_out(t)*pred
+
+class Sampler(BaseSampler):
+    '''
+    Some losses can only be calculated in a specific space. Such as SSIM in x0 space.
+    The model pred need convert to target space.
+
+    :param pred_type: ['x0', 'eps', 'velocity', ..., None]  The output space of the model
+    :param target_type: ['x0', 'eps', 'velocity', ..., None]  The space to calculate the loss
+    '''
+
+    def get_target(self, x_0, x_t, t, eps=None, target_type=None):
+        '''
+        target_type can be specified by the loss. If not specified use self.target_type as default.
+        '''
+        target_type = target_type or self.target_type
+        if target_type == 'x0':
+            raise x_0
+        elif target_type == 'eps':
+            return eps if eps is not None else self.x0_to_eps(eps, x_t, t)
+        elif target_type == 'velocity':
+            return self.x0_to_velocity(x_0, x_t, t, eps)
+        else:
+            return (x_0-self.sigma_scheduler.c_skip(t)*x_t)/self.sigma_scheduler.c_out(t)
+
+    def pred_for_target(self, pred, x_t, t, eps=None, target_type=None):
+        '''
+        target_type can be specified by the loss. If not specified use self.target_type as default.
+        '''
+        target_type = target_type or self.target_type
+        if self.pred_type == target_type:
+            return pred
+        else:
+            cvt_func = getattr(self, f'{self.pred_type}_to_{target_type}', None)
+            if cvt_func is None:
+                if target_type == 'x0':
+                    return self.sigma_scheduler.c_skip(t)*x_t+self.sigma_scheduler.c_out(t)*pred
+                else:
+                    raise ValueError(f'pred_type "{self.pred_type}" can not be convert for target_type "{target_type}"')
+            else:
+                return cvt_func(pred, x_t, t)
+
+    # convert targets
+    def x0_to_eps(self, x_0, x_t, t):
+        return (x_t-self.sigma_scheduler.alpha(t)*x_0)/self.sigma_scheduler.sigma(t)
+
+    def x0_to_velocity(self, x_0, x_t, t, eps=None):
+        d_alpha, d_sigma = self.sigma_scheduler.velocity(t)
+        if eps is None:
+            eps = self.x0_to_eps(x_0, x_t, t)
+        return d_alpha*x_0+d_sigma*eps
+
+    def eps_to_x0(self, eps, x_t, t):
+        return (x_t-self.sigma_scheduler.sigma(t)*eps)/self.sigma_scheduler.alpha(t)
+
+    def eps_to_velocity(self, eps, x_t, t, x_0=None):
+        d_alpha, d_sigma = self.sigma_scheduler.velocity(t)
+        if x_0 is None:
+            x_0 = self.eps_to_x0(eps, x_t, t)
+        return d_alpha*x_0+d_sigma*eps
+
+    def velocity_to_eps(self, v_pred, x_t, t):
+        alpha = self.sigma_scheduler.alpha(t)
+        sigma = self.sigma_scheduler.sigma(t)
+        d_alpha, d_sigma = self.sigma_scheduler.velocity(t)
+        return (alpha*v_pred-d_alpha*x_t)/(d_sigma*alpha-d_alpha*sigma)
+
+    def velocity_to_x0(self, v_pred, x_t, t):
+        alpha = self.sigma_scheduler.alpha(t)
+        sigma = self.sigma_scheduler.sigma(t)
+        d_alpha, d_sigma = self.sigma_scheduler.velocity(t)
+        return (sigma*v_pred-d_sigma*x_t)/(d_alpha*sigma-d_sigma*alpha)

hcpdiff/diffusion/sampler/diffusers.py

@@ -18,31 +18,24 @@ class DiffusersSampler(BaseSampler):
         self.scheduler = scheduler
         self.eta = eta
 
-    def c_in(self, sigma):
-        one = torch.ones_like(sigma)
+        self.sigma_scheduler.c_in = self.c_in
+
+    def c_in(self, t):
+        one = torch.ones_like(t)
         if hasattr(self.scheduler, '_step_index'):
             self.scheduler._step_index = None
-        return self.scheduler.scale_model_input(one, sigma)
-
-    def c_out(self, sigma):
-        return -sigma
-
-    def c_skip(self, sigma):
-        if self.c_in(sigma) == 1.:  # DDPM model
-            return (sigma**2+1).sqrt()  # 1/sqrt(alpha_)
-        else:  # EDM model
-            return 1.
+        return self.scheduler.scale_model_input(one, t)
 
     def get_timesteps(self, N_steps, device='cuda'):
         self.scheduler.set_timesteps(N_steps, device=device)
-        return self.scheduler.timesteps
+        return self.scheduler.timesteps / self.sigma_scheduler.num_timesteps # Normalize timesteps to [0, 1]
 
     def init_noise(self, shape, device='cuda', dtype=torch.float32):
         return randn_tensor(shape, generator=self.generator, device=device, dtype=dtype)*self.scheduler.init_noise_sigma
 
-    def add_noise(self, x, sigma):
+    def add_noise(self, x, t):
         noise = randn_tensor(x.shape, generator=self.generator, device=x.device, dtype=x.dtype)
-        return self.scheduler.add_noise(x, noise, sigma), noise
+        return self.scheduler.add_noise(x, noise, t), noise
 
     def prepare_extra_step_kwargs(self, scheduler, generator, eta):
         # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
@@ -61,6 +54,7 @@ class DiffusersSampler(BaseSampler):
             extra_step_kwargs["generator"] = generator
         return extra_step_kwargs
 
-    def denoise(self, x_t, sigma, eps=None, generator=None):
+    def denoise(self, x_t, t, eps=None, generator=None):
+        t_in = self.sigma_scheduler.c_noise(t)
         extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator, self.eta)
-        return self.scheduler.step(eps, sigma, x_t, **extra_step_kwargs).prev_sample
+        return self.scheduler.step(eps, t_in, x_t, **extra_step_kwargs).prev_sample

hcpdiff/diffusion/sampler/sigma_scheduler/__init__.py

@@ -1,3 +1,5 @@
 from .base import SigmaScheduler
 from .ddpm import DDPMDiscreteSigmaScheduler, DDPMContinuousSigmaScheduler, TimeSigmaScheduler
-from .edm import EDMSigmaScheduler, EDMRefSigmaScheduler
+from .edm import EDMSigmaScheduler, EDMTimeRescaleScheduler
+from .flow import FlowSigmaScheduler
+from .zero_terminal import ZeroTerminalScheduler

hcpdiff/diffusion/sampler/sigma_scheduler/base.py

@@ -3,12 +3,87 @@ from typing import Union, Tuple
 import torch
 
 class SigmaScheduler:
+    def scale_t(self, t):
+        return t
 
-    def get_sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+    def sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
         '''
+        x(t) = \alpha(t)*x(0) + \sigma(t)*eps
         :param t: 0-1, rate of time step
         '''
         raise NotImplementedError
 
-    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)) -> Tuple[torch.Tensor, torch.Tensor]:
+    def alpha(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        '''
+        x(t) = \alpha(t)*x(0) + \sigma(t)*eps
+        :param t: 0-1, rate of time step
+        '''
+        raise NotImplementedError
+
+    def velocity(self, t: Union[float, torch.Tensor], dt=1e-8, normlize=True) -> Tuple[torch.Tensor, torch.Tensor]:
+        '''
+        v(t) = dx(t)/dt = d\alpha(t)/dt * x(0) + d\sigma(t)/dt *eps
+        :param t: 0-1, rate of time step
+        :return: d\alpha(t)/dt, d\sigma(t)/dt
+        '''
+        d_alpha = (self.alpha(t+dt)-self.alpha(t))/dt
+        d_sigma = (self.sigma(t+dt)-self.sigma(t))/dt
+        if normlize:
+            norm = torch.sqrt(d_alpha**2+d_sigma**2)
+            return d_alpha/norm, d_sigma/norm
+        else:
+            return d_alpha, d_sigma
+
+    @property
+    def sigma_start(self):
+        return self.sigma(0)
+
+    @property
+    def sigma_end(self):
+        return self.sigma(1)
+
+    @property
+    def alpha_start(self):
+        return self.alpha(0)
+
+    @property
+    def alpha_end(self):
+        return self.alpha(1)
+
+    def alpha_to_sigma(self, alpha):
         raise NotImplementedError
+
+    def sigma_to_alpha(self, sigma):
+        raise NotImplementedError
+
+    def c_in(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            return 1.
+        else:
+            return torch.ones_like(t, dtype=torch.float32)
+
+    def c_skip(self, t: Union[float, torch.Tensor]):
+        '''
+        \hat{x}(0) = c_skip*x(t) + c_out*f(x(t))
+        :param t: 0-1, rate of time step
+        '''
+        return 1./self.alpha(t)
+
+    def c_out(self, t: Union[float, torch.Tensor]):
+        '''
+        \hat{x}(0) = c_skip*x(t) + c_out*f(x(t))
+        :param t: 0-1, rate of time step
+        '''
+        return -self.sigma(t)/self.alpha(t)
+    
+    def c_noise(self, t: Union[float, torch.Tensor]):
+        return t
+
+    def sample(self, min_t=0.0, max_t=1.0, shape=(1,)) -> torch.Tensor:
+        if isinstance(min_t, float):
+            min_t = torch.full(shape, min_t)
+        if isinstance(max_t, float):
+            max_t = torch.full(shape, max_t)
+
+        t = torch.lerp(min_t, max_t, torch.rand_like(min_t))
+        return t