995 lines
42 KiB
Python
995 lines
42 KiB
Python
import torch
|
|
from torch import Tensor
|
|
from typing import Optional, Callable, Tuple, List, Dict, Any, Union
|
|
|
|
import comfy.model_patcher
|
|
import comfy.supported_models
|
|
|
|
import itertools
|
|
|
|
from .phi_functions import Phi
|
|
from .rk_coefficients_beta import get_implicit_sampler_name_list, get_rk_methods_beta
|
|
from ..helper import ExtraOptions
|
|
from ..latents import get_orthogonal, get_collinear, get_cosine_similarity, tile_latent, untile_latent
|
|
|
|
from ..res4lyf import RESplain
|
|
|
|
MAX_STEPS = 10000
|
|
|
|
|
|
def get_data_from_step (x:Tensor, x_next:Tensor, sigma:Tensor, sigma_next:Tensor) -> Tensor:
|
|
h = sigma_next - sigma
|
|
return (sigma_next * x - sigma * x_next) / h
|
|
|
|
def get_epsilon_from_step(x:Tensor, x_next:Tensor, sigma:Tensor, sigma_next:Tensor) -> Tensor:
|
|
h = sigma_next - sigma
|
|
return (x - x_next) / h
|
|
|
|
|
|
|
|
class RK_Method_Beta:
|
|
def __init__(self,
|
|
model,
|
|
rk_type : str,
|
|
noise_anchor : float,
|
|
noise_boost_normalize : bool = True,
|
|
model_device : str = 'cuda',
|
|
work_device : str = 'cpu',
|
|
dtype : torch.dtype = torch.float64,
|
|
extra_options : str = ""
|
|
):
|
|
|
|
self.work_device = work_device
|
|
self.model_device = model_device
|
|
self.dtype : torch.dtype = dtype
|
|
|
|
self.model = model
|
|
|
|
if hasattr(model, "model"):
|
|
model_sampling = model.model.model_sampling
|
|
elif hasattr(model, "inner_model"):
|
|
model_sampling = model.inner_model.inner_model.model_sampling
|
|
|
|
self.sigma_min : Tensor = model_sampling.sigma_min.to(dtype=dtype, device=work_device)
|
|
self.sigma_max : Tensor = model_sampling.sigma_max.to(dtype=dtype, device=work_device)
|
|
|
|
self.rk_type : str = rk_type
|
|
|
|
self.IMPLICIT : str = rk_type in get_implicit_sampler_name_list(nameOnly=True)
|
|
self.EXPONENTIAL : bool = RK_Method_Beta.is_exponential(rk_type)
|
|
|
|
self.SYNC_SUBSTEP_MEAN_CW : bool = noise_boost_normalize
|
|
|
|
self.A : Optional[Tensor] = None
|
|
self.B : Optional[Tensor] = None
|
|
self.U : Optional[Tensor] = None
|
|
self.V : Optional[Tensor] = None
|
|
|
|
self.rows : int = 0
|
|
self.cols : int = 0
|
|
|
|
self.denoised : Optional[Tensor] = None
|
|
self.uncond : Optional[Tensor] = None
|
|
|
|
self.y0 : Optional[Tensor] = None
|
|
self.y0_inv : Optional[Tensor] = None
|
|
|
|
self.multistep_stages : int = 0
|
|
self.row_offset : Optional[int] = None
|
|
|
|
self.cfg_cw : float = 1.0
|
|
self.extra_args : Optional[Dict[str, Any]] = None
|
|
|
|
self.extra_options : str = extra_options
|
|
self.EO : ExtraOptions = ExtraOptions(extra_options)
|
|
|
|
self.reorder_tableau_indices : list[int] = self.EO("reorder_tableau_indices", [-1])
|
|
|
|
self.LINEAR_ANCHOR_X_0 : float = noise_anchor
|
|
|
|
self.tile_sizes : Optional[List[Tuple[int,int]]] = None
|
|
self.tile_cnt : int = 0
|
|
self.latent_compression_ratio : int = 8
|
|
|
|
@staticmethod
|
|
def is_exponential(rk_type:str) -> bool:
|
|
if rk_type.startswith(( "res",
|
|
"dpmpp",
|
|
"ddim",
|
|
"pec",
|
|
"etdrk",
|
|
"lawson",
|
|
"abnorsett",
|
|
)):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
@staticmethod
|
|
def create(model,
|
|
rk_type : str,
|
|
noise_anchor : float = 1.0,
|
|
noise_boost_normalize : bool = True,
|
|
model_device : str = 'cuda',
|
|
work_device : str = 'cpu',
|
|
dtype : torch.dtype = torch.float64,
|
|
extra_options : str = ""
|
|
) -> "Union[RK_Method_Exponential, RK_Method_Linear]":
|
|
|
|
if RK_Method_Beta.is_exponential(rk_type):
|
|
return RK_Method_Exponential(model, rk_type, noise_anchor, noise_boost_normalize, model_device, work_device, dtype, extra_options)
|
|
else:
|
|
return RK_Method_Linear (model, rk_type, noise_anchor, noise_boost_normalize, model_device, work_device, dtype, extra_options)
|
|
|
|
def __call__(self):
|
|
raise NotImplementedError("This method got clownsharked!")
|
|
|
|
def model_epsilon(self, x:Tensor, sigma:Tensor, **extra_args) -> Tuple[Tensor, Tensor]:
|
|
s_in = x.new_ones([x.shape[0]])
|
|
denoised = self.model(x, sigma * s_in, **extra_args)
|
|
denoised = self.calc_cfg_channelwise(denoised)
|
|
eps = (x - denoised) / (sigma * s_in).view(x.shape[0], 1, 1, 1) #return x0 ###################################THIS WORKS ONLY WITH THE MODEL SAMPLING PATCH
|
|
return eps, denoised
|
|
|
|
def model_denoised(self, x:Tensor, sigma:Tensor, **extra_args) -> Tensor:
|
|
s_in = x.new_ones([x.shape[0]])
|
|
control_tiles = None
|
|
y0_style_pos = self.extra_args['model_options']['transformer_options'].get("y0_style_pos")
|
|
y0_style_neg = self.extra_args['model_options']['transformer_options'].get("y0_style_neg")
|
|
y0_style_pos_tile, sy0_style_neg_tiles = None, None
|
|
|
|
if self.EO("tile_model_calls"):
|
|
tile_h = self.EO("tile_h", 128)
|
|
tile_w = self.EO("tile_w", 128)
|
|
|
|
denoised_tiles = []
|
|
|
|
tiles, orig_shape, grid, strides = tile_latent(x, tile_size=(tile_h,tile_w))
|
|
|
|
for i in range(tiles.shape[0]):
|
|
tile = tiles[i].unsqueeze(0)
|
|
|
|
denoised_tile = self.model(tile, sigma * s_in, **extra_args)
|
|
|
|
denoised_tiles.append(denoised_tile)
|
|
|
|
denoised_tiles = torch.cat(denoised_tiles, dim=0)
|
|
|
|
denoised = untile_latent(denoised_tiles, orig_shape, grid, strides)
|
|
|
|
elif self.tile_sizes is not None:
|
|
tile_h_full = self.tile_sizes[self.tile_cnt % len(self.tile_sizes)][0]
|
|
tile_w_full = self.tile_sizes[self.tile_cnt % len(self.tile_sizes)][1]
|
|
|
|
if tile_h_full == -1:
|
|
tile_h = x.shape[-2]
|
|
tile_h_full = tile_h * self.latent_compression_ratio
|
|
else:
|
|
tile_h = tile_h_full // self.latent_compression_ratio
|
|
|
|
if tile_w_full == -1:
|
|
tile_w = x.shape[-1]
|
|
tile_w_full = tile_w * self.latent_compression_ratio
|
|
else:
|
|
tile_w = tile_w_full // self.latent_compression_ratio
|
|
|
|
#tile_h = tile_h_full // self.latent_compression_ratio
|
|
#tile_w = tile_w_full // self.latent_compression_ratio
|
|
|
|
self.tile_cnt += 1
|
|
|
|
#if len(self.tile_sizes) == 1 and self.tile_cnt % 2 == 1:
|
|
# tile_h, tile_w = tile_w, tile_h
|
|
# tile_h_full, tile_w_full = tile_w_full, tile_h_full
|
|
|
|
if (self.tile_cnt // len(self.tile_sizes)) % 2 == 1 and self.EO("tiles_autorotate"):
|
|
tile_h, tile_w = tile_w, tile_h
|
|
tile_h_full, tile_w_full = tile_w_full, tile_h_full
|
|
|
|
xt_negative = self.model.inner_model.conds.get('xt_negative', self.model.inner_model.conds.get('negative'))
|
|
negative_control = xt_negative[0].get('control')
|
|
|
|
if negative_control is not None and hasattr(negative_control, 'cond_hint_original'):
|
|
negative_cond_hint_init = negative_control.cond_hint.clone() if negative_control.cond_hint is not None else None
|
|
|
|
xt_positive = self.model.inner_model.conds.get('xt_positive', self.model.inner_model.conds.get('positive'))
|
|
positive_control = xt_positive[0].get('control')
|
|
|
|
if positive_control is not None and hasattr(positive_control, 'cond_hint_original'):
|
|
positive_cond_hint_init = positive_control.cond_hint.clone() if positive_control.cond_hint is not None else None
|
|
if positive_control.cond_hint_original.shape[-1] != x.shape[-2] * self.latent_compression_ratio or positive_control.cond_hint_original.shape[-2] != x.shape[-1] * self.latent_compression_ratio:
|
|
positive_control_pretile = comfy.utils.bislerp(positive_control.cond_hint_original.clone().to(torch.float16).to('cuda'), x.shape[-1] * self.latent_compression_ratio, x.shape[-2] * self.latent_compression_ratio)
|
|
positive_control.cond_hint_original = positive_control_pretile.to(positive_control.cond_hint_original)
|
|
positive_control_pretile = positive_control.cond_hint_original.clone().to(torch.float16).to('cuda')
|
|
control_tiles, control_orig_shape, control_grid, control_strides = tile_latent(positive_control_pretile, tile_size=(tile_h_full,tile_w_full))
|
|
control_tiles = control_tiles
|
|
|
|
denoised_tiles = []
|
|
|
|
tiles, orig_shape, grid, strides = tile_latent(x, tile_size=(tile_h,tile_w))
|
|
|
|
if y0_style_pos is not None:
|
|
y0_style_pos_tiles, _, _, _ = tile_latent(y0_style_pos, tile_size=(tile_h,tile_w))
|
|
if y0_style_neg is not None:
|
|
y0_style_neg_tiles, _, _, _ = tile_latent(y0_style_neg, tile_size=(tile_h,tile_w))
|
|
|
|
for i in range(tiles.shape[0]):
|
|
tile = tiles[i].unsqueeze(0)
|
|
|
|
if control_tiles is not None:
|
|
positive_control.cond_hint = control_tiles[i].unsqueeze(0).to(positive_control.cond_hint)
|
|
if negative_control is not None:
|
|
negative_control.cond_hint = control_tiles[i].unsqueeze(0).to(positive_control.cond_hint)
|
|
|
|
if y0_style_pos is not None:
|
|
self.extra_args['model_options']['transformer_options']['y0_style_pos'] = y0_style_pos_tiles[i].unsqueeze(0)
|
|
if y0_style_neg is not None:
|
|
self.extra_args['model_options']['transformer_options']['y0_style_neg'] = y0_style_neg_tiles[i].unsqueeze(0)
|
|
|
|
denoised_tile = self.model(tile, sigma * s_in, **extra_args)
|
|
|
|
denoised_tiles.append(denoised_tile)
|
|
|
|
denoised_tiles = torch.cat(denoised_tiles, dim=0)
|
|
|
|
denoised = untile_latent(denoised_tiles, orig_shape, grid, strides)
|
|
|
|
else:
|
|
denoised = self.model(x, sigma * s_in, **extra_args)
|
|
|
|
if control_tiles is not None:
|
|
positive_control.cond_hint = positive_cond_hint_init
|
|
if negative_control is not None:
|
|
negative_control.cond_hint = negative_cond_hint_init
|
|
|
|
if y0_style_pos is not None:
|
|
self.extra_args['model_options']['transformer_options']['y0_style_pos'] = y0_style_pos
|
|
if y0_style_neg is not None:
|
|
self.extra_args['model_options']['transformer_options']['y0_style_neg'] = y0_style_neg
|
|
|
|
denoised = self.calc_cfg_channelwise(denoised)
|
|
return denoised
|
|
|
|
def update_transformer_options(self,
|
|
transformer_options : Optional[dict] = None,
|
|
):
|
|
|
|
self.extra_args.setdefault("model_options", {}).setdefault("transformer_options", {}).update(transformer_options)
|
|
return
|
|
|
|
def set_coeff(self,
|
|
rk_type : str,
|
|
h : Tensor,
|
|
c1 : float = 0.0,
|
|
c2 : float = 0.5,
|
|
c3 : float = 1.0,
|
|
step : int = 0,
|
|
sigmas : Optional[Tensor] = None,
|
|
sigma_down : Optional[Tensor] = None,
|
|
) -> None:
|
|
|
|
self.rk_type = rk_type
|
|
self.IMPLICIT = rk_type in get_implicit_sampler_name_list(nameOnly=True)
|
|
self.EXPONENTIAL = RK_Method_Beta.is_exponential(rk_type)
|
|
|
|
sigma = sigmas[step]
|
|
sigma_next = sigmas[step+1]
|
|
|
|
h_prev = []
|
|
a, b, u, v, ci, multistep_stages, hybrid_stages, FSAL = get_rk_methods_beta(rk_type,
|
|
h,
|
|
c1,
|
|
c2,
|
|
c3,
|
|
h_prev,
|
|
step,
|
|
sigmas,
|
|
sigma,
|
|
sigma_next,
|
|
sigma_down,
|
|
self.extra_options,
|
|
)
|
|
|
|
self.multistep_stages = multistep_stages
|
|
self.hybrid_stages = hybrid_stages
|
|
|
|
self.A = torch.tensor(a, dtype=h.dtype, device=h.device)
|
|
self.B = torch.tensor(b, dtype=h.dtype, device=h.device)
|
|
self.C = torch.tensor(ci, dtype=h.dtype, device=h.device)
|
|
|
|
self.U = torch.tensor(u, dtype=h.dtype, device=h.device) if u is not None else None
|
|
self.V = torch.tensor(v, dtype=h.dtype, device=h.device) if v is not None else None
|
|
|
|
self.rows = self.A.shape[0]
|
|
self.cols = self.A.shape[1]
|
|
|
|
self.row_offset = 1 if not self.IMPLICIT and self.A[0].sum() == 0 else 0
|
|
|
|
if self.IMPLICIT and self.reorder_tableau_indices[0] != -1:
|
|
self.reorder_tableau(self.reorder_tableau_indices)
|
|
|
|
|
|
|
|
def reorder_tableau(self, indices:list[int]) -> None:
|
|
#if indices[0]:
|
|
self.A = self.A [indices]
|
|
self.B[0] = self.B[0][indices]
|
|
self.C = self.C [indices]
|
|
self.C = torch.cat((self.C, self.C[-1:]))
|
|
return
|
|
|
|
|
|
|
|
def update_substep(self,
|
|
x_0 : Tensor,
|
|
x_ : Tensor,
|
|
eps_ : Tensor,
|
|
eps_prev_ : Tensor,
|
|
row : int,
|
|
row_offset : int,
|
|
h_new : Tensor,
|
|
h_new_orig : Tensor,
|
|
lying_eps_row_factor : float = 1.0,
|
|
) -> Tensor:
|
|
|
|
if row < self.rows - row_offset and self.multistep_stages == 0:
|
|
row_tmp_offset = row + row_offset
|
|
|
|
else:
|
|
row_tmp_offset = row + 1
|
|
|
|
zr_base = self.zum(row+row_offset+self.multistep_stages, eps_, eps_prev_)
|
|
|
|
if self.SYNC_SUBSTEP_MEAN_CW and lying_eps_row_factor != 1.0:
|
|
zr_orig = self.zum(row+row_offset+self.multistep_stages, eps_, eps_prev_)
|
|
x_orig_row = x_0 + h_new * zr_orig
|
|
|
|
#eps_row = eps_ [row].clone()
|
|
#eps_prev_row = eps_prev_[row].clone()
|
|
|
|
eps_ [row] *= lying_eps_row_factor
|
|
eps_prev_[row] *= lying_eps_row_factor
|
|
zr = self.zum(row+row_offset+self.multistep_stages, eps_, eps_prev_)
|
|
|
|
x_[row_tmp_offset] = x_0 + h_new * zr
|
|
|
|
if self.SYNC_SUBSTEP_MEAN_CW and lying_eps_row_factor != 1.0:
|
|
x_[row_tmp_offset] = x_[row_tmp_offset] - x_[row_tmp_offset].mean(dim=(-2,-1), keepdim=True) + x_orig_row.mean(dim=(-2,-1), keepdim=True)
|
|
|
|
#eps_ [row] = eps_row
|
|
#eps_prev_[row] = eps_prev_row
|
|
|
|
if (self.SYNC_SUBSTEP_MEAN_CW and h_new != h_new_orig) or self.EO("sync_mean_noise"):
|
|
if not self.EO("disable_sync_mean_noise"):
|
|
x_row_down = x_0 + h_new_orig * zr
|
|
x_[row_tmp_offset] = x_[row_tmp_offset] - x_[row_tmp_offset].mean(dim=(-2,-1), keepdim=True) + x_row_down.mean(dim=(-2,-1), keepdim=True)
|
|
|
|
return x_
|
|
|
|
|
|
|
|
def a_k_einsum(self, row:int, k :Tensor) -> Tensor:
|
|
return torch.einsum('i, i... -> ...', self.A[row], k[:self.cols])
|
|
|
|
def b_k_einsum(self, row:int, k :Tensor) -> Tensor:
|
|
return torch.einsum('i, i... -> ...', self.B[row], k[:self.cols])
|
|
|
|
def u_k_einsum(self, row:int, k_prev:Tensor) -> Tensor:
|
|
return torch.einsum('i, i... -> ...', self.U[row], k_prev[:self.cols]) if (self.U is not None and k_prev is not None) else 0
|
|
|
|
def v_k_einsum(self, row:int, k_prev:Tensor) -> Tensor:
|
|
return torch.einsum('i, i... -> ...', self.V[row], k_prev[:self.cols]) if (self.V is not None and k_prev is not None) else 0
|
|
|
|
|
|
|
|
def zum(self, row:int, k:Tensor, k_prev:Tensor=None,) -> Tensor:
|
|
if row < self.rows:
|
|
return self.a_k_einsum(row, k) + self.u_k_einsum(row, k_prev)
|
|
else:
|
|
row = row - self.rows
|
|
return self.b_k_einsum(row, k) + self.v_k_einsum(row, k_prev)
|
|
|
|
def zum_tableau(self, k:Tensor, k_prev:Tensor=None,) -> Tensor:
|
|
a_k_sum = torch.einsum('ij, j... -> i...', self.A, k[:self.cols])
|
|
u_k_sum = torch.einsum('ij, j... -> i...', self.U, k_prev[:self.cols]) if (self.U is not None and k_prev is not None) else 0
|
|
return a_k_sum + u_k_sum
|
|
|
|
|
|
|
|
def init_cfg_channelwise(self, x:Tensor, cfg_cw:float=1.0, **extra_args) -> Dict[str, Any]:
|
|
self.uncond = [torch.full_like(x, 0.0)]
|
|
self.cfg_cw = cfg_cw
|
|
if cfg_cw != 1.0:
|
|
def post_cfg_function(args):
|
|
self.uncond[0] = args["uncond_denoised"]
|
|
return args["denoised"]
|
|
model_options = extra_args.get("model_options", {}).copy()
|
|
extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
|
|
return extra_args
|
|
|
|
|
|
def calc_cfg_channelwise(self, denoised:Tensor) -> Tensor:
|
|
if self.cfg_cw != 1.0:
|
|
avg = 0
|
|
for b, c in itertools.product(range(denoised.shape[0]), range(denoised.shape[1])):
|
|
avg += torch.norm(denoised[b][c] - self.uncond[0][b][c])
|
|
avg /= denoised.shape[1]
|
|
|
|
for b, c in itertools.product(range(denoised.shape[0]), range(denoised.shape[1])):
|
|
ratio = torch.nan_to_num(torch.norm(denoised[b][c] - self.uncond[0][b][c]) / avg, 0)
|
|
denoised_new = self.uncond[0] + ratio * self.cfg_cw * (denoised - self.uncond[0])
|
|
return denoised_new
|
|
else:
|
|
return denoised
|
|
|
|
|
|
@staticmethod
|
|
def calculate_res_2m_step(
|
|
x_0 : Tensor,
|
|
denoised_ : Tensor,
|
|
sigma_down : Tensor,
|
|
sigmas : Tensor,
|
|
step : int,
|
|
) -> Tuple[Tensor, Tensor]:
|
|
|
|
if denoised_[2].sum() == 0:
|
|
return None, None
|
|
|
|
sigma = sigmas[step]
|
|
sigma_prev = sigmas[step-1]
|
|
|
|
h_prev = -torch.log(sigma/sigma_prev)
|
|
h = -torch.log(sigma_down/sigma)
|
|
|
|
c1 = 0
|
|
c2 = (-h_prev / h).item()
|
|
|
|
ci = [c1,c2]
|
|
φ = Phi(h, ci, analytic_solution=True)
|
|
|
|
b2 = φ(2)/c2
|
|
b1 = φ(1) - b2
|
|
|
|
eps_2 = denoised_[1] - x_0
|
|
eps_1 = denoised_[0] - x_0
|
|
|
|
h_a_k_sum = h * (b1 * eps_1 + b2 * eps_2)
|
|
|
|
x = torch.exp(-h) * x_0 + h_a_k_sum
|
|
|
|
denoised = x_0 + (sigma / (sigma - sigma_down)) * h_a_k_sum
|
|
|
|
return x, denoised
|
|
|
|
|
|
@staticmethod
|
|
def calculate_res_3m_step(
|
|
x_0 : Tensor,
|
|
denoised_ : Tensor,
|
|
sigma_down : Tensor,
|
|
sigmas : Tensor,
|
|
step : int,
|
|
) -> Tuple[Tensor, Tensor]:
|
|
|
|
if denoised_[3].sum() == 0:
|
|
return None, None
|
|
|
|
sigma = sigmas[step]
|
|
sigma_prev = sigmas[step-1]
|
|
sigma_prev2 = sigmas[step-2]
|
|
|
|
h = -torch.log(sigma_down/sigma)
|
|
h_prev = -torch.log(sigma/sigma_prev)
|
|
h_prev2 = -torch.log(sigma/sigma_prev2)
|
|
|
|
c1 = 0
|
|
c2 = (-h_prev / h).item()
|
|
c3 = (-h_prev2 / h).item()
|
|
|
|
ci = [c1,c2,c3]
|
|
φ = Phi(h, ci, analytic_solution=True)
|
|
|
|
gamma = (3*(c3**3) - 2*c3) / (c2*(2 - 3*c2))
|
|
|
|
b3 = (1 / (gamma * c2 + c3)) * φ(2, -h)
|
|
b2 = gamma * b3
|
|
b1 = φ(1, -h) - b2 - b3
|
|
|
|
eps_3 = denoised_[2] - x_0
|
|
eps_2 = denoised_[1] - x_0
|
|
eps_1 = denoised_[0] - x_0
|
|
|
|
h_a_k_sum = h * (b1 * eps_1 + b2 * eps_2 + b3 * eps_3)
|
|
|
|
x = torch.exp(-h) * x_0 + h_a_k_sum
|
|
|
|
denoised = x_0 + (sigma / (sigma - sigma_down)) * h_a_k_sum
|
|
|
|
return x, denoised
|
|
|
|
def swap_rk_type_at_step_or_threshold(self,
|
|
x_0 : Tensor,
|
|
data_prev_ : Tensor,
|
|
NS,
|
|
sigmas : Tensor,
|
|
step : Tensor,
|
|
rk_swap_step : int,
|
|
rk_swap_threshold : float,
|
|
rk_swap_type : str,
|
|
rk_swap_print : bool,
|
|
) -> str:
|
|
if rk_swap_type == "":
|
|
if self.EXPONENTIAL:
|
|
rk_swap_type = "res_3m"
|
|
else:
|
|
rk_swap_type = "deis_3m"
|
|
|
|
if step > rk_swap_step and self.rk_type != rk_swap_type:
|
|
RESplain("Switching rk_type to:", rk_swap_type)
|
|
self.rk_type = rk_swap_type
|
|
|
|
if RK_Method_Beta.is_exponential(rk_swap_type):
|
|
self.__class__ = RK_Method_Exponential
|
|
else:
|
|
self.__class__ = RK_Method_Linear
|
|
|
|
if rk_swap_type in get_implicit_sampler_name_list(nameOnly=True):
|
|
self.IMPLICIT = True
|
|
self.row_offset = 0
|
|
NS.row_offset = 0
|
|
else:
|
|
self.IMPLICIT = False
|
|
self.row_offset = 1
|
|
NS.row_offset = 1
|
|
NS.h_fn = self.h_fn
|
|
NS.t_fn = self.t_fn
|
|
NS.sigma_fn = self.sigma_fn
|
|
|
|
|
|
|
|
if step > 2 and sigmas[step+1] > 0 and self.rk_type != rk_swap_type and rk_swap_threshold > 0:
|
|
x_res_2m, denoised_res_2m = self.calculate_res_2m_step(x_0, data_prev_, NS.sigma_down, sigmas, step)
|
|
x_res_3m, denoised_res_3m = self.calculate_res_3m_step(x_0, data_prev_, NS.sigma_down, sigmas, step)
|
|
if denoised_res_2m is not None:
|
|
if rk_swap_print:
|
|
RESplain("res_3m - res_2m:", torch.norm(denoised_res_3m - denoised_res_2m).item())
|
|
if rk_swap_threshold > torch.norm(denoised_res_2m - denoised_res_3m):
|
|
RESplain("Switching rk_type to:", rk_swap_type, "at step:", step)
|
|
self.rk_type = rk_swap_type
|
|
|
|
if RK_Method_Beta.is_exponential(rk_swap_type):
|
|
self.__class__ = RK_Method_Exponential
|
|
else:
|
|
self.__class__ = RK_Method_Linear
|
|
|
|
if rk_swap_type in get_implicit_sampler_name_list(nameOnly=True):
|
|
self.IMPLICIT = True
|
|
self.row_offset = 0
|
|
NS.row_offset = 0
|
|
else:
|
|
self.IMPLICIT = False
|
|
self.row_offset = 1
|
|
NS.row_offset = 1
|
|
NS.h_fn = self.h_fn
|
|
NS.t_fn = self.t_fn
|
|
NS.sigma_fn = self.sigma_fn
|
|
|
|
return self.rk_type
|
|
|
|
|
|
def bong_iter(self,
|
|
x_0 : Tensor,
|
|
x_ : Tensor,
|
|
eps_ : Tensor,
|
|
eps_prev_ : Tensor,
|
|
data_ : Tensor,
|
|
sigma : Tensor,
|
|
s_ : Tensor,
|
|
row : int,
|
|
row_offset: int,
|
|
h : Tensor,
|
|
step : int,
|
|
) -> Tuple[Tensor, Tensor, Tensor]:
|
|
|
|
if x_0.ndim == 4:
|
|
norm_dim = (-2,-1)
|
|
elif x_0.ndim == 5:
|
|
norm_dim = (-4,-2,-1)
|
|
|
|
if self.EO("bong_start_step", 0) > step or step > self.EO("bong_stop_step", 10000):
|
|
return x_0, x_, eps_
|
|
|
|
bong_iter_max_row = self.rows - row_offset
|
|
if self.EO("bong_iter_max_row_full"):
|
|
bong_iter_max_row = self.rows
|
|
|
|
if self.EO("bong_iter_lock_x_0_ch_means"):
|
|
x_0_ch_means = x_0.mean(dim=norm_dim, keepdim=True)
|
|
|
|
if self.EO("bong_iter_lock_x_row_ch_means"):
|
|
x_row_means = []
|
|
for rr in range(row+row_offset):
|
|
x_row_mean = x_[rr].mean(dim=norm_dim, keepdim=True)
|
|
x_row_means.append(x_row_mean)
|
|
|
|
if row < bong_iter_max_row and self.multistep_stages == 0:
|
|
bong_strength = self.EO("bong_strength", 1.0)
|
|
|
|
if bong_strength != 1.0:
|
|
x_0_tmp = x_0.clone()
|
|
x_tmp_ = x_.clone()
|
|
eps_tmp_ = eps_.clone()
|
|
|
|
for i in range(100):
|
|
x_0 = x_[row+row_offset] - h * self.zum(row+row_offset, eps_, eps_prev_)
|
|
|
|
if self.EO("bong_iter_lock_x_0_ch_means"):
|
|
x_0 = x_0 - x_0.mean(dim=norm_dim, keepdim=True) + x_0_ch_means
|
|
|
|
for rr in range(row+row_offset):
|
|
x_[rr] = x_0 + h * self.zum(rr, eps_, eps_prev_)
|
|
|
|
if self.EO("bong_iter_lock_x_row_ch_means"):
|
|
for rr in range(row+row_offset):
|
|
x_[rr] = x_[rr] - x_[rr].mean(dim=norm_dim, keepdim=True) + x_row_means[rr]
|
|
|
|
for rr in range(row+row_offset):
|
|
if self.EO("zonkytar"):
|
|
#eps_[rr] = self.get_unsample_epsilon(x_[rr], x_0, data_[rr], sigma, s_[rr])
|
|
eps_[rr] = self.get_epsilon(x_[rr], x_0, data_[rr], sigma, s_[rr])
|
|
else:
|
|
eps_[rr] = self.get_epsilon(x_0, x_[rr], data_[rr], sigma, s_[rr])
|
|
|
|
if bong_strength != 1.0:
|
|
x_0 = x_0_tmp + bong_strength * (x_0 - x_0_tmp)
|
|
x_ = x_tmp_ + bong_strength * (x_ - x_tmp_)
|
|
eps_ = eps_tmp_ + bong_strength * (eps_ - eps_tmp_)
|
|
|
|
return x_0, x_, eps_
|
|
|
|
|
|
def newton_iter(self,
|
|
x_0 : Tensor,
|
|
x_ : Tensor,
|
|
eps_ : Tensor,
|
|
eps_prev_ : Tensor,
|
|
data_ : Tensor,
|
|
s_ : Tensor,
|
|
row : int,
|
|
h : Tensor,
|
|
sigmas : Tensor,
|
|
step : int,
|
|
newton_name: str,
|
|
) -> Tuple[Tensor, Tensor]:
|
|
|
|
newton_iter_name = "newton_iter_" + newton_name
|
|
|
|
default_anchor_x_all = False
|
|
if newton_name == "lying":
|
|
default_anchor_x_all = True
|
|
|
|
newton_iter = self.EO(newton_iter_name, 100)
|
|
newton_iter_skip_last_steps = self.EO(newton_iter_name + "_skip_last_steps", 0)
|
|
newton_iter_mixing_rate = self.EO(newton_iter_name + "_mixing_rate", 1.0)
|
|
|
|
newton_iter_anchor = self.EO(newton_iter_name + "_anchor", 0)
|
|
newton_iter_anchor_x_all = self.EO(newton_iter_name + "_anchor_x_all", default_anchor_x_all)
|
|
newton_iter_type = self.EO(newton_iter_name + "_type", "from_epsilon")
|
|
newton_iter_sequence = self.EO(newton_iter_name + "_sequence", "double")
|
|
|
|
row_b_offset = 0
|
|
if self.EO(newton_iter_name + "_include_row_b"):
|
|
row_b_offset = 1
|
|
|
|
if step >= len(sigmas)-1-newton_iter_skip_last_steps or sigmas[step+1] == 0 or not self.IMPLICIT:
|
|
return x_, eps_
|
|
|
|
sigma = sigmas[step]
|
|
|
|
start, stop = 0, self.rows+row_b_offset
|
|
if newton_name == "pre":
|
|
start = row
|
|
elif newton_name == "post":
|
|
start = row + 1
|
|
|
|
if newton_iter_anchor >= 0:
|
|
eps_anchor = eps_[newton_iter_anchor].clone()
|
|
|
|
if newton_iter_anchor_x_all:
|
|
x_orig_ = x_.clone()
|
|
|
|
for n_iter in range(newton_iter):
|
|
for r in range(start, stop):
|
|
if newton_iter_anchor >= 0:
|
|
eps_[newton_iter_anchor] = eps_anchor.clone()
|
|
if newton_iter_anchor_x_all:
|
|
x_ = x_orig_.clone()
|
|
x_tmp, eps_tmp = x_[r].clone(), eps_[r].clone()
|
|
|
|
seq_start, seq_stop = r, r+1
|
|
|
|
if newton_iter_sequence == "double":
|
|
seq_start, seq_stop = start, stop
|
|
|
|
for r_ in range(seq_start, seq_stop):
|
|
x_[r_] = x_0 + h * self.zum(r_, eps_, eps_prev_)
|
|
|
|
for r_ in range(seq_start, seq_stop):
|
|
if newton_iter_type == "from_data":
|
|
data_[r_] = get_data_from_step(x_0, x_[r_], sigma, s_[r_])
|
|
eps_ [r_] = self.get_epsilon(x_0, x_[r_], data_[r_], sigma, s_[r_])
|
|
elif newton_iter_type == "from_step":
|
|
eps_ [r_] = get_epsilon_from_step(x_0, x_[r_], sigma, s_[r_])
|
|
elif newton_iter_type == "from_alt":
|
|
eps_ [r_] = x_0/sigma - x_[r_]/s_[r_]
|
|
elif newton_iter_type == "from_epsilon":
|
|
eps_ [r_] = self.get_epsilon(x_0, x_[r_], data_[r_], sigma, s_[r_])
|
|
|
|
if self.EO(newton_iter_name + "_opt"):
|
|
opt_timing, opt_type, opt_subtype = self.EO(newton_iter_name+"_opt", [str])
|
|
|
|
opt_start, opt_stop = 0, self.rows+row_b_offset
|
|
if opt_timing == "early":
|
|
opt_stop = row + 1
|
|
elif opt_timing == "late":
|
|
opt_start = row + 1
|
|
|
|
for r2 in range(opt_start, opt_stop):
|
|
if r_ != r2:
|
|
if opt_subtype == "a":
|
|
eps_a = eps_[r2]
|
|
eps_b = eps_[r_]
|
|
elif opt_subtype == "b":
|
|
eps_a = eps_[r_]
|
|
eps_b = eps_[r2]
|
|
|
|
if opt_type == "ortho":
|
|
eps_ [r_] = get_orthogonal(eps_a, eps_b)
|
|
elif opt_type == "collin":
|
|
eps_ [r_] = get_collinear (eps_a, eps_b)
|
|
elif opt_type == "proj":
|
|
eps_ [r_] = get_collinear (eps_a, eps_b) + get_orthogonal(eps_b, eps_a)
|
|
|
|
x_ [r_] = x_tmp + newton_iter_mixing_rate * (x_ [r_] - x_tmp)
|
|
eps_[r_] = eps_tmp + newton_iter_mixing_rate * (eps_[r_] - eps_tmp)
|
|
|
|
if newton_iter_sequence == "double":
|
|
break
|
|
|
|
return x_, eps_
|
|
|
|
|
|
|
|
|
|
class RK_Method_Exponential(RK_Method_Beta):
|
|
def __init__(self,
|
|
model,
|
|
rk_type : str,
|
|
noise_anchor : float,
|
|
noise_boost_normalize : bool,
|
|
|
|
model_device : str = 'cuda',
|
|
work_device : str = 'cpu',
|
|
dtype : torch.dtype = torch.float64,
|
|
extra_options : str = "",
|
|
):
|
|
|
|
super().__init__(model,
|
|
rk_type,
|
|
noise_anchor,
|
|
noise_boost_normalize,
|
|
model_device = model_device,
|
|
work_device = work_device,
|
|
dtype = dtype,
|
|
extra_options = extra_options,
|
|
)
|
|
|
|
@staticmethod
|
|
def alpha_fn(neg_h:Tensor) -> Tensor:
|
|
return torch.exp(neg_h)
|
|
|
|
@staticmethod
|
|
def sigma_fn(t:Tensor) -> Tensor:
|
|
return t.neg().exp()
|
|
|
|
@staticmethod
|
|
def t_fn(sigma:Tensor) -> Tensor:
|
|
return sigma.log().neg()
|
|
|
|
@staticmethod
|
|
def h_fn(sigma_down:Tensor, sigma:Tensor) -> Tensor:
|
|
return -torch.log(sigma_down/sigma)
|
|
|
|
def __call__(self,
|
|
x : Tensor,
|
|
sub_sigma : Tensor,
|
|
x_0 : Optional[Tensor] = None,
|
|
sigma : Optional[Tensor] = None,
|
|
transformer_options : Optional[dict] = None,
|
|
) -> Tuple[Tensor, Tensor]:
|
|
|
|
x_0 = x if x_0 is None else x_0
|
|
sigma = sub_sigma if sigma is None else sigma
|
|
|
|
if transformer_options is not None:
|
|
self.extra_args.setdefault("model_options", {}).setdefault("transformer_options", {}).update(transformer_options)
|
|
|
|
denoised = self.model_denoised(x.to(self.model_device), sub_sigma.to(self.model_device), **self.extra_args).to(sigma.device)
|
|
|
|
eps_anchored = (x_0 - denoised) / sigma
|
|
eps_unmoored = (x - denoised) / sub_sigma
|
|
|
|
eps = eps_unmoored + self.LINEAR_ANCHOR_X_0 * (eps_anchored - eps_unmoored)
|
|
|
|
denoised = x_0 - sigma * eps
|
|
|
|
epsilon = denoised - x_0
|
|
|
|
return epsilon, denoised
|
|
|
|
|
|
|
|
def get_epsilon(self,
|
|
x_0 : Tensor,
|
|
x : Tensor,
|
|
denoised : Tensor,
|
|
sigma : Tensor,
|
|
sub_sigma : Tensor,
|
|
) -> Tensor:
|
|
|
|
eps_anchored = (x_0 - denoised) / sigma
|
|
eps_unmoored = (x - denoised) / sub_sigma
|
|
|
|
eps = eps_unmoored + self.LINEAR_ANCHOR_X_0 * (eps_anchored - eps_unmoored)
|
|
|
|
denoised = x_0 - sigma * eps
|
|
|
|
return denoised - x_0
|
|
|
|
|
|
|
|
def get_epsilon_anchored(self, x_0:Tensor, denoised:Tensor, sigma:Tensor) -> Tensor:
|
|
return denoised - x_0
|
|
|
|
|
|
|
|
def get_guide_epsilon(self,
|
|
x_0 : Tensor,
|
|
x : Tensor,
|
|
y : Tensor,
|
|
sigma : Tensor,
|
|
sigma_cur : Tensor,
|
|
sigma_down : Optional[Tensor] = None,
|
|
epsilon_scale : Optional[Tensor] = None,
|
|
) -> Tensor:
|
|
|
|
sigma_cur = epsilon_scale if epsilon_scale is not None else sigma_cur
|
|
|
|
if sigma_down > sigma:
|
|
eps_unmoored = (sigma_cur/(self.sigma_max - sigma_cur)) * (x - y)
|
|
else:
|
|
eps_unmoored = y - x
|
|
|
|
if self.EO("manually_anchor_unsampler"):
|
|
if sigma_down > sigma:
|
|
eps_anchored = (sigma /(self.sigma_max - sigma)) * (x_0 - y)
|
|
else:
|
|
eps_anchored = y - x_0
|
|
eps_guide = eps_unmoored + self.LINEAR_ANCHOR_X_0 * (eps_anchored - eps_unmoored)
|
|
else:
|
|
eps_guide = eps_unmoored
|
|
|
|
return eps_guide
|
|
|
|
|
|
|
|
|
|
class RK_Method_Linear(RK_Method_Beta):
|
|
def __init__(self,
|
|
model,
|
|
rk_type : str,
|
|
noise_anchor : float,
|
|
noise_boost_normalize : bool,
|
|
model_device : str = 'cuda',
|
|
work_device : str = 'cpu',
|
|
dtype : torch.dtype = torch.float64,
|
|
extra_options : str = "",
|
|
):
|
|
|
|
super().__init__(model,
|
|
rk_type,
|
|
noise_anchor,
|
|
noise_boost_normalize,
|
|
model_device = model_device,
|
|
work_device = work_device,
|
|
dtype = dtype,
|
|
extra_options = extra_options,
|
|
)
|
|
|
|
@staticmethod
|
|
def alpha_fn(neg_h:Tensor) -> Tensor:
|
|
return torch.ones_like(neg_h)
|
|
|
|
@staticmethod
|
|
def sigma_fn(t:Tensor) -> Tensor:
|
|
return t
|
|
|
|
@staticmethod
|
|
def t_fn(sigma:Tensor) -> Tensor:
|
|
return sigma
|
|
|
|
@staticmethod
|
|
def h_fn(sigma_down:Tensor, sigma:Tensor) -> Tensor:
|
|
return sigma_down - sigma
|
|
|
|
def __call__(self,
|
|
x : Tensor,
|
|
sub_sigma : Tensor,
|
|
x_0 : Optional[Tensor] = None,
|
|
sigma : Optional[Tensor] = None,
|
|
transformer_options : Optional[dict] = None,
|
|
) -> Tuple[Tensor, Tensor]:
|
|
|
|
x_0 = x if x_0 is None else x_0
|
|
sigma = sub_sigma if sigma is None else sigma
|
|
|
|
if transformer_options is not None:
|
|
self.extra_args.setdefault("model_options", {}).setdefault("transformer_options", {}).update(transformer_options)
|
|
|
|
denoised = self.model_denoised(x.to(self.model_device), sub_sigma.to(self.model_device), **self.extra_args).to(sigma.device)
|
|
|
|
epsilon_anchor = (x_0 - denoised) / sigma
|
|
epsilon_unmoored = (x - denoised) / sub_sigma
|
|
|
|
epsilon = epsilon_unmoored + self.LINEAR_ANCHOR_X_0 * (epsilon_anchor - epsilon_unmoored)
|
|
|
|
return epsilon, denoised
|
|
|
|
|
|
|
|
def get_epsilon(self,
|
|
x_0 : Tensor,
|
|
x : Tensor,
|
|
denoised : Tensor,
|
|
sigma : Tensor,
|
|
sub_sigma : Tensor,
|
|
) -> Tensor:
|
|
|
|
eps_anchor = (x_0 - denoised) / sigma
|
|
eps_unmoored = (x - denoised) / sub_sigma
|
|
|
|
return eps_unmoored + self.LINEAR_ANCHOR_X_0 * (eps_anchor - eps_unmoored)
|
|
|
|
|
|
|
|
def get_epsilon_anchored(self, x_0:Tensor, denoised:Tensor, sigma:Tensor) -> Tensor:
|
|
return (x_0 - denoised) / sigma
|
|
|
|
|
|
|
|
def get_guide_epsilon(self,
|
|
x_0 : Tensor,
|
|
x : Tensor,
|
|
y : Tensor,
|
|
sigma : Tensor,
|
|
sigma_cur : Tensor,
|
|
sigma_down : Optional[Tensor] = None,
|
|
epsilon_scale : Optional[Tensor] = None,
|
|
) -> Tensor:
|
|
|
|
if sigma_down > sigma:
|
|
sigma_ratio = self.sigma_max - sigma_cur.clone()
|
|
else:
|
|
sigma_ratio = sigma_cur.clone()
|
|
sigma_ratio = epsilon_scale if epsilon_scale is not None else sigma_ratio
|
|
|
|
if sigma_down is None:
|
|
return (x - y) / sigma_ratio
|
|
else:
|
|
if sigma_down > sigma:
|
|
return (y - x) / sigma_ratio
|
|
else:
|
|
return (x - y) / sigma_ratio
|
|
|
|
|