Fix map function + add common variables

should have added them from beginning instead of giving custom names to them. Will not be deleting that.

Author: mcDandy

README.md

@@ -22,6 +22,7 @@ You can also get the node from comfy manager under the name of More math.
 - Math: `+`, `-`, `*`, `/`, `%`, `^`, `||`
 - Boolean: `<`, `<=`, `>`, `>=`, `==`, `!=`
   (`false = 0.0`, `true = 1.0`)
+- Lists: `[v1, v2, ...]` (Vector math supported, only usefull in `conv`)
 
 ## Functions
 
@@ -69,12 +70,22 @@ You can also get the node from comfy manager under the name of More math.
 
 - `tmin(x, y)`: Element-wise minimum of x and y.
 - `tmax(x, y)`: Element-wise maximum of x and y.
-- `smin(x, y, ...)`: **Scalar** minimum. Returns the single smallest value across all input tensors/values.
-- `smax(x, y, ...)`: **Scalar** maximum. Returns the single largest value across all input tensors/values.
+- `smin(x, ...)`: **Scalar** minimum. Returns the single smallest value across all input tensors/values.
+- `smax(x, ...)`: **Scalar** maximum. Returns the single largest value across all input tensors/values.
 - `tnorm(x)`: **Tensor** Normalizes x (L2 norm along last dimension).
 - `snorm(x)`: **Scalar** L2 norm of the entire tensor.
 - `swap(tensor, dim, index1, index2)`: Swaps two slices of a tensor along a specified dimension. (Tensor only)
 
+### Advanced Tensor Operations (Tensor Only)
+
+- `map(tensor, c1, ...)`: Remaps `tensor` using source coordinates.
+  - Up to 3 coordinate mapping functions can be provided which map to the last (up to 3) dimensions of the tensor.
+  - If less than 3 functions are provided and shape of tensor > 3, the remaining dimensions are assumed to be identity functions.
+- `conv(tensor, kw, [kh], [kd], k_expr)`: Applies a convolution to `tensor`.
+  - `k_expr` can be a math expression (using `kX`, `kY`, `kZ`) or a list literal.
+
+- `permute(tensor, [dims])`: Rearranges the dimensions of the tensor. (e.g., `permute(a, [2, 3, 0, 1])`)
+
 ### FFT (Tensor Only)
 
 - `fft(x)`: Fast Fourier Transform (Time to Frequency).
@@ -84,9 +95,13 @@ You can also get the node from comfy manager under the name of More math.
 ### Utility
 
 - `print(x)`: Prints the value of x to the console and returns x.
+- `print_shape(x)` or `pshp`: Prints the shape of x to the console and returns x.
 
 ## Variables
 
+- **Common variables (except FLOAT, MODEL, VAE and CLIP)**:
+  - `D{N}` - position in n-th dimension of tensor (for example D0, D1, D2, ...)
+  - `S{N}` - size of n-th dimension of tensor (for example S0, S1, S2, ...)
 - **common inputs** (matches node input type):
   - `a`, `b`, `c`, `d`
 - **Extra floats**:
@@ -103,8 +118,13 @@ You can also get the node from comfy manager under the name of More math.
   - `W` or `width` - width of image. y/width = 1
   - `H` or `height`- height of image. x/height = 1
   - `B` or 'batch' - position in batch
-  - 'T' or 'batch_count` - number of batches
+  - `T` or `batch_count` - number of batches
   - `N` or `channel_count` - count of channels
+- **IMAGE KERNEL**:
+  - `kX`, `kY` - position in kernel. Centered at 0.0.
+  - `kW`, `kernel_width` - width of kernel.
+  - `kH`, `kernel_height` - height of kernel.
+  - `kD`, `kernel_depth` - depth of kernel.
 
 - **AUDIO**:
   - `B` or 'batch' - position in batch
@@ -124,6 +144,6 @@ You can also get the node from comfy manager under the name of More math.
   - no additional variables
   - `F` or `frequency_count` – frequency count (freq domain, iFFT only)
   - `K` or `frequency` – isotropic frequency (Euclidean norm of indices, iFFT only)
-  - `Kx`, `Ky`, `K_dimN` - frequency index for specific dimension
-  - `Fx`, `Fy`, `F_dimN` - frequency count for specific dimension
+  - `Kx`, `Ky`, `Kz`, `Kw`,`Kv`, `Ku`, `K_dimN` - frequency index for specific dimension
+  - `Fx`, `Fy`, `Fz`, `Fw`,`Fv`,`Fu`, `F_dimN` - frequency count for specific dimension
 - Constants: `e`, `pi`

more_math/AudioMathNode.py

@@ -1,5 +1,5 @@
 import torch
-from .helper_functions import getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like
+from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like
 
 from comfy_api.latest import io
 
@@ -9,12 +9,12 @@
 class AudioMathNode(MathNodeBase):
     """
     Enables math expressions on Audio tensors.
-    
+
     Inputs:
         a, b, c, d: Audio inputs (b, c, d default to zero if not provided)
         w, x, y, z: Float variables for expressions
         AudioExpr: Expression to apply on audio tensors
-    
+
     Outputs:
         AUDIO: Result of applying expression to input audio
     """
@@ -42,31 +42,32 @@ def define_schema(cls) -> io.Schema:
         )
 
     @classmethod
-    def execute(cls, a, AudioExpr, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
-        waveform = a['waveform']
+    def execute(cls, AudioExpr, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
+        av = a['waveform']
         sample_rate = a['sample_rate']
 
         a, b, c, d = cls.prepare_inputs(a, b, c, d)
 
         bv, cv, dv = b['waveform'], c['waveform'], d['waveform']
 
         variables = {
-            'a': waveform, 'b': bv, 'c': cv, 'd': dv,
+            'a': av, 'b': bv, 'c': cv, 'd': dv,
             'w': w, 'x': x, 'y': y, 'z': z,
-            'B': getIndexTensorAlongDim(waveform, 0),
-            'C': getIndexTensorAlongDim(waveform, 1),
-            'S': getIndexTensorAlongDim(waveform, 2),
-            'R': torch.full_like(waveform, sample_rate, dtype=torch.float32),
-            'T': torch.full_like(waveform, waveform.shape[2], dtype=torch.float32),
-            'N': waveform.shape[1],
-            'batch': getIndexTensorAlongDim(waveform, 0),
-            'channel': getIndexTensorAlongDim(waveform, 1),
-            'sample': getIndexTensorAlongDim(waveform, 2),
-            'sample_rate': torch.full_like(waveform, sample_rate, dtype=torch.float32),
-            'sample_count': torch.full_like(waveform, waveform.shape[2], dtype=torch.float32),
-            'channel_count': waveform.shape[1],
-        }
+            'B': getIndexTensorAlongDim(av, 0),
+            'C': getIndexTensorAlongDim(av, 1),
+            'S': getIndexTensorAlongDim(av, 2),
+            'R': torch.full_like(av, sample_rate, dtype=torch.float32),
+            'T': torch.full_like(av, av.shape[2], dtype=torch.float32),
+            'N': av.shape[1],
+            'batch': getIndexTensorAlongDim(av, 0),
+            'channel': getIndexTensorAlongDim(av, 1),
+            'sample': getIndexTensorAlongDim(av, 2),
+            'sample_rate': torch.full_like(av, sample_rate, dtype=torch.float32),
+            'sample_count': torch.full_like(av, av.shape[2], dtype=torch.float32),
+            'channel_count': av.shape[1],
+        } | generate_dim_variables(av)
+
 
-        result_tensor = eval_tensor_expr(AudioExpr, variables, waveform.shape)
+        result_tensor = eval_tensor_expr(AudioExpr, variables, av.shape)
 
         return ({'waveform': result_tensor, 'sample_rate': sample_rate},)

more_math/ConditioningMathNode.py

@@ -1,7 +1,7 @@
 from inspect import cleandoc
 import torch
 
-from .helper_functions import comonLazy, eval_tensor_expr, make_zero_like
+from .helper_functions import comonLazy, eval_tensor_expr, generate_dim_variables, make_zero_like
 from comfy_api.latest import io
 
 
@@ -11,17 +11,17 @@
 class ConditioningMathNode(MathNodeBase):
     """
     Enables math operations on conditionings.
-    
+
     Inputs:
         a, b, c, d: Conditioning inputs (b, c, d default to zero if not provided)
         w, x, y, z: Float variables for expressions
         Tensor: Expression for the tensor part (describes image composition)
         pooled_output: Expression for the pooled output (condensed representation)
-    
+
     Outputs:
         CONDITIONING: Result of applying expressions to input conditionings
     """
-    
+
     @classmethod
     def define_schema(cls) -> io.Schema:
         return io.Schema(
@@ -60,9 +60,9 @@ def execute(cls, Tensor, pooled_output, a, b=None, c=None, d=None, w=0.0, x=0.0,
 
         # Extract tensors
         ta, tb, tc, td = a[0][0], b[0][0], c[0][0], d[0][0]
-        
+
         # Evaluate tensor expression
-        variables = {'a': ta, 'b': tb, 'c': tc, 'd': td, 'w': w, 'x': x, 'y': y, 'z': z}
+        variables = {'a': ta, 'b': tb, 'c': tc, 'd': td, 'w': w, 'x': x, 'y': y, 'z': z} | generate_dim_variables(ta)
         result_tensor = eval_tensor_expr(Tensor, variables, ta.shape)
 
         # Evaluate pooled_output expression if available
@@ -71,7 +71,7 @@ def execute(cls, Tensor, pooled_output, a, b=None, c=None, d=None, w=0.0, x=0.0,
             pb = b[0][1].get("pooled_output")
             pc = c[0][1].get("pooled_output")
             pd = d[0][1].get("pooled_output")
-            variables = {'a': pa, 'b': pb, 'c': pc, 'd': pd, 'w': w, 'x': x, 'y': y, 'z': z}
+            variables = {'a': pa, 'b': pb, 'c': pc, 'd': pd, 'w': w, 'x': x, 'y': y, 'z': z} | generate_dim_variables(pa)
             result_pooled = eval_tensor_expr(pooled_output, variables, pa.shape)
         else:
             result_pooled = None

more_math/ImageMathNode.py

@@ -1,5 +1,5 @@
 import torch
-from .helper_functions import getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like
+from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like
 
 from comfy_api.latest import io
 
@@ -9,12 +9,12 @@
 class ImageMathNode(MathNodeBase):
     """
     Enables math expressions on Images.
-    
+
     Inputs:
         a, b, c, d: Image inputs (b, c, d default to zero if not provided)
         w, x, y, z: Float variables for expressions
         Image: Expression to apply on input images
-    
+
     Outputs:
         IMAGE: Result of applying expression to input images
     """
@@ -45,11 +45,6 @@ def define_schema(cls) -> io.Schema:
     def execute(cls, Image, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
         a, b, c, d = cls.prepare_inputs(a, b, c, d)
 
-        # Permute to B, C, H, W for processing
-        a = a.permute(0, 3, 1, 2)
-        b = b.permute(0, 3, 1, 2)
-        c = c.permute(0, 3, 1, 2)
-        d = d.permute(0, 3, 1, 2)
 
         variables = {
             'a': a, 'b': b, 'c': c, 'd': d,
@@ -62,10 +57,8 @@ def execute(cls, Image, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
             'H': a.shape[2], 'height': a.shape[2],
             'T': a.shape[0], 'batch_count': a.shape[0],
             'N': a.shape[1], 'channel_count': a.shape[1],
-        }
+        } | generate_dim_variables(a)
 
         result = eval_tensor_expr(Image, variables, a.shape)
 
-        # Permute back to B, H, W, C
-        result = result.permute(0, 2, 3, 1)
         return (result,)

more_math/LatentMathNode.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from .helper_functions import getIndexTensorAlongDim, comonLazy, parse_expr, eval_tensor_expr_with_tree, make_zero_like
+from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, parse_expr, eval_tensor_expr_with_tree, make_zero_like
 
 from .MathNodeBase import MathNodeBase
 
@@ -114,7 +114,7 @@ def eval_single_tensor(a_t, b_t, c_t, d_t):
                 '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:

more_math/NoiseMathNode.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from .helper_functions import getIndexTensorAlongDim, comonLazy, parse_expr, eval_tensor_expr_with_tree, make_zero_like
+from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, parse_expr, eval_tensor_expr_with_tree, make_zero_like
 
 from comfy_api.latest import io
 
@@ -183,7 +183,7 @@ def merge_to_tensor(val, ref):
             'I': merged_samples, 'T': frame_count, 'N': merged_samples.shape[channel_dim],
             'batch': B, 'width': merged_samples.shape[width_dim], 'height': merged_samples.shape[height_dim], 'channel': C,
             'batch_count': merged_samples.shape[0], 'channel_count': merged_samples.shape[1], 'input_latent': merged_samples,
-        }
+        } | generate_dim_variables(merged_samples)
         if time_dim is not None:
             F = getIndexTensorAlongDim(merged_samples, time_dim)
             variables.update({'frame': F, 'frame_count': frame_count})

more_math/Parser/MathExpr.g4

@@ -89,6 +89,8 @@ func1
     | SOFTPLUS '(' expr ')'  # SoftplusFunc
     | GELU  '(' expr ')'     # GeluFunc
     | SIGN  '(' expr ')'     # SignFunc
+    | PRINT_SHAPE_L '(' expr ')' # PrintShapeFunc
+    | PRINT_SHAPE   '(' expr ')' # PrintShapeFunc
     ;
 
 // Two-argument functions
@@ -156,6 +158,8 @@ SFFT : 'fft';
 SIFFT : 'ifft';
 ANGL : 'angle';
 PRNT : 'print';
+PRINT_SHAPE_L : 'print_shape';
+PRINT_SHAPE : 'pshp';
 LERP : 'lerp';
 STEP : 'step';
 SMOOTHSTEP : 'smoothstep';