from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable, checkLazyNew, LazyVariableDict, get_tensor_device, move_to_device from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from comfy_api.latest import io, InputImpl, Types from .Stack import MrmthStack from .ParseTree import MrmthParseTree import copy import comfy def _get_video_components(video): if hasattr(video, "get_components"): return video.get_components() if isinstance(video, dict): return video raise TypeError(f"Unsupported video input type: {type(video).__name__}") def _get_video_tensor(video): components = _get_video_components(video) return components["images"] if isinstance(components, dict) else components.images def _get_video_audio(video): components = _get_video_components(video) audio = components.get("audio") if isinstance(components, dict) else components.audio if audio is None: return None return audio["waveform"] if isinstance(audio, dict) else audio.waveform def _get_video_sample_rate(video): components = _get_video_components(video) audio = components.get("audio") if isinstance(components, dict) else components.audio if isinstance(audio, dict): return audio.get("sample_rate", 44100) if audio is not None: return getattr(audio, "sample_rate", 44100) return 44100 def _build_video_output(images, waveform, sample_rate, source_video): components = _get_video_components(source_video) metadata = components.get("metadata") if isinstance(components, dict) else components.metadata frame_rate = components.get("frame_rate") if isinstance(components, dict) else components.frame_rate audio = {"waveform": waveform, "sample_rate": sample_rate} if waveform is not None else None return InputImpl.VideoFromComponents(Types.VideoComponents(images=images, audio=audio, frame_rate=frame_rate, metadata=metadata)) class VideoMathNode(io.ComfyNode): """ Enables math expressions on Audio. Inputs: I: Autogrow image inputs (I0, I1, ...) F: Autogrow float inputs (F0, F1, ...) Image: Expression """ @classmethod def define_schema(cls) -> io.Schema: return io.Schema( node_id="mrmth_ag_VideoMathNode", category="More math/math", display_name="Video math", inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Video.Input("values"), prefix="V", min=1, max=50)), io.Autogrow.Input(id="F", template=io.Autogrow.TemplatePrefix(io.Float.Input("float", default=0.0, optional=True, lazy=True, force_input=True), prefix="F", min=1, max=50)), io.MultiType.Input( io.String.Input("Expression", default="V0", multiline=False), types=[io.String,MrmthParseTree], tooltip="Expression to apply on tensor part of video", ), io.MultiType.Input( io.String.Input("Expression_pi", default="V0", multiline=False), types=[io.String,MrmthParseTree], tooltip="Expression to apply on audio part of video", ) , io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], display_name="on size mismatch", default="error", tooltip="How to handle mismatched video lengths. tile: repeat shorter inputs; error: raise error on mismatch; pad: treat missing frames as zero." ), io.Boolean.Input( id="use_compute_device", default=True, display_name="Move tensors to GPU", tooltip="Temporarily copies video tensors to the compute device for math and moves the result back afterwards.", ), MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True) ], outputs=[ io.Video.Output(), MrmthStack.Output(), ], ) @classmethod def check_lazy_status(cls, Expression,Expression_pi, V, F, length_mismatch="tile",use_compute_device=True,stack=None): d = checkLazyNew(Expression,V,F) b = checkLazyNew(Expression_pi,V,F) return d|b @classmethod def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile",use_compute_device=True,stack=None): tensor_keys = [k for k, v in V.items() if v is not None] if not tensor_keys: raise ValueError("At least one input is required.") needed_vars = checkLazyNew(Expression, V, F) | checkLazyNew(Expression_pi, V, F) needed_tensor_keys = [k for k in tensor_keys if k in needed_vars] if not needed_tensor_keys: needed_tensor_keys = tensor_keys tensors = [_get_video_tensor(V[k]) for k in needed_tensor_keys] original_device = get_tensor_device(tensors[0]) compute_device = comfy.model_management.get_torch_device() if use_compute_device else original_device working_V = move_to_device(V, compute_device) if compute_device is not None and compute_device != original_device else V # Normalize all tensors together to find the common target shape tensors = [_get_video_tensor(working_V[k]) for k in needed_tensor_keys] normalized_tensors = normalize_to_common_shape(*tensors, mode=length_mismatch) V_norm = dict(zip(needed_tensor_keys, normalized_tensors)) # Use first normalized tensor to establish the reference shape ref_tensor = normalized_tensors[0] common_shape = ref_tensor.shape stack = copy.deepcopy(stack) if stack is not None else {} sample_rates = {k + "sr": _get_video_sample_rate(V[k]) for k in needed_tensor_keys} # Setup legacy variables a, b, c, d ae = V_norm.get("V0", make_zero_like(ref_tensor)) be = V_norm.get("V1", make_zero_like(ae)) ce = V_norm.get("V2", make_zero_like(ae)) de = V_norm.get("V3", make_zero_like(ae)) # Ensure legacy variables are normalized in case they were zero-initialized ae, be, ce, de = normalize_to_common_shape(ae, be, ce, de, mode=length_mismatch) if(length_mismatch == "error"): for name, video in V.items(): if video is not None and _get_video_tensor(video).shape[0] != common_shape[0]: raise ValueError(f"Input '{name}' has shape {_get_video_tensor(video).shape[0]}, expected {common_shape[0]} to match input.") variables = LazyVariableDict({ "a": ae, "b": be, "c": ce, "d": de, "w": F.get("F0", 0.0) if F.get("F0") is not None else 0.0, "x": F.get("F1", 0.0) if F.get("F1") is not None else 0.0, "y": F.get("F2", 0.0) if F.get("F2") is not None else 0.0, "z": F.get("F3", 0.0) if F.get("F3") is not None else 0.0, "X": getIndexTensorAlongDim(ae, 2), "Y": getIndexTensorAlongDim(ae, 1), "B": getIndexTensorAlongDim(ae, 0), "batch": getIndexTensorAlongDim(ae, 0), "C": getIndexTensorAlongDim(ae, 3), "channel": getIndexTensorAlongDim(ae, 3), "W": float(ae.shape[2]), "width": float(ae.shape[2]), "H": float(ae.shape[1]), "height": float(ae.shape[1]), "T": float(ae.shape[0]), "batch_count": float(ae.shape[0]), "N": float(ae.shape[3]), "channel_count": float(ae.shape[3]), } | generate_dim_variables(ae) | V_norm | sample_rates) for v_name, v_tensor in V_norm.items(): variables[v_name] = (lambda t=v_tensor: t) variables[v_name].is_lazy_var = True def lazy_v_stack(): return get_v_variable(V_norm)[0] lazy_v_stack.is_lazy_var = True def lazy_v_count(): return float(get_v_variable(V_norm)[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) for k, val in F.items(): variables[k] = val if val is not None else 0.0 tree = None if isinstance(Expression,str): tree = parse_expr(Expression) else: tree = Expression visitor = UnifiedMathVisitor(variables, ae.shape,ae.device,state_storage=stack) result_video = visitor.visit(tree) result_video = as_tensor(result_video, ae.shape) waveforms = {} for k in needed_tensor_keys: waveform = _get_video_audio(working_V[k]) if waveform is None: raise ValueError(f"Input '{k}' does not contain audio data.") waveforms[k] = waveform # Normalize all waveforms together normalized_waveforms = normalize_to_common_shape(*waveforms.values(), mode=length_mismatch) V_norm_waveforms = dict(zip(needed_tensor_keys, normalized_waveforms)) ref_waveform = normalized_waveforms[0] common_shape = ref_waveform.shape sample_rate = _get_video_sample_rate(V[needed_tensor_keys[0]]) if(length_mismatch == "error"): for name in needed_tensor_keys: if waveforms[name].shape != common_shape: raise ValueError(f"Input '{name}' has shape ({waveforms[name].shape[0]}, {waveforms[name].shape[2]}), expected ({common_shape[0]}, {common_shape[2]}) to match input.") # Setup legacy variables a, b, c, d a_w = V_norm_waveforms.get("V0", make_zero_like(ref_waveform)) b_w = V_norm_waveforms.get("V1", make_zero_like(a_w)) c_w = V_norm_waveforms.get("V2", make_zero_like(a_w)) d_w = V_norm_waveforms.get("V3", make_zero_like(a_w)) # Ensure legacy are normalized a_w, b_w, c_w, d_w = normalize_to_common_shape(a_w, b_w, c_w, d_w, mode=length_mismatch) variables = LazyVariableDict({ "a": a_w, "b": b_w, "c": c_w, "d": d_w, "w": F.get("F0", 0.0) if F.get("F0") is not None else 0.0, "x": F.get("F1", 0.0) if F.get("F1") is not None else 0.0, "y": F.get("F2", 0.0) if F.get("F2") is not None else 0.0, "z": F.get("F3", 0.0) if F.get("F3") is not None else 0.0, "B": getIndexTensorAlongDim(a_w, 0), "C": getIndexTensorAlongDim(a_w, 1), "channel": getIndexTensorAlongDim(a_w, 1), "S": getIndexTensorAlongDim(a_w, 2), "sample": getIndexTensorAlongDim(a_w, 2), "R": sample_rate, "sample_rate": sample_rate, "batch": getIndexTensorAlongDim(a_w, 0), "T": float(a_w.shape[0]), "batch_count": float(a_w.shape[0]), } | generate_dim_variables(a_w) | sample_rates) for v_name, v_tensor in V_norm_waveforms.items(): variables[v_name] = (lambda t=v_tensor: t) variables[v_name].is_lazy_var = True def lazy_v_stack(): return get_v_variable(V_norm_waveforms)[0] lazy_v_stack.is_lazy_var = True def lazy_v_count(): return float(get_v_variable(V_norm_waveforms)[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) for k, val in F.items(): variables[k] = val if val is not None else 0.0 tree_pi = None if isinstance(Expression_pi,str): tree_pi = parse_expr(Expression_pi) else: tree_pi = Expression_pi visitor = UnifiedMathVisitor(variables, a_w.shape,a_w.device,state_storage=stack) result_waveform = visitor.visit(tree_pi) result_waveform = as_tensor(result_waveform, a_w.shape) result_waveform = move_to_device(result_waveform, original_device) result_video = move_to_device(result_video, original_device) return (_build_video_output(result_video, result_waveform, sample_rate, V[needed_tensor_keys[0]]), stack)