Flax: Ignore PyTorch, ONNX files when they coexist with Flax weights (h…

…uggingface#5237) Ignore PyTorch, ONNX files when they coexist with Flax weights
yoonseokjin · Oct 2, 2023 · 97b9349 · 97b9349
1 parent 63cee38
commit 97b9349
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pipelines/pipeline_flax_utils.py b/pipelines/pipeline_flax_utils.py
@@ -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__