ComfyUI-MiniMaxH3-Context-Loop/masked_bridge.py

"""Two-ended MiniMax H3 audiovisual latent bridge.

The tail of one decoded source clip and the head of a second clip are encoded
into opposite ends of a normal H3 target latent.  A nested AV denoise mask
protects both source windows while allowing only the middle interval to be
generated.  Mask semantics follow ComfyUI PR #15375: 0 preserves the supplied
target latent and 1 denoises it.

This implementation adapts the GPL-3.0 masked bridge published by seitanism's
ComfyUI-H3-Motion-Context-MultiRef project for this pack's distinct node IDs,
runtime capability gate, categories, and error contract.
"""

from __future__ import annotations

import logging

import torch

from .masking_support import require_h3_mask_support
from .nodes import AUDIO_HZ, FPS, _pixel_frames, _resize


try:
    import torchaudio
except ImportError:  # pragma: no cover - exercised only on minimal installs
    torchaudio = None


_LOG = logging.getLogger("minimax_h3_context_loop.masked_bridge")


def _largest_h3_video_run(frames):
    frames = int(frames)
    if frames < 5:
        return 0
    return 5 + ((frames - 5) // 17) * 17


def _target_streams(latent):
    if not isinstance(latent, dict) or latent.get("samples") is None:
        raise ValueError("h3_masked_bridge: expected a target LATENT.")
    samples = latent["samples"]
    if hasattr(samples, "unbind"):
        parts = list(samples.unbind())
    elif isinstance(samples, (tuple, list)):
        parts = list(samples)
    else:
        raise ValueError(
            "h3_masked_bridge: expected a joint H3 AV latent, got %r." %
            type(samples))
    if len(parts) < 2:
        raise ValueError(
            "h3_masked_bridge: target latent has no joint video/audio streams.")
    video, audio = parts[:2]
    if video.ndim == 4:
        video = video.unsqueeze(0)
    if audio.ndim == 3:
        audio = audio.unsqueeze(0)
    if video.ndim != 5:
        raise ValueError(
            "h3_masked_bridge: video latent must be [B,C,T,H,W], got %s." %
            (tuple(video.shape),))
    if audio.ndim != 4:
        raise ValueError(
            "h3_masked_bridge: audio latent must be [B,C,channels,T], got %s."
            % (tuple(audio.shape),))
    if int(video.shape[0]) != 1 or int(audio.shape[0]) != 1:
        raise ValueError(
            "h3_masked_bridge: only MiniMax H3 batch size 1 is supported.")
    return video, audio


def _cfr_index_map(frame_count, source_fps, device):
    source_fps = float(source_fps)
    if source_fps <= 0.0:
        raise ValueError("h3_masked_bridge: source FPS must be greater than 0.")
    frame_count = int(frame_count)
    if frame_count < 1:
        raise ValueError("h3_masked_bridge: a source video has no frames.")
    output_count = max(1, int(round(frame_count * FPS / source_fps)))
    if output_count == frame_count and abs(source_fps - FPS) < 1e-6:
        return torch.arange(frame_count, device=device, dtype=torch.long)
    frame = torch.arange(output_count, device=device, dtype=torch.float64)
    time = (frame + 0.5) / FPS
    source = torch.round(time * source_fps - 0.5).to(torch.long)
    return source.clamp_(0, frame_count - 1)


def _stereo_audio(audio, target_rate, frame_count, label):
    if audio is None or not isinstance(audio, dict):
        raise ValueError("h3_masked_bridge: %s is required." % label)
    waveform = audio.get("waveform")
    if getattr(waveform, "ndim", 0) != 3:
        raise ValueError(
            "h3_masked_bridge: %s waveform must be [B,C,L], got %s." %
            (label, tuple(getattr(waveform, "shape", ()))))
    waveform = waveform[:1]
    channels = int(waveform.shape[1])
    if channels == 1:
        waveform = waveform.repeat(1, 2, 1)
    elif channels != 2:
        raise ValueError(
            "h3_masked_bridge: %s has %d channels; downmix it to stereo first."
            % (label, channels))
    source_rate = int(audio.get("sample_rate", 0))
    if source_rate <= 0:
        raise ValueError("h3_masked_bridge: %s has no valid sample rate." % label)
    if source_rate != int(target_rate):
        if torchaudio is None:
            raise RuntimeError(
                "h3_masked_bridge: %s is %d Hz but the audio VAE needs %d Hz "
                "and torchaudio is unavailable." %
                (label, source_rate, int(target_rate)))
        waveform = torchaudio.functional.resample(
            waveform, source_rate, int(target_rate))
    wanted = int(round(int(frame_count) / FPS * int(target_rate)))
    have = int(waveform.shape[-1])
    if have > wanted:
        waveform = waveform[..., :wanted]
    elif have < wanted:
        _LOG.warning(
            "h3_masked_bridge: %s is %d samples short of its video; padding "
            "only the source clip tail before selecting the protected window.",
            label, wanted - have)
        waveform = torch.nn.functional.pad(waveform, (0, wanted - have))
    return waveform


def _encode_video_window(vae, frames, width, height, crop, label):
    resized = _resize(frames, width, height, crop)
    encoded = vae.encode(resized)
    if getattr(encoded, "ndim", 0) != 5:
        raise ValueError(
            "h3_masked_bridge: %s video VAE returned %s; expected "
            "[B,C,T,H,W]." % (label, tuple(getattr(encoded, "shape", ()))))
    covered = _pixel_frames(int(encoded.shape[2]))
    if covered != int(frames.shape[0]):
        raise RuntimeError(
            "h3_masked_bridge: %s %d-frame window encoded to %d steps "
            "covering %d frames; refusing a phase-shifted seam." %
            (label, int(frames.shape[0]), int(encoded.shape[2]), covered))
    return encoded[:1]


def _encode_audio_window(audio_vae, waveform, preserve_frames, side, label):
    sample_rate = int(getattr(audio_vae, "audio_sample_rate", 32000))
    samples = int(round(int(preserve_frames) / FPS * sample_rate))
    if int(waveform.shape[-1]) < samples:
        raise ValueError(
            "h3_masked_bridge: %s is shorter than the %d-frame protected "
            "window." % (label, int(preserve_frames)))
    if side == "tail":
        window = waveform[..., -samples:]
    elif side == "head":
        window = waveform[..., :samples]
    else:  # pragma: no cover - private caller contract
        raise ValueError("h3_masked_bridge: invalid audio side %r." % side)
    encoded = audio_vae.encode(window.movedim(1, -1))
    if getattr(encoded, "ndim", 0) != 4:
        raise ValueError(
            "h3_masked_bridge: %s audio VAE returned %s; expected "
            "[B,C,channels,T]." %
            (label, tuple(getattr(encoded, "shape", ()))))
    exact_steps = int(preserve_frames) / FPS * AUDIO_HZ
    expected = int(round(exact_steps))
    if abs(exact_steps - expected) > 1e-9:
        _LOG.warning(
            "h3_masked_bridge: %d protected frames end between H3 audio "
            "ticks (%.6f -> %d steps); 39/90/141/... are exact AV boundaries.",
            int(preserve_frames), exact_steps, expected)
    got = int(encoded.shape[-1])
    if got < expected:
        raise RuntimeError(
            "h3_masked_bridge: %s needs %d audio steps but the VAE produced "
            "%d." % (label, expected, got))
    if got > expected:
        encoded = (encoded[..., -expected:] if side == "tail"
                   else encoded[..., :expected])
    return encoded[:1]


class MiniMaxH3ContexMaskedAVBridge:
    """Protect source AV windows at both ends and generate only the middle."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT", {
                    "tooltip": "Empty joint H3 AV target latent for the full "
                               "bridge duration."}),
                "vae": ("VAE", {"tooltip": "MiniMax H3 video VAE."}),
                "audio_vae": ("VAE", {"tooltip": "MiniMax H3 audio VAE."}),
                "start_frames": ("IMAGE", {
                    "tooltip": "First source clip. Its final protected window "
                               "becomes the bridge prefix."}),
                "end_frames": ("IMAGE", {
                    "tooltip": "Second source clip. Its initial protected "
                               "window becomes the bridge suffix."}),
                "start_fps": ("FLOAT", {
                    "default": 24.0, "min": 0.001, "max": 1000.0,
                    "step": 0.001,
                    "tooltip": "Actual frame rate represented by "
                               "start_frames; it is converted to H3's 24 fps "
                               "timeline before selecting the protected tail.",
                }),
                "end_fps": ("FLOAT", {
                    "default": 24.0, "min": 0.001, "max": 1000.0,
                    "step": 0.001,
                    "tooltip": "Actual frame rate represented by end_frames; "
                               "it is converted to H3's 24 fps timeline before "
                               "selecting the protected head.",
                }),
                "preserve_frames": ("INT", {
                    "default": 39, "min": 5, "max": 9999,
                    "tooltip": "Exact H3 run: 5, 22, 39, 56, ... . Use 39 "
                               "for an exact 65-step AV boundary."}),
                "crop": (["disabled", "center"], {
                    "default": "center",
                    "tooltip": "Resize policy for both endpoint videos: "
                               "disabled stretches to the target canvas; "
                               "center preserves aspect ratio and center-crops.",
                }),
            },
            "optional": {
                "start_audio": ("AUDIO", {
                    "tooltip": "Optional audio synchronized with the first "
                               "clip. When absent, H3 generates the bridge's "
                               "opening audio instead of protecting silence."}),
                "end_audio": ("AUDIO", {
                    "tooltip": "Optional audio synchronized with the second "
                               "clip. When absent, H3 generates the bridge's "
                               "ending audio instead of protecting silence."}),
            },
        }

    RETURN_TYPES = ("LATENT", "INT", "INT")
    RETURN_NAMES = ("latent", "middle_frames", "preserve_frames")
    OUTPUT_TOOLTIPS = (
        "Full bridge target with protected source AV windows at both ends "
        "and a denoised middle.",
        "Number of picture frames left for the model to generate between "
        "the two protected endpoints.",
        "Actual number of protected picture frames copied from each endpoint.",
    )
    FUNCTION = "prepare"
    CATEGORY = "conditioning/minimax/context_loop/masking"
    DESCRIPTION = (
        "Construct a true two-ended H3 AV bridge target: preserve the first "
        "clip's tail and second clip's head, then denoise only the middle.")

    def prepare(
        self,
        latent,
        vae,
        audio_vae,
        start_frames,
        start_audio=None,
        end_frames=None,
        end_audio=None,
        start_fps=24.0,
        end_fps=24.0,
        preserve_frames=39,
        crop="center",
    ):
        require_h3_mask_support("two-ended masked AV bridge")
        target_video, target_audio = _target_streams(latent)
        target_frames = _pixel_frames(int(target_video.shape[2]))
        if _largest_h3_video_run(target_frames) != target_frames:
            raise RuntimeError(
                "h3_masked_bridge: target covers %d frames, not an exact H3 "
                "run." % target_frames)
        expected_audio = int(round(target_frames / FPS * AUDIO_HZ))
        if int(target_audio.shape[-1]) != expected_audio:
            raise RuntimeError(
                "h3_masked_bridge: target has %d audio steps for %d frames; "
                "expected %d." %
                (int(target_audio.shape[-1]), target_frames, expected_audio))

        for frames, label in ((start_frames, "start_frames"),
                              (end_frames, "end_frames")):
            if getattr(frames, "ndim", 0) != 4 or int(frames.shape[0]) < 1:
                raise ValueError(
                    "h3_masked_bridge: %s must be IMAGE [N,H,W,C]." % label)
        start_idx = _cfr_index_map(
            int(start_frames.shape[0]), start_fps, start_frames.device)
        end_idx = _cfr_index_map(
            int(end_frames.shape[0]), end_fps, end_frames.device)

        protected = int(preserve_frames)
        if protected < 5 or _largest_h3_video_run(protected) != protected:
            raise ValueError(
                "h3_masked_bridge: preserve_frames must be an exact H3 video "
                "run (5, 22, 39, 56, ...); got %d." % protected)
        if int(start_idx.numel()) < protected:
            raise ValueError(
                "h3_masked_bridge: first clip has %d canonical frames; need "
                "%d." % (int(start_idx.numel()), protected))
        if int(end_idx.numel()) < protected:
            raise ValueError(
                "h3_masked_bridge: second clip has %d canonical frames; need "
                "%d." % (int(end_idx.numel()), protected))
        if 2 * protected >= target_frames:
            raise ValueError(
                "h3_masked_bridge: protected windows consume the complete "
                "%d-frame target." % target_frames)

        width = int(target_video.shape[4]) * 16
        height = int(target_video.shape[3]) * 16
        start_tail = start_frames.index_select(0, start_idx[-protected:])
        end_head = end_frames.index_select(0, end_idx[:protected])
        start_video = _encode_video_window(
            vae, start_tail, width, height, crop, "first clip tail")
        end_video = _encode_video_window(
            vae, end_head, width, height, crop, "second clip head")
        video_steps = int(start_video.shape[2])
        if int(end_video.shape[2]) != video_steps:
            raise RuntimeError(
                "h3_masked_bridge: protected video windows encoded to "
                "different lengths.")
        suffix_step = int(target_video.shape[2]) - video_steps
        if suffix_step % 5:
            raise RuntimeError(
                "h3_masked_bridge: suffix begins at latent phase %d; refusing "
                "an out-of-phase protected endpoint." % (suffix_step % 5))

        vae_rate = int(getattr(audio_vae, "audio_sample_rate", 32000))
        start_audio_latent = None
        if start_audio is not None:
            start_waveform = _stereo_audio(
                start_audio, vae_rate, int(start_idx.numel()), "start_audio")
            start_audio_latent = _encode_audio_window(
                audio_vae, start_waveform, protected, "tail",
                "first clip tail")
        else:
            _LOG.info(
                "h3_masked_bridge: start_audio is absent; opening audio stays "
                "unmasked for H3 generation.")
        end_audio_latent = None
        if end_audio is not None:
            end_waveform = _stereo_audio(
                end_audio, vae_rate, int(end_idx.numel()), "end_audio")
            end_audio_latent = _encode_audio_window(
                audio_vae, end_waveform, protected, "head", "second clip head")
        else:
            _LOG.info(
                "h3_masked_bridge: end_audio is absent; ending audio stays "
                "unmasked for H3 generation.")
        start_audio_steps = (int(start_audio_latent.shape[-1])
                             if start_audio_latent is not None else 0)
        end_audio_steps = (int(end_audio_latent.shape[-1])
                           if end_audio_latent is not None else 0)
        if (start_audio_latent is not None and end_audio_latent is not None
                and end_audio_steps != start_audio_steps):
            raise RuntimeError(
                "h3_masked_bridge: protected audio windows encoded to "
                "different lengths.")

        out_video = target_video.clone()
        out_audio = target_audio.clone()
        start_video = start_video.to(out_video)
        end_video = end_video.to(out_video)
        if (tuple(start_video.shape[1:2] + start_video.shape[3:]) !=
                tuple(out_video.shape[1:2] + out_video.shape[3:])):
            raise ValueError(
                "h3_masked_bridge: first video latent shape does not match "
                "the target canvas.")
        if (tuple(end_video.shape[1:2] + end_video.shape[3:]) !=
                tuple(out_video.shape[1:2] + out_video.shape[3:])):
            raise ValueError(
                "h3_masked_bridge: second video latent shape does not match "
                "the target canvas.")
        if start_audio_latent is not None:
            start_audio_latent = start_audio_latent.to(out_audio)
            if tuple(start_audio_latent.shape[1:3]) != tuple(
                    out_audio.shape[1:3]):
                raise ValueError(
                    "h3_masked_bridge: first audio latent shape does not match "
                    "the H3 target.")
        if end_audio_latent is not None:
            end_audio_latent = end_audio_latent.to(out_audio)
            if tuple(end_audio_latent.shape[1:3]) != tuple(
                    out_audio.shape[1:3]):
                raise ValueError(
                    "h3_masked_bridge: second audio latent shape does not "
                    "match the H3 target.")

        out_video[:, :, :video_steps] = start_video
        out_video[:, :, -video_steps:] = end_video
        video_mask = torch.ones(
            (1, 1, int(out_video.shape[2]), int(out_video.shape[3]),
             int(out_video.shape[4])),
            device=out_video.device, dtype=torch.float32)
        audio_mask = torch.ones(
            (1, 1, int(out_audio.shape[2]), int(out_audio.shape[3])),
            device=out_audio.device, dtype=torch.float32)
        video_mask[:, :, :video_steps] = 0.0
        video_mask[:, :, -video_steps:] = 0.0
        if start_audio_latent is not None:
            out_audio[..., :start_audio_steps] = start_audio_latent
            audio_mask[..., :start_audio_steps] = 0.0
        if end_audio_latent is not None:
            out_audio[..., -end_audio_steps:] = end_audio_latent
            audio_mask[..., -end_audio_steps:] = 0.0

        import comfy.nested_tensor
        output = latent.copy()
        output["samples"] = comfy.nested_tensor.NestedTensor(
            (out_video, out_audio))
        output["noise_mask"] = comfy.nested_tensor.NestedTensor(
            (video_mask, audio_mask))
        middle_frames = target_frames - 2 * protected
        _LOG.info(
            "h3_masked_bridge: %d-frame target at %dx%d; preserve %d frames "
            "(%d video / %d+%d audio steps) at the two ends; generate %d "
            "frames.",
            target_frames, width, height, protected, video_steps,
            start_audio_steps, end_audio_steps, middle_frames)
        return output, middle_frames, protected


NODE_CLASS_MAPPINGS = {
    "MiniMaxH3ContexMaskedAVBridge": MiniMaxH3ContexMaskedAVBridge,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MiniMaxH3ContexMaskedAVBridge": (
        "MiniMax H3 Masking · Two-Clip AV Bridge"),
}