Class KDiffusionSampler

scripts/webui.py:606–652

class KDiffusionSampler:
    def __init__(self, m, sampler):
        self.model = m
        self.model_wrap = K.external.CompVisDenoiser(m)
        self.schedule = sampler

    def get_sampler_name(self):
        return self.schedule

    def sample(
        self,
        S,
        conditioning,
        batch_size,
        shape,
        verbose,
        unconditional_guidance_scale,
        unconditional_conditioning,
        eta,
        x_T,
        img_callback: Callable = None,
    ):
        sigmas = self.model_wrap.get_sigmas(S)
        x = x_T * sigmas[0]
        model_wrap_cfg = CFGDenoiser(self.model_wrap)

        samples_ddim = K.sampling.__dict__[f"sample_{self.schedule}"](
            model_wrap_cfg,
            x,
            sigmas,
            extra_args={
                "cond": conditioning,
                "uncond": unconditional_conditioning,
                "cond_scale": unconditional_guidance_scale,
            },
            disable=False,
            callback=partial(KDiffusionSampler.img_callback_wrapper, img_callback),
        )

        return samples_ddim, None

    @classmethod
    def img_callback_wrapper(cls, callback: Callable, *args):
        """Converts a KDiffusion callback to the standard img_callback"""
        if callback:
            arg_dict = args[0]
            callback(image_sample=arg_dict["denoised"], iter_num=arg_dict["i"])


def create_random_tensors(shape, seeds):