hcpdiff 2.2__tar.gz → 2.3__tar.gz

{hcpdiff-2.2 → hcpdiff-2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hcpdiff
-Version: 2.2
+Version: 2.3
 Summary: A universal Diffusion toolbox
 Home-page: https://github.com/IrisRainbowNeko/HCP-Diffusion
 Author: Ziyi Dong
@@ -17,7 +17,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: rainbowneko
+Requires-Dist: rainbowneko>=1.9
 Requires-Dist: diffusers
 Requires-Dist: matplotlib
 Requires-Dist: pyarrow
@@ -262,9 +262,13 @@ hcp_run --cfg cfgs/workflow/text2img_cli.py \
     seed=42
 ```
 
-### Tutorials
+### 📚 Tutorials
 
-🚧 In Development
++ 🧠 [Model Training Guide](https://hcpdiff.readthedocs.io/en/latest/user_guides/train.html)
++ 🔧 [LoRA Training Tutorial](https://hcpdiff.readthedocs.io/enlatest/tutorial/lora.html)
++ 🎨 [Image Generation Guide](https://hcpdiff.readthedocs.io/en/latest/user_guides/workflow.html)
++ ⚙️ [Configuration File Explanation](https://hcpdiff.readthedocs.io/en/latest/user_guides/cfg.html)
++ 🧩 [Model Format Explanation](https://hcpdiff.readthedocs.io/en/latest/user_guides/model_format.html)
 
 ---
 

{hcpdiff-2.2 → hcpdiff-2.3}/README.md

@@ -225,9 +225,13 @@ hcp_run --cfg cfgs/workflow/text2img_cli.py \
     seed=42
 ```
 
-### Tutorials
+### 📚 Tutorials
 
-🚧 In Development
++ 🧠 [Model Training Guide](https://hcpdiff.readthedocs.io/en/latest/user_guides/train.html)
++ 🔧 [LoRA Training Tutorial](https://hcpdiff.readthedocs.io/enlatest/tutorial/lora.html)
++ 🎨 [Image Generation Guide](https://hcpdiff.readthedocs.io/en/latest/user_guides/workflow.html)
++ ⚙️ [Configuration File Explanation](https://hcpdiff.readthedocs.io/en/latest/user_guides/cfg.html)
++ 🧩 [Model Format Explanation](https://hcpdiff.readthedocs.io/en/latest/user_guides/model_format.html)
 
 ---
 

{hcpdiff-2.2 → hcpdiff-2.3}/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-2.3/hcpdiff/ckpt_manager/ckpt.py

@@ -0,0 +1,28 @@
+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=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, 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), pt_name, prefix=self.prefix)

{hcpdiff-2.2 → hcpdiff-2.3}/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-2.2 → hcpdiff-2.3}/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-2.2 → hcpdiff-2.3}/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-2.2 → hcpdiff-2.3}/hcpdiff/diffusion/noise/__init__.py

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

hcpdiff-2.3/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-2.3/hcpdiff/diffusion/sampler/__init__.py

@@ -0,0 +1,4 @@
+from .sigma_scheduler import *
+from .base import BaseSampler, Sampler
+from .VP import VPSampler
+from .diffusers import DiffusersSampler

hcpdiff-2.3/hcpdiff/diffusion/sampler/base.py

@@ -0,0 +1,134 @@
+from typing import Tuple
+
+import torch
+
+from .sigma_scheduler import SigmaScheduler
+
+try:
+    from diffusers.utils import randn_tensor
+except:
+    # new version of diffusers
+    from diffusers.utils.torch_utils import randn_tensor
+
+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.
+
+        :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
+        '''
+
+        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'):
+        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 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_end
+        return self.make_nosie(shape, device, dtype)*sigma
+
+    def add_noise(self, x, t) -> Tuple[torch.Tensor, torch.Tensor]:
+        noise = self.make_nosie(x.shape, device=x.device)
+        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]
+        timesteps = self.sigma_scheduler.sample(shape=(bs,))
+        timesteps = timesteps.to(x.device)
+        noisy_x, noise = self.add_noise(x, timesteps)
+
+        # Sample a random timestep for each image
+        return noisy_x, noise, timesteps
+
+    def denoise(self, x, sigma, eps=None, generator=None):
+        raise NotImplementedError
+
+    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-2.2 → hcpdiff-2.3}/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-2.3/hcpdiff/diffusion/sampler/sigma_scheduler/__init__.py

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

hcpdiff-2.3/hcpdiff/diffusion/sampler/sigma_scheduler/base.py

@@ -0,0 +1,89 @@
+from typing import Union, Tuple
+
+import torch
+
+class SigmaScheduler:
+    def scale_t(self, t):
+        return t
+
+    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 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