Update h3_split_blocks.py

h3_split_blocks.py

@@ -12,22 +12,23 @@ Two things the blocks leave to the caller: a keyframe reaches them EXIF-transpos
 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
+import torch
+
+from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3ResizeStep, MiniMaxH3Ref2VASetupStep
 from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
 from diffusers.modular_pipelines.minimax_h3.encoders import (
+    MiniMaxH3KeyframeVaeEncoderStep,
     MiniMaxH3Ref2VAReferenceEncoderStep,
     MiniMaxH3Ref2VATextEncoderStep,
     MiniMaxH3TextEncoderStep,
 )
 from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
-    MiniMaxH3AutoKeyframeVaeEncoderStep,
-    MiniMaxH3AutoResizeStep,
     MiniMaxH3CoreDenoiseStep,
     MiniMaxH3DecodeStep,
+    MiniMaxH3FL2VACoreDenoiseStep,
     MiniMaxH3Ref2VACoreDenoiseStep,
-    _generation_outputs,
 )
-from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
 from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam
 
 
@@ -46,11 +47,95 @@ def _wire_outputs(num_frames: bool = True) -> list[OutputParam]:
     ]
 
 
+class MiniMaxH3SplitBeforeEncodeStep(ConditionalPipelineBlocks):
+    """Media preparation block for the split deployment's `t2va` / `fl2va` half — no `ref2va` branch.
+
+    Narrower than main's own `MiniMaxH3AutoBeforeEncodeStep`, the same way `MiniMaxH3A2VAutoBeforeEncodeStep` is
+    in the audio-conditioned Space: the upstream auto wrapper's `ref2va` branch declares
+    `MiniMaxH3Ref2VASetupStep`, which this half never needs — `MiniMaxH3Ref2VAConditionerBlocks` already owns
+    that, on its own `transformer_ref`-side classes.
+    """
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3ResizeStep]
+    block_names = ["keyframes"]
+    block_trigger_inputs = ["image", "last_image"]
+    default_block_name = None
+
+    def select_block(self, **kwargs) -> str | None:
+        if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
+            return "keyframes"
+        return None
+
+    @property
+    def description(self):
+        return (
+            "Media preparation block.\n"
+            " - `MiniMaxH3ResizeStep` runs when a keyframe is provided (`fl2va`), putting it onto the target "
+            "canvas.\n"
+            " - a text-only request (`t2va`) skips this block, and the layout step falls back to MiniMax-H3's own "
+            "16:9 canvas."
+        )
+
+
+class MiniMaxH3SplitVaeEncoderStep(ConditionalPipelineBlocks):
+    """VAE encoder block for the split deployment's `t2va` / `fl2va` half — no `ref2va` branch, for the same
+    reason `MiniMaxH3SplitBeforeEncodeStep` has none."""
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3KeyframeVaeEncoderStep]
+    block_names = ["keyframes"]
+    block_trigger_inputs = ["image", "last_image"]
+    default_block_name = None
+
+    def select_block(self, **kwargs) -> str | None:
+        if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
+            return "keyframes"
+        return None
+
+    @property
+    def description(self):
+        return (
+            "VAE encoder block.\n"
+            " - `MiniMaxH3KeyframeVaeEncoderStep` runs when a keyframe is provided (`fl2va`).\n"
+            " - a text-only request (`t2va`) skips this block."
+        )
+
+
+class MiniMaxH3SplitDenoiseStep(ConditionalPipelineBlocks):
+    """Denoise block for the split deployment's `t2va` / `fl2va` half.
+
+    `MiniMaxH3CoreDenoiseStep` is `t2va`-only as of the diffusers 0.40.0 refactor (it opens with
+    `MiniMaxH3NoKeyframeAnchorsStep`); keyframe-anchored generation moved to the separate
+    `MiniMaxH3FL2VACoreDenoiseStep`. This selects between them the same way `MiniMaxH3A2VAutoDenoiseStep` selects
+    between its own audio-conditioned pair — no `ref2va` branch, since that would declare `transformer_ref`, the
+    partition this half must never load.
+    """
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3FL2VACoreDenoiseStep, 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 (
+            "Denoise block.\n"
+            " - the `fl2va` core runs when a keyframe is provided, against the `transformer` partition.\n"
+            " - the `t2va` core runs otherwise, against the same partition."
+        )
+
 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_classes = [MiniMaxH3SplitBeforeEncodeStep, MiniMaxH3TextEncoderStep]
     block_names = ["resize", "text_encoder"]
 
     @property
@@ -71,9 +156,9 @@ class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
 
     model_name = "minimax-h3"
     block_classes = [
-        MiniMaxH3AutoResizeStep,
-        MiniMaxH3AutoKeyframeVaeEncoderStep,
-        MiniMaxH3CoreDenoiseStep,
+        MiniMaxH3SplitBeforeEncodeStep,
+        MiniMaxH3SplitVaeEncoderStep,
+        MiniMaxH3SplitDenoiseStep,
         MiniMaxH3AfterDenoiseStep,
         MiniMaxH3DecodeStep,
     ]
@@ -89,7 +174,15 @@ class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
 
     @property
     def outputs(self):
-        return _generation_outputs()
+        return [
+            OutputParam.template("videos", description="The generated video."),
+            OutputParam(
+                "audio",
+                type_hint=torch.Tensor,
+                description="The soundtrack of the packed sequence, of shape `(1, 2, num_samples)`.",
+            ),
+            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack in Hz."),
+        ]
 
 
 class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
@@ -144,4 +237,12 @@ class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
 
     @property
     def outputs(self):
-        return _generation_outputs()
+        return [
+            OutputParam.template("videos", description="The generated video."),
+            OutputParam(
+                "audio",
+                type_hint=torch.Tensor,
+                description="The soundtrack of the packed sequence, of shape `(1, 2, num_samples)`.",
+            ),
+            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack in Hz."),
+        ]