Support LoRA

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

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"},
+    },
 }

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):

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

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