| 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:
|
|
|
|
|
| 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)
|
|
|