From bb50c285fd08950f79b24ea738c48ce05b2d6019 Mon Sep 17 00:00:00 2001
From: mcDandy <danda.martinek@gmail.com>
Date: Mon, 4 Aug 2025 17:49:00 +0200
Subject: [PATCH] added ability to get position on image math node

B = Batch position
C = Color position
W = width
H = height
X = X position of pixel
Y = Y position of pixel
---
 src/more_math/ImageMathNode.py    | 20 ++++++++++++++++----
 src/more_math/helper_functions.py | 15 +++++++++++++++
 src/more_math/nodes.py            |  1 +
 3 files changed, 32 insertions(+), 4 deletions(-)
 create mode 100644 src/more_math/helper_functions.py

diff --git a/src/more_math/ImageMathNode.py b/src/more_math/ImageMathNode.py
index 83d5e26..0dfba31 100644
--- a/src/more_math/ImageMathNode.py
+++ b/src/more_math/ImageMathNode.py
@@ -1,9 +1,9 @@
 from inspect import cleandoc
-from math import e
 
 from antlr4 import CommonTokenStream
 from antlr4.atn.LexerActionExecutor import InputStream
 import torch
+from .helper_functions import getIndexTensorAlongDim
 
 from .Parser.MathExprParser import MathExprParser
 from .Parser.MathExprLexer import MathExprLexer
@@ -25,7 +25,11 @@ class ImageMathNode:
             Returns a LATENT object that contains the result of the math expression applied to the input conditionings.
     """
     def __init__(self):
-        pass
+        shape = []
+        B = None
+        C = None
+        X = None
+        Y = None
 
     @classmethod
     def INPUT_TYPES(s):
@@ -91,10 +95,18 @@ class ImageMathNode:
         b = torch.zeros_like(a) if b is None else b
         c = torch.zeros_like(a) if c is None else c
         d = torch.zeros_like(a) if d is None else d
+        
+        batch = range(0,a.shape[0],1)
+        width = range(0,a.shape[1],1)
+        height = range(0,a.shape[2],1)
+        color = range(0,a.shape[3],1)
 
+        B = getIndexTensorAlongDim(a, 0)
+        W = getIndexTensorAlongDim(a, 1)
+        H = getIndexTensorAlongDim(a, 2)
+        C = getIndexTensorAlongDim(a, 3)
 
-
-        variables = {'a': a, 'b': b, 'c': c, 'd': d, 'w': w, 'x': x, 'y': y, 'z': z}
+        variables = {'a': a, 'b': b, 'c': c, 'd': d, 'w': w, 'x': x, 'y': y, 'z': z,'B':B,'X':W,'Y':H,'C':C,'W':a.shape[1],'H':a.shape[2] }
         input_stream = InputStream(Image)
         lexer = MathExprLexer(input_stream)
         stream = CommonTokenStream(lexer)
diff --git a/src/more_math/helper_functions.py b/src/more_math/helper_functions.py
new file mode 100644
index 0000000..80c0c5f
--- /dev/null
+++ b/src/more_math/helper_functions.py
@@ -0,0 +1,15 @@
+import torch
+
+def getIndexTensorAlongDim(tensor, dim):
+    shape = tensor.shape
+    
+    # Create values: shape (size of dim)
+    values = torch.arange(shape[dim], dtype=torch.float32)
+
+    # Reshape values to align with the target dimension
+    view_shape = [1] * len(shape)
+    view_shape[dim] = shape[dim]
+    values = values.view(*view_shape)
+
+    # Broadcast to full shape
+    return values.expand(*shape)
diff --git a/src/more_math/nodes.py b/src/more_math/nodes.py
index 8c3a151..46c879e 100644
--- a/src/more_math/nodes.py
+++ b/src/more_math/nodes.py
@@ -1,3 +1,4 @@
+import torch
 from .NoiseMathNode import NoiseMathNode
 from .FloatMathNode import FloatMathNode
 from .ConditioningMathNode import ConditioningMathNode