update

Author: DN6

src/diffusers/loaders/single_file_model.py

@@ -205,6 +205,7 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
         revision = kwargs.pop("revision", None)
         torch_dtype = kwargs.pop("torch_dtype", None)
         quantization_config = kwargs.pop("quantization_config", None)
+        device = kwargs.pop("device", None)
 
         if isinstance(pretrained_model_link_or_path_or_dict, dict):
             checkpoint = pretrained_model_link_or_path_or_dict
@@ -326,10 +327,12 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
             )
 
         if is_accelerate_available():
+            param_device = torch.device(device) if device else torch.device("cpu")
             unexpected_keys = load_model_dict_into_meta(
                 model,
                 diffusers_format_checkpoint,
                 dtype=torch_dtype,
+                device=param_device,
                 hf_quantizer=hf_quantizer,
                 keep_in_fp32_modules=keep_in_fp32_modules,
             )

src/diffusers/models/model_loading_utils.py

@@ -184,7 +184,8 @@ def load_model_dict_into_meta(
 ) -> List[str]:
     if device is not None and not isinstance(device, (str, torch.device)):
         raise ValueError(f"Expected device to have type `str` or `torch.device`, but got {type(device)=}.")
-    device = device or torch.device("cpu")
+    if hf_quantizer is None:
+        device = device or torch.device("cpu")
     dtype = dtype or torch.float32
     is_quantized = hf_quantizer is not None