diff --git a/js/mathExpression.js b/js/mathExpression.js
index 28e5088..2409cf1 100644
--- a/js/mathExpression.js
+++ b/js/mathExpression.js
@@ -11,4 +11,25 @@ app.registerExtension({
 			return r;
 		};
 	},
+	beforeRegisterNodeDef(nodeType) {
+		if (nodeType.comfyClass === "MathExpression|pysssss") {
+			const onDrawForeground = nodeType.prototype.onDrawForeground;
+			nodeType.prototype.onDrawForeground = function (ctx) {
+				const r = onDrawForeground?.apply?.(this, arguments);
+
+				const v = app.nodeOutputs?.[this.id + ""];
+				if (!this.flags.collapsed && v) {
+					const text = v.value[0] + "";
+					ctx.save();
+					ctx.font = "bold 12px sans-serif";
+					ctx.fillStyle = "dodgerblue";
+					const sz = ctx.measureText(text);
+					ctx.fillText(text, this.size[0] - sz.width - 5, LiteGraph.NODE_SLOT_HEIGHT * 3);
+					ctx.restore();
+				}
+
+				return r;
+			};
+		}
+	},
 });
diff --git a/py/math_expression.py b/py/math_expression.py
index 3e73dda..d0b7be9 100644
--- a/py/math_expression.py
+++ b/py/math_expression.py
@@ -44,26 +44,25 @@ class MathExpression:
                 "expression": ("STRING", {"multiline": True, "dynamicPrompts": False}),
             },
             "optional": {
-                "a": ("FLOAT", {"forceInput": True}),
-                "b": ("FLOAT", {"forceInput": True}),
-                "c": ("FLOAT", {"forceInput": True}),
+                "a": ("INT,FLOAT,IMAGE,LATENT", {"forceInput": True}),
+                "b": ("INT,FLOAT,IMAGE,LATENT", {"forceInput": True}),
+                "c": ("INT,FLOAT,IMAGE,LATENT", {"forceInput": True}),
             },
             "hidden": {"extra_pnginfo": "EXTRA_PNGINFO",
                        "prompt": "PROMPT"},
         }
 
-    RETURN_TYPES = ("INT", "FLOAT", "STRING", )
+    RETURN_TYPES = ("INT", "FLOAT", )
     FUNCTION = "evaluate"
     CATEGORY = "utils"
+    OUTPUT_NODE = True
 
-    
     @classmethod
     def IS_CHANGED(s, expression, **kwargs):
         if "random" in expression:
             return float("nan")
         return expression
 
-
     def get_widget_value(self, extra_pnginfo, prompt, node_name, widget_name):
         workflow = extra_pnginfo["workflow"]
         node_id = None
@@ -88,10 +87,24 @@ class MathExpression:
             raise NameError(f"Widget not found: {node_name}.{widget_name}")
         raise NameError(f"Node not found: {node_name}.{widget_name}")
 
+    def get_size(self, target, property):
+        if isinstance(target, dict) and "samples" in target:
+            # Latent
+            if property == "width":
+                return target["samples"].shape[3] * 8
+            return target["samples"].shape[2] * 8
+        else:
+            # Image
+            if property == "width":
+                return target.shape[2] 
+            return target.shape[1] 
+
     def evaluate(self, expression, extra_pnginfo, prompt, a=None, b=None, c=None):
         expression = expression.replace('\n', ' ').replace('\r', '')
         node = ast.parse(expression, mode='eval').body
 
+        lookup = {"a": a, "b": b, "c": c}
+
         def eval_expr(node):
             if isinstance(node, ast.Num):
                 return node.n
@@ -100,14 +113,18 @@ class MathExpression:
             elif isinstance(node, ast.UnaryOp):
                 return operators[type(node.op)](eval_expr(node.operand))
             elif isinstance(node, ast.Attribute):
+                if node.value.id in lookup:
+                    if node.attr == "width" or node.attr == "height":
+                        return self.get_size(lookup[node.value.id], node.attr)
+
                 return self.get_widget_value(extra_pnginfo, prompt, node.value.id, node.attr)
             elif isinstance(node, ast.Name):
-                if node.id == "a":
-                    return a
-                if node.id == "b":
-                    return b
-                if node.id == "c":
-                    return c
+                if node.id in lookup:
+                    val = lookup[node.id]
+                    if isinstance(val, (int, float, complex)):
+                        return val
+                    else:
+                        raise TypeError(f"Compex types (LATENT/IMAGE) need to reference their width/height, e.g. {node.id}.width")
                 raise NameError(f"Name not found: {node.id}")
             elif isinstance(node, ast.Call):
                 if node.func.id in functions:
@@ -125,7 +142,7 @@ class MathExpression:
                 raise TypeError(node)
 
         r = eval_expr(node)
-        return (int(r), float(r), str(r))
+        return {"ui": {"value": [r]}, "result": (int(r), float(r),)}
 
 
 NODE_CLASS_MAPPINGS = {