download

Author: hlky

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -28,6 +28,7 @@
 
 from .. import __version__
 from ..utils import (
+    FLASHPACK_WEIGHTS_NAME,
     FLAX_WEIGHTS_NAME,
     ONNX_EXTERNAL_WEIGHTS_NAME,
     ONNX_WEIGHTS_NAME,
@@ -194,6 +195,7 @@ def filter_model_files(filenames):
         FLAX_WEIGHTS_NAME,
         ONNX_WEIGHTS_NAME,
         ONNX_EXTERNAL_WEIGHTS_NAME,
+        FLASHPACK_WEIGHTS_NAME,
     ]
 
     if is_transformers_available():
@@ -1091,6 +1093,7 @@ def _get_ignore_patterns(
     allow_pickle: bool,
     use_onnx: bool,
     is_onnx: bool,
+    use_flashpack: bool,
     variant: Optional[str] = None,
 ) -> List[str]:
     if (
@@ -1116,6 +1119,9 @@ def _get_ignore_patterns(
         if not use_onnx:
             ignore_patterns += ["*.onnx", "*.pb"]
 
+    elif use_flashpack:
+        ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb", "*.msgpack"]
+
     else:
         ignore_patterns = ["*.safetensors", "*.msgpack"]
 

src/diffusers/pipelines/pipeline_utils.py

@@ -1551,6 +1551,9 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
                 Whether or not to allow for custom pipelines and components defined on the Hub in their own files. This
                 option should only be set to `True` for repositories you trust and in which you have read the code, as
                 it will execute code present on the Hub on your local machine.
+            use_onnx (`bool`, *optional*, defaults to `False`):
+                If set to `True`, FlashPack weights will always be downloaded if present. If set to `False`, ONNX weights
+                will never be downloaded.
 
         Returns:
             `os.PathLike`:
@@ -1575,6 +1578,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
         load_connected_pipeline = kwargs.pop("load_connected_pipeline", False)
         trust_remote_code = kwargs.pop("trust_remote_code", False)
         dduf_file: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_file", None)
+        use_flashpack = kwargs.pop("use_flashpack", True)
 
         if dduf_file:
             if custom_pipeline:
@@ -1694,6 +1698,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
                 allow_pickle,
                 use_onnx,
                 pipeline_class._is_onnx,
+                use_flashpack,
                 variant,
             )