From b9818eb910b6ce683c38602c9b8fbd3979d97aaf Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 29 May 2023 02:48:50 -0400
Subject: [PATCH] Add route to get safetensors metadata:

/view_metadata/loras?filename=lora.safetensors
---
 comfy/utils.py  |  9 +++++++++
 folder_paths.py |  2 ++
 server.py       | 25 ++++++++++++++++++++++++-
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/comfy/utils.py b/comfy/utils.py
index 5ed9aaa0..4e84e870 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -1,5 +1,6 @@
 import torch
 import math
+import struct
 
 def load_torch_file(ckpt, safe_load=False):
     if ckpt.lower().endswith(".safetensors"):
@@ -50,6 +51,14 @@ def transformers_convert(sd, prefix_from, prefix_to, number):
                     sd[k_to] = weights[shape_from*x:shape_from*(x + 1)]
     return sd
 
+def safetensors_header(safetensors_path, max_size=100*1024*1024):
+    with open(safetensors_path, "rb") as f:
+        header = f.read(8)
+        length_of_header = struct.unpack('<Q', header)[0]
+        if length_of_header > max_size:
+            return None
+        return f.read(length_of_header)
+
 def bislerp(samples, width, height):
     def slerp(b1, b2, r):
         '''slerps batches b1, b2 according to ratio r, batches should be flat e.g. NxC'''
diff --git a/folder_paths.py b/folder_paths.py
index 20b461c9..19245a61 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -126,11 +126,13 @@ def filter_files_extensions(files, extensions):
 def get_full_path(folder_name, filename):
     global folder_names_and_paths
     folders = folder_names_and_paths[folder_name]
+    filename = os.path.relpath(os.path.join("/", filename), "/")
     for x in folders[0]:
         full_path = os.path.join(x, filename)
         if os.path.isfile(full_path):
             return full_path
 
+    return None
 
 def get_filename_list(folder_name):
     global folder_names_and_paths
diff --git a/server.py b/server.py
index c0f79cbd..72c565a6 100644
--- a/server.py
+++ b/server.py
@@ -22,7 +22,7 @@ except ImportError:
 
 import mimetypes
 from comfy.cli_args import args
-
+import comfy.utils
 
 @web.middleware
 async def cache_control(request: web.Request, handler):
@@ -257,6 +257,29 @@ class PromptServer():
 
             return web.Response(status=404)
 
+        @routes.get("/view_metadata/{folder_name}")
+        async def view_metadata(request):
+            folder_name = request.match_info.get("folder_name", None)
+            if folder_name is None:
+                return web.Response(status=404)
+            if not "filename" in request.rel_url.query:
+                return web.Response(status=404)
+
+            filename = request.rel_url.query["filename"]
+            if not filename.endswith(".safetensors"):
+                return web.Response(status=404)
+
+            safetensors_path = folder_paths.get_full_path(folder_name, filename)
+            if safetensors_path is None:
+                return web.Response(status=404)
+            out = comfy.utils.safetensors_header(safetensors_path, max_size=1024*1024)
+            if out is None:
+                return web.Response(status=404)
+            dt = json.loads(out)
+            if not "__metadata__" in dt:
+                return web.Response(status=404)
+            return web.json_response(dt["__metadata__"])
+
         @routes.get("/prompt")
         async def get_prompt(request):
             return web.json_response(self.get_queue_info())