From 9907a5e4f5b807a910b0969240aa2646d45395c4 Mon Sep 17 00:00:00 2001
From: Macpaul Lin <macpaul@gmail.com>
Date: Thu, 8 Jan 2026 22:52:25 +0800
Subject: [PATCH] fix(quant_ops): add numel, size, shape, dim, and ndim to mock
 QuantizedTensor

Signed-off-by: Macpaul Lin <macpaul@gmail.com>
---
 comfy/quant_ops.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py
index 1e0cc0304..e80e6bcdc 100644
--- a/comfy/quant_ops.py
+++ b/comfy/quant_ops.py
@@ -83,6 +83,30 @@ except ImportError as e:
             self._qdata.requires_grad_(requires_grad)
             return self
 
+        def numel(self):
+            if hasattr(self._layout_params, "orig_shape"):
+                import math
+                return math.prod(self._layout_params.orig_shape)
+            return self._qdata.numel()
+
+        @property
+        def shape(self):
+            if hasattr(self._layout_params, "orig_shape"):
+                return torch.Size(self._layout_params.orig_shape)
+            return self._qdata.shape
+
+        @property
+        def ndim(self):
+            return len(self.shape)
+
+        def size(self, dim=None):
+            if dim is None:
+                return self.shape
+            return self.shape[dim]
+
+        def dim(self):
+            return self.ndim
+
         def __getattr__(self, name):
             if name == "params":
                 return self._layout_params