Merge pull request #12371 from AUTOMATIC1111/refiner

        self.uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, self.steps * self.step_multiplier, [self.cached_uc], self.extra_network_data)

        self.c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, self.steps * self.step_multiplier, [self.cached_c], self.extra_network_data)

    def get_conds(self):

        return self.c, self.uc

    def parse_extra_network_prompts(self):

        self.prompts, self.extra_network_data = extra_networks.parse_prompts(self.prompts)

    stored_opts = {k: opts.data[k] for k in p.override_settings.keys()}

    try:

        # after running refiner, the refiner model is not unloaded - webui swaps back to main model here

        if shared.sd_model.sd_checkpoint_info.title != opts.sd_model_checkpoint:

            sd_models.reload_model_weights()

        # if no checkpoint override or the override checkpoint can't be found, remove override entry and load opts checkpoint

        if sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None:

            p.override_settings.pop('sd_model_checkpoint', None)

            if state.interrupted:

                break

            sd_models.reload_model_weights()  # model can be changed for example by refiner

            p.prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]

            p.negative_prompts = p.all_negative_prompts[n * p.batch_size:(n + 1) * p.batch_size]

            p.seeds = p.all_seeds[n * p.batch_size:(n + 1) * p.batch_size]

                with devices.autocast():

                    extra_networks.activate(self, self.extra_network_data)

    def get_conds(self):

        if self.is_hr_pass:

            return self.hr_c, self.hr_uc

        return super().get_conds()

    def parse_extra_network_prompts(self):

        res = super().parse_extra_network_prompts()

    return res

class SkipWritingToConfig:

    """This context manager prevents load_model_weights from writing checkpoint name to the config when it loads weight."""

    skip = False

    previous = None

    def __enter__(self):

        self.previous = SkipWritingToConfig.skip

        SkipWritingToConfig.skip = True

        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):

        SkipWritingToConfig.skip = self.previous

def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer):

    sd_model_hash = checkpoint_info.calculate_shorthash()

    timer.record("calculate hash")

    if not SkipWritingToConfig.skip:

        shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title

    if state_dict is None:

        timer.record("send model to device")

        model_data.set_sd_model(already_loaded)

        if not SkipWritingToConfig.skip:

            shared.opts.data["sd_model_checkpoint"] = already_loaded.sd_checkpoint_info.title

            shared.opts.data["sd_checkpoint_hash"] = already_loaded.sd_checkpoint_info.sha256

        print(f"Using already loaded model {already_loaded.sd_checkpoint_info.title}: done in {timer.summary()}")

        return model_data.sd_model

    elif shared.opts.sd_checkpoints_limit > 1 and len(model_data.loaded_sd_models) < shared.opts.sd_checkpoints_limit:

    negative prompt.

    """

    def __init__(self, model, sampler):

    def __init__(self, sampler):

        super().__init__()

        self.inner_model = model

        self.model_wrap = None

        self.mask = None

        self.nmask = None

        self.init_latent = None

        self.steps = None

        self.step = 0

        self.image_cfg_scale = None

        self.padded_cond_uncond = False

        self.sampler = sampler

        self.model_wrap = None

        self.p = None

    @property

    def inner_model(self):

        raise NotImplementedError()

    def combine_denoised(self, x_out, conds_list, uncond, cond_scale):

        denoised_uncond = x_out[-uncond.shape[0]:]

    def get_pred_x0(self, x_in, x_out, sigma):

        return x_out

    def update_inner_model(self):

        self.model_wrap = None

        c, uc = self.p.get_conds()

        self.sampler.sampler_extra_args['cond'] = c

        self.sampler.sampler_extra_args['uncond'] = uc

    def forward(self, x, sigma, uncond, cond, cond_scale, s_min_uncond, image_cond):

        if state.interrupted or state.skipped:

            raise sd_samplers_common.InterruptedException

        if sd_samplers_common.apply_refiner(self):

            cond = self.sampler.sampler_extra_args['cond']

            uncond = self.sampler.sampler_extra_args['uncond']

        # at self.image_cfg_scale == 1.0 produced results for edit model are the same as with normal sampling,

        # so is_edit_model is set to False to support AND composition.

        is_edit_model = shared.sd_model.cond_stage_key == "edit" and self.image_cfg_scale is not None and self.image_cfg_scale != 1.0

import numpy as np

import torch

from PIL import Image

from modules import devices, images, sd_vae_approx, sd_samplers, sd_vae_taesd, shared

from modules import devices, images, sd_vae_approx, sd_samplers, sd_vae_taesd, shared, sd_models

from modules.shared import opts, state

import k_diffusion.sampling

replace_torchsde_browinan()

def apply_refiner(sampler):

    completed_ratio = sampler.step / sampler.steps

    if completed_ratio <= shared.opts.sd_refiner_switch_at:

        return False

    if shared.opts.sd_refiner_checkpoint == "None":

        return False

    if shared.sd_model.sd_checkpoint_info.title == shared.opts.sd_refiner_checkpoint:

        return False

    refiner_checkpoint_info = sd_models.get_closet_checkpoint_match(shared.opts.sd_refiner_checkpoint)

    if refiner_checkpoint_info is None:

        raise Exception(f'Could not find checkpoint with name {shared.opts.sd_refiner_checkpoint}')

    sampler.p.extra_generation_params['Refiner'] = refiner_checkpoint_info.short_title

    sampler.p.extra_generation_params['Refiner switch at'] = shared.opts.sd_refiner_switch_at

    with sd_models.SkipWritingToConfig():

        sd_models.reload_model_weights(info=refiner_checkpoint_info)

    devices.torch_gc()

    sampler.p.setup_conds()

    sampler.update_inner_model()

    return True

class TorchHijack:

    """This is here to replace torch.randn_like of k-diffusion.

        self.conditioning_key = shared.sd_model.model.conditioning_key

        self.model_wrap = None

        self.p = None

        self.model_wrap_cfg = None

        self.sampler_extra_args = None

    def callback_state(self, d):

        step = d['i']

        shared.total_tqdm.update()

    def launch_sampling(self, steps, func):

        self.model_wrap_cfg.steps = steps

        state.sampling_steps = steps

        state.sampling_step = 0

        return p.steps

    def initialize(self, p) -> dict:

        self.p = p

        self.model_wrap_cfg.p = p

        self.model_wrap_cfg.mask = p.mask if hasattr(p, 'mask') else None

        self.model_wrap_cfg.nmask = p.nmask if hasattr(p, 'nmask') else None

        self.model_wrap_cfg.step = 0