Ship a modular workflow that carries keyframes and references in one run

combined_blocks.py

@@ -0,0 +1,296 @@
+"""A MiniMax-H3 modular workflow that accepts keyframes **and** references in the same run.
+
+Why this exists. MiniMax-H3 denoises one packed sequence, and that sequence can hold keyframe conditioning rows and
+reference conditioning rows at the same time. `diffusers`' shipped blocks cannot express it: their conditional steps
+dispatch either/or (`select_block` checks `references` first), so a request carrying both is accepted and the
+keyframes are **silently dropped** — no error, no warning, just a reference-only generation.
+
+`MiniMaxH3CombinedBlocks` is `MiniMaxH3Blocks` with three of its conditional steps replaced by ones that know a
+fourth shape. Nothing else changes: `t2va`, `fl2va` and `ref2va` requests take exactly the same path they always did,
+and the denoising loop is untouched, because the conditioning rows are simply the leading rows of the sequence and it
+only ever steps what comes after them.
+
+    from diffusers import ModularPipeline
+
+    pipe = ModularPipeline.from_pretrained(REPO, trust_remote_code=True, workflow="combined")
+    pipe.load_components(dtype=torch.bfloat16, trust_remote_code=True)
+    out = pipe(prompt=..., references=[...], image=first, last_image=last, num_frames=124, ...)
+
+The layout itself lives in `combined_layout.py`, which is pinned by reproducing both shipped builders bit for bit in
+their degenerate cases.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from diffusers.modular_pipelines import ConditionalPipelineBlocks, SequentialPipelineBlocks
+from diffusers.modular_pipelines.minimax_h3.before_denoise import (
+    MiniMaxH3PrepareConditionLatentsStep,
+    MiniMaxH3PrepareLatentsStep,
+    MiniMaxH3Ref2VAPrepareLatentsStep,
+    MiniMaxH3Ref2VAPrepareLayoutStep,
+    MiniMaxH3SetTimestepsStep,
+)
+from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
+from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3Ref2VADenoiseStep
+from diffusers.modular_pipelines.minimax_h3.encoders import (
+    MiniMaxH3Ref2VAReferenceEncoderStep,
+    encode_vae_condition,
+)
+from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
+    MiniMaxH3AutoDenoiseStep,
+    MiniMaxH3AutoTextEncoderStep,
+    MiniMaxH3AutoVaeEncoderStep,
+    MiniMaxH3AutoBeforeEncodeStep,
+    MiniMaxH3Blocks,
+    MiniMaxH3DecodeStep,
+)
+from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+from .combined_layout import build_combined_packed_sequence
+
+
+def _anchors_of(image, last_image) -> tuple[str, ...]:
+    """Which end of the clip each keyframe is anchored to, in packed order."""
+    return tuple(name for name, value in (("first", image), ("last", last_image)) if value is not None)
+
+
+class MiniMaxH3KeyframesOnCanvasStep(ModularPipelineBlocks):
+    """Put the keyframes on the canvas the reference setup already resolved.
+
+    The `fl2va` resize step *derives* the canvas from the keyframe's aspect ratio. Here the references have already
+    settled it, so the keyframes are stretched onto it — which is what `fl2va` does to a keyframe whose aspect does
+    not match the target anyway.
+    """
+
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return "Stretches the keyframes of a combined request onto the canvas the reference setup resolved."
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="image", description="Keyframe the video starts from."),
+            InputParam(name="last_image", description="Keyframe the video ends on."),
+            InputParam(name="height", type_hint=int, required=True, description="Resolved height in pixels."),
+            InputParam(name="width", type_hint=int, required=True, description="Resolved width in pixels."),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("keyframes", type_hint=list, description="The keyframes on the target canvas, packed order."),
+            OutputParam("keyframe_anchors", type_hint=tuple, description="Which end each keyframe is anchored to."),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        size = (block_state.width, block_state.height)
+        block_state.keyframe_anchors = _anchors_of(block_state.image, block_state.last_image)
+        block_state.keyframes = [
+            frame.convert("RGB").resize(size)
+            for frame in (block_state.image, block_state.last_image)
+            if frame is not None
+        ]
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3CombinedKeyframeEncoderStep(ModularPipelineBlocks):
+    """Encode the keyframes and put them **in front of** the reference latents.
+
+    Order is the contract: the combined layout reserves `[text | keyframe cond | reference blocks | targets]`, and the
+    stock prepare-latents step concatenates this list in order — then asserts the rows it produced equal the rows the
+    layout reserved, so a mismatch raises instead of degrading quietly.
+    """
+
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return "Encodes a combined request's keyframes and prepends them to the reference conditioning latents."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("vae")]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="keyframes", type_hint=list, required=True, description="Keyframes on the canvas."),
+            InputParam(name="condition_latents", type_hint=list, required=True,
+                       description="Reference conditioning latents, packed order."),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("condition_latents", type_hint=list,
+                        description="Keyframe latents first, then the reference latents."),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+        keyframe_latents = [
+            encode_vae_condition(
+                components.vae,
+                torch.from_numpy(np.array(image)).to(device).permute(2, 0, 1)[None, :, None],
+                components.pixel_mean,
+                components.pixel_std,
+                components.keyframe_encode_seed,
+            )
+            for image in block_state.keyframes
+        ]
+        block_state.condition_latents = keyframe_latents + list(block_state.condition_latents)
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3CombinedPrepareLayoutStep(MiniMaxH3Ref2VAPrepareLayoutStep):
+    """The `ref2va` layout step, with keyframe rows packed ahead of the reference blocks."""
+
+    @property
+    def description(self) -> str:
+        return (
+            "Resolves the latent shapes of a combined request and builds its packed layout — "
+            "`[text | keyframe conditions | reference blocks | target audio | target video]`. The references push the "
+            "target timeline out, so the keyframe anchors ride on the timeline their spans leave behind."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(name="keyframe_anchors", type_hint=tuple, default=(),
+                       description="Which end of the video each keyframe is anchored to, in packed order."),
+        ]
+
+    # Overrides the parent's `@staticmethod` as a bound method, which is how the anchors reach the builder: the parent
+    # calls `self.build_ref2va_packed_sequence(...)` positionally and knows nothing about keyframes.
+    def build_ref2va_packed_sequence(self, *args, **kwargs):
+        return build_combined_packed_sequence(*args, keyframe_anchors=self._keyframe_anchors, **kwargs)
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        self._keyframe_anchors = tuple(getattr(block_state, "keyframe_anchors", ()) or ())
+        try:
+            return super().__call__(components, state)
+        finally:
+            self._keyframe_anchors = ()
+
+
+class MiniMaxH3CombinedSetupStep(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3Ref2VASetupStep, MiniMaxH3KeyframesOnCanvasStep]
+    block_names = ["references", "keyframes"]
+
+    @property
+    def description(self) -> str:
+        return "Resolves the request plan from the references, then puts the keyframes on the resolved canvas."
+
+
+class MiniMaxH3CombinedVaeEncoderStep(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3Ref2VAReferenceEncoderStep, MiniMaxH3CombinedKeyframeEncoderStep]
+    block_names = ["references", "keyframes"]
+
+    @property
+    def description(self) -> str:
+        return "Encodes the references, then the keyframes, leaving the keyframe latents first in packed order."
+
+
+class MiniMaxH3CombinedCoreDenoiseStep(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3CombinedPrepareLayoutStep,
+        MiniMaxH3PrepareConditionLatentsStep,
+        MiniMaxH3PrepareLatentsStep,
+        MiniMaxH3Ref2VAPrepareLatentsStep,
+        MiniMaxH3SetTimestepsStep,
+        MiniMaxH3Ref2VADenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+    ]
+    block_names = [
+        "prepare_layout",
+        "prepare_condition_latents",
+        "prepare_latents",
+        "prepare_ref_latents",
+        "set_timesteps",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self) -> str:
+        return "Core denoising for a combined request: the `ref2va` chain over the combined packed layout."
+
+
+def _is_combined(kwargs) -> bool:
+    return kwargs.get("references") is not None and (
+        kwargs.get("image") is not None or kwargs.get("last_image") is not None
+    )
+
+
+class MiniMaxH3CombinedAutoBeforeEncodeStep(MiniMaxH3AutoBeforeEncodeStep):
+    block_classes = [MiniMaxH3CombinedSetupStep] + MiniMaxH3AutoBeforeEncodeStep.block_classes
+    block_names = ["combined"] + MiniMaxH3AutoBeforeEncodeStep.block_names
+
+    def select_block(self, **kwargs) -> str | None:
+        return "combined" if _is_combined(kwargs) else super().select_block(**kwargs)
+
+
+class MiniMaxH3CombinedAutoVaeEncoderStep(MiniMaxH3AutoVaeEncoderStep):
+    block_classes = [MiniMaxH3CombinedVaeEncoderStep] + MiniMaxH3AutoVaeEncoderStep.block_classes
+    block_names = ["combined"] + MiniMaxH3AutoVaeEncoderStep.block_names
+
+    def select_block(self, **kwargs) -> str | None:
+        return "combined" if _is_combined(kwargs) else super().select_block(**kwargs)
+
+
+class MiniMaxH3CombinedAutoDenoiseStep(MiniMaxH3AutoDenoiseStep):
+    block_classes = [MiniMaxH3CombinedCoreDenoiseStep] + MiniMaxH3AutoDenoiseStep.block_classes
+    block_names = ["combined"] + MiniMaxH3AutoDenoiseStep.block_names
+
+    def select_block(self, **kwargs) -> str | None:
+        return "combined" if _is_combined(kwargs) else super().select_block(**kwargs)
+
+
+class MiniMaxH3CombinedBlocks(MiniMaxH3Blocks):
+    """`MiniMaxH3Blocks` plus a fourth shape: keyframes and references in the same generation.
+
+    Supported workflows: `t2va`, `fl2va`, `ref2va` — unchanged — and `combined`, which needs `prompt`, `references`
+    and at least one of `image` / `last_image`. A combined request runs against the `transformer_ref` partition, the
+    same one `ref2va` uses.
+    """
+
+    block_classes = [
+        MiniMaxH3CombinedAutoBeforeEncodeStep,
+        MiniMaxH3AutoTextEncoderStep,
+        MiniMaxH3CombinedAutoVaeEncoderStep,
+        MiniMaxH3CombinedAutoDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = ["before_encode", "text_encoder", "vae_encoder", "denoise", "decode"]
+    _workflow_map = dict(
+        MiniMaxH3Blocks._workflow_map,
+        combined=(
+            {"prompt": True, "references": True, "image": True},
+            {"prompt": True, "references": True, "last_image": True},
+        ),
+    )
+
+    @property
+    def description(self) -> str:
+        return (
+            "MiniMax-H3 blocks for joint video + audio generation, with the `t2va`, `fl2va` and `ref2va` workflows "
+            "unchanged and a fourth, `combined`, that carries keyframe *and* reference conditioning in one packed "
+            "sequence instead of dropping one of them."
+        )