Update pk_workflow.py

pk_workflow.py

@@ -370,8 +370,11 @@ def _dpmpp_2m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
         sigma = shift * base / (1 + (shift - 1) * base)
 
     if scheduler._dpmpp_sde_noise_sampler is None:
+        # `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.
         scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
-            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=False
+            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=True
         )
 
     def half_log_snr(s):
@@ -439,8 +442,11 @@ def _dpmpp_3m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
         sigma = shift * base / (1 + (shift - 1) * base)
 
     if scheduler._dpmpp_sde_noise_sampler is None:
+        # `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.
         scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
-            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=False
+            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=True
         )
 
     def half_log_snr(s):