Fix map function + add common variables

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

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 @@ from .MathNodeBase import MathNodeBase
 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 @@ class ImageMathNode(MathNodeBase):
     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 @@ class ImageMathNode(MathNodeBase):
             '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,)