hires fix checkpoint selection

    if "Hires sampler" not in res:

        res["Hires sampler"] = "Use same sampler"

    if "Hires checkpoint" not in res:

        res["Hires checkpoint"] = "Use same checkpoint"

    if "Hires prompt" not in res:

        res["Hires prompt"] = ""

    cached_hr_uc = [None, None]

    cached_hr_c = [None, None]

    def __init__(self, enable_hr: bool = False, denoising_strength: float = 0.75, firstphase_width: int = 0, firstphase_height: int = 0, hr_scale: float = 2.0, hr_upscaler: str = None, hr_second_pass_steps: int = 0, hr_resize_x: int = 0, hr_resize_y: int = 0, hr_sampler_name: str = None, hr_prompt: str = '', hr_negative_prompt: str = '', **kwargs):

    def __init__(self, enable_hr: bool = False, denoising_strength: float = 0.75, firstphase_width: int = 0, firstphase_height: int = 0, hr_scale: float = 2.0, hr_upscaler: str = None, hr_second_pass_steps: int = 0, hr_resize_x: int = 0, hr_resize_y: int = 0, hr_checkpoint_name: str = None, hr_sampler_name: str = None, hr_prompt: str = '', hr_negative_prompt: str = '', **kwargs):

        super().__init__(**kwargs)

        self.enable_hr = enable_hr

        self.denoising_strength = denoising_strength

        self.hr_resize_y = hr_resize_y

        self.hr_upscale_to_x = hr_resize_x

        self.hr_upscale_to_y = hr_resize_y

        self.hr_checkpoint_name = hr_checkpoint_name

        self.hr_checkpoint_info = None

        self.hr_sampler_name = hr_sampler_name

        self.hr_prompt = hr_prompt

        self.hr_negative_prompt = hr_negative_prompt

        self.all_hr_prompts = None

        self.all_hr_negative_prompts = None

        self.latent_scale_mode = None

        if firstphase_width != 0 or firstphase_height != 0:

            self.hr_upscale_to_x = self.width

    def init(self, all_prompts, all_seeds, all_subseeds):

        if self.enable_hr:

            if self.hr_checkpoint_name:

                self.hr_checkpoint_info = sd_models.get_closet_checkpoint_match(self.hr_checkpoint_name)

                if self.hr_checkpoint_info is None:

                    raise Exception(f'Could not find checkpoint with name {self.hr_checkpoint_name}')

                self.extra_generation_params["Hires checkpoint"] = self.hr_checkpoint_info.short_title

            if self.hr_sampler_name is not None and self.hr_sampler_name != self.sampler_name:

                self.extra_generation_params["Hires sampler"] = self.hr_sampler_name

            if tuple(self.hr_negative_prompt) != tuple(self.negative_prompt):

                self.extra_generation_params["Hires negative prompt"] = self.hr_negative_prompt

            self.latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest")

            if self.enable_hr and self.latent_scale_mode is None:

                if not any(x.name == self.hr_upscaler for x in shared.sd_upscalers):

                    raise Exception(f"could not find upscaler named {self.hr_upscaler}")

            if opts.use_old_hires_fix_width_height and self.applied_old_hires_behavior_to != (self.width, self.height):

                self.hr_resize_x = self.width

                self.hr_resize_y = self.height

                    self.truncate_x = (self.hr_upscale_to_x - target_w) // opt_f

                    self.truncate_y = (self.hr_upscale_to_y - target_h) // opt_f

            # special case: the user has chosen to do nothing

            if self.hr_upscale_to_x == self.width and self.hr_upscale_to_y == self.height:

                self.enable_hr = False

                self.denoising_strength = None

                self.extra_generation_params.pop("Hires upscale", None)

                self.extra_generation_params.pop("Hires resize", None)

                return

            if not state.processing_has_refined_job_count:

                if state.job_count == -1:

                    state.job_count = self.n_iter

    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):

        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)

        latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest")

        if self.enable_hr and latent_scale_mode is None:

            if not any(x.name == self.hr_upscaler for x in shared.sd_upscalers):

                raise Exception(f"could not find upscaler named {self.hr_upscaler}")

        x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

        samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))

        if not self.enable_hr:

            return samples

        current = shared.sd_model.sd_checkpoint_info

        try:

            if self.hr_checkpoint_info is not None:

                sd_models.reload_model_weights(info=self.hr_checkpoint_info)

            return self.sample_hr_pass(samples, seeds, subseeds, subseed_strength, prompts)

        finally:

            sd_models.reload_model_weights(info=current)

    def sample_hr_pass(self, samples, seeds, subseeds, subseed_strength, prompts):

        self.is_hr_pass = True

        target_width = self.hr_upscale_to_x

            info = create_infotext(self, self.all_prompts, self.all_seeds, self.all_subseeds, [], iteration=self.iteration, position_in_batch=index)

            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, info=info, p=self, suffix="-before-highres-fix")

        if latent_scale_mode is not None:

        if self.latent_scale_mode is not None:

            for i in range(samples.shape[0]):

                save_intermediate(samples, i)

            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"])

            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode=self.latent_scale_mode["mode"], antialias=self.latent_scale_mode["antialias"])

            # Avoid making the inpainting conditioning unless necessary as

            # this does need some extra compute to decode / encode the image again.

        self.hr_uc = None

        self.hr_c = None

        if self.enable_hr:

        if self.enable_hr and self.hr_checkpoint_info is None:

            if shared.opts.hires_fix_use_firstpass_conds:

                self.calculate_hr_conds()

        self.shorthash = self.sha256[0:10] if self.sha256 else None

        self.title = name if self.shorthash is None else f'{name} [{self.shorthash}]'

        self.short_title = self.name_for_extra if self.shorthash is None else f'{self.name_for_extra} [{self.shorthash}]'

        self.ids = [self.hash, self.model_name, self.title, name, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] if self.shorthash else [])

        checkpoints_list.pop(self.title)

        self.title = f'{self.name} [{self.shorthash}]'

        self.short_title = f'{self.name_for_extra} [{self.shorthash}]'

        self.register()

        return self.shorthash

    enable_midas_autodownload()

def checkpoint_tiles():

    def convert(name):

        return int(name) if name.isdigit() else name.lower()

    def alphanumeric_key(key):

        return [convert(c) for c in re.split('([0-9]+)', key)]

    return sorted([x.title for x in checkpoints_list.values()], key=alphanumeric_key)

def checkpoint_tiles(use_short=False):

    return [x.short_title if use_short else x.title for x in checkpoints_list.values()]

def list_models():

    elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:

        print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)

    for filename in sorted(model_list, key=str.lower):

    for filename in model_list:

        checkpoint_info = CheckpointInfo(filename)

        checkpoint_info.register()

re_strip_checksum = re.compile(r"\s*\[[^]]+]\s*$")

def get_closet_checkpoint_match(search_string):

    checkpoint_info = checkpoint_aliases.get(search_string, None)

    if checkpoint_info is not None:

    if found:

        return found[0]

    search_string_without_checksum = re.sub(re_strip_checksum, '', search_string)

    found = sorted([info for info in checkpoints_list.values() if search_string_without_checksum in info.title], key=lambda x: len(x.title))

    if found:

        return found[0]

    return None

            return

        import modules.sd_samplers

        try:

            if opts.show_progress_grid:

                self.assign_current_image(modules.sd_samplers.samples_to_image_grid(self.current_latent))

            else:

            self.current_image_sampling_step = self.sampling_step

        except Exception:

            # when switching models during genration, VAE would be on CPU, so creating an image will fail.

            # we silently ignore this error

            errors.record_exception()

    def assign_current_image(self, image):

        self.current_image = image

        self.id_live_preview += 1

    "ui_tab_order": OptionInfo([], "UI tab order", ui_components.DropdownMulti, lambda: {"choices": list(tab_names)}).needs_restart(),

    "hidden_tabs": OptionInfo([], "Hidden UI tabs", ui_components.DropdownMulti, lambda: {"choices": list(tab_names)}).needs_restart(),

    "ui_reorder_list": OptionInfo([], "txt2img/img2img UI item order", ui_components.DropdownMulti, lambda: {"choices": list(shared_items.ui_reorder_categories())}).info("selected items appear first").needs_restart(),

    "hires_fix_show_sampler": OptionInfo(False, "Hires fix: show hires sampler selection").needs_restart(),

    "hires_fix_show_sampler": OptionInfo(False, "Hires fix: show hires checkpoint and sampler selection").needs_restart(),

    "hires_fix_show_prompts": OptionInfo(False, "Hires fix: show hires prompt and negative prompt").needs_restart(),

    "disable_token_counters": OptionInfo(False, "Disable prompt token counters").needs_restart(),

}))

import gradio as gr

def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, hr_sampler_index: int, hr_prompt: str, hr_negative_prompt, override_settings_texts, request: gr.Request, *args):

def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, hr_checkpoint_name: str, hr_sampler_index: int, hr_prompt: str, hr_negative_prompt, override_settings_texts, request: gr.Request, *args):

    override_settings = create_override_settings_dict(override_settings_texts)

    p = processing.StableDiffusionProcessingTxt2Img(

        hr_second_pass_steps=hr_second_pass_steps,

        hr_resize_x=hr_resize_x,

        hr_resize_y=hr_resize_y,

        hr_checkpoint_name=None if hr_checkpoint_name == 'Use same checkpoint' else hr_checkpoint_name,

        hr_sampler_name=sd_samplers.samplers_for_img2img[hr_sampler_index - 1].name if hr_sampler_index != 0 else None,

        hr_prompt=hr_prompt,

        hr_negative_prompt=hr_negative_prompt,