Update pk_workflow.py

pk_workflow.py

@@ -81,34 +81,141 @@ def time_shift_sigma(sigma: torch.Tensor, from_shift: float, to_shift: float) ->
     return to_shift * base / (1.0 + (to_shift - 1.0) * base)
 
 
-class use_linear_quadratic:
-    """Force MiniMax-H3's two schedulers onto the `linear_quadratic` grid for one pipeline call.
+# ----------------------------------------------------------------------------------------------------------------
+# BasicScheduler(sgm_uniform / simple / beta / ddim_uniform / normal)
+# ----------------------------------------------------------------------------------------------------------------
+# Five more of ComfyUI's `BasicScheduler` names, ported from `comfy/samplers.py`. Each is computed at the
+# *reference* shift (1.0 — where `time_snr_shift` is the identity, so `sigma(t) == t`) and reprojected onto each
+# scheduler's real shift by `time_shift_sigma`, exactly like `linear_quadratic_sigmas` already is and for the same
+# reason: it keeps the video and audio streams pinned to the same underlying denoising progress at each step,
+# which computing each stream's schedule independently at its own shift would not.
+#
+# `FLOW_TIMESTEPS` mirrors ComfyUI's `ModelSamplingDiscreteFlow`/`ModelSamplingAV` default of 1000 discrete steps
+# (`comfy/model_sampling.py`). Unverified specifically for MiniMax-H3's own `sampling_settings` — if a ported
+# schedule's shape looks visibly different from ComfyUI's own render at the same steps/seed, this is the first
+# thing to check.
+FLOW_TIMESTEPS = 1000
+
+
+def _reference_sigma(index_1based: int) -> float:
+    """`ModelSamplingAV.sigma(timestep)` at shift == 1.0: the shift formula is the identity, so this is just the
+    plain fraction `index / FLOW_TIMESTEPS`. `index_1based` matches ComfyUI's 1-based table construction
+    (`torch.arange(1, timesteps + 1) / timesteps`)."""
+    return index_1based / FLOW_TIMESTEPS
+
+
+def sgm_uniform_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `sgm_uniform`. Uniform in *timestep* space between the max and min sigma, dropping the point
+    that would land exactly on the minimum, then appending an exact 0.0. `steps + 1` sigmas."""
+    steps = int(steps)
+    timesteps = torch.linspace(float(FLOW_TIMESTEPS), 1.0, steps + 1)[:-1]
+    sigmas = (timesteps / FLOW_TIMESTEPS).tolist() + [0.0]
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+def normal_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `normal`. Same idea as `sgm_uniform` but the linspace includes both endpoints (the minimum
+    sigma is reached exactly, not dropped), with 0.0 still appended."""
+    steps = int(steps)
+    timesteps = torch.linspace(float(FLOW_TIMESTEPS), 1.0, steps)
+    sigmas = (timesteps / FLOW_TIMESTEPS).tolist() + [0.0]
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+def simple_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `simple`: evenly-spaced *indices* into the 1000-entry sigma table, walked from the high-noise
+    end, then 0.0 appended."""
+    steps = int(steps)
+    stride = FLOW_TIMESTEPS / steps
+    sigmas = [_reference_sigma(FLOW_TIMESTEPS - int(x * stride)) for x in range(steps)]
+    sigmas.append(0.0)
+    return torch.tensor(sigmas, dtype=torch.float32)
+
 
-    A context manager rather than a pipeline-block subclass on purpose: `MiniMaxH3Scheduler.set_timesteps` already
-    takes a fully-formed `sigmas=` schedule as public API, so nothing here reaches into the modular blocks, and the
-    override lives and dies inside one request.
+def ddim_uniform_sigmas(steps: int) -> torch.Tensor:
+    """ComfyUI's `ddim_uniform`: a fixed-stride walk through the sigma table starting one index in, reversed so
+    the highest sigma comes first, ending at 0.0."""
+    steps = int(steps)
+    stride = max(FLOW_TIMESTEPS // steps, 1)
+    sigmas = [0.0]
+    index = 1
+    while index < FLOW_TIMESTEPS:
+        sigmas.append(_reference_sigma(index))
+        index += stride
+    sigmas.reverse()
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+def beta_sigmas(steps: int, alpha: float = 0.6, beta: float = 0.6) -> torch.Tensor:
+    """ComfyUI's `beta` (arxiv.org/abs/2407.12173): table indices drawn from a Beta(alpha, beta) inverse CDF
+    instead of an even stride, biasing samples toward one end of the trajectory. Needs `scipy`."""
+    import numpy
+    import scipy.stats
+
+    steps = int(steps)
+    total = FLOW_TIMESTEPS - 1
+    positions = 1.0 - numpy.linspace(0.0, 1.0, steps, endpoint=False)
+    indices = numpy.rint(scipy.stats.beta.ppf(positions, alpha, beta) * total)
+    sigmas = []
+    last = -1
+    for value in indices:
+        if value != last:
+            sigmas.append(_reference_sigma(int(value) + 1))
+        last = value
+    sigmas.append(0.0)
+    return torch.tensor(sigmas, dtype=torch.float32)
+
+
+SCHEDULE_SIGMA_FUNCS = {
+    "linear_quadratic": linear_quadratic_sigmas,
+    "sgm_uniform": sgm_uniform_sigmas,
+    "simple": simple_sigmas,
+    "beta": beta_sigmas,
+    "ddim_uniform": ddim_uniform_sigmas,
+    "normal": normal_sigmas,
+}
+
+
+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.
     """
 
-    def __init__(self, pipe, steps: int, threshold_noise: float = 0.025, enabled: bool = True):
-        self.schedulers = [pipe.scheduler, pipe.audio_scheduler] if enabled else []
+    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.shifts = [float(video_shift), float(audio_shift)]
+        self.schedule_name = schedule_name
         self.steps = int(steps)
         self.threshold_noise = float(threshold_noise)
+        self._originals = []
 
     def __enter__(self):
-        video_sigmas = linear_quadratic_sigmas(self.steps, self.threshold_noise)
-        for scheduler in self.schedulers:
-            sigmas = time_shift_sigma(video_sigmas, VIDEO_SHIFT, float(scheduler.shift))
-            unbound = type(scheduler).set_timesteps
+        self._originals = [(scheduler, scheduler.shift) for scheduler in self.schedulers]
+        for scheduler, shift in zip(self.schedulers, self.shifts):
+            scheduler.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:
+                sigmas = time_shift_sigma(base, 1.0, float(scheduler.shift))
+                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, _f=unbound):
+                    return _f(_s, None, device, _grid)
 
-            scheduler.set_timesteps = forced
+                scheduler.set_timesteps = forced
         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
         return False