Skip to content

Commit

Permalink
Merge pull request #16142 from AndreyRGW/dev
Browse files Browse the repository at this point in the history
Add Simple Scheduler
  • Loading branch information
AUTOMATIC1111 committed Jul 6, 2024
2 parents 9cbde79 + 9414309 commit b8c3664
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions modules/sd_schedulers.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,16 @@ def kl_optimal(n, sigma_min, sigma_max, device):
sigmas = torch.tan(step_indices / n * alpha_min + (1.0 - step_indices / n) * alpha_max)
return sigmas


def simple_scheduler(n, sigma_min, sigma_max, inner_model, device):
sigs = []
ss = len(inner_model.sigmas) / n
for x in range(n):
sigs += [float(inner_model.sigmas[-(1 + int(x * ss))])]
sigs += [0.0]
return torch.FloatTensor(sigs).to(device)


def normal_scheduler(n, sigma_min, sigma_max, inner_model, device, sgm=False, floor=False):
start = inner_model.sigma_to_t(torch.tensor(sigma_max))
end = inner_model.sigma_to_t(torch.tensor(sigma_min))
Expand All @@ -92,6 +102,7 @@ def normal_scheduler(n, sigma_min, sigma_max, inner_model, device, sgm=False, fl
sigs += [0.0]
return torch.FloatTensor(sigs).to(device)


def ddim_scheduler(n, sigma_min, sigma_max, inner_model, device):
sigs = []
ss = max(len(inner_model.sigmas) // n, 1)
Expand All @@ -113,6 +124,7 @@ def ddim_scheduler(n, sigma_min, sigma_max, inner_model, device):
Scheduler('sgm_uniform', 'SGM Uniform', sgm_uniform, need_inner_model=True, aliases=["SGMUniform"]),
Scheduler('kl_optimal', 'KL Optimal', kl_optimal),
Scheduler('align_your_steps', 'Align Your Steps', get_align_your_steps_sigmas),
Scheduler('simple', 'Simple', simple_scheduler, need_inner_model=True),
Scheduler('normal', 'Normal', normal_scheduler, need_inner_model=True),
Scheduler('ddim', 'DDIM', ddim_scheduler, need_inner_model=True),
]
Expand Down

0 comments on commit b8c3664

Please sign in to comment.