import math

import torch
from torch import nn, FloatTensor
import torchsde
import kornia
from tqdm.auto import trange, tqdm
import numpy as np

import comfy.sample
import comfy.model_patcher

from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler, PIDStepSizeController, get_ancestral_step, to_d, default_noise_sampler, DPMSolver

# The following function adds the samplers during initialization, in __init__.py
def add_samplers():
    from comfy.samplers import KSampler, k_diffusion_sampling
    if hasattr(KSampler, "DISCARD_PENULTIMATE_SIGMA_SAMPLERS"):
        KSampler.DISCARD_PENULTIMATE_SIGMA_SAMPLERS |= discard_penultimate_sigma_samplers
    added = 0
    for sampler in extra_samplers: #getattr(self, "sample_{}".format(extra_samplers))
        if sampler not in KSampler.SAMPLERS:
            try:
                idx = KSampler.SAMPLERS.index("uni_pc_bh2") # Last item in the samplers list
                KSampler.SAMPLERS.insert(idx+1, sampler) # Add our custom samplers
                setattr(k_diffusion_sampling, "sample_{}".format(sampler), extra_samplers[sampler])
                added += 1
            except ValueError as _err:
                pass
    if added > 0:
        import importlib
        importlib.reload(k_diffusion_sampling)

# The following function adds the samplers during initialization, in __init__.py
def add_schedulers():
    from comfy.samplers import KSampler, k_diffusion_sampling
    added = 0
    for scheduler in extra_schedulers: #getattr(self, "sample_{}".format(extra_samplers))
        if scheduler not in KSampler.SCHEDULERS:
            try:
                idx = KSampler.SCHEDULERS.index("ddim_uniform") # Last item in the samplers list
                KSampler.SCHEDULERS.insert(idx+1, scheduler) # Add our custom samplers
                setattr(k_diffusion_sampling, "get_sigmas_{}".format(scheduler), extra_schedulers[scheduler])
                added += 1
            except ValueError as err:
                pass
    if added > 0:
        import importlib
        importlib.reload(k_diffusion_sampling)


# Noise samplers
IMMISCIBLE_NOISE_NAMES=("gaussian", "perlin")
NOISE_SAMPLER_NAMES=("gaussian", "uniform", "brownian", "highres-pyramid", "pyramid", "perlin", "laplacian", "immiscible_gaussian", "immiscible_gaussian_maximize", "immiscible_perlin", "immiscible_perlin_maximize")

def get_noise_sampler_names(default=None):
    if not default:
        return NOISE_SAMPLER_NAMES
    return (default,) + tuple(n for n in NOISE_SAMPLER_NAMES if n != default)

def get_immiscible_noise_sampler_names(default=None):
    if not default:
        return IMMISCIBLE_NOISE_NAMES
    return (default,) + tuple(n for n in IMMISCIBLE_NOISE_NAMES if n != default)

def mk_noise_sampler(x, fun):
    return lambda _sigma, _sigma_next: fun(x)

def get_noise_sampler(x, sigmas, noise_sampler_type="brownian", extra_args=None, cpu=False):
    if noise_sampler_type == "brownian":
        seed = extra_args.get("seed", None) if extra_args else None
        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
        return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=cpu)
    return mk_noise_sampler(x, NOISE_SAMPLER_HANDLERS.get(noise_sampler_type, uniform_noise_like))

from torch import Generator, Tensor, lerp
from torch.nn.functional import unfold
from typing import Callable, Tuple
from math import pi

def uniform_noise_like(x):
    return (torch.rand_like(x) - 0.5) * 2 * 1.73

from scipy.optimize import linear_sum_assignment
def check_set_immiscible(x, noise_sampler_type, extra_args):
    if noise_sampler_type.startswith("immiscible"):
        match noise_sampler_type:
            case "immiscible_gaussian":
                immiscibility = make_immiscible("gaussian") # FINISH THE REST
                extra_args = immiscibility.set_immiscible_extra_args(extra_args)
                noise_sampler = lambda _sigma, _sigma_next: immiscibility(x)
                return noise_sampler, extra_args
            case "immiscible_gaussian_maximize":
                immiscibility = make_immiscible("gaussian", maximize=True) # FINISH THE REST
                extra_args = immiscibility.set_immiscible_extra_args(extra_args)
                noise_sampler = lambda _sigma, _sigma_next: immiscibility(x)
                return noise_sampler, extra_args
            case "immiscible_perlin":
                immiscibility = make_immiscible("perlin") # FINISH THE REST
                extra_args = immiscibility.set_immiscible_extra_args(extra_args)
                noise_sampler = lambda _sigma, _sigma_next: immiscibility(x)
                return noise_sampler, extra_args
            case "immiscible_perlin_maximize":
                immiscibility = make_immiscible("perlin", maximize=True) # FINISH THE REST
                extra_args = immiscibility.set_immiscible_extra_args(extra_args)
                noise_sampler = lambda _sigma, _sigma_next: immiscibility(x)
                return noise_sampler, extra_args
    return None, extra_args

class make_immiscible:
    def __init__(self, noise_func="gaussian", immiscible_latents=1024, maximize=False):
        self.noise_func = noise_func
        self.n_latents = immiscible_latents
        self.maximize = maximize
        self.updated_latent = None

    """
    def __call__(self, latents):
        # "Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment" (2024) Li et al. arxiv.org/abs/2406.12303
        # Minimize latent-noise pairs over a batch
        # Code from https://github.com/kohya-ss/sd-scripts/pull/1395
        reference_latent = latents
        if self.updated_latent != None:
            reference_latent = self.updated_latent
        reference_latent = self.batch(reference_latent)
        n = self.n_latents # arg is an integer for how many noise tensors to generate
        noise = None
        match self.noise_func:
            case "gaussian_1024":
                #n = 1024
                size = [n] + list(reference_latent.shape[1:])
                noise = torch.randn(size, dtype=reference_latent.dtype, layout=reference_latent.layout, device=reference_latent.device)
            case "perlin":
                #n = n//32
                size = [n] + list(reference_latent.shape[1:])
                noise = torch.randn(size, dtype=reference_latent.dtype, layout=reference_latent.layout, device=reference_latent.device)
                for i in range(n):
                    for j in range(reference_latent.size(dim=1)):
                        noise_values = rand_perlin_2d_octaves((reference_latent.size(dim=-2), reference_latent.size(dim=-1)), (1,1), 1, 1).to(reference_latent.device)
                        result = (1+0/10)*torch.erfinv(2 * noise_values - 1) * (2 ** 0.5)
                        result = torch.where(torch.abs(result) > 5, noise[i, j, :, :], result)
                        noise[i, j, :, :] = result
        latents_expanded = reference_latent.half().unsqueeze(1).expand(-1, n, *reference_latent.shape[1:])
        noise_expanded = noise.half().unsqueeze(0).expand(reference_latent.shape[0], *noise.shape)
        dist = (latents_expanded - noise_expanded)**2
        dist = dist.mean(list(range(2, dist.dim()))).cpu()
        assign_mat = linear_sum_assignment(dist, maximize=self.maximize)
        noise = noise[assign_mat[1]]
        return self.unbatch(noise, latents)

    def batch(self, ref):
        if self.batching == "batch":
            return ref
        rsz = ref.shape
        if len(rsz) != 4:
            raise ValueError("Reference must be four-dimensional")
        if self.batching == "channel":
            ref = ref.view(rsz[0] * rsz[1], *rsz[2:])
            return ref
        if self.batching == "row":
            ref = ref.view(rsz[0] * rsz[1] * rsz[2], rsz[3])
            return ref
        if self.batching == "column":
            ref = ref.permute(0, 1, 3, 2).reshape(rsz[0] * rsz[1] * rsz[3], rsz[2])
            return ref
        raise ValueError("Bad Immmiscible noise batching type")
    
    def unbatch(self, noise, x_ref):
        xsz = x_ref.shape
        if self.batching == "column":
            return noise.view(*xsz[:2], xsz[3], xsz[2]).permute(0, 1, 3, 2)
        return noise.view(*xsz)

    """

    def __call__(self, latents):
        reference_latent = latents
        if self.updated_latent != None:
            reference_latent = self.updated_latent

        batch_size = latents.shape[0] if self.n_latents is None else self.n_latents
        size = [batch_size] + list(latents.shape[1:])
        #noise = torch.randn_like(latents)  # [B, C, H, W]

        match self.noise_func:
            case "gaussian":
                noise = torch.randn(size, dtype=latents.dtype, layout=latents.layout, device=latents.device)
            case "perlin":
                noise = create_noisy_latents_perlin(torch.randn(size, dtype=latents.dtype, layout=latents.layout, device=latents.device))

        # Distance calculation (simplified for single process)
        distance = torch.linalg.vector_norm(
            0.10 * latents.to(torch.float16).flatten(start_dim=1).unsqueeze(1) -
            0.10 * noise.to(torch.float16).flatten(start_dim=1).unsqueeze(0),
            dim=2
        )  # [B, B]

        # Noise Assignment (simplified for single process)
        _, col_ind = linear_sum_assignment(distance.cpu().numpy(), maximize=self.maximize)
        noise = noise[col_ind].to(latents.device)  # Assign the permuted noise

        return noise

    def set_immiscible_extra_args(self, extra_args):
        def immiscible_post_cfg_function(args):
            self.updated_latent = args["cond_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, immiscible_post_cfg_function, disable_cfg1_optimization=True)
        return extra_args

# From https://github.com/Extraltodeus/noise_latent_perlinpinpin/blob/main/latent_noisy_perlin.py
# which was found at https://gist.github.com/vadimkantorov/ac1b097753f217c5c11bc2ff396e0a57
# which was ported from https://github.com/pvigier/perlin-numpy/blob/master/perlin2d.py
def rand_perlin_2d(shape, res, fade = lambda t: 6*t**5 - 15*t**4 + 10*t**3):
    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])
    
    grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim = -1) % 1
    angles = 2*math.pi*torch.rand(res[0]+1, res[1]+1)
    gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim = -1)
    
    tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], 0).repeat_interleave(d[1], 1)
    dot = lambda grad, shift: (torch.stack((grid[:shape[0],:shape[1],0] + shift[0], grid[:shape[0],:shape[1], 1] + shift[1]  ), dim = -1) * grad[:shape[0], :shape[1]]).sum(dim = -1)
    
    n00 = dot(tile_grads([0, -1], [0, -1]), [0,  0])
    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])
    n01 = dot(tile_grads([0, -1],[1, None]), [0, -1])
    n11 = dot(tile_grads([1, None], [1, None]), [-1,-1])
    t = fade(grid[:shape[0], :shape[1]])
    return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1])

def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5):
    noise = torch.zeros(shape)
    frequency = 1
    amplitude = 1
    for _ in range(octaves):
        noise += amplitude * rand_perlin_2d(shape, (frequency*res[0], frequency*res[1]))
        frequency *= 2
        amplitude *= persistence
    noise = torch.remainder(torch.abs(noise)*1000000,11)/11
    # noise = (torch.sin(torch.remainder(noise*1000000,83))+1)/2
    return noise

def create_noisy_latents_perlin(x, detail_level=0):
    batch_size = x.size(dim=0)
    noise = torch.randn((batch_size, x.size(dim=1), x.size(dim=2), x.size(dim=3)), dtype=x.dtype, layout=x.layout, device=x.device)
    for i in range(batch_size):
        for j in range(x.size(dim=1)):
            noise_values = rand_perlin_2d_octaves((x.size(dim=2), x.size(dim=3)), (1,1), 1, 1).to(x.device)
            result = (1+detail_level/10)*torch.erfinv(2 * noise_values - 1) * (2 ** 0.5)
            result = torch.where(torch.abs(result) > 3, noise[i, j, :, :], result)
            noise[i, j, :, :] = result
    return noise

def rand_perlin_like(x): # Even distribution, seemingly produces more information in non-subject areas than the normal (gaussian) noise sampler
    return create_noisy_latents_perlin(x)

def uniform_noise_sampler(x): # Even distribution, seemingly produces more information in non-subject areas than the normal (gaussian) noise sampler
    return lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73

from torch.distributions import StudentT
def studentt_noise_sampler(x): # Produces more subject-focused outputs due to distribution, unsure if this works
    noise = StudentT(loc=0, scale=0.2, df=1).rsample(x.size())
    #noise *= 2 / (torch.max(torch.abs(noise)) + 1e-8)
    s: FloatTensor = torch.quantile(
        noise.flatten(start_dim=1).abs(),
        0.75,
        dim = -1
    )
    #s.clamp_(min = 1.)
    s = s.reshape(*s.shape, 1, 1, 1)
    noise = noise.clamp(-s, s)
    noise = torch.copysign(torch.pow(torch.abs(noise), 0.5), noise)
    print(s)
    return lambda sigma, sigma_next: noise.to(x.device) / (7/3)

from torch.distributions import Laplace
def rand_laplacian_like(x):
    noise = torch.zeros_like(x)#.div_(4.0)
    noise += Laplace(loc=0, scale=2 ** 0.5).rsample(x.size()).to(noise.device)
    return noise / noise.std()

def highres_pyramid_noise_like(x, discount=0.7):
    b, c, h, w = x.shape # EDIT: w and h get over-written, rename for a different variant!
    orig_h = h
    orig_w = w
    u = torch.nn.Upsample(size=(orig_h, orig_w), mode='bilinear')
    noise = (torch.rand_like(x) - 0.5) * 2 * 1.73 # Start with scaled uniform noise
    for i in range(4):
        r = torch.rand(1).item() * 2 + 2 # Rather than always going 2x,
        h, w = min(orig_h*15, int(h*(r**i))), min(orig_w*15, int(w*(r**i)))
        noise += u(torch.randn(b, c, h, w).to(x)) * discount**i
        if h>=orig_h*15 or w>=orig_w*15: break # Lowest resolution is 1x1
    return noise/noise.std() # Scaled back to roughly unit variance

def green_noise_like(x):
    noise = torch.randn_like(x)
    width = noise.size(dim=2)
    height = noise.size(dim=3)
    scale = 1.0 / (width * height)
    fy = torch.fft.fftfreq(width, device=x.device)[:, None] ** 2
    fx = torch.fft.fftfreq(height, device=x.device) ** 2
    f = fy + fx
    power = torch.sqrt(f)
    power[0, 0] = 1
    noise = torch.fft.ifft2(torch.fft.fft2(noise) / torch.sqrt(power))
    noise *= scale / noise.std()
    noise = torch.real(noise).to(x.device)
    return noise / noise.std()

def green_noise_sampler(x): # This doesn't work properly right now
    width = x.size(dim=2)
    height = x.size(dim=3)
    noise = torch.randn(width, height)
    #scale = 1.0 / (width * height)
    fy = torch.fft.fftfreq(width)[:, None] ** 2
    fx = torch.fft.fftfreq(height) ** 2
    f = fy + fx
    power = torch.sqrt(f)
    power[0, 0] = 1
    noise = torch.fft.ifft2(torch.fft.fft2(noise) / torch.sqrt(power))
    #noise *= scale / noise.std()
    noise = torch.real(noise).to(x.device)
    mean = torch.mean(noise)
    std = torch.std(noise)

    noise.sub_(mean).div_(std)
    print(noise)
    return lambda sigma, sigma_next: noise

# I'm not sure how this differs from the other implementation but it doesn't seem to be used at present.
def power_noise_sampler_2(tensor, alpha=2, k=1): # This doesn't work properly right now
    """Generate 1/f noise for a given tensor.

    Args:
        tensor: The tensor to add noise to.
        alpha: The parameter that determines the slope of the spectrum.
        k: A constant.

    Returns:
        A tensor with the same shape as `tensor` containing 1/f noise.
    """
    tensor = torch.randn_like(tensor)
    fft = torch.fft.fft2(tensor)
    freq = torch.arange(1, len(fft) + 1, dtype=torch.float)
    spectral_density = k / freq**alpha
    noise = torch.rand(tensor.shape) * spectral_density
    mean = torch.mean(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
    std = torch.std(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
    noise = noise.to(tensor.device).sub_(mean).div_(std)
    variance = torch.var(noise, dim=(-2, -1), keepdim=True)
    print(variance)
    return lambda sigma, sigma_next: noise / 3

def pyramid_noise_like(size, dtype, layout, generator, device="cpu", discount=0.8):
    b, c, h, w = size
    orig_h = h
    orig_w = w
    noise = torch.zeros(size=size, dtype=dtype, layout=layout, device=device)
    r = 1
    for i in range(5):
        r *= 2 # Rather than always going 2x,
        #w, h = max(1, int(w/(r**i))), max(1, int(h/(r**i)))
        noise += torch.nn.functional.interpolate((torch.normal(mean=0, std=0.5 ** i, size=(b, c, h * r, w * r), dtype=dtype, layout=layout, generator=generator, device=device)), size=(orig_h, orig_w), mode='nearest-exact') * discount**i
        #if w>=orig_w*16 or h>=orig_h*16: break
    return noise

def power_noise_sampler(size, dtype, layout, generator, device="cpu", alpha=2, k=1): # This doesn't work properly right now
    """Generate 1/f noise for a given tensor.

    Args:
        tensor: The tensor to add noise to.
        alpha: The parameter that determines the slope of the spectrum.
        k: A constant.

    Returns:
        A tensor with the same shape as `tensor` containing 1/f noise.
    """
    tensor = torch.randn(size=size, dtype=dtype, layout=layout, generator=generator, device=device)
    fft = torch.fft.fft2(tensor)
    freq = torch.arange(1, len(fft) + 1, dtype=torch.float)
    spectral_density = k / freq**alpha
    noise = torch.rand(size=size, dtype=dtype, layout=layout, generator=generator, device=device) * spectral_density
    mean = torch.mean(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
    std = torch.std(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
    noise = noise.to(tensor.device).sub_(mean).div_(std)
    return noise

def prepare_noise(latent_image, seed, noise_type, noise_inds=None): # From `sample.py`
    """
    creates random noise given a latent image and a seed.
    optional arg skip can be used to skip and discard x number of noise generations for a given seed
    """
    generator = torch.manual_seed(seed)
    match noise_type:
        case "gaussian":
            noise_func = torch.randn
        case "uniform":
            def uniform_rand(*size, **kwargs):
                return (torch.rand(*size, **kwargs) - 0.5) * 2 * 1.73
            noise_func = uniform_rand
        case "pyramid":
            noise_func = pyramid_noise_like
        case "power":
            noise_func = power_noise_sampler
        case _:
            noise_func = torch.randn
    if noise_inds is None:
        return noise_func(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")

    unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
    noises = []
    for i in range(unique_inds[-1]+1):
        noise = noise_func([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
        if i in unique_inds:
            noises.append(noise)
    noises = [noises[i] for i in inverse]
    noises = torch.cat(noises, axis=0)
    return noises

NOISE_SAMPLER_HANDLERS={
    # Brownian is special-cased.
    "gaussian": torch.randn_like,
    "highres-pyramid": highres_pyramid_noise_like,
    "pyramid": lambda x: pyramid_noise_like(x.size(), x.dtype, x.layout, None, device=x.device),
    "perlin": rand_perlin_like,
    "laplacian": rand_laplacian_like,
    "uniform": uniform_noise_like,
}


# Below this point are extra samplers
@torch.no_grad()
def sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1., noise_sampler=None, momentum=0.0):
    """DPM-Solver++(3M) SDE, modified with an extra SDE, and momentumized in both the SDE and ODE(?). 'its a first' - Clybius 2023
    The expression for d1 is derived from the extrapolation formula given in the paper “Diffusion Monte Carlo with stochastic Hamiltonians” by M. Foulkes, L. Mitas, R. Needs, and G. Rajagopal. The formula is given as follows:
    d1 = d1_0 + (d1_0 - d1_1) * r2 / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r2 + r1) * (r0 + r1))
    (if this is an incorrect citing, we blame Google's Bard and OpenAI's ChatGPT for this and NOT me :^) )

    where d1_0, d1_1, and d1_2 are defined as follows:
    d1_0 = (denoised - denoised_1) / r2
    d1_1 = (denoised_1 - denoised_2) / r1
    d1_2 = (denoised_2 - denoised_3) / r0

    The variables r0, r1, and r2 are defined as follows:
    r0 = h_3 / h_2
    r1 = h_2 / h
    r2 = h / h_1
    """

    def momentum_func(diff, velocity, timescale=1.0, offset=-momentum / 2.0): # Diff is current diff, vel is previous diff
        if velocity is None:
            momentum_vel = diff
        else:
            momentum_vel = momentum * (timescale + offset) * velocity + (1 - momentum * (timescale + offset)) * diff
        return momentum_vel

    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
    noise_sampler = rand_perlin_like(x) if noise_sampler is None else noise_sampler

    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    denoised_1, denoised_2, denoised_3 = None, None, None
    h_1, h_2, h_3 = None, None, None
    vel, vel_sde = None, None
    for i in trange(len(sigmas) - 1, disable=disable):
        time = sigmas[i] / sigma_max
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        if sigmas[i + 1] == 0:
            # Denoising step
            x = denoised
        else:
            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
            h = s - t
            h_eta = h * (eta + 1)
            x_diff = momentum_func((-h_eta).expm1().neg() * denoised, vel, time)
            vel = x_diff
            x = torch.exp(-h_eta) * x + vel

            if h_3 is not None:
                r0 = h_3 / h_2
                r1 = h_2 / h
                r2 = h / h_1
                d1_0 = (denoised - denoised_1) / r2
                d1_1 = (denoised_1 - denoised_2) / r1
                d1_2 = (denoised_2 - denoised_3) / r0
                d1 = d1_0 + (d1_0 - d1_1) * r2 / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r2 + r1) * (r0 + r1))
                d2 = (d1_0 - d1_1) / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) / ((r2 + r1) * (r0 + r1))
                phi_3 = h_eta.neg().expm1() / h_eta + 1
                phi_4 = phi_3 / h_eta - 0.5
                sde_diff = momentum_func(phi_3 * d1 - phi_4 * d2, vel_sde, time)
                vel_sde = sde_diff
                x = x + vel_sde
            elif h_2 is not None:
                r0 = h_1 / h
                r1 = h_2 / h
                d1_0 = (denoised - denoised_1) / r0
                d1_1 = (denoised_1 - denoised_2) / r1
                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
                d2 = (d1_0 - d1_1) / (r0 + r1)
                phi_2 = h_eta.neg().expm1() / h_eta + 1
                phi_3 = phi_2 / h_eta - 0.5
                sde_diff = momentum_func(phi_2 * d1 - phi_3 * d2, vel_sde, time)
                vel_sde = sde_diff
                x = x + vel_sde
            elif h_1 is not None:
                r = h_1 / h
                d = (denoised - denoised_1) / r
                phi_2 = h_eta.neg().expm1() / h_eta + 1
                sde_diff = momentum_func(phi_2 * d, vel_sde, time)
                vel_sde = sde_diff
                x = x + vel_sde

            if eta:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise

            denoised_1, denoised_2, denoised_3 = denoised, denoised_1, denoised_2
            h_1, h_2, h_3 = h, h_1, h_2
    return x

# Kat's Truncated Taylor Method sampler, by Katherine Crowson
def sample_ttm_jvp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
    """Second order truncated Taylor method (torch.func.jvp() version)."""

    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])
    model_fn = lambda x, sigma: model(x, sigma * s_in, **extra_args)

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model_fn(x, sigmas[i])

        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})

        if sigmas[i + 1] == 0:
            # Denoising step
            x = denoised
        else:
            # 2nd order truncated Taylor method
            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
            h = s - t
            h_eta = h * (eta + 1)

            eps = to_d(x, sigmas[i], denoised)
            _, denoised_prime = torch.func.jvp(model_fn, (x, sigmas[i]), (eps * -sigmas[i], -sigmas[i]))

            phi_1 = -torch.expm1(-h_eta)
            #phi_2 = torch.expm1(-h_eta) + h_eta
            phi_2 = torch.expm1(-h) + h # seems to work better with eta > 0
            x = torch.exp(-h_eta) * x + phi_1 * denoised + phi_2 * denoised_prime

            if eta:
                phi_1_noise = torch.sqrt(-torch.expm1(-2 * h * eta))
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * phi_1_noise * s_noise

    return x

# Many thanks to Kat + Birch-San for this wonderful sampler implementation! https://github.com/Birch-san/sdxl-play/commits/res/
from .other_samplers.refined_exp_solver import sample_refined_exp_s
def sample_res_solver(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler_type="gaussian", noise_sampler=None, denoise_to_zero=True, simple_phi_calc=False, c2=0.5, ita=torch.Tensor((0.25,)), momentum=0.0):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sample_refined_exp_s(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), denoise_to_zero=denoise_to_zero, simple_phi_calc=simple_phi_calc, c2=c2, ita=ita, momentum=momentum)

@torch.no_grad()
def sample_dpmpp_dualsde_momentum(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1/2, momentum=0.0):
    """DPM-Solver++ (Stochastic with Momentum). Personal modified sampler by Clybius"""
    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
    noise_sampler = rand_perlin_like(x) if noise_sampler is None else noise_sampler
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    sigma_fn = lambda t: t.neg().exp()
    t_fn = lambda sigma: sigma.log().neg()

    denoisedsde_1, denoisedsde_2, denoisedsde_3 = None, None, None # new line
    h_1, h_2, h_3 = None, None, None # new line

    def momentum_func(diff, velocity, timescale=1.0, offset=-momentum / 2.0): # Diff is current diff, vel is previous diff
        if velocity is None:
            momentum_vel = diff
        else:
            momentum_vel = momentum * (timescale + offset) * velocity + (1 - momentum * (timescale + offset)) * diff
        return momentum_vel

    vel = None
    vel_2 = None
    vel_sde = None
    for i in trange(len(sigmas) - 1, disable=disable):
        time = sigmas[i] / sigma_max
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        if sigmas[i + 1] == 0:
            # Euler method
            d = to_d(x, sigmas[i], denoised)
            dt = sigmas[i + 1] - sigmas[i]
            x = x + d * dt
        else:
            # DPM-Solver++
            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
            h = t_next - t
            h_eta = h * (eta + 1)
            s = t + h * r
            fac = 1 / (2 * r)

            # Step 1
            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(s), eta)
            s_ = t_fn(sd)
            diff_2 = momentum_func((t - s_).expm1() * denoised, vel_2, time)
            vel_2 = diff_2
            x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - diff_2
            x_2 = x_2 + noise_sampler(sigma_fn(t), sigma_fn(s)) * s_noise * su
            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)

            # Step 2
            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(t_next), eta)
            t_next_ = t_fn(sd)
            denoised_d = (1 - fac) * denoised + fac * denoised_2
            diff = momentum_func((t - t_next_).expm1() * denoised_d, vel, time)
            vel = diff
            x = (sigma_fn(t_next_) / sigma_fn(t)) * x - diff

            if h_3 is not None:
                r0 = h_3 / h_2
                r1 = h_2 / h
                r2 = h / h_1
                d1_0 = (denoised_d - denoisedsde_1) / r2
                d1_1 = (denoisedsde_1 - denoisedsde_2) / r1
                d1_2 = (denoisedsde_2 - denoisedsde_3) / r0
                d1 = d1_0 + (d1_0 - d1_1) * r2 / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r2 + r1) * (r0 + r1))
                d2 = (d1_0 - d1_1) / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) / ((r2 + r1) * (r0 + r1))
                phi_3 = h_eta.neg().expm1() / h_eta + 1
                phi_4 = phi_3 / h_eta - 0.5
                diff = momentum_func(phi_3 * d1 - phi_4 * d2, vel_sde, time)
                vel_sde = diff
                x = x + diff
            elif h_2 is not None:
                r0 = h_1 / h
                r1 = h_2 / h
                d1_0 = (denoised_d - denoisedsde_1) / r0
                d1_1 = (denoisedsde_1 - denoisedsde_2) / r1
                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
                d2 = (d1_0 - d1_1) / (r0 + r1)
                phi_2 = h_eta.neg().expm1() / h_eta + 1
                phi_3 = phi_2 / h_eta - 0.5
                diff = momentum_func(phi_2 * d1 - phi_3 * d2, vel_sde, time)
                vel_sde = diff
                x = x + diff
            elif h_1 is not None:
                r = h_1 / h
                d = (denoised_d - denoisedsde_1) / r
                phi_2 = h_eta.neg().expm1() / h_eta + 1
                diff = momentum_func(phi_2 * d, vel_sde, time)
                vel_sde = diff
                x = x + diff

            if eta:
                x = x + noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * su
            #if 'denoised_d' in locals():
            denoisedsde_1, denoisedsde_2, denoisedsde_3 = denoised_d, denoisedsde_1, denoisedsde_2 # new line
            #if 'h' in locals():
            h_1, h_2, h_3 = h, h_1, h_2
    return x

def sample_dpmpp_dualsdemomentum(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, r=1/2, momentum=0.0):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sample_dpmpp_dualsde_momentum(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), r=r, momentum=momentum)

from .other_samplers.sample_ttm import sample_ttm_jvp
def sample_ttmcustom(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian",noise_sampler=None):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sample_ttm_jvp(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args))

from comfy.k_diffusion.sampling import sample_lcm
def sample_lcmcustom(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler_type="gaussian", noise_sampler=None):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sample_lcm(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args))

def sample_clyb_4m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="brownian", noise_sampler=None, momentum=0.0):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), momentum=momentum)


# This code works, but I'm currently experimenting with different methods
@torch.no_grad()
def sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, leap=2, eta_dance=1.0):
    #Ancestral sampling with Euler method steps, dancing steps.
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    unsample_noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
    s_in = x.new_ones([x.shape[0]])
    for i in trange(len(sigmas) - 1, disable=disable):
        if i < len(sigmas) - leap:
            is_danceable = sigmas[i + leap] > 0
        else:
            is_danceable = False
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + leap] if is_danceable else sigmas[i + 1], eta=eta)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        d = to_d(x, sigmas[i], denoised)
        # Euler method
        dt = sigma_down - sigmas[i]
        x = x + d * dt
        if sigmas[i + 1] > 0:
            if is_danceable:
                x = x + noise_sampler(sigmas[i], sigmas[i + leap]) * s_noise * sigma_up
                #x = x + noise_sampler(sigmas[i + 2], sigmas[i + 1]) * s_noise * sigma_up
                #denoised2 = model(x, sigmas[i + 2] * s_in, **extra_args)
                sigma_down2, sigma_up2 = get_ancestral_step(sigmas[i + leap], sigmas[i + 1], eta=eta_dance)
                d_2 = to_d(x, sigmas[i + leap], denoised)
                dt_2 = sigma_down2 - sigmas[i + leap]
                x = x + d_2 * dt_2
                x = x + noise_sampler(sigmas[i + leap], sigmas[i + 1]) * s_noise * sigma_up2

                #sigma_down3, sigma_up3 = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
                #x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up3

                #denoised2 = model(x, sigmas[i] * s_in, **extra_args)
                #d_3 = to_d(x, sigmas[i], denoised2)
                #dt_3 = sigma_down3 - sigmas[i]
                #x = x + d_3 * dt_3 + d_2 * dt_2
                #print(dt_3, dt_2)
                #x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up3
                #x = x + d * dt
            else:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up

    return x

def sample_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, leap=2, eta_dance=1.0):
    return sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), leap=leap, eta_dance=eta_dance)


@torch.no_grad()
def sampler_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=None, callback=None, disable=None, eta_max=1.0, eta_min=0.0, s_noise=1., noise_sampler=None):
    """DPM-Solver++(3M) SDE with dynamic eta."""
    def eta_schedule_cosine_annealing(i, n, eta_max=eta_max, eta_min=eta_min):
        """Cosine annealing schedule for eta."""
        progress = i / (n - 1)
        eta = eta_min + 0.5 * (eta_max - eta_min) * (1 + math.cos(math.pi * progress))
        return eta

    seed = extra_args.get("seed", None)
    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    denoised_1, denoised_2 = None, None
    h, h_1, h_2 = None, None, None

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        if sigmas[i + 1] == 0:
            # Denoising step
            x = denoised
        else:
            # DPM-Solver++(3M) SDE
            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
            h = s - t

            # Dynamic eta
            eta = eta_schedule_cosine_annealing(i, len(sigmas))
            h_eta = h * (eta + 1)

            x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised

            if h_2 is not None:
                r0 = h_1 / h
                r1 = h_2 / h
                d1_0 = (denoised - denoised_1) / r0
                d1_1 = (denoised_1 - denoised_2) / r1
                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
                d2 = (d1_0 - d1_1) / (r0 + r1)
                phi_2 = h_eta.neg().expm1() / h_eta + 1
                phi_3 = phi_2 / h_eta - 0.5
                x = x + phi_2 * d1 - phi_3 * d2
            elif h_1 is not None:
                r = h_1 / h
                d = (denoised - denoised_1) / r
                phi_2 = h_eta.neg().expm1() / h_eta + 1
                x = x + phi_2 * d

            if eta:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise

        denoised_1, denoised_2 = denoised, denoised_1
        h_1, h_2 = h, h_1
    return x

def sample_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=None, callback=None, disable=None, eta_max=1.0, eta_min=0.0, s_noise=1., noise_sampler_type="brownian", noise_sampler=None):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta_max=eta_max, eta_min=eta_min, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args))


from .other_samplers.refined_exp_solver import _de_second_order

# Default is 2, so only methods with other values are included here.
SUPREME_ORDER = { "euler": 1, "dpm_1s": 1, "dpm_3s": 3, "rk4": 4, "reversible_heun_1s": 1, "rkf45": 6, "bogacki_shampine": 3, }

@torch.no_grad()
def sampler_supreme(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None, eta=1.0, step_method="euler", substep_method="euler", warmup_method="euler", centralization=0.00, normalization=0.00, edge_enhancement=0.00, perphist=0.25, substeps=2, noise_modulation="none", modulation_strength=2., modulation_dims=3, reversible_eta=1.0, dyneta=True, reversible_dyneta=True, enable_free_reverse=True, free_reverse_eta=0.0, free_reverse_dyneta=True):
    """
    Supreme Sampler, Euler steps. Based on no paper, purely interesting thoughts.

    Args:
        model: Denoising model call.
        x: The initial noisy sample.
        sigmas: The noise schedule.
        extra_args: Additional arguments for the model.
        callback: A callback function for monitoring the sampling process.
        disable: Whether to disable the progress bar.
        s_noise: The noise scale factor.
        noise_sampler: A custom noise sampler function.
        eta: Ancestral-ness.
        centralization: Subtracts mean from the denoised latent.
        normalization: Divides the denoised latent by the standard deviation.
        edge_enhancement: Multiplies the edges by the mean using a laplacian kernel
        perphist: Adds previous denoised variable to the current denoised using perpendicular vector projection
        substeps: Amount of times to iterate over each step and average the results
        noise_modulation: Method of changing the noise based on situations within the sampler
        modulation_strength: Strength of the modulation using a weighted sum between the modulation and noise sampler's noise.
        modulation_dims: Choose between (channel) modulation, (height, width) modulation, or (channels, height, width) modulation
        reversible_eta: Ancestralness in the reversible component of reversible samplers.
        dyneta: Enable a dynamic eta based on sigma. Higher sigmas have a lower eta, while lower sigmas have a higher eta, max clamped to user-chosen eta.
        reversible_dyneta: Enable a dynamic reversible eta based on sigma. Higher sigmas have a lower eta, while lower sigmas have a higher eta, max clamped to user-chosen eta. Good for stability.
    """

    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])

    # Centralization
    def centralize(denoised_sample, centralization, iteration):
        for b in range(len(denoised_sample)):
            for c in range(len(denoised_sample[b])):
                channel = denoised_sample[b][c]
                denoised_sample[b][c] -= channel.mean() * centralization * (sigmas[iteration] ** 0.5)
        return denoised_sample

    # Normalization
    def normalize(denoised_sample, normalization, iteration):
        for b in range(len(denoised_sample)):
            for c in range(len(denoised_sample[b])):
                channel = denoised_sample[b][c]
                denoised_sample[b][c] += ((denoised_sample[b][c] / channel.std()) - denoised_sample[b][c]) * normalization * (sigmas[iteration] ** 0.5)
        return denoised_sample

    # Perp-hist
    def perpadd(denoised_tensor, old_denoised_tensor, x, alpha):
        a_diff = x - (denoised_tensor - x)
        b_diff = x - (old_denoised_tensor - x)
        a_ortho = a_diff * (a_diff / torch.linalg.norm(a_diff) * (b_diff / torch.linalg.norm(a_diff))).sum()
        b_perp = b_diff - a_ortho
        res = denoised_tensor + alpha * b_perp
        return res

    # DynETA
    orig_eta = eta
    orig_reversible_eta = reversible_eta
    orig_free_reverse_eta = free_reverse_eta
    def dyneta_fn(original_eta, sigma, sigma_next):
        return torch.clamp(1 / (sigma**2 - sigma_next**2)**0.5, min=0.0, max=original_eta)

    order, sub_order = SUPREME_ORDER.get(step_method, 2), SUPREME_ORDER.get(substep_method, 2)
    steps_per_sigma = order + sub_order * (substeps - 1)

    def apply_enhancements(x, i, model, sigma_s_in, old_denoised):
        args = extra_args
        denoised = model(x, sigma_s_in, **args)

        if edge_enhancement != 0:
            blur = (kornia.filters.joint_bilateral_blur(x, denoised, (3, 3), 0.1, (1.5, 1.5)) - x) # Blurs non-edges
            denoised += (kornia.filters.unsharp_mask(denoised, (3, 3), (1.5, 1.5)) - denoised) * (sigmas[i] - sigmas[i + 1]) * edge_enhancement # Sharpens everything
            denoised += blur * (sigmas[i] - sigmas[i + 1]) * edge_enhancement # Apply blur to non-edges, thus leaving edges sharpened

        if centralization != 0:
            denoised = centralize(denoised, centralization, i)

        if normalization != 0:
            denoised = normalize(denoised, normalization, i)

        if old_denoised != None and perphist != 0:
            denoised = perpadd(denoised, old_denoised, x, perphist)

        return denoised

    # Dynamic sampling
    dynamic_order_samplers = {
        1: "euler",
        2: "trapezoidal",
        3: "bogacki_shampine",
        4: "rk4",
        6: "rkf45",
    }
    # Adaptive RK order sampling
    adaptive_rk_weights = {
        1: [1],
        2: [0.5, 0.5],
        3: [1/6, 2/3, 1/6],
        4: [1/8, 3/8, 3/8, 1/8],
    }

    def dynamic_step_method(step_method, model, prev_x, denoised, prev_denoised, iteration, substep_iter):
        """
        Step method function, applies cond-error modification, and dynamic step selection if chosen.
        """
        sampler = step_method
        order = 1
        error = 0
        if iteration == 0 or prev_denoised == None: # Warmup with the chosen warmup step, else use substep method for substeps
            if warmup_method == "none":
                return step_method
            if substep_iter > 0:
                return substep_method, 1, error
            order = 2 # Chosen for simplicity
            return warmup_method, order, error

        d = to_d(prev_x, sigmas[iteration - 1], prev_denoised)
        x_pred = prev_x + d * (sigmas[iteration] - sigmas[iteration - 1])

        d_pred = to_d(x_pred, sigmas[iteration], denoised)

        error = torch.linalg.norm(d_pred - d) / torch.linalg.norm(d)

        if substep_iter > 0:
            return substep_method, 1, error
        if step_method != "dynamic" and step_method != "adaptive_rk": # If we're not a dynamic sampler, return the step unmodified step method
            return step_method, order, error

        if (error < 1e-2):
            order = 6
        elif (error < 3.75e-2):
            order = 4
        elif (error < 7.5e-2):
            order = 3
        elif (error < 1.5e-1):
            order = 2
        else:
            order = 1

        if step_method == "adaptive_rk":
            return step_method, min(order, 4), error
        return dynamic_order_samplers[order], order, error

    renoise_weights = torch.ones(substeps, device=x.device) / substeps
    def intensity_based_multiplicative_noise_fn(x, noise, s_noise, sigma_up, intensity, dims):
        """
        Scales noise based on the intensities of the input tensor.
        """
        std = torch.std(x - x.mean(), dim=dims, keepdim=True)  # Average across channels to get intensity
        scaling = (1 / (std * abs(intensity) + 1.0)) # Scale std by intensity, as not doing this leads to more noise being left over, leading to crusty/preceivably extremely oversharpened images
        additive_noise = noise * s_noise * sigma_up
        scaled_noise = noise * s_noise * sigma_up * scaling + additive_noise

        noise_norm = torch.norm(additive_noise)
        scaled_noise_norm = torch.norm(scaled_noise)
        scaled_noise *= noise_norm / scaled_noise_norm # Scale to normal noise strength
        scaled_noise = scaled_noise * intensity + additive_noise * (1 - intensity)
        return scaled_noise

    def frequency_based_noise(z_k, noise, s_noise, sigma_up, intensity, channels):
        """
        Scales the high-frequency components of the noise based on the given intensity.
        """
        additive_noise = noise * s_noise * sigma_up

        std = torch.std(z_k - z_k.mean(), dim=channels, keepdim=True)  # Average across channels to get intensity
        scaling = (1 / (std * abs(intensity) + 1.0))
        # Perform Fast Fourier Transform (FFT)
        z_k_freq = torch.fft.fft2(scaling * additive_noise + additive_noise)

        # Get the magnitudes of the frequency components
        magnitudes = torch.abs(z_k_freq)

        # Create a high-pass filter (emphasize high frequencies)
        h, w = z_k.shape[-2:]
        b = abs(intensity)  # Controls the emphasis of the high pass (higher frequencies are boosted)
        high_pass_filter = 1 - torch.exp(-((torch.arange(h)[:, None] / h)**2 + (torch.arange(w)[None, :] / w)**2) * b**2)
        high_pass_filter = high_pass_filter.to(z_k.device)

        # Apply the filter to the magnitudes
        magnitudes_scaled = magnitudes * (1 + high_pass_filter)

        # Reconstruct the complex tensor with scaled magnitudes
        z_k_freq_scaled = magnitudes_scaled * torch.exp(1j * torch.angle(z_k_freq))

        # Perform Inverse Fast Fourier Transform (IFFT)
        z_k_scaled = torch.fft.ifft2(z_k_freq_scaled)

        # Return the real part of the result
        z_k_scaled = torch.real(z_k_scaled)

        noise_norm = torch.norm(additive_noise)
        scaled_noise_norm = torch.norm(z_k_scaled)

        z_k_scaled *= (noise_norm / scaled_noise_norm) # Scale to normal noise strength

        scaled_noise = z_k_scaled * intensity + additive_noise * (1 - intensity)

        return scaled_noise
    
    def spectral_modulate_noise(z_k, noise, s_noise, sigma_up, intensity, channels, spectral_mod_percentile=5.0): # Modified for soft quantile adjustment using a novel:tm::c::r: method titled linalg.
        additive_noise = noise * s_noise * sigma_up
        # Convert image to Fourier domain
        fourier = torch.fft.fftn(additive_noise, dim=channels)  # Apply FFT along Height and Width dimensions
    
        log_amp = torch.log(torch.sqrt(fourier.real ** 2 + fourier.imag ** 2))

        quantile_low = torch.quantile(
            log_amp.abs().flatten(1),
            spectral_mod_percentile * 0.01,
            dim = 1
        ).unsqueeze(-1).unsqueeze(-1).expand(log_amp.shape)
        
        quantile_high = torch.quantile(
            log_amp.abs().flatten(1),
            1 - (spectral_mod_percentile * 0.01),
            dim = 1
        ).unsqueeze(-1).unsqueeze(-1).expand(log_amp.shape)

        quantile_max = torch.quantile(
            log_amp.abs().flatten(1),
            1,
            dim = 1
        ).unsqueeze(-1).unsqueeze(-1).expand(log_amp.shape)

        # Decrease high-frequency components
        mask_high = log_amp > quantile_high # If we're larger than 95th percentile

        additive_mult_high = torch.where(
            mask_high,
            1 - ((log_amp - quantile_high) / (quantile_max - quantile_high)).clamp_(max=0.5), # (1) - (0-1), where 0 is 95th %ile and 1 is 100%ile
            torch.tensor(1.0)
        )
        

        # Increase low-frequency components
        mask_low = log_amp < quantile_low
        additive_mult_low = torch.where(
            mask_low,
            1 + (1 - (log_amp / quantile_low)).clamp_(max=0.5), # (1) + (0-1), where 0 is 5th %ile and 1 is 0%ile
            torch.tensor(1.0)
        )
        
        mask_mult = ((additive_mult_low * additive_mult_high) ** intensity)
        #print(mask_mult)
        filtered_fourier = fourier * mask_mult
        
        # Inverse transform back to spatial domain
        inverse_transformed = torch.fft.ifftn(filtered_fourier, dim=channels)  # Apply IFFT along Height and Width dimensions
        
        scaled_noise = inverse_transformed.real.to(additive_noise.device)

        #noise_norm = torch.norm(additive_noise)
        #scaled_noise_norm = torch.norm(scaled_noise)

        return scaled_noise# * (noise_norm / scaled_noise_norm)

    dims = (-3, -2, -1)
    match modulation_dims:
        case 1:
            dims = (-3)
        case 2:
            dims = (-2, -1)
        case 3:
            dims = (-3, -2, -1)

    orig_model = model
    old_denoised = None
    prev_denoised = None
    prev_x = x
    for i in trange(len(sigmas) - 1, disable=disable):
        def model(x, sigma_s_in, **extra_args): # Model wrapper to apply enhancements at every call
            nonlocal old_denoised
            denoised = apply_enhancements(x, i, orig_model, sigma_s_in, old_denoised)
            old_denoised = denoised
            if callback is not None:
                callback({'x': z_k, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
            return denoised

        dpm_solver = DPMSolver(model, extra_args)

        # DynETA
        if dyneta: eta = dyneta_fn(orig_eta, sigmas[i], sigmas[i + 1])
        if reversible_dyneta: reversible_eta = dyneta_fn(orig_reversible_eta, sigmas[i], sigmas[i + 1])

        # Renoising iterations
        z_avg = torch.zeros_like(x)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
        sigma_down_reversible, _ = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=reversible_eta)
        for k in range(substeps):
            z_k = x
            orig_zk = z_k
            eps_cache = {}

            denoised = model(z_k, sigmas[i] * s_in, **extra_args)


            eps = (z_k - denoised) / sigmas[i]
            eps_cache = {'eps': eps}

            step_method_dyn, order, error = dynamic_step_method(step_method, model, prev_x, denoised, prev_denoised, i, k) #step_method, model, prev_x, denoised, prev_denoised, i, k

            match step_method_dyn if sigmas[i + 1] != 0 else "euler":
                case "euler": # 1 model call
                    d = to_d(z_k, sigmas[i], denoised)
                    dt = sigma_down - sigmas[i]

                    z_k = z_k + d * dt
                case "dpm_1s": # DPM Family, 1 model call
                    if callback is not None:
                        dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
                    z_k, eps_cache = dpm_solver.dpm_solver_1_step(z_k, dpm_solver.t(sigmas[i]), dpm_solver.t(sigma_down), eps_cache=eps_cache)
                case "dpm_2s": # 2 model calls
                    if callback is not None:
                        dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
                    z_k, eps_cache = dpm_solver.dpm_solver_2_step(z_k, dpm_solver.t(sigmas[i]), dpm_solver.t(sigma_down), eps_cache=eps_cache)
                case "dpm_3s": # 3 model calls
                    if callback is not None:
                        dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
                    z_k, eps_cache = dpm_solver.dpm_solver_3_step(z_k, dpm_solver.t(sigmas[i]), dpm_solver.t(sigma_down), eps_cache=eps_cache)
                case "rk4": # Fourth-order Runge-Kutta method, 4 model calls
                    # Calculate the derivative using the model
                    d = to_d(z_k, sigmas[i], denoised)
                    dt = sigma_down - sigmas[i]

                    # Runge-Kutta steps
                    k1 = d * dt
                    k2 = to_d(z_k + k1 / 2, sigmas[i] + dt / 2, model(z_k + k1 / 2, (sigmas[i] + dt / 2) * s_in, **extra_args)) * dt
                    k3 = to_d(z_k + k2 / 2, sigmas[i] + dt / 2, model(z_k + k2 / 2, (sigmas[i] + dt / 2) * s_in, **extra_args)) * dt
                    k4 = to_d(z_k + k3, sigmas[i] + dt, model(z_k + k3, (sigmas[i] + dt) * s_in, **extra_args)) * dt

                    # Update the sample
                    z_k = z_k + (k1 + 2 * k2 + 2 * k3 + k4) / 6
                case "reversible_heun": # 2 model calls
                    sigma_i, sigma_i_plus_1 = sigmas[i], sigma_down
                    dt = sigma_i_plus_1 - sigma_i
                    dt_reversible = sigma_down_reversible - sigma_i

                    # Calculate the derivative using the model
                    d_i = to_d(z_k, sigma_i, denoised)

                    # Predict the sample at the next sigma using Euler step
                    x_pred = z_k + d_i * dt

                    # Denoised sample at the next sigma
                    denoised_i_plus_1 = model(x_pred, sigma_i_plus_1 * s_in, **extra_args)

                    # Calculate the derivative at the next sigma
                    d_i_plus_1 = to_d(x_pred, sigma_i_plus_1, denoised_i_plus_1)

                    # Update the sample using the Reversible Heun formula
                    z_k = z_k + dt * (d_i + d_i_plus_1) / 2 - dt_reversible**2 * (d_i_plus_1 - d_i) / 4
                case "reversible_heun_1s": # Experimental 1 model call variant, utilizing previous denoised variables to speed up diffusion.
                    # Reversible Heun-inspired update (first-order)
                    sigma_i, sigma_i_plus_1 = sigmas[i], sigma_down
                    dt = sigma_i_plus_1 - sigma_i
                    dt_reversible = sigma_down_reversible - sigma_i

                    # Calculate the derivative using the model
                    d_i_old = to_d(z_k, sigma_i, prev_denoised) if prev_denoised is not None else to_d(z_k, sigma_i, model(z_k, sigma_i * s_in, **extra_args))

                    # Predict the sample at the next sigma using Euler step
                    x_pred = z_k + d_i_old * dt

                    # Calculate the derivative at the next sigma
                    d_i_plus_1 = to_d(x_pred, sigma_i_plus_1, denoised)

                    # Update the sample using the Reversible Heun formula
                    z_k = z_k + dt * (d_i_old + d_i_plus_1) / 2 - dt_reversible**2 * (d_i_plus_1 - d_i_old) / 4
                case "rkf45": # 6 model calls (expensive)
                    sigma_i, sigma_i_plus_1 = sigmas[i], sigma_down
                    dt = sigma_i_plus_1 - sigma_i
                    # Calculate the derivative using the model
                    d_i = to_d(z_k, sigmas[i], denoised)
                    # RKF45 steps
                    k1 = d_i * dt
                    k2 = to_d(z_k + k1 / 4, sigmas[i] + dt / 4, model(z_k + k1 / 4, (sigmas[i] + dt / 4) * s_in, **extra_args)) * dt
                    k3 = to_d(z_k + 3 * k1 / 32 + 9 * k2 / 32, sigmas[i] + 3 * dt / 8, model(z_k + 3 * k1 / 32 + 9 * k2 / 32, (sigmas[i] + 3 * dt / 8) * s_in, **extra_args)) * dt
                    k4 = to_d(z_k + 1932 * k1 / 2197 - 7200 * k2 / 2197 + 7296 * k3 / 2197, sigmas[i] + 12 * dt / 13, model(z_k + 1932 * k1 / 2197 - 7200 * k2 / 2197 + 7296 * k3 / 2197, (sigmas[i] + 12 * dt / 13) * s_in, **extra_args)) * dt
                    k5 = to_d(z_k + 439 * k1 / 216 - 8 * k2 + 3680 * k3 / 513 - 845 * k4 / 4104, sigmas[i] + dt, model(z_k + 439 * k1 / 216 - 8 * k2 + 3680 * k3 / 513 - 845 * k4 / 4104, (sigmas[i] + dt) * s_in, **extra_args)) * dt

                    # Update the sample
                    z_k = z_k + 25 * k1 / 216 + 1408 * k3 / 2565 + 2197 * k4 / 4104 - k5 / 5
                case "adaptive_rk":
                    sigma_i, sigma_i_plus_1 = sigmas[i], sigma_down
                    dt = sigma_i_plus_1 - sigma_i

                    # Calculate the derivative using the model
                    d_i = to_d(z_k, sigma_i, denoised)

                    # Adaptive order Runge-Kutta steps
                    k_values = [d_i * dt]  # Initialize with k1
                    for j in range(1, order):
                        # Calculate intermediate k values based on the current order
                        k_sum = sum(adaptive_rk_weights[order][l] * k_values[l] for l in range(j))
                        k_values.append(to_d(z_k + k_sum, sigma_i + dt * sum(adaptive_rk_weights[order][:j]), model(z_k + k_sum, (sigma_i + dt * sum(adaptive_rk_weights[order][:j])) * s_in, **extra_args)) * dt)

                    # Update the sample using the weighted sum of k values
                    z_k = z_k + sum(adaptive_rk_weights[order][j] * k_values[j] for j in range(order))
                case "bogacki_shampine":
                    sigma_i, sigma_i_plus_1 = sigmas[i], sigma_down
                    dt = sigma_i_plus_1 - sigma_i

                    # Calculate the derivative using the model
                    d_i = to_d(z_k, sigma_i, denoised)

                    # Bogacki-Shampine steps
                    k1 = d_i * dt
                    k2 = to_d(z_k + k1 / 2, sigma_i + dt / 2, model(z_k + k1 / 2, (sigma_i + dt / 2) * s_in, **extra_args)) * dt
                    k3 = to_d(z_k + 3 * k1 / 4 + k2 / 4, sigma_i + 3 * dt / 4, model(z_k + 3 * k1 / 4 + k2 / 4, (sigma_i + 3 * dt / 4) * s_in, **extra_args)) * dt

                    # Update the sample
                    z_k = z_k + 2 * k1 / 9 + k2 / 3 + 4 * k3 / 9
                case "reversible_bogacki_shampine":
                    sigma_i, sigma_i_plus_1 = sigmas[i], sigma_down
                    dt = sigma_i_plus_1 - sigma_i
                    dt_reversible = sigma_down_reversible - sigma_i

                    # Calculate the derivative using the model
                    d_i = to_d(z_k, sigma_i, denoised)

                    # Bogacki-Shampine steps
                    k1 = d_i * dt
                    k2 = to_d(z_k + k1 / 2, sigma_i + dt / 2, model(z_k + k1 / 2, (sigma_i + dt / 2) * s_in, **extra_args)) * dt
                    k3 = to_d(z_k + 3 * k1 / 4 + k2 / 4, sigma_i + 3 * dt / 4, model(z_k + 3 * k1 / 4 + k2 / 4, (sigma_i + 3 * dt / 4) * s_in, **extra_args)) * dt

                    # Reversible correction term (inspired by Reversible Heun)
                    correction = dt_reversible**2 * (k3 - k2) / 6

                    # Update the sample
                    z_k = z_k + 2 * k1 / 9 + k2 / 3 + 4 * k3 / 9 - correction
                case "trapezoidal": # 2 model calls
                    if sigmas[i + 1] > 0:
                        dt = sigmas[i + 1] - sigmas[i]

                        # Calculate the derivative using the model
                        d_i = to_d(z_k, sigmas[i], denoised)

                        # Predict the sample at the next sigma using Euler step
                        x_pred = z_k + d_i * dt

                        # Denoised sample at the next sigma
                        denoised_i_plus_1 = model(x_pred, sigmas[i + 1] * s_in, **extra_args)

                        # Calculate the derivative at the next sigma
                        d_i_plus_1 = to_d(x_pred, sigmas[i + 1], denoised_i_plus_1)

                        dt_2 = sigma_down - sigmas[i]
                        # Update the sample using the Trapezoidal rule
                        z_k = z_k + dt_2 * (d_i + d_i_plus_1) / 2
                    else:
                        z_k = denoised
                case "RES":
                    if sigmas[i + 1] > 0:
                        lam_next = sigma_down.log().neg() if eta != 0 else sigmas[i + 1].log().neg()
                        lam = sigmas[i].log().neg()

                        h = lam_next - lam
                        a2_1, b1, b2 = _de_second_order(h=h, c2=0.5, simple_phi_calc=False)

                        c2_h = 0.5*h

                        x_2 = math.exp(-c2_h)*z_k + a2_1*h*denoised
                        lam_2 = lam + c2_h
                        sigma_2 = lam_2.neg().exp()

                        denoised2 = model(x_2, sigma_2 * s_in, **extra_args)

                        z_k = math.exp(-h)*z_k + h*(b1*denoised + b2*denoised2)
                    else:
                        z_k = denoised
            
            # Free Reverse
            if enable_free_reverse:
                if free_reverse_dyneta: free_reverse_eta = dyneta_fn(orig_free_reverse_eta, sigmas[i], sigmas[i + 1])
                sigma_down_freereversible, _ = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=free_reverse_eta)
                
                d_i = to_d(orig_zk, sigmas[i], denoised)

                dt_reversible = sigma_down_freereversible - sigmas[i]

                d_i_old = to_d(prev_x, sigmas[i], prev_denoised) if prev_denoised is not None else to_d(prev_x, sigmas[i], model(prev_x, sigmas[i] * s_in, **extra_args))

                z_k = z_k + (d_i - d_i_old) / 2 * dt - dt_reversible**2 * (d_i_old - d_i) / 2

            z_avg += renoise_weights[k] * z_k
            if sigmas[i + 1] > 0: # Random noise for variance on ancestral samplers
                noise_mod = noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
                match noise_modulation:
                    case "none":
                        noise_mod = noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
                    case "intensity":
                        noise = noise_sampler(sigmas[i], sigmas[i + 1])
                        noise_mod = intensity_based_multiplicative_noise_fn(z_k, noise, s_noise, sigma_up, modulation_strength, dims)
                    case "frequency":
                        noise = noise_sampler(sigmas[i], sigmas[i + 1])
                        noise_mod = frequency_based_noise(z_k, noise, s_noise, sigma_up, modulation_strength, dims)
                    case "spectral_signum":
                        noise = noise_sampler(sigmas[i], sigmas[i + 1])
                        noise_mod = spectral_modulate_noise(x, noise, s_noise, sigma_up, modulation_strength, dims)
                z_k = z_k + noise_mod

        x = z_avg
        if sigmas[i + 1] > 0:
            noise_mod = noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
            match noise_modulation:
                case "none":
                    noise_mod = noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
                case "intensity":
                    noise = noise_sampler(sigmas[i], sigmas[i + 1])
                    noise_mod = intensity_based_multiplicative_noise_fn(x, noise, s_noise, sigma_up, modulation_strength, dims)
                case "frequency":
                    noise = noise_sampler(sigmas[i], sigmas[i + 1])
                    noise_mod = frequency_based_noise(x, noise, s_noise, sigma_up, modulation_strength, dims)
                case "spectral_signum":
                    noise = noise_sampler(sigmas[i], sigmas[i + 1])
                    noise_mod = spectral_modulate_noise(x, noise, s_noise, sigma_up, modulation_strength, dims)

            x = x + noise_mod

        prev_x = x
        prev_denoised = denoised

    return x

def sample_supreme(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, eta=1.0, step_method="RES", substep_method="euler", warmup_method="euler", centralization=0.00, normalization=0.00, edge_enhancement=0.00, perphist=0.25, substeps=2, noise_modulation="none", modulation_strength=2., modulation_dims=3, reversible_eta=1.0, dyneta=True, reversible_dyneta=True, enable_free_reverse=True, free_reverse_eta=0.0, free_reverse_dyneta=True):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_supreme(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), eta=eta, step_method=step_method, substep_method=substep_method, warmup_method=warmup_method, centralization=centralization, normalization=normalization, edge_enhancement=edge_enhancement, perphist=perphist, substeps=substeps, noise_modulation=noise_modulation, modulation_strength=modulation_strength, modulation_dims=modulation_dims, reversible_eta=reversible_eta, dyneta=dyneta, reversible_dyneta=reversible_dyneta, enable_free_reverse=enable_free_reverse, free_reverse_eta=free_reverse_eta, free_reverse_dyneta=free_reverse_dyneta)

@torch.no_grad()
def sampler_sens(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., rsde_eta=1., tsde_eta=1., s_noise=1., noise_sampler=None, flow=False):
    """SDE-Endowed Nimble Sampler. Based off of DPM-Solver++(2M) SDE and DPM-Solver++(3M) SDE. R-SDE for reversible SDE, T-SDE for tertiary SDE."""
    if len(sigmas) <= 1:
        return x

    seed = extra_args.get("seed", None)
    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    old_denoised, old_denoised_2 = None, None
    h_last, h_last_2 = None, None
    h = None

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        if sigmas[i + 1] == 0:
            # Denoising step
            x = denoised
        else:
            # DPM-Solver++(2M) SDE
            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
            h = s - t
            eta_h = eta * h
            rsde_eta_h = rsde_eta * h
            tsde_eta_h = tsde_eta * h

            # If/for flow model
            downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
            sigma_down = sigmas[i+1] * downstep_ratio
            alpha_ip1 = 1 - sigmas[i+1]
            alpha_down = 1 - sigma_down
            renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5

            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised

            if old_denoised is not None:
                r = h_last / h
                x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised) / 2 - ((-h - rsde_eta_h).expm1().neg() / (-h - rsde_eta_h) + 1)**2 * (1 / r) * (old_denoised - denoised) / 2

            # DPM-Solver++(3M) SDE
            if h_last_2 is not None and tsde_eta:
                r = h_last_2 / h
                d = (old_denoised - old_denoised_2) / r
                d_2 = (old_denoised - denoised) / r
                
                d_rev = (denoised - old_denoised) / r
                d_2_rev = (old_denoised_2 - old_denoised) / r
                
                #phi = eta_h.neg().expm1() / eta_h + 1
                rphi = tsde_eta_h.neg().expm1() / tsde_eta_h + 1
                x = x + rphi * (d + d_2) / 2 - rphi**2 * (d_rev + d_2_rev) / 2

            if eta and not flow:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
            elif eta and flow:
                x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff

        old_denoised, old_denoised_2 = denoised, old_denoised
        h_last, h_last_2 = h, h_last
    return x

@torch.no_grad()
def sample_sens(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., rsde_eta=1., tsde_eta=1., s_noise=1., noise_sampler_type="brownian", noise_sampler=None):
    if len(sigmas) <= 1:
        return x
    flow = False
    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
        flow = True
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_sens(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, rsde_eta=rsde_eta, tsde_eta=tsde_eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), flow=flow)

#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
#under Apache 2 license
def sampler_ipndm_vapp(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4, eta=1., s_noise=1., noise_sampler=None, pp_guidance=1.0):
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler

    temp_uncond = [0]
    temp_cond = [0]
    def post_cfg_function(args):
        temp_uncond[0] = args["uncond_denoised"]
        temp_cond[0] = args["cond_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)

    s_in = x.new_ones([x.shape[0]])

    x_next = x
    t_steps = sigmas

    buffer_model = []
    for i in trange(len(sigmas) - 1, disable=disable):
        t_cur = sigmas[i]
        t_next = sigmas[i + 1]
        sigma_down, sigma_up = get_ancestral_step(t_cur, t_next, eta=eta)

        x_cur = x_next

        denoised = model(x_cur, t_cur * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})

        faux_d_cur = (x_cur - temp_uncond[0]) / t_cur # CFG++
        #d_cur = ((x_cur - temp_cond[0]) - (denoised - temp_uncond[0])) / t_cur # 2x CFG
        d_cur = -temp_cond[0] / t_cur * pp_guidance + (x_cur - denoised) / t_cur + temp_uncond[0] / t_cur * pp_guidance
        # I've found that chhanging x_cur to `denoised` results in over-denoised samples, so we're sticking with this alt method

        order = min(max_order, i+1)
        if order == 1:      # First Euler step.
            x_next = x_cur + (sigma_down - t_cur) * d_cur # Modified t_next to sigma_down for ancestral capability.
        elif order == 2:    # Use one history point.
            h_n = (t_next - t_cur)
            h_n_1 = (t_cur - t_steps[i-1])
            coeff1 = (2 + (h_n / h_n_1)) / 2
            coeff2 = -(h_n / h_n_1) / 2
            x_next = x_cur + (sigma_down - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1])
        elif order == 3:    # Use two history points.
            h_n = (t_next - t_cur)
            h_n_1 = (t_cur - t_steps[i-1])
            h_n_2 = (t_steps[i-1] - t_steps[i-2])
            temp = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
            coeff1 = (2 + (h_n / h_n_1)) / 2 + temp
            coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp
            coeff3 = temp * h_n_1 / h_n_2
            x_next = x_cur + (sigma_down - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2])
        elif order == 4:    # Use three history points.
            h_n = (t_next - t_cur)
            h_n_1 = (t_cur - t_steps[i-1])
            h_n_2 = (t_steps[i-1] - t_steps[i-2])
            h_n_3 = (t_steps[i-2] - t_steps[i-3])
            temp1 = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
            temp2 = ((1 - h_n / (3 * (h_n + h_n_1))) / 2 + (1 - h_n / (2 * (h_n + h_n_1))) * h_n / (6 * (h_n + h_n_1 + h_n_2))) \
                   * (h_n * (h_n + h_n_1) * (h_n + h_n_1 + h_n_2)) / (h_n_1 * (h_n_1 + h_n_2) * (h_n_1 + h_n_2 + h_n_3))
            coeff1 = (2 + (h_n / h_n_1)) / 2 + temp1 + temp2
            coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp1 - (1 + (h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3)))) * temp2
            coeff3 = temp1 * h_n_1 / h_n_2 + ((h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * (1 + h_n_2 / h_n_3)) * temp2
            coeff4 = -temp2 * (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * h_n_1 / h_n_2
            x_next = x_cur + (sigma_down - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2] + coeff4 * buffer_model[-3])
        
        if eta and sigmas[i + 1] > 0:
            x_next = x_next + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up

        if len(buffer_model) == max_order - 1:
            for k in range(max_order - 2):
                buffer_model[k] = buffer_model[k+1]
            buffer_model[-1] = faux_d_cur.detach() # Utilize CFG++ as history points
        else:
            buffer_model.append(faux_d_cur.detach())

    return x_next

@torch.no_grad()
def sample_ipndm_vapp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., max_order=4, noise_sampler_type="gaussian", noise_sampler=None, pp_guidance=1.0):
    if len(sigmas) <= 1:
        return x
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_ipndm_vapp(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, max_order=max_order, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), pp_guidance=pp_guidance)


import functools
import operator
@torch.no_grad()
def sampler_SHIDS(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, order=16, eta_order=1., solver_method="weighted_projection", flow=False):
    """Full ancestral sampling with SHIDS (Stochastic, Historical, Improvised Sampling) steps."""
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler

    temp_uncond = [0]
    temp_cond = [0]
    def post_cfg_function(args):
        temp_uncond[0] = args["uncond_denoised"]
        temp_cond[0] = args["cond_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)

    s_in = x.new_ones([x.shape[0]])
    old_uncond, old_uncond_2 = None, None
    old_cond, old_cond_2 = None, None
    old_dt, old_dt_2 = None, None

    buffer_model_cond = []
    #buffer_model_uncond = []
    #buffer_model_dt = []
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
        _, sigma_up_order = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta_order)

        # If/for flow model
        downstep_ratio = None
        sigma_down_rf = None
        alpha_ip1 = None
        alpha_down = None
        renoise_coeff = None
        if flow:
            downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
            sigma_down_rf = sigmas[i+1] * downstep_ratio
            alpha_ip1 = 1 - sigmas[i+1]
            alpha_down = 1 - sigma_down_rf
            renoise_coeff = (sigmas[i+1]**2 - sigma_down_rf**2*alpha_ip1**2/alpha_down**2)**0.5

        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        d_full = to_d(x, sigmas[i], denoised)
        d = to_d(x, sigmas[i], temp_uncond[0])
        #d_2 = to_d(x, sigmas[i], temp_cond[0])
        # Euler method
        dt = sigma_down - sigmas[i] # Time Difference between now and next step (negative)
        x_full = denoised + d_full * sigma_down
        x_step = denoised + d * sigma_down

        # Project denoised onto a line between (primarily) x_step (cfgpp), and x_full (normal cfg)
        match solver_method:
            case "weighted_projection":
                ba = x_step - denoised
                ca = x_full - denoised
                alpha = (ba * ca) / (ba ** 2 + 1e-8)
                x = (1 - alpha)*denoised + alpha*x_step
            case "qr_decomposition":
                original_shape = x_step.shape
                if not original_shape:
                    shape_2d = (1, 1)
                elif len(original_shape) == 4:
                    shape_2d = (-1, functools.reduce(operator.mul, original_shape[1:]))
                else:
                    shape_2d = (-1, original_shape[-1])

                A = x_step.reshape(shape_2d)
                B = x_full.reshape(shape_2d)
                C = denoised.reshape(shape_2d)
                Q, _ = torch.qr(A - C)
                # Compute the mapping matrix
                mapping_matrix = torch.mm(Q.t(), B - C)
                mapped_tensor = torch.mm(Q, mapping_matrix)
                x = (C + mapped_tensor).reshape(original_shape)
            case "svd_lowrank":
                original_shape = x_step.shape
                if not original_shape:
                    shape_2d = (1, 1)
                elif len(original_shape) == 4:
                    shape_2d = (-1, functools.reduce(operator.mul, original_shape[1:]))
                else:
                    shape_2d = (-1, original_shape[-1])

                A = x_step.reshape(shape_2d)
                B = x_full.reshape(shape_2d)
                C = denoised.reshape(shape_2d)
                Ua, Sa, Va = torch.svd_lowrank(A - C, q=6, niter=2)
                
                A_lowrank = torch.mm(Ua, torch.mm(torch.diag(Sa), Va.t()))
                A_diff = (A - C) - A_lowrank

                Qb, _ = torch.qr(B - C)

                A_diff_projected = torch.mm(Qb, torch.mm(Qb.t(), A_diff))

                x = (B + A_diff_projected).reshape(original_shape)
            case "svd":
                original_shape = x_step.shape
                if not original_shape:
                    shape_2d = (1, 1)
                elif len(original_shape) == 4:
                    shape_2d = (-1, functools.reduce(operator.mul, original_shape[1:]))
                else:
                    shape_2d = (-1, original_shape[-1])

                A = x_step.reshape(shape_2d)
                B = x_full.reshape(shape_2d)
                C = denoised.reshape(shape_2d)
                Ua, Sa, Va = torch.linalg.svd(A - C, full_matrices=False, driver="gesvd")
                Ub, Sb, Vb = torch.linalg.svd(B - C, full_matrices=False, driver="gesvd")#Sb = torch.linalg.svdvals(B - C, driver="gesvd")#
                
                A_lowrank = torch.mm(Ub, torch.mm(torch.diag_embed(Sa), Vb))
                A_diff = (A - C) - A_lowrank

                #Qb, _ = torch.qr(B - C)

                #A_diff_projected = torch.mm(Ua, torch.mm(Ua.t(), A_diff))
                    
                x = (B + A_diff).reshape(original_shape)

        # Create a list of order multipliers
        multipliers = [i for i in range(1, len(buffer_model_cond))]
        # Normalize so that they're summed up to a total of 1
        total = sum(multipliers)
        normalized_multipliers = [m / total for m in multipliers]

        for iteration in range(len(buffer_model_cond) - 1):
            if not flow:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up_order * normalized_multipliers[iteration]
            elif flow and eta_order:
                downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta_order
                sigma_down_rf = sigmas[i+1] * downstep_ratio
                alpha_ip1 = 1 - sigmas[i+1]
                alpha_down = 1 - sigma_down_rf
                renoise_coeff = (sigmas[i+1]**2 - sigma_down_rf**2*alpha_ip1**2/alpha_down**2)**0.5
                x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
            match solver_method:
                case "weighted_projection":
                    ba = x - buffer_model_cond[iteration]
                    ca = x_step - buffer_model_cond[iteration]
                    alpha = (ba * ca) / (ba ** 2 + 1e-8)
                    x = (1 - alpha)*buffer_model_cond[iteration] + alpha*x
                case "qr_decomposition":
                    original_shape = x_step.shape
                    if not original_shape:
                        shape_2d = (1, 1)
                    elif len(original_shape) == 4:
                        shape_2d = (-1, functools.reduce(operator.mul, original_shape[1:]))
                    else:
                        shape_2d = (-1, original_shape[-1])

                    A = x.reshape(shape_2d)
                    B = x_step.reshape(shape_2d)
                    C = buffer_model_cond[iteration].reshape(shape_2d)
                    Q, _ = torch.qr(A - C)
                    # Compute the mapping matrix
                    mapping_matrix = torch.mm(Q.t(), B - C)
                    mapped_tensor = torch.mm(Q, mapping_matrix)
                    x = (C + mapped_tensor).reshape(original_shape)
                case "svd_lowrank":
                    original_shape = x_step.shape
                    if not original_shape:
                        shape_2d = (1, 1)
                    elif len(original_shape) == 4:
                        shape_2d = (-1, functools.reduce(operator.mul, original_shape[1:]))
                    else:
                        shape_2d = (-1, original_shape[-1])

                    A = x.reshape(shape_2d)
                    B = x_step.reshape(shape_2d)
                    C = buffer_model_cond[iteration].reshape(shape_2d)
                    Ua, Sa, Va = torch.svd_lowrank(A - C, q=6, niter=2)
                    
                    A_lowrank = torch.mm(Ua, torch.mm(torch.diag(Sa), Va.t()))
                    A_diff = (A - C) - A_lowrank

                    #Qb, _ = torch.qr(B - C)

                    #A_diff_projected = torch.mm(Qb, torch.mm(Qb.t(), A_diff))
                    
                    x = (B + A_diff).reshape(original_shape)
                case "svd":
                    original_shape = x.shape
                    if not original_shape:
                        shape_2d = (1, 1)
                    elif len(original_shape) == 4:
                        shape_2d = (-1, functools.reduce(operator.mul, original_shape[1:]))
                    else:
                        shape_2d = (-1, original_shape[-1])

                    A = x.reshape(shape_2d)
                    B = x_step.reshape(shape_2d)
                    C = buffer_model_cond[iteration].reshape(shape_2d)
                    Ua, Sa, Va = torch.linalg.svd(A - C, full_matrices=False, driver="gesvd")
                    Ub, Sb, Vb = torch.linalg.svd(B - C, full_matrices=False, driver="gesvd")#Sb = torch.linalg.svdvals(B - C, driver="gesvd")#
                    
                    A_lowrank = torch.mm(Ub, torch.mm(torch.diag_embed(Sa), Vb))
                    A_diff = (A - C) - A_lowrank

                    #Qb, _ = torch.qr(B - C)

                    #A_diff_projected = torch.mm(Ua, torch.mm(Ua.t(), A_diff))
                    
                    x = (B + A_diff).reshape(original_shape)

        if len(buffer_model_cond) == max(order - 1, 1):
            for k in range(order - 2):
                buffer_model_cond[k] = buffer_model_cond[k+1]
                #buffer_model_uncond[k] = buffer_model_uncond[k+1]
                #buffer_model_dt[k] = buffer_model_dt[k+1]
            buffer_model_cond[-1] = denoised.detach()
            #buffer_model_uncond[-1] = temp_uncond[0].detach()
            #buffer_model_dt[-1] = dt.detach()
        else:
            buffer_model_cond.append(denoised.detach())
            #buffer_model_uncond.append(temp_uncond[0].detach())
            #buffer_model_dt.append(dt.detach())
        #if old_uncond is not None and old_cond is not None and order >= 2:
        #    x = x + (old_cond - old_uncond) / (old_dt / dt)
        #if old_uncond_2 is not None and old_cond_2 is not None and order >= 3:
        #    x = x + (old_cond_2 - old_uncond_2) / (old_dt_2 / old_dt) / (old_dt / dt)
        if sigmas[i + 1] > 0 and not flow:
            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
        elif sigmas[i + 1] > 0 and flow:
            x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
        #old_uncond, old_uncond_2 = temp[0], old_uncond
        #old_cond, old_cond_2 = temp_cond[0], old_cond
        #old_dt, old_dt_2 = dt, old_dt
    return x

@torch.no_grad()
def sample_SHIDS(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, order=16, eta_order=1., solver_method="weighted_projection"):
    if len(sigmas) <= 1:
        return x
    flow = False
    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
        flow = True
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_SHIDS(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), order=order, eta_order=eta_order, solver_method=solver_method, flow=flow)

@torch.no_grad()
def sampler_dpmpp_2m_sde_ema(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, amp_fac=2., beta1=0.8, beta2=0.95, weight_decay=0.1, centralization=1.0, normalization=1.0, flow=False):
    """DPM-Solver++(2M) SDE, with EMA uncond."""
    if len(sigmas) <= 1:
        return x

    seed = extra_args.get("seed", None)
    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    old_denoised = None
    h_last = None
    h = None

    ema = torch.zeros_like(x)
    ema_squared = torch.zeros_like(x)

    grad = None
    temp_cond = [0]
    temp_uncond = [0]
    #alpha = torch.linspace(1.0, 0.0, steps=len(sigmas)) ** amp_fac
    alpha = [0]
    def ema_retrieve_uncond_alpha(args):
        temp_cond[0] = args["cond_denoised"]
        temp_uncond[0] = args["uncond_denoised"]
        alpha[0] = model.inner_model.inner_model.model_sampling.timestep(args["sigma"]) / 999.0
        #alpha[0] = args["sigma"]
        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, ema_retrieve_uncond_alpha, disable_cfg1_optimization=True)

    ema = torch.zeros_like(x)
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        if sigmas[i + 1] == 0:
            # Denoising step
            x = denoised
        else:
            # DPM-Solver++(2M) SDE
            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
            h = s - t
            eta_h = eta * h

            # If/for flow model
            downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
            sigma_down = sigmas[i+1] * downstep_ratio
            alpha_ip1 = 1 - sigmas[i+1]
            alpha_down = 1 - sigma_down
            renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5

            grad = denoised

            # Centralization
            if centralization != 0:
                grad.sub_(
                    grad.mean(dim=tuple(range(1, grad.dim())), keepdim=True).mul_(centralization)
                )
            # Lerp EMA
            ema.lerp_(grad, 1. - beta1)
            # Normalization
            ema.lerp_(ema.div(ema.std(dim=tuple(range(1, grad.dim())), keepdim=True)), weight=normalization)
            # Apply EMA onto grad (denoised)
            grad.lerp_(ema, beta2)

            if weight_decay != 0:
                # Perform stepweight decay
                wd_mult = 1 / (1 + weight_decay * (sigmas[i] - sigmas[i + 1]))
                grad.mul_(wd_mult)

            ema += (x - temp_uncond[0]) / sigmas[i] * amp_fac * (sigmas[i] - sigmas[i + 1])
            ema -= (x - temp_cond[0]) / sigmas[i] * amp_fac * (sigmas[i] - sigmas[i + 1])

            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * grad

            if old_denoised is not None:
                r = h_last / h
                x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (grad - old_denoised)
            
            if eta and not flow:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
            elif eta and flow:
                x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff

        old_denoised = denoised
        h_last = h
    return x

@torch.no_grad()
def sample_dpmpp_2m_sde_ema(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="brownian", noise_sampler=None, amp_fac=2., beta1=0.8, beta2=0.95, weight_decay=0.1, centralization=1.0, normalization=1.0):
    if len(sigmas) <= 1:
        return x
    flow = False
    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
        flow = True
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_dpmpp_2m_sde_ema(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), amp_fac=amp_fac, beta1=beta1, beta2=beta2, weight_decay=weight_decay, centralization=centralization, normalization=normalization, flow=flow)

@torch.no_grad()
def sampler_biscope(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, amp_fac=2.0, local_smoothing_fac=4, smoothing_fac=0.75, ema_fac=0.9, flow=False):
    """Solving for the Compass model's noise problem using the Compass-like training procedure as an inference sampler."""
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])

    local_smoothing = []
    smoothing = None
    smoothing_diff = None
    ema = None
    prev_denoised = None
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)

        grad = denoised

        if len(local_smoothing) == max(local_smoothing_fac, 1):
            for k in range(local_smoothing_fac - 1):
                local_smoothing[k] = local_smoothing[k+1]
            local_smoothing[-1] = grad.detach()
        else:
            local_smoothing.append(grad.detach())
        #print(local_smoothing)
        local_grad = torch.mean(torch.stack(local_smoothing), dim=0)# if len(local_smoothing) > 1 else grad

        if smoothing is None:
            smoothing = local_grad

        smoothing.mul_(smoothing_fac).add_(local_grad, alpha=1 - smoothing_fac)

        diff_grad = local_grad - smoothing

        if smoothing_diff is None:
            smoothing_diff = diff_grad
        smoothing_diff.mul_(smoothing_fac).add_(diff_grad, alpha=1 - smoothing_fac)

        local_grad.add_(smoothing_diff, alpha=amp_fac)

        if ema is None:
            ema = local_grad
        ema.mul_(ema_fac).add_(local_grad, alpha=1 - ema_fac)

        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)

        # Flow
        downstep_ratio = None
        alpha_ip1 = None
        alpha_down = None
        renoise_coeff = None
        if flow:
            # If/for flow model
            downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
            sigma_down = sigmas[i+1] * downstep_ratio
            alpha_ip1 = 1 - sigmas[i+1]
            alpha_down = 1 - sigma_down
            renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': ema})
        d = to_d(x, sigmas[i], ema)
        
        # Euler method
        dt = sigma_down - sigmas[i]
        x = x + d * dt
        if sigmas[i + 1] > 0 and not flow:
            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
        elif sigmas[i + 1] > 0 and flow:
            x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff

        prev_denoised = denoised
    return x

@torch.no_grad()
def sample_biscope(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, amp_fac=2.0, local_smoothing_fac=4, smoothing_fac=0.5, ema_fac=0.5):
    if len(sigmas) <= 1:
        return x
    flow = False
    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
        flow = True
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_biscope(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), amp_fac=amp_fac, local_smoothing_fac=local_smoothing_fac, smoothing_fac=smoothing_fac, ema_fac=ema_fac, flow=flow)

def gaussian_kernel_2d(kernel_size, sigma):
    """Generates a 2D Gaussian kernel."""
    k = kernel_size // 2
    x, y = torch.meshgrid(torch.arange(-k, k + 1, dtype=torch.float32), torch.arange(-k, k + 1, dtype=torch.float32))
    gaussian = torch.exp(-(x**2 + y**2) / (2 * sigma**2))
    return gaussian / gaussian.sum()


@torch.no_grad()
def sampler_euler_g(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, g_eta=1.0, sigma=5.0, order=2, flow=False):
    """Ancestral sampling with Euler method steps."""
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])

    neighborhood_size = min(x.shape[-2], x.shape[-1]) * 2 + 1
    padding = neighborhood_size // 2
    kernel = gaussian_kernel_2d(neighborhood_size, sigma).unsqueeze(0).unsqueeze(0).repeat(x.shape[1], 1, 1, 1).to(x.device)

    x_buffer = []
    denoised_buffer = []
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if sigmas[i + 1] == 0:
            return denoised
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
        # Flow
        downstep_ratio = None
        alpha_ip1 = None
        alpha_down = None
        renoise_coeff = None
        if flow:
            # If/for flow model
            downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
            sigma_down = sigmas[i+1] * downstep_ratio
            alpha_ip1 = 1 - sigmas[i+1]
            alpha_down = 1 - sigma_down
            renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        d = to_d(x, sigmas[i], denoised)
        # Euler method
        dt = sigma_down - sigmas[i]
        x = x + d * dt

        """
        for curr_order in range(1, order):
            if sigmas[i + 1] > 0 and not flow:
                faux_x = torch.nn.functional.conv2d(x, kernel, padding=padding, groups=x.shape[1]) + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * (sigma_up)
            elif sigmas[i + 1] > 0 and flow:
                #x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
                #faux_x = (alpha_ip1/(1 - sigmas[i])) * torch.nn.functional.conv2d(x, kernel, padding=padding, groups=x.shape[1]) + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * (sigmas[i+1]**2 - sigmas[i]**2*alpha_ip1**2/(1 - sigmas[i])**2)**0.5
                faux_x = (alpha_ip1/alpha_down) * torch.nn.functional.conv2d(x, kernel, padding=padding, groups=x.shape[1]) + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
            #gauss_x = F.conv2d(faux_x, kernel, padding=padding, groups=x.shape[1])
            faux_denoised = model(faux_x, sigmas[i + 1] * s_in, **extra_args)
            faux_d = to_d(faux_x, sigmas[i + 1], faux_denoised)
            x = x - faux_d * (sigmas[i + 1] - sigmas[i]) * g_eta / (order - 1)
        """

        if sigmas[i + 1] > 0:
            # Create a list of order multipliers
            multipliers = [i for i in range(1, len(x_buffer))]
            # Normalize so that they're summed up to a total of 1
            total = sum(multipliers)
            normalized_multipliers = [m / total for m in multipliers]

            for iteration in range(len(x_buffer) - 1):
                if not flow:
                    faux_x = torch.nn.functional.conv2d(x_buffer[iteration], kernel, padding=padding, groups=x.shape[1]) + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
                else:
                    #x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
                    #faux_x = (alpha_ip1/(1 - sigmas[i])) * torch.nn.functional.conv2d(x, kernel, padding=padding, groups=x.shape[1]) + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * (sigmas[i+1]**2 - sigmas[i]**2*alpha_ip1**2/(1 - sigmas[i])**2)**0.5
                    faux_x = (alpha_ip1/alpha_down) * torch.nn.functional.conv2d(x_buffer[iteration], kernel, padding=padding, groups=x.shape[1]) + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
                #gauss_x = F.conv2d(faux_x, kernel, padding=padding, groups=x.shape[1])
                faux_denoised = model(faux_x, sigmas[i + 1] * s_in, **extra_args)
                faux_d = to_d(faux_x, sigmas[i + 1], faux_denoised)
                x = x - faux_d * (sigmas[i + 1] - sigmas[i]) * g_eta * normalized_multipliers[iteration]

            if len(x_buffer) == max(order - 1, 1):
                for k in range(order - 2):
                    x_buffer[k] = x_buffer[k+1]
                    denoised_buffer[k] = denoised_buffer[k+1]
                x_buffer[-1] = x.detach()
                denoised_buffer[-1] = denoised.detach()
            else:
                x_buffer.append(x.detach())
                denoised_buffer.append(denoised.detach())

            if flow:
                x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
            else:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
    return x

@torch.no_grad()
def sample_euler_g(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, g_eta=1.0, sigma=5.0, order=2):
    if len(sigmas) <= 1:
        return x
    flow = False
    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
        flow = True
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_euler_g(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), g_eta=g_eta, sigma=sigma, order=order, flow=flow)

@torch.no_grad()
def sampler_leaping_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, leap=1, eta=1., s_noise=1., noise_sampler=None, flow=False):
    #if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
    #    return sample_euler_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler)
    """Ancestral sampling with Euler method steps."""
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)

        do_dance = i < (len(sigmas) - (2 + leap))
        if not do_dance:
            leap -= 1
            do_dance = True

        sigma_next = sigmas[i + (1 + leap)] if do_dance else sigmas[i + 1]
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)

        # Flow
        downstep_ratio = None
        alpha_ip1 = None
        alpha_down = None
        renoise_coeff = None
        if flow:
            # If/for flow model
            downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
            sigma_down = sigmas[i+1] * downstep_ratio
            alpha_ip1 = 1 - sigmas[i+1]
            alpha_down = 1 - sigma_down
            renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5

        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        d = to_d(x, sigmas[i], denoised)
        # Euler method
        dt = sigma_next - sigmas[i]
        x_2 = x + d * dt
        
        if do_dance:
            reverse_denoised = model(x_2, sigma_next * s_in, **extra_args)
            _, r_sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
            r_d = to_d(x_2, sigma_next, reverse_denoised)
            r_dt = sigma_down - sigma_next
            x_2 = x + d * dt + r_d * r_dt
            if sigmas[i + 1] > 0 and not flow:
                x_2 = x_2 + noise_sampler(sigma_next, sigmas[i+1]) * s_noise * sigma_up
            elif flow:
                x_2 = (alpha_ip1/alpha_down) * x_2 + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
        x = x_2

    return x

@torch.no_grad()
def sample_leaping_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, leap=1, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None):
    if len(sigmas) <= 1:
        return x
    flow = False
    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
        flow = True
    noise_sampler, extra_args = check_set_immiscible(x, noise_sampler_type, extra_args)
    return sampler_leaping_euler(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, leap=leap, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler if noise_sampler is not None else get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), flow=flow)

# Add your personal samplers below here, just for formatting purposes ;3

# Add any extra samplers to the following dictionary
extra_samplers = {
    "res_momentumized": sample_res_solver,
    "dpmpp_dualsde_momentumized": sample_dpmpp_dualsdemomentum,
    "clyb_4m_sde_momentumized": sample_clyb_4m_sde,
    "ttm": sample_ttmcustom,
    "lcm_custom_noise": sample_lcmcustom,
    "euler_ancestral_dancing": sample_euler_ancestral_dancing,
    "dpmpp_3m_sde_dynamic_eta": sample_dpmpp_3m_sde_dynamic_eta,
    "supreme": sample_supreme,
    "sens": sample_sens,
    "ipndm_vapp": sample_ipndm_vapp,
    "SHIDS": sample_SHIDS,
    "dpmpp_2m_sde_ema": sample_dpmpp_2m_sde_ema,
    "biscope": sample_biscope,
    "euler_g": sample_euler_g,
    "leaping_euler": sample_leaping_euler,
}

discard_penultimate_sigma_samplers = set((
    "dpmpp_dualsde_momentumized",
    "clyb_4m_sde_momentumized"
))

def get_sigmas_simple_exponential(model, steps):
    s = model.model_sampling
    sigs = []
    ss = len(s.sigmas) / steps
    for x in range(steps):
        sigs += [float(s.sigmas[-(1 + int(x * ss))])]
    sigs += [0.0]
    sigs = torch.FloatTensor(sigs)
    exp = torch.exp(torch.log(torch.linspace(1, 0, steps + 1)))
    return sigs * exp

def get_sigmas_kl_optimal(model, steps):
    s = model.model_sampling
    sigs = []
    alpha_min = torch.arctan(s.sigma_min).item()
    alpha_max = torch.arctan(s.sigma_max).item()
    for x in range(steps+1):
        sigs += [torch.tan(torch.tensor(((x/steps) * alpha_min + (1.0-x/steps) * alpha_max)))]
    return torch.FloatTensor(sigs)

def get_sigmas_simple_kl_optimal(model, steps):
    s = model.model_sampling
    sigs = []
    idx_list = []
    ss = len(s.sigmas) / steps
    for x in range(steps):
        step = (x/steps) * math.atan(len(s.sigmas) / steps) + (x/steps) * math.atan(1 / steps)
        idx = int(len(s.sigmas) * (1.0 - math.atan(step))) - 1
        idx_list += [idx]
        sigs += [float(s.sigmas[idx])]
    #print(idx_list)
    sigs += [0.0]
    return torch.FloatTensor(sigs)

extra_schedulers = {
    "simple_exponential": get_sigmas_simple_exponential,
    "kl_optimal": get_sigmas_kl_optimal,
    "simple_kl_optimal": get_sigmas_simple_kl_optimal,
}