diff --git a/extensions-builtin/Lora/extra_networks_lora.py b/extensions-builtin/Lora/extra_networks_lora.py index ccb249ac7..b5fea4d2e 100644 --- a/extensions-builtin/Lora/extra_networks_lora.py +++ b/extensions-builtin/Lora/extra_networks_lora.py @@ -23,5 +23,23 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): lora.load_loras(names, multipliers) + if shared.opts.lora_add_hashes_to_infotext: + lora_hashes = [] + for item in lora.loaded_loras: + shorthash = item.lora_on_disk.shorthash + if not shorthash: + continue + + alias = item.mentioned_name + if not alias: + continue + + alias = alias.replace(":", "").replace(",", "") + + lora_hashes.append(f"{alias}: {shorthash}") + + if lora_hashes: + p.extra_generation_params["Lora hashes"] = ", ".join(lora_hashes) + def deactivate(self, p): pass diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py index fa57d466e..eec147122 100644 --- a/extensions-builtin/Lora/lora.py +++ b/extensions-builtin/Lora/lora.py @@ -3,7 +3,7 @@ import re import torch from typing import Union -from modules import shared, devices, sd_models, errors, scripts, sd_hijack +from modules import shared, devices, sd_models, errors, scripts, sd_hijack, hashes metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20} @@ -76,9 +76,9 @@ class LoraOnDisk: self.name = name self.filename = filename self.metadata = {} + self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors" - _, ext = os.path.splitext(filename) - if ext.lower() == ".safetensors": + if self.is_safetensors: try: self.metadata = sd_models.read_metadata_from_safetensors(filename) except Exception as e: @@ -94,14 +94,43 @@ class LoraOnDisk: 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 self.alias = self.metadata.get('ss_output_name', self.name) + self.hash = None + self.shorthash = None + self.set_hash( + self.metadata.get('sshs_model_hash') or + hashes.sha256_from_cache(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or + '' + ) + + def set_hash(self, v): + self.hash = v + self.shorthash = self.hash[0:12] + + if self.shorthash: + available_lora_hash_lookup[self.shorthash] = self + + def read_hash(self): + if not self.hash: + self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '') + + def get_alias(self): + if shared.opts.lora_preferred_name == "Filename" or self.alias.lower() in forbidden_lora_aliases: + return self.name + else: + return self.alias + class LoraModule: - def __init__(self, name): + def __init__(self, name, lora_on_disk: LoraOnDisk): self.name = name + self.lora_on_disk = lora_on_disk self.multiplier = 1.0 self.modules = {} self.mtime = None + self.mentioned_name = None + """the text that was used to add lora to prompt - can be either name or an alias""" + class LoraUpDownModule: def __init__(self): @@ -126,11 +155,11 @@ def assign_lora_names_to_compvis_modules(sd_model): sd_model.lora_layer_mapping = lora_layer_mapping -def load_lora(name, filename): - lora = LoraModule(name) - lora.mtime = os.path.getmtime(filename) +def load_lora(name, lora_on_disk): + lora = LoraModule(name, lora_on_disk) + lora.mtime = os.path.getmtime(lora_on_disk.filename) - sd = sd_models.read_state_dict(filename) + sd = sd_models.read_state_dict(lora_on_disk.filename) # this should not be needed but is here as an emergency fix for an unknown error people are experiencing in 1.2.0 if not hasattr(shared.sd_model, 'lora_layer_mapping'): @@ -191,7 +220,7 @@ def load_lora(name, filename): raise AssertionError(f"Bad Lora layer name: {key_diffusers} - must end in lora_up.weight, lora_down.weight or alpha") if len(keys_failed_to_match) > 0: - print(f"Failed to match keys when loading Lora {filename}: {keys_failed_to_match}") + print(f"Failed to match keys when loading Lora {lora_on_disk.filename}: {keys_failed_to_match}") return lora @@ -217,14 +246,19 @@ def load_loras(names, multipliers=None): lora = already_loaded.get(name, None) lora_on_disk = loras_on_disk[i] + if lora_on_disk is not None: if lora is None or os.path.getmtime(lora_on_disk.filename) > lora.mtime: try: - lora = load_lora(name, lora_on_disk.filename) + lora = load_lora(name, lora_on_disk) except Exception as e: errors.display(e, f"loading Lora {lora_on_disk.filename}") continue + lora.mentioned_name = name + + lora_on_disk.read_hash() + if lora is None: failed_to_load_loras.append(name) print(f"Couldn't find Lora with name {name}") @@ -403,7 +437,8 @@ def list_available_loras(): available_loras.clear() available_lora_aliases.clear() forbidden_lora_aliases.clear() - forbidden_lora_aliases.update({"none": 1}) + available_lora_hash_lookup.clear() + forbidden_lora_aliases.update({"none": 1, "Addams": 1}) os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True) @@ -457,8 +492,10 @@ def infotext_pasted(infotext, params): if added: params["Prompt"] += "\n" + "".join(added) + available_loras = {} available_lora_aliases = {} +available_lora_hash_lookup = {} forbidden_lora_aliases = {} loaded_loras = [] diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py index a6b340eec..e650f469f 100644 --- a/extensions-builtin/Lora/scripts/lora_script.py +++ b/extensions-builtin/Lora/scripts/lora_script.py @@ -1,3 +1,5 @@ +import re + import torch import gradio as gr from fastapi import FastAPI @@ -54,7 +56,8 @@ script_callbacks.on_infotext_pasted(lora.infotext_pasted) shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), { "sd_lora": shared.OptionInfo("None", "Add Lora to prompt", gr.Dropdown, lambda: {"choices": ["None", *lora.available_loras]}, refresh=lora.list_available_loras), - "lora_preferred_name": shared.OptionInfo("Alias from file", "When adding to prompt, refer to lora by", gr.Radio, {"choices": ["Alias from file", "Filename"]}), + "lora_preferred_name": shared.OptionInfo("Alias from file", "When adding to prompt, refer to Lora by", gr.Radio, {"choices": ["Alias from file", "Filename"]}), + "lora_add_hashes_to_infotext": shared.OptionInfo(True, "Add Lora hashes to infotext"), })) @@ -84,3 +87,30 @@ def api_loras(_: gr.Blocks, app: FastAPI): script_callbacks.on_app_started(api_loras) +re_lora = re.compile("