Upload h3_dpmpp_2s_ancestral.py

h3_dpmpp_2s_ancestral.py

@@ -0,0 +1,267 @@
+"""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.
+"""
+
+import inspect
+
+import torch
+
+from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3SetTimestepsStep
+from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3DenoiseLoopWrapper
+from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
+from diffusers.modular_pipelines.modular_pipeline import BlockState, ModularPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam
+from diffusers.models import MiniMaxH3Transformer3DModel
+from diffusers.schedulers import MiniMaxH3Scheduler
+
+
+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
+    intermediate sigma/sample to evaluate (`sigma_s`, `u`) and the scalars the combine step needs once the second
+    call's `D_i` comes back. `x`/`denoised` are this stream's generated rows only (conditioning rows excluded,
+    same slice `MiniMaxH3LoopSchedulerStep` already steps). Mirrors k-diffusion's `sample_dpmpp_2s_ancestral_RF`
+    body exactly, split at the point the real function makes its second `model(...)` call.
+    """
+    downstep_ratio = 1 + (sigma_ip1 / sigma_i - 1) * eta
+    sigma_down = sigma_ip1 * downstep_ratio
+    alpha_ip1 = 1 - sigma_ip1
+    alpha_down = 1 - sigma_down
+    renoise_coeff = (sigma_ip1**2 - sigma_down**2 * alpha_ip1**2 / alpha_down**2).clamp_min(0).sqrt()
+
+    if float(sigma_i) >= 1.0:
+        # k-diffusion's own first-step special case: `lambda_fn` divides by `sigma_i` inside a `log`, which blows
+        # up at exactly 1.0 (MiniMax-H3's real first sigma). Same nudge `_er_sde_step` uses for its own division.
+        sigma_s = sigma_i.new_tensor(0.9999)
+    else:
+        lambda_fn = lambda sigma: ((1 - sigma) / sigma).log()
+        t_i, t_down = lambda_fn(sigma_i), lambda_fn(sigma_down)
+        h = t_down - t_i
+        s = t_i + 0.5 * h
+        sigma_s = (s.exp() + 1) ** -1
+
+    sigma_s_i_ratio = sigma_s / sigma_i
+    u = sigma_s_i_ratio * x + (1 - sigma_s_i_ratio) * denoised
+    return u, sigma_s, sigma_down, alpha_ip1, alpha_down, renoise_coeff
+
+
+def _dpmpp_2s_ancestral_combine(
+    generator, x: torch.Tensor, sigma_ip1: torch.Tensor, d_i: torch.Tensor, sigma_i: torch.Tensor,
+    sigma_down: torch.Tensor, alpha_ip1: torch.Tensor, alpha_down: torch.Tensor, renoise_coeff: torch.Tensor,
+    eta: float = 1.0, s_noise: float = 1.0,
+):
+    """The DPM-Solver++(2S) combine and ancestral-noise-injection half, once the second call's `d_i` (already
+    converted to a denoised/x0 estimate, same as the first call) is in hand."""
+    sigma_down_i_ratio = sigma_down / sigma_i
+    x = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * d_i
+    if float(sigma_ip1) > 0 and eta > 0:
+        noise = torch.randn(x.shape, dtype=x.dtype, device=x.device, generator=generator)
+        x = (alpha_ip1 / alpha_down) * x + noise * s_noise * renoise_coeff
+    return x
+
+
+class MiniMaxH3DPMpp2SAncestralStep(ModularPipelineBlocks):
+    """One `dpmpp_2s_ancestral` iteration: real forward pass, intermediate forward pass, DPM-Solver++(2S) combine
+    with ancestral noise — video and audio computed independently, each still one packed transformer call.
+    """
+
+    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 `dpmpp_2s_ancestral` 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++(2S) combine with ancestral noise injection."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(self.transformer_name, MiniMaxH3Transformer3DModel),
+            ComponentSpec("scheduler", MiniMaxH3Scheduler),
+            ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="latents", type_hint=torch.Tensor, required=True),
+            InputParam(name="audio_latents", type_hint=torch.Tensor, required=True),
+            InputParam.template("prompt_embeds"),
+            InputParam(name="row_timestep_plan", type_hint=list, required=True),
+            InputParam(name="audio_timesteps", type_hint=torch.Tensor, required=True),
+            InputParam(name="video_indices", type_hint=torch.Tensor, required=True),
+            InputParam(name="audio_indices", type_hint=torch.Tensor, required=True),
+            InputParam(name="text_indices", type_hint=torch.Tensor, required=True),
+            InputParam(name="num_condition_video_rows", type_hint=int, default=0),
+            InputParam(name="num_condition_audio_rows", type_hint=int, default=0),
+            InputParam(kwargs_type="denoiser_input_fields"),
+            InputParam.template("attention_kwargs"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list:
+        return []  # writes block_state.latents / audio_latents in place, same as MiniMaxH3LoopSchedulerStep
+
+    def _forward(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, hidden_states, audio_hidden_states, unique_timesteps, timestep_indices):
+        transformer = getattr(components, self.transformer_name)
+        layout_kwargs = {
+            name: value
+            for name, value in block_state.denoiser_input_fields.items()
+            if name in inspect.signature(transformer.forward).parameters
+        }
+        return transformer(
+            hidden_states=hidden_states[None],
+            audio_hidden_states=audio_hidden_states[None],
+            encoder_hidden_states=block_state.prompt_embeds,
+            timestep=unique_timesteps,
+            timestep_indices=timestep_indices,
+            attention_kwargs=block_state.attention_kwargs,
+            return_dict=False,
+            **layout_kwargs,
+        )
+
+    @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
+
+        # --- first (real) evaluation, identical to MiniMaxH3LoopDenoiser ---
+        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)
+        # velocity -> denoised (x0 estimate), same conversion `_euler_ancestral_step`/`_er_sde_step` already use.
+        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:
+            # k-diffusion's own Euler fallback (`sigmas[i+1] == 0`) collapses to `x = denoised` at eta=1, the only
+            # eta this port exposes — skip the second call entirely rather than compute a no-op intermediate step.
+            new_video, new_audio = denoised_video, denoised_audio
+        else:
+            u_video, sigma_s_video, sigma_down_video, alpha_ip1_video, alpha_down_video, renoise_video = (
+                _dpmpp_2s_ancestral_prepare(x_video_c, sigma_i_video, sigma_ip1_video, denoised_video)
+            )
+            u_audio, sigma_s_audio, sigma_down_audio, alpha_ip1_audio, alpha_down_audio, renoise_audio = (
+                _dpmpp_2s_ancestral_prepare(x_audio_c, sigma_i_audio, sigma_ip1_audio, denoised_audio)
+            )
+
+            # --- second (intermediate) evaluation: full sequence, only the generated rows replaced by u ---
+            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)
+
+            # Conditioning rows stay pinned exactly as the real step pinned them (not re-derived from the
+            # intermediate sigma) — the fixed floor `MiniMaxH3SetTimestepsStep.__call__` itself uses. Text rows
+            # inherit whichever video timestep is passed, here the intermediate one, same rule as the real step.
+            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_2s_ancestral_combine(
+                self.video_generator, x_video_c, sigma_ip1_video, d_i_video, sigma_i_video,
+                sigma_down_video, alpha_ip1_video, alpha_down_video, renoise_video,
+            )
+            new_audio = _dpmpp_2s_ancestral_combine(
+                self.audio_generator, x_audio_c, sigma_ip1_audio, d_i_audio, sigma_i_audio,
+                sigma_down_audio, alpha_ip1_audio, alpha_down_audio, renoise_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
+    manager from `use_schedule` — nest them (`use_schedule` still owns the sigma grid; this only owns which block
+    walks it). No-op when `enabled=False`, so it's safe to wrap every request unconditionally.
+
+    Reaches the running block tree at `pipe._blocks.sub_blocks["denoise"].sub_blocks["denoise"]`:
+    `MiniMaxH3GeneratorBlocks`'s own "denoise" step is `MiniMaxH3CoreDenoiseStep`, whose own "denoise" step is
+    `MiniMaxH3DenoiseStep` — the `LoopSequentialPipelineBlocks` this class temporarily replaces outright.
+    """
+
+    def __init__(self, pipe, seed: int = 0, transformer_name: str = "transformer", enabled: bool = True):
+        self.pipe = pipe
+        self.seed = int(seed)
+        self.transformer_name = transformer_name
+        self.enabled = enabled
+        self._core_denoise = None
+        self._original = None
+
+    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 = MiniMaxH3DPMpp2SAncestralStep(
+            transformer_name=self.transformer_name, video_generator=video_generator, audio_generator=audio_generator
+        )
+        self._core_denoise.sub_blocks["denoise"] = MiniMaxH3DenoiseLoopWrapper.from_blocks_dict(
+            {"dpmpp_2s_ancestral": step}
+        )
+        return self
+
+    def __exit__(self, *_):
+        if self._core_denoise is not None and self._original is not None:
+            self._core_denoise.sub_blocks["denoise"] = self._original
+        return False