| """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.") |
|
|
| |
| |
| |
| |
| 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, |
| ) |
|
|
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| |
| |
| 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"), |
| } |
|
|
