Skip to content

Commit

Permalink
Add files via upload
Browse files Browse the repository at this point in the history 
Proof-of-concept utility for retrieving a named JSON object from `__metadata__` by passing it and the safetensors filename.
  • Loading branch information
@duanemoody
duanemoody authored Dec 1, 2023
1 parent ae554ca commit ce53e31
Showing 1 changed file with 22 additions and 0 deletions.
22 changes: 22 additions & 0 deletions safetensors_m.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
import json, sys
from safetensors_file import SafeTensorsFile
from safetensors_worker import _ParseMore

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

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)

1 comment on commit ce53e31

@duanemoody
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

$ 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.

Currently depends on the context manager protocol modification to safetensors_file.py but could be rewritten to use the conventional SafeTensorsFile.open_file() / SafeTensorsFile.close() methods if that PR isn't approved/merged. I thought I'd reverted that before submitting it.

Please sign in to comment.