Blame for h3_keyframe_control.py

"""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."),
        ]