CavinHuang-comfyui-nodes-docs/docs/ModelSamplingDiscrete.md

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()