Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Performance] Keep sigmas on CPU #15823

Merged
merged 5 commits into from
Jun 9, 2024
Merged

[Performance] Keep sigmas on CPU #15823

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
01491d3
Keep sigmas on CPU
drhead May 17, 2024
58b24ea
Merge branch 'AUTOMATIC1111:master' into patch-3
drhead Jun 8, 2024
428975e
Merge pull request #1 from AUTOMATIC1111/dev
drhead Jun 8, 2024
39a6d56
patch k_diffusion to_d and strip device from schedulers
drhead Jun 8, 2024
d52a1e1
lint
drhead Jun 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion modules/sd_samplers_kdiffusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ def get_sigmas(self, p, steps):
if scheduler.need_inner_model:
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:
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
Expand Down
24 changes: 15 additions & 9 deletions modules/sd_schedulers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@

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
class Scheduler:
name: str
Expand All @@ -19,21 +25,21 @@ class Scheduler:
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)


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))
end = inner_model.sigma_to_t(torch.tensor(sigma_min))
sigs = [
inner_model.t_to_sigma(ts)
for ts in torch.linspace(start, end, n + 1)[:-1]
]
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
def loglinear_interp(t_steps, num_steps):
"""
Expand All @@ -59,12 +65,12 @@ def loglinear_interp(t_steps, num_steps):
else:
sigmas.append(0.0)

return torch.FloatTensor(sigmas).to(device)
return torch.FloatTensor(sigmas)

def kl_optimal(n, sigma_min, sigma_max, device):
alpha_min = torch.arctan(torch.tensor(sigma_min, device=device))
alpha_max = torch.arctan(torch.tensor(sigma_max, device=device))
step_indices = torch.arange(n + 1, device=device)
def kl_optimal(n, sigma_min, sigma_max):
alpha_min = torch.arctan(torch.tensor(sigma_min))
alpha_max = torch.arctan(torch.tensor(sigma_max))
step_indices = torch.arange(n + 1)
sigmas = torch.tan(step_indices / n * alpha_min + (1.0 - step_indices / n) * alpha_max)
return sigmas

Expand Down