Merge pull request #9312 from space-nuko/save-merge-recipe

Embed model merge metadata in .safetensors file

1 files changed, 10 insertions, 1 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 6ea874df..4f7613a1 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -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
\ No newline at end of file
+    return sd_model