Add warning, change _ to default

Author: hlky

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -555,7 +555,7 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic
 
         else:
             sub_model_dtype = (
-                torch_dtype.get(name, torch_dtype.get("_", torch.float32))
+                torch_dtype.get(name, torch_dtype.get("default", torch.float32))
                 if isinstance(torch_dtype, dict)
                 else torch_dtype
             )
@@ -585,7 +585,7 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic
     module_sizes = {
         module_name: compute_module_sizes(
             module,
-            dtype=torch_dtype.get(module_name, torch_dtype.get("_", torch.float32))
+            dtype=torch_dtype.get(module_name, torch_dtype.get("default", torch.float32))
             if isinstance(torch_dtype, dict)
             else torch_dtype,
         )[""]

src/diffusers/pipelines/pipeline_utils.py

@@ -534,8 +534,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 Override the default `torch.dtype` and load the model with another dtype. If "auto" is passed, the
                 dtype is automatically derived from the model's weights. To load submodels with different dtype pass a
                 `dict` (for example `{'transformer': torch.bfloat16, 'vae': torch.float16}`). Set the default dtype for
-                unspecified components with `_` (for example `{'transformer': torch.bfloat16, '_': torch.float16}`). If
-                a component is not specifed and no default is set, `torch.float32` is used.
+                unspecified components with `default` (for example `{'transformer': torch.bfloat16, 'default': torch.float16}`).
+                If a component is not specified and no default is set, `torch.float32` is used.
             custom_pipeline (`str`, *optional*):
 
                 <Tip warning={true}>
@@ -858,6 +858,20 @@ def load_module(name, value):
                     f"Expected types for {key}: {_expected_class_types}, got {class_obj.__class__.__name__}."
                 )
 
+        # Check `torch_dtype` map for unused keys
+        if isinstance(torch_dtype, dict):
+            extra_keys_dtype = set(torch_dtype.keys()) - set(passed_class_obj.keys())
+            extra_keys_obj = set(passed_class_obj.keys()) - set(torch_dtype.keys())
+            if len(extra_keys_dtype) > 0:
+                logger.warning(
+                    f"Expected `{list(passed_class_obj.keys())}`, got extra `torch_dtype` keys `{extra_keys_dtype}`."
+                )
+            if len(extra_keys_obj) > 0:
+                logger.warning(
+                    f"Expected `{list(passed_class_obj.keys())}`, missing `torch_dtype` keys `{extra_keys_dtype}`."
+                    " using `default` or `torch.float32`."
+                )
+
         # Special case: safety_checker must be loaded separately when using `from_flax`
         if from_flax and "safety_checker" in init_dict and "safety_checker" not in passed_class_obj:
             raise NotImplementedError(
@@ -925,7 +939,7 @@ def load_module(name, value):
             else:
                 # load sub model
                 sub_model_dtype = (
-                    torch_dtype.get(name, torch_dtype.get("_", torch.float32))
+                    torch_dtype.get(name, torch_dtype.get("default", torch.float32))
                     if isinstance(torch_dtype, dict)
                     else torch_dtype
                 )