From d2a9a3a1fadf0fd3451e5928dafadb2da68ed126 Mon Sep 17 00:00:00 2001
From: Christian Bastian <80225746+cdb-boop@users.noreply.github.com>
Date: Mon, 19 Aug 2024 09:56:02 -0400
Subject: [PATCH] Swapped to using built-in safetensors header loader. -
 Build-in loader avoids worst-case slow load using max length.

---
 __init__.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/__init__.py b/__init__.py
index 93511fc..6aadb9b 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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 {}