Port h3_split_blocks.py from mrfakename/minimax-h3-ultra-fast

h3_split_blocks.py

@@ -0,0 +1,210 @@
+"""Composable conditioner and generator halves of MiniMax-H3, for both checkpoint partitions.
+
+The blocks cut `MiniMaxH3Blocks` at its `text_encoder` step. They can run in separate Spaces with `prompt_embeds` and
+`text_token_tags` as a wire format, or sequentially in one GPU worker when a compact local conditioner fits beside
+the generator.
+
+`resize` / `setup` run on **both** sides: they own no pretrained component, and each half needs the canvas and the
+prepared keyframes or normalized references. Both conditioner halves also return the resolved `height` / `width` /
+`num_frames`, which the generating half pins rather than re-deriving.
+
+Two things the blocks leave to the caller: a keyframe reaches them EXIF-transposed and in RGB, and the `t2va` / `fl2va`
+frame count is aligned to `17 * n + 5` before the call, since that arithmetic lives on the denoising side of the cut.
+"""
+
+from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
+from diffusers.modular_pipelines.minimax_h3.before_denoise import (
+    MiniMaxH3PrepareLatentsStep,
+    MiniMaxH3PrepareLayoutStep,
+    MiniMaxH3Ref2VAPrepareLayoutStep,
+    MiniMaxH3SetTimestepsStep,
+)
+from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+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,
+    MiniMaxH3AutoResizeStep,
+    MiniMaxH3DecodeStep,
+    _generation_outputs,
+)
+from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
+
+
+class _PreviewLoopMixin:
+    """The pinned Diffusers loop plus four non-authoritative TAE preview emissions."""
+
+    @property
+    def loop_inputs(self):
+        # The stock loop has no reason to carry geometry into its inner BlockState. TAE has to invert the packed rows,
+        # so explicitly retain the three dimensions produced by `prepare_layout` alongside the ordinary timesteps.
+        return [
+            *super().loop_inputs,
+            InputParam("num_latent_frames", type_hint=int, required=True),
+            InputParam("latent_height", type_hint=int, required=True),
+            InputParam("latent_width", type_hint=int, required=True),
+        ]
+
+    def __call__(self, components, state):
+        from h3_tae import maybe_emit_preview
+
+        block_state = self.get_block_state(state)
+        total = len(block_state.timesteps)
+        with self.progress_bar(total=total) as progress_bar:
+            for index, timestep in enumerate(block_state.timesteps):
+                components, block_state = self.loop_step(components, block_state, i=index, t=timestep)
+                maybe_emit_preview(components, block_state, index, total)
+                progress_bar.update()
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3PreviewDenoiseStep(_PreviewLoopMixin, MiniMaxH3DenoiseStep):
+    pass
+
+
+class MiniMaxH3Ref2VAPreviewDenoiseStep(_PreviewLoopMixin, MiniMaxH3Ref2VADenoiseStep):
+    pass
+
+
+class MiniMaxH3PreviewCoreDenoiseStep(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3PrepareLayoutStep,
+        MiniMaxH3PrepareLatentsStep,
+        MiniMaxH3SetTimestepsStep,
+        MiniMaxH3PreviewDenoiseStep,
+    ]
+    block_names = ["prepare_layout", "prepare_latents", "set_timesteps", "denoise"]
+
+
+class MiniMaxH3Ref2VAPreviewCoreDenoiseStep(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3Ref2VAPrepareLayoutStep,
+        MiniMaxH3PrepareLatentsStep,
+        MiniMaxH3SetTimestepsStep,
+        MiniMaxH3Ref2VAPreviewDenoiseStep,
+    ]
+    block_names = ["prepare_layout", "prepare_latents", "set_timesteps", "denoise"]
+
+
+def _wire_outputs(num_frames: bool = True) -> list[OutputParam]:
+    """The wire format of the split. `num_frames` is declared by the `ref2va` half alone, whose setup resolves one."""
+    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.")]
+            if num_frames
+            else []
+        ),
+    ]
+
+
+class MiniMaxH3ConditionerBlocks(SequentialPipelineBlocks):
+    """The conditioner half of a split MiniMax-H3: the keyframes on the canvas plus the Qwen3-VL read at layer 50."""
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3AutoResizeStep, MiniMaxH3TextEncoderStep]
+    block_names = ["resize", "text_encoder"]
+
+    @property
+    def description(self):
+        return (
+            "The conditioner half of a split MiniMax-H3 deployment: puts the keyframes onto the target canvas and "
+            "encodes MiniMax-H3's presentation of the request into the `prompt_embeds` / `text_token_tags` pair the "
+            "denoising half consumes. The frame count is the caller's to align."
+        )
+
+    @property
+    def outputs(self):
+        return _wire_outputs(num_frames=False)
+
+
+class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
+    """The denoising half of a split MiniMax-H3: `MiniMaxH3Blocks` with its `text_encoder` step removed."""
+
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3AutoResizeStep,
+        MiniMaxH3AutoKeyframeVaeEncoderStep,
+        MiniMaxH3PreviewCoreDenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = ["resize", "vae_encoder", "denoise", "after_denoise", "decode"]
+
+    @property
+    def description(self):
+        return (
+            "The denoising half of a split MiniMax-H3 deployment: the `t2va` / `fl2va` branch of `MiniMaxH3Blocks` "
+            "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
+            "conditioner is supplied by the caller or by the preceding local conditioner half."
+        )
+
+    @property
+    def outputs(self):
+        return _generation_outputs()
+
+
+class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
+    """The conditioner half of a split `ref2va`: the resolved plan plus the Qwen3-VL read at its 50th layer.
+
+    Component for component this is `MiniMaxH3ConditionerBlocks`, so one conditioner Space serves both partitions.
+    What differs is the presentation: `ref2va` prepends a label per reference and a vision block per image and per
+    merged video frame pair, so the references themselves have to reach this half.
+    """
+
+    model_name = "minimax-h3"
+    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, references normalized onto MiniMax-H3's own rates and 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`: the `ref2va` branch with its `text_encoder` step removed.
+
+    `reference_encoder` stays here, next to the two autoencoders it runs: its output shapes are where every reference
+    block's geometry in the packed layout comes from.
+    """
+
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3Ref2VASetupStep,
+        MiniMaxH3Ref2VAReferenceEncoderStep,
+        MiniMaxH3Ref2VAPreviewCoreDenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = ["setup", "reference_encoder", "denoise", "after_denoise", "decode"]
+
+    @property
+    def description(self):
+        return (
+            "The denoising half of a split MiniMax-H3 `ref2va` deployment: the `ref2va` branch of `MiniMaxH3Blocks` "
+            "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
+            "conditioner is supplied by the caller or preceding local half. The transformer is the `transformer_ref` partition."
+        )
+
+    @property
+    def outputs(self):
+        return _generation_outputs()