Skip to content

Commit

Permalink
Flax: Ignore PyTorch, ONNX files when they coexist with Flax weights (#…
Browse files Browse the repository at this point in the history
…5237)

Ignore PyTorch, ONNX files when they coexist with Flax weights
  • Loading branch information
pcuenca authored Oct 2, 2023
1 parent 84e5cc5 commit 0c7cb9a
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/diffusers/pipelines/pipeline_flax_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -341,8 +341,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
allow_patterns = [os.path.join(k, "*") for k in folder_names]
allow_patterns += [FLAX_WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, cls.config_name]

# make sure we don't download PyTorch weights, unless when using from_pt
ignore_patterns = "*.bin" if not from_pt else []
ignore_patterns = ["*.bin", "*.safetensors"] if not from_pt else []
ignore_patterns += ["*.onnx", "*.onnx_data", "*.xml", "*.pb"]

if cls != FlaxDiffusionPipeline:
requested_pipeline_class = cls.__name__
Expand Down

0 comments on commit 0c7cb9a

Please sign in to comment.