From b9a0e23aeee83274a095622d6228be9293074f9d Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 15 Aug 2023 10:17:31 +0530 Subject: [PATCH 1/8] relax lora loading logic. --- src/diffusers/loaders.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 0907cbfd163a..885991825536 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -25,7 +25,7 @@ import safetensors import torch import torch.nn.functional as F -from huggingface_hub import hf_hub_download +from huggingface_hub import hf_hub_download, model_info from torch import nn from .utils import ( @@ -1018,6 +1018,22 @@ def lora_state_dict( if (use_safetensors and weight_name is None) or ( weight_name is not None and weight_name.endswith(".safetensors") ): + # Here we're relaxing the loading check to enable more Inference API + # friendliness where sometimes, it's not at all possible to automatically + # determine `weight_name`. + if weight_name is None: + if os.path.isdir(pretrained_model_name_or_path_or_dict): + all_safetensors = [ + f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(".safetensors") + ] + else: + files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings + all_safetensors = [f.rfilename for f in files_in_repo if f.rfilename.endswith(".safetensors")] + if all_safetensors and len(all_safetensors) > 1: + raise ValueError( + f"Provided path contains more than one `.safetensors` file. This makes the loading process ambiguous. `weight_name` wasn't specified. So, we tried to pick a `.safetensors` from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single `.safetensors` file." + ) + weight_name = all_safetensors[0] try: model_file = _get_model_file( pretrained_model_name_or_path_or_dict, @@ -1038,6 +1054,7 @@ def lora_state_dict( raise e # try loading non-safetensors weights pass + if model_file is None: model_file = _get_model_file( pretrained_model_name_or_path_or_dict, From 1e4efbe53ea13b715092a1d679eea50bc1be9ce0 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 15 Aug 2023 10:50:22 +0530 Subject: [PATCH 2/8] cater to the other cases too. --- src/diffusers/loaders.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 885991825536..623e30818380 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -1022,18 +1022,9 @@ def lora_state_dict( # friendliness where sometimes, it's not at all possible to automatically # determine `weight_name`. if weight_name is None: - if os.path.isdir(pretrained_model_name_or_path_or_dict): - all_safetensors = [ - f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(".safetensors") - ] - else: - files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings - all_safetensors = [f.rfilename for f in files_in_repo if f.rfilename.endswith(".safetensors")] - if all_safetensors and len(all_safetensors) > 1: - raise ValueError( - f"Provided path contains more than one `.safetensors` file. This makes the loading process ambiguous. `weight_name` wasn't specified. So, we tried to pick a `.safetensors` from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single `.safetensors` file." - ) - weight_name = all_safetensors[0] + weight_name = cls._best_guess_weight_name( + pretrained_model_name_or_path_or_dict, file_extension=".safetensors" + ) try: model_file = _get_model_file( pretrained_model_name_or_path_or_dict, @@ -1056,6 +1047,7 @@ def lora_state_dict( pass if model_file is None: + weight_name = cls._best_guess_weight_name(pretrained_model_name_or_path_or_dict, file_extension=".bin") model_file = _get_model_file( pretrained_model_name_or_path_or_dict, weights_name=weight_name or LORA_WEIGHT_NAME, @@ -1091,6 +1083,22 @@ def lora_state_dict( return state_dict, network_alphas + @classmethod + def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_extension): + if os.path.isdir(pretrained_model_name_or_path_or_dict): + all_safetensors = [ + f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(file_extension) + ] + else: + files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings + all_safetensors = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)] + if all_safetensors and len(all_safetensors) > 1: + raise ValueError( + f"Provided path contains more than one `.safetensors` file. This makes the loading process ambiguous. `weight_name` wasn't specified. So, we tried to pick a `.safetensors` from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single `.safetensors` file." + ) + weight_name = all_safetensors[0] + return weight_name + @classmethod def _map_sgm_blocks_to_diffusers(cls, state_dict, unet_config, delimiter="_", block_slice_pos=5): is_all_unet = all(k.startswith("lora_unet") for k in state_dict) From e6b4f7cd335c8e7a2895a67f3753fa48380bd0af Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 15 Aug 2023 10:51:26 +0530 Subject: [PATCH 3/8] fix: variable name --- src/diffusers/loaders.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 623e30818380..109acc3e9182 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -1086,17 +1086,17 @@ def lora_state_dict( @classmethod def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_extension): if os.path.isdir(pretrained_model_name_or_path_or_dict): - all_safetensors = [ + targeted_files = [ f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(file_extension) ] else: files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings - all_safetensors = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)] - if all_safetensors and len(all_safetensors) > 1: + targeted_files = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)] + if targeted_files and len(targeted_files) > 1: raise ValueError( f"Provided path contains more than one `.safetensors` file. This makes the loading process ambiguous. `weight_name` wasn't specified. So, we tried to pick a `.safetensors` from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single `.safetensors` file." ) - weight_name = all_safetensors[0] + weight_name = targeted_files[0] return weight_name @classmethod From 77fb6de41beaf5c219bc410e9def63109bb42559 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 15 Aug 2023 11:47:46 +0530 Subject: [PATCH 4/8] bring the chaos down. --- src/diffusers/loaders.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 109acc3e9182..ff4d5926ad2a 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -1018,14 +1018,14 @@ def lora_state_dict( if (use_safetensors and weight_name is None) or ( weight_name is not None and weight_name.endswith(".safetensors") ): - # Here we're relaxing the loading check to enable more Inference API - # friendliness where sometimes, it's not at all possible to automatically - # determine `weight_name`. - if weight_name is None: - weight_name = cls._best_guess_weight_name( - pretrained_model_name_or_path_or_dict, file_extension=".safetensors" - ) try: + # Here we're relaxing the loading check to enable more Inference API + # friendliness where sometimes, it's not at all possible to automatically + # determine `weight_name`. + if weight_name is None: + weight_name = cls._best_guess_weight_name( + pretrained_model_name_or_path_or_dict, file_extension=".safetensors" + ) model_file = _get_model_file( pretrained_model_name_or_path_or_dict, weights_name=weight_name or LORA_WEIGHT_NAME_SAFE, @@ -1047,7 +1047,10 @@ def lora_state_dict( pass if model_file is None: - weight_name = cls._best_guess_weight_name(pretrained_model_name_or_path_or_dict, file_extension=".bin") + if weight_name is None: + weight_name = cls._best_guess_weight_name( + pretrained_model_name_or_path_or_dict, file_extension=".bin" + ) model_file = _get_model_file( pretrained_model_name_or_path_or_dict, weights_name=weight_name or LORA_WEIGHT_NAME, @@ -1084,7 +1087,8 @@ def lora_state_dict( return state_dict, network_alphas @classmethod - def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_extension): + def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_extension=".safetensors"): + targeted_files = [] if os.path.isdir(pretrained_model_name_or_path_or_dict): targeted_files = [ f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(file_extension) @@ -1092,10 +1096,12 @@ def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_ext else: files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings targeted_files = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)] - if targeted_files and len(targeted_files) > 1: + if len(targeted_files) > 1: raise ValueError( - f"Provided path contains more than one `.safetensors` file. This makes the loading process ambiguous. `weight_name` wasn't specified. So, we tried to pick a `.safetensors` from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single `.safetensors` file." + f"Provided path contains more than one weights file in the {file_extension} format. Also, `weight_name` wasn't specified. This makes the loading process ambiguous. So, we tried to pick a file having the `.safetensors` or the `.bin` extension from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single weights file having either `.safetensors` or `.bin` extension." ) + if len(targeted_files) == 0: + return weight_name = targeted_files[0] return weight_name From 787b2c2f5cc2fef6d7753e7f45abf1470006cb02 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 15 Aug 2023 12:50:18 +0530 Subject: [PATCH 5/8] check --- src/diffusers/loaders.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index ff4d5926ad2a..5c49c56f223b 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -1096,6 +1096,7 @@ def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_ext else: files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings targeted_files = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)] + print(f"From LoRA loading: {targeted_files}") if len(targeted_files) > 1: raise ValueError( f"Provided path contains more than one weights file in the {file_extension} format. Also, `weight_name` wasn't specified. This makes the loading process ambiguous. So, we tried to pick a file having the `.safetensors` or the `.bin` extension from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single weights file having either `.safetensors` or `.bin` extension." From 45a5f7341fef28ff7d9878ca3a11d171bbdde2a8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 15 Aug 2023 13:20:28 +0530 Subject: [PATCH 6/8] deal with checkpointed files. --- src/diffusers/loaders.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 5c49c56f223b..e1fde1a4cdaf 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -1096,13 +1096,15 @@ def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_ext else: files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings targeted_files = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)] - print(f"From LoRA loading: {targeted_files}") + + if len(targeted_files) == 0: + return + + targeted_files = list(filter(lambda x: "scheduler" not in x and "optimizer" not in x, targeted_files)) if len(targeted_files) > 1: raise ValueError( f"Provided path contains more than one weights file in the {file_extension} format. Also, `weight_name` wasn't specified. This makes the loading process ambiguous. So, we tried to pick a file having the `.safetensors` or the `.bin` extension from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single weights file having either `.safetensors` or `.bin` extension." ) - if len(targeted_files) == 0: - return weight_name = targeted_files[0] return weight_name From c688689d278f49c28b1927b0d9181aee4a170f20 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 16 Aug 2023 15:10:24 +0530 Subject: [PATCH 7/8] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: apolinário --- src/diffusers/loaders.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index e1fde1a4cdaf..e54f7f05a04a 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -1103,7 +1103,7 @@ def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_ext targeted_files = list(filter(lambda x: "scheduler" not in x and "optimizer" not in x, targeted_files)) if len(targeted_files) > 1: raise ValueError( - f"Provided path contains more than one weights file in the {file_extension} format. Also, `weight_name` wasn't specified. This makes the loading process ambiguous. So, we tried to pick a file having the `.safetensors` or the `.bin` extension from the provided path: {pretrained_model_name_or_path_or_dict}. Please either pass `weight_name` or ensure {pretrained_model_name_or_path_or_dict} only has a single weights file having either `.safetensors` or `.bin` extension." + f"Provided path contains more than one weights file in the {file_extension} format. Either specify `weight_name` in `load_lora_weights` or make sure there's only one `.safetensors` or `.bin` file in {pretrained_model_name_or_path_or_dict}." ) weight_name = targeted_files[0] return weight_name From 5253e72cf102d066449046d1940038298202389a Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 24 Aug 2023 16:56:34 +0530 Subject: [PATCH 8/8] style --- src/diffusers/loaders.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 4bfe4174e479..d50a855c83c0 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -1092,7 +1092,10 @@ def lora_state_dict( @classmethod def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_extension=".safetensors"): targeted_files = [] - if os.path.isdir(pretrained_model_name_or_path_or_dict): + + if os.path.isfile(pretrained_model_name_or_path_or_dict): + return + elif os.path.isdir(pretrained_model_name_or_path_or_dict): targeted_files = [ f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(file_extension) ]