diff --git a/py/capture.py b/py/capture.py index 795420b..fb01d30 100644 --- a/py/capture.py +++ b/py/capture.py @@ -41,6 +41,18 @@ class Capture: inputs[meta].append((node_id, value)) continue + selector = field_data.get("selector") + if selector is not None: + v = selector( + node_id, obj, prompt, extra_data, outputs, input_data + ) + if isinstance(v, list): + for x in v: + inputs[meta].append((node_id, x)) + elif v is not None: + inputs[meta].append((node_id, v)) + continue + field_name = field_data["field_name"] value = input_data.get(field_name) if value is not None: diff --git a/py/defs/ext/efficiency_nodes.py b/py/defs/ext/efficiency_nodes.py index c91f30e..bf15b86 100644 --- a/py/defs/ext/efficiency_nodes.py +++ b/py/defs/ext/efficiency_nodes.py @@ -1,6 +1,41 @@ # https://github.com/jags111/efficiency-nodes-comfyui from ..meta import MetaField -from ..formatters import calc_model_hash, convert_skip_clip +from ..formatters import calc_model_hash, calc_lora_hash, convert_skip_clip + + +def get_lora_model_name_stack(node_id, obj, prompt, extra_data, outputs, input_data): + return get_lora_data_stack(input_data, "lora_name") + + +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_name") + ] + + +def get_lora_strength_model_stack( + node_id, obj, prompt, extra_data, outputs, input_data +): + if input_data["input_mode"][0] == "advanced": + return get_lora_data_stack(input_data, "model_str") + return get_lora_data_stack(input_data, "lora_wt") + + +def get_lora_strength_clip_stack(node_id, obj, prompt, extra_data, outputs, input_data): + if input_data["input_mode"][0] == "advanced": + return get_lora_data_stack(input_data, "clip_str") + return get_lora_data_stack(input_data, "lora_wt") + + +def get_lora_data_stack(input_data, attribute): + lora_count = input_data["lora_count"][0] + return [ + v[0] + for k, v in input_data.items() + if k.startswith(attribute) and v[0] != "None" + ][:lora_count] + SAMPLERS = { "KSampler (Efficient)": { @@ -63,4 +98,10 @@ CAPTURE_FIELD_LIST = { MetaField.SAMPLER_NAME: {"field_name": "sampler_name"}, MetaField.SCHEDULER: {"field_name": "scheduler"}, }, + "LoRA Stacker": { + 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_model_stack}, + MetaField.LORA_STRENGTH_CLIP: {"selector": get_lora_strength_clip_stack}, + }, } diff --git a/py/defs/ext/rgthree.py b/py/defs/ext/rgthree.py new file mode 100644 index 0000000..aad025f --- /dev/null +++ b/py/defs/ext/rgthree.py @@ -0,0 +1,65 @@ +# https://github.com/rgthree/rgthree-comfy +from ..meta import MetaField +from ..formatters import calc_lora_hash + + +def get_lora_model_name(node_id, obj, prompt, extra_data, outputs, input_data): + return get_lora_data(input_data, "lora") + + +def get_lora_model_hash(node_id, obj, prompt, extra_data, outputs, input_data): + return [ + calc_lora_hash(model_name, input_data) + for model_name in get_lora_data(input_data, "lora") + ] + + +def get_lora_strength(node_id, obj, prompt, extra_data, outputs, input_data): + return get_lora_data(input_data, "strength") + + +def get_lora_data(input_data, attribute): + return [ + v[0][attribute] + for k, v in input_data.items() + if k.startswith("lora_") and v[0]["on"] + ] + + +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}, + MetaField.LORA_MODEL_HASH: {"selector": get_lora_model_hash}, + 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}, + }, +}