From 3ca654d25606bb850cc75742befd9cfa27088396 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 13 Mar 2023 23:27:29 -0400 Subject: [PATCH 1/3] speculative fix for alternative vaes --- .../convert_ckpt_to_diffusers.py | 27 ++++++++++-- .../backend/model_management/model_manager.py | 42 +++---------------- 2 files changed, 28 insertions(+), 41 deletions(-) diff --git a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py index ae5550880a..6559ec6a0c 100644 --- a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py +++ b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py @@ -1026,6 +1026,14 @@ def convert_open_clip_checkpoint(checkpoint): return text_model +def replace_checkpoint_vae(checkpoint, vae_path:str): + if vae_path.endswith(".safetensors"): + vae_ckpt = load_file(vae_path) + else: + vae_ckpt = torch.load(vae_path, map_location="cpu") + for vae_key in vae_ckpt['state_dict']: + new_key = f'first_stage_model.{vae_key}' + checkpoint[new_key] = vae_ckpt['state_dict'][vae_key] def load_pipeline_from_original_stable_diffusion_ckpt( checkpoint_path: str, @@ -1038,6 +1046,7 @@ def load_pipeline_from_original_stable_diffusion_ckpt( extract_ema: bool = True, upcast_attn: bool = False, vae: AutoencoderKL = None, + vae_path: str = None, precision: torch.dtype = torch.float32, return_generator_pipeline: bool = False, ) -> Union[StableDiffusionPipeline, StableDiffusionGeneratorPipeline]: @@ -1067,6 +1076,8 @@ def load_pipeline_from_original_stable_diffusion_ckpt( :param precision: precision to use - torch.float16, torch.float32 or torch.autocast :param upcast_attention: Whether the attention computation should always be upcasted. This is necessary when running stable diffusion 2.1. + :param vae: A diffusers VAE to load into the pipeline. + :param vae_path: Path to a checkpoint VAE that will be converted into diffusers and loaded into the pipeline. """ with warnings.catch_warnings(): @@ -1201,9 +1212,19 @@ def load_pipeline_from_original_stable_diffusion_ckpt( unet.load_state_dict(converted_unet_checkpoint) - # Convert the VAE model, or use the one passed - if not vae: + # If a replacement VAE path was specified, we'll incorporate that into + # the checkpoint model and then convert it + if vae_path: + print(f" | Converting VAE {vae_path}") + replace_checkpoint_vae(checkpoint,vae_path) + # otherwise we use the original VAE, provided that + # an externally loaded diffusers VAE was not passed + elif not vae: print(" | Using checkpoint model's original VAE") + + if vae: + print(" | Using replacement diffusers VAE") + else: # convert the original or replacement VAE vae_config = create_vae_diffusers_config( original_config, image_size=image_size ) @@ -1213,8 +1234,6 @@ def load_pipeline_from_original_stable_diffusion_ckpt( vae = AutoencoderKL(**vae_config) vae.load_state_dict(converted_vae_checkpoint) - else: - print(" | Using external VAE specified in config") # Convert the text model. model_type = pipeline_type diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index 84e2ab378b..6545b4592a 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -45,9 +45,6 @@ class SDLegacyType(Enum): UNKNOWN = 99 DEFAULT_MAX_MODELS = 2 -VAE_TO_REPO_ID = { # hack, see note in convert_and_import() - "vae-ft-mse-840000-ema-pruned": "stabilityai/sd-vae-ft-mse", -} class ModelManager(object): ''' @@ -458,14 +455,15 @@ class ModelManager(object): from . import load_pipeline_from_original_stable_diffusion_ckpt self.offload_model(self.current_model) - if vae_config := self._choose_diffusers_vae(model_name): - vae = self._load_vae(vae_config) + vae_path = None + if vae: + vae_path = vae if os.path.isabs(vae) else os.path.normpath(os.path.join(Globals.root, vae)) if self._has_cuda(): torch.cuda.empty_cache() pipeline = load_pipeline_from_original_stable_diffusion_ckpt( checkpoint_path=weights, original_config_file=config, - vae=vae, + vae_path=vae_path, return_generator_pipeline=True, precision=torch.float16 if self.precision == "float16" else torch.float32, ) @@ -519,7 +517,7 @@ class ModelManager(object): def scan_model(self, model_name, checkpoint): """ - Apply picklescanner to the indicated checkpoint and issue a warning +v Apply picklescanner to the indicated checkpoint and issue a warning and option to exit if an infected file is identified. """ # scan model @@ -879,36 +877,6 @@ class ModelManager(object): return search_folder, found_models - def _choose_diffusers_vae( - self, model_name: str, vae: str = None - ) -> Union[dict, str]: - # In the event that the original entry is using a custom ckpt VAE, we try to - # map that VAE onto a diffuser VAE using a hard-coded dictionary. - # I would prefer to do this differently: We load the ckpt model into memory, swap the - # VAE in memory, and then pass that to convert_ckpt_to_diffuser() so that the swapped - # VAE is built into the model. However, when I tried this I got obscure key errors. - if vae: - return vae - if model_name in self.config and ( - vae_ckpt_path := self.model_info(model_name).get("vae", None) - ): - vae_basename = Path(vae_ckpt_path).stem - diffusers_vae = None - if diffusers_vae := VAE_TO_REPO_ID.get(vae_basename, None): - print( - f">> {vae_basename} VAE corresponds to known {diffusers_vae} diffusers version" - ) - vae = {"repo_id": diffusers_vae} - else: - print( - f'** Custom VAE "{vae_basename}" found, but corresponding diffusers model unknown' - ) - print( - '** Using "stabilityai/sd-vae-ft-mse"; If this isn\'t right, please edit the model config' - ) - vae = {"repo_id": "stabilityai/sd-vae-ft-mse"} - return vae - def _make_cache_room(self) -> None: num_loaded_models = len(self.models) if num_loaded_models >= self.max_loaded_models: From 4e0b5d85badcf62f1fa691c5b5c9fe51a5762544 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 23 Mar 2023 13:14:19 -0400 Subject: [PATCH 2/3] convert custom VAEs into diffusers - When a legacy checkpoint model is loaded via --convert_ckpt and its models.yaml stanza refers to a custom VAE path (using the 'vae:' key), the custom VAE will be converted and used within the diffusers model. Otherwise the VAE contained within the legacy model will be used. - Note that the heuristic_import() method, which imports arbitrary legacy files on disk and URLs, will continue to default to the the standard stabilityai/sd-vae-ft-mse VAE. This can be fixed after the fact by editing the models.yaml stanza using the Web or CLI UIs. - Fixes issue #2917 --- .../convert_ckpt_to_diffusers.py | 2 +- .../backend/model_management/model_manager.py | 31 +++++++++++++------ invokeai/frontend/CLI/CLI.py | 10 ++---- 3 files changed, 24 insertions(+), 19 deletions(-) diff --git a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py index a7095ceba5..abbde3d16f 100644 --- a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py +++ b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py @@ -1033,7 +1033,7 @@ def replace_checkpoint_vae(checkpoint, vae_path:str): vae_ckpt = torch.load(vae_path, map_location="cpu") for vae_key in vae_ckpt['state_dict']: new_key = f'first_stage_model.{vae_key}' - checkpoint[new_key] = vae_ckpt['state_dict'][vae_key] + checkpoint[new_key] = vae_ckpt['state_dict'][vae_key] def load_pipeline_from_original_stable_diffusion_ckpt( checkpoint_path: str, diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index c5f982ce2d..05c8703901 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -454,7 +454,12 @@ class ModelManager(object): from . import load_pipeline_from_original_stable_diffusion_ckpt - self.offload_model(self.current_model) + try: + if self.list_models()[self.current_model]['status'] == 'active': + self.offload_model(self.current_model) + except Exception as e: + pass + vae_path = None if vae: vae_path = vae if os.path.isabs(vae) else os.path.normpath(os.path.join(Globals.root, vae)) @@ -510,6 +515,7 @@ class ModelManager(object): print(f">> Offloading {model_name} to CPU") model = self.models[model_name]["model"] model.offload_all() + self.current_model = None gc.collect() if self._has_cuda(): @@ -790,14 +796,15 @@ v Apply picklescanner to the indicated checkpoint and issue a warning return model_name def convert_and_import( - self, - ckpt_path: Path, - diffusers_path: Path, - model_name=None, - model_description=None, - vae=None, - original_config_file: Path = None, - commit_to_conf: Path = None, + self, + ckpt_path: Path, + diffusers_path: Path, + model_name=None, + model_description=None, + vae:dict=None, + vae_path:Path=None, + original_config_file: Path = None, + commit_to_conf: Path = None, ) -> str: """ Convert a legacy ckpt weights file to diffuser model and import @@ -825,13 +832,17 @@ v Apply picklescanner to the indicated checkpoint and issue a warning try: # By passing the specified VAE to the conversion function, the autoencoder # will be built into the model rather than tacked on afterward via the config file - vae_model = self._load_vae(vae) if vae else None + vae_model=None + if vae: + vae_model=self._load_vae(vae) + vae_path=None convert_ckpt_to_diffusers( ckpt_path, diffusers_path, extract_ema=True, original_config_file=original_config_file, vae=vae_model, + vae_path=vae_path, ) print( f" | Success. Optimized model is now located at {str(diffusers_path)}" diff --git a/invokeai/frontend/CLI/CLI.py b/invokeai/frontend/CLI/CLI.py index 17e1c314f7..173ff3ecc1 100644 --- a/invokeai/frontend/CLI/CLI.py +++ b/invokeai/frontend/CLI/CLI.py @@ -772,16 +772,10 @@ def convert_model(model_name_or_path: Union[Path, str], gen, opt, completer): original_config_file = Path(model_info["config"]) model_name = model_name_or_path model_description = model_info["description"] - vae = model_info["vae"] + vae_path = model_info.get("vae") else: print(f"** {model_name_or_path} is not a legacy .ckpt weights file") return - if vae_repo := invokeai.backend.model_management.model_manager.VAE_TO_REPO_ID.get( - Path(vae).stem - ): - vae_repo = dict(repo_id=vae_repo) - else: - vae_repo = None model_name = manager.convert_and_import( ckpt_path, diffusers_path=Path( @@ -790,7 +784,7 @@ def convert_model(model_name_or_path: Union[Path, str], gen, opt, completer): model_name=model_name, model_description=model_description, original_config_file=original_config_file, - vae=vae_repo, + vae_path=vae_path, ) else: try: From a97107bd909c5da99ecaa167f5800f68f2ded9f6 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 23 Mar 2023 15:11:29 -0400 Subject: [PATCH 3/3] handle VAEs that do not have a "state_dict" key --- .../backend/model_management/convert_ckpt_to_diffusers.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py index abbde3d16f..722ce0aeaa 100644 --- a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py +++ b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py @@ -1031,9 +1031,10 @@ def replace_checkpoint_vae(checkpoint, vae_path:str): vae_ckpt = load_file(vae_path) else: vae_ckpt = torch.load(vae_path, map_location="cpu") - for vae_key in vae_ckpt['state_dict']: + state_dict = vae_ckpt['state_dict'] if "state_dict" in vae_ckpt else vae_ckpt + for vae_key in state_dict: new_key = f'first_stage_model.{vae_key}' - checkpoint[new_key] = vae_ckpt['state_dict'][vae_key] + checkpoint[new_key] = state_dict[vae_key] def load_pipeline_from_original_stable_diffusion_ckpt( checkpoint_path: str,