add sdxl support

modules/api/api.py

@@ -406,17 +406,15 @@ class Api:
 
     def interruptapi(self):
         shared.state.interrupt()
-
         return {}
 
     def unloadapi(self):
-        unload_model_weights()
-
+        unload_model_weights(op='model')
+        unload_model_weights(op='refiner')
         return {}
 
     def reloadapi(self):
         reload_model_weights()
-
         return {}
 
     def skip(self):

modules/processing.py

@@ -223,7 +223,8 @@ class StableDiffusionProcessing:
         source_image = devices.cond_cast_float(source_image)
         # HACK: Using introspection as the Depth2Image model doesn't appear to uniquely
         # identify itself with a field common to all models. The conditioning_key is also hybrid.
-        if backend == Backend.DIFFUSERS: # TODO: Diffusers img2img_image_conditioning
+        if backend == Backend.DIFFUSERS:
+            log.warning('Diffusers not implemented: img2img_image_conditioning')
             return None
         if isinstance(self.sd_model, LatentDepth2ImageDiffusion):
             return self.depth2img_image_conditioning(source_image)
@@ -682,11 +683,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 x_samples_ddim = torch.stack(x_samples_ddim).float()
                 x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                 del samples_ddim
-                if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
-                    lowvram.send_everything_to_cpu()
-                    devices.torch_gc()
-                if p.scripts is not None:
-                    p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)
 
             elif backend == Backend.DIFFUSERS:
                 generator_device = 'cpu' if shared.opts.diffusers_generator_device == "cpu" else shared.device
@@ -695,7 +691,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     sampler = sd_samplers.all_samplers_map.get(p.sampler_name, None)
                     if sampler is None:
                         sampler = sd_samplers.all_samplers_map.get("UniPC")
-                    shared.sd_model.scheduler = sd_samplers.create_sampler(sampler.name, shared.sd_model) # TODO(Patrick): For wrapped pipelines this is currently a no-op
+                    # shared.sd_model.scheduler = sd_samplers.create_sampler(sampler.name, shared.sd_model) # TODO(Patrick): For wrapped pipelines this is currently a no-op
 
                 cross_attention_kwargs={}
                 if lora_state['active']:
@@ -708,23 +704,48 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.INPAINTING:
                     # TODO(PVP): change out to latents once possible with `diffusers`
                     task_specific_kwargs = {"image": p.init_images[0], "mask_image": p.image_mask, "strength": p.denoising_strength}
-                output = shared.sd_model(
+
+                output = shared.sd_model( # pylint: disable=not-callable
                     prompt=prompts,
                     negative_prompt=negative_prompts,
                     num_inference_steps=p.steps,
                     guidance_scale=p.cfg_scale,
                     generator=generator,
-                    output_type="np",
+                    output_type='np' if shared.sd_refiner is None else 'latent',
                     cross_attention_kwargs=cross_attention_kwargs,
                     **task_specific_kwargs
                 )
+
+                if shared.sd_refiner is not None:
+                    init_image = output.images[0]
+                    output = shared.sd_refiner( # pylint: disable=not-callable
+                        prompt=prompts,
+                        negative_prompt=negative_prompts,
+                        num_inference_steps=p.steps,
+                        guidance_scale=p.cfg_scale,
+                        generator=generator,
+                        output_type='np',
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        image=init_image
+                    )
+
                 x_samples_ddim = output.images
+
+                if p.enable_hr:
+                    log.warning('Diffusers not implemented: hires fix')
+
                 if lora_state['active']:
                     unload_diffusers_lora()
 
             else:
                 raise ValueError(f"Unknown backend {backend}")
 
+            if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+                lowvram.send_everything_to_cpu()
+                devices.torch_gc()
+            if p.scripts is not None:
+                p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)
+
             for i, x_sample in enumerate(x_samples_ddim):
                 p.batch_index = i
                 if backend == Backend.ORIGINAL:
@@ -898,7 +919,23 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             if self.hr_upscaler is not None:
                 self.extra_generation_params["Hires upscaler"] = self.hr_upscaler
 
-    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): # TODO this is majority of processing time
+    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
+
+        def save_intermediate(image, index):
+            """saves image before applying hires fix, if enabled in options; takes as an argument either an image or batch with latent space images"""
+            if not opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix:
+                return
+            if not isinstance(image, Image.Image):
+                image = sd_samplers.sample_to_image(image, index, approximation=0)
+            orig1 = self.extra_generation_params
+            orig2 = self.restore_faces
+            self.extra_generation_params = {}
+            self.restore_faces = False
+            info = create_infotext(self, self.all_prompts, self.all_seeds, self.all_subseeds, [], iteration=self.iteration, position_in_batch=index)
+            self.extra_generation_params = orig1
+            self.restore_faces = orig2
+            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, info=info, suffix="-before-highres-fix")
+
         if backend == Backend.DIFFUSERS:
             sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
 
@@ -916,21 +953,6 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         target_width = self.hr_upscale_to_x
         target_height = self.hr_upscale_to_y
 
-        def save_intermediate(image, index):
-            """saves image before applying hires fix, if enabled in options; takes as an argument either an image or batch with latent space images"""
-            if not opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix:
-                return
-            if not isinstance(image, Image.Image):
-                image = sd_samplers.sample_to_image(image, index, approximation=0)
-            orig1 = self.extra_generation_params
-            orig2 = self.restore_faces
-            self.extra_generation_params = {}
-            self.restore_faces = False
-            info = create_infotext(self, self.all_prompts, self.all_seeds, self.all_subseeds, [], iteration=self.iteration, position_in_batch=index)
-            self.extra_generation_params = orig1
-            self.restore_faces = orig2
-            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, info=info, suffix="-before-highres-fix")
-
         if latent_scale_mode is not None:
             for i in range(samples.shape[0]):
                 save_intermediate(samples, i)

modules/sd_models.py

@@ -23,6 +23,7 @@ from modules.timer import Timer
 from modules.memstats import memory_stats
 from modules.paths_internal import models_path
 
+
 transformers_logging.set_verbosity_error()
 model_dir = "Stable-diffusion"
 model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
@@ -127,7 +128,9 @@ def list_models():
     checkpoints_list.clear()
     checkpoint_aliases.clear()
     ext_filter=[".safetensors"] if shared.opts.sd_disable_ckpt else [".ckpt", ".safetensors"]
-    model_list = modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])
+    model_list = []
+    if shared.backend == shared.Backend.ORIGINAL or shared.opts.diffusers_pipeline == shared.pipelines[0]:
+        model_list += modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])
     if shared.backend == shared.Backend.DIFFUSERS:
         model_list += modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir)
 
@@ -210,11 +213,18 @@ def model_hash(filename):
         return 'NOHASH'
 
 
-def select_checkpoint(model=True):
-    model_checkpoint = shared.opts.sd_model_checkpoint if model else shared.opts.sd_model_dict
+def select_checkpoint(op='model'):
+    if op == 'model':
+        model_checkpoint = shared.opts.sd_model_checkpoint
+    elif op == 'dict':
+        model_checkpoint = shared.opts.sd_model_dict
+    elif op == 'refiner':
+        model_checkpoint = shared.opts.data['sd_model_refiner']
+    if model_checkpoint is None or model_checkpoint == 'None':
+        return None
     checkpoint_info = get_closet_checkpoint_match(model_checkpoint)
     if checkpoint_info is not None:
-        shared.log.debug(f'Select checkpoint: {checkpoint_info.title if checkpoint_info is not None else None}')
+        shared.log.debug(f'Select checkpoint: {op} {checkpoint_info.title if checkpoint_info is not None else None}')
         return checkpoint_info
     if len(checkpoints_list) == 0:
         shared.log.error("Cannot run without a checkpoint")
@@ -458,9 +468,10 @@ sd1_clip_weight = 'cond_stage_model.transformer.text_model.embeddings.token_embe
 sd2_clip_weight = 'cond_stage_model.model.transformer.resblocks.0.attn.in_proj_weight'
 
 
-class SdModelData:
+class ModelData:
     def __init__(self):
         self.sd_model = None
+        self.sd_refiner = None
         self.sd_dict = 'None'
         self.initial = True
         self.lock = threading.Lock()
@@ -470,9 +481,9 @@ class SdModelData:
             with self.lock:
                 try:
                     if shared.backend == shared.Backend.ORIGINAL:
-                        reload_model_weights()
+                        reload_model_weights(op='model')
                     elif shared.backend == shared.Backend.DIFFUSERS:
-                        load_diffuser()
+                        load_diffuser(op='model')
                     else:
                         shared.log.error(f"Unknown Stable Diffusion backend: {shared.backend}")
                     self.initial = False
@@ -483,11 +494,31 @@ class SdModelData:
         return self.sd_model
 
     def set_sd_model(self, v):
+        shared.log.debug(f"Class model: {v}")
         self.sd_model = v
 
+    def get_sd_refiner(self):
+        if self.sd_model is None:
+            with self.lock:
+                try:
+                    if shared.backend == shared.Backend.ORIGINAL:
+                        reload_model_weights(op='refiner')
+                    elif shared.backend == shared.Backend.DIFFUSERS:
+                        load_diffuser(op='refiner')
+                    else:
+                        shared.log.error(f"Unknown Stable Diffusion backend: {shared.backend}")
+                    self.initial = False
+                except Exception as e:
+                    shared.log.error("Failed to load stable diffusion model")
+                    errors.display(e, "loading stable diffusion model")
+                    self.sd_refiner = None
+        return self.sd_refiner
 
-model_data = SdModelData()
+    def set_sd_refiner(self, v):
+        shared.log.debug(f"Class refiner: {v}")
+        self.sd_refiner = v
 
+model_data = ModelData()
 
 class PriorPipeline:
     def __init__(self, prior, main):
@@ -531,7 +562,7 @@ class PriorPipeline:
         return result
 
 
-def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=None): # pylint: disable=unused-argument
+def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=None, op='model'): # pylint: disable=unused-argument
     if timer is None:
         timer = Timer()
     import logging
@@ -548,27 +579,61 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
 
     if shared.opts.data['sd_model_checkpoint'] == 'model.ckpt':
         shared.opts.data['sd_model_checkpoint'] = "runwayml/stable-diffusion-v1-5"
+
+    if op == 'model' or op == 'dict':
+        if model_data.sd_model is not None and (checkpoint_info.hash == model_data.sd_model.sd_checkpoint_info.hash): # trying to load the same model
+            return
+    else:
+        if model_data.sd_refiner is not None and (checkpoint_info.hash == model_data.sd_refiner.sd_checkpoint_info.hash): # trying to load the same model
+            return
+
     sd_model = None
     try:
-        devices.set_cuda_params() # todo
+        devices.set_cuda_params()
         if shared.cmd_opts.ckpt is not None and model_data.initial: # initial load
             model_name = modelloader.find_diffuser(shared.cmd_opts.ckpt)
             if model_name is not None:
-                shared.log.info(f'Loading diffuser model: {model_name}')
+                shared.log.info(f'Loading diffuser {op}: {model_name}')
                 model_file = modelloader.download_diffusers_model(hub_id=model_name)
                 sd_model = diffusers.DiffusionPipeline.from_pretrained(model_file, **diffusers_load_config)
                 list_models() # rescan for downloaded model
                 checkpoint_info = CheckpointInfo(model_name)
 
         if sd_model is None:
-            checkpoint_info = checkpoint_info or select_checkpoint()
-            shared.log.info(f'Loading diffuser model: {checkpoint_info.filename}')
+            checkpoint_info = checkpoint_info or select_checkpoint(op=op)
+            if checkpoint_info is None:
+                unload_model_weights(op=op)
+                return
+            shared.log.info(f'Loading diffuser {op}: {checkpoint_info.filename}')
             if not os.path.isfile(checkpoint_info.path):
                 sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, **diffusers_load_config)
             else:
                 diffusers_load_config["local_files_only "] = True
                 diffusers_load_config["extract_ema"] = shared.opts.diffusers_extract_ema
-                sd_model = diffusers.StableDiffusionPipeline.from_ckpt(checkpoint_info.path, **diffusers_load_config)
+                try:
+                    # pipelines = ['Stable Diffusion', 'Stable Diffusion XL', 'Kandinsky V1', 'Kandinsky V2', 'DeepFloyd IF', 'Shap-E']
+                    if shared.opts.diffusers_pipeline == shared.pipelines[0]:
+                        pipeline = diffusers.StableDiffusionPipeline
+                    elif shared.opts.diffusers_pipeline == shared.pipelines[1]:
+                        pipeline = diffusers.StableDiffusionXLPipeline
+                    elif shared.opts.diffusers_pipeline == shared.pipelines[2]:
+                        pipeline = diffusers.KandinskyPipeline
+                    elif shared.opts.diffusers_pipeline == shared.pipelines[3]:
+                        pipeline = diffusers.KandinskyV22Pipeline
+                    elif shared.opts.diffusers_pipeline == shared.pipelines[4]:
+                        pipeline = diffusers.IFPipeline
+                    elif shared.opts.diffusers_pipeline == shared.pipelines[5]:
+                        pipeline = diffusers.ShapEPipeline
+                    else:
+                        shared.log.error(f'Diffusers unknown pipeline: {shared.opts.diffusers_pipeline}')
+                except Exception as e:
+                    shared.log.error(f'Diffusers failed initializing pipeline: {shared.opts.diffusers_pipeline} {e}')
+                    return
+                try:
+                    sd_model = pipeline.from_ckpt(checkpoint_info.path, **diffusers_load_config)
+                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
@@ -615,7 +680,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             sd_model.unet.to(memory_format=torch.channels_last)
         if shared.opts.cuda_compile and torch.cuda.is_available():
             sd_model.to(devices.device)
-            import torch._dynamo # pylint: disable=unused-import
+            import torch._dynamo # pylint: disable=unused-import,redefined-outer-name
             log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access
             torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access
             torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access
@@ -632,7 +697,11 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
     except Exception as e:
         shared.log.error("Failed to load diffusers model")
         errors.display(e, "loading Diffusers model")
-    shared.sd_model = sd_model
+
+    if op == 'refiner':
+        model_data.sd_refiner = sd_model
+    else:
+        model_data.sd_model = sd_model
 
     from modules.textual_inversion import textual_inversion
     embedding_db = textual_inversion.EmbeddingDatabase()
@@ -685,7 +754,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     new_pipe.sd_model_checkpoint = sd_model_checkpoint
     new_pipe.sd_model_hash = sd_model_hash
 
-    shared.sd_model = new_pipe
+    model_data.sd_model = new_pipe
     shared.log.info(f"Pipeline class changed from {pipe.__class__.__name__} to {new_pipe_cls.__name__}")
 
 
@@ -700,21 +769,32 @@ def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
     return DiffusersTaskType.TEXT_2_IMAGE
 
 
-def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None):
+def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None, op='model'):
     from modules import lowvram, sd_hijack
-    checkpoint_info = checkpoint_info or select_checkpoint()
+    checkpoint_info = checkpoint_info or select_checkpoint(op=op)
     if checkpoint_info is None:
         return
-    if model_data.sd_model is not None and (checkpoint_info.hash == model_data.sd_model.sd_checkpoint_info.hash): # trying to load the same model
-        return
-    shared.log.debug(f'Load model: name={checkpoint_info.filename} dict={already_loaded_state_dict is not None}')
+    if op == 'model' or op == 'dict':
+        if model_data.sd_model is not None and (checkpoint_info.hash == model_data.sd_model.sd_checkpoint_info.hash): # trying to load the same model
+            return
+    else:
+        if model_data.sd_refiner is not None and (checkpoint_info.hash == model_data.sd_refiner.sd_checkpoint_info.hash): # trying to load the same model
+            return
+    shared.log.debug(f'Load {op}: name={checkpoint_info.filename} dict={already_loaded_state_dict is not None}')
     if timer is None:
         timer = Timer()
     current_checkpoint_info = None
-    if model_data.sd_model is not None:
-        sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
-        current_checkpoint_info = model_data.sd_model.sd_checkpoint_info
-        unload_model_weights()
+    if op == 'model' or op == 'dict':
+        if model_data.sd_model is not None:
+            sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
+            current_checkpoint_info = model_data.sd_model.sd_checkpoint_info
+            unload_model_weights(op=op)
+    else:
+        if model_data.sd_refiner is not None:
+            sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner)
+            current_checkpoint_info = model_data.sd_refiner.sd_checkpoint_info
+            unload_model_weights(op=op)
+
     do_inpainting_hijack()
     devices.set_cuda_params()
     if already_loaded_state_dict is not None:
@@ -760,7 +840,10 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None)
         sd_model = torch.xpu.optimize(sd_model, dtype=devices.dtype, auto_kernel_selection=True, optimize_lstm=True,
         graph_mode=True if shared.opts.cuda_compile and shared.opts.cuda_compile_mode == 'ipex' else False)
         shared.log.info("Applied IPEX Optimize")
-    model_data.sd_model = sd_model
+    if op == 'refiner':
+        model_data.sd_refiner = sd_model
+    else:
+        model_data.sd_model = sd_model
     sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True)  # Reload embeddings after model load as they may or may not fit the model
     timer.record("embeddings")
     script_callbacks.model_loaded_callback(sd_model)
@@ -771,7 +854,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None)
     shared.log.info(f'Model load finished: {memory_stats()} cached={len(checkpoints_loaded.keys())}')
 
 
-def reload_model_weights(sd_model=None, info=None, reuse_dict=False):
+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:
@@ -779,15 +862,18 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False):
         skip_next_load = False
         return
     from modules import lowvram, sd_hijack
-    checkpoint_info = info or select_checkpoint(model=not load_dict) # are we selecting model or dictionary
-    next_checkpoint_info = info or select_checkpoint(model=load_dict) if load_dict else None
+    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
+    if checkpoint_info is None:
+        unload_model_weights(op=op)
+        return
     if load_dict:
         shared.log.debug(f'Model dict: existing={sd_model is not None} target={checkpoint_info.filename} info={info}')
     else:
         model_data.sd_dict = 'None'
         shared.log.debug(f'Load model weights: existing={sd_model is not None} target={checkpoint_info.filename} info={info}')
     if not sd_model:
-        sd_model = model_data.sd_model
+        sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner
     if sd_model is None:  # previous model load failed
         current_checkpoint_info = None
     else:
@@ -802,7 +888,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False):
         shared.log.info('Reusing previous model dictionary')
         sd_hijack.model_hijack.undo_hijack(sd_model)
     else:
-        unload_model_weights()
+        unload_model_weights(op=op)
         sd_model = None
     timer = Timer()
     state_dict = get_checkpoint_state_dict(checkpoint_info, timer)
@@ -811,14 +897,14 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False):
     if sd_model is None or checkpoint_config != sd_model.used_config:
         del sd_model
         if shared.backend == shared.Backend.ORIGINAL:
-            load_model(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer)
+            load_model(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer, op=op)
         else:
-            load_diffuser(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer)
+            load_diffuser(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer, op=op)
         if load_dict and next_checkpoint_info is not None:
             model_data.sd_dict = shared.opts.sd_model_dict
             shared.opts.data["sd_model_checkpoint"] = next_checkpoint_info.title
             reload_model_weights(reuse_dict=True) # ok we loaded dict now lets redo and load model on top of it
-        return model_data.sd_model
+        return model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner
     try:
         load_model_weights(sd_model, checkpoint_info, state_dict, timer)
     except Exception:
@@ -835,17 +921,22 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False):
     shared.log.info(f"Weights loaded in {timer.summary()}")
 
 
-def unload_model_weights(sd_model=None, _info=None):
+def unload_model_weights(op='model'):
     from modules import sd_hijack
-    if model_data.sd_model:
-        model_data.sd_model.to(devices.cpu)
-        if shared.backend == shared.Backend.ORIGINAL:
-            sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
-        sd_model = None
-        model_data.sd_model = None
-        devices.torch_gc(force=True)
-        shared.log.debug(f'Model weights unloaded: {memory_stats()}')
-    return sd_model
+    if op == 'model' or op == 'dict':
+        if model_data.sd_model:
+            model_data.sd_model.to(devices.cpu)
+            if shared.backend == shared.Backend.ORIGINAL:
+                sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
+            model_data.sd_model = None
+    else:
+        if model_data.sd_refiner:
+            model_data.sd_refiner.to(devices.cpu)
+            if shared.backend == shared.Backend.ORIGINAL:
+                sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner)
+            model_data.sd_refiner = None
+    shared.log.debug(f'Weights unloaded {op}: {memory_stats()}')
+    devices.torch_gc(force=True)
 
 
 def apply_token_merging(sd_model, token_merging_ratio):

modules/shared.py

@@ -38,6 +38,7 @@ hypernetworks = {}
 loaded_hypernetworks = []
 gradio_theme = gr.themes.Base()
 settings_components = None
+pipelines = ['Stable Diffusion', 'Stable Diffusion XL', 'Kandinsky V1', 'Kandinsky V2', 'DeepFloyd IF', 'Shap-E']
 latent_upscale_default_mode = "Latent"
 latent_upscale_modes = {
     "Latent": {"mode": "bilinear", "antialias": False},
@@ -217,6 +218,10 @@ class OptionInfo:
         self.comment_after += f"<span class='info'>({info})</span>"
         return self
 
+    def html(self, info):
+        self.comment_after += f"<span class='info'>{info}</span>"
+        return self
+
     def needs_restart(self):
         self.comment_after += " <span class='info'>(requires restart)</span>"
         return self
@@ -295,8 +300,9 @@ else: # cuda
     cross_attention_optimization_default ="Scaled-Dot-Product"
 
 options_templates.update(options_section(('sd', "Stable Diffusion"), {
-    "sd_model_checkpoint": OptionInfo(default_checkpoint, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),
     "sd_checkpoint_autoload": OptionInfo(True, "Stable Diffusion checkpoint autoload on server start"),
+    "sd_model_checkpoint": OptionInfo(default_checkpoint, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),
+    "sd_model_refiner": OptionInfo('None', "Stable Diffusion refiner", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_tiles()}, refresh=refresh_checkpoints),
     "sd_checkpoint_cache": OptionInfo(0, "Number of cached model checkpoints", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
     "sd_vae_checkpoint_cache": OptionInfo(0, "Number of cached VAE checkpoints", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
     "sd_vae": OptionInfo("Automatic", "Select VAE", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list),
@@ -341,10 +347,11 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "cuda_compile_fullgraph": OptionInfo(False, "Model compile fullgraph"),
     "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"),
     "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"),
-    "disable_gc": OptionInfo(False, "Disable Torch memory garbage collection"),
+    "disable_gc": OptionInfo(True, "Disable Torch memory garbage collection on each generation"),
 }))
 
 options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
+    "diffusers_pipeline": OptionInfo(pipelines[0], 'Diffuser Pipeline', gr.Dropdown, lambda: {"choices": pipelines}),
     "diffusers_extract_ema": OptionInfo(True, "Use model EMA weights when possible"),
     "diffusers_generator_device": OptionInfo("default", "Generator device", gr.Radio, lambda: {"choices": ["default", "cpu"]}),
     "diffusers_seq_cpu_offload": OptionInfo(False, "Enable sequential CPU offload"),
@@ -905,7 +912,6 @@ class Shared(sys.modules[__name__].__class__):
     @property
     def sd_model(self):
         import modules.sd_models # pylint: disable=W0621
-        # return modules.sd_models.model_data.sd_model
         return modules.sd_models.model_data.get_sd_model()
 
     @sd_model.setter
@@ -913,6 +919,17 @@ class Shared(sys.modules[__name__].__class__):
         import modules.sd_models # pylint: disable=W0621
         modules.sd_models.model_data.set_sd_model(value)
 
-# sd_model: LatentDiffusion = None  # this var is here just for IDE's type checking; it cannot be accessed because the class field above will be accessed instead
+    @property
+    def sd_refiner(self):
+        import modules.sd_models # pylint: disable=W0621
+        return modules.sd_models.model_data.get_sd_refiner()
+
+    @sd_refiner.setter
+    def sd_refiner(self, value):
+        import modules.sd_models # pylint: disable=W0621
+        modules.sd_models.model_data.set_sd_refiner(value)
+
+
 sd_model = None
+sd_refiner = None
 sys.modules[__name__].__class__ = Shared

modules/textual_inversion/textual_inversion.py

@@ -145,8 +145,12 @@ class EmbeddingDatabase:
             self.word_embeddings[name] = embedding
         except Exception:
             self.skipped_embeddings[name] = embedding
-        text_inv_tokens = pipe.tokenizer.added_tokens_encoder.keys()
-        text_inv_tokens = [t for t in text_inv_tokens if not (len(t.split("_")) > 1 and t.split("_")[-1].isdigit())]
+        try:
+            text_inv_tokens = pipe.tokenizer.added_tokens_encoder.keys()
+            text_inv_tokens = [t for t in text_inv_tokens if not (len(t.split("_")) > 1 and t.split("_")[-1].isdigit())]
+        except Exception:
+            text_inv_tokens = []
+            pass
 
     def load_from_file(self, path, filename):
         name, ext = os.path.splitext(filename)

modules/ui.py

@@ -1031,7 +1031,8 @@ def create_ui(startup_timer = None):
                 create_dirty_indicator("show_all_pages", [], interactive=False)
 
         def unload_sd_weights():
-            modules.sd_models.unload_model_weights()
+            modules.sd_models.unload_model_weights(op='model')
+            modules.sd_models.unload_model_weights(op='refiner')
 
         def reload_sd_weights():
             modules.sd_models.reload_model_weights()