-
Notifications
You must be signed in to change notification settings - Fork 6.1k
Adding use_safetensors
argument to give more control to users
#2123
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Adding use_safetensors
argument to give more control to users
#2123
Conversation
The documentation is not available anymore as the PR was closed or merged. |
I understand the need for the PR, but I wonder whether we could use the soon-to-be-added "variant" kwarg for this instead to not add too many kwargs. The idea with I think it might be better to handle all of this with the "variant" kwarg instead. E.g. from_pretrained(variant="safetensors") then safetensors has to be installed and the files have to be there or an error is thrown. We need to also align this with transformers cc @sgugger @LysandreJik . |
For Transformers, safetensors is the default if available, there is no special flag for it and we are not planning on adding one. We just look at |
@@ -350,6 +350,11 @@ 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. | |||
use_safetensors (`bool`, *optional*, defaults to `None`): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
use_safetensors (`bool`, *optional*, defaults to `None`): | |
use_safetensors (`bool`, *optional*): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actually upon section reflection, I think this makes sense.
Can we also add tests though?
And also would maybe save_loading
suit better (since we have safe_serialization
?)
@@ -350,6 +350,11 @@ 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. | |||
use_safetensors (`bool`, *optional*, defaults to `None`): | |||
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to | |
If set to `True`, the pipeline will forcibly load the model from `safetensors` weights. If set to |
@@ -350,6 +350,11 @@ 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. | |||
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are | |
`None` (the default). The pipeline will load using `safetensors` if safetensors weights are |
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use | |
available *and* if `safetensors` is installed. If set to `False` the pipeline will *not* use |
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. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
`safetensors` at all. | |
`safetensors`. |
@@ -400,6 +400,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`): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
use_safetensors (`bool`, *optional*, defaults to `None`): | |
use_safetensors (`bool`, *optional*): |
@@ -400,6 +400,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 | |||
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are | |
`None` (the default), the pipeline will be loaded from `safetensors` if the safetensors weights are |
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use | |
available *and* if `safetensors` is installed. If the to `False` the pipeline will *not* use |
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. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
`safetensors` at all. | |
`safetensors`. |
@@ -505,13 +511,13 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P | |||
|
|||
user_agent = http_user_agent(user_agent) | |||
|
|||
if is_safetensors_available(): | |||
if use_safetensors in {None, True}: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if use_safetensors in {None, True}: | |
if use_safetensors is not False: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
actually I would like to avoid the extra call to the Hub if safetensors
are not installed. Can we make sure that we still check for whether safetensors are installed?
info = model_info( | ||
pretrained_model_name_or_path, | ||
use_auth_token=use_auth_token, | ||
revision=revision, | ||
) | ||
if is_safetensors_compatible(info): | ||
if use_safetensors is True or is_safetensors_compatible(info): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if use_safetensors is True or is_safetensors_compatible(info): | |
if use_safetensors or is_safetensors_compatible(info): |
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread. Please note that issues that do not follow the contributing guidelines are likely to be ignored. |
Unstale |
@@ -463,7 +469,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}: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if use_safetensors in {None, True}: | |
if use_safetensors is not False: |
Still very much interested in merging this soon, but before merging this we probably need the same functionality in |
huggingface/transformers#22083 is in - think we can continue with this PR :-) cc @Narsil |
about which weights they use.
M1 failure not linked to this PR is it ? |
Thanks a lot @Narsil - looks good to me! |
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`. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Cool!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks!
Cool test cases.
src/diffusers/loaders.py
Outdated
if (is_safetensors_available() and weight_name is None) or ( | ||
if (use_safetensors is not False and weight_name is None) or ( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
hmm, is not False
is pretty confusing. With this and the above line,
use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False)
I'm a bit confused about what the default behavior of the function is supposed to be
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah ok I see the comment later on in the PR. Could we change the code (here and in the other places where we take use_safetensors as an argument with a default) to approximately
use_safetensors = kwargs.pop("use_safetensors", None)
if use_safetensors is None:
use_safetensors = is_safetensors_available()
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")
# in the rest of the function just use `use_safetensors` as a boolean
The rule of thumb here is that values which can be booleans or None can be hard to read/code with, especially when used in conditionals
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
But that's not it here.
If safetensors
is installed, but the weights do not exist on the remote, we don't want to fail, we want to use seamlessly the PT weights.
That's why it gets more complex.
False -> Use only PT (easy)
None -> Use safetensors if both available on system and remote (which we don't know when entering the function and without checking the remote)
True -> USe only safetensors (easy)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah ok, I think the same logic still mostly applies
We can move the error check first and add another boolean for the case when we are allowed to fallback to non-safetensors
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")
if use_safetensors is None:
use_safetensors = is_safetensors_available()
allow_pickle = True
# in the rest of the function just use `use_safetensors` as a boolean
# somewhere else
if use_safetensors and not_safetensors_compatible(...):
if not allow_pickle:
raise ValueError("... pickle not allowed ...")
# do the pickle loading
Is this correct or am I still missing something?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
2 booleans work, Not sure it's easier to read ( I to avoid variable that create nonsensical combos, for instance allow_pickle=True and use_safetensors=False doesn't make sense)
False
None
True
Seems easier to reason about, but indeed creates some unpythoness.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'll move to 2 booleans if you think it's clearer
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
t'would be much appreciated :)
src/diffusers/loaders.py
Outdated
except EnvironmentError: | ||
except EnvironmentError as e: | ||
if use_safetensors is True: | ||
raise e |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It's potentially not the job of this PR because EnvironmentError
was here before, but EnvironmentError
is a very broad error e.g. anything io related. Can we add some docs on why it's ok to catch it here if we're using safe tensors?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I can add some comment, but I don't know why it's EnvironmentError
, I'm guessing it could be IOError since it's a file missing error, but maybe we also want to catch other types of "missing error".
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If it's an IOError because a file is missing, should we just explicitly check for the file not being present? cc @patrickvonplaten I think you committed this, do you have context?
with self.assertRaises(EnvironmentError): | ||
new_model.load_attn_procs(tmpdirname, use_safetensors=True) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Super nit: Similar as above, EnvironmentError is really broad. Should we be checking against a specific error class or error message for this test case that safetensors could not be used?
Love this and think it makes a lot of sense :) Just a few nits on errors and how we handle default cases |
Co-authored-by: Will Berman <[email protected]>
@williamberman adressed all your comments I think. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
love it!
except EnvironmentError: | ||
except IOError as e: | ||
if not allow_pickle: | ||
raise e |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@patrickvonplaten before you merge can you double check this error change is ok?
Tbh, I thought it was better before, also because the code was more in line how we've merged it to Ok to merge now |
…ingface#2123) * Adding `use_safetensors` argument to give more control to users about which weights they use. * Doc style. * Rebased (not functional). * Rebased and functional with tests. * Style. * Apply suggestions from code review * Style. * Addressing comments. * Update tests/test_pipelines.py Co-authored-by: Will Berman <[email protected]> * Black ??? --------- Co-authored-by: Patrick von Platen <[email protected]> Co-authored-by: Will Berman <[email protected]>
…ingface#2123) * Adding `use_safetensors` argument to give more control to users about which weights they use. * Doc style. * Rebased (not functional). * Rebased and functional with tests. * Style. * Apply suggestions from code review * Style. * Addressing comments. * Update tests/test_pipelines.py Co-authored-by: Will Berman <[email protected]> * Black ??? --------- Co-authored-by: Patrick von Platen <[email protected]> Co-authored-by: Will Berman <[email protected]>
…ingface#2123) * Adding `use_safetensors` argument to give more control to users about which weights they use. * Doc style. * Rebased (not functional). * Rebased and functional with tests. * Style. * Apply suggestions from code review * Style. * Addressing comments. * Update tests/test_pipelines.py Co-authored-by: Will Berman <[email protected]> * Black ??? --------- Co-authored-by: Patrick von Platen <[email protected]> Co-authored-by: Will Berman <[email protected]>
about which weights they use.