Use the ComfyUI version of k_diffusion.sampling

This commit is contained in:
blepping
2024-02-14 03:38:26 -07:00
parent e58d971bec
commit 50799c9474
2 changed files with 8 additions and 9 deletions
+2 -3
View File
@@ -8,8 +8,7 @@ from tqdm.auto import trange, tqdm
import comfy.sample
import k_diffusion.sampling
from k_diffusion.sampling import BrownianTreeNoiseSampler, PIDStepSizeController, get_ancestral_step, to_d, default_noise_sampler
from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler, PIDStepSizeController, get_ancestral_step, to_d, default_noise_sampler
import random
# The following function adds the samplers during initialization, in __init__.py
@@ -596,7 +595,7 @@ def sample_ttmcustom(model, x, sigmas, extra_args=None, callback=None, disable=N
noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
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)
from k_diffusion.sampling import sample_lcm
from comfy.k_diffusion.sampling import sample_lcm
def sample_lcmcustom(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
seed = extra_args.get("seed", None)
+1 -1
View File
@@ -1,4 +1,4 @@
from k_diffusion.sampling import default_noise_sampler, to_d
from comfy.k_diffusion.sampling import default_noise_sampler, to_d
from tqdm import trange
import torch
from torch import enable_grad