Update pk_workflow.py

pk_workflow.py

@@ -371,8 +371,13 @@ def _dpmpp_2m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
         # `cpu=True` runs the Brownian-bridge recursion on CPU rather than the GPU — negligible cost next to
         # the transformer forward pass, but noticeably more numerically stable than `cpu=False`, which is what
         # ComfyUI's own non-`_gpu`-suffixed variants default to for exactly this reason.
+        # Padded a hair beyond the real [min, max] span, not built exactly at it: querying a BrownianTree
+        # exactly on its own construction bound is a known torchsde precision edge (`tb<=t1`-style warnings).
+        # The pad only widens the tree's internal span — every query below still uses the real, un-padded sigma.
+        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+        pad = (sigma_max - sigma_min).clamp_min(1e-6) * 1e-4
         scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
-            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=True
+            x, sigma_min - pad, sigma_max + pad, seed=scheduler._dpmpp_sde_seed, cpu=True
         )
 
     def half_log_snr(s):
@@ -443,8 +448,13 @@ def _dpmpp_3m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
         # `cpu=True` runs the Brownian-bridge recursion on CPU rather than the GPU — negligible cost next to
         # the transformer forward pass, but noticeably more numerically stable than `cpu=False`, which is what
         # ComfyUI's own non-`_gpu`-suffixed variants default to for exactly this reason.
+        # Padded a hair beyond the real [min, max] span, not built exactly at it: querying a BrownianTree
+        # exactly on its own construction bound is a known torchsde precision edge (`tb<=t1`-style warnings).
+        # The pad only widens the tree's internal span — every query below still uses the real, un-padded sigma.
+        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+        pad = (sigma_max - sigma_min).clamp_min(1e-6) * 1e-4
         scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
-            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=True
+            x, sigma_min - pad, sigma_max + pad, seed=scheduler._dpmpp_sde_seed, cpu=True
         )
 
     def half_log_snr(s):