improve vae reload

CHANGELOG.md

@@ -22,16 +22,18 @@
     - add model **precompile** option (when model compile is enbled)  
     - **extra network** folder info caching  
       results in much faster startup when you have large number of extra networks  
-  - fixes:
-    - fix extra networks previews  
-    - fix gradio gallery  
-    - css fixes  
-    - improved extensions compatibility (e.g. sd-cn-animation)  
+    - faster **xyz grid** switching  
+      especially when using different checkpoints  
   - add **openvino** (experimental) cpu optimized model compile and inference  
     enable with `--use-openvino`  
     thanks @disty0  
   - enable batch **img2img** scale-by workflows  
     now you can batch process with rescaling based on eah individual original image size  
+  - fixes:
+    - fix extra networks previews  
+    - css fixes  
+    - improved extensions compatibility (e.g. *sd-cn-animation*)  
+    - allow changing **vae** on-the-fly for both original and diffusers backend
 
 ## Update for 2023-08-20
 

javascript/setHints.js

@@ -85,7 +85,7 @@ async function setHints() {
   }
   const t1 = performance.now();
   log('setHints', { type: localeData.type, elements: elements.length, localized, hints, data: localeData.data.length, time: t1 - t0 });
-  sortUIElements();
+  // sortUIElements();
   removeSplash();
   // validateHints(elements, localeData.data)
 }

modules/sd_models.py

@@ -1046,6 +1046,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
         sd_model = None
         if shared.backend == shared.Backend.ORIGINAL:
             load_model(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer, op=op)
+            model_data.sd_dict = shared.opts.sd_model_dict
         else:
             load_diffuser(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer, op=op)
         if load_dict and next_checkpoint_info is not None:

modules/sd_vae.py

@@ -254,8 +254,12 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
         script_callbacks.model_loaded_callback(sd_model)
         if vae_file is not None:
             shared.log.info(f"VAE weights loaded: {vae_file}")
-    # else:
-    #    load_vae_diffusers(model_file, vae_file, vae_source)
+    else:
+        if hasattr(shared.sd_model, "vae") and hasattr(shared.sd_model, "sd_checkpoint_info"):
+            vae = load_vae_diffusers(shared.sd_model.sd_checkpoint_info.filename, vae_file, vae_source)
+            if vae is not None:
+                if vae is not None:
+                    sd_model.vae = vae
 
     if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and not sd_model.has_accelerate:
         sd_model.to(devices.device)

modules/shared.py

@@ -519,7 +519,7 @@ options_templates.update(options_section(('ui', "User Interface"), {
     "keyedit_delimiters": OptionInfo(".,\/!?%^*;:{}=`~()", "Ctrl+up/down word delimiters"), # pylint: disable=anomalous-backslash-in-string
     "quicksettings_list": OptionInfo(["sd_model_checkpoint"], "Quicksettings list", ui_components.DropdownMulti, lambda: {"choices": list(opts.data_labels.keys())}),
     "hidden_tabs": OptionInfo([], "Hidden UI tabs", ui_components.DropdownMulti, lambda: {"choices": list(tab_names)}),
-    "ui_tab_reorder": OptionInfo("From Text, From Image, Process Image", "UI tabs order"),
+    # "ui_tab_reorder": OptionInfo("From Text, From Image, Process Image", "UI tabs order"),
     "ui_scripts_reorder": OptionInfo("Enable Dynamic Thresholding, ControlNet", "UI scripts order"),
 }))
 

scripts/xyz_grid.py

@@ -600,13 +600,13 @@ class Script(scripts.Script):
             if shared.state.interrupted:
                 return Processed(p, [], p.seed, "")
             pc = copy(p)
+            pc.override_settings_restore_afterwards = False
             pc.styles = pc.styles[:]
             x_opt.apply(pc, x, xs)
             y_opt.apply(pc, y, ys)
             z_opt.apply(pc, z, zs)
             res = process_images(pc)
-            # Sets subgrid infotexts
-            subgrid_index = 1 + iz
+            subgrid_index = 1 + iz # Sets subgrid infotexts
             if grid_infotext[subgrid_index] is None and ix == 0 and iy == 0:
                 pc.extra_generation_params = copy(pc.extra_generation_params)
                 pc.extra_generation_params['Script'] = self.title()
@@ -621,10 +621,8 @@ class Script(scripts.Script):
                     if y_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
                         pc.extra_generation_params["Fixed Y Values"] = ", ".join([str(y) for y in ys])
                 grid_infotext[subgrid_index] = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds)
-            # Sets main grid infotext
-            if grid_infotext[0] is None and ix == 0 and iy == 0 and iz == 0:
+            if grid_infotext[0] is None and ix == 0 and iy == 0 and iz == 0: # Sets main grid infotext
                 pc.extra_generation_params = copy(pc.extra_generation_params)
-
                 if z_opt.label != 'Nothing':
                     pc.extra_generation_params["Z Type"] = z_opt.label
                     pc.extra_generation_params["Z Values"] = z_values
@@ -653,22 +651,17 @@ class Script(scripts.Script):
             )
 
         if not processed.images:
-            # It broke, no further handling needed.
-            return processed
+            return processed # It broke, no further handling needed.
         z_count = len(zs)
-        # Set the grid infotexts to the real ones with extra_generation_params (1 main grid + z_count sub-grids)
-        processed.infotexts[:1+z_count] = grid_infotext[:1+z_count]
+        processed.infotexts[:1+z_count] = grid_infotext[:1+z_count] # Set the grid infotexts to the real ones with extra_generation_params (1 main grid + z_count sub-grids)
         if not include_lone_images:
-            # Don't need sub-images anymore, drop from list:
-            processed.images = processed.images[:z_count+1]
-        if shared.opts.grid_save:
-            # Auto-save main and sub-grids:
+            processed.images = processed.images[:z_count+1] # Don't need sub-images anymore, drop from list:
+        if shared.opts.grid_save: # Auto-save main and sub-grids:
             grid_count = z_count + 1 if z_count > 1 else 1
             for g in range(grid_count):
                 adj_g = g-1 if g > 0 else g
                 images.save_image(processed.images[g], p.outpath_grids, "xyz_grid", info=processed.infotexts[g], extension=shared.opts.grid_format, prompt=processed.all_prompts[adj_g], seed=processed.all_seeds[adj_g], grid=True, p=processed)
-        if not include_sub_grids:
-            # Done with sub-grids, drop all related information:
+        if not include_sub_grids: # Done with sub-grids, drop all related information:
             for _sg in range(z_count):
                 del processed.images[1]
                 del processed.all_prompts[1]

webui.py

@@ -196,6 +196,7 @@ def load_model():
     shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='model')), call=False)
     shared.opts.onchange("sd_model_refiner", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='refiner')), call=False)
     shared.opts.onchange("sd_model_dict", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='dict')), call=False)
+    shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False)
     shared.opts.onchange("sd_backend", wrap_queued_call(lambda: modules.sd_models.change_backend()), call=False)
     startup_timer.record("checkpoint")