Stable Diffusion EDM Sampler详细解释

EDM (Euler Discretization with Momentum)： EDM代表了欧拉离散化并带有动量的方法，它通常是对连续时间扩散过程进行数值积分的一种变体，通过引入动量项来改进收敛性和稳定性。
    
   在 EDMSampler 中，重要原理可以概括为以下几点：
   

• 扩散过程：
  
  
  
  • 在训练阶段，扩散模子学习怎样将带有不同水平噪声的数据徐徐还原至无噪声的状态。
    
  • 这个过程可以视为一系列连续的概率转换，每个步调对应一个噪声水平（即 sigma）。
    
  
  
• Euler方法：
  
  
  
  • 该类中的 sampler_step 方法使用了Euler方法进行离散化处置惩罚，这是一种数值积分技能，用于近似解决微分方程。
    
  • 对于扩散模子来说，这个离散化的一步就是根据当前噪声水平和下一步的目标噪声水平，盘算并应用相应的更新到采样状态上。
    
  
  
• 动态调整：
  
  
  
  • EDMSampler 包含了一些自定义参数如 s_churn, s_tmin, s_tmax，它们影响着每一步中加入额外随机性的程度（通过 gamma 参数控制）。
    
  • 当噪声水平在特定范围内时（由 s_tmin 和 s_tmax 定义），会引入额外的随机扰动，这有助于进步采样的多样性和避免陷入局部最优解。
    
  
  
• 去噪与步进：
  
  
  
  • 每一步迭代包括：盘算修正后的噪声水平、给当前状态注入噪声、使用预训练的扩散模子进行去噪，并执行欧拉步进更新状态。
    
  • 可能还包含一个校正步调来进一步优化采样结果的质量。
    
  
  
• 整体流程：
  
  
  
  • __call__ 方法被计划成可以直接调用启动采样过程的形式，它按照预定的噪声水平序列逐级降低噪声强度，直至最终生成高质量的新样本。
    
  
  

   决定能否看懂代码的重点！！！
   euler_step = self.euler_step(x, d, dt) 这段代码是在求解微分方程
   
   

• x 是当前状态变量，即带有某个噪声级别的样本。
  
• d 通常表示在这个噪声级别下对样本进行去噪操作后的结果与原噪声样本之间的差异。
  
• dt 是时间步长或者说噪声水平的变革量，即 next_sigma - sigma_hat。
  

   

   说得简单点就是，花个坐标轴x,y; 现在x轴上某个点的值已知（当前带有噪声的样本），现在x轴方向变革了dt也已知，导致了y轴的变革量d也已知，求斜率（斜率即微分方程的解）
  

1. class EDMSampler(SingleStepDiffusionSampler):
2.     def __init__(self, s_churn=0.0, s_tmin=0.0, s_tmax=float("inf"), s_noise=1.0, *args, **kwargs):
3.         super().__init__(*args, **kwargs)
5.         self.s_churn = s_churn
6.         self.s_tmin = s_tmin
7.         self.s_tmax = s_tmax
8.         self.s_noise = s_noise
10.     def sampler_step(self, sigma, next_sigma, model, x, cond, uc=None, gamma=0.0, **kwargs):
11.         sigma_hat = sigma * (gamma + 1.0)
12.         if gamma > 0:
13.             eps = Tensor(np.random.randn(*x.shape), x.dtype) * self.s_noise
14.             x = x + eps * append_dims(sigma_hat**2 - sigma**2, x.ndim) ** 0.5
16.         denoised = self.denoise(x, model, sigma_hat, cond, uc, **kwargs)
17.         d = to_d(x, sigma_hat, denoised)
18.         dt = append_dims(next_sigma - sigma_hat, x.ndim)
20.         euler_step = self.euler_step(x, d, dt) #核心，在解微分方程
21.         x = self.possible_correction_step(euler_step, x, d, dt, next_sigma, model, cond, uc)
22.         return x
24.     def __call__(self, model, x, cond, uc=None, num_steps=None, **kwargs):
25.         x = ops.cast(x, ms.float32)
26.         x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(x, cond, uc, num_steps)
28.         for i in self.get_sigma_gen(num_sigmas):
29.             gamma = (
30.                 min(self.s_churn / (num_sigmas - 1), 2**0.5 - 1) if self.s_tmin <= sigmas[i] <= self.s_tmax else 0.0
31.             )
32.             x = self.sampler_step(s_in * sigmas[i], s_in * sigmas[i + 1], model, x, cond, uc, gamma, **kwargs)
34.         return x
39. #只想搞懂原理的话，下面的依赖可以不看
42. class SingleStepDiffusionSampler(BaseDiffusionSampler):
43.     def sampler_step(self, sigma, next_sigma, model, x, cond, uc=None, gamma=0.0, **kwargs):
44.         sigma_hat = sigma * (gamma + 1.0)
45.         if gamma > 0:
46.             eps = Tensor(np.random.randn(*x.shape), x.dtype) * self.s_noise
47.             x = x + eps * append_dims(sigma_hat**2 - sigma**2, x.ndim) ** 0.5
49.         denoised = self.denoise(x, model, sigma_hat, cond, uc, **kwargs)
50.         d = to_d(x, sigma_hat, denoised)
51.         dt = append_dims(next_sigma - sigma_hat, x.ndim)
53.         euler_step = self.euler_step(x, d, dt)
54.         x = euler_step
55.         return x
57.     def euler_step(self, x, d, dt):
58.         return x + dt * d
61. class BaseDiffusionSampler:
62.     def __init__(
63.         self,
64.         discretization_config: Union[Dict, ListConfig, OmegaConf],
65.         num_steps: Union[int, None] = None,
66.         guider_config: Union[Dict, ListConfig, OmegaConf, None] = None,
67.         verbose: bool = False,
68.     ):
69.         self.num_steps = num_steps
70.         self.discretization = instantiate_from_config(discretization_config)
71.         self.guider = instantiate_from_config(
72.             default(
73.                 guider_config,
74.                 DEFAULT_GUIDER,
75.             )
76.         )
77.         self.verbose = verbose
79.     def prepare_sampling_loop(self, x, cond, uc=None, num_steps=None):
80.         sigmas = self.discretization(self.num_steps if num_steps is None else num_steps)
82.         uc = default(uc, cond)
84.         x *= Tensor(np.sqrt(1.0 + sigmas[0] ** 2.0), x.dtype)
85.         num_sigmas = len(sigmas)
87.         s_in = ops.ones((x.shape[0],), x.dtype)
89.         return x, s_in, sigmas, num_sigmas, cond, uc
91.     def denoise(self, x, model, sigma, cond, uc, **kwargs):
92.         noised_input, sigmas, cond = self.guider.prepare_inputs(x, sigma, cond, uc)
93.         cond = model.openai_input_warpper(cond)
94.         c_skip, c_out, c_in, c_noise = model.denoiser(sigmas, noised_input.ndim)
95.         model_output = model.model(noised_input * c_in, c_noise, **cond, **kwargs)
96.         model_output = model_output.astype(ms.float32)
97.         denoised = model_output * c_out + noised_input * c_skip
98.         denoised = self.guider(denoised, sigma)
99.         return denoised
101.     def get_sigma_gen(self, num_sigmas):
102.         sigma_generator = range(num_sigmas - 1)
103.         if self.verbose:
104.             print("#" * 30, " Sampling setting ", "#" * 30)
105.             print(f"Sampler: {self.__class__.__name__}")
106.             print(f"Discretization: {self.discretization.__class__.__name__}")
107.             print(f"Guider: {self.guider.__class__.__name__}")
108.             sigma_generator = tqdm(
109.                 sigma_generator,
110.                 total=(num_sigmas - 1),
111.                 desc=f"Sampling with {self.__class__.__name__} for {(num_sigmas - 1)} steps",
112.             )
113.         return sigma_generator

复制代码

   
   
    

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。