From 5680bae108620c31e574dab45f091da7006eeb84 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Sun, 18 Oct 2026 12:53:59 +0200 Subject: [PATCH] Fix guider math with batched sigma and step tracking Co-Authored-By: Claude Opus 5.5 --- more_math/GuiderMathNode.py | 15 ++++++++++----- tests/test_guider_math.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/more_math/GuiderMathNode.py b/more_math/GuiderMathNode.py index 70e9d63..26d7114 100644 --- a/more_math/GuiderMathNode.py +++ b/more_math/GuiderMathNode.py @@ -92,7 +92,6 @@ class MathGuider: 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 @@ class MathGuider: 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 @@ class MathGuider: 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 @@ class MathGuider: "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 @@ class MathGuider: 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] diff --git a/tests/test_guider_math.py b/tests/test_guider_math.py index a00b099..56892ad 100644 --- a/tests/test_guider_math.py +++ b/tests/test_guider_math.py @@ -102,6 +102,16 @@ class TestMathGuider(unittest.TestCase): # At step 2 with 2 steps: 2 / 2 = 1.0 self.assertTrue(torch.allclose(res2, torch.tensor(1.0))) + def test_math_guider_batched_sigma_and_substeps(self): + # Samplers pass sigma with one entry per batch item; heun-like samplers also call the model between sigmas + math_guider = MathGuider({"V0": MockGuider(1.0)}, {}, "current_step + sigma * 0", "V0", {}) + math_guider.sigmas = torch.tensor([10.0, 5.0, 0.0]) + math_guider.steps = 2 + x = torch.zeros((2, 1, 1, 1)) + self.assertTrue(torch.allclose(math_guider(x, torch.tensor([5.0, 5.0])), torch.tensor(1.0))) + self.assertTrue(torch.allclose(math_guider(x, torch.tensor([7.0, 7.0])), torch.tensor(0.0))) + self.assertTrue(torch.allclose(math_guider(x, torch.tensor([10.0, 10.0])), torch.tensor(0.0))) + def test_math_guider_crop_function(self): """Test crop function within guider context""" g0 = MockGuider(1.0)