229 lines
8.3 KiB
Python
229 lines
8.3 KiB
Python
from .helper_functions import generate_dim_variables, parse_expr, as_tensor, get_v_variable, get_f_variable, get_tensor_device, move_to_device, LazyVariableDict, checkLazyNew
|
|
from .Parser.UnifiedMathVisitor import MathDict, UnifiedMathVisitor
|
|
import comfy.model_management
|
|
import torch
|
|
from collections.abc import Mapping
|
|
|
|
|
|
def calculate_lora_dict_autogrow(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False):
|
|
|
|
"""
|
|
Calculate a new LORA_MODEL dict[str, Tensor] from math expressions applied
|
|
over the UNION of keys from all input LoRA dicts, using Autogrow inputs.
|
|
|
|
Unlike calculate_patches_autogrow (Model/VAE/CLIP), there is no
|
|
ModelPatcher/backup involved here - a LORA_MODEL is just raw tensors
|
|
(lora_up.weight, lora_down.weight, alpha, diff, ...). So V{n} is simply
|
|
that tensor as-is (zero-filled when a given key is missing from that
|
|
input), and there is no V{n}_d - these tensors already *are* deltas.
|
|
"""
|
|
if V is None: V = {}
|
|
if F is None: F = {}
|
|
if mapping is None: mapping = {}
|
|
|
|
needed_vars = checkLazyNew(Expr, V, F)
|
|
needed_v_names = [v_name for v_name in V.keys() if v_name in needed_vars]
|
|
if not needed_v_names:
|
|
needed_v_names = [v_name for v_name in V.keys() if V.get(v_name) is not None]
|
|
|
|
if not any(V.get(v_name) is not None for v_name in needed_v_names):
|
|
return {}
|
|
|
|
# Collect all unique keys from all inputs, preserving order from first input
|
|
all_keys_list = []
|
|
seen_keys = set()
|
|
for d in V.values():
|
|
if d is None:
|
|
continue
|
|
for key in d.keys():
|
|
if key not in seen_keys:
|
|
seen_keys.add(key)
|
|
all_keys_list.append(key)
|
|
|
|
tree = parse_expr(Expr) if isinstance(Expr, str) else Expr
|
|
|
|
original_device = None
|
|
compute_device = comfy.model_management.get_torch_device() if use_compute_device else None
|
|
|
|
const_f_vars = {k: val if val is not None else 0.0 for k, val in F.items()}
|
|
const_f_alias_vars = {
|
|
alias: (F[target] if F[target] is not None else 0.0)
|
|
for alias, target in mapping.items()
|
|
if target in F
|
|
}
|
|
mapping_items = tuple(mapping.items())
|
|
|
|
input_dicts = {}
|
|
for v_name in needed_v_names:
|
|
d = V.get(v_name)
|
|
if d is None:
|
|
continue
|
|
if original_device is None:
|
|
original_device = get_tensor_device(d)
|
|
if compute_device is not None and original_device is not None and compute_device != original_device:
|
|
d = move_to_device(d, compute_device)
|
|
input_dicts[v_name] = d
|
|
|
|
layer_count = len(all_keys_list)
|
|
|
|
def make_lazy(value_fn):
|
|
def lazy_eval():
|
|
return value_fn()
|
|
lazy_eval.is_lazy_var = True
|
|
return lazy_eval
|
|
|
|
result = {}
|
|
|
|
for layer_idx, key in enumerate(all_keys_list):
|
|
variables = LazyVariableDict(const_f_vars)
|
|
variables.update(const_f_alias_vars)
|
|
|
|
variables["L"] = float(layer_idx)
|
|
variables["layer"] = float(layer_idx)
|
|
variables["LC"] = float(layer_count)
|
|
variables["layer_count"] = float(layer_count)
|
|
variables["K"] = key
|
|
variables["key"] = key
|
|
|
|
ref_tensor = None
|
|
present = {}
|
|
for v_name in V.keys():
|
|
d = input_dicts.get(v_name)
|
|
if d is None:
|
|
continue
|
|
t = d.get(key, None)
|
|
if t is not None:
|
|
present[v_name] = t
|
|
if ref_tensor is None:
|
|
ref_tensor = t
|
|
|
|
if ref_tensor is None:
|
|
continue
|
|
|
|
for v_name in V.keys():
|
|
if v_name in present:
|
|
t = present[v_name]
|
|
variables[v_name] = make_lazy(lambda t=t: t)
|
|
else:
|
|
variables[v_name] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
|
|
|
|
# Populate aliases for V (a, b, c, d)
|
|
for alias, target in mapping_items:
|
|
if target in variables:
|
|
variables[alias] = variables[target]
|
|
elif target.startswith("V") and target[1:].isdigit():
|
|
variables[alias] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
|
|
|
|
v_cache = [None]
|
|
|
|
def lazy_v_stack():
|
|
if v_cache[0] is None:
|
|
v_cache[0] = get_v_variable(variables)
|
|
return v_cache[0][0]
|
|
|
|
lazy_v_stack.is_lazy_var = True
|
|
|
|
def lazy_v_count():
|
|
if v_cache[0] is None:
|
|
v_cache[0] = get_v_variable(variables)
|
|
return float(v_cache[0][1])
|
|
|
|
lazy_v_count.is_lazy_var = True
|
|
|
|
variables["V"] = lazy_v_stack
|
|
variables["Vcnt"] = lazy_v_count
|
|
variables["V_count"] = lazy_v_count
|
|
|
|
f_stacked, f_cnt = get_f_variable(F)
|
|
if f_stacked is not None:
|
|
variables["F"] = f_stacked
|
|
variables["Fcnt"] = float(f_cnt)
|
|
variables["F_count"] = float(f_cnt)
|
|
|
|
variables = variables | generate_dim_variables(ref_tensor)
|
|
|
|
visitor = UnifiedMathVisitor(variables, ref_tensor.shape, state_storage=stack)
|
|
res = visitor.visit(tree)
|
|
res = as_tensor(res, ref_tensor.shape)
|
|
|
|
if compute_device is not None and original_device is not None and original_device != compute_device:
|
|
res = res.to(device=original_device)
|
|
|
|
result[key] = res
|
|
|
|
if pbar is not None:
|
|
pbar.update(1)
|
|
|
|
return result
|
|
|
|
|
|
def calculate_lora_dict_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False):
|
|
if V is None:
|
|
V = {}
|
|
if F is None:
|
|
F = {}
|
|
if mapping is None:
|
|
mapping = {}
|
|
|
|
needed_vars = checkLazyNew(Expr, V, F)
|
|
needed_v_names = [name for name in V if name in needed_vars]
|
|
if "V0" in V and "V0" not in needed_v_names:
|
|
needed_v_names.append("V0")
|
|
if not needed_v_names:
|
|
needed_v_names = [name for name, value in V.items() if value is not None]
|
|
|
|
input_dicts = {name: V[name] for name in needed_v_names if V.get(name) is not None}
|
|
if not input_dicts:
|
|
return {}
|
|
|
|
all_keys = []
|
|
seen_keys = set()
|
|
for value in input_dicts.values():
|
|
for key in value:
|
|
if key not in seen_keys:
|
|
seen_keys.add(key)
|
|
all_keys.append(key)
|
|
|
|
compute_device = comfy.model_management.get_torch_device() if use_compute_device else None
|
|
original_device = get_tensor_device(next(iter(input_dicts.values())))
|
|
if compute_device is not None and original_device is not None and compute_device != original_device:
|
|
input_dicts = {name: move_to_device(value, compute_device) for name, value in input_dicts.items()}
|
|
|
|
references = {}
|
|
for key in all_keys:
|
|
for value in input_dicts.values():
|
|
tensor = value.get(key)
|
|
if tensor is not None:
|
|
references[key] = tensor
|
|
break
|
|
|
|
def weight_dict(name):
|
|
value = input_dicts.get(name, {})
|
|
return MathDict({key: value.get(key, torch.zeros_like(references[key])) for key in all_keys})
|
|
|
|
variables = LazyVariableDict({name: value if value is not None else 0.0 for name, value in F.items()})
|
|
for alias, target in mapping.items():
|
|
if target in F:
|
|
variables[alias] = F[target] if F[target] is not None else 0.0
|
|
for name in needed_v_names:
|
|
variables[name] = weight_dict(name)
|
|
for alias, target in mapping.items():
|
|
if target in variables:
|
|
variables[alias] = variables[target]
|
|
|
|
visitor_device = compute_device if compute_device is not None else next(iter(references.values())).device
|
|
result = UnifiedMathVisitor(variables, device=visitor_device, state_storage=stack).visit(parse_expr(Expr) if isinstance(Expr, str) else Expr)
|
|
if not isinstance(result, Mapping):
|
|
if isinstance(result, (float, int)):
|
|
result = MathDict({key: torch.full_like(reference, result) for key, reference in references.items()})
|
|
else:
|
|
raise ValueError("Dictionary mode expressions must return a dictionary")
|
|
if tuple(result) != tuple(all_keys):
|
|
raise ValueError("Dictionary mode expressions must preserve all LoRA keys")
|
|
|
|
if compute_device is not None and original_device is not None and original_device != compute_device:
|
|
result = {key: value.to(device=original_device) for key, value in result.items()}
|
|
if pbar is not None:
|
|
pbar.update(len(result))
|
|
return dict(result)
|