Upload h3_momentum.py

h3_momentum.py

@@ -0,0 +1,341 @@
+"""Video-momentum-conditioned `t2va` for Chunked Generation: the previous chunk's trailing frames, imposed on the
+opening latent frames of the next chunk's own target — real motion continuity at the seam, not a single-frame
+keyframe anchor.
+
+Modeled directly on `h3_a2v_blocks.py`'s `MiniMaxH3AudioConditionStep` (from the `multimodalart/minimax-h3-audio-
+to-video` Space): a *given* signal is encoded and imposed on the rows the model would otherwise draw fresh noise
+for, in one of two ways — `locked` (clean from the first forward on, presented at `t = 1.0`) or `blended` (re-
+noised to each step's own sigma, ordinary diffusion inpainting). The difference from that file's audio block is
+scope: audio conditions the *entire* track; this conditions only a *prefix* of the generated video rows — the
+carried head — leaving the rest on the normal free noise-draw path, since a chunk's own new content still needs
+to come from the model, not from what was carried in.
+
+Continuation chunks run keyframe-free (`t2va`-shaped, no `image`/`last_image`): the momentum block already
+determines what the opening frames are, more directly than a keyframe ever could, so there is nothing for a
+keyframe to add. Only chunk one, with no prior chunk to carry from, uses the ordinary keyframe-capable generator.
+
+Two things this has not been validated against yet, flagged rather than assumed away: whether encoding a short
+trailing clip standalone reproduces what those frames would have encoded to as part of the original longer clip
+(a causal-VAE boundary effect, if the encoder has significant temporal receptive field), and whether the video
+VAE's encoder is tolerant of an arbitrary frame count for a short sub-clip the way a full request's `17 * n + 5`
+alignment matters there. Worth checking directly before trusting this at high momentum durations.
+"""
+
+import torch
+
+from diffusers.models import AutoencoderKLMiniMaxH3
+from diffusers.modular_pipelines.minimax_h3.before_denoise import (
+    MiniMaxH3NoKeyframeAnchorsStep,
+    MiniMaxH3PrepareLatentsStep,
+    MiniMaxH3PrepareLayoutStep,
+    MiniMaxH3SetTimestepsStep,
+)
+from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+from diffusers.modular_pipelines.minimax_h3.denoise import (
+    MiniMaxH3DenoiseLoopWrapper,
+    MiniMaxH3LoopDenoiser,
+    MiniMaxH3LoopSchedulerStep,
+)
+from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import MiniMaxH3DecodeStep
+from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
+from diffusers.modular_pipelines.modular_pipeline import BlockState, ModularPipelineBlocks, PipelineState, SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+LOCKED, BLENDED, OFF = "locked", "blended", "off"
+VIDEO_CONDITION_MODES = (LOCKED, BLENDED, OFF)
+
+
+def pack_video_rows(latents: torch.Tensor, patch_size: tuple[int, int, int], channels: int) -> torch.Tensor:
+    """The forward patchify `MiniMaxH3PrepareLayoutStep` applies to fresh noise, run here on a real encoded clip
+    instead. The exact inverse of `MiniMaxH3AfterDenoiseStep`'s own unpack (`reshape` -> `permute(0,4,1,5,2,6,3,7)`
+    -> `reshape`), so a prefix of the rows this returns lines up with a whole number of leading latent frames —
+    with `patch_size`'s temporal component at 1 (no temporal sub-patching), each latent frame maps to an exact,
+    contiguous, non-interleaved block of `(latent_height // patch_h) * (latent_width // patch_w)` rows.
+    """
+    patch_t, patch_h, patch_w = patch_size
+    b, c, f, h, w = latents.shape
+    rows = latents.reshape(b, c, f // patch_t, patch_t, h // patch_h, patch_h, w // patch_w, patch_w)
+    rows = rows.permute(0, 2, 4, 6, 1, 3, 5, 7)
+    return rows.reshape(-1, channels * patch_t * patch_h * patch_w)
+
+
+class MiniMaxH3MomentumConditionStep(ModularPipelineBlocks):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Encodes the given trailing clip into the opening video rows of the packed sequence, and keeps the "
+            "noise those rows were drawn from. Runs after the noise draw so the request's generator is "
+            "untouched: video noise is still the generator's first draw regardless of this block, which only "
+            "overwrites — never skips — that draw for the imposed prefix."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("vae", AutoencoderKLMiniMaxH3)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="given_video",
+                type_hint=torch.Tensor,
+                description="The carried clip's pixel frames, `(num_frames, 3, H, W)` in `[0, 1]`. None runs "
+                "plain `t2va`.",
+            ),
+            InputParam(
+                name="video_condition_mode",
+                type_hint=str,
+                default=BLENDED,
+                description="`locked`, `blended` or `off`.",
+            ),
+            InputParam(name="latents", type_hint=torch.Tensor, required=True, description="The video rows of the "
+                       "packed sequence, as drawn from the request's generator."),
+            InputParam(name="num_condition_video_rows", type_hint=int, default=0),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("latents", type_hint=torch.Tensor, description="The video rows the loop starts from."),
+            OutputParam(
+                "given_video_rows", type_hint=torch.Tensor, description="The encoded, packed carried clip's rows, "
+                "or None.",
+            ),
+            OutputParam(
+                "video_noise_rows", type_hint=torch.Tensor, description="The noise the imposed rows were drawn "
+                "from, which `blended` re-noises against.",
+            ),
+            OutputParam(
+                "num_momentum_rows", type_hint=int, description="How many leading generated video rows are "
+                "imposed — 0 when there is nothing to carry.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+
+        if block_state.video_condition_mode not in VIDEO_CONDITION_MODES:
+            raise ValueError(
+                f"`video_condition_mode` must be one of {VIDEO_CONDITION_MODES}, got "
+                f"{block_state.video_condition_mode!r}."
+            )
+
+        block_state.video_noise_rows = block_state.latents.clone()
+        block_state.given_video_rows = None
+        block_state.num_momentum_rows = 0
+
+        if block_state.given_video is not None and block_state.video_condition_mode != OFF:
+            channels = components.vae_latent_channels
+            patch_size = components.patch_size
+
+            # (num_frames, 3, H, W) in [0, 1] -> (1, 3, num_frames, H, W), the video VAE's own input convention.
+            frames = block_state.given_video.to(device=device, dtype=components.vae.dtype)
+            frames = frames.permute(1, 0, 2, 3)[None]
+            posterior = components.vae.encode(frames, return_dict=False)[0]
+            momentum_latents = posterior.mode()  # (1, channels, k, latent_height, latent_width)
+
+            given_rows = pack_video_rows(momentum_latents, patch_size, channels).to(block_state.latents.dtype)
+            num_momentum_rows = given_rows.shape[0]
+
+            start = block_state.num_condition_video_rows
+            available = block_state.latents.shape[0] - start
+            if num_momentum_rows > available:
+                raise ValueError(
+                    f"The carried clip encodes to {num_momentum_rows} rows, more than the {available} generated "
+                    f"rows available to impose on — shorten the momentum duration."
+                )
+
+            block_state.given_video_rows = given_rows
+            block_state.num_momentum_rows = num_momentum_rows
+
+            if block_state.video_condition_mode == LOCKED:
+                # Clean from the first forward on, matching the `t = 1.0` the timestep plan will claim for these
+                # rows — the same presentation `ref2va` gives a reference, applied here to a target prefix.
+                block_state.latents = block_state.latents.clone()
+                block_state.latents[start : start + num_momentum_rows] = given_rows.clone()
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3MomentumSetTimestepsStep(MiniMaxH3SetTimestepsStep):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "The `t2va` timestep plan, with the imposed video prefix presented as clean (`t = 1.0`) when the "
+            "carried clip is locked. `blended` and `off` leave those rows on their own schedule, where their "
+            "content really is at the level their timestep claims."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(name="video_condition_mode", type_hint=str, default=BLENDED, description="See the block above."),
+            InputParam(name="given_video_rows", type_hint=torch.Tensor, description="The encoded carried clip, or None."),
+            InputParam(name="num_momentum_rows", type_hint=int, default=0),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+
+        components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+        components.audio_scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+        block_state.timesteps = components.scheduler.timesteps
+        block_state.audio_timesteps = components.audio_scheduler.timesteps
+
+        locked = block_state.video_condition_mode == LOCKED and block_state.given_video_rows is not None
+        # The rows are *given*, not *prepended*: `num_condition_video_rows` stays whatever the (keyframe-free)
+        # layout already set it to everywhere else, so the loop still steps them and the unpack step still keeps
+        # them — only this plan's own clean-row count changes.
+        num_clean_video_rows = int(block_state.num_momentum_rows) if locked else 0
+
+        block_state.row_timestep_plan = [
+            tuple(
+                tensor.to(device)
+                for tensor in self.build_row_timesteps(
+                    block_state.video_indices,
+                    block_state.audio_indices,
+                    num_clean_video_rows,
+                    block_state.num_condition_audio_rows,
+                    block_state.text_indices.numel(),
+                    float(timestep),
+                    float(audio_timestep),
+                    1.0,
+                    max(float(audio_timestep), components.keyframe_noise_aug),
+                )
+            )
+            for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps)
+        ]
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3MomentumLoopSchedulerStep(MiniMaxH3LoopSchedulerStep):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Steps both modalities, then imposes the carried clip on the leading video rows: clean when locked, "
+            "re-noised to the video schedule's next sigma when blended. The scheduler is still stepped over "
+            "those rows so its step index keeps up with the loop; the write-back is what the next forward reads."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(name="video_condition_mode", type_hint=str, default=BLENDED, description="See the block above."),
+            InputParam(name="given_video_rows", type_hint=torch.Tensor, description="The encoded carried clip, or None."),
+            InputParam(name="video_noise_rows", type_hint=torch.Tensor, description="The noise the imposed rows were drawn from."),
+            InputParam(name="num_momentum_rows", type_hint=int, default=0),
+            InputParam(name="num_condition_video_rows", type_hint=int, default=0),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: BlockState, i: int, t: torch.Tensor):
+        components, block_state = super().__call__(components, state, i, t)
+
+        given = block_state.given_video_rows
+        if given is None or block_state.video_condition_mode == OFF:
+            return components, block_state
+
+        start = block_state.num_condition_video_rows
+        end = start + block_state.num_momentum_rows
+
+        if block_state.video_condition_mode == LOCKED:
+            # `.clone()` is load-bearing — see the identical note on the audio version of this step: without it
+            # the imposed slice *is* `given`, and the next step's write overwrites the carried clip itself.
+            block_state.latents[start:end] = given.to(block_state.latents).clone()
+        else:
+            timesteps = block_state.timesteps
+            next_timestep = float(timesteps[i + 1]) if i + 1 < timesteps.numel() else 1.0
+            block_state.latents[start:end] = components.scheduler.scale_noise(
+                given.to(block_state.latents),
+                next_timestep,
+                block_state.video_noise_rows[start:end].to(block_state.latents),
+            )
+        return components, block_state
+
+
+class MiniMaxH3MomentumDenoiseStep(MiniMaxH3DenoiseLoopWrapper):
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3LoopDenoiser, MiniMaxH3MomentumLoopSchedulerStep]
+    block_names = ["denoiser", "update"]
+
+    @property
+    def description(self) -> str:
+        return "Runs the `t2va` denoising loop with the carried clip imposed on the leading video rows every step."
+
+
+class MiniMaxH3MomentumCoreDenoiseStep(SequentialPipelineBlocks):
+    """`MiniMaxH3CoreDenoiseStep`'s `t2va`-shaped core, with the momentum-conditioning blocks inserted at the
+    same points `h3_a2v_blocks.py` inserts its audio ones: after the noise draw, before the timestep plan."""
+
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3NoKeyframeAnchorsStep,
+        MiniMaxH3PrepareLayoutStep,
+        MiniMaxH3PrepareLatentsStep,
+        MiniMaxH3MomentumConditionStep,
+        MiniMaxH3MomentumSetTimestepsStep,
+        MiniMaxH3MomentumDenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+    ]
+    block_names = [
+        "no_keyframe_anchors",
+        "prepare_layout",
+        "prepare_latents",
+        "momentum_condition",
+        "set_timesteps",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Core denoising workflow for momentum-conditioned `t2va`: `MiniMaxH3CoreDenoiseStep` with the "
+            "carried clip encoded into the leading video rows and imposed on every step."
+        )
+
+
+class MiniMaxH3MomentumGeneratorBlocks(SequentialPipelineBlocks):
+    """The denoising half of the split deployment, momentum-conditioned — for Chunked Generation continuation
+    chunks only. No keyframe branch at all: the momentum block already determines the opening frames more
+    directly than a keyframe could, so continuation chunks never pass `image`/`last_image`. Chunk one, with
+    nothing to carry, keeps using the ordinary keyframe-capable `MiniMaxH3GeneratorBlocks`.
+    """
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3MomentumCoreDenoiseStep, MiniMaxH3DecodeStep]
+    block_names = ["denoise", "decode"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "The denoising half of a split MiniMax-H3 deployment, conditioned on a carried trailing clip: "
+            "always the `t2va` shape, without a text-encoder step, with the leading video rows imposed rather "
+            "than generated."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("videos", description="The generated video."),
+            OutputParam(
+                "audio",
+                type_hint=torch.Tensor,
+                description="The soundtrack of the packed sequence, of shape `(1, 2, num_samples)`.",
+            ),
+            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack in Hz."),
+        ]