fix refiner reload/unload

modules/processing_diffusers.py

@@ -27,7 +27,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
 
     def hires_resize(latents): # input=latents output=pil
         latent_upscaler = shared.latent_upscale_modes.get(p.hr_upscaler, None)
-        shared.log.info(f'Diffusers Hires: upscaler={p.hr_upscaler} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}')
+        shared.log.info(f'Hires: upscaler={p.hr_upscaler} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}')
         if latent_upscaler is not None:
             latents = torch.nn.functional.interpolate(latents, size=(p.hr_upscale_to_y // 8, p.hr_upscale_to_x // 8), mode=latent_upscaler["mode"], antialias=latent_upscaler["antialias"])
         first_pass_images = vae_decode(latents=latents, model=shared.sd_model, full_quality=True, output_type='pil')
@@ -54,9 +54,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         shared.state.current_latent = latents
 
     def full_vae_decode(latents, model):
-        shared.log.debug(f'Diffusers VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)} images={latents.shape[0]}')
+        shared.log.debug(f'VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)} images={latents.shape[0]}')
         if shared.opts.diffusers_move_unet and not model.has_accelerate:
-            shared.log.debug('Diffusers: Moving UNet to CPU')
+            shared.log.debug('Moving to CPU: model=UNet')
             unet_device = model.unet.device
             model.unet.to(devices.cpu)
             devices.torch_gc()
@@ -69,7 +69,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         return decoded
 
     def taesd_vae_decode(latents):
-        shared.log.debug(f'Diffusers VAE decode: name=TAESD images={latents.shape[0]}')
+        shared.log.debug(f'VAE decode: name=TAESD images={latents.shape[0]}')
         decoded = torch.zeros((len(latents), 3, p.height, p.width), dtype=devices.dtype_vae, device=devices.device)
         for i in range(len(output.images)):
             decoded[i] = (sd_vae_taesd.decode(latents[i]) * 2.0) - 1.0
@@ -181,13 +181,13 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         if 'negative_prompt' in clean:
             clean['negative_prompt'] = len(clean['negative_prompt'])
         if 'prompt_embeds' in clean:
-            clean['prompt_embeds'] = clean['prompt_embeds'].shape
+            clean['prompt_embeds'] = clean['prompt_embeds'].shape if torch.is_tensor(clean['prompt_embeds']) else type(clean['prompt_embeds'])
         if 'pooled_prompt_embeds' in clean:
-            clean['pooled_prompt_embeds'] = clean['pooled_prompt_embeds'].shape
+            clean['pooled_prompt_embeds'] = clean['pooled_prompt_embeds'].shape if torch.is_tensor(clean['pooled_prompt_embeds']) else type(clean['pooled_prompt_embeds'])
         if 'negative_prompt_embeds' in clean:
-            clean['negative_prompt_embeds'] = clean['negative_prompt_embeds'].shape
+            clean['negative_prompt_embeds'] = clean['negative_prompt_embeds'].shape if torch.is_tensor(clean['negative_prompt_embeds']) else type(clean['negative_prompt_embeds'])
         if 'negative_pooled_prompt_embeds' in clean:
-            clean['negative_pooled_prompt_embeds'] = clean['negative_pooled_prompt_embeds'].shape
+            clean['negative_pooled_prompt_embeds'] = clean['negative_pooled_prompt_embeds'].shape if torch.is_tensor(clean['negative_pooled_prompt_embeds']) else type(clean['negative_pooled_prompt_embeds'])
         clean['generator'] = generator_device
         shared.log.debug(f'Diffuser pipeline: {pipeline.__class__.__name__} task={sd_models.get_diffusers_task(model)} set={clean}')
         return args
@@ -318,7 +318,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'):
             save_intermediate(latents=output.images, suffix="-before-refiner")
         if shared.opts.diffusers_move_base and not shared.sd_model.has_accelerate:
-            shared.log.debug('Diffusers: Moving base model to CPU')
+            shared.log.debug('Moving to CPU: model=base')
             shared.sd_model.to(devices.cpu)
             devices.torch_gc()
 
@@ -363,7 +363,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                     results.append(refiner_image)
 
         if shared.opts.diffusers_move_refiner and not shared.sd_refiner.has_accelerate:
-            shared.log.debug('Diffusers: Moving refiner model to CPU')
+            shared.log.debug('Moving to CPU: model=refiner')
             shared.sd_refiner.to(devices.cpu)
             devices.torch_gc()
 

modules/prompt_parser_diffusers.py

@@ -57,10 +57,13 @@ def compel_encode_prompts(
         negative_embeds.append(negative_embed)
         negative_pooleds.append(negative_pooled)
 
-    prompt_embeds = torch.cat(prompt_embeds, dim=0)
-    negative_embeds = torch.cat(negative_embeds, dim=0)
-    if shared.sd_model_type == "sdxl":
+    if prompt_embeds is not None:
+        prompt_embeds = torch.cat(prompt_embeds, dim=0)
+    if negative_embeds is not None:
+        negative_embeds = torch.cat(negative_embeds, dim=0)
+    if positive_pooleds is not None and shared.sd_model_type == "sdxl":
         positive_pooleds = torch.cat(positive_pooleds, dim=0)
+    if negative_pooleds is not None and shared.sd_model_type == "sdxl":
         negative_pooleds = torch.cat(negative_pooleds, dim=0)
     return prompt_embeds, positive_pooleds, negative_embeds, negative_pooleds
 

modules/sd_models.py

@@ -35,7 +35,6 @@ model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
 checkpoints_list = {}
 checkpoint_aliases = {}
 checkpoints_loaded = collections.OrderedDict()
-skip_next_load = False
 sd_metadata_file = os.path.join(paths.data_path, "metadata.json")
 sd_metadata = None
 sd_metadata_pending = 0
@@ -512,7 +511,7 @@ class ModelData:
         self.sd_model = v
 
     def get_sd_refiner(self):
-        if self.sd_model is None:
+        if self.sd_refiner is None:
             with self.lock:
                 try:
                     if shared.backend == shared.Backend.ORIGINAL:
@@ -568,9 +567,9 @@ def detect_pipeline(f: str, op: str = 'model'):
                 else:
                     guess = 'Stable Diffusion XL'
             else:
-                shared.log.error(f'Diffusers autodetect failed, set diffuser pipeline manually: {f}')
+                shared.log.error(f'Model autodetect failed, set diffuser pipeline manually: {f}')
                 return None, None
-            shared.log.debug(f'Diffusers autodetect {op}: {f} pipeline={guess} size={size} GB')
+            shared.log.debug(f'Model autodetect {op}: {f} pipeline={guess} size={size} GB')
         except Exception as e:
             shared.log.error(f'Error detecting diffusers pipeline: model={f} {e}')
             return None, None
@@ -618,7 +617,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         "safety_checker": None,
         "requires_safety_checker": False,
         "load_safety_checker": False,
-        "load_connected_pipeline": True  # always load end-to-end / connected pipelines
+        "load_connected_pipeline": True,
         # "use_safetensors": True,  # TODO(PVP) - we can't enable this for all checkpoints just yet
     }
     if shared.opts.diffusers_model_load_variant == 'default':
@@ -646,58 +645,57 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             ckpt_basename = os.path.basename(shared.cmd_opts.ckpt)
             model_name = modelloader.find_diffuser(ckpt_basename)
             if model_name is not None:
-                shared.log.info(f'Loading diffuser {op}: {model_name}')
+                shared.log.info(f'Loading model {op}: {model_name}')
                 model_file = modelloader.download_diffusers_model(hub_id=model_name)
                 try:
-                    shared.log.debug(f'Diffusers load {op} config: {diffusers_load_config}')
+                    shared.log.debug(f'Model load {op} config: {diffusers_load_config}')
                     sd_model = diffusers.DiffusionPipeline.from_pretrained(model_file, **diffusers_load_config)
                 except Exception as e:
-                    shared.log.error(f'Diffusers failed loading model: {model_file} {e}')
+                    shared.log.error(f'Failed loading model: {model_file} {e}')
                 list_models() # rescan for downloaded model
                 checkpoint_info = CheckpointInfo(model_name)
 
-        if sd_model is None:
-            checkpoint_info = checkpoint_info or select_checkpoint(op=op)
-            if checkpoint_info is None:
-                unload_model_weights(op=op)
+        checkpoint_info = checkpoint_info or select_checkpoint(op=op)
+        if checkpoint_info is None:
+            unload_model_weights(op=op)
+            return
+
+        vae = None
+        sd_vae.loaded_vae_file = None
+        if op == 'model' or op == 'refiner':
+            vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)
+            vae = sd_vae.load_vae_diffusers(checkpoint_info.path, vae_file, vae_source)
+            if vae is not None:
+                diffusers_load_config["vae"] = vae
+
+        shared.log.info(f'Loading diffuser {op}: {checkpoint_info.filename}')
+        if not os.path.isfile(checkpoint_info.path):
+            try:
+                # shared.log.debug(f'Diffusers load {op} config: {diffusers_load_config}')
+                sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, **diffusers_load_config)
+            except Exception as e:
+                shared.log.error(f'Failed loading model {op}: {checkpoint_info.path} {e}')
+        else:
+            diffusers_load_config["local_files_only "] = True
+            diffusers_load_config["extract_ema"] = shared.opts.diffusers_extract_ema
+            pipeline, _model_type = detect_pipeline(checkpoint_info.path, op)
+            if pipeline is None:
+                shared.log.error(f'Diffusers {op} pipeline not initialized: {shared.opts.diffusers_pipeline}')
                 return
-
-            vae = None
-            sd_vae.loaded_vae_file = None
-            if op == 'model' or op == 'refiner':
-                vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename)
-                vae = sd_vae.load_vae_diffusers(checkpoint_info.path, vae_file, vae_source)
-                if vae is not None:
-                    diffusers_load_config["vae"] = vae
-
-            shared.log.info(f'Loading diffuser {op}: {checkpoint_info.filename}')
-            if not os.path.isfile(checkpoint_info.path):
-                try:
-                    # shared.log.debug(f'Diffusers load {op} config: {diffusers_load_config}')
-                    sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, **diffusers_load_config)
-                except Exception as e:
-                    shared.log.error(f'Diffusers {op} failed loading model: {checkpoint_info.path} {e}')
-            else:
-                diffusers_load_config["local_files_only "] = True
-                diffusers_load_config["extract_ema"] = shared.opts.diffusers_extract_ema
-                pipeline, _model_type = detect_pipeline(checkpoint_info.path, op)
-                if pipeline is None:
-                    shared.log.error(f'Diffusers {op} pipeline not initialized: {shared.opts.diffusers_pipeline}')
-                    return
-                try:
-                    if hasattr(pipeline, 'from_single_file'):
-                        diffusers_load_config['use_safetensors'] = True
-                        sd_model = pipeline.from_single_file(checkpoint_info.path, **diffusers_load_config)
-                    elif hasattr(pipeline, 'from_ckpt'):
-                        sd_model = pipeline.from_ckpt(checkpoint_info.path, **diffusers_load_config)
-                    else:
-                        shared.log.error(f'Diffusers {op} cannot load safetensor model: {checkpoint_info.path} {shared.opts.diffusers_pipeline}')
-                        return
-                    if sd_model is not None:
-                        shared.log.debug(f'Diffusers {op}: pipeline={sd_model.__class__.__name__}') # pylint: disable=protected-access
-                except Exception as e:
-                    shared.log.error(f'Diffusers failed loading model using pipeline: {checkpoint_info.path} {shared.opts.diffusers_pipeline} {e}')
+            try:
+                if hasattr(pipeline, 'from_single_file'):
+                    diffusers_load_config['use_safetensors'] = True
+                    sd_model = pipeline.from_single_file(checkpoint_info.path, **diffusers_load_config)
+                elif hasattr(pipeline, 'from_ckpt'):
+                    sd_model = pipeline.from_ckpt(checkpoint_info.path, **diffusers_load_config)
+                else:
+                    shared.log.error(f'Diffusers {op} cannot load safetensor model: {checkpoint_info.path} {shared.opts.diffusers_pipeline}')
                     return
+                if sd_model is not None:
+                    shared.log.debug(f'Model {op}: pipeline={sd_model.__class__.__name__}') # pylint: disable=protected-access
+            except Exception as e:
+                shared.log.error(f'Diffusers failed loading model using pipeline: {checkpoint_info.path} {shared.opts.diffusers_pipeline} {e}')
+                return
 
         if "StableDiffusion" in sd_model.__class__.__name__:
             pass # scheduler is created on first use
@@ -705,8 +703,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             sd_model.scheduler.name = 'DDIM'
 
         if (shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) and (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram):
-            shared.log.warning(f'Diffusers {op}: Model CPU offload (--medvram) and Sequential CPU offload (--lowvram) are not compatible')
-            shared.log.debug(f'Diffusers {op}: disabling model CPU offload and --medvram')
+            shared.log.warning(f'Model {op}: Model CPU offload (--medvram) and Sequential CPU offload (--lowvram) are not compatible')
+            shared.log.debug(f'Model {op}: disabling model CPU offload and --medvram')
             shared.opts.diffusers_model_cpu_offload=False
             shared.cmd_opts.medvram=False
 
@@ -715,7 +713,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         sd_model.has_accelerate = False
         if hasattr(sd_model, "enable_model_cpu_offload"):
             if (shared.cmd_opts.medvram and devices.backend != "directml") or shared.opts.diffusers_model_cpu_offload:
-                shared.log.debug(f'Diffusers {op}: enable model CPU offload')
+                shared.log.debug(f'Model {op}: enable model CPU offload')
                 if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
                     shared.opts.diffusers_move_base = False
                     shared.opts.diffusers_move_unet = False
@@ -725,7 +723,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                 sd_model.has_accelerate = True
         if hasattr(sd_model, "enable_sequential_cpu_offload"):
             if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload:
-                shared.log.debug(f'Diffusers {op}: enable sequential CPU offload')
+                shared.log.debug(f'Model {op}: enable sequential CPU offload')
                 if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
                     shared.opts.diffusers_move_base = False
                     shared.opts.diffusers_move_unet = False
@@ -735,19 +733,19 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                 sd_model.has_accelerate = True
         if hasattr(sd_model, "enable_vae_slicing"):
             if shared.cmd_opts.lowvram or shared.opts.diffusers_vae_slicing:
-                shared.log.debug(f'Diffusers {op}: enable VAE slicing')
+                shared.log.debug(f'Model {op}: enable VAE slicing')
                 sd_model.enable_vae_slicing()
             else:
                 sd_model.disable_vae_slicing()
         if hasattr(sd_model, "enable_vae_tiling"):
             if shared.cmd_opts.lowvram or shared.opts.diffusers_vae_tiling:
-                shared.log.debug(f'Diffusers {op}: enable VAE tiling')
+                shared.log.debug(f'Model {op}: enable VAE tiling')
                 sd_model.enable_vae_tiling()
             else:
                 sd_model.disable_vae_tiling()
         if hasattr(sd_model, "enable_attention_slicing"):
             if shared.cmd_opts.lowvram or shared.opts.diffusers_attention_slicing:
-                shared.log.debug(f'Diffusers {op}: enable attention slicing')
+                shared.log.debug(f'Model {op}: enable attention slicing')
                 sd_model.enable_attention_slicing()
             else:
                 sd_model.disable_attention_slicing()
@@ -761,11 +759,11 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                 else:
                     sd_model.vae.config["force_upcast"] = False
                     sd_model.vae.config.force_upcast = False
-            shared.log.debug(f'Diffusers {op} VAE: name={sd_vae.loaded_vae_file} upcast={sd_model.vae.config.get("force_upcast", None)}')
+            shared.log.debug(f'Model {op} VAE: name={sd_vae.loaded_vae_file} upcast={sd_model.vae.config.get("force_upcast", None)}')
         if shared.opts.cross_attention_optimization == "xFormers" and hasattr(sd_model, 'enable_xformers_memory_efficient_attention'):
             sd_model.enable_xformers_memory_efficient_attention()
         if shared.opts.opt_channelslast:
-            shared.log.debug(f'Diffusers {op}: enable channels last')
+            shared.log.debug(f'Model {op}: enable channels last')
             sd_model.unet.to(memory_format=torch.channels_last)
 
         base_sent_to_cpu=False
@@ -806,7 +804,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     shared.log.info(f"Compiling pipeline={sd_model.__class__.__name__} shape={8 * sd_model.unet.config.sample_size} mode={shared.opts.cuda_compile_backend}")
                     import torch._dynamo # pylint: disable=unused-import,redefined-outer-name
                     if shared.opts.cuda_compile_backend == "openvino_fx":
-                        torch._dynamo.reset()
+                        torch._dynamo.reset() # pylint: disable=protected-access
                         from modules.intel.openvino import openvino_fx, openvino_clear_caches, ModelState # pylint: disable=unused-import
                         openvino_clear_caches()
                         sd_model.compiled_model_state = ModelState()
@@ -995,11 +993,6 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None,
 
 def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model'):
     load_dict = shared.opts.sd_model_dict != model_data.sd_dict
-    global skip_next_load # pylint: disable=global-statement
-    if skip_next_load:
-        shared.log.debug('Load model weights skip')
-        skip_next_load = False
-        return
     from modules import lowvram, sd_hijack
     checkpoint_info = info or select_checkpoint(op=op) # are we selecting model or dictionary
     next_checkpoint_info = info or select_checkpoint(op='dict' if load_dict else 'model') if load_dict else None

modules/shared.py

@@ -813,7 +813,7 @@ else:
 opts.data['sd_backend'] = 'diffusers' if backend == Backend.DIFFUSERS else 'original'
 opts.data['uni_pc_lower_order_final'] = opts.schedulers_use_loworder
 opts.data['uni_pc_order'] = opts.schedulers_solver_order
-log.info(f'Pipeline: {backend}')
+log.info(f'Engine: backend={backend}')
 
 
 prompt_styles = modules.styles.StyleDatabase(opts.styles_dir)

modules/ui.py

@@ -88,12 +88,12 @@ def add_style(name: str, prompt: str, negative_prompt: str):
     return [gr.Dropdown.update(visible=True, choices=list(modules.shared.prompt_styles.styles)) for _ in range(2)]
 
 
-def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y):
+def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y, hr_upscaler):
     from modules import processing, devices
     if not enable:
         return ""
-    # if modules.shared.backend == modules.shared.Backend.DIFFUSERS:
-    #   return "Hires resize: disabled"
+    if hr_upscaler == "None":
+        return "Hires resize: None"
     p = processing.StableDiffusionProcessingTxt2Img(width=width, height=height, enable_hr=True, hr_scale=hr_scale, hr_resize_x=hr_resize_x, hr_resize_y=hr_resize_y)
     p.init_hr()
     with devices.autocast():
@@ -105,9 +105,7 @@ def resize_from_to_html(width, height, scale_by):
     target_width = int(width * scale_by)
     target_height = int(height * scale_by)
     if not target_width or not target_height:
-        return "no image selected"
-    # if modules.shared.backend == modules.shared.Backend.DIFFUSERS:
-    #   return "Hires resize: disabled"
+        return "Hires resize: no image selected"
     return f"Hires resize: from <span class='resolution'>{width}x{height}</span> to <span class='resolution'>{target_width}x{target_height}</span>"
 
 
@@ -413,7 +411,7 @@ def create_ui(startup_timer = None):
                 with FormGroup(elem_id="txt2img_script_container"):
                     custom_inputs = modules.scripts.scripts_txt2img.setup_ui()
 
-            hr_resolution_preview_inputs = [show_second_pass, width, height, hr_scale, hr_resize_x, hr_resize_y]
+            hr_resolution_preview_inputs = [show_second_pass, width, height, hr_scale, hr_resize_x, hr_resize_y, hr_upscaler]
             for preview_input in hr_resolution_preview_inputs:
                 preview_input.change(
                     fn=calc_resolution_hires,
@@ -460,7 +458,7 @@ def create_ui(startup_timer = None):
             submit.click(**txt2img_args)
 
             def enable_hr_change(visible: bool):
-                return {"visible": visible, "__type__": "update"}, f'Refiner{": disabled" if modules.shared.sd_refiner is None else ""}'
+                return {"visible": visible, "__type__": "update"}, f'Refiner: {"disabled" if modules.shared.opts.sd_model_refiner == "None" else "enabled"}'
 
             res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)
             batch_switch_btn.click(lambda w, h: (h, w), inputs=[batch_count, batch_size], outputs=[batch_count, batch_size], show_progress=False)