hcpdiff 2.2.1__py3-none-any.whl → 2.3.1__py3-none-any.whl

hcpdiff/diffusion/sampler/sigma_scheduler/ddpm.py

@@ -1,17 +1,22 @@
-import torch
 import math
-from typing import Union, Tuple
-from hcpdiff.utils import linear_interp
+from typing import Union, Tuple, Callable
+
+import torch
+
+from hcpdiff.utils import invert_func
 from .base import SigmaScheduler
 
 class DDPMDiscreteSigmaScheduler(SigmaScheduler):
-    def __init__(self, beta_schedule: str = "scaled_linear", linear_start=0.00085, linear_end=0.0120, num_timesteps=1000):
+    def __init__(self, beta_schedule: str = "scaled_linear", linear_start=0.00085, linear_end=0.0120, num_timesteps=1000, pred_type='eps'):
         super().__init__()
         self.num_timesteps = num_timesteps
         self.betas = self.make_betas(beta_schedule, linear_start, linear_end, num_timesteps)
         alphas = 1.0-self.betas
         self.alphas_cumprod = torch.cumprod(alphas, dim=0)
-        self.sigmas = ((1-self.alphas_cumprod)/self.alphas_cumprod).sqrt()
+
+        self.alphas = self.alphas_cumprod.sqrt()
+        self.sigmas = (1-self.alphas_cumprod).sqrt()
+        self.pred_type = pred_type
 
         # for VLB calculation
         self.alphas_cumprod_prev = torch.cat([alphas.new_tensor([1.0]), self.alphas_cumprod[:-1]])
@@ -22,37 +27,73 @@ class DDPMDiscreteSigmaScheduler(SigmaScheduler):
         # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
         self.posterior_log_variance_clipped = torch.log(torch.cat([self.posterior_variance[1:2], self.posterior_variance[1:]]))
 
+    # def scale_t(self, t):
+    #     return t*(self.num_timesteps-1)
 
     @property
-    def sigma_min(self):
+    def sigma_start(self):
         return self.sigmas[0]
 
     @property
-    def sigma_max(self):
+    def sigma_end(self):
         return self.sigmas[-1]
 
-    def get_sigma(self, t: Union[float, torch.Tensor]):
+    @property
+    def alpha_start(self):
+        return self.alphas[0]
+
+    @property
+    def alpha_end(self):
+        return self.alphas[-1]
+
+    def sigma(self, t: Union[float, torch.Tensor]):
         if isinstance(t, float):
             t = torch.tensor(t)
-        return self.sigmas[(t*len(self.sigmas)).long()]
+        self.sigmas = self.sigmas.to(t.device)
+        return self.sigmas[((t*self.num_timesteps).round().long()).clip(min=0, max=self.num_timesteps-1)]
 
-    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
-        if isinstance(min_rate, float):
-            min_rate = torch.full(shape, min_rate)
-        if isinstance(max_rate, float):
-            max_rate = torch.full(shape, max_rate)
+    def alpha(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            t = torch.tensor(t)
+        self.alphas = self.alphas.to(t.device)
+        return self.alphas[((t*self.num_timesteps).round().long()).clip(min=0, max=self.num_timesteps-1)]
+    
+    def c_noise(self, t: Union[float, torch.Tensor]):
+        return (t*self.num_timesteps).round()
 
-        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
-        t_scale = (t*(self.num_timesteps-1e-5)).long()  # [0, num_timesteps-1)
-        return self.sigmas[t_scale], t
+    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.sigma(t)
+        d_sigma = self.alpha(t)
+        if normlize:
+            norm = torch.sqrt(d_alpha**2+d_sigma**2)
+            return d_alpha/norm, d_sigma/norm
+        else:
+            return d_alpha, d_sigma
 
     def sigma_to_t(self, sigma: Union[float, torch.Tensor]):
-        t = (self.sigmas-sigma).abs().argmin()
-        return t/self.num_timesteps
+        ref_t = np.linspace(0, 1, len(self.sigmas))
+        t = torch.tensor(np.interp(sigma.cpu().clip(min=1e-8).log().numpy(), self.sigmas, ref_t))
+        return t
+
+    def alpha_to_t(self, alpha: Union[float, torch.Tensor]):
+        ref_t = np.linspace(0, 1, len(self.alphas))
+        t = torch.tensor(np.interp(alpha.cpu().clip(min=1e-8).log().numpy(), self.alphas, ref_t))
+        return t
+
+    def alpha_to_sigma(self, alpha):
+        return torch.sqrt(1 - alpha**2)
+
+    def sigma_to_alpha(self, sigma):
+        return torch.sqrt(1 - sigma**2)
 
     def get_post_mean(self, t, x_0, x_t):
         t = (t*len(self.sigmas)).long()
-        return self.posterior_mean_coef1[t].view(-1, 1, 1, 1).to(t.device)*x_0 + self.posterior_mean_coef2[t].view(-1, 1, 1, 1).to(t.device)*x_t
+        return self.posterior_mean_coef1[t].view(-1, 1, 1, 1).to(t.device)*x_0+self.posterior_mean_coef2[t].view(-1, 1, 1, 1).to(t.device)*x_t
 
     def get_post_log_var(self, t, x_t_var=None):
         t = (t*len(self.sigmas)).long()
@@ -66,7 +107,6 @@ class DDPMDiscreteSigmaScheduler(SigmaScheduler):
             model_log_variance = frac*max_log+(1-frac)*min_log
             return model_log_variance
 
-
     @staticmethod
     def betas_for_alpha_bar(
         num_diffusion_timesteps,
@@ -130,50 +170,154 @@ class DDPMDiscreteSigmaScheduler(SigmaScheduler):
         else:
             raise NotImplementedError(f"{beta_schedule} does is not implemented.")
 
-class DDPMContinuousSigmaScheduler(DDPMDiscreteSigmaScheduler):
+class DDPMContinuousSigmaScheduler(SigmaScheduler):
+    def __init__(self, beta_schedule: str = "scaled_linear", linear_start=0.00085, linear_end=0.0120, t_base=1000):
+        self.alpha_bar_fn = self.make_alpha_bar_fn(beta_schedule, linear_start, linear_end)
+        self.t_base = t_base  # base time step for continuous product
+
+    def continuous_product(self, alpha_fn: Callable[[torch.Tensor], torch.Tensor], t: torch.Tensor):
+        '''
+
+        :param alpha_fn: alpha function
+        :param t: timesteps with shape [B]
+        :return: [B]
+        '''
+        bins = torch.linspace(0, 1, self.t_base, dtype=torch.float32).unsqueeze(0)
+        t_grid = bins*t.float().unsqueeze(1)  # [B, num_bins]
+        alpha_vals = alpha_fn(t_grid)
+
+        if torch.any(alpha_vals<=0):
+            raise ValueError("alpha(t) must > 0 to avoid log(≤0).")
+
+        log_term = torch.log(alpha_vals)  # [B, num_bins]
+        dt = t_grid[:, 1]-t_grid[:, 0]  # [B]
+        integral = torch.cumsum((log_term[:, -1]+log_term[:, 1:])/2*dt.unsqueeze(1), dim=1)  # [B]
+        x_vals = torch.exp(integral)
+        return x_vals
+
+    @staticmethod
+    def alpha_bar_linear(beta_s, beta_e, t, N=1000):
+        A = beta_e-beta_s
+        B = 1-beta_s
+        B_At = B-A*t
+
+        # 避免数值不稳定
+        eps = 1e-12
+        B = torch.clamp(B, min=eps)
+        B_At = torch.clamp(B_At, min=eps)
+
+        term = (B*torch.log(B)-B_At*torch.log(B_At)-A*t)
+        return torch.exp(N*term/A)
+
+    @staticmethod
+    def alpha_bar_scaled_linear(beta_s, beta_e, t, N=1000):
+        sqrt_bs = torch.sqrt(beta_s)
+        sqrt_be = torch.sqrt(beta_e)
+        a = sqrt_be-sqrt_bs
+        b = sqrt_bs
+        u0 = b
+        u1 = a*t+b
+
+        eps = 1e-12
+
+        def safe_log1m(u2):
+            return torch.log(torch.clamp(1-u2, min=eps))
+
+        def safe_log_frac(u):
+            return torch.log(torch.clamp(1+u, min=eps))-torch.log(torch.clamp(1-u, min=eps))
 
-    def get_sigma(self, t: Union[float, torch.Tensor]):
+        term1 = u1*safe_log1m(u1**2)
+        term2 = 0.5*safe_log_frac(u1)
+        term3 = u0*safe_log1m(u0**2)
+        term4 = 0.5*safe_log_frac(u0)
+
+        return torch.exp(N*(term1+term2-term3-term4)/a)
+
+    def make_alpha_bar_fn(self, beta_schedule, beta_start, beta_end, alpha_fn=None):
+        if alpha_fn is not None:
+            return lambda t, alpha_fn_=alpha_fn:self.continuous_product(alpha_fn_(t), t)
+        elif beta_schedule == "linear":
+            return lambda t:self.alpha_bar_linear(beta_start, beta_end, t)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            return lambda t:self.alpha_bar_scaled_linear(beta_start, beta_end, t)
+        elif beta_schedule == "squaredcos_cap_v2":
+            return lambda t:torch.cos((t+0.008)/1.008*math.pi/2)**2
+        elif beta_schedule == "sigmoid":
+            # GeoDiff sigmoid schedule
+            alpha_fn = lambda t:1-torch.sigmoid(torch.lerp(torch.full_like(t, -6), torch.full_like(t, 6), t))*(beta_end-beta_start)+beta_start
+            return lambda t, alpha_fn_=alpha_fn:self.continuous_product(alpha_fn_(t), t)
+        else:
+            raise NotImplementedError(f"{beta_schedule} does is not implemented.")
+
+    def sigma(self, t: Union[float, torch.Tensor]):
         if isinstance(t, float):
-            t = torch.tensor(t)
-        return linear_interp(self.sigmas, t)
+            t = torch.tensor([t])
+        alpha_cumprod = self.alpha_bar_fn(t)
+        return torch.sqrt(1-alpha_cumprod)
 
-    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
-        if isinstance(min_rate, float):
-            min_rate = torch.full(shape, min_rate)
-        if isinstance(max_rate, float):
-            max_rate = torch.full(shape, max_rate)
+    def alpha(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            t = torch.tensor([t])
+        alpha_cumprod = self.alpha_bar_fn(t)
+        return torch.sqrt(alpha_cumprod)
 
-        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
-        t_scale = (t*(self.num_timesteps-1-1e-5))  # [0, num_timesteps-1)
+    def c_noise(self, t: Union[float, torch.Tensor]):
+        return t*self.t_base
 
-        return linear_interp(self.sigmas, t_scale), t
+    @property
+    def sigma_start(self):
+        return self.sigma(0)
 
-    def sigma_to_t(self, sigma: Union[float, torch.Tensor]):
-        diff = self.sigmas-sigma
-        diff[diff<0] = float('inf')
-        t0 = diff.argmin().clamp(0, self.num_timesteps-2)
-        return t0 + diff.min()/(self.sigmas[t0+1]-self.sigmas[t0])
+    @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_t(self, alpha, t_min=0.0, t_max=1.0, tol=1e-5, max_iter=100):
+        """
+        alpha: [B]
+        :return: t [B]
+        """
+        return invert_func(self.alpha, alpha, t_min, t_max, tol, max_iter)
+
+    def sigma_to_t(self, sigma, t_min=0.0, t_max=1.0, tol=1e-5, max_iter=100):
+        """
+        sigma: [B]
+        :return: t [B]
+        """
+        return invert_func(self.sigma, sigma, t_min, t_max, tol, max_iter)
 
 class TimeSigmaScheduler(SigmaScheduler):
     def __init__(self, num_timesteps=1000):
         super().__init__()
         self.num_timesteps = num_timesteps
 
-    def get_sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+    def sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
         '''
         :param t: 0-1, rate of time step
         '''
-        return t
-
-    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)) -> Tuple[torch.Tensor, torch.Tensor]:
-        if isinstance(min_rate, float):
-            min_rate = torch.full(shape, min_rate)
-        if isinstance(max_rate, float):
-            max_rate = torch.full(shape, max_rate)
+        if isinstance(t, float):
+            t = torch.tensor(t)
+        return ((t*self.num_timesteps).round().long()).clip(min=0, max=self.num_timesteps-1)
 
-        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
-        t_scale = (t*(self.num_timesteps-1e-5)).long()  # [0, num_timesteps-1)
-        return t_scale, t
+    def alpha(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        '''
+        :param t: 0-1, rate of time step
+        '''
+        if isinstance(t, float):
+            t = torch.tensor(t)
+        return ((t*self.num_timesteps).round().long()).clip(min=0, max=self.num_timesteps-1)
+    
+    def c_noise(self, t: Union[float, torch.Tensor]):
+        return (t*self.num_timesteps).round()
 
 if __name__ == '__main__':
     from matplotlib import pyplot as plt

hcpdiff/diffusion/sampler/sigma_scheduler/edm.py

@@ -1,19 +1,18 @@
-from typing import Union
+from typing import Union, Tuple
 
-import torch
 import numpy as np
+import torch
 
 from .base import SigmaScheduler
 
 class EDMSigmaScheduler(SigmaScheduler):
-    def __init__(self, sigma_min=0.002, sigma_max=80.0, rho=7.0, num_timesteps=1000):
-        self.sigma_min = torch.tensor(sigma_min)
-        self.sigma_max = torch.tensor(sigma_max)
+    def __init__(self, sigma_min=0.002, sigma_max=80.0, sigma_data=0.5, rho=7.0):
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma_data = sigma_data
         self.rho = rho
 
-        self.num_timesteps=num_timesteps
-
-    def get_sigma(self, t: Union[float, torch.Tensor]):
+    def sigma_edm(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
         if isinstance(t, float):
             t = torch.tensor(t)
 
@@ -21,28 +20,106 @@ class EDMSigmaScheduler(SigmaScheduler):
         max_inv_rho = self.sigma_max**(1/self.rho)
         return torch.lerp(min_inv_rho, max_inv_rho, t)**self.rho
 
-    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
-        if isinstance(min_rate, float):
-            min_rate = torch.full(shape, min_rate)
-        if isinstance(max_rate, float):
-            max_rate = torch.full(shape, max_rate)
-
-        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
-        return self.get_sigma(t), t
-
-class EDMRefSigmaScheduler(EDMSigmaScheduler):
-    def __init__(self, ref_scheduler, sigma_min=0.002, sigma_max=80.0, rho=7.0, num_timesteps=1000):
-        super().__init__(sigma_min, sigma_max, rho, num_timesteps=num_timesteps)
-        self.ref_sigmas = ref_scheduler.sigmas.cpu().clip(min=1e-8).log().numpy()
-        self.ref_t = np.linspace(0, 1, len(self.ref_sigmas))
-
-    def sample_sigma(self, min_rate=0.0, max_rate=1.0, shape=(1,)):
-        if isinstance(min_rate, float):
-            min_rate = torch.full(shape, min_rate)
-        if isinstance(max_rate, float):
-            max_rate = torch.full(shape, max_rate)
-
-        t = torch.lerp(min_rate, max_rate, torch.rand_like(min_rate))
-        sigma = self.get_sigma(t)
-        t_rect = torch.tensor(np.interp(sigma.cpu().clip(min=1e-8).log().numpy(), self.ref_sigmas, self.ref_t))
-        return sigma, t_rect
+    def sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        '''
+        x_t = c_in(t) * (x(0) + \sigma(t)*eps), eps~N(0,I)
+        '''
+        if isinstance(t, float):
+            t = torch.tensor(t)
+
+        sigma_edm = self.sigma_edm(t)
+        return sigma_edm/torch.sqrt(sigma_edm**2+self.sigma_data**2)
+
+    def alpha(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        '''
+        x_t = c_in(t) * (x(0) + \sigma(t)*eps), eps~N(0,I)
+        '''
+        if isinstance(t, float):
+            t = torch.tensor(t)
+
+        sigma_edm = self.sigma_edm(t)
+        return 1./torch.sqrt(sigma_edm**2+self.sigma_data**2)
+
+    def c_skip(self, t: Union[float, torch.Tensor]):
+        '''
+        \hat{x}(0) = c_skip(t)*(x(t)/c_in(t)) + c_out(t)*f(x(t))
+        :param t: 0-1, rate of time step
+        '''
+        sigma_edm = self.sigma_edm(t)
+        return self.sigma_data**2/torch.sqrt(sigma_edm**2+self.sigma_data**2)
+
+    def c_out(self, t: Union[float, torch.Tensor]):
+        '''
+        \hat{x}(0) = c_skip(t)*(x(t)/c_in(t)) + c_out(t)*f(x(t))
+        :param t: 0-1, rate of time step
+        '''
+        sigma_edm = self.sigma_edm(t)
+        return (self.sigma_data*sigma_edm)/torch.sqrt(sigma_edm**2+self.sigma_data**2)
+
+    def c_noise(self, t: Union[float, torch.Tensor]):
+        sigma_edm = self.sigma_edm(t)
+        return sigma_edm.log()/4
+
+    @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):
+        return torch.sqrt(1 - (alpha*self.sigma_data)**2)
+
+    def sigma_to_alpha(self, sigma):
+        return torch.sqrt(1 - sigma**2)/self.sigma_data
+
+class EDMTimeRescaleScheduler(EDMSigmaScheduler):
+    def __init__(self, ref_scheduler: SigmaScheduler, sigma_min=0.002, sigma_max=80.0, rho=7.0):
+        super().__init__(sigma_min, sigma_max, rho)
+        self.ref_scheduler = ref_scheduler
+
+    def scale_t(self, t):
+        ref_t = torch.linspace(0, 1, 1000)
+        alphas = self.alpha(ref_t)
+        sigmas = self.sigma(ref_t)
+        sigmas_edm = sigmas/alphas
+        sigma_edm = self.sigma_edm(t)
+        t = np.interp(sigma_edm.cpu().clip(min=1e-8).log().numpy(), sigmas_edm, ref_t.numpy())
+        return torch.tensor(t)
+
+    def sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        return self.ref_scheduler.sigma(t)
+
+    def alpha(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        return self.ref_scheduler.alpha(t)
+
+    def velocity(self, t: Union[float, torch.Tensor], dt=1e-8, normlize=True) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.ref_scheduler.velocity(t, dt=dt, normlize=normlize)
+
+    def c_skip(self, t: Union[float, torch.Tensor]):
+        return self.ref_scheduler.c_skip(t)
+
+    def c_out(self, t: Union[float, torch.Tensor]):
+        return self.ref_scheduler.c_out(t)
+    
+    def c_noise(self, t: Union[float, torch.Tensor]):
+        return self.ref_scheduler.c_noise(t)
+
+    def sample(self, min_t=0.0, max_t=1.0, shape=(1,)):
+        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))
+        t = self.scale_t(t)
+        return t

hcpdiff/diffusion/sampler/sigma_scheduler/flow.py

@@ -0,0 +1,74 @@
+from typing import Union, Tuple
+
+import torch
+
+from .base import SigmaScheduler
+
+class FlowSigmaScheduler(SigmaScheduler):
+    def __init__(self, t_start=0, t_end=1):
+        super().__init__()
+        self.t_start = t_start
+        self.t_end = t_end
+
+    def sigma(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            t = torch.tensor([t])
+        t = (self.t_end-self.t_start)*t+self.t_start
+        return t
+
+    def alpha(self, t: Union[float, torch.Tensor]):
+        if isinstance(t, float):
+            t = torch.tensor([t])
+        t = (self.t_end-self.t_start)*t+self.t_start
+        return 1-t
+
+    def velocity(self, t: Union[float, torch.Tensor], dt=1e-8, normlize=False) -> 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
+        '''
+        if isinstance(t, float):
+            t = torch.tensor([t])
+        d_alpha = -torch.ones_like(t)
+        d_sigma = torch.ones_like(t)
+        if normlize:
+            norm = torch.sqrt(d_alpha**2+d_sigma**2)
+            return d_alpha/norm, d_sigma/norm
+        else:
+            return d_alpha, d_sigma
+
+    def alpha_to_t(self, alphas):
+        """
+        alphas: [B]
+        :return: t [B]
+        """
+        return alphas
+
+    def sigma_to_t(self, sigmas):
+        """
+        sigmas: [B]
+        :return: t [B]
+        """
+        return 1-sigmas
+
+    def alpha_to_sigma(self, alpha):
+        return 1-alpha
+
+    def sigma_to_alpha(self, sigma):
+        return 1-sigma
+
+    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.
+
+    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
+        '''
+        sigma = self.sigma(t)
+        return -sigma

hcpdiff/diffusion/sampler/sigma_scheduler/zero_terminal.py

@@ -0,0 +1,22 @@
+from typing import Union
+
+import torch
+from .base import SigmaScheduler
+
+class ZeroTerminalScheduler(SigmaScheduler):
+    def __init__(self, ref_scheduler: SigmaScheduler, eps=1e-4):
+        self.ref_scheduler = ref_scheduler
+        self.eps = eps
+
+    def alpha(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        alpha_0 = self.ref_scheduler.alpha_start
+        alpha_T = self.ref_scheduler.alpha_end
+        alpha = self.ref_scheduler.alpha(t)
+        return (alpha - alpha_T)*(alpha_0-self.eps)/(alpha_0 - alpha_T) + self.eps
+
+    def sigma(self, t: Union[float, torch.Tensor]) -> torch.Tensor:
+        try:
+            alpha = self.alpha(t)
+            return self.ref_scheduler.alpha_to_sigma(alpha)
+        except NotImplementedError:
+            raise NotImplementedError(f'{type(self.ref_scheduler)} cannot be a "ZeroTerminalScheduler"!')

hcpdiff/easy/cfg/sd15_train.py

@@ -1,18 +1,17 @@
 import torch
-from rainbowneko.ckpt_manager import ckpt_saver, LAYERS_TRAINABLE, NekoPluginSaver, SafeTensorFormat
-from rainbowneko.data import RatioBucket, FixedBucket
-from rainbowneko.parser import CfgWDPluginParser, neko_cfg, CfgWDModelParser, disable_neko_cfg
-from rainbowneko.utils import ConstantLR, Path_Like
-
 from hcpdiff.ckpt_manager import LoraWebuiFormat
 from hcpdiff.data import TextImagePairDataset, Text2ImageSource, StableDiffusionHandler
 from hcpdiff.data import VaeCache
 from hcpdiff.easy import SD15_auto_loader
 from hcpdiff.models import SD15Wrapper, TEHookCFG
 from hcpdiff.models.lora_layers_patch import LoraLayer
+from rainbowneko.ckpt_manager import ckpt_saver, NekoOptimizerSaver, LAYERS_TRAINABLE, NekoPluginSaver, SafeTensorFormat
+from rainbowneko.data import RatioBucket, FixedBucket
+from rainbowneko.parser import CfgWDPluginParser, neko_cfg, CfgWDModelParser, disable_neko_cfg
+from rainbowneko.utils import ConstantLR, Path_Like
 
 @neko_cfg
-def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int = 500, lr: float = 1e-5, clip_skip: int = 0,
+def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int = 500, save_optimizer=False, lr: float = 1e-5, clip_skip: int = 0,
                     dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0, name: str = 'SD15'):
     if low_vram:
         from bitsandbytes.optim import AdamW8bit
@@ -20,6 +19,17 @@ def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
     else:
         optimizer = torch.optim.AdamW(_partial_=True)
 
+    ckpt_saver_dict = dict(
+        SD15=ckpt_saver(
+            ckpt_type='safetensors',
+            target_module='denoiser',
+            layers=LAYERS_TRAINABLE,
+        )
+    )
+
+    if save_optimizer:
+        ckpt_saver_dict['optimizer'] = NekoOptimizerSaver()
+
     from cfgs.train.py import train_base, tuning_base
 
     return dict(
@@ -34,11 +44,7 @@ def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
         ], weight_decay=1e-2),
 
         ckpt_saver=dict(
-            SD15=ckpt_saver(
-                ckpt_type='safetensors',
-                target_module='denoiser',
-                layers=LAYERS_TRAINABLE,
-            )
+            SD15=ckpt_saver_dict
         ),
 
         train=dict(
@@ -68,9 +74,9 @@ def SD15_finetuning(base_model: str, train_steps: int, dataset, save_step: int =
     )
 
 @neko_cfg
-def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int = 200, lr: float = 1e-4, rank: int = 4, alpha: float = None,
-                    clip_skip: int = 0, with_conv: bool = False, dtype: str = 'fp16', low_vram: bool = False, warmup_steps: int = 0,
-                    name: str = 'SD15', save_webui_format=False):
+def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int = 200, save_optimizer=False, lr: float = 1e-4, rank: int = 4,
+                    alpha: float = None, clip_skip: int = 0, with_conv: bool = False, dtype: str = 'fp16', low_vram: bool = False,
+                    warmup_steps: int = 0, name: str = 'SD15', save_webui_format=False):
     with disable_neko_cfg:
         if alpha is None:
             alpha = rank
@@ -101,6 +107,17 @@ def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
     else:
         lora_format = SafeTensorFormat()
 
+    ckpt_saver_dict = dict(
+        _replace_=True,
+        lora_unet=NekoPluginSaver(
+            format=lora_format,
+            target_plugin='lora1',
+        )
+    )
+
+    if save_optimizer:
+        ckpt_saver_dict['optimizer'] = NekoOptimizerSaver()
+
     from cfgs.train.py.examples import SD_FT
 
     return dict(
@@ -118,13 +135,7 @@ def SD15_lora_train(base_model: str, train_steps: int, dataset, save_step: int =
             )
         ), weight_decay=0.1),
 
-        ckpt_saver=dict(
-            _replace_ = True,
-            lora_unet=NekoPluginSaver(
-                format=lora_format,
-                target_plugin='lora1',
-            )
-        ),
+        ckpt_saver=ckpt_saver_dict,
 
         train=dict(
             train_steps=train_steps,
@@ -181,7 +192,7 @@ def cfg_data_SD_ARB(img_root: Path_Like, batch_size: int = 4, trigger_word: str
     )
 
 @neko_cfg
-def cfg_data_SD_resize_crop(img_root: Path_Like, batch_size: int = 4, trigger_word: str = '', target_size = (512, 512), word_names=None,
+def cfg_data_SD_resize_crop(img_root: Path_Like, batch_size: int = 4, trigger_word: str = '', target_size=(512, 512), word_names=None,
                             prompt_dropout: float = 0, prompt_template: Path_Like = 'prompt_template/caption.txt', loss_weight=1.0):
     if word_names is None:
         word_names = dict(pt1=trigger_word)