i

modules/samplers_advanced.py

@@ -23,8 +23,8 @@ class KSamplerAdvanced:
             sampler = self.SAMPLERS[0]
         self.scheduler = scheduler
         self.sampler = sampler
-        self.sigma_min=float(self.model_wrap.sigma_min)
-        self.sigma_max=float(self.model_wrap.sigma_max)
+        self.sigma_min = float(self.model_wrap.sigma_min)
+        self.sigma_max = float(self.model_wrap.sigma_max)
         self.set_steps(steps, denoise)
         self.denoise = denoise
         self.model_options = model_options
@@ -40,7 +40,8 @@ class KSamplerAdvanced:
         if self.scheduler == "karras":
             sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
         elif self.scheduler == "exponential":
-            sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
+            sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=self.sigma_min,
+                                                                 sigma_max=self.sigma_max)
         elif self.scheduler == "normal":
             sigmas = self.model_wrap.get_sigmas(steps)
         elif self.scheduler == "simple":
@@ -59,11 +60,12 @@ class KSamplerAdvanced:
         if denoise is None or denoise > 0.9999:
             self.sigmas = self.calculate_sigmas(steps).to(self.device)
         else:
-            new_steps = int(steps/denoise)
+            new_steps = int(steps / denoise)
             sigmas = self.calculate_sigmas(new_steps).to(self.device)
             self.sigmas = sigmas[-(steps + 1):]
 
-    def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
+    def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None,
+               force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
         if sigmas is None:
             sigmas = self.sigmas
         sigma_min = self.sigma_min
@@ -92,7 +94,7 @@ class KSamplerAdvanced:
         calculate_start_end_timesteps(self.model_wrap, negative)
         calculate_start_end_timesteps(self.model_wrap, positive)
 
-        #make sure each cond area has an opposite one with the same area
+        # make sure each cond area has an opposite one with the same area
         for c in positive:
             create_cond_with_same_area_if_none(negative, c)
         for c in negative:
@@ -100,30 +102,36 @@ class KSamplerAdvanced:
 
         pre_run_control(self.model_wrap, negative + positive)
 
-        apply_empty_x_to_equal_area(list(filter(lambda c: c[1].get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
+        apply_empty_x_to_equal_area(
+            list(filter(lambda c: c[1].get('control_apply_to_uncond', False) == True, positive)), negative, 'control',
+            lambda cond_cnets, x: cond_cnets[x])
         apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
 
         if self.model.is_adm():
-            positive = encode_adm(self.model, positive, noise.shape[0], noise.shape[3], noise.shape[2], self.device, "positive")
-            negative = encode_adm(self.model, negative, noise.shape[0], noise.shape[3], noise.shape[2], self.device, "negative")
+            positive = encode_adm(self.model, positive, noise.shape[0], noise.shape[3], noise.shape[2], self.device,
+                                  "positive")
+            negative = encode_adm(self.model, negative, noise.shape[0], noise.shape[3], noise.shape[2], self.device,
+                                  "negative")
 
         if latent_image is not None:
             latent_image = self.model.process_latent_in(latent_image)
 
-        extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": self.model_options, "seed":seed}
+        extra_args = {"cond": positive, "uncond": negative, "cond_scale": cfg, "model_options": self.model_options,
+                      "seed": seed}
 
         cond_concat = None
-        if hasattr(self.model, 'concat_keys'): #inpaint
+        if hasattr(self.model, 'concat_keys'):  # inpaint
             cond_concat = []
             for ck in self.model.concat_keys:
                 if denoise_mask is not None:
                     if ck == "mask":
-                        cond_concat.append(denoise_mask[:,:1])
+                        cond_concat.append(denoise_mask[:, :1])
                     elif ck == "masked_image":
-                        cond_concat.append(latent_image) #NOTE: the latent_image should be masked by the mask in pixel space
+                        cond_concat.append(
+                            latent_image)  # NOTE: the latent_image should be masked by the mask in pixel space
                 else:
                     if ck == "mask":
-                        cond_concat.append(torch.ones_like(noise)[:,:1])
+                        cond_concat.append(torch.ones_like(noise)[:, :1])
                     elif ck == "masked_image":
                         cond_concat.append(blank_inpaint_image_like(noise))
             extra_args["cond_concat"] = cond_concat
@@ -133,11 +141,16 @@ class KSamplerAdvanced:
         else:
             max_denoise = True
 
-
         if self.sampler == "uni_pc":
-            samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask, callback=callback, disable=disable_pbar)
+            samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas,
+                                          sampling_function=sampling_function, max_denoise=max_denoise,
+                                          extra_args=extra_args, noise_mask=denoise_mask, callback=callback,
+                                          disable=disable_pbar)
         elif self.sampler == "uni_pc_bh2":
-            samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar)
+            samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas,
+                                          sampling_function=sampling_function, max_denoise=max_denoise,
+                                          extra_args=extra_args, noise_mask=denoise_mask, callback=callback,
+                                          variant='bh2', disable=disable_pbar)
         elif self.sampler == "ddim":
             timesteps = []
             for s in range(sigmas.shape[0]):
@@ -153,24 +166,26 @@ class KSamplerAdvanced:
 
             sampler = DDIMSampler(self.model, device=self.device)
             sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False)
-            z_enc = sampler.stochastic_encode(latent_image, torch.tensor([len(timesteps) - 1] * noise.shape[0]).to(self.device), noise=noise, max_denoise=max_denoise)
+            z_enc = sampler.stochastic_encode(latent_image,
+                                              torch.tensor([len(timesteps) - 1] * noise.shape[0]).to(self.device),
+                                              noise=noise, max_denoise=max_denoise)
             samples, _ = sampler.sample_custom(ddim_timesteps=timesteps,
-                                                    conditioning=positive,
-                                                    batch_size=noise.shape[0],
-                                                    shape=noise.shape[1:],
-                                                    verbose=False,
-                                                    unconditional_guidance_scale=cfg,
-                                                    unconditional_conditioning=negative,
-                                                    eta=0.0,
-                                                    x_T=z_enc,
-                                                    x0=latent_image,
-                                                    img_callback=ddim_callback,
-                                                    denoise_function=self.model_wrap.predict_eps_discrete_timestep,
-                                                    extra_args=extra_args,
-                                                    mask=noise_mask,
-                                                    to_zero=sigmas[-1]==0,
-                                                    end_step=sigmas.shape[0] - 1,
-                                                    disable_pbar=disable_pbar)
+                                               conditioning=positive,
+                                               batch_size=noise.shape[0],
+                                               shape=noise.shape[1:],
+                                               verbose=False,
+                                               unconditional_guidance_scale=cfg,
+                                               unconditional_conditioning=negative,
+                                               eta=0.0,
+                                               x_T=z_enc,
+                                               x0=latent_image,
+                                               img_callback=ddim_callback,
+                                               denoise_function=self.model_wrap.predict_eps_discrete_timestep,
+                                               extra_args=extra_args,
+                                               mask=noise_mask,
+                                               to_zero=sigmas[-1] == 0,
+                                               end_step=sigmas.shape[0] - 1,
+                                               disable_pbar=disable_pbar)
 
         else:
             extra_args["denoise_mask"] = denoise_mask
@@ -190,11 +205,17 @@ class KSamplerAdvanced:
             if latent_image is not None:
                 noise += latent_image
             if self.sampler == "dpm_fast":
-                samples = k_diffusion_sampling.sample_dpm_fast(self.model_k, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
+                samples = k_diffusion_sampling.sample_dpm_fast(self.model_k, noise, sigma_min, sigmas[0], total_steps,
+                                                               extra_args=extra_args, callback=k_callback,
+                                                               disable=disable_pbar)
             elif self.sampler == "dpm_adaptive":
-                samples = k_diffusion_sampling.sample_dpm_adaptive(self.model_k, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=k_callback, disable=disable_pbar)
+                samples = k_diffusion_sampling.sample_dpm_adaptive(self.model_k, noise, sigma_min, sigmas[0],
+                                                                   extra_args=extra_args, callback=k_callback,
+                                                                   disable=disable_pbar)
             else:
-                samples = getattr(k_diffusion_sampling, "sample_{}".format(self.sampler))(self.model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
+                samples = getattr(k_diffusion_sampling, "sample_{}".format(self.sampler))(self.model_k, noise, sigmas,
+                                                                                          extra_args=extra_args,
+                                                                                          callback=k_callback,
+                                                                                          disable=disable_pbar)
 
         return self.model.process_latent_out(samples.to(torch.float32))
-