mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
47 lines
1.4 KiB
Python
Executable File
47 lines
1.4 KiB
Python
Executable File
#!/usr/bin/env python
|
|
import os
|
|
import sys
|
|
import json
|
|
from rich import print # pylint: disable=redefined-builtin
|
|
|
|
|
|
def read_metadata(fn):
|
|
res = {}
|
|
if not fn.lower().endswith(".safetensors"):
|
|
return
|
|
with open(fn, mode="rb") as f:
|
|
try:
|
|
metadata_len = f.read(8)
|
|
metadata_len = int.from_bytes(metadata_len, "little")
|
|
json_start = f.read(2)
|
|
if metadata_len <= 2 or json_start not in (b'{"', b"{'"):
|
|
print(f"Not a valid safetensors file: {fn}")
|
|
json_data = json_start + f.read(metadata_len-2)
|
|
json_obj = json.loads(json_data)
|
|
for k, v in json_obj.get("__metadata__", {}).items():
|
|
res[k] = v
|
|
if isinstance(v, str) and v[0:1] == '{':
|
|
try:
|
|
res[k] = json.loads(v)
|
|
except Exception:
|
|
pass
|
|
print(f"{fn}: {json.dumps(res, indent=4)}")
|
|
except Exception:
|
|
print(f"{fn}: cannot read metadata")
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) == 0:
|
|
print('metadata:', 'no files specified')
|
|
for fn in sys.argv:
|
|
if os.path.isfile(fn):
|
|
read_metadata(fn)
|
|
elif os.path.isdir(fn):
|
|
for root, _dirs, files in os.walk(fn):
|
|
for file in files:
|
|
read_metadata(os.path.join(root, file))
|
|
|
|
if __name__ == '__main__':
|
|
sys.argv.pop(0)
|
|
main()
|