convert VAEs to diffusers format automatically

- If the user enters a VAE .ckpt path into the VAE field of a
  diffusers model, the VAE will be automatically converted behind
  the scenes into a diffusers version, then loaded.

- This commit is untested (done on an airplane).

Author: lstein

ldm/invoke/ckpt_to_diffuser.py

@@ -620,7 +620,10 @@ def convert_ldm_vae_checkpoint(checkpoint, config):
     for key in keys:
         if key.startswith(vae_key):
             vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
+    new_checkpoint = _convert_ldm_vae_checkpoint(vae_state_dict,config)
+    return new_checkpoint
 
+def _convert_ldm_vae_checkpoint(vae_state_dict, config):
     new_checkpoint = {}
 
     new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]

ldm/invoke/model_manager.py

@@ -9,7 +9,6 @@
 import contextlib
 import gc
 import hashlib
-import io
 import os
 import re
 import sys
@@ -32,10 +31,15 @@
 from omegaconf.dictconfig import DictConfig
 from picklescan.scanner import scan_file_path
 
+from .ckpt_to_diffuser import (
+    load_pipeline_from_original_stable_diffusion_ckpt,
+    create_vae_diffusers_config,
+    convert_ldm_vae_checkpoint,
+    )
 from ldm.invoke.devices import CPU_DEVICE
 from ldm.invoke.generator.diffusers_pipeline import StableDiffusionGeneratorPipeline
 from ldm.invoke.globals import Globals, global_cache_dir
-from ldm.util import ask_user, download_with_resume, instantiate_from_config, url_attachment_name
+from ldm.util import ask_user, download_with_resume, url_attachment_name
 
 
 class SDLegacyType(Enum):
@@ -370,14 +374,7 @@ def _load_ckpt_model(self, model_name, mconfig):
         print(
             f">> Converting legacy checkpoint {model_name} into a diffusers model..."
         )
-        from ldm.invoke.ckpt_to_diffuser import load_pipeline_from_original_stable_diffusion_ckpt
-
-#        try:
-#            if self.list_models()[self.current_model]['status'] == 'active':
-#                self.offload_model(self.current_model)
-#        except Exception:
-#            pass
-
+        
         if self._has_cuda():
             torch.cuda.empty_cache()
         pipeline = load_pipeline_from_original_stable_diffusion_ckpt(
@@ -1236,6 +1233,13 @@ def _scan_for_matching_file(
         return vae_path
 
     def _load_vae(self, vae_config) -> AutoencoderKL:
+
+        # Handle the common case of a user shoving a VAE .ckpt into
+        # the vae field for a diffusers. We convert it into diffusers
+        # format and use it.
+        if type(vae_config) in [str,Path]:
+            return self.convert_vae(vae_config)
+            
         vae_args = {}
         try:
             name_or_path = self.model_name_or_path(vae_config)
@@ -1283,6 +1287,21 @@ def _load_vae(self, vae_config) -> AutoencoderKL:
 
         return vae
 
+    def convert_vae(vae_path: Union[Path,str])->AutoencoderKL:
+        vae_path = Path(vae_path)
+        if vae_path.suffix in ['.pt','.ckpt']:
+            vae_state_dict = torch.load(vae_path)
+        else:
+            vae_state_dict = safetensors.torch.load_file(vae_path)
+        # TODO: see if this works with 1.x inpaint models and 2.x models
+        config_file_path = Path(Globals.root,"configs/stable-diffusion/v1-inference.yaml")
+        original_conf = OmegaConf.load(config_file_path)
+        vae_config = create_vae_diffusers_config(original_conf, image_size=512) # TODO: fix
+        diffusers_vae = convert_ldm_vae_checkpoint(vae_state_dict,vae_config)
+        vae = AutoencoderKL(**vae_config)
+        vae.load_state_dict(diffusers_vae)
+        return vae
+
     @staticmethod
     def _delete_model_from_cache(repo_id):
         cache_info = scan_cache_dir(global_cache_dir("diffusers"))