Merge branch 'dev' into dev

1 files changed, 6 insertions, 7 deletions

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 5552a8dc..d72c1b5f 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -2,7 +2,7 @@ from collections import deque
 import torch

 import inspect

 import k_diffusion.sampling

-from modules import prompt_parser, devices, sd_samplers_common

+from modules import prompt_parser, devices, sd_samplers_common, sd_samplers_extra

 

 from modules.shared import opts, state

 import modules.shared as shared

@@ -30,12 +30,14 @@ samplers_k_diffusion = [
     ('DPM++ 2M Karras', 'sample_dpmpp_2m', ['k_dpmpp_2m_ka'], {'scheduler': 'karras'}),

     ('DPM++ SDE Karras', 'sample_dpmpp_sde', ['k_dpmpp_sde_ka'], {'scheduler': 'karras', "second_order": True, "brownian_noise": True}),

     ('DPM++ 2M SDE Karras', 'sample_dpmpp_2m_sde', ['k_dpmpp_2m_sde_ka'], {'scheduler': 'karras', "brownian_noise": True}),

+    ('Restart', sd_samplers_extra.restart_sampler, ['restart'], {'scheduler': 'karras'}),

 ]

 

+

 samplers_data_k_diffusion = [

     sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)

     for label, funcname, aliases, options in samplers_k_diffusion

-    if hasattr(k_diffusion.sampling, funcname)

+    if callable(funcname) or hasattr(k_diffusion.sampling, funcname)

 ]

 

 sampler_extra_params = {

@@ -258,10 +260,7 @@ class TorchHijack:
             if noise.shape == x.shape:

                 return noise

 

-        if opts.randn_source == "CPU" or x.device.type == 'mps':

-            return torch.randn_like(x, device=devices.cpu).to(x.device)

-        else:

-            return torch.randn_like(x)

+        return devices.randn_like(x)

 

 

 class KDiffusionSampler:

@@ -270,7 +269,7 @@ class KDiffusionSampler:
 

         self.model_wrap = denoiser(sd_model, quantize=shared.opts.enable_quantization)

         self.funcname = funcname

-        self.func = getattr(k_diffusion.sampling, self.funcname)

+        self.func = funcname if callable(funcname) else getattr(k_diffusion.sampling, self.funcname)

         self.extra_params = sampler_extra_params.get(funcname, [])

         self.model_wrap_cfg = CFGDenoiser(self.model_wrap)

         self.sampler_noises = None