From e577e2819f9e550bb72dbd917b18352a6bdfc146 Mon Sep 17 00:00:00 2001 From: Kangle <147648913+RCAKangle@users.noreply.github.com> Date: Fri, 30 Jan 2026 11:36:55 +0000 Subject: [PATCH] Support interleave up to 6 inputs --- nodes.py | 106 +++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 71 insertions(+), 35 deletions(-) diff --git a/nodes.py b/nodes.py index 6f22df6..2f70790 100644 --- a/nodes.py +++ b/nodes.py @@ -26,21 +26,32 @@ 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]: +def _interleave_many_tensors(batches: list[torch.Tensor]) -> torch.Tensor: + if len(batches) < 2: + raise ValueError("At least two batches are required.") + for batch in batches: + if batch.dim() < 1: + raise ValueError("Inputs must include a batch dimension.") + base_shape = batches[0].shape[1:] + for batch in batches[1:]: + if batch.shape[1:] != base_shape: + raise ValueError("All inputs must have the same shape (except batch size).") + + lengths = [batch.shape[0] for batch in batches] + if len(set(lengths)) > 1: logging.warning( - "Interleave Frames: batch sizes differ (%s vs %s); truncating.", - batch_a.shape[0], - batch_b.shape[0], + "Interleave Frames: batch sizes differ %s; truncating to shortest.", + lengths, ) - count = min(batch_a.shape[0], batch_b.shape[0]) + count = min(lengths) 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:]) + trimmed = [batch[:count] for batch in batches] + return torch.stack(trimmed, dim=1).reshape(count * len(trimmed), *base_shape) + + +def _interleave_tensors(batch_a: torch.Tensor, batch_b: torch.Tensor) -> torch.Tensor: + return _interleave_many_tensors([batch_a, batch_b]) def _detect_batch_type(value: Any) -> str: @@ -53,22 +64,29 @@ def _detect_batch_type(value: Any) -> str: 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.") +def _interleave_latent_batches(batches: list[dict]) -> dict: + samples = [batch["samples"] for batch in batches] + for sample in samples: + if not isinstance(sample, 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: + output = {"samples": _interleave_many_tensors(samples)} + common_keys = set(batches[0].keys()) + for batch in batches[1:]: + common_keys &= set(batch.keys()) + common_keys.discard("samples") + + for key in common_keys: + values = [batch[key] for batch in batches] + if not all(isinstance(value, torch.Tensor) for value in values): continue - value_b = batch_b[key] - if not isinstance(value_a, torch.Tensor) or not isinstance(value_b, torch.Tensor): + sample_lengths = [sample.shape[0] for sample in samples] + if any(value.shape[0] != length for value, length in zip(values, sample_lengths)): continue - if value_a.shape[0] != samples_a.shape[0] or value_b.shape[0] != samples_b.shape[0]: + try: + output[key] = _interleave_many_tensors(values) + except ValueError: continue - output[key] = _interleave_tensors(value_a, value_b) return output @@ -214,8 +232,14 @@ class InterleaveFrames: def INPUT_TYPES(cls): return { "required": { - "batch_a": ("IMAGE,LATENT", {"tooltip": "First batch to interleave."}), - "batch_b": ("IMAGE,LATENT", {"tooltip": "Second batch to interleave."}), + "batch1": ("IMAGE,LATENT", {"tooltip": "First batch to interleave."}), + "batch2": ("IMAGE,LATENT", {"tooltip": "Second batch to interleave."}), + }, + "optional": { + "batch3": ("IMAGE,LATENT", {"tooltip": "Third batch to interleave."}), + "batch4": ("IMAGE,LATENT", {"tooltip": "Fourth batch to interleave."}), + "batch5": ("IMAGE,LATENT", {"tooltip": "Fifth batch to interleave."}), + "batch6": ("IMAGE,LATENT", {"tooltip": "Sixth batch to interleave."}), } } @@ -224,17 +248,29 @@ class InterleaveFrames: 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).") + def interleave( + self, + batch1, + batch2, + batch3=None, + batch4=None, + batch5=None, + batch6=None, + ): + batches = [batch for batch in [batch1, batch2, batch3, batch4, batch5, batch6] if batch is not None] + if len(batches) < 2: + raise ValueError("At least two batches must be provided.") - if type_a == "IMAGE": - if not isinstance(batch_a, torch.Tensor) or not isinstance(batch_b, torch.Tensor): + types = {_detect_batch_type(batch) for batch in batches} + if len(types) != 1: + raise ValueError("All inputs must have the same type (IMAGE or LATENT).") + batch_type = types.pop() + + if batch_type == "IMAGE": + if not all(isinstance(batch, torch.Tensor) for batch in batches): raise ValueError("IMAGE inputs must be torch tensors.") - return (_interleave_tensors(batch_a, batch_b),) + return (_interleave_many_tensors(batches),) - if not isinstance(batch_a, dict) or not isinstance(batch_b, dict): + if not all(isinstance(batch, dict) for batch in batches): raise ValueError("LATENT inputs must be dicts.") - return (_interleave_latents(batch_a, batch_b),) + return (_interleave_latent_batches(batches),)