# ПЕРВЫЙ ФАЙЛ

# SPDX-License-Identifier: GPL-3.0-or-later
"""Runtime patches for MiniMax H3 continuation conditioning.

Inspired by NikoDemon80/ComfyUI-H3-Motion-Context.

The patch is deliberately narrow:
  * allow motion-context keyframes at arbitrary pixel-frame positions by
    rewriting only PackedLayout.position_ids after stock construction;
  * keep all keyframe + reference visual latents in payload order;
  * optionally move one marked audio-reference block onto the target timeline.

No ComfyUI files are modified on disk.
"""
from __future__ import annotations

import logging
from typing import Any, Iterable

import torch

import comfy.ldm.minimax.model as mm
import comfy.model_base as model_base

LOG = logging.getLogger("h3_continuation_context")

MC_INDEX = "h3_continuation_pixel_index"
MC_AUDIO_END = "h3_continuation_audio_end_frame"

_layout_original = None
_payload_original = None
_layout_applied = False
_payload_applied = False


def _ref_cursor_advance(refs: Iterable[dict] | None) -> float:
    """Reproduce PackedLayout's reference cursor advance exactly."""
    if not refs:
        return 0.0
    cursor = 0.0
    for blk in refs:
        kind = blk.get("kind")
        if kind == "image":
            cursor += 1.0
        elif kind == "audio":
            cursor += float(blk.get("ref_audio_t", 0))
        elif kind in ("video", "video_audio"):
            rt = float(blk.get("ref_audio_t", 0))
            vt = int(blk.get("latent_t", 0))
            cursor += max(rt, sum(mm._video_t_spans(vt)))
    return cursor


def _cond_time(text_len: int, latent_t: int, frame_count: int | None, pixel_index: float) -> float:
    """Time coordinate for a visual condition at a pixel-frame index."""
    p = float(pixel_index)
    if p == 0.0:
        return float(text_len)
    if frame_count is not None and p == float(frame_count - 1):
        return float(text_len) + sum(mm._video_t_spans(latent_t)) - mm.FRAME_RESCALE
    return float(text_len) + mm.FRAME_RESCALE * p


def _fix_visual_positions(layout, text_len, latent_t, frame_count, keyframes, refs=None):
    """Align all keyframe condition rows to the target timeline.

    Motion-context keyframes carry MC_INDEX. Stock first/last keyframes are
    left at their stock coordinate when there are no refs; when refs coexist,
    they receive the same reference-cursor offset as the target so they remain
    aligned instead of sliding backward relative to target frame 0.
    """
    cond_spans = [(a, b) for a, b, kind in layout.segments if kind == "cond"]
    if len(cond_spans) != len(keyframes):
        raise RuntimeError(
            f"h3_continuation_context: expected {len(keyframes)} cond segments, "
            f"layout has {len(cond_spans)}"
        )

    offset = _ref_cursor_advance(refs)
    for (a, b), kf in zip(cond_spans, keyframes):
        custom = kf.get(MC_INDEX)
        if custom is not None:
            t = _cond_time(text_len, latent_t, frame_count, custom) + offset
            layout.position_ids[a:b, 0] = t
        elif offset:
            # Stock first/last keyframe + refs: preserve its stock position
            # relative to the target by shifting it by the ref cursor advance.
            layout.position_ids[a:b, 0] += offset


def _ref_segment_spans(layout, refs):
    """Map each ref object to the PackedLayout segment(s) it created.

    Returns list[(ref, [(a,b,kind), ...])]. This avoids selecting audio rows by
    a global coordinate range and therefore works when normal REF2VA image/video
    refs coexist with the continuation audio ref.
    """
    segs = [(a, b, kind) for a, b, kind in layout.segments if kind in ("ref_img", "ref_audio")]
    idx = 0
    mapped = []
    for ref in refs or []:
        kind = ref.get("kind")
        expect: list[str] = []
        if kind == "image":
            expect = ["ref_img"]
        elif kind == "audio":
            if int(ref.get("ref_audio_t", 0)) > 0:
                expect = ["ref_audio"]
        elif kind in ("video", "video_audio"):
            if int(ref.get("ref_audio_t", 0)) > 0:
                expect.append("ref_audio")
            expect.append("ref_img")
        owned = []
        for wanted in expect:
            if idx >= len(segs):
                raise RuntimeError("h3_continuation_context: ref segment table ended early")
            seg = segs[idx]
            idx += 1
            if seg[2] != wanted:
                raise RuntimeError(
                    f"h3_continuation_context: expected ref segment {wanted}, got {seg[2]}"
                )
            owned.append(seg)
        mapped.append((ref, owned))
    if idx != len(segs):
        raise RuntimeError(
            f"h3_continuation_context: layout has {len(segs) - idx} unmapped ref segment(s)"
        )
    return mapped


def _fix_audio_positions(layout, text_len, refs):
    """Move marked continuation audio from reference time onto target time."""
    if not refs:
        return
    marked = [r for r in refs if r.get(MC_AUDIO_END) is not None]
    if not marked:
        return
    if len(marked) != 1:
        raise RuntimeError(
            f"h3_continuation_context: expected one marked continuation audio ref, got {len(marked)}"
        )
    mark = marked[0]
    if mark.get("kind") != "audio":
        raise RuntimeError("h3_continuation_context: timeline marker may only be used on an audio ref")
    rt = int(mark.get("ref_audio_t", 0))
    if rt <= 0:
        return

    owned = None
    for ref, spans in _ref_segment_spans(layout, refs):
        if ref is mark:
            owned = spans
            break
    if not owned or len(owned) != 1 or owned[0][2] != "ref_audio":
        raise RuntimeError("h3_continuation_context: could not locate marked audio ref segment")

    a, b, _ = owned[0]
    t = layout.position_ids[a:b, 0]
    old_start = float(t.min())
    target_origin = float(text_len) + _ref_cursor_advance(refs)
    end_frame = float(mark[MC_AUDIO_END])
    new_start = target_origin + mm.FRAME_RESCALE * end_frame - float(rt)
    layout.position_ids[a:b, 0] = t + (new_start - old_start)


def _patched_layout_init(self, text_len, latent_t, latent_h, latent_w, audio_t,
                         keyframes=None, refs=None, frame_count=None):
    _layout_original(
        self, text_len, latent_t, latent_h, latent_w, audio_t,
        keyframes=keyframes, refs=refs, frame_count=frame_count,
    )
    if keyframes:
        has_custom = any(kf.get(MC_INDEX) is not None for kf in keyframes)
        # Ref coexistence also needs stock first/last anchors shifted with target.
        if has_custom or refs:
            _fix_visual_positions(self, text_len, latent_t, frame_count, keyframes, refs)
    if refs and any(r.get(MC_AUDIO_END) is not None for r in refs):
        _fix_audio_positions(self, text_len, refs)


_patched_layout_init._h3_continuation_context_patch = True


def _layout_self_test():
    text_len, latent_t, lh, lw, audio_t = 7, 7, 22, 38, 37
    frame_count = sum(mm.FRAME_PER_TOKEN[k % 5] for k in range(latent_t))

    # Endpoint equivalence: our general mechanism must reproduce stock exactly.
    stock = [
        {"resolved_frame_index": 0},
        {"resolved_frame_index": frame_count - 1},
    ]
    ours = [
        {"resolved_frame_index": 0, MC_INDEX: 0},
        {"resolved_frame_index": 0, MC_INDEX: frame_count - 1},
    ]
    a = mm.PackedLayout.__new__(mm.PackedLayout)
    _layout_original(a, text_len, latent_t, lh, lw, audio_t,
                     keyframes=stock, frame_count=frame_count)
    b = mm.PackedLayout.__new__(mm.PackedLayout)
    _layout_original(b, text_len, latent_t, lh, lw, audio_t,
                     keyframes=ours, frame_count=frame_count)
    _fix_visual_positions(b, text_len, latent_t, frame_count, ours)
    if not torch.equal(a.position_ids, b.position_ids):
        raise RuntimeError("endpoint position self-test failed")

    # Interior anchors must be strictly ordered and follow FRAME_RESCALE * frame.
    run_indices = [0, 1, 5, 9]
    run = [{"resolved_frame_index": 0, MC_INDEX: p} for p in run_indices]
    c = mm.PackedLayout.__new__(mm.PackedLayout)
    _layout_original(c, text_len, latent_t, lh, lw, audio_t,
                     keyframes=run, frame_count=frame_count)
    _fix_visual_positions(c, text_len, latent_t, frame_count, run)
    spans = [(x, y) for x, y, k in c.segments if k == "cond"]
    got = [float(c.position_ids[x, 0]) for x, _ in spans]
    want = [_cond_time(text_len, latent_t, frame_count, p) for p in run_indices]
    if any(abs(g - w) > 1e-9 for g, w in zip(got, want)):
        raise RuntimeError(f"interior position self-test failed: {got} vs {want}")

    # Refs shift target origin. Both custom and stock anchors must shift equally.
    refs = [{"kind": "image", "latent_h": 8, "latent_w": 8}]
    d = mm.PackedLayout.__new__(mm.PackedLayout)
    _layout_original(d, text_len, latent_t, lh, lw, audio_t,
                     keyframes=run, refs=refs, frame_count=frame_count)
    _fix_visual_positions(d, text_len, latent_t, frame_count, run, refs)
    shifted = [float(d.position_ids[x, 0]) for x, _ in [(x,y) for x,y,k in d.segments if k=="cond"]]
    off = _ref_cursor_advance(refs)
    if any(abs((s - g) - off) > 1e-9 for s, g in zip(shifted, got)):
        raise RuntimeError("reference cursor compensation self-test failed")

    # Audio marker may coexist with ordinary image refs.
    refs2 = [
        {"kind": "image", "latent_h": 8, "latent_w": 8},
        {"kind": "audio", "ref_audio_t": 8, MC_AUDIO_END: 22.0},
        {"kind": "image", "latent_h": 8, "latent_w": 8},
    ]
    e = mm.PackedLayout.__new__(mm.PackedLayout)
    _layout_original(e, text_len, latent_t, lh, lw, audio_t,
                     keyframes=run, refs=refs2, frame_count=frame_count)
    before = e.position_ids.clone()
    _fix_visual_positions(e, text_len, latent_t, frame_count, run, refs2)
    after_visual = e.position_ids.clone()
    _fix_audio_positions(e, text_len, refs2)
    mapped = _ref_segment_spans(e, refs2)
    a0, b0, _ = mapped[1][1][0]
    changed = (after_visual[:, 0] != e.position_ids[:, 0]).nonzero().flatten().tolist()
    if changed != list(range(a0, b0)):
        raise RuntimeError("audio timeline self-test moved unexpected rows")


def apply_layout_patch() -> bool:
    global _layout_original, _layout_applied
    if _layout_applied:
        return True
    current = getattr(mm.PackedLayout, "__init__", None)
    if current is None:
        LOG.warning("PackedLayout.__init__ not found; continuation patch disabled")
        return False
    if getattr(current, "_h3_continuation_context_patch", False):
        _layout_applied = True
        return True

    # Avoid silently stacking on a foreign layout patch. This is especially
    # important if the original H3-Motion-Context package is installed too.
    if getattr(current, "__module__", "") != mm.__name__:
        LOG.warning(
            "PackedLayout.__init__ is already patched by %s; refusing to stack "
            "H3 continuation layout patches. Disable the other H3 layout patch.",
            getattr(current, "__module__", "unknown"),
        )
        return False

    _layout_original = current
    try:
        _layout_self_test()
    except Exception as exc:
        _layout_original = None
        LOG.exception("H3 continuation layout self-test failed: %s", exc)
        return False
    mm.PackedLayout.__init__ = _patched_layout_init
    _layout_applied = True
    LOG.info("h3_continuation_context: interior visual anchors enabled")
    return True


def _patched_extra_conds(self, **kwargs: Any):
    out = _payload_original(self, **kwargs)
    keyframes = kwargs.get("minimax_keyframes")
    refs = kwargs.get("minimax_refs")
    if not keyframes or not refs:
        return out
    holder = out.get("minimax_payload")
    payload = getattr(holder, "cond", None) if holder is not None else None
    if not isinstance(payload, dict):
        LOG.warning("could not reach MiniMax H3 payload; keyframe/ref merge skipped")
        return out

    payload["cond_video_latents"] = (
        [kf["latent"] for kf in keyframes if kf.get("latent") is not None]
        + [r["latent"] for r in refs if r.get("latent") is not None]
    )
    payload["cond_audio_latents"] = [
        r["audio_latent"] for r in refs if r.get("audio_latent") is not None
    ]
    fc = kwargs.get("minimax_frame_count")
    if fc is not None:
        payload["frame_count"] = fc
    return out


_patched_extra_conds._h3_continuation_payload_patch = True


def apply_payload_patch() -> bool:
    global _payload_original, _payload_applied
    if _payload_applied:
        return True
    cls = getattr(model_base, "MiniMaxH3", None)
    if cls is None or not hasattr(cls, "extra_conds"):
        LOG.warning("MiniMaxH3.extra_conds not found; keyframe/ref merge disabled")
        return False
    current = cls.extra_conds
    if getattr(current, "_h3_continuation_payload_patch", False):
        _payload_applied = True
        return True
    _payload_original = current
    cls.extra_conds = _patched_extra_conds
    _payload_applied = True
    LOG.info("h3_continuation_context: keyframe/reference coexistence enabled")
    return True


def layout_patch_applied() -> bool:
    return _layout_applied


def payload_patch_applied() -> bool:
    return _payload_applied


# ВТОРОЙ ФАЙЛ

# SPDX-License-Identifier: GPL-3.0-or-later
"""MiniMax H3 continuation-context nodes.

The main node turns the tail of a previous clip into H3's native never-denoised
visual conditioning rows at consecutive positions on the *new* clip timeline.
This is fundamentally different from writing values into the target latent.
"""
from __future__ import annotations

import logging
from typing import Any

import torch

import comfy.nested_tensor
import comfy.utils

from .h3_patches import (
    MC_AUDIO_END,
    MC_INDEX,
    layout_patch_applied,
    payload_patch_applied,
)

LOG = logging.getLogger("h3_continuation_context")

FPS = 24
AUDIO_HZ = 40.0
FRAME_RESCALE = 5.0 / 3.0
FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
VIDEO_RUN_GRID = (39, 22, 5, 1)


def _pixel_frames(latent_t: int) -> int:
    return sum(FRAME_PER_TOKEN[k % 5] for k in range(int(latent_t)))


def _step_offsets(latent_t: int) -> list[int]:
    out, acc = [], 0
    for k in range(int(latent_t)):
        out.append(acc)
        acc += FRAME_PER_TOKEN[k % 5]
    return out


def _active_run(requested: int, available: int, allow_one: bool = True) -> int:
    n = min(int(requested), int(available))
    grid = VIDEO_RUN_GRID if allow_one else VIDEO_RUN_GRID[:-1]
    for g in grid:
        if g <= n:
            return g
    raise ValueError(
        "h3_continuation_context: not enough context frames for a valid run; "
        f"need at least {grid[-1]}, have {available}"
    )


def _streams(latent: dict) -> list[torch.Tensor]:
    if not isinstance(latent, dict) or "samples" not in latent:
        raise TypeError("h3_continuation_context: expected a ComfyUI LATENT dict")
    samples = latent["samples"]
    if hasattr(samples, "unbind"):
        parts = list(samples.unbind())
    elif isinstance(samples, (tuple, list)):
        parts = list(samples)
    else:
        raise TypeError(
            "h3_continuation_context: expected H3 NestedTensor(video,audio), got "
            + type(samples).__name__
        )
    if len(parts) < 2:
        raise ValueError("h3_continuation_context: H3 latent has fewer than 2 streams")
    return parts


def _video(latent: dict) -> torch.Tensor:
    z = _streams(latent)[0]
    if z.ndim == 4:
        z = z.unsqueeze(0)
    if z.ndim != 5 or z.shape[1] != 24:
        raise ValueError(
            f"h3_continuation_context: expected video [B,24,T,H,W], got {tuple(z.shape)}"
        )
    return z


def _audio(latent: dict) -> torch.Tensor:
    z = _streams(latent)[1]
    if z.ndim == 3:
        z = z.unsqueeze(0)
    if z.ndim != 4 or z.shape[1] != 32 or z.shape[2] != 2:
        raise ValueError(
            f"h3_continuation_context: expected audio [B,32,2,T], got {tuple(z.shape)}"
        )
    return z


def _resize_frames(images: torch.Tensor, width: int, height: int, crop: str) -> torch.Tensor:
    if images.ndim != 4:
        raise ValueError(f"h3_continuation_context: IMAGE must be [N,H,W,C], got {tuple(images.shape)}")
    x = images[..., :3].movedim(-1, 1)
    x = comfy.utils.common_upscale(x, int(width), int(height), "lanczos", crop)
    return x.movedim(1, -1)


def _encode_video_context(vae, frames: torch.Tensor, target_video: torch.Tensor, requested: int, crop: str):
    if vae is None:
        raise ValueError("h3_continuation_context: video_source=frames requires vae")
    available = int(frames.shape[0])
    run = _active_run(requested, available, allow_one=True)
    if run != min(int(requested), available):
        LOG.warning(
            "h3_continuation_context: context length snapped down to %d frames "
            "(usable video-VAE runs: 1, 5, 22, 39)", run,
        )
    width = int(target_video.shape[-1]) * 16
    height = int(target_video.shape[-2]) * 16
    tail = _resize_frames(frames[available - run:], width, height, crop)
    enc = vae.encode(tail)
    if getattr(enc, "ndim", 0) != 5 or int(enc.shape[1]) != 24:
        raise ValueError(
            "h3_continuation_context: VAE video encode returned "
            f"{tuple(getattr(enc, 'shape', ()))}, expected [B,24,T,H,W]"
        )
    if tuple(enc.shape[-2:]) != tuple(target_video.shape[-2:]):
        raise ValueError(
            "h3_continuation_context: encoded context spatial grid does not match target: "
            f"context={tuple(enc.shape[-2:])}, target={tuple(target_video.shape[-2:])}"
        )
    steps = int(enc.shape[2])
    covered = _pixel_frames(steps)
    if covered != run:
        raise RuntimeError(
            "h3_continuation_context: H3 VAE temporal grid changed: "
            f"{run} RGB frames encoded to {steps} latent steps covering {covered} frames"
        )
    blocks = [enc[:, :, k:k + 1] for k in range(steps)]
    return blocks, _step_offsets(steps), run, f"frames/VAE {run}f->{steps}t"


def _latent_steps_for_run(run: int) -> int:
    # Valid continuation runs end on the same 5-token temporal phase as a full H3 clip.
    mapping = {5: 2, 22: 7, 39: 12}
    if run not in mapping:
        raise ValueError(
            "h3_continuation_context: direct latent video context supports 5, 22 or 39 frames"
        )
    return mapping[run]


def _encode_latent_context(context_latent: dict, target_video: torch.Tensor, requested: int):
    prev = _video(context_latent)
    if tuple(prev.shape[-2:]) != tuple(target_video.shape[-2:]):
        raise ValueError(
            "h3_continuation_context: previous/target video latent spatial mismatch: "
            f"previous={tuple(prev.shape[-2:])}, target={tuple(target_video.shape[-2:])}"
        )
    available_frames = _pixel_frames(int(prev.shape[2]))
    run = _active_run(requested, available_frames, allow_one=False)
    steps = _latent_steps_for_run(run)
    if int(prev.shape[2]) < steps:
        raise ValueError(
            f"h3_continuation_context: previous latent has only {prev.shape[2]} temporal steps, need {steps}"
        )
    # H3 target latents have T = 5k+2. Taking 2/7/12 tail steps therefore starts
    # at phase 0, matching a fresh 5/22/39-frame VAE run.
    phase = (int(prev.shape[2]) - steps) % 5
    if phase != 0:
        raise ValueError(
            "h3_continuation_context: previous latent temporal phase is unexpected; "
            f"tail would start at phase {phase}, refusing direct latent context"
        )
    tail = prev[:1, :, -steps:].clone()
    blocks = [tail[:, :, k:k + 1] for k in range(steps)]
    return blocks, _step_offsets(steps), run, f"latent-direct {run}f/{steps}t"


def _audio_tail_from_latent(latent: dict, a_frames: int):
    video = _video(latent)
    audio = _audio(latent)
    total_t = int(audio.shape[-1])
    frames = _pixel_frames(int(video.shape[2]))
    overhang = total_t - FRAME_RESCALE * frames
    if not (0.0 <= overhang < 1.0):
        LOG.warning(
            "h3_continuation_context: unexpected audio grid (%d steps for %d frames); "
            "assuming no overhang", total_t, frames,
        )
        overhang = 0.0
    rt = int(round(float(a_frames) / FPS * AUDIO_HZ))
    rt = min(rt, total_t)
    if rt < 1:
        raise ValueError("h3_continuation_context: audio context window is empty")
    return audio[:1, ..., total_t - rt:].clone(), rt, float(overhang)


def _encode_audio_tail(audio_vae, audio: dict, a_frames: int):
    if audio_vae is None:
        raise ValueError("h3_continuation_context: context_audio requires audio_vae")
    waveform = audio["waveform"]
    sr = int(audio["sample_rate"])
    vae_sr = int(getattr(audio_vae, "audio_sample_rate", 32000))
    if sr != vae_sr:
        try:
            import torchaudio
        except ImportError as exc:
            raise RuntimeError(
                f"h3_continuation_context: audio resample {sr}->{vae_sr} needs torchaudio"
            ) from exc
        waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
    want = int(round(a_frames / FPS * vae_sr))
    if waveform.shape[-1] > want:
        waveform = waveform[..., -want:]
    z = audio_vae.encode(waveform[:1].movedim(1, -1))
    return z, int(z.shape[-1]), 0.0


def _merge_conditioning(conditioning, keyframes, frame_count, audio_ref=None,
                        preserve_existing=True):
    """Add H3 keyframes without silently deleting upstream FL2VA anchors.

    The first version of this fork replaced ``minimax_keyframes`` outright.
    That meant putting this node after MiniMaxH3ImageToVideo discarded the
    stock first-frame keyframe before adding ours.  For first-frame work that
    made the new path effectively no stronger than native FL2VA.

    Existing keyframes are retained by default.  The layout patch already
    compensates both stock and custom keyframes when references coexist, so
    mixed endpoint/interior anchors remain aligned to the target timeline.
    """
    out = []
    for entry in conditioning:
        if not isinstance(entry, (list, tuple)) or len(entry) != 2:
            raise TypeError("h3_continuation_context: malformed CONDITIONING entry")
        tensor, meta = entry
        meta = dict(meta)

        existing = list(meta.get("minimax_keyframes") or []) if preserve_existing else []
        old_fc = meta.get("minimax_frame_count")
        if old_fc is not None and int(old_fc) != int(frame_count):
            raise ValueError(
                "h3_continuation_context: upstream minimax_frame_count does not match "
                f"target latent: conditioning={old_fc}, target={frame_count}"
            )

        meta["minimax_keyframes"] = existing + list(keyframes)
        meta["minimax_frame_count"] = int(frame_count)
        if audio_ref is not None:
            refs = list(meta.get("minimax_refs") or [])
            refs.append(audio_ref)
            meta["minimax_refs"] = refs
        out.append([tensor, meta])
    return out


def _conditioning_keyframe_count(conditioning) -> int:
    if not conditioning:
        return 0
    entry = conditioning[0]
    if not isinstance(entry, (list, tuple)) or len(entry) != 2:
        return 0
    meta = entry[1] if isinstance(entry[1], dict) else {}
    return len(meta.get("minimax_keyframes") or [])


def _first_frame_blocks(vae, image: torch.Tensor, target_video: torch.Tensor,
                        mode: str, repeats: int, crop: str):
    """Build one-image H3 conditioning blocks for a strong frame-0 anchor.

    Modes:
      single      one ordinary H3 first-frame block (native-equivalent)
      reinforce   duplicate that same never-denoised block at time 0
      static_5    encode five copies as one H3 video run -> 2 blocks at 0,1
      static_22   encode 22 copies as one H3 video run -> 7 blocks

    ``reinforce`` does not consume output time: all copies have the exact same
    frame-0 position.  ``static_*`` deliberately pins a short static prefix and
    therefore is much stronger but can create an initial pause.
    """
    if vae is None:
        raise ValueError("h3_continuation_context: first-frame context requires vae")
    if image is None or getattr(image, "ndim", 0) != 4 or int(image.shape[0]) < 1:
        raise ValueError("h3_continuation_context: first_frame must contain at least one IMAGE")

    image = image[:1]
    if mode in ("single", "reinforce"):
        blocks, _, _, _ = _encode_video_context(vae, image, target_video, 1, crop)
        block = blocks[0]
        count = 1 if mode == "single" else max(1, int(repeats))
        # Same clean condition repeated at the same temporal coordinate.  This
        # increases its representation in attention without freezing later
        # target frames or writing into the denoised target latent.
        return [block for _ in range(count)], [0 for _ in range(count)], 1

    if mode == "static_5":
        run = 5
    elif mode == "static_22":
        run = 22
    else:
        raise ValueError(f"h3_continuation_context: unknown first_frame_mode {mode!r}")

    repeated = image.repeat(run, 1, 1, 1)
    blocks, offsets, span, _ = _encode_video_context(
        vae, repeated, target_video, run, crop
    )
    return blocks, offsets, span


class MiniMaxH3ContinuationContext:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "conditioning": ("CONDITIONING",),
                "latent": ("LATENT",),
                "video_source": (["frames", "latent", "first_frame"], {"default": "frames"}),
                "context_length": ("INT", {"default": 22, "min": 1, "max": 39, "step": 1}),
                "anchor_mode": (["head", "before"], {"default": "head"}),
                "crop": (["center", "disabled"], {"default": "center"}),
                "audio_mode": (["timeline", "ref", "off"], {"default": "timeline"}),
                "audio_context_length": ("INT", {"default": 22, "min": 0, "max": 240, "step": 1}),
            },
            "optional": {
                "vae": ("VAE",),
                "context_frames": ("IMAGE",),
                "first_frame": ("IMAGE",),
                "context_latent": ("LATENT",),
                "context_audio": ("AUDIO",),
                "audio_vae": ("VAE",),
            },
        }

    RETURN_TYPES = ("CONDITIONING", "INT", "STRING")
    RETURN_NAMES = ("conditioning", "trim_frames", "info")
    FUNCTION = "apply"
    CATEGORY = "MiniMax H3/Continuation"
    DESCRIPTION = (
        "Pins a first frame or a consecutive tail from the previous H3 clip as native "
        "never-denoised visual conditioning rows. Unlike a latent noise mask, this uses "
        "H3's own keyframe timeline conditioning."
    )

    def apply(self, conditioning, latent, video_source, context_length, anchor_mode,
              crop, audio_mode, audio_context_length,
              vae=None, context_frames=None, first_frame=None, context_latent=None,
              context_audio=None, audio_vae=None):
        if not layout_patch_applied():
            raise RuntimeError(
                "h3_continuation_context: layout patch is not active. Check startup logs; "
                "another H3 layout patch may already be installed."
            )

        target_video = _video(latent)
        latent_t = int(target_video.shape[2])
        frame_count = _pixel_frames(latent_t)
        width = int(target_video.shape[-1]) * 16
        height = int(target_video.shape[-2]) * 16

        first_frame_source = video_source == "first_frame"
        effective_anchor_mode = anchor_mode

        if video_source == "latent":
            if int(context_length) < 5:
                raise ValueError(
                    "h3_continuation_context: direct latent context cannot represent a clean "
                    "single RGB frame from the tail of a normal H3 clip. Use video_source=frames "
                    "with context_length=1, or video_source=first_frame with an IMAGE + VAE."
                )
            if context_latent is None:
                raise ValueError("h3_continuation_context: video_source=latent requires context_latent")
            blocks, offsets, span, source_info = _encode_latent_context(
                context_latent, target_video, int(context_length)
            )
        elif first_frame_source:
            if first_frame is None:
                raise ValueError("h3_continuation_context: video_source=first_frame requires first_frame")
            if int(first_frame.shape[0]) > 1:
                LOG.warning(
                    "h3_continuation_context: first_frame contains %d images; using only the first",
                    int(first_frame.shape[0]),
                )
            blocks, offsets, span, source_info = _encode_video_context(
                vae, first_frame[:1], target_video, 1, crop
            )
            # A first-frame anchor always belongs at target frame 0. Unlike a
            # continuation head, it is part of the requested output and must not
            # be advertised for trimming.
            effective_anchor_mode = "head"
            source_info = "first-frame/VAE 1f->1t"
        else:
            if context_frames is None:
                raise ValueError("h3_continuation_context: video_source=frames requires context_frames")
            blocks, offsets, span, source_info = _encode_video_context(
                vae, context_frames, target_video, int(context_length), crop
            )

        if span >= frame_count:
            raise ValueError(
                f"h3_continuation_context: context span {span} must be smaller than target clip {frame_count}"
            )

        indices = [o - span for o in offsets] if effective_anchor_mode == "before" else list(offsets)
        keyframes = [
            {
                "resolved_frame_index": 0,  # always legal to stock constructor
                MC_INDEX: int(pixel_index),
                "latent": block,
            }
            for pixel_index, block in zip(indices, blocks)
        ]

        audio_ref = None
        audio_info = "off"
        if first_frame_source and audio_mode != "off":
            audio_info = "off (first_frame source)"
        elif audio_mode != "off" and (context_latent is not None or context_audio is not None):
            if not payload_patch_applied():
                raise RuntimeError(
                    "h3_continuation_context: payload patch is not active; audio refs would "
                    "overwrite visual keyframe latents"
                )
            a_frames = int(audio_context_length) or span
            if context_latent is not None:
                audio_latent, rt, overhang = _audio_tail_from_latent(context_latent, a_frames)
                audio_info = f"latent {a_frames}f->{rt}t"
            else:
                audio_latent, rt, overhang = _encode_audio_tail(audio_vae, context_audio, a_frames)
                audio_info = f"decoded/VAE {a_frames}f->{rt}t"
            audio_ref = {
                "kind": "audio",
                "ref_audio_t": int(rt),
                "audio_latent": audio_latent,
            }
            if audio_mode == "timeline":
                end_frame = float(span if effective_anchor_mode == "head" else 0)
                end_frame += float(overhang) / FRAME_RESCALE
                audio_ref[MC_AUDIO_END] = end_frame
                audio_info += f" timeline@end={end_frame:.3f}f"
            else:
                audio_info += " stock-ref"

        out = _merge_conditioning(conditioning, keyframes, frame_count, audio_ref)
        trim = 0 if first_frame_source else int(span if effective_anchor_mode == "head" else 0)
        info = (
            f"H3 continuation: {source_info}; anchor={effective_anchor_mode}; "
            f"indices={indices[0]}..{indices[-1]}; target={width}x{height}, "
            f"{frame_count}f/{latent_t}t; trim={trim}; audio={audio_info}"
        )
        LOG.info(info)
        return out, trim, info


class MiniMaxH3FirstFrameContext:
    """Strengthen a supplied first image using H3's native conditioning rows.

    This is intentionally separate from continuation context.  It keeps any
    existing FL2VA first/last keyframes already present in CONDITIONING and
    adds extra never-denoised frame-0 context instead of replacing them.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "conditioning": ("CONDITIONING",),
                "latent": ("LATENT",),
                "vae": ("VAE",),
                "first_frame": ("IMAGE",),
                "first_frame_mode": (["reinforce", "single", "static_5", "static_22"],
                                     {"default": "reinforce"}),
                "reinforce_copies": ("INT", {"default": 2, "min": 1, "max": 8, "step": 1}),
                "crop": (["center", "disabled"], {"default": "center"}),
                "preserve_existing_keyframes": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("CONDITIONING", "STRING")
    RETURN_NAMES = ("conditioning", "info")
    FUNCTION = "apply"
    CATEGORY = "MiniMax H3/Continuation"
    DESCRIPTION = (
        "Strengthens H3 first-frame conditioning. Reinforce duplicates the same "
        "never-denoised frame-0 condition without consuming timeline frames; "
        "static modes encode one image as a short repeated video context."
    )

    def apply(self, conditioning, latent, vae, first_frame, first_frame_mode,
              reinforce_copies=2, crop="center", preserve_existing_keyframes=True):
        if not layout_patch_applied():
            raise RuntimeError(
                "h3_continuation_context: layout patch is not active. Check startup logs."
            )

        target_video = _video(latent)
        latent_t = int(target_video.shape[2])
        frame_count = _pixel_frames(latent_t)
        width = int(target_video.shape[-1]) * 16
        height = int(target_video.shape[-2]) * 16

        if int(first_frame.shape[0]) > 1:
            LOG.warning(
                "h3_continuation_context: first_frame batch has %d images; using the first",
                int(first_frame.shape[0]),
            )

        blocks, indices, span = _first_frame_blocks(
            vae, first_frame, target_video, first_frame_mode,
            int(reinforce_copies), crop,
        )
        if span >= frame_count:
            raise ValueError(
                f"h3_continuation_context: first-frame context span {span} must be "
                f"smaller than target clip {frame_count}"
            )

        keyframes = [
            {
                "resolved_frame_index": 0,
                MC_INDEX: int(pixel_index),
                "latent": block,
            }
            for pixel_index, block in zip(indices, blocks)
        ]

        before = _conditioning_keyframe_count(conditioning)
        out = _merge_conditioning(
            conditioning, keyframes, frame_count, audio_ref=None,
            preserve_existing=bool(preserve_existing_keyframes),
        )
        after = _conditioning_keyframe_count(out)

        info = (
            f"H3 first frame: mode={first_frame_mode}; added={len(keyframes)}; "
            f"indices={indices}; existing_before={before}; total_after={after}; "
            f"preserve_existing={bool(preserve_existing_keyframes)}; "
            f"target={width}x{height}, {frame_count}f/{latent_t}t"
        )
        if first_frame_mode.startswith("static_"):
            info += f"; static_prefix={span}f (not trimmed)"
        elif first_frame_mode == "reinforce":
            info += "; all added blocks share frame-0 time (no frozen prefix)"
        LOG.info(info)
        return out, info


class MiniMaxH3ContinuationTrim:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "trim_frames": ("INT", {"default": 0, "min": 0, "max": 4096}),
            },
            "optional": {
                "audio": ("AUDIO",),
                "fps": ("FLOAT", {"default": 24.0, "min": 1.0, "max": 240.0, "step": 0.001}),
                "match_tail": ("BOOLEAN", {"default": True}),
            },
        }
    RETURN_TYPES = ("IMAGE", "AUDIO", "STRING")
    RETURN_NAMES = ("images", "audio", "info")
    FUNCTION = "trim"
    CATEGORY = "MiniMax H3/Continuation"

    def trim(self, images, trim_frames, audio=None, fps=24.0, match_tail=True):
        n = max(0, int(trim_frames))
        total = int(images.shape[0])
        if n >= total:
            raise ValueError(f"h3_continuation_context: cannot trim {n} from {total} frames")
        out_images = images[n:] if n else images
        out_audio = audio
        msg = f"trimmed {n} head frames; {total - n} remain"
        if audio is not None:
            waveform = audio["waveform"]
            sr = int(audio["sample_rate"])
            cut = int(round(n / float(fps) * sr))
            if cut >= waveform.shape[-1]:
                raise ValueError("h3_continuation_context: audio head trim would leave no samples")
            waveform = waveform[..., cut:]
            tail_trimmed = 0
            if match_tail:
                want = int(round((total - n) / float(fps) * sr))
                if waveform.shape[-1] > want:
                    tail_trimmed = int(waveform.shape[-1] - want)
                    waveform = waveform[..., :want]
            out_audio = {"waveform": waveform, "sample_rate": sr}
            msg += f"; audio={waveform.shape[-1]/sr:.4f}s; tail_cut={tail_trimmed} samples"
        LOG.info("h3_continuation_context: %s", msg)
        return out_images, out_audio, msg


class MiniMaxH3ContinuationInspect:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"latent": ("LATENT",)}}
    RETURN_TYPES = ("LATENT", "STRING")
    RETURN_NAMES = ("latent", "info")
    FUNCTION = "inspect"
    CATEGORY = "MiniMax H3/Continuation"

    def inspect(self, latent):
        v, a = _video(latent), _audio(latent)
        frames = _pixel_frames(int(v.shape[2]))
        info = (
            f"video={tuple(v.shape)} -> RGB {v.shape[-1]*16}x{v.shape[-2]*16}, "
            f"frames={frames}; audio={tuple(a.shape)}; "
            f"patches(layout/payload)={layout_patch_applied()}/{payload_patch_applied()}"
        )
        LOG.info("h3_continuation_context: %s", info)
        return latent, info