Merge pull request #7730 from CCRcmcpe/fix-dpm-sde-batch

Fix DPM++ SDE not deterministic across different batch sizes (#5210)

1 files changed, 30 insertions, 8 deletions

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index f076fc55..86d657e2 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -269,6 +269,16 @@ class KDiffusionSampler:
 

         return sigmas

 

+    def create_noise_sampler(self, x, sigmas, p):

+        """For DPM++ SDE: manually create noise sampler to enable deterministic results across different batch sizes"""

+        if shared.opts.no_dpmpp_sde_batch_determinism:

+            return None

+

+        from k_diffusion.sampling import BrownianTreeNoiseSampler

+        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()

+        current_iter_seeds = p.all_seeds[p.iteration * p.batch_size:(p.iteration + 1) * p.batch_size]

+        return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=current_iter_seeds)

+

     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):

         steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)

 

@@ -278,18 +288,24 @@ class KDiffusionSampler:
         xi = x + noise * sigma_sched[0]

         

         extra_params_kwargs = self.initialize(p)

-        if 'sigma_min' in inspect.signature(self.func).parameters:

+        parameters = inspect.signature(self.func).parameters

+

+        if 'sigma_min' in parameters:

             ## last sigma is zero which isn't allowed by DPM Fast & Adaptive so taking value before last

             extra_params_kwargs['sigma_min'] = sigma_sched[-2]

-        if 'sigma_max' in inspect.signature(self.func).parameters:

+        if 'sigma_max' in parameters:

             extra_params_kwargs['sigma_max'] = sigma_sched[0]

-        if 'n' in inspect.signature(self.func).parameters:

+        if 'n' in parameters:

             extra_params_kwargs['n'] = len(sigma_sched) - 1

-        if 'sigma_sched' in inspect.signature(self.func).parameters:

+        if 'sigma_sched' in parameters:

             extra_params_kwargs['sigma_sched'] = sigma_sched

-        if 'sigmas' in inspect.signature(self.func).parameters:

+        if 'sigmas' in parameters:

             extra_params_kwargs['sigmas'] = sigma_sched

 

+        if self.funcname == 'sample_dpmpp_sde':

+            noise_sampler = self.create_noise_sampler(x, sigmas, p)

+            extra_params_kwargs['noise_sampler'] = noise_sampler

+

         self.model_wrap_cfg.init_latent = x

         self.last_latent = x

         extra_args={

@@ -303,7 +319,7 @@ class KDiffusionSampler:
 

         return samples

 

-    def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning = None):

+    def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):

         steps = steps or p.steps

 

         sigmas = self.get_sigmas(p, steps)

@@ -311,14 +327,20 @@ class KDiffusionSampler:
         x = x * sigmas[0]

 

         extra_params_kwargs = self.initialize(p)

-        if 'sigma_min' in inspect.signature(self.func).parameters:

+        parameters = inspect.signature(self.func).parameters

+

+        if 'sigma_min' in parameters:

             extra_params_kwargs['sigma_min'] = self.model_wrap.sigmas[0].item()

             extra_params_kwargs['sigma_max'] = self.model_wrap.sigmas[-1].item()

-            if 'n' in inspect.signature(self.func).parameters:

+            if 'n' in parameters:

                 extra_params_kwargs['n'] = steps

         else:

             extra_params_kwargs['sigmas'] = sigmas

 

+        if self.funcname == 'sample_dpmpp_sde':

+            noise_sampler = self.create_noise_sampler(x, sigmas, p)

+            extra_params_kwargs['noise_sampler'] = noise_sampler

+

         self.last_latent = x

         samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={

             'cond': conditioning,