ComfyUI-MiniMaxH3-Context-Loop/master_audio_context.py

"""Exact master-audio masking with an optional previous-video prefix.

This is the timeline-audio specialization of the general H3 masked-target
path: the target audio stream is replaced with one exact master-audio interval
and protected for the complete raw clip. Continuations may additionally
protect the final native H3 run from either the preceding sampled AV latent or
the legacy decoded-video path.

The design is adapted from seitanism's GPL-3.0
ComfyUI-H3-Motion-Context-MultiRef Update 4 implementation, including Update
5's target-audio-grid boundary correction. This pack keeps a distinct public
node id and reuses its shared native-first mask capability gate so both packs
can be installed together.
"""

from __future__ import annotations

import logging
import math

import torch

from .masked_context import (
    _existing_mask_streams,
    _generated_video_tail,
    _snap_prefix_length,
    _validate_target_streams,
)
from .masking_support import require_h3_mask_support
from .nodes import AUDIO_HZ, FPS, _pixel_frames, _resize, _streams_from_latent
from .av_timing import sample_boundary_from_seconds
from .contracts_v05 import masked_song_options as _contract_masked_song_options


_LOG = logging.getLogger("minimax_h3_context_loop.master_audio_context")
_VOICE_GATE_RELEASE_SECONDS = 0.15


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


def _stereo_first_batch(waveform):
    if getattr(waveform, "ndim", 0) != 3:
        raise ValueError(
            "h3_master_audio_mask: master audio waveform must be [B,C,L], "
            "got %s." %
            (tuple(getattr(waveform, "shape", ())),)
        )
    waveform = waveform[:1]
    channels = int(waveform.shape[1])
    if channels == 1:
        return waveform.repeat(1, 2, 1)
    if channels == 2:
        return waveform
    raise ValueError(
        "h3_master_audio_mask: master audio has %d channels; downmix it to stereo "
        "before this node." % channels
    )


def _resample(waveform, source_rate, target_rate):
    source_rate = int(source_rate)
    target_rate = int(target_rate)
    if source_rate == target_rate:
        return waveform
    try:
        import torchaudio
    except ImportError as exc:
        raise RuntimeError(
            "h3_master_audio_mask: master audio is %d Hz but the H3 audio VAE needs "
            "%d Hz and torchaudio is unavailable." %
            (source_rate, target_rate)
        ) from exc
    return torchaudio.functional.resample(waveform, source_rate, target_rate)


def _fit_audio_slice(waveform, samples):
    wanted = int(samples)
    available = int(waveform.shape[-1])
    if available > wanted:
        return waveform[..., :wanted]
    if available < wanted:
        _LOG.warning(
            "h3_master_audio_mask: master-audio slice is %d samples short; padding "
            "silence at the tail.", wanted - available)
        return torch.nn.functional.pad(waveform, (0, wanted - available))
    return waveform


def _voice_activity_ticks(
        voice, start_seconds, steps, target_rate, hold_seconds,
        force_active_prefix_seconds=0.0):
    """Return a conservative 40 Hz vocal-presence gate for one target span."""
    if not isinstance(voice, dict):
        raise ValueError(
            "h3_master_audio_mask: voice is not a Comfy AUDIO value.")
    waveform = _stereo_first_batch(voice.get("waveform"))
    waveform = _resample(
        waveform, int(voice.get("sample_rate", 0)), int(target_rate))
    ticks = int(steps)
    rate = int(target_rate)
    wanted_samples = int(math.ceil(ticks / float(AUDIO_HZ) * rate))
    start_sample = sample_boundary_from_seconds(float(start_seconds), rate)
    leading_silence = max(0, -start_sample)
    source_start = max(0, start_sample)
    window = torch.zeros(
        (1, int(waveform.shape[1]), wanted_samples),
        dtype=waveform.dtype, device=waveform.device)
    copy_start = min(wanted_samples, leading_silence)
    available = max(0, min(
        wanted_samples - copy_start,
        int(waveform.shape[-1]) - source_start))
    if available:
        window[..., copy_start:copy_start + available] = (
            waveform[..., source_start:source_start + available])

    mono = window.abs().mean(dim=1)[0]
    envelope = torch.zeros((ticks,), dtype=torch.float32, device=mono.device)
    for tick in range(ticks):
        left = int(round(tick / float(AUDIO_HZ) * rate))
        right = int(round((tick + 1) / float(AUDIO_HZ) * rate))
        right = min(wanted_samples, max(left + 1, right))
        if left < wanted_samples:
            envelope[tick] = mono[left:right].mean().to(torch.float32)
    peak = float(envelope.max().item()) if ticks else 0.0
    threshold = max(1e-4, peak * 0.08)
    active = envelope >= threshold if peak >= 1e-4 else torch.zeros(
        (ticks,), dtype=torch.bool, device=envelope.device)

    hold = int(math.ceil(float(hold_seconds) * AUDIO_HZ))
    release = int(math.ceil(_VOICE_GATE_RELEASE_SECONDS * AUDIO_HZ))
    expanded = torch.zeros_like(active)
    for tick in torch.nonzero(active, as_tuple=False).flatten().tolist():
        left = max(0, int(tick) - hold)
        right = min(ticks, int(tick) + hold + release + 1)
        expanded[left:right] = True
    forced = min(ticks, int(math.ceil(
        float(force_active_prefix_seconds) * AUDIO_HZ)))
    if forced:
        expanded[:forced] = True
    return expanded


class MiniMaxH3ContexMasterAudioMaskedAV:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT", {
                    "tooltip": "Target AV latent from the stock MiniMax H3 "
                               "conditioning node.",
                }),
                "audio_vae": ("VAE", {
                    "tooltip": "MiniMax H3 audio VAE used to encode the exact "
                               "master-audio slice into the target latent.",
                }),
                "master_audio": ("AUDIO", {
                    "tooltip": "Full prerecorded audio timeline (music, "
                               "dialogue, narration, or effects). The exact "
                               "current interval is inserted into the target "
                               "with the selected denoise mask.",
                }),
                "clip_start_seconds": ("FLOAT", {
                    "default": 0.0, "min": 0.0, "max": 99999.0,
                    "step": 0.001,
                    "tooltip": "Start of this raw H3 clip on the master-audio "
                               "timeline.",
                }),
                "context_length": ("INT", {
                    "default": 39, "min": 0, "max": 9999,
                    "tooltip": "Previous-video prefix request. Native runs "
                               "such as 5, 22, 39, 56... are used; 0 disables "
                               "video prefixing when neither source input is "
                               "connected.",
                }),
                "source_fps": ("FLOAT", {
                    "default": 24.0, "min": 1.0, "max": 240.0,
                    "step": 0.001,
                    "tooltip": "Frame rate of the legacy decoded "
                               "source_frames path. It is ignored when "
                               "source_latent is used.",
                }),
                "crop": (["disabled", "center"], {
                    "default": "disabled",
                    "tooltip": "Resize policy for legacy source_frames: "
                               "disabled stretches to the target canvas; "
                               "center preserves aspect ratio and center-crops. "
                               "Ignored when source_latent is used.",
                }),
                "preroll_seconds": ("FLOAT", {
                    "default": 0.0, "min": 0.0, "max": 4.0,
                    "step": 0.05, "round": 0.01,
                    "tooltip": "Optional real-song context encoded before "
                               "the kept clip ticks. 0 preserves legacy "
                               "hard-cut behavior."}),
                "lookahead_seconds": ("FLOAT", {
                    "default": 0.0, "min": 0.0, "max": 2.0,
                    "step": 0.05, "round": 0.01,
                    "tooltip": "Optional real-song context encoded after the "
                               "kept clip ticks. 0 preserves legacy behavior."}),
                "audio_denoise": ("FLOAT", {
                    "default": 0.0, "min": 0.0, "max": 1.0,
                    "step": 0.01, "round": 0.01,
                    "tooltip": "Denoise-mask value for the exact song, or for "
                               "detected vocal regions when voice is connected. "
                               "0 freezes the latent exactly."}),
                "gap_denoise": ("FLOAT", {
                    "default": 0.15, "min": 0.0, "max": 1.0,
                    "step": 0.01, "round": 0.01,
                    "tooltip": "Denoise-mask value between vocal phrases. "
                               "Ignored when voice is disconnected."}),
                "gate_hold_seconds": ("FLOAT", {
                    "default": 0.2, "min": 0.0, "max": 2.0,
                    "step": 0.05, "round": 0.01,
                    "tooltip": "Frozen margin around detected vocal ticks."}),
            },
            "optional": {
                "vae": ("VAE", {
                    "tooltip": "MiniMax H3 video VAE; required when previous "
                               "source_frames are connected.",
                }),
                "source_frames": ("IMAGE", {
                    "tooltip": "Legacy decoded continuation path. Its final "
                               "native H3 context run is VAE-encoded into the "
                               "protected video prefix. Prefer source_latent "
                               "for live H3 chaining.",
                }),
                "source_latent": ("LATENT", {
                    "tooltip": "Preferred live continuation path. The "
                               "phase-aligned tail of the previous sampled H3 "
                               "video latent is copied directly; its audio is "
                               "ignored because master_audio remains "
                               "authoritative.",
                }),
                "voice": ("AUDIO", {
                    "tooltip": "Optional isolated vocal stem aligned with "
                               "master_audio. Vocal ticks use audio_denoise; "
                               "between-phrase ticks use gap_denoise."}),
            },
        }

    RETURN_TYPES = ("LATENT", "INT", "AUDIO")
    RETURN_NAMES = ("latent", "trim_frames", "clip_audio")
    OUTPUT_TOOLTIPS = (
        "Sampler target with exact master audio, its selected denoise mask, "
        "and optional "
        "protected previous-video prefix.",
        "Actual protected visual prefix length; trim this many decoded frames.",
        "Exact master-audio interval represented by this raw target.",
    )
    FUNCTION = "prepare"
    CATEGORY = "conditioning/minimax/context_loop/masking"
    DESCRIPTION = (
        "Insert an exact master-audio interval into the complete H3 audio "
        "target with optional real-song encode context and vocal-aware "
        "denoising; optionally preserve a previous sampled-latent or "
        "decoded-video prefix while generating only future video rows."
    )

    def prepare(
        self,
        latent,
        audio_vae,
        master_audio,
        clip_start_seconds=0.0,
        context_length=39,
        source_fps=24.0,
        crop="disabled",
        preroll_seconds=0.0,
        lookahead_seconds=0.0,
        audio_denoise=0.0,
        gap_denoise=0.15,
        gate_hold_seconds=0.2,
        vae=None,
        source_frames=None,
        source_latent=None,
        voice=None,
        voice_clip_start_seconds=None,
        force_active_voice_prefix_seconds=0.0,
    ):
        require_h3_mask_support("exact master-audio latent masking")
        target_video, target_audio, target_frames = _validate_target_streams(
            latent, strict_audio_grid=False)
        if source_latent is not None and source_frames is not None:
            raise ValueError(
                "h3_master_audio_mask: connect either source_latent or "
                "source_frames, not both.")

        # The stock H3 target owns the audio-grid length. It can round up one
        # 40 Hz step beyond what a floor-style audio VAE produces from the
        # exact picture-duration waveform (for example 124 frames -> 207
        # target steps while a picture-only encode can return 206).
        expected_audio_steps = int(target_audio.shape[-1])
        nominal_audio_steps = int(round(
            target_frames / float(FPS) * AUDIO_HZ))
        if expected_audio_steps != nominal_audio_steps:
            _LOG.warning(
                "h3_master_audio_mask: target contains %d audio steps for "
                "%d frames; nominal 40 Hz calculation gives %d. Using the "
                "target length.", expected_audio_steps, target_frames,
                nominal_audio_steps)
        vae_rate = int(getattr(audio_vae, "audio_sample_rate", 32000))
        if not isinstance(master_audio, dict):
            raise ValueError(
                "h3_master_audio_mask: master_audio is not a Comfy AUDIO.")
        waveform = _stereo_first_batch(master_audio.get("waveform"))
        waveform = _resample(
            waveform, int(master_audio.get("sample_rate", 0)), vae_rate)

        options = _contract_masked_song_options(
            preroll_seconds=preroll_seconds,
            lookahead_seconds=lookahead_seconds,
            audio_denoise=audio_denoise,
            gap_denoise=gap_denoise,
            gate_hold_seconds=gate_hold_seconds,
        )
        start_seconds = float(clip_start_seconds)
        if start_seconds < 0.0:
            raise ValueError(
                "h3_master_audio_mask: clip_start_seconds must be >= 0.")
        start_sample = sample_boundary_from_seconds(start_seconds, vae_rate)
        picture_end_sample = sample_boundary_from_seconds(
            start_seconds + target_frames / float(FPS), vae_rate)
        picture_samples = picture_end_sample - start_sample
        if picture_samples < 1:
            raise RuntimeError(
                "h3_master_audio_mask: computed an empty master-audio slice.")
        audio_slice = _fit_audio_slice(
            waveform[..., start_sample:picture_end_sample],
            picture_samples,
        )

        # Encode optional real-song context around the kept interval, then crop
        # by complete 40 Hz ticks. With both values at zero this follows the
        # exact legacy start/grid path byte for byte.
        requested_preroll_samples = sample_boundary_from_seconds(
            options["preroll_seconds"], vae_rate)
        encode_start_sample = max(
            0, start_sample - requested_preroll_samples)
        actual_preroll_samples = start_sample - encode_start_sample
        preroll_steps = int(round(
            actual_preroll_samples / float(vae_rate) * AUDIO_HZ))
        lookahead_steps = int(math.ceil(
            options["lookahead_seconds"] * AUDIO_HZ))
        lookahead_samples = sample_boundary_from_seconds(
            options["lookahead_seconds"], vae_rate)
        required_encoded_steps = (
            preroll_steps + expected_audio_steps + lookahead_steps)
        grid_samples = int(math.ceil(
            required_encoded_steps / float(AUDIO_HZ) * vae_rate))
        contextual_samples = (
            picture_end_sample + lookahead_samples - encode_start_sample)
        encode_samples = max(contextual_samples, grid_samples)
        encode_slice = _fit_audio_slice(
            waveform[..., encode_start_sample:
                     encode_start_sample + encode_samples],
            encode_samples,
        )
        encoded_audio = audio_vae.encode(encode_slice.movedim(1, -1))
        if getattr(encoded_audio, "ndim", 0) != 4:
            raise ValueError(
                "h3_master_audio_mask: audio VAE returned %s; expected "
                "[B,C,2,T]." %
                (tuple(getattr(encoded_audio, "shape", ())),)
            )
        got_audio_steps = int(encoded_audio.shape[-1])
        if got_audio_steps < required_encoded_steps:
            missing = required_encoded_steps - got_audio_steps
            guard_samples = int(math.ceil(
                (missing + 1) * vae_rate / float(AUDIO_HZ)))
            retry_samples = encode_samples + guard_samples
            retry_slice = _fit_audio_slice(
                waveform[..., encode_start_sample:
                         encode_start_sample + retry_samples],
                retry_samples,
            )
            retry_audio = audio_vae.encode(retry_slice.movedim(1, -1))
            if getattr(retry_audio, "ndim", 0) != 4:
                raise ValueError(
                    "h3_master_audio_mask: audio VAE retry returned %s; "
                    "expected [B,C,2,T]." %
                    (tuple(getattr(retry_audio, "shape", ())),)
                )
            retry_steps = int(retry_audio.shape[-1])
            _LOG.warning(
                "h3_master_audio_mask: audio VAE initially produced %d/%d "
                "context steps; retried with %.2f ms of real-audio grid "
                "lookahead and got %d.", got_audio_steps,
                required_encoded_steps,
                guard_samples / float(vae_rate) * 1000.0, retry_steps)
            encoded_audio = retry_audio
            got_audio_steps = retry_steps

        if got_audio_steps < required_encoded_steps:
            raise RuntimeError(
                "h3_master_audio_mask: contextual encode requires %d audio "
                "steps but the "
                "audio VAE produced %d even after grid lookahead." %
                (required_encoded_steps, got_audio_steps)
            )
        if got_audio_steps > required_encoded_steps:
            _LOG.info(
                "h3_master_audio_mask: audio VAE produced %d steps for a %d-step "
                "context window; retaining the aligned scene interval.",
                got_audio_steps, required_encoded_steps)
        retained_end = preroll_steps + expected_audio_steps
        encoded_audio = encoded_audio[..., preroll_steps:retained_end]

        out_video = target_video.clone()
        out_audio = target_audio.clone()
        encoded_audio = encoded_audio[:1].to(
            device=out_audio.device, dtype=out_audio.dtype)
        if tuple(encoded_audio.shape) != tuple(out_audio.shape):
            raise ValueError(
                "h3_master_audio_mask: encoded master audio shape %s does not match "
                "target %s." %
                (tuple(encoded_audio.shape), tuple(out_audio.shape))
            )
        out_audio.copy_(encoded_audio)

        prefix_frames = 0
        video_steps = 0
        if source_latent is not None:
            if int(context_length) <= 0:
                raise ValueError(
                    "h3_master_audio_mask: context_length must be positive "
                    "when source_latent is connected.")
            source_parts = _streams_from_latent(source_latent)
            if len(source_parts) < 2:
                raise ValueError(
                    "h3_master_audio_mask: source_latent must be a joint H3 "
                    "video/audio latent.")
            source_video = source_parts[0]
            if source_video.ndim == 4:
                source_video = source_video.unsqueeze(0)
            if source_video.ndim != 5:
                raise ValueError(
                    "h3_master_audio_mask: source video latent must be "
                    "[B,C,T,H,W].")
            available = _pixel_frames(int(source_video.shape[2]))
            prefix_frames = _snap_prefix_length(
                context_length, available, target_frames)
            prefix, video_steps = _generated_video_tail(
                source_latent, prefix_frames, target_video)
            if video_steps >= int(out_video.shape[2]):
                raise ValueError(
                    "h3_master_audio_mask: video prefix consumes the complete "
                    "target.")
            out_video[:, :, :video_steps] = prefix.to(
                device=out_video.device, dtype=out_video.dtype)
        elif source_frames is not None:
            if vae is None:
                raise ValueError(
                    "h3_master_audio_mask: connect the video VAE when source_frames "
                    "is connected.")
            if (getattr(source_frames, "ndim", 0) != 4
                    or int(source_frames.shape[0]) < 1):
                raise ValueError(
                    "h3_master_audio_mask: source_frames must be IMAGE "
                    "[N,H,W,C].")
            if int(context_length) <= 0:
                raise ValueError(
                    "h3_master_audio_mask: context_length must be positive when "
                    "source_frames is connected.")
            indices = _cfr_index_map(
                int(source_frames.shape[0]), source_fps,
                source_frames.device)
            prefix_frames = _snap_prefix_length(
                context_length, int(indices.numel()), target_frames)
            tail = source_frames.index_select(0, indices[-prefix_frames:])
            width = int(target_video.shape[4]) * 16
            height = int(target_video.shape[3]) * 16
            prefix = vae.encode(_resize(tail, width, height, crop))
            if getattr(prefix, "ndim", 0) != 5:
                raise ValueError(
                    "h3_master_audio_mask: video VAE returned %s; expected "
                    "[B,C,T,H,W]." %
                    (tuple(getattr(prefix, "shape", ())),)
                )
            video_steps = int(prefix.shape[2])
            covered = _pixel_frames(video_steps)
            if covered != prefix_frames:
                raise RuntimeError(
                    "h3_master_audio_mask: %d source frames encoded to %d video "
                    "steps covering %d frames; refusing a shifted seam." %
                    (prefix_frames, video_steps, covered)
                )
            if video_steps >= int(out_video.shape[2]):
                raise ValueError(
                    "h3_master_audio_mask: video prefix consumes the complete "
                    "target.")
            prefix = prefix[:1].to(
                device=out_video.device, dtype=out_video.dtype)
            if (int(prefix.shape[1]) != int(out_video.shape[1])
                    or tuple(prefix.shape[3:]) != tuple(out_video.shape[3:])):
                raise ValueError(
                    "h3_master_audio_mask: video prefix shape %s does not match "
                    "target %s." %
                    (tuple(prefix.shape), tuple(out_video.shape))
                )
            out_video[:, :, :video_steps] = prefix

        # Compose with an upstream masked target instead of replacing it. This
        # is what lets Chain Context own the protected continuation-video
        # prefix while this node replaces and protects the complete current
        # master-audio interval. It also keeps both standalone video-prefix
        # paths composable with any future spatial/video mask.
        video_mask, _ = _existing_mask_streams(
            latent, out_video, out_audio)
        if video_steps:
            video_mask[:, :, :video_steps] = 0.0
        audio_mask = torch.full(
            (1, 1, int(out_audio.shape[2]), int(out_audio.shape[3])),
            float(options["audio_denoise"]),
            device=out_audio.device, dtype=torch.float32)
        gate_status = "complete song"
        if voice is not None:
            activity = _voice_activity_ticks(
                voice,
                (start_seconds if voice_clip_start_seconds is None
                 else float(voice_clip_start_seconds)),
                int(out_audio.shape[-1]), vae_rate,
                options["gate_hold_seconds"],
                force_active_prefix_seconds=(
                    force_active_voice_prefix_seconds),
            ).to(device=audio_mask.device)
            audio_mask.fill_(float(options["gap_denoise"]))
            audio_mask[..., activity] = float(options["audio_denoise"])
            gate_status = "%d/%d vocal ticks" % (
                int(activity.sum().item()), int(activity.numel()))

        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))
        clip_audio = {"waveform": audio_slice, "sample_rate": vae_rate}
        _LOG.info(
            "h3_master_audio_mask: target %d frames / %.3fs; master "
            "%.3f..%.3fs "
            "encoded with %.3fs pre-roll / %.3fs lookahead to %d audio steps; "
            "mask %.3f voice / %.3f gaps (%s); video prefix %d frames / %d "
            "steps.",
            target_frames, target_frames / float(FPS), start_seconds,
            start_seconds + target_frames / float(FPS),
            options["preroll_seconds"], options["lookahead_seconds"],
            expected_audio_steps, options["audio_denoise"],
            options["gap_denoise"], gate_status, prefix_frames, video_steps)
        return output, prefix_frames, clip_audio


NODE_CLASS_MAPPINGS = {
    "MiniMaxH3ContexMasterAudioMaskedAV": (
        MiniMaxH3ContexMasterAudioMaskedAV),
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MiniMaxH3ContexMasterAudioMaskedAV": (
        "MiniMax H3 Masking · Master Audio + Video Prefix"),
}