Update pk_workflow.py

pk_workflow.py

@@ -552,10 +552,21 @@ class use_schedule:
                 scheduler = getattr(self.pipe, attr_name)
                 sigmas_full = time_shift_sigma(base, 1.0, float(scheduler.shift))
                 sigmas = sigmas_full[self.stage_from : self.stage_from + self.steps + 1]
-                unbound = type(scheduler).set_timesteps
 
-                def forced(num_inference_steps=None, device=None, sigmas=None, _s=scheduler, _grid=sigmas, _f=unbound):
-                    return _f(_s, None, device, _grid)
+                def forced(num_inference_steps=None, device=None, sigmas=None, _s=scheduler, _grid=sigmas):
+                    # Replicates `MiniMaxH3Scheduler.set_timesteps`'s own `sigmas=` branch exactly
+                    # (`scheduling_minimax_h3.py`) rather than calling through to it — its validation requires
+                    # the array to end at exactly 0.0, correct for a complete trajectory but wrong for an
+                    # intermediate Staged Denoising slice, which legitimately ends at whatever sigma this stage
+                    # stops at. Still requires strictly decreasing, at least two points.
+                    grid = torch.as_tensor(_grid, dtype=torch.float32).flatten().cpu()
+                    if grid.numel() < 2 or not bool((grid[1:] < grid[:-1]).all()):
+                        raise ValueError("`sigmas` must hold at least two strictly decreasing values.")
+                    _s.sigmas = grid.to(device=device)
+                    _s.timesteps = (1.0 - grid[:-1]).to(device=device)
+                    _s.num_inference_steps = int(_s.timesteps.numel())
+                    _s._step_index = None
+                    _s._begin_index = None
 
                 scheduler.set_timesteps = forced