Hiresfix checkpoint (#2100)

restores checkpoint selection for highres fix

modules/processing.py

@@ -1383,7 +1383,14 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                 decoded_samples = None
 
         with sd_models.SkipWritingToConfig():
-            sd_models.reload_model_weights(info=self.hr_checkpoint_info)
+            if self.hr_checkpoint_name and self.hr_checkpoint_name != 'Use same checkpoint':
+                firstpass_checkpoint = getattr(shared.opts, 'sd_model_checkpoint')
+                if firstpass_checkpoint != self.hr_checkpoint_name:
+                    try:
+                        main_entry.checkpoint_change(self.hr_checkpoint_name, save=False)
+                        sd_models.forge_model_reload();
+                    finally:
+                        main_entry.checkpoint_change(firstpass_checkpoint, save=False)
 
         return self.sample_hr_pass(samples, decoded_samples, seeds, subseeds, subseed_strength, prompts)
 
@@ -1543,7 +1550,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         steps = self.hr_second_pass_steps or self.steps
         total_steps = sampler_config.total_steps(steps) if sampler_config else steps
 
-        if self.cfg_scale == 1:
+        if self.hr_cfg == 1:
             self.hr_uc = None
             print('Skipping unconditional conditioning (HR pass) when CFG = 1. Negative Prompts are ignored.')
         else:

modules/txt2img.py

@@ -92,6 +92,7 @@ def txt2img_upscale_function(id_task: str, request: gr.Request, gallery, gallery
     p.extra_generation_params['Original Size'] = f'{args[8]}x{args[7]}'
 
     p.override_settings['save_images_before_highres_fix'] = False
+    p.override_settings['sd_model_checkpoint'] = p.hr_checkpoint_name
 
     with closing(p):
         processed = modules.scripts.scripts_txt2img.run(p, *p.script_args)

modules/ui.py

@@ -348,9 +348,8 @@ def create_ui():
                                     hr_cfg = gr.Slider(minimum=1.0, maximum=30.0, step=0.1, label="Hires CFG Scale", value=7.0, elem_id="txt2img_hr_cfg")
    
                                 with FormRow(elem_id="txt2img_hires_fix_row3", variant="compact", visible=shared.opts.hires_fix_show_sampler) as hr_sampler_container:
-
-                                    hr_checkpoint_name = gr.Dropdown(label='Checkpoint', elem_id="hr_checkpoint", choices=["Use same checkpoint"], value="Use same checkpoint", visible=False, interactive=False)
-                                    # create_refresh_button(hr_checkpoint_name, modules.sd_models.list_models, lambda: {"choices": ["Use same checkpoint"] + modules.sd_models.checkpoint_tiles(use_short=True)}, "hr_checkpoint_refresh")
+                                    hr_checkpoint_name = gr.Dropdown(label='Checkpoint', elem_id="hr_checkpoint", choices=["Use same checkpoint"] + modules.sd_models.checkpoint_tiles(use_short=True), value="Use same checkpoint")
+                                    create_refresh_button(hr_checkpoint_name, modules.sd_models.list_models, lambda: {"choices": ["Use same checkpoint"] + modules.sd_models.checkpoint_tiles(use_short=True)}, "hr_checkpoint_refresh")
 
                                     hr_sampler_name = gr.Dropdown(label='Hires sampling method', elem_id="hr_sampler", choices=["Use same sampler"] + sd_samplers.visible_sampler_names(), value="Use same sampler")
                                     hr_scheduler = gr.Dropdown(label='Hires schedule type', elem_id="hr_scheduler", choices=["Use same scheduler"] + [x.label for x in sd_schedulers.schedulers], value="Use same scheduler")