diff --git a/py/defs/ext/rgthree.py b/py/defs/ext/rgthree.py index bdac763..aad025f 100644 --- a/py/defs/ext/rgthree.py +++ b/py/defs/ext/rgthree.py @@ -26,6 +26,29 @@ def get_lora_data(input_data, attribute): ] +def get_lora_model_name_stack(node_id, obj, prompt, extra_data, outputs, input_data): + return get_lora_data_stack(input_data, "lora") + + +def get_lora_model_hash_stack(node_id, obj, prompt, extra_data, outputs, input_data): + return [ + calc_lora_hash(model_name, input_data) + for model_name in get_lora_data_stack(input_data, "lora") + ] + + +def get_lora_strength_stack(node_id, obj, prompt, extra_data, outputs, input_data): + return get_lora_data_stack(input_data, "strength") + + +def get_lora_data_stack(input_data, attribute): + return [ + v[0] + for k, v in input_data.items() + if k.startswith(attribute + "_") and v[0] != "None" + ] + + CAPTURE_FIELD_LIST = { "Power Lora Loader (rgthree)": { MetaField.LORA_MODEL_NAME: {"selector": get_lora_model_name}, @@ -33,4 +56,10 @@ CAPTURE_FIELD_LIST = { MetaField.LORA_STRENGTH_MODEL: {"selector": get_lora_strength}, MetaField.LORA_STRENGTH_CLIP: {"selector": get_lora_strength}, }, + "Lora Loader Stack (rgthree)": { + MetaField.LORA_MODEL_NAME: {"selector": get_lora_model_name_stack}, + MetaField.LORA_MODEL_HASH: {"selector": get_lora_model_hash_stack}, + MetaField.LORA_STRENGTH_MODEL: {"selector": get_lora_strength_stack}, + MetaField.LORA_STRENGTH_CLIP: {"selector": get_lora_strength_stack}, + }, }