File size: 3,227 Bytes
c336648 |
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 |
# 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 k_diffusion.sampling
import torch
@torch.no_grad()
def sample_dpmpp_2m_alt(model, x, sigmas, extra_args=None, callback=None, disable=None):
"""DPM-Solver++(2M) alternative sampler
Source: https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457
"""
extra_args = {} if extra_args is None else extra_args
s_in = x.new_ones([x.shape[0]])
sigma_fn = lambda t: t.neg().exp()
t_fn = lambda sigma: sigma.log().neg()
old_denoised = None
for i in k_diffusion.sampling.trange(len(sigmas) - 1, disable=disable):
denoised = model(x, sigmas[i] * s_in, **extra_args)
if callback is not None:
callback(
{
"x": x,
"i": i,
"sigma": sigmas[i],
"sigma_hat": sigmas[i],
"denoised": denoised,
}
)
t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
h = t_next - t
if old_denoised is None or sigmas[i + 1] == 0:
x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
else:
h_last = t - t_fn(sigmas[i - 1])
r = h_last / h
denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
sigma_progress = i / len(sigmas)
adjustment_factor = 1 + (0.15 * (sigma_progress * sigma_progress))
old_denoised = denoised * adjustment_factor
return x
def add_sample_dpmpp_2m_alt_webui() -> None:
"""Adds DPM-Solver++(2M) alternative sampler to the list of available samplers."""
try:
from modules import ( # type: ignore
sd_samplers,
sd_samplers_common,
sd_samplers_kdiffusion,
)
except ImportError:
return
samplers_dpmpp_2m_alt = [
(
"DPM++ 2M alt",
sample_dpmpp_2m_alt,
["k_dpmpp_2m_alt"],
{"scheduler": "karras"},
)
]
samplers_data_dpmpp_2m_alt = [
sd_samplers_common.SamplerData(
label,
lambda model, funcname=funcname: sd_samplers_kdiffusion.KDiffusionSampler(
funcname, model
),
aliases,
options,
)
for label, funcname, aliases, options in samplers_dpmpp_2m_alt
]
sd_samplers.all_samplers.extend(samplers_data_dpmpp_2m_alt)
for x in samplers_data_dpmpp_2m_alt:
sd_samplers.all_samplers_map[x.name] = x
sd_samplers.set_samplers()
add_sample_dpmpp_2m_alt_webui()
|