Swapped to using built-in safetensors header loader.

- Build-in loader avoids worst-case slow load using max length.
This commit is contained in:
Christian Bastian
2024-08-19 09:56:02 -04:00
parent e945be8e3e
commit d2a9a3a1fa

View File

@@ -18,6 +18,7 @@ import json
import requests
requests.packages.urllib3.disable_warnings()
import comfy.utils
import folder_paths
comfyui_model_uri = folder_paths.models_dir
@@ -142,11 +143,9 @@ def search_path_to_system_path(model_path):
def get_safetensor_header(path):
try:
with open(path, "rb") as f:
length_of_header = struct.unpack("<Q", f.read(8))[0]
header_bytes = f.read(length_of_header)
header_json = json.loads(header_bytes)
return header_json
header_bytes = comfy.utils.safetensors_header(path)
header_json = json.loads(header_bytes)
return header_json if header_json is not None else {}
except:
return {}