Merge pull request #15823 from drhead/patch-3

[Performance] Keep sigmas on CPU
This commit is contained in:
AUTOMATIC1111 2024-06-09 21:18:48 +03:00 committed by GitHub
commit 1d0bb39797
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 10 deletions

View File

@ -115,7 +115,7 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
if scheduler.need_inner_model: if scheduler.need_inner_model:
sigmas_kwargs['inner_model'] = self.model_wrap sigmas_kwargs['inner_model'] = self.model_wrap
sigmas = scheduler.function(n=steps, **sigmas_kwargs, device=shared.device) sigmas = scheduler.function(n=steps, **sigmas_kwargs)
if discard_next_to_last_sigma: if discard_next_to_last_sigma:
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]]) sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])

View File

@ -8,6 +8,12 @@ import numpy as np
from modules import shared from modules import shared
def to_d(x, sigma, denoised):
"""Converts a denoiser output to a Karras ODE derivative."""
return (x - denoised) / sigma
k_diffusion.sampling.to_d = to_d
@dataclasses.dataclass @dataclasses.dataclass
class Scheduler: class Scheduler:
name: str name: str
@ -19,11 +25,11 @@ class Scheduler:
aliases: list = None aliases: list = None
def uniform(n, sigma_min, sigma_max, inner_model, device): def uniform(n, sigma_min, sigma_max, inner_model):
return inner_model.get_sigmas(n) return inner_model.get_sigmas(n)
def sgm_uniform(n, sigma_min, sigma_max, inner_model, device): def sgm_uniform(n, sigma_min, sigma_max, inner_model):
start = inner_model.sigma_to_t(torch.tensor(sigma_max)) start = inner_model.sigma_to_t(torch.tensor(sigma_max))
end = inner_model.sigma_to_t(torch.tensor(sigma_min)) end = inner_model.sigma_to_t(torch.tensor(sigma_min))
sigs = [ sigs = [
@ -31,9 +37,9 @@ def sgm_uniform(n, sigma_min, sigma_max, inner_model, device):
for ts in torch.linspace(start, end, n + 1)[:-1] for ts in torch.linspace(start, end, n + 1)[:-1]
] ]
sigs += [0.0] sigs += [0.0]
return torch.FloatTensor(sigs).to(device) return torch.FloatTensor(sigs)
def get_align_your_steps_sigmas(n, sigma_min, sigma_max, device='cpu'): def get_align_your_steps_sigmas(n, sigma_min, sigma_max):
# https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html # https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html
def loglinear_interp(t_steps, num_steps): def loglinear_interp(t_steps, num_steps):
""" """
@ -59,12 +65,12 @@ def get_align_your_steps_sigmas(n, sigma_min, sigma_max, device='cpu'):
else: else:
sigmas.append(0.0) sigmas.append(0.0)
return torch.FloatTensor(sigmas).to(device) return torch.FloatTensor(sigmas)
def kl_optimal(n, sigma_min, sigma_max, device): def kl_optimal(n, sigma_min, sigma_max):
alpha_min = torch.arctan(torch.tensor(sigma_min, device=device)) alpha_min = torch.arctan(torch.tensor(sigma_min))
alpha_max = torch.arctan(torch.tensor(sigma_max, device=device)) alpha_max = torch.arctan(torch.tensor(sigma_max))
step_indices = torch.arange(n + 1, device=device) step_indices = torch.arange(n + 1)
sigmas = torch.tan(step_indices / n * alpha_min + (1.0 - step_indices / n) * alpha_max) sigmas = torch.tan(step_indices / n * alpha_min + (1.0 - step_indices / n) * alpha_max)
return sigmas return sigmas