cleanup

Author: danhipke

src/diffusers/loaders/single_file.py

@@ -106,7 +106,7 @@ def load_single_file_sub_model(
             subfolder=name,
             torch_dtype=torch_dtype,
             local_files_only=local_files_only,
-            no_mmap=True,
+            no_mmap=no_mmap,
             **kwargs,
         )
 
@@ -310,7 +310,8 @@ def from_single_file(cls, pretrained_model_link_or_path, **kwargs):
                     - A path to a *directory* (for example `./my_pipeline_directory/`) containing the pipeline
                       component configs in Diffusers format.
             no_mmap ('bool', *optional*, defaults to 'False'):
-                Whether to use mmap when loading the model.
+                Whether to disable mmap when loading the model. This option can perform better when the model is on
+                a network mount or hard drive, which may not handle the seeky-ness of mmap very well.
             kwargs (remaining dictionary of keyword arguments, *optional*):
                 Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline
                 class). The overwritten components are passed directly to the pipelines `__init__` method. See example
@@ -509,6 +510,7 @@ def load_module(name, value):
                         original_config=original_config,
                         local_files_only=local_files_only,
                         is_legacy_loading=is_legacy_loading,
+                        no_mmap=no_mmap,
                         **kwargs,
                     )
                 except SingleFileComponentError as e:

src/diffusers/loaders/single_file_model.py

@@ -171,7 +171,8 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
                 The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
                 allowed by Git.
             no_mmap ('bool', *optional*, defaults to 'False'):
-                Whether to use mmap when loading the model.
+                Whether to disable mmap when loading the model. This option can perform better when the model is on
+                a network mount or hard drive, which may not handle the seeky-ness of mmap very well.
             kwargs (remaining dictionary of keyword arguments, *optional*):
                 Can be used to overwrite load and saveable variables (for example the pipeline components of the
                 specific pipeline class). The overwritten components are directly passed to the pipelines `__init__`
@@ -233,6 +234,7 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
             )
 
         mapping_functions = SINGLE_FILE_LOADABLE_CLASSES[mapping_class_name]
+
         checkpoint_mapping_fn = mapping_functions["checkpoint_mapping_fn"]
         if original_config is not None:
             if "config_mapping_fn" in mapping_functions: