Skip to content

Commit

Permalink
Individual tensor metadata object query tool
Browse files Browse the repository at this point in the history
```
$ python3 safetensors_m.py ss_network_module /path/to/file/summer_dress.safetensors
"networks.lora"

$ python3 safetensors_m.py nonexistent_module /path/to/file/summer_dress.safetensors
Error: Metadata does not contain a `nonexistent_module` item, did you spell it right?

$ python3 safetensors_m.py ss_network_module /path/to/file/weird_file.safetensors
Error: File header does not contain a `__metadata__` item

$ python3 safetensors_m.py ss_tag_frequency /path/to/file/trina.safetensors
{
    "6_trina": {
        "trina": 26, 
        " black hair": 20, 
        " hands on hips": 1, 
        " looking at viewer": 7
    }
}
```

This is a proof of concept utility for retrieving the value of a named object inside `__metadata__` from a safetensors file, with the intention of eventually moving the code into `safetensors_util.py` as its own separate `cli.command()` block. I wrote this because so far I haven't seen any other utilities for retrieving these values from safetensors files, e.g. checking `ss_network_module` to tell whether a file is LoRA (`networks.lora`) or Lycoris (`lycoris.kohya`) before trying to install it into a potentially incompatible SD client, or just debugging troublesome safetensors files without having to dump/parse the entire metadata block. Applies `_ParseMore()` to all output for readability's sake but this could be optionally disabled.
  • Loading branch information
duanemoody authored Dec 1, 2023
1 parent f6bc80a commit e320547
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions safetensors_object.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
import json, sys
from safetensors_file import SafeTensorsFile
from safetensors_worker import _ParseMore

def get_object(tensorsfile: str) -> str:
s = SafeTensorsFile.open_file(tensorsfile, quiet=True)
js = s.get_header()
s.close_file()

if "__metadata__" not in js:
return "Error: File header does not contain a `__metadata__` item"
md = js["__metadata__"]
if md_object not in md:
return f'Error: Metadata does not contain a `{md_object}` item, did you spell it right?'
_ParseMore(md) # pretty print the metadata
stf = md[md_object]
return json.dumps(stf, ensure_ascii=False, separators=(', ', ': '), indent=4)

md_object = sys.argv[1]
tensorsfile = sys.argv[2]
hdata = get_object(tensorsfile)

print(hdata)

0 comments on commit e320547

Please sign in to comment.