[diffusion] feat: support hybrid conditioning for minimax h3 (#36080)

2026-08-23 21:51:25 +08:00
parent 8014d9d062
commit de6a1dbd7a
10 changed files with 293 additions and 130 deletions
@@ -105,40 +105,42 @@ def _keyframe_canvas_size(shape: Any) -> tuple[int, int]:
    geometry = str(shape["geometry"])
    if geometry != "resolved_v2":
        raise ValueError(
-            "fl2va keyframe preparation requires pre-queue resolved_v2 "
+            "keyframe preparation requires pre-queue resolved_v2 "
            f"geometry, got {geometry!r}"
        )
    return int(shape["width"]), int(shape["height"])


 def _validate_keyframe_materials(plan: Any, keyframes: list[Any]) -> tuple[int, ...]:
-    if str(plan.task) != "fl2va":
-        raise ValueError("keyframe target-canvas materials require plan.task='fl2va'")
+    if str(plan.task) not in {"fl2va", "ref2va"}:
+        raise ValueError(
+            "keyframe target-canvas materials require plan.task='fl2va' or 'ref2va'"
+        )
    semantic_indices = tuple(material.frame_index for material in keyframes)
    if semantic_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
        raise ValueError(
-            "fl2va keyframes must use one of the ordered frame_index signatures "
+            "MiniMax H3 keyframes must use one of the ordered frame_index signatures "
            f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {semantic_indices!r}"
        )
    frame_count = plan.shape.get("frame_count")
    if isinstance(frame_count, bool) or not isinstance(frame_count, int):
-        raise ValueError("fl2va keyframe preparation requires an integer frame_count")
+        raise ValueError("keyframe preparation requires an integer frame_count")
    if frame_count <= 1:
-        raise ValueError("fl2va keyframe preparation requires frame_count > 1")
+        raise ValueError("keyframe preparation requires frame_count > 1")
    expected_pixels = tuple(
        frame_count - 1 if index == -1 else index for index in semantic_indices
    )
    resolved_pixels = tuple(material.resolved_frame_index for material in keyframes)
    if resolved_pixels != expected_pixels:
        raise ValueError(
-            "fl2va keyframe resolved_frame_index values disagree with semantic "
+            "keyframe resolved_frame_index values disagree with semantic "
            f"anchors: expected {expected_pixels!r}, got {resolved_pixels!r}"
        )
    return semantic_indices


 def minimax_h3_prepared_keyframes(batch: Any, plan: Any) -> dict[str, Any]:
-    """Resolve + prepare one or two fl2va keyframes once per request.
+    """Resolve + prepare one or two first/last keyframes once per request.

    The target canvas is shared across keyframes and must already be frozen by
    the pre-queue probe/resolve hook.
@@ -154,7 +156,7 @@ def minimax_h3_prepared_keyframes(batch: Any, plan: Any) -> dict[str, Any]:
        cached_images = cached.get("images") or ()
        if cached_indices != semantic_indices or len(cached_images) != len(keyframes):
            raise ValueError(
-                "cached fl2va keyframe preparation disagrees with the resolved plan"
+                "cached keyframe preparation disagrees with the resolved plan"
            )
        return cached

@@ -279,6 +279,8 @@ def minimax_h3_packed_sequence_ref2va_blocks(
    latent_w: int,
    audio_t: int,
    ref_blocks: Sequence[Mapping[str, object]],
+    keyframe_frame_indices: list[int] | tuple[int, ...] | None = None,
+    frame_count: int | None = None,
    audio_channel: int = 2,
    seq_len: int | None = None,
 ) -> dict[str, Any]:
@@ -290,10 +292,12 @@ def minimax_h3_packed_sequence_ref2va_blocks(
    - ``{"kind": "video"|"video_audio", "ref_audio_t": T,
       "latent_t": RT, "latent_h": RH, "latent_w": RW}``

-    Video-bearing blocks pack their audio rows immediately before their video
-    rows; both share the same temporal origin and advance by the longer of the
-    audio and video spans. Standalone audio advances the target origin by its
-    own T, and image blocks advance it by one integer slot.
+    Optional first/last keyframes are packed immediately after text, matching
+    Comfy's hybrid Ref2VA + guide layout. Video-bearing blocks pack their audio
+    rows immediately before their video rows; both share the same temporal
+    origin and advance by the longer of the audio and video spans. Standalone
+    audio advances the target origin by its own T, and image blocks advance it
+    by one integer slot.
    """
    if not isinstance(ref_blocks, Sequence) or isinstance(ref_blocks, (str, bytes)):
        raise ValueError("ref_blocks must be a sequence")
@@ -345,10 +349,19 @@ def minimax_h3_packed_sequence_ref2va_blocks(

    ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
    frame_rows = ph * pw
+    keyframe_indices = _keyframe_cond_frame_indices(
+        include_keyframe_cond=keyframe_frame_indices is not None,
+        keyframe_frame_indices=keyframe_frame_indices,
+    )
+    resolved_keyframe_indices = _resolve_keyframe_frame_indices(
+        keyframe_indices,
+        frame_count=frame_count,
+    )
+    keyframe_rows = len(keyframe_indices) * frame_rows
    video_rows = latent_t * frame_rows
    audio_rows = audio_t * audio_channel
    ref_rows = ref_visual_rows + ref_audio_rows
-    used = text_len + ref_rows + audio_rows + video_rows
+    used = text_len + keyframe_rows + ref_rows + audio_rows + video_rows
    if seq_len is None:
        seq_len = (
            (used + MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT - 1)
@@ -359,7 +372,8 @@ def minimax_h3_packed_sequence_ref2va_blocks(
        raise ValueError(f"seq_len {seq_len} < used rows {used}")

    text_sl = slice(0, text_len)
-    cursor = text_len
+    keyframe_sl = slice(text_len, text_len + keyframe_rows)
+    cursor = keyframe_sl.stop
    block_slices: list[dict[str, object]] = []
    for item in parsed:
        kind = str(item["kind"])
@@ -466,13 +480,31 @@ def minimax_h3_packed_sequence_ref2va_blocks(
    video_g[:, :, 0] = _video_t_grid(latent_t, t_cursor)[:, None]
    video_g[:, :, 1:] = target_frame[None]

+    for block_index, pixel_index in enumerate(resolved_keyframe_indices):
+        sl = slice(
+            keyframe_sl.start + block_index * frame_rows,
+            keyframe_sl.start + (block_index + 1) * frame_rows,
+        )
+        if pixel_index == 0:
+            cond_t = t_cursor
+        elif frame_count is not None and pixel_index == frame_count - 1:
+            cond_t = t_cursor + _temporal_position_span(latent_t) - _FRAME_RESCALE
+        else:
+            raise ValueError(
+                "hybrid ref2va layout only supports first/last keyframe anchors, "
+                f"got resolved frame index {pixel_index}"
+            )
+        g[sl, 0] = cond_t
+        g[sl, 1:] = target_frame
+
+    keyframe_img_pos = _range_for_slice(keyframe_sl)
    target_img_pos = _range_for_slice(video_sl)
    target_audio_pos = _range_for_slice(audio_sl)
-    img_pos = _cat_ranges(ref_img_pos_parts + [target_img_pos])
+    img_pos = _cat_ranges([keyframe_img_pos] + ref_img_pos_parts + [target_img_pos])
    audio_pos = _cat_ranges(ref_audio_pos_parts + [target_audio_pos])

    update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
-    update_mask[ref_visual_rows:] = True
+    update_mask[keyframe_rows + ref_visual_rows :] = True
    audio_update_mask = torch.zeros(audio_pos.shape[0], dtype=torch.bool)
    audio_update_mask[ref_audio_rows:] = True
    text_pos = torch.arange(0, text_len)
@@ -248,16 +248,31 @@ def _validate_conditions(
    return normalized


-def _validate_fl2va_conditions(conditions: Sequence[Mapping[str, Any]]) -> None:
-    """Enforce the public FL contract after per-entry schema validation."""
+def _validate_keyframe_conditions(
+    conditions: Sequence[Mapping[str, Any]], *, task: str
+) -> None:
+    """Enforce the shared first/last-frame contract after schema validation."""

-    frame_indices = tuple(condition.get("frame_index") for condition in conditions)
+    keyframes = [
+        condition
+        for condition in conditions
+        if condition["role"] == MINIMAX_H3_CONDITION_ROLE_KEYFRAME
+    ]
+    frame_indices = tuple(condition.get("frame_index") for condition in keyframes)
    if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
        raise ValueError(
-            "conditions for task 'fl2va' must be one or two ordered "
+            f"conditions for task {task!r} must include one or two ordered "
            "image/keyframe entries with frame_index [0], [-1], or [0, -1], "
            f"got {list(frame_indices)!r}"
        )
+    if task == MINIMAX_H3_TASK_REF2VA and not any(
+        condition["role"] == MINIMAX_H3_CONDITION_ROLE_REFERENCE
+        for condition in conditions
+    ):
+        raise ValueError(
+            "ref2va keyframes require at least one reference condition; "
+            "use task 'fl2va' for keyframe-only generation"
+        )


 def minimax_h3_validate_canonical_request(
@@ -300,11 +315,16 @@ def minimax_h3_validate_canonical_request(
        frame_count=requested_frame_count,
    )
    if profile.task == MINIMAX_H3_TASK_FL2VA:
-        _validate_fl2va_conditions(normalized_conditions)
-    # ref2va accepts ordered material streams containing any mix of
-    # image/audio/video/video_audio references. Type admission is handled by
-    # the task profile; temporal ambiguity is validated later when target
-    # duration is omitted.
+        _validate_keyframe_conditions(normalized_conditions, task=profile.task)
+    elif profile.task == MINIMAX_H3_TASK_REF2VA and any(
+        condition["role"] == MINIMAX_H3_CONDITION_ROLE_KEYFRAME
+        for condition in normalized_conditions
+    ):
+        _validate_keyframe_conditions(normalized_conditions, task=profile.task)
+    # ref2va accepts ordered reference streams and, for hybrid checkpoints,
+    # one first/last keyframe signature. Type admission is handled by the task
+    # profile; temporal ambiguity is validated later when target duration is
+    # omitted.
    if not profile.video_reference_supported:
        for index, cond in enumerate(normalized_conditions):
            if cond["type"] in ("video", "video_audio"):
@@ -276,8 +276,17 @@ def minimax_h3_resolve_plan(canonical: Mapping[str, Any]) -> MiniMaxH3ResolvedPl
        if key not in canonical:
            raise ValueError(f"canonical request missing {key!r}")
    profile = minimax_h3_task_profile(str(canonical["task"]))
-    if profile.task == "fl2va":
-        conditions = canonical["conditions"]
+    conditions = canonical["conditions"]
+    keyframe_conditions = (
+        [
+            condition
+            for condition in conditions
+            if isinstance(condition, Mapping) and condition.get("role") == "keyframe"
+        ]
+        if isinstance(conditions, (list, tuple))
+        else []
+    )
+    if profile.task == "fl2va" or keyframe_conditions:
        signatures = (
            [
                (
@@ -285,10 +294,9 @@ def minimax_h3_resolve_plan(canonical: Mapping[str, Any]) -> MiniMaxH3ResolvedPl
                    condition.get("role"),
                    condition.get("frame_index"),
                )
-                for condition in conditions
+                for condition in keyframe_conditions
            ]
-            if isinstance(conditions, (list, tuple))
-            and all(isinstance(condition, Mapping) for condition in conditions)
+            if all(isinstance(condition, Mapping) for condition in keyframe_conditions)
            else []
        )
        frame_signature = tuple(signature[2] for signature in signatures)
@@ -298,7 +306,8 @@ def minimax_h3_resolve_plan(canonical: Mapping[str, Any]) -> MiniMaxH3ResolvedPl
            or frame_signature not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES
        ):
            raise ValueError(
-                "fl2va ResolvedPlan requires one or two ordered image/keyframe "
+                f"{profile.task} ResolvedPlan requires one or two ordered "
+                "image/keyframe "
                "conditions with frame_index [0], [-1], or [0, -1], got "
                f"{signatures!r}"
            )
@@ -363,10 +372,15 @@ def minimax_h3_resolve_plan(canonical: Mapping[str, Any]) -> MiniMaxH3ResolvedPl
        if rule.audio_tokenizer_encode:
            audio_encode.append(index)

+    qwen_condition_indices = [
+        index
+        for index, condition in enumerate(canonical["conditions"])
+        if profile.task != "ref2va" or condition["role"] == "reference"
+    ]
    encoders = {
        "qwen": {
            "prompt": canonical["prompt"],
-            "ordered_condition_indices": list(range(len(canonical["conditions"]))),
+            "ordered_condition_indices": qwen_condition_indices,
        },
        "visual": visual_encode,
        "audio": audio_encode,
@@ -76,38 +76,42 @@ def minimax_h3_condition_noise_aug(sampling: Any) -> tuple[float, float]:
    return float(imgvid_noise_aug), float(audio_noise_aug)


-def _validate_fl2va_keyframe_payload(plan: Any, keyframe: Any) -> None:
+def _validate_keyframe_payload(plan: Any, keyframe: Any) -> None:
    """Reject stale/middle/reordered keyframe payloads at the DiT sink."""

    task = None if plan is None else str(plan.task)
-    if task != "fl2va":
+    if task not in {"fl2va", "ref2va"}:
        if keyframe is not None:
            raise ValueError(
-                "keyframe condition rows are only valid for plan.task='fl2va'"
+                "keyframe condition rows require plan.task='fl2va' or 'ref2va'"
            )
        return
+    if keyframe is None:
+        if task == "fl2va":
+            raise ValueError("fl2va denoising requires encoded keyframe condition rows")
+        return
    if not isinstance(keyframe, Mapping):
-        raise ValueError("fl2va denoising requires encoded keyframe condition rows")
+        raise ValueError("encoded keyframe condition rows must be a mapping")

    semantic_indices = tuple(keyframe.get("semantic_frame_indices") or ())
    if semantic_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
        raise ValueError(
-            "fl2va denoising requires semantic_frame_indices in "
+            "keyframe denoising requires semantic_frame_indices in "
            f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, "
            f"got {semantic_indices!r}"
        )
    frame_count = keyframe.get("frame_count")
    if isinstance(frame_count, bool) or not isinstance(frame_count, int):
-        raise ValueError("fl2va keyframe payload requires an integer frame_count")
+        raise ValueError("keyframe payload requires an integer frame_count")
    if frame_count <= 1:
-        raise ValueError("fl2va keyframe payload frame_count must be greater than one")
+        raise ValueError("keyframe payload frame_count must be greater than one")
    pixel_indices = keyframe.get("pixel_frame_indices")
    expected_pixel_indices = [
        frame_count - 1 if index == -1 else index for index in semantic_indices
    ]
    if pixel_indices != expected_pixel_indices:
        raise ValueError(
-            "fl2va denoising requires pixel_frame_indices resolved from the "
+            "keyframe denoising requires pixel_frame_indices resolved from the "
            "semantic anchors, "
            f"got {pixel_indices!r} for frame_count={frame_count}"
        )
@@ -119,15 +123,15 @@ def _validate_fl2va_keyframe_payload(plan: Any, keyframe: Any) -> None:
        or any(not isinstance(entry, Mapping) for entry in entries)
    ):
        raise ValueError(
-            "fl2va denoising requires one encoded keyframe per semantic anchor"
+            "keyframe denoising requires one encoded keyframe per semantic anchor"
        )
    if [entry.get("frame_index") for entry in entries] != list(semantic_indices):
-        raise ValueError("fl2va encoded keyframes must remain in semantic anchor order")
+        raise ValueError("encoded keyframes must remain in semantic anchor order")
    if [
        entry.get("resolved_frame_index") for entry in entries
    ] != expected_pixel_indices:
        raise ValueError(
-            "fl2va encoded keyframes must carry matching resolved_frame_index values"
+            "encoded keyframes must carry matching resolved_frame_index values"
        )

    latent_h = int(keyframe.get("latent_h") or 0)
@@ -142,7 +146,7 @@ def _validate_fl2va_keyframe_payload(plan: Any, keyframe: Any) -> None:
    ):
        actual_rows = None if not isinstance(rows, torch.Tensor) else int(rows.shape[0])
        raise ValueError(
-            "fl2va encoded keyframe rows do not match target-canvas blocks: "
+            "encoded keyframe rows do not match target-canvas blocks: "
            f"expected={expected_rows}, actual={actual_rows}"
        )

@@ -155,8 +159,26 @@ def _imgvid_condition_shapes(
 ) -> list[tuple[int, int, int]]:
    """Return visual-condition ``(T,H,W)`` in packed anchor-row order."""

+    shapes = []
+    if isinstance(keyframe, Mapping):
+        entries = keyframe.get("keyframes")
+        if isinstance(entries, list) and entries:
+            shapes.extend(
+                (1, int(entry["latent_h"]), int(entry["latent_w"])) for entry in entries
+            )
+        else:
+            latent_h = int(keyframe["latent_h"])
+            latent_w = int(keyframe["latent_w"])
+            frame_rows = (latent_h // 2) * (latent_w // 2)
+            rows = keyframe["rows"]
+            if frame_rows <= 0 or int(rows.shape[0]) % frame_rows:
+                raise ValueError(
+                    "legacy keyframe rows cannot be split into visual-condition frames"
+                )
+            shapes.extend(
+                [(1, latent_h, latent_w)] * (int(rows.shape[0]) // frame_rows)
+            )
    if ref2va_blocks is not None:
-        shapes = []
        for block in ref2va_blocks:
            kind = str(block["kind"])
            if kind == "image":
@@ -176,23 +198,7 @@ def _imgvid_condition_shapes(
        # supplied above; reaching here indicates an upstream bug.
        raise ValueError("ref2va visual-condition shapes require ordered blocks")

-    if not isinstance(keyframe, Mapping):
-        return []
-    entries = keyframe.get("keyframes")
-    if isinstance(entries, list) and entries:
-        return [
-            (1, int(entry["latent_h"]), int(entry["latent_w"])) for entry in entries
-        ]
-
-    latent_h = int(keyframe["latent_h"])
-    latent_w = int(keyframe["latent_w"])
-    frame_rows = (latent_h // 2) * (latent_w // 2)
-    rows = keyframe["rows"]
-    if frame_rows <= 0 or int(rows.shape[0]) % frame_rows:
-        raise ValueError(
-            "legacy keyframe rows cannot be split into visual-condition frames"
-        )
-    return [(1, latent_h, latent_w)] * (int(rows.shape[0]) // frame_rows)
+    return shapes


 def _ref2va_payload_entry(
@@ -246,6 +252,8 @@ def _ref2va_ordered_blocks_and_rows(
    for material in plan.materials:
        chain = str(material.material_chain)
        condition_index = int(material.condition_index)
+        if chain == "image.target_canvas":
+            continue
        if chain == "image.reference_preserve":
            entry = _ref2va_payload_entry(
                ref_image,
@@ -789,8 +797,7 @@ class _FullLoopContext:
 def _resolve_full_loop_context(batch: Req) -> _FullLoopContext:
    """Read/validate extras and the denoise state into a loop context.

-    Enforces the task-payload exclusivity rules (keyframe vs reference
-    exclusivity) and cross-checks the resolved latent dims.
+    Validates task payloads and cross-checks the resolved latent dims.
    """
    from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
        MINIMAX_H3_DENOISE_STATE_EXTRA_KEY,
@@ -819,9 +826,7 @@ def _resolve_full_loop_context(batch: Req) -> _FullLoopContext:
        or ctx.ref_audio is not None
        or ctx.ref_video is not None
    )
-    if ctx.is_ref2va and ctx.keyframe is not None:
-        raise ValueError("keyframe and reference extras are mutually exclusive")
-    _validate_fl2va_keyframe_payload(ctx.plan, ctx.keyframe)
+    _validate_keyframe_payload(ctx.plan, ctx.keyframe)

    ctx.latent_t = int(ctx.state["latent_t"])
    ctx.latent_h = int(ctx.state["latent_h"])
@@ -844,7 +849,7 @@ def _assemble_condition_rows(ctx: _FullLoopContext) -> None:
                "ref2va reference extras require a resolved plan; "
                "plan-less ref2va requests are unsupported"
            )
-        ctx.ref2va_positive_blocks, ctx.cond_rows, ctx.audio_ref_rows = (
+        ctx.ref2va_positive_blocks, ref_rows, ctx.audio_ref_rows = (
            _ref2va_ordered_blocks_and_rows(
                plan=ctx.plan,
                ref_image=ctx.ref_image,
@@ -852,6 +857,12 @@ def _assemble_condition_rows(ctx: _FullLoopContext) -> None:
                ref_video=ctx.ref_video,
            )
        )
+        condition_parts = []
+        if ctx.keyframe is not None:
+            condition_parts.append(ctx.keyframe["rows"])
+        if ref_rows is not None:
+            condition_parts.append(ref_rows)
+        ctx.cond_rows = _cat_optional(condition_parts)
        ctx.include_cond = ctx.cond_rows is not None
    else:
        ctx.cond_rows = ctx.keyframe["rows"] if ctx.include_cond else None
@@ -885,6 +896,8 @@ def _build_packed_layout(
            latent_w=ctx.latent_w,
            audio_t=ctx.audio_t,
            ref_blocks=ctx.ref2va_positive_blocks,
+            keyframe_frame_indices=ctx.keyframe_frame_indices,
+            frame_count=ctx.keyframe_frame_count,
        )
    else:
        packed = minimax_h3_packed_sequence(
@@ -245,11 +245,11 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
        keyframes = [
            m for m in plan.materials if m.material_chain == "image.target_canvas"
        ]
-        if plan.task == "fl2va":
+        if plan.task in {"fl2va", "ref2va"} and keyframes:
            frame_indices = tuple(material.frame_index for material in keyframes)
            if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
                raise ValueError(
-                    "fl2va text encoding requires an ordered keyframe signature "
+                    "MiniMax H3 text encoding requires an ordered keyframe signature "
                    f"in {MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got "
                    f"{frame_indices!r}"
                )
@@ -360,7 +360,8 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):

        Per condition in order — image i: '<Picture i>: ' label +
        vision block (prepared reference image); audio j: '<Audio j>: ' label
-        only — then the verbatim prompt.
+        only — then the verbatim prompt. Hybrid keyframes are deliberately
+        omitted: they are guide latents appended after reference presentation.
        """
        from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.presentation import (
            minimax_h3_ref2va_presentation,
@@ -407,6 +408,8 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
        has_image = False
        has_video = False
        for material in plan.materials:
+            if material.material_chain == "image.target_canvas":
+                continue
            if material.material_chain == "image.reference_preserve":
                counters["image"] += 1
                condition_labels.append(("image", counters["image"]))
@@ -166,7 +166,6 @@ class MiniMaxH3VisualEncodingStage(ConditionEncodingStage):
        """Direct keyframe encode: seeded sampled encode_images ->
        normalized [n,96] cond rows in batch.extra."""
        from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.keyframe_encoding import (
-            minimax_h3_encode_keyframe_cond_rows,
            minimax_h3_scoped_encode_fp32,
        )

@@ -177,13 +176,13 @@ class MiniMaxH3VisualEncodingStage(ConditionEncodingStage):
            for material in materials
            if material.material_chain == "image.target_canvas"
        ]
-        if str(plan.task) == "fl2va":
+        if keyframe_materials and str(plan.task) in {"fl2va", "ref2va"}:
            frame_indices = tuple(
                material.frame_index for material in keyframe_materials
            )
            if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
                raise ValueError(
-                    "fl2va visual encoding requires an ordered keyframe signature "
+                    "MiniMax H3 visual encoding requires an ordered keyframe signature "
                    f"in {MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got "
                    f"{frame_indices!r}"
                )
@@ -191,37 +190,39 @@ class MiniMaxH3VisualEncodingStage(ConditionEncodingStage):
            raise ValueError(
                f"task {plan.task!r} cannot carry image.target_canvas materials"
            )
-        if MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY in batch.extra:
-            return
-        if chains == {"image.reference_preserve"}:
-            with minimax_h3_scoped_encode_fp32(self.video_vae):
-                self._encode_reference_image(batch, plan)
-            return
        video_chains = {
            "video.reference_preserve",
            "video_audio.reference_preserve",
        }
-        if chains and chains <= {"image.reference_preserve", *video_chains}:
-            # One VAE dtype toggle for both encodes below, not one each.
-            with minimax_h3_scoped_encode_fp32(self.video_vae):
-                if "image.reference_preserve" in chains:
-                    self._encode_reference_image(batch, plan)
-                if chains & video_chains:
-                    self._encode_reference_video(batch, plan)
-            return
-        unsupported = [
-            m.material_chain
-            for m in materials
-            if m.material_chain != "image.target_canvas"
-        ]
+        supported_chains = {
+            "image.target_canvas",
+            "image.reference_preserve",
+            *video_chains,
+        }
+        unsupported = sorted(chains - supported_chains)
        if unsupported:
            raise NotImplementedError(
-                "MiniMaxH3VisualEncodingStage direct encode only supports "
-                f"image.target_canvas / image.reference_preserve, got {unsupported}"
+                "MiniMaxH3VisualEncodingStage cannot encode material chains "
+                f"{unsupported}"
            )
+        # One VAE dtype toggle for every visual condition in the request.
+        with minimax_h3_scoped_encode_fp32(self.video_vae):
+            if keyframe_materials:
+                self._encode_target_keyframes(batch, plan)
+            if "image.reference_preserve" in chains:
+                self._encode_reference_image(batch, plan)
+            if chains & video_chains:
+                self._encode_reference_video(batch, plan)
+
+    def _encode_target_keyframes(self, batch: Req, plan) -> None:
+        if MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY in batch.extra:
+            return
        from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.canvas import (
            minimax_h3_prepared_keyframes,
        )
+        from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.keyframe_encoding import (
+            minimax_h3_encode_keyframe_cond_rows,
+        )

        # Parallel tiling gives each replicated rank complete tiles, then gathers
        # them before the seeded posterior sample.
@@ -231,38 +232,34 @@ class MiniMaxH3VisualEncodingStage(ConditionEncodingStage):
            prepared.get("images") or ()
        ) != len(prepared_indices):
            raise ValueError(
-                "fl2va visual preparation requires one or two ordered images "
+                "keyframe visual preparation requires one or two ordered images "
                "with a supported semantic_frame_indices signature"
            )
        encoded = []
        rows_list = []
-        # One VAE dtype toggle for the whole signature (up to two keyframes),
-        # not one per keyframe.
-        with minimax_h3_scoped_encode_fp32(self.video_vae):
-            for item in prepared["images"]:
-                image = item["image"]
-                width, height = item["canvas_width"], item["canvas_height"]
-                # The encode sampling seed is pinned at 42 (the VAE sample
-                # seed is part of the contract), independent of the
-                # request seed.
-                rows = minimax_h3_encode_keyframe_cond_rows(
-                    self.video_vae,
-                    image,
-                    self.vae_arch_config,
-                )
-                encoded.append(
-                    {
-                        "rows": rows,
-                        "latent_h": height // 16,
-                        "latent_w": width // 16,
-                        "canvas_height": height,
-                        "canvas_width": width,
-                        "frame_index": item.get("frame_index"),
-                        "resolved_frame_index": item.get("resolved_frame_index"),
-                        "condition_index": item.get("condition_index"),
-                    }
-                )
-                rows_list.append(rows)
+        for item in prepared["images"]:
+            image = item["image"]
+            width, height = item["canvas_width"], item["canvas_height"]
+            # The encode sampling seed is pinned at 42 (the VAE sample
+            # seed is part of the contract), independent of the request seed.
+            rows = minimax_h3_encode_keyframe_cond_rows(
+                self.video_vae,
+                image,
+                self.vae_arch_config,
+            )
+            encoded.append(
+                {
+                    "rows": rows,
+                    "latent_h": height // 16,
+                    "latent_w": width // 16,
+                    "canvas_height": height,
+                    "canvas_width": width,
+                    "frame_index": item.get("frame_index"),
+                    "resolved_frame_index": item.get("resolved_frame_index"),
+                    "condition_index": item.get("condition_index"),
+                }
+            )
+            rows_list.append(rows)
        rows = rows_list[0] if len(rows_list) == 1 else torch.cat(rows_list, dim=0)
        first = encoded[0]
        batch.extra[MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY] = {
@@ -6,6 +6,9 @@ tasks (t2va / fl2va / ref2va). One row per task; stages and the request
 projector consume rows instead of branching on task names.

 Design summary: keyframes bind target geometry; references remain independent.
+Ref2VA also admits keyframes for hybrid checkpoints, matching the released
+Comfy workflow where reference presentation is encoded before guide latents
+are appended.
 """

 from __future__ import annotations
@@ -184,6 +187,13 @@ MINIMAX_H3_TASK_PROFILES: dict[str, MiniMaxH3TaskProfile] = {
        task=MINIMAX_H3_TASK_REF2VA,
        conditions_required=True,
        condition_rules=(
+            MiniMaxH3ConditionRule(
+                role=MINIMAX_H3_CONDITION_ROLE_KEYFRAME,
+                condition_type="image",
+                material_chain="image.target_canvas",
+                requires_frame_index=True,
+                visual_tokenizer_encode=True,
+            ),
            MiniMaxH3ConditionRule(
                role=MINIMAX_H3_CONDITION_ROLE_REFERENCE,
                condition_type="image",
@@ -246,11 +256,11 @@ def _validate_profiles() -> None:
    for task, profile in MINIMAX_H3_TASK_PROFILES.items():
        if profile.task != task:
            raise ValueError(f"profile key/task mismatch: {task} vs {profile.task}")
-        roles = {rule.role for rule in profile.condition_rules}
-        if len(roles) > 1:
-            raise ValueError(
-                f"task {task}: condition roles must not mix, got {sorted(roles)}"
-            )
+        rule_keys = [
+            (rule.role, rule.condition_type) for rule in profile.condition_rules
+        ]
+        if len(rule_keys) != len(set(rule_keys)):
+            raise ValueError(f"task {task}: condition rules must be unique")
        if profile.min_condition_count is not None and profile.min_condition_count <= 0:
            raise ValueError(f"task {task}: min_condition_count must be positive")
        if profile.max_condition_count is not None and profile.max_condition_count <= 0:
@@ -112,6 +112,26 @@ TARGET = {
                "video_audio.reference_preserve",
            ],
        ),
+        (
+            "ref2va",
+            [
+                {
+                    "type": "image",
+                    "uri": "file:///first.png",
+                    "role": "keyframe",
+                    "frame_index": 0,
+                },
+                {
+                    "type": "image",
+                    "uri": "file:///subject.png",
+                    "role": "reference",
+                },
+            ],
+            "ref2va",
+            [0, 1],
+            [],
+            ["image.target_canvas", "image.reference_preserve"],
+        ),
    ],
 )
 def test_public_tasks_resolve_to_exact_partition_and_encoder_plan(
@@ -132,11 +152,35 @@ def test_public_tasks_resolve_to_exact_partition_and_encoder_plan(
    assert plan.encoders["audio"] == audio
    assert [material.material_chain for material in plan.materials] == chains
    if task == "ref2va":
-        assert plan.materials[1].start_time_seconds == 12.5
+        assert plan.encoders["qwen"]["ordered_condition_indices"] == [
+            index
+            for index, condition in enumerate(conditions)
+            if condition["role"] == "reference"
+        ]
+    for index, condition in enumerate(conditions):
+        if condition.get("start_time_seconds") is not None:
+            assert plan.materials[index].start_time_seconds == 12.5
    assert plan.shape["frame_count"] == 124
    assert plan.shape["video_latent_t"] == 37


+def test_ref2va_rejects_keyframes_without_a_reference():
+    with pytest.raises(ValueError, match="at least one reference"):
+        minimax_h3_validate_canonical_request(
+            task="ref2va",
+            prompt="contract",
+            conditions=[
+                {
+                    "type": "image",
+                    "uri": "file:///first.png",
+                    "role": "keyframe",
+                    "frame_index": 0,
+                }
+            ],
+            target=TARGET,
+        )
+
+
@pytest.mark.parametrize(
    ("partition", "tasks"),
    [("fl2va", ["t2va", "fl2va"]), ("ref2va", ["ref2va"])],
@@ -154,3 +154,31 @@ class TestMiniMaxH3PackedSequence(unittest.TestCase):
        target_video_t0 = built["img_position_ids"][built["img_pos"][12], 0]
        target_audio_t0 = built["img_position_ids"][built["audio_pos"][8], 0]
        self.assertEqual(float(target_audio_t0), float(target_video_t0))
+
+    def test_ref2va_hybrid_packs_keyframes_before_references(self):
+        built = minimax_h3_packed_sequence_ref2va_blocks(
+            text_len=5,
+            latent_t=2,
+            latent_h=4,
+            latent_w=4,
+            audio_t=5,
+            ref_blocks=[{"kind": "image", "latent_h": 4, "latent_w": 4}],
+            keyframe_frame_indices=[0, -1],
+            frame_count=5,
+        )
+
+        frame_rows = 4
+        frozen_rows = 2 * frame_rows + frame_rows
+        self.assertEqual(int((~built["update_mask"]).sum()), frozen_rows)
+        keyframe_positions = built["img_pos"][: 2 * frame_rows].reshape(2, -1)
+        keyframe_times = [
+            float(built["img_position_ids"][positions, 0].unique().item())
+            for positions in keyframe_positions
+        ]
+        target_origin = 6.0  # text origin 5 + one reference-image time slot
+        self.assertEqual(keyframe_times[0], target_origin)
+        self.assertAlmostEqual(
+            keyframe_times[1], target_origin + 4 * (5.0 / 3.0), places=12
+        )
+        target_t0 = built["img_position_ids"][built["img_pos"][frozen_rows], 0]
+        self.assertEqual(float(target_t0), target_origin)