Upload h3_keyframe_control.py

h3_keyframe_control.py

@@ -0,0 +1,305 @@
+"""Per-keyframe injection frame and strength for the `fl2va` branch of the staged-denoising Space.
+
+Two independent knobs per keyframe, both optional (defaults reproduce the stock pipeline exactly):
+
+* **Injection frame.** A keyframe's conditioning rows carry one rotary time, `anchor_time`. Stock supports exactly
+  two values: "first" = `T` and "last" = `T + 5/3 * (N - 1)`, where `T` is the text-token count and `N` the number of
+  pixel frames the latent grid spans. The rotary clock runs at 5/3 units per *pixel* frame, so injecting at pixel
+  frame `p` is just `anchor_time = T + 5/3 * p`. Nothing else about the keyframe changes.
+  NOTE: the released model was only trained with the two stock positions. Other positions are untested.
+
+* **Strength.** A keyframe is held at `t = 0.999` ("just short of clean": `x_t = t*x0 + (1-t)*noise`). Strength `s`
+  in `[0, 1]` maps to that `t`: `s = 1` is stock (0.999), `s = 0` is pure noise (`t = 0`, the keyframe carries no
+  information). The same `t` is what the rows are presented to the transformer at, so the label stays honest.
+
+Blocks:
+    MiniMaxH3KeyframePlacementStep   - rewrites the keyframe rows' rotary time after `prepare_layout`
+    MiniMaxH3KeyframeConditionStep   - noises each keyframe to its own level instead of the shared 0.999
+    (build_row_timesteps patch)      - presents each keyframe's rows at its own level, in every sampler
+
+The patch of `MiniMaxH3SetTimestepsStep.build_row_timesteps` is deliberate: the dpmpp_2s / SDE / seeds_2 samplers
+build their own second-evaluation row timesteps through that same static method, so patching it keeps all of them
+consistent with the plan without touching their files. It is a no-op for anything with no conditioning rows
+(`t2va`, the momentum pipeline) and whenever no per-keyframe levels were set for the request.
+"""
+
+import torch
+
+from diffusers.modular_pipelines.minimax_h3.before_denoise import (
+    _ROPE_FRAME_RESCALE,
+    _ROPE_FRAMES_PER_LATENT,
+    MiniMaxH3FL2VAPrepareLatentsStep,
+    MiniMaxH3PrepareConditionLatentsStep,
+    MiniMaxH3PrepareLayoutStep,
+    MiniMaxH3SetTimestepsStep,
+    patchify_video_latents,
+)
+from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3PrepareLatentsStep
+from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3DenoiseStep
+from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
+    MiniMaxH3CoreDenoiseStep,
+    MiniMaxH3DecodeStep,
+)
+from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
+from diffusers.modular_pipelines.modular_pipeline import (
+    ConditionalPipelineBlocks,
+    ModularPipelineBlocks,
+    PipelineState,
+    SequentialPipelineBlocks,
+)
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
+from diffusers.utils.torch_utils import randn_tensor
+
+from h3_split_blocks import MiniMaxH3SplitBeforeEncodeStep, MiniMaxH3SplitVaeEncoderStep
+
+
+# ---------------------------------------------------------------------------------------------------------------
+# Per-request state shared with build_row_timesteps. Set by the condition step, read by the patched static method.
+# ---------------------------------------------------------------------------------------------------------------
+_LEVELS: tuple[float, ...] | None = None
+_AUG: float = 0.999
+
+
+def strength_to_level(strength: float, aug: float) -> float:
+    """Slider strength in [0, 1] -> the keyframe's `t`. 1 (or anything above `aug`) is stock."""
+    return max(0.0, min(float(strength), float(aug)))
+
+
+_original_build_row_timesteps = MiniMaxH3SetTimestepsStep.build_row_timesteps
+
+
+def _build_row_timesteps(
+    video_indices,
+    audio_indices,
+    num_condition_video_rows,
+    num_condition_audio_rows,
+    num_text_tokens,
+    video_timestep,
+    audio_timestep,
+    condition_video_timestep,
+    condition_audio_timestep,
+):
+    levels = _LEVELS
+    n = len(levels) if levels else 0
+    if n == 0 or num_condition_video_rows <= 0 or num_condition_video_rows % n != 0:
+        return _original_build_row_timesteps(
+            video_indices, audio_indices, num_condition_video_rows, num_condition_audio_rows, num_text_tokens,
+            video_timestep, audio_timestep, condition_video_timestep, condition_audio_timestep,
+        )
+
+    sequence_length = int(video_indices.numel() + audio_indices.numel() + num_text_tokens)
+    row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32)
+    rows_per_keyframe = num_condition_video_rows // n
+    for i, level in enumerate(levels):
+        block = video_indices[i * rows_per_keyframe : (i + 1) * rows_per_keyframe]
+        # At or above the stock level the keyframe is presented exactly as stock does (max with the current step's
+        # timestep). Below it, the keyframe really is that noisy and stays labelled that way for every step.
+        row_timesteps[block] = condition_video_timestep if level >= _AUG else level
+    row_timesteps[audio_indices[num_condition_audio_rows:]] = audio_timestep
+    row_timesteps[audio_indices[:num_condition_audio_rows]] = condition_audio_timestep
+    return torch.unique(row_timesteps, sorted=True, return_inverse=True)
+
+
+MiniMaxH3SetTimestepsStep.build_row_timesteps = staticmethod(_build_row_timesteps)
+
+
+# ---------------------------------------------------------------------------------------------------------------
+# Blocks
+# ---------------------------------------------------------------------------------------------------------------
+class MiniMaxH3KeyframePlacementStep(ModularPipelineBlocks):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Moves each keyframe's conditioning rows to a chosen pixel frame on the rotary time axis. The stock "
+            "layout only knows 'first' (frame 0) and 'last' (the final frame); `anchor_time = T + 5/3 * frame`."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="position_ids", type_hint=torch.Tensor, required=True),
+            InputParam(name="text_indices", type_hint=torch.Tensor, required=True),
+            InputParam(name="num_latent_frames", type_hint=int, required=True),
+            InputParam(name="latent_height", type_hint=int, required=True),
+            InputParam(name="latent_width", type_hint=int, required=True),
+            InputParam(name="keyframe_anchors", type_hint=tuple, default=()),
+            InputParam(
+                name="keyframe_frames",
+                type_hint=tuple,
+                default=None,
+                description="Requested pixel frame per keyframe, in packed order. None / missing entries keep stock.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [OutputParam("position_ids", type_hint=torch.Tensor)]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState):
+        block_state = self.get_block_state(state)
+        anchors = tuple(block_state.keyframe_anchors or ())
+        frames = block_state.keyframe_frames
+        if anchors and frames is not None:
+            _, patch_h, patch_w = components.patch_size
+            rows_per_frame = (block_state.latent_height // patch_h) * (block_state.latent_width // patch_w)
+            num_text = int(block_state.text_indices.numel())
+            pattern = _ROPE_FRAMES_PER_LATENT
+            total_pixel_frames = sum(pattern[i % len(pattern)] for i in range(block_state.num_latent_frames))
+            last_frame = total_pixel_frames - 1
+            position_ids = block_state.position_ids
+            for index, (anchor, requested) in enumerate(zip(anchors, frames)):
+                if requested is None:
+                    continue
+                frame = max(0, min(int(requested), last_frame))
+                # Defaults reproduce the stock anchors bit for bit: leave those rows untouched.
+                if (anchor == "first" and frame <= 0) or (anchor == "last" and frame >= last_frame):
+                    continue
+                rows = slice(num_text + index * rows_per_frame, num_text + (index + 1) * rows_per_frame)
+                position_ids[rows, 0] = float(num_text) + _ROPE_FRAME_RESCALE * frame
+            block_state.position_ids = position_ids
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3KeyframeConditionStep(MiniMaxH3PrepareConditionLatentsStep):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "`MiniMaxH3PrepareConditionLatentsStep` with a noise level per keyframe: `x_t = t*x0 + (1-t)*noise`, "
+            "`t = strength` capped at the stock 0.999. Draw order is unchanged (one draw per condition)."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(
+                name="keyframe_strengths",
+                type_hint=tuple,
+                default=None,
+                description="Strength in [0, 1] per keyframe, in packed order. None keeps the stock 0.999.",
+            )
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return super().intermediate_outputs
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState):
+        global _LEVELS, _AUG
+
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+        patch_size = components.patch_size
+        aug = float(components.keyframe_noise_aug)
+        strengths = block_state.keyframe_strengths
+
+        levels = []
+        packed = []
+        for index, condition in enumerate(block_state.condition_latents):
+            noise = randn_tensor(condition.shape, generator=block_state.generator, device=device, dtype=torch.float32)
+            strength = 1.0 if strengths is None or index >= len(strengths) or strengths[index] is None else strengths[index]
+            level = strength_to_level(strength, aug)
+            levels.append(level)
+            noised = components.scheduler.scale_noise(condition.to(device), level, noise)
+            packed.append(patchify_video_latents(noised, patch_size))
+        block_state.condition_rows = torch.cat(packed)
+
+        if block_state.condition_rows.shape[0] != block_state.num_condition_video_rows:
+            raise ValueError(
+                f"The layout reserved {block_state.num_condition_video_rows} conditioning rows but the encoded "
+                f"conditioning latents pack into {block_state.condition_rows.shape[0]}."
+            )
+
+        _AUG = aug
+        _LEVELS = tuple(levels) if any(level < aug for level in levels) else None
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3KeyframeControlCoreDenoiseStep(SequentialPipelineBlocks):
+    """`MiniMaxH3FL2VACoreDenoiseStep` with the placement step after the layout and the per-keyframe condition step."""
+
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3PrepareLayoutStep,
+        MiniMaxH3KeyframePlacementStep,
+        MiniMaxH3KeyframeConditionStep,
+        MiniMaxH3PrepareLatentsStep,
+        MiniMaxH3FL2VAPrepareLatentsStep,
+        MiniMaxH3SetTimestepsStep,
+        MiniMaxH3DenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+    ]
+    block_names = [
+        "prepare_layout",
+        "keyframe_placement",
+        "prepare_condition_latents",
+        "prepare_latents",
+        "prepare_latents_fl2va",
+        "set_timesteps",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self) -> str:
+        return "Core denoising workflow for `fl2va` with per-keyframe injection frame and strength."
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam("latents", type_hint=torch.Tensor, description="The generated video latents."),
+            OutputParam("audio_latents", type_hint=torch.Tensor, description="The generated audio latents."),
+        ]
+
+
+class MiniMaxH3KeyframeControlDenoiseStep(ConditionalPipelineBlocks):
+    """Split-deployment denoise selector with the keyframe-control `fl2va` core in place of the stock one."""
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3KeyframeControlCoreDenoiseStep, MiniMaxH3CoreDenoiseStep]
+    block_names = ["fl2va", "t2va"]
+    block_trigger_inputs = ["image", "last_image"]
+    default_block_name = "t2va"
+
+    def select_block(self, **kwargs) -> str | None:
+        if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
+            return "fl2va"
+        return None
+
+    @property
+    def description(self):
+        return "`fl2va` (keyframe control) when a keyframe is provided, `t2va` otherwise."
+
+
+class MiniMaxH3KeyframeGeneratorBlocks(SequentialPipelineBlocks):
+    """Drop-in for `h3_split_blocks.MiniMaxH3GeneratorBlocks`, with keyframe injection frame / strength."""
+
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3SplitBeforeEncodeStep,
+        MiniMaxH3SplitVaeEncoderStep,
+        MiniMaxH3KeyframeControlDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = ["resize", "vae_encoder", "denoise", "decode"]
+
+    @property
+    def description(self):
+        return "Denoising half of the split deployment with per-keyframe injection frame and strength."
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("videos", description="The generated video."),
+            OutputParam("audio", type_hint=torch.Tensor, description="The soundtrack, `(1, 2, num_samples)`."),
+            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack in Hz."),
+        ]