File size: 3,405 Bytes
bb7f1f4 |
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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
# Copyright 2024 SLAPaper
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import numpy as np
import torch
def loglinear_interp(t_steps: list[float], num_steps: int) -> np.ndarray:
"""
Performs log-linear interpolation of a given array of decreasing numbers.
"""
xs = np.linspace(0, 1, len(t_steps))
ys = np.log(t_steps[::-1])
new_xs = np.linspace(0, 1, num_steps)
new_ys = np.interp(new_xs, xs, ys)
interped_ys = np.exp(new_ys)[::-1].copy()
return interped_ys
def align_your_step_scheduler_v15(
n: int, sigma_min: float, sigma_max: float, device
) -> torch.Tensor:
"""SD15 AYS scheduler from https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html"""
NOISE_LEVELS = [
14.615,
6.475,
3.861,
2.697,
1.886,
1.396,
0.963,
0.652,
0.399,
0.152,
0.029,
]
TIMESTEP_INDICES = [999, 850, 736, 645, 545, 455, 343, 233, 124, 24, 0]
sigs = [sigma for sigma in loglinear_interp(NOISE_LEVELS, n)]
sigs.append(0.0)
logging.info(f"AYS scheduler: {sigs=}")
return torch.FloatTensor(sigs).to(device)
def align_your_step_scheduler_xl(
n: int, sigma_min: float, sigma_max: float, device
) -> torch.Tensor:
"""SDXL AYS scheduler from https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html"""
NOISE_LEVELS = [
14.615,
6.315,
3.771,
2.181,
1.342,
0.862,
0.555,
0.380,
0.234,
0.113,
0.029,
]
TIMESTEP_INDICES = [999, 845, 730, 587, 443, 310, 193, 116, 53, 13, 0]
sigs = [sigma for sigma in loglinear_interp(NOISE_LEVELS, n)]
sigs.append(0.0)
logging.info(f"AYS scheduler: {sigs=}")
return torch.FloatTensor(sigs).to(device)
def add_align_your_step_scheduler() -> None:
"""Add AYS scheduler to the list of schedulers"""
try:
from modules import sd_schedulers # type: ignore
except ImportError:
return
# latest dev webui have already added these scheduler: https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15751
if 'align_your_steps' in sd_schedulers.schedulers_map:
return
scheduler_v15 = sd_schedulers.Scheduler(
"ays_v15", "Align Your Step SD15", align_your_step_scheduler_v15
)
scheduler_xl = sd_schedulers.Scheduler(
"ays_xl", "Align Your Step SDXL", align_your_step_scheduler_xl
)
sd_schedulers.schedulers.append(scheduler_v15)
sd_schedulers.schedulers.append(scheduler_xl)
sd_schedulers.schedulers_map[scheduler_v15.name] = scheduler_v15
sd_schedulers.schedulers_map[scheduler_v15.label] = scheduler_v15
sd_schedulers.schedulers_map[scheduler_xl.name] = scheduler_xl
sd_schedulers.schedulers_map[scheduler_xl.label] = scheduler_xl
add_align_your_step_scheduler()
|