cleanup

more_math/LatentMathNode.py

@@ -1,20 +1,10 @@
 from inspect import cleandoc
-
+import comfy.nested_tensor
 from comfy_api.latest import io
-
 import torch
-
 from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, parse_expr, eval_tensor_expr_with_tree, make_zero_like
-
 from .MathNodeBase import MathNodeBase
 
-# try to import NestedTensor type if available
-try:
-    import comfy.nested_tensor as _nested_tensor_module
-    _NESTED_TENSOR_AVAILABLE = True
-except Exception:
-    _nested_tensor_module = None
-    _NESTED_TENSOR_AVAILABLE = False
 
 
 class LatentMathNode(MathNodeBase):
@@ -75,18 +65,14 @@ class LatentMathNode(MathNodeBase):
         # parse expression once
         tree = parse_expr(Latent)
 
-        # Helper to evaluate for a single tensor
         def eval_single_tensor(a_t, b_t, c_t, d_t):
-            # support tensors with >=4 dims (e.g. 4D latents [B,C,H,W] or 5D [B,T,C,H,W])
             ndim = a_t.ndim
-            # use negative indexing so that channel/height/width mapping works for 4D and 5D
             batch_dim = 0
             channel_dim = -3
             height_dim = -2
             width_dim = -1
             time_dim = None
             if ndim >= 5:
-                # time/frame dim is the one before channels when present
                 time_dim = -4
 
             B = getIndexTensorAlongDim(a_t, batch_dim)
@@ -94,7 +80,6 @@ class LatentMathNode(MathNodeBase):
             H = getIndexTensorAlongDim(a_t, height_dim)
             W = getIndexTensorAlongDim(a_t, width_dim)
 
-            # fill scalar/value tensors
             width_val = a_t.shape[width_dim]
             height_val = a_t.shape[height_dim]
             channel_count = a_t.shape[channel_dim]
@@ -102,37 +87,30 @@ class LatentMathNode(MathNodeBase):
             frame_count = a_t.shape[time_dim] if time_dim is not None else a_t.shape[batch_dim]
 
             variables = {
-                # core tensors and floats
                 'a': a_t, 'b': b_t, 'c': c_t, 'd': d_t,
                 'w': w, 'x': x, 'y': y, 'z': z,
 
                  'X': W, 'Y': H,
                 'B': B, 'batch': B,
                 'C': C, 'channel': C,
-                # scalar dims and counts
                 'W': width_val, 'width': width_val,
                 'H': height_val, 'height': height_val,
                 'T': frame_count, 'batch_count': batch_count,
                 'N': channel_count, 'channel_count': channel_count,
             } | generate_dim_variables(a_t)
 
-            # expose time/frame if present
             if time_dim is not None:
                 F = getIndexTensorAlongDim(a_t, time_dim)
                 variables.update({'frame_idx': F, 'frame': F, 'frame_count': frame_count})
 
             return eval_tensor_expr_with_tree(tree, variables, a_t.shape)
 
-        # If input is a NestedTensor (from comfy), evaluate per-subtensor and return NestedTensor result
         if hasattr(a_in, 'is_nested') and getattr(a_in, 'is_nested'):
-            # get underlying lists
             a_list = a_in.unbind()
             sizes = [t.shape[0] for t in a_list]
-            # merge all a subtensors along batch (dim=0)
             merged_a = torch.cat(a_list, dim=0)
 
             def merge_to_tensor(val, ref):
-                # ref is merged_a
                 if val is None:
                     return make_zero_like(ref)
                 if hasattr(val, 'is_nested') and getattr(val, 'is_nested'):
@@ -141,39 +119,27 @@ class LatentMathNode(MathNodeBase):
                 if isinstance(val, (list, tuple)):
                     return torch.cat(list(val), dim=0)
                 if torch.is_tensor(val):
-                    # if val already matches merged shape
                     if val.shape == ref.shape:
                         return val
-                    # if val is per-subtensor with same per-subtensor batch, replicate
                     try:
                         if val.shape[0] in sizes and val.shape[1:] == a_list[0].shape[1:]:
-                            # broadcast by concatenating copies
                             return torch.cat([val for _ in a_list], dim=0)
                     except Exception:
                         pass
-                    # if val has batch equal to combined, return as is
                     if val.shape[0] == sum(sizes):
                         return val
-                # fallback
                 return make_zero_like(ref)
 
             merged_b = merge_to_tensor(b_in, merged_a)
             merged_c = merge_to_tensor(c_in, merged_a)
             merged_d = merge_to_tensor(d_in, merged_a)
-
-            # evaluate once on merged tensors
             merged_result = eval_single_tensor(merged_a, merged_b, merged_c, merged_d)
 
-            # split back into list
             split_results = list(merged_result.split(sizes, dim=0))
-            if _NESTED_TENSOR_AVAILABLE and _nested_tensor_module is not None:
-                out_samples = _nested_tensor_module.NestedTensor(split_results)
-            else:
-                out_samples = split_results
+            out_samples = comfy.nested_tensor.nested_tensor.NestedTensor(split_results)
+
             return ({"samples": out_samples},)
 
-        # Non-nested (single tensor) path
-        # ensure b/c/d are set appropriately (zeros_like if None)
         def to_tensor(val, ref):
             if val is None:
                 return make_zero_like(ref)