Update h3_dpmpp_2s_ancestral.py

h3_dpmpp_2s_ancestral.py

@@ -1,21 +1,25 @@
-"""A `dpmpp_2s_ancestral` denoise block for MiniMax-H3, replacing `MiniMaxH3LoopDenoiser` + `MiniMaxH3LoopSchedulerStep`.
-
-`dpmpp_2s_ancestral` is a two-evaluation-per-step sampler: it runs the real forward pass exactly as
-`MiniMaxH3LoopDenoiser` does, builds an intermediate sample partway to the next noise level, runs a *second*
-forward pass against it, and combines both predictions (plus fresh ancestral noise) into the step's real output.
-`scheduler.step()` never sees the second call — nothing routed through it could reach `block_state.latents` (full,
-conditioning rows included), `row_timestep_plan`, or the indices `build_row_timesteps` needs. So this is one
-combined `ModularPipelineBlocks`, declared with those as real inputs, standing in for both original sub-blocks at
-once — not a `scheduler.step()` monkeypatch like `euler_ancestral`/`er_sde`/the `dpmpp_*m_sde` pair.
-
-Ports k-diffusion's `sample_dpmpp_2s_ancestral_RF` (the flow-matching branch `sample_dpmpp_2s_ancestral` dispatches
-to for `CONST`-style model sampling — MiniMax-H3's `[0, 1]` sigma space, same family `_euler_ancestral_step` and
-`_er_sde_step` already port) onto the video and audio streams independently, each in one packed transformer call
-per evaluation, since one forward always serves every modality's rows at once.
-
-Wired in via `use_dpmpp_2s_ancestral`, a context manager that temporarily swaps `pipe`'s live `denoise` sub-block
-for the request's duration and restores the original after — the same enter/exit shape `use_schedule` already uses
-for `set_timesteps`/`step`, just swapping a block reference instead.
+"""Denoise blocks for MiniMax-H3's two-evaluation-per-step samplers, replacing `MiniMaxH3LoopDenoiser` +
+`MiniMaxH3LoopSchedulerStep`: `dpmpp_2s_ancestral`, `dpmpp_sde_gpu`, and `seeds_2`.
+
+Each of these runs the real forward pass exactly as `MiniMaxH3LoopDenoiser` does, builds an intermediate sample
+partway to the next noise level, runs a *second* forward pass against it, and combines both predictions (plus
+fresh noise, ancestral or Brownian-tree depending on the sampler) into the step's real output. `scheduler.step()`
+never sees the second call — nothing routed through it could reach `block_state.latents` (full, conditioning rows
+included), `row_timestep_plan`, or the indices `build_row_timesteps` needs. So each is one combined
+`ModularPipelineBlocks`, declared with those as real inputs, standing in for both original sub-blocks at once — not
+a `scheduler.step()` monkeypatch like `euler_ancestral`/`er_sde`/the `dpmpp_*m_sde` pair.
+
+Ports, one function pair per sampler: k-diffusion's `sample_dpmpp_2s_ancestral_RF` (the flow-matching branch
+`sample_dpmpp_2s_ancestral` dispatches to for `CONST`-style model sampling — MiniMax-H3's `[0, 1]` sigma space,
+same family `_euler_ancestral_step`/`_er_sde_step` already port); `sample_dpmpp_sde`/`sample_dpmpp_sde_gpu`, written
+generically against `CONST`'s half-log-SNR specialization rather than a hand-derived `_RF` branch, reused here via
+`_lambda_const`/`_sigma_from_lambda_const`; and `sample_seeds_2`'s `phi_1` path, the same half-log-SNR family.
+Each runs the video and audio streams independently, each still one packed transformer call per evaluation, since
+one forward always serves every modality's rows at once.
+
+Wired in via three context managers (`use_dpmpp_2s_ancestral`, `use_dpmpp_sde_gpu`, `use_seeds_2`) that temporarily
+swap `pipe`'s live `denoise` sub-block for the request's duration and restore the original after — the same
+enter/exit shape `use_schedule` already uses for `set_timesteps`/`step`, just swapping a block reference instead.
 """
 
 import inspect
@@ -30,6 +34,8 @@ from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, In
 from diffusers.models import MiniMaxH3Transformer3DModel
 from diffusers.schedulers import MiniMaxH3Scheduler
 
+from pk_workflow import _BrownianTreeNoiseSampler
+
 
 def _dpmpp_2s_ancestral_prepare(x: torch.Tensor, sigma_i: torch.Tensor, sigma_ip1: torch.Tensor, denoised: torch.Tensor, eta: float = 1.0):
     """The half of one stream's `dpmpp_2s_ancestral_RF` step computable before the second model call: the
@@ -227,6 +233,334 @@ class MiniMaxH3DPMpp2SAncestralStep(ModularPipelineBlocks):
         return components, block_state
 
 
+def _nudge_first_sigma(sigma: torch.Tensor, shift: float) -> torch.Tensor:
+    """ComfyUI's `offset_first_sigma_for_snr`, `CONST` branch, for one already-known-to-be-`>= 1.0` sigma: nudges
+    it to `percent_to_sigma(1e-4)` — the same `0.9999`-below-1.0 value `_er_sde_step`/`_dpmpp_2m_sde_step` already
+    nudge to, re-derived here through the shift formula since `dpmpp_sde_gpu`/`seeds_2` need it as a tensor
+    they can keep differentiating through rather than a hardcoded constant.
+    """
+    base = sigma.new_tensor(1.0 - 1e-4)
+    return shift * base / (1 + (shift - 1) * base)
+
+
+def _lambda_const(sigma: torch.Tensor) -> torch.Tensor:
+    """`sigma_to_half_log_snr`'s `CONST` branch: `log((1 - sigma) / sigma)`."""
+    return ((1 - sigma) / sigma).log()
+
+
+def _sigma_from_lambda_const(half_log_snr: torch.Tensor) -> torch.Tensor:
+    """`half_log_snr_to_sigma`'s `CONST` branch: `1 / (1 + exp(half_log_snr))`."""
+    return (half_log_snr.exp() + 1) ** -1
+
+
+def _get_ancestral_step(sigma_from: torch.Tensor, sigma_to: torch.Tensor, eta: float = 1.0):
+    """k-diffusion's `get_ancestral_step`, unmodified — operates on whatever space the caller hands it (each
+    sampler below calls it in `exp(-half_log_snr)` space, per its own source)."""
+    sigma_up = torch.minimum(sigma_to, eta * (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2).clamp_min(0).sqrt())
+    sigma_down = (sigma_to**2 - sigma_up**2).clamp_min(0).sqrt()
+    return sigma_down, sigma_up
+
+
+def _dpmpp_sde_gpu_prepare(x: torch.Tensor, sigma_i: torch.Tensor, sigma_ip1: torch.Tensor, denoised: torch.Tensor, shift: float, is_first_step: bool, noise_sampler):
+    """Step 1 of `sample_dpmpp_sde`'s `CONST` (half-log-SNR) path for one stream: the intermediate sigma/sample
+    to evaluate (`sigma_s_1`, `u`) plus what `_dpmpp_sde_gpu_combine` needs once the second call's `denoised_2`
+    is in hand. `eta = 1.0`, `r = 0.5` — the only values this port exposes, matching ComfyUI's own `dpmpp_sde_gpu`
+    registration. `noise_sampler` is the request-lifetime, per-stream `_BrownianTreeNoiseSampler` — both of a
+    step's noise draws start from the same `sigma_i`, which is what correlates them.
+    """
+    sigma_i_dpm = _nudge_first_sigma(sigma_i, shift) if is_first_step and float(sigma_i) >= 1.0 else sigma_i
+    lambda_s = _lambda_const(sigma_i_dpm)
+    lambda_t = _lambda_const(sigma_ip1)
+    h = lambda_t - lambda_s
+    lambda_s_1 = lambda_s + 0.5 * h
+    sigma_s_1 = _sigma_from_lambda_const(lambda_s_1)
+    alpha_s, alpha_s_1 = 1 - sigma_i_dpm, 1 - sigma_s_1
+
+    w_i, w_s1 = (-lambda_s).exp(), (-lambda_s_1).exp()
+    sd, su = _get_ancestral_step(w_i, w_s1)
+    h_ = -sd.log() - lambda_s
+    u = (alpha_s_1 / alpha_s) * (-h_).exp() * x - alpha_s_1 * torch.expm1(-h_) * denoised
+    u = u + alpha_s_1 * noise_sampler(sigma_i_dpm, sigma_s_1) * su
+    return u, sigma_s_1, sigma_i_dpm, lambda_s
+
+
+def _dpmpp_sde_gpu_combine(noise_sampler, x: torch.Tensor, sigma_i_dpm: torch.Tensor, lambda_s: torch.Tensor, sigma_ip1: torch.Tensor, denoised_2: torch.Tensor):
+    """Step 2, once the second call's `denoised_2` is in hand. The `fac = 1` collapse at `r = 0.5` means the
+    combine uses `denoised_2` alone — the same simplification `dpmpp_2s_ancestral`'s port already relies on."""
+    lambda_t = _lambda_const(sigma_ip1)
+    alpha_s, alpha_t = 1 - sigma_i_dpm, 1 - sigma_ip1
+    w_i, w_t = (-lambda_s).exp(), (-lambda_t).exp()
+    sd, su = _get_ancestral_step(w_i, w_t)
+    h_ = -sd.log() - lambda_s
+    x = (alpha_t / alpha_s) * (-h_).exp() * x - alpha_t * torch.expm1(-h_) * denoised_2
+    x = x + alpha_t * noise_sampler(sigma_i_dpm, sigma_ip1) * su
+    return x
+
+
+def _seeds2_prepare(x: torch.Tensor, sigma_i: torch.Tensor, sigma_ip1: torch.Tensor, denoised: torch.Tensor, shift: float, is_first_step: bool, generator):
+    """Step 1 of `sample_seeds_2`'s `phi_1` path (its registered default) for one stream. Unlike `dpmpp_sde_gpu`,
+    `seeds_2` draws fresh, independent Gaussian noise per call rather than a correlated Brownian-tree draw —
+    `default_noise_sampler` in the k-diffusion source ignores its sigma arguments entirely, so `generator` here is
+    just the request-lifetime per-stream `torch.Generator`, same pattern `euler_ancestral`/`er_sde` already use.
+    """
+    sigma_i_dpm = _nudge_first_sigma(sigma_i, shift) if is_first_step and float(sigma_i) >= 1.0 else sigma_i
+    lambda_s = _lambda_const(sigma_i_dpm)
+    lambda_t = _lambda_const(sigma_ip1)
+    h = lambda_t - lambda_s
+    lambda_s_1 = lambda_s + 0.5 * h
+    sigma_s_1 = _sigma_from_lambda_const(lambda_s_1)
+    alpha_s_1 = 1 - sigma_s_1
+
+    u = (sigma_s_1 / sigma_i_dpm) * (-0.5 * h).exp() * x - alpha_s_1 * torch.expm1(-h) * denoised
+    noise_1 = torch.randn(x.shape, dtype=x.dtype, device="cpu", generator=generator).to(x.device)
+    sde_noise = torch.sqrt(-torch.expm1(-h)) * noise_1
+    u = u + sde_noise * sigma_s_1
+    return u, sigma_s_1, sde_noise, h, sigma_i_dpm
+
+
+def _seeds2_combine(generator, x: torch.Tensor, sigma_i_dpm: torch.Tensor, sigma_ip1: torch.Tensor, denoised_2: torch.Tensor, sde_noise: torch.Tensor, h: torch.Tensor):
+    """Step 2 of `sample_seeds_2`'s `phi_1` path, once `denoised_2` is in hand — including the second, independent
+    noise draw the source layers on top of the first (`sde_noise` carries the first draw's contribution forward,
+    scaled, rather than being discarded)."""
+    alpha_t = 1 - sigma_ip1
+    x = (sigma_ip1 / sigma_i_dpm) * (-h).exp() * x - alpha_t * torch.expm1(-2 * h) * denoised_2
+
+    segment_factor = -0.5 * h  # (r - 1) * h * eta, r = 0.5, eta = 1.0
+    sde_noise = sde_noise * segment_factor.exp()
+    noise_2 = torch.randn(x.shape, dtype=x.dtype, device="cpu", generator=generator).to(x.device)
+    sde_noise = sde_noise + torch.sqrt(-torch.expm1(2 * segment_factor)) * noise_2
+    x = x + sde_noise * sigma_ip1
+    return x
+
+
+class MiniMaxH3DPMppSdeStep(ModularPipelineBlocks):
+    """One `dpmpp_sde_gpu` iteration, same role as `MiniMaxH3DPMpp2SAncestralStep`: real forward pass, intermediate
+    forward pass, DPM-Solver++(SDE) combine with Brownian-tree-correlated noise."""
+
+    model_name = "minimax-h3"
+
+    def __init__(self, transformer_name: str = "transformer", seed: int = 0):
+        self.transformer_name = transformer_name
+        self.seed = int(seed)
+        self._video_noise_sampler = None
+        self._audio_noise_sampler = None
+        super().__init__()
+
+    @property
+    def description(self) -> str:
+        return (
+            "Runs one `dpmpp_sde_gpu` denoising iteration in place of `MiniMaxH3LoopDenoiser` + "
+            "`MiniMaxH3LoopSchedulerStep`: the real forward pass, a second forward pass at a constructed "
+            "intermediate sample/sigma, and the DPM-Solver++(SDE) combine with Brownian-tree noise."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(self.transformer_name, MiniMaxH3Transformer3DModel),
+            ComponentSpec("scheduler", MiniMaxH3Scheduler),
+            ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+        ]
+
+    inputs = MiniMaxH3DPMpp2SAncestralStep.inputs  # same input contract as the dpmpp_2s_ancestral block
+
+    @property
+    def intermediate_outputs(self) -> list:
+        return []
+
+    _forward = MiniMaxH3DPMpp2SAncestralStep._forward
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        ncv = block_state.num_condition_video_rows
+        nca = block_state.num_condition_audio_rows
+        is_last_step = i == len(block_state.row_timestep_plan) - 1
+        is_first_step = i == 0
+
+        unique_timesteps, timestep_indices = block_state.row_timestep_plan[i]
+        noise_pred, audio_noise_pred = self._forward(
+            components, block_state, block_state.latents, block_state.audio_latents, unique_timesteps, timestep_indices
+        )
+
+        video_timestep_i = float(t)
+
+        x_video = block_state.latents[ncv:]
+        x_audio = block_state.audio_latents[nca:]
+        compute_dtype = torch.float32 if x_video.dtype in (torch.float16, torch.bfloat16) else x_video.dtype
+
+        sigma_i_video = components.scheduler.sigmas[i].to(device=x_video.device, dtype=compute_dtype)
+        sigma_ip1_video = components.scheduler.sigmas[i + 1].to(device=x_video.device, dtype=compute_dtype)
+        sigma_i_audio = components.audio_scheduler.sigmas[i].to(device=x_audio.device, dtype=compute_dtype)
+        sigma_ip1_audio = components.audio_scheduler.sigmas[i + 1].to(device=x_audio.device, dtype=compute_dtype)
+
+        x_video_c = x_video.to(dtype=compute_dtype)
+        x_audio_c = x_audio.to(dtype=compute_dtype)
+        denoised_video = x_video_c + sigma_i_video * noise_pred[0, ncv:].to(dtype=compute_dtype)
+        denoised_audio = x_audio_c + sigma_i_audio * audio_noise_pred[0, nca:].to(dtype=compute_dtype)
+
+        if is_last_step:
+            new_video, new_audio = denoised_video, denoised_audio
+        else:
+            if self._video_noise_sampler is None:
+                # Lazily built on the first step, same lifetime as the request (this block instance is discarded
+                # after) — the real, un-nudged schedule bounds are the tree's span; per-step nudging happens
+                # only in the query sigmas passed to it, same convention `_dpmpp_2m_sde_step` already uses.
+                video_sigmas = components.scheduler.sigmas.to(device=x_video.device, dtype=compute_dtype)
+                audio_sigmas = components.audio_scheduler.sigmas.to(device=x_audio.device, dtype=compute_dtype)
+                self._video_noise_sampler = _BrownianTreeNoiseSampler(
+                    x_video_c, video_sigmas[video_sigmas > 0].min(), video_sigmas.max(), seed=self.seed, cpu=False
+                )
+                self._audio_noise_sampler = _BrownianTreeNoiseSampler(
+                    x_audio_c, audio_sigmas[audio_sigmas > 0].min(), audio_sigmas.max(), seed=self.seed + 1, cpu=False
+                )
+
+            u_video, sigma_s_video, sigma_i_dpm_video, lambda_s_video = _dpmpp_sde_gpu_prepare(
+                x_video_c, sigma_i_video, sigma_ip1_video, denoised_video, float(components.scheduler.shift), is_first_step, self._video_noise_sampler
+            )
+            u_audio, sigma_s_audio, sigma_i_dpm_audio, lambda_s_audio = _dpmpp_sde_gpu_prepare(
+                x_audio_c, sigma_i_audio, sigma_ip1_audio, denoised_audio, float(components.audio_scheduler.shift), is_first_step, self._audio_noise_sampler
+            )
+
+            latents_2 = block_state.latents.clone()
+            latents_2[ncv:] = u_video.to(dtype=block_state.latents.dtype)
+            audio_latents_2 = block_state.audio_latents.clone()
+            audio_latents_2[nca:] = u_audio.to(dtype=block_state.audio_latents.dtype)
+
+            unique_timesteps_2, timestep_indices_2 = tuple(
+                tensor.to(block_state.latents.device)
+                for tensor in MiniMaxH3SetTimestepsStep.build_row_timesteps(
+                    block_state.video_indices,
+                    block_state.audio_indices,
+                    ncv,
+                    nca,
+                    block_state.text_indices.numel(),
+                    1.0 - float(sigma_s_video),
+                    1.0 - float(sigma_s_audio),
+                    max(video_timestep_i, components.keyframe_noise_aug),
+                    1.0,
+                )
+            )
+            noise_pred_2, audio_noise_pred_2 = self._forward(
+                components, block_state, latents_2, audio_latents_2, unique_timesteps_2, timestep_indices_2
+            )
+            d_i_video = u_video + sigma_s_video * noise_pred_2[0, ncv:].to(dtype=compute_dtype)
+            d_i_audio = u_audio + sigma_s_audio * audio_noise_pred_2[0, nca:].to(dtype=compute_dtype)
+
+            new_video = _dpmpp_sde_gpu_combine(self._video_noise_sampler, x_video_c, sigma_i_dpm_video, lambda_s_video, sigma_ip1_video, d_i_video)
+            new_audio = _dpmpp_sde_gpu_combine(self._audio_noise_sampler, x_audio_c, sigma_i_dpm_audio, lambda_s_audio, sigma_ip1_audio, d_i_audio)
+
+        block_state.latents[ncv:] = new_video.to(dtype=block_state.latents.dtype)
+        block_state.audio_latents[nca:] = new_audio.to(dtype=block_state.audio_latents.dtype)
+        return components, block_state
+
+
+class MiniMaxH3Seeds2Step(ModularPipelineBlocks):
+    """One `seeds_2` iteration, same role as `MiniMaxH3DPMpp2SAncestralStep`: real forward pass, intermediate
+    forward pass, SEEDS-2 (`phi_1`) combine with independent Gaussian noise."""
+
+    model_name = "minimax-h3"
+
+    def __init__(self, transformer_name: str = "transformer", video_generator=None, audio_generator=None):
+        self.transformer_name = transformer_name
+        self.video_generator = video_generator
+        self.audio_generator = audio_generator
+        super().__init__()
+
+    @property
+    def description(self) -> str:
+        return (
+            "Runs one `seeds_2` denoising iteration in place of `MiniMaxH3LoopDenoiser` + "
+            "`MiniMaxH3LoopSchedulerStep`: the real forward pass, a second forward pass at a constructed "
+            "intermediate sample/sigma, and the SEEDS-2 (`phi_1`) combine with independent Gaussian noise."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(self.transformer_name, MiniMaxH3Transformer3DModel),
+            ComponentSpec("scheduler", MiniMaxH3Scheduler),
+            ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+        ]
+
+    inputs = MiniMaxH3DPMpp2SAncestralStep.inputs  # same input contract as the dpmpp_2s_ancestral block
+
+    @property
+    def intermediate_outputs(self) -> list:
+        return []
+
+    _forward = MiniMaxH3DPMpp2SAncestralStep._forward
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        ncv = block_state.num_condition_video_rows
+        nca = block_state.num_condition_audio_rows
+        is_last_step = i == len(block_state.row_timestep_plan) - 1
+        is_first_step = i == 0
+
+        unique_timesteps, timestep_indices = block_state.row_timestep_plan[i]
+        noise_pred, audio_noise_pred = self._forward(
+            components, block_state, block_state.latents, block_state.audio_latents, unique_timesteps, timestep_indices
+        )
+
+        video_timestep_i = float(t)
+
+        x_video = block_state.latents[ncv:]
+        x_audio = block_state.audio_latents[nca:]
+        compute_dtype = torch.float32 if x_video.dtype in (torch.float16, torch.bfloat16) else x_video.dtype
+
+        sigma_i_video = components.scheduler.sigmas[i].to(device=x_video.device, dtype=compute_dtype)
+        sigma_ip1_video = components.scheduler.sigmas[i + 1].to(device=x_video.device, dtype=compute_dtype)
+        sigma_i_audio = components.audio_scheduler.sigmas[i].to(device=x_audio.device, dtype=compute_dtype)
+        sigma_ip1_audio = components.audio_scheduler.sigmas[i + 1].to(device=x_audio.device, dtype=compute_dtype)
+
+        x_video_c = x_video.to(dtype=compute_dtype)
+        x_audio_c = x_audio.to(dtype=compute_dtype)
+        denoised_video = x_video_c + sigma_i_video * noise_pred[0, ncv:].to(dtype=compute_dtype)
+        denoised_audio = x_audio_c + sigma_i_audio * audio_noise_pred[0, nca:].to(dtype=compute_dtype)
+
+        if is_last_step:
+            new_video, new_audio = denoised_video, denoised_audio
+        else:
+            u_video, sigma_s_video, sde_noise_video, h_video, sigma_i_dpm_video = _seeds2_prepare(
+                x_video_c, sigma_i_video, sigma_ip1_video, denoised_video, float(components.scheduler.shift), is_first_step, self.video_generator
+            )
+            u_audio, sigma_s_audio, sde_noise_audio, h_audio, sigma_i_dpm_audio = _seeds2_prepare(
+                x_audio_c, sigma_i_audio, sigma_ip1_audio, denoised_audio, float(components.audio_scheduler.shift), is_first_step, self.audio_generator
+            )
+
+            latents_2 = block_state.latents.clone()
+            latents_2[ncv:] = u_video.to(dtype=block_state.latents.dtype)
+            audio_latents_2 = block_state.audio_latents.clone()
+            audio_latents_2[nca:] = u_audio.to(dtype=block_state.audio_latents.dtype)
+
+            unique_timesteps_2, timestep_indices_2 = tuple(
+                tensor.to(block_state.latents.device)
+                for tensor in MiniMaxH3SetTimestepsStep.build_row_timesteps(
+                    block_state.video_indices,
+                    block_state.audio_indices,
+                    ncv,
+                    nca,
+                    block_state.text_indices.numel(),
+                    1.0 - float(sigma_s_video),
+                    1.0 - float(sigma_s_audio),
+                    max(video_timestep_i, components.keyframe_noise_aug),
+                    1.0,
+                )
+            )
+            noise_pred_2, audio_noise_pred_2 = self._forward(
+                components, block_state, latents_2, audio_latents_2, unique_timesteps_2, timestep_indices_2
+            )
+            d_i_video = u_video + sigma_s_video * noise_pred_2[0, ncv:].to(dtype=compute_dtype)
+            d_i_audio = u_audio + sigma_s_audio * audio_noise_pred_2[0, nca:].to(dtype=compute_dtype)
+
+            new_video = _seeds2_combine(self.video_generator, x_video_c, sigma_i_dpm_video, sigma_ip1_video, d_i_video, sde_noise_video, h_video)
+            new_audio = _seeds2_combine(self.audio_generator, x_audio_c, sigma_i_dpm_audio, sigma_ip1_audio, d_i_audio, sde_noise_audio, h_audio)
+
+        block_state.latents[ncv:] = new_video.to(dtype=block_state.latents.dtype)
+        block_state.audio_latents[nca:] = new_audio.to(dtype=block_state.audio_latents.dtype)
+        return components, block_state
+
+
 class use_dpmpp_2s_ancestral:
     """Swaps `pipe`'s live `denoise` sub-block for `MiniMaxH3DPMpp2SAncestralStep` for one request, restoring the
     original block on exit. A block swap rather than a `scheduler.step()` patch, and so a separate context
@@ -265,3 +599,35 @@ class use_dpmpp_2s_ancestral:
         if self._core_denoise is not None and self._original is not None:
             self._core_denoise.sub_blocks["denoise"] = self._original
         return False
+
+
+class use_dpmpp_sde_gpu(use_dpmpp_2s_ancestral):
+    """Same swap as `use_dpmpp_2s_ancestral`, for `MiniMaxH3DPMppSdeStep` instead. Its own subclass rather than a
+    shared parameterized manager, since `MiniMaxH3DPMppSdeStep` takes `seed` directly (its noise samplers are
+    lazily built from `x`'s shape/device, unlike the ancestral samplers' plain `torch.Generator`)."""
+
+    def __enter__(self):
+        if not self.enabled:
+            return self
+        self._core_denoise = self.pipe._blocks.sub_blocks["denoise"]
+        self._original = self._core_denoise.sub_blocks["denoise"]
+        step = MiniMaxH3DPMppSdeStep(transformer_name=self.transformer_name, seed=self.seed)
+        self._core_denoise.sub_blocks["denoise"] = MiniMaxH3DenoiseLoopWrapper.from_blocks_dict({"dpmpp_sde_gpu": step})
+        return self
+
+
+class use_seeds_2(use_dpmpp_2s_ancestral):
+    """Same swap as `use_dpmpp_2s_ancestral`, for `MiniMaxH3Seeds2Step` instead."""
+
+    def __enter__(self):
+        if not self.enabled:
+            return self
+        self._core_denoise = self.pipe._blocks.sub_blocks["denoise"]
+        self._original = self._core_denoise.sub_blocks["denoise"]
+        video_generator = torch.Generator(device="cpu").manual_seed(self.seed)
+        audio_generator = torch.Generator(device="cpu").manual_seed(self.seed + 1)
+        step = MiniMaxH3Seeds2Step(
+            transformer_name=self.transformer_name, video_generator=video_generator, audio_generator=audio_generator
+        )
+        self._core_denoise.sub_blocks["denoise"] = MiniMaxH3DenoiseLoopWrapper.from_blocks_dict({"seeds_2": step})
+        return self