diff --git a/py/capture.py b/py/capture.py index 16306be..1c38335 100644 --- a/py/capture.py +++ b/py/capture.py @@ -1,4 +1,5 @@ import json +import os from . import hook from .defs.captures import CAPTURE_FIELD_LIST @@ -109,8 +110,16 @@ class Capture: @classmethod def get_hashes_for_civitai(cls, inputs): resource_hashes = {} - x = inputs.get(MetaField.MODEL_HASH, []) - if len(x) > 0: - resource_hashes["model"] = x[0][1] + model_hashes = inputs.get(MetaField.MODEL_HASH, []) + if len(model_hashes) > 0: + resource_hashes["model"] = model_hashes[0][1] + + lora_model_names = inputs.get(MetaField.LORA_MODEL_NAME, []) + lora_model_hashes = inputs.get(MetaField.LORA_MODEL_HASH, []) + for lora_model_name, lora_model_hash in zip( + lora_model_names, lora_model_hashes + ): + lora_model_name = os.path.splitext(os.path.basename(lora_model_name[1]))[0] + resource_hashes[f"lora:{lora_model_name}"] = lora_model_hash[1] return resource_hashes diff --git a/py/defs/captures.py b/py/defs/captures.py index 0cd0e0c..88f1422 100644 --- a/py/defs/captures.py +++ b/py/defs/captures.py @@ -2,6 +2,7 @@ from .meta import MetaField from .validators import is_positive_prompt, is_negative_prompt from .formatters import ( calc_model_hash, + calc_lora_hash, convert_skip_clip, get_scaled_width, get_scaled_height, @@ -58,4 +59,21 @@ CAPTURE_FIELD_LIST = { "format": get_scaled_height, }, }, + "LoraLoader": { + MetaField.LORA_MODEL_NAME: {"field_name": "lora_name"}, + MetaField.LORA_MODEL_HASH: { + "field_name": "lora_name", + "format": calc_lora_hash, + }, + MetaField.LORA_STRENGTH_MODEL: {"field_name": "strength_model"}, + MetaField.LORA_STRENGTH_CLIP: {"field_name": "strength_clip"}, + }, + "LoraLoaderModelOnly": { + MetaField.LORA_MODEL_NAME: {"field_name": "lora_name"}, + MetaField.LORA_MODEL_HASH: { + "field_name": "lora_name", + "format": calc_lora_hash, + }, + MetaField.LORA_STRENGTH_MODEL: {"field_name": "strength_model"}, + }, } diff --git a/py/defs/formatters.py b/py/defs/formatters.py index 4a183b4..7f7997a 100644 --- a/py/defs/formatters.py +++ b/py/defs/formatters.py @@ -1,24 +1,18 @@ -import hashlib - import folder_paths +from ..utils.hash import calc_hash + cache_model_hash = {} def calc_model_hash(model_name, input_data): - if model_name in cache_model_hash: - return cache_model_hash[model_name] - filename = folder_paths.get_full_path("checkpoints", model_name) - sha256_hash = hashlib.sha256() + return calc_hash(filename) - with open(filename, "rb") as f: - for byte_block in iter(lambda: f.read(4096), b""): - sha256_hash.update(byte_block) - model_hash = sha256_hash.hexdigest()[:10] - cache_model_hash[model_name] = model_hash - return model_hash +def calc_lora_hash(model_name, input_data): + filename = folder_paths.get_full_path("loras", model_name) + return calc_hash(filename) def convert_skip_clip(stop_at_clip_layer, input_data): diff --git a/py/defs/meta.py b/py/defs/meta.py index b18e2c9..b6dc108 100644 --- a/py/defs/meta.py +++ b/py/defs/meta.py @@ -14,3 +14,8 @@ class MetaField(IntEnum): SCHEDULER = 24 IMAGE_WIDTH = 30 IMAGE_HEIGHT = 31 + EMBEDDING_NAME = 40 + LORA_MODEL_NAME = 50 + LORA_MODEL_HASH = 51 + LORA_STRENGTH_MODEL = 52 + LORA_STRENGTH_CLIP = 53 diff --git a/py/utils/hash.py b/py/utils/hash.py new file mode 100644 index 0000000..6cae680 --- /dev/null +++ b/py/utils/hash.py @@ -0,0 +1,18 @@ +import hashlib + + +cache_model_hash = {} + + +def calc_hash(filename): + if filename in cache_model_hash: + return cache_model_hash[filename] + sha256_hash = hashlib.sha256() + + with open(filename, "rb") as f: + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + model_hash = sha256_hash.hexdigest()[:10] + + cache_model_hash[filename] = model_hash + return model_hash