from __future__ import annotations from comfy.samplers import ksampler from pytorch_wavelets import DTCWTForward, DTCWTInverse, DWTForward, DWTInverse from .tensor_image_ops import ( BLENDING_MODES, Sharpen, ) from .upscale import Upscale from .utils import fallback from .vae import VAEHelper class Config: _overridable_fields = { # noqa: RUF012 "blend_by_mode", "blend_mode", "denoised_wavelet_multiplier", "dtcwt_biort", "dtcwt_mode", "dtcwt_qshift", "dwt_flip_filters", "dwt_level", "dwt_mode", "dwt_wave", "fadeout_factor", "guidance_factor", "guidance_mode", "guidance_restart_s_noise", "guidance_restart", "guidance_steps", "reference_wavelet_multiplier", "renoise_factor", "resample_mode", "rescale_increment", "scale_factor", "sharpen_gaussian_kernel_size", "sharpen_gaussian_sigma", "sharpen_mode", "sharpen_reference", "sharpen_strength", "skip_callback", "sigma_offset", "vae_decode_kwargs", "vae_encode_kwargs", "vae_mode", } _dict_exclude_keys = { # noqa: RUF012 "as_dict", "blend_function", "dwt", "get_iteration_config", "idwt", "iteration_override", "sharpen", "upscale", "vae", } def __init__( self, device, dtype, latent_format, *, blend_mode="lerp", blend_by_mode="image", 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", fadeout_factor=0.0, guidance_factor=1.0, guidance_mode="image", guidance_restart_s_noise=1.0, guidance_restart=0, guidance_sampler=None, guidance_steps=5, iteration_override=None, iterations=1, reference_sampler=None, reference_wavelet_multiplier=1.0, renoise_factor=1.0, resample_mode="bicubic", rescale_increment=64, sampler=None, scale_factor=2.0, sharpen_gaussian_kernel_size=3, sharpen_gaussian_sigma=(0.1, 2.0), sharpen_mode="gaussian", sharpen_reference=True, sharpen_strength=1.0, skip_callback=False, sigma_offset=0, upscale_model=None, vae_decode_kwargs=None, vae_encode_kwargs=None, vae_mode="normal", vae=None, ): sampler = fallback( sampler, lambda: ksampler("euler"), default_is_fun=True, ) self.sigma_offset = sigma_offset 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, ) 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.blend_function = BLENDING_MODES[blend_mode] 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, **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"] = self.vae.vae 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["upscale_model"] = self.upscale.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