Embed model merge metadata in .safetensors file

This commit is contained in:
space-nuko
2023-04-02 17:41:55 -05:00
parent 22bcc7be42
commit d132481058
3 changed files with 55 additions and 4 deletions

View File

@@ -52,6 +52,15 @@ class CheckpointInfo:
self.ids = [self.hash, self.model_name, self.title, name, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] if self.shorthash else [])
self.metadata = {}
_, ext = os.path.splitext(self.filename)
if ext.lower() == ".safetensors":
try:
self.metadata = read_metadata_from_safetensors(filename)
except Exception as e:
errors.display(e, f"reading checkpoint metadata: {filename}")
def register(self):
checkpoints_list[self.title] = self
for id in self.ids:
@@ -544,4 +553,4 @@ def unload_model_weights(sd_model=None, info=None):
print(f"Unloaded weights {timer.summary()}.")
return sd_model
return sd_model