| """Frame-locked source AV targets for recursive MiniMax H3 editing. |
| |
| The stock H3 conditioning node owns the exact joint video/audio target grid. |
| This node selects the current Chain Loop interval from a source video, encodes |
| its picture and sound, and copies both streams into that authoritative target. |
| It deliberately avoids LTX AV concat/separate helpers: independently encoded |
| streams can differ by one audio token at fractional 24-fps/40-Hz boundaries. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| import math |
|
|
| import torch |
|
|
| from .masked_context import _validate_target_streams |
| from .masking_ops import normalize_comfy_mask |
| from .nodes import AUDIO_HZ, FPS, _resize |
|
|
|
|
| try: |
| import torchaudio |
| except ImportError: |
| torchaudio = None |
|
|
|
|
| _LOG = logging.getLogger("minimax_h3_context_loop.source_av_target") |
| STATE_TYPE = "H3_CHAIN_STATE" |
|
|
|
|
| def _canonical_indices(frame_count, source_fps, device): |
| source_fps = float(source_fps) |
| if not math.isfinite(source_fps) or source_fps <= 0.0: |
| raise ValueError("h3_source_av_target: source_fps must be positive.") |
| frame_count = int(frame_count) |
| if frame_count < 1: |
| raise ValueError("h3_source_av_target: 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) |
| source = torch.round((frame + 0.5) * source_fps / FPS - 0.5) |
| return source.to(torch.long).clamp_(0, frame_count - 1) |
|
|
|
|
| def _stereo_waveform(audio, target_rate): |
| if not isinstance(audio, dict): |
| raise ValueError( |
| "h3_source_av_target: source_audio must be a ComfyUI AUDIO.") |
| waveform = audio.get("waveform") |
| if getattr(waveform, "ndim", 0) != 3: |
| raise ValueError( |
| "h3_source_av_target: source audio must be [B,C,L], got %s." % |
| (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_source_av_target: source audio has %d channels; downmix it " |
| "to stereo first." % channels) |
| source_rate = int(audio.get("sample_rate", 0)) |
| target_rate = int(target_rate) |
| if source_rate <= 0: |
| raise ValueError( |
| "h3_source_av_target: source audio has no valid sample rate.") |
| if source_rate != target_rate: |
| if torchaudio is None: |
| raise RuntimeError( |
| "h3_source_av_target: source audio is %d Hz but the H3 audio " |
| "VAE needs %d Hz and torchaudio is unavailable." % |
| (source_rate, target_rate)) |
| waveform = torchaudio.functional.resample( |
| waveform, source_rate, target_rate) |
| return waveform |
|
|
|
|
| def _fit_slice(waveform, start, samples, allow_tail_padding, label): |
| start = int(start) |
| samples = int(samples) |
| if start < 0 or samples < 0: |
| raise ValueError("h3_source_av_target: invalid audio slice.") |
| available = max(0, int(waveform.shape[-1]) - start) |
| if available < samples and not allow_tail_padding: |
| raise ValueError( |
| "h3_source_av_target: source audio is too short for %s; need %d " |
| "samples from offset %d but only %d remain." % |
| (label, samples, start, available)) |
| result = waveform[..., start:start + samples] |
| missing = samples - int(result.shape[-1]) |
| if missing > 0: |
| result = torch.nn.functional.pad(result, (0, missing)) |
| return result, missing |
|
|
|
|
| def _encode_audio_to_target( |
| audio_vae, |
| waveform, |
| start_frame, |
| frame_count, |
| target_audio, |
| ): |
| sample_rate = int(getattr(audio_vae, "audio_sample_rate", 32000)) |
| start_sample = int(round(int(start_frame) / FPS * sample_rate)) |
| picture_samples = int(round(int(frame_count) / FPS * sample_rate)) |
| |
| |
| picture, _ = _fit_slice( |
| waveform, start_sample, picture_samples, False, "the picture interval") |
| target_steps = int(target_audio.shape[-1]) |
| grid_samples = int(math.ceil( |
| target_steps / float(AUDIO_HZ) * sample_rate)) |
| encode_samples = max(picture_samples, grid_samples) |
| encoded_source, padded = _fit_slice( |
| waveform, start_sample, encode_samples, True, "the H3 audio grid") |
| encoded = audio_vae.encode(encoded_source.movedim(1, -1)) |
| if getattr(encoded, "ndim", 0) != 4: |
| raise ValueError( |
| "h3_source_av_target: audio VAE returned %s; expected " |
| "[B,C,2,T]." % (tuple(getattr(encoded, "shape", ())),)) |
| if int(encoded.shape[-1]) < target_steps: |
| missing_steps = target_steps - int(encoded.shape[-1]) |
| guard_samples = int(math.ceil( |
| (missing_steps + 1) * sample_rate / float(AUDIO_HZ))) |
| encoded_source, retry_padding = _fit_slice( |
| waveform, start_sample, encode_samples + guard_samples, True, |
| "the H3 audio-grid retry") |
| padded = max(padded, retry_padding) |
| encoded = audio_vae.encode(encoded_source.movedim(1, -1)) |
| if getattr(encoded, "ndim", 0) != 4 or int(encoded.shape[-1]) < target_steps: |
| raise RuntimeError( |
| "h3_source_av_target: audio VAE cannot fill the stock %d-step " |
| "target (produced %d)." % |
| (target_steps, int(getattr(encoded, "shape", (0, 0, 0, 0))[-1]))) |
| encoded = encoded[:1, ..., :target_steps].to( |
| device=target_audio.device, dtype=target_audio.dtype) |
| if tuple(encoded.shape) != tuple(target_audio.shape): |
| raise ValueError( |
| "h3_source_av_target: encoded audio shape %s does not match the " |
| "stock target %s." % |
| (tuple(encoded.shape), tuple(target_audio.shape))) |
| clip_audio = { |
| "waveform": picture, |
| "sample_rate": sample_rate, |
| } |
| return encoded, clip_audio, padded |
|
|
|
|
| class MiniMaxH3ContexLoopSourceAVTarget: |
| """Replace the current loop target with its exact source AV interval.""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "state": (STATE_TYPE, { |
| "tooltip": "Current state from H3 Chain Current Shot."}), |
| "latent": ("LATENT", { |
| "tooltip": "Current stock H3 joint target, normally the " |
| "latent output from Chain Context."}), |
| "vae": ("VAE", {"tooltip": "MiniMax H3 video VAE."}), |
| "audio_vae": ("VAE", {"tooltip": "MiniMax H3 audio VAE."}), |
| "source_frames": ("IMAGE", { |
| "tooltip": "Complete decoded source-video timeline."}), |
| "source_audio": ("AUDIO", { |
| "tooltip": "Audio from the same source video."}), |
| "source_fps": ("FLOAT", { |
| "default": 24.0, "min": 0.001, "max": 1000.0, |
| "step": 0.001, |
| "tooltip": "Actual FPS of source_frames. Connect the FPS " |
| "output from Get Video Components."}), |
| "crop": (["disabled", "center"], { |
| "default": "center", |
| "tooltip": "Resize policy for the selected source-video " |
| "interval: disabled stretches to the target " |
| "canvas; center preserves aspect ratio and " |
| "center-crops.", |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("LATENT", "IMAGE", "AUDIO", "STRING") |
| RETURN_NAMES = ("source_target", "scene_frames", "scene_audio", "status") |
| OUTPUT_TOOLTIPS = ( |
| "Current stock H3 target with the exact source-video and source-audio " |
| "scene interval encoded into it.", |
| "The resized 24 fps source frames selected for the current scene.", |
| "The synchronized source-audio interval selected for the current scene.", |
| "Summary of the selected timeline range, target grid, and any audio " |
| "tail padding.", |
| ) |
| FUNCTION = "prepare" |
| CATEGORY = "conditioning/minimax/context_loop/masking" |
| DESCRIPTION = ( |
| "Select the current Chain Loop interval from a source video and copy " |
| "frame-locked video/audio encodes into the stock H3 joint target.") |
|
|
| def prepare( |
| self, |
| state, |
| latent, |
| vae, |
| audio_vae, |
| source_frames, |
| source_audio, |
| source_fps=24.0, |
| crop="center", |
| ): |
| target_video, target_audio, target_frames = _validate_target_streams( |
| latent, strict_audio_grid=False) |
| if getattr(source_frames, "ndim", 0) != 4: |
| raise ValueError( |
| "h3_source_av_target: source_frames must be IMAGE " |
| "[N,H,W,C].") |
| try: |
| index = int(state["index"]) |
| shot = state["plan"]["shots"][index - 1] |
| start_frame = int(shot["generation_start_frame"]) |
| planned_frames = int(shot["raw_frames"]) |
| except (KeyError, IndexError, TypeError, ValueError) as exc: |
| raise ValueError( |
| "h3_source_av_target: invalid Chain Current Shot state.") from exc |
| if planned_frames != target_frames: |
| raise RuntimeError( |
| "h3_source_av_target: scene %d plans %d frames but the stock " |
| "target covers %d." % (index, planned_frames, target_frames)) |
| if start_frame < 0: |
| raise ValueError( |
| "h3_source_av_target: scene %d begins before the source " |
| "timeline; external-prelude editing is unsupported." % index) |
|
|
| indices = _canonical_indices( |
| int(source_frames.shape[0]), source_fps, source_frames.device) |
| end_frame = start_frame + target_frames |
| if end_frame > int(indices.numel()): |
| raise ValueError( |
| "h3_source_av_target: scene %d needs canonical source frames " |
| "%d..%d, but the video has only %d frames at 24 fps." % |
| (index, start_frame, end_frame - 1, int(indices.numel()))) |
| selected = source_frames.index_select( |
| 0, indices[start_frame:end_frame]) |
| width = int(target_video.shape[4]) * 16 |
| height = int(target_video.shape[3]) * 16 |
| scene_frames = _resize(selected, width, height, crop) |
| encoded_video = vae.encode(scene_frames) |
| if getattr(encoded_video, "ndim", 0) != 5: |
| raise ValueError( |
| "h3_source_av_target: video VAE returned %s; expected " |
| "[B,C,T,H,W]." % |
| (tuple(getattr(encoded_video, "shape", ())),)) |
| encoded_video = encoded_video[:1].to( |
| device=target_video.device, dtype=target_video.dtype) |
| if tuple(encoded_video.shape) != tuple(target_video.shape): |
| raise ValueError( |
| "h3_source_av_target: encoded video shape %s does not match " |
| "the stock target %s." % |
| (tuple(encoded_video.shape), tuple(target_video.shape))) |
|
|
| audio_rate = int(getattr(audio_vae, "audio_sample_rate", 32000)) |
| waveform = _stereo_waveform(source_audio, audio_rate) |
| encoded_audio, scene_audio, padded = _encode_audio_to_target( |
| audio_vae, waveform, start_frame, target_frames, target_audio) |
|
|
| import comfy.nested_tensor |
|
|
| output = latent.copy() |
| output["samples"] = comfy.nested_tensor.NestedTensor(( |
| encoded_video.clone(), encoded_audio.clone())) |
| status = ( |
| "H3 loop source target scene %d: source frames %d..%d -> %d " |
| "frames at %dx%d; stock target %d video / %d audio steps; audio " |
| "grid tail padding %d samples." |
| % (index, start_frame, end_frame - 1, target_frames, width, height, |
| int(target_video.shape[2]), int(target_audio.shape[-1]), padded)) |
| _LOG.info(status) |
| return output, scene_frames, scene_audio, status |
|
|
|
|
| class MiniMaxH3ContexLoopMaskSlice: |
| """Select the tracked mask interval matching the current loop scene.""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "state": (STATE_TYPE, { |
| "tooltip": "Current state from H3 Chain Current Shot."}), |
| "mask": ("MASK", { |
| "tooltip": "One static mask or a complete tracked mask " |
| "timeline. White means generate."}), |
| "source_fps": ("FLOAT", { |
| "default": 24.0, "min": 0.001, "max": 1000.0, |
| "step": 0.001, |
| "tooltip": "FPS represented by a tracked mask batch. A " |
| "single mask is broadcast and ignores FPS."}), |
| }, |
| } |
|
|
| RETURN_TYPES = ("MASK", "INT", "STRING") |
| RETURN_NAMES = ("scene_mask", "scene_mask_frames", "status") |
| OUTPUT_TOOLTIPS = ( |
| "Static mask broadcast or tracked-mask slice aligned to the current " |
| "scene's complete raw frame interval.", |
| "Number of mask frames returned for the current scene.", |
| "Summary of the selected scene interval and whether the mask was " |
| "broadcast or sliced.", |
| ) |
| FUNCTION = "slice" |
| CATEGORY = "conditioning/minimax/context_loop/masking" |
| DESCRIPTION = ( |
| "Broadcast one static mask or slice the exact tracked-mask interval " |
| "for the current Chain Loop scene and its continuation overlap.") |
|
|
| def slice(self, state, mask, source_fps=24.0): |
| timeline = normalize_comfy_mask(mask) |
| try: |
| index = int(state["index"]) |
| shot = state["plan"]["shots"][index - 1] |
| start_frame = int(shot["generation_start_frame"]) |
| scene_frames = int(shot["raw_frames"]) |
| except (KeyError, IndexError, TypeError, ValueError) as exc: |
| raise ValueError( |
| "h3_loop_mask_slice: invalid Chain Current Shot state.") from exc |
| if start_frame < 0: |
| raise ValueError( |
| "h3_loop_mask_slice: scene %d begins before the mask " |
| "timeline." % index) |
|
|
| mask_frames = int(timeline.shape[0]) |
| if mask_frames == 1: |
| selected = timeline.expand( |
| scene_frames, int(timeline.shape[1]), int(timeline.shape[2]) |
| ).clone() |
| mode = "static mask broadcast" |
| else: |
| indices = _canonical_indices( |
| mask_frames, source_fps, timeline.device) |
| end_frame = start_frame + scene_frames |
| if end_frame > int(indices.numel()): |
| raise ValueError( |
| "h3_loop_mask_slice: scene %d needs canonical mask frames " |
| "%d..%d, but the tracked mask has only %d frames at 24 " |
| "fps." % |
| (index, start_frame, end_frame - 1, |
| int(indices.numel()))) |
| selected = timeline.index_select( |
| 0, indices[start_frame:end_frame]) |
| mode = "tracked mask slice" |
|
|
| status = ( |
| "H3 loop %s scene %d: source mask batch %d frames -> timeline " |
| "%d..%d (%d scene masks)." % |
| (mode, index, mask_frames, start_frame, |
| start_frame + scene_frames - 1, scene_frames)) |
| _LOG.info(status) |
| return selected.contiguous(), scene_frames, status |
|
|
|
|
| NODE_CLASS_MAPPINGS = { |
| "MiniMaxH3ContexLoopSourceAVTarget": MiniMaxH3ContexLoopSourceAVTarget, |
| "MiniMaxH3ContexLoopMaskSlice": MiniMaxH3ContexLoopMaskSlice, |
| } |
|
|
| NODE_DISPLAY_NAME_MAPPINGS = { |
| "MiniMaxH3ContexLoopSourceAVTarget": ( |
| "MiniMax H3 Masking · Loop Source AV Target"), |
| "MiniMaxH3ContexLoopMaskSlice": ( |
| "MiniMax H3 Masking · Loop Mask Slice"), |
| } |
|
|
