From aaad09ddaef6ea2827e8189de5aa999862c63fdb Mon Sep 17 00:00:00 2001 From: Kangle <147648913+RCAKangle@users.noreply.github.com> Date: Fri, 23 Jan 2026 13:08:25 +0000 Subject: [PATCH] Add interleave frames node --- README.md | 18 +++++++++++++ __init__.py | 4 ++- nodes.py | 78 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index fa52013..63d65fb 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ batches. The focus is reliability and compatibility with real-world workflows. - Split videos into ComfyUI IMAGE batches using FFmpeg - Merge IMAGE batches into MP4 (H.264) videos using FFmpeg +- Interleave two frame batches (IMAGE or LATENT) to alternate frames - Optional FPS override for both split and merge - Outputs are compatible with Preview Image, VAE Encode, Save Image, and Save Video @@ -41,6 +42,20 @@ Notes: - If a VIDEO input is provided, it is passed through as `video` and its decoded frames are returned on `frames`. +### Interleave Frames + +Inputs: +- `batch_a` (IMAGE or LATENT) +- `batch_b` (IMAGE or LATENT) + +Outputs: +- `batch` (IMAGE or LATENT) + +Notes: +- Inputs must be the same type (IMAGE+IMAGE or LATENT+LATENT). +- If batch sizes differ, the node truncates to the shorter length and logs a warning. +- Output ordering is `A1, B1, A2, B2, ...`. + ## Requirements - FFmpeg installed and available on PATH @@ -63,6 +78,9 @@ Notes: - Frames to VHS Video Combine: - `IMAGE batch` -> `FFmpeg Merge Frames` -> `frames` -> `VHS Video Combine` +- Interleave two batches: + - `batch_a` + `batch_b` -> `Interleave Frames` -> `batch` + ## Troubleshooting - "FFmpeg not found": install FFmpeg and make sure `ffmpeg` is in PATH. diff --git a/__init__.py b/__init__.py index ca17b68..1b282c8 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,13 @@ -from .nodes import FFmpegMergeFrames, FFmpegSplitVideo +from .nodes import FFmpegMergeFrames, FFmpegSplitVideo, InterleaveFrames NODE_CLASS_MAPPINGS = { "FFmpegSplitVideo": FFmpegSplitVideo, "FFmpegMergeFrames": FFmpegMergeFrames, + "InterleaveFrames": InterleaveFrames, } NODE_DISPLAY_NAME_MAPPINGS = { "FFmpegSplitVideo": "FFmpeg Split Video", "FFmpegMergeFrames": "FFmpeg Merge Frames", + "InterleaveFrames": "Interleave Frames", } diff --git a/nodes.py b/nodes.py index a096268..6f22df6 100644 --- a/nodes.py +++ b/nodes.py @@ -1,4 +1,5 @@ import io +import logging import os import shutil import tempfile @@ -25,6 +26,52 @@ def _is_video_input(value: Any) -> bool: return hasattr(value, "get_stream_source") and hasattr(value, "get_components") +def _interleave_tensors(batch_a: torch.Tensor, batch_b: torch.Tensor) -> torch.Tensor: + if batch_a.dim() < 1 or batch_b.dim() < 1: + raise ValueError("Inputs must include a batch dimension.") + if batch_a.shape[0] != batch_b.shape[0]: + logging.warning( + "Interleave Frames: batch sizes differ (%s vs %s); truncating.", + batch_a.shape[0], + batch_b.shape[0], + ) + count = min(batch_a.shape[0], batch_b.shape[0]) + if count == 0: + raise ValueError("Inputs must contain at least one frame.") + batch_a = batch_a[:count] + batch_b = batch_b[:count] + return torch.stack([batch_a, batch_b], dim=1).reshape(2 * count, *batch_a.shape[1:]) + + +def _detect_batch_type(value: Any) -> str: + if isinstance(value, dict): + if "samples" not in value: + raise ValueError("LATENT input is missing 'samples'.") + return "LATENT" + if isinstance(value, torch.Tensor): + return "IMAGE" + raise ValueError("Input must be IMAGE or LATENT.") + + +def _interleave_latents(batch_a: dict, batch_b: dict) -> dict: + samples_a = batch_a["samples"] + samples_b = batch_b["samples"] + if not isinstance(samples_a, torch.Tensor) or not isinstance(samples_b, torch.Tensor): + raise ValueError("LATENT samples must be torch tensors.") + + output = {"samples": _interleave_tensors(samples_a, samples_b)} + for key, value_a in batch_a.items(): + if key == "samples" or key not in batch_b: + continue + value_b = batch_b[key] + if not isinstance(value_a, torch.Tensor) or not isinstance(value_b, torch.Tensor): + continue + if value_a.shape[0] != samples_a.shape[0] or value_b.shape[0] != samples_b.shape[0]: + continue + output[key] = _interleave_tensors(value_a, value_b) + return output + + def _ensure_video_path(video: Any) -> tuple[str, bool]: if isinstance(video, str): return video, False @@ -160,3 +207,34 @@ class FFmpegMergeFrames: finally: if os.path.isdir(frames_dir): shutil.rmtree(frames_dir, ignore_errors=True) + + +class InterleaveFrames: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "batch_a": ("IMAGE,LATENT", {"tooltip": "First batch to interleave."}), + "batch_b": ("IMAGE,LATENT", {"tooltip": "Second batch to interleave."}), + } + } + + RETURN_TYPES = ("IMAGE,LATENT",) + RETURN_NAMES = ("batch",) + FUNCTION = "interleave" + CATEGORY = "ffmpeg" + + def interleave(self, batch_a, batch_b): + type_a = _detect_batch_type(batch_a) + type_b = _detect_batch_type(batch_b) + if type_a != type_b: + raise ValueError("Input types must match (IMAGE+IMAGE or LATENT+LATENT).") + + if type_a == "IMAGE": + if not isinstance(batch_a, torch.Tensor) or not isinstance(batch_b, torch.Tensor): + raise ValueError("IMAGE inputs must be torch tensors.") + return (_interleave_tensors(batch_a, batch_b),) + + if not isinstance(batch_a, dict) or not isinstance(batch_b, dict): + raise ValueError("LATENT inputs must be dicts.") + return (_interleave_latents(batch_a, batch_b),)