Repository: huggingface/diffusers
File: src/diffusers/models/modeling_utils.py (Python)

Code before:
    # Taken from `transformers`.
    def float(self, *args):
        # Checks if the model is quantized
        if getattr(self, "is_quantized", False):
            raise ValueError(
                "`.float()` is not supported for quantized model. Please use the model as it is, since the"
                " model has already been cast to the correct `dtype`."
            )
        else:
            return super().float(*args)

    @classmethod
    def _load_pretrained_model(
        cls,
        model,
        state_dict: OrderedDict,
        resolved_model_file: List[str],
        pretrained_model_name_or_path: Union[str, os.PathLike],
        loaded_keys: List[str],
        ignore_mismatched_sizes: bool = False,
        assign_to_params_buffers: bool = False,
        hf_quantizer: Optional[DiffusersQuantizer] = None,
        low_cpu_mem_usage: bool = True,
        dtype: Optional[Union[str, torch.dtype]] = None,
        keep_in_fp32_modules: Optional[List[str]] = None,
        device_map: Dict[Union[str, int, torch.device], Union[int, str, torch.device]] = None,
        offload_state_dict: Optional[bool] = None,
        offload_folder: Optional[Union[str, os.PathLike]] = None,
        dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
    ):
        model_state_dict = model.state_dict()
        expected_keys = list(model_state_dict.keys())
        missing_keys = list(set(expected_keys) - set(loaded_keys))
        if hf_quantizer is not None:
            missing_keys = hf_quantizer.update_missing_keys(model, missing_keys, prefix="")
        unexpected_keys = list(set(loaded_keys) - set(expected_keys))
        # Some models may have keys that are not in the state by design, removing them before needlessly warning
        # the user.
        if cls._keys_to_ignore_on_load_unexpected is not None:
            for pat in cls._keys_to_ignore_on_load_unexpected:
                unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

        mismatched_keys = []

        assign_to_params_buffers = None
        error_msgs = []

        # Deal with offload
        if device_map is not None and "disk" in device_map.values():
            if offload_folder is None:
                raise ValueError(

Review comment (suggestion):
Just confirming if this change is correct. Can the dictionary key be a str/int/torch.device? I thought only the values can be that

Code after:
                "`.half()` is not supported for quantized model. Please use the model as it is, since the"
                " model has already been cast to the correct `dtype`."
            )
        else:
            return super().half(*args)

    # Taken from `transformers`.
    def float(self, *args):
        # Checks if the model is quantized
        if getattr(self, "is_quantized", False):
            raise ValueError(
                "`.float()` is not supported for quantized model. Please use the model as it is, since the"
                " model has already been cast to the correct `dtype`."
            )
        else:
            return super().float(*args)

    @classmethod
    def _load_pretrained_model(
        cls,
        model,
        state_dict: OrderedDict,
        resolved_model_file: List[str],
        pretrained_model_name_or_path: Union[str, os.PathLike],
        loaded_keys: List[str],
        ignore_mismatched_sizes: bool = False,
        assign_to_params_buffers: bool = False,
        hf_quantizer: Optional[DiffusersQuantizer] = None,
        low_cpu_mem_usage: bool = True,
        dtype: Optional[Union[str, torch.dtype]] = None,
        keep_in_fp32_modules: Optional[List[str]] = None,
        device_map: Union[str, int, torch.device, Dict[str, Union[int, str, torch.device]]] = None,
        offload_state_dict: Optional[bool] = None,
        offload_folder: Optional[Union[str, os.PathLike]] = None,
        dduf_entries: Optional[Dict[str, DDUFEntry]] = None,
    ):
        model_state_dict = model.state_dict()
        expected_keys = list(model_state_dict.keys())
        missing_keys = list(set(expected_keys) - set(loaded_keys))
        if hf_quantizer is not None:
            missing_keys = hf_quantizer.update_missing_keys(model, missing_keys, prefix="")
        unexpected_keys = list(set(loaded_keys) - set(expected_keys))
        # Some models may have keys that are not in the state by design, removing them before needlessly warning
        # the user.
        if cls._keys_to_ignore_on_load_unexpected is not None:
            for pat in cls._keys_to_ignore_on_load_unexpected:
                unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

        mismatched_keys = []

        assign_to_params_buffers = None