Update pk_workflow.py

pk_workflow.py

@@ -217,6 +217,83 @@ def _euler_ancestral_step(scheduler, generator, model_output, timestep, sample,
     scheduler._step_index += 1
     return prev_sample
 
+def _er_sde_step(scheduler, generator, model_output, timestep, sample, s_noise: float = 1.0, max_stage: int = 3):
+    """Ports k-diffusion's `sample_er_sde` (VP ER-SDE-Solver-3, arXiv:2309.06169) onto one
+    `MiniMaxH3Scheduler.step()` call. Single model evaluation per step — second/third-order accuracy comes from
+    the previous one or two steps' denoised estimates, not an extra evaluation this step — so it carries history
+    on the scheduler instance across calls, reset each request by `use_schedule` alongside `_step_index`.
+    """
+    if scheduler._step_index is None:
+        scheduler._step_index = scheduler.index_for_timestep(timestep) if scheduler._begin_index is None else scheduler._begin_index
+    i = scheduler._step_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
+    sigmas = scheduler.sigmas.to(device=sample.device, dtype=compute_dtype)
+    sigma, sigma_next = sigmas[i], sigmas[i + 1]
+    x = sample.to(dtype=compute_dtype)
+    denoised = denoised.to(dtype=compute_dtype)
+
+    if i == 0 and float(sigma) >= 1.0:
+        # `1 - sigma` sits in a denominator below; MiniMax-H3's first sigma is exactly 1.0, so nudge it a hair
+        # under 1.0 for this sampler's math only, matching ComfyUI's `offset_first_sigma_for_snr`. Does not
+        # touch `sigma_from_timestep` above — the model was still conditioned on the real timestep.
+        base = torch.tensor(1.0 - 1e-4, dtype=compute_dtype, device=sample.device)
+        shift = float(scheduler.shift)
+        sigma = shift * base / (1 + (shift - 1) * base)
+
+    def er_lambda(s):
+        return s / (1 - s)
+
+    def noise_scaler(v):
+        return v * (v**0.3).exp() + v * 10.0
+
+    if sigma_next == 0:
+        prev_sample = denoised
+    else:
+        er_lambda_s, er_lambda_t = er_lambda(sigma), er_lambda(sigma_next)
+        alpha_s, alpha_t = 1 - sigma, 1 - sigma_next
+        r_alpha = alpha_t / alpha_s
+        r = noise_scaler(er_lambda_t) / noise_scaler(er_lambda_s)
+
+        prev_sample = r_alpha * r * x + alpha_t * (1 - r) * denoised
+
+        stage_used = min(max_stage, i + 1)
+        if stage_used >= 2:
+            num_points = 200
+            dt = er_lambda_t - er_lambda_s
+            step_size = -dt / num_points
+            positions = er_lambda_t + torch.arange(num_points, device=x.device, dtype=compute_dtype) * step_size
+            scaled = noise_scaler(positions)
+
+            s_term = torch.sum(1 / scaled) * step_size
+            er_lambda_prev = er_lambda(sigmas[i - 1])
+            denoised_d = (denoised - scheduler._er_sde_old_denoised) / (er_lambda_s - er_lambda_prev)
+            prev_sample = prev_sample + alpha_t * (dt + s_term * noise_scaler(er_lambda_t)) * denoised_d
+
+            if stage_used >= 3:
+                s_u_term = torch.sum((positions - er_lambda_s) / scaled) * step_size
+                er_lambda_prev2 = er_lambda(sigmas[i - 2])
+                denoised_u = (denoised_d - scheduler._er_sde_old_denoised_d) / ((er_lambda_s - er_lambda_prev2) / 2)
+                prev_sample = prev_sample + alpha_t * ((dt**2) / 2 + s_u_term * noise_scaler(er_lambda_t)) * denoised_u
+            scheduler._er_sde_old_denoised_d = denoised_d
+
+        if s_noise > 0:
+            noise = torch.randn(x.shape, dtype=x.dtype, device="cpu", generator=generator).to(x.device)
+            spread = (er_lambda_t**2 - er_lambda_s**2 * r**2).clamp_min(0).sqrt()
+            prev_sample = prev_sample + alpha_t * noise * s_noise * spread
+
+    scheduler._er_sde_old_denoised = denoised
+    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
     one of `SCHEDULE_SIGMA_FUNCS`'s named schedules.
@@ -269,6 +346,17 @@ class use_schedule:
                 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
+        elif self.sampler_name == "er_sde":
+            for offset, attr_name in enumerate(self.attr_names):
+                scheduler = getattr(self.pipe, attr_name)
+                scheduler._er_sde_old_denoised = None
+                scheduler._er_sde_old_denoised_d = None
+                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 (_er_sde_step(_s, _g, model_output, timestep, sample),)
+
                 scheduler.step = stepped
         return self