Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,565 Bytes
3d5837a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import hashlib
import json
from aiohttp import web
from server import PromptServer
import folder_paths
import os
def get_metadata(filepath):
with open(filepath, "rb") as file:
# https://github.com/huggingface/safetensors#format
# 8 bytes: N, an unsigned little-endian 64-bit integer, containing the size of the header
header_size = int.from_bytes(file.read(8), "little", signed=False)
if header_size <= 0:
raise BufferError("Invalid header size")
header = file.read(header_size)
if header_size <= 0:
raise BufferError("Invalid header")
header_json = json.loads(header)
return header_json["__metadata__"] if "__metadata__" in header_json else None
@PromptServer.instance.routes.post("/pysssss/metadata/notes/{name}")
async def save_notes(request):
name = request.match_info["name"]
pos = name.index("/")
type = name[0:pos]
name = name[pos+1:]
file_path = None
if type == "embeddings" or type == "loras":
name = name.lower()
files = folder_paths.get_filename_list(type)
for f in files:
lower_f = f.lower()
if lower_f == name:
file_path = folder_paths.get_full_path(type, f)
else:
n = os.path.splitext(f)[0].lower()
if n == name:
file_path = folder_paths.get_full_path(type, f)
if file_path is not None:
break
else:
file_path = folder_paths.get_full_path(
type, name)
if not file_path:
return web.Response(status=404)
file_no_ext = os.path.splitext(file_path)[0]
info_file = file_no_ext + ".txt"
with open(info_file, "w") as f:
f.write(await request.text())
return web.Response(status=200)
@PromptServer.instance.routes.get("/pysssss/metadata/{name}")
async def load_metadata(request):
name = request.match_info["name"]
pos = name.index("/")
type = name[0:pos]
name = name[pos+1:]
file_path = None
if type == "embeddings" or type == "loras":
name = name.lower()
files = folder_paths.get_filename_list(type)
for f in files:
lower_f = f.lower()
if lower_f == name:
file_path = folder_paths.get_full_path(type, f)
else:
n = os.path.splitext(f)[0].lower()
if n == name:
file_path = folder_paths.get_full_path(type, f)
if file_path is not None:
break
else:
file_path = folder_paths.get_full_path(
type, name)
if not file_path:
return web.Response(status=404)
try:
meta = get_metadata(file_path)
except:
meta = None
if meta is None:
meta = {}
file_no_ext = os.path.splitext(file_path)[0]
info_file = file_no_ext + ".txt"
if os.path.isfile(info_file):
with open(info_file, "r") as f:
meta["pysssss.notes"] = f.read()
hash_file = file_no_ext + ".sha256"
if os.path.isfile(hash_file):
with open(hash_file, "rt") as f:
meta["pysssss.sha256"] = f.read()
else:
with open(file_path, "rb") as f:
meta["pysssss.sha256"] = hashlib.sha256(f.read()).hexdigest()
with open(hash_file, "wt") as f:
f.write(meta["pysssss.sha256"])
return web.json_response(meta)
|