87 lines
3.5 KiB
Markdown
87 lines
3.5 KiB
Markdown
# Documentation
|
|
- Class name: ModelSamplingDiscrete
|
|
- Category: sampling
|
|
- Output node: False
|
|
- Repo Ref: https://github.com/comfyanonymous/ComfyUI
|
|
|
|
ModelSamplingDiscrete 是一个用于生成模型中离散采样的 PyTorch 模块类。它通过注册一个 beta 时间表来管理采样过程，这对于确定采样期间每个时间步的方差至关重要。该类提供了将 sigma 值与时间步转换的方法，促进了离散采样过程的进行。
|
|
|
|
# Input types
|
|
## Optional
|
|
- model_config
|
|
- 模型配置为采样提供了必要的设置，例如 beta 时间表和线性起始/结束值。它对于正确初始化采样参数并影响采样行为至关重要。
|
|
- Comfy dtype: Optional[Dict]
|
|
- Python dtype: Optional[Dict[str, Any]]
|
|
|
|
# Output types
|
|
- sigmas
|
|
- sigmas 参数代表每个时间步的标准差，由注册的 beta 时间表派生而来。它是 ModelSamplingDiscrete 类的关键输出，在采样过程中用于控制噪声水平。
|
|
- Comfy dtype: torch.Tensor
|
|
- Python dtype: torch.Tensor
|
|
|
|
# Usage tips
|
|
- Infra type: GPU
|
|
|
|
# Source code
|
|
```
|
|
class ModelSamplingDiscrete(torch.nn.Module):
|
|
|
|
def __init__(self, model_config=None):
|
|
super().__init__()
|
|
if model_config is not None:
|
|
sampling_settings = model_config.sampling_settings
|
|
else:
|
|
sampling_settings = {}
|
|
beta_schedule = sampling_settings.get('beta_schedule', 'linear')
|
|
linear_start = sampling_settings.get('linear_start', 0.00085)
|
|
linear_end = sampling_settings.get('linear_end', 0.012)
|
|
self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=1000, linear_start=linear_start, linear_end=linear_end, cosine_s=0.008)
|
|
self.sigma_data = 1.0
|
|
|
|
def _register_schedule(self, given_betas=None, beta_schedule='linear', timesteps=1000, linear_start=0.0001, linear_end=0.02, cosine_s=0.008):
|
|
if given_betas is not None:
|
|
betas = given_betas
|
|
else:
|
|
betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
|
|
alphas = 1.0 - betas
|
|
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
|
(timesteps,) = betas.shape
|
|
self.num_timesteps = int(timesteps)
|
|
self.linear_start = linear_start
|
|
self.linear_end = linear_end
|
|
sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
|
|
self.set_sigmas(sigmas)
|
|
|
|
def set_sigmas(self, sigmas):
|
|
self.register_buffer('sigmas', sigmas.float())
|
|
self.register_buffer('log_sigmas', sigmas.log().float())
|
|
|
|
@property
|
|
def sigma_min(self):
|
|
return self.sigmas[0]
|
|
|
|
@property
|
|
def sigma_max(self):
|
|
return self.sigmas[-1]
|
|
|
|
def timestep(self, sigma):
|
|
log_sigma = sigma.log()
|
|
dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]
|
|
return dists.abs().argmin(dim=0).view(sigma.shape).to(sigma.device)
|
|
|
|
def sigma(self, timestep):
|
|
t = torch.clamp(timestep.float().to(self.log_sigmas.device), min=0, max=len(self.sigmas) - 1)
|
|
low_idx = t.floor().long()
|
|
high_idx = t.ceil().long()
|
|
w = t.frac()
|
|
log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]
|
|
return log_sigma.exp().to(timestep.device)
|
|
|
|
def percent_to_sigma(self, percent):
|
|
if percent <= 0.0:
|
|
return 999999999.9
|
|
if percent >= 1.0:
|
|
return 0.0
|
|
percent = 1.0 - percent
|
|
return self.sigma(torch.tensor(percent * 999.0)).item()
|
|
``` |