Update pk_workflow.py

pk_workflow.py

@@ -514,7 +514,7 @@ class use_schedule:
     are shared, request-spanning objects that must not carry one request's shift into the next.
     """
 
-    def __init__(self, pipe, steps: int, schedule_name: str, video_shift: float, audio_shift: float, sampler_name: str = "euler", seed: int = 0, threshold_noise: float = 0.025):
+    def __init__(self, pipe, steps: int, schedule_name: str, video_shift: float, audio_shift: float, sampler_name: str = "euler", seed: int = 0, threshold_noise: float = 0.025, total_steps: int | None = None, stage_from: int = 0):
         self.pipe = pipe
         self.attr_names = ["scheduler", "audio_scheduler"]
         self.shifts = [float(video_shift), float(audio_shift)]
@@ -523,6 +523,15 @@ class use_schedule:
         self.seed = int(seed)
         self.steps = int(steps)
         self.threshold_noise = float(threshold_noise)
+        # Staged Denoising: `total_steps` builds the schedule at the eventual target length rather than `steps`,
+        # and `stage_from` slices this stage's own `steps`-length span out of it
+        # (`sigmas_full[stage_from : stage_from + steps + 1]`) rather than the schedule's own start. Every
+        # schedule and every custom-step sampler is a pure function of its sigma array, so a slice of a longer
+        # schedule is mathematically indistinguishable, from the sampler's perspective, from a complete schedule
+        # of that length — nothing below this class needs to know a stage boundary exists. Defaults reproduce
+        # the unstaged behavior exactly (`total_steps=None` falls back to `steps`, `stage_from=0`).
+        self.total_steps = int(total_steps) if total_steps is not None else int(steps)
+        self.stage_from = int(stage_from)
         self._originals: dict = {}
 
     def __enter__(self):
@@ -534,10 +543,15 @@ class use_schedule:
 
         if self.schedule_name != "native":
             sigma_func = SCHEDULE_SIGMA_FUNCS[self.schedule_name]
-            base = sigma_func(self.steps, self.threshold_noise) if sigma_func is linear_quadratic_sigmas else sigma_func(self.steps)
+            base = (
+                sigma_func(self.total_steps, self.threshold_noise)
+                if sigma_func is linear_quadratic_sigmas
+                else sigma_func(self.total_steps)
+            )
             for attr_name in self.attr_names:
                 scheduler = getattr(self.pipe, attr_name)
-                sigmas = time_shift_sigma(base, 1.0, float(scheduler.shift))
+                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):