diff --git a/modules/samplers_advanced.py b/modules/samplers_advanced.py index 957e27a3..e56562b8 100644 --- a/modules/samplers_advanced.py +++ b/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)) -