111 lines
3.1 KiB
Python
111 lines
3.1 KiB
Python
# 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
|
|
|
|
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()
|