propagate no_mmap

danhipke · danhipke · commit 7231c282e0e7 · 2024-12-18T17:39:20.000Z
diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py
@@ -318,7 +318,7 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
             unexpected_keys = load_model_dict_into_meta(model, diffusers_format_checkpoint, dtype=torch_dtype)
 
         else:
-            _, unexpected_keys = model.load_state_dict(diffusers_format_checkpoint, strict=False)
+            _, unexpected_keys = model.load_state_dict(diffusers_format_checkpoint, strict=False, no_mmap=no_mmap)
 
         if model._keys_to_ignore_on_load_unexpected is not None:
             for pat in model._keys_to_ignore_on_load_unexpected:
diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py
@@ -382,7 +382,8 @@ def load_single_file_checkpoint(
             revision=revision,
         )
 
-    checkpoint = load_state_dict(pretrained_model_link_or_path, no_mmap)
+    logger.warn("before load_state_dict. no_mmap: " + str(no_mmap))
+    checkpoint = load_state_dict(pretrained_model_link_or_path, no_mmap=no_mmap)
 
     # some checkpoints contain the model state dict under a "state_dict" key
     while "state_dict" in checkpoint:

Original file line number	Diff line number	Diff line change
`@@ -382,7 +382,8 @@ def load_single_file_checkpoint(`
`382`	`382`	`revision=revision,`
`383`	`383`	`)`
`384`	`384`
`385`		`- checkpoint = load_state_dict(pretrained_model_link_or_path, no_mmap)`
	`385`	`+ logger.warn("before load_state_dict. no_mmap: " + str(no_mmap))`
	`386`	`+ checkpoint = load_state_dict(pretrained_model_link_or_path, no_mmap=no_mmap)`
`386`	`387`
`387`	`388`	`# some checkpoints contain the model state dict under a "state_dict" key`
`388`	`389`	`while "state_dict" in checkpoint:`