Update pk_workflow.py

pk_workflow.py

@@ -179,29 +179,35 @@ SCHEDULE_SIGMA_FUNCS = {
 class use_schedule:
     """Set each scheduler's shift for one request, and — for anything but `native` — force its sigma grid onto
     one of `SCHEDULE_SIGMA_FUNCS`'s named schedules.
-    Shift is applied unconditionally, including under `native`, so the shift sliders affect the pipeline's own
-    default schedule too and not just the custom ones — and it is always restored on exit, since
-    `pipe.scheduler`/`pipe.audio_scheduler` are shared, request-spanning objects that must not carry one
-    request's shift into the next.
+    `MiniMaxH3Scheduler.shift` is a read-only property, so a different shift means swapping in a freshly built
+    scheduler via `from_config(..., shift=...)` rather than mutating one in place — the standard diffusers idiom
+    for changing a `ConfigMixin` parameter after construction, and correct regardless of exactly how `shift` is
+    stored internally. Applied unconditionally, including under `native`, so the shift sliders affect the
+    pipeline's own default schedule too — and always restored on exit, since `pipe.scheduler`/`pipe.audio_scheduler`
+    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, threshold_noise: float = 0.025):
-        self.schedulers = [pipe.scheduler, pipe.audio_scheduler]
+        self.pipe = pipe
+        self.attr_names = ["scheduler", "audio_scheduler"]
         self.shifts = [float(video_shift), float(audio_shift)]
         self.schedule_name = schedule_name
         self.steps = int(steps)
         self.threshold_noise = float(threshold_noise)
-        self._originals = []
+        self._originals: dict = {}
 
     def __enter__(self):
-        self._originals = [(scheduler, scheduler.shift) for scheduler in self.schedulers]
-        for scheduler, shift in zip(self.schedulers, self.shifts):
-            scheduler.shift = shift
+        for attr_name, shift in zip(self.attr_names, self.shifts):
+            original = getattr(self.pipe, attr_name)
+            self._originals[attr_name] = original
+            if float(original.shift) != shift:
+                setattr(self.pipe, attr_name, type(original).from_config(original.config, shift=shift))
 
         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)
-            for scheduler in self.schedulers:
+            for attr_name in self.attr_names:
+                scheduler = getattr(self.pipe, attr_name)
                 sigmas = time_shift_sigma(base, 1.0, float(scheduler.shift))
                 unbound = type(scheduler).set_timesteps
 
@@ -212,10 +218,9 @@ class use_schedule:
         return self
 
     def __exit__(self, *_):
-        for scheduler in self.schedulers:
-            scheduler.__dict__.pop("set_timesteps", None)
-        for scheduler, original_shift in self._originals:
-            scheduler.shift = original_shift
+        for attr_name, original in self._originals.items():
+            getattr(self.pipe, attr_name).__dict__.pop("set_timesteps", None)
+            setattr(self.pipe, attr_name, original)
         return False