355 lines
12 KiB
Python
355 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from comfy.samplers import ksampler
|
|
from pytorch_wavelets import DTCWTForward, DTCWTInverse, DWTForward, DWTInverse
|
|
|
|
from .tensor_image_ops import Sharpen
|
|
from .upscale import Upscale
|
|
from .utils import fallback
|
|
from .vae import VAEHelper
|
|
|
|
|
|
class Config:
|
|
_overridable_fields = { # noqa: RUF012
|
|
"blend_by_mode",
|
|
"blend_mode",
|
|
"chunked_sampling",
|
|
"custom_noise_name",
|
|
"custom_noise_params",
|
|
"denoised_wavelet_multiplier",
|
|
"dtcwt_biort",
|
|
"dtcwt_mode",
|
|
"dtcwt_qshift",
|
|
"dwt_flip_filters",
|
|
"dwt_level",
|
|
"dwt_mode",
|
|
"dwt_wave",
|
|
"enable_cache_clearing",
|
|
"enable_gc",
|
|
"fadeout_factor",
|
|
"force_upscale_model",
|
|
"guidance_mask_blend_mode",
|
|
"guidance_mask_name",
|
|
"guidance_factor",
|
|
"guidance_mode",
|
|
"guidance_restart_s_noise",
|
|
"guidance_restart",
|
|
"guidance_sampler_name",
|
|
"guidance_steps",
|
|
"highres_sigmas_name",
|
|
"mask_blend_mode",
|
|
"mask_name",
|
|
"reference_image_name",
|
|
"reference_sampler_name",
|
|
"reference_wavelet_multiplier",
|
|
"renoise_factor",
|
|
"resample_mode",
|
|
"rescale_increment",
|
|
"restart_custom_noise_name",
|
|
"sampler_name",
|
|
"scale_factor",
|
|
"schedule_override",
|
|
"sharpen_gaussian_kernel_size",
|
|
"sharpen_gaussian_sigma",
|
|
"sharpen_mode",
|
|
"sharpen_reference",
|
|
"sharpen_strength",
|
|
"sigma_dishonesty_factor_guidance",
|
|
"sigma_dishonesty_factor",
|
|
"skip",
|
|
"skip_callback",
|
|
"upscale_model_name",
|
|
"use_upscale_model",
|
|
"vae_decode_kwargs",
|
|
"vae_encode_kwargs",
|
|
"vae_mode",
|
|
"vae_name",
|
|
}
|
|
|
|
_dict_exclude_keys = { # noqa: RUF012
|
|
"as_dict",
|
|
"dwt",
|
|
"get_iteration_config",
|
|
"guidance_sampler",
|
|
"highres_sigmas",
|
|
"idwt",
|
|
"iteration_override",
|
|
"reference_sampler",
|
|
"sampler",
|
|
"sharpen",
|
|
"upscale",
|
|
"vae",
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
device,
|
|
dtype,
|
|
latent_format,
|
|
*,
|
|
_params,
|
|
blend_by_mode="image",
|
|
blend_mode="lerp",
|
|
chunked_sampling=True,
|
|
custom_noise_name="",
|
|
custom_noise_params=None,
|
|
denoised_wavelet_multiplier=1.0,
|
|
dtcwt_biort="near_sym_a",
|
|
dtcwt_mode=False,
|
|
dtcwt_qshift="qshift_a",
|
|
dwt_flip_filters=False,
|
|
dwt_level=1,
|
|
dwt_mode="symmetric",
|
|
dwt_wave="db4",
|
|
enable_cache_clearing=True,
|
|
enable_gc=True,
|
|
fadeout_factor=0.0,
|
|
force_upscale_model=False,
|
|
guidance_factor=1.0,
|
|
guidance_mask_blend_mode="lerp",
|
|
guidance_mask_name="guidance",
|
|
guidance_mode="image",
|
|
guidance_restart_s_noise=1.0,
|
|
guidance_restart=0,
|
|
guidance_sampler_name="guidance",
|
|
guidance_steps=5,
|
|
highres_sigmas_name="highres",
|
|
iteration_override=None,
|
|
iterations=1,
|
|
mask_blend_mode="lerp",
|
|
mask_name="",
|
|
reference_image_name="reference",
|
|
reference_sampler_name="reference",
|
|
reference_wavelet_multiplier=1.0,
|
|
renoise_factor=1.0,
|
|
resample_mode="bicubic",
|
|
rescale_increment=64,
|
|
restart_custom_noise_name="restart",
|
|
sampler_name="",
|
|
scale_factor=2.0,
|
|
schedule_override=None,
|
|
seed_rng_offset=1,
|
|
seed_rng=True,
|
|
sharpen_gaussian_kernel_size=3,
|
|
sharpen_gaussian_sigma=(0.1, 2.0),
|
|
sharpen_mode="gaussian",
|
|
sharpen_reference=True,
|
|
sharpen_strength=1.0,
|
|
sigma_dishonesty_factor_guidance: float | None = None,
|
|
sigma_dishonesty_factor=0.0,
|
|
skip_callback=False,
|
|
skip=False,
|
|
upscale_model_name="",
|
|
use_upscale_model=True,
|
|
vae_decode_kwargs=None,
|
|
vae_encode_kwargs=None,
|
|
vae_mode="normal",
|
|
vae_name="",
|
|
ensure_model_mode: str | None = "normal",
|
|
):
|
|
self.skip = skip
|
|
if ensure_model_mode not in {
|
|
None,
|
|
"disable",
|
|
"normal",
|
|
"normal_unload",
|
|
"lowvram",
|
|
"novram",
|
|
}:
|
|
raise ValueError(
|
|
"Bad ensure_model_mode - must null or one of disable, normal, normal_unload, lowvram, novram",
|
|
)
|
|
self.ensure_model_mode = ensure_model_mode
|
|
|
|
self.vae_name = vae_name
|
|
self.upscale_model_name = upscale_model_name
|
|
self.highres_sigmas_name = highres_sigmas_name
|
|
self.reference_image_name = reference_image_name
|
|
self.reference_sampler_name = reference_sampler_name
|
|
self.guidance_sampler_name = guidance_sampler_name
|
|
self.sampler_name = sampler_name
|
|
self.custom_noise_name = custom_noise_name
|
|
self.custom_noise_params = (
|
|
{} if not isinstance(custom_noise_params, dict) else custom_noise_params
|
|
)
|
|
self.restart_custom_noise_name = restart_custom_noise_name
|
|
self.mask_name = mask_name
|
|
self.guidance_mask_name = guidance_mask_name
|
|
|
|
sampler = _params.get_item("sampler", name=sampler_name)
|
|
guidance_sampler = _params.get_item("sampler", name=guidance_sampler_name)
|
|
reference_sampler = _params.get_item("sampler", name=reference_sampler_name)
|
|
upscale_model = _params.get_item("upscale_model", name=upscale_model_name)
|
|
vae = _params.get_item("vae", name=vae_name)
|
|
highres_sigmas = _params.get_item("sigmas", name=highres_sigmas_name)
|
|
self.highres_sigmas = (
|
|
None if highres_sigmas is None else highres_sigmas.detach().clone()
|
|
)
|
|
|
|
self.mask_blend_mode = mask_blend_mode
|
|
self.guidance_mask_blend_mode = guidance_mask_blend_mode
|
|
|
|
sampler = fallback(
|
|
sampler,
|
|
lambda: ksampler("euler"),
|
|
default_is_fun=True,
|
|
)
|
|
self.seed_rng = seed_rng
|
|
self.seed_rng_offset = seed_rng_offset
|
|
self.sigma_dishonesty_factor = sigma_dishonesty_factor
|
|
self.sigma_dishonesty_factor_guidance = fallback(
|
|
sigma_dishonesty_factor_guidance,
|
|
sigma_dishonesty_factor,
|
|
)
|
|
self.skip_callback = skip_callback
|
|
self.fadeout_factor = fadeout_factor
|
|
self.scale_factor = scale_factor
|
|
self.guidance_factor = guidance_factor
|
|
self.renoise_factor = renoise_factor
|
|
self.iterations = iterations
|
|
self.guidance_steps = guidance_steps
|
|
self.guidance_mode = guidance_mode
|
|
self.guidance_restart = guidance_restart
|
|
self.guidance_restart_s_noise = guidance_restart_s_noise
|
|
self.sampler = sampler
|
|
self.guidance_sampler = fallback(guidance_sampler, sampler)
|
|
self.reference_sampler = fallback(reference_sampler, sampler)
|
|
self.vae = VAEHelper(
|
|
vae_mode,
|
|
latent_format,
|
|
device=device,
|
|
dtype=dtype,
|
|
vae=vae,
|
|
encode_kwargs=fallback(vae_encode_kwargs, {}),
|
|
decode_kwargs=fallback(vae_decode_kwargs, {}),
|
|
)
|
|
self.sharpen = Sharpen(
|
|
mode=sharpen_mode,
|
|
strength=sharpen_strength if sharpen_reference else 0,
|
|
gaussian_kernel_size=sharpen_gaussian_kernel_size,
|
|
gaussian_sigma=sharpen_gaussian_sigma,
|
|
)
|
|
self.upscale = Upscale(
|
|
resample_mode=resample_mode,
|
|
rescale_increment=rescale_increment,
|
|
upscale_model=upscale_model,
|
|
force_upscale_model=force_upscale_model,
|
|
)
|
|
if schedule_override is not None and not isinstance(schedule_override, dict):
|
|
raise TypeError("Bad type for schedule_override: must be null or object")
|
|
self.schedule_override = schedule_override
|
|
self.use_upscale_model = use_upscale_model
|
|
self.dwt_mode = dwt_mode
|
|
self.dwt_level = dwt_level
|
|
self.dwt_wave = dwt_wave
|
|
self.dtcwt_mode = dtcwt_mode
|
|
self.dtcwt_biort = dtcwt_biort
|
|
self.dtcwt_qshift = dtcwt_qshift
|
|
if dtcwt_mode:
|
|
self.dwt = DTCWTForward(
|
|
J=dwt_level,
|
|
mode=dwt_mode,
|
|
biort=dtcwt_biort,
|
|
qshift=dtcwt_qshift,
|
|
).to(device)
|
|
self.idwt = DTCWTInverse(
|
|
mode=dwt_mode,
|
|
biort=dtcwt_biort,
|
|
qshift=dtcwt_qshift,
|
|
).to(device)
|
|
else:
|
|
self.dwt = DWTForward(J=dwt_level, wave=dwt_wave, mode=dwt_mode).to(device)
|
|
self.idwt = DWTInverse(wave=dwt_wave, mode=dwt_mode).to(device)
|
|
self.dwt_flip_filters = dwt_flip_filters
|
|
self.reference_wavelet_multiplier = reference_wavelet_multiplier
|
|
self.denoised_wavelet_multiplier = denoised_wavelet_multiplier
|
|
self.blend_mode = blend_mode
|
|
if blend_by_mode not in {"image", "latent", "wavelet"}:
|
|
raise ValueError("Bad blend_by_mode: must be one of image, latent, wavelet")
|
|
self.blend_by_mode = blend_by_mode
|
|
self.enable_gc = enable_gc
|
|
self.enable_cache_clearing = enable_cache_clearing
|
|
self.chunked_sampling = chunked_sampling
|
|
self.iteration_override = {}
|
|
if iteration_override is None or iteration_override == {}:
|
|
return
|
|
if not isinstance(iteration_override, dict):
|
|
raise TypeError("Iteration override must be an object")
|
|
selfdict = self.as_dict()
|
|
overrides = self.iteration_override
|
|
for k, v in iteration_override.items():
|
|
if not isinstance(k, (int, str)) or not isinstance(v, dict):
|
|
raise TypeError(
|
|
"Bad type for override item: key must be integer or string, value must be an object",
|
|
)
|
|
okwargs = selfdict | {
|
|
ok: ov for ok, ov in v.items() if ok in self._overridable_fields
|
|
}
|
|
overrides[k] = self.__class__(
|
|
device,
|
|
dtype,
|
|
latent_format,
|
|
_params=_params,
|
|
**okwargs,
|
|
)
|
|
|
|
def as_dict(self) -> dict:
|
|
result = {
|
|
k: getattr(self, k)
|
|
for k in dir(self)
|
|
if not k.startswith("_") and k not in self._dict_exclude_keys
|
|
}
|
|
result["vae_mode"] = self.vae.mode.name.lower()
|
|
result["vae_encode_kwargs"] = self.vae.encode_kwargs
|
|
result["vae_decode_kwargs"] = self.vae.decode_kwargs
|
|
result["sharpen_reference"] = self.sharpen.strength != 0
|
|
result["sharpen_strength"] = self.sharpen.strength
|
|
result["sharpen_gaussian_kernel_size"] = self.sharpen.gaussian_kernel_size
|
|
result["sharpen_gaussian_sigma"] = self.sharpen.gaussian_sigma
|
|
result["resample_mode"] = self.upscale.resample_mode
|
|
result["rescale_increment"] = self.upscale.rescale_increment
|
|
result["force_upscale_model"] = self.upscale.force_upscale_model
|
|
return result
|
|
|
|
def get_iteration_config(self, iteration):
|
|
override = self.iteration_override.get(iteration)
|
|
return override.get_iteration_config(iteration) if override else self
|
|
|
|
|
|
class ParamGroup:
|
|
def __init__(self, items=None):
|
|
self.items = {} if items is None else items
|
|
|
|
def clone(self):
|
|
return self.__class__(items=self.items.copy())
|
|
|
|
def append(self, item):
|
|
self.items.append(item)
|
|
return item
|
|
|
|
def __getitem__(self, key):
|
|
return self.items[key]
|
|
|
|
def __setitem__(self, key, value):
|
|
self.items[key] = value
|
|
|
|
def __len__(self):
|
|
return len(self.items)
|
|
|
|
def __iter__(self):
|
|
return self.items.__iter__()
|
|
|
|
def get_item(
|
|
self,
|
|
type_name: str,
|
|
*,
|
|
name: str | None = "",
|
|
param_mode: bool = False,
|
|
default: Any | None = None, # noqa: ANN401
|
|
) -> Any: # noqa: ANN401
|
|
name = name if name is not None else ""
|
|
key = (type_name, name) if not param_mode else (type_name, name, "params")
|
|
return self.items.get(key, default)
|