693 lines
24 KiB
Python
693 lines
24 KiB
Python
from __future__ import annotations
|
|
|
|
import ast
|
|
import os
|
|
import warnings
|
|
from collections import namedtuple
|
|
|
|
import comfy
|
|
import latent_preview
|
|
import torch
|
|
from comfy.sample import prepare_noise, sample_custom
|
|
from comfy.samplers import KSAMPLER, KSampler, sampler_object
|
|
from comfy.utils import ProgressBar
|
|
from tqdm.auto import trange
|
|
|
|
from .restart_schedulers import NORMAL_SCHEDULER_MAPPING, RESTART_SCHEDULER_MAPPING
|
|
|
|
VERBOSE = os.environ.get("COMFYUI_VERBOSE_RESTART_SAMPLING", "").strip() == "1"
|
|
|
|
DEFAULT_SEGMENTS = "[3,2,0.06,0.30],[3,1,0.30,0.59]"
|
|
|
|
|
|
def add_restart_segment(restart_segments, n_restart, k, t_min, t_max):
|
|
if restart_segments is None:
|
|
restart_segments = []
|
|
restart_segments.append({"n": n_restart, "k": k, "t_min": t_min, "t_max": t_max})
|
|
return restart_segments
|
|
|
|
|
|
def resolve_t_value(val, ms):
|
|
if isinstance(val, (float, int)):
|
|
if val >= 0.0:
|
|
return val
|
|
if val >= -1000:
|
|
return ms.sigma(torch.FloatTensor([abs(int(val))], device="cpu")).item()
|
|
if isinstance(val, str) and val.endswith("%"):
|
|
try:
|
|
val = float(val[:-1])
|
|
if val >= 0 and val <= 100:
|
|
return ms.percent_to_sigma(1.0 - val / 100.0)
|
|
except ValueError:
|
|
pass
|
|
raise ValueError("bad t_min or t_max value")
|
|
|
|
|
|
def prepare_restart_segments(restart_info, ms, sigmas):
|
|
def get_a1111_segment():
|
|
# Emulate A1111 WebUI's restart sampler behavior.
|
|
steps = len(sigmas) - 1
|
|
if steps < 20:
|
|
# Less than 20 steps - no restarts.
|
|
return []
|
|
a1111_t_max = sigmas[int(torch.argmin(abs(sigmas - 2.0), dim=0))].item()
|
|
if steps < 36:
|
|
# Less than 36 steps - one restart with 9 steps.
|
|
return [10, 1, 0.1, a1111_t_max]
|
|
# Otherwise two restarts with steps // 4 steps.
|
|
return [(steps // 4) + 1, 2, 0.1, a1111_t_max]
|
|
|
|
restart_info = restart_info.strip().lower()
|
|
if restart_info == "":
|
|
# No restarts.
|
|
return []
|
|
restart_arrays = None
|
|
if restart_info == "default":
|
|
restart_info = DEFAULT_SEGMENTS
|
|
elif restart_info == "a1111":
|
|
restart_arrays = [get_a1111_segment()]
|
|
if restart_arrays == [[]]:
|
|
return []
|
|
if restart_arrays is None:
|
|
try:
|
|
restart_arrays = ast.literal_eval(f"[{restart_info}]")
|
|
except SyntaxError:
|
|
print("Ill-formed restart segments")
|
|
raise
|
|
temp = []
|
|
default_segments = ast.literal_eval(DEFAULT_SEGMENTS)
|
|
# This phase expands any preset strings into actual 4-item restart segments.
|
|
for idx in range(len(restart_arrays)):
|
|
item = restart_arrays[idx]
|
|
if not isinstance(item, str):
|
|
temp.append(item)
|
|
continue
|
|
preset = item.strip().lower()
|
|
if preset == "default":
|
|
temp += default_segments
|
|
elif preset == "a1111":
|
|
temp.append(get_a1111_segment())
|
|
else:
|
|
raise ValueError("Ill-formed restart segment")
|
|
restart_arrays = temp
|
|
restart_segments = []
|
|
# Now we build the actual restart segments.
|
|
for arr in restart_arrays:
|
|
if not isinstance(arr, (list, tuple)) or len(arr) != 4:
|
|
raise ValueError("Restart segment must be a list with 4 values")
|
|
n_restart, k, val_min, val_max = arr
|
|
n_restart, k = int(n_restart), int(k)
|
|
t_min = resolve_t_value(val_min, ms)
|
|
t_max = resolve_t_value(val_max, ms)
|
|
restart_segments = add_restart_segment(
|
|
restart_segments,
|
|
n_restart,
|
|
k,
|
|
t_min,
|
|
t_max,
|
|
)
|
|
return restart_segments
|
|
|
|
|
|
def round_restart_segments(ts, restart_segments):
|
|
"""
|
|
Map nearest timestep/sigma min to the nearest timestep/sigma to segments.
|
|
:param ts: Timesteps or sigmas of the original denoising schedule
|
|
:param restart_segments: Restart segments dict of the form {'t_min': t_min, 'n': n, 'k': k, 't_max': t_max}
|
|
:return: dict of the form {nearest_t_min: {'n': n, 'k': k, 't_max': t_max}}
|
|
"""
|
|
t_min_mapping = {}
|
|
for segment in reversed(
|
|
restart_segments,
|
|
): # Reversed to prioritize segments to the front
|
|
t_min_neighbor = min(ts, key=lambda ts: abs(ts - segment["t_min"])).item()
|
|
if t_min_neighbor == ts[0]:
|
|
warnings.warn(
|
|
f"\n[Restart Sampling] nearest neighbor of segment t_min {segment['t_min']:.4f} is equal to the first t_min in the denoise schedule {ts[0]:.4f}, ignoring segment...",
|
|
stacklevel=2,
|
|
)
|
|
continue
|
|
if t_min_neighbor > segment["t_max"]:
|
|
warnings.warn(
|
|
f"\n[Restart Sampling] t_min neighbor {t_min_neighbor:.4f} is greater than t_max {segment['t_max']:.4f}, ignoring segment...",
|
|
stacklevel=2,
|
|
)
|
|
continue
|
|
if t_min_neighbor in t_min_mapping:
|
|
warnings.warn(
|
|
f"\n[Restart Sampling] Overwriting segment {t_min_mapping[t_min_neighbor]}, nearest neighbor of {segment['t_min']:.4f} is {t_min_neighbor:.4f}",
|
|
stacklevel=2,
|
|
)
|
|
t_min_mapping[t_min_neighbor] = {
|
|
"n": segment["n"],
|
|
"k": segment["k"],
|
|
"t_max": segment["t_max"],
|
|
}
|
|
return t_min_mapping
|
|
|
|
|
|
def calc_sigmas(
|
|
scheduler,
|
|
n,
|
|
sigma_min,
|
|
sigma_max,
|
|
model,
|
|
device,
|
|
restart_segment=True,
|
|
):
|
|
mapping = RESTART_SCHEDULER_MAPPING if restart_segment else NORMAL_SCHEDULER_MAPPING
|
|
return mapping[scheduler](model, n, sigma_min, sigma_max, device)
|
|
|
|
|
|
def restart_sampling(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
restart_info,
|
|
restart_scheduler,
|
|
denoise=1.0,
|
|
disable_noise=False,
|
|
step_range=None,
|
|
force_full_denoise=False,
|
|
output_only=True,
|
|
custom_noise=None,
|
|
chunked_mode=True,
|
|
sigmas=None,
|
|
):
|
|
if isinstance(sampler, str):
|
|
# Only possible to determine this when the sampler is passed by name. When using
|
|
# a custom sampler, the user will need to slice sigmas when desirable.
|
|
discard_penultimate_sigma = sampler in getattr(
|
|
KSampler,
|
|
"DISCARD_PENULTIMATE_SIGMA_SAMPLERS",
|
|
set(),
|
|
)
|
|
sampler = sampler_object(sampler)
|
|
else:
|
|
discard_penultimate_sigma = False
|
|
|
|
plan = RestartPlan(
|
|
model,
|
|
steps,
|
|
scheduler,
|
|
restart_info,
|
|
restart_scheduler,
|
|
denoise=denoise,
|
|
step_range=step_range,
|
|
force_full_denoise=force_full_denoise,
|
|
sigmas=sigmas,
|
|
discard_penultimate_sigma=discard_penultimate_sigma,
|
|
)
|
|
|
|
if VERBOSE:
|
|
plan.explain(chunked_mode)
|
|
|
|
total_steps = plan.total_steps
|
|
sigmas = plan.sigmas().to(model.load_device)
|
|
|
|
latent = latent_image
|
|
latent_image = latent["samples"]
|
|
|
|
if disable_noise:
|
|
torch.manual_seed(
|
|
seed,
|
|
) # workaround for https://github.com/comfyanonymous/ComfyUI/issues/2833
|
|
noise = torch.zeros(
|
|
latent_image.size(),
|
|
dtype=latent_image.dtype,
|
|
layout=latent_image.layout,
|
|
device="cpu",
|
|
)
|
|
else:
|
|
batch_inds = latent.get("batch_index", None)
|
|
noise = prepare_noise(latent_image, seed, batch_inds)
|
|
|
|
noise_mask = None
|
|
if "noise_mask" in latent:
|
|
noise_mask = latent["noise_mask"]
|
|
|
|
x0_output = {}
|
|
callback = latent_preview.prepare_callback(model, plan.total_steps, x0_output)
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
|
|
restart_options = {
|
|
"restart_chunked": chunked_mode,
|
|
"restart_wrapped_sampler": sampler,
|
|
"restart_custom_noise": custom_noise,
|
|
}
|
|
|
|
ksampler = KSAMPLER(
|
|
RestartSampler.sampler_function,
|
|
extra_options=sampler.extra_options | restart_options,
|
|
inpaint_options=sampler.inpaint_options | {},
|
|
)
|
|
|
|
# Add the additional steps to the progress bar
|
|
pbar_update_absolute = ProgressBar.update_absolute
|
|
|
|
def pbar_update_absolute_wrapper(self, value, total=None, preview=None): # noqa: ARG001
|
|
pbar_update_absolute(self, value, total_steps, preview)
|
|
|
|
ProgressBar.update_absolute = pbar_update_absolute_wrapper
|
|
|
|
try:
|
|
samples = sample_custom(
|
|
model,
|
|
noise,
|
|
cfg,
|
|
ksampler,
|
|
sigmas,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
noise_mask=noise_mask,
|
|
callback=callback,
|
|
disable_pbar=disable_pbar,
|
|
seed=seed,
|
|
)
|
|
finally:
|
|
ProgressBar.update_absolute = pbar_update_absolute
|
|
|
|
out = latent.copy()
|
|
out["samples"] = samples
|
|
|
|
if output_only:
|
|
return (out,)
|
|
|
|
if "x0" in x0_output:
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
|
|
else:
|
|
out_denoised = out
|
|
|
|
return (out, out_denoised)
|
|
|
|
|
|
# PlanItem:
|
|
# sigmas: Sigmas for normal (outside of a restart segment) sampling. They start from after the previous PlanItem's steps
|
|
# if there is one or simply the beginning of sampling.
|
|
# k: This is the same as the restart segment definition. Set to 0 if there is no restart segment.
|
|
# restart_sigmas: Sigmas for the restart segment if it exists, otherwise None.
|
|
# Convenience properties:
|
|
# total_steps
|
|
# s_min, s_max: None if no restart sigmas.
|
|
# n_restart: 0 if no restart sigmas.
|
|
|
|
|
|
class PlanItem(
|
|
namedtuple(
|
|
"PlanItem",
|
|
["sigmas", "k", "restart_sigmas"],
|
|
defaults=[None, 0, None],
|
|
),
|
|
):
|
|
__slots__ = ()
|
|
|
|
def __new__(cls, *args: list, **kwargs: dict):
|
|
obj = super().__new__(cls, *args, **kwargs)
|
|
obj.validate()
|
|
return obj
|
|
|
|
def validate(self, threshold=1e-06):
|
|
if len(self.sigmas) < 2:
|
|
raise ValueError("PlanItem: invalid normal sigmas: too short")
|
|
t = self.sigmas.sort(descending=True, stable=True)[0].unique_consecutive()
|
|
if not torch.equal(self.sigmas, t):
|
|
errstr = (
|
|
f"PlanItem: invalid normal sigmas: out of order or contains duplicates: {self}",
|
|
)
|
|
raise ValueError(errstr)
|
|
if self.k == 0:
|
|
return
|
|
if self.k < 0:
|
|
raise ValueError("PlanItem: invalid negative k value")
|
|
if len(self.restart_sigmas) < 2:
|
|
raise ValueError("PlanItem: invalid restart sigmas: too short")
|
|
if self.s_min >= self.s_max:
|
|
raise ValueError("PlanItem: invalid min/max: min >= max")
|
|
if self.sigmas[-1] - self.restart_sigmas[0] > threshold:
|
|
raise ValueError(
|
|
"PlanItem: invalid sigmas: last normal sigma >= first restart sigma",
|
|
)
|
|
if self.sigmas[-1] - self.restart_sigmas[-1] > threshold:
|
|
errstr = (
|
|
f"PlanItem: invalid sigmas: last restart sigma {self.restart_sigmas[-1]} < last normal sigma {self.sigmas[-1]}",
|
|
)
|
|
raise ValueError(errstr)
|
|
t = self.restart_sigmas.sort(descending=True, stable=True)[
|
|
0
|
|
].unique_consecutive()
|
|
if not torch.equal(self.restart_sigmas, t):
|
|
errstr = (
|
|
f"PlanItem: invalid restart sigmas: out of order or contains duplicates: {self}",
|
|
)
|
|
raise ValueError(errstr)
|
|
|
|
@property
|
|
def total_steps(self):
|
|
if self.k < 1:
|
|
return len(self.sigmas) - 1
|
|
return (len(self.sigmas) - 1) + (len(self.restart_sigmas) - 1) * self.k
|
|
|
|
@property
|
|
def s_min(self):
|
|
return None if self.k < 1 else self.restart_sigmas[-1].item()
|
|
|
|
@property
|
|
def s_max(self):
|
|
return None if self.k < 1 else self.restart_sigmas[0].item()
|
|
|
|
@property
|
|
def n_restart(self):
|
|
return 0 if self.k < 1 else len(self.restart_sigmas) - 1
|
|
|
|
|
|
class RestartPlan:
|
|
def __init__(
|
|
self,
|
|
model,
|
|
steps,
|
|
scheduler,
|
|
restart_info,
|
|
restart_scheduler,
|
|
denoise=1.0,
|
|
step_range=None,
|
|
force_full_denoise=False,
|
|
sigmas=None,
|
|
discard_penultimate_sigma=False,
|
|
):
|
|
if (
|
|
denoise <= 0
|
|
or (sigmas is None and steps < 1)
|
|
or (sigmas is not None and len(sigmas) < 2)
|
|
):
|
|
self.plan = []
|
|
self.total_steps = 0
|
|
return
|
|
ms = model.get_model_object("model_sampling")
|
|
|
|
if sigmas is None:
|
|
effective_steps = steps if denoise > 0.9999 else int(steps / denoise)
|
|
sigmas = calc_sigmas(
|
|
scheduler,
|
|
effective_steps + int(discard_penultimate_sigma), # True evaluates to 1
|
|
float(ms.sigma_min),
|
|
float(ms.sigma_max),
|
|
model.model,
|
|
"cpu",
|
|
restart_segment=False,
|
|
)
|
|
if discard_penultimate_sigma:
|
|
sigmas = torch.cat((sigmas[:-2], sigmas[-1:]))
|
|
else:
|
|
steps = effective_steps = len(sigmas) - 1
|
|
steps = steps if denoise > 0.9999 else int(effective_steps * denoise)
|
|
sigmas = sigmas.clone().detach().cpu()
|
|
if effective_steps != steps:
|
|
sigmas = sigmas[-(steps + 1) :]
|
|
if step_range is not None:
|
|
start_step, last_step = step_range
|
|
|
|
if last_step < len(sigmas) - 1:
|
|
sigmas = sigmas[: last_step + 1]
|
|
if force_full_denoise:
|
|
sigmas[-1] = 0
|
|
|
|
if start_step < len(sigmas) - 1:
|
|
sigmas = sigmas[start_step:]
|
|
|
|
restart_segments = prepare_restart_segments(restart_info, ms, sigmas)
|
|
self.plan, self.total_steps = self.build_plan_items(
|
|
model.model,
|
|
restart_segments,
|
|
restart_scheduler,
|
|
sigmas,
|
|
"cpu",
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<RestartPlan: steps={self.total_steps}, plan={self.plan}>"
|
|
|
|
# Builds a list of PlanItems and calculates the total number of steps. See the comments for PlanItem
|
|
# for more information about plans.
|
|
# Returns two values: the plan and the total steps.
|
|
@staticmethod
|
|
@torch.no_grad()
|
|
def build_plan_items(
|
|
model,
|
|
restart_segments,
|
|
restart_scheduler,
|
|
sigmas,
|
|
device,
|
|
) -> tuple[list, int]:
|
|
model_sigma_min = float(model.model_sampling.sigma_min)
|
|
segments = round_restart_segments(sigmas, restart_segments)
|
|
plan = []
|
|
range_start = -1
|
|
for i in range(len(sigmas) - 1):
|
|
if range_start == -1:
|
|
# Starting a new plan item - main sigmas start at the current index of i.
|
|
range_start = i
|
|
s_min = sigmas[i + 1].item()
|
|
seg = segments.get(s_min)
|
|
if seg is None:
|
|
continue
|
|
s_max, k, n_restart = seg["t_max"], seg["k"], seg["n"]
|
|
if k < 1 or n_restart < 2:
|
|
continue
|
|
if s_max <= model_sigma_min:
|
|
errstr = f"Restart: Invalid restart segment t_max {s_max:.05} <= model minimum sigma {model_sigma_min:.05}"
|
|
raise ValueError(errstr)
|
|
normal_sigmas = sigmas[range_start : i + 2]
|
|
restart_sigmas = calc_sigmas(
|
|
restart_scheduler,
|
|
n_restart,
|
|
max(model_sigma_min, sigmas[i + 1]),
|
|
s_max,
|
|
model,
|
|
device=device,
|
|
)
|
|
if normal_sigmas[-1] != 0:
|
|
restart_sigmas = restart_sigmas[:-1]
|
|
restart_sigmas[-1] = s_min # Force the restart segment to end at s_min.
|
|
plan.append(PlanItem(normal_sigmas, k, restart_sigmas))
|
|
range_start = -1
|
|
if range_start != -1:
|
|
# Include sigmas after the last restart segments in the plan.
|
|
plan.append(PlanItem(sigmas[range_start:]))
|
|
return plan, sum(pi.total_steps for pi in plan)
|
|
|
|
def sigmas(self) -> torch.Tensor:
|
|
# Flattens a plan into sigmas. When the first normal sigma matches the last item's
|
|
# final sigma, we strip the first normal sigma to avoid creating duplicates.
|
|
if not self.plan or self.total_steps < 1:
|
|
return torch.FloatTensor([])
|
|
|
|
def sigmas_generator():
|
|
prev_last = None
|
|
for pi in self.plan:
|
|
yield pi.sigmas if prev_last != pi.sigmas[0] else pi.sigmas[1:]
|
|
prev_last = pi.restart_sigmas[-1] if pi.k > 0 else pi.sigmas[-1]
|
|
for _ in range(pi.k):
|
|
yield pi.restart_sigmas
|
|
|
|
return torch.flatten(torch.cat(tuple(sigmas_generator())))
|
|
|
|
# Dumps information about the plan to the console. It uses the normal plan execute
|
|
# logic.
|
|
def explain(self, chunked=True):
|
|
def pretty_sigmas(sigmas):
|
|
return ", ".join(f"{sig:.4}" for sig in sigmas.tolist())
|
|
|
|
def dump_steps(step, sigmas, restart=0):
|
|
rlabel = f"R{restart:>3}" if restart > 0 else " "
|
|
if chunked:
|
|
chunk_size = len(sigmas) - 2
|
|
step += 1
|
|
print(
|
|
f"[{rlabel}] Step {step:>3}..{step+chunk_size:<3}: {pretty_sigmas(sigmas)}",
|
|
)
|
|
step += chunk_size
|
|
return step
|
|
for i in range(len(sigmas) - 1):
|
|
step += 1
|
|
print(f"[{rlabel}] Step {step:>3}: {pretty_sigmas(sigmas[i:i+2])}")
|
|
return step
|
|
|
|
print(f"** Dumping restart sampling plan (total steps {self.total_steps}):")
|
|
step = 0
|
|
for pi in self.plan:
|
|
step = dump_steps(step, pi.sigmas)
|
|
for kidx in range(pi.k):
|
|
step = dump_steps(step, pi.restart_sigmas, kidx + 1)
|
|
print(
|
|
"** Plan legend: [Rn] - steps for restart #n, normal sampling steps otherwise. Ranges are inclusive.",
|
|
)
|
|
|
|
@staticmethod
|
|
def self_test(
|
|
model,
|
|
schedules=None,
|
|
restart_schedules=None,
|
|
segments=None,
|
|
min_steps=2,
|
|
max_steps=100,
|
|
) -> None:
|
|
if schedules is None:
|
|
schedules = NORMAL_SCHEDULER_MAPPING.keys()
|
|
if restart_schedules is None:
|
|
restart_schedules = RESTART_SCHEDULER_MAPPING.keys()
|
|
if segments is None:
|
|
segments = ("default", "a1111")
|
|
for schname in schedules:
|
|
for rschname in restart_schedules:
|
|
for tsegs in segments:
|
|
print(
|
|
f"--- Test: {min_steps}..{max_steps} steps, schedules {schname}/{rschname}, segments {tsegs}",
|
|
)
|
|
for tsteps in range(min_steps, max_steps + 1):
|
|
label = f"** {tsteps:03}: {schname}, {rschname}, {tsegs}:"
|
|
try:
|
|
_plan = RestartPlan(
|
|
model,
|
|
tsteps,
|
|
schname,
|
|
tsegs,
|
|
rschname,
|
|
1.0,
|
|
)
|
|
except ValueError as err:
|
|
print(f"{label}\n\t!! FAIL: {err}")
|
|
raise
|
|
continue
|
|
print("\n|| Done test")
|
|
|
|
|
|
class RestartSampler:
|
|
@staticmethod
|
|
def get_segment(sigmas: torch.Tensor) -> torch.Tensor:
|
|
# A normal segment ends when we either reach the end of the list or
|
|
# encounter a sigma higher than the previous.
|
|
last_sigma = sigmas[0]
|
|
for idx in range(1, len(sigmas)):
|
|
sigma = sigmas[idx]
|
|
if sigma > last_sigma:
|
|
return sigmas[:idx]
|
|
last_sigma = sigma
|
|
return sigmas
|
|
|
|
@classmethod
|
|
def split_sigmas(cls, sigmas):
|
|
# This function just splits the sigmas into chunks that are sorted descending.
|
|
# If the first sigma of a chunk is > the last sigma of the previous chunk then this
|
|
# is a restart segment: noising the restart uses s_min=prev_chunk[-1], s_max=chunk[0].
|
|
# It's a generator that yields tuples of (noise_scale, chunk_sigmas).
|
|
prev_seg = None
|
|
while len(sigmas) > 1:
|
|
seg = cls.get_segment(sigmas)
|
|
sigmas = sigmas[len(seg) :]
|
|
if prev_seg is not None and seg[0] > prev_seg[-1]:
|
|
s_min, s_max = prev_seg[-1], seg[0]
|
|
noise_scale = ((s_max**2 - s_min**2) ** 0.5).item()
|
|
else:
|
|
noise_scale = 0.0
|
|
prev_seg = seg
|
|
yield (noise_scale, seg)
|
|
|
|
# Some extra explanation for a couple of these arguments:
|
|
#
|
|
# restart_chunked:
|
|
# When False, the sampling function is called step-by-step with only two sigmas at a time.
|
|
# When True, the sampling function will be called with sigmas for multiple steps at a time.
|
|
# this means either the steps up to the next restart segment (or the end of sampling) or the steps within
|
|
# a restart segment.
|
|
#
|
|
# restart_custom_noise:
|
|
# If set to None, restart noise will just use torch.randn_like (gaussian) for noise generation. Otherwise
|
|
# this should contain a function that takes x, sigma_min, sigma_max, seed and returns a noise sampler
|
|
# function (which takes sigma, sigma_next) and returns a noisy tensor.
|
|
@classmethod
|
|
@torch.no_grad()
|
|
def sampler_function(
|
|
cls,
|
|
model,
|
|
x,
|
|
sigmas,
|
|
*args: list,
|
|
restart_wrapped_sampler=None,
|
|
restart_chunked=True,
|
|
restart_custom_noise=None,
|
|
callback=None,
|
|
disable=None,
|
|
**kwargs: dict,
|
|
) -> torch.Tensor:
|
|
if not restart_wrapped_sampler:
|
|
raise ValueError("RestartSampler: missing restart_sampler option!")
|
|
|
|
def restart_noise(x, _s_min, _s_max, _seed):
|
|
return lambda _s, _sn: torch.randn_like(x)
|
|
|
|
seed = (kwargs.get("extra_args", {}) or {}).get("seed", 42)
|
|
if restart_custom_noise is not None:
|
|
restart_noise = restart_custom_noise
|
|
|
|
sampler = restart_wrapped_sampler.sampler_function
|
|
|
|
chunks = tuple(cls.split_sigmas(sigmas))
|
|
total_steps = sum(len(chunk) - 1 for _noise, chunk in chunks)
|
|
step = 0
|
|
noise_count = 0
|
|
with trange(total_steps, disable=disable) as pbar:
|
|
last_cb_sigma = None
|
|
|
|
def cb_wrapper(cb_state):
|
|
nonlocal step, last_cb_sigma
|
|
curr_sigma = cb_state.get("sigma")
|
|
curr_sigma = (
|
|
curr_sigma.item()
|
|
if isinstance(curr_sigma, torch.Tensor)
|
|
else curr_sigma
|
|
)
|
|
if last_cb_sigma is not None and curr_sigma == last_cb_sigma:
|
|
# No change since last time we were called, so we won't track it as a step.
|
|
return
|
|
step += 1
|
|
pbar.update(1)
|
|
cb_state["i"] = step
|
|
last_cb_sigma = curr_sigma
|
|
if callback is not None:
|
|
callback(cb_state)
|
|
|
|
def do_sample(x, sigmas):
|
|
return sampler(
|
|
model,
|
|
x,
|
|
sigmas,
|
|
*args,
|
|
callback=cb_wrapper,
|
|
disable=True,
|
|
**kwargs,
|
|
)
|
|
|
|
for noise_scale, chunk_sigmas in chunks:
|
|
if noise_scale != 0:
|
|
s_min, s_max = chunk_sigmas[-1], chunk_sigmas[0]
|
|
x += (
|
|
restart_noise(x, s_min, s_max, seed + noise_count)(s_max, s_min)
|
|
* noise_scale
|
|
)
|
|
noise_count += 1
|
|
if restart_chunked:
|
|
x = do_sample(x, chunk_sigmas)
|
|
continue
|
|
for i in range(len(chunk_sigmas) - 1):
|
|
x = do_sample(x, chunk_sigmas[i : i + 2])
|
|
return x
|