Update pk_workflow.py

pk_workflow.py

@@ -284,6 +284,12 @@ def _er_sde_step(scheduler, generator, model_output, timestep, sample, s_noise:
 
         if s_noise > 0:
             noise = torch.randn(x.shape, dtype=x.dtype, device="cpu", generator=generator).to(x.device)
+            print(
+                f"[noise-debug] er_sde i={scheduler._step_index} x.shape={tuple(x.shape)} "
+                f"noise.shape={tuple(noise.shape)} noise.std={noise.std().item():.4f} "
+                f"noise.mean={noise.mean().item():.4f}",
+                flush=True,
+            )
             spread = (er_lambda_t**2 - er_lambda_s**2 * r**2).clamp_min(0).sqrt()
             prev_sample = prev_sample + alpha_t * noise * s_noise * spread
 
@@ -401,6 +407,12 @@ def _dpmpp_2m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
 
         if eta > 0 and s_noise > 0:
             noise = scheduler._dpmpp_sde_noise_sampler(sigma, sigma_next).to(device=x.device, dtype=compute_dtype)
+            print(
+                f"[noise-debug] dpmpp_2m_sde i={scheduler._step_index} x.shape={tuple(x.shape)} "
+                f"noise.shape={tuple(noise.shape)} noise.std={noise.std().item():.4f} "
+                f"noise.mean={noise.mean().item():.4f}",
+                flush=True,
+            )
             prev_sample = prev_sample + noise * sigma_next * (-2 * h * eta).expm1().neg().sqrt() * s_noise
 
         scheduler._dpmpp_sde_h_last = h
@@ -492,6 +504,12 @@ def _dpmpp_3m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1
 
         if eta > 0 and s_noise > 0:
             noise = scheduler._dpmpp_sde_noise_sampler(sigma, sigma_next).to(device=x.device, dtype=compute_dtype)
+            print(
+                f"[noise-debug] dpmpp_3m_sde i={scheduler._step_index} x.shape={tuple(x.shape)} "
+                f"noise.shape={tuple(noise.shape)} noise.std={noise.std().item():.4f} "
+                f"noise.mean={noise.mean().item():.4f}",
+                flush=True,
+            )
             prev_sample = prev_sample + noise * sigma_next * (-2 * h * eta).expm1().neg().sqrt() * s_noise
 
         scheduler._dpmpp_sde_h_last_2 = h_1