From c19530f1a590d758463f84523dd4c48c34d723e6 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 14 Mar 2023 09:10:26 +0300
Subject: Add view metadata button for Lora cards.

---
 extensions-builtin/Lora/lora.py                   | 21 ++++++++++++++++++++-
 extensions-builtin/Lora/ui_extra_networks_lora.py |  1 +
 2 files changed, 21 insertions(+), 1 deletion(-)

(limited to 'extensions-builtin')

diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py
index cb8f1d36..8937b585 100644
--- a/extensions-builtin/Lora/lora.py
+++ b/extensions-builtin/Lora/lora.py
@@ -3,7 +3,9 @@ import os
 import re
 import torch
 
-from modules import shared, devices, sd_models
+from modules import shared, devices, sd_models, errors
+
+metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20}
 
 re_digits = re.compile(r"\d+")
 re_unet_down_blocks = re.compile(r"lora_unet_down_blocks_(\d+)_attentions_(\d+)_(.+)")
@@ -43,6 +45,23 @@ class LoraOnDisk:
     def __init__(self, name, filename):
         self.name = name
         self.filename = filename
+        self.metadata = {}
+
+        _, ext = os.path.splitext(filename)
+        if ext.lower() == ".safetensors":
+            try:
+                self.metadata = sd_models.read_metadata_from_safetensors(filename)
+            except Exception as e:
+                errors.display(e, f"reading lora {filename}")
+
+        if self.metadata:
+            m = {}
+            for k, v in sorted(self.metadata.items(), key=lambda x: metadata_tags_order.get(x[0], 999)):
+                m[k] = v
+
+            self.metadata = m
+
+        self.ssmd_cover_images = self.metadata.pop('ssmd_cover_images', None)  # those are cover images and they are too big to display in UI as text
 
 
 class LoraModule:
diff --git a/extensions-builtin/Lora/ui_extra_networks_lora.py b/extensions-builtin/Lora/ui_extra_networks_lora.py
index 8d32052e..68b11332 100644
--- a/extensions-builtin/Lora/ui_extra_networks_lora.py
+++ b/extensions-builtin/Lora/ui_extra_networks_lora.py
@@ -23,6 +23,7 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
                 "search_term": self.search_terms_from_path(lora_on_disk.filename),
                 "prompt": json.dumps(f"<lora:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
                 "local_preview": f"{path}.{shared.opts.samples_format}",
+                "metadata": json.dumps(lora_on_disk.metadata, indent=4) if lora_on_disk.metadata else None,
             }
 
     def allowed_directories_for_previews(self):
-- 
cgit v1.2.1