"""H3 Motion Context clip stitcher for ComfyUI. Loads NikoDemon80/ComfyUI-H3-Motion-Context clip archive files (h3_motion_context_av_v1), decodes each approved clip once, crossfades the carried Motion Context head from clips, and concatenates the picture/audio into one IMAGE + AUDIO pair. This intentionally does NOT reconstruct a NestedTensor and feed the saved files back into Motion Context. The archive format is the sampler output, and this node is a final-media assembly tool. """ import glob import logging import os import re import torch import torch.nn.functional as F import folder_paths try: from safetensors.torch import load_file as st_load except Exception: st_load = None try: import torchaudio except Exception: torchaudio = None log_ = logging.getLogger("h3_motion_context_archive") def _resolve_folder(path): p = (path or "").strip().strip('"').strip("'") if not p: p = "h3_context" candidates = [p, os.path.join(folder_paths.get_output_directory(), p)] for c in candidates: if os.path.isdir(c): return os.path.abspath(c) raise FileNotFoundError("H3 Motion Context Clip Stitcher: folder not found: %s\n" "You can use an absolute path or a path relative to " "ComfyUI's output folder." % p) def _clip_number(path): name = os.path.basename(path) # noinspection RegExpUnnecessaryNonCapturingGroup pat = re.compile(r"(?:^|_)(\d{5})(?:\.safetensors)$", re.IGNORECASE) m = pat.search(name) return int(m.group(1)) if m else -1 def _find_files(folder, pattern, first_clip, last_clip): pattern = (pattern or "clip_*.safetensors").strip() paths = [] for p in glob.glob(os.path.join(folder, pattern)): if not os.path.isfile(p): continue if not p.lower().endswith(".safetensors"): continue idx = _clip_number(p) if idx < 0: continue if idx < int(first_clip): continue if 0 < int(last_clip) < idx: continue paths.append((idx, p)) paths.sort(key=lambda x: x[0]) if not paths: raise FileNotFoundError("H3 Motion Context Clip Stitcher: " "no numbered .safetensors files matched '%s' in %s." % (pattern, folder)) # Do not silently skip a missing numbered clip. A gap usually means an # approved clip was not saved, and silently stitching around it would make # a misleading final timeline. expected = paths[0][0] for idx, _ in paths: if idx != expected: raise ValueError("H3 Motion Context Clip Stitcher: missing clip %05d between " "the selected archive files." % expected) expected += 1 return paths def _load_archive(path): if st_load is None: raise RuntimeError("safetensors is unavailable in this " "ComfyUI Python environment.") # noinspection PyCallingNonCallable data = st_load(path, device="cpu") if "video" not in data or "audio" not in data: raise ValueError("%s is not an h3_motion_context_av_v1 archive: " "expected 'video' and 'audio'." % path) video = data["video"] audio = data["audio"] if video.ndim != 5: raise ValueError("%s: expected video [B,C,T,H,W], got %s" % (path, tuple(video.shape))) if audio.ndim != 4: raise ValueError("%s: expected audio [B,C,2,T], got %s" % (path, tuple(audio.shape))) if video.shape[0] != 1 or audio.shape[0] != 1: raise ValueError("%s: only batch size 1 archive clips are supported." % path) return video, audio def _decode_video(vae, video_latent): """ Decode the H3 video stream and normalize to ComfyUI IMAGE format. """ images = vae.decode(video_latent) # H3's VAE normally returns [B,T,H,W,C]. Some VAE implementations can # return [T,H,W,C], so accept both. if images.ndim == 5: images = images.reshape(-1, *images.shape[-3:]) elif images.ndim != 4: raise RuntimeError("H3 video VAE returned unexpected shape %s" % (tuple(images.shape),)) return images.to(torch.float32).clamp(0, 1).cpu() def _decode_audio(audio_vae, audio_latent): """ Decode the H3 audio stream using the same convention as ComfyUI's VAEDecodeAudio. """ audio = audio_vae.decode(audio_latent) # Current ComfyUI audio VAE returns [B,L,C]. Convert to [B,C,L]. if audio.ndim != 3: raise RuntimeError( "H3 audio VAE returned unexpected shape %s" % (tuple(audio.shape),)) audio = audio.movedim(-1, 1) sr = int(getattr(audio_vae, "audio_sample_rate_output", getattr(audio_vae, "audio_sample_rate", 32000))) return {"waveform": audio.to(torch.float32).cpu(), "sample_rate": sr} def _resample_audio(audio, target_sr): if audio is None: return None sr = int(audio["sample_rate"]) if sr == int(target_sr): return audio if torchaudio is None: raise RuntimeError("Audio sample rates differ (%d vs %d), but torchaudio is " "unavailable to resample them." % (sr, int(target_sr))) # noinspection PyUnresolvedReferences waveform = torchaudio.functional.resample(audio["waveform"], sr, int(target_sr)) return {"waveform": waveform, "sample_rate": int(target_sr)} def _crossfade_boundary(prev_tail_images, cur_images, prev_tail_wave, cur_wave, overlap_frames, cross_samples): """ Crossfade the previous clip's tail with the current clip's head. prev_tail_images: [L,H,W,C] cur_images: [T,H,W,C] prev_tail_wave : [1,C,Ls] cur_wave: [1,C,Cs] (or None) Returns (blend_images [L,H,W,C], blend_wave [1,C,Ls] or None). Video uses a linear dissolve ramp; audio uses an equal-power (cos/sin) ramp over the same time window so picture and sound stay in sync. """ L = int(overlap_frames) if L <= 0: return cur_images[:0], None if L == 1: alpha = torch.full((1, 1, 1, 1), 0.5, dtype=prev_tail_images.dtype, device=prev_tail_images.device) else: alpha = torch.linspace(0.0, 1.0, L, dtype=prev_tail_images.dtype, device=prev_tail_images.device).view(L, 1, 1, 1) blend_images = prev_tail_images * (1.0 - alpha) + cur_images[:L] * alpha blend_wave = None if prev_tail_wave is not None and cur_wave is not None: n = int(cross_samples) if n <= 0: blend_wave = prev_tail_wave else: n = min(n, int(prev_tail_wave.shape[-1]), int(cur_wave.shape[-1])) theta = torch.linspace(0.0, 1.5707963267948966, n, dtype=prev_tail_wave.dtype, device=prev_tail_wave.device).view(1, 1, n) blend_wave = (prev_tail_wave[..., :n] * torch.cos(theta) + cur_wave[..., :n] * torch.sin(theta)) return blend_images, blend_wave def _av_from_live_latent(latent): """ Extract (video, audio) tensors from an in-memory H3 AV LATENT, using the same unpacking convention as NikoDemon80's own _streams_from_latent()/save(): latent["samples"] is a NestedTensor (or tuple/list) whose unbind() gives (video, audio) in that order. """ if not isinstance(latent, dict) or "samples" not in latent: raise ValueError("h3_motion_context: expected a MiniMax H3 AV latent dict with " "a 'samples' key, got %r" % type(latent)) samples = latent["samples"] if hasattr(samples, "unbind"): parts = list(samples.unbind()) elif isinstance(samples, (tuple, list)): parts = list(samples) else: raise ValueError("h3_motion_context: expected a MiniMax H3 AV latent (a nested " "video/audio pair), got %r" % type(samples)) if len(parts) < 2: raise ValueError("h3_motion_context: latent has no audio stream; wire the " "sampler output of an H3 AV graph.") video = parts[0].cpu().contiguous() audio = parts[1].cpu().contiguous() return video, audio class H3MotionContextClipStitcher: """ Load, decode, and crossfade approved H3 Motion Context clips. """ @classmethod def INPUT_TYPES(cls): return {"required": {"folder": ("STRING", {"default": "h3_context", "tooltip": "Folder containing clip_00001.safetensors, " "clip_00002.safetensors, etc.\nAbsolute paths and paths relative " "to ComfyUI/output are accepted."}), "pattern": ("STRING", {"default": "clip_*.safetensors", "tooltip": "Filename glob. The final five-digit number is treated as " "the clip index."}), "first_clip": ("INT", {"default": 1, "min": 1, "max": 9999, "tooltip": "First approved clip to include."}), "last_clip": ("INT", {"default": 0, "min": 0, "max": 9999, "tooltip": "Last clip to include. 0 = every clip from first_clip onward."}), "context_length": (["5", "22", "39", "56"], {"default": "22", "tooltip": "Number of decoded frames to crossfade at each clip boundary. " "The normal setting is 22 frames.\n" "This is the overlap length that is dissolved between " "adjacent clips.\n" "5, 22, 39 or 56 are the lengths that are a whole number of " "latent steps, which is why other numbers aren't offered."}), "fps": ("FLOAT", {"default": 24.0, "min": 1.0, "max": 240.0, "step": 0.001, "tooltip": "H3 native output rate. Keep this at 24 unless your workflow " "deliberately changes it."}), }, "optional": {"video_vae": ("VAE", {"tooltip": "MiniMax H3 video VAE " "(FP16 or INT8 ConvRot)."}), "audio_vae": ("VAE", { "tooltip": "MiniMax H3 audio VAE FP32. Required for the AUDIO " "output."}), "latent": ("LATENT", { "tooltip": "Optional: the currently-generated AV latent (from your " "H3 sampler), used in place of the highest-numbered file " "on disk."}), }, } RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING") RETURN_NAMES = ("images", "audio", "frame_count", "report") FUNCTION = "stitch" CATEGORY = "video/minimax" DESCRIPTION = ("Final assembly for NikoDemon80's H3 Motion Context AV clip archives.\n" "Loads numbered h3_motion_context_av_v1 files, decodes one clip at a " "time, dissolves the overlap between adjacent clips (video + synchronized audio), " "and concatenates them to a final video and audio stream.") # noinspection PyUnusedLocal @classmethod def IS_CHANGED(cls, folder, pattern, first_clip, last_clip, context_length, fps, video_vae=None, audio_vae=None,): # noinspection PyBroadException try: d = _resolve_folder(folder) files = _find_files(d, pattern, first_clip, last_clip) # noinspection PyTypeChecker return tuple((p, os.stat(p).st_mtime_ns, os.path.getsize(p)) for _, p in files) + (int(context_length), float(fps),) except Exception: return float("NaN") @staticmethod def stitch(folder, pattern, first_clip, last_clip, context_length, fps, video_vae=None, audio_vae=None, latent=None,): if video_vae is None: raise ValueError("Connect your MiniMax H3 video VAE to 'video_vae'.") if st_load is None: raise RuntimeError( "safetensors is not available in this ComfyUI environment.") d = _resolve_folder(folder) files = _find_files(d, pattern, first_clip, last_clip) live_entry = None if latent is not None: if files: files = files[:-1] # drop the presumed-duplicate on-disk file live_index = files[-1][0] + 1 if files else int(first_clip) else: live_index = int(first_clip) video_latent, audio_latent = _av_from_live_latent(latent) live_entry = ( live_index, None, video_latent, audio_latent) # path=None marks it as live image_parts = [] audio_parts = [] report_lines = [] target_sr = None overlap = int(context_length) prev_tail_img = None prev_tail_wave = None all_entries = [(idx, path, None, None) for idx, path in files] if live_entry is not None: # noinspection PyTypeChecker all_entries.append(live_entry) total_count = len(files) + (1 if live_entry is not None else 0) log_.info("H3 archive stitcher: %d clip(s) selected from %s", total_count, d) # Degenerate crossfade (no overlap or a single clip) falls back to a # plain concatenation, which is exactly what the trim modes would do. crossfade_active = overlap > 0 and len(all_entries) > 1 for pos, (idx, path, live_video, live_audio) in enumerate(all_entries): if path is not None: video_latent, audio_latent = _load_archive(path) else: video_latent, audio_latent = live_video, live_audio log_.info("H3 archive stitcher: clip %05d taken from live latent input", idx) # Decode one clip at a time. The decoded result is immediately moved # to CPU, so a long chain does not keep every VAE result on VRAM. images = _decode_video(video_vae, video_latent) del video_latent audio = None if audio_vae is not None: audio = _decode_audio(audio_vae, audio_latent, # normalize=normalize_audio_per_clip ) del audio_latent decoded_frames = int(images.shape[0]) is_last = pos == len(all_entries) - 1 if crossfade_active: if decoded_frames < 2 * overlap: raise ValueError("Crossfade requires each clip to have at least " "2*context_length (%d) frames; clip %05d has %d." % (overlap, idx, decoded_frames)) # Resample this clip's audio to the shared target rate before # splitting, so the head/tail sample counts line up across clips. if audio is not None: if target_sr is None: target_sr = int(audio["sample_rate"]) audio = _resample_audio(audio, target_sr) if prev_tail_wave is not None: prev_tail_wave = _resample_audio(prev_tail_wave, target_sr) n = 0 if audio is not None: sr = int(audio["sample_rate"]) n = int(round((overlap / float(fps)) * sr)) if n <= 0: n = 1 if n >= audio["waveform"].shape[-1]: raise ValueError("Audio is too short to extract a %d-frame " "(%0.4fs) crossfade head/tail for clip %05d." % (overlap, overlap / float(fps), idx)) head_img = images[:overlap] body_img = images[overlap:-overlap] tail_img = images[-overlap:] head_wave = body_wave = tail_wave = None if audio is not None: wave = audio["waveform"] sr = int(audio["sample_rate"]) head_wave = {"waveform": wave[..., :n], "sample_rate": sr} body_wave = {"waveform": wave[..., n:-n], "sample_rate": sr} tail_wave = {"waveform": wave[..., -n:], "sample_rate": sr} if pos == 0: # First clip: emit head+body raw, buffer the tail for the next boundary. image_parts.append(torch.cat([head_img, body_img], dim=0)) if audio is not None: audio_parts.append(torch.cat([head_wave["waveform"], body_wave["waveform"]], dim=-1)) prev_tail_img = tail_img prev_tail_wave = tail_wave else: blend_img, blend_wave = _crossfade_boundary(prev_tail_img, images, prev_tail_wave[ "waveform"] if prev_tail_wave is not None else None, audio["waveform"] if audio is not None else None, overlap, n) image_parts.append(blend_img) if audio is not None: audio_parts.append(blend_wave) audio_parts.append(body_wave["waveform"]) if is_last: # Last clip: emit body+tail raw after its boundary blend. image_parts.append(torch.cat([body_img, tail_img], dim=0)) if audio is not None: audio_parts.append(tail_wave["waveform"]) else: image_parts.append(body_img) prev_tail_img = tail_img prev_tail_wave = tail_wave kept_frames = decoded_frames - (overlap if not is_last else 0) audio_sec = (0.0 if audio is None else audio["waveform"].shape[-1] / float(audio["sample_rate"])) report_lines.append("clip_%05d: decoded=%d frames, crossfade=%d frames " "(%.4fs), kept=%d, audio=%.4fs" % (idx, decoded_frames, overlap, overlap / float(fps), kept_frames, audio_sec)) del images if audio is not None: del audio final_images = torch.cat(image_parts, dim=0).contiguous() del image_parts final_audio = None if audio_parts: final_waveform = torch.cat(audio_parts, dim=-1).contiguous() del audio_parts final_audio = {"waveform": final_waveform, "sample_rate": int(target_sr)} frame_count = int(final_images.shape[0]) video_seconds = frame_count / float(fps) audio_seconds = (final_audio["waveform"].shape[-1] / float( final_audio["sample_rate"]) if final_audio is not None else 0.0) report_lines.append("TOTAL: %d frames = %.4fs at %.3f fps; audio=%.4fs%s" % (frame_count, video_seconds, float(fps), audio_seconds, "" if final_audio is not None else " (no audio_vae connected)")) report = "\n".join(report_lines) log_.info("H3 archive stitcher finished: %d frames (%.3fs), audio %.3fs", frame_count, video_seconds, audio_seconds) return final_images, final_audio, frame_count, report NODE_CLASS_MAPPINGS = { "H3MotionContextClipStitcher": H3MotionContextClipStitcher, } NODE_DISPLAY_NAME_MAPPINGS = { "H3MotionContextClipStitcher": "H3 Motion Context Clip Stitcher", }