Update pk_workflow.py

pk_workflow.py

@@ -175,6 +175,47 @@ SCHEDULE_SIGMA_FUNCS = {
     "normal": normal_sigmas,
 }
 
+def _euler_ancestral_step(scheduler, generator, model_output, timestep, sample, eta: float = 1.0, s_noise: float = 1.0):
+    """Ports k-diffusion's `sample_euler_ancestral_RF` — the flow-matching branch `sample_euler_ancestral`
+    dispatches to for `CONST`-style model sampling, which is what MiniMax-H3's `[0, 1]` sigma space is — onto one
+    `MiniMaxH3Scheduler.step()` call. Single model evaluation, same shape as `step()` itself, with fresh
+    ancestral noise injected each step instead of a plain Euler blend. Mirrors `step()`'s own care around
+    recomputing `sigma_from_timestep` from `timestep` rather than reading `self.sigmas` at the current index, for
+    the same numerical-consistency reason documented there.
+    """
+    if scheduler._step_index is None:
+        scheduler._step_index = scheduler.index_for_timestep(timestep) if scheduler._begin_index is None else scheduler._begin_index
+
+    if not isinstance(timestep, torch.Tensor):
+        timestep = torch.tensor(timestep, dtype=sample.dtype)
+    sigma_from_timestep = 1 - timestep.to(device=sample.device, dtype=sample.dtype)
+    while sigma_from_timestep.ndim < sample.ndim:
+        sigma_from_timestep = sigma_from_timestep.unsqueeze(-1)
+    denoised = sample + sigma_from_timestep * model_output
+
+    compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype
+    sigma = scheduler.sigmas[scheduler._step_index].to(device=sample.device, dtype=compute_dtype)
+    sigma_next = scheduler.sigmas[scheduler._step_index + 1].to(device=sample.device, dtype=compute_dtype)
+    x = sample.to(dtype=compute_dtype)
+    denoised = denoised.to(dtype=compute_dtype)
+
+    if sigma_next == 0:
+        prev_sample = denoised
+    else:
+        downstep_ratio = 1 + (sigma_next / sigma - 1) * eta
+        sigma_down = sigma_next * downstep_ratio
+        alpha_next = 1 - sigma_next
+        alpha_down = 1 - sigma_down
+        renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp_min(0).sqrt()
+        ratio = sigma_down / sigma
+        prev_sample = ratio * x + (1 - ratio) * denoised
+        if eta > 0:
+            noise = torch.randn(x.shape, dtype=x.dtype, device=x.device, generator=generator)
+            prev_sample = (alpha_next / alpha_down) * prev_sample + noise * s_noise * renoise_coeff
+
+    prev_sample = prev_sample.to(dtype=sample.dtype)
+    scheduler._step_index += 1
+    return prev_sample
 
 class use_schedule:
     """Set each scheduler's shift for one request, and — for anything but `native` — force its sigma grid onto
@@ -187,11 +228,13 @@ 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, 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):
         self.pipe = pipe
         self.attr_names = ["scheduler", "audio_scheduler"]
         self.shifts = [float(video_shift), float(audio_shift)]
         self.schedule_name = schedule_name
+        self.sampler_name = sampler_name
+        self.seed = int(seed)
         self.steps = int(steps)
         self.threshold_noise = float(threshold_noise)
         self._originals: dict = {}
@@ -215,11 +258,25 @@ class use_schedule:
                     return _f(_s, None, device, _grid)
 
                 scheduler.set_timesteps = forced
+
+        if self.sampler_name == "euler_ancestral":
+            # Separate `torch.Generator` per scheduler (offset seeds) so video and audio ancestral noise don't
+            # correlate — each generator advances across every step call to *that* scheduler over the request.
+            for offset, attr_name in enumerate(self.attr_names):
+                scheduler = getattr(self.pipe, attr_name)
+                generator = torch.Generator(device="cpu").manual_seed(self.seed + offset)
+
+                def stepped(model_output, timestep, sample, return_dict=True, _s=scheduler, _g=generator, **_kwargs):
+                    return (_euler_ancestral_step(_s, _g, model_output, timestep, sample),)
+
+                scheduler.step = stepped
         return self
 
     def __exit__(self, *_):
         for attr_name, original in self._originals.items():
-            getattr(self.pipe, attr_name).__dict__.pop("set_timesteps", None)
+            current = getattr(self.pipe, attr_name)
+            current.__dict__.pop("set_timesteps", None)
+            current.__dict__.pop("step", None)
             setattr(self.pipe, attr_name, original)
         return False