From 864c00e65f21cfc665fb24febb305bc1473561da Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Thu, 26 Jan 2023 17:19:07 +0100 Subject: [PATCH 01/10] Adding `use_safetensors` argument to give more control to users about which weights they use. --- src/diffusers/models/modeling_utils.py | 12 ++++++++++-- src/diffusers/pipelines/pipeline_utils.py | 9 +++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index a21e09548a59..19c23c06feca 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -392,6 +392,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P variant (`str`, *optional*): If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is ignored when using `from_flax`. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. + If set to `None` (the default). The pipeline will load using `safetensors` if the safetensors weights + are actually available *and* you have the library installed. + If the to `False` the pipeline will *not* use `safetensors` at all. @@ -423,6 +428,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P device_map = kwargs.pop("device_map", None) low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) if low_cpu_mem_usage and not is_accelerate_available(): low_cpu_mem_usage = False @@ -509,7 +515,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P model = load_flax_checkpoint_in_pytorch_model(model, model_file) else: - if is_safetensors_available(): + if use_safetensors in {None, True}: try: model_file = _get_model_file( pretrained_model_name_or_path, @@ -525,7 +531,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P user_agent=user_agent, commit_hash=commit_hash, ) - except: # noqa: E722 + except Exception as e: + if use_safetensors is True: + raise e pass if model_file is None: model_file = _get_model_file( diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 917d2dca853a..4e2b5a1966d4 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -694,6 +694,13 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, setting this argument to `True` will raise an error. + return_cached_folder (`bool`, *optional*, defaults to `False`): + If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. + If set to `None` (the default). The pipeline will load using `safetensors` if the safetensors weights + are actually available *and* you have the library installed. + If the to `False` the pipeline will *not* use `safetensors` at all. kwargs (remaining dictionary of keyword arguments, *optional*): Can be used to overwrite load - and saveable variables - *i.e.* the pipeline components - of the specific pipeline class. The overwritten components are then directly passed to the pipelines @@ -752,6 +759,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P device_map = kwargs.pop("device_map", None) low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) variant = kwargs.pop("variant", None) + return_cached_folder = kwargs.pop("return_cached_folder", False) + use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) # 1. Download the checkpoints and configs # use snapshot download here to get it working from from_pretrained From b2236b84bbd015c939330b519c174359cde43983 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Thu, 26 Jan 2023 17:23:23 +0100 Subject: [PATCH 02/10] Doc style. --- src/diffusers/models/modeling_utils.py | 8 ++++---- src/diffusers/pipelines/pipeline_utils.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 19c23c06feca..32e353c40963 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -393,10 +393,10 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is ignored when using `from_flax`. use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. - If set to `None` (the default). The pipeline will load using `safetensors` if the safetensors weights - are actually available *and* you have the library installed. - If the to `False` the pipeline will *not* use `safetensors` at all. + If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to + `None` (the default). The pipeline will load using `safetensors` if the safetensors weights are + actually available *and* you have the library installed. If the to `False` the pipeline will *not* use + `safetensors` at all. diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 4e2b5a1966d4..a90e97309c3d 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -697,10 +697,10 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P return_cached_folder (`bool`, *optional*, defaults to `False`): If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline. use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. - If set to `None` (the default). The pipeline will load using `safetensors` if the safetensors weights - are actually available *and* you have the library installed. - If the to `False` the pipeline will *not* use `safetensors` at all. + If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to + `None` (the default). The pipeline will load using `safetensors` if the safetensors weights are + actually available *and* you have the library installed. If the to `False` the pipeline will *not* use + `safetensors` at all. kwargs (remaining dictionary of keyword arguments, *optional*): Can be used to overwrite load - and saveable variables - *i.e.* the pipeline components - of the specific pipeline class. The overwritten components are then directly passed to the pipelines From b2ef1de3bb1468b91afaecdfa8ef63352315deff Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 14 Mar 2023 13:33:20 +0100 Subject: [PATCH 03/10] Rebased (not functional). --- src/diffusers/models/modeling_utils.py | 12 ++++++------ src/diffusers/pipelines/pipeline_utils.py | 10 +++++----- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 32e353c40963..9fc2afefd7d7 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -392,11 +392,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P variant (`str`, *optional*): If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is ignored when using `from_flax`. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to - `None` (the default). The pipeline will load using `safetensors` if the safetensors weights are - actually available *and* you have the library installed. If the to `False` the pipeline will *not* use - `safetensors` at all. + use_safetensors (`bool`, *optional* ): + If set to `True`, the pipeline will forcibly load the models from `safetensors` weights. If set to + `None` (the default). The pipeline will load using `safetensors` if safetensors weights are + available *and* if `safetensors` is installed. If the to `False` the pipeline will *not* use + `safetensors`. @@ -515,7 +515,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P model = load_flax_checkpoint_in_pytorch_model(model, model_file) else: - if use_safetensors in {None, True}: + if use_safetensors is not False: try: model_file = _get_model_file( pretrained_model_name_or_path, diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index a90e97309c3d..11b44fda031a 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -696,11 +696,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P setting this argument to `True` will raise an error. return_cached_folder (`bool`, *optional*, defaults to `False`): If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to + use_safetensors (`bool`, *optional* ): + If set to `True`, the pipeline will be loaded from `safetensors` weights. If set to `None` (the default). The pipeline will load using `safetensors` if the safetensors weights are - actually available *and* you have the library installed. If the to `False` the pipeline will *not* use - `safetensors` at all. + available *and* if `safetensors` is installed. If the to `False` the pipeline will *not* use + `safetensors`. kwargs (remaining dictionary of keyword arguments, *optional*): Can be used to overwrite load - and saveable variables - *i.e.* the pipeline components - of the specific pipeline class. The overwritten components are then directly passed to the pipelines @@ -760,7 +760,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) variant = kwargs.pop("variant", None) return_cached_folder = kwargs.pop("return_cached_folder", False) - use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) + kwargs.pop("use_safetensors", None if is_safetensors_available() else False) # 1. Download the checkpoints and configs # use snapshot download here to get it working from from_pretrained From 8185f889c73207852c062c10ec091a22c683da98 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 14 Mar 2023 13:52:02 +0100 Subject: [PATCH 04/10] Rebased and functional with tests. --- src/diffusers/loaders.py | 7 +++- src/diffusers/pipelines/pipeline_utils.py | 7 +++- tests/models/test_models_unet_2d_condition.py | 40 ++++++++++++++++++- tests/test_pipelines.py | 11 +++++ 4 files changed, 61 insertions(+), 4 deletions(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index b2e91e3f6ac3..2317013f5437 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -142,6 +142,7 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict revision = kwargs.pop("revision", None) subfolder = kwargs.pop("subfolder", None) weight_name = kwargs.pop("weight_name", None) + use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) user_agent = { "file_type": "attn_procs_weights", @@ -150,7 +151,7 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict model_file = None if not isinstance(pretrained_model_name_or_path_or_dict, dict): - if (is_safetensors_available() and weight_name is None) or weight_name.endswith(".safetensors"): + if (use_safetensors is not False and weight_name is None) or weight_name.endswith(".safetensors"): if weight_name is None: weight_name = LORA_WEIGHT_NAME_SAFE try: @@ -168,7 +169,9 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict user_agent=user_agent, ) state_dict = safetensors.torch.load_file(model_file, device="cpu") - except EnvironmentError: + except EnvironmentError as e: + if use_safetensors is True: + raise e if weight_name == LORA_WEIGHT_NAME_SAFE: weight_name = None if model_file is None: diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 11b44fda031a..4f66a770c483 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -1077,6 +1077,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: from_flax = kwargs.pop("from_flax", False) custom_pipeline = kwargs.pop("custom_pipeline", None) variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) pipeline_is_cached = False allow_patterns = None @@ -1132,9 +1133,13 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: CUSTOM_PIPELINE_FILE_NAME, ] + if use_safetensors is True and not is_safetensors_compatible(model_filenames, variant=variant): + raise EnvironmentError( + f"Could not found the necessary `safetensors` weights in {model_filenames} (variant={variant})" + ) if from_flax: ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] - elif is_safetensors_available() and is_safetensors_compatible(model_filenames, variant=variant): + elif use_safetensors is not False and is_safetensors_compatible(model_filenames, variant=variant): ignore_patterns = ["*.bin", "*.msgpack"] safetensors_variant_filenames = set([f for f in variant_filenames if f.endswith(".safetensors")]) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index c1f3bc05d7c6..4ddea0766407 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -445,7 +445,7 @@ def test_lora_save_load_safetensors(self): # LoRA and no LoRA should NOT be the same assert (sample - old_sample).abs().max() > 1e-4 - def test_lora_save_load_safetensors_load_torch(self): + def test_lora_save_safetensors_load_torch(self): # enable deterministic behavior for gradient checkpointing init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() @@ -482,6 +482,44 @@ def test_lora_save_load_safetensors_load_torch(self): new_model.to(torch_device) new_model.load_attn_procs(tmpdirname, weight_name="pytorch_lora_weights.bin") + def test_lora_save_torch_force_load_safetensors_error(self): + # enable deterministic behavior for gradient checkpointing + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + + init_dict["attention_head_dim"] = (8, 16) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.to(torch_device) + + lora_attn_procs = {} + for name in model.attn_processors.keys(): + cross_attention_dim = None if name.endswith("attn1.processor") else model.config.cross_attention_dim + if name.startswith("mid_block"): + hidden_size = model.config.block_out_channels[-1] + elif name.startswith("up_blocks"): + block_id = int(name[len("up_blocks.")]) + hidden_size = list(reversed(model.config.block_out_channels))[block_id] + elif name.startswith("down_blocks"): + block_id = int(name[len("down_blocks.")]) + hidden_size = model.config.block_out_channels[block_id] + + lora_attn_procs[name] = LoRACrossAttnProcessor( + hidden_size=hidden_size, cross_attention_dim=cross_attention_dim + ) + lora_attn_procs[name] = lora_attn_procs[name].to(model.device) + + model.set_attn_processor(lora_attn_procs) + # Saving as torch, properly reloads with directly filename + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_attn_procs(tmpdirname) + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) + torch.manual_seed(0) + new_model = self.model_class(**init_dict) + new_model.to(torch_device) + with self.assertRaises(EnvironmentError): + new_model.load_attn_procs(tmpdirname, use_safetensors=True) + def test_lora_on_off(self): # enable deterministic behavior for gradient checkpointing init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py index b61d097d9e64..1e221b946b1e 100644 --- a/tests/test_pipelines.py +++ b/tests/test_pipelines.py @@ -108,6 +108,17 @@ def test_download_only_pytorch(self): # We need to never convert this tiny model to safetensors for this test to pass assert not any(f.endswith(".safetensors") for f in files) + def test_force_safetensors_error(self): + with tempfile.TemporaryDirectory() as tmpdirname: + # pipeline has Flax weights + with self.assertRaises(EnvironmentError): + tmpdirname = DiffusionPipeline.download( + "hf-internal-testing/tiny-stable-diffusion-pipe", + safety_checker=None, + cache_dir=tmpdirname, + use_safetensors=True, + ) + def test_returned_cached_folder(self): prompt = "hello" pipe = StableDiffusionPipeline.from_pretrained( From 7eb3a48318fea1983a75ac70acd83b1113a3baa5 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 14 Mar 2023 13:57:04 +0100 Subject: [PATCH 05/10] Style. --- src/diffusers/models/modeling_utils.py | 5 ++--- src/diffusers/pipelines/pipeline_utils.py | 7 +++---- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 9fc2afefd7d7..f99fbf93f849 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -394,9 +394,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P ignored when using `from_flax`. use_safetensors (`bool`, *optional* ): If set to `True`, the pipeline will forcibly load the models from `safetensors` weights. If set to - `None` (the default). The pipeline will load using `safetensors` if safetensors weights are - available *and* if `safetensors` is installed. If the to `False` the pipeline will *not* use - `safetensors`. + `None` (the default). The pipeline will load using `safetensors` if safetensors weights are available + *and* if `safetensors` is installed. If the to `False` the pipeline will *not* use `safetensors`. diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 4f66a770c483..ace5ab273dbd 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -697,10 +697,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P return_cached_folder (`bool`, *optional*, defaults to `False`): If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline. use_safetensors (`bool`, *optional* ): - If set to `True`, the pipeline will be loaded from `safetensors` weights. If set to - `None` (the default). The pipeline will load using `safetensors` if the safetensors weights are - available *and* if `safetensors` is installed. If the to `False` the pipeline will *not* use - `safetensors`. + If set to `True`, the pipeline will be loaded from `safetensors` weights. If set to `None` (the + default). The pipeline will load using `safetensors` if the safetensors weights are available *and* if + `safetensors` is installed. If the to `False` the pipeline will *not* use `safetensors`. kwargs (remaining dictionary of keyword arguments, *optional*): Can be used to overwrite load - and saveable variables - *i.e.* the pipeline components - of the specific pipeline class. The overwritten components are then directly passed to the pipelines From 756063faa4954f5493294b74794d4cdce4943aa7 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Tue, 14 Mar 2023 21:28:35 +0100 Subject: [PATCH 06/10] Apply suggestions from code review --- src/diffusers/pipelines/pipeline_utils.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index ace5ab273dbd..f0048de83654 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -694,8 +694,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, setting this argument to `True` will raise an error. - return_cached_folder (`bool`, *optional*, defaults to `False`): - If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline. use_safetensors (`bool`, *optional* ): If set to `True`, the pipeline will be loaded from `safetensors` weights. If set to `None` (the default). The pipeline will load using `safetensors` if the safetensors weights are available *and* if @@ -758,7 +756,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P device_map = kwargs.pop("device_map", None) low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) variant = kwargs.pop("variant", None) - return_cached_folder = kwargs.pop("return_cached_folder", False) kwargs.pop("use_safetensors", None if is_safetensors_available() else False) # 1. Download the checkpoints and configs From f86261bf8a2edc4796d195c71898964e62dd00b7 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Wed, 15 Mar 2023 12:58:13 +0100 Subject: [PATCH 07/10] Style. --- 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 430e6519d03e..5eeefa3bd971 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -170,7 +170,7 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict user_agent=user_agent, ) state_dict = safetensors.torch.load_file(model_file, device="cpu") - except EnvironmentError: + except EnvironmentError as e: if use_safetensors is True: raise e # try loading non-safetensors weights From 3a2899fc064eab1724ddfe764a88dede63af3481 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Thu, 16 Mar 2023 07:48:50 +0100 Subject: [PATCH 08/10] Addressing comments. --- src/diffusers/loaders.py | 18 +++++++++++++---- src/diffusers/models/modeling_utils.py | 18 +++++++++++++---- src/diffusers/pipelines/pipeline_utils.py | 20 ++++++++++++++++--- tests/models/test_models_unet_2d_condition.py | 3 ++- 4 files changed, 47 insertions(+), 12 deletions(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 5eeefa3bd971..bb76c6fe27e2 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -142,7 +142,17 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict revision = kwargs.pop("revision", None) subfolder = kwargs.pop("subfolder", None) weight_name = kwargs.pop("weight_name", None) - use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) + use_safetensors = kwargs.pop("use_safetensors", None) + + if use_safetensors and not is_safetensors_available(): + raise ValueError( + "`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors" + ) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = is_safetensors_available() + allow_pickle = True user_agent = { "file_type": "attn_procs_weights", @@ -152,7 +162,7 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict model_file = None if not isinstance(pretrained_model_name_or_path_or_dict, dict): # Let's first try to load .safetensors weights - if (use_safetensors is not False and weight_name is None) or ( + if (use_safetensors and weight_name is None) or ( weight_name is not None and weight_name.endswith(".safetensors") ): try: @@ -170,8 +180,8 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict user_agent=user_agent, ) state_dict = safetensors.torch.load_file(model_file, device="cpu") - except EnvironmentError as e: - if use_safetensors is True: + except IOError as e: + if not allow_pickle: raise e # try loading non-safetensors weights pass diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index f99fbf93f849..a7c2c750f654 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -427,7 +427,17 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P device_map = kwargs.pop("device_map", None) low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) variant = kwargs.pop("variant", None) - use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) + use_safetensors = kwargs.pop("use_safetensors", None) + + if use_safetensors and not is_safetensors_available(): + raise ValueError( + "`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors" + ) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = is_safetensors_available() + allow_pickle = True if low_cpu_mem_usage and not is_accelerate_available(): low_cpu_mem_usage = False @@ -514,7 +524,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P model = load_flax_checkpoint_in_pytorch_model(model, model_file) else: - if use_safetensors is not False: + if use_safetensors: try: model_file = _get_model_file( pretrained_model_name_or_path, @@ -530,8 +540,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P user_agent=user_agent, commit_hash=commit_hash, ) - except Exception as e: - if use_safetensors is True: + except IOError as e: + if not allow_pickle: raise e pass if model_file is None: diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index f0048de83654..6560f305c18e 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -1073,7 +1073,17 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: from_flax = kwargs.pop("from_flax", False) custom_pipeline = kwargs.pop("custom_pipeline", None) variant = kwargs.pop("variant", None) - use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) + use_safetensors = kwargs.pop("use_safetensors", None) + + if use_safetensors and not is_safetensors_available(): + raise ValueError( + "`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors" + ) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = is_safetensors_available() + allow_pickle = True pipeline_is_cached = False allow_patterns = None @@ -1129,13 +1139,17 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: CUSTOM_PIPELINE_FILE_NAME, ] - if use_safetensors is True and not is_safetensors_compatible(model_filenames, variant=variant): + if ( + use_safetensors + and not allow_pickle + and not is_safetensors_compatible(model_filenames, variant=variant) + ): raise EnvironmentError( f"Could not found the necessary `safetensors` weights in {model_filenames} (variant={variant})" ) if from_flax: ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] - elif use_safetensors is not False and is_safetensors_compatible(model_filenames, variant=variant): + elif use_safetensors and is_safetensors_compatible(model_filenames, variant=variant): ignore_patterns = ["*.bin", "*.msgpack"] safetensors_variant_filenames = set([f for f in variant_filenames if f.endswith(".safetensors")]) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index 4ddea0766407..0588cd70b600 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -517,8 +517,9 @@ def test_lora_save_torch_force_load_safetensors_error(self): torch.manual_seed(0) new_model = self.model_class(**init_dict) new_model.to(torch_device) - with self.assertRaises(EnvironmentError): + with self.assertRaises(IOError) as e: new_model.load_attn_procs(tmpdirname, use_safetensors=True) + self.assertIn("Error no file named pytorch_lora_weights.safetensors", str(e.exception)) def test_lora_on_off(self): # enable deterministic behavior for gradient checkpointing From 9d35f16c42d58fa0df7cfe699f671cb87303baf3 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Thu, 16 Mar 2023 07:50:22 +0100 Subject: [PATCH 09/10] Update tests/test_pipelines.py Co-authored-by: Will Berman --- tests/test_pipelines.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py index 1e221b946b1e..daf88417227f 100644 --- a/tests/test_pipelines.py +++ b/tests/test_pipelines.py @@ -113,7 +113,7 @@ def test_force_safetensors_error(self): # pipeline has Flax weights with self.assertRaises(EnvironmentError): tmpdirname = DiffusionPipeline.download( - "hf-internal-testing/tiny-stable-diffusion-pipe", + "hf-internal-testing/tiny-stable-diffusion-pipe-no-safetensors", safety_checker=None, cache_dir=tmpdirname, use_safetensors=True, From 89dadb77f7e46eb5542a5f79be4dd89f774fd741 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Thu, 16 Mar 2023 07:58:41 +0100 Subject: [PATCH 10/10] Black ??? --- tests/models/test_models_unet_2d_condition.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index 3bbb51d86130..5ca0bae66c02 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -497,9 +497,7 @@ def test_lora_save_torch_force_load_safetensors_error(self): block_id = int(name[len("down_blocks.")]) hidden_size = model.config.block_out_channels[block_id] - lora_attn_procs[name] = LoRACrossAttnProcessor( - hidden_size=hidden_size, cross_attention_dim=cross_attention_dim - ) + lora_attn_procs[name] = LoRAAttnProcessor(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim) lora_attn_procs[name] = lora_attn_procs[name].to(model.device) model.set_attn_processor(lora_attn_procs)