rk_method_beta.py

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