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)