Update pk_workflow.py

pk_workflow.py

@@ -294,6 +294,191 @@ def _er_sde_step(scheduler, generator, model_output, timestep, sample, s_noise:
     scheduler._step_index += 1
     return prev_sample
 
+class _BatchedBrownianTree:
+    """Minimal port of k-diffusion's `BatchedBrownianTree` (single-seed case only — MiniMax-H3 requests run at
+    batch size 1). Wraps `torchsde.BrownianTree` so consecutive noise draws at adjacent sigma pairs are
+    correlated through a shared stochastic path, as `dpmpp_2m_sde`/`dpmpp_3m_sde` require — independent
+    per-step Gaussian noise (as used for `euler_ancestral`/`er_sde` above) is a materially different sampler.
+    """
+
+    def __init__(self, x, t0, t1, seed, cpu=False):
+        import torchsde
+
+        self.cpu_tree = cpu
+        if t0 > t1:
+            t0, t1, self.sign = t1, t0, -1
+        else:
+            self.sign = 1
+        w0 = torch.zeros_like(x)
+        if self.cpu_tree:
+            t0, w0, t1 = t0.detach().cpu(), w0.detach().cpu(), t1.detach().cpu()
+        self.tree = torchsde.BrownianTree(t0, w0, t1, entropy=seed)
+
+    def __call__(self, t0, t1):
+        if t0 > t1:
+            t0, t1, sign = t1, t0, -1
+        else:
+            sign = 1
+        device, dtype = t0.device, t0.dtype
+        if self.cpu_tree:
+            t0, t1 = t0.detach().cpu().float(), t1.detach().cpu().float()
+        return self.tree(t0, t1).to(device=device, dtype=dtype) * (self.sign * sign)
+
+
+class _BrownianTreeNoiseSampler:
+    """Port of k-diffusion's `BrownianTreeNoiseSampler`. `cpu=False` matches the `*_gpu` sampler names — noise
+    is generated directly on the accelerator rather than the CPU-tree variant the non-`_gpu` names use."""
+
+    def __init__(self, x, sigma_min, sigma_max, seed, cpu=False):
+        self.tree = _BatchedBrownianTree(x, torch.as_tensor(sigma_min), torch.as_tensor(sigma_max), seed, cpu=cpu)
+
+    def __call__(self, sigma, sigma_next):
+        t0, t1 = torch.as_tensor(sigma), torch.as_tensor(sigma_next)
+        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
+
+def _dpmpp_2m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1.0, s_noise: float = 1.0):
+    """Ports k-diffusion's `sample_dpmpp_2m_sde` (`solver_type='midpoint'`, matching ComfyUI's `dpmpp_2m_sde_gpu`
+    — the `_heun` variant is a different `solver_type` and is not ported here) onto one `MiniMaxH3Scheduler.step()`
+    call. Single model evaluation per step; second-order accuracy comes from the previous step's denoised
+    estimate. History and the Brownian-tree noise sampler live on the scheduler instance, reset each request.
+    """
+    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:
+        base = torch.tensor(1.0 - 1e-4, dtype=compute_dtype, device=sample.device)
+        shift = float(scheduler.shift)
+        sigma = shift * base / (1 + (shift - 1) * base)
+
+    if scheduler._dpmpp_sde_noise_sampler is None:
+        scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
+            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=False
+        )
+
+    def half_log_snr(s):
+        return torch.log((1 - s) / s)
+
+    if sigma_next == 0:
+        prev_sample = denoised
+    else:
+        lambda_s, lambda_t = half_log_snr(sigma), half_log_snr(sigma_next)
+        h = lambda_t - lambda_s
+        h_eta = h * (eta + 1)
+        alpha_next = 1 - sigma_next
+
+        prev_sample = (sigma_next / sigma) * (-h * eta).exp() * x + alpha_next * (-h_eta).expm1().neg() * denoised
+
+        old_denoised = scheduler._dpmpp_sde_old_denoised
+        h_last = scheduler._dpmpp_sde_h_last
+        if old_denoised is not None:
+            r = h_last / h
+            prev_sample = prev_sample + 0.5 * alpha_next * (-h_eta).expm1().neg() * (1 / r) * (denoised - old_denoised)
+
+        if eta > 0 and s_noise > 0:
+            noise = scheduler._dpmpp_sde_noise_sampler(sigma, sigma_next).to(device=x.device, dtype=compute_dtype)
+            prev_sample = prev_sample + noise * sigma_next * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+        scheduler._dpmpp_sde_h_last = h
+
+    scheduler._dpmpp_sde_old_denoised = denoised
+    prev_sample = prev_sample.to(dtype=sample.dtype)
+    scheduler._step_index += 1
+    return prev_sample
+
+
+def _dpmpp_3m_sde_step(scheduler, model_output, timestep, sample, eta: float = 1.0, s_noise: float = 1.0):
+    """Ports k-diffusion's `sample_dpmpp_3m_sde` (matching ComfyUI's `dpmpp_3m_sde_gpu`) onto one
+    `MiniMaxH3Scheduler.step()` call. Single model evaluation per step; third-order accuracy (once two prior
+    steps exist) comes from history carried on the scheduler instance, plus the same Brownian-tree noise as
+    `_dpmpp_2m_sde_step`.
+    """
+    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:
+        base = torch.tensor(1.0 - 1e-4, dtype=compute_dtype, device=sample.device)
+        shift = float(scheduler.shift)
+        sigma = shift * base / (1 + (shift - 1) * base)
+
+    if scheduler._dpmpp_sde_noise_sampler is None:
+        scheduler._dpmpp_sde_noise_sampler = _BrownianTreeNoiseSampler(
+            x, sigmas[sigmas > 0].min(), sigmas.max(), seed=scheduler._dpmpp_sde_seed, cpu=False
+        )
+
+    def half_log_snr(s):
+        return torch.log((1 - s) / s)
+
+    if sigma_next == 0:
+        prev_sample = denoised
+    else:
+        lambda_s, lambda_t = half_log_snr(sigma), half_log_snr(sigma_next)
+        h = lambda_t - lambda_s
+        h_eta = h * (eta + 1)
+        alpha_next = 1 - sigma_next
+
+        prev_sample = (sigma_next / sigma) * (-h * eta).exp() * x + alpha_next * (-h_eta).expm1().neg() * denoised
+
+        denoised_1 = scheduler._dpmpp_sde_old_denoised
+        denoised_2 = scheduler._dpmpp_sde_old_denoised_2
+        h_1 = scheduler._dpmpp_sde_h_last
+        h_2 = scheduler._dpmpp_sde_h_last_2
+
+        if h_2 is not None:
+            r0, r1 = h_1 / h, h_2 / h
+            d1_0 = (denoised - denoised_1) / r0
+            d1_1 = (denoised_1 - denoised_2) / r1
+            d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
+            d2 = (d1_0 - d1_1) / (r0 + r1)
+            phi_2 = h_eta.neg().expm1() / h_eta + 1
+            phi_3 = phi_2 / h_eta - 0.5
+            prev_sample = prev_sample + (alpha_next * phi_2) * d1 - (alpha_next * phi_3) * d2
+        elif h_1 is not None:
+            r = h_1 / h
+            d = (denoised - denoised_1) / r
+            phi_2 = h_eta.neg().expm1() / h_eta + 1
+            prev_sample = prev_sample + (alpha_next * phi_2) * d
+
+        if eta > 0 and s_noise > 0:
+            noise = scheduler._dpmpp_sde_noise_sampler(sigma, sigma_next).to(device=x.device, dtype=compute_dtype)
+            prev_sample = prev_sample + noise * sigma_next * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+        scheduler._dpmpp_sde_h_last_2 = h_1
+        scheduler._dpmpp_sde_h_last = h
+
+    scheduler._dpmpp_sde_old_denoised_2 = scheduler._dpmpp_sde_old_denoised
+    scheduler._dpmpp_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.
@@ -357,6 +542,21 @@ class use_schedule:
                 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
+        elif self.sampler_name in ("dpmpp_2m_sde_gpu", "dpmpp_3m_sde_gpu"):
+            step_fn = _dpmpp_2m_sde_step if self.sampler_name == "dpmpp_2m_sde_gpu" else _dpmpp_3m_sde_step
+            for offset, attr_name in enumerate(self.attr_names):
+                scheduler = getattr(self.pipe, attr_name)
+                scheduler._dpmpp_sde_old_denoised = None
+                scheduler._dpmpp_sde_old_denoised_2 = None
+                scheduler._dpmpp_sde_h_last = None
+                scheduler._dpmpp_sde_h_last_2 = None
+                scheduler._dpmpp_sde_noise_sampler = None
+                scheduler._dpmpp_sde_seed = self.seed + offset
+
+                def stepped(model_output, timestep, sample, return_dict=True, _s=scheduler, _f=step_fn, **_kwargs):
+                    return (_f(_s, model_output, timestep, sample),)
+
                 scheduler.step = stepped
         return self