modules/samplers_advanced.py

from comfy.samplers import *

import comfy.model_management


class KSamplerWithRefiner:
    SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
    SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
                "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
                "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2"]

    def __init__(self, model, refiner_model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
        self.model_patcher = model
        self.refiner_model_patcher = refiner_model

        self.model = model.model
        self.refiner_model = refiner_model.model

        self.model_denoise = CFGNoisePredictor(self.model)
        self.refiner_model_denoise = CFGNoisePredictor(self.refiner_model)

        if self.model.model_type == model_base.ModelType.V_PREDICTION:
            self.model_wrap = CompVisVDenoiser(self.model_denoise, quantize=True)
        else:
            self.model_wrap = k_diffusion_external.CompVisDenoiser(self.model_denoise, quantize=True)

        if self.refiner_model.model_type == model_base.ModelType.V_PREDICTION:
            self.refiner_model_wrap = CompVisVDenoiser(self.refiner_model_denoise, quantize=True)
        else:
            self.refiner_model_wrap = k_diffusion_external.CompVisDenoiser(self.refiner_model_denoise, quantize=True)

        self.model_k = KSamplerX0Inpaint(self.model_wrap)
        self.refiner_model_k = KSamplerX0Inpaint(self.refiner_model_wrap)

        self.device = device
        if scheduler not in self.SCHEDULERS:
            scheduler = self.SCHEDULERS[0]
        if sampler not in self.SAMPLERS:
            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.set_steps(steps, denoise)
        self.denoise = denoise
        self.model_options = model_options

    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
        if self.sampler in ['dpm_2', 'dpm_2_ancestral']:
            steps += 1
            discard_penultimate_sigma = True

        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)
        elif self.scheduler == "normal":
            sigmas = self.model_wrap.get_sigmas(steps)
        elif self.scheduler == "simple":
            sigmas = simple_scheduler(self.model_wrap, steps)
        elif self.scheduler == "ddim_uniform":
            sigmas = ddim_scheduler(self.model_wrap, steps)
        else:
            print("error invalid scheduler", self.scheduler)

        if discard_penultimate_sigma:
            sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
        return sigmas

    def set_steps(self, steps, denoise=None):
        self.steps = steps
        if denoise is None or denoise > 0.9999:
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
        else:
            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, refiner_positive, refiner_negative, cfg, latent_image=None,
               start_step=None, last_step=None, refiner_switch_step=None,
               force_full_denoise=False, denoise_mask=None, sigmas=None, callback_function=None, disable_pbar=False, seed=None):
        if sigmas is None:
            sigmas = self.sigmas
        sigma_min = self.sigma_min

        if last_step is not None and last_step < (len(sigmas) - 1):
            sigma_min = sigmas[last_step]
            sigmas = sigmas[:last_step + 1]
            if force_full_denoise:
                sigmas[-1] = 0

        if start_step is not None:
            if start_step < (len(sigmas) - 1):
                sigmas = sigmas[start_step:]
            else:
                if latent_image is not None:
                    return latent_image
                else:
                    return torch.zeros_like(noise)

        positive = positive[:]
        negative = negative[:]

        resolve_cond_masks(positive, noise.shape[2], noise.shape[3], self.device)
        resolve_cond_masks(negative, noise.shape[2], noise.shape[3], self.device)

        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
        for c in positive:
            create_cond_with_same_area_if_none(negative, c)
        for c in negative:
            create_cond_with_same_area_if_none(positive, c)

        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(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")

        refiner_positive = refiner_positive[:]
        refiner_negative = refiner_negative[:]

        resolve_cond_masks(refiner_positive, noise.shape[2], noise.shape[3], self.device)
        resolve_cond_masks(refiner_negative, noise.shape[2], noise.shape[3], self.device)

        calculate_start_end_timesteps(self.refiner_model_wrap, refiner_positive)
        calculate_start_end_timesteps(self.refiner_model_wrap, refiner_negative)

        # make sure each cond area has an opposite one with the same area
        for c in refiner_positive:
            create_cond_with_same_area_if_none(refiner_negative, c)
        for c in refiner_negative:
            create_cond_with_same_area_if_none(refiner_positive, c)

        if self.model.is_adm():
            refiner_positive = encode_adm(self.refiner_model, refiner_positive, noise.shape[0],
                                          noise.shape[3], noise.shape[2], self.device, "positive")
            refiner_negative = encode_adm(self.refiner_model, refiner_negative, noise.shape[0],
                                          noise.shape[3], noise.shape[2], self.device, "negative")

        def refiner_switch():
            comfy.model_management.load_model_gpu(self.refiner_model_patcher)
            self.model_denoise.inner_model = self.refiner_model_denoise.inner_model
            for i in range(len(positive)):
                positive[i] = refiner_positive[i]
            for i in range(len(negative)):
                negative[i] = refiner_negative[i]
            print('Refiner swapped.')
            return

        def callback(step, x0, x, total_steps):
            if step == refiner_switch_step:
                refiner_switch()
            if callback_function is not None:
                callback_function(step, x0, x, total_steps)

        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}

        cond_concat = None
        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])
                    elif ck == "masked_image":
                        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])
                    elif ck == "masked_image":
                        cond_concat.append(blank_inpaint_image_like(noise))
            extra_args["cond_concat"] = cond_concat

        if sigmas[0] != self.sigmas[0] or (self.denoise is not None and self.denoise < 1.0):
            max_denoise = False
        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)
        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)
        elif self.sampler == "ddim":
            raise NotImplementedError('Swapped Refiner Does not support DDIM.')
        else:
            extra_args["denoise_mask"] = denoise_mask
            self.model_k.latent_image = latent_image
            self.model_k.noise = noise

            if max_denoise:
                noise = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0)
            else:
                noise = noise * sigmas[0]

            k_callback = None
            total_steps = len(sigmas) - 1
            if callback is not None:
                k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)

            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)
            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)
            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)

        return self.model.process_latent_out(samples.to(torch.float32))