h3_split_blocks.py

"""The halves of a **split** MiniMax-H3 deployment, for both of its checkpoint partitions.

MiniMax-H3 is modular-only, and `MiniMaxH3Blocks` is a `SequentialPipelineBlocks` of eight steps:

    setup -> text_encoder -> vae_encoder -> prepare_layout -> prepare_latents -> set_timesteps -> denoise -> decode

The conditioner (a 62.14 GiB Qwen3-VL) and the denoiser (a 61.73 GiB transformer plus ~20.5 GiB of float32 VAEs) do
not fit on one 95 GiB card unquantized, so this module cuts the sequence in two at the `text_encoder` step:

  * `MiniMaxH3ConditionerBlocks` = `[setup, text_encoder]`  — loads `text_encoder` / `tokenizer` / `processor` only,
    and emits `prompt_embeds` + `text_token_tags`, which is the whole wire format between the two halves.
  * `MiniMaxH3GeneratorBlocks`   = everything else          — loads `transformer` / `vae` / `audio_vae` / the two
    schedulers only, and takes `prompt_embeds` + `text_token_tags` as *inputs*.

`MiniMaxH3Ref2VABlocks`, the `ref2va` task of the Ref2VA partition, is cut at the very same step:

    setup -> text_encoder -> reference_encoder -> prepare_layout -> prepare_latents -> set_timesteps -> denoise
          -> decode

  * `MiniMaxH3Ref2VAConditionerBlocks` = `[setup, text_encoder]` — the same three conditioner components, so one
    conditioner Space serves both partitions, and the same two tensors come back.
  * `MiniMaxH3Ref2VAGeneratorBlocks`   = everything else         — loads `transformer_ref` / `vae` / `audio_vae` /
    the two schedulers.

Only *text* encoding is remote. `reference_encoder` is the `ref2va` half's own encoder step and stays on the
denoising side: it runs the two autoencoders over the references, which the conditioner Space does not hold.

`setup` runs on both sides on purpose. It owns no component (it is PIL, PyAV-decoded media and arithmetic), it
resolves the canvas, the `17 * n + 5` frame count and the latent geometry, and it prepares the keyframes or the
references — which the conditioner needs to build its vision blocks and the generator needs to encode with the VAEs.
Running it twice over the same inputs is deterministic; both conditioner halves return the resolved `height` /
`width` / `num_frames` anyway, so the caller pins them explicitly on the generating half. For `ref2va` that pinning
is not optional: `num_frames` may be left out of a request whose single audio-bearing reference sets the duration,
and it is the conditioner that resolves it.
"""

from diffusers.modular_pipelines.minimax_h3.before_denoise import (
    MiniMaxH3PrepareLatentsStep,
    MiniMaxH3PrepareLayoutStep,
    MiniMaxH3Ref2VAPrepareLayoutStep,
    MiniMaxH3SetTimestepsStep,
)
from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep, MiniMaxH3SetupStep
from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3DenoiseStep, MiniMaxH3Ref2VADenoiseStep
from diffusers.modular_pipelines.minimax_h3.encoders import (
    MiniMaxH3Ref2VAReferenceEncoderStep,
    MiniMaxH3Ref2VATextEncoderStep,
    MiniMaxH3TextEncoderStep,
)
from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
    MiniMaxH3AutoKeyframeVaeEncoderStep,
    MiniMaxH3DecodeStep,
    _generation_outputs,
)
from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam


def _wire_outputs() -> list[OutputParam]:
    """The wire format of the split, plus the plan the caller pins on the generating half."""
    return [
        OutputParam.template("prompt_embeds"),
        OutputParam("text_token_tags", description="The per-row modality tag of every row of `prompt_embeds`."),
        OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
        OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
        OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."),
    ]


class MiniMaxH3ConditionerBlocks(SequentialPipelineBlocks):
    """The conditioner half of a split MiniMax-H3: the request plan plus the Qwen3-VL read at its 50th layer."""

    model_name = "minimax-h3"
    block_classes = [MiniMaxH3SetupStep, MiniMaxH3TextEncoderStep]
    block_names = ["setup", "text_encoder"]

    @property
    def description(self):
        return (
            "The conditioner half of a split MiniMax-H3 deployment: resolves the request plan (canvas, frame count, "
            "latent geometry, keyframes on the canvas) and encodes MiniMax-H3's presentation of it into the "
            "`prompt_embeds` / `text_token_tags` pair the denoising half consumes."
        )

    @property
    def outputs(self):
        return _wire_outputs()


class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
    """The denoising half of a split MiniMax-H3: `MiniMaxH3Blocks` with its `text_encoder` step removed."""

    model_name = "minimax-h3"
    block_classes = [
        MiniMaxH3SetupStep,
        MiniMaxH3AutoKeyframeVaeEncoderStep,
        MiniMaxH3PrepareLayoutStep,
        MiniMaxH3PrepareLatentsStep,
        MiniMaxH3SetTimestepsStep,
        MiniMaxH3DenoiseStep,
        MiniMaxH3DecodeStep,
    ]
    block_names = [
        "setup",
        "vae_encoder",
        "prepare_layout",
        "prepare_latents",
        "set_timesteps",
        "denoise",
        "decode",
    ]

    @property
    def description(self):
        return (
            "The denoising half of a split MiniMax-H3 deployment: `MiniMaxH3Blocks` without its text-encoder step, so "
            "`prompt_embeds` and `text_token_tags` come in as inputs and the 62.14 GiB Qwen3-VL conditioner is never "
            "loaded here."
        )

    @property
    def outputs(self):
        return _generation_outputs()


class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
    """The conditioner half of a split `ref2va`: the request plan plus the Qwen3-VL read at its 50th layer.

    Component for component this is `MiniMaxH3ConditionerBlocks` — `text_encoder`, `tokenizer`, `processor` — which
    is what lets one conditioner Space serve both partitions of the checkpoint out of the weights it already holds.
    What differs is the presentation the Qwen3-VL is shown: `ref2va` prepends a label per reference, numbered per
    modality, and a vision block per image and per merged video frame pair, so the references themselves have to
    reach this half. An audio reference never does — it contributes its `"<Audio j>: "` label and nothing else — but
    it is still part of the request here, because a single audio-bearing reference is what resolves `num_frames`
    when the request left it open.
    """

    model_name = "minimax-h3-ref2va"
    block_classes = [MiniMaxH3Ref2VASetupStep, MiniMaxH3Ref2VATextEncoderStep]
    block_names = ["setup", "text_encoder"]

    @property
    def description(self):
        return (
            "The conditioner half of a split MiniMax-H3 `ref2va` deployment: resolves the request plan (canvas, frame "
            "count, latent geometry, references prepared at their own resolutions) and encodes MiniMax-H3's "
            "presentation of it into the `prompt_embeds` / `text_token_tags` pair the denoising half consumes."
        )

    @property
    def outputs(self):
        return _wire_outputs()


class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
    """The denoising half of a split `ref2va`: `MiniMaxH3Ref2VABlocks` with its `text_encoder` step removed.

    Only the text-encoder step is dropped. `reference_encoder` is this half's own encoder — it runs the video VAE
    over the image and video references and the audio VAE over the soundtracks, and it is where the references'
    latent geometry is resolved, which the packed layout is built from — so it stays here, next to the autoencoders.
    """

    model_name = "minimax-h3-ref2va"
    block_classes = [
        MiniMaxH3Ref2VASetupStep,
        MiniMaxH3Ref2VAReferenceEncoderStep,
        MiniMaxH3Ref2VAPrepareLayoutStep,
        MiniMaxH3PrepareLatentsStep,
        MiniMaxH3SetTimestepsStep,
        MiniMaxH3Ref2VADenoiseStep,
        MiniMaxH3DecodeStep,
    ]
    block_names = [
        "setup",
        "reference_encoder",
        "prepare_layout",
        "prepare_latents",
        "set_timesteps",
        "denoise",
        "decode",
    ]

    @property
    def description(self):
        return (
            "The denoising half of a split MiniMax-H3 `ref2va` deployment: `MiniMaxH3Ref2VABlocks` without its "
            "text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the 62.14 GiB Qwen3-VL "
            "conditioner is never loaded here. The transformer is the `transformer_ref` partition."
        )

    @property
    def outputs(self):
        return _generation_outputs()