Spaces:
Running
Running
| import nodes | |
| from comfy.k_diffusion import sampling as k_diffusion_sampling | |
| from comfy import samplers | |
| from comfy_extras import nodes_custom_sampler | |
| import torch | |
| import math | |
| def calculate_sigmas(model, sampler, scheduler, steps): | |
| discard_penultimate_sigma = False | |
| if sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']: | |
| steps += 1 | |
| discard_penultimate_sigma = True | |
| sigmas = samplers.calculate_sigmas_scheduler(model.model, scheduler, steps) | |
| if discard_penultimate_sigma: | |
| sigmas = torch.cat([sigmas[:-2], sigmas[-1:]]) | |
| return sigmas | |
| def get_noise_sampler(x, cpu, total_sigmas, **kwargs): | |
| if 'extra_args' in kwargs and 'seed' in kwargs['extra_args']: | |
| sigma_min, sigma_max = total_sigmas[total_sigmas > 0].min(), total_sigmas.max() | |
| seed = kwargs['extra_args'].get("seed", None) | |
| return k_diffusion_sampling.BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=cpu) | |
| return None | |
| def ksampler(sampler_name, total_sigmas, extra_options={}, inpaint_options={}): | |
| if sampler_name == "dpmpp_sde": | |
| def sample_dpmpp_sde(model, x, sigmas, **kwargs): | |
| noise_sampler = get_noise_sampler(x, True, total_sigmas, **kwargs) | |
| if noise_sampler is not None: | |
| kwargs['noise_sampler'] = noise_sampler | |
| return k_diffusion_sampling.sample_dpmpp_sde(model, x, sigmas, **kwargs) | |
| sampler_function = sample_dpmpp_sde | |
| elif sampler_name == "dpmpp_sde_gpu": | |
| def sample_dpmpp_sde(model, x, sigmas, **kwargs): | |
| noise_sampler = get_noise_sampler(x, False, total_sigmas, **kwargs) | |
| if noise_sampler is not None: | |
| kwargs['noise_sampler'] = noise_sampler | |
| return k_diffusion_sampling.sample_dpmpp_sde_gpu(model, x, sigmas, **kwargs) | |
| sampler_function = sample_dpmpp_sde | |
| elif sampler_name == "dpmpp_2m_sde": | |
| def sample_dpmpp_sde(model, x, sigmas, **kwargs): | |
| noise_sampler = get_noise_sampler(x, True, total_sigmas, **kwargs) | |
| if noise_sampler is not None: | |
| kwargs['noise_sampler'] = noise_sampler | |
| return k_diffusion_sampling.sample_dpmpp_2m_sde(model, x, sigmas, **kwargs) | |
| sampler_function = sample_dpmpp_sde | |
| elif sampler_name == "dpmpp_2m_sde_gpu": | |
| def sample_dpmpp_sde(model, x, sigmas, **kwargs): | |
| noise_sampler = get_noise_sampler(x, False, total_sigmas, **kwargs) | |
| if noise_sampler is not None: | |
| kwargs['noise_sampler'] = noise_sampler | |
| return k_diffusion_sampling.sample_dpmpp_2m_sde_gpu(model, x, sigmas, **kwargs) | |
| sampler_function = sample_dpmpp_sde | |
| elif sampler_name == "dpmpp_3m_sde": | |
| def sample_dpmpp_sde(model, x, sigmas, **kwargs): | |
| noise_sampler = get_noise_sampler(x, True, total_sigmas, **kwargs) | |
| if noise_sampler is not None: | |
| kwargs['noise_sampler'] = noise_sampler | |
| return k_diffusion_sampling.sample_dpmpp_2m_sde(model, x, sigmas, **kwargs) | |
| sampler_function = sample_dpmpp_sde | |
| elif sampler_name == "dpmpp_3m_sde_gpu": | |
| def sample_dpmpp_sde(model, x, sigmas, **kwargs): | |
| noise_sampler = get_noise_sampler(x, False, total_sigmas, **kwargs) | |
| if noise_sampler is not None: | |
| kwargs['noise_sampler'] = noise_sampler | |
| return k_diffusion_sampling.sample_dpmpp_2m_sde_gpu(model, x, sigmas, **kwargs) | |
| sampler_function = sample_dpmpp_sde | |
| else: | |
| return samplers.ksampler(sampler_name, extra_options, inpaint_options) | |
| return samplers.KSAMPLER(sampler_function, extra_options, inpaint_options) | |
| def separated_sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative, | |
| latent_image, start_at_step, end_at_step, return_with_leftover_noise, sigma_ratio=1.0, sampler_opt=None): | |
| if sampler_opt is None: | |
| total_sigmas = calculate_sigmas(model, sampler_name, scheduler, steps) | |
| else: | |
| total_sigmas = calculate_sigmas(model, "", scheduler, steps) | |
| sigmas = total_sigmas[start_at_step:end_at_step+1] * sigma_ratio | |
| if sampler_opt is None: | |
| impact_sampler = ksampler(sampler_name, total_sigmas) | |
| else: | |
| impact_sampler = sampler_opt | |
| if len(sigmas) == 0 or (len(sigmas) == 1 and sigmas[0] == 0): | |
| return latent_image | |
| res = nodes_custom_sampler.SamplerCustom().sample(model, add_noise, seed, cfg, positive, negative, impact_sampler, sigmas, latent_image) | |
| if return_with_leftover_noise: | |
| return res[0] | |
| else: | |
| return res[1] | |
| def ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, | |
| refiner_ratio=None, refiner_model=None, refiner_clip=None, refiner_positive=None, refiner_negative=None, sigma_factor=1.0): | |
| if refiner_ratio is None or refiner_model is None or refiner_clip is None or refiner_positive is None or refiner_negative is None: | |
| refined_latent = nodes.KSampler().sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise * sigma_factor)[0] | |
| else: | |
| advanced_steps = math.floor(steps / denoise) | |
| start_at_step = advanced_steps - steps | |
| end_at_step = start_at_step + math.floor(steps * (1.0 - refiner_ratio)) | |
| # print(f"pre: {start_at_step} .. {end_at_step} / {advanced_steps}") | |
| temp_latent = separated_sample(model, True, seed, advanced_steps, cfg, sampler_name, scheduler, | |
| positive, negative, latent_image, start_at_step, end_at_step, True, sigma_ratio=sigma_factor) | |
| if 'noise_mask' in latent_image: | |
| # noise_latent = \ | |
| # impact_sampling.separated_sample(refiner_model, "enable", seed, advanced_steps, cfg, sampler_name, | |
| # scheduler, refiner_positive, refiner_negative, latent_image, end_at_step, | |
| # end_at_step, "enable") | |
| latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() | |
| temp_latent = latent_compositor.composite(latent_image, temp_latent, 0, 0, False, latent_image['noise_mask'])[0] | |
| # print(f"post: {end_at_step} .. {advanced_steps + 1} / {advanced_steps}") | |
| refined_latent = separated_sample(refiner_model, False, seed, advanced_steps, cfg, sampler_name, scheduler, | |
| refiner_positive, refiner_negative, temp_latent, end_at_step, advanced_steps + 1, False, sigma_ratio=sigma_factor) | |
| return refined_latent | |
| class KSamplerAdvancedWrapper: | |
| params = None | |
| def __init__(self, model, cfg, sampler_name, scheduler, positive, negative, sampler_opt=None, sigma_factor=1.0): | |
| self.params = model, cfg, sampler_name, scheduler, positive, negative, sigma_factor | |
| self.sampler_opt = sampler_opt | |
| def clone_with_conditionings(self, positive, negative): | |
| model, cfg, sampler_name, scheduler, _, _, _ = self.params | |
| return KSamplerAdvancedWrapper(model, cfg, sampler_name, scheduler, positive, negative, self.sampler_opt) | |
| def sample_advanced(self, add_noise, seed, steps, latent_image, start_at_step, end_at_step, return_with_leftover_noise, hook=None, | |
| recovery_mode="ratio additional", recovery_sampler="AUTO", recovery_sigma_ratio=1.0): | |
| model, cfg, sampler_name, scheduler, positive, negative, sigma_factor = self.params | |
| # steps, start_at_step, end_at_step = self.compensate_denoise(steps, start_at_step, end_at_step) | |
| if hook is not None: | |
| model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent = hook.pre_ksample_advanced(model, add_noise, seed, steps, cfg, sampler_name, scheduler, | |
| positive, negative, latent_image, start_at_step, end_at_step, | |
| return_with_leftover_noise) | |
| if recovery_mode != 'DISABLE' and sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu']: | |
| base_image = latent_image.copy() | |
| if recovery_mode == "ratio between": | |
| sigma_ratio = 1.0 - recovery_sigma_ratio | |
| else: | |
| sigma_ratio = 1.0 | |
| else: | |
| base_image = None | |
| sigma_ratio = 1.0 | |
| try: | |
| if sigma_ratio > 0: | |
| latent_image = separated_sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler, | |
| positive, negative, latent_image, start_at_step, end_at_step, | |
| return_with_leftover_noise, sigma_ratio=sigma_ratio * sigma_factor, sampler_opt=self.sampler_opt) | |
| except ValueError as e: | |
| if str(e) == 'sigma_min and sigma_max must not be 0': | |
| print(f"\nWARN: sampling skipped - sigma_min and sigma_max are 0") | |
| return latent_image | |
| if (recovery_sigma_ratio > 0 and recovery_mode != 'DISABLE' and | |
| sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu']): | |
| compensate = 0 if sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu'] else 2 | |
| if recovery_sampler == "AUTO": | |
| recovery_sampler = 'dpm_fast' if sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu'] else 'dpmpp_2m' | |
| latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() | |
| noise_mask = latent_image['noise_mask'] | |
| if len(noise_mask.shape) == 4: | |
| noise_mask = noise_mask.squeeze(0).squeeze(0) | |
| latent_image = latent_compositor.composite(base_image, latent_image, 0, 0, False, noise_mask)[0] | |
| try: | |
| latent_image = separated_sample(model, add_noise, seed, steps, cfg, recovery_sampler, scheduler, | |
| positive, negative, latent_image, start_at_step-compensate, end_at_step, | |
| return_with_leftover_noise, sigma_ratio=recovery_sigma_ratio * sigma_factor, sampler_opt=self.sampler_opt) | |
| except ValueError as e: | |
| if str(e) == 'sigma_min and sigma_max must not be 0': | |
| print(f"\nWARN: sampling skipped - sigma_min and sigma_max are 0") | |
| return latent_image | |
| class KSamplerWrapper: | |
| params = None | |
| def __init__(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise): | |
| self.params = model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise | |
| def sample(self, latent_image, hook=None): | |
| model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise = self.params | |
| if hook is not None: | |
| model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise = \ | |
| hook.pre_ksample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise) | |
| return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)[0] | |