Update pk_workflow.py

pk_workflow.py

@@ -354,7 +354,12 @@ def _dpmpp_2m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
     denoised = sample + sigma_from_timestep * model_output
 
     compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype
-    sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
+    # Cached once and reused every call — `torchsde.BrownianTree` caches its internal tree keyed to the exact
+    # float value it was first queried with, and re-deriving "the same" sigma via a fresh `.to()` cast on a
+    # later call can land a few ULPs away from what the tree remembers, which it treats as an ordering error.
+    if scheduler._dpmpp_sde_sigmas is None:
+        scheduler._dpmpp_sde_sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
+    sigmas = scheduler._dpmpp_sde_sigmas
     sigma, sigma_next = sigmas[i], sigmas[i + 1]
     x = sample.to(dtype=compute_dtype)
     denoised = denoised.to(dtype=compute_dtype)
@@ -418,7 +423,12 @@ def _dpmpp_3m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
     denoised = sample + sigma_from_timestep * model_output
 
     compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype
-    sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
+    # Cached once and reused every call — `torchsde.BrownianTree` caches its internal tree keyed to the exact
+    # float value it was first queried with, and re-deriving "the same" sigma via a fresh `.to()` cast on a
+    # later call can land a few ULPs away from what the tree remembers, which it treats as an ordering error.
+    if scheduler._dpmpp_sde_sigmas is None:
+        scheduler._dpmpp_sde_sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
+    sigmas = scheduler._dpmpp_sde_sigmas
     sigma, sigma_next = sigmas[i], sigmas[i + 1]
     x = sample.to(dtype=compute_dtype)
     denoised = denoised.to(dtype=compute_dtype)
@@ -552,6 +562,7 @@ class use_schedule:
                 scheduler._dpmpp_sde_h_last = None
                 scheduler._dpmpp_sde_h_last_2 = None
                 scheduler._dpmpp_sde_noise_sampler = None
+                scheduler._dpmpp_sde_sigmas = None
                 scheduler._dpmpp_sde_seed = self.seed + offset
 
                 def stepped(model_output, timestep, sample, return_dict=True, _s=scheduler, _f=step_fn, **_kwargs):