Review fixes

__init__.py

@@ -1,10 +1,4 @@
-import traceback
-
-try:
-    from .more_math.nodes import comfy_entrypoint as comfy_entrypoint
-except ImportError:
-    print("Something is seriously wrong")
-    traceback.print_exc()
+from .more_math.nodes import comfy_entrypoint as comfy_entrypoint
 
 WEB_DIRECTORY = "./web"
-__all__ = ["comfy_entrypoint", "WEB_DIRECTORY"]
+__all__ = ["comfy_entrypoint", "WEB_DIRECTORY"]

more_math/AudioMathNode.py

@@ -135,7 +135,7 @@ def execute(cls, V, F, Expression, length_mismatch="tile",batching=0, remember_s
             "batch_count": float(a_w.shape[0]),
         } | generate_dim_variables(a_w) | V_norm_waveforms | sample_rates
 
-        v_stacked, v_cnt = get_v_variable(V_norm_waveforms, length_mismatch=length_mismatch)
+        v_stacked, v_cnt = get_v_variable(V_norm_waveforms)
         if v_stacked is not None:
              variables["V"] = v_stacked
              variables["Vcnt"] = float(v_cnt)
@@ -167,6 +167,5 @@ def execute(cls, V, F, Expression, length_mismatch="tile",batching=0, remember_s
                 res_list.append({"waveform": result_chunk, "sample_rate": sample_rate})
             return (res_list, stack)
         else:
-            stack = stack if remember_stack else copy.deepcopy(stack)
 
             return ([{"waveform": result, "sample_rate": sample_rate}], stack)

more_math/ConditioningMathNode.py

@@ -147,7 +147,7 @@ def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile", batchi
             "batch_count": a.shape[0],
         } | generate_dim_variables(a) | V_norm_tensors
 
-        v_stacked, v_cnt = get_v_variable(V_norm_tensors, length_mismatch=length_mismatch)
+        v_stacked, v_cnt = get_v_variable(V_norm_tensors)
         if v_stacked is not None:
              variables["V"] = v_stacked
              variables["Vcnt"] = float(v_cnt)
@@ -193,7 +193,7 @@ def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile", batchi
             "batch_count": a_p.shape[0] if a_p.numel() > 0 else 0,
         } | generate_dim_variables(a_p) | V_norm_pooled
 
-        v_stacked, v_cnt = get_v_variable(V_norm_pooled, length_mismatch=length_mismatch)
+        v_stacked, v_cnt = get_v_variable(V_norm_pooled)
         if v_stacked is not None:
              variables_pi["V"] = v_stacked
              variables_pi["Vcnt"] = float(v_cnt)
@@ -251,5 +251,4 @@ def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile", batchi
             base[0] = (rtensor, new_dict)
             res_list = [base]
 
-        stack = stack if remember_stack else copy.deepcopy(stack)
         return (move_to_device(res_list, original_device),stack)

more_math/GuiderMathNode.py

@@ -92,7 +92,6 @@ def __init__(self, V, F, expression,expression1,stack):
             self.tree1 = expression1
         self.inner_model = None  # Will be set during sample
         self.sigmas = None
-        self.current_step = 0
         self.steps = 0
         self.stck = stack
 
@@ -142,7 +141,6 @@ def lazy_eval():
         eval_samples, variables = self.setVars(x, sigma, seed, lazy_g_results)
         visitor = UnifiedMathVisitor(variables, eval_samples.shape,eval_samples.device,state_storage=self.stck)
         result_tensor = visitor.visit(self.tree)
-        self.current_step = self.current_step + 1;
         return as_tensor(result_tensor, eval_samples.shape).to(x.device)
 
 
@@ -164,6 +162,13 @@ def setVars(self, x, sigma, seed, g_results):
 
         frame_count = eval_samples.shape[time_dim] if time_dim is not None else 1
 
+        # sigma has one entry per batch item, all equal
+        sigma = float(sigma.flatten()[0]) if isinstance(sigma, torch.Tensor) else sigma
+        current_step = 0
+        if self.sigmas is not None:
+            # Index of the step whose interval contains sigma, so extra model calls within a step (heun, dpmpp_2s, ...) keep the same index
+            current_step = max((i for i, s in enumerate(self.sigmas.tolist()) if s >= sigma), default=0)
+
         def make_lazy_var(value_fn):
             value_fn.is_lazy_var = True
             return value_fn
@@ -184,10 +189,10 @@ def make_lazy_var(value_fn):
             "batch_count": float(eval_samples.shape[0]),
             "N": float(eval_samples.shape[channel_dim]) if channel_dim < ndim else 0.0,
             "channel_count": float(eval_samples.shape[channel_dim]) if channel_dim < ndim else 0.0,
-            "sigma": sigma.item() if isinstance(sigma,torch.Tensor) else sigma,
+            "sigma": sigma,
             "seed": seed if seed is not None else 0,
             "steps": self.steps,
-            "current_step": self.current_step,
+            "current_step": current_step,
             "sample": x
         })
         if g_results is not None:
@@ -226,7 +231,7 @@ def lazy_v_cnt():
 
     def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
         self.sigmas = sigmas
-        self.steps = len(sigmas)
+        self.steps = len(sigmas) - 1
         if sigmas.shape[-1] == 0:
             return latent_image
         active_guiders = [g for g in self.V.values() if g is not None]

more_math/ImageMathNode.py

@@ -129,7 +129,7 @@ def execute(cls, V, F, Expression, length_mismatch="error",batching=0,remember_s
         # Add all dynamic inputs
         variables.update(V_norm)
 
-        v_stacked, v_cnt = get_v_variable(V_norm, length_mismatch=length_mismatch)
+        v_stacked, v_cnt = get_v_variable(V_norm)
         if v_stacked is not None:
              variables["V"] = v_stacked
              variables["Vcnt"] = float(v_cnt)
@@ -159,8 +159,6 @@ def execute(cls, V, F, Expression, length_mismatch="error",batching=0,remember_s
             res_list = []
             for result_chunk in res:
                 res_list.append(result_chunk)
-            stack = stack if remember_stack else copy.deepcopy(stack)
             return (res_list, stack)
         else:
-            stack = stack if remember_stack else copy.deepcopy(stack)
             return ([result], stack)

more_math/LatentMathNode.py

@@ -215,7 +215,7 @@ def _resolve_alias(base):
         # Add all dynamic inputs
         variables.update(V_norm_samples)
 
-        v_stacked, v_cnt = get_v_variable(V_norm_samples, length_mismatch=length_mismatch)
+        v_stacked, v_cnt = get_v_variable(V_norm_samples)
         if v_stacked is not None:
              variables["V"] = v_stacked
              variables["Vcnt"] = float(v_cnt)
@@ -242,7 +242,6 @@ def _resolve_alias(base):
         if getattr(result_t, "is_nested", False):
             rl = result_latent.copy()
             rl["samples"] = result_t
-            stack = stack if remember_stack else copy.deepcopy(stack)
             return (move_to_device([rl], original_device), stack)
 
         # If the visitor produced a list/tuple, emit each element as a separate latent.
@@ -252,7 +251,6 @@ def _resolve_alias(base):
                 rl = result_latent.copy()
                 rl["samples"] = comp
                 results.append(rl)
-                stack = stack if remember_stack else copy.deepcopy(stack)
             return (move_to_device(results, original_device), stack)
 
         if batching > 0:
@@ -262,9 +260,7 @@ def _resolve_alias(base):
                 rl = result_latent.copy()
                 rl["samples"] = result_tensor
                 results.append(rl)
-                stack = stack if remember_stack else copy.deepcopy(stack)
             return (move_to_device(results, original_device), stack)
         rl = result_latent.copy()
         rl["samples"] = result_t
-        stack = stack if remember_stack else copy.deepcopy(stack)
         return (move_to_device([rl], original_device), stack)

more_math/MaskMathNode.py

@@ -108,7 +108,7 @@ def execute(cls, V, F, Expression, length_mismatch="tile",batching=0,remember_st
             "batch_count": float(ae.shape[0]),
         } | generate_dim_variables(ae)
 
-        v_stacked, v_cnt = get_v_variable(V_norm, length_mismatch=length_mismatch)
+        v_stacked, v_cnt = get_v_variable(V_norm)
         if v_stacked is not None:
              variables["V"] = v_stacked
              variables["Vcnt"] = float(v_cnt)
@@ -140,8 +140,6 @@ def execute(cls, V, F, Expression, length_mismatch="tile",batching=0,remember_st
             res_list = []
             for result_chunk in res:
                 res_list.append(result_chunk)
-                stack = stack if remember_stack else copy.deepcopy(stack)
             return (res_list, stack)
         else:
-            stack = stack if remember_stack else copy.deepcopy(stack)
             return ([result], stack)