Spaces:
Running on Zero
Running on Zero
| from safetensors import safe_open | |
| import torch, hashlib | |
| def load_state_dict(file_path, torch_dtype=None, device="cpu"): | |
| if isinstance(file_path, list): | |
| state_dict = {} | |
| for file_path_ in file_path: | |
| state_dict.update(load_state_dict(file_path_, torch_dtype, device)) | |
| return state_dict | |
| if file_path.endswith(".safetensors"): | |
| return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) | |
| else: | |
| return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device) | |
| def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"): | |
| state_dict = {} | |
| with safe_open(file_path, framework="pt", device=str(device)) as f: | |
| for k in f.keys(): | |
| state_dict[k] = f.get_tensor(k) | |
| if torch_dtype is not None: | |
| state_dict[k] = state_dict[k].to(torch_dtype) | |
| return state_dict | |
| def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"): | |
| state_dict = torch.load(file_path, map_location=device, weights_only=True) | |
| if len(state_dict) == 1: | |
| if "state_dict" in state_dict: | |
| state_dict = state_dict["state_dict"] | |
| elif "module" in state_dict: | |
| state_dict = state_dict["module"] | |
| elif "model_state" in state_dict: | |
| state_dict = state_dict["model_state"] | |
| if torch_dtype is not None: | |
| for i in state_dict: | |
| if isinstance(state_dict[i], torch.Tensor): | |
| state_dict[i] = state_dict[i].to(torch_dtype) | |
| return state_dict | |
| def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): | |
| keys = [] | |
| for key, value in state_dict.items(): | |
| if isinstance(key, str): | |
| if isinstance(value, torch.Tensor): | |
| if with_shape: | |
| shape = "_".join(map(str, list(value.shape))) | |
| keys.append(key + ":" + shape) | |
| keys.append(key) | |
| elif isinstance(value, dict): | |
| keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) | |
| keys.sort() | |
| keys_str = ",".join(keys) | |
| return keys_str | |
| def hash_state_dict_keys(state_dict, with_shape=True): | |
| keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) | |
| keys_str = keys_str.encode(encoding="UTF-8") | |
| return hashlib.md5(keys_str).hexdigest() | |
| def load_keys_dict(file_path): | |
| if isinstance(file_path, list): | |
| state_dict = {} | |
| for file_path_ in file_path: | |
| state_dict.update(load_keys_dict(file_path_)) | |
| return state_dict | |
| if file_path.endswith(".safetensors"): | |
| return load_keys_dict_from_safetensors(file_path) | |
| else: | |
| return load_keys_dict_from_bin(file_path) | |
| def load_keys_dict_from_safetensors(file_path): | |
| keys_dict = {} | |
| with safe_open(file_path, framework="pt", device="cpu") as f: | |
| for k in f.keys(): | |
| keys_dict[k] = f.get_slice(k).get_shape() | |
| return keys_dict | |
| def convert_state_dict_to_keys_dict(state_dict): | |
| keys_dict = {} | |
| for k, v in state_dict.items(): | |
| if isinstance(v, torch.Tensor): | |
| keys_dict[k] = list(v.shape) | |
| else: | |
| keys_dict[k] = convert_state_dict_to_keys_dict(v) | |
| return keys_dict | |
| def load_keys_dict_from_bin(file_path): | |
| state_dict = load_state_dict_from_bin(file_path) | |
| keys_dict = convert_state_dict_to_keys_dict(state_dict) | |
| return keys_dict | |
| def convert_keys_dict_to_single_str(state_dict, with_shape=True): | |
| keys = [] | |
| for key, value in state_dict.items(): | |
| if isinstance(key, str): | |
| if isinstance(value, dict): | |
| keys.append(key + "|" + convert_keys_dict_to_single_str(value, with_shape=with_shape)) | |
| else: | |
| if with_shape: | |
| shape = "_".join(map(str, list(value))) | |
| keys.append(key + ":" + shape) | |
| keys.append(key) | |
| keys.sort() | |
| keys_str = ",".join(keys) | |
| return keys_str | |
| def hash_model_file(path, with_shape=True): | |
| keys_dict = load_keys_dict(path) | |
| keys_str = convert_keys_dict_to_single_str(keys_dict, with_shape=with_shape) | |
| keys_str = keys_str.encode(encoding="UTF-8") | |
| return hashlib.md5(keys_str).hexdigest() | |