-
Notifications
You must be signed in to change notification settings - Fork 4
/
schedulers.py
59 lines (49 loc) · 2.83 KB
/
schedulers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import torch
import comfy.samplers
from comfy.samplers import SCHEDULER_NAMES, simple_scheduler
from comfy_extras.nodes_align_your_steps import loglinear_interp
from comfy_extras.nodes_gits import GITSScheduler
calculate_sigmas_original = comfy.samplers.calculate_sigmas
AYS_SCHEDULER = "ays"
AYS_PLUS_SCHEDULER = "ays+"
AYS_30_SCHEDULER = "ays_30"
AYS_30_PLUS_SCHEDULER = "ays_30+"
GITS_SCHEDULER = "gits"
# Modified AYS by Extraltodeus from https://github.com/Extraltodeus/sigmas_tools_and_the_golden_scheduler/blob/0dc89a264ef346a093d053c0da751f3ece317613/sigmas_merge.py#L203-L233
# 30-step AYS by Koitenshin from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15751#issuecomment-2143648234
def _ays_scheduler(model_sampling, steps, force_sigma_min, model_type="SDXL"):
timestep_indices = {
"SD1": [999, 850, 736, 645, 545, 455, 343, 233, 124, 24, 0],
"SDXL": [999, 845, 730, 587, 443, 310, 193, 116, 53, 13, 0],
"SD1_30": [999, 954, 907, 855, 819, 782, 744, 713, 684, 654, 623, 592, 559, 528, 499, 470, 437, 401, 364, 328, 291, 256, 221, 183, 149, 114, 70, 41, 21, 4, 0],
"SDXL_30": [999, 953, 904, 850, 813, 777, 738, 695, 650, 602, 556, 510, 462, 417, 374, 331, 290, 250, 214, 182, 155, 131, 108, 85, 66, 49, 32, 20, 12, 3, 0],
}
indices = timestep_indices[model_type]
sigmas = simple_scheduler(model_sampling, 1000).flip(0)[1:][indices]
sigmas = loglinear_interp(sigmas.tolist(), steps + 1 if not force_sigma_min else steps)
sigmas = torch.FloatTensor(sigmas)
sigmas = torch.cat([sigmas[:-1] if not force_sigma_min else sigmas, torch.FloatTensor([0.0])])
return sigmas.cpu()
def _calculate_sigmas_hijack(model_sampling, scheduler_name, steps):
if scheduler_name == AYS_SCHEDULER:
sigmas = _ays_scheduler(model_sampling, steps, False)
elif scheduler_name == AYS_PLUS_SCHEDULER:
sigmas = _ays_scheduler(model_sampling, steps, True)
elif scheduler_name == AYS_30_SCHEDULER:
sigmas = _ays_scheduler(model_sampling, steps, False, "SDXL_30")
elif scheduler_name == AYS_30_PLUS_SCHEDULER:
sigmas = _ays_scheduler(model_sampling, steps, True, "SDXL_30")
elif scheduler_name == GITS_SCHEDULER:
sigmas = GITSScheduler().get_sigmas(1.2, steps, 1.0)[0]
else:
sigmas = calculate_sigmas_original(model_sampling, scheduler_name, steps)
return sigmas
def hijack_schedulers():
if calculate_sigmas_original == _calculate_sigmas_hijack:
raise ValueError("Schedulers are already hijacked")
SCHEDULER_NAMES.append(AYS_SCHEDULER)
SCHEDULER_NAMES.append(AYS_PLUS_SCHEDULER)
SCHEDULER_NAMES.append(AYS_30_SCHEDULER)
SCHEDULER_NAMES.append(AYS_30_PLUS_SCHEDULER)
SCHEDULER_NAMES.append(GITS_SCHEDULER)
comfy.samplers.calculate_sigmas = _calculate_sigmas_hijack